feat(agent): maxSteps + lookup decision by id

Author: marslavish

apps/nextjs-chat-demo/src/app/api/chat/route.ts

@@ -68,6 +68,7 @@ export async function POST(req: Request): Promise<Response> {
   const agent = new Agent({
     initialState: { model, tools, systemPrompt: SYSTEM_PROMPT },
     streamFn: (m, ctx, opts) => adapter.stream(m, ctx, opts),
+    maxSteps: 5,
   });
 
   const isResume = lastMessageHasPendingDecision(messages);

packages/agent/__tests__/agent.test.ts

@@ -2,6 +2,7 @@ import {
   type AssistantMessage,
   type Context,
   createAssistantMessageEventStream,
+  type Message,
   type ModelDescriptor,
   type ToolCallContent,
 } from 'agentic-kit';
@@ -346,6 +347,41 @@ 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)}` }],
+      })
+    );
+
+    const agent = new Agent({
+      initialState: { model: makeFakeModel() },
+      streamFn: provider.stream,
+    });
+    agent.setTools([makeApprovalTool(execute)]);
+
+    await agent.prompt('approve thing');
+
+    attachDecision(agent, 'tool_1', { approved: true });
+
+    const trailingNote: Message = {
+      role: 'user',
+      content: 'side note injected by an external queue while paused',
+      timestamp: Date.now(),
+    };
+    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' }],
+    });
+  });
+
   it('abort() while paused stops further work without throwing', async () => {
     const provider = createScriptedProvider({ responses: [pauseResponse()] });
 
@@ -480,3 +516,124 @@ describe('@agentic-kit/agent — pausable tools', () => {
     expect(events.some((e) => e.type === 'agent_end')).toBe(true);
   });
 });
+
+describe('@agentic-kit/agent — maxSteps', () => {
+  function makeEchoTool(): AgentTool {
+    return {
+      name: 'echo',
+      label: 'Echo',
+      description: 'Echo text',
+      parameters: {
+        type: 'object',
+        properties: { text: { type: 'string' } },
+        required: ['text'],
+      },
+      execute: async (_id, params) => ({
+        content: [{ type: 'text', text: String(params.text) }],
+      }),
+    };
+  }
+
+  function toolThenText(toolText = 'one', finalText = 'done') {
+    return [
+      makeFakeAssistantMessage({
+        stopReason: 'toolUse',
+        content: [
+          { type: 'toolCall', id: 'tool_1', name: 'echo', arguments: { text: toolText } },
+        ],
+      }),
+      makeFakeAssistantMessage({
+        stopReason: 'stop',
+        content: [{ type: 'text', text: finalText }],
+      }),
+    ];
+  }
+
+  it('halts after the configured number of model calls and emits agent_end with stopReason=max_steps', async () => {
+    const provider = createScriptedProvider({ responses: toolThenText() });
+    const agent = new Agent({
+      initialState: { model: makeFakeModel() },
+      streamFn: provider.stream,
+      maxSteps: 1,
+    });
+    agent.setTools([makeEchoTool()]);
+
+    const events: AgentEvent[] = [];
+    agent.subscribe((e) => events.push(e));
+
+    await agent.prompt('go');
+
+    expect(agent.state.stepCount).toBe(1);
+    // Tool ran for the first turn, but no second model call.
+    const toolResults = agent.state.messages.filter((m) => m.role === 'toolResult');
+    expect(toolResults).toHaveLength(1);
+    const assistants = agent.state.messages.filter((m) => m.role === 'assistant');
+    expect(assistants).toHaveLength(1);
+
+    const end = events.find((e) => e.type === 'agent_end');
+    expect(end).toMatchObject({ type: 'agent_end', stopReason: 'max_steps' });
+  });
+
+  it('does not enforce a cap when maxSteps is undefined (no behavior change)', async () => {
+    const provider = createScriptedProvider({ responses: toolThenText() });
+    const agent = new Agent({
+      initialState: { model: makeFakeModel() },
+      streamFn: provider.stream,
+    });
+    agent.setTools([makeEchoTool()]);
+
+    const events: AgentEvent[] = [];
+    agent.subscribe((e) => events.push(e));
+
+    await agent.prompt('go');
+
+    expect(agent.state.stepCount).toBe(2);
+    expect(agent.state.messages.at(-1)).toMatchObject({
+      role: 'assistant',
+      content: [{ type: 'text', text: 'done' }],
+    });
+    const end = events.find((e) => e.type === 'agent_end');
+    expect(end).toMatchObject({ stopReason: 'completed' });
+  });
+
+  it('per-call maxSteps overrides the constructor default', async () => {
+    const provider = createScriptedProvider({ responses: toolThenText() });
+    const agent = new Agent({
+      initialState: { model: makeFakeModel() },
+      streamFn: provider.stream,
+      maxSteps: 1, // would cap; per-call override allows the second call
+    });
+    agent.setTools([makeEchoTool()]);
+
+    await agent.prompt('go', { maxSteps: 5 });
+
+    expect(agent.state.stepCount).toBe(2);
+    expect(agent.state.messages.at(-1)).toMatchObject({
+      role: 'assistant',
+      content: [{ type: 'text', text: 'done' }],
+    });
+  });
+
+  it('prompt() resets stepCount; continue() preserves it across turns', async () => {
+    // Two prompt rounds: first one consumes 2 steps; second prompt resets to 0.
+    const responses = [
+      ...toolThenText('first', 'first-done'),
+      makeFakeAssistantMessage({
+        stopReason: 'stop',
+        content: [{ type: 'text', text: 'second-done' }],
+      }),
+    ];
+    const provider = createScriptedProvider({ responses });
+    const agent = new Agent({
+      initialState: { model: makeFakeModel() },
+      streamFn: provider.stream,
+    });
+    agent.setTools([makeEchoTool()]);
+
+    await agent.prompt('first');
+    expect(agent.state.stepCount).toBe(2);
+
+    await agent.prompt('second');
+    expect(agent.state.stepCount).toBe(1);
+  });
+});

packages/agent/src/agent.ts

@@ -32,6 +32,7 @@ export class Agent {
   private readonly transformContext?: AgentOptions['transformContext'];
   private readonly streamFn: NonNullable<AgentOptions['streamFn']>;
   private readonly validateToolArguments: NonNullable<AgentOptions['validateToolArguments']>;
+  private readonly defaultMaxSteps?: number;
   private abortController?: AbortController;
   private running?: Promise<void>;
   private runChannel?: { push: RunChannelPush };
@@ -44,13 +45,15 @@ export class Agent {
       tools: [],
       messages: [],
       isStreaming: false,
+      stepCount: 0,
       streamMessage: null,
       streamOptions: undefined,
       ...options.initialState,
     };
     this.streamFn = options.streamFn ?? stream;
     this.transformContext = options.transformContext;
     this.validateToolArguments = options.validateToolArguments ?? defaultValidateToolArguments;
+    this.defaultMaxSteps = options.maxSteps;
   }
 
   get state(): AgentState {
@@ -106,55 +109,70 @@ export class Agent {
     return this.running ?? Promise.resolve();
   }
 
-  prompt(input: string | Message): AgentRunHandle {
+  prompt(input: string | Message, opts?: { maxSteps?: number }): AgentRunHandle {
     if (this._state.isStreaming) {
       throw new Error('Agent is already processing a prompt');
     }
 
     const message = typeof input === 'string' ? createUserMessage(input) : input;
+    this._state.stepCount = 0;
 
     return new DefaultAgentRunHandle(async (push, signal) =>
       this.runLoop({
         initialMessages: [message],
         externalPush: push ?? undefined,
         externalAbortSignal: signal,
+        maxSteps: opts?.maxSteps ?? this.defaultMaxSteps,
       })
     );
   }
 
-  continue(): AgentRunHandle {
+  continue(opts?: { maxSteps?: number }): AgentRunHandle {
     if (this._state.isStreaming) {
       throw new Error('Agent is already processing');
     }
 
-    const lastMessage = this._state.messages[this._state.messages.length - 1];
-    if (!lastMessage) {
+    if (this._state.messages.length === 0) {
       throw new Error('No messages to continue from');
     }
 
-    if (lastMessage.role === 'assistant') {
-      const pendingDecisions = this.findPendingDecisions(lastMessage);
-      if (pendingDecisions.length === 0) {
-        throw new Error(
-          'Cannot continue from trailing assistant message: no tool calls awaiting a decision'
-        );
-      }
+    const pendingMessage = this.findMostRecentPendingAssistant();
+    if (pendingMessage) {
+      const pendingDecisions = this.findPendingDecisions(pendingMessage);
       for (const { tool, decision } of pendingDecisions) {
         const errors = validateSchema(tool.decision!, decision, 'root');
         if (errors.length > 0) {
           throw new DecisionValidationError(tool.name, errors);
         }
       }
+    } else {
+      const lastMessage = this._state.messages[this._state.messages.length - 1];
+      if (lastMessage.role === 'assistant') {
+        throw new Error(
+          'Cannot continue from trailing assistant message: no tool calls awaiting a decision'
+        );
+      }
     }
 
     return new DefaultAgentRunHandle(async (push, signal) =>
       this.runLoop({
         externalPush: push ?? undefined,
         externalAbortSignal: signal,
+        maxSteps: opts?.maxSteps ?? this.defaultMaxSteps,
       })
     );
   }
 
+  private findMostRecentPendingAssistant(): AssistantMessage | undefined {
+    for (let i = this._state.messages.length - 1; i >= 0; i--) {
+      const msg = this._state.messages[i];
+      if (msg.role !== 'assistant') continue;
+      const pending = this.findPendingDecisions(msg);
+      if (pending.length > 0) return msg;
+    }
+    return undefined;
+  }
+
   private findPendingDecisions(
     message: AssistantMessage
   ): Array<{ toolCall: ToolCallContent; tool: AgentTool; decision: unknown }> {
@@ -188,6 +206,7 @@ export class Agent {
     initialMessages?: Message[];
     externalPush?: RunChannelPush;
     externalAbortSignal?: AbortSignal;
+    maxSteps?: number;
   }): Promise<void> {
     this.running = (async () => {
       this.abortController = new AbortController();
@@ -208,6 +227,8 @@ export class Agent {
         }
       }
 
+      let stopReason: 'completed' | 'max_steps' = 'completed';
+
       try {
         await this.emit({ type: 'agent_start' });
 
@@ -219,20 +240,25 @@ export class Agent {
           }
         }
 
-        let resumingFromTrailingAssistant =
-          this._state.messages[this._state.messages.length - 1]?.role === 'assistant';
+        let resumeAssistant: AssistantMessage | undefined =
+          this.findMostRecentPendingAssistant();
 
         while (true) {
           let assistantMessage: AssistantMessage;
 
-          if (resumingFromTrailingAssistant) {
-            const last = this._state.messages[this._state.messages.length - 1];
-            if (!last || last.role !== 'assistant') {
-              throw new Error('Cannot resume: last message is not an assistant message');
-            }
-            assistantMessage = last;
-            resumingFromTrailingAssistant = false;
+          if (resumeAssistant) {
+            assistantMessage = resumeAssistant;
+            resumeAssistant = undefined;
           } else {
+            if (
+              opts.maxSteps !== undefined &&
+              this._state.stepCount >= opts.maxSteps
+            ) {
+              stopReason = 'max_steps';
+              break;
+            }
+            this._state.stepCount += 1;
+
             await this.emit({ type: 'turn_start' });
             assistantMessage = await this.generateAssistantMessage(localAbortController.signal);
             this.appendMessage(assistantMessage);
@@ -262,7 +288,7 @@ export class Agent {
           await this.emit({ type: 'turn_end', message: assistantMessage, toolResults: outcome.results });
         }
 
-        await this.emit({ type: 'agent_end', messages: [...this._state.messages] });
+        await this.emit({ type: 'agent_end', messages: [...this._state.messages], stopReason });
       } finally {
         if (opts.externalAbortSignal) {
           opts.externalAbortSignal.removeEventListener('abort', onExternalAbort);

packages/agent/src/types.ts

@@ -36,6 +36,7 @@ export interface AgentState {
   isStreaming: boolean;
   messages: Message[];
   model: ModelDescriptor;
+  stepCount: number;
   streamMessage: AssistantMessage | null;
   streamOptions?: Omit<StreamOptions, 'signal'>;
   systemPrompt: string;
@@ -48,7 +49,7 @@ export interface AgentEventBase {
 
 export type AgentEvent =
   | { type: 'agent_start' }
-  | { type: 'agent_end'; messages: Message[] }
+  | { type: 'agent_end'; messages: Message[]; stopReason?: 'completed' | 'max_steps' }
   | { type: 'turn_start' }
   | { type: 'turn_end'; message: AssistantMessage; toolResults: ToolResultMessage[] }
   | { type: 'message_start'; message: Message }
@@ -79,6 +80,12 @@ export type AgentEvent =
 
 export interface AgentOptions {
   initialState: Pick<AgentState, 'model'> & Partial<Omit<AgentState, 'model'>>;
+  /**
+   * Maximum number of model invocations the agent will perform per run.
+   * One model call counts as one step. Counter persists across `continue()`
+   * — it only resets in `prompt()`. Default: unlimited.
+   */
+  maxSteps?: number;
   streamFn?: (
     model: ModelDescriptor,
     context: Context,