llmgw/internal/router/threads.go

package router

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"strings"

	"gitea.mynodes.xyz/ilker/llmgw/internal/auth"

	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// threadHeader names the conversation a request continues.
//
// A header rather than a body field so the JSON stays exactly what an OpenAI
// client would send: some SDKs drop fields they do not know, and every one of
// them lets a caller set a header. It sits with the other x-llmgw-* headers
// the router already reads.
const threadHeader = "x-llmgw-thread"

// maxThreadID bounds a thread id, which is a key in the store chosen by the
// client.
const maxThreadID = 128

// validThreadID refuses an id too long to be a name, or with characters no
// generated id has.
//
// A thread is shared by every key of its project: whoever knows the id reads
// the conversation. So an id should be unguessable — a UUID — and not
// "chat-1".
func validThreadID(id string) error {
	if len(id) > maxThreadID {
		return fmt.Errorf("%s is %d characters; at most %d", threadHeader, len(id), maxThreadID)
	}
	for _, c := range id {
		if !(c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' || c >= '0' && c <= '9' || strings.ContainsRune("-_.:", c)) {
			return fmt.Errorf("%s may hold letters, digits and - _ . : only", threadHeader)
		}
	}
	return nil
}

// trimmedHeader reports how many stored turns were left out to make the
// conversation fit the model that served it.
//
// Dropping turns is the right answer — a thread that outgrows a small model
// would otherwise be permanently unusable on it, and after a budget fallback
// that is not a choice the caller made — but it must not be a silent one.
const trimmedHeader = "x-llmgw-thread-trimmed"

// threadMessage is one turn as it goes to an engine.
type threadMessage struct {
	Role    string          `json:"role"`
	Content json.RawMessage `json:"content"`
}

// spliceThread puts a stored conversation in front of the request's own
// messages and returns the body to send upstream.
//
// The history goes first and the request's messages last, because the request
// carries the new turn: a client using a thread sends what it just said, not
// the whole conversation again. That is the point of the feature.
//
// budget is how many bytes of prompt the serving model can take, measured the
// same way promptBytes measures it so the two cannot disagree. Zero means the
// model declares no context length, and then nothing is trimmed — guessing a
// limit for a model that never stated one would drop turns for no reason.
func spliceThread(body []byte, history []store.Message, budget int) ([]byte, int, error) {
	if len(history) == 0 {
		return body, 0, nil
	}

	// Decoded into a map so every field the client sent survives untouched.
	// The router forwards bodies it does not understand, and a thread must not
	// be the thing that starts dropping them.
	var raw map[string]json.RawMessage
	if err := json.Unmarshal(body, &raw); err != nil {
		return nil, 0, err
	}

	var own []threadMessage
	if m, ok := raw["messages"]; ok {
		if err := json.Unmarshal(m, &own); err != nil {
			return nil, 0, err
		}
	}

	stored := make([]threadMessage, 0, len(history))
	for _, m := range history {
		c, err := json.Marshal(m.Content)
		if err != nil {
			return nil, 0, err
		}
		stored = append(stored, threadMessage{Role: m.Role, Content: c})
	}

	dropped := 0
	if budget > 0 {
		stored, dropped = trimToFit(stored, own, budget)
	}

	merged, err := json.Marshal(append(stored, own...))
	if err != nil {
		return nil, 0, err
	}
	raw["messages"] = merged

	out, err := json.Marshal(raw)
	if err != nil {
		return nil, 0, err
	}
	return out, dropped, nil
}

// trimToFit drops the oldest stored turns until the conversation fits.
//
// The request's own messages are never dropped: they are what the caller just
// said, and a request that cannot fit without them is one to refuse rather
// than to answer with something the caller did not ask.
//
// System messages are kept wherever they are. They are usually instructions
// that shape every answer, and dropping one silently changes what the model is
// for — which is a stranger failure than a forgotten turn.
func trimToFit(stored, own []threadMessage, budget int) ([]threadMessage, int) {
	fixed := messagesBytes(own)
	for _, m := range stored {
		if m.Role == "system" {
			fixed += messageBytes(m)
		}
	}

	// Oldest first, which is the order they are in.
	keep := make([]bool, len(stored))
	for i := range stored {
		keep[i] = true
	}
	used := fixed
	for _, m := range stored {
		if m.Role != "system" {
			used += messageBytes(m)
		}
	}

	dropped := 0
	for i := range stored {
		if used <= budget {
			break
		}
		if stored[i].Role == "system" {
			continue
		}
		used -= messageBytes(stored[i])
		keep[i] = false
		dropped++
	}

	out := make([]threadMessage, 0, len(stored))
	for i, m := range stored {
		if keep[i] {
			out = append(out, m)
		}
	}
	return out, dropped
}

// messageBytes measures one turn the way promptBytes measures it, so a
// conversation trimmed to fit is measured by the same ruler that decides
// whether it fits.
func messageBytes(m threadMessage) int {
	n := len(m.Role) + 2
	var s string
	if err := json.Unmarshal(m.Content, &s); err == nil {
		return n + len(s)
	}
	return n + len(m.Content)
}

func messagesBytes(ms []threadMessage) int {
	n := 0
	for _, m := range ms {
		n += messageBytes(m)
	}
	return n
}

// ownMessages reads the turns a request carries, for storing after a reply.
func ownMessages(body []byte) []store.Message {
	var b struct {
		Messages []threadMessage `json:"messages"`
	}
	if err := json.Unmarshal(body, &b); err != nil {
		return nil
	}
	out := make([]store.Message, 0, len(b.Messages))
	for _, m := range b.Messages {
		var s string
		if err := json.Unmarshal(m.Content, &s); err != nil {
			// Structured content — an image, an array of parts. Stored as the
			// JSON it arrived as, so replaying it gives the model what the
			// client originally sent rather than a flattened version of it.
			s = string(m.Content)
		}
		out = append(out, store.Message{Role: m.Role, Content: s})
	}
	return out
}

// maxStoredReply bounds what one turn may add to a conversation.
//
// A cap rather than no cap because this is memory held on the request path for
// the length of a reply. Generous enough that ordinary answers are stored
// whole, small enough that a runaway generation cannot cost the gateway more
// than this per concurrent request.
const maxStoredReply = 128 << 10

// captureBytes says how much of the reply the proxy should hand back.
//
// Zero for every request not continuing a conversation, which is what keeps
// the cost of this feature off the path of requests that do not use it.
func (p *parsedRequest) captureBytes() int {
	if p.ThreadID == "" {
		return 0
	}
	return maxStoredReply
}

// loadThread splices a stored conversation into the request.
//
// Returns how many turns were dropped to make it fit, or an error to refuse
// the request with. A thread on a project that keeps no conversations is
// refused rather than ignored: ignoring it would send the model a conversation
// with no history, and it would answer confidently about something it was
// never told — which reads as the model being wrong rather than as a setting
// being off.
func (r *Router) loadThread(ctx context.Context, req *parsedRequest) (int, error) {
	if req.ThreadID == "" {
		return 0, nil
	}
	if err := validThreadID(req.ThreadID); err != nil {
		return 0, err
	}
	project, err := r.threadProject(req.ProjectID)
	if err != nil {
		return 0, err
	}
	req.ThreadProject = project

	// Asked before anything else, because a project that keeps no
	// conversations must be told so rather than served. Without this the
	// request succeeds, nothing is stored, and the caller finds out only when
	// the model has no idea what they are referring to — which reads as the
	// model being broken. The e2e test that caught this asserted the refusal
	// and got a cheerful 200.
	retention, err := r.threads.ThreadRetention(ctx, project)
	if err != nil {
		return 0, err
	}
	if retention == 0 {
		return 0, store.ErrThreadsNotEnabled
	}

	// What the caller sent, before the history is spliced in, so the exchange
	// can be stored afterwards without re-reading a body that has changed.
	req.ThreadTurns = ownMessages(req.UpstreamBody)

	history, err := r.threads.Thread(ctx, project, req.ThreadID)
	if err != nil && !errors.Is(err, store.ErrNoThread) {
		return 0, err
	}
	// A thread that does not exist yet is the first turn of a new one, which is
	// the ordinary way to start.
	if len(history) == 0 {
		return 0, nil
	}

	_, entry := r.resolve(req.Model)
	budget := 0
	if entry.ContextLength > 0 {
		// The same four-bytes-per-token estimate the rest of the router uses,
		// less the output the caller asked for, so trimming and the fit check
		// cannot disagree about what fits.
		budget = (entry.ContextLength - req.MaxTokens) * 4
		if budget < 0 {
			budget = 0
		}
	}

	body, dropped, err := spliceThread(req.UpstreamBody, history, budget)
	if err != nil {
		return 0, err
	}
	req.UpstreamBody = body
	req.PromptTokens = len(body) / 4
	return dropped, nil
}

// saveThread stores the exchange once a reply has been produced.
//
// Best effort, and deliberately after the client has been served: a
// conversation that failed to record is a worse outcome than a reply that
// never arrived, but only just, and the caller has their answer either way.
// The failure is logged rather than swallowed, because a thread silently not
// growing is the kind of fault someone discovers days later.
func (r *Router) saveThread(ctx context.Context, req *parsedRequest, reply string, model string) {
	if req.ThreadID == "" || r.threads == nil || req.ThreadProject == "" || reply == "" {
		return
	}
	turns := append([]store.Message(nil), req.ThreadTurns...)
	turns = append(turns, store.Message{Role: "assistant", Content: reply, Model: model})

	if err := r.threads.AppendThread(ctx, req.ThreadProject, req.ThreadID, turns); err != nil {
		r.log.Warn("could not store a conversation turn",
			"thread", req.ThreadID, "project", req.ThreadProject, "err", err)
	}
}

// anonymousProject owns everything kept by a gateway with no authentication.
const anonymousProject = "_anonymous"

// errThreadsNeedAuth is refused rather than ignored, for the same reason a
// disabled project is.
var errThreadsNeedAuth = errors.New(
	"a conversation needs an API key: without one there is no project to keep it in")

// errThreadsPersonal is what a personal key meets. Conversations and notes are
// kept per project, with the project's retention; a personal key has none, and
// the shared pool an unauthenticated gateway uses is not somewhere a signed-in
// person's history belongs.
var errThreadsPersonal = errors.New(
	"conversations and notes are kept per project, and this is a personal key; " +
		"use a key made in a project that keeps them")

// errThreadsOff is what a thread header meets when nothing is configured to
// keep threads. Refused rather than ignored: ignoring it sends the model a
// conversation with no history, and it answers confidently about something it
// was never told.
var errThreadsOff = errors.New(
	"this gateway keeps no conversations: enable auth with a project that keeps " +
		"threads, or set threads.store and threads.retain_hours")

// threadProject says which project a request's threads and notes belong to.
//
// The caller's own project when authenticated. The shared anonymous one only
// when the operator configured it, because with no identity every caller's
// history lands in one pool readable by anyone who knows an id.
func (r *Router) threadProject(callerProject string) (string, error) {
	switch {
	case r.threads == nil:
		return "", errThreadsOff
	case callerProject != "":
		return callerProject, nil
	case r.auth != nil:
		return "", errThreadsPersonal
	case r.anonymous:
		return anonymousProject, nil
	default:
		return "", errThreadsNeedAuth
	}
}

// handleDeleteThread forgets a conversation.
//
// A client endpoint rather than an admin one: the person who said the thing is
// the person who should be able to unsay it, and making them ask an operator
// would mean the operator reads it first.
func (r *Router) handleDeleteThread(w http.ResponseWriter, req *http.Request) {
	id := req.PathValue("id")
	if r.threads == nil {
		writeError(w, http.StatusNotFound, "invalid_request_error", "no such thread")
		return
	}
	var caller string
	if p, ok := auth.FromContext(req.Context()); ok {
		caller = p.ProjectID
	}
	project, err := r.threadProject(caller)
	if err != nil {
		writeError(w, http.StatusUnauthorized, "invalid_request_error",
			"deleting a conversation needs the key that owns it")
		return
	}

	// The store scopes the delete to the project, so a thread belonging to
	// another one is "no such thread" — the same answer as one that never
	// existed, because two answers would let a caller map which ids are real
	// in projects they cannot read.
	if err := r.threads.DeleteThread(req.Context(), project, id); err != nil {
		if errors.Is(err, store.ErrNoThread) {
			writeError(w, http.StatusNotFound, "invalid_request_error", "no such thread")
			return
		}
		writeError(w, http.StatusInternalServerError, "server_error", err.Error())
		return
	}
	w.WriteHeader(http.StatusNoContent)
}