feat: support dict-format server_names with per-server tool allowlists on Agent

[Oxygen56]

What

Fixes #183 — Adds client-side MCP tool allowlist support to the Agent API.

Problem

Agent.server_names only accepted List[str], requiring all tools from every server to be exposed. There was no way to say "this agent can only call these 3 tools from server A and these 2 tools from server B."

Solution

server_names now accepts a dict format:

# Before — all tools from all servers
Agent(server_names=["news-api", "weather-api"])

# After — fine-grained allowlist  
Agent(server_names={
    "news-api": ["search_articles", "get_headlines"],
    "weather-api": ["get_forecast"],
})

The existing list format remains fully backward compatible. The allowlist is extracted and passed to the MCP server registry, which already supports per-server tool filtering in MCPServerSettings.tools.

Summary by CodeRabbit

• New Features
  • Server configuration accepts either a simple list or a per-server mapping to allowed tools, enabling per-server tool restrictions.
  • Runtime tool selection merges per-agent per-server allowlists with any call-time filters, where call-time filters can further narrow but not expand available tools.
  • Per-agent tool allowlists are isolated to each agent instance to prevent cross-instance leakage.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: b921fcb3-5822-48d2-8de0-166984a6fd46

📥 Commits

Reviewing files that changed from the base of the PR and between e7814eb and aae5117.

📒 Files selected for processing (1)

• src/mcp_agent/agents/agent.py

🚧 Files skipped from review as they are similar to previous changes (1)

• src/mcp_agent/agents/agent.py

---

📝 Walkthrough

Walkthrough

Agent.server_names now accepts either a flat list of server names or a dict mapping servers to allowed tool lists. A Pydantic field_validator normalizes dict input to a flat server list while preserving the original mapping; model_post_init builds a per-instance server->allowed-tools map used by list_tools, which intersects this allowlist with any call-time tool_filter.

Changes

Server allowlist configuration

Layer / File(s)	Summary
Server names field with validator src/mcp_agent/agents/agent.py	Typing imports updated to include Union and validator decorators. Agent.server_names now accepts Union[List[str], Dict[str, List[str]]]. A @field_validator normalizes dict input into a flat server list and stashes the raw mapping for later per-instance allowlist construction.
Per-instance allowlist construction (model_post_init) src/mcp_agent/agents/agent.py	model_post_init consumes the stashed raw dict to populate _server_tool_allowlist: Dict[str, Set[str]] on the Agent instance and clears the temporary storage.
Tool listing merge logic src/mcp_agent/agents/agent.py	list_tools merges the instance _server_tool_allowlist with any call-time tool_filter: for a server present in both, the allowed tool sets are intersected; if only the agent allowlist exists, it is used as the effective filter.
AgentSpec typing update src/mcp_agent/agents/agent_spec.py	Typing imports extended (Dict, Union) and AgentSpec.server_names retyped to Union[List[str], Dict[str, List[str]]] with the same Field(default_factory=list) default.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Suggested reviewers

• saqadri

Poem

🐰 A rabbit hopped through code at night,

