storkit/frontend/src/components/Chat.test.tsx

import {
	act,
	fireEvent,
	render,
	screen,
	waitFor,
} from "@testing-library/react";

import { beforeEach, describe, expect, it, vi } from "vitest";
import { api } from "../api/client";
import type { Message } from "../types";
import { Chat } from "./Chat";

// Module-level store for the WebSocket handlers captured during connect().
// Tests in the "message rendering" suite use this to simulate incoming messages.
type WsHandlers = {
	onToken: (content: string) => void;
	onUpdate: (history: Message[]) => void;
	onSessionId: (sessionId: string) => void;
	onError: (message: string) => void;
	onActivity: (toolName: string) => void;
	onReconciliationProgress: (
		storyId: string,
		status: string,
		message: string,
	) => void;
};
let capturedWsHandlers: WsHandlers | null = null;

vi.mock("../api/client", () => {
	const api = {
		getOllamaModels: vi.fn(),
		getAnthropicApiKeyExists: vi.fn(),
		getAnthropicModels: vi.fn(),
		getModelPreference: vi.fn(),
		setModelPreference: vi.fn(),
		cancelChat: vi.fn(),
		setAnthropicApiKey: vi.fn(),
	};
	class ChatWebSocket {
		connect(handlers: WsHandlers) {
			capturedWsHandlers = handlers;
		}
		close() {}
		sendChat() {}
		cancel() {}
	}
	return { api, ChatWebSocket };
});

const mockedApi = {
	getOllamaModels: vi.mocked(api.getOllamaModels),
	getAnthropicApiKeyExists: vi.mocked(api.getAnthropicApiKeyExists),
	getAnthropicModels: vi.mocked(api.getAnthropicModels),
	getModelPreference: vi.mocked(api.getModelPreference),
	setModelPreference: vi.mocked(api.setModelPreference),
	cancelChat: vi.mocked(api.cancelChat),
	setAnthropicApiKey: vi.mocked(api.setAnthropicApiKey),
};

function setupMocks() {
	mockedApi.getOllamaModels.mockResolvedValue(["llama3.1"]);
	mockedApi.getAnthropicApiKeyExists.mockResolvedValue(true);
	mockedApi.getAnthropicModels.mockResolvedValue([]);
	mockedApi.getModelPreference.mockResolvedValue(null);
	mockedApi.setModelPreference.mockResolvedValue(true);
	mockedApi.cancelChat.mockResolvedValue(true);
	mockedApi.setAnthropicApiKey.mockResolvedValue(true);
}

describe("Chat message rendering — unified tool call UI", () => {
	beforeEach(() => {
		capturedWsHandlers = null;
		setupMocks();
	});

	it("renders tool call badge for assistant message with tool_calls (AC3)", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => expect(capturedWsHandlers).not.toBeNull());

		const messages: Message[] = [
			{ role: "user", content: "Read src/main.rs" },
			{
				role: "assistant",
				content: "I'll read that file.",
				tool_calls: [
					{
						id: "toolu_abc",
						type: "function",
						function: {
							name: "Read",
							arguments: '{"file_path":"src/main.rs"}',
						},
					},
				],
			},
		];

		act(() => {
			capturedWsHandlers?.onUpdate(messages);
		});

		expect(await screen.findByText("I'll read that file.")).toBeInTheDocument();
		expect(await screen.findByText("Read(src/main.rs)")).toBeInTheDocument();
	});

	it("renders collapsible tool output for tool role messages (AC3)", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => expect(capturedWsHandlers).not.toBeNull());

		const messages: Message[] = [
			{ role: "user", content: "Check the file" },
			{
				role: "assistant",
				content: "",
				tool_calls: [
					{
						id: "toolu_1",
						type: "function",
						function: { name: "Read", arguments: '{"file_path":"foo.rs"}' },
					},
				],
			},
			{
				role: "tool",
				content: 'fn main() { println!("hello"); }',
				tool_call_id: "toolu_1",
			},
			{ role: "assistant", content: "The file contains a main function." },
		];

		act(() => {
			capturedWsHandlers?.onUpdate(messages);
		});

		expect(await screen.findByText(/Tool Output/)).toBeInTheDocument();
		expect(
			await screen.findByText("The file contains a main function."),
		).toBeInTheDocument();
	});

	it("renders plain assistant message without tool call badges (AC5)", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => expect(capturedWsHandlers).not.toBeNull());

		const messages: Message[] = [
			{ role: "user", content: "Hello" },
			{ role: "assistant", content: "Hi there! How can I help?" },
		];

		act(() => {
			capturedWsHandlers?.onUpdate(messages);
		});

		expect(
			await screen.findByText("Hi there! How can I help?"),
		).toBeInTheDocument();
		expect(screen.queryByText(/Tool Output/)).toBeNull();
	});

	it("renders multiple tool calls in a single assistant turn (AC3)", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => expect(capturedWsHandlers).not.toBeNull());

		const messages: Message[] = [
			{ role: "user", content: "Do some work" },
			{
				role: "assistant",
				content: "I'll do multiple things.",
				tool_calls: [
					{
						id: "id1",
						type: "function",
						function: { name: "Bash", arguments: '{"command":"cargo test"}' },
					},
					{
						id: "id2",
						type: "function",
						function: { name: "Read", arguments: '{"file_path":"Cargo.toml"}' },
					},
				],
			},
		];

		act(() => {
			capturedWsHandlers?.onUpdate(messages);
		});

		expect(await screen.findByText("Bash(cargo test)")).toBeInTheDocument();
		expect(await screen.findByText("Read(Cargo.toml)")).toBeInTheDocument();
	});

	it("does not fetch Anthropic models when no API key exists", async () => {
		mockedApi.getAnthropicApiKeyExists.mockResolvedValue(false);
		mockedApi.getAnthropicModels.mockClear();

		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => {
			expect(mockedApi.getAnthropicApiKeyExists).toHaveBeenCalled();
		});

		expect(mockedApi.getAnthropicModels).not.toHaveBeenCalled();
	});
});

