ModelRouter/internal/gateway/chat.go

package gateway

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"regexp"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

	"llmsproxy/internal/config"
	"llmsproxy/internal/provider"
	"llmsproxy/internal/scheduler"
	"llmsproxy/internal/types"
)

// chatRequest mirrors the OpenAI chat completions request the gateway accepts.
type chatRequest struct {
	Model           string                 `json:"model"`
	Messages        []types.ChatMessage    `json:"messages"`
	Temperature     *float64               `json:"temperature,omitempty"`
	MaxTokens       int                    `json:"max_tokens,omitempty"`
	Stream          bool                   `json:"stream,omitempty"`
	Tools           []interface{}          `json:"tools,omitempty"`
	ToolChoice      interface{}            `json:"tool_choice,omitempty"`
	DisableThinking bool                   `json:"disable_thinking"`
	ExtraBody       map[string]interface{} `json:"extra_body,omitempty"`

	// PromptCacheKey 是 OpenAI 原生的缓存键；部分客户端用它携带会话标识。
	PromptCacheKey string `json:"prompt_cache_key,omitempty"`
}

// ChatCompletion is the non-streaming OpenAI response object.
type ChatCompletion struct {
	ID      string            `json:"id"`
	Object  string            `json:"object"`
	Created int64             `json:"created"`
	Model   string            `json:"model"`
	Choices []ChatChoice      `json:"choices"`
	Usage   *types.TokenUsage `json:"usage,omitempty"`
	// Cost is the upstream-reported charge for this request, passed through
	// verbatim (OpenCode reports it as a decimal string, always "0" on the
	// flat-rate Go subscription). Absent when the upstream reports nothing.
	Cost string `json:"cost,omitempty"`
}

type ChatChoice struct {
	Index        int         `json:"index"`
	Message      RespMessage `json:"message"`
	FinishReason string      `json:"finish_reason"`
}

type RespMessage struct {
	Role             string          `json:"role,omitempty"`
	Content          string          `json:"content"`
	ReasoningContent string          `json:"reasoning_content,omitempty"`
	ToolCalls        json.RawMessage `json:"tool_calls,omitempty"`
}

type ChatChunk struct {
	ID      string        `json:"id"`
	Object  string        `json:"object"`
	Created int64         `json:"created"`
	Model   string        `json:"model"`
	Choices []ChunkChoice `json:"choices"`
	// Usage is sent in the final chunk of a stream (empty choices) so
	// OpenAI-compatible clients can read token usage.
	Usage *types.TokenUsage `json:"usage,omitempty"`
	// Cost is the upstream-reported charge for this request, passed through
	// verbatim (OpenCode reports it as a decimal string). Emitted on the
	// terminal chunk, mirroring OpenCode's own {"choices":[],"cost":"0"}.
	Cost string `json:"cost,omitempty"`
}

type ChunkChoice struct {
	Index        int         `json:"index"`
	Delta        RespMessage `json:"delta"`
	FinishReason *string     `json:"finish_reason"`
}

var seq int64

func newID() string {
	n := atomic.AddInt64(&seq, 1)
	return fmt.Sprintf("chatcmpl-%d", n)
}

func isAuto(m string) bool {
	m = strings.TrimSpace(m)
	return m == "" || strings.EqualFold(m, "AUTO")
}

// resolveCands picks the ordered candidate providers for a requested model.
// toolCalling requests are anchored: they resolve to exactly one provider
// (highest-priority available) so a tool-call round never switches models.
func (g *Gateway) resolveCands(ctx context.Context, req *chatRequest) ([]*provider.Provider, string) {
	model := req.Model
	if model == "" {
		model = g.core.DefaultModel()
	}
	cands, effective := g.resolveByModel(model)
	cands = chatOnly(cands)
	allow := g.allowedModels(ctx)
	if allow != nil {
		cands = filterCandsByModels(cands, allow)
	}
	if !toolRequest(req) {
		return cands, effective
	}
	// tool-call request: pin to one provider (no AUTO fallback across models)
	if len(cands) == 0 {
		return nil, effective
	}
	first := cands[0]
	// use the prefix-stripped id (effective), never the raw "src:model" form:
	// ModelFor does exact matching and would fall back to the source's best
	// chat model for an unknown id
	eff := first.ModelFor(effective)
	if eff == "" {
		eff = firstModel(first)
	}
	return []*provider.Provider{first}, eff
}

// filterCandsByModels keeps only providers exposing at least one model of the
// scope (used for user keys with a restricted model scope). Scope entries with
// a Source pinned to a specific upstream narrow the candidates to that source
// for the matching model.
//
// An "AUTO" scope entry only allows the AUTO routing mode; it does NOT grant
// access to specific models.
func filterCandsByModels(cands []*provider.Provider, allow []config.ModelScope) []*provider.Provider {
	for _, m := range allow {
		if m.Model == "" {
			return cands
		}
	}
	allowed := make(map[string]bool, len(allow))
	byModelSrc := map[string]map[string]bool{}
	for _, m := range allow {
		allowed[m.Model] = true
		if m.Source != "" {
			if byModelSrc[m.Model] == nil {
				byModelSrc[m.Model] = map[string]bool{}
			}
			byModelSrc[m.Model][m.Source] = true
		}
	}
	out := make([]*provider.Provider, 0, len(cands))
	for _, p := range cands {
		for _, id := range p.Models() {
			if !allowed[id] {
				continue
			}
			if srcs := byModelSrc[id]; len(srcs) > 0 && !srcs[p.Name()] {
				continue
			}
			out = append(out, p)
			break
		}
	}
	return out
}

