feat(sdk): route TypeScript SDK exec and forward to the sandbox owner replica

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<SandboxRef> {
     return new Pager(async (pageToken) => {
       try {
@@ -1378,22 +1402,22 @@ export class SandboxClient {
   ): AsyncGenerator<ExecStreamEvent, void, void> {
     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<ExecInteractiveSessionControl> {
+    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<MessageInitShape<typeof ExecSandboxInputSchema>>();
-    input.push({
+    const start: MessageInitShape<typeof ExecSandboxInputSchema> = {
       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<void> {
     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<typeof TcpForwardFrameSchema> = {
         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;

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<Gateway> {
+  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<ServiceImpl<typeof OpenShell>> = {
+    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<void>((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<string[]> {
+  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<string>((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);
+  });
+});