Know how big a model's context is, and what it can do

Three things the gateway did not know about its own models, all of which
surfaced this week as somebody else's bug.

Context length. Nothing read it, /v1/models did not publish it, and an
over-long prompt was not refused: it was admitted, scheduled, given a
replica and a slot, and rejected by the engine in whatever shape that
engine chose. A client sizing prompts from /v1/models had nothing to read.
It is now declared per model, published when known, and checked before
authentication work or a queue slot is spent — with a ten percent margin,
because the estimate is bytes over four and being a few percent out must
not turn a valid request into a refusal.

Capabilities. The gateway forwards a tools array to any model, and one that
cannot use them answers in prose; every OpenAI-compatible client then
behaves as though the model is broken, and the only way to find out was to
try it. Declared rather than discovered, because no engine will say — which
is why these are claims, and why the next commit makes engine-check verify
the tool-calling one instead of trusting it.

And a bug found while wiring the first: the token estimate was taken from
the routing prefix. That prefix stops at 2 KB by design, so every prompt
above it estimated as 512 tokens however long it really was, and TPM limits
undercounted exactly the requests that consume the most. Prefix routing
wants a bounded head of the conversation; a token estimate wants all of it.
They are now separate functions, and structured multimodal content counts
rather than measuring zero.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

internal/control/catalog/catalog.go

@@ -87,6 +87,39 @@ type Model struct {
 	// RoutingProfile selects how the router scores endpoints for this model.
 	RoutingProfile string `koanf:"routing_profile"`
 	SLO            SLO    `koanf:"slo"`
+
+	// ContextLength is how many tokens one request may use, prompt and
+	// completion together, and MaxOutput bounds the completion alone. Zero
+	// means unknown, which is what every model was until now.
+	//
+	// The gateway knew neither, so an over-long prompt was not refused: it was
+	// admitted, given a replica, and failed at the engine, and a client sizing
+	// its prompts from /v1/models had nothing to read. Worse, the number is not
+	// the one an operator writes in the engine's arguments — llama.cpp divides
+	// --ctx-size across its parallel slots, so a pool of 8192 across four slots
+	// is 2048 a request, and nothing said so.
+	ContextLength int `koanf:"context_length"`
+	MaxOutput     int `koanf:"max_output"`
+
+	// Capabilities is what this model can do beyond plain completion.
+	Capabilities Capabilities `koanf:"capabilities"`
+}
+
+// Capabilities describes what a client may ask a model for.
+//
+// Declared by the operator rather than discovered, because there is no way to
+// ask an engine what its model can do. That makes these claims, and a wrong
+// claim is worse than none — which is why `ctl engine-check` verifies tool
+// calling against a running engine rather than trusting this.
+//
+// The reason it exists: the gateway forwards a tools array to any model, and a
+// model that cannot use them answers in prose. Every OpenAI-compatible client
+// then behaves as though the model is broken, and the only way to find out was
+// to try it. That cost an afternoon.
+type Capabilities struct {
+	Tools      bool `koanf:"tools"`
+	Vision     bool `koanf:"vision"`
+	Embeddings bool `koanf:"embeddings"`
 }
 
 // Catalog is the set of models the cluster should serve.

internal/router/context_limit.go

