Skip to content

fix(cli): keep after_run callbacks when /run client disconnects - #7408

Open
CodeAlex52 wants to merge 1 commit into
google:mainfrom
CodeAlex52:fix/api-server-disconnect-after-run-callback
Open

CodeAlex52 wants to merge 1 commit into
google:mainfrom
CodeAlex52:fix/api-server-disconnect-after-run-callback

Conversation

@CodeAlex52

Copy link
Copy Markdown

Fixes #7394.

Problem

When a client disconnects from /run, the disconnect monitor cancels the worker task with a bare Task.cancel(). The runner treats a bare cancellation as an external cancel, which deliberately skips after_run_callback (see test_run_async_cancellation_does_not_execute_after_run_plugin), so a plugin that finalizes a run — writing a terminal event to the session, telemetry, cleanup — gets no callback at all for a disconnected run.

That path already exists: cancelling with the caller-closed-early marker makes the runner treat the run as "the caller stopped consuming events", which does run the after-run callbacks. /run_sse gets this for free because it closes the generator; the non-streaming /run endpoint is the only path that cancels without the marker.

Fix

One line in cli/api_server.py, plus the test that pins it:

worker_task.cancel(_CALLER_CLOSED_EARLY_MSG)

I deliberately did not change the bare-cancellation semantics in runners.py / workflow/_node_runner_utils.py. Flipping those would contradict the existing intentional behavior and its test, and it would change what happens for genuine external cancellations (shutdown, a caller cancelling its own task). The disconnect case is a caller going away, which is exactly what the marker means, so the fix belongs at the call site.

The error path is unaffected: the generator still re-raises CancelledError, the worker task still ends cancelled, and /run still answers 499.

Tests

Added test_agent_run_disconnect_marks_cancellation_as_caller_closed_early next to the existing disconnect test in tests/unittests/cli/test_fast_api.py, reusing the same harness and asserting the CancelledError that reaches run_async carries the marker.

Validation:

  • The new test fails against the current code (worker_task.cancel() reaches run_async with e.args == ()) and passes with the change.
  • The four pre-existing disconnect tests pass unchanged both with and without the change.
  • tests/unittests/test_runners.py, tests/unittests/plugins, tests/unittests/workflow: 1589 passed, 1 skipped, 5 xfailed.
  • tests/unittests/cli/test_fast_api.py: 176 passed with the change vs 175 on the unmodified tree; the 6 failures and 5 errors in that file are pre-existing in my environment (metrics/agent-identity tests and optional-dependency collection errors) and are identical before and after.
  • pyink --check clean on both files.

Note for discussion

@Li-john1021 commented on #7394 proposing to restore after_run_callback for any external cancel of Runner.run_async. That is a different, larger behavior change and conflicts with the deliberate semantics noted above — happy to discuss, but I think the disconnect path is the actual gap worth fixing first.

Signed-off-by: CodeAlex52 <59381946+CodeAlex52@users.noreply.github.com>
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