Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions cli/testdata/coder_server_--help.golden
Original file line number Diff line number Diff line change
Expand Up @@ -158,6 +158,14 @@ AI BRIDGE OPTIONS:
Maximum number of AI Bridge requests per second per replica. Set to 0
to disable (unlimited).

--aibridge-send-actor-headers bool, $CODER_AIBRIDGE_SEND_ACTOR_HEADERS (default: false)
Once enabled, extra headers will be added to upstream requests to
identify the user (actor) making requests to AI Bridge. This is only
needed if you are using a proxy between AI Bridge and an upstream AI
provider. 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).

--aibridge-structured-logging bool, $CODER_AIBRIDGE_STRUCTURED_LOGGING (default: false)
Emit structured logs for AI Bridge interception records. Use this for
exporting these records to external SIEM or observability systems.
Expand Down
7 changes: 7 additions & 0 deletions cli/testdata/server-config.yaml.golden
Original file line number Diff line number Diff line change
Expand Up @@ -782,6 +782,13 @@ aibridge:
# these records to external SIEM or observability systems.
# (default: false, type: bool)
structuredLogging: false
# Once enabled, extra headers will be added to upstream requests to identify the
# user (actor) making requests to AI Bridge. This is only needed if you are using
# a proxy between AI Bridge and an upstream AI provider. 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).
# (default: false, type: bool)
send_actor_headers: false
# Enable the circuit breaker to protect against cascading failures from upstream
# AI provider rate limits (429, 503, 529 overloaded).
# (default: false, type: bool)
Expand Down
3 changes: 3 additions & 0 deletions coderd/apidoc/docs.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 3 additions & 0 deletions coderd/apidoc/swagger.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

