Move .connect to .wsConnect within DO stub (#175)

* Move .connect to .wsConnect within DO stub

* Use `Sandbox` type instead of DurableObjectStub<>

This gets around IDEs not being able to do go-to-definition due to
mapped types not being resolved properly. This is less truthful
perhaps (hides RPC semantics), but more useful for developers.

* Add changeset

* Revert "Use `Sandbox` type instead of DurableObjectStub<>"

This reverts commit 4ed544e.

* Undo unnecessary type changes

Author: ghostwriternr

.changeset/hot-pans-warn.md

@@ -0,0 +1,5 @@
+---
+"@cloudflare/sandbox": patch
+---
+
+Move .connect to .wsConnect within DO stub

.gitignore

@@ -149,6 +149,8 @@ tests/e2e/test-worker/wrangler.jsonc
 # Allow config files and scripts
 !scripts/**/*.js
 !examples/**/*.js
+# Allow generated types in examples
+!examples/**/*.d.ts
 # Allow source in dist directories
 !dist/**
 !node_modules/**

examples/claude-code/src/index.ts

@@ -1,4 +1,4 @@
-import { getSandbox, type Sandbox } from '@cloudflare/sandbox';
+import { getSandbox } from '@cloudflare/sandbox';
 
 interface CmdOutput {
   success: boolean;
@@ -8,11 +8,6 @@ interface CmdOutput {
 // helper to read the outputs from `.exec` results
 const getOutput = (res: CmdOutput) => (res.success ? res.stdout : res.stderr);
 
-type Env = {
-  Sandbox: DurableObjectNamespace<Sandbox>;
-  ANTHROPIC_API_KEY: string;
-};
-
 const EXTRA_SYSTEM =
   'You are an automatic feature-implementer/bug-fixer.' +
   'You apply all necessary changes to achieve the user request. You must ensure you DO NOT commit the changes, ' +

examples/claude-code/worker-configuration.d.ts

@@ -0,0 +1,2 @@
+CLOUDFLARE_ACCOUNT_ID=xxxx
+CLOUDFLARE_API_KEY=xxxx

examples/code-interpreter/.dev.vars.example

@@ -1,11 +1,7 @@
-import { getSandbox, type Sandbox } from '@cloudflare/sandbox';
+import { getSandbox } from '@cloudflare/sandbox';
 
 export { Sandbox } from '@cloudflare/sandbox';
 
-type Env = {
-  Sandbox: DurableObjectNamespace<Sandbox>;
-};
-
 export default {
   async fetch(request: Request, env: Env): Promise<Response> {
     const url = new URL(request.url);

examples/minimal/src/index.ts

@@ -10,7 +10,7 @@ export {
   SandboxClient,
   UtilityClient
 } from './clients';
-export { connect, getSandbox, Sandbox } from './sandbox';
+export { getSandbox, Sandbox } from './sandbox';
 
 // Legacy types are now imported from the new client architecture
 

examples/minimal/worker-configuration.d.ts

@@ -32,7 +32,7 @@ export function getSandbox(
   id: string,
   options?: SandboxOptions
 ): Sandbox {
-  const stub = getContainer(ns, id) as any as Sandbox;
+  const stub = getContainer(ns, id) as unknown as Sandbox;
 
   // Store the name on first access
   stub.setSandboxName?.(id);
@@ -49,41 +49,24 @@ export function getSandbox(
     stub.setKeepAlive(options.keepAlive);
   }
 
-  return stub;
+  return Object.assign(stub, {
+    wsConnect: connect(stub)
+  });
 }
 
-/**
- * Connect an incoming WebSocket request to a specific port inside the container.
- *
- * Note: This is a standalone function (not a Sandbox method) because WebSocket
- * connections cannot be serialized over Durable Object RPC.
- *
- * @param sandbox - The Sandbox instance to route the request through
- * @param request - The incoming WebSocket upgrade request
- * @param port - The port number to connect to (1024-65535)
- * @returns The WebSocket upgrade response
- * @throws {SecurityError} - If port is invalid or in restricted range
- *
- * @example
- * const sandbox = getSandbox(env.Sandbox, 'sandbox-id');
- * if (request.headers.get('Upgrade')?.toLowerCase() === 'websocket') {
- *   return await connect(sandbox, request, 8080);
- * }
- */
-export async function connect(
-  sandbox: Sandbox,
-  request: Request,
-  port: number
-): Promise<Response> {
-  // Validate port before routing
-  if (!validatePort(port)) {
-    throw new SecurityError(
-      `Invalid or restricted port: ${port}. Ports must be in range 1024-65535 and not reserved.`
-    );
-  }
-
-  const portSwitchedRequest = switchPort(request, port);
-  return await sandbox.fetch(portSwitchedRequest);
+export function connect(
+  stub: { fetch: (request: Request) => Promise<Response> }
+) {
+  return async (request: Request, port: number) => {
+    // Validate port before routing
+    if (!validatePort(port)) {
+      throw new SecurityError(
+        `Invalid or restricted port: ${port}. Ports must be in range 1024-65535 and not reserved.`
+      );
+    }
+    const portSwitchedRequest = switchPort(request, port);
+    return await stub.fetch(portSwitchedRequest);
+  };
 }
 
 export class Sandbox<Env = unknown> extends Container<Env> implements ISandbox {
@@ -100,7 +83,7 @@ export class Sandbox<Env = unknown> extends Container<Env> implements ISandbox {
   private logger: ReturnType<typeof createLogger>;
   private keepAliveEnabled: boolean = false;
 
-  constructor(ctx: DurableObject['ctx'], env: Env) {
+  constructor(ctx: DurableObjectState<{}>, env: Env) {
     super(ctx, env);
 
     const envObj = env as any;
@@ -360,6 +343,11 @@ export class Sandbox<Env = unknown> extends Container<Env> implements ISandbox {
     });
   }
 
+  wsConnect(request: Request, port: number): Promise<Response> {
+    // Dummy implementation that will be overridden by the stub
+    throw new Error('Not implemented here to avoid RPC serialization issues');
+  }
+
   private determinePort(url: URL): number {
     // Extract port from proxy requests (e.g., /proxy/8080/*)
     const proxyMatch = url.pathname.match(/^\/proxy\/(\d+)/);

packages/sandbox/src/index.ts

@@ -53,7 +53,7 @@ vi.mock('@cloudflare/containers', () => {
 
 describe('Sandbox - Automatic Session Management', () => {
   let sandbox: Sandbox;
-  let mockCtx: Partial<DurableObjectState>;
+  let mockCtx: Partial<DurableObjectState<{}>>;
   let mockEnv: any;
 
   beforeEach(async () => {
@@ -82,13 +82,17 @@ describe('Sandbox - Automatic Session Management', () => {
     mockEnv = {};
 
     // Create Sandbox instance - SandboxClient is created internally
-    sandbox = new Sandbox(mockCtx as DurableObjectState, mockEnv);
+    const stub = new Sandbox(mockCtx, mockEnv);
 
     // Wait for blockConcurrencyWhile to complete
     await vi.waitFor(() => {
       expect(mockCtx.blockConcurrencyWhile).toHaveBeenCalled();
     });
 
+    sandbox = Object.assign(stub, {
+      wsConnect: connect(stub)
+    });
+
     // Now spy on the client methods that we need for testing
     vi.spyOn(sandbox.client.utils, 'createSession').mockResolvedValue({
       success: true,
@@ -621,7 +625,7 @@ describe('Sandbox - Automatic Session Management', () => {
     });
   });
 
-  describe('connect() function', () => {
+  describe('wsConnect() method', () => {
     it('should route WebSocket request through switchPort to sandbox.fetch', async () => {
       const { switchPort } = await import('@cloudflare/containers');
       const switchPortMock = vi.mocked(switchPort);
@@ -634,7 +638,7 @@ describe('Sandbox - Automatic Session Management', () => {
       });
 
       const fetchSpy = vi.spyOn(sandbox, 'fetch');
-      const response = await connect(sandbox, request, 8080);
+      const response = await sandbox.wsConnect(request, 8080);
 
       // Verify switchPort was called with correct port
       expect(switchPortMock).toHaveBeenCalledWith(request, 8080);
@@ -653,21 +657,21 @@ describe('Sandbox - Automatic Session Management', () => {
       });
 
       // Invalid port values
-      await expect(connect(sandbox, request, -1)).rejects.toThrow(
+      await expect(sandbox.wsConnect(request, -1)).rejects.toThrow(
         'Invalid or restricted port'
       );
-      await expect(connect(sandbox, request, 0)).rejects.toThrow(
+      await expect(sandbox.wsConnect(request, 0)).rejects.toThrow(
         'Invalid or restricted port'
       );
-      await expect(connect(sandbox, request, 70000)).rejects.toThrow(
+      await expect(sandbox.wsConnect(request, 70000)).rejects.toThrow(
         'Invalid or restricted port'
       );
 
       // Privileged ports
-      await expect(connect(sandbox, request, 80)).rejects.toThrow(
+      await expect(sandbox.wsConnect(request, 80)).rejects.toThrow(
         'Invalid or restricted port'
       );
-      await expect(connect(sandbox, request, 443)).rejects.toThrow(
+      await expect(sandbox.wsConnect(request, 443)).rejects.toThrow(
         'Invalid or restricted port'
       );
     });
@@ -685,7 +689,7 @@ describe('Sandbox - Automatic Session Management', () => {
       );
 
       const fetchSpy = vi.spyOn(sandbox, 'fetch');
-      await connect(sandbox, request, 8080);
+      await sandbox.wsConnect(request, 8080);
 
       const calledRequest = fetchSpy.mock.calls[0][0];