fix(desktop): route SSH media through active connection

2026-08-23 22:05:17 +07:00
parent ec44116d59
commit 6eb77df1aa
12 changed files with 286 additions and 35 deletions
@@ -35,6 +35,7 @@ import {
  normalizeRemoteHeaders,
  normalizeSshConfig,
  normAuthMode,
+  pathForRegistryBackendRequest,
  pathWithGlobalRemoteProfile,
  pathWithProfileScope,
  profileHasRemoteConnection,
@@ -484,6 +485,27 @@ test('pathWithProfileScope keeps an explicit profile query and no-ops on empty p
  assert.equal(pathWithProfileScope('/api/cron/jobs', null), '/api/cron/jobs')
 })

+test('pathForRegistryBackendRequest uses the resolved registry backend scope', () => {
+  assert.equal(
+    pathForRegistryBackendRequest('/api/fs/read-data-url?path=%2Fsrv%2Fimage.png', 'research', {
+      sharedRemote: true
+    }),
+    '/api/fs/read-data-url?path=%2Fsrv%2Fimage.png&profile=research'
+  )
+  assert.equal(
+    pathForRegistryBackendRequest('/api/fs/download?path=%2Fsrv%2Freport.pdf&profile=mara', 'mara', {
+      remoteProfile: 'default'
+    }),
+    '/api/fs/download?path=%2Fsrv%2Freport.pdf&profile=default'
+  )
+  assert.equal(
+    pathForRegistryBackendRequest('/api/fs/download?path=%2Fsrv%2Freport.pdf', 'mara', {
+      remoteProfile: 'default'
+    }),
+    '/api/fs/download?path=%2Fsrv%2Freport.pdf'
+  )
+})
+
 // --- pathWithGlobalRemoteProfile ---

 test('pathWithGlobalRemoteProfile appends profile in global remote mode', () => {
@@ -811,6 +811,23 @@ function pathWithProfileScope(path, profile) {
  return `${parsed.pathname}${parsed.search}${parsed.hash}`
 }

+export interface RegistryBackendRequestScope {
+  remoteProfile?: null | string
+  sharedRemote?: boolean
+}
+
+/**
+ * Scope a REST path for a resolved registry backend. Shared remotes serve
+ * multiple profiles from one process and need an explicit profile query;
+ * isolated SSH backends already own one profile but may translate a Desktop
+ * alias in an existing self-profile filter.
+ */
+function pathForRegistryBackendRequest(path, profile, backend: RegistryBackendRequestScope) {
+  return backend.sharedRemote
+    ? pathWithProfileScope(path, profile)
+    : translateSelfProfileQuery(path, profile, backend.remoteProfile)
+}
+
 /**
 * Registry connection a REST request is explicitly pinned to, or null for the
 * legacy profile-routed path. An explicit `local` id must stay registry-scoped:
@@ -985,6 +1002,7 @@ export {
  normalizeRemoteHeaders,
  normalizeSshConfig,
  normAuthMode,
+  pathForRegistryBackendRequest,
  pathWithGlobalRemoteProfile,
  pathWithProfileScope,
  PRIVY_ACCESS_COOKIE_VARIANTS,
@@ -8,7 +8,8 @@ import {
  gatewayFilePath,
  isNotFoundError,
  parseDataUrlToBuffer,
-  pumpStreamToFile
+  pumpStreamToFile,
+  resolveGatewayFileBackend
 } from './gateway-file-download'

 // A Readable-like response driven manually in tests.
@@ -188,3 +189,55 @@ test('isNotFoundError matches only HTTP 404', () => {
  assert.equal(isNotFoundError(new Error('plain')), false)
  assert.equal(isNotFoundError(null), false)
 })
+
+test('resolveGatewayFileBackend pins registered files to their owning connection', async () => {
+  const calls: string[] = []
+
+  const route = await resolveGatewayFileBackend(
+    { connectionId: '  work-ssh  ', profile: ' default ' },
+    {
+      ensureLegacy: async profile => {
+        calls.push(`legacy:${profile}`)
+
+        return { baseUrl: 'http://local.invalid' }
+      },
+      ensureRegistry: async (connectionId, profile) => {
+        calls.push(`registry:${connectionId}:${profile}`)
+
+        return { baseUrl: 'http://ssh.invalid' }
+      }
+    }
+  )
+
+  assert.deepEqual(calls, ['registry:work-ssh:default'])
+  assert.deepEqual(route, {
+    connection: { baseUrl: 'http://ssh.invalid' },
+    connectionId: 'work-ssh',
+    profile: 'default'
+  })
+})
+
+test('resolveGatewayFileBackend preserves the legacy route when no connection owns the file', async () => {
+  const calls: string[] = []
+
+  const route = await resolveGatewayFileBackend(
+    { profile: 'coder' },
+    {
+      ensureLegacy: async profile => {
+        calls.push(`legacy:${profile}`)
+
+        return { baseUrl: 'http://local.invalid' }
+      },
+      ensureRegistry: async connectionId => {
+        calls.push(`registry:${connectionId}`)
+
+        return { baseUrl: 'http://remote.invalid' }
+      }
+    }
+  )
+
+  assert.deepEqual(calls, ['legacy:coder'])
+  assert.equal(route.connectionId, null)
+  assert.equal(route.profile, 'coder')
+  assert.deepEqual(route.connection, { baseUrl: 'http://local.invalid' })
+})
@@ -35,6 +35,34 @@ export interface PumpDeps {
  unlink: (destPath: string) => Promise<unknown>
 }

+export interface GatewayFileBackendDeps<T> {
+  ensureLegacy: (profile: null | string) => Promise<T>
+  ensureRegistry: (connectionId: string, profile: null | string) => Promise<T>
+}
+
+export interface GatewayFileBackendRoute<T> {
+  connection: T
+  connectionId: null | string
+  profile: null | string
+}
+
+/**
+ * Resolve the backend that owns a renderer-requested gateway file. Registered
+ * connections must never fall through to the legacy profile pool: that pool
+ * can point at another machine with another authentication credential.
+ */
+export async function resolveGatewayFileBackend<T>(
+  payload: { connectionId?: unknown; profile?: unknown },
+  deps: GatewayFileBackendDeps<T>
+): Promise<GatewayFileBackendRoute<T>> {
+  const connectionId = String(payload.connectionId ?? '').trim() || null
+  const profile = String(payload.profile ?? '').trim() || null
+
+  const connection = connectionId ? await deps.ensureRegistry(connectionId, profile) : await deps.ensureLegacy(profile)
+
+  return { connection, connectionId, profile }
+}
+
 // Stream `res` into `destPath`, honoring backpressure. On any read/write error
 // the write stream is torn down and the (partial) destination file is removed
 // before the returned promise rejects, so a failed download never leaves a
@@ -87,11 +87,12 @@ import {
  normalizeRemoteHeaders,
  normalizeSshConfig,
  normAuthMode,
+  pathForRegistryBackendRequest,
  pathWithGlobalRemoteProfile,
-  pathWithProfileScope,
  profileHasRemoteConnection,
  profileRemoteOverride,
  profileSshOverride,
+  type RegistryBackendRequestScope,
  remoteRequestMatchesBaseUrl,
  resolveAuthMode,
  resolveProfileApiRequest,
@@ -100,7 +101,6 @@ import {
  resolveTestWsUrl,
  savedProfileSsh,
  tokenPreview,
-  translateSelfProfileQuery,
  withTransientRetries
 } from './connection-config'
 import { applyConnectionConfigAtomically } from './connection-config-apply'
@@ -166,7 +166,8 @@ import {
  gatewayFilePath,
  isNotFoundError,
  parseDataUrlToBuffer,
-  pumpStreamToFile
+  pumpStreamToFile,
+  resolveGatewayFileBackend
 } from './gateway-file-download'
 import { probeGatewayWebSocket } from './gateway-ws-probe'
 import { registerGitIpc } from './git-ipc'
@@ -1292,7 +1293,8 @@ function registerMediaProtocol() {

      return resolvedPath
    },
-    resolveRemoteConnection: profile => ensureBackend(profile)
+    resolveRemoteConnection: ({ connectionId, profile }) =>
+      connectionId ? ensureRegistryBackend(connectionId, profile) : ensureBackend(profile)
  })

  protocol.handle(MEDIA_PROTOCOL, handler)
@@ -7517,30 +7519,63 @@ function readGatewayErrorText(res): Promise<string> {
  })
 }

