feat(llm-api-gateway): add prompt cache key to chat completions

docs/user/llm-gateway.md

@@ -117,6 +117,7 @@ curl -sS -X POST "http://${GATEWAY_ADDR}/v1/chat/completions" \
   -d '{
     "model": "<function-id>/dummy-model",
     "stream": true,
+    "prompt_cache_key": "nvcf-summary-session",
     "messages": [
       {
         "role": "user",
@@ -185,18 +186,44 @@ The LLM Gateway supports sticky routing for multi-turn OpenAI-compatible request
 
 Sticky routing is not supported on `/v1/embeddings`.
 
-To keep later requests routed to the same backend, send the `x-multi-turn-session-id` response header value back as the `x-multi-turn-session-id` request header on the next request.
+To identify related requests, set `prompt_cache_key` in the request body. You
+can also send the `x-multi-turn-session-id` response header value back as the
+`x-multi-turn-session-id` request header on the next request.
 
 The gateway chooses the sticky routing key in this order:
 
 | Endpoint | Precedence |
 | --- | --- |
 | `/v1/responses` | `prompt_cache_key`, `conversation.id`, `x-multi-turn-session-id`, input hash fallback |
-| `/v1/chat/completions` | `x-multi-turn-session-id`, messages hash fallback |
+| `/v1/chat/completions` | `prompt_cache_key`, `x-multi-turn-session-id`, messages hash fallback |
 
 For Responses API follow-up calls, `previous_response_id` does not override the sticky routing key. Continue sending `prompt_cache_key`, `conversation.id`, or the returned `x-multi-turn-session-id` header when the next request needs the same backend affinity.
 
-Sticky routing only affects backend selection when the LLM request router is configured with a cache-affinity-aware routing method for the target model. Clients should only use `x-multi-turn-session-id`. The gateway derives and forwards the internal `x-cache-affinity-key`; clients should not send that header.
+The gateway accepts a nonempty `prompt_cache_key` of up to 256 bytes without
+control characters. An empty value is ignored. The gateway preserves the raw
+value in the upstream request body and returns it in
+`x-multi-turn-session-id`. The gateway derives a SHA-256 value for the
+internal `x-cache-affinity-key` header. Clients must not send
+`x-cache-affinity-key`.
+
+```mermaid
+sequenceDiagram
+    participant Client
+    participant Gateway as LLM API Gateway
+    participant Router as LLM Request Router
+    participant Backend as Model backend
+
+    Client->>Gateway: Chat request with raw prompt_cache_key
+    Note over Gateway: Validate key and derive SHA-256 affinity value
+    Gateway->>Router: Chat JSON and hashed X-Cache-Affinity-Key
+    Router->>Backend: Chat JSON with raw prompt_cache_key
+    Backend-->>Router: Completion response
+    Router-->>Gateway: Completion response
+    Gateway-->>Client: Response and raw x-multi-turn-session-id
+```
+
+Sticky routing only affects backend selection when the LLM request router is
+configured with a cache-affinity-aware routing method for the target model.
 
 ## Metrics
 

docs/user/llm-request-router-load-balancing.md

@@ -121,15 +121,19 @@ filters:
 
 The stock gateway-routes chart does not expose a value for this filter. Use an
 equivalent policy at an external edge or maintain a route override. Preserve
-`x-multi-turn-session-id`; it is the supported client-facing session header.
+`x-multi-turn-session-id`; clients can use it for session affinity. Chat
+Completions and Responses request bodies can also supply `prompt_cache_key`.
 See the
 [Gateway API header modifier guide](https://gateway-api.sigs.k8s.io/guides/user-guides/http-header-modifier/)
 for filter semantics.
 
 The gateway derives `x-cache-affinity-key` for chat-completions and Responses
-requests when affinity applies. It does not derive affinity for embeddings.
-Do not set `require_cache_affinity_key` on a model that serves
-`/v1/embeddings` unless another trusted gateway supplies the key.
+requests when affinity applies. A request body can contain the raw
+`prompt_cache_key`, but only its SHA-256-derived value appears in the internal
+header. The router forwards the request body to the model backend. It does not
+derive affinity for embeddings. Do not set `require_cache_affinity_key` on a
+model that serves `/v1/embeddings` unless another trusted gateway supplies the
+key.
 
 Stargate returns HTTP `400` for a blank, unknown, or configured-but-unavailable
 `x-routing-method`. It also returns HTTP `400` when a required router header is
@@ -204,8 +208,8 @@ algorithm or is present in `request_algorithms`.
 3. Try a method accepted by `nvcf-cli` that is neither the configured
    algorithm nor present in `request_algorithms`; confirm that Stargate returns
    HTTP `400`.
-4. For an affinity-aware method, repeat a supported multi-turn request with
-   the returned `x-multi-turn-session-id`.
+4. For an affinity-aware method, repeat a supported multi-turn request with the
+   same `prompt_cache_key` or the returned `x-multi-turn-session-id`.
 5. Exercise a failed or saturated backend and confirm selection and retry
    counters change.
 

src/invocation-plane-services/llm-api-gateway/README.md

@@ -65,9 +65,12 @@ the selected function/model and estimated prompt size, including
 `x-routing-key`, `x-model`, `x-input-tokens`, and `x-token-estimate`.
 
 For OpenAI-compatible multi-turn stickiness, chat completions and responses
-return `x-multi-turn-session-id`. Clients should persist that value and send it
-on later requests for the same conversation. The gateway forwards only a hashed
-internal `x-cache-affinity-key` to Stargate.
+accept `prompt_cache_key` and return the selected session value in
+`x-multi-turn-session-id`. Clients can send the same `prompt_cache_key` or
+persist the response header and send it on later requests for the same
+conversation. The gateway preserves the raw body field for the model backend.
+It forwards only a SHA-256-derived value in the internal
+`x-cache-affinity-key` header to Stargate.
 
 When `NVCF_GRPC_ADDR` is configured, the gateway authenticates each request
 through the NVCF LLM gRPC auth service, derives the per-caller rate-limit key

src/invocation-plane-services/llm-api-gateway/api/openai_routes_test.go

@@ -195,6 +195,27 @@ func TestOpenAIChatCompletionsReturnsHeaderSessionID(t *testing.T) {
 	}
 }
 
+func TestOpenAIChatCompletionsReturnsPromptCacheKeySessionID(t *testing.T) {
+	t.Parallel()
+
+	e := newTestAPI(config.Default())
+
+	body := `{"model":"fn-alpha/company-name/model-name","messages":[{"role":"user","content":"hello"}],"prompt_cache_key":"chat-prompt-cache-key"}`
+	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(body))
+	req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON)
+	req.Header.Set(HeaderMultiTurnSessionID, "chat-header-session")
+	rec := httptest.NewRecorder()
+
+	e.ServeHTTP(rec, req)
+
+	if rec.Code != http.StatusOK {
+		t.Fatalf("status = %d, want %d: %s", rec.Code, http.StatusOK, rec.Body.String())
+	}
+	if got := rec.Header().Get(HeaderMultiTurnSessionID); got != "chat-prompt-cache-key" {
+		t.Fatalf("%s = %q, want chat-prompt-cache-key", HeaderMultiTurnSessionID, got)
+	}
+}
+
 func TestOpenAIChatCompletionsReturnsGeneratedSessionIDForPayloadFallback(t *testing.T) {
 	t.Parallel()
 
@@ -261,6 +282,51 @@ func TestOpenAIChatCompletionsStreamReturnsSessionHeader(t *testing.T) {
 	}
 }
 
