feat(agent): add post-turn stop callback
This commit is contained in:
@@ -2,6 +2,10 @@
|
|||||||
|
|
||||||
## [Unreleased]
|
## [Unreleased]
|
||||||
|
|
||||||
|
### Added
|
||||||
|
|
||||||
|
- Added `shouldStopAfterTurn` to the low-level agent loop config for gracefully exiting after a completed turn before polling queued messages or starting another LLM call.
|
||||||
|
|
||||||
## [0.71.1] - 2026-05-01
|
## [0.71.1] - 2026-05-01
|
||||||
|
|
||||||
## [0.71.0] - 2026-04-30
|
## [0.71.0] - 2026-04-30
|
||||||
|
|||||||
@@ -112,6 +112,20 @@ The `beforeToolCall` hook runs after `tool_execution_start` and validated argume
|
|||||||
|
|
||||||
Tools can also return `terminate: true` to hint that the automatic follow-up LLM call should be skipped. The loop only stops early when every finalized tool result in that batch sets `terminate: true`. Mixed batches continue normally.
|
Tools can also return `terminate: true` to hint that the automatic follow-up LLM call should be skipped. The loop only stops early when every finalized tool result in that batch sets `terminate: true`. Mixed batches continue normally.
|
||||||
|
|
||||||
|
Low-level loop callers can set `shouldStopAfterTurn` to stop gracefully after the current turn completes:
|
||||||
|
|
||||||
|
```typescript
|
||||||
|
const stream = agentLoop(prompts, context, {
|
||||||
|
model,
|
||||||
|
convertToLlm,
|
||||||
|
shouldStopAfterTurn: async ({ message, toolResults, context, newMessages }) => {
|
||||||
|
return shouldCompactBeforeNextTurn(context.messages);
|
||||||
|
},
|
||||||
|
});
|
||||||
|
```
|
||||||
|
|
||||||
|
`shouldStopAfterTurn` runs after `turn_end` is emitted and after the assistant response and any tool executions have completed normally. If it returns `true`, the loop emits `agent_end` and exits before polling steering or follow-up queues, and before starting another LLM call. It does not abort the provider stream, does not cancel running tools, and does not alter the assistant message stop reason.
|
||||||
|
|
||||||
When you use the `Agent` class, assistant `message_end` processing is treated as a barrier before tool preflight begins. That means `beforeToolCall` sees agent state that already includes the assistant message that requested the tool call.
|
When you use the `Agent` class, assistant `message_end` processing is treated as a barrier before tool preflight begins. That means `beforeToolCall` sees agent state that already includes the assistant message that requested the tool call.
|
||||||
|
|
||||||
### continue() Event Sequence
|
### continue() Event Sequence
|
||||||
|
|||||||
@@ -215,6 +215,18 @@ async function runLoop(
|
|||||||
|
|
||||||
await emit({ type: "turn_end", message, toolResults });
|
await emit({ type: "turn_end", message, toolResults });
|
||||||
|
|
||||||
|
if (
|
||||||
|
await config.shouldStopAfterTurn?.({
|
||||||
|
message,
|
||||||
|
toolResults,
|
||||||
|
context: currentContext,
|
||||||
|
newMessages,
|
||||||
|
})
|
||||||
|
) {
|
||||||
|
await emit({ type: "agent_end", messages: newMessages });
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
pendingMessages = (await config.getSteeringMessages?.()) || [];
|
pendingMessages = (await config.getSteeringMessages?.()) || [];
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -100,6 +100,18 @@ export interface AfterToolCallContext {
|
|||||||
context: AgentContext;
|
context: AgentContext;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/** Context passed to `shouldStopAfterTurn`. */
|
||||||
|
export interface ShouldStopAfterTurnContext {
|
||||||
|
/** The assistant message that completed the turn. */
|
||||||
|
message: AssistantMessage;
|
||||||
|
/** Tool result messages passed to the preceding `turn_end` event. */
|
||||||
|
toolResults: ToolResultMessage[];
|
||||||
|
/** Current agent context after the turn's assistant message and tool results have been appended. */
|
||||||
|
context: AgentContext;
|
||||||
|
/** Messages that this loop invocation will return if it exits at this point. Prompt runs include the initial prompt messages; continuation runs do not include pre-existing context messages. */
|
||||||
|
newMessages: AgentMessage[];
|
||||||
|
}
|
||||||
|
|
||||||
export interface AgentLoopConfig extends SimpleStreamOptions {
|
export interface AgentLoopConfig extends SimpleStreamOptions {
|
||||||
model: Model<any>;
|
model: Model<any>;
|
||||||
|
|
||||||
@@ -163,10 +175,22 @@ export interface AgentLoopConfig extends SimpleStreamOptions {
|
|||||||
*/
|
*/
|
||||||
getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
|
getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Called after each turn fully completes and `turn_end` has been emitted.
|
||||||
|
*
|
||||||
|
* If it returns true, the loop emits `agent_end` and exits before polling steering or follow-up queues,
|
||||||
|
* without starting another LLM call. The current assistant response and any tool executions finish normally.
|
||||||
|
*
|
||||||
|
* Use this to request a graceful stop after the current turn, e.g. before context gets too full.
|
||||||
|
*
|
||||||
|
* Contract: must not throw or reject. Throwing interrupts the low-level agent loop without producing a normal event sequence.
|
||||||
|
*/
|
||||||
|
shouldStopAfterTurn?: (context: ShouldStopAfterTurnContext) => boolean | Promise<boolean>;
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Returns steering messages to inject into the conversation mid-run.
|
* Returns steering messages to inject into the conversation mid-run.
|
||||||
*
|
*
|
||||||
* Called after the current assistant turn finishes executing its tool calls.
|
* Called after the current assistant turn finishes executing its tool calls, unless `shouldStopAfterTurn` exits first.
|
||||||
* If messages are returned, they are added to the context before the next LLM call.
|
* If messages are returned, they are added to the context before the next LLM call.
|
||||||
* Tool calls from the current assistant message are not skipped.
|
* Tool calls from the current assistant message are not skipped.
|
||||||
*
|
*
|
||||||
|
|||||||
@@ -894,6 +894,103 @@ describe("agentLoop with AgentMessage", () => {
|
|||||||
expect(parallelObserved).toBe(true);
|
expect(parallelObserved).toBe(true);
|
||||||
});
|
});
|
||||||
|
|
||||||
|
it("should stop after the current turn when shouldStopAfterTurn returns true", async () => {
|
||||||
|
const toolSchema = Type.Object({ value: Type.String() });
|
||||||
|
const executed: string[] = [];
|
||||||
|
const tool: AgentTool<typeof toolSchema, { value: string }> = {
|
||||||
|
name: "echo",
|
||||||
|
label: "Echo",
|
||||||
|
description: "Echo tool",
|
||||||
|
parameters: toolSchema,
|
||||||
|
async execute(_toolCallId, params) {
|
||||||
|
executed.push(params.value);
|
||||||
|
return {
|
||||||
|
content: [{ type: "text", text: `echoed: ${params.value}` }],
|
||||||
|
details: { value: params.value },
|
||||||
|
};
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
const context: AgentContext = {
|
||||||
|
systemPrompt: "",
|
||||||
|
messages: [],
|
||||||
|
tools: [tool],
|
||||||
|
};
|
||||||
|
|
||||||
|
let steeringPolls = 0;
|
||||||
|
let followUpPolls = 0;
|
||||||
|
let callbackToolResultIds: string[] = [];
|
||||||
|
let callbackContextRoles: string[] = [];
|
||||||
|
const config: AgentLoopConfig = {
|
||||||
|
model: createModel(),
|
||||||
|
convertToLlm: identityConverter,
|
||||||
|
getSteeringMessages: async () => {
|
||||||
|
steeringPolls++;
|
||||||
|
return [];
|
||||||
|
},
|
||||||
|
getFollowUpMessages: async () => {
|
||||||
|
followUpPolls++;
|
||||||
|
return [createUserMessage("follow up should stay queued")];
|
||||||
|
},
|
||||||
|
shouldStopAfterTurn: async ({ message, toolResults, context }) => {
|
||||||
|
expect(message.role).toBe("assistant");
|
||||||
|
callbackToolResultIds = toolResults.map((toolResult) => toolResult.toolCallId);
|
||||||
|
callbackContextRoles = context.messages.map((contextMessage) => contextMessage.role);
|
||||||
|
return true;
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
let llmCalls = 0;
|
||||||
|
const stream = agentLoop([createUserMessage("echo something")], context, config, undefined, () => {
|
||||||
|
llmCalls++;
|
||||||
|
const mockStream = new MockAssistantStream();
|
||||||
|
queueMicrotask(() => {
|
||||||
|
if (llmCalls === 1) {
|
||||||
|
const message = createAssistantMessage(
|
||||||
|
[{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }],
|
||||||
|
"toolUse",
|
||||||
|
);
|
||||||
|
mockStream.push({ type: "done", reason: "toolUse", message });
|
||||||
|
} else {
|
||||||
|
mockStream.push({
|
||||||
|
type: "done",
|
||||||
|
reason: "stop",
|
||||||
|
message: createAssistantMessage([{ type: "text", text: "should not run" }]),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
});
|
||||||
|
return mockStream;
|
||||||
|
});
|
||||||
|
|
||||||
|
const events: AgentEvent[] = [];
|
||||||
|
for await (const event of stream) {
|
||||||
|
events.push(event);
|
||||||
|
}
|
||||||
|
|
||||||
|
const messages = await stream.result();
|
||||||
|
expect(llmCalls).toBe(1);
|
||||||
|
expect(executed).toEqual(["hello"]);
|
||||||
|
expect(steeringPolls).toBe(1);
|
||||||
|
expect(followUpPolls).toBe(0);
|
||||||
|
expect(callbackToolResultIds).toEqual(["tool-1"]);
|
||||||
|
expect(callbackContextRoles).toEqual(["user", "assistant", "toolResult"]);
|
||||||
|
expect(messages.map((message) => message.role)).toEqual(["user", "assistant", "toolResult"]);
|
||||||
|
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",
|
||||||
|
"agent_end",
|
||||||
|
]);
|
||||||
|
});
|
||||||
|
|
||||||
it("should stop after a tool batch when every tool result sets terminate=true", async () => {
|
it("should stop after a tool batch when every tool result sets terminate=true", async () => {
|
||||||
const toolSchema = Type.Object({ value: Type.String() });
|
const toolSchema = Type.Object({ value: Type.String() });
|
||||||
const tool: AgentTool<typeof toolSchema, { value: string }> = {
|
const tool: AgentTool<typeof toolSchema, { value: string }> = {
|
||||||
|
|||||||
Reference in New Issue
Block a user