feat: persistent E2B sandbox #4
Conversation
|
Caution Review failedThe pull request is closed. WalkthroughAdds user-scoped tools and persistent sandbox management to the chat API (WorkOS-aware), introduces read/write file tools and a terminal tool factory, centralizes tool UI via ToolHandler, adds sandbox utilities and safe background wait, tweaks a code-highlight header, and adds @vercel/functions dependency. Minor .env comment fix. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant ChatAPI as Chat API (POST /api/chat)
participant WorkOS as Auth (WorkOS)
participant Tools as createTools
participant Manager as SandboxManager
participant Sandbox
participant Model as Streaming Model
Client->>ChatAPI: POST message
ChatAPI->>WorkOS: get session (if configured)
WorkOS-->>ChatAPI: userID or error
ChatAPI->>Tools: createTools(userID, mode)
Tools->>Manager: init with userID
ChatAPI->>Model: stream({ tools })
Model->>Manager: getSandbox()
Manager->>Sandbox: ensure connect/resume/create
Sandbox-->>Manager: sandbox handle
Model-->>ChatAPI: streamed response + tool calls
ChatAPI-->>Client: streamed response
ChatAPI->>Sandbox: pauseSandbox (onFinish, background)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Poem
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (2)
✨ Finishing Touches
🧪 Generate unit tests
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
|
The latest updates on your projects. Learn more about Vercel for Git ↗︎
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 10
🧹 Nitpick comments (8)
lib/utils/safe-wait-until.ts (1)
3-9: JSDoc parameter description doesn’t match the actual signatureDocs say “work - The async function…”, but the function currently accepts a Promise. Either update the docs or (preferably) change the API to accept a thunk (() => Promise) so synchronous errors are also captured and execution is deferred until waitUntil schedules it.
Here’s a doc-only fix if you keep the current signature:
- * @param work - The async function to execute in the background + * @param promise - The promise to execute in the backgroundlib/ai/tools/types.ts (3)
3-6: Consider simplifying getSandbox’s return typeReturning an object with a sandbox property adds ceremony without clear benefit. Consider returning the Sandbox directly (or Sandbox | null) to reduce boilerplate and improve ergonomics.
If you go that route:
export interface SandboxManager { getSandbox(): Promise<Sandbox | null>; setSandbox(sandbox: Sandbox): void; // optionally: // clearSandbox(): void; }This typically simplifies downstream usage and error handling.
8-11: SandboxContext shape may be too minimal for consumersExposing only setSandbox without a getter or a current value might force callers to re-fetch. Consider adding either a sandbox?: Sandbox field and/or a getSandbox accessor, or rely solely on SandboxManager to avoid duplicated concepts.
13-15: Document ToolContext to guide implementersAdd a brief JSDoc explaining that ToolContext is passed into tool factories and must be server-only (no client usage), since SandboxManager will touch server resources.
Example:
/** * Context passed to tool factories. Must be used server-side only. * Provides access to per-user sandbox lifecycle via sandboxManager. */ export interface ToolContext { sandboxManager: SandboxManager; }lib/ai/tools/index.ts (1)
13-21: Potential memory leak in sandbox callback closure.The callback function
(newSandbox) => { sandbox = newSandbox; }captures the localsandboxvariable, but the closure may not be properly cleaned up. Consider using a more explicit pattern.Apply this diff to use a more explicit callback pattern:
- let sandbox: Sandbox | null = null; + let sandbox: Sandbox | null = null; + + const setSandboxCallback = (newSandbox: Sandbox) => { + sandbox = newSandbox; + }; const sandboxManager = new DefaultSandboxManager( userID, - (newSandbox) => { - sandbox = newSandbox; - }, + setSandboxCallback, sandbox, );lib/ai/tools/utils/sandbox-manager.ts (1)
19-30: Redundant sandbox check after ensureSandboxConnection.Line 26 passes
this.sandbox(which isnullat this point) asinitialSandboxtoensureSandboxConnection, making this parameter effectively unused. The subsequent null check on line 32 is also redundant sinceensureSandboxConnectionshould always return a valid sandbox or throw.Apply this diff to simplify the logic:
if (!this.sandbox) { const result = await ensureSandboxConnection( { userID: this.userID, setSandbox: this.setSandboxCallback, }, - { - initialSandbox: this.sandbox, - }, ); this.sandbox = result.sandbox; } - if (!this.sandbox) { - throw new Error("Failed to initialize sandbox"); - } return { sandbox: this.sandbox };lib/ai/tools/read-file.ts (1)
18-23: Parameter naming consistency with write-file.This tool uses
target_filewhile write-file usesfile_path. Consider unifying to a single name (e.g.,file_path) across tools to minimize model confusion and UI branching.lib/ai/tools/run-terminal-cmd.ts (1)
10-22: Nit: fix contraction in tool description.Minor copy edit improves polish.
-7. Dont include any newlines in the command.`, +7. Don't include any newlines in the command.`,
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (1)
pnpm-lock.yamlis excluded by!**/pnpm-lock.yaml
📒 Files selected for processing (15)
.env.local.example(1 hunks)app/api/chat/route.ts(1 hunks)app/components/CodeHighlight.tsx(1 hunks)app/components/Message.tsx(2 hunks)app/components/ToolHandler.tsx(1 hunks)lib/ai/tools/index.ts(1 hunks)lib/ai/tools/read-file.ts(1 hunks)lib/ai/tools/run-terminal-cmd.ts(2 hunks)lib/ai/tools/types.ts(1 hunks)lib/ai/tools/utils/sandbox-manager.ts(1 hunks)lib/ai/tools/utils/sandbox-utils.ts(1 hunks)lib/ai/tools/utils/sandbox.ts(1 hunks)lib/ai/tools/write-file.ts(1 hunks)lib/utils/safe-wait-until.ts(1 hunks)package.json(1 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (8)
lib/ai/tools/write-file.ts (1)
lib/ai/tools/types.ts (1)
ToolContext(13-15)
lib/ai/tools/utils/sandbox-utils.ts (3)
lib/ai/tools/types.ts (1)
SandboxContext(8-11)lib/ai/tools/utils/sandbox.ts (1)
createOrConnectPersistentTerminal(21-135)lib/ai/tools/utils/sandbox-manager.ts (1)
setSandbox(39-42)
app/api/chat/route.ts (3)
lib/ai/tools/index.ts (1)
createTools(9-40)lib/ai/tools/utils/sandbox-manager.ts (1)
getSandbox(16-37)lib/ai/tools/utils/sandbox.ts (1)
pauseSandbox(150-167)
lib/ai/tools/utils/sandbox.ts (1)
lib/utils/safe-wait-until.ts (1)
safeWaitUntil(10-22)
lib/ai/tools/read-file.ts (1)
lib/ai/tools/types.ts (1)
ToolContext(13-15)
lib/ai/tools/utils/sandbox-manager.ts (2)
lib/ai/tools/types.ts (1)
SandboxManager(3-6)lib/ai/tools/utils/sandbox-utils.ts (1)
ensureSandboxConnection(8-38)
lib/ai/tools/index.ts (5)
lib/ai/tools/utils/sandbox-manager.ts (2)
DefaultSandboxManager(5-43)getSandbox(16-37)lib/ai/tools/types.ts (1)
ToolContext(13-15)lib/ai/tools/run-terminal-cmd.ts (1)
createRunTerminalCmd(6-55)lib/ai/tools/read-file.ts (1)
createReadFile(5-80)lib/ai/tools/write-file.ts (1)
createWriteFile(5-46)
lib/ai/tools/run-terminal-cmd.ts (1)
lib/ai/tools/types.ts (1)
ToolContext(13-15)
🪛 Biome (2.1.2)
app/components/ToolHandler.tsx
[error] 38-40: Other switch clauses can erroneously access this declaration.
Wrap the declaration in a block to restrict its access to the switch clause.
The declaration is defined in this switch clause:
Safe fix: Wrap the declaration in a block.
(lint/correctness/noSwitchDeclarations)
[error] 87-87: Other switch clauses can erroneously access this declaration.
Wrap the declaration in a block to restrict its access to the switch clause.
The declaration is defined in this switch clause:
Safe fix: Wrap the declaration in a block.
(lint/correctness/noSwitchDeclarations)
🔇 Additional comments (15)
package.json (1)
21-21: safeWaitUntil is confined to a Vercel API Route handlerAll invocations of
waitUntilviasafeWaitUntilare wrapped inpauseSandboxand only ever called inside the/app/api/chat/route.tsonFinish handler. No other call sites exist.• lib/utils/safe-wait-until.ts – defines
safeWaitUntil(importswaitUntil)
• lib/ai/tools/utils/sandbox.ts – callssafeWaitUntilinpauseSandbox()
• app/api/chat/route.ts –await pauseSandbox(sandbox)inside the server-side route’sonFinishcallbackNo usages of
safeWaitUntil(orwaitUntil) occur outside of a request context..env.local.example (1)
14-15: LGTM – comment fix improves clarityThe corrected comment is clearer and aligns with WorkOS dashboard-based configuration.
app/components/CodeHighlight.tsx (1)
99-105: LGTM – layout change keeps actions right-aligned even without a language badgeWrapping the left header segment with flex-1 ensures consistent spacing and alignment. Good change.
app/components/Message.tsx (2)
5-5: LGTM - Clean centralization of tool handling.The import of ToolHandler is appropriate and aligns with the centralized tool handling pattern.
50-61: LGTM - Elegant tool handling refactor.The refactor successfully centralizes tool handling by:
- Using a consistent pattern to extract
toolNamefrompart.type- Leveraging
part.toolCallIdas a stable key for React rendering- Delegating all tool-specific UI logic to the ToolHandler component
This approach is extensible and maintains clean separation of concerns.
lib/ai/tools/write-file.ts (1)
31-44: LGTM - Robust error handling and clear feedback.The implementation provides:
- Proper async/await usage with the sandbox API
- Comprehensive error handling with descriptive messages
- Helpful feedback showing the number of lines written
The error handling correctly distinguishes between Error instances and other types.
lib/ai/tools/index.ts (2)
34-35: LGTM - Appropriate tool filtering based on mode.The mode-based tool filtering is well-implemented:
- "ask" mode provides read-only access with just
readFile- Default "agent" mode provides full tool access
- Clean object destructuring maintains type safety
37-39: LGTM - Simple and effective sandbox accessor.The
getSandboxfunction provides a clean way to access the current sandbox state, properly returningnullwhen no sandbox is initialized.lib/ai/tools/utils/sandbox-manager.ts (1)
39-42: LGTM - Clean sandbox update pattern.The
setSandboxmethod properly updates both the internal state and notifies the callback, maintaining consistency across the sandbox management system.lib/ai/tools/utils/sandbox-utils.ts (2)
29-34: LGTM - Comprehensive error handling with context.The error handling properly:
- Logs the original error for debugging
- Provides a descriptive error message
- Preserves the original error details when available
- Maintains the error chain for troubleshooting
5-6: Ensure sandbox template name aligns with persistent sandbox and validate timeoutI only see
SANDBOX_TEMPLATEdefined once (in lib/ai/tools/utils/sandbox-utils.ts) as"temporary-sandbox". If this constant now powers your persistent‐sandbox feature, rename or update it to match the actual template (e.g."persistent-sandbox"). Also confirm that a 15-minute Bash timeout is sufficient for your use case—or bump it if longer‐lived sessions are expected.• lib/ai/tools/utils/sandbox-utils.ts:
const SANDBOX_TEMPLATE = "temporary-sandbox";const BASH_SANDBOX_TIMEOUT = 15 * 60 * 1000;app/api/chat/route.ts (2)
22-35: WorkOS-aware user scoping looks solid with safe fallbacks.Fetching the user ID via
authkit(req)only when WorkOS is configured, and falling back to"anonymous"on errors, is a pragmatic and safe approach.
37-48: Good: user-scoped tool provisioning and clean integration with streamText.Passing
toolsfromcreateTools(userID, mode)directly intostreamTextaligns the API with tool lifecycle ownership and keeps the route lean.lib/ai/tools/run-terminal-cmd.ts (1)
6-8: Good: context-driven sandbox access via SandboxManager.Refactoring to a factory that pulls the sandbox from
sandboxManagerkeeps lifecycle control in one place. This matches the persistent sandbox architecture introduced in this PR.app/components/ToolHandler.tsx (1)
15-21: Type narrowing OK; UI mapping matches tool payloads.The narrowing of
inputandoutputfor each tool path is sensible, and keys are stable viatoolCallId.
Summary by CodeRabbit
New Features
Improvements
Chores