Enable tools to cancel active execution.

- Plumbed abort signals through to tools
- Updated the shell tool to properly cancel active requests by killing the entire child process tree of the underlying shell process and then report that the shell itself was canceled.

Fixes https://b.corp.google.com/issues/416829935

2 files changed, 17 insertions, 7 deletions

diff --git a/packages/server/src/core/client.ts b/packages/server/src/core/client.ts
index 904e944c..46af465a 100644
--- a/packages/server/src/core/client.ts
+++ b/packages/server/src/core/client.ts
@@ -64,10 +64,13 @@ export class GeminiClient {
           .getTool('read_many_files') as ReadManyFilesTool;
         if (readManyFilesTool) {
           // Read all files in the target directory
-          const result = await readManyFilesTool.execute({
-            paths: ['**/*'], // Read everything recursively
-            useDefaultExcludes: true, // Use default excludes
-          });
+          const result = await readManyFilesTool.execute(
+            {
+              paths: ['**/*'], // Read everything recursively
+              useDefaultExcludes: true, // Use default excludes
+            },
+            AbortSignal.timeout(30000),
+          );
           if (result.llmContent) {
             initialParts.push({
               text: `\n--- Full File Context ---\n${result.llmContent}`,
diff --git a/packages/server/src/core/turn.ts b/packages/server/src/core/turn.ts
index 7d8bf7b6..62219938 100644
--- a/packages/server/src/core/turn.ts
+++ b/packages/server/src/core/turn.ts
@@ -36,7 +36,10 @@ export interface ServerTool {
   name: string;
   schema: FunctionDeclaration;
   // The execute method signature might differ slightly or be wrapped
-  execute(params: Record<string, unknown>): Promise<ToolResult>;
+  execute(
+    params: Record<string, unknown>,
+    signal?: AbortSignal,
+  ): Promise<ToolResult>;
   shouldConfirmExecute(
     params: Record<string, unknown>,
   ): Promise<ToolCallConfirmationDetails | false>;
@@ -153,7 +156,7 @@ export class Turn {
           if (confirmationDetails) {
             return { ...pendingToolCall, confirmationDetails };
           }
-          const result = await tool.execute(pendingToolCall.args);
+          const result = await tool.execute(pendingToolCall.args, signal);
           return {
             ...pendingToolCall,
             result,
@@ -199,7 +202,11 @@ export class Turn {
         resultDisplay: outcome.result?.returnDisplay,
         error: outcome.error,
       };
-      yield { type: GeminiEventType.ToolCallResponse, value: responseInfo };
+
+      // If aborted we're already yielding the user cancellations elsewhere.
+      if (!signal?.aborted) {
+        yield { type: GeminiEventType.ToolCallResponse, value: responseInfo };
+      }
     }
   }