Repository navigation
Python: clear function invocation budget state from session on failed turn - #10146
Johnny Wilson Dougherty (JohnnyWilson16) wants to merge 2 commits into
Conversation
|
Johnny Wilson Dougherty (@JohnnyWilson16) please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.
Contributor License AgreementContribution License AgreementThis Contribution License Agreement (“Agreement”) is agreed to by the party signing below (“You”),
|
1 similar comment
|
Johnny Wilson Dougherty (@JohnnyWilson16) please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.
Contributor License AgreementContribution License AgreementThis Contribution License Agreement (“Agreement”) is agreed to by the party signing below (“You”),
|
There was a problem hiding this comment.
🟡 Changes recommended
Closing a partially consumed public stream no longer awaits provider cleanup or propagates provider-close errors because the inner invocation generator is not explicitly closed, violating the existing lifecycle contract.
1 open issue
What changed in this PR
This PR addresses #9677 by clearing failed-turn function-invocation budgets from Python agent sessions, so retries do not inherit depleted limits while pending approvals retain their budget.
Changes:
- Adds
finallycleanup around streaming and non-streaming invocation loops. - Adds a regression test for provider failures in both modes.
- Adds implementation plans for budget cleanup and a separate Windows shell-parsing change.
| File | Description |
|---|---|
python/packages/core/tests/core/test_harness_tool_approval.py |
Tests session budget removal after provider errors in both modes. |
python/packages/core/agent_framework/_tools.py |
Wraps invocation dispatch with budget cleanup. |
docs/superpowers/plans/2026-11-06-issue-9819-windows-shell-parsing.md |
Documents a separate Windows command-line parsing proposal. |
docs/superpowers/plans/2026-11-06-issue-9677-tool-approval-budget-cleanup.md |
Documents the budget cleanup implementation and validation plan. |
Review effort: Balanced
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Motivation & Context
When
ToolApprovalMiddlewareis active, it maintains function-invocation budget tracking (attempt_count,total_function_calls, andstart_time) insession.state["_function_invocation_budget_state"]across approval boundaries.If an execution turn terminates due to a provider error, timeout, or task cancellation,
_clear_budget_state_from_sessionwas not called (it only ran on clean completion or in_response_invalidation_cleanup). Consequently, subsequent retries on the same session inherited the depleted budget from the failed turn, prematurely triggering tool limits or turning off tool execution.Description & Review Guide
_tools.py(FunctionInvocationLayer.get_response):_run_non_streaming()withtry...finally: _clear_budget_state_from_session(invocation_session, budget_state)ensuring budget state is cleaned up after unhandled exceptions and cancellations._run_streaming()withtry...finally: _clear_budget_state_from_session(invocation_session, budget_state)ensuring budget state is cleaned up on generator termination, error, or early close._clear_budget_state_from_sessionchecks for pending approvals insession.state[_TOOL_APPROVAL_STATE_KEY], retaining budget state during legitimate human-in-the-loop pauses while purging it on errors.test_tool_approval_clears_budget_on_turn_failureintest_harness_tool_approval.pyvalidating that transient turn errors clear the session budget across streaming and non-streaming modes.get_responseand parity between streaming and non-streaming cleanup.Related Issue
Fixes #9677
AI Assistance
AI assistance details: Used an AI coding assistant to help reproduce the session budget leak, structure the unwinding cleanup in the invocation layer, write the regression tests, and prepare this description.
Contribution Checklist