ModelRouter/internal/provider/provider.go

// Package provider binds a configured source + Lua adapter and performs the
// HTTP call / stream / image generation against the upstream LLM, with
// per-source concurrency limiting and availability backoff.
package provider

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"llmsproxy/internal/config"
	"llmsproxy/internal/lua"
	"llmsproxy/internal/types"
)

// ---- per-(source,model) scheduling state ----

// ModelState is the scheduling state of one (source, model) pair: a soft
// preference score used to order candidates within a priority tier, a
// consecutive-failure counter driving exponential cooldown, and a hard
// cooldown deadline. All fields are atomic and cooldown expiry is evaluated
// lazily (no timers, no goroutines). A model is never permanently blacklisted:
// cooldown always expires and any success resets the state, so a fixed
// upstream recovers on its own.
type ModelState struct {
	pref          atomic.Int64 // +1 per success / -5 per failure, clamped
	failCount     atomic.Int64
	cooldownUntil atomic.Int64 // unix seconds; 0 = schedulable
}

const (
	prefFailStep = 5
	prefMin      = -20
	prefMax      = 20
	backoffBase  = 5 * time.Second
	backoffCap   = 30 * time.Minute
	// failures at/above this count back off at the capped duration
	backoffCapN = 10
)

// rateLimitCooldown is the fixed cooldown applied on HTTP 429. A rate-limit
// rejection means "too fast", not "broken": unlike transport/5xx failures it
// must NOT escalate exponentially to backoffCap, or a quota-rich source gets
// locked out for up to 30 minutes while its allowance is untouched. A short
// fixed window lets the slot re-enter the rotation quickly.
const rateLimitCooldown = 30 * time.Second

func clampPref(a *atomic.Int64, lo, hi int64) {
	for {
		cur := a.Load()
		if cur < lo {
			if a.CompareAndSwap(cur, lo) {
				return
			}
			continue
		}
		if cur > hi {
			if a.CompareAndSwap(cur, hi) {
				return
			}
			continue
		}
		return
	}
}

// Available reports whether the model may be scheduled right now (cooldown
// expired or not yet set).
func (s *ModelState) Available() bool {
	return s.cooldownUntil.Load() <= time.Now().Unix()
}

// RecordFailure counts one consecutive failure and schedules exponential
// cooldown (5s, 10s, 20s … capped at 30min). auth marks 401/403 credential
// failures: it jumps straight to the capped cooldown and doubles the
// preference penalty, but the model still recovers when the cooldown expires.
func (s *ModelState) RecordFailure(auth bool) {
	n := s.failCount.Add(1)
	if auth && n < backoffCapN {
		n = backoffCapN
		s.failCount.Store(n) // persist the cap so FailCount() reports it too
	}
	var cd time.Duration
	if n >= backoffCapN {
		cd = backoffCap
	} else {
		cd = backoffBase * time.Duration(1<<(n-1))
		if cd > backoffCap {
			cd = backoffCap
		}
	}
	s.cooldownUntil.Store(time.Now().Add(cd).Unix())
	pen := int64(prefFailStep)
	if auth {
		pen *= 2
	}
	s.pref.Add(-pen)
	clampPref(&s.pref, prefMin, prefMax)
}

// RecordRateLimit records an HTTP 429: a short fixed cooldown instead of the
// exponential schedule, so a merely-throttled source recovers in seconds and
// its remaining quota stays usable. The preference penalty still applies so
// other slots are preferred while cooling.
func (s *ModelState) RecordRateLimit() {
	s.failCount.Add(1)
	s.cooldownUntil.Store(time.Now().Add(rateLimitCooldown).Unix())
	s.pref.Add(-prefFailStep)
	clampPref(&s.pref, prefMin, prefMax)
}

// RecordSuccess resets the failure counter and cooldown and bumps the
// preference score by one.
func (s *ModelState) RecordSuccess() {
	s.failCount.Store(0)
	s.cooldownUntil.Store(0)
	s.pref.Add(1)
	clampPref(&s.pref, prefMin, prefMax)
}

func (s *ModelState) reset() {
	s.failCount.Store(0)
	s.cooldownUntil.Store(0)
	s.pref.Store(0)
}

// Pref is the current preference score (higher = preferred).
func (s *ModelState) Pref() int64 { return s.pref.Load() }

// FailCount is the number of consecutive failures.
func (s *ModelState) FailCount() int64 { return s.failCount.Load() }

// CooldownUntil is the unix timestamp until which the model is cooled; 0 when
// schedulable.
func (s *ModelState) CooldownUntil() int64 { return s.cooldownUntil.Load() }

// Provider is a single configured upstream LLM source.
type Provider struct {
	cfg     config.Source
	vm      *lua.VM
	adapter string
	// client bounds a whole non-streaming request (dial+read body). stream
	// is used for SSE: no overall timeout (a long stream must not be cut),
	// only the transport's ResponseHeaderTimeout bounds time-to-first-byte.
	client *http.Client
	stream *http.Client

	mu        sync.Mutex
	sem       chan struct{}
	states    map[string]*ModelState // key = model id
	// proactive rate limiting: nextOK is the earliest unix-nano time the next
	// request may leave; throttle() spaces requests 60s/RPM apart. Zero when
	// cfg.RPM <= 0 (unlimited).
	rateMu   sync.Mutex
	nextOK   int64
	rpmGap   time.Duration
	lastProbe struct {
		ok  bool
		err string
		at  int64
	}
}

func New(cfg config.Source, vm *lua.VM) *Provider {
	if cfg.Timeout <= 0 {
		cfg.Timeout = config.DefaultSourceTimeout
	}
	if cfg.MaxConcurrent <= 0 {
		cfg.MaxConcurrent = config.DefaultSourceConcurrency
	}
	// Shared transport: ResponseHeaderTimeout bounds how long we wait for the
	// first response byte (applies to both paths); the stream client has no
	// client-level Timeout so the SSE body can run past the header timeout.
	tr := &http.Transport{
		Proxy:                 http.ProxyFromEnvironment,
		DialContext:           (&net.Dialer{Timeout: 30 * time.Second, KeepAlive: 30 * time.Second}).DialContext,
		ForceAttemptHTTP2:     true,
		MaxIdleConns:          100,
		IdleConnTimeout:       90 * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
		ResponseHeaderTimeout: cfg.Timeout,
	}
	p := &Provider{
		cfg:     cfg,
		vm:      vm,
		adapter: cfg.Adapter,
		client:  &http.Client{Timeout: cfg.Timeout, Transport: tr},
		stream:  &http.Client{Transport: tr},
		sem:     make(chan struct{}, cfg.MaxConcurrent),
		states:  map[string]*ModelState{},
	}
	if cfg.MaxConcurrent <= 0 {
		p.sem = nil
	}
	if cfg.RPM > 0 {
		p.rpmGap = time.Minute / time.Duration(cfg.RPM)
	}
	for _, m := range cfg.Models {
		p.states[m.ID] = &ModelState{}
	}
	return p
}

func (p *Provider) Name() string           { return p.cfg.Name }
func (p *Provider) Adapter() string        { return p.cfg.Adapter }
func (p *Provider) MaxConcurrent() int     { return p.cfg.MaxConcurrent }
func (p *Provider) Config() *config.Source { return &p.cfg }

// Models returns the model ids exposed by this source.
func (p *Provider) Models() []string {
	out := make([]string, 0, len(p.cfg.Models))
	for _, m := range p.cfg.Models {
		out = append(out, m.ID)
	}
	return out
}

// ModelByID returns the model definition if owned by this source.
func (p *Provider) ModelByID(id string) *config.Model {
	for i := range p.cfg.Models {
		if p.cfg.Models[i].ID == id {
			return &p.cfg.Models[i]
		}
	}
	return nil
}

// ModelIDFold returns the configured model id matching id case-insensitively
// ("" when no model matches). AUTO-chain slot models are normalized through
// this so cooldown state, quota windows and the upstream model id all refer
// to the exact configured spelling.
func (p *Provider) ModelIDFold(id string) string {
	for _, m := range p.cfg.Models {
		if strings.EqualFold(m.ID, id) {
			return m.ID
		}
	}
	return ""
}

// ModelFor resolves the model name this provider should send upstream.
// If the requested model is not owned by this provider (e.g. an AUTO chain
// fallback), it returns this provider's highest-priority chat model instead.
func (p *Provider) ModelFor(reqModel string) string {
	if reqModel == "" || isAutoID(reqModel) {
		return p.bestChatModel()
	}
	if p.ModelByID(reqModel) != nil {
		return reqModel
	}
	return p.bestChatModel()
}

// bestChatModel returns the highest-priority chat-kind model of this source.
func (p *Provider) bestChatModel() string {
	bestID, bestPrio := "", -1
	for _, m := range p.cfg.Models {
		if m.Kind != "" && m.Kind != "chat" {
			continue
		}
		if m.Priority > bestPrio {
			bestPrio = m.Priority
			bestID = m.ID
		}
	}
	if bestID == "" && len(p.cfg.Models) > 0 {
		bestID = p.cfg.Models[0].ID
	}
	return bestID
}