// intersectModels restricts a model list to the scope (preserving order).
// An "AUTO" scope entry only allows the AUTO routing mode; it does NOT grant
// every model.
func intersectModels(models []string, allow []config.ModelScope) []string {
	allowed := make(map[string]bool, len(allow))
	for _, m := range allow {
		if m.Model == "" {
			return models
		}
		if !strings.EqualFold(m.Model, "AUTO") {
			allowed[m.Model] = true
		}
	}
	out := make([]string, 0, len(models))
	seen := map[string]bool{}
	for _, m := range models {
		if allowed[m] && !seen[m] {
			seen[m] = true
			out = append(out, m)
		}
	}
	return out
}

// checkModelScope validates the effective model against the key's model scope
// and token quota. Returns an error message when rejected.
//
// A scope entry with model "AUTO" only allows requests where the effective
// model is AUTO (the routing mode). It does NOT grant access to specific model
// ids — that requires an explicit scope entry for the model.
func (g *Gateway) checkModelScope(ctx context.Context, model string) string {
	if q := g.checkQuota(ctx, model); q != nil {
		return q.msg
	}
	return ""
}

// quotaRejection is a quota verdict: the message plus whether the client
// should retry. A spent quota is a rate limit (429 + Retry-After), not a
// permission failure (403): a client that sees 403 gives up on the key, while
// one that sees 429 with a retry hint waits and resumes when the window rolls
// over.
type quotaRejection struct {
	msg   string
	retry int64 // seconds until the window resets; 0 = unknown
}

func (q *quotaRejection) Error() string { return q.msg }

// checkQuota is checkKeyScope for callers that need the retry hint. It
// separates the quota verdicts (429) from model-permission verdicts (403).
func (g *Gateway) checkQuota(ctx context.Context, model string) *quotaRejection {
	if q := g.checkKeyQuotaRetry(ctx); q != nil {
		return q
	}
	allow := g.allowedModels(ctx)
	if allow == nil {
		return nil
	}
	for _, sc := range allow {
		if sc.Model != model {
			continue
		}
		if sc.TokenQuota > 0 {
			used := g.scopeTokens(ctx, sc)
			if used >= sc.TokenQuota {
				return &quotaRejection{
					msg:   fmt.Sprintf("token quota exceeded for %q (%d/%d)", model, used, sc.TokenQuota),
					retry: AutoSecondsToReset(sc.Period, sc.Hours),
				}
			}
		}
		return nil
	}
	return &quotaRejection{msg: fmt.Sprintf("model %q is not allowed for this key", model)}
}

// checkKeyQuotaRetry enforces the key-wide caps and reports the remaining
// seconds of the reset window so the caller can answer with 429 + Retry-After.
// An admin key is never capped, and a key with no caps set is never rejected.
func (g *Gateway) checkKeyQuotaRetry(ctx context.Context) *quotaRejection {
	rec, ok := g.core.FindKey(reqKey(ctx))
	if !ok || rec.Role == "admin" {
		return nil
	}
	k := keyID(reqKey(ctx))
	win := AutoPeriodSeconds(rec.Period, rec.Hours)
	if rec.TokenQuota > 0 {
		if used := g.stats.KeyWindowTokens(k, win); used >= rec.TokenQuota {
			return &quotaRejection{
				msg:   fmt.Sprintf("key token quota exceeded (%d/%d%s)", used, rec.TokenQuota, quotaWindowSuffix(rec.Period, rec.Hours)),
				retry: AutoSecondsToReset(rec.Period, rec.Hours),
			}
		}
	}
	if rec.ReqQuota > 0 {
		if used := g.stats.KeyWindowReqs(k, win); used >= rec.ReqQuota {
			return &quotaRejection{
				msg:   fmt.Sprintf("key request quota exceeded (%d/%d%s)", used, rec.ReqQuota, quotaWindowSuffix(rec.Period, rec.Hours)),
				retry: AutoSecondsToReset(rec.Period, rec.Hours),
			}
		}
	}
	return nil
}

// quotaWindowSuffix describes a quota's reset window for an error message, so
// a rejected caller can tell a permanent block from one that clears in an hour.
func quotaWindowSuffix(period string, hours int64) string {
	switch {
	case period == "hour":
		return ", resets hourly"
	case period == "week":
		return ", resets weekly"
	case period == "month":
		return ", resets monthly"
	case period == "nhour" && hours > 1:
		return fmt.Sprintf(", resets every %dh", hours)
	}
	return ""
}

// scopeTokens returns the tokens this key consumed within the scope entry's
// reset window, isolated per key. For an AUTO entry the cap covers everything
// the key routed through AUTO; for a model entry it covers that model only.
//
// It reads the per-key hourly buckets rather than the key-blind model
// buckets, so one key's usage can never exhaust another's quota.
func (g *Gateway) scopeTokens(ctx context.Context, sc config.ModelScope) int64 {
	k := keyID(reqKey(ctx))
	win := AutoPeriodSeconds(sc.Period, sc.Hours)
	if sc.Model != "" && strings.EqualFold(sc.Model, "AUTO") {
		return g.stats.KeyWindowTokens(k, win)
	}
	return g.stats.KeyWindowModelTokens(k, sc.Model, sc.Source, win)
}

// hasScopeModel reports whether a model (possibly with a "source-model" /
// "source:model" / "source/model" pinning prefix) is allowed by a key's model
// scope. The prefix is stripped strictly: only when the prefix names a real
// source that actually serves the bare model (via Registry.EffectiveModel), so
// model ids that themselves contain separators (e.g. "deepseek-v4-flash-free")
// are never corrupted (P10-2).
//
// A scope entry with model "AUTO" only matches the literal AUTO routing mode;
// it does NOT grant access to specific model ids.
func (g *Gateway) hasScopeModel(list []config.ModelScope, s string) bool {
	if r := g.core.Registry(); r != nil {
		s = r.EffectiveModel(s)
	}
	for _, x := range list {
		if x.Model == s {
			return true
		}
	}
	return false
}

