fix(coding-agent): sync tool hooks with agent event processing closes #2113

This commit is contained in:
Mario Zechner
2026-03-14 03:09:20 +01:00
parent f04d9bc428
commit 63ac2df24d
15 changed files with 961 additions and 317 deletions

View File

@@ -2,6 +2,14 @@
## [Unreleased]
### Added
- Added `beforeToolCall` and `afterToolCall` hooks to `AgentOptions` and `AgentLoopConfig` for preflight blocking and post-execution tool result mutation.
### Changed
- Added configurable tool execution mode to `Agent` and `agentLoop` via `toolExecution: "parallel" | "sequential"`, with `parallel` as the default. Parallel mode preflights tool calls sequentially, executes allowed tools concurrently, and emits final tool results in assistant source order.
## [0.57.1] - 2026-03-07
## [0.57.0] - 2026-03-07

View File

@@ -99,6 +99,15 @@ prompt("Read config.json")
└─ agent_end
```
Tool execution mode is configurable:
- `parallel` (default): preflight tool calls sequentially, execute allowed tools concurrently, emit final `tool_execution_end` and `toolResult` messages in assistant source order
- `sequential`: execute tool calls one by one, matching the historical behavior
The `beforeToolCall` hook runs after `tool_execution_start` and validated argument parsing. It can block execution. The `afterToolCall` hook runs after tool execution finishes and before `tool_execution_end` and final tool result message events are emitted.
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()` resumes from existing context without adding a new message. Use it for retries after errors.
@@ -159,6 +168,23 @@ const agent = new Agent({
// Dynamic API key resolution (for expiring OAuth tokens)
getApiKey: async (provider) => refreshToken(),
// Tool execution mode: "parallel" (default) or "sequential"
toolExecution: "parallel",
// Preflight each tool call after args are validated. Can block execution.
beforeToolCall: async ({ toolCall, args, context }) => {
if (toolCall.name === "bash") {
return { block: true, reason: "bash is disabled" };
}
},
// Postprocess each tool result before final tool events are emitted.
afterToolCall: async ({ toolCall, result, isError, context }) => {
if (!isError) {
return { details: { ...result.details, audited: true } };
}
},
// Custom thinking budgets for token-based providers
thinkingBudgets: {
minimal: 128,
@@ -214,6 +240,9 @@ agent.setSystemPrompt("New prompt");
agent.setModel(getModel("openai", "gpt-4o"));
agent.setThinkingLevel("medium");
agent.setTools([myTool]);
agent.setToolExecution("sequential");
agent.setBeforeToolCall(async ({ toolCall }) => undefined);
agent.setAfterToolCall(async ({ toolCall, result }) => undefined);
agent.replaceMessages(newMessages);
agent.appendMessage(message);
agent.clearMessages();
@@ -393,6 +422,9 @@ const context: AgentContext = {
const config: AgentLoopConfig = {
model: getModel("openai", "gpt-4o"),
convertToLlm: (msgs) => msgs.filter(m => ["user", "assistant", "toolResult"].includes(m.role)),
toolExecution: "parallel",
beforeToolCall: async ({ toolCall, args, context }) => undefined,
afterToolCall: async ({ toolCall, result, isError, context }) => undefined,
};
const userMessage = { role: "user", content: "Hello", timestamp: Date.now() };
@@ -407,6 +439,8 @@ for await (const event of agentLoopContinue(context, config)) {
}
```
These low-level streams are observational. They preserve event order, but they do not wait for your async event handling to settle before later producer phases continue. If you need message processing to act as a barrier before tool preflight, use the `Agent` class instead of raw `agentLoop()` or `agentLoopContinue()`.
## License
MIT

View File

@@ -17,10 +17,13 @@ import type {
AgentLoopConfig,
AgentMessage,
AgentTool,
AgentToolCall,
AgentToolResult,
StreamFn,
} from "./types.js";
export type AgentEventSink = (event: AgentEvent) => Promise<void> | void;
/**
* Start an agent loop with a new prompt message.
* The prompt is added to the context and events are emitted for it.
@@ -34,22 +37,18 @@ export function agentLoop(
): EventStream<AgentEvent, AgentMessage[]> {
const stream = createAgentStream();
(async () => {
const newMessages: AgentMessage[] = [...prompts];
const currentContext: AgentContext = {
...context,
messages: [...context.messages, ...prompts],
};
stream.push({ type: "agent_start" });
stream.push({ type: "turn_start" });
for (const prompt of prompts) {
stream.push({ type: "message_start", message: prompt });
stream.push({ type: "message_end", message: prompt });
}
await runLoop(currentContext, newMessages, config, signal, stream, streamFn);
})();
void runAgentLoop(
prompts,
context,
config,
async (event) => {
stream.push(event);
},
signal,
streamFn,
).then((messages) => {
stream.end(messages);
});
return stream;
}
@@ -78,19 +77,71 @@ export function agentLoopContinue(
const stream = createAgentStream();
(async () => {
const newMessages: AgentMessage[] = [];
const currentContext: AgentContext = { ...context };
stream.push({ type: "agent_start" });
stream.push({ type: "turn_start" });
await runLoop(currentContext, newMessages, config, signal, stream, streamFn);
})();
void runAgentLoopContinue(
context,
config,
async (event) => {
stream.push(event);
},
signal,
streamFn,
).then((messages) => {
stream.end(messages);
});
return stream;
}
export async function runAgentLoop(
prompts: AgentMessage[],
context: AgentContext,
config: AgentLoopConfig,
emit: AgentEventSink,
signal?: AbortSignal,
streamFn?: StreamFn,
): Promise<AgentMessage[]> {
const newMessages: AgentMessage[] = [...prompts];
const currentContext: AgentContext = {
...context,
messages: [...context.messages, ...prompts],
};
await emit({ type: "agent_start" });
await emit({ type: "turn_start" });
for (const prompt of prompts) {
await emit({ type: "message_start", message: prompt });
await emit({ type: "message_end", message: prompt });
}
await runLoop(currentContext, newMessages, config, signal, emit, streamFn);
return newMessages;
}
export async function runAgentLoopContinue(
context: AgentContext,
config: AgentLoopConfig,
emit: AgentEventSink,
signal?: AbortSignal,
streamFn?: StreamFn,
): Promise<AgentMessage[]> {
if (context.messages.length === 0) {
throw new Error("Cannot continue: no messages in context");
}
if (context.messages[context.messages.length - 1].role === "assistant") {
throw new Error("Cannot continue from message role: assistant");
}
const newMessages: AgentMessage[] = [];
const currentContext: AgentContext = { ...context };
await emit({ type: "agent_start" });
await emit({ type: "turn_start" });
await runLoop(currentContext, newMessages, config, signal, emit, streamFn);
return newMessages;
}
function createAgentStream(): EventStream<AgentEvent, AgentMessage[]> {
return new EventStream<AgentEvent, AgentMessage[]>(
(event: AgentEvent) => event.type === "agent_end",
@@ -106,7 +157,7 @@ async function runLoop(
newMessages: AgentMessage[],
config: AgentLoopConfig,
signal: AbortSignal | undefined,
stream: EventStream<AgentEvent, AgentMessage[]>,
emit: AgentEventSink,
streamFn?: StreamFn,
): Promise<void> {
let firstTurn = true;
@@ -121,7 +172,7 @@ async function runLoop(
// Inner loop: process tool calls and steering messages
while (hasMoreToolCalls || pendingMessages.length > 0) {
if (!firstTurn) {
stream.push({ type: "turn_start" });
await emit({ type: "turn_start" });
} else {
firstTurn = false;
}
@@ -129,8 +180,8 @@ async function runLoop(
// Process pending messages (inject before next assistant response)
if (pendingMessages.length > 0) {
for (const message of pendingMessages) {
stream.push({ type: "message_start", message });
stream.push({ type: "message_end", message });
await emit({ type: "message_start", message });
await emit({ type: "message_end", message });
currentContext.messages.push(message);
newMessages.push(message);
}
@@ -138,13 +189,12 @@ async function runLoop(
}
// Stream assistant response
const message = await streamAssistantResponse(currentContext, config, signal, stream, streamFn);
const message = await streamAssistantResponse(currentContext, config, signal, emit, streamFn);
newMessages.push(message);
if (message.stopReason === "error" || message.stopReason === "aborted") {
stream.push({ type: "turn_end", message, toolResults: [] });
stream.push({ type: "agent_end", messages: newMessages });
stream.end(newMessages);
await emit({ type: "turn_end", message, toolResults: [] });
await emit({ type: "agent_end", messages: newMessages });
return;
}
@@ -154,13 +204,7 @@ async function runLoop(
const toolResults: ToolResultMessage[] = [];
if (hasMoreToolCalls) {
const toolExecution = await executeToolCalls(
currentContext.tools,
message,
signal,
stream,
config.getSteeringMessages,
);
const toolExecution = await executeToolCalls(currentContext, message, config, signal, emit);
toolResults.push(...toolExecution.toolResults);
steeringAfterTools = toolExecution.steeringMessages ?? null;
@@ -170,7 +214,7 @@ async function runLoop(
}
}
stream.push({ type: "turn_end", message, toolResults });
await emit({ type: "turn_end", message, toolResults });
// Get steering messages after turn completes
if (steeringAfterTools && steeringAfterTools.length > 0) {
@@ -193,8 +237,7 @@ async function runLoop(
break;
}
stream.push({ type: "agent_end", messages: newMessages });
stream.end(newMessages);
await emit({ type: "agent_end", messages: newMessages });
}
/**
@@ -205,7 +248,7 @@ async function streamAssistantResponse(
context: AgentContext,
config: AgentLoopConfig,
signal: AbortSignal | undefined,
stream: EventStream<AgentEvent, AgentMessage[]>,
emit: AgentEventSink,
streamFn?: StreamFn,
): Promise<AssistantMessage> {
// Apply context transform if configured (AgentMessage[] → AgentMessage[])
@@ -245,7 +288,7 @@ async function streamAssistantResponse(
partialMessage = event.partial;
context.messages.push(partialMessage);
addedPartial = true;
stream.push({ type: "message_start", message: { ...partialMessage } });
await emit({ type: "message_start", message: { ...partialMessage } });
break;
case "text_start":
@@ -260,7 +303,7 @@ async function streamAssistantResponse(
if (partialMessage) {
partialMessage = event.partial;
context.messages[context.messages.length - 1] = partialMessage;
stream.push({
await emit({
type: "message_update",
assistantMessageEvent: event,
message: { ...partialMessage },
@@ -277,97 +320,87 @@ async function streamAssistantResponse(
context.messages.push(finalMessage);
}
if (!addedPartial) {
stream.push({ type: "message_start", message: { ...finalMessage } });
await emit({ type: "message_start", message: { ...finalMessage } });
}
stream.push({ type: "message_end", message: finalMessage });
await emit({ type: "message_end", message: finalMessage });
return finalMessage;
}
}
}
return await response.result();
const finalMessage = await response.result();
if (addedPartial) {
context.messages[context.messages.length - 1] = finalMessage;
} else {
context.messages.push(finalMessage);
await emit({ type: "message_start", message: { ...finalMessage } });
}
await emit({ type: "message_end", message: finalMessage });
return finalMessage;
}
/**
* Execute tool calls from an assistant message.
*/
async function executeToolCalls(
tools: AgentTool<any>[] | undefined,
currentContext: AgentContext,
assistantMessage: AssistantMessage,
config: AgentLoopConfig,
signal: AbortSignal | undefined,
stream: EventStream<AgentEvent, AgentMessage[]>,
getSteeringMessages?: AgentLoopConfig["getSteeringMessages"],
emit: AgentEventSink,
): Promise<{ toolResults: ToolResultMessage[]; steeringMessages?: AgentMessage[] }> {
const toolCalls = assistantMessage.content.filter((c) => c.type === "toolCall");
if (config.toolExecution === "sequential") {
return executeToolCallsSequential(currentContext, assistantMessage, toolCalls, config, signal, emit);
}
return executeToolCallsParallel(currentContext, assistantMessage, toolCalls, config, signal, emit);
}
async function executeToolCallsSequential(
currentContext: AgentContext,
assistantMessage: AssistantMessage,
toolCalls: AgentToolCall[],
config: AgentLoopConfig,
signal: AbortSignal | undefined,
emit: AgentEventSink,
): Promise<{ toolResults: ToolResultMessage[]; steeringMessages?: AgentMessage[] }> {
const results: ToolResultMessage[] = [];
let steeringMessages: AgentMessage[] | undefined;
for (let index = 0; index < toolCalls.length; index++) {
const toolCall = toolCalls[index];
const tool = tools?.find((t) => t.name === toolCall.name);
stream.push({
await emit({
type: "tool_execution_start",
toolCallId: toolCall.id,
toolName: toolCall.name,
args: toolCall.arguments,
});
let result: AgentToolResult<any>;
let isError = false;
try {
if (!tool) throw new Error(`Tool ${toolCall.name} not found`);
const validatedArgs = validateToolArguments(tool, toolCall);
result = await tool.execute(toolCall.id, validatedArgs, signal, (partialResult) => {
stream.push({
type: "tool_execution_update",
toolCallId: toolCall.id,
toolName: toolCall.name,
args: toolCall.arguments,
partialResult,
});
});
} catch (e) {
result = {
content: [{ type: "text", text: e instanceof Error ? e.message : String(e) }],
details: {},
};
isError = true;
const preparation = await prepareToolCall(currentContext, assistantMessage, toolCall, config, signal);
if (preparation.kind === "immediate") {
results.push(await emitToolCallOutcome(toolCall, preparation.result, preparation.isError, emit));
} else {
const executed = await executePreparedToolCall(preparation, signal, emit);
results.push(
await finalizeExecutedToolCall(
currentContext,
assistantMessage,
preparation,
executed,
config,
signal,
emit,
),
);
}
stream.push({
type: "tool_execution_end",
toolCallId: toolCall.id,
toolName: toolCall.name,
result,
isError,
});
const toolResultMessage: ToolResultMessage = {
role: "toolResult",
toolCallId: toolCall.id,
toolName: toolCall.name,
content: result.content,
details: result.details,
isError,
timestamp: Date.now(),
};
results.push(toolResultMessage);
stream.push({ type: "message_start", message: toolResultMessage });
stream.push({ type: "message_end", message: toolResultMessage });
// Check for steering messages - skip remaining tools if user interrupted
if (getSteeringMessages) {
const steering = await getSteeringMessages();
if (config.getSteeringMessages) {
const steering = await config.getSteeringMessages();
if (steering.length > 0) {
steeringMessages = steering;
const remainingCalls = toolCalls.slice(index + 1);
for (const skipped of remainingCalls) {
results.push(skipToolCall(skipped, stream));
results.push(await skipToolCall(skipped, emit));
}
break;
}
@@ -377,27 +410,241 @@ async function executeToolCalls(
return { toolResults: results, steeringMessages };
}
function skipToolCall(
toolCall: Extract<AssistantMessage["content"][number], { type: "toolCall" }>,
stream: EventStream<AgentEvent, AgentMessage[]>,
): ToolResultMessage {
const result: AgentToolResult<any> = {
content: [{ type: "text", text: "Skipped due to queued user message." }],
async function executeToolCallsParallel(
currentContext: AgentContext,
assistantMessage: AssistantMessage,
toolCalls: AgentToolCall[],
config: AgentLoopConfig,
signal: AbortSignal | undefined,
emit: AgentEventSink,
): Promise<{ toolResults: ToolResultMessage[]; steeringMessages?: AgentMessage[] }> {
const results: ToolResultMessage[] = [];
const runnableCalls: PreparedToolCall[] = [];
let steeringMessages: AgentMessage[] | undefined;
for (let index = 0; index < toolCalls.length; index++) {
const toolCall = toolCalls[index];
await emit({
type: "tool_execution_start",
toolCallId: toolCall.id,
toolName: toolCall.name,
args: toolCall.arguments,
});
const preparation = await prepareToolCall(currentContext, assistantMessage, toolCall, config, signal);
if (preparation.kind === "immediate") {
results.push(await emitToolCallOutcome(toolCall, preparation.result, preparation.isError, emit));
} else {
runnableCalls.push(preparation);
}
if (config.getSteeringMessages) {
const steering = await config.getSteeringMessages();
if (steering.length > 0) {
steeringMessages = steering;
for (const runnable of runnableCalls) {
results.push(await skipToolCall(runnable.toolCall, emit, { emitStart: false }));
}
const remainingCalls = toolCalls.slice(index + 1);
for (const skipped of remainingCalls) {
results.push(await skipToolCall(skipped, emit));
}
return { toolResults: results, steeringMessages };
}
}
}
const runningCalls = runnableCalls.map((prepared) => ({
prepared,
execution: executePreparedToolCall(prepared, signal, emit),
}));
for (const running of runningCalls) {
const executed = await running.execution;
results.push(
await finalizeExecutedToolCall(
currentContext,
assistantMessage,
running.prepared,
executed,
config,
signal,
emit,
),
);
}
if (!steeringMessages && config.getSteeringMessages) {
const steering = await config.getSteeringMessages();
if (steering.length > 0) {
steeringMessages = steering;
}
}
return { toolResults: results, steeringMessages };
}
type PreparedToolCall = {
kind: "prepared";
toolCall: AgentToolCall;
tool: AgentTool<any>;
args: unknown;
};
type ImmediateToolCallOutcome = {
kind: "immediate";
result: AgentToolResult<any>;
isError: boolean;
};
type ExecutedToolCallOutcome = {
result: AgentToolResult<any>;
isError: boolean;
};
async function prepareToolCall(
currentContext: AgentContext,
assistantMessage: AssistantMessage,
toolCall: AgentToolCall,
config: AgentLoopConfig,
signal: AbortSignal | undefined,
): Promise<PreparedToolCall | ImmediateToolCallOutcome> {
const tool = currentContext.tools?.find((t) => t.name === toolCall.name);
if (!tool) {
return {
kind: "immediate",
result: createErrorToolResult(`Tool ${toolCall.name} not found`),
isError: true,
};
}
try {
const validatedArgs = validateToolArguments(tool, toolCall);
if (config.beforeToolCall) {
const beforeResult = await config.beforeToolCall(
{
assistantMessage,
toolCall,
args: validatedArgs,
context: currentContext,
},
signal,
);
if (beforeResult?.block) {
return {
kind: "immediate",
result: createErrorToolResult(beforeResult.reason || "Tool execution was blocked"),
isError: true,
};
}
}
return {
kind: "prepared",
toolCall,
tool,
args: validatedArgs,
};
} catch (error) {
return {
kind: "immediate",
result: createErrorToolResult(error instanceof Error ? error.message : String(error)),
isError: true,
};
}
}
async function executePreparedToolCall(
prepared: PreparedToolCall,
signal: AbortSignal | undefined,
emit: AgentEventSink,
): Promise<ExecutedToolCallOutcome> {
const updateEvents: Promise<void>[] = [];
try {
const result = await prepared.tool.execute(
prepared.toolCall.id,
prepared.args as never,
signal,
(partialResult) => {
updateEvents.push(
Promise.resolve(
emit({
type: "tool_execution_update",
toolCallId: prepared.toolCall.id,
toolName: prepared.toolCall.name,
args: prepared.toolCall.arguments,
partialResult,
}),
),
);
},
);
await Promise.all(updateEvents);
return { result, isError: false };
} catch (error) {
await Promise.all(updateEvents);
return {
result: createErrorToolResult(error instanceof Error ? error.message : String(error)),
isError: true,
};
}
}
async function finalizeExecutedToolCall(
currentContext: AgentContext,
assistantMessage: AssistantMessage,
prepared: PreparedToolCall,
executed: ExecutedToolCallOutcome,
config: AgentLoopConfig,
signal: AbortSignal | undefined,
emit: AgentEventSink,
): Promise<ToolResultMessage> {
let result = executed.result;
let isError = executed.isError;
if (config.afterToolCall) {
const afterResult = await config.afterToolCall(
{
assistantMessage,
toolCall: prepared.toolCall,
args: prepared.args,
result,
isError,
context: currentContext,
},
signal,
);
if (afterResult) {
result = {
content: afterResult.content ?? result.content,
details: afterResult.details ?? result.details,
};
isError = afterResult.isError ?? isError;
}
}
return await emitToolCallOutcome(prepared.toolCall, result, isError, emit);
}
function createErrorToolResult(message: string): AgentToolResult<any> {
return {
content: [{ type: "text", text: message }],
details: {},
};
}
stream.push({
type: "tool_execution_start",
toolCallId: toolCall.id,
toolName: toolCall.name,
args: toolCall.arguments,
});
stream.push({
async function emitToolCallOutcome(
toolCall: AgentToolCall,
result: AgentToolResult<any>,
isError: boolean,
emit: AgentEventSink,
): Promise<ToolResultMessage> {
await emit({
type: "tool_execution_end",
toolCallId: toolCall.id,
toolName: toolCall.name,
result,
isError: true,
isError,
});
const toolResultMessage: ToolResultMessage = {
@@ -405,13 +652,31 @@ function skipToolCall(
toolCallId: toolCall.id,
toolName: toolCall.name,
content: result.content,
details: {},
isError: true,
details: result.details,
isError,
timestamp: Date.now(),
};
stream.push({ type: "message_start", message: toolResultMessage });
stream.push({ type: "message_end", message: toolResultMessage });
await emit({ type: "message_start", message: toolResultMessage });
await emit({ type: "message_end", message: toolResultMessage });
return toolResultMessage;
}
async function skipToolCall(
toolCall: AgentToolCall,
emit: AgentEventSink,
options?: { emitStart?: boolean },
): Promise<ToolResultMessage> {
const result = createErrorToolResult("Skipped due to queued user message.");
if (options?.emitStart !== false) {
await emit({
type: "tool_execution_start",
toolCallId: toolCall.id,
toolName: toolCall.name,
args: toolCall.arguments,
});
}
return await emitToolCallOutcome(toolCall, result, true, emit);
}

View File

@@ -14,16 +14,21 @@ import {
type ThinkingBudgets,
type Transport,
} from "@mariozechner/pi-ai";
import { agentLoop, agentLoopContinue } from "./agent-loop.js";
import { runAgentLoop, runAgentLoopContinue } from "./agent-loop.js";
import type {
AfterToolCallContext,
AfterToolCallResult,
AgentContext,
AgentEvent,
AgentLoopConfig,
AgentMessage,
AgentState,
AgentTool,
BeforeToolCallContext,
BeforeToolCallResult,
StreamFn,
ThinkingLevel,
ToolExecutionMode,
} from "./types.js";
/**
@@ -97,6 +102,15 @@ export interface AgentOptions {
* Default: 60000 (60 seconds). Set to 0 to disable the cap.
*/
maxRetryDelayMs?: number;
/** Tool execution mode. Default: "parallel" */
toolExecution?: ToolExecutionMode;
/** Called before a tool is executed, after arguments have been validated. */
beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise<BeforeToolCallResult | undefined>;
/** Called after a tool finishes executing, before final tool events are emitted. */
afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise<AfterToolCallResult | undefined>;
}
export class Agent {
@@ -129,6 +143,15 @@ export class Agent {
private _thinkingBudgets?: ThinkingBudgets;
private _transport: Transport;
private _maxRetryDelayMs?: number;
private _toolExecution: ToolExecutionMode;
private _beforeToolCall?: (
context: BeforeToolCallContext,
signal?: AbortSignal,
) => Promise<BeforeToolCallResult | undefined>;
private _afterToolCall?: (
context: AfterToolCallContext,
signal?: AbortSignal,
) => Promise<AfterToolCallResult | undefined>;
constructor(opts: AgentOptions = {}) {
this._state = { ...this._state, ...opts.initialState };
@@ -143,6 +166,9 @@ export class Agent {
this._thinkingBudgets = opts.thinkingBudgets;
this._transport = opts.transport ?? "sse";
this._maxRetryDelayMs = opts.maxRetryDelayMs;
this._toolExecution = opts.toolExecution ?? "parallel";
this._beforeToolCall = opts.beforeToolCall;
this._afterToolCall = opts.afterToolCall;
}
/**
@@ -203,6 +229,30 @@ export class Agent {
this._maxRetryDelayMs = value;
}
get toolExecution(): ToolExecutionMode {
return this._toolExecution;
}
setToolExecution(value: ToolExecutionMode) {
this._toolExecution = value;
}
setBeforeToolCall(
value:
| ((context: BeforeToolCallContext, signal?: AbortSignal) => Promise<BeforeToolCallResult | undefined>)
| undefined,
) {
this._beforeToolCall = value;
}
setAfterToolCall(
value:
| ((context: AfterToolCallContext, signal?: AbortSignal) => Promise<AfterToolCallResult | undefined>)
| undefined,
) {
this._afterToolCall = value;
}
get state(): AgentState {
return this._state;
}
@@ -405,6 +455,50 @@ export class Agent {
await this._runLoop(undefined);
}
private _processLoopEvent(event: AgentEvent): void {
switch (event.type) {
case "message_start":
this._state.streamMessage = event.message;
break;
case "message_update":
this._state.streamMessage = event.message;
break;
case "message_end":
this._state.streamMessage = null;
this.appendMessage(event.message);
break;
case "tool_execution_start": {
const pendingToolCalls = new Set(this._state.pendingToolCalls);
pendingToolCalls.add(event.toolCallId);
this._state.pendingToolCalls = pendingToolCalls;
break;
}
case "tool_execution_end": {
const pendingToolCalls = new Set(this._state.pendingToolCalls);
pendingToolCalls.delete(event.toolCallId);
this._state.pendingToolCalls = pendingToolCalls;
break;
}
case "turn_end":
if (event.message.role === "assistant" && (event.message as any).errorMessage) {
this._state.error = (event.message as any).errorMessage;
}
break;
case "agent_end":
this._state.isStreaming = false;
this._state.streamMessage = null;
break;
}
this.emit(event);
}
/**
* Run the agent loop.
* If messages are provided, starts a new conversation turn with those messages.
@@ -441,6 +535,9 @@ export class Agent {
transport: this._transport,
thinkingBudgets: this._thinkingBudgets,
maxRetryDelayMs: this._maxRetryDelayMs,
toolExecution: this._toolExecution,
beforeToolCall: this._beforeToolCall,
afterToolCall: this._afterToolCall,
convertToLlm: this.convertToLlm,
transformContext: this.transformContext,
getApiKey: this.getApiKey,
@@ -454,77 +551,24 @@ export class Agent {
getFollowUpMessages: async () => this.dequeueFollowUpMessages(),
};
let partial: AgentMessage | null = null;
try {
const stream = messages
? agentLoop(messages, context, config, this.abortController.signal, this.streamFn)
: agentLoopContinue(context, config, this.abortController.signal, this.streamFn);
for await (const event of stream) {
// Update internal state based on events
switch (event.type) {
case "message_start":
partial = event.message;
this._state.streamMessage = event.message;
break;
case "message_update":
partial = event.message;
this._state.streamMessage = event.message;
break;
case "message_end":
partial = null;
this._state.streamMessage = null;
this.appendMessage(event.message);
break;
case "tool_execution_start": {
const s = new Set(this._state.pendingToolCalls);
s.add(event.toolCallId);
this._state.pendingToolCalls = s;
break;
}
case "tool_execution_end": {
const s = new Set(this._state.pendingToolCalls);
s.delete(event.toolCallId);
this._state.pendingToolCalls = s;
break;
}
case "turn_end":
if (event.message.role === "assistant" && (event.message as any).errorMessage) {
this._state.error = (event.message as any).errorMessage;
}
break;
case "agent_end":
this._state.isStreaming = false;
this._state.streamMessage = null;
break;
}
// Emit to listeners
this.emit(event);
}
// Handle any remaining partial message
if (partial && partial.role === "assistant" && partial.content.length > 0) {
const onlyEmpty = !partial.content.some(
(c) =>
(c.type === "thinking" && c.thinking.trim().length > 0) ||
(c.type === "text" && c.text.trim().length > 0) ||
(c.type === "toolCall" && c.name.trim().length > 0),
if (messages) {
await runAgentLoop(
messages,
context,
config,
async (event) => this._processLoopEvent(event),
this.abortController.signal,
this.streamFn,
);
} else {
await runAgentLoopContinue(
context,
config,
async (event) => this._processLoopEvent(event),
this.abortController.signal,
this.streamFn,
);
if (!onlyEmpty) {
this.appendMessage(partial);
} else {
if (this.abortController?.signal.aborted) {
throw new Error("Request was aborted");
}
}
}
} catch (err: any) {
const errorMsg: AgentMessage = {

View File

@@ -1,4 +1,5 @@
import type {
AssistantMessage,
AssistantMessageEvent,
ImageContent,
Message,
@@ -25,8 +26,73 @@ export type StreamFn = (
) => ReturnType<typeof streamSimple> | Promise<ReturnType<typeof streamSimple>>;
/**
* Configuration for the agent loop.
* Configuration for how tool calls from a single assistant message are executed.
*
* - "sequential": each tool call is prepared, executed, and finalized before the next one starts.
* - "parallel": tool calls are prepared sequentially, then allowed tools execute concurrently.
* Final tool results are still emitted in assistant source order.
*/
export type ToolExecutionMode = "sequential" | "parallel";
/** A single tool call content block emitted by an assistant message. */
export type AgentToolCall = Extract<AssistantMessage["content"][number], { type: "toolCall" }>;
/**
* Result returned from `beforeToolCall`.
*
* Returning `{ block: true }` prevents the tool from executing. The loop emits an error tool result instead.
* `reason` becomes the text shown in that error result. If omitted, a default blocked message is used.
*/
export interface BeforeToolCallResult {
block?: boolean;
reason?: string;
}
/**
* Partial override returned from `afterToolCall`.
*
* Merge semantics are field-by-field:
* - `content`: if provided, replaces the tool result content array in full
* - `details`: if provided, replaces the tool result details value in full
* - `isError`: if provided, replaces the tool result error flag
*
* Omitted fields keep the original executed tool result values.
* There is no deep merge for `content` or `details`.
*/
export interface AfterToolCallResult {
content?: (TextContent | ImageContent)[];
details?: unknown;
isError?: boolean;
}
/** Context passed to `beforeToolCall`. */
export interface BeforeToolCallContext {
/** The assistant message that requested the tool call. */
assistantMessage: AssistantMessage;
/** The raw tool call block from `assistantMessage.content`. */
toolCall: AgentToolCall;
/** Validated tool arguments for the target tool schema. */
args: unknown;
/** Current agent context at the time the tool call is prepared. */
context: AgentContext;
}
/** Context passed to `afterToolCall`. */
export interface AfterToolCallContext {
/** The assistant message that requested the tool call. */
assistantMessage: AssistantMessage;
/** The raw tool call block from `assistantMessage.content`. */
toolCall: AgentToolCall;
/** Validated tool arguments for the target tool schema. */
args: unknown;
/** The executed tool result before any `afterToolCall` overrides are applied. */
result: AgentToolResult<any>;
/** Whether the executed tool result is currently treated as an error. */
isError: boolean;
/** Current agent context at the time the tool call is finalized. */
context: AgentContext;
}
export interface AgentLoopConfig extends SimpleStreamOptions {
model: Model<any>;
@@ -115,6 +181,36 @@ export interface AgentLoopConfig extends SimpleStreamOptions {
* Contract: must not throw or reject. Return [] when no follow-up messages are available.
*/
getFollowUpMessages?: () => Promise<AgentMessage[]>;
/**
* Tool execution mode.
* - "sequential": execute tool calls one by one
* - "parallel": preflight tool calls sequentially, then execute allowed tools concurrently
*
* Default: "parallel"
*/
toolExecution?: ToolExecutionMode;
/**
* Called before a tool is executed, after arguments have been validated.
*
* Return `{ block: true }` to prevent execution. The loop emits an error tool result instead.
* The hook receives the agent abort signal and is responsible for honoring it.
*/
beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise<BeforeToolCallResult | undefined>;
/**
* Called after a tool finishes executing, before final tool events are emitted.
*
* Return an `AfterToolCallResult` to override parts of the executed tool result:
* - `content` replaces the full content array
* - `details` replaces the full details payload
* - `isError` replaces the error flag
*
* Any omitted fields keep their original values. No deep merge is performed.
* The hook receives the agent abort signal and is responsible for honoring it.
*/
afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise<AfterToolCallResult | undefined>;
}
/**

View File

@@ -307,6 +307,87 @@ describe("agentLoop with AgentMessage", () => {
}
});
it("should execute tool calls in parallel and emit tool results in source order", async () => {
const toolSchema = Type.Object({ value: Type.String() });
let firstResolved = false;
let parallelObserved = false;
let releaseFirst: (() => void) | undefined;
const firstDone = new Promise<void>((resolve) => {
releaseFirst = resolve;
});
const tool: AgentTool<typeof toolSchema, { value: string }> = {
name: "echo",
label: "Echo",
description: "Echo tool",
parameters: toolSchema,
async execute(_toolCallId, params) {
if (params.value === "first") {
await firstDone;
firstResolved = true;
}
if (params.value === "second" && !firstResolved) {
parallelObserved = true;
}
return {
content: [{ type: "text", text: `echoed: ${params.value}` }],
details: { value: params.value },
};
},
};
const context: AgentContext = {
systemPrompt: "",
messages: [],
tools: [tool],
};
const userPrompt: AgentMessage = createUserMessage("echo both");
const config: AgentLoopConfig = {
model: createModel(),
convertToLlm: identityConverter,
toolExecution: "parallel",
};
let callIndex = 0;
const stream = agentLoop([userPrompt], context, config, undefined, () => {
const mockStream = new MockAssistantStream();
queueMicrotask(() => {
if (callIndex === 0) {
const message = createAssistantMessage(
[
{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "first" } },
{ type: "toolCall", id: "tool-2", name: "echo", arguments: { value: "second" } },
],
"toolUse",
);
mockStream.push({ type: "done", reason: "toolUse", message });
setTimeout(() => releaseFirst?.(), 20);
} else {
const message = createAssistantMessage([{ type: "text", text: "done" }]);
mockStream.push({ type: "done", reason: "stop", message });
}
callIndex++;
});
return mockStream;
});
const events: AgentEvent[] = [];
for await (const event of stream) {
events.push(event);
}
const toolResultIds = events.flatMap((event) => {
if (event.type !== "message_end" || event.message.role !== "toolResult") {
return [];
}
return [event.message.toolCallId];
});
expect(parallelObserved).toBe(true);
expect(toolResultIds).toEqual(["tool-1", "tool-2"]);
});
it("should inject queued messages and skip remaining tool calls", async () => {
const toolSchema = Type.Object({ value: Type.String() });
const executed: string[] = [];
@@ -340,6 +421,7 @@ describe("agentLoop with AgentMessage", () => {
const config: AgentLoopConfig = {
model: createModel(),
convertToLlm: identityConverter,
toolExecution: "sequential",
getSteeringMessages: async () => {
// Return steering message after first tool executes
if (executed.length === 1 && !queuedDelivered) {

View File

@@ -5,6 +5,8 @@ import { Agent } from "../src/index.js";
import { hasBedrockCredentials } from "./bedrock-utils.js";
import { calculateTool } from "./utils/calculate.js";
delete process.env.ANTHROPIC_OAUTH_TOKEN;
async function basicPrompt(model: Model<any>) {
const agent = new Agent({
initialState: {
@@ -278,7 +280,7 @@ describe("Agent E2E Tests", () => {
});
});
describe.skipIf(!process.env.CEREBRAS_API_KEY)("Cerebras Provider (gpt-oss-120b)", () => {
/*describe.skipIf(!process.env.CEREBRAS_API_KEY)("Cerebras Provider (gpt-oss-120b)", () => {
const model = getModel("cerebras", "gpt-oss-120b");
it("should handle basic text prompt", async () => {
@@ -300,7 +302,7 @@ describe("Agent E2E Tests", () => {
it("should maintain context across multiple turns", async () => {
await multiTurnConversation(model);
});
});
});*/
describe.skipIf(!process.env.ZAI_API_KEY)("zAI Provider (glm-4.5-air)", () => {
const model = getModel("zai", "glm-4.5-air");
@@ -357,7 +359,7 @@ describe("Agent.continue()", () => {
const agent = new Agent({
initialState: {
systemPrompt: "Test",
model: getModel("anthropic", "claude-haiku-4-5"),
model: getModel("openai", "gpt-5.4"),
},
});
@@ -368,16 +370,16 @@ describe("Agent.continue()", () => {
const agent = new Agent({
initialState: {
systemPrompt: "Test",
model: getModel("anthropic", "claude-haiku-4-5"),
model: getModel("openai", "gpt-5.4"),
},
});
const assistantMessage: AssistantMessage = {
role: "assistant",
content: [{ type: "text", text: "Hello" }],
api: "anthropic-messages",
provider: "anthropic",
model: "claude-haiku-4-5",
api: "openai-responses",
provider: "openai",
model: "gpt-5.4",
usage: {
input: 0,
output: 0,
@@ -395,8 +397,8 @@ describe("Agent.continue()", () => {
});
});
describe.skipIf(!process.env.ANTHROPIC_API_KEY)("continue from user message", () => {
const model = getModel("anthropic", "claude-haiku-4-5");
describe.skipIf(!process.env.OPENAI_API_KEY)("continue from user message", () => {
const model = getModel("openai", "gpt-5.4");
it("should continue and get response when last message is user", async () => {
const agent = new Agent({
@@ -433,8 +435,8 @@ describe("Agent.continue()", () => {
});
});
describe.skipIf(!process.env.ANTHROPIC_API_KEY)("continue from tool result", () => {
const model = getModel("anthropic", "claude-haiku-4-5");
describe.skipIf(!process.env.OPENAI_API_KEY)("continue from tool result", () => {
const model = getModel("openai", "gpt-5.4");
it("should continue and process tool results", async () => {
const agent = new Agent({