[Fix] Prevent Duplicate Tool Registration Attempts for MCP Server #384
Conversation
WalkthroughIntroduces a per-FastMCP registry to track registered workflow tools, updates ServerContext and create_workflow_tools to check and persist this registry, and adds tests validating idempotent tool registration and persistence across recreated ServerContext instances (e.g., SSE lifecycles). Exposes ServerContext and create_workflow_tools for testing. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor Client
participant ServerContext
participant create_workflow_tools
participant FastMCP as FastMCP (registry)
participant ToolReg as Tool Registration
Client->>ServerContext: initialize (SSE lifecycle may recreate)
Note over ServerContext,FastMCP: ServerContext binds to existing FastMCP
Client->>create_workflow_tools: invoke with (mcp, server_context)
create_workflow_tools->>FastMCP: read _registered_workflow_tools
alt workflow not yet registered
create_workflow_tools->>ToolReg: register run/get_status tools for workflow
create_workflow_tools->>FastMCP: update _registered_workflow_tools
else already registered
Note over create_workflow_tools: Skip creating duplicate tools
end
Client->>ServerContext: register_workflow(name)
ServerContext->>FastMCP: check _registered_workflow_tools
alt not registered
ServerContext->>ToolReg: create per-workflow tools
ServerContext->>FastMCP: mark as registered
else already registered
Note over ServerContext: No-op
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
Poem
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. ✨ 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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (6)
src/mcp_agent/server/app_server.py (3)
40-45: Defensive init for _registered_workflow_tools to ensure correct typeGood call persisting the registry on the FastMCP instance. Make it resilient if something pre-attaches a non-set (e.g., list via config/mocking). Normalize to a set once.
Apply this diff:
- if not hasattr(self.mcp, "_registered_workflow_tools"): - setattr(self.mcp, "_registered_workflow_tools", set()) + tools = getattr(self.mcp, "_registered_workflow_tools", None) + if not isinstance(tools, set): + tools = set(tools) if tools else set() + setattr(self.mcp, "_registered_workflow_tools", tools)
71-76: register_workflow skips tool creation when workflow already existsAs written, tool creation only happens when the workflow is newly added to the mapping. If the workflow is already present (e.g., preloaded at app init) but tools weren’t registered (fresh MCP instance, prior cleanup, or external interference), calling register_workflow won’t (re)create the tools.
Move the tool-registration block outside the workflow-addition if to guarantee "ensure tools exist" semantics.
Apply this diff:
- # Create tools for this workflow if not already registered - registered_workflow_tools = _get_registered_workflow_tools(self.mcp) - if workflow_name not in registered_workflow_tools: - create_workflow_specific_tools(self.mcp, workflow_name, workflow_cls) - registered_workflow_tools.add(workflow_name) + # Create tools for this workflow if not already registered + registered_workflow_tools = _get_registered_workflow_tools(self.mcp) + if workflow_name not in registered_workflow_tools: + create_workflow_specific_tools(self.mcp, workflow_name, workflow_cls) + registered_workflow_tools.add(workflow_name)Follow-up: If you keep the stricter "only when adding" behavior intentionally, consider documenting it in the method docstring to avoid future confusion.
98-101: Make _get_registered_workflow_tools idempotent and type-safeThis helper should also ensure the attribute exists and is of the expected type. Today it returns a new ephemeral set when missing; most call sites are fine, but making it self-healing tightens guarantees.
Apply this diff:
-def _get_registered_workflow_tools(mcp: FastMCP) -> Set[str]: - """Return the set of registered workflow tools for the FastMCP server, if any.""" - return getattr(mcp, "_registered_workflow_tools", set()) +def _get_registered_workflow_tools(mcp: FastMCP) -> Set[str]: + """Return (and ensure) the set of registered workflow tools on the FastMCP server.""" + tools = getattr(mcp, "_registered_workflow_tools", None) + if not isinstance(tools, set): + tools = set(tools) if tools else set() + setattr(mcp, "_registered_workflow_tools", tools) + return toolstests/server/test_app_server.py (3)
295-303: Decorator tracker is neat; consider reusing via a fixtureThe inline decorator tracker is duplicated later in this file. Extracting a small pytest fixture for it will keep tests DRY and easier to tweak if the tool signature changes.
Happy to draft a tiny fixture if you want.
261-358: Great coverage of idempotent registration; add a “pre-existing workflow, missing tools” caseThis verifies the common path. One edge case: if server_context.workflows already contains the workflow (preloaded), but _registered_workflow_tools is empty (e.g., new MCP instance), calling register_workflow should still (re)create tools. That’s currently skipped by the surrounding membership check in register_workflow.
You can augment the test suite with a scenario like:
def test_register_workflow_creates_tools_when_workflow_already_present(): mock_mcp = MagicMock() # Intentionally provide an empty/missing registry if hasattr(mock_mcp, "_registered_workflow_tools"): delattr(mock_mcp, "_registered_workflow_tools") tools_created = [] def track_tool_calls(*args, **kwargs): def decorator(func): tools_created.append(kwargs.get("name", args[0] if args else "unknown")) return func return decorator mock_mcp.tool = track_tool_calls mock_workflow_class = MagicMock() mock_workflow_class.__doc__ = "Test workflow" mock_run = MagicMock(); mock_run.__name__ = "run" mock_workflow_class.run = mock_run mock_app = MagicMock() mock_context = MagicMock(app=mock_app) mock_context.workflow_registry = None mock_context.config = MagicMock() mock_context.config.execution_engine = "asyncio" # Pre-populate workflows mapping mock_app.workflows = {"workflow1": mock_workflow_class} server_context = ServerContext(mcp=mock_mcp, context=mock_context) # Now ensure tools get created even though workflow already exists server_context.register_workflow("workflow1", mock_workflow_class) assert set(tools_created) == { "workflows-workflow1-run", "workflows-workflow1-get_status", }If you accept the earlier refactor to register_workflow, this test should pass.
360-433: Persistence across SSE-like lifecycles validated wellThis mirrors the FastMCP behavior and confirms the registry lives on the MCP instance, not the ephemeral ServerContext. Nicely done.
Two small extensions (optional):
- Add a second SSE-cycle that introduces a new workflow (e.g., workflow2) and assert only its tools are created (no duplication for workflow1).
- Assert isinstance(mock_mcp._registered_workflow_tools, set) is maintained even if someone mutates it (type normalization from earlier refactor will help).
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (2)
src/mcp_agent/server/app_server.py(5 hunks)tests/server/test_app_server.py(2 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
src/mcp_agent/server/app_server.py (2)
src/mcp_agent/core/context.py (1)
mcp(98-99)src/mcp_agent/app.py (1)
workflows(166-167)
tests/server/test_app_server.py (1)
src/mcp_agent/server/app_server.py (6)
ServerContext(32-90)create_workflow_tools(394-410)app(78-80)workflows(83-85)workflow_registry(88-90)register_workflow(67-75)
🔇 Additional comments (2)
src/mcp_agent/server/app_server.py (1)
403-411: Idempotent per-workflow tool creation looks solidThe registry check + single final setattr ensures consistent state on the FastMCP instance and avoids duplicates across SSE lifecycles. Nicely aligns with the PR objective.
tests/server/test_app_server.py (1)
4-8: Public API imports validatedImporting ServerContext and create_workflow_tools directly from app_server enforces the new API surface. Good to see tests exercising these entry points.
2 participants
Description
Currently, workflow tools are attempted to be re-registered during each sse request to an app-associated mcp server:
INFO:mcp_agent:[mcp_agent.server.app_server] Workflow BasicAgentWorkflow started with workflow ID BasicAgentWorkflow-05052bde-b3a5-4e4c-8995-90e4d03d9912 and run ID 0198e1c1-edc0-7a64-9572-1ec2d3a43d3c. Parameters: {'input': 'Print the first 2 paragraphs of https://modelcontextprotocol.io/introduction'} INFO: 127.0.0.1:59463 - "POST /messages/?session_id=9e7c609706a24436884f58ee5ad43c8a HTTP/1.1" 202 Accepted INFO:mcp.server.lowlevel.server:Processing request of type ListToolsRequest INFO: 127.0.0.1:59463 - "POST /messages/?session_id=9e7c609706a24436884f58ee5ad43c8a HTTP/1.1" 202 Accepted INFO:mcp.server.lowlevel.server:Processing request of type CallToolRequest INFO: 127.0.0.1:59484 - "POST /messages/?session_id=9e7c609706a24436884f58ee5ad43c8a HTTP/1.1" 202 Accepted INFO:mcp.server.lowlevel.server:Processing request of type CallToolRequest WARNING:mcp.server.fastmcp.tools.tool_manager:Tool already exists: workflows-BasicAgentWorkflow-run WARNING:mcp.server.fastmcp.tools.tool_manager:Tool already exists: workflows-BasicAgentWorkflow-get_status WARNING:mcp.server.fastmcp.tools.tool_manager:Tool already exists: workflows-PauseResumeWorkflow-run WARNING:mcp.server.fastmcp.tools.tool_manager:Tool already exists: workflows-PauseResumeWorkflow-get_status INFO: 127.0.0.1:59503 - "GET /sse HTTP/1.1" 200 OKThe reasoning is because fastmcp has a
handle_ssecallback for sse requests which callsserver.run()(see here) and enters lifespan context per each request. And entering that context callscreate_workflow_toolsagain before yielding server_context, i.e. per-request/per-session app initialization/tools registration.Our app initialization is already idempotent, so this PR updates the workflow tools registration to be similarly idempotent per mcp instance.
Testing
Run the
mcp-agent/examples/mcp_agent_server/temporalexample and spawn two clients. Without this fix, see the above warnings about tools already existing. With this fix, no warnings and the client results are still correct.Summary by CodeRabbit