From a42ff62d060287717eee2b14cb9c354bd77bd32d Mon Sep 17 00:00:00 2001 From: JianFeeeee Date: Wed, 26 Aug 2026 21:02:33 +0800 Subject: [PATCH] feat(scheduler): separate image-generation AUTO chain with UI toggle Image models previously could not be scheduled through a priority chain: the chat AUTO chain explicitly skips image-kind slots, and AUTO image requests fell back to unordered registry discovery. - config: add auto_image rules (auto_image yaml / image_rules json); legacy auto rules keep their meaning as the chat chain - core: buildAutoImageChain mirrors buildAutoChain with inverted kind filter (image-only); SaveAutoImageRules + AutoImageRules/AutoImageChain - scheduler: ChainImage walks the chain tier-by-tier with round-robin and preference ordering, skipping cooling slots - gateway: handleImage AUTO now runs down AutoImageChain when one is configured (falls back to legacy discovery otherwise) and records the actual served model; handleAutoAPI GET returns image_rules and PUT accepts image_rules independently of rules - webui: priority page gains a chat/image toggle editing two independent lane sets; add-slot picker filters by active kind; persistAuto writes only the active chain's field --- .gitignore | 1 + internal/config/config.go | 4 +- internal/core/core.go | 90 ++++++++++++++++++-- internal/gateway/chat.go | 40 ++++++++- internal/gateway/keys.go | 28 +++++-- internal/gateway/ui/index.html | 143 ++++++++++++++++++++++---------- internal/scheduler/scheduler.go | 57 +++++++++++++ 7 files changed, 299 insertions(+), 64 deletions(-) diff --git a/.gitignore b/.gitignore index 856bfa7..94104b3 100644 --- a/.gitignore +++ b/.gitignore @@ -20,3 +20,4 @@ cmd/gui/*.log # local ops scripts (machine-specific, not for sharing) scripts/ +cmd/gui/dist/ diff --git a/internal/config/config.go b/internal/config/config.go index f9269b2..bb1c90d 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -24,7 +24,8 @@ type Config struct { TLSKeyFile string `yaml:"tls_key_file,omitempty"` // PEM private key PublicBaseURL string `yaml:"public_base_url,omitempty"` // external base for generated config snippets; default inferred from request Sources []Source `yaml:"sources"` - Auto []ModelScope `yaml:"auto,omitempty"` // AUTO 调度链规则（WebUI 优先级页编辑） + Auto []ModelScope `yaml:"auto,omitempty"` // AUTO 调度链规则（WebUI 优先级页编辑，chat） + AutoImage []ModelScope `yaml:"auto_image,omitempty"` // AUTO 生图调度链规则（WebUI 优先级页·生图） Keys []GWKey `yaml:"keys,omitempty"` // 网关密钥（WebUI 密钥页管理） } @@ -318,6 +319,7 @@ type RuntimeConfig struct { DeletedAdapters []string `json:"deleted_adapters,omitempty"` Keys []GWKey `json:"keys,omitempty"` Auto []ModelScope `json:"auto,omitempty"` + AutoImage []ModelScope `json:"auto_image,omitempty"` } // GWKey is a gateway API key persisted in the config file. Role is "admin" diff --git a/internal/core/core.go b/internal/core/core.go index 08f5761..1dce2ba 100644 --- a/internal/core/core.go +++ b/internal/core/core.go @@ -32,7 +32,8 @@ type Core struct { store *config.Store scheduler *scheduler.Scheduler registry *provider.Registry - autoChain atomic.Pointer[scheduler.Chain] + autoChain atomic.Pointer[scheduler.Chain] // chat AUTO chain + autoImageChain atomic.Pointer[scheduler.Chain] // image-generation AUTO chain } // New builds the core from a config file plus runtime overlay. @@ -181,6 +182,9 @@ func (c *Core) Registry() *provider.Registry { return c.registry } // AutoChain returns the current AUTO scheduling chain. func (c *Core) AutoChain() *scheduler.Chain { return c.autoChain.Load() } +// AutoImageChain returns the persisted image-generation AUTO chain snapshot. +func (c *Core) AutoImageChain() *scheduler.Chain { return c.autoImageChain.Load() } + func (c *Core) DefaultModel() string { return c.cfg.DefaultModel } func (c *Core) GatewayKeys() []string { return c.cfg.GatewayKeys } @@ -294,6 +298,15 @@ func (c *Core) AutoRules() []config.ModelScope { return out } +// AutoImageRules returns the configured image-generation AUTO chain rules. +func (c *Core) AutoImageRules() []config.ModelScope { + c.mu.Lock() + defer c.mu.Unlock() + out := make([]config.ModelScope, len(c.cfg.AutoImage)) + copy(out, c.cfg.AutoImage) + return out +} + // cleanScopes drops empty model entries and normalizes placeholder source // names; it returns nil when no entries survive so an empty scope means // "unrestricted" (nil) instead of a restrictive-but-empty list — a non-nil @@ -325,16 +338,35 @@ func (c *Core) SaveAutoRules(entries []config.ModelScope) error { return err } c.buildAutoChain() - if ch := c.autoChain.Load(); ch != nil { - for _, tn := range ch.Tiers { - for _, sl := range tn.Slots { - if p := c.registry.ProviderForSlot(sl.Model, sl.Source); p != nil { - p.ResetModelCooldown(sl.Model) - } + c.resetChainCooldowns(c.autoChain.Load()) + return nil +} + +// SaveAutoImageRules persists the image-generation AUTO chain. Image models +// are kept as-is (unlike buildAutoChain, which skips them). +func (c *Core) SaveAutoImageRules(entries []config.ModelScope) error { + c.mu.Lock() + defer c.mu.Unlock() + c.cfg.AutoImage = cleanScopes(entries) + if err := c.saveConfig(); err != nil { + return err + } + c.buildAutoImageChain() + c.resetChainCooldowns(c.autoImageChain.Load()) + return nil +} + +func (c *Core) resetChainCooldowns(ch *scheduler.Chain) { + if ch == nil { + return + } + for _, tn := range ch.Tiers { + for _, sl := range tn.Slots { + if p := c.registry.ProviderForSlot(sl.Model, sl.Source); p != nil { + p.ResetModelCooldown(sl.Model) } } } - return nil } // ResetHealth clears the scheduling backoff state of every provider. @@ -420,6 +452,7 @@ func (c *Core) rebuildRegistry() error { c.registry.Replace(providers) } c.buildAutoChain() + c.buildAutoImageChain() return nil } @@ -466,6 +499,47 @@ func (c *Core) buildAutoChain() { c.autoChain.Store(scheduler.BuildChain(sr, prov)) } +// buildAutoImageChain rebuilds the image-generation AUTO chain snapshot. +// Unlike buildAutoChain, only image-kind models participate: chat models in +// the rules are skipped so a stale chat slot can't receive image traffic. +func (c *Core) buildAutoImageChain() { + prov := func(model, source string) scheduler.Provider { + p := c.registry.ProviderForSlot(model, source) + if p == nil { + return nil + } + if m := p.ModelByID(model); m != nil && m.Kind != "image" { + return nil + } + return p + } + rules := c.cfg.AutoImage + sr := make([]scheduler.Rule, 0, len(rules)) + for _, e := range rules { + model, source := e.Model, e.Source + if p := c.registry.ProviderForSlot(e.Model, e.Source); p != nil { + if exact := p.ModelIDFold(e.Model); exact != "" { + model = exact + } + if m := p.ModelByID(model); m != nil && m.Kind != "image" { + continue // chat-kind models never join the image AUTO chain + } + if source == "" { + source = p.Name() + } + } + sr = append(sr, scheduler.Rule{ + Model: model, + Source: source, + Tier: e.Tier, + Quota: e.TokenQuota, + Period: e.Period, + Hours: e.Hours, + }) + } + c.autoImageChain.Store(scheduler.BuildChain(sr, prov)) +} + // AutoSlotState is the UI-facing health snapshot of one AUTO chain slot. type AutoSlotState struct { Model string `json:"model"` diff --git a/internal/gateway/chat.go b/internal/gateway/chat.go index 5fd9821..3c993ed 100644 --- a/internal/gateway/chat.go +++ b/internal/gateway/chat.go @@ -897,11 +897,45 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) { if model == "" { model = g.core.DefaultModel() } - 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)) + if isAuto(model) { + if chain := g.core.AutoImageChain(); chain != nil && len(chain.Tiers) > 0 { + if msg := g.checkModelScope(r.Context(), "AUTO"); msg != "" { + writeError(w, http.StatusForbidden, "model_not_allowed", msg) + return + } + done := g.stats.Begin() + defer done() + rec := &Req{Key: keyID(reqKey(r.Context())), Type: "image", Model: model, OK: false} + t0 := time.Now() + resp, usedSrc, usedModel, err := g.core.Scheduler().ChainImage(r.Context(), chain, &req) + rec.LatMs = time.Since(t0).Milliseconds() + if err != nil { + rec.Status = upstreamErrStatus(err) + rec.Err = err.Error() + g.writeRec(rec) + writeError(w, rec.Status, "upstream_error", clientUpstreamErr(err)) + return + } + rec.Source = usedSrc + if usedModel != "" { + rec.Model = usedModel // actual image model served, not "AUTO" + } + rec.OK = true + rec.Status = http.StatusOK + rec.Compl = int64(len(resp.ImageData)) + g.writeRec(rec) + writeJSON(w, http.StatusOK, types.ImageGenResponse{ + Created: time.Now().Unix(), + Data: resp.ImageData, + }) return } + // no image chain configured: fall through to legacy discovery (all + // sources exposing an image model, tried in registry order) + } + if allow := g.allowedModels(r.Context()); allow != nil && !g.hasScopeModel(allow, model) { + writeError(w, http.StatusForbidden, "model_not_allowed", fmt.Sprintf("model %q is not allowed for this key", model)) + return } cands, _ := g.resolveByModel(model) cands = imageOnly(cands) diff --git a/internal/gateway/keys.go b/internal/gateway/keys.go index 19a2ae5..2d8c1e7 100644 --- a/internal/gateway/keys.go +++ b/internal/gateway/keys.go @@ -143,8 +143,9 @@ func (g *Gateway) allowedModels(ctx context.Context) []config.ModelScope { func (g *Gateway) handleAutoAPI(w http.ResponseWriter, r *http.Request) { if r.Method == http.MethodGet { writeJSON(w, http.StatusOK, map[string]interface{}{ - "rules": g.core.AutoRules(), - "states": g.core.AutoSlotStates(), + "rules": g.core.AutoRules(), + "image_rules": g.core.AutoImageRules(), + "states": g.core.AutoSlotStates(), }) return } @@ -155,17 +156,30 @@ func (g *Gateway) handleAutoAPI(w http.ResponseWriter, r *http.Request) { switch r.Method { case http.MethodPut, http.MethodPost: var body struct { - Rules []config.ModelScope `json:"rules"` + Rules []config.ModelScope `json:"rules"` + ImageRules []config.ModelScope `json:"image_rules"` } if err := json.NewDecoder(r.Body).Decode(&body); err != nil { writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error()) return } - if err := g.core.SaveAutoRules(body.Rules); err != nil { - writeError(w, http.StatusBadRequest, "auto_error", err.Error()) - return + if body.Rules != nil { + if err := g.core.SaveAutoRules(body.Rules); err != nil { + writeError(w, http.StatusBadRequest, "auto_error", err.Error()) + return + } } - writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "rules": g.core.AutoRules()}) + if body.ImageRules != nil { + if err := g.core.SaveAutoImageRules(body.ImageRules); err != nil { + writeError(w, http.StatusBadRequest, "auto_error", err.Error()) + return + } + } + writeJSON(w, http.StatusOK, map[string]interface{}{ + "ok": true, + "rules": g.core.AutoRules(), + "image_rules": g.core.AutoImageRules(), + }) default: writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "") } diff --git a/internal/gateway/ui/index.html b/internal/gateway/ui/index.html index 8e3bcd0..7e6bac1 100644 --- a/internal/gateway/ui/index.html +++ b/internal/gateway/ui/index.html @@ -365,6 +365,8 @@ #modal-wrap .card{width:100%;max-width:480px;animation:cardIn .28s var(--ease-spring) both} /* ---- scratch-style priority canvas (sakura restyle, same classes) ---- */ + .kind-toggle{display:inline-flex;gap:4px;margin-right:8px} + .kind-toggle button{padding:4px 12px;font-size:11.5px} .sort-canvas{position:relative;min-height:420px;padding:22px 14px 30px;margin:4px 0 10px;background-color:var(--card2); background-image:linear-gradient(var(--line-strong) 1px,transparent 1px),linear-gradient(90deg,var(--line-strong) 1px,transparent 1px); background-size:28px 28px;border:1px solid var(--line);border-radius:16px;overflow-x:auto;backdrop-filter:blur(10px)} @@ -2563,14 +2565,19 @@ function srcShort(name) { return (name || "?").slice(0, 2).toUpperCase(); } - const sortState = { lanes: [], origin: null, drag: null }; + const sortState = { lanes: [], origin: null, drag: null, kind: "chat" }; + // per-kind editor state: chat lanes and image lanes are independent; + // the toggle switches which chain the canvas edits + const sortLanes = { chat: [], image: [] }; let sortStateMap = new Map(); async function renderSort() { const j = await api("/api/sources"); let autoR = []; + let autoImg = []; try { const a = await api("/api/auto"); autoR = a.rules || []; + autoImg = a.image_rules || []; sortStateMap = new Map( (a.states || []).map((st) => [ st.model + "|" + (st.source || "*"), @@ -2590,50 +2597,62 @@ sourceRows.get(it.prio).push({ src: s.name, id: m.id }); }), ); - if (autoR.length) { - const tiers = []; - autoR.forEach((r, i) => { - const src = normSrc(r.source); - const pair = src ? byPair.get(r.model + "|" + src) : null; - if (src && !pair) return; - const ref = pair || - byModel.get(r.model) || { src: "*", id: r.model }; - const ti = Math.max(1, parseInt(r.tier) || i + 1) - 1; - if (!tiers[ti]) tiers[ti] = []; - tiers[ti].push({ - src: ref.src || "*", - id: r.model, - uid: nexUid(), - meta: { - quota: r.token_quota || 0, - period: r.period || "", - hours: r.hours || 0, - }, + // build a lane structure for one kind from its saved rules; when + // rules are empty, fall back to discovery from sourceRows filtered by + // kind so an unconfigured chain still shows all available slots + const buildLanes = (rules, kind) => { + if (rules.length) { + const tiers = []; + rules.forEach((r, i) => { + const src = normSrc(r.source); + const pair = src ? byPair.get(r.model + "|" + src) : null; + if (src && !pair) return; + const ref = pair || + byModel.get(r.model) || { src: "*", id: r.model }; + const ti = Math.max(1, parseInt(r.tier) || i + 1) - 1; + if (!tiers[ti]) tiers[ti] = []; + tiers[ti].push({ + src: ref.src || "*", + id: r.model, + uid: nexUid(), + meta: { + quota: r.token_quota || 0, + period: r.period || "", + hours: r.hours || 0, + }, + }); }); - }); - sortState.lanes = tiers - .filter(Boolean) - .map((models, i) => ({ prio: (tiers.length - i) * 10, models })); - } else { - sortState.lanes = [...sourceRows.entries()] + return tiers + .filter(Boolean) + .map((models, i) => ({ prio: (tiers.length - i) * 10, models })); + } + // discovery: group slots of this kind by their source priority + return [...sourceRows.entries()] .sort((a, b) => b[0] - a[0]) .map(([prio, models]) => { - models.sort((x, y) => - x.src < y.src ? -1 : x.src > y.src ? 1 : 0, - ); - models.forEach((m) => { - m.uid = nexUid(); - m.meta = null; + const same = models.filter((m) => { + const it = byPair.get(m.id + "|" + m.src); + const k = it ? it.kind : "chat"; + return k === kind; }); - return { prio, models }; - }); - } + if (!same.length) return null; + same.sort((x, y) => x.src < y.src ? -1 : x.src > y.src ? 1 : 0); + same.forEach((m) => { m.uid = nexUid(); m.meta = null; }); + return { prio, models: same }; + }) + .filter(Boolean); + }; + sortLanes.chat = buildLanes(autoR, "chat"); + sortLanes.image = buildLanes(autoImg, "image"); + sortState.lanes = sortLanes[sortState.kind] || []; // model picker for the "add slot" control: one option per (source, model) // pair so identical model ids on different sources stay distinguishable await loadModelPairs(); sortState.origin = JSON.stringify(sortState.lanes); + const kindToggle = + `
`; $("#tab-sort").innerHTML = ` -

${t("sortTitle")} +

${t("sortTitle")}${kindToggle}

@@ -2653,6 +2672,25 @@
`; paintSort(); } + function switchSortKind(k) { + if (sortState.kind === k) return; + sortLanes[sortState.kind] = sortState.lanes; + sortState.kind = k; + sortState.lanes = sortLanes[k] || []; + sortState.origin = JSON.stringify(sortState.lanes); + sortState.drag = null; + const h = $("#tab-sort").querySelector("h2"); + if (h) { + const tg = h.querySelector(".kind-toggle"); + if (tg) { + tg.querySelectorAll("button").forEach((b) => { + const isK = (b.getAttribute("onclick") || "").indexOf("'" + k + "'") >= 0; + b.className = isK ? "small" : "ghost small"; + }); + } + } + paintSort(); + } function healthTag(it) { const st = sortStateMap.get(it.id + "|" + (it.src || "*")); if (!st) return ""; @@ -3115,10 +3153,15 @@ }); }), ); + // each kind persists to its own chain; the other kind's saved rules + // are left untouched (omitting the field means "no change") + const body = { rules: undefined, image_rules: undefined }; + if (sortState.kind === "image") body.image_rules = rules; + else body.rules = rules; await api("/api/auto", { method: "PUT", headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ rules }), + body: JSON.stringify(body), }); } function scrAddModal() { @@ -3127,15 +3170,25 @@ wrap.id = "modal-wrap"; wrap.style.cssText = "position:fixed;inset:0;background:rgba(15,22,44,.45);display:flex;align-items:flex-start;justify-content:center;overflow:auto;padding:48px 20px;z-index:50"; - wrap.innerHTML = `

${t("sortAdd")}

+ // filter the model picker to the kind currently edited (image models + // in image mode, chat models in chat mode) + const kind = sortState.kind; + const pickModels = (allModels || []).filter((m) => { + const k = typeof m === "string" ? "chat" : m.kind || "chat"; + return kind === "image" ? k === "image" : k !== "image"; + }); + const optsHtml = pickModels.length + ? pickModels + .map((m) => + typeof m === "string" + ? `` + : ``, + ) + .join("") + : ``; + wrap.innerHTML = `

${t("sortAdd")} ${kind === "image" ? "· 生图" : "· 聊天"}

- + diff --git a/internal/scheduler/scheduler.go b/internal/scheduler/scheduler.go index 854591d..c4b7c8f 100644 --- a/internal/scheduler/scheduler.go +++ b/internal/scheduler/scheduler.go @@ -327,6 +327,63 @@ func (s *Scheduler) ChainChatStream(ctx context.Context, chain *Chain, req *type return chunks, src, model, err } +// ChainImage runs an image-generation AUTO request down the chain: tiers +// ascending (tier 1 highest priority), per-tier round-robin, same-tier order +// by preference. Each slot's model is pinned to its own image id (ModelFor), +// so a fallback switches per source. Cooling-down slots are skipped. Returns +// the response, serving source and the exact model id used; on total failure +// a *ChainErr summarizing every tier. +func (s *Scheduler) ChainImage(ctx context.Context, chain *Chain, req *types.ImageGenRequest) (*types.UnifiedResponse, string, string, error) { + if chain == nil || len(chain.Tiers) == 0 { + return nil, "", "", fmt.Errorf("no image auto slot configured") + } + var ce ChainErr + for _, tn := range chain.Tiers { + var cands []*Slot + for _, sl := range tn.Slots { + if !sl.Prov.ModelAvailable(sl.Model) { + continue + } + cands = append(cands, sl) + } + if len(cands) == 0 { + ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no schedulable slot (cooling)", tn.Tier)) + continue + } + sort.SliceStable(cands, func(i, j int) bool { + return cands[i].Prov.Pref(cands[i].Model) > cands[j].Prov.Pref(cands[j].Model) + }) + base := tn.NextStart() + var hard []TierError + for i := 0; i < len(cands); i++ { + sl := cands[(int(base)+i)%len(cands)] + if !sl.Prov.ModelAvailable(sl.Model) { + continue + } + r := *req + r.Model = sl.Model + resp, err := sl.Prov.Image(ctx, &r) + if ctx.Err() != nil { + return nil, "", "", ctx.Err() + } + if err == nil { + return resp, sl.Source, sl.Model, nil + } + if errors.Is(err, types.ErrBusy) { + continue + } + hard = append(hard, TierError{Tier: tn.Tier, Source: sl.Source, Model: sl.Model, Err: err}) + } + if len(hard) > 0 { + ce.Tiers = append(ce.Tiers, hard...) + } + } + if len(ce.Tiers) == 0 && len(ce.Skipped) == 0 { + return nil, "", "", fmt.Errorf("no image auto slot configured") + } + return nil, "", "", &ce +} + // ---- direct scheduling ---- // Chat runs a chat request across cands, falling back on failure. Each