oh-my-pi/packages/coding-agent/test/streaming-edit-abort.test.ts

/**
 * Streaming edit abort tests.
 */

import { afterEach, beforeEach, expect, it, vi } from "bun:test";
import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";
import { Agent, type AgentTool } from "@oh-my-pi/pi-agent-core";
import type { AssistantMessage, StopReason, ToolCall } from "@oh-my-pi/pi-ai";
import { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream";
import { getBundledModel } from "@oh-my-pi/pi-catalog/models";
import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session";
import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage";
import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager";
import * as autoGeneratedGuard from "@oh-my-pi/pi-coding-agent/tools/auto-generated-guard";
import { ToolError } from "@oh-my-pi/pi-coding-agent/tools/tool-errors";
import { removeSyncWithRetries, Snowflake } from "@oh-my-pi/pi-utils";
import { type } from "arktype";

function createAssistantMessage(content: AssistantMessage["content"], stopReason: StopReason): AssistantMessage {
	return {
		role: "assistant",
		content,
		api: "anthropic-messages",
		provider: "anthropic",
		model: "mock",
		usage: {
			input: 0,
			output: 0,
			cacheRead: 0,
			cacheWrite: 0,
			totalTokens: 0,
			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
		},
		stopReason,
		timestamp: Date.now(),
	};
}

function createToolCall(id: string, args: Record<string, unknown>): ToolCall {
	return {
		type: "toolCall",
		id,
		name: "edit",
		arguments: args,
	};
}

function lastAssistantMessage(messages: Array<{ role: string }>): AssistantMessage | undefined {
	for (let i = messages.length - 1; i >= 0; i--) {
		const msg = messages[i];
		if (msg.role === "assistant") return msg as AssistantMessage;
	}
	return undefined;
}

function createRng(seed: number): () => number {
	let state = seed % 2147483647;
	if (state <= 0) state += 2147483646;
	return () => {
		state = (state * 48271) % 2147483647;
		return state / 2147483647;
	};
}

function chunkStringRandomly(text: string, seed: number): string[] {
	const rand = createRng(seed);
	const chunks: string[] = [];
	let offset = 0;
	while (offset < text.length) {
		const remaining = text.length - offset;
		const maxSize = Math.min(8, remaining);
		const size = Math.max(1, Math.floor(rand() * maxSize) + 1);
		chunks.push(text.slice(offset, offset + size));
		offset += size;
	}
	return chunks;
}

async function createSession(
	tempDir: string,
	streamFn: Agent["streamFn"],
	tool: AgentTool,
): Promise<{ agent: Agent; session: AgentSession; authStorage: AuthStorage }> {
	const model = getBundledModel("anthropic", "claude-sonnet-4-5")!;
	const agent = new Agent({
		getApiKey: () => "test-key",
		initialState: {
			model,
			systemPrompt: ["Test"],
			tools: [tool],
		},
		streamFn,
	});

	const sessionManager = SessionManager.inMemory(tempDir);
	const settings = Settings.isolated({ "edit.streamingAbort": true });
	const authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db"));
	authStorage.setRuntimeApiKey("anthropic", "test-key");
	const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml"));

	return {
		agent,
		session: new AgentSession({
			agent,
			sessionManager,
			settings,
			modelRegistry,
		}),
		authStorage,
	};
}

function buildEditTool(): AgentTool {
	const schema = type({
		path: "string",
		diff: "string",
		op: "string?",
		rename: "string?",
	});

	return {
		name: "edit",
		label: "Edit",
		description: "",
		parameters: schema,
		async execute() {
			return { content: [{ type: "text", text: "ok" }] };
		},
	};
}

function createStreamForDiff(
	path: string,
	chunks: string[],
	abortSignalRef: { current?: AbortSignal },
	streamStateRef?: { deltaCount: number; waitBeforeFirstDelta?: Promise<void> },
): Agent["streamFn"] {
	let callIndex = 0;
	return (_model, _context, options) => {
		abortSignalRef.current = options?.signal;
		const stream = new AssistantMessageEventStream();
		const toolCallId = "call_edit_1";
		let diffSoFar = "";
		let aborted = false;

		const notifyAbort = () => {
			if (aborted) return;
			aborted = true;
			const partialCall = createToolCall(toolCallId, { path, diff: diffSoFar });
			stream.push({
				type: "toolcall_delta",
				contentIndex: 0,
				delta: "",
				partial: createAssistantMessage([partialCall], "stop"),
			});
			stream.push({ type: "error", reason: "aborted", error: createAssistantMessage([], "aborted") });
		};

		options?.signal?.addEventListener("abort", notifyAbort, { once: true });

		queueMicrotask(async () => {
			if (callIndex > 0) {
				const finalMessage = createAssistantMessage([{ type: "text", text: "done" }], "stop");
				stream.push({ type: "done", reason: "stop", message: finalMessage });
				callIndex++;
				return;
			}
			const startMessage = createAssistantMessage([], "stop");
			stream.push({ type: "start", partial: startMessage });

			const startCall = createToolCall(toolCallId, { path, diff: "" });
			stream.push({ type: "toolcall_start", contentIndex: 0, partial: createAssistantMessage([startCall], "stop") });
			if (streamStateRef?.waitBeforeFirstDelta) {
				await streamStateRef.waitBeforeFirstDelta;
			}

			for (const chunk of chunks) {
				if (aborted) return;
				diffSoFar += chunk;
				if (streamStateRef) {
					streamStateRef.deltaCount += 1;
				}
				const partialCall = createToolCall(toolCallId, { path, diff: diffSoFar });
				stream.push({
					type: "toolcall_delta",
					contentIndex: 0,
					delta: chunk,
					partial: createAssistantMessage([partialCall], "stop"),
				});
				await Bun.sleep(0);
			}

			if (aborted) return;

			const finalCall = createToolCall(toolCallId, { path, diff: diffSoFar });
			const finalMessage = createAssistantMessage([finalCall], "toolUse");
			stream.push({ type: "toolcall_end", contentIndex: 0, toolCall: finalCall, partial: finalMessage });
			stream.push({ type: "done", reason: "toolUse", message: finalMessage });
			callIndex++;
		});

		return stream;
	};
}

let tempDir: string;
const editTool = buildEditTool();
// One deterministic seed is enough to exercise the streaming abort decision: seed 7 splits
// each diff into 6 chunks and fragments the decision-critical context line (-beta / -omega)
// across deltas, driving the partial-parse abort logic through intermediate states. Multi-seed
// fan-out re-ran the full session machinery per seed (a fresh AuthStorage SQLite open each time)
// without adding meaningful coverage of the success/fail contracts.
const seeds = [7];
const STREAMING_EDIT_RANDOM_STREAM_TIMEOUT_MS = 20_000;

beforeEach(() => {
	tempDir = path.join(os.tmpdir(), `pi-streaming-edit-${Snowflake.next()}`);
	fs.mkdirSync(tempDir, { recursive: true });
});

afterEach(async () => {
	if (tempDir) {
		removeSyncWithRetries(tempDir);
	}
});

it(
	"does not abort for successful patches across random streams",
	async () => {
		await Bun.write(path.join(tempDir, "sample.txt"), "alpha\nbeta\ngamma\n");
		const diff = "@@\n-beta\n+beta2\n";

		for (const seed of seeds) {
			const chunks = chunkStringRandomly(diff, seed);
			const abortSignalRef: { current?: AbortSignal } = {};
			const streamFn = createStreamForDiff("sample.txt", chunks, abortSignalRef);
			const { session, authStorage } = await createSession(tempDir, streamFn, editTool);

			try {
				await session.prompt("apply patch");

				const lastAssistant = lastAssistantMessage(session.state.messages);
				expect(lastAssistant?.stopReason).not.toBe("aborted");
				expect(abortSignalRef.current?.aborted ?? false).toBe(false);
			} finally {
				try {
					await session.dispose();
				} finally {
					authStorage.close();
				}
			}
		}
	},
	STREAMING_EDIT_RANDOM_STREAM_TIMEOUT_MS,
);

it(
	"aborts for failing patches across random streams",
	async () => {
		await Bun.write(path.join(tempDir, "sample.txt"), "alpha\nbeta\ngamma\n");
		const diff = "@@\n-omega\n+beta2\n";

		for (const seed of seeds) {
			const chunks = chunkStringRandomly(diff, seed);
			const abortSignalRef: { current?: AbortSignal } = {};
			const streamFn = createStreamForDiff("sample.txt", chunks, abortSignalRef);
			const { session, authStorage } = await createSession(tempDir, streamFn, editTool);

			try {
				await session.prompt("apply patch");

				const lastAssistant = lastAssistantMessage(session.state.messages);
				expect(lastAssistant?.stopReason).toBe("aborted");
				expect(abortSignalRef.current?.aborted ?? false).toBe(true);
			} finally {
				try {
					await session.dispose();
				} finally {
					authStorage.close();
				}
			}
		}
	},
	STREAMING_EDIT_RANDOM_STREAM_TIMEOUT_MS,
);

it("does not abort when auto-generated peek fails with ENOENT (non-ToolError)", async () => {
	const checkSpy = vi
		.spyOn(autoGeneratedGuard, "assertEditableFile")
		.mockRejectedValue(Object.assign(new Error("ENOENT"), { code: "ENOENT" }));

	await Bun.write(path.join(tempDir, "sample.txt"), "alpha\nbeta\ngamma\n");
	const diff = "@@\n-beta\n+beta2\n";
	const abortSignalRef: { current?: AbortSignal } = {};
	const streamFn = createStreamForDiff("sample.txt", chunkStringRandomly(diff, 7), abortSignalRef);
	const { agent, session, authStorage } = await createSession(tempDir, streamFn, editTool);
	const abortSpy = vi.spyOn(agent, "abort");

	try {
		await session.prompt("apply patch");

		expect(abortSpy).not.toHaveBeenCalled();
		expect(abortSignalRef.current?.aborted ?? false).toBe(false);
		const lastAssistant = lastAssistantMessage(session.state.messages);
		expect(lastAssistant?.stopReason).not.toBe("aborted");
	} finally {
		checkSpy.mockRestore();
		abortSpy.mockRestore();
		try {
			await session.dispose();
		} finally {
			authStorage.close();
		}
	}
});

it("aborts when auto-generated check rejects with ToolError", async () => {
	const checkSpy = vi
		.spyOn(autoGeneratedGuard, "assertEditableFile")
		.mockRejectedValue(new ToolError("Cannot modify auto-generated file"));

	await Bun.write(path.join(tempDir, "sample.txt"), "alpha\nbeta\ngamma\n");
	const diff = "@@\n-beta\n+beta2\n";
	const abortSignalRef: { current?: AbortSignal } = {};
	const streamFn = createStreamForDiff("sample.txt", chunkStringRandomly(diff, 7), abortSignalRef);
	const { agent, session, authStorage } = await createSession(tempDir, streamFn, editTool);
	const abortSpy = vi.spyOn(agent, "abort");

	try {
		await session.prompt("apply patch");

		expect(abortSpy).toHaveBeenCalled();
		expect(abortSignalRef.current?.aborted ?? false).toBe(true);
		const lastAssistant = lastAssistantMessage(session.state.messages);
		expect(lastAssistant?.stopReason).toBe("aborted");
	} finally {
		checkSpy.mockRestore();
		abortSpy.mockRestore();
		try {
			await session.dispose();
		} finally {
			authStorage.close();
		}
	}
});

it("resolves local:// internal-scheme paths through the protocol handler instead of panicking", async () => {
	// Plan-mode persists the plan file under the synthetic local:// URL scheme.
	// Earlier the streaming pre-cache (#preCacheStreamingEditFile →
	// #getStreamingEditToolCall) called resolveToCwd() unconditionally on the
	// path, which throws for internal-scheme URLs via assertNotInternalUrl().
	// The throw escaped the synchronous interceptor as an Unhandled Rejection
	// and killed the session.
	//
	// The fix routes `local://` paths through resolveLocalUrlToPath() so they
	// map onto the session's on-disk local-artifacts directory; pre-caching,
	// auto-generated detection, and post-edit invalidation all run on the real
	// file. Drive a streaming Edit toolcall with path: 'local://PLAN.md' and
	// confirm the session completes without panicking and that
	// assertEditableFile is invoked with a real (non-internal-scheme) absolute
	// path so the auto-generated guard works for plan-mode edits too.
	const checkSpy = vi.spyOn(autoGeneratedGuard, "assertEditableFile");
	const diff = "@@\n-old\n+new\n";
	const abortSignalRef: { current?: AbortSignal } = {};
	const chunks = chunkStringRandomly(diff, 7);
	const streamFn = createStreamForDiff("local://PLAN.md", chunks, abortSignalRef);
	const { session, authStorage } = await createSession(tempDir, streamFn, editTool);

	try {
		await session.prompt("edit plan file");

		const lastAssistant = lastAssistantMessage(session.state.messages);
		expect(lastAssistant).toBeDefined();
		expect(lastAssistant?.stopReason).not.toBe("aborted");
		expect(abortSignalRef.current?.aborted ?? false).toBe(false);

		// Confirm pre-cache resolved the URL rather than skipping it: the
		// auto-generated guard should have been invoked with a concrete fs path
		// (absolute, no internal scheme) plus the original local:// display path.
		expect(checkSpy).toHaveBeenCalled();
		const [absolutePath, displayPath] = checkSpy.mock.calls[0] ?? [];
		expect(typeof absolutePath).toBe("string");
		expect(absolutePath).toMatch(/^(?:\/|[A-Za-z]:[\\/])/);
		expect(absolutePath).not.toMatch(/^[a-z]+:\/\//);
		expect(displayPath).toBe("local://PLAN.md");
	} finally {
		checkSpy.mockRestore();
		try {
			await session.dispose();
		} finally {
			authStorage.close();
		}
	}
});

it("aborts auto-generated file edits as soon as the path is available", async () => {
	const generatedPath = path.join(tempDir, "generated.ts");
	await Bun.write(generatedPath, "// Code generated by sqlc. DO NOT EDIT.\nexport const foo = 1;\n");
	const abortSignalRef: { current?: AbortSignal } = {};
	const diff = "@@\n-export const foo = 1;\n+export const foo = 2;\n";
	const chunks = chunkStringRandomly(diff, 7);
	const waitBeforeFirstDelta = Promise.withResolvers<void>();
	const streamState = { deltaCount: 0, waitBeforeFirstDelta: waitBeforeFirstDelta.promise };
	const streamFn = createStreamForDiff("generated.ts", chunks, abortSignalRef, streamState);
	const { agent, session, authStorage } = await createSession(tempDir, streamFn, editTool);
	const abortSpy = vi.spyOn(agent, "abort");
	const checkStarted = Promise.withResolvers<void>();
	const releaseCheck = Promise.withResolvers<void>();
	const checkSpy = vi.spyOn(autoGeneratedGuard, "assertEditableFile").mockImplementation(async () => {
		checkStarted.resolve();
		await releaseCheck.promise;
		throw new ToolError("Cannot modify auto-generated file");
	});

	try {
		const promptPromise = session.prompt("apply patch");
		await checkStarted.promise;

		expect(streamState.deltaCount).toBe(0);

		releaseCheck.resolve();
		waitBeforeFirstDelta.resolve();
		await promptPromise;

		expect(checkSpy).toHaveBeenCalledWith(generatedPath, "generated.ts");
		expect(abortSpy).toHaveBeenCalled();
		expect(abortSignalRef.current?.aborted ?? false).toBe(true);
		const lastAssistant = lastAssistantMessage(session.state.messages);
		expect(lastAssistant?.stopReason).toBe("aborted");
	} finally {
		waitBeforeFirstDelta.resolve();
		releaseCheck.resolve();
		checkSpy.mockRestore();
		abortSpy.mockRestore();
		try {
			await session.dispose();
		} finally {
			authStorage.close();
		}
	}
});