feat: add per-token latency tracking for streaming responses (#596)

## Summary

Add per-token latency tracking for streaming responses to improve observability and performance monitoring.

## Changes

- Added latency tracking for each token in streaming responses across all providers (Anthropic, Bedrock, Cohere, Gemini, OpenAI)
- Implemented two types of latency measurements:
  - Per-chunk latency: Time since the last chunk was received
  - Total latency: Time from the start of the request to the final chunk
- Added a new Prometheus metric `bifrost_stream_token_latency_seconds` to track token latency
- Enhanced the telemetry plugin to record these metrics for each streaming chunk
- Improved JSON marshaling for BifrostStream to prevent field conflicts
- Simplified the streaming response handler in HTTP transport

## Type of change

- [x] Feature
- [x] Refactor

## Affected areas

- [x] Core (Go)
- [x] Transports (HTTP)
- [x] Providers/Integrations
- [x] Plugins

## How to test

Test streaming responses with different providers and verify latency metrics are being recorded:

```sh
# Start Bifrost with telemetry plugin enabled
go run cmd/bifrost/main.go

# Make streaming requests to different providers
curl -X POST http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"model":"gpt-3.5-turbo","messages":[{"role":"user","content":"Write a short story"}],"stream":true}'

# Check Prometheus metrics
curl http://localhost:8000/metrics | grep bifrost_stream_token_latency_seconds
```

## Breaking changes

- [x] No

## Related issues

Improves observability for streaming responses, which helps diagnose performance issues.

## Security considerations

No security implications as this only adds internal performance tracking.

## Checklist

- [x] I added/updated tests where appropriate
- [x] I verified builds succeed (Go and UI)

Author: akshaydeo

core/changelog.md

@@ -1,4 +1,8 @@
 <!-- The pattern we follow here is to keep the changelog for the latest version -->
 <!-- Old changelogs are automatically attached to the GitHub releases -->
 
-- Fix: Anthropic tool results aggregation logic.
+- Feat: Stream token latency sent back in extra fields.
+- Feat: Plugin interface extended with TransportInterceptor method.
+- Feat: Add Anthropic thinking parameter
+- Feat: Add Custom key selector logic and send back request latency in extra fields.
+- Bug: Fallbacks not working occasionally.

core/providers/anthropic.go

@@ -444,6 +444,9 @@ func handleAnthropicStreaming(
 		scanner := bufio.NewScanner(resp.Body)
 		chunkIndex := -1
 
+		startTime := time.Now()
+		lastChunkTime := startTime
+
 		// Track minimal state needed for response format
 		var messageID string
 		var modelName string
@@ -527,8 +530,10 @@ func handleAnthropicStreaming(
 								Provider:       providerType,
 								ModelRequested: modelName,
 								ChunkIndex:     chunkIndex,
+								Latency:        time.Since(lastChunkTime).Milliseconds(),
 							},
 						}
+						lastChunkTime = time.Now()
 
 						// Use utility function to process and send response
 						processAndSendResponse(ctx, postHookRunner, streamResponse, responseChan, logger)
@@ -573,8 +578,10 @@ func handleAnthropicStreaming(
 									Provider:       providerType,
 									ModelRequested: modelName,
 									ChunkIndex:     chunkIndex,
+									Latency:        time.Since(lastChunkTime).Milliseconds(),
 								},
 							}
+							lastChunkTime = time.Now()
 
 							// Use utility function to process and send response
 							processAndSendResponse(ctx, postHookRunner, streamResponse, responseChan, logger)
@@ -610,8 +617,10 @@ func handleAnthropicStreaming(
 								Provider:       providerType,
 								ModelRequested: modelName,
 								ChunkIndex:     chunkIndex,
+								Latency:        time.Since(lastChunkTime).Milliseconds(),
 							},
 						}
+						lastChunkTime = time.Now()
 
 						// Use utility function to process and send response
 						processAndSendResponse(ctx, postHookRunner, streamResponse, responseChan, logger)
@@ -646,8 +655,10 @@ func handleAnthropicStreaming(
 									Provider:       providerType,
 									ModelRequested: modelName,
 									ChunkIndex:     chunkIndex,
+									Latency:        time.Since(lastChunkTime).Milliseconds(),
 								},
 							}
