orca/src/shared/remote-runtime-client.test.ts

import type { AddressInfo } from 'node:net'
import { afterEach, describe, expect, it, vi } from 'vitest'
import WebSocketClient, { WebSocketServer, type WebSocket } from 'ws'
import { encodePairingOffer, parsePairingCode, type PairingOffer } from './pairing'
import {
  decrypt,
  decryptBytes,
  deriveSharedKey,
  encrypt,
  generateKeyPair,
  publicKeyFromBase64,
  publicKeyToBase64
} from './e2ee-crypto'
import { sendRemoteRuntimeRequest, subscribeRemoteRuntimeRequest } from './remote-runtime-client'
import { remoteRuntimeClientCapabilities } from './remote-runtime-client-capabilities'
import { MAX_TIMER_DELAY_MS } from './timer-delay'
import {
  BROWSER_CLIENT_HOST_RUNTIME_CAPABILITY,
  BROWSER_CLIENT_PAGE_METADATA_RUNTIME_CAPABILITY,
  BROWSER_NETWORK_TUNNEL_RUNTIME_CAPABILITY,
  ELECTRON_REMOTE_RUNTIME_CLIENT_CAPABILITIES
} from './protocol-version'

const servers: WebSocketServer[] = []

afterEach(async () => {
  await Promise.all(
    servers.splice(0).map(
      (server) =>
        new Promise<void>((resolve) => {
          for (const client of server.clients) {
            client.close()
          }
          server.close(() => resolve())
        })
    )
  )
})

describe('subscribeRemoteRuntimeRequest', () => {
  it('includes WebSocket close details when subscription admission is rejected', async () => {
    const server = await createClosingServer(1013, 'Maximum connections reached')

    await expect(
      subscribeRemoteRuntimeRequest(server.pairing, 'terminal.subscribe', {}, 1000, {
        onResponse: vi.fn(),
        onError: vi.fn()
      })
    ).rejects.toThrow(
      'Remote Orca runtime closed the connection (1013: Maximum connections reached).'
    )
  })

  it('sends encrypted binary frames on an established subscription socket', async () => {
    const server = await createSubscriptionServer()
    const onResponse = vi.fn()
    const onError = vi.fn()

    const subscription = await subscribeRemoteRuntimeRequest(
      server.pairing,
      'terminal.subscribe',
      { terminal: 't1' },
      1000,
      {
        onResponse,
        onError
      }
    )

    await vi.waitFor(() =>
      expect(onResponse).toHaveBeenCalledWith(
        expect.objectContaining({ ok: true, result: { type: 'subscribed' } })
      )
    )
    await expect(server.nextAuth).resolves.toEqual({
      type: 'e2ee_auth',
      deviceToken: 'device-token',
      clientCapabilities: remoteRuntimeClientCapabilities()
    })
    const bytes = new Uint8Array([1, 2, 3])
    expect(subscription.sendBinary(bytes)).toBe(true)
    await expect(server.nextBinary).resolves.toEqual(bytes)
    expect(onError).not.toHaveBeenCalled()
    subscription.close()
  })

  it('binds optional capabilities and rejects a hard outbound queue overflow', async () => {
    const server = await createSubscriptionServer()
    const onResponse = vi.fn()
    const onError = vi.fn()
    const subscription = await subscribeRemoteRuntimeRequest(
      server.pairing,
      'network.browserTunnel',
      {},
      1000,
      { onResponse, onError },
      {
        clientCapabilities: [BROWSER_NETWORK_TUNNEL_RUNTIME_CAPABILITY],
        outboundQueue: { softCapBytes: 0, maxQueuedBytes: 1, maxQueuedFrames: 1 }
      }
    )

    await vi.waitFor(() => expect(onResponse).toHaveBeenCalled())
    await expect(server.nextAuth).resolves.toEqual({
      type: 'e2ee_auth',
      deviceToken: 'device-token',
      clientCapabilities: remoteRuntimeClientCapabilities([
        BROWSER_NETWORK_TUNNEL_RUNTIME_CAPABILITY
      ])
    })
    expect(subscription.sendBinary(new Uint8Array([9]))).toBe(false)
    expect(onError).toHaveBeenCalledWith(
      expect.objectContaining({ code: 'remote_runtime_unavailable' })
    )
  })

  it('accounts encrypted queues and native socket buffers in an injected aggregate budget', async () => {
    const server = await createSubscriptionServer()
    const releaseQueued = vi.fn()
    const releaseSocket = vi.fn()
    const canSend = vi.fn(() => false)
    let readBufferedAmount: (() => number) | undefined
    const outboundMemoryBudget = {
      claimQueuedBytes: vi.fn(() => releaseQueued),
      registerBufferedAmount: vi.fn((read: () => number) => {
        readBufferedAmount = read
        return { canSend, release: releaseSocket }
      })
    }
    const onResponse = vi.fn()
    const onClose = vi.fn()
    const subscription = await subscribeRemoteRuntimeRequest(
      server.pairing,
      'network.browserTunnel',
      {},
      1000,
      { onResponse, onError: vi.fn(), onClose },
      {
        outboundMemoryBudget,
        outboundQueue: { softCapBytes: 1, maxQueuedBytes: 1024, maxQueuedFrames: 8 }
      }
    )
    await vi.waitFor(() => expect(onResponse).toHaveBeenCalled())

    expect(subscription.sendBinary(new Uint8Array([9]))).toBe(true)
    expect(outboundMemoryBudget.registerBufferedAmount).toHaveBeenCalledOnce()
    expect(readBufferedAmount?.()).toBeGreaterThanOrEqual(0)
    expect(canSend).toHaveBeenCalledWith(expect.any(Number), false)
    expect(outboundMemoryBudget.claimQueuedBytes).toHaveBeenCalledWith(expect.any(Number))

    subscription.close()
    subscription.close()
    expect(releaseQueued).toHaveBeenCalledOnce()
    expect(releaseSocket).not.toHaveBeenCalled()
    await vi.waitFor(() => expect(onClose).toHaveBeenCalledOnce())
    expect(releaseSocket).toHaveBeenCalledOnce()
  })

  it('keeps native socket bytes charged through failure cleanup and a later close', async () => {
    const server = await createSubscriptionServer()
    const releaseSocket = vi.fn()
    const onError = vi.fn()
    const subscription = await subscribeRemoteRuntimeRequest(
      server.pairing,
      'network.browserTunnel',
      {},
      1000,
      { onResponse: vi.fn(), onError },
      {
        outboundMemoryBudget: {
          claimQueuedBytes: vi.fn(() => vi.fn()),
          registerBufferedAmount: vi.fn(() => ({ canSend: () => false, release: releaseSocket }))
        },
        outboundQueue: { softCapBytes: 1, maxQueuedBytes: 1, maxQueuedFrames: 1 }
      }
    )

    expect(subscription.sendBinary(new Uint8Array([9]))).toBe(false)
    expect(onError).toHaveBeenCalledWith(
      expect.objectContaining({ code: 'remote_runtime_unavailable' })
    )
    subscription.close()
    expect(releaseSocket).not.toHaveBeenCalled()

    await vi.waitFor(() => expect(releaseSocket).toHaveBeenCalledOnce())
  })

  it('detaches subscription socket listeners after close', async () => {
    const offSpy = vi.spyOn(WebSocketClient.prototype, 'off')
    try {
      const server = await createSubscriptionServer()
      const onResponse = vi.fn()
      const onError = vi.fn()
      const onClose = vi.fn()

      const subscription = await subscribeRemoteRuntimeRequest(
        server.pairing,
        'terminal.subscribe',
        { terminal: 't1' },
        1000,
        {
          onResponse,
          onError,
          onClose
        }
      )

      await vi.waitFor(() => expect(onResponse).toHaveBeenCalled())
      subscription.close()
      await vi.waitFor(() => expect(onClose).toHaveBeenCalledOnce())

      const removedEvents = offSpy.mock.calls.map(([event]) => event)
      expect(removedEvents).toEqual(expect.arrayContaining(['open', 'error', 'close', 'message']))
      expect(subscription.sendBinary(new Uint8Array([9]))).toBe(false)
      expect(onError).not.toHaveBeenCalled()
    } finally {
      offSpy.mockRestore()
    }
  })

  it('closes a half-open subscription socket via client liveness so callers can resubscribe', async () => {
    // Why: dedicated stream sockets (terminal.multiplex, browser.screencast)
    // must not hang forever when a tunnel goes half-open — no close frame, no
    // pongs, no data (#7718). Client liveness surfaces onError/onClose so the
    // renderer's onTransportClose resubscribe path can run.
    const server = await createSubscriptionServer({ disableAutoPong: true })
    const onResponse = vi.fn()
    const onError = vi.fn()
    const onClose = vi.fn()

    const subscription = await subscribeRemoteRuntimeRequest(
      server.pairing,
      'terminal.multiplex',
      {},
      1000,
      { onResponse, onError, onClose },
      { pingIntervalMs: 50, livenessTimeoutMs: 200 }
    )

    await vi.waitFor(() => expect(onResponse).toHaveBeenCalled())
    await vi.waitFor(
      () =>
        expect(onError).toHaveBeenCalledWith(
          expect.objectContaining({ code: 'remote_runtime_unavailable' })
        ),
      { timeout: 5000 }
    )
    expect(onClose).toHaveBeenCalledOnce()
    subscription.close()
  })

  it('closes established subscription sockets after terminal protocol errors', async () => {
    const offSpy = vi.spyOn(WebSocketClient.prototype, 'off')
    try {
      const server = await createSubscriptionServer({ sendMismatchedResponseAfterSubscribe: true })
      const onResponse = vi.fn()
      const onError = vi.fn()
      const onClose = vi.fn()

      const subscription = await subscribeRemoteRuntimeRequest(
        server.pairing,
        'terminal.subscribe',
        { terminal: 't1' },
        1000,
        {
          onResponse,
          onError,
          onClose
        }
      )

      await vi.waitFor(() => expect(onResponse).toHaveBeenCalled())
      await vi.waitFor(() =>
        expect(onError).toHaveBeenCalledWith(
          expect.objectContaining({ code: 'invalid_runtime_response' })
        )
      )
      expect(onClose).toHaveBeenCalledOnce()

      const removedEvents = offSpy.mock.calls.map(([event]) => event)
      expect(removedEvents).toEqual(expect.arrayContaining(['open', 'error', 'close', 'message']))
      expect(subscription.sendBinary(new Uint8Array([9]))).toBe(false)
    } finally {
      offSpy.mockRestore()
    }
  })
})

describe('sendRemoteRuntimeRequest', () => {
  it('keeps generic native authentication free of Electron placement support', async () => {
    let receivedAuth: Record<string, unknown> | null = null
    const server = await createOneShotServer({
      onAuth: (auth) => {
        receivedAuth = auth
      }
    })

    await sendRemoteRuntimeRequest(
      server.pairing,
      'session.tabs.list',
      { worktree: 'id:worktree-a' },
      1000
    )

    expect(receivedAuth).toMatchObject({
      clientCapabilities: remoteRuntimeClientCapabilities()
    })
  })

  it('advertises browser placement support only when the Electron caller opts in', async () => {
    let receivedAuth: Record<string, unknown> | null = null
    const server = await createOneShotServer({
      onAuth: (auth) => {
        receivedAuth = auth
      }
    })

    await sendRemoteRuntimeRequest(
      server.pairing,
      'session.tabs.list',
      { worktree: 'id:worktree-a' },
      1000,
      undefined,
      undefined,
      ELECTRON_REMOTE_RUNTIME_CLIENT_CAPABILITIES
    )

    expect(receivedAuth).toMatchObject({
      clientCapabilities: expect.arrayContaining([
        BROWSER_CLIENT_HOST_RUNTIME_CAPABILITY,
        BROWSER_CLIENT_PAGE_METADATA_RUNTIME_CAPABILITY
      ])
    })
  })

  it.each([-1, 1.5, MAX_TIMER_DELAY_MS + 1, Number.MAX_SAFE_INTEGER + 1])(
    'rejects invalid timer delay %s before reading pairing data',
    async (timeoutMs) => {
      await expect(
        sendRemoteRuntimeRequest({} as PairingOffer, 'status.get', {}, timeoutMs)
      ).rejects.toMatchObject({ code: 'invalid_argument' })
    }
  )

  it('includes WebSocket close details when one-shot admission is rejected', async () => {
    const server = await createClosingServer(1013, 'Maximum connections reached')

    await expect(sendRemoteRuntimeRequest(server.pairing, 'status.get', {}, 1000)).rejects.toThrow(
      'Remote Orca runtime closed the connection (1013: Maximum connections reached).'
    )
  })

  it('classifies a non-Orca handshake as a host identity mismatch', async () => {
    const server = await createInvalidHandshakeServer()

    await expect(
      sendRemoteRuntimeRequest(server.pairing, 'status.get', {}, 1000)
    ).rejects.toMatchObject({
      code: 'invalid_runtime_response',
      pairingStage: 'host-identity'
    })
  })

  it('classifies an undecryptable post-auth frame as a runtime failure', async () => {
    const server = await createOneShotServer({ sendUndecryptableResponse: true })

    await expect(
      sendRemoteRuntimeRequest(server.pairing, 'status.get', {}, 1000)
    ).rejects.toMatchObject({
      code: 'invalid_runtime_response',
      pairingStage: 'runtime'
    })
  })

  it('refreshes the per-call timeout when the runtime sends keepalive frames', async () => {
    const server = await createOneShotServer()

    const response = await sendRemoteRuntimeRequest<{ satisfied: boolean }>(
      server.pairing,
      'terminal.wait',
      { terminal: 't1', for: 'tui-idle', timeoutMs: 550 },
      300
    )

    expect(response).toMatchObject({
      ok: true,
      result: { satisfied: true }
    })
  })

  it('aborts and closes an in-flight one-shot socket', async () => {
    let requestObserved: () => void = () => {}
    const observed = new Promise<void>((resolve) => {
      requestObserved = resolve
    })
    const server = await createOneShotServer({ onRequest: requestObserved })
    const closeSpy = vi.spyOn(WebSocketClient.prototype, 'close')
    const controller = new AbortController()
    try {
      const request = sendRemoteRuntimeRequest(
        server.pairing,
        'skills.install',
        {},
        60_000,
        undefined,
        controller.signal
      )
      await observed

      controller.abort()
      await expect(request).rejects.toMatchObject({ name: 'AbortError' })
      expect(closeSpy).toHaveBeenCalled()
    } finally {
      closeSpy.mockRestore()
    }
  })

  it('preserves structured failure data for remote computer-use recovery hints', async () => {
    const server = await createOneShotServer({
      response: (requestId) => ({
        id: requestId,
        ok: false,
        error: {
          code: 'app_not_found',
          message: 'app not found: Gmail',
          data: {
            nextSteps: ['Target the desktop browser app/window that contains Gmail.']
          }
        },
        _meta: { runtimeId: 'runtime-test' }
      })
    })

    const response = await sendRemoteRuntimeRequest(
      server.pairing,
      'computer.getAppState',
      { app: 'Gmail' },
      1000
    )

    expect(response).toMatchObject({
      ok: false,
      error: {
        code: 'app_not_found',
        data: {
          nextSteps: [expect.stringContaining('desktop browser app/window')]
        }
      }
    })
  })

  it('sends orchestration authentication fields in the admitted encrypted request', async () => {
    let receivedRequest: Record<string, unknown> | null = null
    const server = await createOneShotServer({
      onRequest: (request) => {
        receivedRequest = request
      }
    })

    await sendRemoteRuntimeRequest(
      server.pairing,
      'orchestration.federationControl',
      { dispatch: 'ctx_1' },
      1000,
      {
        orchestrationCapability: 'capability',
        orchestrationContractVersion: 1,
        orchestrationRequestId: 'mutation_1'
      }
    )

    expect(receivedRequest).toMatchObject({
      method: 'orchestration.federationControl',
      params: { dispatch: 'ctx_1' },
      orchestrationCapability: 'capability',
      orchestrationContractVersion: 1,
      orchestrationRequestId: 'mutation_1'
    })
  })

  it('omits optional request fields for hosts that predate them', async () => {
    let receivedRequest: Record<string, unknown> | null = null
    const server = await createOneShotServer({
      onRequest: (request) => {
        receivedRequest = request
      }
    })

    await sendRemoteRuntimeRequest(server.pairing, 'status.get', undefined, 1000)

    expect(receivedRequest).toMatchObject({
      deviceToken: 'device-token',
      method: 'status.get'
    })
    expect(receivedRequest).not.toHaveProperty('params')
    expect(receivedRequest).not.toHaveProperty('orchestrationCapability')
    expect(receivedRequest).not.toHaveProperty('orchestrationContractVersion')
    expect(receivedRequest).not.toHaveProperty('orchestrationRequestId')
  })

  it('detaches one-shot socket listeners after a successful response', async () => {
    const offSpy = vi.spyOn(WebSocketClient.prototype, 'off')
    try {
      const server = await createOneShotServer()

      await sendRemoteRuntimeRequest<{ satisfied: boolean }>(
        server.pairing,
        'terminal.wait',
        { terminal: 't1', for: 'tui-idle', timeoutMs: 550 },
        300
      )

      const removedEvents = offSpy.mock.calls.map(([event]) => event)
      expect(removedEvents).toEqual(expect.arrayContaining(['open', 'error', 'close', 'message']))
    } finally {
      offSpy.mockRestore()
    }
  })
})

async function createSubscriptionServer(
  options: {
    sendMismatchedResponseAfterSubscribe?: boolean
    // Why: half-open simulation — the socket stays open but never answers
    // protocol pings, like a wedged tunnel that swallows frames silently.
    disableAutoPong?: boolean
  } = {}
): Promise<{
  pairing: PairingOffer
  nextBinary: Promise<Uint8Array>
  nextAuth: Promise<unknown>
}> {
  const serverKeyPair = generateKeyPair()
  let resolveBinary: (bytes: Uint8Array) => void = () => {}
  const nextBinary = new Promise<Uint8Array>((resolve) => {
    resolveBinary = resolve
  })
  let resolveAuth: (auth: unknown) => void = () => {}
  const nextAuth = new Promise<unknown>((resolve) => {
    resolveAuth = resolve
  })
  // host must match the 127.0.0.1 clients dial: a wildcard bind lets a foreign loopback listener claim the port and answer here.
  const wss = new WebSocketServer({
    host: '127.0.0.1',
    port: 0,
    autoPong: options.disableAutoPong !== true
  })
  servers.push(wss)

  wss.on('connection', (ws) => {
    let sharedKey: Uint8Array | null = null
    let authenticated = false

    ws.on('message', (data, isBinary) => {
      if (isBinary) {
        if (!sharedKey) {
          return
        }
        const plaintext = decryptBytes(new Uint8Array(data as Buffer), sharedKey)
        if (plaintext) {
          resolveBinary(plaintext)
        }
        return
      }

      const frame = data.toString()
      if (!sharedKey) {
        const hello = JSON.parse(frame) as { publicKeyB64: string }
        sharedKey = deriveSharedKey(
          serverKeyPair.secretKey,
          publicKeyFromBase64(hello.publicKeyB64)
        )
        ws.send(JSON.stringify({ type: 'e2ee_ready' }))
        return
      }

      const plaintext = decrypt(frame, sharedKey)
      if (!plaintext) {
        return
      }
      if (!authenticated) {
        resolveAuth(JSON.parse(plaintext))
        authenticated = true
        sendEncrypted(ws, sharedKey, { type: 'e2ee_authenticated' })
        return
      }

      const request = JSON.parse(plaintext) as { id: string }
      sendEncrypted(ws, sharedKey, {
        id: request.id,
        ok: true,
        streaming: true,
        result: { type: 'subscribed' },
        _meta: { runtimeId: 'runtime-test' }
      })
      if (options.sendMismatchedResponseAfterSubscribe) {
        sendEncrypted(ws, sharedKey, {
          id: `${request.id}-mismatch`,
          ok: true,
          streaming: true,
          result: { type: 'subscribed' },
          _meta: { runtimeId: 'runtime-test' }
        })
      }
    })
  })

  await new Promise<void>((resolve) => wss.once('listening', resolve))
  const address = wss.address() as AddressInfo
  const pairing = parsePairingCode(
    encodePairingOffer({
      v: 2,
      endpoint: `ws://127.0.0.1:${address.port}`,
      deviceToken: 'device-token',
      publicKeyB64: publicKeyToBase64(serverKeyPair.publicKey)
    })
  )
  if (!pairing) {
    throw new Error('Failed to create test pairing')
  }
  return { pairing, nextBinary, nextAuth }
}

function sendEncrypted(ws: WebSocket, sharedKey: Uint8Array, message: unknown): void {
  ws.send(encrypt(JSON.stringify(message), sharedKey))
}

async function createClosingServer(
  code: number,
  reason: string
): Promise<{ pairing: PairingOffer }> {
  const serverKeyPair = generateKeyPair()
  const wss = new WebSocketServer({ host: '127.0.0.1', port: 0 })
  servers.push(wss)
  wss.on('connection', (ws) => {
    ws.close(code, reason)
  })

  await new Promise<void>((resolve) => wss.once('listening', resolve))
  const address = wss.address() as AddressInfo
  const pairing = parsePairingCode(
    encodePairingOffer({
      v: 2,
      endpoint: `ws://127.0.0.1:${address.port}`,
      deviceToken: 'device-token',
      publicKeyB64: publicKeyToBase64(serverKeyPair.publicKey)
    })
  )
  if (!pairing) {
    throw new Error('Failed to create test pairing')
  }
  return { pairing }
}

async function createInvalidHandshakeServer(): Promise<{ pairing: PairingOffer }> {
  const serverKeyPair = generateKeyPair()
  const wss = new WebSocketServer({ host: '127.0.0.1', port: 0 })
  servers.push(wss)
  wss.on('connection', (ws) => {
    ws.once('message', () => ws.send(JSON.stringify({ type: 'not_orca' })))
  })

  await new Promise<void>((resolve) => wss.once('listening', resolve))
  const address = wss.address() as AddressInfo
  const pairing = parsePairingCode(
    encodePairingOffer({
      v: 2,
      endpoint: `ws://127.0.0.1:${address.port}`,
      deviceToken: 'device-token',
      publicKeyB64: publicKeyToBase64(serverKeyPair.publicKey)
    })
  )
  if (!pairing) {
    throw new Error('Failed to create test pairing')
  }
  return { pairing }
}