// bestImageModel returns the highest-priority image-kind model of this source
// (fallback: first configured model). Used for AUTO image generation so a
// mixed source never sends a chat model to /v1/images/generations.
func (p *Provider) bestImageModel() string {
	bestID, bestPrio := "", -1
	for _, m := range p.cfg.Models {
		if m.Kind != "" && m.Kind != "image" {
			continue
		}
		if m.Priority > bestPrio {
			bestPrio = m.Priority
			bestID = m.ID
		}
	}
	if bestID == "" && len(p.cfg.Models) > 0 {
		bestID = p.cfg.Models[0].ID
	}
	return bestID
}

// IsAutoID reports whether s is an AUTO routing placeholder.
func isAutoID(s string) bool {
	s = strings.TrimSpace(s)
	return s == "" || strings.EqualFold(s, "AUTO")
}

// Endpoint resolves the upstream chat path.
func (p *Provider) Endpoint() string {
	if p.cfg.Endpoint != "" {
		return p.cfg.Endpoint
	}
	if ep := p.vm.Endpoint(p.adapter); ep != "" {
		return ep
	}
	return "/chat/completions"
}

// ImageEndpoint resolves the upstream image-generation path.
func (p *Provider) ImageEndpoint() string {
	if p.cfg.ImageEndpoint != "" {
		return p.cfg.ImageEndpoint
	}
	if ep := p.vm.Endpoint(p.adapter + "_image"); ep != "" {
		return ep
	}
	return "/v1/images/generations"
}

func (p *Provider) URL() string {
	return strings.TrimRight(p.cfg.BaseURL, "/") + p.Endpoint()
}

func (p *Provider) ImageURL() string {
	return strings.TrimRight(p.cfg.BaseURL, "/") + p.ImageEndpoint()
}

// ---- availability ----

// ErrBusy is returned when every concurrency slot of a source is in use. It
// is a soft signal: schedulers skip a busy candidate without recording any
// failure (busy is not a failure) and gateways map it to HTTP 429. It aliases
// types.ErrBusy so the scheduler layer (which must not depend on this package)
// can detect busy via the shared sentinel.
var ErrBusy = types.ErrBusy

// Available reports whether the source is schedulable at source level: at
// least one of its models is not cooling down. Per-model scheduling decisions
// must use ModelAvailable instead.
func (p *Provider) Available() bool {
	p.mu.Lock()
	defer p.mu.Unlock()
	for _, s := range p.states {
		if s.Available() {
			return true
		}
	}
	return false
}

// ModelAvailable reports whether the exact model is schedulable right now
// (its cooldown expired). An unknown model id is treated as available.
func (p *Provider) ModelAvailable(model string) bool {
	p.mu.Lock()
	defer p.mu.Unlock()
	if s, ok := p.states[model]; ok {
		return s.Available()
	}
	return true
}

// Pref returns the adaptive preference score of a model (higher = preferred
// within a priority tier). Used by the AUTO chain to order same-tier slots.
func (p *Provider) Pref(model string) int64 {
	return p.state(model).Pref()
}

// ResetModelCooldown clears the cooldown and failure counter of a single
// model while preserving its preference score. Called after AUTO-chain edits
// so edited slots become schedulable immediately (a stored preference for a
// reliably good model is kept).
func (p *Provider) ResetModelCooldown(model string) {
	s := p.state(model)
	s.failCount.Store(0)
	s.cooldownUntil.Store(0)
}

// ModelHealthInfo exposes the per-model scheduling state for the web UI.
// An unknown model id reports zeros.
func (p *Provider) ModelHealthInfo(model string) (pref, failCount, cooldownUntil int64) {
	s := p.state(model)
	return s.Pref(), s.FailCount(), s.CooldownUntil()
}

// Probe performs a lightweight reachability + auth check against the source.
// It first tries GET <base>/models (fast, ~1s for OpenAI-compatible upstreams)
// and only falls back to a 1-token chat call when that endpoint is
// unavailable. It does NOT touch the health/backoff state so probing never
// disables a source.
func (p *Provider) Probe(ctx context.Context) (bool, string) {
	ok, msg := p.probeModels(ctx)
	if !ok && msg == "" {
		ok, msg = p.probeChat(ctx)
	}
	p.mu.Lock()
	p.lastProbe.ok = ok
	p.lastProbe.err = msg
	p.lastProbe.at = time.Now().Unix()
	p.mu.Unlock()
	return ok, msg
}

// probeModels GETs <base>/models. Returns (true,…) when reachable, (false,
// errortext) on an auth/permanent failure, and (false,"") when the endpoint
// simply isn't available so the caller can fall back to a chat probe.
func (p *Provider) probeModels(ctx context.Context) (bool, string) {
	u := strings.TrimRight(p.cfg.BaseURL, "/") + "/models"
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, u, nil)
	if err != nil {
		return false, ""
	}
	if hdrs, herr := p.buildHeaders("{}", u); herr == nil {
		req.Header = hdrs
	}
	resp, err := p.client.Do(req)
	if err != nil {
		return false, ""
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(resp.Body)
	switch {
	case resp.StatusCode == 200:
		return true, ""
	case resp.StatusCode == 404 || resp.StatusCode == 405:
		return false, ""
	default:
		return false, p.apiErrReason(resp.StatusCode, string(raw))
	}
}

