Skip to content

Commit 69ff03b

Browse files
committed
feat: add Jupyter notebook stderr support to stdio_client
- Add _is_jupyter_notebook() to detect Jupyter/IPython environments - Add _print_stderr() to format stderr with HTML in Jupyter (red color) - Add async _stderr_reader() to capture and display stderr - Pipe stderr (subprocess.PIPE) instead of redirecting to file - Update Windows process creation to support piped stderr - Extract stdout/stdin/stderr readers as module-level functions This enables server stderr output to be visible in Jupyter notebooks, addressing issue #156. Fixes #156
1 parent 5983a65 commit 69ff03b

File tree

2 files changed

+194
-57
lines changed

2 files changed

+194
-57
lines changed

src/mcp/client/stdio/__init__.py

Lines changed: 163 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import logging
22
import os
3+
import subprocess
34
import sys
45
from contextlib import asynccontextmanager
56
from pathlib import Path
@@ -48,6 +49,47 @@
4849
PROCESS_TERMINATION_TIMEOUT = 2.0
4950

5051

52+
def _is_jupyter_notebook() -> bool:
53+
"""
54+
Detect if running in a Jupyter notebook or IPython environment.
55+
56+
Returns:
57+
bool: True if running in Jupyter/IPython, False otherwise
58+
"""
59+
try:
60+
from IPython import get_ipython # type: ignore[import-not-found]
61+
62+
ipython = get_ipython() # type: ignore[no-untyped-call]
63+
return ipython is not None and ipython.__class__.__name__ in ("ZMQInteractiveShell", "TerminalInteractiveShell")
64+
except ImportError:
65+
return False
66+
67+
68+
def _print_stderr(line: str, errlog: TextIO) -> None:
69+
"""
70+
Print stderr output, using IPython's display system if in Jupyter notebook.
71+
72+
Args:
73+
line: The line to print
74+
errlog: The fallback TextIO stream (used when not in Jupyter)
75+
"""
76+
if _is_jupyter_notebook():
77+
try:
78+
from IPython.display import HTML, display # type: ignore[import-not-found]
79+
80+
# Use IPython's display system with red color for stderr
81+
# This ensures proper rendering in Jupyter notebooks
82+
display(HTML(f'<pre style="color: red;">{line}</pre>')) # type: ignore[no-untyped-call]
83+
except Exception:
84+
# If IPython display fails, fall back to regular print
85+
# Log the error but continue (non-critical)
86+
logger.debug("Failed to use IPython display for stderr, falling back to print", exc_info=True)
87+
print(line, file=errlog)
88+
else:
89+
# Not in Jupyter, use standard stderr redirection
90+
print(line, file=errlog)
91+
92+
5193
def get_default_environment() -> dict[str, str]:
5294
"""
5395
Returns a default environment object including only environment variables deemed
@@ -102,11 +144,121 @@ class StdioServerParameters(BaseModel):
102144
"""
103145

104146