async function createOneShotServer(
  options: {
    response?: (requestId: string) => unknown
    onAuth?: (auth: Record<string, unknown>) => void
    onRequest?: (request: Record<string, unknown>) => void
    sendUndecryptableResponse?: boolean
  } = {}
): Promise<{ pairing: PairingOffer }> {
  const serverKeyPair = generateKeyPair()
  const wss = new WebSocketServer({ host: '127.0.0.1', port: 0 })
  servers.push(wss)

  wss.on('connection', (ws) => {
    let sharedKey: Uint8Array | null = null
    let authenticated = false

    ws.on('message', (data, isBinary) => {
      if (isBinary) {
        return
      }
      const frame = data.toString()
      if (!sharedKey) {
        const hello = JSON.parse(frame) as { publicKeyB64: string }
        sharedKey = deriveSharedKey(
          serverKeyPair.secretKey,
          publicKeyFromBase64(hello.publicKeyB64)
        )
        ws.send(JSON.stringify({ type: 'e2ee_ready' }))
        return
      }

      const plaintext = decrypt(frame, sharedKey)
      if (!plaintext) {
        return
      }
      if (!authenticated) {
        options.onAuth?.(JSON.parse(plaintext) as Record<string, unknown>)
        authenticated = true
        sendEncrypted(ws, sharedKey, { type: 'e2ee_authenticated' })
        return
      }

      const request = JSON.parse(plaintext) as { id: string } & Record<string, unknown>
      options.onRequest?.(request)
      if (options.sendUndecryptableResponse) {
        ws.send('not-an-encrypted-frame')
        return
      }
      const key = sharedKey
      const keepalive = setInterval(() => {
        sendEncrypted(ws, key, { _keepalive: true })
      }, 100)
      ws.once('close', () => clearInterval(keepalive))
      setTimeout(() => {
        clearInterval(keepalive)
        sendEncrypted(
          ws,
          key,
          options.response?.(request.id) ?? {
            id: request.id,
            ok: true,
            result: { satisfied: true },
            _meta: { runtimeId: 'runtime-test' }
          }
        )
      }, 550)
    })
  })

  await new Promise<void>((resolve) => wss.once('listening', resolve))
  const address = wss.address() as AddressInfo
  const pairing = parsePairingCode(
    encodePairingOffer({
      v: 2,
      endpoint: `ws://127.0.0.1:${address.port}`,
      deviceToken: 'device-token',
      publicKeyB64: publicKeyToBase64(serverKeyPair.publicKey)
    })
  )
  if (!pairing) {
    throw new Error('Failed to create test pairing')
  }
  return { pairing }
}