fix(desktop): remote server switching (#17214)

Co-authored-by: Brendan Allan <[email protected]>

6 files changed, 133 insertions, 89 deletions

diff --git a/packages/app/src/app.tsx b/packages/app/src/app.tsx
index e37086221..857892123 100644
--- a/packages/app/src/app.tsx
+++ b/packages/app/src/app.tsx
@@ -265,6 +265,15 @@ function ConnectionError(props: { onRetry?: () => void; onServerSelected?: (key:
   )
 }
 
+function ServerKey(props: ParentProps) {
+  const server = useServer()
+  return (
+    <Show when={server.key} keyed>
+      {props.children}
+    </Show>
+  )
+}
+
 export function AppInterface(props: {
   children?: JSX.Element
   defaultServer: ServerConnection.Key
@@ -275,20 +284,22 @@ export function AppInterface(props: {
   return (
     <ServerProvider defaultServer={props.defaultServer} servers={props.servers}>
       <ConnectionGate disableHealthCheck={props.disableHealthCheck}>
-        <GlobalSDKProvider>
-          <GlobalSyncProvider>
-            <Dynamic
-              component={props.router ?? Router}
-              root={(routerProps) => <RouterRoot appChildren={props.children}>{routerProps.children}</RouterRoot>}
-            >
-              <Route path="/" component={HomeRoute} />
-              <Route path="/:dir" component={DirectoryLayout}>
-                <Route path="/" component={SessionIndexRoute} />
-                <Route path="/session/:id?" component={SessionRoute} />
-              </Route>
-            </Dynamic>
-          </GlobalSyncProvider>
-        </GlobalSDKProvider>
+        <ServerKey>
+          <GlobalSDKProvider>
+            <GlobalSyncProvider>
+              <Dynamic
+                component={props.router ?? Router}
+                root={(routerProps) => <RouterRoot appChildren={props.children}>{routerProps.children}</RouterRoot>}
+              >
+                <Route path="/" component={HomeRoute} />
+                <Route path="/:dir" component={DirectoryLayout}>
+                  <Route path="/" component={SessionIndexRoute} />
+                  <Route path="/session/:id?" component={SessionRoute} />
+                </Route>
+              </Dynamic>
+            </GlobalSyncProvider>
+          </GlobalSDKProvider>
+        </ServerKey>
       </ConnectionGate>
     </ServerProvider>
   )
diff --git a/packages/app/src/components/dialog-select-server.tsx b/packages/app/src/components/dialog-select-server.tsx
index eb039c14d..f8d14cbb9 100644
--- a/packages/app/src/components/dialog-select-server.tsx
+++ b/packages/app/src/components/dialog-select-server.tsx
@@ -291,8 +291,8 @@ export function DialogSelectServer() {
       navigate("/")
       return
     }
-    server.setActive(ServerConnection.key(conn))
     navigate("/")
+    queueMicrotask(() => server.setActive(ServerConnection.key(conn)))
   }
 
   const handleAddChange = (value: string) => {
diff --git a/packages/app/src/components/status-popover.tsx b/packages/app/src/components/status-popover.tsx
index 61facb84e..063205f0c 100644
--- a/packages/app/src/components/status-popover.tsx
+++ b/packages/app/src/components/status-popover.tsx
@@ -277,8 +277,8 @@ export function StatusPopover() {
                         aria-disabled={isBlocked()}
                         onClick={() => {
                           if (isBlocked()) return
-                          server.setActive(key)
                           navigate("/")
+                          queueMicrotask(() => server.setActive(key))
                         }}
                       >
                         <ServerHealthIndicator health={health[key]} />
diff --git a/packages/app/src/components/terminal.tsx b/packages/app/src/components/terminal.tsx
index 9297d6626..aed46f126 100644
--- a/packages/app/src/components/terminal.tsx
+++ b/packages/app/src/components/terminal.tsx
@@ -165,6 +165,12 @@ export const Terminal = (props: TerminalProps) => {
   const theme = useTheme()
   const language = useLanguage()
   const server = useServer()
+  const directory = sdk.directory
+  const client = sdk.client
+  const url = sdk.url
+  const auth = server.current?.http
+  const username = auth?.username ?? "opencode"
+  const password = auth?.password ?? ""
   let container!: HTMLDivElement
   const [local, others] = splitProps(props, ["pty", "class", "classList", "autoFocus", "onConnect", "onConnectError"])
   const id = local.pty.id
@@ -215,7 +221,7 @@ export const Terminal = (props: TerminalProps) => {
   }
 
   const pushSize = (cols: number, rows: number) => {
-    return sdk.client.pty
+    return client.pty
       .update({
         ptyID: id,
         size: { cols, rows },
@@ -474,7 +480,7 @@ export const Terminal = (props: TerminalProps) => {
       }
 
       const gone = () =>
-        sdk.client.pty
+        client.pty
           .get({ ptyID: id })
           .then(() => false)
           .catch((err) => {
@@ -506,14 +512,14 @@ export const Terminal = (props: TerminalProps) => {
         if (disposed) return
         drop?.()
 
-        const url = new URL(sdk.url + `/pty/${id}/connect`)
-        url.searchParams.set("directory", sdk.directory)
-        url.searchParams.set("cursor", String(seek))
-        url.protocol = url.protocol === "https:" ? "wss:" : "ws:"
-        url.username = server.current?.http.username ?? "opencode"
-        url.password = server.current?.http.password ?? ""
+        const next = new URL(url + `/pty/${id}/connect`)
+        next.searchParams.set("directory", directory)
+        next.searchParams.set("cursor", String(seek))
+        next.protocol = next.protocol === "https:" ? "wss:" : "ws:"
+        next.username = username
+        next.password = password
 
-        const socket = new WebSocket(url)
+        const socket = new WebSocket(next)
         socket.binaryType = "arraybuffer"
         ws = socket
 
diff --git a/packages/app/src/context/terminal.tsx b/packages/app/src/context/terminal.tsx
index e65c16788..17355aab9 100644
--- a/packages/app/src/context/terminal.tsx
+++ b/packages/app/src/context/terminal.tsx
@@ -185,6 +185,60 @@ function createWorkspaceTerminalSession(sdk: ReturnType<typeof useSDK>, dir: str
   })
   onCleanup(unsub)
 
+  const update = (client: ReturnType<typeof useSDK>["client"], pty: Partial<LocalPTY> & { id: string }) => {
+    const index = store.all.findIndex((x) => x.id === pty.id)
+    const previous = index >= 0 ? store.all[index] : undefined
+    if (index >= 0) {
+      setStore("all", index, (item) => ({ ...item, ...pty }))
+    }
+    client.pty
+      .update({
+        ptyID: pty.id,
+        title: pty.title,
+        size: pty.cols && pty.rows ? { rows: pty.rows, cols: pty.cols } : undefined,
+      })
+      .catch((error: unknown) => {
+        if (previous) {
+          const currentIndex = store.all.findIndex((item) => item.id === pty.id)
+          if (currentIndex >= 0) setStore("all", currentIndex, previous)
+        }
+        console.error("Failed to update terminal", error)
+      })
+  }
+
+  const clone = async (client: ReturnType<typeof useSDK>["client"], id: string) => {
+    const index = store.all.findIndex((x) => x.id === id)
+    const pty = store.all[index]
+    if (!pty) return
+    const next = await client.pty
+      .create({
+        title: pty.title,
+      })
+      .catch((error: unknown) => {
+        console.error("Failed to clone terminal", error)
+        return undefined
+      })
+    if (!next?.data) return
+
+    const active = store.active === pty.id
+
+    batch(() => {
+      setStore("all", index, {
+        id: next.data.id,
+        title: next.data.title ?? pty.title,
+        titleNumber: pty.titleNumber,
+        buffer: undefined,
+        cursor: undefined,
+        scrollY: undefined,
+        rows: undefined,
+        cols: undefined,
+      })
+      if (active) {
+        setStore("active", next.data.id)
+      }
+    })
+  }
+
   return {
     ready,
     all: createMemo(() => store.all),
@@ -216,24 +270,7 @@ function createWorkspaceTerminalSession(sdk: ReturnType<typeof useSDK>, dir: str
         })
     },
     update(pty: Partial<LocalPTY> & { id: string }) {
-      const index = store.all.findIndex((x) => x.id === pty.id)
-      const previous = index >= 0 ? store.all[index] : undefined
-      if (index >= 0) {
-        setStore("all", index, (item) => ({ ...item, ...pty }))
-      }
-      sdk.client.pty
-        .update({
-          ptyID: pty.id,
-          title: pty.title,
-          size: pty.cols && pty.rows ? { rows: pty.rows, cols: pty.cols } : undefined,
-        })
-        .catch((error: unknown) => {
-          if (previous) {
-            const currentIndex = store.all.findIndex((item) => item.id === pty.id)
-            if (currentIndex >= 0) setStore("all", currentIndex, previous)
-          }
-          console.error("Failed to update terminal", error)
-        })
+      update(sdk.client, pty)
     },
     trim(id: string) {
       const index = store.all.findIndex((x) => x.id === id)
@@ -248,37 +285,23 @@ function createWorkspaceTerminalSession(sdk: ReturnType<typeof useSDK>, dir: str
       })
     },
     async clone(id: string) {
-      const index = store.all.findIndex((x) => x.id === id)
-      const pty = store.all[index]
-      if (!pty) return
-      const clone = await sdk.client.pty
-        .create({
-          title: pty.title,
-        })
-        .catch((error: unknown) => {
-          console.error("Failed to clone terminal", error)
-          return undefined
-        })
-      if (!clone?.data) return
-
-      const active = store.active === pty.id
-
-      batch(() => {
-        setStore("all", index, {
-          id: clone.data.id,
-          title: clone.data.title ?? pty.title,
-          titleNumber: pty.titleNumber,
-          // New PTY process, so start clean.
-          buffer: undefined,
-          cursor: undefined,
-          scrollY: undefined,
-          rows: undefined,
-          cols: undefined,
-        })
-        if (active) {
-          setStore("active", clone.data.id)
-        }
-      })
+      await clone(sdk.client, id)
+    },
+    bind() {
+      const client = sdk.client
+      return {
+        trim(id: string) {
+          const index = store.all.findIndex((x) => x.id === id)
+          if (index === -1) return
+          setStore("all", index, (pty) => trimTerminal(pty))
+        },
+        update(pty: Partial<LocalPTY> & { id: string }) {
+          update(client, pty)
+        },
+        async clone(id: string) {
+          await clone(client, id)
+        },
+      }
     },
     open(id: string) {
       setStore("active", id)
@@ -403,6 +426,7 @@ export const { use: useTerminal, provider: TerminalProvider } = createSimpleCont
       trim: (id: string) => workspace().trim(id),
       trimAll: () => workspace().trimAll(),
       clone: (id: string) => workspace().clone(id),
+      bind: () => workspace(),
       open: (id: string) => workspace().open(id),
       close: (id: string) => workspace().close(id),
       move: (id: string, to: number) => workspace().move(id, to),
diff --git a/packages/app/src/pages/session/terminal-panel.tsx b/packages/app/src/pages/session/terminal-panel.tsx
index d62d91c19..c663d7d67 100644
--- a/packages/app/src/pages/session/terminal-panel.tsx
+++ b/packages/app/src/pages/session/terminal-panel.tsx
@@ -280,21 +280,24 @@ export function TerminalPanel() {
               </Tabs>
               <div class="flex-1 min-h-0 relative">
                 <Show when={terminal.active()} keyed>
-                  {(id) => (
-                    <Show when={all().find((pty) => pty.id === id)}>
-                      {(pty) => (
-                        <div id={`terminal-wrapper-${id}`} class="absolute inset-0">
-                          <Terminal
-                            pty={pty()}
-                            autoFocus={opened()}
-                            onConnect={() => terminal.trim(id)}
-                            onCleanup={terminal.update}
-                            onConnectError={() => terminal.clone(id)}
-                          />
-                        </div>
-                      )}
-                    </Show>
-                  )}
+                  {(id) => {
+                    const ops = terminal.bind()
+                    return (
+                      <Show when={all().find((pty) => pty.id === id)}>
+                        {(pty) => (
+                          <div id={`terminal-wrapper-${id}`} class="absolute inset-0">
+                            <Terminal
+                              pty={pty()}
+                              autoFocus={opened()}
+                              onConnect={() => ops.trim(id)}
+                              onCleanup={ops.update}
+                              onConnectError={() => ops.clone(id)}
+                            />
+                          </div>
+                        )}
+                      </Show>
+                    )
+                  }}
                 </Show>
               </div>
             </div>