// writeScopeReject answers a model-scope or quota rejection. A spent quota is
// 429 (rate_limit_exceeded) with Retry-After, so a client waits and resumes
// after the reset; a model the key may not use stays 403 (model_not_allowed),
// because retrying cannot help.
func (g *Gateway) writeScopeReject(w http.ResponseWriter, r *http.Request, model string) {
	if q := g.checkQuota(r.Context(), model); q != nil {
		if q.retry > 0 {
			w.Header().Set("Retry-After", strconv.FormatInt(q.retry, 10))
			writeError(w, http.StatusTooManyRequests, "rate_limit_exceeded", q.msg)
			return
		}
		// no window to wait for: the cap is either permanent or key-wide
		// with no period. "rate_limit_exceeded" still says "come back
		// after the operator raises the cap", which 403 would not.
		if strings.Contains(q.msg, "quota exceeded") {
			writeError(w, http.StatusTooManyRequests, "rate_limit_exceeded", q.msg)
			return
		}
		writeError(w, http.StatusForbidden, "model_not_allowed", q.msg)
		return
	}
	// The scope check already passed; reaching here means the state changed
	// between the two calls. Fall back to the pre-existing behaviour.
	writeError(w, http.StatusForbidden, "model_not_allowed", fmt.Sprintf("model %q is not allowed for this key", model))
}

func (g *Gateway) resolveByModel(model string) ([]*provider.Provider, string) {
	if isAuto(model) {
		return g.core.Registry().Resolve("AUTO"), ""
	}
	return g.core.Registry().Resolve(model), g.core.Registry().EffectiveModel(model)
}

// chatOnly keeps providers that expose at least one chat-capable model, so a
// chat/AUTO request never lands on an image-only source (or borrows its image
// model id). Explicit image-kind requests stay on the imageOnly path.
func chatOnly(cands []*provider.Provider) []*provider.Provider {
	out := make([]*provider.Provider, 0, len(cands))
	for _, p := range cands {
		for _, id := range p.Models() {
			if m := p.ModelByID(id); m == nil || m.Kind != "image" {
				out = append(out, p)
				break
			}
		}
	}
	return out
}

// toolRequest reports whether the request participates in a tool-call round.
func toolRequest(req *chatRequest) bool {
	if len(req.Tools) > 0 || req.ToolChoice != nil {
		return true
	}
	for _, m := range req.Messages {
		if m.Role == "tool" || len(m.ToolCalls) > 0 {
			return true
		}
	}
	return false
}

func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use POST")
		return
	}
	var req chatRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
		return
	}
	// 客户端自带的会话标识（若有）比推导出来的更准，见 clientSessionFromRequest。
	clientSession := clientSessionFromRequest(r, req.PromptCacheKey)
	if len(req.Messages) == 0 {
		writeError(w, http.StatusBadRequest, "invalid_request", "messages is required")
		return
	}
	model := req.Model
	if model == "" {
		model = g.core.DefaultModel()
	}
	if isAuto(model) {
		chain := g.core.AutoChain()
		if chain == nil || len(chain.Tiers) == 0 {
			writeError(w, http.StatusServiceUnavailable, "no_provider", "no auto slot configured")
			return
		}
		if msg := g.checkModelScope(r.Context(), "AUTO"); msg != "" {
			g.writeScopeReject(w, r, "AUTO")
			return
		}
		ctx := r.Context()
		done := g.stats.Begin()
		defer done()
		inner := &types.ChatRequest{
			Messages:        req.Messages,
			Temperature:     req.Temperature,
			MaxTokens:       req.MaxTokens,
			Stream:          req.Stream,
			Tools:           req.Tools,
			ToolChoice:      req.ToolChoice,
			DisableThinking: req.DisableThinking,
			ExtraBody:       req.ExtraBody,
			ClientSession:   clientSession,
		}
		rec := &Req{
			Key:  keyID(reqKey(ctx)),
			Type: "chat",
			OK:   false,
		}
		// quotaExhausted reports a slot whose token window has been used up;
		// exhausted slots are dropped from scheduling without penalty.
		quotaExhausted := func(sl *scheduler.Slot) bool {
			if sl.Quota <= 0 {
				return false
			}
			win := AutoPeriodSeconds(sl.Period, sl.Hours)
			return g.stats.WindowTokens(sl.Model, sl.Source, win) >= sl.Quota
		}
		if req.Stream {
			rec.Type = "stream"
			g.streamChatAuto(w, ctx, chain, inner, rec, quotaExhausted)
			return
		}
		g.singleChatAuto(w, ctx, chain, inner, rec, quotaExhausted)
		return
	}
	if !isAuto(model) {
		if allow := g.allowedModels(r.Context()); allow != nil && !g.hasScopeModel(allow, model) {
			writeError(w, http.StatusForbidden, "model_not_allowed", fmt.Sprintf("model %q is not allowed for this key", model))
			return
		}
	}
	cands, effective := g.resolveCands(r.Context(), &req)
	if len(cands) == 0 {
		writeError(w, http.StatusNotFound, "model_not_found", fmt.Sprintf("model %q is not configured", model))
		return
	}
	if effective == "" {
		effective = firstModel(cands[0])
	}
	if msg := g.checkModelScope(r.Context(), effective); msg != "" {
		g.writeScopeReject(w, r, effective)
		return
	}
	ctx := r.Context()
	done := g.stats.Begin()
	defer done()

	inner := &types.ChatRequest{
		Model:           effective,
		Messages:        req.Messages,
		Temperature:     req.Temperature,
		MaxTokens:       req.MaxTokens,
		Stream:          req.Stream,
		Tools:           req.Tools,
		ToolChoice:      req.ToolChoice,
		DisableThinking: req.DisableThinking,
		ExtraBody:       req.ExtraBody,
		ClientSession:   clientSession,
	}
	rec := &Req{
		Key:    keyID(reqKey(ctx)),
		Type:   "chat",
		Model:  effective,
		Source: firstSource(cands),
		OK:     false,
	}
	if req.Stream {
		rec.Type = "stream"
		g.streamChat(w, ctx, cands, inner, effective, rec)
		return
	}
	g.singleChat(w, ctx, cands, inner, effective, rec)
}

