From 4d2b8b407b84db811e9caa5ad5c2b036d4028f26 Mon Sep 17 00:00:00 2001 From: mlsmaycon Date: Tue, 11 Aug 2026 02:56:29 +0000 Subject: [PATCH] [proxy] Keep the Vertex model id out of the count-tokens method segment Vertex hangs token counting off the model as its own path segment, and the parser split the tail on the final colon alone. A count-tokens request therefore reported its model as "claude-sonnet-5/count-tokens", which no route claims, so the request denied as not-routable and the access log recorded a model that does not exist. Stop at the first "/" after the model id so the method segment stays out of it, leaving the client free to price its context against the dedicated endpoint instead of the billable inference one. --- .../builtin/llm_request_parser/middleware.go | 6 +++++ .../llm_request_parser/middleware_test.go | 23 +++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/proxy/internal/middleware/builtin/llm_request_parser/middleware.go b/proxy/internal/middleware/builtin/llm_request_parser/middleware.go index 051e7949b..4d97b2655 100644 --- a/proxy/internal/middleware/builtin/llm_request_parser/middleware.go +++ b/proxy/internal/middleware/builtin/llm_request_parser/middleware.go @@ -257,6 +257,12 @@ func parseVertexPath(reqPath string) (vertexRequest, bool) { if c := strings.LastIndex(rest, ":"); c >= 0 { model, action = rest[:c], rest[c+1:] } + // Token counting hangs off the model as its own path segment + // (".../models/{model}/count-tokens:rawPredict"), so anything past the + // first "/" belongs to the method rather than the model id. + if slash := strings.Index(model, "/"); slash >= 0 { + model = model[:slash] + } model = llm.NormalizeVertexModel(model) if model == "" { return vertexRequest{}, false diff --git a/proxy/internal/middleware/builtin/llm_request_parser/middleware_test.go b/proxy/internal/middleware/builtin/llm_request_parser/middleware_test.go index 7afe545b9..9518a894e 100644 --- a/proxy/internal/middleware/builtin/llm_request_parser/middleware_test.go +++ b/proxy/internal/middleware/builtin/llm_request_parser/middleware_test.go @@ -441,3 +441,26 @@ func TestInvoke_NilInputAllows(t *testing.T) { assert.Equal(t, middleware.DecisionAllow, out.Decision, "nil input still allows") assert.Empty(t, out.Metadata, "nil input emits no metadata") } + +// TestParseVertexPath_CountTokensKeepsModel covers Vertex token counting, +// where the method hangs off the model as its own path segment. Splitting +// only on the final colon swallowed "/count-tokens" into the model id, so +// the router saw a model no route could claim. +func TestParseVertexPath_CountTokensKeepsModel(t *testing.T) { + cases := map[string]struct { + model string + stream bool + }{ + "/v1/projects/p/locations/global/publishers/anthropic/models/claude-sonnet-5:rawPredict": {model: "claude-sonnet-5"}, + "/v1/projects/p/locations/global/publishers/anthropic/models/claude-sonnet-5:streamRawPredict": {model: "claude-sonnet-5", stream: true}, + "/v1/projects/p/locations/global/publishers/anthropic/models/claude-sonnet-5/count-tokens:rawPredict": {model: "claude-sonnet-5"}, + "/v1/projects/p/locations/global/publishers/anthropic/models/claude-sonnet-5@20250929/count-tokens:rawPredict": {model: "claude-sonnet-5"}, + } + for path, want := range cases { + vx, ok := parseVertexPath(path) + require.True(t, ok, "must parse %q", path) + assert.Equal(t, want.model, vx.model, "model for %q", path) + assert.Equal(t, want.stream, vx.stream, "stream flag for %q", path) + assert.Equal(t, "anthropic", vx.publisher, "publisher for %q", path) + } +}