diff --git a/coderd/exp_chats.go b/coderd/exp_chats.go index 37b1927764..75a8cf6d95 100644 --- a/coderd/exp_chats.go +++ b/coderd/exp_chats.go @@ -33,6 +33,7 @@ import ( "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/db2sdk" "github.com/coder/coder/v2/coderd/database/dbauthz" + dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub" "github.com/coder/coder/v2/coderd/externalauth" "github.com/coder/coder/v2/coderd/externalauth/gitprovider" "github.com/coder/coder/v2/coderd/httpapi" @@ -110,6 +111,28 @@ func maybeWriteLimitErr(ctx context.Context, rw http.ResponseWriter, err error) return false } +func publishChatConfigEvent(logger slog.Logger, ps dbpubsub.Pubsub, kind pubsub.ChatConfigEventKind, entityID uuid.UUID) { + payload, err := json.Marshal(pubsub.ChatConfigEvent{ + Kind: kind, + EntityID: entityID, + }) + if err != nil { + logger.Error(context.Background(), "failed to marshal chat config event", + slog.F("kind", kind), + slog.F("entity_id", entityID), + slog.Error(err), + ) + return + } + if err := ps.Publish(pubsub.ChatConfigEventChannel, payload); err != nil { + logger.Error(context.Background(), "failed to publish chat config event", + slog.F("kind", kind), + slog.F("entity_id", entityID), + slog.Error(err), + ) + } +} + // EXPERIMENTAL: this endpoint is experimental and is subject to change. func (api *API) watchChats(rw http.ResponseWriter, r *http.Request) { ctx := r.Context() @@ -3052,6 +3075,8 @@ func (api *API) putUserChatCustomPrompt(rw http.ResponseWriter, r *http.Request) return } + publishChatConfigEvent(api.Logger, api.Pubsub, pubsub.ChatConfigEventUserPrompt, apiKey.UserID) + httpapi.Write(ctx, rw, http.StatusOK, codersdk.UserChatCustomPrompt{ CustomPrompt: updatedConfig.Value, }) @@ -3886,6 +3911,8 @@ func (api *API) createChatProvider(rw http.ResponseWriter, r *http.Request) { } } + publishChatConfigEvent(api.Logger, api.Pubsub, pubsub.ChatConfigEventProviders, uuid.Nil) + httpapi.Write( ctx, rw, @@ -3972,6 +3999,8 @@ func (api *API) updateChatProvider(rw http.ResponseWriter, r *http.Request) { return } + publishChatConfigEvent(api.Logger, api.Pubsub, pubsub.ChatConfigEventProviders, uuid.Nil) + httpapi.Write( ctx, rw, @@ -4026,6 +4055,8 @@ func (api *API) deleteChatProvider(rw http.ResponseWriter, r *http.Request) { return } + publishChatConfigEvent(api.Logger, api.Pubsub, pubsub.ChatConfigEventProviders, uuid.Nil) + rw.WriteHeader(http.StatusNoContent) } @@ -4205,6 +4236,8 @@ func (api *API) createChatModelConfig(rw http.ResponseWriter, r *http.Request) { } } + publishChatConfigEvent(api.Logger, api.Pubsub, pubsub.ChatConfigEventModelConfig, inserted.ID) + httpapi.Write(ctx, rw, http.StatusCreated, convertChatModelConfig(inserted)) } @@ -4376,6 +4409,8 @@ func (api *API) updateChatModelConfig(rw http.ResponseWriter, r *http.Request) { } } + publishChatConfigEvent(api.Logger, api.Pubsub, pubsub.ChatConfigEventModelConfig, updated.ID) + httpapi.Write(ctx, rw, http.StatusOK, convertChatModelConfig(updated)) } @@ -4416,6 +4451,8 @@ func (api *API) deleteChatModelConfig(rw http.ResponseWriter, r *http.Request) { return } + publishChatConfigEvent(api.Logger, api.Pubsub, pubsub.ChatConfigEventModelConfig, modelConfigID) + rw.WriteHeader(http.StatusNoContent) } diff --git a/coderd/pubsub/chatconfigevent.go b/coderd/pubsub/chatconfigevent.go new file mode 100644 index 0000000000..896a2aaf82 --- /dev/null +++ b/coderd/pubsub/chatconfigevent.go @@ -0,0 +1,52 @@ +package pubsub + +import ( + "context" + "encoding/json" + + "github.com/google/uuid" + "golang.org/x/xerrors" +) + +// ChatConfigEventChannel is the pubsub channel for chat config +// changes (providers, model configs, user prompts). All replicas +// subscribe to this channel to invalidate their local caches. +const ChatConfigEventChannel = "chat:config_change" + +// HandleChatConfigEvent wraps a typed callback for ChatConfigEvent +// messages, following the same pattern as HandleChatEvent. +func HandleChatConfigEvent(cb func(ctx context.Context, payload ChatConfigEvent, err error)) func(ctx context.Context, message []byte, err error) { + return func(ctx context.Context, message []byte, err error) { + if err != nil { + cb(ctx, ChatConfigEvent{}, xerrors.Errorf("chat config event pubsub: %w", err)) + return + } + var payload ChatConfigEvent + if err := json.Unmarshal(message, &payload); err != nil { + cb(ctx, ChatConfigEvent{}, xerrors.Errorf("unmarshal chat config event: %w", err)) + return + } + + cb(ctx, payload, err) + } +} + +// ChatConfigEvent is published when chat configuration changes +// (provider CRUD, model config CRUD, or user prompt updates). +// Subscribers use this to invalidate their local caches. +type ChatConfigEvent struct { + Kind ChatConfigEventKind `json:"kind"` + // EntityID carries context for the invalidation: + // - For providers: uuid.Nil (all providers are invalidated). + // - For model configs: the specific config ID. + // - For user prompts: the user ID. + EntityID uuid.UUID `json:"entity_id"` +} + +type ChatConfigEventKind string + +const ( + ChatConfigEventProviders ChatConfigEventKind = "providers" + ChatConfigEventModelConfig ChatConfigEventKind = "model_config" + ChatConfigEventUserPrompt ChatConfigEventKind = "user_prompt" +) diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 4f88792106..48b2fd5032 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -108,6 +108,8 @@ type Server struct { pubsub pubsub.Pubsub webpushDispatcher webpush.Dispatcher providerAPIKeys chatprovider.ProviderAPIKeys + configCache *chatConfigCache + configCacheUnsubscribe func() // chatStreams stores per-chat stream state. Using sync.Map // gives each chat independent locking — concurrent chats @@ -1740,6 +1742,31 @@ func New(cfg Config) *Server { //nolint:gocritic // The chat processor uses a scoped chatd context. ctx = dbauthz.AsChatd(ctx) + + p.configCache = newChatConfigCache(ctx, cfg.Database, clk) + if p.pubsub != nil { + cancelConfigSub, err := p.pubsub.SubscribeWithErr( + coderdpubsub.ChatConfigEventChannel, + coderdpubsub.HandleChatConfigEvent(func(ctx context.Context, ev coderdpubsub.ChatConfigEvent, err error) { + if err != nil { + p.logger.Warn(ctx, "chat config event error", slog.Error(err)) + return + } + switch ev.Kind { + case coderdpubsub.ChatConfigEventProviders: + p.configCache.InvalidateProviders() + case coderdpubsub.ChatConfigEventModelConfig: + p.configCache.InvalidateModelConfig(ev.EntityID) + case coderdpubsub.ChatConfigEventUserPrompt: + p.configCache.InvalidateUserPrompt(ev.EntityID) + } + }), + ) + if err != nil { + p.logger.Error(ctx, "subscribe to chat config events", slog.Error(err)) + } + p.configCacheUnsubscribe = cancelConfigSub + } go p.start(ctx) return p @@ -4017,7 +4044,7 @@ func (p *Server) resolveChatModel( }) g.Go(func() error { var err error - providers, err = p.db.GetEnabledChatProviders(ctx) + providers, err = p.configCache.EnabledProviders(ctx) if err != nil { return xerrors.Errorf("get enabled chat providers: %w", err) } @@ -4061,7 +4088,7 @@ func (p *Server) resolveModelConfig( chat database.Chat, ) (database.ChatModelConfig, error) { if chat.LastModelConfigID != uuid.Nil { - modelConfig, err := p.db.GetChatModelConfigByID( + modelConfig, err := p.configCache.ModelConfigByID( ctx, chat.LastModelConfigID, ) if err == nil { @@ -4076,7 +4103,7 @@ func (p *Server) resolveModelConfig( // Model config was deleted, fall through to default. } - defaultConfig, err := p.db.GetDefaultChatModelConfig(ctx) + defaultConfig, err := p.configCache.DefaultModelConfig(ctx) if err != nil { if xerrors.Is(err, sql.ErrNoRows) { return database.ChatModelConfig{}, xerrors.New( @@ -4286,7 +4313,7 @@ func (p *Server) resolveUserCompactionThreshold(ctx context.Context, userID uuid // database and wraps it in tags. Returns empty // string if no prompt is set. func (p *Server) resolveUserPrompt(ctx context.Context, userID uuid.UUID) string { - raw, err := p.db.GetUserChatCustomPrompt(ctx, userID) + raw, err := p.configCache.UserPrompt(ctx, userID) if err != nil { // sql.ErrNoRows is the normal "not set" case. return "" @@ -4447,6 +4474,10 @@ func (p *Server) dispatchPush( // Close stops the processor and waits for it to finish. func (p *Server) Close() error { + if unsub := p.configCacheUnsubscribe; unsub != nil { + p.configCacheUnsubscribe = nil + unsub() + } p.cancel() <-p.closed p.inflight.Wait() diff --git a/coderd/x/chatd/configcache.go b/coderd/x/chatd/configcache.go new file mode 100644 index 0000000000..4d4876e3b8 --- /dev/null +++ b/coderd/x/chatd/configcache.go @@ -0,0 +1,412 @@ +package chatd + +import ( + "context" + "database/sql" + "errors" + "fmt" + "slices" + "sync" + "time" + + "github.com/ammario/tlru" + "github.com/google/uuid" + "tailscale.com/util/singleflight" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/quartz" +) + +const ( + chatConfigProvidersTTL = 10 * time.Second + chatConfigModelConfigTTL = 10 * time.Second + chatConfigUserPromptTTL = 5 * time.Second + // Bound user-prompt cache cardinality so one-shot users do not + // accumulate forever in long-lived chatd processes. + chatConfigUserPromptEntryLimit = 64 * 1024 +) + +type cachedProviders struct { + providers []database.ChatProvider + expiresAt time.Time +} + +type cachedModelConfig struct { + config database.ChatModelConfig + expiresAt time.Time +} + +type modelConfigSnapshot struct { + epoch uint64 + generation uint64 +} + +// cloneModelConfig returns a shallow copy of cfg with Options +// deep-cloned so the cache owns its own backing array. +func cloneModelConfig(cfg database.ChatModelConfig) database.ChatModelConfig { + cfg.Options = slices.Clone(cfg.Options) + return cfg +} + +type chatConfigCache struct { + db database.Store + clock quartz.Clock + // ctx is the server-scoped context used for all DB fills. + // Cache fills run inside singleflight.Do where one caller + // becomes the leader for all coalesced waiters. Using a + // per-request context would mean the leader's cancellation + // (timeout, user disconnect) fans the error to every waiter. + // Storing the server context here makes that impossible by + // construction — callers cannot pass a request context into + // the shared fill path. + ctx context.Context + + mu sync.RWMutex + + // Providers (singleton). + providers *cachedProviders + providerGeneration uint64 + providerFetches singleflight.Group[string, []database.ChatProvider] + + // Model configs (keyed by ID). + modelTopologyEpoch uint64 + modelConfigs map[uuid.UUID]cachedModelConfig + modelConfigFetches singleflight.Group[string, database.ChatModelConfig] + + // Default model config (singleton). + defaultModelConfig *cachedModelConfig + defaultModelConfigGeneration uint64 + defaultModelConfigFetches singleflight.Group[string, database.ChatModelConfig] + + // User custom prompts (keyed by user ID). + userPromptEpoch uint64 + userPrompts *tlru.Cache[uuid.UUID, string] + userPromptFetches singleflight.Group[string, string] +} + +func newChatConfigCache(ctx context.Context, db database.Store, clock quartz.Clock) *chatConfigCache { + return &chatConfigCache{ + db: db, + clock: clock, + ctx: ctx, + modelConfigs: make(map[uuid.UUID]cachedModelConfig), + userPrompts: tlru.New[uuid.UUID]( + tlru.ConstantCost[string], + chatConfigUserPromptEntryLimit, + ), + } +} + +// singleflightDoChan wraps a singleflight group's DoChan method, +// allowing the caller to abandon the wait if their context is +// canceled while the shared fill continues running to completion. +// This separates two lifetimes: the fill runs under the server-scoped +// context, while each caller waits under its own request-scoped context. +func singleflightDoChan[K comparable, V any]( + ctx context.Context, + group *singleflight.Group[K, V], + key K, + fn func() (V, error), +) (V, error) { + ch := group.DoChan(key, fn) + select { + case <-ctx.Done(): + var zero V + return zero, ctx.Err() + case res := <-ch: + return res.Val, res.Err + } +} + +func (c *chatConfigCache) EnabledProviders(ctx context.Context) ([]database.ChatProvider, error) { + if providers, ok := c.cachedProviders(); ok { + return providers, nil + } + + generation := c.providersGeneration() + providers, err := singleflightDoChan( + ctx, + &c.providerFetches, + fmt.Sprintf("%d:providers", generation), + func() ([]database.ChatProvider, error) { + if cached, ok := c.cachedProviders(); ok { + return cached, nil + } + + fetched, err := c.db.GetEnabledChatProviders(c.ctx) + if err != nil { + return nil, err + } + c.storeProviders(generation, fetched) + return slices.Clone(fetched), nil + }, + ) + if err != nil { + return nil, err + } + + return slices.Clone(providers), nil +} + +func (c *chatConfigCache) cachedProviders() ([]database.ChatProvider, bool) { + c.mu.RLock() + entry := c.providers + c.mu.RUnlock() + if entry == nil { + return nil, false + } + if c.clock.Now().Before(entry.expiresAt) { + return slices.Clone(entry.providers), true + } + + c.mu.Lock() + if current := c.providers; current != nil && !c.clock.Now().Before(current.expiresAt) { + c.providers = nil + } + c.mu.Unlock() + + return nil, false +} + +func (c *chatConfigCache) providersGeneration() uint64 { + c.mu.RLock() + generation := c.providerGeneration + c.mu.RUnlock() + return generation +} + +func (c *chatConfigCache) storeProviders(generation uint64, providers []database.ChatProvider) { + c.mu.Lock() + defer c.mu.Unlock() + + if c.providerGeneration != generation { + return + } + + c.providers = &cachedProviders{ + providers: slices.Clone(providers), + expiresAt: c.clock.Now().Add(chatConfigProvidersTTL), + } +} + +func (c *chatConfigCache) InvalidateProviders() { + c.mu.Lock() + c.providers = nil + c.providerGeneration++ + // Provider topology changed — model selections depend on + // provider existence, so flush all model-config state. + clear(c.modelConfigs) + c.modelTopologyEpoch++ + c.defaultModelConfig = nil + c.defaultModelConfigGeneration++ + c.mu.Unlock() +} + +func (c *chatConfigCache) ModelConfigByID(ctx context.Context, id uuid.UUID) (database.ChatModelConfig, error) { + if config, ok := c.cachedModelConfig(id); ok { + return config, nil + } + + snap := c.modelConfigSnapshot() + config, err := singleflightDoChan(ctx, &c.modelConfigFetches, fmt.Sprintf("%d:%s", snap.epoch, id), func() (database.ChatModelConfig, error) { + if cached, ok := c.cachedModelConfig(id); ok { + return cached, nil + } + + fetched, err := c.db.GetChatModelConfigByID(c.ctx, id) + if err != nil { + return database.ChatModelConfig{}, err + } + c.storeModelConfig(snap, fetched) + return cloneModelConfig(fetched), nil + }) + if err != nil { + return database.ChatModelConfig{}, err + } + + return config, nil +} + +func (c *chatConfigCache) cachedModelConfig(id uuid.UUID) (database.ChatModelConfig, bool) { + c.mu.RLock() + entry, ok := c.modelConfigs[id] + c.mu.RUnlock() + if !ok { + return database.ChatModelConfig{}, false + } + if c.clock.Now().Before(entry.expiresAt) { + return cloneModelConfig(entry.config), true + } + + c.mu.Lock() + if current, ok := c.modelConfigs[id]; ok && !c.clock.Now().Before(current.expiresAt) { + delete(c.modelConfigs, id) + } + c.mu.Unlock() + + return database.ChatModelConfig{}, false +} + +func (c *chatConfigCache) modelConfigSnapshot() modelConfigSnapshot { + c.mu.RLock() + snap := modelConfigSnapshot{epoch: c.modelTopologyEpoch} + c.mu.RUnlock() + return snap +} + +func (c *chatConfigCache) storeModelConfig(snap modelConfigSnapshot, config database.ChatModelConfig) { + c.mu.Lock() + defer c.mu.Unlock() + + if c.modelTopologyEpoch != snap.epoch { + return + } + + c.modelConfigs[config.ID] = cachedModelConfig{ + config: cloneModelConfig(config), + expiresAt: c.clock.Now().Add(chatConfigModelConfigTTL), + } +} + +func (c *chatConfigCache) DefaultModelConfig(ctx context.Context) (database.ChatModelConfig, error) { + if config, ok := c.cachedDefaultModelConfig(); ok { + return config, nil + } + + snap := c.defaultModelConfigSnapshot() + config, err := singleflightDoChan(ctx, &c.defaultModelConfigFetches, fmt.Sprintf("%d:default", snap.epoch), func() (database.ChatModelConfig, error) { + if cached, ok := c.cachedDefaultModelConfig(); ok { + return cached, nil + } + + fetched, err := c.db.GetDefaultChatModelConfig(c.ctx) + if err != nil { + return database.ChatModelConfig{}, err + } + c.storeDefaultModelConfig(snap, fetched) + return cloneModelConfig(fetched), nil + }) + if err != nil { + return database.ChatModelConfig{}, err + } + + return config, nil +} + +func (c *chatConfigCache) cachedDefaultModelConfig() (database.ChatModelConfig, bool) { + c.mu.RLock() + entry := c.defaultModelConfig + c.mu.RUnlock() + if entry == nil { + return database.ChatModelConfig{}, false + } + if c.clock.Now().Before(entry.expiresAt) { + return cloneModelConfig(entry.config), true + } + + c.mu.Lock() + if current := c.defaultModelConfig; current != nil && !c.clock.Now().Before(current.expiresAt) { + c.defaultModelConfig = nil + } + c.mu.Unlock() + + return database.ChatModelConfig{}, false +} + +func (c *chatConfigCache) defaultModelConfigSnapshot() modelConfigSnapshot { + c.mu.RLock() + snap := modelConfigSnapshot{ + epoch: c.modelTopologyEpoch, + generation: c.defaultModelConfigGeneration, + } + c.mu.RUnlock() + return snap +} + +func (c *chatConfigCache) storeDefaultModelConfig(snap modelConfigSnapshot, config database.ChatModelConfig) { + c.mu.Lock() + defer c.mu.Unlock() + + if c.modelTopologyEpoch != snap.epoch { + return + } + if c.defaultModelConfigGeneration != snap.generation { + return + } + + c.defaultModelConfig = &cachedModelConfig{ + config: cloneModelConfig(config), + expiresAt: c.clock.Now().Add(chatConfigModelConfigTTL), + } +} + +func (c *chatConfigCache) UserPrompt(ctx context.Context, userID uuid.UUID) (string, error) { + if prompt, ok := c.cachedUserPrompt(userID); ok { + return prompt, nil + } + + epoch := c.currentUserPromptEpoch() + prompt, err := singleflightDoChan(ctx, &c.userPromptFetches, fmt.Sprintf("%d:%s", epoch, userID), func() (string, error) { + if cached, ok := c.cachedUserPrompt(userID); ok { + return cached, nil + } + + fetched, err := c.db.GetUserChatCustomPrompt(c.ctx, userID) + if err != nil { + if errors.Is(err, sql.ErrNoRows) { + c.storeUserPrompt(epoch, userID, "") + return "", nil + } + return "", err + } + c.storeUserPrompt(epoch, userID, fetched) + return fetched, nil + }) + if err != nil { + return "", err + } + + return prompt, nil +} + +func (c *chatConfigCache) cachedUserPrompt(userID uuid.UUID) (string, bool) { + prompt, _, ok := c.userPrompts.Get(userID) + if !ok { + return "", false + } + return prompt, true +} + +func (c *chatConfigCache) currentUserPromptEpoch() uint64 { + c.mu.RLock() + epoch := c.userPromptEpoch + c.mu.RUnlock() + return epoch +} + +func (c *chatConfigCache) storeUserPrompt(epoch uint64, userID uuid.UUID, prompt string) { + c.mu.Lock() + defer c.mu.Unlock() + + if c.userPromptEpoch != epoch { + return + } + + c.userPrompts.Set(userID, prompt, chatConfigUserPromptTTL) +} + +func (c *chatConfigCache) InvalidateModelConfig(id uuid.UUID) { + c.mu.Lock() + delete(c.modelConfigs, id) + c.modelTopologyEpoch++ + c.defaultModelConfig = nil + c.defaultModelConfigGeneration++ + c.mu.Unlock() +} + +func (c *chatConfigCache) InvalidateUserPrompt(userID uuid.UUID) { + c.mu.Lock() + c.userPrompts.Delete(userID) + c.userPromptEpoch++ + c.mu.Unlock() +} diff --git a/coderd/x/chatd/configcache_test.go b/coderd/x/chatd/configcache_test.go new file mode 100644 index 0000000000..c99a81d75e --- /dev/null +++ b/coderd/x/chatd/configcache_test.go @@ -0,0 +1,978 @@ +package chatd //nolint:testpackage // Uses internal cache state. + +import ( + "context" + "database/sql" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/testutil" + "github.com/coder/quartz" +) + +type stubChatConfigStore struct { + database.Store + + getEnabledChatProviders func(context.Context) ([]database.ChatProvider, error) + getChatModelConfigByID func(context.Context, uuid.UUID) (database.ChatModelConfig, error) + getDefaultChatModelConfig func(context.Context) (database.ChatModelConfig, error) + getUserChatCustomPrompt func(context.Context, uuid.UUID) (string, error) + + enabledProvidersCalls atomic.Int32 + modelConfigByIDCalls atomic.Int32 + defaultModelConfigCall atomic.Int32 + userPromptCalls atomic.Int32 +} + +func (s *stubChatConfigStore) GetEnabledChatProviders(ctx context.Context) ([]database.ChatProvider, error) { + s.enabledProvidersCalls.Add(1) + if s.getEnabledChatProviders == nil { + panic("unexpected GetEnabledChatProviders call") + } + return s.getEnabledChatProviders(ctx) +} + +func (s *stubChatConfigStore) GetChatModelConfigByID(ctx context.Context, id uuid.UUID) (database.ChatModelConfig, error) { + s.modelConfigByIDCalls.Add(1) + if s.getChatModelConfigByID == nil { + panic("unexpected GetChatModelConfigByID call") + } + return s.getChatModelConfigByID(ctx, id) +} + +func (s *stubChatConfigStore) GetDefaultChatModelConfig(ctx context.Context) (database.ChatModelConfig, error) { + s.defaultModelConfigCall.Add(1) + if s.getDefaultChatModelConfig == nil { + panic("unexpected GetDefaultChatModelConfig call") + } + return s.getDefaultChatModelConfig(ctx) +} + +func (s *stubChatConfigStore) GetUserChatCustomPrompt(ctx context.Context, userID uuid.UUID) (string, error) { + s.userPromptCalls.Add(1) + if s.getUserChatCustomPrompt == nil { + panic("unexpected GetUserChatCustomPrompt call") + } + return s.getUserChatCustomPrompt(ctx, userID) +} + +func TestConfigCache_EnabledProviders_CacheHit(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + providers := []database.ChatProvider{testChatProvider("provider-a")} + store := &stubChatConfigStore{ + getEnabledChatProviders: func(context.Context) ([]database.ChatProvider, error) { + return providers, nil + }, + } + cache := newChatConfigCache(ctx, store, clock) + + first, err := cache.EnabledProviders(ctx) + require.NoError(t, err) + second, err := cache.EnabledProviders(ctx) + require.NoError(t, err) + + require.Equal(t, providers, first) + require.Equal(t, providers, second) + require.Equal(t, int32(1), store.enabledProvidersCalls.Load()) +} + +func TestConfigCache_EnabledProviders_TTLExpiry(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + store := &stubChatConfigStore{} + store.getEnabledChatProviders = func(context.Context) ([]database.ChatProvider, error) { + call := store.enabledProvidersCalls.Load() + return []database.ChatProvider{testChatProvider(fmt.Sprintf("provider-%d", call))}, nil + } + cache := newChatConfigCache(ctx, store, clock) + + first, err := cache.EnabledProviders(ctx) + require.NoError(t, err) + clock.Advance(chatConfigProvidersTTL).MustWait(ctx) + second, err := cache.EnabledProviders(ctx) + require.NoError(t, err) + + require.NotEqual(t, first, second) + require.Equal(t, int32(2), store.enabledProvidersCalls.Load()) +} + +func TestConfigCache_EnabledProviders_Invalidation(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + store := &stubChatConfigStore{} + store.getEnabledChatProviders = func(context.Context) ([]database.ChatProvider, error) { + call := store.enabledProvidersCalls.Load() + return []database.ChatProvider{testChatProvider(fmt.Sprintf("provider-%d", call))}, nil + } + cache := newChatConfigCache(ctx, store, clock) + + first, err := cache.EnabledProviders(ctx) + require.NoError(t, err) + cache.InvalidateProviders() + second, err := cache.EnabledProviders(ctx) + require.NoError(t, err) + + require.NotEqual(t, first, second) + require.Equal(t, int32(2), store.enabledProvidersCalls.Load()) +} + +func TestConfigCache_ModelConfigByID_CacheHit(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + configID := uuid.New() + config := testChatModelConfig(configID, "model-a") + store := &stubChatConfigStore{ + getChatModelConfigByID: func(context.Context, uuid.UUID) (database.ChatModelConfig, error) { + return config, nil + }, + } + cache := newChatConfigCache(ctx, store, clock) + + first, err := cache.ModelConfigByID(ctx, configID) + require.NoError(t, err) + second, err := cache.ModelConfigByID(ctx, configID) + require.NoError(t, err) + + require.Equal(t, config, first) + require.Equal(t, config, second) + require.Equal(t, int32(1), store.modelConfigByIDCalls.Load()) +} + +func TestConfigCache_ModelConfigByID_ClonesOptionsForCache(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + configID := uuid.New() + const options = `{"temperature":0.1}` + config := testChatModelConfig(configID, "model-a") + config.Options = []byte(options) + store := &stubChatConfigStore{ + getChatModelConfigByID: func(context.Context, uuid.UUID) (database.ChatModelConfig, error) { + return config, nil + }, + } + cache := newChatConfigCache(ctx, store, clock) + + // First call populates cache via singleflight. + first, err := cache.ModelConfigByID(ctx, configID) + require.NoError(t, err) + first.Options[0] = 'x' // mutate singleflight return + + // Second call is a cache hit. + second, err := cache.ModelConfigByID(ctx, configID) + require.NoError(t, err) + require.Equal(t, options, string(second.Options)) + second.Options[0] = 'y' // mutate cache-hit return + + // Third call is another cache hit — must be unaffected. + third, err := cache.ModelConfigByID(ctx, configID) + require.NoError(t, err) + require.Equal(t, options, string(third.Options)) +} + +func TestConfigCache_ModelConfigByID_NotFound(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + configID := uuid.New() + store := &stubChatConfigStore{ + getChatModelConfigByID: func(context.Context, uuid.UUID) (database.ChatModelConfig, error) { + return database.ChatModelConfig{}, sql.ErrNoRows + }, + } + cache := newChatConfigCache(ctx, store, clock) + + _, err := cache.ModelConfigByID(ctx, configID) + require.ErrorIs(t, err, sql.ErrNoRows) + _, err = cache.ModelConfigByID(ctx, configID) + require.ErrorIs(t, err, sql.ErrNoRows) + + require.Equal(t, int32(2), store.modelConfigByIDCalls.Load()) + _, ok := cache.modelConfigs[configID] + require.False(t, ok) +} + +func TestConfigCache_InvalidateModelConfig_CascadesToDefault(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + configID := uuid.New() + config := testChatModelConfig(configID, "model-a") + store := &stubChatConfigStore{} + store.getChatModelConfigByID = func(context.Context, uuid.UUID) (database.ChatModelConfig, error) { + return config, nil + } + store.getDefaultChatModelConfig = func(context.Context) (database.ChatModelConfig, error) { + call := store.defaultModelConfigCall.Load() + return testChatModelConfig(uuid.New(), fmt.Sprintf("default-model-%d", call)), nil + } + cache := newChatConfigCache(ctx, store, clock) + + _, err := cache.ModelConfigByID(ctx, configID) + require.NoError(t, err) + firstDefault, err := cache.DefaultModelConfig(ctx) + require.NoError(t, err) + + cache.InvalidateModelConfig(configID) + require.Nil(t, cache.defaultModelConfig) + + secondDefault, err := cache.DefaultModelConfig(ctx) + require.NoError(t, err) + + require.NotEqual(t, firstDefault, secondDefault) + require.Equal(t, int32(2), store.defaultModelConfigCall.Load()) +} + +func TestConfigCache_UserPrompt_NegativeCaching(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + userID := uuid.New() + store := &stubChatConfigStore{ + getUserChatCustomPrompt: func(context.Context, uuid.UUID) (string, error) { + return "", sql.ErrNoRows + }, + } + cache := newChatConfigCache(ctx, store, clock) + + first, err := cache.UserPrompt(ctx, userID) + require.NoError(t, err) + second, err := cache.UserPrompt(ctx, userID) + require.NoError(t, err) + + require.Empty(t, first) + require.Empty(t, second) + require.Equal(t, int32(1), store.userPromptCalls.Load()) +} + +func TestConfigCache_UserPrompt_ExpiredEntryRefetches(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + userID := uuid.New() + store := &stubChatConfigStore{} + store.getUserChatCustomPrompt = func(context.Context, uuid.UUID) (string, error) { + call := store.userPromptCalls.Load() + return fmt.Sprintf("prompt-%d", call), nil + } + cache := newChatConfigCache(ctx, store, clock) + cache.userPrompts.Set(userID, "stale", 0) + + first, err := cache.UserPrompt(ctx, userID) + require.NoError(t, err) + second, err := cache.UserPrompt(ctx, userID) + require.NoError(t, err) + + require.Equal(t, "prompt-1", first) + require.Equal(t, first, second) + require.Equal(t, int32(1), store.userPromptCalls.Load()) +} + +func TestConfigCache_InvalidateUserPrompt(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + userID := uuid.New() + store := &stubChatConfigStore{} + store.getUserChatCustomPrompt = func(context.Context, uuid.UUID) (string, error) { + call := store.userPromptCalls.Load() + return fmt.Sprintf("prompt-%d", call), nil + } + cache := newChatConfigCache(ctx, store, clock) + + first, err := cache.UserPrompt(ctx, userID) + require.NoError(t, err) + cache.InvalidateUserPrompt(userID) + second, err := cache.UserPrompt(ctx, userID) + require.NoError(t, err) + + require.NotEqual(t, first, second) + require.Equal(t, int32(2), store.userPromptCalls.Load()) +} + +func TestConfigCache_InvalidateUserPrompt_BlocksStaleInFlightPrompt(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitMedium) + clock := quartz.NewMock(t) + userID := uuid.New() + const stalePrompt = "stale prompt" + const freshPrompt = "fresh prompt" + firstStarted := make(chan struct{}) + secondStarted := make(chan struct{}) + releaseFirst := make(chan struct{}) + releaseSecond := make(chan struct{}) + store := &stubChatConfigStore{} + store.getUserChatCustomPrompt = func(context.Context, uuid.UUID) (string, error) { + switch call := store.userPromptCalls.Load(); call { + case 1: + close(firstStarted) + <-releaseFirst + return stalePrompt, nil + case 2: + close(secondStarted) + <-releaseSecond + return freshPrompt, nil + default: + return "", xerrors.Errorf("unexpected user prompt call %d", call) + } + } + cache := newChatConfigCache(ctx, store, clock) + + type result struct { + prompt string + err error + } + + firstResult := make(chan result, 1) + go func() { + prompt, err := cache.UserPrompt(ctx, userID) + firstResult <- result{prompt: prompt, err: err} + }() + + waitForSignal(t, firstStarted) + cache.InvalidateUserPrompt(userID) + + secondResult := make(chan result, 1) + go func() { + prompt, err := cache.UserPrompt(ctx, userID) + secondResult <- result{prompt: prompt, err: err} + }() + + waitForSignal(t, secondStarted) + close(releaseFirst) + first := <-firstResult + require.NoError(t, first.err) + require.Equal(t, stalePrompt, first.prompt) + _, _, ok := cache.userPrompts.Get(userID) + require.False(t, ok) + + close(releaseSecond) + second := <-secondResult + require.NoError(t, second.err) + require.Equal(t, freshPrompt, second.prompt) + require.Equal(t, int32(2), store.userPromptCalls.Load()) + + third, err := cache.UserPrompt(ctx, userID) + require.NoError(t, err) + require.Equal(t, freshPrompt, third) + require.Equal(t, int32(2), store.userPromptCalls.Load()) +} + +func TestConfigCache_Singleflight(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitMedium) + clock := quartz.NewMock(t) + providers := []database.ChatProvider{testChatProvider("provider-a")} + fetchStarted := make(chan struct{}) + releaseFetch := make(chan struct{}) + var startedOnce sync.Once + store := &stubChatConfigStore{} + store.getEnabledChatProviders = func(context.Context) ([]database.ChatProvider, error) { + startedOnce.Do(func() { close(fetchStarted) }) + <-releaseFetch + return providers, nil + } + cache := newChatConfigCache(ctx, store, clock) + + const callers = 8 + results := make([][]database.ChatProvider, callers) + errs := make([]error, callers) + var wg sync.WaitGroup + start := make(chan struct{}) + for i := 0; i < callers; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + <-start + results[i], errs[i] = cache.EnabledProviders(ctx) + }(i) + } + + close(start) + waitForSignal(t, fetchStarted) + close(releaseFetch) + wg.Wait() + + for i := 0; i < callers; i++ { + require.NoError(t, errs[i]) + require.Equal(t, providers, results[i]) + } + require.Equal(t, int32(1), store.enabledProvidersCalls.Load()) +} + +func TestConfigCache_GenerationPreventsStaleWrite(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitMedium) + clock := quartz.NewMock(t) + firstProviders := []database.ChatProvider{testChatProvider("provider-a")} + secondProviders := []database.ChatProvider{testChatProvider("provider-b")} + fetchStarted := make(chan struct{}) + releaseFetch := make(chan struct{}) + var startedOnce sync.Once + store := &stubChatConfigStore{} + store.getEnabledChatProviders = func(context.Context) ([]database.ChatProvider, error) { + call := store.enabledProvidersCalls.Load() + if call == 1 { + startedOnce.Do(func() { close(fetchStarted) }) + <-releaseFetch + return firstProviders, nil + } + return secondProviders, nil + } + cache := newChatConfigCache(ctx, store, clock) + + resultCh := make(chan []database.ChatProvider, 1) + errCh := make(chan error, 1) + go func() { + providers, err := cache.EnabledProviders(ctx) + if err != nil { + errCh <- err + return + } + resultCh <- providers + }() + + waitForSignal(t, fetchStarted) + cache.InvalidateProviders() + close(releaseFetch) + + select { + case err := <-errCh: + require.NoError(t, err) + case providers := <-resultCh: + require.Equal(t, firstProviders, providers) + case <-time.After(testutil.WaitShort): + t.Fatal("timed out waiting for in-flight fetch") + } + + require.Nil(t, cache.providers) + second, err := cache.EnabledProviders(ctx) + require.NoError(t, err) + require.Equal(t, secondProviders, second) + require.Equal(t, int32(2), store.enabledProvidersCalls.Load()) +} + +func TestConfigCache_InvalidateProviders_BlocksStaleInFlightProviders(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitMedium) + clock := quartz.NewMock(t) + staleProviders := []database.ChatProvider{testChatProvider("provider-stale")} + freshProviders := []database.ChatProvider{testChatProvider("provider-fresh")} + firstStarted := make(chan struct{}) + secondStarted := make(chan struct{}) + releaseFirst := make(chan struct{}) + releaseSecond := make(chan struct{}) + store := &stubChatConfigStore{} + store.getEnabledChatProviders = func(context.Context) ([]database.ChatProvider, error) { + switch call := store.enabledProvidersCalls.Load(); call { + case 1: + close(firstStarted) + <-releaseFirst + return staleProviders, nil + case 2: + close(secondStarted) + <-releaseSecond + return freshProviders, nil + default: + return nil, xerrors.Errorf("unexpected provider call %d", call) + } + } + cache := newChatConfigCache(ctx, store, clock) + + type result struct { + providers []database.ChatProvider + err error + } + + firstResult := make(chan result, 1) + go func() { + providers, err := cache.EnabledProviders(ctx) + firstResult <- result{providers: providers, err: err} + }() + + waitForSignal(t, firstStarted) + cache.InvalidateProviders() + + secondResult := make(chan result, 1) + go func() { + providers, err := cache.EnabledProviders(ctx) + secondResult <- result{providers: providers, err: err} + }() + + waitForSignal(t, secondStarted) + close(releaseFirst) + first := <-firstResult + require.NoError(t, first.err) + require.Equal(t, staleProviders, first.providers) + require.Nil(t, cache.providers) + + close(releaseSecond) + second := <-secondResult + require.NoError(t, second.err) + require.Equal(t, freshProviders, second.providers) + require.Equal(t, int32(2), store.enabledProvidersCalls.Load()) + + third, err := cache.EnabledProviders(ctx) + require.NoError(t, err) + require.Equal(t, freshProviders, third) + require.Equal(t, int32(2), store.enabledProvidersCalls.Load()) +} + +func TestConfigCache_InvalidateProviders_CascadesToModelConfigs(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + configID := uuid.New() + store := &stubChatConfigStore{} + store.getChatModelConfigByID = func(context.Context, uuid.UUID) (database.ChatModelConfig, error) { + call := store.modelConfigByIDCalls.Load() + return testChatModelConfig(configID, fmt.Sprintf("model-%d", call)), nil + } + cache := newChatConfigCache(ctx, store, clock) + + first, err := cache.ModelConfigByID(ctx, configID) + require.NoError(t, err) + cache.InvalidateProviders() + second, err := cache.ModelConfigByID(ctx, configID) + require.NoError(t, err) + + require.NotEqual(t, first, second) + require.Equal(t, int32(2), store.modelConfigByIDCalls.Load()) +} + +func TestConfigCache_InvalidateProviders_CascadesToDefaultModelConfig(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + store := &stubChatConfigStore{} + store.getDefaultChatModelConfig = func(context.Context) (database.ChatModelConfig, error) { + call := store.defaultModelConfigCall.Load() + return testChatModelConfig(uuid.New(), fmt.Sprintf("default-model-%d", call)), nil + } + cache := newChatConfigCache(ctx, store, clock) + + first, err := cache.DefaultModelConfig(ctx) + require.NoError(t, err) + cache.InvalidateProviders() + second, err := cache.DefaultModelConfig(ctx) + require.NoError(t, err) + + require.NotEqual(t, first, second) + require.Equal(t, int32(2), store.defaultModelConfigCall.Load()) +} + +func TestConfigCache_InvalidateProviders_BlocksStaleInFlightModelConfig(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitMedium) + clock := quartz.NewMock(t) + configID := uuid.New() + staleConfig := testChatModelConfig(configID, "stale-model") + freshConfig := testChatModelConfig(configID, "fresh-model") + firstStarted := make(chan struct{}) + secondStarted := make(chan struct{}) + releaseFirst := make(chan struct{}) + releaseSecond := make(chan struct{}) + store := &stubChatConfigStore{} + store.getChatModelConfigByID = func(context.Context, uuid.UUID) (database.ChatModelConfig, error) { + switch call := store.modelConfigByIDCalls.Load(); call { + case 1: + close(firstStarted) + <-releaseFirst + return staleConfig, nil + case 2: + close(secondStarted) + <-releaseSecond + return freshConfig, nil + default: + return database.ChatModelConfig{}, xerrors.Errorf("unexpected model config call %d", call) + } + } + cache := newChatConfigCache(ctx, store, clock) + + type result struct { + config database.ChatModelConfig + err error + } + + firstResult := make(chan result, 1) + go func() { + config, err := cache.ModelConfigByID(ctx, configID) + firstResult <- result{config: config, err: err} + }() + + waitForSignal(t, firstStarted) + cache.InvalidateProviders() + + secondResult := make(chan result, 1) + go func() { + config, err := cache.ModelConfigByID(ctx, configID) + secondResult <- result{config: config, err: err} + }() + + waitForSignal(t, secondStarted) + close(releaseFirst) + first := <-firstResult + require.NoError(t, first.err) + require.Equal(t, staleConfig, first.config) + _, ok := cache.modelConfigs[configID] + require.False(t, ok) + + close(releaseSecond) + second := <-secondResult + require.NoError(t, second.err) + require.Equal(t, freshConfig, second.config) + require.Equal(t, int32(2), store.modelConfigByIDCalls.Load()) + + third, err := cache.ModelConfigByID(ctx, configID) + require.NoError(t, err) + require.Equal(t, freshConfig, third) + require.Equal(t, int32(2), store.modelConfigByIDCalls.Load()) +} + +func testChatProvider(name string) database.ChatProvider { + return database.ChatProvider{ + ID: uuid.New(), + Provider: name, + DisplayName: name, + Enabled: true, + CreatedAt: time.Unix(0, 0).UTC(), + UpdatedAt: time.Unix(0, 0).UTC(), + } +} + +func testChatModelConfig(id uuid.UUID, model string) database.ChatModelConfig { + return database.ChatModelConfig{ + ID: id, + Provider: "openai", + Model: model, + DisplayName: model, + Enabled: true, + CreatedAt: time.Unix(0, 0).UTC(), + UpdatedAt: time.Unix(0, 0).UTC(), + ContextLimit: 128000, + CompressionThreshold: 64000, + } +} + +func waitForSignal(t *testing.T, ch <-chan struct{}) { + t.Helper() + + select { + case <-ch: + case <-time.After(testutil.WaitShort): + t.Fatal("timed out waiting for signal") + } +} + +// TestConfigCache_CallerCancellation verifies the DoChan-based +// cancellation semantics across all four cache methods: +// - A canceled caller returns immediately without waiting for the +// shared fill to complete. +// - One canceled waiter does not poison other coalesced waiters. +// - Server context cancellation propagates through the fill. +func TestConfigCache_CallerCancellation(t *testing.T) { + t.Parallel() + + type cacheMethod struct { + name string + // setupBlocked configures the store to block on release. + // The started channel is closed when the fill enters the + // store. The release channel unblocks the store. + setupBlocked func(store *stubChatConfigStore, started, release chan struct{}) + // setupCtxSensitive configures the store to block until + // its context is canceled (for server-shutdown testing). + setupCtxSensitive func(store *stubChatConfigStore, started chan struct{}) + // call invokes the cache method under test. + call func(ctx context.Context, cache *chatConfigCache) error + // storeCalls returns the number of underlying store calls. + storeCalls func(store *stubChatConfigStore) int32 + } + + configID := uuid.New() + userID := uuid.New() + + methods := []cacheMethod{ + { + name: "EnabledProviders", + setupBlocked: func(store *stubChatConfigStore, started, release chan struct{}) { + var once sync.Once + store.getEnabledChatProviders = func(ctx context.Context) ([]database.ChatProvider, error) { + once.Do(func() { close(started) }) + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-release: + return []database.ChatProvider{testChatProvider("p")}, nil + } + } + }, + setupCtxSensitive: func(store *stubChatConfigStore, started chan struct{}) { + var once sync.Once + store.getEnabledChatProviders = func(ctx context.Context) ([]database.ChatProvider, error) { + once.Do(func() { close(started) }) + <-ctx.Done() + return nil, ctx.Err() + } + }, + call: func(ctx context.Context, cache *chatConfigCache) error { + _, err := cache.EnabledProviders(ctx) + return err + }, + storeCalls: func(store *stubChatConfigStore) int32 { + return store.enabledProvidersCalls.Load() + }, + }, + { + name: "ModelConfigByID", + setupBlocked: func(store *stubChatConfigStore, started, release chan struct{}) { + var once sync.Once + store.getChatModelConfigByID = func(ctx context.Context, id uuid.UUID) (database.ChatModelConfig, error) { + once.Do(func() { close(started) }) + select { + case <-ctx.Done(): + return database.ChatModelConfig{}, ctx.Err() + case <-release: + return testChatModelConfig(id, "model"), nil + } + } + }, + setupCtxSensitive: func(store *stubChatConfigStore, started chan struct{}) { + var once sync.Once + store.getChatModelConfigByID = func(ctx context.Context, _ uuid.UUID) (database.ChatModelConfig, error) { + once.Do(func() { close(started) }) + <-ctx.Done() + return database.ChatModelConfig{}, ctx.Err() + } + }, + call: func(ctx context.Context, cache *chatConfigCache) error { + _, err := cache.ModelConfigByID(ctx, configID) + return err + }, + storeCalls: func(store *stubChatConfigStore) int32 { + return store.modelConfigByIDCalls.Load() + }, + }, + { + name: "DefaultModelConfig", + setupBlocked: func(store *stubChatConfigStore, started, release chan struct{}) { + var once sync.Once + store.getDefaultChatModelConfig = func(ctx context.Context) (database.ChatModelConfig, error) { + once.Do(func() { close(started) }) + select { + case <-ctx.Done(): + return database.ChatModelConfig{}, ctx.Err() + case <-release: + return testChatModelConfig(uuid.New(), "default"), nil + } + } + }, + setupCtxSensitive: func(store *stubChatConfigStore, started chan struct{}) { + var once sync.Once + store.getDefaultChatModelConfig = func(ctx context.Context) (database.ChatModelConfig, error) { + once.Do(func() { close(started) }) + <-ctx.Done() + return database.ChatModelConfig{}, ctx.Err() + } + }, + call: func(ctx context.Context, cache *chatConfigCache) error { + _, err := cache.DefaultModelConfig(ctx) + return err + }, + storeCalls: func(store *stubChatConfigStore) int32 { + return store.defaultModelConfigCall.Load() + }, + }, + { + name: "UserPrompt", + setupBlocked: func(store *stubChatConfigStore, started, release chan struct{}) { + var once sync.Once + store.getUserChatCustomPrompt = func(ctx context.Context, _ uuid.UUID) (string, error) { + once.Do(func() { close(started) }) + select { + case <-ctx.Done(): + return "", ctx.Err() + case <-release: + return "custom prompt", nil + } + } + }, + setupCtxSensitive: func(store *stubChatConfigStore, started chan struct{}) { + var once sync.Once + store.getUserChatCustomPrompt = func(ctx context.Context, _ uuid.UUID) (string, error) { + once.Do(func() { close(started) }) + <-ctx.Done() + return "", ctx.Err() + } + }, + call: func(ctx context.Context, cache *chatConfigCache) error { + _, err := cache.UserPrompt(ctx, userID) + return err + }, + storeCalls: func(store *stubChatConfigStore) int32 { + return store.userPromptCalls.Load() + }, + }, + } + + // Test A: A canceled caller stops waiting immediately; the + // shared fill still completes and populates the cache. + t.Run("CanceledCallerStopsWaiting", func(t *testing.T) { + t.Parallel() + for _, m := range methods { + t.Run(m.name, func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitMedium) + clock := quartz.NewMock(t) + store := &stubChatConfigStore{} + started := make(chan struct{}) + release := make(chan struct{}) + m.setupBlocked(store, started, release) + cache := newChatConfigCache(ctx, store, clock) + + callerCtx, callerCancel := context.WithCancel(ctx) + errCh := make(chan error, 1) + go func() { + errCh <- m.call(callerCtx, cache) + }() + + // Wait for the fill to enter the store, then + // cancel the caller's context. + waitForSignal(t, started) + callerCancel() + + select { + case err := <-errCh: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(testutil.WaitShort): + t.Fatal("canceled caller did not return promptly") + } + + // Release the store so the fill can complete. + close(release) + + // A fresh call must succeed — either a cache + // hit or by joining the still-in-flight fill. + // Only one store call should have occurred. + require.NoError(t, m.call(ctx, cache)) + require.Equal(t, int32(1), m.storeCalls(store)) + }) + } + }) + + // Test B: One canceled waiter does not poison other coalesced + // waiters sharing the same singleflight entry. + t.Run("CanceledWaiterDoesNotPoisonOthers", func(t *testing.T) { + t.Parallel() + for _, m := range methods { + t.Run(m.name, func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitMedium) + clock := quartz.NewMock(t) + store := &stubChatConfigStore{} + started := make(chan struct{}) + release := make(chan struct{}) + m.setupBlocked(store, started, release) + cache := newChatConfigCache(ctx, store, clock) + + cancelCtx, cancel := context.WithCancel(ctx) + cancelErrCh := make(chan error, 1) + survivorErrCh := make(chan error, 1) + + go func() { + cancelErrCh <- m.call(cancelCtx, cache) + }() + go func() { + survivorErrCh <- m.call(ctx, cache) + }() + + waitForSignal(t, started) + cancel() + + select { + case err := <-cancelErrCh: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(testutil.WaitShort): + t.Fatal("canceled caller did not return promptly") + } + + // Release the store; the surviving waiter + // must receive the successful result. + close(release) + + select { + case err := <-survivorErrCh: + require.NoError(t, err) + case <-time.After(testutil.WaitShort): + t.Fatal("survivor caller did not return") + } + + require.Equal(t, int32(1), m.storeCalls(store)) + }) + } + }) + + // Test C: Server context cancellation propagates through the + // fill, ensuring graceful shutdown behavior is preserved. + t.Run("ServerCancellation", func(t *testing.T) { + t.Parallel() + for _, m := range methods { + t.Run(m.name, func(t *testing.T) { + t.Parallel() + clock := quartz.NewMock(t) + store := &stubChatConfigStore{} + started := make(chan struct{}) + m.setupCtxSensitive(store, started) + + serverCtx, serverCancel := context.WithCancel(context.Background()) + defer serverCancel() + cache := newChatConfigCache(serverCtx, store, clock) + + callerCtx := testutil.Context(t, testutil.WaitMedium) + errCh := make(chan error, 1) + go func() { + errCh <- m.call(callerCtx, cache) + }() + + waitForSignal(t, started) + serverCancel() + + select { + case err := <-errCh: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(testutil.WaitShort): + t.Fatal("caller did not return after server cancel") + } + }) + } + }) +} diff --git a/coderd/x/chatd/subagent.go b/coderd/x/chatd/subagent.go index cd1d71a0e1..77cd234efb 100644 --- a/coderd/x/chatd/subagent.go +++ b/coderd/x/chatd/subagent.go @@ -72,7 +72,7 @@ func (p *Server) isAnthropicConfigured(ctx context.Context) bool { if p.providerAPIKeys.APIKey("anthropic") != "" { return true } - dbProviders, err := p.db.GetEnabledChatProviders(ctx) + dbProviders, err := p.configCache.EnabledProviders(ctx) if err != nil { return false }