Skip to content

Commit ec6ae45

Browse files
feat: byok-observability for aibridge
1 parent a011104 commit ec6ae45

File tree

12 files changed

+73
-11
lines changed

12 files changed

+73
-11
lines changed

bridge.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -228,6 +228,8 @@ func newInterceptionProcessor(p provider.Provider, cbs *circuitbreaker.ProviderC
228228
Client: string(client),
229229
ClientSessionID: sessionID,
230230
CorrelatingToolCallID: interceptor.CorrelatingToolCallID(),
231+
CredentialKind: interceptor.CredentialKind(),
232+
CredentialHint: interceptor.CredentialHint(),
231233
}); err != nil {
232234
span.SetStatus(codes.Error, fmt.Sprintf("failed to record interception: %v", err))
233235
logger.Warn(ctx, "failed to record interception", slog.Error(err))

intercept/chatcompletions/base.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,8 @@ type interceptionBase struct {
4040

4141
recorder recorder.Recorder
4242
mcpProxy mcp.ServerProxier
43+
44+
intercept.CredentialFields
4345
}
4446

4547
func (i *interceptionBase) newCompletionsService() openai.ChatCompletionService {

intercept/credential.go

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
package intercept
2+
3+
// Credential kind constants for interception recording.
4+
const (
5+
CredentialKindCentralized = "centralized"
6+
CredentialKindPersonalAPIKey = "personal_api_key"
7+
CredentialKindSubscription = "subscription"
8+
)
9+
10+
// CredentialFields is an embeddable helper that implements the
11+
// SetCredential, CredentialKind, and CredentialHint methods of the
12+
// Interceptor interface.
13+
type CredentialFields struct {
14+
Kind string
15+
Hint string
16+
}
17+
18+
func (c *CredentialFields) SetCredential(kind, hint string) {
19+
c.Kind = kind
20+
c.Hint = hint
21+
}
22+
23+
func (c *CredentialFields) CredentialKind() string {
24+
return c.Kind
25+
}
26+
27+
func (c *CredentialFields) CredentialHint() string {
28+
return c.Hint
29+
}

intercept/interceptor.go

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,15 @@ type Interceptor interface {
2525
Streaming() bool
2626
// TraceAttributes returns tracing attributes for this [Interceptor]
2727
TraceAttributes(*http.Request) []attribute.KeyValue
28+
// SetCredential sets the credential kind and hint for this
29+
// interception. Called by the provider after construction.
30+
SetCredential(kind, hint string)
31+
// CredentialKind returns how the request was authenticated:
32+
// centralized, personal_api_key, or subscription.
33+
CredentialKind() string
34+
// CredentialHint returns a masked credential identifier for audit
35+
// purposes (e.g. sk-a...abc1).
36+
CredentialHint() string
2837
// CorrelatingToolCallID returns the ID of a tool call result submitted
2938
// in the request, if present. This is used to correlate the current
3039
// interception back to the previous interception that issued those tool

intercept/messages/base.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,8 @@ type interceptionBase struct {
7979

8080
recorder recorder.Recorder
8181
mcpProxy mcp.ServerProxier
82+
83+
intercept.CredentialFields
8284
}
8385

8486
func (i *interceptionBase) ID() uuid.UUID {

intercept/responses/base.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,8 @@ type responsesInterceptionBase struct {
4949

5050
logger slog.Logger
5151
tracer trace.Tracer
52+
53+
intercept.CredentialFields
5254
}
5355

5456
func (i *responsesInterceptionBase) newResponsesService() responses.ResponseService {

provider/anthropic.go

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,14 +130,20 @@ func (p *Anthropic) CreateInterceptor(w http.ResponseWriter, r *http.Request, tr
130130
// set BYOKBearerToken and clear the centralized key.
131131
// When both are present, X-Api-Key takes priority to match
132132
// claude-code behavior.
133+
credKind := intercept.CredentialKindCentralized
134+
credHint := utils.MaskSecret(cfg.Key)
133135
authHeaderName := p.AuthHeader()
134136
if apiKey := r.Header.Get("X-Api-Key"); apiKey != "" {
135137
cfg.Key = apiKey
136138
authHeaderName = "X-Api-Key"
139+
credKind = intercept.CredentialKindPersonalAPIKey
140+
credHint = utils.MaskSecret(apiKey)
137141
} else if token := utils.ExtractBearerToken(r.Header.Get("Authorization")); token != "" {
138142
cfg.BYOKBearerToken = token
139143
cfg.Key = ""
140144
authHeaderName = "Authorization"
145+
credKind = intercept.CredentialKindSubscription
146+
credHint = utils.MaskSecret(token)
141147
}
142148

143149
var interceptor intercept.Interceptor
@@ -146,6 +152,7 @@ func (p *Anthropic) CreateInterceptor(w http.ResponseWriter, r *http.Request, tr
146152
} else {
147153
interceptor = messages.NewBlockingInterceptor(id, reqPayload, p.Name(), cfg, p.bedrockCfg, r.Header, authHeaderName, tracer)
148154
}
155+
interceptor.SetCredential(credKind, credHint)
149156
span.SetAttributes(interceptor.TraceAttributes(r)...)
150157
return interceptor, nil
151158
}

provider/copilot.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -182,6 +182,7 @@ func (p *Copilot) CreateInterceptor(_ http.ResponseWriter, r *http.Request, trac
182182
return nil, UnknownRoute
183183
}
184184

185+
interceptor.SetCredential(intercept.CredentialKindSubscription, utils.MaskSecret(key))
185186
span.SetAttributes(interceptor.TraceAttributes(r)...)
186187
return interceptor, nil
187188
}

provider/openai.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,8 +113,10 @@ func (p *OpenAI) CreateInterceptor(w http.ResponseWriter, r *http.Request, trace
113113
//
114114
// In BYOK mode the user's credential is in Authorization. Replace
115115
// the centralized key with it so it is forwarded upstream.
116+
credKind := intercept.CredentialKindCentralized
116117
if token := utils.ExtractBearerToken(r.Header.Get("Authorization")); token != "" {
117118
cfg.Key = token
119+
credKind = intercept.CredentialKindPersonalAPIKey
118120
}
119121

120122
path := strings.TrimPrefix(r.URL.Path, p.RoutePrefix())
@@ -150,6 +152,7 @@ func (p *OpenAI) CreateInterceptor(w http.ResponseWriter, r *http.Request, trace
150152
span.SetStatus(codes.Error, "unknown route: "+r.URL.Path)
151153
return nil, UnknownRoute
152154
}
155+
interceptor.SetCredential(credKind, utils.MaskSecret(cfg.Key))
153156
span.SetAttributes(interceptor.TraceAttributes(r)...)
154157
return interceptor, nil
155158
}

recorder/types.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,8 @@ type InterceptionRecord struct {
3939
Client string
4040
UserAgent string
4141
CorrelatingToolCallID *string
42+
CredentialKind string
43+
CredentialHint string
4244
}
4345

4446
type InterceptionRecordEnded struct {

0 commit comments

Comments
 (0)