Allow a custom MCP server definition with MCPApp #377
Conversation
WalkthroughAdds FastMCP as a first-class, injectable server instance to MCPApp and Context. Updates app_server to support both external and internal FastMCP startup modes with unified resolution of workflows and context. Introduces a globally registered get_token_usage tool in the example server and refactors tests to new context access patterns. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant FastMCP
participant MCPApp
participant ServerContext
participant WorkflowRegistry
rect rgb(240,240,240)
Note over Client,WorkflowRegistry: Startup and workflow resolution (dual mode)
Client->>FastMCP: Connect / start
alt External FastMCP provided
MCPApp->>FastMCP: attach(app)
FastMCP->>ServerContext: ensure/set attached ServerContext
else Internal FastMCP managed
MCPApp->>FastMCP: create and attach(app)
FastMCP->>ServerContext: init lifespan ServerContext
end
Client->>FastMCP: list_workflows / run_workflow
FastMCP->>MCPApp: resolve workflows + context
MCPApp->>WorkflowRegistry: fetch registry
WorkflowRegistry-->>Client: workflows / run status
end
sequenceDiagram
participant Client
participant FastMCP as FastMCP (tool)
participant Context
participant TokenCounter
participant WorkflowNode
Note over Client,FastMCP: Tool: get_token_usage
Client->>FastMCP: call get_token_usage(workflow_id/run_id/name)
FastMCP->>Context: access token_counter
alt Counter available and workflow found
Context->>TokenCounter: get_workflow_node(...)
TokenCounter-->>Context: WorkflowNode
Context->>WorkflowNode: aggregate_usage()
WorkflowNode-->>Context: totals
Context->>TokenCounter: calculate_cost / per-model cost
Context-->>FastMCP: structured usage/cost/model_breakdown/tree
FastMCP-->>Client: result
else Error
Context-->>FastMCP: error structure
FastMCP-->>Client: error
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Poem
✨ 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: 1
🧹 Nitpick comments (3)
src/mcp_agent/server/app_server.py (2)
84-87: Consider adding docstring standards.While the helper function is correctly implemented, consider expanding the docstring to follow the project's documentation standards with Args and Returns sections.
def _get_attached_app(mcp: FastMCP) -> MCPApp | None: - """Return the MCPApp instance attached to the FastMCP server, if any.""" + """Return the MCPApp instance attached to the FastMCP server, if any. + + Args: + mcp: The FastMCP server instance to check for attached app. + + Returns: + The attached MCPApp instance or None if not attached. + """ return getattr(mcp, "_mcp_agent_app", None)
89-92: Consider adding full docstring.Similar to the previous helper, consider expanding the docstring.
def _get_attached_server_context(mcp: FastMCP) -> ServerContext | None: - """Return the ServerContext attached to the FastMCP server, if any.""" + """Return the ServerContext attached to the FastMCP server, if any. + + Args: + mcp: The FastMCP server instance to check for attached context. + + Returns: + The attached ServerContext instance or None if not attached. + """ return getattr(mcp, "_mcp_agent_server_context", None)examples/mcp_agent_server/asyncio/basic_agent_server.py (1)
204-296: Consider adding error handling for edge cases.The implementation is thorough and well-structured. However, consider adding validation for potential None values when accessing nested attributes.
def collect_model_usage(node: TokenNode): """Recursively collect model usage from a node tree""" if node.usage.model_name: model_name = node.usage.model_name - provider = node.usage.model_info.provider if node.usage.model_info else None + provider = getattr(node.usage.model_info, 'provider', None) if node.usage.model_info else NoneAlso consider handling the case where
aggregate_usage()might return None or raise an exception:# Get the aggregated usage for this workflow - workflow_usage = workflow_node.aggregate_usage() + try: + workflow_usage = workflow_node.aggregate_usage() + except Exception as e: + return { + "error": "Failed to aggregate usage", + "message": str(e), + }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
examples/mcp_agent_server/asyncio/basic_agent_server.py(3 hunks)src/mcp_agent/app.py(3 hunks)src/mcp_agent/core/context.py(2 hunks)src/mcp_agent/server/app_server.py(8 hunks)tests/server/test_app_server.py(1 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-07-22T18:59:49.368Z
Learnt from: CR
PR: lastmile-ai/mcp-agent#0
File: examples/usecases/reliable_conversation/CLAUDE.md:0-0
Timestamp: 2025-07-22T18:59:49.368Z
Learning: Applies to examples/usecases/reliable_conversation/examples/reliable_conversation/src/**/*.py : Use mcp-agent's Agent abstraction for ALL LLM interactions, including quality evaluation, to ensure consistent tool access, logging, and error handling.
Applied to files:
examples/mcp_agent_server/asyncio/basic_agent_server.py
🧬 Code Graph Analysis (5)
src/mcp_agent/core/context.py (1)
src/mcp_agent/server/app_server.py (1)
app(69-71)
tests/server/test_app_server.py (3)
src/mcp_agent/server/app_server.py (2)
_workflow_run(479-527)workflows(74-76)src/mcp_agent/app.py (4)
executor(146-147)workflow(326-359)workflows(162-163)context(130-135)src/mcp_agent/executor/workflow.py (2)
executor(114-119)WorkflowExecution(62-69)
src/mcp_agent/server/app_server.py (4)
src/mcp_agent/core/context.py (1)
mcp(93-94)src/mcp_agent/app.py (4)
MCPApp(36-583)workflows(162-163)context(130-135)workflow(326-359)src/mcp_agent/executor/workflow_registry.py (1)
WorkflowRegistry(20-147)src/mcp_agent/executor/workflow.py (1)
create(137-156)
src/mcp_agent/app.py (7)
src/mcp_agent/core/context.py (1)
mcp(93-94)src/mcp_agent/config.py (1)
Settings(533-611)tests/test_app.py (2)
human_input_callback(41-49)signal_notification(52-58)src/mcp_agent/human_input/types.py (1)
HumanInputCallback(45-58)src/mcp_agent/elicitation/types.py (1)
ElicitationCallback(14-26)src/mcp_agent/executor/workflow_signal.py (1)
SignalWaitCallback(339-359)src/mcp_agent/workflows/llm/llm_selector.py (1)
ModelSelector(97-448)
examples/mcp_agent_server/asyncio/basic_agent_server.py (4)
src/mcp_agent/core/context.py (1)
mcp(93-94)src/mcp_agent/server/app_server.py (1)
app(69-71)src/mcp_agent/app.py (2)
MCPApp(36-583)context(130-135)src/mcp_agent/tracing/token_counter.py (9)
get_workflow_node(1420-1455)aggregate_usage(180-210)_calculate_node_cost(1274-1307)input_tokens(382-384)output_tokens(387-389)total_tokens(377-379)to_dict(212-250)TokenNode(95-346)calculate_cost(923-987)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: checks / test
🔇 Additional comments (19)
src/mcp_agent/core/context.py (1)
13-13: LGTM! Clean integration of FastMCP dependency.The import is correctly placed and follows the existing import organization pattern.
src/mcp_agent/app.py (4)
10-10: LGTM! Import follows existing pattern.The FastMCP import is correctly placed alongside other MCP-related imports.
61-67: Excellent type annotation updates!The type annotations have been modernized to use the union operator (
|) syntax consistently throughout, which is more readable and aligns with Python 3.10+ best practices.
77-79: Well-documented new parameter!The
mcpparameter documentation clearly explains both usage scenarios (external vs. internal FastMCP) and the implications of each approach.
87-87: LGTM! Proper storage of the MCP instance.The mcp parameter is correctly stored as an instance attribute for later use.
tests/server/test_app_server.py (2)
3-3: LGTM! Appropriate import for test context construction.The addition of
SimpleNamespaceimport supports the refactored mock context structure.
11-19: Clean and explicit mock context setup!The refactored fixture creates a more maintainable and explicit context structure using
SimpleNamespaceinstead of nestedMagicMockchains. This makes the test setup clearer and aligns well with the dual startup-mode support in the server.src/mcp_agent/server/app_server.py (8)
9-9: LGTM! Required type hints for new helper functions.The addition of
OptionalandTupleimports supports the type annotations in the new resolver functions.
94-116: Excellent dual-mode resolution logic!The function cleanly handles both startup modes (lifespan-based and attached app) with clear precedence and proper fallback behavior. The implementation correctly checks for the required attributes before accessing them.
118-129: LGTM! Consistent resolution pattern.The workflow registry resolver follows the same clean pattern as the workflows/context resolver.
163-190: Well-structured dual startup mode implementation!The implementation elegantly handles both external and internal MCP server scenarios:
- External MCP: Properly attaches the app and creates a ServerContext when needed
- Internal MCP: Creates and configures a new FastMCP with lifespan management
The code correctly stores the server reference and ensures tools are registered in both modes.
201-204: Clean adaptation to dual-mode support.The list_workflows function now correctly uses the resolver to work with both startup modes.
236-240: Proper error handling for server context resolution.The function correctly attempts to resolve the server context from both sources and provides a clear error message when neither is available.
485-499: Excellent refactoring of workflow execution!The _workflow_run function has been cleanly refactored to:
- Use the dual-mode resolver for workflows and context
- Properly validate both are available before proceeding
- Use the resolved app_context for workflow creation
This ensures consistent behavior regardless of startup mode.
533-533: Consistent use of resolver pattern.The _workflow_status function correctly uses the registry resolver for dual-mode support.
examples/mcp_agent_server/asyncio/basic_agent_server.py (4)
16-16: LGTM! Required import for external MCP integration.The FastMCP import is correctly placed for creating an external server instance.
33-42: Good demonstration of external MCP wiring!The example clearly shows how to wire an external FastMCP instance to MCPApp. The comments effectively explain that this is optional and what happens when not provided.
183-203: Comprehensive and well-documented token usage tool!The tool provides excellent functionality with:
- Clear, detailed documentation explaining what information is returned
- Proper parameter descriptions for different lookup methods
- Structured output format
283-294: Good cost calculation and display formatting.The code properly calculates costs per model and formats the display key to include provider information when available. The rounding to 4 decimal places is appropriate for cost display.
| @property | ||
| def mcp(self) -> FastMCP | None: | ||
| return self.app.mcp if self.app else None |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Consider adding type annotation consistency check.
The property implementation is correct. However, ensure that the app attribute's type annotation in the class definition is updated to explicitly indicate that when set, app.mcp will be of type FastMCP | None.
Let me verify that the app attribute type is properly defined and that MCPApp has the mcp attribute:
🏁 Script executed:
#!/bin/bash
# Check MCPApp class definition for mcp attribute
ast-grep --pattern 'class MCPApp {
$$$
def __init__($$$) {
$$$
}
$$$
}'
# Also check for mcp attribute definition
rg -A 3 'class MCPApp' src/mcp_agent/app.py | head -20Length of output: 242
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Print context.py around the referenced lines with numbers
echo "== src/mcp_agent/core/context.py =="
nl -ba src/mcp_agent/core/context.py | sed -n '1,240p' || true
echo
echo "== src/mcp_agent/app.py =="
nl -ba src/mcp_agent/app.py | sed -n '1,240p' || true
echo
echo "== Search for 'mcp' mentions across repo =="
rg -n "mcp" || true
echo
echo "== Search for 'class MCPApp' with context =="
rg -n --hidden -n "class MCPApp" -S || trueLength of output: 264217
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== src/mcp_agent/core/context.py (around mcp property) =="
awk 'NR>=1 && NR<=220 { printf("%4d: %s\n", NR, $0) }' src/mcp_agent/core/context.py || true
echo
echo "== src/mcp_agent/app.py (around class and __init__) =="
awk 'NR>=1 && NR<=260 { printf("%4d: %s\n", NR, $0) }' src/mcp_agent/app.py || true
echo
echo "== grep for 'self.mcp' and 'mcp:' annotations in app.py =="
rg -n --hidden --glob '!**/node_modules/**' "self.mcp|mcp: FastMCP" src/mcp_agent/app.py || trueLength of output: 20622
Add explicit class-level type annotation for MCPApp.mcp
Verified: Context.app is Optional["MCPApp"] and Context.mcp returns self.app.mcp if present. MCPApp.init accepts mcp: FastMCP | None and assigns self.mcp = mcp, but MCPApp does not declare a class-level annotation for the attribute. Add an explicit annotation so static type-checkers reliably know app.mcp is FastMCP | None.
Files to update:
- src/mcp_agent/core/context.py — Context.app: Optional["MCPApp"] = None (OK; keeps using app.mcp)
- src/mcp_agent/app.py — add class-level annotation for mcp
Suggested small patch (add inside the MCPApp class, e.g. after the docstring):
class MCPApp:
"""
Main application class ...
"""
mcp: FastMCP | None = None
def __init__(...
🤖 Prompt for AI Agents
In src/mcp_agent/core/context.py around lines 92-94, Context.app is
Optional["MCPApp"] and Context.mcp returns self.app.mcp if present; ensure
static type-checkers recognize MCPApp.mcp by adding an explicit class-level
annotation in src/mcp_agent/app.py inside the MCPApp class (e.g., add "mcp:
FastMCP | None = None" immediately after the class docstring) so that app.mcp is
typed as FastMCP | None.
2 participants
The overall idea is to let a developer define their own
FastMCPserver definition, in order to define custom tools they may want on their agent server.MCPAppnow optionally takes theFastMCPobject as part of its constructor arguments.create_mcp_server_for_appnow handles both cases:Tool handlers now prefer the attached app for workflow discovery and registry access when there’s no
ServerContext:ServerContext.ServerContext.Summary by CodeRabbit