feat(rust): plumb CancellationToken through ToolInvocation for handler cancellation#1706
Open
gimenete wants to merge 4 commits into
Open
feat(rust): plumb CancellationToken through ToolInvocation for handler cancellation#1706gimenete wants to merge 4 commits into
gimenete wants to merge 4 commits into
Conversation
…r cancellation Add `cancellation_token: CancellationToken` to `ToolInvocation` so that handlers can cooperatively stop work when `Session::abort()` is called. Changes: - Add `cancellation_token` field to `ToolInvocation` in `types.rs`; skipped in serde so it has no effect on JSON serialization. - Add `tool_abort: Arc<ParkingLotMutex<CancellationToken>>` to `Session`; initialized as a child of `shutdown` on create/resume. - `Session::abort()` now cancels the current `tool_abort` token (firing all child tokens held by in-flight handlers) and replaces it with a fresh child of `shutdown` so subsequent tool calls are not pre-cancelled. - Thread `tool_abort` through `spawn_event_loop` and `handle_notification`; each `ToolInvocation` receives a fresh child token. - Update existing tests to use `..Default::default()` for the new field. - Add two unit tests verifying token cancellation semantics. Backwards compatible: handlers that don't consume the token continue to work unchanged. Handlers that opt in can check `invocation.cancellation_token.is_cancelled()` or `select!` on `invocation.cancellation_token.cancelled()`. Closes github#1433 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
This PR adds cooperative cancellation for tool handlers by wiring a CancellationToken into each ToolInvocation and having Session::abort() signal all in-flight tool work.
Changes:
- Added
cancellation_token: CancellationTokentoToolInvocation(skipped for serde) so handlers can stop early on abort. - Implemented session-level tool-abort broadcast token that
Session::abort()cancels and the event loop uses to derive per-invocation child tokens. - Updated and extended Rust tests to use
Defaultstruct update syntax and to validate token cancellation behavior.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| rust/src/types.rs | Extends ToolInvocation with a cancellation token and documents intended cancellation semantics. |
| rust/src/session.rs | Introduces a shared abort token, cancels it on abort(), and injects per-tool child tokens during dispatch. |
| rust/src/tool.rs | Updates tests for new ToolInvocation field and adds cancellation-token unit tests. |
Comments suppressed due to low confidence (1)
rust/src/types.rs:1
- Adding a new public field to a public struct is a semver-breaking change for downstream callers that construct
ToolInvocationvia struct literals. If this crate has external consumers, consider mitigating by providing (and documenting) a constructor (e.g.,ToolInvocation::new(...)) and encouraging struct update syntax, or by making the new field private with an accessor so future additions don’t force call-site changes.
//! Protocol types shared between the SDK and the GitHub Copilot CLI.
Replace the single shared `tool_abort` token with a per-invocation map `in_flight_tool_calls: Arc<ParkingLotMutex<HashMap<String, CancellationToken>>>` keyed by `tool_call_id`. Changes: - `Session`: replace `tool_abort` field with `in_flight_tool_calls` (`Arc<ParkingLotMutex<HashMap<String, CancellationToken>>>`). - `handle_notification`: create a fresh `shutdown.child_token()` per dispatch, insert under `tool_call_id` before awaiting the handler, remove after the handler resolves (success, error, or cancellation). - `Session::abort()`: iterate the map and cancel every token; no longer replaces a shared token. - New `Session::cancel_tool_call(&self, tool_call_id: &str) -> bool`: cancels only the named in-flight handler, returns false for unknown IDs. - Add three unit tests for the map mechanics: cancel_tool_call_cancels_only_the_targeted_handler, cancel_tool_call_returns_false_for_unknown_id, abort_cancels_all_in_flight_tokens. - Add README section documenting cancellation behavior and cancel_tool_call with a code example. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Remove stale doc paragraph from ToolInvocation::cancellation_token that incorrectly claimed post-abort dispatches receive a pre-cancelled token. With per-call tracking each invocation gets a fresh child of shutdown, so that statement was never true under the new implementation. - Add async end-to-end test abort_unblocks_handler_awaiting_cancellation: simulates the real dispatch path (insert token into map, pass to handler task that awaits cancelled(), cancel map entries via abort logic) and asserts the handler unblocks within 1 second. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Change the field type from `CancellationToken` to `Option<CancellationToken>` so that constructing a `ToolInvocation` never requires supplying a token. Combined with the existing `#[non_exhaustive]` attribute, this keeps the addition backwards compatible: external callers and tests can build the struct without referencing the new field, and `Default` yields `None`. The SDK wraps the per-call child token in `Some(...)` at dispatch; handlers treat `None` as 'no cancellation signal available'. Updated the README example and tests accordingly, and added a test asserting the field defaults to `None`. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Implements issue #1433 for the Rust SDK: propagates cancellation to in-flight tool handlers when
Session::abort()is called.Changes
types.rs: Addcancellation_token: CancellationTokenfield toToolInvocation. Skipped during serde serialization/deserialization so it has no effect on the JSON wire format. Defaults to a freshCancellationToken::new()(never-cancelled) so existing code constructing invocations directly continues to compile.session.rs:tool_abort: Arc<ParkingLotMutex<CancellationToken>>toSession, initialized as a child of the session'sshutdowntoken.Session::abort()now cancels the currenttool_aborttoken (which fires all child tokens held by in-flight handlers) and replaces it with a fresh child ofshutdown, so subsequent tool invocations are not pre-cancelled.tool_abortthroughspawn_event_loop→handle_notification. EachToolInvocationgets a fresh child token at dispatch time.tool.rs: Update testToolInvocationconstructions to use..Default::default()for the new field; add two unit tests covering token cancellation semantics.API
Handlers that opt in can check
invocation.cancellation_token.is_cancelled()orselect!oninvocation.cancellation_token.cancelled():Backwards compatibility
Handlers that don't consume the token continue to work unchanged.
ToolInvocationis#[non_exhaustive]so adding the field is non-breaking for external consumers.