oh-my-pi/packages/coding-agent/test/agent-session-interrupted-thinking.test.ts

import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test";
import * as path from "node:path";
import { Agent } from "@oh-my-pi/pi-agent-core";
import type { Api, AssistantMessage, Model, ThinkingContent } from "@oh-my-pi/pi-ai";
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 {
	type CustomMessage,
	convertToLlm,
	INTERRUPTED_THINKING_MESSAGE_TYPE,
	USER_INTERRUPT_LABEL,
} from "@oh-my-pi/pi-coding-agent/session/messages";
import type { SessionEntry } from "@oh-my-pi/pi-coding-agent/session/session-entries";
import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager";
import { TempDir } from "@oh-my-pi/pi-utils";

const REASONING_TEXT = "I have partly reasoned through the implementation and should preserve this.";
const VISIBLE_TEXT = "visible interrupted text";

function emptyUsage(): AssistantMessage["usage"] {
	return {
		input: 0,
		output: 0,
		cacheRead: 0,
		cacheWrite: 0,
		totalTokens: 0,
		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
	};
}

function baseAssistant(model: Model<Api>, content: AssistantMessage["content"]): AssistantMessage {
	return {
		role: "assistant",
		content,
		api: model.api,
		provider: model.provider,
		model: model.id,
		usage: emptyUsage(),
		stopReason: "aborted",
		timestamp: Date.now(),
	};
}

function thinkingAssistant(model: Model<Api>, errorMessage: string): AssistantMessage {
	const thinking: ThinkingContent = { type: "thinking", thinking: REASONING_TEXT };
	return { ...baseAssistant(model, [thinking]), errorMessage };
}

function textAssistant(model: Model<Api>): AssistantMessage {
	return { ...baseAssistant(model, [{ type: "text", text: VISIBLE_TEXT }]), errorMessage: USER_INTERRUPT_LABEL };
}

function isAssistantMessage(message: unknown): message is AssistantMessage {
	return typeof message === "object" && message !== null && "role" in message && message.role === "assistant";
}

function isInterruptedThinkingMessage(message: unknown): message is CustomMessage {
	return (
		typeof message === "object" &&
		message !== null &&
		"role" in message &&
		message.role === "custom" &&
		"customType" in message &&
		message.customType === INTERRUPTED_THINKING_MESSAGE_TYPE
	);
}

async function emitAssistantEnd(
	session: AgentSession,
	sessionManager: SessionManager,
	message: AssistantMessage,
	waitFor: (entry: SessionEntry) => boolean,
): Promise<void> {
	const existing = sessionManager.getBranch().find(waitFor);
	if (existing) return;
	const appended = Promise.withResolvers<void>();
	const previous = sessionManager.onEntryAppended;
	sessionManager.onEntryAppended = entry => {
		previous?.(entry);
		if (!waitFor(entry)) return;
		sessionManager.onEntryAppended = previous;
		appended.resolve();
	};
	session.agent.emitExternalEvent({ type: "message_start", message });
	session.agent.emitExternalEvent({ type: "message_end", message });
	await appended.promise;
}

describe("AgentSession interrupted thinking persistence", () => {
	let tempDir: TempDir;
	let authStorage: AuthStorage;
	let session: AgentSession | undefined;

	beforeEach(async () => {
		tempDir = TempDir.createSync("@pi-interrupted-thinking-");
		authStorage = await AuthStorage.create(path.join(tempDir.path(), "auth.db"));
		authStorage.setRuntimeApiKey("anthropic", "anthropic-test-key");
	});

	afterEach(async () => {
		try {
			await session?.dispose();
		} finally {
			session = undefined;
			authStorage.close();
			tempDir.removeSync();
			vi.restoreAllMocks();
		}
	});

	function createSession(): { model: Model<Api>; sessionManager: SessionManager; session: AgentSession } {
		const model = getBundledModel("anthropic", "claude-sonnet-4-5");
		const agent = new Agent({
			getApiKey: () => "anthropic-test-key",
			initialState: { model, systemPrompt: ["Test"], tools: [], messages: [] },
			convertToLlm,
		});
		const settings = Settings.isolated({
			"advisor.enabled": false,
			"compaction.enabled": false,
			"retry.enabled": false,
			"todo.enabled": false,
		});
		settings.setModelRole("default", `${model.provider}/${model.id}`);
		const sessionManager = SessionManager.inMemory();
		session = new AgentSession({
			agent,
			sessionManager,
			settings,
			modelRegistry: new ModelRegistry(authStorage, path.join(tempDir.path(), "models.yml")),
		});
		return { model, sessionManager, session };
	}

	it("demotes native thinking from a user-interrupted assistant into hidden durable context", async () => {
		const harness = createSession();
		await emitAssistantEnd(
			harness.session,
			harness.sessionManager,
			thinkingAssistant(harness.model, USER_INTERRUPT_LABEL),
			entry => entry.type === "custom_message" && entry.customType === INTERRUPTED_THINKING_MESSAGE_TYPE,
		);

		const messages = harness.session.agent.state.messages;
		const assistant = messages.find(isAssistantMessage);
		expect(assistant).toBeDefined();
		expect(assistant?.content.some(block => block.type === "thinking")).toBe(false);
		const hidden = messages.find(isInterruptedThinkingMessage);
		expect(hidden).toBeDefined();
		expect(hidden?.display).toBe(false);
		expect(hidden?.attribution).toBe("agent");
		expect(typeof hidden?.content === "string" ? hidden.content : JSON.stringify(hidden?.content)).toContain(
			REASONING_TEXT,
		);
		expect(hidden?.details).toMatchObject({
			provider: "anthropic",
			model: "claude-sonnet-4-5",
			blockCount: 1,
		});
		expect(typeof (hidden?.details as { interruptedAt?: unknown } | undefined)?.interruptedAt).toBe("number");

		const branch = harness.sessionManager.getBranch();
		const assistantEntryIndex = branch.findIndex(
			entry => entry.type === "message" && entry.message.role === "assistant",
		);
		const hiddenEntryIndex = branch.findIndex(
			entry => entry.type === "custom_message" && entry.customType === INTERRUPTED_THINKING_MESSAGE_TYPE,
		);
		expect(assistantEntryIndex).toBeGreaterThanOrEqual(0);
		expect(hiddenEntryIndex).toBe(assistantEntryIndex + 1);
		const assistantEntry = branch[assistantEntryIndex];
		if (assistantEntry?.type !== "message" || assistantEntry.message.role !== "assistant") {
			throw new Error("assistant entry was not persisted");
		}
		expect(assistantEntry.message.content.some(block => block.type === "thinking")).toBe(false);
		const hiddenEntry = branch[hiddenEntryIndex];
		if (hiddenEntry?.type !== "custom_message") throw new Error("interrupted-thinking entry was not persisted");
		expect(hiddenEntry.display).toBe(false);
		expect(hiddenEntry.attribution).toBe("agent");
		expect(
			typeof hiddenEntry.content === "string" ? hiddenEntry.content : JSON.stringify(hiddenEntry.content),
		).toContain(REASONING_TEXT);
	});

	it("leaves native thinking on non-user aborts and does not append hidden context", async () => {
		const harness = createSession();
		await emitAssistantEnd(
			harness.session,
			harness.sessionManager,
			thinkingAssistant(harness.model, "Request was aborted"),
			entry => entry.type === "message" && entry.message.role === "assistant",
		);

		const messages = harness.session.agent.state.messages;
		const assistant = messages.find(isAssistantMessage);
		expect(assistant?.content.some(block => block.type === "thinking")).toBe(true);
		expect(messages.some(isInterruptedThinkingMessage)).toBe(false);
		expect(
			harness.sessionManager
				.getBranch()
				.some(entry => entry.type === "custom_message" && entry.customType === INTERRUPTED_THINKING_MESSAGE_TYPE),
		).toBe(false);
	});

	it("leaves user-interrupted text-only assistant content unchanged", async () => {
		const harness = createSession();
		await emitAssistantEnd(
			harness.session,
			harness.sessionManager,
			textAssistant(harness.model),
			entry => entry.type === "message" && entry.message.role === "assistant",
		);

		const messages = harness.session.agent.state.messages;
		const assistant = messages.find(isAssistantMessage);
		expect(assistant?.content).toEqual([{ type: "text", text: VISIBLE_TEXT }]);
		expect(messages.some(isInterruptedThinkingMessage)).toBe(false);
		expect(
			harness.sessionManager
				.getBranch()
				.some(entry => entry.type === "custom_message" && entry.customType === INTERRUPTED_THINKING_MESSAGE_TYPE),
		).toBe(false);
	});
});