147+
async def _stdout_reader(
148+
process: Process | FallbackProcess,
149+
read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception],
150+
encoding: str,
151+
encoding_error_handler: str,
152+
):
153+
"""Read stdout from the process and parse JSONRPC messages."""
154+
assert process.stdout, "Opened process is missing stdout"
155+
156+
try:
157+
async with read_stream_writer:
158+
buffer = ""
159+
async for chunk in TextReceiveStream(
160+
process.stdout,
161+
encoding=encoding,
162+
errors=encoding_error_handler,
163+
):
164+
lines = (buffer + chunk).split("\n")
165+
buffer = lines.pop()
166+
167+
for line in lines:
168+
try:
169+
message = types.JSONRPCMessage.model_validate_json(line)
170+
except Exception as exc: # pragma: no cover
171+
logger.exception("Failed to parse JSONRPC message from server")
172+
await read_stream_writer.send(exc)
173+
continue
174+
175+
session_message = SessionMessage(message)
176+
await read_stream_writer.send(session_message)
177+
except anyio.ClosedResourceError: # pragma: no cover
178+
await anyio.lowlevel.checkpoint()
179+
180+
181+
async def _stdin_writer(
182+
process: Process | FallbackProcess,
183+
write_stream_reader: MemoryObjectReceiveStream[SessionMessage],
184+
encoding: str,
185+
encoding_error_handler: str,
186+
):
187+
"""Write session messages to the process stdin."""
188+
assert process.stdin, "Opened process is missing stdin"
189+
190+
try:
191+
async with write_stream_reader:
192+
async for session_message in write_stream_reader:
193+
json = session_message.message.model_dump_json(by_alias=True, exclude_none=True)
194+
await process.stdin.send(
195+
(json + "\n").encode(
196+
encoding=encoding,
197+
errors=encoding_error_handler,
198+
)
199+
)
200+
except anyio.ClosedResourceError: # pragma: no cover
201+
await anyio.lowlevel.checkpoint()
202+
203+
204+
async def _stderr_reader(
205+
process: Process | FallbackProcess,
206+
errlog: TextIO,
207+
encoding: str,
208+
encoding_error_handler: str,
209+
):
210+
"""Read stderr from the process and display it appropriately."""
211+
if not process.stderr:
212+
return
213+
214+
try:
215+
buffer = ""
216+
async for chunk in TextReceiveStream(
217+
process.stderr,
218+
encoding=encoding,
219+
errors=encoding_error_handler,
220+
):
221+
lines = (buffer + chunk).split("\n")
222+
buffer = lines.pop()
223+
224+
for line in lines:
225+
if line.strip(): # Only print non-empty lines
226+
try:
227+
_print_stderr(line, errlog)
228+
except Exception:
229+
# Log errors but continue (non-critical)
230+
logger.debug("Failed to print stderr line", exc_info=True)
231+
232+
# Print any remaining buffer content
233+
if buffer.strip():
234+
try:
235+
_print_stderr(buffer, errlog)
236+
except Exception:
237+
logger.debug("Failed to print final stderr buffer", exc_info=True)
238+
except anyio.ClosedResourceError: # pragma: no cover
239+
await anyio.lowlevel.checkpoint()
240+
except Exception:
241+
# Log errors but continue (non-critical)
242+
logger.debug("Error reading stderr", exc_info=True)
243+
244+
105245
@asynccontextmanager
106246
async def stdio_client(server: StdioServerParameters, errlog: TextIO = sys.stderr):
107247
"""
108248
Client transport for stdio: this will connect to a server by spawning a
109249
process and communicating with it over stdin/stdout.
250+
251+
This function automatically handles stderr output in a way that is compatible
252+
with Jupyter notebook environments. When running in Jupyter, stderr output
253+
is displayed using IPython's display system with red color formatting.
254+
When not in Jupyter, stderr is redirected to the provided errlog stream
255+
(defaults to sys.stderr).
256+
257+
Args:
258+
server: Parameters for the server process to spawn
259+
errlog: TextIO stream for stderr output when not in Jupyter (defaults to sys.stderr).
260+
This parameter is kept for backward compatibility but may be ignored
261+
when running in Jupyter notebook environments.
110262
"""
111263
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
112264
read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
@@ -136,55 +288,14 @@ async def stdio_client(server: StdioServerParameters, errlog: TextIO = sys.stder
136288
await write_stream_reader.aclose()
137289
raise
138290

139-
async def stdout_reader():
140-
assert process.stdout, "Opened process is missing stdout"
141-
142-
try:
143-
async with read_stream_writer:
144-
buffer = ""
145-
async for chunk in TextReceiveStream(
146-
process.stdout,
147-
encoding=server.encoding,
148-
errors=server.encoding_error_handler,
149-
):
150-
lines = (buffer + chunk).split("\n")
151-
buffer = lines.pop()
152-
153-
for line in lines:
154-
try:
155-
message = types.JSONRPCMessage.model_validate_json(line)
156-
except Exception as exc: # pragma: no cover
157-
logger.exception("Failed to parse JSONRPC message from server")
158-
await read_stream_writer.send(exc)
159-
continue
160-
161-
session_message = SessionMessage(message)
162-
await read_stream_writer.send(session_message)
163-
except anyio.ClosedResourceError: # pragma: no cover
164-
await anyio.lowlevel.checkpoint()
165-
166-
async def stdin_writer():
167-
assert process.stdin, "Opened process is missing stdin"
168-
169-
try:
170-
async with write_stream_reader:
171-
async for session_message in write_stream_reader:
172-
json = session_message.message.model_dump_json(by_alias=True, exclude_none=True)
173-
await process.stdin.send(
174-
(json + "\n").encode(
175-
encoding=server.encoding,
176-
errors=server.encoding_error_handler,
177-
)
178-
)
179-
except anyio.ClosedResourceError: # pragma: no cover
180-
await anyio.lowlevel.checkpoint()
181-
182291
async with (
183292
anyio.create_task_group() as tg,
184293
process,
185294
):
186-
tg.start_soon(stdout_reader)
187-
tg.start_soon(stdin_writer)
295+
tg.start_soon(_stdout_reader, process, read_stream_writer, server.encoding, server.encoding_error_handler)
296+
tg.start_soon(_stdin_writer, process, write_stream_reader, server.encoding, server.encoding_error_handler)
297+
if process.stderr:
298+
tg.start_soon(_stderr_reader, process, errlog, server.encoding, server.encoding_error_handler)
188299
try:
189300
yield read_stream, write_stream
190301
finally:
@@ -244,14 +355,19 @@ async def _create_platform_compatible_process(
244355
245356
Unix: Creates process in a new session/process group for killpg support
246357
Windows: Creates process in a Job Object for reliable child termination
358+
359+
Note: stderr is piped (not redirected) to allow async reading for Jupyter
360+
notebook compatibility. The errlog parameter is kept for backward compatibility
361+
but is only used when not in Jupyter environments.
247362
"""
248363
if sys.platform == "win32": # pragma: no cover
249-
process = await create_windows_process(command, args, env, errlog, cwd)
364+
process = await create_windows_process(command, args, env, errlog, cwd, pipe_stderr=True)
250365
else:
366+
# Pipe stderr instead of redirecting to allow async reading
251367
process = await anyio.open_process(
252368
[command, *args],
253369
env=env,
254-
stderr=errlog,
370+
stderr=subprocess.PIPE,
255371
cwd=cwd,
256372
start_new_session=True,
257373
) # pragma: no cover

src/mcp/os/win32/utilities.py

Lines changed: 31 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@ class FallbackProcess:
7070
A fallback process wrapper for Windows to handle async I/O
7171
when using subprocess.Popen, which provides sync-only FileIO objects.
7272
73-
This wraps stdin and stdout into async-compatible
73+
This wraps stdin, stdout, and stderr into async-compatible
7474
streams (FileReadStream, FileWriteStream),
7575
so that MCP clients expecting async streams can work properly.
7676
"""
@@ -79,10 +79,12 @@ def __init__(self, popen_obj: subprocess.Popen[bytes]):
7979
self.popen: subprocess.Popen[bytes] = popen_obj
8080
self.stdin_raw = popen_obj.stdin # type: ignore[assignment]
8181
self.stdout_raw = popen_obj.stdout # type: ignore[assignment]
82-
self.stderr = popen_obj.stderr # type: ignore[assignment]
82+
self.stderr_raw = popen_obj.stderr # type: ignore[assignment]
8383

8484
self.stdin = FileWriteStream(cast(BinaryIO, self.stdin_raw)) if self.stdin_raw else None
8585
self.stdout = FileReadStream(cast(BinaryIO, self.stdout_raw)) if self.stdout_raw else None
86+
# Wrap stderr in async stream if it's piped (for Jupyter compatibility)
87+
self.stderr = FileReadStream(cast(BinaryIO, self.stderr_raw)) if self.stderr_raw else None
8688

8789
async def __aenter__(self):
8890
"""Support async context manager entry."""
@@ -103,12 +105,14 @@ async def __aexit__(
103105
await self.stdin.aclose()
104106
if self.stdout:
105107
await self.stdout.aclose()
108+
if self.stderr:
109+
await self.stderr.aclose()
106110
if self.stdin_raw:
107111
self.stdin_raw.close()
108112
if self.stdout_raw:
109113
self.stdout_raw.close()
110-
if self.stderr:
111-
self.stderr.close()
114+
if self.stderr_raw:
115+
self.stderr_raw.close()
112116

113117
async def wait(self):
114118
"""Async wait for process completion."""
@@ -139,6 +143,7 @@ async def create_windows_process(
139143
env: dict[str, str] | None = None,
140144
errlog: TextIO | None = sys.stderr,
141145
cwd: Path | str | None = None,
146+
pipe_stderr: bool = False,
142147
) -> Process | FallbackProcess:
143148
"""
144149
Creates a subprocess in a Windows-compatible way with Job Object support.
@@ -155,15 +160,20 @@ async def create_windows_process(
155160
command (str): The executable to run
156161
args (list[str]): List of command line arguments
157162
env (dict[str, str] | None): Environment variables
158-
errlog (TextIO | None): Where to send stderr output (defaults to sys.stderr)
163+
errlog (TextIO | None): Where to send stderr output (defaults to sys.stderr).
164+
Only used when pipe_stderr is False.
159165
cwd (Path | str | None): Working directory for the subprocess
166+
pipe_stderr (bool): If True, pipe stderr instead of redirecting to errlog.
167+
This allows async reading of stderr for Jupyter compatibility.
160168
161169
Returns:
162170
Process | FallbackProcess: Async-compatible subprocess with stdin and stdout streams
163171
"""
164172
job = _create_job_object()
165173
process = None
166174

175+
stderr_target = subprocess.PIPE if pipe_stderr else errlog
176+
167177
try:
168178
# First try using anyio with Windows-specific flags to hide console window
169179
process = await anyio.open_process(
@@ -173,18 +183,18 @@ async def create_windows_process(
173183
creationflags=subprocess.CREATE_NO_WINDOW # type: ignore
174184
if hasattr(subprocess, "CREATE_NO_WINDOW")
175185
else 0,
176-
stderr=errlog,
186+
stderr=stderr_target,
177187
cwd=cwd,
178188
)
179189
except NotImplementedError:
180190
# If Windows doesn't support async subprocess creation, use fallback
181-
process = await _create_windows_fallback_process(command, args, env, errlog, cwd)
191+
process = await _create_windows_fallback_process(command, args, env, errlog, cwd, pipe_stderr=pipe_stderr)
182192
except Exception:
183193
# Try again without creation flags
184194
process = await anyio.open_process(
185195
[command, *args],
186196
env=env,
187-
stderr=errlog,
197+
stderr=stderr_target,
188198
cwd=cwd,
189199
)
190200

@@ -198,19 +208,30 @@ async def _create_windows_fallback_process(
198208
env: dict[str, str] | None = None,
199209
errlog: TextIO | None = sys.stderr,
200210
cwd: Path | str | None = None,
211+
pipe_stderr: bool = False,
201212
) -> FallbackProcess:
202213
"""
203214
Create a subprocess using subprocess.Popen as a fallback when anyio fails.
204215
205216
This function wraps the sync subprocess.Popen in an async-compatible interface.
217+
218+
Args:
219+
command: The executable to run
220+
args: List of command line arguments
221+
env: Environment variables
222+
errlog: Where to send stderr output (only used when pipe_stderr is False)
223+
cwd: Working directory for the subprocess
224+
pipe_stderr: If True, pipe stderr instead of redirecting to errlog
206225
"""
226+
stderr_target = subprocess.PIPE if pipe_stderr else errlog
227+
207228
try:
208229
# Try launching with creationflags to avoid opening a new console window
209230
popen_obj = subprocess.Popen(
210231
[command, *args],
211232
stdin=subprocess.PIPE,
212233
stdout=subprocess.PIPE,
213-
stderr=errlog,
234+
stderr=stderr_target,
214235
env=env,
215236
cwd=cwd,
216237
bufsize=0, # Unbuffered output
@@ -222,7 +243,7 @@ async def _create_windows_fallback_process(
222243
[command, *args],
223244
stdin=subprocess.PIPE,
224245
stdout=subprocess.PIPE,
225-
stderr=errlog,
246+
stderr=stderr_target,
226247
env=env,
227248
cwd=cwd,
228249
bufsize=0,

0 commit comments

Comments
 (0)