perf(coderd/x/chatd): add process-wide config cache for hot DB queries (#23272)

## Summary

Adds a process-wide cache for three hot database queries in `chatd` that
were hitting Postgres on **every chat turn** despite returning
rarely-changing configuration data:

| Query | Before (50k turns) | After | Reduction |
|---|---|---|---|
| `GetEnabledChatProviders` | ~98.6k calls | ~500-1000 | ~99% |
| `GetChatModelConfigByID` | ~49.2k calls | ~500-1000 | ~98% |
| `GetUserChatCustomPrompt` | ~46.7k calls | ~1000-2000 | ~97% |

These were identified via `coder exp scaletest chat` (5000 concurrent
chats × 10 turns) as the dominant source of Postgres load during chat
processing.

## Design

Follows the established **webpush subscription cache pattern**
(`coderd/webpush/webpush.go`):
- `sync.RWMutex` + `tailscale.com/util/singleflight` (generic) +
generation-based stale prevention + TTL
- 10s TTL for provider/model config, 5s TTL for user prompts
- Negative caching for `sql.ErrNoRows` on user prompts (the common case
— most users don't set custom prompts)
- Deep-clones `ChatModelConfig.Options` (`json.RawMessage` = `[]byte`)
on both store and read paths

### Invalidation

Single pubsub channel (`chat:config_change`) with kind discriminator for
cross-replica cache invalidation. Seven publish points in
`coderd/chats.go` cover all admin mutation endpoints
(create/update/delete for providers and model configs, put for user
prompts).

_This PR was generated with mux and was reviewed by a human_
This commit is contained in:
Ethan
2026-03-26 18:04:53 +11:00
committed by GitHub
parent 2ff329b68a
commit 15f2fa55c6
6 changed files with 1515 additions and 5 deletions
+37
View File
@@ -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)
}
+52
View File
@@ -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"
)
+35 -4
View File
@@ -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 <user-instructions> 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()
+412
View File
@@ -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()
}
+978
View File
@@ -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")
}
})
}
})
}
+1 -1
View File
@@ -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
}