go: plumb context cancellation through ToolInvocation for session.Abort()

go/README.md

@@ -355,6 +355,65 @@ session, _ := client.CreateSession(context.Background(), &copilot.SessionConfig{
 
 When the model selects a tool, the SDK automatically runs your handler (in parallel with other calls) and responds to the CLI's `tool.call` with the handler's result.
 
+#### Cooperative Cancellation via session.Abort
+
+`ToolInvocation.Context` is a `context.Context` that is cancelled when `session.Abort` is called. Pass it to any cancellable operation (HTTP requests, DB queries, sleeps) so the handler stops promptly when the session is aborted:
+
+```go
+lookupIssue := copilot.DefineTool("lookup_issue", "Fetch issue details from our tracker",
+    func(params LookupIssueParams, inv copilot.ToolInvocation) (any, error) {
+        // Pass inv.Context so the HTTP request is cancelled on session.Abort.
+        req, err := http.NewRequestWithContext(inv.Context, "GET",
+            "https://api.example.com/issues/"+params.ID, nil)
+        if err != nil {
+            return nil, err
+        }
+        resp, err := http.DefaultClient.Do(req)
+        if err != nil {
+            return nil, err // returns context.Canceled when aborted
+        }
+        defer resp.Body.Close()
+        // ...
+        return summary, nil
+    })
+```
+
+Handlers that don't use `inv.Context` are unaffected; they run to completion as before.
+
+#### Cancelling a single tool call
+
+Use `session.CancelToolCall(toolCallID)` to cancel one specific in-flight handler without aborting the session or any other concurrent handlers. It returns `true` if the tool call was found and cancelled, `false` if it was not in flight.
+
+```go
+// Cancel a specific tool call by its ID.
+if cancelled := session.CancelToolCall(toolCallID); !cancelled {
+    log.Println("tool call was not in flight")
+}
+```
+
+`toolCallID` is available inside the handler as `inv.ToolCallID`. You can capture it to enable external cancellation of a specific operation:
+
+```go
+var mu sync.Mutex
+activeCalls := map[string]string{} // label → toolCallID
+
+slowTool := copilot.DefineTool("slow_op", "A long-running operation",
+    func(params SlowOpParams, inv copilot.ToolInvocation) (any, error) {
+        mu.Lock()
+        activeCalls[params.Label] = inv.ToolCallID
+        mu.Unlock()
+
+        // ... do work, checking inv.Context.Done() ...
+        return result, nil
+    })
+
+// Elsewhere, cancel by label:
+mu.Lock()
+id := activeCalls["my-label"]
+mu.Unlock()
+session.CancelToolCall(id)
+```
+
 #### Overriding Built-in Tools
 
 If you register a tool with the same name as a built-in CLI tool (e.g. `edit_file`, `read_file`), the SDK will throw an error unless you explicitly opt in by setting `OverridesBuiltInTool = true`. This flag signals that you intend to replace the built-in tool with your custom implementation.

go/session.go

@@ -82,6 +82,11 @@ type Session struct {
 	capabilities          SessionCapabilities
 	capabilitiesMu        sync.RWMutex
 
+	// toolCallCancels tracks cancel functions for in-flight tool calls so that
+	// Abort can propagate cancellation into handler contexts.
+	toolCallCancels   map[string]context.CancelFunc
+	toolCallCancelsMu sync.Mutex
+
 	// eventCh serializes user event handler dispatch. dispatchEvent enqueues;
 	// a single goroutine (processEvents) dequeues and invokes handlers in FIFO order.
 	eventCh   chan SessionEvent
@@ -1337,11 +1342,35 @@ func (s *Session) handleBroadcastEvent(event SessionEvent) {
 
 // executeToolAndRespond executes a tool handler and sends the result back via RPC.
 func (s *Session) executeToolAndRespond(requestID, toolName, toolCallID string, arguments any, handler ToolHandler, traceparent, tracestate string) {
-	ctx := contextWithTraceParent(context.Background(), traceparent, tracestate)
+	// traceCtx carries OTel trace propagation but is not subject to abort cancellation.
+	// It is used for administrative RPC calls that must complete regardless of abort.
+	traceCtx := contextWithTraceParent(context.Background(), traceparent, tracestate)
+	// ctx is passed to the tool handler and is cancelled when session.Abort is called,
+	// giving handlers a cooperative cancellation signal.
+	ctx, cancel := context.WithCancel(traceCtx)
+
+	s.toolCallCancelsMu.Lock()
+	if s.toolCallCancels == nil {
+		s.toolCallCancels = make(map[string]context.CancelFunc)
+	}
+	s.toolCallCancels[toolCallID] = cancel
+	s.toolCallCancelsMu.Unlock()
+
+	// Cleanup runs last (registered first). Removes the cancel from the in-flight map
+	// and releases context resources.
+	defer func() {
+		s.toolCallCancelsMu.Lock()
+		delete(s.toolCallCancels, toolCallID)
+		s.toolCallCancelsMu.Unlock()
+		cancel()
+	}()
+
+	// Panic recovery runs first (registered second, LIFO). Uses traceCtx to ensure
+	// the error response is sent even if ctx was already cancelled by Abort.
 	defer func() {
 		if r := recover(); r != nil {
 			errMsg := fmt.Sprintf("tool panic: %v", r)
-			s.RPC.Tools.HandlePendingToolCall(ctx, &rpc.HandlePendingToolCallRequest{
+			s.RPC.Tools.HandlePendingToolCall(traceCtx, &rpc.HandlePendingToolCallRequest{
 				RequestID: requestID,
 				Error:     &errMsg,
 			})
@@ -1353,13 +1382,14 @@ func (s *Session) executeToolAndRespond(requestID, toolName, toolCallID string,
 		ToolCallID:   toolCallID,
 		ToolName:     toolName,
 		Arguments:    arguments,
-		TraceContext: ctx,
+		Context:      ctx,
+		TraceContext: traceCtx,
 	}
 
 	result, err := handler(invocation)
 	if err != nil {
 		errMsg := err.Error()
-		s.RPC.Tools.HandlePendingToolCall(ctx, &rpc.HandlePendingToolCallRequest{
+		s.RPC.Tools.HandlePendingToolCall(traceCtx, &rpc.HandlePendingToolCallRequest{
 			RequestID: requestID,
 			Error:     &errMsg,
 		})
@@ -1389,7 +1419,7 @@ func (s *Session) executeToolAndRespond(requestID, toolName, toolCallID string,
 	if result.Error != "" {
 		rpcResult.Error = &result.Error
 	}
-	s.RPC.Tools.HandlePendingToolCall(ctx, &rpc.HandlePendingToolCallRequest{
+	s.RPC.Tools.HandlePendingToolCall(traceCtx, &rpc.HandlePendingToolCallRequest{
 		RequestID: requestID,
 		Result:    rpcResult,
 	})
@@ -1555,9 +1585,44 @@ func (s *Session) Abort(ctx context.Context) error {
 		return fmt.Errorf("failed to abort session: %w", err)
 	}
 
+	s.toolCallCancelsMu.Lock()
+	for id, cancel := range s.toolCallCancels {
+		cancel()
+		delete(s.toolCallCancels, id)
+	}
+	s.toolCallCancelsMu.Unlock()
+
 	return nil
 }
 
+// CancelToolCall cancels a single in-flight tool handler identified by toolCallID
+// without aborting the agentic loop or any other concurrent tool handlers.
+//
+// It looks up the cancel func registered when the handler was dispatched, calls it
+// (cancelling the context passed to that handler via ToolInvocation.Context), removes
+// the entry from the in-flight map, and returns true. If no handler with the given
+// toolCallID is currently executing, CancelToolCall is a no-op and returns false.
+//
+// Example:
+//
+//	// Start a session with a long-running tool registered.
+//	// Later, cancel only a specific tool call without aborting the session:
+//	if cancelled := session.CancelToolCall("tool-call-id-123"); !cancelled {
+//	    log.Println("tool call was not in flight")
+//	}
+func (s *Session) CancelToolCall(toolCallID string) bool {
+	s.toolCallCancelsMu.Lock()
+	defer s.toolCallCancelsMu.Unlock()
+
+	cancel, ok := s.toolCallCancels[toolCallID]
+	if !ok {
+		return false
+	}
+	cancel()
+	delete(s.toolCallCancels, toolCallID)
+	return true
+}
+
 // SetModelOptions configures optional parameters for SetModel.
 type SetModelOptions struct {
 	// ReasoningEffort sets the reasoning effort level for the new model (e.g., "low", "medium", "high", "xhigh").