| 1 | // Drives the context-maintenance E2E scenarios against the real DeepSeek API: |
| 2 | // seed → (idle past cache TTL) → resume A/B-compares cold-restart miss tokens with and without pruning. |
| 3 | package main |
| 4 | |
| 5 | import ( |
| 6 | "context" |
| 7 | "encoding/json" |
| 8 | "flag" |
| 9 | "fmt" |
| 10 | "os" |
| 11 | "path/filepath" |
| 12 | "strings" |
| 13 | "time" |
| 14 | |
| 15 | "reasonix/internal/agent" |
| 16 | "reasonix/internal/event" |
| 17 | "reasonix/internal/provider" |
| 18 | _ "reasonix/internal/provider/openai" |
| 19 | "reasonix/internal/tool" |
| 20 | "reasonix/internal/tool/builtin" |
| 21 | ) |
| 22 | |
| 23 | const ( |
| 24 | model = "deepseek-v4-flash" |
| 25 | baseURL = "https://api.deepseek.com" |
| 26 | fatResults = 20 |
| 27 | fatBytes = 12_000 |
| 28 | ) |
| 29 | |
| 30 | func prov() provider.Provider { |
| 31 | key := os.Getenv("DEEPSEEK_API_KEY") |
| 32 | if key == "" { |
| 33 | fmt.Fprintln(os.Stderr, "DEEPSEEK_API_KEY not set") |
| 34 | os.Exit(1) |
| 35 | } |
| 36 | p, err := provider.New("openai", provider.Config{Name: "e2e", BaseURL: baseURL, Model: model, APIKey: key}) |
| 37 | if err != nil { |
| 38 | fmt.Fprintln(os.Stderr, err) |
| 39 | os.Exit(1) |
| 40 | } |
| 41 | return p |
| 42 | } |
| 43 | |
| 44 | func fakeGoFile(nonce string, i, size int) string { |
| 45 | var b strings.Builder |
| 46 | fmt.Fprintf(&b, "// module %s file%02d\npackage stress\n\n", nonce, i) |
| 47 | line := 0 |
| 48 | for b.Len() < size { |
| 49 | fmt.Fprintf(&b, "func helper_%s_%02d_%04d(x int) int { return x*%d + %d }\n", nonce, i, line, line+3, line*7) |
| 50 | line++ |
| 51 | } |
| 52 | return b.String() |
| 53 | } |
| 54 | |
| 55 | func seedSession(nonce string) *agent.Session { |
| 56 | s := agent.NewSession("You are a terse coding agent reviewing a Go codebase.") |
| 57 | s.Add(provider.Message{Role: provider.RoleUser, Content: "Review every file in module " + nonce + " one by one. Keep notes short."}) |
| 58 | for i := 0; i < fatResults; i++ { |
| 59 | id := fmt.Sprintf("c%02d", i) |
| 60 | name := fmt.Sprintf("src/file%02d.go", i) |
| 61 | s.Add(provider.Message{Role: provider.RoleAssistant, ToolCalls: []provider.ToolCall{{ID: id, Name: "read_file", Arguments: fmt.Sprintf(`{"path":%q}`, name)}}}) |
| 62 | s.Add(provider.Message{Role: provider.RoleTool, ToolCallID: id, Name: "read_file", Content: fakeGoFile(nonce, i, fatBytes)}) |
| 63 | s.Add(provider.Message{Role: provider.RoleAssistant, Content: fmt.Sprintf("Reviewed %s.", name)}) |
| 64 | } |
| 65 | return s |
| 66 | } |
| 67 | |
| 68 | // oneShot appends a user message and runs a single completion, returning usage. |
| 69 | func oneShot(p provider.Provider, msgs []provider.Message, question string) (provider.Usage, error) { |
| 70 | req := provider.Request{ |
| 71 | Messages: append(append([]provider.Message(nil), msgs...), provider.Message{Role: provider.RoleUser, Content: question}), |
| 72 | MaxTokens: 32, |
| 73 | } |
| 74 | ch, err := p.Stream(context.Background(), req) |
| 75 | if err != nil { |
| 76 | return provider.Usage{}, err |
| 77 | } |
| 78 | var u provider.Usage |
| 79 | for c := range ch { |
| 80 | switch c.Type { |
| 81 | case provider.ChunkUsage: |
| 82 | u = *c.Usage |
| 83 | case provider.ChunkError: |
| 84 | return u, c.Err |
| 85 | } |
| 86 | } |
| 87 | return u, nil |
| 88 | } |
| 89 | |
| 90 | type meta struct { |
| 91 | SeededAt time.Time `json:"seeded_at"` |
| 92 | Nonces map[string]string `json:"nonces"` |
| 93 | SeedUse map[string]provider.Usage `json:"seed_usage"` |
| 94 | } |
| 95 | |
| 96 | func seed(dir string) { |
| 97 | p := prov() |
| 98 | if err := os.MkdirAll(dir, 0o755); err != nil { |
| 99 | fmt.Fprintln(os.Stderr, err) |
| 100 | os.Exit(1) |
| 101 | } |
| 102 | m := meta{SeededAt: time.Now(), Nonces: map[string]string{}, SeedUse: map[string]provider.Usage{}} |
| 103 | for _, arm := range []string{"pruned", "control"} { |
| 104 | nonce := fmt.Sprintf("%s%d", arm, time.Now().UnixNano()%1_000_000) |
| 105 | s := seedSession(nonce) |
| 106 | u, err := oneShot(p, s.Snapshot(), "How many files have you reviewed so far? Reply with just the number.") |
| 107 | if err != nil { |
| 108 | fmt.Fprintf(os.Stderr, "seed %s: %v\n", arm, err) |
| 109 | os.Exit(1) |
| 110 | } |
| 111 | if err := s.Save(filepath.Join(dir, arm+".jsonl")); err != nil { |
| 112 | fmt.Fprintln(os.Stderr, err) |
| 113 | os.Exit(1) |
| 114 | } |
| 115 | m.Nonces[arm] = nonce |
| 116 | m.SeedUse[arm] = u |
| 117 | fmt.Printf("seeded %-7s prompt=%d hit=%d miss=%d\n", arm, u.PromptTokens, u.CacheHitTokens, u.CacheMissTokens) |
| 118 | } |
| 119 | b, _ := json.MarshalIndent(m, "", " ") |
| 120 | if err := os.WriteFile(filepath.Join(dir, "meta.json"), b, 0o644); err != nil { |
| 121 | fmt.Fprintln(os.Stderr, err) |
| 122 | os.Exit(1) |
| 123 | } |
| 124 | } |
| 125 | |
| 126 | func resume(dir string) { |
| 127 | p := prov() |
| 128 | raw, err := os.ReadFile(filepath.Join(dir, "meta.json")) |
| 129 | if err != nil { |
| 130 | fmt.Fprintln(os.Stderr, err) |
| 131 | os.Exit(1) |
| 132 | } |
| 133 | var m meta |
| 134 | if err := json.Unmarshal(raw, &m); err != nil { |
| 135 | fmt.Fprintln(os.Stderr, err) |
| 136 | os.Exit(1) |
| 137 | } |
| 138 | idle := time.Since(m.SeededAt).Round(time.Minute) |
| 139 | |
| 140 | out := map[string]provider.Usage{} |
| 141 | for _, arm := range []string{"pruned", "control"} { |
| 142 | s, err := agent.LoadSession(filepath.Join(dir, arm+".jsonl")) |
| 143 | if err != nil { |
| 144 | fmt.Fprintln(os.Stderr, err) |
| 145 | os.Exit(1) |
| 146 | } |
| 147 | prunes := 0 |
| 148 | if arm == "pruned" { |
| 149 | a := agent.New(nil, tool.NewRegistry(), s, agent.Options{ContextWindow: 1_000_000, ArchiveDir: filepath.Join(dir, "archive")}, event.Discard) |
| 150 | st, err := a.PruneStaleToolResults() |
| 151 | if err != nil { |
| 152 | fmt.Fprintln(os.Stderr, "prune:", err) |
| 153 | os.Exit(1) |
| 154 | } |
| 155 | prunes = st.Results |
| 156 | } |
| 157 | u, err := oneShot(p, s.Snapshot(), "Which file did you review first? Reply with just the path.") |
| 158 | if err != nil { |
| 159 | fmt.Fprintf(os.Stderr, "resume %s: %v\n", arm, err) |
| 160 | os.Exit(1) |
| 161 | } |
| 162 | out[arm] = u |
| 163 | fmt.Printf("resume %-7s idle=%s pruned=%d prompt=%d hit=%d miss=%d\n", arm, idle, prunes, u.PromptTokens, u.CacheHitTokens, u.CacheMissTokens) |
| 164 | } |
| 165 | b, _ := json.MarshalIndent(map[string]any{"idle": idle.String(), "resume_usage": out}, "", " ") |
| 166 | if err := os.WriteFile(filepath.Join(dir, fmt.Sprintf("resume-%d.json", time.Now().Unix())), b, 0o644); err != nil { |
| 167 | fmt.Fprintln(os.Stderr, err) |
| 168 | os.Exit(1) |
| 169 | } |
| 170 | |
| 171 | c, pr := out["control"], out["pruned"] |
| 172 | if c.CacheMissTokens > 0 { |
| 173 | fmt.Printf("\ncold-restart miss tokens: control=%d pruned=%d (%.0f%% reduction)\n", |
| 174 | c.CacheMissTokens, pr.CacheMissTokens, 100*(1-float64(pr.CacheMissTokens)/float64(c.CacheMissTokens))) |
| 175 | } |
| 176 | } |
| 177 | |
| 178 | // comprehension checks that the agent re-reads a file behind a prune placeholder |
| 179 | // instead of hallucinating: the answer is a number that exists only in the file. |
| 180 | func comprehension(trials int) { |
| 181 | p := prov() |
| 182 | pass := 0 |
| 183 | for t := 0; t < trials; t++ { |
| 184 | dir, err := os.MkdirTemp("", "cm-e2e-") |
| 185 | if err != nil { |
| 186 | fmt.Fprintln(os.Stderr, err) |
| 187 | os.Exit(1) |
| 188 | } |
| 189 | secret := fmt.Sprintf("%d", 1000+time.Now().UnixNano()%9000) |
| 190 | content := "package cfg\n\n// retention floor, milliseconds\nconst cacheRetentionFloor = " + secret + "\n" + strings.Repeat("// padding line filler for prune eligibility\n", 400) |
| 191 | if err := os.WriteFile(filepath.Join(dir, "config.go"), []byte(content), 0o644); err != nil { |
| 192 | fmt.Fprintln(os.Stderr, err) |
| 193 | os.Exit(1) |
| 194 | } |
| 195 | |
| 196 | s := agent.NewSession("You are a terse coding agent. Use tools when you need file contents.") |
| 197 | s.Add(provider.Message{Role: provider.RoleUser, Content: "Read config.go and note its constants."}) |
| 198 | s.Add(provider.Message{Role: provider.RoleAssistant, ToolCalls: []provider.ToolCall{{ID: "r1", Name: "read_file", Arguments: `{"path":"config.go"}`}}}) |
| 199 | s.Add(provider.Message{Role: provider.RoleTool, ToolCallID: "r1", Name: "read_file", Content: content}) |
| 200 | s.Add(provider.Message{Role: provider.RoleAssistant, Content: "Noted the constants in config.go."}) |
| 201 | for i := 0; i < 4; i++ { |
| 202 | s.Add(provider.Message{Role: provider.RoleUser, Content: fmt.Sprintf("ack %d", i)}) |
| 203 | s.Add(provider.Message{Role: provider.RoleAssistant, Content: "ok"}) |
| 204 | } |
| 205 | |
| 206 | reg := tool.NewRegistry() |
| 207 | for _, tl := range (builtin.Workspace{Dir: dir}).Tools("read_file") { |
| 208 | reg.Add(tl) |
| 209 | } |
| 210 | a := agent.New(p, reg, s, agent.Options{ContextWindow: 2000, RecentKeep: 2, MaxSteps: 5}, event.Discard) |
| 211 | st, err := a.PruneStaleToolResults() |
| 212 | if err != nil || st.Results == 0 { |
| 213 | fmt.Fprintf(os.Stderr, "trial %d: prune did not fire (st=%+v err=%v)\n", t, st, err) |
| 214 | os.Exit(1) |
| 215 | } |
| 216 | |
| 217 | err = a.Run(context.Background(), "What is the exact numeric value of cacheRetentionFloor in config.go? Reply with just the number.") |
| 218 | reRead, answered := false, false |
| 219 | for _, msg := range s.Snapshot() { |
| 220 | if msg.Role == provider.RoleTool && strings.Contains(msg.Content, "cacheRetentionFloor = "+secret) { |
| 221 | reRead = true |
| 222 | } |
| 223 | if msg.Role == provider.RoleAssistant && strings.Contains(msg.Content, secret) { |
| 224 | answered = true |
| 225 | } |
| 226 | } |
| 227 | ok := err == nil && reRead && answered |
| 228 | if ok { |
| 229 | pass++ |
| 230 | } |
| 231 | fmt.Printf("trial %d: re-read=%v answered=%v err=%v\n", t, reRead, answered, err) |
| 232 | os.RemoveAll(dir) |
| 233 | } |
| 234 | fmt.Printf("\ncomprehension: %d/%d passed\n", pass, trials) |
| 235 | if pass < trials { |
| 236 | os.Exit(1) |
| 237 | } |
| 238 | } |
| 239 | |
| 240 | func main() { |
| 241 | dir := flag.String("dir", "benchmarks/context-maintenance-e2e/run", "state directory for seed/resume") |
| 242 | trials := flag.Int("trials", 5, "comprehension trials") |
| 243 | flag.Parse() |
| 244 | switch flag.Arg(0) { |
| 245 | case "seed": |
| 246 | seed(*dir) |
| 247 | case "resume": |
| 248 | resume(*dir) |
| 249 | case "comprehension": |
| 250 | comprehension(*trials) |
| 251 | default: |
| 252 | fmt.Fprintln(os.Stderr, "usage: context-maintenance-e2e [seed|resume|comprehension]") |
| 253 | os.Exit(1) |
| 254 | } |
| 255 | } |
| 256 |