free-ide-proxy/internal/proxy/models.go

package proxy

import (
	"bytes"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"net/http"
	"sync"
	"time"
)

// ModelInfo describes a single model exposed by a backend.
type ModelInfo struct {
	ID              string `json:"id"`
	Name            string `json:"name,omitempty"`
	OwnedBy         string `json:"owned_by"`
	ContextWindow   int    `json:"context_window"`
	MaxOutputTokens int    `json:"max_output_tokens"`
	Description     string `json:"description"`
	SupportsVision  bool   `json:"supports_vision"`
}

// healthCheckInterval is how often caged models get re-probed.
const healthCheckInterval = 30 * time.Minute

// healthPingTimeout caps each health probe. Set high enough for slow reasoning
// models (nvidia/nemotron can take ~80s to first token).
const healthPingTimeout = 90 * time.Second

// modelHealth tracks the rate-limit state of a single model.
type modelHealth struct {
	status    string // "active" | "rate_limited"
	lastCheck time.Time
}

// Proxy routes OpenAI-format requests to the right Backend and manages session IDs.
type Proxy struct {
	apiKey string

	backends       []Backend
	defaultBackend Backend

	sessions map[string]*session
	mu       sync.Mutex

	health   map[string]*modelHealth
	healthMu sync.RWMutex
}

type session struct {
	id        string
	expiresAt time.Time
}

// NewProxy creates a Proxy with the default set of free-model backends.
//
// Order matters: more specific backends come first so their model IDs win over
// the catch-all opencode aliases. The default backend is the fallback for
// unknown models (pass-through as-is).
func NewProxy(apiKey string) *Proxy {
	oc := OpenCodeBackend{}
	kl := KiloBackend{}
	p := &Proxy{
		apiKey:         apiKey,
		backends:       []Backend{kl, oc},
		defaultBackend: oc,
		sessions:       make(map[string]*session),
		health:         make(map[string]*modelHealth),
	}
	go p.healthLoop()
	return p
}

// APIKey returns the configured proxy API key (empty = no auth).
func (p *Proxy) APIKey() string { return p.apiKey }

// Models returns the aggregated free-model catalogue across all backends,
// excluding any model currently caged for rate limiting (HTTP 429).
func (p *Proxy) Models() []ModelInfo {
	var all []ModelInfo
	for _, b := range p.backends {
		for _, m := range b.Models() {
			if !p.isRateLimited(m.ID) {
				all = append(all, m)
			}
		}
	}
	return all
}

// BackendFor returns the backend that recognises the model, or the default
// backend (pass-through) when no backend claims it.
func (p *Proxy) BackendFor(model string) Backend {
	b, _, _ := p.resolveModel(model)
	return b
}

// resolveModel walks the backends in order and returns the first match.
// When nothing matches it returns the default backend with the model unchanged
// and ok=false, so callers can forward it as-is (pass-through).
func (p *Proxy) resolveModel(model string) (backend Backend, resolved string, ok bool) {
	for _, b := range p.backends {
		if r, matched := b.Resolve(model); matched {
			return b, r, true
		}
	}
	return p.defaultBackend, model, false
}

// ResolveModel maps a requested model name to its upstream ID (backend-aware).
// Kept for backwards compatibility with existing call sites.
func (p *Proxy) ResolveModel(requested string) (string, bool) {
	_, r, ok := p.resolveModel(requested)
	return r, ok
}

// SessionID returns a session ID for the given user key, rotating every 30 min.
func (p *Proxy) SessionID(userKey string) string {
	p.mu.Lock()
	defer p.mu.Unlock()

	if s, exists := p.sessions[userKey]; exists && time.Now().Before(s.expiresAt) {
		return s.id
	}

	p.sessions[userKey] = &session{
		id:        "ses_" + randomHex(12),
		expiresAt: time.Now().Add(30 * time.Minute),
	}
	return p.sessions[userKey].id
}

// RequestID generates a new unique request ID.
func RequestID() string {
	return "msg_" + randomHex(12)
}

func randomHex(n int) string {
	b := make([]byte, n)
	rand.Read(b)
	return hex.EncodeToString(b)
}

// ---------- Model health / rate-limit tracking ----------

