diff --git a/src/agents/pi-embedded-subscribe.handlers.tools.test.ts b/src/agents/pi-embedded-subscribe.handlers.tools.test.ts index c03eb00da57..96a988e5bc6 100644 --- a/src/agents/pi-embedded-subscribe.handlers.tools.test.ts +++ b/src/agents/pi-embedded-subscribe.handlers.tools.test.ts @@ -88,6 +88,37 @@ describe("handleToolExecutionStart read path checks", () => { expect(warn).toHaveBeenCalledTimes(1); expect(String(warn.mock.calls[0]?.[0] ?? "")).toContain("read tool called without path"); }); + + it("awaits onBlockReplyFlush before continuing tool start processing", async () => { + const { ctx, onBlockReplyFlush } = createTestContext(); + let releaseFlush: (() => void) | undefined; + onBlockReplyFlush.mockImplementation( + () => + new Promise((resolve) => { + releaseFlush = resolve; + }), + ); + + const evt: ToolExecutionStartEvent = { + type: "tool_execution_start", + toolName: "exec", + toolCallId: "tool-await-flush", + args: { command: "echo hi" }, + }; + + const pending = handleToolExecutionStart(ctx, evt); + // Let the async function reach the awaited flush Promise. + await Promise.resolve(); + + // If flush isn't awaited, tool metadata would already be recorded here. + expect(ctx.state.toolMetaById.has("tool-await-flush")).toBe(false); + expect(releaseFlush).toBeTypeOf("function"); + + releaseFlush?.(); + await pending; + + expect(ctx.state.toolMetaById.has("tool-await-flush")).toBe(true); + }); }); describe("handleToolExecutionEnd cron.add commitment tracking", () => { diff --git a/src/agents/pi-embedded-subscribe.handlers.tools.ts b/src/agents/pi-embedded-subscribe.handlers.tools.ts index ea3031a6cc4..18dc11193f0 100644 --- a/src/agents/pi-embedded-subscribe.handlers.tools.ts +++ b/src/agents/pi-embedded-subscribe.handlers.tools.ts @@ -174,7 +174,7 @@ export async function handleToolExecutionStart( // Flush pending block replies to preserve message boundaries before tool execution. ctx.flushBlockReplyBuffer(); if (ctx.params.onBlockReplyFlush) { - void ctx.params.onBlockReplyFlush(); + await ctx.params.onBlockReplyFlush(); } const rawToolName = String(evt.toolName);