ModelRouter/internal/gateway/api.go

package gateway

import (
	"encoding/csv"
	"encoding/json"
	"io"
	"log"
	"net/http"
	"strconv"
	"strings"
	"time"

	"llmsproxy/internal/config"
)

type adapterPayload struct {
	Name string `json:"name"`
	Code string `json:"code"`
}

func (g *Gateway) handleAdaptersAPI(w http.ResponseWriter, r *http.Request) {
	if reqRole(r.Context()) != "admin" {
		writeError(w, http.StatusForbidden, "forbidden", "admin role required")
		return
	}
	path := strings.TrimPrefix(r.URL.Path, "/api/adapters")
	path = strings.Trim(path, "/")

	switch r.Method {
	case http.MethodGet:
		if path == "" {
			writeJSON(w, http.StatusOK, map[string]interface{}{"adapters": g.core.ListAdapters()})
			return
		}
		// A GET on a specific adapter is almost always a client that meant to
		// DELETE it but let fetch default to GET; say so instead of only
		// reporting that the code is not exposed.
		writeError(w, http.StatusNotFound, "not_found",
			"adapter code not exposed; edit in UI (to remove it use DELETE /api/adapters/"+path+")")
	case http.MethodPost:
		var p adapterPayload
		if err := json.NewDecoder(r.Body).Decode(&p); err != nil {
			writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
			return
		}
		if err := g.core.UploadAdapter(p.Name, p.Code); err != nil {
			writeError(w, http.StatusBadRequest, "adapter_error", err.Error())
			return
		}
		writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "name": p.Name})
	case http.MethodDelete:
		if path == "" {
			writeError(w, http.StatusBadRequest, "invalid_request", "adapter name required")
			return
		}
		if err := g.core.RemoveAdapter(path); err != nil {
			writeError(w, http.StatusBadRequest, "adapter_error", err.Error())
			return
		}
		writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true})
	default:
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
	}
}

// sourcePayload mirrors config.Source for JSON web UI editing.
type sourcePayload struct {
	Name          string                 `json:"name"`
	BaseURL       string                 `json:"base_url"`
	APIKey        string                 `json:"api_key"`
	Adapter       string                 `json:"adapter"`
	Endpoint      string                 `json:"endpoint"`
	ImageEndpoint string                 `json:"image_endpoint"`
	Models        []config.Model         `json:"models"`
	Headers       map[string]string      `json:"headers"`
	Meta          map[string]interface{} `json:"meta"`
	Temperature   float64                `json:"temperature"`
	MaxTokens     int                    `json:"max_tokens"`
	MaxConcurrent int                    `json:"max_concurrent"`
	RPM           int                    `json:"rpm"` // optional requests-per-minute cap, 0 = unlimited
	// ModelIDs is an additive-only shorthand: when set, the named models are
	// merged into an existing source instead of Models replacing the list.
	// It exists because "add one model" is the most common scripted edit and a
	// full Models list cannot be written without reading the source first.
	ModelIDs []string `json:"model_ids,omitempty"`
}

// keepExistingAPIKey is the mask a client sends when it means "keep the
// current credential". POST /api/sources is an upsert, so a script that edits
// one field and echoes a placeholder back would otherwise persist that
// placeholder as the live api_key — which fails closed as 401 on the next
// request, long after the script exited successfully. See TestSourceUpsertKeepsAPIKey.
const keepExistingAPIKey = "__KEEP__"

// resolveAPIKey returns the credential to store, inheriting the current one
// when the payload does not supply a new one.
//
// Two distinct "no new credential" signals exist and they must not collapse:
//   - the explicit mask keepExistingAPIKey — always inherit.
//   - the empty string — inherit only for a source that already has a
//     credential. Empty is a legitimate value for a *new* credential-less
//     source (dshcodebuddy uses "dsh-local-noauth", local servers use ""),
//     so a fresh source with an empty key must stay empty rather than inherit
//     nothing and fail differently.
func (g *Gateway) resolveAPIKey(name, want string, exists bool) string {
	if want == keepExistingAPIKey {
		return g.currentAPIKey(name)
	}
	if exists && want == "" {
		return g.currentAPIKey(name)
	}
	return want
}

func (g *Gateway) currentAPIKey(name string) string {
	if s := g.sourceByName(name); s != nil {
		return s.APIKey
	}
	return ""
}

// sourceByName returns a live view of one source, or nil when it does not
// exist. It goes through core.Sources() so the key is already unsealed and
// defaults applied, matching what the routes actually use.
func (g *Gateway) sourceByName(name string) *config.Source {
	for _, s := range g.core.Sources() {
		if s.Name == name {
			ss := s
			return &ss
		}
	}
	return nil
}

func (g *Gateway) sourceExists(name string) bool {
	return g.sourceByName(name) != nil
}

func (g *Gateway) handleSourcesAPI(w http.ResponseWriter, r *http.Request) {
	if reqRole(r.Context()) != "admin" {
		writeError(w, http.StatusForbidden, "forbidden", "admin role required")
		return
	}
	path := strings.TrimPrefix(r.URL.Path, "/api/sources")
	path = strings.Trim(path, "/")

	switch r.Method {
	case http.MethodGet:
		writeJSON(w, http.StatusOK, map[string]interface{}{"sources": g.core.Sources()})
	case http.MethodPost, http.MethodPut:
		var p sourcePayload
		body, _ := io.ReadAll(r.Body)
		if err := json.Unmarshal(body, &p); err != nil {
			writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
			return
		}
		// PUT addresses the source by path; the body name is only a fallback so
		// a single-field edit does not have to echo the name back. A body that
		// disagrees with the path is a caller bug, not something to silently
		// resolve — it usually means the wrong source is about to be written.
		if r.Method == http.MethodPut {
			if path == "" {
				writeError(w, http.StatusBadRequest, "invalid_request", "source name required in path")
				return
			}
			if p.Name != "" && p.Name != path {
				writeError(w, http.StatusBadRequest, "invalid_request",
					"source name in body ("+p.Name+") does not match path ("+path+")")
				return
			}
			p.Name = path
		}
		exists := g.sourceExists(p.Name)
		src := config.Source{
			Name:          p.Name,
			BaseURL:       p.BaseURL,
			APIKey:        g.resolveAPIKey(p.Name, p.APIKey, exists),
			Adapter:       p.Adapter,
			Endpoint:      p.Endpoint,
			ImageEndpoint: p.ImageEndpoint,
			Models:        p.Models,
			Headers:       p.Headers,
			Meta:          p.Meta,
			Temperature:   p.Temperature,
			MaxTokens:     p.MaxTokens,
			MaxConcurrent: p.MaxConcurrent,
			RPM:           p.RPM,
		}
		// model_ids is additive: "add these models" is the common scripted edit
		// and it must not require reading (and echoing) the whole list back.
		// A request that omits models entirely is therefore a pure add, not a
		// request to empty the list -- that distinction is the whole point.
		if len(p.ModelIDs) > 0 {
			if len(src.Models) == 0 {
				if existing := g.sourceByName(p.Name); existing != nil {
					src.Models = append(src.Models, existing.Models...)
				}
			}
			have := map[string]bool{}
			for _, m := range src.Models {
				have[m.ID] = true
			}
			for _, id := range p.ModelIDs {
				id = strings.TrimSpace(id)
				if id == "" || have[id] {
					continue
				}
				have[id] = true
				src.Models = append(src.Models, config.Model{ID: id, Kind: "chat"})
			}
		}
		if err := g.core.AddSource(src); err != nil {
			writeError(w, http.StatusBadRequest, "source_error", err.Error())
			return
		}
		writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true})
	case http.MethodDelete:
		if path == "" {
			writeError(w, http.StatusBadRequest, "invalid_request", "source name required")
			return
		}
		if err := g.core.RemoveSource(path); err != nil {
			writeError(w, http.StatusBadRequest, "source_error", err.Error())
			return
		}
		writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true})
	default:
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
	}
}

// csvHeaders stamps the shared download headers for CSV exports.
func csvHeaders(w http.ResponseWriter, filename string) {
	w.Header().Set("Content-Type", "text/csv; charset=utf-8")
	w.Header().Set("Content-Disposition", "attachment; filename="+filename)
}

