diff --git a/clikernel/base.py b/clikernel/base.py index 6765d48..8bc4290 100644 --- a/clikernel/base.py +++ b/clikernel/base.py @@ -128,8 +128,8 @@ def serve_stream( class _Worker: - def __init__(self, argv=None): - self.argv = argv or [sys.executable, "-m", "clikernel.cli"] + def __init__(self, argv=None, media=False): + self.argv = argv or [sys.executable, "-m", "clikernel.cli"] + (["--media"] if media else []) self.proc,self.delim,self.started,self.busy,self.desynced = None,None,False,False,False self.startup_info = "" self.lock = asyncio.Lock() @@ -189,6 +189,17 @@ async def run(self, code): def _errbox(): return "\n" + traceback.format_exc() + "" +_MEDIA_RE = re.compile(r'\n?\n(.*?)\n', re.S) +_IMG_MIMES = {'image/png','image/jpeg','image/gif','image/webp'} + +def _split_media(body): + "Split `` elements out of a worker response into `(text, content blocks)`; images only, the rest is dropped" + from mcp.types import ImageContent + parts = _MEDIA_RE.findall(body) + blocks = [ImageContent(type='image', data=d, mimeType=m) for m,d in parts if m in _IMG_MIMES] + return (_MEDIA_RE.sub('', body) if parts else body), blocks + + def _kill_worker(w, grace=3): "Reap the worker child so it can never outlive the supervisor (signal-handler safe: no await); SIGTERM first so it can clean up, SIGKILL if it lingers" p = w.proc @@ -230,7 +241,7 @@ async def _serve(w, name, docs, instructions=None, eager=False): died = docs['died'] async def execute(code:str # Code to run in the persistent session - )->str: # Rendered outputs (stdout, display data, last-expression result, errors) + ): # Rendered outputs (stdout, display data, last-expression result, errors), plus any rich media blocks try: async with w.lock: note = "" @@ -245,7 +256,8 @@ async def execute(code:str # Code to run in the persistent session acked, body = await w.run(code) if body is None: return note + "\nkernel process died while executing this request; a fresh kernel will be started on the next call, with all session state lost\n" - return note + body + text, blocks = _split_media(note + body) + return [text, *blocks] if blocks else text except Exception: w.desynced = True return _errbox() @@ -275,9 +287,10 @@ def run_mcp( name='clikernel', # MCP server name docs=None, # Overrides for `TOOL_DOCS` entries (tool descriptions and the state-lost note) instructions=None, # Static MCP `instructions` text, for when the worker isn't started eagerly - eager=False # Start the worker at server startup, forwarding its banner (with startup output) as `instructions`? + eager=False, # Start the worker at server startup, forwarding its banner (with startup output) as `instructions`? + media=True, # Forward rich display outputs (images) from the worker as MCP content blocks? ): "Run an MCP server supervising a persistent stream-protocol worker subprocess." - w = _Worker(argv) + w = _Worker(argv, media=media) _install_signal_guards(w) asyncio.run(_serve(w, name, {**TOOL_DOCS, **(docs or {})}, instructions, eager)) diff --git a/clikernel/cli.py b/clikernel/cli.py index 7364ac0..2d8e3fe 100644 --- a/clikernel/cli.py +++ b/clikernel/cli.py @@ -2,7 +2,7 @@ from pathlib import Path from fastcore.xdg import xdg_config_home from clikernel import INSTRUCTIONS -from clikernel.base import fmt_error,init_worker,run_startup,serve_stream +from clikernel.base import _IMG_MIMES,fmt_error,init_worker,run_startup,serve_stream def _state_root(): if d := os.environ.get("CLIKERNEL_STATE_DIR"): return Path(d).expanduser() @@ -59,11 +59,14 @@ def _inspect(shell, inspectors, code): return "".join(note for f in inspectors if (note := _call_inspector(f, tree, code))) -def _execute(shell, inspectors, code): - from fastcore.nbio import render_text +def _execute(shell, inspectors, code, media=False): + from fastcore.nbio import render_text,render_media try: note = _inspect(shell, inspectors, code) except Exception as e: return fmt_error("blocked", f"{type(e).__name__}: {e}") - return note + render_text(shell.run(code)) + res = shell.run(code) + out = note + render_text(res) + if media: out += ''.join(f'\n\n{d}\n' for m,d in render_media(res) if m in _IMG_MIMES) + return out def _request_exit(shell): shell._clikernel_exit = True @@ -108,7 +111,7 @@ def main(): shell = _make_shell() block = _startup_block(shell) inspectors = _load_inspectors() - serve_stream(lambda code: _execute(shell, inspectors, code), + serve_stream(lambda code: _execute(shell, inspectors, code, media="--media" in sys.argv[1:]), INSTRUCTIONS + ("\n\n" + block if block else ""), should_exit=lambda: _should_exit(shell)) diff --git a/clikernel/skill.py b/clikernel/skill.py index 77a17e6..24003fa 100644 --- a/clikernel/skill.py +++ b/clikernel/skill.py @@ -73,7 +73,7 @@ def f(x): 42 -`display_data`/`execute_result` prefer a non-image, markdown-over-HTML representation; images are ignored. Exceptions come back as a single clean `` traceback -- no color codes, not duplicated. +`display_data`/`execute_result` prefer a non-image, markdown-over-HTML representation; image display outputs are forwarded to MCP clients as media content blocks (text-only for other consumers, unless the worker runs with `--media`). Exceptions come back as a single clean `` traceback -- no color codes, not duplicated. # Interaction rules diff --git a/tests/test_cli.py b/tests/test_cli.py index a5e188a..596de4d 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -67,12 +67,12 @@ def _env(tmp_path): return env -def start_kernel(tmp_path, extra_env=None): +def start_kernel(tmp_path, extra_env=None, args=()): cmd = shutil.which("clikernel") assert cmd, "clikernel console script is not on PATH" env = _env(tmp_path) if extra_env: env.update(extra_env) - proc = subprocess.Popen([cmd], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env) + proc = subprocess.Popen([cmd, *args], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env) proc._stdout_buffer = bytearray() return proc @@ -369,3 +369,22 @@ def test_cli_profile(tmp_path): assert send(proc, "prof_ran\n")[0] == "7\n" assert send(proc, "ck_ext\n")[0] == "1\n" finally: stop_kernel(proc) + + +PNG1 = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/q842iQAAAABJRU5ErkJggg==' +SHOW_IMG = f"import base64; from IPython.display import Image, display; display(Image(base64.b64decode('{PNG1}'))); 'done'" + +def test_media_flag(tmp_path): + "Rich outputs stay text-only by default (existing consumers see no change); --media appends elements after the rendered text." + proc = start_kernel(tmp_path) + try: + read_until_ready(proc) + out, _ = send(proc, SHOW_IMG + "\n") + assert "done" in out and "\n{PNG1}\n' in out + finally: stop_kernel(proc) diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 1d77a34..0cf3f64 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -161,3 +161,16 @@ async def test_graceful_kill(tmp_path): _kill_worker(w) # sync path (signal handler, atexit) await w.proc.wait() assert m2.read_text() == "cleaned" + + +PNG1 = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/q842iQAAAABJRU5ErkJggg==' + +async def test_mcp_media(): + "Rich display outputs come back as ImageContent blocks after the text part; text-only results are unchanged." + async with stdio_client(_server_params()) as (r, w), ClientSession(r, w) as s: + await s.initialize() + res = await s.call_tool("execute", {"code": f"import base64; from IPython.display import Image, display; display(Image(base64.b64decode('{PNG1}'))); 'done'"}) + assert res.content[0].type == "text" and "done" in res.content[0].text and "