shumengya/sproutclaw
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

8
packages/agent/CHANGELOG.md
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

34
packages/agent/README.md
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

511
packages/agent/src/agent-loop.ts
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);
}

184
packages/agent/src/agent.ts
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 = {

98
packages/agent/src/types.ts
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>;
}
/**

82
packages/agent/test/agent-loop.test.ts
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) {

24
packages/agent/test/e2e.test.ts
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({

5
packages/coding-agent/CHANGELOG.md
View File

@@ -2,8 +2,13 @@
## [Unreleased]
### Changed
- Changed extension tool interception to use agent-core `beforeToolCall` and `afterToolCall` hooks instead of wrapper-based interception. Tool calls now execute in parallel by default, extension `tool_call` preflight still runs sequentially, and final tool results are emitted in assistant source order.
### Fixed
- Fixed `tool_call` extension handlers observing stale `sessionManager` state during multi-tool turns by draining queued agent events before each `tool_call` preflight. In parallel tool mode this guarantees state through the current assistant tool-calling message, but not sibling tool results from the same assistant message.
- Fixed interactive input fields backed by the TUI `Input` component to scroll by visual column width for wide Unicode text (CJK, fullwidth characters), preventing rendered line overflow and TUI crashes in places like search and filter inputs ([#1982](https://github.com/badlogic/pi-mono/issues/1982))
- Fixed `shift+tab` and other modified Tab bindings in tmux when `extended-keys-format` is left at the default `xterm`
- Fixed EXIF orientation not being applied during image convert and resize, causing JPEG and WebP images from phone cameras to display rotated or mirrored ([#2105](https://github.com/badlogic/pi-mono/pull/2105) by [@melihmucuk](https://github.com/melihmucuk))

21
packages/coding-agent/docs/extensions.md
View File

@@ -247,11 +247,11 @@ user sends prompt ────────────────────
│ ├─► before_provider_request (can inspect or replace payload)
│ │ │ │
│ │ LLM responds, may call tools: │ │
│ │ ├─► tool_call (can block) │ │
│ │ ├─► tool_execution_start │ │
│ │ ├─► tool_call (can block) │ │
│ │ ├─► tool_execution_update │ │
│ │ ├─► tool_execution_end │ │
│ │ └─► tool_result (can modify) │ │
│ │ ├─► tool_result (can modify) │ │
│ │ └─► tool_execution_end │ │
│ │ │ │
│ └─► turn_end │ │
│ │
@@ -485,6 +485,11 @@ pi.on("message_end", async (event, ctx) => {
Fired for tool execution lifecycle updates.
In parallel tool mode:
- `tool_execution_start` is emitted in assistant source order during the preflight phase
- `tool_execution_update` events may interleave across tools
- `tool_execution_end` is emitted in assistant source order, matching final tool result message order
```typescript
pi.on("tool_execution_start", async (event, ctx) => {
// event.toolCallId, event.toolName, event.args
@@ -553,7 +558,11 @@ Use this to update UI elements (status bars, footers) or perform model-specific
#### tool_call
Fired before tool executes. **Can block.** Use `isToolCallEventType` to narrow and get typed inputs.
Fired after `tool_execution_start`, before the tool executes. **Can block.** Use `isToolCallEventType` to narrow and get typed inputs.
Before `tool_call` runs, pi waits for previously emitted Agent events to finish draining through `AgentSession`. This means `ctx.sessionManager` is up to date through the current assistant tool-calling message.
In the default parallel tool execution mode, sibling tool calls from the same assistant message are preflighted sequentially, then executed concurrently. `tool_call` is not guaranteed to see sibling tool results from that same assistant message in `ctx.sessionManager`.
```typescript
import { isToolCallEventType } from "@mariozechner/pi-coding-agent";
@@ -602,7 +611,7 @@ pi.on("tool_call", (event) => {
#### tool_result
Fired after tool executes. **Can modify result.**
Fired after tool execution finishes and before `tool_execution_end` plus the final tool result message events are emitted. **Can modify result.**
`tool_result` handlers chain like middleware:
- Handlers run in extension load order
@@ -715,6 +724,8 @@ Current working directory.
Read-only access to session state. See [session.md](session.md) for the full SessionManager API and entry types.
For `tool_call`, this state is synchronized through the current assistant message before handlers run. In parallel tool execution mode it is still not guaranteed to include sibling tool results from the same assistant message.
```typescript
ctx.sessionManager.getEntries() // All entries
ctx.sessionManager.getBranch() // Current branch

69
packages/coding-agent/src/core/agent-session.ts
View File

@@ -67,7 +67,6 @@ import {
type TurnEndEvent,
type TurnStartEvent,
wrapRegisteredTools,
wrapToolsWithExtensions,
} from "./extensions/index.js";
import type { BashExecutionMessage, CustomMessage } from "./messages.js";
import type { ModelRegistry } from "./model-registry.js";
@@ -291,6 +290,7 @@ export class AgentSession {
// Always subscribe to agent events for internal handling
// (session persistence, extensions, auto-compaction, retry logic)
this._unsubscribeAgent = this.agent.subscribe(this._handleAgentEvent);
this._installAgentToolHooks();
this._buildRuntime({
activeToolNames: this._initialActiveToolNames,
@@ -303,6 +303,65 @@ export class AgentSession {
return this._modelRegistry;
}
/**
* Install tool hooks once on the Agent instance.
*
* The callbacks read `this._extensionRunner` at execution time, so extension reload swaps in the
* new runner without reinstalling hooks. Extension-specific tool wrappers are still used to adapt
* registered tool execution to the extension context. Tool call and tool result interception now
* happens here instead of in wrappers.
*/
private _installAgentToolHooks(): void {
this.agent.setBeforeToolCall(async ({ toolCall, args }) => {
const runner = this._extensionRunner;
if (!runner?.hasHandlers("tool_call")) {
return undefined;
}
await this._agentEventQueue;
try {
return await runner.emitToolCall({
type: "tool_call",
toolName: toolCall.name,
toolCallId: toolCall.id,
input: args as Record<string, unknown>,
});
} catch (err) {
if (err instanceof Error) {
throw err;
}
throw new Error(`Extension failed, blocking execution: ${String(err)}`);
}
});
this.agent.setAfterToolCall(async ({ toolCall, args, result, isError }) => {
const runner = this._extensionRunner;
if (!runner?.hasHandlers("tool_result")) {
return undefined;
}
const hookResult = await runner.emitToolResult({
type: "tool_result",
toolName: toolCall.name,
toolCallId: toolCall.id,
input: args as Record<string, unknown>,
content: result.content,
details: isError ? undefined : result.details,
isError,
});
if (!hookResult || isError) {
return undefined;
}
return {
content: hookResult.content,
details: hookResult.details,
};
});
}
// =========================================================================
// Event Subscription
// =========================================================================
@@ -2144,13 +2203,7 @@ export class AgentSession {
for (const tool of wrappedExtensionTools as AgentTool[]) {
toolRegistry.set(tool.name, tool);
}
if (this._extensionRunner) {
const wrappedAllTools = wrapToolsWithExtensions(Array.from(toolRegistry.values()), this._extensionRunner);
this._toolRegistry = new Map(wrappedAllTools.map((tool) => [tool.name, tool]));
} else {
this._toolRegistry = toolRegistry;
}
this._toolRegistry = toolRegistry;
const nextActiveToolNames = options?.activeToolNames
? [...options.activeToolNames]

7
packages/coding-agent/src/core/extensions/index.ts
View File

@@ -163,9 +163,4 @@ export {
isToolCallEventType,
isWriteToolResult,
} from "./types.js";
export {
wrapRegisteredTool,
wrapRegisteredTools,
wrapToolsWithExtensions,
wrapToolWithExtensions,
} from "./wrapper.js";
export { wrapRegisteredTool, wrapRegisteredTools } from "./wrapper.js";

96
packages/coding-agent/src/core/extensions/wrapper.ts
View File

@@ -1,10 +1,13 @@
/**
* Tool wrappers for extensions.
* Tool wrappers for extension-registered tools.
*
* These wrappers only adapt tool execution so extension tools receive the runner context.
* Tool call and tool result interception is handled by AgentSession via agent-core hooks.
*/
import type { AgentTool, AgentToolUpdateCallback } from "@mariozechner/pi-agent-core";
import type { AgentTool } from "@mariozechner/pi-agent-core";
import type { ExtensionRunner } from "./runner.js";
import type { RegisteredTool, ToolCallEventResult } from "./types.js";
import type { RegisteredTool } from "./types.js";
/**
* Wrap a RegisteredTool into an AgentTool.
@@ -29,90 +32,3 @@ export function wrapRegisteredTool(registeredTool: RegisteredTool, runner: Exten
export function wrapRegisteredTools(registeredTools: RegisteredTool[], runner: ExtensionRunner): AgentTool[] {
return registeredTools.map((rt) => wrapRegisteredTool(rt, runner));
}
/**
* Wrap a tool with extension callbacks for interception.
* - Emits tool_call event before execution (can block)
* - Emits tool_result event after execution (can modify result)
*/
export function wrapToolWithExtensions<T>(tool: AgentTool<any, T>, runner: ExtensionRunner): AgentTool<any, T> {
return {
...tool,
execute: async (
toolCallId: string,
params: Record<string, unknown>,
signal?: AbortSignal,
onUpdate?: AgentToolUpdateCallback<T>,
) => {
// Emit tool_call event - extensions can block execution
if (runner.hasHandlers("tool_call")) {
try {
const callResult = (await runner.emitToolCall({
type: "tool_call",
toolName: tool.name,
toolCallId,
input: params,
})) as ToolCallEventResult | undefined;
if (callResult?.block) {
const reason = callResult.reason || "Tool execution was blocked by an extension";
throw new Error(reason);
}
} catch (err) {
if (err instanceof Error) {
throw err;
}
throw new Error(`Extension failed, blocking execution: ${String(err)}`);
}
}
// Execute the actual tool
try {
const result = await tool.execute(toolCallId, params, signal, onUpdate);
// Emit tool_result event - extensions can modify the result
if (runner.hasHandlers("tool_result")) {
const resultResult = await runner.emitToolResult({
type: "tool_result",
toolName: tool.name,
toolCallId,
input: params,
content: result.content,
details: result.details,
isError: false,
});
if (resultResult) {
return {
content: resultResult.content ?? result.content,
details: (resultResult.details ?? result.details) as T,
};
}
}
return result;
} catch (err) {
// Emit tool_result event for errors
if (runner.hasHandlers("tool_result")) {
await runner.emitToolResult({
type: "tool_result",
toolName: tool.name,
toolCallId,
input: params,
content: [{ type: "text", text: err instanceof Error ? err.message : String(err) }],
details: undefined,
isError: true,
});
}
throw err;
}
},
};
}
/**
* Wrap all tools with extension callbacks.
*/
export function wrapToolsWithExtensions<T>(tools: AgentTool<any, T>[], runner: ExtensionRunner): AgentTool<any, T>[] {
return tools.map((tool) => wrapToolWithExtensions(tool, runner));
}

1
packages/coding-agent/src/core/index.ts
View File

@@ -57,5 +57,4 @@ export {
type ToolResultEvent,
type TurnEndEvent,
type TurnStartEvent,
wrapToolsWithExtensions,
} from "./extensions/index.js";

2
packages/coding-agent/src/index.ts
View File

@@ -137,8 +137,6 @@ export {
isWriteToolResult,
wrapRegisteredTool,
wrapRegisteredTools,
wrapToolsWithExtensions,
wrapToolWithExtensions,
} from "./core/extensions/index.js";
// Footer data provider (git branch + extension statuses - data not otherwise available to extensions)
export type { ReadonlyFooterDataProvider } from "./core/footer-data-provider.js";

136
packages/coding-agent/test/agent-session-concurrent.test.ts
View File

@@ -216,6 +216,142 @@ describe("AgentSession concurrent prompt guard", () => {
await expect(session.prompt("Second message")).resolves.not.toThrow();
});
it("should wait for queued agent events before emitting tool_call", async () => {
const model = getModel("anthropic", "claude-sonnet-4-5")!;
const tool = {
name: "dummy",
description: "Dummy tool",
label: "dummy",
parameters: Type.Object({ q: Type.String() }),
execute: async (_toolCallId: string, params: unknown) => {
const q =
typeof params === "object" && params !== null && "q" in params
? String((params as { q: unknown }).q)
: "";
return {
content: [{ type: "text" as const, text: `result:${q}` }],
details: {},
};
},
};
const agent = new Agent({
getApiKey: () => "test-key",
initialState: {
model,
systemPrompt: "Test",
tools: [tool],
},
streamFn: async (_model, context) => {
const stream = new MockAssistantStream();
queueMicrotask(() => {
const toolResultCount = context.messages.filter((message) => message.role === "toolResult").length;
if (toolResultCount > 0) {
const message: AssistantMessage = {
role: "assistant",
content: [{ type: "text", text: "done" }],
api: "anthropic-messages",
provider: "anthropic",
model: "mock",
usage: {
input: 1,
output: 1,
cacheRead: 0,
cacheWrite: 0,
totalTokens: 2,
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
},
stopReason: "stop",
timestamp: Date.now(),
};
stream.push({ type: "start", partial: { ...message, content: [] } });
stream.push({ type: "done", reason: "stop", message });
return;
}
const message: AssistantMessage = {
role: "assistant",
content: [
{ type: "toolCall", id: "toolu_1", name: "dummy", arguments: { q: "x" } },
{ type: "toolCall", id: "toolu_2", name: "dummy", arguments: { q: "y" } },
],
api: "anthropic-messages",
provider: "anthropic",
model: "mock",
usage: {
input: 1,
output: 1,
cacheRead: 0,
cacheWrite: 0,
totalTokens: 2,
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
},
stopReason: "toolUse",
timestamp: Date.now(),
};
stream.push({ type: "start", partial: { ...message, content: [] } });
stream.push({ type: "done", reason: "toolUse", message });
});
return stream;
},
});
const sessionManager = SessionManager.inMemory();
const settingsManager = SettingsManager.create(tempDir, tempDir);
const authStorage = AuthStorage.create(join(tempDir, "auth.json"));
const modelRegistry = new ModelRegistry(authStorage, tempDir);
authStorage.setRuntimeApiKey("anthropic", "test-key");
session = new AgentSession({
agent,
sessionManager,
settingsManager,
cwd: tempDir,
modelRegistry,
resourceLoader: createTestResourceLoader(),
baseToolsOverride: { dummy: tool },
});
const snapshots: string[][] = [];
const sessionWithRunner = session as unknown as {
_extensionRunner?: {
hasHandlers: (eventType: string) => boolean;
emit: (event: { type: string; message?: { role?: string } }) => Promise<void>;
emitToolCall: (event: { type: string; toolCallId: string }) => Promise<undefined>;
emitInput: (
text: string,
images: unknown,
source: "interactive" | "rpc" | "extension",
) => Promise<{ action: "continue" }>;
emitBeforeAgentStart: (prompt: string, images: unknown, systemPrompt: string) => Promise<undefined>;
};
};
sessionWithRunner._extensionRunner = {
hasHandlers: (eventType) => eventType === "tool_call",
emit: async () => {},
emitToolCall: async () => {
snapshots.push(
sessionManager
.getEntries()
.filter((entry) => entry.type === "message")
.map((entry) => entry.message.role),
);
return undefined;
},
emitInput: async () => ({ action: "continue" }),
emitBeforeAgentStart: async () => undefined,
};
await session.prompt("hi");
await session.agent.waitForIdle();
expect(snapshots).toEqual([
["user", "assistant"],
["user", "assistant"],
]);
});
it("should persist message_end events in order with slow extension handlers", async () => {
const model = getModel("anthropic", "claude-sonnet-4-5")!;
const tool = {