coder · pawbana · Nov 10, 2025 · Oct 27, 2025 · Oct 28, 2025 · Oct 28, 2025
diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go
diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json
diff --git a/coderd/database/db2sdk/db2sdk.go b/coderd/database/db2sdk/db2sdk.go
@@ -974,6 +974,9 @@ func AIBridgeInterception(interception database.AIBridgeInterception, initiator
 		UserPrompts: sdkUserPrompts,
 		ToolUsages:  sdkToolUsages,
 	}
+	if interception.APIKeyID.Valid {
+		intc.APIKeyID = &interception.APIKeyID.String
+	}
 	if interception.EndedAt.Valid {
 		intc.EndedAt = &interception.EndedAt.Time
 	}

diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go
@@ -1496,6 +1496,7 @@ func ClaimPrebuild(
 func AIBridgeInterception(t testing.TB, db database.Store, seed database.InsertAIBridgeInterceptionParams, endedAt *time.Time) database.AIBridgeInterception {
 	interception, err := db.InsertAIBridgeInterception(genCtx, database.InsertAIBridgeInterceptionParams{
 		ID:          takeFirst(seed.ID, uuid.New()),
+		APIKeyID:    seed.APIKeyID,
 		InitiatorID: takeFirst(seed.InitiatorID, uuid.New()),
 		Provider:    takeFirst(seed.Provider, "provider"),
 		Model:       takeFirst(seed.Model, "model"),

diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql
diff --git a/coderd/database/migrations/000396_add_aibridge_interceptions_api_key_id.down.sql b/coderd/database/migrations/000396_add_aibridge_interceptions_api_key_id.down.sql
@@ -0,0 +1 @@
+ALTER TABLE aibridge_interceptions DROP COLUMN api_key_id;
diff --git a/coderd/database/migrations/000396_add_aibridge_interceptions_api_key_id.up.sql b/coderd/database/migrations/000396_add_aibridge_interceptions_api_key_id.up.sql
@@ -0,0 +1,2 @@
+  -- column is nullable to not break interceptions recorded before this column was added
+ALTER TABLE aibridge_interceptions ADD COLUMN api_key_id text;
diff --git a/coderd/database/modelqueries.go b/coderd/database/modelqueries.go
@@ -805,6 +805,7 @@ func (q *sqlQuerier) ListAuthorizedAIBridgeInterceptions(ctx context.Context, ar
 			&i.AIBridgeInterception.StartedAt,
 			&i.AIBridgeInterception.Metadata,
 			&i.AIBridgeInterception.EndedAt,
+			&i.AIBridgeInterception.APIKeyID,
 			&i.VisibleUser.ID,
 			&i.VisibleUser.Username,
 			&i.VisibleUser.Name,

diff --git a/coderd/database/models.go b/coderd/database/models.go
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/aibridge.sql b/coderd/database/queries/aibridge.sql
@@ -1,8 +1,8 @@
 -- name: InsertAIBridgeInterception :one
 INSERT INTO aibridge_interceptions (
-	id, initiator_id, provider, model, metadata, started_at
+	id, api_key_id, initiator_id, provider, model, metadata, started_at
 ) VALUES (
-	@id, @initiator_id, @provider, @model, COALESCE(@metadata::jsonb, '{}'::jsonb), @started_at
+	@id, @api_key_id, @initiator_id, @provider, @model, COALESCE(@metadata::jsonb, '{}'::jsonb), @started_at
 )
 RETURNING *;
 

diff --git a/codersdk/aibridge.go b/codersdk/aibridge.go
@@ -13,6 +13,7 @@ import (
 
 type AIBridgeInterception struct {
 	ID          uuid.UUID            `json:"id" format:"uuid"`
+	APIKeyID    *string              `json:"api_key_id"`
 	Initiator   MinimalUser          `json:"initiator"`
 	Provider    string               `json:"provider"`
 	Model       string               `json:"model"`

diff --git a/docs/reference/api/aibridge.md b/docs/reference/api/aibridge.md
diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md
diff --git a/enterprise/aibridged/aibridged_integration_test.go b/enterprise/aibridged/aibridged_integration_test.go
@@ -18,6 +18,7 @@ import (
 	"github.com/coder/coder/v2/coderd/database/dbtestutil"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
 	"github.com/coder/coder/v2/coderd/externalauth"
+	"github.com/coder/coder/v2/coderd/httpmw"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/enterprise/aibridged"
 	"github.com/coder/coder/v2/enterprise/coderd/coderdenttest"
@@ -221,6 +222,18 @@ func TestIntegration(t *testing.T) {
 	require.NoError(t, err)
 	require.Len(t, interceptions, 1)
 
+	intc0 := interceptions[0]
+	keyID, _, err := httpmw.SplitAPIToken(apiKey.Key)
+	require.NoError(t, err)
+	require.Equal(t, user.ID, intc0.InitiatorID)
+	require.True(t, intc0.APIKeyID.Valid)
+	require.Equal(t, keyID, intc0.APIKeyID.String)
+	require.Equal(t, "openai", intc0.Provider)
+	require.Equal(t, "gpt-4.1", intc0.Model)
+	require.True(t, intc0.EndedAt.Valid)
+	require.True(t, intc0.StartedAt.Before(intc0.EndedAt.Time))
+	require.Less(t, intc0.EndedAt.Time.Sub(intc0.StartedAt), 5*time.Second)
+
 	prompts, err := db.GetAIBridgeUserPromptsByInterceptionID(ctx, interceptions[0].ID)
 	require.NoError(t, err)
 	require.Len(t, prompts, 1)

diff --git a/enterprise/aibridged/http.go b/enterprise/aibridged/http.go
@@ -68,6 +68,7 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
 
 	handler, err := s.GetRequestHandler(ctx, Request{
 		SessionKey:  key,
+		APIKeyID:    resp.ApiKeyId,
 		InitiatorID: id,
 	})
 	if err != nil {

diff --git a/enterprise/aibridged/pool.go b/enterprise/aibridged/pool.go
@@ -111,17 +111,9 @@ func (p *CachedBridgePool) Acquire(ctx context.Context, req Request, clientFn Cl
 	// may visit the slow path unnecessarily.
 	defer p.cache.Wait()
 
-	recorder := aibridge.NewRecorder(p.logger.Named("recorder"), func() (aibridge.Recorder, error) {
-		client, err := clientFn()
-		if err != nil {
-			return nil, xerrors.Errorf("acquire client: %w", err)
-		}
-
-		return &recorderTranslation{client: client}, nil
-	})
-
 	// Fast path.
-	bridge, ok := p.cache.Get(req.InitiatorID.String())
+	cacheKey := req.InitiatorID.String() + "|" + req.APIKeyID
+	bridge, ok := p.cache.Get(cacheKey)
 	if ok && bridge != nil {
 		// TODO: future improvement:
 		// Once we can detect token expiry against an MCP server, we no longer need to let these instances
@@ -131,6 +123,15 @@ func (p *CachedBridgePool) Acquire(ctx context.Context, req Request, clientFn Cl
 		return bridge, nil
 	}
 
+	recorder := aibridge.NewRecorder(p.logger.Named("recorder"), func() (aibridge.Recorder, error) {
+		client, err := clientFn()
+		if err != nil {
+			return nil, xerrors.Errorf("acquire client: %w", err)
+		}
+
+		return &recorderTranslation{apiKeyID: req.APIKeyID, client: client}, nil
+	})
+
 	// Slow path.
 	// Creating an *aibridge.RequestBridge may take some time, so gate all subsequent callers behind the initial request and return the resulting value.
 	// TODO: track startup time since it adds latency to first request (histogram count will also help us see how often this occurs).
@@ -158,7 +159,7 @@ func (p *CachedBridgePool) Acquire(ctx context.Context, req Request, clientFn Cl
 			return nil, xerrors.Errorf("create new request bridge: %w", err)
 		}
 
-		p.cache.SetWithTTL(req.InitiatorID.String(), bridge, cacheCost, p.options.TTL)
+		p.cache.SetWithTTL(cacheKey, bridge, cacheCost, p.options.TTL)
 
 		return bridge, nil
 	})

diff --git a/enterprise/aibridged/pool_test.go b/enterprise/aibridged/pool_test.go
@@ -34,7 +34,7 @@ func TestPool(t *testing.T) {
 	require.NoError(t, err)
 	t.Cleanup(func() { pool.Shutdown(context.Background()) })
 
-	id, id2 := uuid.New(), uuid.New()
+	id, id2, apiKeyID1, apiKeyID2 := uuid.New(), uuid.New(), uuid.New(), uuid.New()
 	clientFn := func() (aibridged.DRPCClient, error) {
 		return client, nil
 	}
@@ -50,13 +50,15 @@ func TestPool(t *testing.T) {
 	inst, err := pool.Acquire(t.Context(), aibridged.Request{
 		SessionKey:  "key",
 		InitiatorID: id,
+		APIKeyID:    apiKeyID1.String(),
 	}, clientFn, newMockMCPFactory(mcpProxy))
 	require.NoError(t, err, "acquire pool instance")
 
 	// ...and it will return it when acquired again.
 	instB, err := pool.Acquire(t.Context(), aibridged.Request{
 		SessionKey:  "key",
 		InitiatorID: id,
+		APIKeyID:    apiKeyID1.String(),
 	}, clientFn, newMockMCPFactory(mcpProxy))
 	require.NoError(t, err, "acquire pool instance")
 	require.Same(t, inst, instB)
@@ -74,6 +76,7 @@ func TestPool(t *testing.T) {
 	inst2, err := pool.Acquire(t.Context(), aibridged.Request{
 		SessionKey:  "key",
 		InitiatorID: id2,
+		APIKeyID:    apiKeyID1.String(),
 	}, clientFn, newMockMCPFactory(mcpProxy))
 	require.NoError(t, err, "acquire pool instance")
 	require.NotSame(t, inst, inst2)
@@ -84,6 +87,24 @@ func TestPool(t *testing.T) {
 	require.EqualValues(t, 1, metrics.Hits())
 	require.EqualValues(t, 2, metrics.Misses())
 
+	// This will get called again because a new instance will be created.
+	mcpProxy.EXPECT().Init(gomock.Any()).Times(1).Return(nil)
+
+	// New instance is created for different api key id
+	inst2B, err := pool.Acquire(t.Context(), aibridged.Request{
+		SessionKey:  "key",
+		InitiatorID: id2,
+		APIKeyID:    apiKeyID2.String(),
+	}, clientFn, newMockMCPFactory(mcpProxy))
+	require.NoError(t, err, "acquire pool instance 2B")
+	require.NotSame(t, inst2, inst2B)
+
+	metrics = pool.Metrics()
+	require.EqualValues(t, 3, metrics.KeysAdded())
+	require.EqualValues(t, 2, metrics.KeysEvicted())
+	require.EqualValues(t, 1, metrics.Hits())
+	require.EqualValues(t, 3, metrics.Misses())
+
 	// TODO: add test for expiry.
 	// This requires Go 1.25's [synctest](https://pkg.go.dev/testing/synctest) since the
 	// internal cache lib cannot be tested using coder/quartz.
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		ALTER TABLE aibridge_interceptions DROP COLUMN api_key_id;
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		-- column is nullable to not break interceptions recorded before this column was added
		ALTER TABLE aibridge_interceptions ADD COLUMN api_key_id text;
pawbana marked this conversation as resolved. Show resolved Hide resolved