Thread context.Context through modelsdev store API

Replace the lazy-loaded closure-based Store with a simple struct
that holds a cache file path and loads the database on first
GetDatabase call. Pass caller-provided context to the HTTP request
in fetchFromAPI instead of using context.Background().

NewStore() remains context-free (no I/O). All Store methods
(GetDatabase, GetProvider, GetModel, ResolveModelAlias) and
ModelSupportsReasoning accept context for the network call.
Updated all callsites, using context.Background() only in TUI
callbacks where no caller context is available.

Assisted-By: cagent

Author: rumpl

pkg/config/examples_test.go

@@ -70,7 +70,7 @@ func TestParseExamples(t *testing.T) {
 					continue
 				}
 
-				model, err := modelsStore.GetModel(model.Provider + "/" + model.Model)
+				model, err := modelsStore.GetModel(t.Context(), model.Provider+"/"+model.Model)
 				require.NoError(t, err)
 				require.NotNil(t, model)
 			}

pkg/config/model_alias.go

@@ -1,6 +1,7 @@
 package config
 
 import (
+	"context"
 	"log/slog"
 	"strings"
 
@@ -16,7 +17,7 @@ import (
 // either set directly on the model or inherited from a custom provider definition.
 // This is necessary because external providers (like Azure Foundry) may use the alias
 // names directly as deployment names rather than the pinned version names.
-func ResolveModelAliases(cfg *latest.Config, store *modelsdev.Store) {
+func ResolveModelAliases(ctx context.Context, cfg *latest.Config, store *modelsdev.Store) {
 	// Resolve model aliases in the models section
 	for name, modelCfg := range cfg.Models {
 		// Skip alias resolution for models with custom base_url (direct or via provider)
@@ -27,15 +28,15 @@ func ResolveModelAliases(cfg *latest.Config, store *modelsdev.Store) {
 			continue
 		}
 
-		if resolved := store.ResolveModelAlias(modelCfg.Provider, modelCfg.Model); resolved != modelCfg.Model {
+		if resolved := store.ResolveModelAlias(ctx, modelCfg.Provider, modelCfg.Model); resolved != modelCfg.Model {
 			modelCfg.Model = resolved
 			cfg.Models[name] = modelCfg
 		}
 
 		// Resolve model aliases in routing rules
 		for i, rule := range modelCfg.Routing {
 			if provider, model, ok := strings.Cut(rule.Model, "/"); ok {
-				if resolved := store.ResolveModelAlias(provider, model); resolved != model {
+				if resolved := store.ResolveModelAlias(ctx, provider, model); resolved != model {
 					modelCfg.Routing[i].Model = provider + "/" + resolved
 				}
 			}
@@ -52,7 +53,7 @@ func ResolveModelAliases(cfg *latest.Config, store *modelsdev.Store) {
 		var resolvedModels []string
 		for modelRef := range strings.SplitSeq(agent.Model, ",") {
 			if provider, model, ok := strings.Cut(modelRef, "/"); ok {
-				if resolved := store.ResolveModelAlias(provider, model); resolved != model {
+				if resolved := store.ResolveModelAlias(ctx, provider, model); resolved != model {
 					resolvedModels = append(resolvedModels, provider+"/"+resolved)
 					continue
 				}

pkg/config/model_alias_test.go

@@ -237,7 +237,7 @@ func TestResolveModelAliases(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			ResolveModelAliases(tt.cfg, store)
+			ResolveModelAliases(t.Context(), tt.cfg, store)
 			assert.Equal(t, tt.expected, tt.cfg)
 		})
 	}

pkg/model/provider/bedrock/client.go

@@ -112,7 +112,7 @@ func NewClient(ctx context.Context, cfg *latest.ModelConfig, env environment.Pro
 
 	// Detect prompt caching capability at init time for efficiency.
 	// Uses models.dev cache pricing as proxy for capability detection.
-	cachingSupported := detectCachingSupport(cfg.Model)
+	cachingSupported := detectCachingSupport(ctx, cfg.Model)
 
 	slog.Debug("Bedrock client created successfully",
 		"model", cfg.Model,
@@ -133,15 +133,15 @@ func NewClient(ctx context.Context, cfg *latest.ModelConfig, env environment.Pro
 // detectCachingSupport checks if a model supports prompt caching using models.dev data.
 // Models with non-zero CacheRead/CacheWrite costs support prompt caching.
 // Returns false on lookup failure (safe default for unsupported models).
-func detectCachingSupport(model string) bool {
+func detectCachingSupport(ctx context.Context, model string) bool {
 	store, err := modelsdev.NewStore()
 	if err != nil {
 		slog.Debug("Bedrock models store unavailable, prompt caching disabled", "error", err)
 		return false
 	}
 
 	modelID := "amazon-bedrock/" + model
-	m, err := store.GetModel(modelID)
+	m, err := store.GetModel(ctx, modelID)
 	if err != nil {
 		slog.Debug("Bedrock prompt caching disabled: model not found in models.dev",
 			"model_id", modelID, "error", err)

pkg/model/provider/bedrock/client_test.go

@@ -1249,23 +1249,23 @@ func TestDetectCachingSupport_SupportedModel(t *testing.T) {
 	t.Parallel()
 
 	// Uses real models.dev lookup to verify Claude models support caching
-	supported := detectCachingSupport("anthropic.claude-3-5-sonnet-20241022-v2:0")
+	supported := detectCachingSupport(t.Context(), "anthropic.claude-3-5-sonnet-20241022-v2:0")
 	assert.True(t, supported)
 }
 
 func TestDetectCachingSupport_UnsupportedModel(t *testing.T) {
 	t.Parallel()
 
 	// Llama doesn't have cache pricing in models.dev
-	supported := detectCachingSupport("meta.llama3-8b-instruct-v1:0")
+	supported := detectCachingSupport(t.Context(), "meta.llama3-8b-instruct-v1:0")
 	assert.False(t, supported)
 }
 
 func TestDetectCachingSupport_UnknownModel(t *testing.T) {
 	t.Parallel()
 
 	// Unknown model should gracefully return false, not panic
-	supported := detectCachingSupport("nonexistent.model.that.does.not.exist:v1")
+	supported := detectCachingSupport(t.Context(), "nonexistent.model.that.does.not.exist:v1")
 	assert.False(t, supported)
 }
 

pkg/modelsdev/store.go

@@ -21,25 +21,16 @@ const (
 )
 
 // Store manages access to the models.dev data.
-// The database is loaded lazily on first access and cached for the
-// lifetime of the Store. All methods are safe for concurrent use.
+// All methods are safe for concurrent use.
 type Store struct {
-	db func() (*Database, error)
+	cacheFile string
+	mu        sync.Mutex
+	db        *Database
 }
 
-// defaultStore is a cached singleton store instance for repeated access.
-var defaultStore = sync.OnceValues(newStoreInternal)
-
-// NewStore returns the cached default store instance.
-// The underlying database is fetched lazily on first access
-// from a local cache file or the models.dev API.
+// NewStore creates a new models.dev store.
+// The database is loaded on first access via GetDatabase.
 func NewStore() (*Store, error) {
-	return defaultStore()
-}
-
-// newStoreInternal creates a new models.dev store that loads data
-// from the filesystem cache or the network on first access.
-func newStoreInternal() (*Store, error) {
 	homeDir, err := os.UserHomeDir()
 	if err != nil {
 		return nil, fmt.Errorf("failed to get user home directory: %w", err)
@@ -50,12 +41,8 @@ func newStoreInternal() (*Store, error) {
 		return nil, fmt.Errorf("failed to create cache directory: %w", err)
 	}
 
-	cacheFile := filepath.Join(cacheDir, CacheFileName)
-
 	return &Store{
-		db: sync.OnceValues(func() (*Database, error) {
-			return loadDatabase(cacheFile)
-		}),
+		cacheFile: filepath.Join(cacheDir, CacheFileName),
 	}, nil
 }
 
@@ -64,19 +51,30 @@ func newStoreInternal() (*Store, error) {
 // from the network or touches the filesystem, making it suitable for
 // tests and any scenario where the provider data is already known.
 func NewDatabaseStore(db *Database) *Store {
-	return &Store{
-		db: func() (*Database, error) { return db, nil },
-	}
+	return &Store{db: db}
 }
 
 // GetDatabase returns the models.dev database, fetching from cache or API as needed.
-func (s *Store) GetDatabase() (*Database, error) {
-	return s.db()
+func (s *Store) GetDatabase(ctx context.Context) (*Database, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if s.db != nil {
+		return s.db, nil
+	}
+
+	db, err := loadDatabase(ctx, s.cacheFile)
+	if err != nil {
+		return nil, err
+	}
+
+	s.db = db
+	return db, nil
 }
 
 // GetProvider returns a specific provider by ID.
-func (s *Store) GetProvider(providerID string) (*Provider, error) {
-	db, err := s.db()
+func (s *Store) GetProvider(ctx context.Context, providerID string) (*Provider, error) {
+	db, err := s.GetDatabase(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -90,15 +88,15 @@ func (s *Store) GetProvider(providerID string) (*Provider, error) {
 }
 
 // GetModel returns a specific model by provider ID and model ID.
-func (s *Store) GetModel(id string) (*Model, error) {
+func (s *Store) GetModel(ctx context.Context, id string) (*Model, error) {
 	parts := strings.SplitN(id, "/", 2)
 	if len(parts) != 2 {
 		return nil, fmt.Errorf("invalid model ID: %q", id)
 	}
 	providerID := parts[0]
 	modelID := parts[1]
 
-	provider, err := s.GetProvider(providerID)
+	provider, err := s.GetProvider(ctx, providerID)
 	if err != nil {
 		return nil, err
 	}
@@ -130,15 +128,15 @@ func (s *Store) GetModel(id string) (*Model, error) {
 
 // loadDatabase loads the database from the local cache file or
 // falls back to fetching from the models.dev API.
-func loadDatabase(cacheFile string) (*Database, error) {
+func loadDatabase(ctx context.Context, cacheFile string) (*Database, error) {
 	// Try to load from cache first
 	cached, err := loadFromCache(cacheFile)
 	if err == nil && time.Since(cached.LastRefresh) < refreshInterval {
 		return &cached.Database, nil
 	}
 
 	// Cache is invalid or doesn't exist, fetch from API
-	database, fetchErr := fetchFromAPI()
+	database, fetchErr := fetchFromAPI(ctx)
 	if fetchErr != nil {
 		// If API fetch fails, but we have cached data, use it
 		if cached != nil {
@@ -156,8 +154,8 @@ func loadDatabase(cacheFile string) (*Database, error) {
 	return database, nil
 }
 
-func fetchFromAPI() (*Database, error) {
-	req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, ModelsDevAPIURL, http.NoBody)
+func fetchFromAPI(ctx context.Context) (*Database, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, ModelsDevAPIURL, http.NoBody)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create request: %w", err)
 	}
@@ -225,7 +223,7 @@ var datePattern = regexp.MustCompile(`-\d{4}-?\d{2}-?\d{2}$`)
 // For example, ("anthropic", "claude-sonnet-4-5") might resolve to "claude-sonnet-4-5-20250929".
 // If the model is not an alias (already pinned or unknown), the original model name is returned.
 // This method uses the models.dev database to find the corresponding pinned version.
-func (s *Store) ResolveModelAlias(providerID, modelName string) string {
+func (s *Store) ResolveModelAlias(ctx context.Context, providerID, modelName string) string {
 	if providerID == "" || modelName == "" {
 		return modelName
 	}
@@ -236,7 +234,7 @@ func (s *Store) ResolveModelAlias(providerID, modelName string) string {
 	}
 
 	// Get the provider from the database
-	provider, err := s.GetProvider(providerID)
+	provider, err := s.GetProvider(ctx, providerID)
 	if err != nil {
 		return modelName
 	}
@@ -285,7 +283,7 @@ func isBedrockRegionPrefix(prefix string) bool {
 //   - If modelID is empty or not in "provider/model" format, returns true (fail-open)
 //   - If models.dev lookup fails for any reason, returns true (fail-open)
 //   - If lookup succeeds, returns the model's Reasoning field value
-func ModelSupportsReasoning(modelID string) bool {
+func ModelSupportsReasoning(ctx context.Context, modelID string) bool {
 	// Fail-open for empty model ID
 	if modelID == "" {
 		return true
@@ -303,7 +301,7 @@ func ModelSupportsReasoning(modelID string) bool {
 		return true
 	}
 
-	model, err := store.GetModel(modelID)
+	model, err := store.GetModel(ctx, modelID)
 	if err != nil {
 		slog.Debug("Failed to lookup model in models.dev, assuming reasoning supported to allow user choice", "model_id", modelID, "error", err)
 		return true

pkg/modelsdev/store_test.go

@@ -57,7 +57,7 @@ func TestResolveModelAlias(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			result := store.ResolveModelAlias(tt.provider, tt.model)
+			result := store.ResolveModelAlias(t.Context(), tt.provider, tt.model)
 			assert.Equal(t, tt.expected, result)
 		})
 	}

pkg/rag/strategy/semantic_embeddings.go

@@ -506,7 +506,7 @@ func calculateSemanticUsageCost(modelsStore modelStore, modelID string, usage *c
 		return 0
 	}
 
-	model, err := modelsStore.GetModel(modelID)
+	model, err := modelsStore.GetModel(context.Background(), modelID)
 	if err != nil {
 		slog.Debug("Failed to get semantic model pricing from models.dev, cost will be 0",
 			"model_id", modelID,

pkg/rag/strategy/vector_store.go

@@ -87,7 +87,7 @@ type VectorStore struct {
 }
 
 type modelStore interface {
-	GetModel(modelID string) (*modelsdev.Model, error)
+	GetModel(ctx context.Context, modelID string) (*modelsdev.Model, error)
 }
 
 // EmbeddingInputBuilder builds the string that will be sent to the embedding model
@@ -174,7 +174,7 @@ func (s *VectorStore) calculateCost(tokens int64) float64 {
 		return 0
 	}
 
-	model, err := s.modelsStore.GetModel(s.modelID)
+	model, err := s.modelsStore.GetModel(context.Background(), s.modelID)
 	if err != nil {
 		slog.Debug("Failed to get model pricing from models.dev, cost will be 0",
 			"model_id", s.modelID,

pkg/runtime/model_switcher.go

@@ -286,7 +286,7 @@ func (r *LocalRuntime) AvailableModels(ctx context.Context) []ModelChoice {
 // buildCatalogChoices builds ModelChoice entries from the models.dev catalog,
 // filtered by supported providers and available credentials.
 func (r *LocalRuntime) buildCatalogChoices(ctx context.Context) []ModelChoice {
-	db, err := r.modelsStore.GetDatabase()
+	db, err := r.modelsStore.GetDatabase(ctx)
 	if err != nil {
 		slog.Debug("Failed to get models.dev database for catalog", "error", err)
 		return nil
@@ -446,7 +446,7 @@ func (r *LocalRuntime) createProviderFromConfig(ctx context.Context, cfg *latest
 	if cfg.MaxTokens != nil {
 		opts = append(opts, options.WithMaxTokens(*cfg.MaxTokens))
 	} else if r.modelsStore != nil {
-		m, err := r.modelsStore.GetModel(cfg.Provider + "/" + cfg.Model)
+		m, err := r.modelsStore.GetModel(ctx, cfg.Provider+"/"+cfg.Model)
 		if err == nil && m != nil {
 			opts = append(opts, options.WithMaxTokens(m.Limit.Output))
 		}