Add streamable-http transport with required bearer-token auth

clikernel/base.py

@@ -222,7 +222,29 @@ def _shutdown(signum, frame):
     atexit.register(_kill_worker, w)
 
 
-async def _serve(w, name, docs, instructions=None, eager=False):
+def _auth_app(app, token):
+    "Wrap ASGI `app` to reject any request whose bearer token doesn't match `token`"
+    async def _f(scope, receive, send):
+        if scope['type'] == 'http':
+            hdr = dict(scope.get('headers') or ()).get(b'authorization', b'').decode()
+            if not (hdr.startswith('Bearer ') and secrets.compare_digest(hdr[7:], token)):
+                await send(dict(type='http.response.start', status=401, headers=[(b'www-authenticate', b'Bearer')]))
+                return await send(dict(type='http.response.body', body=b'unauthorized'))
+        await app(scope, receive, send)
+    return _f
+
+
+def _http_app(mcp):
+    "The streamable-http ASGI app, with the SDK's localhost-only Host check off: the bearer token is what guards this server, and it serves through proxies and public domains"
+    from mcp.server.transport_security import TransportSecuritySettings
+    sec = TransportSecuritySettings(enable_dns_rebinding_protection=False)
+    try: return mcp.streamable_http_app(transport_security=sec)  # mcp>=2
+    except TypeError:                                            # mcp 1.x reads it off the server instead
+        mcp.settings.transport_security = sec
+        return mcp.streamable_http_app()
+
+
+async def _serve(w, name, docs, instructions=None, eager=False, transport='stdio', host='127.0.0.1', port=8000, token=None):
     try: from mcp.server.mcpserver import MCPServer as FastMCP  # mcp>=2 renamed FastMCP
     except ImportError: from mcp.server.fastmcp import FastMCP  # mcp 1.x
     # When `eager`, start the worker before building the server so its banner (including startup output) is
@@ -273,17 +295,27 @@ async def interrupt()->str:
     for f in (execute, restart, interrupt):
         f.__doc__ = docs[f.__name__]
         mcp.tool(structured_output=False)(f)
-    await mcp.run_stdio_async()
+    if transport == 'stdio': return await mcp.run_stdio_async()
+    import uvicorn
+    cfg = uvicorn.Config(_auth_app(_http_app(mcp), token), host=host, port=port, log_level='warning')
+    await uvicorn.Server(cfg).serve()
 
 
 def run_mcp(
     argv=None,         # Worker command line (default: the Python clikernel worker)
     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`?
+    transport='stdio', # 'stdio', or 'streamable-http' to serve over HTTP
+    host='127.0.0.1',  # Interface to bind, for `streamable-http`
+    port=8000,         # Port to bind, for `streamable-http`
+    token=None,        # Bearer token `streamable-http` clients must send (default: `$CLIKERNEL_TOKEN`)
 ):
     "Run an MCP server supervising a persistent stream-protocol worker subprocess."
+    if transport != 'stdio' and not token:
+        token = os.getenv('CLIKERNEL_TOKEN')
+        if not token: raise ValueError(f'{transport} needs a bearer token: pass `token` or set $CLIKERNEL_TOKEN')
     w = _Worker(argv)
     _install_signal_guards(w)
-    asyncio.run(_serve(w, name, {**TOOL_DOCS, **(docs or {})}, instructions, eager))
+    asyncio.run(_serve(w, name, {**TOOL_DOCS, **(docs or {})}, instructions, eager, transport, host, port, token))

clikernel/mcp.py

@@ -1,10 +1,18 @@
 "MCP server supervising a persistent `clikernel` CLI worker subprocess."
 from pathlib import Path
+from fastcore.script import call_parse
 from clikernel import INSTRUCTIONS
 from clikernel.base import run_mcp
 
 
-def main(): run_mcp(instructions=INSTRUCTIONS, eager=Path('pyproject.toml').exists())
+@call_parse
+def main(
+    transport:str='stdio', # 'stdio', or 'streamable-http' to serve over HTTP (needs `$CLIKERNEL_TOKEN`)
+    host:str='127.0.0.1',  # Interface to bind, for `streamable-http`
+    port:int=8000,         # Port to bind, for `streamable-http`
+):
+    "Run the clikernel MCP server"
+    run_mcp(instructions=INSTRUCTIONS, eager=Path('pyproject.toml').exists(), transport=transport, host=host, port=port)
 
 
 if __name__ == "__main__": main()

tests/test_mcp.py

@@ -1,4 +1,4 @@
-import asyncio, json, os, shutil, signal, sys, tempfile
+import asyncio, json, os, shutil, signal, socket, subprocess, sys, tempfile, urllib.error, urllib.request
 import pytest
 from mcp import ClientSession, StdioServerParameters
 from mcp.client.stdio import stdio_client
@@ -161,3 +161,50 @@ async def test_graceful_kill(tmp_path):
     _kill_worker(w)                                  # sync path (signal handler, atexit)
     await w.proc.wait()
     assert m2.read_text() == "cleaned"
+
+
+def _free_port():
+    with socket.socket() as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]
+
+
+def _post(url, token=None, host=None):
+    "POST to `url`, returning the HTTP status (auth is checked before MCP parses anything)"
+    hdrs = {"Content-Type": "application/json"}
+    if token: hdrs["Authorization"] = f"Bearer {token}"
+    if host: hdrs["Host"] = host
+    req = urllib.request.Request(url, data=b"{}", headers=hdrs)
+    try: return urllib.request.urlopen(req, timeout=10).status
+    except urllib.error.HTTPError as e: return e.code
+
+
+async def test_http_auth(tmp_path):
+    "The streamable-http transport serves the same tools, behind a bearer token that unauthorized requests can't skip."
+    from mcp.client.streamable_http import streamable_http_client, create_mcp_http_client
+    port, tok = _free_port(), "test-token"
+    cmd = shutil.which("clikernel-mcp")
+    argv = ([cmd] if cmd else [sys.executable, "-m", "clikernel.mcp"]) + ["--transport", "streamable-http", "--port", str(port)]
+    env = dict(os.environ, CLIKERNEL_TOKEN=tok, CLIKERNEL_STATE_DIR=str(tmp_path))
+    proc = subprocess.Popen(argv, env=env, cwd=tmp_path)
+    url = f"http://127.0.0.1:{port}/mcp"
+    try:
+        for _ in range(100):
+            try:
+                with socket.create_connection(("127.0.0.1", port), timeout=0.5): break
+            except OSError: await asyncio.sleep(0.1)
+        assert await asyncio.to_thread(_post, url) == 401                   # no token
+        assert await asyncio.to_thread(_post, url, "wrong-token") == 401    # wrong token
+        assert await asyncio.to_thread(_post, url, tok, "example.com") \
+            == await asyncio.to_thread(_post, url, tok)                 # a foreign Host is not the server's business: it binds a public interface
+        async with create_mcp_http_client(headers={"Authorization": f"Bearer {tok}"}) as hc, \
+                streamable_http_client(url, http_client=hc) as streams, \
+                ClientSession(*streams[:2]) as s:   # mcp 2.x drops the session-id callback from the tuple
+            await s.initialize()
+            assert {t.name for t in (await s.list_tools()).tools} == {"execute", "restart", "interrupt"}
+            await _text(s, "execute", code="x = 41")
+            assert await _text(s, "execute", code="x + 1") == "42"      # state persists across HTTP calls
+    finally:
+        proc.terminate()
+        proc.wait(timeout=10)
+