From df37f3a095b0e3ee500e0bda75065b50147dddd7 Mon Sep 17 00:00:00 2001 From: roroghost17 Date: Tue, 30 Jun 2026 12:27:00 +0530 Subject: [PATCH] fix: forwards request id and trace id through telemetry --- docs/providers/request-options.mdx | 11 ++ .../bifrost-http/handlers/middlewares.go | 8 ++ .../bifrost-http/handlers/middlewares_test.go | 128 ++++++++++++++++++ 3 files changed, 147 insertions(+) diff --git a/docs/providers/request-options.mdx b/docs/providers/request-options.mdx index a5f2bec8e51..1df10f1470e 100644 --- a/docs/providers/request-options.mdx +++ b/docs/providers/request-options.mdx @@ -342,6 +342,17 @@ Input: messages, +#### Correlation response headers + +Every traced gateway request echoes the correlation IDs so you can pivot from a single request into its logs and trace in tools like Grafana (Loki + Tempo): + +| Response header | Description | +| --- | --- | +| `x-request-id` | The request ID you sent, or the UUID Bifrost generated when none was provided. | +| `x-bifrost-trace-id` | The trace ID linking all spans for the request. Inherited from an incoming W3C `traceparent` when present, otherwise generated. | + +The same `trace_id` and `request_id` appear as fields on Bifrost's structured access logs, so the value returned to the caller can be searched directly in your log store. + ### Send Back Raw Request **Context Key:** `BifrostContextKeySendBackRawRequest` diff --git a/transports/bifrost-http/handlers/middlewares.go b/transports/bifrost-http/handlers/middlewares.go index 2b8bdf212c0..c52c97a0516 100644 --- a/transports/bifrost-http/handlers/middlewares.go +++ b/transports/bifrost-http/handlers/middlewares.go @@ -150,6 +150,10 @@ func (c *CorsMiddleware) Middleware() schemas.BifrostHTTPMiddleware { if traceID, ok := ctx.UserValue(schemas.BifrostContextKeyTraceID).(string); ok && traceID != "" { logBuilder = logBuilder.Str("trace_id", traceID) } + // Emit the request ID alongside trace_id + if requestID := string(ctx.Request.Header.Peek("x-request-id")); requestID != "" { + logBuilder = logBuilder.Str("request_id", requestID) + } if cfg.dumpErrorsInConsoleLogs { if statusCode >= 400 && !ctx.Response.IsBodyStream() { if body := ctx.Response.Body(); len(body) > 0 { @@ -1211,6 +1215,10 @@ func (m *TracingMiddleware) Middleware() schemas.BifrostHTTPMiddleware { inheritedTraceID := tracing.ExtractParentID(&ctx.Request.Header) // Create trace in store - only ID returned (trace data stays in store) traceID := tracer.CreateTrace(inheritedTraceID, requestID) + // Surface correlation IDs back to the caller so a request can be pivoted + // into its logs (Loki) and trace (Tempo) in Grafana and similar stacks. + ctx.Response.Header.Set("x-request-id", requestID) + ctx.Response.Header.Set("x-bifrost-trace-id", traceID) // Store dimensions and session ID at the trace level (not as span // attributes) so connectors like BigQuery can export them without // changing the OTEL/Datadog span payloads. diff --git a/transports/bifrost-http/handlers/middlewares_test.go b/transports/bifrost-http/handlers/middlewares_test.go index c149a70b1e9..97df8a12705 100644 --- a/transports/bifrost-http/handlers/middlewares_test.go +++ b/transports/bifrost-http/handlers/middlewares_test.go @@ -2227,3 +2227,131 @@ func TestCollectDimensionHeaders(t *testing.T) { t.Errorf("collectDimensionHeaders(nil) = %v, want nil", got) } } + +// TestTracingMiddleware_SetsCorrelationHeaders asserts that every traced response +// carries x-request-id and x-bifrost-trace-id so callers can pivot a request into +// its logs and trace in Grafana/Tempo/Loki (BF-1041). +func TestTracingMiddleware_SetsCorrelationHeaders(t *testing.T) { + SetLogger(&mockLogger{}) + + store := tracing.NewTraceStore(5*time.Minute, nil) + defer store.Stop() + tracer := tracing.NewTracer(store, nil, nil) + defer tracer.Stop() + mw := NewTracingMiddleware(tracer).Middleware() + + newCtx := func() *fasthttp.RequestCtx { + ctx := &fasthttp.RequestCtx{} + ctx.Request.SetRequestURI("/openai/v1/chat/completions") + ctx.Request.Header.SetMethod("POST") + return ctx + } + + t.Run("generates request id when absent", func(t *testing.T) { + ctx := newCtx() + mw(func(*fasthttp.RequestCtx) {})(ctx) + + if got := string(ctx.Response.Header.Peek("x-bifrost-trace-id")); got == "" { + t.Error("expected x-bifrost-trace-id response header to be set") + } + if got := string(ctx.Response.Header.Peek("x-request-id")); got == "" { + t.Error("expected x-request-id response header to be set") + } + }) + + t.Run("echoes caller-supplied request id", func(t *testing.T) { + ctx := newCtx() + ctx.Request.Header.Set("x-request-id", "req-abc-123") + mw(func(*fasthttp.RequestCtx) {})(ctx) + + if got := string(ctx.Response.Header.Peek("x-request-id")); got != "req-abc-123" { + t.Errorf("x-request-id = %q, want req-abc-123", got) + } + if got := string(ctx.Response.Header.Peek("x-bifrost-trace-id")); got == "" { + t.Error("expected x-bifrost-trace-id response header to be set") + } + }) + + t.Run("headers survive the error path", func(t *testing.T) { + ctx := newCtx() + mw(func(c *fasthttp.RequestCtx) { + SendError(c, fasthttp.StatusBadGateway, "boom") + })(ctx) + + if ctx.Response.StatusCode() != fasthttp.StatusBadGateway { + t.Fatalf("status = %d, want %d", ctx.Response.StatusCode(), fasthttp.StatusBadGateway) + } + if got := string(ctx.Response.Header.Peek("x-bifrost-trace-id")); got == "" { + t.Error("expected x-bifrost-trace-id to survive the error path") + } + if got := string(ctx.Response.Header.Peek("x-request-id")); got == "" { + t.Error("expected x-request-id to survive the error path") + } + }) +} + +// captureLogEvent records the structured string fields emitted on the access log so +// a test can assert which correlation keys were written. +type captureLogEvent struct { + strFields map[string]string +} + +func (c *captureLogEvent) Str(key, val string) schemas.LogEventBuilder { + c.strFields[key] = val + return c +} +func (c *captureLogEvent) Int(string, int) schemas.LogEventBuilder { return c } +func (c *captureLogEvent) Int64(string, int64) schemas.LogEventBuilder { return c } +func (c *captureLogEvent) Send() {} + +type captureLogger struct { + mockLogger + events []*captureLogEvent +} + +func (l *captureLogger) LogHTTPRequest(schemas.LogLevel, string) schemas.LogEventBuilder { + e := &captureLogEvent{strFields: map[string]string{}} + l.events = append(l.events, e) + return e +} + +// TestTracingMiddleware_AccessLogIncludesRequestID asserts the stdout access log +// carries both trace_id and request_id, so Loki can index on either (BF-1041). +func TestTracingMiddleware_AccessLogIncludesRequestID(t *testing.T) { + logger := &captureLogger{} + SetLogger(logger) + defer SetLogger(&mockLogger{}) + + config := &lib.Config{ + ClientConfig: &configstore.ClientConfig{ + AllowedOrigins: []string{}, + }, + } + cors := NewCorsMiddleware(config).Middleware() + + store := tracing.NewTraceStore(5*time.Minute, nil) + defer store.Stop() + tracer := tracing.NewTracer(store, nil, nil) + defer tracer.Stop() + tm := NewTracingMiddleware(tracer).Middleware() + + ctx := &fasthttp.RequestCtx{} + ctx.Request.SetRequestURI("/openai/v1/chat/completions") + ctx.Request.Header.SetMethod("POST") + ctx.Request.Header.Set("x-request-id", "req-xyz") + + // CORS owns the access-log defer and wraps TracingMiddleware, so the trace_id + // UserValue and x-request-id header set by tracing are visible when it runs. + cors(tm(func(*fasthttp.RequestCtx) {}))(ctx) + + if len(logger.events) != 1 { + t.Fatalf("access log events = %d, want 1", len(logger.events)) + } + fields := logger.events[0].strFields + if got := fields["request_id"]; got != "req-xyz" { + t.Errorf("access log request_id = %q, want req-xyz", got) + } + if got := fields["trace_id"]; got == "" { + t.Error("expected access log to include a non-empty trace_id") + } +}