func firstModel(p *provider.Provider) string {
	ms := p.Models()
	if len(ms) > 0 {
		return ms[0]
	}
	return "auto"
}

// imageOnly keeps providers exposing at least one image-kind model.
func imageOnly(cands []*provider.Provider) []*provider.Provider {
	var out []*provider.Provider
	for _, p := range cands {
		for _, id := range p.Models() {
			if m := p.ModelByID(id); m != nil && m.Kind == "image" {
				out = append(out, p)
				break
			}
		}
	}
	return out
}

// toolCallsWire converts unified tool calls to the OpenAI wire format:
// tool_calls:[{id,type,function:{name,arguments:StringJSON}}]. Clients expect
// arguments to be a JSON string, not an object.
func toolCallsWire(tcs []types.ToolCall) json.RawMessage {
	wire := make([]map[string]interface{}, 0, len(tcs))
	for _, tc := range tcs {
		args := "{}"
		if tc.Arguments != nil {
			if b, err := json.Marshal(tc.Arguments); err == nil {
				args = string(b)
			}
		}
		wire = append(wire, map[string]interface{}{
			"id":   tc.ID,
			"type": tc.Type,
			"function": map[string]interface{}{
				"name":      tc.Name,
				"arguments": args,
			},
		})
	}
	b, _ := json.Marshal(wire)
	return b
}

func firstSource(cands []*provider.Provider) string {
	if len(cands) > 0 {
		return cands[0].Name()
	}
	return ""
}

// effectiveImageModel picks the image model id that will actually be used so
// quota checks can target it (AUTO resolves to the first image candidate).
func effectiveImageModel(model string, cands []*provider.Provider) string {
	if !isAuto(model) && model != "" {
		return model
	}
	if len(cands) > 0 {
		for _, id := range cands[0].Models() {
			if m := cands[0].ModelByID(id); m != nil && m.Kind == "image" {
				return id
			}
		}
	}
	return model
}

func estimatePromptTokens(req *types.ChatRequest) int64 {
	if req == nil {
		return 0
	}
	b, _ := json.Marshal(struct {
		Messages []types.ChatMessage `json:"messages"`
		Tools    []interface{}       `json:"tools,omitempty"`
	}{Messages: req.Messages, Tools: req.Tools})
	if len(b) == 0 {
		return 0
	}
	return int64(len(b)/3 + 1)
}

func estimateTextTokens(parts ...interface{}) int64 {
	var n int
	for _, p := range parts {
		switch v := p.(type) {
		case string:
			n += len(v)
		case json.RawMessage:
			n += len(v)
		case []types.ToolCall:
			b, _ := json.Marshal(v)
			n += len(b)
		}
	}
	if n == 0 {
		return 0
	}
	return int64(n/3 + 1)
}

// toScheduler adapts concrete providers to the scheduler.Provider interface.
// It lives here (not in the scheduler package) so scheduler tests do not pull
// in the provider package and with it the Lua runtime's link requirements.
func toScheduler(cands []*provider.Provider) []scheduler.Provider {
	out := make([]scheduler.Provider, len(cands))
	for i, p := range cands {
		out[i] = p
	}
	return out
}

// upstreamErrStatus maps a scheduling error to its HTTP status: a failed
// AUTO chain answers 503 with its per-tier summary, a busy source (every
// concurrency slot in use) is a transient capacity condition answered with
// 429 so clients fail fast, while other upstream failures stay 502.
func upstreamErrStatus(err error) int {
	var ce *scheduler.ChainErr
	if errors.As(err, &ce) {
		return http.StatusServiceUnavailable
	}
	if errors.Is(err, provider.ErrBusy) {
		return http.StatusTooManyRequests
	}
	return http.StatusBadGateway
}

// overflowMarkers 匹配上游「上下文超窗」类措辞。
//
// 上游写法五花八门，且**不在 pi 客户端的识别列表里**。pi 靠
// @earendil-works/pi-ai 的 OVERFLOW_PATTERNS 判断超窗并据此触发压缩重试，
// 而 justworker 返回的是「请精简对话历史…(Context window is full…)」——
// 与那 25 条正则一条都不匹配，于是 pi 既不压缩也不重试，只把它当成一条
// 普通上游错误。这里把可识别的超窗措辞归一化成 pi 一定认得的标记。
var overflowMarkers = []*regexp.Regexp{
	regexp.MustCompile(`(?i)context[ _-]?window is full`),
	regexp.MustCompile(`(?i)context[_ ]length[_ ]exceeded`),
	regexp.MustCompile(`(?i)exceeds? the context window`),
	regexp.MustCompile(`(?i)maximum context length`),
	regexp.MustCompile(`(?i)reduce the length of the messages`),
	regexp.MustCompile(`(?i)too many tokens`),
	regexp.MustCompile(`(?i)token limit exceeded`),
	regexp.MustCompile(`请精简对话历史`),
	regexp.MustCompile(`上下文(长度)?超(出|限)`),
	regexp.MustCompile(`对话历史过长`),
}

// overflowCanonical 命中 pi 的 /context[_ ]length[_ ]exceeded/i。
//
// 注意：这个前缀只是必要条件，不是充分条件——pi 还会先用
// NON_OVERFLOW_PATTERNS 排除整条消息，见 overflowClientMessage。
const overflowCanonical = "context_length_exceeded"

// looksLikeOverflow 判断错误文本是否属于上下文超窗。
func looksLikeOverflow(s string) bool {
	for _, re := range overflowMarkers {
		if re.MatchString(s) {
			return true
		}
	}
	return false
}

// clientUpstreamErr collapses upstream failure details into a short message
// for the client: per-tier bodies (WAF HTML pages, quota payloads, ...) stay
// in rec.Err / the stats API and the server log instead of the response.
// Per-tier one-line reasons are kept (quota/cooling skips carry no error body
// and are the actionable part); each is capped so HTML dumps can't leak.
//
// 超窗会被归一化成「干净」的 overflowCanonical 消息，见 overflowClientMessage。
func clientUpstreamErr(err error) string {
	log.Printf("[gateway] upstream failure surfaced to client: %v", err)
	if looksLikeOverflow(err.Error()) {
		return overflowClientMessage(err)
	}
	return upstreamErrSummary(err)
}

// overflowHint 是给客户端的超窗说明，刻意只用 pi 认得的措辞。
const overflowHint = "context window is full; reduce the length of the messages"

// clientSessionFromRequest 取客户端自带的会话标识，拿不到时返回 ""。
//
// 通用 OpenAI 客户端默认不带会话 id，但 pi 支持：只要 provider 的 compat 里打开
// sendSessionAffinityHeaders，pi 就会把**平台真实的会话 id**（uuidv7，整个会话
// 恒定）放到 x-session-affinity / x-client-request-id 上（sessionAffinityFormat
// 为 openrouter 时是 x-session-id，为 openai 时是 session_id）。
// 有了它就不必再从首条 user 消息推导会话指纹（那只在客户端不发会话 id 时才作为
// 退路，且历史压缩后会漂移）。prompt_cache_key 是 OpenAI 原生的缓存键，
// 客户端若带也一并当会话用。
//
// 刻意不采纳 x-client-request-id：名字含 request，部分客户端每请求都换，
// 拿它当会话会让上游缓存永不命中。pi 总会同时发 x-session-affinity，够用。
func clientSessionFromRequest(r *http.Request, bodyKey string) string {
	for _, h := range []string{"X-Session-Affinity", "X-Session-Id", "Session-Id"} {
		if v := strings.TrimSpace(r.Header.Get(h)); v != "" {
			return v
		}
	}
	return strings.TrimSpace(bodyKey)
}

// overflowClientMessage 为超窗失败生成「干净」的客户端消息。
//
// 为什么不能沿用 upstreamErrSummary：pi 的 isContextOverflow 先查
// NON_OVERFLOW_PATTERNS（/rate limit/、/too many requests/、Bedrock 前缀），
// 一旦命中就直接判为「非超窗」——**哪怕消息里已经有 context_length_exceeded**，
// pi 也不会压缩重试。而 AUTO 链的失败消息天生是多 tier 原因的拼接：
// 超窗 tier（gozen 400 maximum context length）常与配额/限流 tier
// （429 token plan exhausted、cooling、no free slot）同时出现。
// 把明细原样带出去，等于让一条限流 tier 的措辞反过来封杀超窗识别。
// 所以这里只保留「超窗」措辞 + 是哪个源超的窗，其余一律不带。
func overflowClientMessage(err error) string {
	var ce *scheduler.ChainErr
	if errors.As(err, &ce) {
		for _, t := range ce.Tiers {
			if looksLikeOverflow(t.Err.Error()) {
				return fmt.Sprintf("%s: %s (%s/%s)", overflowCanonical, overflowHint, t.Source, t.Model)
			}
		}
	}
	return fmt.Sprintf("%s: %s", overflowCanonical, overflowHint)
}

// upstreamErrSummary 把上游失败压成一行短消息。
func upstreamErrSummary(err error) string {
	var ce *scheduler.ChainErr
	if errors.As(err, &ce) {
		parts := make([]string, 0, len(ce.Tiers)+len(ce.Skipped))
		for _, t := range ce.Tiers {
			// 160 而不是 80：短诊断词（"Context window is full"）常落在尾部，
			// 80 字节按字节截断正好会把它切掉，超窗就再也认不出来。
			parts = append(parts, types.OneLine(fmt.Sprintf("%s/%s: %v", t.Source, t.Model, t.Err), 160))
		}
		for _, sk := range ce.Skipped {
			parts = append(parts, types.OneLine(sk, 160))
		}
		msg := strings.Join(parts, "; ")
		if len(msg) > 300 {
			msg = msg[:300] + "..."
		}
		return fmt.Sprintf("all %d auto providers failed: %s", len(parts), msg)
	}
	return types.OneLine(err.Error(), 160)
}

