Skip to content

fix: track concurrent streaming tool calls - #7059

Open
CorgiBoyG wants to merge 1 commit into
google:mainfrom
CorgiBoyG:fix/streaming-tool-task-tracking
Open

fix: track concurrent streaming tool calls#7059
CorgiBoyG wants to merge 1 commit into
google:mainfrom
CorgiBoyG:fix/streaming-tool-task-tracking

Conversation

@CorgiBoyG

Copy link
Copy Markdown

Summary

  • track every active call when a live response invokes the same streaming tool concurrently
  • give each input-streaming call its own LiveRequestQueue and broadcast live input to all active call streams
  • cancel and release all matching tasks on stop_streaming and live-run teardown while preserving later registrations

Fixes #7058

Problem

_process_function_live_helper() stored one task per tool name in InvocationContext.active_streaming_tools. If a live model response invoked the same async-generator tool more than once, each registration replaced the previous task reference. The earlier call kept running but was no longer reachable by stop_streaming() or run teardown, so it could continue writing responses after its agent run ended. Input-streaming calls also shared or replaced the single stream slot.

This violated the lifecycle invariant that every background tool call started by a live run remains tracked until it completes or the run stops it.

Solution

ActiveStreamingTool now maintains a private task-to-stream registry while preserving its existing task and stream compatibility fields. Each call gets an independent input queue. Completion callbacks release task and stream references, stop_streaming operates on a snapshot so it cannot remove calls registered later, and run teardown enumerates every active task.

Testing

  • Added a regression test through handle_function_calls_live() with two same-name calls and distinct input streams. It fails on main because one task remains pending after teardown.
  • Covered both run teardown and stop_streaming cancellation paths.
  • Covered natural completion and error-path reference release.
  • 206 passed across the complete streaming suite and directly related flow/tool tests.
  • 13,971 passed, 85 skipped, 27 xfailed, 2 xpassed in the full unit suite. Two import-loading tests fail identically on unchanged main because this host Python preloads sitecustomize; the same two tests were rerun against upstream/main to confirm the baseline.
  • The concurrent regression test passed 30 repeated runs with PYTHONASYNCIODEBUG=1.
  • All pre-commit hooks passed except update-constraints, which was stopped after more than 18 minutes without output; this change does not modify dependencies or constraint files.
  • Mypy reports no issues in the changed task-tracking and live-flow source files. A separate pre-existing error in function_tool.py reproduces unchanged on upstream/main.

Risk

The change is limited to live async-generator tool bookkeeping and input-stream injection. Existing name-based stop_streaming behavior remains intact: it stops all active calls for that tool name. Empty registry placeholders and the public compatibility fields retain their existing observable behavior, while completed tasks and streams are released.

@google-cla

google-cla Bot commented Sep 8, 2026

Copy link
Copy Markdown

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

@CorgiBoyG
CorgiBoyG force-pushed the fix/streaming-tool-task-tracking branch from 4f68ffe to 33602a5 Compare September 8, 2026 10:26
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Parallel calls to the same streaming tool leave an orphan task after live teardown

2 participants