13 changes: 13 additions & 0 deletions codersdk/deployment.go
Original file line number Diff line number Diff line change
Expand Up @@ -3509,6 +3509,18 @@ Write out the current server config as YAML to stdout.`,
Group: &deploymentGroupAIBridge,
YAML: "structuredLogging",
},
{
Name: "AI Bridge Send Actor Headers",
Description: "Once enabled, extra headers will be added to upstream requests to identify the user (actor) making requests to AI Bridge. " +
"This is only needed if you are using a proxy between AI Bridge and an upstream AI provider. " +
"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).",
Flag: "aibridge-send-actor-headers",
Env: "CODER_AIBRIDGE_SEND_ACTOR_HEADERS",
Value: &c.AI.BridgeConfig.SendActorHeaders,
Default: "false",
Group: &deploymentGroupAIBridge,
YAML: "send_actor_headers",
Comment thread
dannykopping marked this conversation as resolved.
},
{
Name: "AI Bridge Circuit Breaker Enabled",
Description: "Enable the circuit breaker to protect against cascading failures from upstream AI provider rate limits (429, 503, 529 overloaded).",
Expand Down Expand Up @@ -3722,6 +3734,7 @@ type AIBridgeConfig struct {
MaxConcurrency serpent.Int64 `json:"max_concurrency" typescript:",notnull"`
RateLimit serpent.Int64 `json:"rate_limit" typescript:",notnull"`
StructuredLogging serpent.Bool `json:"structured_logging" typescript:",notnull"`
SendActorHeaders serpent.Bool `json:"send_actor_headers" typescript:",notnull"`
// Circuit breaker protects against cascading failures from upstream AI
// provider rate limits (429, 503, 529 overloaded).
CircuitBreakerEnabled serpent.Bool `json:"circuit_breaker_enabled" typescript:",notnull"`
Expand Down
1 change: 1 addition & 0 deletions docs/reference/api/general.md

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

5 changes: 5 additions & 0 deletions docs/reference/api/schemas.md

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

11 changes: 11 additions & 0 deletions docs/reference/cli/server.md

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

100 changes: 100 additions & 0 deletions enterprise/aibridged/aibridged_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ import (

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

// TestServeHTTP_ActorHeaders validates that actor headers are correctly forwarded to
// upstream AI providers when SendActorHeaders is enabled in the provider configuration.
// These headers allow upstream providers to identify the user making the request for
// tracking and auditing purposes.
func TestServeHTTP_ActorHeaders(t *testing.T) {
t.Parallel()

testUsername := "testuser"
testUserID := uuid.New()

cases := []struct {
path string
}{
// Not a complete set of paths; we're not testing the specific APIs - just the provider configs.
{
path: "/openai/v1/chat/completions",
},
{
path: "/anthropic/v1/messages",
},
}

for _, tc := range cases {
t.Run(tc.path, func(t *testing.T) {
t.Parallel()

// Setup mock upstream AI server that captures headers.
var receivedHeaders http.Header
upstreamSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
receivedHeaders = r.Header.Clone()
w.WriteHeader(http.StatusTeapot)
_, _ = w.Write([]byte(`i am a teapot`))
}))
t.Cleanup(upstreamSrv.Close)

// Setup with SendActorHeaders enabled.
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
ctrl := gomock.NewController(t)
client := mock.NewMockDRPCClient(ctrl)

// Create providers with SendActorHeaders=true.
providers := []aibridge.Provider{
aibridge.NewOpenAIProvider(aibridge.OpenAIConfig{
BaseURL: upstreamSrv.URL,
SendActorHeaders: true,
}),
aibridge.NewAnthropicProvider(aibridge.AnthropicConfig{
BaseURL: upstreamSrv.URL,
SendActorHeaders: true,
}, nil),
}

pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, providers, logger, nil, testTracer)
require.NoError(t, err)
conn := &mockDRPCConn{}
client.EXPECT().DRPCConn().AnyTimes().Return(conn)

// Return authorization response with user ID and username.
client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{
OwnerId: testUserID.String(),
Username: testUsername,
}, nil)
client.EXPECT().GetMCPServerConfigs(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.GetMCPServerConfigsResponse{}, nil)
client.EXPECT().RecordInterception(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.RecordInterceptionResponse{}, nil)
client.EXPECT().RecordInterceptionEnded(gomock.Any(), gomock.Any()).AnyTimes()

// Given: aibridged is started.
srv, err := aibridged.New(t.Context(), pool, func(ctx context.Context) (aibridged.DRPCClient, error) {
return client, nil
}, logger, testTracer)
require.NoError(t, err, "create new aibridged")
t.Cleanup(func() {
_ = srv.Shutdown(testutil.Context(t, testutil.WaitShort))
})

// When: a request is made to aibridged.
ctx := testutil.Context(t, testutil.WaitShort)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, tc.path, bytes.NewBufferString(`{}`))
require.NoError(t, err, "make request to test server")
req.Header.Add("Authorization", "Bearer key")
req.Header.Add("Accept", "application/json")

// When: aibridged handles the request.
rec := httptest.NewRecorder()
srv.ServeHTTP(rec, req)

// Then: the actor headers should be present in the upstream request.
require.NotEmpty(t, receivedHeaders, "upstream server should have received headers")

// Verify the actor ID header is present with the correct value.
actorIDHeader := receivedHeaders.Get(intercept.ActorIDHeader())
assert.Equal(t, testUserID.String(), actorIDHeader, "actor ID header should contain user ID")
// Verify the actor metadata header for username is present.
usernameHeader := receivedHeaders.Get(intercept.ActorMetadataHeader("Username"))
assert.Equal(t, testUsername, usernameHeader, "actor metadata username header should contain username")
})
}
}

// TestRouting validates that a request which originates with aibridged will be handled
// by coder/aibridge's handling logic in a provider-specific manner.
// We must validate that logic that pertains to coder/coder is exercised.
Expand Down
9 changes: 8 additions & 1 deletion enterprise/aibridged/http.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import (

"cdr.dev/slog/v3"
"github.com/coder/aibridge"
"github.com/coder/aibridge/recorder"
agplaibridge "github.com/coder/coder/v2/coderd/aibridge"
"github.com/coder/coder/v2/enterprise/aibridged/proto"
)
Expand Down Expand Up @@ -61,7 +62,13 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
}

// Rewire request context to include actor.
r = r.WithContext(aibridge.AsActor(ctx, resp.GetOwnerId(), nil))
//
// [NOTE]
// The metadata provided here must NOT be sensitive as it could be included
// in requests to upstream services.
Comment thread
dannykopping marked this conversation as resolved.
r = r.WithContext(aibridge.AsActor(ctx, resp.GetOwnerId(), recorder.Metadata{
"Username": resp.GetUsername(),
}))
Comment thread
dannykopping marked this conversation as resolved.

id, err := uuid.Parse(resp.GetOwnerId())
if err != nil {
Expand Down
Loading
Loading