describe("Chat two-column layout", () => {
	beforeEach(() => {
		capturedWsHandlers = null;
		setupMocks();
	});

	it("renders left and right column containers (AC1, AC2)", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		expect(await screen.findByTestId("chat-content-area")).toBeInTheDocument();
		expect(await screen.findByTestId("chat-left-column")).toBeInTheDocument();
		expect(await screen.findByTestId("chat-right-column")).toBeInTheDocument();
	});

	it("renders chat input inside the left column (AC2, AC5)", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		const leftColumn = await screen.findByTestId("chat-left-column");
		const input = screen.getByPlaceholderText("Send a message...");
		expect(leftColumn).toContainElement(input);
	});

	it("renders panels inside the right column (AC2)", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		const rightColumn = await screen.findByTestId("chat-right-column");
		const agentsPanel = await screen.findByText("Agents");
		expect(rightColumn).toContainElement(agentsPanel);
	});

	it("uses row flex-direction on wide screens (AC3)", async () => {
		Object.defineProperty(window, "innerWidth", {
			writable: true,
			configurable: true,
			value: 1200,
		});
		window.dispatchEvent(new Event("resize"));

		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		const contentArea = await screen.findByTestId("chat-content-area");
		expect(contentArea).toHaveStyle({ flexDirection: "row" });
	});

	it("uses column flex-direction on narrow screens (AC4)", async () => {
		Object.defineProperty(window, "innerWidth", {
			writable: true,
			configurable: true,
			value: 600,
		});
		window.dispatchEvent(new Event("resize"));

		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		const contentArea = await screen.findByTestId("chat-content-area");
		expect(contentArea).toHaveStyle({ flexDirection: "column" });

		// Restore wide width for subsequent tests
		Object.defineProperty(window, "innerWidth", {
			writable: true,
			configurable: true,
			value: 1024,
		});
	});
});

describe("Chat input Shift+Enter behavior", () => {
	beforeEach(() => {
		capturedWsHandlers = null;
		setupMocks();
	});

	it("renders a textarea element for the chat input (AC3)", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);
		const input = screen.getByPlaceholderText("Send a message...");
		expect(input.tagName.toLowerCase()).toBe("textarea");
	});

	it("sends message on Enter key press without Shift (AC2)", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);
		const input = screen.getByPlaceholderText("Send a message...");

		await act(async () => {
			fireEvent.change(input, { target: { value: "Hello" } });
		});
		await act(async () => {
			fireEvent.keyDown(input, { key: "Enter", shiftKey: false });
		});

		await waitFor(() => {
			expect((input as HTMLTextAreaElement).value).toBe("");
		});
	});

	it("does not send message on Shift+Enter (AC1)", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);
		const input = screen.getByPlaceholderText("Send a message...");

		await act(async () => {
			fireEvent.change(input, { target: { value: "Hello" } });
		});
		await act(async () => {
			fireEvent.keyDown(input, { key: "Enter", shiftKey: true });
		});

		expect((input as HTMLTextAreaElement).value).toBe("Hello");
	});
});

