diff --git a/src-tauri/src/commands/workspace_windows.rs b/src-tauri/src/commands/workspace_windows.rs index 4484f94540..72870f156e 100644 --- a/src-tauri/src/commands/workspace_windows.rs +++ b/src-tauri/src/commands/workspace_windows.rs @@ -897,7 +897,7 @@ mod tests { pending: vec![remote(1), remote(2)], ..SessionState::default() }; - // Remote 2 is opened by hand before its (failed) check is processed. + // Remote 2 is opened by hand before its build failure is processed. state.mark_open(remote(2)); state.complete_restore(); assert_eq!(state.snapshot(), vec![remote(1), remote(2)]); @@ -1025,6 +1025,40 @@ mod tests { assert_eq!(restore.plan.front, remote(b.id)); } + #[tokio::test] + async fn an_offline_remote_shell_survives_quit_but_an_explicit_close_does_not() { + let db = fresh_in_memory_db().await; + // Restoring a saved shell must not contact this unreachable host. + let connection = remote_workspace_connection_service::create( + &db.conn, + "offline", + "http://127.0.0.1:1", + "token", + &[], + ) + .await + .unwrap(); + let window = remote(connection.id); + let session = WorkspaceWindowSession::new(); + session.apply(|state| state.mark_open(window)); + session.freeze(); + session.write_latest(&db.conn).await; + let restore = load_restore(&db.conn, false, true).await; + assert!(!restore.show_local()); + assert_eq!(restore.plan.remotes, vec![connection.id]); + + // A shell closing deliberately while the server remains unavailable + // still leaves the session. We do not revive every saved connection. + let closed_session = WorkspaceWindowSession::new(); + closed_session.apply(|state| state.mark_open(window)); + closed_session.apply(|state| state.mark_closed(window)); + closed_session.freeze(); + closed_session.write_latest(&db.conn).await; + let restore = load_restore(&db.conn, false, true).await; + assert!(restore.show_local()); + assert!(restore.plan.remotes.is_empty()); + } + #[tokio::test] async fn an_unchanged_list_is_not_written_again() { let db = fresh_in_memory_db().await; diff --git a/src/components/connection/remote-connection-status.test.tsx b/src/components/connection/remote-connection-status.test.tsx new file mode 100644 index 0000000000..39e57cd52d --- /dev/null +++ b/src/components/connection/remote-connection-status.test.tsx @@ -0,0 +1,145 @@ +import { act, cleanup, fireEvent, render, screen } from "@testing-library/react" +import { NextIntlClientProvider } from "next-intl" +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest" +import enMessages from "@/i18n/messages/en.json" +import type { RemoteConnectionState } from "@/lib/transport/remote-desktop-transport" +import { RemoteConnectionStatus } from "./remote-connection-status" + +function createHealth() { + let state: RemoteConnectionState = "reconnecting" + const listeners = new Set<() => void>() + return { + listeners, + getConnectionSnapshot: () => state, + subscribeConnection: (callback: () => void) => { + listeners.add(callback) + return () => { + listeners.delete(callback) + } + }, + eventStream: vi.fn(() => ({ attach: vi.fn() })), + reconnectNow: vi.fn(), + setState(next: RemoteConnectionState) { + state = next + for (const listener of listeners) listener() + }, + } +} + +function renderStatus(transport = createHealth()) { + const view = render( + +
Workspace stays open
+ + + ) + return { ...view, transport } +} + +beforeEach(() => vi.useFakeTimers()) +afterEach(() => { + cleanup() + vi.useRealTimers() +}) + +describe("RemoteConnectionStatus", () => { + it("starts the shared proxy but hides brief initial connects", () => { + const { transport } = renderStatus() + expect(transport.eventStream).toHaveBeenCalledTimes(1) + expect(screen.queryByRole("status")).not.toBeInTheDocument() + act(() => { + vi.advanceTimersByTime(3_999) + transport.setState("connected") + }) + act(() => vi.advanceTimersByTime(10_000)) + expect(screen.queryByRole("status")).not.toBeInTheDocument() + }) + + it("shows initial offline retries, then an honest terminal state and reconnect action", () => { + const { transport } = renderStatus() + act(() => vi.advanceTimersByTime(4_000)) + expect(screen.getByRole("status")).toHaveTextContent("Connection lost") + expect( + screen.getByText(enMessages.WebConnection.reconnectingDescription) + ).toBeInTheDocument() + + act(() => transport.setState("disconnected")) + expect( + screen.queryByText(enMessages.WebConnection.reconnectingDescription) + ).not.toBeInTheDocument() + expect(screen.queryByText("Session expired")).not.toBeInTheDocument() + fireEvent.click(screen.getByRole("button", { name: "Reconnect now" })) + expect(transport.reconnectNow).toHaveBeenCalledTimes(1) + expect(screen.getByText("Workspace stays open")).toBeInTheDocument() + }) + + it("shows terminal failures immediately and clears when the connection recovers", () => { + const { transport } = renderStatus() + act(() => transport.setState("disconnected")) + expect(screen.getByRole("status")).toHaveTextContent("Connection lost") + act(() => transport.setState("connected")) + expect(screen.queryByRole("status")).not.toBeInTheDocument() + }) + + it("starts a fresh grace window for each later outage", () => { + const { transport } = renderStatus() + act(() => vi.advanceTimersByTime(4_000)) + expect(screen.getByRole("status")).toBeInTheDocument() + act(() => transport.setState("connected")) + act(() => transport.setState("reconnecting")) + act(() => vi.advanceTimersByTime(3_999)) + expect(screen.queryByRole("status")).not.toBeInTheDocument() + act(() => vi.advanceTimersByTime(1)) + expect(screen.getByRole("status")).toBeInTheDocument() + }) + + it("shows a requested retry's progress at once, without the grace window", () => { + const { transport } = renderStatus() + act(() => transport.setState("disconnected")) + transport.reconnectNow.mockImplementation(() => + transport.setState("reconnecting") + ) + + fireEvent.click(screen.getByRole("button", { name: "Reconnect now" })) + expect(screen.getByRole("status")).toHaveTextContent("Connection lost") + expect( + screen.getByText(enMessages.WebConnection.reconnectingDescription) + ).toBeInTheDocument() + expect( + screen.queryByRole("button", { name: "Reconnect now" }) + ).not.toBeInTheDocument() + + // The retry gives up again: the action is back. + act(() => transport.setState("disconnected")) + fireEvent.click(screen.getByRole("button", { name: "Reconnect now" })) + expect(screen.getByRole("status")).toBeInTheDocument() + + // Once a retry has connected, a later outage gets its grace again. + act(() => transport.setState("connected")) + act(() => transport.setState("reconnecting")) + expect(screen.queryByRole("status")).not.toBeInTheDocument() + act(() => vi.advanceTimersByTime(4_000)) + expect(screen.getByRole("status")).toBeInTheDocument() + }) + + it("keeps the grace window after a retry the transport did not start", () => { + const { transport } = renderStatus() + act(() => transport.setState("disconnected")) + + // The mock ignores the retry, as the transport does while one is running. + fireEvent.click(screen.getByRole("button", { name: "Reconnect now" })) + expect(transport.reconnectNow).toHaveBeenCalledTimes(1) + act(() => transport.setState("connected")) + act(() => transport.setState("reconnecting")) + + expect(screen.queryByRole("status")).not.toBeInTheDocument() + }) + + it("unsubscribes and cancels the grace timer on unmount", () => { + const { transport, unmount } = renderStatus() + expect(transport.listeners.size).toBe(1) + unmount() + expect(transport.listeners.size).toBe(0) + expect(vi.getTimerCount()).toBe(0) + }) +}) diff --git a/src/components/connection/remote-connection-status.tsx b/src/components/connection/remote-connection-status.tsx new file mode 100644 index 0000000000..6f12d5fb1f --- /dev/null +++ b/src/components/connection/remote-connection-status.tsx @@ -0,0 +1,98 @@ +"use client" + +import { useCallback, useEffect, useState, useSyncExternalStore } from "react" +import { Loader2, WifiOff } from "lucide-react" +import { useTranslations } from "next-intl" +import { Button } from "@/components/ui/button" +import type { RemoteDesktopTransport } from "@/lib/transport/remote-desktop-transport" + +type ConnectionHealth = Pick< + RemoteDesktopTransport, + | "getConnectionSnapshot" + | "subscribeConnection" + | "eventStream" + | "reconnectNow" +> + +/** Mounted after the remote transport is configured, so it cannot subscribe + * to the local shell by mistake. Keep the workspace mounted during an outage, + * which the proxy keeps retrying, and show a manual recovery action once this + * window's subscription to it has stopped. */ +export function RemoteConnectionStatus({ + transport, +}: { + transport: ConnectionHealth +}) { + const t = useTranslations("WebConnection") + const subscribe = useCallback( + (callback: () => void) => transport.subscribeConnection(callback), + [transport] + ) + const getSnapshot = useCallback( + () => transport.getConnectionSnapshot(), + [transport] + ) + const state = useSyncExternalStore(subscribe, getSnapshot, () => "connected") + const [graceElapsed, setGraceElapsed] = useState(false) + // A retry the user just asked for shows its progress at once: the pill is + // already up, and hiding it for the grace window reads as "nothing + // happened". + const [retryRequested, setRetryRequested] = useState(false) + + useEffect(() => { + // Settings and other restored remote shells may have no conversation + // subscribers yet. Start the same shared proxy, not a second socket. + transport.eventStream() + }, [transport]) + + useEffect(() => { + if (state !== "reconnecting") return + const timer = setTimeout(() => setGraceElapsed(true), 4_000) + return () => { + clearTimeout(timer) + setGraceElapsed(false) + setRetryRequested(false) + } + }, [state]) + + const retry = () => { + transport.reconnectNow() + // Only a retry that started: the transport ignores one while another + // is still under way. + setRetryRequested(transport.getConnectionSnapshot() === "reconnecting") + } + + if ( + state === "connected" || + (state === "reconnecting" && !graceElapsed && !retryRequested) + ) { + return null + } + + const retrying = state === "reconnecting" + return ( +
+ {retrying ? ( + + ) : ( + + )} +
+

{t("disconnectedTitle")}

+ {retrying && ( +

+ {t("reconnectingDescription")} +

+ )} +
+ {!retrying && ( + + )} +
+ ) +} diff --git a/src/contexts/remote-connection-context.tsx b/src/contexts/remote-connection-context.tsx index 3c08a7170a..8a1c7ad82f 100644 --- a/src/contexts/remote-connection-context.tsx +++ b/src/contexts/remote-connection-context.tsx @@ -15,7 +15,10 @@ import { useTranslations } from "next-intl" import { clearRemoteDesktopTransport, configureRemoteDesktopTransport, + getTransport, } from "@/lib/transport" +import type { RemoteDesktopTransport } from "@/lib/transport/remote-desktop-transport" +import { RemoteConnectionStatus } from "@/components/connection/remote-connection-status" import { resetBackendScopedStores } from "@/stores/backend-scoped-store-reset" import { getRemoteWorkspaceConnection } from "@/lib/remote-workspace" import { toErrorMessage } from "@/lib/app-error" @@ -107,6 +110,8 @@ export function RemoteConnectionGate({ children }: { children: ReactNode }) { baseUrl: next.base_url, token: next.token, windowInstanceId: remoteWindowId, + // Only a refused token. A socket that is down keeps the workspace + // mounted while it retries; `RemoteConnectionStatus` shows that. onUnauthorized: () => setState((prev) => ({ ...prev, expired: true })), }) @@ -213,6 +218,11 @@ export function RemoteConnectionGate({ children }: { children: ReactNode }) { return ( {children} + {hasRemoteConnection && ( + + )} ) } diff --git a/src/contexts/remote-connection-gate.test.tsx b/src/contexts/remote-connection-gate.test.tsx new file mode 100644 index 0000000000..cdffa1a467 --- /dev/null +++ b/src/contexts/remote-connection-gate.test.tsx @@ -0,0 +1,89 @@ +import { act, cleanup, render, screen } from "@testing-library/react" +import { NextIntlClientProvider } from "next-intl" +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest" +import enMessages from "@/i18n/messages/en.json" +import type { RemoteTransportConfig } from "@/lib/transport/types" + +const mocks = vi.hoisted(() => ({ + configure: vi.fn<(config: RemoteTransportConfig) => void>(), + clear: vi.fn(), + readConnection: vi.fn(), + status: vi.fn(), + transport: { id: "remote-transport" }, +})) + +vi.mock("next/navigation", () => ({ + useSearchParams: () => + new URLSearchParams("remoteConnectionId=17&remoteWindowId=window-instance"), +})) +vi.mock("@/lib/transport", () => ({ + configureRemoteDesktopTransport: mocks.configure, + clearRemoteDesktopTransport: mocks.clear, + getTransport: () => mocks.transport, +})) +vi.mock("@/lib/remote-workspace", () => ({ + getRemoteWorkspaceConnection: mocks.readConnection, +})) +vi.mock("@/stores/backend-scoped-store-reset", () => ({ + resetBackendScopedStores: vi.fn(), +})) +vi.mock("@/components/connection/remote-connection-status", () => ({ + RemoteConnectionStatus: (props: unknown) => { + mocks.status(props) + return
Remote status
+ }, +})) + +import { RemoteConnectionGate } from "./remote-connection-context" + +beforeEach(() => { + vi.clearAllMocks() + mocks.readConnection.mockResolvedValue({ + id: 17, + name: "Remote workspace", + base_url: "http://remote.test:3210", + token: "test-token", + }) +}) + +afterEach(() => { + cleanup() + vi.unstubAllGlobals() +}) + +async function renderGate() { + await act(async () => { + render( + + +
Workspace stays open
+ + + ) + }) + return mocks.configure.mock.calls[mocks.configure.mock.calls.length - 1][0] +} + +describe("RemoteConnectionGate failure recovery", () => { + // A socket that is merely down never reaches the gate: the transport + // keeps retrying it in its own state, which the status pill shows. + it("shows credential expiry once the transport reports a refused token", async () => { + const config = await renderGate() + expect(screen.queryByText(/expired/i)).not.toBeInTheDocument() + + act(() => config.onUnauthorized?.()) + + expect(screen.queryByText("Workspace stays open")).not.toBeInTheDocument() + expect(screen.queryByText("Remote status")).not.toBeInTheDocument() + expect(screen.getByText(/expired/i)).toBeInTheDocument() + }) + + it("binds recovery to the existing transport instead of reloading the window", async () => { + await renderGate() + const props = mocks.status.mock.calls[mocks.status.mock.calls.length - 1][0] + expect(props.transport).toBe(mocks.transport) + expect(props).not.toHaveProperty("onReconnect") + expect(mocks.configure).toHaveBeenCalledTimes(1) + expect(screen.getByText("Workspace stays open")).toBeInTheDocument() + }) +}) diff --git a/src/lib/transport/remote-desktop-transport-lifecycle.test.ts b/src/lib/transport/remote-desktop-transport-lifecycle.test.ts new file mode 100644 index 0000000000..ba4f37fbff --- /dev/null +++ b/src/lib/transport/remote-desktop-transport-lifecycle.test.ts @@ -0,0 +1,1002 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest" +import type { EventEnvelope, LiveSessionSnapshot } from "@/lib/types" +import { RemoteDesktopTransport } from "./remote-desktop-transport" +import type { RemoteTransportConfig } from "./types" + +const { invokeMock, listenMock } = vi.hoisted(() => ({ + invokeMock: + vi.fn< + (command: string, args?: Record) => Promise + >(), + listenMock: + vi.fn< + ( + event: string, + handler: (event: { payload: unknown }) => void + ) => Promise<() => void> + >(), +})) + +vi.mock("@tauri-apps/api/core", () => ({ invoke: invokeMock })) +vi.mock("@tauri-apps/api/event", () => ({ listen: listenMock })) + +let receiveFrame: (event: { payload: unknown }) => void +let unlisten: ReturnType +let transports: RemoteDesktopTransport[] + +beforeEach(() => { + vi.useFakeTimers() + vi.resetAllMocks() + unlisten = vi.fn() + transports = [] + invokeMock.mockResolvedValue(undefined) + listenMock.mockImplementation(async (_event, handler) => { + receiveFrame = handler + return unlisten + }) +}) + +afterEach(() => { + for (const transport of transports) transport.destroy() + vi.clearAllTimers() + vi.restoreAllMocks() + vi.useRealTimers() +}) + +function createTransport(overrides: Partial = {}) { + const transport = new RemoteDesktopTransport({ + id: 17, + name: "Remote workspace", + baseUrl: "http://remote.test:3210/", + token: "test-token", + windowInstanceId: "window-instance", + ...overrides, + }) + transports.push(transport) + return transport +} + +function emit(frame: unknown) { + receiveFrame({ payload: frame }) +} + +function ready() { + emit({ channel: "__ready__", payload: null }) +} + +function disconnect() { + emit({ channel: "__disconnected__", payload: null }) +} + +/** + * The server refused the token on the socket's handshake. Network failures + * never stop the proxy, so this is how its task ends and drops the window's + * subscription, leaving the transport "disconnected". + */ +function tokenRefused() { + emit({ channel: "__unauthorized__", payload: null }) +} + +async function flush() { + await vi.advanceTimersByTimeAsync(0) +} + +async function connectReady() { + const transport = createTransport() + const stream = transport.eventStream() + await flush() + ready() + return { transport, stream } +} + +function sentFrames() { + return invokeMock.mock.calls + .filter(([command]) => command === "remote_ws_send_text") + .map(([, args]) => JSON.parse(args!.text as string)) +} + +function lastSentFrame() { + const frames = sentFrames() + return frames[frames.length - 1] +} + +function commandCalls(command: string) { + return invokeMock.mock.calls.filter(([name]) => name === command) +} + +function attachHandlers() { + return { + onSnapshot: vi.fn(), + onReplay: vi.fn(), + onEvent: vi.fn(), + onDetached: vi.fn(), + } +} + +function snapshot(eventSeq: number): LiveSessionSnapshot { + return { + connection_id: "agent-connection", + conversation_id: 3, + folder_id: 2, + status: "prompting", + external_id: null, + live_message: null, + active_tool_calls: [], + pending_permission: null, + modes: null, + current_mode: null, + config_options: null, + prompt_capabilities: null, + usage: null, + fork_supported: false, + available_commands: [], + selectors_ready: true, + event_seq: eventSeq, + } +} + +function envelope(seq: number): EventEnvelope { + return { + type: "content_delta", + connection_id: "agent-connection", + seq, + text: `chunk ${seq}`, + } +} + +describe("RemoteDesktopTransport lifecycle", () => { + it("registers the window listener before subscribing and waits for __ready__", async () => { + const transport = createTransport() + const handler = vi.fn() + const resolved = vi.fn() + const subscription = transport.subscribe("acp://event", handler) + void subscription.then(resolved) + await flush() + + expect(listenMock).toHaveBeenCalledWith( + "remote-ws-event-17", + expect.any(Function) + ) + expect(invokeMock).toHaveBeenCalledWith("remote_ws_subscribe", { + connectionId: 17, + subscriptionId: expect.any(String), + windowInstanceId: "window-instance", + }) + expect(listenMock.mock.invocationCallOrder[0]).toBeLessThan( + invokeMock.mock.invocationCallOrder[0] + ) + expect(resolved).not.toHaveBeenCalled() + + ready() + const unsubscribe = await subscription + emit({ channel: "acp://event", payload: { text: "hello" } }) + expect(handler).toHaveBeenCalledWith({ text: "hello" }) + + unsubscribe() + emit({ channel: "acp://event", payload: { text: "after unsubscribe" } }) + expect(handler).toHaveBeenCalledTimes(1) + }) + + it("notifies reconnect callbacks only after readiness is restored", async () => { + const transport = createTransport() + const onReconnect = vi.fn() + const unsubscribe = transport.onReconnect(onReconnect) + transport.eventStream() + await flush() + + ready() + expect(onReconnect).not.toHaveBeenCalled() + disconnect() + const resolved = vi.fn() + void transport.waitForReady().then(resolved) + await flush() + expect(resolved).not.toHaveBeenCalled() + expect(onReconnect).not.toHaveBeenCalled() + + ready() + await flush() + expect(resolved).toHaveBeenCalledTimes(1) + expect(onReconnect).toHaveBeenCalledTimes(1) + + unsubscribe() + disconnect() + ready() + expect(onReconnect).toHaveBeenCalledTimes(1) + }) + + it("keeps sibling reconnect callbacks running when one throws", async () => { + const { transport } = await connectReady() + vi.spyOn(console, "error").mockImplementation(() => {}) + transport.onReconnect(() => { + throw new Error("consumer failed") + }) + const sibling = vi.fn() + transport.onReconnect(sibling) + + disconnect() + ready() + + expect(sibling).toHaveBeenCalledTimes(1) + }) + + it("preserves initial readiness waiters when connecting attempts disconnect", async () => { + const transport = createTransport() + const resolved = vi.fn() + void transport.subscribe("acp://event", vi.fn()).then(resolved) + await flush() + + disconnect() + disconnect() + ready() + await flush() + + // No 5s timeout should be needed to release the original promise. + expect(resolved).toHaveBeenCalledTimes(1) + expect(vi.getTimerCount()).toBe(0) + }) + + it("preserves reconnect waiters across repeated disconnect notifications", async () => { + const { transport } = await connectReady() + disconnect() + const subscribed = vi.fn() + const readyWaiter = vi.fn() + void transport.subscribe("acp://event", vi.fn()).then(subscribed) + void transport.waitForReady().then(readyWaiter) + await flush() + + disconnect() + await flush() + expect(subscribed).not.toHaveBeenCalled() + expect(readyWaiter).not.toHaveBeenCalled() + + ready() + await flush() + + expect(subscribed).toHaveBeenCalledTimes(1) + expect(readyWaiter).toHaveBeenCalledTimes(1) + expect(vi.getTimerCount()).toBe(0) + }) + + it("bounds readiness waits for older servers that never emit __ready__", async () => { + vi.spyOn(console, "warn").mockImplementation(() => {}) + const transport = createTransport() + const resolved = vi.fn() + void transport.subscribe("acp://event", vi.fn()).then(resolved) + await flush() + + await vi.advanceTimersByTimeAsync(4_999) + expect(resolved).not.toHaveBeenCalled() + await vi.advanceTimersByTimeAsync(1) + expect(resolved).toHaveBeenCalledTimes(1) + expect(console.warn).toHaveBeenCalledTimes(1) + }) + + it("keeps reconnecting without claiming expiry through a long outage", async () => { + const onUnauthorized = vi.fn() + const transport = createTransport({ onUnauthorized }) + const stream = transport.eventStream() + await flush() + ready() + + // The proxy reports every failed attempt and never gives up on its own. + for (let attempt = 0; attempt < 10; attempt++) disconnect() + + expect(transport.getConnectionSnapshot()).toBe("reconnecting") + expect(onUnauthorized).not.toHaveBeenCalled() + stream.attach("agent-connection", {}, attachHandlers()) + expect(sentFrames()).toEqual([]) + }) + + it("reports a token the socket handshake refused as expired credentials", async () => { + const onUnauthorized = vi.fn() + const transport = createTransport({ onUnauthorized }) + const stream = transport.eventStream() + await flush() + ready() + + disconnect() + expect(onUnauthorized).not.toHaveBeenCalled() + tokenRefused() + + expect(transport.getConnectionSnapshot()).toBe("disconnected") + // The same callback, with the same (absent) arguments, as an HTTP 401. + expect(onUnauthorized.mock.calls).toEqual([[]]) + stream.attach("agent-connection", {}, attachHandlers()) + expect(sentFrames()).toEqual([]) + }) + + it("refreshes once on first ready when initial HTTP reads failed offline", async () => { + const transport = createTransport() + const onRestored = vi.fn() + const onReconnect = vi.fn(() => { + void transport.call("list_all_folder_details").then(onRestored) + }) + transport.onReconnect(onReconnect) + transport.eventStream() + await flush() + const error = { code: "network_error", message: "Server unreachable" } + invokeMock.mockRejectedValueOnce(error) + await expect(transport.call("list_all_folder_details")).rejects.toBe(error) + expect(onReconnect).not.toHaveBeenCalled() + const restoredFolders = [{ id: 2, path: "/remote/project" }] + invokeMock.mockResolvedValue(restoredFolders) + + ready() + await flush() + + expect(onReconnect).toHaveBeenCalledTimes(1) + expect(onRestored).toHaveBeenCalledWith(restoredFolders) + expect(commandCalls("remote_http_call")).toHaveLength(2) + expect(transport.getConnectionSnapshot()).toBe("connected") + expect(vi.getTimerCount()).toBe(0) + }) + + it.each(["__disconnected__", "__unauthorized__"])( + "refreshes initial state after an initial %s signal recovers", + async (channel) => { + const transport = createTransport() + const onReconnect = vi.fn() + const subscribed = vi.fn() + transport.onReconnect(onReconnect) + void transport.subscribe("acp://event", vi.fn()).then(subscribed) + await flush() + emit({ channel, payload: null }) + if (channel === "__unauthorized__") { + transport.reconnectNow() + await flush() + // The ready below answers this retry's subscription, not the first. + expect(commandCalls("remote_ws_subscribe")).toHaveLength(2) + } + await flush() + expect(onReconnect).not.toHaveBeenCalled() + expect(subscribed).not.toHaveBeenCalled() + + ready() + await flush() + + expect(onReconnect).toHaveBeenCalledTimes(1) + expect(subscribed).toHaveBeenCalledTimes(1) + expect(vi.getTimerCount()).toBe(0) + } + ) + + it("refreshes after an initial HTTP failure settles later than first ready", async () => { + const transport = createTransport() + const onReconnect = vi.fn() + transport.onReconnect(onReconnect) + transport.eventStream() + await flush() + let rejectRead!: (reason: unknown) => void + invokeMock.mockReturnValueOnce( + new Promise((_resolve, reject) => { + rejectRead = reject + }) + ) + const error = { code: "network_error", message: "Server was offline" } + const failedRead = expect(transport.call("list_folders")).rejects.toBe( + error + ) + ready() + expect(onReconnect).not.toHaveBeenCalled() + + rejectRead(error) + await failedRead + expect(onReconnect).not.toHaveBeenCalled() + await flush() + + expect(onReconnect).toHaveBeenCalledTimes(1) + expect(vi.getTimerCount()).toBe(0) + }) + + it("cancels a late initial-read recovery callback when destroyed", async () => { + const transport = createTransport() + const onReconnect = vi.fn() + transport.onReconnect(onReconnect) + transport.eventStream() + await flush() + let rejectRead!: (reason: unknown) => void + invokeMock.mockReturnValueOnce( + new Promise((_resolve, reject) => { + rejectRead = reject + }) + ) + const error = { code: "network_error", message: "Server was offline" } + const failedRead = expect(transport.call("list_folders")).rejects.toBe( + error + ) + ready() + rejectRead(error) + await failedRead + expect(vi.getTimerCount()).toBe(1) + + transport.destroy() + await flush() + + expect(onReconnect).not.toHaveBeenCalled() + expect(vi.getTimerCount()).toBe(0) + }) +}) + +// A retry in place starts from a "disconnected" transport. Here a refused +// handshake puts it there, which also ends the proxy task that held the +// subscription. Behind the status pill the usual way there is a subscription +// the IPC layer failed to set up. +describe("RemoteDesktopTransport manual retry", () => { + it("reuses handlers, stream, replay cursor, and readiness waiters when retrying in place", async () => { + const { transport, stream } = await connectReady() + const legacyHandler = vi.fn() + await transport.subscribe("acp://event", legacyHandler) + const handlers = attachHandlers() + const sub = stream.attach("agent-connection", {}, handlers) + emit({ + type: "snapshot", + subscription_id: sub.subscriptionId, + connection_id: "agent-connection", + snapshot: snapshot(10), + event_seq: 10, + }) + const onReconnect = vi.fn() + transport.onReconnect(onReconnect) + disconnect() + const readyWaiter = vi.fn() + void transport.waitForReady().then(readyWaiter) + tokenRefused() + const originalSubscription = commandCalls("remote_ws_subscribe")[0][1] + + transport.reconnectNow() + await flush() + + expect(transport.getConnectionSnapshot()).toBe("reconnecting") + expect(transport.eventStream()).toBe(stream) + expect(readyWaiter).not.toHaveBeenCalled() + expect(sentFrames()).toHaveLength(1) + expect(commandCalls("remote_ws_subscribe")).toHaveLength(2) + expect(commandCalls("remote_ws_subscribe")[1][1]).toEqual( + originalSubscription + ) + expect(listenMock).toHaveBeenCalledTimes(1) + expect(unlisten).not.toHaveBeenCalled() + + ready() + await flush() + expect(readyWaiter).toHaveBeenCalledTimes(1) + expect(onReconnect).toHaveBeenCalledTimes(1) + expect(lastSentFrame()).toEqual({ + action: "attach", + subscription_id: sub.subscriptionId, + connection_id: "agent-connection", + since_seq: 10, + }) + const replay = [envelope(11), envelope(12)] + emit({ + type: "replay", + subscription_id: sub.subscriptionId, + connection_id: "agent-connection", + events: replay, + high_water_seq: 12, + }) + emit({ channel: "acp://event", payload: "recovered event" }) + expect(handlers.onReplay).toHaveBeenCalledWith(replay, 12) + expect(legacyHandler).toHaveBeenCalledWith("recovered event") + expect(vi.getTimerCount()).toBe(0) + }) + + it("awaits unsubscribe and coalesces repeated retry clicks into one subscribe", async () => { + const { transport } = await connectReady() + tokenRefused() + let finishUnsubscribe!: () => void + const unsubscribePending = new Promise((resolve) => { + finishUnsubscribe = resolve + }) + invokeMock.mockImplementation((command) => { + return command === "remote_ws_unsubscribe" + ? unsubscribePending + : Promise.resolve(undefined) + }) + + transport.reconnectNow() + transport.reconnectNow() + transport.reconnectNow() + await flush() + + expect(commandCalls("remote_ws_unsubscribe")).toHaveLength(1) + expect(commandCalls("remote_ws_subscribe")).toHaveLength(1) + expect(transport.getConnectionSnapshot()).toBe("reconnecting") + + finishUnsubscribe() + await flush() + transport.reconnectNow() + expect(commandCalls("remote_ws_subscribe")).toHaveLength(2) + expect(commandCalls("remote_ws_unsubscribe")).toHaveLength(1) + }) + + it("does not replace a healthy or automatically reconnecting proxy", async () => { + const transport = createTransport() + transport.reconnectNow() + expect(invokeMock).not.toHaveBeenCalled() + transport.eventStream() + await flush() + ready() + transport.reconnectNow() + disconnect() + transport.reconnectNow() + await flush() + + expect(commandCalls("remote_ws_subscribe")).toHaveLength(1) + expect(commandCalls("remote_ws_unsubscribe")).toHaveLength(0) + }) + + it("does not resubscribe when destroyed while retry unsubscribe is pending", async () => { + const { transport } = await connectReady() + tokenRefused() + let finishUnsubscribe!: () => void + invokeMock.mockReturnValueOnce( + new Promise((resolve) => { + finishUnsubscribe = resolve + }) + ) + transport.reconnectNow() + await flush() + // The retry, not destroy() below, is what is waiting on the unsubscribe. + expect(commandCalls("remote_ws_unsubscribe")).toHaveLength(1) + transport.destroy() + finishUnsubscribe() + await flush() + transport.reconnectNow() + await flush() + + expect(commandCalls("remote_ws_subscribe")).toHaveLength(1) + expect(listenMock).toHaveBeenCalledTimes(1) + expect(unlisten).toHaveBeenCalledTimes(1) + expect(vi.getTimerCount()).toBe(0) + }) + + it.each(["remote_ws_unsubscribe", "remote_ws_subscribe"])( + "retains active streams and allows another retry after %s fails", + async (failedCommand) => { + const { transport, stream } = await connectReady() + vi.spyOn(console, "warn").mockImplementation(() => {}) + const handlers = attachHandlers() + const sub = stream.attach("agent-connection", { sinceSeq: 22 }, handlers) + tokenRefused() + let failOnce = true + invokeMock.mockImplementation((command) => { + if (command === failedCommand && failOnce) { + failOnce = false + return Promise.reject(new Error("IPC unavailable")) + } + return Promise.resolve(undefined) + }) + + transport.reconnectNow() + await flush() + + // The retry reached the failing command, and the failure ended it. + expect(failOnce).toBe(false) + expect(transport.getConnectionSnapshot()).toBe("disconnected") + expect(transport.eventStream()).toBe(stream) + expect(sentFrames()).toHaveLength(1) + const subscribesBefore = commandCalls("remote_ws_subscribe").length + transport.reconnectNow() + await flush() + expect(transport.getConnectionSnapshot()).toBe("reconnecting") + expect(commandCalls("remote_ws_subscribe")).toHaveLength( + subscribesBefore + 1 + ) + ready() + expect(transport.getConnectionSnapshot()).toBe("connected") + expect(lastSentFrame()).toEqual({ + action: "attach", + subscription_id: sub.subscriptionId, + connection_id: "agent-connection", + since_seq: 22, + }) + emit({ + type: "event", + subscription_id: sub.subscriptionId, + envelope: envelope(23), + }) + expect(handlers.onEvent).toHaveBeenCalledWith(envelope(23)) + } + ) +}) + +describe("RemoteDesktopTransport connection health", () => { + it("publishes lifecycle snapshots once per state change", async () => { + const transport = createTransport() + const states: string[] = [] + transport.subscribeConnection(() => { + states.push(transport.getConnectionSnapshot()) + }) + expect(transport.getConnectionSnapshot()).toBe("reconnecting") + transport.eventStream() + await flush() + + disconnect() + expect(states).toEqual([]) + ready() + ready() + expect(transport.getConnectionSnapshot()).toBe("connected") + expect(states).toEqual(["connected"]) + + disconnect() + disconnect() + expect(transport.getConnectionSnapshot()).toBe("reconnecting") + ready() + tokenRefused() + tokenRefused() + + expect(transport.getConnectionSnapshot()).toBe("disconnected") + expect(states).toEqual([ + "connected", + "reconnecting", + "connected", + "disconnected", + ]) + }) + + it("stops notifying a health subscriber after unsubscribe", async () => { + const { transport } = await connectReady() + const listener = vi.fn() + const unsubscribe = transport.subscribeConnection(listener) + + disconnect() + expect(listener).toHaveBeenCalledTimes(1) + unsubscribe() + unsubscribe() + ready() + + expect(transport.getConnectionSnapshot()).toBe("connected") + expect(listener).toHaveBeenCalledTimes(1) + }) + + it("isolates failing health listeners and ignores late lifecycle events after destroy", async () => { + const { transport } = await connectReady() + vi.spyOn(console, "error").mockImplementation(() => {}) + transport.subscribeConnection(() => { + throw new Error("health listener failed") + }) + const sibling = vi.fn() + transport.subscribeConnection(sibling) + disconnect() + expect(sibling).toHaveBeenCalledTimes(1) + + transport.destroy() + ready() + emit({ channel: "__unauthorized__", payload: null }) + + expect(transport.getConnectionSnapshot()).toBe("reconnecting") + expect(sibling).toHaveBeenCalledTimes(1) + }) + + it("marks confirmed HTTP credential rejection disconnected with the default callback source", async () => { + const onUnauthorized = vi.fn() + const transport = createTransport({ onUnauthorized }) + const stream = transport.eventStream() + await flush() + ready() + const listener = vi.fn() + transport.subscribeConnection(listener) + const error = { + code: "authentication_failed", + message: "Invalid remote token", + } + invokeMock.mockRejectedValueOnce(error) + + await expect(transport.call("get_status")).rejects.toBe(error) + + expect(transport.getConnectionSnapshot()).toBe("disconnected") + expect(listener).toHaveBeenCalledTimes(1) + expect(onUnauthorized.mock.calls).toEqual([[]]) + stream.attach("agent-connection", {}, attachHandlers()) + expect(sentFrames()).toEqual([]) + }) + + it("does not label an ordinary HTTP failure as rejected credentials", async () => { + const onUnauthorized = vi.fn() + const transport = createTransport({ onUnauthorized }) + transport.eventStream() + await flush() + ready() + const listener = vi.fn() + transport.subscribeConnection(listener) + const error = { code: "network_error", message: "Request timed out" } + invokeMock.mockRejectedValueOnce(error) + + await expect(transport.call("get_status")).rejects.toBe(error) + + expect(transport.getConnectionSnapshot()).toBe("connected") + expect(listener).not.toHaveBeenCalled() + expect(onUnauthorized).not.toHaveBeenCalled() + }) +}) + +describe("RemoteDesktopTransport attach stream", () => { + it("queues initial attaches and catches up using snapshot, event, and replay cursors", async () => { + const transport = createTransport() + const stream = transport.eventStream() + expect(transport.eventStream()).toBe(stream) + const handlers = attachHandlers() + const sub = stream.attach("agent-connection", {}, handlers) + await flush() + expect(sentFrames()).toEqual([]) + + ready() + expect(sentFrames()).toEqual([ + { + action: "attach", + subscription_id: sub.subscriptionId, + connection_id: "agent-connection", + }, + ]) + + const initialSnapshot = snapshot(10) + emit({ + type: "snapshot", + subscription_id: sub.subscriptionId, + connection_id: "agent-connection", + snapshot: initialSnapshot, + event_seq: 10, + }) + expect(handlers.onSnapshot).toHaveBeenCalledWith(initialSnapshot, 10) + disconnect() + ready() + expect(lastSentFrame()).toMatchObject({ since_seq: 10 }) + + const events = [envelope(11), envelope(12)] + emit({ + type: "replay", + subscription_id: sub.subscriptionId, + connection_id: "agent-connection", + events, + high_water_seq: 12, + }) + expect(handlers.onReplay).toHaveBeenCalledWith(events, 12) + emit({ + type: "event", + subscription_id: sub.subscriptionId, + envelope: envelope(13), + }) + expect(handlers.onEvent).toHaveBeenCalledWith(envelope(13)) + disconnect() + ready() + expect(lastSentFrame()).toMatchObject({ since_seq: 13 }) + + // Even an empty replay may advance the server's high-water mark. + emit({ + type: "replay", + subscription_id: sub.subscriptionId, + connection_id: "agent-connection", + events: [], + high_water_seq: 18, + }) + disconnect() + ready() + expect(lastSentFrame()).toMatchObject({ since_seq: 18 }) + expect( + invokeMock.mock.calls.filter(([c]) => c === "remote_ws_subscribe") + ).toHaveLength(1) + }) + + it("routes attach frames by subscription and removes detached subscriptions", async () => { + const { stream } = await connectReady() + const firstHandlers = attachHandlers() + const secondHandlers = attachHandlers() + const first = stream.attach( + "agent-connection", + { sinceSeq: 3 }, + firstHandlers + ) + const second = stream.attach("other-connection", {}, secondHandlers) + + emit({ + type: "event", + subscription_id: first.subscriptionId, + envelope: envelope(4), + }) + emit({ type: "event", subscription_id: "unknown", envelope: envelope(5) }) + emit({ type: "pong" }) + emit(null) + emit({ unexpected: true }) + expect(firstHandlers.onEvent).toHaveBeenCalledTimes(1) + expect(secondHandlers.onEvent).not.toHaveBeenCalled() + + first.detach() + first.detach() + expect(sentFrames().filter((frame) => frame.action === "detach")).toEqual([ + { action: "detach", subscription_id: first.subscriptionId }, + ]) + emit({ + type: "event", + subscription_id: first.subscriptionId, + envelope: envelope(6), + }) + emit({ + type: "detached", + subscription_id: second.subscriptionId, + reason: "connection_gone", + }) + expect(firstHandlers.onEvent).toHaveBeenCalledTimes(1) + expect(secondHandlers.onDetached).toHaveBeenCalledWith("connection_gone") + + const sendCount = sentFrames().length + disconnect() + ready() + expect(sentFrames()).toHaveLength(sendCount) + }) + + it("debounces rejected sends and reattaches active subscriptions with their latest cursor", async () => { + const { stream } = await connectReady() + vi.spyOn(console, "warn").mockImplementation(() => {}) + invokeMock.mockRejectedValueOnce(new Error("queue full")) + invokeMock.mockRejectedValueOnce(new Error("queue full")) + const first = stream.attach( + "agent-connection", + { sinceSeq: 7 }, + attachHandlers() + ) + const second = stream.attach( + "other-connection", + { sinceSeq: 20 }, + attachHandlers() + ) + await flush() + emit({ + type: "event", + subscription_id: first.subscriptionId, + envelope: envelope(8), + }) + + await vi.advanceTimersByTimeAsync(199) + expect(sentFrames()).toHaveLength(2) + await vi.advanceTimersByTimeAsync(1) + expect(sentFrames().slice(2)).toEqual([ + { + action: "attach", + subscription_id: first.subscriptionId, + connection_id: "agent-connection", + since_seq: 8, + }, + { + action: "attach", + subscription_id: second.subscriptionId, + connection_id: "other-connection", + since_seq: 20, + }, + ]) + expect(vi.getTimerCount()).toBe(0) + }) + + it("leaves recovery to the next ready event when a rejected send is followed by disconnect", async () => { + const { stream } = await connectReady() + vi.spyOn(console, "warn").mockImplementation(() => {}) + invokeMock.mockRejectedValueOnce(new Error("socket closed")) + stream.attach("agent-connection", { sinceSeq: 9 }, attachHandlers()) + await flush() + + disconnect() + await vi.advanceTimersByTimeAsync(200) + expect(sentFrames()).toHaveLength(1) + ready() + expect(sentFrames()).toHaveLength(2) + expect(lastSentFrame()).toMatchObject({ since_seq: 9 }) + }) +}) + +describe("RemoteDesktopTransport destruction", () => { + it("unsubscribes once, releases readiness waiters, and ignores late frames", async () => { + const transport = createTransport() + const onEvent = vi.fn() + const onReconnect = vi.fn() + transport.onReconnect(onReconnect) + const subscribed = vi.fn() + void transport.subscribe("acp://event", onEvent).then(subscribed) + await flush() + const subscriptionArgs = invokeMock.mock.calls.find( + ([command]) => command === "remote_ws_subscribe" + )![1]! + + transport.destroy() + transport.destroy() + await flush() + + expect(unlisten).toHaveBeenCalledTimes(1) + expect(subscribed).toHaveBeenCalledTimes(1) + expect( + invokeMock.mock.calls.filter(([c]) => c === "remote_ws_unsubscribe") + ).toEqual([ + [ + "remote_ws_unsubscribe", + { + connectionId: 17, + subscriptionId: subscriptionArgs.subscriptionId, + }, + ], + ]) + ready() + emit({ channel: "acp://event", payload: "late event" }) + expect(onEvent).not.toHaveBeenCalled() + expect(onReconnect).not.toHaveBeenCalled() + expect(vi.getTimerCount()).toBe(0) + }) + + it("cancels pending send-failure retries when destroyed", async () => { + const { transport, stream } = await connectReady() + vi.spyOn(console, "warn").mockImplementation(() => {}) + invokeMock.mockRejectedValueOnce(new Error("queue full")) + const handlers = attachHandlers() + const sub = stream.attach("agent-connection", {}, handlers) + await flush() + expect(vi.getTimerCount()).toBe(1) + + transport.destroy() + await vi.advanceTimersByTimeAsync(1_000) + emit({ + type: "event", + subscription_id: sub.subscriptionId, + envelope: envelope(1), + }) + + expect(sentFrames()).toHaveLength(1) + expect(handlers.onEvent).not.toHaveBeenCalled() + expect(vi.getTimerCount()).toBe(0) + }) + + it("cleans up an in-flight listener without subscribing after destruction", async () => { + let finishListen!: (listener: () => void) => void + listenMock.mockReturnValueOnce( + new Promise((resolve) => { + finishListen = resolve + }) + ) + const transport = createTransport() + const resolved = vi.fn() + void transport.subscribe("acp://event", vi.fn()).then(resolved) + + transport.destroy() + finishListen(unlisten) + await flush() + + expect(unlisten).toHaveBeenCalledTimes(1) + expect( + invokeMock.mock.calls.filter(([c]) => c === "remote_ws_subscribe") + ).toHaveLength(0) + expect(resolved).toHaveBeenCalledTimes(1) + expect(vi.getTimerCount()).toBe(0) + }) + + it("reissues cleanup with the same subscription ID after an in-flight subscribe completes", async () => { + let finishSubscribe!: () => void + invokeMock.mockImplementation((command) => { + if (command === "remote_ws_subscribe") { + return new Promise((resolve) => { + finishSubscribe = resolve + }) + } + return Promise.resolve(undefined) + }) + const transport = createTransport() + const subscription = transport.subscribe("acp://event", vi.fn()) + await flush() + const subscriptionArgs = invokeMock.mock.calls.find( + ([command]) => command === "remote_ws_subscribe" + )![1]! + + transport.destroy() + finishSubscribe() + await subscription + + const unsubscribeCalls = invokeMock.mock.calls.filter( + ([command]) => command === "remote_ws_unsubscribe" + ) + expect(unsubscribeCalls).toHaveLength(2) + for (const [, args] of unsubscribeCalls) { + expect(args).toEqual({ + connectionId: 17, + subscriptionId: subscriptionArgs.subscriptionId, + }) + } + expect(unlisten).toHaveBeenCalledTimes(1) + expect(vi.getTimerCount()).toBe(0) + }) +}) diff --git a/src/lib/transport/remote-desktop-transport.ts b/src/lib/transport/remote-desktop-transport.ts index 860f813208..023441c863 100644 --- a/src/lib/transport/remote-desktop-transport.ts +++ b/src/lib/transport/remote-desktop-transport.ts @@ -21,6 +21,7 @@ const READY_TIMEOUT_MS = 5_000 /// the wire format exactly; we re-export `__ready__` via `WS_READY_CHANNEL` /// to keep the shared-constants invariant with `WebTransport`. const WS_DISCONNECTED_CHANNEL = "__disconnected__" +/// The server refused the token on the socket's handshake (HTTP 401). const WS_UNAUTHORIZED_CHANNEL = "__unauthorized__" // Two wire shapes flow on the same `remote-ws-event-{id}` Tauri event: @@ -34,6 +35,11 @@ const WS_UNAUTHORIZED_CHANNEL = "__unauthorized__" // call site rather than silently destructuring `undefined`. type WsFrame = unknown +export type RemoteConnectionState = + | "connected" + | "reconnecting" + | "disconnected" + const ATTACH_FRAME_TYPES = new Set([ "snapshot", "replay", @@ -94,12 +100,17 @@ export class RemoteDesktopTransport implements Transport { /// Tracks whether `__ready__` has arrived since this transport was /// constructed. The first arrival is the initial connect; subsequent /// arrivals are reconnects (after `__disconnected__` reset the promise). - /// A first ready after a failed initial load also triggers state resync. + /// A clean initial connect needs no refetch, but one preceded by failed + /// startup reads must refresh the workspace on its first ready too. private hasReadiedOnce = false // Initial loads can fail while a restored remote is offline. Its first // eventual ready must refetch that state too, even without an earlier ready. - private needsResync = false + private recoveryNeeded = false + private recoveryTimer: ReturnType | null = null + private resubscribing = false private reconnectCallbacks = new Set<() => void>() + private connectionState: RemoteConnectionState = "reconnecting" + private connectionListeners = new Set<() => void>() /// Listener handle for `remote-ws-event-{id}`. Null when not subscribed. private unlistenWsEvent: UnlistenFn | null = null /// Opaque ID generated at construction time, passed to `remote_ws_subscribe` @@ -156,7 +167,7 @@ export class RemoteDesktopTransport implements Transport { ]) if (timeoutId !== undefined) clearTimeout(timeoutId) if (result === "timeout") { - this.needsResync = true + this.recoveryNeeded = true console.warn( `[RemoteDesktopTransport] WS __ready__ frame did not arrive within ${READY_TIMEOUT_MS}ms; ` + "proceeding without server-side subscribe confirmation (initial-connect race may reopen)." @@ -169,6 +180,7 @@ export class RemoteDesktopTransport implements Transport { args?: Record, options?: CallOptions ): Promise { + const initialRequest = !this.hasReadiedOnce try { // Forward `timeoutMs` through to `remote_http_call`. Without this, // the Rust client's 30s default fires before the backend can answer @@ -182,12 +194,33 @@ export class RemoteDesktopTransport implements Transport { }) return result } catch (err) { + if (this.destroyed) throw err // The Rust proxy returns 401 from the remote server as // `AppErrorCode::AuthenticationFailed`. Surface the connection-expired // UI in just the calling window (the rest stay live until they // themselves hit a 401 — per design we don't broadcast). if (isAuthenticationFailed(err)) { + if (this.wsOpen) this.resetReady() + this.wsOpen = false + this.setConnectionState("disconnected") this.config.onUnauthorized?.() + } else if ( + initialRequest && + err && + typeof err === "object" && + (err as RustError).code === "network_error" + ) { + // Initial folder/conversation reads are not gated on the WS. If + // they failed offline, the first ready must retry their hydration. + this.recoveryNeeded = true + if (this.wsOpen && this.recoveryTimer === null) { + // A slow failed HTTP request can settle AFTER the first ready. + // Let its caller finish recording the failure before refreshing. + this.recoveryTimer = setTimeout(() => { + this.recoveryTimer = null + if (!this.destroyed && this.wsOpen) this.notifyReconnect() + }, 0) + } } throw err } @@ -229,6 +262,82 @@ export class RemoteDesktopTransport implements Transport { } } + getConnectionSnapshot(): RemoteConnectionState { + return this.connectionState + } + + subscribeConnection(callback: () => void): UnsubscribeFn { + this.connectionListeners.add(callback) + return () => { + this.connectionListeners.delete(callback) + } + } + + /** Retry a stopped proxy in place, preserving unsaved UI state and streams. */ + reconnectNow(): void { + if ( + this.destroyed || + this.resubscribing || + this.connectionState !== "disconnected" + ) { + return + } + this.resubscribing = true + this.recoveryNeeded = true + this.wsStarted = true + this.setConnectionState("reconnecting") + void (async () => { + try { + // Await removal before reusing the opaque ID, or a delayed + // unsubscribe could remove the freshly inserted registration. + await invoke("remote_ws_unsubscribe", { + connectionId: this.config.id, + subscriptionId: this.subscriptionId, + }) + if (this.destroyed) return + if (this.unlistenWsEvent) await this.subscribeWs() + else await this.startWs() + } catch (err) { + if (!this.destroyed) { + this.setConnectionState("disconnected") + console.warn("[RemoteDesktopTransport] reconnect failed:", err) + } + } finally { + this.resubscribing = false + } + })() + } + + private notifyReconnect(): void { + this.recoveryNeeded = false + if (this.recoveryTimer !== null) { + clearTimeout(this.recoveryTimer) + this.recoveryTimer = null + } + for (const callback of this.reconnectCallbacks) { + try { + callback() + } catch (err) { + console.error("[RemoteDesktopTransport] reconnect callback threw:", err) + } + } + } + + private setConnectionState(state: RemoteConnectionState): void { + if (this.connectionState === state) return + this.connectionState = state + for (const callback of this.connectionListeners) { + try { + callback() + } catch (err) { + console.error( + "[RemoteDesktopTransport] connection listener threw:", + err + ) + } + } + } + eventStream(): EventStream { if (!this.eventStreamInstance) { const host: AttachTransportHost = { @@ -300,7 +409,6 @@ export class RemoteDesktopTransport implements Transport { private async startWs() { this.wsStarted = true - try { const unlisten = await listen( `remote-ws-event-${this.config.id}`, @@ -314,9 +422,15 @@ export class RemoteDesktopTransport implements Transport { this.unlistenWsEvent = unlisten } catch (err) { this.wsStarted = false + this.recoveryNeeded = true + this.setConnectionState("disconnected") throw err } + await this.subscribeWs() + } + + private async subscribeWs() { try { await invoke("remote_ws_subscribe", { connectionId: this.config.id, @@ -336,6 +450,8 @@ export class RemoteDesktopTransport implements Transport { this.unlistenWsEvent?.() this.unlistenWsEvent = null this.wsStarted = false + this.recoveryNeeded = true + this.setConnectionState("disconnected") throw err } } @@ -362,11 +478,11 @@ export class RemoteDesktopTransport implements Transport { // Rust may acknowledge another subscription in the same webview while // the shared socket is already ready. This is not another reconnect. if (this.wsOpen) return - const resync = this.hasReadiedOnce || this.needsResync + const shouldRefresh = this.hasReadiedOnce || this.recoveryNeeded this.hasReadiedOnce = true - this.needsResync = false this.wsOpen = true this.readyResolve() + this.setConnectionState("connected") // Notify EventStream so it can re-issue attach frames for any // active subscriptions. Fires on initial connect AND every reconnect. for (const cb of this.wsReadyCallbacks) { @@ -376,26 +492,18 @@ export class RemoteDesktopTransport implements Transport { console.error("[RemoteDesktopTransport] wsReady callback threw:", err) } } - if (resync) { + if (shouldRefresh) { // Reconnect path: server-side receiver_count was 0 during the // disconnect window, so any event fired in that gap was dropped. // Notify consumers to recover state (e.g. refetch snapshots). // Errors in user callbacks must not break sibling callbacks. - for (const cb of this.reconnectCallbacks) { - try { - cb() - } catch (err) { - console.error( - "[RemoteDesktopTransport] reconnect callback threw:", - err - ) - } - } + this.notifyReconnect() } return } if (channel === WS_DISCONNECTED_CHANNEL) { - this.needsResync = true + this.recoveryNeeded = true + this.setConnectionState("reconnecting") // Repeated failed handshakes share the same waiters. Replacing their // promise again would strand them until the readiness timeout. if (this.wsOpen) this.resetReady() @@ -403,8 +511,13 @@ export class RemoteDesktopTransport implements Transport { return } if (channel === WS_UNAUTHORIZED_CHANNEL) { + // Only a definitive authentication rejection ends Rust's retries, and + // it drops this window's subscription: the same expiry an HTTP 401 + // reports. + this.recoveryNeeded = true + if (this.wsOpen) this.resetReady() this.wsOpen = false - // Only a definitive authentication rejection ends Rust's retries. + this.setConnectionState("disconnected") this.config.onUnauthorized?.() return } @@ -420,6 +533,10 @@ export class RemoteDesktopTransport implements Transport { if (this.destroyed) return this.destroyed = true this.wsOpen = false + if (this.recoveryTimer !== null) { + clearTimeout(this.recoveryTimer) + this.recoveryTimer = null + } if (this.sendFailRetryTimer !== null) { clearTimeout(this.sendFailRetryTimer) this.sendFailRetryTimer = null @@ -442,6 +559,7 @@ export class RemoteDesktopTransport implements Transport { } this.handlers.clear() this.reconnectCallbacks.clear() + this.connectionListeners.clear() this.wsReadyCallbacks.clear() this.eventStreamInstance?.destroy() this.eventStreamInstance = null diff --git a/src/lib/transport/types.ts b/src/lib/transport/types.ts index 562bc89994..16c7100424 100644 --- a/src/lib/transport/types.ts +++ b/src/lib/transport/types.ts @@ -8,6 +8,11 @@ export interface RemoteTransportConfig { baseUrl: string token: string windowInstanceId: string + /** + * The remote server refused the connection's token: a 401 on an HTTP call + * or on the event socket's handshake. Network failures never call this: + * the event socket keeps retrying them. + */ onUnauthorized?: () => void } @@ -121,7 +126,9 @@ export interface Transport { * `receiver_count == 0`, so anything fired between `onclose` and the next * `__ready__` is lost. Re-fetching backend snapshots is the recovery path. * - * Not fired on the initial connect (consumers handle that separately). + * Not fired on a clean initial connect (consumers handle that separately). + * A transport may notify on its first successful connection if offline + * startup reads failed and the consumers need to retry hydration. * Returns an unsubscribe function. Optional — IPC-only transports (e.g. * Tauri) leave this undefined. */ diff --git a/src/stores/app-workspace-store-offline.test.tsx b/src/stores/app-workspace-store-offline.test.tsx new file mode 100644 index 0000000000..f5c3bea0aa --- /dev/null +++ b/src/stores/app-workspace-store-offline.test.tsx @@ -0,0 +1,149 @@ +import { act, cleanup, render, waitFor } from "@testing-library/react" +import { afterEach, beforeEach, expect, it, vi } from "vitest" +import { TabProvider } from "@/contexts/tab-context" +import { + buildNewConversationDraftStorageKey, + loadMessageInputDraftV2, + saveMessageInputDraftV2, +} from "@/lib/message-input-draft" +import type { FolderDetail } from "@/lib/types" +import { + resetAppWorkspaceStore, + useAppWorkspaceStore, +} from "./app-workspace-store" +import { resetTabStore, useTabStore } from "./tab-store" + +const mocks = vi.hoisted(() => ({ + listOpenFolders: vi.fn(), + listAllFolders: vi.fn(), + listGroups: vi.fn(), + listTabs: vi.fn(), + saveTabs: vi.fn(), + activate: vi.fn(), + disconnect: vi.fn(), + translate: (key: string) => key, +})) + +vi.mock("next-intl", () => ({ useTranslations: () => mocks.translate })) +vi.mock("@/lib/api", () => ({ + listOpenFolderDetails: mocks.listOpenFolders, + listAllFolderDetails: mocks.listAllFolders, + listFolderGroups: mocks.listGroups, + listOpenedTabs: mocks.listTabs, + saveOpenedTabs: mocks.saveTabs, +})) +vi.mock("@/lib/platform", () => ({ + subscribe: async () => () => {}, + onTransportReconnect: () => () => {}, +})) +vi.mock("@/contexts/workspace-context", () => ({ + useWorkspaceActions: () => ({ activateConversationPane: mocks.activate }), +})) +vi.mock("@/contexts/acp-connections-context", () => ({ + useAcpActions: () => ({ disconnect: mocks.disconnect }), +})) +vi.mock("@/hooks/use-sorted-available-agents", () => ({ + useSortedAvailableAgents: () => ({ sortedTypes: [], fresh: false }), +})) + +const draftId = "new-offline-draft" +const draftKey = buildNewConversationDraftStorageKey(draftId) +let errors: ReturnType + +beforeEach(() => { + vi.resetAllMocks() + localStorage.clear() + resetAppWorkspaceStore() + localStorage.setItem( + "workspace:tab-groups:v1", + JSON.stringify({ + layout: { type: "group", id: "g-main" }, + assignments: {}, + selection: { "g-main": draftId }, + tileByGroup: {}, + drafts: [ + { + id: draftId, + group: "g-main", + index: 0, + folderId: 1, + workingDir: "/repo", + agentType: "codex", + }, + ], + activeDraft: draftId, + }) + ) + saveMessageInputDraftV2(draftKey, { + type: "doc", + content: [ + { + type: "paragraph", + content: [{ type: "text", text: "Unsent work survives offline" }], + }, + ], + }) + resetTabStore() + mocks.listOpenFolders.mockRejectedValue(new Error("remote offline")) + mocks.listAllFolders.mockRejectedValue(new Error("remote offline")) + mocks.listGroups.mockRejectedValue(new Error("remote offline")) + mocks.listTabs.mockRejectedValue(new Error("remote offline")) + errors = vi.spyOn(console, "error").mockImplementation(() => {}) +}) + +afterEach(() => { + cleanup() + resetTabStore() + errors.mockRestore() + localStorage.clear() +}) + +it("preserves a restored composer draft until an authoritative folder snapshot arrives", async () => { + const savedText = loadMessageInputDraftV2(draftKey) + const storageKey = `codeg:message-input-draft:v2:${draftKey}` + const persistedText = localStorage.getItem(storageKey) + expect(persistedText).not.toBeNull() + render( + +
Offline workspace
+ + ) + await act(async () => { + await useAppWorkspaceStore.getState().fetchFolders() + }) + await waitFor(() => expect(useTabStore.getState().tabsHydrated).toBe(true)) + + // Both HTTP reads failed, but the real TabProvider must not interpret an + // unknown folder set as deletion and clear the persisted composer key. + expect(useAppWorkspaceStore.getState().foldersHydrated).toBe(false) + expect(useTabStore.getState().rawTabs.map((tab) => tab.id)).toContain(draftId) + expect(loadMessageInputDraftV2(draftKey)).toEqual(savedText) + expect(localStorage.getItem(storageKey)).toBe(persistedText) + + const folder: FolderDetail = { + id: 1, + name: "repo", + path: "/repo", + git_branch: null, + default_agent_type: "codex", + last_opened_at: "2026-01-01T00:00:00.000Z", + sort_order: 0, + color: "inherit", + parent_id: null, + kind: "regular", + alias: null, + group_id: null, + } + mocks.listOpenFolders.mockResolvedValue([folder]) + mocks.listAllFolders.mockResolvedValue([folder]) + mocks.listGroups.mockResolvedValue([]) + await act(async () => { + await useAppWorkspaceStore.getState().fetchFolders() + }) + + expect(useAppWorkspaceStore.getState().foldersHydrated).toBe(true) + expect(useTabStore.getState().rawTabs.map((tab) => tab.id)).toContain(draftId) + expect(loadMessageInputDraftV2(draftKey)).toEqual(savedText) + expect(localStorage.getItem(storageKey)).toBe(persistedText) + expect(mocks.saveTabs).not.toHaveBeenCalled() +}) diff --git a/src/stores/app-workspace-store.test.ts b/src/stores/app-workspace-store.test.ts index 53065bedd8..73c82252a7 100644 --- a/src/stores/app-workspace-store.test.ts +++ b/src/stores/app-workspace-store.test.ts @@ -1,4 +1,4 @@ -import { beforeEach, describe, expect, it, vi } from "vitest" +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest" import { resetAppWorkspaceStore, useAppWorkspaceStore, @@ -23,12 +23,18 @@ vi.mock("@/lib/api", () => ({ setFolderGroup: vi.fn(), })) -const { getFolder, getGitHead, listAllFolderDetails, listOpenFolderDetails } = - await import("@/lib/api") +const { + getFolder, + getGitHead, + listAllFolderDetails, + listOpenFolderDetails, + listFolderGroups: readFolderGroups, +} = await import("@/lib/api") const mockGetFolder = vi.mocked(getFolder) const mockGetGitHead = vi.mocked(getGitHead) const mockListAllFolders = vi.mocked(listAllFolderDetails) const mockListOpenFolders = vi.mocked(listOpenFolderDetails) +const mockListFolderGroups = vi.mocked(readFolderGroups) function makeSummary( overrides: Partial & { id: number } @@ -146,6 +152,82 @@ function makeFolder( } } +describe("fetchFolders — offline hydration safety", () => { + let errors: ReturnType + + beforeEach(() => { + mockListOpenFolders.mockReset().mockResolvedValue([]) + mockListAllFolders.mockReset().mockResolvedValue([]) + mockListFolderGroups.mockReset().mockResolvedValue([]) + errors = vi.spyOn(console, "error").mockImplementation(() => {}) + }) + + afterEach(() => errors.mockRestore()) + + it.each(["open", "all", "groups"] as const)( + "keeps folders unknown when the initial %s read fails, then hydrates on recovery", + async (failedRead) => { + const reads = { + open: mockListOpenFolders, + all: mockListAllFolders, + groups: mockListFolderGroups, + } + reads[failedRead].mockRejectedValueOnce(new Error("remote offline")) + + await useAppWorkspaceStore.getState().fetchFolders() + + // TabProvider gates orphan-draft pruning on foldersHydrated. A failed + // snapshot must keep that gate closed, even though loading has settled. + expect(useAppWorkspaceStore.getState()).toMatchObject({ + foldersLoading: false, + foldersHydrated: false, + allFolders: [], + }) + + const folder = makeFolder({ id: 1 }) + mockListOpenFolders.mockResolvedValueOnce([folder]) + mockListAllFolders.mockResolvedValueOnce([folder]) + await useAppWorkspaceStore.getState().fetchFolders() + + expect(useAppWorkspaceStore.getState()).toMatchObject({ + foldersLoading: false, + foldersHydrated: true, + folders: [folder], + allFolders: [folder], + }) + } + ) + + it("keeps the last authoritative snapshot hydrated when a refresh fails", async () => { + const folder = makeFolder({ id: 1 }) + mockListOpenFolders.mockResolvedValueOnce([folder]) + mockListAllFolders.mockResolvedValueOnce([folder]) + await useAppWorkspaceStore.getState().fetchFolders() + const before = useAppWorkspaceStore.getState() + + mockListAllFolders.mockRejectedValueOnce(new Error("remote offline")) + await useAppWorkspaceStore.getState().fetchFolders() + + const after = useAppWorkspaceStore.getState() + expect(after.foldersLoading).toBe(false) + expect(after.foldersHydrated).toBe(true) + expect(after.folders).toBe(before.folders) + expect(after.allFolders).toBe(before.allFolders) + expect(after.folderGroups).toBe(before.folderGroups) + }) + + it("accepts a successful empty snapshot as authoritative", async () => { + await useAppWorkspaceStore.getState().fetchFolders() + + expect(useAppWorkspaceStore.getState()).toMatchObject({ + foldersLoading: false, + foldersHydrated: true, + folders: [], + allFolders: [], + }) + }) +}) + describe("refreshFolder — branch null-guard", () => { it("keeps the poll-resolved branch when the refreshed row's git_branch is null", async () => { // Git-head polling has populated the display branch; the folder row's diff --git a/src/stores/app-workspace-store.ts b/src/stores/app-workspace-store.ts index 315aac9f89..727911385a 100644 --- a/src/stores/app-workspace-store.ts +++ b/src/stores/app-workspace-store.ts @@ -52,6 +52,8 @@ export interface AppWorkspaceStoreState { * group hasn't arrived — would render a half-built sidebar. */ folderGroups: FolderGroupDetail[] + /** True only after a successful folder snapshot, so an offline start cannot + * mistake an unknown folder list for deleted folders and prune drafts. */ foldersHydrated: boolean foldersLoading: boolean @@ -388,11 +390,12 @@ export const useAppWorkspaceStore = create()( allFolders: allList, folderGroups: groupList, branches, + foldersHydrated: true, }) } catch (err) { console.error("[AppWorkspace] fetchFolders failed:", err) } finally { - set({ foldersLoading: false, foldersHydrated: true }) + set({ foldersLoading: false }) } },