From 049761fc9e498d35a45ed41a699b350fb4379f04 Mon Sep 17 00:00:00 2001 From: Danny Kopping Date: Fri, 26 Jun 2026 15:43:28 +0200 Subject: [PATCH 1/6] feat(coderd/aibridged): push provider changes to gateways via WatchAIProviders Add a WatchAIProviders streaming RPC to the ProviderConfigurator service so a running standalone AI Gateway refetches its provider set when the provider configuration changes. The server forwards AIProvidersChangedChannel events (published by the provider CRUD endpoints) as payload-free signals; the gateway refetches via GetAIProviders on each signal. Bumps the aibridged API to v1.2. --- cli/aibridged_internal_test.go | 2 +- coderd/aibridged.go | 2 +- coderd/aibridged/aibridgedmock/clientmock.go | 15 + coderd/aibridged/proto/aibridged.pb.go | 403 +++++++++++------- coderd/aibridged/proto/aibridged.proto | 12 + coderd/aibridged/proto/aibridged_drpc.pb.go | 70 ++- coderd/aibridged/proto/version.go | 7 +- coderd/aibridged/reload.go | 68 +++ coderd/aibridged/reload_test.go | 106 +++++ coderd/aibridgedserver/aibridgedserver.go | 62 ++- .../aibridgedserver/aibridgedserver_test.go | 86 +++- enterprise/cli/aigatewaystart.go | 18 +- enterprise/coderd/aibridgeserve.go | 1 + 13 files changed, 687 insertions(+), 165 deletions(-) diff --git a/cli/aibridged_internal_test.go b/cli/aibridged_internal_test.go index 5c2d8f6d88155..e2b913dcd5bfb 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 add55e2097d3a..0749bb2cec7c5 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 2ae8d283e2f52..4170ef1b6e607 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/proto/aibridged.pb.go b/coderd/aibridged/proto/aibridged.pb.go index b4be124d9729b..0fc3791288068 100644 --- a/coderd/aibridged/proto/aibridged.pb.go +++ b/coderd/aibridged/proto/aibridged.pb.go @@ -1353,6 +1353,84 @@ 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. The client +// refetches the authoritative provider set via GetAIProviders on receipt. +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 +1450,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 +1463,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 +1476,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 +1538,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 +1551,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 +1564,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 +1877,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 +1993,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 +2016,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 +2063,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 +2341,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 +2353,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 +2397,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 8dc78e7a894eb..0829c6b457f2f 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,12 @@ message GetAIProvidersResponse { repeated AIProvider providers = 1; } +message WatchAIProvidersRequest {} + +// WatchAIProvidersResponse is an intentionally empty change signal. The client +// refetches the authoritative provider set via GetAIProviders on receipt. +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 5939b888bb9db..ee2273bd6e95a 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 bb484eeb1b73c..631ba3af06ca3 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 7e48fd45d26f9..b703f7ee9d788 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. @@ -55,3 +58,68 @@ func SubscribeProviderReload( } return unsubscribe, nil } + +// WatchProviderReload opens a coderd WatchAIProviders stream via client and +// calls reloader.Reload on each change signal the server emits. The stream is +// re-established with exponential backoff whenever it drops. It runs until ctx +// is canceled, then returns ctx.Err(). It does not perform an initial load; the +// caller is responsible for any blocking load before serving. +func WatchProviderReload( + ctx context.Context, + client ClientFunc, + reloader ProviderReloader, + logger slog.Logger, +) error { + if client == 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 { + connected, err := watchProviderReloadOnce(ctx, client, reloader, logger) + if ctx.Err() != nil { + return ctx.Err() + } + logger.Warn(ctx, "ai provider watch stream ended; reconnecting", slog.Error(err)) + // A stream that opened resets the backoff so the next reconnect starts + // from the floor. + if connected { + r.Reset() + } + if !r.Wait(ctx) { + return ctx.Err() + } + } +} + +// watchProviderReloadOnce opens a single WatchAIProviders stream and reloads on +// each signal until the stream fails. connected reports whether the stream +// opened before the error. +func watchProviderReloadOnce(ctx context.Context, client ClientFunc, reloader ProviderReloader, logger slog.Logger) (connected bool, err error) { + // client() blocks until the daemon is connected to coderd. + c, err := client() + 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 true, xerrors.Errorf("receive ai providers change signal: %w", err) + } + if err := reloader.Reload(ctx); err != nil { + logger.Warn(ctx, "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 a604aaa6925c8..464e923a2b877 100644 --- a/coderd/aibridged/reload_test.go +++ b/coderd/aibridged/reload_test.go @@ -2,14 +2,19 @@ package aibridged_test import ( "context" + "io" "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" @@ -99,6 +104,107 @@ func TestSubscribeProviderReloadIgnoresEventError(t *testing.T) { require.Equal(t, 2, 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() (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") + + // Canceling the watch context unblocks Recv and ends the loop with ctx.Err(). + 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() (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) }() + + // One reload from the first stream and at least one more after reconnect. + 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) +} + +// 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 { diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index 8b539153e9809..90717be4b36ca 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" + coderpubsub "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,62 @@ func (s *Server) GetAIProviders(ctx context.Context, _ *proto.GetAIProvidersRequ return &proto.GetAIProvidersResponse{Providers: providers}, nil } +// WatchAIProviders streams a payload-free change signal to the AI Gateway daemon +// on each AIProvidersChangedChannel event, published by the provider CRUD +// endpoints. The client refetches the authoritative set via GetAIProviders on +// receipt. One signal is sent immediately on subscribe, and pubsub drop errors +// are delivered as a signal rather than failing the stream. The stream runs +// until its 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 ends when either the stream or the server lifecycle is canceled. + ctx, cancel := context.WithCancel(stream.Context()) + defer cancel() + go func() { + select { + case <-s.lifecycleCtx.Done(): + cancel() + case <-ctx.Done(): + } + }() + + // 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(coderpubsub.AIProvidersChangedChannel, func(_ context.Context, _ []byte, _ error) { + notify() + }) + if err != nil { + return xerrors.Errorf("subscribe to %s: %w", coderpubsub.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 725f4e954af28..1340d3198da36 100644 --- a/coderd/aibridgedserver/aibridgedserver_test.go +++ b/coderd/aibridgedserver/aibridgedserver_test.go @@ -23,6 +23,7 @@ import ( "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 +40,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" + coderpubsub "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 +207,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 +369,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 +463,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 +503,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 +1610,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 +1969,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 +2286,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 +2330,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 +2483,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 +2546,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 +2614,66 @@ 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) + + // Each publish yields a signal. + require.NoError(t, ps.Publish(coderpubsub.AIProvidersChangedChannel, nil)) + testutil.TryReceive(ctx, t, stream.sent) + + require.NoError(t, ps.Publish(coderpubsub.AIProvidersChangedChannel, nil)) + testutil.TryReceive(ctx, t, stream.sent) + + // Canceling the stream context ends the handler without error. + streamCancel() + require.NoError(t, testutil.TryReceive(ctx, t, watchErr)) +} + +// 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/enterprise/cli/aigatewaystart.go b/enterprise/cli/aigatewaystart.go index 14b984ba0bdac..58a88b1485d37 100644 --- a/enterprise/cli/aigatewaystart.go +++ b/enterprise/cli/aigatewaystart.go @@ -101,9 +101,8 @@ 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. + // success. Subsequent changes are delivered by the watch loop + // started below. clientFn := func() (aibridged.DRPCClient, error) { return srv.ClientContext(signalCtx) } @@ -119,6 +118,14 @@ func (r *RootCmd) aiGatewayStart() *serpent.Command { mw := coderd.AIGatewayDataPlaneMiddleware(vals.AI.BridgeConfig) + // Watch coderd for provider changes and refresh the pool on each + // signal, until signalCtx is canceled at shutdown. + go func() { + if err := aibridged.WatchProviderReload(signalCtx, srv.Client, reloader, providerLogger); err != nil && signalCtx.Err() == nil { + providerLogger.Warn(signalCtx, "ai provider watch loop exited", slog.Error(err)) + } + }() + // 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 +259,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 7cd8586a0d7e8..08b86f4cd6201 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, From 94a444b627af570d532f018cf5e20ddeabc7bd4a Mon Sep 17 00:00:00 2001 From: Danny Kopping Date: Fri, 26 Jun 2026 16:54:40 +0200 Subject: [PATCH 2/6] review: address coder-agents-review feedback on WatchAIProviders - Use context.AfterFunc to bind stream lifecycle to the server (CRF-2) - Log pubsub delivery errors in the watch callback (CRF-7) - Clarify the reload-failure warn log (CRF-3) - Drain the standalone gateway watch goroutine before srv.Close() (CRF-4) - Add a test that a failed reload does not stop the watch loop (CRF-1) - Rename import alias coderpubsub to coderdpubsub (CRF-6) - Drop comments that restate the code (CRF-5) --- coderd/aibridged/reload.go | 2 +- coderd/aibridged/reload_test.go | 53 ++++++++++++++++++- coderd/aibridgedserver/aibridgedserver.go | 20 ++++--- .../aibridgedserver/aibridgedserver_test.go | 8 ++- enterprise/cli/aigatewaystart.go | 16 ++++-- 5 files changed, 77 insertions(+), 22 deletions(-) diff --git a/coderd/aibridged/reload.go b/coderd/aibridged/reload.go index b703f7ee9d788..c51bdc7c34cb8 100644 --- a/coderd/aibridged/reload.go +++ b/coderd/aibridged/reload.go @@ -117,7 +117,7 @@ func watchProviderReloadOnce(ctx context.Context, client ClientFunc, reloader Pr return true, xerrors.Errorf("receive ai providers change signal: %w", err) } if err := reloader.Reload(ctx); err != nil { - logger.Warn(ctx, "reload ai provider snapshot from watch signal", slog.Error(err)) + 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 464e923a2b877..71692e7f5cd50 100644 --- a/coderd/aibridged/reload_test.go +++ b/coderd/aibridged/reload_test.go @@ -133,7 +133,6 @@ func TestWatchProviderReload(t *testing.T) { require.Eventually(t, func() bool { return calls.count() >= 2 }, testutil.WaitShort, testutil.IntervalFast, "each change signal must trigger a reload") - // Canceling the watch context unblocks Recv and ends the loop with ctx.Err(). watchCancel() require.ErrorIs(t, testutil.TryReceive(ctx, t, done), context.Canceled) } @@ -170,7 +169,6 @@ func TestWatchProviderReloadReconnects(t *testing.T) { done := make(chan error, 1) go func() { done <- aibridged.WatchProviderReload(watchCtx, clientFunc, calls, logger) }() - // One reload from the first stream and at least one more after reconnect. require.Eventually(t, func() bool { return calls.count() >= 2 }, testutil.WaitShort, testutil.IntervalFast, "reload must continue after the stream drops and reconnects") @@ -178,6 +176,39 @@ func TestWatchProviderReloadReconnects(t *testing.T) { 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, 8) + for range 4 { + 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() (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 @@ -224,6 +255,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 90717be4b36ca..6c56aac3695ce 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -28,7 +28,7 @@ import ( "github.com/coder/coder/v2/coderd/externalauth" "github.com/coder/coder/v2/coderd/httpmw" codermcp "github.com/coder/coder/v2/coderd/mcp" - coderpubsub "github.com/coder/coder/v2/coderd/pubsub" + coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" "github.com/coder/coder/v2/coderd/util/ptr" "github.com/coder/coder/v2/codersdk" ) @@ -793,16 +793,11 @@ func (s *Server) WatchAIProviders(_ *proto.WatchAIProvidersRequest, stream proto return xerrors.New("pubsub not configured") } - // ctx ends when either the stream or the server lifecycle is canceled. ctx, cancel := context.WithCancel(stream.Context()) defer cancel() - go func() { - select { - case <-s.lifecycleCtx.Done(): - cancel() - case <-ctx.Done(): - } - }() + // 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. @@ -815,11 +810,14 @@ func (s *Server) WatchAIProviders(_ *proto.WatchAIProvidersRequest, stream proto } // Every event signals, including dropped-message errors. - unsubscribe, err := s.pubsub.SubscribeWithErr(coderpubsub.AIProvidersChangedChannel, func(_ context.Context, _ []byte, _ error) { + 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", coderpubsub.AIProvidersChangedChannel, err) + return xerrors.Errorf("subscribe to %s: %w", coderdpubsub.AIProvidersChangedChannel, err) } defer unsubscribe() diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go index 1340d3198da36..d866bcf3fe46b 100644 --- a/coderd/aibridgedserver/aibridgedserver_test.go +++ b/coderd/aibridgedserver/aibridgedserver_test.go @@ -43,7 +43,7 @@ import ( "github.com/coder/coder/v2/coderd/database/pubsub" "github.com/coder/coder/v2/coderd/externalauth" codermcp "github.com/coder/coder/v2/coderd/mcp" - coderpubsub "github.com/coder/coder/v2/coderd/pubsub" + 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" @@ -2644,14 +2644,12 @@ func TestWatchAIProviders(t *testing.T) { // initial signal. testutil.TryReceive(ctx, t, stream.sent) - // Each publish yields a signal. - require.NoError(t, ps.Publish(coderpubsub.AIProvidersChangedChannel, nil)) + require.NoError(t, ps.Publish(coderdpubsub.AIProvidersChangedChannel, nil)) testutil.TryReceive(ctx, t, stream.sent) - require.NoError(t, ps.Publish(coderpubsub.AIProvidersChangedChannel, nil)) + require.NoError(t, ps.Publish(coderdpubsub.AIProvidersChangedChannel, nil)) testutil.TryReceive(ctx, t, stream.sent) - // Canceling the stream context ends the handler without error. streamCancel() require.NoError(t, testutil.TryReceive(ctx, t, watchErr)) } diff --git a/enterprise/cli/aigatewaystart.go b/enterprise/cli/aigatewaystart.go index 58a88b1485d37..3fa7da81ae120 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" @@ -119,12 +120,21 @@ func (r *RootCmd) aiGatewayStart() *serpent.Command { mw := coderd.AIGatewayDataPlaneMiddleware(vals.AI.BridgeConfig) // Watch coderd for provider changes and refresh the pool on each - // signal, until signalCtx is canceled at shutdown. + // signal. The deferred cancel+wait drains this goroutine before + // srv.Close() runs at shutdown. + watchCtx, watchCancel := context.WithCancel(signalCtx) + var watchWG sync.WaitGroup + watchWG.Add(1) go func() { - if err := aibridged.WatchProviderReload(signalCtx, srv.Client, reloader, providerLogger); err != nil && signalCtx.Err() == nil { - providerLogger.Warn(signalCtx, "ai provider watch loop exited", slog.Error(err)) + defer watchWG.Done() + if err := aibridged.WatchProviderReload(watchCtx, srv.Client, 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 From b1f9597458e9a6f494e38b3ef0791714bbd4523d Mon Sep 17 00:00:00 2001 From: Danny Kopping Date: Mon, 29 Jun 2026 09:42:05 +0200 Subject: [PATCH 3/6] review: fix shutdown deadlock in standalone gateway watch drain srv.Close cancels the daemon lifecycle context, the only context WatchProviderReload's blocking Client() call observes. Call it before watchWG.Wait() so a watch goroutine waiting to reconnect unblocks on the HTTP server error path instead of hanging. Close is idempotent (CRF-8). --- enterprise/cli/aigatewaystart.go | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/enterprise/cli/aigatewaystart.go b/enterprise/cli/aigatewaystart.go index 3fa7da81ae120..2a16d9508f4fe 100644 --- a/enterprise/cli/aigatewaystart.go +++ b/enterprise/cli/aigatewaystart.go @@ -120,8 +120,7 @@ func (r *RootCmd) aiGatewayStart() *serpent.Command { mw := coderd.AIGatewayDataPlaneMiddleware(vals.AI.BridgeConfig) // Watch coderd for provider changes and refresh the pool on each - // signal. The deferred cancel+wait drains this goroutine before - // srv.Close() runs at shutdown. + // signal. watchCtx, watchCancel := context.WithCancel(signalCtx) var watchWG sync.WaitGroup watchWG.Add(1) @@ -133,6 +132,12 @@ func (r *RootCmd) aiGatewayStart() *serpent.Command { }() defer func() { watchCancel() + // srv.Close cancels the daemon lifecycle context, which is the + // only context WatchProviderReload's blocking Client() call + // observes. Without it, a watch goroutine waiting to reconnect + // would not unblock and Wait would hang on the HTTP error path. + // Close is idempotent, so the deferred srv.Close above is safe. + _ = srv.Close() watchWG.Wait() }() From 1aa5d2ed58eebb311213ae10e143a86fde332774 Mon Sep 17 00:00:00 2001 From: Danny Kopping Date: Mon, 29 Jun 2026 13:39:07 +0200 Subject: [PATCH 4/6] review: address coder-agents-review feedback on provider reload - Reset watch backoff only after a signal is received, avoiding a tight reconnect loop when the server fails before the first Recv. - Document that WatchProviderReload cancellation needs both ctx and the ClientFunc to unblock. - Reload on dropped-message delivery errors in SubscribeProviderReload to match WatchAIProviders reconverge semantics. - Rename the ClientFunc parameter to clientFn for package consistency. - Use sync.WaitGroup.Go in the standalone gateway watch goroutine. - Trim the WatchAIProviders doc comment and proto response comment. - Add tests for the pubsub-error signal path and the dial-failure backoff path. --- coderd/aibridged/proto/aibridged.pb.go | 3 +- coderd/aibridged/proto/aibridged.proto | 3 +- coderd/aibridged/reload.go | 49 +++++++++------ coderd/aibridged/reload_test.go | 51 +++++++++++++++- coderd/aibridgedserver/aibridgedserver.go | 10 ++-- .../aibridgedserver/aibridgedserver_test.go | 59 +++++++++++++++++++ enterprise/cli/aigatewaystart.go | 6 +- 7 files changed, 146 insertions(+), 35 deletions(-) diff --git a/coderd/aibridged/proto/aibridged.pb.go b/coderd/aibridged/proto/aibridged.pb.go index 0fc3791288068..f5db2828bddc2 100644 --- a/coderd/aibridged/proto/aibridged.pb.go +++ b/coderd/aibridged/proto/aibridged.pb.go @@ -1391,8 +1391,7 @@ func (*WatchAIProvidersRequest) Descriptor() ([]byte, []int) { return file_coderd_aibridged_proto_aibridged_proto_rawDescGZIP(), []int{21} } -// WatchAIProvidersResponse is an intentionally empty change signal. The client -// refetches the authoritative provider set via GetAIProviders on receipt. +// WatchAIProvidersResponse is an intentionally empty change signal. type WatchAIProvidersResponse struct { state protoimpl.MessageState sizeCache protoimpl.SizeCache diff --git a/coderd/aibridged/proto/aibridged.proto b/coderd/aibridged/proto/aibridged.proto index 0829c6b457f2f..da4f086716f6e 100644 --- a/coderd/aibridged/proto/aibridged.proto +++ b/coderd/aibridged/proto/aibridged.proto @@ -182,8 +182,7 @@ message GetAIProvidersResponse { message WatchAIProvidersRequest {} -// WatchAIProvidersResponse is an intentionally empty change signal. The client -// refetches the authoritative provider set via GetAIProviders on receipt. +// WatchAIProvidersResponse is an intentionally empty change signal. message WatchAIProvidersResponse {} message AIProvider { diff --git a/coderd/aibridged/reload.go b/coderd/aibridged/reload.go index c51bdc7c34cb8..be1a600fd602b 100644 --- a/coderd/aibridged/reload.go +++ b/coderd/aibridged/reload.go @@ -24,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, @@ -40,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)) @@ -59,18 +63,24 @@ func SubscribeProviderReload( return unsubscribe, nil } -// WatchProviderReload opens a coderd WatchAIProviders stream via client and +// 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 runs until ctx -// is canceled, then returns ctx.Err(). It does not perform an initial load; the -// caller is responsible for any blocking load before serving. +// 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(). Cancellation requires +// both ctx and the underlying ClientFunc to unblock: during reconnection +// clientFn may block on its own lifecycle (e.g. Server.Client waits on the +// daemon lifecycle context), so canceling ctx alone does not unblock a pending +// clientFn call. func WatchProviderReload( ctx context.Context, - client ClientFunc, + clientFn ClientFunc, reloader ProviderReloader, logger slog.Logger, ) error { - if client == nil { + if clientFn == nil { return xerrors.New("client is required") } if reloader == nil { @@ -79,14 +89,16 @@ func WatchProviderReload( r := retry.New(50*time.Millisecond, 10*time.Second) for { - connected, err := watchProviderReloadOnce(ctx, client, reloader, logger) + 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)) - // A stream that opened resets the backoff so the next reconnect starts - // from the floor. - if connected { + // 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) { @@ -96,11 +108,11 @@ func WatchProviderReload( } // watchProviderReloadOnce opens a single WatchAIProviders stream and reloads on -// each signal until the stream fails. connected reports whether the stream -// opened before the error. -func watchProviderReloadOnce(ctx context.Context, client ClientFunc, reloader ProviderReloader, logger slog.Logger) (connected bool, err error) { - // client() blocks until the daemon is connected to coderd. - c, err := client() +// each signal until the stream fails. received reports whether at least one +// signal was received before the error. +func watchProviderReloadOnce(ctx context.Context, clientFn ClientFunc, reloader ProviderReloader, logger slog.Logger) (received bool, err error) { + // clientFn() blocks until the daemon is connected to coderd. + c, err := clientFn() if err != nil { return false, xerrors.Errorf("get ai-gateway client: %w", err) } @@ -114,8 +126,9 @@ func watchProviderReloadOnce(ctx context.Context, client ClientFunc, reloader Pr for { if _, err := stream.Recv(); err != nil { - return true, xerrors.Errorf("receive ai providers change signal: %w", err) + 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 diff --git a/coderd/aibridged/reload_test.go b/coderd/aibridged/reload_test.go index 71692e7f5cd50..44f3f1678ad6c 100644 --- a/coderd/aibridged/reload_test.go +++ b/coderd/aibridged/reload_test.go @@ -82,7 +82,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) @@ -97,11 +97,13 @@ 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) { @@ -176,6 +178,49 @@ func TestWatchProviderReloadReconnects(t *testing.T) { 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. This exercises the + // connected=false backoff path: a failed dial must not reset the backoff + // and must not trigger a reload, but the loop must keep retrying until the + // dial succeeds. + var attempt atomic.Int32 + clientFunc := func() (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() diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index 6c56aac3695ce..9ecc5bdf7e502 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -782,12 +782,10 @@ func (s *Server) GetAIProviders(ctx context.Context, _ *proto.GetAIProvidersRequ return &proto.GetAIProvidersResponse{Providers: providers}, nil } -// WatchAIProviders streams a payload-free change signal to the AI Gateway daemon -// on each AIProvidersChangedChannel event, published by the provider CRUD -// endpoints. The client refetches the authoritative set via GetAIProviders on -// receipt. One signal is sent immediately on subscribe, and pubsub drop errors -// are delivered as a signal rather than failing the stream. The stream runs -// until its context or the server lifecycle is canceled. +// 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") diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go index d866bcf3fe46b..52f22395a1fde 100644 --- a/coderd/aibridgedserver/aibridgedserver_test.go +++ b/coderd/aibridgedserver/aibridgedserver_test.go @@ -19,6 +19,7 @@ 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" @@ -2654,6 +2655,64 @@ func TestWatchAIProviders(t *testing.T) { 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)) +} + +// 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 { diff --git a/enterprise/cli/aigatewaystart.go b/enterprise/cli/aigatewaystart.go index 2a16d9508f4fe..ba63ebbc6efed 100644 --- a/enterprise/cli/aigatewaystart.go +++ b/enterprise/cli/aigatewaystart.go @@ -123,13 +123,11 @@ func (r *RootCmd) aiGatewayStart() *serpent.Command { // signal. watchCtx, watchCancel := context.WithCancel(signalCtx) var watchWG sync.WaitGroup - watchWG.Add(1) - go func() { - defer watchWG.Done() + watchWG.Go(func() { if err := aibridged.WatchProviderReload(watchCtx, srv.Client, reloader, providerLogger); err != nil && watchCtx.Err() == nil { providerLogger.Warn(watchCtx, "ai provider watch loop exited", slog.Error(err)) } - }() + }) defer func() { watchCancel() // srv.Close cancels the daemon lifecycle context, which is the From 8bc4680f3722fc590764f175d37500f5811cebda Mon Sep 17 00:00:00 2001 From: Danny Kopping Date: Mon, 29 Jun 2026 14:45:50 +0200 Subject: [PATCH 5/6] test(coderd/aibridgedserver): cover lifecycle-cancel and pubsub interface - Add TestWatchAIProvidersStopsOnLifecycleCancel to lock that canceling the server lifecycle context stops the handler while the stream context stays open, guarding against a shutdown goroutine leak. - Add a compile-time interface assertion for captureListenerPubsub to match the sibling test helpers. --- .../aibridgedserver/aibridgedserver_test.go | 40 +++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go index 52f22395a1fde..0032162558a8b 100644 --- a/coderd/aibridgedserver/aibridgedserver_test.go +++ b/coderd/aibridgedserver/aibridgedserver_test.go @@ -2692,6 +2692,46 @@ func TestWatchAIProvidersSignalsOnDeliveryError(t *testing.T) { 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 { From 3b848269aff20a03f8f3a7f058c1371b4007edb6 Mon Sep 17 00:00:00 2001 From: Danny Kopping Date: Tue, 7 Jul 2026 14:35:24 +0200 Subject: [PATCH 6/6] fix: make provider reload client acquisition context-aware Add ClientFuncWithContext and use it for WatchProviderReload and PoolRPCReloader, so blocking client acquisition (waiting for the daemon to connect to coderd) unblocks when the caller's context is canceled. Previously the reloader acquired its client via a constructor-time closure that ignored Reload's context. On the standalone gateway's serve-error exit path, a watch goroutine blocked inside Reload would not observe watchCancel, and watchWG.Wait deadlocked unless srv.Close ran first. The srv.Close drain workaround is now unnecessary and removed; TestWatchProviderReloadCancelUnblocksClient locks the cancellation contract. --- cli/aibridged.go | 16 ++++++---- coderd/aibridged/client.go | 5 +++ coderd/aibridged/reload.go | 17 +++++----- coderd/aibridged/reload_test.go | 46 +++++++++++++++++++++------ coderd/aibridgedtest/aibridgedtest.go | 2 +- enterprise/cli/aigatewaystart.go | 20 +++++------- 6 files changed, 67 insertions(+), 39 deletions(-) diff --git a/cli/aibridged.go b/cli/aibridged.go index addeb3f95987b..e898c28247778 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/coderd/aibridged/client.go b/coderd/aibridged/client.go index 2fcd3733965a2..3144b1a8744d1 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/reload.go b/coderd/aibridged/reload.go index be1a600fd602b..305606b47faf5 100644 --- a/coderd/aibridged/reload.go +++ b/coderd/aibridged/reload.go @@ -69,14 +69,13 @@ func SubscribeProviderReload( // perform an initial load; the caller is responsible for any blocking load // before serving. // -// It runs until ctx is canceled, then returns ctx.Err(). Cancellation requires -// both ctx and the underlying ClientFunc to unblock: during reconnection -// clientFn may block on its own lifecycle (e.g. Server.Client waits on the -// daemon lifecycle context), so canceling ctx alone does not unblock a pending -// clientFn call. +// 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 ClientFunc, + clientFn ClientFuncWithContext, reloader ProviderReloader, logger slog.Logger, ) error { @@ -110,9 +109,9 @@ func WatchProviderReload( // 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 ClientFunc, reloader ProviderReloader, logger slog.Logger) (received bool, err error) { - // clientFn() blocks until the daemon is connected to coderd. - c, err := clientFn() +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) } diff --git a/coderd/aibridged/reload_test.go b/coderd/aibridged/reload_test.go index 44f3f1678ad6c..abb15a44cc46a 100644 --- a/coderd/aibridged/reload_test.go +++ b/coderd/aibridged/reload_test.go @@ -3,6 +3,7 @@ package aibridged_test import ( "context" "io" + "sync" "sync/atomic" "testing" @@ -126,7 +127,7 @@ func TestWatchProviderReload(t *testing.T) { }).AnyTimes() calls := &recordingReloader{} - clientFunc := func() (aibridged.DRPCClient, error) { return mockClient, nil } + clientFunc := func(context.Context) (aibridged.DRPCClient, error) { return mockClient, nil } watchCtx, watchCancel := context.WithCancel(ctx) done := make(chan error, 1) @@ -165,7 +166,7 @@ func TestWatchProviderReloadReconnects(t *testing.T) { }).AnyTimes() calls := &recordingReloader{} - clientFunc := func() (aibridged.DRPCClient, error) { return mockClient, nil } + clientFunc := func(context.Context) (aibridged.DRPCClient, error) { return mockClient, nil } watchCtx, watchCancel := context.WithCancel(ctx) done := make(chan error, 1) @@ -178,6 +179,33 @@ func TestWatchProviderReloadReconnects(t *testing.T) { 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() @@ -197,12 +225,10 @@ func TestWatchProviderReloadRetriesDialFailure(t *testing.T) { calls := &recordingReloader{} - // The first dial fails; the second succeeds. This exercises the - // connected=false backoff path: a failed dial must not reset the backoff - // and must not trigger a reload, but the loop must keep retrying until the - // dial succeeds. + // 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() (aibridged.DRPCClient, error) { + clientFunc := func(context.Context) (aibridged.DRPCClient, error) { if attempt.Add(1) == 1 { return nil, xerrors.New("dial failed") } @@ -230,8 +256,8 @@ func TestWatchProviderReloadContinuesAfterReloadError(t *testing.T) { ctrl := gomock.NewController(t) mockClient := aibridgedmock.NewMockDRPCClient(ctrl) - events := make(chan error, 8) - for range 4 { + events := make(chan error, 3) + for range 3 { events <- nil } mockClient.EXPECT().WatchAIProviders(gomock.Any(), gomock.Any()).DoAndReturn( @@ -241,7 +267,7 @@ func TestWatchProviderReloadContinuesAfterReloadError(t *testing.T) { // Fails its first two reloads, then succeeds. reloader := &failNReloader{n: 2} - clientFunc := func() (aibridged.DRPCClient, error) { return mockClient, nil } + clientFunc := func(context.Context) (aibridged.DRPCClient, error) { return mockClient, nil } watchCtx, watchCancel := context.WithCancel(ctx) done := make(chan error, 1) diff --git a/coderd/aibridgedtest/aibridgedtest.go b/coderd/aibridgedtest/aibridgedtest.go index 615212cc67a4b..7c577c982ed0d 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 ba63ebbc6efed..03827c2185e41 100644 --- a/enterprise/cli/aigatewaystart.go +++ b/enterprise/cli/aigatewaystart.go @@ -103,12 +103,11 @@ func (r *RootCmd) aiGatewayStart() *serpent.Command { // Fetch the initial provider set from coderd, retrying until // success. Subsequent changes are delivered by the watch loop - // started below. - clientFn := func() (aibridged.DRPCClient, error) { - return srv.ClientContext(signalCtx) - } + // 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") @@ -124,18 +123,15 @@ func (r *RootCmd) aiGatewayStart() *serpent.Command { watchCtx, watchCancel := context.WithCancel(signalCtx) var watchWG sync.WaitGroup watchWG.Go(func() { - if err := aibridged.WatchProviderReload(watchCtx, srv.Client, reloader, providerLogger); err != nil && watchCtx.Err() == nil { + // 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() - // srv.Close cancels the daemon lifecycle context, which is the - // only context WatchProviderReload's blocking Client() call - // observes. Without it, a watch goroutine waiting to reconnect - // would not unblock and Wait would hang on the HTTP error path. - // Close is idempotent, so the deferred srv.Close above is safe. - _ = srv.Close() watchWG.Wait() }()