feat: tool allow_list in request params #437
Conversation
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. WalkthroughAdds per-request, per-server tool filtering: Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Client
participant AugmentedLLM
participant Agent
participant ToolRegistry as "Tool Sources"
Client->>AugmentedLLM: generate(input, params{tool_filter})
AugmentedLLM->>Agent: list_tools(server_name?, tool_filter)
Agent->>ToolRegistry: retrieve server-scoped tools & non-namespaced tools
Agent-->>Agent: evaluate non-namespaced via _should_include_non_namespaced_tool
Agent-->>Agent: apply per-server & wildcard exact-name filtering
Agent-->>Agent: collect filtered_out_tools (name + reason) up to 20 for telemetry
Agent-->>AugmentedLLM: ListToolsResult(tools, filtered_out_tools, span attrs)
AugmentedLLM-->>Client: continue generation with filtered tools
note right of Agent #e6f7ff: span attrs set: tool_filter_applied, tools_included_count, tools_filtered_out_count, filtered_tool.i.name/reason
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
Poem
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/mcp_agent/workflows/llm/augmented_llm_anthropic.py (1)
218-224: Only include tools in Anthropic request when non-empty.Prevents sending an empty list, which can be ambiguous across provider SDKs.
- arguments = { - "model": model, - "max_tokens": params.maxTokens, - "messages": messages, - "stop_sequences": params.stopSequences or [], - "tools": available_tools, - } + arguments = { + "model": model, + "max_tokens": params.maxTokens, + "messages": messages, + "stop_sequences": params.stopSequences or [], + } + if available_tools: + arguments["tools"] = available_tools
🧹 Nitpick comments (7)
src/mcp_agent/agents/agent.py (2)
546-559: Enrich tracing when filtering; flag empty results explicitly.Add a couple of span attributes to aid debugging when filters over-restrict.
# Apply tool filtering if specified if tool_filter is not None: original_count = len(result.tools) result.tools = [tool for tool in result.tools if tool.name in tool_filter] filtered_count = original_count - len(result.tools) if filtered_count > 0: logger.debug( f"Tool filter applied: {filtered_count} tools filtered out, {len(result.tools)} tools remaining" ) + span.set_attribute("tool_filter_size", len(tool_filter)) span.set_attribute("tool_filter_applied", True) span.set_attribute("tools_filtered_count", filtered_count) span.set_attribute("tools_remaining_count", len(result.tools)) + if len(result.tools) == 0: + span.set_attribute("tools_empty_after_filter", True)
461-464: Clarify tool_filter semantics and accept Iterable inputs
- Update
list_toolsdocstring to specify thattool_filterexpects fully-qualified (namespaced) tool names and that an empty set returns no tools.- Change the
tool_filter: Set[str] | Noneparameter to accept anyIterable[str]and normalize it internally to aSet[str].
No positional-argument usage detected at call sites—all invocations use named parameters.src/mcp_agent/workflows/llm/augmented_llm_openai.py (1)
192-204: Rename local variable to avoid confusion with later “response”.The variable named response becomes a ChatCompletion later; renaming improves readability.
- response: ListToolsResult = await self.agent.list_tools(tool_filter=params.tool_filter) + tools_result: ListToolsResult = await self.agent.list_tools(tool_filter=params.tool_filter) available_tools: List[ChatCompletionToolParam] = [ ChatCompletionToolParam( type="function", function={ "name": tool.name, "description": tool.description, "parameters": tool.inputSchema, # TODO: saqadri - determine if we should specify "strict" to True by default }, ) - for tool in response.tools + for tool in tools_result.tools ]src/mcp_agent/workflows/llm/augmented_llm_anthropic.py (1)
183-191: Guard for empty allow-list results and rename variable.Mirror the OpenAI path: if no tools after filtering, don’t send the tools field; also rename for clarity.
- list_tools_result = await self.agent.list_tools(tool_filter=params.tool_filter) - available_tools: List[ToolParam] = [ + tools_result = await self.agent.list_tools(tool_filter=params.tool_filter) + available_tools: List[ToolParam] = [ { "name": tool.name, "description": tool.description, "input_schema": tool.inputSchema, } - for tool in list_tools_result.tools + for tool in tools_result.tools ]src/mcp_agent/workflows/llm/augmented_llm_google.py (1)
92-105: Forwarding tool_filter looks good; add a guardrail log when it filters to zero.Great to see
tool_filterpassed through. Add a small warning to help diagnose empty tool sets at runtime.Apply this diff:
- response = await self.agent.list_tools(tool_filter=params.tool_filter) + response = await self.agent.list_tools(tool_filter=params.tool_filter) + + # Warn if a provided filter yields no tools + if params.tool_filter is not None and not response.tools: + self.logger.warning( + "No tools match the provided tool_filter", + data={"tool_filter": list(params.tool_filter)}, + )src/mcp_agent/workflows/llm/augmented_llm_azure.py (1)
163-174: Forwarding tool_filter is correct; avoid sending an empty tools list.Azure may accept an empty list, but conditionally including the key mirrors Bedrock and reduces surface for 400s.
Apply this diff:
- response = await self.agent.list_tools(tool_filter=params.tool_filter) + response = await self.agent.list_tools(tool_filter=params.tool_filter) @@ - arguments = { - "messages": messages, - "temperature": params.temperature, - "model": model, - "max_tokens": params.maxTokens, - "stop": params.stopSequences, - "tools": tools, - } + arguments = { + "messages": messages, + "temperature": params.temperature, + "model": model, + "max_tokens": params.maxTokens, + "stop": params.stopSequences, + } + if tools: + arguments["tools"] = toolsOptionally also log when
tool_filterproduces zero tools (similar to the Google change).Also applies to: 189-198
src/mcp_agent/workflows/llm/augmented_llm.py (1)
548-551: Pass-through to agent is correct; consider tracing the filter at the LLM layer.Agent-level spans capture filtering, but adding a count at the LLM layer can ease correlation.
For example, in
annotate_span_with_request_params, add:
- request_params.tool_filter_count
- request_params.tool_filter_sample (first few names)
Let me patch that if you want.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (8)
src/mcp_agent/agents/agent.py(3 hunks)src/mcp_agent/workflows/llm/augmented_llm.py(3 hunks)src/mcp_agent/workflows/llm/augmented_llm_anthropic.py(1 hunks)src/mcp_agent/workflows/llm/augmented_llm_azure.py(1 hunks)src/mcp_agent/workflows/llm/augmented_llm_bedrock.py(1 hunks)src/mcp_agent/workflows/llm/augmented_llm_google.py(1 hunks)src/mcp_agent/workflows/llm/augmented_llm_openai.py(1 hunks)tests/workflows/llm/test_request_params_tool_filter.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (8)
src/mcp_agent/workflows/llm/augmented_llm_anthropic.py (3)
src/mcp_agent/agents/agent.py (1)
list_tools(461-563)src/mcp_agent/workflows/llm/augmented_llm.py (1)
list_tools(548-550)src/mcp_agent/mcp/mcp_aggregator.py (1)
list_tools(645-688)
src/mcp_agent/workflows/llm/augmented_llm_openai.py (3)
src/mcp_agent/agents/agent.py (1)
list_tools(461-563)src/mcp_agent/workflows/llm/augmented_llm.py (1)
list_tools(548-550)src/mcp_agent/mcp/mcp_aggregator.py (1)
list_tools(645-688)
tests/workflows/llm/test_request_params_tool_filter.py (2)
src/mcp_agent/workflows/llm/augmented_llm.py (6)
RequestParams(125-181)get(95-96)get(118-119)list_tools(548-550)AugmentedLLM(236-824)get_request_params(388-409)src/mcp_agent/agents/agent.py (1)
list_tools(461-563)
src/mcp_agent/workflows/llm/augmented_llm_bedrock.py (3)
src/mcp_agent/agents/agent.py (1)
list_tools(461-563)src/mcp_agent/workflows/llm/augmented_llm.py (1)
list_tools(548-550)src/mcp_agent/mcp/mcp_aggregator.py (1)
list_tools(645-688)
src/mcp_agent/workflows/llm/augmented_llm_azure.py (3)
src/mcp_agent/agents/agent.py (1)
list_tools(461-563)src/mcp_agent/workflows/llm/augmented_llm.py (1)
list_tools(548-550)src/mcp_agent/mcp/mcp_aggregator.py (1)
list_tools(645-688)
src/mcp_agent/workflows/llm/augmented_llm_google.py (2)
src/mcp_agent/agents/agent.py (1)
list_tools(461-563)src/mcp_agent/workflows/llm/augmented_llm.py (1)
list_tools(548-550)
src/mcp_agent/agents/agent.py (2)
src/mcp_agent/workflows/llm/augmented_llm.py (1)
list_tools(548-550)src/mcp_agent/mcp/mcp_aggregator.py (1)
list_tools(645-688)
src/mcp_agent/workflows/llm/augmented_llm.py (3)
src/mcp_agent/agents/agent.py (1)
list_tools(461-563)src/mcp_agent/mcp/mcp_aggregator.py (1)
list_tools(645-688)tests/mcp/test_mcp_aggregator.py (1)
list_tools(787-794)
🔇 Additional comments (7)
src/mcp_agent/agents/agent.py (2)
4-4: Import update looks good.Set typing added appropriately for the new parameter.
532-543: Conditional human_input tool injection is correct.Good change to avoid early return; tool remains subject to the allow-list as intended.
src/mcp_agent/workflows/llm/augmented_llm_openai.py (1)
211-213: Nice: empty tool set handled.Setting available_tools to None when empty avoids sending an empty tools list to OpenAI.
src/mcp_agent/workflows/llm/augmented_llm_bedrock.py (1)
110-124: LGTM: tool_filter is correctly forwarded and toolConfig is omitted when empty.This aligns with Bedrock’s expectations and avoids sending an empty tool spec.
src/mcp_agent/workflows/llm/augmented_llm.py (1)
9-9: Importing Set for the new field is correct.No issues with the typing change.
tests/workflows/llm/test_request_params_tool_filter.py (2)
12-26: Solid coverage for defaults, mixed params, and model_dump behavior.These tests pin important backwards-compat guarantees.
Also applies to: 53-88
114-126: Nice guard to ensure legacy calls remain argument-free.Keeps compatibility with existing
list_tools()usage.
![graphite-app[bot]](https://avatars.githubusercontent.com/in/158384?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/mcp_agent/agents/agent.py (1)
493-556: Wildcard filter not applied when server_name is provided.If server_name is specified and no explicit entry exists for it, you currently ignore the "*" rules. This diverges from the all‑servers path and can leak tools unintentionally.
Apply:
- # Check if we should apply filtering for this specific server - if tool_filter is not None and server_name in tool_filter: + # Check if we should apply filtering for this specific server + if tool_filter is not None and (server_name in tool_filter or "*" in tool_filter): # Server is explicitly in filter dict - apply its filter rules - # If tool_filter[server_name] is empty set, no tools will pass - # If tool_filter[server_name] has tools, only those will pass - allowed_tools = tool_filter[server_name] + # Choose server-specific or fallback to wildcard set + allowed_tools = tool_filter.get(server_name, tool_filter.get("*", set())) result_tools = [] for namespaced_tool in server_tools: if namespaced_tool.tool.name in allowed_tools: result_tools.append( namespaced_tool.tool.model_copy( update={"name": namespaced_tool.namespaced_tool_name} ) ) else: filtered_out_tools.append( - (namespaced_tool.namespaced_tool_name, f"Not in tool_filter[{server_name}]") + ( + namespaced_tool.namespaced_tool_name, + ( + f"Not in tool_filter[{server_name}]" + if server_name in tool_filter + else "Not in tool_filter[*]" + ), + ) ) result = ListToolsResult(tools=result_tools)Add a unit test for list_tools(server_name=..., tool_filter={"*": {...}}).
♻️ Duplicate comments (1)
src/mcp_agent/workflows/llm/augmented_llm.py (1)
178-196: Harden tool_filter parsing to avoid string→set‑of‑chars footgun.Pydantic will coerce {"": "toolA"} into {"": {"t","o","o","l","A"}}. Validate and coerce dict values to sets of strings; reject plain strings.
-from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict, Field, field_validator @@ class RequestParams(CreateMessageRequestParams): @@ - tool_filter: Dict[str, Set[str]] | None = None + tool_filter: Dict[str, Set[str]] | None = None @@ Tool names should match exactly as they appear in the server's tool list. """ + + @field_validator("tool_filter", mode="before") + @classmethod + def _validate_tool_filter(cls, v): + if v is None: + return None + if not isinstance(v, dict): + raise TypeError("tool_filter must be a dict[str, set[str]]") + normalized: dict[str, set[str]] = {} + for k, val in v.items(): + if isinstance(val, str): + raise TypeError(f"tool_filter['{k}'] must be an iterable of strings, not a string") + try: + s = set(val) + except TypeError: + raise TypeError(f"tool_filter['{k}'] must be a set or iterable of strings") + if not all(isinstance(x, str) for x in s): + raise TypeError(f"tool_filter['{k}'] elements must be strings") + normalized[str(k)] = s + return normalizedAlso update the example to use "human_input" (exact tool name) unless you accept the alias.
- - {"non_namespaced_tools": {"human_input", "func1"}} - Allow specific non-namespaced tools + - {"non_namespaced_tools": {"__human_input__", "func1"}} - Allow specific non-namespaced tools
🧹 Nitpick comments (4)
src/mcp_agent/agents/agent.py (4)
462-492: Non‑namespaced tool filter helper: solid; consider alias for human tool name.Helper is clear and correctly prioritizes non_namespaced_tools over "*". Optional: accept "human_input" as an alias for "human_input" to avoid config mismatches, or document the exact required name everywhere.
Apply this minimal alias support:
def _should_include_non_namespaced_tool( self, tool_name: str, tool_filter: Dict[str, Set[str]] | None ) -> tuple[bool, str | None]: @@ - if "non_namespaced_tools" in tool_filter: + if "non_namespaced_tools" in tool_filter: + allowed = tool_filter["non_namespaced_tools"] + # Back-compat alias for human input + if tool_name == HUMAN_INPUT_TOOL_NAME and "human_input" in allowed: + return True, None - if tool_name in tool_filter["non_namespaced_tools"]: + if tool_name in allowed: return True, None else: return False, f"{tool_name} not in tool_filter[non_namespaced_tools]"
557-593: Concurrent read safety: copy before iterating shared maps._result may change during re-initialization. Copy items to avoid RuntimeError/tearing.
- for namespaced_tool_name, namespaced_tool in self._namespaced_tool_map.items(): + for namespaced_tool_name, namespaced_tool in list(self._namespaced_tool_map.items()):
647-669: Human input tool: align naming in docs/configs.Tool name is "human_input". Ensure examples and tests don’t instruct "human_input". If keeping the alias, reflect it in docs; otherwise, update examples to the exact name.
670-699: Use logging constants, not numeric level 10.10 is DEBUG in stdlib; there’s no TRACE. Use logging.DEBUG or add a custom TRACE level if supported.
- if logger.isEnabledFor(10): # TRACE level is usually 10 + import logging + if logger.isEnabledFor(logging.DEBUG):
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (8)
src/mcp_agent/agents/agent.py(4 hunks)src/mcp_agent/workflows/llm/augmented_llm.py(3 hunks)src/mcp_agent/workflows/llm/augmented_llm_anthropic.py(1 hunks)src/mcp_agent/workflows/llm/augmented_llm_azure.py(1 hunks)src/mcp_agent/workflows/llm/augmented_llm_bedrock.py(1 hunks)src/mcp_agent/workflows/llm/augmented_llm_google.py(1 hunks)src/mcp_agent/workflows/llm/augmented_llm_openai.py(1 hunks)tests/workflows/llm/test_request_params_tool_filter.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (5)
- src/mcp_agent/workflows/llm/augmented_llm_azure.py
- src/mcp_agent/workflows/llm/augmented_llm_openai.py
- src/mcp_agent/workflows/llm/augmented_llm_bedrock.py
- src/mcp_agent/workflows/llm/augmented_llm_anthropic.py
- tests/workflows/llm/test_request_params_tool_filter.py
🧰 Additional context used
🧬 Code graph analysis (3)
src/mcp_agent/workflows/llm/augmented_llm_google.py (3)
src/mcp_agent/agents/agent.py (1)
list_tools(493-703)src/mcp_agent/workflows/llm/augmented_llm.py (1)
list_tools(563-565)src/mcp_agent/mcp/mcp_aggregator.py (1)
list_tools(636-679)
src/mcp_agent/agents/agent.py (2)
src/mcp_agent/workflows/llm/augmented_llm.py (3)
list_tools(563-565)get(96-97)get(119-120)src/mcp_agent/mcp/mcp_aggregator.py (1)
list_tools(636-679)
src/mcp_agent/workflows/llm/augmented_llm.py (3)
src/mcp_agent/agents/agent.py (1)
list_tools(493-703)src/mcp_agent/mcp/mcp_aggregator.py (1)
list_tools(636-679)tests/mcp/test_mcp_aggregator.py (1)
list_tools(787-794)
🔇 Additional comments (4)
src/mcp_agent/agents/agent.py (1)
605-622: Function tools inclusion: LGTM.src/mcp_agent/workflows/llm/augmented_llm_google.py (1)
92-92: Propagating tool_filter into list_tools: LGTM.src/mcp_agent/workflows/llm/augmented_llm.py (2)
5-11: Typing imports updated: LGTM.
563-566: AugmentedLLM.list_tools forwards tool_filter correctly: LGTM.Please confirm all provider backends call agent.list_tools with tool_filter (Anthropic, Azure, Bedrock, OpenAI). If any remain, add the param.
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (1)
tests/workflows/llm/test_request_params_tool_filter.py (1)
626-645: Tighten invalid‑type expectations: plain string must raise; list values should coerce to set.Match the model’s validator semantics and avoid permissive try/except.
def test_request_params_with_invalid_tool_filter_type(self): """Test that RequestParams handles invalid tool_filter types gracefully.""" - # Test with string (should cause type error) - try: - params = RequestParams(tool_filter="invalid_string") - # If no exception, it's being converted somehow - assert isinstance(params.tool_filter, dict) or params.tool_filter is None - except (ValueError, TypeError): - pass # This is expected behavior + # String should be rejected by the validator + with pytest.raises((TypeError, ValueError)): + RequestParams(tool_filter="invalid_string") - # Test with dict having non-set values (should convert or error) - try: - params_with_list = RequestParams(tool_filter={"server1": ["tool1", "tool2"]}) - # Pydantic should convert list to set - if params_with_list.tool_filter: - assert isinstance(params_with_list.tool_filter["server1"], set) - assert params_with_list.tool_filter["server1"] == {"tool1", "tool2"} - except (ValueError, TypeError): - pass # This is also acceptable behavior + # Dict with list values should coerce to set + params_with_list = RequestParams(tool_filter={"server1": ["tool1", "tool2"]}) + assert isinstance(params_with_list.tool_filter["server1"], set) + assert params_with_list.tool_filter["server1"] == {"tool1", "tool2"}
🧹 Nitpick comments (8)
tests/workflows/llm/test_request_params_tool_filter.py (8)
2-2: Fix Ruff F401: remove unused import 'patch'.
patchis never used; CI flags this. Drop it to unblock the pipeline.-from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock
211-244: Don’t re‑implement Agent.list_tools in tests; exercise the real method.This helper duplicates production logic, so tests can pass while the real implementation regresses. Call Agent.list_tools on a minimally constructed Agent instead.
Example pattern (no diff; illustrative):
@pytest.mark.asyncio async def test_specific_tools_filter_integration(monkeypatch, mock_agent_with_tools): # Build a real Agent instance without running __init__ from mcp_agent.agents.agent import Agent agent = Agent.__new__(Agent) agent.name = "test_agent" agent.initialized = True agent.context = type("Ctx", (), {"tracing_enabled": False})() agent._server_to_tool_map = mock_agent_with_tools._server_to_tool_map agent._namespaced_tool_map = {} agent._function_tool_map = {} agent.human_input_callback = None # Stub tracer to a no‑op class _Span: def __enter__(self): return self def __exit__(self, *a): pass def set_attribute(self, *a, **k): pass class _Tracer: def start_as_current_span(self, *_): return _Span() monkeypatch.setattr("mcp_agent.agents.agent.get_tracer", lambda _ctx: _Tracer()) result = await Agent.list_tools(agent, server_name="server1", tool_filter={"server1": {"tool1", "tool3"}}) assert {t.name for t in result.tools} == {"server1:tool1", "server1:tool3"}
374-417: Same issue: remove duplicated “all servers” filtering helper and hit production code.Mirror of the above—call Agent.list_tools(server_name=None, …) on a real Agent. This ensures wildcard and per‑server precedence match the shipping logic.
500-512: Low‑value test: asserting a mocked return doesn’t validate compatibility.This only verifies your AsyncMock returns two tools. Prefer a real call (or at least assert the call signature under different argument combinations).
592-616: Edge‑case test isn’t exercising behavior.You’re asserting dict contents, not the filtering result. Convert to an assertion on the actual names returned by list_tools for both servers.
618-625: Same: test doesn’t hit code under test.Call Agent.list_tools(server_name="nonexistent") and assert an empty result instead of checking a plain dict .get().
420-474: Good coverage of non‑namespaced filtering rules. Consider adding human‑input gating tests.Add cases proving “human_input” is only exposed when human_input_callback is set, and that tool_filter still applies.
Example:
@pytest.mark.asyncio async def test_human_input_included_only_with_callback(monkeypatch): from mcp_agent.agents.agent import Agent agent = Agent.__new__(Agent) agent.name = "test_agent" agent.initialized = True agent.context = type("Ctx", (), {"tracing_enabled": False})() agent._namespaced_tool_map = {} agent._server_to_tool_map = {} agent._function_tool_map = {} # no function tools # tracer stub class _Span: def __enter__(self): return self def __exit__(self,*a): pass def set_attribute(self,*a,**k): pass class _Tracer: def start_as_current_span(self,*_): return _Span() monkeypatch.setattr("mcp_agent.agents.agent.get_tracer", lambda _ctx: _Tracer()) # No callback → no human_input agent.human_input_callback = None res = await Agent.list_tools(agent, tool_filter={"non_namespaced_tools": {"human_input"}}) assert all(t.name != "human_input" for t in res.tools) # With callback → human_input allowed by filter agent.human_input_callback = lambda prompt: None res = await Agent.list_tools(agent, tool_filter={"non_namespaced_tools": {"human_input"}}) assert any(t.name == "human_input" for t in res.tools)
175-186: Add a server‑scoped + wildcard test to lock in precedence.When listing a specific server not present in tool_filter, wildcard should NOT apply (by design). Add an assertion for server_name="server2" with tool_filter={"*": {"tool1"}} → include all server2 tools.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
tests/workflows/llm/test_request_params_tool_filter.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
tests/workflows/llm/test_request_params_tool_filter.py (3)
src/mcp_agent/workflows/llm/augmented_llm.py (8)
RequestParams(126-196)get(96-97)get(119-120)append(93-94)append(116-117)list_tools(563-565)AugmentedLLM(251-839)get_request_params(403-424)src/mcp_agent/agents/agent.py (2)
_should_include_non_namespaced_tool(462-491)list_tools(493-703)src/mcp_agent/mcp/mcp_aggregator.py (1)
NamespacedTool(52-59)
🪛 GitHub Actions: Pull Request Checks
tests/workflows/llm/test_request_params_tool_filter.py
[error] 2-2: Ruff: F401 'patch' imported but unused in tests/workflows/llm/test_request_params_tool_filter.py. Remove the import or run 'ruff --fix'.
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (1)
tests/workflows/llm/test_request_params_tool_filter.py (1)
626-635: Invalid string input should raise, not be coerced.Align this test with the validator’s expected behavior.
Apply this diff:
- # Test with string (should cause type error) - try: - params = RequestParams(tool_filter="invalid_string") - # If no exception, it's being converted somehow - assert isinstance(params.tool_filter, dict) or params.tool_filter is None - except (ValueError, TypeError): - pass # This is expected behavior + # Test with string must raise (validator rejects plain strings) + with pytest.raises((TypeError, ValueError)): + RequestParams(tool_filter="invalid_string")
🧹 Nitpick comments (8)
tests/workflows/llm/test_request_params_tool_filter.py (8)
96-103: Tighten exclude_unset expectation.exclude_unset should omit the key entirely when the field isn’t provided.
Apply this diff:
- # tool_filter should not be in dumped output if not set - assert "tool_filter" not in dumped1 or dumped1.get("tool_filter") is None + # tool_filter should not be in dumped output if not set + assert "tool_filter" not in dumped1
199-205: Fix docstring to match the asserted filter.Docstring says {"tool1","tool2"} but test uses {"tool1","tool3"}.
Apply this diff:
- """Test: tool_filter[server_name] = {"tool1", "tool2"} → Only include those tools.""" + """Test: tool_filter[server_name] = {"tool1", "tool3"} → Only include those tools."""
211-244: Remove unused variable and dead accumulation.filtered_out_tools is never asserted; keep the test lean.
Apply this diff:
- filtered_out_tools = [] + # no filtered_out_tools tracking needed in this unit test @@ - else: - filtered_out_tools.append( - (namespaced_tool.namespaced_tool_name, f"Not in tool_filter[{server_name}]") - ) + else: + pass # filtered out
374-417: Same here: drop unused filtered_out_tools.Removes noise and potential lint warnings.
Apply this diff:
- filtered_out_tools = [] + # no filtered_out_tools tracking needed in this unit test @@ - else: - filtered_out_tools.append( - (namespaced_tool_name, f"Not in tool_filter[{namespaced_tool.server_name}]") - ) + else: + pass # filtered out @@ - else: - filtered_out_tools.append( - (namespaced_tool_name, "Not in tool_filter[*]") - ) + else: + pass # filtered out by wildcard
559-565: Avoid variable name collision (‘result3’ reused).Reusing result3 makes the section harder to read and can mask failures.
Apply this diff:
- # Test 3: Override with existing params only + # Test 4: Override with existing params only override_params2 = RequestParams(temperature=0.9) - result3 = llm.get_request_params(request_params=override_params2) - assert result3.maxTokens == 1000 # From default - assert result3.temperature == 0.9 # From override - assert result3.tool_filter is None # Default + result4 = llm.get_request_params(request_params=override_params2) + assert result4.maxTokens == 1000 # From default + assert result4.temperature == 0.9 # From override + assert result4.tool_filter is None # Default
636-645: Assert list→set coercion explicitly (no try/except).This makes the expectation crisp and fails loudly if coercion changes.
Apply this diff:
- # Test with dict having non-set values (should convert or error) - try: - params_with_list = RequestParams(tool_filter={"server1": ["tool1", "tool2"]}) - # Pydantic should convert list to set - if params_with_list.tool_filter: - assert isinstance(params_with_list.tool_filter["server1"], set) - assert params_with_list.tool_filter["server1"] == {"tool1", "tool2"} - except (ValueError, TypeError): - pass # This is also acceptable behavior + # Test with dict having list values is coerced to set + params_with_list = RequestParams(tool_filter={"server1": ["tool1", "tool2"]}) + assert isinstance(params_with_list.tool_filter["server1"], set) + assert params_with_list.tool_filter["server1"] == {"tool1", "tool2"}
592-616: Make this test actually exercise the filtering logic.Currently it only inspects the filter dict, not outcomes.
Apply this diff:
- # Filter should work independently for each server - tool_filter = {"server1": {"tool1"}, "server2": set()} - - # server1:tool1 should be included, server2:tool1 should not - assert "server1" in tool_filter - assert "tool1" in tool_filter["server1"] - assert "server2" in tool_filter - assert len(tool_filter["server2"]) == 0 + # Filter should work independently for each server + tool_filter = {"server1": {"tool1"}, "server2": set()} + + # Compute included tools per the all-servers filtering semantics + included = [] + for namespaced_name, nt in agent._namespaced_tool_map.items(): + if nt.server_name in tool_filter: + if nt.tool.name in tool_filter[nt.server_name]: + included.append(namespaced_name) + elif "*" in tool_filter: + if nt.tool.name in tool_filter["*"]: + included.append(namespaced_name) + else: + included.append(namespaced_name) + + assert set(included) == {"server1:tool1"}
617-625: Assert the actual result object for a missing server.Exercise the same shape as list_tools(server_name=...) would produce.
Apply this diff:
- # Should return empty list, not error - server_tools = agent._server_to_tool_map.get("nonexistent", []) - assert server_tools == [] + # Should return empty tools list, not error + server_tools = agent._server_to_tool_map.get("nonexistent", []) + result = ListToolsResult( + tools=[ + nt.tool.model_copy(update={"name": nt.namespaced_tool_name}) + for nt in server_tools + ] + ) + assert result.tools == []
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
tests/workflows/llm/test_request_params_tool_filter.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
tests/workflows/llm/test_request_params_tool_filter.py (4)
src/mcp_agent/core/context.py (2)
mcp(102-103)Context(57-103)src/mcp_agent/workflows/llm/augmented_llm.py (8)
RequestParams(126-196)get(96-97)get(119-120)append(93-94)append(116-117)list_tools(563-565)AugmentedLLM(251-839)get_request_params(403-424)src/mcp_agent/agents/agent.py (3)
Agent(62-1142)_should_include_non_namespaced_tool(462-491)list_tools(493-703)src/mcp_agent/mcp/mcp_aggregator.py (1)
NamespacedTool(52-59)
🔇 Additional comments (1)
tests/workflows/llm/test_request_params_tool_filter.py (1)
36-41: Nice coverage of non_namespaced_tools behavior.These assertions clearly validate inclusion/exclusion semantics for human_input and function tools.
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
tests/workflows/llm/test_request_params_tool_filter.py (1)
626-635: Invalid string input must raise; don’t silently coerce.Aligns with prior feedback; assert a validation error.
-# Test with string (should cause type error) -try: - params = RequestParams(tool_filter="invalid_string") - # If no exception, it's being converted somehow - assert isinstance(params.tool_filter, dict) or params.tool_filter is None -except (ValueError, TypeError): - pass # This is expected behavior +# Test with string must raise (validator rejects plain strings) +with pytest.raises((TypeError, ValueError)): + RequestParams(tool_filter="invalid_string")
🧹 Nitpick comments (5)
tests/workflows/llm/test_request_params_tool_filter.py (5)
6-10: Fix incorrect NamespacedTool import (use agents.agent).Importing NamespacedTool from mcp_aggregator is inconsistent with production (it’s defined in mcp_agent.agents.agent). Use the same source as Agent to avoid drift.
from mcp_agent.workflows.llm.augmented_llm import RequestParams -from mcp_agent.agents.agent import Agent -from mcp_agent.mcp.mcp_aggregator import NamespacedTool +from mcp_agent.agents.agent import Agent, NamespacedTool from mcp_agent.core.context import Context
96-103: Tighten exclude_unset assertion.Be explicit: when unset, tool_filter should be entirely absent from the dump.
dumped1 = params1.model_dump(exclude_unset=True) -# tool_filter should not be in dumped output if not set -assert "tool_filter" not in dumped1 or dumped1.get("tool_filter") is None +# tool_filter should not be present if not set +assert "tool_filter" not in dumped1
420-474: Non-namespaced precedence test missing.Add a test that when both non_namespaced_tools and "*" are present, non_namespaced_tools takes precedence (as per implementation priority).
+def test_non_namespaced_precedence_over_wildcard(self): + from mcp_agent.agents.agent import Agent + agent = MagicMock(spec=Agent) + agent._should_include_non_namespaced_tool = Agent._should_include_non_namespaced_tool.__get__(agent) + should_include, reason = agent._should_include_non_namespaced_tool( + "funcX", {"non_namespaced_tools": set(), "*": {"funcX"}} + ) + # non_namespaced_tools key should control; funcX excluded despite wildcard + assert should_include is False + assert "not in tool_filter[non_namespaced_tools]" in reason
500-523: Add a propagation test for AugmentedLLM.list_tools.We only check signature. Add a test that verifies tool_filter is forwarded to agent.list_tools.
@pytest.mark.asyncio async def test_existing_code_with_server_name_still_works(self, mock_agent): ... mock_agent.list_tools.assert_called_with(server_name="test_server") +@pytest.mark.asyncio +async def test_augmented_llm_list_tools_propagates_filter(self, mock_context): + from mcp_agent.workflows.llm.augmented_llm import AugmentedLLM + # Create a minimal instance without calling __init__ + llm = object.__new__(AugmentedLLM) + llm.agent = MagicMock() + llm.agent.list_tools = AsyncMock(return_value=ListToolsResult(tools=[])) + tool_filter = {"*": {"t1"}} + res = await AugmentedLLM.list_tools(llm, server_name="s1", tool_filter=tool_filter) + llm.agent.list_tools.assert_called_with(server_name="s1", tool_filter=tool_filter) + assert res.tools == []
637-645: Assert list→set coercion deterministically.Make the expected coercion explicit; avoid permissive try/except.
-# Test with dict having non-set values (should convert or error) -try: - params_with_list = RequestParams(tool_filter={"server1": ["tool1", "tool2"]}) - # Pydantic should convert list to set - if params_with_list.tool_filter: - assert isinstance(params_with_list.tool_filter["server1"], set) - assert params_with_list.tool_filter["server1"] == {"tool1", "tool2"} -except (ValueError, TypeError): - pass # This is also acceptable behavior +# Test dict with list values coerces to set +params_with_list = RequestParams(tool_filter={"server1": ["tool1", "tool2"]}) +assert isinstance(params_with_list.tool_filter["server1"], set) +assert params_with_list.tool_filter["server1"] == {"tool1", "tool2"}
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
tests/workflows/llm/test_request_params_tool_filter.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
tests/workflows/llm/test_request_params_tool_filter.py (3)
src/mcp_agent/core/context.py (2)
mcp(102-103)Context(57-103)src/mcp_agent/workflows/llm/augmented_llm.py (8)
RequestParams(126-196)get(96-97)get(119-120)append(93-94)append(116-117)list_tools(563-565)AugmentedLLM(251-839)get_request_params(403-424)src/mcp_agent/agents/agent.py (2)
_should_include_non_namespaced_tool(462-491)list_tools(493-703)
🔇 Additional comments (2)
tests/workflows/llm/test_request_params_tool_filter.py (2)
554-565: LGTM: variable collision fixed.The duplicate result3 assignment noted earlier is resolved by introducing result4; assertions now target the correct object.
211-244: ```shell
#!/bin/bash
set -euo pipefail
echo "=== repo root: $(pwd) ==="
echo "--- git branch ---"
git rev-parse --abbrev-ref HEAD || true
echoecho "== Locate test file =="
rg -n --hidden --no-ignore "test_request_params_tool_filter.py" || true
echoecho "== Show tests/workflows/llm/test_request_params_tool_filter.py (first 400 lines) =="
sed -n '1,400p' tests/workflows/llm/test_request_params_tool_filter.py || true
echoecho "== Search for _apply_list_tools_logic definition/usages across repo =="
rg -n --hidden --no-ignore "def _apply_list_tools_logic|async def _apply_list_tools_logic" -S || true
echoecho "== Search for 'tool_filter' usage across repo (context lines) =="
rg -n --hidden --no-ignore "tool_filter" -S -C2 || true
echoecho "== Search for wildcard '' literal in code/tests (looks for '' keys or string '') =="
rg -n --hidden --no-ignore "'\'|"\"|\:" -S -C2 || true
echoecho "== Search for ListToolsResult occurrences =="
rg -n --hidden --no-ignore "ListToolsResult" -S -C2 || true
echoecho "== Search for server_to_tool_map occurrences =="
rg -n --hidden --no-ignore "server_to_tool_map" -S -C2 || true
echo</blockquote></details> </blockquote></details> </details> <!-- This is an auto-generated comment by CodeRabbit for review status -->
| class TestAgentToolFilteringAllServers: | ||
| """Test cases when server_name is NOT provided (listing all tools).""" | ||
|
|
||
| @pytest.fixture | ||
| def mock_agent_all_servers(self): | ||
| """Create a mock agent with test data.""" | ||
| agent = MagicMock(spec=Agent) | ||
| agent.initialized = True | ||
| agent.context = MagicMock(spec=Context) | ||
| agent.context.tracing_enabled = False | ||
|
|
||
| # Setup namespaced tool map | ||
| agent._namespaced_tool_map = { | ||
| "server1:tool1": NamespacedTool( | ||
| tool=Tool(name="tool1", description="Tool 1", inputSchema={}), | ||
| server_name="server1", | ||
| namespaced_tool_name="server1:tool1" | ||
| ), | ||
| "server1:tool2": NamespacedTool( | ||
| tool=Tool(name="tool2", description="Tool 2", inputSchema={}), | ||
| server_name="server1", | ||
| namespaced_tool_name="server1:tool2" | ||
| ), | ||
| "server2:tool1": NamespacedTool( | ||
| tool=Tool(name="tool1", description="Tool 1", inputSchema={}), | ||
| server_name="server2", | ||
| namespaced_tool_name="server2:tool1" | ||
| ), | ||
| "server2:tool3": NamespacedTool( | ||
| tool=Tool(name="tool3", description="Tool 3", inputSchema={}), | ||
| server_name="server2", | ||
| namespaced_tool_name="server2:tool3" | ||
| ), | ||
| "server3:tool4": NamespacedTool( | ||
| tool=Tool(name="tool4", description="Tool 4", inputSchema={}), | ||
| server_name="server3", | ||
| namespaced_tool_name="server3:tool4" | ||
| ), | ||
| } | ||
|
|
||
| agent._function_tool_map = {} | ||
| agent.human_input_callback = None | ||
|
|
||
| return agent | ||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_server_in_filter_applies_filter(self, mock_agent_all_servers): | ||
| """Test: X in tool_filter → Apply filter for server X.""" | ||
| result = await self._apply_list_tools_logic_all_servers( | ||
| mock_agent_all_servers, | ||
| tool_filter={"server1": {"tool1"}, "server2": {"tool3"}} | ||
| ) | ||
|
|
||
| # server1: only tool1, server2: only tool3, server3: all tools (no filter) | ||
| assert len(result.tools) == 3 | ||
| tool_names = {tool.name for tool in result.tools} | ||
| assert tool_names == {"server1:tool1", "server2:tool3", "server3:tool4"} | ||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_wildcard_applies_to_unfiltered_servers(self, mock_agent_all_servers): | ||
| """Test: X not in tool_filter and '*' in tool_filter → Apply wildcard filter.""" | ||
| result = await self._apply_list_tools_logic_all_servers( | ||
| mock_agent_all_servers, | ||
| tool_filter={ | ||
| "server1": {"tool1"}, # Explicit filter for server1 | ||
| "*": {"tool3", "tool4"} # Wildcard for others | ||
| } | ||
| ) | ||
|
|
||
| # server1: only tool1 (explicit filter) | ||
| # server2: only tool3 (from wildcard) | ||
| # server3: only tool4 (from wildcard) | ||
| assert len(result.tools) == 3 | ||
| tool_names = {tool.name for tool in result.tools} | ||
| assert tool_names == {"server1:tool1", "server2:tool3", "server3:tool4"} | ||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_no_filter_no_wildcard_includes_tool(self, mock_agent_all_servers): | ||
| """Test: X not in tool_filter and '*' not in tool_filter → Include tool (no filter).""" | ||
| result = await self._apply_list_tools_logic_all_servers( | ||
| mock_agent_all_servers, | ||
| tool_filter={"server1": {"tool1"}} # Only server1 has filter | ||
| ) | ||
|
|
||
| # server1: only tool1 (explicit filter) | ||
| # server2: all tools (no filter) | ||
| # server3: all tools (no filter) | ||
| assert len(result.tools) == 4 | ||
| tool_names = {tool.name for tool in result.tools} | ||
| assert tool_names == {"server1:tool1", "server2:tool1", "server2:tool3", "server3:tool4"} | ||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_empty_filter_dict_includes_all(self, mock_agent_all_servers): | ||
| """Test: tool_filter = {} → All tools included (no explicit filters defined).""" | ||
| result = await self._apply_list_tools_logic_all_servers( | ||
| mock_agent_all_servers, | ||
| tool_filter={} | ||
| ) | ||
|
|
||
| # Empty dict means no explicit filters are defined | ||
| # Since no server is explicitly listed and there's no wildcard, | ||
| # the logic falls through to include all tools by default | ||
| assert len(result.tools) == 5 # All 5 tools from the fixture should be included | ||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_wildcard_only_filter(self, mock_agent_all_servers): | ||
| """Test: Only wildcard filter applies to all servers.""" | ||
| result = await self._apply_list_tools_logic_all_servers( | ||
| mock_agent_all_servers, | ||
| tool_filter={"*": {"tool1"}} |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Empty dict semantics conflict with RequestParams docstring. Pick one and align tests/docs.
This test asserts {} includes all tools, but RequestParams docstring example says {} means “No tools allowed from any server.” Decide which behavior is canonical and update either the production docstring or this test (and filtering logic) accordingly.
To verify the docstring reference in the repo:
🏁 Script executed:
#!/bin/bash
# Show the RequestParams.tool_filter examples mentioning '{}'
rg -n -C2 'RequestParams|tool_filter|No tools allowed from any server|\{\}\s*-\s*No tools allowed' src/mcp_agent/workflows/llm/augmented_llm.pyLength of output: 4255
Pick canonical semantics for RequestParams.tool_filter ('{}' behavior) and align docs, tests, and filtering logic.
- Docstring: src/mcp_agent/workflows/llm/augmented_llm.py (lines ~176–194) currently documents "- {} - No tools allowed from any server".
- Test: tests/workflows/llm/test_request_params_tool_filter.py — TestAgentToolFilteringAllServers.test_empty_filter_dict_includes_all asserts {} includes all tools (expects 5 tools).
Choose one behavior and update accordingly: either change the docstring to say "{} means include all" or change the test + filtering implementation so "{}" means "no tools" (and adjust expectations).
2 participants
Add tool filtering capability to LLM requests
This PR adds the ability to filter which tools are exposed to LLMs on a per-request basis. This provides fine-grained control over which tools an LLM can access during a specific interaction.
Key features:
tool_filterparameter toRequestParamsto control which tools are exposed to LLMs"*": Wildcard filter for servers without explicit filters"non_namespaced_tools": Filter for function tools and human inputUsage example: