pi-ext

import {
	createServer,
	type IncomingMessage,
	type ServerResponse,
} from "node:http";
import {
	createTestSession,
	type TestSession,
} from "@marcfargas/pi-test-harness";
import { Type } from "typebox";
import { afterEach, describe, expect, it } from "vitest";
import klaus from "../index";
import { type KlausQueryHandle, startSdkQuery } from "../src/agent-sdk";
import { MemorySessionStore } from "../src/cache";
import type { KlausContentEvent, KlausRequest } from "../src/protocol";

interface CapturedRequest {
	url: string;
	headers: IncomingMessage["headers"];
	body: unknown;
}

const requests: CapturedRequest[] = [];
let closeServer: (() => Promise<void>) | undefined;
let queryHandle: KlausQueryHandle | undefined;
let piSession: TestSession | undefined;

function messageRequests(): CapturedRequest[] {
	return requests
		.filter((request) => request.url.startsWith("/v1/messages"))
		.filter((request) => !request.url.includes("count_tokens"));
}

async function createKlausTestSession(): Promise<TestSession> {
	return createTestSession({ extensionFactories: [klaus] });
}

function sse(response: ServerResponse, frames: unknown[]): void {
	response.writeHead(200, {
		"content-type": "text/event-stream",
		"cache-control": "no-cache",
		connection: "keep-alive",
	});
	for (const frame of frames) {
		const type = (frame as { type: string }).type;
		response.write(`event: ${type}\ndata: ${JSON.stringify(frame)}\n\n`);
	}
	response.end();
}

async function startStalledAnthropic(): Promise<string> {
	const server = createServer(async (request, response) => {
		const chunks: Buffer[] = [];
		for await (const chunk of request) chunks.push(Buffer.from(chunk));
		const text = Buffer.concat(chunks).toString("utf8");
		requests.push({
			url: request.url ?? "",
			headers: request.headers,
			body: text ? JSON.parse(text) : undefined,
		});
		response.writeHead(200, {
			"content-type": "text/event-stream",
			"cache-control": "no-cache",
			connection: "keep-alive",
		});
		response.write(
			`event: message_start\ndata: ${JSON.stringify({
				type: "message_start",
				message: {
					id: "msg_stalled",
					type: "message",
					role: "assistant",
					model: "claude-sonnet-5",
					content: [],
					stop_reason: null,
					stop_sequence: null,
					usage: {
						input_tokens: 1,
						output_tokens: 0,
						cache_creation_input_tokens: 0,
						cache_read_input_tokens: 0,
					},
				},
			})}\n\n`,
		);
	});
	await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
	const address = server.address();
	if (!address || typeof address === "string")
		throw new Error("Fake Anthropic server failed to bind.");
	closeServer = async () => {
		server.closeAllConnections();
		await new Promise<void>((resolve, reject) =>
			server.close((error) => (error ? reject(error) : resolve())),
		);
	};
	return `http://127.0.0.1:${address.port}`;
}

async function startFakeAnthropic(toolScenario = false): Promise<string> {
	let messageCount = 0;
	const server = createServer(async (request, response) => {
		const chunks: Buffer[] = [];
		for await (const chunk of request) chunks.push(Buffer.from(chunk));
		const text = Buffer.concat(chunks).toString("utf8");
		requests.push({
			url: request.url ?? "",
			headers: request.headers,
			body: text ? JSON.parse(text) : undefined,
		});
		if (request.url?.startsWith("/v1/messages/count_tokens")) {
			response.writeHead(200, { "content-type": "application/json" });
			response.end(JSON.stringify({ input_tokens: 8 }));
			return;
		}
		if (!request.url?.startsWith("/v1/messages")) {
			response.writeHead(404).end();
			return;
		}
		messageCount += 1;
		const start = {
			type: "message_start",
			message: {
				id: `msg_klaus_${messageCount}`,
				type: "message",
				role: "assistant",
				model: "claude-sonnet-5",
				content: [],
				stop_reason: null,
				stop_sequence: null,
				usage: {
					input_tokens: 8,
					output_tokens: 0,
					cache_creation_input_tokens: 0,
					cache_read_input_tokens: 0,
				},
			},
		};
		if (toolScenario && messageCount === 1) {
			sse(response, [
				start,
				{
					type: "content_block_start",
					index: 0,
					content_block: {
						type: "tool_use",
						id: "toolu_klaus_1",
						name: "echo",
						input: {},
					},
				},
				{
					type: "content_block_delta",
					index: 0,
					delta: {
						type: "input_json_delta",
						partial_json: '{"value":"hello"}',
					},
				},
				{ type: "content_block_stop", index: 0 },
				{
					type: "message_delta",
					delta: { stop_reason: "tool_use", stop_sequence: null },
					usage: { output_tokens: 4 },
				},
				{ type: "message_stop" },
			]);
			return;
		}
		sse(response, [
			start,
			{
				type: "content_block_start",
				index: 0,
				content_block: { type: "text", text: "" },
			},
			{
				type: "content_block_delta",
				index: 0,
				delta: { type: "text_delta", text: "Klaus works." },
			},
			{ type: "content_block_stop", index: 0 },
			{
				type: "message_delta",
				delta: { stop_reason: "end_turn", stop_sequence: null },
				usage: { output_tokens: 3 },
			},
			{ type: "message_stop" },
		]);
	});
	await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
	const address = server.address();
	if (!address || typeof address === "string")
		throw new Error("Fake Anthropic server failed to bind.");
	closeServer = async () =>
		await new Promise<void>((resolve, reject) =>
			server.close((error) => (error ? reject(error) : resolve())),
		);
	return `http://127.0.0.1:${address.port}`;
}

afterEach(async () => {
	await queryHandle?.close("E2E finished.");
	queryHandle = undefined;
	if (piSession) {
		await piSession.session.extensionRunner.emit({
			type: "session_shutdown",
			reason: "quit",
		});
		piSession.dispose();
		piSession = undefined;
	}
	await closeServer?.();
	closeServer = undefined;
	requests.length = 0;
	delete process.env.KLAUS_E2E_BASE_URL;
});

describe("Klaus real Agent SDK E2E", () => {
	it(
		"streams through the registered Pi provider route",
		async () => {
			process.env.KLAUS_E2E_BASE_URL = await startFakeAnthropic();
			piSession = await createKlausTestSession();
			const runtime = piSession.session.modelRuntime;
			Object.defineProperty(runtime, "isUsingOAuth", {
				value: () => true,
			});
			Object.defineProperty(runtime, "getAuth", {
				value: async () => ({
					auth: { apiKey: "fake" },
					source: "OAuth",
				}),
			});
			const model = runtime.getModel("klaus", "claude-sonnet-5");
			if (!model) throw new Error("Klaus model was not registered.");

			const result = await runtime.completeSimple(
				model,
				{
					systemPrompt: "Answer exactly as scripted.",
					messages: [{ role: "user", content: "hello", timestamp: Date.now() }],
				},
				{
					sessionId: piSession.session.sessionManager.getSessionId(),
					onPayload: () => undefined,
				},
			);

			expect(result.stopReason).toBe("stop");
			expect(result.content).toContainEqual({
				type: "text",
				text: "Klaus works.",
			});
			expect(result.usage.input).toBe(8);
			expect(result.usage.output).toBe(3);
			expect(result.usage.cost.total).toBeGreaterThan(0);
			expect(messageRequests()).not.toHaveLength(0);
		},
		10 * 60_000,
	);

	it(
		"aborts a stalled real Claude request on Pi activity timeout",
		async () => {
			process.env.KLAUS_E2E_BASE_URL = await startStalledAnthropic();
			piSession = await createKlausTestSession();
			const runtime = piSession.session.modelRuntime;
			Object.defineProperty(runtime, "isUsingOAuth", { value: () => true });
			Object.defineProperty(runtime, "getAuth", {
				value: async () => ({
					auth: { apiKey: "fake" },
					source: "OAuth",
				}),
			});
			const model = runtime.getModel("klaus", "claude-sonnet-5");
			if (!model) throw new Error("Klaus model was not registered.");

			const result = await runtime.completeSimple(
				model,
				{
					messages: [{ role: "user", content: "stall", timestamp: Date.now() }],
				},
				{
					sessionId: piSession.session.sessionManager.getSessionId(),
					timeoutMs: 100,
				},
			);

			expect(result.stopReason).toBe("error");
			expect(result.errorMessage).toContain("timed out");
		},
		10 * 60_000,
	);

	it(
		"streams text through the real Claude child and fake Anthropic endpoint",
		async () => {
			process.env.KLAUS_E2E_BASE_URL = await startFakeAnthropic();
			const content: KlausContentEvent[] = [];
			const request: KlausRequest = {
				modelId: "claude-haiku-4-5",
				selector: "haiku",
				systemPrompt: "Answer exactly as scripted.",
				prompt: "hello",
				tools: [],
				thinking: "high",
				thinkingBudget: 1234,
				cwd: process.cwd(),
				headers: { "x-pi-e2e": "present", "x-klaus-private": "blocked" },
			};
			const finished = Promise.withResolvers<void>();
			queryHandle = await startSdkQuery(request, "fake", {
				onReady: async () => undefined,
				onActivity: () => undefined,
				onNotice: () => undefined,
				onContent: (event) => content.push(event),
				onToolBoundary: () =>
					finished.reject(new Error("Unexpected tool boundary.")),
				onResult: () => finished.resolve(),
				onError: (error) => finished.reject(error),
			});
			await finished.promise;

			expect(content).toContainEqual({
				type: "text-delta",
				index: 0,
				delta: "Klaus works.",
			});
			const messageRequest = messageRequests()[0];
			expect(messageRequest).toBeDefined();
			expect(messageRequest?.headers.authorization).toBe("Bearer fake");
			expect(messageRequest?.headers["x-pi-e2e"]).toBe("present");
			expect(
				Object.keys(messageRequest?.headers ?? {}).some((name) =>
					name.startsWith("x-klaus-"),
				),
			).toBe(false);
			expect(JSON.stringify(messageRequest?.body)).not.toContain("x-klaus-");
			expect(messageRequest?.body).toMatchObject({
				model: "claude-haiku-4-5-20251001",
				thinking: { type: "enabled", budget_tokens: 1234 },
			});
		},
		10 * 60_000,
	);

	it(
		"sends Pi images as native Agent SDK image blocks",
		async () => {
			process.env.KLAUS_E2E_BASE_URL = await startFakeAnthropic();
			const finished = Promise.withResolvers<void>();
			queryHandle = await startSdkQuery(
				{
					modelId: "claude-sonnet-5",
					selector: "sonnet",
					systemPrompt: "Describe the image.",
					prompt: "image",
					images: [
						{
							data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAusB9WlXz9sAAAAASUVORK5CYII=",
							mimeType: "image/png",
						},
					],
					tools: [],
					thinking: undefined,
					cwd: process.cwd(),
					headers: {},
				},
				"fake",
				{
					onReady: async () => undefined,
					onActivity: () => undefined,
					onNotice: () => undefined,
					onContent: () => undefined,
					onToolBoundary: () =>
						finished.reject(new Error("Unexpected tool boundary.")),
					onResult: () => finished.resolve(),
					onError: (error) => finished.reject(error),
				},
			);
			await finished.promise;

			const body = JSON.stringify(messageRequests()[0]?.body);
			expect(body).toContain('"type":"image"');
			expect(body).toContain('"media_type":"image/png"');
		},
		10 * 60_000,
	);

	it(
		"parks an MCP tool in Claude and resumes it with a Pi result",
		async () => {
			process.env.KLAUS_E2E_BASE_URL = await startFakeAnthropic(true);
			const boundary = Promise.withResolvers<void>();
			const finished = Promise.withResolvers<void>();
			const toolEnds: KlausContentEvent[] = [];
			const request: KlausRequest = {
				modelId: "claude-sonnet-5",
				selector: "sonnet",
				systemPrompt: "Call echo once.",
				prompt: "echo hello",
				tools: [
					{
						name: "echo",
						description: "Echo a value.",
						inputSchema: {
							type: "object",
							properties: { value: { type: "string" } },
							required: ["value"],
							additionalProperties: false,
						},
					},
				],
				thinking: undefined,
				cwd: process.cwd(),
				sessionId: "klaus-e2e-tools",
				headers: {},
			};
			queryHandle = await startSdkQuery(request, "fake", {
				onReady: async () => undefined,
				onActivity: () => undefined,
				onNotice: () => undefined,
				onContent: (event) => {
					if (event.type === "tool-end") toolEnds.push(event);
				},
				onToolBoundary: () => boundary.resolve(),
				onResult: () => finished.resolve(),
				onError: (error) => finished.reject(error),
			});
			await boundary.promise;
			expect(toolEnds).toContainEqual({
				type: "tool-end",
				index: 0,
				id: "toolu_klaus_1",
				name: "echo",
				arguments: { value: "hello" },
			});
			await queryHandle.bridge.waitForPending(["toolu_klaus_1"]);
			expect(
				queryHandle.bridge.settle({
					id: "toolu_klaus_1",
					content: [{ type: "text", text: "hello" }],
					isError: false,
				}),
			).toBe(true);
			await finished.promise;

			const messages = messageRequests();
			expect(messages).toHaveLength(2);
			expect(JSON.stringify(messages[0]?.body)).toContain('"name":"echo"');
			expect(JSON.stringify(messages[1]?.body)).toContain("hello");
		},
		10 * 60_000,
	);

	it(
		"restarts canonically when Pi changes tools across a parked call",
		async () => {
			process.env.KLAUS_E2E_BASE_URL = await startFakeAnthropic(true);
			piSession = await createKlausTestSession();
			const runtime = piSession.session.modelRuntime;
			Object.defineProperty(runtime, "isUsingOAuth", { value: () => true });
			Object.defineProperty(runtime, "getAuth", {
				value: async () => ({ auth: { apiKey: "fake" }, source: "OAuth" }),
			});
			const model = runtime.getModel("klaus", "claude-sonnet-5");
			if (!model) throw new Error("Klaus model was not registered.");
			const sessionId = piSession.session.sessionManager.getSessionId();
			const user = {
				role: "user" as const,
				content: "echo hello",
				timestamp: 1,
			};
			const first = await runtime.completeSimple(
				model,
				{
					messages: [user],
					tools: [
						{
							name: "echo",
							description: "Echo a value.",
							parameters: Type.Object({ value: Type.String() }),
						},
					],
				},
				{ sessionId },
			);
			const call = first.content.find((item) => item.type === "toolCall");
			if (call?.type !== "toolCall")
				throw new Error("Claude did not call echo.");
			const second = await runtime.completeSimple(
				model,
				{
					messages: [
						user,
						first,
						{
							role: "toolResult",
							toolCallId: call.id,
							toolName: call.name,
							content: [{ type: "text", text: "hello" }],
							isError: false,
							timestamp: 2,
						},
					],
					tools: [
						{
							name: "new_tool",
							description: "A newly active tool.",
							parameters: Type.Object({}),
						},
					],
				},
				{ sessionId },
			);

			expect(second.content).toContainEqual({
				type: "text",
				text: "Klaus works.",
			});
			const messages = messageRequests();
			expect(messages.length).toBeGreaterThanOrEqual(2);
			const replayBody = JSON.stringify(messages.at(-1)?.body);
			expect(replayBody).toContain('"name":"new_tool"');
			expect(replayBody).not.toContain('"name":"echo"');
		},
		10 * 60_000,
	);

	it(
		"resumes an opaque mirrored session through the public SessionStore",
		async () => {
			process.env.KLAUS_E2E_BASE_URL = await startFakeAnthropic();
			const store = new MemorySessionStore();
			let sdkSessionId = "";
			let firstPosition = "";
			const firstFinished = Promise.withResolvers<void>();
			queryHandle = await startSdkQuery(
				{
					modelId: "claude-sonnet-5",
					selector: "sonnet",
					systemPrompt: "Continue linearly.",
					prompt: "first",
					tools: [],
					thinking: undefined,
					cwd: process.cwd(),
					sessionId: "pi-session",
					headers: {},
					sessionStore: store,
				},
				"fake",
				{
					onReady: async () => undefined,
					onActivity: () => undefined,
					onNotice: () => undefined,
					onContent: () => undefined,
					onToolBoundary: () =>
						firstFinished.reject(new Error("Unexpected tool boundary.")),
					onResult: (result) => {
						sdkSessionId = result.sdkSessionId ?? "";
						firstPosition = result.position ?? "";
						firstFinished.resolve();
					},
					onError: (error) => firstFinished.reject(error),
				},
			);
			await firstFinished.promise;
			await queryHandle.close("First turn complete.");
			expect(sdkSessionId).not.toBe("");

			const secondFinished = Promise.withResolvers<void>();
			queryHandle = await startSdkQuery(
				{
					modelId: "claude-sonnet-5",
					selector: "sonnet",
					systemPrompt: "Continue linearly.",
					prompt: "second",
					tools: [],
					thinking: undefined,
					cwd: process.cwd(),
					sessionId: "pi-session",
					headers: {},
					sessionStore: store,
					resume: sdkSessionId,
				},
				"fake",
				{
					onReady: async () => undefined,
					onActivity: () => undefined,
					onNotice: () => undefined,
					onContent: () => undefined,
					onToolBoundary: () =>
						secondFinished.reject(new Error("Unexpected tool boundary.")),
					onResult: () => secondFinished.resolve(),
					onError: (error) => secondFinished.reject(error),
				},
			);
			await secondFinished.promise;
			await queryHandle.close("Second turn complete.");

			const branchFinished = Promise.withResolvers<void>();
			queryHandle = await startSdkQuery(
				{
					modelId: "claude-sonnet-5",
					selector: "sonnet",
					systemPrompt: "Continue linearly.",
					prompt: "branch",
					tools: [],
					thinking: undefined,
					cwd: process.cwd(),
					sessionId: "pi-session",
					headers: {},
					sessionStore: store,
					resume: sdkSessionId,
					resumeAt: firstPosition,
					forkSession: true,
				},
				"fake",
				{
					onReady: async () => undefined,
					onActivity: () => undefined,
					onNotice: () => undefined,
					onContent: () => undefined,
					onToolBoundary: () =>
						branchFinished.reject(new Error("Unexpected tool boundary.")),
					onResult: () => branchFinished.resolve(),
					onError: (error) => branchFinished.reject(error),
				},
			);
			await branchFinished.promise;

			const messages = messageRequests();
			expect(messages).toHaveLength(3);
			expect(JSON.stringify(messages[1]?.body)).toContain("second");
			expect(JSON.stringify(messages[2]?.body)).toContain("branch");
		},
		10 * 60_000,
	);
});