diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index f2dee9a6b3b..55f50d4fcc6 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 c3940583f9d..203cfeb5fa8 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 1f0454cabdd..7be7b900dd5 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 adc9f99cafb..eeb9d480593 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 4fbc7cc388e..f5430fc2fae 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 8c61cab1468..582c2051ee6 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 df238f98d8c..4567c34e078 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",