ModelRouter/internal/gateway/discover_api.go

package gateway

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"modelrouter/internal/config"
	"net/http"
	"strings"
	"time"
)

// HTTP surface for model discovery.
//
// Two endpoints, deliberately separate:
//
//	POST /api/sources/discover        — probe a URL+key, return the catalog.
//	                                   Nothing is written.
//	POST /api/sources/discover/apply  — probe, then write the source.
//
// The split exists because discovery mutates nothing and is therefore safe to
// call repeatedly while the operator is still typing a key, whereas apply
// rewrites credentials and the AUTO chain's inputs. Making them one endpoint
// with a flag would mean a "just check it" click can silently create a source.
//
// Both are admin-only, matching /api/sources: the request body carries an API
// key, and a model catalog reveals which upstreams an operator uses.

// discoverRequest is the body of both endpoints.
type discoverRequest struct {
	// BaseURL is what the operator pasted. It may or may not include a
	// version segment; discovery works that out rather than making the
	// caller know this gateway's path-composition rules.
	BaseURL string `json:"base_url"`
	APIKey  string `json:"api_key"`

	// Name is required for apply, ignored for discover. It defaults from the
	// base URL's host when omitted, because "https://api.vendor.com/v1" is a
	// better starting name than "new-source".
	Name string `json:"name"`

	// Adapter is optional. Discovery assumes "openai" because that is what a
	// /v1/models catalog implies; a non-OpenAI upstream (Anthropic, Gemini)
	// has no such endpoint, so discovery will legitimately fail for those and
	// the operator configures them by hand.
	Adapter string `json:"adapter"`

	// Temperature and MaxConcurrent are carried into the written source.
	Temperature   float64 `json:"temperature"`
	MaxConcurrent int     `json:"max_concurrent"`

	// ReplaceModels drops the source's existing model list before applying the
	// discovered one. It defaults to true for a NEW source and is required to
	// be spelled out for an existing one, so a re-discovery cannot silently
	// delete hand-tuned priorities.
	ReplaceModels *bool `json:"replace_models,omitempty"`

	// Select narrows what gets written. Empty means everything discovered.
	// Ids that are not in the catalog are an error rather than a silent skip:
	// a typo here would produce a source that 404s on every request.
	Include []string `json:"include,omitempty"`
	Exclude []string `json:"exclude,omitempty"`

	// ImageModels names ids that are **kind=image** but whose upstream catalog
	// does not say so.
	//
	// 为什么需要（真实缺口）：discoveredModelKind 只看上游的 modality 字段，
	// 并且刻意**不从模型名推断**（那会把 "gpt-image-1" 这类同名聊天模型
	// 错归成生图，而错误的 kind 会毒死 AUTO 链）。
	// 但有些网关的 `/v1/models` 只返回 `{id,object,created,owned_by}` ——
	// 零 modality（实测：某聚合网关的 3 个 gpt-image-* 全部落成 chat，
	// 于是生图端点永远报 `no image source configured`）。
	//
	// 两条原则都不破：不猜（默认仍完全按元数据），但允许**显式声明**。
	// 只影响写入的 kind，不影响目录里有没有这个模型。
	ImageModels []string `json:"image_models,omitempty"`
}

// discoverResponse is the JSON the WebUI binds to.
type discoverResponse struct {
	OK bool `json:"ok"`

	// Normalized source fields, ready to be echoed back to /api/sources.
	BaseURL     string   `json:"base_url"`
	Adapter     string   `json:"adapter"`
	Models      []string `json:"models"`
	ImageModels []string `json:"image_models"`

	// Endpoint is set only when discovery had to move a version segment.
	Endpoint string `json:"endpoint"`
	// ImageEndpoint is required whenever base_url carries a version segment:
	// the gateway's default image path is the ABSOLUTE "/v1/images/generations"
	// while chat's default is the relative "/chat/completions", so the same
	// base_url that yields a working chat URL yields a 404 image URL.
	ImageEndpoint string `json:"image_endpoint"`

	// Diagnostics for the operator.
	ModelsPath string   `json:"models_path"`
	Tried      []string `json:"tried"`
	LatencyMS  int64    `json:"latency_ms"`
	Note       string   `json:"note,omitempty"`
	Warnings   []string `json:"warnings,omitempty"`

	// Set by apply only.
	SourceName string `json:"source_name,omitempty"`
	Applied    bool   `json:"applied,omitempty"`
}