// probeChat sends a minimal single-token chat request to the chat endpoint.
func (p *Provider) probeChat(ctx context.Context) (bool, string) {
	ok := false
	msg := ""
	model := p.bestChatModel()
	if pm := p.ModelByID(model); pm != nil && pm.Kind == "image" {
		model = ""
	}
	if model == "" {
		if ms := p.Models(); len(ms) > 0 {
			model = ms[0]
		}
	}
	if model == "" {
		msg = "no chat model configured"
	} else {
		probe := map[string]interface{}{
			"model":      model,
			"messages":   []map[string]interface{}{{"role": "user", "content": "hi"}},
			"max_tokens": 1,
		}
		body, err := json.Marshal(probe)
		if err == nil {
			var hdr http.Header
			if hdrs, herr := p.buildHeaders(string(body), p.URL()); herr == nil {
				hdr = hdrs
			}
			raw, status, derr := p.do(ctx, p.URL(), string(body), hdr)
			if derr != nil {
				msg = derr.Error()
			} else if status == 200 {
				ok = true
			} else {
				msg = p.apiErrReason(status, raw)
			}
		} else {
			msg = err.Error()
		}
	}
	return ok, msg
}

func (p *Provider) LastProbe() (bool, string, int64) {
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.lastProbe.ok, p.lastProbe.err, p.lastProbe.at
}

// state returns the ModelState for a model id, creating it on first use so
// dynamically requested models are still tracked. The registry keeps states
// alive across provider rebuilds only for configured models; a lazily created
// state simply lives for the provider's lifetime.
func (p *Provider) state(model string) *ModelState {
	p.mu.Lock()
	defer p.mu.Unlock()
	s, ok := p.states[model]
	if !ok {
		s = &ModelState{}
		p.states[model] = s
	}
	return s
}

// RecordFailure records a failed downstream attempt on model: consecutive
// failure count +1 and exponential cooldown (5s·2^n, capped at 30min).
// code 401/403 is treated as a credential problem: the cooldown jumps to the
// cap and the preference penalty doubles, but the model still recovers when
// the cooldown expires (no permanent blacklist). code 0 = transport failure.
func (p *Provider) RecordFailure(model string, code int) {
	p.state(model).RecordFailure(code == 401 || code == 403)
}

// RecordSuccess resets the model's failure counter / cooldown and bumps its
// preference by one.
func (p *Provider) RecordSuccess(model string) {
	p.state(model).RecordSuccess()
}

// ReportStatus records an upstream HTTP status for the given model and drives
// the (source, model) backoff state. 401/403 → capped self-healing cooldown
// with doubled penalty; 5xx → normal exponential backoff; 429 → short fixed
// cooldown (rateLimitCooldown): a rate rejection means "too fast", not
// "broken", so the slot must re-enter rotation quickly instead of escalating
// to a 30-minute lockout that wastes remaining quota.
func (p *Provider) ReportStatus(model string, code int) {
	if code == 401 || code == 403 {
		p.RecordFailure(model, code)
		return
	}
	if code == 429 {
		p.state(model).RecordRateLimit()
		return
	}
	if code >= 500 {
		p.RecordFailure(model, code)
	}
}

// ResetHealth resets the scheduling state of every model of this source
// (cooldown and preference to zero), so the source becomes fully schedulable
// again. Called after AUTO-chain edits and from the admin UI.
func (p *Provider) ResetHealth() {
	p.mu.Lock()
	defer p.mu.Unlock()
	for _, s := range p.states {
		s.reset()
	}
}

// HealthInfo exposes the source-level backoff state for the status page: the
// highest failure count and the latest cooldown deadline across all models of
// this source. permanent is always false — the permanent-blacklist semantics
// were removed; every cooldown expires on its own.
func (p *Provider) HealthInfo() (failCount int, until time.Time, permanent bool) {
	p.mu.Lock()
	defer p.mu.Unlock()
	now := time.Now().Unix()
	for _, s := range p.states {
		if n := int(s.FailCount()); n > failCount {
			failCount = n
		}
		if t := s.CooldownUntil(); t > now && t > until.Unix() {
			until = time.Unix(t, 0)
		}
	}
	return failCount, until, false
}

// ---- concurrency limiting ----

