diff --git a/ui/src/components/transcript/useLiveRunTranscripts.test.tsx b/ui/src/components/transcript/useLiveRunTranscripts.test.tsx new file mode 100644 index 00000000..6a5105f0 --- /dev/null +++ b/ui/src/components/transcript/useLiveRunTranscripts.test.tsx @@ -0,0 +1,118 @@ +// @vitest-environment jsdom + +import { act } from "react"; +import { createRoot } from "react-dom/client"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { useLiveRunTranscripts } from "./useLiveRunTranscripts"; + +const { useQueryMock, logMock } = vi.hoisted(() => ({ + useQueryMock: vi.fn(() => ({ data: { censorUsernameInLogs: false } })), + logMock: vi.fn(async () => ({ runId: "run-1", store: "memory", logRef: "log-1", content: "", nextOffset: 0 })), +})); + +vi.mock("@tanstack/react-query", () => ({ + useQuery: useQueryMock, +})); + +vi.mock("../../api/instanceSettings", () => ({ + instanceSettingsApi: { + getGeneral: vi.fn(), + }, +})); + +vi.mock("../../api/heartbeats", () => ({ + heartbeatsApi: { + log: logMock, + }, +})); + +vi.mock("../../adapters", () => ({ + buildTranscript: (chunks: unknown[]) => chunks, + getUIAdapter: () => null, + onAdapterChange: () => () => {}, +})); + +class FakeWebSocket { + static readonly CONNECTING = 0; + static readonly OPEN = 1; + static readonly CLOSING = 2; + static readonly CLOSED = 3; + static instances: FakeWebSocket[] = []; + + readonly url: string; + readyState = FakeWebSocket.CONNECTING; + onopen: ((event: Event) => void) | null = null; + onmessage: ((event: MessageEvent) => void) | null = null; + onerror: ((event: Event) => void) | null = null; + onclose: ((event: CloseEvent) => void) | null = null; + closeCalls: Array<{ code?: number; reason?: string }> = []; + + constructor(url: string) { + this.url = url; + FakeWebSocket.instances.push(this); + } + + close(code?: number, reason?: string) { + this.closeCalls.push({ code, reason }); + this.readyState = FakeWebSocket.CLOSING; + } + + triggerOpen() { + this.readyState = FakeWebSocket.OPEN; + this.onopen?.(new Event("open")); + } +} + +// eslint-disable-next-line @typescript-eslint/no-explicit-any +(globalThis as any).IS_REACT_ACT_ENVIRONMENT = true; + +describe("useLiveRunTranscripts", () => { + const OriginalWebSocket = globalThis.WebSocket; + + beforeEach(() => { + FakeWebSocket.instances = []; + useQueryMock.mockClear(); + logMock.mockClear(); + globalThis.WebSocket = FakeWebSocket as unknown as typeof WebSocket; + }); + + afterEach(() => { + globalThis.WebSocket = OriginalWebSocket; + }); + + it("waits for a connecting socket to open before closing it during cleanup", async () => { + function Harness() { + useLiveRunTranscripts({ + companyId: "company-1", + runs: [{ id: "run-1", status: "running", adapterType: "codex_local" }], + }); + return null; + } + + const container = document.createElement("div"); + document.body.appendChild(container); + const root = createRoot(container); + + await act(async () => { + root.render(); + await Promise.resolve(); + }); + + expect(FakeWebSocket.instances).toHaveLength(1); + const socket = FakeWebSocket.instances[0]; + expect(socket.closeCalls).toHaveLength(0); + + act(() => { + root.unmount(); + }); + + expect(socket.closeCalls).toHaveLength(0); + + act(() => { + socket.triggerOpen(); + }); + + expect(socket.closeCalls).toEqual([{ code: 1000, reason: "live_run_transcripts_unmount" }]); + container.remove(); + }); +}); diff --git a/ui/src/components/transcript/useLiveRunTranscripts.ts b/ui/src/components/transcript/useLiveRunTranscripts.ts index 868645d6..6cccb64a 100644 --- a/ui/src/components/transcript/useLiveRunTranscripts.ts +++ b/ui/src/components/transcript/useLiveRunTranscripts.ts @@ -281,7 +281,16 @@ export function useLiveRunTranscripts({ socket.onmessage = null; socket.onerror = null; socket.onclose = null; - socket.close(1000, "live_run_transcripts_unmount"); + if (socket.readyState === WebSocket.CONNECTING) { + // Defer the close until the handshake completes so the browser + // does not emit a noisy "closed before the connection is established" + // warning during rapid run teardown. + socket.onopen = () => { + socket?.close(1000, "live_run_transcripts_unmount"); + }; + } else if (socket.readyState === WebSocket.OPEN) { + socket.close(1000, "live_run_transcripts_unmount"); + } } }; }, [activeRunIds, companyId, runById]); diff --git a/ui/src/lib/issue-chat-messages.test.ts b/ui/src/lib/issue-chat-messages.test.ts index f3ae87d7..92585bed 100644 --- a/ui/src/lib/issue-chat-messages.test.ts +++ b/ui/src/lib/issue-chat-messages.test.ts @@ -270,7 +270,7 @@ describe("buildIssueChatMessages", () => { "system:activity:event-1", "user:comment-1", "assistant:comment-2", - "assistant:live-run:run-live-1", + "assistant:run-assistant:run-live-1", ]); const liveRunMessage = messages.at(-1); @@ -316,7 +316,7 @@ describe("buildIssueChatMessages", () => { expect(messages).toHaveLength(1); expect(messages[0]).toMatchObject({ - id: "historical-run:run-history-1", + id: "run-assistant:run-history-1", role: "assistant", status: { type: "complete", reason: "stop" }, metadata: { @@ -333,6 +333,64 @@ describe("buildIssueChatMessages", () => { ]); }); + it("keeps the same assistant message id when a live run becomes a cancelled historical run", () => { + const liveMessages = buildIssueChatMessages({ + comments: [], + timelineEvents: [], + linkedRuns: [], + liveRuns: [ + { + id: "run-1", + status: "running", + invocationSource: "manual", + triggerDetail: null, + startedAt: "2026-04-06T12:01:00.000Z", + finishedAt: null, + createdAt: "2026-04-06T12:01:00.000Z", + agentId: "agent-1", + agentName: "CodexCoder", + adapterType: "codex_local", + }, + ], + transcriptsByRunId: new Map([ + ["run-1", [{ kind: "assistant", ts: "2026-04-06T12:01:05.000Z", text: "Working on it." }]], + ]), + hasOutputForRun: (runId) => runId === "run-1", + currentUserId: "user-1", + }); + + const cancelledMessages = buildIssueChatMessages({ + comments: [], + timelineEvents: [], + linkedRuns: [ + { + runId: "run-1", + status: "cancelled", + agentId: "agent-1", + agentName: "CodexCoder", + createdAt: new Date("2026-04-06T12:01:00.000Z"), + startedAt: new Date("2026-04-06T12:01:00.000Z"), + finishedAt: new Date("2026-04-06T12:01:08.000Z"), + }, + ], + liveRuns: [], + transcriptsByRunId: new Map([ + ["run-1", [{ kind: "assistant", ts: "2026-04-06T12:01:05.000Z", text: "Working on it." }]], + ]), + hasOutputForRun: (runId) => runId === "run-1", + currentUserId: "user-1", + }); + + expect(liveMessages).toHaveLength(1); + expect(cancelledMessages).toHaveLength(1); + expect(liveMessages[0]).toMatchObject({ id: "run-assistant:run-1", status: { type: "running" } }); + expect(cancelledMessages[0]).toMatchObject({ + id: "run-assistant:run-1", + status: { type: "complete", reason: "stop" }, + metadata: { custom: { runStatus: "cancelled" } }, + }); + }); + it("can keep succeeded runs without transcript output for embedded run feeds", () => { const messages = buildIssueChatMessages({ comments: [], diff --git a/ui/src/lib/issue-chat-messages.ts b/ui/src/lib/issue-chat-messages.ts index 82086320..4a5e9a30 100644 --- a/ui/src/lib/issue-chat-messages.ts +++ b/ui/src/lib/issue-chat-messages.ts @@ -410,7 +410,7 @@ function createHistoricalTranscriptMessage(args: { : []; const message: ThreadAssistantMessage = { - id: `historical-run:${run.runId}`, + id: `run-assistant:${run.runId}`, role: "assistant", createdAt: toDate(run.startedAt ?? run.createdAt), content, @@ -606,7 +606,7 @@ function createLiveRunMessage(args: { const content = parts; const message: ThreadAssistantMessage = { - id: `live-run:${run.id}`, + id: `run-assistant:${run.id}`, role: "assistant", createdAt: toDate(run.startedAt ?? run.createdAt), content,