Squashed 'third_party/pi/' content from commit d981de122

git-subtree-dir: third_party/pi
git-subtree-split: d981de1229ef899957bbe968bc8dcda02a21f477

packages/agent/test/proxy.test.ts

@@ -0,0 +1,129 @@
+import type { AssistantMessage, AssistantMessageEvent, Model } from "@earendil-works/pi-ai";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { type ProxyAssistantMessageEvent, streamProxy } from "../src/proxy.ts";
+
+const model: Model<"openai-responses"> = {
+	id: "gpt-5.4",
+	name: "GPT-5.4",
+	api: "openai-responses",
+	provider: "openai",
+	baseUrl: "https://api.openai.com/v1",
+	reasoning: true,
+	input: ["text"],
+	cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+	contextWindow: 400000,
+	maxTokens: 128000,
+};
+
+const usage: AssistantMessage["usage"] = {
+	input: 0,
+	output: 0,
+	cacheRead: 0,
+	cacheWrite: 0,
+	totalTokens: 0,
+	cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+};
+
+afterEach(() => {
+	vi.unstubAllGlobals();
+});
+
+describe("streamProxy", () => {
+	it("preserves tool-call metadata received only on toolcall_end", async () => {
+		const proxyEvents: ProxyAssistantMessageEvent[] = [
+			{ type: "start" },
+			{ type: "toolcall_start", contentIndex: 0, id: "call_test|fc_test", toolName: "lookup" },
+			{ type: "toolcall_delta", contentIndex: 0, delta: '{"value":"hello"}' },
+			{
+				type: "toolcall_end",
+				contentIndex: 0,
+				toolCall: {
+					type: "toolCall",
+					id: "call_test|fc_test",
+					name: "lookup",
+					arguments: { value: "hello" },
+					namespace: "dynamic_tools",
+				},
+			},
+			{ type: "done", reason: "toolUse", usage },
+		];
+		const body = proxyEvents.map((event) => `data: ${JSON.stringify(event)}\n\n`).join("");
+		vi.stubGlobal(
+			"fetch",
+			vi.fn(async () => new Response(body, { status: 200 })),
+		);
+
+		const stream = streamProxy(
+			model,
+			{ systemPrompt: "", messages: [] },
+			{
+				authToken: "test-token",
+				proxyUrl: "https://proxy.example.com",
+			},
+		);
+		const events: AssistantMessageEvent[] = [];
+		for await (const event of stream) events.push(event);
+		const result = await stream.result();
+		const endEvent = events.find((event) => event.type === "toolcall_end");
+
+		expect(endEvent).toMatchObject({
+			type: "toolcall_end",
+			toolCall: { namespace: "dynamic_tools" },
+		});
+		expect(result.content[0]).toMatchObject({
+			type: "toolCall",
+			arguments: { value: "hello" },
+			namespace: "dynamic_tools",
+		});
+	});
+
+	// Regression tests for https://github.com/earendil-works/pi/issues/8996
+	it("processes terminal metadata when the event is not newline-terminated", async () => {
+		const start = `data: ${JSON.stringify({ type: "start" })}\n\n`;
+		const done = `data: ${JSON.stringify({ type: "done", reason: "stop", usage, providerThinkingLevel: "high" })}`;
+		vi.stubGlobal(
+			"fetch",
+			vi.fn(async () => new Response(start + done, { status: 200 })),
+		);
+
+		const stream = streamProxy(
+			model,
+			{ systemPrompt: "", messages: [] },
+			{
+				authToken: "test-token",
+				proxyUrl: "https://proxy.example.com",
+			},
+		);
+		const events: AssistantMessageEvent[] = [];
+		for await (const event of stream) events.push(event);
+		const result = await stream.result();
+
+		expect(events.map((event) => event.type)).toEqual(["start", "done"]);
+		expect(result.stopReason).toBe("stop");
+		expect(result.providerThinkingLevel).toBe("high");
+	});
+
+	it("emits an error instead of hanging when the stream ends without a terminal event", async () => {
+		const body = `data: ${JSON.stringify({ type: "start" })}\n\n`;
+		vi.stubGlobal(
+			"fetch",
+			vi.fn(async () => new Response(body, { status: 200 })),
+		);
+
+		const stream = streamProxy(
+			model,
+			{ systemPrompt: "", messages: [] },
+			{
+				authToken: "test-token",
+				proxyUrl: "https://proxy.example.com",
+			},
+		);
+		const events: AssistantMessageEvent[] = [];
+		for await (const event of stream) events.push(event);
+		const result = await stream.result();
+
+		expect(events.map((event) => event.type)).toEqual(["start", "error"]);
+		expect(result.stopReason).toBe("error");
+		expect(result.errorMessage).toContain("Connection closed by proxy server");
+	});
+});