fix: notify active MCP sessions after HMR remount#80
Conversation
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey - I've found 1 issue, and left some high level feedback:
- Overriding
ServerSession.__init__globally is fragile (e.g., in concurrent or nested usage) and may affect other code paths usingServerSession; consider using a factory, subclass, or a hook on session creation instead of monkey-patching the constructor. - Sessions added to
active_sessionsare never removed, so closed or failed sessions may remain referenced; consider hooking into session teardown/close to discard them and avoid notifying dead sessions unnecessarily. _notify_list_changedcatches all exceptions and sends them asedirectly tosend_log_message, which may not be serializable or informative; consider converting tostr(e)(or including a traceback) and narrowing the exception handling to avoid masking unexpected errors.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- Overriding `ServerSession.__init__` globally is fragile (e.g., in concurrent or nested usage) and may affect other code paths using `ServerSession`; consider using a factory, subclass, or a hook on session creation instead of monkey-patching the constructor.
- Sessions added to `active_sessions` are never removed, so closed or failed sessions may remain referenced; consider hooking into session teardown/close to discard them and avoid notifying dead sessions unnecessarily.
- `_notify_list_changed` catches all exceptions and sends them as `e` directly to `send_log_message`, which may not be serializable or informative; consider converting to `str(e)` (or including a traceback) and narrowing the exception handling to avoid masking unexpected errors.
## Individual Comments
### Comment 1
<location path="packages/mcp-hmr/mcp_hmr.py" line_range="65-69" />
<code_context>
+ with suppress(Exception):
+ await session.send_log_message("warning", e)
+
+ original_session_init = ServerSession.__init__
+
+ def _capture_session_init(self: ServerSession, *args, **kwargs):
+ original_session_init(self, *args, **kwargs)
+ active_sessions.add(self)
+
if Path(module).is_file(): # module:attr
</code_context>
<issue_to_address>
**suggestion (bug_risk):** Monkey-patching `ServerSession.__init__` globally can have surprising side effects in concurrent or long-lived processes.
Because the class is patched globally, every `ServerSession` in the process will use this wrapped `__init__` while the context is active, including sessions from other servers or tools sharing the same process. It also assumes `__aenter__`/`__aexit__` always run in pairs to restore the original `__init__`, which may not hold in failure scenarios. Prefer an approach that avoids mutating the global class (e.g., a per-server factory or hook in the session creation path) so `active_sessions` tracking is confined to this app’s lifecycle.
Suggested implementation:
```python
def create_tracked_session(*args, **kwargs) -> ServerSession:
"""Create a ServerSession that is tracked in active_sessions.
This avoids monkey-patching ServerSession.__init__ globally and confines
tracking to sessions explicitly created through this module.
"""
session = ServerSession(*args, **kwargs)
active_sessions.add(session)
return session
```
```python
async def __aenter__(self):
# Avoid mutating ServerSession.__init__ globally; rely on session
# creation going through create_tracked_session instead.
call_pre_reload_hooks()
try:
await main()
tg.create_task(self.start_watching())
```
To fully implement the suggestion and keep `active_sessions` working without global monkey-patching, you should also:
1. Identify all places in this module (and related HMR code paths) where `ServerSession` instances for this app are created directly, e.g.:
- `ServerSession(...)`
- `session_cls = ServerSession` / `session_cls(...)`
Replace those with `create_tracked_session(...)` (or with a factory that delegates to `create_tracked_session`), so only this app’s sessions are added to `active_sessions`.
2. Remove any corresponding cleanup/restoration of `ServerSession.__init__` that may exist in `__aexit__` or other teardown logic, since we no longer patch the class globally.
3. If there is a server or session factory object (e.g. something like `Server(..., session_factory=ServerSession)`), prefer injecting `create_tracked_session` there instead of calling it ad‑hoc, so session tracking remains localized and composable.
These changes will confine `active_sessions` tracking to this app’s lifecycle and avoid side effects in other `ServerSession` users sharing the same process.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces session tracking for the MCP HMR server by monkey-patching ServerSession.init and storing instances in a WeakSet. It also adds logic to notify active sessions of changes to tools, resources, and prompts. A significant issue was identified where concurrent execution of multiple HMR servers could lead to race conditions and incorrect restoration of the patched init method; a module-level lock was suggested as a fix.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: fe761a850a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
Warning Rate limit exceeded
Your organization is not enrolled in usage-based pricing. Contact your admin to enable usage-based pricing to continue reviews beyond the rate limit, or try again in 7 minutes and 57 seconds. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughAdds module-level tracking of live mcp.server.session.ServerSession instances via a WeakSet and a new exported patch_session_init() that monkey-patches ServerSession.init (returns an unpatch() to restore the original). During each hot-reload cycle, while holding the HMR lock and inside the mounted app context, the code iterates captured sessions and invokes per-session notifiers to call send_tool_list_changed(), send_resource_list_changed(), and send_prompt_list_changed(), catching notifier errors and emitting a "warning" via send_log_message when possible. The CLI now installs the patch before running HMR and ensures unpatch() in a finally block. README documents the new API. Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. 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: 1
🧹 Nitpick comments (1)
packages/mcp-hmr/mcp_hmr.py (1)
50-50: Redundant task scheduling pattern.
tg.create_task()schedules tasks in theTaskGroup, whilegather()independently awaits coroutines. Combining them is redundant —gatheralready handles concurrent execution.♻️ Proposed fix
- await gather(*(tg.create_task(_notify_list_changed(session)) for session in [*active_sessions])) + await gather(*(_notify_list_changed(session) for session in [*active_sessions]))🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/mcp-hmr/mcp_hmr.py` at line 50, The code is redundantly wrapping TaskGroup.create_task() calls in asyncio.gather(); change to using the TaskGroup API properly by creating tasks directly in the TaskGroup and removing gather. Specifically, inside the TaskGroup where tg is used, replace the await gather(*(tg.create_task(_notify_list_changed(session)) for session in [*active_sessions])) pattern with a simple loop that calls tg.create_task(_notify_list_changed(session)) for each session in active_sessions (referencing _notify_list_changed and active_sessions) and let the TaskGroup scope manage awaiting their completion.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@packages/mcp-hmr/mcp_hmr.py`:
- Around line 56-63: In _notify_list_changed, the except block passes the raw
Exception object to session.send_log_message which expects JSON-serializable
data; convert the exception to a string (e.g., use str(e)) before calling
send_log_message so the data parameter is serializable and the call won't raise;
update the call inside the suppress block that currently does await
session.send_log_message("warning", e) to await
session.send_log_message("warning", str(e)) (refer to function
_notify_list_changed and method send_log_message).
---
Nitpick comments:
In `@packages/mcp-hmr/mcp_hmr.py`:
- Line 50: The code is redundantly wrapping TaskGroup.create_task() calls in
asyncio.gather(); change to using the TaskGroup API properly by creating tasks
directly in the TaskGroup and removing gather. Specifically, inside the
TaskGroup where tg is used, replace the await
gather(*(tg.create_task(_notify_list_changed(session)) for session in
[*active_sessions])) pattern with a simple loop that calls
tg.create_task(_notify_list_changed(session)) for each session in
active_sessions (referencing _notify_list_changed and active_sessions) and let
the TaskGroup scope manage awaiting their completion.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 62b15118-2b7f-4339-abd5-440e8a029cb1
📒 Files selected for processing (1)
packages/mcp-hmr/mcp_hmr.py
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 78488bb7c1
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9693c09382
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
♻️ Duplicate comments (1)
packages/mcp-hmr/mcp_hmr.py (1)
79-81:⚠️ Potential issue | 🟡 MinorUse a serializable payload in
send_log_message.Line 81 passes the raw exception object; fallback logging can fail silently if serialization fails. Convert to
str(e).🔧 Proposed fix
except Exception as e: with suppress(Exception): - await session.send_log_message("warning", e) + await session.send_log_message("warning", str(e))For the MCP Python SDK version used here, confirm whether `ServerSession.send_log_message(..., data=...)` requires JSON-serializable payloads and whether passing a Python `Exception` object is valid.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/mcp-hmr/mcp_hmr.py` around lines 79 - 81, The fallback logger passes the raw Exception object to session.send_log_message which may not be JSON-serializable; change the call in the exception handler to send a serializable payload (e.g., convert the exception to a string and/or include its type) by calling session.send_log_message("warning", str(e)) or a small dict like {"error": str(e), "type": type(e).__name__} so the SDK's send_log_message receives JSON-serializable data; update the except block around session.send_log_message accordingly (refer to session.send_log_message).
🧹 Nitpick comments (1)
packages/mcp-hmr/mcp_hmr.py (1)
16-30: Makepatch_session_init()idempotent to avoid stacked monkey-patches.If
patch_session_init()is called more than once, unpatch order can restore a wrapped__init__instead of the real original.♻️ Suggested hardening
active_sessions = WeakSet() +_session_init_unpatch: callable | None = None def patch_session_init(): + global _session_init_unpatch + if _session_init_unpatch is not None: + return _session_init_unpatch + from mcp.server.session import ServerSession original_session_init = ServerSession.__init__ @@ - def unpatch(): + def unpatch(): + global _session_init_unpatch ServerSession.__init__ = original_session_init + _session_init_unpatch = None - return unpatch + _session_init_unpatch = unpatch + return unpatch🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/mcp-hmr/mcp_hmr.py` around lines 16 - 30, Make patch_session_init idempotent by detecting and preserving the real original __init__ on first patch and avoiding wrapping an already-wrapped initializer; e.g., in patch_session_init(), check for a sentinel like ServerSession._mcp_original_init — if present, return a no-op (or the existing unpatch) immediately; otherwise store the current ServerSession.__init__ into ServerSession._mcp_original_init, replace __init__ with the wrapper (_capture_session_init) that calls ServerSession._mcp_original_init, and implement unpatch to restore ServerSession.__init__ from ServerSession._mcp_original_init and then delete that sentinel attribute so subsequent patches behave correctly.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@packages/mcp-hmr/mcp_hmr.py`:
- Around line 79-81: The fallback logger passes the raw Exception object to
session.send_log_message which may not be JSON-serializable; change the call in
the exception handler to send a serializable payload (e.g., convert the
exception to a string and/or include its type) by calling
session.send_log_message("warning", str(e)) or a small dict like {"error":
str(e), "type": type(e).__name__} so the SDK's send_log_message receives
JSON-serializable data; update the except block around session.send_log_message
accordingly (refer to session.send_log_message).
---
Nitpick comments:
In `@packages/mcp-hmr/mcp_hmr.py`:
- Around line 16-30: Make patch_session_init idempotent by detecting and
preserving the real original __init__ on first patch and avoiding wrapping an
already-wrapped initializer; e.g., in patch_session_init(), check for a sentinel
like ServerSession._mcp_original_init — if present, return a no-op (or the
existing unpatch) immediately; otherwise store the current
ServerSession.__init__ into ServerSession._mcp_original_init, replace __init__
with the wrapper (_capture_session_init) that calls
ServerSession._mcp_original_init, and implement unpatch to restore
ServerSession.__init__ from ServerSession._mcp_original_init and then delete
that sentinel attribute so subsequent patches behave correctly.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 6be22c35-102d-4fe5-9185-2f31cc26d67d
📒 Files selected for processing (1)
packages/mcp-hmr/mcp_hmr.py
Co-Authored-By: Richard Powell <richard@powell.dev>
|
|
||
| try: | ||
| run(run_with_hmr(**kwargs)) | ||
| except KeyboardInterrupt: |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: eb1efd725d
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| for session in active_sessions: | ||
| tg.create_task(_notify_list_changed(session)) |
There was a problem hiding this comment.
using() now notifies every object in the process-wide active_sessions set, so when multiple mcp_server(...) contexts run in one process, a reload in one server will send *_list_changed notifications to sessions that belong to other servers too. In that setup, unrelated clients receive spurious refresh signals and slow/broken sessions from another server can accumulate long-lived notification tasks in this server’s task group; session tracking should be scoped or filtered per server instance before broadcasting.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@packages/mcp-hmr/mcp_hmr.py`:
- Around line 13-30: The global active_sessions WeakSet causes
cross-installation sharing and incorrect unpatching; to fix, make the session
registry and captured original initializer local to each install by creating
active_sessions inside patch_session_init and closing over the specific
original_session_init there (keep ServerSession.__init__ replaced by the local
_capture_session_init), and ensure unpatch restores that captured
original_session_init; reference the active_sessions symbol, the
patch_session_init function, the ServerSession.__init__ assignment, the
_capture_session_init wrapper, and the unpatch closure when making this change.
- Around line 220-227: The CLI currently calls
patch_session_init()/unpatch_session_init() around run_with_hmr(), but direct
callers of run_with_hmr() don't get the session-init patch so active_sessions
stays empty and reloads won't emit tools/list_changed; modify run_with_hmr() to
call patch_session_init() at its start (store the returned unpatch function) and
ensure it is always undone on exit (use try/finally or context management) so
any direct callers get the same patching behavior as cli(); reference
run_with_hmr(), patch_session_init(), unpatch_session_init(), active_sessions
and the tools/list_changed notifier when updating the function.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 6be22c35-102d-4fe5-9185-2f31cc26d67d
📒 Files selected for processing (1)
packages/mcp-hmr/mcp_hmr.py
1 participant
Resolves #79
Notify active MCP sessions about list changes after an HMR remount.