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

import { act, 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;
};
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,
		});
	});
});