api: fix MCP lifespan, streaming parser, --no-continue flag, MCP session protocol

Author: psyb0t

api_server.py

@@ -14,7 +14,22 @@
 from pydantic import BaseModel, ConfigDict, Field
 
 
-app = FastAPI()
+_mcp_lifespan_cm = None
+
+
+from contextlib import asynccontextmanager  # noqa: E402
+
+
+@asynccontextmanager
+async def _lifespan(app):
+    if _mcp_lifespan_cm:
+        async with _mcp_lifespan_cm:
+            yield
+    else:
+        yield
+
+
+app = FastAPI(lifespan=_lifespan)
 
 
 def _to_camel(name: str) -> str:
@@ -166,8 +181,8 @@ async def _run_claude_text(
     args = ["claude", "--dangerously-skip-permissions"]
     if resume:
         args += ["--resume", resume]
-    elif no_continue:
-        args.append("--no-continue")
+    elif not no_continue:
+        args.append("--continue")
     args += ["-p", prompt, "--output-format", "json"]
     if model:
         args += ["--model", model]
@@ -488,8 +503,8 @@ def _build_oai_run_args(
     no_continue: bool = True,
 ) -> list[str]:
     args = ["claude", "--dangerously-skip-permissions"]
-    if no_continue:
-        args.append("--no-continue")
+    if not no_continue:
+        args.append("--continue")
     if streaming:
         args += ["-p", prompt, "--output-format", "stream-json", "--verbose"]
     else:
@@ -602,18 +617,20 @@ def _chunk(delta: dict, fr=None) -> str:
                         except json.JSONDecodeError:
                             continue
                         etype = event.get("type", "")
-                        if etype == "message_start":
-                            model_name = event.get("message", {}).get("model", model_name)
-                        elif etype == "content_block_delta":
-                            delta = event.get("delta", {})
-                            if delta.get("type") == "text_delta":
-                                text = delta.get("text", "")
-                                if text:
-                                    yield _chunk({"content": text})
-                        elif etype == "message_delta":
-                            stop = event.get("delta", {}).get("stop_reason")
-                            if stop:
-                                finish_reason = stop
+                        if etype == "assistant":
+                            msg = event.get("message", {})
+                            m = msg.get("model")
+                            if m:
+                                model_name = m
+                            for block in msg.get("content", []):
+                                if block.get("type") == "text":
+                                    text = block.get("text", "")
+                                    if text:
+                                        yield _chunk({"content": text})
+                        elif etype == "result":
+                            sr = event.get("stop_reason", "")
+                            if sr:
+                                finish_reason = sr
 
                 await proc.wait()
                 yield _chunk({}, finish_reason)
@@ -628,31 +645,7 @@ def _chunk(delta: dict, fr=None) -> str:
 
 from mcp.server.fastmcp import FastMCP  # noqa: E402
 
-
-class _MCPBearerAuth:
-    """Wrap an ASGI app with simple Bearer token auth."""
-
-    def __init__(self, asgi_app):
-        self._app = asgi_app
-
-    async def __call__(self, scope, receive, send):
-        if not API_TOKEN or scope["type"] not in ("http", "websocket"):
-            await self._app(scope, receive, send)
-            return
-        headers = {k: v for k, v in scope.get("headers", [])}
-        auth = headers.get(b"authorization", b"").decode()
-        if auth != f"Bearer {API_TOKEN}":
-            await send({
-                "type": "http.response.start",
-                "status": 401,
-                "headers": [[b"content-type", b"application/json"]],
-            })
-            await send({"type": "http.response.body", "body": b'{"detail":"unauthorized"}'})
-            return
-        await self._app(scope, receive, send)
-
-
-_mcp = FastMCP("claude-code")
+_mcp = FastMCP("claude-code", streamable_http_path="/")
 
 
 @_mcp.tool()
@@ -774,7 +767,46 @@ async def delete_file(path: str) -> str:
     return json.dumps({"status": "ok", "path": path})
 
 
-app.mount("/mcp", _MCPBearerAuth(_mcp.http_app()))
+_mcp_app = _mcp.streamable_http_app()
+_mcp_lifespan_cm = _mcp_app.router.lifespan_context(_mcp_app)
+
+
+def _mcp_auth_check(scope) -> bool:
+    if not API_TOKEN:
+        return True
+    headers = {k: v for k, v in scope.get("headers", [])}
+    auth = headers.get(b"authorization", b"").decode()
+    if auth == f"Bearer {API_TOKEN}":
+        return True
+    qs = scope.get("query_string", b"").decode()
+    for part in qs.split("&"):
+        if part.startswith("apiToken=") and part[9:] == API_TOKEN:
+            return True
+    return False
+
+
+class _MCPWithAuth:
+    """Route /mcp requests to the MCP ASGI handler with auth."""
+
+    def __init__(self, mcp_app):
+        self._app = mcp_app
+
+    async def __call__(self, scope, receive, send):
+        if scope["type"] not in ("http", "websocket"):
+            await self._app(scope, receive, send)
+            return
+        if not _mcp_auth_check(scope):
+            await send({
+                "type": "http.response.start",
+                "status": 401,
+                "headers": [[b"content-type", b"application/json"]],
+            })
+            await send({"type": "http.response.body", "body": b'{"detail":"unauthorized"}'})
+            return
+        await self._app(scope, receive, send)
+
+
+app.mount("/mcp", _MCPWithAuth(_mcp_app))
 
 
 if __name__ == "__main__":

tests/test_api.sh

@@ -391,14 +391,33 @@ test_api_openai_reasoning_effort() {
 
 # ── MCP server ─────────────────────────────────────────────────────────────────
 
+_MCP_ACCEPT="Accept: application/json, text/event-stream"
+_MCP_INIT='{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}'
+
+# init MCP session: sets MCP_SESSION var
+_mcp_init() {
+    local url="$1"
+    MCP_SESSION=$(curl -s -D - -X POST "$url" \
+        -H "Content-Type: application/json" -H "$_MCP_ACCEPT" \
+        -d "$_MCP_INIT" | grep -i "mcp-session-id" | awk '{print $2}' | tr -d '\r\n')
+}
+
+# send MCP JSON-RPC with session
+_mcp_call() {
+    local url="$1" data="$2"
+    curl -s -X POST "$url" \
+        -H "Content-Type: application/json" -H "$_MCP_ACCEPT" \
+        -H "mcp-session-id: $MCP_SESSION" \
+        -d "$data"
+}
+
 test_api_mcp_init() {
     _api_start "${API_CONTAINER}-mcp" || return 1
 
-    # MCP streamable HTTP: POST JSON-RPC initialize to /mcp
-    local out code
-    code=$(curl -s -o /dev/null -w "%{http_code}" -X POST "$API_BASE/mcp" \
-        -H "Content-Type: application/json" \
-        -d '{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}')
+    local code
+    code=$(curl -s -o /dev/null -w "%{http_code}" -X POST "$API_BASE/mcp/" \
+        -H "Content-Type: application/json" -H "$_MCP_ACCEPT" \
+        -d "$_MCP_INIT")
     assert_eq "$code" "200" "mcp initialize returns 200" || { _api_stop "${API_CONTAINER}-mcp"; return 1; }
 
     echo "OK: api_mcp_init"
@@ -408,10 +427,11 @@ test_api_mcp_init() {
 test_api_mcp_tools_list() {
     _api_start "${API_CONTAINER}-mcp-tl" || return 1
 
+    _mcp_init "$API_BASE/mcp/"
+    assert_not_empty "$MCP_SESSION" "mcp session id" || { _api_stop "${API_CONTAINER}-mcp-tl"; return 1; }
+
     local out
-    out=$(curl -sf -X POST "$API_BASE/mcp" \
-        -H "Content-Type: application/json" \
-        -d '{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}')
+    out=$(_mcp_call "$API_BASE/mcp/" '{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}')
     assert_contains "$out" "claude_run" "mcp tools/list has claude_run" || { _api_stop "${API_CONTAINER}-mcp-tl"; return 1; }
     assert_contains "$out" "list_files" "mcp tools/list has list_files" || { _api_stop "${API_CONTAINER}-mcp-tl"; return 1; }
     assert_contains "$out" "read_file" "mcp tools/list has read_file" || { _api_stop "${API_CONTAINER}-mcp-tl"; return 1; }
@@ -425,10 +445,11 @@ test_api_mcp_tools_list() {
 test_api_mcp_claude_run() {
     _api_start "${API_CONTAINER}-mcp-run" || return 1
 
+    _mcp_init "$API_BASE/mcp/"
+
     local out
-    out=$(curl -sf -X POST "$API_BASE/mcp" \
-        -H "Content-Type: application/json" \
-        -d '{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"claude_run","arguments":{"prompt":"respond with exactly MCPTEST","model":"'"$TEST_MODEL"'","no_continue":true}}}')
+    out=$(_mcp_call "$API_BASE/mcp/" \
+        '{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"claude_run","arguments":{"prompt":"respond with exactly MCPTEST","model":"'"$TEST_MODEL"'","no_continue":true}}}')
     assert_contains "$out" "MCPTEST" "mcp claude_run returns response" || { _api_stop "${API_CONTAINER}-mcp-run"; return 1; }
 
     echo "OK: api_mcp_claude_run"
@@ -438,30 +459,28 @@ test_api_mcp_claude_run() {
 test_api_mcp_file_ops() {
     _api_start "${API_CONTAINER}-mcp-f" || return 1
 
+    _mcp_init "$API_BASE/mcp/"
+
     local out
 
     # write
-    out=$(curl -sf -X POST "$API_BASE/mcp" \
-        -H "Content-Type: application/json" \
-        -d '{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"write_file","arguments":{"path":"mcptest.txt","content":"hello mcp"}}}')
+    out=$(_mcp_call "$API_BASE/mcp/" \
+        '{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"write_file","arguments":{"path":"mcptest.txt","content":"hello mcp"}}}')
     assert_contains "$out" "ok" "mcp write_file ok" || { _api_stop "${API_CONTAINER}-mcp-f"; return 1; }
 
     # read
-    out=$(curl -sf -X POST "$API_BASE/mcp" \
-        -H "Content-Type: application/json" \
-        -d '{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"read_file","arguments":{"path":"mcptest.txt"}}}')
+    out=$(_mcp_call "$API_BASE/mcp/" \
+        '{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"read_file","arguments":{"path":"mcptest.txt"}}}')
     assert_contains "$out" "hello mcp" "mcp read_file returns content" || { _api_stop "${API_CONTAINER}-mcp-f"; return 1; }
 
     # list
-    out=$(curl -sf -X POST "$API_BASE/mcp" \
-        -H "Content-Type: application/json" \
-        -d '{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"list_files","arguments":{}}}')
+    out=$(_mcp_call "$API_BASE/mcp/" \
+        '{"jsonrpc":"2.0","id":4,"method":"tools/call","params":{"name":"list_files","arguments":{}}}')
     assert_contains "$out" "mcptest.txt" "mcp list_files shows written file" || { _api_stop "${API_CONTAINER}-mcp-f"; return 1; }
 
     # delete
-    out=$(curl -sf -X POST "$API_BASE/mcp" \
-        -H "Content-Type: application/json" \
-        -d '{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"delete_file","arguments":{"path":"mcptest.txt"}}}')
+    out=$(_mcp_call "$API_BASE/mcp/" \
+        '{"jsonrpc":"2.0","id":5,"method":"tools/call","params":{"name":"delete_file","arguments":{"path":"mcptest.txt"}}}')
     assert_contains "$out" "ok" "mcp delete_file ok" || { _api_stop "${API_CONTAINER}-mcp-f"; return 1; }
 
     echo "OK: api_mcp_file_ops"
@@ -473,17 +492,23 @@ test_api_mcp_auth() {
 
     # no token → 401
     local code
-    code=$(curl -s -o /dev/null -w "%{http_code}" -X POST "$API_BASE/mcp" \
-        -H "Content-Type: application/json" \
-        -d '{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}')
+    code=$(curl -s -o /dev/null -w "%{http_code}" -X POST "$API_BASE/mcp/" \
+        -H "Content-Type: application/json" -H "$_MCP_ACCEPT" \
+        -d "$_MCP_INIT")
     assert_eq "$code" "401" "mcp no token returns 401" || { _api_stop "${API_CONTAINER}-mcp-auth"; return 1; }
 
-    # correct token → 200
-    code=$(curl -s -o /dev/null -w "%{http_code}" -X POST "$API_BASE/mcp" \
-        -H "Content-Type: application/json" \
+    # correct token via header → 200
+    code=$(curl -s -o /dev/null -w "%{http_code}" -X POST "$API_BASE/mcp/" \
+        -H "Content-Type: application/json" -H "$_MCP_ACCEPT" \
         -H "Authorization: Bearer mcpsecret" \
-        -d '{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}')
-    assert_eq "$code" "200" "mcp correct token returns 200" || { _api_stop "${API_CONTAINER}-mcp-auth"; return 1; }
+        -d "$_MCP_INIT")
+    assert_eq "$code" "200" "mcp correct token via header returns 200" || { _api_stop "${API_CONTAINER}-mcp-auth"; return 1; }
+
+    # correct token via query param → 200
+    code=$(curl -s -o /dev/null -w "%{http_code}" -X POST "$API_BASE/mcp/?apiToken=mcpsecret" \
+        -H "Content-Type: application/json" -H "$_MCP_ACCEPT" \
+        -d "$_MCP_INIT")
+    assert_eq "$code" "200" "mcp correct token via query param returns 200" || { _api_stop "${API_CONTAINER}-mcp-auth"; return 1; }
 
     echo "OK: api_mcp_auth"
     _api_stop "${API_CONTAINER}-mcp-auth"