refactor: unify API type handling across sidecar and coordinator

pkg/common/request/apitype.go

@@ -0,0 +1,128 @@
+/*
+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"
+	"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 was sent to. A value outside the
+// constants below degrades to APITypeChatCompletions.
+type APIType int
+
+const (
+	// APITypeChatCompletions is the Chat Completions API (/v1/chat/completions).
+	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
+	// APITypeMessages is the Anthropic Messages API (/v1/messages).
+	APITypeMessages
+)
+
+// 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"
+	case APITypeMessages:
+		return "messages"
+	default:
+		return fmt.Sprintf("APIType(%d)", int(a))
+	}
+}
+
+// Path returns the canonical request path for the API.
+func (a APIType) Path() string {
+	switch a {
+	case APITypeCompletions:
+		return PathCompletions
+	case APITypeResponses:
+		return PathResponses
+	case APITypeGenerate:
+		return PathGenerate
+	case APITypeMessages:
+		return PathMessages
+	default:
+		return PathChatCompletions
+	}
+}
+
+// DetectAPIType classifies a request path. An unrecognized path maps to
+// APITypeChatCompletions: callers that route only known paths never reach the
+// fallback.
+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 APITypeMessages
+	case strings.Contains(path, PathGenerate):
+		return APITypeGenerate
+	default:
+		return APITypeChatCompletions
+	}
+}
+
+// JSON request field names that cap output tokens, by API. Chat completions caps
+// both max_tokens and max_completion_tokens: vLLM and SGLang accept the two
+// together and prefer max_completion_tokens, so capping both bounds the request
+// regardless of which field the engine consults. The Completions, Messages, and
+// generate APIs share a list: none of them 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 APITypeResponses:
+		return responsesTokenLimitFields
+	case APITypeCompletions, APITypeGenerate, APITypeMessages:
+		return maxTokensOnlyTokenLimitFields
+	default:
+		return chatCompletionTokenLimitFields
+	}
+}

pkg/common/request/apitype_test.go

@@ -0,0 +1,84 @@
+/*
+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},
+		APITypeMessages:        {"messages", PathMessages},
+		APIType(7):             {"APIType(7)", PathChatCompletions},
+	}
+	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", path: PathMessages, want: APITypeMessages},
+		{name: "generate", path: PathGenerate, want: APITypeGenerate},
+		{name: "prefixed chat completions", path: "/prefix" + PathChatCompletions, want: APITypeChatCompletions},
+		{name: "prefixed completions", path: "/prefix" + PathCompletions, want: APITypeCompletions},
+		{name: "prefixed messages", path: "/prefix" + PathMessages, want: APITypeMessages},
+		{name: "prefixed generate", path: "/prefix" + PathGenerate, want: APITypeGenerate},
+		{name: "unknown path falls back to chat completions", path: "/v1/embeddings", want: APITypeChatCompletions},
+		{name: "empty path falls back to chat completions", path: "", want: APITypeChatCompletions},
+	}
+	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},
+		APITypeMessages:        {FieldMaxTokens},
+		APIType(7):             {FieldMaxTokens, FieldMaxCompletionTokens},
+	}
+	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)
+		}
+	}
+}

pkg/common/request/tokens.go

@@ -16,26 +16,37 @@ 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)
-}
+import "maps"
 
-// 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: sampling_params for the generate API, body itself
+// otherwise. The generate API also expects transfer params in that map, 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).
+//
+// body is rewritten in place, so the caller passes its own copy. A one-level
+// copy is enough: the generate sampling_params is always replaced with a map
+// body owns, so the rewrite never reaches a nested map the body was cloned from.
+func CapSingleToken(body map[string]any, apiType APIType) map[string]any {
+	limits := body
+	if apiType == APITypeGenerate {
+		sp, _ := body[FieldSamplingParams].(map[string]any)
+		limits = make(map[string]any, len(sp)+1)
+		maps.Copy(limits, sp)
+		body[FieldSamplingParams] = limits
+	}
+	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
 }