// handleSourceTemplatesAPI manages reusable source templates (a Source minus
// name and api_key) stored in the runtime file so the WebUI can spin up
// multiple key-bearing sources from one shared template.
func (g *Gateway) handleSourceTemplatesAPI(w http.ResponseWriter, r *http.Request) {
	if reqRole(r.Context()) != "admin" {
		writeError(w, http.StatusForbidden, "forbidden", "admin role required")
		return
	}
	path := strings.TrimPrefix(r.URL.Path, "/api/source_templates")
	path = strings.Trim(path, "/")
	switch r.Method {
	case http.MethodGet:
		writeJSON(w, http.StatusOK, map[string]interface{}{"templates": g.core.Templates()})
	case http.MethodPost:
		var t config.SourceTemplate
		if err := json.NewDecoder(r.Body).Decode(&t); err != nil {
			writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
			return
		}
		if err := g.core.SaveTemplate(t); err != nil {
			writeError(w, http.StatusBadRequest, "template_error", err.Error())
			return
		}
		writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "templates": g.core.Templates()})
	case http.MethodDelete:
		if path == "" {
			writeError(w, http.StatusBadRequest, "invalid_request", "template name required")
			return
		}
		if err := g.core.RemoveTemplate(path); err != nil {
			writeError(w, http.StatusBadRequest, "template_error", err.Error())
			return
		}
		writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true})
	default:
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
	}
}

// exportKey resolves which masked key id an export covers: admins may pass
// any key filter, user keys are always scoped to themselves (records and
// aggregates are keyed by the masked keyID form).
func exportKey(r *http.Request) string {
	k := r.URL.Query().Get("key")
	if reqRole(r.Context()) != "admin" {
		k = keyID(reqKey(r.Context()))
	}
	return k
}