describe("Chat reconciliation banner", () => {
	beforeEach(() => {
		capturedWsHandlers = null;
		setupMocks();
	});

	it("shows banner when a non-done reconciliation event is received", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => expect(capturedWsHandlers).not.toBeNull());

		act(() => {
			capturedWsHandlers?.onReconciliationProgress(
				"42_story_test",
				"checking",
				"Checking for committed work in 2_current/",
			);
		});

		expect(
			await screen.findByTestId("reconciliation-banner"),
		).toBeInTheDocument();
		expect(
			await screen.findByText("Reconciling startup state..."),
		).toBeInTheDocument();
	});

	it("shows event message in the banner", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => expect(capturedWsHandlers).not.toBeNull());

		act(() => {
			capturedWsHandlers?.onReconciliationProgress(
				"42_story_test",
				"gates_running",
				"Running acceptance gates…",
			);
		});

		expect(
			await screen.findByText(/Running acceptance gates/),
		).toBeInTheDocument();
	});

	it("dismisses banner when done event is received", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => expect(capturedWsHandlers).not.toBeNull());

		act(() => {
			capturedWsHandlers?.onReconciliationProgress(
				"42_story_test",
				"checking",
				"Checking for committed work",
			);
		});

		expect(
			await screen.findByTestId("reconciliation-banner"),
		).toBeInTheDocument();

		act(() => {
			capturedWsHandlers?.onReconciliationProgress(
				"",
				"done",
				"Startup reconciliation complete.",
			);
		});

		await waitFor(() => {
			expect(
				screen.queryByTestId("reconciliation-banner"),
			).not.toBeInTheDocument();
		});
	});
});

describe("Chat activity status indicator (Bug 140)", () => {
	beforeEach(() => {
		capturedWsHandlers = null;
		setupMocks();
	});

	it("shows activity label when tool activity fires during streaming content", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => expect(capturedWsHandlers).not.toBeNull());

		// Simulate sending a message to set loading=true
		const input = screen.getByPlaceholderText("Send a message...");
		await act(async () => {
			fireEvent.change(input, { target: { value: "Read my file" } });
		});
		await act(async () => {
			fireEvent.keyDown(input, { key: "Enter", shiftKey: false });
		});

		// Simulate tokens arriving (streamingContent becomes non-empty)
		act(() => {
			capturedWsHandlers?.onToken("I'll read that file for you.");
		});

		// Now simulate a tool activity event while streamingContent is non-empty
		act(() => {
			capturedWsHandlers?.onActivity("read_file");
		});

		// The activity indicator should be visible with the tool activity label
		const indicator = await screen.findByTestId("activity-indicator");
		expect(indicator).toBeInTheDocument();
		expect(indicator).toHaveTextContent("Reading file...");
	});

	it("shows Thinking... fallback when loading with no streaming and no activity", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => expect(capturedWsHandlers).not.toBeNull());

		// Simulate sending a message to set loading=true
		const input = screen.getByPlaceholderText("Send a message...");
		await act(async () => {
			fireEvent.change(input, { target: { value: "Hello" } });
		});
		await act(async () => {
			fireEvent.keyDown(input, { key: "Enter", shiftKey: false });
		});

		// No tokens, no activity — should show "Thinking..."
		const indicator = await screen.findByTestId("activity-indicator");
		expect(indicator).toBeInTheDocument();
		expect(indicator).toHaveTextContent("Thinking...");
	});

	it("hides Thinking... when streaming content is present but no tool activity", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => expect(capturedWsHandlers).not.toBeNull());

		// Simulate sending a message to set loading=true
		const input = screen.getByPlaceholderText("Send a message...");
		await act(async () => {
			fireEvent.change(input, { target: { value: "Hello" } });
		});
		await act(async () => {
			fireEvent.keyDown(input, { key: "Enter", shiftKey: false });
		});

		// Tokens arrive — streamingContent is non-empty, no activity
		act(() => {
			capturedWsHandlers?.onToken("Here is my response...");
		});

		// The activity indicator should NOT be visible (just streaming bubble)
		expect(screen.queryByTestId("activity-indicator")).not.toBeInTheDocument();
	});

	it("shows activity label for Claude Code tool names (Read, Bash, etc.)", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => expect(capturedWsHandlers).not.toBeNull());

		// Simulate sending a message to set loading=true
		const input = screen.getByPlaceholderText("Send a message...");
		await act(async () => {
			fireEvent.change(input, { target: { value: "Read my file" } });
		});
		await act(async () => {
			fireEvent.keyDown(input, { key: "Enter", shiftKey: false });
		});

		// Simulate tokens arriving
		act(() => {
			capturedWsHandlers?.onToken("Let me read that.");
		});

		// Claude Code sends tool name "Read" (not "read_file")
		act(() => {
			capturedWsHandlers?.onActivity("Read");
		});

		const indicator = await screen.findByTestId("activity-indicator");
		expect(indicator).toBeInTheDocument();
		expect(indicator).toHaveTextContent("Reading file...");
	});

	it("shows activity label for Claude Code Bash tool", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => expect(capturedWsHandlers).not.toBeNull());

		const input = screen.getByPlaceholderText("Send a message...");
		await act(async () => {
			fireEvent.change(input, { target: { value: "Run the tests" } });
		});
		await act(async () => {
			fireEvent.keyDown(input, { key: "Enter", shiftKey: false });
		});

		act(() => {
			capturedWsHandlers?.onToken("Running tests now.");
		});

		act(() => {
			capturedWsHandlers?.onActivity("Bash");
		});

		const indicator = await screen.findByTestId("activity-indicator");
		expect(indicator).toBeInTheDocument();
		expect(indicator).toHaveTextContent("Executing command...");
	});

	it("shows generic label for unknown tool names", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => expect(capturedWsHandlers).not.toBeNull());

		const input = screen.getByPlaceholderText("Send a message...");
		await act(async () => {
			fireEvent.change(input, { target: { value: "Do something" } });
		});
		await act(async () => {
			fireEvent.keyDown(input, { key: "Enter", shiftKey: false });
		});

		act(() => {
			capturedWsHandlers?.onToken("Working on it.");
		});

		act(() => {
			capturedWsHandlers?.onActivity("SomeCustomTool");
		});

		const indicator = await screen.findByTestId("activity-indicator");
		expect(indicator).toBeInTheDocument();
		expect(indicator).toHaveTextContent("Using SomeCustomTool...");
	});
});

