websockets support

[whoiskatrin]

No description provided.

[changeset-bot]

⚠️ No Changeset found

Latest commit: 7867acd

Merging this PR will not cause a version bump for any packages. If these changes should not result in a new version, you're good to go. If these changes should result in a version bump, you need to add a changeset.

This PR includes no changesets

When changesets are added to this PR, you'll see the packages that this PR includes changesets for and the associated semver types

Click here to learn what changesets are, and how to add one.

Click here if you're a maintainer who wants to add a changeset to this PR

[claude]

WebSocket Support PR Review

Thank you for adding WebSocket support to the Sandbox SDK! This is a valuable feature that enables real-time bidirectional communication. Below is a comprehensive review covering code quality, potential issues, and recommendations.

---

✅ Strengths

1. Excellent API Design

• The SandboxWebSocket wrapper class provides a clean, type-safe abstraction over raw WebSockets
• The createWebSocketHandler function handles all boilerplate setup elegantly
• Helper methods like sendReady(), sendError(), sendOutput() make the API intuitive
• The separation between raw WebSocket operations and sandbox-aware helpers is well-designed

2. Comprehensive Documentation

• The README is extremely thorough with multiple use cases and examples
• Code examples cover various languages (JavaScript, Python, Node.js)
• Best practices section is valuable (error handling, heartbeat, rate limiting, security)

3. Good Code Organization

• Clear separation of concerns with typed message interfaces
• Proper exports in index.ts
• Both basic and helper-based examples provided for comparison

---

🐛 Potential Bugs & Issues

Critical

1. Type Safety Issue in index-with-helpers.ts:70

const result = await ws.raw.sandbox.exec(message.command);

Problem: ws.raw returns the WebSocket instance, not a wrapper with a sandbox property. This will cause a runtime error.

Fix: Access the sandbox from the handler result:

async function handleEchoWebSocket(request: Request, env: Env): Promise<Response> {
  const { response, websocket, sandbox } = await createWebSocketHandler(request, env.Sandbox, {
    onMessage: async (ws, message) => {
      if (message.type === "execute") {
        const result = await sandbox.exec(message.command);
        ws.sendResult({...});
      }
    }
  });
  return response;
}

2. Missing Error Handling in Stream Processing

In websocket.ts:145-155, the SSE stream iteration could fail unexpectedly:

for await (const event of parseSSEStream<LogEvent>(logStream)) {
  // No try-catch here - errors will propagate
}

Recommendation: Add granular error handling within the loop to continue processing other events if one fails.

3. Race Condition in startProcessWithStreaming

// Stream logs in the background
this.streamProcessLogs(process).catch(error => {
  console.error('Error streaming process logs:', error);
});

Problem: If the process exits immediately or the stream fails to start, the client may miss output. The catch only logs to console - client isn't notified.

Fix: Send error to client:

this.streamProcessLogs(process).catch(error => {
  this.sendError(error);
});

Medium Priority

4. Potential Memory Leak in Examples

In index.ts:299-330, the async IIFE for streaming logs has no cleanup mechanism if the WebSocket closes before the stream completes:

(async () => {
  try {
    const logStream = await sandbox.streamProcessLogs(currentProcess.id);
    for await (const event of parseSSEStream<LogEvent>(logStream)) {
      server.send(...); // This could fail if WebSocket is closed
    }
  }
})();

Fix: Track the promise and abort it on WebSocket close:

let streamAbortController = new AbortController();

server.addEventListener("close", async () => {
  streamAbortController.abort();
  // ... cleanup
});

5. Missing WebSocket State Checks

In websocket.ts:66-70, server.send() doesn't check if the WebSocket is still open:

send(message: WebSocketMessage): void {
  try {
    this.server.send(JSON.stringify(message));
  } catch (error) {
    console.error('Failed to send WebSocket message:', error);
  }
}

Recommendation: Check readyState before sending:

send(message: WebSocketMessage): void {
  if (this.server.readyState !== WebSocket.OPEN) {
    console.warn('WebSocket not open, skipping message');
    return;
  }
  try {
    this.server.send(JSON.stringify(message));
  } catch (error) {
    console.error('Failed to send WebSocket message:', error);
  }
}

6. Inconsistent Error Handling Pattern

Some methods throw errors after sending them to the client (e.g., runCodeWithStreaming:198), while others don't. This inconsistency can lead to unhandled promise rejections.

} catch (error) {
  this.sendError(error as Error);
  throw error; // Should callers expect this to throw?
}

Recommendation: Document the throw behavior or make it consistent.

---

🔒 Security Concerns

1. Command Injection Risk in Examples

The README shows examples without input validation:

case "execute":
  const result = await sandbox.exec(message.command); // Unsanitized user input

While the sandbox provides isolation, the README should emphasize input validation in production code.

Add to README:

// Validate and sanitize inputs
if (!message.command || typeof message.command !== 'string') {
  ws.sendError('Invalid command format');
  return;
}

// Consider whitelisting allowed commands
const allowedCommands = ['python', 'node', 'npm'];
const cmd = message.command.split(' ')[0];
if (!allowedCommands.includes(cmd)) {
  ws.sendError('Command not allowed');
  return;
}

2. Missing Authentication in Helper

The createWebSocketHandler function doesn't enforce authentication. The README mentions it in "Security Considerations" but the code doesn't demonstrate it.

Recommendation: Add an authenticate callback to options:

export interface WebSocketHandlerOptions {
  authenticate?: (request: Request) => Promise<boolean> | boolean;
  // ...
}

