Skip to content

fix: include completed child usage in lead budget checks - #6218

Merged
WillemJiang merged 4 commits into
bytedance:mainfrom
quifox:fix/lead-budget
Oct 4, 2026
Merged

WillemJiang merged 4 commits into
bytedance:mainfrom
quifox:fix/lead-budget

Conversation

@quifox

@quifox quifox commented Oct 3, 2026

Copy link
Copy Markdown
Contributor

Why

Completed subtask usage could be missing from the lead's running total when its budget was checked, causing the configured budget to undercount completed work.

What changed

Completed subtask usage is now backfilled directly into the dispatch message before budget enforcement. Current-step attribution remains after the guards that can change tool calls, and terminal tool results are marked to prevent duplicate attribution. The behavior remains gated by the existing usage and budget settings.

Surface area

  • Agents / LangGraph - lead usage accounting
  • Default behavior change

Bug fix verification

  • Reproducer: backend/tests/test_lead_token_budget.py
  • The regression was reproduced on the unmodified code and is covered by the fixed test.

Validation

  • cd backend && .venv/bin/python -m pytest -q tests/test_lead_token_budget.py tests/test_token_budget_middleware.py tests/test_token_usage_middleware.py tests/test_lead_agent_model_resolution.py tests/test_loop_detection_middleware.py tests/test_agent_guidance_check.py: 268 passed.
  • Ruff check, Ruff format check, and repository guidance tests: passed.
  • The full backend suite is not clean; this result is limited to the targeted suites above.
  • Targeted tests: backend/tests/test_lead_token_budget.py, backend/tests/test_token_budget_middleware.py, backend/tests/test_token_usage_middleware.py, backend/tests/test_lead_agent_model_resolution.py, backend/tests/test_loop_detection_middleware.py, backend/tests/test_agent_guidance_check.py.
  • Boundary: controlled graph and usage cases cover sync and async runs, same-run continuation, guard ordering, and usage/budget gating; they do not measure provider billing.

@CLAassistant

CLAassistant commented Oct 3, 2026 •

Copy link
Copy Markdown

CLA assistant check
All committers have signed the CLA.

@github-actions github-actions Bot added area:agents Agents, subagents, graph wiring, prompts, langgraph.json area:docs Documentation and Markdown only needs-validation Touches front/back contract surface; needs real-path validation risk:high High risk: backend API, agents, sandbox, auth, deps, CI size/L PR changes 300-700 lines labels Oct 3, 2026

@willem-bd willem-bd left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed the head revision. The core mechanics hold up: langchain 1.3.14 compiles each middleware's after_model into its own graph node dispatched in reverse list order with state applied between nodes, so registering CompletedSubagentUsageMiddleware after TokenBudgetMiddleware really does backfill terminal child usage before enforcement, and the budget's per-message seen-map (rebuilt in before_agent) counts the bumped usage_metadata exactly once across same-run continuations. The _attribute_completed_subagent_usage extraction is a faithful move and the attribution flag keeps the second invocation (from TokenUsageMiddleware) idempotent. Omitting the middleware from the subagent chain built in tool_error_handling_middleware.py is correct since subagents disallow the task tool. Two suggestions below, neither blocking.

Comment thread backend/tests/test_lead_token_budget.py Resolved
@WillemJiang

Copy link
Copy Markdown
Collaborator

@quifox, thanks for your contribution. Please check out the review comment first.

@quifox

quifox commented Oct 3, 2026

Copy link
Copy Markdown
Contributor Author

Thanks for the review. The test now imports the new middleware directly, and the SDK factory gap is kept separate from this lead-agent fix. The related targeted suites pass (268 tests). Further contributions will be discussed individually after the review on the existing PRs is addressed.

@WillemJiang WillemJiang added this to the 2.2.0 milestone Oct 4, 2026
Condense the subagent, loop, and token budget guard descriptions while retaining their limits, lifecycle rules, and completed child usage ordering.
@WillemJiang
WillemJiang merged commit e8de758 into bytedance:main Oct 4, 2026
15 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:agents Agents, subagents, graph wiring, prompts, langgraph.json area:docs Documentation and Markdown only needs-validation Touches front/back contract surface; needs real-path validation risk:high High risk: backend API, agents, sandbox, auth, deps, CI size/L PR changes 300-700 lines

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants