orca/src/shared/remote-runtime-shared-control-connection.test.ts

import path from 'node:path'
import { afterEach, describe, expect, it, vi } from 'vitest'
import {
  REMOTE_RUNTIME_MAX_PENDING_RPC_BYTES,
  retainedRemoteRuntimeJsonStringBytes,
  serializeRemoteRuntimeRpcRequest
} from './remote-runtime-memory-limits'
import { getRemoteRuntimeRequestAdmissionEvidence } from './remote-runtime-prepared-request-admission'
import { remoteRuntimeClientCapabilities } from './remote-runtime-client-capabilities'
import { RemoteRuntimeSharedControlConnection } from './remote-runtime-shared-control-connection'
import * as sharedControlProtocol from './remote-runtime-shared-control-protocol'
import { isRuntimeSubscriptionReplayResponse } from './runtime-subscription-replay'
import {
  closeSharedControlTestServers,
  createSharedControlTestServer as createServer
} from './remote-runtime-shared-control-test-server'

const TEST_PROJECT_PATH = path.join('tmp', 'project')

afterEach(closeSharedControlTestServers)

describe('RemoteRuntimeSharedControlConnection', () => {
  it('routes multiple one-shot RPCs over one authenticated WebSocket', async () => {
    const server = await createServer()
    const states: string[] = []
    const connection = new RemoteRuntimeSharedControlConnection(server.pairing, {
      onDiagnosticsChanged: ({ state }) => states.push(state)
    })

    const first = await connection.request('worktree.ps', undefined, 1000)
    const second = await connection.request('session.tabs.listAll', null, 1000)

    expect(first).toMatchObject({ ok: true, result: { method: 'worktree.ps' } })
    expect(second).toMatchObject({ ok: true, result: { method: 'session.tabs.listAll' } })
    expect(server.connectionCount()).toBe(1)
    expect(server.auths).toContainEqual({
      type: 'e2ee_auth',
      deviceToken: 'device-token',
      clientCapabilities: remoteRuntimeClientCapabilities()
    })
    expect(server.requests.map((request) => request.method)).toEqual([
      'worktree.ps',
      'session.tabs.listAll'
    ])
    expect((connection.close(), states)).toEqual(
      expect.arrayContaining(['awaiting_ready', 'ready', 'closed'])
    )
  })

  it('preserves orchestration authority fields on shared-control requests', async () => {
    const server = await createServer()
    const connection = new RemoteRuntimeSharedControlConnection(server.pairing)
    const envelope = {
      orchestrationCapability: 'capability',
      orchestrationContractVersion: 1,
      orchestrationRequestId: 'request-1',
      compatibilityInvocationId: 'compatibility-1',
      orchestrationCompatibilityEvidence: {
        terminalHandle: 'term-1',
        paneKey: 'pane-1',
        launchToken: 'launch-1'
      },
      id: 'forged-id',
      deviceToken: 'forged-token',
      method: 'orchestration.federationAck',
      params: { dispatchId: 'forged-dispatch' }
    }

    await connection.request('orchestration.federationPull', {}, 1000, envelope)

    expect(server.requests).toContainEqual({
      ...envelope,
      id: expect.any(String),
      deviceToken: 'device-token',
      method: 'orchestration.federationPull',
      params: {}
    })
    connection.close()
  })

  it('does not expose a binary sender on the shared control protocol surface', () => {
    expect('sendSharedControlEncryptedBinary' in sharedControlProtocol).toBe(false)
  })

  it('releases a pending request when the socket send throws', async () => {
    const connection = new RemoteRuntimeSharedControlConnection({
      v: 2,
      endpoint: 'ws://127.0.0.1:1',
      deviceToken: 'token',
      publicKeyB64: Buffer.from(new Uint8Array(32).fill(1)).toString('base64')
    })
    const unsafe = connection as unknown as {
      state: string
      ws: { readyState: number; send: () => void; close: () => void } | null
      sharedKey: Uint8Array | null
      pendingRequests: Map<string, unknown>
    }
    unsafe.state = 'ready'
    unsafe.ws = {
      readyState: 1,
      send: () => {
        throw new Error('send failed')
      },
      close: vi.fn()
    }
    unsafe.sharedKey = new Uint8Array(32).fill(2)

    await expect(connection.request('worktree.ps', undefined, 1000)).rejects.toMatchObject({
      code: 'remote_runtime_unavailable'
    })
    expect(unsafe.pendingRequests.size).toBe(0)
    expect(getRemoteRuntimeRequestAdmissionEvidence()).toEqual({
      pendingRequestCount: 0,
      retainedBytes: 0
    })
    connection.close()
  })

  it('replaces a stuck pre-ready socket when a one-shot probe proves reachability', () => {
    const connection = new RemoteRuntimeSharedControlConnection({
      v: 2,
      endpoint: 'ws://127.0.0.1:1',
      deviceToken: 'token',
      publicKeyB64: Buffer.from(new Uint8Array(32).fill(1)).toString('base64')
    })
    const close = vi.fn()
    const cleanup = vi.fn()
    const open = vi.fn()
    const unsafe = connection as unknown as {
      state: string
      ws: { readyState: number; close: () => void } | null
      socketCleanup: (() => void) | null
      open: () => void
    }
    unsafe.state = 'awaiting_ready'
    unsafe.ws = { readyState: 0, close }
    unsafe.socketCleanup = cleanup
    unsafe.open = open

    connection.reconnectNow()

    expect(cleanup).toHaveBeenCalledOnce()
    expect(close).toHaveBeenCalledOnce()
    expect(open).toHaveBeenCalledOnce()
  })

  it('keeps a waiting request alive when a reachability probe replaces its pre-ready socket', async () => {
    const server = await createServer({ suppressReadyFrameCount: 1 })
    const connection = new RemoteRuntimeSharedControlConnection(server.pairing)

    const response = connection.request('worktree.ps', undefined, 1000)
    await vi.waitFor(() => expect(server.connectionCount()).toBe(1))

    connection.reconnectNow()

    await expect(response).resolves.toMatchObject({
      ok: true,
      result: { method: 'worktree.ps' }
    })
    expect(server.connectionCount()).toBe(2)
    expect(server.requests.map(({ method }) => method)).toEqual(['worktree.ps'])
    expect(connection.getDiagnostics().pendingRequestCount).toBe(0)
    expect(getRemoteRuntimeRequestAdmissionEvidence()).toEqual({
      pendingRequestCount: 0,
      retainedBytes: 0
    })
    connection.close()
  })

  it('logs unknown response ids without breaking pending requests', async () => {
    const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined)
    const server = await createServer({ sendUnknownResponseBeforeResponse: true })
    const connection = new RemoteRuntimeSharedControlConnection(server.pairing, {
      environmentId: 'env-test'
    })

    const response = await connection.request('worktree.ps', undefined, 1000)

    expect(response).toMatchObject({ ok: true, result: { method: 'worktree.ps' } })
    expect(warn).toHaveBeenCalledWith(
      '[remote-runtime.shared-control] unknown response id',
      expect.objectContaining({
        environmentId: 'env-test',
        responseId: 'unknown-response-id',
        pendingMethods: ['worktree.ps']
      })
    )
    connection.close()
    warn.mockRestore()
  })

  it('routes multiple logical subscriptions over one socket and cleans them up explicitly', async () => {
    const server = await createServer()
    const connection = new RemoteRuntimeSharedControlConnection(server.pairing)
    const onAccounts = vi.fn()
    const onEvents = vi.fn()

    const accounts = await connection.subscribe('accounts.subscribe', null, 1000, {
      onResponse: onAccounts,
      onError: vi.fn()
    })
    await connection.subscribe('runtime.clientEvents.subscribe', null, 1000, {
      onResponse: onEvents,
      onError: vi.fn()
    })

    await vi.waitFor(() => expect(onAccounts).toHaveBeenCalled())
    await vi.waitFor(() => expect(onEvents).toHaveBeenCalled())
    accounts.close()
    await vi.waitFor(() =>
      expect(server.requests.map((request) => request.method)).toContain('accounts.unsubscribe')
    )

    expect(server.connectionCount()).toBe(1)
    expect(server.requests.map((request) => request.method)).toEqual([
      'accounts.subscribe',
      'runtime.clientEvents.subscribe',
      'accounts.unsubscribe'
    ])

    connection.close()
  })

  it('cleans up one all-session-tabs subscription by logical request id', async () => {
    const server = await createServer()
    const connection = new RemoteRuntimeSharedControlConnection(server.pairing)
    const subscription = await connection.subscribe('session.tabs.subscribeAll', null, 1000, {
      onResponse: vi.fn(),
      onError: vi.fn()
    })
    await vi.waitFor(() =>
      expect(server.requests.map((request) => request.method)).toEqual([
        'session.tabs.subscribeAll'
      ])
    )
    const subscribeRequestId = server.requests[0]!.id

    subscription.close()

    await vi.waitFor(() =>
      expect(server.requests.map((request) => request.method)).toEqual([
        'session.tabs.subscribeAll',
        'session.tabs.unsubscribeAll'
      ])
    )
    expect(server.requests[1]).toMatchObject({
      params: { subscriptionId: subscribeRequestId }
    })
    connection.close()
  })

  it('keeps many logical subscriptions on one authenticated WebSocket', async () => {
    const server = await createServer()
    const connection = new RemoteRuntimeSharedControlConnection(server.pairing)
    const subscriptions = await Promise.all(
      Array.from({ length: 35 }, (_value, index) =>
        connection.subscribe('runtime.clientEvents.subscribe', { index }, 1000, {
          onResponse: vi.fn(),
          onError: vi.fn()
        })
      )
    )

    await vi.waitFor(() => expect(server.requests).toHaveLength(35))

    expect(server.connectionCount()).toBe(1)
    expect(
      server.requests.every((request) => request.method === 'runtime.clientEvents.subscribe')
    ).toBe(true)

    subscriptions.forEach((subscription) => subscription.close())
    connection.close()
  })

  it('reconnects and replays passive subscriptions without closing them', async () => {
    const server = await createServer({ closeAfterFirstStreamingResponse: true })
    const connection = new RemoteRuntimeSharedControlConnection(server.pairing)
    const onClose = vi.fn()
    const onError = vi.fn()

    await connection.subscribe('runtime.clientEvents.subscribe', null, 1000, {
      onResponse: vi.fn(),
      onError,
      onClose
    })

    await vi.waitFor(() => expect(server.connectionCount()).toBe(2))
    await vi.waitFor(() =>
      expect(server.requests.map((request) => request.method)).toEqual([
        'runtime.clientEvents.subscribe',
        'runtime.clientEvents.subscribe'
      ])
    )
    expect(onError).toHaveBeenCalledTimes(1)
    expect(onClose).not.toHaveBeenCalled()
    expect(connection.getDiagnostics()).toMatchObject({
      state: 'ready',
      lastError: null,
      lastClose: null
    })

    connection.close()
  })

  it('keeps passive subscriptions alive after reaching the capped reconnect delay', async () => {
    const server = await createServer()
    const connection = new RemoteRuntimeSharedControlConnection(server.pairing)
    const onClose = vi.fn()

    const unsafe = connection as unknown as {
      reconnect: {
        attempt: number
        scheduleWithDefaultBackoff: (intentionallyClosed: boolean, open: () => void) => void
      }
      subscriptions: Map<string, unknown>
    }
    unsafe.reconnect.attempt = 7
    unsafe.subscriptions.set('sub-1', {
      requestId: 'sub-1',
      method: 'runtime.clientEvents.subscribe',
      params: null,
      callbacks: { onResponse: vi.fn(), onError: vi.fn(), onClose },
      sent: false,
      closed: false,
      closeAfterReady: false,
      remoteSubscriptionId: null
    })

    unsafe.reconnect.scheduleWithDefaultBackoff(false, () => {})

    expect(onClose).not.toHaveBeenCalled()
    expect(connection.getDiagnostics()).toMatchObject({
      state: 'reconnecting',
      reconnectAttempt: 8,
      subscriptionCount: 1
    })

    connection.close()
  })

  it('resets reconnect attempts after a stable authenticated ready period', async () => {
    const server = await createServer()
    const connection = new RemoteRuntimeSharedControlConnection(server.pairing, {
      reconnectStableResetMs: 50
    })

    await expect(connection.request('worktree.ps', undefined, 1000)).resolves.toMatchObject({
      ok: true
    })
    ;(connection as unknown as { reconnect: { attempt: number } }).reconnect.attempt = 3

    await vi.waitFor(() =>
      expect(connection.getDiagnostics()).toMatchObject({ reconnectAttempt: 0 })
    )
    connection.close()
  })

  it('removes ready waiters when a one-shot request times out during handshake', async () => {
    const server = await createServer({ suppressReadyFrame: true })
    const connection = new RemoteRuntimeSharedControlConnection(server.pairing)
    const unsafe = connection as unknown as {
      readyWaiters: unknown[]
      pendingRequests: Map<string, unknown>
    }

    await expect(connection.request('worktree.ps', undefined, 25)).rejects.toThrow('Timed out')

    await vi.waitFor(() => expect(unsafe.readyWaiters).toHaveLength(0))
    expect(unsafe.pendingRequests.size).toBe(0)
    connection.close()
  })

  it('cleans up an id-scoped subscription closed before its ready response', async () => {
    const server = await createServer({ delaySubscriptionReady: true })
    const connection = new RemoteRuntimeSharedControlConnection(server.pairing)
    const onAccounts = vi.fn()

    const accounts = await connection.subscribe('accounts.subscribe', null, 1000, {
      onResponse: onAccounts,
      onError: vi.fn()
    })
    await vi.waitFor(() =>
      expect(server.requests.map((request) => request.method)).toEqual(['accounts.subscribe'])
    )

    accounts.close()
    server.flushDelayedResponses()

    await vi.waitFor(() =>
      expect(server.requests.map((request) => request.method)).toEqual([
        'accounts.subscribe',
        'accounts.unsubscribe'
      ])
    )
    expect(onAccounts).not.toHaveBeenCalled()

    connection.close()
  })

  it.each([
    ['session.tabs.subscribeAll', undefined, 'session.tabs.unsubscribeAll'],
    ['runtime.clientEvents.subscribe', null, 'runtime.clientEvents.unsubscribe'],
    ['files.watch', { path: TEST_PROJECT_PATH }, 'files.unwatch']
  ])('cleans up %s explicitly on close', async (method, params, cleanupMethod) => {
    const server = await createServer()
    const connection = new RemoteRuntimeSharedControlConnection(server.pairing)
    const onResponse = vi.fn()

    const subscription = await connection.subscribe(method, params, 1000, {
      onResponse,
      onError: vi.fn()
    })
    await vi.waitFor(() => expect(onResponse).toHaveBeenCalled())

    subscription.close()

    await vi.waitFor(() =>
      expect(server.requests.map((request) => request.method)).toContain(cleanupMethod)
    )
    connection.close()
  })

  it('sends file watch cleanup at most once when a subscription closes repeatedly', async () => {
    const server = await createServer()
    const connection = new RemoteRuntimeSharedControlConnection(server.pairing)
    const onResponse = vi.fn()

    const subscription = await connection.subscribe(
      'files.watch',
      { path: TEST_PROJECT_PATH },
      1000,
      {
        onResponse,
        onError: vi.fn()
      }
    )
    await vi.waitFor(() => expect(onResponse).toHaveBeenCalled())

    subscription.close()
    subscription.close()

    await vi.waitFor(() =>
      expect(server.requests.filter((request) => request.method === 'files.unwatch')).toHaveLength(
        1
      )
    )
    connection.close()
  })

  it('ignores encrypted keepalive frames while waiting for a response', async () => {
    const server = await createServer({ sendKeepaliveBeforeResponse: true })
    const connection = new RemoteRuntimeSharedControlConnection(server.pairing)

    await expect(connection.request('worktree.ps', undefined, 1000)).resolves.toMatchObject({
      ok: true,
      result: { method: 'worktree.ps' }
    })

    connection.close()
  })

  it('times out a stuck short RPC on its absolute deadline despite keepalive frames', async () => {
    // Why: a keepalive on the shared socket is armed by an unrelated long-poll,
    // not by this request. It must NOT extend a stuck short RPC's deadline —
    // otherwise a hung server call hangs the caller forever (#7948).
    const server = await createServer({
      silentMethods: ['worktree.hang'],
      sendKeepaliveBeforeResponse: true,
      keepaliveDelayMs: 20
    })
    const connection = new RemoteRuntimeSharedControlConnection(server.pairing)

    await expect(connection.request('worktree.hang', undefined, 60)).rejects.toThrow('Timed out')

    connection.close()
  })

  it('sends explicit subscription cleanup before graceful connection close', async () => {
    const server = await createServer()
    const connection = new RemoteRuntimeSharedControlConnection(server.pairing)
    const onResponse = vi.fn()

    await connection.subscribe('runtime.clientEvents.subscribe', null, 1000, {
      onResponse,
      onError: vi.fn()
    })
    await vi.waitFor(() => expect(onResponse).toHaveBeenCalled())

    connection.close()

    await vi.waitFor(() =>
      expect(server.requests.map((request) => request.method)).toContain(
        'runtime.clientEvents.unsubscribe'
      )
    )
  })

  it('treats remote binary frames as unsupported on the shared control lane', async () => {
    const server = await createServer({ sendBinaryAfterAuth: true })
    const connection = new RemoteRuntimeSharedControlConnection(server.pairing)

    await expect(connection.request('worktree.ps', undefined, 1000)).rejects.toThrow(
      'unexpected binary frame'
    )

    connection.close()
  })

  it('does not send outbound binary frames on the shared control lane', async () => {
    const server = await createServer()
    const connection = new RemoteRuntimeSharedControlConnection(server.pairing)

    const subscription = await connection.subscribe('runtime.clientEvents.subscribe', null, 1000, {
      onResponse: vi.fn(),
      onError: vi.fn()
    })
    await vi.waitFor(() => expect(server.requests).toHaveLength(1))

    expect(subscription.sendBinary(new Uint8Array([1, 2, 3]))).toBe(false)
    connection.close()
  })

  it('detects a half-open socket via client liveness, reconnects, and tags the replayed response', async () => {
    // Why: the server keeps the TCP connection open but stops answering —
    // no close frame, no pongs (autoPong disabled), no responses. This is the
    // half-open devtunnel scenario from #7718: edge-triggered reconnect never
    // fires, so client liveness must terminate the socket itself.
    const server = await createServer({ disableAutoPong: true })
    const connection = new RemoteRuntimeSharedControlConnection(server.pairing, {
      liveness: { pingIntervalMs: 50, livenessTimeoutMs: 200 }
    })
    const onResponse = vi.fn()
    const onClose = vi.fn()

    await connection.subscribe('runtime.clientEvents.subscribe', null, 1000, {
      onResponse,
      onError: vi.fn(),
      onClose
    })
    await vi.waitFor(() => expect(onResponse).toHaveBeenCalled())
    expect(isRuntimeSubscriptionReplayResponse(onResponse.mock.calls[0]?.[0])).toBe(false)

    // Liveness terminates the silent socket and the reconnect path replays
    // the subscription on a fresh connection.
    await vi.waitFor(() => expect(server.connectionCount()).toBeGreaterThanOrEqual(2), {
      timeout: 5000
    })
    await vi.waitFor(
      () =>
        expect(
          server.requests.filter((request) => request.method === 'runtime.clientEvents.subscribe')
            .length
        ).toBeGreaterThanOrEqual(2),
      { timeout: 5000 }
    )
    // The first response after the reconnect replay carries the replay tag so
    // snapshot freshness gates can accept the re-emitted snapshot.
    await vi.waitFor(
      () =>
        expect(
          onResponse.mock.calls.some(([response]) => isRuntimeSubscriptionReplayResponse(response))
        ).toBe(true),
      { timeout: 5000 }
    )
    expect(onClose).not.toHaveBeenCalled()

    connection.close()
  })

  it('keeps unrelated pending requests alive when one request times out', async () => {
    const server = await createServer({
      silentMethods: ['worktree.hang'],
      delayedMethods: ['worktree.ps']
    })
    const connection = new RemoteRuntimeSharedControlConnection(server.pairing)
    const unsafe = connection as unknown as {
      pendingRequests: Map<
        string,
        {
          method: string
          preparedRequest?: { retainedBytes: number; serializedRequest: string | null } | null
        }
      >
    }

    const timedOut = connection.request('worktree.hang', undefined, 250)
    void timedOut.catch(() => undefined)
    await vi.waitFor(() =>
      expect(server.requests.map(({ method }) => method)).toContain('worktree.hang')
    )
    const survivor = connection.request('worktree.ps', undefined, 1000).then(
      (response) => ({ ok: true as const, response }),
      (error: unknown) => ({ ok: false as const, error })
    )
    await vi.waitFor(() =>
      expect(server.requests.map(({ method }) => method)).toContain('worktree.ps')
    )
    expect(
      Array.from(unsafe.pendingRequests.values()).every(
        (pending) =>
          pending.preparedRequest?.serializedRequest === null &&
          pending.preparedRequest.retainedBytes > 0
      )
    ).toBe(true)
    expect(getRemoteRuntimeRequestAdmissionEvidence().pendingRequestCount).toBe(2)

    await expect(timedOut).rejects.toThrow('Timed out')
    // Why: a single slow method is not evidence that a shared socket is dead;
    // liveness monitoring owns connection-wide failure detection.
    expect(connection.getDiagnostics()).toMatchObject({ state: 'ready', pendingRequestCount: 1 })
    expect(getRemoteRuntimeRequestAdmissionEvidence().pendingRequestCount).toBe(1)

    server.flushDelayedResponses()
    await expect(survivor).resolves.toMatchObject({
      ok: true,
      response: { ok: true, result: { method: 'worktree.ps' } }
    })
    expect(unsafe.pendingRequests.size).toBe(0)
    expect(getRemoteRuntimeRequestAdmissionEvidence()).toEqual({
      pendingRequestCount: 0,
      retainedBytes: 0
    })
    expect(server.connectionCount()).toBe(1)

    connection.close()
  })

  it('keeps sent request bytes admitted while a ready socket stops responding', async () => {
    const server = await createServer({ silentMethods: ['worktree.large'] })
    const connection = new RemoteRuntimeSharedControlConnection(server.pairing)
    const params = { value: 'x'.repeat(3 * 1024 * 1024) }
    const retainedBytes = retainedRemoteRuntimeJsonStringBytes(
      serializeRemoteRuntimeRpcRequest({
        requestId: '00000000-0000-4000-8000-000000000000',
        deviceToken: server.pairing.deviceToken,
        method: 'worktree.large',
        params
      })
    )
    const admittedCount = Math.floor(REMOTE_RUNTIME_MAX_PENDING_RPC_BYTES / retainedBytes)
    const pendingRequests = (
      connection as unknown as {
        pendingRequests: Map<
          string,
          { preparedRequest?: { serializedRequest: string | null } | null }
        >
      }
    ).pendingRequests
    const requests = Array.from({ length: admittedCount }, () =>
      connection.request('worktree.large', params, 60_000).catch(() => undefined)
    )
    await vi.waitFor(() => expect(server.requests).toHaveLength(admittedCount))

    expect(
      Array.from(pendingRequests.values()).every(
        (pending) => pending.preparedRequest?.serializedRequest === null
      )
    ).toBe(true)
    await expect(connection.request('worktree.large', params, 60_000)).rejects.toMatchObject({
      code: 'remote_runtime_busy'
    })
    expect(getRemoteRuntimeRequestAdmissionEvidence().retainedBytes).toBeLessThanOrEqual(
      REMOTE_RUNTIME_MAX_PENDING_RPC_BYTES
    )

    connection.close()
    await Promise.all(requests)
    expect(getRemoteRuntimeRequestAdmissionEvidence()).toEqual({
      pendingRequestCount: 0,
      retainedBytes: 0
    })
  })

  it('rejects pending requests and schedules standing recovery when the socket closes', async () => {
    const server = await createServer({ closeBeforeResponse: true })
    const connection = new RemoteRuntimeSharedControlConnection(server.pairing)

    await expect(connection.request('worktree.ps', undefined, 1000)).rejects.toThrow(
      'Remote Orca runtime closed the connection'
    )
    expect(connection.getDiagnostics()).toMatchObject({
      state: 'reconnecting',
      pendingRequestCount: 0,
      lastClose: { code: 4001, reason: 'test close' }
    })
    connection.pauseStandingRetry()
    expect(connection.getDiagnostics()).toMatchObject({ state: 'closed' })
    connection.close()
  })
})