Deterministic Pi cold-restore overflow assistant reproduction for earendil-works/pi#7874

7874-cold-restore-overflow-assistant.test.ts

import { existsSync, mkdirSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import type { AssistantMessage, Context, Message } from "@earendil-works/pi-ai/compat";
import {
	type FauxProviderRegistration,
	fauxAssistantMessage,
	registerFauxProvider,
} from "@earendil-works/pi-ai/compat";
import { afterEach, describe, expect, it } from "vitest";
import { AuthStorage } from "../../../src/core/auth-storage.ts";
import { convertToLlm } from "../../../src/core/messages.ts";
import { createAgentSession } from "../../../src/core/sdk.ts";
import { SessionManager } from "../../../src/core/session-manager.ts";
import { SettingsManager } from "../../../src/core/settings-manager.ts";
import { createInMemoryModelRegistry, getModelRuntime } from "../../model-runtime-test-utils.ts";
import { createTestResourceLoader } from "../../utilities.ts";

type Scenario = {
	name: string;
	stopReason: "error" | "length";
	errorMessage?: string;
	contextWindow: number;
	seed: boolean;
};

function captureResponse(
	requests: Message[][],
	content: string,
	options: { stopReason?: AssistantMessage["stopReason"]; errorMessage?: string } = {},
) {
	return (context: Context) => {
		requests.push(structuredClone(context.messages));
		return fauxAssistantMessage(content, options);
	};
}

function isOverflowAssistant(message: Message, scenario: Scenario): boolean {
	return (
		message.role === "assistant" &&
		message.stopReason === scenario.stopReason &&
		message.content.some((part) => part.type === "text" && part.text === `${scenario.name} overflow response`)
	);
}

describe("#7874 cold restore replays overflow assistants removed by live recovery", () => {
	const cleanups: Array<() => void> = [];

	afterEach(() => {
		while (cleanups.length > 0) {
			cleanups.pop()?.();
		}
	});

	it.each<Scenario>([
		{
			name: "length",
			stopReason: "length",
			contextWindow: 1000,
			seed: false,
		},
		{
			name: "error",
			stopReason: "error",
			errorMessage: "prompt is too long",
			contextWindow: 100_000,
			seed: true,
		},
	])("reproduces the $name history mismatch without extensions or real APIs", async (scenario) => {
		const tempDir = join(tmpdir(), `pi-7874-${scenario.name}-${Date.now()}-${Math.random().toString(36).slice(2)}`);
		mkdirSync(tempDir, { recursive: true });
		cleanups.push(() => {
			if (existsSync(tempDir)) rmSync(tempDir, { recursive: true, force: true });
		});

		const faux: FauxProviderRegistration = registerFauxProvider({
			models: [{ id: "faux-1", contextWindow: scenario.contextWindow, maxTokens: 100 }],
		});
		cleanups.push(() => faux.unregister());
		const model = faux.getModel();
		const authStorage = AuthStorage.inMemory();
		await authStorage.modify(model.provider, async () => ({ type: "api_key", key: "faux-key" }));
		const modelRegistry = await createInMemoryModelRegistry(authStorage);
		modelRegistry.registerProvider(model.provider, {
			baseUrl: model.baseUrl,
			apiKey: "faux-key",
			api: faux.api,
			models: faux.models,
		});
		const modelRuntime = getModelRuntime(modelRegistry);
		const settings = { compaction: { keepRecentTokens: 1, reserveTokens: 0 } };
		const requests: Message[][] = [];
		const responses = [
			captureResponse(requests, `${scenario.name} overflow response`, {
				stopReason: scenario.stopReason,
				errorMessage: scenario.errorMessage,
			}),
			captureResponse(requests, `${scenario.name} compaction summary`),
			captureResponse(requests, `${scenario.name} recovered response`),
			captureResponse(requests, `${scenario.name} cold response`),
		];
		if (scenario.seed) {
			responses.unshift(captureResponse(requests, `${scenario.name} seed response`));
			responses.splice(3, 0, captureResponse(requests, `${scenario.name} split-turn summary`));
		}
		faux.setResponses(responses);

		const liveManager = SessionManager.create(tempDir, tempDir);
		const { session: liveSession } = await createAgentSession({
			cwd: tempDir,
			agentDir: tempDir,
			model,
			modelRuntime,
			noTools: "all",
			resourceLoader: createTestResourceLoader(),
			sessionManager: liveManager,
			settingsManager: SettingsManager.inMemory(settings),
		});
		cleanups.push(() => liveSession.dispose());
		let completedOverflowCompaction = false;
		liveSession.subscribe((event) => {
			if (event.type === "compaction_end" && event.reason === "overflow" && !event.aborted && event.willRetry) {
				completedOverflowCompaction = true;
			}
		});

		if (scenario.seed) {
			await liveSession.prompt("seed turn");
		}
		await liveSession.prompt("x".repeat(5000));

		expect(completedOverflowCompaction).toBe(true);
		const liveHistory = structuredClone(convertToLlm(liveSession.messages));
		expect(liveHistory.some((message) => isOverflowAssistant(message, scenario))).toBe(false);
		expect(
			liveManager
				.getEntries()
				.some(
					(entry) =>
						entry.type === "message" &&
						entry.message.role === "assistant" &&
						entry.message.stopReason === scenario.stopReason,
				),
		).toBe(true);

		const sessionFile = liveManager.getSessionFile();
		expect(sessionFile).toBeDefined();
		liveSession.dispose();

		const coldManager = SessionManager.open(sessionFile!, tempDir);
		const { session: coldSession } = await createAgentSession({
			cwd: tempDir,
			agentDir: tempDir,
			model,
			modelRuntime,
			noTools: "all",
			resourceLoader: createTestResourceLoader(),
			sessionManager: coldManager,
			settingsManager: SettingsManager.inMemory(settings),
		});
		cleanups.push(() => coldSession.dispose());
		coldSession.subscribe(() => {});
		await coldSession.prompt("after cold restore");

		const coldRequest = requests.at(-1);
		expect(coldRequest).toBeDefined();
		const coldHistory = coldRequest!.slice(0, -1);
		expect(coldHistory.filter((message) => isOverflowAssistant(message, scenario))).toHaveLength(1);
		expect(coldHistory).not.toEqual(liveHistory);
		expect(coldHistory.filter((message) => !isOverflowAssistant(message, scenario))).toEqual(liveHistory);
	});
});

README.md

Deterministic cold-restore overflow reproduction

This reproduces earendil-works/pi#7874 on Pi commit 2a9b4ebc680053c64e31f635b0b22d5e22564001 without extensions, network access, or provider credentials.

The two cases cover stopReason: "length" and an explicit context-overflow stopReason: "error". Each case follows the normal file-backed path:

AgentSession prompt -> overflow -> auto-compaction -> successful retry
-> dispose -> SessionManager.open -> createAgentSession -> next provider request

Copy 7874-cold-restore-overflow-assistant.test.ts to:

packages/coding-agent/test/suite/regressions/

Then run from packages/coding-agent:

node "$(git rev-parse --show-toplevel)/node_modules/vitest/dist/cli.js" --run test/suite/regressions/7874-cold-restore-overflow-assistant.test.ts

Expected result on the affected version: two passing characterization cases. They assert that the overflow assistant is absent from live provider history, remains in canonical JSONL, reappears in the first cold provider request, and is the only difference between the two histories.