llmgw/internal/router/threads.go

package router

import (
	"context"
	"encoding/json"
	"errors"

	"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"

// 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 == "" || r.store == nil {
		return 0, nil
	}
	if req.ProjectID == "" {
		// No identity, so no project to scope the thread to. Storing
		// conversations for an unauthenticated gateway would put every
		// caller's history in one pool, readable by anyone who guesses an id.
		return 0, errThreadsNeedAuth
	}

	// 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.store.ThreadRetention(ctx, req.ProjectID)
	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.store.Thread(ctx, req.ProjectID, 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.store == nil || reply == "" {
		return
	}
	turns := append([]store.Message(nil), req.ThreadTurns...)
	turns = append(turns, store.Message{Role: "assistant", Content: reply, Model: model})

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

// 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")