// TryAcquire takes one concurrency slot without blocking: it returns nil when
// a slot is free and ErrBusy when the source is at capacity. A nil semaphore
// (MaxConcurrent <= 0) means unlimited and always succeeds. TryAcquire is the
// single busy/idle signal for schedulers; a busy source is skipped, never
// penalized.
func (p *Provider) TryAcquire(ctx context.Context) error {
	if p.sem == nil {
		return nil
	}
	select {
	case p.sem <- struct{}{}:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	default:
		return ErrBusy
	}
}

// Acquire waits for a free concurrency slot (bounded by cfg.QueueTimeout),
// or context cancel. The HTTP call itself is not truncated. Direct requests
// historically queued here; the scheduler now prefers TryAcquire so a full
// source fails fast instead of blocking the whole chain.
func (p *Provider) Acquire(ctx context.Context) error {
	if p.sem == nil {
		return nil
	}
	var qCtx context.Context
	var cancel context.CancelFunc
	if p.cfg.QueueTimeout > 0 {
		qCtx, cancel = context.WithTimeout(ctx, p.cfg.QueueTimeout)
	} else {
		qCtx, cancel = context.WithCancel(ctx)
	}
	defer cancel()
	select {
	case p.sem <- struct{}{}:
		return nil
	case <-qCtx.Done():
		return qCtx.Err()
	}
}

func (p *Provider) Release() {
	if p.sem == nil {
		return
	}
	<-p.sem
}

// ---- proactive rate limiting ----

// Throttle blocks until this source's rate-limit window allows the next
// request, spacing outbound requests at least rpmGap apart (60s/RPM). It is a
// no-op when no rpm cap is configured. The wait happens AFTER the concurrency
// slot is taken, so a queued request still counts against max_concurrent —
// bounded by the caller's ctx (queue timeout / client disconnect).
func (p *Provider) Throttle(ctx context.Context) error {
	if p.rpmGap <= 0 {
		return nil
	}
	p.rateMu.Lock()
	now := time.Now().UnixNano()
	wait := p.nextOK - now
	if wait > 0 {
		p.nextOK += int64(p.rpmGap) // reserve the next window for the follower
		p.rateMu.Unlock()
		t := time.NewTimer(time.Duration(wait))
		defer t.Stop()
		select {
		case <-t.C:
			return nil
		case <-ctx.Done():
			return ctx.Err()
		}
	}
	p.nextOK = now + int64(p.rpmGap)
	p.rateMu.Unlock()
	return nil
}

// ---- request construction ----

func (p *Provider) buildHeaders(body, url string) (http.Header, error) {
	meta := map[string]interface{}{
		"url":       url,
		"method":    http.MethodPost,
		"body":      body,
		"api_key":   p.cfg.APIKey,
		"timestamp": types.Now(),
		"source": map[string]interface{}{
			"name": p.cfg.Name,
			"meta": p.cfg.Meta,
		},
	}
	hdrs, err := p.vm.BuildHeaders(p.adapter, meta)
	if err != nil {
		return nil, err
	}
	h := http.Header{}
	h.Set("Content-Type", "application/json")
	for k, v := range p.cfg.Headers {
		h.Set(k, v)
	}
	for k, v := range hdrs {
		if _, ok := p.cfg.Headers[k]; !ok {
			h.Set(k, v)
		}
	}
	if h.Get("Authorization") == "" && p.cfg.APIKey != "" {
		h.Set("Authorization", "Bearer "+p.cfg.APIKey)
	}
	return h, nil
}

// ---- chat ----

// Chat performs a non-streaming round trip and returns the unified response.
// It fails fast with ErrBusy when the source is at capacity; success/failure
// is recorded against the resolved (source, model) scheduling state.
func (p *Provider) Chat(ctx context.Context, req *types.ChatRequest) (*types.UnifiedResponse, error) {
	if err := p.TryAcquire(ctx); err != nil {
		return nil, err
	}
	defer p.Release()
	if err := p.Throttle(ctx); err != nil {
		return nil, err
	}
	model := p.ModelFor(req.Model)

	body, err := marshalTransform(p.vm, p.adapter, "transform_request", req)
	if err != nil {
		return nil, err
	}
	hdrs, err := p.buildHeaders(body, p.URL())
	if err != nil {
		return nil, err
	}
	raw, status, err := p.do(ctx, p.URL(), body, hdrs)
	if err != nil {
		// a client disconnect or cancelled context is neither a success nor
		// a failure for scheduling purposes — only upstream errors count
		if ctx.Err() == nil {
			p.RecordFailure(model, 0)
		}
		return nil, err
	}
	if status != 200 {
		p.ReportStatus(model, status)
		return nil, fmt.Errorf("%s", p.apiErrReason(status, raw))
	}
	unified, err := p.vm.Transform(p.adapter, "transform_response", raw)
	if err != nil {
		// adapter produced unusable output: a real failure the slot must
		// back off from, otherwise a broken adapter source is retried at
		// full latency forever
		if ctx.Err() == nil {
			p.RecordFailure(model, 0)
		}
		return nil, err
	}
	var out types.UnifiedResponse
	if err := json.Unmarshal([]byte(unified), &out); err != nil {
		if ctx.Err() == nil {
			p.RecordFailure(model, 0)
		}
		return nil, fmt.Errorf("unmarshal unified response: %w (body: %s)", err, unified)
	}
	p.RecordSuccess(model)
	return &out, nil
}