@@ -0,0 +1,60 @@
+package router
+
+import (
+	"fmt"
+
+	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
+)
+
+// checkContextFits reports why a request cannot fit in a model's context, or
+// "" when it can.
+//
+// Refused here rather than at the engine. An over-long prompt used to be
+// admitted, scheduled, given a replica and a slot, and only then rejected —
+// paying the whole cost of the request to learn something knowable from the
+// request itself. The engine's error also arrives in whatever shape that engine
+// chose, which is not something a client can act on.
+//
+// The prompt estimate is bytes over four, the same crude stand-in used for
+// scoring, so the comparison is approximate in one direction: it can let
+// through a request the engine will refuse. That is the right way round. A
+// tokenizer would make it exact, and needing one is why plan §7.5 tier 2 is v2.
+// To keep an estimate that reads slightly low from refusing valid requests, a
+// margin is allowed before anything is rejected.
+func checkContextFits(m catalog.Model, promptTokens, maxTokens int) string {
+	if m.ContextLength <= 0 {
+		// Unknown, which is every model that has not declared one. Saying
+		// nothing beats guessing a limit and refusing work over it.
+		return ""
+	}
+
+	if m.MaxOutput > 0 && maxTokens > m.MaxOutput {
+		return fmt.Sprintf("max_tokens is %d, but %q allows at most %d",
+			maxTokens, m.Name, m.MaxOutput)
+	}
+
+	// Only the prompt is estimated; max_tokens is exact when given. A request
+	// that does not say how much it wants back is judged on its prompt alone,
+	// because assuming the whole remaining window would refuse most requests.
+	need := promptTokens + maxTokens
+	allowed := m.ContextLength + contextEstimateMargin(m.ContextLength)
+	if need <= allowed {
+		return ""
+	}
+	if maxTokens > 0 {
+		return fmt.Sprintf(
+			"about %d prompt tokens plus max_tokens %d exceeds the %d-token context of %q",
+			promptTokens, maxTokens, m.ContextLength, m.Name)
+	}
+	return fmt.Sprintf("the prompt is about %d tokens, over the %d-token context of %q",
+		promptTokens, m.ContextLength, m.Name)
+}
+
+// contextEstimateMargin is the slack allowed for the bytes-over-four estimate.
+//
+// That estimate is wrong in both directions depending on the text — code and
+// non-Latin scripts tokenize very differently from English prose — and being
+// wrong by a few percent must not turn a valid request into a refusal. Ten
+// percent is enough to absorb that while still catching a prompt that is
+// genuinely too big, which is the case worth catching.
+func contextEstimateMargin(context int) int { return context / 10 }

internal/router/context_limit_test.go

@@ -0,0 +1,105 @@
+package router
+
+import (
+	"encoding/json"
+	"strings"
+	"testing"
+
+	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
+)
+
+func TestAModelWithNoDeclaredContextRefusesNothing(t *testing.T) {
+	// Every model was in this state until now, and guessing a limit would turn
+	// working requests into errors.
+	m := catalog.Model{Name: "m"}
+	if why := checkContextFits(m, 1_000_000, 1_000_000); why != "" {
+		t.Errorf("refused a request against an unknown context: %q", why)
+	}
+}
+
+func TestAPromptOverTheContextIsRefused(t *testing.T) {
+	m := catalog.Model{Name: "coder", ContextLength: 2048}
+	why := checkContextFits(m, 5000, 0)
+	if why == "" {
+		t.Fatal("a prompt well over the context was allowed")
+	}
+	for _, want := range []string{"5000", "2048", "coder"} {
+		if !strings.Contains(why, want) {
+			t.Errorf("the message does not mention %s: %q", want, why)
+		}
+	}
+}
+
+func TestPromptPlusMaxTokensMustFitTogether(t *testing.T) {
+	// The failure people actually hit: a prompt that fits, asking for a
+	// completion that cannot follow it.
+	m := catalog.Model{Name: "coder", ContextLength: 2048}
+	if why := checkContextFits(m, 1800, 1000); why == "" {
+		t.Error("prompt plus max_tokens over the window was allowed")
+	}
+	if why := checkContextFits(m, 1000, 500); why != "" {
+		t.Errorf("a request that fits was refused: %q", why)
+	}
+}
+
+func TestTheEstimateIsGivenRoomToBeWrong(t *testing.T) {
+	// Bytes over four is wrong in both directions — code and non-Latin scripts
+	// tokenize nothing like English prose — and being a few percent out must
+	// not turn a valid request into a refusal.
+	m := catalog.Model{Name: "m", ContextLength: 1000}
+	if why := checkContextFits(m, 1050, 0); why != "" {
+		t.Errorf("a request just over the estimate was refused: %q", why)
+	}
+	if why := checkContextFits(m, 2000, 0); why == "" {
+		t.Error("a prompt at twice the context was allowed; the margin is too wide")
+	}
+}
+
+func TestMaxOutputIsCheckedOnItsOwn(t *testing.T) {
+	// Asking for more output than the model can produce is worth catching even
+	// when the total would fit.
+	m := catalog.Model{Name: "m", ContextLength: 32768, MaxOutput: 4096}
+	why := checkContextFits(m, 10, 8000)
+	if why == "" {
+		t.Fatal("max_tokens over max_output was allowed")
+	}
+	if !strings.Contains(why, "4096") {
+		t.Errorf("the message does not say the limit: %q", why)
+	}
+}
+
+func TestTheTokenEstimateCoversTheWholePrompt(t *testing.T) {
+	// It used to be taken from the routing prefix, which stops at prefixBytes,
+	// so every prompt over 2 KB estimated as 512 tokens however long it was.
+	// TPM limits therefore undercounted exactly the requests that consume most,
+	// and no context check could work.
+	long := strings.Repeat("x", 40_000)
+	b := &inferenceBody{Messages: []struct {
+		Role    string          `json:"role"`
+		Content json.RawMessage `json:"content"`
+	}{{Role: "user", Content: json.RawMessage(`"` + long + `"`)}}}
+
+	got := promptBytes(b)
+	if got < 40_000 {
+		t.Fatalf("promptBytes = %d for a 40,000 byte prompt", got)
+	}
+	if prefix := len(promptPrefix(b)); prefix > prefixBytes {
+		t.Errorf("the routing prefix grew past its bound: %d", prefix)
+	}
+	// The two must stay different things: bounded for routing, whole for limits.
+	if got <= prefixBytes {
+		t.Error("the estimate is still bounded by the routing prefix")
+	}
+}
+
+func TestStructuredContentIsCounted(t *testing.T) {
+	// Multimodal messages carry an array of parts rather than a string. Counting
+	// zero for those would let a large request through unmeasured.
+	b := &inferenceBody{Messages: []struct {
+		Role    string          `json:"role"`
+		Content json.RawMessage `json:"content"`
+	}{{Role: "user", Content: json.RawMessage(`[{"type":"text","text":"` + strings.Repeat("y", 5000) + `"}]`)}}}
+	if got := promptBytes(b); got < 5000 {
+		t.Errorf("structured content counted as %d bytes", got)
+	}
+}

internal/router/handlers.go

@@ -195,6 +195,17 @@ func (r *Router) handleInference(path string) http.HandlerFunc {
 			return
 		}
 
+		// Before authentication work and before admission: a request that
+		// cannot fit should not consume a quota reservation or a queue slot on
+		// its way to being refused.
+		if _, entry := r.resolve(parsed.Model); entry.ContextLength > 0 {
+			if why := checkContextFits(entry, parsed.PromptTokens, parsed.MaxTokens); why != "" {
+				writeError(w, http.StatusBadRequest, "invalid_request_error", why)
+				r.m.observeRequest(parsed.Model, "", http.StatusBadRequest, time.Since(start), nil, 0)
+				return
+			}
+		}
+
 		if p, ok := auth.FromContext(req.Context()); ok {
 			if !r.permits(p, parsed.Model) {
 				auth.WriteForbiddenModel(w, parsed.Model)
@@ -373,8 +384,10 @@ func (r *Router) parse(req *http.Request, path string) (*parsedRequest, error) {
 	}
 
 	// Bytes over four: a crude stand-in for a tokenizer, good enough for
-	// scoring. It becomes exact only with plan §7.5 tier 2 (see plan §14).
-	p.PromptTokens = len(p.PromptPrefix) / 4
+	// scoring and for limits. It becomes exact only with plan §7.5 tier 2 (see
+	// plan §14). Measured over the whole prompt, not the routing prefix —
+	// see promptBytes.
+	p.PromptTokens = promptBytes(&b) / 4
 
 	// Ask the engine for usage the client did not request, so accounting works
 	// regardless of client behaviour; the extra chunk is dropped on the way back.
@@ -432,6 +445,40 @@ func promptPrefix(b *inferenceBody) string {
 	return s
 }
 
+// promptBytes measures the whole prompt, which promptPrefix deliberately does
+// not.
+//
+// The two were the same function for too long, and the estimate was taken from
+// the prefix: bytes over four of at most prefixBytes, so every prompt above
+// 2 KB counted as 512 tokens however long it really was. Prefix routing wants a
+// bounded head of the conversation; a token estimate wants all of it, and using
+// one for the other meant TPM limits undercounted exactly the requests that
+// consume the most, and nothing could tell whether a prompt fit in a context.
+func promptBytes(b *inferenceBody) int {
+	n := 0
+	add := func(raw json.RawMessage) {
+		if len(raw) == 0 {
+			return
+		}
+		var str string
+		if err := json.Unmarshal(raw, &str); err == nil {
+			n += len(str)
+			return
+		}
+		// Not a plain string: structured content, an array of parts, an image.
+		// The raw JSON overstates the text slightly, which is the safe
+		// direction for a limit.
+		n += len(raw)
+	}
+	for _, m := range b.Messages {
+		n += len(m.Role) + 2
+		add(m.Content)
+	}
+	add(b.Prompt)
+	add(b.Input)
+	return n
+}
+
 // injectIncludeUsage sets stream_options.include_usage on the outgoing body.
 //
 // This is the one place the router rewrites a request, so it is done by
@@ -556,9 +603,23 @@ func (r *Router) handleModels(w http.ResponseWriter, _ *http.Request) {
 		if r.models[n].Simulated() {
 			owner = "llmgw-sim"
 		}
-		data = append(data, map[string]any{
+		entry := map[string]any{
 			"id": n, "object": "model", "created": time.Now().Unix(), "owned_by": owner,
-		})
+		}
+		// Only when known. A context_length of 0 would read as "no context",
+		// which is a worse answer than saying nothing.
+		m := r.models[n]
+		if m.ContextLength > 0 {
+			entry["context_length"] = m.ContextLength
+		}
+		if m.MaxOutput > 0 {
+			entry["max_output_tokens"] = m.MaxOutput
+		}
+		entry["capabilities"] = map[string]bool{
+			"tools": m.Capabilities.Tools, "vision": m.Capabilities.Vision,
+			"embeddings": m.Capabilities.Embeddings,
+		}
+		data = append(data, entry)
 	}
 	writeJSON(w, http.StatusOK, map[string]any{"object": "list", "data": data})
 }

internal/router/router_test.go

@@ -415,3 +415,89 @@ func TestTheErrorIsHonestWhenNothingIsServed(t *testing.T) {
 		t.Errorf("an empty fleet is not explained: %q", msg)
 	}
 }
