Skip to content

Commit 964be4f

Browse files
codefionncodefionn
committed
fix: Limit sandbox output
1 parent 1eaba21 commit 964be4f

File tree

6 files changed

+238
-4
lines changed

6 files changed

+238
-4
lines changed

‎internal/session/session.go‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,9 @@ type Session struct {
3333
PlanningActive bool // whether planning phase is currently running
3434
PlanningObjective string // objective of current planning phase
3535
PlanningStartTime time.Time // when current planning phase started
36+
LastSandboxExitCode int // exit code from last sandbox execution
37+
LastSandboxStdout string // stdout from last sandbox execution
38+
LastSandboxStderr string // stderr from last sandbox execution
3639
mu sync.RWMutex
3740
CreatedAt time.Time
3841
UpdatedAt time.Time
@@ -358,3 +361,20 @@ func (s *Session) CompactWithSummary(original []*Message, summary string) bool {
358361

359362
return true
360363
}
364+
365+
// SetLastSandboxOutput stores the output from the last sandbox execution
366+
func (s *Session) SetLastSandboxOutput(exitCode int, stdout, stderr string) {
367+
s.mu.Lock()
368+
defer s.mu.Unlock()
369+
s.LastSandboxExitCode = exitCode
370+
s.LastSandboxStdout = stdout
371+
s.LastSandboxStderr = stderr
372+
s.UpdatedAt = time.Now()
373+
}
374+
375+
// GetLastSandboxOutput retrieves the output from the last sandbox execution
376+
func (s *Session) GetLastSandboxOutput() (exitCode int, stdout, stderr string) {
377+
s.mu.RLock()
378+
defer s.mu.RUnlock()
379+
return s.LastSandboxExitCode, s.LastSandboxStdout, s.LastSandboxStderr
380+
}

‎internal/tools/sandbox.go‎

Lines changed: 31 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -167,8 +167,37 @@ func (t *SandboxTool) Description() string {
167167
b.WriteString("Basic standard library packages available (don't use the `os`, `ioutil, `net`, `exec` package instead use methods provided below). Timeout enforced.\n")
168168
b.WriteString("Every program **must** declare `package main`, define `func main()`, and print results (e.g., via `fmt.Println`) so the orchestrator receives the output.\n\n")
169169
b.WriteString("Try to reduce the output of shell programs by e.g. only searching and outputting errors.\n\n")
170-
b.WriteString("Don't use this tool calls for just outputing the summary text at the end.")
171-
b.WriteString("Seven custom functions are automatically available in your code:\n\n")
170+
b.WriteString("Output is limited to 4096 lines. When truncated, consider parsing specific parts with Go (e.g., only output lines around error messages).\n\n")
171+
b.WriteString("Don't use this tool calls for just outputing the summary text at the end.\n\n")
172+
b.WriteString("# Previous Execution State\n\n")
173+
b.WriteString("Three global variables are automatically available that preserve output from the previous sandbox execution:\n")
174+
b.WriteString("- `last_exit_code` (int): Exit code from the last sandbox run (0 on first run)\n")
175+
b.WriteString("- `last_stdout` (string): Standard output from the last sandbox run (empty on first run)\n")
176+
b.WriteString("- `last_stderr` (string): Standard error from the last sandbox run (empty on first run)\n\n")
177+
b.WriteString("These variables are useful when output is truncated and you want to process specific parts in subsequent runs.\n")
178+
b.WriteString("Example:\n")
179+
b.WriteString("```go\n")
180+
b.WriteString("package main\n\n")
181+
b.WriteString("import (\n")
182+
b.WriteString(" \"fmt\"\n")
183+
b.WriteString(" \"strings\"\n")
184+
b.WriteString(")\n\n")
185+
b.WriteString("func main() {\n")
186+
b.WriteString(" if last_exit_code != 0 {\n")
187+
b.WriteString(" fmt.Printf(\"Previous run failed with exit code %d\\n\", last_exit_code)\n")
188+
b.WriteString(" // Parse errors from last_stderr\n")
189+
b.WriteString(" for _, line := range strings.Split(last_stderr, \"\\n\") {\n")
190+
b.WriteString(" if strings.Contains(line, \"error\") {\n")
191+
b.WriteString(" fmt.Println(line)\n")
192+
b.WriteString(" }\n")
193+
b.WriteString(" }\n")
194+
b.WriteString(" } else if len(last_stdout) > 0 {\n")
195+
b.WriteString(" fmt.Printf(\"Processing previous output (%d bytes)\\n\", len(last_stdout))\n")
196+
b.WriteString(" }\n")
197+
b.WriteString("}\n")
198+
b.WriteString("```\n\n")
199+
b.WriteString("# Available Host Functions\n\n")
200+
b.WriteString("Custom functions are automatically available in your code:\n\n")
172201

173202
b.WriteString("1. Fetch(method, url, body string) (responseBody string, statusCode int)\n")
174203
b.WriteString(" - Make HTTP requests (GET, POST, PUT, DELETE, etc.)\n")

‎internal/tools/sandbox_host_functions.go‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -738,3 +738,70 @@ func (t *SandboxTool) registerRemoveDirHostFunction(envBuilder HostModuleBuilder
738738
}).
739739
Export("remove_dir")
740740
}
741+
742+
// registerGetLastExitCodeHostFunction registers the get_last_exit_code host function
743+
func (t *SandboxTool) registerGetLastExitCodeHostFunction(envBuilder HostModuleBuilder) {
744+
// get_last_exit_code() -> exit_code
745+
envBuilder.NewFunctionBuilder().
746+
WithFunc(func(ctx context.Context, m api.Module) int32 {
747+
if t.session == nil {
748+
return 0 // No session, return 0
749+
}
750+
751+
exitCode, _, _ := t.session.GetLastSandboxOutput()
752+
return int32(exitCode)
753+
}).
754+
Export("get_last_exit_code")
755+
}
756+
757+
// registerGetLastStdoutHostFunction registers the get_last_stdout host function
758+
func (t *SandboxTool) registerGetLastStdoutHostFunction(envBuilder HostModuleBuilder) {
759+
// get_last_stdout(buffer_ptr, buffer_cap) -> length
760+
envBuilder.NewFunctionBuilder().
761+
WithFunc(func(ctx context.Context, m api.Module, bufferPtr, bufferCap uint32) int32 {
762+
if t.session == nil {
763+
return 0 // No session, return empty
764+
}
765+
766+
_, stdout, _ := t.session.GetLastSandboxOutput()
767+
stdoutBytes := []byte(stdout)
768+
769+
memory := m.Memory()
770+
if uint32(len(stdoutBytes)) > bufferCap {
771+
stdoutBytes = stdoutBytes[:bufferCap]
772+
}
773+
774+
if bufferCap > 0 && len(stdoutBytes) > 0 {
775+
memory.Write(bufferPtr, stdoutBytes)
776+
}
777+
778+
return int32(len(stdoutBytes))
779+
}).
780+
Export("get_last_stdout")
781+
}
782+
783+
// registerGetLastStderrHostFunction registers the get_last_stderr host function
784+
func (t *SandboxTool) registerGetLastStderrHostFunction(envBuilder HostModuleBuilder) {
785+
// get_last_stderr(buffer_ptr, buffer_cap) -> length
786+
envBuilder.NewFunctionBuilder().
787+
WithFunc(func(ctx context.Context, m api.Module, bufferPtr, bufferCap uint32) int32 {
788+
if t.session == nil {
789+
return 0 // No session, return empty
790+
}
791+
792+
_, _, stderr := t.session.GetLastSandboxOutput()
793+
stderrBytes := []byte(stderr)
794+
795+
memory := m.Memory()
796+
if uint32(len(stderrBytes)) > bufferCap {
797+
stderrBytes = stderrBytes[:bufferCap]
798+
}
799+
800+
if bufferCap > 0 && len(stderrBytes) > 0 {
801+
memory.Write(bufferPtr, stderrBytes)
802+
}
803+
804+
return int32(len(stderrBytes))
805+
}).
806+
Export("get_last_stderr")
807+
}

