diff --git a/clikernel/base.py b/clikernel/base.py index 02b2c42..412aa62 100644 --- a/clikernel/base.py +++ b/clikernel/base.py @@ -132,8 +132,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() @@ -193,6 +193,17 @@ async def run(self, code): def _errbox(): return "\n" + traceback.format_exc() + "" +_IMG_MIMES = ('image/png','image/jpeg') # mimes forwarded as MCP image blocks; others (e.g. svg) stay in the text +_MEDIA_RE = re.compile(r'\n?<(display_data|execute_result) mime="(%s)">\n(.*?)\n' % '|'.join(_IMG_MIMES), re.S) + +def _split_media(body): + "Split image output elements out of a worker response into `(text, content blocks)`" + from mcp.types import ImageContent + parts = _MEDIA_RE.findall(body) + blocks = [ImageContent(type='image', data=d, mimeType=m) for _,m,d in parts] + 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 @@ -234,7 +245,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 = "" @@ -249,7 +260,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] if text.strip() else []), *blocks] if blocks else text except Exception: w.desynced = True return _errbox() @@ -279,9 +291,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 09f42b0..13e33ef 100644 --- a/clikernel/cli.py +++ b/clikernel/cli.py @@ -61,12 +61,12 @@ def _inspect(shell, inspectors, code): return "".join(note for f in inspectors if (note := _call_inspector(f, tree, code))) -def _execute(shell, inspectors, code): +def _execute(shell, inspectors, code, media=False): from fastcore.nbio import render_text try: note = _inspect(shell, inspectors, code) except RuleBlock as e: return fmt_error("blocked", str(e)) except Exception as e: note = f"\ninspector crashed, check skipped ({type(e).__name__}: {e})\n\n" - return note + render_text(shell.run(code)) + return note + render_text(shell.run(code), include_imgs=media) def _request_exit(shell): shell._clikernel_exit = True @@ -123,7 +123,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 36513f6..07b9f85 100644 --- a/clikernel/skill.py +++ b/clikernel/skill.py @@ -72,7 +72,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 @@ -88,6 +88,7 @@ def f(x): - `importlib.reload`ing a module is not always enough to pick up a change: other already-imported modules that did `from x import *`, or that monkeypatched one of its classes (e.g. via fastcore's `@patch`), hold stale references that a targeted reload won't fix. If you hit a stale-class symptom (a class missing a method you know it has, `isinstance` mysteriously failing), you need a fresh interpreter process: the MCP `restart` tool, or in CLI mode, exit and start a new process. - Everything a cell outputs lands in the conversation and stays there. Be surgical: inspect only what's needed, and don't dump large values -- `print(len(v))` first, then decide whether to print in full or filter down. - The kernel is scoped to one client session and shared by any subagents within it. If the server restarts or exits, in-memory state is gone -- redo imports and setup. +- If the host conversation survived a kernel restart and the relevant `doc()` output is still visible, call `doced(name1, name2, ...)` to restore doc state without reprinting it. After context compaction, call `forget_doced()` and read the docs again. # Pyskills diff --git a/tests/test_cli.py b/tests/test_cli.py index cf8dede..96e19b5 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 @@ -377,3 +377,25 @@ 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'" +MULTI_IMG = f"display({{'image/png': base64.b64decode('{PNG1}'), 'image/jpeg': base64.b64decode('{PNG1}')}}, raw=True); 'done'" + +def test_media_flag(tmp_path): + "Rich outputs stay text-only by default (existing consumers see no change); --media renders image outputs inline with a `mime` attribute." + proc = start_kernel(tmp_path) + try: + read_until_ready(proc) + out, _ = send(proc, SHOW_IMG + "\n") + assert "done" in out and 'mime=' not in out + finally: stop_kernel(proc) + proc = start_kernel(tmp_path, args=("--media",)) + try: + read_until_ready(proc) + out, _ = send(proc, SHOW_IMG + "\n") + assert "done" in out and f'\n{PNG1}\n' in out + out, _ = send(proc, MULTI_IMG + "\n") + assert out.count('mime=') == 1 and 'mime="image/jpeg"' in out # one image per output: fastcore's preferred mime wins + finally: stop_kernel(proc) diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 0a863dc..9c5bc9d 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 PNG1 not in res.content[0].text + img = next(c for c in res.content if c.type == "image") + assert img.mimeType == "image/png" and img.data == PNG1 + assert await _text(s, "execute", code="40+2") == "42"