Skip to content

feat(proxy): add key_alias, key_hash, requested_model DD APM span tags - #22710

Merged
ishaan-jaff merged 3 commits into
mainfrom
worktree-hidden-petting-candy
Mar 4, 2026
Merged

feat(proxy): add key_alias, key_hash, requested_model DD APM span tags#22710
ishaan-jaff merged 3 commits into
mainfrom
worktree-hidden-petting-candy

Conversation

@ishaan-jaff

Copy link
Copy Markdown
Contributor

Relevant issues

Pre-Submission checklist

  • I have Added testing in the tests/test_litellm/ directory, Adding at least 1 test is a hard requirement - see details
  • My PR passes all unit tests on make test-unit
  • My PR's scope is as isolated as possible, it only solves 1 specific problem
  • I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

CI (LiteLLM team)

  • Branch creation CI run
    Link:

  • CI run for the last commit
    Link:

  • Merge / cherry-pick CI run
    Links:

Type

🆕 New Feature

Changes

Adds _add_dd_apm_tags_for_request to common_request_processing.py, called on every authenticated request alongside the existing call_id tag.

Three new tags on each DD APM span:

  • litellm.key_alias — human-readable API key alias (e.g. "prod-app")
  • litellm.key_hash — hashed API key (already hashed before storage, safe to log)
  • litellm.requested_model — model as sent by the client

Useful for:

  • tracing all calls from a specific user/key by filtering on litellm.key_alias or litellm.key_hash
  • tracing all calls to a specific model by filtering on litellm.requested_model

Key info is skipped when not available (e.g. 401 unauthenticated requests). All tagging is best-effort and never affects request processing.

@vercel

vercel Bot commented Mar 3, 2026

Copy link
Copy Markdown

The latest updates on your projects. Learn more about Vercel for GitHub.

Project Deployment Actions Updated (UTC)
litellm Error Error Mar 3, 2026 10:03pm

Request Review

@greptile-apps

greptile-apps Bot commented Mar 3, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

Refactors Datadog APM span tagging into a new DDSpanTagger class (litellm/proxy/dd_span_tagger.py) and adds three new span tags (litellm.key_alias, litellm.key_hash, litellm.requested_model) to every authenticated proxy request. The existing call_id tagging logic is moved from an inline function in common_request_processing.py into the new class. All tagging is best-effort and wrapped in try/except, so it never impacts request processing. The token field used for key_hash is always pre-hashed by the UserAPIKeyAuth model validator, so no raw secrets are leaked. Also includes new dashscope/qwen3 model pricing entries.

  • New DDSpanTagger class consolidates all DD APM tagging logic in one place
  • Tags key_alias, key_hash, and requested_model on each authenticated request span
  • 3 new mock-only unit tests cover the tagging behavior
  • Model pricing backup updated with new dashscope/qwen3 model entries

Confidence Score: 4/5

  • This PR is safe to merge — the changes are additive, best-effort, and well-isolated from request processing logic.
  • The new DDSpanTagger class is cleanly extracted, all tagging is wrapped in try/except and never affects request processing, the token field is pre-hashed (safe to log), tests use mocks only, and no new DB queries or expensive objects are created in the critical path. The only minor concern (already flagged in a prior thread) is that requested_model captures the post-alias model rather than the original client-sent model.
  • No files require special attention beyond the previously flagged requested_model semantics in litellm/proxy/common_request_processing.py.

Important Files Changed

Filename Overview
litellm/proxy/dd_span_tagger.py New file extracting DD span tagging into a dedicated DDSpanTagger class. Adds tag_call_id (moved from common_request_processing.py) and new tag_request for key_alias, key_hash, and requested_model tags. All tagging is best-effort with proper exception handling.
litellm/proxy/common_request_processing.py Removes inline _add_dd_apm_tags_for_litellm_call_id function, replaces with DDSpanTagger.tag_call_id and adds DDSpanTagger.tag_request call. Note: requested_model is tagged after alias resolution, which was flagged in a prior thread.
tests/test_litellm/proxy/test_common_request_processing.py Adds 3 well-structured mock-only tests for DDSpanTagger.tag_request covering: all tags present, no tags when key absent, and model-only tagging. Existing tag_call_id test updated to use new class. No network calls.
litellm/model_prices_and_context_window_backup.json Adds new dashscope/qwen3 model entries with pricing and capability metadata. Standard data addition, no code logic changes.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Incoming Proxy Request] --> B[Authentication & Key Resolution]
    B --> C[Model Alias Mapping]
    C --> D[Generate litellm_call_id]
    D --> E[DDSpanTagger.tag_call_id]
    D --> F[DDSpanTagger.tag_request]
    E --> G["Set litellm.call_id tag"]
    F --> H{key_alias present?}
    H -->|Yes| I["Set litellm.key_alias tag"]
    H -->|No| J{token present?}
    I --> J
    J -->|Yes| K["Set litellm.key_hash tag"]
    J -->|No| L{requested_model present?}
    K --> L
    L -->|Yes| M["Set litellm.requested_model tag"]
    L -->|No| N[Continue Request Processing]
    M --> N
    G --> N
