From f2ea78d0e4e5d25ab3cc25dc9f1492135630c9be Mon Sep 17 00:00:00 2001
From: "N. Taylor Mullen" <ntaylormullen@google.com>
Date: Sun, 8 Jun 2025 15:42:49 -0700
Subject: fix(tool-scheduler): Correctly pipe cancellation signal to tool calls
 (#852)

---
 packages/cli/src/ui/hooks/useReactToolScheduler.ts | 20 ++++++++------------
 1 file changed, 8 insertions(+), 12 deletions(-)

(limited to 'packages/cli/src/ui/hooks/useReactToolScheduler.ts')

diff --git a/packages/cli/src/ui/hooks/useReactToolScheduler.ts b/packages/cli/src/ui/hooks/useReactToolScheduler.ts
index 32a3e77f..ae58ed38 100644
--- a/packages/cli/src/ui/hooks/useReactToolScheduler.ts
+++ b/packages/cli/src/ui/hooks/useReactToolScheduler.ts
@@ -32,8 +32,8 @@ import {
 
 export type ScheduleFn = (
   request: ToolCallRequestInfo | ToolCallRequestInfo[],
+  signal: AbortSignal,
 ) => void;
-export type CancelFn = (reason?: string) => void;
 export type MarkToolsAsSubmittedFn = (callIds: string[]) => void;
 
 export type TrackedScheduledToolCall = ScheduledToolCall & {
@@ -69,7 +69,7 @@ export function useReactToolScheduler(
   setPendingHistoryItem: React.Dispatch<
     React.SetStateAction<HistoryItemWithoutId | null>
   >,
-): [TrackedToolCall[], ScheduleFn, CancelFn, MarkToolsAsSubmittedFn] {
+): [TrackedToolCall[], ScheduleFn, MarkToolsAsSubmittedFn] {
   const [toolCallsForDisplay, setToolCallsForDisplay] = useState<
     TrackedToolCall[]
   >([]);
@@ -172,15 +172,11 @@ export function useReactToolScheduler(
   );
 
   const schedule: ScheduleFn = useCallback(
-    async (request: ToolCallRequestInfo | ToolCallRequestInfo[]) => {
-      scheduler.schedule(request);
-    },
-    [scheduler],
-  );
-
-  const cancel: CancelFn = useCallback(
-    (reason: string = 'unspecified') => {
-      scheduler.cancelAll(reason);
+    async (
+      request: ToolCallRequestInfo | ToolCallRequestInfo[],
+      signal: AbortSignal,
+    ) => {
+      scheduler.schedule(request, signal);
     },
     [scheduler],
   );
@@ -198,7 +194,7 @@ export function useReactToolScheduler(
     [],
   );
 
-  return [toolCallsForDisplay, schedule, cancel, markToolsAsSubmitted];
+  return [toolCallsForDisplay, schedule, markToolsAsSubmitted];
 }
 
 /**
-- 
cgit v1.2.3