Files
ModelRouter/internal/gateway/chat.go
JianFeeeee 0121d23f91 fix(tokens): 流式统计改用上游真实 usage,图片不再记 token
两处 token 单位错误,均影响 per-model 配额计费:

1. 流式路径的 prompt/completion 只是「字节÷3」估算。
   pumpStream 明明收到了上游最后一帧的真实 usage,却只发给客户端、
   从不回写审计记录,于是配额按估算值扣。生产实测同一请求:
   上游 prompt=37/completion=179 → 记账 27/262,prompt 低估 1.4x、
   completion 高估 1.5x(双向失真)。同模型流式 completion 中位数
   是非流式的 4-27 倍。非流式路径本就用真实值,两路不一致。
   修法:lastUsage 非零时写回 rec.Prompt/rec.Compl,估算降为兜底
   (上游不报 usage 时仍保留原估算行为)。

2. 图片请求把「图片张数」记成 completion_tokens。
   rec.Compl = int64(len(resp.ImageData)),len 是切片长度即张数
   (生产 38 条 image 记录全是 1),且被计入 token 总量。
   图片生成无 token 概念 ⇒ 新增 Req.ImageCount 独立字段,
   Prompt/Compl 归 0;UI 记录表 image 行改显示张数(新增 i18n thImgs)。

顺带补 TestUILocaleKeyParity:此前无人校验 zh/en 键集合一致,
单边加键不会报错,只会显示原始键名。

新增 token_units_test.go(定值上游 6 项),做过变异验证:
回退修复实测复现 stream=16/173 vs chat=44/100、image completion=3。
2026-09-28 22:19:12 +08:00

1201 lines
40 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

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.
// checkModelScope keeps its string-returning signature for callers that only
// need to know whether the request may proceed.
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 validates the effective model against the key's model scope and
// that entry's quota. Returns nil when the request may proceed.
//
// Quotas are per scope entry, never key-wide: a model whose budget is spent
// is refused on its own while the key's other models keep working. The verdict
// carries the remaining seconds of the reset window so a spent budget answers
// 429 + Retry-After (come back when it rolls over) instead of 403 (which reads
// as "this key may never use this model" and makes clients give up).
func (g *Gateway) checkQuota(ctx context.Context, model string) *quotaRejection {
allow := g.allowedModels(ctx)
if allow == nil {
return nil
}
for _, sc := range allow {
if sc.Model != model {
continue
}
win := AutoPeriodSeconds(sc.Period, sc.Hours)
k := keyID(reqKey(ctx))
if sc.TokenQuota > 0 {
if used := g.scopeTokens(ctx, sc); 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),
}
}
}
if sc.ReqQuota > 0 {
if used := g.stats.KeyWindowReqs(k, win); used >= sc.ReqQuota {
return &quotaRejection{
msg: fmt.Sprintf("request quota exceeded for %q (%d/%d)", model, used, sc.ReqQuota),
retry: AutoSecondsToReset(sc.Period, sc.Hours),
}
}
}
return nil
}
return &quotaRejection{msg: fmt.Sprintf("model %q is not allowed for this key", model)}
}
// 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 on the scope entry's model
// within its 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) writeReject(w http.ResponseWriter, q *quotaRejection) {
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)
}
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 q := g.checkQuota(r.Context(), "AUTO"); q != nil {
g.writeReject(w, q)
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 q := g.checkQuota(r.Context(), effective); q != nil {
g.writeReject(w, q)
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.
// Write the upstream's exact numbers back onto the audit record. Without
// this the streamed path kept only the per-chunk byte estimate, so the same
// request recorded ~1.4x its real completion tokens (measured: upstream 100,
// audit 145) while the non-streaming path recorded 100. Two paths, two
// different numbers for one request is a reporting bug, not a rounding one.
if lastUsage != nil {
if lastUsage.Prompt > 0 {
rec.Prompt = int64(lastUsage.Prompt)
}
if lastUsage.Completion > 0 {
rec.Compl = int64(lastUsage.Completion)
}
}
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 q := g.checkQuota(r.Context(), "AUTO"); q != nil {
g.writeReject(w, q)
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
// Image generation has no token concept. Recording len(ImageData)
// (the image COUNT) in completion_tokens mislabels image count as
// tokens and feeds it into the token totals; leave it 0.
rec.ImageCount = 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 q := g.checkQuota(r.Context(), effectiveImageModel(model, cands)); q != nil {
g.writeReject(w, q)
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
// Image generation has no token concept — see the AUTO path above.
rec.ImageCount = len(resp.ImageData)
g.writeRec(rec)
writeJSON(w, http.StatusOK, types.ImageGenResponse{
Created: time.Now().Unix(),
Data: resp.ImageData,
})
}