mirror of
https://github.com/netbirdio/netbird.git
synced 2026-08-24 16:41:30 +02:00
[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:
@@ -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
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user