feat(llm-gateway): forward priority (#277)

* feat(llm-api-gateway): forward resolved x-priority to Stargate

NVCF API now resolves a caller priority at auth time and returns it on
AuthLlmInvokeResponse. Carry that value through the gateway and forward it
to Stargate as the X-Priority header so Stargate can apply per-priority
routing and queue estimation.

- Add optional uint32 priority to the AuthLlmInvokeResponse proto and
  regenerate the generated Go.
- Carry the nullable priority on InvocationAuthResponse and copy it onto
  RequestContext in the invocation auth middleware.
- Emit X-Priority on all three LLM paths (chat/completions, responses, and
  the embeddings/generic proxy) only when a value is resolved. An explicit 0
  (highest priority) is forwarded; an unset value omits the header so
  Stargate applies its own default, matching how the gateway already omits
  x-routing-method and x-cache-affinity-key.

The clone-based proxy path still forwards a client-supplied X-Priority when
the gateway resolves none. Rejecting client-supplied X-Priority is handled
separately; that change should land with or before this one.

Signed-off-by: along <along@nvidia.com>

* test(llm-api-gateway): pin priority unset and override semantics

Close two review gaps, both mutation-verified: a client mapping that turns
an unset proto priority into 0 and a proxy that unconditionally emits
X-Priority 0 previously passed the whole suite.

- Assert AuthorizeInvocation maps an unset proto priority to nil, and add a
  round-trip case for explicit 0 and the uint32 max as set values.
- Assert Proxy omits X-Priority when no priority is resolved, and that a
  resolved value replaces a client-supplied header from the clone with
  exactly one value rather than appending.

Signed-off-by: along <along@nvidia.com>

* feat(llm-api-gateway): strip client-supplied x-priority on proxied paths

Signed-off-by: along <along@nvidia.com>

* test(llm-api-gateway): cover priority propagation on all LLM entries

Address review feedback: the middleware priority tests exercised only
/v1/chat/completions. The auth middleware is registered globally and the
routing key comes from the request body model field, so propagation is
path-independent; make the resolved and unset cases table-driven over
/v1/chat/completions, /v1/responses, and /v1/embeddings to pin that for
every entry that forwards X-Priority.

Signed-off-by: along <along@nvidia.com>

---------

Signed-off-by: along <along@nvidia.com>

Author: along-2017

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

@@ -133,6 +133,7 @@ func applyInvocationAuth(
 		reqCtx.RoutingKey = authRoutingKey
 	}
 	reqCtx.ModelSpecs = authResponse.ModelSpecs
+	reqCtx.Priority = authResponse.Priority
 
 	return nil
 }

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

@@ -110,6 +110,125 @@ func TestNVCFAuthMiddlewareEnrichesRequestContext(t *testing.T) {
 	}
 }
 
