fix(agent-core): stop loop after aborted tool run (#94412)

szsip239 · vincentkoc · web-flow · commit 7fe287b0d30d · 2026-06-22T18:04:50.000+08:00
Merged via squash. Prepared head SHA: e11d971 Co-authored-by: szsip239 <88223778+szsip239@users.noreply.github.com> Co-authored-by: vincentkoc <25068+vincentkoc@users.noreply.github.com> Reviewed-by: @vincentkoc
diff --git a/packages/agent-core/src/agent-loop.test.ts b/packages/agent-core/src/agent-loop.test.ts
@@ -767,6 +767,154 @@ describe("agentLoop tool termination", () => {
     expect(events.filter((event) => event.type === "tool_execution_start")).toHaveLength(1);
     expect(events.at(-1)).toMatchObject({ type: "agent_end" });
   });
+
+  it("does not request another model turn after a tool aborts the run", async () => {
+    const controller = new AbortController();
+    let streamCalls = 0;
+    const streamFn: StreamFn = () => {
+      streamCalls += 1;
+      if (streamCalls > 1) {
+        throw new Error("model was called after abort");
+      }
+      const stream = createAssistantMessageEventStream();
+      queueMicrotask(() => {
+        const message = makeAssistantMessage([
+          { type: "toolCall", id: "call-abort", name: "abort_tool", arguments: {} },
+        ]);
+        stream.push({ type: "done", reason: "toolUse", message });
+        stream.end();
+      });
+      return stream;
+    };
+    const abortTool: AgentTool = {
+      name: "abort_tool",
+      label: "abort_tool",
+      description: "Abort the active run",
+      parameters: Type.Object({}, { additionalProperties: false }),
+      execute: async () => {
+        controller.abort(new Error("user aborted"));
+        return {
+          content: [{ type: "text", text: "aborted" }],
+          details: { aborted: true },
+        };
+      },
+    };
+    const events: AgentEvent[] = [];
+
+    const messages = await runAgentLoop(
+      [{ role: "user", content: "abort during tool", timestamp: 1 }],
+      {
+        systemPrompt: "",
+        messages: [],
+        tools: [abortTool],
+      },
+      config,
+      (event) => {
+        events.push(event);
+      },
+      controller.signal,
+      streamFn,
+    );
+
+    expect(streamCalls).toBe(1);
+    expect(messages.map((message) => message.role)).toEqual([
+      "user",
+      "assistant",
+      "toolResult",
+      "assistant",
+    ]);
+    expect(messages.at(-1)).toMatchObject({ role: "assistant", stopReason: "aborted" });
+    expect(events.map((event) => event.type)).toEqual([
+      "agent_start",
+      "turn_start",
+      "message_start",
+      "message_end",
+      "message_start",
+      "message_end",
+      "tool_execution_start",
+      "tool_execution_end",
+      "message_start",
+      "message_end",
+      "turn_end",
+      "turn_start",
+      "message_start",
+      "message_end",
+      "turn_end",
+      "agent_end",
+    ]);
+    expect(events.at(-1)).toMatchObject({ type: "agent_end" });
+  });
+
+  it("does not request another model turn when an async turn hook aborts the run", async () => {
+    const controller = new AbortController();
+    let streamCalls = 0;
+    const streamFn: StreamFn = () => {
+      streamCalls += 1;
+      if (streamCalls > 1) {
+        throw new Error("model was called after abort");
+      }
+      const stream = createAssistantMessageEventStream();
+      queueMicrotask(() => {
+        const message = makeAssistantMessage([
+          { type: "toolCall", id: "call-hook-abort", name: "hook_abort", arguments: {} },
+        ]);
+        stream.push({ type: "done", reason: "toolUse", message });
+        stream.end();
+      });
+      return stream;
+    };
+    const events: AgentEvent[] = [];
+
+    const messages = await runAgentLoop(
+      [{ role: "user", content: "abort from hook", timestamp: 1 }],
+      {
+        systemPrompt: "",
+        messages: [],
+        tools: [makeTool("hook_abort", [])],
+      },
+      {
+        ...config,
+        prepareNextTurn: async () => {
+          await Promise.resolve();
+          controller.abort(new Error("user aborted"));
+          return undefined;
+        },
+      },
+      (event) => {
+        events.push(event);
+      },
+      controller.signal,
+      streamFn,
+    );
+
+    expect(streamCalls).toBe(1);
+    expect(messages.map((message) => message.role)).toEqual([
+      "user",
+      "assistant",
+      "toolResult",
+      "assistant",
+    ]);
+    expect(messages.at(-1)).toMatchObject({ role: "assistant", stopReason: "aborted" });
+    expect(events.map((event) => event.type)).toEqual([
+      "agent_start",
+      "turn_start",
+      "message_start",
+      "message_end",
+      "message_start",
+      "message_end",
+      "tool_execution_start",
+      "tool_execution_end",
+      "message_start",
+      "message_end",
+      "turn_end",
+      "turn_start",
+      "message_start",
+      "message_end",
+      "turn_end",
+      "agent_end",
+    ]);
+    expect(events.at(-1)).toMatchObject({ type: "agent_end" });
+  });
 });
 
 describe("agentLoop thinking state", () => {
diff --git a/packages/agent-core/src/agent-loop.ts b/packages/agent-core/src/agent-loop.ts
@@ -267,17 +267,46 @@ async function runLoop(
   let currentContext = initialContext;
   let config = initialConfig;
   let firstTurn = true;
+  let turnOpen = true;
   // Check for steering messages at start (user may have typed while waiting)
   let pendingMessages: AgentMessage[] = (await config.getSteeringMessages?.()) || [];
+  const stopIfAborted = async (): Promise<boolean> => {
+    if (!signal?.aborted) {
+      return false;
+    }
+    // Persist an aborted assistant outcome so session post-processing does not
+    // compact or continue from the preceding toolUse message.
+    const abortedMessage = createLoopFailureMessage(
+      config,
+      signal.reason instanceof Error ? signal.reason : new Error("Agent run aborted"),
+      true,
+    );
+    newMessages.push(abortedMessage);
+    if (!turnOpen) {
+      await emit({ type: "turn_start" });
+      turnOpen = true;
+    }
+    await emit({ type: "message_start", message: abortedMessage });
+    await emit({ type: "message_end", message: abortedMessage });
+    await emit({ type: "turn_end", message: abortedMessage, toolResults: [] });
+    turnOpen = false;
+    await emit({ type: "agent_end", messages: newMessages });
+    return true;
+  };
 
   // Outer loop: continues when queued follow-up messages arrive after agent would stop
   while (true) {
     let hasMoreToolCalls = true;
 
     // Inner loop: process tool calls and steering messages
     while (hasMoreToolCalls || pendingMessages.length > 0) {
+      if (await stopIfAborted()) {
+        return;
+      }
+
       if (!firstTurn) {
         await emit({ type: "turn_start" });
+        turnOpen = true;
       } else {
         firstTurn = false;
       }
@@ -292,6 +321,10 @@ async function runLoop(
         }
       }
 
+      if (await stopIfAborted()) {
+        return;
+      }
+
       // Stream assistant response
       const message = await streamAssistantResponse(
         currentContext,
@@ -332,6 +365,10 @@ async function runLoop(
       }
 
       await emit({ type: "turn_end", message, toolResults });
+      turnOpen = false;
+      if (await stopIfAborted()) {
+        return;
+      }
 
       const nextTurnContext = {
         message,
@@ -357,6 +394,9 @@ async function runLoop(
           reasoning: nextReasoning,
         });
       }
+      if (await stopIfAborted()) {
+        return;
+      }
 
       if (
         await config.shouldStopAfterTurn?.({
@@ -371,6 +411,9 @@ async function runLoop(
       }
 
       pendingMessages = (await config.getSteeringMessages?.()) || [];
+      if (await stopIfAborted()) {
+        return;
+      }
     }
 
     const followUpMessages = (await config.getFollowUpMessages?.()) || [];
