Skip to content

fix: keep the results of finished parallel tool calls on abort - #7992

Open
AkshayShah03 wants to merge 1 commit into
google:mainfrom
AkshayShah03:fix/abort-keep-finished-tool-results
Open

AkshayShah03 wants to merge 1 commit into
google:mainfrom
AkshayShah03:fix/abort-keep-finished-tool-results

Conversation

@AkshayShah03

Copy link
Copy Markdown
Contributor

Link to Issue or Description of Change

1. Link to an existing issue (if applicable):

Problem:
When an invocation was aborted while one call of a parallel batch was still running, the calls that had already finished were sealed as aborted along with it. That covers abort_signal.set(), a client disconnect on /run or /run_sse, and cancelling the task driving run_async. Their real results and state_delta were dropped, so the model was told a tool that had run did not, and could call it again (for example, send the same email twice).

_gather_or_cancel re-raises CancelledError after cancelling the unfinished calls. _execute_prepared_function_calls only kept finished results for Exception / NodeInterruptedError (#7428), so on cancellation they were discarded, and _build_abort_events sealed every call in the batch.

Solution:
As suggested on #7958:

  • _execute_prepared_function_calls also catches asyncio.CancelledError. It merges the plain results of the calls that finished, using the same _is_plain_result filter and merge as the error path, and stores them on the invocation's _AbortState. Then it re-raises. Live mode is left out, as in the error path.
  • Runner._synthesize_abort_events_if_needed appends those events through _process_and_append_event before calling _build_abort_events. That function only seals calls without a response, so it now covers just the unfinished ones.

The kept event is the normal merged response event, so its state_delta and artifact_delta are persisted too. It's also yielded to the caller along with the abort events, as the synthetic responses already are.

If the caller closes the stream early (aclose() / break), nothing is sealed, as before, so nothing is recorded either.

Behaviour change: an aborted batch now records finished calls as finished instead of aborted. Per the issue, this is for maintainers to confirm.

Question: the error path still skips aborted invocations (_keep_completed_responses returns early when is_aborted). So when a sibling raises in the same step an abort lands, the finished call is still sealed as aborted. test_run_async_tool_error_on_abort_leaves_calls_to_abort_sealing asserts that. I left it unchanged to keep this PR to the cancellation case. Should that path record the finished result too, for consistency?

Testing Plan

Unit Tests:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.

New test in tests/unittests/test_runners.py: test_run_async_abort_keeps_results_of_finished_parallel_calls. It is parametrized over the abort route (abort_signal.set() / task.cancel()) and over a plain LlmAgent or one inside a Workflow. It checks:

  • the finished call's real response is persisted and only the running call is sealed;
  • the finished tool's state_delta reaches session state;
  • for the plain agent, the next turn's model request contains the real result.

All 4 cases fail on main.

Full unit suite: the same 6 failures as main (test_google_llm continuation-token tests, test_release_dependencies, and test_app_details, which also fails on main in a full run). Nothing new.

Manual End-to-End (E2E) Tests:

The repro from #7958 (send_email + a slow lookup_order, aborted with abort_signal.set()):

send_email actually ran: ['a@example.com']
state emails_sent: 1
recorded response: send_email {'status': 'sent', 'to': 'a@example.com'}
recorded response: lookup_order {'error': 'Invocation was aborted by client.'}

The same holds when the consumer task is cancelled instead.

When an invocation was aborted (abort_signal, a client disconnect, or
cancelling the task driving run_async) while one call of a parallel batch was
still running, the calls that had already finished were sealed as aborted
along with it. Their real results and state_delta were dropped, so the model
was told a tool that had run did not, and could call it again.

On cancellation the batch executor now keeps the merged response event of the
plain results that finished, as it already does when a sibling call raises,
on the invocation's abort state. Abort sealing records it before building the
synthetic responses, so only the calls that were still running are sealed.

Fixes google#7958
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.

Aborting during parallel tool calls records tools that already finished as "aborted" and drops their state changes

2 participants