+func TestOpenAIChatCompletionsStreamReturnsPromptCacheKeySessionHeader(t *testing.T) {
+	t.Parallel()
+
+	cfg := config.Default()
+	handlers := NewHandlers(
+		cfg,
+		&stubResponsesProvider{
+			streamEvents: []provider.StreamEvent{
+				{
+					Chunk: &models.ChatCompletionChunk{
+						Choices: []models.ChatCompletionChunkChoice{
+							{
+								Delta: models.ChatCompletionChunkDelta{
+									Content: ptr.To("hello"),
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+		nil,
+		nil,
+	)
+
+	e := echo.New()
+	e.Use(NewContextMiddleware(cfg))
+	handlers.AsOpenAIChatHandlers().RegisterRoutes(e.Group(""))
+
+	body := `{"model":"fn-alpha/company-name/model-name","messages":[{"role":"user","content":"hello"}],"prompt_cache_key":"chat-stream-prompt-cache-key","stream":true}`
+	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(body))
+	req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON)
+	req.Header.Set(HeaderMultiTurnSessionID, "chat-stream-header-session")
+	rec := httptest.NewRecorder()
+
+	e.ServeHTTP(rec, req)
+
+	if rec.Code != http.StatusOK {
+		t.Fatalf("status = %d, want %d: %s", rec.Code, http.StatusOK, rec.Body.String())
+	}
+	if got := rec.Header().Get(HeaderMultiTurnSessionID); got != "chat-stream-prompt-cache-key" {
+		t.Fatalf("%s = %q, want chat-stream-prompt-cache-key", HeaderMultiTurnSessionID, got)
+	}
+}
+
 func TestOpenAIChatCompletionsRejectsInvalidSessionHeader(t *testing.T) {
 	t.Parallel()
 

src/invocation-plane-services/llm-api-gateway/api/session_affinity.go

@@ -96,6 +96,9 @@ func applyChatSessionAffinity(
 	if err != nil {
 		return err
 	}
+	if value := stringPtrValue(request.PromptCacheKey); value != "" {
+		return setSessionAffinity(reqCtx, sessionAffinitySourcePrompt, value)
+	}
 	if headerSessionID != "" {
 		return setSessionAffinity(reqCtx, sessionAffinitySourceHeader, headerSessionID)
 	}
@@ -132,7 +135,12 @@ func setSessionAffinity(
 	}
 	reqCtx.SessionID = sessionID
 	reqCtx.SessionSource = source
-	reqCtx.CacheAffinityKey = affinityKeyForSessionID(sessionID)
+	reqCtx.CacheAffinityKey = affinityKey(sessionAffinitySourceSession, []byte(sessionID))
+	// A header may contain a derived key returned by an earlier request. Body
+	// identifiers are always hashed so their raw values never become router headers.
+	if source == sessionAffinitySourceHeader {
+		reqCtx.CacheAffinityKey = affinityKeyForSessionID(sessionID)
+	}
 	return nil
 }