import { existsSync, mkdirSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { Agent } from "@earendil-works/pi-agent-core";
import {
	type AssistantMessage,
	type AssistantMessageEvent,
	EventStream,
	getModel,
	type ImageContent,
	type TextContent,
} from "@earendil-works/pi-ai";
import { Type } from "typebox";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import { createAgentSessionMessage } from "../src/core/agent-messages.js";
import { AgentSession } from "../src/core/agent-session.js";
import { AuthStorage } from "../src/core/auth-storage.js";
import type { CompactionResult } from "../src/core/compaction/index.js";
import type { AgentCronJob } from "../src/core/cron-jobs.js";
import { ModelRegistry } from "../src/core/model-registry.js";
import { SessionManager } from "../src/core/session-manager.js";
import { SettingsManager } from "../src/core/settings-manager.js";
import type { BuildSystemPromptOptions } from "../src/core/system-prompt.js";
import { createTestExtensionsResult, createTestResourceLoader } from "./utilities.js";

class MockAssistantStream extends EventStream<AssistantMessageEvent, AssistantMessage> {
	constructor() {
		super(
			(event) => event.type === "done" || event.type === "error",
			(event) => {
				if (event.type === "done") return event.message;
				if (event.type === "error") return event.error;
				throw new Error("Unexpected event type");
			},
		);
	}
}

function createAssistantMessage(text: string): AssistantMessage {
	return {
		role: "assistant",
		content: [{ type: "text", text }],
		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: "stop",
		timestamp: Date.now(),
	};
}

describe("AgentSession concurrent prompt guard", () => {
	let session: AgentSession;
	let tempDir: string;

	beforeEach(() => {
		tempDir = join(tmpdir(), `pi-concurrent-test-${Date.now()}`);
		mkdirSync(tempDir, { recursive: true });
	});

	afterEach(async () => {
		delete (globalThis as typeof globalThis & { testExtensionApi?: unknown }).testExtensionApi;
		delete (globalThis as typeof globalThis & { testCommandRuns?: unknown }).testCommandRuns;
		if (session) {
			session.dispose();
		}
		if (tempDir && existsSync(tempDir)) {
			rmSync(tempDir, { recursive: true });
		}
	});

	function createSession() {
		const model = getModel("anthropic", "claude-sonnet-4-5")!;
		let abortSignal: AbortSignal | undefined;

		const agent = new Agent({
			getApiKey: () => "test-key",
			initialState: {
				model,
				systemPrompt: "Test",
				tools: [],
			},
			streamFn: (_model, _context, options) => {
				abortSignal = options?.signal;
				const stream = new MockAssistantStream();
				queueMicrotask(() => {
					stream.push({ type: "start", partial: createAssistantMessage("") });
					const checkAbort = () => {
						if (abortSignal?.aborted) {
							stream.push({ type: "error", reason: "aborted", error: createAssistantMessage("Aborted") });
						} else {
							setTimeout(checkAbort, 5);
						}
					};
					checkAbort();
				});
				return stream;
			},
		});

		const sessionManager = SessionManager.inMemory();
		const settingsManager = SettingsManager.create(tempDir, tempDir);
		const authStorage = AuthStorage.create(join(tempDir, "auth.json"));
		const modelRegistry = ModelRegistry.create(authStorage, tempDir);
		authStorage.setRuntimeApiKey("anthropic", "test-key");

		session = new AgentSession({
			agent,
			sessionManager,
			settingsManager,
			cwd: tempDir,
			modelRegistry,
			resourceLoader: createTestResourceLoader(),
		});

		return session;
	}

	it("awaits asynchronous dispose callbacks during graceful disposal", async () => {
		createSession();
		let releaseCallback: () => void = () => {};
		const callbackGate = new Promise<void>((resolve) => {
			releaseCallback = resolve;
		});
		let callbackStarted = false;
		session.registerDisposeCallback(async () => {
			callbackStarted = true;
			await callbackGate;
		});

		let disposed = false;
		const disposal = session.disposeAsync().then(() => {
			disposed = true;
		});
		await new Promise<void>((resolve) => setTimeout(resolve, 0));
		expect(callbackStarted).toBe(true);
		expect(disposed).toBe(false);
		releaseCallback();
		await disposal;
		expect(disposed).toBe(true);
	});

	it("awaits dispose callbacks when synchronous disposal wins during the refinement drain", async () => {
		createSession();
		let releaseDrain: () => void = () => {};
		const drainGate = new Promise<void>((resolve) => {
			releaseDrain = resolve;
		});
		let drainStarted = false;
		const internals = session as unknown as {
			_drainPendingRefinementForDisposal: () => Promise<void>;
		};
		vi.spyOn(internals, "_drainPendingRefinementForDisposal").mockImplementation(async () => {
			drainStarted = true;
			await drainGate;
		});

		let releaseCallback: () => void = () => {};
		const callbackGate = new Promise<void>((resolve) => {
			releaseCallback = resolve;
		});
		let callbackStarted = false;
		session.registerDisposeCallback(async () => {
			callbackStarted = true;
			await callbackGate;
		});

		let disposed = false;
		const disposal = session.disposeAsync().then(() => {
			disposed = true;
		});
		await vi.waitFor(() => expect(drainStarted).toBe(true));
		session.dispose();
		expect(callbackStarted).toBe(true);

		releaseDrain();
		await new Promise<void>((resolve) => setTimeout(resolve, 0));
		expect(disposed).toBe(false);
		releaseCallback();
		await disposal;
		expect(disposed).toBe(true);
	});

	it("forwards kernelSnapshot: false to the kernel provisioner during disposal", async () => {
		createSession();
		const dispose = vi.fn(async () => {});
		Reflect.set(session, "_ipythonKernelProvisioner", { dispose });

		await session.disposeAsync({ kernelSnapshot: false });
		expect(dispose).toHaveBeenCalledWith({ snapshot: false });
	});

	it("should throw when prompt() called while streaming", async () => {
		createSession();

		const firstPrompt = session.prompt("First message");
		await new Promise((resolve) => setTimeout(resolve, 10));
		expect(session.isStreaming).toBe(true);

		await expect(session.prompt("Second message")).rejects.toThrow(
			"Agent is already processing. Specify streamingBehavior ('steer' or 'followUp') to queue the message.",
		);

		await session.abort();
		await firstPrompt.catch(() => {});
	});

	it("should allow steer() while streaming", async () => {
		createSession();

		const firstPrompt = session.prompt("First message");
		await new Promise((resolve) => setTimeout(resolve, 10));

		expect(() => session.steer("Steering message")).not.toThrow();
		expect(session.queuedActionCount).toBe(1);

		await session.abort();
		await firstPrompt.catch(() => {});
	});

	it("should allow followUp() while streaming", async () => {
		createSession();

		const firstPrompt = session.prompt("First message");
		await new Promise((resolve) => setTimeout(resolve, 10));

		expect(() => session.followUp("Follow-up message")).not.toThrow();
		expect(session.queuedActionCount).toBe(1);

		await session.abort();
		await firstPrompt.catch(() => {});
	});

	it("should queue extension-origin steering messages while streaming", async () => {
		const model = getModel("anthropic", "claude-sonnet-4-5")!;
		let abortSignal: AbortSignal | undefined;
		let sawSteeringMessage = false;
		let lastInputSource: string | undefined;
		const queueEvents: Array<{ steering: readonly string[]; followUp: readonly string[] }> = [];

		const agent = new Agent({
			getApiKey: () => "test-key",
			initialState: {
				model,
				systemPrompt: "Test",
				tools: [],
			},
			streamFn: (_model, context, options) => {
				abortSignal = options?.signal;
				const stream = new MockAssistantStream();
				queueMicrotask(() => {
					const userTexts = context.messages
						.filter((message) => message.role === "user")
						.map((message) => {
							if (typeof message.content === "string") {
								return message.content;
							}
							return message.content
								.filter((part): part is TextContent | ImageContent => typeof part === "object" && part !== null)
								.filter((part): part is TextContent => part.type === "text")
								.map((part) => part.text)
								.join("\n");
						});

					if (userTexts.includes("Steer from extension")) {
						sawSteeringMessage = true;
						stream.push({ type: "start", partial: createAssistantMessage("") });
						stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Steered") });
						return;
					}

					stream.push({ type: "start", partial: createAssistantMessage("") });
					const checkAbort = () => {
						if (abortSignal?.aborted) {
							stream.push({ type: "error", reason: "aborted", error: createAssistantMessage("Aborted") });
						} else {
							setTimeout(checkAbort, 5);
						}
					};
					checkAbort();
				});
				return stream;
			},
		});

		const sessionManager = SessionManager.inMemory();
		const settingsManager = SettingsManager.create(tempDir, tempDir);
		const authStorage = AuthStorage.create(join(tempDir, "auth.json"));
		const modelRegistry = ModelRegistry.create(authStorage, tempDir);
		authStorage.setRuntimeApiKey("anthropic", "test-key");

		const extensionsResult = await createTestExtensionsResult([
			(pi) => {
				(globalThis as typeof globalThis & { testExtensionApi?: unknown }).testExtensionApi = pi;
			},
			(pi) => {
				pi.on("input", async (event) => {
					lastInputSource = event.source;
				});
			},
		]);

		session = new AgentSession({
			agent,
			sessionManager,
			settingsManager,
			cwd: tempDir,
			modelRegistry,
			resourceLoader: createTestResourceLoader({ extensionsResult }),
		});
		session.subscribe((event) => {
			if (event.type === "session_action_update") {
				queueEvents.push({ steering: event.actions.steering, followUp: event.actions.followUps });
			}
		});

		const firstPrompt = session.prompt("First message");
		await new Promise((resolve) => setTimeout(resolve, 10));
		expect(session.isStreaming).toBe(true);

		const pi = (
			globalThis as typeof globalThis & {
				testExtensionApi?: {
					sendUserMessage: (content: string, options?: { deliverAs?: "steer" | "followUp" }) => void;
				};
			}
		).testExtensionApi;
		expect(pi).toBeDefined();

		pi!.sendUserMessage("Steer from extension", { deliverAs: "steer" });
		await new Promise((resolve) => setTimeout(resolve, 25));

		expect(session.queuedActionCount).toBe(1);
		expect(session.getSteeringMessages()).toContain("Steer from extension");
		expect(lastInputSource).toBe("extension");
		expect(queueEvents.some((event) => event.steering.includes("Steer from extension"))).toBe(true);

		await session.abort();
		await firstPrompt.catch(() => {});

		expect(sawSteeringMessage).toBe(false);
		expect(session.getSteeringMessages()).toContain("Steer from extension");
		expect(session.queuedActionCount).toBe(1);

		await session.prompt("After abort");

		expect(sawSteeringMessage).toBe(true);
		expect(session.queuedActionCount).toBe(0);
	});

	it("delivers accepted agent messages without extension input interception", async () => {
		const model = getModel("anthropic", "claude-sonnet-4-5")!;
		let inputCalls = 0;
		let receivedUserText: string | undefined;
		const agent = new Agent({
			getApiKey: () => "test-key",
			initialState: {
				model,
				systemPrompt: "Test",
				tools: [],
			},
			streamFn: (_model, context) => {
				const stream = new MockAssistantStream();
				queueMicrotask(() => {
					const userMessages = context.messages.filter((message) => message.role === "user");
					const user = userMessages.at(-1);
					if (user && typeof user.content !== "string") {
						receivedUserText = user.content
							.filter(
								(part): part is TextContent =>
									typeof part === "object" && part !== null && part.type === "text",
							)
							.map((part) => part.text)
							.join("\n");
					}
					stream.push({ type: "start", partial: createAssistantMessage("") });
					stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Delivered") });
				});
				return stream;
			},
		});

		const sessionManager = SessionManager.inMemory();
		const settingsManager = SettingsManager.create(tempDir, tempDir);
		const authStorage = AuthStorage.create(join(tempDir, "auth.json"));
		const modelRegistry = ModelRegistry.create(authStorage, tempDir);
		authStorage.setRuntimeApiKey("anthropic", "test-key");
		const extensionsResult = await createTestExtensionsResult([
			(pi) => {
				pi.on("input", async () => {
					inputCalls++;
					return { action: "handled" };
				});
			},
		]);

		session = new AgentSession({
			agent,
			sessionManager,
			settingsManager,
			cwd: tempDir,
			modelRegistry,
			resourceLoader: createTestResourceLoader({ extensionsResult }),
		});

		await session.acceptAgentMessagePrompt("agent-to-agent payload", { expandPromptTemplates: false });
		await session.agent.waitForIdle();

		expect(inputCalls).toBe(0);
		expect(receivedUserText).toBe("agent-to-agent payload");
	});

	it("queues an agent message behind an active turn", async () => {
		createSession();
		const activeTurn = session.prompt("active turn");
		await vi.waitFor(() => expect(session.isStreaming).toBe(true));
		let queued: boolean | undefined;

		await session.acceptAgentMessagePrompt("message during turn", {
			queueIfBusy: true,
			streamingBehavior: "steer",
			preflightResult: (accepted, didQueue) => {
				expect(accepted).toBe(true);
				queued = didQueue;
			},
		});

		expect(queued).toBe(true);
		expect(session.getSessionActionSnapshot().steering).toContain("message during turn");
		await session.abort();
		await activeTurn.catch(() => undefined);
	});

	it("serializes an agent message behind cron admission", async () => {
		createSession();
		const now = new Date().toISOString();
		const job: AgentCronJob = {
			id: "cron-overlap",
			status: "active",
			source: "cron",
			activeSessionId: "active-1",
			sessionId: session.sessionId,
			sessionFile: "/tmp/session.jsonl",
			cwd: tempDir,
			prompt: "cron prompt",
			schedule: { kind: "interval", expression: "every 5m", intervalMs: 300_000 },
			createdAt: now,
			updatedAt: now,
			runCount: 0,
		};
		const cronTurn = session.promptHeartbeat(job, { streamingBehavior: "followUp", source: "rpc" });
		let queued: boolean | undefined;

		await session.acceptAgentMessagePrompt("message during cron admission", {
			queueIfBusy: true,
			streamingBehavior: "steer",
			preflightResult: (accepted, didQueue) => {
				expect(accepted).toBe(true);
				queued = didQueue;
			},
		});

		expect(queued).toBe(true);
		expect(session.getSessionActionSnapshot().steering).toContain("message during cron admission");
		await session.abort();
		await cronTurn.catch(() => undefined);
	});

	it("does not admit a cron prompt invalidated while async input handlers ran", async () => {
		const model = getModel("anthropic", "claude-sonnet-4-5")!;
		let deliveredUserText: string | undefined;
		const agent = new Agent({
			getApiKey: () => "test-key",
			initialState: {
				model,
				systemPrompt: "Test",
				tools: [],
			},
			streamFn: (_model, context) => {
				const stream = new MockAssistantStream();
				queueMicrotask(() => {
					const user = context.messages.filter((message) => message.role === "user").at(-1);
					if (user && typeof user.content !== "string") {
						deliveredUserText = user.content
							.filter(
								(part): part is TextContent =>
									typeof part === "object" && part !== null && part.type === "text",
							)
							.map((part) => part.text)
							.join("\n");
					}
					stream.push({ type: "start", partial: createAssistantMessage("") });
					stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Delivered") });
				});
				return stream;
			},
		});

		const sessionManager = SessionManager.inMemory();
		const settingsManager = SettingsManager.create(tempDir, tempDir);
		const authStorage = AuthStorage.create(join(tempDir, "auth.json"));
		const modelRegistry = ModelRegistry.create(authStorage, tempDir);
		authStorage.setRuntimeApiKey("anthropic", "test-key");
		let jobInvalidated = false;
		const extensionsResult = await createTestExtensionsResult([
			(pi) => {
				pi.on("input", async () => {
					jobInvalidated = true;
					return undefined;
				});
			},
		]);

		session = new AgentSession({
			agent,
			sessionManager,
			settingsManager,
			cwd: tempDir,
			modelRegistry,
			resourceLoader: createTestResourceLoader({ extensionsResult }),
		});

		// Mirrors the daemon cron path: admissionCommitted re-checks the job and
		// throws when it was cancelled or updated before the prompt is admitted.
		await expect(
			session.promptUntilAccepted("stale cron prompt", {
				streamingBehavior: "followUp",
				source: "rpc",
				admissionCommitted: () => {
					if (jobInvalidated) throw new Error("Cron job became unrunnable before admission");
				},
			}),
		).rejects.toThrow("Cron job became unrunnable before admission");

		await session.agent.waitForIdle();
		expect(deliveredUserText).toBeUndefined();
		expect(session.isStreaming).toBe(false);
	});

	it("serializes concurrent agent messages", async () => {
		createSession();
		const dispositions = new Map<string, boolean | undefined>();
		const accept = (message: string) =>
			session.acceptAgentMessagePrompt(message, {
				queueIfBusy: true,
				streamingBehavior: "steer",
				preflightResult: (accepted, queued) => {
					expect(accepted).toBe(true);
					dispositions.set(message, queued);
				},
			});

		await Promise.all([accept("first concurrent message"), accept("second concurrent message")]);

		expect(dispositions.get("first concurrent message")).toBe(false);
		expect(dispositions.get("second concurrent message")).toBe(true);
		expect(session.getSessionActionSnapshot().steering).toContain("second concurrent message");
		await session.abort();
	});

	it("enforces the agent message queue cap inside core admission", async () => {
		createSession();
		const accept = (index: number) => {
			const message = createAgentSessionMessage({
				id: `agentmsg-${index}`,
				source: "agent_message",
				message: `message ${index}`,
				from: { clientId: "test" },
				target: {
					activeSessionId: "target",
					sessionId: session.sessionId,
					runtimeKind: "subagent",
				},
			});
			return session.queueAgentMessagePrompt(message.content, "steer", message);
		};

		for (let index = 0; index < 20; index++) {
			await accept(index);
		}
		await expect(accept(20)).rejects.toThrow("Target session has too many pending messages");
		expect(session.unfinishedActionCount).toBe(20);
		const snapshot = session.getSessionActionRecoverySnapshot();
		session.dispose();

		createSession();
		await expect(session.restoreSessionActions(snapshot)).resolves.toBe(20);
		expect(session.unfinishedActionCount).toBe(20);
	});

	it("rejects an agent message when clear wins core admission", async () => {
		createSession();
		const internals = session as unknown as {
			_acquireSessionActionCommitFence(): Promise<{ release(): void }>;
		};
		const fence = await internals._acquireSessionActionCommitFence();
		const message = createAgentSessionMessage({
			id: "agentmsg-clear-race",
			source: "agent_message",
			message: "clear race",
			from: { clientId: "test" },
			target: {
				activeSessionId: "target",
				sessionId: session.sessionId,
				runtimeKind: "subagent",
			},
		});
		const accepted = session.acceptAgentMessagePrompt(message.content, {
			customMessage: message,
			queueIfBusy: true,
			streamingBehavior: "steer",
		});
		await Promise.resolve();

		session.clearQueuedAgentMessages();
		fence.release();

		await expect(accepted).rejects.toThrow("Agent message was cleared before admission");
		expect(session.unfinishedActionCount).toBe(0);
	});

	it("accepts agent messages during and after direct compaction", async () => {
		createSession();
		let releaseCompaction: (result: CompactionResult) => void = () => {};
		const compactionGate = new Promise<CompactionResult>((resolve) => {
			releaseCompaction = resolve;
		});
		const internals = session as unknown as {
			_performCompaction: () => Promise<CompactionResult>;
		};
		vi.spyOn(internals, "_performCompaction").mockImplementation(() => compactionGate);
		const compaction = session.compact();
		await vi.waitFor(() => expect(session.isCompacting).toBe(true));
		let queued: boolean | undefined;

		await session.acceptAgentMessagePrompt("message during direct compaction", {
			queueIfBusy: true,
			streamingBehavior: "steer",
			preflightResult: (accepted, didQueue) => {
				expect(accepted).toBe(true);
				queued = didQueue;
			},
		});

		expect(queued).toBe(true);
		expect(session.getSessionActionSnapshot().steering).toContain("message during direct compaction");
		releaseCompaction({ summary: "compacted", firstKeptEntryId: "entry-1", tokensBefore: 1 });
		await compaction;
		let postCompactionQueued: boolean | undefined;
		await session.acceptAgentMessagePrompt("message after direct compaction", {
			queueIfBusy: true,
			streamingBehavior: "steer",
			preflightResult: (accepted, didQueue) => {
				expect(accepted).toBe(true);
				postCompactionQueued = didQueue;
			},
		});
		expect(postCompactionQueued).toBe(true);
		const deliveryCount = (text: string) =>
			session.messages.filter(
				(message) =>
					message.role === "user" &&
					typeof message.content !== "string" &&
					message.content.some((part) => part.type === "text" && part.text === text),
			).length;
		expect(deliveryCount("message during direct compaction")).toBe(0);
		expect(deliveryCount("message after direct compaction")).toBe(0);

		session.resumeQueuedWork();
		await vi.waitFor(() => expect(deliveryCount("message during direct compaction")).toBe(1));
		expect(session.getSessionActionSnapshot().steering).toContain("message after direct compaction");
	});

	it("delivers an accepted agent message once after scheduler-owned compaction", async () => {
		createSession();
		let releaseCompaction: (result: CompactionResult) => void = () => {};
		const compactionGate = new Promise<CompactionResult>((resolve) => {
			releaseCompaction = resolve;
		});
		const internals = session as unknown as {
			_performCompaction: () => Promise<CompactionResult>;
		};
		vi.spyOn(internals, "_performCompaction").mockImplementation(() => compactionGate);
		const compaction = session.compact(undefined, { skipAbort: true });
		await vi.waitFor(() => expect(session.isCompacting).toBe(true));
		let queued: boolean | undefined;

		await session.acceptAgentMessagePrompt("message during scheduled compaction", {
			queueIfBusy: true,
			streamingBehavior: "steer",
			preflightResult: (accepted, didQueue) => {
				expect(accepted).toBe(true);
				queued = didQueue;
			},
		});

		expect(queued).toBe(true);
		expect(session.getSessionActionSnapshot().steering).toContain("message during scheduled compaction");
		releaseCompaction({ summary: "compacted", firstKeptEntryId: "entry-1", tokensBefore: 1 });
		await compaction;
		const deliveryCount = () =>
			session.messages.filter(
				(message) =>
					message.role === "user" &&
					typeof message.content !== "string" &&
					message.content.some(
						(part) => part.type === "text" && part.text === "message during scheduled compaction",
					),
			).length;

		await vi.waitFor(() => expect(deliveryCount()).toBe(1));
		expect(deliveryCount()).toBe(1);
	});

	it("delivers internal prompts without extension input interception", async () => {
		const model = getModel("anthropic", "claude-sonnet-4-5")!;
		let inputCalls = 0;
		let receivedUserText: string | undefined;
		const agent = new Agent({
			getApiKey: () => "test-key",
			initialState: {
				model,
				systemPrompt: "Test",
				tools: [],
			},
			streamFn: (_model, context) => {
				const stream = new MockAssistantStream();
				queueMicrotask(() => {
					const userMessages = context.messages.filter((message) => message.role === "user");
					const user = userMessages.at(-1);
					if (user && typeof user.content !== "string") {
						receivedUserText = user.content
							.filter(
								(part): part is TextContent =>
									typeof part === "object" && part !== null && part.type === "text",
							)
							.map((part) => part.text)
							.join("\n");
					}
					stream.push({ type: "start", partial: createAssistantMessage("") });
					stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Delivered") });
				});
				return stream;
			},
		});

		const sessionManager = SessionManager.inMemory();
		const settingsManager = SettingsManager.create(tempDir, tempDir);
		const authStorage = AuthStorage.create(join(tempDir, "auth.json"));
		const modelRegistry = ModelRegistry.create(authStorage, tempDir);
		authStorage.setRuntimeApiKey("anthropic", "test-key");
		const extensionsResult = await createTestExtensionsResult([
			(pi) => {
				pi.on("input", async () => {
					inputCalls++;
					return { action: "handled" };
				});
			},
		]);

		session = new AgentSession({
			agent,
			sessionManager,
			settingsManager,
			cwd: tempDir,
			modelRegistry,
			resourceLoader: createTestResourceLoader({ extensionsResult }),
		});

		await session.prompt("host gate follow-up", { internalPrompt: true });
		await session.agent.waitForIdle();

		expect(inputCalls).toBe(0);
		expect(receivedUserText).toBe("host gate follow-up");
	});

	it("should allow prompt() after previous completes", async () => {
		const model = getModel("anthropic", "claude-sonnet-4-5")!;
		const agent = new Agent({
			getApiKey: () => "test-key",
			initialState: {
				model,
				systemPrompt: "Test",
				tools: [],
			},
			streamFn: () => {
				const stream = new MockAssistantStream();
				queueMicrotask(() => {
					stream.push({ type: "start", partial: createAssistantMessage("") });
					stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Done") });
				});
				return stream;
			},
		});

		const sessionManager = SessionManager.inMemory();
		const settingsManager = SettingsManager.create(tempDir, tempDir);
		const authStorage = AuthStorage.create(join(tempDir, "auth.json"));
		const modelRegistry = ModelRegistry.create(authStorage, tempDir);
		authStorage.setRuntimeApiKey("anthropic", "test-key");

		session = new AgentSession({
			agent,
			sessionManager,
			settingsManager,
			cwd: tempDir,
			modelRegistry,
			resourceLoader: createTestResourceLoader(),
		});

		await session.prompt("First message");
		expect(session.isStreaming).toBe(false);
		await expect(session.prompt("Second message")).resolves.not.toThrow();
	});

	it("should wait for queued agent events before emitting tool_call", async () => {
		const model = getModel("anthropic", "claude-sonnet-4-5")!;
		const tool = {
			name: "dummy",
			description: "Dummy tool",
			label: "dummy",
			parameters: Type.Object({ q: Type.String() }),
			execute: async (_toolCallId: string, params: unknown) => {
				const q =
					typeof params === "object" && params !== null && "q" in params
						? String((params as { q: unknown }).q)
						: "";
				return {
					content: [{ type: "text" as const, text: `result:${q}` }],
					details: {},
				};
			},
		};

		const agent = new Agent({
			getApiKey: () => "test-key",
			initialState: {
				model,
				systemPrompt: "Test",
				tools: [tool],
			},
			streamFn: async (_model, context) => {
				const stream = new MockAssistantStream();
				queueMicrotask(() => {
					const toolResultCount = context.messages.filter((message) => message.role === "toolResult").length;
					if (toolResultCount > 0) {
						const message: AssistantMessage = {
							role: "assistant",
							content: [{ type: "text", text: "done" }],
							api: "anthropic-messages",
							provider: "anthropic",
							model: "mock",
							usage: {
								input: 1,
								output: 1,
								cacheRead: 0,
								cacheWrite: 0,
								totalTokens: 2,
								cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
							},
							stopReason: "stop",
							timestamp: Date.now(),
						};
						stream.push({ type: "start", partial: { ...message, content: [] } });
						stream.push({ type: "done", reason: "stop", message });
						return;
					}

					const message: AssistantMessage = {
						role: "assistant",
						content: [
							{ type: "toolCall", id: "toolu_1", name: "dummy", arguments: { q: "x" } },
							{ type: "toolCall", id: "toolu_2", name: "dummy", arguments: { q: "y" } },
						],
						api: "anthropic-messages",
						provider: "anthropic",
						model: "mock",
						usage: {
							input: 1,
							output: 1,
							cacheRead: 0,
							cacheWrite: 0,
							totalTokens: 2,
							cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
						},
						stopReason: "toolUse",
						timestamp: Date.now(),
					};

					stream.push({ type: "start", partial: { ...message, content: [] } });
					stream.push({ type: "done", reason: "toolUse", message });
				});
				return stream;
			},
		});

		const sessionManager = SessionManager.inMemory();
		const settingsManager = SettingsManager.create(tempDir, tempDir);
		const authStorage = AuthStorage.create(join(tempDir, "auth.json"));
		const modelRegistry = ModelRegistry.create(authStorage, tempDir);
		authStorage.setRuntimeApiKey("anthropic", "test-key");

		session = new AgentSession({
			agent,
			sessionManager,
			settingsManager,
			cwd: tempDir,
			modelRegistry,
			resourceLoader: createTestResourceLoader(),
			baseToolsOverride: { dummy: tool },
		});

		const snapshots: string[][] = [];
		const sessionWithRunner = session as unknown as {
			_extensionRunner?: {
				hasHandlers: (eventType: string) => boolean;
				emit: (event: { type: string; message?: { role?: string } }) => Promise<void>;
				emitMessageEnd: (event: { type: string; message?: { role?: string } }) => Promise<undefined>;
				emitToolCall: (event: { type: string; toolCallId: string }) => Promise<undefined>;
				emitInput: (
					text: string,
					images: unknown,
					source: "interactive" | "rpc" | "extension",
				) => Promise<{ action: "continue" }>;
				emitBeforeAgentStart: (
					prompt: string,
					images: unknown,
					systemPrompt: string,
					systemPromptOptions: BuildSystemPromptOptions,
				) => Promise<undefined>;
				invalidate: (message?: string) => void;
			};
		};
		sessionWithRunner._extensionRunner = {
			hasHandlers: (eventType) => eventType === "tool_call",
			emit: async () => {},
			emitMessageEnd: async () => undefined,
			emitToolCall: async () => {
				snapshots.push(
					sessionManager
						.getEntries()
						.filter((entry) => entry.type === "message")
						.map((entry) => entry.message.role),
				);
				return undefined;
			},
			emitInput: async () => ({ action: "continue" }),
			emitBeforeAgentStart: async () => undefined,
			invalidate: () => {},
		};

		await session.prompt("hi");
		await session.agent.waitForIdle();

		expect(snapshots).toEqual([
			["user", "assistant"],
			["user", "assistant"],
		]);
	});

	it("should persist message_end events in order with slow extension handlers", async () => {
		const model = getModel("anthropic", "claude-sonnet-4-5")!;
		const tool = {
			name: "dummy",
			description: "Dummy tool",
			label: "dummy",
			parameters: Type.Object({ q: Type.String() }),
			execute: async (_toolCallId: string, params: unknown) => {
				const q =
					typeof params === "object" && params !== null && "q" in params
						? String((params as { q: unknown }).q)
						: "";
				return {
					content: [{ type: "text" as const, text: `result:${q}` }],
					details: {},
				};
			},
		};

		const agent = new Agent({
			getApiKey: () => "test-key",
			initialState: {
				model,
				systemPrompt: "Test",
				tools: [tool],
			},
			streamFn: async (_model, context) => {
				const stream = new MockAssistantStream();
				queueMicrotask(() => {
					const hasToolResult = context.messages.some((message) => message.role === "toolResult");

					if (hasToolResult) {
						const message: AssistantMessage = {
							role: "assistant",
							content: [{ type: "text", text: "done" }],
							api: "anthropic-messages",
							provider: "anthropic",
							model: "mock",
							usage: {
								input: 1,
								output: 1,
								cacheRead: 0,
								cacheWrite: 0,
								totalTokens: 2,
								cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
							},
							stopReason: "stop",
							timestamp: Date.now(),
						};
						stream.push({ type: "start", partial: { ...message, content: [] } });
						stream.push({ type: "done", reason: "stop", message });
						return;
					}

					const message: AssistantMessage = {
						role: "assistant",
						content: [
							{ type: "text", text: "calling tool" },
							{ type: "toolCall", id: "toolu_1", name: "dummy", arguments: { q: "x" } },
						],
						api: "anthropic-messages",
						provider: "anthropic",
						model: "mock",
						usage: {
							input: 1,
							output: 1,
							cacheRead: 0,
							cacheWrite: 0,
							totalTokens: 2,
							cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
						},
						stopReason: "toolUse",
						timestamp: Date.now(),
					};

					stream.push({ type: "start", partial: { ...message, content: [] } });
					stream.push({ type: "done", reason: "toolUse", message });
				});
				return stream;
			},
		});

		const sessionManager = SessionManager.inMemory();
		const settingsManager = SettingsManager.create(tempDir, tempDir);
		const authStorage = AuthStorage.create(join(tempDir, "auth.json"));
		const modelRegistry = ModelRegistry.create(authStorage, tempDir);
		authStorage.setRuntimeApiKey("anthropic", "test-key");

		session = new AgentSession({
			agent,
			sessionManager,
			settingsManager,
			cwd: tempDir,
			modelRegistry,
			resourceLoader: createTestResourceLoader(),
			baseToolsOverride: { dummy: tool },
		});

		const sessionWithRunner = session as unknown as {
			_extensionRunner?: {
				hasHandlers: (eventType: string) => boolean;
				emit: (event: { type: string; message?: { role?: string } }) => Promise<void>;
				emitMessageEnd: (event: { type: string; message?: { role?: string } }) => Promise<undefined>;
				emitInput: (
					text: string,
					images: unknown,
					source: "interactive" | "rpc" | "extension",
				) => Promise<{ action: "continue" }>;
				emitBeforeAgentStart: (
					prompt: string,
					images: unknown,
					systemPrompt: string,
					systemPromptOptions: BuildSystemPromptOptions,
				) => Promise<undefined>;
				invalidate: (message?: string) => void;
			};
		};
		sessionWithRunner._extensionRunner = {
			hasHandlers: () => false,
			emit: async () => {},
			emitMessageEnd: async (event) => {
				if (event.type === "message_end" && event.message?.role === "assistant") {
					await new Promise((resolve) => setTimeout(resolve, 40));
				}
				return undefined;
			},
			emitInput: async () => ({ action: "continue" }),
			emitBeforeAgentStart: async () => undefined,
			invalidate: () => {},
		};

		await session.prompt("hi");
		await session.agent.waitForIdle();
		await new Promise((resolve) => setTimeout(resolve, 100));

		const messageEntries = sessionManager.getEntries().filter((entry) => entry.type === "message");
		expect(messageEntries.map((entry) => entry.message.role)).toEqual([
			"user",
			"assistant",
			"toolResult",
			"assistant",
		]);
	});
});