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

import {
  pairedRuntimeConnectionFromOffer,
  type PairedRuntimeConnection
} from './paired-runtime-connection'
import type { AddressInfo } from 'node:net'
import { afterEach, describe, expect, it, vi } from 'vitest'
import WebSocketClient, { WebSocketServer, type WebSocket } from 'ws'
import { encodePairingOffer, parsePairingCode } 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_NETWORK_TUNNEL_RUNTIME_CAPABILITY } 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('closes a half-open subscription socket via client liveness so callers can resubscribe', async () => {
    // Half-open dedicated streams get no close/pongs/data; client liveness must trigger onError/onClose so renderer resubscription runs (#7718).
    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.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(
          pairedRuntimeConnectionFromOffer({ endpoint: '', deviceToken: '', publicKeyB64: '' }),
          '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.each(['refused', 'timeout'])(
    'observes a slow restart response finitely (%s)',
    async (mode) => {
      const observed = Promise.withResolvers<void>()
      const requests = vi.fn((_request: Record<string, unknown>) => observed.resolve())
      const close = vi.spyOn(WebSocketClient.prototype, 'close')
      const server = await createOneShotServer({
        onRequest: requests,
        keepalive: false,
        responseDelayMs: mode === 'refused' ? 20_000 : null,
        response: (id) => ({
          id,
          ok: true,
          result: { continued: [{ sessionId: 'a', outcome: 'refused' }] },
          _meta: { runtimeId: 'runtime-test' }
        })
      })
      vi.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout', 'setInterval', 'clearInterval'] })
      const timers = vi.spyOn(globalThis, 'setTimeout')
      const clearTimer = vi.spyOn(globalThis, 'clearTimeout')
      try {
        const settled = vi.fn()
        const request = sendRemoteRuntimeRequest(
          server.pairing,
          'agentSession.restartContinue',
          { sessionIds: ['a'] },
          600_000
        )
        const result = request.catch((error: unknown) => error)
        void result.then(settled)
        await observed.promise
        await vi.advanceTimersByTimeAsync(15_001)
        expect(settled).not.toHaveBeenCalled()
        await vi.advanceTimersByTimeAsync(mode === 'refused' ? 4_999 : 584_998)
        if (mode === 'refused') {
          await expect(result).resolves.toMatchObject({
            ok: true,
            result: { continued: [{ sessionId: 'a', outcome: 'refused' }] }
          })
        } else {
          expect(settled).not.toHaveBeenCalled()
          await vi.advanceTimersByTimeAsync(1)
          await expect(result).resolves.toMatchObject({ code: 'runtime_timeout' })
        }
        expect(requests).toHaveBeenCalledOnce()
        expect(requests.mock.calls[0]?.[0]).toMatchObject({
          method: 'agentSession.restartContinue',
          params: { sessionIds: ['a'] }
        })
        expect(close).toHaveBeenCalled()
        const timeoutIndex = timers.mock.calls.findIndex(([, delay]) => delay === 600_000)
        expect(timeoutIndex).toBeGreaterThanOrEqual(0)
        expect(clearTimer).toHaveBeenCalledWith(timers.mock.results[timeoutIndex].value)
      } finally {
        timers.mockRestore()
        clearTimer.mockRestore()
        vi.useRealTimers()
        close.mockRestore()
      }
    }
  )

  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')
  })
})

async function createSubscriptionServer(
  options: {
    sendMismatchedResponseAfterSubscribe?: boolean
    // Simulate half-open tunnels by leaving sockets open without pong replies.
    disableAutoPong?: boolean
  } = {}
): Promise<{
  pairing: PairedRuntimeConnection
  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: pairedRuntimeConnectionFromOffer(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: PairedRuntimeConnection }> {
  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: pairedRuntimeConnectionFromOffer(pairing) }
}

async function createInvalidHandshakeServer(): Promise<{ pairing: PairedRuntimeConnection }> {
  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: pairedRuntimeConnectionFromOffer(pairing) }
}

async function createOneShotServer(
  options: {
    response?: (requestId: string) => unknown
    onAuth?: (auth: Record<string, unknown>) => void
    onRequest?: (request: Record<string, unknown>) => void
    sendUndecryptableResponse?: boolean
    keepalive?: boolean
    responseDelayMs?: number | null
  } = {}
): Promise<{ pairing: PairedRuntimeConnection }> {
  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
      }
      if (options.responseDelayMs === null) {
        return
      }
      const key = sharedKey
      const keepalive =
        options.keepalive === false
          ? undefined
          : 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' }
          }
        )
      }, options.responseDelayMs ?? 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: pairedRuntimeConnectionFromOffer(pairing) }
}