Thanks to visit codestin.com
Credit goes to github.com

Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions coderd/aibridgedserver/aibridgedserver.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
}

Expand Down
72 changes: 72 additions & 0 deletions coderd/aibridgedserver/aibridgedserver_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import (
"net"
"net/url"
"strconv"
"sync/atomic"
"testing"
"time"

Expand Down Expand Up @@ -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())
})
}
}
4 changes: 4 additions & 0 deletions coderd/apidoc/docs.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 4 additions & 0 deletions coderd/apidoc/swagger.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 4 additions & 0 deletions codersdk/deployment.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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:
Expand All @@ -5313,6 +5316,7 @@ var ExperimentsKnown = Experiments{
ExperimentWorkspaceBuildUpdates,
ExperimentMinimumImplicitMember,
ExperimentWorkspaceCapableLicensing,
ExperimentAIGatewaySeatExclusion,
ExperimentAIGatewayCostControl,
ExperimentChatAdvisor,
ExperimentChatVirtualDesktop,
Expand Down
6 changes: 3 additions & 3 deletions docs/reference/api/schemas.md

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions site/src/api/typesGenerated.ts

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading