Skip to content

fix(runners): run after_run_callback when an aborted run is cancelled - #7403

Open
abhayjoshi201 wants to merge 1 commit into
google:mainfrom
abhayjoshi201:fix/after-run-callback-on-cancel
Open

abhayjoshi201 wants to merge 1 commit into
google:mainfrom
abhayjoshi201:fix/after-run-callback-on-cancel

Conversation

@abhayjoshi201

Copy link
Copy Markdown
Contributor

Link to Issue or Description of Change

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

Problem:
When an invocation is aborted via abort_signal and the task driving Runner.run_async is cancelled (such as on POST /run_sse client disconnect, where abort_signal.set() and producer_task.cancel() are called in the same turn), Runner._exec_with_plugin and run_node_async treated asyncio.CancelledError as an unhandled run_error unless _CALLER_CLOSED_EARLY_MSG was passed. As a result, after_run_callback was skipped in the finally block even though invocation_context.is_aborted was True and dangling function calls were sealed.

Additionally, POST /run in src/google/adk/cli/api_server.py did not pass or trip abort_signal in its http.disconnect monitor before calling worker_task.cancel(), so client disconnects on /run neither sealed dangling function calls (INVOCATION_ABORTED) nor executed after_run_callback.

Solution:

  1. In Runner._exec_with_plugin (src/google/adk/runners.py) and run_node_async (src/google/adk/workflow/_node_runner_utils.py), treat asyncio.CancelledError when invocation_context.is_aborted (ic.is_aborted) is True as an early close (closing_early = True, leaving run_error = None) so after_run_callback executes after sealing any dangling function calls.
  2. In POST /run (src/google/adk/cli/api_server.py), pass abort_signal to runner.run_async(...) and call abort_signal.set() before worker_task.cancel() on http.disconnect, matching POST /run_sse.

Testing Plan

Unit Tests:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.
pytest tests/unittests/test_runners.py tests/unittests/cli/test_fast_api.py tests/unittests/workflow/test_workflow.py tests/unittests/plugins/test_notification_error_callbacks.py
460 passed, 5 skipped, 1 xfailed in 36.77s

Manual End-to-End (E2E) Tests:

Verified with SlowToolAgent and _AfterRunPlugin across Runner.run_async, POST /run, and POST /run_sse that disconnecting a client or cancelling an aborted invocation task seals dangling FunctionCalls with INVOCATION_ABORTED and executes after_run_callback.

Checklist

  • I have read the CONTRIBUTING.md document.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes.
  • I have manually tested my changes end-to-end.
  • Any dependent changes have been merged and published in downstream modules.

When an invocation is aborted via abort_signal and the task driving Runner.run_async is cancelled (such as on POST /run_sse or POST /run client disconnect), Runner._exec_with_plugin and run_node_async previously treated asyncio.CancelledError as an unhandled run_error unless _CALLER_CLOSED_EARLY_MSG was passed, skipping after_run_callback in finally.

Treat asyncio.CancelledError when invocation_context.is_aborted is set as an early close so after_run_callback runs after sealing any dangling function calls, and wire abort_signal into POST /run's disconnect monitor to match POST /run_sse.

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

after_run_callback is not called when the task driving Runner.run_async is cancelled (regression in 2.11.0)

2 participants