mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-11-07 09:47:23 +00:00
用户连续报「codeart 还是不可用 / empty stream」。把上游原始响应打出来逐条对齐, 三个症状同源,且都不是"模型坏了": | 场景 | 上游实际返回 | |---|---| | 流式、短时间连续请求 | HTTP **200** + 只带 error 的帧 ⇒ 内核判 `empty stream` | | 非流式、同一条件 | HTTP 400 `TM.00001041: 并发会话数已达上限(3个)` | | 突发请求 | HTTP 429 `TPM limit … 500000 tokens per minute` | 实测:最后一个请求结束后约 **44~60 秒**释放;同一账号连续请求 3~4 次即占满。 参考实现(示例插件 llm-adapter.ts)的做法是**每 10s 轮询后重试 chat**(最多 30 分钟),所以用户那边"看起来没事"。 ## 我们的改法(内核只做通用机制,厂商知识留在插件) 1. **内核新增声明式重试策略** `source.retry`(纯数据:状态码列表 + 响应体子串 + 间隔 + 次数)。内核不认识任何厂商;命中即在同一网关请求内每 10s 重试一次。 流式只在**握手阶段**重试(首帧尚未下发,不会重复输出),非流式整体重试。 2. **插件为 codearts 声明**:400/429 + 命中 `TM.00001041` / `TPM limit` / `并发会话数已达上限` → 10s × 最多 18 次(≈3 分钟)。 3. **适配器补两个会话细节**:`Connection: close`(keep-alive 空闲连接会被上游 计入并发会话)与稳定 `Session-Id`(同一账号恒定,避免每次请求新建会话)。 4. 错误文案补可行动提示(并发会话上限 / TPM 限流,约 1 分钟释放)。 ## 验证(生产) - 连续 6 次**流式**:6/6 成功(撞上限时自动等 10s 再试,客户端只感觉慢): `+0s / +12s / +34s / +86s / +100s / +144s` - 连续 5 次**非流式**:5/5 成功(`+0 / +11 / +20 / +28 / +38s`) - 修复前:第 4 个请求起必然 TM.00001041 / empty stream。 > 注:`TM.00001005 无正常可用套餐` 属于账号权益问题(不重试、如实报出), > 与限流要分开处理。
1790 lines
59 KiB
Go
1790 lines
59 KiB
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"
|
||
"net/url"
|
||
"strconv"
|
||
"strings"
|
||
"sync"
|
||
"sync/atomic"
|
||
"time"
|
||
|
||
"llmsproxy/internal/config"
|
||
"llmsproxy/internal/lua"
|
||
"llmsproxy/internal/types"
|
||
)
|
||
|
||
// proxyFor picks the transport proxy for one source. An explicit per-source
|
||
// proxy_url wins (e.g. clash on 127.0.0.1:7890 for a geo-blocked upstream);
|
||
// otherwise fall back to the process-wide env proxy, which is direct by
|
||
// default.
|
||
func proxyFor(cfg config.Source) func(*http.Request) (*url.URL, error) {
|
||
if cfg.ProxyURL != "" {
|
||
proxyURL, err := url.Parse(cfg.ProxyURL)
|
||
if err == nil {
|
||
return http.ProxyURL(proxyURL)
|
||
}
|
||
}
|
||
return http.ProxyFromEnvironment
|
||
}
|
||
|
||
// ---- 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 // +prefSuccessStep per success / -prefFailStep per failure, clamped
|
||
failCount atomic.Int64
|
||
cooldownUntil atomic.Int64 // unix seconds; 0 = schedulable
|
||
// prefTouched is the unix second pref was last moved by an outcome. It is
|
||
// the reference point for decay: a negative score with no traffic since
|
||
// prefDecayAfter drifts one step toward 0 per elapsed interval, so a slot
|
||
// cannot stay written off forever just because nothing selected it.
|
||
prefTouched atomic.Int64
|
||
// cooldownFrom is the unix second the current cooldown window opened. It
|
||
// is the reference point for the half-cooldown probe gate: probing is
|
||
// allowed only in the second half of [cooldownFrom, cooldownUntil), so a
|
||
// freshly failed slot stays fully silent for a while instead of being
|
||
// hammered immediately.
|
||
cooldownFrom atomic.Int64
|
||
// probeInFlight is the single-token probe permit. At most one probe
|
||
// request may be in flight per (source, model) while cooling; the token is
|
||
// released by ProbeDone regardless of outcome.
|
||
probeInFlight atomic.Bool
|
||
}
|
||
|
||
const (
|
||
// prefFailStep is the penalty for a REAL failure (5xx, transport error,
|
||
// unusable adapter output): something is wrong with the upstream or the
|
||
// model, so two of them are enough to push a slot to the back of its tier.
|
||
prefFailStep = 5
|
||
// prefQuotaStep is the penalty for running out of allowance (quota/TPM
|
||
// exhausted). Deliberately much smaller than prefFailStep: an exhausted
|
||
// quota means "no budget right now", not "this model is bad", and for
|
||
// allowance-metered sources it is an everyday event rather than a fault.
|
||
// Charging it the full failure penalty made such sources drift to a deeply
|
||
// negative score with ZERO visible failures - RecordQuotaExhausted
|
||
// intentionally does not touch failCount, so the WebUI showed fails=0 and
|
||
// cooling=false while the score sat at -10.
|
||
prefQuotaStep = 1
|
||
// prefSuccessStep is the reward for a successful request. Recovery used to
|
||
// be +1 against a -5 penalty, a 5:1 asymmetry that made a slot at -10 need
|
||
// ten consecutive successes just to reach neutral. Rewarding +2 keeps
|
||
// failures meaningful (a real failure still costs more than one success
|
||
// earns) while letting a recovered slot return to normal rotation quickly.
|
||
prefSuccessStep = 2
|
||
prefMin = -20
|
||
prefMax = 20
|
||
// prefDecayAfter is how long a slot may sit at a negative score without any
|
||
// traffic before the score drifts back toward neutral. Without this a slot
|
||
// that was penalised and then simply not selected again (the score itself
|
||
// makes it unattractive) stays penalised indefinitely: it has no way to earn
|
||
// the successes that would rehabilitate it. Applied lazily on read, so there
|
||
// are no timers.
|
||
prefDecayAfter = 2 * time.Minute
|
||
backoffBase = 5 * time.Second
|
||
// backoffCap bounds the exponential schedule. It used to be 30min, which
|
||
// in practice meant "a slot whose upstream blipped is unusable for half an
|
||
// hour even though its quota recovered in seconds". Combined with the
|
||
// half-cooldown probe gate below a much shorter cap is safe: a still-dead
|
||
// upstream costs one probe request per cap/2, while a recovered one is back
|
||
// in full rotation on its first probe success.
|
||
backoffCap = 5 * time.Minute
|
||
// failures at/above this count back off at the capped duration
|
||
backoffCapN = 10
|
||
)
|
||
|
||
// authCooldown is the cooldown applied to 401/403 credential failures. A bad
|
||
// key does not fix itself in seconds, so it cools longer than a transport
|
||
// blip — but it is still a self-healing cooldown (a rotated key is picked up by
|
||
// the next probe), never a permanent blacklist.
|
||
const authCooldown = 10 * time.Minute
|
||
|
||
// maxQuotaCooldown bounds a quota-exhausted cooldown. Even when the reported
|
||
// reset window is huge (monthly quotas), the slot is re-probed at least this
|
||
// often so a manually topped-up allowance is noticed.
|
||
const maxQuotaCooldown = 30 * time.Minute
|
||
|
||
// quotaDefaultCooldown is used when the source gives no hint about its quota
|
||
// window length. It is long enough not to hammer an exhausted allowance and
|
||
// short enough that the half-cooldown probe retries within minutes.
|
||
const quotaDefaultCooldown = 10 * time.Minute
|
||
|
||
// 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 backoffCap). auth marks 401/403
|
||
// credential failures: they get the longer fixed authCooldown and a doubled
|
||
// preference penalty, but the model still recovers when the cooldown expires.
|
||
func (s *ModelState) RecordFailure(auth bool) {
|
||
n := s.failCount.Add(1)
|
||
var cd time.Duration
|
||
switch {
|
||
case auth:
|
||
// Persist the cap so FailCount() reports "this is a hard problem" and a
|
||
// following non-auth failure does not restart the ladder from the bottom.
|
||
if n < backoffCapN {
|
||
n = backoffCapN
|
||
s.failCount.Store(n)
|
||
}
|
||
cd = authCooldown
|
||
case n >= backoffCapN:
|
||
cd = backoffCap
|
||
default:
|
||
cd = backoffBase * time.Duration(1<<(n-1))
|
||
if cd > backoffCap {
|
||
cd = backoffCap
|
||
}
|
||
}
|
||
s.openCooldown(cd)
|
||
pen := int64(prefFailStep)
|
||
if auth {
|
||
pen *= 2
|
||
}
|
||
s.addPref(-pen)
|
||
}
|
||
|
||
// addPref moves the preference score and stamps prefTouched, which restarts the
|
||
// decay clock. Every outcome-driven score change must go through here so decay
|
||
// never fires on a slot that is actively being scored.
|
||
func (s *ModelState) addPref(delta int64) {
|
||
s.pref.Add(delta)
|
||
clampPref(&s.pref, prefMin, prefMax)
|
||
s.prefTouched.Store(time.Now().Unix())
|
||
}
|
||
|
||
// decayPref drifts a negative score back toward neutral once a slot has been
|
||
// idle for prefDecayAfter. This closes a starvation loop: a penalised slot is
|
||
// unattractive, so it stops being selected, so it never earns the successes that
|
||
// would rehabilitate it, so it stays penalised. One step per elapsed interval,
|
||
// never past 0, and positive scores are left alone (a proven-good slot should
|
||
// not be dragged down for being idle).
|
||
//
|
||
// Called lazily from Pref(); no timers and no goroutines.
|
||
func (s *ModelState) decayPref() {
|
||
cur := s.pref.Load()
|
||
if cur >= 0 {
|
||
return
|
||
}
|
||
last := s.prefTouched.Load()
|
||
if last <= 0 {
|
||
// no reference point yet: start the clock instead of decaying blindly
|
||
s.prefTouched.CompareAndSwap(0, time.Now().Unix())
|
||
return
|
||
}
|
||
step := int64(prefDecayAfter / time.Second)
|
||
if step <= 0 {
|
||
return
|
||
}
|
||
elapsed := time.Now().Unix() - last
|
||
if elapsed < step {
|
||
return
|
||
}
|
||
steps := elapsed / step
|
||
target := cur + steps
|
||
if target > 0 {
|
||
target = 0
|
||
}
|
||
// Advance prefTouched by exactly the consumed intervals so leftover time
|
||
// still counts toward the next step.
|
||
if s.pref.CompareAndSwap(cur, target) {
|
||
s.prefTouched.Store(last + steps*step)
|
||
}
|
||
}
|
||
|
||
// openCooldown starts a fresh cooldown window of length cd. Recording both
|
||
// ends of the window (not just the deadline) is what makes the half-cooldown
|
||
// probe gate possible, and resetting cooldownFrom on every new failure means a
|
||
// failed probe pushes the next probe to the middle of the NEW window instead of
|
||
// retrying immediately.
|
||
func (s *ModelState) openCooldown(cd time.Duration) {
|
||
now := time.Now()
|
||
s.cooldownFrom.Store(now.Unix())
|
||
s.cooldownUntil.Store(now.Add(cd).Unix())
|
||
}
|
||
|
||
// 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.openCooldown(rateLimitCooldown)
|
||
// A 429 is "too fast", not "broken": charge it like an allowance event, not
|
||
// like a fault, or a throttled-but-healthy source sinks purely for being
|
||
// popular.
|
||
s.addPref(-prefQuotaStep)
|
||
}
|
||
|
||
// RecordQuotaExhausted marks the model as out of allowance until its quota
|
||
// window resets. Unlike a failure this is not "the upstream is broken", so it
|
||
// must not feed the exponential ladder: the slot is unusable until the window
|
||
// boundary and immediately usable again after it. resetIn is the time left in
|
||
// the current quota window; non-positive values fall back to the short
|
||
// rate-limit window (the caller could not determine a period).
|
||
func (s *ModelState) RecordQuotaExhausted(resetIn time.Duration) {
|
||
if resetIn <= 0 {
|
||
resetIn = rateLimitCooldown
|
||
}
|
||
if resetIn > maxQuotaCooldown {
|
||
resetIn = maxQuotaCooldown
|
||
}
|
||
s.openCooldown(resetIn)
|
||
// No failCount bump: an exhausted quota is not an error streak, and
|
||
// letting it inflate failCount would make the NEXT real failure jump
|
||
// straight to the capped backoff.
|
||
//
|
||
// The score penalty is prefQuotaStep, not prefFailStep, for the same
|
||
// reason: cooldown alone already keeps the slot out of rotation until the
|
||
// window resets, so the score only needs to express a mild preference for
|
||
// slots with budget left. Charging the full failure penalty here is what
|
||
// drove allowance-metered sources to -10 and beyond with fails=0.
|
||
s.addPref(-prefQuotaStep)
|
||
}
|
||
|
||
// RecordSuccess resets the failure counter and cooldown and rewards the
|
||
// preference score. A single success lifts a slot off the prefMin floor, which
|
||
// is what lets one successful probe return a written-off slot to normal
|
||
// rotation.
|
||
func (s *ModelState) RecordSuccess() {
|
||
s.failCount.Store(0)
|
||
s.cooldownUntil.Store(0)
|
||
s.cooldownFrom.Store(0)
|
||
s.addPref(prefSuccessStep)
|
||
}
|
||
|
||
func (s *ModelState) reset() {
|
||
s.failCount.Store(0)
|
||
s.cooldownUntil.Store(0)
|
||
s.cooldownFrom.Store(0)
|
||
s.probeInFlight.Store(false)
|
||
s.pref.Store(0)
|
||
s.prefTouched.Store(0)
|
||
}
|
||
|
||
// TryProbe attempts to claim the single probe permit for a model that normal
|
||
// scheduling refuses. It returns true only when a probe is both needed and due:
|
||
//
|
||
// - cooling: allowed once the window passed its midpoint (the first half stays
|
||
// completely silent).
|
||
// - cooldown expired but the preference score sank to prefMin: ModelAvailable
|
||
// still refuses the slot, so a probe is the ONLY way back. Without this a
|
||
// slot that failed prefMin/prefFailStep times in a row would be blacklisted
|
||
// forever despite every cooldown having long expired.
|
||
//
|
||
// A healthy, non-cooling model returns false: the caller should schedule it
|
||
// normally. On true the caller MUST call ProbeDone exactly once.
|
||
func (s *ModelState) TryProbe() bool {
|
||
now := time.Now().Unix()
|
||
until := s.cooldownUntil.Load()
|
||
if until <= now {
|
||
s.decayPref() // a naturally-recovered slot needs no probe
|
||
if s.pref.Load() > prefMin {
|
||
return false // schedulable normally, no probe needed
|
||
}
|
||
// preference floor: probe to earn the way out
|
||
return s.probeInFlight.CompareAndSwap(false, true)
|
||
}
|
||
from := s.cooldownFrom.Load()
|
||
if from <= 0 || from >= until {
|
||
return false // no measurable window (legacy/degenerate state)
|
||
}
|
||
if now < from+(until-from)/2 {
|
||
return false // first half of the window: stay silent
|
||
}
|
||
return s.probeInFlight.CompareAndSwap(false, true)
|
||
}
|
||
|
||
// ProbeDone releases the probe permit claimed by TryProbe.
|
||
func (s *ModelState) ProbeDone() { s.probeInFlight.Store(false) }
|
||
|
||
// Probing reports whether a probe request currently holds the permit.
|
||
func (s *ModelState) Probing() bool { return s.probeInFlight.Load() }
|
||
|
||
// CooldownFrom is the unix timestamp the current cooldown window opened; 0
|
||
// when the model is not cooling.
|
||
func (s *ModelState) CooldownFrom() int64 { return s.cooldownFrom.Load() }
|
||
|
||
// Pref is the current preference score (higher = preferred). Reading it also
|
||
// applies idle decay, so a slot penalised long ago and never selected since
|
||
// drifts back toward neutral instead of being written off permanently.
|
||
func (s *ModelState) Pref() int64 {
|
||
s.decayPref()
|
||
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.
|
||
// A per-source ProxyURL (e.g. clash on 127.0.0.1:7890) overrides the
|
||
// process-wide env proxy for upstreams that are geo/IP-blocked; sources
|
||
// without one keep http.ProxyFromEnvironment (direct by default).
|
||
tr := &http.Transport{
|
||
Proxy: proxyFor(cfg),
|
||
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
|
||
}
|
||
|
||
// ModelsWithKind returns the same ids as Models but paired with each model's
|
||
// kind, in the same order.
|
||
//
|
||
// Models() cannot answer "which of these are image models", and that question
|
||
// has to be answerable: the chat test page lists models from one source list,
|
||
// and an image model offered there fails on every call because /api/chat is the
|
||
// chat endpoint.
|
||
func (p *Provider) ModelsWithKind() []ModelKindInfo {
|
||
out := make([]ModelKindInfo, 0, len(p.cfg.Models))
|
||
for _, m := range p.cfg.Models {
|
||
kind := m.Kind
|
||
if kind == "" {
|
||
kind = "chat" // config.Model documents "" as the chat default
|
||
}
|
||
out = append(out, ModelKindInfo{ID: m.ID, Kind: kind})
|
||
}
|
||
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 first chat-kind model of this source.
|
||
//
|
||
// ★ 曾经叫 "highest-priority chat model",但 `config.Model.Priority` 已删除
|
||
//
|
||
// (它从未参与跨源调度,只在这里用;而这里的语义其实只是"选一个能发 chat 的
|
||
// 模型",选第一个就够了 —— 源内模型顺序由配置/发现顺序决定)。
|
||
func (p *Provider) bestChatModel() string {
|
||
for _, m := range p.cfg.Models {
|
||
if m.Kind != "" && m.Kind != "chat" {
|
||
continue
|
||
}
|
||
return m.ID
|
||
}
|
||
if len(p.cfg.Models) > 0 {
|
||
return p.cfg.Models[0].ID
|
||
}
|
||
return ""
|
||
}
|
||
|
||
// bestImageModel returns the first 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 {
|
||
for _, m := range p.cfg.Models {
|
||
if m.Kind != "" && m.Kind != "image" {
|
||
continue
|
||
}
|
||
return m.ID
|
||
}
|
||
if len(p.cfg.Models) > 0 {
|
||
return p.cfg.Models[0].ID
|
||
}
|
||
return ""
|
||
}
|
||
|
||
// 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. It may contain the placeholder
|
||
// "{model}"; use URL (or ChatURL) rather than calling this directly when the
|
||
// path has to be usable.
|
||
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"
|
||
}
|
||
|
||
// modelPlaceholder is the marker an adapter puts in its endpoint template when
|
||
// the upstream carries the model id in the PATH rather than in the body.
|
||
// Every other adapter's endpoint is a fixed path, so substituting is a no-op
|
||
// for them.
|
||
const modelPlaceholder = "{model}"
|
||
|
||
// ChatURL resolves the full chat URL for a specific model.
|
||
//
|
||
// Two substitutions happen here, both driven by endpoint declarations:
|
||
//
|
||
// - "{model}" is replaced by the model id actually being sent. Without this
|
||
// the gateway would POST to Gemini's model-LIST endpoint, which answers 405.
|
||
// - a streaming request may use a DIFFERENT endpoint. This is deliberately
|
||
// not a vendor rule: the kernel does not know that one vendor's verb
|
||
// changes, only that an endpoint declaration said "use this other path
|
||
// when streaming".
|
||
//
|
||
// Precedence (a source that spells its own path takes full responsibility for
|
||
// it — the old behaviour, which must not regress):
|
||
//
|
||
// source.endpoint_stream > source.endpoint > adapter.endpoint_stream > adapter.endpoint
|
||
//
|
||
// Note the third slot only applies when the source did NOT spell its own
|
||
// chat endpoint. A source with `endpoint: /custom/path` keeps that path for
|
||
// streaming too, because the operator overrode the adapter's template and the
|
||
// kernel has no business re-adding the adapter's stream path on top.
|
||
func (p *Provider) ChatURL(model string, stream bool) string {
|
||
ep := p.Endpoint()
|
||
if stream {
|
||
switch {
|
||
case p.cfg.StreamEndpoint != "":
|
||
ep = p.cfg.StreamEndpoint
|
||
case p.cfg.Endpoint == "":
|
||
// Source did not override the chat path, so the adapter's streaming
|
||
// endpoint (if any) is authoritative.
|
||
if se := p.vm.StreamEndpoint(p.adapter); se != "" {
|
||
ep = se
|
||
}
|
||
}
|
||
}
|
||
if strings.Contains(ep, modelPlaceholder) {
|
||
// A model id is put in a URL path, so it must be escaped: an id with a
|
||
// slash would otherwise silently address a different resource.
|
||
ep = strings.ReplaceAll(ep, modelPlaceholder, url.PathEscape(model))
|
||
}
|
||
return strings.TrimRight(p.cfg.BaseURL, "/") + ep
|
||
}
|
||
|
||
func (p *Provider) URL() string {
|
||
return p.ChatURL("", false)
|
||
}
|
||
|
||
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() && s.Pref() > prefMin
|
||
}
|
||
return true
|
||
}
|
||
|
||
// ModelSchedulable is the scheduling gate that replaces a bare ModelAvailable
|
||
// check. It reports whether the model may take a request right now and whether
|
||
// doing so is a cooldown PROBE rather than normal traffic:
|
||
//
|
||
// - (true, false): fully available, schedule normally.
|
||
// - (true, true): cooling, but past the half-cooldown mark and this caller
|
||
// claimed the probe permit. The caller MUST call ProbeDone(model) once the
|
||
// attempt finished, and SHOULD prefer non-probe candidates first.
|
||
// - (false, false): cooling and not probeable (first half of the window, the
|
||
// permit is taken, or the slot is persistently bad).
|
||
//
|
||
// Probing is what stops a recovered upstream from sitting out a whole cooldown:
|
||
// the probe is a real request, so success immediately clears the cooldown via
|
||
// RecordSuccess and the slot is back in full rotation.
|
||
func (p *Provider) ModelSchedulable(model string) (ok bool, isProbe bool) {
|
||
p.mu.Lock()
|
||
s, known := p.states[model]
|
||
p.mu.Unlock()
|
||
if !known {
|
||
return true, false
|
||
}
|
||
if s.Available() && s.Pref() > prefMin {
|
||
return true, false
|
||
}
|
||
if s.TryProbe() {
|
||
return true, true
|
||
}
|
||
return false, false
|
||
}
|
||
|
||
// ProbeDone releases a probe permit claimed through ModelSchedulable.
|
||
func (p *Provider) ProbeDone(model string) { p.state(model).ProbeDone() }
|
||
|
||
// ProbeSlots is how many concurrent probe requests this source tolerates while
|
||
// a model cools: one per 10 configured concurrency slots, at least 1 and never
|
||
// more than 2. A source configured for max_concurrent=10 therefore lets exactly
|
||
// one request through to probe, which is the intended "small-scale probe"
|
||
// behaviour; a large source still cannot flood a cooling upstream.
|
||
func (p *Provider) ProbeSlots() int {
|
||
n := p.cfg.MaxConcurrent / 10
|
||
if n < 1 {
|
||
return 1
|
||
}
|
||
if n > 2 {
|
||
return 2
|
||
}
|
||
return n
|
||
}
|
||
|
||
// 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()
|
||
}
|
||
|
||
// ModelProbeInfo exposes the probe-gate state for the web UI: when the current
|
||
// cooldown window opened, when probing becomes allowed (its midpoint), and
|
||
// whether a probe is in flight right now. All zero when not cooling.
|
||
func (p *Provider) ModelProbeInfo(model string) (from, probeAfter int64, probing bool) {
|
||
s := p.state(model)
|
||
until := s.CooldownUntil()
|
||
from = s.CooldownFrom()
|
||
if until <= 0 || from <= 0 || from >= until {
|
||
return 0, 0, s.Probing()
|
||
}
|
||
return from, from + (until-from)/2, s.Probing()
|
||
}
|
||
|
||
// 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
|
||
probeURL := p.ChatURL(model, false)
|
||
if hdrs, herr := p.buildHeaders(string(body), probeURL, ""); herr == nil {
|
||
hdr = hdrs
|
||
}
|
||
raw, status, derr := p.do(ctx, probeURL, 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 → long self-healing authCooldown
|
||
// 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 long lockout that wastes remaining quota.
|
||
//
|
||
// Prefer ReportStatusReason when the upstream error text is available: a
|
||
// quota-exhausted rejection is cooled to its window boundary instead of being
|
||
// treated as a failure streak.
|
||
func (p *Provider) ReportStatus(model string, code int) {
|
||
p.ReportStatusReason(model, code, "")
|
||
}
|
||
|
||
// ReportStatusReason is ReportStatus plus the upstream error text, which lets a
|
||
// quota/allowance rejection be distinguished from a generic rate limit. Quota
|
||
// exhaustion is not an error streak: it is cooled until the allowance window
|
||
// can plausibly have reset, and never inflates the exponential ladder.
|
||
func (p *Provider) ReportStatusReason(model string, code int, reason string) {
|
||
if code == 401 || code == 403 {
|
||
p.RecordFailure(model, code)
|
||
return
|
||
}
|
||
if code == 429 || code == 402 {
|
||
if quotaExhaustedReason(reason) {
|
||
p.state(model).RecordQuotaExhausted(p.quotaResetIn(model))
|
||
return
|
||
}
|
||
p.state(model).RecordRateLimit()
|
||
return
|
||
}
|
||
if code >= 500 {
|
||
p.RecordFailure(model, code)
|
||
}
|
||
}
|
||
|
||
// quotaResetIn estimates how long the model's allowance window still has to
|
||
// run. Upstreams do not report this, so it is taken from the source's optional
|
||
// `quota_hours` meta hint when present (aligned to the epoch, matching the
|
||
// gateway's own quota-window convention) and otherwise falls back to
|
||
// quotaDefaultCooldown. Either way the value is capped by maxQuotaCooldown, so
|
||
// a wrong hint can only delay recovery to that bound — and the half-cooldown
|
||
// probe fires at half of it.
|
||
func (p *Provider) quotaResetIn(model string) time.Duration {
|
||
hours := p.quotaHours()
|
||
if hours <= 0 {
|
||
return quotaDefaultCooldown
|
||
}
|
||
win := int64(hours * float64(time.Hour/time.Second))
|
||
if win <= 0 {
|
||
return quotaDefaultCooldown
|
||
}
|
||
elapsed := time.Now().Unix() % win
|
||
return time.Duration(win-elapsed) * time.Second
|
||
}
|
||
|
||
// quotaHours reads the optional per-source `quota_hours` meta hint describing
|
||
// how long its allowance window is. Absent or unparseable = 0 (unknown).
|
||
func (p *Provider) quotaHours() float64 {
|
||
v, ok := p.cfg.Meta["quota_hours"]
|
||
if !ok {
|
||
return 0
|
||
}
|
||
switch n := v.(type) {
|
||
case float64:
|
||
return n
|
||
case int:
|
||
return float64(n)
|
||
case int64:
|
||
return float64(n)
|
||
case string:
|
||
f, err := strconv.ParseFloat(strings.TrimSpace(n), 64)
|
||
if err != nil {
|
||
return 0
|
||
}
|
||
return f
|
||
}
|
||
return 0
|
||
}
|
||
|
||
// quotaKeywords are the upstream phrases that mean "you are out of allowance"
|
||
// rather than "you are going too fast". Matching is case-insensitive substring
|
||
// matching against the adapter-normalized error reason.
|
||
var quotaKeywords = []string{
|
||
"insufficient_quota",
|
||
"insufficient quota",
|
||
"quota exceeded",
|
||
"quota_exceeded",
|
||
"exceeded your current quota",
|
||
"out of credit",
|
||
"insufficient balance",
|
||
"insufficient_user_quota",
|
||
"billing_hard_limit_reached",
|
||
"credit balance is too low",
|
||
"余额不足",
|
||
"额度不足",
|
||
"额度已用完",
|
||
}
|
||
|
||
func quotaExhaustedReason(reason string) bool {
|
||
if reason == "" {
|
||
return false
|
||
}
|
||
low := strings.ToLower(reason)
|
||
for _, k := range quotaKeywords {
|
||
if strings.Contains(low, k) {
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
}
|
||
|
||
// 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, clientSession string) (http.Header, error) {
|
||
meta := map[string]interface{}{
|
||
"url": url,
|
||
"method": http.MethodPost,
|
||
"body": body,
|
||
"api_key": p.cfg.APIKey,
|
||
"timestamp": types.Now(),
|
||
// client_session 是客户端自带的会话标识(请求头),为空时适配器自行推导。
|
||
"client_session": clientSession,
|
||
"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
|
||
}
|
||
chatURL := p.ChatURL(model, false)
|
||
hdrs, err := p.buildHeaders(body, chatURL, req.ClientSession)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
raw, status, err := p.doWithRetry(ctx, chatURL, 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 {
|
||
reason := p.apiErrReason(status, raw)
|
||
p.ReportStatusReason(model, status, reason)
|
||
return nil, fmt.Errorf("%s", reason)
|
||
}
|
||
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
|
||
}
|
||
// stream=true so a path-carried model plus the streaming verb is resolved
|
||
// for THIS request's model, not the source default.
|
||
streamURL := p.ChatURL(model, true)
|
||
hdrs, err := p.buildHeaders(body, streamURL, req.ClientSession)
|
||
if err != nil {
|
||
p.Release()
|
||
return nil, err
|
||
}
|
||
|
||
type respOrErr struct {
|
||
resp *http.Response
|
||
err error
|
||
}
|
||
rc := make(chan respOrErr, 1)
|
||
go func() {
|
||
// Retry only the HANDSHAKE: a policy-matched rejection (e.g. a
|
||
// per-account concurrency cap answered with a business 4xx) can clear
|
||
// within seconds, and nothing has been sent to the client yet — the
|
||
// kernel holds the first chunk back. Once a 200 arrives and streaming
|
||
// starts, no retry happens (that would duplicate output).
|
||
rp := p.cfg.Retry
|
||
attempts := 1
|
||
if rp != nil && rp.MaxAttempts > 1 {
|
||
attempts = rp.MaxAttempts
|
||
}
|
||
for i := 0; i < attempts; i++ {
|
||
resp, err := p.doRawStream(ctx, streamURL, body, hdrs)
|
||
if err != nil {
|
||
rc <- respOrErr{resp, err}
|
||
return
|
||
}
|
||
if resp.StatusCode == 200 {
|
||
rc <- respOrErr{resp, nil}
|
||
return
|
||
}
|
||
// A rejection can only be classified after the body is read: the
|
||
// policy matches on body substrings, and a source may use the same
|
||
// status (400) for unrelated errors that must NOT be retried.
|
||
raw, _ := io.ReadAll(resp.Body)
|
||
resp.Body.Close()
|
||
if !rp.ShouldRetry(resp.StatusCode, string(raw)) {
|
||
// Not the retryable body: hand the rejection back verbatim so
|
||
// the caller reports the real upstream error.
|
||
resp2 := &http.Response{StatusCode: resp.StatusCode, Body: io.NopCloser(strings.NewReader(string(raw)))}
|
||
rc <- respOrErr{resp2, nil}
|
||
return
|
||
}
|
||
if i == attempts-1 {
|
||
resp2 := &http.Response{StatusCode: resp.StatusCode, Body: io.NopCloser(strings.NewReader(string(raw)))}
|
||
rc <- respOrErr{resp2, nil}
|
||
return
|
||
}
|
||
if serr := sleepCtx(ctx, rp.DelayDuration()); serr != nil {
|
||
rc <- respOrErr{nil, serr}
|
||
return
|
||
}
|
||
}
|
||
}()
|
||
|
||
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()
|
||
reason := p.apiErrReason(sel.resp.StatusCode, string(raw))
|
||
p.ReportStatusReason(model, sel.resp.StatusCode, reason)
|
||
p.Release()
|
||
return nil, fmt.Errorf("%s", reason)
|
||
}
|
||
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
|
||
var sawToolCalls bool
|
||
// toolIdx maps a per-stream tool fragment's upstream index (Anthropic
|
||
// content-block ordinal, which is sparse when thinking/text blocks
|
||
// precede tool_use) to a dense 0-based OpenAI tool_call ordinal so
|
||
// clients that accumulate fragments by index reassemble multi-tool
|
||
// responses correctly.
|
||
toolIdx := newToolIndexRemapper()
|
||
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
|
||
}
|
||
// A few OpenAI-compatible upstreams emit a final `stop` marker
|
||
// after streaming tool-call fragments. Once a real fragment was
|
||
// observed, the gateway must preserve tool-call semantics for the
|
||
// agent instead of making it treat the turn as ordinary text.
|
||
if len(ck.ToolCalls) > 0 {
|
||
trimmed := strings.TrimSpace(string(ck.ToolCalls))
|
||
if trimmed != "null" && trimmed != "[]" && trimmed != "{}" {
|
||
sawToolCalls = true
|
||
}
|
||
}
|
||
if sawToolCalls && ck.FinishReason == "stop" {
|
||
ck.FinishReason = "tool_calls"
|
||
}
|
||
ck.ToolCalls = toolIdx.remap(ck.ToolCalls)
|
||
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.
|
||
|
||
// toolIndexRemapper maps per-stream tool_call fragment indices from a sparse
|
||
// upstream ordinal (e.g. Anthropic content_block index, which skips over
|
||
// thinking/text blocks) to a dense 0-based OpenAI tool_call ordinal.
|
||
type toolIndexRemapper struct {
|
||
seen map[int]int // upstream index -> dense ordinal
|
||
next int // next dense ordinal to assign
|
||
}
|
||
|
||
func newToolIndexRemapper() toolIndexRemapper {
|
||
return toolIndexRemapper{seen: make(map[int]int)}
|
||
}
|
||
|
||
func (t *toolIndexRemapper) remap(raw json.RawMessage) json.RawMessage {
|
||
if len(raw) == 0 {
|
||
return raw
|
||
}
|
||
// Decode into generic maps so every upstream field is preserved verbatim;
|
||
// only the index value is rewritten.
|
||
var frags []map[string]json.RawMessage
|
||
if err := json.Unmarshal(raw, &frags); err != nil || len(frags) == 0 {
|
||
return raw
|
||
}
|
||
changed := false
|
||
for i := range frags {
|
||
orig := 0
|
||
if v, ok := frags[i]["index"]; ok {
|
||
if err := json.Unmarshal(v, &orig); err != nil {
|
||
continue
|
||
}
|
||
}
|
||
mapped, ok := t.seen[orig]
|
||
if !ok {
|
||
mapped = t.next
|
||
t.next++
|
||
t.seen[orig] = mapped
|
||
}
|
||
if mapped != orig {
|
||
b, err := json.Marshal(mapped)
|
||
if err != nil {
|
||
continue
|
||
}
|
||
frags[i]["index"] = b
|
||
changed = true
|
||
}
|
||
}
|
||
if !changed {
|
||
return raw
|
||
}
|
||
out, err := json.Marshal(frags)
|
||
if err != nil {
|
||
return raw
|
||
}
|
||
return out
|
||
}
|
||
|
||
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))
|
||
}
|
||
}
|
||
// ⚠️ 措辞必须**区分两种原因**(实测教训):适配器**有** transform_error 钩子
|
||
// 但它返回 nil(不认识的错误体)是**常见**情况,而"真的没这个钩子"是另一回事。
|
||
// 原先一律说 "lacks transform_error" —— 我照着这句话去适配器里找那个函数,
|
||
// 查了几轮才发现函数明明在、只是抽不出消息。
|
||
// 这里直接问内核要事实,不再猜(判据也据此断言)。
|
||
if p.vm.HasFn(p.adapter, "transform_error") {
|
||
log.Printf("[provider] unhandled upstream error body (adapter %q: transform_error returned no reason for this shape): status=%d body=%.300s",
|
||
p.adapter, status, types.OneLine(raw, 300))
|
||
} else {
|
||
log.Printf("[provider] unhandled upstream error body (adapter %q defines no 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 {
|
||
reason := p.apiErrReason(status, raw)
|
||
p.ReportStatusReason(model, status, reason)
|
||
return nil, fmt.Errorf("%s", reason)
|
||
}
|
||
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
|
||
}
|
||
|
||
// doWithRetry runs one non-stream request, retrying per the source's
|
||
// declarative RetryPolicy. The policy is pure data (status list + body
|
||
// substrings), so the kernel stays vendor-agnostic: a provider that answers a
|
||
// per-account concurrency cap with a business 4xx can clear by itself while the
|
||
// client just waits a few seconds instead of seeing an error.
|
||
func (p *Provider) doWithRetry(ctx context.Context, url, body string, hdr http.Header) (string, int, error) {
|
||
rp := p.cfg.Retry
|
||
attempts := 1
|
||
if rp != nil && rp.MaxAttempts > 1 {
|
||
attempts = rp.MaxAttempts
|
||
}
|
||
var raw string
|
||
var status int
|
||
for i := 0; i < attempts; i++ {
|
||
r, st, err := p.do(ctx, url, body, hdr)
|
||
if err != nil {
|
||
return "", 0, err
|
||
}
|
||
raw, status = r, st
|
||
if status == 200 || !rp.ShouldRetry(status, raw) || i == attempts-1 {
|
||
break
|
||
}
|
||
if serr := sleepCtx(ctx, rp.DelayDuration()); serr != nil {
|
||
return "", 0, serr
|
||
}
|
||
}
|
||
return raw, status, nil
|
||
}
|
||
|
||
// sleepCtx sleeps for d unless the request context ends first.
|
||
func sleepCtx(ctx context.Context, d time.Duration) error {
|
||
if d <= 0 {
|
||
return ctx.Err()
|
||
}
|
||
t := time.NewTimer(d)
|
||
defer t.Stop()
|
||
select {
|
||
case <-ctx.Done():
|
||
return ctx.Err()
|
||
case <-t.C:
|
||
return 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"`
|
||
CompletionTokensDetails *struct {
|
||
ReasoningTokens int `json:"reasoning_tokens"`
|
||
} `json:"completion_tokens_details"`
|
||
} `json:"usage"`
|
||
// Cost is a top-level decimal string on OpenCode Zen/Go responses.
|
||
Cost json.RawMessage `json:"cost"`
|
||
}
|
||
if err := json.Unmarshal([]byte(data), &raw); err != nil {
|
||
return ""
|
||
}
|
||
if len(raw.Choices) == 0 && raw.UpstreamUsage.Total == 0 && raw.UpstreamUsage.TotalTokens == 0 {
|
||
// Cost-only chunk: OpenCode sends the charge on its own frame
|
||
// {"choices":[],"cost":"0"} with no usage at all. Dropping it here
|
||
// would lose the only cost signal the client can see.
|
||
if c := rawCostString(raw.Cost); c != "" {
|
||
out, _ := json.Marshal(types.UnifiedChunk{Cost: c})
|
||
return string(out)
|
||
}
|
||
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,
|
||
}
|
||
}
|
||
if pu.CompletionTokensDetails != nil {
|
||
usage.CompletionTokensDetails = &types.CompletionTokensDetails{
|
||
ReasoningTokens: pu.CompletionTokensDetails.ReasoningTokens,
|
||
}
|
||
}
|
||
}
|
||
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)
|
||
}
|
||
|
||
// rawCostString normalizes an upstream cost field into a plain string. The
|
||
// value is a decimal **string** on OpenCode Zen/Go ("0", "0.0012"), but other
|
||
// upstreams have been seen to send a bare number, so both are accepted.
|
||
// Returns "" when absent, null, or empty.
|
||
func rawCostString(raw json.RawMessage) string {
|
||
s := strings.TrimSpace(string(raw))
|
||
if s == "" || s == "null" {
|
||
return ""
|
||
}
|
||
if s[0] == '"' {
|
||
var unq string
|
||
if json.Unmarshal(raw, &unq) == nil {
|
||
return strings.TrimSpace(unq)
|
||
}
|
||
}
|
||
return s
|
||
}
|
||
|
||
// 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
|
||
}
|