mirror of
https://github.com/Twigpine/zero.git
synced 2026-11-08 09:49:34 -10:00
Compare commits
15
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a5553394fe | ||
|
|
7cebc447ca | ||
|
|
e93519e2d1 | ||
|
|
ea27648d0d | ||
|
|
b9e811e5d9 | ||
|
|
c90b984aec | ||
|
|
a68e8f3e42 | ||
|
|
4e7d7b3bba | ||
|
|
ca33fc5feb | ||
|
|
57e26e6042 | ||
|
|
ebc5b0fac7 | ||
|
|
4377aa9693 | ||
|
|
7bc135249e | ||
|
|
eeda3de2ad | ||
|
|
41aeb30016 |
No files matched your search
+16
-3
@@ -272,7 +272,8 @@ sequenceDiagram
|
||||
Agent->>Tool: execute(ctx, args)
|
||||
Tool-->>Agent: output, error, metadata
|
||||
Agent->>Hooks: afterTool hooks
|
||||
Agent->>Agent: redact, truncate, classify success/failure
|
||||
Agent->>Agent: redact secrets, apply semantic output budget, enforce byte ceiling
|
||||
Agent->>Agent: classify success/failure
|
||||
Agent-->>Surface: OnToolResult callback
|
||||
Agent->>Transcript: append tool result message
|
||||
end
|
||||
@@ -604,10 +605,22 @@ flowchart TD
|
||||
Prompt -- Yes --> Decision[Permission callback]
|
||||
Prompt -- No --> Run[Registry.RunWithOptions]
|
||||
Decision --> Run
|
||||
Run --> Redact[Redact secrets + enforce output ceiling]
|
||||
Redact --> Message[Return tool result to model]
|
||||
Run --> Redact[Redact secrets]
|
||||
Redact --> Budget[Token-aware semantic output budget]
|
||||
Budget --> Ceiling[Existing hard byte ceiling]
|
||||
Ceiling --> Message[Return one tool result to model]
|
||||
```
|
||||
|
||||
Oversized results use deterministic, provider-neutral estimated-token budgets.
|
||||
Policies retain useful structure for files, search matches, tests, process logs,
|
||||
diffs, and worker conclusions; tools without a declared category use a UTF-8-safe
|
||||
head/tail fallback. The estimate is intentionally conservative for non-ASCII
|
||||
text and is not exact provider tokenization. Existing byte ceilings remain the
|
||||
authoritative safety limit. When the existing spill mechanism can persist the
|
||||
complete redacted text received by the budgeting layer, the result includes its
|
||||
safe spill reference. This does not imply capture of subprocess bytes already
|
||||
discarded by a tool's established internal buffer.
|
||||
|
||||
Core tool groups include:
|
||||
|
||||
- **Read-only tools**: file reads, directory listing, glob, grep, LSP navigation,
|
||||
|
||||
@@ -141,9 +141,11 @@ func TestRunAskUserCancellationAbortsRun(t *testing.T) {
|
||||
registry := registryWithAskUser()
|
||||
args := `{"questions":[{"question":"Which framework?"}]}`
|
||||
provider := providerCallingAskUserThenAnswer(args, "done")
|
||||
var toolResults []ToolResult
|
||||
|
||||
result, err := Run(context.Background(), "clarify", provider, Options{
|
||||
Registry: registry,
|
||||
Registry: registry,
|
||||
OnToolResult: func(result ToolResult) { toolResults = append(toolResults, result) },
|
||||
OnAskUser: func(_ context.Context, _ AskUserRequest) (AskUserResponse, error) {
|
||||
return AskUserResponse{}, context.Canceled
|
||||
},
|
||||
@@ -158,6 +160,9 @@ func TestRunAskUserCancellationAbortsRun(t *testing.T) {
|
||||
if len(provider.requests) != 1 {
|
||||
t.Fatalf("expected the run to stop after the canceled ask_user (1 turn), got %d", len(provider.requests))
|
||||
}
|
||||
if len(toolResults) != 1 || toolResults[0].ToolCallID == "" {
|
||||
t.Fatalf("cancellation must emit exactly one result for the call, got %#v", toolResults)
|
||||
}
|
||||
// The recorded tool result must reflect cancellation, not a synthetic answer.
|
||||
var toolMsg string
|
||||
for _, m := range result.Messages {
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"fmt"
|
||||
"io"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
@@ -656,6 +657,7 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
|
||||
toolSpan.End()
|
||||
}
|
||||
options.Trace.Counter(trace.CounterToolCalls, 1)
|
||||
recordOutputBudgetTrace(options.Trace, toolResult)
|
||||
if options.OnToolResult != nil {
|
||||
options.OnToolResult(toolResult)
|
||||
}
|
||||
@@ -841,6 +843,28 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func recordOutputBudgetTrace(recorder *trace.Recorder, result ToolResult) {
|
||||
if recorder == nil || result.Meta["output_budget_category"] == "" {
|
||||
return
|
||||
}
|
||||
parseInt := func(key string) int {
|
||||
value, _ := strconv.Atoi(result.Meta[key])
|
||||
return value
|
||||
}
|
||||
spillCreated, _ := strconv.ParseBool(result.Meta["output_budget_spill_created"])
|
||||
recorder.EmitOutputBudget(trace.OutputBudgetEvent{
|
||||
Tool: result.Name,
|
||||
Category: result.Meta["output_budget_category"],
|
||||
OriginalBytes: parseInt("output_budget_original_bytes"),
|
||||
RetainedBytes: parseInt("output_budget_retained_bytes"),
|
||||
EstimatedOriginalTokens: parseInt("output_budget_estimated_original_tokens"),
|
||||
EstimatedRetainedTokens: parseInt("output_budget_estimated_retained_tokens"),
|
||||
Truncated: result.Truncated,
|
||||
Reason: result.Meta["output_budget_reason"],
|
||||
SpillCreated: spillCreated,
|
||||
})
|
||||
}
|
||||
|
||||
func finalAnswerAfterMaxTurns(ctx context.Context, provider Provider, messages []zeroruntime.Message, toolDefs []zeroruntime.ToolDefinition, options Options) (string, []zeroruntime.Message, string) {
|
||||
finalMessages := copyMessages(messages)
|
||||
finalMessages = append(finalMessages, zeroruntime.Message{
|
||||
@@ -1295,6 +1319,7 @@ func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCal
|
||||
if didRedact {
|
||||
result.Redacted = true
|
||||
}
|
||||
result = registry.RebudgetAfterHook(call.Name, args, result)
|
||||
}
|
||||
}
|
||||
// Secret scrubbing happens at the registry boundary (the single point both
|
||||
@@ -1305,6 +1330,7 @@ func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCal
|
||||
Name: call.Name,
|
||||
Status: result.Status,
|
||||
Output: result.Output,
|
||||
Truncated: result.Truncated,
|
||||
Meta: result.Meta,
|
||||
Redacted: result.Redacted,
|
||||
ChangedFiles: result.ChangedFiles,
|
||||
@@ -1615,6 +1641,7 @@ func toolResultFromPrePermissionReject(call ToolCall, result tools.Result) ToolR
|
||||
Name: call.Name,
|
||||
Status: result.Status,
|
||||
Output: output,
|
||||
Truncated: result.Truncated,
|
||||
Meta: meta,
|
||||
Redacted: result.Redacted || outputRedacted || summaryRedacted || metaRedacted,
|
||||
ChangedFiles: result.ChangedFiles,
|
||||
@@ -1884,6 +1911,7 @@ func askUserFallbackResult(ctx context.Context, registry *tools.Registry, call T
|
||||
Name: call.Name,
|
||||
Status: result.Status,
|
||||
Output: result.Output,
|
||||
Truncated: result.Truncated,
|
||||
Meta: result.Meta,
|
||||
Redacted: result.Redacted,
|
||||
ChangedFiles: result.ChangedFiles,
|
||||
|
||||
@@ -0,0 +1,167 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/hooks"
|
||||
"github.com/Gitlawb/zero/internal/tools"
|
||||
"github.com/Gitlawb/zero/internal/trace"
|
||||
"github.com/Gitlawb/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
type propagationOutputTool struct {
|
||||
output string
|
||||
}
|
||||
|
||||
func TestOutputBudgetHookHelperProcess(t *testing.T) {
|
||||
for index, arg := range os.Args {
|
||||
if arg != "--zero-output-budget-hook" || index+1 >= len(os.Args) {
|
||||
continue
|
||||
}
|
||||
if _, err := os.Stdout.WriteString(os.Args[index+1]); err != nil {
|
||||
os.Exit(2)
|
||||
}
|
||||
os.Exit(0)
|
||||
}
|
||||
}
|
||||
|
||||
func largeOutputBudgetHookDispatcher() *hooks.Dispatcher {
|
||||
feedback := strings.Repeat("hook feedback ", 200)
|
||||
return hooks.NewDispatcher(hooks.DispatcherOptions{Config: hooks.Config{
|
||||
Enabled: true,
|
||||
Hooks: []hooks.Definition{{
|
||||
ID: "large-feedback",
|
||||
Event: hooks.EventAfterTool,
|
||||
Matcher: "propagation_output",
|
||||
Command: os.Args[0],
|
||||
Args: []string{
|
||||
"-test.run=TestOutputBudgetHookHelperProcess",
|
||||
"--",
|
||||
"--zero-output-budget-hook",
|
||||
feedback,
|
||||
},
|
||||
Enabled: true,
|
||||
}},
|
||||
}})
|
||||
}
|
||||
|
||||
func (tool propagationOutputTool) Name() string { return "propagation_output" }
|
||||
func (tool propagationOutputTool) Description() string { return "returns output for propagation tests" }
|
||||
func (tool propagationOutputTool) Parameters() tools.Schema {
|
||||
return tools.Schema{Type: "object", AdditionalProperties: false}
|
||||
}
|
||||
func (tool propagationOutputTool) Safety() tools.Safety {
|
||||
return tools.Safety{SideEffect: tools.SideEffectRead, Permission: tools.PermissionAllow, Reason: "test read"}
|
||||
}
|
||||
func (tool propagationOutputTool) Run(context.Context, map[string]any) tools.Result {
|
||||
return tools.Result{Status: tools.StatusOK, Output: tool.output}
|
||||
}
|
||||
|
||||
func TestExecuteToolCallPropagatesOutputTruncation(t *testing.T) {
|
||||
t.Setenv("TMPDIR", t.TempDir())
|
||||
t.Setenv("ZERO_TOOL_OUTPUT_CEILING_TOKENS", "80")
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(propagationOutputTool{output: strings.Repeat("large output\n", 1000)})
|
||||
|
||||
result, abortErr := executeToolCall(context.Background(), registry, ToolCall{
|
||||
ID: "call-budget",
|
||||
Name: "propagation_output",
|
||||
Arguments: `{}`,
|
||||
}, PermissionModeAuto, Options{Cwd: t.TempDir()})
|
||||
if abortErr != nil {
|
||||
t.Fatalf("executeToolCall abort error: %v", abortErr)
|
||||
}
|
||||
if !result.Truncated {
|
||||
t.Fatalf("agent ToolResult lost tools.Result.Truncated: %#v", result)
|
||||
}
|
||||
if result.Meta["spill_path"] == "" {
|
||||
t.Fatalf("agent ToolResult lost spill metadata: %#v", result.Meta)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecordOutputBudgetTraceUsesOnlyCompactMetadata(t *testing.T) {
|
||||
recorder := trace.NewRecorder("session", "run", "")
|
||||
recorder.Start()
|
||||
recordOutputBudgetTrace(recorder, ToolResult{
|
||||
Name: "grep",
|
||||
Truncated: true,
|
||||
Output: "SECRET OUTPUT MUST NOT ENTER TRACE",
|
||||
Meta: map[string]string{
|
||||
"output_budget_category": "search",
|
||||
"output_budget_original_bytes": "1000",
|
||||
"output_budget_retained_bytes": "100",
|
||||
"output_budget_estimated_original_tokens": "250",
|
||||
"output_budget_estimated_retained_tokens": "25",
|
||||
"output_budget_reason": "semantic_search_budget",
|
||||
"output_budget_spill_created": "true",
|
||||
"spill_path": "/secret/path/not-for-trace",
|
||||
},
|
||||
})
|
||||
events := recorder.Finish().OutputBudgets
|
||||
if len(events) != 1 {
|
||||
t.Fatalf("events = %#v", events)
|
||||
}
|
||||
event := events[0]
|
||||
if event.Tool != "grep" || event.Category != "search" || event.OriginalBytes != 1000 || event.RetainedBytes != 100 || !event.SpillCreated {
|
||||
t.Fatalf("unexpected trace event: %#v", event)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteToolCallRebudgetsOversizedAfterToolFeedback(t *testing.T) {
|
||||
t.Setenv("ZERO_TOOL_OUTPUT_CEILING_TOKENS", "80")
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(propagationOutputTool{output: "tool output"})
|
||||
dispatcher := largeOutputBudgetHookDispatcher()
|
||||
|
||||
result, abortErr := executeToolCall(context.Background(), registry, ToolCall{
|
||||
ID: "call-hook-budget",
|
||||
Name: "propagation_output",
|
||||
Arguments: `{}`,
|
||||
}, PermissionModeAuto, Options{Hooks: dispatcher})
|
||||
if abortErr != nil {
|
||||
t.Fatalf("executeToolCall abort error: %v", abortErr)
|
||||
}
|
||||
if !result.Truncated || len(result.Output) > 80*4 {
|
||||
t.Fatalf("afterTool feedback bypassed output budget: truncated=%t bytes=%d meta=%#v", result.Truncated, len(result.Output), result.Meta)
|
||||
}
|
||||
if result.Meta["output_budget_category"] == "" || result.Meta["output_budget_retained_bytes"] != strconv.Itoa(len(result.Output)) {
|
||||
t.Fatalf("post-hook budget metadata does not describe final output: %#v", result.Meta)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunTraceReflectsPostHookBudget(t *testing.T) {
|
||||
t.Setenv("ZERO_TOOL_OUTPUT_CEILING_TOKENS", "80")
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(propagationOutputTool{output: "tool output"})
|
||||
dispatcher := largeOutputBudgetHookDispatcher()
|
||||
provider := &mockProvider{turns: [][]zeroruntime.StreamEvent{
|
||||
{
|
||||
{Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "call-hook-trace", ToolName: "propagation_output"},
|
||||
{Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "call-hook-trace", ArgumentsFragment: `{}`},
|
||||
{Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "call-hook-trace"},
|
||||
{Type: zeroruntime.StreamEventDone},
|
||||
},
|
||||
{{Type: zeroruntime.StreamEventText, Content: "done"}, {Type: zeroruntime.StreamEventDone}},
|
||||
}}
|
||||
recorder := trace.NewRecorder("session", "run", "")
|
||||
var toolResults []ToolResult
|
||||
if _, err := Run(context.Background(), "budget hook", provider, Options{
|
||||
Registry: registry,
|
||||
Hooks: dispatcher,
|
||||
Trace: recorder,
|
||||
OnToolResult: func(result ToolResult) { toolResults = append(toolResults, result) },
|
||||
}); err != nil {
|
||||
t.Fatalf("Run: %v", err)
|
||||
}
|
||||
if len(toolResults) != 1 || !toolResults[0].Truncated {
|
||||
t.Fatalf("tool result = %#v, want one truncated post-hook result", toolResults)
|
||||
}
|
||||
events := recorder.Finish().OutputBudgets
|
||||
if len(events) != 1 || !events[0].Truncated || events[0].RetainedBytes != len(toolResults[0].Output) {
|
||||
t.Fatalf("trace does not describe final post-hook output: events=%#v result=%#v", events, toolResults[0])
|
||||
}
|
||||
}
|
||||
@@ -61,10 +61,13 @@ const (
|
||||
)
|
||||
|
||||
type ToolResult struct {
|
||||
ToolCallID string
|
||||
Name string
|
||||
Status tools.Status
|
||||
Output string
|
||||
ToolCallID string
|
||||
Name string
|
||||
Status tools.Status
|
||||
Output string
|
||||
// Truncated reports that the tool's model-visible output omitted content.
|
||||
// The full result may be recoverable through Meta["spill_path"].
|
||||
Truncated bool
|
||||
Meta map[string]string
|
||||
Redacted bool
|
||||
ChangedFiles []string
|
||||
|
||||
@@ -636,6 +636,9 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in
|
||||
if len(result.Meta) > 0 {
|
||||
payload["meta"] = result.Meta
|
||||
}
|
||||
if result.Truncated {
|
||||
payload["truncated"] = true
|
||||
}
|
||||
if result.Redacted {
|
||||
payload["redacted"] = true
|
||||
}
|
||||
|
||||
@@ -153,6 +153,9 @@ func (writer *execEventWriter) toolResult(result agent.ToolResult) {
|
||||
if len(result.Meta) > 0 {
|
||||
payload["meta"] = result.Meta
|
||||
}
|
||||
if result.Truncated {
|
||||
payload["truncated"] = true
|
||||
}
|
||||
if result.Redacted {
|
||||
payload["redacted"] = true
|
||||
}
|
||||
@@ -166,7 +169,8 @@ func (writer *execEventWriter) toolResult(result agent.ToolResult) {
|
||||
return
|
||||
}
|
||||
if writer.format == execOutputStreamJSON {
|
||||
output, truncated := truncateForStreamJSONOutput(result.Output)
|
||||
output, surfaceTruncated := truncateForStreamJSONOutput(result.Output)
|
||||
truncated := result.Truncated || surfaceTruncated
|
||||
event := streamjson.Event{
|
||||
Type: streamjson.EventToolResult,
|
||||
RunID: writer.runID,
|
||||
|
||||
@@ -1,8 +1,12 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/agent"
|
||||
"github.com/Gitlawb/zero/internal/tools"
|
||||
)
|
||||
|
||||
@@ -16,3 +20,35 @@ func TestStreamJSONSideEffectReportsNoneForControlTool(t *testing.T) {
|
||||
t.Fatalf("streamJSONSideEffect(escalate_model) = %q, want none", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecWriterPropagatesToolResultTruncation(t *testing.T) {
|
||||
for _, format := range []execOutputFormat{execOutputJSON, execOutputStreamJSON} {
|
||||
t.Run(string(format), func(t *testing.T) {
|
||||
var stdout, stderr bytes.Buffer
|
||||
writer := execEventWriter{
|
||||
stdout: &stdout,
|
||||
stderr: &stderr,
|
||||
format: format,
|
||||
runID: "run_budget",
|
||||
streamedText: &strings.Builder{},
|
||||
}
|
||||
writer.toolResult(agent.ToolResult{
|
||||
ToolCallID: "call_budget",
|
||||
Name: "read_file",
|
||||
Status: tools.StatusOK,
|
||||
Output: "bounded output",
|
||||
Truncated: true,
|
||||
})
|
||||
if writer.err != nil {
|
||||
t.Fatalf("toolResult: %v", writer.err)
|
||||
}
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal(bytes.TrimSpace(stdout.Bytes()), &payload); err != nil {
|
||||
t.Fatalf("decode output %q: %v", stdout.String(), err)
|
||||
}
|
||||
if payload["truncated"] != true {
|
||||
t.Fatalf("truncated = %#v, want true; payload=%#v", payload["truncated"], payload)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
+41
-6
@@ -23,6 +23,15 @@ type bashTool struct {
|
||||
scope PathScope
|
||||
}
|
||||
|
||||
func (bashTool) outputCategory(args map[string]any) outputCategory {
|
||||
command, _ := bashCommandArg(args)
|
||||
return shellOutputCategory(command)
|
||||
}
|
||||
|
||||
func bashCommandArg(args map[string]any) (string, error) {
|
||||
return aliasedStringArg(args, []string{"command", "cmd", "script", "shell"}, "", true, false)
|
||||
}
|
||||
|
||||
func NewBashTool(workspaceRoot string) Tool {
|
||||
return NewScopedBashTool(workspaceRoot, nil)
|
||||
}
|
||||
@@ -63,15 +72,19 @@ func NewScopedBashTool(workspaceRoot string, scope PathScope) Tool {
|
||||
}
|
||||
|
||||
func (tool bashTool) Run(ctx context.Context, args map[string]any) Result {
|
||||
return tool.run(ctx, args, nil)
|
||||
return tool.run(ctx, args, nil, true)
|
||||
}
|
||||
|
||||
func (tool bashTool) RunWithSandbox(ctx context.Context, args map[string]any, engine *zeroSandbox.Engine) Result {
|
||||
return tool.run(ctx, args, engine)
|
||||
return tool.run(ctx, args, engine, true)
|
||||
}
|
||||
|
||||
func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroSandbox.Engine) Result {
|
||||
commandText, err := aliasedStringArg(args, []string{"command", "cmd", "script", "shell"}, "", true, false)
|
||||
func (tool bashTool) RunWithOptions(ctx context.Context, args map[string]any, options RunOptions) Result {
|
||||
return tool.run(ctx, args, options.Sandbox, false)
|
||||
}
|
||||
|
||||
func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroSandbox.Engine, directBudget bool) Result {
|
||||
commandText, err := bashCommandArg(args)
|
||||
if err != nil {
|
||||
return errorResult("Error: Invalid arguments for bash: " + err.Error())
|
||||
}
|
||||
@@ -172,7 +185,7 @@ func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroS
|
||||
}
|
||||
}
|
||||
markLikelySandboxDenial(meta, plan, exitCode, stdoutText, stderrText)
|
||||
outText, errText, truncated := budgetBashCapture(stdoutText, stdout.total, stderrText, stderrTotal, meta)
|
||||
outText, errText, truncated := prepareBashOutput(stdoutText, stdout.total, stderrText, stderrTotal, meta, directBudget)
|
||||
return Result{
|
||||
Status: StatusError,
|
||||
Output: formatBashOutputWithShellHint(outText, errText, exitCode, meta),
|
||||
@@ -182,7 +195,7 @@ func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroS
|
||||
}
|
||||
|
||||
markLikelySandboxDenial(meta, plan, exitCode, stdoutText, stderrText)
|
||||
outText, errText, truncated := budgetBashCapture(stdoutText, stdout.total, stderrText, stderrTotal, meta)
|
||||
outText, errText, truncated := prepareBashOutput(stdoutText, stdout.total, stderrText, stderrTotal, meta, directBudget)
|
||||
if meta[SandboxLikelyDeniedMeta] == "true" {
|
||||
return Result{
|
||||
Status: StatusError,
|
||||
@@ -415,6 +428,28 @@ func budgetBashOutput(stdout string, stderr string, meta map[string]string) (str
|
||||
return budgetBashCapture(stdout, len(stdout), stderr, len(stderr), meta)
|
||||
}
|
||||
|
||||
// prepareBashOutput keeps direct callers on the established positional budget.
|
||||
// Registry calls retain the existing bounded capture but leave final semantic
|
||||
// reduction and spill creation to the post-redaction registry boundary.
|
||||
func prepareBashOutput(out string, outTotal int, errStr string, errTotal int, meta map[string]string, directBudget bool) (string, string, bool) {
|
||||
if directBudget {
|
||||
return budgetBashCapture(out, outTotal, errStr, errTotal, meta)
|
||||
}
|
||||
outText := sectionWithCaptureGap(out, outTotal)
|
||||
errText := sectionWithCaptureGap(errStr, errTotal)
|
||||
truncated := outTotal > len(out) || errTotal > len(errStr)
|
||||
if meta != nil {
|
||||
meta["raw_bytes"] = strconv.Itoa(outTotal + errTotal)
|
||||
meta["emitted_bytes"] = strconv.Itoa(len(outText) + len(errText))
|
||||
meta["estimated_tokens"] = strconv.Itoa(estimatedTokensFromBytes(len(outText) + len(errText)))
|
||||
if truncated {
|
||||
meta["truncated"] = "true"
|
||||
meta["truncation_reason"] = "capture_budget"
|
||||
}
|
||||
}
|
||||
return outText, errText, truncated
|
||||
}
|
||||
|
||||
// budgetBashCapture is budgetBashOutput for the streaming-capture path: outTotal
|
||||
// and errTotal are the true byte counts (from boundedBuffer.total), which may
|
||||
// exceed the retained strings when the middle was dropped during capture. Meta's
|
||||
|
||||
@@ -37,6 +37,18 @@ func TestBudgetBashOutputSmallPassesThrough(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestPrepareBashOutputDefersRegistrySemanticBudget(t *testing.T) {
|
||||
output := "START\n" + strings.Repeat("progress\n", 8_000) + "ERROR: retained\nEND"
|
||||
registryOutput, _, registryTruncated := prepareBashOutput(output, len(output), "", 0, map[string]string{}, false)
|
||||
if registryTruncated || registryOutput != output {
|
||||
t.Fatalf("registry preparation must preserve capture-bounded output for semantic reduction: truncated=%t bytes=%d", registryTruncated, len(registryOutput))
|
||||
}
|
||||
directOutput, _, directTruncated := prepareBashOutput(output, len(output), "", 0, map[string]string{}, true)
|
||||
if !directTruncated || len(directOutput) > bashOutputBudgetBytes+512 {
|
||||
t.Fatalf("direct path must preserve legacy positional budget: truncated=%t bytes=%d", directTruncated, len(directOutput))
|
||||
}
|
||||
}
|
||||
|
||||
// Oversized stdout is truncated head+tail: both the first and last lines survive,
|
||||
// the middle is dropped behind a marker, meta is flagged, and the captured text
|
||||
// is spilled to a re-readable file.
|
||||
|
||||
@@ -427,6 +427,15 @@ type execCommandTool struct {
|
||||
manager *execSessionManager
|
||||
}
|
||||
|
||||
func (execCommandTool) outputCategory(args map[string]any) outputCategory {
|
||||
command, _ := execCommandArg(args)
|
||||
return shellOutputCategory(command)
|
||||
}
|
||||
|
||||
func execCommandArg(args map[string]any) (string, error) {
|
||||
return aliasedStringArg(args, []string{"cmd", "command", "script", "shell"}, "", true, false)
|
||||
}
|
||||
|
||||
func NewExecCommandTool(workspaceRoot string, manager *execSessionManager) Tool {
|
||||
return NewScopedExecCommandTool(workspaceRoot, nil, manager)
|
||||
}
|
||||
@@ -475,11 +484,15 @@ func NewScopedExecCommandTool(workspaceRoot string, scope PathScope, manager *ex
|
||||
}
|
||||
|
||||
func (tool execCommandTool) Run(ctx context.Context, args map[string]any) Result {
|
||||
return tool.run(ctx, args, nil)
|
||||
return tool.run(ctx, args, nil, true)
|
||||
}
|
||||
|
||||
func (tool execCommandTool) RunWithSandbox(ctx context.Context, args map[string]any, engine *zeroSandbox.Engine) Result {
|
||||
return tool.run(ctx, args, engine)
|
||||
return tool.run(ctx, args, engine, true)
|
||||
}
|
||||
|
||||
func (tool execCommandTool) RunWithOptions(ctx context.Context, args map[string]any, options RunOptions) Result {
|
||||
return tool.run(ctx, args, options.Sandbox, false)
|
||||
}
|
||||
|
||||
func (tool execCommandTool) ExecSessions() []ExecSessionSnapshot {
|
||||
@@ -494,8 +507,8 @@ 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)
|
||||
func (tool execCommandTool) run(ctx context.Context, args map[string]any, engine *zeroSandbox.Engine, directBudget bool) Result {
|
||||
commandText, err := execCommandArg(args)
|
||||
if err != nil {
|
||||
return errorResult("Error: Invalid arguments for exec_command: " + err.Error())
|
||||
}
|
||||
@@ -550,7 +563,7 @@ func (tool execCommandTool) run(ctx context.Context, args map[string]any, engine
|
||||
if exited {
|
||||
tool.manager.remove(session.id)
|
||||
}
|
||||
return execToolResult(execToolResultInput{
|
||||
return execToolResultWithBudget(execToolResultInput{
|
||||
commandText: commandText,
|
||||
output: output,
|
||||
outputBufferTruncated: outputTruncated,
|
||||
@@ -561,7 +574,7 @@ func (tool execCommandTool) run(ctx context.Context, args map[string]any, engine
|
||||
tty: session.tty,
|
||||
plan: session.plan,
|
||||
maxOutputTokens: maxOutputTokens,
|
||||
})
|
||||
}, directBudget)
|
||||
}
|
||||
|
||||
func (tool execCommandTool) startSession(commandText string, absoluteCwd string, relativeCwd string, ttyRequested bool, engine *zeroSandbox.Engine, sandboxPermissions SandboxPermissionOverride) (*execSession, error) {
|
||||
@@ -694,6 +707,10 @@ type writeStdinTool struct {
|
||||
manager *execSessionManager
|
||||
}
|
||||
|
||||
func (writeStdinTool) outputCategory(map[string]any) outputCategory {
|
||||
return outputCategoryProcess
|
||||
}
|
||||
|
||||
func NewWriteStdinTool(manager *execSessionManager) Tool {
|
||||
if manager == nil {
|
||||
manager = defaultExecSessionManager
|
||||
@@ -843,7 +860,15 @@ type execToolResultInput struct {
|
||||
}
|
||||
|
||||
func execToolResult(input execToolResultInput) Result {
|
||||
output, truncated := truncateExecOutput(input.output, input.maxOutputTokens)
|
||||
return execToolResultWithBudget(input, true)
|
||||
}
|
||||
|
||||
func execToolResultWithBudget(input execToolResultInput, directBudget bool) Result {
|
||||
output := input.output
|
||||
truncated := false
|
||||
if directBudget {
|
||||
output, truncated = truncateExecOutput(input.output, input.maxOutputTokens)
|
||||
}
|
||||
meta := map[string]string{
|
||||
"cwd": input.relativeCwd,
|
||||
"tty": strconv.FormatBool(input.tty),
|
||||
|
||||
+23
-14
@@ -18,6 +18,8 @@ type globTool struct {
|
||||
scope PathScope
|
||||
}
|
||||
|
||||
func (globTool) outputCategory(map[string]any) outputCategory { return outputCategorySearch }
|
||||
|
||||
func NewGlobTool(workspaceRoot string) Tool {
|
||||
return NewScopedGlobTool(workspaceRoot, nil)
|
||||
}
|
||||
@@ -47,17 +49,25 @@ func NewScopedGlobTool(workspaceRoot string, scope PathScope) Tool {
|
||||
}
|
||||
|
||||
func (tool globTool) Run(ctx context.Context, args map[string]any) Result {
|
||||
return tool.runWith(ctx, args, readExcluder{})
|
||||
return tool.runWith(ctx, args, readExcluder{}, true)
|
||||
}
|
||||
|
||||
func (tool globTool) RunWithOptions(ctx context.Context, args map[string]any, options RunOptions) Result {
|
||||
exclude := readExcluder{}
|
||||
if options.Sandbox != nil {
|
||||
exclude = sandboxReadExcluder(options.Sandbox)
|
||||
}
|
||||
return tool.runWith(ctx, args, exclude, false)
|
||||
}
|
||||
|
||||
// RunWithSandbox runs glob while skipping subtrees the sandbox policy denies
|
||||
// reads to (DenyRead). With no DenyRead configured the excluder is a no-op and
|
||||
// behavior is unchanged.
|
||||
func (tool globTool) RunWithSandbox(ctx context.Context, args map[string]any, engine *sandbox.Engine) Result {
|
||||
return tool.runWith(ctx, args, sandboxReadExcluder(engine))
|
||||
return tool.runWith(ctx, args, sandboxReadExcluder(engine), true)
|
||||
}
|
||||
|
||||
func (tool globTool) runWith(ctx context.Context, args map[string]any, exclude readExcluder) Result {
|
||||
func (tool globTool) runWith(ctx context.Context, args map[string]any, exclude readExcluder, directBudget bool) Result {
|
||||
pattern, err := aliasedStringArg(args, []string{"pattern", "glob", "match", "query", "expression"}, "", true, false)
|
||||
if err != nil {
|
||||
return errorResult("Error: Invalid arguments for glob: " + err.Error())
|
||||
@@ -110,24 +120,23 @@ func (tool globTool) runWith(ctx context.Context, args map[string]any, exclude r
|
||||
if truncated {
|
||||
output += fmt.Sprintf("\n\n[truncated: showing first %d of %d matches; increase limit or narrow cwd/pattern]", len(matches), totalMatches)
|
||||
}
|
||||
budgeted := applyOutputBudget(output, searchOutputBudgetBytes, "increase limit or narrow cwd/pattern")
|
||||
meta := outputBudgetMeta(budgeted)
|
||||
meta := map[string]string{}
|
||||
meta["pattern"] = pattern
|
||||
if truncated || budgeted.Truncated {
|
||||
if truncated {
|
||||
meta["truncated"] = "true"
|
||||
if budgeted.Truncated {
|
||||
meta["truncation_reason"] = "byte_budget"
|
||||
} else {
|
||||
meta["truncation_reason"] = "limit"
|
||||
}
|
||||
meta["truncation_reason"] = "limit"
|
||||
}
|
||||
|
||||
return Result{
|
||||
result := Result{
|
||||
Status: StatusOK,
|
||||
Output: budgeted.Output,
|
||||
Truncated: truncated || budgeted.Truncated,
|
||||
Output: output,
|
||||
Truncated: truncated,
|
||||
Meta: meta,
|
||||
}
|
||||
if directBudget {
|
||||
return applyLegacyByteBudgetToResult(result, searchOutputBudgetBytes, "increase limit or narrow cwd/pattern")
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func scanGlob(ctx context.Context, root string, displayRoot string, matcher *regexp.Regexp, includeDirs bool, exclude readExcluder) ([]string, error) {
|
||||
|
||||
+29
-23
@@ -21,6 +21,8 @@ type grepTool struct {
|
||||
scope PathScope
|
||||
}
|
||||
|
||||
func (grepTool) outputCategory(map[string]any) outputCategory { return outputCategorySearch }
|
||||
|
||||
type grepMatch struct {
|
||||
file string
|
||||
line int
|
||||
@@ -62,7 +64,15 @@ func NewScopedGrepTool(workspaceRoot string, scope PathScope) Tool {
|
||||
}
|
||||
|
||||
func (tool grepTool) Run(ctx context.Context, args map[string]any) Result {
|
||||
return tool.runWith(ctx, args, readExcluder{})
|
||||
return tool.runWith(ctx, args, readExcluder{}, true)
|
||||
}
|
||||
|
||||
func (tool grepTool) RunWithOptions(ctx context.Context, args map[string]any, options RunOptions) Result {
|
||||
exclude := readExcluder{}
|
||||
if options.Sandbox != nil {
|
||||
exclude = sandboxReadExcluder(options.Sandbox)
|
||||
}
|
||||
return tool.runWith(ctx, args, exclude, false)
|
||||
}
|
||||
|
||||
// RunWithSandbox runs the search while skipping subtrees the sandbox policy
|
||||
@@ -70,10 +80,10 @@ func (tool grepTool) Run(ctx context.Context, args map[string]any) Result {
|
||||
// path. With no DenyRead configured the excluder is a no-op and behavior is
|
||||
// unchanged.
|
||||
func (tool grepTool) RunWithSandbox(ctx context.Context, args map[string]any, engine *sandbox.Engine) Result {
|
||||
return tool.runWith(ctx, args, sandboxReadExcluder(engine))
|
||||
return tool.runWith(ctx, args, sandboxReadExcluder(engine), true)
|
||||
}
|
||||
|
||||
func (tool grepTool) runWith(ctx context.Context, args map[string]any, exclude readExcluder) Result {
|
||||
func (tool grepTool) runWith(ctx context.Context, args map[string]any, exclude readExcluder, directBudget bool) Result {
|
||||
pattern, err := aliasedStringArg(args, []string{"pattern", "query", "regex", "search", "expression"}, "", true, false)
|
||||
if err != nil {
|
||||
return errorResult("Error: Invalid arguments for grep: " + err.Error())
|
||||
@@ -154,7 +164,7 @@ func (tool grepTool) runWith(ctx context.Context, args map[string]any, exclude r
|
||||
}
|
||||
return errorResult("Error running grep: " + err.Error())
|
||||
}
|
||||
return collector.result()
|
||||
return applyDirectSearchBudget(collector.result(), directBudget, "narrow path/glob/pattern to continue")
|
||||
case "files_with_matches":
|
||||
collector := &grepFileListCollector{}
|
||||
if err := scanGrepMatches(ctx, resolvedRoot, target, globMatcher, exclude, absolutePaths, presenceGrepLineMatcher(compiled), collector.collect); err != nil {
|
||||
@@ -163,7 +173,7 @@ func (tool grepTool) runWith(ctx context.Context, args map[string]any, exclude r
|
||||
}
|
||||
return errorResult("Error running grep: " + err.Error())
|
||||
}
|
||||
return collector.result()
|
||||
return applyDirectSearchBudget(collector.result(), directBudget, "narrow path/glob/pattern to continue")
|
||||
default:
|
||||
collector := &grepContentCollector{headLimit: headLimit}
|
||||
if err := scanGrepMatches(ctx, resolvedRoot, target, globMatcher, exclude, absolutePaths, presenceGrepLineMatcher(compiled), collector.collect); err != nil {
|
||||
@@ -172,10 +182,17 @@ func (tool grepTool) runWith(ctx context.Context, args map[string]any, exclude r
|
||||
}
|
||||
return errorResult("Error running grep: " + err.Error())
|
||||
}
|
||||
return collector.result()
|
||||
return applyDirectSearchBudget(collector.result(), directBudget, "narrow path/glob/pattern or increase head_limit")
|
||||
}
|
||||
}
|
||||
|
||||
func applyDirectSearchBudget(result Result, directBudget bool, hint string) Result {
|
||||
if !directBudget {
|
||||
return result
|
||||
}
|
||||
return applyLegacyByteBudgetToResult(result, searchOutputBudgetBytes, hint)
|
||||
}
|
||||
|
||||
// resolveGrepRoot picks the scope root whose EvalSymlinks-resolved path contains
|
||||
// the already-resolved target, so that confineGrepFile computes correct
|
||||
// workspace-relative paths even when the target lives in an extra root.
|
||||
@@ -455,13 +472,7 @@ func (collector *grepFileListCollector) result() Result {
|
||||
return okResult("No matches found.")
|
||||
}
|
||||
sort.Strings(collector.files)
|
||||
budgeted := applyOutputBudget(strings.Join(collector.files, "\n"), searchOutputBudgetBytes, "narrow path/glob/pattern to continue")
|
||||
meta := outputBudgetMeta(budgeted)
|
||||
if budgeted.Truncated {
|
||||
meta["truncated"] = "true"
|
||||
meta["truncation_reason"] = "byte_budget"
|
||||
}
|
||||
return Result{Status: StatusOK, Output: budgeted.Output, Truncated: budgeted.Truncated, Meta: meta}
|
||||
return Result{Status: StatusOK, Output: strings.Join(collector.files, "\n")}
|
||||
}
|
||||
|
||||
type grepContentCollector struct {
|
||||
@@ -494,20 +505,15 @@ func (collector *grepContentCollector) result() Result {
|
||||
if truncated {
|
||||
output += fmt.Sprintf("\n\n[truncated: showing first %d matches; narrow path/glob/pattern or increase head_limit]", len(lines))
|
||||
}
|
||||
budgeted := applyOutputBudget(output, searchOutputBudgetBytes, "narrow path/glob/pattern or increase head_limit")
|
||||
meta := outputBudgetMeta(budgeted)
|
||||
if truncated || budgeted.Truncated {
|
||||
meta := map[string]string{}
|
||||
if truncated {
|
||||
meta["truncated"] = "true"
|
||||
if budgeted.Truncated {
|
||||
meta["truncation_reason"] = "byte_budget"
|
||||
} else {
|
||||
meta["truncation_reason"] = "head_limit"
|
||||
}
|
||||
meta["truncation_reason"] = "head_limit"
|
||||
}
|
||||
return Result{
|
||||
Status: StatusOK,
|
||||
Output: budgeted.Output,
|
||||
Truncated: truncated || budgeted.Truncated,
|
||||
Output: output,
|
||||
Truncated: truncated,
|
||||
Meta: meta,
|
||||
}
|
||||
}
|
||||
@@ -15,6 +15,10 @@ type listDirectoryTool struct {
|
||||
scope PathScope
|
||||
}
|
||||
|
||||
func (listDirectoryTool) outputCategory(map[string]any) outputCategory {
|
||||
return outputCategorySearch
|
||||
}
|
||||
|
||||
func NewListDirectoryTool(workspaceRoot string) Tool {
|
||||
return NewScopedListDirectoryTool(workspaceRoot, nil)
|
||||
}
|
||||
@@ -42,6 +46,14 @@ func NewScopedListDirectoryTool(workspaceRoot string, scope PathScope) Tool {
|
||||
}
|
||||
|
||||
func (tool listDirectoryTool) Run(_ context.Context, args map[string]any) Result {
|
||||
return tool.run(args, true)
|
||||
}
|
||||
|
||||
func (tool listDirectoryTool) RunWithOptions(_ context.Context, args map[string]any, _ RunOptions) Result {
|
||||
return tool.run(args, false)
|
||||
}
|
||||
|
||||
func (tool listDirectoryTool) run(args map[string]any, directBudget bool) Result {
|
||||
// Optional with a "." default: treat an explicit empty path (a common
|
||||
// weak-model quirk) the same as the key being absent rather than erroring.
|
||||
requestedPath, err := aliasedStringArg(args, []string{"path", "directory", "dir"}, ".", false, true)
|
||||
@@ -76,13 +88,11 @@ func (tool listDirectoryTool) Run(_ context.Context, args map[string]any) Result
|
||||
return okResult("Directory is empty: " + relativePath)
|
||||
}
|
||||
output := "Contents of " + relativePath + ":\n\n" + strings.Join(entries, "\n")
|
||||
budgeted := applyOutputBudget(output, searchOutputBudgetBytes, "use path, recursive=false, or a smaller max_depth to narrow the listing")
|
||||
meta := outputBudgetMeta(budgeted)
|
||||
if budgeted.Truncated {
|
||||
meta["truncated"] = "true"
|
||||
meta["truncation_reason"] = "byte_budget"
|
||||
result := Result{Status: StatusOK, Output: output}
|
||||
if directBudget {
|
||||
return applyLegacyByteBudgetToResult(result, searchOutputBudgetBytes, "use path, recursive=false, or a smaller max_depth to narrow the listing")
|
||||
}
|
||||
return Result{Status: StatusOK, Output: budgeted.Output, Truncated: budgeted.Truncated, Meta: meta}
|
||||
return result
|
||||
}
|
||||
|
||||
func listDirectoryEntries(path string, depth int, maxDepth int) ([]string, error) {
|
||||
|
||||
@@ -0,0 +1,254 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
const (
|
||||
outputBudgetCategoryMeta = "output_budget_category"
|
||||
outputBudgetOriginalBytesMeta = "output_budget_original_bytes"
|
||||
outputBudgetRetainedBytesMeta = "output_budget_retained_bytes"
|
||||
outputBudgetEstimatedOriginalTokensMeta = "output_budget_estimated_original_tokens"
|
||||
outputBudgetEstimatedRetainedTokensMeta = "output_budget_estimated_retained_tokens"
|
||||
outputBudgetReasonMeta = "output_budget_reason"
|
||||
outputBudgetSpillCreatedMeta = "output_budget_spill_created"
|
||||
)
|
||||
|
||||
// applyRegistryOutputBudget is the common post-redaction semantic budgeting
|
||||
// boundary for tools that do not already own a deliberate output budget.
|
||||
func applyRegistryOutputBudget(tool Tool, toolName string, args map[string]any, result Result) Result {
|
||||
budget := registryOutputBudget(toolName)
|
||||
if budget.maxEstimatedTokens <= 0 && budget.hardMaxBytes <= 0 {
|
||||
return result // preserve ZERO_TOOL_OUTPUT_CEILING_TOKENS=0 semantics
|
||||
}
|
||||
|
||||
category := resolveOutputCategory(tool, toolName, args)
|
||||
budgeted := budgetSemanticOutput(result.Output, category, budget)
|
||||
if budgeted.truncated {
|
||||
budgeted = attachExistingSpill(toolName, result.Output, budget, budgeted)
|
||||
}
|
||||
if result.Truncated {
|
||||
budgeted.truncated = true
|
||||
if budgeted.reason == "" {
|
||||
budgeted.reason = result.Meta["truncation_reason"]
|
||||
if budgeted.reason == "" {
|
||||
budgeted.reason = "upstream_tool_budget"
|
||||
}
|
||||
}
|
||||
if budgeted.spillPath == "" {
|
||||
budgeted.spillPath = result.Meta["spill_path"]
|
||||
}
|
||||
}
|
||||
result.Output = budgeted.text
|
||||
result.Truncated = result.Truncated || budgeted.truncated
|
||||
result.Meta = addOutputBudgetMetadata(result.Meta, budgeted)
|
||||
return result
|
||||
}
|
||||
|
||||
// applySelfManagedOutputBudget applies semantic retention to tools that retain
|
||||
// output with their own capture-aware limits. Their explicit budgets remain
|
||||
// authoritative, rather than being replaced by the registry default ceiling.
|
||||
func applySelfManagedOutputBudget(tool Tool, toolName string, args map[string]any, result Result) Result {
|
||||
budget := selfManagedOutputBudget(toolName, args)
|
||||
if budget.maxEstimatedTokens <= 0 && budget.hardMaxBytes <= 0 {
|
||||
return annotateSelfBudgetedOutput(tool, toolName, args, result)
|
||||
}
|
||||
|
||||
category := resolveOutputCategory(tool, toolName, args)
|
||||
budgeted := budgetSemanticOutput(result.Output, category, budget)
|
||||
if result.Truncated && !budgeted.truncated {
|
||||
budgeted.truncated = true
|
||||
budgeted.reason = result.Meta["truncation_reason"]
|
||||
if budgeted.reason == "" {
|
||||
budgeted.reason = "upstream_tool_budget"
|
||||
}
|
||||
}
|
||||
if budgeted.truncated {
|
||||
budgeted = attachExistingSpill(toolName, result.Output, budget, budgeted)
|
||||
}
|
||||
result.Output = budgeted.text
|
||||
result.Truncated = result.Truncated || budgeted.truncated
|
||||
result.Meta = addOutputBudgetMetadata(result.Meta, budgeted)
|
||||
return result
|
||||
}
|
||||
|
||||
func selfManagedOutputBudget(toolName string, args map[string]any) outputBudget {
|
||||
switch toolName {
|
||||
case "bash":
|
||||
// bash previously exposed at most 32 KiB from each stream. Keep that
|
||||
// combined 64 KiB ceiling while selecting its contents semantically.
|
||||
return outputBudget{maxEstimatedTokens: (bashOutputBudgetBytes * 2) / 4, hardMaxBytes: bashOutputBudgetBytes * 2}
|
||||
case ExecCommandToolName:
|
||||
maxOutputTokens := defaultMaxOutputTokens
|
||||
if parsed, err := intArg(args, "max_output_tokens", defaultMaxOutputTokens, 1, maxExecOutputTokenRequest); err == nil {
|
||||
maxOutputTokens = parsed
|
||||
}
|
||||
return outputBudget{maxEstimatedTokens: maxOutputTokens, hardMaxBytes: maxOutputTokens * 4}
|
||||
default:
|
||||
return outputBudget{}
|
||||
}
|
||||
}
|
||||
|
||||
// RebudgetAfterHook reapplies the registry's redaction and output limits after
|
||||
// an afterTool hook appends model-visible feedback to a completed result. The
|
||||
// initial registry pass has already run; this second pass is limited to the
|
||||
// newly combined result so hooks cannot bypass the established safety ceiling.
|
||||
func (registry *Registry) RebudgetAfterHook(toolName string, args map[string]any, result Result) Result {
|
||||
result = scrubResultSecrets(result)
|
||||
tool, _ := registry.Get(toolName)
|
||||
if _, ok := tool.(selfBudgeting); ok {
|
||||
// Match the primary registry boundary: self-managed tools keep their
|
||||
// call-specific capture/output budget instead of being tightened or
|
||||
// loosened to the generic registry ceiling after hook feedback.
|
||||
return applySelfManagedOutputBudget(tool, toolName, args, result)
|
||||
}
|
||||
result = applyRegistryOutputBudget(tool, toolName, args, result)
|
||||
return enforceOutputCeiling(toolName, result)
|
||||
}
|
||||
|
||||
func registryOutputBudget(toolName string) outputBudget {
|
||||
switch toolName {
|
||||
case "read_file", "read_minified_file":
|
||||
return outputBudget{maxEstimatedTokens: readOutputBudgetBytes / 4, hardMaxBytes: readOutputBudgetBytes}
|
||||
case "grep", "glob", "list_directory":
|
||||
return outputBudget{maxEstimatedTokens: searchOutputBudgetBytes / 4, hardMaxBytes: searchOutputBudgetBytes}
|
||||
default:
|
||||
ceilingTokens := resolveOutputCeilingTokens()
|
||||
if ceilingTokens <= 0 {
|
||||
return outputBudget{}
|
||||
}
|
||||
return outputBudget{maxEstimatedTokens: ceilingTokens, hardMaxBytes: ceilingTokens * 4}
|
||||
}
|
||||
}
|
||||
|
||||
func resolveOutputCategory(tool Tool, toolName string, args map[string]any) outputCategory {
|
||||
if provider, ok := tool.(outputPolicyProvider); ok {
|
||||
if category := provider.outputCategory(args); category != "" {
|
||||
return category
|
||||
}
|
||||
}
|
||||
switch toolName {
|
||||
case "Task", "swarm_collect":
|
||||
return outputCategoryWorker
|
||||
case "apply_patch":
|
||||
return outputCategoryDiff
|
||||
case "write_stdin":
|
||||
return outputCategoryProcess
|
||||
default:
|
||||
return outputCategoryDefault
|
||||
}
|
||||
}
|
||||
|
||||
// annotateSelfBudgetedOutput records the same compact metadata for tools whose
|
||||
// existing capture/budget implementation remains authoritative in PR11. It
|
||||
// does not re-budget their text or claim that raw_bytes represents every byte a
|
||||
// subprocess produced beyond its established capture limits.
|
||||
func annotateSelfBudgetedOutput(tool Tool, toolName string, args map[string]any, result Result) Result {
|
||||
retainedBytes := len(result.Output)
|
||||
originalBytes := retainedBytes
|
||||
if parsed, err := strconv.Atoi(result.Meta["raw_bytes"]); err == nil && parsed > originalBytes {
|
||||
originalBytes = parsed
|
||||
}
|
||||
retainedTokens := estimateOutputTokens(result.Output)
|
||||
originalTokens := retainedTokens
|
||||
if originalBytes > retainedBytes {
|
||||
originalTokens = max(originalTokens, estimatedTokensFromBytes(originalBytes))
|
||||
}
|
||||
reason := result.Meta["truncation_reason"]
|
||||
if result.Truncated && reason == "" {
|
||||
reason = "upstream_tool_budget"
|
||||
}
|
||||
observed := budgetedOutput{
|
||||
text: result.Output,
|
||||
originalBytes: originalBytes,
|
||||
retainedBytes: retainedBytes,
|
||||
estimatedOriginalTokens: originalTokens,
|
||||
estimatedRetainedTokens: retainedTokens,
|
||||
truncated: result.Truncated,
|
||||
category: resolveOutputCategory(tool, toolName, args),
|
||||
reason: reason,
|
||||
spillPath: result.Meta["spill_path"],
|
||||
}
|
||||
result.Meta = addOutputBudgetMetadata(result.Meta, observed)
|
||||
return result
|
||||
}
|
||||
|
||||
func shellOutputCategory(command string) outputCategory {
|
||||
normalized := strings.ToLower(strings.TrimSpace(command))
|
||||
if containsAny(normalized,
|
||||
"go test", "pytest", "python -m pytest", "cargo test", "npm test", "npm run test",
|
||||
"pnpm test", "yarn test", "bun test", "dotnet test", "mvn test", "gradle test", "phpunit") {
|
||||
return outputCategoryTest
|
||||
}
|
||||
if containsAny(normalized, "git diff", "git show", "git format-patch", " diff -u", "diff --git") || strings.HasPrefix(normalized, "diff ") {
|
||||
return outputCategoryDiff
|
||||
}
|
||||
return outputCategoryProcess
|
||||
}
|
||||
|
||||
// attachExistingSpill reuses Zero's hardened per-user spill directory. output
|
||||
// is the already-redacted text received by this layer; it may itself be a
|
||||
// capture-bounded view produced by a subprocess tool.
|
||||
func attachExistingSpill(toolName, output string, budget outputBudget, current budgetedOutput) budgetedOutput {
|
||||
path := spillTruncatedOutput(toolName, output)
|
||||
if path == "" {
|
||||
return current
|
||||
}
|
||||
notice := "[zero] full output received by budgeting layer saved to " + path + " (grep or read_file it instead of re-running)"
|
||||
reduced := outputBudget{
|
||||
maxEstimatedTokens: max(1, budget.maxEstimatedTokens-estimateOutputTokens("\n"+notice)),
|
||||
hardMaxBytes: max(1, budget.hardMaxBytes-len("\n"+notice)),
|
||||
}
|
||||
base := budgetSemanticOutput(output, current.category, reduced)
|
||||
text := strings.TrimRight(base.text, "\n") + "\n" + notice
|
||||
if !fitsOutputBudget(text, budget) {
|
||||
// An unusually long temp path or tiny configured ceiling can leave no
|
||||
// room for the full notice. Keep the safe bounded result; the spill still
|
||||
// exists but is intentionally not advertised with a chopped reference.
|
||||
current.spillPath = path
|
||||
return current
|
||||
}
|
||||
base.text = text
|
||||
base.retainedBytes = len(text)
|
||||
base.estimatedRetainedTokens = estimateOutputTokens(text)
|
||||
base.truncated = true
|
||||
if base.reason == "" {
|
||||
base.reason = current.reason
|
||||
}
|
||||
base.category = current.category
|
||||
base.spillPath = path
|
||||
return base
|
||||
}
|
||||
|
||||
func addOutputBudgetMetadata(meta map[string]string, output budgetedOutput) map[string]string {
|
||||
if meta == nil {
|
||||
meta = map[string]string{}
|
||||
}
|
||||
meta[outputBudgetCategoryMeta] = string(output.category)
|
||||
meta[outputBudgetOriginalBytesMeta] = strconv.Itoa(output.originalBytes)
|
||||
meta[outputBudgetRetainedBytesMeta] = strconv.Itoa(output.retainedBytes)
|
||||
meta[outputBudgetEstimatedOriginalTokensMeta] = strconv.Itoa(output.estimatedOriginalTokens)
|
||||
meta[outputBudgetEstimatedRetainedTokensMeta] = strconv.Itoa(output.estimatedRetainedTokens)
|
||||
meta[outputBudgetSpillCreatedMeta] = strconv.FormatBool(output.spillPath != "")
|
||||
// These established fields describe the final model-visible output. Refresh
|
||||
// them on every pass because redaction and after-tool hooks may change the
|
||||
// text even when this pass does not newly truncate it.
|
||||
meta["emitted_bytes"] = strconv.Itoa(output.retainedBytes)
|
||||
meta["estimated_tokens"] = strconv.Itoa(output.estimatedRetainedTokens)
|
||||
if output.reason != "" {
|
||||
meta[outputBudgetReasonMeta] = output.reason
|
||||
}
|
||||
if output.truncated {
|
||||
// Preserve the existing metadata vocabulary used by callers and tests.
|
||||
if _, exists := meta["raw_bytes"]; !exists {
|
||||
meta["raw_bytes"] = strconv.Itoa(output.originalBytes)
|
||||
}
|
||||
meta["truncated"] = "true"
|
||||
meta["truncation_reason"] = output.reason
|
||||
}
|
||||
if output.spillPath != "" {
|
||||
meta["spill_path"] = output.spillPath
|
||||
}
|
||||
return meta
|
||||
}
|
||||
@@ -0,0 +1,381 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
func TestRegistryBudgetRunsAfterRedactionAndSpillsRedactedOutput(t *testing.T) {
|
||||
setTestTempDir(t)
|
||||
secret := "ghp_" + strings.Repeat("a", 36)
|
||||
big := "HEAD\n" + secret + "\n" + strings.Repeat("🙂 noisy output\n", 20_000) + "TAIL"
|
||||
registry := NewRegistry()
|
||||
registry.Register(newCeilingFakeTool("redacted_big", big))
|
||||
|
||||
result := registry.Run(context.Background(), "redacted_big", map[string]any{})
|
||||
if !result.Truncated || result.Meta[outputBudgetSpillCreatedMeta] != "true" {
|
||||
t.Fatalf("unexpected budget result: truncated=%t meta=%#v", result.Truncated, result.Meta)
|
||||
}
|
||||
if strings.Contains(result.Output, secret) || !utf8.ValidString(result.Output) {
|
||||
t.Fatalf("exposed output leaked secret or invalid UTF-8: %q", result.Output)
|
||||
}
|
||||
spillPath := result.Meta["spill_path"]
|
||||
content, err := os.ReadFile(spillPath)
|
||||
if err != nil {
|
||||
t.Fatalf("read spill: %v", err)
|
||||
}
|
||||
if strings.Contains(string(content), secret) {
|
||||
t.Fatal("spill contains unredacted secret")
|
||||
}
|
||||
if !strings.Contains(string(content), "HEAD") || !strings.Contains(string(content), "TAIL") {
|
||||
t.Fatal("spill does not contain the complete redacted output received by the budget layer")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistryBudgetSpillFailureFallsBackToBoundedOutput(t *testing.T) {
|
||||
temp := t.TempDir()
|
||||
blockedTemp := filepath.Join(temp, "not-a-directory")
|
||||
if err := os.WriteFile(blockedTemp, []byte("x"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Setenv("TMPDIR", blockedTemp)
|
||||
t.Setenv("TMP", blockedTemp)
|
||||
t.Setenv("TEMP", blockedTemp)
|
||||
t.Setenv(outputCeilingEnv, "100")
|
||||
|
||||
registry := NewRegistry()
|
||||
registry.Register(newCeilingFakeTool("spill_failure", strings.Repeat("large output\n", 1000)))
|
||||
result := registry.Run(context.Background(), "spill_failure", map[string]any{})
|
||||
if !result.Truncated || len(result.Output) > 400 {
|
||||
t.Fatalf("fallback is not bounded: truncated=%t bytes=%d", result.Truncated, len(result.Output))
|
||||
}
|
||||
if result.Meta[outputBudgetSpillCreatedMeta] != "false" || result.Meta["spill_path"] != "" {
|
||||
t.Fatalf("spill failure incorrectly advertised: %#v", result.Meta)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistryMissingPolicyUsesDefaultAndExactHardCeiling(t *testing.T) {
|
||||
setTestTempDir(t)
|
||||
t.Setenv(outputCeilingEnv, "80")
|
||||
input := "HEAD\n" + strings.Repeat("x", 3000) + "\nTAIL"
|
||||
registry := NewRegistry()
|
||||
registry.Register(newCeilingFakeTool("no_policy", input))
|
||||
result := registry.Run(context.Background(), "no_policy", map[string]any{})
|
||||
if got := result.Meta[outputBudgetCategoryMeta]; got != string(outputCategoryDefault) {
|
||||
t.Fatalf("category = %q, want default", got)
|
||||
}
|
||||
if len(result.Output) > 80*4 {
|
||||
t.Fatalf("output = %d bytes, hard ceiling %d", len(result.Output), 80*4)
|
||||
}
|
||||
if result.Meta[outputBudgetOriginalBytesMeta] != strconv.Itoa(len(input)) {
|
||||
t.Fatalf("original size metadata = %q", result.Meta[outputBudgetOriginalBytesMeta])
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistrySmallOutputRemainsByteIdentical(t *testing.T) {
|
||||
registry := NewRegistry()
|
||||
registry.Register(newCeilingFakeTool("small_identity", "hello\nworld\n"))
|
||||
result := registry.Run(context.Background(), "small_identity", map[string]any{})
|
||||
if result.Output != "hello\nworld\n" || result.Truncated {
|
||||
t.Fatalf("small output changed: %#v", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDirectReadAndSearchToolsKeepLegacyByteBudgets(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
if err := os.WriteFile(filepath.Join(root, "large-read.txt"), []byte(strings.Repeat("source line with enough content\n", 7000)), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result := NewReadFileTool(root).Run(context.Background(), map[string]any{"path": "large-read.txt"}); !result.Truncated || len(result.Output) > readOutputBudgetBytes {
|
||||
t.Fatalf("direct read_file output is not bounded: %#v", result)
|
||||
}
|
||||
|
||||
for index := 0; index < 800; index++ {
|
||||
name := fmt.Sprintf("%04d-%s.txt", index, strings.Repeat("n", 80))
|
||||
if err := os.WriteFile(filepath.Join(root, name), []byte("match\n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(root, "large-content.txt"), []byte(strings.Repeat("match "+strings.Repeat("x", 2000)+"\n", 50)), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
for name, result := range map[string]Result{
|
||||
"list_directory": NewListDirectoryTool(root).Run(context.Background(), map[string]any{"recursive": false}),
|
||||
"glob": NewGlobTool(root).Run(context.Background(), map[string]any{"pattern": "*.txt", "limit": 1000}),
|
||||
"grep files": NewGrepTool(root).Run(context.Background(), map[string]any{"pattern": "match", "output_mode": "files_with_matches"}),
|
||||
"grep content": NewGrepTool(root).Run(context.Background(), map[string]any{"pattern": "match", "path": "large-content.txt", "head_limit": 50}),
|
||||
} {
|
||||
if !result.Truncated || len(result.Output) > searchOutputBudgetBytes {
|
||||
t.Fatalf("direct %s output is not bounded: truncated=%t bytes=%d meta=%#v", name, result.Truncated, len(result.Output), result.Meta)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistryLargeFileUsesSemanticFilePolicy(t *testing.T) {
|
||||
setTestTempDir(t)
|
||||
root := t.TempDir()
|
||||
var content strings.Builder
|
||||
content.WriteString("HEAD_MARK\n")
|
||||
for line := 0; line < 7000; line++ {
|
||||
content.WriteString("ordinary source line with enough text to fill the output budget\n")
|
||||
}
|
||||
content.WriteString("TAIL_MARK\n")
|
||||
if err := os.WriteFile(filepath.Join(root, "large.txt"), []byte(content.String()), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
registry := NewRegistry()
|
||||
registry.Register(NewReadFileTool(root))
|
||||
result := registry.Run(context.Background(), "read_file", map[string]any{"path": "large.txt"})
|
||||
if !result.Truncated || result.Meta[outputBudgetCategoryMeta] != string(outputCategoryFile) {
|
||||
t.Fatalf("large file was not semantically budgeted: truncated=%t meta=%#v", result.Truncated, result.Meta)
|
||||
}
|
||||
if len(result.Output) > readOutputBudgetBytes {
|
||||
t.Fatalf("large file output = %d bytes, ceiling %d", len(result.Output), readOutputBudgetBytes)
|
||||
}
|
||||
for _, want := range []string{"File: large.txt", "HEAD_MARK", "TAIL_MARK"} {
|
||||
if !strings.Contains(result.Output, want) {
|
||||
t.Fatalf("large file output missing %q", want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistryLargeSingleLineFileRetainsHeadAndTail(t *testing.T) {
|
||||
setTestTempDir(t)
|
||||
root := t.TempDir()
|
||||
content := "HEAD_SINGLE_LINE_MARK" + strings.Repeat("x", readOutputBudgetBytes*2) + "TAIL_SINGLE_LINE_MARK"
|
||||
if err := os.WriteFile(filepath.Join(root, "large.min.js"), []byte(content), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
registry := NewRegistry()
|
||||
registry.Register(NewReadFileTool(root))
|
||||
result := registry.Run(context.Background(), "read_file", map[string]any{"path": "large.min.js"})
|
||||
if !result.Truncated || result.Meta[outputBudgetCategoryMeta] != string(outputCategoryFile) {
|
||||
t.Fatalf("single-line file was not semantically budgeted: truncated=%t meta=%#v", result.Truncated, result.Meta)
|
||||
}
|
||||
if len(result.Output) > readOutputBudgetBytes || !utf8.ValidString(result.Output) {
|
||||
t.Fatalf("single-line file output is not safely bounded: bytes=%d valid=%t", len(result.Output), utf8.ValidString(result.Output))
|
||||
}
|
||||
if rawBytes, err := strconv.Atoi(result.Meta["raw_bytes"]); err != nil || rawBytes <= len(result.Output) {
|
||||
t.Fatalf("single-line file capture was not bounded head/tail output: raw=%q emitted=%d", result.Meta["raw_bytes"], len(result.Output))
|
||||
}
|
||||
for _, want := range []string{"File: large.min.js", "HEAD_SINGLE_LINE_MARK", "TAIL_SINGLE_LINE_MARK"} {
|
||||
if !strings.Contains(result.Output, want) {
|
||||
t.Fatalf("single-line file output missing %q", want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistryMinifiedFileUpdatesEmittedBytesAfterBudgeting(t *testing.T) {
|
||||
setTestTempDir(t)
|
||||
root := t.TempDir()
|
||||
content := "HEAD_MINIFIED_MARK" + strings.Repeat("x", readOutputBudgetBytes*2) + "TAIL_MINIFIED_MARK"
|
||||
if err := os.WriteFile(filepath.Join(root, "large.txt"), []byte(content), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
registry := NewRegistry()
|
||||
registry.Register(NewReadMinifiedFileTool(root))
|
||||
result := registry.Run(context.Background(), "read_minified_file", map[string]any{"path": "large.txt"})
|
||||
if !result.Truncated || result.Meta["emitted_bytes"] != strconv.Itoa(len(result.Output)) {
|
||||
t.Fatalf("minified metadata does not describe final output: truncated=%t emitted=%q bytes=%d", result.Truncated, result.Meta["emitted_bytes"], len(result.Output))
|
||||
}
|
||||
for _, want := range []string{"HEAD_MINIFIED_MARK", "TAIL_MINIFIED_MARK"} {
|
||||
if !strings.Contains(result.Output, want) {
|
||||
t.Fatalf("minified single-line file output missing %q", want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistryGrepUsesSemanticMultiFileCoverage(t *testing.T) {
|
||||
setTestTempDir(t)
|
||||
root := t.TempDir()
|
||||
for _, name := range []string{"a.txt", "b.txt", "c.txt", "d.txt"} {
|
||||
body := strings.Repeat("needle "+strings.Repeat(name, 15)+"\n", 350)
|
||||
if err := os.WriteFile(filepath.Join(root, name), []byte(body), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
registry := NewRegistry()
|
||||
registry.Register(NewGrepTool(root))
|
||||
result := registry.Run(context.Background(), "grep", map[string]any{
|
||||
"pattern": "needle",
|
||||
"path": ".",
|
||||
"head_limit": 1400,
|
||||
})
|
||||
if !result.Truncated || result.Meta[outputBudgetCategoryMeta] != string(outputCategorySearch) {
|
||||
t.Fatalf("grep was not semantically budgeted: truncated=%t meta=%#v", result.Truncated, result.Meta)
|
||||
}
|
||||
for _, name := range []string{"a.txt:", "b.txt:", "c.txt:", "d.txt:"} {
|
||||
if !strings.Contains(result.Output, name) {
|
||||
t.Fatalf("search output lost %s coverage", name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestShellOutputCategoryClassification(t *testing.T) {
|
||||
tests := []struct {
|
||||
command string
|
||||
want outputCategory
|
||||
}{
|
||||
{"go test ./...", outputCategoryTest},
|
||||
{"python -m pytest -q", outputCategoryTest},
|
||||
{"git diff --stat main...HEAD", outputCategoryDiff},
|
||||
{"git show HEAD", outputCategoryDiff},
|
||||
{"make build", outputCategoryProcess},
|
||||
}
|
||||
for _, test := range tests {
|
||||
if got := shellOutputCategory(test.command); got != test.want {
|
||||
t.Errorf("shellOutputCategory(%q) = %q, want %q", test.command, got, test.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestShellToolOutputCategoryUsesAllCommandAliases(t *testing.T) {
|
||||
providers := []struct {
|
||||
name string
|
||||
provider outputPolicyProvider
|
||||
}{
|
||||
{name: "bash", provider: NewBashTool(t.TempDir()).(outputPolicyProvider)},
|
||||
{name: "exec_command", provider: NewExecCommandTool(t.TempDir(), newExecSessionManager()).(outputPolicyProvider)},
|
||||
}
|
||||
for _, provider := range providers {
|
||||
for _, alias := range []string{"command", "cmd", "script", "shell"} {
|
||||
t.Run(provider.name+"/"+alias, func(t *testing.T) {
|
||||
if got := provider.provider.outputCategory(map[string]any{alias: "go test ./..."}); got != outputCategoryTest {
|
||||
t.Fatalf("category = %q, want test", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
if got := providers[0].provider.outputCategory(map[string]any{"command": "make build", "cmd": "go test ./..."}); got != outputCategoryProcess {
|
||||
t.Fatalf("bash alias precedence category = %q, want process", got)
|
||||
}
|
||||
if got := providers[1].provider.outputCategory(map[string]any{"cmd": "go test ./...", "command": "make build"}); got != outputCategoryTest {
|
||||
t.Fatalf("exec_command alias precedence category = %q, want test", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistryBudgetPreservesExistingTruncationAndRefreshesMetadata(t *testing.T) {
|
||||
output := "post-hook output"
|
||||
result := applyRegistryOutputBudget(newCeilingFakeTool("existing_truncation", output), "existing_truncation", map[string]any{}, Result{
|
||||
Status: StatusOK,
|
||||
Output: output,
|
||||
Truncated: true,
|
||||
Meta: map[string]string{
|
||||
"emitted_bytes": "999",
|
||||
"estimated_tokens": "999",
|
||||
"truncation_reason": "capture_budget",
|
||||
},
|
||||
})
|
||||
if !result.Truncated || result.Meta["truncated"] != "true" || result.Meta[outputBudgetReasonMeta] != "capture_budget" {
|
||||
t.Fatalf("prior truncation state was lost: %#v", result)
|
||||
}
|
||||
if result.Meta["emitted_bytes"] != strconv.Itoa(len(output)) || result.Meta["estimated_tokens"] != strconv.Itoa(estimateOutputTokens(output)) {
|
||||
t.Fatalf("final output metadata is stale: %#v", result.Meta)
|
||||
}
|
||||
if result.Meta[outputBudgetSpillCreatedMeta] != "false" {
|
||||
t.Fatalf("existing truncation unexpectedly created a new spill: %#v", result.Meta)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelfManagedOutputBudgetUsesSemanticTestPolicy(t *testing.T) {
|
||||
setTestTempDir(t)
|
||||
input := "stdout:\n" + strings.Repeat("PASS progress\n", 8_000) +
|
||||
"--- FAIL: TestImportant (0.01s)\nthing_test.go:42: expected 7, got 9\nFAIL\nexit_code: 1"
|
||||
tests := []struct {
|
||||
name string
|
||||
tool Tool
|
||||
args map[string]any
|
||||
}{
|
||||
{name: "bash", tool: NewBashTool(t.TempDir()), args: map[string]any{"command": "go test ./..."}},
|
||||
{name: "exec", tool: NewExecCommandTool(t.TempDir(), newExecSessionManager()), args: map[string]any{"cmd": "go test ./...", "max_output_tokens": 16_000}},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
got := applySelfManagedOutputBudget(test.tool, test.tool.Name(), test.args, Result{Status: StatusError, Output: input, Meta: map[string]string{}})
|
||||
if !got.Truncated || got.Meta[outputBudgetCategoryMeta] != string(outputCategoryTest) {
|
||||
t.Fatalf("semantic budget was not applied: %#v", got)
|
||||
}
|
||||
for _, want := range []string{"TestImportant", "expected 7", "FAIL", "exit_code: 1"} {
|
||||
if !strings.Contains(got.Output, want) {
|
||||
t.Fatalf("semantic test output missing %q:\n%s", want, got.Output)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
type semanticBashFakeTool struct{ ceilingFakeTool }
|
||||
|
||||
func (semanticBashFakeTool) managesOutputBudget() {}
|
||||
|
||||
func (semanticBashFakeTool) outputCategory(args map[string]any) outputCategory {
|
||||
command, _ := args["command"].(string)
|
||||
return shellOutputCategory(command)
|
||||
}
|
||||
|
||||
func TestRegistryAppliesSemanticBudgetToSelfManagedShellOutput(t *testing.T) {
|
||||
setTestTempDir(t)
|
||||
input := "startup\n" + strings.Repeat("PASS progress\n", 8_000) +
|
||||
"--- FAIL: TestRegistryImportant (0.01s)\nexpected 7, got 9\nFAIL\nexit_code: 1"
|
||||
registry := NewRegistry()
|
||||
registry.Register(semanticBashFakeTool{newCeilingFakeTool("bash", input)})
|
||||
result := registry.Run(context.Background(), "bash", map[string]any{"command": "go test ./..."})
|
||||
if !result.Truncated || result.Meta[outputBudgetReasonMeta] != "semantic_test_budget" {
|
||||
t.Fatalf("registry did not apply semantic test budget: %#v", result.Meta)
|
||||
}
|
||||
for _, want := range []string{"TestRegistryImportant", "expected 7", "FAIL", "exit_code: 1"} {
|
||||
if !strings.Contains(result.Output, want) {
|
||||
t.Fatalf("registry semantic test output missing %q:\n%s", want, result.Output)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistryAppliesSemanticBudgetToSelfManagedProcessOutput(t *testing.T) {
|
||||
setTestTempDir(t)
|
||||
input := "starting build\n" + strings.Repeat("progress\n", 8_000) +
|
||||
"WARNING: cache is cold\nERROR: compilation failed\nexit_code: 1"
|
||||
registry := NewRegistry()
|
||||
registry.Register(semanticBashFakeTool{newCeilingFakeTool("bash", input)})
|
||||
result := registry.Run(context.Background(), "bash", map[string]any{"command": "make build"})
|
||||
if !result.Truncated || result.Meta[outputBudgetReasonMeta] != "semantic_process_budget" {
|
||||
t.Fatalf("registry did not apply semantic process budget: %#v", result.Meta)
|
||||
}
|
||||
for _, want := range []string{"starting build", "WARNING", "ERROR", "exit_code: 1"} {
|
||||
if !strings.Contains(result.Output, want) {
|
||||
t.Fatalf("registry semantic process output missing %q:\n%s", want, result.Output)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRebudgetAfterHookPreservesSelfManagedExecBudget(t *testing.T) {
|
||||
setTestTempDir(t)
|
||||
registry := NewRegistry()
|
||||
registry.Register(NewExecCommandTool(t.TempDir(), newExecSessionManager()))
|
||||
|
||||
t.Run("tight call budget remains authoritative", func(t *testing.T) {
|
||||
args := map[string]any{"cmd": "make build", "max_output_tokens": 10}
|
||||
result := registry.RebudgetAfterHook(ExecCommandToolName, args, Result{Status: StatusOK, Output: strings.Repeat("x", 1_000)})
|
||||
if !result.Truncated || len(result.Output) > 40 {
|
||||
t.Fatalf("post-hook output escaped call budget: truncated=%t bytes=%d meta=%#v", result.Truncated, len(result.Output), result.Meta)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("raised call budget is not replaced by generic ceiling", func(t *testing.T) {
|
||||
args := map[string]any{"cmd": "make build", "max_output_tokens": 20_000}
|
||||
output := strings.Repeat("x", 70_000)
|
||||
result := registry.RebudgetAfterHook(ExecCommandToolName, args, Result{Status: StatusOK, Output: output})
|
||||
if result.Truncated || result.Output != output {
|
||||
t.Fatalf("post-hook output was incorrectly limited by generic ceiling: truncated=%t bytes=%d", result.Truncated, len(result.Output))
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
const (
|
||||
@@ -11,6 +12,167 @@ const (
|
||||
searchOutputBudgetBytes = 64 * 1024
|
||||
)
|
||||
|
||||
// outputCategory selects the deterministic retention policy used for a tool
|
||||
// result. It is deliberately internal: categories are an implementation detail
|
||||
// of the tools boundary, not part of Zero's public tool protocol.
|
||||
type outputCategory string
|
||||
|
||||
const (
|
||||
outputCategoryDefault outputCategory = "default"
|
||||
outputCategoryFile outputCategory = "file"
|
||||
outputCategorySearch outputCategory = "search"
|
||||
outputCategoryTest outputCategory = "test"
|
||||
outputCategoryProcess outputCategory = "process"
|
||||
outputCategoryDiff outputCategory = "diff"
|
||||
outputCategoryWorker outputCategory = "worker"
|
||||
)
|
||||
|
||||
// outputBudget combines a provider-neutral estimated-token target with the
|
||||
// existing byte ceiling. The byte ceiling is always authoritative: semantic
|
||||
// retention may use fewer bytes, but can never use more.
|
||||
type outputBudget struct {
|
||||
maxEstimatedTokens int
|
||||
hardMaxBytes int
|
||||
}
|
||||
|
||||
// budgetedOutput describes the text retained by a semantic output policy. The
|
||||
// original size is the output received by this layer; it does not claim to be
|
||||
// every byte an underlying subprocess may have produced before its own capture
|
||||
// limits were applied.
|
||||
type budgetedOutput struct {
|
||||
text string
|
||||
originalBytes int
|
||||
retainedBytes int
|
||||
estimatedOriginalTokens int
|
||||
estimatedRetainedTokens int
|
||||
truncated bool
|
||||
category outputCategory
|
||||
reason string
|
||||
spillPath string
|
||||
}
|
||||
|
||||
// outputPolicyProvider optionally assigns a semantic category to a tool call.
|
||||
// Tools that do not implement it use outputCategoryDefault.
|
||||
type outputPolicyProvider interface {
|
||||
outputCategory(args map[string]any) outputCategory
|
||||
}
|
||||
|
||||
// estimateOutputTokens is a deterministic, provider-neutral estimate used only
|
||||
// for output budgeting. It is not exact provider tokenization. ASCII non-space
|
||||
// text uses the repository's established four-bytes-per-token approximation;
|
||||
// each byte of non-ASCII UTF-8 is counted as a token, intentionally
|
||||
// overestimating multilingual text and emoji rather than letting them bypass a
|
||||
// budget. Existing hard byte ceilings remain the final safety limit.
|
||||
func estimateOutputTokens(value string) int {
|
||||
asciiNonSpace, nonASCIIBytes := outputTokenComponents(value)
|
||||
return (asciiNonSpace+3)/4 + nonASCIIBytes
|
||||
}
|
||||
|
||||
func outputTokenComponents(value string) (asciiNonSpace int, nonASCIIBytes int) {
|
||||
for index := 0; index < len(value); {
|
||||
r, size := utf8.DecodeRuneInString(value[index:])
|
||||
if r == utf8.RuneError && size == 1 {
|
||||
// Invalid input is charged conservatively one token per byte. Budget
|
||||
// slicing itself remains rune-safe for valid UTF-8 tool output.
|
||||
nonASCIIBytes++
|
||||
index++
|
||||
continue
|
||||
}
|
||||
if r < utf8.RuneSelf {
|
||||
switch r {
|
||||
case ' ', '\t', '\n', '\r', '\f', '\v':
|
||||
default:
|
||||
asciiNonSpace++
|
||||
}
|
||||
} else {
|
||||
nonASCIIBytes += size
|
||||
}
|
||||
index += size
|
||||
}
|
||||
return asciiNonSpace, nonASCIIBytes
|
||||
}
|
||||
|
||||
const defaultSemanticTruncationNotice = "\n[zero] output truncated\n"
|
||||
|
||||
// budgetDefaultOutput applies the safe fallback policy: retain a rune-safe head
|
||||
// and tail around a stable truncation notice. Small output is returned
|
||||
// byte-identically. Spill creation is intentionally handled by the shared
|
||||
// boundary integration so every semantic policy uses the existing spill path.
|
||||
func budgetDefaultOutput(output string, budget outputBudget) budgetedOutput {
|
||||
result := budgetedOutput{
|
||||
text: output,
|
||||
originalBytes: len(output),
|
||||
retainedBytes: len(output),
|
||||
estimatedOriginalTokens: estimateOutputTokens(output),
|
||||
estimatedRetainedTokens: estimateOutputTokens(output),
|
||||
category: outputCategoryDefault,
|
||||
}
|
||||
if fitsOutputBudget(output, budget) {
|
||||
return result
|
||||
}
|
||||
|
||||
reason := outputBudgetReason(result.originalBytes, result.estimatedOriginalTokens, budget)
|
||||
maxContentBytes := len(output)
|
||||
if budget.hardMaxBytes > 0 {
|
||||
maxContentBytes = min(maxContentBytes, max(0, budget.hardMaxBytes-len(defaultSemanticTruncationNotice)))
|
||||
}
|
||||
|
||||
// Find the largest deterministic head+tail window that satisfies both the
|
||||
// estimate and the hard byte ceiling. fitsOutputBudget is monotonic as this
|
||||
// window grows, so binary search avoids repeatedly trimming one rune at a time.
|
||||
low, high := 0, maxContentBytes
|
||||
best := ""
|
||||
if fitsOutputBudget(defaultSemanticTruncationNotice, budget) {
|
||||
best = defaultSemanticTruncationNotice
|
||||
}
|
||||
for low <= high {
|
||||
window := low + (high-low)/2
|
||||
candidate := defaultHeadTail(output, window)
|
||||
if fitsOutputBudget(candidate, budget) {
|
||||
best = candidate
|
||||
low = window + 1
|
||||
} else {
|
||||
high = window - 1
|
||||
}
|
||||
}
|
||||
|
||||
result.text = best
|
||||
result.retainedBytes = len(best)
|
||||
result.estimatedRetainedTokens = estimateOutputTokens(best)
|
||||
result.truncated = true
|
||||
result.reason = reason
|
||||
return result
|
||||
}
|
||||
|
||||
func defaultHeadTail(output string, contentBytes int) string {
|
||||
if contentBytes <= 0 {
|
||||
return defaultSemanticTruncationNotice
|
||||
}
|
||||
headBytes := contentBytes / 2
|
||||
tailBytes := contentBytes - headBytes
|
||||
return utf8Prefix(output, headBytes) + defaultSemanticTruncationNotice + utf8Suffix(output, tailBytes)
|
||||
}
|
||||
|
||||
func fitsOutputBudget(output string, budget outputBudget) bool {
|
||||
if budget.hardMaxBytes > 0 && len(output) > budget.hardMaxBytes {
|
||||
return false
|
||||
}
|
||||
return budget.maxEstimatedTokens <= 0 || estimateOutputTokens(output) <= budget.maxEstimatedTokens
|
||||
}
|
||||
|
||||
func outputBudgetReason(originalBytes int, estimatedTokens int, budget outputBudget) string {
|
||||
overTokens := budget.maxEstimatedTokens > 0 && estimatedTokens > budget.maxEstimatedTokens
|
||||
overBytes := budget.hardMaxBytes > 0 && originalBytes > budget.hardMaxBytes
|
||||
switch {
|
||||
case overTokens && overBytes:
|
||||
return "token_and_byte_budget"
|
||||
case overTokens:
|
||||
return "estimated_token_budget"
|
||||
default:
|
||||
return "hard_byte_ceiling"
|
||||
}
|
||||
}
|
||||
|
||||
type outputBudgetResult struct {
|
||||
Output string
|
||||
Truncated bool
|
||||
@@ -56,6 +218,7 @@ func estimatedTokensFromBytes(bytes int) int {
|
||||
|
||||
type outputBudgetBuilder struct {
|
||||
builder strings.Builder
|
||||
capture *boundedBuffer
|
||||
rawBytes int
|
||||
maxBytes int
|
||||
hint string
|
||||
@@ -65,17 +228,60 @@ func newOutputBudgetBuilder(maxBytes int, hint string) *outputBudgetBuilder {
|
||||
return &outputBudgetBuilder{maxBytes: maxBytes, hint: hint}
|
||||
}
|
||||
|
||||
// newHeadTailOutputBudgetBuilder bounds capture while retaining both ends of
|
||||
// an input for a later semantic policy. It is used for registry-driven file
|
||||
// reads; direct Tool.Run calls keep the legacy prefix-only builder above.
|
||||
func newHeadTailOutputBudgetBuilder(maxBytes int, hint string) *outputBudgetBuilder {
|
||||
headBytes := maxBytes / 2
|
||||
return &outputBudgetBuilder{
|
||||
capture: newBoundedBuffer(headBytes, maxBytes-headBytes),
|
||||
maxBytes: maxBytes,
|
||||
hint: hint,
|
||||
}
|
||||
}
|
||||
|
||||
func (builder *outputBudgetBuilder) WriteString(value string) {
|
||||
builder.rawBytes += len(value)
|
||||
if builder.maxBytes <= 0 || builder.builder.Len() >= builder.maxBytes {
|
||||
if builder.capture != nil {
|
||||
_, _ = builder.capture.Write([]byte(value))
|
||||
return
|
||||
}
|
||||
if builder.maxBytes <= 0 {
|
||||
builder.builder.WriteString(value)
|
||||
return
|
||||
}
|
||||
if builder.builder.Len() >= builder.maxBytes {
|
||||
return
|
||||
}
|
||||
remaining := builder.maxBytes - builder.builder.Len()
|
||||
builder.builder.WriteString(utf8Prefix(value, remaining))
|
||||
}
|
||||
|
||||
// applyLegacyByteBudgetToResult preserves direct Tool.Run behavior for callers
|
||||
// that intentionally bypass Registry.RunWithOptions. Agent/MCP execution uses
|
||||
// the registry's post-redaction semantic boundary instead.
|
||||
func applyLegacyByteBudgetToResult(result Result, maxBytes int, hint string) Result {
|
||||
budgeted := applyOutputBudget(result.Output, maxBytes, hint)
|
||||
result.Output = budgeted.Output
|
||||
result.Truncated = result.Truncated || budgeted.Truncated
|
||||
if result.Meta == nil {
|
||||
result.Meta = map[string]string{}
|
||||
}
|
||||
for key, value := range outputBudgetMeta(budgeted) {
|
||||
result.Meta[key] = value
|
||||
}
|
||||
if budgeted.Truncated {
|
||||
result.Meta["truncated"] = "true"
|
||||
result.Meta["truncation_reason"] = "byte_budget"
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func (builder *outputBudgetBuilder) Result() outputBudgetResult {
|
||||
output := builder.builder.String()
|
||||
if builder.capture != nil {
|
||||
output = builder.capture.retained()
|
||||
}
|
||||
result := outputBudgetResult{
|
||||
Output: output,
|
||||
RawBytes: builder.rawBytes,
|
||||
@@ -90,7 +296,12 @@ func (builder *outputBudgetBuilder) Result() outputBudgetResult {
|
||||
if budget < 0 {
|
||||
budget = 0
|
||||
}
|
||||
result.Output = utf8Prefix(output, budget) + marker
|
||||
if builder.capture != nil {
|
||||
headBytes := budget / 2
|
||||
result.Output = utf8Prefix(output, headBytes) + marker + utf8Suffix(output, budget-headBytes)
|
||||
} else {
|
||||
result.Output = utf8Prefix(output, budget) + marker
|
||||
}
|
||||
result.Truncated = true
|
||||
result.EmittedBytes = len(result.Output)
|
||||
return result
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
func TestEstimateOutputTokensASCIIAndUnicode(t *testing.T) {
|
||||
if got := estimateOutputTokens("abcd efgh"); got != 2 {
|
||||
t.Fatalf("ASCII estimate = %d, want 2", got)
|
||||
}
|
||||
unicodeText := "🙂界"
|
||||
if got := estimateOutputTokens(unicodeText); got < len([]byte(unicodeText)) {
|
||||
t.Fatalf("Unicode estimate = %d, want conservative estimate >= %d", got, len([]byte(unicodeText)))
|
||||
}
|
||||
if first, second := estimateOutputTokens(unicodeText), estimateOutputTokens(unicodeText); first != second {
|
||||
t.Fatalf("estimator is not deterministic: %d != %d", first, second)
|
||||
}
|
||||
if input := strings.Repeat("\u0080", 4); estimateOutputTokens(input) < len([]byte(input)) {
|
||||
t.Fatalf("U+0080 estimate = %d, want conservative estimate >= %d", estimateOutputTokens(input), len([]byte(input)))
|
||||
}
|
||||
}
|
||||
|
||||
func TestBudgetDefaultOutputLeavesSmallOutputByteIdentical(t *testing.T) {
|
||||
input := "small output\nwith spacing\n"
|
||||
got := budgetDefaultOutput(input, outputBudget{maxEstimatedTokens: 100, hardMaxBytes: 1024})
|
||||
if got.text != input {
|
||||
t.Fatalf("small output changed: got %q want %q", got.text, input)
|
||||
}
|
||||
if got.truncated || got.originalBytes != len(input) || got.retainedBytes != len(input) {
|
||||
t.Fatalf("unexpected small-output metadata: %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBudgetDefaultOutputKeepsUTF8HeadTailWithinHardCeiling(t *testing.T) {
|
||||
input := "HEAD🙂\n" + strings.Repeat("界", 200) + "\nTAIL🙂"
|
||||
const hardMax = 120
|
||||
got := budgetDefaultOutput(input, outputBudget{maxEstimatedTokens: 10_000, hardMaxBytes: hardMax})
|
||||
if !got.truncated {
|
||||
t.Fatal("large output was not truncated")
|
||||
}
|
||||
if len(got.text) > hardMax {
|
||||
t.Fatalf("retained %d bytes, hard ceiling %d", len(got.text), hardMax)
|
||||
}
|
||||
if !utf8.ValidString(got.text) {
|
||||
t.Fatalf("budgeted output is invalid UTF-8: %q", got.text)
|
||||
}
|
||||
for _, want := range []string{"HEAD", "TAIL", "output truncated"} {
|
||||
if !strings.Contains(got.text, want) {
|
||||
t.Fatalf("budgeted output missing %q: %q", want, got.text)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBudgetDefaultOutputHonorsEstimatedTokenBudget(t *testing.T) {
|
||||
input := strings.Repeat("🙂", 100)
|
||||
got := budgetDefaultOutput(input, outputBudget{maxEstimatedTokens: 60, hardMaxBytes: 1024})
|
||||
if !got.truncated || got.reason != "estimated_token_budget" {
|
||||
t.Fatalf("unexpected token-budget result: %#v", got)
|
||||
}
|
||||
if got.estimatedRetainedTokens > 60 {
|
||||
t.Fatalf("retained estimate = %d, budget 60", got.estimatedRetainedTokens)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBudgetDefaultOutputDeterministic(t *testing.T) {
|
||||
input := "head\n" + strings.Repeat("same line\n", 1000) + "tail\n"
|
||||
budget := outputBudget{maxEstimatedTokens: 80, hardMaxBytes: 512}
|
||||
want := budgetDefaultOutput(input, budget)
|
||||
for iteration := 0; iteration < 20; iteration++ {
|
||||
got := budgetDefaultOutput(input, budget)
|
||||
if got != want {
|
||||
t.Fatalf("iteration %d differs:\n got %#v\nwant %#v", iteration, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBudgetDefaultOutputTinyBudgetsNeverExceedTheirCeilings(t *testing.T) {
|
||||
input := strings.Repeat("oversized output ", 100)
|
||||
for _, budget := range []outputBudget{
|
||||
{hardMaxBytes: len(defaultSemanticTruncationNotice) - 1},
|
||||
{maxEstimatedTokens: estimateOutputTokens(defaultSemanticTruncationNotice) - 1},
|
||||
} {
|
||||
got := budgetDefaultOutput(input, budget)
|
||||
if !got.truncated || !fitsOutputBudget(got.text, budget) {
|
||||
t.Fatalf("tiny budget result does not fit: budget=%#v result=%#v", budget, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -25,23 +25,19 @@ const defaultOutputCeilingTokens = 16_000
|
||||
// disables the ceiling entirely; unset or unparsable keeps the default.
|
||||
const outputCeilingEnv = "ZERO_TOOL_OUTPUT_CEILING_TOKENS"
|
||||
|
||||
// selfBudgeting marks a tool that enforces its own deliberate output budget —
|
||||
// possibly model-raisable (exec_command) — which the registry ceiling must not
|
||||
// second-guess. The method is unexported on purpose: only tools in this
|
||||
// package can opt out, so an MCP-served tool can never exempt itself.
|
||||
// selfBudgeting marks a tool with a deliberate capture-aware output budget —
|
||||
// possibly model-raisable (exec_command). The registry applies semantic
|
||||
// retention within that explicit budget rather than replacing it with the
|
||||
// default ceiling. The method is unexported so an MCP-served tool cannot opt
|
||||
// into this path.
|
||||
type selfBudgeting interface{ managesOutputBudget() }
|
||||
|
||||
// The exemption list, kept in one place. Each of these applies its own budget
|
||||
// before returning: bash (bashOutputBudgetBytes per stream + spill),
|
||||
// exec_command (model-raisable token budget + spill), read tools (128 KiB),
|
||||
// search tools (64 KiB).
|
||||
func (bashTool) managesOutputBudget() {}
|
||||
func (execCommandTool) managesOutputBudget() {}
|
||||
func (readFileTool) managesOutputBudget() {}
|
||||
func (readMinifiedFileTool) managesOutputBudget() {}
|
||||
func (grepTool) managesOutputBudget() {}
|
||||
func (globTool) managesOutputBudget() {}
|
||||
func (listDirectoryTool) managesOutputBudget() {}
|
||||
// The list is kept in one place. Shell/process tools retain their established
|
||||
// capture-aware budgets: bash (per-stream bounded capture) and exec_command
|
||||
// (model-raisable bounded session output). File/search tools use the standard
|
||||
// shared post-redaction boundary.
|
||||
func (bashTool) managesOutputBudget() {}
|
||||
func (execCommandTool) managesOutputBudget() {}
|
||||
|
||||
func resolveOutputCeilingTokens() int {
|
||||
raw := strings.TrimSpace(os.Getenv(outputCeilingEnv))
|
||||
@@ -61,6 +57,12 @@ func resolveOutputCeilingTokens() int {
|
||||
// and the spill file agree on what was hidden.
|
||||
func enforceOutputCeiling(toolName string, result Result) Result {
|
||||
ceiling := resolveOutputCeilingTokens()
|
||||
switch toolName {
|
||||
case "read_file", "read_minified_file":
|
||||
ceiling = readOutputBudgetBytes / 4
|
||||
case "grep", "glob", "list_directory":
|
||||
ceiling = searchOutputBudgetBytes / 4
|
||||
}
|
||||
if ceiling <= 0 {
|
||||
return result
|
||||
}
|
||||
|
||||
@@ -55,10 +55,10 @@ func TestOutputCeilingCapsUnbudgetedTool(t *testing.T) {
|
||||
if result.Meta["raw_bytes"] != strconv.Itoa(len(big)) {
|
||||
t.Fatalf("raw_bytes = %s, want %d", result.Meta["raw_bytes"], len(big))
|
||||
}
|
||||
if !strings.Contains(result.Output, "full output saved to ") {
|
||||
if !strings.Contains(result.Output, "full output received by budgeting layer saved to ") {
|
||||
t.Fatalf("ceiling truncation must include a spill hint: %q", result.Output[:200])
|
||||
}
|
||||
start := strings.Index(result.Output, "full output saved to ") + len("full output saved to ")
|
||||
start := strings.Index(result.Output, "full output received by budgeting layer saved to ") + len("full output received by budgeting layer saved to ")
|
||||
end := strings.Index(result.Output[start:], " (grep")
|
||||
content, err := os.ReadFile(result.Output[start : start+end])
|
||||
if err != nil {
|
||||
@@ -123,11 +123,6 @@ func TestSelfBudgetingExemptionList(t *testing.T) {
|
||||
exempt := []Tool{
|
||||
NewBashTool(dir),
|
||||
NewExecCommandTool(dir, newExecSessionManager()),
|
||||
NewReadFileTool(dir),
|
||||
NewReadMinifiedFileTool(dir),
|
||||
NewGrepTool(dir),
|
||||
NewGlobTool(dir),
|
||||
NewListDirectoryTool(dir),
|
||||
}
|
||||
for _, tool := range exempt {
|
||||
if _, ok := tool.(selfBudgeting); !ok {
|
||||
@@ -137,4 +132,9 @@ func TestSelfBudgetingExemptionList(t *testing.T) {
|
||||
if _, ok := NewWebFetchTool().(selfBudgeting); ok {
|
||||
t.Error("web_fetch must NOT be exempt — the ceiling is its only budget")
|
||||
}
|
||||
for _, tool := range []Tool{NewReadFileTool(dir), NewReadMinifiedFileTool(dir), NewGrepTool(dir), NewGlobTool(dir), NewListDirectoryTool(dir)} {
|
||||
if _, ok := tool.(selfBudgeting); ok {
|
||||
t.Errorf("%s must use the shared post-redaction budget boundary", tool.Name())
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,447 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// budgetSemanticOutput applies the deterministic policy for category. The
|
||||
// caller supplies output after secret redaction. Small results pass through
|
||||
// byte-identically; only oversized results are semantically reduced.
|
||||
func budgetSemanticOutput(output string, category outputCategory, budget outputBudget) budgetedOutput {
|
||||
if category == "" {
|
||||
category = outputCategoryDefault
|
||||
}
|
||||
if fitsOutputBudget(output, budget) {
|
||||
return unchangedBudgetedOutput(output, category)
|
||||
}
|
||||
|
||||
var retained string
|
||||
switch category {
|
||||
case outputCategoryFile:
|
||||
retained = budgetFileLines(output, budget)
|
||||
case outputCategorySearch:
|
||||
retained = budgetSearchLines(output, budget)
|
||||
case outputCategoryTest:
|
||||
retained = budgetTestLines(output, budget)
|
||||
case outputCategoryProcess:
|
||||
retained = budgetProcessLines(output, budget)
|
||||
case outputCategoryDiff:
|
||||
retained = budgetDiffUnits(output, budget)
|
||||
case outputCategoryWorker:
|
||||
retained = budgetWorkerLines(output, budget)
|
||||
default:
|
||||
return budgetDefaultOutput(output, budget)
|
||||
}
|
||||
|
||||
if retained == "" || !fitsOutputBudget(retained, budget) {
|
||||
fallback := budgetDefaultOutput(output, budget)
|
||||
fallback.category = category
|
||||
return fallback
|
||||
}
|
||||
return truncatedBudgetedOutput(output, retained, category, "semantic_"+string(category)+"_budget")
|
||||
}
|
||||
|
||||
func unchangedBudgetedOutput(output string, category outputCategory) budgetedOutput {
|
||||
estimated := estimateOutputTokens(output)
|
||||
return budgetedOutput{
|
||||
text: output,
|
||||
originalBytes: len(output),
|
||||
retainedBytes: len(output),
|
||||
estimatedOriginalTokens: estimated,
|
||||
estimatedRetainedTokens: estimated,
|
||||
category: category,
|
||||
}
|
||||
}
|
||||
|
||||
func truncatedBudgetedOutput(original, retained string, category outputCategory, reason string) budgetedOutput {
|
||||
return budgetedOutput{
|
||||
text: retained,
|
||||
originalBytes: len(original),
|
||||
retainedBytes: len(retained),
|
||||
estimatedOriginalTokens: estimateOutputTokens(original),
|
||||
estimatedRetainedTokens: estimateOutputTokens(retained),
|
||||
truncated: true,
|
||||
category: category,
|
||||
reason: reason,
|
||||
}
|
||||
}
|
||||
|
||||
func budgetFileLines(output string, budget outputBudget) string {
|
||||
lines := outputLines(output)
|
||||
priorities := make([]int, 0, len(lines))
|
||||
// Keep the existing file/range header first, then alternate from the start
|
||||
// and end so both requested-location boundaries survive.
|
||||
if len(lines) > 0 {
|
||||
priorities = append(priorities, 0)
|
||||
}
|
||||
for left, right := 1, len(lines)-1; left <= right; left, right = left+1, right-1 {
|
||||
priorities = append(priorities, left)
|
||||
if right != left {
|
||||
priorities = append(priorities, right)
|
||||
}
|
||||
}
|
||||
units := lineDiffUnits(lines)
|
||||
selected := selectPrioritizedUnits(units, priorities, budget)
|
||||
for index := range selected {
|
||||
// Index zero is the file/range header. If no non-blank requested
|
||||
// content line fits as a complete line, let the common caller fall
|
||||
// back to its UTF-8-safe head/tail policy rather than returning only
|
||||
// a header and omission marker for a minified single-line file.
|
||||
if index > 0 && strings.TrimSpace(lines[index]) != "" {
|
||||
return renderSelectedUnits(units, selected)
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func budgetSearchLines(output string, budget outputBudget) string {
|
||||
lines := collapseConsecutiveDuplicateLines(outputLines(output))
|
||||
priorities := make([]int, 0, len(lines))
|
||||
for index, line := range lines {
|
||||
lower := strings.ToLower(line)
|
||||
// A match body naturally contains the search term (often literally
|
||||
// "match"); treat only non-result lines as summaries so one busy file
|
||||
// cannot consume the budget before cross-file representatives are chosen.
|
||||
if searchResultFile(line) == "" && (strings.Contains(lower, "match") || strings.Contains(lower, "truncated") || strings.Contains(lower, "result")) {
|
||||
priorities = append(priorities, index)
|
||||
}
|
||||
}
|
||||
if len(lines) > 0 {
|
||||
priorities = append(priorities, 0, len(lines)-1)
|
||||
}
|
||||
seenFiles := map[string]bool{}
|
||||
for index, line := range lines {
|
||||
if file := searchResultFile(line); file != "" && !seenFiles[file] {
|
||||
seenFiles[file] = true
|
||||
priorities = append(priorities, index)
|
||||
}
|
||||
}
|
||||
priorities = append(priorities, sequence(len(lines))...)
|
||||
return retainPrioritizedLines(lines, priorities, budget)
|
||||
}
|
||||
|
||||
func budgetTestLines(output string, budget outputBudget) string {
|
||||
lines := collapseConsecutiveDuplicateLines(outputLines(output))
|
||||
priorities := make([]int, 0, len(lines))
|
||||
for index, line := range lines {
|
||||
if isTestFailureLine(line) {
|
||||
for contextIndex := max(0, index-2); contextIndex <= min(len(lines)-1, index+3); contextIndex++ {
|
||||
priorities = append(priorities, contextIndex)
|
||||
}
|
||||
}
|
||||
}
|
||||
// Final summaries and process status usually live at the tail.
|
||||
for index := max(0, len(lines)-16); index < len(lines); index++ {
|
||||
priorities = append(priorities, index)
|
||||
}
|
||||
for index := 0; index < min(8, len(lines)); index++ {
|
||||
priorities = append(priorities, index)
|
||||
}
|
||||
priorities = append(priorities, sequence(len(lines))...)
|
||||
return retainPrioritizedLines(lines, priorities, budget)
|
||||
}
|
||||
|
||||
func budgetProcessLines(output string, budget outputBudget) string {
|
||||
lines := collapseConsecutiveDuplicateLines(outputLines(output))
|
||||
priorities := make([]int, 0, len(lines))
|
||||
for index := 0; index < min(10, len(lines)); index++ {
|
||||
priorities = append(priorities, index)
|
||||
}
|
||||
seenDiagnostic := map[string]bool{}
|
||||
for index, line := range lines {
|
||||
lower := strings.ToLower(line)
|
||||
if containsAny(lower, "error", "warning", "warn:", "fatal", "panic", "failed", "denied") && !seenDiagnostic[line] {
|
||||
seenDiagnostic[line] = true
|
||||
priorities = append(priorities, index)
|
||||
}
|
||||
}
|
||||
for index := max(0, len(lines)-16); index < len(lines); index++ {
|
||||
priorities = append(priorities, index)
|
||||
}
|
||||
priorities = append(priorities, sequence(len(lines))...)
|
||||
return retainPrioritizedLines(lines, priorities, budget)
|
||||
}
|
||||
|
||||
func budgetWorkerLines(output string, budget outputBudget) string {
|
||||
lines := collapseConsecutiveDuplicateLines(outputLines(output))
|
||||
priorities := make([]int, 0, len(lines))
|
||||
for index, line := range lines {
|
||||
lower := strings.ToLower(line)
|
||||
if containsAny(lower, "status", "error", "failed", "failure", "session_id", "changed", "files", "tools executed", "conclusion", "result") {
|
||||
priorities = append(priorities, index)
|
||||
}
|
||||
}
|
||||
// A specialist's final conclusion is conventionally at the end.
|
||||
for index := max(0, len(lines)-20); index < len(lines); index++ {
|
||||
priorities = append(priorities, index)
|
||||
}
|
||||
for index := 0; index < min(6, len(lines)); index++ {
|
||||
priorities = append(priorities, index)
|
||||
}
|
||||
priorities = append(priorities, sequence(len(lines))...)
|
||||
return retainPrioritizedLines(lines, priorities, budget)
|
||||
}
|
||||
|
||||
type diffUnit struct {
|
||||
order int
|
||||
text string
|
||||
key bool
|
||||
}
|
||||
|
||||
func budgetDiffUnits(output string, budget outputBudget) string {
|
||||
units := splitDiffUnits(output)
|
||||
if len(units) == 0 {
|
||||
return ""
|
||||
}
|
||||
priority := make([]int, 0, len(units))
|
||||
for index, unit := range units {
|
||||
if unit.key {
|
||||
priority = append(priority, index)
|
||||
}
|
||||
}
|
||||
priority = append(priority, sequence(len(units))...)
|
||||
return retainPrioritizedUnits(units, priority, budget)
|
||||
}
|
||||
|
||||
// splitDiffUnits keeps each hunk indivisible. File headers/stat text are key
|
||||
// units so broad file coverage is selected before secondary hunks.
|
||||
func splitDiffUnits(output string) []diffUnit {
|
||||
lines := outputLines(output)
|
||||
units := make([]diffUnit, 0)
|
||||
var current []string
|
||||
currentKey := false
|
||||
flush := func() {
|
||||
if len(current) == 0 {
|
||||
return
|
||||
}
|
||||
units = append(units, diffUnit{order: len(units), text: strings.Join(current, "\n"), key: currentKey})
|
||||
current = nil
|
||||
currentKey = false
|
||||
}
|
||||
for _, line := range lines {
|
||||
if strings.HasPrefix(line, "diff --git ") {
|
||||
flush()
|
||||
currentKey = true
|
||||
current = append(current, line)
|
||||
continue
|
||||
}
|
||||
if strings.HasPrefix(line, "@@ ") {
|
||||
flush()
|
||||
current = append(current, line)
|
||||
continue
|
||||
}
|
||||
if len(current) == 0 {
|
||||
currentKey = true // diff stat/summary or standalone file headers
|
||||
}
|
||||
current = append(current, line)
|
||||
}
|
||||
flush()
|
||||
return units
|
||||
}
|
||||
|
||||
func retainPrioritizedLines(lines []string, priorities []int, budget outputBudget) string {
|
||||
return retainPrioritizedUnits(lineDiffUnits(lines), priorities, budget)
|
||||
}
|
||||
|
||||
func lineDiffUnits(lines []string) []diffUnit {
|
||||
units := make([]diffUnit, 0, len(lines))
|
||||
for index, line := range lines {
|
||||
units = append(units, diffUnit{order: index, text: line})
|
||||
}
|
||||
return units
|
||||
}
|
||||
|
||||
func retainPrioritizedUnits(units []diffUnit, priorities []int, budget outputBudget) string {
|
||||
return renderSelectedUnits(units, selectPrioritizedUnits(units, priorities, budget))
|
||||
}
|
||||
|
||||
func selectPrioritizedUnits(units []diffUnit, priorities []int, budget outputBudget) map[int]bool {
|
||||
selected := map[int]bool{}
|
||||
indexes := make([]int, 0, len(priorities))
|
||||
cost := retainedUnitCost{}
|
||||
for _, index := range priorities {
|
||||
if index < 0 || index >= len(units) || selected[index] {
|
||||
continue
|
||||
}
|
||||
position := sort.SearchInts(indexes, index)
|
||||
previous, next := -1, len(units)
|
||||
if position > 0 {
|
||||
previous = indexes[position-1]
|
||||
}
|
||||
if position < len(indexes) {
|
||||
next = indexes[position]
|
||||
}
|
||||
|
||||
candidate := cost
|
||||
candidate.add(units[index].text)
|
||||
if len(indexes) == 0 {
|
||||
candidate.addOmission(index)
|
||||
candidate.addOmission(len(units) - index - 1)
|
||||
} else {
|
||||
candidate.removeOmission(next - previous - 1)
|
||||
candidate.addOmission(index - previous - 1)
|
||||
candidate.addOmission(next - index - 1)
|
||||
}
|
||||
if !candidate.fits(budget) {
|
||||
continue
|
||||
}
|
||||
selected[index] = true
|
||||
indexes = append(indexes, 0)
|
||||
copy(indexes[position+1:], indexes[position:])
|
||||
indexes[position] = index
|
||||
cost = candidate
|
||||
}
|
||||
return selected
|
||||
}
|
||||
|
||||
// retainedUnitCost tracks the exact rendered size of selected units without
|
||||
// repeatedly rebuilding the full candidate text. Newlines between rendered
|
||||
// parts add bytes but no estimated tokens, so the token components remain
|
||||
// additive even as priorities select units out of source order.
|
||||
type retainedUnitCost struct {
|
||||
textBytes int
|
||||
asciiNonSpace int
|
||||
nonASCIIBytes int
|
||||
parts int
|
||||
}
|
||||
|
||||
func (cost *retainedUnitCost) add(text string) {
|
||||
ascii, nonASCII := outputTokenComponents(text)
|
||||
cost.textBytes += len(text)
|
||||
cost.asciiNonSpace += ascii
|
||||
cost.nonASCIIBytes += nonASCII
|
||||
cost.parts++
|
||||
}
|
||||
|
||||
func (cost *retainedUnitCost) remove(text string) {
|
||||
ascii, nonASCII := outputTokenComponents(text)
|
||||
cost.textBytes -= len(text)
|
||||
cost.asciiNonSpace -= ascii
|
||||
cost.nonASCIIBytes -= nonASCII
|
||||
cost.parts--
|
||||
}
|
||||
|
||||
func (cost *retainedUnitCost) addOmission(count int) {
|
||||
if count > 0 {
|
||||
cost.add(omittedSectionsMarker(count))
|
||||
}
|
||||
}
|
||||
|
||||
func (cost *retainedUnitCost) removeOmission(count int) {
|
||||
if count > 0 {
|
||||
cost.remove(omittedSectionsMarker(count))
|
||||
}
|
||||
}
|
||||
|
||||
func (cost retainedUnitCost) fits(budget outputBudget) bool {
|
||||
bytes := cost.textBytes + max(0, cost.parts-1)
|
||||
tokens := (cost.asciiNonSpace+3)/4 + cost.nonASCIIBytes
|
||||
return (budget.hardMaxBytes <= 0 || bytes <= budget.hardMaxBytes) &&
|
||||
(budget.maxEstimatedTokens <= 0 || tokens <= budget.maxEstimatedTokens)
|
||||
}
|
||||
|
||||
func omittedSectionsMarker(count int) string {
|
||||
return fmt.Sprintf("[zero] ... %d section(s) omitted ...", count)
|
||||
}
|
||||
|
||||
func renderSelectedUnits(units []diffUnit, selected map[int]bool) string {
|
||||
indexes := make([]int, 0, len(selected))
|
||||
for index := range selected {
|
||||
indexes = append(indexes, index)
|
||||
}
|
||||
sort.Ints(indexes)
|
||||
if len(indexes) == 0 {
|
||||
return ""
|
||||
}
|
||||
parts := make([]string, 0, len(indexes)*2)
|
||||
previous := -1
|
||||
for _, index := range indexes {
|
||||
if previous >= 0 && index != previous+1 {
|
||||
omitted := index - previous - 1
|
||||
parts = append(parts, omittedSectionsMarker(omitted))
|
||||
} else if previous < 0 && index > 0 {
|
||||
parts = append(parts, omittedSectionsMarker(index))
|
||||
}
|
||||
parts = append(parts, units[index].text)
|
||||
previous = index
|
||||
}
|
||||
if previous < len(units)-1 {
|
||||
parts = append(parts, omittedSectionsMarker(len(units)-previous-1))
|
||||
}
|
||||
return strings.Join(parts, "\n")
|
||||
}
|
||||
|
||||
func outputLines(output string) []string {
|
||||
return strings.Split(strings.TrimRight(output, "\r\n"), "\n")
|
||||
}
|
||||
|
||||
func collapseConsecutiveDuplicateLines(lines []string) []string {
|
||||
if len(lines) < 2 {
|
||||
return lines
|
||||
}
|
||||
result := make([]string, 0, len(lines))
|
||||
for index := 0; index < len(lines); {
|
||||
end := index + 1
|
||||
for end < len(lines) && lines[end] == lines[index] {
|
||||
end++
|
||||
}
|
||||
result = append(result, lines[index])
|
||||
if count := end - index; count > 1 {
|
||||
result = append(result, fmt.Sprintf("[zero] previous line repeated %d more time(s)", count-1))
|
||||
}
|
||||
index = end
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func searchResultFile(line string) string {
|
||||
// Search output is path:line: text. Locate the first numeric field from
|
||||
// the left after an optional Windows drive prefix. Match text may itself
|
||||
// contain numeric colon-delimited fields, which are not record structure.
|
||||
searchFrom := 0
|
||||
if len(line) >= 2 && line[1] == ':' && ((line[0] >= 'A' && line[0] <= 'Z') || (line[0] >= 'a' && line[0] <= 'z')) {
|
||||
searchFrom = 2
|
||||
}
|
||||
for searchFrom < len(line) {
|
||||
lineStartOffset := strings.IndexByte(line[searchFrom:], ':')
|
||||
if lineStartOffset < 0 {
|
||||
return ""
|
||||
}
|
||||
lineStart := searchFrom + lineStartOffset
|
||||
lineEndOffset := strings.IndexByte(line[lineStart+1:], ':')
|
||||
if lineEndOffset < 0 {
|
||||
return ""
|
||||
}
|
||||
lineEnd := lineStart + 1 + lineEndOffset
|
||||
if _, err := strconv.Atoi(line[lineStart+1 : lineEnd]); err == nil {
|
||||
return strings.TrimSpace(line[:lineStart])
|
||||
}
|
||||
searchFrom = lineStart + 1
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func isTestFailureLine(line string) bool {
|
||||
lower := strings.ToLower(line)
|
||||
return containsAny(lower, "--- fail:", "failed", "failure", "panic", "assert", "error", "fatal", "expected", "actual:")
|
||||
}
|
||||
|
||||
func containsAny(value string, needles ...string) bool {
|
||||
for _, needle := range needles {
|
||||
if strings.Contains(value, needle) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func sequence(length int) []int {
|
||||
result := make([]int, length)
|
||||
for index := range result {
|
||||
result[index] = index
|
||||
}
|
||||
return result
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
func semanticTestBudget() outputBudget {
|
||||
return outputBudget{maxEstimatedTokens: 140, hardMaxBytes: 900}
|
||||
}
|
||||
|
||||
func TestFileOutputPolicyKeepsCompleteHeadAndTailLines(t *testing.T) {
|
||||
lines := []string{"File: huge.go (500 lines)", "1 | package main"}
|
||||
for index := 2; index < 500; index++ {
|
||||
lines = append(lines, fmt.Sprintf("%d | value := %d", index, index))
|
||||
}
|
||||
lines = append(lines, "500 | // END_MARK")
|
||||
got := budgetSemanticOutput(strings.Join(lines, "\n"), outputCategoryFile, semanticTestBudget())
|
||||
if !got.truncated || !strings.Contains(got.text, "File: huge.go") || !strings.Contains(got.text, "END_MARK") {
|
||||
t.Fatalf("file policy lost boundaries: %#v", got)
|
||||
}
|
||||
if !utf8.ValidString(got.text) || strings.Contains(got.text, "value := 49\ufffd") {
|
||||
t.Fatalf("file policy produced invalid or split text: %q", got.text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearchOutputPolicyPreservesMultiFileCoverage(t *testing.T) {
|
||||
var lines []string
|
||||
for _, file := range []string{"a.go", "b.go", "c.go", "d.go"} {
|
||||
for line := 1; line <= 30; line++ {
|
||||
lines = append(lines, fmt.Sprintf("%s:%d: match value %d", file, line, line))
|
||||
}
|
||||
}
|
||||
lines = append(lines, "120 matches found")
|
||||
got := budgetSemanticOutput(strings.Join(lines, "\n"), outputCategorySearch, semanticTestBudget())
|
||||
for _, want := range []string{"a.go:", "b.go:", "c.go:", "d.go:", "120 matches"} {
|
||||
if !strings.Contains(got.text, want) {
|
||||
t.Fatalf("search policy missing %q:\n%s", want, got.text)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearchOutputPolicyPreservesWindowsPathCoverage(t *testing.T) {
|
||||
var lines []string
|
||||
for _, file := range []string{"a.go", "b.go", "c.go", "d.go"} {
|
||||
for line := 1; line <= 30; line++ {
|
||||
lines = append(lines, fmt.Sprintf(`C:\workspace\%s:%d: match value %d`, file, line, line))
|
||||
}
|
||||
}
|
||||
lines = append(lines, "120 matches found")
|
||||
got := budgetSemanticOutput(strings.Join(lines, "\n"), outputCategorySearch, semanticTestBudget())
|
||||
for _, want := range []string{`C:\workspace\a.go:`, `C:\workspace\b.go:`, `C:\workspace\c.go:`, `C:\workspace\d.go:`, "120 matches"} {
|
||||
if !strings.Contains(got.text, want) {
|
||||
t.Fatalf("search policy missing %q:\n%s", want, got.text)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearchResultFileIgnoresNumericMatchContent(t *testing.T) {
|
||||
tests := []struct {
|
||||
line string
|
||||
want string
|
||||
}{
|
||||
{line: `C:\workspace\a.go:12: error code:123: failed`, want: `C:\workspace\a.go`},
|
||||
{line: `/workspace/a.go:12: error code:123: failed`, want: `/workspace/a.go`},
|
||||
{line: `/workspace/name:with-colon.go:12: match`, want: `/workspace/name:with-colon.go`},
|
||||
}
|
||||
for _, test := range tests {
|
||||
if got := searchResultFile(test.line); got != test.want {
|
||||
t.Errorf("searchResultFile(%q) = %q, want %q", test.line, got, test.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessOutputPolicyCollapsesRepetitiveLogsAndKeepsDiagnostics(t *testing.T) {
|
||||
input := "starting server\n" + strings.Repeat("polling...\n", 300) + "WARNING: queue slow\nERROR: request failed\nshutdown complete"
|
||||
got := budgetSemanticOutput(input, outputCategoryProcess, semanticTestBudget())
|
||||
for _, want := range []string{"starting server", "repeated", "WARNING", "ERROR", "shutdown complete"} {
|
||||
if !strings.Contains(got.text, want) {
|
||||
t.Fatalf("process policy missing %q:\n%s", want, got.text)
|
||||
}
|
||||
}
|
||||
if strings.Count(got.text, "polling...") > 1 {
|
||||
t.Fatalf("repetitive line was not collapsed:\n%s", got.text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTestOutputPolicyKeepsFailureAndFinalSummary(t *testing.T) {
|
||||
input := "=== RUN TestSuite\n" + strings.Repeat("ok progress\n", 300) +
|
||||
"--- FAIL: TestImportant (0.01s)\n thing_test.go:42: expected 7, got 9\nFAIL\nexit status 1"
|
||||
got := budgetSemanticOutput(input, outputCategoryTest, semanticTestBudget())
|
||||
for _, want := range []string{"TestImportant", "expected 7", "FAIL", "exit status 1"} {
|
||||
if !strings.Contains(got.text, want) {
|
||||
t.Fatalf("test policy missing %q:\n%s", want, got.text)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiffOutputPolicyKeepsFilesAndCompleteHunks(t *testing.T) {
|
||||
var diff strings.Builder
|
||||
for file := 1; file <= 4; file++ {
|
||||
fmt.Fprintf(&diff, "diff --git a/f%d.go b/f%d.go\n--- a/f%d.go\n+++ b/f%d.go\n", file, file, file, file)
|
||||
for hunk := 1; hunk <= 8; hunk++ {
|
||||
fmt.Fprintf(&diff, "@@ -%d,2 +%d,2 @@\n-old%d_%d\n+new%d_%d\n", hunk, hunk, file, hunk, file, hunk)
|
||||
}
|
||||
}
|
||||
got := budgetSemanticOutput(diff.String(), outputCategoryDiff, outputBudget{maxEstimatedTokens: 220, hardMaxBytes: 1400})
|
||||
for file := 1; file <= 4; file++ {
|
||||
if !strings.Contains(got.text, fmt.Sprintf("diff --git a/f%d.go", file)) {
|
||||
t.Fatalf("diff policy lost file %d header:\n%s", file, got.text)
|
||||
}
|
||||
}
|
||||
for _, line := range strings.Split(got.text, "\n") {
|
||||
if strings.HasPrefix(line, "@@ ") {
|
||||
// Each retained hunk is an indivisible three-line unit.
|
||||
if !strings.Contains(got.text, line+"\n-old") {
|
||||
t.Fatalf("hunk header was sliced from its body: %q", line)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestWorkerOutputPolicyKeepsStatusErrorsAndConclusion(t *testing.T) {
|
||||
input := "session_id: child_1\nstatus: error\n" + strings.Repeat("worker progress\n", 200) +
|
||||
"tools executed: grep, read_file\nchanged files: a.go, b.go\nerror: verification failed\nConclusion: fix the nil check before merging."
|
||||
got := budgetSemanticOutput(input, outputCategoryWorker, semanticTestBudget())
|
||||
for _, want := range []string{"session_id", "status: error", "tools executed", "changed files", "verification failed", "Conclusion"} {
|
||||
if !strings.Contains(got.text, want) {
|
||||
t.Fatalf("worker policy missing %q:\n%s", want, got.text)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSemanticPoliciesAreDeterministic(t *testing.T) {
|
||||
input := "start\n" + strings.Repeat("tick\n", 500) + "ERROR: boom\nend"
|
||||
want := budgetSemanticOutput(input, outputCategoryProcess, semanticTestBudget())
|
||||
for iteration := 0; iteration < 20; iteration++ {
|
||||
if got := budgetSemanticOutput(input, outputCategoryProcess, semanticTestBudget()); got != want {
|
||||
t.Fatalf("iteration %d differs:\n got %#v\nwant %#v", iteration, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -18,6 +18,8 @@ type readFileTool struct {
|
||||
scope PathScope
|
||||
}
|
||||
|
||||
func (readFileTool) outputCategory(map[string]any) outputCategory { return outputCategoryFile }
|
||||
|
||||
func NewReadFileTool(workspaceRoot string) Tool {
|
||||
return NewScopedReadFileTool(workspaceRoot, nil)
|
||||
}
|
||||
@@ -50,10 +52,14 @@ func NewScopedReadFileTool(workspaceRoot string, scope PathScope) Tool {
|
||||
}
|
||||
|
||||
func (tool readFileTool) Run(ctx context.Context, args map[string]any) Result {
|
||||
return tool.RunWithOptions(ctx, args, RunOptions{})
|
||||
return tool.run(args, RunOptions{}, true)
|
||||
}
|
||||
|
||||
func (tool readFileTool) RunWithOptions(_ context.Context, args map[string]any, options RunOptions) Result {
|
||||
return tool.run(args, options, false)
|
||||
}
|
||||
|
||||
func (tool readFileTool) run(args map[string]any, options RunOptions, directBudget bool) Result {
|
||||
requestedPath, err := aliasedStringArg(args, []string{"path", "file", "file_path", "filepath", "filename"}, "", true, false)
|
||||
if err != nil {
|
||||
return errorResult("Error: Invalid arguments for read_file: " + err.Error())
|
||||
@@ -86,10 +92,14 @@ func (tool readFileTool) RunWithOptions(_ context.Context, args map[string]any,
|
||||
// not the authoritative content hash.
|
||||
options.FileTracker.RecordHash(absolutePath, stats.hash, stats.info)
|
||||
|
||||
return renderReadFileRange(absolutePath, relativePath, stats.lines, startLine, endLine, maxLines)
|
||||
maxBytes := 0
|
||||
if directBudget {
|
||||
maxBytes = readOutputBudgetBytes
|
||||
}
|
||||
return renderReadFileRange(absolutePath, relativePath, stats.lines, startLine, endLine, maxLines, maxBytes)
|
||||
}
|
||||
|
||||
func renderReadFileRange(absolutePath string, relativePath string, total int, startLine int, endLine int, maxLines int) Result {
|
||||
func renderReadFileRange(absolutePath string, relativePath string, total int, startLine int, endLine int, maxLines int, maxBytes int) Result {
|
||||
if startLine > total {
|
||||
return okResult(fmt.Sprintf("File: %s\n(start_line %d is past the end of the file, which has %d lines)", relativePath, startLine, total))
|
||||
}
|
||||
@@ -114,7 +124,16 @@ func renderReadFileRange(absolutePath string, relativePath string, total int, st
|
||||
header = fmt.Sprintf("File: %s (lines %d-%d of %d)", relativePath, startLine, lastLine, total)
|
||||
}
|
||||
|
||||
budgetedOutput := newOutputBudgetBuilder(readOutputBudgetBytes, "use start_line/end_line or max_lines to continue with a smaller range")
|
||||
// The shared registry boundary applies the file policy after redaction. Its
|
||||
// input still needs a bounded capture: retain the requested range's head and
|
||||
// tail instead of materializing an arbitrarily large file in memory. Direct
|
||||
// Tool.Run calls preserve the legacy prefix-only byte budget.
|
||||
var budgetedOutput *outputBudgetBuilder
|
||||
if maxBytes > 0 {
|
||||
budgetedOutput = newOutputBudgetBuilder(maxBytes, "use start_line/end_line or max_lines to continue with a smaller range")
|
||||
} else {
|
||||
budgetedOutput = newHeadTailOutputBudgetBuilder(readOutputBudgetBytes, "use start_line/end_line or max_lines to continue with a smaller range")
|
||||
}
|
||||
budgetedOutput.WriteString(header)
|
||||
budgetedOutput.WriteString("\n\n")
|
||||
if err := appendReadFileRange(budgetedOutput, absolutePath, startLine, selectedLines, width); err != nil {
|
||||
@@ -128,10 +147,19 @@ func renderReadFileRange(absolutePath string, relativePath string, total int, st
|
||||
}
|
||||
|
||||
budgeted := budgetedOutput.Result()
|
||||
meta := outputBudgetMeta(budgeted)
|
||||
if budgeted.Truncated {
|
||||
meta := map[string]string{}
|
||||
if maxBytes > 0 || budgeted.Truncated {
|
||||
for key, value := range outputBudgetMeta(budgeted) {
|
||||
meta[key] = value
|
||||
}
|
||||
}
|
||||
if truncated || budgeted.Truncated {
|
||||
meta["truncated"] = "true"
|
||||
meta["truncation_reason"] = "byte_budget"
|
||||
if budgeted.Truncated {
|
||||
meta["truncation_reason"] = "byte_budget"
|
||||
} else {
|
||||
meta["truncation_reason"] = "max_lines"
|
||||
}
|
||||
}
|
||||
return Result{
|
||||
Status: StatusOK,
|
||||
|
||||
@@ -16,6 +16,10 @@ type readMinifiedFileTool struct {
|
||||
scope PathScope
|
||||
}
|
||||
|
||||
func (readMinifiedFileTool) outputCategory(map[string]any) outputCategory {
|
||||
return outputCategoryFile
|
||||
}
|
||||
|
||||
func NewReadMinifiedFileTool(workspaceRoot string) Tool {
|
||||
return NewScopedReadMinifiedFileTool(workspaceRoot, nil)
|
||||
}
|
||||
@@ -42,10 +46,14 @@ func NewScopedReadMinifiedFileTool(workspaceRoot string, scope PathScope) Tool {
|
||||
}
|
||||
|
||||
func (tool readMinifiedFileTool) Run(ctx context.Context, args map[string]any) Result {
|
||||
return tool.RunWithOptions(ctx, args, RunOptions{})
|
||||
return tool.run(args, RunOptions{}, true)
|
||||
}
|
||||
|
||||
func (tool readMinifiedFileTool) RunWithOptions(_ context.Context, args map[string]any, options RunOptions) Result {
|
||||
return tool.run(args, options, false)
|
||||
}
|
||||
|
||||
func (tool readMinifiedFileTool) run(args map[string]any, options RunOptions, directBudget bool) Result {
|
||||
requestedPath, err := aliasedStringArg(args, []string{"path", "file", "file_path", "filepath", "filename"}, "", true, false)
|
||||
if err != nil {
|
||||
return errorResult("Error: Invalid arguments for read_minified_file: " + err.Error())
|
||||
@@ -92,23 +100,21 @@ func (tool readMinifiedFileTool) RunWithOptions(_ context.Context, args map[stri
|
||||
savedTokens = estimatedTokensFromBytes(savedBytes)
|
||||
}
|
||||
output := header + "\n\n" + result.Content
|
||||
budgeted := applyOutputBudget(output, readOutputBudgetBytes, "use read_file with start_line/end_line or max_lines for a smaller exact range")
|
||||
meta := outputBudgetMeta(budgeted)
|
||||
meta := map[string]string{}
|
||||
meta["path"] = relativePath
|
||||
meta["mode"] = result.Language
|
||||
meta["compacted"] = strconv.FormatBool(result.Applied)
|
||||
meta["raw_bytes"] = strconv.Itoa(rawBytes)
|
||||
meta["compact_bytes"] = strconv.Itoa(compactBytes)
|
||||
meta["emitted_bytes"] = strconv.Itoa(budgeted.EmittedBytes)
|
||||
meta["emitted_bytes"] = strconv.Itoa(len(output))
|
||||
meta["raw_lines"] = strconv.Itoa(rawLines)
|
||||
meta["emitted_lines"] = strconv.Itoa(minLines)
|
||||
meta["estimated_tokens_saved"] = strconv.Itoa(savedTokens)
|
||||
if budgeted.Truncated {
|
||||
meta["truncated"] = "true"
|
||||
meta["truncation_reason"] = "byte_budget"
|
||||
toolResult := Result{Status: StatusOK, Output: output, Meta: meta}
|
||||
if directBudget {
|
||||
return applyLegacyByteBudgetToResult(toolResult, readOutputBudgetBytes, "use read_file with start_line/end_line or max_lines for a smaller exact range")
|
||||
}
|
||||
|
||||
return Result{Status: StatusOK, Output: budgeted.Output, Truncated: budgeted.Truncated, Meta: meta}
|
||||
return toolResult
|
||||
}
|
||||
|
||||
// lineCount reports the number of newline-separated lines in s (an empty string
|
||||
|
||||
@@ -126,20 +126,25 @@ func (registry *Registry) RunWithOptions(ctx context.Context, name string, args
|
||||
// args/paths) are redacted at the boundary just like tool output. The output
|
||||
// ceiling runs after the scrub so the transcript and the spill file agree on
|
||||
// what was hidden.
|
||||
ceilingExempt := false
|
||||
selfManagedOutput := false
|
||||
var tool Tool
|
||||
var ok bool
|
||||
defer func() {
|
||||
result = scrubResultSecrets(result)
|
||||
if !ceilingExempt {
|
||||
if selfManagedOutput {
|
||||
result = applySelfManagedOutputBudget(tool, name, args, result)
|
||||
} else {
|
||||
result = applyRegistryOutputBudget(tool, name, args, result)
|
||||
result = enforceOutputCeiling(name, result)
|
||||
}
|
||||
}()
|
||||
|
||||
tool, ok := registry.Get(name)
|
||||
tool, ok = registry.Get(name)
|
||||
if !ok {
|
||||
return errorResult(`Error: Unknown tool "` + name + `".`)
|
||||
}
|
||||
if _, ok := tool.(selfBudgeting); ok {
|
||||
ceilingExempt = true
|
||||
selfManagedOutput = true
|
||||
}
|
||||
if rejecter, ok := tool.(PrePermissionRejecter); ok {
|
||||
if res, rejected := rejecter.RejectBeforePermission(args); rejected {
|
||||
|
||||
@@ -117,6 +117,25 @@ func WriteNDJSON(w io.Writer, t *TurnTrace) error {
|
||||
return err
|
||||
}
|
||||
}
|
||||
// Output budget events stay in tool-result emission order. Sorting them by
|
||||
// tool/category would destroy correlation with concurrent calls whose results
|
||||
// are deliberately consumed in original call order.
|
||||
for _, event := range t.OutputBudgets {
|
||||
if err := enc.Encode(map[string]any{
|
||||
"type": "output_budget",
|
||||
"tool": event.Tool,
|
||||
"category": event.Category,
|
||||
"original_bytes": event.OriginalBytes,
|
||||
"retained_bytes": event.RetainedBytes,
|
||||
"estimated_original_tokens": event.EstimatedOriginalTokens,
|
||||
"estimated_retained_tokens": event.EstimatedRetainedTokens,
|
||||
"truncated": event.Truncated,
|
||||
"reason": event.Reason,
|
||||
"spill_created": event.SpillCreated,
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -177,6 +196,15 @@ func WriteText(w io.Writer, t *TurnTrace) error {
|
||||
for _, c := range counters {
|
||||
write(" %-22s %d\n", c.Name, c.Value)
|
||||
}
|
||||
if len(t.OutputBudgets) > 0 {
|
||||
write("output budgets:\n")
|
||||
for _, event := range t.OutputBudgets {
|
||||
write(" tool=%s category=%s bytes=%d/%d tokens=%d/%d truncated=%t reason=%s spill=%t\n",
|
||||
event.Tool, event.Category, event.RetainedBytes, event.OriginalBytes,
|
||||
event.EstimatedRetainedTokens, event.EstimatedOriginalTokens,
|
||||
event.Truncated, event.Reason, event.SpillCreated)
|
||||
}
|
||||
}
|
||||
return firstErr
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,100 @@
|
||||
package trace
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestOutputBudgetTraceRoundTripContainsNoOutput(t *testing.T) {
|
||||
recorder := NewRecorder("session", "run", "")
|
||||
recorder.Start()
|
||||
event := OutputBudgetEvent{
|
||||
Tool: "grep",
|
||||
Category: "search",
|
||||
OriginalBytes: 10000,
|
||||
RetainedBytes: 1000,
|
||||
EstimatedOriginalTokens: 2500,
|
||||
EstimatedRetainedTokens: 250,
|
||||
Truncated: true,
|
||||
Reason: "semantic_search_budget",
|
||||
SpillCreated: true,
|
||||
}
|
||||
recorder.EmitOutputBudget(event)
|
||||
|
||||
var encoded bytes.Buffer
|
||||
if err := WriteNDJSON(&encoded, recorder.Finish()); err != nil {
|
||||
t.Fatalf("WriteNDJSON: %v", err)
|
||||
}
|
||||
const secretOutput = "secret output body"
|
||||
if strings.Contains(encoded.String(), secretOutput) {
|
||||
t.Fatal("output_budget trace contains raw secret output")
|
||||
}
|
||||
allowedKeys := map[string]bool{
|
||||
"type": true, "tool": true, "category": true, "original_bytes": true, "retained_bytes": true,
|
||||
"estimated_original_tokens": true, "estimated_retained_tokens": true, "truncated": true,
|
||||
"reason": true, "spill_created": true,
|
||||
}
|
||||
findUndocumentedKey := func(record map[string]any) string {
|
||||
for key := range record {
|
||||
if !allowedKeys[key] {
|
||||
return key
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
decoder := json.NewDecoder(strings.NewReader(encoded.String()))
|
||||
var outputBudgetRecord map[string]any
|
||||
for {
|
||||
var record map[string]any
|
||||
if err := decoder.Decode(&record); err != nil {
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
t.Fatalf("decode trace record: %v", err)
|
||||
}
|
||||
if record["type"] == "output_budget" {
|
||||
outputBudgetRecord = record
|
||||
break
|
||||
}
|
||||
}
|
||||
if outputBudgetRecord == nil {
|
||||
t.Fatal("missing output_budget trace record")
|
||||
}
|
||||
if key := findUndocumentedKey(outputBudgetRecord); key != "" {
|
||||
t.Fatalf("output_budget trace contains undocumented key %q", key)
|
||||
}
|
||||
for key := range allowedKeys {
|
||||
if _, exists := outputBudgetRecord[key]; !exists {
|
||||
t.Fatalf("output_budget trace missing documented key %q", key)
|
||||
}
|
||||
}
|
||||
recordJSON, err := json.Marshal(outputBudgetRecord)
|
||||
if err != nil {
|
||||
t.Fatalf("encode output_budget trace record: %v", err)
|
||||
}
|
||||
var serializedEvent OutputBudgetEvent
|
||||
if err := json.Unmarshal(recordJSON, &serializedEvent); err != nil {
|
||||
t.Fatalf("decode output_budget trace event: %v", err)
|
||||
}
|
||||
if serializedEvent != event {
|
||||
t.Fatalf("serialized output_budget event = %#v, want %#v", serializedEvent, event)
|
||||
}
|
||||
contaminated := make(map[string]any, len(outputBudgetRecord)+1)
|
||||
for key, value := range outputBudgetRecord {
|
||||
contaminated[key] = value
|
||||
}
|
||||
contaminated["output"] = secretOutput
|
||||
if key := findUndocumentedKey(contaminated); key == "" {
|
||||
t.Fatal("trace key validation would allow raw secret output")
|
||||
}
|
||||
parsed, err := ReadNDJSON(strings.NewReader(encoded.String()))
|
||||
if err != nil {
|
||||
t.Fatalf("ReadNDJSON: %v", err)
|
||||
}
|
||||
if len(parsed.OutputBudgets) != 1 || parsed.OutputBudgets[0] != event {
|
||||
t.Fatalf("unexpected round trip: %#v", parsed.OutputBudgets)
|
||||
}
|
||||
}
|
||||
+21
-1
@@ -116,6 +116,21 @@ func ReadNDJSON(r io.Reader) (*TurnTrace, error) {
|
||||
SchemaHash: stringField(obj, "schema"),
|
||||
CompletePrefixHash: stringField(obj, "complete_prefix"),
|
||||
})
|
||||
case "output_budget":
|
||||
if !sawTraceHeader {
|
||||
return nil, errors.New("parse trace: not a valid NDJSON trace (no type:trace header)")
|
||||
}
|
||||
t.OutputBudgets = append(t.OutputBudgets, OutputBudgetEvent{
|
||||
Tool: stringField(obj, "tool"),
|
||||
Category: stringField(obj, "category"),
|
||||
OriginalBytes: int(parseInt64(obj["original_bytes"])),
|
||||
RetainedBytes: int(parseInt64(obj["retained_bytes"])),
|
||||
EstimatedOriginalTokens: int(parseInt64(obj["estimated_original_tokens"])),
|
||||
EstimatedRetainedTokens: int(parseInt64(obj["estimated_retained_tokens"])),
|
||||
Truncated: boolField(obj, "truncated"),
|
||||
Reason: stringField(obj, "reason"),
|
||||
SpillCreated: boolField(obj, "spill_created"),
|
||||
})
|
||||
default:
|
||||
// Unknown event type: tolerate (forward-compat) but only after a
|
||||
// header has been seen.
|
||||
@@ -136,12 +151,17 @@ func ReadNDJSON(r io.Reader) (*TurnTrace, error) {
|
||||
if !sawTraceHeader {
|
||||
return nil, errors.New("parse trace: non-empty input had no type:trace header")
|
||||
}
|
||||
if len(t.Spans) == 0 && len(t.Counters) == 0 && len(t.PrefixHashes) == 0 {
|
||||
if len(t.Spans) == 0 && len(t.Counters) == 0 && len(t.PrefixHashes) == 0 && len(t.OutputBudgets) == 0 {
|
||||
return nil, errors.New("parse trace: header present but no spans, counters, or prefix hashes recovered (corrupt or truncated)")
|
||||
}
|
||||
return t, nil
|
||||
}
|
||||
|
||||
func boolField(obj map[string]any, key string) bool {
|
||||
value, _ := obj[key].(bool)
|
||||
return value
|
||||
}
|
||||
|
||||
// stringField returns obj[key] as a string, or "" if the key is missing or
|
||||
// the value is not a string. JSON-marshaled trace events always emit
|
||||
// string fields as JSON strings, so the type assertion is the right
|
||||
|
||||
@@ -188,6 +188,21 @@ func (r *Recorder) EmitPrefixHash(p PrefixHash) {
|
||||
r.tr.PrefixHashes = append(r.tr.PrefixHashes, p)
|
||||
}
|
||||
|
||||
// EmitOutputBudget records one content-free output budgeting decision. Calls
|
||||
// are retained in transcript emission order, including when the underlying
|
||||
// tools executed concurrently.
|
||||
func (r *Recorder) EmitOutputBudget(event OutputBudgetEvent) {
|
||||
if r == nil {
|
||||
return
|
||||
}
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if r.finished {
|
||||
return
|
||||
}
|
||||
r.tr.OutputBudgets = append(r.tr.OutputBudgets, event)
|
||||
}
|
||||
|
||||
// Finish stamps CompletedAt, derives each span's parent (by interval
|
||||
// containment) and exclusive time, and returns a snapshot of the trace. Calling
|
||||
// Finish more than once returns the same snapshot.
|
||||
@@ -207,6 +222,7 @@ func (r *Recorder) Finish() *TurnTrace {
|
||||
snap.Spans = append([]Span(nil), r.tr.Spans...)
|
||||
snap.Counters = append([]Counter(nil), r.tr.Counters...)
|
||||
snap.PrefixHashes = append([]PrefixHash(nil), r.tr.PrefixHashes...)
|
||||
snap.OutputBudgets = append([]OutputBudgetEvent(nil), r.tr.OutputBudgets...)
|
||||
return &snap
|
||||
}
|
||||
|
||||
|
||||
+27
-11
@@ -79,20 +79,36 @@ type Counter struct {
|
||||
Value int64 `json:"value"`
|
||||
}
|
||||
|
||||
// OutputBudgetEvent is compact, content-free metadata for one tool result's
|
||||
// output budgeting decision. It deliberately carries no output text, paths, or
|
||||
// arguments so tracing cannot become a secret-bearing side channel.
|
||||
type OutputBudgetEvent struct {
|
||||
Tool string `json:"tool"`
|
||||
Category string `json:"category"`
|
||||
OriginalBytes int `json:"original_bytes"`
|
||||
RetainedBytes int `json:"retained_bytes"`
|
||||
EstimatedOriginalTokens int `json:"estimated_original_tokens"`
|
||||
EstimatedRetainedTokens int `json:"estimated_retained_tokens"`
|
||||
Truncated bool `json:"truncated"`
|
||||
Reason string `json:"reason,omitempty"`
|
||||
SpillCreated bool `json:"spill_created"`
|
||||
}
|
||||
|
||||
// TurnTrace is the finished record for one agent.Run. It is the value
|
||||
// emitters serialize; it is not mutated after Finish returns a snapshot.
|
||||
type TurnTrace struct {
|
||||
SessionID string `json:"session_id"`
|
||||
RunID string `json:"run_id"`
|
||||
Profile string `json:"profile,omitempty"`
|
||||
StartedAt time.Time `json:"started_at"`
|
||||
FirstVisibleEventAt time.Time `json:"first_visible_event_at,omitempty"`
|
||||
FirstUsefulActionAt time.Time `json:"first_useful_action_at,omitempty"`
|
||||
FirstTokenAt time.Time `json:"first_token_at,omitempty"`
|
||||
CompletedAt time.Time `json:"completed_at"`
|
||||
Spans []Span `json:"spans"`
|
||||
Counters []Counter `json:"counters"`
|
||||
PrefixHashes []PrefixHash `json:"prefix_hashes,omitempty"`
|
||||
SessionID string `json:"session_id"`
|
||||
RunID string `json:"run_id"`
|
||||
Profile string `json:"profile,omitempty"`
|
||||
StartedAt time.Time `json:"started_at"`
|
||||
FirstVisibleEventAt time.Time `json:"first_visible_event_at,omitempty"`
|
||||
FirstUsefulActionAt time.Time `json:"first_useful_action_at,omitempty"`
|
||||
FirstTokenAt time.Time `json:"first_token_at,omitempty"`
|
||||
CompletedAt time.Time `json:"completed_at"`
|
||||
Spans []Span `json:"spans"`
|
||||
Counters []Counter `json:"counters"`
|
||||
PrefixHashes []PrefixHash `json:"prefix_hashes,omitempty"`
|
||||
OutputBudgets []OutputBudgetEvent `json:"output_budgets,omitempty"`
|
||||
}
|
||||
|
||||
// PrefixHash is one fingerprint of the prompt prefix emitted by an agent run.
|
||||
|
||||
Reference in new issue
Block a user