+
+func TestAnOverLongRequestIsRefusedBeforeItReachesAnEngine(t *testing.T) {
+	// The whole point of checking here: an over-long prompt used to be
+	// admitted, scheduled, given a replica and a slot, and refused by the
+	// engine — paying the cost of the request to learn something knowable from
+	// the request itself.
+	e := newFakeEngine(t, "m")
+	cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}})
+	m := catModel("m", nil)
+	m.ContextLength = 2048
+	cfg.Models = []catalog.Model{m}
+	rt := newTestRouter(t, cfg)
+
+	before := e.calls.Load()
+	rec := post(t, rt.Handler(), map[string]any{
+		"model":    "m",
+		"messages": []any{map[string]any{"role": "user", "content": strings.Repeat("x", 40_000)}},
+	}, nil)
+
+	if rec.Code != http.StatusBadRequest {
+		t.Fatalf("status %d, expected 400; body %s", rec.Code, rec.Body.String())
+	}
+	if got := e.calls.Load(); got != before {
+		t.Errorf("the engine was called %d times for a request that could not fit", got-before)
+	}
+	if !strings.Contains(rec.Body.String(), "2048") {
+		t.Errorf("the error does not name the limit: %s", rec.Body.String())
+	}
+}
+
+func TestModelsEndpointPublishesContextAndCapabilities(t *testing.T) {
+	// A client sizing its prompts had nothing to read, and a client deciding
+	// whether to send tools had nothing either.
+	e := newFakeEngine(t, "m")
+	cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}})
+	m := catModel("m", nil)
+	m.ContextLength, m.MaxOutput = 32768, 4096
+	m.Capabilities = catalog.Capabilities{Tools: true}
+	cfg.Models = []catalog.Model{m}
+	rt := newTestRouter(t, cfg)
+
+	req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
+	rec := httptest.NewRecorder()
+	rt.Handler().ServeHTTP(rec, req)
+
+	var out struct {
+		Data []struct {
+			ID           string          `json:"id"`
+			Context      int             `json:"context_length"`
+			MaxOutput    int             `json:"max_output_tokens"`
+			Capabilities map[string]bool `json:"capabilities"`
+		} `json:"data"`
+	}
+	if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil {
+		t.Fatal(err)
+	}
+	var found bool
+	for _, d := range out.Data {
+		if d.ID != "m" {
+			continue
+		}
+		found = true
+		if d.Context != 32768 || d.MaxOutput != 4096 {
+			t.Errorf("context %d, max output %d", d.Context, d.MaxOutput)
+		}
+		if !d.Capabilities["tools"] || d.Capabilities["vision"] {
+			t.Errorf("capabilities wrong: %v", d.Capabilities)
+		}
+	}
+	if !found {
+		t.Fatal("the model was not listed")
+	}
+}
+
+func TestAModelWithNoContextPublishesNoContextField(t *testing.T) {
+	// Zero would read as "no context", which is worse than saying nothing.
+	e := newFakeEngine(t, "m")
+	rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}}))
+
+	req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
+	rec := httptest.NewRecorder()
+	rt.Handler().ServeHTTP(rec, req)
+	if strings.Contains(rec.Body.String(), "context_length") {
+		t.Errorf("an unknown context was published as a number: %s", rec.Body.String())
+	}
+}