[Feat] Add Tool Policies for AI Gateway - #22732
Merged
Merged
Conversation
- 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>
…owlist 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>
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Contributor
Greptile SummaryThis PR introduces a comprehensive Tool Policy system for the AI Gateway, splitting the previous single
Confidence Score: 4/5
|
| Filename | Overview |
|---|---|
| litellm/proxy/db/tool_registry_writer.py | Complete overhaul: raw SQL → Prisma ORM, new ToolPolicyRegistry in-memory singleton for hot-path enforcement. Read-modify-write race condition in blocked_tools management; full table scan on every sync cycle may not scale. |
| litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py | Two-stage guardrail: blocks tools by input_policy, then enforces trust chain on response path. Clean architecture using in-memory registry. Minor: import inside hot path. |
| litellm/proxy/auth/auth_checks.py | New check_tools_allowlist function added as step 12 in common_checks. No DB calls in hot path — reads from pre-loaded token metadata. Clean implementation. |
| litellm/proxy/guardrails/tool_name_extraction.py | New centralized tool name extraction module supporting OpenAI, Anthropic, Google, and MCP formats. Handler instantiation per call is wasteful but not a bug. |
| litellm/proxy/management_endpoints/tool_management_endpoints.py | Comprehensive new endpoints for tool policy CRUD, overrides, logs, and options. Key hash detection uses in instead of startswith which can cause false matches on hashed tokens containing "sk-". |
| litellm/proxy/db/spend_log_tool_index.py | New module for tracking tool usage in SpendLogToolIndex. Clean implementation with skip_duplicates for idempotency and non-fatal error handling. |
| litellm/types/tool_management.py | Clean Pydantic model definitions. Proper split of call_policy into input_policy/output_policy with correct Literal types. New models for overrides, options, and usage logs. |
| schema.prisma | ToolTable updated with input_policy/output_policy split, new metadata fields. New SpendLogToolIndex table. blocked_tools added to ObjectPermissionTable. All three schema files are in sync. |
| tests/test_litellm/proxy/db/test_tool_registry_writer.py | Comprehensive tests for Prisma ORM operations and ToolPolicyRegistry. Tests cover upsert, list, get, update, policy map, and registry sync. Uses global singleton in one test but isolation is maintained. |
| tests/test_litellm/proxy/test_tools_allowlist_enforcement.py | New test file covering tool name extraction across all API formats and allowlist enforcement. All tests are mock-based with no network calls. Good edge case coverage. |
| ui/litellm-dashboard/src/components/ToolPolicies.tsx | Significant UI overhaul: metric cards, needs-review banner, split input/output policy columns, tool detail navigation. Clean component structure. |
| ui/litellm-dashboard/src/components/ToolDetail.tsx | New detailed tool view with policy management, overrides per team/key, and usage logs. Good use of React Query for caching. |
| ui/litellm-dashboard/src/components/networking.tsx | New API functions for tool detail, policy options, usage logs, and override deletion. Type definitions updated for input/output policy split. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[Incoming Request] --> B{Auth Layer}
B -->|Tools allowed| C{Guardrail Layer}
B -->|Tool not in allowlist| D[Reject 403]
C -->|Stage 1 Block Check| E{Any tool blocked}
E -->|Yes| F[Reject 400]
E -->|No| G{Response path}
G -->|No| H[Pass through]
G -->|Yes| I{Stage 2 Trust Chain}
I -->|Untrusted source detected| J[Reject 400]
I -->|All clear| H
K[ToolPolicyRegistry] -.->|reads| E
K -.->|reads| I
L[Periodic DB Sync] --> K
Last reviewed commit: f21f739
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Comment on lines
921
to
+932
| @@index([policy_id, start_time]) | ||
| } | ||
|
|
||
| // Index for fast "last N logs for tool" from SpendLogs – see how a tool is called in production | ||
| model LiteLLM_SpendLogToolIndex { | ||
| request_id String | ||
| tool_name String // matches LiteLLM_ToolTable.tool_name; join for call_policy etc. | ||
| start_time DateTime | ||
|
|
||
| @@id([request_id, tool_name]) | ||
| @@index([tool_name, start_time]) | ||
| } |
Contributor
There was a problem hiding this comment.
Root schema.prisma not updated to input_policy/output_policy
The LiteLLM_ToolTable model in the root schema.prisma still uses the old call_policy column, while litellm/proxy/schema.prisma and litellm-proxy-extras/.../schema.prisma have been correctly migrated to input_policy + output_policy. This drift means:
- The root schema is out of sync with the other two schema files
- The
SpendLogToolIndexcomment on line 927 still says "join for call_policy" while the other schemas say "join for input_policy" - The root schema is also missing the new
user_agentandlast_used_atcolumns
All three schema files should stay identical for the LiteLLM_ToolTable model definition.
Contributor
Additional Comments (1)
Multiple tests still use the old
All |
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
…put/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
- schema.prisma (all 3): whitespace alignment on agent_id field - proxy_server.py: keep main's new imports (LITELLM_UI_SESSION_DURATION, InFlightRequestsMiddleware) with consistent formatting - ToolPolicies.tsx: keep branch's extracted PolicySelect component and onSelectTool prop, discard main's inline PolicySelect duplicate Made-with: Cursor
Contributor
Author
|
@greptile review again |
Comment on lines
+384
to
+411
| async def add_tool_to_object_permission_blocked( | ||
| prisma_client: "PrismaClient", | ||
| object_permission_id: str, | ||
| tool_name: str, | ||
| ) -> bool: | ||
| """Add tool_name to the permission's blocked_tools if not already present.""" | ||
| if not object_permission_id or not tool_name: | ||
| return False | ||
| try: | ||
| row = await prisma_client.db.litellm_objectpermissiontable.find_unique( | ||
| where={"object_permission_id": object_permission_id}, | ||
| ) | ||
| if row is None: | ||
| return False | ||
| current = list(getattr(row, "blocked_tools", []) or []) | ||
| if tool_name in current: | ||
| return True | ||
| current.append(tool_name) | ||
| await prisma_client.db.litellm_objectpermissiontable.update( | ||
| where={"object_permission_id": object_permission_id}, | ||
| data={"blocked_tools": current}, | ||
| ) | ||
| return True | ||
| except Exception as e: | ||
| verbose_proxy_logger.error( | ||
| "tool_registry_writer add_tool_to_object_permission_blocked error: %s", e | ||
| ) | ||
| return False |
Contributor
There was a problem hiding this comment.
Race condition in read-modify-write
add_tool_to_object_permission_blocked (and the corresponding remove_… function) performs a non-atomic read-modify-write: it reads blocked_tools, appends/removes in Python, then writes the full list back. If two concurrent requests modify the same permission's blocked_tools, one write silently overwrites the other.
Consider using Prisma's push operation for adds or a database-level array operation to make this atomic:
Suggested change
| async def add_tool_to_object_permission_blocked( | |
| prisma_client: "PrismaClient", | |
| object_permission_id: str, | |
| tool_name: str, | |
| ) -> bool: | |
| """Add tool_name to the permission's blocked_tools if not already present.""" | |
| if not object_permission_id or not tool_name: | |
| return False | |
| try: | |
| row = await prisma_client.db.litellm_objectpermissiontable.find_unique( | |
| where={"object_permission_id": object_permission_id}, | |
| ) | |
| if row is None: | |
| return False | |
| current = list(getattr(row, "blocked_tools", []) or []) | |
| if tool_name in current: | |
| return True | |
| current.append(tool_name) | |
| await prisma_client.db.litellm_objectpermissiontable.update( | |
| where={"object_permission_id": object_permission_id}, | |
| data={"blocked_tools": current}, | |
| ) | |
| return True | |
| except Exception as e: | |
| verbose_proxy_logger.error( | |
| "tool_registry_writer add_tool_to_object_permission_blocked error: %s", e | |
| ) | |
| return False | |
| async def add_tool_to_object_permission_blocked( | |
| prisma_client: "PrismaClient", | |
| object_permission_id: str, | |
| tool_name: str, | |
| ) -> bool: | |
| """Add tool_name to the permission's blocked_tools if not already present.""" | |
| if not object_permission_id or not tool_name: | |
| return False | |
| try: | |
| row = await prisma_client.db.litellm_objectpermissiontable.find_unique( | |
| where={"object_permission_id": object_permission_id}, | |
| ) | |
| if row is None: | |
| return False | |
| current = list(getattr(row, "blocked_tools", []) or []) | |
| if tool_name in current: | |
| return True | |
| await prisma_client.db.litellm_objectpermissiontable.update( | |
| where={"object_permission_id": object_permission_id}, | |
| data={"blocked_tools": {"push": [tool_name]}}, | |
| ) | |
| return True | |
| except Exception as e: | |
| verbose_proxy_logger.error( | |
| "tool_registry_writer add_tool_to_object_permission_blocked error: %s", e | |
| ) | |
| return False |
Comment on lines
+307
to
+337
| async def sync_tool_policy_from_db(self, prisma_client: "PrismaClient") -> None: | ||
| """Load all tool policies and object-permission blocked_tools from DB.""" | ||
| try: | ||
| tools = await prisma_client.db.litellm_tooltable.find_many() | ||
| self._tool_input_policies = { | ||
| row.tool_name: getattr(row, "input_policy", "untrusted") or "untrusted" | ||
| for row in tools | ||
| } | ||
| self._tool_output_policies = { | ||
| row.tool_name: getattr(row, "output_policy", "untrusted") or "untrusted" | ||
| for row in tools | ||
| } | ||
|
|
||
| perms = await prisma_client.db.litellm_objectpermissiontable.find_many() | ||
| self._blocked_tools_by_op_id = {} | ||
| for row in perms: | ||
| op_id = getattr(row, "object_permission_id", None) | ||
| blocked = getattr(row, "blocked_tools", None) or [] | ||
| if op_id: | ||
| self._blocked_tools_by_op_id[op_id] = list(blocked) | ||
|
|
||
| self._initialized = True | ||
| verbose_proxy_logger.info( | ||
| "ToolPolicyRegistry: synced %d tool policies and %d object permissions from DB", | ||
| len(self._tool_input_policies), | ||
| len(self._blocked_tools_by_op_id), | ||
| ) | ||
| except Exception as e: | ||
| verbose_proxy_logger.exception( | ||
| "ToolPolicyRegistry sync_tool_policy_from_db error: %s", e | ||
| ) |
Contributor
There was a problem hiding this comment.
Full table scan on every sync
sync_tool_policy_from_db loads the entire LiteLLM_ToolTable and the entire LiteLLM_ObjectPermissionTable into memory on every sync cycle. This runs periodically via _init_non_llm_objects_in_db. For deployments with thousands of tools and permissions, this can cause significant memory pressure and DB load spikes.
Consider adding pagination, delta-based syncing (e.g. using updated_at timestamps), or at least a select clause to only fetch the columns needed (tool_name, input_policy, output_policy for tools; object_permission_id, blocked_tools for permissions).
| """Resolve key (hash or raw) to object_permission_id; create permission if key has none.""" | ||
| from litellm.proxy.proxy_server import hash_token | ||
|
|
||
| hashed = key_hash if "sk-" not in (key_hash or "") else hash_token(key_hash) |
Contributor
There was a problem hiding this comment.
Missing created_at field in _row_to_model fallback for ToolTable schema
The batch_upsert_tools function in tool_registry_writer.py sets "created_at" in the create block but the _row_to_model function (line 23-67 of tool_registry_writer.py) correctly handles it. However, in this endpoint the update_tool_policy route also sets "created_at": now in create_data (line 183 of tool_registry_writer.py), which is redundant since the schema has @default(now()). Not a bug, just slightly redundant.
More importantly: the _resolve_key_hash_to_object_permission_id function at line 349 hashes the key with hash_token(key_hash) only when the input contains "sk-". This means if the caller passes an already-hashed token, it won't be re-hashed. However, if a user passes a raw API key that doesn't start with "sk-" (e.g., custom key formats), it will be used as-is for the DB lookup and won't match.
hashed = key_hash if "sk-" not in (key_hash or "") else hash_token(key_hash)Consider using startswith for clarity and to avoid false positives (e.g., a hash that happens to contain "sk-" somewhere in the middle):
Suggested change
| hashed = key_hash if "sk-" not in (key_hash or "") else hash_token(key_hash) | |
| hashed = key_hash if not (key_hash or "").startswith("sk-") else hash_token(key_hash) |
| object_permission_id, team_object_permission_id = ( | ||
| _get_request_object_permission_ids(request_data) | ||
| ) | ||
| from litellm.proxy.db.tool_registry_writer import get_tool_policy_registry |
Contributor
There was a problem hiding this comment.
Import inside hot path on every call
get_tool_policy_registry is imported from litellm.proxy.db.tool_registry_writer inside apply_guardrail on every invocation (line 160). While Python caches module imports, this pattern adds unnecessary overhead to the guardrail hot path. Move the import to the top of the file or to __init__.
| return STANDALONE_EXTRACTORS[call_type.value](data) | ||
| handler_cls = mappings.get(call_type) | ||
| if handler_cls is not None: | ||
| names = handler_cls().extract_request_tool_names(data) |
Contributor
There was a problem hiding this comment.
New handler instance created on every extraction
On line 82, handler_cls() instantiates a new guardrail translation handler object for every call to extract_request_tool_names. If this runs in the auth hot path via check_tools_allowlist (called on every request), this creates unnecessary GC pressure. Consider caching handler instances or making extract_request_tool_names a class method / static method on the handler.
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>
jquinter
added a commit
that referenced
this pull request
Mar 4, 2026
PR #22732 changed the ToolTable schema (renamed call_policy to input_policy, added output_policy/user_agent/last_used_at columns, updated indexes) but didn't include a migration for these changes. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
3 tasks
3 tasks
7 tasks
fzowl
pushed a commit
to fzowl/litellm
that referenced
this pull request
Jun 24, 2026
* 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>
fzowl
pushed a commit
to fzowl/litellm
that referenced
this pull request
Jun 24, 2026
PR BerriAI#22732 changed the ToolTable schema (renamed call_policy to input_policy, added output_policy/user_agent/last_used_at columns, updated indexes) but didn't include a migration for these changes.
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
[Feat] Add Tool Policies for AI Gateway
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
tests/test_litellm/directory, Adding at least 1 test is a hard requirement - see detailsmake test-unit@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewCI (LiteLLM team)
Branch creation CI run
Link:
CI run for the last commit
Link:
Merge / cherry-pick CI run
Links:
Type
🆕 New Feature
🐛 Bug Fix
🧹 Refactoring
📖 Documentation
🚄 Infrastructure
✅ Test
Changes