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

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 = {