From 2b14e56b9b454d40d1d78ba891d6e4ec50b68897 Mon Sep 17 00:00:00 2001 From: Scott Miller Date: Mon, 27 Jul 2026 14:45:14 +0000 Subject: [PATCH] feat: skip AI Governance seat recording for AI Bridge usage under a new experiment Under the ai-gateway-seat-exclusion experiment, RecordInterception no longer records ai_seat_state usage for the initiator: AI Gateway access is licensed by the AI Governance add-on rather than per seat. Task workspace builds still claim AI Governance seats. --- coderd/aibridgedserver/aibridgedserver.go | 8 ++- .../aibridgedserver/aibridgedserver_test.go | 72 +++++++++++++++++++ coderd/apidoc/docs.go | 4 ++ coderd/apidoc/swagger.json | 4 ++ codersdk/deployment.go | 4 ++ docs/reference/api/schemas.md | 6 +- site/src/api/typesGenerated.ts | 2 + 7 files changed, 95 insertions(+), 5 deletions(-) diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index f2dee9a6b3b4c..55f50d4fcc6de 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -115,6 +115,7 @@ type Server struct { coderMCPConfig *proto.MCPServerConfig // may be nil if not available structuredLogging bool aiSeatTracker aiseats.SeatTracker + experiments codersdk.Experiments // budgetPolicy selects the effective group when a user belongs to multiple // budgeted groups, used for cost attribution on token usage records. budgetPolicy codersdk.AIBudgetPolicy @@ -167,6 +168,7 @@ func NewServer(lifecycleCtx context.Context, opts Options) (*Server, error) { externalAuthConfigs: eac, structuredLogging: opts.GatewayCfg.StructuredLogging.Value(), aiSeatTracker: opts.AISeatTracker, + experiments: opts.Experiments, budgetPolicy: codersdk.NewAIBudgetPolicyFromString(opts.GatewayCfg.BudgetPolicy), budgetPeriod: codersdk.NewAIBudgetPeriodFromString(opts.GatewayCfg.BudgetPeriod), clock: opts.Clock, @@ -269,8 +271,10 @@ func (s *Server) RecordInterception(ctx context.Context, in *proto.RecordInterce return nil, xerrors.Errorf("start interception: %w", err) } - reason := aiseats.ReasonAIBridge("provider=" + in.Provider + ", model=" + in.Model) - s.aiSeatTracker.RecordUsage(ctx, initID, reason) + if !s.experiments.Enabled(codersdk.ExperimentAIGatewaySeatExclusion) { + reason := aiseats.ReasonAIBridge("provider=" + in.Provider + ", model=" + in.Model) + s.aiSeatTracker.RecordUsage(ctx, initID, reason) + } return &proto.RecordInterceptionResponse{}, nil } diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go index c3940583f9d37..203cfeb5fa85c 100644 --- a/coderd/aibridgedserver/aibridgedserver_test.go +++ b/coderd/aibridgedserver/aibridgedserver_test.go @@ -10,6 +10,7 @@ import ( "net" "net/url" "strconv" + "sync/atomic" "testing" "time" @@ -4141,3 +4142,74 @@ func (*fakeWatchProvidersStream) MsgSend(drpc.Message, drpc.Encoding) error { re func (*fakeWatchProvidersStream) MsgRecv(drpc.Message, drpc.Encoding) error { return nil } func (*fakeWatchProvidersStream) CloseSend() error { return nil } func (*fakeWatchProvidersStream) Close() error { return nil } + +// countingSeatTracker records the number of RecordUsage calls. +type countingSeatTracker struct { + calls atomic.Int64 +} + +func (c *countingSeatTracker) RecordUsage(context.Context, uuid.UUID, agplaiseats.Reason) { + c.calls.Add(1) +} + +// TestRecordInterceptionAISeat verifies that bridge usage claims an AI +// Governance seat only when the seat exclusion experiment is disabled. +func TestRecordInterceptionAISeat(t *testing.T) { + t.Parallel() + + newRequest := func() *proto.RecordInterceptionRequest { + return &proto.RecordInterceptionRequest{ + Id: uuid.NewString(), + ApiKeyId: uuid.NewString(), + InitiatorId: uuid.NewString(), + Provider: "anthropic", + Model: "claude-4-opus", + StartedAt: timestamppb.Now(), + } + } + + cases := []struct { + name string + experiments []codersdk.Experiment + expectedCalls int64 + }{ + { + name: "experiment off records a seat", + experiments: requiredExperiments, + expectedCalls: 1, + }, + { + name: "seat exclusion skips the seat", + experiments: append([]codersdk.Experiment{codersdk.ExperimentAIGatewaySeatExclusion}, requiredExperiments...), + expectedCalls: 0, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + db.EXPECT().InsertAIBridgeInterception(gomock.Any(), gomock.Any()). + Return(database.AIBridgeInterception{}, nil) + + tracker := &countingSeatTracker{} + ctx := testutil.Context(t, testutil.WaitLong) + srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{ + Store: db, + AISeatTracker: tracker, + AccessURL: "/", + GatewayCfg: codersdk.AIBridgeConfig{}, + Experiments: tc.experiments, + Logger: testutil.Logger(t), + Clock: quartz.NewReal(), + }) + require.NoError(t, err) + + _, err = srv.RecordInterception(ctx, newRequest()) + require.NoError(t, err) + require.Equal(t, tc.expectedCalls, tracker.calls.Load()) + }) + } +} diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 1f0454cabddf1..7be7b900dd586 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -20126,12 +20126,14 @@ const docTemplate = `{ "nats_pubsub", "minimum-implicit-member", "workspace-capable-licensing", + "ai-gateway-seat-exclusion", "ai-gateway-cost-control", "chat-advisor", "chat-virtual-desktop" ], "x-enum-comments": { "ExperimentAIGatewayCostControl": "Enables AI Gateway cost control functionality.", + "ExperimentAIGatewaySeatExclusion": "Excludes AI Gateway (AI Bridge) usage from AI Governance seat consumption.", "ExperimentAutoFillParameters": "This should not be taken out of experiments until we have redesigned the feature.", "ExperimentChatAdvisor": "Enables the advisor tool for root agent chats.", "ExperimentChatVirtualDesktop": "Enables virtual desktop and computer use provider for agents.", @@ -20156,6 +20158,7 @@ const docTemplate = `{ "Enables embedded NATS pubsub.", "Allows organizations to deviate from the default organization-member roles, in support of Gateway Accounts.", "Counts only users holding the workspace-create permission toward the license seat limit.", + "Excludes AI Gateway (AI Bridge) usage from AI Governance seat consumption.", "Enables AI Gateway cost control functionality.", "Enables the advisor tool for root agent chats.", "Enables virtual desktop and computer use provider for agents." @@ -20171,6 +20174,7 @@ const docTemplate = `{ "ExperimentNATSPubsub", "ExperimentMinimumImplicitMember", "ExperimentWorkspaceCapableLicensing", + "ExperimentAIGatewaySeatExclusion", "ExperimentAIGatewayCostControl", "ExperimentChatAdvisor", "ExperimentChatVirtualDesktop" diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index adc9f99cafbe5..eeb9d48059344 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -18278,12 +18278,14 @@ "nats_pubsub", "minimum-implicit-member", "workspace-capable-licensing", + "ai-gateway-seat-exclusion", "ai-gateway-cost-control", "chat-advisor", "chat-virtual-desktop" ], "x-enum-comments": { "ExperimentAIGatewayCostControl": "Enables AI Gateway cost control functionality.", + "ExperimentAIGatewaySeatExclusion": "Excludes AI Gateway (AI Bridge) usage from AI Governance seat consumption.", "ExperimentAutoFillParameters": "This should not be taken out of experiments until we have redesigned the feature.", "ExperimentChatAdvisor": "Enables the advisor tool for root agent chats.", "ExperimentChatVirtualDesktop": "Enables virtual desktop and computer use provider for agents.", @@ -18308,6 +18310,7 @@ "Enables embedded NATS pubsub.", "Allows organizations to deviate from the default organization-member roles, in support of Gateway Accounts.", "Counts only users holding the workspace-create permission toward the license seat limit.", + "Excludes AI Gateway (AI Bridge) usage from AI Governance seat consumption.", "Enables AI Gateway cost control functionality.", "Enables the advisor tool for root agent chats.", "Enables virtual desktop and computer use provider for agents." @@ -18323,6 +18326,7 @@ "ExperimentNATSPubsub", "ExperimentMinimumImplicitMember", "ExperimentWorkspaceCapableLicensing", + "ExperimentAIGatewaySeatExclusion", "ExperimentAIGatewayCostControl", "ExperimentChatAdvisor", "ExperimentChatVirtualDesktop" diff --git a/codersdk/deployment.go b/codersdk/deployment.go index 4fbc7cc388e6d..f5430fc2faee4 100644 --- a/codersdk/deployment.go +++ b/codersdk/deployment.go @@ -5260,6 +5260,7 @@ const ( ExperimentNATSPubsub Experiment = "nats_pubsub" // Enables embedded NATS pubsub. ExperimentMinimumImplicitMember Experiment = "minimum-implicit-member" // Allows organizations to deviate from the default organization-member roles, in support of Gateway Accounts. ExperimentWorkspaceCapableLicensing Experiment = "workspace-capable-licensing" // Counts only users holding the workspace-create permission toward the license seat limit. + ExperimentAIGatewaySeatExclusion Experiment = "ai-gateway-seat-exclusion" // Excludes AI Gateway (AI Bridge) usage from AI Governance seat consumption. ExperimentAIGatewayCostControl Experiment = "ai-gateway-cost-control" // Enables AI Gateway cost control functionality. ExperimentChatAdvisor Experiment = "chat-advisor" // Enables the advisor tool for root agent chats. ExperimentChatVirtualDesktop Experiment = "chat-virtual-desktop" // Enables virtual desktop and computer use provider for agents. @@ -5287,6 +5288,8 @@ func (e Experiment) DisplayName() string { return "Gateway Accounts (minimum implicit member)" case ExperimentWorkspaceCapableLicensing: return "Workspace-Capable Licensing" + case ExperimentAIGatewaySeatExclusion: + return "AI Gateway Seat Exclusion" case ExperimentAIGatewayCostControl: return "AI Gateway Cost Control" case ExperimentChatAdvisor: @@ -5313,6 +5316,7 @@ var ExperimentsKnown = Experiments{ ExperimentWorkspaceBuildUpdates, ExperimentMinimumImplicitMember, ExperimentWorkspaceCapableLicensing, + ExperimentAIGatewaySeatExclusion, ExperimentAIGatewayCostControl, ExperimentChatAdvisor, ExperimentChatVirtualDesktop, diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index 8c61cab1468f7..582c2051ee600 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -7242,9 +7242,9 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o #### Enumerated Values -| Value(s) | -|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `ai-gateway-cost-control`, `auto-fill-parameters`, `chat-advisor`, `chat-virtual-desktop`, `example`, `mcp-server-http`, `minimum-implicit-member`, `nats_pubsub`, `notifications`, `oauth2`, `workspace-build-updates`, `workspace-capable-licensing`, `workspace-usage` | +| Value(s) | +|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `ai-gateway-cost-control`, `ai-gateway-seat-exclusion`, `auto-fill-parameters`, `chat-advisor`, `chat-virtual-desktop`, `example`, `mcp-server-http`, `minimum-implicit-member`, `nats_pubsub`, `notifications`, `oauth2`, `workspace-build-updates`, `workspace-capable-licensing`, `workspace-usage` | ## codersdk.ExternalAPIKeyScopes diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index df238f98d8cf9..4567c34e07849 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -4832,6 +4832,7 @@ export const EntitlementsWarningHeader = "X-Coder-Entitlements-Warning"; // From codersdk/deployment.go export type Experiment = | "ai-gateway-cost-control" + | "ai-gateway-seat-exclusion" | "auto-fill-parameters" | "chat-advisor" | "chat-virtual-desktop" @@ -4847,6 +4848,7 @@ export type Experiment = export const Experiments: Experiment[] = [ "ai-gateway-cost-control", + "ai-gateway-seat-exclusion", "auto-fill-parameters", "chat-advisor", "chat-virtual-desktop",