// ChatStream performs a streaming round trip, emitting unified chunks. It
// fails fast with ErrBusy when the source is at capacity. Only a failure
// before the first chunk (connect error or non-200 status) is recorded
// against the (source, model) state; afterwards the stream is pinned. A clean
// end ([DONE] or EOF without read errors, and no client disconnect) counts as
// success and resets the cooldown.
func (p *Provider) ChatStream(ctx context.Context, req *types.ChatRequest) (<-chan types.UnifiedChunk, error) {
	if err := p.TryAcquire(ctx); err != nil {
		return nil, err
	}
	if err := p.Throttle(ctx); err != nil {
		p.Release()
		return nil, err
	}
	model := p.ModelFor(req.Model)
	req.Stream = true
	body, err := marshalTransform(p.vm, p.adapter, "transform_request", req)
	if err != nil {
		p.Release()
		return nil, err
	}
	hdrs, err := p.buildHeaders(body, p.URL())
	if err != nil {
		p.Release()
		return nil, err
	}

	type respOrErr struct {
		resp *http.Response
		err  error
	}
	rc := make(chan respOrErr, 1)
	go func() {
		resp, err := p.doRawStream(ctx, p.URL(), body, hdrs)
		rc <- respOrErr{resp, err}
	}()

	inner := make(chan types.UnifiedChunk, 64)
	sel := <-rc
	if sel.err != nil {
		// client disconnect/cancel before the first byte: not a scheduling
		// failure (a healthy source must not be cooled by client cancellations)
		if ctx.Err() == nil {
			p.RecordFailure(model, 0)
		}
		p.Release()
		return nil, sel.err
	}
	if sel.resp.StatusCode != 200 {
		raw, _ := io.ReadAll(sel.resp.Body)
		sel.resp.Body.Close()
		p.ReportStatus(model, sel.resp.StatusCode)
		p.Release()
		return nil, fmt.Errorf("%s", p.apiErrReason(sel.resp.StatusCode, string(raw)))
	}
	go func() {
		defer p.Release()
		defer close(inner)
		defer sel.resp.Body.Close()
		scanner := bufio.NewScanner(sel.resp.Body)
		scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
		var chunks int
		var realChunks int
		var doneSeen bool
		var doneSent bool
		for scanner.Scan() {
			line := strings.TrimSpace(scanner.Text())
			if line == "" || !strings.HasPrefix(line, "data:") {
				continue
			}
			data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
			if data == "" {
				continue
			}
			if data == "[DONE]" {
				doneSeen = true
				// adapters that already emitted their terminating done chunk
				// (with the real finish reason) must not get a second,
				// reason-less done from [DONE] — it would override the true
				// finish_reason downstream.
				if !doneSent {
					select {
					case inner <- types.UnifiedChunk{Done: true}:
					case <-ctx.Done():
					}
				}
				continue
			}
			unified, err := p.vm.Transform(p.adapter, "transform_stream_chunk", data)
			if err != nil || unified == "" {
				continue
			}
			if unified == data {
				unified = standardSSEChunk(data)
				if unified == "" {
					continue
				}
			}
			var ck types.UnifiedChunk
			if err := json.Unmarshal([]byte(unified), &ck); err != nil {
				continue
			}
			if ck.Done {
				doneSent = true
			}
			if !errorOnlyChunk(ck) {
				realChunks++
			}
			chunks++
			select {
			case inner <- ck:
			case <-ctx.Done():
				return
			}
		}
		// The stream ended cleanly ([DONE] seen or EOF without an upstream
		// read error): record success so a previously cooled model can be
		// retried. A client disconnect or mid-stream read error is neither
		// success nor failure for scheduling purposes. A 200 that produced
		// zero chunks and no [DONE] is an empty stream, i.e. a failure
		// before the first chunk — record it so the slot can fall back.
		// A stream whose only payload was error-only done chunks (e.g.
		// finish_reason:"network_error") is likewise a failure, not a success.
		if ctx.Err() == nil && scanner.Err() == nil {
			if realChunks > 0 || (doneSeen && chunks == 0) {
				p.RecordSuccess(model)
			} else {
				p.RecordFailure(model, 0)
			}
		}
	}()

	// Hold back the first chunk to validate the stream actually carries
	// content: some upstreams answer HTTP 200 with a degenerate stream whose
	// only payload is an error finish reason (zen free pool sends
	// finish_reason:"network_error" with empty delta). Failing the candidate
	// here — before any byte reaches the gateway — lets the scheduler fall
	// through to the next source instead of serving the client an empty reply.
	type heldChunk struct {
		ck types.UnifiedChunk
		ok bool
	}
	var first heldChunk
	select {
	case ck, ok := <-inner:
		first = heldChunk{ck, ok}
	case <-ctx.Done():
		return nil, ctx.Err()
	}
	if !first.ok {
		return nil, fmt.Errorf("provider %s: empty stream", p.Name())
	}
	if errorOnlyChunk(first.ck) {
		go func() {
			for range inner {
			}
		}()
		return nil, fmt.Errorf("provider %s: upstream returned %q stream",
			p.Name(), first.ck.FinishReason)
	}
	out := make(chan types.UnifiedChunk, 64)
	go func() {
		defer close(out)
		out <- first.ck
		for ck := range inner {
			out <- ck
		}
	}()
	return out, nil
}