// In createWebSocketHandler:
if (options.authenticate && !(await options.authenticate(request))) {
  return {
    response: new Response('Unauthorized', { status: 401 }),
    // ... null values for other fields
  };
}

3. No Rate Limiting in SDK

While the README discusses rate limiting for clients, the SDK doesn't provide server-side rate limiting. Consider adding per-connection rate limiting to the wrapper class.

4. Unrestricted Sandbox ID Generation

const sandboxId =
  options.sandboxId ||
  url.searchParams.get('id') ||
  crypto.randomUUID();

Allowing arbitrary sandbox IDs from URL params could enable unauthorized access to existing sandboxes. Consider validating or hashing the ID.

---

⚡ Performance Considerations

1. JSON Parsing in Hot Path

In createWebSocketHandler:368, every message is parsed twice (once to check type, again in handler):

server.addEventListener('message', async (event) => {
  try {
    const message = JSON.parse(event.data as string);
    await options.onMessage!(ws, message, event);
  }
})

Impact: Low for typical workloads, but consider if handling high-frequency messages.

2. Unbounded Memory in Stream Buffering

The SSE stream parsing doesn't limit buffer size. A malicious client could send huge chunks.

Recommendation: Add size limits to parseSSEStream or implement backpressure.

3. No Connection Limiting

The examples don't show how to limit concurrent WebSocket connections per sandbox, which could lead to resource exhaustion.

---

🧪 Test Coverage

Missing Tests

The PR adds new functionality but I don't see corresponding test files:

Required Tests:

1. Unit tests for websocket.ts:

   • SandboxWebSocket class methods
   • createWebSocketHandler with various options
   • Error handling paths
   • Message serialization/deserialization

2. Integration tests:

   • End-to-end WebSocket connection lifecycle
   • Process streaming with WebSocket
   • Code execution with streaming output
   • Concurrent connections to same sandbox

3. Error scenario tests:

   • WebSocket closes during streaming
   • Sandbox process exits unexpectedly
   • Invalid message formats
   • Network interruptions

Recommendation: Add packages/sandbox/tests/websocket.test.ts with comprehensive test coverage following the patterns in existing test files.

---

📝 Code Quality & Best Practices

Minor Issues

1. Type Assertion Overhead

Multiple places use as casts that could be avoided:

ws.send({
  type: 'ready',
  message,
  sandboxId,
} as ReadyMessage);

Since TypeScript can infer the type, consider restructuring to avoid casts.

2. Inconsistent Naming

• runCodeWithStreaming vs startProcessWithStreaming - mixing "run" and "start"
• Consider: executeCodeWithStreaming and startProcessWithStreaming for consistency

3. Magic Numbers

webSocket: client, // Position [1] in WebSocketPair

Consider destructuring with named variables for clarity:

const [clientWebSocket, serverWebSocket] = Object.values(pair);

4. Console Logging in Production Code

Multiple instances of console.error() and console.log() that should use a proper logging system:

console.error('Failed to send WebSocket message:', error);

Consider using the SDK's logging infrastructure if available.

Documentation Improvements

1. Add JSDoc Examples

The methods have good descriptions but lack usage examples:

/**
 * Stream process logs to the WebSocket client
 * 
 * @example
 * ```typescript
 * const process = await sandbox.startProcess('npm test');
 * await ws.streamProcessLogs(process);
 * ```
 * 
 * @param process - The process to stream logs from
 * @returns Promise that resolves when streaming completes
 */
async streamProcessLogs(process: Process): Promise<void>

2. README: Terminal Example Issue

Line 189 in index-with-helpers.ts uses a non-portable approach:

await sandbox.exec(`echo "${message.data}" | /proc/${shellProcess.pid}/fd/0`);

This won't work (stdin isn't writable via /proc) and should be clearly marked as "not yet implemented" or removed.

---

📋 Additional Recommendations

1. Add WebSocket Close Codes

Use proper WebSocket close codes throughout:

ws.close(1000); // Normal closure
ws.close(1008); // Policy violation  
ws.close(1011); // Internal server error

The examples use these but the wrapper class doesn't expose them conveniently.

2. Implement Ping/Pong

Add built-in heartbeat mechanism to SandboxWebSocket:

private startHeartbeat(interval: number = 30000) {
  this.heartbeatInterval = setInterval(() => {
    if (this.server.readyState === WebSocket.OPEN) {
      this.send({ type: 'ping' });
    }
  }, interval);
}

3. Add TypeScript Strict Mode

Ensure tsconfig.json in the websocket example has strict mode enabled:

{
  "compilerOptions": {
    "strict": true,
    "noUncheckedIndexedAccess": true
  }
}

4. Export Types More Granularly

The index.ts exports everything with export * which can cause namespace pollution. Consider explicit exports.

5. Add Migration Guide

For users upgrading from basic WebSocket usage to the helper-based approach, add a migration section to the README.

---

🎯 Summary

This is a well-designed feature that adds significant value to the Sandbox SDK. The API is clean, the documentation is thorough, and the examples are helpful.

Before merging, please address:

1. ✋ Critical: Fix the ws.raw.sandbox bug in index-with-helpers.ts:70
2. ⚠️ Improve error handling in stream processing
3. 🧪 Add comprehensive test coverage
4. 🔒 Strengthen security guidance in README
5. 📚 Document error throwing behavior

Nice to have:

• Server-side rate limiting
• Built-in heartbeat mechanism
• Authentication callback in helper
• Connection limiting examples

Overall Assessment: ✅ Approve with requested changes

This PR demonstrates strong engineering and thoughtful design. With the bug fixes and additional test coverage, it will be a great addition to the SDK. Great work! 🎉