diff --git a/docs/user/llm-gateway.md b/docs/user/llm-gateway.md index c56dc40dc..863aed2b7 100644 --- a/docs/user/llm-gateway.md +++ b/docs/user/llm-gateway.md @@ -117,6 +117,7 @@ curl -sS -X POST "http://${GATEWAY_ADDR}/v1/chat/completions" \ -d '{ "model": "/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 diff --git a/docs/user/llm-request-router-load-balancing.md b/docs/user/llm-request-router-load-balancing.md index 306ab4245..2b568affc 100644 --- a/docs/user/llm-request-router-load-balancing.md +++ b/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. diff --git a/src/invocation-plane-services/llm-api-gateway/README.md b/src/invocation-plane-services/llm-api-gateway/README.md index d07385e42..a04c2fb3b 100644 --- a/src/invocation-plane-services/llm-api-gateway/README.md +++ b/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 diff --git a/src/invocation-plane-services/llm-api-gateway/api/openai_routes_test.go b/src/invocation-plane-services/llm-api-gateway/api/openai_routes_test.go index 759e51219..ab93ee5d0 100644 --- a/src/invocation-plane-services/llm-api-gateway/api/openai_routes_test.go +++ b/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() diff --git a/src/invocation-plane-services/llm-api-gateway/api/session_affinity.go b/src/invocation-plane-services/llm-api-gateway/api/session_affinity.go index 9b7641e88..7e904651a 100644 --- a/src/invocation-plane-services/llm-api-gateway/api/session_affinity.go +++ b/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 } diff --git a/src/invocation-plane-services/llm-api-gateway/api/session_affinity_test.go b/src/invocation-plane-services/llm-api-gateway/api/session_affinity_test.go index 39a703f61..4c6b2b89e 100644 --- a/src/invocation-plane-services/llm-api-gateway/api/session_affinity_test.go +++ b/src/invocation-plane-services/llm-api-gateway/api/session_affinity_test.go @@ -18,6 +18,7 @@ limitations under the License. package api import ( + "fmt" "net/http" "net/http/httptest" "strings" @@ -92,6 +93,32 @@ func TestApplyResponsesSessionAffinityPrefersPromptCacheKeyOverConversation(t *t } } +func TestApplyResponsesSessionAffinityHashesPromptCacheKeyThatLooksLikeAffinityKey(t *testing.T) { + t.Parallel() + + ctx := newSessionAffinityTestContext() + promptCacheKey := "mt:v1:payload:" + strings.Repeat("a", 64) + request := &openairesponses.CreateRequest{ + PromptCacheKey: &promptCacheKey, + } + + if err := applyResponsesSessionAffinity(ctx, request); err != nil { + t.Fatalf("applyResponsesSessionAffinity: %v", err) + } + + reqCtx := ctx.RequestContext() + if reqCtx.SessionID != promptCacheKey { + t.Fatalf("SessionID = %q, want %q", reqCtx.SessionID, promptCacheKey) + } + if reqCtx.CacheAffinityKey == promptCacheKey { + t.Fatalf("CacheAffinityKey contains unmodified prompt cache key: %q", reqCtx.CacheAffinityKey) + } + want := affinityKey(sessionAffinitySourceSession, []byte(promptCacheKey)) + if reqCtx.CacheAffinityKey != want { + t.Fatalf("CacheAffinityKey = %q, want %q", reqCtx.CacheAffinityKey, want) + } +} + func TestApplyResponsesSessionAffinityReusesReturnedBodyIDFromHeader(t *testing.T) { t.Parallel() @@ -174,6 +201,167 @@ func TestApplyChatSessionAffinityReusesGeneratedPayloadIDFromHeader(t *testing.T } } +func TestApplyChatSessionAffinityPrefersPromptCacheKeyOverHeader(t *testing.T) { + t.Parallel() + + ctx := newSessionAffinityTestContext() + ctx.Request().Header.Set(HeaderMultiTurnSessionID, "header-session") + promptCacheKey := "prompt-cache-session" + request := &models.ChatCompletionRequest{ + PromptCacheKey: &promptCacheKey, + } + + if err := applyChatSessionAffinity(ctx, request); err != nil { + t.Fatalf("applyChatSessionAffinity: %v", err) + } + + reqCtx := ctx.RequestContext() + if reqCtx.SessionID != promptCacheKey { + t.Fatalf("SessionID = %q, want %q", reqCtx.SessionID, promptCacheKey) + } + if reqCtx.SessionSource != sessionAffinitySourcePrompt { + t.Fatalf("SessionSource = %q, want %q", reqCtx.SessionSource, sessionAffinitySourcePrompt) + } + if !strings.HasPrefix(reqCtx.CacheAffinityKey, "mt:v1:session:") { + t.Fatalf("CacheAffinityKey = %q, want session source", reqCtx.CacheAffinityKey) + } + if strings.Contains(reqCtx.CacheAffinityKey, promptCacheKey) || + strings.Contains(reqCtx.CacheAffinityKey, "header-session") { + t.Fatalf("CacheAffinityKey leaks raw session value: %q", reqCtx.CacheAffinityKey) + } + + second := newSessionAffinityTestContext() + secondRequest := &models.ChatCompletionRequest{PromptCacheKey: &promptCacheKey} + if err := applyChatSessionAffinity(second, secondRequest); err != nil { + t.Fatalf("second applyChatSessionAffinity: %v", err) + } + if second.RequestContext().CacheAffinityKey != reqCtx.CacheAffinityKey { + t.Fatalf("second CacheAffinityKey = %q, want %q", second.RequestContext().CacheAffinityKey, reqCtx.CacheAffinityKey) + } +} + +func TestApplyChatSessionAffinityHashesPromptCacheKeyThatLooksLikeAffinityKey(t *testing.T) { + t.Parallel() + + ctx := newSessionAffinityTestContext() + promptCacheKey := "mt:v1:payload:" + strings.Repeat("a", 64) + request := &models.ChatCompletionRequest{ + PromptCacheKey: &promptCacheKey, + } + + if err := applyChatSessionAffinity(ctx, request); err != nil { + t.Fatalf("applyChatSessionAffinity: %v", err) + } + + reqCtx := ctx.RequestContext() + if reqCtx.SessionID != promptCacheKey { + t.Fatalf("SessionID = %q, want %q", reqCtx.SessionID, promptCacheKey) + } + if reqCtx.CacheAffinityKey == promptCacheKey { + t.Fatalf("CacheAffinityKey contains unmodified prompt cache key: %q", reqCtx.CacheAffinityKey) + } + want := affinityKey(sessionAffinitySourceSession, []byte(promptCacheKey)) + if reqCtx.CacheAffinityKey != want { + t.Fatalf("CacheAffinityKey = %q, want %q", reqCtx.CacheAffinityKey, want) + } +} + +func TestApplyChatSessionAffinityEmptyPromptCacheKeyFallsBackToHeader(t *testing.T) { + t.Parallel() + + ctx := newSessionAffinityTestContext() + ctx.Request().Header.Set(HeaderMultiTurnSessionID, "header-session") + promptCacheKey := "" + request := &models.ChatCompletionRequest{ + PromptCacheKey: &promptCacheKey, + } + + if err := applyChatSessionAffinity(ctx, request); err != nil { + t.Fatalf("applyChatSessionAffinity: %v", err) + } + + reqCtx := ctx.RequestContext() + if reqCtx.SessionID != "header-session" { + t.Fatalf("SessionID = %q, want header-session", reqCtx.SessionID) + } + if reqCtx.SessionSource != sessionAffinitySourceHeader { + t.Fatalf("SessionSource = %q, want %q", reqCtx.SessionSource, sessionAffinitySourceHeader) + } +} + +func TestApplyChatSessionAffinityRejectsInvalidPromptCacheKey(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + promptCacheKey string + wantMessage string + }{ + { + name: "oversized", + promptCacheKey: strings.Repeat("a", sessionIDMaxLen+1), + wantMessage: "must be at most 256 bytes", + }, + { + name: "control character", + promptCacheKey: "prompt\ncache", + wantMessage: "must not contain control characters", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + ctx := newSessionAffinityTestContext() + request := &models.ChatCompletionRequest{ + PromptCacheKey: &test.promptCacheKey, + } + + err := applyChatSessionAffinity(ctx, request) + if err == nil { + t.Fatal("applyChatSessionAffinity returned nil error") + } + httpErr, ok := err.(*echo.HTTPError) + if !ok { + t.Fatalf("error = %T, want *echo.HTTPError", err) + } + if httpErr.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want %d", httpErr.Code, http.StatusBadRequest) + } + if !strings.Contains(fmt.Sprint(httpErr.Message), test.wantMessage) { + t.Fatalf("message = %q, want substring %q", httpErr.Message, test.wantMessage) + } + }) + } +} + +func TestApplyChatSessionAffinityRejectsInvalidHeaderBeforePromptCacheKey(t *testing.T) { + t.Parallel() + + ctx := newSessionAffinityTestContext() + ctx.Request().Header.Set(HeaderMultiTurnSessionID, "bad\nsession") + promptCacheKey := "valid-prompt-cache-session" + request := &models.ChatCompletionRequest{ + PromptCacheKey: &promptCacheKey, + } + + err := applyChatSessionAffinity(ctx, request) + if err == nil { + t.Fatal("applyChatSessionAffinity returned nil error") + } + httpErr, ok := err.(*echo.HTTPError) + if !ok { + t.Fatalf("error = %T, want *echo.HTTPError", err) + } + if httpErr.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want %d", httpErr.Code, http.StatusBadRequest) + } + if !strings.Contains(fmt.Sprint(httpErr.Message), HeaderMultiTurnSessionID+" is invalid") { + t.Fatalf("message = %q, want invalid header error", httpErr.Message) + } +} + func TestApplyChatSessionAffinityDoesNotOverwriteExistingResponsesAffinity(t *testing.T) { t.Parallel() diff --git a/src/invocation-plane-services/llm-api-gateway/models/openai.go b/src/invocation-plane-services/llm-api-gateway/models/openai.go index 018c957c1..322ff554e 100644 --- a/src/invocation-plane-services/llm-api-gateway/models/openai.go +++ b/src/invocation-plane-services/llm-api-gateway/models/openai.go @@ -322,6 +322,7 @@ type ChatCompletionRequest struct { ToolChoice ChatCompletionToolChoiceField `json:"tool_choice"` FunctionChoice ChatCompletionFunctionChoiceField `json:"function_call"` ParallelToolCalls *bool `json:"parallel_tool_calls"` + PromptCacheKey *string `json:"prompt_cache_key"` User *string `json:"user"` ReasoningFormat *string `json:"reasoning_format"` ReasoningEffort *string `json:"reasoning_effort"` diff --git a/src/invocation-plane-services/llm-api-gateway/provider/stargate_test.go b/src/invocation-plane-services/llm-api-gateway/provider/stargate_test.go index 80c9b5b27..a1cccf766 100644 --- a/src/invocation-plane-services/llm-api-gateway/provider/stargate_test.go +++ b/src/invocation-plane-services/llm-api-gateway/provider/stargate_test.go @@ -19,6 +19,7 @@ package provider import ( "context" + "crypto/sha256" "encoding/json" "errors" "fmt" @@ -134,9 +135,12 @@ func metricHasFunctionID(data metricdata.Aggregation, want string) bool { func TestStargateProviderCompleteForwardsChatPayloadAndRoutingHeaders(t *testing.T) { t.Parallel() + promptCacheKey := "chat-prompt-cache-key" + cacheAffinityKey := fmt.Sprintf("mt:v1:session:%x", sha256.Sum256([]byte(promptCacheKey))) request := &NormalizedRequest{ ChatRequest: &models.ChatCompletionRequest{ - Model: "upstream-model", + Model: "upstream-model", + PromptCacheKey: &promptCacheKey, Messages: &[]models.ChatMessage{ { Role: models.ChatCompletionRoleUser, @@ -161,7 +165,7 @@ func TestStargateProviderCompleteForwardsChatPayloadAndRoutingHeaders(t *testing Model: "upstream-model", RoutingMethod: "experimental_method", TargetRegion: "us-west1", - CacheAffinityKey: "mt:v1:header:hash", + CacheAffinityKey: cacheAffinityKey, } wantEstimate := routingTokenEstimate(request) @@ -182,7 +186,8 @@ func TestStargateProviderCompleteForwardsChatPayloadAndRoutingHeaders(t *testing require.Equal(t, "fn-abc", r.Header.Get(headerRoutingKey)) require.Equal(t, "upstream-model", r.Header.Get(headerModel)) require.Equal(t, "experimental_method", r.Header.Get(headerRoutingMethod)) - require.Equal(t, "mt:v1:header:hash", r.Header.Get(headerCacheAffinityKey)) + require.Equal(t, cacheAffinityKey, r.Header.Get(headerCacheAffinityKey)) + require.NotEqual(t, promptCacheKey, r.Header.Get(headerCacheAffinityKey)) require.Equal(t, fmt.Sprintf("%d", wantEstimate), r.Header.Get(headerInputTokens)) require.Equal(t, fmt.Sprintf("%d", wantEstimate), r.Header.Get(headerTokenEstimate)) @@ -193,6 +198,8 @@ func TestStargateProviderCompleteForwardsChatPayloadAndRoutingHeaders(t *testing require.NotNil(t, payload.StreamOptions) require.NotNil(t, payload.StreamOptions.IncludeUsage) require.True(t, ptr.Deref(payload.StreamOptions.IncludeUsage)) + require.NotNil(t, payload.PromptCacheKey) + require.Equal(t, promptCacheKey, ptr.Deref(payload.PromptCacheKey)) require.NotNil(t, payload.Messages) require.Len(t, *payload.Messages, 1) require.Equal(t, models.ChatCompletionRoleUser, (*payload.Messages)[0].Role)