Previously, when block streaming was disabled (the default), text generated between tool calls would only appear after all tools completed. This was because onBlockReply wasn't passed to the subscription when block streaming was off, so flushBlockReplyBuffer() before tool execution did nothing. Now onBlockReply is always passed, and when block streaming is disabled, block replies are sent directly during tool flush. Directly sent payloads are tracked to avoid duplicates in final payloads. Also fixes a race condition where tool summaries could be emitted before the typing indicator started by awaiting onAgentEvent in tool handlers.
135 lines
3.8 KiB
TypeScript
135 lines
3.8 KiB
TypeScript
import { describe, expect, it, vi } from "vitest";
|
|
import { subscribeEmbeddedPiSession } from "./pi-embedded-subscribe.js";
|
|
|
|
type StubSession = {
|
|
subscribe: (fn: (evt: unknown) => void) => () => void;
|
|
};
|
|
|
|
describe("subscribeEmbeddedPiSession", () => {
|
|
const _THINKING_TAG_CASES = [
|
|
{ tag: "think", open: "<think>", close: "</think>" },
|
|
{ tag: "thinking", open: "<thinking>", close: "</thinking>" },
|
|
{ tag: "thought", open: "<thought>", close: "</thought>" },
|
|
{ tag: "antthinking", open: "<antthinking>", close: "</antthinking>" },
|
|
] as const;
|
|
|
|
it("waits for multiple compaction retries before resolving", async () => {
|
|
const listeners: SessionEventHandler[] = [];
|
|
const session = {
|
|
subscribe: (listener: SessionEventHandler) => {
|
|
listeners.push(listener);
|
|
return () => {};
|
|
},
|
|
} as unknown as Parameters<typeof subscribeEmbeddedPiSession>[0]["session"];
|
|
|
|
const subscription = subscribeEmbeddedPiSession({
|
|
session,
|
|
runId: "run-3",
|
|
});
|
|
|
|
for (const listener of listeners) {
|
|
listener({ type: "auto_compaction_end", willRetry: true });
|
|
listener({ type: "auto_compaction_end", willRetry: true });
|
|
}
|
|
|
|
let resolved = false;
|
|
const waitPromise = subscription.waitForCompactionRetry().then(() => {
|
|
resolved = true;
|
|
});
|
|
|
|
await Promise.resolve();
|
|
expect(resolved).toBe(false);
|
|
|
|
for (const listener of listeners) {
|
|
listener({ type: "agent_end" });
|
|
}
|
|
|
|
await Promise.resolve();
|
|
expect(resolved).toBe(false);
|
|
|
|
for (const listener of listeners) {
|
|
listener({ type: "agent_end" });
|
|
}
|
|
|
|
await waitPromise;
|
|
expect(resolved).toBe(true);
|
|
});
|
|
it("emits tool summaries at tool start when verbose is on", async () => {
|
|
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" },
|
|
});
|
|
|
|
// Wait for async handler to complete
|
|
await Promise.resolve();
|
|
|
|
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("includes browser action metadata in tool summaries", async () => {
|
|
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-browser-tool",
|
|
verboseLevel: "on",
|
|
onToolResult,
|
|
});
|
|
|
|
handler?.({
|
|
type: "tool_execution_start",
|
|
toolName: "browser",
|
|
toolCallId: "tool-browser-1",
|
|
args: { action: "snapshot", targetUrl: "https://example.com" },
|
|
});
|
|
|
|
// Wait for async handler to complete
|
|
await Promise.resolve();
|
|
|
|
expect(onToolResult).toHaveBeenCalledTimes(1);
|
|
const payload = onToolResult.mock.calls[0][0];
|
|
expect(payload.text).toContain("🌐");
|
|
expect(payload.text).toContain("browser");
|
|
expect(payload.text).toContain("snapshot");
|
|
expect(payload.text).toContain("https://example.com");
|
|
});
|
|
});
|