// failChat maps a scheduling failure onto the audit record and answers the
// client. A failed AUTO chain additionally pins its first failed tier onto
// the record; direct-path errors never match *scheduler.ChainErr, so the
// extraction is safely shared by all four entry points. Callers own the
// writeRec call (inline for non-stream paths, deferred for stream paths).
func (g *Gateway) failChat(w http.ResponseWriter, rec *Req, err error) {
	rec.OK = false
	rec.Status = upstreamErrStatus(err)
	rec.Err = err.Error()
	var ce *scheduler.ChainErr
	if errors.As(err, &ce) && len(ce.Tiers) > 0 {
		rec.Source = ce.Tiers[0].Source
		rec.Model = ce.Tiers[0].Model
	}
	writeError(w, rec.Status, "upstream_error", clientUpstreamErr(err))
}

// recordChatUsage fills token accounting for a finished non-streaming
// request: exact upstream numbers win, byte estimates fill the gaps.
func recordChatUsage(rec *Req, req *types.ChatRequest, resp *types.UnifiedResponse) {
	rec.Prompt = int64(resp.TokenUsage.Prompt)
	if rec.Prompt == 0 {
		rec.Prompt = estimatePromptTokens(req)
	}
	rec.Compl = int64(resp.TokenUsage.Completion)
	if rec.Compl == 0 {
		rec.Compl = estimateTextTokens(resp.Content, resp.ReasoningContent, resp.ToolCalls)
	}
	// Cache accounting: whichever source format the adapter normalized into
	// (prompt_tokens_details.cached_tokens or legacy prompt_cache_hit_tokens),
	// read it back so the request record carries the hit/miss split.
	if d := resp.TokenUsage.PromptTokensDetails; d != nil && d.CachedTokens > 0 {
		rec.CacheHit = int64(d.CachedTokens)
	}
	if resp.TokenUsage.PromptTokensDetails != nil {
		// Upstream reported cache details (even a 0 hit) — tag the row so
		// the UI can show 0% rather than “—”.
		rec.CacheReported = true
	} else if resp.TokenUsage.PromptCacheHit > 0 {
		rec.CacheHit = int64(resp.TokenUsage.PromptCacheHit)
		rec.CacheReported = true
	}
	if resp.TokenUsage.PromptCacheMiss > 0 {
		rec.CacheMiss = int64(resp.TokenUsage.PromptCacheMiss)
	}
}

// writeChatCompletion renders a unified response as an OpenAI
// chat.completion object. modelName is the id clients see as the serving
// model: the requested id for direct routes, the exact slot model for AUTO.
func writeChatCompletion(w http.ResponseWriter, resp *types.UnifiedResponse, modelName string) {
	msg := RespMessage{Role: "assistant", Content: resp.Content}
	if resp.ReasoningContent != "" {
		msg.ReasoningContent = resp.ReasoningContent
	}
	if len(resp.ToolCalls) > 0 {
		msg.ToolCalls = toolCallsWire(resp.ToolCalls)
	}
	out := ChatCompletion{
		ID:      newID(),
		Object:  "chat.completion",
		Created: time.Now().Unix(),
		Model:   modelName,
		Choices: []ChatChoice{{Index: 0, Message: msg, FinishReason: resp.FinishReason}},
		Cost:    resp.Cost,
	}
	if resp.TokenUsage.Total > 0 || resp.TokenUsage.Prompt > 0 || resp.TokenUsage.Completion > 0 {
		out.Usage = &resp.TokenUsage
	}
	writeJSON(w, http.StatusOK, out)
}

// singleChat runs a direct (model-pinned) non-streaming request across the
// candidate list, falling back on failure.
func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) {
	rec.LatMs = 0
	t0 := time.Now()
	resp, usedSrc, usedModel, err := g.core.Scheduler().Chat(ctx, toScheduler(cands), req)
	rec.LatMs = time.Since(t0).Milliseconds()
	if err != nil {
		g.failChat(w, rec, err)
		g.writeRec(rec)
		return
	}
	rec.OK = true
	rec.Status = http.StatusOK
	recordChatUsage(rec, req, resp)
	rec.Source = usedSrc
	rec.Model = usedModel
	// Non-streaming: the whole response arrives at once, so TTFB equals
	// the total latency.
	rec.FirstByteMs = rec.LatMs
	g.writeRec(rec)
	writeChatCompletion(w, resp, effective)
}

// writeRec records a finished request (audit + aggregates).
func (g *Gateway) writeRec(rec *Req) {
	if rec == nil {
		return
	}
	if rec.Time == 0 {
		rec.Time = time.Now().UnixMilli()
	}
	g.stats.Record(*rec)
}

// mergeUsage combines token usage across stream chunks additively. Some
// providers split usage across chunks (e.g. Anthropic reports prompt tokens
// in message_start and the final completion tokens in message_delta); a plain
// "last non-nil wins" would discard the prompt half. Non-zero fields from cur
// override prev; total is recomputed from the merged parts so a partial later
// chunk can't shrink it. For the common single-chunk case (OpenAI's terminal
// empty-choices+usage chunk) upstream totals are preserved exactly.
func mergeUsage(prev, cur *types.TokenUsage) *types.TokenUsage {
	if prev == nil {
		u := *cur
		if u.Total == 0 && (u.Prompt > 0 || u.Completion > 0) {
			u.Total = u.Prompt + u.Completion
		}
		return &u
	}
	out := *prev
	if cur.Prompt > 0 {
		out.Prompt = cur.Prompt
	}
	if cur.Completion > 0 {
		out.Completion = cur.Completion
	}
	out.Total = out.Prompt + out.Completion
	if cur.PromptTokensDetails != nil {
		// Keep the details object even when CachedTokens is 0: a reported
		// zero-hit is meaningful ("cache missed") and must stay
		// distinguishable from "upstream never reported cache info".
		// Dropping it here made streaming rows show “—” instead of 0%.
		out.PromptTokensDetails = cur.PromptTokensDetails
	}
	if cur.PromptCacheHit > 0 {
		out.PromptCacheHit = cur.PromptCacheHit
	}
	if cur.PromptCacheMiss > 0 {
		out.PromptCacheMiss = cur.PromptCacheMiss
	}
	return &out
}