+							lastChunkTime = time.Now()
 
 							// Use utility function to process and send response
 							processAndSendResponse(ctx, postHookRunner, streamResponse, responseChan, logger)
@@ -683,8 +694,10 @@ func handleAnthropicStreaming(
 									Provider:       providerType,
 									ModelRequested: modelName,
 									ChunkIndex:     chunkIndex,
+									Latency:        time.Since(lastChunkTime).Milliseconds(),
 								},
 							}
+							lastChunkTime = time.Now()
 
 							// Use utility function to process and send response
 							processAndSendResponse(ctx, postHookRunner, streamResponse, responseChan, logger)
@@ -713,8 +726,10 @@ func handleAnthropicStreaming(
 									Provider:       providerType,
 									ModelRequested: modelName,
 									ChunkIndex:     chunkIndex,
+									Latency:        time.Since(lastChunkTime).Milliseconds(),
 								},
 							}
+							lastChunkTime = time.Now()
 
 							// Use utility function to process and send response
 							processAndSendResponse(ctx, postHookRunner, streamResponse, responseChan, logger)
@@ -774,6 +789,7 @@ func handleAnthropicStreaming(
 			processAndSendError(ctx, postHookRunner, err, responseChan, schemas.ChatCompletionStreamRequest, providerType, modelName, logger)
 		} else {
 			response := createBifrostChatCompletionChunkResponse(messageID, usage, finishReason, chunkIndex, schemas.ChatCompletionStreamRequest, providerType, modelName)
+			response.ExtraFields.Latency = time.Since(startTime).Milliseconds()
 			handleStreamEndWithSuccess(ctx, response, postHookRunner, responseChan, logger)
 		}
 	}()

core/providers/bedrock.go

@@ -644,14 +644,16 @@ func (provider *BedrockProvider) ChatCompletionStream(ctx context.Context, postH
 		reader := bufio.NewReader(resp.Body)
 		buffer := make([]byte, 1024*1024) // 1MB buffer
 		var accumulator []byte            // Accumulate data across reads
+		startTime := time.Now()
+		lastChunkTime := startTime
 
 		for {
 			n, err := reader.Read(buffer)
 			if err != nil {
 				if err == io.EOF {
 					// Process any remaining data in the accumulator
 					if len(accumulator) > 0 {
-						_ = provider.processAWSEventStreamData(ctx, postHookRunner, accumulator, &messageID, &chunkIndex, &usage, &finishReason, request.Model, providerName, responseChan)
+						_ = provider.processAWSEventStreamData(ctx, postHookRunner, accumulator, &messageID, &chunkIndex, &usage, &finishReason, request.Model, providerName, responseChan, &lastChunkTime)
 					}
 					break
 				}
@@ -668,14 +670,15 @@ func (provider *BedrockProvider) ChatCompletionStream(ctx context.Context, postH
 			accumulator = append(accumulator, buffer[:n]...)
 
 			// Process the accumulated data and get the remaining unprocessed part
-			remaining := provider.processAWSEventStreamData(ctx, postHookRunner, accumulator, &messageID, &chunkIndex, &usage, &finishReason, request.Model, providerName, responseChan)
+			remaining := provider.processAWSEventStreamData(ctx, postHookRunner, accumulator, &messageID, &chunkIndex, &usage, &finishReason, request.Model, providerName, responseChan, &lastChunkTime)
 
 			// Reset accumulator with remaining data
 			accumulator = remaining
 		}
 
 		// Send final response
 		response := createBifrostChatCompletionChunkResponse(messageID, usage, finishReason, chunkIndex, schemas.ChatCompletionStreamRequest, providerName, request.Model)
+		response.ExtraFields.Latency = time.Since(startTime).Milliseconds()
 		handleStreamEndWithSuccess(ctx, response, postHookRunner, responseChan, provider.logger)
 	}()
 
@@ -695,6 +698,7 @@ func (provider *BedrockProvider) processAWSEventStreamData(
 	model string,
 	providerName schemas.ModelProvider,
 	responseChan chan *schemas.BifrostStream,
+	lastChunkTime *time.Time,
 ) []byte {
 	lastProcessed := 0
 	depth := 0
@@ -741,7 +745,7 @@ func (provider *BedrockProvider) processAWSEventStreamData(
 						bytes.Contains(jsonBytes, []byte(`metadata`))
 
 					if hasQuotes && hasRelevantContent {
-						provider.processEventBuffer(ctx, postHookRunner, jsonBytes, messageID, chunkIndex, usage, finishReason, model, providerName, responseChan)
+						provider.processEventBuffer(ctx, postHookRunner, jsonBytes, messageID, chunkIndex, usage, finishReason, model, providerName, responseChan, lastChunkTime)
 						lastProcessed = i + 1
 					}
 					objStart = -1
@@ -759,7 +763,7 @@ func (provider *BedrockProvider) processAWSEventStreamData(
 }
 
 // processEventBuffer processes AWS Event Stream JSON payloads using typed Bedrock stream events
-func (provider *BedrockProvider) processEventBuffer(ctx context.Context, postHookRunner schemas.PostHookRunner, eventBuffer []byte, messageID *string, chunkIndex *int, usage **schemas.LLMUsage, finishReason **string, model string, providerName schemas.ModelProvider, responseChan chan *schemas.BifrostStream) {
+func (provider *BedrockProvider) processEventBuffer(ctx context.Context, postHookRunner schemas.PostHookRunner, eventBuffer []byte, messageID *string, chunkIndex *int, usage **schemas.LLMUsage, finishReason **string, model string, providerName schemas.ModelProvider, responseChan chan *schemas.BifrostStream, lastChunkTime *time.Time) {
 	// Parse the JSON event into our typed structure
 	var streamEvent bedrock.BedrockStreamEvent
 	if err := sonic.Unmarshal(eventBuffer, &streamEvent); err != nil {
@@ -798,9 +802,11 @@ func (provider *BedrockProvider) processEventBuffer(ctx context.Context, postHoo
 				Provider:       providerName,
 				ModelRequested: model,
 				ChunkIndex:     *chunkIndex,
+				Latency:        time.Since(*lastChunkTime).Milliseconds(),
 			},
 		}
 
+		*lastChunkTime = time.Now()
 		processAndSendResponse(ctx, postHookRunner, streamResponse, responseChan, provider.logger)
 
 	case streamEvent.Start != nil && streamEvent.Start.ToolUse != nil:
@@ -838,9 +844,11 @@ func (provider *BedrockProvider) processEventBuffer(ctx context.Context, postHoo
 				Provider:       providerName,
 				ModelRequested: model,
 				ChunkIndex:     *chunkIndex,
+				Latency:        time.Since(*lastChunkTime).Milliseconds(),
 			},
 		}
 
+		*lastChunkTime = time.Now()
 		processAndSendResponse(ctx, postHookRunner, streamResponse, responseChan, provider.logger)
 
 	case streamEvent.ContentBlockIndex != nil && streamEvent.Delta != nil:
@@ -872,9 +880,11 @@ func (provider *BedrockProvider) processEventBuffer(ctx context.Context, postHoo
 						Provider:       providerName,
 						ModelRequested: model,
 						ChunkIndex:     *chunkIndex,
+						Latency:        time.Since(*lastChunkTime).Milliseconds(),
 					},
 				}
 
+				*lastChunkTime = time.Now()
 				processAndSendResponse(ctx, postHookRunner, streamResponse, responseChan, provider.logger)
 			}
 
@@ -909,9 +919,11 @@ func (provider *BedrockProvider) processEventBuffer(ctx context.Context, postHoo
 					Provider:       providerName,
 					ModelRequested: model,
 					ChunkIndex:     *chunkIndex,
+					Latency:        time.Since(*lastChunkTime).Milliseconds(),
 				},
 			}
 
+			*lastChunkTime = time.Now()
 			processAndSendResponse(ctx, postHookRunner, streamResponse, responseChan, provider.logger)
 		}
 

core/providers/cohere.go

@@ -436,6 +436,8 @@ func (provider *CohereProvider) ChatCompletionStream(ctx context.Context, postHo
 
 		scanner := bufio.NewScanner(resp.Body)
 		var responseID string
+		startTime := time.Now()
+		lastChunkTime := startTime
 
 		for scanner.Scan() {
 			line := scanner.Text()
@@ -487,8 +489,10 @@ func (provider *CohereProvider) ChatCompletionStream(ctx context.Context, postHo
 						Provider:       providerName,
 						ModelRequested: request.Model,
 						ChunkIndex:     chunkIndex,
+						Latency:        time.Since(lastChunkTime).Milliseconds(),
 					},
 				}
+				lastChunkTime = time.Now()
 
 				switch event.Type {
 				case cohere.StreamEventMessageStart:
@@ -569,6 +573,7 @@ func (provider *CohereProvider) ChatCompletionStream(ctx context.Context, postHo
 						}
 
 						ctx = context.WithValue(ctx, schemas.BifrostContextKeyStreamEndIndicator, true)
+						response.ExtraFields.Latency = time.Since(startTime).Milliseconds()
 					}
 
 				case cohere.StreamEventToolCallEnd, cohere.StreamEventContentEnd:

core/providers/gemini.go

@@ -363,6 +363,8 @@ func (provider *GeminiProvider) SpeechStream(ctx context.Context, postHookRunner
 		scanner.Buffer(buf, 1024*1024)  // Allow up to 1MB tokens
 		chunkIndex := -1
 		usage := &schemas.AudioLLMUsage{}
+		startTime := time.Now()
+		lastChunkTime := startTime
 
 		for scanner.Scan() {
 			line := scanner.Text()
@@ -452,8 +454,10 @@ func (provider *GeminiProvider) SpeechStream(ctx context.Context, postHookRunner
 						Provider:       providerName,
 						ModelRequested: request.Model,
 						ChunkIndex:     chunkIndex,
+						Latency:        time.Since(lastChunkTime).Milliseconds(),
 					},
 				}
+				lastChunkTime = time.Now()
 
 				// Process response through post-hooks and send to channel
 				processAndSendResponse(ctx, postHookRunner, response, responseChan, provider.logger)
@@ -475,9 +479,11 @@ func (provider *GeminiProvider) SpeechStream(ctx context.Context, postHookRunner
 					Provider:       providerName,
 					ModelRequested: request.Model,
 					ChunkIndex:     chunkIndex + 1,
+					Latency:        time.Since(startTime).Milliseconds(),
 				},
 			}
 
+			ctx = context.WithValue(ctx, schemas.BifrostContextKeyStreamEndIndicator, true)
 			handleStreamEndWithSuccess(ctx, response, postHookRunner, responseChan, provider.logger)
 		}
 	}()
