Cap coordinator prefill/encode legs to a single token per schema

pkg/coordinator/gateway/client.go

@@ -92,7 +92,7 @@ func (c *Client) Request(ctx context.Context, method, path string, body []byte,
 	logger := log.FromContext(ctx).WithName("gateway")
 	if body != nil {
 		if v := logger.V(logutil.TRACE); v.Enabled() {
-			v.Info("request body", "method", method, "path", path, "headers", httplog.RedactedHeaders(req.Header), "body", redactBody(body))
+			v.Info("request body", "method", method, "path", path, "headers", httplog.RedactedHeaders(req.Header), "body", RedactBody(body))
 		}
 	}
 
@@ -109,7 +109,7 @@ func (c *Client) Request(ctx context.Context, method, path string, body []byte,
 		if err != nil {
 			return nil, fmt.Errorf("reading response from gateway: %w", err)
 		}
-		v.Info("response body", "status", resp.StatusCode, "body", redactBody(respBody))
+		v.Info("response body", "status", resp.StatusCode, "body", RedactBody(respBody))
 		resp.Body = io.NopCloser(bytes.NewReader(respBody))
 	}
 
@@ -146,10 +146,11 @@ const (
 	maxRedactDepth = 32
 )
 
-// redactBody parses JSON and redacts oversized string values (see
+// RedactBody parses JSON and redacts oversized string values (see
 // maxRedactStringLen); a body that is not valid JSON is returned verbatim, or
-// truncated at maxRedactRawBodyLen.
-func redactBody(data []byte) any {
+// truncated at maxRedactRawBodyLen. It is exported for the decode reverse-proxy
+// path, which bypasses Client.Request and so redacts its own body before logging.
+func RedactBody(data []byte) any {
 	var v any
 	if err := json.Unmarshal(data, &v); err != nil {
 		if len(data) > maxRedactRawBodyLen {

pkg/coordinator/gateway/client_test.go

@@ -147,9 +147,9 @@ func TestRedactBody(t *testing.T) {
 			t.Fatal(err)
 		}
 
-		got, ok := redactBody(body).(map[string]any)
+		got, ok := RedactBody(body).(map[string]any)
 		if !ok {
-			t.Fatalf("expected map[string]any, got %T", redactBody(body))
+			t.Fatalf("expected map[string]any, got %T", RedactBody(body))
 		}
 		if got["model"] != "m" {
 			t.Errorf("short value changed: %v", got["model"])
@@ -160,23 +160,23 @@ func TestRedactBody(t *testing.T) {
 	})
 
 	t.Run("non-JSON under max length returned verbatim", func(t *testing.T) {
-		if got := redactBody([]byte("not json")); got != "not json" {
+		if got := RedactBody([]byte("not json")); got != "not json" {
 			t.Errorf("redactBody = %v, want verbatim string", got)
 		}
 	})
 
 	t.Run("non-JSON exactly max length returned verbatim", func(t *testing.T) {
 		raw := strings.Repeat("q", maxRedactRawBodyLen)
-		if got := redactBody([]byte(raw)); got != raw {
+		if got := RedactBody([]byte(raw)); got != raw {
 			t.Errorf("redactBody = %v, want verbatim (boundary is > not >=)", got)
 		}
 	})
 
 	t.Run("non-JSON over max length truncated", func(t *testing.T) {
 		raw := strings.Repeat("q", maxRedactRawBodyLen+50)
-		got, ok := redactBody([]byte(raw)).(string)
+		got, ok := RedactBody([]byte(raw)).(string)
 		if !ok {
-			t.Fatalf("expected string, got %T", redactBody([]byte(raw)))
+			t.Fatalf("expected string, got %T", RedactBody([]byte(raw)))
 		}
 		if want := raw[:maxRedactRawBodyLen] + "..."; got != want {
 			t.Errorf("redactBody truncation = %q, want %q", got, want)

pkg/coordinator/steps/decode_proxy.go

@@ -78,6 +78,14 @@ func newDecodeProxyRequest(ctx context.Context, logger logr.Logger, step string,
 		v.Info("request body", "method", "POST", "path", reqCtx.OriginalPath, "bodyLen", len(bodyBytes), "headers", httplog.RedactedHeaders(proxyReq.Header))
 	}
 
+	// The decode leg is proxied straight to the gateway and never passes through
+	// gateway.Client.Request, which is what logs full request bodies at TRACE for
+	// the prefill and encode legs. Log the redacted body here so the decode leg's
+	// fields (min_tokens, max_tokens, ...) are observable at the same verbosity.
+	if v := logger.V(logutil.TRACE); v.Enabled() {
+		v.Info("request body", "method", "POST", "path", reqCtx.OriginalPath, "headers", httplog.RedactedHeaders(proxyReq.Header), "body", gateway.RedactBody(bodyBytes))
+	}
+
 	return proxyReq, nil
 }
 

pkg/coordinator/steps/encode.go

@@ -215,19 +215,20 @@ func (s *EncodeStep) buildEncodeBody(reqCtx *pipeline.RequestContext, tokenIDs [
 				},
 			},
 		}
-		reqcommon.PrimeSingleTokenRequest(body, reqCtx.Body)
+		capSingleTokenOutput(body, format)
 		return body
 	default:
-		return map[string]any{
+		body := map[string]any{
 			"model":     reqCtx.Model,
 			"token_ids": tokenIDs,
 			"features": map[string]any{
 				"mm_hashes":       map[string][]string{ModalityImage: {entry.Hash}},
 				"mm_placeholders": map[string][]any{ModalityImage: {map[string]any{"offset": 1, "length": entry.Placeholder.Length}}},
 				"kwargs_data":     mmKwargsField([]string{entry.KwargsData}),
 			},
-			reqcommon.FieldSamplingParams: map[string]any{reqcommon.FieldMaxTokens: 1},
 		}
+		capSingleTokenOutput(body, format)
+		return body
 	}
 }
 

pkg/coordinator/steps/encode_test.go

@@ -352,10 +352,10 @@ func TestEncodeStep_ChatCompletionsFormat(t *testing.T) {
 }
 
 // TestEncodeStep_ChatCompletionsFormat_CapsMaxCompletionTokens is a
-// regression test: buildEncodeBody used to always emit max_tokens=1 without
-// regard to which token-limit field the client used, so a reasoning-model
-// client's max_completion_tokens was never capped in the encoder sub-request.
-func TestEncodeStep_ChatCompletionsFormat_CapsMaxCompletionTokens(t *testing.T) {
+// The encode chat sub-request is built fresh from the request context and does
+// not carry the client's sampling fields, so max_completion_tokens is not
+// propagated and is never injected: max_tokens=1 alone caps output.
+func TestEncodeStep_ChatCompletionsFormat_OmitsMaxCompletionTokens(t *testing.T) {
 	var receivedBody map[string]any
 
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -401,8 +401,11 @@ func TestEncodeStep_ChatCompletionsFormat_CapsMaxCompletionTokens(t *testing.T)
 		t.Fatalf("unexpected error: %v", err)
 	}
 
-	if receivedBody["max_completion_tokens"] != float64(1) {
-		t.Fatalf("expected encode sub-request max_completion_tokens capped to 1, got %v", receivedBody["max_completion_tokens"])
+	if receivedBody["max_tokens"] != float64(1) {
+		t.Fatalf("expected encode sub-request max_tokens capped to 1, got %v", receivedBody["max_tokens"])
+	}
+	if _, ok := receivedBody["max_completion_tokens"]; ok {
+		t.Fatalf("expected encode sub-request to omit max_completion_tokens, got %v", receivedBody["max_completion_tokens"])
 	}
 }
 
@@ -575,3 +578,44 @@ func TestEncodeStep_BuildsCorrectTokenIDs(t *testing.T) {
 		}
 	}
 }
+
+// TestEncodeStep_GenerateFormat_CapsSingleToken verifies the generate-format
+// encoder sub-request caps output to a single token: sampling_params carries
+// max_tokens=1 and strips min_tokens (it defaults to 0, keeping min_tokens <=
+// max_tokens).
+func TestEncodeStep_GenerateFormat_CapsSingleToken(t *testing.T) {
+	var samplingParams map[string]any
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		body, _ := io.ReadAll(r.Body)
+		var parsed map[string]any
+		_ = json.Unmarshal(body, &parsed)
+		samplingParams, _ = parsed["sampling_params"].(map[string]any)
+		_ = json.NewEncoder(w).Encode(map[string]any{
+			"ec_transfer_params": map[string]any{"h1": map[string]any{"peer_port": 5501}},
+		})
+	}))
+	defer server.Close()
+
+	gwClient := gateway.New(config.GatewayConfig{Address: server.URL})
+	step, _ := NewEncodeStep(gwClient, map[string]any{"use_openai_format": false})
+
+	reqCtx := &pipeline.RequestContext{
+		RequestID: "req-gen-cap",
+		Model:     "test",
+		TokenIDs:  []int{1, 32000, 32000, 32000, 2345},
+		MultimodalEntries: []pipeline.MultimodalEntry{
+			{Index: 0, Hash: "h1", KwargsData: "dGVzdA==", Placeholder: pipeline.PlaceholderRange{Offset: 1, Length: 3}},
+		},
+	}
+
+	if err := step.Execute(context.Background(), reqCtx); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if samplingParams["max_tokens"] != float64(1) {
+		t.Fatalf("expected sampling_params.max_tokens=1, got %v", samplingParams["max_tokens"])
+	}
+	if _, ok := samplingParams["min_tokens"]; ok {
+		t.Fatalf("expected sampling_params.min_tokens to be stripped, got %v", samplingParams["min_tokens"])
+	}
+}

pkg/coordinator/steps/prefill.go

@@ -137,7 +137,7 @@ func (s *PrefillStep) buildPrefillBody(ctx context.Context, reqCtx *pipeline.Req
 	switch format {
 	case gateway.FormatChatCompletions:
 		body := maps.Clone(reqCtx.Body)
-		reqcommon.PrimeSingleTokenRequest(body, reqCtx.Body)
+		capSingleTokenOutput(body, format)
 		tokens := map[string]any{
 			"token_ids": reqCtx.TokenIDs,
 		}
@@ -164,9 +164,9 @@ func (s *PrefillStep) buildPrefillBody(ctx context.Context, reqCtx *pipeline.Req
 			"request_id":                    reqCtx.RequestID,
 			"model":                         reqCtx.Model,
 			"prompt":                        prompt,
-			reqcommon.FieldMaxTokens:        1,
 			reqcommon.FieldKVTransferParams: kvParams,
 		}
+		capSingleTokenOutput(body, format)
 		if features != nil {
 			body["features"] = features
 		}
@@ -186,6 +186,7 @@ func (s *PrefillStep) buildPrefillBody(ctx context.Context, reqCtx *pipeline.Req
 			"model":                       reqCtx.Model,
 			reqcommon.FieldSamplingParams: sampling,
 		}
+		capSingleTokenOutput(body, format)
 		if features != nil {
 			body["features"] = features
 		}

pkg/coordinator/steps/prefill_test.go

@@ -133,6 +133,9 @@ func TestPrefillStep_SendsCorrectGenerateRequest(t *testing.T) {
 	if samplingParams["max_tokens"] != float64(1) {
 		t.Fatalf("expected sampling_params.max_tokens=1, got %v", samplingParams["max_tokens"])
 	}
+	if _, ok := samplingParams["min_tokens"]; ok {
+		t.Fatalf("expected sampling_params.min_tokens to be stripped, got %v", samplingParams["min_tokens"])
+	}
 	extraArgs, ok := samplingParams["extra_args"].(map[string]any)
 	if !ok {
 		t.Fatal("expected sampling_params.extra_args in generate format")
@@ -222,6 +225,14 @@ func TestPrefillStep_CompletionsFormat(t *testing.T) {
 	if prefillBody["request_id"] != "req-compl" {
 		t.Fatalf("expected request_id, got %v", prefillBody["request_id"])
 	}
+	// Prefill leg caps output to a single token: max_tokens is pinned to 1 and
+	// min_tokens is stripped (it defaults to 0, keeping min_tokens <= max_tokens).
+	if prefillBody["max_tokens"] != float64(1) {
+		t.Fatalf("expected max_tokens=1, got %v", prefillBody["max_tokens"])
+	}
+	if _, ok := prefillBody["min_tokens"]; ok {
+		t.Fatalf("expected min_tokens to be stripped, got %v", prefillBody["min_tokens"])
+	}
 	// Completions format has top-level kv_transfer_params
 	kvParams, ok := prefillBody["kv_transfer_params"].(map[string]any)
 	if !ok {
@@ -470,6 +481,56 @@ func TestPrefillStep_ChatCompletionsFormat_CapsMaxCompletionTokens(t *testing.T)
 	}
 }
 
+// TestPrefillStep_ChatCompletionsFormat_StripsClientMinTokens is a regression
+// test for the one coordinator path where a client-supplied min_tokens survives
+// into the capped body: chat-completions clones reqCtx.Body, so a client
+// min_tokens > 1 would leave min_tokens > max_tokens=1 and vLLM rejects the leg.
+// The generate and completions legs build fresh bodies that never carry it.
+func TestPrefillStep_ChatCompletionsFormat_StripsClientMinTokens(t *testing.T) {
+	var prefillBody map[string]any
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		body, _ := io.ReadAll(r.Body)
+		_ = json.Unmarshal(body, &prefillBody)
+		_ = json.NewEncoder(w).Encode(map[string]any{
+			"kv_transfer_params": map[string]any{"block_id": "block-5"},
+		})
+	}))
+	defer server.Close()
+
+	gwClient := gateway.New(config.GatewayConfig{Address: server.URL})
+	step, err := NewPrefillStep(gwClient, map[string]any{})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	reqCtx := &pipeline.RequestContext{
+		RequestID:    "req-chat-min-tokens",
+		OriginalPath: gateway.PathChatCompletions,
+		Model:        "test-model",
+		Body: map[string]any{
+			"model":      "test-model",
+			"min_tokens": 50,
+			"max_tokens": 200,
+			"messages": []any{
+				map[string]any{"role": "user", "content": "hello"},
+			},
+		},
+		KVTransferParams: make(map[string]any),
+	}
+
+	if err := step.Execute(context.Background(), reqCtx); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	if prefillBody["max_tokens"] != float64(1) {
+		t.Fatalf("expected prefill request max_tokens=1, got %v", prefillBody["max_tokens"])
+	}
+	if _, ok := prefillBody["min_tokens"]; ok {
+		t.Fatalf("expected client min_tokens to be stripped, got %v", prefillBody["min_tokens"])
+	}
+}
+
 // TestSharedStorage_OmitsECTransferParams_InPrefillBody verifies that the
 // ec-shared-storage EC connector never emits an ec_transfer_params field on
 // the prefill body, in every prefill wire format (chat-completions,

pkg/coordinator/steps/utils.go

@@ -74,6 +74,40 @@ func resolveFormat(useOpenAIFormat bool, path string) gateway.RequestFormat {
 	return detected
 }
 
+// capSingleTokenOutput rewrites body into a single-output-token, non-streaming
+// request for the synthetic prefill and encode legs.
+//
+// Intentionally distinct from the sidecar's reqcommon.PrimeSingleTokenRequest
+// for now.
+// TODO: unify the two into one shared single-token helper in a future refactor.
+func capSingleTokenOutput(body map[string]any, format gateway.RequestFormat) {
+	target := body
+	if format == gateway.FormatGenerate {
+		sp, ok := body[reqcommon.FieldSamplingParams].(map[string]any)
+		if !ok {
+			sp = map[string]any{}
+			body[reqcommon.FieldSamplingParams] = sp
+		}
+		target = sp
+	}
+
+	target[reqcommon.FieldMaxTokens] = 1
+	// Strip rather than clamp min_tokens: it defaults to 0 in vLLM, so removing it
+	// keeps min_tokens <= max_tokens=1 without raising the floor above the cap.
+	delete(target, reqcommon.FieldMinTokens)
+
+	if _, ok := body[reqcommon.FieldMaxCompletionTokens]; ok {
+		body[reqcommon.FieldMaxCompletionTokens] = 1
+	}
+
+	// TODO: max_output_tokens is another client-supplied output cap (Responses
+	// API) that a client can send instead of max_tokens/max_completion_tokens; it
+	// should be capped to 1 here as well so the synthetic legs stay single-token.
+
+	body[reqcommon.FieldStream] = false
+	delete(body, reqcommon.FieldStreamOptions)
+}
+
 // buildMMFeatures builds the multimodal features map (mm_hashes, mm_placeholders,
 // and optionally kwargs_data) from the request's multimodal entries. It returns
 // nil when there are no entries.