llmgw/internal/router/chat.go

package router

import (
	"bytes"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"slices"
	"strings"

	"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
	"gitea.mynodes.xyz/ilker/llmgw/internal/version"
	"gitea.mynodes.xyz/ilker/llmgw/internal/web"
)

// registerChat mounts the in-browser chat for signed-in people.
//
// It exists so someone can try a model without setting up a client, and — more
// often — so they can tell whether a problem is the gateway or their code.
//
// The browser never sees an API key. The page carries a session, the server
// turns that into the same identity an API key would produce, and the request
// goes down the ordinary path: admission, model allowlist, rate limits, budget,
// accounting. Anything else would make this page a hole in every limit the rest
// of the gateway enforces.
func (r *Router) registerChat(mux *http.ServeMux, pages *web.AccountPages) {
	mux.HandleFunc("GET /chat", pages.Guard(func(w http.ResponseWriter, req *http.Request, user store.User) {
		web.Chat(user, r.modelsFor(req, user), version.Version)(w, req)
	}))
	mux.HandleFunc("POST /chat/send", pages.Guard(r.chatSend))
	// Behind the same guard as the rest of the page: it takes a POST body and
	// returns a file, so it wants a session and the CSRF check for the same
	// reasons /chat/send does.
	mux.HandleFunc("POST /chat/download", pages.Guard(
		func(w http.ResponseWriter, req *http.Request, _ store.User) {
			web.CodeArchive(r.cfg.MaxBodyBytes)(w, req)
		}))
}

// chatSend dispatches one turn on behalf of a signed-in person.
func (r *Router) chatSend(w http.ResponseWriter, req *http.Request, user store.User) {
	principal, err := r.principalFor(req, user)
	if err != nil {
		http.Error(w, "your account has no project to bill this to", http.StatusForbidden)
		return
	}

	body, err := io.ReadAll(io.LimitReader(req.Body, r.cfg.MaxBodyBytes))
	if err != nil {
		http.Error(w, "could not read the request", http.StatusBadRequest)
		return
	}

	var in struct {
		Model    string            `json:"model"`
		Messages []json.RawMessage `json:"messages"`
	}
	if err := json.Unmarshal(body, &in); err != nil || in.Model == "" || len(in.Messages) == 0 {
		http.Error(w, "a model and at least one message are required", http.StatusBadRequest)
		return
	}

	// Rebuilt rather than forwarded, so the page cannot set fields it has no
	// business setting — a stream:false that would hang the reader, a
	// max_tokens large enough to occupy a replica for an hour.
	upstream, err := json.Marshal(map[string]any{
		"model":      in.Model,
		"messages":   in.Messages,
		"stream":     true,
		"max_tokens": chatMaxTokens,
	})
	if err != nil {
		http.Error(w, "could not build the request", http.StatusInternalServerError)
		return
	}

	next := req.Clone(auth.WithPrincipal(req.Context(), principal))
	next.Body = io.NopCloser(bytes.NewReader(upstream))
	next.ContentLength = int64(len(upstream))
	next.Header.Set("Content-Type", "application/json")

	// The same handler the API serves, so there is exactly one request path and
	// no second place for its rules to drift.
	r.handleInference("/v1/chat/completions")(w, next)
}

// chatMaxTokens bounds one turn from the page.
//
// A playground should not be the thing that ties up a replica for ten minutes,
// and someone who wants a longer answer has an API key for that.
//
// 512 was chosen when every model here was a 0.5B answering questions in a
// paragraph. It is about a screenful of code, so a coding model cut off
// mid-function on the first real request — and the page said nothing, because
// it never read finish_reason. The cap being too low was a nuisance; the reply
// ending silently was the actual defect, and both are fixed.
const chatMaxTokens = 2048

// principalFor builds the identity a signed-in person acts under.
//
// The key id is synthetic and marked as a session, because there is no key: it
// keeps usage rows attributable and rate limits per person, while making it
// obvious in the usage table that these requests came from the page rather than
// from a credential somebody holds.
func (r *Router) principalFor(req *http.Request, user store.User) (auth.Principal, error) {
	if r.store == nil || user.ProjectID == "" {
		return auth.Principal{}, errors.New("no project")
	}
	eff, err := r.store.ResolveProject(req.Context(), user.ProjectID)
	if err != nil {
		return auth.Principal{}, err
	}
	return auth.Principal{
		KeyID:          "session:" + user.ID,
		KeyPrefix:      "session",
		ProjectID:      eff.ProjectID,
		TeamID:         eff.TeamID,
		OrgID:          eff.OrgID,
		Band:           eff.Band,
		RoutingProfile: eff.RoutingProfile,
		Limits:         eff,
	}, nil
}

// modelsFor lists what this person may actually ask for.
//
// Served rather than catalogued: offering a model with no replica produces a
// 503 the reader cannot interpret, and the allowlist is applied so the picker
// never shows a model the request would refuse.
//
// Real models are ordered ahead of simulated ones so the browser's default
// selection — the first option — is always a model that answers. Sorted purely
// by name, "big-model" came first and every visitor who pressed Send without
// touching the picker got placeholder text.
func (r *Router) modelsFor(req *http.Request, user store.User) []web.ChatModel {
	principal, err := r.principalFor(req, user)
	if err != nil {
		return nil
	}
	seen := map[string]struct{}{}
	var out []web.ChatModel
	for _, s := range r.reg.Snapshots("") {
		for _, m := range s.Models {
			if _, dup := seen[m]; dup {
				continue
			}
			seen[m] = struct{}{}
			if r.permits(principal, m) {
				out = append(out, web.ChatModel{Name: m, Simulated: r.modelMap()[m].Simulated()})
			}
		}
	}
	slices.SortFunc(out, func(a, b web.ChatModel) int {
		if a.Simulated != b.Simulated {
			if b.Simulated {
				return -1
			}
			return 1
		}
		return strings.Compare(a.Name, b.Name)
	})
	return out
}