1package provider
2
3import (
4 "context"
5 "encoding/base64"
6 "encoding/json"
7 "io"
8 "net/http"
9 "net/http/httptest"
10 "os"
11 "strings"
12 "testing"
13)
14
15// TestGeminiStreamHappyPath drives the gemini client end-to-end against
16// a fake SSE server speaking the Gemini Generative Language wire format.
17// We assert text deltas accumulate, usage rolls in from usageMetadata,
18// and the final EventDone carries StopEnd.
19func TestGeminiStreamHappyPath(t *testing.T) {
20 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
21 if got := r.Header.Get("x-goog-api-key"); got != "k" {
22 t.Errorf("missing api key header: %q", got)
23 }
24 if !strings.Contains(r.URL.RawQuery, "alt=sse") {
25 t.Errorf("missing alt=sse: %q", r.URL.RawQuery)
26 }
27 w.Header().Set("content-type", "text/event-stream")
28 w.WriteHeader(http.StatusOK)
29 fl, _ := w.(http.Flusher)
30 write := func(s string) {
31 _, _ = w.Write([]byte(s))
32 if fl != nil {
33 fl.Flush()
34 }
35 }
36 // Two text chunks, then a usage-only finish chunk.
37 write("data: " + `{"candidates":[{"content":{"role":"model","parts":[{"text":"hel"}]}}]}` + "\n\n")
38 write("data: " + `{"candidates":[{"content":{"role":"model","parts":[{"text":"lo"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":12,"candidatesTokenCount":2,"totalTokenCount":14}}` + "\n\n")
39 }))
40 defer srv.Close()
41
42 c := NewGemini("k", srv.URL)
43 evs, err := c.Stream(context.Background(), Request{Model: "gemini-2.5-pro"})
44 if err != nil {
45 t.Fatal(err)
46 }
47 var gotText string
48 var done EventDone
49 var usage Usage
50 for ev := range evs {
51 switch e := ev.(type) {
52 case EventTextDelta:
53 gotText += e.Delta
54 case EventUsage:
55 usage = e.Usage
56 case EventDone:
57 done = e
58 }
59 }
60 if gotText != "hello" {
61 t.Fatalf("text=%q", gotText)
62 }
63 if done.Stop != StopEnd {
64 t.Fatalf("stop=%v", done.Stop)
65 }
66 if usage.InputTokens != 12 || usage.OutputTokens != 2 {
67 t.Fatalf("usage=%+v", usage)
68 }
69}
70
71func TestGeminiStreamInlineImage(t *testing.T) {
72 dir := t.TempDir()
73 prev, err := os.Getwd()
74 if err != nil {
75 t.Fatal(err)
76 }
77 if err := os.Chdir(dir); err != nil {
78 t.Fatal(err)
79 }
80 t.Cleanup(func() { _ = os.Chdir(prev) })
81 img := base64.StdEncoding.EncodeToString([]byte("png-bytes"))
82 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
83 w.Header().Set("content-type", "text/event-stream")
84 w.WriteHeader(http.StatusOK)
85 _, _ = w.Write([]byte("data: " + `{"candidates":[{"content":{"role":"model","parts":[{"text":"here"},{"inlineData":{"mimeType":"image/png","data":"` + img + `"}}]},"finishReason":"STOP"}]}` + "\n\n"))
86 }))
87 defer srv.Close()
88
89 c := NewGemini("k", srv.URL)
90 evs, err := c.Stream(context.Background(), Request{Model: "gemini-2.5-flash-image"})
91 if err != nil {
92 t.Fatal(err)
93 }
94 var done EventDone
95 for ev := range evs {
96 if e, ok := ev.(EventDone); ok {
97 done = e
98 }
99 }
100 if len(done.Message.Content) != 3 {
101 t.Fatalf("content count=%d: %+v", len(done.Message.Content), done.Message.Content)
102 }
103 if tb, ok := done.Message.Content[0].(TextBlock); !ok || tb.Text != "here" {
104 t.Fatalf("text block=%T %+v", done.Message.Content[0], done.Message.Content[0])
105 }
106 ib, ok := done.Message.Content[1].(ImageBlock)
107 if !ok {
108 t.Fatalf("want image block, got %T", done.Message.Content[1])
109 }
110 if ib.MimeType != "image/png" || string(ib.Data) != "png-bytes" {
111 t.Fatalf("image=%q %q", ib.MimeType, string(ib.Data))
112 }
113 saved, ok := done.Message.Content[2].(TextBlock)
114 if !ok || !strings.Contains(saved.Text, "zut-gemini-image-") || !strings.Contains(saved.Text, ".png") {
115 t.Fatalf("saved path block=%T %+v", done.Message.Content[2], done.Message.Content[2])
116 }
117 path := strings.TrimPrefix(saved.Text, "Saved image: `")
118 path = strings.TrimSuffix(path, "`")
119 if _, err := os.Stat(path); err != nil {
120 t.Fatalf("saved image missing at %q: %v", path, err)
121 }
122}
123
124// TestGeminiToolCall covers the tool-call branch: a single
125// functionCall part should produce ToolStart/Args/End and the final
126// stop reason should be StopToolUse.
127func TestGeminiToolCall(t *testing.T) {
128 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
129 w.Header().Set("content-type", "text/event-stream")
130 w.WriteHeader(http.StatusOK)
131 fl, _ := w.(http.Flusher)
132 write := func(s string) {
133 _, _ = w.Write([]byte(s))
134 if fl != nil {
135 fl.Flush()
136 }
137 }
138 write("data: " + `{"candidates":[{"content":{"role":"model","parts":[{"functionCall":{"name":"read","args":{"path":"a"}}}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":3,"candidatesTokenCount":7,"totalTokenCount":10}}` + "\n\n")
139 }))
140 defer srv.Close()
141
142 c := NewGemini("k", srv.URL)
143 evs, err := c.Stream(context.Background(), Request{Model: "gemini-2.5-pro"})
144 if err != nil {
145 t.Fatal(err)
146 }
147 var startName, argsBuf string
148 var done EventDone
149 for ev := range evs {
150 switch e := ev.(type) {
151 case EventToolStart:
152 startName = e.Name
153 case EventToolArgs:
154 argsBuf += e.Delta
155 case EventDone:
156 done = e
157 }
158 }
159 if startName != "read" {
160 t.Fatalf("tool name=%q", startName)
161 }
162 if !strings.Contains(argsBuf, `"path":"a"`) {
163 t.Fatalf("tool args=%q", argsBuf)
164 }
165 if done.Stop != StopToolUse {
166 t.Fatalf("stop=%v", done.Stop)
167 }
168 // The assembled message should contain the tool call as the first content block.
169 if len(done.Message.Content) != 1 {
170 t.Fatalf("message content count=%d", len(done.Message.Content))
171 }
172 tc, ok := done.Message.Content[0].(ToolCallBlock)
173 if !ok {
174 t.Fatalf("expected ToolCallBlock, got %T", done.Message.Content[0])
175 }
176 if tc.Name != "read" {
177 t.Fatalf("tool block name=%q", tc.Name)
178 }
179}
180
181// TestGeminiBuildRequestSystemAndTools confirms the wire payload puts
182// the system prompt under systemInstruction and tool defs under tools[0].
183func TestGeminiBuildRequestSystemAndTools(t *testing.T) {
184 c := NewGemini("k", "https://example.invalid").(*geminiClient)
185 wire, _, err := c.buildRequest(Request{
186 Model: "gemini-2.5-pro",
187 System: "you are zut",
188 Tools: []Tool{
189 {Name: "read", Description: "read a file", Schema: json.RawMessage(`{"type":"object","properties":{"path":{"type":"string"}}}`)},
190 },
191 Messages: []Message{{Role: RoleUser, Content: []Content{TextBlock{Text: "hi"}}}},
192 })
193 if err != nil {
194 t.Fatal(err)
195 }
196 if wire.SystemInstruction == nil || wire.SystemInstruction.Parts[0].Text != "you are zut" {
197 t.Fatalf("system: %+v", wire.SystemInstruction)
198 }
199 if len(wire.Tools) != 1 || len(wire.Tools[0].FunctionDeclarations) != 1 || wire.Tools[0].FunctionDeclarations[0].Name != "read" {
200 t.Fatalf("tools: %+v", wire.Tools)
201 }
202 if len(wire.Contents) != 1 || wire.Contents[0].Role != "user" {
203 t.Fatalf("contents: %+v", wire.Contents)
204 }
205}
206
207// TestGeminiBuildRequestImageModelOmitsTools confirms image-generation
208// models receive direct multimodal prompts without function declarations.
209func TestGeminiBuildRequestStripsUnsupportedSchemaFields(t *testing.T) {
210 c := NewGemini("k", "https://example.invalid").(*geminiClient)
211 wire, _, err := c.buildRequest(Request{
212 Model: "gemini-2.5-pro",
213 Tools: []Tool{{
214 Name: "edit",
215 Description: "edit a file",
216 Schema: json.RawMessage(`{
217 "$schema":"http://json-schema.org/draft-07/schema#",
218 "type":"object",
219 "additionalProperties":false,
220 "properties":{
221 "edits":{
222 "type":"array",
223 "items":{
224 "type":"object",
225 "additionalProperties":false,
226 "properties":{"oldText":{"type":"string"},"newText":{"type":"string"}}
227 }
228 }
229 }
230 }`),
231 }},
232 })
233 if err != nil {
234 t.Fatal(err)
235 }
236 got := string(wire.Tools[0].FunctionDeclarations[0].Parameters)
237 if strings.Contains(got, "additionalProperties") || strings.Contains(got, "$schema") {
238 t.Fatalf("Gemini schema should strip unsupported fields, got %s", got)
239 }
240 if !strings.Contains(got, `"oldText"`) || !strings.Contains(got, `"newText"`) {
241 t.Fatalf("Gemini schema lost nested properties, got %s", got)
242 }
243}
244
245func TestGeminiBuildRequestImageModelOmitsTools(t *testing.T) {
246 c := NewGemini("k", "https://example.invalid").(*geminiClient)
247 wire, _, err := c.buildRequest(Request{
248 Model: "gemini-2.5-flash-image",
249 Tools: []Tool{
250 {Name: "read", Description: "read a file", Schema: json.RawMessage(`{"type":"object","properties":{"path":{"type":"string"}}}`)},
251 },
252 Messages: []Message{
253 {Role: RoleUser, Content: []Content{TextBlock{Text: "edit this image"}, ImageBlock{MimeType: "image/png", Data: []byte("png")}}},
254 {Role: RoleAssistant, Content: []Content{TextBlock{Text: "checking"}, ToolCallBlock{ID: "1", Name: "read", Arguments: json.RawMessage(`{"path":"x"}`)}}},
255 {Role: RoleTool, Content: []Content{ToolResultBlock{CallID: "1", Content: []Content{TextBlock{Text: "tool output"}}}}},
256 },
257 })
258 if err != nil {
259 t.Fatal(err)
260 }
261 if len(wire.Tools) != 0 {
262 t.Fatalf("image model should omit tools, got %+v", wire.Tools)
263 }
264 if len(wire.Contents) != 2 {
265 t.Fatalf("contents count=%d: %+v", len(wire.Contents), wire.Contents)
266 }
267 if len(wire.Contents[0].Parts) != 2 || wire.Contents[0].Parts[1].InlineData == nil {
268 t.Fatalf("user image parts not preserved: %+v", wire.Contents[0].Parts)
269 }
270 if len(wire.Contents[1].Parts) != 1 || wire.Contents[1].Parts[0].FunctionCall != nil {
271 t.Fatalf("assistant tool call not stripped: %+v", wire.Contents[1].Parts)
272 }
273}
274
275// TestGeminiErrorStatus confirms HTTP error bodies bubble up.
276func TestGeminiErrorStatus(t *testing.T) {
277 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
278 w.WriteHeader(http.StatusUnauthorized)
279 _, _ = io.WriteString(w, `{"error":{"code":401,"message":"bad key"}}`)
280 }))
281 defer srv.Close()
282
283 c := NewGemini("k", srv.URL)
284 _, err := c.Stream(context.Background(), Request{Model: "gemini-2.5-pro"})
285 if err == nil || !strings.Contains(err.Error(), "401") {
286 t.Fatalf("want 401 err, got %v", err)
287 }
288}
289
290// TestGeminiThinkingConfig spot-checks the level → wire mapping for
291// representative model families. The full table is exercised via
292// integration; this guards the routing logic.
293func TestGeminiThinkingConfig(t *testing.T) {
294 cases := []struct {
295 modelID string
296 level string
297 wantLvl string
298 wantBud int
299 }{
300 {"gemini-3-pro", "low", "LOW", 0},
301 {"gemini-3-pro", "medium", "HIGH", 0}, // Pro can't go below LOW; medium → HIGH
302 {"gemini-3-flash", "medium", "MEDIUM", 0},
303 {"gemini-2.5-pro", "high", "", 16384},
304 {"gemini-2.5-pro", "maximum", "", 32768},
305 {"gemini-2.5-flash", "minimum", "", 1024},
306 {"gemini-2.5-flash", "low", "", 2048},
307 {"gemini-2.0-flash", "high", "", 0}, // 2.0 has no thinking → nil
308 }
309 for _, tc := range cases {
310 got := geminiThinkingConfig(tc.modelID, tc.level)
311 if tc.wantLvl == "" && tc.wantBud == 0 {
312 if got != nil {
313 t.Errorf("%s/%s: want nil, got %+v", tc.modelID, tc.level, got)
314 }
315 continue
316 }
317 if got == nil {
318 t.Errorf("%s/%s: got nil", tc.modelID, tc.level)
319 continue
320 }
321 if tc.wantLvl != "" && got.ThinkingLevel != tc.wantLvl {
322 t.Errorf("%s/%s: level=%q want %q", tc.modelID, tc.level, got.ThinkingLevel, tc.wantLvl)
323 }
324 if tc.wantBud != 0 {
325 if got.ThinkingBudget == nil || *got.ThinkingBudget != tc.wantBud {
326 t.Errorf("%s/%s: budget=%v want %d", tc.modelID, tc.level, got.ThinkingBudget, tc.wantBud)
327 }
328 }
329 }
330}
331
332func TestGeminiThoughtSignaturesStreamAndReplay(t *testing.T) {
333 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
334 w.Header().Set("content-type", "text/event-stream")
335 _, _ = io.WriteString(w, "data: "+`{"candidates":[{"content":{"parts":[{"text":"think","thought":true}]}}]}`+"\n\n")
336 _, _ = io.WriteString(w, "data: "+`{"candidates":[{"content":{"parts":[{"text":"ing","thought":true,"thoughtSignature":"reasoning-sig"}]}}]}`+"\n\n")
337 _, _ = io.WriteString(w, "data: "+`{"candidates":[{"content":{"parts":[{"text":"ans"}]}}]}`+"\n\n")
338 _, _ = io.WriteString(w, "data: "+`{"candidates":[{"content":{"parts":[{"text":"wer","thoughtSignature":"text-sig"}]}}]}`+"\n\n")
339 _, _ = io.WriteString(w, "data: "+`{"candidates":[{"content":{"parts":[{"functionCall":{"name":"read","args":{"path":"a"}},"thoughtSignature":"tool-sig"}]},"finishReason":"STOP"}]}`+"\n\n")
340 }))
341 defer srv.Close()
342
343 client := NewGemini("k", srv.URL).(*geminiClient)
344 events, err := client.Stream(context.Background(), Request{Model: "gemini-3-flash"})
345 if err != nil {
346 t.Fatal(err)
347 }
348 var done EventDone
349 for event := range events {
350 if event, ok := event.(EventDone); ok {
351 done = event
352 }
353 }
354 if len(done.Message.Content) != 3 {
355 t.Fatalf("content = %#v", done.Message.Content)
356 }
357 reasoning, ok := done.Message.Content[0].(ReasoningBlock)
358 if !ok || reasoning.Summary != "thinking" || reasoning.Encrypted != "reasoning-sig" {
359 t.Fatalf("reasoning = %#v", done.Message.Content[0])
360 }
361 text, ok := done.Message.Content[1].(TextBlock)
362 if !ok || text.Text != "answer" || text.ThoughtSignature != "text-sig" {
363 t.Fatalf("text = %#v", done.Message.Content[1])
364 }
365 tool, ok := done.Message.Content[2].(ToolCallBlock)
366 if !ok || tool.ThoughtSignature != "tool-sig" {
367 t.Fatalf("tool = %#v", done.Message.Content[2])
368 }
369
370 wire, _, err := client.buildRequest(Request{
371 Model: "gemini-3-flash",
372 Messages: []Message{done.Message},
373 })
374 if err != nil {
375 t.Fatal(err)
376 }
377 parts := wire.Contents[0].Parts
378 if len(parts) != 3 || parts[0].ThoughtSignature != "reasoning-sig" ||
379 parts[1].ThoughtSignature != "text-sig" || parts[2].ThoughtSignature != "tool-sig" {
380 t.Fatalf("replayed parts = %#v", parts)
381 }
382}
383
384func TestGeminiAssistantImageReplaysThoughtSignature(t *testing.T) {
385 client := NewGemini("k", "https://example.invalid").(*geminiClient)
386 wire, _, err := client.buildRequest(Request{
387 Model: "gemini-3-flash",
388 Messages: []Message{{Role: RoleAssistant, Content: []Content{
389 ImageBlock{MimeType: "image/png", Data: []byte("png"), ThoughtSignature: "image-sig"},
390 }}},
391 })
392 if err != nil {
393 t.Fatal(err)
394 }
395 part := wire.Contents[0].Parts[0]
396 if part.InlineData == nil || part.ThoughtSignature != "image-sig" {
397 t.Fatalf("image part = %#v", part)
398 }
399}
400
401// TestDiscoverGoogle exercises the discovery helper against a fake
402// /v1beta/models endpoint, confirming pagination plus filtering of
403// non-chat ids (embedding, aqa).
404func TestDiscoverGoogle(t *testing.T) {
405 calls := 0
406 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
407 if r.Header.Get("x-goog-api-key") != "k" {
408 t.Errorf("missing api key")
409 }
410 calls++
411 if calls == 1 {
412 _, _ = io.WriteString(w, `{
413 "models": [
414 {"name":"models/gemini-2.5-pro","displayName":"Gemini 2.5 Pro","supportedGenerationMethods":["generateContent","streamGenerateContent"],"inputTokenLimit":1048576,"outputTokenLimit":65536},
415 {"name":"models/text-embedding-004","displayName":"Text Embedding","supportedGenerationMethods":["embedContent"]},
416 {"name":"models/aqa","displayName":"AQA","supportedGenerationMethods":["generateAnswer"]}
417 ],
418 "nextPageToken": "p2"
419 }`)
420 } else {
421 _, _ = io.WriteString(w, `{
422 "models": [
423 {"name":"models/gemini-2.5-flash","displayName":"Gemini 2.5 Flash","supportedGenerationMethods":["generateContent"]}
424 ]
425 }`)
426 }
427 }))
428 defer srv.Close()
429
430 got, err := DiscoverGoogle(context.Background(), "k", srv.URL)
431 if err != nil {
432 t.Fatal(err)
433 }
434 if calls != 2 {
435 t.Errorf("want 2 calls, got %d", calls)
436 }
437 if len(got) != 2 {
438 t.Fatalf("want 2 models, got %d: %+v", len(got), got)
439 }
440 if got[0].ID != "gemini-2.5-pro" || got[0].Provider != "google" {
441 t.Errorf("first model wrong: %+v", got[0])
442 }
443 if got[1].ID != "gemini-2.5-flash" {
444 t.Errorf("second model wrong: %+v", got[1])
445 }
446}
447
448// TestSanitizeGeminiToolSchemaEmptyOrInvalid covers the pre-existing
449// fallback behavior for empty/invalid input, to make sure the new logic
450// didn't regress it.
451func TestSanitizeGeminiToolSchemaEmptyOrInvalid(t *testing.T) {
452 for _, in := range []json.RawMessage{nil, json.RawMessage(``), json.RawMessage(`not json`)} {
453 out := sanitizeGeminiToolSchema(in)
454 var got map[string]any
455 if err := json.Unmarshal(out, &got); err != nil {
456 t.Fatalf("fallback schema is not valid JSON: %v (%s)", err, out)
457 }
458 if got["type"] != "object" {
459 t.Fatalf("want type=object fallback, got %+v", got)
460 }
461 if _, ok := got["properties"].(map[string]any); !ok {
462 t.Fatalf("want properties object fallback, got %+v", got)
463 }
464 }
465}
466
467// TestGeminiStreamReasoning confirms that when the Gemini SSE stream
468// contains thought blocks with a signature, they are correctly mapped to
469// ReasoningBlock structures inside the final assistant message.
470func TestGeminiStreamReasoning(t *testing.T) {
471 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
472 w.Header().Set("content-type", "text/event-stream")
473 w.WriteHeader(http.StatusOK)
474 fl, _ := w.(http.Flusher)
475 write := func(s string) {
476 _, _ = w.Write([]byte(s))
477 if fl != nil {
478 fl.Flush()
479 }
480 }
481 // Write a thought part with signature, then text, then usage containing thoughtsTokenCount.
482 write("data: " + `{"candidates":[{"content":{"role":"model","parts":[{"text":"Thinking hard","thought":true,"thoughtSignature":"sig123"}]}}]}` + "\n\n")
483 write("data: " + `{"candidates":[{"content":{"role":"model","parts":[{"text":"Answer text"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":5,"candidatesTokenCount":2,"thoughtsTokenCount":3,"totalTokenCount":10}}` + "\n\n")
484 }))
485 defer srv.Close()
486
487 c := NewGemini("k", srv.URL)
488 evs, err := c.Stream(context.Background(), Request{Model: "gemini-2.5-pro"})
489 if err != nil {
490 t.Fatal(err)
491 }
492
493 var gotText string
494 var done EventDone
495 var usage Usage
496 for ev := range evs {
497 switch e := ev.(type) {
498 case EventTextDelta:
499 gotText += e.Delta
500 case EventUsage:
501 usage = e.Usage
502 case EventDone:
503 done = e
504 }
505 }
506
507 if gotText != "Answer text" {
508 t.Fatalf("text=%q", gotText)
509 }
510 if usage.OutputTokens != 5 { // candidatesTokenCount (2) + thoughtsTokenCount (3)
511 t.Fatalf("usage output tokens=%d, want 5", usage.OutputTokens)
512 }
513 if usage.ReasoningTokens != 3 || !usage.ReasoningTokensKnown {
514 t.Fatalf("usage reasoning tokens=%d known=%v, want 3 true", usage.ReasoningTokens, usage.ReasoningTokensKnown)
515 }
516
517 // Verify that the final message contains both the ReasoningBlock and the TextBlock
518 msg := done.Message
519 if len(msg.Content) != 2 {
520 t.Fatalf("want 2 content blocks, got %d: %+v", len(msg.Content), msg.Content)
521 }
522
523 rb, ok := msg.Content[0].(ReasoningBlock)
524 if !ok {
525 t.Fatalf("first block should be ReasoningBlock, got %T", msg.Content[0])
526 }
527 if rb.Summary != "Thinking hard" || rb.Encrypted != "sig123" {
528 t.Fatalf("invalid ReasoningBlock: %+v", rb)
529 }
530
531 tb, ok := msg.Content[1].(TextBlock)
532 if !ok {
533 t.Fatalf("second block should be TextBlock, got %T", msg.Content[1])
534 }
535 if tb.Text != "Answer text" {
536 t.Fatalf("invalid TextBlock: %+v", tb)
537 }
538}
539
540// TestGeminiBuildRequestWithReasoningReplay verifies that when zut sends a
541// historical assistant message containing a ReasoningBlock, buildRequest
542// serializes it into the verbatim Gemini "thought: true" wire representation
543// with its corresponding thoughtSignature.
544func TestGeminiBuildRequestWithReasoningReplay(t *testing.T) {
545 c := NewGemini("k", "https://example.invalid").(*geminiClient)
546 wire, _, err := c.buildRequest(Request{
547 Model: "gemini-2.5-pro",
548 Messages: []Message{
549 {Role: RoleUser, Content: []Content{TextBlock{Text: "Why is the sky blue?"}}},
550 {Role: RoleAssistant, Content: []Content{
551 ReasoningBlock{Summary: "Evaluating Rayleigh scattering", Encrypted: "sky-sig-999"},
552 TextBlock{Text: "It scatters blue light more than other colors."},
553 }},
554 },
555 })
556 if err != nil {
557 t.Fatal(err)
558 }
559
560 if len(wire.Contents) != 2 {
561 t.Fatalf("want 2 messages in wire, got %d", len(wire.Contents))
562 }
563
564 assistant := wire.Contents[1]
565 if assistant.Role != "model" {
566 t.Fatalf("assistant role should be 'model', got %q", assistant.Role)
567 }
568
569 if len(assistant.Parts) != 2 {
570 t.Fatalf("want 2 parts, got %d: %+v", len(assistant.Parts), assistant.Parts)
571 }
572
573 // Check the serialized ReasoningBlock
574 p0 := assistant.Parts[0]
575 if !p0.Thought {
576 t.Errorf("part 0: expected Thought=true, got %+v", p0)
577 }
578 if p0.Text != "Evaluating Rayleigh scattering" {
579 t.Errorf("part 0: expected text to match summary, got %q", p0.Text)
580 }
581 if p0.ThoughtSignature != "sky-sig-999" {
582 t.Errorf("part 0: expected thought signature 'sky-sig-999', got %q", p0.ThoughtSignature)
583 }
584
585 // Check the serialized TextBlock
586 p1 := assistant.Parts[1]
587 if p1.Thought {
588 t.Errorf("part 1: expected Thought=false, got %+v", p1)
589 }
590 if p1.Text != "It scatters blue light more than other colors." {
591 t.Errorf("part 1: expected text to match, got %q", p1.Text)
592 }
593}
594
595// TestGeminiStreamToolCallWithThoughtSignature confirms that a tool call
596// featuring a thoughtSignature in the SSE stream successfully extracts and
597// attaches the thought signature onto the final ToolCallBlock.
598func TestGeminiStreamToolCallWithThoughtSignature(t *testing.T) {
599 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
600 w.Header().Set("content-type", "text/event-stream")
601 w.WriteHeader(http.StatusOK)
602 _, _ = w.Write([]byte("data: " + `{"candidates":[{"content":{"role":"model","parts":[{"functionCall":{"name":"read","args":{"path":"a"}},"thoughtSignature":"tool-sig-456"}]},"finishReason":"STOP"}]}` + "\n\n"))
603 }))
604 defer srv.Close()
605
606 c := NewGemini("k", srv.URL)
607 evs, err := c.Stream(context.Background(), Request{Model: "gemini-2.5-pro"})
608 if err != nil {
609 t.Fatal(err)
610 }
611
612 var done EventDone
613 for ev := range evs {
614 if e, ok := ev.(EventDone); ok {
615 done = e
616 }
617 }
618
619 if len(done.Message.Content) != 1 {
620 t.Fatalf("want 1 content block, got %d", len(done.Message.Content))
621 }
622
623 tc, ok := done.Message.Content[0].(ToolCallBlock)
624 if !ok {
625 t.Fatalf("expected ToolCallBlock, got %T", done.Message.Content[0])
626 }
627
628 if tc.Name != "read" {
629 t.Fatalf("tool name=%q, want 'read'", tc.Name)
630 }
631
632 if tc.ThoughtSignature != "tool-sig-456" {
633 t.Fatalf("tool block thought signature=%q, want 'tool-sig-456'", tc.ThoughtSignature)
634 }
635}