‎internal/tools/sandbox_utils.go‎

Lines changed: 62 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -125,11 +125,22 @@ func formatSandboxUIResult(result map[string]interface{}) string {
125125

126126
var sections []string
127127

128+
// Apply 4096 line limit to stdout and stderr
129+
const maxLines = 4096
130+
128131
if stdout != "" {
129-
sections = append(sections, fmt.Sprintf("stdout:\n%s", stdout))
132+
truncatedStdout, stdoutTruncMsg := truncateToLines(stdout, maxLines)
133+
sections = append(sections, fmt.Sprintf("stdout:\n%s", truncatedStdout))
134+
if stdoutTruncMsg != "" {
135+
sections = append(sections, stdoutTruncMsg)
136+
}
130137
}
131138
if stderr != "" {
132-
sections = append(sections, fmt.Sprintf("stderr:\n%s", stderr))
139+
truncatedStderr, stderrTruncMsg := truncateToLines(stderr, maxLines)
140+
sections = append(sections, fmt.Sprintf("stderr:\n%s", truncatedStderr))
141+
if stderrTruncMsg != "" {
142+
sections = append(sections, stderrTruncMsg)
143+
}
133144
}
134145
if errorMsg != "" {
135146
sections = append(sections, fmt.Sprintf("error: %s", errorMsg))
@@ -182,3 +193,52 @@ func CalculateOutputStats(output string) (int, int) {
182193

183194
return bytes, lines
184195
}
196+
197+
// truncateToLines limits text to a maximum number of lines.
198+
// If truncated, returns the truncated text and a truncation message.
199+
func truncateToLines(text string, maxLines int) (string, string) {
200+
if text == "" {
201+
return "", ""
202+
}
203+
204+
// Count lines
205+
lineCount := 0
206+
for _, char := range text {
207+
if char == '\n' {
208+
lineCount++
209+
}
210+
}
211+
212+
// Count the last line if it doesn't end with newline
213+
if len(text) > 0 && text[len(text)-1] != '\n' {
214+
lineCount++
215+
}
216+
217+
// If under the limit, return as-is
218+
if lineCount <= maxLines {
219+
return text, ""
220+
}
221+
222+
// Find the truncation point (end of maxLines-th line)
223+
currentLine := 0
224+
truncIndex := len(text) // default to end of string
225+
for i, char := range text {
226+
if char == '\n' {
227+
currentLine++
228+
if currentLine >= maxLines {
229+
truncIndex = i
230+
break
231+
}
232+
}
233+
}
234+
235+
// If we didn't find enough newlines, truncate at maxLines character count
236+
if currentLine < maxLines {
237+
truncIndex = len(text)
238+
}
239+
240+
truncated := text[:truncIndex]
241+
truncationMsg := fmt.Sprintf("\n\n[Output truncated: showed %d of %d lines. The full output is preserved in last_stdout/last_stderr variables for the next sandbox execution. Consider parsing specific parts with Go to reduce output.]", maxLines, lineCount)
242+
243+
return truncated, truncationMsg
244+
}

‎internal/tools/sandbox_wazero.go‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,9 @@ func (t *SandboxTool) executeWASM(ctx context.Context, wasmBytes []byte, sandbox
7474
t.registerListFilesHostFunction(envBuilder, callTracker)
7575
t.registerRemoveFileHostFunction(envBuilder, callTracker)
7676
t.registerRemoveDirHostFunction(envBuilder, callTracker)
77+
t.registerGetLastExitCodeHostFunction(envBuilder)
78+
t.registerGetLastStdoutHostFunction(envBuilder)
79+
t.registerGetLastStderrHostFunction(envBuilder)
7780

7881
_, err = envBuilder.Instantiate(ctx)
7982
if err != nil {
@@ -211,6 +214,12 @@ func (t *SandboxTool) executeWASM(ctx context.Context, wasmBytes []byte, sandbox
211214
stderr = "Runtime error: " + runtimeErr.Error()
212215
}
213216
metadata := t.buildSandboxMetadata(startTime, commandSummary, timeoutSeconds, finalExitCode, stdout, stderr, false, callTracker)
217+
218+
// Store the output in session even on error
219+
if t.session != nil {
220+
t.session.SetLastSandboxOutput(finalExitCode, stdout, stderr)
221+
}
222+
214223
return attachExecutionMetadata(map[string]interface{}{
215224
"stdout": stdout,
216225
"stderr": stderr,
@@ -221,6 +230,12 @@ func (t *SandboxTool) executeWASM(ctx context.Context, wasmBytes []byte, sandbox
221230
}
222231

223232
metadata := t.buildSandboxMetadata(startTime, commandSummary, timeoutSeconds, finalExitCode, stdout, stderr, false, callTracker)
233+
234+
// Store the output in session for the next sandbox execution
235+
if t.session != nil {
236+
t.session.SetLastSandboxOutput(finalExitCode, stdout, stderr)
237+
}
238+
224239
return attachExecutionMetadata(map[string]interface{}{
225240
"stdout": stdout,
226241
"stderr": stderr,

‎internal/wasi/code_wrapper.go‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,49 @@ func WrapGoCodeWithAuthorization(userCode string) string {
3232
wrapped.WriteString(`// STATCODE_AI_INTERNAL: Authorization system
3333
// This is injected by scriptschnell to enforce network authorization
3434
35+
//go:wasmimport env get_last_exit_code
36+
func getLastExitCodeHost() int32
37+
38+
//go:wasmimport env get_last_stdout
39+
func getLastStdoutHost(bufferPtr *byte, bufferCap int32) int32
40+
41+
//go:wasmimport env get_last_stderr
42+
func getLastStderrHost(bufferPtr *byte, bufferCap int32) int32
43+
44+
// Global variables accessible in user code
45+
var (
46+
last_exit_code int
47+
last_stdout string
48+
last_stderr string
49+
)
50+
51+
func init() {
52+
// Initialize last_* variables from previous sandbox execution
53+
last_exit_code = int(getLastExitCodeHost())
54+
55+
// Get last stdout
56+
stdoutBuffer := make([]byte, 1024*1024) // 1MB buffer
57+
var stdoutPtr *byte
58+
if len(stdoutBuffer) > 0 {
59+
stdoutPtr = &stdoutBuffer[0]
60+
}
61+
stdoutLen := getLastStdoutHost(stdoutPtr, int32(len(stdoutBuffer)))
62+
if stdoutLen > 0 {
63+
last_stdout = string(stdoutBuffer[:stdoutLen])
64+
}
65+
66+
// Get last stderr
67+
stderrBuffer := make([]byte, 1024*1024) // 1MB buffer
68+
var stderrPtr *byte
69+
if len(stderrBuffer) > 0 {
70+
stderrPtr = &stderrBuffer[0]
71+
}
72+
stderrLen := getLastStderrHost(stderrPtr, int32(len(stderrBuffer)))
73+
if stderrLen > 0 {
74+
last_stderr = string(stderrBuffer[:stderrLen])
75+
}
76+
}
77+
3578
//go:wasmimport env authorize_domain
3679
func authorizeDomainHost(domainPtr *byte, domainLen int32) int32
3780

0 commit comments

Comments
 (0)