Repository navigation
fix: include completed child usage in lead budget checks - #6218
Conversation
willem-bd
left a comment
There was a problem hiding this comment.
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.
|
@quifox, thanks for your contribution. Please check out the review comment first. |
|
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. |
Condense the subagent, loop, and token budget guard descriptions while retaining their limits, lifecycle rules, and completed child usage ordering.
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
Bug fix verification
backend/tests/test_lead_token_budget.pyValidation
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.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.