// pumpStream writes the full SSE sequence for a started stream: role
// preamble, one chunk per unified delta, the terminating finish_reason, the
// OpenAI-standard final usage chunk (empty choices) and [DONE]. modelName
// follows writeChatCompletion's rule (requested id for direct routes, exact
// slot model for AUTO).
func (g *Gateway) pumpStream(w http.ResponseWriter, rec *Req, chunks <-chan types.UnifiedChunk, modelName string, t0 time.Time) {
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.WriteHeader(http.StatusOK)

	flusher, _ := w.(http.Flusher)
	id := newID()
	created := time.Now().Unix()
	send := func(obj interface{}) bool {
		b, err := json.Marshal(obj)
		if err != nil {
			return false
		}
		if _, err := fmt.Fprintf(w, "data: %s\n\n", b); err != nil {
			return false
		}
		if flusher != nil {
			flusher.Flush()
		}
		return true
	}

	if !send(ChatChunk{
		ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
		Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{Role: "assistant"}}},
	}) {
		return
	}
	// First SSE byte sent to the client: record time-to-first-byte for the
	// source's status-page latency average.
	rec.FirstByteMs = time.Since(t0).Milliseconds()
	var lastUsage *types.TokenUsage
	var lastCost string
	lastFinish := ""
	for ck := range chunks {
		if ck.Usage != nil {
			lastUsage = mergeUsage(lastUsage, ck.Usage)
		}
		if ck.Cost != "" {
			lastCost = ck.Cost
		}
		chunk := ChatChunk{
			ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
		}
		delta := RespMessage{Role: "assistant", Content: ck.Content}
		if ck.ReasoningContent != "" {
			delta.ReasoningContent = ck.ReasoningContent
		}
		if len(ck.ToolCalls) > 0 {
			delta.ToolCalls = ck.ToolCalls
		}
		choice := ChunkChoice{Index: 0, Delta: delta}
		if ck.Done {
			fin := ck.FinishReason
			if fin == "" {
				fin = "stop"
			}
			lastFinish = fin
			choice.FinishReason = &fin
		}
		chunk.Choices = []ChunkChoice{choice}
		rec.Compl += int64(len(ck.Content)+len(ck.ReasoningContent)+len(ck.ToolCalls)) / 3
		if !send(chunk) {
			return
		}
	}
	finalFinish := lastFinish
	if finalFinish == "" {
		finalFinish = "stop"
	}
	send(ChatChunk{
		ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
		Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &finalFinish}},
	})
	// Final usage chunk (OpenAI standard: empty choices + usage before [DONE]).
	// Prefer the upstream's exact usage if the stream carried it; fall back to
	// the gateway's estimate otherwise.
	var tut *types.TokenUsage
	if lastUsage != nil {
		tut = lastUsage
		// Write the upstream's cache accounting back onto the request record
		// so the audit trail carries the hit/miss split for streaming too.
		if d := tut.PromptTokensDetails; d != nil {
			rec.CacheReported = true
			if d.CachedTokens > 0 {
				rec.CacheHit = int64(d.CachedTokens)
			}
		} else if tut.PromptCacheHit > 0 {
			rec.CacheHit = int64(tut.PromptCacheHit)
			rec.CacheReported = true
		}
		if tut.PromptCacheMiss > 0 {
			rec.CacheMiss = int64(tut.PromptCacheMiss)
		}
	} else if rec.Prompt+rec.Compl > 0 {
		u := types.TokenUsage{
			Prompt:     int(rec.Prompt),
			Completion: int(rec.Compl),
			Total:      int(rec.Prompt + rec.Compl),
		}
		tut = &u
	}
	if tut != nil {
		send(ChatChunk{
			ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
			Choices: []ChunkChoice{},
			Usage:   tut,
			Cost:    lastCost,
		})
	} else if lastCost != "" {
		// Cost arrived without any usage (OpenCode sends it on its own frame).
		send(ChatChunk{
			ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
			Choices: []ChunkChoice{},
			Cost:    lastCost,
		})
	}
	fmt.Fprintf(w, "data: [DONE]\n\n")
	if flusher != nil {
		flusher.Flush()
	}
}

// streamChat runs a direct (model-pinned) streaming request across the
// candidate list, falling back early on connect errors.
func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) {
	rec.LatMs = 0
	t0 := time.Now()
	rec.OK = true
	rec.Status = http.StatusOK
	defer func() {
		rec.LatMs = time.Since(t0).Milliseconds()
		g.writeRec(rec)
	}()
	chunks, usedSrc, usedModel, err := g.core.Scheduler().ChatStream(ctx, toScheduler(cands), req)
	if err != nil {
		g.failChat(w, rec, err)
		return
	}
	if usedModel != "" {
		rec.Model = usedModel
	}
	// Audit accuracy: pin the source that actually served the stream (after a
	// failover it differs from the first candidate). Direct streams previously
	// discarded it.
	rec.Source = usedSrc
	rec.Prompt = estimatePromptTokens(req)
	g.pumpStream(w, rec, chunks, effective, t0)
}

