fix(sdk/ts): revoke forward tokens on close

sdk/typescript/src/client.test.ts

@@ -12,6 +12,7 @@ import type { MessageInitShape } from '@bufbuild/protobuf';
 import { Code, ConnectError, createRouterTransport, type ServiceImpl, type Transport } from '@connectrpc/connect';
 import { afterEach, describe, expect, it } from 'vitest';
 import {
+  __setSshCleanupTimeoutMsForTests,
   __setWatchLogsBackoffSleepForTests,
   errorCode,
   Pager,
@@ -2236,6 +2237,81 @@ describe('forward', () => {
     expect(revokedToken).toBe('fwd-tok');
   });
 
+  it('revokes the token of a connection still open at close()', async () => {
+    const revoked: string[] = [];
+    const sandbox = client({
+      getSandbox: () => readySandbox('sb', 'sb-id-open'),
+      createSshSession: () => ({
+        sandboxId: 'sb-id-open',
+        token: 'open-tok',
+        gatewayHost: 'gw',
+        gatewayPort: 443,
+        gatewayScheme: 'https',
+        hostKeyFingerprint: '',
+        expirationTime: undefined,
+      }),
+      // A real transport never delivers a call whose signal is already aborted.
+      revokeSshSession: (req, ctx) => {
+        if (!ctx.signal.aborted) revoked.push(req.token);
+        return { outcome: 1 };
+      },
+      forwardTcp: async function* (requests) {
+        for await (const frame of requests) {
+          if (frame.payload.case === 'data') yield { payload: { case: 'data', value: frame.payload.value } };
+        }
+      },
+    });
+
+    const handle = await sandbox.forward('sb', { targetPort: 9000 });
+    const socket = net.connect(handle.localPort, handle.localHost);
+    socket.on('error', () => {});
+    await new Promise<void>((resolve) => {
+      socket.once('data', () => resolve());
+      socket.write('ping');
+    });
+
+    await handle.close();
+    expect(revoked).toEqual(['open-tok']);
+  });
+
+  it('close() resolves within the cleanup bound when revoke never resolves', async () => {
+    __setSshCleanupTimeoutMsForTests(100);
+    try {
+      const sandbox = client({
+        getSandbox: () => readySandbox('sb', 'sb-id-stall'),
+        createSshSession: () => ({
+          sandboxId: 'sb-id-stall',
+          token: 'stall-tok',
+          gatewayHost: 'gw',
+          gatewayPort: 443,
+          gatewayScheme: 'https',
+          hostKeyFingerprint: '',
+          expirationTime: undefined,
+        }),
+        revokeSshSession: () => new Promise(() => {}),
+        forwardTcp: async function* (requests) {
+          for await (const frame of requests) {
+            if (frame.payload.case === 'data') yield { payload: { case: 'data', value: frame.payload.value } };
+          }
+        },
+      });
+
+      const handle = await sandbox.forward('sb', { targetPort: 9000 });
+      const socket = net.connect(handle.localPort, handle.localHost);
+      socket.on('error', () => {});
+      await new Promise<void>((resolve) => {
+        socket.once('data', () => resolve());
+        socket.write('ping');
+      });
+
+      const start = Date.now();
+      await handle.close();
+      expect(Date.now() - start).toBeLessThan(2_000);
+    } finally {
+      __setSshCleanupTimeoutMsForTests();
+    }
+  });
+
   it('rejects when the sandbox is not ready', async () => {
     const sandbox = client({
       getSandbox: () => ({

sdk/typescript/src/client.ts

@@ -485,7 +485,8 @@ export interface ForwardOptions extends SandboxWorkspaceOptions {
 
 // A process-lifetime local listener that tunnels each accepted connection into
 // the sandbox. Call `close()` on teardown; `closed` resolves once the listener
-// is fully torn down. An in-process forward cannot outlive the Node process.
+// is fully torn down. `close()` may wait up to 5s for best-effort revokes of
+// open connections' tokens. An in-process forward cannot outlive the Node process.
 export interface ForwardHandle {
   localHost: string;
   localPort: number;
@@ -723,6 +724,15 @@ function versionPin(value: string | undefined): bigint {
 
 const FORWARD_CHUNK = 64 * 1024;
 
+// Bound for revoking a forward's SSH session token during cleanup.
+const SSH_CLEANUP_TIMEOUT_MS = 5_000;
+let sshCleanupTimeoutMs = SSH_CLEANUP_TIMEOUT_MS;
+
+// Test hook; not re-exported from index.ts. No argument restores the default.
+export function __setSshCleanupTimeoutMsForTests(ms?: number): void {
+  sshCleanupTimeoutMs = ms ?? SSH_CLEANUP_TIMEOUT_MS;
+}
+
 function workspaceName(options?: SandboxWorkspaceOptions | null): string {
   const workspace = options?.workspace ?? 'default';
   if (workspace.trim() === '') throw new SdkError('invalid_config', 'workspace must be non-empty');
@@ -1894,7 +1904,11 @@ export class SandboxClient {
       input.end();
       if (token !== undefined) {
         try {
-          await this.grpc.revokeSshSession({ token, allowMissing: true }, { signal });
+          // Not the connection signal: close() aborts it before this runs.
+          await this.grpc.revokeSshSession(
+            { token, allowMissing: true },
+            { signal: AbortSignal.timeout(sshCleanupTimeoutMs) },
+          );
         } catch {
           // Best-effort revoke; the token expires on its own regardless.
         }