go: plumb context cancellation through ToolInvocation for session.Abort()#1705
Open
gimenete wants to merge 3 commits into
Open
go: plumb context cancellation through ToolInvocation for session.Abort()#1705gimenete wants to merge 3 commits into
gimenete wants to merge 3 commits into
Conversation
…rt() Add a Context field to ToolInvocation that is cancelled when session.Abort is called. This lets tool handlers cooperatively stop in-flight work (HTTP requests, DB queries, sleeps) when the session is aborted, without requiring OS-level process kills. Changes: - Add Context context.Context to ToolInvocation; it carries OTel trace propagation AND is cancelled on Abort() - Deprecate TraceContext in favour of Context (same value, new name that better communicates its dual role) - Track per-tool-call cancel funcs in Session.toolCallCancels; clean up after each call completes - In Abort(), cancel all in-flight cancel funcs after the RPC call - RPC response calls (HandlePendingToolCall) use traceCtx (non-cancellable) so they succeed even when the handler context was aborted - Add TestToolInvocation_ContextCancelledOnAbort and TestToolInvocation_ContextPopulated unit tests - Document the cancellation pattern in go/README.md Fixes 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 introduces cooperative cancellation for tool handlers by propagating session.Abort() into a new ToolInvocation.Context, while preserving trace propagation for OpenTelemetry.
Changes:
- Add
ToolInvocation.Context(cancelled onSession.Abort) and deprecate direct usage ofToolInvocation.TraceContext. - Track in-flight tool call cancel functions in
Sessionand cancel them duringAbort. - Add tests and documentation describing cancellation behavior for tool handlers.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| go/types.go | Adds ToolInvocation.Context and deprecation guidance for TraceContext. |
| go/session.go | Implements per-tool-call cancellation tracking and abort-triggered cancellation; uses a non-cancelled context for RPC responses. |
| go/session_test.go | Adds tests validating handler context cancellation on abort and context population. |
| go/README.md | Documents cooperative cancellation via ToolInvocation.Context. |
Add Session.CancelToolCall(toolCallID string) bool — cancels a single in-flight tool handler by its tool call ID without affecting other concurrent handlers or the agentic loop. - CancelToolCall looks up the cancel func in toolCallCancels, calls it, removes the entry, and returns true; returns false if not in flight - Abort() now also removes entries after cancelling them, so a subsequent CancelToolCall on the same ID correctly returns false - Extract newRPCDrainSession test helper to eliminate boilerplate across the three tool-context tests - Add TestCancelToolCall_cancelsTargetedHandlerOnly covering: unknown ID returns false, targeted handler context is cancelled, sibling handler context stays live, idempotent false on second call - Document CancelToolCall in go/README.md alongside Abort cancellation Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Set TraceContext to traceCtx (non-cancellable) to preserve backward compatibility; TraceContext is no longer affected by session.Abort() or CancelToolCall(). Only Context (the new field) is cancelled. - Fix newRPCDrainSession drain goroutine: replace bufio.Scanner + io.ReadFull on the same reader (which causes buffer-ahead corruption) with a single bufio.Reader used for both header line reads and body ReadFull calls. - Update TestToolInvocation_ContextPopulated assertion: Context and TraceContext are now different instances (Context is a cancellable child of TraceContext), so assert they differ rather than being equal. - Clarify TraceContext deprecation comment: document it is never cancelled. 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
Closes #1433 (Go SDK portion)
When
session.Abort()is called, in-flight tool handlers now receive a cooperative cancellation signal viaToolInvocation.Context. A newsession.CancelToolCall(toolCallID)method lets callers cancel a single handler without affecting others or the agentic loop.Changes
go/types.goContext context.ContexttoToolInvocation— carries W3C Trace Context (OTel) and is cancelled whensession.Abort()orsession.CancelToolCall()is called.TraceContextin favour ofContext(same value; existing code usingTraceContextcontinues to work unchanged).go/session.gotoolCallCancels map[string]context.CancelFunc+ mutex toSessionto track in-flight tool call cancel funcs.executeToolAndRespondcreates acontext.WithCancelchild of the trace context and registers its cancel func; cleans up (deregisters + callscancel()) when the handler returns or panics.HandlePendingToolCall) usetraceCtx(non-cancellable) so they succeed even ifAbort()already fired.Abort()cancels all stored cancel funcs and removes them from the map after the RPC acknowledgment.CancelToolCall(toolCallID string) bool— cancels a single in-flight handler by ID. Returnstrueif found and cancelled,falseif not in flight. Idempotent.go/session_test.gonewRPCDrainSessionhelper (pipe-based mock RPC, drains all requests with empty success responses).TestToolInvocation_ContextCancelledOnAbort— handler context cancelled by abort simulation; map cleaned up.TestToolInvocation_ContextPopulated—ContextandTraceContextboth set and equal.TestCancelToolCall_cancelsTargetedHandlerOnly— unknown ID returns false; targeted handler cancelled; sibling handler stays live; second call on same ID returns false.go/README.mdCancelToolCallexample andToolCallIDcapture pattern.Backwards compatibility
Existing handlers that ignore
inv.Context/inv.TraceContextare unaffected.TraceContextholds the same value as before (additive change).