diff --git a/cli/aibridged.go b/cli/aibridged.go index addeb3f9598..e898c282477 100644 --- a/cli/aibridged.go +++ b/cli/aibridged.go @@ -61,7 +61,7 @@ func newAIBridgeDaemon(coderAPI *coderd.API, cfg codersdk.AIBridgeConfig, reg pr reg.MustRegister(keypool.NewStateCollector(pool.KeyPools)) // Create daemon. Construct it before subscribing so the reloader can use - // srv.Client() to fetch providers over the in-memory RPC. + // srv.ClientContext to fetch providers over the in-memory RPC. srv, err := aibridged.New(ctx, pool, func(dialCtx context.Context) (aibridged.DRPCClient, error) { return coderAPI.CreateInMemoryAIBridgeServer(dialCtx) }, logger, tracer) @@ -72,7 +72,7 @@ func newAIBridgeDaemon(coderAPI *coderd.API, cfg codersdk.AIBridgeConfig, reg pr // Subscribe to ai_providers change events so the pool tracks the database // without a restart, and perform the initial reload. The reload data path // is the in-memory RPC. - reloader := NewPoolRPCReloader(pool, srv.Client, cfg, logger.Named("provider-loader"), metrics, providerMetrics) + reloader := NewPoolRPCReloader(pool, srv.ClientContext, cfg, logger.Named("provider-loader"), metrics, providerMetrics) unsubscribe, err := aibridged.SubscribeProviderReload(ctx, coderAPI.Pubsub, reloader, logger.Named("provider-reload")) if err != nil { // Without the subscription the pool can never track provider changes, @@ -91,7 +91,7 @@ func newAIBridgeDaemon(coderAPI *coderd.API, cfg codersdk.AIBridgeConfig, reg pr // build, replace, and reload-metric accounting live in one place. type poolRPCReloader struct { pool *aibridged.CachedBridgePool - client func() (aibridged.DRPCClient, error) + client aibridged.ClientFuncWithContext cfg codersdk.AIBridgeConfig logger slog.Logger aibridgeMetrics *aibridge.Metrics @@ -100,10 +100,12 @@ type poolRPCReloader struct { // NewPoolRPCReloader builds an [aibridged.ProviderReloader] that fetches the // provider set over the DRPC client returned by client and replaces pool's -// providers, recording reload metrics against providerMetrics. +// providers, recording reload metrics against providerMetrics. client receives +// Reload's context, so a blocking acquisition unblocks when that context is +// canceled. func NewPoolRPCReloader( pool *aibridged.CachedBridgePool, - client func() (aibridged.DRPCClient, error), + client aibridged.ClientFuncWithContext, cfg codersdk.AIBridgeConfig, logger slog.Logger, aibridgeMetrics *aibridge.Metrics, @@ -121,8 +123,8 @@ func NewPoolRPCReloader( func (r *poolRPCReloader) Reload(ctx context.Context) error { r.providerMetrics.RecordReloadAttempt() - // r.client() blocks until the daemon is connected to coderd. - client, err := r.client() + // r.client blocks until the daemon connects to coderd or ctx is canceled. + client, err := r.client(ctx) if err != nil { return xerrors.Errorf("get ai-gateway client: %w", err) } diff --git a/cli/aibridged_internal_test.go b/cli/aibridged_internal_test.go index 5c2d8f6d881..e2b913dcd5b 100644 --- a/cli/aibridged_internal_test.go +++ b/cli/aibridged_internal_test.go @@ -53,7 +53,7 @@ func buildFromEnv(t *testing.T, cfg codersdk.AIBridgeConfig) ([]aibridge.Provide // (providers, outcomes) the embedded reloader would observe. func buildFromDB(ctx context.Context, t *testing.T, db database.Store, cfg codersdk.AIBridgeConfig, logger slog.Logger) ([]aibridge.Provider, []aibridged.ProviderOutcome, error) { t.Helper() - srv, err := aibridgedserver.NewServer(ctx, db, logger, "/", cfg, nil, nil, agplaiseats.Noop{}) + srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", cfg, nil, nil, agplaiseats.Noop{}) if err != nil { return nil, nil, err } diff --git a/coderd/aibridged.go b/coderd/aibridged.go index add55e2097d..0749bb2cec7 100644 --- a/coderd/aibridged.go +++ b/coderd/aibridged.go @@ -65,7 +65,7 @@ func (api *API) CreateInMemoryAIBridgeServer(dialCtx context.Context) (client ai }() mux := drpcmux.New() - srv, err := aibridgedserver.NewServer(api.ctx, api.Database, api.Logger.Named("aibridgedserver"), + srv, err := aibridgedserver.NewServer(api.ctx, api.Database, api.Pubsub, api.Logger.Named("aibridgedserver"), api.AccessURL.String(), api.DeploymentValues.AI.BridgeConfig, api.ExternalAuthConfigs, api.Experiments, api.AISeatTracker) if err != nil { return nil, err diff --git a/coderd/aibridged/aibridgedmock/clientmock.go b/coderd/aibridged/aibridgedmock/clientmock.go index 2ae8d283e2f..4170ef1b6e6 100644 --- a/coderd/aibridged/aibridgedmock/clientmock.go +++ b/coderd/aibridged/aibridgedmock/clientmock.go @@ -205,3 +205,18 @@ func (mr *MockDRPCClientMockRecorder) RecordToolUsage(ctx, in any) *gomock.Call mr.mock.ctrl.T.Helper() return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RecordToolUsage", reflect.TypeOf((*MockDRPCClient)(nil).RecordToolUsage), ctx, in) } + +// WatchAIProviders mocks base method. +func (m *MockDRPCClient) WatchAIProviders(ctx context.Context, in *proto.WatchAIProvidersRequest) (proto.DRPCProviderConfigurator_WatchAIProvidersClient, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "WatchAIProviders", ctx, in) + ret0, _ := ret[0].(proto.DRPCProviderConfigurator_WatchAIProvidersClient) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// WatchAIProviders indicates an expected call of WatchAIProviders. +func (mr *MockDRPCClientMockRecorder) WatchAIProviders(ctx, in any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WatchAIProviders", reflect.TypeOf((*MockDRPCClient)(nil).WatchAIProviders), ctx, in) +} diff --git a/coderd/aibridged/client.go b/coderd/aibridged/client.go index 2fcd3733965..3144b1a8744 100644 --- a/coderd/aibridged/client.go +++ b/coderd/aibridged/client.go @@ -12,6 +12,11 @@ type Dialer func(ctx context.Context) (DRPCClient, error) type ClientFunc func() (DRPCClient, error) +// ClientFuncWithContext acquires a DRPCClient, honoring the passed context so a +// blocking acquisition (e.g. waiting for the daemon to connect to coderd) +// unblocks when the context is canceled. Server.ClientContext satisfies it. +type ClientFuncWithContext func(context.Context) (DRPCClient, error) + // DRPCClient is the union of various service interfaces the client must support. type DRPCClient interface { proto.DRPCRecorderClient diff --git a/coderd/aibridged/proto/aibridged.pb.go b/coderd/aibridged/proto/aibridged.pb.go index b4be124d972..f5db2828bdd 100644 --- a/coderd/aibridged/proto/aibridged.pb.go +++ b/coderd/aibridged/proto/aibridged.pb.go @@ -1353,6 +1353,83 @@ func (x *GetAIProvidersResponse) GetProviders() []*AIProvider { return nil } +type WatchAIProvidersRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields +} + +func (x *WatchAIProvidersRequest) Reset() { + *x = WatchAIProvidersRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_coderd_aibridged_proto_aibridged_proto_msgTypes[21] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *WatchAIProvidersRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*WatchAIProvidersRequest) ProtoMessage() {} + +func (x *WatchAIProvidersRequest) ProtoReflect() protoreflect.Message { + 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 { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use WatchAIProvidersRequest.ProtoReflect.Descriptor instead. +func (*WatchAIProvidersRequest) Descriptor() ([]byte, []int) { + return file_coderd_aibridged_proto_aibridged_proto_rawDescGZIP(), []int{21} +} + +// WatchAIProvidersResponse is an intentionally empty change signal. +type WatchAIProvidersResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields +} + +func (x *WatchAIProvidersResponse) Reset() { + *x = WatchAIProvidersResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_coderd_aibridged_proto_aibridged_proto_msgTypes[22] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *WatchAIProvidersResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*WatchAIProvidersResponse) ProtoMessage() {} + +func (x *WatchAIProvidersResponse) ProtoReflect() protoreflect.Message { + 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 { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use WatchAIProvidersResponse.ProtoReflect.Descriptor instead. +func (*WatchAIProvidersResponse) Descriptor() ([]byte, []int) { + return file_coderd_aibridged_proto_aibridged_proto_rawDescGZIP(), []int{22} +} + type AIProvider struct { state protoimpl.MessageState sizeCache protoimpl.SizeCache @@ -1372,7 +1449,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 +1462,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 +1475,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 +1537,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 +1550,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 +1563,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 { @@ -1799,99 +1876,108 @@ var file_coderd_aibridged_proto_aibridged_proto_rawDesc = []byte{ 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, + 0x64, 0x65, 0x72, 0x73, 0x22, 0x19, 0x0a, 0x17, 0x57, 0x61, 0x74, 0x63, 0x68, 0x41, 0x49, 0x50, + 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x22, + 0x1a, 0x0a, 0x18, 0x57, 0x61, 0x74, 0x63, 0x68, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, + 0x65, 0x72, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 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, 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, 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, - 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, + 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, 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, 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, 0xbc, 0x01, + 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, 0x12, 0x55, 0x0a, 0x10, 0x57, 0x61, 0x74, 0x63, 0x68, 0x41, 0x49, + 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x12, 0x1e, 0x2e, 0x70, 0x72, 0x6f, 0x74, + 0x6f, 0x2e, 0x57, 0x61, 0x74, 0x63, 0x68, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, + 0x72, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1f, 0x2e, 0x70, 0x72, 0x6f, 0x74, + 0x6f, 0x2e, 0x57, 0x61, 0x74, 0x63, 0x68, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, + 0x72, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x30, 0x01, 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 +1992,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 @@ -1929,41 +2015,43 @@ var file_coderd_aibridged_proto_aibridged_proto_goTypes = []interface{}{ (*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 + (*WatchAIProvidersRequest)(nil), // 21: proto.WatchAIProvidersRequest + (*WatchAIProvidersResponse)(nil), // 22: proto.WatchAIProvidersResponse + (*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 + 30, // 13: proto.GetMCPServerAccessTokensBatchResponse.access_tokens:type_name -> proto.GetMCPServerAccessTokensBatchResponse.AccessTokensEntry + 31, // 14: proto.GetMCPServerAccessTokensBatchResponse.errors:type_name -> proto.GetMCPServerAccessTokensBatchResponse.ErrorsEntry + 23, // 15: proto.GetAIProvidersResponse.providers:type_name -> proto.AIProvider + 24, // 16: proto.AIProvider.bedrock:type_name -> proto.AIProviderKindBedrock + 33, // 17: proto.RecordInterceptionRequest.MetadataEntry.value:type_name -> google.protobuf.Any + 33, // 18: proto.RecordTokenUsageRequest.MetadataEntry.value:type_name -> google.protobuf.Any + 33, // 19: proto.RecordPromptUsageRequest.MetadataEntry.value:type_name -> google.protobuf.Any + 33, // 20: proto.RecordToolUsageRequest.MetadataEntry.value:type_name -> google.protobuf.Any + 33, // 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 @@ -1974,18 +2062,20 @@ var file_coderd_aibridged_proto_aibridged_proto_depIdxs = []int32{ 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 + 21, // 32: proto.ProviderConfigurator.WatchAIProviders:input_type -> proto.WatchAIProvidersRequest + 1, // 33: proto.Recorder.RecordInterception:output_type -> proto.RecordInterceptionResponse + 3, // 34: proto.Recorder.RecordInterceptionEnded:output_type -> proto.RecordInterceptionEndedResponse + 5, // 35: proto.Recorder.RecordTokenUsage:output_type -> proto.RecordTokenUsageResponse + 7, // 36: proto.Recorder.RecordPromptUsage:output_type -> proto.RecordPromptUsageResponse + 9, // 37: proto.Recorder.RecordToolUsage:output_type -> proto.RecordToolUsageResponse + 11, // 38: proto.Recorder.RecordModelThought:output_type -> proto.RecordModelThoughtResponse + 13, // 39: proto.MCPConfigurator.GetMCPServerConfigs:output_type -> proto.GetMCPServerConfigsResponse + 16, // 40: proto.MCPConfigurator.GetMCPServerAccessTokensBatch:output_type -> proto.GetMCPServerAccessTokensBatchResponse + 18, // 41: proto.Authorizer.IsAuthorized:output_type -> proto.IsAuthorizedResponse + 20, // 42: proto.ProviderConfigurator.GetAIProviders:output_type -> proto.GetAIProvidersResponse + 22, // 43: proto.ProviderConfigurator.WatchAIProviders:output_type -> proto.WatchAIProvidersResponse + 33, // [33:44] is the sub-list for method output_type + 22, // [22:33] 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 @@ -2250,7 +2340,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.(*WatchAIProvidersRequest); i { case 0: return &v.state case 1: @@ -2262,6 +2352,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.(*WatchAIProvidersResponse); 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 @@ -2282,7 +2396,7 @@ func file_coderd_aibridged_proto_aibridged_proto_init() { 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..da4f086716f 100644 --- a/coderd/aibridged/proto/aibridged.proto +++ b/coderd/aibridged/proto/aibridged.proto @@ -42,6 +42,12 @@ service ProviderConfigurator { // GetAIProviders returns the full provider set (enabled and disabled). // It synchronizes with provider seeding so the response is never raced. rpc GetAIProviders(GetAIProvidersRequest) returns (GetAIProvidersResponse); + // WatchAIProviders streams a signal whenever the provider set changes + // (env seed completion or a CRUD add/update/delete). The signal carries no + // payload; on each message the client refetches via GetAIProviders. The + // server emits one signal immediately on subscribe so a client that + // connected after the last change still converges. + rpc WatchAIProviders(WatchAIProvidersRequest) returns (stream WatchAIProvidersResponse); } message RecordInterceptionRequest { @@ -174,6 +180,11 @@ message GetAIProvidersResponse { repeated AIProvider providers = 1; } +message WatchAIProvidersRequest {} + +// WatchAIProvidersResponse is an intentionally empty change signal. +message WatchAIProvidersResponse {} + message AIProvider { string name = 1; string type = 2; diff --git a/coderd/aibridged/proto/aibridged_drpc.pb.go b/coderd/aibridged/proto/aibridged_drpc.pb.go index 5939b888bb9..ee2273bd6e9 100644 --- a/coderd/aibridged/proto/aibridged_drpc.pb.go +++ b/coderd/aibridged/proto/aibridged_drpc.pb.go @@ -504,6 +504,7 @@ type DRPCProviderConfiguratorClient interface { DRPCConn() drpc.Conn GetAIProviders(ctx context.Context, in *GetAIProvidersRequest) (*GetAIProvidersResponse, error) + WatchAIProviders(ctx context.Context, in *WatchAIProvidersRequest) (DRPCProviderConfigurator_WatchAIProvidersClient, error) } type drpcProviderConfiguratorClient struct { @@ -525,8 +526,49 @@ func (c *drpcProviderConfiguratorClient) GetAIProviders(ctx context.Context, in return out, nil } +func (c *drpcProviderConfiguratorClient) WatchAIProviders(ctx context.Context, in *WatchAIProvidersRequest) (DRPCProviderConfigurator_WatchAIProvidersClient, error) { + stream, err := c.cc.NewStream(ctx, "/proto.ProviderConfigurator/WatchAIProviders", drpcEncoding_File_coderd_aibridged_proto_aibridged_proto{}) + if err != nil { + return nil, err + } + x := &drpcProviderConfigurator_WatchAIProvidersClient{stream} + if err := x.MsgSend(in, drpcEncoding_File_coderd_aibridged_proto_aibridged_proto{}); err != nil { + return nil, err + } + if err := x.CloseSend(); err != nil { + return nil, err + } + return x, nil +} + +type DRPCProviderConfigurator_WatchAIProvidersClient interface { + drpc.Stream + Recv() (*WatchAIProvidersResponse, error) +} + +type drpcProviderConfigurator_WatchAIProvidersClient struct { + drpc.Stream +} + +func (x *drpcProviderConfigurator_WatchAIProvidersClient) GetStream() drpc.Stream { + return x.Stream +} + +func (x *drpcProviderConfigurator_WatchAIProvidersClient) Recv() (*WatchAIProvidersResponse, error) { + m := new(WatchAIProvidersResponse) + if err := x.MsgRecv(m, drpcEncoding_File_coderd_aibridged_proto_aibridged_proto{}); err != nil { + return nil, err + } + return m, nil +} + +func (x *drpcProviderConfigurator_WatchAIProvidersClient) RecvMsg(m *WatchAIProvidersResponse) error { + return x.MsgRecv(m, drpcEncoding_File_coderd_aibridged_proto_aibridged_proto{}) +} + type DRPCProviderConfiguratorServer interface { GetAIProviders(context.Context, *GetAIProvidersRequest) (*GetAIProvidersResponse, error) + WatchAIProviders(*WatchAIProvidersRequest, DRPCProviderConfigurator_WatchAIProvidersStream) error } type DRPCProviderConfiguratorUnimplementedServer struct{} @@ -535,9 +577,13 @@ func (s *DRPCProviderConfiguratorUnimplementedServer) GetAIProviders(context.Con return nil, drpcerr.WithCode(errors.New("Unimplemented"), drpcerr.Unimplemented) } +func (s *DRPCProviderConfiguratorUnimplementedServer) WatchAIProviders(*WatchAIProvidersRequest, DRPCProviderConfigurator_WatchAIProvidersStream) error { + return drpcerr.WithCode(errors.New("Unimplemented"), drpcerr.Unimplemented) +} + type DRPCProviderConfiguratorDescription struct{} -func (DRPCProviderConfiguratorDescription) NumMethods() int { return 1 } +func (DRPCProviderConfiguratorDescription) NumMethods() int { return 2 } func (DRPCProviderConfiguratorDescription) Method(n int) (string, drpc.Encoding, drpc.Receiver, interface{}, bool) { switch n { @@ -550,6 +596,15 @@ func (DRPCProviderConfiguratorDescription) Method(n int) (string, drpc.Encoding, in1.(*GetAIProvidersRequest), ) }, DRPCProviderConfiguratorServer.GetAIProviders, true + case 1: + return "/proto.ProviderConfigurator/WatchAIProviders", drpcEncoding_File_coderd_aibridged_proto_aibridged_proto{}, + func(srv interface{}, ctx context.Context, in1, in2 interface{}) (drpc.Message, error) { + return nil, srv.(DRPCProviderConfiguratorServer). + WatchAIProviders( + in1.(*WatchAIProvidersRequest), + &drpcProviderConfigurator_WatchAIProvidersStream{in2.(drpc.Stream)}, + ) + }, DRPCProviderConfiguratorServer.WatchAIProviders, true default: return "", nil, nil, nil, false } @@ -574,3 +629,16 @@ func (x *drpcProviderConfigurator_GetAIProvidersStream) SendAndClose(m *GetAIPro } return x.CloseSend() } + +type DRPCProviderConfigurator_WatchAIProvidersStream interface { + drpc.Stream + Send(*WatchAIProvidersResponse) error +} + +type drpcProviderConfigurator_WatchAIProvidersStream struct { + drpc.Stream +} + +func (x *drpcProviderConfigurator_WatchAIProvidersStream) Send(m *WatchAIProvidersResponse) error { + return x.MsgSend(m, drpcEncoding_File_coderd_aibridged_proto_aibridged_proto{}) +} diff --git a/coderd/aibridged/proto/version.go b/coderd/aibridged/proto/version.go index bb484eeb1b7..631ba3af06c 100644 --- a/coderd/aibridged/proto/version.go +++ b/coderd/aibridged/proto/version.go @@ -12,9 +12,14 @@ import "github.com/coder/coder/v2/apiversion" // - Adds the ProviderConfigurator service with the GetAIProviders unary RPC, // letting embedded and standalone gateways fetch provider configuration // over DRPC instead of reading the database directly. +// +// API v1.2: +// - Adds the ProviderConfigurator.WatchAIProviders streaming RPC, pushing a +// change signal to gateways so a running standalone gateway refetches its +// provider set when the provider configuration changes. const ( CurrentMajor = 1 - CurrentMinor = 1 + CurrentMinor = 2 ) // VersionQueryParam is the URL query parameter the standalone AI Gateway diff --git a/coderd/aibridged/reload.go b/coderd/aibridged/reload.go index 7e48fd45d26..305606b47fa 100644 --- a/coderd/aibridged/reload.go +++ b/coderd/aibridged/reload.go @@ -2,12 +2,15 @@ package aibridged import ( "context" + "time" "golang.org/x/xerrors" "cdr.dev/slog/v3" + "github.com/coder/coder/v2/coderd/aibridged/proto" dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub" "github.com/coder/coder/v2/coderd/pubsub" + "github.com/coder/retry" ) // ProviderReloader refreshes a component's provider snapshot. @@ -21,7 +24,10 @@ type ProviderReloader interface { // event is missed. // // A subscription failure returns an error without reloading. The initial -// reload is best-effort: a reload failure is logged and not returned. +// reload is best-effort: a reload failure is logged and not returned. A +// dropped-message delivery error triggers a reload too, matching +// WatchAIProviders: a drop may have masked a change, so the snapshot must +// reconverge. func SubscribeProviderReload( ctx context.Context, ps dbpubsub.Pubsub, @@ -37,8 +43,9 @@ func SubscribeProviderReload( unsubscribe, err := ps.SubscribeWithErr(pubsub.AIProvidersChangedChannel, func(cbCtx context.Context, _ []byte, err error) { if err != nil { + // A dropped message may have masked a change, so reload anyway to + // reconverge rather than skipping. logger.Warn(cbCtx, "ai providers changed event delivered with error", slog.Error(err)) - return } if err := reloader.Reload(cbCtx); err != nil { logger.Warn(cbCtx, "reload ai provider snapshot from pubsub event", slog.Error(err)) @@ -55,3 +62,76 @@ func SubscribeProviderReload( } return unsubscribe, nil } + +// WatchProviderReload opens a coderd WatchAIProviders stream via clientFn and +// calls reloader.Reload on each change signal the server emits. The stream is +// re-established with exponential backoff whenever it drops. It does not +// perform an initial load; the caller is responsible for any blocking load +// before serving. +// +// It runs until ctx is canceled, then returns ctx.Err(). clientFn receives ctx, +// so a client acquisition that blocks (e.g. Server.ClientContext waiting for +// the daemon to connect to coderd) unblocks when ctx is canceled, leaving no +// goroutine behind. +func WatchProviderReload( + ctx context.Context, + clientFn ClientFuncWithContext, + reloader ProviderReloader, + logger slog.Logger, +) error { + if clientFn == nil { + return xerrors.New("client is required") + } + if reloader == nil { + return xerrors.New("reloader is required") + } + + r := retry.New(50*time.Millisecond, 10*time.Second) + for { + received, err := watchProviderReloadOnce(ctx, clientFn, reloader, logger) + if ctx.Err() != nil { + return ctx.Err() + } + logger.Warn(ctx, "ai provider watch stream ended; reconnecting", slog.Error(err)) + // Only reset the backoff once a signal was actually received. A stream + // that opens but errors before any Recv (e.g. the server fails during + // subscribe) would otherwise reset to the floor and reconnect at + // network-RTT speed. + if received { + r.Reset() + } + if !r.Wait(ctx) { + return ctx.Err() + } + } +} + +// watchProviderReloadOnce opens a single WatchAIProviders stream and reloads on +// each signal until the stream fails. received reports whether at least one +// signal was received before the error. +func watchProviderReloadOnce(ctx context.Context, clientFn ClientFuncWithContext, reloader ProviderReloader, logger slog.Logger) (received bool, err error) { + // clientFn blocks until the daemon connects to coderd or ctx is canceled. + c, err := clientFn(ctx) + if err != nil { + return false, xerrors.Errorf("get ai-gateway client: %w", err) + } + stream, err := c.WatchAIProviders(ctx, &proto.WatchAIProvidersRequest{}) + if err != nil { + return false, xerrors.Errorf("open ai providers watch stream: %w", err) + } + defer func() { + _ = stream.Close() + }() + + for { + if _, err := stream.Recv(); err != nil { + return received, xerrors.Errorf("receive ai providers change signal: %w", err) + } + received = true + if err := reloader.Reload(ctx); err != nil { + logger.Warn(ctx, "failed to reload ai provider snapshot from watch signal", slog.Error(err)) + continue + } + logger.Debug(ctx, "reloaded ai provider snapshot from watch signal") + } +} diff --git a/coderd/aibridged/reload_test.go b/coderd/aibridged/reload_test.go index a604aaa6925..abb15a44cc4 100644 --- a/coderd/aibridged/reload_test.go +++ b/coderd/aibridged/reload_test.go @@ -2,14 +2,20 @@ package aibridged_test import ( "context" + "io" + "sync" "sync/atomic" "testing" "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" "golang.org/x/xerrors" + "storj.io/drpc" "cdr.dev/slog/v3/sloggers/slogtest" "github.com/coder/coder/v2/coderd/aibridged" + "github.com/coder/coder/v2/coderd/aibridged/aibridgedmock" + "github.com/coder/coder/v2/coderd/aibridged/proto" dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub" "github.com/coder/coder/v2/coderd/pubsub" "github.com/coder/coder/v2/testutil" @@ -77,7 +83,7 @@ func TestSubscribeProviderReloadFailsWhenSubscribeFails(t *testing.T) { require.Equal(t, 0, calls.count()) } -func TestSubscribeProviderReloadIgnoresEventError(t *testing.T) { +func TestSubscribeProviderReloadReloadsOnEventError(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitMedium) @@ -92,13 +98,215 @@ func TestSubscribeProviderReloadIgnoresEventError(t *testing.T) { require.Equal(t, 1, calls.count()) + // A dropped-message delivery error may have masked a change, so it must + // still trigger a reload to reconverge. ps.listener(ctx, nil, errPubsubDelivery) - require.Equal(t, 1, calls.count()) + require.Equal(t, 2, calls.count()) ps.listener(ctx, nil, nil) - require.Equal(t, 2, calls.count()) + require.Equal(t, 3, calls.count()) +} + +func TestWatchProviderReload(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitMedium) + logger := slogtest.Make(t, nil) + + ctrl := gomock.NewController(t) + mockClient := aibridgedmock.NewMockDRPCClient(ctrl) + + // A single stream delivers two change signals, then blocks on its context + // until the watch is canceled. + events := make(chan error, 2) + events <- nil + events <- nil + mockClient.EXPECT().WatchAIProviders(gomock.Any(), gomock.Any()).DoAndReturn( + func(rpcCtx context.Context, _ *proto.WatchAIProvidersRequest) (proto.DRPCProviderConfigurator_WatchAIProvidersClient, error) { + return &fakeWatchClientStream{ctx: rpcCtx, events: events}, nil + }).AnyTimes() + + calls := &recordingReloader{} + clientFunc := func(context.Context) (aibridged.DRPCClient, error) { return mockClient, nil } + + watchCtx, watchCancel := context.WithCancel(ctx) + done := make(chan error, 1) + go func() { done <- aibridged.WatchProviderReload(watchCtx, clientFunc, calls, logger) }() + + require.Eventually(t, func() bool { return calls.count() >= 2 }, testutil.WaitShort, testutil.IntervalFast, + "each change signal must trigger a reload") + + watchCancel() + require.ErrorIs(t, testutil.TryReceive(ctx, t, done), context.Canceled) +} + +func TestWatchProviderReloadReconnects(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitMedium) + logger := slogtest.Make(t, nil) + + ctrl := gomock.NewController(t) + mockClient := aibridgedmock.NewMockDRPCClient(ctrl) + + // The first stream delivers one signal then drops; subsequent streams + // deliver one signal then block. WatchProviderReload must reconnect after + // the drop and keep reloading. + var attempt atomic.Int32 + mockClient.EXPECT().WatchAIProviders(gomock.Any(), gomock.Any()).DoAndReturn( + func(rpcCtx context.Context, _ *proto.WatchAIProvidersRequest) (proto.DRPCProviderConfigurator_WatchAIProvidersClient, error) { + ev := make(chan error, 2) + if attempt.Add(1) == 1 { + ev <- nil + ev <- io.EOF + } else { + ev <- nil + } + return &fakeWatchClientStream{ctx: rpcCtx, events: ev}, nil + }).AnyTimes() + + calls := &recordingReloader{} + clientFunc := func(context.Context) (aibridged.DRPCClient, error) { return mockClient, nil } + + watchCtx, watchCancel := context.WithCancel(ctx) + done := make(chan error, 1) + go func() { done <- aibridged.WatchProviderReload(watchCtx, clientFunc, calls, logger) }() + + require.Eventually(t, func() bool { return calls.count() >= 2 }, testutil.WaitShort, testutil.IntervalFast, + "reload must continue after the stream drops and reconnects") + + watchCancel() + require.ErrorIs(t, testutil.TryReceive(ctx, t, done), context.Canceled) } +func TestWatchProviderReloadCancelUnblocksClient(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitMedium) + logger := slogtest.Make(t, nil) + + // clientFn blocks until its context is canceled, modeling + // Server.ClientContext waiting for the daemon to connect to coderd. Only + // watchCancel is exercised (no stream activity, no daemon close), so the + // loop can return only if clientFn honors the context it receives. + var once sync.Once + entered := make(chan struct{}) + clientFunc := func(clientCtx context.Context) (aibridged.DRPCClient, error) { + once.Do(func() { close(entered) }) + <-clientCtx.Done() + return nil, clientCtx.Err() + } + + watchCtx, watchCancel := context.WithCancel(ctx) + done := make(chan error, 1) + go func() { done <- aibridged.WatchProviderReload(watchCtx, clientFunc, &recordingReloader{}, logger) }() + + testutil.TryReceive(ctx, t, entered) + watchCancel() + require.ErrorIs(t, testutil.TryReceive(ctx, t, done), context.Canceled) +} + +func TestWatchProviderReloadRetriesDialFailure(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitMedium) + logger := slogtest.Make(t, nil) + + ctrl := gomock.NewController(t) + mockClient := aibridgedmock.NewMockDRPCClient(ctrl) + + // Once dialed, the stream delivers one signal then blocks on its context. + mockClient.EXPECT().WatchAIProviders(gomock.Any(), gomock.Any()).DoAndReturn( + func(rpcCtx context.Context, _ *proto.WatchAIProvidersRequest) (proto.DRPCProviderConfigurator_WatchAIProvidersClient, error) { + ev := make(chan error, 1) + ev <- nil + return &fakeWatchClientStream{ctx: rpcCtx, events: ev}, nil + }).AnyTimes() + + calls := &recordingReloader{} + + // The first dial fails; the second succeeds, and the loop must keep + // retrying until the dial succeeds and a reload fires. + var attempt atomic.Int32 + clientFunc := func(context.Context) (aibridged.DRPCClient, error) { + if attempt.Add(1) == 1 { + return nil, xerrors.New("dial failed") + } + return mockClient, nil + } + + watchCtx, watchCancel := context.WithCancel(ctx) + done := make(chan error, 1) + go func() { done <- aibridged.WatchProviderReload(watchCtx, clientFunc, calls, logger) }() + + require.Eventually(t, func() bool { return calls.count() >= 1 }, testutil.WaitShort, testutil.IntervalFast, + "reload must fire only after a failed dial is retried successfully") + require.GreaterOrEqual(t, int(attempt.Load()), 2, "the first dial must have failed and been retried") + + watchCancel() + require.ErrorIs(t, testutil.TryReceive(ctx, t, done), context.Canceled) +} + +func TestWatchProviderReloadContinuesAfterReloadError(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitMedium) + logger := slogtest.Make(t, nil) + + ctrl := gomock.NewController(t) + mockClient := aibridgedmock.NewMockDRPCClient(ctrl) + + events := make(chan error, 3) + for range 3 { + events <- nil + } + mockClient.EXPECT().WatchAIProviders(gomock.Any(), gomock.Any()).DoAndReturn( + func(rpcCtx context.Context, _ *proto.WatchAIProvidersRequest) (proto.DRPCProviderConfigurator_WatchAIProvidersClient, error) { + return &fakeWatchClientStream{ctx: rpcCtx, events: events}, nil + }).AnyTimes() + + // Fails its first two reloads, then succeeds. + reloader := &failNReloader{n: 2} + clientFunc := func(context.Context) (aibridged.DRPCClient, error) { return mockClient, nil } + + watchCtx, watchCancel := context.WithCancel(ctx) + done := make(chan error, 1) + go func() { done <- aibridged.WatchProviderReload(watchCtx, clientFunc, reloader, logger) }() + + require.Eventually(t, func() bool { return reloader.count() >= 3 }, testutil.WaitShort, testutil.IntervalFast, + "a failed reload must not stop the watch loop") + + watchCancel() + require.ErrorIs(t, testutil.TryReceive(ctx, t, done), context.Canceled) +} + +// fakeWatchClientStream is a minimal +// proto.DRPCProviderConfigurator_WatchAIProvidersClient. Each value popped from +// events either yields a change signal (nil) or returns the given error; when +// events is empty Recv blocks until the stream context is canceled. +type fakeWatchClientStream struct { + ctx context.Context + events chan error +} + +func (s *fakeWatchClientStream) Recv() (*proto.WatchAIProvidersResponse, error) { + select { + case err := <-s.events: + if err != nil { + return nil, err + } + return &proto.WatchAIProvidersResponse{}, nil + case <-s.ctx.Done(): + return nil, s.ctx.Err() + } +} + +func (s *fakeWatchClientStream) Context() context.Context { return s.ctx } +func (*fakeWatchClientStream) MsgSend(drpc.Message, drpc.Encoding) error { return nil } +func (*fakeWatchClientStream) MsgRecv(drpc.Message, drpc.Encoding) error { return nil } +func (*fakeWatchClientStream) CloseSend() error { return nil } +func (*fakeWatchClientStream) Close() error { return nil } + // recordingReloader is a minimal [aibridged.ProviderReloader] that // counts calls. type recordingReloader struct { @@ -118,6 +326,24 @@ func (r *recordingReloader) count() int { return int(r.n.Load()) } +// failNReloader fails its first n Reload calls, then succeeds, counting all +// calls. +type failNReloader struct { + n int32 + calls atomic.Int32 +} + +func (r *failNReloader) Reload(_ context.Context) error { + if r.calls.Add(1) <= r.n { + return errReloadFailed + } + return nil +} + +func (r *failNReloader) count() int { + return int(r.calls.Load()) +} + var ( errReloadFailed = stubError("reload failed") errPubsubDelivery = stubError("pubsub delivery failed") diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index 8b539153e98..9ecc5bdf7e5 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -24,9 +24,11 @@ import ( "github.com/coder/coder/v2/coderd/database/db2sdk" "github.com/coder/coder/v2/coderd/database/dbauthz" "github.com/coder/coder/v2/coderd/database/dbtime" + "github.com/coder/coder/v2/coderd/database/pubsub" "github.com/coder/coder/v2/coderd/externalauth" "github.com/coder/coder/v2/coderd/httpmw" codermcp "github.com/coder/coder/v2/coderd/mcp" + coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" "github.com/coder/coder/v2/coderd/util/ptr" "github.com/coder/coder/v2/codersdk" ) @@ -97,6 +99,7 @@ type Server struct { // long-running operations. lifecycleCtx context.Context store store + pubsub pubsub.Pubsub logger slog.Logger externalAuthConfigs map[string]*externalauth.Config @@ -108,7 +111,7 @@ type Server struct { budgetPolicy codersdk.AIBudgetPolicy } -func NewServer(lifecycleCtx context.Context, store store, logger slog.Logger, accessURL string, +func NewServer(lifecycleCtx context.Context, store store, ps pubsub.Pubsub, logger slog.Logger, accessURL string, bridgeCfg codersdk.AIBridgeConfig, externalAuthConfigs []*externalauth.Config, experiments codersdk.Experiments, aiSeatTracker aiseats.SeatTracker, ) (*Server, error) { @@ -125,6 +128,7 @@ func NewServer(lifecycleCtx context.Context, store store, logger slog.Logger, ac srv := &Server{ lifecycleCtx: lifecycleCtx, store: store, + pubsub: ps, logger: logger, externalAuthConfigs: eac, structuredLogging: bridgeCfg.StructuredLogging.Value(), @@ -778,6 +782,58 @@ func (s *Server) GetAIProviders(ctx context.Context, _ *proto.GetAIProvidersRequ return &proto.GetAIProvidersResponse{Providers: providers}, nil } +// WatchAIProviders streams a payload-free change signal on each +// AIProvidersChangedChannel event, plus one immediately on subscribe. Pubsub +// drop errors produce a signal rather than failing the stream. Blocks until the +// stream context or the server lifecycle is canceled. +func (s *Server) WatchAIProviders(_ *proto.WatchAIProvidersRequest, stream proto.DRPCProviderConfigurator_WatchAIProvidersStream) error { + if s.pubsub == nil { + return xerrors.New("pubsub not configured") + } + + ctx, cancel := context.WithCancel(stream.Context()) + defer cancel() + // Cancel when the server lifecycle ends, not just when the stream closes. + stop := context.AfterFunc(s.lifecycleCtx, cancel) + defer stop() + + // Buffered to one so a burst of events collapses into a single pending + // signal. + signals := make(chan struct{}, 1) + notify := func() { + select { + case signals <- struct{}{}: + default: + } + } + + // Every event signals, including dropped-message errors. + unsubscribe, err := s.pubsub.SubscribeWithErr(coderdpubsub.AIProvidersChangedChannel, func(cbCtx context.Context, _ []byte, err error) { + if err != nil { + s.logger.Warn(cbCtx, "ai providers changed event delivered with error", slog.Error(err)) + } + notify() + }) + if err != nil { + return xerrors.Errorf("subscribe to %s: %w", coderdpubsub.AIProvidersChangedChannel, err) + } + defer unsubscribe() + + // Initial signal on subscribe. + notify() + + for { + select { + case <-ctx.Done(): + return nil + case <-signals: + if err := stream.Send(&proto.WatchAIProvidersResponse{}); err != nil { + return xerrors.Errorf("send ai providers change signal: %w", err) + } + } + } +} + // Deprecated: Injected MCP in AI Bridge is deprecated and will be removed in a future release. func getCoderMCPServerConfig(experiments codersdk.Experiments, accessURL string) (*proto.MCPServerConfig, error) { // Both the MCP & OAuth2 experiments are currently required in order to use our diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go index 725f4e954af..0032162558a 100644 --- a/coderd/aibridgedserver/aibridgedserver_test.go +++ b/coderd/aibridgedserver/aibridgedserver_test.go @@ -19,10 +19,12 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + "golang.org/x/xerrors" protobufproto "google.golang.org/protobuf/proto" "google.golang.org/protobuf/types/known/anypb" "google.golang.org/protobuf/types/known/structpb" "google.golang.org/protobuf/types/known/timestamppb" + "storj.io/drpc" "cdr.dev/slog/v3" "cdr.dev/slog/v3/sloggers/slogjson" @@ -39,8 +41,10 @@ import ( "github.com/coder/coder/v2/coderd/database/dbmock" "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/coderd/database/dbtime" + "github.com/coder/coder/v2/coderd/database/pubsub" "github.com/coder/coder/v2/coderd/externalauth" codermcp "github.com/coder/coder/v2/coderd/mcp" + coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" "github.com/coder/coder/v2/coderd/rbac" "github.com/coder/coder/v2/coderd/util/ptr" "github.com/coder/coder/v2/codersdk" @@ -204,7 +208,7 @@ func TestAuthorization(t *testing.T) { tc.mocksFn(db, apiKey, user) } - srv, err := aibridgedserver.NewServer(t.Context(), db, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) + srv, err := aibridgedserver.NewServer(t.Context(), db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) require.NoError(t, err) require.NotNil(t, srv) @@ -366,7 +370,7 @@ func TestAuthorization_Delegated(t *testing.T) { tc.mocksFn(db, apiKey, user) } - srv, err := aibridgedserver.NewServer(t.Context(), db, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) + srv, err := aibridgedserver.NewServer(t.Context(), db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) require.NoError(t, err) require.NotNil(t, srv) @@ -460,7 +464,7 @@ func TestGetMCPServerConfigs(t *testing.T) { logger := testutil.Logger(t) accessURL := "https://my-cool-deployment.com" - srv, err := aibridgedserver.NewServer(t.Context(), db, logger, accessURL, codersdk.AIBridgeConfig{ + srv, err := aibridgedserver.NewServer(t.Context(), db, nil, logger, accessURL, codersdk.AIBridgeConfig{ InjectCoderMCPTools: serpent.Bool(!tc.disableCoderMCPInjection), }, tc.externalAuthConfigs, tc.experiments, agplaiseats.Noop{}) require.NoError(t, err) @@ -500,7 +504,7 @@ func TestGetMCPServerAccessTokensBatch(t *testing.T) { logger := testutil.Logger(t) // Given: 2 external auth configured with MCP and 1 without. - srv, err := aibridgedserver.NewServer(t.Context(), db, logger, "/", codersdk.AIBridgeConfig{}, []*externalauth.Config{ + srv, err := aibridgedserver.NewServer(t.Context(), db, nil, logger, "/", codersdk.AIBridgeConfig{}, []*externalauth.Config{ { ID: "1", MCPURL: "1.com/mcp", @@ -1607,7 +1611,7 @@ func TestRecordTokenUsageAuthorized(t *testing.T) { }, nil) // The server runs every store call as subjectAibridged via the authzDB. - srv, err := aibridgedserver.NewServer(ctx, authzDB, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) + srv, err := aibridgedserver.NewServer(ctx, authzDB, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) require.NoError(t, err) _, err = srv.RecordTokenUsage(ctx, &proto.RecordTokenUsageRequest{ @@ -1966,7 +1970,7 @@ func testRecordMethod[Req any, Resp any]( } ctx := testutil.Context(t, testutil.WaitLong) - srv, err := aibridgedserver.NewServer(ctx, db, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) + srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) require.NoError(t, err) resp, err := callMethod(srv, ctx, tc.request) @@ -2283,7 +2287,7 @@ func TestStructuredLogging(t *testing.T) { tc.setupMocks(db, interceptionID) ctx := testutil.Context(t, testutil.WaitLong) - srv, err := aibridgedserver.NewServer(ctx, db, logger, "/", codersdk.AIBridgeConfig{ + srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", codersdk.AIBridgeConfig{ StructuredLogging: serpent.Bool(tc.structuredLogging), }, nil, requiredExperiments, agplaiseats.Noop{}) require.NoError(t, err) @@ -2327,7 +2331,7 @@ func TestInferredThreadsByToolCalls(t *testing.T) { user := dbgen.User(t, db, database.User{}) - srv, err := aibridgedserver.NewServer(ctx, db, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) + srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) require.NoError(t, err) aID := uuid.New() @@ -2480,7 +2484,7 @@ func TestGetAIProviders(t *testing.T) { Settings: sql.NullString{String: "{not valid json", Valid: true}, }) - srv, err := aibridgedserver.NewServer(ctx, db, logger, "/", codersdk.AIBridgeConfig{}, nil, nil, agplaiseats.Noop{}) + srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, nil, agplaiseats.Noop{}) require.NoError(t, err) resp, err := srv.GetAIProviders(ctx, &proto.GetAIProvidersRequest{}) @@ -2543,7 +2547,7 @@ func TestGetAIProvidersBlocksOnSeedLock(t *testing.T) { BaseUrl: "https://api.openai.com/", }, "sk-openai") - srv, err := aibridgedserver.NewServer(ctx, db, logger, "/", codersdk.AIBridgeConfig{}, nil, nil, agplaiseats.Noop{}) + srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, nil, agplaiseats.Noop{}) require.NoError(t, err) // Simulate an in-flight env seed holding the advisory lock until released. @@ -2611,3 +2615,162 @@ func TestGetAIProvidersBlocksOnSeedLock(t *testing.T) { assert.Equal(t, "openai", resp.GetProviders()[0].GetName()) assert.Equal(t, []string{"sk-openai"}, resp.GetProviders()[0].GetKeys()) } + +// TestWatchAIProviders asserts that the WatchAIProviders handler emits an +// initial signal on subscribe, one signal per AIProvidersChangedChannel publish, +// and returns cleanly when the stream context is canceled. +func TestWatchAIProviders(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + logger := slogtest.Make(t, nil) + // In-memory pubsub delivers Publish synchronously for deterministic signals. + ps := pubsub.NewInMemory() + + srv, err := aibridgedserver.NewServer(ctx, db, ps, logger, "/", codersdk.AIBridgeConfig{}, nil, nil, agplaiseats.Noop{}) + require.NoError(t, err) + + streamCtx, streamCancel := context.WithCancel(ctx) + defer streamCancel() + stream := &fakeWatchProvidersStream{ctx: streamCtx, sent: make(chan struct{}, 16)} + + watchErr := make(chan error, 1) + go func() { + watchErr <- srv.WatchAIProviders(&proto.WatchAIProvidersRequest{}, stream) + }() + + // The handler sends an initial signal immediately on subscribe. Draining it + // before publishing guarantees the next publish is not coalesced into the + // initial signal. + testutil.TryReceive(ctx, t, stream.sent) + + require.NoError(t, ps.Publish(coderdpubsub.AIProvidersChangedChannel, nil)) + testutil.TryReceive(ctx, t, stream.sent) + + require.NoError(t, ps.Publish(coderdpubsub.AIProvidersChangedChannel, nil)) + testutil.TryReceive(ctx, t, stream.sent) + + streamCancel() + require.NoError(t, testutil.TryReceive(ctx, t, watchErr)) +} + +// TestWatchAIProvidersSignalsOnDeliveryError asserts that a dropped-message +// delivery error is forwarded as a change signal rather than failing the +// stream, so the gateway reconverges after a pubsub drop. +func TestWatchAIProvidersSignalsOnDeliveryError(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + logger := slogtest.Make(t, nil) + ps := &captureListenerPubsub{listenerC: make(chan pubsub.ListenerWithErr, 1)} + + srv, err := aibridgedserver.NewServer(ctx, db, ps, logger, "/", codersdk.AIBridgeConfig{}, nil, nil, agplaiseats.Noop{}) + require.NoError(t, err) + + streamCtx, streamCancel := context.WithCancel(ctx) + defer streamCancel() + stream := &fakeWatchProvidersStream{ctx: streamCtx, sent: make(chan struct{}, 16)} + + watchErr := make(chan error, 1) + go func() { + watchErr <- srv.WatchAIProviders(&proto.WatchAIProvidersRequest{}, stream) + }() + + // Capture the registered listener and drain the initial subscribe signal so + // the delivery-error signal that follows is not coalesced into it. + listener := testutil.TryReceive(ctx, t, ps.listenerC) + testutil.TryReceive(ctx, t, stream.sent) + + // A delivery error must still produce a signal, exercising the pubsub-error + // branch of the handler. + listener(ctx, nil, pubsub.ErrDroppedMessages) + testutil.TryReceive(ctx, t, stream.sent) + + streamCancel() + require.NoError(t, testutil.TryReceive(ctx, t, watchErr)) +} + +// TestWatchAIProvidersStopsOnLifecycleCancel asserts the handler returns when +// the server lifecycle context is canceled even though the stream context +// remains open, so a stream that outlives the server does not leak a goroutine +// on shutdown. +func TestWatchAIProvidersStopsOnLifecycleCancel(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + logger := slogtest.Make(t, nil) + ps := pubsub.NewInMemory() + + // The lifecycle context is independent of the stream context so it can be + // canceled while the stream stays open. + lifecycleCtx, lifecycleCancel := context.WithCancel(ctx) + defer lifecycleCancel() + srv, err := aibridgedserver.NewServer(lifecycleCtx, db, ps, logger, "/", codersdk.AIBridgeConfig{}, nil, nil, agplaiseats.Noop{}) + require.NoError(t, err) + + streamCtx, streamCancel := context.WithCancel(ctx) + defer streamCancel() + stream := &fakeWatchProvidersStream{ctx: streamCtx, sent: make(chan struct{}, 16)} + + watchErr := make(chan error, 1) + go func() { + watchErr <- srv.WatchAIProviders(&proto.WatchAIProvidersRequest{}, stream) + }() + + // Drain the initial subscribe signal to confirm the handler is running + // before the lifecycle is canceled. + testutil.TryReceive(ctx, t, stream.sent) + + // Canceling only the lifecycle context must stop the handler even though + // the stream context is still open. + lifecycleCancel() + require.NoError(t, testutil.TryReceive(ctx, t, watchErr)) +} + +var _ pubsub.Pubsub = (*captureListenerPubsub)(nil) + +// captureListenerPubsub captures the ListenerWithErr registered via +// SubscribeWithErr so a test can drive delivery (including errors) directly. +type captureListenerPubsub struct { + listenerC chan pubsub.ListenerWithErr +} + +func (*captureListenerPubsub) Subscribe(string, pubsub.Listener) (func(), error) { + return nil, xerrors.New("Subscribe not implemented") +} + +func (p *captureListenerPubsub) SubscribeWithErr(_ string, listener pubsub.ListenerWithErr) (func(), error) { + p.listenerC <- listener + return func() {}, nil +} + +func (*captureListenerPubsub) Publish(string, []byte) error { + return xerrors.New("Publish not implemented") +} + +func (*captureListenerPubsub) Close() error { return nil } + +// fakeWatchProvidersStream is a minimal proto.DRPCProviderConfigurator_WatchAIProvidersStream +// that records Send calls on a channel. +type fakeWatchProvidersStream struct { + ctx context.Context + sent chan struct{} +} + +func (s *fakeWatchProvidersStream) Send(*proto.WatchAIProvidersResponse) error { + select { + case s.sent <- struct{}{}: + return nil + case <-s.ctx.Done(): + return s.ctx.Err() + } +} + +func (s *fakeWatchProvidersStream) Context() context.Context { return s.ctx } +func (*fakeWatchProvidersStream) MsgSend(drpc.Message, drpc.Encoding) error { return nil } +func (*fakeWatchProvidersStream) MsgRecv(drpc.Message, drpc.Encoding) error { return nil } +func (*fakeWatchProvidersStream) CloseSend() error { return nil } +func (*fakeWatchProvidersStream) Close() error { return nil } diff --git a/coderd/aibridgedtest/aibridgedtest.go b/coderd/aibridgedtest/aibridgedtest.go index 615212cc67a..7c577c982ed 100644 --- a/coderd/aibridgedtest/aibridgedtest.go +++ b/coderd/aibridgedtest/aibridgedtest.go @@ -62,7 +62,7 @@ func StartTestAIBridgeDaemon( // The reloader fetches providers from coderd over srv's DRPC client; the // subscription drives an initial load and refreshes on change events. - reloader := cli.NewPoolRPCReloader(pool, srv.Client, cfg, logger.Named("reloader"), nil, metrics) + reloader := cli.NewPoolRPCReloader(pool, srv.ClientContext, cfg, logger.Named("reloader"), nil, metrics) unsubscribe, err := aibridged.SubscribeProviderReload(ctx, api.Pubsub, reloader, logger.Named("subscriber")) if err != nil { t.Fatalf("subscribe provider reload: %v", err) diff --git a/enterprise/cli/aigatewaystart.go b/enterprise/cli/aigatewaystart.go index 14b984ba0bd..03827c2185e 100644 --- a/enterprise/cli/aigatewaystart.go +++ b/enterprise/cli/aigatewaystart.go @@ -7,6 +7,7 @@ import ( "errors" "net" "net/http" + "sync" "time" "github.com/prometheus/client_golang/prometheus" @@ -101,14 +102,12 @@ func (r *RootCmd) aiGatewayStart() *serpent.Command { defer srv.Close() // Fetch the initial provider set from coderd, retrying until - // success. - // TODO(AIGOV-465): the standalone gateway has no refresh trigger - // yet, so this runs once on startup. - clientFn := func() (aibridged.DRPCClient, error) { - return srv.ClientContext(signalCtx) - } + // success. Subsequent changes are delivered by the watch loop + // started below. The reloader's client acquisition honors the + // context of each Reload call, so loadProviders is bounded by + // signalCtx and the watch loop by watchCtx. providerLogger := logger.Named("aibridge.providers") - reloader := agpl.NewPoolRPCReloader(pool, clientFn, vals.AI.BridgeConfig, providerLogger, metrics, providerMetrics) + reloader := agpl.NewPoolRPCReloader(pool, srv.ClientContext, vals.AI.BridgeConfig, providerLogger, metrics, providerMetrics) if err := loadProviders(signalCtx, reloader, providerLogger, srv.Done()); err != nil { if signalCtx.Err() != nil { logger.Info(signalCtx, "shutting down standalone AI Gateway") @@ -119,6 +118,23 @@ func (r *RootCmd) aiGatewayStart() *serpent.Command { mw := coderd.AIGatewayDataPlaneMiddleware(vals.AI.BridgeConfig) + // Watch coderd for provider changes and refresh the pool on each + // signal. + watchCtx, watchCancel := context.WithCancel(signalCtx) + var watchWG sync.WaitGroup + watchWG.Go(func() { + // srv.ClientContext observes watchCtx, so watchCancel below + // unblocks a pending client acquisition and drains this + // goroutine without relying on srv.Close. + if err := aibridged.WatchProviderReload(watchCtx, srv.ClientContext, reloader, providerLogger); err != nil && watchCtx.Err() == nil { + providerLogger.Warn(watchCtx, "ai provider watch loop exited", slog.Error(err)) + } + }) + defer func() { + watchCancel() + watchWG.Wait() + }() + // The standalone listener is dedicated to Gateway traffic, so // the daemon is served at the root. The /api/v2/ai-gateway // and /api/v2/aibridge/ aliases are added for compatibility @@ -252,9 +268,8 @@ func (r *RootCmd) aiGatewayStart() *serpent.Command { // reload is retried with backoff. A successful empty provider list is a valid // result and ends the loop. // -// TODO(AIGOV-465): the standalone gateway has no provider-change refresh -// trigger yet, so this runs once on startup; provider add/enable will not -// propagate to a running standalone gateway. +// Subsequent provider changes are delivered by WatchProviderReload, started +// after this initial load returns. func loadProviders(ctx context.Context, reloader aibridged.ProviderReloader, logger slog.Logger, aibridgedDone <-chan struct{}) error { for r := retry.New(50*time.Millisecond, 10*time.Second); r.Wait(ctx); { if err := reloader.Reload(ctx); err != nil { diff --git a/enterprise/coderd/aibridgeserve.go b/enterprise/coderd/aibridgeserve.go index 7cd8586a0d7..08b86f4cd62 100644 --- a/enterprise/coderd/aibridgeserve.go +++ b/enterprise/coderd/aibridgeserve.go @@ -137,6 +137,7 @@ func (api *API) aiGatewayServe(rw http.ResponseWriter, r *http.Request) { srv, err := aibridgedserver.NewServer( connCtx, api.Database, + api.AGPL.Pubsub, logger, api.AccessURL.String(), api.DeploymentValues.AI.BridgeConfig,