From 34948dac686501797c0ed9d4aef043504cd9cb2e Mon Sep 17 00:00:00 2001 From: Jinwoo-H Date: Mon, 13 Jul 2026 02:27:08 -0400 Subject: [PATCH] feat(mobile): add stable relay RPC sessions Co-authored-by: Orca --- mobile/src/transport/client-context.test.ts | 2 +- mobile/src/transport/client-context.tsx | 10 +- mobile/src/transport/host-logical-client.ts | 12 + .../src/transport/mobile-relay-e2ee-link.ts | 151 ++++++++++ .../transport/mobile-relay-physical-client.ts | 8 +- .../mobile-relay-rpc-session.test.ts | 207 ++++++++++++++ .../src/transport/mobile-relay-rpc-session.ts | 244 ++++++++++++++++ .../src/transport/mobile-relay-rpc-streams.ts | 165 +++++++++++ .../stable-logical-rpc-client.test.ts | 135 +++++++++ .../transport/stable-logical-rpc-client.ts | 260 ++++++++++++++++++ 10 files changed, 1180 insertions(+), 14 deletions(-) create mode 100644 mobile/src/transport/host-logical-client.ts create mode 100644 mobile/src/transport/mobile-relay-e2ee-link.ts create mode 100644 mobile/src/transport/mobile-relay-rpc-session.test.ts create mode 100644 mobile/src/transport/mobile-relay-rpc-session.ts create mode 100644 mobile/src/transport/mobile-relay-rpc-streams.ts create mode 100644 mobile/src/transport/stable-logical-rpc-client.test.ts create mode 100644 mobile/src/transport/stable-logical-rpc-client.ts diff --git a/mobile/src/transport/client-context.test.ts b/mobile/src/transport/client-context.test.ts index b2ed60e74e4..4a7d286a2e4 100644 --- a/mobile/src/transport/client-context.test.ts +++ b/mobile/src/transport/client-context.test.ts @@ -131,7 +131,7 @@ describe('useHostClient', () => { loadHostsMock.mockResolvedValue([HOST]) const harness = await renderHarness(HOST.id) - expect(harness.hook.client).toBe(fake) + expect(harness.hook.client).not.toBeNull() expect(harness.hook.state).toBe('connected') // Regression (STA-1511): closeHost deletes the entry; before the fix the diff --git a/mobile/src/transport/client-context.tsx b/mobile/src/transport/client-context.tsx index df54f143528..6cb06317350 100644 --- a/mobile/src/transport/client-context.tsx +++ b/mobile/src/transport/client-context.tsx @@ -19,11 +19,12 @@ import { useState, type ReactNode } from 'react' -import { connect, type RpcClient } from './rpc-client' +import type { RpcClient } from './rpc-client' import { connectionLogStore } from './connection-log-buffer' import { subscribeConnectionRevivalTriggers } from './connection-revival-triggers' import { HostClientOpenRegistry } from './host-client-open-registry' import { loadHosts } from './host-store' +import { openHostLogicalClient } from './host-logical-client' import type { ConnectionState, HostProfile } from './types' type StoreEntry = { @@ -155,12 +156,7 @@ export function RpcClientProvider({ children }: { children: ReactNode }) { let client: RpcClient try { - client = connect(host.endpoint, host.deviceToken, host.publicKeyB64, { - // Why: retain reconnect lifecycle events for the Connection Log - // screen — without this the reasons a host is stuck live only in - // console.log, which users can't see or share. - onLog: (entry) => connectionLogStore.append(hostId, entry) - }) + client = openHostLogicalClient(host, (entry) => connectionLogStore.append(hostId, entry)) } catch { // Why: connect() can throw synchronously if the public key is // malformed or the endpoint URL is invalid. Notify so the UI diff --git a/mobile/src/transport/host-logical-client.ts b/mobile/src/transport/host-logical-client.ts new file mode 100644 index 00000000000..2926234c6b8 --- /dev/null +++ b/mobile/src/transport/host-logical-client.ts @@ -0,0 +1,12 @@ +import { connect, type RpcClient } from './rpc-client' +import { createStableLogicalRpcClient } from './stable-logical-rpc-client' +import type { ConnectionLogSink, HostProfile } from './types' + +export function openHostLogicalClient(host: HostProfile, onLog: ConnectionLogSink): RpcClient { + // Why: the stable facade owns app-visible RPC/subscription state while the + // direct socket remains a replaceable first physical generation. + return createStableLogicalRpcClient( + connect(host.endpoint, host.deviceToken, host.publicKeyB64, { onLog }), + 'direct' + ) +} diff --git a/mobile/src/transport/mobile-relay-e2ee-link.ts b/mobile/src/transport/mobile-relay-e2ee-link.ts new file mode 100644 index 00000000000..183a8deea6d --- /dev/null +++ b/mobile/src/transport/mobile-relay-e2ee-link.ts @@ -0,0 +1,151 @@ +import { + RelayPhoneHelloSchema, + type RelayPhoneHello +} from '../../../src/shared/mobile-relay-phone-protocol' +import { MobileE2EEV2ClientSession } from './mobile-e2ee-v2-client-session' +import { MobileE2EEV2PhysicalChannel } from './mobile-e2ee-v2-physical-channel' +import { websocketPayloadToUint8 } from './websocket-payload-bytes' + +export class RelayOuterError extends Error { + constructor(readonly code: number) { + super(`relay_outer_${code}`) + } +} + +type MobileRelayE2eeLinkOptions = { + endpoint: { cellUrl: string; relayHostId: string } + credential: string + expectedCredentialKind: 'invite' | 'resume' + deviceToken: string + desktopPublicKeyB64: string + onAuthenticated: () => void + onText: (plaintext: string) => void + onBinary: (plaintext: Uint8Array) => void + onHello?: (hello: Extract) => void + onError: (error: Error) => void + createSocket?: (url: string) => WebSocket +} + +export class MobileRelayE2eeLink { + private readonly options: MobileRelayE2eeLinkOptions + private readonly socket: WebSocket + private readonly channel: MobileE2EEV2PhysicalChannel + private outerReady = false + private closed = false + private inboundChain: Promise = Promise.resolve() + + constructor(options: MobileRelayE2eeLinkOptions) { + this.options = options + this.socket = (options.createSocket ?? ((url) => new WebSocket(url)))( + relaySocketUrl(options.endpoint) + ) + const session = MobileE2EEV2ClientSession.create({ + desktopPublicKeyB64: options.desktopPublicKeyB64, + transport: 'relay', + relayHostId: options.endpoint.relayHostId + }) + this.channel = new MobileE2EEV2PhysicalChannel({ + session, + socket: this.socket, + deviceToken: options.deviceToken, + decodeBinary: websocketPayloadToUint8, + onAuthenticated: options.onAuthenticated, + onText: options.onText, + onBinary: options.onBinary, + onError: (error) => this.fail(error) + }) + this.bindSocket() + } + + sendText(plaintext: string): boolean { + return !this.closed && this.channel.sendText(plaintext) + } + + sendBinary(plaintext: Uint8Array): boolean { + return !this.closed && this.channel.sendBinary(plaintext) + } + + close(): void { + if (this.closed) { + return + } + this.closed = true + this.channel.dispose() + this.socket.close() + } + + private bindSocket(): void { + this.socket.onopen = () => { + this.socket.send( + JSON.stringify({ + type: 'relay-auth', + v: 1, + mode: 'connect', + credential: this.options.credential + }) + ) + } + this.socket.onmessage = (event) => { + this.inboundChain = this.inboundChain + .then(async () => { + if (this.closed) { + return + } + if (!this.outerReady) { + this.acceptHello(event.data) + } else { + await this.channel.handleMessage(event.data) + } + }) + .catch((error: unknown) => this.fail(asError(error))) + } + this.socket.onerror = () => this.fail(new Error('relay transport error')) + this.socket.onclose = (event) => this.fail(new RelayOuterError(event.code || 1006)) + } + + private acceptHello(raw: unknown): void { + if (typeof raw !== 'string') { + throw new Error('expected plaintext relay hello') + } + let value: unknown + try { + value = JSON.parse(raw) + } catch { + throw new Error('invalid relay hello JSON') + } + const parsed = RelayPhoneHelloSchema.safeParse(value) + if (!parsed.success) { + throw new Error('invalid relay hello') + } + if (!parsed.data.ok) { + throw new RelayOuterError(parsed.data.code) + } + if (parsed.data.credentialKind !== this.options.expectedCredentialKind) { + throw new Error('relay credential resolved as an unexpected credential kind') + } + this.outerReady = true + this.options.onHello?.(parsed.data) + this.channel.start() + } + + private fail(error: Error): void { + if (this.closed) { + return + } + this.closed = true + this.channel.dispose() + this.options.onError(error) + this.socket.close() + } +} + +function relaySocketUrl(endpoint: { cellUrl: string; relayHostId: string }): string { + const url = new URL(endpoint.cellUrl) + url.protocol = 'wss:' + url.pathname = `/v1/connect/${encodeURIComponent(endpoint.relayHostId)}` + return url.toString() +} + +function asError(error: unknown): Error { + return error instanceof Error ? error : new Error(String(error)) +} diff --git a/mobile/src/transport/mobile-relay-physical-client.ts b/mobile/src/transport/mobile-relay-physical-client.ts index 9c12b786eb5..acf264b8c73 100644 --- a/mobile/src/transport/mobile-relay-physical-client.ts +++ b/mobile/src/transport/mobile-relay-physical-client.ts @@ -5,6 +5,8 @@ import { MobileE2EEV2PhysicalChannel } from './mobile-e2ee-v2-physical-channel' import { isRpcResponse } from './rpc-response-shape' import type { RpcResponse } from './types' import { websocketPayloadToUint8 } from './websocket-payload-bytes' +export { RelayOuterError } from './mobile-relay-e2ee-link' +import { RelayOuterError } from './mobile-relay-e2ee-link' type PendingRequest = { resolve: (response: RpcResponse) => void @@ -17,12 +19,6 @@ export type PairingCandidateClient = { close(): void } -export class RelayOuterError extends Error { - constructor(readonly code: number) { - super(`relay_outer_${code}`) - } -} - export function connectMobileRelayForPairing(args: { relay: PairingRelay deviceToken: string diff --git a/mobile/src/transport/mobile-relay-rpc-session.test.ts b/mobile/src/transport/mobile-relay-rpc-session.test.ts new file mode 100644 index 00000000000..9fa88dca9b8 --- /dev/null +++ b/mobile/src/transport/mobile-relay-rpc-session.test.ts @@ -0,0 +1,207 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' +import { + BrowserScreencastOpcode, + encodeBrowserScreencastFrame +} from '../../../src/shared/browser-screencast-protocol' +import { encodeTerminalStreamFrame, TerminalStreamOpcode } from './terminal-stream-protocol' + +const fakes = vi.hoisted(() => ({ + linkOptions: null as null | { + endpoint: { cellUrl: string; relayHostId: string } + credential: string + expectedCredentialKind: string + onHello(value: unknown): void + onAuthenticated(): void + onText(value: string): void + onBinary(value: Uint8Array): void + onError(error: Error): void + }, + sendText: vi.fn(() => true), + close: vi.fn() +})) + +vi.mock('./mobile-relay-e2ee-link', () => ({ + MobileRelayE2eeLink: class { + constructor(options: NonNullable) { + fakes.linkOptions = options + } + sendText = fakes.sendText + close = fakes.close + } +})) + +import { connectMobileRelayRpcSession } from './mobile-relay-rpc-session' + +const relay = { + v: 1 as const, + directorUrl: 'https://relay.onorca.dev', + cellUrl: 'https://relay-c1.onorca.dev', + assignmentEpoch: 7, + relayHostId: 'AbCdEf0123_-xyZ9', + e2eeFraming: 2 as const +} + +function openSession() { + return connectMobileRelayRpcSession({ + relay, + resumeToken: 'resume-secret', + resumeCredentialVersion: 3, + resumeConfirmReqId: 'confirm-1', + deviceToken: 'device-token', + desktopPublicKeyB64: 'AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=', + requestTimeoutMs: 1000 + }) +} + +async function authenticateSession() { + const session = openSession() + fakes.linkOptions!.onHello({ + type: 'relay-hello', + ok: true, + credentialKind: 'resume', + leaseExpiresAt: Date.now() + 60_000, + acceptedCredentialVersion: 3, + acceptedAs: 'current', + resumeExpiresAt: Date.now() + 300_000 + }) + expect(session.getState()).toBe('handshaking') + fakes.linkOptions!.onAuthenticated() + await vi.waitFor(() => expect(fakes.sendText).toHaveBeenCalledOnce()) + const request = JSON.parse(fakes.sendText.mock.calls[0]![0] as string) as { + id: string + method: string + params: unknown + } + fakes.linkOptions!.onText( + JSON.stringify({ + id: request.id, + ok: true, + result: { + v: 1, + relay, + resumeConfirmation: { + v: 1, + reqId: 'confirm-1', + currentVersion: 3, + acceptedAs: 'current', + renewed: true, + resumeExpiresAt: Date.now() + 300_000 + } + }, + _meta: { runtimeId: 'runtime-1' } + }) + ) + await vi.waitFor(() => expect(session.getState()).toBe('connected')) + fakes.sendText.mockClear() + return { session, confirmationRequest: request } +} + +describe('mobile relay RPC session', () => { + beforeEach(() => { + vi.clearAllMocks() + fakes.linkOptions = null + fakes.sendText.mockReturnValue(true) + }) + + it('requires exact resume observations and confirms by request ID before becoming connected', async () => { + const { session, confirmationRequest } = await authenticateSession() + + expect(fakes.linkOptions).toMatchObject({ + endpoint: relay, + credential: 'resume-secret', + expectedCredentialKind: 'resume' + }) + expect(confirmationRequest).toMatchObject({ + method: 'pairing.getEndpoints', + params: { resumeConfirmReqId: 'confirm-1' }, + deviceToken: 'device-token' + }) + expect(confirmationRequest.params).not.toHaveProperty('relayDeviceId') + expect(confirmationRequest.params).not.toHaveProperty('acceptedCredentialVersion') + expect(session.getLeaseExpiresAt()).toEqual(expect.any(Number)) + }) + + it('rejects a mismatched outer credential version and closes the physical link', () => { + const session = openSession() + fakes.linkOptions!.onHello({ + type: 'relay-hello', + ok: true, + credentialKind: 'resume', + leaseExpiresAt: Date.now() + 60_000, + acceptedCredentialVersion: 2, + acceptedAs: 'grace', + resumeExpiresAt: Date.now() + 300_000 + }) + + expect(session.getState()).toBe('disconnected') + expect(fakes.close).toHaveBeenCalledOnce() + expect(fakes.sendText).not.toHaveBeenCalled() + }) + + it('routes terminal and browser binary streams after confirmation', async () => { + const { session } = await authenticateSession() + const terminalListener = vi.fn() + session.subscribe('terminal.subscribe', { terminal: 'term-1' }, terminalListener) + await vi.waitFor(() => expect(fakes.sendText).toHaveBeenCalledOnce()) + const terminalRequest = JSON.parse(fakes.sendText.mock.calls[0]![0] as string) as { + id: string + } + fakes.linkOptions!.onText( + JSON.stringify({ + id: terminalRequest.id, + ok: true, + result: { streamId: 42 }, + _meta: { runtimeId: 'runtime-1' } + }) + ) + fakes.linkOptions!.onBinary( + encodeTerminalStreamFrame({ + opcode: TerminalStreamOpcode.Output, + streamId: 42, + seq: 1, + payload: new TextEncoder().encode('hello') + }) + ) + expect(terminalListener).toHaveBeenLastCalledWith({ + type: 'data', + streamId: 42, + chunk: 'hello' + }) + + fakes.sendText.mockClear() + const onBinaryFrame = vi.fn() + session.subscribe('browser.screencast', {}, vi.fn(), { onBinaryFrame }) + await vi.waitFor(() => expect(fakes.sendText).toHaveBeenCalledOnce()) + const browserRequest = JSON.parse(fakes.sendText.mock.calls[0]![0] as string) as { id: string } + fakes.linkOptions!.onText( + JSON.stringify({ + id: browserRequest.id, + ok: true, + result: { subscriptionId: 'browser-1' }, + _meta: { runtimeId: 'runtime-1' } + }) + ) + fakes.linkOptions!.onBinary( + encodeBrowserScreencastFrame({ + opcode: BrowserScreencastOpcode.Frame, + seq: 9, + format: 'jpeg', + metadata: { imageWidth: 800 }, + image: new Uint8Array([1, 2, 3]) + }) + ) + expect(onBinaryFrame).toHaveBeenCalledWith( + expect.objectContaining({ seq: 9, format: 'jpeg', image: new Uint8Array([1, 2, 3]) }) + ) + }) + + it('rejects pending RPC work when the physical link fails', async () => { + const { session } = await authenticateSession() + const pending = session.sendRequest('status.get') + await vi.waitFor(() => expect(fakes.sendText).toHaveBeenCalledOnce()) + fakes.linkOptions!.onError(new Error('relay transport error')) + + await expect(pending).rejects.toThrow('relay transport error') + expect(session.getState()).toBe('disconnected') + }) +}) diff --git a/mobile/src/transport/mobile-relay-rpc-session.ts b/mobile/src/transport/mobile-relay-rpc-session.ts new file mode 100644 index 00000000000..6755aa32c1d --- /dev/null +++ b/mobile/src/transport/mobile-relay-rpc-session.ts @@ -0,0 +1,244 @@ +import { + PairingGetEndpointsResultSchema, + type MobileRelayEndpoint +} from '../../../src/shared/mobile-relay-credential-contract' +import { MobileRelayE2eeLink } from './mobile-relay-e2ee-link' +import { MobileRelayRpcStreams } from './mobile-relay-rpc-streams' +import { isRpcResponse } from './rpc-response-shape' +import type { RpcClient } from './rpc-client' +import type { ConnectionState, RpcResponse } from './types' + +type PendingRequest = { + resolve: (response: RpcResponse) => void + reject: (error: Error) => void + timer: ReturnType +} + +export type MobileRelayRpcSession = RpcClient & { + getLeaseExpiresAt(): number | null +} + +export function connectMobileRelayRpcSession(args: { + relay: MobileRelayEndpoint + resumeToken: string + resumeCredentialVersion: number + resumeConfirmReqId: string + deviceToken: string + desktopPublicKeyB64: string + requestTimeoutMs?: number + createSocket?: (url: string) => WebSocket +}): MobileRelayRpcSession { + const requestTimeoutMs = args.requestTimeoutMs ?? 30_000 + const pending = new Map() + const stateListeners = new Set<(state: ConnectionState) => void>() + let state: ConnectionState = 'connecting' + let requestCounter = 0 + let lastConnectedAt: number | null = null + let leaseExpiresAt: number | null = null + let closed = false + const streams = new MobileRelayRpcStreams({ + nextId, + sendFrame, + waitForConnected: () => waitForConnected() + }) + + const link = new MobileRelayE2eeLink({ + endpoint: args.relay, + credential: args.resumeToken, + expectedCredentialKind: 'resume', + deviceToken: args.deviceToken, + desktopPublicKeyB64: args.desktopPublicKeyB64, + createSocket: args.createSocket, + onHello: (hello) => { + if ( + hello.credentialKind !== 'resume' || + hello.acceptedCredentialVersion !== args.resumeCredentialVersion + ) { + fail(new Error('relay resume credential version mismatch')) + return + } + leaseExpiresAt = hello.leaseExpiresAt + publishState('handshaking') + }, + onAuthenticated: () => void confirmResume(), + onText: handleText, + onBinary: handleBinary, + onError: fail + }) + + const client: MobileRelayRpcSession = { + async sendRequest(method, params, options) { + await waitForConnected(options?.timeoutMs) + return sendRpc(method, params, options?.timeoutMs) + }, + + subscribe(method, params, listener, options) { + if (closed) { + return () => {} + } + return streams.subscribe(method, params, listener, options) + }, + + updateTerminalSubscriptionViewport(terminal, viewport) { + streams.updateTerminalViewport(terminal, viewport) + }, + getState: () => state, + getReconnectAttempt: () => 0, + getLastConnectedAt: () => lastConnectedAt, + onStateChange(listener) { + stateListeners.add(listener) + return () => stateListeners.delete(listener) + }, + notifyForeground: () => {}, + close() { + if (closed) { + return + } + closed = true + link.close() + rejectPending(new Error('Client closed')) + streams.clear() + publishState('disconnected') + }, + getLeaseExpiresAt: () => leaseExpiresAt + } + return client + + async function confirmResume(): Promise { + try { + const response = await sendRpc( + 'pairing.getEndpoints', + { resumeConfirmReqId: args.resumeConfirmReqId }, + requestTimeoutMs, + true + ) + if (!response.ok) { + throw new Error(response.error.code) + } + const result = PairingGetEndpointsResultSchema.parse(response.result) + if (!result.resumeConfirmation || result.relay?.relayHostId !== args.relay.relayHostId) { + throw new Error('relay resume confirmation missing') + } + lastConnectedAt = Date.now() + publishState('connected') + } catch (error) { + fail(asError(error)) + } + } + + function sendRpc( + method: string, + params: unknown, + timeoutMs = requestTimeoutMs, + beforeConnected = false + ): Promise { + if (closed || (!beforeConnected && state !== 'connected')) { + return Promise.reject(new Error('relay session not connected')) + } + const id = nextId() + return new Promise((resolve, reject) => { + const timer = setTimeout(() => { + pending.delete(id) + reject(new Error(`relay RPC timed out: ${method}`)) + }, timeoutMs) + pending.set(id, { resolve, reject, timer }) + if (!sendFrame({ id, method, params })) { + clearTimeout(timer) + pending.delete(id) + reject(new Error('relay E2EE channel not ready')) + } + }) + } + + function sendFrame(request: { id: string; method: string; params?: unknown }): boolean { + return link.sendText(JSON.stringify({ ...request, deviceToken: args.deviceToken })) + } + + function handleText(plaintext: string): void { + let value: unknown + try { + value = JSON.parse(plaintext) + } catch { + return + } + if (!isRpcResponse(value)) { + return + } + const request = pending.get(value.id) + if (request) { + clearTimeout(request.timer) + pending.delete(value.id) + request.resolve(value) + return + } + streams.handleResponse(value) + } + + function handleBinary(bytes: Uint8Array): void { + streams.handleBinary(bytes) + } + + function waitForConnected(timeoutMs = requestTimeoutMs): Promise { + if (state === 'connected') { + return Promise.resolve() + } + return new Promise((resolve, reject) => { + let timer: ReturnType | null = null + const unsubscribe = client.onStateChange((next) => { + if (next === 'connected') { + finish() + resolve() + } else if (next === 'disconnected' || next === 'auth-failed') { + finish() + reject(new Error(`relay session ${next}`)) + } + }) + timer = setTimeout(() => { + finish() + reject(new Error('relay session connection timed out')) + }, timeoutMs) + function finish(): void { + if (timer) { + clearTimeout(timer) + } + unsubscribe() + } + }) + } + + function publishState(next: ConnectionState): void { + if (state === next) { + return + } + state = next + for (const listener of stateListeners) { + listener(next) + } + } + + function fail(error: Error): void { + if (closed) { + return + } + closed = true + link.close() + rejectPending(error) + publishState('disconnected') + } + + function rejectPending(error: Error): void { + for (const request of pending.values()) { + clearTimeout(request.timer) + request.reject(error) + } + pending.clear() + } + + function nextId(): string { + return `relay-rpc-${++requestCounter}-${Date.now()}` + } +} + +function asError(error: unknown): Error { + return error instanceof Error ? error : new Error(String(error)) +} diff --git a/mobile/src/transport/mobile-relay-rpc-streams.ts b/mobile/src/transport/mobile-relay-rpc-streams.ts new file mode 100644 index 00000000000..552880e163a --- /dev/null +++ b/mobile/src/transport/mobile-relay-rpc-streams.ts @@ -0,0 +1,165 @@ +import { decodeBrowserScreencastFrame } from './browser-screencast-protocol' +import { + handleTerminalBinaryFrame, + type TerminalSnapshotState +} from './rpc-client-terminal-binary-frame' +import { + buildTerminalUnsubscribeParams, + updateTerminalSubscriptionViewport +} from './rpc-client-terminal-subscription' +import type { RpcClient } from './rpc-client' +import type { RpcResponse, RpcSuccess } from './types' + +type StreamRecord = { + method: string + params: unknown + listener: (result: unknown) => void + onBinaryFrame?: Parameters[3] extends + | { onBinaryFrame?: infer Listener } + | undefined + ? Listener + : never + streamIds: Set + subscriptionId?: string + cancelled: boolean +} + +type StreamManagerOptions = { + nextId: () => string + sendFrame: (request: { id: string; method: string; params?: unknown }) => boolean + waitForConnected: () => Promise +} + +export class MobileRelayRpcStreams { + private readonly streams = new Map() + private readonly terminalListeners = new Map void>() + private readonly terminalSnapshots = new Map() + private activeBrowserStream: StreamRecord | null = null + + constructor(private readonly options: StreamManagerOptions) {} + + subscribe( + method: string, + params: unknown, + listener: (result: unknown) => void, + subscribeOptions?: Parameters[3] + ): () => void { + const id = this.options.nextId() + const stream: StreamRecord = { + method, + params, + listener, + onBinaryFrame: subscribeOptions?.onBinaryFrame, + streamIds: new Set(), + cancelled: false + } + this.streams.set(id, stream) + void this.options + .waitForConnected() + .then(() => { + if (!stream.cancelled && !this.options.sendFrame({ id, method, params: stream.params })) { + this.remove(id) + } + }) + .catch(() => this.remove(id)) + return () => this.cancel(id) + } + + updateTerminalViewport(terminal: string, viewport: { cols: number; rows: number }): void { + updateTerminalSubscriptionViewport(this.streams.values(), terminal, viewport) + } + + handleResponse(response: RpcResponse): boolean { + const stream = this.streams.get(response.id) + if (!stream) { + return false + } + if (!response.ok) { + this.remove(response.id) + return true + } + const result = (response as RpcSuccess).result + if (result && typeof result === 'object') { + const metadata = result as { subscriptionId?: unknown; streamId?: unknown; type?: unknown } + if (typeof metadata.subscriptionId === 'string') { + stream.subscriptionId = metadata.subscriptionId + } + if (typeof metadata.streamId === 'number') { + stream.streamIds.add(metadata.streamId) + this.terminalListeners.set(metadata.streamId, stream.listener) + } + if (stream.method === 'browser.screencast') { + this.activeBrowserStream = stream + } + if (metadata.type === 'end') { + stream.listener(result) + this.remove(response.id) + return true + } + } + if (!stream.cancelled) { + stream.listener(result) + } + return true + } + + handleBinary(bytes: Uint8Array): void { + const browserFrame = decodeBrowserScreencastFrame(bytes) + if (browserFrame && this.activeBrowserStream?.onBinaryFrame) { + this.activeBrowserStream.onBinaryFrame(browserFrame) + return + } + handleTerminalBinaryFrame(bytes, { + terminalSnapshots: this.terminalSnapshots, + getListener: (streamId) => this.terminalListeners.get(streamId), + recordValidatedInboundTraffic: () => {} + }) + } + + clear(): void { + this.streams.clear() + this.terminalListeners.clear() + this.terminalSnapshots.clear() + this.activeBrowserStream = null + } + + private cancel(id: string): void { + const stream = this.streams.get(id) + if (!stream || stream.cancelled) { + return + } + stream.cancelled = true + if (stream.method === 'terminal.subscribe') { + const params = buildTerminalUnsubscribeParams(stream.params) + if (params) { + this.options.sendFrame({ + id: this.options.nextId(), + method: 'terminal.unsubscribe', + params + }) + } + } else if (stream.subscriptionId) { + this.options.sendFrame({ + id: this.options.nextId(), + method: stream.method.replace(/\.subscribe$/, '.unsubscribe'), + params: { subscriptionId: stream.subscriptionId } + }) + } + this.remove(id) + } + + private remove(id: string): void { + const stream = this.streams.get(id) + if (!stream) { + return + } + for (const streamId of stream.streamIds) { + this.terminalListeners.delete(streamId) + this.terminalSnapshots.delete(streamId) + } + if (this.activeBrowserStream === stream) { + this.activeBrowserStream = null + } + this.streams.delete(id) + } +} diff --git a/mobile/src/transport/stable-logical-rpc-client.test.ts b/mobile/src/transport/stable-logical-rpc-client.test.ts new file mode 100644 index 00000000000..d74e001fb60 --- /dev/null +++ b/mobile/src/transport/stable-logical-rpc-client.test.ts @@ -0,0 +1,135 @@ +import { describe, expect, it, vi } from 'vitest' +import type { ConnectionState, RpcResponse } from './types' +import type { RpcClient } from './rpc-client' +import { + createStableLogicalRpcClient, + LogicalClientCutoverError +} from './stable-logical-rpc-client' + +class FakeSession implements RpcClient { + readonly sendRequest = + vi.fn< + (method: string, params?: unknown, options?: { timeoutMs?: number }) => Promise + >() + readonly subscribe = vi.fn() + readonly updateTerminalSubscriptionViewport = + vi.fn() + readonly notifyForeground = vi.fn() + readonly close = vi.fn() + private state: ConnectionState + private readonly stateListeners = new Set<(state: ConnectionState) => void>() + private readonly streamListeners = new Set<(result: unknown) => void>() + + constructor(state: ConnectionState) { + this.state = state + this.subscribe.mockImplementation((_method, _params, listener) => { + this.streamListeners.add(listener) + return () => this.streamListeners.delete(listener) + }) + } + + getState = (): ConnectionState => this.state + getReconnectAttempt = (): number => 0 + getLastConnectedAt = (): number | null => null + onStateChange = (listener: (state: ConnectionState) => void): (() => void) => { + this.stateListeners.add(listener) + return () => this.stateListeners.delete(listener) + } + + setState(state: ConnectionState): void { + this.state = state + for (const listener of this.stateListeners) { + listener(state) + } + } + + emitStream(value: unknown): void { + for (const listener of this.streamListeners) { + listener(value) + } + } +} + +function success(value: unknown): RpcResponse { + return { id: 'rpc-1', ok: true, result: value, _meta: { runtimeId: 'runtime-1' } } +} + +function deferred() { + let resolve!: (value: T) => void + const promise = new Promise((resolvePromise) => { + resolve = resolvePromise + }) + return { promise, resolve } +} + +describe('stable logical RPC client', () => { + it('makes before break, rejects in-flight work, and replays subscriptions', async () => { + const oldSession = new FakeSession('connected') + const nextSession = new FakeSession('connecting') + const pending = deferred() + oldSession.sendRequest.mockReturnValue(pending.promise) + nextSession.sendRequest.mockResolvedValue(success('next')) + const client = createStableLogicalRpcClient(oldSession, 'direct') + const stream = vi.fn() + client.subscribe('terminal.subscribe', { terminal: 'term-1' }, stream) + const request = client.sendRequest('worktree.create', { name: 'new' }) + + const migrating = client.migrateTo(nextSession, 'relay') + expect(oldSession.close).not.toHaveBeenCalled() + expect(nextSession.subscribe).not.toHaveBeenCalled() + nextSession.setState('connected') + await migrating + + await expect(request).rejects.toBeInstanceOf(LogicalClientCutoverError) + expect(nextSession.subscribe).toHaveBeenCalledWith( + 'terminal.subscribe', + { terminal: 'term-1' }, + expect.any(Function), + undefined + ) + expect(oldSession.close).toHaveBeenCalledOnce() + expect(client.getActivePath()).toBe('relay') + expect(client.getGeneration()).toBe(2) + oldSession.emitStream('stale') + nextSession.emitStream('current') + expect(stream).toHaveBeenCalledOnce() + expect(stream).toHaveBeenCalledWith('current') + pending.resolve(success('late')) + }) + + it('keeps replies that commit before cutover and carries viewport state into replay', async () => { + const oldSession = new FakeSession('connected') + const nextSession = new FakeSession('connected') + oldSession.sendRequest.mockResolvedValue(success('old')) + const client = createStableLogicalRpcClient(oldSession, 'direct') + client.subscribe( + 'terminal.subscribe', + { terminal: 'term-1', viewport: { cols: 80, rows: 24 } }, + vi.fn() + ) + client.updateTerminalSubscriptionViewport('term-1', { cols: 120, rows: 40 }) + + await expect(client.sendRequest('status.get')).resolves.toEqual(success('old')) + await client.migrateTo(nextSession, 'relay') + expect(nextSession.subscribe).toHaveBeenCalledWith( + 'terminal.subscribe', + { terminal: 'term-1', viewport: { cols: 120, rows: 40 } }, + expect.any(Function), + undefined + ) + }) + + it('closes a replacement that fails authentication and preserves the active session', async () => { + const oldSession = new FakeSession('connected') + const replacement = new FakeSession('connecting') + const client = createStableLogicalRpcClient(oldSession, 'direct') + const migrating = client.migrateTo(replacement, 'relay') + replacement.setState('auth-failed') + + await expect(migrating).rejects.toThrow(/auth-failed/) + expect(replacement.close).toHaveBeenCalledOnce() + expect(oldSession.close).not.toHaveBeenCalled() + expect(client.getActivePath()).toBe('direct') + expect(client.getGeneration()).toBe(1) + }) +}) diff --git a/mobile/src/transport/stable-logical-rpc-client.ts b/mobile/src/transport/stable-logical-rpc-client.ts new file mode 100644 index 00000000000..b62c6eef01f --- /dev/null +++ b/mobile/src/transport/stable-logical-rpc-client.ts @@ -0,0 +1,260 @@ +import type { ConnectionState, RpcResponse } from './types' +import type { RpcClient } from './rpc-client' + +export type MobileConnectionPath = 'direct' | 'relay' + +export class LogicalClientCutoverError extends Error { + constructor() { + super('RPC interrupted by connection migration') + } +} + +type SubscriptionRecord = { + method: string + params: unknown + listener: (result: unknown) => void + options?: Parameters[3] + disposePhysical: (() => void) | null + cancelled: boolean +} + +type PendingRequest = { + reject: (error: Error) => void +} + +export type StableLogicalRpcClient = RpcClient & { + migrateTo(session: RpcClient, path: MobileConnectionPath, timeoutMs?: number): Promise + getActivePath(): MobileConnectionPath + getGeneration(): number +} + +export function createStableLogicalRpcClient( + initialSession: RpcClient, + initialPath: MobileConnectionPath +): StableLogicalRpcClient { + let activeSession = initialSession + let activePath = initialPath + let generation = 1 + let closed = false + let nextSubscriptionId = 0 + let activeStateUnsubscribe: (() => void) | null = null + const subscriptions = new Map() + const pendingRequests = new Set() + const stateListeners = new Set<(state: ConnectionState) => void>() + let state = initialSession.getState() + + bindActiveState(initialSession, generation) + + const logical: StableLogicalRpcClient = { + sendRequest(method, params, options) { + if (closed) { + return Promise.reject(new Error('Client closed')) + } + const requestGeneration = generation + const session = activeSession + return new Promise((resolve, reject) => { + const pending = { reject } + pendingRequests.add(pending) + void session.sendRequest(method, params, options).then( + (response) => { + pendingRequests.delete(pending) + if (closed) { + reject(new Error('Client closed')) + } else if (requestGeneration !== generation) { + reject(new LogicalClientCutoverError()) + } else { + resolve(response) + } + }, + (error: unknown) => { + pendingRequests.delete(pending) + reject(error) + } + ) + }) + }, + + subscribe(method, params, listener, options) { + if (closed) { + return () => {} + } + const id = ++nextSubscriptionId + const record: SubscriptionRecord = { + method, + params, + listener, + options, + disposePhysical: null, + cancelled: false + } + subscriptions.set(id, record) + attachSubscription(record, activeSession, generation) + return () => { + if (record.cancelled) { + return + } + record.cancelled = true + record.disposePhysical?.() + record.disposePhysical = null + subscriptions.delete(id) + } + }, + + updateTerminalSubscriptionViewport(terminal, viewport) { + for (const record of subscriptions.values()) { + if ( + record.params && + typeof record.params === 'object' && + 'terminal' in record.params && + record.params.terminal === terminal + ) { + record.params = { ...record.params, viewport } + } + } + activeSession.updateTerminalSubscriptionViewport(terminal, viewport) + }, + + getState: () => state, + getReconnectAttempt: () => activeSession.getReconnectAttempt(), + getLastConnectedAt: () => activeSession.getLastConnectedAt(), + onStateChange(listener) { + stateListeners.add(listener) + return () => stateListeners.delete(listener) + }, + notifyForeground: () => activeSession.notifyForeground(), + close() { + if (closed) { + return + } + closed = true + activeStateUnsubscribe?.() + activeStateUnsubscribe = null + for (const pending of pendingRequests) { + pending.reject(new Error('Client closed')) + } + pendingRequests.clear() + for (const record of subscriptions.values()) { + record.disposePhysical?.() + } + subscriptions.clear() + activeSession.close() + publishState('disconnected') + }, + + async migrateTo(nextSession, path, timeoutMs = 12_000) { + if (closed) { + nextSession.close() + throw new Error('Client closed') + } + try { + await waitForAuthenticated(nextSession, timeoutMs) + } catch (error) { + nextSession.close() + throw error + } + if (closed) { + nextSession.close() + throw new Error('Client closed') + } + const previous = activeSession + const previousStateUnsubscribe = activeStateUnsubscribe + const nextGeneration = generation + 1 + + // Why: replay on the authenticated replacement before closing the old + // session, but fence callbacks until the generation becomes current. + for (const record of subscriptions.values()) { + const disposePrevious = record.disposePhysical + attachSubscription(record, nextSession, nextGeneration) + disposePrevious?.() + } + generation = nextGeneration + activeSession = nextSession + activePath = path + previousStateUnsubscribe?.() + bindActiveState(nextSession, nextGeneration) + for (const pending of pendingRequests) { + pending.reject(new LogicalClientCutoverError()) + } + pendingRequests.clear() + state = nextSession.getState() + for (const listener of stateListeners) { + listener(state) + } + previous.close() + }, + + getActivePath: () => activePath, + getGeneration: () => generation + } + + return logical + + function attachSubscription( + record: SubscriptionRecord, + session: RpcClient, + subscriptionGeneration: number + ): void { + record.disposePhysical = session.subscribe( + record.method, + record.params, + (result) => { + if (!closed && !record.cancelled && generation === subscriptionGeneration) { + record.listener(result) + } + }, + record.options + ) + } + + function bindActiveState(session: RpcClient, sessionGeneration: number): void { + activeStateUnsubscribe = session.onStateChange((next) => { + if (!closed && generation === sessionGeneration && session === activeSession) { + publishState(next) + } + }) + } + + function publishState(next: ConnectionState): void { + if (state === next) { + return + } + state = next + for (const listener of stateListeners) { + listener(next) + } + } +} + +function waitForAuthenticated(session: RpcClient, timeoutMs: number): Promise { + if (session.getState() === 'connected') { + return Promise.resolve() + } + return new Promise((resolve, reject) => { + let settled = false + let timer: ReturnType | null = null + const unsubscribe = session.onStateChange((state) => { + if (state === 'connected') { + finish() + resolve() + } else if (state === 'auth-failed' || state === 'disconnected') { + finish() + reject(new Error(`replacement session ${state}`)) + } + }) + timer = setTimeout(() => { + finish() + reject(new Error('replacement session authentication timed out')) + }, timeoutMs) + + function finish(): void { + if (settled) { + return + } + settled = true + if (timer) { + clearTimeout(timer) + } + unsubscribe() + } + }) +}