mirror of
https://github.com/coder/coder.git
synced 2026-10-16 14:48:56 +08:00
refactor: use typed AI Gateway actor identity (#30258)
Adds typed AI Gateway actor identity, including the API key ID used by recording. Recorders now read the authenticated API key ID from each call context instead of capturing request identity at construction time. This allows one recorder to be reused safely across sequential and concurrent requests. Includes context helper coverage and recorder regression tests. Generated by Coder Agents on behalf of @pawbana.
This commit is contained in:
+10
-5
@@ -27,6 +27,9 @@ type (
|
||||
|
||||
Provider = provider.Provider
|
||||
|
||||
// Actor is the authenticated identity of an AI Gateway request.
|
||||
Actor = aibcontext.Actor
|
||||
|
||||
InterceptionRecord = recorder.InterceptionRecord
|
||||
InterceptionRecordEnded = recorder.InterceptionRecordEnded
|
||||
TokenUsageRecord = recorder.TokenUsageRecord
|
||||
@@ -44,8 +47,9 @@ type (
|
||||
CopilotConfig = config.Copilot
|
||||
)
|
||||
|
||||
func AsActor(ctx context.Context, actorID, email string, metadata recorder.Metadata) context.Context {
|
||||
return aibcontext.AsActor(ctx, actorID, email, metadata)
|
||||
// AsActor attaches the authenticated identity to an AI Gateway request context.
|
||||
func AsActor(ctx context.Context, actor Actor) context.Context {
|
||||
return aibcontext.AsActor(ctx, actor)
|
||||
}
|
||||
|
||||
// NewAnthropicProvider constructs the Anthropic provider. At most one of
|
||||
@@ -80,7 +84,8 @@ func NewMetrics(reg prometheus.Registerer) *metrics.Metrics {
|
||||
}
|
||||
|
||||
// NewRecorder creates a [Recorder] which logs each record and refuses
|
||||
// malformed ones before handing it to base.
|
||||
// malformed ones before handing it to base. Interception start records require
|
||||
// an actor with an API-key ID on their context.
|
||||
//
|
||||
// middleware is inserted below the logging and validating middleware, so that
|
||||
// every record is logged and checked before any of it runs, and above the
|
||||
@@ -88,9 +93,9 @@ func NewMetrics(reg prometheus.Registerer) *metrics.Metrics {
|
||||
// measure. Policy that drops records, such as [recorder.WithoutRecords],
|
||||
// belongs here: it keeps NewRecorder to its own concerns and leaves the choice
|
||||
// to the caller.
|
||||
func NewRecorder(logger slog.Logger, tracer trace.Tracer, apiKeyID string, structured bool, base Recorder, middleware ...recorder.Middleware) Recorder {
|
||||
func NewRecorder(logger slog.Logger, tracer trace.Tracer, structured bool, base Recorder, middleware ...recorder.Middleware) Recorder {
|
||||
chain := make([]recorder.Middleware, 0, len(middleware)+3)
|
||||
chain = append(chain, recorder.WithLogging(logger, apiKeyID, structured))
|
||||
chain = append(chain, recorder.WithLogging(logger, structured))
|
||||
chain = append(chain, recorder.WithValidation(logger))
|
||||
chain = append(chain, middleware...)
|
||||
chain = append(chain, recorder.WithTracing(tracer))
|
||||
|
||||
+10
-5
@@ -200,6 +200,11 @@ func newInterceptionProcessor(p provider.Provider, cbs *circuitbreaker.ProviderC
|
||||
return
|
||||
}
|
||||
|
||||
actorID := actor.ID.String()
|
||||
var metadata recorder.Metadata
|
||||
if actor.Username != "" {
|
||||
metadata = recorder.Metadata{"Username": actor.Username}
|
||||
}
|
||||
cred := interceptor.Credential()
|
||||
traceAttrs := interceptor.TraceAttributes(r)
|
||||
span.SetAttributes(traceAttrs...)
|
||||
@@ -217,15 +222,15 @@ func newInterceptionProcessor(p provider.Provider, cbs *circuitbreaker.ProviderC
|
||||
asyncRecorder.WithMetrics(m)
|
||||
asyncRecorder.WithProvider(p.Name())
|
||||
asyncRecorder.WithModel(interceptor.Model())
|
||||
asyncRecorder.WithInitiatorID(actor.ID)
|
||||
asyncRecorder.WithInitiatorID(actorID)
|
||||
asyncRecorder.WithClient(string(client))
|
||||
interceptor.Setup(logger, asyncRecorder, mcpProxy)
|
||||
|
||||
if err := rec.RecordInterception(ctx, &recorder.InterceptionRecord{
|
||||
StartedAt: time.Now().UTC(),
|
||||
ID: interceptor.ID().String(),
|
||||
InitiatorID: actor.ID,
|
||||
Metadata: actor.Metadata,
|
||||
InitiatorID: actorID,
|
||||
Metadata: metadata,
|
||||
Model: interceptor.Model(),
|
||||
Provider: p.Type(),
|
||||
ProviderName: p.Name(),
|
||||
@@ -273,13 +278,13 @@ func newInterceptionProcessor(p provider.Provider, cbs *circuitbreaker.ProviderC
|
||||
errType, errMsg := interceptionerror.Categorize(p, execErr)
|
||||
if execErr != nil {
|
||||
if m != nil {
|
||||
m.InterceptionCount.WithLabelValues(p.Name(), interceptor.Model(), metrics.InterceptionCountStatusFailed, route, r.Method, actor.ID, string(client)).Add(1)
|
||||
m.InterceptionCount.WithLabelValues(p.Name(), interceptor.Model(), metrics.InterceptionCountStatusFailed, route, r.Method, actorID, string(client)).Add(1)
|
||||
}
|
||||
span.SetStatus(codes.Error, fmt.Sprintf("interception failed: %v", execErr))
|
||||
log.Warn(credCtx, "interception failed", slog.Error(execErr), slog.F("error_type", string(errType)))
|
||||
} else {
|
||||
if m != nil {
|
||||
m.InterceptionCount.WithLabelValues(p.Name(), interceptor.Model(), metrics.InterceptionCountStatusCompleted, route, r.Method, actor.ID, string(client)).Add(1)
|
||||
m.InterceptionCount.WithLabelValues(p.Name(), interceptor.Model(), metrics.InterceptionCountStatusCompleted, route, r.Method, actorID, string(client)).Add(1)
|
||||
}
|
||||
log.Debug(credCtx, "interception ended")
|
||||
}
|
||||
|
||||
+28
-10
@@ -3,25 +3,33 @@ package context
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/coder/coder/v2/aibridge/recorder"
|
||||
"github.com/google/uuid"
|
||||
"golang.org/x/xerrors"
|
||||
)
|
||||
|
||||
type (
|
||||
actorContextKey struct{}
|
||||
)
|
||||
|
||||
// ErrMissingAPIKeyID indicates that the context has no actor with an API-key ID.
|
||||
var ErrMissingAPIKeyID = xerrors.New("missing authenticated API-key ID")
|
||||
|
||||
// Actor is the authenticated identity attached to an AI Gateway request.
|
||||
type Actor struct {
|
||||
ID string
|
||||
// Email is kept out of Metadata so it is never recorded or logged. It is
|
||||
// only read when forwarding actor headers upstream.
|
||||
Email string
|
||||
Metadata recorder.Metadata
|
||||
ID uuid.UUID
|
||||
APIKeyID string
|
||||
Username string
|
||||
// Email is only used for the opt-in email actor header forwarding. It is kept
|
||||
// out of interception recordings.
|
||||
Email string
|
||||
}
|
||||
|
||||
func AsActor(ctx context.Context, actorID, email string, metadata recorder.Metadata) context.Context {
|
||||
return context.WithValue(ctx, actorContextKey{}, &Actor{ID: actorID, Email: email, Metadata: metadata})
|
||||
// AsActor returns a context containing a copy of the authenticated actor.
|
||||
func AsActor(ctx context.Context, actor Actor) context.Context {
|
||||
return context.WithValue(ctx, actorContextKey{}, &actor)
|
||||
}
|
||||
|
||||
// ActorFromContext returns the authenticated actor, or nil when absent.
|
||||
func ActorFromContext(ctx context.Context) *Actor {
|
||||
a, ok := ctx.Value(actorContextKey{}).(*Actor)
|
||||
if !ok {
|
||||
@@ -34,8 +42,18 @@ func ActorFromContext(ctx context.Context) *Actor {
|
||||
// ActorIDFromContext safely extracts the actor ID from the context.
|
||||
// Returns an empty string if no actor is found.
|
||||
func ActorIDFromContext(ctx context.Context) string {
|
||||
if actor := ActorFromContext(ctx); actor != nil {
|
||||
return actor.ID
|
||||
if actor := ActorFromContext(ctx); actor != nil && actor.ID != uuid.Nil {
|
||||
return actor.ID.String()
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// APIKeyIDFromContext returns the actor's API-key ID, or ErrMissingAPIKeyID if
|
||||
// the actor is absent or its API-key ID is empty.
|
||||
func APIKeyIDFromContext(ctx context.Context) (string, error) {
|
||||
actor := ActorFromContext(ctx)
|
||||
if actor == nil || actor.APIKeyID == "" {
|
||||
return "", ErrMissingAPIKeyID
|
||||
}
|
||||
return actor.APIKeyID, nil
|
||||
}
|
||||
|
||||
@@ -4,28 +4,30 @@ import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
aibcontext "github.com/coder/coder/v2/aibridge/context"
|
||||
"github.com/coder/coder/v2/aibridge/recorder"
|
||||
)
|
||||
|
||||
func TestAsActor(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// Given: a metadata map
|
||||
metadata := recorder.Metadata{"key": "value"}
|
||||
|
||||
// When: storing an actor in the context
|
||||
ctx := aibcontext.AsActor(context.Background(), "actor-123", "actor@example.com", metadata)
|
||||
|
||||
// Then: the actor should be retrievable with correct ID, email, and metadata
|
||||
want := aibcontext.Actor{
|
||||
ID: uuid.New(),
|
||||
APIKeyID: "test-key-id",
|
||||
Username: "actor",
|
||||
Email: "actor@example.com",
|
||||
}
|
||||
ctx := aibcontext.AsActor(t.Context(), want)
|
||||
actor := aibcontext.ActorFromContext(ctx)
|
||||
require.NotNil(t, actor)
|
||||
assert.Equal(t, "actor-123", actor.ID)
|
||||
assert.Equal(t, "actor@example.com", actor.Email)
|
||||
assert.Equal(t, "value", actor.Metadata["key"])
|
||||
assert.Equal(t, want, *actor)
|
||||
|
||||
want.Username = "changed"
|
||||
assert.Equal(t, "actor", actor.Username, "the context owns a copy of the actor")
|
||||
assert.Same(t, actor, aibcontext.ActorFromContext(context.WithoutCancel(ctx)))
|
||||
}
|
||||
|
||||
func TestActorFromContext(t *testing.T) {
|
||||
@@ -35,14 +37,15 @@ func TestActorFromContext(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// Given: a context with an actor
|
||||
ctx := aibcontext.AsActor(context.Background(), "test-id", "", recorder.Metadata{})
|
||||
id := uuid.New()
|
||||
ctx := aibcontext.AsActor(t.Context(), aibcontext.Actor{ID: id})
|
||||
|
||||
// When: extracting the actor from context
|
||||
actor := aibcontext.ActorFromContext(ctx)
|
||||
|
||||
// Then: the actor should be returned with correct ID
|
||||
require.NotNil(t, actor)
|
||||
assert.Equal(t, "test-id", actor.ID)
|
||||
assert.Equal(t, id, actor.ID)
|
||||
})
|
||||
|
||||
t.Run("returns nil when no actor", func(t *testing.T) {
|
||||
@@ -59,6 +62,52 @@ func TestActorFromContext(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestAPIKeyIDFromContext(t *testing.T) {
|
||||
t.Parallel()
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
actor *aibcontext.Actor
|
||||
detached bool
|
||||
want string
|
||||
wantErr error
|
||||
}{
|
||||
{
|
||||
name: "Present",
|
||||
actor: &aibcontext.Actor{APIKeyID: "test-key-id"},
|
||||
want: "test-key-id",
|
||||
},
|
||||
{
|
||||
name: "MissingActor",
|
||||
wantErr: aibcontext.ErrMissingAPIKeyID,
|
||||
},
|
||||
{
|
||||
name: "EmptyAPIKeyID",
|
||||
actor: &aibcontext.Actor{},
|
||||
wantErr: aibcontext.ErrMissingAPIKeyID,
|
||||
},
|
||||
{
|
||||
name: "DetachedContext",
|
||||
actor: &aibcontext.Actor{APIKeyID: "test-key-id"},
|
||||
detached: true,
|
||||
want: "test-key-id",
|
||||
},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := t.Context()
|
||||
if tc.actor != nil {
|
||||
ctx = aibcontext.AsActor(ctx, *tc.actor)
|
||||
}
|
||||
if tc.detached {
|
||||
ctx = context.WithoutCancel(ctx)
|
||||
}
|
||||
got, err := aibcontext.APIKeyIDFromContext(ctx)
|
||||
require.ErrorIs(t, err, tc.wantErr)
|
||||
assert.Equal(t, tc.want, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestActorIDFromContext(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
@@ -66,13 +115,20 @@ func TestActorIDFromContext(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// Given: a context with an actor
|
||||
ctx := aibcontext.AsActor(context.Background(), "test-actor-id", "", recorder.Metadata{})
|
||||
id := uuid.New()
|
||||
ctx := aibcontext.AsActor(t.Context(), aibcontext.Actor{ID: id})
|
||||
|
||||
// When: extracting the actor ID from context
|
||||
got := aibcontext.ActorIDFromContext(ctx)
|
||||
|
||||
// Then: the actor ID should be returned
|
||||
assert.Equal(t, "test-actor-id", got)
|
||||
assert.Equal(t, id.String(), got)
|
||||
})
|
||||
|
||||
t.Run("returns empty string for nil ID", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := aibcontext.AsActor(t.Context(), aibcontext.Actor{})
|
||||
assert.Empty(t, aibcontext.ActorIDFromContext(ctx))
|
||||
})
|
||||
|
||||
t.Run("returns empty string when no actor", func(t *testing.T) {
|
||||
|
||||
@@ -107,7 +107,8 @@ func IsActorHeader(name string) bool {
|
||||
}
|
||||
|
||||
// headersFromActor maps supported actor attributes to configured header names.
|
||||
// Attributes with no destination and unknown mapping keys are ignored.
|
||||
// Attributes with no destination, empty values (including a nil ID), and
|
||||
// unknown mapping keys are ignored. The API key ID is never forwarded.
|
||||
// No headers are returned if actor is nil.
|
||||
func headersFromActor(actor *aibcontext.Actor, actorHeaderNames map[string]string) map[string]string {
|
||||
if actor == nil {
|
||||
@@ -115,13 +116,11 @@ func headersFromActor(actor *aibcontext.Actor, actorHeaderNames map[string]strin
|
||||
}
|
||||
|
||||
headers := make(map[string]string, len(actorHeaderNames))
|
||||
if name := actorHeaderNames[ActorAttributeID]; name != "" {
|
||||
headers[name] = actor.ID
|
||||
if name := actorHeaderNames[ActorAttributeID]; name != "" && actor.ID != uuid.Nil {
|
||||
headers[name] = actor.ID.String()
|
||||
}
|
||||
if name := actorHeaderNames[ActorAttributeUsername]; name != "" {
|
||||
if username, ok := actor.Metadata["Username"].(string); ok && username != "" {
|
||||
headers[name] = username
|
||||
}
|
||||
if name := actorHeaderNames[ActorAttributeUsername]; name != "" && actor.Username != "" {
|
||||
headers[name] = actor.Username
|
||||
}
|
||||
if name := actorHeaderNames[ActorAttributeEmail]; name != "" && actor.Email != "" {
|
||||
headers[name] = actor.Email
|
||||
|
||||
@@ -3,15 +3,17 @@ package headers
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/coder/coder/v2/aibridge/context"
|
||||
"github.com/coder/coder/v2/aibridge/recorder"
|
||||
)
|
||||
|
||||
func TestHeadersFromActor(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
actorID := uuid.MustParse("6f1d2c3b-4a5e-4f60-8a7b-9c0d1e2f3a4b")
|
||||
const apiKeyID = "api-key-id"
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
actor *context.Actor
|
||||
@@ -21,95 +23,93 @@ func TestHeadersFromActor(t *testing.T) {
|
||||
{name: "nil actor", names: map[string]string{"id": ActorIDHeader}},
|
||||
{
|
||||
name: "id only",
|
||||
actor: &context.Actor{ID: "user-123"},
|
||||
actor: &context.Actor{ID: actorID},
|
||||
names: map[string]string{"id": "X-Downstream-User-Id"},
|
||||
want: map[string]string{"X-Downstream-User-Id": "user-123"},
|
||||
want: map[string]string{"X-Downstream-User-Id": actorID.String()},
|
||||
},
|
||||
{
|
||||
name: "configured supported attributes only",
|
||||
actor: &context.Actor{
|
||||
ID: "user-123",
|
||||
Email: "alice@example.com",
|
||||
Metadata: recorder.Metadata{
|
||||
"Username": "alice",
|
||||
"Count": 42,
|
||||
"Role": "admin",
|
||||
},
|
||||
ID: actorID,
|
||||
APIKeyID: apiKeyID,
|
||||
Username: "alice",
|
||||
Email: "alice@example.com",
|
||||
},
|
||||
names: map[string]string{
|
||||
"id": "X-Downstream-User-Id",
|
||||
"username": "X-Downstream-Username",
|
||||
"email": "X-Downstream-Email",
|
||||
"Role": "X-Downstream-Role",
|
||||
"id": "X-Downstream-User-Id",
|
||||
"username": "X-Downstream-Username",
|
||||
"email": "X-Downstream-Email",
|
||||
"Role": "X-Downstream-Role",
|
||||
"api_key_id": "X-Downstream-Api-Key-Id",
|
||||
"APIKeyID": "X-Downstream-Api-Key",
|
||||
},
|
||||
want: map[string]string{
|
||||
"X-Downstream-User-Id": "user-123",
|
||||
"X-Downstream-User-Id": actorID.String(),
|
||||
"X-Downstream-Username": "alice",
|
||||
"X-Downstream-Email": "alice@example.com",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "id omitted",
|
||||
actor: &context.Actor{ID: "user-123", Email: "alice@example.com", Metadata: recorder.Metadata{"Username": "alice"}},
|
||||
actor: &context.Actor{ID: actorID, Email: "alice@example.com", Username: "alice"},
|
||||
names: map[string]string{"username": "X-Downstream-Username", "email": "X-Downstream-Email"},
|
||||
want: map[string]string{"X-Downstream-Username": "alice", "X-Downstream-Email": "alice@example.com"},
|
||||
},
|
||||
{
|
||||
name: "username omitted",
|
||||
actor: &context.Actor{ID: "user-123", Email: "alice@example.com", Metadata: recorder.Metadata{"Username": "alice"}},
|
||||
actor: &context.Actor{ID: actorID, Email: "alice@example.com", Username: "alice"},
|
||||
names: map[string]string{"id": "X-Downstream-User-Id", "email": "X-Downstream-Email"},
|
||||
want: map[string]string{"X-Downstream-User-Id": "user-123", "X-Downstream-Email": "alice@example.com"},
|
||||
want: map[string]string{"X-Downstream-User-Id": actorID.String(), "X-Downstream-Email": "alice@example.com"},
|
||||
},
|
||||
{
|
||||
name: "email omitted",
|
||||
actor: &context.Actor{ID: "user-123", Email: "alice@example.com", Metadata: recorder.Metadata{"Username": "alice"}},
|
||||
actor: &context.Actor{ID: actorID, Email: "alice@example.com", Username: "alice"},
|
||||
names: map[string]string{"id": "X-Downstream-User-Id", "username": "X-Downstream-Username"},
|
||||
want: map[string]string{"X-Downstream-User-Id": "user-123", "X-Downstream-Username": "alice"},
|
||||
want: map[string]string{"X-Downstream-User-Id": actorID.String(), "X-Downstream-Username": "alice"},
|
||||
},
|
||||
{
|
||||
name: "missing username",
|
||||
actor: &context.Actor{ID: "user-123"},
|
||||
actor: &context.Actor{ID: actorID},
|
||||
names: map[string]string{"username": "X-Username"},
|
||||
want: map[string]string{},
|
||||
},
|
||||
{
|
||||
name: "empty username",
|
||||
actor: &context.Actor{Metadata: recorder.Metadata{"Username": ""}},
|
||||
names: map[string]string{"username": "X-Username"},
|
||||
name: "nil id",
|
||||
actor: &context.Actor{Username: "alice"},
|
||||
names: map[string]string{"id": "X-User-Id"},
|
||||
want: map[string]string{},
|
||||
},
|
||||
{
|
||||
name: "empty email",
|
||||
actor: &context.Actor{ID: "user-123"},
|
||||
actor: &context.Actor{ID: actorID},
|
||||
names: map[string]string{"email": "X-Email"},
|
||||
want: map[string]string{},
|
||||
},
|
||||
{
|
||||
name: "non-string username",
|
||||
actor: &context.Actor{Metadata: recorder.Metadata{"Username": 42}},
|
||||
names: map[string]string{"username": "X-Username"},
|
||||
name: "unknown attribute keys forward nothing",
|
||||
actor: &context.Actor{ID: actorID, APIKeyID: apiKeyID, Username: "alice", Email: "alice@example.com"},
|
||||
names: map[string]string{"Username": "X-Username", "ID": "X-Id", "api_key_id": "X-Api-Key-Id"},
|
||||
want: map[string]string{},
|
||||
},
|
||||
{
|
||||
name: "empty map forwards no actor headers",
|
||||
actor: &context.Actor{
|
||||
ID: "user-123",
|
||||
Metadata: recorder.Metadata{
|
||||
"Username": "alice",
|
||||
},
|
||||
},
|
||||
name: "empty map forwards no actor headers",
|
||||
actor: &context.Actor{ID: actorID, Username: "alice"},
|
||||
names: map[string]string{},
|
||||
want: map[string]string{},
|
||||
},
|
||||
{
|
||||
name: "nil map forwards no actor headers",
|
||||
actor: &context.Actor{ID: "user-123"},
|
||||
actor: &context.Actor{ID: actorID},
|
||||
want: map[string]string{},
|
||||
},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
require.Equal(t, tc.want, headersFromActor(tc.actor, tc.names))
|
||||
got := headersFromActor(tc.actor, tc.names)
|
||||
require.Equal(t, tc.want, got)
|
||||
for name, value := range got {
|
||||
require.NotEqual(t, apiKeyID, value, "the API key ID must never be forwarded in %s", name)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4,15 +4,19 @@ import (
|
||||
"net/http"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/coder/coder/v2/aibridge/context"
|
||||
"github.com/coder/coder/v2/aibridge/headers"
|
||||
"github.com/coder/coder/v2/aibridge/recorder"
|
||||
agplaibridge "github.com/coder/coder/v2/coderd/aibridge"
|
||||
)
|
||||
|
||||
// actorID is the authenticated actor in tests. Client and SDK values use
|
||||
// distinct strings so a forwarded spoofed value cannot match it.
|
||||
var actorID = uuid.MustParse("6f1d2c3b-4a5e-4f60-8a7b-9c0d1e2f3a4b")
|
||||
|
||||
func TestActorHeaders(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
@@ -539,14 +543,14 @@ func TestBuildUpstreamHeaders(t *testing.T) {
|
||||
names: map[string]string{"id": headers.ActorIDHeader},
|
||||
want: http.Header{
|
||||
"Authorization": {"Bearer provider-key"},
|
||||
"X-Ai-Bridge-Actor-Id": {"user-123"},
|
||||
"X-Ai-Bridge-Actor-Id": {actorID.String()},
|
||||
"X-Unrelated": {"preserved"},
|
||||
},
|
||||
},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
actor := &context.Actor{ID: "user-123", Metadata: recorder.Metadata{"Username": "alice"}}
|
||||
actor := &context.Actor{ID: actorID, APIKeyID: "api-key-id", Username: "alice"}
|
||||
result := headers.BuildUpstreamHeaders(sdkHeaders, clientHeaders, "Authorization", tc.names, actor)
|
||||
require.Equal(t, tc.want, result)
|
||||
require.Equal(t, sdkCopy, sdkHeaders)
|
||||
@@ -565,18 +569,23 @@ func TestBuildUpstreamHeaders(t *testing.T) {
|
||||
}{
|
||||
{
|
||||
name: "configured actor",
|
||||
actor: &context.Actor{ID: "user-123", Email: "alice@example.com", Metadata: recorder.Metadata{"Username": "alice"}},
|
||||
want: http.Header{"X-Downstream-User-Id": {"user-123"}, "X-Downstream-Username": {"alice"}, "X-Downstream-Email": {"alice@example.com"}},
|
||||
actor: &context.Actor{ID: actorID, APIKeyID: "api-key-id", Username: "alice", Email: "alice@example.com"},
|
||||
want: http.Header{"X-Downstream-User-Id": {actorID.String()}, "X-Downstream-Username": {"alice"}, "X-Downstream-Email": {"alice@example.com"}},
|
||||
},
|
||||
{
|
||||
name: "missing email",
|
||||
actor: &context.Actor{ID: "user-123", Metadata: recorder.Metadata{"Username": "alice"}},
|
||||
want: http.Header{"X-Downstream-User-Id": {"user-123"}, "X-Downstream-Username": {"alice"}},
|
||||
actor: &context.Actor{ID: actorID, Username: "alice"},
|
||||
want: http.Header{"X-Downstream-User-Id": {actorID.String()}, "X-Downstream-Username": {"alice"}},
|
||||
},
|
||||
{
|
||||
name: "empty username",
|
||||
actor: &context.Actor{ID: "user-123", Metadata: recorder.Metadata{"Username": ""}},
|
||||
want: http.Header{"X-Downstream-User-Id": {"user-123"}},
|
||||
actor: &context.Actor{ID: actorID},
|
||||
want: http.Header{"X-Downstream-User-Id": {actorID.String()}},
|
||||
},
|
||||
{
|
||||
name: "nil id",
|
||||
actor: &context.Actor{Username: "alice"},
|
||||
want: http.Header{"X-Downstream-Username": {"alice"}},
|
||||
},
|
||||
{
|
||||
name: "nil actor",
|
||||
@@ -622,9 +631,9 @@ func TestBuildUpstreamHeaders(t *testing.T) {
|
||||
clientHeaders.Set("X-Unrelated", "preserved")
|
||||
sdkCopy, clientCopy := sdkHeaders.Clone(), clientHeaders.Clone()
|
||||
|
||||
result := headers.BuildUpstreamHeaders(sdkHeaders, clientHeaders, "Authorization", map[string]string{"id": "X-Downstream-User-Id"}, &context.Actor{ID: "user-123", Metadata: recorder.Metadata{"Username": "alice"}})
|
||||
result := headers.BuildUpstreamHeaders(sdkHeaders, clientHeaders, "Authorization", map[string]string{"id": "X-Downstream-User-Id"}, &context.Actor{ID: actorID, Username: "alice"})
|
||||
|
||||
require.Equal(t, "user-123", result.Get("X-Downstream-User-Id"))
|
||||
require.Equal(t, actorID.String(), result.Get("X-Downstream-User-Id"))
|
||||
require.NotContains(t, result, http.CanonicalHeaderKey(headers.ActorIDHeader))
|
||||
require.NotContains(t, result, http.CanonicalHeaderKey(headers.ActorMetadataHeader("Username")))
|
||||
require.NotContains(t, result, http.CanonicalHeaderKey(headers.ActorMetadataHeader("Email")))
|
||||
@@ -637,7 +646,7 @@ func TestBuildUpstreamHeaders(t *testing.T) {
|
||||
t.Run("missing actor email does not reuse stale header", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
result := headers.BuildUpstreamHeaders(http.Header{"X-Email": {"sdk-email"}}, http.Header{"X-Email": {"client-email"}}, "Authorization", map[string]string{"email": "X-Email"}, &context.Actor{ID: "user-123"})
|
||||
result := headers.BuildUpstreamHeaders(http.Header{"X-Email": {"sdk-email"}}, http.Header{"X-Email": {"client-email"}}, "Authorization", map[string]string{"email": "X-Email"}, &context.Actor{ID: actorID})
|
||||
require.NotContains(t, result, "X-Email")
|
||||
})
|
||||
|
||||
@@ -655,14 +664,14 @@ func TestBuildUpstreamHeaders(t *testing.T) {
|
||||
}
|
||||
clientHeaders.Set(headers.ActorMetadataHeader("Custom"), "client-custom")
|
||||
sdkCopy, clientCopy := sdkHeaders.Clone(), clientHeaders.Clone()
|
||||
actor := &context.Actor{ID: "user-123", Metadata: recorder.Metadata{"Username": "alice"}}
|
||||
actor := &context.Actor{ID: actorID, Username: "alice"}
|
||||
|
||||
result := headers.BuildUpstreamHeaders(sdkHeaders, clientHeaders, "Authorization", map[string]string{
|
||||
"id": headers.ActorIDHeader,
|
||||
"username": headers.ActorMetadataHeader("Username"),
|
||||
}, actor)
|
||||
|
||||
require.Equal(t, []string{"user-123"}, result.Values(headers.ActorIDHeader))
|
||||
require.Equal(t, []string{actorID.String()}, result.Values(headers.ActorIDHeader))
|
||||
require.Equal(t, []string{"alice"}, result.Values(headers.ActorMetadataHeader("Username")))
|
||||
require.NotContains(t, result, http.CanonicalHeaderKey(headers.ActorMetadataHeader("Custom")))
|
||||
require.Equal(t, "Bearer provider-key", result.Get("Authorization"))
|
||||
@@ -673,14 +682,14 @@ func TestBuildUpstreamHeaders(t *testing.T) {
|
||||
t.Run("actor forwarding with nil client headers", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
actor := &context.Actor{ID: "user-123"}
|
||||
actor := &context.Actor{ID: actorID}
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
actor *context.Actor
|
||||
names map[string]string
|
||||
want string
|
||||
}{
|
||||
{name: "enabled", names: map[string]string{"id": headers.ActorIDHeader}, actor: actor, want: actor.ID},
|
||||
{name: "enabled", names: map[string]string{"id": headers.ActorIDHeader}, actor: actor, want: actorID.String()},
|
||||
{name: "disabled", actor: actor},
|
||||
{name: "nil actor", names: map[string]string{"id": headers.ActorIDHeader}},
|
||||
} {
|
||||
|
||||
@@ -3145,9 +3145,7 @@ func TestActorHeaders(t *testing.T) {
|
||||
metadataKey := "Username"
|
||||
bridgeServer := newBridgeTestServer(ctx, t, upstream.URL,
|
||||
withCustomProvider(tc.createProviderFn(upstream.URL, apiKey, send)),
|
||||
withActor(defaultActorID, recorder.Metadata{
|
||||
metadataKey: actorUsername,
|
||||
}),
|
||||
withActor(defaultActorUUID, actorUsername),
|
||||
)
|
||||
|
||||
// Add the stream param to the request.
|
||||
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
promtest "github.com/prometheus/client_golang/prometheus/testutil"
|
||||
"github.com/stretchr/testify/assert"
|
||||
@@ -30,6 +31,9 @@ const (
|
||||
openAIOverloadedError = `{"error":{"message":"Service Unavailable.","type":"cf_service_unavailable","code":503}}`
|
||||
)
|
||||
|
||||
// circuitBreakerActorID is a fixed actor distinct from defaultActorUUID.
|
||||
var circuitBreakerActorID = uuid.MustParse("8d0c2f43-6b7e-4a51-9c3d-2e5f1a7b9c04")
|
||||
|
||||
func anthropicSuccessResponse(model string) string {
|
||||
return fmt.Sprintf(`{"id":"msg_01","type":"message","role":"assistant","content":[{"type":"text","text":"Hello!"}],"model":%q,"stop_reason":"end_turn","usage":{"input_tokens":10,"output_tokens":5}}`, model)
|
||||
}
|
||||
@@ -136,7 +140,7 @@ func TestCircuitBreaker_FullRecoveryCycle(t *testing.T) {
|
||||
bridgeServer := newBridgeTestServer(ctx, t, mockUpstream.URL,
|
||||
withCustomProvider(tc.createProvider(mockUpstream.URL, cbConfig)),
|
||||
withMetrics(m),
|
||||
withActor("test-user-id", nil),
|
||||
withActor(circuitBreakerActorID, ""),
|
||||
)
|
||||
|
||||
doRequest := func() int {
|
||||
@@ -293,7 +297,7 @@ func TestCircuitBreaker_HalfOpenFailure(t *testing.T) {
|
||||
bridgeServer := newBridgeTestServer(ctx, t, mockUpstream.URL,
|
||||
withCustomProvider(tc.createProvider(mockUpstream.URL, cbConfig)),
|
||||
withMetrics(m),
|
||||
withActor("test-user-id", nil),
|
||||
withActor(circuitBreakerActorID, ""),
|
||||
)
|
||||
|
||||
doRequest := func() int {
|
||||
@@ -441,7 +445,7 @@ func TestCircuitBreaker_HalfOpenMaxRequests(t *testing.T) {
|
||||
bridgeServer := newBridgeTestServer(ctx, t, mockUpstream.URL,
|
||||
withCustomProvider(tc.createProvider(mockUpstream.URL, cbConfig)),
|
||||
withMetrics(m),
|
||||
withActor("test-user-id", nil),
|
||||
withActor(circuitBreakerActorID, ""),
|
||||
)
|
||||
|
||||
doRequest := func() int {
|
||||
@@ -560,7 +564,7 @@ func TestCircuitBreaker_PerModelIsolation(t *testing.T) {
|
||||
CircuitBreaker: cbConfig,
|
||||
}, nil)),
|
||||
withMetrics(m),
|
||||
withActor("test-user-id", nil),
|
||||
withActor(circuitBreakerActorID, ""),
|
||||
)
|
||||
|
||||
doRequest := func(model string) int {
|
||||
|
||||
@@ -84,7 +84,7 @@ func TestClaudePlatformLive(t *testing.T) {
|
||||
WorkspaceID: workspaceID,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
bridgeServer := newBridgeTestServer(ctx, t, "", withCustomProvider(p), withActor(defaultActorID, nil))
|
||||
bridgeServer := newBridgeTestServer(ctx, t, "", withCustomProvider(p), withActor(defaultActorUUID, ""))
|
||||
|
||||
body := []byte(fmt.Sprintf(`{"model":%q,"max_tokens":16,"messages":[{"role":"user","content":"Say OK"}]}`, model))
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, bridgeServer.URL+pathAnthropicMessages, bytes.NewReader(body))
|
||||
|
||||
@@ -76,7 +76,7 @@ func TestMantleLive(t *testing.T) {
|
||||
BaseURL: "https://bedrock-mantle." + region + ".api.aws/anthropic",
|
||||
Protocol: config.BedrockProtocolMantle,
|
||||
})),
|
||||
withActor(defaultActorID, nil),
|
||||
withActor(defaultActorUUID, ""),
|
||||
)
|
||||
|
||||
for _, model := range models {
|
||||
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/sjson"
|
||||
"go.opentelemetry.io/otel"
|
||||
@@ -39,11 +40,17 @@ const (
|
||||
providerBedrock = "bedrock"
|
||||
|
||||
// defaults
|
||||
apiKey = "api-key"
|
||||
defaultActorID = "ae235cc1-9f8f-417d-a636-a7b170bac62e"
|
||||
apiKey = "api-key"
|
||||
defaultActorID = "ae235cc1-9f8f-417d-a636-a7b170bac62e"
|
||||
defaultAPIKeyID = "integration-test-api-key-id"
|
||||
)
|
||||
|
||||
var defaultTracer = otel.Tracer("integrationtest")
|
||||
var (
|
||||
defaultTracer = otel.Tracer("integrationtest")
|
||||
// defaultActorUUID is defaultActorID parsed for the request actor; records
|
||||
// and labels compare against the string form.
|
||||
defaultActorUUID = uuid.MustParse(defaultActorID)
|
||||
)
|
||||
|
||||
type bridgeConfig struct {
|
||||
providerBuilders []func(t *testing.T, upstreamURL string) aibridge.Provider
|
||||
@@ -52,11 +59,13 @@ type bridgeConfig struct {
|
||||
mcpProxy mcp.ServerProxier
|
||||
// noMCPProxy leaves the proxier nil instead of falling back to
|
||||
// NoopMCPManager, which is non-nil and reports zero tools.
|
||||
noMCPProxy bool
|
||||
userID string
|
||||
metadata recorder.Metadata
|
||||
logger slog.Logger
|
||||
apiKeyID string
|
||||
noMCPProxy bool
|
||||
actorID uuid.UUID
|
||||
actorUsername string
|
||||
logger slog.Logger
|
||||
// apiKeyID is stamped on the request actor. Recorders require it, so
|
||||
// [newBridgeTestServer] defaults it to defaultAPIKeyID.
|
||||
apiKeyID string
|
||||
// structuredLogging makes the bridge emit AI Gateway interception
|
||||
// records in the format described by [recorder.InterceptionLogMarker].
|
||||
structuredLogging bool
|
||||
@@ -130,9 +139,9 @@ func withMCP(p mcp.ServerProxier) bridgeOption {
|
||||
return func(c *bridgeConfig) { c.mcpProxy = p }
|
||||
}
|
||||
|
||||
// withActor sets the actor ID and metadata for the BaseContext.
|
||||
func withActor(id string, md recorder.Metadata) bridgeOption {
|
||||
return func(c *bridgeConfig) { c.userID = id; c.metadata = md }
|
||||
// withActor sets the actor ID and username for the BaseContext.
|
||||
func withActor(id uuid.UUID, username string) bridgeOption {
|
||||
return func(c *bridgeConfig) { c.actorID = id; c.actorUsername = username }
|
||||
}
|
||||
|
||||
// newBridgeTestServer creates a fully configured test server running
|
||||
@@ -141,7 +150,8 @@ func withActor(id string, md recorder.Metadata) bridgeOption {
|
||||
// - NoopMCPManager (unless withMCP)
|
||||
// - slogtest debug logger
|
||||
// - defaultTracer (unless withTracer)
|
||||
// - defaultActorID (unless withActor)
|
||||
// - defaultActorID with no username (unless withActor)
|
||||
// - defaultAPIKeyID (unless apiKeyID is set)
|
||||
func newBridgeTestServer(
|
||||
ctx context.Context,
|
||||
t *testing.T,
|
||||
@@ -151,11 +161,14 @@ func newBridgeTestServer(
|
||||
t.Helper()
|
||||
|
||||
cfg := &bridgeConfig{
|
||||
userID: defaultActorID,
|
||||
actorID: defaultActorUUID,
|
||||
}
|
||||
for _, o := range opts {
|
||||
o(cfg)
|
||||
}
|
||||
if cfg.apiKeyID == "" {
|
||||
cfg.apiKeyID = defaultAPIKeyID
|
||||
}
|
||||
if cfg.tracer == nil {
|
||||
cfg.tracer = defaultTracer
|
||||
}
|
||||
@@ -179,7 +192,7 @@ func newBridgeTestServer(
|
||||
}
|
||||
|
||||
mockRec := &testutil.MockRecorder{}
|
||||
rec := aibridge.NewRecorder(cfg.logger, cfg.tracer, cfg.apiKeyID, cfg.structuredLogging, mockRec, cfg.recorderMiddleware...)
|
||||
rec := aibridge.NewRecorder(cfg.logger, cfg.tracer, cfg.structuredLogging, mockRec, cfg.recorderMiddleware...)
|
||||
|
||||
bridge, err := aibridge.NewRequestBridge(
|
||||
ctx, providers, rec, cfg.mcpProxy,
|
||||
@@ -187,10 +200,14 @@ func newBridgeTestServer(
|
||||
)
|
||||
require.NoError(t, err)
|
||||
|
||||
actorID, md := cfg.userID, cfg.metadata
|
||||
actor := aibcontext.Actor{
|
||||
ID: cfg.actorID,
|
||||
APIKeyID: cfg.apiKeyID,
|
||||
Username: cfg.actorUsername,
|
||||
}
|
||||
srv := httptest.NewUnstartedServer(bridge)
|
||||
srv.Config.BaseContext = func(_ net.Listener) context.Context {
|
||||
return aibcontext.AsActor(ctx, actorID, "", md)
|
||||
return aibcontext.AsActor(ctx, actor)
|
||||
}
|
||||
srv.Start()
|
||||
t.Cleanup(srv.Close)
|
||||
@@ -235,7 +252,7 @@ func setupInjectedToolTest(
|
||||
allOpts := []bridgeOption{
|
||||
withMCP(mockMCP),
|
||||
withTracer(tracer),
|
||||
withActor(defaultActorID, nil),
|
||||
withActor(defaultActorUUID, ""),
|
||||
}
|
||||
allOpts = append(allOpts, opts...)
|
||||
bridgeServer := newBridgeTestServer(ctx, t, upstream.URL, allOpts...)
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"go.opentelemetry.io/otel"
|
||||
@@ -16,7 +17,6 @@ import (
|
||||
aibheaders "github.com/coder/coder/v2/aibridge/headers"
|
||||
aibtestutil "github.com/coder/coder/v2/aibridge/internal/testutil"
|
||||
"github.com/coder/coder/v2/aibridge/keypool"
|
||||
"github.com/coder/coder/v2/aibridge/recorder"
|
||||
testutil "github.com/coder/coder/v2/testutil"
|
||||
)
|
||||
|
||||
@@ -138,7 +138,7 @@ func TestCopilot_CreateInterceptor(t *testing.T) {
|
||||
req.Header.Set("Authorization", "Bearer test-token")
|
||||
req.Header.Set("Editor-Version", "vscode/1.85.0")
|
||||
req.Header.Set("Copilot-Integration-Id", "test-integration")
|
||||
req = req.WithContext(aibcontext.AsActor(req.Context(), "actor-id", "", recorder.Metadata{"Username": "actor-username"}))
|
||||
req = req.WithContext(aibcontext.AsActor(req.Context(), aibcontext.Actor{ID: uuid.MustParse("6f1d2c3b-4a5e-4f60-8a7b-9c0d1e2f3a4b"), Username: "actor-username"}))
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
interceptor, err := provider.CreateInterceptor(w, req, testTracer)
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"google.golang.org/protobuf/types/known/structpb"
|
||||
"google.golang.org/protobuf/types/known/timestamppb"
|
||||
|
||||
aibcontext "github.com/coder/coder/v2/aibridge/context"
|
||||
agplaibridge "github.com/coder/coder/v2/coderd/aibridge"
|
||||
"github.com/coder/coder/v2/coderd/aibridged/proto"
|
||||
)
|
||||
@@ -18,7 +19,6 @@ var _ Recorder = &DRPCRecorder{}
|
||||
|
||||
// DRPCRecorder satisfies the Recorder interface and translates calls into dRPC calls to aibridgedserver.
|
||||
type DRPCRecorder struct {
|
||||
apiKeyID string
|
||||
clientFn func(context.Context) (proto.DRPCRecorderClient, error)
|
||||
}
|
||||
|
||||
@@ -35,6 +35,10 @@ func (t *DRPCRecorder) client(ctx context.Context) (proto.DRPCRecorderClient, er
|
||||
}
|
||||
|
||||
func (t *DRPCRecorder) RecordInterception(ctx context.Context, req *InterceptionRecord) error {
|
||||
apiKeyID, err := aibcontext.APIKeyIDFromContext(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
client, err := t.client(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -42,7 +46,7 @@ func (t *DRPCRecorder) RecordInterception(ctx context.Context, req *Interception
|
||||
|
||||
in := &proto.RecordInterceptionRequest{
|
||||
Id: req.ID,
|
||||
ApiKeyId: t.apiKeyID,
|
||||
ApiKeyId: apiKeyID,
|
||||
InitiatorId: req.InitiatorID,
|
||||
Provider: req.Provider,
|
||||
ProviderName: req.ProviderName,
|
||||
@@ -216,10 +220,8 @@ func marshalForProto(in Metadata) map[string]*anypb.Any {
|
||||
}
|
||||
|
||||
// NewDRPCRecorder creates a [DRPCRecorder]. clientFn receives the context of
|
||||
// the record call it serves.
|
||||
func NewDRPCRecorder(aPIKeyID string, clientFn func(context.Context) (proto.DRPCRecorderClient, error)) *DRPCRecorder {
|
||||
return &DRPCRecorder{
|
||||
apiKeyID: aPIKeyID,
|
||||
clientFn: clientFn,
|
||||
}
|
||||
// the record call it serves. Each interception reads its authenticated API-key
|
||||
// ID from that context rather than storing request identity on the recorder.
|
||||
func NewDRPCRecorder(clientFn func(context.Context) (proto.DRPCRecorderClient, error)) *DRPCRecorder {
|
||||
return &DRPCRecorder{clientFn: clientFn}
|
||||
}
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"maps"
|
||||
|
||||
"cdr.dev/slog/v3"
|
||||
aibcontext "github.com/coder/coder/v2/aibridge/context"
|
||||
"github.com/coder/coder/v2/coderd/util/ptr"
|
||||
)
|
||||
|
||||
@@ -40,23 +41,22 @@ const (
|
||||
// most records. This is suboptimal, but a requirement for backwards compatibility.
|
||||
type LogRecorder struct {
|
||||
logger slog.Logger
|
||||
apiKeyID string
|
||||
structured bool
|
||||
wrapped Recorder
|
||||
}
|
||||
|
||||
// NewLogRecorder creates a [LogRecorder] which logs each record and then
|
||||
// delegates it to wrapped. wrapped may be nil, in which case records are only
|
||||
// logged.
|
||||
func NewLogRecorder(logger slog.Logger, apiKeyID string, structured bool, wrapped Recorder) *LogRecorder {
|
||||
return &LogRecorder{logger: logger, apiKeyID: apiKeyID, structured: structured, wrapped: wrapped}
|
||||
// logged. Interception identity is read from each call's actor context.
|
||||
func NewLogRecorder(logger slog.Logger, structured bool, wrapped Recorder) *LogRecorder {
|
||||
return &LogRecorder{logger: logger, structured: structured, wrapped: wrapped}
|
||||
}
|
||||
|
||||
// WithLogging returns a [Middleware] which wraps a [Recorder] in a
|
||||
// [LogRecorder], for composition by [ChainMiddleware].
|
||||
func WithLogging(logger slog.Logger, apiKeyID string, structured bool) Middleware {
|
||||
func WithLogging(logger slog.Logger, structured bool) Middleware {
|
||||
return func(next Recorder) Recorder {
|
||||
return NewLogRecorder(logger, apiKeyID, structured, next)
|
||||
return NewLogRecorder(logger, structured, next)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -99,6 +99,10 @@ func (r *LogRecorder) marshalToolArgs(ctx context.Context, args ToolArgs) string
|
||||
}
|
||||
|
||||
func (r *LogRecorder) RecordInterception(ctx context.Context, req *InterceptionRecord) error {
|
||||
apiKeyID, err := aibcontext.APIKeyIDFromContext(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// thread_parent_id and thread_root_id are deliberately absent: they are
|
||||
// resolved from recorded tool usages by a database lookup which only
|
||||
// coderd can perform. An absent key fails a consumer loudly; a nil UUID
|
||||
@@ -106,7 +110,7 @@ func (r *LogRecorder) RecordInterception(ctx context.Context, req *InterceptionR
|
||||
r.logStructured(ctx, RecordTypeInterceptionStart,
|
||||
slog.F("interception_id", req.ID),
|
||||
slog.F("initiator_id", req.InitiatorID),
|
||||
slog.F("api_key_id", r.apiKeyID),
|
||||
slog.F("api_key_id", apiKeyID),
|
||||
slog.F("provider", req.Provider),
|
||||
slog.F("model", req.Model),
|
||||
slog.F("client", req.Client),
|
||||
@@ -116,7 +120,6 @@ func (r *LogRecorder) RecordInterception(ctx context.Context, req *InterceptionR
|
||||
slog.F("correlating_tool_call_id", ptr.NilToEmpty(req.CorrelatingToolCallID)),
|
||||
)
|
||||
|
||||
var err error
|
||||
if r.wrapped != nil {
|
||||
err = r.wrapped.RecordInterception(ctx, req)
|
||||
}
|
||||
|
||||
@@ -13,11 +13,11 @@ type Middleware func(next Recorder) Recorder
|
||||
// logging will happen first, then tracing, and then the call to the recorder.
|
||||
//
|
||||
// middleware := ChainMiddleware(
|
||||
// WithLogging(logger),
|
||||
// WithLogging(logger, true),
|
||||
// WithTracing(tracer),
|
||||
// )
|
||||
//
|
||||
// recorder := NewDRPCRecorder(apiKeyID, clientFn)
|
||||
// recorder := NewDRPCRecorder(clientFn)
|
||||
//
|
||||
// wrappedRecorder := middleware(recorder)
|
||||
//
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"go.opentelemetry.io/otel/trace/noop"
|
||||
@@ -148,8 +149,9 @@ func TestBuildProviderFromProtoSetsAPIDumpDir(t *testing.T) {
|
||||
func TestBuildProviderActorHeaders(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
actorID := uuid.MustParse("6f1d2c3b-4a5e-4f60-8a7b-9c0d1e2f3a4b")
|
||||
const (
|
||||
actorID = "authenticated-user"
|
||||
actorAPIKeyID = "authenticated-api-key-id"
|
||||
actorUsername = "authenticated-username"
|
||||
actorEmail = "authenticated@example.com"
|
||||
clientID = "client-user"
|
||||
@@ -226,14 +228,14 @@ func TestBuildProviderActorHeaders(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
|
||||
request := httptest.NewRequest(http.MethodPost, provider.RoutePrefix()+"/chat/completions", bytes.NewBufferString(`{"model":"gpt-4","messages":[],"stream":false}`))
|
||||
request = request.WithContext(aibridge.AsActor(request.Context(), actorID, actorEmail, aibridge.Metadata{"Username": actorUsername}))
|
||||
request = request.WithContext(aibridge.AsActor(request.Context(), aibridge.Actor{ID: actorID, APIKeyID: actorAPIKeyID, Username: actorUsername, Email: actorEmail}))
|
||||
request.Header.Set("Authorization", "Bearer client-key")
|
||||
request.Header.Set(headers.ActorIDHeader, clientID)
|
||||
request.Header.Set(headers.ActorMetadataHeader("Username"), clientName)
|
||||
|
||||
interceptor, err := provider.CreateInterceptor(httptest.NewRecorder(), request, noop.NewTracerProvider().Tracer("test"))
|
||||
require.NoError(t, err)
|
||||
interceptor.Setup(slog.Make(), recorder.NewLogRecorder(slog.Make(), "", false, nil), nil)
|
||||
interceptor.Setup(slog.Make(), recorder.NewLogRecorder(slog.Make(), false, nil), nil)
|
||||
|
||||
processRequest := httptest.NewRequest(http.MethodPost, provider.RoutePrefix()+"/chat/completions", bytes.NewBufferString(`{"model":"gpt-4","messages":[],"stream":false}`)).WithContext(request.Context())
|
||||
response := httptest.NewRecorder()
|
||||
@@ -241,7 +243,7 @@ func TestBuildProviderActorHeaders(t *testing.T) {
|
||||
|
||||
receivedHeaders := testutil.TryReceive(testutil.Context(t, testutil.WaitShort), t, upstreamHeaders)
|
||||
if tt.wantCustomHeaders {
|
||||
assert.Equal(t, actorID, receivedHeaders.Get(customIDHeader))
|
||||
assert.Equal(t, actorID.String(), receivedHeaders.Get(customIDHeader))
|
||||
assert.Equal(t, actorUsername, receivedHeaders.Get(customNameHeader))
|
||||
assert.Equal(t, actorEmail, receivedHeaders.Get(customMailHeader))
|
||||
} else {
|
||||
@@ -252,6 +254,9 @@ func TestBuildProviderActorHeaders(t *testing.T) {
|
||||
// Client actor headers never reach the upstream.
|
||||
assert.NotContains(t, receivedHeaders, http.CanonicalHeaderKey(headers.ActorIDHeader))
|
||||
assert.NotContains(t, receivedHeaders, http.CanonicalHeaderKey(headers.ActorMetadataHeader("Username")))
|
||||
for name, values := range receivedHeaders {
|
||||
assert.NotContains(t, values, actorAPIKeyID, "the actor API key ID must never be forwarded in %s", name)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -381,10 +381,11 @@ func (s *Server) ReplaceProviders(ctx context.Context, providers []aibridge.Prov
|
||||
return nil
|
||||
}
|
||||
|
||||
// newRecorder builds the recorder for one API key.
|
||||
// newRecorder builds a reusable recorder that reads identity from each call's
|
||||
// actor context and acquires the DRPC client with that call's context.
|
||||
//
|
||||
// revive:disable-next-line:flag-parameter // Constructor configuration flags.
|
||||
func newRecorder(logger slog.Logger, tracer trace.Tracer, apiKeyID string, structuredLogging bool, disableContentRecording bool, clientFn ClientFunc) recorder.Recorder {
|
||||
func newRecorder(logger slog.Logger, tracer trace.Tracer, structuredLogging bool, disableContentRecording bool, clientFn ClientFunc) recorder.Recorder {
|
||||
var middleware []recorder.Middleware
|
||||
if disableContentRecording {
|
||||
middleware = append(middleware, recorder.WithoutRecords(recorder.DisabledRecords{
|
||||
@@ -394,10 +395,8 @@ func newRecorder(logger slog.Logger, tracer trace.Tracer, apiKeyID string, struc
|
||||
}))
|
||||
}
|
||||
return aibridge.NewRecorder(
|
||||
logger.Named("recorder"), tracer, apiKeyID, structuredLogging,
|
||||
recorder.NewDRPCRecorder(apiKeyID, func(ctx context.Context) (proto.DRPCRecorderClient, error) {
|
||||
// The recorder outlives its caller, so acquire the client with each
|
||||
// record call's context.
|
||||
logger.Named("recorder"), tracer, structuredLogging,
|
||||
recorder.NewDRPCRecorder(func(ctx context.Context) (proto.DRPCRecorderClient, error) {
|
||||
return clientFn(ctx)
|
||||
}),
|
||||
middleware...,
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"slices"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
@@ -18,6 +19,7 @@ import (
|
||||
"cdr.dev/slog/v3"
|
||||
"cdr.dev/slog/v3/sloggers/slogtest"
|
||||
"github.com/coder/coder/v2/aibridge"
|
||||
aibcontext "github.com/coder/coder/v2/aibridge/context"
|
||||
"github.com/coder/coder/v2/aibridge/recorder"
|
||||
"github.com/coder/coder/v2/coderd/aibridged/proto"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
@@ -26,7 +28,7 @@ import (
|
||||
|
||||
func TestNewRecorder(t *testing.T) {
|
||||
t.Parallel()
|
||||
drpcErr := xerrors.New("DRPC recording was attempted")
|
||||
drpcErr := xerrors.New("client unavailable")
|
||||
for _, tc := range []struct {
|
||||
structuredLogging bool
|
||||
disableContentRecording bool
|
||||
@@ -57,24 +59,24 @@ func TestNewRecorder(t *testing.T) {
|
||||
} {
|
||||
t.Run(fmt.Sprintf("Structured=%t/DisableContent=%t", tc.structuredLogging, tc.disableContentRecording), func(t *testing.T) {
|
||||
t.Parallel()
|
||||
apiKeyID, interceptionID := uuid.NewString(), uuid.NewString()
|
||||
actor := aibridge.Actor{ID: uuid.New(), APIKeyID: uuid.NewString()}
|
||||
id := uuid.NewString()
|
||||
now := time.Now().UTC()
|
||||
type contextKey struct{}
|
||||
sink := testutil.NewFakeSink(t)
|
||||
clientCalls := 0
|
||||
rec := newRecorder(sink.Logger(), noop.NewTracerProvider().Tracer(t.Name()), apiKeyID, tc.structuredLogging, tc.disableContentRecording, func(ctx context.Context) (DRPCClient, error) {
|
||||
rec := newRecorder(sink.Logger(), noop.NewTracerProvider().Tracer(t.Name()), tc.structuredLogging, tc.disableContentRecording, func(ctx context.Context) (DRPCClient, error) {
|
||||
clientCalls++
|
||||
require.Equal(t, "record context", ctx.Value(contextKey{}))
|
||||
require.Equal(t, &actor, aibcontext.ActorFromContext(ctx), "the client must be acquired with the record call's actor")
|
||||
return nil, drpcErr
|
||||
})
|
||||
require.Zero(t, clientCalls, "creating a recorder must not acquire a client")
|
||||
ctx := context.WithValue(t.Context(), contextKey{}, "record context")
|
||||
require.ErrorIs(t, rec.RecordInterception(ctx, &recorder.InterceptionRecord{ID: interceptionID, StartedAt: now}), drpcErr)
|
||||
require.ErrorIs(t, rec.RecordInterceptionEnded(ctx, &recorder.InterceptionRecordEnded{ID: interceptionID, EndedAt: now}), drpcErr)
|
||||
require.ErrorIs(t, rec.RecordTokenUsage(ctx, &recorder.TokenUsageRecord{InterceptionID: interceptionID, CreatedAt: now, Input: 1}), drpcErr)
|
||||
require.ErrorIs(t, rec.RecordPromptUsage(ctx, &recorder.PromptUsageRecord{InterceptionID: interceptionID, CreatedAt: now, Prompt: "prompt"}), tc.wantContentErr)
|
||||
require.ErrorIs(t, rec.RecordToolUsage(ctx, &recorder.ToolUsageRecord{InterceptionID: interceptionID, CreatedAt: now, Tool: "tool"}), tc.wantContentErr)
|
||||
require.ErrorIs(t, rec.RecordModelThought(ctx, &recorder.ModelThoughtRecord{InterceptionID: interceptionID, CreatedAt: now, Content: "thought"}), tc.wantContentErr)
|
||||
ctx := aibridge.AsActor(t.Context(), actor)
|
||||
require.ErrorIs(t, rec.RecordInterception(ctx, &recorder.InterceptionRecord{ID: id, InitiatorID: actor.ID.String(), StartedAt: now}), drpcErr)
|
||||
require.ErrorIs(t, rec.RecordInterceptionEnded(ctx, &recorder.InterceptionRecordEnded{ID: id, EndedAt: now}), drpcErr)
|
||||
require.ErrorIs(t, rec.RecordTokenUsage(ctx, &recorder.TokenUsageRecord{InterceptionID: id, CreatedAt: now, Input: 1}), drpcErr)
|
||||
require.ErrorIs(t, rec.RecordPromptUsage(ctx, &recorder.PromptUsageRecord{InterceptionID: id, CreatedAt: now, Prompt: "prompt"}), tc.wantContentErr)
|
||||
require.ErrorIs(t, rec.RecordToolUsage(ctx, &recorder.ToolUsageRecord{InterceptionID: id, CreatedAt: now, Tool: "tool"}), tc.wantContentErr)
|
||||
require.ErrorIs(t, rec.RecordModelThought(ctx, &recorder.ModelThoughtRecord{InterceptionID: id, CreatedAt: now, Content: "thought"}), tc.wantContentErr)
|
||||
require.Equal(t, tc.wantClientCalls, clientCalls, "each forwarded record acquires its own client")
|
||||
logs := sink.Entries(func(entry slog.SinkEntry) bool { return entry.Message == recorder.InterceptionLogMarker })
|
||||
require.Len(t, logs, tc.wantLogs, "logging must precede content filtering")
|
||||
@@ -82,29 +84,79 @@ func TestNewRecorder(t *testing.T) {
|
||||
require.Equal(t, tc.wantClientCalls, clientCalls)
|
||||
})
|
||||
}
|
||||
t.Run("APIKeyID", func(t *testing.T) {
|
||||
|
||||
t.Run("RecorderSharedAcrossRequestsDoesNotMixActorData", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
apiKeyID, interceptionID := uuid.NewString(), uuid.NewString()
|
||||
client := &interceptionRecordingClient{}
|
||||
rec := newRecorder(testutil.NewFakeSink(t).Logger(), noop.NewTracerProvider().Tracer(t.Name()), apiKeyID, false, false, func(context.Context) (DRPCClient, error) {
|
||||
sink := testutil.NewFakeSink(t)
|
||||
client := &recordingClient{}
|
||||
rec := newRecorder(sink.Logger(), noop.NewTracerProvider().Tracer(t.Name()), true, false, func(context.Context) (DRPCClient, error) {
|
||||
return client, nil
|
||||
})
|
||||
require.NoError(t, rec.RecordInterception(t.Context(), &recorder.InterceptionRecord{ID: interceptionID, StartedAt: time.Now().UTC()}))
|
||||
require.Len(t, client.requests, 1)
|
||||
require.Equal(t, apiKeyID, client.requests[0].GetApiKeyId())
|
||||
require.Equal(t, interceptionID, client.requests[0].GetId())
|
||||
record := func(t *testing.T) {
|
||||
t.Helper()
|
||||
actor := aibridge.Actor{ID: uuid.New(), APIKeyID: uuid.NewString()}
|
||||
in := &recorder.InterceptionRecord{ID: uuid.NewString(), InitiatorID: actor.ID.String(), StartedAt: time.Now().UTC()}
|
||||
require.NoError(t, rec.RecordInterception(aibridge.AsActor(t.Context(), actor), in))
|
||||
|
||||
value, ok := client.requests.Load(in.ID)
|
||||
require.True(t, ok, "interception must reach the DRPC client")
|
||||
rpc := value.(*proto.RecordInterceptionRequest)
|
||||
require.Equal(t, in.InitiatorID, rpc.GetInitiatorId())
|
||||
require.Equal(t, actor.APIKeyID, rpc.GetApiKeyId())
|
||||
logs := sink.Entries(func(entry slog.SinkEntry) bool {
|
||||
return entry.Message == recorder.InterceptionLogMarker && slices.ContainsFunc(entry.Fields, func(field slog.Field) bool {
|
||||
return field.Name == "interception_id" && field.Value == in.ID
|
||||
})
|
||||
})
|
||||
require.Len(t, logs, 1)
|
||||
require.Contains(t, logs[0].Fields, slog.F("initiator_id", in.InitiatorID))
|
||||
require.Contains(t, logs[0].Fields, slog.F("api_key_id", actor.APIKeyID))
|
||||
}
|
||||
|
||||
t.Run("Sequential", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
record(t)
|
||||
record(t)
|
||||
})
|
||||
t.Run("Concurrent", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
for i := range 2 {
|
||||
t.Run(fmt.Sprintf("Request%d", i), func(t *testing.T) {
|
||||
t.Parallel()
|
||||
record(t)
|
||||
})
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
t.Run("DetachedContext", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
actor := aibridge.Actor{ID: uuid.New(), APIKeyID: uuid.NewString()}
|
||||
clientCalls := 0
|
||||
rec := newRecorder(testutil.NewFakeSink(t).Logger(), noop.NewTracerProvider().Tracer(t.Name()), false, false, func(ctx context.Context) (DRPCClient, error) {
|
||||
clientCalls++
|
||||
require.NoError(t, ctx.Err(), "client acquisition must not inherit request cancellation")
|
||||
require.Equal(t, &actor, aibcontext.ActorFromContext(ctx))
|
||||
return nil, drpcErr
|
||||
})
|
||||
ctx, cancel := context.WithCancel(aibridge.AsActor(t.Context(), actor))
|
||||
cancel()
|
||||
require.ErrorIs(t, rec.RecordInterception(context.WithoutCancel(ctx), &recorder.InterceptionRecord{ID: uuid.NewString(), InitiatorID: actor.ID.String(), StartedAt: time.Now().UTC()}), drpcErr)
|
||||
require.Equal(t, 1, clientCalls)
|
||||
})
|
||||
}
|
||||
|
||||
// interceptionRecordingClient captures interception records. Its nil embedded
|
||||
// client makes any other call fail loudly.
|
||||
type interceptionRecordingClient struct {
|
||||
// recordingClient captures interception RPCs by ID for concurrent assertions.
|
||||
// Its nil embedded client makes unexpected calls fail.
|
||||
type recordingClient struct {
|
||||
DRPCClient
|
||||
requests []*proto.RecordInterceptionRequest
|
||||
requests sync.Map
|
||||
}
|
||||
|
||||
func (c *interceptionRecordingClient) RecordInterception(_ context.Context, in *proto.RecordInterceptionRequest) (*proto.RecordInterceptionResponse, error) {
|
||||
c.requests = append(c.requests, in)
|
||||
func (c *recordingClient) RecordInterception(_ context.Context, in *proto.RecordInterceptionRequest) (*proto.RecordInterceptionResponse, error) {
|
||||
if _, loaded := c.requests.LoadOrStore(in.GetId(), in); loaded {
|
||||
return nil, xerrors.New("interception recorded more than once")
|
||||
}
|
||||
return &proto.RecordInterceptionResponse{}, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -20,6 +20,7 @@ import (
|
||||
"cdr.dev/slog/v3/sloggers/slogtest"
|
||||
"github.com/coder/coder/v2/aibridge"
|
||||
"github.com/coder/coder/v2/aibridge/aibridgetest"
|
||||
aibcontext "github.com/coder/coder/v2/aibridge/context"
|
||||
aibheaders "github.com/coder/coder/v2/aibridge/headers"
|
||||
"github.com/coder/coder/v2/aibridge/keypool"
|
||||
agplaibridge "github.com/coder/coder/v2/coderd/aibridge"
|
||||
@@ -376,6 +377,54 @@ func TestServeHTTP_FailureModes(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestServeHTTP_ActorIdentity(t *testing.T) {
|
||||
t.Parallel()
|
||||
for _, delegated := range []bool{false, true} {
|
||||
name := "Direct"
|
||||
if delegated {
|
||||
name = "Delegated"
|
||||
}
|
||||
t.Run(name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
srv, client, pool := newTestServer(t)
|
||||
ownerID := uuid.New()
|
||||
verified := aibridge.Actor{
|
||||
ID: ownerID,
|
||||
APIKeyID: uuid.NewString(),
|
||||
Username: "verified-user",
|
||||
Email: "verified@example.test",
|
||||
}
|
||||
client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).Return(&proto.IsAuthorizedResponse{
|
||||
OwnerId: ownerID.String(),
|
||||
ApiKeyId: verified.APIKeyID,
|
||||
Username: verified.Username,
|
||||
Email: verified.Email,
|
||||
}, nil)
|
||||
client.EXPECT().IsBudgetExceeded(gomock.Any(), gomock.Any()).Return(&proto.IsBudgetExceededResponse{}, nil)
|
||||
pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
|
||||
func(ctx context.Context, request aibridged.Request, _ aibridged.ClientFunc, _ aibridged.MCPProxyBuilder) (http.Handler, error) {
|
||||
require.Equal(t, verified, *aibcontext.ActorFromContext(ctx))
|
||||
require.Equal(t, verified.APIKeyID, request.APIKeyID)
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
require.Equal(t, verified, *aibcontext.ActorFromContext(r.Context()))
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
}), nil
|
||||
})
|
||||
ctx := aibridge.AsActor(t.Context(), aibridge.Actor{ID: uuid.New(), APIKeyID: "unverified-key", Username: "unverified"})
|
||||
if delegated {
|
||||
ctx = agplaibridge.WithDelegatedAPIKeyID(ctx, "delegated-key")
|
||||
}
|
||||
req := httptest.NewRequest(http.MethodPost, "/openai/v1/chat/completions", nil).WithContext(ctx)
|
||||
if !delegated {
|
||||
req.Header.Set("Authorization", "Bearer coder-secret")
|
||||
}
|
||||
rw := httptest.NewRecorder()
|
||||
srv.ServeHTTP(rw, req)
|
||||
require.Equal(t, http.StatusNoContent, rw.Code)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// When the request context carries a delegated API key ID (set by the
|
||||
// in-process transport on behalf of a trusted caller like chatd), the handler
|
||||
// must authenticate via the key_id field, skipping the header-based key
|
||||
@@ -873,8 +922,12 @@ func TestServeHTTP_ActorHeaders(t *testing.T) {
|
||||
conn := &mockDRPCConn{}
|
||||
client.EXPECT().DRPCConn().AnyTimes().Return(conn)
|
||||
|
||||
// Return authorization response with user ID, username, and email.
|
||||
authResponse := &proto.IsAuthorizedResponse{OwnerId: testUserID.String(), Username: testUsername}
|
||||
// Return authorization response with user ID, API key ID, username, and email.
|
||||
authResponse := &proto.IsAuthorizedResponse{
|
||||
OwnerId: testUserID.String(),
|
||||
ApiKeyId: uuid.NewString(),
|
||||
Username: testUsername,
|
||||
}
|
||||
if !tc.legacy {
|
||||
authResponse.Email = testEmail
|
||||
}
|
||||
@@ -988,7 +1041,10 @@ func TestRouting(t *testing.T) {
|
||||
conn := &mockDRPCConn{}
|
||||
client.EXPECT().DRPCConn().AnyTimes().Return(conn)
|
||||
|
||||
client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId: uuid.NewString()}, nil)
|
||||
client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{
|
||||
OwnerId: uuid.NewString(),
|
||||
ApiKeyId: uuid.NewString(),
|
||||
}, nil)
|
||||
client.EXPECT().IsBudgetExceeded(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsBudgetExceededResponse{}, nil)
|
||||
client.EXPECT().GetMCPServerConfigs(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.GetMCPServerConfigsResponse{}, nil)
|
||||
// This is the only recording we really care about in this test. This is called before the provider-specific logic processes
|
||||
|
||||
@@ -10,7 +10,6 @@ import (
|
||||
|
||||
"cdr.dev/slog/v3"
|
||||
"github.com/coder/coder/v2/aibridge"
|
||||
"github.com/coder/coder/v2/aibridge/recorder"
|
||||
agplaibridge "github.com/coder/coder/v2/coderd/aibridge"
|
||||
"github.com/coder/coder/v2/coderd/aibridged/proto"
|
||||
)
|
||||
@@ -178,15 +177,15 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
// Rewire request context to include actor.
|
||||
//
|
||||
// [NOTE]
|
||||
// The metadata provided here must NOT be sensitive as it is recorded and
|
||||
// could be included in requests to upstream services. Email is passed
|
||||
// separately so it is only forwarded when configured.
|
||||
r = r.WithContext(aibridge.AsActor(ctx, resp.GetOwnerId(), resp.GetEmail(), recorder.Metadata{
|
||||
"Username": resp.GetUsername(),
|
||||
}))
|
||||
// Attach only the verified identity. Key secrets remain outside the
|
||||
// request context.
|
||||
ctx = aibridge.AsActor(ctx, aibridge.Actor{
|
||||
ID: id,
|
||||
APIKeyID: resp.GetApiKeyId(),
|
||||
Username: resp.GetUsername(),
|
||||
Email: resp.GetEmail(),
|
||||
})
|
||||
r = r.WithContext(ctx)
|
||||
|
||||
handler, err := s.GetRequestHandler(ctx, Request{
|
||||
SessionKey: key,
|
||||
|
||||
@@ -258,7 +258,7 @@ func (p *CachedBridgePool) Acquire(ctx context.Context, req Request, clientFn Cl
|
||||
|
||||
span.AddEvent("cache_miss")
|
||||
providerVersion := p.providerVersion.Load()
|
||||
rec := newRecorder(p.logger, p.tracer, req.APIKeyID, p.options.StructuredLogging, p.options.DisableContentRecording, clientFn)
|
||||
rec := newRecorder(p.logger, p.tracer, p.options.StructuredLogging, p.options.DisableContentRecording, clientFn)
|
||||
|
||||
// Slow path.
|
||||
// Creating an *aibridge.RequestBridge may take some time, so gate all subsequent callers behind the initial request and return the resulting value.
|
||||
|
||||
Reference in New Issue
Block a user