+// The auth middleware is registered globally, so priority propagation is
+// path-independent; exercise every LLM entry that forwards X-Priority.
+var priorityMiddlewarePaths = []struct {
+	path string
+	body string
+}{
+	{
+		path: "/v1/chat/completions",
+		body: `{"model":"fn-chat/company-name/model-name","messages":[{"role":"user","content":"hello"}]}`,
+	},
+	{
+		path: "/v1/responses",
+		body: `{"model":"fn-chat/company-name/model-name","input":"hello"}`,
+	},
+	{
+		path: "/v1/embeddings",
+		body: `{"model":"fn-chat/company-name/model-name","input":"hello"}`,
+	},
+}
+
+func TestNVCFAuthMiddlewarePropagatesResolvedPriority(t *testing.T) {
+	t.Parallel()
+
+	for _, tc := range priorityMiddlewarePaths {
+		t.Run(tc.path, func(t *testing.T) {
+			t.Parallel()
+
+			wantPriority := uint32(3)
+			authClient := &stubInvocationAuthClient{
+				authResponse: &nvcf.InvocationAuthResponse{
+					RoutingKey:   "fn-chat",
+					ClientAuthID: "subject-123",
+					AuthContext:  map[string]string{"ncaId": "nca-456"},
+					RateLimitKey: "nca-456",
+					Priority:     &wantPriority,
+				},
+			}
+
+			cfg := config.Default()
+
+			e := echo.New()
+			e.Use(NewContextMiddleware(cfg))
+			e.Use(NewNVCFAuthMiddleware(authClient))
+			e.POST(tc.path, func(ec echo.Context) error {
+				gc := ec.(*GatewayContext)
+				reqCtx := gc.RequestContext()
+				if reqCtx == nil {
+					t.Fatal("request context was not set")
+				}
+				if reqCtx.Priority == nil {
+					t.Fatal("priority was not propagated to request context")
+				}
+				if *reqCtx.Priority != 3 {
+					t.Fatalf("priority = %d, want 3", *reqCtx.Priority)
+				}
+				return gc.NoContent(http.StatusNoContent)
+			})
+
+			req := httptest.NewRequest(http.MethodPost, tc.path, strings.NewReader(tc.body))
+			req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON)
+			req.Header.Set(echo.HeaderAuthorization, "Bearer sk-live")
+			rec := httptest.NewRecorder()
+
+			e.ServeHTTP(rec, req)
+
+			if rec.Code != http.StatusNoContent {
+				t.Fatalf("status = %d, want %d: %s", rec.Code, http.StatusNoContent, rec.Body.String())
+			}
+		})
+	}
+}
+
+func TestNVCFAuthMiddlewareLeavesPriorityUnsetWhenAbsent(t *testing.T) {
+	t.Parallel()
+
+	for _, tc := range priorityMiddlewarePaths {
+		t.Run(tc.path, func(t *testing.T) {
+			t.Parallel()
+
+			authClient := &stubInvocationAuthClient{
+				authResponse: &nvcf.InvocationAuthResponse{
+					RoutingKey:   "fn-chat",
+					ClientAuthID: "subject-123",
+					AuthContext:  map[string]string{"ncaId": "nca-456"},
+					RateLimitKey: "nca-456",
+				},
+			}
+
+			cfg := config.Default()
+
+			e := echo.New()
+			e.Use(NewContextMiddleware(cfg))
+			e.Use(NewNVCFAuthMiddleware(authClient))
+			e.POST(tc.path, func(ec echo.Context) error {
+				gc := ec.(*GatewayContext)
+				reqCtx := gc.RequestContext()
+				if reqCtx == nil {
+					t.Fatal("request context was not set")
+				}
+				if reqCtx.Priority != nil {
+					t.Fatalf("priority = %d, want unset", *reqCtx.Priority)
+				}
+				return gc.NoContent(http.StatusNoContent)
+			})
+
+			req := httptest.NewRequest(http.MethodPost, tc.path, strings.NewReader(tc.body))
+			req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON)
+			req.Header.Set(echo.HeaderAuthorization, "Bearer sk-live")
+			rec := httptest.NewRecorder()
+
+			e.ServeHTTP(rec, req)
+
+			if rec.Code != http.StatusNoContent {
+				t.Fatalf("status = %d, want %d: %s", rec.Code, http.StatusNoContent, rec.Body.String())
+			}
+		})
+	}
+}
+
 func TestNVCFAuthMiddlewareRejectsMissingBearerToken(t *testing.T) {
 	t.Parallel()
 

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

@@ -53,6 +53,7 @@ const (
 	headerResponsesMethod       = "X-Routing-Method"
 	headerResponsesModel        = "X-Model"
 	headerResponsesAffinity     = "X-Cache-Affinity-Key"
+	headerResponsesPriority     = "X-Priority"
 )
 
 func (h *ResponsesHandlers) RegisterRoutes(group *echo.Group) {
@@ -293,6 +294,12 @@ func setResponsesProxyContextHeaders(headers http.Header, reqCtx *requestctx.Req
 	if reqCtx.CacheAffinityKey != "" {
 		headers.Set(headerResponsesAffinity, reqCtx.CacheAffinityKey)
 	}
+	// X-Priority is gateway-owned: strip any client-supplied value from the
+	// cloned inbound headers, then set it only when a priority resolved.
+	headers.Del(headerResponsesPriority)
+	if reqCtx.Priority != nil {
+		headers.Set(headerResponsesPriority, strconv.FormatUint(uint64(*reqCtx.Priority), 10))
+	}
 }
 
 func (h *ResponsesHandlers) relayNativeResponsesStream(

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

@@ -167,6 +167,46 @@ func TestSetResponsesProxyContextHeadersSetsTargetRegionCompatibilityHeaders(t *
 	}
 }
 
+func TestSetResponsesProxyContextHeadersForwardsResolvedPriority(t *testing.T) {
+	t.Parallel()
+
+	priority := uint32(3)
+	headers := http.Header{}
+	setResponsesProxyContextHeaders(headers, &requestctx.RequestContext{
+		Priority: &priority,
+	})
+
+	if got := headers.Get(headerResponsesPriority); got != "3" {
+		t.Fatalf("%s = %q, want 3", headerResponsesPriority, got)
+	}
+}
+
+func TestSetResponsesProxyContextHeadersOmitsUnsetPriority(t *testing.T) {
+	t.Parallel()
+
+	headers := http.Header{}
+	setResponsesProxyContextHeaders(headers, &requestctx.RequestContext{})
+
+	if got := headers.Get(headerResponsesPriority); got != "" {
+		t.Fatalf("%s = %q, want empty", headerResponsesPriority, got)
+	}
+}
+
+func TestSetResponsesProxyContextHeadersStripsClientSuppliedPriorityWhenUnset(t *testing.T) {
+	t.Parallel()
+
+	// The responses path clones inbound client headers before applying
+	// context headers; a client-supplied X-Priority must be stripped when no
+	// priority resolved.
+	headers := http.Header{}
+	headers.Set(headerResponsesPriority, "9")
+	setResponsesProxyContextHeaders(headers, &requestctx.RequestContext{})
+
+	if got := headers.Get(headerResponsesPriority); got != "" {
+		t.Fatalf("%s = %q, want empty", headerResponsesPriority, got)
+	}
+}
+
 func TestCreateResponseReusesReturnedSessionIDForAffinity(t *testing.T) {
 	t.Parallel()
 

src/invocation-plane-services/llm-api-gateway/nvcf/client.go

@@ -262,9 +262,18 @@ func (c *GRPCClient) AuthorizeInvocation(
 		AuthContext:  authContext,
 		RateLimitKey: deriveRateLimitKey(authContext),
 		ModelSpecs:   modelSpecsFromProto(resp.GetModelSpecs()),
+		Priority:     priorityFromProto(resp),
 	}, nil
 }
 
+func priorityFromProto(resp *llmgatewaypb.AuthLlmInvokeResponse) *uint32 {
+	if resp == nil || resp.Priority == nil {
+		return nil
+	}
+	value := resp.GetPriority()
+	return &value
+}
+
 func modelSpecsFromProto(specs map[string]*llmgatewaypb.AuthLlmInvokeResponse_ModelSpec) map[string]ModelSpec {
 	if specs == nil {
 		return nil

src/invocation-plane-services/llm-api-gateway/nvcf/client_test.go

@@ -110,6 +110,61 @@ func TestGRPCClientAuthorizeInvocation(t *testing.T) {
 	if len(spec.URIs) != 1 || spec.URIs[0] != "https://example.com/model" {
 		t.Fatalf("uris = %#v, want [https://example.com/model]", spec.URIs)
 	}
+	if authResponse.Priority != nil {
+		t.Fatalf("priority = %d, want unset", *authResponse.Priority)
+	}
+}
+
+func TestGRPCClientAuthorizeInvocationMapsResolvedPriority(t *testing.T) {
+	t.Parallel()
+
+	// Explicit 0 (highest priority) and the uint32 max (lowest) must both
+	// round-trip as set values, distinct from an unset priority.
+	for _, priority := range []uint32{0, 4294967295} {
+		invocationService := &stubInvocationService{
+			t:        t,
+			priority: uint32Ptr(priority),
+		}
+
+		listener := bufconn.Listen(1024 * 1024)
+		server := grpc.NewServer()
+		llmgatewaypb.RegisterLlmGatewayServer(server, invocationService)
+
+		go func() {
+			_ = server.Serve(listener)
+		}()
+		t.Cleanup(func() {
+			server.Stop()
+			_ = listener.Close()
+		})
+
+		conn, err := grpc.NewClient(
+			"passthrough:///bufnet",
+			grpc.WithTransportCredentials(insecure.NewCredentials()),
+			grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) {
+				return listener.Dial()
+			}),
+		)
+		if err != nil {
+			t.Fatalf("create client conn: %v", err)
+		}
+		t.Cleanup(func() {
+			_ = conn.Close()
+		})
+
+		client := NewClientWithConn(conn, func() string { return "service-token" }, time.Second)
+
+		authResponse, err := client.AuthorizeInvocation(context.Background(), "client-token", "fn-123")
+		if err != nil {
+			t.Fatalf("authorize invocation: %v", err)
+		}
+		if authResponse.Priority == nil {
+			t.Fatalf("priority = nil, want %d", priority)
+		}
+		if *authResponse.Priority != priority {
+			t.Fatalf("priority = %d, want %d", *authResponse.Priority, priority)
+		}
+	}
 }
 
 func TestGRPCClientAuthorizeInvocationDoesNotFallbackRateLimitKey(t *testing.T) {
@@ -306,6 +361,7 @@ type stubInvocationService struct {
 	clientAuthID       string
 	clientNCAID        string
 	clientProjectID    string
+	priority           *uint32
 	expectedAuthHeader string
 	traceparent        string
 }
@@ -346,6 +402,7 @@ func (s *stubInvocationService) AuthLlmInvocation(
 	resp := &llmgatewaypb.AuthLlmInvokeResponse{
 		RoutingKey:        "fn-123",
 		ClientAuthSubject: clientAuthID,
+		Priority:          s.priority,
 		ModelSpecs: map[string]*llmgatewaypb.AuthLlmInvokeResponse_ModelSpec{
 			"gateway-model": {
 				Uris:           []string{"https://example.com/model"},
@@ -385,3 +442,7 @@ func incomingMetadataValue(ctx context.Context, key string) string {
 func stringPtr(value string) *string {
 	return &value
 }
+
+func uint32Ptr(value uint32) *uint32 {
+	return &value
+}

src/invocation-plane-services/llm-api-gateway/nvcf/pb/llm_gateway.pb.go

@@ -35,6 +35,9 @@ message AuthLlmInvokeResponse {
     string clientAuthSubject = 2;
     map<string, string> authContext = 3;
     map<string, ModelSpec> modelSpecs = 4;
+    // Resolved caller priority; lower value = higher priority, 0 is highest.
+    // Unset when no priority config applies; absent is not 0, check hasPriority().
+    optional uint32 priority = 5;
 
     message ModelSpec {
         reserved 3;