Loading

Last reviewed commit: 7d0aca2

Comment on lines +699 to +702
_add_dd_apm_tags_for_request(
user_api_key_dict=user_api_key_dict,
requested_model=self.data.get("model"),
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

requested_model tag captures post-alias model, not client-sent model

The _add_dd_apm_tags_for_request call here uses self.data.get("model") which has already been transformed by model alias mapping (lines 680-693). However, the docstring (line 278) and PR description say this tag captures "model name as sent by the client."

The actual client-sent model is captured later in base_process as requested_model_from_client (line 869-870) before common_processing_pre_call_logic is called. So the tag litellm.requested_model will contain the alias-resolved name, not the original request model.

If the intent is to trace by the original client-sent model name, the tag should be set before alias resolution or use the original value. If the post-alias model is intentional, the docstring/tag name should be updated to avoid confusion (e.g., litellm.resolved_model).

@ishaan-jaff
ishaan-jaff merged commit 7befe3c into main Mar 4, 2026
29 of 41 checks passed
ghost pushed a commit that referenced this pull request Mar 4, 2026
* add explicit caching to litellm proxy for gemini models via injection

* fix: add missing `supports_function_calling` for deepinfra models

All 55 deepinfra models that had `supports_tool_choice: true` were
missing the `supports_function_calling` flag, causing
`litellm.supports_function_calling()` to incorrectly return False.

Fixes #22619

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* Managed batches - Address PR bot comments from #22464

* feat(togetherai): add support for TogetherAI Qwen3.5-397B-A17B model

* Agent Tracing - support context_id based trace id propogation + nested llm calls  (#22626)

* style(ui/): distinguish agent calls from llm calls on ui

* feat: initial grouping working

* feat: set stable contextid for a2a calls - allows for easily passing to downstream llm/mcp calls

* feat(a2a_endpoints.py): fix tracing to avoid recreating logging objects for the same call

allows stable trace id usage

* fix(guardrail_endpoints): handle string ui_type values in _build_field_dict

_build_field_dict unconditionally called .value on ui_type, which crashes
for guardrail configs that use plain strings (e.g. BlockCodeExecutionGuardrailConfigModel
uses "multiselect" and "percentage"). Now checks with hasattr before calling .value.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: propagate trace/session id from headers in MCP server calls

Cherry-picked mcp_server/server.py fixes from 6feb9ba: adds
get_chain_id_from_headers to extract x-litellm-trace-id /
x-litellm-session-id from raw headers, and uses it in call_tool
and list_tools to keep spend logs and tracing consistent with A2A.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

* [Feat] UI - Add Open in New Tab on leftnav Bar (#22731)

* Add minimal dev_config.yaml for proxy development

Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com>

* feat(ui): wrap left nav items in <a> tags for open-in-new-tab support

Nav items are now rendered as <a> elements with proper href attributes,
enabling right-click → 'Open in new tab', Ctrl/Cmd+click, and
middle-click to open any sidebar page in a new browser tab.

Normal clicks continue to use SPA navigation (no full page reload).

Applied to both leftnav.tsx (query-param routing) and Sidebar2.tsx
(Next.js file-based routing).

Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com>

---------

Co-authored-by: Cursor Agent <cursoragent@cursor.com>
Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com>

* [Feat] Add Tool Policies for AI Gateway  (#22732)

* fix: fix ui render

* fix: fix minor bugs

* refactor: use prisma functions instead of raw sql (safer)

* fix(add-new-tiles-to-tool-policies): allow developer to see what's available

* feat: ensure tool allowlist runs correctly for tool names + mcp's

* refactor: more ui improvements

* feat: working key tool blocking

* feat(tools): show tool logs

* refactor: backend code improvements

* refactor: improve log viewer for tools

* fix: address PR review feedback for tool access control

- Add missing blocked_tools column to root schema.prisma (schema drift)
- Invalidate ToolPolicyRegistry after policy mutations so changes take effect immediately
- Remove dead code: unused get_effective_policies, get_tool_policies_cached, and helpers

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: race condition in permission resolution and remove duplicate allowlist check

- Use atomic update_many with object_permission_id=None to prevent concurrent
  requests from creating orphaned permission rows and losing tool blocks
- Remove duplicate allowed_tools enforcement from guardrail (already enforced
  in auth layer via check_tools_allowlist)
- Move inline uuid import to module level

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* update to account for  userAgent

* UI - Add ToolDetails

* input/output policy

* LiteLLM_PolicyAttachmentTable

* LiteLLM_PolicyAttachmentTable

* fix: add _enqueue_tool_registry_upsert

* fix: tool mgmt endpoints

* tool mgmt endpoints

* Update tests/test_litellm/proxy/db/test_tool_registry_writer.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

* Update tests/test_litellm/proxy/db/test_tool_registry_writer.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

* Update tests/test_litellm/proxy/db/test_tool_registry_writer.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

* fix: sync root schema.prisma and fix test_tool_registry_writer for input/output policy

- Migrate root schema.prisma LiteLLM_ToolTable from call_policy to
  input_policy/output_policy, add missing user_agent and last_used_at columns
  (now consistent with litellm/proxy/schema.prisma and litellm-proxy-extras)
- Fix SpendLogToolIndex comment across all three schema files
- Fix all call_policy references in test_tool_registry_writer.py:
  swapped update_tool_policy arguments, wrong get_tools_by_names return type
  assertions, _mock_tool_row setting call_policy instead of input_policy

Addresses Greptile review feedback on PR #22732.

Made-with: Cursor

---------

Co-authored-by: Krrish Dholakia <krrishdholakia@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

* feat(proxy): add key_alias, key_hash, requested_model DD APM span tags (#22710)

* feat(proxy): add key_alias, key_hash, requested_model tags to DD APM spans

* refactor(proxy): consolidate DD APM tag helpers into DDSpanTagger class

* refactor(proxy): move DDSpanTagger to its own file litellm/proxy/dd_span_tagger.py

---------

Co-authored-by: liweiguang <codingpunk@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: Ephrim Stanley <ephrim.stanley@point72.com>
Co-authored-by: Varad Khonde <varadkhonde@gmail.com>
Co-authored-by: Krish Dholakia <krrishdholakia@gmail.com>
Co-authored-by: Sameer Kankute <sameer@berri.ai>
Co-authored-by: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Co-authored-by: Cursor Agent <cursoragent@cursor.com>
Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
fzowl pushed a commit to fzowl/litellm that referenced this pull request Jun 24, 2026
BerriAI#22710)

* feat(proxy): add key_alias, key_hash, requested_model tags to DD APM spans

* refactor(proxy): consolidate DD APM tag helpers into DDSpanTagger class

* refactor(proxy): move DDSpanTagger to its own file litellm/proxy/dd_span_tagger.py
fzowl pushed a commit to fzowl/litellm that referenced this pull request Jun 24, 2026
…I#21881)

* add explicit caching to litellm proxy for gemini models via injection

* fix: add missing `supports_function_calling` for deepinfra models

All 55 deepinfra models that had `supports_tool_choice: true` were
missing the `supports_function_calling` flag, causing
`litellm.supports_function_calling()` to incorrectly return False.

Fixes BerriAI#22619


* Managed batches - Address PR bot comments from BerriAI#22464

* feat(togetherai): add support for TogetherAI Qwen3.5-397B-A17B model

* Agent Tracing - support context_id based trace id propogation + nested llm calls  (BerriAI#22626)

* style(ui/): distinguish agent calls from llm calls on ui

* feat: initial grouping working

* feat: set stable contextid for a2a calls - allows for easily passing to downstream llm/mcp calls

* feat(a2a_endpoints.py): fix tracing to avoid recreating logging objects for the same call

allows stable trace id usage

* fix(guardrail_endpoints): handle string ui_type values in _build_field_dict

_build_field_dict unconditionally called .value on ui_type, which crashes
for guardrail configs that use plain strings (e.g. BlockCodeExecutionGuardrailConfigModel
uses "multiselect" and "percentage"). Now checks with hasattr before calling .value.


* fix: propagate trace/session id from headers in MCP server calls

Cherry-picked mcp_server/server.py fixes from 6feb9ba: adds
get_chain_id_from_headers to extract x-litellm-trace-id /
x-litellm-session-id from raw headers, and uses it in call_tool
and list_tools to keep spend logs and tracing consistent with A2A.


---------


* [Feat] UI - Add Open in New Tab on leftnav Bar (BerriAI#22731)

* Add minimal dev_config.yaml for proxy development

Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com>

* feat(ui): wrap left nav items in <a> tags for open-in-new-tab support

Nav items are now rendered as <a> elements with proper href attributes,
enabling right-click → 'Open in new tab', Ctrl/Cmd+click, and
middle-click to open any sidebar page in a new browser tab.

Normal clicks continue to use SPA navigation (no full page reload).

Applied to both leftnav.tsx (query-param routing) and Sidebar2.tsx
(Next.js file-based routing).

Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com>

---------

Co-authored-by: Cursor Agent <cursoragent@cursor.com>
Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com>

* [Feat] Add Tool Policies for AI Gateway  (BerriAI#22732)

* fix: fix ui render

* fix: fix minor bugs

* refactor: use prisma functions instead of raw sql (safer)

* fix(add-new-tiles-to-tool-policies): allow developer to see what's available

* feat: ensure tool allowlist runs correctly for tool names + mcp's

* refactor: more ui improvements

* feat: working key tool blocking

* feat(tools): show tool logs

* refactor: backend code improvements

* refactor: improve log viewer for tools

* fix: address PR review feedback for tool access control

- Add missing blocked_tools column to root schema.prisma (schema drift)
- Invalidate ToolPolicyRegistry after policy mutations so changes take effect immediately
- Remove dead code: unused get_effective_policies, get_tool_policies_cached, and helpers


* fix: race condition in permission resolution and remove duplicate allowlist check

- Use atomic update_many with object_permission_id=None to prevent concurrent
  requests from creating orphaned permission rows and losing tool blocks
- Remove duplicate allowed_tools enforcement from guardrail (already enforced
  in auth layer via check_tools_allowlist)
- Move inline uuid import to module level


* update to account for  userAgent

* UI - Add ToolDetails

* input/output policy

* LiteLLM_PolicyAttachmentTable

* LiteLLM_PolicyAttachmentTable

* fix: add _enqueue_tool_registry_upsert

* fix: tool mgmt endpoints

* tool mgmt endpoints

* Update tests/test_litellm/proxy/db/test_tool_registry_writer.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

* Update tests/test_litellm/proxy/db/test_tool_registry_writer.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

* Update tests/test_litellm/proxy/db/test_tool_registry_writer.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

* fix: sync root schema.prisma and fix test_tool_registry_writer for input/output policy

- Migrate root schema.prisma LiteLLM_ToolTable from call_policy to
  input_policy/output_policy, add missing user_agent and last_used_at columns
  (now consistent with litellm/proxy/schema.prisma and litellm-proxy-extras)
- Fix SpendLogToolIndex comment across all three schema files
- Fix all call_policy references in test_tool_registry_writer.py:
  swapped update_tool_policy arguments, wrong get_tools_by_names return type
  assertions, _mock_tool_row setting call_policy instead of input_policy

Addresses Greptile review feedback on PR BerriAI#22732.

Made-with: Cursor

---------

Co-authored-by: Krrish Dholakia <krrishdholakia@gmail.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

* feat(proxy): add key_alias, key_hash, requested_model DD APM span tags (BerriAI#22710)

* feat(proxy): add key_alias, key_hash, requested_model tags to DD APM spans

* refactor(proxy): consolidate DD APM tag helpers into DDSpanTagger class

* refactor(proxy): move DDSpanTagger to its own file litellm/proxy/dd_span_tagger.py

---------

Co-authored-by: liweiguang <codingpunk@gmail.com>
Co-authored-by: Ephrim Stanley <ephrim.stanley@point72.com>
Co-authored-by: Varad Khonde <varadkhonde@gmail.com>
Co-authored-by: Krish Dholakia <krrishdholakia@gmail.com>
Co-authored-by: Sameer Kankute <sameer@berri.ai>
Co-authored-by: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Co-authored-by: Cursor Agent <cursoragent@cursor.com>
Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant