feat(agent): harden run lifecycle and abort handling

Author: marslavish

packages/agent/__tests__/agent.test.ts

@@ -347,13 +347,9 @@ describe('@agentic-kit/agent — pausable tools', () => {
     expect(() => agent.continue()).toThrow(/no tool calls awaiting a decision/);
   });
 
-  it('continue() resumes from a non-trailing assistant when a later message was appended after the pause', async () => {
-    const provider = createScriptedProvider({ responses: [pauseResponse(), finalResponse()] });
-    const execute = jest.fn(
-      async (_id: string, _params: Record<string, unknown>, decision: unknown) => ({
-        content: [{ type: 'text' as const, text: `decision=${JSON.stringify(decision)}` }],
-      })
-    );
+  it('continue() rejects when non-toolResult messages have been appended after the pending assistant', async () => {
+    const provider = createScriptedProvider({ responses: [pauseResponse()] });
+    const execute = jest.fn();
 
     const agent = new Agent({
       initialState: { model: makeFakeModel() },
@@ -372,14 +368,8 @@ describe('@agentic-kit/agent — pausable tools', () => {
     };
     agent.replaceMessages([...agent.state.messages, trailingNote]);
 
-    await agent.continue();
-
-    expect(execute).toHaveBeenCalledTimes(1);
-    expect(execute.mock.calls[0]?.[2]).toEqual({ approved: true });
-    expect(agent.state.messages.at(-1)).toMatchObject({
-      role: 'assistant',
-      content: [{ type: 'text', text: 'finalized' }],
-    });
+    expect(() => agent.continue()).toThrow(/non-toolResult messages have been appended/);
+    expect(execute).not.toHaveBeenCalled();
   });
 
   it('abort() while paused stops further work without throwing', async () => {
@@ -469,6 +459,67 @@ describe('@agentic-kit/agent — pausable tools', () => {
     });
   });
 
+  it('abort() during tool execution stops the loop and does not invoke the model again', async () => {
+    const provider = createScriptedProvider({
+      responses: [
+        makeFakeAssistantMessage({
+          stopReason: 'toolUse',
+          content: [
+            { type: 'toolCall', id: 'tool_slow', name: 'slow', arguments: {} },
+          ],
+        }),
+        makeFakeAssistantMessage({
+          stopReason: 'stop',
+          content: [{ type: 'text', text: 'should never reach here' }],
+        }),
+      ],
+    });
+
+    const streamCalls = jest.fn(provider.stream);
+    let abortAgent: (() => void) | undefined;
+    const slowExecute = jest.fn<
+      ReturnType<AgentTool['execute']>,
+      Parameters<AgentTool['execute']>
+    >((_id, _params, _decision, signal) => {
+      return new Promise((_resolve, reject) => {
+        signal?.addEventListener(
+          'abort',
+          () => reject(new Error('aborted')),
+          { once: true }
+        );
+        queueMicrotask(() => abortAgent?.());
+      });
+    });
+
+    const agent = new Agent({
+      initialState: { model: makeFakeModel() },
+      streamFn: streamCalls,
+    });
+    abortAgent = () => agent.abort();
+    agent.setTools([
+      {
+        name: 'slow',
+        label: 'Slow',
+        description: 'Slow tool',
+        parameters: { type: 'object', properties: {} },
+        execute: slowExecute,
+      },
+    ]);
+
+    const events: AgentEvent[] = [];
+    agent.subscribe((e) => events.push(e));
+
+    await agent.prompt('go');
+
+    expect(streamCalls).toHaveBeenCalledTimes(1);
+    expect(slowExecute).toHaveBeenCalledTimes(1);
+    const end = events.find(
+      (e): e is Extract<AgentEvent, { type: 'agent_end' }> => e.type === 'agent_end'
+    );
+    expect(end?.stopReason).toBe('aborted');
+    expect(agent.state.isStreaming).toBe(false);
+  });
+
   it('regression: a tool without a decision schema runs without pausing', async () => {
     const provider = createScriptedProvider({
       responses: [

packages/agent/src/agent.ts

@@ -36,6 +36,7 @@ export class Agent {
   private abortController?: AbortController;
   private running?: Promise<void>;
   private runChannel?: { push: RunChannelPush };
+  private outstandingHandle?: AgentRunHandle;
 
   private _state: AgentState;
 
@@ -103,41 +104,51 @@ export class Agent {
 
   abort(): void {
     this.abortController?.abort();
+    this.outstandingHandle = undefined;
   }
 
   waitForIdle(): Promise<void> {
     return this.running ?? Promise.resolve();
   }
 
   prompt(input: string | Message, opts?: { maxSteps?: number }): AgentRunHandle {
-    if (this._state.isStreaming) {
-      throw new Error('Agent is already processing a prompt');
-    }
+    this.assertIdle('prompt');
 
     const message = typeof input === 'string' ? createUserMessage(input) : input;
     this._state.stepCount = 0;
 
-    return new DefaultAgentRunHandle(async (push, signal) =>
-      this.runLoop({
+    const handle: AgentRunHandle = new DefaultAgentRunHandle(async (push, signal) => {
+      if (this.outstandingHandle === handle) {
+        this.outstandingHandle = undefined;
+      }
+      return this.runLoop({
         initialMessages: [message],
         externalPush: push ?? undefined,
         externalAbortSignal: signal,
         maxSteps: opts?.maxSteps ?? this.defaultMaxSteps,
-      })
-    );
+      });
+    });
+    this.outstandingHandle = handle;
+    return handle;
   }
 
   continue(opts?: { maxSteps?: number }): AgentRunHandle {
-    if (this._state.isStreaming) {
-      throw new Error('Agent is already processing');
-    }
+    this.assertIdle('continue');
 
     if (this._state.messages.length === 0) {
       throw new Error('No messages to continue from');
     }
 
     const pendingMessage = this.findMostRecentPendingAssistant();
     if (pendingMessage) {
+      const pendingIndex = this._state.messages.indexOf(pendingMessage);
+      const trailing = this._state.messages.slice(pendingIndex + 1);
+      const hasNonToolResultTrailing = trailing.some((m) => m.role !== 'toolResult');
+      if (hasNonToolResultTrailing) {
+        throw new Error(
+          'Cannot continue() with a pending decision when non-toolResult messages have been appended after the pending assistant. Use injectDeferralResults() + prompt() instead — see the agentic-kit deferral docs.'
+        );
+      }
       const pendingDecisions = this.findPendingDecisions(pendingMessage);
       for (const { tool, decision } of pendingDecisions) {
         const errors = validateSchema(tool.decision!, decision, 'root');
@@ -154,13 +165,29 @@ export class Agent {
       }
     }
 
-    return new DefaultAgentRunHandle(async (push, signal) =>
-      this.runLoop({
+    const handle: AgentRunHandle = new DefaultAgentRunHandle(async (push, signal) => {
+      if (this.outstandingHandle === handle) {
+        this.outstandingHandle = undefined;
+      }
+      return this.runLoop({
         externalPush: push ?? undefined,
         externalAbortSignal: signal,
         maxSteps: opts?.maxSteps ?? this.defaultMaxSteps,
-      })
-    );
+      });
+    });
+    this.outstandingHandle = handle;
+    return handle;
+  }
+
+  private assertIdle(method: 'prompt' | 'continue'): void {
+    if (this._state.isStreaming) {
+      throw new Error(`Agent is already processing; cannot call ${method}() while a run is active`);
+    }
+    if (this.outstandingHandle) {
+      throw new Error(
+        `Agent has an unconsumed run handle from a previous ${method}()/prompt()/continue() call; consume it (events / toReadableStream / toResponse / wait) or abort the agent before issuing another`
+      );
+    }
   }
 
   private findMostRecentPendingAssistant(): AssistantMessage | undefined {
@@ -227,7 +254,7 @@ export class Agent {
         }
       }
 
-      let stopReason: 'completed' | 'max_steps' = 'completed';
+      let stopReason: 'completed' | 'max_steps' | 'aborted' = 'completed';
 
       try {
         await this.emit({ type: 'agent_start' });
@@ -286,6 +313,11 @@ export class Agent {
           }
 
           await this.emit({ type: 'turn_end', message: assistantMessage, toolResults: outcome.results });
+
+          if (localAbortController.signal.aborted) {
+            stopReason = 'aborted';
+            break;
+          }
         }
 
         await this.emit({ type: 'agent_end', messages: [...this._state.messages], stopReason });
@@ -372,6 +404,10 @@ export class Agent {
         continue;
       }
 
+      if (signal.aborted) {
+        break;
+      }
+
       const tool = this._state.tools.find((candidate) => candidate.name === toolCall.name);
       const args = toolCall.arguments as Record<string, unknown>;
       const decisionAttached = 'decision' in toolCall && toolCall.decision !== undefined;

packages/agent/src/types.ts

@@ -49,7 +49,7 @@ export interface AgentEventBase {
 
 export type AgentEvent =
   | { type: 'agent_start' }
-  | { type: 'agent_end'; messages: Message[]; stopReason?: 'completed' | 'max_steps' }
+  | { type: 'agent_end'; messages: Message[]; stopReason?: 'completed' | 'max_steps' | 'aborted' }
   | { type: 'turn_start' }
   | { type: 'turn_end'; message: AssistantMessage; toolResults: ToolResultMessage[] }
   | { type: 'message_start'; message: Message }