describe("Chat message queue (Story 155)", () => {
	beforeEach(() => {
		capturedWsHandlers = null;
		setupMocks();
	});

	it("shows queued message indicator when submitting while loading (AC1, AC2)", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => expect(capturedWsHandlers).not.toBeNull());

		// Send first message to put the chat in loading state
		const input = screen.getByPlaceholderText("Send a message...");
		await act(async () => {
			fireEvent.change(input, { target: { value: "First message" } });
		});
		await act(async () => {
			fireEvent.keyDown(input, { key: "Enter", shiftKey: false });
		});

		// Now type and submit a second message while loading is true
		await act(async () => {
			fireEvent.change(input, { target: { value: "Queued message" } });
		});
		await act(async () => {
			fireEvent.keyDown(input, { key: "Enter", shiftKey: false });
		});

		// The queued message indicator should appear
		const indicator = await screen.findByTestId("queued-message-indicator");
		expect(indicator).toBeInTheDocument();
		expect(indicator).toHaveTextContent("Queued");
		expect(indicator).toHaveTextContent("Queued message");

		// Input should be cleared after queuing
		expect((input as HTMLTextAreaElement).value).toBe("");
	});

	it("auto-sends queued message when agent response completes (AC4)", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => expect(capturedWsHandlers).not.toBeNull());

		const input = screen.getByPlaceholderText("Send a message...");

		// Send first message
		await act(async () => {
			fireEvent.change(input, { target: { value: "First" } });
		});
		await act(async () => {
			fireEvent.keyDown(input, { key: "Enter", shiftKey: false });
		});

		// Queue a second message while loading
		await act(async () => {
			fireEvent.change(input, { target: { value: "Auto-send this" } });
		});
		await act(async () => {
			fireEvent.keyDown(input, { key: "Enter", shiftKey: false });
		});

		// Verify it's queued
		expect(
			await screen.findByTestId("queued-message-indicator"),
		).toBeInTheDocument();

		// Simulate agent response completing (loading → false)
		act(() => {
			capturedWsHandlers?.onUpdate([
				{ role: "user", content: "First" },
				{ role: "assistant", content: "Done." },
			]);
		});

		// The queued indicator should disappear (message was sent)
		await waitFor(() => {
			expect(
				screen.queryByTestId("queued-message-indicator"),
			).not.toBeInTheDocument();
		});
	});

	it("cancel button discards the queued message (AC3, AC6)", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => expect(capturedWsHandlers).not.toBeNull());

		const input = screen.getByPlaceholderText("Send a message...");

		// Send first message to start loading
		await act(async () => {
			fireEvent.change(input, { target: { value: "First" } });
		});
		await act(async () => {
			fireEvent.keyDown(input, { key: "Enter", shiftKey: false });
		});

		// Queue a second message
		await act(async () => {
			fireEvent.change(input, { target: { value: "Discard me" } });
		});
		await act(async () => {
			fireEvent.keyDown(input, { key: "Enter", shiftKey: false });
		});

		const indicator = await screen.findByTestId("queued-message-indicator");
		expect(indicator).toBeInTheDocument();

		// Click the ✕ cancel button
		const cancelBtn = screen.getByTitle("Cancel queued message");
		await act(async () => {
			fireEvent.click(cancelBtn);
		});

		// Indicator should be gone
		expect(
			screen.queryByTestId("queued-message-indicator"),
		).not.toBeInTheDocument();
	});

	it("edit button puts queued message back into input (AC3)", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => expect(capturedWsHandlers).not.toBeNull());

		const input = screen.getByPlaceholderText("Send a message...");

		// Send first message to start loading
		await act(async () => {
			fireEvent.change(input, { target: { value: "First" } });
		});
		await act(async () => {
			fireEvent.keyDown(input, { key: "Enter", shiftKey: false });
		});

		// Queue a second message
		await act(async () => {
			fireEvent.change(input, { target: { value: "Edit me back" } });
		});
		await act(async () => {
			fireEvent.keyDown(input, { key: "Enter", shiftKey: false });
		});

		await screen.findByTestId("queued-message-indicator");

		// Click the Edit button
		const editBtn = screen.getByTitle("Edit queued message");
		await act(async () => {
			fireEvent.click(editBtn);
		});

		// Indicator should be gone and message back in input
		expect(
			screen.queryByTestId("queued-message-indicator"),
		).not.toBeInTheDocument();
		expect((input as HTMLTextAreaElement).value).toBe("Edit me back");
	});

	it("subsequent submissions replace the queued message (AC5)", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => expect(capturedWsHandlers).not.toBeNull());

		const input = screen.getByPlaceholderText("Send a message...");

		// Send first message to start loading
		await act(async () => {
			fireEvent.change(input, { target: { value: "First" } });
		});
		await act(async () => {
			fireEvent.keyDown(input, { key: "Enter", shiftKey: false });
		});

		// Queue first replacement
		await act(async () => {
			fireEvent.change(input, { target: { value: "Original queue" } });
		});
		await act(async () => {
			fireEvent.keyDown(input, { key: "Enter", shiftKey: false });
		});

		await screen.findByTestId("queued-message-indicator");

		// Queue second replacement — should overwrite the first
		await act(async () => {
			fireEvent.change(input, { target: { value: "Replaced queue" } });
		});
		await act(async () => {
			fireEvent.keyDown(input, { key: "Enter", shiftKey: false });
		});

		const indicator = await screen.findByTestId("queued-message-indicator");
		expect(indicator).toHaveTextContent("Replaced queue");
		expect(indicator).not.toHaveTextContent("Original queue");
	});

	it("does not auto-send queued message when generation is cancelled (AC6)", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => expect(capturedWsHandlers).not.toBeNull());

		const input = screen.getByPlaceholderText("Send a message...");

		// Send first message to start loading
		await act(async () => {
			fireEvent.change(input, { target: { value: "First" } });
		});
		await act(async () => {
			fireEvent.keyDown(input, { key: "Enter", shiftKey: false });
		});

		// Queue a second message
		await act(async () => {
			fireEvent.change(input, { target: { value: "Should not send" } });
		});
		await act(async () => {
			fireEvent.keyDown(input, { key: "Enter", shiftKey: false });
		});

		await screen.findByTestId("queued-message-indicator");

		// Click the stop button (■) — but input is empty so button is stop
		// Actually simulate cancel by clicking the stop button (which requires empty input)
		// We need to use the send button when input is empty (stop mode)
		// Simulate cancel via the cancelGeneration path: the button when loading && !input
		// At this point input is empty (was cleared after queuing)
		const stopButton = screen.getByRole("button", { name: "■" });
		await act(async () => {
			fireEvent.click(stopButton);
		});

		// Queued indicator should be gone (cancelled)
		await waitFor(() => {
			expect(
				screen.queryByTestId("queued-message-indicator"),
			).not.toBeInTheDocument();
		});
	});
});