Skip to content

Commit 5738557

Browse files
authored
Merge pull request #16 from andylim-duo/upstream-03202026
Upstream 03202026
2 parents 436a108 + 581b207 commit 5738557

6 files changed

Lines changed: 357 additions & 99 deletions

File tree

src/mcp/cli/claude.py

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ def get_claude_config_path() -> Path | None: # pragma: no cover
3333
def get_uv_path() -> str:
3434
"""Get the full path to the uv executable."""
3535
uv_path = shutil.which("uv")
36-
if not uv_path: # pragma: no cover
36+
if not uv_path:
3737
logger.error(
3838
"uv executable not found in PATH, falling back to 'uv'. Please ensure uv is installed and in your PATH"
3939
)
@@ -65,7 +65,7 @@ def update_claude_config(
6565
"""
6666
config_dir = get_claude_config_path()
6767
uv_path = get_uv_path()
68-
if not config_dir: # pragma: no cover
68+
if not config_dir:
6969
raise RuntimeError(
7070
"Claude Desktop config directory not found. Please ensure Claude Desktop"
7171
" is installed and has been run at least once to initialize its config."
@@ -90,7 +90,7 @@ def update_claude_config(
9090
config["mcpServers"] = {}
9191

9292
# Always preserve existing env vars and merge with new ones
93-
if server_name in config["mcpServers"] and "env" in config["mcpServers"][server_name]: # pragma: no cover
93+
if server_name in config["mcpServers"] and "env" in config["mcpServers"][server_name]:
9494
existing_env = config["mcpServers"][server_name]["env"]
9595
if env_vars:
9696
# New vars take precedence over existing ones
@@ -103,22 +103,26 @@ def update_claude_config(
103103

104104
# Collect all packages in a set to deduplicate
105105
packages = {MCP_PACKAGE}
106-
if with_packages: # pragma: no cover
106+
if with_packages:
107107
packages.update(pkg for pkg in with_packages if pkg)
108108

109109
# Add all packages with --with
110110
for pkg in sorted(packages):
111111
args.extend(["--with", pkg])
112112

113-
if with_editable: # pragma: no cover
113+
if with_editable:
114114
args.extend(["--with-editable", str(with_editable)])
115115

116116
# Convert file path to absolute before adding to command
117117
# Split off any :object suffix first
118-
if ":" in file_spec:
118+
# First check if we have a Windows path (e.g., C:\...)
119+
has_windows_drive = len(file_spec) > 1 and file_spec[1] == ":"
120+
121+
# Split on the last colon, but only if it's not part of the Windows drive letter
122+
if ":" in (file_spec[2:] if has_windows_drive else file_spec):
119123
file_path, server_object = file_spec.rsplit(":", 1)
120124
file_spec = f"{Path(file_path).resolve()}:{server_object}"
121-
else: # pragma: no cover
125+
else:
122126
file_spec = str(Path(file_spec).resolve())
123127

124128
# Add mcp run command
@@ -127,7 +131,7 @@ def update_claude_config(
127131
server_config: dict[str, Any] = {"command": uv_path, "args": args}
128132

129133
# Add environment variables if specified
130-
if env_vars: # pragma: no cover
134+
if env_vars:
131135
server_config["env"] = env_vars
132136

133137
config["mcpServers"][server_name] = server_config

src/mcp/server/lowlevel/server.py

Lines changed: 37 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -388,16 +388,23 @@ async def run(
388388
await stack.enter_async_context(task_support.run())
389389

390390
async with anyio.create_task_group() as tg:
391-
async for message in session.incoming_messages:
392-
logger.debug("Received message: %s", message)
393-
394-
tg.start_soon(
395-
self._handle_message,
396-
message,
397-
session,
398-
lifespan_context,
399-
raise_exceptions,
400-
)
391+
try:
392+
async for message in session.incoming_messages:
393+
logger.debug("Received message: %s", message)
394+
395+
tg.start_soon(
396+
self._handle_message,
397+
message,
398+
session,
399+
lifespan_context,
400+
raise_exceptions,
401+
)
402+
finally:
403+
# Transport closed: cancel in-flight handlers. Without this the
404+
# TG join waits for them, and when they eventually try to
405+
# respond they hit a closed write stream (the session's
406+
# _receive_loop closed it when the read stream ended).
407+
tg.cancel_scope.cancel()
401408

402409
async def _handle_message(
403410
self,
@@ -475,16 +482,32 @@ async def _handle_request(
475482
except MCPError as err:
476483
response = err.error
477484
except anyio.get_cancelled_exc_class():
478-
logger.info("Request %s cancelled - duplicate response suppressed", message.request_id)
479-
return
485+
if message.cancelled:
486+
# Client sent CancelledNotification; responder.cancel() already
487+
# sent an error response, so skip the duplicate.
488+
logger.info("Request %s cancelled - duplicate response suppressed", message.request_id)
489+
return
490+
# Transport-close cancellation from the TG in run(); re-raise so the
491+
# TG swallows its own cancellation.
492+
raise
480493
except Exception as err:
481494
if raise_exceptions: # pragma: no cover
482495
raise err
483496
response = types.ErrorData(code=0, message=str(err))
497+
else: # pragma: no cover
498+
response = types.ErrorData(code=types.METHOD_NOT_FOUND, message="Method not found")
484499

500+
try:
485501
await message.respond(response)
486-
else: # pragma: no cover
487-
await message.respond(types.ErrorData(code=types.METHOD_NOT_FOUND, message="Method not found"))
502+
except (anyio.BrokenResourceError, anyio.ClosedResourceError):
503+
# Transport closed between handler unblocking and respond. Happens
504+
# when _receive_loop's finally wakes a handler blocked on
505+
# send_request: the handler runs to respond() before run()'s TG
506+
# cancel fires, but after the write stream closed. Closed if our
507+
# end closed (_receive_loop's async-with exit); Broken if the peer
508+
# end closed first (streamable_http terminate()).
509+
logger.debug("Response for %s dropped - transport closed", message.request_id)
510+
return
488511

489512
logger.debug("Response sent")
490513

src/mcp/shared/session.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,7 @@ def __exit__(
105105
) -> None:
106106
"""Exit the context manager, performing cleanup and notifying completion."""
107107
try:
108-
if self._completed: # pragma: no branch
108+
if self._completed:
109109
self._on_complete(self)
110110
finally:
111111
self._entered = False
@@ -418,7 +418,9 @@ async def _receive_loop(self) -> None:
418418
finally:
419419
# after the read stream is closed, we need to send errors
420420
# to any pending requests
421-
for id, stream in self._response_streams.items():
421+
# Snapshot: stream.send() wakes the waiter, whose finally pops
422+
# from _response_streams before the next __next__() call.
423+
for id, stream in list(self._response_streams.items()):
422424
error = ErrorData(code=CONNECTION_CLOSED, message="Connection closed")
423425
try:
424426
await stream.send(JSONRPCError(jsonrpc="2.0", id=id, error=error))

tests/cli/test_claude.py

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
"""Tests for mcp.cli.claude — Claude Desktop config file generation."""
2+
3+
import json
4+
from pathlib import Path
5+
from typing import Any
6+
7+
import pytest
8+
9+
from mcp.cli.claude import get_uv_path, update_claude_config
10+
11+
12+
@pytest.fixture
13+
def config_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
14+
"""Temp Claude config dir with get_claude_config_path and get_uv_path mocked."""
15+
claude_dir = tmp_path / "Claude"
16+
claude_dir.mkdir()
17+
monkeypatch.setattr("mcp.cli.claude.get_claude_config_path", lambda: claude_dir)
18+
monkeypatch.setattr("mcp.cli.claude.get_uv_path", lambda: "/fake/bin/uv")
19+
return claude_dir
20+
21+
22+
def _read_server(config_dir: Path, name: str) -> dict[str, Any]:
23+
config = json.loads((config_dir / "claude_desktop_config.json").read_text())
24+
return config["mcpServers"][name]
25+
26+
27+
def test_generates_uv_run_command(config_dir: Path):
28+
"""Should write a uv run command that invokes mcp run on the resolved file spec."""
29+
assert update_claude_config(file_spec="server.py:app", server_name="my_server")
30+
31+
resolved = Path("server.py").resolve()
32+
assert _read_server(config_dir, "my_server") == {
33+
"command": "/fake/bin/uv",
34+
"args": ["run", "--frozen", "--with", "mcp[cli]", "mcp", "run", f"{resolved}:app"],
35+
}
36+
37+
38+
def test_file_spec_without_object_suffix(config_dir: Path):
39+
"""File specs without :object should still resolve to an absolute path."""
40+
assert update_claude_config(file_spec="server.py", server_name="s")
41+
42+
assert _read_server(config_dir, "s")["args"][-1] == str(Path("server.py").resolve())
43+
44+
45+
def test_with_packages_sorted_and_deduplicated(config_dir: Path):
46+
"""Extra packages should appear as --with flags, sorted and deduplicated with mcp[cli]."""
47+
assert update_claude_config(file_spec="s.py:app", server_name="s", with_packages=["zebra", "aardvark", "zebra"])
48+
49+
args = _read_server(config_dir, "s")["args"]
50+
assert args[:8] == ["run", "--frozen", "--with", "aardvark", "--with", "mcp[cli]", "--with", "zebra"]
51+
52+
53+
def test_with_editable_adds_flag(config_dir: Path, tmp_path: Path):
54+
"""with_editable should add --with-editable after the --with flags."""
55+
editable = tmp_path / "project"
56+
assert update_claude_config(file_spec="s.py:app", server_name="s", with_editable=editable)
57+
58+
args = _read_server(config_dir, "s")["args"]
59+
assert args[4:6] == ["--with-editable", str(editable)]
60+
61+
62+
def test_env_vars_written(config_dir: Path):
63+
"""env_vars should be written under the server's env key."""
64+
assert update_claude_config(file_spec="s.py:app", server_name="s", env_vars={"KEY": "val"})
65+
66+
assert _read_server(config_dir, "s")["env"] == {"KEY": "val"}
67+
68+
69+
def test_existing_env_vars_merged_new_wins(config_dir: Path):
70+
"""Re-installing should merge env vars, with new values overriding existing ones."""
71+
(config_dir / "claude_desktop_config.json").write_text(
72+
json.dumps({"mcpServers": {"s": {"env": {"OLD": "keep", "KEY": "old"}}}})
73+
)
74+
75+
assert update_claude_config(file_spec="s.py:app", server_name="s", env_vars={"KEY": "new"})
76+
77+
assert _read_server(config_dir, "s")["env"] == {"OLD": "keep", "KEY": "new"}
78+
79+
80+
def test_existing_env_vars_preserved_without_new(config_dir: Path):
81+
"""Re-installing without env_vars should keep the existing env block intact."""
82+
(config_dir / "claude_desktop_config.json").write_text(json.dumps({"mcpServers": {"s": {"env": {"KEEP": "me"}}}}))
83+
84+
assert update_claude_config(file_spec="s.py:app", server_name="s")
85+
86+
assert _read_server(config_dir, "s")["env"] == {"KEEP": "me"}
87+
88+
89+
def test_other_servers_preserved(config_dir: Path):
90+
"""Installing a new server should not clobber existing mcpServers entries."""
91+
(config_dir / "claude_desktop_config.json").write_text(json.dumps({"mcpServers": {"other": {"command": "x"}}}))
92+
93+
assert update_claude_config(file_spec="s.py:app", server_name="s")
94+
95+
config = json.loads((config_dir / "claude_desktop_config.json").read_text())
96+
assert set(config["mcpServers"]) == {"other", "s"}
97+
assert config["mcpServers"]["other"] == {"command": "x"}
98+
99+
100+
def test_raises_when_config_dir_missing(monkeypatch: pytest.MonkeyPatch):
101+
"""Should raise RuntimeError when Claude Desktop config dir can't be found."""
102+
monkeypatch.setattr("mcp.cli.claude.get_claude_config_path", lambda: None)
103+
monkeypatch.setattr("mcp.cli.claude.get_uv_path", lambda: "/fake/bin/uv")
104+
105+
with pytest.raises(RuntimeError, match="Claude Desktop config directory not found"):
106+
update_claude_config(file_spec="s.py:app", server_name="s")
107+
108+
109+
@pytest.mark.parametrize("which_result, expected", [("/usr/local/bin/uv", "/usr/local/bin/uv"), (None, "uv")])
110+
def test_get_uv_path(monkeypatch: pytest.MonkeyPatch, which_result: str | None, expected: str):
111+
"""Should return shutil.which's result, or fall back to bare 'uv' when not on PATH."""
112+
113+
def fake_which(cmd: str) -> str | None:
114+
return which_result
115+
116+
monkeypatch.setattr("shutil.which", fake_which)
117+
assert get_uv_path() == expected
118+
119+
120+
@pytest.mark.parametrize(
121+
"file_spec, expected_last_arg",
122+
[
123+
("C:\\Users\\server.py", "C:\\Users\\server.py"),
124+
("C:\\Users\\server.py:app", "C:\\Users\\server.py:app"),
125+
],
126+
)
127+
def test_windows_drive_letter_not_split(
128+
config_dir: Path, monkeypatch: pytest.MonkeyPatch, file_spec: str, expected_last_arg: str
129+
):
130+
"""Drive-letter paths like 'C:\\server.py' must not be split on the drive colon.
131+
132+
Before the fix, a bare 'C:\\path\\server.py' would hit rsplit(":", 1) and yield
133+
("C", "\\path\\server.py"), calling resolve() on Path("C") instead of the full path.
134+
"""
135+
seen: list[str] = []
136+
137+
def fake_resolve(self: Path) -> Path:
138+
seen.append(str(self))
139+
return self
140+
141+
monkeypatch.setattr(Path, "resolve", fake_resolve)
142+
143+
assert update_claude_config(file_spec=file_spec, server_name="s")
144+
145+
assert seen == ["C:\\Users\\server.py"]
146+
assert _read_server(config_dir, "s")["args"][-1] == expected_last_arg

tests/client/test_config.py

Lines changed: 0 additions & 75 deletions
This file was deleted.

0 commit comments

Comments
 (0)