core/server/handlers/model_sync.go

package handlers

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"os"
	"strconv"
	"time"

	"switchboard-core/models"
	"switchboard-core/providers"
	"switchboard-core/store"
)

// providerSyncTimeout is the maximum time allowed for an outbound provider
// HTTP call during model sync. Configurable via PROVIDER_SYNC_TIMEOUT env
// var (seconds). Default 30s.
var providerSyncTimeout = func() time.Duration {
	if v := os.Getenv("PROVIDER_SYNC_TIMEOUT"); v != "" {
		if n, err := strconv.Atoi(v); err == nil && n > 0 {
			return time.Duration(n) * time.Second
		}
	}
	return 30 * time.Second
}()

// ErrUpstreamTimeout is returned when a provider API call exceeds the
// configured timeout. Handlers use this to distinguish upstream timeouts
// from other errors and return 504 instead of 502.
var ErrUpstreamTimeout = errors.New("upstream timeout")

// syncResult holds the outcome of a model sync operation.
type syncResult struct {
	Added   int `json:"added"`
	Updated int `json:"updated"`
	Total   int `json:"total"`
}

// syncProviderModels fetches models from a provider's API and syncs them into the catalog.
// New models default to 'disabled' visibility (admin must explicitly enable for global providers).
// apiKeyPlain is the decrypted API key — callers are responsible for decryption.
func syncProviderModels(ctx context.Context, stores store.Stores, cfg *models.ProviderConfig, apiKeyPlain string) (syncResult, error) {
	prov, err := providers.Get(cfg.Provider)
	if err != nil {
		return syncResult{}, fmt.Errorf("unknown provider type: %s", cfg.Provider)
	}

	proxyURL := ""
	if cfg.ProxyURL != nil {
		proxyURL = *cfg.ProxyURL
	}
	provCfg := providers.ProviderConfig{
		Endpoint:  cfg.Endpoint,
		APIKey:    apiKeyPlain,
		ProxyMode: cfg.ProxyMode,
		ProxyURL:  proxyURL,
	}

	if cfg.Headers != nil {
		customHeaders := make(map[string]string)
		for k, v := range cfg.Headers {
			if s, ok := v.(string); ok {
				customHeaders[k] = s
			}
		}
		provCfg.CustomHeaders = customHeaders
	}

	// Parse config for any extra settings the provider needs
	if cfg.Config != nil {
		raw, _ := json.Marshal(cfg.Config)
		var extra map[string]string
		if json.Unmarshal(raw, &extra) == nil {
			if provCfg.CustomHeaders == nil {
				provCfg.CustomHeaders = make(map[string]string)
			}
			for k, v := range extra {
				provCfg.CustomHeaders[k] = v
			}
		}
	}

	// Wrap context with timeout for the outbound provider call.
	// All provider ListModels implementations use http.NewRequestWithContext,
	// so the deadline propagates to the HTTP transport automatically.
	syncCtx, cancel := context.WithTimeout(ctx, providerSyncTimeout)
	defer cancel()

	provModels, err := prov.ListModels(syncCtx, provCfg)
	if err != nil {
		if errors.Is(err, context.DeadlineExceeded) {
			return syncResult{}, fmt.Errorf("%w: %s", ErrUpstreamTimeout, cfg.Provider)
		}
		return syncResult{}, err
	}

	syncEntries := make([]store.CatalogSyncEntry, len(provModels))
	for i, m := range provModels {
		syncEntries[i] = store.CatalogSyncEntry{
			ModelID:      m.ID,
			DisplayName:  m.Name,
			ModelType:    m.Type,
			Capabilities: m.Capabilities,
			Pricing:      m.Pricing,
		}
	}

	added, updated, err := stores.Catalog.UpsertFromSync(ctx, cfg.ID, syncEntries)
	if err != nil {
		return syncResult{}, fmt.Errorf("failed to sync: %w", err)
	}

	// Sync pricing from provider catalog (won't overwrite manual admin overrides)
	if stores.Pricing != nil {
		for _, m := range provModels {
			if m.Pricing != nil {
				if err := stores.Pricing.UpsertFromCatalog(ctx, cfg.ID, m.ID, m.Pricing); err != nil {
					log.Printf("warn: pricing sync for %s/%s failed: %v", cfg.ID, m.ID, err)
				}
			}
		}
	}

	return syncResult{Added: added, Updated: updated, Total: len(provModels)}, nil
}

// syncAndEnableProviderModels fetches models and auto-enables them all.
// Used for personal (BYOK) providers — the user explicitly added this provider to use it.
func syncAndEnableProviderModels(ctx context.Context, stores store.Stores, cfg *models.ProviderConfig, apiKeyPlain string) (syncResult, error) {
	result, err := syncProviderModels(ctx, stores, cfg, apiKeyPlain)
	if err != nil {
		return result, err
	}

	// Auto-enable: user added this key to USE it, not to stare at disabled models
	if result.Total > 0 {
		if err := stores.Catalog.BulkSetVisibility(ctx, cfg.ID, "enabled"); err != nil {
			log.Printf("warn: auto-enable models for provider %s failed: %v", cfg.ID, err)
		}
	}

	return result, nil
}