From 14f8a1ab90419b69d0e25f91fa252102147fd2b8 Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 8 Aug 2026 12:36:39 +0530 Subject: [PATCH 1/3] Centralize agent context planning --- internal/agent/compaction.go | 12 +- internal/agent/context_measurement.go | 36 ++++++ internal/agent/context_measurement_test.go | 3 + internal/agent/context_planner.go | 115 ++++++++++++++++++ internal/agent/context_planner_test.go | 135 +++++++++++++++++++++ internal/agent/loop.go | 91 ++++++-------- internal/agent/loop_test.go | 12 ++ internal/trace/emit.go | 1 + internal/trace/parse.go | 1 + internal/trace/trace.go | 1 + internal/trace/trace_test.go | 9 +- 11 files changed, 351 insertions(+), 65 deletions(-) create mode 100644 internal/agent/context_planner.go create mode 100644 internal/agent/context_planner_test.go diff --git a/internal/agent/compaction.go b/internal/agent/compaction.go index 3a19dcdaf..981ebc187 100644 --- a/internal/agent/compaction.go +++ b/internal/agent/compaction.go @@ -624,13 +624,11 @@ func summarizeWithFallback(ctx context.Context, provider Provider, messages []ze // summarizeMessagesOnce performs a single tool-less summarization call. func summarizeMessagesOnce(ctx context.Context, provider Provider, messages []zeroruntime.Message, onUsage func(Usage)) (string, error) { - request := zeroruntime.CompletionRequest{ - Messages: []zeroruntime.Message{ - {Role: zeroruntime.MessageRoleSystem, Content: CompactionSummaryInstructions}, - {Role: zeroruntime.MessageRoleUser, Content: "Summarize this conversation:\n\n" + renderTranscript(messages)}, - }, - // No tools: this is a plain text summarization call. - } + planner := newContextPlanner(contextPlannerConfig{}) + request := planner.Plan([]zeroruntime.Message{ + {Role: zeroruntime.MessageRoleSystem, Content: CompactionSummaryInstructions}, + {Role: zeroruntime.MessageRoleUser, Content: "Summarize this conversation:\n\n" + renderTranscript(messages)}, + }, nil, "").Request stream, err := provider.StreamCompletion(ctx, request) if err != nil { return "", err diff --git a/internal/agent/context_measurement.go b/internal/agent/context_measurement.go index 98a537a00..b46903189 100644 --- a/internal/agent/context_measurement.go +++ b/internal/agent/context_measurement.go @@ -23,6 +23,24 @@ type ContextBreakdown struct { TotalTokens int // SystemTokens + ToolTokens + MessageTokens ContextWindow int // model context window; 0 when unknown UsedFraction float64 // TotalTokens / ContextWindow; 0 when window unknown + // Blocks explain why each model-visible category is present without exposing + // its content. The planner currently preserves all three categories. + Blocks []ContextBlock + // CompletePrefixHash identifies the provider-visible system+tool prefix. + CompletePrefixHash string + // PrefixInvalidationReason is "initial", "unchanged", or a stable comma- + // separated list of changed prefix categories. + PrefixInvalidationReason string +} + +// ContextBlock is content-free evidence for one model-visible context category. +type ContextBlock struct { + Kind string + Tokens int + CacheClass string + Authority string + Recoverable bool + Reason string } // MeasureContext estimates the per-category token footprint of a request: the @@ -41,6 +59,24 @@ func MeasureContext(messages []zeroruntime.Message, tools []zeroruntime.ToolDefi ContextWindow: contextWindow, } breakdown.TotalTokens = breakdown.SystemTokens + breakdown.ToolTokens + breakdown.MessageTokens + if breakdown.SystemTokens > 0 { + breakdown.Blocks = append(breakdown.Blocks, ContextBlock{ + Kind: "system", Tokens: breakdown.SystemTokens, CacheClass: "stable_prefix", + Authority: "required", Reason: "instructions, workspace policy, and safety policy", + }) + } + if breakdown.ToolTokens > 0 { + breakdown.Blocks = append(breakdown.Blocks, ContextBlock{ + Kind: "tools", Tokens: breakdown.ToolTokens, CacheClass: "stable_prefix", + Authority: "capability", Reason: "currently permitted and selected tool definitions", + }) + } + if breakdown.MessageTokens > 0 { + breakdown.Blocks = append(breakdown.Blocks, ContextBlock{ + Kind: "conversation", Tokens: breakdown.MessageTokens, CacheClass: "dynamic_tail", + Authority: "transcript", Reason: "valid provider replay history and current task context", + }) + } if contextWindow > 0 { breakdown.UsedFraction = float64(breakdown.TotalTokens) / float64(contextWindow) } diff --git a/internal/agent/context_measurement_test.go b/internal/agent/context_measurement_test.go index 7ad1b347f..c46520fac 100644 --- a/internal/agent/context_measurement_test.go +++ b/internal/agent/context_measurement_test.go @@ -32,6 +32,9 @@ func TestMeasureContextSplitsByCategory(t *testing.T) { if breakdown.UsedFraction != wantFraction { t.Fatalf("UsedFraction = %v, want %v", breakdown.UsedFraction, wantFraction) } + if len(breakdown.Blocks) != 3 { + t.Fatalf("Blocks = %#v, want system/tools/conversation", breakdown.Blocks) + } } func TestMeasureContextUnknownWindowHasZeroFraction(t *testing.T) { diff --git a/internal/agent/context_planner.go b/internal/agent/context_planner.go new file mode 100644 index 000000000..567e28ea4 --- /dev/null +++ b/internal/agent/context_planner.go @@ -0,0 +1,115 @@ +package agent + +import ( + "strings" + + "github.com/Gitlawb/zero/internal/zeroruntime" +) + +type contextPlannerConfig struct { + contextWindow int + promptCacheKey string + promptParts systemPromptParts +} + +// contextPlanner is the single seam for constructing model-visible requests. +// Planning is deterministic and does not retrieve content, execute tools, or +// change permissions. The initial implementation deliberately preserves every +// message and tool definition while making composition and cache drift +// inspectable; later selection policies must continue to cross this seam. +type contextPlanner struct { + config contextPlannerConfig + previousPrefix prefixFingerprint + hasPrevious bool +} + +type contextPlan struct { + Request zeroruntime.CompletionRequest + Breakdown ContextBreakdown + PrefixFingerprint prefixFingerprint +} + +func newContextPlanner(config contextPlannerConfig) *contextPlanner { + return &contextPlanner{config: config} +} + +// Plan returns a provider request snapshot plus content-free accounting. +// It intentionally performs no relevance filtering: preserving current model +// capability is the baseline contract for future planner policies. +func (planner *contextPlanner) Plan(messages []zeroruntime.Message, toolDefs []zeroruntime.ToolDefinition, reasoningEffort string) contextPlan { + request := zeroruntime.CompletionRequest{ + Messages: copyMessages(messages), + Tools: toolDefs, + ReasoningEffort: reasoningEffort, + PromptCacheKey: planner.config.promptCacheKey, + } + parts := planner.config.promptParts + if parts.prompt == "" { + parts.prompt = leadingSystemContent(request.Messages) + } + fingerprint := computePrefixFingerprint(buildPromptSubstringsFromParts(parts, request.Tools)) + breakdown := MeasureContext(request.Messages, request.Tools, planner.config.contextWindow) + breakdown.CompletePrefixHash = fingerprint.CompletePrefixHash + var previous *prefixFingerprint + if planner.hasPrevious { + previous = &planner.previousPrefix + } + breakdown.PrefixInvalidationReason = explainPrefixChange(previous, fingerprint) + planner.previousPrefix = fingerprint + planner.hasPrevious = true + return contextPlan{ + Request: request, + Breakdown: breakdown, + PrefixFingerprint: fingerprint, + } +} + +func leadingSystemContent(messages []zeroruntime.Message) string { + contents := make([]string, 0, 1) + for _, message := range messages { + if message.Role != zeroruntime.MessageRoleSystem { + break + } + contents = append(contents, message.Content) + } + return strings.Join(contents, "\n\n") +} + +func explainPrefixChange(previous *prefixFingerprint, current prefixFingerprint) string { + if previous == nil { + return "initial" + } + if previous.CompletePrefixHash == current.CompletePrefixHash { + return "unchanged" + } + reasons := make([]string, 0, 3) + if previous.SystemPromptHash != current.SystemPromptHash { + parts := make([]string, 0, 4) + if previous.BaseInstructionsHash != current.BaseInstructionsHash { + parts = append(parts, "base_instructions") + } + if previous.ConfirmationPolicyHash != current.ConfirmationPolicyHash { + parts = append(parts, "confirmation_policy") + } + if previous.ProjectContextHash != current.ProjectContextHash { + parts = append(parts, "project_context") + } + if previous.SkillsHash != current.SkillsHash { + parts = append(parts, "skills") + } + if len(parts) == 0 { + parts = append(parts, "system_prompt") + } + reasons = append(reasons, parts...) + } + if previous.ToolsHash != current.ToolsHash { + reasons = append(reasons, "tools") + } + if previous.SchemaHash != current.SchemaHash { + reasons = append(reasons, "schema") + } + if len(reasons) == 0 { + return "prefix_changed" + } + return strings.Join(reasons, ",") +} diff --git a/internal/agent/context_planner_test.go b/internal/agent/context_planner_test.go new file mode 100644 index 000000000..e47ae85d7 --- /dev/null +++ b/internal/agent/context_planner_test.go @@ -0,0 +1,135 @@ +package agent + +import ( + "reflect" + "testing" + + "github.com/Gitlawb/zero/internal/zeroruntime" +) + +func TestContextPlannerPreservesProviderRequest(t *testing.T) { + messages := []zeroruntime.Message{ + {Role: zeroruntime.MessageRoleSystem, Content: "system"}, + {Role: zeroruntime.MessageRoleUser, Content: "inspect this", Images: []zeroruntime.ImageBlock{{MediaType: "image/png", Data: []byte{1, 2, 3}}}}, + {Role: zeroruntime.MessageRoleAssistant, Content: "working", ToolCalls: []zeroruntime.ToolCall{{ID: "call-1", Name: "read_file", Arguments: `{"path":"main.go"}`}}, Reasoning: []zeroruntime.ReasoningBlock{{Provider: "test", Type: "thinking", Signature: "sig"}}}, + } + toolDefs := []zeroruntime.ToolDefinition{{Name: "read_file", Description: "Read a file", Parameters: map[string]any{"type": "object"}}} + planner := newContextPlanner(contextPlannerConfig{ + contextWindow: 128_000, + promptCacheKey: "session-1", + promptParts: systemPromptParts{prompt: "system", baseInstructions: "system"}, + }) + + plan := planner.Plan(messages, toolDefs, "medium") + want := zeroruntime.CompletionRequest{ + Messages: copyMessages(messages), + Tools: toolDefs, + ReasoningEffort: "medium", + PromptCacheKey: "session-1", + } + if !reflect.DeepEqual(plan.Request, want) { + t.Fatalf("planned request changed provider input:\n got: %#v\nwant: %#v", plan.Request, want) + } + + // A plan owns its request snapshot; later caller mutation cannot change it. + messages[1].Content = "changed" + messages[1].Images[0].Data[0] = 9 + if plan.Request.Messages[1].Content != "inspect this" || plan.Request.Messages[1].Images[0].Data[0] != 1 { + t.Fatalf("planned messages alias caller state: %#v", plan.Request.Messages[1]) + } +} + +func TestContextPlannerReportsInspectableBlocks(t *testing.T) { + planner := newContextPlanner(contextPlannerConfig{contextWindow: 100_000}) + plan := planner.Plan([]zeroruntime.Message{ + {Role: zeroruntime.MessageRoleSystem, Content: filler(400)}, + {Role: zeroruntime.MessageRoleUser, Content: filler(200)}, + }, []zeroruntime.ToolDefinition{{Name: "read_file", Description: "Read a file"}}, "") + + if plan.Breakdown.PrefixInvalidationReason != "initial" || plan.Breakdown.CompletePrefixHash == "" { + t.Fatalf("prefix evidence = %#v", plan.Breakdown) + } + if len(plan.Breakdown.Blocks) != 3 { + t.Fatalf("blocks = %#v, want system/tools/conversation", plan.Breakdown.Blocks) + } + wantKinds := []string{"system", "tools", "conversation"} + total := 0 + for index, block := range plan.Breakdown.Blocks { + if block.Kind != wantKinds[index] || block.Tokens <= 0 || block.Reason == "" || block.Authority == "" || block.CacheClass == "" { + t.Fatalf("block %d = %#v", index, block) + } + total += block.Tokens + } + if total != plan.Breakdown.TotalTokens { + t.Fatalf("block tokens = %d, want total %d", total, plan.Breakdown.TotalTokens) + } +} + +func TestContextPlannerExplainsPrefixInvalidation(t *testing.T) { + planner := newContextPlanner(contextPlannerConfig{promptParts: systemPromptParts{ + prompt: "system", baseInstructions: "base", projectContext: "project", + }}) + messages := []zeroruntime.Message{{Role: zeroruntime.MessageRoleSystem, Content: "system"}, {Role: zeroruntime.MessageRoleUser, Content: "task"}} + tools := []zeroruntime.ToolDefinition{{Name: "read_file", Description: "Read", Parameters: map[string]any{"type": "object"}}} + + if got := planner.Plan(messages, tools, "").Breakdown.PrefixInvalidationReason; got != "initial" { + t.Fatalf("first reason = %q, want initial", got) + } + if got := planner.Plan(messages, tools, "").Breakdown.PrefixInvalidationReason; got != "unchanged" { + t.Fatalf("stable reason = %q, want unchanged", got) + } + tools[0].Description = "Read an exact file" + if got := planner.Plan(messages, tools, "").Breakdown.PrefixInvalidationReason; got != "tools" { + t.Fatalf("tool reason = %q, want tools", got) + } + tools[0].Parameters["required"] = []any{"path"} + if got := planner.Plan(messages, tools, "").Breakdown.PrefixInvalidationReason; got != "schema" { + t.Fatalf("schema reason = %q, want schema", got) + } + planner.config.promptParts.prompt = "system with changed project" + planner.config.promptParts.projectContext = "changed project" + if got := planner.Plan(messages, tools, "").Breakdown.PrefixInvalidationReason; got != "project_context" { + t.Fatalf("project reason = %q, want project_context", got) + } +} + +func TestExplainPrefixChangeNamesEveryChangedComponent(t *testing.T) { + baseline := prefixFingerprint{ + SystemPromptHash: "system", + BaseInstructionsHash: "base", + ConfirmationPolicyHash: "policy", + ProjectContextHash: "project", + SkillsHash: "skills", + ToolsHash: "tools", + SchemaHash: "schema", + CompletePrefixHash: "complete", + } + tests := []struct { + name string + mutate func(*prefixFingerprint) + want string + }{ + {name: "base instructions", mutate: func(value *prefixFingerprint) { + value.SystemPromptHash, value.BaseInstructionsHash = "changed", "changed" + }, want: "base_instructions"}, + {name: "confirmation policy", mutate: func(value *prefixFingerprint) { + value.SystemPromptHash, value.ConfirmationPolicyHash = "changed", "changed" + }, want: "confirmation_policy"}, + {name: "project context", mutate: func(value *prefixFingerprint) { + value.SystemPromptHash, value.ProjectContextHash = "changed", "changed" + }, want: "project_context"}, + {name: "skills", mutate: func(value *prefixFingerprint) { value.SystemPromptHash, value.SkillsHash = "changed", "changed" }, want: "skills"}, + {name: "unclassified system prompt", mutate: func(value *prefixFingerprint) { value.SystemPromptHash = "changed" }, want: "system_prompt"}, + {name: "tools and schema", mutate: func(value *prefixFingerprint) { value.ToolsHash, value.SchemaHash = "changed", "changed" }, want: "tools,schema"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + current := baseline + test.mutate(¤t) + current.CompletePrefixHash = "changed" + if got := explainPrefixChange(&baseline, current); got != test.want { + t.Fatalf("reason = %q, want %q", got, test.want) + } + }) + } +} diff --git a/internal/agent/loop.go b/internal/agent/loop.go index 4991fd8e0..8f7bc734d 100644 --- a/internal/agent/loop.go +++ b/internal/agent/loop.go @@ -183,6 +183,11 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) defer runPermissions.cleanup() promptParts := buildSystemPromptParts(options) + planner := newContextPlanner(contextPlannerConfig{ + contextWindow: options.ContextWindow, + promptCacheKey: options.SessionID, + promptParts: promptParts, + }) messages := zeroruntime.SeedMessagesWithImages(promptParts.prompt, prompt, options.Images) guards := newGuardState() @@ -282,42 +287,15 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) exposed, _ := partitionToolsCached(registry, permissionMode, options, loaded, toolDefCache) toolPartitionSpan.End() - // Fingerprint the exact system prompt built at run start plus this turn's - // provider-visible tool definitions in their emitted order. Reusing the - // retained prompt parts avoids rebuilding workspace context while ensuring - // the trace describes the same bytes carried by the request. - if options.Trace != nil { - fp := computePrefixFingerprint(buildPromptSubstringsFromParts(promptParts, exposed)) - options.Trace.EmitPrefixHash(trace.PrefixHash{ - SystemPromptHash: fp.SystemPromptHash, - BaseInstructionsHash: fp.BaseInstructionsHash, - ConfirmationPolicyHash: fp.ConfirmationPolicyHash, - ProjectContextHash: fp.ProjectContextHash, - SkillsHash: fp.SkillsHash, - ToolsHash: fp.ToolsHash, - SchemaHash: fp.SchemaHash, - CompletePrefixHash: fp.CompletePrefixHash, - }) - } - // PROACTIVE compaction: if the history is approaching the model's // context window, summarize the oldest middle before building the // request. A no-op when ContextWindow == 0 (compaction disabled). compactionSpan := options.Trace.Span(trace.SpanCompaction) messages = compactor.maybeCompact(ctx, provider, messages, exposed) compactionSpan.End() - request := zeroruntime.CompletionRequest{ - Messages: copyMessages(messages), - Tools: exposed, - ReasoningEffort: options.ReasoningEffort, - PromptCacheKey: options.SessionID, - } - - // Report the per-category context budget for this turn so a surface can - // show utilization. Opt-in: a no-op when OnContext is unset. - if options.OnContext != nil { - options.OnContext(MeasureContext(messages, request.Tools, options.ContextWindow)) - } + plan := planner.Plan(messages, exposed, options.ReasoningEffort) + request := plan.Request + recordContextPlan(options, plan) // A transient upstream disconnect on the initial connect is retried with // backoff (before any content is forwarded, so no OnText is duplicated); @@ -342,12 +320,7 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) // on registry+loaded, not on the messages, so they stay valid after // compaction. Using the bare toolDefinitions here would route through an // empty-loaded partition, re-hiding every already-loaded deferred tool. - request = zeroruntime.CompletionRequest{ - Messages: copyMessages(messages), - Tools: exposed, - ReasoningEffort: options.ReasoningEffort, - PromptCacheKey: options.SessionID, - } + request = planner.Plan(messages, exposed, options.ReasoningEffort).Request // Pre-content connect after a context-limit compaction: route through the // reconnect helper so a transient upstream hiccup here doesn't fail the // whole run and re-burn every token (AUDIT-L1). @@ -431,12 +404,7 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) // Reuse the SAME active-mode partition (exposed) from this turn rather // than the bare toolDefinitions: exposed depends on registry+loaded (not // the messages), so it stays valid after compaction. - retryRequest := zeroruntime.CompletionRequest{ - Messages: copyMessages(messages), - Tools: exposed, - ReasoningEffort: options.ReasoningEffort, - PromptCacheKey: options.SessionID, - } + retryRequest := planner.Plan(messages, exposed, options.ReasoningEffort).Request retryStream, retryStreamErr := streamWithReconnect(ctx, provider, retryRequest, reconnectNoticeFor(options)) if retryStreamErr != nil { return collected, retryStreamErr @@ -494,12 +462,7 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) result.Messages = copyMessages(messages) return result, err } - retryRequest := zeroruntime.CompletionRequest{ - Messages: copyMessages(messages), - Tools: exposed, - ReasoningEffort: options.ReasoningEffort, - PromptCacheKey: options.SessionID, - } + retryRequest := planner.Plan(messages, exposed, options.ReasoningEffort).Request retryStream, retryErr := streamWithReconnect(ctx, provider, retryRequest, reconnectNoticeFor(options)) if retryErr != nil { result.Messages = copyMessages(messages) @@ -929,7 +892,7 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) }) } finalExposed, _ := partitionToolsCached(registry, permissionMode, options, loaded, toolDefCache) - if answer, finalMessages, finishReason := finalAnswerAfterMaxTurns(ctx, provider, messages, finalExposed, options); strings.TrimSpace(answer) != "" { + if answer, finalMessages, finishReason := finalAnswerAfterMaxTurns(ctx, provider, planner, messages, finalExposed, options); strings.TrimSpace(answer) != "" { result.FinalAnswer = answer result.FinishReason = finishReason result.Messages = copyMessages(finalMessages) @@ -971,7 +934,27 @@ func recordOutputBudgetTrace(recorder *trace.Recorder, result ToolResult) { }) } -func finalAnswerAfterMaxTurns(ctx context.Context, provider Provider, messages []zeroruntime.Message, toolDefs []zeroruntime.ToolDefinition, options Options) (string, []zeroruntime.Message, string) { +func recordContextPlan(options Options, plan contextPlan) { + if options.Trace != nil { + fingerprint := plan.PrefixFingerprint + options.Trace.EmitPrefixHash(trace.PrefixHash{ + SystemPromptHash: fingerprint.SystemPromptHash, + BaseInstructionsHash: fingerprint.BaseInstructionsHash, + ConfirmationPolicyHash: fingerprint.ConfirmationPolicyHash, + ProjectContextHash: fingerprint.ProjectContextHash, + SkillsHash: fingerprint.SkillsHash, + ToolsHash: fingerprint.ToolsHash, + SchemaHash: fingerprint.SchemaHash, + CompletePrefixHash: fingerprint.CompletePrefixHash, + InvalidationReason: plan.Breakdown.PrefixInvalidationReason, + }) + } + if options.OnContext != nil { + options.OnContext(plan.Breakdown) + } +} + +func finalAnswerAfterMaxTurns(ctx context.Context, provider Provider, planner *contextPlanner, messages []zeroruntime.Message, toolDefs []zeroruntime.ToolDefinition, options Options) (string, []zeroruntime.Message, string) { finalMessages := copyMessages(messages) finalMessages = append(finalMessages, zeroruntime.Message{ Role: zeroruntime.MessageRoleUser, @@ -980,12 +963,8 @@ func finalAnswerAfterMaxTurns(ctx context.Context, provider Provider, messages [ // The max-turns final-answer call is a pre-content connect, often after a long // autonomous/cron run — route it through the reconnect helper so a single // transient hiccup doesn't drop the final summary (AUDIT-L1). - stream, err := streamWithReconnect(ctx, provider, zeroruntime.CompletionRequest{ - Messages: copyMessages(finalMessages), - Tools: toolDefs, - ReasoningEffort: options.ReasoningEffort, - PromptCacheKey: options.SessionID, - }, reconnectNoticeFor(options)) + request := planner.Plan(finalMessages, toolDefs, options.ReasoningEffort).Request + stream, err := streamWithReconnect(ctx, provider, request, reconnectNoticeFor(options)) if err != nil { return "", messages, "" } diff --git a/internal/agent/loop_test.go b/internal/agent/loop_test.go index 86647da84..46f0affb1 100644 --- a/internal/agent/loop_test.go +++ b/internal/agent/loop_test.go @@ -3729,6 +3729,8 @@ func TestRunTracingWrapperStampsUsage(t *testing.T) { {Type: zeroruntime.StreamEventDone}, }}} onUsageCalls := 0 + onContextCalls := 0 + var contextPlan ContextBreakdown rec := trace.NewRecorder("tracing-session", "run-1", "test") if _, err := Run(context.Background(), "hi", provider, Options{ SessionID: "tracing-session", @@ -3737,6 +3739,10 @@ func TestRunTracingWrapperStampsUsage(t *testing.T) { Model: "test-model", Trace: rec, OnUsage: func(Usage) { onUsageCalls++ }, + OnContext: func(breakdown ContextBreakdown) { + onContextCalls++ + contextPlan = breakdown + }, }); err != nil { t.Fatalf("Run: %v", err) } @@ -3763,6 +3769,12 @@ func TestRunTracingWrapperStampsUsage(t *testing.T) { if onUsageCalls == 0 { t.Fatal("wrapped OnUsage did not forward to the caller's callback") } + if onContextCalls != 1 || len(contextPlan.Blocks) != 2 || contextPlan.PrefixInvalidationReason != "initial" { + t.Fatalf("context plan callback = calls %d, plan %#v", onContextCalls, contextPlan) + } + if len(tr.PrefixHashes) != 1 || tr.PrefixHashes[0].InvalidationReason != "initial" || tr.PrefixHashes[0].CompletePrefixHash != contextPlan.CompletePrefixHash { + t.Fatalf("trace context evidence = %#v, plan %#v", tr.PrefixHashes, contextPlan) + } } // TestRunNilTraceForwardsUsage verifies a nil recorder leaves the loop diff --git a/internal/trace/emit.go b/internal/trace/emit.go index d8f35cb4d..62c4a43be 100644 --- a/internal/trace/emit.go +++ b/internal/trace/emit.go @@ -105,6 +105,7 @@ func WriteNDJSON(w io.Writer, t *TurnTrace) error { "tools": p.ToolsHash, "schema": p.SchemaHash, "complete_prefix": p.CompletePrefixHash, + "invalidation_reason": p.InvalidationReason, }); err != nil { return err } diff --git a/internal/trace/parse.go b/internal/trace/parse.go index 7204e6906..2ffa098d7 100644 --- a/internal/trace/parse.go +++ b/internal/trace/parse.go @@ -116,6 +116,7 @@ func ReadNDJSON(r io.Reader) (*TurnTrace, error) { ToolsHash: stringField(obj, "tools"), SchemaHash: stringField(obj, "schema"), CompletePrefixHash: stringField(obj, "complete_prefix"), + InvalidationReason: stringField(obj, "invalidation_reason"), }) case "output_budget": if !sawTraceHeader { diff --git a/internal/trace/trace.go b/internal/trace/trace.go index 37b28927a..868449e42 100644 --- a/internal/trace/trace.go +++ b/internal/trace/trace.go @@ -159,6 +159,7 @@ type PrefixHash struct { ToolsHash string `json:"tools"` SchemaHash string `json:"schema"` CompletePrefixHash string `json:"complete_prefix"` + InvalidationReason string `json:"invalidation_reason,omitempty"` } // WallDuration is the total traced wall time of the run. diff --git a/internal/trace/trace_test.go b/internal/trace/trace_test.go index fffa8c695..66289f10e 100644 --- a/internal/trace/trace_test.go +++ b/internal/trace/trace_test.go @@ -382,6 +382,7 @@ func TestReadNDJSONRoundTrip(t *testing.T) { ToolsHash: "t1", SchemaHash: "x1", CompletePrefixHash: "complete1", + InvalidationReason: "initial", }) r.EmitPrefixHash(PrefixHash{ SystemPromptHash: "system2", @@ -392,6 +393,7 @@ func TestReadNDJSONRoundTrip(t *testing.T) { ToolsHash: "t2", SchemaHash: "x2", CompletePrefixHash: "complete2", + InvalidationReason: "tools", }) original := r.Finish() @@ -421,8 +423,8 @@ func TestReadNDJSONRoundTrip(t *testing.T) { if parsed.FirstTokenAt.IsZero() { t.Fatal("first_token_at lost in round-trip") } - // prefix_hash round-trip: two events, in insertion order, with all - // seven sub-hash fields preserved exactly. + // prefix_hash round-trip: two events, in insertion order, with all hash + // fields and the invalidation reason preserved exactly. if len(parsed.PrefixHashes) != 2 { t.Fatalf("expected 2 prefix_hash events after round-trip, got %d", len(parsed.PrefixHashes)) } @@ -435,6 +437,9 @@ func TestReadNDJSONRoundTrip(t *testing.T) { if parsed.PrefixHashes[1].SystemPromptHash != "system2" || parsed.PrefixHashes[1].BaseInstructionsHash != "b2" || parsed.PrefixHashes[1].SchemaHash != "x2" { t.Fatalf("prefix_hash sub-hashes lost on second event: got %+v", parsed.PrefixHashes[1]) } + if parsed.PrefixHashes[0].InvalidationReason != "initial" || parsed.PrefixHashes[1].InvalidationReason != "tools" { + t.Fatalf("prefix invalidation reasons lost: %+v", parsed.PrefixHashes) + } } func TestReadNDJSONRejectsNonTrace(t *testing.T) { From 0368c2bd3b47114a5df4556800cb9eeb48c49b48 Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 8 Aug 2026 13:02:10 +0530 Subject: [PATCH 2/3] Fix context plan request accounting --- internal/agent/compaction.go | 39 ++++++++++------ internal/agent/compaction_test.go | 63 ++++++++++++++++++++++++++ internal/agent/context_planner.go | 19 ++++++-- internal/agent/context_planner_test.go | 19 ++++++++ internal/agent/loop.go | 50 ++++++++++++-------- internal/agent/types.go | 8 ++-- internal/trace/emit.go | 19 ++++---- internal/trace/recorder.go | 6 +-- internal/trace/trace_test.go | 29 ++++++++++++ 9 files changed, 203 insertions(+), 49 deletions(-) diff --git a/internal/agent/compaction.go b/internal/agent/compaction.go index 981ebc187..25fc9208a 100644 --- a/internal/agent/compaction.go +++ b/internal/agent/compaction.go @@ -381,6 +381,8 @@ type compactionState struct { // but its token COST must still be counted so usage reports and budgets include it. onUsage func(Usage) task *taskState + planner *contextPlanner + trace *trace.Recorder // calibrationRatio scales the raw byte/4 token estimate toward the provider's // real prompt-token count. ApproxTextTokens over-counts code-heavy content by @@ -427,6 +429,8 @@ func newCompactionState(options Options, task *taskState) *compactionState { preserveLast: options.CompactionPreserveLast, onUsage: options.OnUsage, task: task, + planner: newContextPlanner(contextPlannerConfig{contextWindow: options.ContextWindow}), + trace: options.Trace, } return state } @@ -482,7 +486,7 @@ func (state *compactionState) maybeCompact( compacted, err := Compact(messages, CompactionOptions{ PreserveLast: state.preserveLast, - Summarize: summarizeClosure(ctx, provider, state.onUsage), + Summarize: state.summarizeClosure(ctx, provider), taskState: state.task.snapshotForCompaction(messages), }) if err != nil { @@ -533,7 +537,7 @@ func (state *compactionState) recover( result, compactErr := Compact(messages, CompactionOptions{ PreserveLast: state.preserveLast, - Summarize: summarizeClosure(ctx, provider, state.onUsage), + Summarize: state.summarizeClosure(ctx, provider), taskState: state.task.snapshotForCompaction(messages), }) if compactErr != nil { @@ -567,9 +571,12 @@ func (state *compactionState) recover( // provider call. The summary stream intentionally does NOT forward OnText (so // compaction stays invisible on the user-facing surface), but it DOES forward // OnUsage so the summarizer's token cost is still counted by usage/budgeting. -func summarizeClosure(ctx context.Context, provider Provider, onUsage func(Usage)) func([]zeroruntime.Message) (string, error) { +func (state *compactionState) summarizeClosure(ctx context.Context, provider Provider) func([]zeroruntime.Message) (string, error) { + if state.planner == nil { + state.planner = newContextPlanner(contextPlannerConfig{}) + } return func(toSummarize []zeroruntime.Message) (string, error) { - return summarizeWithFallback(ctx, provider, toSummarize, onUsage) + return summarizeWithPlanner(ctx, provider, toSummarize, state.onUsage, state.planner, state.trace) } } @@ -581,7 +588,11 @@ func summarizeClosure(ctx context.Context, provider Provider, onUsage func(Usage // Non-context-limit errors (and a single message that still won't fit) surface // to the caller unchanged. func summarizeWithFallback(ctx context.Context, provider Provider, messages []zeroruntime.Message, onUsage func(Usage)) (string, error) { - summary, err := summarizeMessagesOnce(ctx, provider, messages, onUsage) + return summarizeWithPlanner(ctx, provider, messages, onUsage, newContextPlanner(contextPlannerConfig{}), nil) +} + +func summarizeWithPlanner(ctx context.Context, provider Provider, messages []zeroruntime.Message, onUsage func(Usage), planner *contextPlanner, recorder *trace.Recorder) (string, error) { + summary, err := summarizeMessagesOnce(ctx, provider, messages, onUsage, planner, recorder) if err == nil { return summary, nil } @@ -590,11 +601,11 @@ func summarizeWithFallback(ctx context.Context, provider Provider, messages []ze } mid := len(messages) / 2 - left, leftErr := summarizeWithFallback(ctx, provider, messages[:mid], onUsage) + left, leftErr := summarizeWithPlanner(ctx, provider, messages[:mid], onUsage, planner, recorder) if leftErr != nil { return "", leftErr } - right, rightErr := summarizeWithFallback(ctx, provider, messages[mid:], onUsage) + right, rightErr := summarizeWithPlanner(ctx, provider, messages[mid:], onUsage, planner, recorder) if rightErr != nil { return "", rightErr } @@ -607,7 +618,7 @@ func summarizeWithFallback(ctx context.Context, provider Provider, messages []ze combined := strings.TrimSpace(left + "\n\n" + right) reduced, reduceErr := summarizeMessagesOnce(ctx, provider, []zeroruntime.Message{ {Role: zeroruntime.MessageRoleUser, Content: combined}, - }, onUsage) + }, onUsage, planner, recorder) if reduceErr != nil { if isContextLimitError(reduceErr.Error()) { // Even the two combined partials don't fit (extreme): fall back to the @@ -623,13 +634,15 @@ func summarizeWithFallback(ctx context.Context, provider Provider, messages []ze } // summarizeMessagesOnce performs a single tool-less summarization call. -func summarizeMessagesOnce(ctx context.Context, provider Provider, messages []zeroruntime.Message, onUsage func(Usage)) (string, error) { - planner := newContextPlanner(contextPlannerConfig{}) - request := planner.Plan([]zeroruntime.Message{ +func summarizeMessagesOnce(ctx context.Context, provider Provider, messages []zeroruntime.Message, onUsage func(Usage), planner *contextPlanner, recorder *trace.Recorder) (string, error) { + requestMessages := []zeroruntime.Message{ {Role: zeroruntime.MessageRoleSystem, Content: CompactionSummaryInstructions}, {Role: zeroruntime.MessageRoleUser, Content: "Summarize this conversation:\n\n" + renderTranscript(messages)}, - }, nil, "").Request - stream, err := provider.StreamCompletion(ctx, request) + } + parts := systemPromptParts{prompt: CompactionSummaryInstructions, baseInstructions: CompactionSummaryInstructions} + plan := planner.planWithPromptParts(requestMessages, nil, "", parts) + recordContextPlanTrace(recorder, plan) + stream, err := provider.StreamCompletion(ctx, plan.Request) if err != nil { return "", err } diff --git a/internal/agent/compaction_test.go b/internal/agent/compaction_test.go index 2b38e5ea8..3db3db1f8 100644 --- a/internal/agent/compaction_test.go +++ b/internal/agent/compaction_test.go @@ -7,6 +7,7 @@ import ( "testing" "github.com/Gitlawb/zero/internal/tools" + "github.com/Gitlawb/zero/internal/trace" "github.com/Gitlawb/zero/internal/zeroruntime" ) @@ -428,6 +429,7 @@ type reactiveProvider struct { summarizeCalls int turnRequests int failedOnce bool + connectError bool bigText string finalText string } @@ -449,6 +451,9 @@ func (provider *reactiveProvider) StreamCompletion(ctx context.Context, request return streamEvents(toolTurnWithText(provider.bigText, "1", "read_file", `{"path":"x"}`)), nil case provider.turnRequests == 2 && !provider.failedOnce: provider.failedOnce = true + if provider.connectError { + return nil, errors.New("context_length_exceeded") + } return streamEvents([]zeroruntime.StreamEvent{ {Type: zeroruntime.StreamEventError, Error: "This model's maximum context length is 1000 tokens. Please reduce the length of the messages."}, }), nil @@ -457,6 +462,39 @@ func (provider *reactiveProvider) StreamCompletion(ctx context.Context, request } } +func TestRunConnectErrorCompactionRecordsReplacementPlan(t *testing.T) { + provider := &reactiveProvider{ + connectError: true, + bigText: strings.Repeat("b", 6000), + finalText: "recovered", + } + registry := tools.NewRegistry() + registry.Register(tools.NewScopedReadFileTool(t.TempDir(), nil)) + recorder := trace.NewRecorder("connect-compaction", "run-1", "test") + var contextPlans []ContextBreakdown + + result, err := Run(context.Background(), strings.Repeat("z", 6000), provider, Options{ + Registry: registry, + PermissionMode: PermissionModeUnsafe, + ContextWindow: 10_000_000, + CompactionPreserveLast: 2, + Trace: recorder, + OnContext: func(breakdown ContextBreakdown) { + contextPlans = append(contextPlans, breakdown) + }, + }) + if err != nil || result.FinalAnswer != "recovered" { + t.Fatalf("connect-error compaction failed: result=%#v err=%v", result, err) + } + if len(contextPlans) != 3 || contextPlans[2].MessageTokens >= contextPlans[1].MessageTokens { + t.Fatalf("replacement request context not reported: %#v", contextPlans) + } + gotTrace := recorder.Finish() + if len(gotTrace.PrefixHashes) != 4 || gotTrace.PrefixHashes[3].InvalidationReason != "unchanged" { + t.Fatalf("replacement request trace not recorded: %#v", gotTrace.PrefixHashes) + } +} + func TestRunReactiveCompactionRecovers(t *testing.T) { provider := &reactiveProvider{ bigText: strings.Repeat("b", 6000), @@ -468,6 +506,8 @@ func TestRunReactiveCompactionRecovers(t *testing.T) { // bloats the history. registry := tools.NewRegistry() registry.Register(tools.NewScopedReadFileTool(t.TempDir(), nil)) + recorder := trace.NewRecorder("reactive-compaction", "run-1", "test") + var contextPlans []ContextBreakdown // ContextWindow large enough that proactive compaction never triggers, so // only the reactive path can save the run. @@ -476,6 +516,10 @@ func TestRunReactiveCompactionRecovers(t *testing.T) { PermissionMode: PermissionModeUnsafe, ContextWindow: 10_000_000, CompactionPreserveLast: 2, + Trace: recorder, + OnContext: func(breakdown ContextBreakdown) { + contextPlans = append(contextPlans, breakdown) + }, }) if err != nil { t.Fatalf("expected reactive compaction to recover the run, got error: %v", err) @@ -486,6 +530,25 @@ func TestRunReactiveCompactionRecovers(t *testing.T) { if provider.summarizeCalls == 0 { t.Fatal("expected reactive compaction to invoke the summarizer") } + if len(contextPlans) != 3 { + t.Fatalf("main request context callbacks = %d, want initial turns plus compacted retry", len(contextPlans)) + } + if contextPlans[2].MessageTokens >= contextPlans[1].MessageTokens { + t.Fatalf("retry context was not updated after compaction: before=%d after=%d", contextPlans[1].MessageTokens, contextPlans[2].MessageTokens) + } + gotTrace := recorder.Finish() + if len(gotTrace.PrefixHashes) != 4 { + t.Fatalf("prefix evidence count = %d, want two turns, compaction, and retry: %#v", len(gotTrace.PrefixHashes), gotTrace.PrefixHashes) + } + if gotTrace.PrefixHashes[0].InvalidationReason != "initial" || + gotTrace.PrefixHashes[1].InvalidationReason != "unchanged" || + gotTrace.PrefixHashes[2].InvalidationReason != "initial" || + gotTrace.PrefixHashes[3].InvalidationReason != "unchanged" { + t.Fatalf("unexpected request-specific invalidation evidence: %#v", gotTrace.PrefixHashes) + } + if gotTrace.PrefixHashes[2].CompletePrefixHash == gotTrace.PrefixHashes[1].CompletePrefixHash { + t.Fatalf("compaction fingerprint reused the main request prefix: %#v", gotTrace.PrefixHashes) + } } // midStreamReactiveProvider forwards some text BEFORE surfacing a context-limit diff --git a/internal/agent/context_planner.go b/internal/agent/context_planner.go index 567e28ea4..c62fcfbe7 100644 --- a/internal/agent/context_planner.go +++ b/internal/agent/context_planner.go @@ -37,15 +37,28 @@ func newContextPlanner(config contextPlannerConfig) *contextPlanner { // It intentionally performs no relevance filtering: preserving current model // capability is the baseline contract for future planner policies. func (planner *contextPlanner) Plan(messages []zeroruntime.Message, toolDefs []zeroruntime.ToolDefinition, reasoningEffort string) contextPlan { + return planner.plan(messages, toolDefs, reasoningEffort, planner.config.promptParts) +} + +// planWithPromptParts plans a request whose stable system sections differ from +// the planner's configured main-run prompt, such as a compaction summary call. +func (planner *contextPlanner) planWithPromptParts(messages []zeroruntime.Message, toolDefs []zeroruntime.ToolDefinition, reasoningEffort string, parts systemPromptParts) contextPlan { + return planner.plan(messages, toolDefs, reasoningEffort, parts) +} + +func (planner *contextPlanner) plan(messages []zeroruntime.Message, toolDefs []zeroruntime.ToolDefinition, reasoningEffort string, parts systemPromptParts) contextPlan { request := zeroruntime.CompletionRequest{ Messages: copyMessages(messages), Tools: toolDefs, ReasoningEffort: reasoningEffort, PromptCacheKey: planner.config.promptCacheKey, } - parts := planner.config.promptParts - if parts.prompt == "" { - parts.prompt = leadingSystemContent(request.Messages) + requestSystemPrompt := leadingSystemContent(request.Messages) + if parts.prompt != requestSystemPrompt { + // Component boundaries from a different request would make the detailed + // invalidation reason misleading. Retain the exact request prefix and + // classify it conservatively as an unstructured system prompt instead. + parts = systemPromptParts{prompt: requestSystemPrompt} } fingerprint := computePrefixFingerprint(buildPromptSubstringsFromParts(parts, request.Tools)) breakdown := MeasureContext(request.Messages, request.Tools, planner.config.contextWindow) diff --git a/internal/agent/context_planner_test.go b/internal/agent/context_planner_test.go index e47ae85d7..fb5d0f910 100644 --- a/internal/agent/context_planner_test.go +++ b/internal/agent/context_planner_test.go @@ -65,6 +65,24 @@ func TestContextPlannerReportsInspectableBlocks(t *testing.T) { } } +func TestContextPlannerFingerprintsActualRequestSystemPrompt(t *testing.T) { + planner := newContextPlanner(contextPlannerConfig{promptParts: systemPromptParts{ + prompt: "configured system", baseInstructions: "configured system", + }}) + messages := []zeroruntime.Message{ + {Role: zeroruntime.MessageRoleSystem, Content: "request-specific system"}, + {Role: zeroruntime.MessageRoleUser, Content: "task"}, + } + + plan := planner.Plan(messages, nil, "") + want := computePrefixFingerprint(buildPromptSubstringsFromParts( + systemPromptParts{prompt: "request-specific system"}, nil, + )) + if plan.PrefixFingerprint != want { + t.Fatalf("fingerprint describes configured prompt instead of request:\n got: %#v\nwant: %#v", plan.PrefixFingerprint, want) + } +} + func TestContextPlannerExplainsPrefixInvalidation(t *testing.T) { planner := newContextPlanner(contextPlannerConfig{promptParts: systemPromptParts{ prompt: "system", baseInstructions: "base", projectContext: "project", @@ -88,6 +106,7 @@ func TestContextPlannerExplainsPrefixInvalidation(t *testing.T) { } planner.config.promptParts.prompt = "system with changed project" planner.config.promptParts.projectContext = "changed project" + messages[0].Content = "system with changed project" if got := planner.Plan(messages, tools, "").Breakdown.PrefixInvalidationReason; got != "project_context" { t.Fatalf("project reason = %q, want project_context", got) } diff --git a/internal/agent/loop.go b/internal/agent/loop.go index 8f7bc734d..7c3d17813 100644 --- a/internal/agent/loop.go +++ b/internal/agent/loop.go @@ -320,7 +320,9 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) // on registry+loaded, not on the messages, so they stay valid after // compaction. Using the bare toolDefinitions here would route through an // empty-loaded partition, re-hiding every already-loaded deferred tool. - request = planner.Plan(messages, exposed, options.ReasoningEffort).Request + retryPlan := planner.Plan(messages, exposed, options.ReasoningEffort) + request = retryPlan.Request + recordContextPlan(options, retryPlan) // Pre-content connect after a context-limit compaction: route through the // reconnect helper so a transient upstream hiccup here doesn't fail the // whole run and re-burn every token (AUDIT-L1). @@ -404,7 +406,9 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) // Reuse the SAME active-mode partition (exposed) from this turn rather // than the bare toolDefinitions: exposed depends on registry+loaded (not // the messages), so it stays valid after compaction. - retryRequest := planner.Plan(messages, exposed, options.ReasoningEffort).Request + retryPlan := planner.Plan(messages, exposed, options.ReasoningEffort) + retryRequest := retryPlan.Request + recordContextPlan(options, retryPlan) retryStream, retryStreamErr := streamWithReconnect(ctx, provider, retryRequest, reconnectNoticeFor(options)) if retryStreamErr != nil { return collected, retryStreamErr @@ -462,7 +466,9 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) result.Messages = copyMessages(messages) return result, err } - retryRequest := planner.Plan(messages, exposed, options.ReasoningEffort).Request + retryPlan := planner.Plan(messages, exposed, options.ReasoningEffort) + retryRequest := retryPlan.Request + recordContextPlan(options, retryPlan) retryStream, retryErr := streamWithReconnect(ctx, provider, retryRequest, reconnectNoticeFor(options)) if retryErr != nil { result.Messages = copyMessages(messages) @@ -935,25 +941,30 @@ func recordOutputBudgetTrace(recorder *trace.Recorder, result ToolResult) { } func recordContextPlan(options Options, plan contextPlan) { - if options.Trace != nil { - fingerprint := plan.PrefixFingerprint - options.Trace.EmitPrefixHash(trace.PrefixHash{ - SystemPromptHash: fingerprint.SystemPromptHash, - BaseInstructionsHash: fingerprint.BaseInstructionsHash, - ConfirmationPolicyHash: fingerprint.ConfirmationPolicyHash, - ProjectContextHash: fingerprint.ProjectContextHash, - SkillsHash: fingerprint.SkillsHash, - ToolsHash: fingerprint.ToolsHash, - SchemaHash: fingerprint.SchemaHash, - CompletePrefixHash: fingerprint.CompletePrefixHash, - InvalidationReason: plan.Breakdown.PrefixInvalidationReason, - }) - } + recordContextPlanTrace(options.Trace, plan) if options.OnContext != nil { options.OnContext(plan.Breakdown) } } +func recordContextPlanTrace(recorder *trace.Recorder, plan contextPlan) { + if recorder == nil { + return + } + fingerprint := plan.PrefixFingerprint + recorder.EmitPrefixHash(trace.PrefixHash{ + SystemPromptHash: fingerprint.SystemPromptHash, + BaseInstructionsHash: fingerprint.BaseInstructionsHash, + ConfirmationPolicyHash: fingerprint.ConfirmationPolicyHash, + ProjectContextHash: fingerprint.ProjectContextHash, + SkillsHash: fingerprint.SkillsHash, + ToolsHash: fingerprint.ToolsHash, + SchemaHash: fingerprint.SchemaHash, + CompletePrefixHash: fingerprint.CompletePrefixHash, + InvalidationReason: plan.Breakdown.PrefixInvalidationReason, + }) +} + func finalAnswerAfterMaxTurns(ctx context.Context, provider Provider, planner *contextPlanner, messages []zeroruntime.Message, toolDefs []zeroruntime.ToolDefinition, options Options) (string, []zeroruntime.Message, string) { finalMessages := copyMessages(messages) finalMessages = append(finalMessages, zeroruntime.Message{ @@ -963,8 +974,9 @@ func finalAnswerAfterMaxTurns(ctx context.Context, provider Provider, planner *c // The max-turns final-answer call is a pre-content connect, often after a long // autonomous/cron run — route it through the reconnect helper so a single // transient hiccup doesn't drop the final summary (AUDIT-L1). - request := planner.Plan(finalMessages, toolDefs, options.ReasoningEffort).Request - stream, err := streamWithReconnect(ctx, provider, request, reconnectNoticeFor(options)) + plan := planner.Plan(finalMessages, toolDefs, options.ReasoningEffort) + recordContextPlan(options, plan) + stream, err := streamWithReconnect(ctx, provider, plan.Request, reconnectNoticeFor(options)) if err != nil { return "", messages, "" } diff --git a/internal/agent/types.go b/internal/agent/types.go index 64035435e..7e6cab927 100644 --- a/internal/agent/types.go +++ b/internal/agent/types.go @@ -333,9 +333,11 @@ type Options struct { // specialist child process emits while running. The toolCallID identifies // which Task tool call the progress belongs to. nil is a no-op. OnToolProgress func(toolCallID string, event streamjson.Event) - // OnContext, when set, is called once per turn with the per-category context - // budget of the request about to be sent, so a surface (TUI/CLI) can show - // context utilization. Opt-in like the other callbacks; nil is a no-op. + // OnContext, when set, is called for each main agent request with its + // per-category context budget, including a replacement request after + // compaction or a stall retry. Internal summarizer requests are excluded so + // surfaces keep showing the active conversation budget. Opt-in like the other + // callbacks; nil is a no-op. OnContext func(ContextBreakdown) // ModelSwitcher, when set, lets a tool escalate the run to a stronger model // mid-run: the loop calls it with the requested model id and, on success, diff --git a/internal/trace/emit.go b/internal/trace/emit.go index 62c4a43be..dbca36eb9 100644 --- a/internal/trace/emit.go +++ b/internal/trace/emit.go @@ -86,16 +86,16 @@ func WriteNDJSON(w io.Writer, t *TurnTrace) error { } } - // Prefix fingerprints are emitted after counters in insertion (turn) + // Prefix fingerprints are emitted after counters in insertion (request) // order. The order is the order EmitPrefixHash was called, which is the - // order the agent loop computed each turn's fingerprint, which is the - // order a downstream consumer needs to correlate a prefix_hash event - // with the cached_input_tokens counter for that turn. Sorting by - // complete_prefix hash would destroy that correlation, so we do not + // order the agent computed each request's fingerprint, which is the + // order a downstream consumer needs to reconstruct prefix changes across + // main, compaction, and retry requests. Sorting by complete_prefix hash would + // destroy that sequence, so we do not // sort. The slice is already a deep copy from Finish (see // Recorder.Finish) so it is safe to range over without copying. for _, p := range t.PrefixHashes { - if err := enc.Encode(map[string]any{ + obj := map[string]any{ "type": "prefix_hash", "system_prompt": p.SystemPromptHash, "base_instructions": p.BaseInstructionsHash, @@ -105,8 +105,11 @@ func WriteNDJSON(w io.Writer, t *TurnTrace) error { "tools": p.ToolsHash, "schema": p.SchemaHash, "complete_prefix": p.CompletePrefixHash, - "invalidation_reason": p.InvalidationReason, - }); err != nil { + } + if p.InvalidationReason != "" { + obj["invalidation_reason"] = p.InvalidationReason + } + if err := enc.Encode(obj); err != nil { return err } } diff --git a/internal/trace/recorder.go b/internal/trace/recorder.go index 137e504fc..9546bb153 100644 --- a/internal/trace/recorder.go +++ b/internal/trace/recorder.go @@ -177,9 +177,9 @@ func (r *Recorder) StampFirstUsefulAction() { } // EmitPrefixHash records one prompt-prefix fingerprint on the trace. Multiple -// calls are allowed within a run (one per turn, typically) and accumulate in -// order. The first call after Finish is a no-op; later calls are also no-ops -// because the trace has been sealed. +// calls are allowed within a run (one per planned provider request) and +// accumulate in request order. The first call after Finish is a no-op; later +// calls are also no-ops because the trace has been sealed. func (r *Recorder) EmitPrefixHash(p PrefixHash) { if r == nil { return diff --git a/internal/trace/trace_test.go b/internal/trace/trace_test.go index 66289f10e..9aac80489 100644 --- a/internal/trace/trace_test.go +++ b/internal/trace/trace_test.go @@ -442,6 +442,35 @@ func TestReadNDJSONRoundTrip(t *testing.T) { } } +func TestWriteNDJSONOmitsEmptyPrefixInvalidationReason(t *testing.T) { + recorder := NewRecorder("session", "run", "test") + recorder.Start() + recorder.EmitPrefixHash(PrefixHash{CompletePrefixHash: "stable"}) + tr := recorder.Finish() + + var output bytes.Buffer + if err := WriteNDJSON(&output, tr); err != nil { + t.Fatalf("WriteNDJSON: %v", err) + } + found := false + for _, line := range strings.Split(strings.TrimSpace(output.String()), "\n") { + var event map[string]any + if err := json.Unmarshal([]byte(line), &event); err != nil { + t.Fatalf("decode event: %v", err) + } + if event["type"] != "prefix_hash" { + continue + } + found = true + if _, exists := event["invalidation_reason"]; exists { + t.Fatalf("empty invalidation reason should be omitted: %s", line) + } + } + if !found { + t.Fatal("prefix_hash event not written") + } +} + func TestReadNDJSONRejectsNonTrace(t *testing.T) { // A file with content but no type:trace header is never a valid empty trace. if _, err := ReadNDJSON(strings.NewReader("not json at all\n")); err == nil { From bea0c21472113512f1c3ac5d0f5d3f7acfdde1d6 Mon Sep 17 00:00:00 2001 From: anandh8x Date: Sat, 8 Aug 2026 13:22:05 +0530 Subject: [PATCH 3/3] Freeze planned tool schemas --- internal/agent/context_planner.go | 77 ++++++++++++++++++++++---- internal/agent/context_planner_test.go | 35 +++++++++++- 2 files changed, 99 insertions(+), 13 deletions(-) diff --git a/internal/agent/context_planner.go b/internal/agent/context_planner.go index c62fcfbe7..16fb920e7 100644 --- a/internal/agent/context_planner.go +++ b/internal/agent/context_planner.go @@ -18,9 +18,11 @@ type contextPlannerConfig struct { // message and tool definition while making composition and cache drift // inspectable; later selection policies must continue to cross this seam. type contextPlanner struct { - config contextPlannerConfig - previousPrefix prefixFingerprint - hasPrevious bool + config contextPlannerConfig + previousPrefix prefixFingerprint + hasPrevious bool + toolSnapshotKey string + toolSnapshot []zeroruntime.ToolDefinition } type contextPlan struct { @@ -47,20 +49,21 @@ func (planner *contextPlanner) planWithPromptParts(messages []zeroruntime.Messag } func (planner *contextPlanner) plan(messages []zeroruntime.Message, toolDefs []zeroruntime.ToolDefinition, reasoningEffort string, parts systemPromptParts) contextPlan { - request := zeroruntime.CompletionRequest{ - Messages: copyMessages(messages), - Tools: toolDefs, - ReasoningEffort: reasoningEffort, - PromptCacheKey: planner.config.promptCacheKey, - } - requestSystemPrompt := leadingSystemContent(request.Messages) + requestMessages := copyMessages(messages) + requestSystemPrompt := leadingSystemContent(requestMessages) if parts.prompt != requestSystemPrompt { // Component boundaries from a different request would make the detailed // invalidation reason misleading. Retain the exact request prefix and // classify it conservatively as an unstructured system prompt instead. parts = systemPromptParts{prompt: requestSystemPrompt} } - fingerprint := computePrefixFingerprint(buildPromptSubstringsFromParts(parts, request.Tools)) + fingerprint := computePrefixFingerprint(buildPromptSubstringsFromParts(parts, toolDefs)) + request := zeroruntime.CompletionRequest{ + Messages: requestMessages, + Tools: planner.snapshotTools(toolDefs, fingerprint), + ReasoningEffort: reasoningEffort, + PromptCacheKey: planner.config.promptCacheKey, + } breakdown := MeasureContext(request.Messages, request.Tools, planner.config.contextWindow) breakdown.CompletePrefixHash = fingerprint.CompletePrefixHash var previous *prefixFingerprint @@ -77,6 +80,58 @@ func (planner *contextPlanner) plan(messages []zeroruntime.Message, toolDefs []z } } +// snapshotTools freezes provider-visible tool definitions when their semantic +// fingerprint changes. Stable turns reuse the frozen schemas instead of +// recursively cloning every map and slice for every request. +func (planner *contextPlanner) snapshotTools(toolDefs []zeroruntime.ToolDefinition, fingerprint prefixFingerprint) []zeroruntime.ToolDefinition { + key := fingerprint.ToolsHash + "\x00" + fingerprint.SchemaHash + if planner.toolSnapshotKey != key { + planner.toolSnapshot = copyToolDefinitions(toolDefs) + planner.toolSnapshotKey = key + } + return planner.toolSnapshot +} + +func copyToolDefinitions(toolDefs []zeroruntime.ToolDefinition) []zeroruntime.ToolDefinition { + if toolDefs == nil { + return nil + } + copied := make([]zeroruntime.ToolDefinition, len(toolDefs)) + for index, definition := range toolDefs { + copied[index] = definition + copied[index].Parameters = copySchemaMap(definition.Parameters) + } + return copied +} + +func copySchemaMap(schema map[string]any) map[string]any { + if schema == nil { + return nil + } + copied := make(map[string]any, len(schema)) + for key, value := range schema { + copied[key] = copySchemaValue(value) + } + return copied +} + +func copySchemaValue(value any) any { + switch typed := value.(type) { + case map[string]any: + return copySchemaMap(typed) + case []any: + copied := make([]any, len(typed)) + for index, item := range typed { + copied[index] = copySchemaValue(item) + } + return copied + case []string: + return append([]string(nil), typed...) + default: + return value + } +} + func leadingSystemContent(messages []zeroruntime.Message) string { contents := make([]string, 0, 1) for _, message := range messages { diff --git a/internal/agent/context_planner_test.go b/internal/agent/context_planner_test.go index fb5d0f910..19a6bc1f1 100644 --- a/internal/agent/context_planner_test.go +++ b/internal/agent/context_planner_test.go @@ -13,7 +13,15 @@ func TestContextPlannerPreservesProviderRequest(t *testing.T) { {Role: zeroruntime.MessageRoleUser, Content: "inspect this", Images: []zeroruntime.ImageBlock{{MediaType: "image/png", Data: []byte{1, 2, 3}}}}, {Role: zeroruntime.MessageRoleAssistant, Content: "working", ToolCalls: []zeroruntime.ToolCall{{ID: "call-1", Name: "read_file", Arguments: `{"path":"main.go"}`}}, Reasoning: []zeroruntime.ReasoningBlock{{Provider: "test", Type: "thinking", Signature: "sig"}}}, } - toolDefs := []zeroruntime.ToolDefinition{{Name: "read_file", Description: "Read a file", Parameters: map[string]any{"type": "object"}}} + toolDefs := []zeroruntime.ToolDefinition{{ + Name: "read_file", Description: "Read a file", + Parameters: map[string]any{ + "type": "object", + "required": []string{"path"}, + "properties": map[string]any{"path": map[string]any{"type": "string"}}, + "anyOf": []any{map[string]any{"type": "object"}}, + }, + }} planner := newContextPlanner(contextPlannerConfig{ contextWindow: 128_000, promptCacheKey: "session-1", @@ -30,13 +38,36 @@ func TestContextPlannerPreservesProviderRequest(t *testing.T) { if !reflect.DeepEqual(plan.Request, want) { t.Fatalf("planned request changed provider input:\n got: %#v\nwant: %#v", plan.Request, want) } + stable := planner.Plan(messages, toolDefs, "medium") + if &stable.Request.Tools[0] != &plan.Request.Tools[0] { + t.Fatal("stable tool fingerprint rebuilt the immutable schema snapshot") + } - // A plan owns its request snapshot; later caller mutation cannot change it. + // A plan owns its request snapshot; later caller mutation cannot change its + // messages or its recursively nested tool schema. messages[1].Content = "changed" messages[1].Images[0].Data[0] = 9 + toolDefs[0].Name = "changed_tool" + toolDefs[0].Parameters["required"].([]string)[0] = "other" + toolDefs[0].Parameters["properties"].(map[string]any)["path"].(map[string]any)["type"] = "integer" + toolDefs[0].Parameters["anyOf"].([]any)[0].(map[string]any)["type"] = "array" if plan.Request.Messages[1].Content != "inspect this" || plan.Request.Messages[1].Images[0].Data[0] != 1 { t.Fatalf("planned messages alias caller state: %#v", plan.Request.Messages[1]) } + if plan.Request.Tools[0].Name != "read_file" || + plan.Request.Tools[0].Parameters["required"].([]string)[0] != "path" || + plan.Request.Tools[0].Parameters["properties"].(map[string]any)["path"].(map[string]any)["type"] != "string" || + plan.Request.Tools[0].Parameters["anyOf"].([]any)[0].(map[string]any)["type"] != "object" { + t.Fatalf("planned tools alias caller state: %#v", plan.Request.Tools[0]) + } + + next := planner.Plan(messages, toolDefs, "medium") + if next.PrefixFingerprint == plan.PrefixFingerprint || next.Request.Tools[0].Name != "changed_tool" { + t.Fatalf("changed tool schema did not refresh snapshot: %#v", next) + } + if plan.Request.Tools[0].Name != "read_file" { + t.Fatalf("refresh mutated earlier request: %#v", plan.Request.Tools[0]) + } } func TestContextPlannerReportsInspectableBlocks(t *testing.T) {