feat(auto): quota-aware AUTO slot scheduling (model tiers, hour/week/month resets) + key scope periods; fix key/slot create/delete UX

internal/gateway/chat.go

@@ -9,6 +9,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"llmsproxy/internal/config"
 	"llmsproxy/internal/provider"
 	"llmsproxy/internal/scheduler"
 	"llmsproxy/internal/types"
@@ -106,11 +107,15 @@ func (g *Gateway) resolveCands(ctx context.Context, req *chatRequest) ([]*provid
 }
 
 // filterCandsByModels keeps only providers exposing at least one model of the
-// whitelist (used for user keys with a restricted model scope).
-func filterCandsByModels(cands []*provider.Provider, allow []string) []*provider.Provider {
+// scope (used for user keys with a restricted model scope). An "AUTO" scope
+// entry means the key is allowed to use any model.
+func filterCandsByModels(cands []*provider.Provider, allow []config.ModelScope) []*provider.Provider {
 	allowed := make(map[string]bool, len(allow))
 	for _, m := range allow {
-		allowed[m] = true
+		if m.Model == "" || strings.EqualFold(m.Model, "AUTO") {
+			return cands
+		}
+		allowed[m.Model] = true
 	}
 	out := make([]*provider.Provider, 0, len(cands))
 	for _, p := range cands {
@@ -124,11 +129,20 @@ func filterCandsByModels(cands []*provider.Provider, allow []string) []*provider
 	return out
 }
 
-// intersectModels restricts a model list to the whitelist (preserving order).
-func intersectModels(models, allow []string) []string {
+// intersectModels restricts a model list to the scope (preserving order). An
+// "AUTO" scope entry grants every model.
+func intersectModels(models []string, allow []config.ModelScope) []string {
 	allowed := make(map[string]bool, len(allow))
+	any := false
 	for _, m := range allow {
-		allowed[m] = true
+		if m.Model == "" || strings.EqualFold(m.Model, "AUTO") {
+			any = true
+			break
+		}
+		allowed[m.Model] = true
+	}
+	if any {
+		return models
 	}
 	out := make([]string, 0, len(models))
 	seen := map[string]bool{}
@@ -141,6 +155,60 @@ func intersectModels(models, allow []string) []string {
 	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" grants all models; its quota caps the key's total tokens.
+func (g *Gateway) checkModelScope(ctx context.Context, model string) string {
+	allow := g.allowedModels(ctx)
+	if allow == nil {
+		return ""
+	}
+	for _, sc := range allow {
+		if sc.Model != "" && strings.EqualFold(sc.Model, "AUTO") {
+			if sc.TokenQuota > 0 {
+				used := g.scopeTokens(ctx, sc)
+				if used >= sc.TokenQuota {
+					return fmt.Sprintf("token quota exceeded (%d/%d)", used, sc.TokenQuota)
+				}
+			}
+			return ""
+		}
+	}
+	for _, sc := range allow {
+		if sc.Model != model {
+			continue
+		}
+		if sc.TokenQuota > 0 {
+			used := g.scopeTokens(ctx, sc)
+			if used >= sc.TokenQuota {
+				return fmt.Sprintf("token quota exceeded for %q (%d/%d)", model, used, sc.TokenQuota)
+			}
+		}
+		return ""
+	}
+	return fmt.Sprintf("model %q is not allowed for this key", model)
+}
+
+// scopeTokens returns the tokens a scope entry has consumed within its reset
+// window (total for AUTO / per model otherwise).
+func (g *Gateway) scopeTokens(ctx context.Context, sc config.ModelScope) int64 {
+	k := keyID(reqKey(ctx))
+	if sc.Model != "" && strings.EqualFold(sc.Model, "AUTO") {
+		return g.stats.KeyTokens(k)
+	}
+	win := AutoPeriodSeconds(sc.Period, sc.Hours)
+	return g.stats.WindowTokens(sc.Model, win)
+}
+
+func hasScopeModel(list []config.ModelScope, s string) bool {
+	for _, x := range list {
+		if x.Model == s || (x.Model != "" && strings.EqualFold(x.Model, "AUTO")) {
+			return true
+		}
+	}
+	return false
+}
+
 func (g *Gateway) resolveByModel(model string) ([]*provider.Provider, string) {
 	if isAuto(model) {
 		return g.core.Registry().Resolve("AUTO"), ""
@@ -195,8 +263,44 @@ func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) {
 	if model == "" {
 		model = g.core.DefaultModel()
 	}
+	if strings.EqualFold(strings.TrimSpace(req.Model), "AUTO") {
+		plans := g.autoPlans()
+		if len(plans) == 0 {
+			writeError(w, http.StatusServiceUnavailable, "no_provider", "no auto slot available (quota exhausted or none configured)")
+			return
+		}
+		if msg := g.checkModelScope(r.Context(), "AUTO"); msg != "" {
+			writeError(w, http.StatusForbidden, "model_not_allowed", msg)
+			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,
+		}
+		rec := &Req{
+			Key:  keyID(reqKey(ctx)),
+			Type: "chat",
+			OK:   false,
+		}
+		if req.Stream {
+			rec.Type = "stream"
+			g.streamChatAuto(w, ctx, plans, inner, rec)
+			return
+		}
+		g.singleChatAuto(w, ctx, plans, inner, rec)
+		return
+	}
 	if !isAuto(model) {
-		if allow := g.allowedModels(r.Context()); allow != nil && !containsStr(allow, model) {
+		if allow := g.allowedModels(r.Context()); allow != nil && !hasScopeModel(allow, model) {
 			writeError(w, http.StatusForbidden, "model_not_allowed", fmt.Sprintf("model %q is not allowed for this key", model))
 			return
 		}
@@ -209,6 +313,10 @@ func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) {
 	if effective == "" {
 		effective = firstModel(cands[0])
 	}
+	if msg := g.checkModelScope(r.Context(), effective); msg != "" {
+		writeError(w, http.StatusForbidden, "model_not_allowed", msg)
+		return
+	}
 	ctx := r.Context()
 	done := g.stats.Begin()
 	defer done()
@@ -300,13 +408,20 @@ func firstSource(cands []*provider.Provider) string {
 	return ""
 }
 
-func containsStr(list []string, s string) bool {
-	for _, x := range list {
-		if x == s {
-			return true
+// 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 false
+	return model
 }
 
 func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) {
@@ -441,6 +556,193 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
 	}
 }
 
+// autoPlan is one schedulable AUTO slot: a model id pinned to its provider
+// with an optional token quota window. Quota-exhausted slots are skipped.
+type autoPlan struct {
+	p     *provider.Provider
+	model string
+	quota int64
+	win   int64
+}
+
+// autoPlans builds the schedulable AUTO slots from the persisted rules. A
+// slot is schedulable while its model is available and (when quota > 0) the
+// tokens used within its reset window are below the quota.
+func (g *Gateway) autoPlans() []autoPlan {
+	rules := g.core.AutoRules()
+	if len(rules) == 0 {
+		return nil
+	}
+	plans := make([]autoPlan, 0, len(rules))
+	for _, e := range rules {
+		p := g.core.ProviderForModel(e.Model)
+		if p == nil {
+			continue
+		}
+		if m := p.ModelByID(e.Model); m != nil && m.Kind == "image" {
+			continue
+		}
+		win := AutoPeriodSeconds(e.Period, e.Hours)
+		if e.TokenQuota > 0 && g.stats.WindowTokens(e.Model, win) >= e.TokenQuota {
+			continue
+		}
+		plans = append(plans, autoPlan{p: p, model: e.Model, quota: e.TokenQuota, win: win})
+	}
+	return plans
+}
+
+// singleChatAuto runs a non-streaming AUTO request slot by slot: each slot
+// pins its own model; a slot whose provider errors out is skipped. The first
+// slot to answer wins; when every slot fails, the recorded error is from the
+// last one.
+func (g *Gateway) singleChatAuto(w http.ResponseWriter, ctx context.Context, plans []autoPlan, req *types.ChatRequest, rec *Req) {
+	rec.LatMs = 0
+	t0 := time.Now()
+	var lastErr error
+	for _, pl := range plans {
+		if !pl.p.Available() {
+			continue
+		}
+		r := *req
+		r.Model = pl.model
+		resp, usedSrc, usedModel, err := g.core.Scheduler().Chat(ctx, scheduler.FromRegistry([]*provider.Provider{pl.p}), &r)
+		if err != nil {
+			lastErr = err
+			continue
+		}
+		rec.LatMs = time.Since(t0).Milliseconds()
+		rec.OK = true
+		rec.Status = http.StatusOK
+		rec.Prompt = int64(resp.TokenUsage.Prompt)
+		rec.Compl = int64(resp.TokenUsage.Completion)
+		rec.Source = usedSrc
+		rec.Model = usedModel
+		g.writeRec(rec)
+		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:   usedModel,
+			Choices: []ChatChoice{{Index: 0, Message: msg, FinishReason: resp.FinishReason}},
+		}
+		if resp.TokenUsage.Total > 0 || resp.TokenUsage.Prompt > 0 || resp.TokenUsage.Completion > 0 {
+			out.Usage = &resp.TokenUsage
+		}
+		writeJSON(w, http.StatusOK, out)
+		return
+	}
+	rec.LatMs = time.Since(t0).Milliseconds()
+	if lastErr == nil {
+		lastErr = fmt.Errorf("no provider available")
+	}
+	rec.OK = false
+	rec.Status = http.StatusBadGateway
+	rec.Err = lastErr.Error()
+	g.writeRec(rec)
+	writeError(w, http.StatusBadGateway, "upstream_error", lastErr.Error())
+}
+
+// streamChatAuto streams an AUTO request. It stays pinned to the first slot
+// whose stream begins; a slot that fails to connect is skipped.
+func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, plans []autoPlan, req *types.ChatRequest, 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)
+	}()
+	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: "auto",
+		Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{Role: "assistant"}}},
+	}) {
+		return
+	}
+	var lastErr error
+	for _, pl := range plans {
+		if !pl.p.Available() {
+			continue
+		}
+		r := *req
+		r.Model = pl.model
+		chunks, usedSrc, usedModel, err := g.core.Scheduler().ChatStream(ctx, scheduler.FromRegistry([]*provider.Provider{pl.p}), &r)
+		if err != nil {
+			lastErr = err
+			continue
+		}
+		if usedModel != "" {
+			rec.Model = usedModel
+			rec.Source = usedSrc
+		}
+		for ck := range chunks {
+			chunk := ChatChunk{
+				ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
+			}
+			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 {
+				stop := "stop"
+				choice.FinishReason = &stop
+			}
+			chunk.Choices = []ChunkChoice{choice}
+			rec.Compl += int64(len(ck.Content)+len(ck.ReasoningContent)) / 3
+			if !send(chunk) {
+				return
+			}
+		}
+		return
+	}
+	if lastErr == nil {
+		lastErr = fmt.Errorf("no provider available")
+	}
+	rec.OK = false
+	rec.Status = http.StatusBadGateway
+	rec.Err = lastErr.Error()
+	stop := "stop"
+	send(ChatChunk{
+		ID: id, Object: "chat.completion.chunk", Created: created, Model: "auto",
+		Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &stop}},
+	})
+	fmt.Fprintf(w, "data: [DONE]\n\n")
+	if flusher != nil {
+		flusher.Flush()
+	}
+}
+
 func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
 	if r.Method != http.MethodPost {
 		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use POST")
@@ -460,7 +762,7 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
 		model = g.core.DefaultModel()
 	}
 	if !isAuto(model) {
-		if allow := g.allowedModels(r.Context()); allow != nil && !containsStr(allow, model) {
+		if allow := g.allowedModels(r.Context()); allow != nil && !hasScopeModel(allow, model) {
 			writeError(w, http.StatusForbidden, "model_not_allowed", fmt.Sprintf("model %q is not allowed for this key", model))
 			return
 		}
@@ -474,6 +776,10 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
 		writeError(w, http.StatusServiceUnavailable, "no_provider", "no image source configured")
 		return
 	}
+	if msg := g.checkModelScope(r.Context(), effectiveImageModel(model, cands)); 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, Source: firstSource(cands), OK: false}