| 1 | package agent |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "errors" |
| 7 | "strings" |
| 8 | "testing" |
| 9 | "time" |
| 10 | |
| 11 | "reasonix/internal/event" |
| 12 | "reasonix/internal/provider" |
| 13 | "reasonix/internal/tool" |
| 14 | ) |
| 15 | |
| 16 | type sharedWindowTestProvider struct { |
| 17 | budget int |
| 18 | shared bool |
| 19 | policy provider.SharedWindowInputPolicy |
| 20 | last provider.Request |
| 21 | calls int |
| 22 | finish string |
| 23 | } |
| 24 | |
| 25 | type outputLimitRetryProvider struct { |
| 26 | calls []provider.Request |
| 27 | } |
| 28 | |
| 29 | type namedOutputBudgetProvider struct{ name string } |
| 30 | |
| 31 | func (p *namedOutputBudgetProvider) Name() string { return p.name } |
| 32 | |
| 33 | func (*namedOutputBudgetProvider) Stream(context.Context, provider.Request) (<-chan provider.Chunk, error) { |
| 34 | return nil, errors.New("unused") |
| 35 | } |
| 36 | |
| 37 | func (*outputLimitRetryProvider) Name() string { return "output-limit-retry" } |
| 38 | |
| 39 | func (p *outputLimitRetryProvider) Stream(_ context.Context, req provider.Request) (<-chan provider.Chunk, error) { |
| 40 | p.calls = append(p.calls, req) |
| 41 | if len(p.calls) == 1 { |
| 42 | return nil, &provider.OutputLimitError{ |
| 43 | APIError: &provider.APIError{Status: 400, Body: "max_tokens is too large"}, |
| 44 | RequestedTokens: req.MaxTokens, |
| 45 | MaxOutputTokens: 131_072, |
| 46 | } |
| 47 | } |
| 48 | ch := make(chan provider.Chunk, 2) |
| 49 | ch <- provider.Chunk{Type: provider.ChunkText, Text: "ok"} |
| 50 | ch <- provider.Chunk{Type: provider.ChunkDone} |
| 51 | close(ch) |
| 52 | return ch, nil |
| 53 | } |
| 54 | |
| 55 | func TestStreamProviderRequestRetriesOutputLimitBeforeAnyOutput(t *testing.T) { |
| 56 | prov := &outputLimitRetryProvider{} |
| 57 | a := &Agent{svc: agentServices{prov: prov}, sess: sessionRuntime{}} |
| 58 | ch, err := a.streamProviderRequest(context.Background(), provider.Request{MaxTokens: 384_000}) |
| 59 | if err != nil { |
| 60 | t.Fatalf("streamProviderRequest: %v", err) |
| 61 | } |
| 62 | var text strings.Builder |
| 63 | for chunk := range ch { |
| 64 | if chunk.Type == provider.ChunkText { |
| 65 | text.WriteString(chunk.Text) |
| 66 | } |
| 67 | if chunk.Type == provider.ChunkError { |
| 68 | t.Fatalf("retry stream emitted error: %v", chunk.Err) |
| 69 | } |
| 70 | } |
| 71 | if text.String() != "ok" || len(prov.calls) != 2 || prov.calls[1].MaxTokens != 131_072 { |
| 72 | t.Fatalf("retry calls = %+v, text=%q", prov.calls, text.String()) |
| 73 | } |
| 74 | if got := a.learnedCompletionBudget(); got != 131_072 { |
| 75 | t.Fatalf("learned completion budget = %d, want 131072", got) |
| 76 | } |
| 77 | } |
| 78 | |
| 79 | func TestLearnedOutputBudgetCacheIsScopedToProviderRouteAndModel(t *testing.T) { |
| 80 | providerName := "output-budget-cache-provider-unique" |
| 81 | modelRef := "opencode-go/deepseek-v4-flash-cache-unique" |
| 82 | first := &Agent{ |
| 83 | agentConfig: agentConfig{modelRef: modelRef}, |
| 84 | svc: agentServices{prov: &namedOutputBudgetProvider{name: providerName}}, |
| 85 | sess: sessionRuntime{}, |
| 86 | } |
| 87 | first.learnOutputBudget(131_072) |
| 88 | |
| 89 | sameModel := &Agent{ |
| 90 | agentConfig: agentConfig{modelRef: modelRef}, |
| 91 | svc: agentServices{prov: &namedOutputBudgetProvider{name: providerName}}, |
| 92 | sess: sessionRuntime{}, |
| 93 | } |
| 94 | if got := sameModel.learnedCompletionBudget(); got != 131_072 { |
| 95 | t.Fatalf("same provider/route/model budget = %d, want 131072", got) |
| 96 | } |
| 97 | |
| 98 | differentRoute := &Agent{ |
| 99 | agentConfig: agentConfig{modelRef: modelRef}, |
| 100 | svc: agentServices{prov: &namedOutputBudgetProvider{name: providerName + "-responses"}}, |
| 101 | sess: sessionRuntime{}, |
| 102 | } |
| 103 | if got := differentRoute.learnedCompletionBudget(); got != 0 { |
| 104 | t.Fatalf("different route inherited budget %d", got) |
| 105 | } |
| 106 | |
| 107 | differentModel := &Agent{ |
| 108 | agentConfig: agentConfig{modelRef: modelRef + "-other"}, |
| 109 | svc: agentServices{prov: &namedOutputBudgetProvider{name: providerName}}, |
| 110 | sess: sessionRuntime{}, |
| 111 | } |
| 112 | if got := differentModel.learnedCompletionBudget(); got != 0 { |
| 113 | t.Fatalf("different model inherited budget %d", got) |
| 114 | } |
| 115 | |
| 116 | key := outputBudgetCacheKey(first) |
| 117 | learnedOutputBudgetCache.Lock() |
| 118 | entry := learnedOutputBudgetCache.entries[key] |
| 119 | entry.expiresAt = time.Now().Add(-time.Minute) |
| 120 | learnedOutputBudgetCache.entries[key] = entry |
| 121 | learnedOutputBudgetCache.Unlock() |
| 122 | if got := (&Agent{ |
| 123 | agentConfig: agentConfig{modelRef: modelRef}, |
| 124 | svc: agentServices{prov: &namedOutputBudgetProvider{name: providerName}}, |
| 125 | sess: sessionRuntime{}, |
| 126 | }).learnedCompletionBudget(); got != 0 { |
| 127 | t.Fatalf("expired budget = %d, want 0", got) |
| 128 | } |
| 129 | } |
| 130 | |
| 131 | func (*sharedWindowTestProvider) Name() string { return "shared-window-test" } |
| 132 | |
| 133 | func (p *sharedWindowTestProvider) Stream(_ context.Context, req provider.Request) (<-chan provider.Chunk, error) { |
| 134 | p.last = req |
| 135 | p.calls++ |
| 136 | ch := make(chan provider.Chunk, 3) |
| 137 | ch <- provider.Chunk{Type: provider.ChunkText, Text: "summary"} |
| 138 | if p.finish != "" { |
| 139 | ch <- provider.Chunk{Type: provider.ChunkUsage, Usage: &provider.Usage{FinishReason: p.finish}} |
| 140 | } |
| 141 | ch <- provider.Chunk{Type: provider.ChunkDone} |
| 142 | close(ch) |
| 143 | return ch, nil |
| 144 | } |
| 145 | |
| 146 | // Before any usage calibrates the session, the estimate compared against the |
| 147 | // context window must still be tokens. It used to be characters, which reads |
| 148 | // 3-4x high and compacted long before the configured ratio. |
| 149 | func TestEstimatedPromptTokensStayInTokenUnitBeforeCalibration(t *testing.T) { |
| 150 | a := &Agent{agentConfig: agentConfig{contextWindow: 1_000_000}} |
| 151 | cases := []struct { |
| 152 | name string |
| 153 | text string |
| 154 | realish, upper int |
| 155 | }{ |
| 156 | // DeepSeek bills Chinese near 0.6 tokens per han rune, English near 0.25 |
| 157 | // per character. The cold estimate may be conservative, never 3x. |
| 158 | {"chinese", strings.Repeat("上下文压缩策略", 8_000), 33_600, 50_000}, |
| 159 | {"english", strings.Repeat("compact the context window ", 8_000), 54_000, 70_000}, |
| 160 | } |
| 161 | for _, tc := range cases { |
| 162 | t.Run(tc.name, func(t *testing.T) { |
| 163 | got := a.estimatedPromptTokens([]provider.Message{{Role: provider.RoleUser, Content: tc.text}}) |
| 164 | if got < tc.realish/2 || got > tc.upper { |
| 165 | t.Fatalf("cold estimate = %d tokens, want between %d and %d (real ~%d)", |
| 166 | got, tc.realish/2, tc.upper, tc.realish) |
| 167 | } |
| 168 | }) |
| 169 | } |
| 170 | } |
| 171 | |
| 172 | // Desktop rebinds sessions constantly — tab switches, forks, and the snapshot |
| 173 | // conflict path that adopts the newer disk transcript. Each rebind used to drop |
| 174 | // the calibration and put the next turn back on the cold estimate. |
| 175 | func TestSessionSwapKeepsPromptCalibration(t *testing.T) { |
| 176 | a := &Agent{agentConfig: agentConfig{contextWindow: 200_000}} |
| 177 | msgs := []provider.Message{{Role: provider.RoleUser, Content: strings.Repeat("字", 60_000)}} |
| 178 | a.setPromptTokenCalibration(36_000, requestCalibrationShapeOf(provider.Request{Messages: msgs})) |
| 179 | |
| 180 | before := a.estimatedPromptTokens(msgs) |
| 181 | a.SetSession(NewSession("system")) |
| 182 | after := a.estimatedPromptTokens(msgs) |
| 183 | |
| 184 | if before != after { |
| 185 | t.Fatalf("estimate moved across a session swap: %d -> %d", before, after) |
| 186 | } |
| 187 | if after > 40_000 { |
| 188 | t.Fatalf("estimate = %d, want the calibrated ~36000 rather than a cold fallback", after) |
| 189 | } |
| 190 | } |
| 191 | |
| 192 | func TestSharedWindowFoldDoesNotPrivatelyShortenOversizedInput(t *testing.T) { |
| 193 | prov := &sharedWindowTestProvider{budget: 128 * 1024, shared: true} |
| 194 | a := &Agent{agentConfig: agentConfig{contextWindow: 100_000}, svc: agentServices{prov: prov, sink: event.Discard}, sess: sessionRuntime{output: outputBudgetState{outputBudget: prov.budget}}} |
| 195 | // Manual summary input is not privately shortened. An unfittable request is |
| 196 | // rejected before the provider call. |
| 197 | toolBody := strings.Repeat("file line content here. ", 20_000) // ~480K chars |
| 198 | fold := []provider.Message{ |
| 199 | {Role: provider.RoleUser, Content: "read large files"}, |
| 200 | {Role: provider.RoleAssistant, ToolCalls: []provider.ToolCall{{ID: "1", Name: "read_file", Arguments: "{}"}}}, |
| 201 | {Role: provider.RoleTool, ToolCallID: "1", Name: "read_file", Content: toolBody}, |
| 202 | {Role: provider.RoleAssistant, ToolCalls: []provider.ToolCall{{ID: "2", Name: "read_file", Arguments: "{}"}}}, |
| 203 | {Role: provider.RoleTool, ToolCallID: "2", Name: "read_file", Content: toolBody}, |
| 204 | } |
| 205 | |
| 206 | if _, err := a.foldToSummary(context.Background(), fold, ""); !errors.Is(err, ErrCompactionRequired) { |
| 207 | t.Fatalf("foldToSummary = %v, want admission failure", err) |
| 208 | } |
| 209 | if prov.calls != 0 { |
| 210 | t.Fatalf("unfittable fold called provider %d times", prov.calls) |
| 211 | } |
| 212 | } |
| 213 | |
| 214 | // A single unshortenable fold that still exceeds the single-request budget after |
| 215 | // all deterministic shorteners must fail once (no multi-span split). |
| 216 | func TestSharedWindowFoldRejectsUnshortenableOverBudgetInput(t *testing.T) { |
| 217 | prov := &sharedWindowTestProvider{budget: 128 * 1024, shared: true} |
| 218 | a := &Agent{agentConfig: agentConfig{contextWindow: 100_000}, svc: agentServices{prov: prov, sink: event.Discard}, sess: sessionRuntime{output: outputBudgetState{outputBudget: prov.budget}}} |
| 219 | fold := []provider.Message{{Role: provider.RoleUser, Content: strings.Repeat("字", 200_000)}} |
| 220 | _, err := a.foldToSummary(context.Background(), fold, "") |
| 221 | if !errors.Is(err, ErrCompactionRequired) { |
| 222 | t.Fatalf("foldToSummary err = %v, want context admission failure", err) |
| 223 | } |
| 224 | if prov.calls != 0 { |
| 225 | t.Fatalf("over-budget unshortenable fold still called summarizer %d times", prov.calls) |
| 226 | } |
| 227 | } |
| 228 | |
| 229 | func (p *sharedWindowTestProvider) OutputBudget() int { return p.budget } |
| 230 | func (p *sharedWindowTestProvider) SharesContextWindow() bool { return p.shared } |
| 231 | |
| 232 | func (p *sharedWindowTestProvider) SharedWindowInputPolicy() provider.SharedWindowInputPolicy { |
| 233 | return p.policy |
| 234 | } |
| 235 | |
| 236 | func TestEffectiveOutputBudgetClipsSharedWindowRequest(t *testing.T) { |
| 237 | prov := &sharedWindowTestProvider{budget: 128 * 1024, shared: true} |
| 238 | a := &Agent{agentConfig: agentConfig{contextWindow: 1_048_576}, svc: agentServices{prov: prov}, sess: sessionRuntime{output: outputBudgetState{outputBudget: prov.budget}}} |
| 239 | msgs := []provider.Message{{Role: provider.RoleUser, Content: strings.Repeat("字", 950_000)}} |
| 240 | // Calibrate this session at one token per rune. The 950K prompt fits, but |
| 241 | // not beside the provider's full 128K output default. |
| 242 | a.sess.output.lastUsage.Store(&provider.Usage{PromptTokens: 950_000}) |
| 243 | a.setPromptTokenCalibration(950_000, requestCalibrationShapeOf(provider.Request{Messages: msgs})) |
| 244 | |
| 245 | got, clipped, err := a.effectiveOutputBudget(provider.Request{Messages: msgs}) |
| 246 | if err != nil { |
| 247 | t.Fatalf("effectiveOutputBudget: %v", err) |
| 248 | } |
| 249 | if !clipped { |
| 250 | t.Fatal("near-window request kept the provider's full output budget") |
| 251 | } |
| 252 | if got <= 0 || got >= prov.budget { |
| 253 | t.Fatalf("clipped budget = %d, want 0 < budget < %d", got, prov.budget) |
| 254 | } |
| 255 | if got+950_000 > a.contextWindow-outputBudgetReserve { |
| 256 | t.Fatalf("input + output = %d, exceeds reserved shared window %d", got+950_000, a.contextWindow-outputBudgetReserve) |
| 257 | } |
| 258 | } |
| 259 | |
| 260 | func TestCalibratedOutputBudgetIncludesReplayedReasoning(t *testing.T) { |
| 261 | prov := &sharedWindowTestProvider{budget: 128 * 1024, shared: true} |
| 262 | a := &Agent{agentConfig: agentConfig{contextWindow: 200_000}, svc: agentServices{prov: prov}, sess: sessionRuntime{output: outputBudgetState{outputBudget: prov.budget}}} |
| 263 | previous := []provider.Message{{Role: provider.RoleUser, Content: strings.Repeat("x", 300_000)}} |
| 264 | a.setPromptTokenCalibration(75_000, requestCalibrationShapeOf(provider.Request{Messages: previous})) |
| 265 | current := append(previous, provider.Message{ |
| 266 | Role: provider.RoleAssistant, |
| 267 | ReasoningContent: strings.Repeat("r", 400_000), |
| 268 | ToolCalls: []provider.ToolCall{{ID: "call_1", Name: "bash", Arguments: `{}`}}, |
| 269 | }) |
| 270 | |
| 271 | before := a.estimatedPromptTokens(previous) |
| 272 | after := a.estimatedPromptTokens(current) |
| 273 | if after < before+99_000 { |
| 274 | t.Fatalf("400K replayed reasoning was not calibrated: before=%d after=%d", before, after) |
| 275 | } |
| 276 | budget, clipped, err := a.effectiveOutputBudget(provider.Request{Messages: current}) |
| 277 | if err != nil { |
| 278 | t.Fatalf("effectiveOutputBudget: %v", err) |
| 279 | } |
| 280 | adm := a.lastAdmission() |
| 281 | if !clipped || budget <= 0 || budget >= prov.budget || |
| 282 | budget+adm.PromptTokens+adm.ReserveTokens > a.contextWindow { |
| 283 | t.Fatalf("replayed reasoning budget = %d clipped=%v admission=%+v, want a clipped request within the shared window", budget, clipped, adm) |
| 284 | } |
| 285 | } |
| 286 | |
| 287 | func TestCalibratedOutputBudgetKeepsCJKConservativeFloor(t *testing.T) { |
| 288 | prov := &sharedWindowTestProvider{budget: 128 * 1024, shared: true} |
| 289 | a := &Agent{agentConfig: agentConfig{contextWindow: 1_048_576}, svc: agentServices{prov: prov}, sess: sessionRuntime{output: outputBudgetState{outputBudget: prov.budget}}} |
| 290 | previous := []provider.Message{{Role: provider.RoleUser, Content: strings.Repeat("x", 300_000)}} |
| 291 | a.setPromptTokenCalibration(75_000, requestCalibrationShapeOf(provider.Request{Messages: previous})) |
| 292 | // Enough unrepresented CJK that the reply no longer fits beside it: at the |
| 293 | // corrected unit 430K runes leave most of a 1M window free. |
| 294 | current := append(append([]provider.Message(nil), previous...), provider.Message{ |
| 295 | Role: provider.RoleAssistant, |
| 296 | ReasoningContent: strings.Repeat("字", 1_200_000), |
| 297 | ToolCalls: []provider.ToolCall{{ID: "call_1", Name: "bash", Arguments: `{}`}}, |
| 298 | }) |
| 299 | |
| 300 | // The unrepresented CJK runes are priced at the cold rate: 3 bytes each at |
| 301 | // ~4 chars per token, i.e. 0.75 tokens per rune against a real ~0.6. |
| 302 | calibrated := a.estimatedPromptTokens(current) |
| 303 | wantFloor := 75_000 + 1_200_000*3/4 |
| 304 | if calibrated < wantFloor { |
| 305 | t.Fatalf("calibrated estimate %d fell below mixed-script safety floor %d", calibrated, wantFloor) |
| 306 | } |
| 307 | |
| 308 | budget, clipped, err := a.effectiveOutputBudget(provider.Request{Messages: current}) |
| 309 | if err != nil { |
| 310 | t.Fatalf("effectiveOutputBudget: %v", err) |
| 311 | } |
| 312 | if !clipped || budget <= 0 || budget >= prov.budget { |
| 313 | t.Fatalf("mixed-script request budget = %d clipped=%v, want a clipped positive budget below %d", budget, clipped, prov.budget) |
| 314 | } |
| 315 | } |
| 316 | |
| 317 | func TestCalibrationIgnoresNonReplayableOrdinaryReasoning(t *testing.T) { |
| 318 | a := &Agent{} |
| 319 | previous := []provider.Message{ |
| 320 | {Role: provider.RoleUser, Content: strings.Repeat("x", 300_000)}, |
| 321 | {Role: provider.RoleAssistant, ReasoningContent: strings.Repeat("hidden", 150_000)}, |
| 322 | } |
| 323 | a.setPromptTokenCalibration(75_000, requestCalibrationShapeOf(provider.Request{Messages: previous})) |
| 324 | current := append(append([]provider.Message(nil), previous...), provider.Message{ |
| 325 | Role: provider.RoleAssistant, |
| 326 | ReasoningContent: strings.Repeat("r", 400_000), |
| 327 | ToolCalls: []provider.ToolCall{{ID: "call_1", Name: "bash", Arguments: `{}`}}, |
| 328 | }) |
| 329 | |
| 330 | if got := a.estimatedPromptTokens(current); got < 160_000 { |
| 331 | t.Fatalf("replayable reasoning estimate = %d, want ordinary local reasoning excluded from calibration denominator", got) |
| 332 | } |
| 333 | } |
| 334 | |
| 335 | func TestCalibratedResponsesBudgetIncludesNewOrdinaryReasoning(t *testing.T) { |
| 336 | prov := &sharedWindowTestProvider{budget: 128 * 1024, shared: true, |
| 337 | policy: provider.SharedWindowInputPolicy{ReplaysOrdinaryReasoning: true}} |
| 338 | a := &Agent{agentConfig: agentConfig{contextWindow: 200_000}, svc: agentServices{prov: prov}, sess: sessionRuntime{output: outputBudgetState{outputBudget: prov.budget}}} |
| 339 | previous := provider.Request{Messages: []provider.Message{{ |
| 340 | Role: provider.RoleUser, Content: strings.Repeat("x", 300_000), |
| 341 | }}} |
| 342 | a.setPromptTokenCalibration(75_000, a.requestCalibrationShape(previous)) |
| 343 | current := previous |
| 344 | current.Messages = append(append([]provider.Message(nil), previous.Messages...), provider.Message{ |
| 345 | Role: provider.RoleAssistant, ReasoningContent: strings.Repeat("r", 400_000), |
| 346 | }) |
| 347 | |
| 348 | if got := a.estimatedRequestTokens(current); got < 174_000 { |
| 349 | t.Fatalf("Responses ordinary reasoning estimate = %d, want newly replayed reasoning included", got) |
| 350 | } |
| 351 | if budget, clipped, err := a.effectiveOutputBudget(current); err != nil || !clipped || budget >= prov.budget { |
| 352 | t.Fatalf("Responses ordinary reasoning budget = %d clipped=%v err=%v, want a clipped budget", budget, clipped, err) |
| 353 | } |
| 354 | } |
| 355 | |
| 356 | func TestCalibratedResponsesBudgetIncludesNewReplayItems(t *testing.T) { |
| 357 | prov := &sharedWindowTestProvider{budget: 128 * 1024, shared: true, |
| 358 | policy: provider.SharedWindowInputPolicy{ReplaysResponsesItems: true}} |
| 359 | a := &Agent{agentConfig: agentConfig{contextWindow: 200_000}, svc: agentServices{prov: prov}, sess: sessionRuntime{output: outputBudgetState{outputBudget: prov.budget}}} |
| 360 | previous := provider.Request{Messages: []provider.Message{{ |
| 361 | Role: provider.RoleUser, Content: strings.Repeat("x", 300_000), |
| 362 | }}} |
| 363 | a.setPromptTokenCalibration(75_000, a.requestCalibrationShape(previous)) |
| 364 | item := json.RawMessage(`{"id":"ws_1","type":"web_search_call","status":"completed","action":{"query":"` + strings.Repeat("q", 400_000) + `"}}`) |
| 365 | current := previous |
| 366 | current.Messages = append(append([]provider.Message(nil), previous.Messages...), provider.Message{ |
| 367 | Role: provider.RoleAssistant, ResponsesItems: []json.RawMessage{item}, |
| 368 | }) |
| 369 | |
| 370 | if got := a.estimatedRequestTokens(current); got < 174_000 { |
| 371 | t.Fatalf("Responses replay-item estimate = %d, want newly replayed item included", got) |
| 372 | } |
| 373 | if budget, clipped, err := a.effectiveOutputBudget(current); err != nil || !clipped || budget >= prov.budget { |
| 374 | t.Fatalf("Responses replay-item budget = %d clipped=%v err=%v, want a clipped budget", budget, clipped, err) |
| 375 | } |
| 376 | } |
| 377 | |
| 378 | func TestCalibratedOutputBudgetCountsToolSchemasOnce(t *testing.T) { |
| 379 | a := &Agent{} |
| 380 | req := provider.Request{ |
| 381 | Messages: []provider.Message{{Role: provider.RoleUser, Content: strings.Repeat("x", 100_000)}}, |
| 382 | Tools: []provider.ToolSchema{{ |
| 383 | Name: "lookup", Description: strings.Repeat("y", 100_000), Parameters: []byte(`{"type":"object"}`), |
| 384 | }}, |
| 385 | } |
| 386 | a.setPromptTokenCalibration(60_000, requestCalibrationShapeOf(req)) |
| 387 | |
| 388 | if got := a.estimatedRequestTokens(req); got != 60_000 { |
| 389 | t.Fatalf("calibrated request tokens = %d, want tool schema counted once in 60000", got) |
| 390 | } |
| 391 | } |
| 392 | |
| 393 | func TestPrepareSamplingRequestClipsSharedWindowOutput(t *testing.T) { |
| 394 | prov := &sharedWindowTestProvider{budget: 128 * 1024, shared: true} |
| 395 | msgs := []provider.Message{{Role: provider.RoleUser, Content: strings.Repeat("字", 950_000)}} |
| 396 | sess := NewSession("") |
| 397 | sess.Replace(msgs) |
| 398 | // compactRatio 2 disables auto maintenance for this output-clip test |
| 399 | a := &Agent{agentConfig: agentConfig{contextWindow: 1_048_576, compactRatio: 2}, svc: agentServices{prov: prov, tools: tool.NewRegistry()}, |
| 400 | sess: sessionRuntime{conversation: sess, output: outputBudgetState{outputBudget: prov.budget}}} |
| 401 | a.sess.output.lastUsage.Store(&provider.Usage{PromptTokens: 950_000}) |
| 402 | a.setPromptTokenCalibration(950_000, requestCalibrationShapeOf(provider.Request{Messages: msgs})) |
| 403 | |
| 404 | prepared, err := a.prepareSamplingRequest(context.Background()) |
| 405 | if err != nil { |
| 406 | t.Fatalf("prepareSamplingRequest: %v", err) |
| 407 | } |
| 408 | if prepared.req.MaxTokens <= 0 || prepared.req.MaxTokens >= prov.budget { |
| 409 | t.Fatalf("prepared MaxTokens = %d, want a clipped positive budget below %d", prepared.req.MaxTokens, prov.budget) |
| 410 | } |
| 411 | } |
| 412 | |
| 413 | func TestEffectiveOutputBudgetRejectsExhaustedSharedWindow(t *testing.T) { |
| 414 | prov := &sharedWindowTestProvider{budget: 128 * 1024, shared: true} |
| 415 | a := &Agent{agentConfig: agentConfig{contextWindow: 1_048_576}, svc: agentServices{prov: prov}, sess: sessionRuntime{output: outputBudgetState{outputBudget: prov.budget}}} |
| 416 | msgs := []provider.Message{{Role: provider.RoleUser, Content: strings.Repeat("字", 1_045_000)}} |
| 417 | a.sess.output.lastUsage.Store(&provider.Usage{PromptTokens: 1_045_000}) |
| 418 | a.setPromptTokenCalibration(1_045_000, requestCalibrationShapeOf(provider.Request{Messages: msgs})) |
| 419 | |
| 420 | _, _, err := a.effectiveOutputBudget(provider.Request{Messages: msgs}) |
| 421 | if !errors.Is(err, ErrCompactionRequired) { |
| 422 | t.Fatalf("effectiveOutputBudget error = %v, want ErrCompactionRequired", err) |
| 423 | } |
| 424 | } |
| 425 | |
| 426 | func TestEffectiveOutputBudgetLeavesIndependentProviderUnchanged(t *testing.T) { |
| 427 | prov := &sharedWindowTestProvider{budget: 128 * 1024, shared: false} |
| 428 | a := &Agent{agentConfig: agentConfig{contextWindow: 1_048_576}, svc: agentServices{prov: prov}, sess: sessionRuntime{output: outputBudgetState{outputBudget: prov.budget}}} |
| 429 | got, clipped, err := a.effectiveOutputBudget(provider.Request{ |
| 430 | Messages: []provider.Message{{Role: provider.RoleUser, Content: strings.Repeat("字", 950_000)}}, |
| 431 | }) |
| 432 | if err != nil || clipped || got != 0 { |
| 433 | t.Fatalf("independent provider changed: budget=%d clipped=%v err=%v", got, clipped, err) |
| 434 | } |
| 435 | } |
| 436 | |
| 437 | func TestEffectiveOutputBudgetHonorsExplicitOmit(t *testing.T) { |
| 438 | prov := &sharedWindowTestProvider{budget: 128 * 1024, shared: true} |
| 439 | a := &Agent{agentConfig: agentConfig{contextWindow: 1_048_576}, svc: agentServices{prov: prov}, sess: sessionRuntime{output: outputBudgetState{outputBudget: prov.budget}}} |
| 440 | got, clipped, err := a.effectiveOutputBudget(provider.Request{ |
| 441 | Messages: []provider.Message{{Role: provider.RoleUser, Content: strings.Repeat("字", 950_000)}}, |
| 442 | MaxTokens: -1, |
| 443 | }) |
| 444 | if err != nil || clipped || got != 0 { |
| 445 | t.Fatalf("explicit omit changed: budget=%d clipped=%v err=%v", got, clipped, err) |
| 446 | } |
| 447 | } |
| 448 | |
| 449 | func TestSummarizeClipsSharedWindowOutputBudget(t *testing.T) { |
| 450 | prov := &sharedWindowTestProvider{budget: 128 * 1024, shared: true} |
| 451 | a := &Agent{agentConfig: agentConfig{contextWindow: 100_000}, svc: agentServices{prov: prov, sink: event.Discard}, sess: sessionRuntime{output: outputBudgetState{outputBudget: prov.budget}}} |
| 452 | region := []provider.Message{{Role: provider.RoleUser, Content: strings.Repeat("字", 50_000)}} |
| 453 | a.sess.output.lastUsage.Store(&provider.Usage{PromptTokens: 50_000}) |
| 454 | a.setPromptTokenCalibration(50_000, requestCalibrationShapeOf(provider.Request{Messages: region})) |
| 455 | |
| 456 | if _, _, err := a.summarize(context.Background(), region, ""); err != nil { |
| 457 | t.Fatalf("summarize: %v", err) |
| 458 | } |
| 459 | if prov.last.MaxTokens <= 0 || prov.last.MaxTokens >= prov.budget { |
| 460 | t.Fatalf("summarizer MaxTokens = %d, want a clipped positive budget below %d", prov.last.MaxTokens, prov.budget) |
| 461 | } |
| 462 | } |
| 463 | |
| 464 | func TestSummarizeRejectsLengthTruncation(t *testing.T) { |
| 465 | prov := &sharedWindowTestProvider{budget: 128 * 1024, shared: true, finish: "length"} |
| 466 | a := &Agent{agentConfig: agentConfig{contextWindow: 1_048_576}, svc: agentServices{prov: prov, sink: event.Discard}, sess: sessionRuntime{output: outputBudgetState{outputBudget: prov.budget}}} |
| 467 | |
| 468 | _, _, err := a.summarizeOnce(context.Background(), []provider.Message{{ |
| 469 | Role: provider.RoleUser, Content: "retain every durable fact", |
| 470 | }}, "") |
| 471 | if err == nil || !strings.Contains(err.Error(), "truncated") { |
| 472 | t.Fatalf("summarizeOnce error = %v, want truncation failure", err) |
| 473 | } |
| 474 | if prov.calls != 1 { |
| 475 | t.Fatalf("length-truncated summary calls = %d, want no identical retry", prov.calls) |
| 476 | } |
| 477 | } |
| 478 | |
| 479 | func TestSetSessionResetsPerTranscriptUsageState(t *testing.T) { |
| 480 | a := &Agent{} |
| 481 | a.sess.output.lastUsage.Store(&provider.Usage{PromptTokens: 200_000}) |
| 482 | active := requestCalibrationShape{requestChars: 900_000, compactChars: 850_000} |
| 483 | a.sess.output.activeReqShape.Store(&active) |
| 484 | a.setPromptTokenCalibration(200_000, requestCalibrationShape{requestChars: 1_000_000, compactChars: 950_000}) |
| 485 | a.learnContextBudget(1_048_576, 384_000, true) |
| 486 | a.storeAdmission(contextAdmission{ |
| 487 | WindowMode: provider.ContextWindowShared.String(), WindowTokens: 1_048_576, |
| 488 | PromptTokens: 810_882, LastRecovery: contextRecoveryLearnedRetry, |
| 489 | }) |
| 490 | a.SetSession(NewSession("new")) |
| 491 | |
| 492 | if got := a.sess.output.lastUsage.Load(); got != nil { |
| 493 | t.Fatalf("lastUsage survived session switch: %+v", got) |
| 494 | } |
| 495 | if got := a.sess.output.activeReqShape.Load(); got != nil { |
| 496 | t.Fatalf("activeReqShape survived session switch: %+v", got) |
| 497 | } |
| 498 | if got := a.sess.output.promptCalibration.Load(); got == nil { |
| 499 | t.Fatal("promptCalibration was dropped on session switch; the tokenizer ratio outlives the transcript") |
| 500 | } |
| 501 | if got := a.sess.output.learned.Load(); got == nil || got.windowTokens != 1_048_576 || got.completionBudget != 384_000 { |
| 502 | t.Fatalf("learned provider budget was dropped on session switch: %+v", got) |
| 503 | } |
| 504 | if got := a.sess.output.admission.Load(); got != nil { |
| 505 | t.Fatalf("context admission survived session switch: %+v", got) |
| 506 | } |
| 507 | if got := a.ContextMaintenanceSnapshot().ContextBudget; got != nil { |
| 508 | t.Fatalf("new transcript exposed the previous context budget: %+v", got) |
| 509 | } |
| 510 | } |
| 511 | |
| 512 | func TestLatestUsagePairsWithActiveRequestSize(t *testing.T) { |
| 513 | a := &Agent{} |
| 514 | active := requestCalibrationShape{requestChars: 222, compactChars: 111, cjkRunes: 22, cjkBytes: 66} |
| 515 | a.sess.output.activeReqShape.Store(&active) |
| 516 | a.storeLatestRequestUsage(&provider.Usage{PromptTokens: 100}) |
| 517 | |
| 518 | if got := a.sess.output.promptCalibration.Load(); got == nil || got.promptTokens != 100 || got.requestChars != 222 || got.compactChars != 111 || got.cjkRunes != 22 || got.cjkBytes != 66 { |
| 519 | t.Fatalf("promptCalibration = %+v, want promptTokens=100 requestChars=222 compactChars=111 cjkRunes=22 cjkBytes=66", got) |
| 520 | } |
| 521 | } |
| 522 | |
| 523 | func TestEstimatedUsageDoesNotReplacePromptCalibration(t *testing.T) { |
| 524 | a := &Agent{} |
| 525 | active := requestCalibrationShape{requestChars: 200_000, compactChars: 100_000} |
| 526 | a.sess.output.activeReqShape.Store(&active) |
| 527 | a.setPromptTokenCalibration(50_000, requestCalibrationShape{requestChars: 100_000, compactChars: 80_000}) |
| 528 | |
| 529 | a.storeLatestRequestUsage(&provider.Usage{ |
| 530 | PromptTokens: 10_000, |
| 531 | TotalTokens: 10_100, |
| 532 | Estimated: true, |
| 533 | }) |
| 534 | |
| 535 | got := a.sess.output.promptCalibration.Load() |
| 536 | if got == nil || got.promptTokens != 50_000 || got.requestChars != 100_000 || got.compactChars != 80_000 { |
| 537 | t.Fatalf("estimated usage replaced provider calibration: %+v", got) |
| 538 | } |
| 539 | if latest := a.sess.output.lastUsage.Load(); latest == nil || !latest.Estimated { |
| 540 | t.Fatalf("estimated usage was not retained for accounting: %+v", latest) |
| 541 | } |
| 542 | } |
| 543 | |
| 544 | func TestCalibratedBudgetIgnoresEncryptedSearchRaw(t *testing.T) { |
| 545 | prov := &sharedWindowTestProvider{budget: 128 * 1024, shared: true} |
| 546 | a := &Agent{agentConfig: agentConfig{contextWindow: 200_000}, svc: agentServices{prov: prov}, sess: sessionRuntime{output: outputBudgetState{outputBudget: prov.budget}}} |
| 547 | previous := provider.Request{Messages: []provider.Message{{ |
| 548 | Role: provider.RoleUser, Content: strings.Repeat("x", 300_000), |
| 549 | }}} |
| 550 | a.setPromptTokenCalibration(75_000, a.requestCalibrationShape(previous)) |
| 551 | visible := provider.ServerSearchCall{ |
| 552 | ID: "s1", Query: "latest", |
| 553 | Results: []provider.ServerSearchHit{{Title: "Change Log", URL: "https://api-docs.deepseek.com/updates/"}}, |
| 554 | } |
| 555 | withRaw := previous |
| 556 | withRaw.Messages = append(append([]provider.Message(nil), previous.Messages...), provider.Message{ |
| 557 | Role: provider.RoleAssistant, Content: "answer", |
| 558 | ServerSearch: []provider.ServerSearchCall{{ |
| 559 | ID: visible.ID, Query: visible.Query, Results: visible.Results, |
| 560 | Raw: json.RawMessage(`[{"encrypted_content":"` + strings.Repeat("E", 400_000) + `"}]`), |
| 561 | }}, |
| 562 | }) |
| 563 | withoutRaw := previous |
| 564 | withoutRaw.Messages = append(append([]provider.Message(nil), previous.Messages...), provider.Message{ |
| 565 | Role: provider.RoleAssistant, |
| 566 | Content: "answer", |
| 567 | ServerSearch: []provider.ServerSearchCall{visible}, |
| 568 | }) |
| 569 | if got, want := a.estimatedRequestTokens(withRaw), a.estimatedRequestTokens(withoutRaw); got != want { |
| 570 | t.Fatalf("estimate with encrypted raw = %d, without = %d", got, want) |
| 571 | } |
| 572 | wantBudget, wantClipped, wantErr := a.effectiveOutputBudget(withoutRaw) |
| 573 | gotBudget, gotClipped, gotErr := a.effectiveOutputBudget(withRaw) |
| 574 | if gotBudget != wantBudget || gotClipped != wantClipped || (gotErr != nil) != (wantErr != nil) { |
| 575 | t.Fatalf("encrypted raw changed output budget: got %d clipped=%v err=%v, want %d clipped=%v err=%v", |
| 576 | gotBudget, gotClipped, gotErr, wantBudget, wantClipped, wantErr) |
| 577 | } |
| 578 | } |
| 579 | |
| 580 | func TestForkCaptureProviderPreservesOutputBudgetCapabilities(t *testing.T) { |
| 581 | t.Setenv("REASONIX_EXPERIMENT_FORK_CAPTURE_DIR", t.TempDir()) |
| 582 | prov := &sharedWindowTestProvider{budget: 128 * 1024, shared: true, |
| 583 | policy: provider.SharedWindowInputPolicy{ReplaysOrdinaryReasoning: true, ReplaysResponsesItems: true}} |
| 584 | a := New(prov, tool.NewRegistry(), NewSession(""), Options{}, event.Discard) |
| 585 | |
| 586 | if !sharesContextWindow(a.svc.prov) { |
| 587 | t.Fatal("fork capture wrapper erased shared-window output capability") |
| 588 | } |
| 589 | if got := outputBudgetOf(a.svc.prov); got != prov.budget { |
| 590 | t.Fatalf("wrapped output budget = %d, want %d", got, prov.budget) |
| 591 | } |
| 592 | if got := sharedWindowInputPolicyOf(a.svc.prov); got != prov.policy { |
| 593 | t.Fatalf("wrapped input policy = %+v, want %+v", got, prov.policy) |
| 594 | } |
| 595 | } |
| 596 |