diff --git a/lib/chat/__tests__/setupToolsForRequest.test.ts b/lib/chat/__tests__/setupToolsForRequest.test.ts index eab5488..633f837 100644 --- a/lib/chat/__tests__/setupToolsForRequest.test.ts +++ b/lib/chat/__tests__/setupToolsForRequest.test.ts @@ -10,6 +10,13 @@ vi.mock("@/lib/composio/toolRouter", () => ({ getComposioTools: vi.fn(), })); +vi.mock("@/lib/chat/tools/createPromptSandboxStreamingTool", () => ({ + createPromptSandboxStreamingTool: vi.fn(() => ({ + description: "Mock streaming sandbox tool", + parameters: {}, + })), +})); + // Import after mocks import { setupToolsForRequest } from "../setupToolsForRequest"; import { getMcpTools } from "@/lib/mcp/getMcpTools"; @@ -293,4 +300,50 @@ describe("setupToolsForRequest", () => { expect(mockGetComposioTools).toHaveBeenCalledTimes(1); }); }); + + describe("local streaming tool override", () => { + it("includes prompt_sandbox when authToken is provided", async () => { + const body: ChatRequestBody = { + accountId: "account-123", + orgId: null, + authToken: "test-token-123", + messages: [{ id: "1", role: "user", content: "Hello" }], + }; + + const result = await setupToolsForRequest(body); + + expect(result).toHaveProperty("prompt_sandbox"); + }); + + it("overrides MCP prompt_sandbox with local streaming version", async () => { + mockGetMcpTools.mockResolvedValue({ + prompt_sandbox: { description: "MCP version", parameters: {} }, + }); + + const body: ChatRequestBody = { + accountId: "account-123", + orgId: null, + authToken: "test-token-123", + messages: [{ id: "1", role: "user", content: "Hello" }], + }; + + const result = await setupToolsForRequest(body); + + expect(result.prompt_sandbox).toEqual( + expect.objectContaining({ description: "Mock streaming sandbox tool" }), + ); + }); + + it("does not include prompt_sandbox when authToken is not provided", async () => { + const body: ChatRequestBody = { + accountId: "account-123", + orgId: null, + messages: [{ id: "1", role: "user", content: "Hello" }], + }; + + const result = await setupToolsForRequest(body); + + expect(result).not.toHaveProperty("prompt_sandbox"); + }); + }); }); diff --git a/lib/chat/setupToolsForRequest.ts b/lib/chat/setupToolsForRequest.ts index ba586e1..71e0295 100644 --- a/lib/chat/setupToolsForRequest.ts +++ b/lib/chat/setupToolsForRequest.ts @@ -3,12 +3,14 @@ import { filterExcludedTools } from "./filterExcludedTools"; import { ChatRequestBody } from "./validateChatRequest"; import { getMcpTools } from "@/lib/mcp/getMcpTools"; import { getComposioTools } from "@/lib/composio/toolRouter"; +import { createPromptSandboxStreamingTool } from "@/lib/chat/tools/createPromptSandboxStreamingTool"; /** * Sets up and filters tools for a chat request. * Aggregates tools from: * - MCP server (via HTTP transport to /api/mcp for proper auth) * - Composio Tool Router (Google Sheets, Google Drive, Google Docs, TikTok) + * - Local streaming tools (override MCP versions for real-time output) * * @param body - The chat request body * @returns Filtered tool set ready for use @@ -22,10 +24,18 @@ export async function setupToolsForRequest(body: ChatRequestBody): Promise ({ + promptSandboxStreaming: (...args: unknown[]) => + mockPromptSandboxStreaming(...args), +})); + +// Helper to drain an async iterable into yields + return value +async function drainGenerator(iterable: AsyncIterable) { + const yields: unknown[] = []; + for await (const value of iterable) { + yields.push(value); + } + return yields; +} + +describe("createPromptSandboxStreamingTool", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("yields booting → streaming → complete statuses in order", async () => { + const finalResult = { + sandboxId: "sbx_123", + stdout: "Hello world", + stderr: "", + exitCode: 0, + created: false, + }; + + async function* fakeStreaming() { + yield { data: "Hello ", stream: "stdout" as const }; + yield { data: "world", stream: "stdout" as const }; + return finalResult; + } + + mockPromptSandboxStreaming.mockReturnValue(fakeStreaming()); + + const tool = createPromptSandboxStreamingTool("acc_123", "api-key-123"); + const iterable = tool.execute!({ prompt: "say hello" }, { + abortSignal: new AbortController().signal, + toolCallId: "tc_1", + messages: [], + }) as AsyncIterable; + + const yields = await drainGenerator(iterable); + + // First yield: booting + expect(yields[0]).toEqual({ + status: "booting", + output: "", + }); + + // Middle yields: streaming with accumulated stdout + expect(yields[1]).toEqual({ + status: "streaming", + output: "Hello ", + }); + expect(yields[2]).toEqual({ + status: "streaming", + output: "Hello world", + }); + + // Last yield: complete + expect(yields[3]).toEqual({ + status: "complete", + output: "Hello world", + stderr: "", + exitCode: 0, + }); + }); + + it("passes accountId, apiKey, and prompt to promptSandboxStreaming", async () => { + async function* fakeStreaming() { + return { + sandboxId: "sbx_123", + stdout: "", + stderr: "", + exitCode: 0, + created: false, + }; + } + + mockPromptSandboxStreaming.mockReturnValue(fakeStreaming()); + + const tool = createPromptSandboxStreamingTool("acc_456", "key_789"); + const iterable = tool.execute!({ prompt: "do stuff" }, { + abortSignal: new AbortController().signal, + toolCallId: "tc_2", + messages: [], + }) as AsyncIterable; + + await drainGenerator(iterable); + + expect(mockPromptSandboxStreaming).toHaveBeenCalledWith({ + accountId: "acc_456", + apiKey: "key_789", + prompt: "do stuff", + abortSignal: expect.any(AbortSignal), + }); + }); + + it("yields only stderr chunks in streaming status", async () => { + async function* fakeStreaming() { + yield { data: "warning!", stream: "stderr" as const }; + return { + sandboxId: "sbx_123", + stdout: "", + stderr: "warning!", + exitCode: 1, + created: false, + }; + } + + mockPromptSandboxStreaming.mockReturnValue(fakeStreaming()); + + const tool = createPromptSandboxStreamingTool("acc_1", "key_1"); + const iterable = tool.execute!({ prompt: "fail" }, { + abortSignal: new AbortController().signal, + toolCallId: "tc_3", + messages: [], + }) as AsyncIterable; + + const yields = await drainGenerator(iterable); + + // booting + expect(yields[0]).toEqual({ status: "booting", output: "" }); + + // streaming — stderr doesn't change output (only stdout does) + expect(yields[1]).toEqual({ status: "streaming", output: "" }); + + // complete — stderr is included + expect(yields[2]).toEqual({ + status: "complete", + output: "", + stderr: "warning!", + exitCode: 1, + }); + }); + + it("has the correct tool description and input schema", () => { + const tool = createPromptSandboxStreamingTool("acc_1", "key_1"); + + expect(tool.description).toContain("sandbox"); + expect(tool.inputSchema).toBeDefined(); + }); +}); diff --git a/lib/chat/tools/createPromptSandboxStreamingTool.ts b/lib/chat/tools/createPromptSandboxStreamingTool.ts new file mode 100644 index 0000000..2e95595 --- /dev/null +++ b/lib/chat/tools/createPromptSandboxStreamingTool.ts @@ -0,0 +1,88 @@ +import { z } from "zod"; +import type { Tool } from "ai"; +import { promptSandboxStreaming } from "@/lib/sandbox/promptSandboxStreaming"; + +const promptSandboxSchema = z.object({ + prompt: z + .string() + .min(1) + .describe("The prompt to send to OpenClaw running in the sandbox."), +}); + +interface SandboxStreamProgress { + status: "booting" | "streaming" | "complete"; + output: string; + stderr?: string; + exitCode?: number; +} + +interface PromptSandboxFinalResult { + sandboxId: string; + stdout: string; + stderr: string; + exitCode: number; + created: boolean; +} + +/** + * Creates a local AI SDK generator tool that streams sandbox output to the UI. + * Overrides the MCP prompt_sandbox tool with real-time streaming support. + * + * @param accountId - The account ID for sandbox lookup + * @param apiKey - The API key passed as RECOUP_API_KEY to the sandbox + * @returns An AI SDK tool with generator-based execute function + */ +export function createPromptSandboxStreamingTool( + accountId: string, + apiKey: string, +): Tool, SandboxStreamProgress> { + return { + description: + "Send a prompt to OpenClaw running in a persistent sandbox. " + + "Reuses the account's existing running sandbox or creates one from the latest snapshot. " + + "Streams output in real-time. The sandbox stays alive for follow-up prompts.", + inputSchema: promptSandboxSchema, + execute: async function* ({ prompt }, { abortSignal }) { + yield { status: "booting" as const, output: "" }; + + const gen = promptSandboxStreaming({ + accountId, + apiKey, + prompt, + abortSignal, + }); + + let stdout = ""; + let finalResult: PromptSandboxFinalResult | undefined; + + while (true) { + const iterResult = await gen.next(); + + if (iterResult.done) { + finalResult = iterResult.value as PromptSandboxFinalResult; + break; + } + + const chunk = iterResult.value as { + data: string; + stream: "stdout" | "stderr"; + }; + + if (chunk.stream === "stdout") { + stdout += chunk.data; + } + + yield { status: "streaming" as const, output: stdout }; + } + + yield { + status: "complete" as const, + output: finalResult!.stdout, + stderr: finalResult!.stderr, + exitCode: finalResult!.exitCode, + }; + + return finalResult as never; + }, + }; +} diff --git a/lib/mcp/tools/sandbox/__tests__/registerPromptSandboxTool.test.ts b/lib/mcp/tools/sandbox/__tests__/registerPromptSandboxTool.test.ts deleted file mode 100644 index df9a8e1..0000000 --- a/lib/mcp/tools/sandbox/__tests__/registerPromptSandboxTool.test.ts +++ /dev/null @@ -1,192 +0,0 @@ -import { describe, it, expect, vi, beforeEach } from "vitest"; -import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; -import type { RequestHandlerExtra } from "@modelcontextprotocol/sdk/shared/protocol.js"; -import type { ServerRequest, ServerNotification } from "@modelcontextprotocol/sdk/types.js"; - -import { registerPromptSandboxTool } from "../registerPromptSandboxTool"; - -const mockPromptSandbox = vi.fn(); -const mockResolveAccountId = vi.fn(); - -vi.mock("@/lib/sandbox/promptSandbox", () => ({ - promptSandbox: (...args: unknown[]) => mockPromptSandbox(...args), -})); - -vi.mock("@/lib/mcp/resolveAccountId", () => ({ - resolveAccountId: (...args: unknown[]) => mockResolveAccountId(...args), -})); - -type ServerRequestHandlerExtra = RequestHandlerExtra; - -function createMockExtra(authInfo?: { - accountId?: string; - orgId?: string | null; -}): ServerRequestHandlerExtra { - return { - authInfo: authInfo - ? { - token: "test-token", - scopes: ["mcp:tools"], - clientId: authInfo.accountId, - extra: { - accountId: authInfo.accountId, - orgId: authInfo.orgId ?? null, - }, - } - : undefined, - } as unknown as ServerRequestHandlerExtra; -} - -describe("registerPromptSandboxTool", () => { - let mockServer: McpServer; - let registeredHandler: (args: unknown, extra: ServerRequestHandlerExtra) => Promise; - - beforeEach(() => { - vi.clearAllMocks(); - - mockServer = { - registerTool: vi.fn((name, config, handler) => { - registeredHandler = handler; - }), - } as unknown as McpServer; - - registerPromptSandboxTool(mockServer); - }); - - it("registers the prompt_sandbox tool", () => { - expect(mockServer.registerTool).toHaveBeenCalledWith( - "prompt_sandbox", - expect.objectContaining({ - description: expect.any(String), - }), - expect.any(Function), - ); - }); - - it("returns error when resolveAccountId returns an error", async () => { - mockResolveAccountId.mockResolvedValue({ - accountId: null, - error: "Authentication required. Provide an API key via Authorization: Bearer header, or provide account_id from the system prompt context.", - }); - - const result = await registeredHandler( - { prompt: "hello" }, - createMockExtra(), - ); - - expect(result).toEqual({ - content: [ - { - type: "text", - text: expect.stringContaining("Authentication required"), - }, - ], - }); - }); - - it("returns error when resolveAccountId returns null accountId without error", async () => { - mockResolveAccountId.mockResolvedValue({ - accountId: null, - error: null, - }); - - const result = await registeredHandler( - { prompt: "hello" }, - createMockExtra(), - ); - - expect(result).toEqual({ - content: [ - { - type: "text", - text: expect.stringContaining("Failed to resolve account ID"), - }, - ], - }); - }); - - it("calls promptSandbox and returns success", async () => { - mockResolveAccountId.mockResolvedValue({ - accountId: "acc_123", - error: null, - }); - mockPromptSandbox.mockResolvedValue({ - sandboxId: "sbx_123", - stdout: "Hello world", - stderr: "", - exitCode: 0, - created: false, - }); - - const result = await registeredHandler( - { prompt: "say hello" }, - createMockExtra({ accountId: "acc_123" }), - ); - - expect(mockPromptSandbox).toHaveBeenCalledWith({ - accountId: "acc_123", - apiKey: "test-token", - prompt: "say hello", - }); - expect(result).toEqual({ - content: [ - { - type: "text", - text: expect.stringContaining('"sandboxId":"sbx_123"'), - }, - ], - }); - expect(result).toEqual({ - content: [ - { - type: "text", - text: expect.stringContaining('"stdout":"Hello world"'), - }, - ], - }); - }); - - it("resolves accountId from auth only without override", async () => { - mockResolveAccountId.mockResolvedValue({ - accountId: "acc_123", - error: null, - }); - mockPromptSandbox.mockResolvedValue({ - sandboxId: "sbx_123", - stdout: "", - stderr: "", - exitCode: 0, - created: false, - }); - - const extra = createMockExtra({ accountId: "acc_123" }); - await registeredHandler({ prompt: "test" }, extra); - - expect(mockResolveAccountId).toHaveBeenCalledWith({ - authInfo: extra.authInfo, - accountIdOverride: undefined, - }); - }); - - it("returns error when promptSandbox throws", async () => { - mockResolveAccountId.mockResolvedValue({ - accountId: "acc_123", - error: null, - }); - mockPromptSandbox.mockRejectedValue(new Error("Sandbox timed out")); - - const result = await registeredHandler( - { prompt: "hello" }, - createMockExtra({ accountId: "acc_123" }), - ); - - expect(result).toEqual({ - content: [ - { - type: "text", - text: expect.stringContaining("Sandbox timed out"), - }, - ], - }); - }); -}); diff --git a/lib/mcp/tools/sandbox/index.ts b/lib/mcp/tools/sandbox/index.ts index bedaf2a..5587ea0 100644 --- a/lib/mcp/tools/sandbox/index.ts +++ b/lib/mcp/tools/sandbox/index.ts @@ -1,13 +1,12 @@ import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; -import { registerPromptSandboxTool } from "./registerPromptSandboxTool"; import { registerRunSandboxCommandTool } from "./registerRunSandboxCommandTool"; /** * Registers all sandbox-related MCP tools on the server. + * Note: prompt_sandbox is now a local streaming tool in setupToolsForRequest. * * @param server - The MCP server instance to register tools on. */ export const registerAllSandboxTools = (server: McpServer): void => { - registerPromptSandboxTool(server); registerRunSandboxCommandTool(server); }; diff --git a/lib/mcp/tools/sandbox/registerPromptSandboxTool.ts b/lib/mcp/tools/sandbox/registerPromptSandboxTool.ts deleted file mode 100644 index 402f45e..0000000 --- a/lib/mcp/tools/sandbox/registerPromptSandboxTool.ts +++ /dev/null @@ -1,62 +0,0 @@ -import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; -import type { RequestHandlerExtra } from "@modelcontextprotocol/sdk/shared/protocol.js"; -import type { ServerRequest, ServerNotification } from "@modelcontextprotocol/sdk/types.js"; -import { z } from "zod"; -import type { McpAuthInfo } from "@/lib/mcp/verifyApiKey"; -import { resolveAccountId } from "@/lib/mcp/resolveAccountId"; -import { getToolResultSuccess } from "@/lib/mcp/getToolResultSuccess"; -import { getToolResultError } from "@/lib/mcp/getToolResultError"; -import { promptSandbox } from "@/lib/sandbox/promptSandbox"; - -const promptSandboxSchema = z.object({ - prompt: z - .string() - .min(1) - .describe("The prompt to send to OpenClaw running in the sandbox."), -}); - -/** - * Registers the "prompt_sandbox" tool on the MCP server. - * Sends a prompt to OpenClaw in a persistent per-account sandbox and returns the output. - * - * @param server - The MCP server instance to register the tool on. - */ -export function registerPromptSandboxTool(server: McpServer): void { - server.registerTool( - "prompt_sandbox", - { - description: - "Send a prompt to OpenClaw running in a persistent sandbox. Reuses the account's existing running sandbox or creates one from the latest snapshot. Returns raw stdout/stderr. The sandbox stays alive for follow-up prompts.", - inputSchema: promptSandboxSchema, - }, - async (args, extra: RequestHandlerExtra) => { - const authInfo = extra.authInfo as McpAuthInfo | undefined; - const { accountId, error } = await resolveAccountId({ - authInfo, - accountIdOverride: undefined, - }); - - if (error) { - return getToolResultError(error); - } - - if (!accountId) { - return getToolResultError("Failed to resolve account ID"); - } - - try { - const result = await promptSandbox({ - accountId, - apiKey: authInfo!.token, - prompt: args.prompt, - }); - - return getToolResultSuccess(result); - } catch (error) { - const message = - error instanceof Error ? error.message : "Failed to prompt sandbox"; - return getToolResultError(message); - } - }, - ); -} diff --git a/lib/sandbox/__tests__/promptSandbox.test.ts b/lib/sandbox/__tests__/promptSandbox.test.ts deleted file mode 100644 index 06c98c3..0000000 --- a/lib/sandbox/__tests__/promptSandbox.test.ts +++ /dev/null @@ -1,134 +0,0 @@ -import { describe, it, expect, vi, beforeEach } from "vitest"; -import type { Sandbox } from "@vercel/sandbox"; - -import { promptSandbox } from "../promptSandbox"; - -const mockGetOrCreateSandbox = vi.fn(); - -vi.mock("../getOrCreateSandbox", () => ({ - getOrCreateSandbox: (...args: unknown[]) => mockGetOrCreateSandbox(...args), -})); - -describe("promptSandbox", () => { - const mockRunCommand = vi.fn(); - const mockSandbox = { - sandboxId: "sbx_123", - status: "running", - runCommand: mockRunCommand, - } as unknown as Sandbox; - - beforeEach(() => { - vi.clearAllMocks(); - }); - - it("returns stdout/stderr from successful prompt", async () => { - mockGetOrCreateSandbox.mockResolvedValue({ - sandbox: mockSandbox, - sandboxId: "sbx_123", - created: false, - }); - - const mockResult = { - stdout: vi.fn().mockResolvedValue("Hello from sandbox"), - stderr: vi.fn().mockResolvedValue(""), - exitCode: 0, - }; - mockRunCommand.mockResolvedValue(mockResult); - - const result = await promptSandbox({ - accountId: "acc_1", - apiKey: "key_abc", - prompt: "say hello", - }); - - expect(result).toEqual({ - sandboxId: "sbx_123", - stdout: "Hello from sandbox", - stderr: "", - exitCode: 0, - created: false, - }); - }); - - it("passes caller's API key as RECOUP_API_KEY env var", async () => { - mockGetOrCreateSandbox.mockResolvedValue({ - sandbox: mockSandbox, - sandboxId: "sbx_123", - created: false, - }); - - const mockResult = { - stdout: vi.fn().mockResolvedValue(""), - stderr: vi.fn().mockResolvedValue(""), - exitCode: 0, - }; - mockRunCommand.mockResolvedValue(mockResult); - - await promptSandbox({ - accountId: "acc_1", - apiKey: "caller-api-key", - prompt: "do something", - }); - - expect(mockRunCommand).toHaveBeenCalledWith({ - cmd: "openclaw", - args: ["agent", "--agent", "main", "--message", "do something"], - env: { - RECOUP_API_KEY: "caller-api-key", - }, - }); - }); - - it("reports created=true when sandbox was newly created", async () => { - mockGetOrCreateSandbox.mockResolvedValue({ - sandbox: mockSandbox, - sandboxId: "sbx_new", - created: true, - }); - - const mockResult = { - stdout: vi.fn().mockResolvedValue("setup done"), - stderr: vi.fn().mockResolvedValue(""), - exitCode: 0, - }; - mockRunCommand.mockResolvedValue(mockResult); - - const result = await promptSandbox({ - accountId: "acc_1", - apiKey: "key_abc", - prompt: "setup", - }); - - expect(result.created).toBe(true); - expect(result.sandboxId).toBe("sbx_new"); - }); - - it("handles non-zero exit code", async () => { - mockGetOrCreateSandbox.mockResolvedValue({ - sandbox: mockSandbox, - sandboxId: "sbx_123", - created: false, - }); - - const mockResult = { - stdout: vi.fn().mockResolvedValue(""), - stderr: vi.fn().mockResolvedValue("command not found"), - exitCode: 127, - }; - mockRunCommand.mockResolvedValue(mockResult); - - const result = await promptSandbox({ - accountId: "acc_1", - apiKey: "key_abc", - prompt: "bad command", - }); - - expect(result).toEqual({ - sandboxId: "sbx_123", - stdout: "", - stderr: "command not found", - exitCode: 127, - created: false, - }); - }); -}); diff --git a/lib/sandbox/__tests__/promptSandboxStreaming.test.ts b/lib/sandbox/__tests__/promptSandboxStreaming.test.ts new file mode 100644 index 0000000..ef5e07b --- /dev/null +++ b/lib/sandbox/__tests__/promptSandboxStreaming.test.ts @@ -0,0 +1,231 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import type { Sandbox } from "@vercel/sandbox"; + +import { promptSandboxStreaming } from "../promptSandboxStreaming"; + +const mockGetOrCreateSandbox = vi.fn(); + +vi.mock("../getOrCreateSandbox", () => ({ + getOrCreateSandbox: (...args: unknown[]) => mockGetOrCreateSandbox(...args), +})); + +describe("promptSandboxStreaming", () => { + const mockRunCommand = vi.fn(); + const mockSandbox = { + sandboxId: "sbx_123", + status: "running", + runCommand: mockRunCommand, + } as unknown as Sandbox; + + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("yields log chunks in order and returns final result", async () => { + mockGetOrCreateSandbox.mockResolvedValue({ + sandbox: mockSandbox, + sandboxId: "sbx_123", + created: false, + }); + + async function* fakeLogs() { + yield { data: "Hello ", stream: "stdout" as const }; + yield { data: "world", stream: "stdout" as const }; + } + + const mockCmd = { + logs: () => fakeLogs(), + wait: vi.fn().mockResolvedValue({ exitCode: 0 }), + }; + mockRunCommand.mockResolvedValue(mockCmd); + + const chunks: Array<{ data: string; stream: "stdout" | "stderr" }> = []; + let finalResult; + + const gen = promptSandboxStreaming({ + accountId: "acc_1", + apiKey: "key_abc", + prompt: "say hello", + }); + + while (true) { + const result = await gen.next(); + if (result.done) { + finalResult = result.value; + break; + } + chunks.push( + result.value as { data: string; stream: "stdout" | "stderr" }, + ); + } + + expect(chunks).toEqual([ + { data: "Hello ", stream: "stdout" }, + { data: "world", stream: "stdout" }, + ]); + + expect(finalResult).toEqual({ + sandboxId: "sbx_123", + stdout: "Hello world", + stderr: "", + exitCode: 0, + created: false, + }); + }); + + it("accumulates stderr separately", async () => { + mockGetOrCreateSandbox.mockResolvedValue({ + sandbox: mockSandbox, + sandboxId: "sbx_123", + created: false, + }); + + async function* fakeLogs() { + yield { data: "output", stream: "stdout" as const }; + yield { data: "warn: something", stream: "stderr" as const }; + } + + const mockCmd = { + logs: () => fakeLogs(), + wait: vi.fn().mockResolvedValue({ exitCode: 0 }), + }; + mockRunCommand.mockResolvedValue(mockCmd); + + let finalResult; + + const gen = promptSandboxStreaming({ + accountId: "acc_1", + apiKey: "key_abc", + prompt: "test", + }); + + while (true) { + const result = await gen.next(); + if (result.done) { + finalResult = result.value; + break; + } + } + + expect(finalResult).toEqual({ + sandboxId: "sbx_123", + stdout: "output", + stderr: "warn: something", + exitCode: 0, + created: false, + }); + }); + + it("uses detached mode with runCommand", async () => { + mockGetOrCreateSandbox.mockResolvedValue({ + sandbox: mockSandbox, + sandboxId: "sbx_123", + created: false, + }); + + async function* fakeLogs() { + yield { data: "done", stream: "stdout" as const }; + } + + const mockCmd = { + logs: () => fakeLogs(), + wait: vi.fn().mockResolvedValue({ exitCode: 0 }), + }; + mockRunCommand.mockResolvedValue(mockCmd); + + const gen = promptSandboxStreaming({ + accountId: "acc_1", + apiKey: "key_abc", + prompt: "run", + }); + + // Drain the generator + for await (const _ of gen) { + // consume + } + + expect(mockRunCommand).toHaveBeenCalledWith({ + cmd: "openclaw", + args: ["agent", "--agent", "main", "--message", "run"], + env: { RECOUP_API_KEY: "key_abc" }, + detached: true, + }); + }); + + it("reports created=true when sandbox was newly created", async () => { + mockGetOrCreateSandbox.mockResolvedValue({ + sandbox: mockSandbox, + sandboxId: "sbx_new", + created: true, + }); + + async function* fakeLogs() { + yield { data: "setup done", stream: "stdout" as const }; + } + + const mockCmd = { + logs: () => fakeLogs(), + wait: vi.fn().mockResolvedValue({ exitCode: 0 }), + }; + mockRunCommand.mockResolvedValue(mockCmd); + + const gen = promptSandboxStreaming({ + accountId: "acc_1", + apiKey: "key_abc", + prompt: "setup", + }); + + let finalResult; + while (true) { + const result = await gen.next(); + if (result.done) { + finalResult = result.value; + break; + } + } + + expect(finalResult!.created).toBe(true); + expect(finalResult!.sandboxId).toBe("sbx_new"); + }); + + it("handles non-zero exit code", async () => { + mockGetOrCreateSandbox.mockResolvedValue({ + sandbox: mockSandbox, + sandboxId: "sbx_123", + created: false, + }); + + async function* fakeLogs() { + yield { data: "error output", stream: "stderr" as const }; + } + + const mockCmd = { + logs: () => fakeLogs(), + wait: vi.fn().mockResolvedValue({ exitCode: 1 }), + }; + mockRunCommand.mockResolvedValue(mockCmd); + + const gen = promptSandboxStreaming({ + accountId: "acc_1", + apiKey: "key_abc", + prompt: "bad command", + }); + + let finalResult; + while (true) { + const result = await gen.next(); + if (result.done) { + finalResult = result.value; + break; + } + } + + expect(finalResult).toEqual({ + sandboxId: "sbx_123", + stdout: "", + stderr: "error output", + exitCode: 1, + created: false, + }); + }); +}); diff --git a/lib/sandbox/promptSandbox.ts b/lib/sandbox/promptSandbox.ts deleted file mode 100644 index e02f6cd..0000000 --- a/lib/sandbox/promptSandbox.ts +++ /dev/null @@ -1,51 +0,0 @@ -import { getOrCreateSandbox } from "./getOrCreateSandbox"; - -interface PromptSandboxInput { - accountId: string; - apiKey: string; - prompt: string; -} - -interface PromptSandboxResult { - sandboxId: string; - stdout: string; - stderr: string; - exitCode: number; - created: boolean; -} - -/** - * Sends a prompt to OpenClaw running inside a persistent per-account sandbox. - * Reuses an existing running sandbox or creates one from the latest snapshot. - * The sandbox stays alive after the prompt for follow-up prompts. - * - * @param input - The account ID and prompt to run - * @returns The sandbox ID, command output, and whether the sandbox was newly created - */ -export async function promptSandbox( - input: PromptSandboxInput, -): Promise { - const { accountId, apiKey, prompt } = input; - - const { sandbox, sandboxId, created } = - await getOrCreateSandbox(accountId); - - const commandResult = await sandbox.runCommand({ - cmd: "openclaw", - args: ["agent", "--agent", "main", "--message", prompt], - env: { - RECOUP_API_KEY: apiKey, - }, - }); - - const stdout = (await commandResult.stdout()) || ""; - const stderr = (await commandResult.stderr()) || ""; - - return { - sandboxId, - stdout, - stderr, - exitCode: commandResult.exitCode, - created, - }; -} diff --git a/lib/sandbox/promptSandboxStreaming.ts b/lib/sandbox/promptSandboxStreaming.ts new file mode 100644 index 0000000..1b9ec52 --- /dev/null +++ b/lib/sandbox/promptSandboxStreaming.ts @@ -0,0 +1,68 @@ +import { getOrCreateSandbox } from "./getOrCreateSandbox"; + +interface PromptSandboxStreamingInput { + accountId: string; + apiKey: string; + prompt: string; + abortSignal?: AbortSignal; +} + +interface PromptSandboxStreamingResult { + sandboxId: string; + stdout: string; + stderr: string; + exitCode: number; + created: boolean; +} + +/** + * Streams output from OpenClaw running inside a persistent per-account sandbox. + * Yields log chunks as they arrive, then returns the full result. + * + * @param input - The account ID, API key, prompt, and optional abort signal + * @yields Log chunks with data and stream type (stdout/stderr) + * @returns The sandbox ID, accumulated output, exit code, and whether the sandbox was newly created + */ +export async function* promptSandboxStreaming( + input: PromptSandboxStreamingInput, +): AsyncGenerator< + { data: string; stream: "stdout" | "stderr" }, + PromptSandboxStreamingResult, + undefined +> { + const { accountId, apiKey, prompt, abortSignal } = input; + + const { sandbox, sandboxId, created } = + await getOrCreateSandbox(accountId); + + const cmd = await sandbox.runCommand({ + cmd: "openclaw", + args: ["agent", "--agent", "main", "--message", prompt], + env: { + RECOUP_API_KEY: apiKey, + }, + detached: true, + }); + + let stdout = ""; + let stderr = ""; + + for await (const log of cmd.logs({ signal: abortSignal })) { + if (log.stream === "stdout") { + stdout += log.data; + } else { + stderr += log.data; + } + yield log; + } + + const { exitCode } = await cmd.wait(); + + return { + sandboxId, + stdout, + stderr, + exitCode, + created, + }; +}