Oct 27, 2025 · Oct 28, 2025 · Oct 28, 2025 · Oct 28, 2025 · Oct 28, 2025 · Oct 29, 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
 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
 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/000394_add_aibridge_interceptions_api_key_id.down.sql b/coderd/database/migrations/000394_add_aibridge_interceptions_api_key_id.down.sql
 ALTER TABLE aibridge_interceptions DROP COLUMN api_key_id;
diff --git a/coderd/database/migrations/000394_add_aibridge_interceptions_api_key_id.up.sql b/coderd/database/migrations/000394_add_aibridge_interceptions_api_key_id.up.sql
  -- 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
 &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
 -- 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

 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
 "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"
 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

 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
 // 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
 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).
 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
 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
 }
 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)
 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)
 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
Expand Up		@@ -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
		}
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -1498,6 +1498,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"),
Expand Down
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;
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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,
Expand Down
Original file line number	Diff line number	Diff line change
		@@ -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*;

Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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"`
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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"
Expand DownExpand Up		@@ -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)
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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 {
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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
Expand All		@@ -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).
Expand DownExpand Up		@@ -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
		})
Expand Down