Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 46 additions & 4 deletions src/mcp/client/stdio.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import logging
import os
import subprocess
import sys
from contextlib import asynccontextmanager
from pathlib import Path
Expand All @@ -24,6 +25,22 @@

logger = logging.getLogger(__name__)


def _is_jupyter_notebook() -> bool:
"""Check if running in a Jupyter notebook environment.

Returns True when running inside Jupyter Notebook, JupyterLab, or
any IPython kernel that uses ZMQ for communication.
"""
try:
from ipykernel.zmqshell import ZMQInteractiveShell # type: ignore
from IPython import get_ipython # type: ignore

return isinstance(get_ipython(), ZMQInteractiveShell) # type: ignore[no-untyped-def]
except Exception:
return False


# Environment variables to inherit by default
DEFAULT_INHERITED_ENV_VARS = (
[
Expand Down Expand Up @@ -123,7 +140,6 @@ async def stdio_client(server: StdioServerParameters, errlog: TextIO = sys.stder
command=command,
args=server.args,
env=({**get_default_environment(), **server.env} if server.env is not None else get_default_environment()),
errlog=errlog,
cwd=server.cwd,
)
except OSError:
Expand Down Expand Up @@ -177,9 +193,32 @@ async def stdin_writer():
except anyio.ClosedResourceError: # pragma: no cover
await anyio.lowlevel.checkpoint()

async def stderr_reader():
"""Read stderr from the process and forward to errlog.

In Jupyter notebook environments, sys.stderr cannot reliably be used
as a subprocess file descriptor. By always piping stderr and reading
it asynchronously, we ensure output is visible in all environments:
Jupyter gets ANSI-colored print output; other environments write
directly to the errlog stream (defaulting to sys.stderr).
"""
assert process.stderr, "Opened process is missing stderr"

in_jupyter = _is_jupyter_notebook()
try:
async for line in process.stderr:
text = line.decode(errors="replace").rstrip("\n")
if in_jupyter:
print(f"\033[91m{text}\033[0m")
else:
print(text, file=errlog)
except anyio.ClosedResourceError: # pragma: lax no cover
await anyio.lowlevel.checkpoint()

async with anyio.create_task_group() as tg, process:
tg.start_soon(stdout_reader)
tg.start_soon(stdin_writer)
tg.start_soon(stderr_reader)
try:
yield read_stream, write_stream
finally:
Expand Down Expand Up @@ -230,21 +269,24 @@ async def _create_platform_compatible_process(
command: str,
args: list[str],
env: dict[str, str] | None = None,
errlog: TextIO = sys.stderr,
cwd: Path | str | None = None,
):
"""Creates a subprocess in a platform-compatible way.

Unix: Creates process in a new session/process group for killpg support
Windows: Creates process in a Job Object for reliable child termination

Stderr is always piped so that the caller can read it asynchronously.
This is required because sys.stderr may not be a valid subprocess file
descriptor in environments such as Jupyter notebooks.
"""
if sys.platform == "win32": # pragma: no cover
process = await create_windows_process(command, args, env, errlog, cwd)
process = await create_windows_process(command, args, env, cwd=cwd)
else: # pragma: lax no cover
process = await anyio.open_process(
[command, *args],
env=env,
stderr=errlog,
stderr=subprocess.PIPE,
cwd=cwd,
start_new_session=True,
)
Expand Down
8 changes: 4 additions & 4 deletions src/mcp/os/win32/utilities.py
Original file line number Diff line number Diff line change
Expand Up @@ -173,7 +173,7 @@ async def create_windows_process(
creationflags=subprocess.CREATE_NO_WINDOW # type: ignore
if hasattr(subprocess, "CREATE_NO_WINDOW")
else 0,
stderr=errlog,
stderr=subprocess.PIPE,
cwd=cwd,
)
except NotImplementedError:
Expand All @@ -184,7 +184,7 @@ async def create_windows_process(
process = await anyio.open_process(
[command, *args],
env=env,
stderr=errlog,
stderr=subprocess.PIPE,
cwd=cwd,
)

Expand All @@ -209,7 +209,7 @@ async def _create_windows_fallback_process(
[command, *args],
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=errlog,
stderr=subprocess.PIPE,
env=env,
cwd=cwd,
bufsize=0, # Unbuffered output
Expand All @@ -221,7 +221,7 @@ async def _create_windows_fallback_process(
[command, *args],
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=errlog,
stderr=subprocess.PIPE,
env=env,
cwd=cwd,
bufsize=0,
Expand Down
65 changes: 65 additions & 0 deletions tests/client/test_stdio.py
Original file line number Diff line number Diff line change
Expand Up @@ -558,3 +558,68 @@ def sigterm_handler(signum, frame):
f"stdio_client cleanup took {elapsed:.1f} seconds for stdin-ignoring process. "
f"Expected between 2-4 seconds (2s stdin timeout + termination time)."
)


@pytest.mark.anyio
async def test_is_jupyter_notebook_returns_false_without_ipykernel():
"""_is_jupyter_notebook returns False when ipykernel is not importable."""
from mcp.client.stdio import _is_jupyter_notebook

# ipykernel is not installed in the test environment
assert _is_jupyter_notebook() is False


@pytest.mark.anyio
async def test_is_jupyter_notebook_returns_true_with_mocked_jupyter(monkeypatch):
"""_is_jupyter_notebook returns True when ZMQInteractiveShell is active."""
import types

# Create a fake ZMQInteractiveShell class
fake_zmq_module = types.ModuleType("ipykernel.zmqshell")

class FakeZMQShell:
pass

fake_zmq_module.ZMQInteractiveShell = FakeZMQShell # type: ignore[attr-defined]

# Create a fake ipykernel module
fake_ipykernel = types.ModuleType("ipykernel")
fake_ipykernel.zmqshell = fake_zmq_module # type: ignore[attr-defined]

# Create a fake IPython module with get_ipython returning our shell
fake_ipython = types.ModuleType("IPython")
fake_ipython.get_ipython = lambda: FakeZMQShell() # type: ignore[attr-defined]

monkeypatch.setitem(sys.modules, "ipykernel", fake_ipykernel)
monkeypatch.setitem(sys.modules, "ipykernel.zmqshell", fake_zmq_module)
monkeypatch.setitem(sys.modules, "IPython", fake_ipython)

from mcp.client.stdio import _is_jupyter_notebook

assert _is_jupyter_notebook() is True


@pytest.mark.anyio
async def test_stderr_reader_writes_ansi_when_in_jupyter(monkeypatch):
"""stderr_reader uses ANSI escape codes when running in Jupyter."""
captured_lines: list[str] = []

def fake_print(*args, **kwargs):
captured_lines.append(str(args[0]) if args else "")

monkeypatch.setattr("builtins.print", fake_print)
monkeypatch.setattr("mcp.client.stdio._is_jupyter_notebook", lambda: True)

server = StdioServerParameters(
command=sys.executable,
args=["-c", "import sys; sys.stderr.write('errline\\n'); sys.stderr.flush()"],
)

async with stdio_client(server) as (read_stream, _write_stream):
with anyio.fail_after(3.0):
async for _ in read_stream:
pass

# At least one line should contain ANSI red escape codes
ansi_lines = [l for l in captured_lines if "\033[91m" in l]
assert len(ansi_lines) > 0, f"Expected ANSI-colored stderr, got: {captured_lines}"
2 changes: 1 addition & 1 deletion tests/issues/test_1027_win_unreachable_cleanup.py
Original file line number Diff line number Diff line change
Expand Up @@ -195,7 +195,7 @@ def echo(text: str) -> str:
# This test manually manages the process to verify stdin-based shutdown
# Start the server process
process = await _create_platform_compatible_process(
command=sys.executable, args=[server_script], env=None, errlog=sys.stderr, cwd=None
command=sys.executable, args=[server_script], env=None, cwd=None
)

# Wait for server to start
Expand Down