Forward rich image display outputs as MCP ImageContent blocks

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 "<internal-error>\n" + traceback.format_exc() + "</internal-error>"
 
 
+_MEDIA_RE = re.compile(r'\n?<media mime="([^"]+)">\n(.*?)\n</media>', re.S)
+_IMG_MIMES = {'image/png','image/jpeg','image/gif','image/webp'}
+
+def _split_media(body):
+    "Split `<media>` 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 + "<internal-error>\nkernel process died while executing this request; a fresh kernel will be started on the next call, with all session state lost\n</internal-error>"
-                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))

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<media mime="{m}">\n{d}\n</media>' 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))
 
 

clikernel/skill.py

@@ -73,7 +73,7 @@ def f(x):
     42
     </execute_result>
 
-`display_data`/`execute_result` prefer a non-image, markdown-over-HTML representation; images are ignored. Exceptions come back as a single clean `<error>` 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 `<error>` traceback -- no color codes, not duplicated.
 
 # Interaction rules
 

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 <media> 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 "<media" 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'<media mime="image/png">\n{PNG1}\n</media>' in out
+    finally: stop_kernel(proc)

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 "<media" 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"