@cloudflare/sandbox 0.0.0-0b4cc05 → 0.0.0-102fc4f

package/tests/sandbox.test.ts

@@ -1,27 +1,59 @@
+import { Container } from '@cloudflare/containers';
 import type { DurableObjectState } from '@cloudflare/workers-types';
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
-import { Sandbox } from '../src/sandbox';
+import { connect, Sandbox } from '../src/sandbox';
 
 // Mock dependencies before imports
 vi.mock('./interpreter', () => ({
-  CodeInterpreter: vi.fn().mockImplementation(() => ({})),
+  CodeInterpreter: vi.fn().mockImplementation(() => ({}))
 }));
 
-vi.mock('@cloudflare/containers', () => ({
-  Container: class Container {
+vi.mock('@cloudflare/containers', () => {
+  const mockSwitchPort = vi.fn((request: Request, port: number) => {
+    // Create a new request with the port in the URL path
+    const url = new URL(request.url);
+    url.pathname = `/proxy/${port}${url.pathname}`;
+    return new Request(url, request);
+  });
+
+  const MockContainer = class Container {
     ctx: any;
     env: any;
     constructor(ctx: any, env: any) {
       this.ctx = ctx;
       this.env = env;
     }
-  },
-  getContainer: vi.fn(),
-}));
+    async fetch(request: Request): Promise<Response> {
+      // Mock implementation - will be spied on in tests
+      const upgradeHeader = request.headers.get('Upgrade');
+      if (upgradeHeader?.toLowerCase() === 'websocket') {
+        return new Response('WebSocket Upgraded', {
+          status: 200,
+          headers: {
+            'X-WebSocket-Upgraded': 'true',
+            Upgrade: 'websocket',
+            Connection: 'Upgrade'
+          }
+        });
+      }
+      return new Response('Mock Container fetch');
+    }
+    async containerFetch(request: Request, port: number): Promise<Response> {
+      // Mock implementation for HTTP path
+      return new Response('Mock Container HTTP fetch');
+    }
+  };
+
+  return {
+    Container: MockContainer,
+    getContainer: vi.fn(),
+    switchPort: mockSwitchPort
+  };
+});
 
 describe('Sandbox - Automatic Session Management', () => {
   let sandbox: Sandbox;
-  let mockCtx: Partial<DurableObjectState>;
+  let mockCtx: Partial<DurableObjectState<{}>>;
   let mockEnv: any;
 
   beforeEach(async () => {
@@ -33,31 +65,39 @@ describe('Sandbox - Automatic Session Management', () => {
         get: vi.fn().mockResolvedValue(null),
         put: vi.fn().mockResolvedValue(undefined),
         delete: vi.fn().mockResolvedValue(undefined),
-        list: vi.fn().mockResolvedValue(new Map()),
+        list: vi.fn().mockResolvedValue(new Map())
       } as any,
-      blockConcurrencyWhile: vi.fn((fn: () => Promise<void>) => fn()),
+      blockConcurrencyWhile: vi
+        .fn()
+        .mockImplementation(
+          <T>(callback: () => Promise<T>): Promise<T> => callback()
+        ),
       id: {
         toString: () => 'test-sandbox-id',
         equals: vi.fn(),
-        name: 'test-sandbox',
-      } as any,
+        name: 'test-sandbox'
+      } as any
     };
 
     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,
       id: 'sandbox-default',
-      message: 'Created',
+      message: 'Created'
     } as any);
 
     vi.spyOn(sandbox.client.commands, 'execute').mockResolvedValue({
@@ -66,13 +106,13 @@ describe('Sandbox - Automatic Session Management', () => {
       stderr: '',
       exitCode: 0,
       command: '',
-      timestamp: new Date().toISOString(),
+      timestamp: new Date().toISOString()
     } as any);
 
     vi.spyOn(sandbox.client.files, 'writeFile').mockResolvedValue({
       success: true,
       path: '/test.txt',
-      timestamp: new Date().toISOString(),
+      timestamp: new Date().toISOString()
     } as any);
   });
 
@@ -88,7 +128,7 @@ describe('Sandbox - Automatic Session Management', () => {
         stderr: '',
         exitCode: 0,
         command: 'echo test',
-        timestamp: new Date().toISOString(),
+        timestamp: new Date().toISOString()
       } as any);
 
       await sandbox.exec('echo test');
@@ -97,7 +137,7 @@ describe('Sandbox - Automatic Session Management', () => {
       expect(sandbox.client.utils.createSession).toHaveBeenCalledWith({
         id: expect.stringMatching(/^sandbox-/),
         env: {},
-        cwd: '/workspace',
+        cwd: '/workspace'
       });
 
       expect(sandbox.client.commands.execute).toHaveBeenCalledWith(
@@ -113,9 +153,12 @@ describe('Sandbox - Automatic Session Management', () => {
 
       expect(sandbox.client.utils.createSession).toHaveBeenCalledTimes(1);
 
-      const firstSessionId = vi.mocked(sandbox.client.commands.execute).mock.calls[0][1];
-      const fileSessionId = vi.mocked(sandbox.client.files.writeFile).mock.calls[0][2];
-      const secondSessionId = vi.mocked(sandbox.client.commands.execute).mock.calls[1][1];
+      const firstSessionId = vi.mocked(sandbox.client.commands.execute).mock
+        .calls[0][1];
+      const fileSessionId = vi.mocked(sandbox.client.files.writeFile).mock
+        .calls[0][2];
+      const secondSessionId = vi.mocked(sandbox.client.commands.execute).mock
+        .calls[1][1];
 
       expect(firstSessionId).toBe(fileSessionId);
       expect(firstSessionId).toBe(secondSessionId);
@@ -127,19 +170,21 @@ describe('Sandbox - Automatic Session Management', () => {
         processId: 'proc-1',
         pid: 1234,
         command: 'sleep 10',
-        timestamp: new Date().toISOString(),
+        timestamp: new Date().toISOString()
       } as any);
 
       vi.spyOn(sandbox.client.processes, 'listProcesses').mockResolvedValue({
         success: true,
-        processes: [{
-          id: 'proc-1',
-          pid: 1234,
-          command: 'sleep 10',
-          status: 'running',
-          startTime: new Date().toISOString(),
-        }],
-        timestamp: new Date().toISOString(),
+        processes: [
+          {
+            id: 'proc-1',
+            pid: 1234,
+            command: 'sleep 10',
+            status: 'running',
+            startTime: new Date().toISOString()
+          }
+        ],
+        timestamp: new Date().toISOString()
       } as any);
 
       const process = await sandbox.startProcess('sleep 10');
@@ -148,11 +193,14 @@ describe('Sandbox - Automatic Session Management', () => {
       expect(sandbox.client.utils.createSession).toHaveBeenCalledTimes(1);
 
       // startProcess uses sessionId (to start process in that session)
-      const startSessionId = vi.mocked(sandbox.client.processes.startProcess).mock.calls[0][1];
+      const startSessionId = vi.mocked(sandbox.client.processes.startProcess)
+        .mock.calls[0][1];
       expect(startSessionId).toMatch(/^sandbox-/);
 
       // listProcesses is sandbox-scoped - no sessionId parameter
-      const listProcessesCall = vi.mocked(sandbox.client.processes.listProcesses).mock.calls[0];
+      const listProcessesCall = vi.mocked(
+        sandbox.client.processes.listProcesses
+      ).mock.calls[0];
       expect(listProcessesCall).toEqual([]);
 
       // Verify the started process appears in the list
@@ -168,10 +216,12 @@ describe('Sandbox - Automatic Session Management', () => {
         stderr: '',
         branch: 'main',
         targetDir: '/workspace/repo',
-        timestamp: new Date().toISOString(),
+        timestamp: new Date().toISOString()
       } as any);
 
-      await sandbox.gitCheckout('https://github.com/test/repo.git', { branch: 'main' });
+      await sandbox.gitCheckout('https://github.com/test/repo.git', {
+        branch: 'main'
+      });
 
       expect(sandbox.client.utils.createSession).toHaveBeenCalledTimes(1);
       expect(sandbox.client.git.checkout).toHaveBeenCalledWith(
@@ -189,7 +239,7 @@ describe('Sandbox - Automatic Session Management', () => {
       expect(sandbox.client.utils.createSession).toHaveBeenCalledWith({
         id: 'sandbox-my-sandbox',
         env: {},
-        cwd: '/workspace',
+        cwd: '/workspace'
       });
     });
   });
@@ -199,19 +249,19 @@ describe('Sandbox - Automatic Session Management', () => {
       vi.mocked(sandbox.client.utils.createSession).mockResolvedValueOnce({
         success: true,
         id: 'custom-session-123',
-        message: 'Created',
+        message: 'Created'
       } as any);
 
       const session = await sandbox.createSession({
         id: 'custom-session-123',
         env: { NODE_ENV: 'test' },
-        cwd: '/test',
+        cwd: '/test'
       });
 
       expect(sandbox.client.utils.createSession).toHaveBeenCalledWith({
         id: 'custom-session-123',
         env: { NODE_ENV: 'test' },
-        cwd: '/test',
+        cwd: '/test'
       });
 
       expect(session.id).toBe('custom-session-123');
@@ -225,7 +275,7 @@ describe('Sandbox - Automatic Session Management', () => {
       vi.mocked(sandbox.client.utils.createSession).mockResolvedValueOnce({
         success: true,
         id: 'isolated-session',
-        message: 'Created',
+        message: 'Created'
       } as any);
 
       const session = await sandbox.createSession({ id: 'isolated-session' });
@@ -240,8 +290,16 @@ describe('Sandbox - Automatic Session Management', () => {
 
     it('should isolate multiple explicit sessions', async () => {
       vi.mocked(sandbox.client.utils.createSession)
-        .mockResolvedValueOnce({ success: true, id: 'session-1', message: 'Created' } as any)
-        .mockResolvedValueOnce({ success: true, id: 'session-2', message: 'Created' } as any);
+        .mockResolvedValueOnce({
+          success: true,
+          id: 'session-1',
+          message: 'Created'
+        } as any)
+        .mockResolvedValueOnce({
+          success: true,
+          id: 'session-2',
+          message: 'Created'
+        } as any);
 
       const session1 = await sandbox.createSession({ id: 'session-1' });
       const session2 = await sandbox.createSession({ id: 'session-2' });
@@ -249,8 +307,10 @@ describe('Sandbox - Automatic Session Management', () => {
       await session1.exec('echo build');
       await session2.exec('echo test');
 
-      const session1Id = vi.mocked(sandbox.client.commands.execute).mock.calls[0][1];
-      const session2Id = vi.mocked(sandbox.client.commands.execute).mock.calls[1][1];
+      const session1Id = vi.mocked(sandbox.client.commands.execute).mock
+        .calls[0][1];
+      const session2Id = vi.mocked(sandbox.client.commands.execute).mock
+        .calls[1][1];
 
       expect(session1Id).toBe('session-1');
       expect(session2Id).toBe('session-2');
@@ -259,19 +319,32 @@ describe('Sandbox - Automatic Session Management', () => {
 
     it('should not interfere with default session', async () => {
       vi.mocked(sandbox.client.utils.createSession)
-        .mockResolvedValueOnce({ success: true, id: 'sandbox-default', message: 'Created' } as any)
-        .mockResolvedValueOnce({ success: true, id: 'explicit-session', message: 'Created' } as any);
+        .mockResolvedValueOnce({
+          success: true,
+          id: 'sandbox-default',
+          message: 'Created'
+        } as any)
+        .mockResolvedValueOnce({
+          success: true,
+          id: 'explicit-session',
+          message: 'Created'
+        } as any);
 
       await sandbox.exec('echo default');
 
-      const explicitSession = await sandbox.createSession({ id: 'explicit-session' });
+      const explicitSession = await sandbox.createSession({
+        id: 'explicit-session'
+      });
       await explicitSession.exec('echo explicit');
 
       await sandbox.exec('echo default-again');
 
-      const defaultSessionId1 = vi.mocked(sandbox.client.commands.execute).mock.calls[0][1];
-      const explicitSessionId = vi.mocked(sandbox.client.commands.execute).mock.calls[1][1];
-      const defaultSessionId2 = vi.mocked(sandbox.client.commands.execute).mock.calls[2][1];
+      const defaultSessionId1 = vi.mocked(sandbox.client.commands.execute).mock
+        .calls[0][1];
+      const explicitSessionId = vi.mocked(sandbox.client.commands.execute).mock
+        .calls[1][1];
+      const defaultSessionId2 = vi.mocked(sandbox.client.commands.execute).mock
+        .calls[2][1];
 
       expect(defaultSessionId1).toBe('sandbox-default');
       expect(explicitSessionId).toBe('explicit-session');
@@ -284,7 +357,7 @@ describe('Sandbox - Automatic Session Management', () => {
       vi.mocked(sandbox.client.utils.createSession).mockResolvedValueOnce({
         success: true,
         id: 'session-generated-123',
-        message: 'Created',
+        message: 'Created'
       } as any);
 
       await sandbox.createSession();
@@ -292,7 +365,7 @@ describe('Sandbox - Automatic Session Management', () => {
       expect(sandbox.client.utils.createSession).toHaveBeenCalledWith({
         id: expect.stringMatching(/^session-/),
         env: undefined,
-        cwd: undefined,
+        cwd: undefined
       });
     });
   });
@@ -304,7 +377,7 @@ describe('Sandbox - Automatic Session Management', () => {
       vi.mocked(sandbox.client.utils.createSession).mockResolvedValueOnce({
         success: true,
         id: 'test-session',
-        message: 'Created',
+        message: 'Created'
       } as any);
 
       session = await sandbox.createSession({ id: 'test-session' });
@@ -312,7 +385,10 @@ describe('Sandbox - Automatic Session Management', () => {
 
     it('should execute command with session context', async () => {
       await session.exec('pwd');
-      expect(sandbox.client.commands.execute).toHaveBeenCalledWith('pwd', 'test-session');
+      expect(sandbox.client.commands.execute).toHaveBeenCalledWith(
+        'pwd',
+        'test-session'
+      );
     });
 
     it('should start process with session context', async () => {
@@ -323,8 +399,8 @@ describe('Sandbox - Automatic Session Management', () => {
           pid: 1234,
           command: 'sleep 10',
           status: 'running',
-          startTime: new Date().toISOString(),
-        },
+          startTime: new Date().toISOString()
+        }
       } as any);
 
       await session.startProcess('sleep 10');
@@ -340,7 +416,7 @@ describe('Sandbox - Automatic Session Management', () => {
       vi.spyOn(sandbox.client.files, 'writeFile').mockResolvedValue({
         success: true,
         path: '/test.txt',
-        timestamp: new Date().toISOString(),
+        timestamp: new Date().toISOString()
       } as any);
 
       await session.writeFile('/test.txt', 'content');
@@ -360,7 +436,7 @@ describe('Sandbox - Automatic Session Management', () => {
         stderr: '',
         branch: 'main',
         targetDir: '/workspace/repo',
-        timestamp: new Date().toISOString(),
+        timestamp: new Date().toISOString()
       } as any);
 
       await session.gitCheckout('https://github.com/test/repo.git');
@@ -379,7 +455,9 @@ describe('Sandbox - Automatic Session Management', () => {
         new Error('Session creation failed')
       );
 
-      await expect(sandbox.exec('echo test')).rejects.toThrow('Session creation failed');
+      await expect(sandbox.exec('echo test')).rejects.toThrow(
+        'Session creation failed'
+      );
     });
 
     it('should initialize with empty environment when not set', async () => {
@@ -388,7 +466,7 @@ describe('Sandbox - Automatic Session Management', () => {
       expect(sandbox.client.utils.createSession).toHaveBeenCalledWith({
         id: expect.any(String),
         env: {},
-        cwd: '/workspace',
+        cwd: '/workspace'
       });
     });
 
@@ -400,7 +478,7 @@ describe('Sandbox - Automatic Session Management', () => {
       expect(sandbox.client.utils.createSession).toHaveBeenCalledWith({
         id: expect.any(String),
         env: { NODE_ENV: 'production', DEBUG: 'true' },
-        cwd: '/workspace',
+        cwd: '/workspace'
       });
     });
   });
@@ -412,7 +490,7 @@ describe('Sandbox - Automatic Session Management', () => {
         success: true,
         port: 8080,
         name: 'test-service',
-        exposedAt: new Date().toISOString(),
+        exposedAt: new Date().toISOString()
       } as any);
     });
 
@@ -446,7 +524,10 @@ describe('Sandbox - Automatic Session Management', () => {
       ];
 
       for (const { hostname } of testCases) {
-        const result = await sandbox.exposePort(8080, { name: 'test', hostname });
+        const result = await sandbox.exposePort(8080, {
+          name: 'test',
+          hostname
+        });
         expect(result.url).toContain(hostname);
         expect(result.port).toBe(8080);
       }
@@ -462,4 +543,197 @@ describe('Sandbox - Automatic Session Management', () => {
       expect(sandbox.client.ports.exposePort).toHaveBeenCalled();
     });
   });
+
+  describe('fetch() override - WebSocket detection', () => {
+    let superFetchSpy: any;
+
+    beforeEach(async () => {
+      await sandbox.setSandboxName('test-sandbox');
+
+      // Spy on Container.prototype.fetch to verify WebSocket routing
+      superFetchSpy = vi
+        .spyOn(Container.prototype, 'fetch')
+        .mockResolvedValue(new Response('WebSocket response'));
+    });
+
+    afterEach(() => {
+      superFetchSpy?.mockRestore();
+    });
+
+    it('should detect WebSocket upgrade header and route to super.fetch', async () => {
+      const request = new Request('https://example.com/ws', {
+        headers: {
+          Upgrade: 'websocket',
+          Connection: 'Upgrade'
+        }
+      });
+
+      const response = await sandbox.fetch(request);
+
+      // Should route through super.fetch() for WebSocket
+      expect(superFetchSpy).toHaveBeenCalledTimes(1);
+      expect(await response.text()).toBe('WebSocket response');
+    });
+
+    it('should route non-WebSocket requests through containerFetch', async () => {
+      // GET request
+      const getRequest = new Request('https://example.com/api/data');
+      await sandbox.fetch(getRequest);
+      expect(superFetchSpy).not.toHaveBeenCalled();
+
+      vi.clearAllMocks();
+
+      // POST request
+      const postRequest = new Request('https://example.com/api/data', {
+        method: 'POST',
+        body: JSON.stringify({ data: 'test' }),
+        headers: { 'Content-Type': 'application/json' }
+      });
+      await sandbox.fetch(postRequest);
+      expect(superFetchSpy).not.toHaveBeenCalled();
+
+      vi.clearAllMocks();
+
+      // SSE request (should not be detected as WebSocket)
+      const sseRequest = new Request('https://example.com/events', {
+        headers: { Accept: 'text/event-stream' }
+      });
+      await sandbox.fetch(sseRequest);
+      expect(superFetchSpy).not.toHaveBeenCalled();
+    });
+
+    it('should preserve WebSocket request unchanged when calling super.fetch()', async () => {
+      const request = new Request('https://example.com/ws', {
+        headers: {
+          Upgrade: 'websocket',
+          Connection: 'Upgrade',
+          'Sec-WebSocket-Key': 'test-key-123',
+          'Sec-WebSocket-Version': '13'
+        }
+      });
+
+      await sandbox.fetch(request);
+
+      expect(superFetchSpy).toHaveBeenCalledTimes(1);
+      const passedRequest = superFetchSpy.mock.calls[0][0] as Request;
+      expect(passedRequest.headers.get('Upgrade')).toBe('websocket');
+      expect(passedRequest.headers.get('Connection')).toBe('Upgrade');
+      expect(passedRequest.headers.get('Sec-WebSocket-Key')).toBe(
+        'test-key-123'
+      );
+      expect(passedRequest.headers.get('Sec-WebSocket-Version')).toBe('13');
+    });
+  });
+
+  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);
+
+      const request = new Request('http://localhost/ws/echo', {
+        headers: {
+          Upgrade: 'websocket',
+          Connection: 'Upgrade'
+        }
+      });
+
+      const fetchSpy = vi.spyOn(sandbox, 'fetch');
+      const response = await sandbox.wsConnect(request, 8080);
+
+      // Verify switchPort was called with correct port
+      expect(switchPortMock).toHaveBeenCalledWith(request, 8080);
+
+      // Verify fetch was called with the switched request
+      expect(fetchSpy).toHaveBeenCalledOnce();
+
+      // Verify response indicates WebSocket upgrade
+      expect(response.status).toBe(200);
+      expect(response.headers.get('X-WebSocket-Upgraded')).toBe('true');
+    });
+
+    it('should reject invalid ports with SecurityError', async () => {
+      const request = new Request('http://localhost/ws/test', {
+        headers: { Upgrade: 'websocket', Connection: 'Upgrade' }
+      });
+
+      // Invalid port values
+      await expect(sandbox.wsConnect(request, -1)).rejects.toThrow(
+        'Invalid or restricted port'
+      );
+      await expect(sandbox.wsConnect(request, 0)).rejects.toThrow(
+        'Invalid or restricted port'
+      );
+      await expect(sandbox.wsConnect(request, 70000)).rejects.toThrow(
+        'Invalid or restricted port'
+      );
+
+      // Privileged ports
+      await expect(sandbox.wsConnect(request, 80)).rejects.toThrow(
+        'Invalid or restricted port'
+      );
+      await expect(sandbox.wsConnect(request, 443)).rejects.toThrow(
+        'Invalid or restricted port'
+      );
+    });
+
+    it('should preserve request properties through routing', async () => {
+      const request = new Request(
+        'http://localhost/ws/test?token=abc&room=lobby',
+        {
+          headers: {
+            Upgrade: 'websocket',
+            Connection: 'Upgrade',
+            'X-Custom-Header': 'custom-value'
+          }
+        }
+      );
+
+      const fetchSpy = vi.spyOn(sandbox, 'fetch');
+      await sandbox.wsConnect(request, 8080);
+
+      const calledRequest = fetchSpy.mock.calls[0][0];
+
+      // Verify headers are preserved
+      expect(calledRequest.headers.get('Upgrade')).toBe('websocket');
+      expect(calledRequest.headers.get('X-Custom-Header')).toBe('custom-value');
+
+      // Verify query parameters are preserved
+      const url = new URL(calledRequest.url);
+      expect(url.searchParams.get('token')).toBe('abc');
+      expect(url.searchParams.get('room')).toBe('lobby');
+    });
+  });
+
+  describe('deleteSession', () => {
+    it('should prevent deletion of default session', async () => {
+      // Trigger creation of default session
+      await sandbox.exec('echo "test"');
+
+      // Verify default session exists
+      expect((sandbox as any).defaultSession).toBeTruthy();
+      const defaultSessionId = (sandbox as any).defaultSession;
+
+      // Attempt to delete default session should throw
+      await expect(sandbox.deleteSession(defaultSessionId)).rejects.toThrow(
+        `Cannot delete default session '${defaultSessionId}'. Use sandbox.destroy() to terminate the sandbox.`
+      );
+    });
+
+    it('should allow deletion of non-default sessions', async () => {
+      // Mock the deleteSession API response
+      vi.spyOn(sandbox.client.utils, 'deleteSession').mockResolvedValue({
+        success: true,
+        sessionId: 'custom-session',
+        timestamp: new Date().toISOString()
+      });
+
+      // Create a custom session
+      await sandbox.createSession({ id: 'custom-session' });
+
+      // Should successfully delete non-default session
+      const result = await sandbox.deleteSession('custom-session');
+      expect(result.success).toBe(true);
+      expect(result.sessionId).toBe('custom-session');
+    });
+  });
 });