From 711ac48d487fba1ecfcf8bd4eef4f9ca739eafa0 Mon Sep 17 00:00:00 2001 From: divesh Date: Fri, 23 Oct 2026 10:10:48 -0700 Subject: [PATCH] feat(sdk): route TypeScript SDK exec and forward to the sandbox owner replica Ask GetSandbox for the owner replica and send it as x-openshell-replica on ExecSandbox, ExecSandboxInteractive and ForwardTcp. Retry once without the hint when the per-replica route has no healthy endpoint, holding stream input until response headers arrive so the retry only resends the first frame. Part of #4223 Signed-off-by: divesh --- sdk/typescript/src/client.ts | 83 +++++++++++---- sdk/typescript/src/routing.test.ts | 162 +++++++++++++++++++++++++++++ sdk/typescript/src/routing.ts | 65 ++++++++++++ 3 files changed, 289 insertions(+), 21 deletions(-) create mode 100644 sdk/typescript/src/routing.test.ts create mode 100644 sdk/typescript/src/routing.ts diff --git a/sdk/typescript/src/client.ts b/sdk/typescript/src/client.ts index 74ade2297d..49e8897c37 100644 --- a/sdk/typescript/src/client.ts +++ b/sdk/typescript/src/client.ts @@ -40,6 +40,7 @@ import { import type { EffectiveSetting, GetSandboxConfigResponse, SandboxPolicy, SettingValue } from './gen/sandbox_pb.js'; import { PolicySource, type SandboxPolicySchema, SettingScope, type SettingValueSchema } from './gen/sandbox_pb.js'; import { ProviderClient } from './providers.js'; +import { heldUntilReady, OWNER_REPLICA_HEADER, OWNER_REQUEST_HEADER, routedStream } from './routing.js'; import { validateSshResponse } from './ssh-validate.js'; import { buildTransport, type ConnectOptions } from './transport.js'; @@ -1278,6 +1279,29 @@ export class SandboxClient { } } + // get() that also asks the gateway which replica owns the supervisor session. + private async getWithOwner( + name: string, + options?: SandboxCallOptions | null, + ): Promise<{ ref: SandboxRef; owner?: string }> { + let owner: string | undefined; + try { + const resp = await this.grpc.getSandbox( + { ...namedTarget(name, options) }, + { + ...requestCallOptions(options), + headers: { [OWNER_REQUEST_HEADER]: '1' }, + onHeader: (header) => { + owner = header.get(OWNER_REPLICA_HEADER) || undefined; + }, + }, + ); + return { ref: sandboxRef(resp.sandbox), owner }; + } catch (e) { + throw fromConnect(e); + } + } + list(options?: ListOptions | null): Pager { return new Pager(async (pageToken) => { try { @@ -1378,22 +1402,22 @@ export class SandboxClient { ): AsyncGenerator { try { // Preserve the existing preflight so lookup failures surface before the stream starts. - await this.get(name, { + const { owner } = await this.getWithOwner(name, { workspace: options?.workspace, ...(options?.signal ? { signal: options.signal } : {}), }); - const stream = this.grpc.execSandbox( - { - ...sandboxTarget(name, options), - command, - workdir: options?.workdir ?? '', - environment: options?.environment ?? {}, - executionTimeout: durationFromSeconds(options?.timeoutSecs ?? 0), - stdin: options?.stdin ? new Uint8Array(options.stdin) : new Uint8Array(), - tty: false, - noLoginShell: options?.noLoginShell ?? false, - }, - { signal: options?.signal }, + const request = { + ...sandboxTarget(name, options), + command, + workdir: options?.workdir ?? '', + environment: options?.environment ?? {}, + executionTimeout: durationFromSeconds(options?.timeoutSecs ?? 0), + stdin: options?.stdin ? new Uint8Array(options.stdin) : new Uint8Array(), + tty: false, + noLoginShell: options?.noLoginShell ?? false, + }; + const stream = routedStream(owner, ({ headers, onHeader }) => + this.grpc.execSandbox(request, { signal: options?.signal, headers, onHeader }), ); let sawExit = false; @@ -1534,14 +1558,18 @@ export class SandboxClient { command: string[], options?: ExecInteractiveOptions | null, ): Promise { + let owner: string | undefined; try { - await this.get(name, { workspace: options?.workspace, ...(options?.signal ? { signal: options.signal } : {}) }); + ({ owner } = await this.getWithOwner(name, { + workspace: options?.workspace, + ...(options?.signal ? { signal: options.signal } : {}), + })); } catch (e) { throw e instanceof SdkError ? e : fromConnect(e); } const input = new Pushable>(); - input.push({ + const start: MessageInitShape = { payload: { case: 'start', value: { @@ -1557,7 +1585,7 @@ export class SandboxClient { noLoginShell: options?.noLoginShell ?? false, }, }, - }); + }; const controller = new AbortController(); const signal = options?.signal ? AbortSignal.any([options.signal, controller.signal]) : controller.signal; @@ -1611,7 +1639,9 @@ export class SandboxClient { if (signal.aborted) throw new SdkError('canceled', 'exec cancelled'); // Start and observe the transport immediately, independently of output // consumption. Backpressure bounds the queue to 16 chunks of 64 KiB. - const stream = grpc.execSandboxInteractive(input, { signal }); + const stream = routedStream(owner, ({ headers, onHeader, ready }) => + grpc.execSandboxInteractive(heldUntilReady(start, input, ready), { signal, headers, onHeader }), + ); for await (const event of stream) { if (signal.aborted) throw new SdkError('canceled', 'exec cancelled'); if (exitCode !== undefined) { @@ -1706,12 +1736,18 @@ export class SandboxClient { const localPort = opts.localPort ?? 0; let sandboxId: string; + let owner: string | undefined; try { - const ref = await this.get(name, { workspace: opts.workspace, ...(opts.signal ? { signal: opts.signal } : {}) }); + const found = await this.getWithOwner(name, { + workspace: opts.workspace, + ...(opts.signal ? { signal: opts.signal } : {}), + }); + const ref = found.ref; if (ref.phase !== 'ready') { throw new SdkError('connect', `sandbox '${name}' is not ready (phase: ${ref.phase})`); } sandboxId = ref.id; + owner = found.owner; } catch (e) { throw e instanceof SdkError ? e : fromConnect(e); } @@ -1737,6 +1773,7 @@ export class SandboxClient { opts.workspace, targetHost, targetPort, + owner, controller.signal, ) .catch((error: unknown) => { @@ -1821,6 +1858,7 @@ export class SandboxClient { workspace: string | undefined, targetHost: string, targetPort: number, + owner: string | undefined, signal: AbortSignal, ): Promise { let token: string | undefined; @@ -1833,7 +1871,7 @@ export class SandboxClient { // tears down this one socket via the catch below. validateSshResponse(session, sandboxId); token = session.token; - input.push({ + const init: MessageInitShape = { payload: { case: 'init', value: { @@ -1846,7 +1884,7 @@ export class SandboxClient { authorizationToken: token, }, }, - }); + }; socket.on('data', (chunk: Buffer) => { for (let off = 0; off < chunk.length; off += FORWARD_CHUNK) { @@ -1868,7 +1906,10 @@ export class SandboxClient { signal.addEventListener('abort', onAbort, { once: true }); try { - for await (const frame of this.grpc.forwardTcp(input, { signal })) { + const frames = routedStream(owner, ({ headers, onHeader, ready }) => + this.grpc.forwardTcp(heldUntilReady(init, input, ready), { signal, headers, onHeader }), + ); + for await (const frame of frames) { if (frame.payload.case !== 'data') continue; const data = frame.payload.value; if (data.length === 0) continue; diff --git a/sdk/typescript/src/routing.test.ts b/sdk/typescript/src/routing.test.ts new file mode 100644 index 0000000000..5b215e0bb7 --- /dev/null +++ b/sdk/typescript/src/routing.test.ts @@ -0,0 +1,162 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// Replica routing over real gRPC/HTTP2, with a front proxy that answers routed +// calls the way Envoy does when the owner's route has no ready endpoints. + +import * as http2 from 'node:http2'; +import type { AddressInfo } from 'node:net'; +import * as net from 'node:net'; +import { Code, ConnectError, type ServiceImpl } from '@connectrpc/connect'; +import { connectNodeAdapter, createGrpcTransport } from '@connectrpc/connect-node'; +import { afterEach, describe, expect, it } from 'vitest'; +import { SandboxClient } from './client.js'; +import { OpenShell, SandboxPhase } from './gen/openshell_pb.js'; +import { retryUnrouted } from './routing.js'; + +const enc = (s: string) => new TextEncoder().encode(s); +const closers: Array<() => void> = []; +afterEach(() => { + for (const close of closers.splice(0)) close(); +}); + +interface Gateway { + sandbox: SandboxClient; + wantedOwner: string[]; + // Replica header seen by the gateway on each streaming call ('' when unrouted). + routes: string[]; + rejectedByProxy: number; + interactiveFrames: string[]; +} + +async function gateway(opts: { owner?: string; rejectRouted?: boolean; execError?: ConnectError }): Promise { + const g: Gateway = { + sandbox: undefined as never, + wantedOwner: [], + routes: [], + rejectedByProxy: 0, + interactiveFrames: [], + }; + const route = (header: Headers): void => { + g.routes.push(header.get('x-openshell-replica') ?? ''); + }; + const impl: Partial> = { + getSandbox: (_req, ctx) => { + g.wantedOwner.push(ctx.requestHeader.get('x-openshell-want-owner') ?? ''); + if (opts.owner) ctx.responseHeader.set('x-openshell-owner', opts.owner); + return { sandbox: { metadata: { id: 'sb-id', name: 'sb' }, status: { phase: SandboxPhase.READY } } }; + }, + execSandbox: async function* (_req, ctx) { + route(ctx.requestHeader); + if (opts.execError) throw opts.execError; + yield { payload: { case: 'stdout', value: { data: enc('hi') } } }; + yield { payload: { case: 'exit', value: { exitCode: 0 } } }; + }, + // The first (empty) message flushes response headers, like the gateway does. + execSandboxInteractive: async function* (requests, ctx) { + route(ctx.requestHeader); + yield { payload: { case: 'stdout', value: { data: new Uint8Array() } } }; + for await (const frame of requests) { + g.interactiveFrames.push(frame.payload.case ?? ''); + if (frame.payload.case === 'stdin') { + yield { payload: { case: 'stdout', value: { data: frame.payload.value } } }; + yield { payload: { case: 'exit', value: { exitCode: 0 } } }; + return; + } + } + }, + createSshSession: () => ({ + sandboxId: 'sb-id', + token: 'tok', + gatewayHost: 'gw', + gatewayPort: 443, + gatewayScheme: 'https', + }), + revokeSshSession: () => ({ outcome: 1 }), + forwardTcp: async function* (requests, ctx) { + route(ctx.requestHeader); + yield { payload: { case: 'data', value: new Uint8Array() } }; + for await (const frame of requests) { + if (frame.payload.case === 'data') yield { payload: { case: 'data', value: frame.payload.value } }; + } + }, + }; + const adapter = connectNodeAdapter({ routes: (router) => router.service(OpenShell, impl) }); + const server = http2.createServer((req, res) => { + if (opts.rejectRouted && req.headers['x-openshell-replica']) { + g.rejectedByProxy++; + res.writeHead(200, { 'content-type': 'application/grpc', 'grpc-status': '14' }); + res.end(); + return; + } + adapter(req, res); + }); + await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)); + closers.push(() => server.close()); + const { port } = server.address() as AddressInfo; + g.sandbox = new SandboxClient(createGrpcTransport({ baseUrl: `http://127.0.0.1:${port}` })); + return g; +} + +// Run exec, an interactive exec that writes stdin right away, and one forwarded connection. +async function runAll(g: Gateway): Promise { + const exec = await g.sandbox.exec('sb', ['true']); + + const session = await g.sandbox.execInteractive('sb', ['cat']); + session.write(Buffer.from('typed')); + let typed = ''; + for await (const event of session.output) if (!('type' in event)) typed += event.data.toString(); + + const handle = await g.sandbox.forward('sb', { targetPort: 9000 }); + const echoed = await new Promise((resolve, reject) => { + const socket = net.connect(handle.localPort, handle.localHost, () => socket.write('ping')); + socket.on('data', (d) => { + resolve(d.toString()); + socket.end(); + }); + socket.on('error', reject); + }); + await handle.close(); + return [exec.stdout.toString(), typed, echoed]; +} + +describe('replica routing', () => { + it('asks for the owner and routes exec, interactive exec and forward to it', async () => { + const g = await gateway({ owner: 'openshell-1' }); + expect(await runAll(g)).toEqual(['hi', 'typed', 'ping']); + expect(g.wantedOwner).toEqual(['1', '1', '1']); + expect(g.routes).toEqual(['openshell-1', 'openshell-1', 'openshell-1']); + }); + + it('sends no hint when the gateway reports no owner', async () => { + const g = await gateway({}); + expect(await runAll(g)).toEqual(['hi', 'typed', 'ping']); + expect(g.routes).toEqual(['', '', '']); + }); + + it('retries once unrouted when the proxy rejects the owner route, without losing input', async () => { + const g = await gateway({ owner: 'openshell-1', rejectRouted: true }); + expect(await runAll(g)).toEqual(['hi', 'typed', 'ping']); + expect(g.rejectedByProxy).toBe(3); + expect(g.routes).toEqual(['', '', '']); + expect(g.interactiveFrames).toEqual(['start', 'stdin']); + }); + + it('does not retry a gateway error', async () => { + const g = await gateway({ + owner: 'openshell-1', + execError: new ConnectError('supervisor not connected', Code.Unavailable), + }); + await expect(g.sandbox.exec('sb', ['true'])).rejects.toThrow('supervisor not connected'); + expect(g.routes).toEqual(['openshell-1']); + }); + + it('matches only the proxy no-endpoints answer', () => { + expect(retryUnrouted(new ConnectError('', Code.Unavailable))).toBe(true); + expect(retryUnrouted(new ConnectError('no healthy upstream', Code.Unavailable))).toBe(true); + expect(retryUnrouted(new ConnectError('upstream connect error: connection failure', Code.Unavailable))).toBe(true); + expect(retryUnrouted(new ConnectError('supervisor not connected', Code.Unavailable))).toBe(false); + expect(retryUnrouted(new ConnectError('', Code.NotFound))).toBe(false); + expect(retryUnrouted(new Error(''))).toBe(false); + }); +}); diff --git a/sdk/typescript/src/routing.ts b/sdk/typescript/src/routing.ts new file mode 100644 index 0000000000..c8cb046f10 --- /dev/null +++ b/sdk/typescript/src/routing.ts @@ -0,0 +1,65 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// Metadata that lets an edge proxy send a sandbox's long-lived calls straight +// to the gateway replica holding its supervisor session. + +import { Code, ConnectError } from '@connectrpc/connect'; + +export const OWNER_REPLICA_HEADER = 'x-openshell-owner'; +export const OWNER_REQUEST_HEADER = 'x-openshell-want-owner'; +const ROUTE_REPLICA_HEADER = 'x-openshell-replica'; + +export interface RoutedAttempt { + headers?: Record; + onHeader: () => void; + // Resolves true once the response headers arrive, false if the attempt fails first. + ready: Promise; +} + +// Envoy Gateway answers a replica route with no ready endpoints with +// UNAVAILABLE and no message; gateway errors always carry one. +export function retryUnrouted(error: unknown): boolean { + if (!(error instanceof ConnectError) || error.code !== Code.Unavailable) return false; + const message = error.rawMessage; + return message === '' || message.includes('no healthy upstream') || message.includes('connection failure'); +} + +// Open a stream tagged with the owner replica, and once more untagged when the +// proxy rejects it before sending headers because the owner's route has no +// ready endpoints. A stale owner is safe: any replica relays the call. +export async function* routedStream( + owner: string | undefined, + open: (attempt: RoutedAttempt) => AsyncIterable, +): AsyncGenerator { + let replica = owner; + for (;;) { + let headersSeen = !replica; + let settle!: (ok: boolean) => void; + const ready = new Promise((resolve) => { + settle = resolve; + }); + if (headersSeen) settle(true); + const onHeader = (): void => { + headersSeen = true; + settle(true); + }; + try { + const headers = replica ? { [ROUTE_REPLICA_HEADER]: replica } : undefined; + yield* open({ headers, onHeader, ready }); + return; + } catch (e) { + if (headersSeen || !retryUnrouted(e)) throw e; + replica = undefined; + } finally { + settle(false); + } + } +} + +// Request frames for one attempt: `first`, then `rest` once the attempt is +// ready, so a retried attempt only has to resend `first`. +export async function* heldUntilReady(first: T, rest: AsyncIterable, ready: Promise): AsyncGenerator { + yield first; + if (await ready) yield* rest; +}