Passages of one organization never share a vLLM prefix cache with another's

A request given passages carries a cache_salt derived from its organization
when the replica chosen is vLLM, so how fast an answer starts cannot say what
another tenant's documents held. Other engines are not sent the field.

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

BACKLOG.md

@@ -975,10 +975,8 @@ the model's context; added to the one leading system message (or a new
 one); `x-llmgw-sources`; before guardInput; never cached; metrics. Default collections built (kb_defaults on org/team/project/key, narrowest
 wins, 30 s cache, `ctl kb default`, console action; never refuse); the
 portal chat's Documents picker (collection, default or none; sources listed
-under the answer; a header the page sends itself is dropped). Left:
-`cache_salt` per organization (not sent: llama.cpp may refuse an unknown
-field, and the router does not know the engine before choosing the
-replica); returned passages
+under the answer; a header the page sends itself is dropped). `cache_salt` per organization on requests given passages, sent to vLLM
+replicas only (chosen replica's engine). Left: returned passages
 are not audited per request (the header and the kept request record them).
 
 ### F4. Adapter training (L, after F1–F3)

docs/knowledge.md

@@ -316,6 +316,10 @@ Behaviour to know about:
 - **Other APIs.** Requests through `/v1/responses` and `/v1/messages`
   (Codex, Claude Code) work the same way, since they are answered as chat
   completions.
+- **Separate prefix caches.** On vLLM replicas, a request given passages
+  carries a `cache_salt` derived from its organization. One organization's
+  passages then never share cached blocks with another's, so how fast an
+  answer starts says nothing about what another tenant asked.
 - **No answer cache.** Answers from knowledge bases are never served from
   the answer cache, because two keys of one project can read different
   collections.

internal/router/handlers.go

@@ -409,6 +409,9 @@ type parsedRequest struct {
 	guardIn guard.Verdict
 	// knowledge is set when passages from knowledge bases were added.
 	knowledge bool
+	// kbOrg is the organization whose passages were added, for the prefix
+	// cache's salt.
+	kbOrg string
 	// noStore is a policy that keeps nothing of this request.
 	noStore bool
 	// inCluster is a project whose requests may not leave the cluster.

internal/router/rag.go

@@ -3,6 +3,8 @@ package router
 import (
 	"bytes"
 	"context"
+	"crypto/sha256"
+	"encoding/hex"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -223,6 +225,7 @@ func (r *Router) retrieve(ctx context.Context, w http.ResponseWriter, req *parse
 	req.UpstreamBody = body
 	req.PromptTokens = len(body) / 4
 	req.knowledge = true
+	req.kbOrg = e.OrgID
 	// Set now: the proxy writes the headers when the answer starts.
 	w.Header().Set(sourcesHeader, strings.Join(sources, ","))
 	result := "found"
@@ -425,3 +428,19 @@ func (r *Router) kbDefaults(ctx context.Context, p auth.Principal, project strin
 	r.kbDefMu.Unlock()
 	return names
 }
+
+// withCacheSalt sets vLLM's cache_salt to one derived from an organization:
+// the same for its requests, so they still share a cache among themselves.
+func withCacheSalt(body []byte, org string) []byte {
+	var top map[string]json.RawMessage
+	if json.Unmarshal(body, &top) != nil {
+		return body
+	}
+	sum := sha256.Sum256([]byte("llmgw-kb\x00" + org))
+	top["cache_salt"], _ = json.Marshal(hex.EncodeToString(sum[:16]))
+	out, err := json.Marshal(top)
+	if err != nil {
+		return body
+	}
+	return out
+}

internal/router/rag_test.go

@@ -9,6 +9,7 @@ import (
 	"strings"
 	"testing"
 
+	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
 	"gitea.mynodes.xyz/ilker/llmgw/internal/kb/kbsync"
 	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
 )
@@ -51,6 +52,20 @@ func TestPassagesJoinTheLeadingSystemMessage(t *testing.T) {
 	}
 }
 
+func TestTheCacheSaltIsAnOrganizations(t *testing.T) {
+	body := []byte(`{"model":"m","messages":[],"temperature":0}`)
+	a1, a2, b := withCacheSalt(body, "org-a"), withCacheSalt(body, "org-a"), withCacheSalt(body, "org-b")
+	var x, y struct {
+		Salt string  `json:"cache_salt"`
+		T    float64 `json:"temperature"`
+	}
+	_ = json.Unmarshal(a1, &x)
+	_ = json.Unmarshal(b, &y)
+	if string(a1) != string(a2) || x.Salt == "" || x.Salt == y.Salt || len(x.Salt) != 32 || !strings.Contains(string(a1), `"temperature":0`) {
+		t.Errorf("salts %s / %s", a1, b)
+	}
+}
+
 func TestTheLastUserMessageIsTheQuestion(t *testing.T) {
 	body := `{"messages":[{"role":"user","content":"first"},{"role":"assistant","content":"a"},
 		{"role":"user","content":[{"type":"text","text":"yıllık izin"},{"type":"image_url","image_url":{"url":"x"}}]}]}`
@@ -142,6 +157,11 @@ func TestAChatIsAnsweredFromItsCollections(t *testing.T) {
 		t.Errorf("the model was sent:\n%s", sent)
 	}
 
+	// The engine here is no vLLM: no salt is sent to it.
+	if strings.Contains(string(sent), "cache_salt") {
+		t.Errorf("a salt was sent to an engine that is not vLLM:\n%s", sent)
+	}
+
 	// A project default answers a request that names nothing, and none
 	// turns it off.
 	if err := s.SetKBDefault(ctx, store.KBDefault{Scope: store.ScopeProject, ScopeID: proj,
@@ -181,3 +201,55 @@ func mustID(t *testing.T, s *store.DB, name string) string {
 	t.Fatalf("no collection %s", name)
 	return ""
 }
+
+func TestAVLLMReplicaGetsTheOrganizationsSaltOnlyWithPassages(t *testing.T) {
+	dsn := os.Getenv("LLMGW_TEST_POSTGRES")
+	if dsn == "" {
+		t.Skip("set LLMGW_TEST_POSTGRES to a pgvector database to run this")
+	}
+	e := newFakeEngine(t, "qwen2.5-14b-instruct")
+	cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Engine: "vllm", Models: []string{"qwen2.5-14b-instruct"}})
+	cfg.Auth.Enabled = true
+	cfg.Auth.Store = dsn
+	rt := newTestRouter(t, cfg)
+	if err := rt.store.KnowledgeAvailable(); err != nil {
+		t.Skip(err)
+	}
+	ctx := t.Context()
+	s := rt.store
+	sfx := strings.ToLower(store.NewID(""))
+	org, _ := s.CreateOrg(ctx, "salt-"+sfx)
+	team, _ := s.CreateTeam(ctx, org, "t")
+	proj, _ := s.CreateProject(ctx, team, "p")
+	c, err := s.CreateCollection(ctx, store.KBCollection{OwnerScope: store.ScopeProject, OwnerID: proj,
+		Name: "c-" + sfx, EmbedModel: "qwen2.5-14b-instruct", Dims: 3})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, _, err := s.PutDocument(ctx, c, store.KBDocument{Name: "a.md", Digest: "d"},
+		[]store.KBChunk{{Body: "yıllık izin yirmi gün", Embedding: []float32{1, 1, 1}}}); err != nil {
+		t.Fatal(err)
+	}
+	key := kbKey(t, s, proj, "")
+	ask := func(knowledge string) string {
+		req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(
+			`{"model":"qwen2.5-14b-instruct","messages":[{"role":"user","content":"yıllık izin"}]}`))
+		req.Header.Set("Authorization", "Bearer "+key)
+		if knowledge != "" {
+			req.Header.Set(knowledgeHeader, knowledge)
+		}
+		rec := httptest.NewRecorder()
+		rt.Handler().ServeHTTP(rec, req)
+		if rec.Code != http.StatusOK {
+			t.Fatalf("%d %s", rec.Code, rec.Body.String())
+		}
+		b, _ := e.last.Load().([]byte)
+		return string(b)
+	}
+	if got := ask(c.Name); !strings.Contains(got, `"cache_salt":"`) {
+		t.Errorf("no salt with passages:\n%s", got)
+	}
+	if got := ask(""); strings.Contains(got, "cache_salt") {
+		t.Errorf("a salt without passages:\n%s", got)
+	}
+}

internal/router/router.go

@@ -1079,6 +1079,14 @@ func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *pa
 		if isAdapter && !onDemand && dec.Endpoint.Engine != "vllm" {
 			body = withAdapter(body, model, ad.index, ad.scale)
 		}
+		// Passages of one organization's documents never share cached
+		// blocks with another's: a prefix found warm would say what another
+		// tenant asked about, by how fast its answer started. vLLM keys its
+		// prefix cache by this salt; other engines are not sent a field
+		// they may refuse.
+		if req.kbOrg != "" && dec.Endpoint.Engine == "vllm" && req.ContentType == "" {
+			body = withCacheSalt(body, req.kbOrg)
+		}
 		// What computes the answer, said before the proxy writes the headers.
 		// A model can have replicas on cards and on CPUs at once, and the two
 		// differ tenfold in speed: a reply that came slowly is otherwise