diff --git a/src/features/chat/ui/AddRemoteHostDialog.tsx b/src/features/chat/ui/AddRemoteHostDialog.tsx new file mode 100644 index 000000000..f25f5f602 --- /dev/null +++ b/src/features/chat/ui/AddRemoteHostDialog.tsx @@ -0,0 +1,148 @@ +import { useEffect, useRef, useState, type FormEvent } from "react"; +import { useTranslation } from "react-i18next"; +import { useRemoteHostStore } from "@/features/remoteHosts/stores/remoteHostStore"; +import { isRemoteBackendError } from "@/shared/api/remoteHosts"; +import { Button } from "@/shared/ui/button"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/shared/ui/dialog"; +import { Input } from "@/shared/ui/input"; +import { Label } from "@/shared/ui/label"; + +interface AddRemoteHostDialogProps { + open: boolean; + onOpenChange: (open: boolean) => void; + onConnected: (host: string) => void; +} + +export function AddRemoteHostDialog({ + open, + onOpenChange, + onConnected, +}: AddRemoteHostDialogProps) { + const { t } = useTranslation("chat"); + const [hostDraft, setHostDraft] = useState(""); + const [error, setError] = useState(null); + const [pending, setPending] = useState(false); + const attemptRef = useRef(0); + + useEffect( + () => () => { + attemptRef.current += 1; + }, + [], + ); + + const handleOpenChange = (nextOpen: boolean) => { + onOpenChange(nextOpen); + if (!nextOpen) { + // The backend connection may still finish, but closing the dialog must + // keep that late result from changing the composer's selected host. + attemptRef.current += 1; + setHostDraft(""); + setError(null); + setPending(false); + } + }; + + const handleSubmit = async (event: FormEvent) => { + event.preventDefault(); + if (pending) return; + + const host = hostDraft.trim(); + if (!host) { + setError(t("toolbar.remoteHost.add.emptyHost")); + return; + } + + setError(null); + setPending(true); + const attempt = ++attemptRef.current; + try { + const outcome = await useRemoteHostStore + .getState() + .ensureHostConnected(host); + if (attemptRef.current !== attempt) return; + if (outcome === "superseded") { + setError(t("toolbar.remoteHost.add.superseded")); + return; + } + onConnected(host); + handleOpenChange(false); + } catch (connectionError) { + if (attemptRef.current !== attempt) return; + setError( + isRemoteBackendError(connectionError) + ? connectionError.message + : String(connectionError), + ); + } finally { + if (attemptRef.current === attempt) { + setPending(false); + } + } + }; + + return ( + + +
+ + {t("toolbar.remoteHost.add.title")} + + {t("toolbar.remoteHost.add.description")} + + +
+ + { + setHostDraft(event.target.value); + if (error) setError(null); + }} + /> + {error ? ( + + ) : null} +
+ + + + +
+ + + ); +} diff --git a/src/features/chat/ui/RemoteHostSelector.tsx b/src/features/chat/ui/RemoteHostSelector.tsx index 5b7043d15..bbb1b338b 100644 --- a/src/features/chat/ui/RemoteHostSelector.tsx +++ b/src/features/chat/ui/RemoteHostSelector.tsx @@ -1,12 +1,20 @@ -import { Laptop, Server } from "lucide-react"; +import { useState } from "react"; +import { Laptop, Plus, Server } from "lucide-react"; import { useTranslation } from "react-i18next"; +import { AddRemoteHostDialog } from "./AddRemoteHostDialog"; import { ChatInputSelector, type ChatInputSelectorItem, } from "./ChatInputSelector"; import { useRemoteHostStore } from "@/features/remoteHosts/stores/remoteHostStore"; -const LOCAL_HOST_VALUE = "__local__"; +const LOCAL_HOST_VALUE = "action:local"; +const ADD_HOST_VALUE = "action:add-ssh-environment"; +const HOST_VALUE_PREFIX = "host:"; + +function hostValue(host: string): string { + return `${HOST_VALUE_PREFIX}${host}`; +} interface RemoteHostSelectorProps { selectedHost?: string | null; @@ -33,6 +41,7 @@ export function RemoteHostSelector({ modal, }: RemoteHostSelectorProps) { const { t } = useTranslation("chat"); + const [addDialogOpen, setAddDialogOpen] = useState(false); const configHosts = useRemoteHostStore((state) => state.configHosts); const manualHosts = useRemoteHostStore((state) => state.manualHosts); const statusByHost = useRemoteHostStore((state) => state.statusByHost); @@ -58,14 +67,24 @@ export function RemoteHostSelector({ }; const hostItems: ChatInputSelectorItem[] = listedHosts.map((host) => ({ - value: host, + value: hostValue(host), label: host, description: statusDescription(host), icon: , })); const handleValueChange = (value: string) => { - onHostChange?.(value === LOCAL_HOST_VALUE ? null : value); + if (value === ADD_HOST_VALUE) { + setAddDialogOpen(true); + return; + } + if (value === LOCAL_HOST_VALUE) { + onHostChange?.(null); + return; + } + if (value.startsWith(HOST_VALUE_PREFIX)) { + onHostChange?.(value.slice(HOST_VALUE_PREFIX.length)); + } }; const handleOpenChange = (nextOpen: boolean) => { @@ -77,52 +96,72 @@ export function RemoteHostSelector({ }; return ( - - ) : ( - - ) - } - open={open} - onOpenChange={handleOpenChange} - onRequestComposerFocus={onRequestComposerFocus} - triggerIconOnly={triggerIconOnly} - triggerVariant="toolbar" - menuLabel={t("toolbar.remoteHost.chooseHost")} - contentWidth="wide" - disabled={disabled} - modal={modal} - sections={[ - { - items: [ - { - value: LOCAL_HOST_VALUE, - label: t("toolbar.remoteHost.thisComputer"), - description: t("toolbar.remoteHost.thisComputerDescription"), - icon: , - }, - ], - }, - ...(hostItems.length > 0 - ? [ + <> + + ) : ( + + ) + } + open={open} + onOpenChange={handleOpenChange} + onRequestComposerFocus={onRequestComposerFocus} + triggerIconOnly={triggerIconOnly} + triggerVariant="toolbar" + menuLabel={t("toolbar.remoteHost.chooseHost")} + contentWidth="wide" + disabled={disabled} + modal={modal} + sections={[ + { + items: [ + { + value: LOCAL_HOST_VALUE, + label: t("toolbar.remoteHost.thisComputer"), + description: t("toolbar.remoteHost.thisComputerDescription"), + icon: , + }, + ], + }, + ...(hostItems.length > 0 + ? [ + { + label: t("toolbar.remoteHost.sshHosts"), + items: hostItems, + }, + ] + : []), + { + items: [ { - label: t("toolbar.remoteHost.sshHosts"), - items: hostItems, + value: ADD_HOST_VALUE, + label: t("toolbar.remoteHost.add.action"), + icon: , }, - ] - : []), - ]} - onValueChange={handleValueChange} - /> + ], + }, + ]} + onValueChange={handleValueChange} + preservesExternalFocus={(value) => value === ADD_HOST_VALUE} + /> + + onHostChange?.(host)} + /> + ); } diff --git a/src/features/chat/ui/__tests__/RemoteHostConnectionBanner.test.tsx b/src/features/chat/ui/__tests__/RemoteHostConnectionBanner.test.tsx index aab2712b5..a0ed3f08a 100644 --- a/src/features/chat/ui/__tests__/RemoteHostConnectionBanner.test.tsx +++ b/src/features/chat/ui/__tests__/RemoteHostConnectionBanner.test.tsx @@ -78,6 +78,7 @@ describe("RemoteHostConnectionBanner", () => { it("reconnects the host and reloads the session transcript", async () => { const ensureHostConnected = vi.fn(async () => { setHostState("ready"); + return "connected" as const; }); const original = useRemoteHostStore.getState().ensureHostConnected; useRemoteHostStore.setState({ ensureHostConnected }); diff --git a/src/features/chat/ui/__tests__/RemoteHostSelector.test.tsx b/src/features/chat/ui/__tests__/RemoteHostSelector.test.tsx index 33e3879d8..74c4e1b18 100644 --- a/src/features/chat/ui/__tests__/RemoteHostSelector.test.tsx +++ b/src/features/chat/ui/__tests__/RemoteHostSelector.test.tsx @@ -1,14 +1,15 @@ -import { render, screen } from "@testing-library/react"; +import { render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { RemoteHostSelector } from "../RemoteHostSelector"; import { useRemoteHostStore } from "@/features/remoteHosts/stores/remoteHostStore"; const mockListSshConfigHosts = vi.fn(); +const mockConnectRemoteHost = vi.fn(); vi.mock("@/shared/api/remoteHosts", () => ({ listSshConfigHosts: (...args: unknown[]) => mockListSshConfigHosts(...args), - connectRemoteHost: vi.fn(), + connectRemoteHost: (...args: unknown[]) => mockConnectRemoteHost(...args), disconnectRemoteHost: vi.fn(), shutdownRemoteHost: vi.fn(), listRemoteBackends: vi.fn().mockResolvedValue([]), @@ -23,9 +24,16 @@ describe("RemoteHostSelector", () => { // Opening the selector refreshes hosts from the SSH config, so the mock // must agree with the seeded store state. mockListSshConfigHosts.mockReset().mockResolvedValue(["devbox", "gpu-box"]); + mockConnectRemoteHost.mockReset().mockResolvedValue({ + incarnation: "slot-1", + generation: 1, + }); useRemoteHostStore.setState({ configHosts: ["devbox", "gpu-box"], + manualHosts: [], statusByHost: { devbox: { state: "ready" } }, + forgottenHosts: {}, + lifecycleByHost: {}, }); }); @@ -70,6 +78,33 @@ describe("RemoteHostSelector", () => { expect(onHostChange).toHaveBeenCalledWith(null); }); + it("treats aliases matching the former action sentinels as SSH hosts", async () => { + const aliases = ["__local__", "__add_ssh_environment__"]; + mockListSshConfigHosts.mockResolvedValue(aliases); + useRemoteHostStore.setState({ configHosts: aliases }); + const user = userEvent.setup(); + const onHostChange = vi.fn(); + const { unmount } = render( + , + ); + + await user.click(screen.getByRole("button", { name: /select computer/i })); + await user.click(screen.getByRole("menuitem", { name: "__local__" })); + expect(onHostChange).toHaveBeenLastCalledWith("__local__"); + unmount(); + + onHostChange.mockClear(); + render( + , + ); + await user.click(screen.getByRole("button", { name: /select computer/i })); + await user.click( + screen.getByRole("menuitem", { name: "__add_ssh_environment__" }), + ); + expect(onHostChange).toHaveBeenLastCalledWith("__add_ssh_environment__"); + expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); + }); + it("still lists a selected host that is missing from the SSH config", async () => { mockListSshConfigHosts.mockResolvedValue(["gpu-box"]); useRemoteHostStore.setState({ configHosts: ["gpu-box"] }); @@ -82,4 +117,157 @@ describe("RemoteHostSelector", () => { screen.getByRole("menuitem", { name: /devbox/i }), ).toBeInTheDocument(); }); + + it("offers an add SSH host action even when no hosts are configured", async () => { + mockListSshConfigHosts.mockResolvedValue([]); + useRemoteHostStore.setState({ configHosts: [], statusByHost: {} }); + const user = userEvent.setup(); + render(); + + await user.click(screen.getByRole("button", { name: /select computer/i })); + + expect( + screen.getByRole("menuitem", { name: /add ssh host/i }), + ).toBeInTheDocument(); + }); + + it("connects and selects a host added from the environment dialog", async () => { + const user = userEvent.setup(); + const onHostChange = vi.fn(); + render( + , + ); + + await user.click(screen.getByRole("button", { name: /select computer/i })); + await user.click(screen.getByRole("menuitem", { name: /add ssh host/i })); + await user.type(screen.getByRole("textbox", { name: /ssh host/i }), "blox"); + await user.click(screen.getByRole("button", { name: /^connect$/i })); + + await waitFor(() => { + expect(mockConnectRemoteHost).toHaveBeenCalledWith("blox"); + expect(onHostChange).toHaveBeenCalledWith("blox"); + }); + expect( + screen.queryByRole("dialog", { name: /add ssh host/i }), + ).not.toBeInTheDocument(); + }); + + it("keeps the add dialog open with feedback when connecting fails", async () => { + mockConnectRemoteHost.mockRejectedValue(new Error("SSH host unavailable")); + const user = userEvent.setup(); + const onHostChange = vi.fn(); + render( + , + ); + + await user.click(screen.getByRole("button", { name: /select computer/i })); + await user.click(screen.getByRole("menuitem", { name: /add ssh host/i })); + await user.type( + screen.getByRole("textbox", { name: /ssh host/i }), + "offline-box", + ); + await user.click(screen.getByRole("button", { name: /^connect$/i })); + + expect(await screen.findByRole("alert")).toHaveTextContent( + "SSH host unavailable", + ); + expect(onHostChange).not.toHaveBeenCalled(); + expect( + screen.getByRole("dialog", { name: /add ssh host/i }), + ).toBeInTheDocument(); + }); + + it("dismisses a pending connection and ignores its late success", async () => { + let resolveConnection: (value: { + incarnation: string; + generation: number; + }) => void = () => {}; + mockConnectRemoteHost.mockImplementation( + () => + new Promise((resolve) => { + resolveConnection = resolve; + }), + ); + const user = userEvent.setup(); + const onHostChange = vi.fn(); + render( + , + ); + + await user.click(screen.getByRole("button", { name: /select computer/i })); + await user.click(screen.getByRole("menuitem", { name: /add ssh host/i })); + await user.type( + screen.getByRole("textbox", { name: /ssh host/i }), + "slow-box", + ); + await user.click(screen.getByRole("button", { name: /^connect$/i })); + await user.click(screen.getByRole("button", { name: /^cancel$/i })); + + expect( + screen.queryByRole("dialog", { name: /add ssh host/i }), + ).not.toBeInTheDocument(); + + resolveConnection({ incarnation: "slot-slow", generation: 1 }); + await waitFor(() => { + expect( + useRemoteHostStore.getState().statusByHost["slow-box"]?.state, + ).toBe("ready"); + }); + expect(onHostChange).not.toHaveBeenCalled(); + }); + + it("does not select a connection lifecycle superseded while the dialog waits", async () => { + const resolvers: Array< + (value: { incarnation: string; generation: number }) => void + > = []; + mockConnectRemoteHost.mockImplementation( + () => + new Promise((resolve) => { + resolvers.push(resolve); + }), + ); + const user = userEvent.setup(); + const onHostChange = vi.fn(); + render( + , + ); + + await user.click(screen.getByRole("button", { name: /select computer/i })); + await user.click(screen.getByRole("menuitem", { name: /add ssh host/i })); + await user.type( + screen.getByRole("textbox", { name: /ssh host/i }), + "superseded-box", + ); + await user.click(screen.getByRole("button", { name: /^connect$/i })); + + const replacement = useRemoteHostStore + .getState() + .ensureHostConnected("superseded-box"); + await waitFor(() => expect(resolvers).toHaveLength(2)); + resolvers[1]?.({ incarnation: "slot-new", generation: 2 }); + await expect(replacement).resolves.toBe("connected"); + resolvers[0]?.({ incarnation: "slot-old", generation: 1 }); + + await waitFor(() => { + expect( + screen.getByRole("button", { name: /^connect$/i }), + ).not.toBeDisabled(); + }); + expect(onHostChange).not.toHaveBeenCalled(); + expect( + screen.getByRole("dialog", { name: /add ssh host/i }), + ).toBeInTheDocument(); + expect(screen.getByRole("alert")).toHaveTextContent( + "This SSH connection changed while connecting. Try again or cancel.", + ); + expect(screen.getByRole("button", { name: /^connect$/i })).toBeEnabled(); + expect(screen.getByRole("button", { name: /^cancel$/i })).toBeEnabled(); + expect( + useRemoteHostStore.getState().statusByHost["superseded-box"], + ).toEqual({ + state: "ready", + incarnation: "slot-new", + generation: 2, + }); + }); }); diff --git a/src/features/remoteHosts/stores/remoteHostStore.test.ts b/src/features/remoteHosts/stores/remoteHostStore.test.ts index 3e14743db..e2ae715bc 100644 --- a/src/features/remoteHosts/stores/remoteHostStore.test.ts +++ b/src/features/remoteHosts/stores/remoteHostStore.test.ts @@ -251,6 +251,23 @@ describe("ensureHostConnected", () => { expect(mocks.connectRemoteHost).not.toHaveBeenCalled(); }); + it("remembers a manually entered host when its backend is already ready", async () => { + useRemoteHostStore.setState({ configHosts: ["configured"] }); + useRemoteHostStore.getState().applyStatusEvent({ + host: "workstation.blox", + ...backendIdentity, + state: "ready", + }); + + await useRemoteHostStore.getState().ensureHostConnected("workstation.blox"); + + expect(mocks.connectRemoteHost).not.toHaveBeenCalled(); + expect(useRemoteHostStore.getState().manualHosts).toEqual([ + "workstation.blox", + ]); + expect(loadPersistedManualHosts()).toEqual(["workstation.blox"]); + }); + it("connects and marks the host ready", async () => { let resolveConnect: (value: RemoteBackendConnection) => void = () => {}; mocks.connectRemoteHost.mockImplementation( @@ -292,9 +309,9 @@ describe("ensureHostConnected", () => { generation: 4, }; resolvers[1]?.(newestConnection); - await newer; + await expect(newer).resolves.toBe("connected"); resolvers[0]?.({ ...connection, incarnation: "slot-old", generation: 9 }); - await older; + await expect(older).resolves.toBe("superseded"); expect(useRemoteHostStore.getState().statusByHost.devbox).toEqual({ state: "ready", @@ -721,6 +738,31 @@ describe("manual host persistence", () => { expect(loadPersistedManualHosts()).toEqual(["adhoc.blox"]); }); + it("retains concurrent already-ready manual hosts in state and storage", async () => { + useRemoteHostStore.setState({ + statusByHost: { + "alpha.blox": { ...backendIdentity, state: "ready" }, + "beta.blox": { + incarnation: "slot-beta", + generation: 2, + state: "ready", + }, + }, + }); + + const accepted = await Promise.all([ + useRemoteHostStore.getState().ensureHostConnected("alpha.blox"), + useRemoteHostStore.getState().ensureHostConnected("beta.blox"), + ]); + + expect(accepted).toEqual(["connected", "connected"]); + expect(useRemoteHostStore.getState().manualHosts).toEqual([ + "beta.blox", + "alpha.blox", + ]); + expect(loadPersistedManualHosts()).toEqual(["beta.blox", "alpha.blox"]); + }); + it("does not record ssh-config hosts as manual", async () => { mocks.connectRemoteHost.mockResolvedValue(connection); useRemoteHostStore.setState({ configHosts: ["configured"] }); diff --git a/src/features/remoteHosts/stores/remoteHostStore.ts b/src/features/remoteHosts/stores/remoteHostStore.ts index 88b1c245a..6945c3e30 100644 --- a/src/features/remoteHosts/stores/remoteHostStore.ts +++ b/src/features/remoteHosts/stores/remoteHostStore.ts @@ -38,6 +38,8 @@ export interface RemoteHostStatus { error?: RemoteBackendErrorLike; } +export type RemoteHostConnectOutcome = "connected" | "superseded"; + function backendStatus( payload: RemoteBackendStatusPayload | RemoteBackendSnapshotEntry, ): RemoteHostStatus { @@ -173,7 +175,8 @@ export interface RemoteHostStore { refreshConfigHosts: () => Promise; syncBackendSnapshot: () => Promise; applyStatusEvent: (payload: RemoteBackendStatusPayload) => void; - ensureHostConnected: (host: string) => Promise; + /** Connect the host and report whether this exact lifecycle became current. */ + ensureHostConnected: (host: string) => Promise; disconnect: (host: string) => Promise; shutdownHost: (host: string, expectedInstanceToken?: string) => Promise; runDoctor: (host: string) => Promise; @@ -267,12 +270,38 @@ export const useRemoteHostStore = create((set, get) => ({ }, ensureHostConnected: async (host) => { - const current = get(); + let current = get(); if ( current.statusByHost[host]?.state === "ready" && !current.forgottenHosts[host] ) { - return; + // A manually entered host can already be ready when it was restored + // from the backend snapshot. Remember it even though no new connect is + // required, otherwise it disappears from the selector after restart. + let accepted = false; + set((state) => { + if ( + state.statusByHost[host]?.state !== "ready" || + state.forgottenHosts[host] + ) { + return state; + } + accepted = true; + if ( + state.configHosts.includes(host) || + state.manualHosts.includes(host) + ) { + return state; + } + const manualHosts = [host, ...state.manualHosts].slice( + 0, + MAX_MANUAL_HOSTS, + ); + persistManualHosts(manualHosts); + return { manualHosts }; + }); + if (accepted) return "connected"; + current = get(); } // An explicit connection starts a new local lifecycle. This is the only @@ -313,6 +342,7 @@ export const useRemoteHostStore = create((set, get) => ({ }); try { const connection = await connectRemoteHost(host); + let accepted = false; set((state) => { if ( state.forgottenHosts[host] || @@ -321,6 +351,7 @@ export const useRemoteHostStore = create((set, get) => ({ ) { return state; } + accepted = true; // A host that connected but isn't in ~/.ssh/config was typed in // manually; remember it across restarts. const isKnown = @@ -348,7 +379,9 @@ export const useRemoteHostStore = create((set, get) => ({ }, }; }); + return accepted ? "connected" : "superseded"; } catch (error) { + let accepted = false; set((state) => { if ( state.forgottenHosts[host] || @@ -357,6 +390,7 @@ export const useRemoteHostStore = create((set, get) => ({ ) { return state; } + accepted = true; const connectPendingLifecycleByHost = { ...state.connectPendingLifecycleByHost, }; @@ -378,6 +412,7 @@ export const useRemoteHostStore = create((set, get) => ({ }, }; }); + if (!accepted) return "superseded"; throw error; } }, @@ -588,7 +623,10 @@ export const useRemoteHostStore = create((set, get) => ({ * action for callers outside React (e.g. session routing in chat). */ export function ensureHostConnected(host: string): Promise { - return useRemoteHostStore.getState().ensureHostConnected(host); + return useRemoteHostStore + .getState() + .ensureHostConnected(host) + .then(() => undefined); } let remoteHostStoreInitStarted = false; diff --git a/src/features/remoteHosts/ui/__tests__/RemoteHostsSettings.test.tsx b/src/features/remoteHosts/ui/__tests__/RemoteHostsSettings.test.tsx index 71dc93f56..a207cf45c 100644 --- a/src/features/remoteHosts/ui/__tests__/RemoteHostsSettings.test.tsx +++ b/src/features/remoteHosts/ui/__tests__/RemoteHostsSettings.test.tsx @@ -11,7 +11,7 @@ import { import { useRemoteHostStore } from "@/features/remoteHosts/stores/remoteHostStore"; import { RemoteHostsSettings } from "../RemoteHostsSettings"; -const ensureHostConnected = vi.fn(async () => {}); +const ensureHostConnected = vi.fn(async () => "connected" as const); const disconnect = vi.fn(async () => {}); const shutdownHost = vi.fn(async () => {}); const runDoctor = vi.fn(async () => {}); diff --git a/src/shared/i18n/locales/en/chat.json b/src/shared/i18n/locales/en/chat.json index 103dc0536..a1a783d40 100644 --- a/src/shared/i18n/locales/en/chat.json +++ b/src/shared/i18n/locales/en/chat.json @@ -536,6 +536,18 @@ "thisComputer": "This computer", "thisComputerDescription": "Run the session locally", "sshHosts": "SSH hosts", + "add": { + "action": "Add SSH host", + "title": "Add SSH host", + "description": "Connect using an SSH config alias or a user@host address.", + "hostLabel": "SSH host", + "hostPlaceholder": "user@host", + "emptyHost": "Enter an SSH host.", + "superseded": "This SSH connection changed while connecting. Try again or cancel.", + "cancel": "Cancel", + "connect": "Connect", + "close": "Close add SSH host dialog" + }, "localTriggerTitle": "Runs on this computer", "remoteTriggerTitle": "Runs on {{host}} over SSH", "missingDirectory": "Choose a folder on the remote host before sending",