mirror of
https://github.com/Twigpine/zero.git
synced 2026-11-13 06:51:27 -10:00
* fix(tui): show the context-fill gauge for custom/local Ollama models
The colored context gauge (used/window · N%) only renders when Zero knows
a model's context window. It resolves that two ways: a curated remote
catalog (how Grok/hosted models are known), or live /v1/models discovery —
but parseModelsResponse only ever extracts id/description from that
response; OpenAI-compatible model listings don't carry context-window
metadata at all. A custom/local Ollama model tag (e.g. a locally-pulled
":cloud"-tagged model) has no curated-catalog entry to borrow one from
either, so it always resolved to an unknown window and the gauge silently
fell back to a bare token count with no color grading.
Ollama does expose this, just via a different, Ollama-native endpoint
(/api/show, not /v1/models) whose model_info map reports it under an
architecture-prefixed key like "qwen2.context_length". Added
DiscoverOllamaContextWindow (providermodeldiscovery) to probe it, wired
through an async command scoped to the local "ollama" catalog ID only
(not "ollama-cloud", a different hosted service this endpoint isn't
assumed to exist on) — fired from Init() for the provider active at
launch, and from switchProviderModel for a provider switched to
mid-session (which previously warmed no discovery at all for the new
provider, a separate gap this also fixes).
New tests: DiscoverOllamaContextWindow's endpoint construction and
model_info parsing (including a mixed-type model_info and a missing
*.context_length key), the discovery command's scoping to local Ollama
only, the Update() handler populating modelContextWindow's fallback, and
switchProviderModel warming discovery for the newly active provider.
Local only — not pushed. I can't run a live Ollama daemon in this
sandbox, so the actual /api/show request/response shape (verified against
Ollama's documented API, not a live instance) needs confirming against a
real daemon before this goes out.
* fix(tui): stop showing the token count twice next to the context gauge
contextWindowSegment (the gauge) and usageStatusSegment (the plain "N tok"
figure) both read the exact same value (latestUsageTokens) — the code
already suppressed the plain figure when the sidebar shows it ("the token
figure shows on both sides"), but never extended that same dedup to the
gauge. So as soon as a model's context window becomes known (e.g. after
the previous commit's Ollama fix), the footer showed the identical number
twice: once inside the gauge, once again right next to it.
Fixed by falling back to cost-only whenever the gauge is shown, same as
the sidebar-active branch already does.
* fix: check collect()'s deadline even while output keeps draining
The wait deadline in execSession.collect() was only checked in the branch
reached when drainString() returns empty. A background process producing
output continuously enough to keep the buffer perpetually non-empty could
theoretically starve that check forever, turning the caller's requested
timeout into an unbounded block regardless of yield_time_ms.
Investigated as part of a cross-provider "stuck" report (multiple sessions
across different models/providers all hanging on the same prompt, each
with a background terminal running and no forward progress). Could not
empirically reproduce a hang via a synthetic continuous-writer test under
Go's real scheduler — the reader won the race often enough in practice
that this doesn't appear to be the actual cause of those reports. Keeping
the fix regardless since it closes a real, if hard-to-trigger, gap: the
deadline is now checked in both branches, so continuous output can never
prevent collect() from returning within its requested window.
964 lines
30 KiB
Go
964 lines
30 KiB
Go
package tools
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"os/exec"
|
|
"runtime"
|
|
"sort"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
"unicode/utf8"
|
|
|
|
zeroSandbox "github.com/Gitlawb/zero/internal/sandbox"
|
|
)
|
|
|
|
const (
|
|
ExecCommandToolName = "exec_command"
|
|
WriteStdinToolName = "write_stdin"
|
|
defaultExecYieldTimeMS = 10000
|
|
defaultPollYieldTimeMS = 5000
|
|
maxExecYieldTimeMS = 30000
|
|
maxPollYieldTimeMS = 300000
|
|
defaultMaxOutputTokens = 10000
|
|
maxExecOutputTokenRequest = 200000
|
|
completedSessionRetention = 30 * time.Second
|
|
maxExecSessions = 64
|
|
recentExecOutputBytes = 4096
|
|
// 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 {
|
|
mu sync.Mutex
|
|
nextID int
|
|
sessions map[int]*execSession
|
|
completedRetention time.Duration
|
|
maxSessions int
|
|
}
|
|
|
|
func newExecSessionManager() *execSessionManager {
|
|
return &execSessionManager{
|
|
nextID: 1000,
|
|
sessions: make(map[int]*execSession),
|
|
completedRetention: completedSessionRetention,
|
|
maxSessions: maxExecSessions,
|
|
}
|
|
}
|
|
|
|
var defaultExecSessionManager = newExecSessionManager()
|
|
|
|
func (manager *execSessionManager) allocateID() int {
|
|
manager.mu.Lock()
|
|
defer manager.mu.Unlock()
|
|
id := manager.nextID
|
|
manager.nextID++
|
|
return id
|
|
}
|
|
|
|
func (manager *execSessionManager) store(session *execSession) {
|
|
manager.mu.Lock()
|
|
var pruned *execSession
|
|
removePruned := false
|
|
if manager.maxSessions > 0 && len(manager.sessions) >= manager.maxSessions {
|
|
pruned = manager.sessionToPruneLocked()
|
|
if pruned != nil {
|
|
removePruned = pruned.doneClosed()
|
|
if removePruned {
|
|
delete(manager.sessions, pruned.id)
|
|
}
|
|
}
|
|
}
|
|
manager.sessions[session.id] = session
|
|
manager.mu.Unlock()
|
|
if pruned != nil && !removePruned {
|
|
pruned.output.Write([]byte(execSessionEvictedMessage))
|
|
pruned.terminate()
|
|
}
|
|
}
|
|
|
|
func (manager *execSessionManager) get(id int) (*execSession, bool) {
|
|
manager.mu.Lock()
|
|
defer manager.mu.Unlock()
|
|
session, ok := manager.sessions[id]
|
|
return session, ok
|
|
}
|
|
|
|
func (manager *execSessionManager) remove(id int) {
|
|
manager.mu.Lock()
|
|
defer manager.mu.Unlock()
|
|
delete(manager.sessions, id)
|
|
}
|
|
|
|
func (manager *execSessionManager) removeCompletedLater(session *execSession) {
|
|
retention := manager.completedRetention
|
|
go func() {
|
|
<-session.done
|
|
if retention > 0 {
|
|
timer := time.NewTimer(retention)
|
|
<-timer.C
|
|
}
|
|
manager.remove(session.id)
|
|
}()
|
|
}
|
|
|
|
func (manager *execSessionManager) len() int {
|
|
manager.mu.Lock()
|
|
defer manager.mu.Unlock()
|
|
return len(manager.sessions)
|
|
}
|
|
|
|
func (manager *execSessionManager) sessionToPruneLocked() *execSession {
|
|
if len(manager.sessions) == 0 {
|
|
return nil
|
|
}
|
|
// Snapshot each session's lastUsedAt UNDER session.mu before sorting: touch()
|
|
// writes lastUsedAt under session.mu, so reading it here (under manager.mu only)
|
|
// was a data race on a multi-word time.Time. Lock order stays manager.mu →
|
|
// session.mu (no path takes them the other way). (AUDIT-L15)
|
|
type sessionAge struct {
|
|
session *execSession
|
|
last time.Time
|
|
}
|
|
ages := make([]sessionAge, 0, len(manager.sessions))
|
|
for _, session := range manager.sessions {
|
|
ages = append(ages, sessionAge{session: session, last: session.lastUsed()})
|
|
}
|
|
sort.Slice(ages, func(i, j int) bool {
|
|
return ages[i].last.Before(ages[j].last)
|
|
})
|
|
for _, a := range ages {
|
|
if a.session.doneClosed() {
|
|
return a.session
|
|
}
|
|
}
|
|
if len(ages) <= 8 {
|
|
return nil
|
|
}
|
|
return ages[0].session
|
|
}
|
|
|
|
// lastUsed returns the session's last-used time under its own lock, so the prune
|
|
// comparator never races touch()'s write to lastUsedAt.
|
|
func (session *execSession) lastUsed() time.Time {
|
|
session.mu.Lock()
|
|
defer session.mu.Unlock()
|
|
return session.lastUsedAt
|
|
}
|
|
|
|
func (manager *execSessionManager) list() []ExecSessionSnapshot {
|
|
manager.mu.Lock()
|
|
sessions := make([]*execSession, 0, len(manager.sessions))
|
|
for _, session := range manager.sessions {
|
|
if !session.doneClosed() {
|
|
sessions = append(sessions, session)
|
|
}
|
|
}
|
|
manager.mu.Unlock()
|
|
|
|
snapshots := make([]ExecSessionSnapshot, 0, len(sessions))
|
|
for _, session := range sessions {
|
|
snapshots = append(snapshots, session.snapshot())
|
|
}
|
|
sort.Slice(snapshots, func(i, j int) bool {
|
|
return snapshots[i].ID < snapshots[j].ID
|
|
})
|
|
return snapshots
|
|
}
|
|
|
|
func (manager *execSessionManager) stop(id int) bool {
|
|
session, ok := manager.get(id)
|
|
if !ok {
|
|
return false
|
|
}
|
|
session.terminate()
|
|
return true
|
|
}
|
|
|
|
func (manager *execSessionManager) stopAll() []int {
|
|
manager.mu.Lock()
|
|
sessions := make([]*execSession, 0, len(manager.sessions))
|
|
for _, session := range manager.sessions {
|
|
if !session.doneClosed() {
|
|
sessions = append(sessions, session)
|
|
}
|
|
}
|
|
manager.mu.Unlock()
|
|
ids := make([]int, 0, len(sessions))
|
|
for _, session := range sessions {
|
|
session.terminate()
|
|
ids = append(ids, session.id)
|
|
}
|
|
waitForExecSessions(sessions, execSessionStopTimeout)
|
|
sort.Ints(ids)
|
|
return ids
|
|
}
|
|
|
|
func waitForExecSessions(sessions []*execSession, timeout time.Duration) {
|
|
if timeout <= 0 || len(sessions) == 0 {
|
|
return
|
|
}
|
|
deadline := time.Now().Add(timeout)
|
|
for _, session := range sessions {
|
|
remaining := time.Until(deadline)
|
|
if remaining <= 0 {
|
|
return
|
|
}
|
|
timer := time.NewTimer(remaining)
|
|
select {
|
|
case <-session.done:
|
|
case <-timer.C:
|
|
return
|
|
}
|
|
if !timer.Stop() {
|
|
select {
|
|
case <-timer.C:
|
|
default:
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
type ExecSessionSnapshot struct {
|
|
ID int
|
|
Command string
|
|
Cwd string
|
|
RelativeCwd string
|
|
StartedAt time.Time
|
|
LastUsedAt time.Time
|
|
TTY bool
|
|
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 {
|
|
ExecSessions() []ExecSessionSnapshot
|
|
StopExecSession(id int) bool
|
|
StopAllExecSessions() []int
|
|
}
|
|
|
|
type execSession struct {
|
|
id int
|
|
commandText string
|
|
cwd string
|
|
relativeCwd string
|
|
startedAt time.Time
|
|
lastUsedAt time.Time
|
|
tty bool
|
|
command *exec.Cmd
|
|
plan zeroSandbox.CommandPlan
|
|
cancel context.CancelFunc
|
|
stdin io.WriteCloser
|
|
cleanup func()
|
|
output *execOutputBuffer
|
|
|
|
doneOnce sync.Once
|
|
done chan struct{}
|
|
mu sync.Mutex
|
|
exitCode *int
|
|
waitErr error
|
|
}
|
|
|
|
func (session *execSession) markDone(err error, exitCode int) {
|
|
session.mu.Lock()
|
|
session.waitErr = err
|
|
session.exitCode = &exitCode
|
|
session.mu.Unlock()
|
|
session.doneOnce.Do(func() { close(session.done) })
|
|
}
|
|
|
|
func (session *execSession) doneClosed() bool {
|
|
select {
|
|
case <-session.done:
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
func (session *execSession) exitStatus() (int, bool) {
|
|
session.mu.Lock()
|
|
defer session.mu.Unlock()
|
|
if session.exitCode == nil {
|
|
return 0, false
|
|
}
|
|
return *session.exitCode, true
|
|
}
|
|
|
|
func (session *execSession) touch() {
|
|
session.mu.Lock()
|
|
session.lastUsedAt = time.Now()
|
|
session.mu.Unlock()
|
|
}
|
|
|
|
func (session *execSession) terminate() {
|
|
if session.cancel != nil {
|
|
session.cancel()
|
|
}
|
|
}
|
|
|
|
func (session *execSession) snapshot() ExecSessionSnapshot {
|
|
session.mu.Lock()
|
|
startedAt := session.startedAt
|
|
lastUsedAt := session.lastUsedAt
|
|
exitCode := session.exitCode
|
|
var copiedExit *int
|
|
if exitCode != nil {
|
|
value := *exitCode
|
|
copiedExit = &value
|
|
}
|
|
session.mu.Unlock()
|
|
status := "running"
|
|
if copiedExit != nil {
|
|
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(),
|
|
OutputTruncated: session.output.peekTruncated(),
|
|
}
|
|
}
|
|
|
|
type execOutputBuffer struct {
|
|
mu sync.Mutex
|
|
data []byte
|
|
recent []byte
|
|
truncated bool
|
|
notify chan struct{}
|
|
}
|
|
|
|
func newExecOutputBuffer() *execOutputBuffer {
|
|
return &execOutputBuffer{notify: make(chan struct{}, 1)}
|
|
}
|
|
|
|
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:]
|
|
}
|
|
buffer.mu.Unlock()
|
|
select {
|
|
case buffer.notify <- struct{}{}:
|
|
default:
|
|
}
|
|
return len(p), nil
|
|
}
|
|
|
|
func (buffer *execOutputBuffer) recentString() string {
|
|
buffer.mu.Lock()
|
|
defer buffer.mu.Unlock()
|
|
return string(buffer.recent)
|
|
}
|
|
|
|
func (buffer *execOutputBuffer) drainString() string {
|
|
buffer.mu.Lock()
|
|
defer buffer.mu.Unlock()
|
|
if len(buffer.data) == 0 {
|
|
return ""
|
|
}
|
|
out := string(buffer.data)
|
|
buffer.data = nil
|
|
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
|
|
scope PathScope
|
|
manager *execSessionManager
|
|
}
|
|
|
|
func NewExecCommandTool(workspaceRoot string, manager *execSessionManager) Tool {
|
|
return NewScopedExecCommandTool(workspaceRoot, nil, manager)
|
|
}
|
|
|
|
func NewScopedExecCommandTool(workspaceRoot string, scope PathScope, manager *execSessionManager) Tool {
|
|
if manager == nil {
|
|
manager = defaultExecSessionManager
|
|
}
|
|
description := "Runs a command in a PTY, returning output or a session ID for ongoing interaction."
|
|
if runtimeGOOS() == "windows" {
|
|
description += "\n\n" + shellGuidanceForGOOS(runtimeGOOS())
|
|
}
|
|
return execCommandTool{
|
|
baseTool: baseTool{
|
|
name: ExecCommandToolName,
|
|
description: description,
|
|
parameters: Schema{
|
|
Type: "object",
|
|
Properties: map[string]PropertySchema{
|
|
"cmd": {Type: "string", Description: "Shell command to execute."},
|
|
"workdir": {Type: "string", Description: "Working directory for the command. Defaults to the turn cwd.", Default: "."},
|
|
"cwd": {Type: "string", Description: "Alias for workdir. Prefer workdir.", Default: "."},
|
|
"yield_time_ms": {Type: "integer", Description: "Wait before yielding output. Defaults to 10000 ms; effective range is 250-30000 ms.", Default: defaultExecYieldTimeMS, Minimum: intPtr(1), Maximum: intPtr(maxExecYieldTimeMS)},
|
|
"max_output_tokens": {Type: "integer", Description: "Output token budget. Defaults to 10000 tokens; larger requests may be capped by policy.", Default: defaultMaxOutputTokens, Minimum: intPtr(1), Maximum: intPtr(maxExecOutputTokenRequest)},
|
|
"sandbox_permissions": {Type: "string", Enum: []string{string(SandboxPermissionsUseDefault), string(SandboxPermissionsWithAdditionalPermissions), string(SandboxPermissionsRequireEscalated)}, Description: "Per-command sandbox override. Defaults to `use_default`; use `with_additional_permissions` with `additional_permissions` for sandboxed file/network access, or `require_escalated` only when the command must run outside the sandbox, such as host/global process, socket, service, or desktop state hidden by sandbox namespaces.", Default: string(SandboxPermissionsUseDefault)},
|
|
"additional_permissions": {
|
|
Type: "object",
|
|
Description: "Sandboxed filesystem or network access for this command; only with `sandbox_permissions: \"with_additional_permissions\"`.",
|
|
Properties: additionalPermissionsProperties(),
|
|
},
|
|
"justification": {Type: "string", Description: "User-facing approval question for `require_escalated`; omit otherwise."},
|
|
"prefix_rule": {Type: "array", Items: &PropertySchema{Type: "string"}, Description: "Reusable approval prefix for this command, only with `sandbox_permissions: \"require_escalated\"`; keep it narrow, for example [\"git\", \"pull\"]."},
|
|
"tty": {Type: "boolean", Description: "True allocates a PTY for the command; false or omitted uses plain pipes.", Default: false},
|
|
},
|
|
Required: []string{"cmd"},
|
|
AdditionalProperties: false,
|
|
},
|
|
safety: promptSafety(SideEffectShell, "Shell commands can read, write, or execute programs."),
|
|
},
|
|
workspaceRoot: normalizeWorkspaceRoot(workspaceRoot),
|
|
scope: scope,
|
|
manager: manager,
|
|
}
|
|
}
|
|
|
|
func (tool execCommandTool) Run(ctx context.Context, args map[string]any) Result {
|
|
return tool.run(ctx, args, nil)
|
|
}
|
|
|
|
func (tool execCommandTool) RunWithSandbox(ctx context.Context, args map[string]any, engine *zeroSandbox.Engine) Result {
|
|
return tool.run(ctx, args, engine)
|
|
}
|
|
|
|
func (tool execCommandTool) ExecSessions() []ExecSessionSnapshot {
|
|
return tool.manager.list()
|
|
}
|
|
|
|
func (tool execCommandTool) StopExecSession(id int) bool {
|
|
return tool.manager.stop(id)
|
|
}
|
|
|
|
func (tool execCommandTool) StopAllExecSessions() []int {
|
|
return tool.manager.stopAll()
|
|
}
|
|
|
|
func (tool execCommandTool) run(ctx context.Context, args map[string]any, engine *zeroSandbox.Engine) Result {
|
|
commandText, err := aliasedStringArg(args, []string{"cmd", "command", "script", "shell"}, "", true, false)
|
|
if err != nil {
|
|
return errorResult("Error: Invalid arguments for exec_command: " + err.Error())
|
|
}
|
|
workdir, err := aliasedStringArg(args, []string{"workdir", "cwd", "dir", "directory"}, ".", false, true)
|
|
if err != nil {
|
|
return errorResult("Error: Invalid arguments for exec_command: " + err.Error())
|
|
}
|
|
yieldTimeMS, err := intArg(args, "yield_time_ms", defaultExecYieldTimeMS, 1, maxExecYieldTimeMS)
|
|
if err != nil {
|
|
return errorResult("Error: Invalid arguments for exec_command: " + err.Error())
|
|
}
|
|
maxOutputTokens, err := intArg(args, "max_output_tokens", defaultMaxOutputTokens, 1, maxExecOutputTokenRequest)
|
|
if err != nil {
|
|
return errorResult("Error: Invalid arguments for exec_command: " + err.Error())
|
|
}
|
|
ttyRequested, err := boolArg(args, "tty", false)
|
|
if err != nil {
|
|
return errorResult("Error: Invalid arguments for exec_command: " + err.Error())
|
|
}
|
|
sandboxPermissions, err := sandboxPermissionsArg(args)
|
|
if err != nil {
|
|
return errorResult("Error: Invalid arguments for exec_command: " + err.Error())
|
|
}
|
|
if issue := detectShellCommandIssue(commandText, runtimeGOOS()); issue != nil {
|
|
return shellIssueBlockResult(*issue)
|
|
}
|
|
if interactive := zeroSandbox.DetectInteractiveCommand(commandText, runtimeGOOS()); interactive.Interactive {
|
|
return interactiveBlockResult(interactive)
|
|
}
|
|
absoluteCwd, relativeCwd, err := resolveScopedPath(tool.workspaceRoot, tool.scope, workdir)
|
|
if err != nil {
|
|
return errorResult("Error running exec_command: " + err.Error())
|
|
}
|
|
|
|
session, err := tool.startSession(commandText, absoluteCwd, relativeCwd, ttyRequested, engine, sandboxPermissions)
|
|
if err != nil {
|
|
return errorResult("Error starting exec_command: " + err.Error())
|
|
}
|
|
output, outputTruncated := session.collect(ctx, time.Duration(yieldTimeMS)*time.Millisecond)
|
|
if ctx != nil && ctx.Err() != nil && !session.doneClosed() {
|
|
session.terminate()
|
|
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,
|
|
outputBufferTruncated: outputTruncated,
|
|
sessionID: session.id,
|
|
exitCode: exitCode,
|
|
exited: exited,
|
|
relativeCwd: relativeCwd,
|
|
tty: session.tty,
|
|
plan: session.plan,
|
|
maxOutputTokens: maxOutputTokens,
|
|
})
|
|
}
|
|
|
|
func (tool execCommandTool) startSession(commandText string, absoluteCwd string, relativeCwd string, ttyRequested bool, engine *zeroSandbox.Engine, sandboxPermissions SandboxPermissionOverride) (*execSession, error) {
|
|
id := tool.manager.allocateID()
|
|
commandCtx, cancel := context.WithCancel(context.Background())
|
|
commandEngine := commandEngineForSandboxPermissions(engine, sandboxPermissions)
|
|
command, plan, err := buildBashCommand(commandCtx, commandText, absoluteCwd, commandEngine)
|
|
if err != nil {
|
|
cancel()
|
|
return nil, err
|
|
}
|
|
output := newExecOutputBuffer()
|
|
monitor := zeroSandbox.StartDenialMonitor(context.Background(), plan.MonitorTag)
|
|
stdin, tty, cleanup, err := startExecProcess(command, output, ttyRequested)
|
|
if err != nil {
|
|
_ = monitor.Stop()
|
|
plan.Cleanup()
|
|
cancel()
|
|
return nil, err
|
|
}
|
|
session := &execSession{
|
|
id: id,
|
|
commandText: commandText,
|
|
cwd: absoluteCwd,
|
|
relativeCwd: relativeCwd,
|
|
startedAt: time.Now(),
|
|
lastUsedAt: time.Now(),
|
|
tty: tty,
|
|
command: command,
|
|
plan: plan,
|
|
cancel: cancel,
|
|
stdin: stdin,
|
|
cleanup: cleanup,
|
|
output: output,
|
|
done: make(chan struct{}),
|
|
}
|
|
tool.manager.store(session)
|
|
tool.manager.removeCompletedLater(session)
|
|
go func() {
|
|
err := command.Wait()
|
|
if blocks := monitor.Stop(); len(blocks) > 0 {
|
|
output.Write([]byte(appendSandboxBlocks("", blocks)))
|
|
}
|
|
if session.cleanup != nil {
|
|
session.cleanup()
|
|
}
|
|
plan.Cleanup()
|
|
cancel()
|
|
session.markDone(err, commandExitCode(err))
|
|
}()
|
|
return session, nil
|
|
}
|
|
|
|
// 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
|
|
}
|
|
// A background process that keeps writing output faster than this
|
|
// loop can drain it (a chatty dev server, a watch-mode rebuild loop,
|
|
// a crash-restart spew) would otherwise never let drainString return
|
|
// empty, so the deadline check below — only reached on an empty
|
|
// drain — would never fire. That turns `wait` into an unbounded
|
|
// block regardless of what the caller asked for. Check it here too,
|
|
// so continuous output can never starve the timeout.
|
|
if time.Now().After(deadline) {
|
|
return finish()
|
|
}
|
|
continue
|
|
}
|
|
if session.doneClosed() {
|
|
if chunk := session.output.drainString(); chunk != "" {
|
|
builder.WriteString(chunk)
|
|
}
|
|
return finish()
|
|
}
|
|
remaining := time.Until(deadline)
|
|
if remaining <= 0 {
|
|
return finish()
|
|
}
|
|
timer := time.NewTimer(remaining)
|
|
select {
|
|
case <-session.output.notify:
|
|
if !timer.Stop() {
|
|
select {
|
|
case <-timer.C:
|
|
default:
|
|
}
|
|
}
|
|
case <-session.done:
|
|
if !timer.Stop() {
|
|
select {
|
|
case <-timer.C:
|
|
default:
|
|
}
|
|
}
|
|
case <-ctx.Done():
|
|
if !timer.Stop() {
|
|
select {
|
|
case <-timer.C:
|
|
default:
|
|
}
|
|
}
|
|
return finish()
|
|
case <-timer.C:
|
|
return finish()
|
|
}
|
|
}
|
|
}
|
|
|
|
type writeStdinTool struct {
|
|
baseTool
|
|
manager *execSessionManager
|
|
}
|
|
|
|
func NewWriteStdinTool(manager *execSessionManager) Tool {
|
|
if manager == nil {
|
|
manager = defaultExecSessionManager
|
|
}
|
|
return writeStdinTool{
|
|
baseTool: baseTool{
|
|
name: WriteStdinToolName,
|
|
description: "Writes characters to an existing unified exec session and returns recent output.",
|
|
parameters: Schema{
|
|
Type: "object",
|
|
Properties: map[string]PropertySchema{
|
|
"session_id": {Type: "integer", Description: "Identifier of the running unified exec session."},
|
|
"chars": {Type: "string", Description: "Bytes to write to stdin. Defaults to empty, which polls without writing.", Default: ""},
|
|
"yield_time_ms": {Type: "integer", Description: "Wait before yielding output. Non-empty writes default to 250 ms and cap at 30000 ms; empty polls wait 5000-300000 ms by default.", Default: defaultPollYieldTimeMS, Minimum: intPtr(1), Maximum: intPtr(maxPollYieldTimeMS)},
|
|
"max_output_tokens": {Type: "integer", Description: "Output token budget. Defaults to 10000 tokens; larger requests may be capped by policy.", Default: defaultMaxOutputTokens, Minimum: intPtr(1), Maximum: intPtr(maxExecOutputTokenRequest)},
|
|
},
|
|
Required: []string{"session_id"},
|
|
AdditionalProperties: false,
|
|
},
|
|
safety: Safety{
|
|
SideEffect: SideEffectShell,
|
|
Permission: PermissionPrompt,
|
|
Reason: "Sending stdin can drive an existing shell process beyond the original command; empty polling and Ctrl-C interrupts are allowed automatically.",
|
|
AdvertiseInAuto: true,
|
|
},
|
|
},
|
|
manager: manager,
|
|
}
|
|
}
|
|
|
|
func (tool writeStdinTool) PermissionForArgs(args map[string]any) Permission {
|
|
raw, ok := args["chars"]
|
|
if !ok || raw == nil {
|
|
return PermissionAllow
|
|
}
|
|
chars, ok := raw.(string)
|
|
if !ok {
|
|
return PermissionPrompt
|
|
}
|
|
if chars == "" || chars == "\x03" {
|
|
return PermissionAllow
|
|
}
|
|
return PermissionPrompt
|
|
}
|
|
|
|
func (tool writeStdinTool) Run(ctx context.Context, args map[string]any) Result {
|
|
return tool.RunWithOptions(ctx, args, RunOptions{})
|
|
}
|
|
|
|
func (tool writeStdinTool) RunWithOptions(ctx context.Context, args map[string]any, _ RunOptions) Result {
|
|
if value, ok := args["session_id"]; !ok || value == nil {
|
|
return errorResult("Error: Invalid arguments for write_stdin: session_id is required")
|
|
}
|
|
sessionID, err := intArg(args, "session_id", 0, 1, 0)
|
|
if err != nil {
|
|
return errorResult("Error: Invalid arguments for write_stdin: " + err.Error())
|
|
}
|
|
chars, err := stringArgWithEmpty(args, "chars", "", false, true)
|
|
if err != nil {
|
|
return errorResult("Error: Invalid arguments for write_stdin: " + err.Error())
|
|
}
|
|
yieldTimeMS, err := intArg(args, "yield_time_ms", defaultPollYieldTimeMS, 1, maxPollYieldTimeMS)
|
|
if err != nil {
|
|
return errorResult("Error: Invalid arguments for write_stdin: " + err.Error())
|
|
}
|
|
maxOutputTokens, err := intArg(args, "max_output_tokens", defaultMaxOutputTokens, 1, maxExecOutputTokenRequest)
|
|
if err != nil {
|
|
return errorResult("Error: Invalid arguments for write_stdin: " + err.Error())
|
|
}
|
|
session, ok := tool.manager.get(sessionID)
|
|
if !ok {
|
|
return errorResult(fmt.Sprintf("Error: Unknown exec session_id %d.", sessionID))
|
|
}
|
|
session.touch()
|
|
interrupted := false
|
|
if chars != "" {
|
|
if shouldInterruptExecSession(chars, session.tty) {
|
|
interrupted = true
|
|
session.terminate()
|
|
} else if !session.tty {
|
|
return errorResult(fmt.Sprintf("Error: exec session_id %d does not accept stdin. Use empty chars to poll, or send chars \"\\u0003\" to interrupt/stop it.", sessionID))
|
|
} else if session.stdin != nil {
|
|
if _, err := io.WriteString(session.stdin, chars); err != nil && !session.doneClosed() {
|
|
return errorResult("Error writing to exec session: " + err.Error())
|
|
}
|
|
}
|
|
}
|
|
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,
|
|
outputBufferTruncated: outputTruncated,
|
|
sessionID: session.id,
|
|
exitCode: exitCode,
|
|
exited: exited,
|
|
relativeCwd: session.relativeCwd,
|
|
tty: session.tty,
|
|
interrupted: interrupted,
|
|
plan: session.plan,
|
|
maxOutputTokens: maxOutputTokens,
|
|
})
|
|
}
|
|
|
|
func shouldInterruptExecSession(chars string, tty bool) bool {
|
|
if strings.Contains(chars, "\x03") {
|
|
return true
|
|
}
|
|
normalized := strings.ToLower(strings.TrimSpace(chars))
|
|
normalizedNoSpace := strings.ReplaceAll(normalized, " ", "")
|
|
switch normalizedNoSpace {
|
|
case `\u0003`, `\\u0003`, `\x03`, `\\x03`, "^c", "ctrl-c", "control-c", "sigint", "interrupt":
|
|
return true
|
|
}
|
|
if tty {
|
|
return false
|
|
}
|
|
switch normalized {
|
|
case "q", "quit", "exit", "stop", "kill", "terminate":
|
|
return true
|
|
}
|
|
return false
|
|
}
|
|
|
|
type execToolResultInput struct {
|
|
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 {
|
|
output, truncated := truncateExecOutput(input.output, input.maxOutputTokens)
|
|
meta := map[string]string{
|
|
"cwd": input.relativeCwd,
|
|
"tty": strconv.FormatBool(input.tty),
|
|
}
|
|
addSandboxMeta(meta, input.plan)
|
|
if input.exited {
|
|
meta["exit_code"] = strconv.Itoa(input.exitCode)
|
|
if input.interrupted {
|
|
meta["interrupted"] = "true"
|
|
}
|
|
if !input.interrupted {
|
|
markLikelySandboxDenial(meta, input.plan, input.exitCode, output)
|
|
}
|
|
} 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 || input.outputBufferTruncated,
|
|
Meta: meta,
|
|
Display: Display{
|
|
Summary: execDisplaySummary(input.commandText, input.sessionID, input.exited, input.exitCode),
|
|
Kind: "shell",
|
|
},
|
|
}
|
|
}
|
|
|
|
func formatExecCommandOutput(output string, sessionID int, exited bool, exitCode int, interrupted bool) string {
|
|
output = strings.TrimRight(output, "\r\n")
|
|
parts := []string{}
|
|
if output != "" {
|
|
parts = append(parts, "output:\n"+output)
|
|
}
|
|
if exited {
|
|
if output == "" {
|
|
if interrupted {
|
|
parts = append(parts, "Command interrupted.")
|
|
} else {
|
|
parts = append(parts, "Command completed with no output.")
|
|
}
|
|
}
|
|
if interrupted {
|
|
parts = append(parts, "interrupted: true")
|
|
}
|
|
parts = append(parts, fmt.Sprintf("exit_code: %d", exitCode))
|
|
} else {
|
|
if output == "" {
|
|
parts = append(parts, "Command is still running.")
|
|
}
|
|
parts = append(parts, fmt.Sprintf("session_id: %d", sessionID))
|
|
parts = append(parts, fmt.Sprintf("Use write_stdin with session_id %d and empty chars to poll; send chars \"\\u0003\" to interrupt/stop it.", sessionID))
|
|
}
|
|
return strings.Join(parts, "\n")
|
|
}
|
|
|
|
func truncateExecOutput(output string, maxOutputTokens int) (string, bool) {
|
|
if maxOutputTokens <= 0 {
|
|
maxOutputTokens = defaultMaxOutputTokens
|
|
}
|
|
maxBytes := maxOutputTokens * 4
|
|
if len(output) <= maxBytes {
|
|
return output, false
|
|
}
|
|
head := maxBytes / 2
|
|
tail := maxBytes - head
|
|
return utf8Prefix(output, head) + "\n[zero] output truncated\n" + utf8Suffix(output, tail), true
|
|
}
|
|
|
|
func utf8Prefix(value string, maxBytes int) string {
|
|
if maxBytes <= 0 {
|
|
return ""
|
|
}
|
|
if len(value) <= maxBytes {
|
|
return value
|
|
}
|
|
for maxBytes > 0 && !utf8.RuneStart(value[maxBytes]) {
|
|
maxBytes--
|
|
}
|
|
return value[:maxBytes]
|
|
}
|
|
|
|
func utf8Suffix(value string, maxBytes int) string {
|
|
if maxBytes <= 0 {
|
|
return ""
|
|
}
|
|
if len(value) <= maxBytes {
|
|
return value
|
|
}
|
|
start := len(value) - maxBytes
|
|
for start < len(value) && !utf8.RuneStart(value[start]) {
|
|
start++
|
|
}
|
|
return value[start:]
|
|
}
|
|
|
|
func execDisplaySummary(commandText string, sessionID int, exited bool, exitCode int) string {
|
|
commandText = strings.TrimSpace(commandText)
|
|
if commandText == "" {
|
|
commandText = "command"
|
|
}
|
|
if exited {
|
|
return fmt.Sprintf("%s exited with code %d", commandText, exitCode)
|
|
}
|
|
return fmt.Sprintf("%s still running as session %d", commandText, sessionID)
|
|
}
|
|
|
|
func runtimeGOOS() string {
|
|
return runtime.GOOS
|
|
}
|