1package core
2
3import (
4 "context"
5 "reflect"
6 "strings"
7 "testing"
8 "unicode/utf8"
9
10 "github.com/bnema/zut/packages/provider"
11)
12
13func TestProjectToolResultMessagesBoundsAndUTF8(t *testing.T) {
14 msgs := make([]provider.Message, 0, 10)
15 for i := 0; i < 5; i++ {
16 id := "call-" + string(rune('1'+i))
17 msgs = append(msgs,
18 provider.Message{
19 Role: provider.RoleAssistant,
20 Content: []provider.Content{provider.ToolCallBlock{ID: id, Name: "read"}},
21 },
22 provider.Message{
23 Role: provider.RoleTool,
24 Content: []provider.Content{provider.ToolResultBlock{
25 CallID: id,
26 Content: []provider.Content{provider.TextBlock{Text: strings.Repeat("界", maxToolResultTextBytes+100)}},
27 }},
28 },
29 )
30 }
31
32 projected := projectToolResultMessages(msgs)
33 total := 0
34 for _, message := range projected {
35 for _, content := range message.Content {
36 result, ok := content.(provider.ToolResultBlock)
37 if !ok {
38 continue
39 }
40 retainedBytes := retainedToolResultTextBytes(result.Content)
41 if retainedBytes > maxToolResultTextBytes {
42 t.Fatalf("result %q retains %d text bytes, want <= %d", result.CallID, retainedBytes, maxToolResultTextBytes)
43 }
44 if !hasToolResultOmissionMarker(result.Content) {
45 t.Fatalf("truncated result %q is missing the omission marker", result.CallID)
46 }
47 for _, inner := range result.Content {
48 if text, ok := inner.(provider.TextBlock); ok && !utf8.ValidString(text.Text) {
49 t.Fatalf("result %q contains invalid UTF-8", result.CallID)
50 }
51 }
52 total += retainedBytes
53 }
54 }
55 if total > maxToolResultTotalTextBytes {
56 t.Fatalf("projected tool-result text = %d bytes, want <= %d", total, maxToolResultTotalTextBytes)
57 }
58}
59
60func TestProjectToolResultMessagesRetainsNewestResultsFirst(t *testing.T) {
61 msgs := make([]provider.Message, 0, 10)
62 for i := 1; i <= 5; i++ {
63 id := "call-" + string(rune('0'+i))
64 text := strings.Repeat("result-"+id+" ", maxToolResultTextBytes/4)
65 msgs = append(msgs,
66 provider.Message{
67 Role: provider.RoleAssistant,
68 Content: []provider.Content{provider.ToolCallBlock{ID: id, Name: "read"}},
69 },
70 provider.Message{
71 Role: provider.RoleTool,
72 Content: []provider.Content{provider.ToolResultBlock{CallID: id, Content: []provider.Content{provider.TextBlock{Text: text}}}},
73 },
74 )
75 }
76
77 projected := projectToolResultMessages(msgs)
78 oldest := toolResultTextForCall(projected, "call-1")
79 newest := toolResultTextForCall(projected, "call-5")
80 if oldest != toolResultOmissionMarker {
81 t.Fatalf("oldest result = %q, want omission marker after newer results filled budget", oldest)
82 }
83 if !strings.HasPrefix(newest, "result-call-5 ") || !strings.Contains(newest, toolResultOmissionMarker) {
84 t.Fatalf("newest result was not retained and truncated first: %q", newest[:min(len(newest), 80)])
85 }
86}
87
88func TestProjectToolResultMessagesMarksResultsAfterAggregateBudgetIsExhausted(t *testing.T) {
89 const resultCount = maxToolResultTotalTextBytes/maxToolResultTextBytes + 1
90 msgs := make([]provider.Message, 0, resultCount)
91 for i := 0; i < resultCount; i++ {
92 id := "call-" + string(rune('1'+i))
93 msgs = append(msgs, provider.Message{
94 Role: provider.RoleTool,
95 Content: []provider.Content{provider.ToolResultBlock{
96 CallID: id,
97 Content: []provider.Content{provider.TextBlock{Text: strings.Repeat("x", maxToolResultTextBytes+1)}},
98 }},
99 })
100 }
101
102 projected := projectToolResultMessages(msgs)
103 retainedTotal := 0
104 for _, message := range projected {
105 for _, content := range message.Content {
106 result, ok := content.(provider.ToolResultBlock)
107 if !ok {
108 continue
109 }
110 if !hasToolResultOmissionMarker(result.Content) {
111 t.Fatalf("result %q is missing the omission marker", result.CallID)
112 }
113 retainedTotal += retainedToolResultTextBytes(result.Content)
114 }
115 }
116 if retainedTotal != maxToolResultTotalTextBytes {
117 t.Fatalf("retained tool-result text = %d bytes, want exactly %d", retainedTotal, maxToolResultTotalTextBytes)
118 }
119 if got := toolResultTextForCall(projected, "call-1"); got != toolResultOmissionMarker {
120 t.Fatalf("aggregate-exhausted result = %q, want omission marker", got)
121 }
122}
123
124func TestProjectToolResultMessagesPreservesToolResultBlocksAndRawInput(t *testing.T) {
125 msgs := []provider.Message{
126 {Role: provider.RoleAssistant, Content: []provider.Content{
127 provider.ToolCallBlock{ID: "call-1", Name: "read"},
128 provider.ToolCallBlock{ID: "call-2", Name: "bash"},
129 }},
130 {Role: provider.RoleTool, Content: []provider.Content{
131 provider.ToolResultBlock{CallID: "call-1", Content: []provider.Content{provider.TextBlock{Text: strings.Repeat("one", maxToolResultTextBytes)}}},
132 provider.ToolResultBlock{CallID: "call-2", Content: []provider.Content{provider.TextBlock{Text: strings.Repeat("two", maxToolResultTextBytes)}}},
133 }},
134 }
135 wantRaw := append([]provider.Message(nil), msgs...)
136
137 projected := projectToolResultMessages(msgs)
138 if !reflect.DeepEqual(msgs, wantRaw) {
139 t.Fatal("projection mutated the raw input transcript")
140 }
141 var resultIDs []string
142 for _, content := range projected[1].Content {
143 result, ok := content.(provider.ToolResultBlock)
144 if !ok {
145 t.Fatalf("projected content = %T, want ToolResultBlock", content)
146 }
147 resultIDs = append(resultIDs, result.CallID)
148 }
149 if !reflect.DeepEqual(resultIDs, []string{"call-1", "call-2"}) {
150 t.Fatalf("projected result IDs = %v, want both tool results in order", resultIDs)
151 }
152 if got := toolResultTextForCall(msgs, "call-1"); got != strings.Repeat("one", maxToolResultTextBytes) {
153 t.Fatal("raw tool result text changed after projection")
154 }
155}
156
157func TestAgentNormalRequestProjectsToolResultsAfterPairRepair(t *testing.T) {
158 client := &captureClient{}
159 agent := NewAgent(client, "model", "system", Registry{})
160 full := strings.Repeat("full-normal-result ", maxToolResultTextBytes)
161 agent.SetMessages([]provider.Message{
162 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "continue"}}},
163 {Role: provider.RoleAssistant, Content: []provider.Content{provider.ToolCallBlock{ID: "call-1", Name: "read"}}},
164 {Role: provider.RoleTool, Content: []provider.Content{provider.ToolResultBlock{
165 CallID: "call-1",
166 Content: []provider.Content{provider.TextBlock{Text: full}},
167 }}},
168 {Role: provider.RoleAssistant, Content: []provider.Content{provider.ToolCallBlock{ID: "call-2", Name: "read"}}},
169 })
170
171 if err := agent.Continue(context.Background(), nil); err != nil {
172 t.Fatalf("Continue: %v", err)
173 }
174 projectedText := toolResultTextForCall(client.lastReq.Messages, "call-1")
175 if retainedToolResultTextBytesForCall(client.lastReq.Messages, "call-1") > maxToolResultTextBytes || !strings.Contains(projectedText, toolResultOmissionMarker) {
176 t.Fatalf("normal request result was not bounded: %d bytes, %q", len(projectedText), projectedText[:min(len(projectedText), 80)])
177 }
178 if !hasToolResultCall(client.lastReq.Messages, "call-2") {
179 t.Fatal("pair repair result was not preserved in normal request")
180 }
181 if got := toolResultTextForCall(agent.Messages(), "call-1"); got != full {
182 t.Fatalf("raw Agent.Messages tool result changed: got %d bytes, want %d", len(got), len(full))
183 }
184 if hasToolResultCall(agent.Messages(), "call-2") {
185 t.Fatal("pair-repair stub leaked into raw Agent.Messages")
186 }
187}
188
189func TestCompactionProjectsSummaryInputWithoutChangingPersistedTail(t *testing.T) {
190 client := &compactLifecycleClient{}
191 agent := NewAgent(client, "model", "system", Registry{})
192 oldText := strings.Repeat("full-compaction-result ", maxToolResultTextBytes)
193 tailText := strings.Repeat("full-tail-result ", maxToolResultTextBytes)
194 agent.SetMessages([]provider.Message{
195 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "summarize this"}}},
196 {Role: provider.RoleAssistant, Content: []provider.Content{provider.ToolCallBlock{ID: "old-call", Name: "read"}}},
197 {Role: provider.RoleTool, Content: []provider.Content{provider.ToolResultBlock{
198 CallID: "old-call",
199 Content: []provider.Content{provider.TextBlock{Text: oldText}},
200 }}},
201 {Role: provider.RoleAssistant, Content: []provider.Content{provider.ToolCallBlock{ID: "tail-call", Name: "read"}}},
202 {Role: provider.RoleTool, Content: []provider.Content{provider.ToolResultBlock{
203 CallID: "tail-call",
204 Content: []provider.Content{provider.TextBlock{Text: tailText}},
205 }}},
206 })
207 var persisted []provider.Message
208 agent.OnTranscriptCompacted = func(messages []provider.Message) {
209 persisted = messages
210 }
211
212 if _, err := agent.Compact(context.Background(), 2, nil); err != nil {
213 t.Fatalf("Compact: %v", err)
214 }
215 requestText := client.req.Messages[0].Content[0].(provider.TextBlock).Text
216 if !strings.Contains(requestText, toolResultOmissionMarker) {
217 t.Fatal("compaction request is missing the tool-result omission marker")
218 }
219 if got := toolResultTextForCall(agent.Messages(), "tail-call"); got != tailText {
220 t.Fatalf("compacted Agent.Messages tail changed: got %d bytes, want %d", len(got), len(tailText))
221 }
222 if got := toolResultTextForCall(persisted, "tail-call"); got != tailText {
223 t.Fatalf("persisted compacted tail changed: got %d bytes, want %d", len(got), len(tailText))
224 }
225 if hasToolResultCall(agent.Messages(), "old-call") {
226 t.Fatal("summarized tool result remained in the compacted transcript")
227 }
228}
229
230func retainedToolResultTextBytes(content []provider.Content) int {
231 total := 0
232 for _, block := range content {
233 if text, ok := block.(provider.TextBlock); ok && text.Text != toolResultOmissionMarker {
234 total += len(text.Text)
235 }
236 }
237 return total
238}
239
240func hasToolResultOmissionMarker(content []provider.Content) bool {
241 for _, block := range content {
242 if text, ok := block.(provider.TextBlock); ok && text.Text == toolResultOmissionMarker {
243 return true
244 }
245 }
246 return false
247}
248
249func retainedToolResultTextBytesForCall(messages []provider.Message, callID string) int {
250 for _, message := range messages {
251 for _, content := range message.Content {
252 result, ok := content.(provider.ToolResultBlock)
253 if ok && result.CallID == callID {
254 return retainedToolResultTextBytes(result.Content)
255 }
256 }
257 }
258 return 0
259}
260
261func toolResultTextForCall(messages []provider.Message, callID string) string {
262 var text strings.Builder
263 for _, message := range messages {
264 for _, content := range message.Content {
265 result, ok := content.(provider.ToolResultBlock)
266 if !ok || result.CallID != callID {
267 continue
268 }
269 for _, inner := range result.Content {
270 if block, ok := inner.(provider.TextBlock); ok {
271 text.WriteString(block.Text)
272 }
273 }
274 }
275 }
276 return text.String()
277}
278
279func hasToolResultCall(messages []provider.Message, callID string) bool {
280 for _, message := range messages {
281 for _, content := range message.Content {
282 if result, ok := content.(provider.ToolResultBlock); ok && result.CallID == callID {
283 return true
284 }
285 }
286 }
287 return false
288}
289
290func min(a, b int) int {
291 if a < b {
292 return a
293 }
294 return b
295}