diff --git a/internal/providers/gemini_test.go b/internal/providers/gemini_test.go index cb1522b7..d450ba36 100644 --- a/internal/providers/gemini_test.go +++ b/internal/providers/gemini_test.go @@ -3,6 +3,7 @@ package providers import ( "encoding/json" "reflect" + "strings" "testing" "gophergate/internal/models" @@ -209,3 +210,57 @@ func TestEmitGeminiChunk_ToolCalls(t *testing.T) { t.Fatalf("expected response on channel") } } + +func TestStreamGeminiJSONArrayStream(t *testing.T) { + input := `[ + { + "candidates": [ + { + "content": { + "parts": [ + {"text": "Hello"} + ] + } + } + ] + }, + { + "candidates": [ + { + "content": { + "parts": [ + {"text": " world"} + ] + } + } + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 2, + "totalTokenCount": 7 + } + } + ]` + + ch := make(chan *models.ChatCompletionStreamResponse, 5) + defer close(ch) + + streamGeminiJSONArrayStream(strings.NewReader(input), ch, "gemini-2.5-flash") + + var texts []string + for i := 0; i < 2; i++ { + select { + case resp := <-ch: + if len(resp.Choices) > 0 && resp.Choices[0].Delta.Content != nil { + texts = append(texts, *resp.Choices[0].Delta.Content) + } + default: + t.Fatalf("expected chunk %d", i+1) + } + } + + joined := strings.Join(texts, "") + if joined != "Hello world" { + t.Errorf("expected 'Hello world', got %q", joined) + } +} diff --git a/internal/providers/helpers.go b/internal/providers/helpers.go index 84e99877..159c054e 100644 --- a/internal/providers/helpers.go +++ b/internal/providers/helpers.go @@ -523,9 +523,12 @@ func StreamGemini(ctx io.ReadCloser, model string) (<-chan *models.ChatCompletio first := string(peek[:n]) if first[0] == '[' { - // JSON array format - rest, _ := io.ReadAll(ctx) - streamGeminiJSONArray(append([]byte(first), rest...), ch, model) + // JSON array format — stream parse it in real-time + combined := io.MultiReader( + strings.NewReader(string(peek[:n])), + ctx, + ) + streamGeminiJSONArrayStream(combined, ch, model) return } else if strings.HasPrefix(first, "data:") || strings.HasPrefix(first, "data: ") { // SSE format — pre-pend the peeked bytes then run SSE scanner @@ -548,35 +551,48 @@ func StreamGemini(ctx io.ReadCloser, model string) (<-chan *models.ChatCompletio return ch, nil } -// readAll reads remaining bytes from a reader (keeps the function signature simple -// for the JSON array fallback path). -func readAll(r io.Reader) []byte { - b, _ := io.ReadAll(r) - return b -} +func streamGeminiJSONArrayStream(r io.Reader, ch chan<- *models.ChatCompletionStreamResponse, model string) { + dec := json.NewDecoder(r) -func streamGeminiJSONArray(data []byte, ch chan<- *models.ChatCompletionStreamResponse, model string) { - var chunks []geminiStreamChunk - if err := json.Unmarshal(data, &chunks); err != nil { - fmt.Printf("[Gemini-Stream] JSON array parse error: %v\n", err) + // Read open bracket '[' + t, err := dec.Token() + if err != nil { + fmt.Printf("[Gemini-Stream] JSON array token error: %v\n", err) return } - // Track the last chunk with usage for the final emission + delim, ok := t.(json.Delim) + if !ok || delim != '[' { + fmt.Printf("[Gemini-Stream] JSON array expected '[', got %v\n", t) + return + } + var lastUsage *geminiStreamChunk - for i := range chunks { - if chunks[i].UsageMetadata.TotalTokenCount > 0 { - lastUsage = &chunks[i] + + // Read array elements + for dec.More() { + var chunk geminiStreamChunk + if err := dec.Decode(&chunk); err != nil { + fmt.Printf("[Gemini-Stream] JSON array decode error: %v\n", err) + return + } + + if chunk.UsageMetadata.TotalTokenCount > 0 { + temp := chunk + lastUsage = &temp + } + + // Emit content-bearing chunks immediately + if len(chunk.Candidates) > 0 { + emitGeminiChunk(ch, &chunk, model) } } - if lastUsage != nil { - // Emit a synthetic final chunk with usage data - if len(lastUsage.Candidates) == 0 && lastUsage.UsageMetadata.TotalTokenCount > 0 { - emitGeminiChunk(ch, lastUsage, model) - } - } - // Also emit each content-bearing chunk - for i := range chunks { - emitGeminiChunk(ch, &chunks[i], model) + + // Read close bracket ']' + _, _ = dec.Token() + + // Emit synthetic final chunk with usage if we collected it and it was not yet emitted + if lastUsage != nil && len(lastUsage.Candidates) == 0 && lastUsage.UsageMetadata.TotalTokenCount > 0 { + emitGeminiChunk(ch, lastUsage, model) } }