Make mcp-agent Context a derivative of FastMCP context #504
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. WalkthroughContext now subclasses MCPContext and adds FastMCP-aware properties and request binding. It introduces request-scoped fallbacks for fastmcp/mcp, session, logger, async log/report_progress/read_resource methods, and typing updates (Literal, TYPE_CHECKING types, Logger alias). Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant App as App Context
participant Ctx as Context (MCPContext)
participant Req as Request Context
participant FMCP as FastMCP
participant Sess as ServerSession
participant Log as Logger
Note over App,Ctx: app-level Context created
App->>Ctx: instantiate Context(app_state)
Note over Req,Ctx: request binding
Req->>Ctx: bind_request(request_ctx, fastmcp?)
Ctx->>Ctx: set _request_context, _fastmcp
Note over Ctx,FMCP: session resolution
Ctx->>Ctx: session lookup (upstream_session -> request -> app/FMCP)
alt upstream_session present
Ctx-->>Sess: return upstream_session
else request or app FMCP
Ctx->>FMCP: request session()
FMCP-->>Ctx: ServerSession or None
end
Note over Ctx,Log: logging & progress
Ctx->>Log: log(level, message, logger_name?)
Ctx->>Req: report_progress(progress, total?, message?) (no-op if no request)
Note over Ctx,FMCP: resource reading
alt request-bound provider
Req-->>Ctx: read_resource(uri) -> data
else app/FMCP provider
Ctx->>FMCP: read(uri)
FMCP-->>Ctx: resource data
else
Ctx-->>Ctx: raise ValueError
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 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 |
![graphite-app[bot]](https://avatars.githubusercontent.com/in/158384?s=60&v=4){kind=small}
| def logger(self) -> "Logger": | ||
| return self.app.logger if self.app else None |
There was a problem hiding this comment.
The logger property's return type annotation is Logger, but it returns None when self.app is None. This type mismatch could lead to runtime errors when calling methods on what's expected to be a Logger object. Consider either:
- Updating the return type annotation to
Logger | Noneto accurately reflect the possible return values, or - Providing a fallback logger implementation when
self.appisNone
This would help prevent potential NoneType errors during execution.
| def logger(self) -> "Logger": | |
| return self.app.logger if self.app else None | |
| def logger(self) -> "Logger | None": | |
| return self.app.logger if self.app else None |
Spotted by Diamond
Is this helpful? React 👍 or 👎 to let us know.
![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 (5)
src/mcp_agent/core/context.py (5)
10-10: Fix Ruff F401: remove unused BaseModel import
BaseModelis no longer used after switching toMCPContext.-from pydantic import BaseModel, ConfigDict +from pydantic import ConfigDict
126-163: Narrow exception handling in session resolutionCatching bare
Exceptioncan hide real issues. Limit to the expected failures (e.g.,AttributeError,RuntimeError) for clearer diagnostics.- try: - return super().session # type: ignore[misc] - except Exception: + try: + return super().session # type: ignore[misc] + except (AttributeError, RuntimeError): pass @@ - try: + try: mcp = self.mcp if mcp is not None: ctx = mcp.get_context() # FastMCP.get_context returns a Context that raises outside a request; # guard accordingly. - try: + try: return getattr(ctx, "session", None) - except Exception: + except (AttributeError, RuntimeError): return None - except Exception: + except (AttributeError, RuntimeError): passConsider applying the same pattern to other helpers in this file that currently use
except Exception.
165-167: Fix return type for logger propertyThis can return
Nonewhenappis absent; reflect that in the annotation.- def logger(self) -> "Logger": + def logger(self) -> "Logger | None": return self.app.logger if self.app else None
209-244: Honor logger_name in local fallbackWhen falling back to local logging, use
logger_nameif provided to keep parity with FastMCP’s namespacing.- _logger = self.logger + _logger = self.logger + if logger_name: + _logger = get_logger(logger_name, session_id=self.session_id, context=self)
71-71: Verified — upstream_session typing OK; unify client_session_factory annotations (optional)Runtime: OK — start_server tolerates factories with/without context (tries with context then falls back), MCPAgentClientSession accepts an optional context, and SessionProxy subclasses ServerSession so Context.upstream_session: Optional[ServerSession] is correct.
Recommended (optional) refactor — align type hints for client_session_factory to avoid static-type confusion:
- src/mcp_agent/core/context.py — upstream_session remains at line ~71 (Optional[ServerSession]).
- src/mcp_agent/mcp/mcp_server_registry.py — start_server signature (≈ lines 108–116).
- src/mcp_agent/mcp/gen_client.py — gen_client/connect signatures (≈ lines 17–31).
- src/mcp_agent/mcp/mcp_connection_manager.py — launch_server/get_server/get_server_capabilities signatures (≈ lines 64–70, 410–416, 519–526).
No critical fixes required.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
src/mcp_agent/core/context.py(5 hunks)
🧰 Additional context used
🧠 Learnings (3)
📓 Common learnings
Learnt from: saqadri
PR: lastmile-ai/mcp-agent#386
File: src/mcp_agent/mcp/mcp_server_registry.py:110-116
Timestamp: 2025-08-28T15:07:10.015Z
Learning: In MCP server registry methods, when client_session_factory parameters are updated to accept additional context parameters, ensure the type hints match what is actually passed (Context instance vs ServerSession) and that the default factory (MCPAgentClientSession) can handle the number of arguments being passed to avoid TypeError at runtime.
Learnt from: rholinshead
PR: lastmile-ai/mcp-agent#241
File: src/mcp_agent/mcp/mcp_agent_client_session.py:174-187
Timestamp: 2025-05-27T19:10:22.911Z
Learning: In the MCP Agent codebase, all request parameter types inherit from RequestParams or NotificationParams with consistent Meta class structures, so trace context injection using the base Meta classes doesn't cause type-mismatch issues.
📚 Learning: 2025-08-28T15:07:10.015Z
Learnt from: saqadri
PR: lastmile-ai/mcp-agent#386
File: src/mcp_agent/mcp/mcp_server_registry.py:110-116
Timestamp: 2025-08-28T15:07:10.015Z
Learning: In MCP server registry methods, when client_session_factory parameters are updated to accept additional context parameters, ensure the type hints match what is actually passed (Context instance vs ServerSession) and that the default factory (MCPAgentClientSession) can handle the number of arguments being passed to avoid TypeError at runtime.
Applied to files:
src/mcp_agent/core/context.py
📚 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:
src/mcp_agent/core/context.py
🧬 Code graph analysis (1)
src/mcp_agent/core/context.py (3)
src/mcp_agent/server/app_server.py (2)
app(135-137)workflow_registry(145-147)src/mcp_agent/app.py (5)
MCPApp(42-967)logger(193-210)upstream_session(173-174)upstream_session(177-178)session_id(189-190)src/mcp_agent/logging/logger.py (2)
Logger(34-311)get_logger(513-541)
🪛 GitHub Actions: Pull Request Checks
src/mcp_agent/core/context.py
[error] 10-10: Ruff: F401 'pydantic.BaseModel' imported but unused. Remove unused import. (Command: 'ruff check')
🔇 Additional comments (3)
src/mcp_agent/core/context.py (3)
108-125: Nice FastMCP fallback behaviorGraceful fallback to app-managed FastMCP when no request-bound context exists looks good.
245-259: Good no-op behavior for report_progress outside requestsMatches API while avoiding spurious errors. LGTM.
260-284: Resource read fallback is sensibleDelegating to app-level FastMCP when not request-bound is a solid UX improvement, with a clear error otherwise.
src/mcp_agent/core/context.py
Outdated
| def bind_request(self, request_context: Any, fastmcp: FastMCP | None = None) -> "Context": | ||
| """Return a shallow-copied Context bound to a specific FastMCP request. | ||
| - Shares app-wide state (config, registries, token counter, etc.) with the original Context | ||
| - Attaches `_request_context` and `_fastmcp` so FastMCP Context APIs work during the request | ||
| - Does not mutate the original Context (safe for concurrent requests) | ||
| """ | ||
| # Shallow copy to preserve references to registries/loggers while keeping isolation | ||
| bound: Context = self.model_copy(deep=False) | ||
| try: | ||
| setattr(bound, "_request_context", request_context) | ||
| except Exception: | ||
| pass | ||
| try: | ||
| if fastmcp is None: | ||
| fastmcp = getattr(self, "_fastmcp", None) or self.mcp | ||
| setattr(bound, "_fastmcp", fastmcp) | ||
| except Exception: | ||
| pass | ||
| return bound | ||
|
|
There was a problem hiding this comment.
Avoid Pydantic-only dependency in bind_request
self.model_copy assumes a Pydantic base. If MCPContext isn’t a Pydantic model, this will raise. Use model_copy when available, else fall back to copy.copy.
- # Shallow copy to preserve references to registries/loggers while keeping isolation
- bound: Context = self.model_copy(deep=False)
+ # Shallow copy to preserve references to registries/loggers while keeping isolation
+ # Prefer Pydantic's model_copy when available; otherwise use copy.copy
+ bound: Context = (
+ self.model_copy(deep=False) # type: ignore[attr-defined]
+ if hasattr(self, "model_copy")
+ else copy.copy(self)
+ )Add the missing import (outside this hunk):
# at top-level imports
import copy🤖 Prompt for AI Agents
In src/mcp_agent/core/context.py around lines 170 to 190, bind_request currently
calls self.model_copy(deep=False) which assumes the Context is a Pydantic model
and will raise if not; change this to use model_copy only when present (e.g., if
hasattr(self, "model_copy")), otherwise fall back to using copy.copy(self) to
perform a shallow copy; also add the missing top-level import "import copy" to
the module imports so the fallback works.
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/core/context.py (1)
76-78: Fix mutable class-level default: loaded_subagents is shared across all Context instancesThis list will be shared (and can bleed across requests/apps). Use a per-instance default.
Apply this diff:
- loaded_subagents: List["AgentSpec"] = [] + loaded_subagents: List["AgentSpec"] | None = NoneAnd initialize it during context setup (see below).
@@ - context = Context() + context = Context() + if context.loaded_subagents is None: + context.loaded_subagents = []
♻️ Duplicate comments (2)
src/mcp_agent/core/context.py (2)
179-181: Don’t assume Pydantic: model_copy may not exist on MCPContextMCPContext may not expose Pydantic’s model_copy; this will raise at runtime in bind_request.
Apply this diff to use Pydantic when available, else fall back to copy.copy:
@@ - bound: Context = self.model_copy(deep=False) + bound: Context = ( + self.model_copy(deep=False) # type: ignore[attr-defined] + if hasattr(self, "model_copy") + else copy.copy(self) + )Add import:
@@ -import asyncio +import asyncio +import copyAlso applies to: 5-6
164-167: Type mismatch: logger can be NoneReturn type is Logger but None is returned when app is None. Align the annotation (or provide a fallback logger).
- def logger(self) -> "Logger": + def logger(self) -> "Logger | None": return self.app.logger if self.app else None
🧹 Nitpick comments (2)
src/mcp_agent/core/context.py (2)
211-246: Honor logger_name in fallback logging pathWhen not in a FastMCP request, the logger_name parameter is ignored. Respect it to keep parity with FastMCP behavior and improve observability.
- # Fall back to local logger if available + # Fall back to local logger if available try: - _logger = self.logger + _logger = ( + get_logger(logger_name, session_id=self.session_id, context=self) + if logger_name + else self.logger + ) if _logger is not None: if level == "debug": _logger.debug(message) elif level == "warning": _logger.warning(message) elif level == "error": _logger.error(message) else: _logger.info(message)
99-103: Pydantic config likely inert after switching to MCPContextIf MCPContext isn’t a Pydantic model, model_config/ConfigDict won’t have effect. Consider removing the pydantic dependency here or add a brief comment clarifying why it remains.
Also applies to: 10-10
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
src/mcp_agent/core/context.py(5 hunks)
🧰 Additional context used
🧠 Learnings (3)
📓 Common learnings
Learnt from: saqadri
PR: lastmile-ai/mcp-agent#386
File: src/mcp_agent/mcp/mcp_server_registry.py:110-116
Timestamp: 2025-08-28T15:07:10.015Z
Learning: In MCP server registry methods, when client_session_factory parameters are updated to accept additional context parameters, ensure the type hints match what is actually passed (Context instance vs ServerSession) and that the default factory (MCPAgentClientSession) can handle the number of arguments being passed to avoid TypeError at runtime.
Learnt from: rholinshead
PR: lastmile-ai/mcp-agent#241
File: src/mcp_agent/mcp/mcp_agent_client_session.py:174-187
Timestamp: 2025-05-27T19:10:22.911Z
Learning: In the MCP Agent codebase, all request parameter types inherit from RequestParams or NotificationParams with consistent Meta class structures, so trace context injection using the base Meta classes doesn't cause type-mismatch issues.
📚 Learning: 2025-08-28T15:07:10.015Z
Learnt from: saqadri
PR: lastmile-ai/mcp-agent#386
File: src/mcp_agent/mcp/mcp_server_registry.py:110-116
Timestamp: 2025-08-28T15:07:10.015Z
Learning: In MCP server registry methods, when client_session_factory parameters are updated to accept additional context parameters, ensure the type hints match what is actually passed (Context instance vs ServerSession) and that the default factory (MCPAgentClientSession) can handle the number of arguments being passed to avoid TypeError at runtime.
Applied to files:
src/mcp_agent/core/context.py
📚 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:
src/mcp_agent/core/context.py
🧬 Code graph analysis (1)
src/mcp_agent/core/context.py (3)
src/mcp_agent/server/app_server.py (2)
app(135-137)workflow_registry(145-147)src/mcp_agent/app.py (5)
MCPApp(53-1001)logger(204-221)upstream_session(184-185)upstream_session(188-189)session_id(200-201)src/mcp_agent/logging/logger.py (7)
Logger(34-311)get_logger(513-541)debug(261-269)warning(281-289)error(291-299)info(271-279)progress(301-311)
🔇 Additional comments (5)
src/mcp_agent/core/context.py (5)
108-125: Behavior change: fastmcp now returns None outside requests—confirm downstream expectationsThe override relaxes the base contract (likely non-Optional) and may surface None at call sites that previously assumed a value or an exception. Verify callers handle Optional[FastMCP].
126-163: Session resolution fallback looks soundPrioritization (explicit upstream_session → active FastMCP request → app.get_context()) is reasonable and defensive.
262-286: read_resource fallback and error messaging look goodGraceful delegation to FastMCP when available; clear error otherwise.
193-210: client_id/request_id overrides are safe and defensiveGood fallbacks that avoid raising outside of requests.
7-15: Imports, TYPE_CHECKING aliases, base-class switch, and mcp property: LGTMAlso applies to: 39-56, 60-75, 104-107
This makes it much easier to interface with the MCP server session everywhere that
Contextis available.It's especially useful for @app.tool decorated functions:
@app.tool
def my_agent_fn(params: Any, app_ctx: Context) -> str:
app_ctx.elicit(...)
app_ctx.log(...)
return "Mission accomplished"
Summary by CodeRabbit
New Features
Improvements
Refactor