| 1 | package agent |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "errors" |
| 7 | "reflect" |
| 8 | "sync" |
| 9 | "testing" |
| 10 | |
| 11 | "reasonix/internal/event" |
| 12 | "reasonix/internal/provider" |
| 13 | "reasonix/internal/tool" |
| 14 | ) |
| 15 | |
| 16 | type budgetTestTool struct{} |
| 17 | |
| 18 | func (budgetTestTool) Name() string { return "budget_fixture" } |
| 19 | func (budgetTestTool) Description() string { return "Budget recovery request fixture." } |
| 20 | func (budgetTestTool) Schema() json.RawMessage { |
| 21 | return json.RawMessage(`{"type":"object","properties":{"q":{"type":"string"}}}`) |
| 22 | } |
| 23 | func (budgetTestTool) ReadOnly() bool { return true } |
| 24 | func (budgetTestTool) Execute(context.Context, json.RawMessage) (string, error) { |
| 25 | return "ok", nil |
| 26 | } |
| 27 | |
| 28 | func sameProviderRequestExceptMaxTokens(a, b provider.Request) bool { |
| 29 | a.MaxTokens = 0 |
| 30 | b.MaxTokens = 0 |
| 31 | return reflect.DeepEqual(a, b) |
| 32 | } |
| 33 | |
| 34 | type scriptedBudgetProvider struct { |
| 35 | mu sync.Mutex |
| 36 | policy provider.ContextBudgetPolicy |
| 37 | errs []error |
| 38 | reqs []provider.Request |
| 39 | texts []string |
| 40 | } |
| 41 | |
| 42 | func (p *scriptedBudgetProvider) Name() string { return "scripted-budget" } |
| 43 | func (p *scriptedBudgetProvider) ContextBudgetPolicy() provider.ContextBudgetPolicy { |
| 44 | return p.policy |
| 45 | } |
| 46 | func (p *scriptedBudgetProvider) Stream(_ context.Context, req provider.Request) (<-chan provider.Chunk, error) { |
| 47 | p.mu.Lock() |
| 48 | p.reqs = append(p.reqs, req) |
| 49 | idx := len(p.reqs) - 1 |
| 50 | var err error |
| 51 | if idx < len(p.errs) { |
| 52 | err = p.errs[idx] |
| 53 | } |
| 54 | text := "ok" |
| 55 | if idx < len(p.texts) && p.texts[idx] != "" { |
| 56 | text = p.texts[idx] |
| 57 | } |
| 58 | p.mu.Unlock() |
| 59 | if err != nil { |
| 60 | return nil, err |
| 61 | } |
| 62 | ch := make(chan provider.Chunk, 2) |
| 63 | ch <- provider.Chunk{Type: provider.ChunkText, Text: text} |
| 64 | ch <- provider.Chunk{Type: provider.ChunkDone} |
| 65 | close(ch) |
| 66 | return ch, nil |
| 67 | } |
| 68 | |
| 69 | func issue8909Limit() *provider.ContextLimitError { |
| 70 | return &provider.ContextLimitError{ |
| 71 | APIError: &provider.APIError{Provider: "p", Status: 400, Body: "context"}, |
| 72 | WindowTokens: 1_048_576, |
| 73 | RequestedTokens: 1_165_351, |
| 74 | PromptTokens: 810_882, |
| 75 | CompletionTokens: 354_469, |
| 76 | } |
| 77 | } |
| 78 | |
| 79 | func newBudgetAgent(t *testing.T, p provider.Provider) *Agent { |
| 80 | t.Helper() |
| 81 | sess := NewSession("") |
| 82 | sess.Replace([]provider.Message{{Role: provider.RoleUser, Content: "continue"}}) |
| 83 | registry := tool.NewRegistry() |
| 84 | registry.Add(budgetTestTool{}) |
| 85 | return New(p, registry, sess, Options{ContextWindow: 1_048_576, CompactRatio: 2, MaxOutputTokens: 0, Temperature: 0.25}, event.Discard) |
| 86 | } |
| 87 | |
| 88 | func TestContextLimitRecoveryChangesOnlyOutputField(t *testing.T) { |
| 89 | prov := &scriptedBudgetProvider{ |
| 90 | policy: provider.ContextBudgetPolicy{ |
| 91 | WindowMode: provider.ContextWindowShared, AutoOutputTokens: 384_000, |
| 92 | MaxOutputTokens: 384_000, LimitMode: provider.OutputLimitOmitWhenSafe, |
| 93 | }, |
| 94 | errs: []error{issue8909Limit(), nil}, |
| 95 | } |
| 96 | a := newBudgetAgent(t, prov) |
| 97 | a.sess.conversation.Replace([]provider.Message{ |
| 98 | { |
| 99 | Role: provider.RoleAssistant, Content: "tool preface", ReasoningContent: "provider reasoning", |
| 100 | ReasoningSignature: "reasoning-signature", ReasoningID: "reasoning-id", ReasoningStatus: "completed", |
| 101 | ToolCalls: []provider.ToolCall{{ID: "call-1", Name: "budget_fixture", Arguments: `{"q":"status"}`, ThoughtSignature: "thought-signature"}}, |
| 102 | ResponsesItems: []json.RawMessage{json.RawMessage(`{"type":"reasoning","id":"item-1"}`)}, |
| 103 | ServerSearch: []provider.ServerSearchCall{{ |
| 104 | ID: "search-1", Query: "context budgets", |
| 105 | Results: []provider.ServerSearchHit{{Title: "Result", URL: "https://example.test"}}, |
| 106 | Raw: json.RawMessage(`{"query":"context budgets"}`), |
| 107 | }}, |
| 108 | }, |
| 109 | {Role: provider.RoleTool, Name: "budget_fixture", ToolCallID: "call-1", Content: "done"}, |
| 110 | {Role: provider.RoleUser, Content: "continue", Images: []string{"data:image/png;base64,AA=="}}, |
| 111 | }) |
| 112 | beforeMessages := a.sess.conversation.Snapshot() |
| 113 | got := a.streamWithSamplingRecovery(WithResponseFormat(context.Background(), "json_object"), 1) |
| 114 | if got.err != nil { |
| 115 | t.Fatalf("recovery failed: %v", got.err) |
| 116 | } |
| 117 | prov.mu.Lock() |
| 118 | defer prov.mu.Unlock() |
| 119 | if len(prov.reqs) != 2 { |
| 120 | t.Fatalf("requests = %d, want 2", len(prov.reqs)) |
| 121 | } |
| 122 | if !sameProviderRequestExceptMaxTokens(prov.reqs[0], prov.reqs[1]) { |
| 123 | t.Fatalf("provider request changed outside MaxTokens:\nfirst=%+v\nretry=%+v", prov.reqs[0], prov.reqs[1]) |
| 124 | } |
| 125 | if prov.reqs[1].MaxTokens != 229_502 { |
| 126 | t.Fatalf("retry MaxTokens = %d, want 229502", prov.reqs[1].MaxTokens) |
| 127 | } |
| 128 | if a.lastAdmission().LastRecovery != contextRecoveryLearnedRetry { |
| 129 | t.Fatalf("last recovery = %s", a.lastAdmission().LastRecovery) |
| 130 | } |
| 131 | budget := a.ContextMaintenanceSnapshot().ContextBudget |
| 132 | if budget == nil { |
| 133 | t.Fatal("missing context budget snapshot after learned retry") |
| 134 | } |
| 135 | if budget.Source != provider.ContextBudgetSourceLearned || budget.WindowMode != provider.ContextWindowShared.String() { |
| 136 | t.Fatalf("retry source/window = %s/%s, want learned/shared", budget.Source, budget.WindowMode) |
| 137 | } |
| 138 | if budget.RequestedOutputTokens != 384_000 || budget.EffectiveOutputTokens != 229_502 || budget.PhysicalRemaining != 229_502 || !budget.Clipped { |
| 139 | t.Fatalf("retry budget = %+v, want requested=384000 effective=physical=229502 clipped", budget) |
| 140 | } |
| 141 | if budget.ObservedWindow != 1_048_576 || budget.ObservedPrompt != 810_882 || budget.ObservedCompletion != 354_469 { |
| 142 | t.Fatalf("retry observations = %+v", budget) |
| 143 | } |
| 144 | if after := a.sess.conversation.Snapshot(); !reflect.DeepEqual(after, beforeMessages) { |
| 145 | t.Fatalf("recovery mutated the transcript:\nbefore=%+v\nafter=%+v", beforeMessages, after) |
| 146 | } |
| 147 | } |
| 148 | |
| 149 | // An overflow without token numbers (Zhipu GLM 1261) must not be answered |
| 150 | // with the same prompt under a clipped output cap: the estimate that admitted |
| 151 | // the request is exactly what the provider rejected, so recovery goes straight |
| 152 | // to overflow compaction and retries the rebuilt request. |
| 153 | func TestUnnumberedContextLimitSkipsIdenticalRetry(t *testing.T) { |
| 154 | prov := &scriptedBudgetProvider{ |
| 155 | policy: provider.ContextBudgetPolicy{ |
| 156 | WindowMode: provider.ContextWindowShared, AutoOutputTokens: 384_000, |
| 157 | MaxOutputTokens: 384_000, LimitMode: provider.OutputLimitOmitWhenSafe, |
| 158 | }, |
| 159 | errs: []error{&provider.ContextLimitError{APIError: &provider.APIError{ |
| 160 | Provider: "glm", Status: 400, Body: `{"error":{"code":"1261","message":"Prompt exceeds max length"}}`, |
| 161 | }}}, |
| 162 | } |
| 163 | a := newBudgetAgent(t, prov) |
| 164 | a.sess.conversation.Replace(foldableSessionOverForce(6).Messages) |
| 165 | |
| 166 | got := a.streamWithSamplingRecovery(context.Background(), 1) |
| 167 | if got.err != nil { |
| 168 | t.Fatalf("recovery failed: %v", got.err) |
| 169 | } |
| 170 | prov.mu.Lock() |
| 171 | defer prov.mu.Unlock() |
| 172 | if len(prov.reqs) != 3 { |
| 173 | t.Fatalf("requests = %d, want rejected sampling, one summary, and the rebuilt sampling", len(prov.reqs)) |
| 174 | } |
| 175 | if !requestContains(prov.reqs[1], "Compact the preceding conversation prefix") { |
| 176 | t.Fatal("second request must be the overflow compaction summary, not a retry of the rejected prompt") |
| 177 | } |
| 178 | if sameProviderRequestExceptMaxTokens(prov.reqs[0], prov.reqs[2]) { |
| 179 | t.Fatal("the retried request must be rebuilt on the compacted view, not the rejected prompt") |
| 180 | } |
| 181 | if a.lastAdmission().LastRecovery != contextRecoveryCompacted { |
| 182 | t.Fatalf("last recovery = %s, want compacted", a.lastAdmission().LastRecovery) |
| 183 | } |
| 184 | } |
| 185 | |
| 186 | func TestContextLimitRecoveryPublishesUnknownGatewayBudget(t *testing.T) { |
| 187 | limit := &provider.ContextLimitError{ |
| 188 | APIError: &provider.APIError{Provider: "compatible", Status: 400, Body: "context"}, |
| 189 | WindowTokens: 20_000, |
| 190 | RequestedTokens: 25_000, |
| 191 | PromptTokens: 10_000, |
| 192 | CompletionTokens: 15_000, |
| 193 | } |
| 194 | prov := &scriptedBudgetProvider{ |
| 195 | policy: provider.ContextBudgetPolicy{WindowMode: provider.ContextWindowUnknown, LimitMode: provider.OutputLimitOmitWhenSafe}, |
| 196 | errs: []error{limit, nil}, |
| 197 | } |
| 198 | a := newBudgetAgent(t, prov) |
| 199 | got := a.streamWithSamplingRecovery(context.Background(), 1) |
| 200 | if got.err != nil { |
| 201 | t.Fatalf("unknown gateway recovery failed: %v", got.err) |
| 202 | } |
| 203 | prov.mu.Lock() |
| 204 | if len(prov.reqs) != 2 || prov.reqs[0].MaxTokens != 0 || prov.reqs[1].MaxTokens != 1_808 { |
| 205 | t.Fatalf("unknown gateway requests = %+v, want omitted then 1808", prov.reqs) |
| 206 | } |
| 207 | prov.mu.Unlock() |
| 208 | budget := a.ContextMaintenanceSnapshot().ContextBudget |
| 209 | if budget == nil { |
| 210 | t.Fatal("missing learned unknown-gateway budget") |
| 211 | } |
| 212 | if budget.Source != provider.ContextBudgetSourceLearned || budget.WindowMode != provider.ContextWindowShared.String() || |
| 213 | budget.AutoOutputTokens != 15_000 || budget.RequestedOutputTokens != 15_000 || |
| 214 | budget.EffectiveOutputTokens != 1_808 || budget.PhysicalRemaining != 1_808 || !budget.Clipped || |
| 215 | budget.LastRecovery != contextRecoveryLearnedRetry { |
| 216 | t.Fatalf("unknown gateway retry budget = %+v", budget) |
| 217 | } |
| 218 | } |
| 219 | |
| 220 | func TestContextLimitRecoveryRetriesOriginalRequestOnlyOnce(t *testing.T) { |
| 221 | limit := issue8909Limit() |
| 222 | limit.PromptTokens = 1_040_000 |
| 223 | limit.CompletionTokens = 20_000 |
| 224 | limit.RequestedTokens = 1_060_000 |
| 225 | prov := &scriptedBudgetProvider{ |
| 226 | policy: provider.ContextBudgetPolicy{ |
| 227 | WindowMode: provider.ContextWindowShared, AutoOutputTokens: 384_000, |
| 228 | LimitMode: provider.OutputLimitOmitWhenSafe, |
| 229 | }, |
| 230 | errs: []error{limit, limit, limit}, |
| 231 | } |
| 232 | a := newBudgetAgent(t, prov) |
| 233 | got := a.streamWithSamplingRecovery(context.Background(), 1) |
| 234 | if got.err == nil { |
| 235 | t.Fatal("expected terminal context overflow") |
| 236 | } |
| 237 | if a.lastAdmission().LastRecovery != contextRecoveryFailed { |
| 238 | t.Fatalf("last recovery = %s, want failed", a.lastAdmission().LastRecovery) |
| 239 | } |
| 240 | if provider.AsContextLimitError(got.err) == nil && !errors.Is(got.err, ErrCompactionRequired) { |
| 241 | t.Fatalf("terminal err = %v", got.err) |
| 242 | } |
| 243 | prov.mu.Lock() |
| 244 | defer prov.mu.Unlock() |
| 245 | if got := len(prov.reqs); got != 2 { |
| 246 | t.Fatalf("provider requests = %d, want initial request plus one retry", got) |
| 247 | } |
| 248 | } |
| 249 | |
| 250 | func TestContextBudgetLearnAndSnapshotRace(t *testing.T) { |
| 251 | a := &Agent{agentConfig: agentConfig{contextWindow: 1_048_576}, sess: sessionRuntime{conversation: NewSession("")}} |
| 252 | done := make(chan struct{}) |
| 253 | go func() { |
| 254 | defer close(done) |
| 255 | for i := range 200 { |
| 256 | a.learnContextBudget(1_000_000-i, 1000+i, true) |
| 257 | a.setLastRecovery(contextRecoveryLearnedRetry) |
| 258 | _ = a.ContextMaintenanceSnapshot() |
| 259 | _ = a.effectiveContextWindow() |
| 260 | } |
| 261 | }() |
| 262 | for range 200 { |
| 263 | a.learnContextBudget(900_000, 2000, true) |
| 264 | _ = a.ContextMaintenanceSnapshot() |
| 265 | _ = a.lastAdmission() |
| 266 | } |
| 267 | <-done |
| 268 | } |
| 269 | |
| 270 | func TestThreeStateMaxOutputTokens(t *testing.T) { |
| 271 | prov := &policyWindowProvider{policy: provider.ContextBudgetPolicy{ |
| 272 | WindowMode: provider.ContextWindowShared, AutoOutputTokens: 384_000, |
| 273 | MaxOutputTokens: 384_000, LimitMode: provider.OutputLimitOmitWhenSafe, |
| 274 | }} |
| 275 | a := &Agent{agentConfig: agentConfig{contextWindow: 1_048_576}, svc: agentServices{prov: prov}} |
| 276 | msgs := []provider.Message{{Role: provider.RoleUser, Content: "hi"}} |
| 277 | pos := provider.Request{Messages: msgs, MaxTokens: 8192} |
| 278 | if err := a.applyAdmissionToRequest(&pos); err != nil || pos.MaxTokens != 8192 { |
| 279 | t.Fatalf("positive cap = %d err=%v", pos.MaxTokens, err) |
| 280 | } |
| 281 | zero := provider.Request{Messages: msgs, MaxTokens: 0} |
| 282 | if err := a.applyAdmissionToRequest(&zero); err != nil || zero.MaxTokens != 0 { |
| 283 | t.Fatalf("auto omit = %d err=%v", zero.MaxTokens, err) |
| 284 | } |
| 285 | neg := provider.Request{Messages: msgs, MaxTokens: -1} |
| 286 | if err := a.applyAdmissionToRequest(&neg); err != nil || neg.MaxTokens != -1 { |
| 287 | t.Fatalf("explicit omit = %d err=%v", neg.MaxTokens, err) |
| 288 | } |
| 289 | } |
| 290 |