| 1 | // CompactionBench measures what repeated compaction costs and what it loses. |
| 2 | // Both arms drive the real agent compaction path over a session that grows one |
| 3 | // generation at a time: |
| 4 | // |
| 5 | // -mode=cost offline: what each fold costs and whether any single |
| 6 | // summarizer call can still overflow the window |
| 7 | // -mode=fidelity real provider: which planted facts survive N folds, |
| 8 | // scored against a full-history control |
| 9 | package main |
| 10 | |
| 11 | import ( |
| 12 | "context" |
| 13 | "encoding/json" |
| 14 | "flag" |
| 15 | "fmt" |
| 16 | "os" |
| 17 | "path/filepath" |
| 18 | "strings" |
| 19 | "time" |
| 20 | "unicode/utf8" |
| 21 | |
| 22 | "reasonix/internal/ablation" |
| 23 | "reasonix/internal/agent" |
| 24 | "reasonix/internal/event" |
| 25 | "reasonix/internal/provider" |
| 26 | _ "reasonix/internal/provider/openai" |
| 27 | "reasonix/internal/tool" |
| 28 | ) |
| 29 | |
| 30 | const ( |
| 31 | realModel = "deepseek-v4-flash" |
| 32 | realBaseURL = "https://api.deepseek.com" |
| 33 | // probeAnswerTokens must cover a thinking model's reasoning plus the short |
| 34 | // answer; too small and every probe scores as lost. |
| 35 | probeAnswerTokens = 2048 |
| 36 | ) |
| 37 | |
| 38 | func main() { |
| 39 | mode := flag.String("mode", "cost", "cost | fidelity") |
| 40 | gens := flag.Int("gens", 8, "generations of work+compaction to run") |
| 41 | report := flag.String("report", "1,2,4,8", "generations to report on") |
| 42 | window := flag.Int("window", 128_000, "context window in tokens") |
| 43 | control := flag.Bool("control", true, "fidelity: also score probes against full history") |
| 44 | arm := flag.String("arm", "full", "full | incremental: re-derive each digest from canonical, or fold the previous projection") |
| 45 | snip := flag.Bool("snip", false, "legacy no-op: automatic snip projections are gone; kept so old scripts do not fail") |
| 46 | out := flag.String("out", "", "write the JSON report here") |
| 47 | flag.Parse() |
| 48 | |
| 49 | var ( |
| 50 | res []genResult |
| 51 | err error |
| 52 | ) |
| 53 | a := arms{incremental: *arm == "incremental", snip: *snip} |
| 54 | switch { |
| 55 | case *arm != "full" && *arm != "incremental": |
| 56 | err = fmt.Errorf("unknown arm %q", *arm) |
| 57 | case *mode == "cost": |
| 58 | res, err = runCost(*gens, *window, a) |
| 59 | case *mode == "fidelity": |
| 60 | res, err = runFidelity(*gens, *window, *control, a) |
| 61 | default: |
| 62 | err = fmt.Errorf("unknown mode %q", *mode) |
| 63 | } |
| 64 | if err != nil { |
| 65 | fmt.Fprintln(os.Stderr, err) |
| 66 | os.Exit(1) |
| 67 | } |
| 68 | printReport(*mode+" / "+*arm, res, reportAt(*report)) |
| 69 | if *out != "" { |
| 70 | b, _ := json.MarshalIndent(map[string]any{"mode": *mode, "arm": *arm, "window": *window, "generations": res}, "", " ") |
| 71 | if werr := os.WriteFile(*out, append(b, '\n'), 0o644); werr != nil { |
| 72 | fmt.Fprintln(os.Stderr, werr) |
| 73 | os.Exit(1) |
| 74 | } |
| 75 | } |
| 76 | } |
| 77 | |
| 78 | // genResult is one generation: the fold that ran and what it cost or lost. |
| 79 | type genResult struct { |
| 80 | Gen int `json:"gen"` |
| 81 | CanonicalTokens int `json:"canonical_tokens"` |
| 82 | ProjectionTokens int `json:"projection_tokens"` |
| 83 | SummarizerCalls int `json:"summarizer_calls"` |
| 84 | SummarizerInput int `json:"summarizer_input_tokens"` |
| 85 | LargestCall int `json:"largest_call_tokens"` |
| 86 | Mode string `json:"mode,omitempty"` |
| 87 | SnippedResults int `json:"snipped_results,omitempty"` |
| 88 | SnippedChars int `json:"snipped_chars,omitempty"` |
| 89 | Seconds float64 `json:"seconds"` |
| 90 | Error string `json:"error,omitempty"` |
| 91 | Survived map[string]int `json:"survived,omitempty"` // probe class -> 1 kept, 0 lost |
| 92 | ControlOK map[string]int `json:"control_ok,omitempty"` // same probes against full history |
| 93 | // What the model actually said, so a score can be audited rather than trusted. |
| 94 | Answers map[string]string `json:"answers,omitempty"` |
| 95 | ControlAnswers map[string]string `json:"control_answers,omitempty"` |
| 96 | } |
| 97 | |
| 98 | type harness struct { |
| 99 | sess *agent.Session |
| 100 | agentA *agent.Agent |
| 101 | path string |
| 102 | calls *callRecorder |
| 103 | snip bool |
| 104 | } |
| 105 | |
| 106 | func newHarness(t *testingDir, p provider.Provider, window int, rec *callRecorder, arm arms) *harness { |
| 107 | sess := newSession() |
| 108 | path := filepath.Join(t.dir, "session.jsonl") |
| 109 | a := agent.New(p, tool.NewRegistry(), sess, agent.Options{ |
| 110 | ContextWindow: window, |
| 111 | ArchiveDir: filepath.Join(t.dir, "archive"), |
| 112 | SessionPath: path, |
| 113 | RecentKeep: 4, |
| 114 | // boot's default when cfg.Agent.Keep is unset; without it the bench |
| 115 | // would measure a configuration no real session runs. |
| 116 | KeepPolicy: agent.KeepErrors, |
| 117 | Ablation: foldArm(arm.incremental), |
| 118 | }, rec.sink()) |
| 119 | return &harness{sess: sess, agentA: a, path: path, calls: rec, snip: arm.snip} |
| 120 | } |
| 121 | |
| 122 | // foldArm switches full re-derivation off, which is what makes a fold read the |
| 123 | // previous projection instead of the canonical transcript. |
| 124 | // arms selects the maintenance behaviour under test. Snipping is off by default |
| 125 | // so a run stays comparable with baselines recorded before it existed. |
| 126 | type arms struct { |
| 127 | incremental bool |
| 128 | snip bool |
| 129 | } |
| 130 | |
| 131 | func foldArm(incremental bool) ablation.Set { |
| 132 | if incremental { |
| 133 | return ablation.New(ablation.FullFold) |
| 134 | } |
| 135 | return ablation.Set{} |
| 136 | } |
| 137 | |
| 138 | // runGeneration grows the session and folds it, returning what that fold cost. |
| 139 | func (h *harness) runGeneration(ctx context.Context, gen int, probes []probe) genResult { |
| 140 | growSession(h.sess, gen, probes) |
| 141 | r := genResult{Gen: gen, CanonicalTokens: estimateTokens(renderAll(h.sess.Snapshot()))} |
| 142 | |
| 143 | h.calls.reset() |
| 144 | start := time.Now() |
| 145 | if h.snip { |
| 146 | // SnipStaleToolResults is intentionally a no-op; record zeros for |
| 147 | // report schema compatibility with pre-content-driven baselines. |
| 148 | st, serr := h.agentA.SnipStaleToolResults() |
| 149 | if serr != nil { |
| 150 | r.Error = serr.Error() |
| 151 | } |
| 152 | r.SnippedResults, r.SnippedChars = st.Results, st.SavedChars |
| 153 | } |
| 154 | err := h.agentA.CompactNow(ctx, "") |
| 155 | r.Seconds = time.Since(start).Seconds() |
| 156 | if err != nil { |
| 157 | r.Error = err.Error() |
| 158 | } |
| 159 | r.SummarizerCalls = len(h.calls.calls) |
| 160 | for _, c := range h.calls.calls { |
| 161 | r.SummarizerInput += c.tokens |
| 162 | r.LargestCall = max(r.LargestCall, c.tokens) |
| 163 | } |
| 164 | if st, ok, sterr := agent.LoadCompactionState(h.path); sterr == nil && ok { |
| 165 | r.ProjectionTokens = st.Projection.ProjectionTokens |
| 166 | if st.LastReceipt != nil && st.LastReceipt.Action == "summary" { |
| 167 | r.Mode = agent.CompactionModeSummarized |
| 168 | } else if st.LastMode != "" { |
| 169 | r.Mode = st.LastMode |
| 170 | } |
| 171 | } |
| 172 | return r |
| 173 | } |
| 174 | |
| 175 | func runCost(gens, window int, a arms) ([]genResult, error) { |
| 176 | dir, cleanup, err := tempDir() |
| 177 | if err != nil { |
| 178 | return nil, err |
| 179 | } |
| 180 | defer cleanup() |
| 181 | |
| 182 | rec := &callRecorder{} |
| 183 | p := &scriptedProvider{rec: rec, reply: syntheticDigest, window: window} |
| 184 | h := newHarness(dir, p, window, rec, a) |
| 185 | |
| 186 | var out []genResult |
| 187 | for gen := range gens { |
| 188 | out = append(out, h.runGeneration(context.Background(), gen, probeSuite())) |
| 189 | } |
| 190 | return out, nil |
| 191 | } |
| 192 | |
| 193 | func runFidelity(gens, window int, control bool, a arms) ([]genResult, error) { |
| 194 | key := os.Getenv("DEEPSEEK_API_KEY") |
| 195 | if key == "" { |
| 196 | return nil, fmt.Errorf("fidelity mode needs DEEPSEEK_API_KEY") |
| 197 | } |
| 198 | p, err := provider.New("openai", provider.Config{Name: "compactionbench", BaseURL: realBaseURL, Model: realModel, APIKey: key}) |
| 199 | if err != nil { |
| 200 | return nil, err |
| 201 | } |
| 202 | dir, cleanup, cerr := tempDir() |
| 203 | if cerr != nil { |
| 204 | return nil, cerr |
| 205 | } |
| 206 | defer cleanup() |
| 207 | |
| 208 | rec := &callRecorder{} |
| 209 | h := newHarness(dir, &recordingProvider{inner: p, rec: rec}, window, rec, a) |
| 210 | probes := probeSuite() |
| 211 | |
| 212 | ctx := context.Background() |
| 213 | var out []genResult |
| 214 | for gen := range gens { |
| 215 | r := h.runGeneration(ctx, gen, probes) |
| 216 | r.Survived, r.ControlOK = map[string]int{}, map[string]int{} |
| 217 | r.Answers, r.ControlAnswers = map[string]string{}, map[string]string{} |
| 218 | visible, verr := visibleContext(h.path, h.sess) |
| 219 | if verr != nil { |
| 220 | return nil, verr |
| 221 | } |
| 222 | for _, probe := range probes { |
| 223 | if probe.settledAt() > gen { |
| 224 | continue |
| 225 | } |
| 226 | answer, aerr := ask(ctx, p, visible, probe.question) |
| 227 | if aerr != nil { |
| 228 | return nil, fmt.Errorf("probe %s: %w", probe, aerr) |
| 229 | } |
| 230 | r.Survived[probe.class], r.Answers[probe.class] = boolToInt(probe.score(answer)), answer |
| 231 | if control { |
| 232 | full, ferr := ask(ctx, p, h.sess.Snapshot(), probe.question) |
| 233 | if ferr != nil { |
| 234 | return nil, fmt.Errorf("control %s: %w", probe, ferr) |
| 235 | } |
| 236 | r.ControlOK[probe.class], r.ControlAnswers[probe.class] = boolToInt(probe.score(full)), full |
| 237 | } |
| 238 | } |
| 239 | out = append(out, r) |
| 240 | } |
| 241 | return out, nil |
| 242 | } |
| 243 | |
| 244 | // ask puts one probe question to the model on top of the given context. The |
| 245 | // budget has to clear the model's reasoning as well as its answer: a thinking |
| 246 | // model spends its first tokens reasoning, and a budget sized for the one-word |
| 247 | // answer alone comes back empty and scores as a fact compaction never lost. |
| 248 | func ask(ctx context.Context, p provider.Provider, msgs []provider.Message, question string) (string, error) { |
| 249 | answer, reasoning, err := askOnce(ctx, p, msgs, question, probeAnswerTokens) |
| 250 | if err != nil { |
| 251 | return "", err |
| 252 | } |
| 253 | if answer == "" || strings.Contains(answer, toolCallMarker) { |
| 254 | // One retry with room to think: a reply cut off mid-reasoning says |
| 255 | // nothing about whether the fold kept the fact. |
| 256 | answer, reasoning, err = askOnce(ctx, p, msgs, question, probeAnswerTokens*4) |
| 257 | if err != nil { |
| 258 | return "", err |
| 259 | } |
| 260 | } |
| 261 | switch { |
| 262 | case strings.Contains(answer, toolCallMarker): |
| 263 | return toolCallInvalid, nil |
| 264 | case answer == "": |
| 265 | return fmt.Sprintf("%s: %d reasoning chars>", noAnswerMarker, reasoning), nil |
| 266 | } |
| 267 | return answer, nil |
| 268 | } |
| 269 | |
| 270 | func askOnce(ctx context.Context, p provider.Provider, msgs []provider.Message, question string, budget int) (string, int, error) { |
| 271 | req := provider.Request{ |
| 272 | Messages: append(append([]provider.Message(nil), provider.ModelMessages(msgs)...), |
| 273 | provider.Message{Role: provider.RoleUser, Content: question + "\n\n" + probeAnswerContract}), |
| 274 | MaxTokens: budget, |
| 275 | } |
| 276 | ch, err := p.Stream(ctx, req) |
| 277 | if err != nil { |
| 278 | return "", 0, err |
| 279 | } |
| 280 | var answer, reasoning strings.Builder |
| 281 | for c := range ch { |
| 282 | switch c.Type { |
| 283 | case provider.ChunkText: |
| 284 | answer.WriteString(c.Text) |
| 285 | case provider.ChunkReasoning: |
| 286 | reasoning.WriteString(c.Text) |
| 287 | case provider.ChunkError: |
| 288 | return "", reasoning.Len(), c.Err |
| 289 | } |
| 290 | } |
| 291 | return strings.TrimSpace(answer.String()), reasoning.Len(), nil |
| 292 | } |
| 293 | |
| 294 | func printReport(mode string, res []genResult, at map[int]bool) { |
| 295 | fmt.Printf("\n## CompactionBench (%s)\n\n", mode) |
| 296 | fmt.Println("| gen | canonical tok | fold calls | fold input tok | largest call | projection tok | s | result |") |
| 297 | fmt.Println("| ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- |") |
| 298 | for _, r := range res { |
| 299 | status := r.Mode |
| 300 | if r.Error != "" { |
| 301 | status = "ERROR: " + firstLine(r.Error) |
| 302 | } |
| 303 | fmt.Printf("| %d | %d | %d | %d | %d | %d | %.1f | %s |\n", |
| 304 | r.Gen+1, r.CanonicalTokens, r.SummarizerCalls, r.SummarizerInput, r.LargestCall, r.ProjectionTokens, r.Seconds, status) |
| 305 | } |
| 306 | if !strings.HasPrefix(mode, "fidelity") { |
| 307 | return |
| 308 | } |
| 309 | classes := probeSuite() |
| 310 | fmt.Printf("\n### Probe survival (compacted / full-history control)\n\n| probe | %s |\n", joinGens(res, at)) |
| 311 | fmt.Printf("| --- | %s |\n", strings.Repeat(" ---: |", countGens(res, at))) |
| 312 | for _, p := range classes { |
| 313 | row := []string{} |
| 314 | for _, r := range res { |
| 315 | if !at[r.Gen+1] { |
| 316 | continue |
| 317 | } |
| 318 | if _, asked := r.Survived[p.class]; !asked { |
| 319 | row = append(row, "–") |
| 320 | continue |
| 321 | } |
| 322 | row = append(row, fmt.Sprintf("%s/%s", mark(r.Survived[p.class], r.Answers[p.class]), mark(r.ControlOK[p.class], r.ControlAnswers[p.class]))) |
| 323 | } |
| 324 | fmt.Printf("| %s | %s |\n", p.class, strings.Join(row, " | ")) |
| 325 | } |
| 326 | printMeasurementQuality(res) |
| 327 | } |
| 328 | |
| 329 | // printMeasurementQuality reports how many probes never got an answer at all. |
| 330 | // A survival rate quoted without it would read harness noise as fact loss. |
| 331 | func printMeasurementQuality(res []genResult) { |
| 332 | asked, bad, badControl := 0, 0, 0 |
| 333 | for _, r := range res { |
| 334 | for _, a := range r.Answers { |
| 335 | asked++ |
| 336 | if invalidAnswer(a) { |
| 337 | bad++ |
| 338 | } |
| 339 | } |
| 340 | for _, a := range r.ControlAnswers { |
| 341 | if invalidAnswer(a) { |
| 342 | badControl++ |
| 343 | } |
| 344 | } |
| 345 | } |
| 346 | fmt.Printf("\nUnanswered probes (excluded from the rates above): %d of %d compacted, %d of %d control.\n", bad, asked, badControl, asked) |
| 347 | if bad > 0 || badControl > 0 { |
| 348 | fmt.Println("A run with unanswered probes measures the harness as much as the compactor; see answers in the JSON report.") |
| 349 | } |
| 350 | } |
| 351 | |
| 352 | func mark(v int, answer string) string { |
| 353 | switch { |
| 354 | case invalidAnswer(answer): |
| 355 | return "n/a" |
| 356 | case v == 1: |
| 357 | return "ok" |
| 358 | } |
| 359 | return "LOST" |
| 360 | } |
| 361 | |
| 362 | func joinGens(res []genResult, at map[int]bool) string { |
| 363 | var s []string |
| 364 | for _, r := range res { |
| 365 | if at[r.Gen+1] { |
| 366 | s = append(s, fmt.Sprintf("gen %d", r.Gen+1)) |
| 367 | } |
| 368 | } |
| 369 | return strings.Join(s, " | ") |
| 370 | } |
| 371 | |
| 372 | func countGens(res []genResult, at map[int]bool) int { |
| 373 | n := 0 |
| 374 | for _, r := range res { |
| 375 | if at[r.Gen+1] { |
| 376 | n++ |
| 377 | } |
| 378 | } |
| 379 | return n |
| 380 | } |
| 381 | |
| 382 | func reportAt(spec string) map[int]bool { |
| 383 | at := map[int]bool{} |
| 384 | for part := range strings.SplitSeq(spec, ",") { |
| 385 | var n int |
| 386 | if _, err := fmt.Sscanf(strings.TrimSpace(part), "%d", &n); err == nil { |
| 387 | at[n] = true |
| 388 | } |
| 389 | } |
| 390 | return at |
| 391 | } |
| 392 | |
| 393 | // estimateTokens mirrors the kernel's own estimator so bench numbers and |
| 394 | // compaction telemetry are read in the same unit. |
| 395 | func estimateTokens(s string) int { |
| 396 | if s == "" { |
| 397 | return 0 |
| 398 | } |
| 399 | if runes := utf8.RuneCountInString(s); runes > (len(s)+3)/4 { |
| 400 | return runes |
| 401 | } |
| 402 | return (len(s) + 3) / 4 |
| 403 | } |
| 404 | |
| 405 | func renderAll(msgs []provider.Message) string { |
| 406 | var b strings.Builder |
| 407 | for _, m := range msgs { |
| 408 | b.WriteString(m.Content) |
| 409 | for _, tc := range m.ToolCalls { |
| 410 | b.WriteString(tc.Name) |
| 411 | b.WriteString(tc.Arguments) |
| 412 | } |
| 413 | b.WriteByte('\n') |
| 414 | } |
| 415 | return b.String() |
| 416 | } |
| 417 | |
| 418 | func firstLine(s string) string { |
| 419 | first, _, _ := strings.Cut(s, "\n") |
| 420 | return first |
| 421 | } |
| 422 | |
| 423 | func boolToInt(b bool) int { |
| 424 | if b { |
| 425 | return 1 |
| 426 | } |
| 427 | return 0 |
| 428 | } |
| 429 | |
| 430 | type testingDir struct{ dir string } |
| 431 | |
| 432 | func tempDir() (*testingDir, func(), error) { |
| 433 | dir, err := os.MkdirTemp("", "compactionbench-") |
| 434 | if err != nil { |
| 435 | return nil, nil, err |
| 436 | } |
| 437 | return &testingDir{dir: dir}, func() { _ = os.RemoveAll(dir) }, nil |
| 438 | } |
| 439 | |
| 440 | // callRecorder captures every summarizer request the fold issued, which is the |
| 441 | // measurement the cost arm exists for: how many calls, and how large the |
| 442 | // largest one got. |
| 443 | type callRecorder struct{ calls []recordedCall } |
| 444 | |
| 445 | type recordedCall struct { |
| 446 | tokens int |
| 447 | system string |
| 448 | } |
| 449 | |
| 450 | func (r *callRecorder) reset() { r.calls = nil } |
| 451 | |
| 452 | func (r *callRecorder) note(req provider.Request) { |
| 453 | c := recordedCall{} |
| 454 | for _, m := range req.Messages { |
| 455 | c.tokens += estimateTokens(m.Content) |
| 456 | if m.Role == provider.RoleSystem { |
| 457 | c.system = m.Content |
| 458 | } |
| 459 | } |
| 460 | r.calls = append(r.calls, c) |
| 461 | } |
| 462 | |
| 463 | func (r *callRecorder) sink() event.Sink { return event.Discard } |
| 464 | |
| 465 | const syntheticDigest = `## Standing facts & constraints |
| 466 | - never modify config/schema.sql |
| 467 | ## Goal |
| 468 | Fix the config round-trip formatting bug. |
| 469 | ## Pending & next step |
| 470 | Re-run TestRoundTrip after the latest edit.` |
| 471 | |
| 472 | // scriptedProvider answers every summarizer call with a fixed digest so the |
| 473 | // cost arm is deterministic and needs no API key. It refuses an input larger |
| 474 | // than the window the way a real provider does, so the bench observes the |
| 475 | // wedge — a fold that can no longer be summarized at all — instead of |
| 476 | // inferring it from the input size. |
| 477 | type scriptedProvider struct { |
| 478 | rec *callRecorder |
| 479 | reply string |
| 480 | window int |
| 481 | } |
| 482 | |
| 483 | func (p *scriptedProvider) Name() string { return "scripted" } |
| 484 | |
| 485 | func (p *scriptedProvider) Stream(_ context.Context, req provider.Request) (<-chan provider.Chunk, error) { |
| 486 | p.rec.note(req) |
| 487 | ch := make(chan provider.Chunk, 2) |
| 488 | if in := p.rec.calls[len(p.rec.calls)-1].tokens; p.window > 0 && in > p.window { |
| 489 | ch <- provider.Chunk{Type: provider.ChunkError, Err: fmt.Errorf("this model's maximum context length is %d tokens, however you requested %d tokens", p.window, in)} |
| 490 | close(ch) |
| 491 | return ch, nil |
| 492 | } |
| 493 | ch <- provider.Chunk{Type: provider.ChunkText, Text: p.reply} |
| 494 | ch <- provider.Chunk{Type: provider.ChunkDone} |
| 495 | close(ch) |
| 496 | return ch, nil |
| 497 | } |
| 498 | |
| 499 | // recordingProvider measures the same thing against a real provider. |
| 500 | type recordingProvider struct { |
| 501 | inner provider.Provider |
| 502 | rec *callRecorder |
| 503 | } |
| 504 | |
| 505 | func (p *recordingProvider) Name() string { return p.inner.Name() } |
| 506 | |
| 507 | func (p *recordingProvider) Stream(ctx context.Context, req provider.Request) (<-chan provider.Chunk, error) { |
| 508 | p.rec.note(req) |
| 509 | return p.inner.Stream(ctx, req) |
| 510 | } |
| 511 |