mirror of
https://github.com/Twigpine/zero.git
synced 2026-11-08 09:49:34 -10:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a766163401 | ||
|
|
16ce70b7a9 |
No files matched your search
+134
-54
@@ -28,8 +28,17 @@ const (
|
||||
completedSessionRetention = 30 * time.Second
|
||||
maxExecSessions = 64
|
||||
recentExecOutputBytes = 4096
|
||||
execSessionStopTimeout = 3 * time.Second
|
||||
execSessionEvictedMessage = "[zero] session evicted: too many background terminals\n"
|
||||
// maxExecOutputBufferBytes caps the undrained output an unpolled session can
|
||||
// accumulate. Without a cap, a long-lived background session nobody polls
|
||||
// again (e.g. a dev server left running after its initiating run was
|
||||
// cancelled) grows this buffer forever as long as the process keeps writing
|
||||
// output, with no ceiling — this previously ran a session's memory into the
|
||||
// tens of gigabytes over several hours and got the whole zero process
|
||||
// OOM-killed by the OS.
|
||||
maxExecOutputBufferBytes = 2 * 1024 * 1024
|
||||
execSessionStopTimeout = 3 * time.Second
|
||||
execSessionEvictedMessage = "[zero] session evicted: too many background terminals\n"
|
||||
execOutputBufferTruncatedMessage = "[zero] output buffer truncated: undrained output exceeded 2MiB, oldest output dropped"
|
||||
)
|
||||
|
||||
type execSessionManager struct {
|
||||
@@ -233,6 +242,10 @@ type ExecSessionSnapshot struct {
|
||||
Status string
|
||||
ExitCode *int
|
||||
RecentOutput string
|
||||
// OutputTruncated reflects the buffer's last-known truncation state (see
|
||||
// execOutputBuffer.peekTruncated), so a session listing (/ps) still shows
|
||||
// this even if the session is reaped before ever being polled again.
|
||||
OutputTruncated bool
|
||||
}
|
||||
|
||||
type ExecSessionController interface {
|
||||
@@ -317,24 +330,26 @@ func (session *execSession) snapshot() ExecSessionSnapshot {
|
||||
status = "exited"
|
||||
}
|
||||
return ExecSessionSnapshot{
|
||||
ID: session.id,
|
||||
Command: session.commandText,
|
||||
Cwd: session.cwd,
|
||||
RelativeCwd: session.relativeCwd,
|
||||
StartedAt: startedAt,
|
||||
LastUsedAt: lastUsedAt,
|
||||
TTY: session.tty,
|
||||
Status: status,
|
||||
ExitCode: copiedExit,
|
||||
RecentOutput: session.output.recentString(),
|
||||
ID: session.id,
|
||||
Command: session.commandText,
|
||||
Cwd: session.cwd,
|
||||
RelativeCwd: session.relativeCwd,
|
||||
StartedAt: startedAt,
|
||||
LastUsedAt: lastUsedAt,
|
||||
TTY: session.tty,
|
||||
Status: status,
|
||||
ExitCode: copiedExit,
|
||||
RecentOutput: session.output.recentString(),
|
||||
OutputTruncated: session.output.peekTruncated(),
|
||||
}
|
||||
}
|
||||
|
||||
type execOutputBuffer struct {
|
||||
mu sync.Mutex
|
||||
data []byte
|
||||
recent []byte
|
||||
notify chan struct{}
|
||||
mu sync.Mutex
|
||||
data []byte
|
||||
recent []byte
|
||||
truncated bool
|
||||
notify chan struct{}
|
||||
}
|
||||
|
||||
func newExecOutputBuffer() *execOutputBuffer {
|
||||
@@ -344,6 +359,10 @@ func newExecOutputBuffer() *execOutputBuffer {
|
||||
func (buffer *execOutputBuffer) Write(p []byte) (int, error) {
|
||||
buffer.mu.Lock()
|
||||
buffer.data = append(buffer.data, p...)
|
||||
if len(buffer.data) > maxExecOutputBufferBytes {
|
||||
buffer.data = buffer.data[len(buffer.data)-maxExecOutputBufferBytes:]
|
||||
buffer.truncated = true
|
||||
}
|
||||
buffer.recent = append(buffer.recent, p...)
|
||||
if len(buffer.recent) > recentExecOutputBytes {
|
||||
buffer.recent = buffer.recent[len(buffer.recent)-recentExecOutputBytes:]
|
||||
@@ -373,6 +392,34 @@ func (buffer *execOutputBuffer) drainString() string {
|
||||
return out
|
||||
}
|
||||
|
||||
// consumeTruncated reports whether the buffer has dropped output to stay
|
||||
// within maxExecOutputBufferBytes since the last call, resetting the flag.
|
||||
// Kept as an out-of-band signal rather than text embedded in drainString's
|
||||
// result: an earlier version prefixed a notice directly into the drained
|
||||
// string, but that notice always sat ~maxExecOutputBufferBytes (2MiB) before
|
||||
// the end of the combined output, far outside the byte-budget head/tail
|
||||
// window truncateExecOutput keeps afterward (at most 400KB even at the
|
||||
// tool's maximum max_output_tokens) — so the notice was reliably swallowed
|
||||
// or chopped by that second, unrelated truncation pass. Surfacing it as a
|
||||
// separate bool lets the caller report it regardless of where in the byte
|
||||
// stream the overflow happened.
|
||||
func (buffer *execOutputBuffer) consumeTruncated() bool {
|
||||
buffer.mu.Lock()
|
||||
defer buffer.mu.Unlock()
|
||||
truncated := buffer.truncated
|
||||
buffer.truncated = false
|
||||
return truncated
|
||||
}
|
||||
|
||||
// peekTruncated reports the buffer's last-known truncation state without
|
||||
// clearing it, for status views (e.g. /ps) that shouldn't consume the signal
|
||||
// a subsequent poll is still meant to report via consumeTruncated.
|
||||
func (buffer *execOutputBuffer) peekTruncated() bool {
|
||||
buffer.mu.Lock()
|
||||
defer buffer.mu.Unlock()
|
||||
return buffer.truncated
|
||||
}
|
||||
|
||||
type execCommandTool struct {
|
||||
baseTool
|
||||
workspaceRoot string
|
||||
@@ -485,25 +532,28 @@ func (tool execCommandTool) run(ctx context.Context, args map[string]any, engine
|
||||
if err != nil {
|
||||
return errorResult("Error starting exec_command: " + err.Error())
|
||||
}
|
||||
output := session.collect(ctx, time.Duration(yieldTimeMS)*time.Millisecond)
|
||||
output, outputTruncated := session.collect(ctx, time.Duration(yieldTimeMS)*time.Millisecond)
|
||||
if ctx != nil && ctx.Err() != nil && !session.doneClosed() {
|
||||
session.terminate()
|
||||
output += session.collect(context.Background(), time.Second)
|
||||
more, moreTruncated := session.collect(context.Background(), time.Second)
|
||||
output += more
|
||||
outputTruncated = outputTruncated || moreTruncated
|
||||
}
|
||||
exitCode, exited := session.exitStatus()
|
||||
if exited {
|
||||
tool.manager.remove(session.id)
|
||||
}
|
||||
return execToolResult(execToolResultInput{
|
||||
commandText: commandText,
|
||||
output: output,
|
||||
sessionID: session.id,
|
||||
exitCode: exitCode,
|
||||
exited: exited,
|
||||
relativeCwd: relativeCwd,
|
||||
tty: session.tty,
|
||||
plan: session.plan,
|
||||
maxOutputTokens: maxOutputTokens,
|
||||
commandText: commandText,
|
||||
output: output,
|
||||
outputBufferTruncated: outputTruncated,
|
||||
sessionID: session.id,
|
||||
exitCode: exitCode,
|
||||
exited: exited,
|
||||
relativeCwd: relativeCwd,
|
||||
tty: session.tty,
|
||||
plan: session.plan,
|
||||
maxOutputTokens: maxOutputTokens,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -558,26 +608,39 @@ func (tool execCommandTool) startSession(commandText string, absoluteCwd string,
|
||||
return session, nil
|
||||
}
|
||||
|
||||
func (session *execSession) collect(ctx context.Context, wait time.Duration) string {
|
||||
// collect drains the session's output buffer, returning the accumulated text
|
||||
// and whether the buffer ever had to drop output to stay within
|
||||
// maxExecOutputBufferBytes since the last collect call.
|
||||
func (session *execSession) collect(ctx context.Context, wait time.Duration) (string, bool) {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
deadline := time.Now().Add(wait)
|
||||
var builder strings.Builder
|
||||
truncated := false
|
||||
finish := func() (string, bool) {
|
||||
if session.output.consumeTruncated() {
|
||||
truncated = true
|
||||
}
|
||||
return builder.String(), truncated
|
||||
}
|
||||
for {
|
||||
if chunk := session.output.drainString(); chunk != "" {
|
||||
builder.WriteString(chunk)
|
||||
if session.output.consumeTruncated() {
|
||||
truncated = true
|
||||
}
|
||||
continue
|
||||
}
|
||||
if session.doneClosed() {
|
||||
if chunk := session.output.drainString(); chunk != "" {
|
||||
builder.WriteString(chunk)
|
||||
}
|
||||
return builder.String()
|
||||
return finish()
|
||||
}
|
||||
remaining := time.Until(deadline)
|
||||
if remaining <= 0 {
|
||||
return builder.String()
|
||||
return finish()
|
||||
}
|
||||
timer := time.NewTimer(remaining)
|
||||
select {
|
||||
@@ -602,9 +665,9 @@ func (session *execSession) collect(ctx context.Context, wait time.Duration) str
|
||||
default:
|
||||
}
|
||||
}
|
||||
return builder.String()
|
||||
return finish()
|
||||
case <-timer.C:
|
||||
return builder.String()
|
||||
return finish()
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -701,22 +764,23 @@ func (tool writeStdinTool) RunWithOptions(ctx context.Context, args map[string]a
|
||||
}
|
||||
}
|
||||
}
|
||||
output := session.collect(ctx, time.Duration(yieldTimeMS)*time.Millisecond)
|
||||
output, outputTruncated := session.collect(ctx, time.Duration(yieldTimeMS)*time.Millisecond)
|
||||
exitCode, exited := session.exitStatus()
|
||||
if exited {
|
||||
tool.manager.remove(session.id)
|
||||
}
|
||||
return execToolResult(execToolResultInput{
|
||||
commandText: session.commandText,
|
||||
output: output,
|
||||
sessionID: session.id,
|
||||
exitCode: exitCode,
|
||||
exited: exited,
|
||||
relativeCwd: session.relativeCwd,
|
||||
tty: session.tty,
|
||||
interrupted: interrupted,
|
||||
plan: session.plan,
|
||||
maxOutputTokens: maxOutputTokens,
|
||||
commandText: session.commandText,
|
||||
output: output,
|
||||
outputBufferTruncated: outputTruncated,
|
||||
sessionID: session.id,
|
||||
exitCode: exitCode,
|
||||
exited: exited,
|
||||
relativeCwd: session.relativeCwd,
|
||||
tty: session.tty,
|
||||
interrupted: interrupted,
|
||||
plan: session.plan,
|
||||
maxOutputTokens: maxOutputTokens,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -741,16 +805,22 @@ func shouldInterruptExecSession(chars string, tty bool) bool {
|
||||
}
|
||||
|
||||
type execToolResultInput struct {
|
||||
commandText string
|
||||
output string
|
||||
sessionID int
|
||||
exitCode int
|
||||
exited bool
|
||||
relativeCwd string
|
||||
tty bool
|
||||
interrupted bool
|
||||
plan zeroSandbox.CommandPlan
|
||||
maxOutputTokens int
|
||||
commandText string
|
||||
output string
|
||||
// outputBufferTruncated is true when the session's undrained output buffer
|
||||
// had to drop bytes to stay within maxExecOutputBufferBytes since it was
|
||||
// last collected — data that is gone for good, unlike the head/tail
|
||||
// truncation below (which the caller can recover by polling again or
|
||||
// raising max_output_tokens).
|
||||
outputBufferTruncated bool
|
||||
sessionID int
|
||||
exitCode int
|
||||
exited bool
|
||||
relativeCwd string
|
||||
tty bool
|
||||
interrupted bool
|
||||
plan zeroSandbox.CommandPlan
|
||||
maxOutputTokens int
|
||||
}
|
||||
|
||||
func execToolResult(input execToolResultInput) Result {
|
||||
@@ -771,16 +841,26 @@ func execToolResult(input execToolResultInput) Result {
|
||||
} else {
|
||||
meta["session_id"] = strconv.Itoa(input.sessionID)
|
||||
}
|
||||
if input.outputBufferTruncated {
|
||||
meta["output_buffer_truncated"] = "true"
|
||||
}
|
||||
|
||||
status := StatusOK
|
||||
if input.exited && ((input.exitCode != 0 && !input.interrupted) || meta[SandboxLikelyDeniedMeta] == "true") {
|
||||
status = StatusError
|
||||
}
|
||||
body := formatExecCommandOutput(output, input.sessionID, input.exited, input.exitCode, input.interrupted)
|
||||
if input.outputBufferTruncated {
|
||||
// Appended after truncateExecOutput's own head/tail slicing, not
|
||||
// embedded in the text that goes through it — a marker inside that
|
||||
// text can land in the discarded middle or get chopped at a small
|
||||
// max_output_tokens budget. Appending here guarantees it survives.
|
||||
body += "\n" + execOutputBufferTruncatedMessage
|
||||
}
|
||||
return Result{
|
||||
Status: status,
|
||||
Output: body,
|
||||
Truncated: truncated,
|
||||
Truncated: truncated || input.outputBufferTruncated,
|
||||
Meta: meta,
|
||||
Display: Display{
|
||||
Summary: execDisplaySummary(input.commandText, input.sessionID, input.exited, input.exitCode),
|
||||
|
||||
@@ -363,6 +363,93 @@ func TestStartExecProcessFallsBackAfterPTYStartMutation(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestExecOutputBufferCapsUndrainedData: a session nobody polls must not grow
|
||||
// its undrained buffer without bound — a long-lived background process that
|
||||
// keeps writing while unpolled previously ran a session's memory into the
|
||||
// tens of gigabytes and got the whole zero process OOM-killed by the OS.
|
||||
func TestExecOutputBufferCapsUndrainedData(t *testing.T) {
|
||||
buffer := newExecOutputBuffer()
|
||||
|
||||
chunk := strings.Repeat("x", 1024)
|
||||
writes := (maxExecOutputBufferBytes / len(chunk)) + 10 // comfortably over the cap
|
||||
for i := 0; i < writes; i++ {
|
||||
if _, err := buffer.Write([]byte(chunk)); err != nil {
|
||||
t.Fatalf("Write: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
buffer.mu.Lock()
|
||||
dataLen := len(buffer.data)
|
||||
buffer.mu.Unlock()
|
||||
if dataLen > maxExecOutputBufferBytes {
|
||||
t.Fatalf("undrained buffer grew to %d bytes, want <= %d", dataLen, maxExecOutputBufferBytes)
|
||||
}
|
||||
|
||||
out := buffer.drainString()
|
||||
// drainString itself carries no truncation marker: a marker embedded in
|
||||
// the drained text would always sit ~maxExecOutputBufferBytes before the
|
||||
// end of the string, past any realistic head/tail truncation window a
|
||||
// caller applies afterward (see execToolResult), so it's reliably lost.
|
||||
// The signal lives out of band on consumeTruncated/peekTruncated instead.
|
||||
if !strings.HasSuffix(out, chunk) {
|
||||
t.Fatal("drained output should keep the most recent bytes, not the oldest")
|
||||
}
|
||||
if len(out) > maxExecOutputBufferBytes {
|
||||
t.Fatalf("drained output = %d bytes, want <= %d (no marker text mixed in)", len(out), maxExecOutputBufferBytes)
|
||||
}
|
||||
|
||||
if !buffer.peekTruncated() {
|
||||
t.Fatal("peekTruncated should report the overflow without clearing it")
|
||||
}
|
||||
if !buffer.peekTruncated() {
|
||||
t.Fatal("a second peekTruncated call should still report it — peek must not consume")
|
||||
}
|
||||
if !buffer.consumeTruncated() {
|
||||
t.Fatal("consumeTruncated should report the overflow")
|
||||
}
|
||||
if buffer.consumeTruncated() {
|
||||
t.Fatal("consumeTruncated should reset after being read once")
|
||||
}
|
||||
if buffer.peekTruncated() {
|
||||
t.Fatal("peekTruncated should reflect the reset state after consumeTruncated")
|
||||
}
|
||||
|
||||
// A second drain with no intervening writes must be empty.
|
||||
if got := buffer.drainString(); got != "" {
|
||||
t.Fatalf("drainString after a full drain = %q, want empty", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestExecToolResultSurfacesBufferTruncationOutsideByteBudget: an earlier
|
||||
// version embedded the truncation notice directly in the drained text, which
|
||||
// a review found sits ~maxExecOutputBufferBytes before the end of the
|
||||
// string — far past the head/tail window truncateExecOutput keeps even at
|
||||
// the tool's smallest allowed max_output_tokens, so the notice was reliably
|
||||
// swallowed or chopped. The notice must survive regardless of how small the
|
||||
// byte budget is, and it must not count against that budget.
|
||||
func TestExecToolResultSurfacesBufferTruncationOutsideByteBudget(t *testing.T) {
|
||||
hugeOutput := strings.Repeat("y", maxExecOutputBufferBytes+1024)
|
||||
|
||||
result := execToolResult(execToolResultInput{
|
||||
commandText: "cmd",
|
||||
output: hugeOutput,
|
||||
outputBufferTruncated: true,
|
||||
sessionID: 1,
|
||||
exited: false,
|
||||
maxOutputTokens: 1, // the schema's own declared minimum
|
||||
})
|
||||
|
||||
if !strings.Contains(result.Output, execOutputBufferTruncatedMessage) {
|
||||
t.Fatalf("result output should contain the buffer-truncation notice even at the smallest max_output_tokens, got %q", result.Output[:min(len(result.Output), 200)])
|
||||
}
|
||||
if result.Meta["output_buffer_truncated"] != "true" {
|
||||
t.Fatalf("result meta should flag output_buffer_truncated, got %#v", result.Meta)
|
||||
}
|
||||
if !result.Truncated {
|
||||
t.Fatal("result should report Truncated when the buffer dropped output")
|
||||
}
|
||||
}
|
||||
|
||||
// resilientTempDir is like t.TempDir() but tolerates the Windows handle-release
|
||||
// lag: a SIGKILL'd child process that had the dir as its cwd may not have
|
||||
// released it the instant it is reaped, so the immediate RemoveAll t.TempDir()
|
||||
|
||||
@@ -146,6 +146,9 @@ func formatBackgroundTerminalRow(session tools.ExecSessionSnapshot, now time.Tim
|
||||
}
|
||||
preview := compactCommandOutputText(session.RecentOutput)
|
||||
prefix := fmt.Sprintf("%d · %s · %s · %s", session.ID, session.Status, age, cwd)
|
||||
if session.OutputTruncated {
|
||||
prefix += " · output truncated"
|
||||
}
|
||||
if preview != "" {
|
||||
return prefix + " · " + command + " · " + preview
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user