Found server lists reshaped just right,
Dicts became keys and lists unfurled,
Per-server tools tucked safe in a world,
Hooray — agents now guard the tool-filled site!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 75.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The PR title clearly and concisely describes the main change: adding support for dict-format server_names with per-server tool allowlists on the Agent class.
Linked Issues check	✅ Passed	The code changes implement the requirements from issue #183: Agent.server_names now accepts Dict[str, List[str]] format alongside List[str], enabling per-server tool allowlisting with integration into the MCP server registry.
Out of Scope Changes check	✅ Passed	All changes are scoped to supporting dict-format server_names and per-server tool allowlists. No unrelated modifications were introduced outside the stated objectives.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with 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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 3

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/mcp_agent/agents/agent.py`:
- Around line 78-90: AgentSpec.server_names and the config-driven constructor
agent_from_spec still assume a List[str], which prevents the new
Agent.server_names union type (Union[List[str], Dict[str, List[str]]]) from
working; update AgentSpec.server_names to use the same Union[List[str],
Dict[str, List[str]]] type (and Pydantic Field/defaults if used) and modify
agent_from_spec to detect and pass through dict-style server_name mappings
unchanged (and validate allowed tool lists) instead of coercing or rejecting
them so the dict allowlist format is preserved when constructing Agent
instances; touch AgentSpec, agent_from_spec, and any
serialization/deserialization helpers to accept/emit both formats consistently.
- Around line 92-105: The validator _normalize_server_names is dropping
per-server allowlist by setting _raw_server_dict on the class
(object.__setattr__(cls,...)) during a field_validator(pre) and not populating
the instance-level _server_tool_allowlist used by
InitAggregatorRequest/MCPAggregator; replace this with a
root_validator(pre=True) that, when server_names is a dict, extracts the
allowlist mapping into the incoming values (e.g. set
values["_server_tool_allowlist"] = {server: set(allowlist)}) and replaces
server_names with list(values["server_names"].keys()), so the instance gets the
per-server allowlist available for later initialization in InitAggregatorRequest
and MCPAggregator; apply the same pattern to the other validator sites
referenced (around lines noted) to ensure allowlists are preserved.
- Around line 95-105: The validator _normalize_server_names must not mutate the
model class or drop the allowlist: remove object.__setattr__ calls and instead
preserve the dict allowlist on the agent instance (e.g., set an instance
attribute like _server_tool_allowlist) and return the list of server names;
change the validator to run in an appropriate mode (after or use a
root_validator) so instance state can be set safely. Update
AgentSpec.server_names to accept the dict-form (Union[List[str], Dict[str,
List[str]]]) so the allowlist can be parsed from spec loading and ensure
InitAggregatorRequest carries both server_names and the allowlist (e.g., a new
field server_tool_allowlist). Finally, modify MCPAggregator.load_server to
consult the passed server_tool_allowlist (or agent._server_tool_allowlist) when
filtering tools instead of relying only on
context.server_registry.get_server_config(...).allowed_tools so the
agent-provided allowlist is actually applied.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: e1fd1cb0-e4bd-418c-ae59-3c8037d6b009

📥 Commits

Reviewing files that changed from the base of the PR and between f62d849 and 65c706b.

📒 Files selected for processing (1)

• src/mcp_agent/agents/agent.py

[coderabbitai]

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)

95-106: ⚠️ Potential issue | 🔴 Critical | ⚡ Quick win

Race condition: class-level storage corrupts concurrent Agent instantiations.

The validator stores the raw dict on the class (cls._server_names_raw_input), which is shared across all instances. When multiple Agent instances are created concurrently (common in async workflows or multi-threaded Temporal workers), they race:

1. Thread A validator: cls._server_names_raw_input = {"server1": ["tool1"]}
2. Thread B validator: cls._server_names_raw_input = {"server2": ["tool2"]}
3. Thread A model_post_init: reads {"server2": ["tool2"]} (wrong data!)

The comment "Pydantic validation runs synchronously within a single model construction" is true per-instance but does not prevent concurrent constructions from racing.

🔧 Suggested fix: use a root-level model_validator to inject instance data

+    _raw_server_names_input: Optional[Dict[str, List[str]]] = PrivateAttr(default=None)
+
+    `@model_validator`(mode="before")
+    `@classmethod`
+    def _extract_server_allowlist(cls, data: Any) -> Any:
+        if isinstance(data, dict) and "server_names" in data:
+            server_names = data.get("server_names")
+            if isinstance(server_names, dict):
+                # Preserve raw dict for model_post_init via a dedicated field
+                data = dict(data)  # shallow copy to avoid mutating input
+                data["_raw_server_names_input"] = server_names
+                data["server_names"] = list(server_names.keys())
+        return data
+
-    `@field_validator`("server_names", mode="before")
-    `@classmethod`
-    def _normalize_server_names(cls, v):
-        """Normalize server_names to a flat list, extracting allowlist if dict."""
-        if isinstance(v, dict):
-            # Temporarily store the raw dict on the class so model_post_init
-            # can move it to the instance.  This is safe because Pydantic
-            # validation runs synchronously within a single model construction.
-            cls._server_names_raw_input = v  # type: ignore[attr-defined]
-            return list(v.keys())
-        cls._server_names_raw_input = None  # type: ignore[attr-defined]
-        return v

Then in model_post_init:

     def model_post_init(self, __context) -> None:
-        # Consume the raw server_names dict (if any) from the validator
-        # and populate the per-instance allowlist.
-        raw = getattr(self.__class__, "_server_names_raw_input", None)
+        raw = self._raw_server_names_input
         if raw is not None:
             self._server_tool_allowlist = {k: set(v) for k, v in raw.items()}
-            # Clear class-level temp storage so it does not leak to the
-            # next Agent instance.
-            self.__class__._server_names_raw_input = None  # type: ignore[attr-defined]
+            self._raw_server_names_input = None  # allow GC

This approach requires importing model_validator and Any, but eliminates shared mutable state.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/mcp_agent/agents/agent.py` around lines 95 - 106, The class-level storage
of raw input in _normalize_server_names via cls._server_names_raw_input causes
races across concurrent Agent constructions; change the approach to avoid
class-shared state by removing cls._server_names_raw_input usage and instead
implement a model-level validator (use pydantic's `@model_validator`(mode="after")
or a root-level validator) that receives the original input and moves the
allowlist into the instance within model_post_init logic; update Agent to import
model_validator and typing.Any as needed, remove any writes to
cls._server_names_raw_input in _normalize_server_names, and ensure
_normalize_server_names returns the flat list while the new model_validator
handles attaching the per-instance raw dict to self in a thread-safe manner
(refer to functions _normalize_server_names and model_post_init for where to
apply the change).

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Outside diff comments:
In `@src/mcp_agent/agents/agent.py`:
- Around line 95-106: The class-level storage of raw input in
_normalize_server_names via cls._server_names_raw_input causes races across
concurrent Agent constructions; change the approach to avoid class-shared state
by removing cls._server_names_raw_input usage and instead implement a
model-level validator (use pydantic's `@model_validator`(mode="after") or a
root-level validator) that receives the original input and moves the allowlist
into the instance within model_post_init logic; update Agent to import
model_validator and typing.Any as needed, remove any writes to
cls._server_names_raw_input in _normalize_server_names, and ensure
_normalize_server_names returns the flat list while the new model_validator
handles attaching the per-instance raw dict to self in a thread-safe manner
(refer to functions _normalize_server_names and model_post_init for where to
apply the change).

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 6d53ea58-a4bc-4f53-8089-ad435f6a733c

📥 Commits

Reviewing files that changed from the base of the PR and between 65c706b and e7814eb.

📒 Files selected for processing (2)

• src/mcp_agent/agents/agent.py
• src/mcp_agent/agents/agent_spec.py