feat: byok-observability for aibridge

Author: evgeniy-scherbina

bridge.go

@@ -203,6 +203,8 @@ func newInterceptionProcessor(p provider.Provider, cbs *circuitbreaker.ProviderC
 			Client:                string(client),
 			ClientSessionID:       sessionID,
 			CorrelatingToolCallID: interceptor.CorrelatingToolCallID(),
+			CredentialKind:        interceptor.CredentialKind(),
+			CredentialHint:        interceptor.CredentialHint(),
 		}); err != nil {
 			span.SetStatus(codes.Error, fmt.Sprintf("failed to record interception: %v", err))
 			logger.Warn(ctx, "failed to record interception", slog.Error(err))

intercept/chatcompletions/base.go

@@ -39,6 +39,8 @@ type interceptionBase struct {
 
 	recorder recorder.Recorder
 	mcpProxy mcp.ServerProxier
+
+	intercept.CredentialFields
 }
 
 func (i *interceptionBase) newCompletionsService() openai.ChatCompletionService {

intercept/credential.go

@@ -0,0 +1,39 @@
+package intercept
+
+// Credential kind constants for interception recording.
+const (
+	CredentialKindCentralized    = "centralized"
+	CredentialKindPersonalAPIKey = "personal_api_key"
+	CredentialKindSubscription   = "subscription"
+)
+
+// MaskCredential returns a masked version of a credential suitable for
+// audit logging. It preserves a short prefix and suffix separated by
+// an ellipsis. Returns an empty string for empty input.
+func MaskCredential(secret string) string {
+	if len(secret) <= 8 {
+		return "***"
+	}
+	return secret[:4] + "..." + secret[len(secret)-4:]
+}
+
+// CredentialFields is an embeddable helper that implements the
+// CredentialKind, CredentialHint, and SetCredential methods of the
+// Interceptor interface.
+type CredentialFields struct {
+	Kind string
+	Hint string
+}
+
+func (c *CredentialFields) CredentialKind() string {
+	return c.Kind
+}
+
+func (c *CredentialFields) CredentialHint() string {
+	return c.Hint
+}
+
+func (c *CredentialFields) SetCredential(kind, hint string) {
+	c.Kind = kind
+	c.Hint = hint
+}

intercept/interceptor.go

@@ -25,6 +25,15 @@ type Interceptor interface {
 	Streaming() bool
 	// TraceAttributes returns tracing attributes for this [Interceptor]
 	TraceAttributes(*http.Request) []attribute.KeyValue
+	// SetCredential sets the credential kind and hint for this
+	// interception. Called by the provider after construction.
+	SetCredential(kind, hint string)
+	// CredentialKind returns how the request was authenticated:
+	// centralized, personal_api_key, or subscription.
+	CredentialKind() string
+	// CredentialHint returns a masked credential identifier for audit
+	// purposes (e.g. sk-...abc1).
+	CredentialHint() string
 	// CorrelatingToolCallID returns the ID of a tool call result submitted
 	// in the request, if present. This is used to correlate the current
 	// interception back to the previous interception that issued those tool

intercept/messages/base.go

@@ -78,6 +78,8 @@ type interceptionBase struct {
 
 	recorder recorder.Recorder
 	mcpProxy mcp.ServerProxier
+
+	intercept.CredentialFields
 }
 
 func (i *interceptionBase) ID() uuid.UUID {

intercept/responses/base.go

@@ -48,6 +48,8 @@ type responsesInterceptionBase struct {
 
 	logger slog.Logger
 	tracer trace.Tracer
+
+	intercept.CredentialFields
 }
 
 func (i *responsesInterceptionBase) newResponsesService() responses.ResponseService {

provider/anthropic.go

@@ -123,14 +123,17 @@ func (p *Anthropic) CreateInterceptor(w http.ResponseWriter, r *http.Request, tr
 		// set BYOKBearerToken and clear the centralized key.
 		// When both are present, X-Api-Key takes priority to match
 		// claude-code behavior.
+		credKind := intercept.CredentialKindCentralized
 		authHeaderName := p.AuthHeader()
 		if apiKey := r.Header.Get("X-Api-Key"); apiKey != "" {
 			cfg.Key = apiKey
 			authHeaderName = "X-Api-Key"
+			credKind = intercept.CredentialKindPersonalAPIKey
 		} else if token := utils.ExtractBearerToken(r.Header.Get("Authorization")); token != "" {
 			cfg.BYOKBearerToken = token
 			cfg.Key = ""
 			authHeaderName = "Authorization"
+			credKind = intercept.CredentialKindPersonalAPIKey
 		}
 
 		var interceptor intercept.Interceptor
@@ -139,6 +142,12 @@ func (p *Anthropic) CreateInterceptor(w http.ResponseWriter, r *http.Request, tr
 		} else {
 			interceptor = messages.NewBlockingInterceptor(id, reqPayload, cfg, p.bedrockCfg, r.Header, authHeaderName, tracer)
 		}
+		// Determine the credential hint from the key actually used.
+		credHint := intercept.MaskCredential(cfg.Key)
+		if cfg.BYOKBearerToken != "" {
+			credHint = intercept.MaskCredential(cfg.BYOKBearerToken)
+		}
+		interceptor.SetCredential(credKind, credHint)
 		span.SetAttributes(interceptor.TraceAttributes(r)...)
 		return interceptor, nil
 	}

provider/chatgpt.go

@@ -145,6 +145,7 @@ func (p *ChatGPT) CreateInterceptor(w http.ResponseWriter, r *http.Request, trac
 		span.SetStatus(codes.Error, "unknown route: "+r.URL.Path)
 		return nil, UnknownRoute
 	}
+	interceptor.SetCredential(intercept.CredentialKindSubscription, intercept.MaskCredential(token))
 	span.SetAttributes(interceptor.TraceAttributes(r)...)
 	return interceptor, nil
 }

provider/copilot.go

@@ -175,6 +175,7 @@ func (p *Copilot) CreateInterceptor(_ http.ResponseWriter, r *http.Request, trac
 		return nil, UnknownRoute
 	}
 
+	interceptor.SetCredential(intercept.CredentialKindSubscription, intercept.MaskCredential(key))
 	span.SetAttributes(interceptor.TraceAttributes(r)...)
 	return interceptor, nil
 }

provider/openai.go

@@ -106,8 +106,10 @@ func (p *OpenAI) CreateInterceptor(w http.ResponseWriter, r *http.Request, trace
 	//
 	// In BYOK mode the user's credential is in Authorization. Replace
 	// the centralized key with it so it is forwarded upstream.
+	credKind := intercept.CredentialKindCentralized
 	if token := utils.ExtractBearerToken(r.Header.Get("Authorization")); token != "" {
 		cfg.Key = token
+		credKind = intercept.CredentialKindPersonalAPIKey
 	}
 
 	path := strings.TrimPrefix(r.URL.Path, p.RoutePrefix())
@@ -143,6 +145,7 @@ func (p *OpenAI) CreateInterceptor(w http.ResponseWriter, r *http.Request, trace
 		span.SetStatus(codes.Error, "unknown route: "+r.URL.Path)
 		return nil, UnknownRoute
 	}
+	interceptor.SetCredential(credKind, intercept.MaskCredential(cfg.Key))
 	span.SetAttributes(interceptor.TraceAttributes(r)...)
 	return interceptor, nil
 }