| 1 | package agent |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "reflect" |
| 7 | "strings" |
| 8 | "testing" |
| 9 | |
| 10 | "reasonix/internal/event" |
| 11 | "reasonix/internal/extension" |
| 12 | "reasonix/internal/extension/protocol" |
| 13 | "reasonix/internal/provider" |
| 14 | "reasonix/internal/tool" |
| 15 | ) |
| 16 | |
| 17 | func pinnedTestState(t *testing.T, snapshot PinnedContextSnapshot) pinnedContextState { |
| 18 | t.Helper() |
| 19 | state, err := normalizePinnedContextSnapshot(snapshot) |
| 20 | if err != nil { |
| 21 | t.Fatal(err) |
| 22 | } |
| 23 | return state |
| 24 | } |
| 25 | |
| 26 | func pinnedTestMessage(t *testing.T, next, previous pinnedContextState, kind string) provider.Message { |
| 27 | t.Helper() |
| 28 | encoded, err := encodePinnedContextRevision(next, previous, kind) |
| 29 | if err != nil { |
| 30 | t.Fatal(err) |
| 31 | } |
| 32 | return provider.Message{Role: provider.RoleUser, Origin: provider.MessageOriginHost, Content: string(encoded)} |
| 33 | } |
| 34 | |
| 35 | func TestPinnedContextRevisionWaitsForAcceptedTurn(t *testing.T) { |
| 36 | client := &fakeDispatchClient{interceptFn: func(ev protocol.InterceptEvent, _ json.RawMessage) (protocol.InterceptResult, error) { |
| 37 | if ev == protocol.EventAgentBeforeStart { |
| 38 | return blockWith("no runs today"), nil |
| 39 | } |
| 40 | return protocol.InterceptResult{Decision: protocol.DecisionContinue}, nil |
| 41 | }} |
| 42 | dispatcher := newExtDispatcher(client, true, nil, extension.PointAgentBeforeStart) |
| 43 | prov := &mockProvider{name: "p"} |
| 44 | session := NewSession("sys") |
| 45 | a := New(prov, tool.NewRegistry(), session, Options{Extensions: dispatcher}, event.Discard) |
| 46 | if err := a.StagePinnedContext(PinnedContextSnapshot{Files: []PinnedContextFile{{Path: "blocked.md", Content: "must not append"}}}); err != nil { |
| 47 | t.Fatal(err) |
| 48 | } |
| 49 | if err := a.Run(context.Background(), "hello"); err == nil || !strings.Contains(err.Error(), "no runs today") { |
| 50 | t.Fatalf("Run err = %v, want block reason", err) |
| 51 | } |
| 52 | if len(prov.requests) != 0 || len(session.Messages) != 1 { |
| 53 | t.Fatalf("blocked turn persisted revision or reached provider: messages=%d requests=%d", len(session.Messages), len(prov.requests)) |
| 54 | } |
| 55 | } |
| 56 | |
| 57 | func TestPinnedContextRevisionDeterministicEscapedXML(t *testing.T) { |
| 58 | first := pinnedTestState(t, PinnedContextSnapshot{Files: []PinnedContextFile{ |
| 59 | {Path: "z & quote.md", Content: "before </pinned_context_revision>\x01 after"}, |
| 60 | {Path: "a.md", Content: "A"}, |
| 61 | }}) |
| 62 | second := pinnedTestState(t, PinnedContextSnapshot{Files: []PinnedContextFile{ |
| 63 | {Path: "a.md", Content: "A"}, |
| 64 | {Path: "z & quote.md", Content: "before </pinned_context_revision>\x01 after"}, |
| 65 | }}) |
| 66 | one, err := encodePinnedContextRevision(first, emptyPinnedContextState(), "checkpoint") |
| 67 | if err != nil { |
| 68 | t.Fatal(err) |
| 69 | } |
| 70 | two, err := encodePinnedContextRevision(second, emptyPinnedContextState(), "checkpoint") |
| 71 | if err != nil { |
| 72 | t.Fatal(err) |
| 73 | } |
| 74 | if string(one) != string(two) || first.Revision != second.Revision { |
| 75 | t.Fatal("equivalent snapshots produced different revision bytes") |
| 76 | } |
| 77 | encoded := string(one) |
| 78 | if strings.Contains(encoded, "</pinned_context_revision>\x01") || !strings.Contains(encoded, "</pinned_context_revision>") || |
| 79 | !strings.Contains(encoded, "z & quote.md") || !strings.Contains(encoded, "�") { |
| 80 | t.Fatalf("revision XML was not safely encoded: %s", encoded) |
| 81 | } |
| 82 | applied := applyPinnedContextRevision(emptyPinnedContextState(), pinnedTestMessage(t, first, emptyPinnedContextState(), "checkpoint")) |
| 83 | if applied.Broken || applied.Revision != first.Revision || applied.Files["z & quote.md"].Content != first.Files["z & quote.md"].Content { |
| 84 | t.Fatalf("round-trip state = %+v", applied) |
| 85 | } |
| 86 | } |
| 87 | |
| 88 | func TestPinnedContextSnapshotRejectsMismatchedDigestAndInstruction(t *testing.T) { |
| 89 | if err := ValidatePinnedContextSnapshot(PinnedContextSnapshot{Files: []PinnedContextFile{{ |
| 90 | Path: "a.md", Content: "A", SHA256: "wrong", |
| 91 | }}}); err == nil { |
| 92 | t.Fatal("snapshot with mismatched digest was accepted") |
| 93 | } |
| 94 | state := pinnedTestState(t, PinnedContextSnapshot{Files: []PinnedContextFile{{Path: "a.md", Content: "A"}}}) |
| 95 | message := pinnedTestMessage(t, state, emptyPinnedContextState(), "checkpoint") |
| 96 | message.Content = strings.Replace(message.Content, pinnedContextRevisionInstruction, "changed instruction", 1) |
| 97 | if got := applyPinnedContextRevision(emptyPinnedContextState(), message); !got.Broken { |
| 98 | t.Fatalf("revision with changed instruction was accepted: %+v", got) |
| 99 | } |
| 100 | } |
| 101 | |
| 102 | func TestPinnedContextDeltaAddChangeRemoveUnavailableRecover(t *testing.T) { |
| 103 | empty := emptyPinnedContextState() |
| 104 | initial := pinnedTestState(t, PinnedContextSnapshot{Files: []PinnedContextFile{ |
| 105 | {Path: "a.md", Content: "A1"}, {Path: "b.md", Content: "B"}, |
| 106 | }}) |
| 107 | applied := applyPinnedContextRevision(empty, pinnedTestMessage(t, initial, empty, "checkpoint")) |
| 108 | |
| 109 | unavailable := pinnedTestState(t, PinnedContextSnapshot{ |
| 110 | Files: []PinnedContextFile{{Path: "a.md", Content: "A2"}, {Path: "c.md", Content: "C"}}, |
| 111 | Issues: []PinnedContextIssue{{Path: "b.md", Reason: PinnedContextIssueReadFailed}}, |
| 112 | }) |
| 113 | delta := pinnedTestMessage(t, unavailable, initial, "delta") |
| 114 | if !strings.Contains(delta.Content, `<remove path="b.md"></remove>`) || !strings.Contains(delta.Content, `path="c.md"`) { |
| 115 | t.Fatalf("delta does not carry change and tombstone: %s", delta.Content) |
| 116 | } |
| 117 | applied = applyPinnedContextRevision(applied, delta) |
| 118 | if applied.Broken || applied.Revision != unavailable.Revision || len(applied.Files) != 2 || applied.Issues["b.md"] != PinnedContextIssueReadFailed { |
| 119 | t.Fatalf("unavailable state = %+v", applied) |
| 120 | } |
| 121 | |
| 122 | recovered := pinnedTestState(t, PinnedContextSnapshot{Files: []PinnedContextFile{{Path: "b.md", Content: "B2"}}}) |
| 123 | applied = applyPinnedContextRevision(applied, pinnedTestMessage(t, recovered, unavailable, "delta")) |
| 124 | checkpointApplied := applyPinnedContextRevision(empty, pinnedTestMessage(t, recovered, empty, "checkpoint")) |
| 125 | if applied.Broken || checkpointApplied.Broken || applied.Revision != checkpointApplied.Revision || |
| 126 | !reflect.DeepEqual(applied.Files, checkpointApplied.Files) || !reflect.DeepEqual(applied.Issues, checkpointApplied.Issues) { |
| 127 | t.Fatalf("delta and checkpoint disagree: delta=%+v checkpoint=%+v", applied, checkpointApplied) |
| 128 | } |
| 129 | wrongBase := delta |
| 130 | wrongBase.Content = strings.Replace(wrongBase.Content, `base_revision="`+initial.Revision, `base_revision="sha256:wrong`, 1) |
| 131 | if got := applyPinnedContextRevision(applyPinnedContextRevision(empty, pinnedTestMessage(t, initial, empty, "checkpoint")), wrongBase); !got.Broken { |
| 132 | t.Fatalf("delta with mismatched base was accepted: %+v", got) |
| 133 | } |
| 134 | } |
| 135 | |
| 136 | func TestPinnedContextRevisionTrustAndSelfHealing(t *testing.T) { |
| 137 | desired := pinnedTestState(t, PinnedContextSnapshot{Files: []PinnedContextFile{{Path: "a.md", Content: "A"}}}) |
| 138 | valid := pinnedTestMessage(t, desired, emptyPinnedContextState(), "checkpoint") |
| 139 | spoof := valid |
| 140 | spoof.Origin = provider.MessageOriginUser |
| 141 | if state := pinnedContextStateFromMessages([]provider.Message{spoof}); state.Seen || state.Broken { |
| 142 | t.Fatalf("user-authored spoof changed state: %+v", state) |
| 143 | } |
| 144 | |
| 145 | broken := valid |
| 146 | broken.Content = strings.Replace(broken.Content, `revision="`+desired.Revision, `revision="sha256:broken`, 1) |
| 147 | session := NewSession("system") |
| 148 | session.Add(broken) |
| 149 | a := New(nil, nil, session, Options{}, event.Discard) |
| 150 | if err := a.StagePinnedContext(PinnedContextSnapshot{Files: []PinnedContextFile{{Path: "a.md", Content: "A"}}}); err != nil { |
| 151 | t.Fatal(err) |
| 152 | } |
| 153 | plan, err := a.preparePinnedRevision() |
| 154 | if err != nil { |
| 155 | t.Fatal(err) |
| 156 | } |
| 157 | if plan.message == nil || !strings.Contains(plan.message.Content, `kind="checkpoint"`) { |
| 158 | t.Fatalf("broken chain did not self-heal with checkpoint: %+v", plan.message) |
| 159 | } |
| 160 | unknown := valid |
| 161 | unknown.Content = strings.Replace(unknown.Content, `schema_version="1"`, `schema_version="99"`, 1) |
| 162 | if state := pinnedContextStateFromMessages([]provider.Message{unknown}); !state.Broken { |
| 163 | t.Fatalf("unknown schema did not damage the derived chain: %+v", state) |
| 164 | } |
| 165 | } |
| 166 | |
| 167 | func TestPinnedContextDerivedStateRescansAfterRewindAndSessionSwitch(t *testing.T) { |
| 168 | empty := emptyPinnedContextState() |
| 169 | stateA := pinnedTestState(t, PinnedContextSnapshot{Files: []PinnedContextFile{{Path: "a.md", Content: "A"}}}) |
| 170 | stateB := pinnedTestState(t, PinnedContextSnapshot{Files: []PinnedContextFile{{Path: "a.md", Content: "B"}}}) |
| 171 | checkpointA := pinnedTestMessage(t, stateA, empty, "checkpoint") |
| 172 | session := NewSession("system") |
| 173 | session.Add(checkpointA) |
| 174 | a := New(nil, nil, session, Options{}, event.Discard) |
| 175 | if err := a.StagePinnedContext(PinnedContextSnapshot{Files: []PinnedContextFile{{Path: "a.md", Content: "B"}}}); err != nil { |
| 176 | t.Fatal(err) |
| 177 | } |
| 178 | plan, err := a.preparePinnedRevision() |
| 179 | if err != nil || plan.message == nil { |
| 180 | t.Fatalf("initial delta plan = %+v, %v", plan, err) |
| 181 | } |
| 182 | session.AddBatch(*plan.message, provider.Message{Role: provider.RoleUser, Origin: provider.MessageOriginUser, Content: "turn"}) |
| 183 | a.commitPinnedRevisionPlan(plan) |
| 184 | |
| 185 | session.Rewrite([]provider.Message{{Role: provider.RoleSystem, Content: "system"}, checkpointA}, "rewind_truncate") |
| 186 | if err := a.StagePinnedContext(PinnedContextSnapshot{Files: []PinnedContextFile{{Path: "a.md", Content: "B"}}}); err != nil { |
| 187 | t.Fatal(err) |
| 188 | } |
| 189 | restored, err := a.preparePinnedRevision() |
| 190 | if err != nil || restored.message == nil { |
| 191 | t.Fatalf("rewind did not restore desired pinned state: %+v, %v", restored, err) |
| 192 | } |
| 193 | restoredState := applyPinnedContextRevision(applyPinnedContextRevision(empty, checkpointA), *restored.message) |
| 194 | if restoredState.Broken || restoredState.Revision != stateB.Revision { |
| 195 | t.Fatalf("rewind restoration event applied as %+v", restoredState) |
| 196 | } |
| 197 | |
| 198 | switched := NewSession("system") |
| 199 | switched.Add(pinnedTestMessage(t, stateB, empty, "checkpoint")) |
| 200 | a.SetSession(switched) |
| 201 | if err := a.StagePinnedContext(PinnedContextSnapshot{Files: []PinnedContextFile{{Path: "a.md", Content: "B"}}}); err != nil { |
| 202 | t.Fatal(err) |
| 203 | } |
| 204 | stable, err := a.preparePinnedRevision() |
| 205 | if err != nil || stable.message != nil { |
| 206 | t.Fatalf("session switch did not rebuild applied state: %+v, %v", stable, err) |
| 207 | } |
| 208 | } |
| 209 | |
| 210 | func TestPinnedContextSnapshotLimits(t *testing.T) { |
| 211 | tooMany := PinnedContextSnapshot{} |
| 212 | for i := range MaxPinnedContextFiles + 1 { |
| 213 | tooMany.Files = append(tooMany.Files, PinnedContextFile{Path: string(rune('a'+i)) + ".md", Content: "x"}) |
| 214 | } |
| 215 | if err := ValidatePinnedContextSnapshot(tooMany); err == nil { |
| 216 | t.Fatal("33 files were accepted") |
| 217 | } |
| 218 | if err := ValidatePinnedContextSnapshot(PinnedContextSnapshot{Files: []PinnedContextFile{{Path: "large.md", Content: strings.Repeat("x", MaxPinnedContextFileBytes+1)}}}); err == nil { |
| 219 | t.Fatal("oversized file was accepted") |
| 220 | } |
| 221 | total := PinnedContextSnapshot{Files: []PinnedContextFile{ |
| 222 | {Path: "a.md", Content: strings.Repeat("a", MaxPinnedContextFileBytes)}, |
| 223 | {Path: "b.md", Content: strings.Repeat("b", MaxPinnedContextFileBytes)}, |
| 224 | {Path: "c.md", Content: strings.Repeat("c", MaxPinnedContextFileBytes)}, |
| 225 | {Path: "d.md", Content: strings.Repeat("d", MaxPinnedContextFileBytes)}, |
| 226 | }} |
| 227 | if err := ValidatePinnedContextSnapshot(total); err == nil { |
| 228 | t.Fatal("oversized serialized checkpoint was accepted") |
| 229 | } |
| 230 | } |
| 231 | |
| 232 | func TestPinnedContextRevisionIsMarkedInDisplayIndex(t *testing.T) { |
| 233 | snapshot := PinnedContextSnapshot{Files: []PinnedContextFile{{Path: "guide.md", Content: "guide"}}} |
| 234 | state := pinnedTestState(t, snapshot) |
| 235 | message := pinnedTestMessage(t, state, pinnedContextState{}, "checkpoint") |
| 236 | idx := BuildSessionDisplayIndex([]provider.Message{ |
| 237 | {Role: provider.RoleSystem, Content: "system"}, |
| 238 | message, |
| 239 | {Role: provider.RoleUser, Origin: provider.MessageOriginUser, Content: "question"}, |
| 240 | }, 0, false, [32]byte{}) |
| 241 | if idx == nil || len(idx.Entries) != 3 { |
| 242 | t.Fatalf("display index = %+v", idx) |
| 243 | } |
| 244 | if !idx.Entries[1].PinnedContextRevision || idx.Entries[2].PinnedContextRevision { |
| 245 | t.Fatalf("pinned revision flags = %+v", idx.Entries) |
| 246 | } |
| 247 | } |
| 248 | |
| 249 | func TestSessionAddBatchCommitsOneTranscriptVersion(t *testing.T) { |
| 250 | session := NewSession("system") |
| 251 | before := session.TranscriptVersion() |
| 252 | session.AddBatch( |
| 253 | provider.Message{Role: provider.RoleUser, Origin: provider.MessageOriginHost, Content: "revision"}, |
| 254 | provider.Message{Role: provider.RoleUser, Origin: provider.MessageOriginUser, Content: "question"}, |
| 255 | ) |
| 256 | if got := session.TranscriptVersion(); got != before+1 { |
| 257 | t.Fatalf("transcript version = %d, want %d", got, before+1) |
| 258 | } |
| 259 | if got := session.Snapshot(); len(got) != 3 || got[1].Content != "revision" || got[2].Content != "question" { |
| 260 | t.Fatalf("atomic batch snapshot = %+v", got) |
| 261 | } |
| 262 | } |
| 263 | |
| 264 | func TestPinnedContextProjectionRebasesAtCanonicalCoverage(t *testing.T) { |
| 265 | empty := emptyPinnedContextState() |
| 266 | stateA := pinnedTestState(t, PinnedContextSnapshot{Files: []PinnedContextFile{{Path: "a.md", Content: "A"}}}) |
| 267 | stateB := pinnedTestState(t, PinnedContextSnapshot{Files: []PinnedContextFile{{Path: "a.md", Content: "B"}}}) |
| 268 | checkpointA := pinnedTestMessage(t, stateA, empty, "checkpoint") |
| 269 | deltaB := pinnedTestMessage(t, stateB, stateA, "delta") |
| 270 | canonical := []provider.Message{ |
| 271 | {Role: provider.RoleSystem, Content: "system"}, checkpointA, |
| 272 | {Role: provider.RoleUser, Origin: provider.MessageOriginUser, Content: "first"}, |
| 273 | {Role: provider.RoleAssistant, Content: "answer"}, deltaB, |
| 274 | {Role: provider.RoleUser, Origin: provider.MessageOriginUser, Content: "second"}, |
| 275 | } |
| 276 | projected := []provider.Message{ |
| 277 | canonical[0], checkpointA, formatSummaryMessage("summary"), |
| 278 | } |
| 279 | rebased, ok, err := rebasePinnedContextProjection(projected, canonical, 4) |
| 280 | if err != nil || !ok { |
| 281 | t.Fatalf("rebase = %v, ok=%v", err, ok) |
| 282 | } |
| 283 | if len(rebased) != 3 || !IsPinnedContextRevision(rebased[1]) || !isCompactionSummary(rebased[2]) { |
| 284 | t.Fatalf("rebased projection = %+v", rebased) |
| 285 | } |
| 286 | visible := append(append([]provider.Message(nil), rebased...), canonical[4:]...) |
| 287 | state := pinnedContextStateFromMessages(visible) |
| 288 | if state.Broken || state.Revision != stateB.Revision || state.Files["a.md"].Content != "B" { |
| 289 | t.Fatalf("checkpoint plus tail delta = %+v", state) |
| 290 | } |
| 291 | |
| 292 | a := &Agent{} |
| 293 | _, fold, retention := a.partitionFoldForProjection([]provider.Message{checkpointA, canonical[2]}) |
| 294 | if len(fold) != 1 || fold[0].Content != "first" || retention.Dropped != 1 { |
| 295 | t.Fatalf("summary fold contains pinned context: fold=%+v retention=%+v", fold, retention) |
| 296 | } |
| 297 | } |
| 298 | |
| 299 | func TestProjectionV4RejectsPinnedRevisionProvenanceTampering(t *testing.T) { |
| 300 | state := pinnedTestState(t, PinnedContextSnapshot{Files: []PinnedContextFile{{Path: "a.md", Content: "A"}}}) |
| 301 | checkpoint := pinnedTestMessage(t, state, emptyPinnedContextState(), "checkpoint") |
| 302 | canonical := []provider.Message{{Role: provider.RoleSystem, Content: "system"}, checkpoint} |
| 303 | projection := CompactionState{ |
| 304 | SchemaVersion: compactionStateSchemaV4, |
| 305 | Projection: ContextProjection{ |
| 306 | Messages: append([]provider.Message(nil), canonical...), |
| 307 | CoveredCount: len(canonical), |
| 308 | CoveredPrefixHash: coveredPrefixHash(canonical, len(canonical)), |
| 309 | PinnedContextHash: pinnedContextCoverageHash(canonical, len(canonical)), |
| 310 | }, |
| 311 | } |
| 312 | if !projectionContentValid(projection, canonical) { |
| 313 | t.Fatal("valid v4 pinned projection was rejected") |
| 314 | } |
| 315 | spoofed := append([]provider.Message(nil), canonical...) |
| 316 | spoofed[1].Origin = provider.MessageOriginUser |
| 317 | if coveredPrefixHash(spoofed, len(spoofed)) != projection.Projection.CoveredPrefixHash { |
| 318 | t.Fatal("test setup expected provider-visible hash to ignore local origin") |
| 319 | } |
| 320 | if projectionContentValid(projection, spoofed) { |
| 321 | t.Fatal("v4 projection accepted pinned revision with user provenance") |
| 322 | } |
| 323 | } |
| 324 | |
| 325 | func TestPinnedContextExplicitCompactionCheckpointsWithoutSummarizingBodies(t *testing.T) { |
| 326 | empty := emptyPinnedContextState() |
| 327 | stateA := pinnedTestState(t, PinnedContextSnapshot{Files: []PinnedContextFile{{Path: "a.md", Content: "PINNED_SECRET_A"}}}) |
| 328 | stateB := pinnedTestState(t, PinnedContextSnapshot{Files: []PinnedContextFile{{Path: "a.md", Content: "PINNED_SECRET_B"}}}) |
| 329 | checkpointA := pinnedTestMessage(t, stateA, empty, "checkpoint") |
| 330 | deltaB := pinnedTestMessage(t, stateB, stateA, "delta") |
| 331 | session := &Session{Messages: []provider.Message{ |
| 332 | {Role: provider.RoleSystem, Content: "system"}, checkpointA, |
| 333 | {Role: provider.RoleUser, Origin: provider.MessageOriginUser, Content: "old request"}, |
| 334 | {Role: provider.RoleAssistant, Content: strings.Repeat("old work ", 400)}, deltaB, |
| 335 | {Role: provider.RoleUser, Origin: provider.MessageOriginUser, Content: "retained boundary"}, |
| 336 | {Role: provider.RoleAssistant, Content: strings.Repeat("retained work ", 400)}, |
| 337 | }} |
| 338 | canonical := session.Snapshot() |
| 339 | prov := &fakeProvider{reply: "summary"} |
| 340 | a := New(prov, tool.NewRegistry(), session, Options{}, event.Discard) |
| 341 | |
| 342 | for index, anchor := range []string{"retained boundary", "new boundary"} { |
| 343 | if index == 1 { |
| 344 | stateC := pinnedTestState(t, PinnedContextSnapshot{Files: []PinnedContextFile{{Path: "a.md", Content: "PINNED_SECRET_C"}}}) |
| 345 | session.AddBatch( |
| 346 | pinnedTestMessage(t, stateC, stateB, "delta"), |
| 347 | provider.Message{Role: provider.RoleUser, Origin: provider.MessageOriginUser, Content: anchor}, |
| 348 | provider.Message{Role: provider.RoleAssistant, Content: strings.Repeat("new work ", 400)}, |
| 349 | ) |
| 350 | canonical = session.Snapshot() |
| 351 | } |
| 352 | result, err := a.CompressContext(context.Background(), tool.CompressRequest{Direction: "before", Anchor: anchor}) |
| 353 | if err != nil || result.Status != "ok" { |
| 354 | t.Fatalf("compression %d = %+v, %v", index+1, result, err) |
| 355 | } |
| 356 | if !reflect.DeepEqual(session.Snapshot(), canonical) { |
| 357 | t.Fatalf("compression %d changed canonical transcript", index+1) |
| 358 | } |
| 359 | if body := joinContents(prov.got); strings.Contains(body, "PINNED_SECRET_") { |
| 360 | t.Fatalf("compression %d copied pinned body into summary input: %s", index+1, body) |
| 361 | } |
| 362 | visible := a.modelVisibleMessages() |
| 363 | revisions := 0 |
| 364 | for _, message := range visible { |
| 365 | if IsPinnedContextRevision(message) { |
| 366 | revisions++ |
| 367 | } |
| 368 | } |
| 369 | state := pinnedContextStateFromMessages(visible) |
| 370 | want := stateB |
| 371 | if index == 1 { |
| 372 | want = pinnedTestState(t, PinnedContextSnapshot{Files: []PinnedContextFile{{Path: "a.md", Content: "PINNED_SECRET_C"}}}) |
| 373 | } |
| 374 | if revisions != 1 || state.Broken || state.Revision != want.Revision { |
| 375 | t.Fatalf("compression %d visible pinned state = %+v, revisions=%d", index+1, state, revisions) |
| 376 | } |
| 377 | if got := a.sess.compactionState.SchemaVersion; got != compactionStateSchemaV4 { |
| 378 | t.Fatalf("compression %d schema = %d, want v4", index+1, got) |
| 379 | } |
| 380 | } |
| 381 | } |
| 382 | |
| 383 | func TestSafeSummaryPrefixBudgetExcludesPinnedRevisionBody(t *testing.T) { |
| 384 | state := pinnedTestState(t, PinnedContextSnapshot{Files: []PinnedContextFile{{ |
| 385 | Path: "large.md", Content: strings.Repeat("p", 48*1024), |
| 386 | }}}) |
| 387 | checkpoint := pinnedTestMessage(t, state, emptyPinnedContextState(), "checkpoint") |
| 388 | messages := []provider.Message{ |
| 389 | {Role: provider.RoleSystem, Content: "system"}, checkpoint, |
| 390 | {Role: provider.RoleUser, Origin: provider.MessageOriginUser, Content: "old task"}, |
| 391 | {Role: provider.RoleAssistant, Content: strings.Repeat("work ", 100)}, |
| 392 | {Role: provider.RoleUser, Origin: provider.MessageOriginUser, Content: "recent task"}, |
| 393 | } |
| 394 | a := &Agent{ |
| 395 | agentConfig: agentConfig{contextWindow: 10_000}, |
| 396 | svc: agentServices{prov: &overflowSummaryProvider{}}, |
| 397 | sess: sessionRuntime{conversation: &Session{Messages: messages}}, |
| 398 | } |
| 399 | if end := a.maximumSafeSummaryPrefixEnd(messages, 1, 4, ""); end != 4 { |
| 400 | t.Fatalf("safe summary prefix ended at %d, want 4 after excluding pinned body", end) |
| 401 | } |
| 402 | } |
| 403 | |
| 404 | type pinnedPrefixProvider struct { |
| 405 | requests []provider.Request |
| 406 | } |
| 407 | |
| 408 | func (p *pinnedPrefixProvider) Name() string { return "pinned-prefix" } |
| 409 | |
| 410 | func (p *pinnedPrefixProvider) Stream(_ context.Context, req provider.Request) (<-chan provider.Chunk, error) { |
| 411 | p.requests = append(p.requests, req) |
| 412 | ch := make(chan provider.Chunk, 2) |
| 413 | ch <- provider.Chunk{Type: provider.ChunkText, Text: "ok"} |
| 414 | ch <- provider.Chunk{Type: provider.ChunkDone} |
| 415 | close(ch) |
| 416 | return ch, nil |
| 417 | } |
| 418 | |
| 419 | func TestPinnedContextProviderRequestsStayAppendOnly(t *testing.T) { |
| 420 | for _, strict := range []bool{false, true} { |
| 421 | t.Run(map[bool]string{false: "ordinary", true: "strict-alternating"}[strict], func(t *testing.T) { |
| 422 | prov := &pinnedPrefixProvider{} |
| 423 | session := NewSession("system") |
| 424 | a := New(prov, nil, session, Options{StrictAlternatingRoles: strict}, event.Discard) |
| 425 | content := "A" |
| 426 | for _, input := range []string{"first", "second", "third"} { |
| 427 | if input == "third" { |
| 428 | content = "B" |
| 429 | } |
| 430 | if err := a.StagePinnedContext(PinnedContextSnapshot{Files: []PinnedContextFile{{Path: "a.md", Content: content}}}); err != nil { |
| 431 | t.Fatal(err) |
| 432 | } |
| 433 | if err := a.Run(context.Background(), input); err != nil { |
| 434 | t.Fatal(err) |
| 435 | } |
| 436 | } |
| 437 | if len(prov.requests) != 3 { |
| 438 | t.Fatalf("requests = %d", len(prov.requests)) |
| 439 | } |
| 440 | for i := 1; i < len(prov.requests); i++ { |
| 441 | previous := prov.requests[i-1].Messages |
| 442 | current := prov.requests[i].Messages |
| 443 | if len(current) < len(previous) || !reflect.DeepEqual(current[:len(previous)], previous) { |
| 444 | t.Fatalf("request %d does not preserve request %d as exact prefix", i, i-1) |
| 445 | } |
| 446 | } |
| 447 | if session.RewriteVersion() != 0 || len(session.DrainContentRewriteReasons()) != 0 { |
| 448 | t.Fatalf("revision append recorded a rewrite: version=%d", session.RewriteVersion()) |
| 449 | } |
| 450 | }) |
| 451 | } |
| 452 | } |
| 453 |