From ae7f0fbb74df61585df14ca062898bb8bdadf426 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Wed, 1 Jul 2026 12:08:11 +0000 Subject: [PATCH 1/6] feat(coderd): enforce ai budget on pre-request path --- coderd/aibridged/aibridged_test.go | 38 ++- coderd/aibridged/aibridgedmock/clientmock.go | 15 ++ coderd/aibridged/http.go | 41 ++- coderd/aibridged/proto/aibridged.proto | 22 ++ coderd/aibridged/proto/aibridged_drpc.pb.go | 42 ++- coderd/aibridgedserver/aibridgedserver.go | 93 +++++++ .../aibridgedserver/aibridgedserver_test.go | 255 ++++++++++++++++++ coderd/database/dbtime/dbtime.go | 6 + 8 files changed, 503 insertions(+), 9 deletions(-) diff --git a/coderd/aibridged/aibridged_test.go b/coderd/aibridged/aibridged_test.go index e1fd056075372..24481ee4255e6 100644 --- a/coderd/aibridged/aibridged_test.go +++ b/coderd/aibridged/aibridged_test.go @@ -128,12 +128,40 @@ func TestServeHTTP_FailureModes(t *testing.T) { // TODO: coderd connection-related failures. + // Budget-related failures. + { + name: "budget exceeded", + applyMocksFn: func(client *mock.MockDRPCClient, _ *mock.MockPooler) { + // Authorization passes. + client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId: uuid.NewString()}, nil) + client.EXPECT().GetUserAISpendStatus(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.GetUserAISpendStatusResponse{ + Exceeded: true, + EffectiveGroupId: uuid.NewString(), + SpendLimitMicros: 1_000, + CurrentSpendMicros: 1_500, + }, nil) + }, + expectedErr: aibridged.ErrBudgetExceeded, + expectedStatus: http.StatusPaymentRequired, + }, + { + name: "budget check failed", + applyMocksFn: func(client *mock.MockDRPCClient, _ *mock.MockPooler) { + // Authorization passes. + client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId: uuid.NewString()}, nil) + client.EXPECT().GetUserAISpendStatus(gomock.Any(), gomock.Any()).AnyTimes().Return(nil, xerrors.New("oops")) + }, + expectedErr: aibridged.ErrInternalServerError, + expectedStatus: http.StatusInternalServerError, + }, + // Pool-related failures. { name: "pool instance", applyMocksFn: func(client *mock.MockDRPCClient, pool *mock.MockPooler) { - // Should pass authorization. + // Should pass authorization and budget check. client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId: uuid.NewString()}, nil) + client.EXPECT().GetUserAISpendStatus(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.GetUserAISpendStatusResponse{}, nil) // But fail when acquiring a pool instance. pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().Return(nil, xerrors.New("oops")) }, @@ -223,6 +251,7 @@ func TestServeHTTP_DelegatedAPIKey(t *testing.T) { Username: "u", }, nil }) + client.EXPECT().GetUserAISpendStatus(gomock.Any(), gomock.Any()).Return(&proto.GetUserAISpendStatusResponse{}, nil) pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( func(_ context.Context, req aibridged.Request, _ aibridged.ClientFunc, _ aibridged.MCPProxyBuilder) (http.Handler, error) { assert.Empty(t, req.SessionKey, @@ -255,6 +284,7 @@ func TestServeHTTP_DelegatedAPIKey(t *testing.T) { ApiKeyId: testKeyID, Username: "u", }, nil) + client.EXPECT().GetUserAISpendStatus(gomock.Any(), gomock.Any()).Return(&proto.GetUserAISpendStatusResponse{}, nil) pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( func(_ context.Context, req aibridged.Request, _ aibridged.ClientFunc, _ aibridged.MCPProxyBuilder) (http.Handler, error) { assert.Equal(t, "coder-token-byok", req.SessionKey, @@ -345,6 +375,7 @@ func TestServeHTTP_DelegatedAPIKey_BYOK_Integration(t *testing.T) { Username: "u", }, nil }) + client.EXPECT().GetUserAISpendStatus(gomock.Any(), gomock.Any()).Return(&proto.GetUserAISpendStatusResponse{}, nil) pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(mockH, nil) factory := aibridged.NewTransportFactory(srv) @@ -397,6 +428,7 @@ func TestServeHTTP_DelegatedAPIKey_Integration(t *testing.T) { Username: "u", }, nil }) + client.EXPECT().GetUserAISpendStatus(gomock.Any(), gomock.Any()).Return(&proto.GetUserAISpendStatusResponse{}, nil) pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(mockH, nil) factory := aibridged.NewTransportFactory(srv) @@ -483,6 +515,7 @@ func TestServeHTTP_StripCoderToken(t *testing.T) { conn := &mockDRPCConn{} client.EXPECT().DRPCConn().AnyTimes().Return(conn) client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId: uuid.NewString()}, nil) + client.EXPECT().GetUserAISpendStatus(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.GetUserAISpendStatusResponse{}, nil) pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().Return(mockH, nil) httpSrv := httptest.NewServer(srv) @@ -677,6 +710,7 @@ func TestServeHTTP_ActorHeaders(t *testing.T) { OwnerId: testUserID.String(), Username: testUsername, }, nil) + client.EXPECT().GetUserAISpendStatus(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.GetUserAISpendStatusResponse{}, nil) client.EXPECT().GetMCPServerConfigs(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.GetMCPServerConfigsResponse{}, nil) client.EXPECT().RecordInterception(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.RecordInterceptionResponse{}, nil) client.EXPECT().RecordInterceptionEnded(gomock.Any(), gomock.Any()).AnyTimes() @@ -775,6 +809,7 @@ func TestRouting(t *testing.T) { client.EXPECT().DRPCConn().AnyTimes().Return(conn) client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId: uuid.NewString()}, nil) + client.EXPECT().GetUserAISpendStatus(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.GetUserAISpendStatusResponse{}, nil) client.EXPECT().GetMCPServerConfigs(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.GetMCPServerConfigsResponse{}, nil) // This is the only recording we really care about in this test. This is called before the provider-specific logic processes // the incoming request, and anything beyond that is the responsibility of coder/aibridge to test. @@ -850,6 +885,7 @@ func TestServeHTTP_StripInternalHeaders(t *testing.T) { conn := &mockDRPCConn{} client.EXPECT().DRPCConn().AnyTimes().Return(conn) client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId: uuid.NewString()}, nil) + client.EXPECT().GetUserAISpendStatus(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.GetUserAISpendStatusResponse{}, nil) pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().Return(mockH, nil) httpSrv := httptest.NewServer(srv) diff --git a/coderd/aibridged/aibridgedmock/clientmock.go b/coderd/aibridged/aibridgedmock/clientmock.go index 2ae8d283e2f52..c036fae84e8e1 100644 --- a/coderd/aibridged/aibridgedmock/clientmock.go +++ b/coderd/aibridged/aibridgedmock/clientmock.go @@ -101,6 +101,21 @@ func (mr *MockDRPCClientMockRecorder) GetMCPServerConfigs(ctx, in any) *gomock.C return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetMCPServerConfigs", reflect.TypeOf((*MockDRPCClient)(nil).GetMCPServerConfigs), ctx, in) } +// GetUserAISpendStatus mocks base method. +func (m *MockDRPCClient) GetUserAISpendStatus(ctx context.Context, in *proto.GetUserAISpendStatusRequest) (*proto.GetUserAISpendStatusResponse, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetUserAISpendStatus", ctx, in) + ret0, _ := ret[0].(*proto.GetUserAISpendStatusResponse) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetUserAISpendStatus indicates an expected call of GetUserAISpendStatus. +func (mr *MockDRPCClientMockRecorder) GetUserAISpendStatus(ctx, in any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserAISpendStatus", reflect.TypeOf((*MockDRPCClient)(nil).GetUserAISpendStatus), ctx, in) +} + // IsAuthorized mocks base method. func (m *MockDRPCClient) IsAuthorized(ctx context.Context, in *proto.IsAuthorizedRequest) (*proto.IsAuthorizedResponse, error) { m.ctrl.T.Helper() diff --git a/coderd/aibridged/http.go b/coderd/aibridged/http.go index 640716f37a9c2..a263359875d8d 100644 --- a/coderd/aibridged/http.go +++ b/coderd/aibridged/http.go @@ -1,17 +1,20 @@ package aibridged import ( + "fmt" "net/http" "strings" "github.com/google/uuid" "golang.org/x/xerrors" + "google.golang.org/protobuf/types/known/timestamppb" "cdr.dev/slog/v3" "github.com/coder/coder/v2/aibridge" "github.com/coder/coder/v2/aibridge/recorder" agplaibridge "github.com/coder/coder/v2/coderd/aibridge" "github.com/coder/coder/v2/coderd/aibridged/proto" + "github.com/coder/coder/v2/coderd/database/dbtime" ) var _ http.Handler = &Server{} @@ -21,6 +24,8 @@ var ( ErrConnect = xerrors.New("could not connect to coderd") ErrUnauthorized = xerrors.New("unauthorized") ErrAcquireRequestHandler = xerrors.New("failed to acquire request handler") + ErrBudgetExceeded = xerrors.New("ai budget exceeded") + ErrInternalServerError = xerrors.New("internal server error") ) // ServeHTTP is the entrypoint for requests which will be intercepted by AI Bridge. @@ -135,6 +140,35 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) { return } + id, err := uuid.Parse(resp.GetOwnerId()) + if err != nil { + logger.Warn(ctx, "failed to parse user ID", slog.Error(err), slog.F("id", resp.GetOwnerId())) + http.Error(rw, ErrUnauthorized.Error(), http.StatusForbidden) + return + } + + // TODO: make the budget period configurable; monthly for now. + periodStart := dbtime.StartOfMonth(dbtime.Now().UTC()) + spendStatus, err := client.GetUserAISpendStatus(ctx, &proto.GetUserAISpendStatusRequest{ + UserId: id.String(), + PeriodStart: timestamppb.New(periodStart), + }) + if err != nil { + logger.Warn(ctx, "ai spend status check failed", slog.Error(err)) + http.Error(rw, ErrInternalServerError.Error(), http.StatusInternalServerError) + return + } + if spendStatus.GetExceeded() { + http.Error(rw, fmt.Sprintf( + "%s: spent %.2f USD of %.2f USD in group %s", + ErrBudgetExceeded.Error(), + float64(spendStatus.GetCurrentSpendMicros())/1_000_000, + float64(spendStatus.GetSpendLimitMicros())/1_000_000, + spendStatus.GetEffectiveGroupId(), + ), http.StatusPaymentRequired) + return + } + // Rewire request context to include actor. // // [NOTE] @@ -144,13 +178,6 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) { "Username": resp.GetUsername(), })) - id, err := uuid.Parse(resp.GetOwnerId()) - if err != nil { - logger.Warn(ctx, "failed to parse user ID", slog.Error(err), slog.F("id", resp.GetOwnerId())) - http.Error(rw, ErrUnauthorized.Error(), http.StatusForbidden) - return - } - handler, err := s.GetRequestHandler(ctx, Request{ SessionKey: key, APIKeyID: resp.ApiKeyId, diff --git a/coderd/aibridged/proto/aibridged.proto b/coderd/aibridged/proto/aibridged.proto index 8dc78e7a894eb..4152c1c729a8e 100644 --- a/coderd/aibridged/proto/aibridged.proto +++ b/coderd/aibridged/proto/aibridged.proto @@ -34,6 +34,9 @@ service Authorizer { // IsAuthorized validates that a given Coder key is valid and the user is authorized to use AI Bridge. // TODO: add authorization; currently only key validation takes place. rpc IsAuthorized(IsAuthorizedRequest) returns (IsAuthorizedResponse); + // GetUserAISpendStatus returns the user's AI spend status aggregated over + // [period_start, now]. + rpc GetUserAISpendStatus(GetUserAISpendStatusRequest) returns (GetUserAISpendStatusResponse); } // ProviderConfigurator serves AI provider configuration to embedded and @@ -168,6 +171,25 @@ message IsAuthorizedResponse { string username = 3; } +message GetUserAISpendStatusRequest { + string user_id = 1; // UUID + // The spend aggregation window is [period_start, now]. + google.protobuf.Timestamp period_start = 2; +} + +message GetUserAISpendStatusResponse { + // exceeded is true when the user's aggregated spend has reached the + // effective limit. False when no budget is configured for the user OR + // they are within their limit. + bool exceeded = 1; + string user_id = 2; // UUID + // effective_group_id is empty when no budget is configured for the user. + // spend_limit_micros and current_spend_micros are 0 in that case. + string effective_group_id = 3; + int64 spend_limit_micros = 4; + int64 current_spend_micros = 5; +} + message GetAIProvidersRequest {} message GetAIProvidersResponse { diff --git a/coderd/aibridged/proto/aibridged_drpc.pb.go b/coderd/aibridged/proto/aibridged_drpc.pb.go index 5939b888bb9db..093286db2988d 100644 --- a/coderd/aibridged/proto/aibridged_drpc.pb.go +++ b/coderd/aibridged/proto/aibridged_drpc.pb.go @@ -429,6 +429,7 @@ type DRPCAuthorizerClient interface { DRPCConn() drpc.Conn IsAuthorized(ctx context.Context, in *IsAuthorizedRequest) (*IsAuthorizedResponse, error) + GetUserAISpendStatus(ctx context.Context, in *GetUserAISpendStatusRequest) (*GetUserAISpendStatusResponse, error) } type drpcAuthorizerClient struct { @@ -450,8 +451,18 @@ func (c *drpcAuthorizerClient) IsAuthorized(ctx context.Context, in *IsAuthorize return out, nil } +func (c *drpcAuthorizerClient) GetUserAISpendStatus(ctx context.Context, in *GetUserAISpendStatusRequest) (*GetUserAISpendStatusResponse, error) { + out := new(GetUserAISpendStatusResponse) + err := c.cc.Invoke(ctx, "/proto.Authorizer/GetUserAISpendStatus", drpcEncoding_File_coderd_aibridged_proto_aibridged_proto{}, in, out) + if err != nil { + return nil, err + } + return out, nil +} + type DRPCAuthorizerServer interface { IsAuthorized(context.Context, *IsAuthorizedRequest) (*IsAuthorizedResponse, error) + GetUserAISpendStatus(context.Context, *GetUserAISpendStatusRequest) (*GetUserAISpendStatusResponse, error) } type DRPCAuthorizerUnimplementedServer struct{} @@ -460,9 +471,13 @@ func (s *DRPCAuthorizerUnimplementedServer) IsAuthorized(context.Context, *IsAut return nil, drpcerr.WithCode(errors.New("Unimplemented"), drpcerr.Unimplemented) } +func (s *DRPCAuthorizerUnimplementedServer) GetUserAISpendStatus(context.Context, *GetUserAISpendStatusRequest) (*GetUserAISpendStatusResponse, error) { + return nil, drpcerr.WithCode(errors.New("Unimplemented"), drpcerr.Unimplemented) +} + type DRPCAuthorizerDescription struct{} -func (DRPCAuthorizerDescription) NumMethods() int { return 1 } +func (DRPCAuthorizerDescription) NumMethods() int { return 2 } func (DRPCAuthorizerDescription) Method(n int) (string, drpc.Encoding, drpc.Receiver, interface{}, bool) { switch n { @@ -475,6 +490,15 @@ func (DRPCAuthorizerDescription) Method(n int) (string, drpc.Encoding, drpc.Rece in1.(*IsAuthorizedRequest), ) }, DRPCAuthorizerServer.IsAuthorized, true + case 1: + return "/proto.Authorizer/GetUserAISpendStatus", drpcEncoding_File_coderd_aibridged_proto_aibridged_proto{}, + func(srv interface{}, ctx context.Context, in1, in2 interface{}) (drpc.Message, error) { + return srv.(DRPCAuthorizerServer). + GetUserAISpendStatus( + ctx, + in1.(*GetUserAISpendStatusRequest), + ) + }, DRPCAuthorizerServer.GetUserAISpendStatus, true default: return "", nil, nil, nil, false } @@ -500,6 +524,22 @@ func (x *drpcAuthorizer_IsAuthorizedStream) SendAndClose(m *IsAuthorizedResponse return x.CloseSend() } +type DRPCAuthorizer_GetUserAISpendStatusStream interface { + drpc.Stream + SendAndClose(*GetUserAISpendStatusResponse) error +} + +type drpcAuthorizer_GetUserAISpendStatusStream struct { + drpc.Stream +} + +func (x *drpcAuthorizer_GetUserAISpendStatusStream) SendAndClose(m *GetUserAISpendStatusResponse) error { + if err := x.MsgSend(m, drpcEncoding_File_coderd_aibridged_proto_aibridged_proto{}); err != nil { + return err + } + return x.CloseSend() +} + type DRPCProviderConfiguratorClient interface { DRPCConn() drpc.Conn diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index 88fdc8e090012..8cae676ce3171 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -8,6 +8,7 @@ import ( "slices" "strings" "sync" + "time" "github.com/google/uuid" "github.com/hashicorp/go-multierror" @@ -16,6 +17,7 @@ import ( "google.golang.org/protobuf/types/known/structpb" "cdr.dev/slog/v3" + "github.com/coder/coder/v2/coderd/aibridge/budget" "github.com/coder/coder/v2/coderd/aibridged" "github.com/coder/coder/v2/coderd/aibridged/proto" "github.com/coder/coder/v2/coderd/aiseats" @@ -75,6 +77,7 @@ type store interface { GetAIModelPriceByProviderModel(ctx context.Context, arg database.GetAIModelPriceByProviderModelParams) (database.AIModelPrice, error) GetUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (database.UserAIBudgetOverride, error) GetHighestGroupAIBudgetByUser(ctx context.Context, userID uuid.UUID) (database.GetHighestGroupAIBudgetByUserRow, error) + GetUserAISpendSince(ctx context.Context, arg database.GetUserAISpendSinceParams) (database.GetUserAISpendSinceRow, error) // MCPConfigurator-related queries. GetExternalAuthLinksByUserID(ctx context.Context, userID uuid.UUID) ([]database.ExternalAuthLink, error) @@ -727,6 +730,96 @@ func (s *Server) IsAuthorized(ctx context.Context, in *proto.IsAuthorizedRequest }, nil } +// GetUserAISpendStatus returns the user's AI spend status aggregated over +// [period_start, now]. +func (s *Server) GetUserAISpendStatus(ctx context.Context, in *proto.GetUserAISpendStatusRequest) (*proto.GetUserAISpendStatusResponse, error) { + //nolint:gocritic // AIBridged has specific authz rules. + ctx = dbauthz.AsAIBridged(ctx) + + userID, err := uuid.Parse(in.GetUserId()) + if err != nil { + return nil, xerrors.Errorf("invalid user_id %q: %w", in.GetUserId(), err) + } + // An unset period_start deserializes to time.Unix(0, 0), which would + // aggregate the user's lifetime spend against a period budget. + if in.PeriodStart == nil { + return nil, xerrors.New("period_start is required") + } + periodStart := in.GetPeriodStart().AsTime() + + status, err := s.resolveUserAISpendStatus(ctx, userID, periodStart) + if err != nil { + return nil, err + } + resp := &proto.GetUserAISpendStatusResponse{ + Exceeded: status.Exceeded, + UserId: userID.String(), + SpendLimitMicros: status.SpendLimitMicros, + CurrentSpendMicros: status.CurrentSpendMicros, + } + if status.EffectiveGroupID != uuid.Nil { + resp.EffectiveGroupId = status.EffectiveGroupID.String() + } + return resp, nil +} + +// userAISpendStatus is a snapshot of a user's AI spend against their effective +// budget over [PeriodStart, now]. When no budget is configured for the user, +// all fields are zero-valued. +type userAISpendStatus struct { + EffectiveGroupID uuid.UUID + PeriodStart time.Time + SpendLimitMicros int64 + CurrentSpendMicros int64 + Exceeded bool +} + +// resolveUserAISpendStatus computes the user's AI spend status aggregated from +// periodStart. Returns a zero-valued status (EffectiveGroupID = uuid.Nil) when +// no budget applies to the user. +func (s *Server) resolveUserAISpendStatus(ctx context.Context, userID uuid.UUID, periodStart time.Time) (userAISpendStatus, error) { + effectiveBudget, ok, err := budget.ResolveUserAIBudget(ctx, s.store, userID, s.budgetPolicy) + if err != nil { + return userAISpendStatus{}, xerrors.Errorf("resolve effective AI budget for user %q with budget policy %q: %w", userID, s.budgetPolicy, err) + } + if !ok { + // No budget configured for the user; return zero-valued status. + return userAISpendStatus{}, nil + } + + spend, err := s.store.GetUserAISpendSince(ctx, database.GetUserAISpendSinceParams{ + UserID: userID, + EffectiveGroupID: effectiveBudget.GroupID, + PeriodStart: periodStart, + }) + if err != nil { + return userAISpendStatus{}, xerrors.Errorf("get user AI spend for user %q in group %q: %w", userID, effectiveBudget.GroupID, err) + } + + status := userAISpendStatus{ + EffectiveGroupID: effectiveBudget.GroupID, + PeriodStart: periodStart, + SpendLimitMicros: effectiveBudget.SpendLimitMicros, + CurrentSpendMicros: spend.SpendMicros, + Exceeded: spend.SpendMicros >= effectiveBudget.SpendLimitMicros, + } + + logger := s.logger.With( + slog.F("user_id", userID), + slog.F("effective_group_id", status.EffectiveGroupID), + slog.F("period_start", status.PeriodStart), + slog.F("current_spend_micros", status.CurrentSpendMicros), + slog.F("spend_limit_micros", status.SpendLimitMicros), + slog.F("exceeded", status.Exceeded), + ) + logger.Debug(ctx, "user AI spend status") + if status.Exceeded { + logger.Info(ctx, "user AI budget exceeded") + } + + return status, nil +} + // GetAIProviders returns the full AI provider set (enabled and disabled) from // the database, which is the single source of truth seeded from coderd's // environment. Embedded and standalone AI Gateway daemons call this over DRPC diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go index b886f56961537..769ecfc0f5cc3 100644 --- a/coderd/aibridgedserver/aibridgedserver_test.go +++ b/coderd/aibridgedserver/aibridgedserver_test.go @@ -391,6 +391,261 @@ func TestAuthorization_Delegated(t *testing.T) { } } +func TestGetUserAISpendStatus(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + userIDStr string + omitPeriodStart bool + setupMocks func(db *dbmock.MockStore, userID uuid.UUID) *proto.GetUserAISpendStatusResponse + wantErr bool + }{ + { + // Invalid UUID short-circuits before any store call. + name: "invalid user_id", + userIDStr: "not-a-uuid", + wantErr: true, + }, + { + // Missing period_start is rejected before any store call. + name: "missing period_start", + omitPeriodStart: true, + wantErr: true, + }, + { + // No override and no group budget resolves: pass-through. All + // spend/limit fields are zero-valued. + name: "no budget configured returns not exceeded", + setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.GetUserAISpendStatusResponse { + db.EXPECT().GetUserAIBudgetOverride(gomock.Any(), userID). + Return(database.UserAIBudgetOverride{}, sql.ErrNoRows) + db.EXPECT().GetHighestGroupAIBudgetByUser(gomock.Any(), userID). + Return(database.GetHighestGroupAIBudgetByUserRow{}, sql.ErrNoRows) + return &proto.GetUserAISpendStatusResponse{ + Exceeded: false, + UserId: userID.String(), + EffectiveGroupId: "", + SpendLimitMicros: 0, + CurrentSpendMicros: 0, + } + }, + }, + { + // Group budget resolves, spend below limit (spend 500 < limit 1000). + name: "under limit returns not exceeded", + setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.GetUserAISpendStatusResponse { + groupID := uuid.New() + db.EXPECT().GetUserAIBudgetOverride(gomock.Any(), userID). + Return(database.UserAIBudgetOverride{}, sql.ErrNoRows) + db.EXPECT().GetHighestGroupAIBudgetByUser(gomock.Any(), userID). + Return(database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 1_000}, nil) + db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()). + Return(database.GetUserAISpendSinceRow{SpendMicros: 500}, nil) + return &proto.GetUserAISpendStatusResponse{ + Exceeded: false, + UserId: userID.String(), + EffectiveGroupId: groupID.String(), + SpendLimitMicros: 1_000, + CurrentSpendMicros: 500, + } + }, + }, + { + // spend == limit is treated as exceeded (spend 1000 >= limit 1000). + name: "at limit returns exceeded", + setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.GetUserAISpendStatusResponse { + groupID := uuid.New() + db.EXPECT().GetUserAIBudgetOverride(gomock.Any(), userID). + Return(database.UserAIBudgetOverride{}, sql.ErrNoRows) + db.EXPECT().GetHighestGroupAIBudgetByUser(gomock.Any(), userID). + Return(database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 1_000}, nil) + db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()). + Return(database.GetUserAISpendSinceRow{SpendMicros: 1_000}, nil) + return &proto.GetUserAISpendStatusResponse{ + Exceeded: true, + UserId: userID.String(), + EffectiveGroupId: groupID.String(), + SpendLimitMicros: 1_000, + CurrentSpendMicros: 1_000, + } + }, + }, + { + // Spend above limit (spend 1500 > limit 1000). + name: "over limit returns exceeded", + setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.GetUserAISpendStatusResponse { + groupID := uuid.New() + db.EXPECT().GetUserAIBudgetOverride(gomock.Any(), userID). + Return(database.UserAIBudgetOverride{}, sql.ErrNoRows) + db.EXPECT().GetHighestGroupAIBudgetByUser(gomock.Any(), userID). + Return(database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 1_000}, nil) + db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()). + Return(database.GetUserAISpendSinceRow{SpendMicros: 1_500}, nil) + return &proto.GetUserAISpendStatusResponse{ + Exceeded: true, + UserId: userID.String(), + EffectiveGroupId: groupID.String(), + SpendLimitMicros: 1_000, + CurrentSpendMicros: 1_500, + } + }, + }, + { + // User override wins: group lookup is skipped, spend is aggregated + // against the override's group (spend 600 > limit 500). + name: "user override wins over group budget", + setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.GetUserAISpendStatusResponse { + overrideGroupID := uuid.New() + db.EXPECT().GetUserAIBudgetOverride(gomock.Any(), userID). + Return(database.UserAIBudgetOverride{ + UserID: userID, + GroupID: overrideGroupID, + SpendLimitMicros: 500, + }, nil) + db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Cond(func(p database.GetUserAISpendSinceParams) bool { + return assert.Equal(t, overrideGroupID, p.EffectiveGroupID, "spend aggregated against override group") + })).Return(database.GetUserAISpendSinceRow{SpendMicros: 600}, nil) + return &proto.GetUserAISpendStatusResponse{ + Exceeded: true, + UserId: userID.String(), + EffectiveGroupId: overrideGroupID.String(), + SpendLimitMicros: 500, + CurrentSpendMicros: 600, + } + }, + }, + { + // Unexpected error from budget override lookup propagates. + name: "budget resolution error propagates", + setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.GetUserAISpendStatusResponse { + db.EXPECT().GetUserAIBudgetOverride(gomock.Any(), userID). + Return(database.UserAIBudgetOverride{}, sql.ErrConnDone) + return nil + }, + wantErr: true, + }, + { + // Error from spend aggregation propagates (fail-closed). + name: "spend aggregation error propagates", + setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.GetUserAISpendStatusResponse { + db.EXPECT().GetUserAIBudgetOverride(gomock.Any(), userID). + Return(database.UserAIBudgetOverride{}, sql.ErrNoRows) + db.EXPECT().GetHighestGroupAIBudgetByUser(gomock.Any(), userID). + Return(database.GetHighestGroupAIBudgetByUserRow{GroupID: uuid.New(), SpendLimitMicros: 1_000}, nil) + db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()). + Return(database.GetUserAISpendSinceRow{}, sql.ErrConnDone) + return nil + }, + wantErr: true, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + logger := testutil.Logger(t) + + userID := uuid.New() + userIDStr := tc.userIDStr + if userIDStr == "" { + userIDStr = userID.String() + } + + var wantResp *proto.GetUserAISpendStatusResponse + if tc.setupMocks != nil { + wantResp = tc.setupMocks(db, userID) + } + + srv, err := aibridgedserver.NewServer(t.Context(), db, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) + require.NoError(t, err) + + req := &proto.GetUserAISpendStatusRequest{UserId: userIDStr} + if !tc.omitPeriodStart { + req.PeriodStart = timestamppb.New(dbtime.StartOfMonth(dbtime.Now().UTC())) + } + resp, err := srv.GetUserAISpendStatus(t.Context(), req) + if tc.wantErr { + require.Error(t, err) + require.Nil(t, resp) + return + } + require.NoError(t, err) + require.NotNil(t, resp) + require.Equal(t, wantResp.GetExceeded(), resp.GetExceeded(), "exceeded") + require.Equal(t, wantResp.GetUserId(), resp.GetUserId(), "user_id") + require.Equal(t, wantResp.GetEffectiveGroupId(), resp.GetEffectiveGroupId(), "effective_group_id") + require.Equal(t, wantResp.GetSpendLimitMicros(), resp.GetSpendLimitMicros(), "spend_limit_micros") + require.Equal(t, wantResp.GetCurrentSpendMicros(), resp.GetCurrentSpendMicros(), "current_spend_micros") + }) + } +} + +// TestGetUserAISpendStatus_PeriodBoundary exercises the aggregation window at +// a UTC month boundary. A spend row in January pushes the user over their +// limit; the same row is excluded when the query moves to the next period. +func TestGetUserAISpendStatus_PeriodBoundary(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + logger := testutil.Logger(t) + + rawDB, _ := dbtestutil.NewDB(t) + authzDB := dbauthz.New(rawDB, rbac.NewStrictAuthorizer(prometheus.NewRegistry()), logger, coderdtest.AccessControlStorePointer()) + + org := dbgen.Organization(t, rawDB, database.Organization{}) + user := dbgen.User(t, rawDB, database.User{}) + dbgen.OrganizationMember(t, rawDB, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID}) + group := dbgen.Group(t, rawDB, database.Group{OrganizationID: org.ID}) + dbgen.GroupMember(t, rawDB, database.GroupMemberTable{UserID: user.ID, GroupID: group.ID}) + + const limitMicros = 1_000_000 + _, err := rawDB.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ + GroupID: group.ID, + SpendLimitMicros: limitMicros, + }) + require.NoError(t, err) + + prevMonth := time.Date(2026, time.January, 1, 0, 0, 0, 0, time.UTC) + newMonth := time.Date(2026, time.February, 1, 0, 0, 0, 0, time.UTC) + + // Insert a spend row in the previous month that pushes the user over the limit. + _, err = rawDB.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, + EffectiveGroupID: group.ID, + Day: prevMonth.AddDate(0, 0, 14), + CostMicros: 1_500_000, + }) + require.NoError(t, err) + + srv, err := aibridgedserver.NewServer(ctx, authzDB, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) + require.NoError(t, err) + + // Same UTC month as the spend row: aggregation includes it, user exceeded. + prevMonthResp, err := srv.GetUserAISpendStatus(ctx, &proto.GetUserAISpendStatusRequest{ + UserId: user.ID.String(), + PeriodStart: timestamppb.New(prevMonth), + }) + require.NoError(t, err) + require.True(t, prevMonthResp.GetExceeded()) + require.Equal(t, int64(1_500_000), prevMonthResp.GetCurrentSpendMicros()) + require.Equal(t, int64(limitMicros), prevMonthResp.GetSpendLimitMicros()) + + // After the boundary: the previous month's row is outside the aggregation + // window and the user is no longer exceeded. + newMonthResp, err := srv.GetUserAISpendStatus(ctx, &proto.GetUserAISpendStatusRequest{ + UserId: user.ID.String(), + PeriodStart: timestamppb.New(newMonth), + }) + require.NoError(t, err) + require.False(t, newMonthResp.GetExceeded()) + require.Equal(t, int64(0), newMonthResp.GetCurrentSpendMicros()) + require.Equal(t, int64(limitMicros), newMonthResp.GetSpendLimitMicros()) +} + func TestGetMCPServerConfigs(t *testing.T) { t.Parallel() diff --git a/coderd/database/dbtime/dbtime.go b/coderd/database/dbtime/dbtime.go index bda5a2263ce2b..700a79abe794f 100644 --- a/coderd/database/dbtime/dbtime.go +++ b/coderd/database/dbtime/dbtime.go @@ -22,3 +22,9 @@ func StartOfDay(t time.Time) time.Time { year, month, day := t.Date() return time.Date(year, month, day, 0, 0, 0, 0, t.Location()) } + +// StartOfMonth returns the first timestamp of the month of the input timestamp in its location. +func StartOfMonth(t time.Time) time.Time { + year, month, _ := t.Date() + return time.Date(year, month, 1, 0, 0, 0, 0, t.Location()) +} From 4f3a9fef0a54a140e047a5940ccb67e1ef45bcf4 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Wed, 1 Jul 2026 18:37:00 +0000 Subject: [PATCH 2/6] chore: address comments --- coderd/aibridged/aibridged_test.go | 2 +- coderd/aibridged/http.go | 8 +- coderd/aibridgedserver/aibridgedserver.go | 16 +- .../aibridgedserver/aibridgedserver_test.go | 193 ++++++++++++------ 4 files changed, 153 insertions(+), 66 deletions(-) diff --git a/coderd/aibridged/aibridged_test.go b/coderd/aibridged/aibridged_test.go index 24481ee4255e6..f421580a6c6c0 100644 --- a/coderd/aibridged/aibridged_test.go +++ b/coderd/aibridged/aibridged_test.go @@ -151,7 +151,7 @@ func TestServeHTTP_FailureModes(t *testing.T) { client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId: uuid.NewString()}, nil) client.EXPECT().GetUserAISpendStatus(gomock.Any(), gomock.Any()).AnyTimes().Return(nil, xerrors.New("oops")) }, - expectedErr: aibridged.ErrInternalServerError, + expectedErr: aibridged.ErrSpendStatusCheck, expectedStatus: http.StatusInternalServerError, }, diff --git a/coderd/aibridged/http.go b/coderd/aibridged/http.go index a263359875d8d..fd15eef965244 100644 --- a/coderd/aibridged/http.go +++ b/coderd/aibridged/http.go @@ -25,7 +25,7 @@ var ( ErrUnauthorized = xerrors.New("unauthorized") ErrAcquireRequestHandler = xerrors.New("failed to acquire request handler") ErrBudgetExceeded = xerrors.New("ai budget exceeded") - ErrInternalServerError = xerrors.New("internal server error") + ErrSpendStatusCheck = xerrors.New("internal server error checking user ai spend") ) // ServeHTTP is the entrypoint for requests which will be intercepted by AI Bridge. @@ -146,6 +146,7 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) { http.Error(rw, ErrUnauthorized.Error(), http.StatusForbidden) return } + logger = logger.With(slog.F("user_id", id)) // TODO: make the budget period configurable; monthly for now. periodStart := dbtime.StartOfMonth(dbtime.Now().UTC()) @@ -155,16 +156,15 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) { }) if err != nil { logger.Warn(ctx, "ai spend status check failed", slog.Error(err)) - http.Error(rw, ErrInternalServerError.Error(), http.StatusInternalServerError) + http.Error(rw, ErrSpendStatusCheck.Error(), http.StatusInternalServerError) return } if spendStatus.GetExceeded() { http.Error(rw, fmt.Sprintf( - "%s: spent %.2f USD of %.2f USD in group %s", + "%s: spent US$%.2f of US$%.2f", ErrBudgetExceeded.Error(), float64(spendStatus.GetCurrentSpendMicros())/1_000_000, float64(spendStatus.GetSpendLimitMicros())/1_000_000, - spendStatus.GetEffectiveGroupId(), ), http.StatusPaymentRequired) return } diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index 8cae676ce3171..5507f0648deab 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -731,7 +731,7 @@ func (s *Server) IsAuthorized(ctx context.Context, in *proto.IsAuthorizedRequest } // GetUserAISpendStatus returns the user's AI spend status aggregated over -// [period_start, now]. +// [PeriodStart, now]. func (s *Server) GetUserAISpendStatus(ctx context.Context, in *proto.GetUserAISpendStatusRequest) (*proto.GetUserAISpendStatusResponse, error) { //nolint:gocritic // AIBridged has specific authz rules. ctx = dbauthz.AsAIBridged(ctx) @@ -740,8 +740,8 @@ func (s *Server) GetUserAISpendStatus(ctx context.Context, in *proto.GetUserAISp if err != nil { return nil, xerrors.Errorf("invalid user_id %q: %w", in.GetUserId(), err) } - // An unset period_start deserializes to time.Unix(0, 0), which would - // aggregate the user's lifetime spend against a period budget. + // An unset PeriodStart deserializes to time.Unix(0, 0), which would + // incorrectly aggregate the user's lifetime spend against a period budget. if in.PeriodStart == nil { return nil, xerrors.New("period_start is required") } @@ -777,6 +777,14 @@ type userAISpendStatus struct { // resolveUserAISpendStatus computes the user's AI spend status aggregated from // periodStart. Returns a zero-valued status (EffectiveGroupID = uuid.Nil) when // no budget applies to the user. +// +// Note: there is a potential race condition where two concurrent requests +// from the same user can both pass the check if processed in parallel, +// allowing brief overage. This is acceptable because: +// - Cost is only known after the LLM API returns. +// - Overage is bounded by request cost × concurrency; once the accumulated +// spend crosses the limit, subsequent requests are blocked. +// - Fail-open is acceptable for this case. func (s *Server) resolveUserAISpendStatus(ctx context.Context, userID uuid.UUID, periodStart time.Time) (userAISpendStatus, error) { effectiveBudget, ok, err := budget.ResolveUserAIBudget(ctx, s.store, userID, s.budgetPolicy) if err != nil { @@ -814,7 +822,7 @@ func (s *Server) resolveUserAISpendStatus(ctx context.Context, userID uuid.UUID, ) logger.Debug(ctx, "user AI spend status") if status.Exceeded { - logger.Info(ctx, "user AI budget exceeded") + logger.Warn(ctx, "user AI budget exceeded") } return status, nil diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go index 769ecfc0f5cc3..54d6c6ea09ec2 100644 --- a/coderd/aibridgedserver/aibridgedserver_test.go +++ b/coderd/aibridgedserver/aibridgedserver_test.go @@ -399,19 +399,19 @@ func TestGetUserAISpendStatus(t *testing.T) { userIDStr string omitPeriodStart bool setupMocks func(db *dbmock.MockStore, userID uuid.UUID) *proto.GetUserAISpendStatusResponse - wantErr bool + wantErrContains string }{ { // Invalid UUID short-circuits before any store call. - name: "invalid user_id", - userIDStr: "not-a-uuid", - wantErr: true, + name: "invalid user_id", + userIDStr: "not-a-uuid", + wantErrContains: "invalid user_id", }, { // Missing period_start is rejected before any store call. name: "missing period_start", omitPeriodStart: true, - wantErr: true, + wantErrContains: "period_start is required", }, { // No override and no group budget resolves: pass-through. All @@ -471,6 +471,28 @@ func TestGetUserAISpendStatus(t *testing.T) { } }, }, + { + // A configured limit of 0 is a valid "block-all" setting, distinct + // from "no budget configured". The >= comparison means spend of 0 + // against a limit of 0 is exceeded. + name: "zero limit blocks all requests", + setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.GetUserAISpendStatusResponse { + groupID := uuid.New() + db.EXPECT().GetUserAIBudgetOverride(gomock.Any(), userID). + Return(database.UserAIBudgetOverride{}, sql.ErrNoRows) + db.EXPECT().GetHighestGroupAIBudgetByUser(gomock.Any(), userID). + Return(database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 0}, nil) + db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()). + Return(database.GetUserAISpendSinceRow{SpendMicros: 0}, nil) + return &proto.GetUserAISpendStatusResponse{ + Exceeded: true, + UserId: userID.String(), + EffectiveGroupId: groupID.String(), + SpendLimitMicros: 0, + CurrentSpendMicros: 0, + } + }, + }, { // Spend above limit (spend 1500 > limit 1000). name: "over limit returns exceeded", @@ -523,7 +545,7 @@ func TestGetUserAISpendStatus(t *testing.T) { Return(database.UserAIBudgetOverride{}, sql.ErrConnDone) return nil }, - wantErr: true, + wantErrContains: "resolve effective AI budget", }, { // Error from spend aggregation propagates (fail-closed). @@ -537,7 +559,7 @@ func TestGetUserAISpendStatus(t *testing.T) { Return(database.GetUserAISpendSinceRow{}, sql.ErrConnDone) return nil }, - wantErr: true, + wantErrContains: "get user AI spend", }, } @@ -568,9 +590,10 @@ func TestGetUserAISpendStatus(t *testing.T) { req.PeriodStart = timestamppb.New(dbtime.StartOfMonth(dbtime.Now().UTC())) } resp, err := srv.GetUserAISpendStatus(t.Context(), req) - if tc.wantErr { + if tc.wantErrContains != "" { require.Error(t, err) require.Nil(t, resp) + assert.ErrorContains(t, err, tc.wantErrContains) return } require.NoError(t, err) @@ -584,66 +607,122 @@ func TestGetUserAISpendStatus(t *testing.T) { } } -// TestGetUserAISpendStatus_PeriodBoundary exercises the aggregation window at -// a UTC month boundary. A spend row in January pushes the user over their -// limit; the same row is excluded when the query moves to the next period. -func TestGetUserAISpendStatus_PeriodBoundary(t *testing.T) { +// TestGetUserAISpendStatus_Enforcement exercises real-DB scenarios that drive +// enforcement decisions. +func TestGetUserAISpendStatus_Enforcement(t *testing.T) { t.Parallel() - ctx := testutil.Context(t, testutil.WaitLong) - logger := testutil.Logger(t) + const groupLimitMicros = 1_000_000 - rawDB, _ := dbtestutil.NewDB(t) - authzDB := dbauthz.New(rawDB, rbac.NewStrictAuthorizer(prometheus.NewRegistry()), logger, coderdtest.AccessControlStorePointer()) + // setup provisions a user in an organization with a single budgeted group. + setup := func(t *testing.T) (context.Context, database.Store, *aibridgedserver.Server, database.User, database.Group) { + t.Helper() - org := dbgen.Organization(t, rawDB, database.Organization{}) - user := dbgen.User(t, rawDB, database.User{}) - dbgen.OrganizationMember(t, rawDB, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID}) - group := dbgen.Group(t, rawDB, database.Group{OrganizationID: org.ID}) - dbgen.GroupMember(t, rawDB, database.GroupMemberTable{UserID: user.ID, GroupID: group.ID}) + ctx := testutil.Context(t, testutil.WaitLong) + logger := testutil.Logger(t) - const limitMicros = 1_000_000 - _, err := rawDB.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ - GroupID: group.ID, - SpendLimitMicros: limitMicros, - }) - require.NoError(t, err) + rawDB, _ := dbtestutil.NewDB(t) + authzDB := dbauthz.New(rawDB, rbac.NewStrictAuthorizer(prometheus.NewRegistry()), logger, coderdtest.AccessControlStorePointer()) - prevMonth := time.Date(2026, time.January, 1, 0, 0, 0, 0, time.UTC) - newMonth := time.Date(2026, time.February, 1, 0, 0, 0, 0, time.UTC) + org := dbgen.Organization(t, rawDB, database.Organization{}) + user := dbgen.User(t, rawDB, database.User{}) + dbgen.OrganizationMember(t, rawDB, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID}) + group := dbgen.Group(t, rawDB, database.Group{OrganizationID: org.ID}) + dbgen.GroupMember(t, rawDB, database.GroupMemberTable{UserID: user.ID, GroupID: group.ID}) - // Insert a spend row in the previous month that pushes the user over the limit. - _, err = rawDB.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ - UserID: user.ID, - EffectiveGroupID: group.ID, - Day: prevMonth.AddDate(0, 0, 14), - CostMicros: 1_500_000, - }) - require.NoError(t, err) + _, err := rawDB.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{ + GroupID: group.ID, + SpendLimitMicros: groupLimitMicros, + }) + require.NoError(t, err) - srv, err := aibridgedserver.NewServer(ctx, authzDB, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) - require.NoError(t, err) + srv, err := aibridgedserver.NewServer(ctx, authzDB, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) + require.NoError(t, err) + + return ctx, rawDB, srv, user, group + } + + t.Run("period boundary excludes prior period spend", func(t *testing.T) { + t.Parallel() + ctx, rawDB, srv, user, group := setup(t) + + prevMonth := time.Date(2026, time.January, 1, 0, 0, 0, 0, time.UTC) + newMonth := time.Date(2026, time.February, 1, 0, 0, 0, 0, time.UTC) + + // User spend on 2026-01-15. + _, err := rawDB.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, + EffectiveGroupID: group.ID, + Day: prevMonth.AddDate(0, 0, 14), + CostMicros: 1_500_000, + }) + require.NoError(t, err) - // Same UTC month as the spend row: aggregation includes it, user exceeded. - prevMonthResp, err := srv.GetUserAISpendStatus(ctx, &proto.GetUserAISpendStatusRequest{ - UserId: user.ID.String(), - PeriodStart: timestamppb.New(prevMonth), + // Query with period_start 2026-01-01: includes the 2026-01-15 spend, user exceeded. + prevMonthResp, err := srv.GetUserAISpendStatus(ctx, &proto.GetUserAISpendStatusRequest{ + UserId: user.ID.String(), + PeriodStart: timestamppb.New(prevMonth), + }) + require.NoError(t, err) + require.True(t, prevMonthResp.GetExceeded()) + require.Equal(t, int64(1_500_000), prevMonthResp.GetCurrentSpendMicros()) + require.Equal(t, int64(groupLimitMicros), prevMonthResp.GetSpendLimitMicros()) + + // Query with period_start 2026-02-01: excludes the 2026-01-15 spend, user not exceeded. + newMonthResp, err := srv.GetUserAISpendStatus(ctx, &proto.GetUserAISpendStatusRequest{ + UserId: user.ID.String(), + PeriodStart: timestamppb.New(newMonth), + }) + require.NoError(t, err) + require.False(t, newMonthResp.GetExceeded()) + require.Equal(t, int64(0), newMonthResp.GetCurrentSpendMicros()) + require.Equal(t, int64(groupLimitMicros), newMonthResp.GetSpendLimitMicros()) }) - require.NoError(t, err) - require.True(t, prevMonthResp.GetExceeded()) - require.Equal(t, int64(1_500_000), prevMonthResp.GetCurrentSpendMicros()) - require.Equal(t, int64(limitMicros), prevMonthResp.GetSpendLimitMicros()) - - // After the boundary: the previous month's row is outside the aggregation - // window and the user is no longer exceeded. - newMonthResp, err := srv.GetUserAISpendStatus(ctx, &proto.GetUserAISpendStatusRequest{ - UserId: user.ID.String(), - PeriodStart: timestamppb.New(newMonth), + + t.Run("new user override unblocks user", func(t *testing.T) { + t.Parallel() + ctx, rawDB, srv, user, group := setup(t) + + periodStart := dbtime.StartOfMonth(dbtime.Now().UTC()) + + // User spend today, within the current period. + _, err := rawDB.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, + EffectiveGroupID: group.ID, + Day: dbtime.Now().UTC(), + CostMicros: 1_500_000, + }) + require.NoError(t, err) + + // User's spend exceeds the group limit. + beforeResp, err := srv.GetUserAISpendStatus(ctx, &proto.GetUserAISpendStatusRequest{ + UserId: user.ID.String(), + PeriodStart: timestamppb.New(periodStart), + }) + require.NoError(t, err) + require.True(t, beforeResp.GetExceeded()) + require.Equal(t, int64(groupLimitMicros), beforeResp.GetSpendLimitMicros()) + require.Equal(t, int64(1_500_000), beforeResp.GetCurrentSpendMicros()) + + // Add user override with a higher limit on the same group. The override + // wins, so the user's spend is now under the effective limit. + const overrideLimitMicros = 2_000_000 + _, err = rawDB.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{ + UserID: user.ID, + GroupID: group.ID, + SpendLimitMicros: overrideLimitMicros, + }) + require.NoError(t, err) + + afterResp, err := srv.GetUserAISpendStatus(ctx, &proto.GetUserAISpendStatusRequest{ + UserId: user.ID.String(), + PeriodStart: timestamppb.New(periodStart), + }) + require.NoError(t, err) + require.False(t, afterResp.GetExceeded()) + require.Equal(t, int64(overrideLimitMicros), afterResp.GetSpendLimitMicros()) + require.Equal(t, int64(1_500_000), afterResp.GetCurrentSpendMicros()) }) - require.NoError(t, err) - require.False(t, newMonthResp.GetExceeded()) - require.Equal(t, int64(0), newMonthResp.GetCurrentSpendMicros()) - require.Equal(t, int64(limitMicros), newMonthResp.GetSpendLimitMicros()) } func TestGetMCPServerConfigs(t *testing.T) { From 6e34d549b5fe6d4796fe229ef2d35305f0d5344f Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Thu, 2 Jul 2026 10:43:28 +0000 Subject: [PATCH 3/6] refactor: rename to IsBudgetExceeded and simplify RPC contract --- coderd/aibridged/aibridged_test.go | 35 +++--- coderd/aibridged/aibridgedmock/clientmock.go | 14 +-- coderd/aibridged/http.go | 20 ++-- coderd/aibridged/proto/aibridged.proto | 20 ++-- coderd/aibridged/proto/aibridged_drpc.pb.go | 28 ++--- coderd/aibridgedserver/aibridgedserver.go | 30 +++-- .../aibridgedserver/aibridgedserver_test.go | 105 +++++++----------- 7 files changed, 108 insertions(+), 144 deletions(-) diff --git a/coderd/aibridged/aibridged_test.go b/coderd/aibridged/aibridged_test.go index f421580a6c6c0..52975e033b821 100644 --- a/coderd/aibridged/aibridged_test.go +++ b/coderd/aibridged/aibridged_test.go @@ -24,6 +24,7 @@ import ( "github.com/coder/coder/v2/coderd/aibridged" mock "github.com/coder/coder/v2/coderd/aibridged/aibridgedmock" "github.com/coder/coder/v2/coderd/aibridged/proto" + "github.com/coder/coder/v2/coderd/util/ptr" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/testutil" "github.com/coder/quartz" @@ -134,24 +135,22 @@ func TestServeHTTP_FailureModes(t *testing.T) { applyMocksFn: func(client *mock.MockDRPCClient, _ *mock.MockPooler) { // Authorization passes. client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId: uuid.NewString()}, nil) - client.EXPECT().GetUserAISpendStatus(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.GetUserAISpendStatusResponse{ - Exceeded: true, - EffectiveGroupId: uuid.NewString(), - SpendLimitMicros: 1_000, - CurrentSpendMicros: 1_500, + client.EXPECT().IsBudgetExceeded(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsBudgetExceededResponse{ + Exceeded: true, + SpendLimitMicros: ptr.Ref(int64(1_000)), }, nil) }, - expectedErr: aibridged.ErrBudgetExceeded, - expectedStatus: http.StatusPaymentRequired, + expectedErr: xerrors.New("ai budget of"), + expectedStatus: http.StatusForbidden, }, { name: "budget check failed", applyMocksFn: func(client *mock.MockDRPCClient, _ *mock.MockPooler) { // Authorization passes. client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId: uuid.NewString()}, nil) - client.EXPECT().GetUserAISpendStatus(gomock.Any(), gomock.Any()).AnyTimes().Return(nil, xerrors.New("oops")) + client.EXPECT().IsBudgetExceeded(gomock.Any(), gomock.Any()).AnyTimes().Return(nil, xerrors.New("oops")) }, - expectedErr: aibridged.ErrSpendStatusCheck, + expectedErr: aibridged.ErrBudgetCheck, expectedStatus: http.StatusInternalServerError, }, @@ -161,7 +160,7 @@ func TestServeHTTP_FailureModes(t *testing.T) { applyMocksFn: func(client *mock.MockDRPCClient, pool *mock.MockPooler) { // Should pass authorization and budget check. client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId: uuid.NewString()}, nil) - client.EXPECT().GetUserAISpendStatus(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.GetUserAISpendStatusResponse{}, nil) + client.EXPECT().IsBudgetExceeded(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsBudgetExceededResponse{}, nil) // But fail when acquiring a pool instance. pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().Return(nil, xerrors.New("oops")) }, @@ -251,7 +250,7 @@ func TestServeHTTP_DelegatedAPIKey(t *testing.T) { Username: "u", }, nil }) - client.EXPECT().GetUserAISpendStatus(gomock.Any(), gomock.Any()).Return(&proto.GetUserAISpendStatusResponse{}, nil) + client.EXPECT().IsBudgetExceeded(gomock.Any(), gomock.Any()).Return(&proto.IsBudgetExceededResponse{}, nil) pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( func(_ context.Context, req aibridged.Request, _ aibridged.ClientFunc, _ aibridged.MCPProxyBuilder) (http.Handler, error) { assert.Empty(t, req.SessionKey, @@ -284,7 +283,7 @@ func TestServeHTTP_DelegatedAPIKey(t *testing.T) { ApiKeyId: testKeyID, Username: "u", }, nil) - client.EXPECT().GetUserAISpendStatus(gomock.Any(), gomock.Any()).Return(&proto.GetUserAISpendStatusResponse{}, nil) + client.EXPECT().IsBudgetExceeded(gomock.Any(), gomock.Any()).Return(&proto.IsBudgetExceededResponse{}, nil) pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( func(_ context.Context, req aibridged.Request, _ aibridged.ClientFunc, _ aibridged.MCPProxyBuilder) (http.Handler, error) { assert.Equal(t, "coder-token-byok", req.SessionKey, @@ -375,7 +374,7 @@ func TestServeHTTP_DelegatedAPIKey_BYOK_Integration(t *testing.T) { Username: "u", }, nil }) - client.EXPECT().GetUserAISpendStatus(gomock.Any(), gomock.Any()).Return(&proto.GetUserAISpendStatusResponse{}, nil) + client.EXPECT().IsBudgetExceeded(gomock.Any(), gomock.Any()).Return(&proto.IsBudgetExceededResponse{}, nil) pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(mockH, nil) factory := aibridged.NewTransportFactory(srv) @@ -428,7 +427,7 @@ func TestServeHTTP_DelegatedAPIKey_Integration(t *testing.T) { Username: "u", }, nil }) - client.EXPECT().GetUserAISpendStatus(gomock.Any(), gomock.Any()).Return(&proto.GetUserAISpendStatusResponse{}, nil) + client.EXPECT().IsBudgetExceeded(gomock.Any(), gomock.Any()).Return(&proto.IsBudgetExceededResponse{}, nil) pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(mockH, nil) factory := aibridged.NewTransportFactory(srv) @@ -515,7 +514,7 @@ func TestServeHTTP_StripCoderToken(t *testing.T) { conn := &mockDRPCConn{} client.EXPECT().DRPCConn().AnyTimes().Return(conn) client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId: uuid.NewString()}, nil) - client.EXPECT().GetUserAISpendStatus(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.GetUserAISpendStatusResponse{}, nil) + client.EXPECT().IsBudgetExceeded(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsBudgetExceededResponse{}, nil) pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().Return(mockH, nil) httpSrv := httptest.NewServer(srv) @@ -710,7 +709,7 @@ func TestServeHTTP_ActorHeaders(t *testing.T) { OwnerId: testUserID.String(), Username: testUsername, }, nil) - client.EXPECT().GetUserAISpendStatus(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.GetUserAISpendStatusResponse{}, nil) + client.EXPECT().IsBudgetExceeded(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsBudgetExceededResponse{}, nil) client.EXPECT().GetMCPServerConfigs(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.GetMCPServerConfigsResponse{}, nil) client.EXPECT().RecordInterception(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.RecordInterceptionResponse{}, nil) client.EXPECT().RecordInterceptionEnded(gomock.Any(), gomock.Any()).AnyTimes() @@ -809,7 +808,7 @@ func TestRouting(t *testing.T) { client.EXPECT().DRPCConn().AnyTimes().Return(conn) client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId: uuid.NewString()}, nil) - client.EXPECT().GetUserAISpendStatus(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.GetUserAISpendStatusResponse{}, nil) + client.EXPECT().IsBudgetExceeded(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsBudgetExceededResponse{}, nil) client.EXPECT().GetMCPServerConfigs(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.GetMCPServerConfigsResponse{}, nil) // This is the only recording we really care about in this test. This is called before the provider-specific logic processes // the incoming request, and anything beyond that is the responsibility of coder/aibridge to test. @@ -885,7 +884,7 @@ func TestServeHTTP_StripInternalHeaders(t *testing.T) { conn := &mockDRPCConn{} client.EXPECT().DRPCConn().AnyTimes().Return(conn) client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId: uuid.NewString()}, nil) - client.EXPECT().GetUserAISpendStatus(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.GetUserAISpendStatusResponse{}, nil) + client.EXPECT().IsBudgetExceeded(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsBudgetExceededResponse{}, nil) pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().Return(mockH, nil) httpSrv := httptest.NewServer(srv) diff --git a/coderd/aibridged/aibridgedmock/clientmock.go b/coderd/aibridged/aibridgedmock/clientmock.go index c036fae84e8e1..4aa0dd801b750 100644 --- a/coderd/aibridged/aibridgedmock/clientmock.go +++ b/coderd/aibridged/aibridgedmock/clientmock.go @@ -101,19 +101,19 @@ func (mr *MockDRPCClientMockRecorder) GetMCPServerConfigs(ctx, in any) *gomock.C return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetMCPServerConfigs", reflect.TypeOf((*MockDRPCClient)(nil).GetMCPServerConfigs), ctx, in) } -// GetUserAISpendStatus mocks base method. -func (m *MockDRPCClient) GetUserAISpendStatus(ctx context.Context, in *proto.GetUserAISpendStatusRequest) (*proto.GetUserAISpendStatusResponse, error) { +// IsBudgetExceeded mocks base method. +func (m *MockDRPCClient) IsBudgetExceeded(ctx context.Context, in *proto.IsBudgetExceededRequest) (*proto.IsBudgetExceededResponse, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetUserAISpendStatus", ctx, in) - ret0, _ := ret[0].(*proto.GetUserAISpendStatusResponse) + ret := m.ctrl.Call(m, "IsBudgetExceeded", ctx, in) + ret0, _ := ret[0].(*proto.IsBudgetExceededResponse) ret1, _ := ret[1].(error) return ret0, ret1 } -// GetUserAISpendStatus indicates an expected call of GetUserAISpendStatus. -func (mr *MockDRPCClientMockRecorder) GetUserAISpendStatus(ctx, in any) *gomock.Call { +// IsBudgetExceeded indicates an expected call of IsBudgetExceeded. +func (mr *MockDRPCClientMockRecorder) IsBudgetExceeded(ctx, in any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserAISpendStatus", reflect.TypeOf((*MockDRPCClient)(nil).GetUserAISpendStatus), ctx, in) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsBudgetExceeded", reflect.TypeOf((*MockDRPCClient)(nil).IsBudgetExceeded), ctx, in) } // IsAuthorized mocks base method. diff --git a/coderd/aibridged/http.go b/coderd/aibridged/http.go index fd15eef965244..f35b895f7813c 100644 --- a/coderd/aibridged/http.go +++ b/coderd/aibridged/http.go @@ -24,8 +24,7 @@ var ( ErrConnect = xerrors.New("could not connect to coderd") ErrUnauthorized = xerrors.New("unauthorized") ErrAcquireRequestHandler = xerrors.New("failed to acquire request handler") - ErrBudgetExceeded = xerrors.New("ai budget exceeded") - ErrSpendStatusCheck = xerrors.New("internal server error checking user ai spend") + ErrBudgetCheck = xerrors.New("internal server error checking user ai budget") ) // ServeHTTP is the entrypoint for requests which will be intercepted by AI Bridge. @@ -148,24 +147,21 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) { } logger = logger.With(slog.F("user_id", id)) - // TODO: make the budget period configurable; monthly for now. periodStart := dbtime.StartOfMonth(dbtime.Now().UTC()) - spendStatus, err := client.GetUserAISpendStatus(ctx, &proto.GetUserAISpendStatusRequest{ + budgetStatus, err := client.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{ UserId: id.String(), PeriodStart: timestamppb.New(periodStart), }) if err != nil { - logger.Warn(ctx, "ai spend status check failed", slog.Error(err)) - http.Error(rw, ErrSpendStatusCheck.Error(), http.StatusInternalServerError) + logger.Warn(ctx, "ai budget check failed", slog.Error(err)) + http.Error(rw, ErrBudgetCheck.Error(), http.StatusInternalServerError) return } - if spendStatus.GetExceeded() { + if budgetStatus.GetExceeded() { http.Error(rw, fmt.Sprintf( - "%s: spent US$%.2f of US$%.2f", - ErrBudgetExceeded.Error(), - float64(spendStatus.GetCurrentSpendMicros())/1_000_000, - float64(spendStatus.GetSpendLimitMicros())/1_000_000, - ), http.StatusPaymentRequired) + "ai budget of US$%.2f exceeded. Please contact an administrator for more details.", + float64(budgetStatus.GetSpendLimitMicros())/1_000_000, + ), http.StatusForbidden) return } diff --git a/coderd/aibridged/proto/aibridged.proto b/coderd/aibridged/proto/aibridged.proto index 4152c1c729a8e..d0fd26803fd75 100644 --- a/coderd/aibridged/proto/aibridged.proto +++ b/coderd/aibridged/proto/aibridged.proto @@ -34,9 +34,9 @@ service Authorizer { // IsAuthorized validates that a given Coder key is valid and the user is authorized to use AI Bridge. // TODO: add authorization; currently only key validation takes place. rpc IsAuthorized(IsAuthorizedRequest) returns (IsAuthorizedResponse); - // GetUserAISpendStatus returns the user's AI spend status aggregated over - // [period_start, now]. - rpc GetUserAISpendStatus(GetUserAISpendStatusRequest) returns (GetUserAISpendStatusResponse); + // IsBudgetExceeded reports whether the user's AI spend has reached their + // effective limit over [period_start, now]. + rpc IsBudgetExceeded(IsBudgetExceededRequest) returns (IsBudgetExceededResponse); } // ProviderConfigurator serves AI provider configuration to embedded and @@ -171,23 +171,21 @@ message IsAuthorizedResponse { string username = 3; } -message GetUserAISpendStatusRequest { +message IsBudgetExceededRequest { string user_id = 1; // UUID // The spend aggregation window is [period_start, now]. google.protobuf.Timestamp period_start = 2; } -message GetUserAISpendStatusResponse { +message IsBudgetExceededResponse { // exceeded is true when the user's aggregated spend has reached the // effective limit. False when no budget is configured for the user OR // they are within their limit. bool exceeded = 1; - string user_id = 2; // UUID - // effective_group_id is empty when no budget is configured for the user. - // spend_limit_micros and current_spend_micros are 0 in that case. - string effective_group_id = 3; - int64 spend_limit_micros = 4; - int64 current_spend_micros = 5; + // spend_limit_micros is the effective spend limit in micro-units. + // Unset when no budget is configured for the user (unlimited). + // 0 when a group is explicitly configured with a 0 limit (blocked). + optional int64 spend_limit_micros = 2; } message GetAIProvidersRequest {} diff --git a/coderd/aibridged/proto/aibridged_drpc.pb.go b/coderd/aibridged/proto/aibridged_drpc.pb.go index 093286db2988d..723e670dad3e2 100644 --- a/coderd/aibridged/proto/aibridged_drpc.pb.go +++ b/coderd/aibridged/proto/aibridged_drpc.pb.go @@ -429,7 +429,7 @@ type DRPCAuthorizerClient interface { DRPCConn() drpc.Conn IsAuthorized(ctx context.Context, in *IsAuthorizedRequest) (*IsAuthorizedResponse, error) - GetUserAISpendStatus(ctx context.Context, in *GetUserAISpendStatusRequest) (*GetUserAISpendStatusResponse, error) + IsBudgetExceeded(ctx context.Context, in *IsBudgetExceededRequest) (*IsBudgetExceededResponse, error) } type drpcAuthorizerClient struct { @@ -451,9 +451,9 @@ func (c *drpcAuthorizerClient) IsAuthorized(ctx context.Context, in *IsAuthorize return out, nil } -func (c *drpcAuthorizerClient) GetUserAISpendStatus(ctx context.Context, in *GetUserAISpendStatusRequest) (*GetUserAISpendStatusResponse, error) { - out := new(GetUserAISpendStatusResponse) - err := c.cc.Invoke(ctx, "/proto.Authorizer/GetUserAISpendStatus", drpcEncoding_File_coderd_aibridged_proto_aibridged_proto{}, in, out) +func (c *drpcAuthorizerClient) IsBudgetExceeded(ctx context.Context, in *IsBudgetExceededRequest) (*IsBudgetExceededResponse, error) { + out := new(IsBudgetExceededResponse) + err := c.cc.Invoke(ctx, "/proto.Authorizer/IsBudgetExceeded", drpcEncoding_File_coderd_aibridged_proto_aibridged_proto{}, in, out) if err != nil { return nil, err } @@ -462,7 +462,7 @@ func (c *drpcAuthorizerClient) GetUserAISpendStatus(ctx context.Context, in *Get type DRPCAuthorizerServer interface { IsAuthorized(context.Context, *IsAuthorizedRequest) (*IsAuthorizedResponse, error) - GetUserAISpendStatus(context.Context, *GetUserAISpendStatusRequest) (*GetUserAISpendStatusResponse, error) + IsBudgetExceeded(context.Context, *IsBudgetExceededRequest) (*IsBudgetExceededResponse, error) } type DRPCAuthorizerUnimplementedServer struct{} @@ -471,7 +471,7 @@ func (s *DRPCAuthorizerUnimplementedServer) IsAuthorized(context.Context, *IsAut return nil, drpcerr.WithCode(errors.New("Unimplemented"), drpcerr.Unimplemented) } -func (s *DRPCAuthorizerUnimplementedServer) GetUserAISpendStatus(context.Context, *GetUserAISpendStatusRequest) (*GetUserAISpendStatusResponse, error) { +func (s *DRPCAuthorizerUnimplementedServer) IsBudgetExceeded(context.Context, *IsBudgetExceededRequest) (*IsBudgetExceededResponse, error) { return nil, drpcerr.WithCode(errors.New("Unimplemented"), drpcerr.Unimplemented) } @@ -491,14 +491,14 @@ func (DRPCAuthorizerDescription) Method(n int) (string, drpc.Encoding, drpc.Rece ) }, DRPCAuthorizerServer.IsAuthorized, true case 1: - return "/proto.Authorizer/GetUserAISpendStatus", drpcEncoding_File_coderd_aibridged_proto_aibridged_proto{}, + return "/proto.Authorizer/IsBudgetExceeded", drpcEncoding_File_coderd_aibridged_proto_aibridged_proto{}, func(srv interface{}, ctx context.Context, in1, in2 interface{}) (drpc.Message, error) { return srv.(DRPCAuthorizerServer). - GetUserAISpendStatus( + IsBudgetExceeded( ctx, - in1.(*GetUserAISpendStatusRequest), + in1.(*IsBudgetExceededRequest), ) - }, DRPCAuthorizerServer.GetUserAISpendStatus, true + }, DRPCAuthorizerServer.IsBudgetExceeded, true default: return "", nil, nil, nil, false } @@ -524,16 +524,16 @@ func (x *drpcAuthorizer_IsAuthorizedStream) SendAndClose(m *IsAuthorizedResponse return x.CloseSend() } -type DRPCAuthorizer_GetUserAISpendStatusStream interface { +type DRPCAuthorizer_IsBudgetExceededStream interface { drpc.Stream - SendAndClose(*GetUserAISpendStatusResponse) error + SendAndClose(*IsBudgetExceededResponse) error } -type drpcAuthorizer_GetUserAISpendStatusStream struct { +type drpcAuthorizer_IsBudgetExceededStream struct { drpc.Stream } -func (x *drpcAuthorizer_GetUserAISpendStatusStream) SendAndClose(m *GetUserAISpendStatusResponse) error { +func (x *drpcAuthorizer_IsBudgetExceededStream) SendAndClose(m *IsBudgetExceededResponse) error { if err := x.MsgSend(m, drpcEncoding_File_coderd_aibridged_proto_aibridged_proto{}); err != nil { return err } diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index 5507f0648deab..edc521afb4dd4 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -730,9 +730,9 @@ func (s *Server) IsAuthorized(ctx context.Context, in *proto.IsAuthorizedRequest }, nil } -// GetUserAISpendStatus returns the user's AI spend status aggregated over -// [PeriodStart, now]. -func (s *Server) GetUserAISpendStatus(ctx context.Context, in *proto.GetUserAISpendStatusRequest) (*proto.GetUserAISpendStatusResponse, error) { +// IsBudgetExceeded reports whether the user's AI spend has reached their +// effective limit over [PeriodStart, now]. +func (s *Server) IsBudgetExceeded(ctx context.Context, in *proto.IsBudgetExceededRequest) (*proto.IsBudgetExceededResponse, error) { //nolint:gocritic // AIBridged has specific authz rules. ctx = dbauthz.AsAIBridged(ctx) @@ -751,25 +751,19 @@ func (s *Server) GetUserAISpendStatus(ctx context.Context, in *proto.GetUserAISp if err != nil { return nil, err } - resp := &proto.GetUserAISpendStatusResponse{ - Exceeded: status.Exceeded, - UserId: userID.String(), - SpendLimitMicros: status.SpendLimitMicros, - CurrentSpendMicros: status.CurrentSpendMicros, - } - if status.EffectiveGroupID != uuid.Nil { - resp.EffectiveGroupId = status.EffectiveGroupID.String() - } - return resp, nil + return &proto.IsBudgetExceededResponse{ + Exceeded: status.Exceeded, + SpendLimitMicros: status.SpendLimitMicros, + }, nil } // userAISpendStatus is a snapshot of a user's AI spend against their effective // budget over [PeriodStart, now]. When no budget is configured for the user, -// all fields are zero-valued. +// SpendLimitMicros is nil and the other fields are zero-valued. type userAISpendStatus struct { EffectiveGroupID uuid.UUID PeriodStart time.Time - SpendLimitMicros int64 + SpendLimitMicros *int64 CurrentSpendMicros int64 Exceeded bool } @@ -784,6 +778,8 @@ type userAISpendStatus struct { // - Cost is only known after the LLM API returns. // - Overage is bounded by request cost × concurrency; once the accumulated // spend crosses the limit, subsequent requests are blocked. +// - Cost accounting is advisory, not strict. The goal is to prevent +// overages, not build an accounting system. // - Fail-open is acceptable for this case. func (s *Server) resolveUserAISpendStatus(ctx context.Context, userID uuid.UUID, periodStart time.Time) (userAISpendStatus, error) { effectiveBudget, ok, err := budget.ResolveUserAIBudget(ctx, s.store, userID, s.budgetPolicy) @@ -807,7 +803,7 @@ func (s *Server) resolveUserAISpendStatus(ctx context.Context, userID uuid.UUID, status := userAISpendStatus{ EffectiveGroupID: effectiveBudget.GroupID, PeriodStart: periodStart, - SpendLimitMicros: effectiveBudget.SpendLimitMicros, + SpendLimitMicros: ptr.Ref(effectiveBudget.SpendLimitMicros), CurrentSpendMicros: spend.SpendMicros, Exceeded: spend.SpendMicros >= effectiveBudget.SpendLimitMicros, } @@ -817,7 +813,7 @@ func (s *Server) resolveUserAISpendStatus(ctx context.Context, userID uuid.UUID, slog.F("effective_group_id", status.EffectiveGroupID), slog.F("period_start", status.PeriodStart), slog.F("current_spend_micros", status.CurrentSpendMicros), - slog.F("spend_limit_micros", status.SpendLimitMicros), + slog.F("spend_limit_micros", *status.SpendLimitMicros), slog.F("exceeded", status.Exceeded), ) logger.Debug(ctx, "user AI spend status") diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go index 54d6c6ea09ec2..1411c7e68944d 100644 --- a/coderd/aibridgedserver/aibridgedserver_test.go +++ b/coderd/aibridgedserver/aibridgedserver_test.go @@ -391,14 +391,14 @@ func TestAuthorization_Delegated(t *testing.T) { } } -func TestGetUserAISpendStatus(t *testing.T) { +func TestIsBudgetExceeded(t *testing.T) { t.Parallel() cases := []struct { name string userIDStr string omitPeriodStart bool - setupMocks func(db *dbmock.MockStore, userID uuid.UUID) *proto.GetUserAISpendStatusResponse + setupMocks func(db *dbmock.MockStore, userID uuid.UUID) *proto.IsBudgetExceededResponse wantErrContains string }{ { @@ -414,27 +414,23 @@ func TestGetUserAISpendStatus(t *testing.T) { wantErrContains: "period_start is required", }, { - // No override and no group budget resolves: pass-through. All - // spend/limit fields are zero-valued. + // No override and no group budget resolves: pass-through. name: "no budget configured returns not exceeded", - setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.GetUserAISpendStatusResponse { + setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.IsBudgetExceededResponse { db.EXPECT().GetUserAIBudgetOverride(gomock.Any(), userID). Return(database.UserAIBudgetOverride{}, sql.ErrNoRows) db.EXPECT().GetHighestGroupAIBudgetByUser(gomock.Any(), userID). Return(database.GetHighestGroupAIBudgetByUserRow{}, sql.ErrNoRows) - return &proto.GetUserAISpendStatusResponse{ - Exceeded: false, - UserId: userID.String(), - EffectiveGroupId: "", - SpendLimitMicros: 0, - CurrentSpendMicros: 0, + return &proto.IsBudgetExceededResponse{ + Exceeded: false, + SpendLimitMicros: nil, } }, }, { // Group budget resolves, spend below limit (spend 500 < limit 1000). name: "under limit returns not exceeded", - setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.GetUserAISpendStatusResponse { + setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.IsBudgetExceededResponse { groupID := uuid.New() db.EXPECT().GetUserAIBudgetOverride(gomock.Any(), userID). Return(database.UserAIBudgetOverride{}, sql.ErrNoRows) @@ -442,19 +438,16 @@ func TestGetUserAISpendStatus(t *testing.T) { Return(database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 1_000}, nil) db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()). Return(database.GetUserAISpendSinceRow{SpendMicros: 500}, nil) - return &proto.GetUserAISpendStatusResponse{ - Exceeded: false, - UserId: userID.String(), - EffectiveGroupId: groupID.String(), - SpendLimitMicros: 1_000, - CurrentSpendMicros: 500, + return &proto.IsBudgetExceededResponse{ + Exceeded: false, + SpendLimitMicros: ptr.Ref(int64(1_000)), } }, }, { // spend == limit is treated as exceeded (spend 1000 >= limit 1000). name: "at limit returns exceeded", - setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.GetUserAISpendStatusResponse { + setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.IsBudgetExceededResponse { groupID := uuid.New() db.EXPECT().GetUserAIBudgetOverride(gomock.Any(), userID). Return(database.UserAIBudgetOverride{}, sql.ErrNoRows) @@ -462,12 +455,9 @@ func TestGetUserAISpendStatus(t *testing.T) { Return(database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 1_000}, nil) db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()). Return(database.GetUserAISpendSinceRow{SpendMicros: 1_000}, nil) - return &proto.GetUserAISpendStatusResponse{ - Exceeded: true, - UserId: userID.String(), - EffectiveGroupId: groupID.String(), - SpendLimitMicros: 1_000, - CurrentSpendMicros: 1_000, + return &proto.IsBudgetExceededResponse{ + Exceeded: true, + SpendLimitMicros: ptr.Ref(int64(1_000)), } }, }, @@ -476,7 +466,7 @@ func TestGetUserAISpendStatus(t *testing.T) { // from "no budget configured". The >= comparison means spend of 0 // against a limit of 0 is exceeded. name: "zero limit blocks all requests", - setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.GetUserAISpendStatusResponse { + setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.IsBudgetExceededResponse { groupID := uuid.New() db.EXPECT().GetUserAIBudgetOverride(gomock.Any(), userID). Return(database.UserAIBudgetOverride{}, sql.ErrNoRows) @@ -484,19 +474,17 @@ func TestGetUserAISpendStatus(t *testing.T) { Return(database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 0}, nil) db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()). Return(database.GetUserAISpendSinceRow{SpendMicros: 0}, nil) - return &proto.GetUserAISpendStatusResponse{ - Exceeded: true, - UserId: userID.String(), - EffectiveGroupId: groupID.String(), - SpendLimitMicros: 0, - CurrentSpendMicros: 0, + // Blocked user: SpendLimitMicros is explicitly set to 0. + return &proto.IsBudgetExceededResponse{ + Exceeded: true, + SpendLimitMicros: ptr.Ref(int64(0)), } }, }, { // Spend above limit (spend 1500 > limit 1000). name: "over limit returns exceeded", - setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.GetUserAISpendStatusResponse { + setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.IsBudgetExceededResponse { groupID := uuid.New() db.EXPECT().GetUserAIBudgetOverride(gomock.Any(), userID). Return(database.UserAIBudgetOverride{}, sql.ErrNoRows) @@ -504,12 +492,9 @@ func TestGetUserAISpendStatus(t *testing.T) { Return(database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 1_000}, nil) db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()). Return(database.GetUserAISpendSinceRow{SpendMicros: 1_500}, nil) - return &proto.GetUserAISpendStatusResponse{ - Exceeded: true, - UserId: userID.String(), - EffectiveGroupId: groupID.String(), - SpendLimitMicros: 1_000, - CurrentSpendMicros: 1_500, + return &proto.IsBudgetExceededResponse{ + Exceeded: true, + SpendLimitMicros: ptr.Ref(int64(1_000)), } }, }, @@ -517,7 +502,7 @@ func TestGetUserAISpendStatus(t *testing.T) { // User override wins: group lookup is skipped, spend is aggregated // against the override's group (spend 600 > limit 500). name: "user override wins over group budget", - setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.GetUserAISpendStatusResponse { + setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.IsBudgetExceededResponse { overrideGroupID := uuid.New() db.EXPECT().GetUserAIBudgetOverride(gomock.Any(), userID). Return(database.UserAIBudgetOverride{ @@ -528,19 +513,16 @@ func TestGetUserAISpendStatus(t *testing.T) { db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Cond(func(p database.GetUserAISpendSinceParams) bool { return assert.Equal(t, overrideGroupID, p.EffectiveGroupID, "spend aggregated against override group") })).Return(database.GetUserAISpendSinceRow{SpendMicros: 600}, nil) - return &proto.GetUserAISpendStatusResponse{ - Exceeded: true, - UserId: userID.String(), - EffectiveGroupId: overrideGroupID.String(), - SpendLimitMicros: 500, - CurrentSpendMicros: 600, + return &proto.IsBudgetExceededResponse{ + Exceeded: true, + SpendLimitMicros: ptr.Ref(int64(500)), } }, }, { // Unexpected error from budget override lookup propagates. name: "budget resolution error propagates", - setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.GetUserAISpendStatusResponse { + setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.IsBudgetExceededResponse { db.EXPECT().GetUserAIBudgetOverride(gomock.Any(), userID). Return(database.UserAIBudgetOverride{}, sql.ErrConnDone) return nil @@ -550,7 +532,7 @@ func TestGetUserAISpendStatus(t *testing.T) { { // Error from spend aggregation propagates (fail-closed). name: "spend aggregation error propagates", - setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.GetUserAISpendStatusResponse { + setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.IsBudgetExceededResponse { db.EXPECT().GetUserAIBudgetOverride(gomock.Any(), userID). Return(database.UserAIBudgetOverride{}, sql.ErrNoRows) db.EXPECT().GetHighestGroupAIBudgetByUser(gomock.Any(), userID). @@ -577,7 +559,7 @@ func TestGetUserAISpendStatus(t *testing.T) { userIDStr = userID.String() } - var wantResp *proto.GetUserAISpendStatusResponse + var wantResp *proto.IsBudgetExceededResponse if tc.setupMocks != nil { wantResp = tc.setupMocks(db, userID) } @@ -585,11 +567,11 @@ func TestGetUserAISpendStatus(t *testing.T) { srv, err := aibridgedserver.NewServer(t.Context(), db, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) require.NoError(t, err) - req := &proto.GetUserAISpendStatusRequest{UserId: userIDStr} + req := &proto.IsBudgetExceededRequest{UserId: userIDStr} if !tc.omitPeriodStart { req.PeriodStart = timestamppb.New(dbtime.StartOfMonth(dbtime.Now().UTC())) } - resp, err := srv.GetUserAISpendStatus(t.Context(), req) + resp, err := srv.IsBudgetExceeded(t.Context(), req) if tc.wantErrContains != "" { require.Error(t, err) require.Nil(t, resp) @@ -599,17 +581,14 @@ func TestGetUserAISpendStatus(t *testing.T) { require.NoError(t, err) require.NotNil(t, resp) require.Equal(t, wantResp.GetExceeded(), resp.GetExceeded(), "exceeded") - require.Equal(t, wantResp.GetUserId(), resp.GetUserId(), "user_id") - require.Equal(t, wantResp.GetEffectiveGroupId(), resp.GetEffectiveGroupId(), "effective_group_id") - require.Equal(t, wantResp.GetSpendLimitMicros(), resp.GetSpendLimitMicros(), "spend_limit_micros") - require.Equal(t, wantResp.GetCurrentSpendMicros(), resp.GetCurrentSpendMicros(), "current_spend_micros") + require.Equal(t, wantResp.SpendLimitMicros, resp.SpendLimitMicros, "spend_limit_micros") }) } } -// TestGetUserAISpendStatus_Enforcement exercises real-DB scenarios that drive +// TestIsBudgetExceeded_Enforcement exercises real-DB scenarios that drive // enforcement decisions. -func TestGetUserAISpendStatus_Enforcement(t *testing.T) { +func TestIsBudgetExceeded_Enforcement(t *testing.T) { t.Parallel() const groupLimitMicros = 1_000_000 @@ -659,23 +638,21 @@ func TestGetUserAISpendStatus_Enforcement(t *testing.T) { require.NoError(t, err) // Query with period_start 2026-01-01: includes the 2026-01-15 spend, user exceeded. - prevMonthResp, err := srv.GetUserAISpendStatus(ctx, &proto.GetUserAISpendStatusRequest{ + prevMonthResp, err := srv.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{ UserId: user.ID.String(), PeriodStart: timestamppb.New(prevMonth), }) require.NoError(t, err) require.True(t, prevMonthResp.GetExceeded()) - require.Equal(t, int64(1_500_000), prevMonthResp.GetCurrentSpendMicros()) require.Equal(t, int64(groupLimitMicros), prevMonthResp.GetSpendLimitMicros()) // Query with period_start 2026-02-01: excludes the 2026-01-15 spend, user not exceeded. - newMonthResp, err := srv.GetUserAISpendStatus(ctx, &proto.GetUserAISpendStatusRequest{ + newMonthResp, err := srv.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{ UserId: user.ID.String(), PeriodStart: timestamppb.New(newMonth), }) require.NoError(t, err) require.False(t, newMonthResp.GetExceeded()) - require.Equal(t, int64(0), newMonthResp.GetCurrentSpendMicros()) require.Equal(t, int64(groupLimitMicros), newMonthResp.GetSpendLimitMicros()) }) @@ -695,14 +672,13 @@ func TestGetUserAISpendStatus_Enforcement(t *testing.T) { require.NoError(t, err) // User's spend exceeds the group limit. - beforeResp, err := srv.GetUserAISpendStatus(ctx, &proto.GetUserAISpendStatusRequest{ + beforeResp, err := srv.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{ UserId: user.ID.String(), PeriodStart: timestamppb.New(periodStart), }) require.NoError(t, err) require.True(t, beforeResp.GetExceeded()) require.Equal(t, int64(groupLimitMicros), beforeResp.GetSpendLimitMicros()) - require.Equal(t, int64(1_500_000), beforeResp.GetCurrentSpendMicros()) // Add user override with a higher limit on the same group. The override // wins, so the user's spend is now under the effective limit. @@ -714,14 +690,13 @@ func TestGetUserAISpendStatus_Enforcement(t *testing.T) { }) require.NoError(t, err) - afterResp, err := srv.GetUserAISpendStatus(ctx, &proto.GetUserAISpendStatusRequest{ + afterResp, err := srv.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{ UserId: user.ID.String(), PeriodStart: timestamppb.New(periodStart), }) require.NoError(t, err) require.False(t, afterResp.GetExceeded()) require.Equal(t, int64(overrideLimitMicros), afterResp.GetSpendLimitMicros()) - require.Equal(t, int64(1_500_000), afterResp.GetCurrentSpendMicros()) }) } From a12c024a3bd7afeb7c170a1887bac9503915eeee Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Thu, 2 Jul 2026 11:03:28 +0000 Subject: [PATCH 4/6] chore: address comments --- coderd/aibridged/aibridged_test.go | 2 +- coderd/aibridged/aibridgedmock/clientmock.go | 28 ++++++++++---------- coderd/aibridged/http.go | 6 ++--- 3 files changed, 18 insertions(+), 18 deletions(-) diff --git a/coderd/aibridged/aibridged_test.go b/coderd/aibridged/aibridged_test.go index 52975e033b821..4ef3d603e6089 100644 --- a/coderd/aibridged/aibridged_test.go +++ b/coderd/aibridged/aibridged_test.go @@ -140,7 +140,7 @@ func TestServeHTTP_FailureModes(t *testing.T) { SpendLimitMicros: ptr.Ref(int64(1_000)), }, nil) }, - expectedErr: xerrors.New("ai budget of"), + expectedErr: xerrors.New("AI budget of"), expectedStatus: http.StatusForbidden, }, { diff --git a/coderd/aibridged/aibridgedmock/clientmock.go b/coderd/aibridged/aibridgedmock/clientmock.go index 4aa0dd801b750..2d08d54bec21c 100644 --- a/coderd/aibridged/aibridgedmock/clientmock.go +++ b/coderd/aibridged/aibridgedmock/clientmock.go @@ -101,34 +101,34 @@ func (mr *MockDRPCClientMockRecorder) GetMCPServerConfigs(ctx, in any) *gomock.C return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetMCPServerConfigs", reflect.TypeOf((*MockDRPCClient)(nil).GetMCPServerConfigs), ctx, in) } -// IsBudgetExceeded mocks base method. -func (m *MockDRPCClient) IsBudgetExceeded(ctx context.Context, in *proto.IsBudgetExceededRequest) (*proto.IsBudgetExceededResponse, error) { +// IsAuthorized mocks base method. +func (m *MockDRPCClient) IsAuthorized(ctx context.Context, in *proto.IsAuthorizedRequest) (*proto.IsAuthorizedResponse, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "IsBudgetExceeded", ctx, in) - ret0, _ := ret[0].(*proto.IsBudgetExceededResponse) + ret := m.ctrl.Call(m, "IsAuthorized", ctx, in) + ret0, _ := ret[0].(*proto.IsAuthorizedResponse) ret1, _ := ret[1].(error) return ret0, ret1 } -// IsBudgetExceeded indicates an expected call of IsBudgetExceeded. -func (mr *MockDRPCClientMockRecorder) IsBudgetExceeded(ctx, in any) *gomock.Call { +// IsAuthorized indicates an expected call of IsAuthorized. +func (mr *MockDRPCClientMockRecorder) IsAuthorized(ctx, in any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsBudgetExceeded", reflect.TypeOf((*MockDRPCClient)(nil).IsBudgetExceeded), ctx, in) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsAuthorized", reflect.TypeOf((*MockDRPCClient)(nil).IsAuthorized), ctx, in) } -// IsAuthorized mocks base method. -func (m *MockDRPCClient) IsAuthorized(ctx context.Context, in *proto.IsAuthorizedRequest) (*proto.IsAuthorizedResponse, error) { +// IsBudgetExceeded mocks base method. +func (m *MockDRPCClient) IsBudgetExceeded(ctx context.Context, in *proto.IsBudgetExceededRequest) (*proto.IsBudgetExceededResponse, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "IsAuthorized", ctx, in) - ret0, _ := ret[0].(*proto.IsAuthorizedResponse) + ret := m.ctrl.Call(m, "IsBudgetExceeded", ctx, in) + ret0, _ := ret[0].(*proto.IsBudgetExceededResponse) ret1, _ := ret[1].(error) return ret0, ret1 } -// IsAuthorized indicates an expected call of IsAuthorized. -func (mr *MockDRPCClientMockRecorder) IsAuthorized(ctx, in any) *gomock.Call { +// IsBudgetExceeded indicates an expected call of IsBudgetExceeded. +func (mr *MockDRPCClientMockRecorder) IsBudgetExceeded(ctx, in any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsAuthorized", reflect.TypeOf((*MockDRPCClient)(nil).IsAuthorized), ctx, in) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsBudgetExceeded", reflect.TypeOf((*MockDRPCClient)(nil).IsBudgetExceeded), ctx, in) } // RecordInterception mocks base method. diff --git a/coderd/aibridged/http.go b/coderd/aibridged/http.go index f35b895f7813c..3bada27d5852e 100644 --- a/coderd/aibridged/http.go +++ b/coderd/aibridged/http.go @@ -24,7 +24,7 @@ var ( ErrConnect = xerrors.New("could not connect to coderd") ErrUnauthorized = xerrors.New("unauthorized") ErrAcquireRequestHandler = xerrors.New("failed to acquire request handler") - ErrBudgetCheck = xerrors.New("internal server error checking user ai budget") + ErrBudgetCheck = xerrors.New("internal server error checking user AI budget") ) // ServeHTTP is the entrypoint for requests which will be intercepted by AI Bridge. @@ -153,13 +153,13 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) { PeriodStart: timestamppb.New(periodStart), }) if err != nil { - logger.Warn(ctx, "ai budget check failed", slog.Error(err)) + logger.Warn(ctx, "user AI budget check failed", slog.Error(err)) http.Error(rw, ErrBudgetCheck.Error(), http.StatusInternalServerError) return } if budgetStatus.GetExceeded() { http.Error(rw, fmt.Sprintf( - "ai budget of US$%.2f exceeded. Please contact an administrator for more details.", + "AI budget of US$%.2f exceeded. Please contact an administrator for more details.", float64(budgetStatus.GetSpendLimitMicros())/1_000_000, ), http.StatusForbidden) return From 0027c24ea649383af9a8efc70d2d3473d6b797af Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Thu, 2 Jul 2026 11:26:35 +0000 Subject: [PATCH 5/6] chore: minor improvements --- coderd/aibridged/http.go | 6 +- coderd/aibridgedserver/aibridgedserver.go | 58 ++++++++----------- .../aibridgedserver/aibridgedserver_test.go | 23 ++++---- 3 files changed, 39 insertions(+), 48 deletions(-) diff --git a/coderd/aibridged/http.go b/coderd/aibridged/http.go index 3bada27d5852e..c17c26587dee4 100644 --- a/coderd/aibridged/http.go +++ b/coderd/aibridged/http.go @@ -148,7 +148,7 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) { logger = logger.With(slog.F("user_id", id)) periodStart := dbtime.StartOfMonth(dbtime.Now().UTC()) - budgetStatus, err := client.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{ + budgetResp, err := client.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{ UserId: id.String(), PeriodStart: timestamppb.New(periodStart), }) @@ -157,10 +157,10 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) { http.Error(rw, ErrBudgetCheck.Error(), http.StatusInternalServerError) return } - if budgetStatus.GetExceeded() { + if budgetResp.GetExceeded() { http.Error(rw, fmt.Sprintf( "AI budget of US$%.2f exceeded. Please contact an administrator for more details.", - float64(budgetStatus.GetSpendLimitMicros())/1_000_000, + float64(budgetResp.GetSpendLimitMicros())/1_000_000, ), http.StatusForbidden) return } diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index edc521afb4dd4..4a63b14412f51 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -747,30 +747,25 @@ func (s *Server) IsBudgetExceeded(ctx context.Context, in *proto.IsBudgetExceede } periodStart := in.GetPeriodStart().AsTime() - status, err := s.resolveUserAISpendStatus(ctx, userID, periodStart) + userBudget, err := s.checkUserAIBudget(ctx, userID, periodStart) if err != nil { return nil, err } return &proto.IsBudgetExceededResponse{ - Exceeded: status.Exceeded, - SpendLimitMicros: status.SpendLimitMicros, + Exceeded: userBudget.Exceeded, + SpendLimitMicros: userBudget.SpendLimitMicros, }, nil } -// userAISpendStatus is a snapshot of a user's AI spend against their effective -// budget over [PeriodStart, now]. When no budget is configured for the user, -// SpendLimitMicros is nil and the other fields are zero-valued. -type userAISpendStatus struct { - EffectiveGroupID uuid.UUID - PeriodStart time.Time - SpendLimitMicros *int64 - CurrentSpendMicros int64 - Exceeded bool +// userAIBudget is a snapshot of a user's AI budget status. SpendLimitMicros +// is nil when no budget is configured for the user (unlimited). +type userAIBudget struct { + Exceeded bool + SpendLimitMicros *int64 } -// resolveUserAISpendStatus computes the user's AI spend status aggregated from -// periodStart. Returns a zero-valued status (EffectiveGroupID = uuid.Nil) when -// no budget applies to the user. +// checkUserAIBudget evaluates the user's AI budget status aggregated over +// [periodStart, now]. // // Note: there is a potential race condition where two concurrent requests // from the same user can both pass the check if processed in parallel, @@ -781,14 +776,14 @@ type userAISpendStatus struct { // - Cost accounting is advisory, not strict. The goal is to prevent // overages, not build an accounting system. // - Fail-open is acceptable for this case. -func (s *Server) resolveUserAISpendStatus(ctx context.Context, userID uuid.UUID, periodStart time.Time) (userAISpendStatus, error) { +func (s *Server) checkUserAIBudget(ctx context.Context, userID uuid.UUID, periodStart time.Time) (userAIBudget, error) { effectiveBudget, ok, err := budget.ResolveUserAIBudget(ctx, s.store, userID, s.budgetPolicy) if err != nil { - return userAISpendStatus{}, xerrors.Errorf("resolve effective AI budget for user %q with budget policy %q: %w", userID, s.budgetPolicy, err) + return userAIBudget{}, xerrors.Errorf("resolve effective AI budget for user %q with budget policy %q: %w", userID, s.budgetPolicy, err) } if !ok { // No budget configured for the user; return zero-valued status. - return userAISpendStatus{}, nil + return userAIBudget{}, nil } spend, err := s.store.GetUserAISpendSince(ctx, database.GetUserAISpendSinceParams{ @@ -797,31 +792,28 @@ func (s *Server) resolveUserAISpendStatus(ctx context.Context, userID uuid.UUID, PeriodStart: periodStart, }) if err != nil { - return userAISpendStatus{}, xerrors.Errorf("get user AI spend for user %q in group %q: %w", userID, effectiveBudget.GroupID, err) + return userAIBudget{}, xerrors.Errorf("get user AI spend for user %q in group %q: %w", userID, effectiveBudget.GroupID, err) } - status := userAISpendStatus{ - EffectiveGroupID: effectiveBudget.GroupID, - PeriodStart: periodStart, - SpendLimitMicros: ptr.Ref(effectiveBudget.SpendLimitMicros), - CurrentSpendMicros: spend.SpendMicros, - Exceeded: spend.SpendMicros >= effectiveBudget.SpendLimitMicros, - } + exceeded := spend.SpendMicros >= effectiveBudget.SpendLimitMicros logger := s.logger.With( slog.F("user_id", userID), - slog.F("effective_group_id", status.EffectiveGroupID), - slog.F("period_start", status.PeriodStart), - slog.F("current_spend_micros", status.CurrentSpendMicros), - slog.F("spend_limit_micros", *status.SpendLimitMicros), - slog.F("exceeded", status.Exceeded), + slog.F("effective_group_id", effectiveBudget.GroupID), + slog.F("period_start", periodStart), + slog.F("current_spend_micros", spend.SpendMicros), + slog.F("spend_limit_micros", effectiveBudget.SpendLimitMicros), + slog.F("exceeded", exceeded), ) logger.Debug(ctx, "user AI spend status") - if status.Exceeded { + if exceeded { logger.Warn(ctx, "user AI budget exceeded") } - return status, nil + return userAIBudget{ + Exceeded: exceeded, + SpendLimitMicros: ptr.Ref(effectiveBudget.SpendLimitMicros), + }, nil } // GetAIProviders returns the full AI provider set (enabled and disabled) from diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go index 1411c7e68944d..9d4f5b7fb98df 100644 --- a/coderd/aibridgedserver/aibridgedserver_test.go +++ b/coderd/aibridgedserver/aibridgedserver_test.go @@ -428,7 +428,7 @@ func TestIsBudgetExceeded(t *testing.T) { }, }, { - // Group budget resolves, spend below limit (spend 500 < limit 1000). + // Group budget resolves, spend below limit (spend 500 < limit 1000): pass-through. name: "under limit returns not exceeded", setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.IsBudgetExceededResponse { groupID := uuid.New() @@ -445,7 +445,7 @@ func TestIsBudgetExceeded(t *testing.T) { }, }, { - // spend == limit is treated as exceeded (spend 1000 >= limit 1000). + // Group budget resolves, spend at limit (spend 1000 == limit 1000): blocked. name: "at limit returns exceeded", setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.IsBudgetExceededResponse { groupID := uuid.New() @@ -462,9 +462,8 @@ func TestIsBudgetExceeded(t *testing.T) { }, }, { - // A configured limit of 0 is a valid "block-all" setting, distinct - // from "no budget configured". The >= comparison means spend of 0 - // against a limit of 0 is exceeded. + // Limit of 0 is a valid "block-all" setting, distinct from + // "no budget configured": blocked. name: "zero limit blocks all requests", setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.IsBudgetExceededResponse { groupID := uuid.New() @@ -474,7 +473,6 @@ func TestIsBudgetExceeded(t *testing.T) { Return(database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 0}, nil) db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()). Return(database.GetUserAISpendSinceRow{SpendMicros: 0}, nil) - // Blocked user: SpendLimitMicros is explicitly set to 0. return &proto.IsBudgetExceededResponse{ Exceeded: true, SpendLimitMicros: ptr.Ref(int64(0)), @@ -482,7 +480,7 @@ func TestIsBudgetExceeded(t *testing.T) { }, }, { - // Spend above limit (spend 1500 > limit 1000). + // Group budget resolves, spend above limit (spend 1500 > limit 1000): blocked. name: "over limit returns exceeded", setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.IsBudgetExceededResponse { groupID := uuid.New() @@ -499,8 +497,8 @@ func TestIsBudgetExceeded(t *testing.T) { }, }, { - // User override wins: group lookup is skipped, spend is aggregated - // against the override's group (spend 600 > limit 500). + // User override wins, group lookup skipped, spend aggregated against + // the override's group (spend 600 > limit 500): blocked. name: "user override wins over group budget", setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) *proto.IsBudgetExceededResponse { overrideGroupID := uuid.New() @@ -660,13 +658,14 @@ func TestIsBudgetExceeded_Enforcement(t *testing.T) { t.Parallel() ctx, rawDB, srv, user, group := setup(t) - periodStart := dbtime.StartOfMonth(dbtime.Now().UTC()) + // Use fixed dates to keep the test deterministic. + periodStart := time.Date(2026, time.March, 1, 0, 0, 0, 0, time.UTC) - // User spend today, within the current period. + // User spend on 2026-03-15. _, err := rawDB.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, - Day: dbtime.Now().UTC(), + Day: periodStart.AddDate(0, 0, 14), CostMicros: 1_500_000, }) require.NoError(t, err) From e753f18d3f13b884d85e360f50406e57b0d4cc05 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Thu, 2 Jul 2026 15:12:09 +0000 Subject: [PATCH 6/6] fix: run make gen --- coderd/aibridged/proto/aibridged.pb.go | 512 ++++++++++++++++--------- 1 file changed, 340 insertions(+), 172 deletions(-) diff --git a/coderd/aibridged/proto/aibridged.pb.go b/coderd/aibridged/proto/aibridged.pb.go index b4be124d9729b..f45ef36084c02 100644 --- a/coderd/aibridged/proto/aibridged.pb.go +++ b/coderd/aibridged/proto/aibridged.pb.go @@ -1268,6 +1268,123 @@ func (x *IsAuthorizedResponse) GetUsername() string { return "" } +type IsBudgetExceededRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + UserId string `protobuf:"bytes,1,opt,name=user_id,json=userId,proto3" json:"user_id,omitempty"` // UUID + // The spend aggregation window is [period_start, now]. + PeriodStart *timestamppb.Timestamp `protobuf:"bytes,2,opt,name=period_start,json=periodStart,proto3" json:"period_start,omitempty"` +} + +func (x *IsBudgetExceededRequest) Reset() { + *x = IsBudgetExceededRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_coderd_aibridged_proto_aibridged_proto_msgTypes[19] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *IsBudgetExceededRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*IsBudgetExceededRequest) ProtoMessage() {} + +func (x *IsBudgetExceededRequest) ProtoReflect() protoreflect.Message { + mi := &file_coderd_aibridged_proto_aibridged_proto_msgTypes[19] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use IsBudgetExceededRequest.ProtoReflect.Descriptor instead. +func (*IsBudgetExceededRequest) Descriptor() ([]byte, []int) { + return file_coderd_aibridged_proto_aibridged_proto_rawDescGZIP(), []int{19} +} + +func (x *IsBudgetExceededRequest) GetUserId() string { + if x != nil { + return x.UserId + } + return "" +} + +func (x *IsBudgetExceededRequest) GetPeriodStart() *timestamppb.Timestamp { + if x != nil { + return x.PeriodStart + } + return nil +} + +type IsBudgetExceededResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + // exceeded is true when the user's aggregated spend has reached the + // effective limit. False when no budget is configured for the user OR + // they are within their limit. + Exceeded bool `protobuf:"varint,1,opt,name=exceeded,proto3" json:"exceeded,omitempty"` + // spend_limit_micros is the effective spend limit in micro-units. + // Unset when no budget is configured for the user (unlimited). + // 0 when a group is explicitly configured with a 0 limit (blocked). + SpendLimitMicros *int64 `protobuf:"varint,2,opt,name=spend_limit_micros,json=spendLimitMicros,proto3,oneof" json:"spend_limit_micros,omitempty"` +} + +func (x *IsBudgetExceededResponse) Reset() { + *x = IsBudgetExceededResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_coderd_aibridged_proto_aibridged_proto_msgTypes[20] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *IsBudgetExceededResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*IsBudgetExceededResponse) ProtoMessage() {} + +func (x *IsBudgetExceededResponse) ProtoReflect() protoreflect.Message { + mi := &file_coderd_aibridged_proto_aibridged_proto_msgTypes[20] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use IsBudgetExceededResponse.ProtoReflect.Descriptor instead. +func (*IsBudgetExceededResponse) Descriptor() ([]byte, []int) { + return file_coderd_aibridged_proto_aibridged_proto_rawDescGZIP(), []int{20} +} + +func (x *IsBudgetExceededResponse) GetExceeded() bool { + if x != nil { + return x.Exceeded + } + return false +} + +func (x *IsBudgetExceededResponse) GetSpendLimitMicros() int64 { + if x != nil && x.SpendLimitMicros != nil { + return *x.SpendLimitMicros + } + return 0 +} + type GetAIProvidersRequest struct { state protoimpl.MessageState sizeCache protoimpl.SizeCache @@ -1277,7 +1394,7 @@ type GetAIProvidersRequest struct { func (x *GetAIProvidersRequest) Reset() { *x = GetAIProvidersRequest{} if protoimpl.UnsafeEnabled { - mi := &file_coderd_aibridged_proto_aibridged_proto_msgTypes[19] + mi := &file_coderd_aibridged_proto_aibridged_proto_msgTypes[21] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1290,7 +1407,7 @@ func (x *GetAIProvidersRequest) String() string { func (*GetAIProvidersRequest) ProtoMessage() {} func (x *GetAIProvidersRequest) ProtoReflect() protoreflect.Message { - mi := &file_coderd_aibridged_proto_aibridged_proto_msgTypes[19] + mi := &file_coderd_aibridged_proto_aibridged_proto_msgTypes[21] if protoimpl.UnsafeEnabled && x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1303,7 +1420,7 @@ func (x *GetAIProvidersRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use GetAIProvidersRequest.ProtoReflect.Descriptor instead. func (*GetAIProvidersRequest) Descriptor() ([]byte, []int) { - return file_coderd_aibridged_proto_aibridged_proto_rawDescGZIP(), []int{19} + return file_coderd_aibridged_proto_aibridged_proto_rawDescGZIP(), []int{21} } type GetAIProvidersResponse struct { @@ -1317,7 +1434,7 @@ type GetAIProvidersResponse struct { func (x *GetAIProvidersResponse) Reset() { *x = GetAIProvidersResponse{} if protoimpl.UnsafeEnabled { - mi := &file_coderd_aibridged_proto_aibridged_proto_msgTypes[20] + mi := &file_coderd_aibridged_proto_aibridged_proto_msgTypes[22] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1330,7 +1447,7 @@ func (x *GetAIProvidersResponse) String() string { func (*GetAIProvidersResponse) ProtoMessage() {} func (x *GetAIProvidersResponse) ProtoReflect() protoreflect.Message { - mi := &file_coderd_aibridged_proto_aibridged_proto_msgTypes[20] + mi := &file_coderd_aibridged_proto_aibridged_proto_msgTypes[22] if protoimpl.UnsafeEnabled && x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1343,7 +1460,7 @@ func (x *GetAIProvidersResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use GetAIProvidersResponse.ProtoReflect.Descriptor instead. func (*GetAIProvidersResponse) Descriptor() ([]byte, []int) { - return file_coderd_aibridged_proto_aibridged_proto_rawDescGZIP(), []int{20} + return file_coderd_aibridged_proto_aibridged_proto_rawDescGZIP(), []int{22} } func (x *GetAIProvidersResponse) GetProviders() []*AIProvider { @@ -1372,7 +1489,7 @@ type AIProvider struct { func (x *AIProvider) Reset() { *x = AIProvider{} if protoimpl.UnsafeEnabled { - mi := &file_coderd_aibridged_proto_aibridged_proto_msgTypes[21] + mi := &file_coderd_aibridged_proto_aibridged_proto_msgTypes[23] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1385,7 +1502,7 @@ func (x *AIProvider) String() string { func (*AIProvider) ProtoMessage() {} func (x *AIProvider) ProtoReflect() protoreflect.Message { - mi := &file_coderd_aibridged_proto_aibridged_proto_msgTypes[21] + mi := &file_coderd_aibridged_proto_aibridged_proto_msgTypes[23] if protoimpl.UnsafeEnabled && x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1398,7 +1515,7 @@ func (x *AIProvider) ProtoReflect() protoreflect.Message { // Deprecated: Use AIProvider.ProtoReflect.Descriptor instead. func (*AIProvider) Descriptor() ([]byte, []int) { - return file_coderd_aibridged_proto_aibridged_proto_rawDescGZIP(), []int{21} + return file_coderd_aibridged_proto_aibridged_proto_rawDescGZIP(), []int{23} } func (x *AIProvider) GetName() string { @@ -1460,7 +1577,7 @@ type AIProviderKindBedrock struct { func (x *AIProviderKindBedrock) Reset() { *x = AIProviderKindBedrock{} if protoimpl.UnsafeEnabled { - mi := &file_coderd_aibridged_proto_aibridged_proto_msgTypes[22] + mi := &file_coderd_aibridged_proto_aibridged_proto_msgTypes[24] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1473,7 +1590,7 @@ func (x *AIProviderKindBedrock) String() string { func (*AIProviderKindBedrock) ProtoMessage() {} func (x *AIProviderKindBedrock) ProtoReflect() protoreflect.Message { - mi := &file_coderd_aibridged_proto_aibridged_proto_msgTypes[22] + mi := &file_coderd_aibridged_proto_aibridged_proto_msgTypes[24] if protoimpl.UnsafeEnabled && x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1486,7 +1603,7 @@ func (x *AIProviderKindBedrock) ProtoReflect() protoreflect.Message { // Deprecated: Use AIProviderKindBedrock.ProtoReflect.Descriptor instead. func (*AIProviderKindBedrock) Descriptor() ([]byte, []int) { - return file_coderd_aibridged_proto_aibridged_proto_rawDescGZIP(), []int{22} + return file_coderd_aibridged_proto_aibridged_proto_rawDescGZIP(), []int{24} } func (x *AIProviderKindBedrock) GetRegion() string { @@ -1793,105 +1910,126 @@ var file_coderd_aibridged_proto_aibridged_proto_rawDesc = []byte{ 0x79, 0x5f, 0x69, 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x61, 0x70, 0x69, 0x4b, 0x65, 0x79, 0x49, 0x64, 0x12, 0x1a, 0x0a, 0x08, 0x75, 0x73, 0x65, 0x72, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x75, 0x73, 0x65, 0x72, 0x6e, 0x61, 0x6d, 0x65, - 0x22, 0x17, 0x0a, 0x15, 0x47, 0x65, 0x74, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, - 0x72, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x22, 0x49, 0x0a, 0x16, 0x47, 0x65, 0x74, - 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, - 0x6e, 0x73, 0x65, 0x12, 0x2f, 0x0a, 0x09, 0x70, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, - 0x18, 0x01, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x11, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x41, - 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x52, 0x09, 0x70, 0x72, 0x6f, 0x76, 0x69, - 0x64, 0x65, 0x72, 0x73, 0x22, 0xb5, 0x01, 0x0a, 0x0a, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, - 0x64, 0x65, 0x72, 0x12, 0x12, 0x0a, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, - 0x09, 0x52, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x12, 0x12, 0x0a, 0x04, 0x74, 0x79, 0x70, 0x65, 0x18, - 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x74, 0x79, 0x70, 0x65, 0x12, 0x18, 0x0a, 0x07, 0x65, - 0x6e, 0x61, 0x62, 0x6c, 0x65, 0x64, 0x18, 0x03, 0x20, 0x01, 0x28, 0x08, 0x52, 0x07, 0x65, 0x6e, - 0x61, 0x62, 0x6c, 0x65, 0x64, 0x12, 0x19, 0x0a, 0x08, 0x62, 0x61, 0x73, 0x65, 0x5f, 0x75, 0x72, - 0x6c, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x07, 0x62, 0x61, 0x73, 0x65, 0x55, 0x72, 0x6c, - 0x12, 0x12, 0x0a, 0x04, 0x6b, 0x65, 0x79, 0x73, 0x18, 0x05, 0x20, 0x03, 0x28, 0x09, 0x52, 0x04, - 0x6b, 0x65, 0x79, 0x73, 0x12, 0x36, 0x0a, 0x07, 0x62, 0x65, 0x64, 0x72, 0x6f, 0x63, 0x6b, 0x18, - 0x06, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1c, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x41, 0x49, - 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x4b, 0x69, 0x6e, 0x64, 0x42, 0x65, 0x64, 0x72, - 0x6f, 0x63, 0x6b, 0x52, 0x07, 0x62, 0x65, 0x64, 0x72, 0x6f, 0x63, 0x6b, 0x22, 0xf6, 0x01, 0x0a, - 0x15, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x4b, 0x69, 0x6e, 0x64, 0x42, - 0x65, 0x64, 0x72, 0x6f, 0x63, 0x6b, 0x12, 0x16, 0x0a, 0x06, 0x72, 0x65, 0x67, 0x69, 0x6f, 0x6e, - 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x72, 0x65, 0x67, 0x69, 0x6f, 0x6e, 0x12, 0x1d, - 0x0a, 0x0a, 0x61, 0x63, 0x63, 0x65, 0x73, 0x73, 0x5f, 0x6b, 0x65, 0x79, 0x18, 0x02, 0x20, 0x01, - 0x28, 0x09, 0x52, 0x09, 0x61, 0x63, 0x63, 0x65, 0x73, 0x73, 0x4b, 0x65, 0x79, 0x12, 0x2a, 0x0a, - 0x11, 0x61, 0x63, 0x63, 0x65, 0x73, 0x73, 0x5f, 0x6b, 0x65, 0x79, 0x5f, 0x73, 0x65, 0x63, 0x72, - 0x65, 0x74, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0f, 0x61, 0x63, 0x63, 0x65, 0x73, 0x73, - 0x4b, 0x65, 0x79, 0x53, 0x65, 0x63, 0x72, 0x65, 0x74, 0x12, 0x14, 0x0a, 0x05, 0x6d, 0x6f, 0x64, - 0x65, 0x6c, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, 0x6d, 0x6f, 0x64, 0x65, 0x6c, 0x12, - 0x28, 0x0a, 0x10, 0x73, 0x6d, 0x61, 0x6c, 0x6c, 0x5f, 0x66, 0x61, 0x73, 0x74, 0x5f, 0x6d, 0x6f, - 0x64, 0x65, 0x6c, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0e, 0x73, 0x6d, 0x61, 0x6c, 0x6c, - 0x46, 0x61, 0x73, 0x74, 0x4d, 0x6f, 0x64, 0x65, 0x6c, 0x12, 0x19, 0x0a, 0x08, 0x72, 0x6f, 0x6c, - 0x65, 0x5f, 0x61, 0x72, 0x6e, 0x18, 0x06, 0x20, 0x01, 0x28, 0x09, 0x52, 0x07, 0x72, 0x6f, 0x6c, - 0x65, 0x41, 0x72, 0x6e, 0x12, 0x1f, 0x0a, 0x0b, 0x65, 0x78, 0x74, 0x65, 0x72, 0x6e, 0x61, 0x6c, - 0x5f, 0x69, 0x64, 0x18, 0x07, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0a, 0x65, 0x78, 0x74, 0x65, 0x72, - 0x6e, 0x61, 0x6c, 0x49, 0x64, 0x32, 0xa9, 0x04, 0x0a, 0x08, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, - 0x65, 0x72, 0x12, 0x59, 0x0a, 0x12, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, 0x74, 0x65, - 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x20, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, - 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, 0x74, 0x65, 0x72, 0x63, 0x65, 0x70, 0x74, - 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x21, 0x2e, 0x70, 0x72, 0x6f, - 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, 0x74, 0x65, 0x72, 0x63, 0x65, - 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x68, 0x0a, - 0x17, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, 0x74, 0x65, 0x72, 0x63, 0x65, 0x70, 0x74, - 0x69, 0x6f, 0x6e, 0x45, 0x6e, 0x64, 0x65, 0x64, 0x12, 0x25, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, - 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, 0x74, 0x65, 0x72, 0x63, 0x65, 0x70, 0x74, - 0x69, 0x6f, 0x6e, 0x45, 0x6e, 0x64, 0x65, 0x64, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, - 0x26, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, - 0x74, 0x65, 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x45, 0x6e, 0x64, 0x65, 0x64, 0x52, - 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x53, 0x0a, 0x10, 0x52, 0x65, 0x63, 0x6f, 0x72, - 0x64, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x55, 0x73, 0x61, 0x67, 0x65, 0x12, 0x1e, 0x2e, 0x70, 0x72, - 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x55, - 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1f, 0x2e, 0x70, 0x72, - 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x55, - 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x56, 0x0a, 0x11, + 0x22, 0x71, 0x0a, 0x17, 0x49, 0x73, 0x42, 0x75, 0x64, 0x67, 0x65, 0x74, 0x45, 0x78, 0x63, 0x65, + 0x65, 0x64, 0x65, 0x64, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x17, 0x0a, 0x07, 0x75, + 0x73, 0x65, 0x72, 0x5f, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x75, 0x73, + 0x65, 0x72, 0x49, 0x64, 0x12, 0x3d, 0x0a, 0x0c, 0x70, 0x65, 0x72, 0x69, 0x6f, 0x64, 0x5f, 0x73, + 0x74, 0x61, 0x72, 0x74, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1a, 0x2e, 0x67, 0x6f, 0x6f, + 0x67, 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x54, 0x69, 0x6d, + 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x52, 0x0b, 0x70, 0x65, 0x72, 0x69, 0x6f, 0x64, 0x53, 0x74, + 0x61, 0x72, 0x74, 0x22, 0x80, 0x01, 0x0a, 0x18, 0x49, 0x73, 0x42, 0x75, 0x64, 0x67, 0x65, 0x74, + 0x45, 0x78, 0x63, 0x65, 0x65, 0x64, 0x65, 0x64, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, + 0x12, 0x1a, 0x0a, 0x08, 0x65, 0x78, 0x63, 0x65, 0x65, 0x64, 0x65, 0x64, 0x18, 0x01, 0x20, 0x01, + 0x28, 0x08, 0x52, 0x08, 0x65, 0x78, 0x63, 0x65, 0x65, 0x64, 0x65, 0x64, 0x12, 0x31, 0x0a, 0x12, + 0x73, 0x70, 0x65, 0x6e, 0x64, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x5f, 0x6d, 0x69, 0x63, 0x72, + 0x6f, 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x03, 0x48, 0x00, 0x52, 0x10, 0x73, 0x70, 0x65, 0x6e, + 0x64, 0x4c, 0x69, 0x6d, 0x69, 0x74, 0x4d, 0x69, 0x63, 0x72, 0x6f, 0x73, 0x88, 0x01, 0x01, 0x42, + 0x15, 0x0a, 0x13, 0x5f, 0x73, 0x70, 0x65, 0x6e, 0x64, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x5f, + 0x6d, 0x69, 0x63, 0x72, 0x6f, 0x73, 0x22, 0x17, 0x0a, 0x15, 0x47, 0x65, 0x74, 0x41, 0x49, 0x50, + 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x22, + 0x49, 0x0a, 0x16, 0x47, 0x65, 0x74, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, + 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x2f, 0x0a, 0x09, 0x70, 0x72, 0x6f, + 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x18, 0x01, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x11, 0x2e, 0x70, + 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x52, + 0x09, 0x70, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x22, 0xb5, 0x01, 0x0a, 0x0a, 0x41, + 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x12, 0x12, 0x0a, 0x04, 0x6e, 0x61, 0x6d, + 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x12, 0x12, 0x0a, + 0x04, 0x74, 0x79, 0x70, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x74, 0x79, 0x70, + 0x65, 0x12, 0x18, 0x0a, 0x07, 0x65, 0x6e, 0x61, 0x62, 0x6c, 0x65, 0x64, 0x18, 0x03, 0x20, 0x01, + 0x28, 0x08, 0x52, 0x07, 0x65, 0x6e, 0x61, 0x62, 0x6c, 0x65, 0x64, 0x12, 0x19, 0x0a, 0x08, 0x62, + 0x61, 0x73, 0x65, 0x5f, 0x75, 0x72, 0x6c, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x07, 0x62, + 0x61, 0x73, 0x65, 0x55, 0x72, 0x6c, 0x12, 0x12, 0x0a, 0x04, 0x6b, 0x65, 0x79, 0x73, 0x18, 0x05, + 0x20, 0x03, 0x28, 0x09, 0x52, 0x04, 0x6b, 0x65, 0x79, 0x73, 0x12, 0x36, 0x0a, 0x07, 0x62, 0x65, + 0x64, 0x72, 0x6f, 0x63, 0x6b, 0x18, 0x06, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1c, 0x2e, 0x70, 0x72, + 0x6f, 0x74, 0x6f, 0x2e, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x4b, 0x69, + 0x6e, 0x64, 0x42, 0x65, 0x64, 0x72, 0x6f, 0x63, 0x6b, 0x52, 0x07, 0x62, 0x65, 0x64, 0x72, 0x6f, + 0x63, 0x6b, 0x22, 0xf6, 0x01, 0x0a, 0x15, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, + 0x72, 0x4b, 0x69, 0x6e, 0x64, 0x42, 0x65, 0x64, 0x72, 0x6f, 0x63, 0x6b, 0x12, 0x16, 0x0a, 0x06, + 0x72, 0x65, 0x67, 0x69, 0x6f, 0x6e, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x72, 0x65, + 0x67, 0x69, 0x6f, 0x6e, 0x12, 0x1d, 0x0a, 0x0a, 0x61, 0x63, 0x63, 0x65, 0x73, 0x73, 0x5f, 0x6b, + 0x65, 0x79, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x09, 0x61, 0x63, 0x63, 0x65, 0x73, 0x73, + 0x4b, 0x65, 0x79, 0x12, 0x2a, 0x0a, 0x11, 0x61, 0x63, 0x63, 0x65, 0x73, 0x73, 0x5f, 0x6b, 0x65, + 0x79, 0x5f, 0x73, 0x65, 0x63, 0x72, 0x65, 0x74, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0f, + 0x61, 0x63, 0x63, 0x65, 0x73, 0x73, 0x4b, 0x65, 0x79, 0x53, 0x65, 0x63, 0x72, 0x65, 0x74, 0x12, + 0x14, 0x0a, 0x05, 0x6d, 0x6f, 0x64, 0x65, 0x6c, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, + 0x6d, 0x6f, 0x64, 0x65, 0x6c, 0x12, 0x28, 0x0a, 0x10, 0x73, 0x6d, 0x61, 0x6c, 0x6c, 0x5f, 0x66, + 0x61, 0x73, 0x74, 0x5f, 0x6d, 0x6f, 0x64, 0x65, 0x6c, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, + 0x0e, 0x73, 0x6d, 0x61, 0x6c, 0x6c, 0x46, 0x61, 0x73, 0x74, 0x4d, 0x6f, 0x64, 0x65, 0x6c, 0x12, + 0x19, 0x0a, 0x08, 0x72, 0x6f, 0x6c, 0x65, 0x5f, 0x61, 0x72, 0x6e, 0x18, 0x06, 0x20, 0x01, 0x28, + 0x09, 0x52, 0x07, 0x72, 0x6f, 0x6c, 0x65, 0x41, 0x72, 0x6e, 0x12, 0x1f, 0x0a, 0x0b, 0x65, 0x78, + 0x74, 0x65, 0x72, 0x6e, 0x61, 0x6c, 0x5f, 0x69, 0x64, 0x18, 0x07, 0x20, 0x01, 0x28, 0x09, 0x52, + 0x0a, 0x65, 0x78, 0x74, 0x65, 0x72, 0x6e, 0x61, 0x6c, 0x49, 0x64, 0x32, 0xa9, 0x04, 0x0a, 0x08, + 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x65, 0x72, 0x12, 0x59, 0x0a, 0x12, 0x52, 0x65, 0x63, 0x6f, + 0x72, 0x64, 0x49, 0x6e, 0x74, 0x65, 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x20, + 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, 0x74, + 0x65, 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, + 0x1a, 0x21, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, + 0x6e, 0x74, 0x65, 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, + 0x6e, 0x73, 0x65, 0x12, 0x68, 0x0a, 0x17, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, 0x74, + 0x65, 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x45, 0x6e, 0x64, 0x65, 0x64, 0x12, 0x25, + 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, 0x74, + 0x65, 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x45, 0x6e, 0x64, 0x65, 0x64, 0x52, 0x65, + 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x26, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, + 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, 0x74, 0x65, 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, + 0x45, 0x6e, 0x64, 0x65, 0x64, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x53, 0x0a, + 0x10, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x55, 0x73, 0x61, 0x67, + 0x65, 0x12, 0x1e, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, + 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, + 0x74, 0x1a, 0x1f, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, + 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, + 0x73, 0x65, 0x12, 0x56, 0x0a, 0x11, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x50, 0x72, 0x6f, 0x6d, + 0x70, 0x74, 0x55, 0x73, 0x61, 0x67, 0x65, 0x12, 0x1f, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x50, 0x72, 0x6f, 0x6d, 0x70, 0x74, 0x55, 0x73, 0x61, 0x67, - 0x65, 0x12, 0x1f, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, - 0x50, 0x72, 0x6f, 0x6d, 0x70, 0x74, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, - 0x73, 0x74, 0x1a, 0x20, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, - 0x64, 0x50, 0x72, 0x6f, 0x6d, 0x70, 0x74, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x73, 0x70, - 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x50, 0x0a, 0x0f, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, - 0x6f, 0x6c, 0x55, 0x73, 0x61, 0x67, 0x65, 0x12, 0x1d, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, - 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6f, 0x6c, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, - 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1e, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, - 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6f, 0x6c, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, - 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x59, 0x0a, 0x12, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, - 0x4d, 0x6f, 0x64, 0x65, 0x6c, 0x54, 0x68, 0x6f, 0x75, 0x67, 0x68, 0x74, 0x12, 0x20, 0x2e, 0x70, - 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x4d, 0x6f, 0x64, 0x65, 0x6c, - 0x54, 0x68, 0x6f, 0x75, 0x67, 0x68, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x21, - 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x4d, 0x6f, 0x64, - 0x65, 0x6c, 0x54, 0x68, 0x6f, 0x75, 0x67, 0x68, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, - 0x65, 0x32, 0xeb, 0x01, 0x0a, 0x0f, 0x4d, 0x43, 0x50, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x75, - 0x72, 0x61, 0x74, 0x6f, 0x72, 0x12, 0x5c, 0x0a, 0x13, 0x47, 0x65, 0x74, 0x4d, 0x43, 0x50, 0x53, - 0x65, 0x72, 0x76, 0x65, 0x72, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x73, 0x12, 0x21, 0x2e, 0x70, - 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x4d, 0x43, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, - 0x72, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, - 0x22, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x4d, 0x43, 0x50, 0x53, 0x65, - 0x72, 0x76, 0x65, 0x72, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, - 0x6e, 0x73, 0x65, 0x12, 0x7a, 0x0a, 0x1d, 0x47, 0x65, 0x74, 0x4d, 0x43, 0x50, 0x53, 0x65, 0x72, - 0x76, 0x65, 0x72, 0x41, 0x63, 0x63, 0x65, 0x73, 0x73, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x42, - 0x61, 0x74, 0x63, 0x68, 0x12, 0x2b, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, + 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x20, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, + 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x50, 0x72, 0x6f, 0x6d, 0x70, 0x74, 0x55, 0x73, 0x61, + 0x67, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x50, 0x0a, 0x0f, 0x52, 0x65, + 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6f, 0x6c, 0x55, 0x73, 0x61, 0x67, 0x65, 0x12, 0x1d, 0x2e, + 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6f, 0x6c, + 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1e, 0x2e, 0x70, + 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6f, 0x6c, 0x55, + 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x59, 0x0a, 0x12, + 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x4d, 0x6f, 0x64, 0x65, 0x6c, 0x54, 0x68, 0x6f, 0x75, 0x67, + 0x68, 0x74, 0x12, 0x20, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, + 0x64, 0x4d, 0x6f, 0x64, 0x65, 0x6c, 0x54, 0x68, 0x6f, 0x75, 0x67, 0x68, 0x74, 0x52, 0x65, 0x71, + 0x75, 0x65, 0x73, 0x74, 0x1a, 0x21, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, + 0x6f, 0x72, 0x64, 0x4d, 0x6f, 0x64, 0x65, 0x6c, 0x54, 0x68, 0x6f, 0x75, 0x67, 0x68, 0x74, 0x52, + 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x32, 0xeb, 0x01, 0x0a, 0x0f, 0x4d, 0x43, 0x50, 0x43, + 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x75, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x12, 0x5c, 0x0a, 0x13, 0x47, + 0x65, 0x74, 0x4d, 0x43, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x43, 0x6f, 0x6e, 0x66, 0x69, + 0x67, 0x73, 0x12, 0x21, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x4d, 0x43, + 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x73, 0x52, 0x65, + 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x22, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, + 0x74, 0x4d, 0x43, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, + 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x7a, 0x0a, 0x1d, 0x47, 0x65, 0x74, 0x4d, 0x43, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x41, 0x63, 0x63, 0x65, 0x73, 0x73, 0x54, - 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x42, 0x61, 0x74, 0x63, 0x68, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, - 0x74, 0x1a, 0x2c, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x4d, 0x43, 0x50, - 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x41, 0x63, 0x63, 0x65, 0x73, 0x73, 0x54, 0x6f, 0x6b, 0x65, - 0x6e, 0x73, 0x42, 0x61, 0x74, 0x63, 0x68, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x32, - 0x55, 0x0a, 0x0a, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, 0x69, 0x7a, 0x65, 0x72, 0x12, 0x47, 0x0a, - 0x0c, 0x49, 0x73, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, 0x69, 0x7a, 0x65, 0x64, 0x12, 0x1a, 0x2e, - 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x49, 0x73, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, 0x69, 0x7a, - 0x65, 0x64, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1b, 0x2e, 0x70, 0x72, 0x6f, 0x74, - 0x6f, 0x2e, 0x49, 0x73, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, 0x69, 0x7a, 0x65, 0x64, 0x52, 0x65, - 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x32, 0x65, 0x0a, 0x14, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, - 0x65, 0x72, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x75, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x12, 0x4d, - 0x0a, 0x0e, 0x47, 0x65, 0x74, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, - 0x12, 0x1c, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x41, 0x49, 0x50, 0x72, - 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1d, - 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, - 0x69, 0x64, 0x65, 0x72, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x42, 0x32, 0x5a, - 0x30, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x63, 0x6f, 0x64, 0x65, - 0x72, 0x2f, 0x63, 0x6f, 0x64, 0x65, 0x72, 0x2f, 0x76, 0x32, 0x2f, 0x63, 0x6f, 0x64, 0x65, 0x72, - 0x64, 0x2f, 0x61, 0x69, 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, 0x64, 0x2f, 0x70, 0x72, 0x6f, 0x74, - 0x6f, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, + 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x42, 0x61, 0x74, 0x63, 0x68, 0x12, 0x2b, 0x2e, 0x70, 0x72, 0x6f, + 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x4d, 0x43, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x41, + 0x63, 0x63, 0x65, 0x73, 0x73, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x42, 0x61, 0x74, 0x63, 0x68, + 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x2c, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, + 0x47, 0x65, 0x74, 0x4d, 0x43, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x41, 0x63, 0x63, 0x65, + 0x73, 0x73, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x42, 0x61, 0x74, 0x63, 0x68, 0x52, 0x65, 0x73, + 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x32, 0xaa, 0x01, 0x0a, 0x0a, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, + 0x69, 0x7a, 0x65, 0x72, 0x12, 0x47, 0x0a, 0x0c, 0x49, 0x73, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, + 0x69, 0x7a, 0x65, 0x64, 0x12, 0x1a, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x49, 0x73, 0x41, + 0x75, 0x74, 0x68, 0x6f, 0x72, 0x69, 0x7a, 0x65, 0x64, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, + 0x1a, 0x1b, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x49, 0x73, 0x41, 0x75, 0x74, 0x68, 0x6f, + 0x72, 0x69, 0x7a, 0x65, 0x64, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x53, 0x0a, + 0x10, 0x49, 0x73, 0x42, 0x75, 0x64, 0x67, 0x65, 0x74, 0x45, 0x78, 0x63, 0x65, 0x65, 0x64, 0x65, + 0x64, 0x12, 0x1e, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x49, 0x73, 0x42, 0x75, 0x64, 0x67, + 0x65, 0x74, 0x45, 0x78, 0x63, 0x65, 0x65, 0x64, 0x65, 0x64, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, + 0x74, 0x1a, 0x1f, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x49, 0x73, 0x42, 0x75, 0x64, 0x67, + 0x65, 0x74, 0x45, 0x78, 0x63, 0x65, 0x65, 0x64, 0x65, 0x64, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, + 0x73, 0x65, 0x32, 0x65, 0x0a, 0x14, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x43, 0x6f, + 0x6e, 0x66, 0x69, 0x67, 0x75, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x12, 0x4d, 0x0a, 0x0e, 0x47, 0x65, + 0x74, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x12, 0x1c, 0x2e, 0x70, + 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, + 0x65, 0x72, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1d, 0x2e, 0x70, 0x72, 0x6f, + 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, + 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x42, 0x32, 0x5a, 0x30, 0x67, 0x69, 0x74, + 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x63, 0x6f, 0x64, 0x65, 0x72, 0x2f, 0x63, 0x6f, + 0x64, 0x65, 0x72, 0x2f, 0x76, 0x32, 0x2f, 0x63, 0x6f, 0x64, 0x65, 0x72, 0x64, 0x2f, 0x61, 0x69, + 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, 0x64, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x06, 0x70, + 0x72, 0x6f, 0x74, 0x6f, 0x33, } var ( @@ -1906,7 +2044,7 @@ func file_coderd_aibridged_proto_aibridged_proto_rawDescGZIP() []byte { return file_coderd_aibridged_proto_aibridged_proto_rawDescData } -var file_coderd_aibridged_proto_aibridged_proto_msgTypes = make([]protoimpl.MessageInfo, 30) +var file_coderd_aibridged_proto_aibridged_proto_msgTypes = make([]protoimpl.MessageInfo, 32) var file_coderd_aibridged_proto_aibridged_proto_goTypes = []interface{}{ (*RecordInterceptionRequest)(nil), // 0: proto.RecordInterceptionRequest (*RecordInterceptionResponse)(nil), // 1: proto.RecordInterceptionResponse @@ -1927,68 +2065,73 @@ var file_coderd_aibridged_proto_aibridged_proto_goTypes = []interface{}{ (*GetMCPServerAccessTokensBatchResponse)(nil), // 16: proto.GetMCPServerAccessTokensBatchResponse (*IsAuthorizedRequest)(nil), // 17: proto.IsAuthorizedRequest (*IsAuthorizedResponse)(nil), // 18: proto.IsAuthorizedResponse - (*GetAIProvidersRequest)(nil), // 19: proto.GetAIProvidersRequest - (*GetAIProvidersResponse)(nil), // 20: proto.GetAIProvidersResponse - (*AIProvider)(nil), // 21: proto.AIProvider - (*AIProviderKindBedrock)(nil), // 22: proto.AIProviderKindBedrock - nil, // 23: proto.RecordInterceptionRequest.MetadataEntry - nil, // 24: proto.RecordTokenUsageRequest.MetadataEntry - nil, // 25: proto.RecordPromptUsageRequest.MetadataEntry - nil, // 26: proto.RecordToolUsageRequest.MetadataEntry - nil, // 27: proto.RecordModelThoughtRequest.MetadataEntry - nil, // 28: proto.GetMCPServerAccessTokensBatchResponse.AccessTokensEntry - nil, // 29: proto.GetMCPServerAccessTokensBatchResponse.ErrorsEntry - (*timestamppb.Timestamp)(nil), // 30: google.protobuf.Timestamp - (*anypb.Any)(nil), // 31: google.protobuf.Any + (*IsBudgetExceededRequest)(nil), // 19: proto.IsBudgetExceededRequest + (*IsBudgetExceededResponse)(nil), // 20: proto.IsBudgetExceededResponse + (*GetAIProvidersRequest)(nil), // 21: proto.GetAIProvidersRequest + (*GetAIProvidersResponse)(nil), // 22: proto.GetAIProvidersResponse + (*AIProvider)(nil), // 23: proto.AIProvider + (*AIProviderKindBedrock)(nil), // 24: proto.AIProviderKindBedrock + nil, // 25: proto.RecordInterceptionRequest.MetadataEntry + nil, // 26: proto.RecordTokenUsageRequest.MetadataEntry + nil, // 27: proto.RecordPromptUsageRequest.MetadataEntry + nil, // 28: proto.RecordToolUsageRequest.MetadataEntry + nil, // 29: proto.RecordModelThoughtRequest.MetadataEntry + nil, // 30: proto.GetMCPServerAccessTokensBatchResponse.AccessTokensEntry + nil, // 31: proto.GetMCPServerAccessTokensBatchResponse.ErrorsEntry + (*timestamppb.Timestamp)(nil), // 32: google.protobuf.Timestamp + (*anypb.Any)(nil), // 33: google.protobuf.Any } var file_coderd_aibridged_proto_aibridged_proto_depIdxs = []int32{ - 23, // 0: proto.RecordInterceptionRequest.metadata:type_name -> proto.RecordInterceptionRequest.MetadataEntry - 30, // 1: proto.RecordInterceptionRequest.started_at:type_name -> google.protobuf.Timestamp - 30, // 2: proto.RecordInterceptionEndedRequest.ended_at:type_name -> google.protobuf.Timestamp - 24, // 3: proto.RecordTokenUsageRequest.metadata:type_name -> proto.RecordTokenUsageRequest.MetadataEntry - 30, // 4: proto.RecordTokenUsageRequest.created_at:type_name -> google.protobuf.Timestamp - 25, // 5: proto.RecordPromptUsageRequest.metadata:type_name -> proto.RecordPromptUsageRequest.MetadataEntry - 30, // 6: proto.RecordPromptUsageRequest.created_at:type_name -> google.protobuf.Timestamp - 26, // 7: proto.RecordToolUsageRequest.metadata:type_name -> proto.RecordToolUsageRequest.MetadataEntry - 30, // 8: proto.RecordToolUsageRequest.created_at:type_name -> google.protobuf.Timestamp - 27, // 9: proto.RecordModelThoughtRequest.metadata:type_name -> proto.RecordModelThoughtRequest.MetadataEntry - 30, // 10: proto.RecordModelThoughtRequest.created_at:type_name -> google.protobuf.Timestamp + 25, // 0: proto.RecordInterceptionRequest.metadata:type_name -> proto.RecordInterceptionRequest.MetadataEntry + 32, // 1: proto.RecordInterceptionRequest.started_at:type_name -> google.protobuf.Timestamp + 32, // 2: proto.RecordInterceptionEndedRequest.ended_at:type_name -> google.protobuf.Timestamp + 26, // 3: proto.RecordTokenUsageRequest.metadata:type_name -> proto.RecordTokenUsageRequest.MetadataEntry + 32, // 4: proto.RecordTokenUsageRequest.created_at:type_name -> google.protobuf.Timestamp + 27, // 5: proto.RecordPromptUsageRequest.metadata:type_name -> proto.RecordPromptUsageRequest.MetadataEntry + 32, // 6: proto.RecordPromptUsageRequest.created_at:type_name -> google.protobuf.Timestamp + 28, // 7: proto.RecordToolUsageRequest.metadata:type_name -> proto.RecordToolUsageRequest.MetadataEntry + 32, // 8: proto.RecordToolUsageRequest.created_at:type_name -> google.protobuf.Timestamp + 29, // 9: proto.RecordModelThoughtRequest.metadata:type_name -> proto.RecordModelThoughtRequest.MetadataEntry + 32, // 10: proto.RecordModelThoughtRequest.created_at:type_name -> google.protobuf.Timestamp 14, // 11: proto.GetMCPServerConfigsResponse.coder_mcp_config:type_name -> proto.MCPServerConfig 14, // 12: proto.GetMCPServerConfigsResponse.external_auth_mcp_configs:type_name -> proto.MCPServerConfig - 28, // 13: proto.GetMCPServerAccessTokensBatchResponse.access_tokens:type_name -> proto.GetMCPServerAccessTokensBatchResponse.AccessTokensEntry - 29, // 14: proto.GetMCPServerAccessTokensBatchResponse.errors:type_name -> proto.GetMCPServerAccessTokensBatchResponse.ErrorsEntry - 21, // 15: proto.GetAIProvidersResponse.providers:type_name -> proto.AIProvider - 22, // 16: proto.AIProvider.bedrock:type_name -> proto.AIProviderKindBedrock - 31, // 17: proto.RecordInterceptionRequest.MetadataEntry.value:type_name -> google.protobuf.Any - 31, // 18: proto.RecordTokenUsageRequest.MetadataEntry.value:type_name -> google.protobuf.Any - 31, // 19: proto.RecordPromptUsageRequest.MetadataEntry.value:type_name -> google.protobuf.Any - 31, // 20: proto.RecordToolUsageRequest.MetadataEntry.value:type_name -> google.protobuf.Any - 31, // 21: proto.RecordModelThoughtRequest.MetadataEntry.value:type_name -> google.protobuf.Any - 0, // 22: proto.Recorder.RecordInterception:input_type -> proto.RecordInterceptionRequest - 2, // 23: proto.Recorder.RecordInterceptionEnded:input_type -> proto.RecordInterceptionEndedRequest - 4, // 24: proto.Recorder.RecordTokenUsage:input_type -> proto.RecordTokenUsageRequest - 6, // 25: proto.Recorder.RecordPromptUsage:input_type -> proto.RecordPromptUsageRequest - 8, // 26: proto.Recorder.RecordToolUsage:input_type -> proto.RecordToolUsageRequest - 10, // 27: proto.Recorder.RecordModelThought:input_type -> proto.RecordModelThoughtRequest - 12, // 28: proto.MCPConfigurator.GetMCPServerConfigs:input_type -> proto.GetMCPServerConfigsRequest - 15, // 29: proto.MCPConfigurator.GetMCPServerAccessTokensBatch:input_type -> proto.GetMCPServerAccessTokensBatchRequest - 17, // 30: proto.Authorizer.IsAuthorized:input_type -> proto.IsAuthorizedRequest - 19, // 31: proto.ProviderConfigurator.GetAIProviders:input_type -> proto.GetAIProvidersRequest - 1, // 32: proto.Recorder.RecordInterception:output_type -> proto.RecordInterceptionResponse - 3, // 33: proto.Recorder.RecordInterceptionEnded:output_type -> proto.RecordInterceptionEndedResponse - 5, // 34: proto.Recorder.RecordTokenUsage:output_type -> proto.RecordTokenUsageResponse - 7, // 35: proto.Recorder.RecordPromptUsage:output_type -> proto.RecordPromptUsageResponse - 9, // 36: proto.Recorder.RecordToolUsage:output_type -> proto.RecordToolUsageResponse - 11, // 37: proto.Recorder.RecordModelThought:output_type -> proto.RecordModelThoughtResponse - 13, // 38: proto.MCPConfigurator.GetMCPServerConfigs:output_type -> proto.GetMCPServerConfigsResponse - 16, // 39: proto.MCPConfigurator.GetMCPServerAccessTokensBatch:output_type -> proto.GetMCPServerAccessTokensBatchResponse - 18, // 40: proto.Authorizer.IsAuthorized:output_type -> proto.IsAuthorizedResponse - 20, // 41: proto.ProviderConfigurator.GetAIProviders:output_type -> proto.GetAIProvidersResponse - 32, // [32:42] is the sub-list for method output_type - 22, // [22:32] is the sub-list for method input_type - 22, // [22:22] is the sub-list for extension type_name - 22, // [22:22] is the sub-list for extension extendee - 0, // [0:22] is the sub-list for field type_name + 30, // 13: proto.GetMCPServerAccessTokensBatchResponse.access_tokens:type_name -> proto.GetMCPServerAccessTokensBatchResponse.AccessTokensEntry + 31, // 14: proto.GetMCPServerAccessTokensBatchResponse.errors:type_name -> proto.GetMCPServerAccessTokensBatchResponse.ErrorsEntry + 32, // 15: proto.IsBudgetExceededRequest.period_start:type_name -> google.protobuf.Timestamp + 23, // 16: proto.GetAIProvidersResponse.providers:type_name -> proto.AIProvider + 24, // 17: proto.AIProvider.bedrock:type_name -> proto.AIProviderKindBedrock + 33, // 18: proto.RecordInterceptionRequest.MetadataEntry.value:type_name -> google.protobuf.Any + 33, // 19: proto.RecordTokenUsageRequest.MetadataEntry.value:type_name -> google.protobuf.Any + 33, // 20: proto.RecordPromptUsageRequest.MetadataEntry.value:type_name -> google.protobuf.Any + 33, // 21: proto.RecordToolUsageRequest.MetadataEntry.value:type_name -> google.protobuf.Any + 33, // 22: proto.RecordModelThoughtRequest.MetadataEntry.value:type_name -> google.protobuf.Any + 0, // 23: proto.Recorder.RecordInterception:input_type -> proto.RecordInterceptionRequest + 2, // 24: proto.Recorder.RecordInterceptionEnded:input_type -> proto.RecordInterceptionEndedRequest + 4, // 25: proto.Recorder.RecordTokenUsage:input_type -> proto.RecordTokenUsageRequest + 6, // 26: proto.Recorder.RecordPromptUsage:input_type -> proto.RecordPromptUsageRequest + 8, // 27: proto.Recorder.RecordToolUsage:input_type -> proto.RecordToolUsageRequest + 10, // 28: proto.Recorder.RecordModelThought:input_type -> proto.RecordModelThoughtRequest + 12, // 29: proto.MCPConfigurator.GetMCPServerConfigs:input_type -> proto.GetMCPServerConfigsRequest + 15, // 30: proto.MCPConfigurator.GetMCPServerAccessTokensBatch:input_type -> proto.GetMCPServerAccessTokensBatchRequest + 17, // 31: proto.Authorizer.IsAuthorized:input_type -> proto.IsAuthorizedRequest + 19, // 32: proto.Authorizer.IsBudgetExceeded:input_type -> proto.IsBudgetExceededRequest + 21, // 33: proto.ProviderConfigurator.GetAIProviders:input_type -> proto.GetAIProvidersRequest + 1, // 34: proto.Recorder.RecordInterception:output_type -> proto.RecordInterceptionResponse + 3, // 35: proto.Recorder.RecordInterceptionEnded:output_type -> proto.RecordInterceptionEndedResponse + 5, // 36: proto.Recorder.RecordTokenUsage:output_type -> proto.RecordTokenUsageResponse + 7, // 37: proto.Recorder.RecordPromptUsage:output_type -> proto.RecordPromptUsageResponse + 9, // 38: proto.Recorder.RecordToolUsage:output_type -> proto.RecordToolUsageResponse + 11, // 39: proto.Recorder.RecordModelThought:output_type -> proto.RecordModelThoughtResponse + 13, // 40: proto.MCPConfigurator.GetMCPServerConfigs:output_type -> proto.GetMCPServerConfigsResponse + 16, // 41: proto.MCPConfigurator.GetMCPServerAccessTokensBatch:output_type -> proto.GetMCPServerAccessTokensBatchResponse + 18, // 42: proto.Authorizer.IsAuthorized:output_type -> proto.IsAuthorizedResponse + 20, // 43: proto.Authorizer.IsBudgetExceeded:output_type -> proto.IsBudgetExceededResponse + 22, // 44: proto.ProviderConfigurator.GetAIProviders:output_type -> proto.GetAIProvidersResponse + 34, // [34:45] is the sub-list for method output_type + 23, // [23:34] is the sub-list for method input_type + 23, // [23:23] is the sub-list for extension type_name + 23, // [23:23] is the sub-list for extension extendee + 0, // [0:23] is the sub-list for field type_name } func init() { file_coderd_aibridged_proto_aibridged_proto_init() } @@ -2226,7 +2369,7 @@ func file_coderd_aibridged_proto_aibridged_proto_init() { } } file_coderd_aibridged_proto_aibridged_proto_msgTypes[19].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*GetAIProvidersRequest); i { + switch v := v.(*IsBudgetExceededRequest); i { case 0: return &v.state case 1: @@ -2238,7 +2381,7 @@ func file_coderd_aibridged_proto_aibridged_proto_init() { } } file_coderd_aibridged_proto_aibridged_proto_msgTypes[20].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*GetAIProvidersResponse); i { + switch v := v.(*IsBudgetExceededResponse); i { case 0: return &v.state case 1: @@ -2250,7 +2393,7 @@ func file_coderd_aibridged_proto_aibridged_proto_init() { } } file_coderd_aibridged_proto_aibridged_proto_msgTypes[21].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*AIProvider); i { + switch v := v.(*GetAIProvidersRequest); i { case 0: return &v.state case 1: @@ -2262,6 +2405,30 @@ func file_coderd_aibridged_proto_aibridged_proto_init() { } } file_coderd_aibridged_proto_aibridged_proto_msgTypes[22].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*GetAIProvidersResponse); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_coderd_aibridged_proto_aibridged_proto_msgTypes[23].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*AIProvider); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_coderd_aibridged_proto_aibridged_proto_msgTypes[24].Exporter = func(v interface{}, i int) interface{} { switch v := v.(*AIProviderKindBedrock); i { case 0: return &v.state @@ -2276,13 +2443,14 @@ func file_coderd_aibridged_proto_aibridged_proto_init() { } file_coderd_aibridged_proto_aibridged_proto_msgTypes[0].OneofWrappers = []interface{}{} file_coderd_aibridged_proto_aibridged_proto_msgTypes[8].OneofWrappers = []interface{}{} + file_coderd_aibridged_proto_aibridged_proto_msgTypes[20].OneofWrappers = []interface{}{} type x struct{} out := protoimpl.TypeBuilder{ File: protoimpl.DescBuilder{ GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: file_coderd_aibridged_proto_aibridged_proto_rawDesc, NumEnums: 0, - NumMessages: 30, + NumMessages: 32, NumExtensions: 0, NumServices: 4, },