Skip to content

Commit f72a795

Browse files
evgeniy-scherbinassncferreira
evgeniy-scherbina
and
ssncferreira
authored
feat: byok-observability for aibridge (#239)
* feat: byok-observability for aibridge * refactor: re-org credential configuring * Update intercept/credential.go Co-authored-by: Susana Ferreira <susana@coder.com> * refactor: minor refactor * refactor: make secret is always called * refactor: add CredentialKind custom type * test: enhance tests * test: enhance tests * ci: make fmt * fix: revert MaskSecret changes * fix: add small comment to enum * feat: log credential info on interception errors * refactor: update CredentialKind enum * refactor: MaskSecret should be fixed-length * feat: keep secret length in logs * Update provider/anthropic_test.go Co-authored-by: Susana Ferreira <susana@coder.com> * refactor: minor change in MaskSecret * refactor: minor change in MaskSecret --------- Co-authored-by: Susana Ferreira <susana@coder.com>
1 parent a011104 commit f72a795

File tree

21 files changed

+189
-76
lines changed

21 files changed

+189
-76
lines changed

‎bridge.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -217,6 +217,7 @@ func newInterceptionProcessor(p provider.Provider, cbs *circuitbreaker.ProviderC
217217
asyncRecorder.WithClient(string(client))
218218
interceptor.Setup(logger, asyncRecorder, mcpProxy)
219219

220+
cred := interceptor.Credential()
220221
if err := rec.RecordInterception(ctx, &recorder.InterceptionRecord{
221222
ID: interceptor.ID().String(),
222223
InitiatorID: actor.ID,
@@ -228,6 +229,8 @@ func newInterceptionProcessor(p provider.Provider, cbs *circuitbreaker.ProviderC
228229
Client: string(client),
229230
ClientSessionID: sessionID,
230231
CorrelatingToolCallID: interceptor.CorrelatingToolCallID(),
232+
CredentialKind: string(cred.Kind),
233+
CredentialHint: cred.Hint,
231234
}); err != nil {
232235
span.SetStatus(codes.Error, fmt.Sprintf("failed to record interception: %v", err))
233236
logger.Warn(ctx, "failed to record interception", slog.Error(err))
@@ -242,6 +245,9 @@ func newInterceptionProcessor(p provider.Provider, cbs *circuitbreaker.ProviderC
242245
slog.F("interception_id", interceptor.ID()),
243246
slog.F("user_agent", r.UserAgent()),
244247
slog.F("streaming", interceptor.Streaming()),
248+
slog.F("credential_kind", string(cred.Kind)),
249+
slog.F("credential_hint", cred.Hint),
250+
slog.F("credential_length", cred.Length),
245251
)
246252

247253
log.Debug(ctx, "interception started")

‎intercept/chatcompletions/base.go‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,8 +38,9 @@ type interceptionBase struct {
3838
logger slog.Logger
3939
tracer trace.Tracer
4040

41-
recorder recorder.Recorder
42-
mcpProxy mcp.ServerProxier
41+
recorder recorder.Recorder
42+
mcpProxy mcp.ServerProxier
43+
credential intercept.CredentialInfo
4344
}
4445

4546
func (i *interceptionBase) newCompletionsService() openai.ChatCompletionService {
@@ -74,6 +75,10 @@ func (i *interceptionBase) ID() uuid.UUID {
7475
return i.id
7576
}
7677

78+
func (i *interceptionBase) Credential() intercept.CredentialInfo {
79+
return i.credential
80+
}
81+
7782
func (i *interceptionBase) Setup(logger slog.Logger, recorder recorder.Recorder, mcpProxy mcp.ServerProxier) {
7883
i.logger = logger
7984
i.recorder = recorder

‎intercept/chatcompletions/blocking.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ func NewBlockingInterceptor(
3636
clientHeaders http.Header,
3737
authHeaderName string,
3838
tracer trace.Tracer,
39+
cred intercept.CredentialInfo,
3940
) *BlockingInterception {
4041
return &BlockingInterception{interceptionBase: interceptionBase{
4142
id: id,
@@ -45,6 +46,7 @@ func NewBlockingInterceptor(
4546
clientHeaders: clientHeaders,
4647
authHeaderName: authHeaderName,
4748
tracer: tracer,
49+
credential: cred,
4850
}}
4951
}
5052

‎intercept/chatcompletions/streaming.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ func NewStreamingInterceptor(
4141
clientHeaders http.Header,
4242
authHeaderName string,
4343
tracer trace.Tracer,
44+
cred intercept.CredentialInfo,
4445
) *StreamingInterception {
4546
return &StreamingInterception{interceptionBase: interceptionBase{
4647
id: id,
@@ -50,6 +51,7 @@ func NewStreamingInterceptor(
5051
clientHeaders: clientHeaders,
5152
authHeaderName: authHeaderName,
5253
tracer: tracer,
54+
credential: cred,
5355
}}
5456
}
5557

‎intercept/chatcompletions/streaming_test.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"cdr.dev/slog/v3"
1010
"cdr.dev/slog/v3/sloggers/slogtest"
1111
"github.com/coder/aibridge/config"
12+
"github.com/coder/aibridge/intercept"
1213
"github.com/coder/aibridge/internal/testutil"
1314
"github.com/google/uuid"
1415
"github.com/openai/openai-go/v3"
@@ -86,7 +87,7 @@ func TestStreamingInterception_RelaysUpstreamErrorToClient(t *testing.T) {
8687
httpReq := httptest.NewRequest(http.MethodPost, "/chat/completions", nil)
8788

8889
tracer := otel.Tracer("test")
89-
interceptor := NewStreamingInterceptor(uuid.New(), req, config.ProviderOpenAI, cfg, httpReq.Header, "Authorization", tracer)
90+
interceptor := NewStreamingInterceptor(uuid.New(), req, config.ProviderOpenAI, cfg, httpReq.Header, "Authorization", tracer, intercept.CredentialInfo{})
9091

9192
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: false}).Leveled(slog.LevelDebug)
9293
interceptor.Setup(logger, &testutil.MockRecorder{}, nil)

‎intercept/credential.go‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
package intercept
2+
3+
import "github.com/coder/aibridge/utils"
4+
5+
// CredentialKind identifies how a request was authenticated.
6+
// Keep in sync with the credential_kind enum in coderd's database.
7+
type CredentialKind string
8+
9+
// Credential kind constants for interception recording.
10+
const (
11+
CredentialKindCentralized CredentialKind = "centralized"
12+
CredentialKindBYOK CredentialKind = "byok"
13+
)
14+
15+
// CredentialInfo holds credential metadata for an interception.
16+
type CredentialInfo struct {
17+
Kind CredentialKind
18+
Hint string
19+
Length int
20+
}
21+
22+
// NewCredentialInfo creates a CredentialInfo from a raw credential.
23+
// The credential is automatically masked before storage so that the
24+
// original secret is never retained.
25+
func NewCredentialInfo(kind CredentialKind, credential string) CredentialInfo {
26+
return CredentialInfo{
27+
Kind: kind,
28+
Hint: utils.MaskSecret(credential),
29+
Length: len(credential),
30+
}
31+
}

‎intercept/interceptor.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@ type Interceptor interface {
2525
Streaming() bool
2626
// TraceAttributes returns tracing attributes for this [Interceptor]
2727
TraceAttributes(*http.Request) []attribute.KeyValue
28+
// Credential returns the credential metadata for this interception.
29+
Credential() CredentialInfo
2830
// CorrelatingToolCallID returns the ID of a tool call result submitted
2931
// in the request, if present. This is used to correlate the current
3032
// interception back to the previous interception that issued those tool

‎intercept/messages/base.go‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -77,14 +77,19 @@ type interceptionBase struct {
7777
tracer trace.Tracer
7878
logger slog.Logger
7979

80-
recorder recorder.Recorder
81-
mcpProxy mcp.ServerProxier
80+
recorder recorder.Recorder
81+
mcpProxy mcp.ServerProxier
82+
credential intercept.CredentialInfo
8283
}
8384

8485
func (i *interceptionBase) ID() uuid.UUID {
8586
return i.id
8687
}
8788

89+
func (i *interceptionBase) Credential() intercept.CredentialInfo {
90+
return i.credential
91+
}
92+
8893
func (i *interceptionBase) Setup(logger slog.Logger, recorder recorder.Recorder, mcpProxy mcp.ServerProxier) {
8994
i.logger = logger
9095
i.recorder = recorder

‎intercept/messages/blocking.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ func NewBlockingInterceptor(
3737
clientHeaders http.Header,
3838
authHeaderName string,
3939
tracer trace.Tracer,
40+
cred intercept.CredentialInfo,
4041
) *BlockingInterception {
4142
return &BlockingInterception{interceptionBase: interceptionBase{
4243
id: id,
@@ -47,6 +48,7 @@ func NewBlockingInterceptor(
4748
clientHeaders: clientHeaders,
4849
authHeaderName: authHeaderName,
4950
tracer: tracer,
51+
credential: cred,
5052
}}
5153
}
5254

‎intercept/messages/streaming.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ func NewStreamingInterceptor(
4343
clientHeaders http.Header,
4444
authHeaderName string,
4545
tracer trace.Tracer,
46+
cred intercept.CredentialInfo,
4647
) *StreamingInterception {
4748
return &StreamingInterception{interceptionBase: interceptionBase{
4849
id: id,
@@ -53,6 +54,7 @@ func NewStreamingInterceptor(
5354
clientHeaders: clientHeaders,
5455
authHeaderName: authHeaderName,
5556
tracer: tracer,
57+
credential: cred,
5658
}}
5759
}
5860

0 commit comments

Comments
 (0)