feat(proxy): tool policies - auto-discover tools + policy enforcement guardrail - #22041
Merged
Conversation
…ardrail enforcement
- New LiteLLM_ToolTable in schema.prisma to store discovered tools
- Auto-discovery: tools seen in LLM responses get upserted via ToolDiscoveryQueue
(hooks into DBSpendUpdateWriter, same pipeline as spend tracking)
- Management endpoints: GET /v1/tool/list, GET /v1/tool/{name}, POST /v1/tool/policy
- ToolPolicyGuardrail: blocks tool_calls in responses based on policy setting
- UI: Tool Policies page under Guardrails section with policy selector,
filters by policy/team/key, live tail, sortable table
- Unit tests for queue, writer, endpoints, guardrail
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
…ody and /messages API - Add call_count column to LiteLLM_ToolTable; incremented on every flush - Extract tools from request body too (not just response tool_calls): - OpenAI /chat/completions: tools[].function.name - Anthropic /messages pass-through: request_body.tools[].name - Show call_count column in UI table (sortable) - UI: drop dual_llm option, keep only trusted/blocked
Contributor
Greptile SummaryThis PR adds a full Tool Policies system to the LiteLLM proxy, allowing admins to auto-discover tools from LLM responses and enforce per-tool call policies (
Confidence Score: 3/5
|
| Filename | Overview |
|---|---|
| litellm/proxy/db/db_spend_update_writer.py | Integrates tool discovery into the spend update writer with _enqueue_tool_registry_upsert and _flush_tool_discovery_queue. Extracts tool names from 4 sources (MCP, OpenAI request, Anthropic passthrough, response). Errors are non-blocking. |
| litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py | New queue with per-flush-cycle deduplication. Simple and correct implementation; seen-set clears on flush. |
| litellm/proxy/db/tool_registry_writer.py | Raw SQL-based tool registry CRUD. Multiple previous-thread concerns (N+1 upserts, PostgreSQL-specific SQL). Uses parameterized queries so no injection risk. |
| litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py | Tool policy enforcement guardrail with per-tool-name caching. Has DB query on cache miss in request path (noted in previous threads). Cache design is per-tool-name which is good for avoiding unnecessary invalidation. |
| litellm/proxy/management_endpoints/tool_management_endpoints.py | Three CRUD endpoints for tool management. Uses inline imports (standard pattern for avoiding circular deps). No pagination (noted in previous threads). |
| litellm/proxy/schema.prisma | New LiteLLM_ToolTable model. Has redundant index on tool_name (noted in previous threads). Schema design is reasonable. |
| litellm/types/tool_management.py | Clean Pydantic models for tool management API. No fastapi imports (correctly placed outside proxy/). |
| ui/litellm-dashboard/src/components/ToolPolicies.tsx | New UI component for tool policies. Has a logic bug: policyStyle fallback renders untrusted tools (the default) with blocked styling. Uses deprecated Tremor Switch component. |
Sequence Diagram
sequenceDiagram
participant Client
participant Proxy as LiteLLM Proxy
participant Guard as ToolPolicyGuardrail
participant Cache as DualCache
participant DB as LiteLLM_ToolTable
participant LLM
participant Writer as DBSpendUpdateWriter
participant Queue as ToolDiscoveryQueue
Client->>Proxy: POST /chat/completions (with tools)
Proxy->>Guard: pre_call (check request tools)
Guard->>Cache: lookup tool policies
alt Cache miss
Cache->>DB: SELECT call_policy WHERE tool_name IN (...)
DB-->>Cache: {tool: policy}
end
Cache-->>Guard: policy map
alt Any tool blocked
Guard-->>Client: HTTP 400 (tool blocked)
end
Proxy->>LLM: Forward request
LLM-->>Proxy: Response with tool_calls
Proxy->>Guard: post_call (check response tool_calls)
Guard->>Cache: lookup tool policies
Cache-->>Guard: policy map
alt Any tool_call blocked
Guard-->>Client: HTTP 400 (tool blocked)
end
Proxy-->>Client: Response
Proxy->>Writer: update_database (async)
Writer->>Queue: _enqueue_tool_registry_upsert
Note over Queue: Dedup by tool_name per flush cycle
Queue-->>Writer: flush() → batch items
Writer->>DB: INSERT ... ON CONFLICT (upsert tools)
Last reviewed commit: d4fac55
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
- Remove redundant @@index([tool_name]) from schema.prisma (tool_name has @unique which already creates an index) - Replace gen_random_uuid()::text with str(uuid.uuid4()) for portability - Rewrite test_tool_registry_writer.py to mock execute_raw/query_raw (actual implementation) instead of Prisma model methods - Fix test patches in test_tool_management_endpoints.py to target source modules since imports are inside function bodies - Add "Tool Policies" page title to ToolPolicies.tsx
Contributor
Author
|
@greptile review again |
- Replace NOW() with Python datetime parameter in tool_registry_writer (SQLite portability) - Fix cache key collision in tool_policy_guardrail: use null-byte separator instead of colon - Remove type==function filter from request-side tool extraction to match response-side behavior - Clear seen_tool_names on flush so call_count increments per batch cycle not per pod lifetime
Contributor
Author
|
@greptile-apps please re-review |
Comment on lines
+56
to
+62
| def test_seen_names_persist_across_flushes(queue): | ||
| """Process-local dedup should prevent re-queuing even after a flush.""" | ||
| queue.add_update({"tool_name": "tool_a"}) | ||
| queue.flush() | ||
| queue.add_update({"tool_name": "tool_a"}) # already seen | ||
| items = queue.flush() | ||
| assert len(items) == 0 |
Contributor
There was a problem hiding this comment.
Test contradicts implementation — will fail
This test asserts that _seen_tool_names persists across flushes (expecting 0 items after re-adding tool_a), but the flush() method at tool_discovery_queue.py:52 calls self._seen_tool_names.clear(). After flush(), the seen-set is empty, so tool_a will be re-queued and the second flush() will return 1 item, not 0.
Either the test expectation is wrong (should be assert len(items) == 1) or flush() should not clear the seen-set. The class docstring says "The seen-set is cleared on flush so subsequent batches can re-count the same tool", which aligns with the implementation but contradicts this test.
Suggested change
| def test_seen_names_persist_across_flushes(queue): | |
| """Process-local dedup should prevent re-queuing even after a flush.""" | |
| queue.add_update({"tool_name": "tool_a"}) | |
| queue.flush() | |
| queue.add_update({"tool_name": "tool_a"}) # already seen | |
| items = queue.flush() | |
| assert len(items) == 0 | |
| def test_seen_names_reset_on_flush_allows_requeue(queue): | |
| """After flush clears the seen-set, the same tool can be re-queued.""" | |
| queue.add_update({"tool_name": "tool_a"}) | |
| queue.flush() | |
| queue.add_update({"tool_name": "tool_a"}) # seen-set was cleared by flush | |
| items = queue.flush() | |
| assert len(items) == 1 |
Comment on lines
+35
to
+58
| async def list_tools( | ||
| call_policy: Optional[ToolCallPolicy] = None, | ||
| user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), | ||
| ): | ||
| """ | ||
| List all auto-discovered tools and their call policies. | ||
|
|
||
| Parameters: | ||
| - call_policy: Optional filter — one of "trusted", "untrusted", "dual_llm", "blocked" | ||
| """ | ||
| from litellm.proxy.db.tool_registry_writer import list_tools as db_list_tools | ||
| from litellm.proxy.proxy_server import prisma_client | ||
|
|
||
| if prisma_client is None: | ||
| raise HTTPException( | ||
| status_code=500, detail=CommonProxyErrors.db_not_connected_error.value | ||
| ) | ||
|
|
||
| try: | ||
| tools = await db_list_tools(prisma_client=prisma_client, call_policy=call_policy) | ||
| return ToolListResponse(tools=tools, total=len(tools)) | ||
| except Exception as e: | ||
| verbose_proxy_logger.exception("Error listing tools: %s", e) | ||
| raise HTTPException(status_code=500, detail=str(e)) |
Contributor
There was a problem hiding this comment.
list_tools has no pagination — unbounded result set
The list_tools endpoint returns all rows from LiteLLM_ToolTable with no LIMIT/OFFSET. In deployments where many tools are auto-discovered (e.g., via MCP servers that expose hundreds of tools), this will return an increasingly large response, potentially causing memory pressure and slow API responses.
Consider adding page and page_size query parameters (consistent with other management endpoints in the codebase) and passing them through to the SQL query with LIMIT $N OFFSET $M.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
Comment on lines
+127
to
+148
| from litellm.proxy.db.tool_registry_writer import get_tools_by_names | ||
| from litellm.proxy.proxy_server import prisma_client | ||
|
|
||
| if not tool_names or prisma_client is None: | ||
| return {} | ||
|
|
||
| cache_key = f"tool_policies:{chr(0).join(sorted(tool_names))}" | ||
| cached = await self._policy_cache.async_get_cache(cache_key) | ||
| if cached is not None and isinstance(cached, dict): | ||
| verbose_proxy_logger.debug( | ||
| "ToolPolicyGuardrail: cache hit for tools %s", tool_names | ||
| ) | ||
| return cached | ||
|
|
||
| policy_map = await get_tools_by_names( | ||
| prisma_client=prisma_client, tool_names=tool_names | ||
| ) | ||
| await self._policy_cache.async_set_cache( | ||
| key=cache_key, | ||
| value=policy_map, | ||
| ttl=TOOL_POLICY_CACHE_TTL_SECONDS, | ||
| ) |
Contributor
There was a problem hiding this comment.
Direct DB query in the critical request path on cache miss
_get_policies_cached calls get_tools_by_names() (a raw SQL query) when the cache misses. While the DualCache with TTL mitigates this for steady-state traffic, the first request for each distinct set of tool names will hit the database directly in the request path.
Per project rules, the critical request path should avoid direct DB queries. Consider pre-warming the cache during startup or using the existing get_key/get_team helper pattern, or at minimum documenting that this is an intentional tradeoff (cache-guarded DB read on miss only).
Context Used: Rule from dashboard - What: In critical path of request, there should be no direct db queries. Only allow them to be made ... (source)
Comment on lines
+1
to
+7
| """ | ||
| In-memory buffer for tool registry upserts. | ||
|
|
||
| Unlike SpendUpdateQueue (which aggregates increments), ToolDiscoveryQueue | ||
| uses set-deduplication: each unique tool_name is only queued once per pod | ||
| lifetime, so DB upserts stop entirely after warmup. | ||
| """ |
Contributor
There was a problem hiding this comment.
Module docstring contradicts class docstring and implementation
The module-level docstring (lines 1–7) says "each unique tool_name is only queued once per pod lifetime, so DB upserts stop entirely after warmup." However, the class docstring (lines 17–24) says "The seen-set is cleared on flush so subsequent batches can re-count the same tool," and flush() at line 52 indeed calls self._seen_tool_names.clear().
This inconsistency also causes the test test_seen_names_persist_across_flushes to fail (see related comment). Please update the module docstring to reflect the actual per-flush-cycle dedup behavior.
Comment on lines
+66
to
+81
| await prisma_client.db.execute_raw( | ||
| 'INSERT INTO "LiteLLM_ToolTable" ' | ||
| "(tool_id, tool_name, origin, call_policy, call_count, created_by, updated_by, key_hash, team_id, key_alias) " | ||
| "VALUES ($7, $1, $2, 'untrusted', 1, $3, $3, $4, $5, $6) " | ||
| "ON CONFLICT (tool_name) DO UPDATE SET " | ||
| "call_count = \"LiteLLM_ToolTable\".call_count + 1, " | ||
| "updated_at = $8", | ||
| tool_name, | ||
| origin, | ||
| created_by, | ||
| key_hash, | ||
| team_id, | ||
| key_alias, | ||
| str(uuid.uuid4()), | ||
| now, | ||
| ) |
Contributor
There was a problem hiding this comment.
Raw SQL INSERT missing created_at and updated_at columns
The INSERT statement specifies columns (tool_id, tool_name, origin, call_policy, call_count, created_by, updated_by, key_hash, team_id, key_alias) but does not include created_at or updated_at. While the Prisma schema has @default(now()) for these fields, raw SQL bypasses Prisma's defaults entirely — PostgreSQL will only apply its own column defaults if created_at/updated_at have a DEFAULT at the DB level.
Since the ON CONFLICT branch sets updated_at = $8 (the Python-generated timestamp), the INSERT branch should also explicitly include created_at and updated_at for consistency:
Suggested change
| await prisma_client.db.execute_raw( | |
| 'INSERT INTO "LiteLLM_ToolTable" ' | |
| "(tool_id, tool_name, origin, call_policy, call_count, created_by, updated_by, key_hash, team_id, key_alias) " | |
| "VALUES ($7, $1, $2, 'untrusted', 1, $3, $3, $4, $5, $6) " | |
| "ON CONFLICT (tool_name) DO UPDATE SET " | |
| "call_count = \"LiteLLM_ToolTable\".call_count + 1, " | |
| "updated_at = $8", | |
| tool_name, | |
| origin, | |
| created_by, | |
| key_hash, | |
| team_id, | |
| key_alias, | |
| str(uuid.uuid4()), | |
| now, | |
| ) | |
| await prisma_client.db.execute_raw( | |
| 'INSERT INTO "LiteLLM_ToolTable" ' | |
| "(tool_id, tool_name, origin, call_policy, call_count, created_by, updated_by, key_hash, team_id, key_alias, created_at, updated_at) " | |
| "VALUES ($7, $1, $2, 'untrusted', 1, $3, $3, $4, $5, $6, $8, $8) " | |
| "ON CONFLICT (tool_name) DO UPDATE SET " | |
| "call_count = \"LiteLLM_ToolTable\".call_count + 1, " | |
| "updated_at = $8", | |
| tool_name, | |
| origin, | |
| created_by, | |
| key_hash, | |
| team_id, | |
| key_alias, | |
| str(uuid.uuid4()), | |
| now, | |
| ) |
- Fix test_seen_names_persist_across_flushes to match actual per-flush-cycle behavior - Update module docstring in tool_discovery_queue.py to accurately describe flush behavior - Add created_at/updated_at to raw SQL INSERT in batch_upsert_tools and update_tool_policy
Contributor
Author
|
@greptile-apps please re-review |
Previously the cache key was built from the full set of tool names in a request, so each unique combination of tools got its own cold cache entry and triggered a separate DB query. With N distinct tools across requests this was effectively a DB hit on every request. Now each tool name is cached individually. Cache hits are checked per tool, only missing tools are fetched from DB in a single batch query, and each result is cached separately. Once a tool's policy is warm, any subsequent request using that tool benefits from the cache regardless of what other tools are in the request.
Contributor
Author
|
@greptile-apps please re-review |
Comment on lines
+27
to
+28
| const policyStyle = (p: string) => | ||
| POLICY_OPTIONS.find((o) => o.value === p) ?? POLICY_OPTIONS[1]; |
Contributor
There was a problem hiding this comment.
untrusted tools display as blocked in UI
Newly discovered tools default to call_policy = "untrusted" (per the schema), but POLICY_OPTIONS only contains "trusted" and "blocked". The fallback ?? POLICY_OPTIONS[1] returns the blocked style object, so every untrusted tool will render with red "blocked" styling and show "blocked" in the dropdown — even though its actual policy is "untrusted".
This will mislead admins into thinking all newly discovered tools are blocked when they're actually untrusted (allowed through).
Add "untrusted" to POLICY_OPTIONS or change the fallback to a neutral default:
Suggested change
| const policyStyle = (p: string) => | |
| POLICY_OPTIONS.find((o) => o.value === p) ?? POLICY_OPTIONS[1]; | |
| const policyStyle = (p: string) => | |
| POLICY_OPTIONS.find((o) => o.value === p) ?? { value: p, label: p, color: "#6b7280", bg: "#f3f4f6", border: "#d1d5db" }; |
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Sameerlite
pushed a commit
that referenced
this pull request
Mar 3, 2026
… guardrail (#22041) * feat(proxy): tool policies - auto-discover tools, manage policies, guardrail enforcement - New LiteLLM_ToolTable in schema.prisma to store discovered tools - Auto-discovery: tools seen in LLM responses get upserted via ToolDiscoveryQueue (hooks into DBSpendUpdateWriter, same pipeline as spend tracking) - Management endpoints: GET /v1/tool/list, GET /v1/tool/{name}, POST /v1/tool/policy - ToolPolicyGuardrail: blocks tool_calls in responses based on policy setting - UI: Tool Policies page under Guardrails section with policy selector, filters by policy/team/key, live tail, sortable table - Unit tests for queue, writer, endpoints, guardrail * feat(tool-policies): track call_count + discover tools from request body and /messages API - Add call_count column to LiteLLM_ToolTable; incremented on every flush - Extract tools from request body too (not just response tool_calls): - OpenAI /chat/completions: tools[].function.name - Anthropic /messages pass-through: request_body.tools[].name - Show call_count column in UI table (sortable) - UI: drop dual_llm option, keep only trusted/blocked * fix: address greptile review feedback - Remove redundant @@index([tool_name]) from schema.prisma (tool_name has @unique which already creates an index) - Replace gen_random_uuid()::text with str(uuid.uuid4()) for portability - Rewrite test_tool_registry_writer.py to mock execute_raw/query_raw (actual implementation) instead of Prisma model methods - Fix test patches in test_tool_management_endpoints.py to target source modules since imports are inside function bodies - Add "Tool Policies" page title to ToolPolicies.tsx * fix: address greptile review round 2 - Replace NOW() with Python datetime parameter in tool_registry_writer (SQLite portability) - Fix cache key collision in tool_policy_guardrail: use null-byte separator instead of colon - Remove type==function filter from request-side tool extraction to match response-side behavior - Clear seen_tool_names on flush so call_count increments per batch cycle not per pod lifetime * fix: address greptile review round 3 - Fix test_seen_names_persist_across_flushes to match actual per-flush-cycle behavior - Update module docstring in tool_discovery_queue.py to accurately describe flush behavior - Add created_at/updated_at to raw SQL INSERT in batch_upsert_tools and update_tool_policy * fix: cache tool policies per tool name not per combination Previously the cache key was built from the full set of tool names in a request, so each unique combination of tools got its own cold cache entry and triggered a separate DB query. With N distinct tools across requests this was effectively a DB hit on every request. Now each tool name is cached individually. Cache hits are checked per tool, only missing tools are fetched from DB in a single batch query, and each result is cached separately. Once a tool's policy is warm, any subsequent request using that tool benefits from the cache regardless of what other tools are in the request. * Update ui/litellm-dashboard/src/components/ToolPolicies.tsx 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>
7 tasks
fzowl
pushed a commit
to fzowl/litellm
that referenced
this pull request
Jun 24, 2026
… guardrail (BerriAI#22041) * feat(proxy): tool policies - auto-discover tools, manage policies, guardrail enforcement - New LiteLLM_ToolTable in schema.prisma to store discovered tools - Auto-discovery: tools seen in LLM responses get upserted via ToolDiscoveryQueue (hooks into DBSpendUpdateWriter, same pipeline as spend tracking) - Management endpoints: GET /v1/tool/list, GET /v1/tool/{name}, POST /v1/tool/policy - ToolPolicyGuardrail: blocks tool_calls in responses based on policy setting - UI: Tool Policies page under Guardrails section with policy selector, filters by policy/team/key, live tail, sortable table - Unit tests for queue, writer, endpoints, guardrail * feat(tool-policies): track call_count + discover tools from request body and /messages API - Add call_count column to LiteLLM_ToolTable; incremented on every flush - Extract tools from request body too (not just response tool_calls): - OpenAI /chat/completions: tools[].function.name - Anthropic /messages pass-through: request_body.tools[].name - Show call_count column in UI table (sortable) - UI: drop dual_llm option, keep only trusted/blocked * fix: address greptile review feedback - Remove redundant @@index([tool_name]) from schema.prisma (tool_name has @unique which already creates an index) - Replace gen_random_uuid()::text with str(uuid.uuid4()) for portability - Rewrite test_tool_registry_writer.py to mock execute_raw/query_raw (actual implementation) instead of Prisma model methods - Fix test patches in test_tool_management_endpoints.py to target source modules since imports are inside function bodies - Add "Tool Policies" page title to ToolPolicies.tsx * fix: address greptile review round 2 - Replace NOW() with Python datetime parameter in tool_registry_writer (SQLite portability) - Fix cache key collision in tool_policy_guardrail: use null-byte separator instead of colon - Remove type==function filter from request-side tool extraction to match response-side behavior - Clear seen_tool_names on flush so call_count increments per batch cycle not per pod lifetime * fix: address greptile review round 3 - Fix test_seen_names_persist_across_flushes to match actual per-flush-cycle behavior - Update module docstring in tool_discovery_queue.py to accurately describe flush behavior - Add created_at/updated_at to raw SQL INSERT in batch_upsert_tools and update_tool_policy * fix: cache tool policies per tool name not per combination Previously the cache key was built from the full set of tool names in a request, so each unique combination of tools got its own cold cache entry and triggered a separate DB query. With N distinct tools across requests this was effectively a DB hit on every request. Now each tool name is cached individually. Cache hits are checked per tool, only missing tools are fetched from DB in a single batch query, and each result is cached separately. Once a tool's policy is warm, any subsequent request using that tool benefits from the cache regardless of what other tools are in the request. * Update ui/litellm-dashboard/src/components/ToolPolicies.tsx 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>
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.
Relevant issues
Pre-Submission checklist
tests/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
Changes
Adds a Tool Policies system so admins can see which tools LLMs are calling and block specific ones.
Auto-discovery
tool_callsgets the tool names extracted and upserted into a newLiteLLM_ToolTableDBSpendUpdateWritervia a newToolDiscoveryQueue— same batch/dedup pattern as spend trackingManagement endpoints
GET /v1/tool/list— list all discovered tools with their policiesGET /v1/tool/{name}— get a single toolPOST /v1/tool/policy— update a tool's policy (trustedorblocked)Policy enforcement
ToolPolicyGuardrailcheckstool_callsin LLM responses against stored policiesblockedtools raise HTTP 400 before the response is returned to the callerDualCachewith configurable TTL to avoid DB reads on every requestUI
Schema
LiteLLM_ToolTableinschema.prismaTests — 28 unit tests across:
tests/test_litellm/proxy/db/db_transaction_queue/test_tool_discovery_queue.pytests/test_litellm/proxy/db/test_tool_registry_writer.pytests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.pytests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py