refactor: unify API type handling across sidecar and coordinator

pkg/common/request/apitype.go

@@ -0,0 +1,135 @@
+/*
+Copyright 2026 The llm-d Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package request
+
+import (
+	"fmt"
+	"maps"
+	"strings"
+)
+
+// Inference API paths served by the sidecar and the coordinator.
+const (
+	PathChatCompletions = "/v1/chat/completions"
+	PathCompletions     = "/v1/completions"
+	PathResponses       = "/v1/responses"
+	PathMessages        = "/v1/messages"
+	PathGenerate        = "/inference/v1/generate"
+)
+
+// APIType is the inference API a request speaks. It selects the JSON field
+// names a request carries and the path a synthesized request is sent to.
+type APIType int
+
+const (
+	// APITypeChatCompletions is the Chat Completions API (/v1/chat/completions)
+	// and the Anthropic Messages API (/v1/messages), which share its field names.
+	APITypeChatCompletions APIType = iota
+	// APITypeCompletions is the legacy Completions API (/v1/completions).
+	APITypeCompletions
+	// APITypeResponses is the Responses API (/v1/responses).
+	APITypeResponses
+	// APITypeGenerate is vLLM's token-in generate API (/inference/v1/generate).
+	APITypeGenerate
+)
+
+// String implements fmt.Stringer so structured logs show readable API names.
+func (a APIType) String() string {
+	switch a {
+	case APITypeChatCompletions:
+		return "chat_completions"
+	case APITypeCompletions:
+		return "completions"
+	case APITypeResponses:
+		return "responses"
+	case APITypeGenerate:
+		return "generate"
+	default:
+		return fmt.Sprintf("APIType(%d)", int(a))
+	}
+}
+
+// Path returns the canonical request path for the API. APITypeChatCompletions
+// maps to PathChatCompletions; PathMessages shares its field names but is not
+// a synthesis target.
+func (a APIType) Path() string {
+	switch a {
+	case APITypeChatCompletions:
+		return PathChatCompletions
+	case APITypeCompletions:
+		return PathCompletions
+	case APITypeResponses:
+		return PathResponses
+	default:
+		return PathGenerate
+	}
+}
+
+// DetectAPIType classifies a request path. An unrecognized path maps to
+// APITypeGenerate: callers that route only known paths never reach the
+// fallback, and a path the router does not register is not a client fault.
+func DetectAPIType(path string) APIType {
+	switch {
+	case strings.Contains(path, PathChatCompletions):
+		return APITypeChatCompletions
+	case strings.Contains(path, PathCompletions):
+		return APITypeCompletions
+	case strings.Contains(path, PathResponses):
+		return APITypeResponses
+	case strings.Contains(path, PathMessages):
+		return APITypeChatCompletions
+	default:
+		return APITypeGenerate
+	}
+}
+
+// JSON request field names that cap output tokens, by API. The Completions and
+// generate APIs share a list: neither defines max_completion_tokens, so capping
+// it would put a field on the wire that a strict server is free to reject.
+var (
+	chatCompletionTokenLimitFields = []string{FieldMaxTokens, FieldMaxCompletionTokens}
+	maxTokensOnlyTokenLimitFields  = []string{FieldMaxTokens}
+	responsesTokenLimitFields      = []string{FieldMaxOutputTokens}
+)
+
+// TokenLimitFields returns the output token cap field names the API uses.
+// The returned slices are shared package-level vars; callers must not mutate them.
+func (a APIType) TokenLimitFields() []string {
+	switch a {
+	case APITypeCompletions, APITypeGenerate:
+		return maxTokensOnlyTokenLimitFields
+	case APITypeResponses:
+		return responsesTokenLimitFields
+	default:
+		return chatCompletionTokenLimitFields
+	}
+}
+
+// TokenLimitMap returns the map inside body that holds the token limit fields:
+// sampling_params for the generate API, body itself otherwise. The generate map
+// is always replaced with one body owns, so a caller that writes into the result
+// never reaches a nested map the body was cloned from.
+func (a APIType) TokenLimitMap(body map[string]any) map[string]any {
+	if a != APITypeGenerate {
+		return body
+	}
+	sp, _ := body[FieldSamplingParams].(map[string]any)
+	owned := make(map[string]any, len(sp)+1)
+	maps.Copy(owned, sp)
+	body[FieldSamplingParams] = owned
+	return owned
+}

pkg/common/request/apitype_test.go

@@ -0,0 +1,131 @@
+/*
+Copyright 2026 The llm-d Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package request
+
+import (
+	"reflect"
+	"testing"
+)
+
+func TestAPIType_StringAndPath(t *testing.T) {
+	cases := map[APIType]struct{ name, path string }{
+		APITypeChatCompletions: {"chat_completions", PathChatCompletions},
+		APITypeCompletions:     {"completions", PathCompletions},
+		APITypeResponses:       {"responses", PathResponses},
+		APITypeGenerate:        {"generate", PathGenerate},
+		APIType(7):             {"APIType(7)", PathGenerate},
+	}
+	for apiType, want := range cases {
+		if got := apiType.String(); got != want.name {
+			t.Errorf("APIType(%d).String() = %q, want %q", int(apiType), got, want.name)
+		}
+		if got := apiType.Path(); got != want.path {
+			t.Errorf("APIType(%d).Path() = %q, want %q", int(apiType), got, want.path)
+		}
+	}
+}
+
+func TestDetectAPIType(t *testing.T) {
+	tests := []struct {
+		name string
+		path string
+		want APIType
+	}{
+		{name: "chat completions", path: PathChatCompletions, want: APITypeChatCompletions},
+		{name: "completions", path: PathCompletions, want: APITypeCompletions},
+		{name: "responses", path: PathResponses, want: APITypeResponses},
+		{name: "messages shares chat completions fields", path: PathMessages, want: APITypeChatCompletions},
+		{name: "generate", path: PathGenerate, want: APITypeGenerate},
+		{name: "prefixed chat completions", path: "/prefix" + PathChatCompletions, want: APITypeChatCompletions},
+		{name: "prefixed completions", path: "/prefix" + PathCompletions, want: APITypeCompletions},
+		{name: "unknown path falls back to generate", path: "/v1/embeddings", want: APITypeGenerate},
+		{name: "empty path falls back to generate", path: "", want: APITypeGenerate},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := DetectAPIType(tt.path); got != tt.want {
+				t.Errorf("DetectAPIType(%q) = %v, want %v", tt.path, got, tt.want)
+			}
+		})
+	}
+}
+
+func TestAPIType_TokenLimitFields(t *testing.T) {
+	cases := map[APIType][]string{
+		APITypeChatCompletions: {FieldMaxTokens, FieldMaxCompletionTokens},
+		APITypeCompletions:     {FieldMaxTokens},
+		APITypeResponses:       {FieldMaxOutputTokens},
+		APITypeGenerate:        {FieldMaxTokens},
+	}
+	for apiType, want := range cases {
+		if got := apiType.TokenLimitFields(); !reflect.DeepEqual(got, want) {
+			t.Errorf("APIType(%d).TokenLimitFields() = %v, want %v", int(apiType), got, want)
+		}
+	}
+}
+
+func TestAPIType_TokenLimitMap(t *testing.T) {
+	t.Run("non-generate returns the body itself", func(t *testing.T) {
+		body := map[string]any{"model": "m"}
+		got := APITypeChatCompletions.TokenLimitMap(body)
+		if !reflect.DeepEqual(got, body) {
+			t.Errorf("got %v, want the body %v", got, body)
+		}
+		if _, ok := body[FieldSamplingParams]; ok {
+			t.Error("sampling_params was added to a non-generate body")
+		}
+	})
+
+	// A generate body may share sampling_params with the body it was cloned from,
+	// so the caller gets a copy the body owns and the original stays intact.
+	t.Run("generate copies an existing sampling_params", func(t *testing.T) {
+		sp := map[string]any{FieldMaxTokens: 100}
+		body := map[string]any{FieldSamplingParams: sp}
+
+		got := APITypeGenerate.TokenLimitMap(body)
+
+		if !reflect.DeepEqual(got, sp) {
+			t.Errorf("got %v, want the entries of %v", got, sp)
+		}
+		got[FieldMaxTokens] = 1
+		if sp[FieldMaxTokens] != 100 {
+			t.Errorf("caller's sampling_params was written through: %v", sp)
+		}
+		if body[FieldSamplingParams].(map[string]any)[FieldMaxTokens] != 1 {
+			t.Errorf("body sampling_params = %v, want the returned map", body[FieldSamplingParams])
+		}
+	})
+
+	for name, value := range map[string]any{"absent": nil, "not a map": "not-a-map"} {
+		t.Run("generate replaces a sampling_params that is "+name, func(t *testing.T) {
+			body := map[string]any{"model": "m"}
+			if value != nil {
+				body[FieldSamplingParams] = value
+			}
+
+			got := APITypeGenerate.TokenLimitMap(body)
+
+			if len(got) != 0 {
+				t.Errorf("got %v, want an empty map", got)
+			}
+			got[FieldMaxTokens] = 1
+			if sp, ok := body[FieldSamplingParams].(map[string]any); !ok || sp[FieldMaxTokens] != 1 {
+				t.Errorf("body sampling_params = %v, want the returned map", body[FieldSamplingParams])
+			}
+		})
+	}
+}

pkg/common/request/tokens.go

@@ -16,26 +16,33 @@ limitations under the License.
 
 package request
 
-// CapMaxTokensField caps target's max_tokens to 1 and strips min_tokens.
-// min_tokens is stripped rather than clamped: it defaults to 0 in vLLM, so
-// removing it keeps min_tokens <= max_tokens=1 without raising the floor
-// above the cap (vLLM's SamplingParams rejects min_tokens > max_tokens).
-func CapMaxTokensField(target map[string]any) {
-	target[FieldMaxTokens] = 1
-	delete(target, FieldMinTokens)
-}
-
-// PrimeSingleTokenRequest mutates target in place into a synthetic,
-// non-streaming, single-output-token chat-completions or completions
-// request. max_completion_tokens is unconditionally capped to 1 alongside
-// max_tokens: vLLM and SGLang both accept the two fields together
-// (max_completion_tokens takes precedence over max_tokens when present),
-// so setting both guarantees the cap regardless of which field the serving
-// engine consults.
-func PrimeSingleTokenRequest(target map[string]any) {
-	CapMaxTokensField(target)
-	target[FieldMaxCompletionTokens] = 1
+// CapSingleToken rewrites body into a synthetic, non-streaming,
+// single-output-token request for a prefill or encode leg. It returns the map
+// the caps were written into, which is where the generate API also expects
+// transfer params, so a caller adding them needs no second lookup.
+//
+// The caps to rewrite come from APIType.TokenLimitFields, so each API's output
+// caps are named in one place. min_tokens is a floor rather than a cap, so it is
+// stripped instead of capped: it defaults to 0 in vLLM, so removing it keeps
+// min_tokens <= max_tokens=1 without raising the floor above the cap (vLLM's
+// SamplingParams rejects min_tokens > max_tokens).
+//
+// Chat completions lists both max_tokens and max_completion_tokens: vLLM and
+// SGLang accept the two together and prefer max_completion_tokens, so capping
+// both bounds the leg regardless of which field the engine consults.
+//
+// body is rewritten in place, so the caller passes its own copy. A one-level
+// copy is enough: the generate API caps inside sampling_params, and
+// TokenLimitMap replaces that nested map rather than writing through it, so a
+// body that still shares it with the decode leg keeps the client's limits.
+func CapSingleToken(body map[string]any, apiType APIType) map[string]any {
+	limits := apiType.TokenLimitMap(body)
+	for _, field := range apiType.TokenLimitFields() {
+		limits[field] = 1
+	}
+	delete(limits, FieldMinTokens)
 
-	target[FieldStream] = false
-	delete(target, FieldStreamOptions)
+	body[FieldStream] = false
+	delete(body, FieldStreamOptions)
+	return limits
 }