[proxy] Skip OpenAI-shape identity injection on non-OpenAI bodies

Gateway records enable body-level identity so LiteLLM's tag-budget check
can read it, and the injector wrote "user" and "metadata.tags" into every
JSON object regardless of dialect. Claude Code reaches those same records
on /v1/messages, where "user" is not a permitted top-level field and
metadata accepts only "user_id", so the upstream rejected the request with
a 400 naming a field the client never sent. Rewriting the body also
changed the bytes a gateway-side prompt cache keys on.

Gate the body write on the surface llm_request_parser resolved from the
path. Header stamping is untouched, so spend tracking and per-end-user
budgets keep working on the surfaces that lose the body path.
This commit is contained in:
mlsmaycon
2026-08-11 02:52:11 +00:00
committed by GitHub
parent 119b70c664
commit 00a4f7dee1
2 changed files with 85 additions and 0 deletions

View File

@@ -217,6 +217,32 @@ func applyHeaderPair(rule *HeaderPairRule, in *middleware.Input) *middleware.Mut
return mutations
}
// bodyInjectableSurfaces are the request-body dialects that accept the
// OpenAI-standard identity fields this middleware writes. A surface
// outside this set gets header-only stamping: "user" and "metadata.tags"
// are not part of the Anthropic Messages schema, which rejects unknown
// top-level fields and permits only "user_id" under metadata, so writing
// them into an Anthropic-shaped body turns a working request into a 400.
// Claude Code speaks that shape through gateway records pinned to the
// OpenAI parser, so the check keys on the detected surface rather than
// on the provider record.
var bodyInjectableSurfaces = map[string]struct{}{
"openai": {},
// An empty surface means no parser claimed the path (a custom gateway
// base). Those upstreams are OpenAI-compatible by convention, so keep
// the long-standing behaviour rather than silently dropping identity.
"": {},
}
// bodyAcceptsOpenAIIdentity reports whether the request body may carry the
// OpenAI-standard identity fields, read from the surface llm_request_parser
// resolved from the request path.
func bodyAcceptsOpenAIIdentity(in *middleware.Input) bool {
surface, _ := lookupMetadata(in.Metadata, middleware.KeyLLMProvider)
_, ok := bodyInjectableSurfaces[surface]
return ok
}
// injectIntoBody parses the request body and writes the supplied
// identity dimensions into it. Tags land at metadata.tags (creating
// the metadata object when absent); the user identity lands at the
@@ -225,6 +251,8 @@ func applyHeaderPair(rule *HeaderPairRule, in *middleware.Input) *middleware.Mut
// was written. Returns ok=false (no mutation) when:
//
// - both inputs are empty (nothing to write);
// - the body speaks a dialect without these fields (see
// bodyInjectableSurfaces);
// - the body is empty or truncated (we don't have the full document
// to safely round-trip);
// - the body isn't a JSON object (skip silently — this middleware
@@ -245,6 +273,9 @@ func injectIntoBody(in *middleware.Input, tags []string, userID string) ([]byte,
if in == nil || len(in.Body) == 0 || in.BodyTruncated {
return nil, false
}
if !bodyAcceptsOpenAIIdentity(in) {
return nil, false
}
var doc map[string]any
if err := json.Unmarshal(in.Body, &doc); err != nil {
return nil, false

View File

@@ -704,3 +704,57 @@ func TestInject_ExtraHeaders_EmptyValueSkipped(t *testing.T) {
"empty extra value must not be stamped")
}
}
// TestInject_AnthropicBodyIsNotRewritten pins the shape gate. Claude Code
// reaches a LiteLLM record on /v1/messages, where "user" is not a
// permitted top-level field and metadata accepts only "user_id", so
// writing the OpenAI-standard fields would turn a working request into a
// 400 naming a field the client never sent. Header stamping still runs, so
// spend tracking and per-end-user budgets keep working.
func TestInject_AnthropicBodyIsNotRewritten(t *testing.T) {
rule := liteLLMRuleWithBody()
rule.HeaderPair.EndUserIDInBody = true
mw := New(Config{Providers: []ProviderInjection{rule}})
in := newInput(litellmProvider, "alice", []string{"grp-eng"})
in.UserEmail = "alice@example.com"
in.Metadata = append(in.Metadata, middleware.KV{Key: middleware.KeyLLMProvider, Value: "anthropic"})
in.Body = []byte(`{"model":"claude-sonnet-5","messages":[]}`)
out, err := mw.Invoke(context.Background(), in)
require.NoError(t, err)
require.NotNil(t, out.Mutations)
assert.Empty(t, out.Mutations.BodyReplace,
"an Anthropic-shaped body must reach the upstream unmodified")
var endUser string
for _, kv := range out.Mutations.HeadersAdd {
if kv.Key == "x-litellm-end-user-id" {
endUser = kv.Value
}
}
assert.Equal(t, "alice@example.com", endUser,
"header stamping must still carry identity when body inject is skipped")
}
// TestInject_OpenAIBodyStillRewritten guards the gate against
// over-reaching: the OpenAI surface must keep its body-level identity,
// which is the only path LiteLLM's tag-budget check reads.
func TestInject_OpenAIBodyStillRewritten(t *testing.T) {
mw := New(Config{Providers: []ProviderInjection{liteLLMRuleWithBody()}})
in := newInput(litellmProvider, "alice", []string{"grp-eng"})
in.Metadata = append(in.Metadata, middleware.KV{Key: middleware.KeyLLMProvider, Value: "openai"})
in.Body = []byte(`{"model":"gpt-4o-mini","messages":[]}`)
out, err := mw.Invoke(context.Background(), in)
require.NoError(t, err)
require.NotNil(t, out.Mutations)
require.NotEmpty(t, out.Mutations.BodyReplace, "the OpenAI surface still gets body tags")
var doc map[string]any
require.NoError(t, json.Unmarshal(out.Mutations.BodyReplace, &doc))
meta, ok := doc["metadata"].(map[string]any)
require.True(t, ok, "metadata must be an object")
assert.NotEmpty(t, meta["tags"], "metadata.tags must still be written")
}