// singleChatAuto runs a non-streaming AUTO request down the chain (see
// scheduler.ChainChat): tiers ascending (tier 1 = highest priority first),
// per-tier round-robin ordered by preference, cooldown as the only hard skip,
// busy slots skipped without penalty and a bounded busy wait. When every
// tier fails, the response is a 503 carrying the per-tier error summary
// (which source/model failed why).
func (g *Gateway) singleChatAuto(w http.ResponseWriter, ctx context.Context, chain *scheduler.Chain, req *types.ChatRequest, rec *Req, quotaExhausted func(*scheduler.Slot) bool) {
	rec.LatMs = 0
	t0 := time.Now()
	resp, usedSrc, usedModel, err := g.core.Scheduler().ChainChat(ctx, chain, req, quotaExhausted)
	rec.LatMs = time.Since(t0).Milliseconds()
	if err != nil {
		g.failChat(w, rec, err)
		g.writeRec(rec)
		return
	}
	rec.OK = true
	rec.Status = http.StatusOK
	recordChatUsage(rec, req, resp)
	rec.Source = usedSrc
	rec.Model = usedModel
	rec.FirstByteMs = rec.LatMs
	g.writeRec(rec)
	writeChatCompletion(w, resp, usedModel)
}

// streamChatAuto streams an AUTO request down the chain. A slot is abandoned
// only before its first chunk (connect error / non-200 / busy); once a stream
// starts it stays pinned. Total failure writes a JSON 503 (with the per-tier
// summary) before any SSE byte is sent.
func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, chain *scheduler.Chain, req *types.ChatRequest, rec *Req, quotaExhausted func(*scheduler.Slot) bool) {
	rec.LatMs = 0
	t0 := time.Now()
	rec.OK = true
	rec.Status = http.StatusOK
	defer func() {
		rec.LatMs = time.Since(t0).Milliseconds()
		g.writeRec(rec)
	}()
	chunks, usedSrc, usedModel, err := g.core.Scheduler().ChainChatStream(ctx, chain, req, quotaExhausted)
	if err != nil {
		g.failChat(w, rec, err)
		return
	}
	if usedModel != "" {
		rec.Model = usedModel
	}
	rec.Source = usedSrc
	rec.Prompt = estimatePromptTokens(req)
	g.pumpStream(w, rec, chunks, usedModel, t0)
}

func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use POST")
		return
	}
	var req types.ImageGenRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
		return
	}
	if req.Prompt == "" {
		writeError(w, http.StatusBadRequest, "invalid_request", "prompt is required")
		return
	}
	model := req.Model
	if model == "" {
		model = g.core.DefaultModel()
	}
	if isAuto(model) {
		if chain := g.core.AutoImageChain(); chain != nil && len(chain.Tiers) > 0 {
			if msg := g.checkModelScope(r.Context(), "AUTO"); msg != "" {
				g.writeScopeReject(w, r, "AUTO")
				return
			}
			done := g.stats.Begin()
			defer done()
			rec := &Req{Key: keyID(reqKey(r.Context())), Type: "image", Model: model, OK: false}
			t0 := time.Now()
			resp, usedSrc, usedModel, err := g.core.Scheduler().ChainImage(r.Context(), chain, &req)
			rec.LatMs = time.Since(t0).Milliseconds()
			if err != nil {
				rec.Status = upstreamErrStatus(err)
				rec.Err = err.Error()
				g.writeRec(rec)
				writeError(w, rec.Status, "upstream_error", clientUpstreamErr(err))
				return
			}
			rec.Source = usedSrc
			if usedModel != "" {
				rec.Model = usedModel // actual image model served, not "AUTO"
			}
			rec.OK = true
			rec.Status = http.StatusOK
			rec.Compl = int64(len(resp.ImageData))
			g.writeRec(rec)
			writeJSON(w, http.StatusOK, types.ImageGenResponse{
				Created: time.Now().Unix(),
				Data:    resp.ImageData,
			})
			return
		}
		// no image chain configured: fall through to legacy discovery (all
		// sources exposing an image model, tried in registry order)
	}
	if allow := g.allowedModels(r.Context()); allow != nil && !g.hasScopeModel(allow, model) {
		writeError(w, http.StatusForbidden, "model_not_allowed", fmt.Sprintf("model %q is not allowed for this key", model))
		return
	}
	cands, _ := g.resolveByModel(model)
	cands = imageOnly(cands)
	if allow := g.allowedModels(r.Context()); allow != nil {
		cands = filterCandsByModels(cands, allow)
	}
	if len(cands) == 0 {
		writeError(w, http.StatusServiceUnavailable, "no_provider", "no image source configured")
		return
	}
	if msg := g.checkModelScope(r.Context(), effectiveImageModel(model, cands)); msg != "" {
		g.writeScopeReject(w, r, effectiveImageModel(model, cands))
		return
	}
	done := g.stats.Begin()
	defer done()
	rec := &Req{Key: keyID(reqKey(r.Context())), Type: "image", Model: model, Source: firstSource(cands), OK: false}
	t0 := time.Now()
	resp, usedSrc, err := g.core.Scheduler().Image(r.Context(), toScheduler(cands), &req)
	rec.LatMs = time.Since(t0).Milliseconds()
	if err != nil {
		rec.Status = upstreamErrStatus(err)
		rec.Err = err.Error()
		g.writeRec(rec)
		writeError(w, rec.Status, "upstream_error", clientUpstreamErr(err))
		return
	}
	if usedSrc != "" {
		rec.Source = usedSrc
	}
	if resp.Model != "" {
		rec.Model = resp.Model // record the actual model served, not the raw request id
	}
	rec.OK = true
	rec.Status = http.StatusOK
	rec.Compl = int64(len(resp.ImageData))
	g.writeRec(rec)
	writeJSON(w, http.StatusOK, types.ImageGenResponse{
		Created: time.Now().Unix(),
		Data:    resp.ImageData,
	})
}