diff --git a/proxy/internal/middleware/builtin/llm_request_parser/middleware.go b/proxy/internal/middleware/builtin/llm_request_parser/middleware.go index 4d97b2655..7129c2298 100644 --- a/proxy/internal/middleware/builtin/llm_request_parser/middleware.go +++ b/proxy/internal/middleware/builtin/llm_request_parser/middleware.go @@ -61,6 +61,8 @@ func (middlewareImpl) MetadataKeys() []string { middleware.KeyLLMRequestPromptRaw, middleware.KeyLLMCaptureTruncated, middleware.KeyLLMSessionID, + middleware.KeyLLMAgentID, + middleware.KeyLLMParentAgentID, } } @@ -121,9 +123,9 @@ func (m middlewareImpl) Invoke(_ context.Context, in *middleware.Input) (*middle } appendSessionID := func(md []middleware.KV) []middleware.KV { if sessionID != "" { - return append(md, middleware.KV{Key: middleware.KeyLLMSessionID, Value: sessionID}) + md = append(md, middleware.KV{Key: middleware.KeyLLMSessionID, Value: sessionID}) } - return md + return appendAgentIDs(md, in.Headers) } facts, err := parser.ParseRequest(in.Body) @@ -165,6 +167,41 @@ func (m middlewareImpl) Invoke(_ context.Context, in *middleware.Input) (*middle return out, nil } +// agentIDHeader and parentAgentIDHeader carry sub-agent attribution: a +// coding agent that spawns helpers stamps the spawned agent's id, plus the +// spawning agent's when that helper is itself nested. Both are opaque +// identifiers rather than content, so they're emitted regardless of the +// prompt-collection toggle, the same way the session id is. +const ( + agentIDHeader = "x-claude-code-agent-id" + parentAgentIDHeader = "x-claude-code-parent-agent-id" +) + +// appendAgentIDs stamps the sub-agent attribution headers onto the metadata +// bag, skipping either one the request doesn't carry. +func appendAgentIDs(md []middleware.KV, headers []middleware.KV) []middleware.KV { + for _, pair := range []struct{ key, header string }{ + {middleware.KeyLLMAgentID, agentIDHeader}, + {middleware.KeyLLMParentAgentID, parentAgentIDHeader}, + } { + if v := headerValue(headers, pair.header); v != "" { + md = append(md, middleware.KV{Key: pair.key, Value: v}) + } + } + return md +} + +// headerValue returns the first non-empty value for the named header. +// Headers arrive in canonical form, so the match is case-insensitive. +func headerValue(headers []middleware.KV, want string) string { + for _, kv := range headers { + if strings.EqualFold(kv.Key, want) && kv.Value != "" { + return kv.Value + } + } + return "" +} + // sessionIDHeaders are request header names that may carry a client // session identifier, checked in order, case-insensitively. Matching is // against Go's canonical header form, so use the hyphenated names the @@ -178,10 +215,8 @@ var sessionIDHeaders = []string{"x-claude-code-session-id", "session-id", "x-ses // canonical form, so the match is case-insensitive. func sessionIDFromHeaders(headers []middleware.KV) string { for _, want := range sessionIDHeaders { - for _, kv := range headers { - if strings.EqualFold(kv.Key, want) && kv.Value != "" { - return kv.Value - } + if v := headerValue(headers, want); v != "" { + return v } } return "" @@ -309,6 +344,7 @@ func (m middlewareImpl) invokeVertex(in *middleware.Input, vx vertexRequest) *mi if sessionID != "" { md = append(md, middleware.KV{Key: middleware.KeyLLMSessionID, Value: sessionID}) } + md = appendAgentIDs(md, in.Headers) promptTruncated := false if parser != nil && m.capturePrompt { @@ -410,6 +446,7 @@ func (m middlewareImpl) invokeBedrock(in *middleware.Input, br bedrockRequest) * if sessionID != "" { md = append(md, middleware.KV{Key: middleware.KeyLLMSessionID, Value: sessionID}) } + md = appendAgentIDs(md, in.Headers) promptTruncated := false if parser != nil && m.capturePrompt { 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 9518a894e..8d8517860 100644 --- a/proxy/internal/middleware/builtin/llm_request_parser/middleware_test.go +++ b/proxy/internal/middleware/builtin/llm_request_parser/middleware_test.go @@ -45,6 +45,8 @@ func TestMiddleware_StaticSurface(t *testing.T) { middleware.KeyLLMRequestPromptRaw, middleware.KeyLLMCaptureTruncated, middleware.KeyLLMSessionID, + middleware.KeyLLMAgentID, + middleware.KeyLLMParentAgentID, } assert.Equal(t, expected, keys, "metadata key allowlist must match the spec") } @@ -464,3 +466,58 @@ func TestParseVertexPath_CountTokensKeepsModel(t *testing.T) { assert.Equal(t, "anthropic", vx.publisher, "publisher for %q", path) } } + +// TestInvoke_EmitsAgentIDs covers sub-agent attribution: several agents run +// in parallel inside one session, and without their ids every request in +// the session attributes to the session alone. +func TestInvoke_EmitsAgentIDs(t *testing.T) { + mw := newMiddleware(t) + + t.Run("spawned agent", func(t *testing.T) { + out, err := mw.Invoke(context.Background(), &middleware.Input{ + URL: "/v1/messages", + Body: []byte(`{"model":"claude-sonnet-5","messages":[]}`), + Headers: []middleware.KV{ + {Key: "X-Claude-Code-Session-Id", Value: "sess-1"}, + {Key: "X-Claude-Code-Agent-Id", Value: "agent-7"}, + }, + }) + require.NoError(t, err) + + agent, ok := metaValue(t, out.Metadata, middleware.KeyLLMAgentID) + require.True(t, ok, "the spawned agent's id must be emitted") + assert.Equal(t, "agent-7", agent) + + _, ok = metaValue(t, out.Metadata, middleware.KeyLLMParentAgentID) + assert.False(t, ok, "a top-level agent has no parent to emit") + }) + + t.Run("nested agent", func(t *testing.T) { + out, err := mw.Invoke(context.Background(), &middleware.Input{ + URL: "/v1/messages", + Body: []byte(`{"model":"claude-sonnet-5","messages":[]}`), + Headers: []middleware.KV{ + {Key: "X-Claude-Code-Agent-Id", Value: "agent-9"}, + {Key: "X-Claude-Code-Parent-Agent-Id", Value: "agent-7"}, + }, + }) + require.NoError(t, err) + + agent, _ := metaValue(t, out.Metadata, middleware.KeyLLMAgentID) + assert.Equal(t, "agent-9", agent) + parent, ok := metaValue(t, out.Metadata, middleware.KeyLLMParentAgentID) + require.True(t, ok, "a nested agent must carry the spawning agent's id") + assert.Equal(t, "agent-7", parent) + }) + + t.Run("absent on a plain request", func(t *testing.T) { + out, err := mw.Invoke(context.Background(), &middleware.Input{ + URL: "/v1/messages", + Body: []byte(`{"model":"claude-sonnet-5","messages":[]}`), + }) + require.NoError(t, err) + + _, ok := metaValue(t, out.Metadata, middleware.KeyLLMAgentID) + assert.False(t, ok, "no key is emitted when the client sends no agent id") + }) +} diff --git a/proxy/internal/middleware/keys.go b/proxy/internal/middleware/keys.go index 53233c925..eff3fa756 100644 --- a/proxy/internal/middleware/keys.go +++ b/proxy/internal/middleware/keys.go @@ -22,6 +22,15 @@ const ( // body. Empty for clients that don't send one. KeyLLMSessionID = "llm.session_id" + // Sub-agent attribution (emitted by llm_request_parser from the + // client's request headers). A coding agent that spawns helpers + // stamps the spawned agent's id, and the spawning agent's id when + // the helper is itself nested, so cost within one session can be + // split across the agents that ran in parallel. These identify an + // agent, not a person or a device: never treat them as a user id. + KeyLLMAgentID = "llm.agent_id" + KeyLLMParentAgentID = "llm.parent_agent_id" + // LLM response-side metadata (emitted by llm_response_parser). //nolint:gosec // metadata key name, not a credential KeyLLMInputTokens = "llm.input_tokens"