| 1 | package agent |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "strings" |
| 7 | "testing" |
| 8 | |
| 9 | "reasonix/internal/event" |
| 10 | "reasonix/internal/provider" |
| 11 | "reasonix/internal/tool" |
| 12 | ) |
| 13 | |
| 14 | func TestSplitExtractChunksCountsReplayResponsesItems(t *testing.T) { |
| 15 | largeItem := json.RawMessage(`"` + strings.Repeat("r", extractChunkNewestBytes) + `"`) |
| 16 | msgs := []provider.Message{ |
| 17 | {Role: provider.RoleUser, Content: "old"}, |
| 18 | {Role: provider.RoleAssistant, Content: "middle", ResponsesItems: []json.RawMessage{largeItem}}, |
| 19 | {Role: provider.RoleUser, Content: "new"}, |
| 20 | } |
| 21 | |
| 22 | replayPolicy := provider.SharedWindowInputPolicy{ReplaysResponsesItems: true} |
| 23 | if chunks := splitExtractChunks(msgs, extractChunkOverlapBytes, replayPolicy); len(chunks) < 2 { |
| 24 | t.Fatalf("replay-aware chunks = %d, want at least 2", len(chunks)) |
| 25 | } |
| 26 | if chunks := splitExtractChunks(msgs, extractChunkOverlapBytes, provider.SharedWindowInputPolicy{}); len(chunks) != 1 { |
| 27 | t.Fatalf("non-replay chunks = %d, want 1", len(chunks)) |
| 28 | } |
| 29 | } |
| 30 | |
| 31 | func TestMessageWireBytesCountsProviderVisibleReplayFields(t *testing.T) { |
| 32 | msg := provider.Message{ |
| 33 | Role: provider.RoleAssistant, |
| 34 | Content: "content", |
| 35 | Images: []string{"image-ref"}, |
| 36 | ReasoningContent: "reasoning", |
| 37 | ReasoningID: "reasoning-id", |
| 38 | ReasoningStatus: "completed", |
| 39 | ReasoningSignature: "reasoning-signature", |
| 40 | ToolCalls: []provider.ToolCall{{ |
| 41 | ID: "call-id", |
| 42 | Name: "tool-name", |
| 43 | Arguments: `{"key":"value"}`, |
| 44 | ThoughtSignature: "thought-signature", |
| 45 | }}, |
| 46 | ResponsesItems: []json.RawMessage{json.RawMessage(`{"type":"reasoning"}`)}, |
| 47 | ServerSearch: []provider.ServerSearchCall{{ |
| 48 | ID: "search-id", |
| 49 | Query: "search-query", |
| 50 | Results: []provider.ServerSearchHit{{ |
| 51 | Title: "result-title", |
| 52 | URL: "https://example.test/result", |
| 53 | }}, |
| 54 | Raw: json.RawMessage(`{"encrypted_content":"must-not-count"}`), |
| 55 | }}, |
| 56 | } |
| 57 | policy := provider.SharedWindowInputPolicy{ReplaysResponsesItems: true} |
| 58 | want := 4 + len(string(msg.Role)) + len(msg.Content) + len(msg.ReasoningContent) + |
| 59 | 8 + len(msg.ToolCalls[0].ID) + len(msg.ToolCalls[0].Name) + len(msg.ToolCalls[0].Arguments) + |
| 60 | len(msg.ResponsesItems[0]) + len(msg.ServerSearch[0].ID) + len(msg.ServerSearch[0].Query) + |
| 61 | len(msg.ServerSearch[0].Results[0].Title) + len(msg.ServerSearch[0].Results[0].URL) + |
| 62 | len(msg.ReasoningID) + len(msg.ReasoningStatus) + len(msg.ReasoningSignature) + |
| 63 | len(msg.ToolCalls[0].ThoughtSignature) + len(msg.Images[0]) |
| 64 | if got := messageWireBytes(msg, policy); got != want { |
| 65 | t.Fatalf("wire bytes = %d, want %d", got, want) |
| 66 | } |
| 67 | |
| 68 | withoutRaw := msg |
| 69 | withoutRaw.ServerSearch = append([]provider.ServerSearchCall(nil), msg.ServerSearch...) |
| 70 | withoutRaw.ServerSearch[0].Raw = nil |
| 71 | if got, want := messageWireBytes(msg, policy), messageWireBytes(withoutRaw, policy); got != want { |
| 72 | t.Fatalf("ServerSearch.Raw changed wire estimate: with=%d without=%d", got, want) |
| 73 | } |
| 74 | |
| 75 | withoutItems := messageWireBytes(msg, provider.SharedWindowInputPolicy{}) |
| 76 | if got := messageWireBytes(msg, policy) - withoutItems; got != len(msg.ResponsesItems[0]) { |
| 77 | t.Fatalf("ResponsesItems replay bytes = %d, want %d", got, len(msg.ResponsesItems[0])) |
| 78 | } |
| 79 | } |
| 80 | |
| 81 | func TestSummarizeExtractChunksPreflightsMinimumPlan(t *testing.T) { |
| 82 | chunk := []provider.Message{{Role: provider.RoleUser, Content: "fragment"}} |
| 83 | tests := []struct { |
| 84 | name string |
| 85 | count int |
| 86 | wantCalls int |
| 87 | wantErr bool |
| 88 | }{ |
| 89 | {name: "63 chunks fit 64 calls", count: 63, wantCalls: 64}, |
| 90 | {name: "64 chunks require 65 calls", count: 64, wantCalls: 0, wantErr: true}, |
| 91 | } |
| 92 | for _, tc := range tests { |
| 93 | t.Run(tc.name, func(t *testing.T) { |
| 94 | prov := &extractStubProvider{reply: "digest"} |
| 95 | a := New(prov, tool.NewRegistry(), extractStubSession(), Options{}, event.Discard) |
| 96 | chunks := make([][]provider.Message, tc.count) |
| 97 | for i := range chunks { |
| 98 | chunks[i] = chunk |
| 99 | } |
| 100 | _, err := a.summarizeExtractChunks(context.Background(), chunks, "", nil, newChunkedSummaryRun(a)) |
| 101 | if (err != nil) != tc.wantErr { |
| 102 | t.Fatalf("summarizeExtractChunks error = %v, wantErr %v", err, tc.wantErr) |
| 103 | } |
| 104 | if prov.calls != tc.wantCalls { |
| 105 | t.Fatalf("provider calls = %d, want %d", prov.calls, tc.wantCalls) |
| 106 | } |
| 107 | }) |
| 108 | } |
| 109 | } |
| 110 | |
| 111 | func TestChunkedSummaryRunReservesCallsAfterRequest(t *testing.T) { |
| 112 | prov := &extractStubProvider{reply: "digest"} |
| 113 | a := New(prov, tool.NewRegistry(), extractStubSession(), Options{}, event.Discard) |
| 114 | run := newChunkedSummaryRun(a) |
| 115 | run.calls = maxChunkedSummaryCalls - 1 |
| 116 | _, err := run.summarize(context.Background(), []provider.Message{{Role: provider.RoleUser, Content: "x"}}, extractMergeInstruction, 1) |
| 117 | if err == nil || !strings.Contains(err.Error(), "call budget exhausted") { |
| 118 | t.Fatalf("reservation error = %v", err) |
| 119 | } |
| 120 | if prov.calls != 0 { |
| 121 | t.Fatalf("provider calls = %d, want 0", prov.calls) |
| 122 | } |
| 123 | } |
| 124 | |
| 125 | func TestRecursiveRecoveryPreservesOuterCallBudget(t *testing.T) { |
| 126 | t.Run("fragment split", func(t *testing.T) { |
| 127 | prov := &extractStubProvider{failFirst: 1, reply: "digest"} |
| 128 | a := New(prov, tool.NewRegistry(), extractStubSession(), Options{}, event.Discard) |
| 129 | run := newChunkedSummaryRun(a) |
| 130 | run.calls = maxChunkedSummaryCalls - 4 |
| 131 | chunk := []provider.Message{ |
| 132 | {Role: provider.RoleUser, Content: "left"}, |
| 133 | {Role: provider.RoleUser, Content: "right"}, |
| 134 | } |
| 135 | _, err := a.extractFragmentResilient(context.Background(), chunk, extractFragmentInstruction(1, 1, ""), extractMergeInstruction, func(bool) {}, run, 1) |
| 136 | if err == nil || !strings.Contains(err.Error(), "call budget exhausted") { |
| 137 | t.Fatalf("fragment recovery error = %v", err) |
| 138 | } |
| 139 | if prov.calls != 1 { |
| 140 | t.Fatalf("provider calls = %d, want 1 failed root and no doomed split", prov.calls) |
| 141 | } |
| 142 | }) |
| 143 | |
| 144 | t.Run("merge split", func(t *testing.T) { |
| 145 | prov := &extractStubProvider{failFirst: 1, reply: "digest"} |
| 146 | a := New(prov, tool.NewRegistry(), extractStubSession(), Options{}, event.Discard) |
| 147 | run := newChunkedSummaryRun(a) |
| 148 | run.calls = maxChunkedSummaryCalls - 4 |
| 149 | _, err := a.mergeGroup(context.Background(), []string{"left", "right"}, extractMergeInstruction, run, 0, 1) |
| 150 | if err == nil || !strings.Contains(err.Error(), "call budget exhausted") { |
| 151 | t.Fatalf("merge recovery error = %v", err) |
| 152 | } |
| 153 | if prov.calls != 1 { |
| 154 | t.Fatalf("provider calls = %d, want 1 failed root and no doomed split", prov.calls) |
| 155 | } |
| 156 | }) |
| 157 | } |
| 158 | |
| 159 | func TestMergeTreeRechecksBudgetBeforeNewRound(t *testing.T) { |
| 160 | prov := &extractStubProvider{reply: strings.Repeat("digest ", 320)} |
| 161 | a := New(prov, tool.NewRegistry(), extractStubSession(), Options{ContextWindow: 2000}, event.Discard) |
| 162 | run := newChunkedSummaryRun(a) |
| 163 | run.calls = maxChunkedSummaryCalls - 3 |
| 164 | parts := []string{ |
| 165 | strings.Repeat("one ", 320), |
| 166 | strings.Repeat("two ", 320), |
| 167 | strings.Repeat("three ", 320), |
| 168 | strings.Repeat("four ", 320), |
| 169 | strings.Repeat("five ", 320), |
| 170 | } |
| 171 | _, err := a.mergeFragmentsWithRun(context.Background(), parts, extractMergeInstruction, run, 0) |
| 172 | if err == nil || !strings.Contains(err.Error(), "call budget exhausted") { |
| 173 | t.Fatalf("merge tree error = %v", err) |
| 174 | } |
| 175 | if prov.calls != 2 { |
| 176 | t.Fatalf("provider calls = %d, want 2 first-round pairs and no doomed second round", prov.calls) |
| 177 | } |
| 178 | } |
| 179 |