[Core] fix(gemini): improve Gemini 3/3.5 thinking level mapping and split stream usage

core/providers/gemini/chat.go

@@ -386,6 +386,8 @@ func (response *GenerateContentResponse) ToBifrostChatResponse() *schemas.Bifros
 type GeminiStreamState struct {
 	nextToolCallIndex int
 	hadToolCalls      bool // true if any tool calls were seen in this stream
+	hadFinishReason   bool // true if a candidate finish reason was seen in this stream
+	groundingMetadata *GroundingMetadata
 }
 
 // NewGeminiStreamState returns initialised stream state for one streaming response.
@@ -521,8 +523,33 @@ func (response *GenerateContentResponse) toBifrostChatCompletionStreamDelta(stat
 		state = NewGeminiStreamState()
 	}
 
-	// Handle empty candidates (filtered/malformed responses)
+	// Handle empty candidates (filtered/malformed responses, or a trailing usage-only chunk)
 	if len(response.Candidates) == 0 {
+		if response.UsageMetadata != nil && state.hadFinishReason {
+			streamResponse := &schemas.BifrostChatResponse{
+				ID:     response.ResponseID,
+				Model:  response.ModelVersion,
+				Object: "chat.completion.chunk",
+				Choices: []schemas.BifrostResponseChoice{{
+					Index: 0,
+					ChatStreamResponseChoice: &schemas.ChatStreamResponseChoice{
+						Delta: &schemas.ChatStreamResponseChoiceDelta{},
+					},
+				}},
+				Usage: ConvertGeminiUsageMetadataToChatUsage(response.UsageMetadata),
+			}
+			if !response.CreateTime.IsZero() {
+				streamResponse.Created = int(response.CreateTime.Unix())
+			}
+			applyGeminiSearchQueryChatUsage(streamResponse.Usage, state.groundingMetadata, response.ModelVersion)
+			if t := mapGeminiTrafficTypeToBifrost(response.UsageMetadata.TrafficType); t != nil {
+				streamResponse.ServiceTier = t
+			} else if response.UsageMetadata.ServiceTier != "" {
+				tier := mapGeminiServiceTierToBifrost(response.UsageMetadata.ServiceTier)
+				streamResponse.ServiceTier = &tier
+			}
+			return streamResponse, nil, true
+		}
 		finishReason := ConvertGeminiFinishReasonToBifrost(FinishReasonMalformedFunctionCall)
 		if response.promptBlockError() != nil {
 			finishReason = ConvertGeminiFinishReasonToBifrost(FinishReasonSafety)
@@ -538,8 +565,15 @@ func (response *GenerateContentResponse) toBifrostChatCompletionStreamDelta(stat
 		return createErrorResponse(response, finishReason, true), nil, true
 	}
 
+	if candidate.FinishReason != "" {
+		state.hadFinishReason = true
+	}
+	if candidate.GroundingMetadata != nil {
+		state.groundingMetadata = candidate.GroundingMetadata
+	}
+
 	// Determine if this is the last chunk based on finish reason and usage metadata
-	isLastChunk := candidate.FinishReason != "" && response.UsageMetadata != nil
+	isLastChunk := state.hadFinishReason && response.UsageMetadata != nil
 
 	// Create the streaming response
 	streamResponse := &schemas.BifrostChatResponse{
@@ -712,14 +746,14 @@ func (response *GenerateContentResponse) toBifrostChatCompletionStreamDelta(stat
 	}
 
 	// Check if delta has any content - if not and it's not the last chunk, skip it
-	hasDeltaContent := delta.Role != nil || delta.Content != nil || delta.Audio != nil || len(delta.ToolCalls) > 0 || len(delta.ReasoningDetails) > 0 || len(delta.Annotations) > 0
+	hasDeltaContent := delta.Role != nil || delta.Content != nil || delta.Audio != nil || len(delta.ToolCalls) > 0 || len(delta.ReasoningDetails) > 0 || len(delta.Annotations) > 0 || candidate.FinishReason != "" || response.UsageMetadata != nil
 	if !hasDeltaContent && !isLastChunk {
 		return nil, nil, false
 	}
 
 	// Build the choice
 	var finishReason *string
-	if isLastChunk && candidate.FinishReason != "" {
+	if candidate.FinishReason != "" {
 		reason := ConvertGeminiFinishReasonToBifrost(candidate.FinishReason)
 		// Gemini uses "STOP" for both text completions and tool call responses.
 		// Override to "tool_calls" when tool calls were seen in this stream for uniformity.
@@ -740,10 +774,15 @@ func (response *GenerateContentResponse) toBifrostChatCompletionStreamDelta(stat
 
 	streamResponse.Choices = []schemas.BifrostResponseChoice{choice}
 
-	// Add usage information if this is the last chunk
-	if isLastChunk && response.UsageMetadata != nil {
+	// Add usage information whenever present so running stream usage stays current
+	// and split finishReason / usageMetadata chunks both survive.
+	if response.UsageMetadata != nil {
 		streamResponse.Usage = ConvertGeminiUsageMetadataToChatUsage(response.UsageMetadata)
-		applyGeminiSearchQueryChatUsage(streamResponse.Usage, candidate.GroundingMetadata, response.ModelVersion)
+		grounding := candidate.GroundingMetadata
+		if grounding == nil {
+			grounding = state.groundingMetadata
+		}
+		applyGeminiSearchQueryChatUsage(streamResponse.Usage, grounding, response.ModelVersion)
 		if t := mapGeminiTrafficTypeToBifrost(response.UsageMetadata.TrafficType); t != nil {
 			streamResponse.ServiceTier = t
 		} else if response.UsageMetadata.ServiceTier != "" {

core/providers/gemini/gemini_test.go

@@ -6141,3 +6141,231 @@ func TestGenAIPartMediaResolution_SurvivesBatchConversion(t *testing.T) {
 	require.NotNil(t, got, "batch conversion must preserve per-part media resolution")
 	assert.Equal(t, "MEDIA_RESOLUTION_ULTRA_HIGH", got.Level)
 }
+
+func TestGemini35ThinkingAndStreamingUsage(t *testing.T) {
+	t.Run("EffortAndMaxTokensOnGemini3EmitsThinkingLevel", func(t *testing.T) {
+		for _, tc := range []struct {
+			model     string
+			effort    string
+			maxTokens int
+			summary   *string
+			display   *string
+			wantLevel string
+			wantInc   bool
+		}{
+			{
+				model:     "gemini-3.5-flash",
+				effort:    "low",
+				maxTokens: 4096,
+				wantLevel: "low",
+				wantInc:   true,
+			},
+			{
+				model:     "au.gemini-3.5-flash",
+				effort:    "minimal",
+				maxTokens: 1024,
+				summary:   schemas.Ptr("none"),
+				display:   schemas.Ptr("omitted"),
+				wantLevel: "minimal",
+				wantInc:   false,
+			},
+			{
+				model:     "gemini-3.5-pro",
+				effort:    "minimal",
+				maxTokens: 2048,
+				wantLevel: "low",
+				wantInc:   true,
+			},
+			{
+				model:     "gemini-3.1-pro",
+				effort:    "low",
+				maxTokens: 4096,
+				summary:   schemas.Ptr("none"),
+				display:   schemas.Ptr("omitted"),
+				wantLevel: "low",
+				wantInc:   false,
+			},
+		} {
+			t.Run(tc.model, func(t *testing.T) {
+				chatReq := &schemas.BifrostChatRequest{
+					Model: tc.model,
+					Input: minimalChatInput(),
+					Params: &schemas.ChatParameters{
+						Reasoning: &schemas.ChatReasoning{
+							Effort:    schemas.Ptr(tc.effort),
+							MaxTokens: schemas.Ptr(tc.maxTokens),
+							Display:   tc.display,
+						},
+					},
+				}
+				chatOut, err := gemini.ToGeminiChatCompletionRequest(nil, chatReq)
+				require.NoError(t, err)
+				require.NotNil(t, chatOut.GenerationConfig.ThinkingConfig)
+				require.NotNil(t, chatOut.GenerationConfig.ThinkingConfig.ThinkingLevel)
+				assert.Equal(t, tc.wantLevel, *chatOut.GenerationConfig.ThinkingConfig.ThinkingLevel)
+				assert.Nil(t, chatOut.GenerationConfig.ThinkingConfig.ThinkingBudget)
+				assert.Equal(t, tc.wantInc, chatOut.GenerationConfig.ThinkingConfig.IncludeThoughts)
+
+				respReq := &schemas.BifrostResponsesRequest{
+					Model: tc.model,
+					Input: []schemas.ResponsesMessage{{
+						Role:    schemas.Ptr(schemas.ResponsesInputMessageRoleUser),
+						Content: &schemas.ResponsesMessageContent{ContentStr: schemas.Ptr("Solve this")},
+					}},
+					Params: &schemas.ResponsesParameters{
+						Reasoning: &schemas.ResponsesParametersReasoning{
+							Effort:    schemas.Ptr(tc.effort),
+							MaxTokens: schemas.Ptr(tc.maxTokens),
+							Summary:   tc.summary,
+						},
+					},
+				}
+				respOut, err := gemini.ToGeminiResponsesRequest(nil, respReq)
+				require.NoError(t, err)
+				require.NotNil(t, respOut.GenerationConfig.ThinkingConfig)
+				require.NotNil(t, respOut.GenerationConfig.ThinkingConfig.ThinkingLevel)
+				assert.Equal(t, tc.wantLevel, *respOut.GenerationConfig.ThinkingConfig.ThinkingLevel)
+				assert.Nil(t, respOut.GenerationConfig.ThinkingConfig.ThinkingBudget)
+				assert.Equal(t, tc.wantInc, respOut.GenerationConfig.ThinkingConfig.IncludeThoughts)
+			})
+		}
+	})
+
+	t.Run("RegionalPrefixSupportsToolCombinationOnVertex", func(t *testing.T) {
+		req := &schemas.BifrostResponsesRequest{
+			Provider: schemas.Vertex,
+			Model:    "au.gemini-3.5-flash",
+			Input: []schemas.ResponsesMessage{{
+				Role:    schemas.Ptr(schemas.ResponsesInputMessageRoleUser),
+				Content: &schemas.ResponsesMessageContent{ContentStr: schemas.Ptr("weather in sydney?")},
+			}},
+			Params: &schemas.ResponsesParameters{
+				Tools: []schemas.ResponsesTool{
+					{Type: schemas.ResponsesToolTypeWebSearch, ResponsesToolWebSearch: &schemas.ResponsesToolWebSearch{}},
+					{
+						Type:                  schemas.ResponsesToolTypeFunction,
+						Name:                  schemas.Ptr("get_weather"),
+						ResponsesToolFunction: &schemas.ResponsesToolFunction{Parameters: &schemas.ToolFunctionParameters{Type: "object"}},
+					},
+				},
+			},
+		}
+		out, err := gemini.ToGeminiResponsesRequest(nil, req)
+		require.NoError(t, err)
+		require.Len(t, out.Tools, 2, "regional Gemini 3+ model must support combining function declarations and Google Search")
+		require.Len(t, out.Tools[0].FunctionDeclarations, 1)
+		assert.NotNil(t, out.Tools[1].GoogleSearch)
+	})
+
+	t.Run("SplitFinishReasonAndUsageMetadataStreamChunks", func(t *testing.T) {
+		state := gemini.NewGeminiStreamState()
+
+		chunk1 := &gemini.GenerateContentResponse{
+			ResponseID:   "resp-split-1",
+			ModelVersion: "gemini-3.5-flash",
+			Candidates: []*gemini.Candidate{{
+				Index:        0,
+				FinishReason: gemini.FinishReasonStop,
+				Content: &gemini.Content{
+					Role:  string(gemini.RoleModel),
+					Parts: []*gemini.Part{{Text: "Hello"}},
+				},
+			}},
+		}
+		deltas1, bifrostErr, isLast1 := chunk1.ToBifrostChatCompletionStream(state)
+		require.Nil(t, bifrostErr)
+		assert.False(t, isLast1, "finishReason chunk without usageMetadata must not terminate stream early")
+		require.Len(t, deltas1, 1)
+		require.Len(t, deltas1[0].Choices, 1)
+		require.NotNil(t, deltas1[0].Choices[0].FinishReason)
+		assert.Equal(t, "stop", *deltas1[0].Choices[0].FinishReason)
+
+		chunk2 := &gemini.GenerateContentResponse{
+			ResponseID:   "resp-split-1",
+			ModelVersion: "gemini-3.5-flash",
+			Candidates:   []*gemini.Candidate{},
+			UsageMetadata: &gemini.GenerateContentResponseUsageMetadata{
+				PromptTokenCount:        2500,
+				CachedContentTokenCount: 2000,
+				CandidatesTokenCount:    10,
+				ThoughtsTokenCount:      120,
+				TotalTokenCount:         2630,
+			},
+		}
+		deltas2, bifrostErr, isLast2 := chunk2.ToBifrostChatCompletionStream(state)
+		require.Nil(t, bifrostErr)
+		assert.True(t, isLast2, "trailing usageMetadata chunk after finishReason must close stream")
+		require.Len(t, deltas2, 1)
+		require.NotNil(t, deltas2[0].Usage)
+		assert.Equal(t, 2500, deltas2[0].Usage.PromptTokens)
+		require.NotNil(t, deltas2[0].Usage.PromptTokensDetails)
+		assert.Equal(t, 2000, deltas2[0].Usage.PromptTokensDetails.CachedReadTokens)
+		require.NotNil(t, deltas2[0].Usage.CompletionTokensDetails)
+		assert.Equal(t, 120, deltas2[0].Usage.CompletionTokensDetails.ReasoningTokens)
+		assert.Equal(t, 2630, deltas2[0].Usage.TotalTokens)
+	})
+}
+
+func BenchmarkGemini3ThinkingAndStreamUsage(b *testing.B) {
+	effort := "medium"
+	budget := 2048
+	summary := "none"
+	chatReq := &schemas.BifrostChatRequest{
+		Provider: schemas.Vertex,
+		Model:    "au.gemini-3.5-flash",
+		Input: []schemas.ChatMessage{
+			{
+				Role:    schemas.ChatMessageRoleUser,
+				Content: &schemas.ChatMessageContent{ContentStr: schemas.Ptr("Hello")},
+			},
+		},
+		Params: &schemas.ChatParameters{
+			Reasoning: &schemas.ChatReasoning{
+				Effort:    &effort,
+				MaxTokens: &budget,
+				Summary:   &summary,
+			},
+		},
+	}
+
+	b.Run("ToGeminiChatCompletionRequest", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			_, _ = gemini.ToGeminiChatCompletionRequest(nil, chatReq)
+		}
+	})
+
+	chunk1 := &gemini.GenerateContentResponse{
+		ResponseID:   "resp-split-1",
+		ModelVersion: "gemini-3.5-flash",
+		Candidates: []*gemini.Candidate{{
+			Index:        0,
+			FinishReason: gemini.FinishReasonStop,
+			Content: &gemini.Content{
+				Role:  string(gemini.RoleModel),
+				Parts: []*gemini.Part{{Text: "Hello"}},
+			},
+		}},
+	}
+	chunk2 := &gemini.GenerateContentResponse{
+		ResponseID:   "resp-split-1",
+		ModelVersion: "gemini-3.5-flash",
+		Candidates:   []*gemini.Candidate{},
+		UsageMetadata: &gemini.GenerateContentResponseUsageMetadata{
+			PromptTokenCount:        2500,
+			CachedContentTokenCount: 2000,
+			CandidatesTokenCount:    10,
+			ThoughtsTokenCount:      120,
+			TotalTokenCount:         2630,
+		},
+	}
+
+	b.Run("SplitStreamChunks", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			state := gemini.NewGeminiStreamState()
+			_, _, _ = chunk1.ToBifrostChatCompletionStream(state)
+			_, _, _ = chunk2.ToBifrostChatCompletionStream(state)
+		}
+	})
+}

core/providers/gemini/passthrough_usage.go

@@ -190,6 +190,7 @@ func extractGeminiGenerateContentUsage(body []byte) *schemas.BifrostPassthroughU
 	}
 	if ru.OutputTokensDetails != nil {
 		usage.CompletionTokensDetails = &schemas.ChatCompletionTokensDetails{
+			TextTokens:      ru.OutputTokensDetails.TextTokens,
 			ReasoningTokens: ru.OutputTokensDetails.ReasoningTokens,
 			AudioTokens:     ru.OutputTokensDetails.AudioTokens,
 		}