func (g *Gateway) handleSourceDiscoverAPI(w http.ResponseWriter, r *http.Request) {
	if reqRole(r.Context()) != "admin" {
		writeError(w, http.StatusForbidden, "forbidden", "admin role required")
		return
	}
	apply := strings.HasSuffix(r.URL.Path, "/apply")
	if r.Method != http.MethodPost {
		w.Header().Set("Allow", "POST")
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed",
			"use POST with {\"base_url\":…,\"api_key\":…}")
		return
	}

	body, err := io.ReadAll(io.LimitReader(r.Body, 1<<20))
	if err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request", "cannot read body")
		return
	}
	var req discoverRequest
	if err := json.Unmarshal(body, &req); err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
		return
	}
	req.BaseURL = strings.TrimSpace(req.BaseURL)
	req.Name = strings.TrimSpace(req.Name)

	// ★ 只给 `name`（或两者都给但 key 为空）⇒ **从已存源取 base_url 与凭据**。
	//
	// 为什么必须有这条路径（真实痛点）：刷新一个**已存在**源的模型目录，此前
	// 必须让运维把明文 key 再粘一次 —— 而库里的 key 是 `enc:v1:` 密文，
	// 运维手里未必还有原文；WebUI 也不回显明文（这是正确的安全设计）。
	// 结果就是「想刷新目录 → 必须重新申请/粘贴一次 key」，很多人干脆不刷新，
	// 于是源的模型目录越放越旧（实例：某源的 3 个生图模型从未被登记）。
	//
	// 安全边界：① 已要求 admin 角色；② 只能探测**本网关已配置的**源，
	// 不能借此拿任意 URL + 任意 key 做探测（那样就成了 SSRF/密钥外带渠道）；
	// ③ 响应不回显 key（只回目录）。
	if req.Name != "" {
		if old := g.sourceByName(req.Name); old != nil {
			if req.BaseURL == "" {
				req.BaseURL = old.BaseURL
			}
			if req.APIKey == "" {
				req.APIKey = old.APIKey
			}
			if req.Adapter == "" {
				req.Adapter = old.Adapter
			}
		}
	}

	if req.BaseURL == "" {
		writeError(w, http.StatusBadRequest, "invalid_request",
			"base_url required (or pass name of an existing source to reuse its base_url and credential)")
		return
	}
	if req.Adapter == "" {
		req.Adapter = "openai"
	}

	// Bounded so a hung upstream cannot pin the handler; discovery has its own
	// per-request timeout on top.
	ctx, cancel := context.WithTimeout(r.Context(), 60*time.Second)
	defer cancel()

	res, err := DiscoverModels(ctx, req.BaseURL, req.APIKey)
	if err != nil {
		// A rejected key must not be reported as an upstream outage: the two
		// need opposite actions from the operator, and 502 sends them to the
		// wrong one.
		status := http.StatusBadGateway
		var f *discoveryFailure
		if errors.As(err, &f) {
			status = f.StatusOrDefault()
		}
		writeError(w, status, "discover_failed", err.Error())
		return
	}

	// ★ 显式生图声明：把 operator 点名的 id 从 chat 挪到 image。
	//
	// 放在 selectModels **之前**：后者按 include/exclude 裁剪两个列表，
	// 挪类必须在裁剪前完成，否则一个"只声明成生图、又被 include 选中"的
	// 模型会先在 chat 里被保留、再被这里挪走 —— 顺序反了会得到两份。
	//
	// 只认**目录里确实存在的** id：写错一个字不该静默造出一个不存在的模型
	// （那会让源在每次照片请求时 404），而是明确报错。
	if len(req.ImageModels) > 0 {
		have := map[string]bool{}
		for _, m := range res.Models {
			have[m] = true
		}
		for _, m := range res.ImageModels {
			have[m] = true
		}
		var missing []string
		force := map[string]bool{}
		for _, id := range req.ImageModels {
			if !have[id] {
				missing = append(missing, id)
				continue
			}
			force[id] = true
		}
		if len(missing) > 0 {
			writeError(w, http.StatusBadRequest, "invalid_request",
				"image_models contains ids that are not in the upstream catalog: "+
					strings.Join(missing, ", "))
			return
		}
		var keepChat []string
		for _, m := range res.Models {
			if !force[m] {
				keepChat = append(keepChat, m)
			}
		}
		seenImg := map[string]bool{}
		var newImage []string
		for _, m := range res.ImageModels {
			seenImg[m] = true
			newImage = append(newImage, m)
		}
		for id := range force {
			if !seenImg[id] {
				newImage = append(newImage, id)
			}
		}
		res.Models = keepChat
		res.ImageModels = newImage
	}

	chat, image, warns, err := selectModels(res.Models, res.ImageModels, req.Include, req.Exclude)
	if err != nil {
		// 400: the operator asked for a model the catalog does not have.
		writeError(w, http.StatusBadRequest, "invalid_request", err.Error())
		return
	}

	out := &discoverResponse{
		OK:            true,
		BaseURL:       res.BaseURL,
		Adapter:       req.Adapter,
		Models:        chat,
		ImageModels:   image,
		Endpoint:      res.Endpoint,
		ImageEndpoint: res.ImageEndpoint,
		ModelsPath:    res.ModelsPath,
		Tried:         res.Tried,
		LatencyMS:     res.LatencyMS,
		Note:          res.Note,
		Warnings:      warns,
	}
	if len(chat) == 0 && len(image) == 0 {
		out.Warnings = append(out.Warnings,
			"every discovered model was excluded; applying this would create a source with no models")
	}

	if !apply {
		writeJSON(w, http.StatusOK, out)
		return
	}

	name := strings.TrimSpace(req.Name)
	if name == "" {
		name = defaultSourceName(req.BaseURL)
	}
	if name == "" {
		writeError(w, http.StatusBadRequest, "invalid_request",
			"name required: cannot derive one from base_url")
		return
	}
	if strings.ContainsAny(name, "/\\ \t") || name == "." || name == ".." {
		writeError(w, http.StatusBadRequest, "invalid_request",
			"source name must not contain slashes or whitespace")
		return
	}

	exists := g.sourceExists(name)
	if exists && req.ReplaceModels == nil {
		// Refusing here is deliberate. Re-running discovery on an existing
		// source and having it silently replace a hand-tuned model list —
		// including priorities chosen for the AUTO chain — is the kind of
		// "helpful" behaviour that loses an afternoon.
		writeError(w, http.StatusConflict, "conflict",
			"source "+name+" already exists; set replace_models=true to overwrite its model list, "+
				"or choose a different name")
		return
	}

	// Build the source directly rather than round-tripping through the
	// /api/sources payload: resolveAPIKey carries the "an empty api_key means
	// keep the stored one" rule, and going through a second decoder here would
	// duplicate (and eventually diverge from) that logic.
	src := config.Source{
		Name:          name,
		BaseURL:       res.BaseURL,
		APIKey:        g.resolveAPIKey(name, req.APIKey, exists),
		Adapter:       req.Adapter,
		Endpoint:      res.Endpoint,
		ImageEndpoint: res.ImageEndpoint,
		Models:        modelsToConfig(chat, image),
		Temperature:   req.Temperature,
	}
	if req.MaxConcurrent > 0 {
		src.MaxConcurrent = req.MaxConcurrent
	}

	// An existing source may carry fields discovery does not know about
	// (proxy_url, headers, api_key_env, timeout). Rebuilding from scratch would
	// drop them, so they are inherited unless the request overrides them.
	if exists {
		old := g.sourceByName(name)
		if old != nil {
			src.ProxyURL = old.ProxyURL
			src.APIKeyEnv = old.APIKeyEnv
			src.Timeout = old.Timeout
			src.QueueTimeout = old.QueueTimeout
			src.Headers = old.Headers
			src.Meta = old.Meta
			src.MaxTokens = old.MaxTokens
			src.RPM = old.RPM
			// image_endpoint is NOT inherited when discovery supplied one:
			// the stored value was correct for the OLD base_url, and the base
			// is exactly what may have changed. Inheriting it would keep a
			// stale value that 404s against the new base.
			if res.ImageEndpoint == "" {
				src.ImageEndpoint = old.ImageEndpoint
			}
			if src.Adapter == "" {
				src.Adapter = old.Adapter
			}
			if req.MaxConcurrent <= 0 {
				src.MaxConcurrent = old.MaxConcurrent
			}
			// replace_models=false keeps the curated list. Discovery still
			// refreshed base_url/endpoint/credential, which is the point of
			// re-running it.
			if req.ReplaceModels != nil && !*req.ReplaceModels {
				src.Models = old.Models
				out.Warnings = append(out.Warnings,
					"model list left unchanged (replace_models=false)")
			}
		}
	}

	if err := g.core.AddSource(src); err != nil {
		writeError(w, http.StatusBadRequest, "source_error", err.Error())
		return
	}

	out.SourceName = name
	out.Applied = true
	writeJSON(w, http.StatusOK, out)
}

// selectModels applies the include/exclude filters.
//
// An include entry the catalog does not have is an error: silently dropping it
// produces a source that looks configured and 404s at request time, which is
// far harder to diagnose than a rejected apply. An exclude entry that is not
// present is fine — the operator is allowed to be defensive.
func selectModels(chat, image, include, exclude []string) ([]string, []string, []string, error) {
	have := map[string]bool{}
	for _, m := range chat {
		have[m] = true
	}
	for _, m := range image {
		have[m] = true
	}

	if len(include) > 0 {
		var missing []string
		for _, want := range include {
			if !have[want] {
				missing = append(missing, want)
			}
		}
		if len(missing) > 0 {
			return nil, nil, nil, &discoverError{
				msg: "these models are not in the upstream catalog: " + strings.Join(missing, ", "),
			}
		}
		keep := map[string]bool{}
		for _, want := range include {
			keep[want] = true
		}
		var c, i []string
		for _, m := range chat {
			if keep[m] {
				c = append(c, m)
			}
		}
		for _, m := range image {
			if keep[m] {
				i = append(i, m)
			}
		}
		chat, image = c, i
	}

	if len(exclude) > 0 {
		drop := map[string]bool{}
		for _, m := range exclude {
			drop[m] = true
		}
		var c, i []string
		for _, m := range chat {
			if !drop[m] {
				c = append(c, m)
			}
		}
		for _, m := range image {
			if !drop[m] {
				i = append(i, m)
			}
		}
		chat, image = c, i
	}

	// chat and image must stay non-nil: the JSON contract forbids null here.
	if chat == nil {
		chat = []string{}
	}
	if image == nil {
		image = []string{}
	}

	var warns []string
	if len(chat) == 0 && len(image) == 0 {
		// Caller can only get here by excluding everything, or from a catalog
		// that was already empty. Either way "silence" would read as success.
		warns = append(warns,
			"no models selected; applying this would create a source that cannot serve any request")
		return chat, image, warns, nil
	}
	if len(image) > 0 && len(chat) == 0 {
		warns = append(warns,
			"every model is an image model; AUTO chat requests have nothing to select")
	}
	if len(chat) > 0 && len(image) > 0 {
		warns = append(warns,
			"image models are only reachable through /v1/images/generations, not AUTO chat")
	}
	return chat, image, warns, nil
}

// discoverError is a client mistake rather than an upstream failure, which is
// why it gets its own type instead of fmt.Errorf.
type discoverError struct{ msg string }

func (e *discoverError) Error() string { return e.msg }

// defaultSourceName derives a source name from a base URL host, stripping the
// version segment and any "api." prefix. It returns "" for input that is not a
// URL at all — the caller must not fall back to echoing garbage into a source
// name, since the name becomes a config key and an API path segment.
func defaultSourceName(baseURL string) string {
	s := strings.TrimSpace(baseURL)
	if !strings.Contains(s, "://") {
		return ""
	}
	s = s[strings.Index(s, "://")+3:]
	if i := strings.IndexAny(s, "/?#"); i >= 0 {
		s = s[:i]
	}
	if i := strings.LastIndex(s, ":"); i > 0 && !strings.Contains(s[i+1:], "]") {
		s = s[:i] // strip port
	}
	// A host must still look like one after trimming the cosmetic prefixes.
	s = strings.TrimPrefix(s, "api.")
	s = strings.TrimPrefix(s, "www.")
	if s == "" || strings.ContainsAny(s, " \t\"'<>") {
		return ""
	}
	for _, r := range s {
		switch {
		case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9':
		case r == '.' || r == '-' || r == '_':
		default:
			return ""
		}
	}
	if !strings.Contains(s, ".") {
		// A bare label is far more likely to be a mistake than a hostname.
		return ""
	}
	return s
}