fix: stream tool summaries early and tool output

This commit is contained in:
Peter Steinberger
2026-01-03 20:37:57 +01:00
parent 03c1599544
commit a15cffb7de
15 changed files with 379 additions and 122 deletions

View File

@@ -48,6 +48,7 @@ import {
} from "./pi-embedded-subscribe.js";
import { extractAssistantText } from "./pi-embedded-utils.js";
import { createClawdisCodingTools } from "./pi-tools.js";
import { resolveSandboxContext } from "./sandbox.js";
import {
applySkillEnvOverrides,
applySkillEnvOverridesFromSnapshot,
@@ -362,7 +363,13 @@ export async function runEmbeddedPiAgent(params: {
return enqueueCommandInLane(sessionLane, () =>
enqueueGlobal(async () => {
const started = Date.now();
const resolvedWorkspace = resolveUserPath(params.workspaceDir);
const sandbox = await resolveSandboxContext({
config: params.config,
sessionKey: params.sessionKey,
workspaceDir: params.workspaceDir,
});
const workspaceDir = sandbox?.workspaceDir ?? params.workspaceDir;
const resolvedWorkspace = resolveUserPath(workspaceDir);
const prevCwd = process.cwd();
const provider =
@@ -425,6 +432,7 @@ export async function runEmbeddedPiAgent(params: {
const tools = createClawdisCodingTools({
bash: params.config?.agent?.bash,
surface: params.surface,
sandbox,
});
const machineName = await getMachineDisplayName();
const runtimeInfo = {
@@ -497,7 +505,6 @@ export async function runEmbeddedPiAgent(params: {
assistantTexts,
toolMetas,
unsubscribe,
flush: flushToolDebouncer,
waitForCompactionRetry,
} = subscribeEmbeddedPiSession({
session,
@@ -571,7 +578,6 @@ export async function runEmbeddedPiAgent(params: {
abortWarnTimer = undefined;
}
unsubscribe();
flushToolDebouncer();
if (ACTIVE_EMBEDDED_RUNS.get(params.sessionId) === queueHandle) {
ACTIVE_EMBEDDED_RUNS.delete(params.sessionId);
}

View File

@@ -630,4 +630,107 @@ describe("subscribeEmbeddedPiSession", () => {
await waitPromise;
expect(resolved).toBe(true);
});
it("emits tool summaries at tool start when verbose is on", () => {
let handler: ((evt: unknown) => void) | undefined;
const session: StubSession = {
subscribe: (fn) => {
handler = fn;
return () => {};
},
};
const onToolResult = vi.fn();
subscribeEmbeddedPiSession({
session: session as unknown as Parameters<
typeof subscribeEmbeddedPiSession
>[0]["session"],
runId: "run-tool",
verboseLevel: "on",
onToolResult,
});
handler?.({
type: "tool_execution_start",
toolName: "read",
toolCallId: "tool-1",
args: { path: "/tmp/a.txt" },
});
expect(onToolResult).toHaveBeenCalledTimes(1);
const payload = onToolResult.mock.calls[0][0];
expect(payload.text).toContain("/tmp/a.txt");
handler?.({
type: "tool_execution_end",
toolName: "read",
toolCallId: "tool-1",
isError: false,
result: "ok",
});
expect(onToolResult).toHaveBeenCalledTimes(1);
});
it("skips tool summaries when shouldEmitToolResult is false", () => {
let handler: ((evt: unknown) => void) | undefined;
const session: StubSession = {
subscribe: (fn) => {
handler = fn;
return () => {};
},
};
const onToolResult = vi.fn();
subscribeEmbeddedPiSession({
session: session as unknown as Parameters<
typeof subscribeEmbeddedPiSession
>[0]["session"],
runId: "run-tool-off",
shouldEmitToolResult: () => false,
onToolResult,
});
handler?.({
type: "tool_execution_start",
toolName: "read",
toolCallId: "tool-2",
args: { path: "/tmp/b.txt" },
});
expect(onToolResult).not.toHaveBeenCalled();
});
it("emits tool summaries when shouldEmitToolResult overrides verbose", () => {
let handler: ((evt: unknown) => void) | undefined;
const session: StubSession = {
subscribe: (fn) => {
handler = fn;
return () => {};
},
};
const onToolResult = vi.fn();
subscribeEmbeddedPiSession({
session: session as unknown as Parameters<
typeof subscribeEmbeddedPiSession
>[0]["session"],
runId: "run-tool-override",
verboseLevel: "off",
shouldEmitToolResult: () => true,
onToolResult,
});
handler?.({
type: "tool_execution_start",
toolName: "read",
toolCallId: "tool-3",
args: { path: "/tmp/c.txt" },
});
expect(onToolResult).toHaveBeenCalledTimes(1);
});
});

View File

@@ -2,10 +2,7 @@ import type { AgentEvent, AgentMessage } from "@mariozechner/pi-agent-core";
import type { AssistantMessage } from "@mariozechner/pi-ai";
import type { AgentSession } from "@mariozechner/pi-coding-agent";
import {
createToolDebouncer,
formatToolAggregate,
} from "../auto-reply/tool-meta.js";
import { formatToolAggregate } from "../auto-reply/tool-meta.js";
import { emitAgentEvent } from "../infra/agent-events.js";
import { splitMediaFromOutput } from "../media/parse.js";
import { defaultRuntime } from "../runtime.js";
@@ -113,6 +110,7 @@ export function subscribeEmbeddedPiSession(params: {
const assistantTexts: string[] = [];
const toolMetas: Array<{ toolName?: string; meta?: string }> = [];
const toolMetaById = new Map<string, string | undefined>();
const toolSummaryById = new Set<string>();
const blockReplyBreak = params.blockReplyBreak ?? "text_end";
let deltaBuffer = "";
let blockBuffer = "";
@@ -176,17 +174,25 @@ export function subscribeEmbeddedPiSession(params: {
return afterStart.slice(0, endIndex);
};
const toolDebouncer = createToolDebouncer((toolName, metas) => {
if (!params.onPartialReply) return;
const text = formatToolAggregate(toolName, metas);
const { text: cleanedText, mediaUrls } = splitMediaFromOutput(text);
void params.onPartialReply({
text: cleanedText,
mediaUrls: mediaUrls?.length ? mediaUrls : undefined,
});
});
const blockChunking = params.blockReplyChunking;
const shouldEmitToolResult = () =>
typeof params.shouldEmitToolResult === "function"
? params.shouldEmitToolResult()
: params.verboseLevel === "on";
const emitToolSummary = (toolName?: string, meta?: string) => {
if (!params.onToolResult) return;
const agg = formatToolAggregate(toolName, meta ? [meta] : undefined);
const { text: cleanedText, mediaUrls } = splitMediaFromOutput(agg);
if (!cleanedText && (!mediaUrls || mediaUrls.length === 0)) return;
try {
void params.onToolResult({
text: cleanedText,
mediaUrls: mediaUrls?.length ? mediaUrls : undefined,
});
} catch {
// ignore tool result delivery failures
}
};
const findSentenceBreak = (window: string, minChars: number): number => {
if (!window) return -1;
@@ -298,12 +304,12 @@ export function subscribeEmbeddedPiSession(params: {
assistantTexts.length = 0;
toolMetas.length = 0;
toolMetaById.clear();
toolSummaryById.clear();
deltaBuffer = "";
blockBuffer = "";
lastStreamedAssistant = undefined;
lastBlockReplyText = undefined;
assistantTextBaseline = 0;
toolDebouncer.flush();
};
const unsubscribe = params.session.subscribe(
@@ -336,6 +342,15 @@ export function subscribeEmbeddedPiSession(params: {
stream: "tool",
data: { phase: "start", name: toolName, toolCallId },
});
if (
params.onToolResult &&
shouldEmitToolResult() &&
!toolSummaryById.has(toolCallId)
) {
toolSummaryById.add(toolCallId);
emitToolSummary(toolName, meta);
}
}
if (evt.type === "tool_execution_update") {
@@ -382,7 +397,8 @@ export function subscribeEmbeddedPiSession(params: {
const sanitizedResult = sanitizeToolResult(result);
const meta = toolMetaById.get(toolCallId);
toolMetas.push({ toolName, meta });
toolDebouncer.push(toolName, meta);
toolMetaById.delete(toolCallId);
toolSummaryById.delete(toolCallId);
emitAgentEvent({
runId: params.runId,
@@ -406,25 +422,6 @@ export function subscribeEmbeddedPiSession(params: {
isError,
},
});
const emitToolResult =
typeof params.shouldEmitToolResult === "function"
? params.shouldEmitToolResult()
: params.verboseLevel === "on";
if (emitToolResult && params.onToolResult) {
const agg = formatToolAggregate(toolName, meta ? [meta] : undefined);
const { text: cleanedText, mediaUrls } = splitMediaFromOutput(agg);
if (cleanedText || (mediaUrls && mediaUrls.length > 0)) {
try {
void params.onToolResult({
text: cleanedText,
mediaUrls: mediaUrls?.length ? mediaUrls : undefined,
});
} catch {
// ignore tool result delivery failures
}
}
}
}
if (evt.type === "message_update") {
@@ -626,7 +623,6 @@ export function subscribeEmbeddedPiSession(params: {
if (evt.type === "agent_end") {
defaultRuntime.log?.(`embedded run agent end: runId=${params.runId}`);
toolDebouncer.flush();
if (pendingCompactionRetry > 0) {
resolveCompactionRetry();
} else {
@@ -640,7 +636,6 @@ export function subscribeEmbeddedPiSession(params: {
assistantTexts,
toolMetas,
unsubscribe,
flush: () => toolDebouncer.flush(),
waitForCompactionRetry: () => {
if (compactionInFlight || pendingCompactionRetry > 0) {
ensureCompactionPromise();