// markRateLimited cages a model so it disappears from /v1/models until it recovers.
// Only logs the transition to avoid spam when many 429s arrive in a row.
func (p *Proxy) markRateLimited(modelID string) {
	p.healthMu.Lock()
	defer p.healthMu.Unlock()
	h := p.health[modelID]
	if h == nil {
		h = &modelHealth{}
		p.health[modelID] = h
	}
	if h.status != "rate_limited" {
		fmt.Printf("[health] caging model '%s' (rate limited)\n", modelID)
	}
	h.status = "rate_limited"
	h.lastCheck = time.Now()
}

// markActive restores a caged model so it reappears in /v1/models.
func (p *Proxy) markActive(modelID string) {
	p.healthMu.Lock()
	defer p.healthMu.Unlock()
	h := p.health[modelID]
	if h == nil {
		h = &modelHealth{}
		p.health[modelID] = h
	}
	h.status = "active"
	h.lastCheck = time.Now()
}

// isRateLimited reports whether a model is currently caged for rate limiting.
func (p *Proxy) isRateLimited(modelID string) bool {
	p.healthMu.RLock()
	defer p.healthMu.RUnlock()
	h, ok := p.health[modelID]
	return ok && h.status == "rate_limited"
}

// healthLoop runs an initial full scan on startup, then re-probes caged models
// every healthCheckInterval (30 min). Released models reappear in /v1/models.
func (p *Proxy) healthLoop() {
	p.scanAllModels()
	ticker := time.NewTicker(healthCheckInterval)
	defer ticker.Stop()
	for range ticker.C {
		p.recheckCagedModels()
	}
}

// scanAllModels probes every known model once. Models returning non-200 (429)
// are caged; the rest stay active. Used for the startup scan.
func (p *Proxy) scanAllModels() {
	var ids []string
	for _, b := range p.backends {
		for _, m := range b.Models() {
			ids = append(ids, m.ID)
		}
	}
	fmt.Printf("[health] initial scan: probing %d models\n", len(ids))
	for _, id := range ids {
		ok := p.pingModel(id)
		p.healthMu.Lock()
		h := p.health[id]
		if h == nil {
			h = &modelHealth{}
			p.health[id] = h
		}
		h.lastCheck = time.Now()
		if ok {
			h.status = "active"
		} else {
			h.status = "rate_limited"
			fmt.Printf("[health] caging model '%s' (rate limited on initial scan)\n", id)
		}
		p.healthMu.Unlock()
	}
	active := 0
	p.healthMu.RLock()
	for _, h := range p.health {
		if h.status == "active" {
			active++
		}
	}
	p.healthMu.RUnlock()
	fmt.Printf("[health] initial scan complete: %d active, %d caged\n", active, len(ids)-active)
}

// recheckCagedModels re-probes only the models currently caged. A 200 releases
// them; a 429 or timeout keeps them caged for another interval.
func (p *Proxy) recheckCagedModels() {
	p.healthMu.RLock()
	var caged []string
	for id, h := range p.health {
		if h.status == "rate_limited" {
			caged = append(caged, id)
		}
	}
	p.healthMu.RUnlock()

	if len(caged) == 0 {
		return
	}
	fmt.Printf("[health] re-checking %d caged models\n", len(caged))
	for _, id := range caged {
		ok := p.pingModel(id)
		p.healthMu.Lock()
		h := p.health[id]
		if h == nil {
			h = &modelHealth{}
			p.health[id] = h
		}
		h.lastCheck = time.Now()
		if ok {
			h.status = "active"
			fmt.Printf("[health] releasing model '%s' (recovered)\n", id)
		} else {
			fmt.Printf("[health] model '%s' still rate limited, retry in 30m\n", id)
		}
		p.healthMu.Unlock()
	}
}

// pingModel sends a minimal chat completion to the model's backend and reports
// whether the upstream answered HTTP 200. Timeouts and non-200 count as
// unhealthy so the model stays caged (conservative).
func (p *Proxy) pingModel(modelID string) bool {
	backend, resolved, known := p.resolveModel(modelID)
	if !known {
		resolved = modelID
	}

	body := fmt.Sprintf(`{"model":%q,"messages":[{"role":"user","content":"OK"}],"max_tokens":5,"temperature":0}`, resolved)
	req, err := http.NewRequest("POST", backend.ChatURL(), bytes.NewReader([]byte(body)))
	if err != nil {
		return false
	}
	for k, v := range backend.Headers(RequestID(), "ses_health") {
		req.Header.Set(k, v)
	}

	client := &http.Client{Timeout: healthPingTimeout}
	resp, err := client.Do(req)
	if err != nil {
		return false
	}
	resp.Body.Close()
	return resp.StatusCode == 200
}