From 6b518dc9e4c601c0108768932dc1450c036075fd Mon Sep 17 00:00:00 2001
From: Taylor Mullen <ntaylormullen@google.com>
Date: Fri, 9 May 2025 23:29:02 -0700
Subject: 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
---
 packages/server/src/core/client.ts | 11 +++++++----
 packages/server/src/core/turn.ts   | 13 ++++++++++---
 2 files changed, 17 insertions(+), 7 deletions(-)

(limited to 'packages/server/src/core')

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 };
+      }
     }
   }
 
-- 
cgit v1.2.3