import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { act, fireEvent, render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import type { AgentEntry, AgentTask } from "@/bindings"; const ipc = vi.hoisted(() => ({ rollback: vi.fn(), prompt: vi.fn(), cancel: vi.fn() })); const entries: AgentEntry[] = [ { type: "user", text: "Original message", attachments: [], sentMs: 1, doneMs: 2, steered: false }, { type: "agent", text: "Original answer" }, ]; vi.mock("@/platform/ipc", async (original) => ({ ...(await original()), commands: { agentEntries: async () => ({ status: "ok", data: entries }), agentChanges: async () => ({ status: "ok", data: [] }), vcsBranchInfo: async () => ({ status: "ok", data: { head: "main" } }), agentRollback: ipc.rollback, agentPrompt: ipc.prompt, agentCancel: ipc.cancel, }, })); vi.mock("@tauri-apps/api/webview", () => ({ getCurrentWebview: () => ({ onDragDropEvent: () => Promise.resolve(() => {}) }), })); import { SessionEditor } from "./SessionEditor"; import { updateComposer, useAgents } from "./store"; const task: AgentTask = { id: "editing", title: "Test task", provider: "codex", status: "ready", isolation: "local", model: null, cwd: "/project", branch: null, base: null, baseCommit: null, error: null, createdMs: 0, updatedMs: 0, config: [], commands: [], usage: null, queue: [], capabilities: { fork: true, rollback: true, steer: true, images: true }, }; beforeEach(() => { ipc.rollback.mockReset().mockResolvedValue({ status: "ok", data: null }); ipc.prompt.mockReset().mockResolvedValue({ status: "ok", data: null }); ipc.cancel.mockReset().mockResolvedValue({ status: "ok", data: null }); useAgents.setState({ tasks: new Map([[task.id, task]]), entries: new Map(), composers: new Map(), focus: null }); }); async function edit() { render(); const user = userEvent.setup(); await user.click(await screen.findByRole("button", { name: "Edit" })); const box = screen.getAllByRole("textbox", { name: "Message" })[0]!; fireEvent.change(box, { target: { value: "Edited message" } }); return { user, box, send: screen.getAllByRole("button", { name: "Send" })[0]! }; } describe("a sent message with images", () => { afterEach(() => { entries[0] = { type: "user", text: "Original message", attachments: [], sentMs: 1, doneMs: 2, steered: false }; }); it("shows the images and lists only the files as chips", async () => { entries[0] = { type: "user", text: "Look", attachments: ["a.ts", "shot.png"], images: [{ mimeType: "image/png", data: "AAA", name: "shot.png" }], sentMs: 1, doneMs: 2, steered: false, }; render(); const img = await screen.findByRole("img", { name: "shot.png" }); expect(img).toHaveAttribute("src", "data:image/png;base64,AAA"); expect(screen.getByText("a.ts")).toBeInTheDocument(); expect(screen.queryByText("shot.png")).not.toBeInTheDocument(); entries[0] = { type: "user", text: "Original message", attachments: [], sentMs: 1, doneMs: 2, steered: false }; }); it("opens an image in a dialog and keeps the images when editing, minus the removed ones", async () => { entries[0] = { type: "user", text: "Look", attachments: ["a.png", "b.png"], images: [ { mimeType: "image/png", data: "AAA", name: "a.png" }, { mimeType: "image/png", data: "BBB", name: "b.png" }, ], sentMs: 1, doneMs: 2, steered: false, }; render(); const user = userEvent.setup(); await user.click(await screen.findByRole("button", { name: "a.png" })); expect(await screen.findByRole("dialog")).toBeInTheDocument(); expect(screen.getAllByRole("img", { name: "a.png" })).toHaveLength(2); await user.click(screen.getAllByRole("button", { name: "Close" })[0]!); await user.click(screen.getByRole("button", { name: "Edit" })); expect(screen.getAllByRole("button", { name: "Remove Image" })).toHaveLength(2); await user.click(screen.getAllByRole("button", { name: "Remove Image" })[0]!); expect(screen.getAllByRole("button", { name: "Remove Image" })).toHaveLength(1); expect(screen.queryByRole("img", { name: "a.png" })).not.toBeInTheDocument(); await user.click(screen.getAllByRole("button", { name: "Send" })[0]!); await waitFor(() => expect(ipc.prompt).toHaveBeenCalledTimes(1)); expect(ipc.prompt.mock.calls[0]![1].images).toEqual([{ mimeType: "image/png", data: "BBB", name: "b.png" }]); entries[0] = { type: "user", text: "Original message", attachments: [], sentMs: 1, doneMs: 2, steered: false }; }); }); describe("editing a sent agent message", () => { it("keeps Enter, Shift+Enter and IME confirmation for editing", async () => { const { box } = await edit(); fireEvent.keyDown(box, { key: "Enter" }); fireEvent.keyDown(box, { key: "Enter", shiftKey: true }); fireEvent.keyDown(box, { key: "Enter", metaKey: true, isComposing: true }); fireEvent.keyDown(box, { key: "Enter", metaKey: true, keyCode: 229 }); expect(ipc.rollback).not.toHaveBeenCalled(); expect(box).toHaveValue("Edited message"); }); it("keeps the draft and shows a rollback failure beside it", async () => { ipc.rollback.mockResolvedValue({ status: "error", error: { kind: "internal", message: "Rollback failed" } }); const { user, box, send } = await edit(); await user.click(send); expect(await screen.findByRole("alert")).toHaveTextContent("Rollback failed"); expect(box).toHaveValue("Edited message"); expect(box).toBeInTheDocument(); expect(send).toBeEnabled(); expect(ipc.prompt).not.toHaveBeenCalled(); ipc.rollback.mockResolvedValue({ status: "ok", data: null }); await user.click(send); await waitFor(() => expect(ipc.prompt).toHaveBeenCalledTimes(1)); expect(screen.queryByRole("alert")).not.toBeInTheDocument(); }); it("rolls back before sending the edited text when Send is clicked", async () => { let finish!: (value: unknown) => void; ipc.rollback.mockReturnValue( new Promise((resolve) => { finish = resolve; }), ); const { user, box, send } = await edit(); await user.click(send); expect(ipc.rollback).toHaveBeenCalledWith(task.id, 0); expect(ipc.prompt).not.toHaveBeenCalled(); expect(send).toBeDisabled(); fireEvent.keyDown(box, { key: "Enter", metaKey: true }); expect(ipc.rollback).toHaveBeenCalledTimes(1); await act(async () => finish({ status: "ok", data: null })); await waitFor(() => expect(ipc.prompt).toHaveBeenCalledWith(task.id, { text: "Edited message", files: [], images: [] }, false), ); }); it.each(["", "Another draft"])("recovers the edit after rollback if sending fails (draft: %j)", async (draft) => { ipc.rollback.mockImplementation(async () => { // The backend emits Truncated before replying, removing the inline editor. useAgents.setState({ entries: new Map([[task.id, []]]) }); return { status: "ok", data: null }; }); let finish!: (value: unknown) => void; ipc.prompt.mockReturnValue( new Promise((resolve) => { finish = resolve; }), ); const { user, box, send } = await edit(); act(() => updateComposer(task.id, { text: draft })); await user.click(send); expect(box).not.toBeInTheDocument(); await act(async () => finish({ status: "error", error: { kind: "internal", message: "Send failed" } })); expect(screen.getByRole("textbox", { name: "Message" })).toHaveValue( draft ? `Edited message\n${draft}` : "Edited message", ); expect(screen.getByRole("textbox", { name: "Message" })).toHaveFocus(); ipc.prompt.mockResolvedValue({ status: "ok", data: null }); await user.click(screen.getByRole("button", { name: "Send" })); expect(ipc.prompt).toHaveBeenCalledTimes(2); expect(ipc.rollback).toHaveBeenCalledTimes(1); }); it.each([{ metaKey: true }, { ctrlKey: true }])("still sends with modified Enter (%j)", async (modifier) => { const { box } = await edit(); fireEvent.keyDown(box, { key: "Enter", ...modifier }); await waitFor(() => expect(ipc.prompt).toHaveBeenCalledTimes(1)); }); it("waits for a running turn to stop before rolling it back", async () => { useAgents.setState({ tasks: new Map([[task.id, { ...task, status: "running" }]]) }); const { user, send } = await edit(); await user.click(send); expect(ipc.cancel).toHaveBeenCalledWith(task.id, null); expect(ipc.rollback).not.toHaveBeenCalled(); act(() => useAgents.setState({ tasks: new Map([[task.id, task]]) })); await waitFor(() => expect(ipc.prompt).toHaveBeenCalledTimes(1)); }); it.each(["Cancel", "Escape"])( "%s restores the message and discards edits without touching the conversation", async (method) => { const { user, box } = await edit(); if (method === "Cancel") await user.click(screen.getByRole("button", { name: "Cancel" })); else fireEvent.keyDown(box, { key: "Escape" }); expect(box).not.toBeInTheDocument(); expect(screen.getByText("Original message")).toBeInTheDocument(); expect(screen.getByText("Original answer")).toBeInTheDocument(); expect(ipc.rollback).not.toHaveBeenCalled(); expect(ipc.prompt).not.toHaveBeenCalled(); await user.click(screen.getByRole("button", { name: "Edit" })); expect(screen.getAllByRole("textbox", { name: "Message" })[0]).toHaveValue("Original message"); }, ); });