Skip to content

fix: close nested agent streams when consumers stop early - #1042

Draft
u7-u7 wants to merge 1 commit into
modelscope:mainfrom
u7-u7:fix/close-nested-agent-streams-20261031
Draft

u7-u7 wants to merge 1 commit into
modelscope:mainfrom
u7-u7:fix/close-nested-agent-streams-20261031

Conversation

@u7-u7

@u7-u7 u7-u7 commented Oct 31, 2026

Copy link
Copy Markdown

Change Summary

Closing an SDK agent stream after consuming a chunk currently closes only the outer async generator. Python's async for does not automatically close the iterator it consumes. As a result, the nested retry/step/model generators remain suspended: the HTTP model stream stays open and run_loop's existing interrupted-round persistence is deferred to eventual async-generator finalization.

Use contextlib.aclosing at the three ownership boundaries (run(stream=True) → run_loop → retried step) so explicit consumer shutdown reaches the model stream before aclose() returns. This does not change retries, payloads, or normal streaming output.

Validation

  • New regression exercises step, run_loop, and the public run(stream=True) against a real OpenAI-compatible SDK stream using httpx.MockTransport. All three failed on the base (model stream outlived its consumer), and pass after the fix; run-loop/public-run cases also verify interrupted-round persistence is invoked.
  • A focused retry-wrapper regression verifies that closing runs the active attempt's finalizer before returning and does not start another retry.
  • pytest -q tests/utils/test_retry_classification.py tests/agent/test_llm_io.py tests/agent/test_partial_round.py: 59 passed, 15 skipped.
  • Flake8 for changed production files, YAPF for changed sections, and git diff --check: passed.

Not tested: live providers, a full initialized tool/MCP session, the entire repository suite, or other Python versions. Optional Anthropic cases were skipped because its dependency is not installed. Tests use the real stream/agent loop with unrelated setup mocked; these are local results, not upstream CI.

AI assistance

Implementation and automated verification were performed with OpenAI Codex. No human-review attestation is made.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

1 participant