@@ -629,6 +635,8 @@ func (provider *GeminiProvider) TranscriptionStream(ctx context.Context, postHoo
 		scanner := bufio.NewScanner(resp.Body)
 		chunkIndex := -1
 		usage := &schemas.TranscriptionUsage{}
+		startTime := time.Now()
+		lastChunkTime := startTime
 
 		var fullTranscriptionText string
 
@@ -727,8 +735,10 @@ func (provider *GeminiProvider) TranscriptionStream(ctx context.Context, postHoo
 						Provider:       providerName,
 						ModelRequested: request.Model,
 						ChunkIndex:     chunkIndex,
+						Latency:        time.Since(lastChunkTime).Milliseconds(),
 					},
 				}
+				lastChunkTime = time.Now()
 
 				// Process response through post-hooks and send to channel
 				processAndSendResponse(ctx, postHookRunner, response, responseChan, provider.logger)
@@ -756,9 +766,11 @@ func (provider *GeminiProvider) TranscriptionStream(ctx context.Context, postHoo
 					Provider:       providerName,
 					ModelRequested: request.Model,
 					ChunkIndex:     chunkIndex + 1,
+					Latency:        time.Since(startTime).Milliseconds(),
 				},
 			}
 
+			ctx = context.WithValue(ctx, schemas.BifrostContextKeyStreamEndIndicator, true)
 			handleStreamEndWithSuccess(ctx, response, postHookRunner, responseChan, provider.logger)
 		}
 	}()