Skip to content

Commit 7123518

Browse files
dannykoppingdannykopping
authored
feat: conditionally send aibridge actor headers (coder#21643)
Also passes along the authenticated username as actor metadata. Closes coder/aibridge#135 Depends on coder/aibridge#142 **Replace aibridge tag with merge commit once coder/aibridge#142 lands.** --------- Signed-off-by: Danny Kopping <danny@coder.com>
1 parent bb186b8 commit 7123518

19 files changed

Lines changed: 252 additions & 69 deletions

File tree

‎cli/testdata/coder_server_--help.golden‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,14 @@ AI BRIDGE OPTIONS:
158158
Maximum number of AI Bridge requests per second per replica. Set to 0
159159
to disable (unlimited).
160160

161+
--aibridge-send-actor-headers bool, $CODER_AIBRIDGE_SEND_ACTOR_HEADERS (default: false)
162+
Once enabled, extra headers will be added to upstream requests to
163+
identify the user (actor) making requests to AI Bridge. This is only
164+
needed if you are using a proxy between AI Bridge and an upstream AI
165+
provider. This will send X-Ai-Bridge-Actor-Id (the ID of the user
166+
making the request) and X-Ai-Bridge-Actor-Metadata-Username (their
167+
username).
168+
161169
--aibridge-structured-logging bool, $CODER_AIBRIDGE_STRUCTURED_LOGGING (default: false)
162170
Emit structured logs for AI Bridge interception records. Use this for
163171
exporting these records to external SIEM or observability systems.

‎cli/testdata/server-config.yaml.golden‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -782,6 +782,13 @@ aibridge:
782782
# these records to external SIEM or observability systems.
783783
# (default: false, type: bool)
784784
structuredLogging: false
785+
# Once enabled, extra headers will be added to upstream requests to identify the
786+
# user (actor) making requests to AI Bridge. This is only needed if you are using
787+
# a proxy between AI Bridge and an upstream AI provider. This will send
788+
# X-Ai-Bridge-Actor-Id (the ID of the user making the request) and
789+
# X-Ai-Bridge-Actor-Metadata-Username (their username).
790+
# (default: false, type: bool)
791+
send_actor_headers: false
785792
# Enable the circuit breaker to protect against cascading failures from upstream
786793
# AI provider rate limits (429, 503, 529 overloaded).
787794
# (default: false, type: bool)

‎coderd/apidoc/docs.go‎

Lines changed: 3 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/apidoc/swagger.json‎

Lines changed: 3 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎codersdk/deployment.go‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3509,6 +3509,18 @@ Write out the current server config as YAML to stdout.`,
35093509
Group: &deploymentGroupAIBridge,
35103510
YAML: "structuredLogging",
35113511
},
3512+
{
3513+
Name: "AI Bridge Send Actor Headers",
3514+
Description: "Once enabled, extra headers will be added to upstream requests to identify the user (actor) making requests to AI Bridge. " +
3515+
"This is only needed if you are using a proxy between AI Bridge and an upstream AI provider. " +
3516+
"This will send X-Ai-Bridge-Actor-Id (the ID of the user making the request) and X-Ai-Bridge-Actor-Metadata-Username (their username).",
3517+
Flag: "aibridge-send-actor-headers",
3518+
Env: "CODER_AIBRIDGE_SEND_ACTOR_HEADERS",
3519+
Value: &c.AI.BridgeConfig.SendActorHeaders,
3520+
Default: "false",
3521+
Group: &deploymentGroupAIBridge,
3522+
YAML: "send_actor_headers",
3523+
},
35123524
{
35133525
Name: "AI Bridge Circuit Breaker Enabled",
35143526
Description: "Enable the circuit breaker to protect against cascading failures from upstream AI provider rate limits (429, 503, 529 overloaded).",
@@ -3722,6 +3734,7 @@ type AIBridgeConfig struct {
37223734
MaxConcurrency serpent.Int64 `json:"max_concurrency" typescript:",notnull"`
37233735
RateLimit serpent.Int64 `json:"rate_limit" typescript:",notnull"`
37243736
StructuredLogging serpent.Bool `json:"structured_logging" typescript:",notnull"`
3737+
SendActorHeaders serpent.Bool `json:"send_actor_headers" typescript:",notnull"`
37253738
// Circuit breaker protects against cascading failures from upstream AI
37263739
// provider rate limits (429, 503, 529 overloaded).
37273740
CircuitBreakerEnabled serpent.Bool `json:"circuit_breaker_enabled" typescript:",notnull"`

‎docs/reference/api/general.md‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎docs/reference/api/schemas.md‎

Lines changed: 5 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎docs/reference/cli/server.md‎

Lines changed: 11 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎enterprise/aibridged/aibridged_test.go‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import (
1717

1818
"cdr.dev/slog/v3/sloggers/slogtest"
1919
"github.com/coder/aibridge"
20+
"github.com/coder/aibridge/intercept"
2021
agplaibridge "github.com/coder/coder/v2/coderd/aibridge"
2122
"github.com/coder/coder/v2/codersdk"
2223
"github.com/coder/coder/v2/enterprise/aibridged"
@@ -312,6 +313,105 @@ func (h *mockHandler) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
312313
_, _ = rw.Write([]byte(r.URL.Path))
313314
}
314315

316+
// TestServeHTTP_ActorHeaders validates that actor headers are correctly forwarded to
317+
// upstream AI providers when SendActorHeaders is enabled in the provider configuration.
318+
// These headers allow upstream providers to identify the user making the request for
319+
// tracking and auditing purposes.
320+
func TestServeHTTP_ActorHeaders(t *testing.T) {
321+
t.Parallel()
322+
323+
testUsername := "testuser"
324+
testUserID := uuid.New()
325+
326+
cases := []struct {
327+
path string
328+
}{
329+
// Not a complete set of paths; we're not testing the specific APIs - just the provider configs.
330+
{
331+
path: "/openai/v1/chat/completions",
332+
},
333+
{
334+
path: "/anthropic/v1/messages",
335+
},
336+
}
337+
338+
for _, tc := range cases {
339+
t.Run(tc.path, func(t *testing.T) {
340+
t.Parallel()
341+
342+
// Setup mock upstream AI server that captures headers.
343+
var receivedHeaders http.Header
344+
upstreamSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
345+
receivedHeaders = r.Header.Clone()
346+
w.WriteHeader(http.StatusTeapot)
347+
_, _ = w.Write([]byte(`i am a teapot`))
348+
}))
349+
t.Cleanup(upstreamSrv.Close)
350+
351+
// Setup with SendActorHeaders enabled.
352+
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
353+
ctrl := gomock.NewController(t)
354+
client := mock.NewMockDRPCClient(ctrl)
355+
356+
// Create providers with SendActorHeaders=true.
357+
providers := []aibridge.Provider{
358+
aibridge.NewOpenAIProvider(aibridge.OpenAIConfig{
359+
BaseURL: upstreamSrv.URL,
360+
SendActorHeaders: true,
361+
}),
362+
aibridge.NewAnthropicProvider(aibridge.AnthropicConfig{
363+
BaseURL: upstreamSrv.URL,
364+
SendActorHeaders: true,
365+
}, nil),
366+
}
367+
368+
pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, providers, logger, nil, testTracer)
369+
require.NoError(t, err)
370+
conn := &mockDRPCConn{}
371+
client.EXPECT().DRPCConn().AnyTimes().Return(conn)
372+
373+
// Return authorization response with user ID and username.
374+
client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{
375+
OwnerId: testUserID.String(),
376+
Username: testUsername,
377+
}, nil)
378+
client.EXPECT().GetMCPServerConfigs(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.GetMCPServerConfigsResponse{}, nil)
379+
client.EXPECT().RecordInterception(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.RecordInterceptionResponse{}, nil)
380+
client.EXPECT().RecordInterceptionEnded(gomock.Any(), gomock.Any()).AnyTimes()
381+
382+
// Given: aibridged is started.
383+
srv, err := aibridged.New(t.Context(), pool, func(ctx context.Context) (aibridged.DRPCClient, error) {
384+
return client, nil
385+
}, logger, testTracer)
386+
require.NoError(t, err, "create new aibridged")
387+
t.Cleanup(func() {
388+
_ = srv.Shutdown(testutil.Context(t, testutil.WaitShort))
389+
})
390+
391+
// When: a request is made to aibridged.
392+
ctx := testutil.Context(t, testutil.WaitShort)
393+
req, err := http.NewRequestWithContext(ctx, http.MethodPost, tc.path, bytes.NewBufferString(`{}`))
394+
require.NoError(t, err, "make request to test server")
395+
req.Header.Add("Authorization", "Bearer key")
396+
req.Header.Add("Accept", "application/json")
397+
398+
// When: aibridged handles the request.
399+
rec := httptest.NewRecorder()
400+
srv.ServeHTTP(rec, req)
401+
402+
// Then: the actor headers should be present in the upstream request.
403+
require.NotEmpty(t, receivedHeaders, "upstream server should have received headers")
404+
405+
// Verify the actor ID header is present with the correct value.
406+
actorIDHeader := receivedHeaders.Get(intercept.ActorIDHeader())
407+
assert.Equal(t, testUserID.String(), actorIDHeader, "actor ID header should contain user ID")
408+
// Verify the actor metadata header for username is present.
409+
usernameHeader := receivedHeaders.Get(intercept.ActorMetadataHeader("Username"))
410+
assert.Equal(t, testUsername, usernameHeader, "actor metadata username header should contain username")
411+
})
412+
}
413+
}
414+
315415
// TestRouting validates that a request which originates with aibridged will be handled
316416
// by coder/aibridge's handling logic in a provider-specific manner.
317417
// We must validate that logic that pertains to coder/coder is exercised.

‎enterprise/aibridged/http.go‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99

1010
"cdr.dev/slog/v3"
1111
"github.com/coder/aibridge"
12+
"github.com/coder/aibridge/recorder"
1213
agplaibridge "github.com/coder/coder/v2/coderd/aibridge"
1314
"github.com/coder/coder/v2/enterprise/aibridged/proto"
1415
)
@@ -61,7 +62,13 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
6162
}
6263

6364
// Rewire request context to include actor.
64-
r = r.WithContext(aibridge.AsActor(ctx, resp.GetOwnerId(), nil))
65+
//
66+
// [NOTE]
67+
// The metadata provided here must NOT be sensitive as it could be included
68+
// in requests to upstream services.
69+
r = r.WithContext(aibridge.AsActor(ctx, resp.GetOwnerId(), recorder.Metadata{
70+
"Username": resp.GetUsername(),
71+
}))
6572

6673
id, err := uuid.Parse(resp.GetOwnerId())
6774
if err != nil {

0 commit comments

Comments
 (0)