mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: exclude AI Bridge usage from AI Governance seat counting (#27280)
Under the new `ai-gateway-seat-exclusion` experiment, AI Bridge usage stops counting toward AI Governance seats. ## Seat recording Under the 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. This experiment is independent of `workspace-capable-licensing` (#27279) so the two licensing behaviors can be enabled separately. Task workspace builds still claim AI Governance seats. ## Manual verification Verified live on a dev deployment (provider chained to dev.coder.com's gateway, model `gpt-5.6-luna`): with the experiment off, the first bridge request from each identity type (admin, plain member, service account) wrote an `ai_seat_state` row (`aibridge` reason); with it on, requests recorded interceptions but left seat state untouched — no new rows, and existing rows' `last_used_at` did not advance. Part of the gateway-accounts feature. ## Stack Part 2 of the gateway-accounts stack: 1. **#27279**: permission-based license seat counting. Behind the `workspace-capable-licensing` experiment and gated on the AI Governance add-on, `user_limit` counts only users the RBAC engine authorizes to create workspaces. 2. **This PR**: stops AI Bridge usage from claiming AI Governance seats under the new `ai-gateway-seat-exclusion` experiment. 3. ~~**#27281**: adds a `use_shared` capability precondition for workspace ACL grants, so workspace sharing is ineffective for (and rejected toward) users without workspace capabilities, evaluated live on every authorization.~~ This will be done in follow-up work when we have time to look into the performance impact. Related but independent: **#27278** hides the Workspaces page create CTAs for users without workspace-create permission.
This commit is contained in:
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
Generated
+4
@@ -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"
|
||||
|
||||
Generated
+4
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user