// errorOnlyChunk reports whether ck carries nothing but an upstream error
// signal: a done chunk with a non-standard finish reason and zero content,
// tool calls, reasoning text or usage. Standard OpenAI finish reasons are
// never classified as errors, so legitimate instant-empty completions
// (finish_reason:"stop", no output) still reach the client.
func errorOnlyChunk(ck types.UnifiedChunk) bool {
	if !ck.Done || ck.FinishReason == "" {
		return false
	}
	switch ck.FinishReason {
	case "stop", "length", "tool_calls", "function_call", "content_filter":
		return false
	}
	return ck.Content == "" && len(ck.ToolCalls) == 0 &&
		ck.ReasoningContent == "" && ck.Usage == nil
}

// apiErrReason builds the client-facing reason for a non-200 upstream
// response: the adapter's optional transform_error hook wins (per-source
// protocol knowledge lives in Lua), otherwise clients get a uniform
// "unknown error" while the raw body stays in the server log for debugging.
func (p *Provider) apiErrReason(status int, raw string) string {
	if reason, ok, err := p.vm.TransformError(p.adapter, status, raw); err == nil && ok {
		if trimmed := strings.TrimSpace(reason); trimmed != "" {
			return fmt.Sprintf("api error %d: %s", status, types.OneLine(trimmed, 200))
		}
	}
	log.Printf("[provider] unhandled upstream error body (adapter %q lacks transform_error): status=%d body=%.300s",
		p.adapter, status, types.OneLine(raw, 300))
	return fmt.Sprintf("api error %d: unknown error", status)
}

// Image generates images via /v1/images/generations. Same scheduling-state
// accounting as Chat: fail fast on busy, record per (source, model).
func (p *Provider) Image(ctx context.Context, req *types.ImageGenRequest) (*types.UnifiedResponse, error) {
	if err := p.TryAcquire(ctx); err != nil {
		return nil, err
	}
	defer p.Release()
	// AUTO/unknown model: resolve to this source's image model, never to the
	// best chat model (a mixed source would otherwise send a chat id to
	// /v1/images/generations). An explicitly pinned model is honored as-is.
	if isAutoID(req.Model) || p.ModelByID(req.Model) == nil {
		r := *req
		r.Model = p.bestImageModel()
		req = &r
	}
	model := req.Model

	b, _ := json.Marshal(req)
	transformed, err := p.vm.Transform(p.adapter+"_image", "transform_request", string(b))
	if err != nil {
		// fall back to passthrough adapter (openai-style)
		transformed = string(b)
	}
	hdrs, err := p.buildHeaders(transformed, p.ImageURL())
	if err != nil {
		return nil, err
	}
	raw, status, err := p.do(ctx, p.ImageURL(), transformed, hdrs)
	if err != nil {
		// client disconnect/cancel: not a scheduling failure
		if ctx.Err() == nil {
			p.RecordFailure(model, 0)
		}
		return nil, err
	}
	if status != 200 {
		p.ReportStatus(model, status)
		return nil, fmt.Errorf("%s", p.apiErrReason(status, raw))
	}
	var out types.UnifiedResponse
	// try adapter transform_response; if missing, parse standard openai image format
	unified, terr := p.vm.Transform(p.adapter+"_image", "transform_response", raw)
	if terr == nil && unified != raw {
		if err := json.Unmarshal([]byte(unified), &out); err == nil {
			out.Model = model // actual model served (AUTO resolved to bestImageModel)
			p.RecordSuccess(model)
			return &out, nil
		}
	}
	var img types.ImageGenResponse
	if err := json.Unmarshal([]byte(raw), &img); err != nil {
		if ctx.Err() == nil {
			p.RecordFailure(model, 0)
		}
		return nil, fmt.Errorf("unmarshal image response: %w", err)
	}
	out.ImageData = img.Data
	out.Model = model // actual model served (AUTO resolved to bestImageModel)
	p.RecordSuccess(model)
	return &out, nil
}

// ---- http helpers ----

func (p *Provider) do(ctx context.Context, url, body string, hdr http.Header) (string, int, error) {
	resp, err := p.doRaw(ctx, url, body, hdr)
	if err != nil {
		return "", 0, err
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(resp.Body)
	return string(raw), resp.StatusCode, nil
}

func (p *Provider) doRaw(ctx context.Context, url, body string, hdr http.Header) (*http.Response, error) {
	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader([]byte(body)))
	if err != nil {
		return nil, err
	}
	httpReq.Header = hdr
	return p.client.Do(httpReq)
}

// doRawStream is the streaming variant of doRaw: it uses the no-overall-timeout
// stream client so a long SSE body is not cut by client.Timeout. The transport
// still enforces ResponseHeaderTimeout on time-to-first-byte.
func (p *Provider) doRawStream(ctx context.Context, url, body string, hdr http.Header) (*http.Response, error) {
	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader([]byte(body)))
	if err != nil {
		return nil, err
	}
	httpReq.Header = hdr
	return p.stream.Do(httpReq)
}

func marshalTransform(vm *lua.VM, adapter, fn string, v interface{}) (string, error) {
	b, err := json.Marshal(v)
	if err != nil {
		return "", err
	}
	out, err := vm.Transform(adapter, fn, string(b))
	if err != nil {
		return "", err
	}
	return out, nil
}

func standardSSEChunk(data string) string {
	var raw struct {
		Choices []struct {
			Delta struct {
				Content string `json:"content"`
			} `json:"delta"`
			FinishReason *string `json:"finish_reason"`
		} `json:"choices"`
		UpstreamUsage struct {
			PromptTokens        int `json:"prompt_tokens"`
			CompletionTokens    int `json:"completion_tokens"`
			TotalTokens         int `json:"total_tokens"`
			Prompt              int `json:"prompt"`
			Completion          int `json:"completion"`
			Total               int `json:"total"`
			PromptCacheHit      int `json:"prompt_cache_hit_tokens"`
			PromptCacheMiss     int `json:"prompt_cache_miss_tokens"`
			PromptTokensDetails *struct {
				CachedTokens int `json:"cached_tokens"`
			} `json:"prompt_tokens_details"`
		} `json:"usage"`
	}
	if err := json.Unmarshal([]byte(data), &raw); err != nil {
		return ""
	}
	if len(raw.Choices) == 0 && raw.UpstreamUsage.Total == 0 && raw.UpstreamUsage.TotalTokens == 0 {
		return ""
	}
	var usage *types.TokenUsage
	pu := raw.UpstreamUsage
	if pu.Total > 0 || pu.TotalTokens > 0 {
		usage = &types.TokenUsage{
			Prompt:          pickFirst(pu.PromptTokens, pu.Prompt),
			Completion:      pickFirst(pu.CompletionTokens, pu.Completion),
			Total:           pickFirst(pu.TotalTokens, pu.Total),
			PromptCacheHit:  pu.PromptCacheHit,
			PromptCacheMiss: pu.PromptCacheMiss,
		}
		if pu.PromptTokensDetails != nil && pu.PromptTokensDetails.CachedTokens > 0 {
			usage.PromptTokensDetails = &types.PromptTokensDetails{
				CachedTokens: pu.PromptTokensDetails.CachedTokens,
			}
		}
	}
	finish := ""
	done := false
	if len(raw.Choices) > 0 && raw.Choices[0].FinishReason != nil {
		// empty-string finish reasons (sensenova sends "" on every chunk)
		// are not a finish signal
		if fr := *raw.Choices[0].FinishReason; fr != "" {
			finish, done = fr, true
		}
	}
	out, _ := json.Marshal(types.UnifiedChunk{
		Content: func() string {
			if len(raw.Choices) > 0 {
				return raw.Choices[0].Delta.Content
			}
			return ""
		}(),
		Done:         done,
		FinishReason: finish,
		Usage:        usage,
	})
	return string(out)
}

// pickFirst returns a if non-zero, else b (for usage keys that may appear in
// either standard *_tokens or legacy short form).
func pickFirst(a, b int) int {
	if a != 0 {
		return a
	}
	return b
}