diff --git a/coderd/aibridged/aibridged_test.go b/coderd/aibridged/aibridged_test.go index e1fd0560753..4ef3d603e60 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" @@ -128,12 +129,38 @@ 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().IsBudgetExceeded(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsBudgetExceededResponse{ + Exceeded: true, + SpendLimitMicros: ptr.Ref(int64(1_000)), + }, nil) + }, + 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().IsBudgetExceeded(gomock.Any(), gomock.Any()).AnyTimes().Return(nil, xerrors.New("oops")) + }, + expectedErr: aibridged.ErrBudgetCheck, + 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().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")) }, @@ -223,6 +250,7 @@ func TestServeHTTP_DelegatedAPIKey(t *testing.T) { Username: "u", }, 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, @@ -255,6 +283,7 @@ func TestServeHTTP_DelegatedAPIKey(t *testing.T) { ApiKeyId: testKeyID, Username: "u", }, 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, @@ -345,6 +374,7 @@ func TestServeHTTP_DelegatedAPIKey_BYOK_Integration(t *testing.T) { Username: "u", }, 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) @@ -397,6 +427,7 @@ func TestServeHTTP_DelegatedAPIKey_Integration(t *testing.T) { Username: "u", }, 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) @@ -483,6 +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().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) @@ -677,6 +709,7 @@ func TestServeHTTP_ActorHeaders(t *testing.T) { OwnerId: testUserID.String(), Username: testUsername, }, 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() @@ -775,6 +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().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. @@ -850,6 +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().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 2ae8d283e2f..2d08d54bec2 100644 --- a/coderd/aibridged/aibridgedmock/clientmock.go +++ b/coderd/aibridged/aibridgedmock/clientmock.go @@ -116,6 +116,21 @@ func (mr *MockDRPCClientMockRecorder) IsAuthorized(ctx, in any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsAuthorized", reflect.TypeOf((*MockDRPCClient)(nil).IsAuthorized), ctx, in) } +// 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, "IsBudgetExceeded", ctx, in) + ret0, _ := ret[0].(*proto.IsBudgetExceededResponse) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// 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, "IsBudgetExceeded", reflect.TypeOf((*MockDRPCClient)(nil).IsBudgetExceeded), ctx, in) +} + // RecordInterception mocks base method. func (m *MockDRPCClient) RecordInterception(ctx context.Context, in *proto.RecordInterceptionRequest) (*proto.RecordInterceptionResponse, error) { m.ctrl.T.Helper() diff --git a/coderd/aibridged/http.go b/coderd/aibridged/http.go index 640716f37a9..c17c26587de 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,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") ) // ServeHTTP is the entrypoint for requests which will be intercepted by AI Bridge. @@ -135,6 +139,32 @@ 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 + } + logger = logger.With(slog.F("user_id", id)) + + periodStart := dbtime.StartOfMonth(dbtime.Now().UTC()) + budgetResp, err := client.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{ + UserId: id.String(), + PeriodStart: timestamppb.New(periodStart), + }) + if err != nil { + logger.Warn(ctx, "user AI budget check failed", slog.Error(err)) + http.Error(rw, ErrBudgetCheck.Error(), http.StatusInternalServerError) + return + } + if budgetResp.GetExceeded() { + http.Error(rw, fmt.Sprintf( + "AI budget of US$%.2f exceeded. Please contact an administrator for more details.", + float64(budgetResp.GetSpendLimitMicros())/1_000_000, + ), http.StatusForbidden) + return + } + // Rewire request context to include actor. // // [NOTE] @@ -144,13 +174,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.pb.go b/coderd/aibridged/proto/aibridged.pb.go index b4be124d972..f45ef36084c 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, }, diff --git a/coderd/aibridged/proto/aibridged.proto b/coderd/aibridged/proto/aibridged.proto index 8dc78e7a894..d0fd26803fd 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); + // 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 @@ -168,6 +171,23 @@ message IsAuthorizedResponse { string username = 3; } +message IsBudgetExceededRequest { + string user_id = 1; // UUID + // The spend aggregation window is [period_start, now]. + google.protobuf.Timestamp period_start = 2; +} + +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; + // 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 {} message GetAIProvidersResponse { diff --git a/coderd/aibridged/proto/aibridged_drpc.pb.go b/coderd/aibridged/proto/aibridged_drpc.pb.go index 5939b888bb9..723e670dad3 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) + IsBudgetExceeded(ctx context.Context, in *IsBudgetExceededRequest) (*IsBudgetExceededResponse, error) } type drpcAuthorizerClient struct { @@ -450,8 +451,18 @@ func (c *drpcAuthorizerClient) IsAuthorized(ctx context.Context, in *IsAuthorize return out, nil } +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 + } + return out, nil +} + type DRPCAuthorizerServer interface { IsAuthorized(context.Context, *IsAuthorizedRequest) (*IsAuthorizedResponse, error) + IsBudgetExceeded(context.Context, *IsBudgetExceededRequest) (*IsBudgetExceededResponse, 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) IsBudgetExceeded(context.Context, *IsBudgetExceededRequest) (*IsBudgetExceededResponse, 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/IsBudgetExceeded", drpcEncoding_File_coderd_aibridged_proto_aibridged_proto{}, + func(srv interface{}, ctx context.Context, in1, in2 interface{}) (drpc.Message, error) { + return srv.(DRPCAuthorizerServer). + IsBudgetExceeded( + ctx, + in1.(*IsBudgetExceededRequest), + ) + }, DRPCAuthorizerServer.IsBudgetExceeded, true default: return "", nil, nil, nil, false } @@ -500,6 +524,22 @@ func (x *drpcAuthorizer_IsAuthorizedStream) SendAndClose(m *IsAuthorizedResponse return x.CloseSend() } +type DRPCAuthorizer_IsBudgetExceededStream interface { + drpc.Stream + SendAndClose(*IsBudgetExceededResponse) error +} + +type drpcAuthorizer_IsBudgetExceededStream struct { + drpc.Stream +} + +func (x *drpcAuthorizer_IsBudgetExceededStream) SendAndClose(m *IsBudgetExceededResponse) 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 88fdc8e0900..4a63b14412f 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,92 @@ func (s *Server) IsAuthorized(ctx context.Context, in *proto.IsAuthorizedRequest }, nil } +// 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) + + userID, err := uuid.Parse(in.GetUserId()) + if err != nil { + return nil, xerrors.Errorf("invalid user_id %q: %w", in.GetUserId(), err) + } + // 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") + } + periodStart := in.GetPeriodStart().AsTime() + + userBudget, err := s.checkUserAIBudget(ctx, userID, periodStart) + if err != nil { + return nil, err + } + return &proto.IsBudgetExceededResponse{ + Exceeded: userBudget.Exceeded, + SpendLimitMicros: userBudget.SpendLimitMicros, + }, nil +} + +// 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 +} + +// 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, +// 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. +// - 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) 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 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 userAIBudget{}, nil + } + + spend, err := s.store.GetUserAISpendSince(ctx, database.GetUserAISpendSinceParams{ + UserID: userID, + EffectiveGroupID: effectiveBudget.GroupID, + PeriodStart: periodStart, + }) + if err != nil { + return userAIBudget{}, xerrors.Errorf("get user AI spend for user %q in group %q: %w", userID, effectiveBudget.GroupID, err) + } + + exceeded := spend.SpendMicros >= effectiveBudget.SpendLimitMicros + + logger := s.logger.With( + slog.F("user_id", userID), + 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 exceeded { + logger.Warn(ctx, "user AI budget exceeded") + } + + return userAIBudget{ + Exceeded: exceeded, + SpendLimitMicros: ptr.Ref(effectiveBudget.SpendLimitMicros), + }, 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 b886f569615..9d4f5b7fb98 100644 --- a/coderd/aibridgedserver/aibridgedserver_test.go +++ b/coderd/aibridgedserver/aibridgedserver_test.go @@ -391,6 +391,314 @@ func TestAuthorization_Delegated(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.IsBudgetExceededResponse + wantErrContains string + }{ + { + // Invalid UUID short-circuits before any store call. + 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, + wantErrContains: "period_start is required", + }, + { + // 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.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.IsBudgetExceededResponse{ + Exceeded: false, + SpendLimitMicros: nil, + } + }, + }, + { + // 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() + 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.IsBudgetExceededResponse{ + Exceeded: false, + SpendLimitMicros: ptr.Ref(int64(1_000)), + } + }, + }, + { + // 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() + 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.IsBudgetExceededResponse{ + Exceeded: true, + SpendLimitMicros: ptr.Ref(int64(1_000)), + } + }, + }, + { + // 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() + 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.IsBudgetExceededResponse{ + Exceeded: true, + SpendLimitMicros: ptr.Ref(int64(0)), + } + }, + }, + { + // 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() + 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.IsBudgetExceededResponse{ + Exceeded: true, + SpendLimitMicros: ptr.Ref(int64(1_000)), + } + }, + }, + { + // 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() + 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.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.IsBudgetExceededResponse { + db.EXPECT().GetUserAIBudgetOverride(gomock.Any(), userID). + Return(database.UserAIBudgetOverride{}, sql.ErrConnDone) + return nil + }, + wantErrContains: "resolve effective AI budget", + }, + { + // Error from spend aggregation propagates (fail-closed). + name: "spend aggregation error propagates", + 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{GroupID: uuid.New(), SpendLimitMicros: 1_000}, nil) + db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()). + Return(database.GetUserAISpendSinceRow{}, sql.ErrConnDone) + return nil + }, + wantErrContains: "get user AI spend", + }, + } + + 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.IsBudgetExceededResponse + 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.IsBudgetExceededRequest{UserId: userIDStr} + if !tc.omitPeriodStart { + req.PeriodStart = timestamppb.New(dbtime.StartOfMonth(dbtime.Now().UTC())) + } + resp, err := srv.IsBudgetExceeded(t.Context(), req) + if tc.wantErrContains != "" { + require.Error(t, err) + require.Nil(t, resp) + assert.ErrorContains(t, err, tc.wantErrContains) + return + } + require.NoError(t, err) + require.NotNil(t, resp) + require.Equal(t, wantResp.GetExceeded(), resp.GetExceeded(), "exceeded") + require.Equal(t, wantResp.SpendLimitMicros, resp.SpendLimitMicros, "spend_limit_micros") + }) + } +} + +// TestIsBudgetExceeded_Enforcement exercises real-DB scenarios that drive +// enforcement decisions. +func TestIsBudgetExceeded_Enforcement(t *testing.T) { + t.Parallel() + + const groupLimitMicros = 1_000_000 + + // 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() + + 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}) + + _, 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) + + 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) + + // Query with period_start 2026-01-01: includes the 2026-01-15 spend, user exceeded. + 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(groupLimitMicros), prevMonthResp.GetSpendLimitMicros()) + + // Query with period_start 2026-02-01: excludes the 2026-01-15 spend, user not exceeded. + 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(groupLimitMicros), newMonthResp.GetSpendLimitMicros()) + }) + + t.Run("new user override unblocks user", func(t *testing.T) { + t.Parallel() + ctx, rawDB, srv, user, group := setup(t) + + // Use fixed dates to keep the test deterministic. + periodStart := time.Date(2026, time.March, 1, 0, 0, 0, 0, time.UTC) + + // User spend on 2026-03-15. + _, err := rawDB.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, + EffectiveGroupID: group.ID, + Day: periodStart.AddDate(0, 0, 14), + CostMicros: 1_500_000, + }) + require.NoError(t, err) + + // User's spend exceeds the group limit. + 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()) + + // 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.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()) + }) +} + func TestGetMCPServerConfigs(t *testing.T) { t.Parallel() diff --git a/coderd/database/dbtime/dbtime.go b/coderd/database/dbtime/dbtime.go index bda5a2263ce..700a79abe79 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()) +}