// handleStatsAPI returns per-key / per-model / per-source usage aggregates and
// the recent request audit trail.
func (g *Gateway) handleStatsAPI(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET")
		return
	}
	// Default to ONE SCREEN of records. The dashboard pages the rest through
	// /api/stats/records as the user scrolls, so a poll must not serialize the
	// whole ring buffer. An explicit limit is still honoured (capped by the ring).
	limit := firstScreenRecords
	if v := r.URL.Query().Get("limit"); v != "" {
		if n, err := strconv.Atoi(v); err == nil && n > 0 && n <= 20000 {
			limit = n
		}
	}
	key := exportKey(r)
	if r.URL.Query().Get("export") == "csv" {
		from, _ := strconv.ParseInt(r.URL.Query().Get("from"), 10, 64)
		to, _ := strconv.ParseInt(r.URL.Query().Get("to"), 10, 64)
		if to == 0 {
			to = time.Now().UnixMilli()
		}
		csvHeaders(w, "llmsproxy-requests.csv")
		cw := csv.NewWriter(w)
		names := map[string]string{}
		for _, k := range g.core.ListKeys() {
			names[keyID(k.Key)] = k.Name
		}
		_ = cw.Write([]string{"time", "key", "key_name", "type", "model", "source", "status", "ok", "prompt_tokens", "completion_tokens", "latency_ms", "first_byte_ms", "cache_hit_tokens", "cache_miss_tokens", "error"})
		// Stream the window straight to the client instead of materializing it:
		// an export covering months of audit data must not be bounded by RAM, and
		// nothing is retained once the response is written.
		flushEvery := 1000
		n := 0
		err := g.stats.StreamAuditRecords(from, to, key, func(rec Req) error {
			_ = cw.Write([]string{
				time.UnixMilli(rec.Time).Format(time.RFC3339),
				rec.Key,
				names[rec.Key],
				rec.Type,
				rec.Model,
				rec.Source,
				strconv.Itoa(rec.Status),
				strconv.FormatBool(rec.OK),
				strconv.FormatInt(rec.Prompt, 10),
				strconv.FormatInt(rec.Compl, 10),
				strconv.FormatInt(rec.LatMs, 10),
				strconv.FormatInt(rec.FirstByteMs, 10),
				strconv.FormatInt(rec.CacheHit, 10),
				strconv.FormatInt(rec.CacheMiss, 10),
				rec.Err,
			})
			n++
			if n%flushEvery == 0 {
				cw.Flush()
				if err := cw.Error(); err != nil {
					return err // client went away: stop walking the audit files
				}
				if f, ok := w.(http.Flusher); ok {
					f.Flush()
				}
			}
			return nil
		})
		cw.Flush()
		if err != nil {
			log.Printf("[gateway] csv export aborted after %d rows: %v", n, err)
		}
		return
	}
	if r.URL.Query().Get("export") == "keys-csv" {
		csvHeaders(w, "llmsproxy-keys.csv")
		cw := csv.NewWriter(w)
		_ = cw.Write([]string{"key", "key_name", "role", "models", "total_requests", "success_requests", "failed_requests", "prompt_tokens", "completion_tokens", "total_tokens", "avg_latency_ms", "max_latency_ms", "created_at"})
		keyFilter := exportKey(r)
		// by_key rows are keyed by the masked id (keyID); build a masked-id ->
		// record map so names/roles/models resolve for the export.
		byMasked := map[string]config.GWKey{}
		for _, k := range g.core.ListKeys() {
			byMasked[keyID(k.Key)] = k
		}
		snap := g.stats.Snapshot(0, keyFilter)
		byKeyRaw, ok := snap["by_key"].([]StatsRow)
		if !ok {
			byKeyRaw = []StatsRow{}
		}
		for _, row := range byKeyRaw {
			keyInfo, found := byMasked[row.Name]
			name, role, models := "", "", ""
			if found {
				name = keyInfo.Name
				role = keyInfo.Role
				modelNames := make([]string, 0, len(keyInfo.Models))
				for _, m := range keyInfo.Models {
					modelNames = append(modelNames, m.Model)
				}
				models = strings.Join(modelNames, ",")
			}
			avgLat := int64(0)
			if row.Stat.Reqs > 0 {
				avgLat = row.Stat.LatSum / row.Stat.Reqs
			}
			createdAt := ""
			if found && keyInfo.CreatedAt > 0 {
				createdAt = time.Unix(keyInfo.CreatedAt, 0).Format(time.RFC3339)
			}
			_ = cw.Write([]string{
				row.Name, name, role, models,
				strconv.FormatInt(row.Stat.Reqs, 10),
				strconv.FormatInt(row.Stat.OK, 10),
				strconv.FormatInt(row.Stat.Err, 10),
				strconv.FormatInt(row.Stat.Prompt, 10),
				strconv.FormatInt(row.Stat.Compl, 10),
				strconv.FormatInt(row.Stat.Prompt+row.Stat.Compl, 10),
				strconv.FormatInt(avgLat, 10),
				strconv.FormatInt(row.Stat.LatMax, 10),
				createdAt,
			})
		}
		cw.Flush()
		return
	}
	snap := g.stats.Snapshot(limit, key)
	keyNames := map[string]string{}
	for _, k := range g.core.ListKeys() {
		keyNames[keyID(k.Key)] = k.Name
	}
	snap["key_names"] = keyNames
	writeJSON(w, http.StatusOK, snap)
}

// handleStatsRecordsAPI pages the request records straight off the audit files.
// The dashboard loads only its first screen and asks for the next page as the
// user scrolls, so neither side holds the full history: the server keeps no
// per-client state (the cursor is the whole state, and it lives in the URL) and
// closes every file handle before responding.
//
//	 GET /api/stats/records?before=<cursor>&limit=100[&key=<masked id>]
//	-> { records: [...newest first...], next_cursor, has_more, rotated }
func (g *Gateway) handleStatsRecordsAPI(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET")
		return
	}
	limit := 100
	if v := r.URL.Query().Get("limit"); v != "" {
		if n, err := strconv.Atoi(v); err == nil && n > 0 && n <= 1000 {
			limit = n
		}
	}
	// exportKey pins non-admin callers to their own key, so a user key can never
	// page another key's records.
	page := g.stats.AuditPage(r.URL.Query().Get("before"), limit, exportKey(r))
	keyNames := map[string]string{}
	for _, k := range g.core.ListKeys() {
		keyNames[keyID(k.Key)] = k.Name
	}
	writeJSON(w, http.StatusOK, map[string]interface{}{
		"records":     page.Records,
		"next_cursor": page.Next,
		"has_more":    page.HasMore,
		"rotated":     page.Rotated,
		"key_names":   keyNames,
	})
}