-async function gatedFileAuth(connection) {
+interface GatewayFileConnection extends RegistryBackendRequestScope {
+  authMode?: 'oauth' | 'token'
+  baseUrl: string
+  token?: null | string
+}
+
+interface GatewayFileSaveContext {
+  fallbackName: string
+  suggested: string
+}
+
+interface GatewayFileSavePayload {
+  connectionId?: unknown
+  path?: unknown
+  profile?: unknown
+  suggestedName?: unknown
+}
+
+async function gatedFileAuth(connection: GatewayFileConnection) {
  const nativeAt =
    connection.authMode === 'oauth' ? await ensureNativeAccessToken(connection.baseUrl).catch(() => null) : null

  return resolveGatedDownloadAuth(connection.authMode, nativeAt, connection.token)
 }

-async function saveGatewayFile(payload: any = {}) {
+function gatewayFileRequestPath(
+  connection: GatewayFileConnection,
+  connectionId: null | string,
+  profile: null | string,
+  requestPath: string
+) {
+  return connectionId
+    ? pathForRegistryBackendRequest(requestPath, profile, connection)
+    : pathWithGlobalRemoteProfile(requestPath, profile, profileRouteOptions(profile))
+}
+
+async function saveGatewayFile(payload: GatewayFileSavePayload = {}) {
  const filePath = gatewayFilePath(payload.path)

  if (!filePath) {
    throw new Error('Missing gateway file path')
  }

-  const profile = payload.profile || null
-  const connection = await ensureBackend(profile)
+  const { connection, connectionId, profile } = await resolveGatewayFileBackend<GatewayFileConnection>(payload, {
+    ensureLegacy: ensureBackend,
+    ensureRegistry: ensureRegistryBackend
+  })
+
  const suggested = String(payload.suggestedName || '').trim()
  const fallbackName = path.basename(filePath) || suggested || 'download'
  const ctx = { suggested, fallbackName }

-  const requestPath = pathWithGlobalRemoteProfile(
-    `/api/fs/download?path=${encodeURIComponent(filePath)}`,
+  const requestPath = gatewayFileRequestPath(
+    connection,
+    connectionId,
    profile,
-    profileRouteOptions(profile)
+    `/api/fs/download?path=${encodeURIComponent(filePath)}`
  )

  const url = `${connection.baseUrl}${requestPath}`
@@ -7562,7 +7597,7 @@ async function saveGatewayFile(payload: any = {}) {
    // /api/fs/download 404s here; fall back (ONLY on 404) to the older capped
    // data-URL route so downloads keep working against older backends.
    if (isNotFoundError(error)) {
-      return await saveGatewayFileViaDataUrl(connection, profile, filePath, ctx)
+      return await saveGatewayFileViaDataUrl(connection, connectionId, profile, filePath, ctx)
    }

    throw error
@@ -7573,16 +7608,23 @@ async function saveGatewayFile(payload: any = {}) {
 // `/api/fs/read-data-url` route, decode it, and save. Bounded by the gateway's
 // data-URL cap, so it only serves smaller files — enough to keep older gateways
 // working until they gain the streaming route.
-async function saveGatewayFileViaDataUrl(connection, profile, filePath, ctx: any = {}) {
-  const requestPath = pathWithGlobalRemoteProfile(
-    `/api/fs/read-data-url?path=${encodeURIComponent(filePath)}`,
+async function saveGatewayFileViaDataUrl(
+  connection: GatewayFileConnection,
+  connectionId: null | string,
+  profile: null | string,
+  filePath: string,
+  ctx: GatewayFileSaveContext
+) {
+  const requestPath = gatewayFileRequestPath(
+    connection,
+    connectionId,
    profile,
-    profileRouteOptions(profile)
+    `/api/fs/read-data-url?path=${encodeURIComponent(filePath)}`
  )

  const url = `${connection.baseUrl}${requestPath}`
  const auth = await gatedFileAuth(connection)
-  let json: any
+  let json: unknown

  if (auth.kind === 'bearer') {
    json = await fetchJson(url, null, { bearer: auth.token })
@@ -7592,7 +7634,8 @@ async function saveGatewayFileViaDataUrl(connection, profile, filePath, ctx: any
    json = await fetchJson(url, auth.token)
  }

-  const dataUrl = json?.dataUrl
+  const dataUrl =
+    json && typeof json === 'object' && 'dataUrl' in json && typeof json.dataUrl === 'string' ? json.dataUrl : ''

  if (!dataUrl) {
    throw new Error('Gateway returned no file data')
@@ -13865,9 +13908,7 @@ async function dispatchRegistryApiRequest(
 ) {
  const connection: any = await ensureRegistryBackend(registryConnectionId, routeProfile)

-  const requestPath = connection.sharedRemote
-    ? pathWithProfileScope(request.path, requestProfile)
-    : translateSelfProfileQuery(request.path, requestProfile, connection.remoteProfile)
+  const requestPath = pathForRegistryBackendRequest(request.path, requestProfile, connection)

  return fetchJsonForBackend(connection, requestPath, {
    method: request?.method,
@@ -15,7 +15,7 @@ function dependencies(overrides: Partial<MediaProtocolDependencies> = {}) {
    fetchRemote: vi.fn(async (_url: string, _headers: Headers) => new Response('remote', { status: 206 })),
    fetchRemoteWithCookies: vi.fn(async (_url: string, _headers: Headers) => new Response('cookie', { status: 206 })),
    resolveLocalFile: vi.fn(async (filePath: string) => filePath),
-    resolveRemoteConnection: vi.fn(async (_profile?: string) => ({
+    resolveRemoteConnection: vi.fn(async (_scope: { connectionId?: string; profile?: string }) => ({
      authMode: 'token' as const,
      baseUrl: 'https://gateway.test',
      mode: 'remote' as const,
@@ -105,13 +105,13 @@ describe('createMediaProtocolHandler', () => {
    })

    const response = await createMediaProtocolHandler(deps)(
-      request('hermes-media://remote/%2Froot%2Foutputs%2Frender.mp4?profile=reviewer', {
+      request('hermes-media://remote/%2Froot%2Foutputs%2Frender.mp4?connectionId=work-ssh&profile=reviewer', {
        Range: 'bytes=0-1023'
      })
    )

    expect(response.status).toBe(206)
-    expect(deps.resolveRemoteConnection).toHaveBeenCalledWith('reviewer')
+    expect(deps.resolveRemoteConnection).toHaveBeenCalledWith({ connectionId: 'work-ssh', profile: 'reviewer' })
    expect(deps.fetchRemote).toHaveBeenCalledOnce()
    const [rawUrl, headers] = vi.mocked(deps.fetchRemote).mock.calls[0]
    const url = new URL(rawUrl)
@@ -122,6 +122,28 @@ describe('createMediaProtocolHandler', () => {
    expect(headers.get('range')).toBe('bytes=0-1023')
  })

+  it('adds profile scope when one registry backend serves multiple profiles', async () => {
+    const deps = dependencies({
+      resolveRemoteConnection: vi.fn(async () => ({
+        authMode: 'token' as const,
+        baseUrl: 'https://gateway.test',
+        mode: 'remote' as const,
+        sharedRemote: true,
+        token: 'secret'
+      }))
+    })
+
+    await createMediaProtocolHandler(deps)(
+      request('hermes-media://remote/%2Froot%2Foutputs%2Frender.mp4?connectionId=cloud&profile=research')
+    )
+
+    const [rawUrl] = vi.mocked(deps.fetchRemote).mock.calls[0]
+    const url = new URL(rawUrl)
+
+    expect(url.searchParams.get('path')).toBe('/root/outputs/render.mp4')
+    expect(url.searchParams.get('profile')).toBe('research')
+  })
+
  it('preserves explicit HEAD requests through the token-auth remote proxy', async () => {
    const fetchRemote = vi.fn(async (..._args: unknown[]) => new Response(null, { status: 200 }))

@@ -19,16 +19,23 @@ export const MEDIA_PROTOCOL = 'hermes-media'
 type MediaProtocolMode = 'remote' | 'stream'

 interface MediaProtocolTarget {
+  connectionId?: string
  filePath: string
  mode: MediaProtocolMode
  profile?: string
 }

+export interface MediaRemoteScope {
+  connectionId?: string
+  profile?: string
+}
+
 export interface MediaRemoteConnection {
  authMode?: 'oauth' | 'token'
  baseUrl: string
  mode?: 'local' | 'remote'
  token?: null | string
+  sharedRemote?: boolean
 }

 type MediaRequestMethod = 'GET' | 'HEAD'
@@ -39,7 +46,7 @@ export interface MediaProtocolDependencies {
  fetchRemote: (url: string, headers: Headers, method: MediaRequestMethod) => Promise<Response>
  fetchRemoteWithCookies: (url: string, headers: Headers, method: MediaRequestMethod) => Promise<Response>
  resolveLocalFile: (filePath: string) => Promise<string>
-  resolveRemoteConnection: (profile?: string) => Promise<MediaRemoteConnection>
+  resolveRemoteConnection: (scope: MediaRemoteScope) => Promise<MediaRemoteConnection>
 }

 function parseMediaProtocolTarget(rawUrl: string): MediaProtocolTarget {
@@ -56,9 +63,10 @@ function parseMediaProtocolTarget(rawUrl: string): MediaProtocolTarget {
    throw new Error('Missing media path')
  }

+  const connectionId = url.searchParams.get('connectionId')?.trim() || undefined
  const profile = url.searchParams.get('profile')?.trim() || undefined

-  return { filePath, mode, profile }
+  return { connectionId, filePath, mode, profile }
 }

 export function isStreamableMediaPath(filePath: string): boolean {
@@ -81,7 +89,7 @@ export function mediaRequestHeaders(source: Headers): Headers {
  return forwarded
 }

-export function remoteMediaEndpoint(baseUrl: string, filePath: string): string {
+export function remoteMediaEndpoint(baseUrl: string, filePath: string, profile?: string): string {
  const normalizedBase = baseUrl.replace(/\/+$/, '')
  const url = new URL(`${normalizedBase}/api/files/stream`)

@@ -91,6 +99,10 @@ export function remoteMediaEndpoint(baseUrl: string, filePath: string): string {

  url.searchParams.set('path', filePath)

+  if (profile) {
+    url.searchParams.set('profile', profile)
+  }
+
  return url.toString()
 }

@@ -133,13 +145,20 @@ export function createMediaProtocolHandler(dependencies: MediaProtocolDependenci
    }

    try {
-      const connection = await dependencies.resolveRemoteConnection(target.profile)
+      const connection = await dependencies.resolveRemoteConnection({
+        connectionId: target.connectionId,
+        profile: target.profile
+      })

      if (connection.mode !== 'remote') {
        return new Response('Remote media backend unavailable', { status: 404 })
      }

-      const endpoint = remoteMediaEndpoint(connection.baseUrl, target.filePath)
+      const endpoint = remoteMediaEndpoint(
+        connection.baseUrl,
+        target.filePath,
+        connection.sharedRemote ? target.profile : undefined
+      )

      if (connection.authMode === 'oauth') {
        const bearer = await dependencies.ensureRemoteBearer(connection.baseUrl)
@@ -221,7 +221,12 @@ declare global {
      }) => Promise<null | string>
      writeClipboard: (text: string) => Promise<boolean>
      readClipboard: () => Promise<string>
-      saveGatewayFile?: (payload: { path: string; profile?: null | string; suggestedName?: string }) => Promise<{
+      saveGatewayFile?: (payload: {
+        connectionId?: null | string
+        path: string
+        profile?: null | string
+        suggestedName?: string
+      }) => Promise<{
        canceled?: boolean
        path?: string
        saved: boolean
@@ -1,5 +1,6 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'

+import { setApiRequestConnection } from '@/api/client'
 import { $connection } from '@/store/session'

 import {
@@ -66,12 +67,14 @@ describe('desktop filesystem facade', () => {
  beforeEach(() => {
    stubBridge()
    $connection.set(null)
+    setApiRequestConnection(null)
  })

  afterEach(() => {
    vi.unstubAllGlobals()
    vi.clearAllMocks()
    $connection.set(null)
+    setApiRequestConnection(null)
    setDesktopFsRemotePicker(null)
  })

@@ -144,6 +147,24 @@ describe('desktop filesystem facade', () => {
    expect(api).toHaveBeenCalledWith({ path: '/api/fs/default-cwd', profile: 'remote-docker' })
  })

+  it('pins SSH filesystem reads to the active registry connection', async () => {
+    $connection.set({
+      connectionId: 'work-ssh',
+      mode: 'remote',
+      profile: 'default',
+      remoteKind: 'ssh'
+    } as never)
+    setApiRequestConnection('work-ssh')
+
+    await readDesktopFileDataUrl('/srv/project/image.png')
+
+    expect(api).toHaveBeenCalledWith({
+      connectionId: 'work-ssh',
+      path: '/api/fs/read-data-url?path=%2Fsrv%2Fproject%2Fimage.png',
+      profile: 'default'
+    })
+  })
+
  it('keys SSH filesystem caches by stable host identity instead of the forwarded port', () => {
    $connection.set({
      mode: 'remote',
@@ -1,3 +1,4 @@
+import { hermesApi } from '@/api/client'
 import type {
  HermesConnection,
  HermesReadDirResult,
@@ -58,7 +59,7 @@ function bridge() {
 }

 function remoteFsApi<T>(path: string, body?: Record<string, unknown>): Promise<T> {
-  return bridge().api<T>(
+  return hermesApi<T>(
    body ? { body, method: 'POST', path, profile: desktopFsProfile() } : { path, profile: desktopFsProfile() }
  )
 }
@@ -91,6 +91,19 @@ describe('mediaGatewayStreamUrl', () => {
    $connection.set({ authMode: 'oauth', mode: 'remote', profile: 'voice reviewer', token: null } as never)
    expect(mediaGatewayStreamUrl('/tmp/a.mp4')).toBe('hermes-media://remote/%2Ftmp%2Fa.mp4?profile=voice%20reviewer')
  })
+
+  it('pins remote streams to their registered connection and profile', () => {
+    $connection.set({
+      connectionId: 'studio-ssh',
+      mode: 'remote',
+      profile: 'voice reviewer',
+      remoteKind: 'ssh'
+    } as never)
+
+    expect(mediaGatewayStreamUrl('/tmp/a.mp4')).toBe(
+      'hermes-media://remote/%2Ftmp%2Fa.mp4?connectionId=studio-ssh&profile=voice%20reviewer'
+    )
+  })
 })

 describe('resolveMediaDisplaySrc', () => {
@@ -225,7 +238,7 @@ describe('downloadGatewayMediaFile', () => {
  beforeEach(() => {
    saveGatewayFile.mockClear()
    vi.stubGlobal('window', { hermesDesktop: { saveGatewayFile } })
-    $connection.set({ mode: 'remote', profile: 'docker-gw' } as never)
+    $connection.set({ connectionId: 'work-ssh', mode: 'remote', profile: 'docker-gw' } as never)
  })

  afterEach(() => {
@@ -240,6 +253,7 @@ describe('downloadGatewayMediaFile', () => {
    })

    expect(saveGatewayFile).toHaveBeenCalledWith({
+      connectionId: 'work-ssh',
      path: '/Users/me/project/a b.md',
      profile: 'docker-gw',
      suggestedName: 'a b.md'
@@ -139,9 +139,15 @@ export function mediaGatewayStreamUrl(path: string): string {

  if (isRemoteGateway()) {
    const file = encodeURIComponent(filePathFromMediaPath(path))
-    const profile = conn?.profile ? `?profile=${encodeURIComponent(conn.profile)}` : ''

-    return `hermes-media://remote/${file}${profile}`
+    const scope = [
+      conn?.connectionId ? `connectionId=${encodeURIComponent(conn.connectionId)}` : '',
+      conn?.profile ? `profile=${encodeURIComponent(conn.profile)}` : ''
+    ]
+      .filter(Boolean)
+      .join('&')
+
+    return `hermes-media://remote/${file}${scope ? `?${scope}` : ''}`
  }

  return mediaExternalUrl(path)
@@ -208,6 +214,7 @@ export async function downloadGatewayMediaFile(
  }

  return window.hermesDesktop.saveGatewayFile({
+    connectionId: conn?.connectionId,
    path: file,
    profile: conn?.profile,
    suggestedName: mediaName(file)