diff --git a/src/computer/index.ts b/src/computer/index.ts index 1bc735c5..26cbd00c 100644 --- a/src/computer/index.ts +++ b/src/computer/index.ts @@ -15,6 +15,47 @@ import { DEFAULT_WORKDIR } from "./definitions"; import { newProcess, Process } from "./process"; +const VALID_ENV_VAR_NAME = /^[A-Za-z_][A-Za-z0-9_]*$/; + +function buildShellCommand( + cmd: string, + cwd: string, + env: Record, + emitPidMarker: boolean, +): Result { + for (const key of Object.keys(env)) { + if (!VALID_ENV_VAR_NAME.test(key)) { + return err( + "invalid_parameters_error", + `Invalid environment variable name: ${key}`, + ); + } + } + + const script = [ + "set -e", + 'while [ "$#" -gt 2 ]; do', + ' export "$1=$2"', + ' shift 2', + "done", + 'cd -- "$1"', + emitPidMarker ? 'echo "SRCHD_PID:$$" >&2' : undefined, + 'exec /bin/bash -lc "$2"', + ].filter((line): line is string => line !== undefined).join("\n"); + + const envArgs = Object.entries(env).flatMap(([key, value]) => [key, value]); + + return ok([ + "/bin/bash", + "-lc", + script, + "srchd", + ...envArgs, + cwd, + cmd, + ]); +} + export function computerId( experiment: ExperimentResource, agent: AgentResource, @@ -217,41 +258,30 @@ export class Computer { const env = options?.env ?? {}; const process = newProcess(cmd, cwd, env, options?.tty); - // Build the command with environment variables and working directory - // We wrap the command to capture the PID and run it in the foreground - // This keeps the K8s exec connection alive for the duration of the command - let fullCmd = ""; - if (options?.env) { - const envVars = Object.entries(env) - .map(([k, v]) => `export ${k}="${v.replace(/"/g, '\\"')}"`) - .join("; "); - fullCmd += envVars + "; "; + const shellCommand = buildShellCommand(cmd, cwd, env, true); + if (shellCommand.isErr()) { + return shellCommand; } - // Run the command in foreground to keep the K8s exec connection alive. - // We use a wrapper script that: - // 1. Changes to the specified directory - // 2. Prints the PID to stderr for tracking - // 3. Runs the command directly (not in background) - // This ensures stdin/stdout/stderr remain connected for the duration of the process. - const escapedCmd = cmd.replace(/'/g, "'\\''"); - const escapedCwd = cwd.replace(/'/g, "'\\''"); - fullCmd += `cd '\''${escapedCwd}'\'' && echo "SRCHD_PID:$$" >&2 && ${escapedCmd}`; - const res = await k8sSpawn( - ["/bin/bash", "-lc", fullCmd], + shellCommand.value, this.namespace, this.computerId, process, options?.tty ? undefined : options?.timeoutMs, ); - // Extract PID from captured output - // The PID marker is written to stderr (or stdout in TTY mode) - const outputToSearch = options?.tty ? process.stdout : process.stderr; - const pidMatch = outputToSearch.match(/SRCHD_PID:(\d+)/); - if (pidMatch && pidMatch[1]) { - process.pid = parseInt(pidMatch[1], 10); + // Extract PID from captured output. + // We prefer the PID detected while streaming so it remains available even if + // the output buffer has already truncated older data. + if (process.detectedPid !== undefined) { + process.pid = process.detectedPid; + } else { + const outputToSearch = options?.tty ? process.stdout : process.stderr; + const pidMatch = outputToSearch.match(/SRCHD_PID:(\d+)/); + if (pidMatch?.[1]) { + process.pid = parseInt(pidMatch[1], 10); + } } this.processes.set(process.pid, process); @@ -304,18 +334,13 @@ export class Computer { const startTs = Date.now(); - // Build the command with environment variables and working directory - let fullCmd = ""; - if (options?.env) { - const envVars = Object.entries(options.env) - .map(([k, v]) => `export ${k}="${v.replace(/"/g, '\\"')}"`) - .join("; "); - fullCmd += envVars + "; "; + const shellCommand = buildShellCommand(cmd, cwd, options?.env ?? {}, false); + if (shellCommand.isErr()) { + return shellCommand; } - fullCmd += `cd "${cwd.replace(/"/g, '\\"')}" && ${cmd}`; const execPromise = computerExec( - ["/bin/bash", "-lc", fullCmd], + shellCommand.value, this.namespace, this.computerId, options?.timeoutMs, diff --git a/src/computer/process.ts b/src/computer/process.ts index a624b711..3dab6b90 100644 --- a/src/computer/process.ts +++ b/src/computer/process.ts @@ -1,6 +1,8 @@ -import { PassThrough } from "stream"; +import { PassThrough, Writable } from "stream"; import { Terminal } from "@xterm/xterm"; +const MAX_OUTPUT_CHARS = 1024 * 1024; // Keep only the last 1MB per stream. + export type ProcessStatus = | 'running' | 'terminated'; @@ -9,11 +11,12 @@ export type Process = { pid: number; status: ProcessStatus; stdinStream: PassThrough; - stdoutStream: PassThrough; - stderrStream: PassThrough; + stdoutStream: Writable; + stderrStream: Writable; output: { stdout: string; stderr: string }; get stdout(): string; get stderr(): string; + detectedPid?: number; exitCode?: number; createdAt: Date; command: string; @@ -25,6 +28,40 @@ export type Process = { getTerminalBuffer(): string; } +function appendToBoundedBuffer( + current: string, + incoming: string, + maxChars: number, +): { value: string; truncatedChars: number } { + if (incoming.length >= maxChars) { + return { + value: incoming.slice(-maxChars), + truncatedChars: current.length + incoming.length - maxChars, + }; + } + + const overflow = current.length + incoming.length - maxChars; + if (overflow <= 0) { + return { + value: current + incoming, + truncatedChars: 0, + }; + } + + return { + value: current.slice(overflow) + incoming, + truncatedChars: overflow, + }; +} + +function formatBufferedOutput(value: string, truncatedChars: number): string { + if (truncatedChars <= 0) { + return value; + } + + return `[srchd truncated ${truncatedChars} chars of earlier output; showing the most recent ${value.length} chars]\n${value}`; +} + export function newProcess( command: string, cwd: string, @@ -32,10 +69,12 @@ export function newProcess( tty: boolean = false, ): Process { - const stdoutStream = new PassThrough(); - const stderrStream = new PassThrough(); - // Since JS strings are passed by value, we need to use getters to capture the output - const output = { stdout: "", stderr: "" }; + const output = { + stdout: "", + stderr: "", + stdoutTruncatedChars: 0, + stderrTruncatedChars: 0, + }; // Create terminal instance for TTY mode let terminal: Terminal | undefined; @@ -50,45 +89,79 @@ export function newProcess( }); } - const originalStdoutWrite = stdoutStream.write.bind(stdoutStream); - stdoutStream.write = (chunk: any, ...args: any[]) => { - if (chunk && chunk !== null) { - const text = chunk.toString(); - output.stdout += text; + let process: Process | undefined; - // Write to terminal if in TTY mode - if (terminal) { - terminal.write(text); - } + const detectPid = (text: string) => { + const pidMatch = text.match(/SRCHD_PID:(\d+)/); + if (!pidMatch?.[1] || !process) { + return; } - return originalStdoutWrite(chunk, ...args); + + const pid = parseInt(pidMatch[1], 10); + process.detectedPid = pid; + process.pid = pid; }; - const originalStderrWrite = stderrStream.write.bind(stderrStream); - stderrStream.write = (chunk: any, ...args: any[]) => { - if (chunk && chunk !== null) { - const text = chunk.toString(); - output.stderr += text; + const stdoutStream = new Writable({ + write(chunk: any, _encoding, callback) { + try { + if (chunk !== undefined && chunk !== null) { + const text = chunk.toString(); + const next = appendToBoundedBuffer(output.stdout, text, MAX_OUTPUT_CHARS); + output.stdout = next.value; + output.stdoutTruncatedChars += next.truncatedChars; + detectPid(text); - // In TTY mode, stderr also goes to terminal (like real TTY behavior) - if (terminal) { - terminal.write(text); + // Write to terminal if in TTY mode + if (terminal) { + terminal.write(text); + } + } + callback(); + } catch (e) { + callback(e as Error); } - } - return originalStderrWrite(chunk, ...args); - }; + }, + }); + + const stderrStream = new Writable({ + write(chunk: any, _encoding, callback) { + try { + if (chunk !== undefined && chunk !== null) { + const text = chunk.toString(); + const next = appendToBoundedBuffer(output.stderr, text, MAX_OUTPUT_CHARS); + output.stderr = next.value; + output.stderrTruncatedChars += next.truncatedChars; + detectPid(text); + + // In TTY mode, stderr also goes to terminal (like real TTY behavior) + if (terminal) { + terminal.write(text); + } + } + callback(); + } catch (e) { + callback(e as Error); + } + }, + }); const stdinStream = new PassThrough(); - return { + process = { pid: -1, // Nonexistent when starting process status: 'running', stdinStream, stdoutStream, stderrStream, output, - get stdout() { return output.stdout; }, - get stderr() { return output.stderr; }, + get stdout() { + return formatBufferedOutput(output.stdout, output.stdoutTruncatedChars); + }, + get stderr() { + return formatBufferedOutput(output.stderr, output.stderrTruncatedChars); + }, + detectedPid: undefined, exitCode: undefined, createdAt: new Date(), command, @@ -98,7 +171,7 @@ export function newProcess( terminal, getTerminalBuffer(): string { if (!terminal) { - return output.stdout; + return formatBufferedOutput(output.stdout, output.stdoutTruncatedChars); } // Serialize the terminal buffer @@ -115,4 +188,6 @@ export function newProcess( return lines.join('\n'); }, }; + + return process; } diff --git a/src/models/moonshotai.ts b/src/models/moonshotai.ts index 70a9ee8f..f327a3ef 100644 --- a/src/models/moonshotai.ts +++ b/src/models/moonshotai.ts @@ -48,6 +48,10 @@ const TOKEN_PRICING: Record = { "kimi-k2.6": normalizeTokenPrices(0.95, 4.0, 0.16), }; +function stripNullBytes(value: string): string { + return value.replace(/\u0000/g, ""); +} + export class MoonshotAILLM extends LLM { private client: OpenAI; private model: MoonshotAIModel; @@ -66,7 +70,7 @@ export class MoonshotAILLM extends LLM { messages(prompt: string, messages: Message[]) { const inputItems: ChatCompletionMessageParam[] = [ - { role: "system", content: prompt }, + { role: "system", content: stripNullBytes(prompt) }, ...removeNulls( messages .map((msg) => { @@ -75,14 +79,17 @@ export class MoonshotAILLM extends LLM { return msg.content.map((c) => { switch (c.type) { case "text": - return { role: "user" as const, content: c.text }; + return { + role: "user" as const, + content: stripNullBytes(c.text), + }; case "tool_result": return { role: "tool" as const, - name: c.toolUseName, - tool_call_id: c.toolUseId, - id: c.toolUseId, - content: JSON.stringify(c.content), + name: stripNullBytes(c.toolUseName), + tool_call_id: stripNullBytes(c.toolUseId), + id: stripNullBytes(c.toolUseId), + content: stripNullBytes(JSON.stringify(c.content)), }; default: return undefined; @@ -98,19 +105,19 @@ export class MoonshotAILLM extends LLM { msg.content.forEach((c) => { switch (c.type) { case "text": - message.content = c.text; + message.content = stripNullBytes(c.text); break; case "thinking": - message.reasoning_content = c.thinking; + message.reasoning_content = stripNullBytes(c.thinking); break; case "tool_use": message.tool_calls = message.tool_calls ?? []; message.tool_calls.push({ type: "function" as const, - id: c.id, + id: stripNullBytes(c.id), function: { - name: c.name, - arguments: JSON.stringify(c.input), + name: stripNullBytes(c.name), + arguments: stripNullBytes(JSON.stringify(c.input)), }, }); break; @@ -143,8 +150,8 @@ export class MoonshotAILLM extends LLM { tools: tools.map((tool) => ({ type: "function", function: { - name: tool.name, - description: tool.description, + name: stripNullBytes(tool.name), + description: tool.description ? stripNullBytes(tool.description) : undefined, parameters: tool.inputSchema as any, }, strict: false, @@ -152,10 +159,12 @@ export class MoonshotAILLM extends LLM { }); const message = response.choices[0].message; - const textContent = message.content; + const textContent = message.content + ? stripNullBytes(message.content) + : message.content; const thinkingContent = - "reasoning_content" in message - ? (message.reasoning_content as string) + "reasoning_content" in message && typeof message.reasoning_content === "string" + ? stripNullBytes(message.reasoning_content) : undefined; const toolCalls = message.tool_calls; @@ -184,12 +193,12 @@ export class MoonshotAILLM extends LLM { .map((toolCall) => { return { type: "tool_use" as const, - id: toolCall.id, - name: toolCall.function.name, - input: JSON.parse(toolCall.function.arguments), + id: stripNullBytes(toolCall.id), + name: stripNullBytes(toolCall.function.name), + input: JSON.parse(stripNullBytes(toolCall.function.arguments)), provider: { moonshotai: { - id: toolCall.id, + id: stripNullBytes(toolCall.id), }, }, }; @@ -259,8 +268,8 @@ export class MoonshotAILLM extends LLM { tools: tools.map((tool) => ({ type: "function", function: { - name: tool.name, - description: tool.description, + name: stripNullBytes(tool.name), + description: tool.description ? stripNullBytes(tool.description) : undefined, parameters: tool.inputSchema as any, }, strict: false, diff --git a/src/tools/computer_process.ts b/src/tools/computer_process.ts index c08f919a..a801f584 100644 --- a/src/tools/computer_process.ts +++ b/src/tools/computer_process.ts @@ -335,7 +335,7 @@ Success status, last 100 lines of stdout after sending input, stderr, and exit c Displays the tail of a process's stdout (for long outputs), stderr, current status, and exit code if available. **Background processes:** -When a process is moved to background (either explicitly or via timeout), this command shows any output received up until that point. The output buffer continues to accumulate while the process runs in the background. +When a process is moved to background (either explicitly or via timeout), this command shows the recent output received up until that point. The output buffer is bounded, so older output may be truncated while the process runs in the background. **Use cases:** - Check progress of long-running processes