diff --git a/pkg/agent/llm_call.go b/pkg/agent/llm_call.go index 90bcd26..99d4e97 100644 --- a/pkg/agent/llm_call.go +++ b/pkg/agent/llm_call.go @@ -136,6 +136,9 @@ func (al *AgentLoop) callLLM(ctx context.Context, st *iterationState, msgs []his toolsForCall := al.selectToolsForCall(ctx, st.sessionKey, st.streamChan, msgs) msgsForLLM := al.buildMsgsForLLM(ctx, st.sessionKey, st.iteration, msgs) + if al.RequestInvariant != nil { + st.sentMessages = msgsForLLM + } llmCtx := ctx if al.ThinkingBudget > 0 { diff --git a/pkg/agent/loop_iteration.go b/pkg/agent/loop_iteration.go index e69f529..7d9b8e8 100644 --- a/pkg/agent/loop_iteration.go +++ b/pkg/agent/loop_iteration.go @@ -59,7 +59,9 @@ func (al *AgentLoop) runIteration(ctx context.Context, sessionKey string, stream tracker: tracker, } + snap := al.snapshotRequest(*msgs) finalContent, result, err := al.callLLMWithRetry(ctx, st, msgsForLLM) + al.checkRequestInvariant(ctx, snap, *msgs, st.sentMessages) if err != nil { al.saveSession(ctx, sessionKey, *msgs) // A cancel that lands while the provider stream is in flight surfaces diff --git a/pkg/agent/loop_iteration_helpers.go b/pkg/agent/loop_iteration_helpers.go index 90e9f14..a93aa22 100644 --- a/pkg/agent/loop_iteration_helpers.go +++ b/pkg/agent/loop_iteration_helpers.go @@ -20,39 +20,59 @@ import ( // silence. The event is suppressed when nothing changed, so a run whose // context always fits still emits nothing. func (al *AgentLoop) enforceTokenBudget(ctx context.Context, sessionKey string, streamChan chan<- StreamEvent, iteration int, msgs []history.Message) []history.Message { - orig := msgs - if al.MaxTokenBudget <= 0 { - pruned, changes := pruneContextMessages(msgs, defaultProtectedEnds) - al.emitContextTrace(ctx, sessionKey, streamChan, iteration, ContextPolicyDefault, orig, pruned, changes) - return pruned + d := deriveRequestMessages(msgs, al.MaxTokenBudget) + for _, notice := range d.notices { + al.emit(ctx, sessionKey, streamChan, Event(ThoughtEvent{Message: notice})) + } + al.emitContextTrace(ctx, sessionKey, streamChan, iteration, d.policy, msgs, d.messages, d.changes) + return d.messages +} + +// contextDerivation is one run of the budget policy over stored history: +// the derived message list plus everything the emitting half needs to +// report what it did. +type contextDerivation struct { + messages []history.Message + policy ContextPolicy + changes []ContextRef + notices []string +} + +// deriveRequestMessages turns stored history into the message list a single +// LLM call should carry. It is the pure half of enforceTokenBudget: same +// input, same output, no emit, no receiver state. Purity is what lets the +// request invariant re-run it as an independent check — see +// AgentLoop.checkRequestInvariant. +func deriveRequestMessages(stored []history.Message, maxTokenBudget int) contextDerivation { + if maxTokenBudget <= 0 { + pruned, changes := pruneContextMessages(stored, defaultProtectedEnds) + return contextDerivation{messages: pruned, policy: ContextPolicyDefault, changes: changes} } + msgs := stored estToks := estimateTokens(msgs) - thresh := int(float64(al.MaxTokenBudget) * budgetWarnRatio) - policy := ContextPolicyDefault - var changes []ContextRef + thresh := int(float64(maxTokenBudget) * budgetWarnRatio) + d := contextDerivation{policy: ContextPolicyDefault} - if estToks > thresh && estToks <= al.MaxTokenBudget { - al.emit(ctx, sessionKey, streamChan, Event(ThoughtEvent{Message: fmt.Sprintf("Token budget near threshold (~%d >= %d). Truncating tool arguments.", estToks, thresh)})) + if estToks > thresh && estToks <= maxTokenBudget { + d.notices = append(d.notices, fmt.Sprintf("Token budget near threshold (~%d >= %d). Truncating tool arguments.", estToks, thresh)) var truncated []ContextRef msgs, truncated = truncateToolArguments(msgs) - changes = append(changes, truncated...) - policy = ContextPolicyBudgetWarn + d.changes = append(d.changes, truncated...) + d.policy = ContextPolicyBudgetWarn } depth := defaultProtectedEnds - if postTrim := estimateTokens(msgs); postTrim > al.MaxTokenBudget { - al.emit(ctx, sessionKey, streamChan, Event(ThoughtEvent{ - Message: fmt.Sprintf("Token budget exceeded (~%d est. tokens). Applying emergency context pruning.", postTrim), - })) + if postTrim := estimateTokens(msgs); postTrim > maxTokenBudget { + d.notices = append(d.notices, fmt.Sprintf("Token budget exceeded (~%d est. tokens). Applying emergency context pruning.", postTrim)) depth = emergencyProtectedEnds - policy = ContextPolicyBudgetEmergency + d.policy = ContextPolicyBudgetEmergency } pruned, prunedRefs := pruneContextMessages(msgs, depth) - changes = append(changes, prunedRefs...) - al.emitContextTrace(ctx, sessionKey, streamChan, iteration, policy, orig, pruned, changes) - return pruned + d.messages = pruned + d.changes = append(d.changes, prunedRefs...) + return d } // emitSoftLandingNudge fires the user-visible thought event marking the diff --git a/pkg/agent/loop_state.go b/pkg/agent/loop_state.go index a09cd19..d53a66d 100644 --- a/pkg/agent/loop_state.go +++ b/pkg/agent/loop_state.go @@ -19,6 +19,12 @@ type iterationState struct { specMap map[string]*speculativeExec specMu *sync.Mutex tracker *loopDetector + + // sentMessages records the exact list handed to the provider on the + // most recent attempt, for AgentLoop.checkRequestInvariant to compare + // against. Only populated when RequestInvariant is set. Written and + // read on the loop goroutine, between which no tool goroutine runs. + sentMessages []history.Message } // waveState carries the per-wave shared mutable state used by the diff --git a/pkg/agent/loop_stream.go b/pkg/agent/loop_stream.go index ac09b6e..4a6be71 100644 --- a/pkg/agent/loop_stream.go +++ b/pkg/agent/loop_stream.go @@ -444,6 +444,14 @@ type AgentLoop struct { // and saves the session before returning. MaxToolCallsPerSession int + // RequestInvariant, when set, receives a description of any broken + // request-path invariant: the pipeline writing through to stored + // history, or a request that a re-derivation cannot reproduce. Nil + // (default) disables the check entirely — no snapshot is taken and the + // cost is one nil comparison per iteration. Set it with + // WithRequestInvariant in development and staging. + RequestInvariant RequestViolationFunc + // AutoCacheSystem, when true, stamps Message.CacheHint=true on the first // system message of every LLM call. On Anthropic this promotes the // entire system prompt into the prompt-cache prefix, typically cutting diff --git a/pkg/agent/options.go b/pkg/agent/options.go index 6affbb3..7f92b80 100644 --- a/pkg/agent/options.go +++ b/pkg/agent/options.go @@ -262,6 +262,20 @@ func WithMaxToolCallsPerSession(n int) Option { return func(al *AgentLoop) { al.MaxToolCallsPerSession = n } } +// WithRequestInvariant enables a per-call check at the LLM boundary that +// the request path neither wrote through to stored history nor produced a +// request a re-derivation cannot reproduce. onViolation is called +// synchronously on the loop goroutine with a descriptive error; the turn +// continues either way, so a handler that fails the build (t.Fatal, panic +// in a staging binary) is the point. +// +// Off by default. Enabled, it costs one copy of the message headers plus +// one extra derivation per LLM call — cheap enough for staging, and not +// something to leave on in a latency-sensitive production path. +func WithRequestInvariant(onViolation RequestViolationFunc) Option { + return func(al *AgentLoop) { al.RequestInvariant = onViolation } +} + // WithoutAutoCacheSystem disables the auto-stamp of CacheHint=true on // the first system message of every LLM call. Default behavior is to // stamp (Anthropic prompt-cache prefix). Disable when you manage cache diff --git a/pkg/agent/request_invariant.go b/pkg/agent/request_invariant.go new file mode 100644 index 0000000..cd9d47c --- /dev/null +++ b/pkg/agent/request_invariant.go @@ -0,0 +1,161 @@ +package agent + +import ( + "context" + "fmt" + "strings" + + "github.com/hung12ct/gopheragent/pkg/history" +) + +// The request the loop sends is not stored history: deriveRequestMessages +// applies the budget policy to it, and buildMsgsForLLM then layers on the +// soft-landing, memory-note, tool-chaining, plan-mode and dynamic-context +// injections. Two properties keep that pipeline honest, and both are +// currently held by convention alone: +// +// 1. Stored history is read-only along the request path. Every stage +// returns a derived slice; none may write through to the caller's +// messages, which are the source of truth for what gets persisted. +// 2. The derivation is pure. Re-running it on the same stored history +// yields the same messages, so what a request carried stays +// reconstructable from the log plus the declared injections. +// +// Property 1 has been violated before — buildMsgsForLLM copies before +// stamping CacheHint precisely because the stamp used to leak into the +// caller's session-loaded slice. A comment cannot hold that; this can. + +// RequestViolationFunc receives a description of a broken request-path +// invariant. It runs synchronously on the loop goroutine, so keep it cheap; +// panicking or blocking here stalls the turn. +type RequestViolationFunc func(ctx context.Context, err error) + +// requestSnapshot is a copy of stored history taken before the request +// pipeline runs, kept so the loop can prove the pipeline did not write +// through to it. +type requestSnapshot struct { + stored []history.Message +} + +// snapshotRequest copies the fields the request path could plausibly +// rewrite. It is nil-cheap: with no invariant configured, no copy is taken +// and the whole mechanism costs one nil check per iteration. +// +// Parts are compared by length rather than deep-copied: the pruning and +// injection stages only ever rewrite Content, and copying media bytes on +// every call would cost more than the check is worth. +func (al *AgentLoop) snapshotRequest(stored []history.Message) *requestSnapshot { + if al.RequestInvariant == nil { + return nil + } + cp := make([]history.Message, len(stored)) + copy(cp, stored) + return &requestSnapshot{stored: cp} +} + +// checkRequestInvariant verifies both properties and reports any violation +// through al.RequestInvariant. snap is nil when the invariant is disabled, +// which makes this a single branch on the hot path. +// +// sent is the message list handed to the provider; stored is the loop's +// live history slice after the call returned. +func (al *AgentLoop) checkRequestInvariant(ctx context.Context, snap *requestSnapshot, stored, sent []history.Message) { + if snap == nil || al.RequestInvariant == nil { + return + } + if err := storedUnchanged(snap.stored, stored); err != nil { + al.RequestInvariant(ctx, fmt.Errorf("agent: request path wrote through to stored history: %w", err)) + } + if err := derivationReproduces(snap.stored, al.MaxTokenBudget, sent); err != nil { + al.RequestInvariant(ctx, fmt.Errorf("agent: request diverges from a re-derivation of stored history: %w", err)) + } +} + +// storedUnchanged reports whether the request pipeline left the caller's +// history slice exactly as it found it. +func storedUnchanged(before, after []history.Message) error { + if len(before) != len(after) { + return fmt.Errorf("message count changed from %d to %d", len(before), len(after)) + } + for i := range before { + if err := sameMessage(before[i], after[i]); err != nil { + return fmt.Errorf("message %d (%s): %w", i, before[i].Role, err) + } + } + return nil +} + +// derivationReproduces re-runs the pure derivation over the pre-call +// snapshot and checks that the conversation the provider received is still +// exactly the derived one. +// +// System messages are excluded: they are the declared framing surface, and +// four separate stages legitimately prepend to or extend them. What must +// survive untouched is the conversation itself — every derived user, +// assistant and tool message, in order, byte for byte. A request may also +// carry the dynamic-context injection, which is admitted by its sentinel; +// anything else appearing in conversation position is content reaching the +// model that no re-derivation can account for. +func derivationReproduces(snapshot []history.Message, maxTokenBudget int, sent []history.Message) error { + want := conversationOf(deriveRequestMessages(snapshot, maxTokenBudget).messages) + got := conversationOf(sent) + + next := 0 + for _, m := range got { + if next < len(want) { + if err := sameMessage(want[next], m); err == nil { + next++ + continue + } + } + if strings.Contains(m.Content, dynamicContextSentinel) { + continue + } + if next >= len(want) { + return fmt.Errorf("request carries an extra %s message not derived from stored history and not a declared injection", m.Role) + } + return fmt.Errorf("message %d: expected the derived %s message, got a %s message that diverges: %w", + next, want[next].Role, m.Role, sameMessage(want[next], m)) + } + if next != len(want) { + return fmt.Errorf("request dropped %d of %d derived conversation messages", len(want)-next, len(want)) + } + return nil +} + +// conversationOf returns the non-system messages: the part of a request +// that must be reconstructable from stored history. +func conversationOf(msgs []history.Message) []history.Message { + out := make([]history.Message, 0, len(msgs)) + for _, m := range msgs { + if m.Role != "system" { + out = append(out, m) + } + } + return out +} + +// sameMessage compares the fields that decide what a provider receives and +// what a replay would reconstruct. Deliberately not reflect.DeepEqual: +// Parts can carry megabytes of media, and CacheHint is stamped on the +// request copy by design. +func sameMessage(a, b history.Message) error { + switch { + case a.Role != b.Role: + return fmt.Errorf("role %q became %q", a.Role, b.Role) + case a.Content != b.Content: + return fmt.Errorf("content changed (%d chars became %d)", len(a.Content), len(b.Content)) + case a.ToolCallID != b.ToolCallID: + return fmt.Errorf("tool_call_id %q became %q", a.ToolCallID, b.ToolCallID) + case len(a.ToolCalls) != len(b.ToolCalls): + return fmt.Errorf("tool_calls count %d became %d", len(a.ToolCalls), len(b.ToolCalls)) + case len(a.Parts) != len(b.Parts): + return fmt.Errorf("parts count %d became %d", len(a.Parts), len(b.Parts)) + } + for i := range a.ToolCalls { + if a.ToolCalls[i].ID != b.ToolCalls[i].ID || a.ToolCalls[i].Arguments != b.ToolCalls[i].Arguments { + return fmt.Errorf("tool_call %d changed", i) + } + } + return nil +} diff --git a/pkg/agent/request_invariant_test.go b/pkg/agent/request_invariant_test.go new file mode 100644 index 0000000..081fff5 --- /dev/null +++ b/pkg/agent/request_invariant_test.go @@ -0,0 +1,207 @@ +package agent + +import ( + "context" + "strings" + "sync" + "testing" + + "github.com/hung12ct/gopheragent/pkg/history" +) + +// violationRecorder collects invariant reports so a test can assert on them. +type violationRecorder struct { + mu sync.Mutex + errs []error +} + +func (v *violationRecorder) record(_ context.Context, err error) { + v.mu.Lock() + defer v.mu.Unlock() + v.errs = append(v.errs, err) +} + +func (v *violationRecorder) all() []error { + v.mu.Lock() + defer v.mu.Unlock() + return append([]error(nil), v.errs...) +} + +// A noisy invariant is a useless one. Exercise the paths that legitimately +// reshape a request — plan mode prepends a system message, dynamic context +// appends a user message, the budget policy prunes and truncates — and +// require silence on all of them. +func TestRequestInvariant_CleanRunReportsNothing(t *testing.T) { + cases := []struct { + name string + apply func(*AgentLoop) + }{ + {"defaults", func(*AgentLoop) {}}, + {"budget forces pruning", func(al *AgentLoop) { al.MaxTokenBudget = 1 }}, + {"dynamic context injects a user message", func(al *AgentLoop) { + al.DynamicContext = func(context.Context, string) string { return "extra grounding" } + }}, + {"memory notes rewrite the system message", func(al *AgentLoop) { + al.MaxTokenBudget = 40 + }}, + {"parallel cap binds", func(al *AgentLoop) { + // A wave wider than the cap serialises some calls, so results + // land out of dispatch order. They must still be committed in + // model order for the next turn's derivation to match. + al.MaxParallelToolCalls = 2 + }}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + rec := &violationRecorder{} + ct := &countingTool{name: "counter"} + provider := &scriptProvider{turns: []LLMResult{ + {ToolCalls: fanoutCalls(5)}, + {Content: "final"}, + }} + loop, _ := setup(provider, ct) + WithRequestInvariant(rec.record)(loop) + tc.apply(loop) + + if _, err := loop.RunIteration(context.Background(), "s1", "go"); err != nil { + t.Fatalf("unexpected error: %v", err) + } + for _, got := range rec.all() { + t.Errorf("unexpected violation: %v", got) + } + }) + } +} + +func TestRequestInvariant_DisabledTakesNoSnapshot(t *testing.T) { + loop, _ := setup(&scriptProvider{turns: []LLMResult{{Content: "hi"}}}) + + if snap := loop.snapshotRequest([]history.Message{{Role: "user", Content: "x"}}); snap != nil { + t.Fatal("snapshotRequest allocated with no invariant configured") + } + // The check must also be inert, not panic, on a nil snapshot. + loop.checkRequestInvariant(context.Background(), nil, nil, nil) +} + +func TestStoredUnchanged(t *testing.T) { + base := []history.Message{ + {Role: "system", Content: "sys"}, + {Role: "user", Content: "hello"}, + {Role: "assistant", ToolCalls: []history.ToolCall{{ID: "a", Name: "t", Arguments: `{}`}}}, + {Role: "tool", Content: "result", ToolCallID: "a"}, + } + clone := func() []history.Message { + out := make([]history.Message, len(base)) + copy(out, base) + return out + } + + if err := storedUnchanged(base, clone()); err != nil { + t.Fatalf("identical slices reported a change: %v", err) + } + + t.Run("rewritten content", func(t *testing.T) { + after := clone() + after[3].Content = "truncated" + err := storedUnchanged(base, after) + if err == nil || !strings.Contains(err.Error(), "content changed") { + t.Fatalf("err = %v, want a content-changed report", err) + } + }) + + t.Run("dropped message", func(t *testing.T) { + if err := storedUnchanged(base, clone()[:2]); err == nil { + t.Fatal("a shortened slice reported no change") + } + }) + + t.Run("rewritten tool call", func(t *testing.T) { + after := clone() + after[2].ToolCalls = []history.ToolCall{{ID: "a", Name: "t", Arguments: `{"x":1}`}} + if err := storedUnchanged(base, after); err == nil { + t.Fatal("a rewritten tool call reported no change") + } + }) +} + +func TestDerivationReproduces(t *testing.T) { + stored := []history.Message{ + {Role: "system", Content: "sys"}, + {Role: "user", Content: "hello"}, + {Role: "assistant", Content: "hi"}, + } + derived := deriveRequestMessages(stored, 0).messages + + t.Run("system framing may be reshaped", func(t *testing.T) { + sent := append([]history.Message{{Role: "system", Content: "plan mode"}}, derived...) + sent[1].Content = "sys + memory notes" + if err := derivationReproduces(stored, 0, sent); err != nil { + t.Fatalf("system reshaping reported a violation: %v", err) + } + }) + + t.Run("declared dynamic context is admitted", func(t *testing.T) { + sent := append(append([]history.Message(nil), derived...), + history.Message{Role: "user", Content: dynamicContextSentinel + "\nextra"}) + if err := derivationReproduces(stored, 0, sent); err != nil { + t.Fatalf("declared injection reported a violation: %v", err) + } + }) + + t.Run("undeclared injection is caught", func(t *testing.T) { + sent := append(append([]history.Message(nil), derived...), + history.Message{Role: "user", Content: "smuggled instruction"}) + if err := derivationReproduces(stored, 0, sent); err == nil { + t.Fatal("an unlogged user message reported no violation") + } + }) + + t.Run("rewritten conversation is caught", func(t *testing.T) { + sent := append([]history.Message(nil), derived...) + for i := range sent { + if sent[i].Role == "user" { + sent[i].Content = "tampered" + } + } + if err := derivationReproduces(stored, 0, sent); err == nil { + t.Fatal("a rewritten user message reported no violation") + } + }) + + t.Run("dropped conversation is caught", func(t *testing.T) { + if err := derivationReproduces(stored, 0, derived[:1]); err == nil { + t.Fatal("a dropped message reported no violation") + } + }) +} + +// The derivation must be a pure function of its inputs: same stored history +// in, same messages out, and the caller's slice untouched. Everything the +// invariant reports rests on this. +func TestDeriveRequestMessages_IsPure(t *testing.T) { + long := strings.Repeat("x", 40_000) + stored := []history.Message{ + {Role: "system", Content: "sys"}, + {Role: "user", Content: "hello"}, + {Role: "tool", Content: long, ToolCallID: "a"}, + {Role: "assistant", Content: "done"}, + } + before := make([]history.Message, len(stored)) + copy(before, stored) + + for _, budget := range []int{0, 1, 50, 1_000_000} { + first := deriveRequestMessages(stored, budget) + second := deriveRequestMessages(stored, budget) + + if err := storedUnchanged(before, stored); err != nil { + t.Fatalf("budget %d: derivation mutated its input: %v", budget, err) + } + if err := storedUnchanged(first.messages, second.messages); err != nil { + t.Fatalf("budget %d: derivation is not deterministic: %v", budget, err) + } + if first.policy != second.policy { + t.Fatalf("budget %d: policy %v then %v", budget, first.policy, second.policy) + } + } +}