diff --git a/CLAUDE.md b/CLAUDE.md index 5ccac7f..b279a4b 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,6 +1,6 @@ # Claude Code on Databricks -Welcome! This environment comes pre-configured with 5 AI coding agents, 39 skills, and 2 MCP servers. Hermes Agent is available alongside Claude Code, Codex, Gemini CLI, and OpenCode — launch it with `hermes chat`. +Welcome! This environment comes pre-configured with 5 AI coding agents, 43 skills, and 3 MCP servers. Hermes Agent is available alongside Claude Code, Codex, Gemini CLI, and OpenCode — launch it with `hermes chat`. ## Skills (30 total) @@ -39,6 +39,7 @@ From [obra/superpowers](https://github.com/obra/superpowers): - **DeepWiki** - AI-powered documentation for any GitHub repository - **Exa** - Web search and code context retrieval +- **CoDA** (exposed at `/mcp`) - Delegate coding tasks to AI agents via MCP. Any MCP client (Genie Code, Claude Desktop, Cursor) can call `coda_run`, `coda_inbox`, and `coda_get_result` to submit background tasks, check status, and retrieve results. See `docs/mcp-v2-background-execution.md`. ## Databricks CLI diff --git a/README.md b/README.md index 97a378b..4b4726d 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ [![Use this template](https://img.shields.io/badge/Use%20this%20template-2ea44f?logo=github)](https://github.com/datasciencemonkey/coding-agents-databricks-apps/generate) [![Deploy to Databricks](https://img.shields.io/badge/Deploy-Databricks%20Apps-FF3621?logo=databricks&logoColor=white)](docs/deployment.md) [![Agents](https://img.shields.io/badge/Agents-5%20included-green)](#whats-inside) -[![Skills](https://img.shields.io/badge/Skills-39%20built--in-blue)](#-all-39-skills) +[![Skills](https://img.shields.io/badge/Skills-43%20built--in-blue)](#-all-43-skills) > Run Claude Code, Codex, Gemini CLI, Hermes Agent, and OpenCode in your browser — zero setup, wired to your Databricks workspace. @@ -58,7 +58,7 @@ This isn't just a terminal in the cloud. Running coding agents on Databricks giv | ✂️ **Split Panes** | Run two sessions side by side with a draggable divider | | 🌐 **WebSocket I/O** | Real-time terminal output over WebSocket — zero-latency, eliminates polling delay | | 🔁 **HTTP Polling Fallback** | Automatic fallback via Web Worker when WebSocket is unavailable | -| 🚀 **Parallel Setup** | 7 agent setups run in parallel (~5x faster startup) | +| 🚀 **Parallel Setup** | 6 agent setups run in parallel (~5x faster startup) | | 🔍 **Search** | Find anything in your terminal history (Ctrl+Shift+F) | | 🎤 **Voice Input** | Dictate commands with your mic (Option+V) | | 📋 **Image Paste** | Paste or drag-and-drop images into the terminal — saved to `~/uploads/`, path inserted automatically | @@ -169,7 +169,7 @@ This template repo opens that vision up for every Databricks user — no IDE set ---
-🧠 All 39 Skills +🧠 All 43 Skills ### Databricks Skills (25) — [ai-dev-kit](https://github.com/databricks-solutions/ai-dev-kit) @@ -194,16 +194,100 @@ This template repo opens that vision up for every Databricks user — no IDE set | Ship | finishing-branch, git-worktrees | | Meta | dispatching-agents, writing-skills, using-superpowers | +### BDD Skills (4) + +| Category | Skills | +|----------|--------| +| Testing | bdd-features, bdd-run, bdd-scaffold, bdd-steps | +
-🔌 2 MCP Servers +🔌 MCP Servers + +### Built-in MCP Clients | Server | What it does | |--------|-------------| | **DeepWiki** | Ask questions about any GitHub repo — gets AI-powered answers from the codebase | | **Exa** | Web search and code context retrieval for up-to-date information | +### CoDA MCP Server (exposed at `/mcp`) + +CoDA itself exposes an **MCP server** that any MCP-compatible client can connect to — delegate coding tasks to AI agents running on Databricks, without needing the terminal UI. + +| Tool | Purpose | +|------|---------| +| `coda_run` | Fire-and-forget: submit a coding task, get back immediately | +| `coda_inbox` | Dashboard: see all running/completed/failed tasks at a glance | +| `coda_get_result` | Pull the full structured result of a completed task | + +**Why this matters:** Any tool that speaks MCP can use your Databricks-hosted coding agents — no custom integration needed. + +#### Example: Databricks Genie Code + +Genie Code connects to CoDA's MCP endpoint and delegates coding work to agents running in the background: + +``` +User → Genie Code: "Build me a sales pipeline using the transactions table" + +Genie Code calls coda_run(prompt="Build a sales pipeline...", email="user@company.com", + context='{"tables": ["sales.transactions"]}') + +→ Returns immediately: {task_id: "task-abc", status: "running"} +→ User keeps chatting with Genie Code while the agent works + +User → Genie Code: "How's my pipeline coming?" + +Genie Code calls coda_inbox() +→ {tasks: [{task_id: "task-abc", status: "completed", summary: "Built pipeline.py..."}]} + +Genie Code calls coda_get_result(task_id="task-abc", session_id="sess-123") +→ {summary: "Created pipeline.py with 3 stages", files_changed: ["pipeline.py"], ...} +``` + +#### Connecting MCP Clients (Claude Code, Claude Desktop, Cursor, etc.) + +Databricks Apps use OAuth — not PATs — for authentication. A static `Authorization: Bearer ` header will get a `302` redirect to the OAuth login page. To connect any MCP client, use the **stdio bridge** (`tools/coda-bridge.py`) which injects fresh OAuth tokens automatically via `databricks auth token`. + +**1. Copy the bridge script:** + +```bash +mkdir -p ~/.claude/mcp-bridges +cp tools/coda-bridge.py ~/.claude/mcp-bridges/ +``` + +**2. Add to your MCP client settings** (e.g. `~/.claude/settings.json`): + +```json +"coda-mcp": { + "type": "stdio", + "command": "python3", + "args": ["/path/to/.claude/mcp-bridges/coda-bridge.py"], + "env": { + "CODA_MCP_URL": "https://your-app.databricksapps.com/mcp", + "DATABRICKS_PROFILE": "your-profile" + } +} +``` + +**3. Restart your MCP client.** + +The bridge reads `CODA_MCP_URL` and `DATABRICKS_PROFILE` from environment — no hardcoded values. If you redeploy the app or switch workspaces, just update the `env` block. + +**Prerequisites:** `databricks` CLI installed and authenticated (`databricks auth login -p `), Python 3.8+, no pip dependencies. + +**Troubleshooting:** Bridge logs go to stderr. If you see `Auth failed (302)`, refresh your CLI session with `databricks auth login -p `. See [full setup guide](docs/mcp-client-setup.md) for details. + +#### Task Chaining + +Chain tasks by passing `previous_session_id` — the new agent reads the prior task's results for context: + +``` +coda_run(prompt="Add monitoring to the pipeline", previous_session_id="sess-123") +``` + +See [MCP v2 Design Doc](docs/mcp-v2-background-execution.md) for the full protocol reference.
@@ -211,13 +295,14 @@ This template repo opens that vision up for every Databricks user — no IDE set 🏗️ Architecture ``` -┌─────────────────────┐ WebSocket ┌─────────────────────┐ -│ Browser Client │◄═══════════►│ Gunicorn + Flask │ -│ (xterm.js) │ (primary) │ + Flask-SocketIO │ -│ │───────────►│ (PTY Manager) │ -│ │ HTTP Poll │ │ -│ │ (fallback) │ │ -└─────────────────────┘ └─────────────────────┘ +┌─────────────────────┐ WebSocket ┌──────────────────────────────────┐ +│ Browser Client │◄═══════════►│ uvicorn (ASGI) │ +│ (xterm.js) │ (fallback) │ ├─ python-socketio (Socket.IO) │ +│ │───────────►│ ├─ FastMCP /mcp │ +│ │ HTTP Poll │ └─ WSGIMiddleware(Flask + PTY) │ +│ │ (primary │ │ +│ │ under uvicorn) │ +└─────────────────────┘ └──────────────────────────────────┘ │ │ │ on first load │ on startup ▼ ▼ @@ -235,9 +320,9 @@ This template repo opens that vision up for every Databricks user — no IDE set ### Startup Flow -1. Gunicorn starts, calls `initialize_app()` via `post_worker_init` hook +1. uvicorn starts `coda_mcp.mcp_asgi:app`, which calls `initialize_app()` during ASGI lifespan startup (Flask mounted via `WSGIMiddleware`; MCP mounted at `/mcp` via native ASGI; Socket.IO wraps both) 2. App serves the terminal UI with inline setup progress -3. Background thread runs setup: 5 sequential steps (git config, micro editor, GitHub CLI, Databricks CLI upgrade, content-filter proxy), then 6 agent setups (Claude, Codex, OpenCode, Gemini, Databricks CLI config, MLflow) run in parallel via `ThreadPoolExecutor` +3. Background thread runs setup: 5 sequential steps (git config, micro editor, GitHub CLI, Databricks CLI upgrade, content-filter proxy), then 6 agent setups (`setup/setup_claude.py`, `setup/setup_codex.py`, etc.) run in parallel via `ThreadPoolExecutor` 4. `/api/setup-status` endpoint reports progress to the UI 5. Once complete, the terminal becomes interactive @@ -257,6 +342,7 @@ This template repo opens that vision up for every Databricks user — no IDE set | `/api/resize` | POST | Resize terminal dimensions | | `/api/upload` | POST | Upload file (clipboard image paste) | | `/api/session/close` | POST | Close terminal session | +| `/mcp` | POST | MCP JSON-RPC endpoint (CoDA tools) | ### WebSocket Events (Socket.IO) @@ -293,9 +379,9 @@ This template repo opens that vision up for every Databricks user — no IDE set Single-user app — the owner is resolved via the app's service principal and Apps API (`app.creator`), with no PAT required at deploy time. Authorization checks `X-Forwarded-Email` against `app.creator`. On first terminal session, the user pastes a short-lived PAT interactively. Tokens auto-rotate every 10 minutes (15-minute lifetime), with old tokens proactively revoked. On restart, the user re-pastes (no persistence by design). -### Gunicorn +### Server -Production uses `workers=1` (PTY state is process-local), `threads=16` (concurrent polling + WebSocket), `gthread` worker class, `timeout=60` (long-lived WebSocket connections). +Production uses `uvicorn` (single worker — PTY state is process-local) serving `coda_mcp.mcp_asgi:app`. The ASGI stack composes `python-socketio.ASGIApp` → MCP Streamable HTTP at `/mcp` → `WSGIMiddleware(Flask)` for the terminal UI. WebSocket transport falls back to HTTP polling under uvicorn — the `static/poll-worker.js` Web Worker already handles this transparently. `gunicorn.conf.py` is retained for reference and local WSGI-only dev; it is **not** used in production. @@ -306,27 +392,36 @@ Production uses `workers=1` (PTY state is process-local), `threads=16` (concurre coding-agents-databricks-apps/ ├── app.py # Flask backend + PTY management + setup orchestration ├── app_state.py # Shared app state (setup progress, session registry) -├── app.yaml.template # Databricks Apps deployment config template +├── app.yaml # Databricks Apps deployment config (uvicorn entrypoint) ├── cli_auth.py # Interactive PAT setup + CLI credential writer ├── content_filter_proxy.py # Proxy that sanitises empty-content blocks for OpenCode -├── gunicorn.conf.py # Gunicorn production server config +├── gunicorn.conf.py # Legacy WSGI-only config (unused in production; uvicorn is the entrypoint) ├── pat_rotator.py # Background PAT auto-rotation (10-min cycle) ├── pyproject.toml # Package metadata + uv config (supply-chain guardrails) ├── requirements.txt # Compiled from pyproject.toml (Dependabot compatibility) ├── requirements.lock # Hash-pinned lockfile (auto-regenerated by CI) ├── Makefile # Deploy, redeploy, status, and cleanup targets -├── setup_claude.py # Claude Code CLI + MCP configuration -├── setup_codex.py # Codex CLI configuration -├── setup_gemini.py # Gemini CLI configuration -├── setup_opencode.py # OpenCode configuration -├── setup_databricks.py # Databricks CLI configuration -├── setup_mlflow.py # MLflow tracing auto-configuration -├── setup_proxy.py # Content-filter proxy startup ├── sync_to_workspace.py # Post-commit hook: sync to Workspace -├── install_micro.sh # Micro editor installer -├── install_gh.sh # GitHub CLI installer (OS/arch-aware) -├── install_databricks_cli.sh # Databricks CLI upgrade script -├── utils.py # Utility functions (ensure_https) +├── utils.py # Utility functions (ensure_https, gateway discovery) +├── coda_mcp/ # MCP server package (CoDA — Coding Agents) +│ ├── __init__.py +│ ├── mcp_server.py # FastMCP tool definitions (coda_run, coda_inbox, coda_get_result) +│ ├── mcp_endpoint.py # Flask Blueprint: JSON-RPC /mcp endpoint +│ ├── mcp_asgi.py # ASGI bridge (optional, for native MCP SDK transport) +│ └── task_manager.py # Disk-based session/task state manager +├── setup/ # Agent setup scripts (run at boot) +│ ├── setup_claude.py # Claude Code CLI + MCP configuration +│ ├── setup_codex.py # Codex CLI configuration +│ ├── setup_gemini.py # Gemini CLI configuration +│ ├── setup_opencode.py # OpenCode configuration +│ ├── setup_hermes.py # Hermes Agent configuration +│ ├── setup_databricks.py # Databricks CLI configuration +│ ├── setup_mlflow.py # MLflow tracing auto-configuration +│ └── setup_proxy.py # Content-filter proxy startup +├── scripts/ # Shell scripts +│ ├── install_micro.sh # Micro editor installer +│ ├── install_gh.sh # GitHub CLI installer (OS/arch-aware) +│ └── install_databricks_cli.sh # Databricks CLI upgrade script ├── static/ │ ├── index.html # Terminal UI (xterm.js + split panes + WebSocket) │ ├── favicon.svg # App favicon @@ -340,8 +435,12 @@ coding-agents-databricks-apps/ │ └── workflows/ │ ├── dependency-audit.yml # Weekly CVE audit + lockfile drift check │ └── update-lockfile.yml # Auto-regenerate requirements.lock on push +├── tools/ +│ └── coda-bridge.py # Stdio-to-HTTP MCP bridge (OAuth token injection) └── docs/ ├── deployment.md # Full Databricks Apps deployment guide + ├── mcp-client-setup.md # MCP client setup guide (bridge config) + ├── mcp-v2-background-execution.md # MCP server design doc ├── prd/ # Product requirement documents └── plans/ # Design documentation ``` @@ -352,4 +451,4 @@ coding-agents-databricks-apps/ ## Technologies -Flask · Flask-SocketIO · Socket.IO · Gunicorn · xterm.js · Python PTY · uv · Databricks SDK · Databricks AI Gateway · MLflow +Flask · Flask-SocketIO · Socket.IO · uvicorn · MCP (Streamable HTTP) · xterm.js · Python PTY · uv · Databricks SDK · Databricks AI Gateway · MLflow diff --git a/app.py b/app.py index b5acb65..0386302 100644 --- a/app.py +++ b/app.py @@ -1,3 +1,4 @@ +import asyncio import os import pty import fcntl @@ -46,6 +47,7 @@ CLEANUP_INTERVAL_SECONDS = 900 # Check for stale sessions every 15 min GRACEFUL_SHUTDOWN_WAIT = 3 # Seconds to wait after SIGHUP before SIGKILL MAX_CONCURRENT_SESSIONS = int(os.environ.get("MAX_CONCURRENT_SESSIONS", "5")) +TRANSCRIPT_CAP_BYTES = 10 * 1024 * 1024 # 10 MB soft cap per transcript # Logging setup logging.basicConfig(level=logging.INFO) @@ -58,8 +60,46 @@ app.config['MAX_CONTENT_LENGTH'] = 32 * 1024 * 1024 # 32 MB — aligned with Claude Code's 30 MB file limit # WebSocket support via Flask-SocketIO (simple-websocket transport, threading mode) +# Used for local dev (python app.py). Under uvicorn/ASGI, the AsyncServer in +# mcp_asgi.py intercepts /socket.io/ before WSGIMiddleware, so these handlers +# are only active in WSGI mode. socketio = SocketIO(app, async_mode='threading', cors_allowed_origins=[], logger=False, engineio_logger=False) +# ── ASGI WebSocket support (python-socketio AsyncServer) ───────────── +# Set by mcp_asgi.py at startup. Background threads use _emit_from_thread() +# which routes to the async server (ASGI) or Flask-SocketIO (WSGI) automatically. +_async_sio = None +_event_loop = None + + +def set_async_sio(sio_instance, loop): + """Called by mcp_asgi.py to wire up the ASGI Socket.IO server.""" + global _async_sio, _event_loop + _async_sio = sio_instance + _event_loop = loop + + +def _emit_from_thread(event, data, room=None): + """Thread-safe emit for background threads (PTY reader, cleanup, SIGTERM). + + Routes to AsyncServer (ASGI mode) or Flask-SocketIO (WSGI mode) automatically. + """ + if _async_sio and _event_loop and _event_loop.is_running(): + try: + asyncio.run_coroutine_threadsafe( + _async_sio.emit(event, data, room=room), + _event_loop, + ) + except Exception: + pass + else: + # WSGI mode (local dev) — use Flask-SocketIO directly + try: + socketio.emit(event, data, room=room) + except Exception: + pass + + # Store sessions: {session_id: {"master_fd": fd, "pid": pid, "output_buffer": deque, "lock": Lock, ...}} # sessions_lock guards dict-level ops (add/remove/iterate); each session["lock"] guards per-session state sessions = {} @@ -86,10 +126,7 @@ def handle_sigterm(signum, frame): shutting_down = True logger.info("SIGTERM received — setting shutting_down flag for clients") # Notify WS clients immediately (HTTP poll clients will see shutting_down on next poll) - try: - socketio.emit('shutting_down', {}) - except Exception: - pass + _emit_from_thread('shutting_down', {}) # NOTE: Do not register SIGTERM handler at module level. # It is installed in initialize_app() for gunicorn only. @@ -150,6 +187,11 @@ def _run_step(step_id, command): env.pop("DATABRICKS_CLIENT_ID", None) env.pop("DATABRICKS_CLIENT_SECRET", None) + # Ensure setup scripts can still import from repo root (e.g. `from utils import ...`) + app_dir = os.path.dirname(os.path.abspath(__file__)) + existing_pp = env.get("PYTHONPATH", "") + env["PYTHONPATH"] = f"{app_dir}:{existing_pp}" if existing_pp else app_dir + result = subprocess.run(command, env=env, capture_output=True, text=True, timeout=300) if result.returncode == 0: _update_step(step_id, status="complete", completed_at=time.time()) @@ -370,8 +412,14 @@ def _configure_all_cli_auth(token): # 3. Re-run Codex, OpenCode, Gemini setup scripts with token in env # They are idempotent: detect CLI already installed, just write config files - env = {**os.environ, "DATABRICKS_TOKEN": token} - for script in ["setup_codex.py", "setup_opencode.py", "setup_gemini.py", "setup_hermes.py"]: + app_dir = os.path.dirname(os.path.abspath(__file__)) + existing_pp = os.environ.get("PYTHONPATH", "") + env = { + **os.environ, + "DATABRICKS_TOKEN": token, + "PYTHONPATH": f"{app_dir}:{existing_pp}" if existing_pp else app_dir, + } + for script in ["setup/setup_codex.py", "setup/setup_opencode.py", "setup/setup_gemini.py", "setup/setup_hermes.py"]: try: result = subprocess.run( ["uv", "run", "python", script], @@ -410,26 +458,26 @@ def run_setup(): _update_step("git", status="error", completed_at=time.time(), error=str(e)) _run_step("micro", ["bash", "-c", - "mkdir -p ~/.local/bin && bash install_micro.sh && mv micro ~/.local/bin/ 2>/dev/null || true"]) + "mkdir -p ~/.local/bin && bash scripts/install_micro.sh && mv micro ~/.local/bin/ 2>/dev/null || true"]) - _run_step("gh", ["bash", "install_gh.sh"]) + _run_step("gh", ["bash", "scripts/install_gh.sh"]) # --- Upgrade Databricks CLI (runtime image ships an older version) --- - _run_step("dbcli", ["bash", "install_databricks_cli.sh"]) + _run_step("dbcli", ["bash", "scripts/install_databricks_cli.sh"]) # --- Content-filter proxy (must be running before OpenCode starts) --- # Sanitizes requests/responses between OpenCode and Databricks # (see OpenCode #5028, docs/plans/2026-03-11-litellm-empty-content-blocks-design.md) - _run_step("proxy", ["uv", "run", "python", "setup_proxy.py"]) + _run_step("proxy", ["uv", "run", "python", "setup/setup_proxy.py"]) # --- Parallel agent setup (all independent of each other) --- parallel_steps = [ - ("claude", ["uv", "run", "python", "setup_claude.py"]), - ("codex", ["uv", "run", "python", "setup_codex.py"]), - ("opencode", ["uv", "run", "python", "setup_opencode.py"]), - ("gemini", ["uv", "run", "python", "setup_gemini.py"]), - ("hermes", ["uv", "run", "python", "setup_hermes.py"]), - ("databricks", ["uv", "run", "python", "setup_databricks.py"]), + ("claude", ["uv", "run", "python", "setup/setup_claude.py"]), + ("codex", ["uv", "run", "python", "setup/setup_codex.py"]), + ("opencode", ["uv", "run", "python", "setup/setup_opencode.py"]), + ("gemini", ["uv", "run", "python", "setup/setup_gemini.py"]), + ("hermes", ["uv", "run", "python", "setup/setup_hermes.py"]), + ("databricks", ["uv", "run", "python", "setup/setup_databricks.py"]), ] with ThreadPoolExecutor(max_workers=len(parallel_steps)) as executor: @@ -442,7 +490,7 @@ def run_setup(): # --- MLflow setup runs AFTER claude setup to avoid settings.json race --- # setup_mlflow.py merges env vars into ~/.claude/settings.json which # setup_claude.py also writes; running sequentially prevents clobbering. - _run_step("mlflow", ["uv", "run", "python", "setup_mlflow.py"]) + _run_step("mlflow", ["uv", "run", "python", "setup/setup_mlflow.py"]) # Sync latest token into all CLI configs — covers the race where PAT # rotation happened while a setup script was still installing (the @@ -580,7 +628,132 @@ def _check_ws_authorization(): return True -# ── WebSocket Event Handlers ────────────────────────────────────────────── +def _check_ws_authorization_from_environ(environ): + """Check authorization from WSGI environ dict (for ASGI WebSocket via python-socketio). + + Same logic as _check_ws_authorization() but reads headers from the environ + dict instead of Flask's request context. WSGI environ stores HTTP headers as + HTTP_X_FORWARDED_EMAIL (uppercase, underscores, HTTP_ prefix). + """ + if not app_owner: + if _is_databricks_apps(): + logger.error("SECURITY: app_owner not resolved — denying WebSocket (fail-closed)") + return False + return True # Local dev only + + raw_user = ( + environ.get("HTTP_X_FORWARDED_EMAIL") + or environ.get("HTTP_X_FORWARDED_USER") + or environ.get("HTTP_X_DATABRICKS_USER_EMAIL") + ) + current_user = raw_user.lower() if raw_user else raw_user + + if not current_user: + if _is_databricks_apps(): + logger.warning("No user identity in WebSocket request on Databricks Apps — denying") + return False + return True # Local dev only + + if current_user != app_owner: + logger.warning(f"WebSocket unauthorized: {current_user} (owner: {app_owner})") + return False + return True + + +def register_sio_handlers(sio): + """Register Socket.IO event handlers on an AsyncServer for ASGI mode. + + Called by mcp_asgi.py. The handlers mirror the Flask-SocketIO handlers below + but use python-socketio's async API (explicit sid, enter_room/leave_room, + async def, ConnectionRefusedError for auth denial). + """ + + @sio.on('connect') + async def handle_connect(sid, environ, auth): + # Capture event loop on first connection for _emit_from_thread() + set_async_sio(sio, asyncio.get_running_loop()) + + # Diagnostic: log transport and header presence for debugging proxy behavior + transport = environ.get('QUERY_STRING', '') + has_email = bool(environ.get('HTTP_X_FORWARDED_EMAIL')) + has_user = bool(environ.get('HTTP_X_FORWARDED_USER')) + logger.info(f"WS connect: sid={sid}, qs={transport}, " + f"has_email={has_email}, has_user={has_user}") + + if not _check_ws_authorization_from_environ(environ): + raise ConnectionRefusedError('unauthorized') + logger.info("WebSocket client connected (ASGI)") + + @sio.on('join_session') + async def handle_join_session(sid, data): + session_id = data.get('session_id') + if not session_id: + return {'status': 'error', 'message': 'session_id required'} + sess = _get_session(session_id) + if not sess: + return {'status': 'error', 'message': 'Session not found'} + with sess["lock"]: + sess["last_poll_time"] = time.time() + sess["output_buffer"].clear() + await sio.enter_room(sid, session_id) + logger.info(f"WebSocket client joined session room {session_id}") + return {'status': 'ok'} + + @sio.on('leave_session') + async def handle_leave_session(sid, data): + session_id = data.get('session_id') + if session_id: + await sio.leave_room(sid, session_id) + logger.info(f"WebSocket client left session room {session_id}") + + @sio.on('terminal_input') + async def handle_terminal_input(sid, data): + session_id = data.get('session_id') + input_data = data.get('input', '') + sess = _get_session(session_id) + if not sess: + return + with sess["lock"]: + sess["last_poll_time"] = time.time() + fd = sess["master_fd"] + try: + os.write(fd, input_data.encode()) + except OSError as e: + logger.warning(f"WebSocket input write error for {session_id}: {e}") + + @sio.on('terminal_resize') + async def handle_terminal_resize(sid, data): + session_id = data.get('session_id') + cols = data.get('cols', 80) + rows = data.get('rows', 24) + sess = _get_session(session_id) + if not sess: + return + with sess["lock"]: + sess["last_poll_time"] = time.time() + fd = sess["master_fd"] + try: + winsize = struct.pack("HHHH", rows, cols, 0, 0) + fcntl.ioctl(fd, termios.TIOCSWINSZ, winsize) + except OSError as e: + logger.warning(f"WebSocket resize error for {session_id}: {e}") + + @sio.on('heartbeat') + async def handle_heartbeat(sid, data): + session_ids = data.get('session_ids', []) + now = time.time() + for s_id in session_ids: + sess = _get_session(s_id) + if sess: + with sess["lock"]: + sess["last_poll_time"] = now + + @sio.on('disconnect') + async def handle_disconnect(sid): + logger.info("WebSocket client disconnected (ASGI)") + + +# ── WebSocket Event Handlers (Flask-SocketIO — WSGI/local dev only) ────── @socketio.on('connect') def handle_ws_connect(): @@ -686,6 +859,42 @@ def _get_session(session_id): return sessions.get(session_id) +def _tee_transcript_chunk(session, output: bytes, cap: int = TRANSCRIPT_CAP_BYTES) -> None: + """Append PTY output to the transcript file. Single-writer (read_pty_output). + + All file-handle access is under ``session["lock"]`` so we never race the + Timer-driven close path in ``terminate_session``. The ``ValueError`` catch + is belt-and-suspenders for the tiny window where the handle is closed + between the ``is not None`` check and the actual ``write`` call (the lock + prevents this, but be defensive). + """ + with session["lock"]: + fh = session.get("transcript_fh") + written = session.get("transcript_bytes", 0) + if fh is None: + return + remaining = cap - written + if remaining <= 0: + return + chunk = output[:remaining] + try: + fh.write(chunk) + fh.flush() + session["transcript_bytes"] = written + len(chunk) + if len(chunk) < len(output): + fh.write(b"\n[transcript truncated at %d bytes]\n" % cap) + fh.flush() + fh.close() + session["transcript_fh"] = None + except (OSError, ValueError) as exc: + logger.warning("transcript write failed: %s", exc) + try: + fh.close() + except Exception: + pass + session["transcript_fh"] = None + + def read_pty_output(session_id, fd): """Background thread to read PTY output into buffer and push via WebSocket.""" session = _get_session(session_id) @@ -711,12 +920,11 @@ def read_pty_output(session_id, fd): session["output_buffer"].append(decoded) session["last_poll_time"] = time.time() # Keep session alive during WS output # Push via WebSocket to the session room (AC-8) - try: - socketio.emit('terminal_output', + _emit_from_thread('terminal_output', {'session_id': session_id, 'output': decoded}, room=session_id) - except Exception: - pass # No WebSocket clients — HTTP polling handles it + # Tee to transcript file if enabled for this session + _tee_transcript_chunk(session, output) else: # select timed out — check if process is still alive try: @@ -731,10 +939,7 @@ def read_pty_output(session_id, fd): break # Process exited or fd closed — notify WebSocket clients (AC-9) - try: - socketio.emit('session_exited', {'session_id': session_id}, room=session_id) - except Exception: - pass + _emit_from_thread('session_exited', {'session_id': session_id}, room=session_id) logger.info(f"Session {session_id} process exited") @@ -748,10 +953,21 @@ def terminate_session(session_id, pid, master_fd): logger.info(f"Terminating stale session {session_id} (pid={pid})") # Notify WebSocket clients that the session is closed - try: - socketio.emit('session_closed', {'session_id': session_id}, room=session_id) - except Exception: - pass + _emit_from_thread('session_closed', {'session_id': session_id}, room=session_id) + + # Close transcript handle (if any) under per-session lock; swap-then-close + # outside the lock to avoid blocking on slow filesystems. + with sessions_lock: + sess = sessions.get(session_id) + if sess is not None: + with sess["lock"]: + transcript_fh = sess.get("transcript_fh") + sess["transcript_fh"] = None + if transcript_fh is not None: + try: + transcript_fh.close() + except Exception: + pass try: os.kill(pid, signal.SIGHUP) @@ -858,7 +1074,7 @@ def cleanup_stale_sessions(): def authorize_request(): """Check authorization before processing any request.""" # Skip auth for health check, setup status, and Socket.IO (has own auth via connect event) - if request.path in ("/health", "/api/setup-status", "/api/pat-status", "/api/configure-pat", "/api/app-state") or request.path.startswith("/socket.io"): + if request.path in ("/health", "/api/setup-status", "/api/pat-status", "/api/configure-pat", "/api/app-state") or request.path.startswith("/socket.io") or request.path.startswith("/mcp"): return None authorized, user = check_authorization() @@ -873,6 +1089,10 @@ def authorize_request(): @app.after_request def set_security_headers(response): + # MCP endpoint handles its own CORS/headers — skip security headers + # that might interfere (CSP connect-src, X-Frame-Options, etc.) + if request.path.startswith("/mcp"): + return response response.headers["X-Content-Type-Options"] = "nosniff" response.headers["X-Frame-Options"] = "DENY" response.headers["X-XSS-Protection"] = "1; mode=block" @@ -936,15 +1156,38 @@ def list_sessions(): @app.route("/api/session/attach", methods=["POST"]) def attach_session(): - """Reattach to an existing session — returns buffered output for replay.""" + """Reattach to an existing session — returns buffered output for replay. + + If the live PTY is gone but an on-disk transcript exists for this + pty_session_id, return the transcript as ``output`` with ``replay: True``. + """ data = request.get_json(silent=True) or {} session_id = data.get("session_id", "") sess = _get_session(session_id) if not sess or sess.get("exited"): + # Replay fallback: look up transcript.log by pty_session_id + from coda_mcp import task_manager as _tm + tdir = _tm.find_task_dir_by_pty_session(session_id) + if tdir: + transcript = os.path.join(tdir, "transcript.log") + if os.path.isfile(transcript): + try: + with open(transcript, "rb") as f: + content = f.read() + return jsonify({ + "session_id": session_id, + "label": "hermes-mcp (replay)", + "output": [content.decode("utf-8", errors="replace")], + "replay": True, + "process": None, + "created_at": None, + }) + except OSError: + pass return jsonify({"error": "Session not found or exited"}), 404 - # Reset idle clock so the 24h reaper starts fresh + # Existing live-attach path sess["last_poll_time"] = time.time() return jsonify({ @@ -1083,7 +1326,8 @@ def create_session(): """Create a new terminal session.""" # Quick reject before forking a PTY (approximate — authoritative check below) with sessions_lock: - if len(sessions) >= MAX_CONCURRENT_SESSIONS: + active = sum(1 for s in sessions.values() if not s.get("grace")) + if active >= MAX_CONCURRENT_SESSIONS: return jsonify({"error": f"Maximum {MAX_CONCURRENT_SESSIONS} concurrent sessions reached. Close an existing session first."}), 429 data = request.get_json(silent=True) or {} @@ -1122,7 +1366,8 @@ def create_session(): with sessions_lock: # Authoritative check under the same lock as insertion — prevents # TOCTOU race where two concurrent requests both pass the early check. - if len(sessions) >= MAX_CONCURRENT_SESSIONS: + active = sum(1 for s in sessions.values() if not s.get("grace")) + if active >= MAX_CONCURRENT_SESSIONS: os.close(master_fd) try: os.kill(pid, signal.SIGKILL) @@ -1151,6 +1396,153 @@ def create_session(): return jsonify({"error": str(e)}), 500 +# ── MCP Integration Helpers ────────────────────────────────────────── + + +def mcp_create_pty_session(label: str = "hermes-mcp", transcript_path: str | None = None) -> str: + """Create a PTY session for MCP use. Returns the PTY session_id.""" + with sessions_lock: + active = sum(1 for s in sessions.values() if not s.get("grace")) + if active >= MAX_CONCURRENT_SESSIONS: + raise RuntimeError( + f"Maximum {MAX_CONCURRENT_SESSIONS} concurrent sessions reached." + ) + + master_fd, slave_fd = pty.openpty() + + shell_env = os.environ.copy() + shell_env["TERM"] = "xterm-256color" + shell_env.pop("CLAUDECODE", None) + shell_env.pop("CLAUDE_CODE_SESSION", None) + shell_env.pop("DATABRICKS_TOKEN", None) + shell_env.pop("DATABRICKS_HOST", None) + shell_env.pop("GEMINI_API_KEY", None) + if not shell_env.get("HOME") or shell_env["HOME"] == "/": + shell_env["HOME"] = "/app/python/source_code" + local_bin = f"{shell_env['HOME']}/.local/bin" + shell_env["PATH"] = f"{local_bin}:{shell_env.get('PATH', '')}" + + projects_dir = os.path.join(shell_env["HOME"], "projects") + os.makedirs(projects_dir, exist_ok=True) + + pid = subprocess.Popen( + ["/bin/bash"], + stdin=slave_fd, + stdout=slave_fd, + stderr=slave_fd, + preexec_fn=os.setsid, + env=shell_env, + cwd=projects_dir, + ).pid + os.close(slave_fd) + + # Open transcript file (if requested) before locking the session dict. + transcript_fh = None + if transcript_path: + try: + parent_dir = os.path.dirname(transcript_path) + if parent_dir: + os.makedirs(parent_dir, exist_ok=True) + transcript_fh = open(transcript_path, "ab", buffering=0) + os.fchmod(transcript_fh.fileno(), 0o600) + except OSError as exc: + logger.warning("Could not open transcript at %s: %s", transcript_path, exc) + transcript_fh = None + + session_id = str(uuid.uuid4()) + + try: + with sessions_lock: + active = sum(1 for s in sessions.values() if not s.get("grace")) + if active >= MAX_CONCURRENT_SESSIONS: + os.close(master_fd) + try: + os.kill(pid, signal.SIGKILL) + except OSError: + pass + raise RuntimeError( + f"Maximum {MAX_CONCURRENT_SESSIONS} concurrent sessions reached." + ) + sessions[session_id] = { + "master_fd": master_fd, + "pid": pid, + "output_buffer": deque(maxlen=1000), + "lock": threading.Lock(), + "last_poll_time": time.time(), + "created_at": time.time(), + "label": label, + "transcript_path": transcript_path if transcript_fh else None, + "transcript_fh": transcript_fh, + "transcript_bytes": 0, + "grace": False, + } + + thread = threading.Thread( + target=read_pty_output, args=(session_id, master_fd), daemon=True + ) + thread.start() + except BaseException: + # Roll back transcript open if anything below it raises before the + # session is fully wired. The PTY itself is cleaned up by existing + # error paths; this is just the transcript handle. + if transcript_fh is not None: + try: + transcript_fh.close() + except Exception: + pass + raise + + return session_id + + +def mcp_send_input(session_id: str, data: str): + """Send input to a PTY session.""" + session = _get_session(session_id) + if not session: + raise RuntimeError(f"Session {session_id} not found") + with session["lock"]: + os.write(session["master_fd"], data.encode()) + + +def mcp_close_pty_session(session_id: str): + """Close a PTY session.""" + session = _get_session(session_id) + if not session: + return + terminate_session(session_id, session["pid"], session["master_fd"]) + + +def _mark_grace_for_session(session_id: str) -> None: + """Mark a PTY session as 'in grace period' so it doesn't count toward + MAX_CONCURRENT_SESSIONS. Called by ``_watch_task`` immediately before + scheduling the deferred close Timer. + + No-op if the session does not exist (e.g., already torn down). + """ + with sessions_lock: + sess = sessions.get(session_id) + if sess is None: + return + with sess["lock"]: + sess["grace"] = True + + +def _bump_session_last_poll(session_id: str, delta_s: float) -> None: + """Advance ``last_poll_time`` by ``delta_s`` so the idle reaper can't + preempt the Timer's deferred close. Defensive: at the current 24h + SESSION_TIMEOUT_SECONDS the reaper would never win anyway, but a future + tuning shouldn't break the grace window. + + No-op if the session does not exist. + """ + with sessions_lock: + sess = sessions.get(session_id) + if sess is None: + return + with sess["lock"]: + sess["last_poll_time"] = time.time() + delta_s + + @app.route("/api/input", methods=["POST"]) def send_input(): """Send input to the terminal.""" @@ -1368,6 +1760,20 @@ def initialize_app(local_dev=False): logger.info(f"Started session cleanup thread (timeout={SESSION_TIMEOUT_SECONDS}s, interval={CLEANUP_INTERVAL_SECONDS}s)") +# ── MCP Endpoint ───────────────────────────────────────────────────── +from coda_mcp.mcp_endpoint import mcp_bp +from coda_mcp.mcp_server import set_app_hooks + +app.register_blueprint(mcp_bp) + +# Wire MCP tools to PTY infrastructure +set_app_hooks( + create_session_fn=mcp_create_pty_session, + send_input_fn=mcp_send_input, + close_session_fn=mcp_close_pty_session, +) + + if __name__ == "__main__": # Local dev — no SIGTERM handler (SIG_DFL), no shutting_down flag initialize_app(local_dev=True) diff --git a/app.yaml b/app.yaml index 4d20047..b84a8bc 100644 --- a/app.yaml +++ b/app.yaml @@ -1,6 +1,17 @@ +# Production entrypoint is uvicorn (ASGI), not gunicorn. Required because +# the MCP server at /mcp uses FastMCP.streamable_http_app(), a native ASGI +# transport that cannot be served by gunicorn's WSGI workers. Flask is +# mounted via WSGIMiddleware inside coda_mcp.mcp_asgi alongside MCP and +# Socket.IO. WebSocket transport falls back to HTTP polling under uvicorn — +# acceptable because static/poll-worker.js already implements the fallback. +# gunicorn.conf.py is retained for legacy WSGI-only local dev; not used here. command: - - gunicorn - - app:app + - uvicorn + - coda_mcp.mcp_asgi:app + - --host + - 0.0.0.0 + - --port + - "8000" env: - name: HOME value: /app/python/source_code diff --git a/cli_auth.py b/cli_auth.py index 61c9f25..53c2a25 100644 --- a/cli_auth.py +++ b/cli_auth.py @@ -35,6 +35,7 @@ def _update_claude(token): settings["env"]["ANTHROPIC_AUTH_TOKEN"] = token with open(path, "w") as f: json.dump(settings, f, indent=2) + os.chmod(path, 0o600) except (OSError, json.JSONDecodeError): pass # file doesn't exist yet — initial setup hasn't run @@ -59,6 +60,7 @@ def _update_opencode(token): if changed: with open(path, "w") as f: json.dump(auth, f, indent=2) + os.chmod(path, 0o600) except (OSError, json.JSONDecodeError): pass @@ -84,6 +86,7 @@ def _update_hermes(token): if new_content != content: with open(path, "w") as f: f.write(new_content) + os.chmod(path, 0o600) except OSError: pass @@ -102,5 +105,6 @@ def _replace_dotenv_key(path, key, value): if new_content != content: with open(path, "w") as f: f.write(new_content) + os.chmod(path, 0o600) except OSError: pass diff --git a/coda_mcp/__init__.py b/coda_mcp/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/coda_mcp/mcp_asgi.py b/coda_mcp/mcp_asgi.py new file mode 100644 index 0000000..f745e32 --- /dev/null +++ b/coda_mcp/mcp_asgi.py @@ -0,0 +1,122 @@ +"""Native MCP ASGI app with WebSocket support for terminal I/O. + +Architecture (all on one port, one uvicorn process): + + socketio.ASGIApp ← /socket.io/ → native ASGI WebSocket (terminal) + └── mcp_starlette ← /mcp → FastMCP Streamable HTTP (Genie Code) + └── WSGI(Flask) ← /* → REST API, static files (HTTP only) + +Usage in app.yaml:: + + command: ["uvicorn", "coda_mcp.mcp_asgi:app", "--host", "0.0.0.0", "--port", "8000"] +""" + +import os +import logging +import warnings + +import socketio as socketio_lib +from starlette.middleware.cors import CORSMiddleware + +with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + from starlette.middleware.wsgi import WSGIMiddleware + +from coda_mcp.mcp_server import mcp as mcp_instance, set_app_hooks +from coda_mcp import url_builder +from utils import ensure_https + +logger = logging.getLogger(__name__) + + +class AppUrlCaptureMiddleware: + """Capture X-Forwarded-Host (or Host) from every inbound HTTP request and + populate url_builder._app_url_cache. Used so MCP tools can return a + working viewer_url without manual configuration. + + Caveat: /socket.io/ traffic is intercepted by socketio.ASGIApp *before* + reaching mcp_starlette, so WebSocket connect requests never hit this + middleware. This is fine in practice — every HTTP request to /mcp and to + Flask routes does hit it, which is enough to keep the cache hot. + """ + + def __init__(self, app): + self.app = app + + async def __call__(self, scope, receive, send): + if scope.get("type") == "http": + headers = dict(scope.get("headers") or []) + host_bytes = headers.get(b"x-forwarded-host") or headers.get(b"host") + if host_bytes: + try: + url_builder.capture_from_headers(host_bytes.decode("latin-1")) + except Exception: + pass + await self.app(scope, receive, send) + +# ── Build allowed origins ───────────────────────────────────────── +# The browser connects from the app's own URL (e.g. mcp-test-coda-*.databricksapps.com) +# which differs from DATABRICKS_HOST (workspace URL). Databricks proxy handles auth, +# so Socket.IO CORS can safely allow all origins. Starlette CORSMiddleware below +# uses the same list for MCP/Flask routes. +_databricks_host = os.environ.get("DATABRICKS_HOST", "") +ALLOWED_ORIGINS = [] +if _databricks_host: + ALLOWED_ORIGINS.append(ensure_https(_databricks_host).rstrip("/")) + +# ── Import and initialize Flask app ──────────────────────────────── +from app import ( + app as flask_app, + initialize_app, + mcp_create_pty_session, + mcp_send_input, + mcp_close_pty_session, + register_sio_handlers, +) + +initialize_app() + +# Wire MCP tools to PTY infrastructure +set_app_hooks( + create_session_fn=mcp_create_pty_session, + send_input_fn=mcp_send_input, + close_session_fn=mcp_close_pty_session, +) + +# ── Async Socket.IO server (native ASGI WebSocket) ─────────────── +# python-socketio AsyncServer handles /socket.io/ with real WebSocket, +# eliminating the WSGIMiddleware limitation that forced HTTP polling fallback. +sio = socketio_lib.AsyncServer( + async_mode='asgi', + cors_allowed_origins='*', # App URL differs from DATABRICKS_HOST; proxy handles auth + logger=False, + engineio_logger=False, +) + +# Register terminal I/O event handlers (connect, join_session, terminal_input, etc.) +register_sio_handlers(sio) + +# ── Build the ASGI app per Genie Code docs ───────────────────────── +mcp_starlette = mcp_instance.streamable_http_app() + +# Mount Flask as catch-all via WSGI adapter (HTTP routes only) +flask_asgi = WSGIMiddleware(flask_app.wsgi_app) +mcp_starlette.mount("/", app=flask_asgi) + +# CORS for MCP and Flask routes +mcp_starlette.add_middleware( + CORSMiddleware, + allow_origins=ALLOWED_ORIGINS or ["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + +# Capture X-Forwarded-Host into url_builder cache (for MCP viewer_url). +# Added AFTER CORS so it wraps the CORS-handled request. +mcp_starlette.add_middleware(AppUrlCaptureMiddleware) + +# ── Top-level ASGI app ──────────────────────────────────────────── +# socketio.ASGIApp intercepts /socket.io/ for WebSocket + polling, +# passes everything else to mcp_starlette (MCP at /mcp, Flask at /) +app = socketio_lib.ASGIApp(sio, other_asgi_app=mcp_starlette) diff --git a/coda_mcp/mcp_endpoint.py b/coda_mcp/mcp_endpoint.py new file mode 100644 index 0000000..45b84ad --- /dev/null +++ b/coda_mcp/mcp_endpoint.py @@ -0,0 +1,177 @@ +"""Flask Blueprint fallback for MCP JSON-RPC. + +NOTE: This is NOT the production path. Production deployment uses +`coda_mcp.mcp_asgi:app` served by uvicorn, which mounts the native MCP +SDK Streamable HTTP transport at /mcp. This module is a Flask-native +JSON-RPC fallback used only under WSGI runtimes (gunicorn local dev, +tests that exercise the Flask test client without spinning up ASGI). + +Both paths expose the same three tools (coda_run, coda_inbox, +coda_get_result) and produce equivalent JSON-RPC responses, so switching +between them is transparent to MCP clients. +""" +import asyncio +import json +import logging +from flask import Blueprint, request, jsonify + +logger = logging.getLogger(__name__) + +mcp_bp = Blueprint("mcp", __name__) + +# Import tool functions from mcp_server.py +from coda_mcp.mcp_server import ( + mcp as mcp_instance, + coda_run, + coda_inbox, + coda_get_result, +) + +# Tool function dispatch +_TOOL_DISPATCH = { + "coda_run": coda_run, + "coda_inbox": coda_inbox, + "coda_get_result": coda_get_result, +} + +SERVER_INFO = { + "name": "coda", + "version": "1.0.0", +} + +CAPABILITIES = { + "tools": {"listChanged": False}, +} + + + +def _cors_headers(): + """Build CORS response headers. + + Permissive CORS for /mcp — the Databricks Apps proxy handles auth. + """ + headers = {} + origin = request.headers.get("Origin", "") + if origin: + headers["Access-Control-Allow-Origin"] = origin + headers["Access-Control-Allow-Methods"] = "GET, POST, DELETE, OPTIONS" + # Explicitly list all headers Genie Code might send + # (wildcard * is incompatible with credentials=true per CORS spec) + allowed_headers = ", ".join([ + "Content-Type", "Authorization", "Accept", + "Mcp-Session-Id", "X-Request-Id", "X-Requested-With", + "X-Forwarded-Email", "X-Forwarded-User", "X-Databricks-User-Email", + "Cookie", "Origin", "Referer", + ]) + headers["Access-Control-Allow-Headers"] = allowed_headers + headers["Access-Control-Allow-Credentials"] = "true" + headers["Access-Control-Max-Age"] = "86400" + return headers + + +@mcp_bp.route("/mcp", methods=["POST", "OPTIONS", "GET"]) +def mcp_handler(): + # Handle CORS preflight + if request.method == "OPTIONS": + resp = jsonify({}) + resp.status_code = 204 + for k, v in _cors_headers().items(): + resp.headers[k] = v + return resp + + # Handle GET for SSE (not supported in stateless mode) + if request.method == "GET": + resp = jsonify({"error": "SSE not supported. Use POST."}) + resp.status_code = 405 + return resp + + # Origin validation skipped — Databricks Apps proxy handles auth. + + data = request.get_json(silent=True) or {} + method = data.get("method", "") + req_id = data.get("id") + params = data.get("params", {}) + + # Route by method + if method == "initialize": + result = { + "protocolVersion": params.get("protocolVersion", "2025-03-26"), + "capabilities": CAPABILITIES, + "serverInfo": SERVER_INFO, + "instructions": mcp_instance._instructions if hasattr(mcp_instance, '_instructions') else "", + } + resp = jsonify({"jsonrpc": "2.0", "id": req_id, "result": result}) + + elif method == "notifications/initialized": + # No-op acknowledgment — return empty OK + resp = jsonify({}) + resp.status_code = 200 + + elif method == "tools/list": + tools = _build_tools_list() + resp = jsonify({"jsonrpc": "2.0", "id": req_id, "result": {"tools": tools}}) + + elif method == "tools/call": + tool_name = params.get("name", "") + arguments = params.get("arguments", {}) + tool_fn = _TOOL_DISPATCH.get(tool_name) + if not tool_fn: + resp = jsonify({ + "jsonrpc": "2.0", "id": req_id, + "error": {"code": -32601, "message": f"Unknown tool: {tool_name}"} + }) + else: + try: + # Tool functions are async — run them + result_str = asyncio.run(tool_fn(**arguments)) + result_data = json.loads(result_str) + resp = jsonify({ + "jsonrpc": "2.0", "id": req_id, + "result": { + "content": [{"type": "text", "text": result_str}], + "isError": "error" in result_data, + } + }) + except Exception as e: + resp = jsonify({ + "jsonrpc": "2.0", "id": req_id, + "error": {"code": -32603, "message": str(e)} + }) + + elif method == "ping": + resp = jsonify({"jsonrpc": "2.0", "id": req_id, "result": {}}) + + else: + resp = jsonify({ + "jsonrpc": "2.0", "id": req_id, + "error": {"code": -32601, "message": f"Method not found: {method}"} + }) + + # Add CORS headers + for k, v in _cors_headers().items(): + resp.headers[k] = v + + return resp + + +def _build_tools_list(): + """Extract tool definitions from FastMCP registry.""" + tools = [] + # Access FastMCP's internal tool manager + tool_manager = mcp_instance._tool_manager + for name, tool in tool_manager._tools.items(): + tool_dict = { + "name": tool.name, + "description": tool.description or "", + "inputSchema": tool.parameters if hasattr(tool, 'parameters') else {}, + } + if hasattr(tool, 'annotations') and tool.annotations: + tool_dict["annotations"] = {} + if tool.annotations.readOnlyHint is not None: + tool_dict["annotations"]["readOnlyHint"] = tool.annotations.readOnlyHint + if tool.annotations.destructiveHint is not None: + tool_dict["annotations"]["destructiveHint"] = tool.annotations.destructiveHint + if tool.annotations.idempotentHint is not None: + tool_dict["annotations"]["idempotentHint"] = tool.annotations.idempotentHint + tools.append(tool_dict) + return tools diff --git a/coda_mcp/mcp_server.py b/coda_mcp/mcp_server.py new file mode 100644 index 0000000..2fb98bb --- /dev/null +++ b/coda_mcp/mcp_server.py @@ -0,0 +1,436 @@ +"""MCP server exposing CoDA session/task tools via FastMCP. + +v2: Background execution + inbox pattern. +- ``coda_run`` — fire-and-forget task submission (auto-creates ephemeral session) +- ``coda_inbox`` — dashboard of all background tasks +- ``coda_get_result`` — pull full structured result for a completed task + +Delegates all disk state to ``task_manager.py``. PTY operations are +handled through optional app hooks set via ``set_app_hooks()``. + +Run standalone for testing:: + + python mcp_server.py # stdio transport +""" + +import json +import logging +import os +import threading +import time + +from mcp.server.fastmcp import FastMCP +from mcp.server.fastmcp.server import TransportSecuritySettings +from mcp.types import ToolAnnotations + +from coda_mcp import task_manager +from coda_mcp import url_builder + +logger = logging.getLogger(__name__) + +# ── FastMCP instance ──────────────────────────────────────────────── + +# Build allowed origins from DATABRICKS_HOST for Genie Code requests +_databricks_host = os.environ.get("DATABRICKS_HOST", "") +_allowed_origins = [] +if _databricks_host: + # Ensure https:// prefix, strip trailing slash + origin = _databricks_host if _databricks_host.startswith("https://") else f"https://{_databricks_host}" + _allowed_origins.append(origin.rstrip("/")) + +mcp = FastMCP( + "coda", + instructions=( + "CoDA MCP server — delegate coding tasks to AI agents on Databricks.\n\n" + "CRITICAL — FIRE AND FORGET:\n" + "coda_run submits work and returns IMMEDIATELY. The task runs autonomously " + "in the background. After calling coda_run, DO NOT call coda_inbox or " + "coda_get_result to check on it. Do NOT loop, poll, or wait. Simply tell " + "the user the task was submitted and MOVE ON to their next request.\n\n" + "WHEN TO CHECK INBOX:\n" + "Call coda_inbox ONLY when the user explicitly asks about background tasks " + "(e.g. 'how's my task going?', 'check on that', 'what's in my inbox'). " + "Never call it proactively, automatically, or in a loop.\n\n" + "WORKFLOW:\n" + "1) coda_run — submit work, get back task_id. Tell user it's running. Stop.\n" + "2) Continue chatting about other topics — the task runs independently.\n" + "3) coda_inbox — ONLY when user asks. Shows all tasks from last 24h.\n" + "4) coda_get_result — for completed tasks, get full structured output.\n\n" + "CHAINING: pass previous_session_id from a completed task's session_id " + "to give the new task context of what was done before.\n\n" + "SHARE THE LIVE URL: When coda_run returns a viewer_url field (non-null), " + "mention it to the user in plain text (e.g. \"you can watch progress at " + "\"). The URL is safe to share — it points to the same Databricks App " + "the user is already authenticated against. Do this on the first mention " + "of the task and any time the user asks where the task is or how to see it." + ), + stateless_http=True, + json_response=True, + transport_security=TransportSecuritySettings( + enable_dns_rebinding_protection=False, + ), +) + +# ── App hooks (PTY integration) ───────────────────────────────────── + +_app_create_session = None +_app_send_input = None +_app_close_session = None +_app_mark_grace = None +_app_bump_poll = None + +GRACE_PERIOD_S = 300 # 5 minutes + + +def set_app_hooks( + create_session_fn, + send_input_fn, + close_session_fn, + mark_grace_fn=None, + bump_poll_fn=None, +): + """Wire up Flask app callbacks for PTY operations. + + The two new optional hooks (mark_grace, bump_poll) are used by ``_watch_task`` + to defer PTY close by ``GRACE_PERIOD_S`` after task completion so live viewers + can keep watching for a few minutes. + """ + global _app_create_session, _app_send_input, _app_close_session + global _app_mark_grace, _app_bump_poll + _app_create_session = create_session_fn + _app_send_input = send_input_fn + _app_close_session = close_session_fn + _app_mark_grace = mark_grace_fn + _app_bump_poll = bump_poll_fn + + +# ── Background watcher ────────────────────────────────────────────── + + +def _watch_task(session_id: str, task_id: str, timeout_s: int) -> None: + """Poll for result.json in a daemon thread. + + - Checks every 5 seconds for ``result.json`` in the task directory. + - If found, calls ``task_manager.complete_task()`` (which auto-closes session). + - Tracks last activity from ``status.jsonl`` mtime. + - Timeout: if wall clock exceeds *timeout_s* AND no status update + in the last 5 minutes, writes a timeout result and completes. + - On completion, closes the PTY if hooks are wired. + """ + tdir = task_manager._task_dir(session_id, task_id) + status_path = os.path.join(tdir, "status.jsonl") + start = time.time() + stale_threshold = 300 # 5 minutes + + while True: + time.sleep(5) + + # Check for result.json (may be at root or in results/ subdir) + result_path = task_manager._find_result_json(tdir) + if result_path: + try: + task_manager.complete_task(session_id, task_id) + _schedule_deferred_close(session_id) + logger.info("Watcher: task %s completed (result found)", task_id) + except Exception: + logger.exception("Watcher: error completing task %s", task_id) + return + + # Check timeout + elapsed = time.time() - start + if elapsed > timeout_s: + # Check last activity + try: + last_activity = os.path.getmtime(status_path) + except OSError: + last_activity = start + + if (time.time() - last_activity) > stale_threshold: + # Write timeout result and complete + try: + timeout_result_path = os.path.join(tdir, "result.json") + task_manager._write_json(timeout_result_path, { + "status": "timeout", + "summary": "Task timed out", + "files_changed": [], + "artifacts": [], + "errors": [f"Timeout after {timeout_s}s with no activity for 5 min"], + }) + task_manager.complete_task(session_id, task_id) + _schedule_deferred_close(session_id) + logger.warning("Watcher: task %s timed out", task_id) + except Exception: + logger.exception("Watcher: error timing out task %s", task_id) + return + + +def _close_pty_immediately(session_id: str) -> None: + """Close the PTY associated with a session, skipping the grace window. + + This bypasses the deferred-close grace period that production paths use + via ``_schedule_deferred_close``. Only use from emergency teardown or + tests; production watchers should prefer the deferred variant so users + have a window to deep-link in and view a final replay. + """ + if _app_close_session is None: + return + try: + session = task_manager._read_session(session_id) + pty_session_id = session.get("pty_session_id") + if pty_session_id: + _app_close_session(pty_session_id) + except Exception: + logger.debug("Could not close PTY for session %s", session_id, exc_info=True) + + +def _schedule_deferred_close(session_id: str) -> None: + """Mark the PTY as in-grace and schedule a delayed close. + + Both completion and timeout paths call this in place of the immediate + ``_close_pty_immediately``. The Timer is a daemon thread so it doesn't + block uvicorn shutdown. + """ + if _app_close_session is None: + return + try: + session = task_manager._read_session(session_id) + except task_manager.SessionNotFoundError: + return + pty_session_id = session.get("pty_session_id") + if not pty_session_id: + return + + if _app_mark_grace is not None: + _app_mark_grace(pty_session_id) + if _app_bump_poll is not None: + _app_bump_poll(pty_session_id, GRACE_PERIOD_S) + + t = threading.Timer(GRACE_PERIOD_S, _app_close_session, args=(pty_session_id,)) + t.daemon = True + t.start() + logger.info( + "Watcher: scheduled deferred close for pty %s in %ds", + pty_session_id, GRACE_PERIOD_S, + ) + + +# ── Tool definitions ──────────────────────────────────────────────── + + +@mcp.tool( + annotations=ToolAnnotations( + readOnlyHint=False, + destructiveHint=False, + idempotentHint=False, + ), +) +async def coda_run( + prompt: str, + email: str, + context: str = "{}", + previous_session_id: str = "", + permissions: str = "smart", + timeout_s: int = 3600, +) -> str: + """Submit a coding task — FIRE AND FORGET. + + Returns IMMEDIATELY with a task_id. The task runs autonomously in the + background. After receiving the response, tell the user the task was + submitted and move on. Do NOT follow up with coda_inbox or coda_get_result + unless the user explicitly asks to check status later. + + ``context`` is a JSON string with Unity Catalog metadata (tables, schemas). + ``previous_session_id`` chains to a prior task's session for context continuity. + ``permissions`` can be ``"smart"`` (default, safe) or ``"yolo"`` (auto-approve all). + + Returns JSON with ``task_id``, ``session_id``, and ``status: "running"``. + """ + try: + # Check concurrency limit + running = task_manager.count_running_tasks() + if running >= task_manager.MAX_CONCURRENT_TASKS: + return json.dumps({ + "status": "error", + "error": f"Concurrency limit reached ({task_manager.MAX_CONCURRENT_TASKS} " + f"tasks running). Try again when a task completes.", + }) + + # Parse context JSON + try: + ctx = json.loads(context) if context else None + except json.JSONDecodeError: + return json.dumps({ + "status": "error", + "error": f"Invalid JSON in context parameter: {context!r}", + }) + + # Auto-create ephemeral session + session_result = task_manager.create_session(email, "", label="hermes-mcp") + session_id = session_result["session_id"] + + # Create task first (we need task_id to compute transcript_path). + result = task_manager.create_task( + session_id=session_id, + prompt=prompt, + email=email, + context=ctx, + timeout_s=timeout_s, + permissions=permissions, + previous_session_id=previous_session_id or None, + ) + task_id = result["task_id"] + + pty_session_id = None + if _app_create_session is not None: + transcript_path = os.path.join( + task_manager._task_dir(session_id, task_id), + "transcript.log", + ) + pty_session_id = _app_create_session( + label="hermes-mcp", + transcript_path=transcript_path, + ) + task_manager._update_session_field( + session_id, "pty_session_id", pty_session_id + ) + + # Send to PTY if hooks are wired + if _app_send_input is not None and pty_session_id is not None: + tdir = task_manager._task_dir(session_id, task_id) + prompt_path = os.path.join(tdir, "prompt.txt") + cmd = f'hermes -z "{prompt_path}"' + if permissions == "yolo": + cmd += " --yolo" + cmd += "\n" + _app_send_input(pty_session_id, cmd) + + # Start background watcher + t = threading.Thread( + target=_watch_task, + args=(session_id, task_id, timeout_s), + daemon=True, + ) + t.start() + + return json.dumps({ + "task_id": task_id, + "session_id": session_id, + "status": "running", + "viewer_url": url_builder.build_viewer_url(pty_session_id) if pty_session_id else None, + }) + + except Exception as exc: + return json.dumps({"status": "error", "error": str(exc)}) + + +@mcp.tool( + annotations=ToolAnnotations( + readOnlyHint=True, + destructiveHint=False, + idempotentHint=True, + ), +) +async def coda_inbox( + email: str = "", + status: str = "", +) -> str: + """Check status of all background tasks — your inbox. + + Call this instead of polling — it returns ALL tasks at once. + No need to track individual task_ids; the inbox shows everything + from the last 24 hours: running, completed, and failed tasks. + + By default returns all tasks. Filter by ``status`` to narrow: + ``"running"`` for in-progress only, ``"completed"`` for finished, + ``"failed"`` for errors, or ``""`` (default) for everything. + + Each task includes: ``task_id``, ``session_id``, ``status``, + ``elapsed_s``, ``prompt_summary`` (first 100 chars of what was asked), + ``previous_session_id`` (if chained from prior work). + Completed tasks also include ``summary`` (what was done). + Running tasks also include ``progress`` (latest agent step). + + Returns JSON with ``tasks`` (list sorted most recent first) + and ``counts`` (e.g. ``{"running": 1, "completed": 2, "failed": 0}``). + """ + try: + tasks = task_manager.list_all_tasks(email=email, status_filter=status) + # Decorate each task with its viewer URL (if available). + for t in tasks: + sess = task_manager._read_session_safe(t["session_id"]) + pty = sess.get("pty_session_id") if sess else None + if pty: + vu = url_builder.build_viewer_url(pty) + if vu: + t["viewer_url"] = vu + + counts = {"running": 0, "completed": 0, "failed": 0} + for t in tasks: + s = t.get("status", "") + if s in counts: + counts[s] += 1 + elif s == "done": + counts["completed"] += 1 + elif s == "timeout": + counts["failed"] += 1 + + return json.dumps({"tasks": tasks, "counts": counts}) + except Exception as exc: + return json.dumps({"status": "error", "error": str(exc)}) + + +@mcp.tool( + annotations=ToolAnnotations( + readOnlyHint=True, + destructiveHint=False, + idempotentHint=True, + ), +) +async def coda_get_result( + task_id: str, + session_id: str, +) -> str: + """Retrieve the structured result of a completed task. + + Call this AFTER coda_inbox shows a task as "completed" or "failed". + + Returns JSON with ``task_id``, ``session_id``, ``status``, ``summary`` + (what was done), ``files_changed`` (list of modified files), + ``artifacts`` (job IDs, commit hashes, etc.), and ``errors`` (if any). + """ + try: + result = task_manager.get_task_result(task_id, session_id) + if result is None: + # No result yet — return current status + status = task_manager.get_task_status(task_id, session_id) + return json.dumps({ + "task_id": task_id, + "session_id": session_id, + "status": status.get("status", "unknown"), + "message": "Result not yet available — task is still in progress.", + }) + + result["task_id"] = task_id + result["session_id"] = session_id + # Ensure standard fields exist + result.setdefault("status", "done") + result.setdefault("summary", "") + result.setdefault("files_changed", []) + result.setdefault("artifacts", []) + result.setdefault("errors", []) + # Decorate with viewer_url if known + sess = task_manager._read_session_safe(session_id) + pty = sess.get("pty_session_id") if sess else None + if pty: + vu = url_builder.build_viewer_url(pty) + if vu: + result["viewer_url"] = vu + return json.dumps(result) + except Exception as exc: + return json.dumps({"status": "error", "task_id": task_id, "error": str(exc)}) + + +# ── Standalone entry point ────────────────────────────────────────── + +if __name__ == "__main__": + mcp.run() diff --git a/coda_mcp/task_manager.py b/coda_mcp/task_manager.py new file mode 100644 index 0000000..bad6f28 --- /dev/null +++ b/coda_mcp/task_manager.py @@ -0,0 +1,611 @@ +"""Disk-based state manager for MCP sessions and tasks. + +Pure Python module — no Flask dependency. Just file I/O. + +Layout on disk +-------------- +~/.coda/sessions/{session-id}/ + session.json – session metadata + tasks/{task-id}/ + prompt.txt – wrapped prompt sent to the agent + meta.json – task metadata (email, timestamps, chaining) + status.jsonl – append-only progress log + result.json – final output (written by the agent) +""" + +import json +import os +import secrets +import time +import logging + +logger = logging.getLogger(__name__) + +# ── Root directory (patched in tests) ──────────────────────────────── + +SESSIONS_DIR = os.path.join( + os.environ.get("HOME", "/app/python/source_code"), ".coda", "sessions" +) + +# ── Concurrency limit ─────────────────────────────────────────────── + +MAX_CONCURRENT_TASKS = int(os.environ.get("CODA_MAX_CONCURRENT", "5")) + +# ── Task TTL (seconds) ────────────────────────────────────────────── + +TASK_TTL_S = int(os.environ.get("CODA_TASK_TTL", str(24 * 3600))) # 24h + +# ── PTY → task-dir reverse lookup (used by attach_session replay fallback) ── + +_pty_lookup_cache: dict[str, tuple[str, float]] = {} # pty_id -> (task_dir, ts) +_PTY_LOOKUP_TTL = 60.0 # seconds + +# ── Exceptions ─────────────────────────────────────────────────────── + + +class SessionBusyError(Exception): + """Raised when a task is submitted to a session that already has one running.""" + + +class SessionNotFoundError(Exception): + """Raised when the requested session does not exist or is closed.""" + + +class ConcurrencyLimitError(Exception): + """Raised when MAX_CONCURRENT_TASKS running tasks already exist.""" + + +# ── ID generators ──────────────────────────────────────────────────── + + +def _new_session_id() -> str: + return f"sess-{secrets.token_hex(6)}" + + +def _new_task_id() -> str: + return f"task-{secrets.token_hex(4)}" + + +# ── Low-level I/O ──────────────────────────────────────────────────── + + +def _session_dir(session_id: str) -> str: + return os.path.join(SESSIONS_DIR, session_id) + + +def _session_file(session_id: str) -> str: + return os.path.join(_session_dir(session_id), "session.json") + + +def _task_dir(session_id: str, task_id: str) -> str: + """Return the path to a task's directory.""" + return os.path.join(_session_dir(session_id), "tasks", task_id) + + +def _write_json(path: str, data: dict) -> None: + """Atomic write via tmp-then-rename.""" + os.makedirs(os.path.dirname(path), exist_ok=True) + tmp = path + ".tmp" + with open(tmp, "w") as f: + json.dump(data, f, indent=2) + os.replace(tmp, path) + + +def _read_session(session_id: str) -> dict: + """Read session.json or raise SessionNotFoundError.""" + path = _session_file(session_id) + try: + with open(path) as f: + return json.load(f) + except (OSError, json.JSONDecodeError): + raise SessionNotFoundError(f"Session {session_id} not found or corrupt") + + +def _read_session_safe(session_id: str) -> dict | None: + """Read session.json, returning None on missing/corrupt instead of raising.""" + try: + return _read_session(session_id) + except SessionNotFoundError: + return None + + +def _update_session_field(session_id: str, key: str, value) -> None: + """Update a single field in session.json (read-modify-write).""" + data = _read_session(session_id) + data[key] = value + _write_json(_session_file(session_id), data) + + +# ── Session lifecycle ──────────────────────────────────────────────── + + +def create_session(email: str, user_id: str, label: str = "") -> dict: + """Create a new session directory with session.json. + + Returns ``{"session_id": "sess-…", "status": "ready"}``. + """ + session_id = _new_session_id() + data = { + "session_id": session_id, + "email": email, + "user_id": user_id, + "label": label, + "status": "ready", + "current_task": None, + "completed_tasks": [], + "created_at": time.time(), + } + _write_json(_session_file(session_id), data) + logger.info("Created session %s for %s", session_id, email) + return {"session_id": session_id, "status": "ready"} + + +def close_session(session_id: str) -> None: + """Mark a session as closed. Raises SessionNotFoundError if missing.""" + _read_session(session_id) # existence check + _update_session_field(session_id, "status", "closed") + logger.info("Closed session %s", session_id) + + +# ── Prompt wrapping ────────────────────────────────────────────────── + + +def wrap_prompt( + task_id: str, + session_id: str, + email: str, + prompt: str, + context: dict | None, + results_dir: str, + context_hint: str | None = None, + previous_session_id: str | None = None, +) -> str: + """Build the full prompt string written to ``prompt.txt``. + + Uses the ``---CODA-TASK---`` envelope convention so the agent can + parse metadata from the prompt deterministically. + """ + context_block = "" + if context: + context_block = f"\nCONTEXT:\n{json.dumps(context, indent=2)}\n" + + hint_line = "" + if context_hint: + hint_line = f"context_hint: {context_hint}\n" + + prior_session_block = "" + if previous_session_id: + prior_dir = _session_dir(previous_session_id) + prior_session_block = ( + f"\nPRIOR SESSION: {previous_session_id}\n" + f"Read {prior_dir}/tasks/*/result.json for context on prior work.\n" + ) + + return ( + f"---CODA-TASK---\n" + f"task_id: {task_id}\n" + f"session_id: {session_id}\n" + f"user: {email}\n" + f"{hint_line}" + f"{prior_session_block}" + f"{context_block}\n" + f"TASK:\n" + f"{prompt}\n" + f"\n" + f"INSTRUCTIONS:\n" + f"1. As you work, append progress lines to {results_dir}/status.jsonl\n" + f' Each line must be valid JSON: {{"step": "label", "message": "what you are doing"}}\n' + f"\n" + f"2. When you are COMPLETELY DONE, write a SINGLE FILE at this exact path:\n" + f" {results_dir}/result.json\n" + f" It must contain this JSON structure:\n" + f" {{\n" + f' "status": "completed",\n' + f' "summary": "one paragraph describing what you did",\n' + f' "files_changed": ["list", "of", "file", "paths"],\n' + f' "artifacts": {{}},\n' + f' "errors": []\n' + f" }}\n" + f" If you failed, set status to \"failed\" and describe the error.\n" + f" IMPORTANT: result.json is a FILE not a directory. Write it with:\n" + f" echo '{{...}}' > {results_dir}/result.json\n" + f"\n" + f"3. If you delegate to a sub-agent, update status.jsonl with delegation steps.\n" + f"\n" + f"SAFETY:\n" + f"- Do NOT delete, drop, or truncate tables, schemas, catalogs, or volumes.\n" + f"- Do NOT delete files outside the current project directory.\n" + f"- Do NOT run destructive Databricks CLI commands (e.g. databricks clusters delete, " + f"databricks jobs delete, databricks pipelines delete).\n" + f"- Do NOT modify permissions, grants, or access controls unless explicitly requested.\n" + f"- Prefer CREATE OR REPLACE over DROP+CREATE. Prefer INSERT/MERGE over DELETE+INSERT.\n" + f"- If the task requires a destructive operation, describe what you would do in " + f"result.json with status \"needs_approval\" instead of executing it.\n" + f"---END-CODA-TASK---" + ) + + +# ── Task lifecycle ─────────────────────────────────────────────────── + + +def create_task( + session_id: str, + prompt: str, + email: str, + context: dict | None = None, + context_hint: str | None = None, + timeout_s: int | None = None, + permissions: str | None = None, + previous_session_id: str | None = None, +) -> dict: + """Create a task inside an existing session. + + Raises + ------ + SessionNotFoundError + If the session does not exist or is closed. + SessionBusyError + If the session already has a running task. + + Returns ``{"task_id": "task-…", "status": "running"}``. + """ + session = _read_session(session_id) + + if session.get("status") == "closed": + raise SessionNotFoundError(f"Session {session_id} is closed") + + if session.get("status") == "busy": + raise SessionBusyError( + f"Session {session_id} already has a running task: " + f"{session.get('current_task')}" + ) + + task_id = _new_task_id() + tdir = _task_dir(session_id, task_id) + os.makedirs(tdir, exist_ok=True) + + # Write wrapped prompt + results_dir = os.path.join(tdir, "results") + wrapped = wrap_prompt( + task_id=task_id, + session_id=session_id, + email=email, + prompt=prompt, + context=context, + results_dir=results_dir, + context_hint=context_hint, + previous_session_id=previous_session_id, + ) + with open(os.path.join(tdir, "prompt.txt"), "w") as f: + f.write(wrapped) + + # Write meta.json for inbox scanning + now = time.time() + meta = { + "email": email, + "created_at": now, + "previous_session_id": previous_session_id or "", + "permissions": permissions or "smart", + "timeout_s": timeout_s or 3600, + "prompt_summary": prompt[:100], + } + _write_json(os.path.join(tdir, "meta.json"), meta) + + # Seed status log + with open(os.path.join(tdir, "status.jsonl"), "w") as f: + f.write(json.dumps({"status": "running", "ts": now}) + "\n") + + # Mark session busy + data = _read_session(session_id) + data["status"] = "busy" + data["current_task"] = task_id + _write_json(_session_file(session_id), data) + + logger.info("Created task %s in session %s", task_id, session_id) + return {"task_id": task_id, "status": "running"} + + +# ── Task queries ───────────────────────────────────────────────────── + + +def get_task_status(task_id: str, session_id: str) -> dict: + """Read the last line of status.jsonl for the task. + + Returns ``{"status": "not_found"}`` if the task directory is missing. + """ + status_path = os.path.join(_task_dir(session_id, task_id), "status.jsonl") + try: + last = None + with open(status_path) as f: + for line in f: + line = line.strip() + if line: + last = json.loads(line) + return last or {"status": "not_found"} + except (OSError, json.JSONDecodeError): + return {"status": "not_found"} + + +def _find_result_json(task_dir: str) -> str | None: + """Find result.json — agents may write it at root or in results/ subdir.""" + for candidate in [ + os.path.join(task_dir, "result.json"), + os.path.join(task_dir, "results", "result.json"), + ]: + if os.path.isfile(candidate): + return candidate + return None + + +def get_task_result(task_id: str, session_id: str) -> dict | None: + """Read result.json if it exists; otherwise return None.""" + result_path = _find_result_json(_task_dir(session_id, task_id)) + if not result_path: + return None + try: + with open(result_path) as f: + return json.load(f) + except (OSError, json.JSONDecodeError): + return None + + +# ── Task completion ────────────────────────────────────────────────── + + +def complete_task(session_id: str, task_id: str) -> None: + """Mark a task as done and auto-close the session. + + Appends a ``done`` entry to status.jsonl, adds task_id to + ``completed_tasks``, and closes the session (v2: ephemeral sessions). + """ + session = _read_session(session_id) + + # Append done to status log + status_path = os.path.join(_task_dir(session_id, task_id), "status.jsonl") + with open(status_path, "a") as f: + f.write(json.dumps({"status": "done", "ts": time.time()}) + "\n") + + # Update session — auto-close (v2: sessions are ephemeral) + session["status"] = "closed" + session["current_task"] = None + session["closed_at"] = time.time() + if task_id not in session["completed_tasks"]: + session["completed_tasks"].append(task_id) + _write_json(_session_file(session_id), session) + + logger.info("Completed task %s in session %s (auto-closed)", task_id, session_id) + + +# ── Inbox: list all tasks across sessions ─────────────────────────── + + +def list_all_tasks(email: str = "", status_filter: str = "") -> list[dict]: + """Scan all sessions and return a flat list of tasks for the inbox. + + Returns tasks from the last ``TASK_TTL_S`` seconds, sorted most recent first. + Each entry includes task_id, session_id, status, elapsed_s, prompt_summary, + summary (if completed), progress (if running), previous_session_id, created_at. + """ + now = time.time() + cutoff = now - TASK_TTL_S + tasks = [] + + if not os.path.isdir(SESSIONS_DIR): + return tasks + + for sess_name in os.listdir(SESSIONS_DIR): + sess_dir = os.path.join(SESSIONS_DIR, sess_name) + if not os.path.isdir(sess_dir): + continue + + tasks_dir = os.path.join(sess_dir, "tasks") + if not os.path.isdir(tasks_dir): + continue + + for task_name in os.listdir(tasks_dir): + task_dir = os.path.join(tasks_dir, task_name) + if not os.path.isdir(task_dir): + continue + + # Read meta.json + meta_path = os.path.join(task_dir, "meta.json") + try: + with open(meta_path) as f: + meta = json.load(f) + except (OSError, json.JSONDecodeError): + # Legacy task without meta.json — skip or build minimal entry + meta = {} + + created_at = meta.get("created_at", 0) + if created_at < cutoff: + continue + + # Filter by email + if email and meta.get("email", "") != email: + continue + + # Determine task status from status.jsonl + task_status = _read_last_status(task_dir) + + # Check for result.json to determine completion + result_path = _find_result_json(task_dir) + summary = "" + if result_path: + try: + with open(result_path) as f: + result_data = json.load(f) + task_status = result_data.get("status", "completed") + summary = result_data.get("summary", "") + except (OSError, json.JSONDecodeError): + pass + + # Filter by status + if status_filter and task_status != status_filter: + continue + + # Get progress for running tasks + progress = "" + if task_status == "running": + progress = _read_last_progress(task_dir) + + elapsed_s = round(now - created_at, 1) + + entry = { + "task_id": task_name, + "session_id": sess_name, + "status": task_status, + "elapsed_s": elapsed_s, + "prompt_summary": meta.get("prompt_summary", ""), + "previous_session_id": meta.get("previous_session_id", ""), + "created_at": created_at, + } + if summary: + entry["summary"] = summary + if progress: + entry["progress"] = progress + + tasks.append(entry) + + # Sort most recent first + tasks.sort(key=lambda t: t["created_at"], reverse=True) + return tasks + + +def _read_last_status(task_dir: str) -> str: + """Read the last status from status.jsonl.""" + status_path = os.path.join(task_dir, "status.jsonl") + try: + last = None + with open(status_path) as f: + for line in f: + line = line.strip() + if line: + last = json.loads(line) + return (last or {}).get("status", "unknown") + except (OSError, json.JSONDecodeError): + return "unknown" + + +def _read_last_progress(task_dir: str) -> str: + """Read the last progress message from status.jsonl.""" + status_path = os.path.join(task_dir, "status.jsonl") + try: + last = None + with open(status_path) as f: + for line in f: + line = line.strip() + if line: + last = json.loads(line) + return (last or {}).get("message", "") + except (OSError, json.JSONDecodeError): + return "" + + +# ── Concurrency check ────────────────────────────────────────────── + + +def count_running_tasks() -> int: + """Count tasks currently in 'running' state across all sessions.""" + count = 0 + if not os.path.isdir(SESSIONS_DIR): + return count + + for sess_name in os.listdir(SESSIONS_DIR): + sess_file = os.path.join(SESSIONS_DIR, sess_name, "session.json") + try: + with open(sess_file) as f: + session = json.load(f) + if session.get("status") == "busy": + count += 1 + except (OSError, json.JSONDecodeError): + continue + return count + + +# ── PTY → task-dir reverse lookup ────────────────────────────────── + + +def find_task_dir_by_pty_session(pty_session_id: str) -> str | None: + """Find the task dir whose session.json carries this pty_session_id. + + Returns the path to the active task dir, or — if the session has completed — + the most recently completed task dir. Returns None on no match. + + Cached for ``_PTY_LOOKUP_TTL`` seconds to avoid disk scans on every browser + refresh. + + Invariant: CoDA MCP sessions are ephemeral — one task per session. If the + lifecycle ever changes to allow multiple tasks per session, this function + must be revisited to pick the active or grace-period task rather than + ``completed_tasks[-1]``. + """ + now = time.time() + cached = _pty_lookup_cache.get(pty_session_id) + if cached and (now - cached[1]) < _PTY_LOOKUP_TTL: + return cached[0] + + if not os.path.isdir(SESSIONS_DIR): + return None + + for sess_name in os.listdir(SESSIONS_DIR): + sess_file = os.path.join(SESSIONS_DIR, sess_name, "session.json") + try: + with open(sess_file) as f: + data = json.load(f) + except (OSError, json.JSONDecodeError): + continue + + if data.get("pty_session_id") != pty_session_id: + continue + + candidate = data.get("current_task") or ( + data["completed_tasks"][-1] if data.get("completed_tasks") else None + ) + if candidate: + tdir = os.path.join(SESSIONS_DIR, sess_name, "tasks", candidate) + _pty_lookup_cache[pty_session_id] = (tdir, now) + return tdir + + return None + + +# ── Cleanup expired sessions ──────────────────────────────────────── + + +def cleanup_expired_tasks() -> int: + """Remove session directories older than TASK_TTL_S. Returns count removed.""" + import shutil + + now = time.time() + cutoff = now - TASK_TTL_S + removed = 0 + + if not os.path.isdir(SESSIONS_DIR): + return removed + + for sess_name in os.listdir(SESSIONS_DIR): + sess_dir = os.path.join(SESSIONS_DIR, sess_name) + if not os.path.isdir(sess_dir): + continue + + sess_file = os.path.join(sess_dir, "session.json") + try: + with open(sess_file) as f: + session = json.load(f) + except (OSError, json.JSONDecodeError): + continue + + # Only clean closed sessions past TTL + if session.get("status") != "closed": + continue + + closed_at = session.get("closed_at", session.get("created_at", 0)) + if closed_at < cutoff: + try: + shutil.rmtree(sess_dir) + removed += 1 + logger.info("Cleaned up expired session %s", sess_name) + except OSError: + logger.warning("Failed to clean up session %s", sess_name) + + return removed diff --git a/coda_mcp/url_builder.py b/coda_mcp/url_builder.py new file mode 100644 index 0000000..c08d2ed --- /dev/null +++ b/coda_mcp/url_builder.py @@ -0,0 +1,46 @@ +"""Builds the viewer_url returned by CoDA MCP tools. + +Resolution order: +1. ``CODA_APP_URL`` env var (explicit override for local dev / power users). +2. Module-level cache populated by ``AppUrlCaptureMiddleware`` from the + ``X-Forwarded-Host`` header (officially provided by Databricks Apps). +3. ``None`` — caller omits the field entirely. + +The cache is process-global (single uvicorn worker per app) and refreshed +on every inbound HTTP request. +""" +from __future__ import annotations + +import os +from typing import Optional + +_app_url_cache: Optional[str] = None + + +def capture_from_headers(host: Optional[str]) -> None: + """Called by the ASGI middleware on every inbound HTTP request. + + No-op when ``host`` is falsy (None or empty) to avoid wiping a good + cache value with a missing header on a probe/CORS preflight. + + Strips any accidental ``https://`` / ``http://`` prefix on the way in + so build_viewer_url's unconditional ``https://`` prepend can't produce + a double-scheme URL. + """ + global _app_url_cache + if host: + host = host.removeprefix("https://").removeprefix("http://").strip("/") + if host: + _app_url_cache = host + + +def build_viewer_url(pty_session_id: str) -> Optional[str]: + """Return the full viewer URL for a PTY session, or None if no base is known.""" + override = os.environ.get("CODA_APP_URL", "").strip() + if override: + base = override.rstrip("/") + elif _app_url_cache: + base = f"https://{_app_url_cache}" + else: + return None + return f"{base}/?session={pty_session_id}" diff --git a/docs/mcp-client-setup.md b/docs/mcp-client-setup.md new file mode 100644 index 0000000..f8e1bb6 --- /dev/null +++ b/docs/mcp-client-setup.md @@ -0,0 +1,73 @@ +# CoDA MCP Client Setup + +CoDA exposes an MCP endpoint at `/mcp` on the Databricks App. Databricks Apps use OAuth (not PATs) for authentication, so MCP clients need a stdio bridge that injects fresh OAuth tokens. + +## How it works + +`tools/coda-bridge.py` is a zero-dependency Python script that: + +1. Claude Code launches it as a stdio MCP server +2. It reads JSON-RPC messages from stdin +3. Fetches a fresh OAuth token via `databricks auth token` +4. Forwards requests to the App's HTTP endpoint with the token +5. Returns responses on stdout + +Tokens are cached for 30 minutes (they expire after 60). + +## Setup + +### 1. Copy the bridge script + +```bash +mkdir -p ~/.claude/mcp-bridges +cp tools/coda-bridge.py ~/.claude/mcp-bridges/ +``` + +### 2. Add to Claude Code settings + +Add this to `mcpServers` in `~/.claude/settings.json`: + +```json +"coda-mcp": { + "type": "stdio", + "command": "python3", + "args": ["/path/to/.claude/mcp-bridges/coda-bridge.py"], + "env": { + "CODA_MCP_URL": "https://.databricksapps.com/mcp", + "DATABRICKS_PROFILE": "" + } +} +``` + +### 3. Restart Claude Code + +The MCP server will start automatically on next session. + +## Configuration + +| Environment Variable | Description | Example | +|---------------------|-------------|---------| +| `CODA_MCP_URL` | Full URL to the app's `/mcp` endpoint | `https://mcp-test-coda-747...com/mcp` | +| `DATABRICKS_PROFILE` | Databricks CLI profile name | `9cefok` | + +## Prerequisites + +- `databricks` CLI installed and authenticated (`databricks auth login -p `) +- Python 3.8+ +- No pip dependencies required (stdlib only) + +## Troubleshooting + +Bridge logs go to stderr. Check with: + +```bash +CODA_MCP_URL="https://your-app.databricksapps.com/mcp" \ +DATABRICKS_PROFILE="your-profile" \ +echo '{"jsonrpc":"2.0","method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}},"id":1}' | python3 tools/coda-bridge.py +``` + +If you see `Auth failed (302)`, your Databricks CLI session may have expired. Run: + +```bash +databricks auth login -p +``` diff --git a/docs/mcp-v2-background-execution.md b/docs/mcp-v2-background-execution.md new file mode 100644 index 0000000..3d7557c --- /dev/null +++ b/docs/mcp-v2-background-execution.md @@ -0,0 +1,171 @@ +# CoDA MCP v2 — Background Execution + Inbox Pattern + +## Overview + +CoDA exposes 3 MCP tools so Databricks GenieCode (or any MCP client) can delegate +coding tasks to AI agents running in the background. GenieCode's chat context stays +free while tasks execute — no polling required. + +## Tools + +| Tool | Purpose | +|------|---------| +| `coda_run` | Fire-and-forget task submission | +| `coda_inbox` | Dashboard of all background tasks | +| `coda_get_result` | Pull full structured result | + +## Flow Diagram + +``` +┌─────────────┐ ┌──────────────┐ ┌─────────────┐ +│ GenieCode │ │ CoDA MCP │ │ Hermes │ +│ (caller) │ │ (3 tools) │ │ (executor) │ +└──────┬──────┘ └──────┬───────┘ └──────┬──────┘ + │ │ │ + │ 1. coda_run(prompt) │ │ + │──────────────────────>│ │ + │ │ auto-create session │ + │ │ + PTY + task dir │ + │ │ write prompt.txt │ + │ │ write meta.json │ + │ │ │ + │ {task_id, sess_id, │ hermes -z prompt.txt │ + │ status: "running"} │───────────────────────>│ + │<──────────────────────│ │ + │ │ _watch_task thread │ + │ ✓ context is FREE │ monitors result.json │ + │ user keeps chatting │ │ + │ │ │ works... + │ ... │ │ delegates + │ │ │ to claude/ + │ │ │ codex/gemini + │ │ │ + │ 2. coda_inbox() │ │ writes + │──────────────────────>│ │ status.jsonl + │ │ scan all sessions │ + │ {tasks: [...], │ read meta + status │ + │ counts: {run:1}} │ │ + │<──────────────────────│ │ + │ │ │ + │ ... │ │ writes + │ │ │ result.json + │ │ │ + │ │ _watch_task detects │ + │ │ result.json exists │ + │ │ → complete_task() │ + │ │ → auto-close session │ + │ │ → free PTY │ + │ │ │ + │ 3. coda_inbox() │ │ + │──────────────────────>│ │ + │ {tasks: [{status: │ │ + │ "completed", │ │ + │ summary: "..."}]} │ │ + │<──────────────────────│ │ + │ │ │ + │ 4. coda_get_result() │ │ + │──────────────────────>│ │ + │ {summary, files, │ read result.json │ + │ artifacts, errors} │ │ + │<──────────────────────│ │ + │ │ │ + ├── CHAINING ───────────┤ │ + │ │ │ + │ 5. coda_run(prompt, │ │ + │ previous_session_id) │ new session + PTY │ + │──────────────────────>│ inject PRIOR SESSION │ + │ │ block in prompt │ + │ {new task_id, │───────────────────────>│ + │ new sess_id} │ │ reads prior + │<──────────────────────│ │ result.json + │ │ │ for context +``` + +## Key Design Decisions + +### Sessions are ephemeral, tasks are persistent +- Session = PTY + Hermes instance. Auto-closes when task completes. +- Task state (prompt, status, result) persists on disk for 24 hours. +- Continuity via `previous_session_id`, not long-lived sessions. + +### No polling from GenieCode +- `coda_inbox` replaces `coda_get_status` — shows ALL tasks at once. +- GenieCode checks when the user asks, not on a timer. +- CoDA's internal `_watch_task` thread polls the filesystem (invisible to caller). + +### Task chaining +- `previous_session_id` points to a prior session's disk state. +- Hermes reads `~/.coda/sessions/{prev_id}/tasks/*/result.json` for context. +- Chain depth: one level. Hermes can walk deeper if needed. + +### Concurrency +- `CODA_MAX_CONCURRENT` env var (default: 5). +- Each task gets its own session — no "session busy" errors. +- Exceeding the limit returns a clear error. + +## Data Model + +``` +~/.coda/sessions/{session-id}/ + session.json # metadata + auto-close timestamp + tasks/{task-id}/ + prompt.txt # wrapped prompt sent to Hermes + meta.json # {email, created_at, previous_session_id, permissions} + status.jsonl # append-only progress log + result.json # final structured output +``` + +## Tool Reference + +### `coda_run` + +```python +coda_run( + prompt: str, # what to do + email: str, # who's asking + context: str = "{}", # UC metadata (tables, schemas) + previous_session_id: str = "", # chain from prior work + permissions: str = "smart", # "smart" or "yolo" + timeout_s: int = 3600, # max 1 hour default +) +# Returns: {"task_id", "session_id", "status": "running"} +``` + +### `coda_inbox` + +```python +coda_inbox( + email: str = "", # filter by user + status: str = "", # "running", "completed", "failed", or "" for all +) +# Returns: {"tasks": [...], "counts": {"running": N, "completed": N, "failed": N}} +``` + +Each task entry: `task_id`, `session_id`, `status`, `elapsed_s`, `prompt_summary`, +`summary` (completed), `progress` (running), `previous_session_id`, `created_at`. + +### `coda_get_result` + +```python +coda_get_result(task_id: str, session_id: str) +# Returns: {"task_id", "session_id", "status", "summary", +# "files_changed", "artifacts", "errors"} +``` + +## Migration from v1 + +| v1 Tool | v2 Equivalent | +|---------|--------------| +| `coda_create_session` | Removed — auto-created by `coda_run` | +| `coda_run_task` | `coda_run` (simplified, auto-session) | +| `coda_get_status` | `coda_inbox` (all tasks at once) | +| `coda_get_result` | `coda_get_result` (unchanged) | +| `coda_close_session` | Removed — auto-closed on completion | + +## Limitations + +- **Ephemeral filesystem**: On Databricks Apps, `~/.coda/` is local disk. App + redeployment wipes task state. Real artifacts (git commits, jobs, workspace files) + are unaffected. +- **No push notifications**: GenieCode must call `coda_inbox` to discover completions. + SSE/streaming is a future consideration if polling proves insufficient. diff --git a/docs/plans/2026-05-01-coda-mcp-server.md b/docs/plans/2026-05-01-coda-mcp-server.md new file mode 100644 index 0000000..1e59ba3 --- /dev/null +++ b/docs/plans/2026-05-01-coda-mcp-server.md @@ -0,0 +1,1179 @@ +# CoDA MCP Server Implementation Plan + +> **⚠️ SUPERSEDED — historical reference only.** This was the v1 implementation plan (5 tools, gunicorn + WSGI bridge). The shipped implementation diverged during iteration: the production design is documented in [`docs/mcp-v2-background-execution.md`](../mcp-v2-background-execution.md) (3 tools — `coda_run`, `coda_inbox`, `coda_get_result` — on uvicorn + native ASGI). Kept in the tree so reviewers can see the design evolution; do not follow this plan as-is. + +> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task. + +**Goal:** Add an MCP server endpoint (`/mcp`) to CoDA so Databricks Genie Code can delegate coding tasks to Hermes Agent via the MCP protocol. + +**Architecture:** Python MCP SDK mounted as a stateless HTTP app at `/mcp` alongside the existing Flask app. A new `task_manager.py` module handles session/task state on disk (`~/.coda/sessions/`). The MCP tools call into the existing PTY infrastructure for session creation and input piping. Hermes is always the agent invoked. + +**Tech Stack:** Python MCP SDK (`mcp` package, already installed), Flask, existing PTY session infrastructure, Hermes Agent CLI (`hermes -z`) + +**Design doc:** `.humantokens/coda-mcp-design.md` (full design with all decisions) + +--- + +### Task 1: Create Task Manager Module + +The task manager handles all disk-based state for MCP sessions and tasks. It's a pure Python module with no Flask dependency — just file I/O. + +**Files:** +- Create: `task_manager.py` +- Create: `tests/test_task_manager.py` + +**Step 1: Write the failing tests** + +```python +# tests/test_task_manager.py +import os +import json +import tempfile +import pytest +from unittest.mock import patch + +# All tests use a temp dir instead of ~/.coda +@pytest.fixture +def task_mgr(tmp_path): + with patch("task_manager.SESSIONS_DIR", str(tmp_path / "sessions")): + import task_manager + # Force reimport to pick up patched path + task_manager.SESSIONS_DIR = str(tmp_path / "sessions") + yield task_manager + + +def test_create_session(task_mgr): + result = task_mgr.create_session(email="alice@example.com", user_id="123") + assert "session_id" in result + assert result["status"] == "ready" + + # Verify session.json on disk + session_dir = os.path.join(task_mgr.SESSIONS_DIR, result["session_id"]) + assert os.path.isdir(session_dir) + with open(os.path.join(session_dir, "session.json")) as f: + data = json.load(f) + assert data["created_by"] == "alice@example.com" + assert data["status"] == "idle" + assert data["current_task"] is None + + +def test_create_task(task_mgr): + session = task_mgr.create_session(email="alice@example.com", user_id="123") + sid = session["session_id"] + + result = task_mgr.create_task( + session_id=sid, + prompt="create a pipeline", + email="alice@example.com", + context={"tables": ["sales.transactions"]}, + ) + assert "task_id" in result + assert result["status"] == "running" + + # Verify task dir and files + task_dir = os.path.join(task_mgr.SESSIONS_DIR, sid, "tasks", result["task_id"]) + assert os.path.isfile(os.path.join(task_dir, "prompt.txt")) + + # Session should be busy + with open(os.path.join(task_mgr.SESSIONS_DIR, sid, "session.json")) as f: + data = json.load(f) + assert data["status"] == "busy" + assert data["current_task"] == result["task_id"] + + +def test_create_task_rejects_when_busy(task_mgr): + session = task_mgr.create_session(email="alice@example.com", user_id="123") + sid = session["session_id"] + + task_mgr.create_task(session_id=sid, prompt="task 1", email="alice@example.com") + with pytest.raises(task_mgr.SessionBusyError): + task_mgr.create_task(session_id=sid, prompt="task 2", email="alice@example.com") + + +def test_get_status_running(task_mgr): + session = task_mgr.create_session(email="alice@example.com", user_id="123") + sid = session["session_id"] + task = task_mgr.create_task(session_id=sid, prompt="do work", email="alice@example.com") + + status = task_mgr.get_task_status(task["task_id"], sid) + assert status["status"] == "running" + assert "elapsed_s" in status + assert status.get("progress") is None # no status.jsonl yet + + +def test_get_status_with_progress(task_mgr): + session = task_mgr.create_session(email="alice@example.com", user_id="123") + sid = session["session_id"] + task = task_mgr.create_task(session_id=sid, prompt="do work", email="alice@example.com") + tid = task["task_id"] + + # Simulate agent writing status.jsonl + status_file = os.path.join(task_mgr.SESSIONS_DIR, sid, "tasks", tid, "status.jsonl") + with open(status_file, "a") as f: + f.write(json.dumps({"step": "planning", "message": "Analyzing requirements"}) + "\n") + f.write(json.dumps({"step": "coding", "message": "Writing pipeline"}) + "\n") + + status = task_mgr.get_task_status(tid, sid) + assert status["status"] == "running" + assert status["progress"]["step"] == "coding" + + +def test_get_result_completed(task_mgr): + session = task_mgr.create_session(email="alice@example.com", user_id="123") + sid = session["session_id"] + task = task_mgr.create_task(session_id=sid, prompt="do work", email="alice@example.com") + tid = task["task_id"] + + # Simulate agent writing result.json + result_file = os.path.join(task_mgr.SESSIONS_DIR, sid, "tasks", tid, "result.json") + with open(result_file, "w") as f: + json.dump({ + "status": "completed", + "summary": "Created pipeline", + "files_changed": ["pipeline.py"], + "artifacts": {"job_id": "123"}, + "errors": [] + }, f) + + result = task_mgr.get_task_result(tid, sid) + assert result["status"] == "completed" + assert result["summary"] == "Created pipeline" + + +def test_get_result_not_done(task_mgr): + session = task_mgr.create_session(email="alice@example.com", user_id="123") + sid = session["session_id"] + task = task_mgr.create_task(session_id=sid, prompt="do work", email="alice@example.com") + + result = task_mgr.get_task_result(task["task_id"], sid) + assert result["status"] == "running" + assert result.get("summary") is None + + +def test_complete_task(task_mgr): + session = task_mgr.create_session(email="alice@example.com", user_id="123") + sid = session["session_id"] + task = task_mgr.create_task(session_id=sid, prompt="do work", email="alice@example.com") + tid = task["task_id"] + + # Simulate result.json written by agent + result_file = os.path.join(task_mgr.SESSIONS_DIR, sid, "tasks", tid, "result.json") + with open(result_file, "w") as f: + json.dump({"status": "completed", "summary": "Done", "files_changed": [], "artifacts": {}, "errors": []}, f) + + task_mgr.complete_task(sid, tid) + + # Session should be idle again + with open(os.path.join(task_mgr.SESSIONS_DIR, sid, "session.json")) as f: + data = json.load(f) + assert data["status"] == "idle" + assert data["current_task"] is None + assert tid in data["completed_tasks"] + + +def test_close_session(task_mgr): + session = task_mgr.create_session(email="alice@example.com", user_id="123") + sid = session["session_id"] + + result = task_mgr.close_session(sid) + assert result["status"] == "closed" + + with open(os.path.join(task_mgr.SESSIONS_DIR, sid, "session.json")) as f: + data = json.load(f) + assert data["status"] == "closed" + + +def test_wrap_prompt(task_mgr): + wrapped = task_mgr.wrap_prompt( + task_id="task-007", + session_id="sess-abc", + email="alice@example.com", + prompt="create a pipeline", + context={"tables": ["sales.transactions"]}, + results_dir="/tmp/test" + ) + assert "---CODA-TASK---" in wrapped + assert "task-007" in wrapped + assert "create a pipeline" in wrapped + assert "sales.transactions" in wrapped + assert "result.json" in wrapped + assert "---END-CODA-TASK---" in wrapped +``` + +**Step 2: Run tests to verify they fail** + +Run: `cd /Users/sathish.gangichetty/Documents/xterm-experiment/.worktrees/coda-mcp && uv run pytest tests/test_task_manager.py -v` +Expected: FAIL — `ModuleNotFoundError: No module named 'task_manager'` + +**Step 3: Write the task_manager module** + +```python +# task_manager.py +"""Disk-based state manager for MCP sessions and tasks. + +Manages the lifecycle of sessions (PTY-backed Hermes instances) and tasks +(units of work within a session). All state is persisted to ~/.coda/sessions/ +so the MCP transport can remain stateless. +""" +import json +import os +import time +import uuid + +HOME = os.environ.get("HOME", os.path.expanduser("~")) +SESSIONS_DIR = os.path.join(HOME, ".coda", "sessions") + + +class SessionBusyError(Exception): + """Raised when a task is submitted to a session that's already running one.""" + pass + + +class SessionNotFoundError(Exception): + """Raised when a session_id doesn't exist.""" + pass + + +def _session_dir(session_id: str) -> str: + return os.path.join(SESSIONS_DIR, session_id) + + +def _task_dir(session_id: str, task_id: str) -> str: + return os.path.join(SESSIONS_DIR, session_id, "tasks", task_id) + + +def _read_session(session_id: str) -> dict: + path = os.path.join(_session_dir(session_id), "session.json") + if not os.path.isfile(path): + raise SessionNotFoundError(f"Session {session_id} not found") + with open(path) as f: + return json.load(f) + + +def _write_session(session_id: str, data: dict): + path = os.path.join(_session_dir(session_id), "session.json") + with open(path, "w") as f: + json.dump(data, f, indent=2) + + +def create_session(email: str, user_id: str = "", label: str = "") -> dict: + """Create a new session directory and session.json. Returns {session_id, status}.""" + session_id = f"sess-{uuid.uuid4().hex[:12]}" + session_dir = _session_dir(session_id) + os.makedirs(os.path.join(session_dir, "tasks"), exist_ok=True) + + session_data = { + "created_by": email, + "user_id": user_id, + "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + "status": "idle", + "current_task": None, + "completed_tasks": [], + "label": label, + } + _write_session(session_id, session_data) + + return {"session_id": session_id, "status": "ready"} + + +def create_task( + session_id: str, + prompt: str, + email: str, + context: dict = None, + context_hint: str = None, + timeout_s: int = 3600, + permissions: str = "smart", +) -> dict: + """Create a new task within a session. Returns {task_id, status}. + + Raises SessionBusyError if the session already has a running task. + """ + session_data = _read_session(session_id) + + if session_data["status"] == "busy": + raise SessionBusyError(f"Session {session_id} is busy with task {session_data['current_task']}") + + if session_data["status"] == "closed": + raise SessionNotFoundError(f"Session {session_id} is closed") + + task_id = f"task-{uuid.uuid4().hex[:8]}" + task_dir = _task_dir(session_id, task_id) + os.makedirs(task_dir, exist_ok=True) + + # Write prompt file + results_dir = task_dir + wrapped = wrap_prompt( + task_id=task_id, + session_id=session_id, + email=email, + prompt=prompt, + context=context, + results_dir=results_dir, + context_hint=context_hint, + ) + with open(os.path.join(task_dir, "prompt.txt"), "w") as f: + f.write(wrapped) + + # Write task metadata + with open(os.path.join(task_dir, "meta.json"), "w") as f: + json.dump({ + "task_id": task_id, + "session_id": session_id, + "email": email, + "started_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + "timeout_s": timeout_s, + "permissions": permissions, + "context_hint": context_hint, + }, f, indent=2) + + # Update session state + session_data["status"] = "busy" + session_data["current_task"] = task_id + _write_session(session_id, session_data) + + return {"task_id": task_id, "status": "running"} + + +def get_task_status(task_id: str, session_id: str) -> dict: + """Get current status of a task. Reads status.jsonl for progress.""" + task_dir = _task_dir(session_id, task_id) + + # Check if result.json exists (task completed) + result_path = os.path.join(task_dir, "result.json") + if os.path.isfile(result_path): + with open(result_path) as f: + result = json.load(f) + return { + "task_id": task_id, + "status": result.get("status", "completed"), + "elapsed_s": _elapsed(task_dir), + } + + # Check for progress in status.jsonl + status_path = os.path.join(task_dir, "status.jsonl") + progress = None + if os.path.isfile(status_path): + with open(status_path) as f: + lines = f.readlines() + if lines: + try: + progress = json.loads(lines[-1].strip()) + except json.JSONDecodeError: + pass + + return { + "task_id": task_id, + "status": "running", + "elapsed_s": _elapsed(task_dir), + "progress": progress, + } + + +def get_task_result(task_id: str, session_id: str) -> dict: + """Get the result of a completed task.""" + task_dir = _task_dir(session_id, task_id) + result_path = os.path.join(task_dir, "result.json") + + if not os.path.isfile(result_path): + return { + "task_id": task_id, + "status": "running", + "elapsed_s": _elapsed(task_dir), + } + + with open(result_path) as f: + result = json.load(f) + + result["task_id"] = task_id + result["elapsed_s"] = _elapsed(task_dir) + return result + + +def complete_task(session_id: str, task_id: str): + """Mark a task as completed and update session state back to idle.""" + session_data = _read_session(session_id) + session_data["status"] = "idle" + session_data["current_task"] = None + if task_id not in session_data.get("completed_tasks", []): + session_data.setdefault("completed_tasks", []).append(task_id) + _write_session(session_id, session_data) + + +def close_session(session_id: str) -> dict: + """Mark a session as closed.""" + session_data = _read_session(session_id) + session_data["status"] = "closed" + _write_session(session_id, session_data) + return {"session_id": session_id, "status": "closed"} + + +def wrap_prompt( + task_id: str, + session_id: str, + email: str, + prompt: str, + context: dict = None, + results_dir: str = "", + context_hint: str = None, +) -> str: + """Wrap a user prompt with the CODA-TASK convention.""" + context_block = "" + if context: + context_block = json.dumps(context, indent=2) + + hint_line = "" + if context_hint: + hint_line = f"context_hint: {context_hint}\n" + + return f"""---CODA-TASK--- +task_id: {task_id} +session_id: {session_id} +user: {email} +{hint_line}results_dir: {results_dir} + +CONTEXT: +{context_block} + +TASK: +{prompt} + +INSTRUCTIONS: +1. Append progress to {results_dir}/status.jsonl + Format: {{"step": "label", "message": "description"}} +2. When done, write {results_dir}/result.json with: + {{"status", "summary", "files_changed", "artifacts", "errors"}} +3. If you delegate to a sub-agent (Claude, Codex, Gemini), update + status.jsonl with delegation steps so the caller can track progress. +---END-CODA-TASK---""" + + +def _elapsed(task_dir: str) -> float: + """Calculate elapsed seconds since task started.""" + meta_path = os.path.join(task_dir, "meta.json") + if os.path.isfile(meta_path): + with open(meta_path) as f: + meta = json.load(f) + started = meta.get("started_at", "") + if started: + try: + started_ts = time.mktime(time.strptime(started, "%Y-%m-%dT%H:%M:%SZ")) + return round(time.time() - started_ts, 1) + except ValueError: + pass + # Fallback: use directory creation time + return round(time.time() - os.path.getctime(task_dir), 1) +``` + +**Step 4: Run tests to verify they pass** + +Run: `cd /Users/sathish.gangichetty/Documents/xterm-experiment/.worktrees/coda-mcp && uv run pytest tests/test_task_manager.py -v` +Expected: All 10 tests PASS + +**Step 5: Commit** + +```bash +cd /Users/sathish.gangichetty/Documents/xterm-experiment/.worktrees/coda-mcp +git add task_manager.py tests/test_task_manager.py +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" commit -m "feat: add task manager for MCP session/task state" +``` + +--- + +### Task 2: Create MCP Server Module + +The MCP server registers 5 tools and delegates to `task_manager.py` for state. It also integrates with the existing PTY session infrastructure in `app.py` for creating terminal sessions and piping prompts. + +**Files:** +- Create: `mcp_server.py` +- Create: `tests/test_mcp_server.py` + +**Step 1: Write the failing tests** + +```python +# tests/test_mcp_server.py +import json +import pytest +from unittest.mock import patch, MagicMock + + +def test_mcp_tool_list(): + """Verify all 5 tools are registered.""" + from mcp_server import mcp + # The server should have 5 tools registered + tools = mcp._tool_manager._tools # internal access for testing + tool_names = [t.name for t in tools.values()] + assert "create_session" in tool_names + assert "run_task" in tool_names + assert "get_status" in tool_names + assert "get_result" in tool_names + assert "close_session" in tool_names + assert len(tool_names) == 5 +``` + +**Step 2: Run test to verify it fails** + +Run: `cd /Users/sathish.gangichetty/Documents/xterm-experiment/.worktrees/coda-mcp && uv run pytest tests/test_mcp_server.py -v` +Expected: FAIL — `ModuleNotFoundError: No module named 'mcp_server'` + +**Step 3: Write the MCP server module** + +```python +# mcp_server.py +"""MCP server for CoDA — exposes coding agent capabilities to Genie Code. + +Registers 5 tools: create_session, run_task, get_status, get_result, close_session. +Uses the Python MCP SDK with stateless HTTP transport as required by Genie Code. +""" +import json +import logging +import os +import threading + +from mcp.server.fastmcp import FastMCP + +import task_manager + +logger = logging.getLogger(__name__) + +mcp = FastMCP( + "coda", + stateless_http=True, +) + +# Reference to app.py's session infrastructure — set by mount_mcp() +_app_create_session = None +_app_send_input = None +_app_close_session = None + + +def set_app_hooks(create_session_fn, send_input_fn, close_session_fn): + """Called by app.py to wire MCP tools to the PTY session infrastructure.""" + global _app_create_session, _app_send_input, _app_close_session + _app_create_session = create_session_fn + _app_send_input = send_input_fn + _app_close_session = close_session_fn + + +@mcp.tool() +def create_session( + email: str, + user_id: str = "", + label: str = "", +) -> str: + """Create a new coding agent session backed by Hermes Agent. + + Returns a session_id that can be used with run_task to send work. + Sessions are long-lived — reuse them for follow-up tasks to maintain context. + """ + # Create task manager state on disk + result = task_manager.create_session(email=email, user_id=user_id, label=label) + session_id = result["session_id"] + + # Create the actual PTY session via app.py infrastructure + if _app_create_session: + pty_session_id = _app_create_session(label="hermes-mcp") + # Map our session_id to the PTY session_id + task_manager._update_session_field(session_id, "pty_session_id", pty_session_id) + + return json.dumps(result) + + +@mcp.tool() +def run_task( + session_id: str, + prompt: str, + email: str, + user_id: str = "", + context: str = "{}", + context_hint: str = "", + timeout_s: int = 3600, + permissions: str = "smart", +) -> str: + """Send a coding task to Hermes Agent in an existing session. + + The task runs asynchronously — use get_status to poll progress + and get_result to retrieve the outcome. + + Args: + session_id: From create_session + prompt: Natural language task description + email: User email for audit trail + context: JSON string with Unity Catalog context (tables, schemas, etc.) + context_hint: "new_topic" to signal unrelated work in same session + timeout_s: Max seconds before timeout (default 3600) + permissions: "smart" (default, safe) or "yolo" (full autonomy) + """ + try: + context_dict = json.loads(context) if context else {} + except json.JSONDecodeError: + context_dict = {} + + try: + result = task_manager.create_task( + session_id=session_id, + prompt=prompt, + email=email, + context=context_dict, + context_hint=context_hint or None, + timeout_s=timeout_s, + permissions=permissions, + ) + except task_manager.SessionBusyError as e: + return json.dumps({"error": str(e)}) + except task_manager.SessionNotFoundError as e: + return json.dumps({"error": str(e)}) + + task_id = result["task_id"] + + # Read the wrapped prompt from disk + task_dir = task_manager._task_dir(session_id, task_id) + with open(os.path.join(task_dir, "prompt.txt")) as f: + wrapped_prompt = f.read() + + # Build hermes command + yolo_flag = " --yolo" if permissions == "yolo" else "" + hermes_cmd = f'hermes -z "{task_dir}/prompt.txt"{yolo_flag}\n' + + # Pipe to PTY session in background + if _app_send_input: + session_data = task_manager._read_session(session_id) + pty_session_id = session_data.get("pty_session_id") + if pty_session_id: + # Send the hermes command to the terminal + _app_send_input(pty_session_id, hermes_cmd) + + # Start background watcher for task completion + thread = threading.Thread( + target=_watch_task, + args=(session_id, task_id, timeout_s), + daemon=True, + ) + thread.start() + + return json.dumps(result) + + +@mcp.tool() +def get_status(task_id: str, session_id: str) -> str: + """Check the current status and progress of a running task. + + Returns status (running/completed/failed/timeout), elapsed time, + and the latest progress update from the agent if available. + """ + try: + result = task_manager.get_task_status(task_id, session_id) + return json.dumps(result) + except Exception as e: + return json.dumps({"error": str(e)}) + + +@mcp.tool() +def get_result(task_id: str, session_id: str) -> str: + """Retrieve the structured result of a completed task. + + Returns summary, files changed, artifacts (job IDs, commit hashes, etc.), + and any errors. If the task isn't done yet, returns running status. + """ + try: + result = task_manager.get_task_result(task_id, session_id) + return json.dumps(result) + except Exception as e: + return json.dumps({"error": str(e)}) + + +@mcp.tool() +def close_session(session_id: str) -> str: + """Close a session and clean up resources. + + The PTY process is terminated and session state is marked as closed. + """ + try: + # Close task manager state + result = task_manager.close_session(session_id) + + # Close the PTY session + if _app_close_session: + session_data = task_manager._read_session(session_id) + pty_session_id = session_data.get("pty_session_id") + if pty_session_id: + _app_close_session(pty_session_id) + + return json.dumps(result) + except Exception as e: + return json.dumps({"error": str(e)}) + + +def _watch_task(session_id: str, task_id: str, timeout_s: int): + """Background thread that watches for task completion or timeout.""" + import time + + task_dir = task_manager._task_dir(session_id, task_id) + result_path = os.path.join(task_dir, "result.json") + status_path = os.path.join(task_dir, "status.jsonl") + start = time.time() + last_activity = start + stale_threshold = 300 # 5 minutes with no status update = stale + + while True: + elapsed = time.time() - start + + # Check for result.json (task completed) + if os.path.isfile(result_path): + task_manager.complete_task(session_id, task_id) + logger.info(f"Task {task_id} completed in {elapsed:.0f}s") + return + + # Check for stale (no activity in 5 min) + if os.path.isfile(status_path): + mtime = os.path.getmtime(status_path) + if mtime > last_activity: + last_activity = mtime + + # Timeout: wall clock exceeded AND stale + if elapsed > timeout_s and (time.time() - last_activity) > stale_threshold: + logger.warning(f"Task {task_id} timed out after {elapsed:.0f}s") + # Write a timeout result + with open(result_path, "w") as f: + json.dump({ + "status": "timeout", + "summary": f"Task timed out after {elapsed:.0f} seconds", + "files_changed": [], + "artifacts": {}, + "errors": ["timeout"], + }, f) + task_manager.complete_task(session_id, task_id) + return + + time.sleep(5) # Poll every 5 seconds +``` + +**Step 4: Run tests to verify they pass** + +Run: `cd /Users/sathish.gangichetty/Documents/xterm-experiment/.worktrees/coda-mcp && uv run pytest tests/test_mcp_server.py -v` +Expected: PASS + +**Step 5: Commit** + +```bash +cd /Users/sathish.gangichetty/Documents/xterm-experiment/.worktrees/coda-mcp +git add mcp_server.py tests/test_mcp_server.py +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" commit -m "feat: add MCP server with 5 tools for Genie Code integration" +``` + +--- + +### Task 3: Mount MCP Server in Flask App + +Wire the MCP server into the existing Flask app. Add CORS support, skip auth for `/mcp` (Databricks proxy handles it), and expose helper functions for PTY integration. + +**Files:** +- Modify: `app.py` (add mount + helper functions) +- Modify: `pyproject.toml` (add flask-cors dependency) + +**Step 1: Add flask-cors to dependencies** + +In `pyproject.toml`, add `"flask-cors>=4.0"` to dependencies list. + +**Step 2: Add PTY helper functions to app.py** + +Add these functions after the existing `create_session` route (around line 1081), before the `send_input` route: + +```python +# ── MCP Integration Helpers ────────────────────────────────────────────── + +def mcp_create_pty_session(label: str = "hermes-mcp") -> str: + """Create a PTY session for MCP use. Returns the PTY session_id.""" + master_fd, slave_fd = pty.openpty() + shell_env = os.environ.copy() + shell_env["TERM"] = "xterm-256color" + shell_env.pop("CLAUDECODE", None) + shell_env.pop("CLAUDE_CODE_SESSION", None) + shell_env.pop("DATABRICKS_TOKEN", None) + shell_env.pop("DATABRICKS_HOST", None) + shell_env.pop("GEMINI_API_KEY", None) + if not shell_env.get("HOME") or shell_env["HOME"] == "/": + shell_env["HOME"] = "/app/python/source_code" + local_bin = f"{shell_env['HOME']}/.local/bin" + shell_env["PATH"] = f"{local_bin}:{shell_env.get('PATH', '')}" + projects_dir = os.path.join(shell_env["HOME"], "projects") + os.makedirs(projects_dir, exist_ok=True) + + pid = subprocess.Popen( + ["/bin/bash"], + stdin=slave_fd, stdout=slave_fd, stderr=slave_fd, + preexec_fn=os.setsid, + env=shell_env, + cwd=projects_dir + ).pid + os.close(slave_fd) + + session_id = str(uuid.uuid4()) + with sessions_lock: + if len(sessions) >= MAX_CONCURRENT_SESSIONS: + os.close(master_fd) + try: + os.kill(pid, signal.SIGKILL) + except OSError: + pass + raise RuntimeError(f"Maximum {MAX_CONCURRENT_SESSIONS} concurrent sessions reached") + sessions[session_id] = { + "master_fd": master_fd, + "pid": pid, + "output_buffer": deque(maxlen=1000), + "lock": threading.Lock(), + "last_poll_time": time.time(), + "created_at": time.time(), + "label": label, + } + + thread = threading.Thread(target=read_pty_output, args=(session_id, master_fd), daemon=True) + thread.start() + log_telemetry("agent", label) + return session_id + + +def mcp_send_input(session_id: str, data: str): + """Send input to a PTY session. Used by MCP to pipe hermes commands.""" + sess = _get_session(session_id) + if not sess: + return + with sess["lock"]: + try: + os.write(sess["master_fd"], data.encode()) + except OSError: + pass + + +def mcp_close_pty_session(session_id: str): + """Close a PTY session. Used by MCP close_session tool.""" + sess = _get_session(session_id) + if not sess: + return + terminate_session(session_id, sess["pid"], sess["master_fd"]) +``` + +**Step 3: Mount the MCP app and add CORS** + +At the end of `app.py`, before the `if __name__ == "__main__"` block (around line 1298), add: + +```python +# ── MCP Server Mount ───────────────────────────────────────────────────── +from flask_cors import CORS +from mcp_server import mcp, set_app_hooks + +# CORS for Genie Code cross-origin requests +databricks_host = os.environ.get("DATABRICKS_HOST", "") +if databricks_host: + CORS(app, origins=[ensure_https(databricks_host)], supports_credentials=True) + +# Wire MCP tools to PTY infrastructure +set_app_hooks( + create_session_fn=mcp_create_pty_session, + send_input_fn=mcp_send_input, + close_session_fn=mcp_close_pty_session, +) + +# Mount MCP as ASGI app at /mcp +from werkzeug.middleware.dispatcher import DispatcherMiddleware +from a]syncio import run as arun + +mcp_asgi_app = mcp.streamable_http_app() + +# Bridge ASGI MCP app into Flask's WSGI world +# We use a thin WSGI wrapper since Flask is WSGI and MCP SDK produces ASGI +import asyncio +from io import BytesIO + +def mcp_wsgi_app(environ, start_response): + """WSGI-to-ASGI bridge for the MCP endpoint.""" + # Read request body + content_length = int(environ.get('CONTENT_LENGTH', 0) or 0) + body = environ['wsgi.input'].read(content_length) if content_length else b'' + + async def run_asgi(): + response_started = False + status_code = None + response_headers = None + response_body = BytesIO() + + async def receive(): + return {"type": "http.request", "body": body} + + async def send(message): + nonlocal response_started, status_code, response_headers + if message["type"] == "http.response.start": + status_code = message["status"] + response_headers = [ + (k.decode() if isinstance(k, bytes) else k, + v.decode() if isinstance(v, bytes) else v) + for k, v in message.get("headers", []) + ] + response_started = True + elif message["type"] == "http.response.body": + response_body.write(message.get("body", b"")) + + scope = { + "type": "http", + "asgi": {"version": "3.0"}, + "http_version": "1.1", + "method": environ["REQUEST_METHOD"], + "path": environ.get("PATH_INFO", "/"), + "query_string": environ.get("QUERY_STRING", "").encode(), + "headers": [ + (k.lower().replace("http_", "").replace("_", "-").encode(), + v.encode()) + for k, v in environ.items() + if k.startswith("HTTP_") + ] + ( + [(b"content-type", environ["CONTENT_TYPE"].encode())] + if environ.get("CONTENT_TYPE") else [] + ), + "server": (environ.get("SERVER_NAME", "localhost"), + int(environ.get("SERVER_PORT", 8000))), + } + + await mcp_asgi_app(scope, receive, send) + return status_code, response_headers, response_body.getvalue() + + status_code, headers, body_bytes = asyncio.run(run_asgi()) + status_str = f"{status_code} OK" + start_response(status_str, headers or []) + return [body_bytes] + +app.wsgi_app = DispatcherMiddleware(app.wsgi_app, {"/mcp": mcp_wsgi_app}) +``` + +**Step 4: Update auth bypass for /mcp path** + +In `app.py` line 808, update the auth bypass to include `/mcp`: + +```python +# Before: +if request.path in ("/health", "/api/setup-status", ...): +# After: +if request.path in ("/health", "/api/setup-status", "/api/pat-status", "/api/configure-pat", "/api/app-state") or request.path.startswith("/socket.io") or request.path.startswith("/mcp"): +``` + +Note: `/mcp` auth is handled by the Databricks Apps proxy (same as all other routes), but the Flask `before_request` check would reject because MCP requests from Genie Code may not carry the same headers as browser requests. The Databricks Apps proxy still enforces authentication before the request reaches CoDA. + +**Step 5: Run the app locally to verify mount** + +Run: `cd /Users/sathish.gangichetty/Documents/xterm-experiment/.worktrees/coda-mcp && uv run python -c "from app import app; print('MCP mounted at /mcp'); print([rule.rule for rule in app.url_map.iter_rules()])"` +Expected: No import errors, `/mcp` visible in routes + +**Step 6: Commit** + +```bash +cd /Users/sathish.gangichetty/Documents/xterm-experiment/.worktrees/coda-mcp +git add app.py pyproject.toml mcp_server.py +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" commit -m "feat: mount MCP server at /mcp with CORS and PTY integration" +``` + +--- + +### Task 4: Add _update_session_field to task_manager + +The MCP server needs to store the `pty_session_id` mapping. Add the helper and its test. + +**Files:** +- Modify: `task_manager.py` (add `_update_session_field`) +- Modify: `tests/test_task_manager.py` (add test) + +**Step 1: Add test** + +```python +# Append to tests/test_task_manager.py + +def test_update_session_field(task_mgr): + session = task_mgr.create_session(email="alice@example.com", user_id="123") + sid = session["session_id"] + + task_mgr._update_session_field(sid, "pty_session_id", "pty-abc-123") + + with open(os.path.join(task_mgr.SESSIONS_DIR, sid, "session.json")) as f: + data = json.load(f) + assert data["pty_session_id"] == "pty-abc-123" +``` + +**Step 2: Add the function to task_manager.py** + +After the `_write_session` function: + +```python +def _update_session_field(session_id: str, key: str, value): + """Update a single field in session.json.""" + data = _read_session(session_id) + data[key] = value + _write_session(session_id, data) +``` + +**Step 3: Run tests** + +Run: `cd /Users/sathish.gangichetty/Documents/xterm-experiment/.worktrees/coda-mcp && uv run pytest tests/test_task_manager.py -v` +Expected: All 11 tests PASS + +**Step 4: Commit** + +```bash +cd /Users/sathish.gangichetty/Documents/xterm-experiment/.worktrees/coda-mcp +git add task_manager.py tests/test_task_manager.py +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" commit -m "feat: add _update_session_field helper for PTY mapping" +``` + +--- + +### Task 5: Update requirements.txt + +Regenerate requirements after adding flask-cors. + +**Files:** +- Modify: `pyproject.toml` (already done in Task 3) +- Regenerate: `requirements.txt` + +**Step 1: Regenerate requirements** + +Run: `cd /Users/sathish.gangichetty/Documents/xterm-experiment/.worktrees/coda-mcp && uv pip compile pyproject.toml -o requirements.txt` + +**Step 2: Commit** + +```bash +cd /Users/sathish.gangichetty/Documents/xterm-experiment/.worktrees/coda-mcp +git add pyproject.toml requirements.txt +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" commit -m "chore: add flask-cors dependency" +``` + +--- + +### Task 6: Integration Test — End-to-End MCP Flow + +Test the full flow: create session → run task → check status → get result → close session. + +**Files:** +- Create: `tests/test_mcp_integration.py` + +**Step 1: Write the integration test** + +```python +# tests/test_mcp_integration.py +"""Integration test for MCP server flow (no real PTY, mocked app hooks).""" +import json +import os +import pytest +from unittest.mock import patch, MagicMock + +import task_manager +import mcp_server + + +@pytest.fixture(autouse=True) +def setup_env(tmp_path): + """Redirect all state to temp dir and mock PTY hooks.""" + with patch.object(task_manager, "SESSIONS_DIR", str(tmp_path / "sessions")): + # Mock the app hooks (no real PTY in tests) + mcp_server.set_app_hooks( + create_session_fn=lambda label: "pty-mock-123", + send_input_fn=MagicMock(), + close_session_fn=MagicMock(), + ) + yield tmp_path + + +def test_full_mcp_flow(): + """End-to-end: create → run → status → result → close.""" + # 1. Create session + result = json.loads(mcp_server.create_session(email="alice@test.com", user_id="u1")) + assert result["status"] == "ready" + sid = result["session_id"] + + # 2. Run task + result = json.loads(mcp_server.run_task( + session_id=sid, + prompt="create a sales pipeline", + email="alice@test.com", + context='{"tables": ["sales.transactions"]}', + )) + assert result["status"] == "running" + tid = result["task_id"] + + # 3. Check status (running, no progress yet) + status = json.loads(mcp_server.get_status(task_id=tid, session_id=sid)) + assert status["status"] == "running" + assert status["progress"] is None + + # 4. Simulate agent writing progress + task_dir = task_manager._task_dir(sid, tid) + with open(os.path.join(task_dir, "status.jsonl"), "w") as f: + f.write(json.dumps({"step": "coding", "message": "Writing pipeline"}) + "\n") + + status = json.loads(mcp_server.get_status(task_id=tid, session_id=sid)) + assert status["progress"]["step"] == "coding" + + # 5. Simulate agent writing result + with open(os.path.join(task_dir, "result.json"), "w") as f: + json.dump({ + "status": "completed", + "summary": "Created sales pipeline with 3 stages", + "files_changed": ["pipelines/sales.py"], + "artifacts": {"job_id": "789"}, + "errors": [] + }, f) + + # 6. Get result + result = json.loads(mcp_server.get_result(task_id=tid, session_id=sid)) + assert result["status"] == "completed" + assert result["summary"] == "Created sales pipeline with 3 stages" + assert result["artifacts"]["job_id"] == "789" + + # 7. Complete and close + task_manager.complete_task(sid, tid) + result = json.loads(mcp_server.close_session(session_id=sid)) + assert result["status"] == "closed" + + +def test_busy_session_rejects(): + """Running a second task on a busy session should return error.""" + result = json.loads(mcp_server.create_session(email="bob@test.com")) + sid = result["session_id"] + + # First task + json.loads(mcp_server.run_task(session_id=sid, prompt="task 1", email="bob@test.com")) + + # Second task should fail + result = json.loads(mcp_server.run_task(session_id=sid, prompt="task 2", email="bob@test.com")) + assert "error" in result + assert "busy" in result["error"].lower() +``` + +**Step 2: Run tests** + +Run: `cd /Users/sathish.gangichetty/Documents/xterm-experiment/.worktrees/coda-mcp && uv run pytest tests/test_mcp_integration.py -v` +Expected: All 2 tests PASS + +**Step 3: Run all tests together** + +Run: `cd /Users/sathish.gangichetty/Documents/xterm-experiment/.worktrees/coda-mcp && uv run pytest tests/ -v` +Expected: All tests PASS + +**Step 4: Commit** + +```bash +cd /Users/sathish.gangichetty/Documents/xterm-experiment/.worktrees/coda-mcp +git add tests/test_mcp_integration.py +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" commit -m "test: add end-to-end MCP integration test" +``` + +--- + +## Summary + +| Task | What | Files | +|------|------|-------| +| 1 | Task manager (disk state) | `task_manager.py`, `tests/test_task_manager.py` | +| 2 | MCP server (5 tools) | `mcp_server.py`, `tests/test_mcp_server.py` | +| 3 | Flask mount + CORS + PTY helpers | `app.py`, `pyproject.toml` | +| 4 | Session field helper | `task_manager.py`, `tests/test_task_manager.py` | +| 5 | Dependencies | `pyproject.toml`, `requirements.txt` | +| 6 | Integration test | `tests/test_mcp_integration.py` | + +Total: 4 new files, 2 modified files, ~400 lines of production code, ~250 lines of tests. diff --git a/docs/superpowers/plans/2026-05-27-coda-mcp-live-session-url.md b/docs/superpowers/plans/2026-05-27-coda-mcp-live-session-url.md new file mode 100644 index 0000000..ade3838 --- /dev/null +++ b/docs/superpowers/plans/2026-05-27-coda-mcp-live-session-url.md @@ -0,0 +1,1900 @@ +# CoDA MCP Live Session URL Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add a `viewer_url` to CoDA MCP tool responses so the calling user can watch hermes execute live in a browser, with a 5-minute grace period after task completion and indefinite static replay from an on-disk PTY transcript. + +**Architecture:** Tee PTY bytes to `~/.coda/sessions/{sess}/tasks/{task}/transcript.log` from `read_pty_output`. Replace the immediate post-completion close in `_watch_task` with a `threading.Timer(300, close)`. Mark grace-period PTYs to exempt them from `MAX_CONCURRENT_SESSIONS`. Build `viewer_url` by capturing `X-Forwarded-Host` from inbound requests in an ASGI middleware. The Flask `/api/session/attach` endpoint adds a replay fallback that returns transcript bytes when the live PTY is gone. The SPA reads `?session=` on boot and routes to either the existing `_doAttach` (live) or a new `_doReplay` (static, chunked). + +**Tech Stack:** Python 3 (Flask + FastMCP + python-socketio AsyncServer + Starlette + uvicorn), xterm.js, pytest, `uv` for runs. + +**Spec:** `docs/superpowers/specs/2026-05-27-coda-mcp-live-session-url-design.md` at commit `02431c8` on `feat/coda-mcp-server`. + +--- + +## Conventions used in this plan + +- Worktree: `/Users/sathish.gangichetty/Documents/xterm-experiment/.worktrees/coda-mcp/` +- All `git commit` commands use `-c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty"` (per repo convention). No `Co-authored-by` line. +- All pytest invocations use `uv run pytest ...` (per repo convention). +- All file paths are relative to the worktree root. + +--- + +## Task 1: `coda_mcp/url_builder.py` — base URL resolution module + +**Files:** +- Create: `coda_mcp/url_builder.py` +- Test: `tests/test_url_builder.py` (new) + +- [ ] **Step 1: Write the failing tests** + +Create `tests/test_url_builder.py`: + +```python +"""Tests for url_builder module — base URL resolution for viewer_url.""" +import os +import importlib +from unittest import mock + +import pytest + + +@pytest.fixture(autouse=True) +def _reset_module(): + """Re-import url_builder fresh for each test (module-level cache).""" + from coda_mcp import url_builder + importlib.reload(url_builder) + yield + + +def test_returns_none_when_neither_env_nor_cache(): + from coda_mcp import url_builder + assert url_builder.build_viewer_url("pty-1") is None + + +def test_env_override_wins(): + from coda_mcp import url_builder + with mock.patch.dict(os.environ, {"CODA_APP_URL": "https://override.example.com"}): + assert url_builder.build_viewer_url("pty-1") == \ + "https://override.example.com/?session=pty-1" + + +def test_env_override_strips_trailing_slash(): + from coda_mcp import url_builder + with mock.patch.dict(os.environ, {"CODA_APP_URL": "https://override.example.com/"}): + assert url_builder.build_viewer_url("pty-1") == \ + "https://override.example.com/?session=pty-1" + + +def test_header_capture_used_when_no_env(): + from coda_mcp import url_builder + url_builder.capture_from_headers("app.databricksapps.com") + assert url_builder.build_viewer_url("pty-1") == \ + "https://app.databricksapps.com/?session=pty-1" + + +def test_env_overrides_header_capture(): + from coda_mcp import url_builder + url_builder.capture_from_headers("captured.example.com") + with mock.patch.dict(os.environ, {"CODA_APP_URL": "https://override.example.com"}): + assert url_builder.build_viewer_url("pty-1") == \ + "https://override.example.com/?session=pty-1" + + +def test_header_capture_overwrites_previous(): + from coda_mcp import url_builder + url_builder.capture_from_headers("first.example.com") + url_builder.capture_from_headers("second.example.com") + assert "second.example.com" in url_builder.build_viewer_url("pty-1") + + +def test_capture_empty_string_does_not_overwrite(): + from coda_mcp import url_builder + url_builder.capture_from_headers("good.example.com") + url_builder.capture_from_headers("") + assert "good.example.com" in url_builder.build_viewer_url("pty-1") + + +def test_capture_none_does_not_crash(): + from coda_mcp import url_builder + url_builder.capture_from_headers(None) + assert url_builder.build_viewer_url("pty-1") is None +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `uv run pytest tests/test_url_builder.py -v` +Expected: ImportError on `from coda_mcp import url_builder` — module does not exist yet. + +- [ ] **Step 3: Implement `coda_mcp/url_builder.py`** + +Create `coda_mcp/url_builder.py`: + +```python +"""Builds the viewer_url returned by CoDA MCP tools. + +Resolution order: +1. ``CODA_APP_URL`` env var (explicit override for local dev / power users). +2. Module-level cache populated by ``AppUrlCaptureMiddleware`` from the + ``X-Forwarded-Host`` header (officially provided by Databricks Apps). +3. ``None`` — caller omits the field entirely. + +The cache is process-global (single uvicorn worker per app) and refreshed +on every inbound HTTP request. +""" +from __future__ import annotations + +import os +from typing import Optional + +_app_url_cache: Optional[str] = None + + +def capture_from_headers(host: Optional[str]) -> None: + """Called by the ASGI middleware on every inbound HTTP request. + + No-op when ``host`` is falsy (None or empty) to avoid wiping a good + cache value with a missing header on a probe/CORS preflight. + """ + global _app_url_cache + if host: + _app_url_cache = host + + +def build_viewer_url(pty_session_id: str) -> Optional[str]: + """Return the full viewer URL for a PTY session, or None if no base is known.""" + override = os.environ.get("CODA_APP_URL", "").strip() + if override: + base = override.rstrip("/") + elif _app_url_cache: + base = f"https://{_app_url_cache}" + else: + return None + return f"{base}/?session={pty_session_id}" +``` + +- [ ] **Step 4: Run tests to verify they pass** + +Run: `uv run pytest tests/test_url_builder.py -v` +Expected: 8 passed. + +- [ ] **Step 5: Commit** + +```bash +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" \ + add coda_mcp/url_builder.py tests/test_url_builder.py +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" \ + commit -m "feat(coda-mcp): url_builder module for viewer_url resolution" +``` + +--- + +## Task 2: `task_manager.find_task_dir_by_pty_session` — reverse lookup with TTL cache + +**Files:** +- Modify: `coda_mcp/task_manager.py` (add new function at end, before `cleanup_expired_tasks`) +- Test: `tests/test_task_manager.py` (extend) + +- [ ] **Step 1: Write the failing tests** + +Append to `tests/test_task_manager.py` (locate existing test file; this assumes pytest fixtures `tmp_path` and patching of `SESSIONS_DIR` already exist in the file — confirm pattern, otherwise use the snippet below as a self-contained module): + +```python +import json +import os +import time +from unittest import mock + +import pytest + +from coda_mcp import task_manager + + +@pytest.fixture +def sessions_root(tmp_path, monkeypatch): + monkeypatch.setattr(task_manager, "SESSIONS_DIR", str(tmp_path)) + # Reset the lookup cache between tests + task_manager._pty_lookup_cache.clear() + return tmp_path + + +def _make_session_dir(root, sess_id, pty_id, current_task=None, completed=None): + sdir = root / sess_id + (sdir / "tasks").mkdir(parents=True) + data = { + "session_id": sess_id, + "pty_session_id": pty_id, + "current_task": current_task, + "completed_tasks": completed or [], + "status": "ready", + } + (sdir / "session.json").write_text(json.dumps(data)) + return sdir + + +def test_find_task_dir_hits_current_task(sessions_root): + _make_session_dir(sessions_root, "sess-A", "pty-1", current_task="task-X") + result = task_manager.find_task_dir_by_pty_session("pty-1") + assert result == str(sessions_root / "sess-A" / "tasks" / "task-X") + + +def test_find_task_dir_falls_back_to_last_completed(sessions_root): + _make_session_dir( + sessions_root, "sess-A", "pty-1", + current_task=None, + completed=["task-old", "task-recent"], + ) + result = task_manager.find_task_dir_by_pty_session("pty-1") + assert result == str(sessions_root / "sess-A" / "tasks" / "task-recent") + + +def test_find_task_dir_returns_none_when_no_match(sessions_root): + _make_session_dir(sessions_root, "sess-A", "pty-1", current_task="task-X") + assert task_manager.find_task_dir_by_pty_session("pty-NONEXIST") is None + + +def test_find_task_dir_ignores_corrupt_session_json(sessions_root): + sdir = sessions_root / "sess-bad" + sdir.mkdir() + (sdir / "session.json").write_text("not json {{{") + _make_session_dir(sessions_root, "sess-good", "pty-1", current_task="task-X") + assert task_manager.find_task_dir_by_pty_session("pty-1") == \ + str(sessions_root / "sess-good" / "tasks" / "task-X") + + +def test_find_task_dir_cache_hits_within_ttl(sessions_root): + _make_session_dir(sessions_root, "sess-A", "pty-1", current_task="task-X") + task_manager.find_task_dir_by_pty_session("pty-1") + # Remove session.json — cache should still return the hit + (sessions_root / "sess-A" / "session.json").unlink() + assert task_manager.find_task_dir_by_pty_session("pty-1") == \ + str(sessions_root / "sess-A" / "tasks" / "task-X") + + +def test_find_task_dir_cache_expires(sessions_root, monkeypatch): + monkeypatch.setattr(task_manager, "_PTY_LOOKUP_TTL", 0.01) + _make_session_dir(sessions_root, "sess-A", "pty-1", current_task="task-X") + task_manager.find_task_dir_by_pty_session("pty-1") + (sessions_root / "sess-A" / "session.json").unlink() + time.sleep(0.02) + assert task_manager.find_task_dir_by_pty_session("pty-1") is None + + +def test_find_task_dir_no_sessions_dir(sessions_root, monkeypatch): + monkeypatch.setattr(task_manager, "SESSIONS_DIR", "/nonexistent/path/that/does/not/exist") + assert task_manager.find_task_dir_by_pty_session("pty-1") is None +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `uv run pytest tests/test_task_manager.py -v -k find_task_dir` +Expected: 7 failures with `AttributeError: module 'coda_mcp.task_manager' has no attribute 'find_task_dir_by_pty_session'`. + +- [ ] **Step 3: Add module-level cache and function** + +Edit `coda_mcp/task_manager.py`. Near the top, after the existing module constants (after `TASK_TTL_S = ...`): + +```python +# ── PTY → task-dir reverse lookup (used by attach_session replay fallback) ── + +_pty_lookup_cache: dict[str, tuple[str, float]] = {} # pty_id -> (task_dir, ts) +_PTY_LOOKUP_TTL = 60.0 # seconds +``` + +Then before `def cleanup_expired_tasks()`, add: + +```python +def find_task_dir_by_pty_session(pty_session_id: str) -> str | None: + """Find the task dir whose session.json carries this pty_session_id. + + Returns the path to the active task dir, or — if the session has completed — + the most recently completed task dir. Returns None on no match. + + Cached for ``_PTY_LOOKUP_TTL`` seconds to avoid disk scans on every browser + refresh. + + Invariant: CoDA MCP sessions are ephemeral — one task per session. If the + lifecycle ever changes to allow multiple tasks per session, this function + must be revisited to pick the active or grace-period task rather than + ``completed_tasks[-1]``. + """ + now = time.time() + cached = _pty_lookup_cache.get(pty_session_id) + if cached and (now - cached[1]) < _PTY_LOOKUP_TTL: + return cached[0] + + if not os.path.isdir(SESSIONS_DIR): + return None + + for sess_name in os.listdir(SESSIONS_DIR): + sess_file = os.path.join(SESSIONS_DIR, sess_name, "session.json") + try: + with open(sess_file) as f: + data = json.load(f) + except (OSError, json.JSONDecodeError): + continue + + if data.get("pty_session_id") != pty_session_id: + continue + + candidate = data.get("current_task") or ( + data["completed_tasks"][-1] if data.get("completed_tasks") else None + ) + if candidate: + tdir = os.path.join(SESSIONS_DIR, sess_name, "tasks", candidate) + _pty_lookup_cache[pty_session_id] = (tdir, now) + return tdir + + return None +``` + +- [ ] **Step 4: Run tests to verify they pass** + +Run: `uv run pytest tests/test_task_manager.py -v -k find_task_dir` +Expected: 7 passed. + +- [ ] **Step 5: Commit** + +```bash +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" \ + add coda_mcp/task_manager.py tests/test_task_manager.py +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" \ + commit -m "feat(coda-mcp): find_task_dir_by_pty_session lookup with TTL cache" +``` + +--- + +## Task 3: `app.py::read_pty_output` — tee PTY bytes to transcript with lock-guarded writes + +**Files:** +- Modify: `app.py` (top: new constant; `read_pty_output` function lines 861-910) +- Test: `tests/test_transcript.py` (new — standalone unit tests for the tee logic; integration tested later) + +- [ ] **Step 1: Write the failing tests** + +Create `tests/test_transcript.py`: + +```python +"""Unit tests for the transcript tee in read_pty_output. + +These tests exercise the tee logic directly by simulating output dispatch into +a synthesized session dict and a real on-disk transcript file. The full PTY +read loop is not exercised here — see test_mcp_integration.py for E2E. +""" +import os +import stat +import threading +from pathlib import Path + +import pytest + + +@pytest.fixture +def session_dict(tmp_path): + """Build a minimally valid sessions[pty_id] entry with a real transcript handle.""" + transcript = tmp_path / "transcript.log" + fh = open(transcript, "ab", buffering=0) + os.fchmod(fh.fileno(), 0o600) + return { + "transcript_path": str(transcript), + "transcript_fh": fh, + "transcript_bytes": 0, + "lock": threading.Lock(), + } + + +def _write_chunk(session, output: bytes, cap: int = 10 * 1024 * 1024) -> None: + """Mirror the tee logic from read_pty_output for unit testing.""" + from app import _tee_transcript_chunk + _tee_transcript_chunk(session, output, cap=cap) + + +def test_tee_writes_bytes_and_flushes(session_dict): + _write_chunk(session_dict, b"hello world\n") + assert session_dict["transcript_bytes"] == 12 + assert Path(session_dict["transcript_path"]).read_bytes() == b"hello world\n" + + +def test_tee_chmod_is_0600(session_dict): + mode = stat.S_IMODE(os.stat(session_dict["transcript_path"]).st_mode) + assert mode == 0o600 + + +def test_tee_truncation_at_cap(session_dict): + cap = 16 + _write_chunk(session_dict, b"AAAAAAAAAA", cap=cap) + _write_chunk(session_dict, b"BBBBBBBBBBBBBBBBBBBB", cap=cap) + body = Path(session_dict["transcript_path"]).read_bytes() + # 10 A's, then 6 B's, then truncation marker. + assert body.startswith(b"AAAAAAAAAABBBBBB") + assert b"[transcript truncated at" in body + # Handle is closed after marker + assert session_dict["transcript_fh"] is None + + +def test_tee_no_op_when_fh_is_none(session_dict): + session_dict["transcript_fh"] = None + _write_chunk(session_dict, b"should not write") + assert Path(session_dict["transcript_path"]).read_bytes() == b"" + + +def test_tee_handles_write_error(session_dict, monkeypatch): + # Close the handle out from under the tee — write() will ValueError. + session_dict["transcript_fh"].close() + _write_chunk(session_dict, b"this will fail") + # Handle replaced with None; no crash. + assert session_dict["transcript_fh"] is None +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `uv run pytest tests/test_transcript.py -v` +Expected: ImportError on `from app import _tee_transcript_chunk`. + +- [ ] **Step 3: Add the helper and the constant in `app.py`** + +Near the top of `app.py` (after the existing constants block around line 46-50), add: + +```python +TRANSCRIPT_CAP_BYTES = 10 * 1024 * 1024 # 10 MB soft cap per transcript +``` + +Then add the helper (place it near `read_pty_output`, e.g., immediately above it): + +```python +def _tee_transcript_chunk(session, output: bytes, cap: int = TRANSCRIPT_CAP_BYTES) -> None: + """Append PTY output to the transcript file. Single-writer (read_pty_output). + + All file-handle access is under ``session["lock"]`` so we never race the + Timer-driven close path in ``terminate_session``. The ``ValueError`` catch + is belt-and-suspenders for the tiny window where the handle is closed + between the ``is not None`` check and the actual ``write`` call (the lock + prevents this, but be defensive). + """ + with session["lock"]: + fh = session.get("transcript_fh") + written = session.get("transcript_bytes", 0) + if fh is None: + return + remaining = cap - written + if remaining <= 0: + return + chunk = output[:remaining] + try: + fh.write(chunk) + fh.flush() + session["transcript_bytes"] = written + len(chunk) + if len(chunk) < len(output): + fh.write(b"\n[transcript truncated at %d bytes]\n" % cap) + fh.flush() + fh.close() + session["transcript_fh"] = None + except (OSError, ValueError) as exc: + logger.warning("transcript write failed: %s", exc) + try: + fh.close() + except Exception: + pass + session["transcript_fh"] = None +``` + +- [ ] **Step 4: Wire the tee into `read_pty_output`** + +In `app.py::read_pty_output`, locate the block (currently around line 880-888): + +```python + decoded = output.decode(errors="replace") + with session_lock: + # Buffer for HTTP polling fallback (AC-15) + session["output_buffer"].append(decoded) + session["last_poll_time"] = time.time() # Keep session alive during WS output + # Push via WebSocket to the session room (AC-8) + _emit_from_thread('terminal_output', + {'session_id': session_id, 'output': decoded}, + room=session_id) +``` + +Immediately after the `_emit_from_thread` call (and before the `else:` branch), add: + +```python + # Tee to transcript file if enabled for this session + _tee_transcript_chunk(session, output) +``` + +- [ ] **Step 5: Run unit tests to verify they pass** + +Run: `uv run pytest tests/test_transcript.py -v` +Expected: 5 passed. + +- [ ] **Step 6: Run existing terminal tests to verify no regression** + +Run: `uv run pytest tests/test_terminal_env_strip.py tests/test_session_linger.py tests/test_session_detach.py -v` +Expected: existing pass count unchanged (no failures introduced). + +- [ ] **Step 7: Commit** + +```bash +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" \ + add app.py tests/test_transcript.py +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" \ + commit -m "feat: tee PTY output to transcript.log with lock-guarded writes" +``` + +--- + +## Task 4: `app.py` — open transcript handle in `mcp_create_pty_session` + close in `terminate_session` + +**Files:** +- Modify: `app.py::mcp_create_pty_session` (lines ~1324-1387) +- Modify: `app.py::terminate_session` (lines ~912-936) + +- [ ] **Step 1: Write the failing tests** + +Append to `tests/test_transcript.py`: + +```python +def test_mcp_create_pty_session_opens_transcript_when_path_given(tmp_path, monkeypatch): + monkeypatch.setattr("app.MAX_CONCURRENT_SESSIONS", 5) + transcript = tmp_path / "transcript.log" + from app import mcp_create_pty_session, sessions, mcp_close_pty_session + sid = mcp_create_pty_session(label="test", transcript_path=str(transcript)) + try: + assert transcript.exists() + mode = stat.S_IMODE(os.stat(transcript).st_mode) + assert mode == 0o600 + sess = sessions[sid] + assert sess["transcript_path"] == str(transcript) + assert sess["transcript_fh"] is not None + assert sess["transcript_bytes"] == 0 + finally: + mcp_close_pty_session(sid) + + +def test_mcp_create_pty_session_no_transcript_when_path_none(monkeypatch): + monkeypatch.setattr("app.MAX_CONCURRENT_SESSIONS", 5) + from app import mcp_create_pty_session, sessions, mcp_close_pty_session + sid = mcp_create_pty_session(label="test") + try: + sess = sessions[sid] + assert sess.get("transcript_fh") is None + assert sess.get("transcript_path") is None + finally: + mcp_close_pty_session(sid) + + +def test_terminate_session_closes_transcript_handle(tmp_path, monkeypatch): + monkeypatch.setattr("app.MAX_CONCURRENT_SESSIONS", 5) + transcript = tmp_path / "transcript.log" + from app import mcp_create_pty_session, sessions, mcp_close_pty_session + sid = mcp_create_pty_session(label="test", transcript_path=str(transcript)) + fh = sessions[sid]["transcript_fh"] + mcp_close_pty_session(sid) + assert fh.closed + # Session removed from dict + assert sid not in sessions +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `uv run pytest tests/test_transcript.py -v -k "create_pty or terminate"` +Expected: 3 failures — `mcp_create_pty_session` does not yet accept `transcript_path`. + +- [ ] **Step 3: Modify `mcp_create_pty_session` signature** + +In `app.py`, change the signature (line ~1324): + +```python +def mcp_create_pty_session(label: str = "hermes-mcp", transcript_path: str | None = None) -> str: +``` + +After the `os.close(slave_fd)` line (around line 1358) and before `session_id = str(uuid.uuid4())`, add the transcript open. Place it inside the existing flow so the file handle is constructed before being stored: + +```python + # Open transcript file (if requested) before locking the session dict. + transcript_fh = None + if transcript_path: + try: + os.makedirs(os.path.dirname(transcript_path), exist_ok=True) + transcript_fh = open(transcript_path, "ab", buffering=0) + os.fchmod(transcript_fh.fileno(), 0o600) + except OSError as exc: + logger.warning("Could not open transcript at %s: %s", transcript_path, exc) + transcript_fh = None +``` + +Modify the `sessions[session_id] = { ... }` block to include the new fields: + +```python + sessions[session_id] = { + "master_fd": master_fd, + "pid": pid, + "output_buffer": deque(maxlen=1000), + "lock": threading.Lock(), + "last_poll_time": time.time(), + "created_at": time.time(), + "label": label, + "transcript_path": transcript_path if transcript_fh else None, + "transcript_fh": transcript_fh, + "transcript_bytes": 0, + "grace": False, + } +``` + +- [ ] **Step 4: Modify `terminate_session` to close the transcript handle** + +In `app.py::terminate_session` (line ~912), at the top of the function (right after the `logger.info` and the `_emit_from_thread('session_closed', ...)` call), add: + +```python + # Close transcript handle (if any) under per-session lock; swap-then-close + # outside the lock to avoid blocking on slow filesystems. + with sessions_lock: + sess = sessions.get(session_id) + if sess is not None: + with sess["lock"]: + transcript_fh = sess.get("transcript_fh") + sess["transcript_fh"] = None + if transcript_fh is not None: + try: + transcript_fh.close() + except Exception: + pass +``` + +- [ ] **Step 5: Run tests to verify they pass** + +Run: `uv run pytest tests/test_transcript.py -v -k "create_pty or terminate"` +Expected: 3 passed. + +- [ ] **Step 6: Run full transcript test suite** + +Run: `uv run pytest tests/test_transcript.py -v` +Expected: 8 passed. + +- [ ] **Step 7: Commit** + +```bash +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" \ + add app.py tests/test_transcript.py +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" \ + commit -m "feat: open transcript handle in mcp_create_pty_session; close in terminate_session" +``` + +--- + +## Task 5: `app.py` — grace-period exemption from `MAX_CONCURRENT_SESSIONS` + helper hooks + +**Files:** +- Modify: `app.py` (the two `MAX_CONCURRENT_SESSIONS` check sites + add two new helpers near the bottom near other MCP hook functions) + +- [ ] **Step 1: Write the failing tests** + +Append to `tests/test_transcript.py`: + +```python +def test_grace_period_pty_does_not_count_toward_max(monkeypatch): + monkeypatch.setattr("app.MAX_CONCURRENT_SESSIONS", 2) + from app import mcp_create_pty_session, mcp_close_pty_session, sessions, _mark_grace_for_session + + sid1 = mcp_create_pty_session(label="t1") + sid2 = mcp_create_pty_session(label="t2") + try: + # At cap. A third creation should raise. + with pytest.raises(RuntimeError, match="Maximum"): + mcp_create_pty_session(label="t3") + # Mark one as grace; now we should have headroom. + _mark_grace_for_session(sid1) + assert sessions[sid1]["grace"] is True + sid3 = mcp_create_pty_session(label="t3") + mcp_close_pty_session(sid3) + finally: + for s in [sid1, sid2]: + try: mcp_close_pty_session(s) + except Exception: pass + + +def test_bump_session_last_poll_advances_clock(monkeypatch): + monkeypatch.setattr("app.MAX_CONCURRENT_SESSIONS", 5) + from app import mcp_create_pty_session, mcp_close_pty_session, sessions, _bump_session_last_poll + sid = mcp_create_pty_session(label="t") + try: + baseline = sessions[sid]["last_poll_time"] + _bump_session_last_poll(sid, 300) + assert sessions[sid]["last_poll_time"] >= baseline + 299 + finally: + mcp_close_pty_session(sid) + + +def test_mark_grace_on_missing_session_is_noop(): + from app import _mark_grace_for_session + _mark_grace_for_session("nonexistent-pty-id") # must not raise + + +def test_bump_session_last_poll_missing_is_noop(): + from app import _bump_session_last_poll + _bump_session_last_poll("nonexistent-pty-id", 100) # must not raise +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `uv run pytest tests/test_transcript.py -v -k "grace or bump_session"` +Expected: failures — `_mark_grace_for_session` / `_bump_session_last_poll` don't exist; the cap check still uses raw `len`. + +- [ ] **Step 3: Replace the `MAX_CONCURRENT_SESSIONS` checks** + +There are two checkpoints in `app.py`: + +**Site 1 — `create_session()` (around line 1252):** + +```python + with sessions_lock: + if len(sessions) >= MAX_CONCURRENT_SESSIONS: + return jsonify({"error": f"Maximum {MAX_CONCURRENT_SESSIONS} concurrent sessions reached. Close an existing session first."}), 429 +``` + +Replace with: + +```python + with sessions_lock: + active = sum(1 for s in sessions.values() if not s.get("grace")) + if active >= MAX_CONCURRENT_SESSIONS: + return jsonify({"error": f"Maximum {MAX_CONCURRENT_SESSIONS} concurrent sessions reached. Close an existing session first."}), 429 +``` + +**Site 2 — `mcp_create_pty_session()` (around lines 1326-1330 and again 1362-1371):** + +Both `len(sessions) >= MAX_CONCURRENT_SESSIONS` checks become: + +```python + active = sum(1 for s in sessions.values() if not s.get("grace")) + if active >= MAX_CONCURRENT_SESSIONS: + raise RuntimeError( + f"Maximum {MAX_CONCURRENT_SESSIONS} concurrent sessions reached." + ) +``` + +(Apply at both pre-spawn and post-spawn check sites.) + +- [ ] **Step 4: Add the two helper functions** + +Place near `mcp_close_pty_session` (around line 1399): + +```python +def _mark_grace_for_session(session_id: str) -> None: + """Mark a PTY session as 'in grace period' so it doesn't count toward + MAX_CONCURRENT_SESSIONS. Called by ``_watch_task`` immediately before + scheduling the deferred close Timer. + + No-op if the session does not exist (e.g., already torn down). + """ + with sessions_lock: + sess = sessions.get(session_id) + if sess is None: + return + with sess["lock"]: + sess["grace"] = True + + +def _bump_session_last_poll(session_id: str, delta_s: float) -> None: + """Advance ``last_poll_time`` by ``delta_s`` so the idle reaper can't + preempt the Timer's deferred close. Defensive: at the current 24h + SESSION_TIMEOUT_SECONDS the reaper would never win anyway, but a future + tuning shouldn't break the grace window. + + No-op if the session does not exist. + """ + with sessions_lock: + sess = sessions.get(session_id) + if sess is None: + return + with sess["lock"]: + sess["last_poll_time"] = time.time() + delta_s +``` + +- [ ] **Step 5: Run tests to verify they pass** + +Run: `uv run pytest tests/test_transcript.py -v -k "grace or bump_session"` +Expected: 4 passed. + +- [ ] **Step 6: Run full transcript suite + session limit test for regression** + +Run: `uv run pytest tests/test_transcript.py tests/test_session_limit.py -v` +Expected: all pass. + +- [ ] **Step 7: Commit** + +```bash +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" \ + add app.py tests/test_transcript.py +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" \ + commit -m "feat: exempt grace-period PTYs from MAX_CONCURRENT_SESSIONS" +``` + +--- + +## Task 6: `mcp_server.py` — wire deferred close via `Timer`; update `set_app_hooks` + +**Files:** +- Modify: `coda_mcp/mcp_server.py` (lines 70-90 hook plumbing; lines 94-148 `_watch_task` + helpers) +- Test: `tests/test_mcp_server.py` (extend) + +- [ ] **Step 1: Write the failing tests** + +Append to `tests/test_mcp_server.py`: + +```python +import threading +from unittest import mock + +from coda_mcp import mcp_server, task_manager + + +def test_set_app_hooks_accepts_grace_and_bump_hooks(): + create = mock.MagicMock() + send = mock.MagicMock() + close = mock.MagicMock() + mark_grace = mock.MagicMock() + bump_poll = mock.MagicMock() + mcp_server.set_app_hooks(create, send, close, mark_grace, bump_poll) + assert mcp_server._app_mark_grace is mark_grace + assert mcp_server._app_bump_poll is bump_poll + + +def test_watch_task_schedules_timer_on_completion(tmp_path, monkeypatch): + monkeypatch.setattr(task_manager, "SESSIONS_DIR", str(tmp_path)) + # Create a session + task with a faked result.json + s = task_manager.create_session("u@x", "uid", label="t") + sid = s["session_id"] + task_manager._update_session_field(sid, "pty_session_id", "pty-abc") + t = task_manager.create_task(sid, "do thing", "u@x") + tid = t["task_id"] + tdir = task_manager._task_dir(sid, tid) + task_manager._write_json(tdir + "/result.json", {"status": "completed"}) + + mark = mock.MagicMock() + bump = mock.MagicMock() + closer = mock.MagicMock() + mcp_server.set_app_hooks(mock.MagicMock(), mock.MagicMock(), closer, mark, bump) + + timer_created = [] + real_timer = threading.Timer + + def fake_timer(seconds, fn, args=None, kwargs=None): + timer_created.append((seconds, fn, args)) + t = real_timer(seconds, fn, args=args, kwargs=kwargs) + return t + + monkeypatch.setattr(mcp_server.threading, "Timer", fake_timer) + + # Use a very short watch interval and ensure no real Timer fires + monkeypatch.setattr(mcp_server, "GRACE_PERIOD_S", 0.05) + + # Run one iteration manually + mcp_server._watch_task(sid, tid, timeout_s=10) + + # Timer should be scheduled for GRACE_PERIOD_S seconds with closer + pty_session_id + assert len(timer_created) == 1 + delay, fn, args = timer_created[0] + assert delay == 0.05 + assert fn is closer + assert args == ("pty-abc",) + + # _mark_grace and _bump_session_last_poll should have been called + mark.assert_called_once_with("pty-abc") + bump.assert_called_once_with("pty-abc", 0.05) +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `uv run pytest tests/test_mcp_server.py -v -k "set_app_hooks_accepts or watch_task_schedules"` +Expected: failures — extra params on `set_app_hooks` not accepted; `_watch_task` calls close synchronously. + +- [ ] **Step 3: Extend `set_app_hooks` and module state** + +In `coda_mcp/mcp_server.py`, at the top of the "App hooks" block (around line 70), expand: + +```python +_app_create_session = None +_app_send_input = None +_app_close_session = None +_app_mark_grace = None +_app_bump_poll = None + +GRACE_PERIOD_S = 300 # 5 minutes + + +def set_app_hooks( + create_session_fn, + send_input_fn, + close_session_fn, + mark_grace_fn=None, + bump_poll_fn=None, +): + """Wire up Flask app callbacks for PTY operations. + + The two new optional hooks (mark_grace, bump_poll) are used by ``_watch_task`` + to defer PTY close by ``GRACE_PERIOD_S`` after task completion so live viewers + can keep watching for a few minutes. + """ + global _app_create_session, _app_send_input, _app_close_session + global _app_mark_grace, _app_bump_poll + _app_create_session = create_session_fn + _app_send_input = send_input_fn + _app_close_session = close_session_fn + _app_mark_grace = mark_grace_fn + _app_bump_poll = bump_poll_fn +``` + +- [ ] **Step 4: Replace the immediate close inside `_watch_task`** + +Replace the existing `_close_pty_for_session(session_id)` calls in `_watch_task` (one in the completion branch around line 117, one in the timeout branch around line 144) with the deferred-Timer helper. Add a new helper at the bottom of the existing helper section (right after `_close_pty_for_session` around line 161): + +```python +def _schedule_deferred_close(session_id: str) -> None: + """Mark the PTY as in-grace and schedule a delayed close. + + Both completion and timeout paths call this in place of the immediate + ``_close_pty_for_session``. The Timer is a daemon thread so it doesn't + block uvicorn shutdown. + """ + if _app_close_session is None: + return + try: + session = task_manager._read_session(session_id) + except task_manager.SessionNotFoundError: + return + pty_session_id = session.get("pty_session_id") + if not pty_session_id: + return + + if _app_mark_grace is not None: + _app_mark_grace(pty_session_id) + if _app_bump_poll is not None: + _app_bump_poll(pty_session_id, GRACE_PERIOD_S) + + t = threading.Timer(GRACE_PERIOD_S, _app_close_session, args=(pty_session_id,)) + t.daemon = True + t.start() + logger.info( + "Watcher: scheduled deferred close for pty %s in %ds", + pty_session_id, GRACE_PERIOD_S, + ) +``` + +Then in `_watch_task`, replace both occurrences of `_close_pty_for_session(session_id)` with `_schedule_deferred_close(session_id)`. + +- [ ] **Step 5: Run tests to verify they pass** + +Run: `uv run pytest tests/test_mcp_server.py -v -k "set_app_hooks_accepts or watch_task_schedules"` +Expected: 2 passed. + +- [ ] **Step 6: Run full mcp_server test suite for regression** + +Run: `uv run pytest tests/test_mcp_server.py -v` +Expected: all pass (existing tests should be unaffected since hooks default to None). + +- [ ] **Step 7: Commit** + +```bash +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" \ + add coda_mcp/mcp_server.py tests/test_mcp_server.py +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" \ + commit -m "feat(coda-mcp): defer PTY close by GRACE_PERIOD_S via threading.Timer" +``` + +--- + +## Task 7: `mcp_server.py` — return `viewer_url` from all three tools + pass `transcript_path` to PTY creation + update instructions + +**Files:** +- Modify: `coda_mcp/mcp_server.py` (`coda_run` body, `coda_inbox` body, `coda_get_result` body, `instructions` block) + +- [ ] **Step 1: Write the failing tests** + +Append to `tests/test_mcp_server.py`: + +```python +import asyncio +import json +import os +from unittest import mock + +from coda_mcp import mcp_server, task_manager, url_builder + + +def _run(coro): + return asyncio.get_event_loop().run_until_complete(coro) if not asyncio.iscoroutine(coro) else asyncio.run(coro) + + +def test_coda_run_includes_viewer_url_when_builder_returns_one(tmp_path, monkeypatch): + monkeypatch.setattr(task_manager, "SESSIONS_DIR", str(tmp_path)) + monkeypatch.setattr(url_builder, "_app_url_cache", "app.example.com") + + create = mock.MagicMock(return_value="pty-abc") + send = mock.MagicMock() + closer = mock.MagicMock() + mcp_server.set_app_hooks(create, send, closer, mock.MagicMock(), mock.MagicMock()) + + result_json = asyncio.run(mcp_server.coda_run(prompt="do it", email="u@x")) + result = json.loads(result_json) + assert result["status"] == "running" + assert "?session=pty-abc" in result["viewer_url"] + assert result["viewer_url"].startswith("https://app.example.com") + + +def test_coda_run_omits_viewer_url_when_builder_returns_none(tmp_path, monkeypatch): + monkeypatch.setattr(task_manager, "SESSIONS_DIR", str(tmp_path)) + monkeypatch.setattr(url_builder, "_app_url_cache", None) + monkeypatch.delenv("CODA_APP_URL", raising=False) + + create = mock.MagicMock(return_value="pty-abc") + mcp_server.set_app_hooks(create, mock.MagicMock(), mock.MagicMock(), mock.MagicMock(), mock.MagicMock()) + + result_json = asyncio.run(mcp_server.coda_run(prompt="do it", email="u@x")) + result = json.loads(result_json) + # viewer_url present but None when builder returns None + assert result.get("viewer_url") is None + + +def test_coda_run_passes_transcript_path_to_create_session(tmp_path, monkeypatch): + monkeypatch.setattr(task_manager, "SESSIONS_DIR", str(tmp_path)) + create = mock.MagicMock(return_value="pty-abc") + mcp_server.set_app_hooks(create, mock.MagicMock(), mock.MagicMock(), mock.MagicMock(), mock.MagicMock()) + + asyncio.run(mcp_server.coda_run(prompt="do it", email="u@x")) + # create_session was called with transcript_path=... pointing into ~/.coda/sessions//tasks//transcript.log + kwargs = create.call_args.kwargs + assert "transcript_path" in kwargs + assert kwargs["transcript_path"].endswith("transcript.log") + assert "tasks" in kwargs["transcript_path"] + + +def test_coda_inbox_decorates_each_task_with_viewer_url(tmp_path, monkeypatch): + monkeypatch.setattr(task_manager, "SESSIONS_DIR", str(tmp_path)) + monkeypatch.setattr(url_builder, "_app_url_cache", "app.example.com") + + # Seed one session with one task and a pty_session_id + s = task_manager.create_session("u@x", "uid", label="t") + sid = s["session_id"] + task_manager._update_session_field(sid, "pty_session_id", "pty-xyz") + task_manager.create_task(sid, "prompt", "u@x") + + result_json = asyncio.run(mcp_server.coda_inbox()) + result = json.loads(result_json) + assert len(result["tasks"]) == 1 + assert "viewer_url" in result["tasks"][0] + assert "?session=pty-xyz" in result["tasks"][0]["viewer_url"] + + +def test_coda_get_result_includes_viewer_url(tmp_path, monkeypatch): + monkeypatch.setattr(task_manager, "SESSIONS_DIR", str(tmp_path)) + monkeypatch.setattr(url_builder, "_app_url_cache", "app.example.com") + + s = task_manager.create_session("u@x", "uid", label="t") + sid = s["session_id"] + task_manager._update_session_field(sid, "pty_session_id", "pty-xyz") + t = task_manager.create_task(sid, "prompt", "u@x") + tid = t["task_id"] + tdir = task_manager._task_dir(sid, tid) + task_manager._write_json(tdir + "/result.json", { + "status": "completed", "summary": "ok", + }) + + result_json = asyncio.run(mcp_server.coda_get_result(tid, sid)) + result = json.loads(result_json) + assert "viewer_url" in result + assert "?session=pty-xyz" in result["viewer_url"] +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `uv run pytest tests/test_mcp_server.py -v -k "viewer_url or transcript_path"` +Expected: failures — fields not present, `transcript_path` not passed. + +- [ ] **Step 3: Modify `coda_run`** + +In `coda_mcp/mcp_server.py`, at the top of the file add the import: + +```python +from coda_mcp import url_builder +``` + +In the body of `coda_run` (around line 219), modify the PTY creation block to compute and pass the transcript path: + +```python + # Create PTY if hooks are wired + if _app_create_session is not None: + transcript_path = os.path.join( + task_manager._task_dir(session_id, _new_task_id_preview := task_manager._new_task_id()), + "transcript.log", + ) +``` + +Wait — `task_id` isn't known until after `task_manager.create_task`. Restructure: create the task FIRST (so we have task_id), then create the PTY with transcript path, then send the input. The existing order is: create_session → create_pty → update session with pty_id → create_task → send_input. We need: create_session → create_task → create_pty(transcript_path) → update session with pty_id → send_input. + +Replace the existing PTY-create + create_task block (lines ~218-258) with this restructured version: + +```python + # Create task first (we need task_id to compute transcript_path). + result = task_manager.create_task( + session_id=session_id, + prompt=prompt, + email=email, + context=ctx, + timeout_s=timeout_s, + permissions=permissions, + previous_session_id=previous_session_id or None, + ) + task_id = result["task_id"] + + pty_session_id = None + if _app_create_session is not None: + transcript_path = os.path.join( + task_manager._task_dir(session_id, task_id), + "transcript.log", + ) + pty_session_id = _app_create_session( + label="hermes-mcp", + transcript_path=transcript_path, + ) + task_manager._update_session_field( + session_id, "pty_session_id", pty_session_id + ) + + # Send to PTY if hooks are wired + if _app_send_input is not None and pty_session_id is not None: + tdir = task_manager._task_dir(session_id, task_id) + prompt_path = os.path.join(tdir, "prompt.txt") + cmd = f'hermes -z "{prompt_path}"' + if permissions == "yolo": + cmd += " --yolo" + cmd += "\n" + _app_send_input(pty_session_id, cmd) + + # Start background watcher + t = threading.Thread( + target=_watch_task, + args=(session_id, task_id, timeout_s), + daemon=True, + ) + t.start() + + return json.dumps({ + "task_id": task_id, + "session_id": session_id, + "status": "running", + "viewer_url": url_builder.build_viewer_url(pty_session_id) if pty_session_id else None, + }) +``` + +- [ ] **Step 4: Add `viewer_url` to `coda_inbox` entries** + +In `coda_inbox` (around line 300), after the `list_all_tasks` call, decorate each entry. Replace: + +```python + tasks = task_manager.list_all_tasks(email=email, status_filter=status) +``` + +with: + +```python + tasks = task_manager.list_all_tasks(email=email, status_filter=status) + # Decorate each task with its viewer URL (if available). + for t in tasks: + sess = task_manager._read_session_safe(t["session_id"]) + pty = sess.get("pty_session_id") if sess else None + if pty: + vu = url_builder.build_viewer_url(pty) + if vu: + t["viewer_url"] = vu +``` + +This requires adding `_read_session_safe` to `task_manager.py` — a wrapper that returns `None` instead of raising. Add it now in `coda_mcp/task_manager.py` next to `_read_session`: + +```python +def _read_session_safe(session_id: str) -> dict | None: + """Read session.json, returning None on missing/corrupt instead of raising.""" + try: + return _read_session(session_id) + except SessionNotFoundError: + return None +``` + +- [ ] **Step 5: Add `viewer_url` to `coda_get_result`** + +In `coda_get_result` (around line 327), after the existing field-setting block, add: + +```python + # Decorate with viewer_url if known + sess = task_manager._read_session_safe(session_id) + pty = sess.get("pty_session_id") if sess else None + if pty: + vu = url_builder.build_viewer_url(pty) + if vu: + result["viewer_url"] = vu +``` + +Place this immediately before `return json.dumps(result)`. + +- [ ] **Step 6: Update FastMCP `instructions`** + +In `coda_mcp/mcp_server.py`, modify the `instructions=` argument to FastMCP (around line 42) by appending a paragraph at the end of the existing instructions string: + +```python + "CHAINING: pass previous_session_id from a completed task's session_id " + "to give the new task context of what was done before.\n\n" + "SHARE THE LIVE URL: When coda_run returns a viewer_url field (non-null), " + "mention it to the user in plain text (e.g. \"you can watch progress at " + "\"). The URL is safe to share — it points to the same Databricks App " + "the user is already authenticated against. Do this on the first mention " + "of the task and any time the user asks where the task is or how to see it." +``` + +- [ ] **Step 7: Run tests to verify they pass** + +Run: `uv run pytest tests/test_mcp_server.py -v -k "viewer_url or transcript_path"` +Expected: 5 passed. + +- [ ] **Step 8: Run full mcp test suite for regression** + +Run: `uv run pytest tests/test_mcp_server.py tests/test_mcp_integration.py -v` +Expected: all pass. + +- [ ] **Step 9: Commit** + +```bash +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" \ + add coda_mcp/mcp_server.py coda_mcp/task_manager.py tests/test_mcp_server.py +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" \ + commit -m "feat(coda-mcp): return viewer_url from coda_run/inbox/get_result + transcript wiring" +``` + +--- + +## Task 8: `mcp_asgi.py` — capture `X-Forwarded-Host` via ASGI middleware + +**Files:** +- Modify: `coda_mcp/mcp_asgi.py` (add middleware class + register it on `mcp_starlette`) +- Test: `tests/test_app_url_middleware.py` (new) + +- [ ] **Step 1: Write the failing tests** + +Create `tests/test_app_url_middleware.py`: + +```python +"""Tests for AppUrlCaptureMiddleware — populates url_builder._app_url_cache.""" +import asyncio +import importlib + +import pytest + +from coda_mcp import url_builder + + +@pytest.fixture(autouse=True) +def _reset_cache(): + importlib.reload(url_builder) + yield + + +async def _fake_app(scope, receive, send): + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"", "more_body": False}) + + +def _make_scope(headers: list[tuple[bytes, bytes]]): + return { + "type": "http", + "asgi": {"version": "3.0"}, + "method": "POST", + "path": "/mcp", + "headers": headers, + } + + +async def _drive(middleware, scope): + sent = [] + async def send(msg): sent.append(msg) + async def receive(): return {"type": "http.request", "body": b"", "more_body": False} + await middleware(scope, receive, send) + + +def test_middleware_captures_x_forwarded_host(): + from coda_mcp.mcp_asgi import AppUrlCaptureMiddleware + mw = AppUrlCaptureMiddleware(_fake_app) + scope = _make_scope([(b"x-forwarded-host", b"app.databricksapps.com")]) + asyncio.run(_drive(mw, scope)) + assert url_builder._app_url_cache == "app.databricksapps.com" + + +def test_middleware_falls_back_to_host_when_no_xforwarded(): + from coda_mcp.mcp_asgi import AppUrlCaptureMiddleware + mw = AppUrlCaptureMiddleware(_fake_app) + scope = _make_scope([(b"host", b"localhost:8000")]) + asyncio.run(_drive(mw, scope)) + assert url_builder._app_url_cache == "localhost:8000" + + +def test_middleware_skips_non_http_scope(): + from coda_mcp.mcp_asgi import AppUrlCaptureMiddleware + mw = AppUrlCaptureMiddleware(_fake_app) + scope = {"type": "lifespan"} + async def receive(): return {"type": "lifespan.startup"} + sent = [] + async def send(msg): sent.append(msg) + # Must not crash. Cache stays None. + asyncio.run(mw(scope, receive, send)) + assert url_builder._app_url_cache is None + + +def test_middleware_no_op_when_no_host_header(): + from coda_mcp.mcp_asgi import AppUrlCaptureMiddleware + mw = AppUrlCaptureMiddleware(_fake_app) + scope = _make_scope([]) + asyncio.run(_drive(mw, scope)) + assert url_builder._app_url_cache is None +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `uv run pytest tests/test_app_url_middleware.py -v` +Expected: ImportError on `AppUrlCaptureMiddleware`. + +- [ ] **Step 3: Add the middleware class to `mcp_asgi.py`** + +At the top of `coda_mcp/mcp_asgi.py` (after imports, around line 28), add: + +```python +from coda_mcp import url_builder + + +class AppUrlCaptureMiddleware: + """Capture X-Forwarded-Host (or Host) from every inbound HTTP request and + populate url_builder._app_url_cache. Used so MCP tools can return a + working viewer_url without manual configuration. + + Caveat: /socket.io/ traffic is intercepted by socketio.ASGIApp *before* + reaching mcp_starlette, so WebSocket connect requests never hit this + middleware. This is fine in practice — every HTTP request to /mcp and to + Flask routes does hit it, which is enough to keep the cache hot. + """ + + def __init__(self, app): + self.app = app + + async def __call__(self, scope, receive, send): + if scope.get("type") == "http": + headers = dict(scope.get("headers") or []) + host_bytes = headers.get(b"x-forwarded-host") or headers.get(b"host") + if host_bytes: + try: + url_builder.capture_from_headers(host_bytes.decode("latin-1")) + except Exception: + pass + await self.app(scope, receive, send) +``` + +- [ ] **Step 4: Register the middleware on `mcp_starlette`** + +In the existing block that adds CORS (around lines 80-86): + +```python +# CORS for MCP and Flask routes +mcp_starlette.add_middleware( + CORSMiddleware, + allow_origins=ALLOWED_ORIGINS or ["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) +``` + +Add a second `add_middleware` call immediately after: + +```python +# Capture X-Forwarded-Host into url_builder cache (for MCP viewer_url). +# Added AFTER CORS so it wraps the CORS-handled request. +mcp_starlette.add_middleware(AppUrlCaptureMiddleware) +``` + +- [ ] **Step 5: Run tests to verify they pass** + +Run: `uv run pytest tests/test_app_url_middleware.py -v` +Expected: 4 passed. + +- [ ] **Step 6: Commit** + +```bash +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" \ + add coda_mcp/mcp_asgi.py tests/test_app_url_middleware.py +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" \ + commit -m "feat(coda-mcp): AppUrlCaptureMiddleware seeds url_builder from X-Forwarded-Host" +``` + +--- + +## Task 9: `app.py::attach_session` — replay fallback when PTY is gone + +**Files:** +- Modify: `app.py::attach_session` (lines ~1104-1123) +- Test: `tests/test_replay_attach.py` (new) + +- [ ] **Step 1: Write the failing tests** + +Create `tests/test_replay_attach.py`: + +```python +"""Tests for /api/session/attach replay fallback.""" +import json +import os +from pathlib import Path + +import pytest + +from coda_mcp import task_manager + + +@pytest.fixture +def client(tmp_path, monkeypatch): + monkeypatch.setattr(task_manager, "SESSIONS_DIR", str(tmp_path)) + monkeypatch.setenv("MAX_CONCURRENT_SESSIONS", "5") + from app import app + # Bypass authorization (single-user app pattern used by other tests) + monkeypatch.setattr("app.check_authorization", lambda: True) + with app.test_client() as c: + yield c, tmp_path + + +def _seed_transcript(sessions_root: Path, pty_id: str, content: bytes) -> None: + sess_id = "sess-test" + task_id = "task-test" + sdir = sessions_root / sess_id + tdir = sdir / "tasks" / task_id + tdir.mkdir(parents=True) + (sdir / "session.json").write_text(json.dumps({ + "session_id": sess_id, + "pty_session_id": pty_id, + "current_task": None, + "completed_tasks": [task_id], + "status": "closed", + })) + (tdir / "transcript.log").write_bytes(content) + + +def test_attach_returns_replay_when_pty_gone_and_transcript_exists(client): + c, root = client + _seed_transcript(root, "pty-gone", b"hello\r\nworld\r\n") + resp = c.post("/api/session/attach", json={"session_id": "pty-gone"}) + assert resp.status_code == 200 + data = resp.get_json() + assert data["replay"] is True + assert data["output"] == ["hello\r\nworld\r\n"] + assert data["label"] == "hermes-mcp (replay)" + + +def test_attach_404_when_pty_gone_and_no_transcript(client): + c, root = client + resp = c.post("/api/session/attach", json={"session_id": "pty-nope"}) + assert resp.status_code == 404 +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `uv run pytest tests/test_replay_attach.py -v` +Expected: replay test fails (no fallback); 404 test passes already. + +- [ ] **Step 3: Modify `attach_session`** + +In `app.py::attach_session` (around line 1104), replace the body with: + +```python +@app.route("/api/session/attach", methods=["POST"]) +def attach_session(): + """Reattach to an existing session — returns buffered output for replay. + + If the live PTY is gone but an on-disk transcript exists for this + pty_session_id, return the transcript as ``output`` with ``replay: True``. + """ + data = request.get_json(silent=True) or {} + session_id = data.get("session_id", "") + + sess = _get_session(session_id) + if not sess or sess.get("exited"): + # Replay fallback: look up transcript.log by pty_session_id + from coda_mcp import task_manager as _tm + tdir = _tm.find_task_dir_by_pty_session(session_id) + if tdir: + transcript = os.path.join(tdir, "transcript.log") + if os.path.isfile(transcript): + try: + with open(transcript, "rb") as f: + content = f.read() + return jsonify({ + "session_id": session_id, + "label": "hermes-mcp (replay)", + "output": [content.decode("utf-8", errors="replace")], + "replay": True, + "process": None, + "created_at": None, + }) + except OSError: + pass + return jsonify({"error": "Session not found or exited"}), 404 + + # Existing live-attach path + sess["last_poll_time"] = time.time() + return jsonify({ + "session_id": session_id, + "label": sess.get("label", ""), + "output": list(sess["output_buffer"]), + "process": _get_session_process(sess["pid"]), + "created_at": sess.get("created_at"), + }) +``` + +- [ ] **Step 4: Run tests to verify they pass** + +Run: `uv run pytest tests/test_replay_attach.py -v` +Expected: 2 passed. + +- [ ] **Step 5: Run regression for the existing session-attach tests** + +Run: `uv run pytest tests/test_session_detach.py -v` +Expected: all pass. + +- [ ] **Step 6: Commit** + +```bash +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" \ + add app.py tests/test_replay_attach.py +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" \ + commit -m "feat: attach_session replay fallback reads transcript.log when PTY is gone" +``` + +--- + +## Task 10: `static/index.html` — boot URL parse + `_doReplay` + history hygiene + +**Files:** +- Modify: `static/index.html` + +> **Note**: This is the most "real" change. We add ~50-70 LoC of JS. Tested manually (Playwright not configured in this repo). + +- [ ] **Step 1: Locate the SPA boot path** + +Read `static/index.html` lines 990-1030 (the existing session-picker boot logic) to confirm where pane creation happens after the picker. The new URL-driven branch must run before the picker. + +- [ ] **Step 2: Add boot-time URL parse** + +Find the existing function that runs on `DOMContentLoaded` or the IIFE that initializes the app. Just before it would invoke the session picker, insert: + +```javascript + // ── Deep-link to a CoDA MCP session via ?session= ── + async function _initFromQueryString() { + const params = new URLSearchParams(location.search); + const sessionId = params.get('session'); + if (!sessionId) return false; + + try { + const resp = await fetch('/api/session/attach', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ session_id: sessionId }) + }); + + if (resp.status === 404) { + _renderExpiredPage(sessionId); + return true; // handled, skip picker + } + + const data = await resp.json(); + const term = createTerminalPane({ sessionId, label: data.label || sessionId }); + + if (data.replay) { + const content = (data.output || []).join(''); + await _doReplay(term, sessionId, content); + } else { + await _doAttach(term, sessionId); + if (typeof socket !== 'undefined' && socket) { + socket.emit('join_session', { session_id: sessionId }); + } + } + + return true; // handled, skip picker + } catch (err) { + console.error('deep-link attach failed:', err); + return false; + } + } +``` + +`createTerminalPane({ sessionId, label })` is the name commonly used in this repo for pane creation; if the actual name differs, substitute the local helper. Read the existing pane creation site to confirm and adjust the call site accordingly. + +- [ ] **Step 3: Add `_doReplay`** + +Place near `_doAttach` (around line 1339): + +```javascript + async function _doReplay(term, sessionId, content) { + // Chunk the write to avoid main-thread jank on multi-MB transcripts. + const CHUNK = 64 * 1024; + for (let i = 0; i < content.length; i += CHUNK) { + term.write(content.slice(i, i + CHUNK)); + await new Promise(r => requestAnimationFrame(r)); + } + // Mount a static banner above the pane. + _showReplayBanner(term, sessionId); + // NOTE: do NOT wire term.onData → terminal_input; do NOT include in heartbeat + // session_ids list; do NOT emit join_session. + return sessionId; + } + + function _showReplayBanner(term, sessionId) { + const pane = getAllPanes().find(p => p.sessionId === sessionId); + if (!pane || !pane.element) return; + const banner = document.createElement('div'); + banner.className = 'replay-banner'; + banner.textContent = 'Task completed — viewing replay'; + banner.style.cssText = 'padding:4px 8px;background:#333;color:#aaa;font-size:12px;text-align:center;'; + pane.element.insertBefore(banner, pane.element.firstChild); + } +``` + +- [ ] **Step 4: Add `_renderExpiredPage`** + +Place near `_doReplay`: + +```javascript + function _renderExpiredPage(sessionId) { + const root = document.body; + root.innerHTML = ` +
+

Session expired

+

Session ${sessionId.replace(/[<>]/g, '')} is gone, and no replay is available.

+

The transcript may have aged out after the 24-hour retention window.

+

← Back to terminal

+
+ `; + } +``` + +- [ ] **Step 5: Wire `_initFromQueryString` into the boot path** + +Find where the existing session-picker is shown after `DOMContentLoaded`. Wrap it: + +```javascript + document.addEventListener('DOMContentLoaded', async () => { + // existing init code (sockets, themes, etc.) + + const handled = await _initFromQueryString(); + if (handled) return; + + // existing flow (show session picker, etc.) + }); +``` + +The exact insertion site depends on the existing boot structure — read lines 990-1050 of `static/index.html` to find the right place. + +- [ ] **Step 6: Add history hygiene on pane close** + +Locate the existing pane-close handler. Inside, after the pane is removed, add: + +```javascript + // If this pane was opened via ?session=, drop the query param so a + // refresh doesn't re-attach to a stale id. + const params = new URLSearchParams(location.search); + if (params.get('session') === pane.sessionId) { + history.replaceState({}, '', '/'); + } +``` + +- [ ] **Step 7: Manual smoke test** + +Local dev: + +```bash +uv run uvicorn coda_mcp.mcp_asgi:app --host 0.0.0.0 --port 8000 +``` + +Then open `http://localhost:8000/?session=fake-id` in a browser. Expected: "Session expired" page (404 since no transcript exists). + +Create a fake live session via the regular UI, note its session_id from the picker, then navigate to `http://localhost:8000/?session=` — expected: terminal opens directly attached to that session. + +- [ ] **Step 8: Commit** + +```bash +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" \ + add static/index.html +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" \ + commit -m "feat(spa): deep-link ?session= with live attach + replay rendering" +``` + +--- + +## Task 11: Integration test — E2E grace period + transcript replay + +**Files:** +- Modify: `tests/test_mcp_integration.py` (extend) + +- [ ] **Step 1: Write the failing test** + +Append to `tests/test_mcp_integration.py`: + +```python +import asyncio +import json +import os +import time +from pathlib import Path +from unittest import mock + +import pytest + +from coda_mcp import mcp_server, task_manager, url_builder + + +@pytest.fixture +def mcp_env(tmp_path, monkeypatch): + monkeypatch.setattr(task_manager, "SESSIONS_DIR", str(tmp_path)) + monkeypatch.setattr(url_builder, "_app_url_cache", "app.example.com") + # Shrink grace for the test + monkeypatch.setattr(mcp_server, "GRACE_PERIOD_S", 2) + return tmp_path + + +def test_end_to_end_grace_and_replay(mcp_env, monkeypatch): + """Stub hermes via direct file I/O, then exercise the full coda_run flow.""" + from app import mcp_create_pty_session, mcp_send_input, mcp_close_pty_session + from app import _mark_grace_for_session, _bump_session_last_poll, sessions + + mcp_server.set_app_hooks( + mcp_create_pty_session, mcp_send_input, mcp_close_pty_session, + _mark_grace_for_session, _bump_session_last_poll, + ) + + # Submit a fake task + result_json = asyncio.run(mcp_server.coda_run( + prompt="test", email="u@x", timeout_s=5, + )) + result = json.loads(result_json) + assert result["status"] == "running" + sess_id = result["session_id"] + task_id = result["task_id"] + pty_id = task_manager._read_session(sess_id)["pty_session_id"] + + # viewer_url returned + assert pty_id in result["viewer_url"] + + # Simulate hermes writing to the PTY by sending input that echoes to bash + mcp_send_input(pty_id, "echo HELLO_FROM_HERMES\n") + time.sleep(0.5) + + # Now simulate hermes completion by writing result.json + tdir = task_manager._task_dir(sess_id, task_id) + Path(tdir).joinpath("result.json").write_text(json.dumps({ + "status": "completed", "summary": "stub", "files_changed": [], + "artifacts": {}, "errors": [], + })) + + # Wait for watcher to pick it up (polls every 5s — shorten via patch below if slow) + # In practice, the test patches the poll interval. For now, manually invoke: + mcp_server._schedule_deferred_close(sess_id) + + # PTY still alive immediately after grace scheduling + assert pty_id in sessions + assert sessions[pty_id]["grace"] is True + + # Wait past GRACE_PERIOD_S + time.sleep(2.5) + + # PTY now gone + assert pty_id not in sessions + + # Transcript file exists and contains the echoed line + transcript = Path(tdir) / "transcript.log" + assert transcript.exists() + assert b"HELLO_FROM_HERMES" in transcript.read_bytes() + + # find_task_dir_by_pty_session now returns the task dir from the on-disk record + found = task_manager.find_task_dir_by_pty_session(pty_id) + assert found == str(tdir) +``` + +- [ ] **Step 2: Run the test** + +Run: `uv run pytest tests/test_mcp_integration.py -v -k end_to_end_grace_and_replay` +Expected: pass. + +- [ ] **Step 3: Run the full test suite for regression** + +Run: `uv run pytest tests/ -v --timeout=60` +Expected: prior pass count + the new tests. No failures. + +- [ ] **Step 4: Commit** + +```bash +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" \ + add tests/test_mcp_integration.py +git -c user.email=datasciencemonkey@gmail.com -c user.name="Sathish Gangichetty" \ + commit -m "test: E2E coverage for grace period + transcript replay" +``` + +--- + +## Task 12: Manual smoke + deployment verification + +**Files:** none (verification only) + +- [ ] **Step 1: Deploy the worktree to the test app** + +From the worktree root: + +```bash +databricks bundle deploy --target test-coda +``` + +(Adjust target name to whatever the existing deployment uses — check `databricks.yml` or `app.yaml` notes.) + +- [ ] **Step 2: Verify in Genie Code** + +In the Databricks workspace, open Genie Code, ensure the Custom MCP server `mcp-test-coda` is connected. Submit a simple task: `"List the files in /tmp"`. + +Expected: +- Genie Code's response mentions a `viewer_url` like `https://mcp-test-coda-.aws.databricksapps.com/?session=`. +- Clicking the URL opens the terminal pre-attached to that session. +- Hermes output streams in real time. + +- [ ] **Step 3: Verify replay** + +After the task completes, wait 6+ minutes (grace period + buffer), then reload the same URL. + +Expected: +- Page loads showing the static transcript of what hermes did. +- "Task completed — viewing replay" banner. +- No input is sent when you type. + +- [ ] **Step 4: Verify chmod on transcript** + +From a shell in the deployed app (workspace terminal or `databricks workspace files` API): + +```bash +ls -la ~/.coda/sessions/*/tasks/*/transcript.log +``` + +Expected: files have mode `-rw-------` (0o600). + +- [ ] **Step 5: Verify `viewer_url` absence locally without env** + +```bash +unset CODA_APP_URL +uv run uvicorn coda_mcp.mcp_asgi:app --host 0.0.0.0 --port 8000 & +SERVER_PID=$! + +# Submit a coda_run via curl-formatted JSON-RPC +curl -s http://localhost:8000/mcp \ + -H 'Content-Type: application/json' \ + -d '{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"coda_run","arguments":{"prompt":"test","email":"local@dev"}}}' + +kill $SERVER_PID +``` + +Expected: the JSON response contains `"viewer_url": "http://localhost:8000/?session=..."` (because the inbound `Host: localhost:8000` was captured). + +- [ ] **Step 6: Final commit (if any verification turned up a fix)** + +If smoke tests revealed issues, fix them as separate commits, then update this checklist. + +--- + +## Self-review notes + +- All eight spec decisions covered: §1 viewer mode → Task 10 `_doReplay`; §2 transcript tee → Tasks 3-4; §3 deferred Timer → Task 6; §4 grace exemption → Task 5; §5 URL form → Tasks 1, 7; §6 ASGI middleware → Task 8; §7 attach replay fallback → Task 9; §8 SPA → Task 10. +- No "TODO" / "TBD" / "implement later" / placeholder text — every step has concrete code, exact paths, exact commands. +- Type/method consistency: + - `set_app_hooks` signature in Task 6 matches the call site updated in Task 11 (`mcp_server.set_app_hooks(create, send, close, mark_grace, bump_poll)` with optional defaults). + - `_mark_grace_for_session` / `_bump_session_last_poll` defined in Task 5 used by Task 6 and Task 11. + - `transcript_path` kwarg added to `mcp_create_pty_session` in Task 4 used by `coda_run` in Task 7. + - `find_task_dir_by_pty_session` defined in Task 2 used by `attach_session` in Task 9. + - `url_builder.build_viewer_url` defined in Task 1 used by `coda_run`/`coda_inbox`/`coda_get_result` in Task 7. +- Spec §3 "Architecture" diagram preserved as the mental model; data flows §5.1-5.4 map to Tasks 7, 9, 6, 9 respectively. +- Risks §9 (secrets, grace race, multi-tab) accepted in the spec; surface in the test plan via the chmod-600 verification in Task 12 step 4. diff --git a/docs/superpowers/specs/2026-05-27-coda-mcp-live-session-url-design.md b/docs/superpowers/specs/2026-05-27-coda-mcp-live-session-url-design.md new file mode 100644 index 0000000..c82bdc6 --- /dev/null +++ b/docs/superpowers/specs/2026-05-27-coda-mcp-live-session-url-design.md @@ -0,0 +1,447 @@ +# CoDA MCP Live Session URL — Design + +**Date:** 2026-05-27 +**Branch:** `feat/coda-mcp-server` +**Status:** Spec approved by user; ready for implementation plan +**Related PR:** databrickslabs/coding-agents-databricks-apps#64 (parent feature) + +## 1. Problem + +`coda_run` is fire-and-forget today: it returns `{task_id, session_id, status: "running"}` and the calling MCP client (Genie Code, Claude Desktop, Cursor) has no way to surface progress to the user. The user only sees a structured `result.json` after the task completes via `coda_inbox`/`coda_get_result`. Status messages from `status.jsonl` are coarse-grained. There is no way to watch hermes execute live, intervene mid-task, or reconstruct what happened after the fact. + +The Flask app side already has a fully working real-time terminal UI (xterm.js + Socket.IO + HTTP polling fallback) that knows how to attach to any active PTY by id. The MCP server already spawns those PTYs to run hermes. **The two halves are not connected by a URL.** + +## 2. Goal + +Give every `coda_run` (and existing tasks listed via `coda_inbox` / fetched via `coda_get_result`) a `viewer_url` that: + +- **During execution** — opens the existing terminal UI attached to that task's live PTY. The user can watch hermes work in real time and type into the session if they want to redirect or take over (single-user app; this is intentional). +- **For ~5 minutes after completion** — keeps the PTY alive so a viewer who joined mid-task isn't yanked the instant `result.json` is written. Heartbeats from an active viewer do not extend this window — the grace timer is fixed. +- **Indefinitely after PTY closes** (within the 24h `TASK_TTL_S`) — serves a static "replay" rendering of the captured terminal transcript so a user can scroll the full execution history from `coda_inbox`. + +Out of scope (deferred to separate specs): configurable agent selection (hermes vs claude-code vs codex), multi-user attribution, asciinema-style timed replay. + +## 3. Architecture + +``` +┌────────────────────────────────────────────────────────────────┐ +│ MCP client Browser │ +│ (Genie Code, Claude Desktop) (single user, app URL) │ +└──────────┬──────────────────────────────────┬──────────────────┘ + │ tools/call coda_run │ GET /?session= + ▼ ▼ + ┌───────────────┐ ┌─────────────────────┐ + │ coda_mcp /mcp │ │ Flask /static + WS │ + │ +viewer_url │ │ /api/session/attach│ + └───────┬───────┘ └──────────┬──────────┘ + │ │ + ▼ ▼ + ┌──────────────────────────────────────────────────────┐ + │ Flask app (single process) │ + │ sessions[] → {fd, buffer, transcript_fh, │ + │ grace: bool} │ + │ read_pty_output thread: │ + │ fd → buffer → socketio emit (room=) │ + │ fd → transcript.log (NEW: tee, flush per write) │ + └──────────────────────────────────────────────────────┘ + │ │ + │ writes (chmod 600) │ reads when PTY gone + ▼ ▼ + ~/.coda/sessions/{sess}/tasks/{task}/transcript.log +``` + +Everything between the MCP server and the Flask app already exists. The feature is mostly plumbing: + +1. **Tee PTY output** to `transcript.log` (on disk, per task, chmod 0600, 10 MB soft cap). +2. **Defer PTY close** on task completion by 5 minutes (`threading.Timer`) so live viewers can finish reading. +3. **Build `viewer_url`** in MCP tool responses by capturing `X-Forwarded-Host` from the inbound request. +4. **Teach the SPA** to read `?session=` on load and to render replay mode when the PTY is gone but a transcript exists. + +## 4. Components + +### 4.1 `app.py::sessions[pty_id]` dict (additive) + +Four new keys, all optional/defaulting: + +- `transcript_path: str | None` — absolute path to the tee target. +- `transcript_fh: BinaryIO | None` — open file handle owned by `read_pty_output`. +- `transcript_bytes: int` (default 0) — running count to enforce the 10 MB cap. +- `grace: bool` (default False) — set `True` when `_watch_task` schedules deferred close. Used by the concurrency check to exempt this slot. + +No removals. No semantic changes to existing keys. + +### 4.2 `app.py::mcp_create_pty_session(label, transcript_path=None)` + +New optional kwarg. When provided: + +- `os.makedirs(os.path.dirname(transcript_path), exist_ok=True)` +- Open file: `fh = open(transcript_path, "ab", buffering=0)` (binary append, unbuffered) +- `os.fchmod(fh.fileno(), 0o600)` immediately +- Store `transcript_path` and `transcript_fh` on the session dict +- If open fails: log error, set both to `None`, continue (live PTY still works) + +### 4.3 `app.py::read_pty_output` (additive) + +After the existing buffer append and Socket.IO emit, if a transcript handle is present, write under the per-session lock to prevent races against `terminate_session` (which may close the handle from the Timer thread): + +```python +with session_lock: + fh = session.get("transcript_fh") + written = session.get("transcript_bytes", 0) + if fh is not None: + remaining = TRANSCRIPT_CAP_BYTES - written + if remaining > 0: + chunk = output[:remaining] + try: + fh.write(chunk) + fh.flush() + session["transcript_bytes"] = written + len(chunk) + if len(chunk) < len(output): + fh.write(b"\n[transcript truncated at 10MB]\n") + fh.flush() + fh.close() + session["transcript_fh"] = None + except (OSError, ValueError) as exc: + logger.warning("transcript write failed for %s: %s", session_id, exc) + try: fh.close() + except Exception: pass + session["transcript_fh"] = None +``` + +`TRANSCRIPT_CAP_BYTES = 10 * 1024 * 1024`. + +**Invariants** (documented for future maintainers): + +- `transcript_fh` is opened in `mcp_create_pty_session`, written exclusively by `read_pty_output`, and closed by either (a) `read_pty_output` on cap/error or (b) `terminate_session` on PTY teardown. All three sites operate under `session["lock"]`. +- `transcript_bytes` is incremented only by `read_pty_output`. Single-writer; reads from other threads must hold `session["lock"]`. +- `ValueError` is caught alongside `OSError` to defend against a tiny window where `terminate_session` closes the handle between the spec's `if fh is not None` check and the actual `fh.write` call — the lock prevents this, but the catch is belt-and-suspenders. + +### 4.4 `app.py::terminate_session` (additive) + +Close the transcript file handle under the per-session lock before the existing fd close. The swap-to-`None` is the synchronization point that lets `read_pty_output` notice the handle is gone on its next iteration: + +```python +sess = sessions.get(session_id) +if sess is not None: + with sess["lock"]: + fh = sess.get("transcript_fh") + sess["transcript_fh"] = None # swap first, then close + if fh is not None: + try: fh.close() + except Exception: pass +``` + +(The actual close happens outside the lock to avoid holding it across a potential blocking I/O on a slow filesystem.) + +### 4.5 `app.py::MAX_CONCURRENT_SESSIONS` check (modified) + +At the `if len(sessions) >= MAX_CONCURRENT_SESSIONS` checkpoints in `create_session()` and `mcp_create_pty_session()`, replace the raw length check with a filtered count that excludes grace-period PTYs: + +```python +active = sum(1 for s in sessions.values() if not s.get("grace")) +if active >= MAX_CONCURRENT_SESSIONS: ... +``` + +`cleanup_stale_sessions` itself is **unchanged** — it still treats grace-period PTYs like any other session, but the 24h `SESSION_TIMEOUT_SECONDS` is so long the reaper never wins the race against the 5-min Timer. + +`MAX_CONCURRENT_SESSIONS` default stays at 5. + +### 4.6 `coda_mcp/mcp_server.py::_watch_task` (modified) + +Both completion and timeout paths replace immediate `_close_pty_for_session(session_id)` with: + +```python +session_data = task_manager._read_session(session_id) +pty_session_id = session_data.get("pty_session_id") +if pty_session_id and _app_close_session is not None: + _mark_grace(pty_session_id) # sets sessions[pty_id]["grace"] = True + _bump_last_poll(pty_session_id, GRACE_PERIOD_S) # defensive against reaper + threading.Timer( + GRACE_PERIOD_S, + _app_close_session, + args=(pty_session_id,), + ).start() +``` + +`GRACE_PERIOD_S = 300` (5 minutes), defined as a module constant for testability. `_mark_grace` and `_bump_last_poll` are two new hook callbacks wired through `set_app_hooks()` alongside the existing three — consistent with the current pattern (no direct Flask imports from the MCP module). + +The Timer must be a daemon so it doesn't block uvicorn shutdown: `t = threading.Timer(...); t.daemon = True; t.start()`. + +### 4.7 `coda_mcp/mcp_server.py::coda_run` (additive) + +After `mcp_create_pty_session`, compute the transcript path and pass it in: + +```python +transcript_path = os.path.join( + task_manager._task_dir(session_id, task_id), + "transcript.log", +) +pty_session_id = _app_create_session( + label="hermes-mcp", + transcript_path=transcript_path, +) +``` + +(Note: `_app_create_session` signature gains the kwarg. The implementation in `app.py` already documented above.) + +Then build the response with the new field: + +```python +return json.dumps({ + "task_id": task_id, + "session_id": session_id, + "status": "running", + "viewer_url": _build_viewer_url(pty_session_id), # may be None +}) +``` + +Tools serialize via `json.dumps` so `None` becomes `null`. Clients that don't recognize the field will ignore it. + +### 4.8 `coda_mcp/url_builder.py` (new tiny module) + +```python +import os +from typing import Optional + +_app_url_cache: Optional[str] = None + +def capture_from_headers(host: Optional[str]) -> None: + """Called by middleware on every inbound request.""" + global _app_url_cache + if host: + _app_url_cache = host + +def build_viewer_url(pty_session_id: str) -> Optional[str]: + override = os.environ.get("CODA_APP_URL", "").strip() + if override: + base = override.rstrip("/") + elif _app_url_cache: + base = f"https://{_app_url_cache}" + else: + return None + return f"{base}/?session={pty_session_id}" +``` + +### 4.9 `coda_mcp/mcp_asgi.py` (additive middleware) + +Insert a small ASGI middleware on `mcp_starlette` (via `mcp_starlette.add_middleware(...)`) that extracts `X-Forwarded-Host` (fallback: `Host`) from every HTTP request and calls `url_builder.capture_from_headers(host)`. Both MCP requests AND inbound browser HTTP requests refresh the cache. + +**Coverage caveat** (not a problem in practice): the top-level ASGI app is `socketio.ASGIApp(sio, other_asgi_app=mcp_starlette)`, so `/socket.io/` traffic is intercepted by socketio *before* it reaches `mcp_starlette` and therefore never hits this middleware. This is fine because (a) the user always loads the SPA via plain HTTP first (which refreshes the cache), and (b) every `coda_run` MCP call is a plain HTTP POST to `/mcp` (also through the middleware). The cache is hot by the time any tool needs the URL. + +```python +class AppUrlCaptureMiddleware: + def __init__(self, app): self.app = app + async def __call__(self, scope, receive, send): + if scope["type"] == "http": + headers = dict(scope.get("headers") or []) + host = headers.get(b"x-forwarded-host") or headers.get(b"host") + if host: + url_builder.capture_from_headers(host.decode()) + await self.app(scope, receive, send) +``` + +### 4.10 `coda_mcp/task_manager.py::find_task_dir_by_pty_session` (new) + +```python +_pty_lookup_cache: dict[str, tuple[str, float]] = {} # pty_id -> (task_dir, ts) +_PTY_LOOKUP_TTL = 60.0 # seconds + +def find_task_dir_by_pty_session(pty_session_id: str) -> str | None: + """Find the task dir whose session.json carries this pty_session_id.""" + now = time.time() + cached = _pty_lookup_cache.get(pty_session_id) + if cached and (now - cached[1]) < _PTY_LOOKUP_TTL: + return cached[0] + # Scan SESSIONS_DIR + if not os.path.isdir(SESSIONS_DIR): + return None + for sess_name in os.listdir(SESSIONS_DIR): + sess_file = os.path.join(SESSIONS_DIR, sess_name, "session.json") + try: + with open(sess_file) as f: + data = json.load(f) + except (OSError, json.JSONDecodeError): + continue + if data.get("pty_session_id") != pty_session_id: + continue + # The session has a current_task or completed_tasks; pick the most recent. + candidate = data.get("current_task") or ( + data["completed_tasks"][-1] if data.get("completed_tasks") else None + ) + if candidate: + tdir = os.path.join(SESSIONS_DIR, sess_name, "tasks", candidate) + _pty_lookup_cache[pty_session_id] = (tdir, now) + return tdir + return None +``` + +TTL handles the rename/close case without manual invalidation. + +**Invariant**: CoDA MCP sessions are ephemeral — one task per session (see `task_manager.create_session` then `complete_task` which sets `current_task=None` and appends to `completed_tasks`). This function therefore returns the right task dir for the lifetime of the URL. If the lifecycle ever changes to allow task reuse within a single session, this function must be revisited to pick the *active or grace-period* task rather than `completed_tasks[-1]`. + +### 4.11 `app.py::attach_session` endpoint (additive) + +After the existing `_get_session()` lookup, add a fallback: + +```python +sess = _get_session(session_id) +if not sess or sess.get("exited"): + # NEW: try transcript replay + tdir = task_manager.find_task_dir_by_pty_session(session_id) + if tdir: + transcript = os.path.join(tdir, "transcript.log") + if os.path.isfile(transcript): + with open(transcript, "rb") as f: + content = f.read() + return jsonify({ + "session_id": session_id, + "label": "hermes-mcp (replay)", + "output": [content.decode("utf-8", errors="replace")], + "replay": True, + "process": None, + "created_at": None, + }) + return jsonify({"error": "Session not found or exited"}), 404 +``` + +The response shape (`output: [str]`, `replay: true|absent`, plus existing keys) is **NOT** consumed by the existing `_doAttach` — that function deliberately ignores `data.output` and forces a SIGWINCH redraw of the live application (`static/index.html:1339-1357`, comment at line 1347: "We skip buffer replay because it contains raw escape sequences that produce garbled output"). The replay-mode response is consumed by a new SPA function `_doReplay` described in §4.12, which writes the bytes directly into xterm. + +### 4.12 `static/index.html` (~50-70 LoC) + +Four additions: + +1. **Boot-time URL parse** — before the existing session-picker fetch, check `new URLSearchParams(location.search).get("session")`. If absent → existing flow. If present → call `POST /api/session/attach` once and branch on the response: + - 200 with `replay: true` → call **`_doReplay`** (new, described below). Skip `_doAttach`. Do NOT emit `join_session`. Do NOT wire `terminal_input` to the WS. + - 200 without `replay` → call the existing `_doAttach(term, sessionId)` and the existing `socket.emit('join_session', { session_id })` path. (Reusing `_doAttach` is correct here because the *live* PTY is running an interactive app, and SIGWINCH-redraw is the right behavior.) + - 404 → render a small in-page fallback: "session expired or never existed" + a button to navigate to `/`. + +2. **`_doReplay(term, sessionId, bytes)` — new function** that handles static replay rendering. Cannot route through `_doAttach` because `_doAttach` discards `data.output` (it relies on a running app to redraw via SIGWINCH; replay mode has no running app). Implementation: + + ```js + async function _doReplay(term, sessionId, content) { + // Chunk the write to avoid main-thread jank on multi-MB transcripts. + // xterm.js write() is internally batched, but a single 10MB call + // still blocks until the parser drains. 64KB slices with rAF gives + // the browser a chance to repaint between chunks. + const CHUNK = 64 * 1024; + for (let i = 0; i < content.length; i += CHUNK) { + term.write(content.slice(i, i + CHUNK)); + await new Promise(r => requestAnimationFrame(r)); + } + // Mount a small "Task completed — viewing replay" banner above the pane. + // No input handler, no WS subscription, no heartbeat for this session id. + } + ``` + +3. **Replay-mode pane behavior** — the tab gets a "(replay)" badge. The xterm input handler is not wired. The session is NOT included in the heartbeat session_ids list (the PTY is dead; heartbeats would 404 the lookup). + +4. **History/URL hygiene** — when the user closes a pane that was opened via `?session=`, call `history.replaceState({}, '', '/')` so a refresh doesn't re-attach. + +**Estimate revised**: 50-70 LoC including the new `_doReplay` and the 404 fallback. Architecturally the most "real" change in the spec — the rest of the codebase shifts are mostly additive. + +### 4.13 MCP tool `instructions` update (`coda_mcp/mcp_server.py`) + +Append one paragraph to the existing `instructions` block on the FastMCP instance: + +> SHARE THE LIVE URL: When `coda_run` returns a `viewer_url` field, mention it to the user in plain text (e.g. "you can watch progress at "). The URL is safe to share — it points to the same Databricks App the user is already authenticated against. Do this on the FIRST mention of the task and any time the user asks where the task is or how to see it. + +## 5. Data flow + +### 5.1 Submit + +`MCP client → /mcp coda_run → task_manager.create_session → mcp_create_pty_session(transcript_path) → task_manager.create_task → mcp_send_input("hermes -z ...") → _watch_task thread spawned → return {task_id, session_id, status: "running", viewer_url}`. + +### 5.2 Live view + +`Browser → GET /?session= → SPA reads ?session → POST /api/session/attach → live output buffer returned → WS join_session → live stream from read_pty_output → terminal_input writes to fd → heartbeat keeps the (already non-grace) PTY alive`. + +### 5.3 Grace window + +At T+0 hermes writes `result.json`. `_watch_task` calls `task_manager.complete_task` (disk status → closed), marks the PTY `grace=True`, bumps `last_poll_time`, schedules `Timer(300, _app_close_session)`. A viewer present at T+0 keeps streaming for up to 5 min. At T+300 the Timer SIGHUPs bash, `read_pty_output` sees EOF, flushes and closes the transcript handle, removes the session entry. + +### 5.4 Replay + +`Browser → GET /?session= → POST /api/session/attach → PTY not found → find_task_dir_by_pty_session → read transcript.log → return {output: [bytes], replay: true} → SPA renders bytes, no WS subscription`. + +## 6. Error handling + +| Failure | Behavior | +|---|---| +| `CODA_APP_URL` and `X-Forwarded-Host` both absent | `viewer_url: null`. One startup WARN. | +| Transcript open fails | `transcript_fh = None`. Live PTY works; replay disabled. | +| Transcript write fails mid-stream | Log once per session, close handle, set `transcript_fh = None`, keep reading PTY. | +| 10 MB cap hit | Write marker, close handle, set `transcript_fh = None`. PTY keeps streaming live (no further teeing). | +| Timer fires after manual close | `terminate_session` is re-entrant; `sessions.pop(_, None)` and `os.kill` wrapped in try/except. No-op. | +| uvicorn restart during grace | In-memory state lost; old `viewer_url` falls through to transcript replay (if file exists) or 404. Acceptable. | +| Browser opens URL mid-grace, grace expires while connected | `read_pty_output` emits `session_exited` to the room. SPA shows "session ended" banner. User reloads → replay mode. | +| Browser opens URL after grace AND transcript reaped | 404. SPA shows expired page. | +| `MAX_CONCURRENT_TASKS` reached | Unchanged "concurrency limit" error. Grace PTYs don't count toward this (disk status = closed). | +| `MAX_CONCURRENT_SESSIONS` reached among active (non-grace) | Existing 429. Grace PTYs don't count. | +| Hermes hangs (no `result.json`) | Existing `_watch_task` timeout path now also defers close via the same Timer mechanism. | + +## 7. Testing + +### 7.1 Unit + +- `coda_mcp/url_builder.py`: env override beats header capture; `None` when both absent; trailing slash on override is stripped. +- `coda_run` returns `viewer_url` only when builder returns non-None; same for `coda_inbox` per-entry and `coda_get_result`. +- `find_task_dir_by_pty_session`: hit, miss, TTL expiry, ignores corrupt session.json. +- `_watch_task`: schedules `Timer` (mocked) with correct args on both completion and timeout paths; never calls `_app_close_session` synchronously. +- `_mark_grace` / `_bump_last_poll` set the session dict fields. + +### 7.2 Integration (`tests/test_mcp_integration.py`) + +- E2E with a stub hermes (`bash -c 'echo hello; touch results/result.json; echo done'`): + - `transcript.log` contains "hello". + - At T+1s, PTY still alive (grace). + - At T+(GRACE+1)s (test uses a 2s grace via patched constant), PTY closed; transcript file persists. + - `/api/session/attach` returns `replay: true` after close; live mode before. +- Concurrency: submit `MAX_CONCURRENT_TASKS` tasks, complete them all (grace begins), submit `MAX_CONCURRENT_TASKS` more — all succeed (grace PTYs don't block). +- 10 MB cap: feed a hermes stub that prints `>10MB` of output; transcript file is exactly `10MB + marker`; PTY keeps running. + +### 7.3 SPA + +- New `tests/test_frontend_deeplink.spec.js` (Playwright if available; else manual checklist): + - `/?session=` → live attach, WS room joined, terminal renders. + - `/?session=` → replay rendered, no WS join, banner visible. + - `/?session=` → expired page. + - Closing the pane drops `?session=` from `history`. + +### 7.4 Manual smoke + +- Deploy to `mcp-test-coda` app, connect Genie Code, run a `coda_run`, click `viewer_url` from the chat response, confirm live stream + grace + replay. +- `chmod 600` check: `ls -la ~/.coda/sessions/*/tasks/*/transcript.log` on deployed pod. +- Confirm `viewer_url` absent on a local uvicorn boot without `CODA_APP_URL` and no inbound request yet. + +## 8. Open questions (resolved) + +- ~~Read-only vs interactive viewer?~~ → Interactive (full terminal). +- ~~Grace period mechanism?~~ → `threading.Timer(300, _close)`. +- ~~Replay storage?~~ → Tee to `transcript.log`. +- ~~Configurable agent?~~ → Deferred to a separate spec. +- ~~Base URL resolution?~~ → `CODA_APP_URL` env override → `X-Forwarded-Host` capture (officially provided by Databricks Apps). +- ~~Concurrency under grace?~~ → Exempt grace PTYs from `MAX_CONCURRENT_SESSIONS`. Cap stays at 5. + +## 9. Risks accepted + +- **Transcript on disk contains secrets** if hermes prints them. Single-user app, file is mode 0600, cleaned with the rest of the session at 24h TTL. Documented in `docs/mcp-v2-background-execution.md`. +- **5 min grace + 0 second active task** means a viewer who opens the URL late may still race the close. Acceptable; replay mode covers them. +- **Browser tabs can interact with the same PTY simultaneously.** Already true for the existing terminal UI; no new exposure. + +## 10. Surface summary + +| Surface | LoC est | Risk | +|---|---|---| +| `app.py` (4 functions touched) | ~60 | Low — additive, no semantic shifts | +| `coda_mcp/mcp_server.py` (2 functions + instructions) | ~40 | Low | +| `coda_mcp/url_builder.py` (new) | ~25 | Low | +| `coda_mcp/mcp_asgi.py` (middleware) | ~15 | Low | +| `coda_mcp/task_manager.py` (new lookup) | ~30 | Low | +| `static/index.html` | ~50-70 | Medium — new boot branch + new `_doReplay` rendering path; live attach still reuses `_doAttach` | +| Tests | ~250 | — | + +**Total**: ~235-255 LoC of production code + ~250 LoC of tests. + +## 11. Next step + +Hand to `writing-plans` skill to produce an executable implementation plan with task ordering, dependencies, and verification gates. diff --git a/install_databricks_cli.sh b/scripts/install_databricks_cli.sh similarity index 100% rename from install_databricks_cli.sh rename to scripts/install_databricks_cli.sh diff --git a/install_gh.sh b/scripts/install_gh.sh similarity index 100% rename from install_gh.sh rename to scripts/install_gh.sh diff --git a/install_micro.sh b/scripts/install_micro.sh similarity index 100% rename from install_micro.sh rename to scripts/install_micro.sh diff --git a/setup_claude.py b/setup/setup_claude.py similarity index 83% rename from setup_claude.py rename to setup/setup_claude.py index 125393e..bbe56ed 100644 --- a/setup_claude.py +++ b/setup/setup_claude.py @@ -129,31 +129,35 @@ local_bin = home / ".local" / "bin" claude_bin = local_bin / "claude" -# Honour CLAUDE_INSTALLER_URL for enterprise environments where claude.ai is -# firewalled — defaults to the public installer when unset. The URL is -# validated by enterprise_config to reject shell metacharacters before it -# reaches subprocess. Additionally, we avoid embedding the URL in a shell -# string by piping curl's output into bash via positional args — even if a -# malicious URL somehow slipped through validation, it would land as a curl -# argument, not as shell. -from enterprise_config import claude_installer_url - -installer_url = claude_installer_url() -print(f"Installing/upgrading Claude Code CLI from {installer_url}...") -curl_proc = subprocess.Popen( - ["curl", "-fsSL", installer_url], - stdout=subprocess.PIPE, - env={**os.environ, "HOME": str(home)}, -) -result = subprocess.run( - ["bash"], - stdin=curl_proc.stdout, - env={**os.environ, "HOME": str(home)}, - capture_output=True, - text=True, -) -curl_proc.stdout.close() -curl_proc.wait() +if os.environ.get("SKIP_CLAUDE_INSTALL"): + print("SKIP_CLAUDE_INSTALL set — skipping CLI install") + result = type("R", (), {"returncode": 0, "stderr": ""})() +else: + # Honour CLAUDE_INSTALLER_URL for enterprise environments where claude.ai is + # firewalled — defaults to the public installer when unset. The URL is + # validated by enterprise_config to reject shell metacharacters before it + # reaches subprocess. Additionally, we avoid embedding the URL in a shell + # string by piping curl's output into bash via positional args — even if a + # malicious URL somehow slipped through validation, it would land as a curl + # argument, not as shell. + from enterprise_config import claude_installer_url + + installer_url = claude_installer_url() + print(f"Installing/upgrading Claude Code CLI from {installer_url}...") + curl_proc = subprocess.Popen( + ["curl", "-fsSL", installer_url], + stdout=subprocess.PIPE, + env={**os.environ, "HOME": str(home)}, + ) + result = subprocess.run( + ["bash"], + stdin=curl_proc.stdout, + env={**os.environ, "HOME": str(home)}, + capture_output=True, + text=True, + ) + curl_proc.stdout.close() + curl_proc.wait() if result.returncode == 0: print("Claude Code CLI installed successfully") else: diff --git a/setup_codex.py b/setup/setup_codex.py similarity index 100% rename from setup_codex.py rename to setup/setup_codex.py diff --git a/setup_databricks.py b/setup/setup_databricks.py similarity index 100% rename from setup_databricks.py rename to setup/setup_databricks.py diff --git a/setup_gemini.py b/setup/setup_gemini.py similarity index 100% rename from setup_gemini.py rename to setup/setup_gemini.py diff --git a/setup_hermes.py b/setup/setup_hermes.py similarity index 55% rename from setup_hermes.py rename to setup/setup_hermes.py index 599777e..d533aef 100644 --- a/setup_hermes.py +++ b/setup/setup_hermes.py @@ -241,6 +241,172 @@ def _run(cmd, **kwargs): cli_name="Hermes", ) +# 5b. Append CoDA orchestrator instructions to HERMES.md +CODA_ORCHESTRATOR_INSTRUCTIONS = """ + +## CoDA Constitution (NON-NEGOTIABLE) + +This is the single most important rule. It applies to you AND every sub-agent you delegate to. + +**NO DESTRUCTIVE ACTIONS on pre-existing assets.** Specifically: +- **NEVER delete** files, tables, jobs, notebooks, pipelines, or any resource that was NOT + created during the current session — unless you have EXPLICIT confirmation from the user + or upstream caller. +- **NEVER drop** database tables, schemas, or catalogs that existed before the task started. +- **NEVER overwrite** existing files without confirmation if the content would be lost. +- **NEVER run** destructive CLI commands (`rm -rf`, `databricks jobs delete`, `DROP TABLE`, etc.) + on assets you didn't create. + +**What IS allowed without confirmation:** +- Creating new files, tables, jobs, pipelines, notebooks — building is always OK. +- Modifying files you created during the session. +- Deleting temporary files or artifacts you created during the session. +- Iterating on work in progress — edit, refactor, rebuild freely. +- Overwriting files you created in this session. + +**When in doubt:** Report back to the upstream caller (Genie Code or the user) describing +what you want to delete and why, and ask for confirmation before proceeding. This applies +to you directly AND to any sub-agent you delegate to — pass this rule in every delegation prompt. + +## CoDA Orchestrator Role + +You are Hermes, the primary orchestrator inside **CoDA** (Coding Agents on Databricks Apps). +You are not just a chat assistant — you are the brain that receives tasks and decides how +to execute them, either directly or by delegating to specialized sub-agents. + +### Your Environment + +- You are running inside a Databricks App with full workspace access. +- The Databricks CLI is pre-configured: `databricks` commands work out of the box. +- Unity Catalog, Jobs, Workflows, Notebooks, MLflow — all accessible. +- Projects live at `~/projects/` and sync to `/Workspace/Users/{email}/` on git commit. +- You have 39 Databricks and workflow skills available. + +### Prior Session Context + +When your prompt includes a `PRIOR SESSION:` block, it means this task continues +work from a previous session. The prior session's results are stored on disk: + +``` +~/.coda/sessions/{previous_session_id}/tasks/*/result.json +``` + +**Read those result files** to understand what was done before. Each result.json contains: +- `summary` — what the prior task accomplished +- `files_changed` — which files were created or modified +- `artifacts` — job IDs, commit hashes, dashboard URLs, etc. + +Use this context to continue the work without asking the user to repeat themselves. + +### Sub-Agents Available + +You have three coding agents you can delegate work to. Choose the best one for each subtask: + +**Claude Code** — Deep work, complex implementations, orchestration +```bash +claude -p "your prompt here" --allowedTools "Read,Edit,Bash" --max-turns 50 +``` +- Best for: multi-step implementations, planning, debugging, code review +- Can spawn teams: assign roles, goals, and backstory to parallel workers +- Has access to all 39 skills (Databricks + workflow) +- Use `--max-turns` to bound execution, `--max-budget-usd` for cost control + +**Codex** — Fast edits, refactoring, structured transforms +```bash +codex -q "your prompt here" +``` +- Best for: quick code changes, targeted refactors, code review +- Lightweight and fast — use when the task is well-scoped + +**Gemini** — Research, documentation, large-context analysis +```bash +gemini -p "your prompt here" +``` +- Best for: broad codebase analysis, documentation generation, research tasks +- Large context window — good for understanding big codebases + +### How to Delegate + +1. **Assess the task.** Is it something you can handle directly, or does it need a specialist? +2. **Pick the right agent.** Match the task to the agent's strengths (see above). +3. **Be specific.** Give the sub-agent a clear, self-contained prompt with all context it needs. +4. **Collect results.** Read the sub-agent's output and incorporate it into your response. +5. **Chain when needed.** Plan with Claude, implement with Codex, review with Gemini. + +### For Complex Tasks — Use Claude Code Teams + +When a task is large enough to benefit from parallel work, use Claude Code's team capability: +```bash +claude -p "Create a team of 3 agents to: [task]. Agent 1 handles [X], Agent 2 handles [Y], Agent 3 handles [Z]. Coordinate and merge results." --allowedTools "Read,Edit,Bash" --max-turns 100 +``` + +### Ephemeral Session Model + +Each task runs in its own short-lived session. When the task completes, the session closes +automatically. You will NOT receive follow-up tasks in the same session. + +**What this means for you:** +- **Be self-contained.** Complete the entire task in one go — there is no "next message." +- **Read prior context if provided.** If the prompt has a `PRIOR SESSION:` block, read + those result files to understand what was done before. This is how task chaining works. +- **Write thorough results.** Your `result.json` is the only thing the next task (or the + user) will see. Include a clear summary, all files changed, and any artifacts created. +- **Don't rely on in-memory state.** Anything you want to persist must go to disk — + either in the result files, git commits, or the workspace. + +### Single-User Mode + +You are operating in **single-user mode**. Every task comes from the same person — the app owner. +This means: + +- **Learn their patterns.** Pay attention to how they work, what tools they prefer, what + coding style they use, and what kind of tasks they send. +- **Remember across tasks.** If they always work with certain tables, frameworks, or patterns, + carry that knowledge forward. Use your memory system to persist insights. +- **Be proactive.** If you notice patterns, suggest improvements: + - "I've noticed you frequently create similar pipelines — want me to template this?" + - "Based on your last 3 tasks, you might want to consider..." + - "This task is similar to what you asked last time. Should I reuse that approach?" +- **Adapt your communication style.** Match their level of detail preference, verbosity, + and technical depth. Some users want terse results, others want explanations. +- **Build a profile over time.** Track their preferred tools, common workflows, recurring + patterns, and pain points. The longer you work together, the better you should get. + +### Task Protocol (CODA-TASK Convention) + +When you receive a task wrapped in `---CODA-TASK---` markers, follow this protocol: + +1. **Read the envelope.** Extract task_id, session_id, user, context, and the actual task. +2. **Write progress.** As you work, append lines to `{results_dir}/status.jsonl`: + ```json + {"step": "planning", "message": "Analyzing task requirements"} + {"step": "delegating", "message": "Sending implementation to Claude Code"} + {"step": "complete", "message": "Pipeline created successfully"} + ``` +3. **Write result.** When done, write `{results_dir}/result.json`: + ```json + { + "status": "completed", + "summary": "One paragraph of what was done", + "files_changed": ["path/to/file1.py"], + "artifacts": {"job_id": "123", "commit": "abc123"}, + "errors": [] + } + ``` + IMPORTANT: `result.json` must be a FILE, not a directory. + +4. **If you delegate,** update `status.jsonl` with delegation steps so the caller can track + which sub-agent is doing what. +""" + +if hermes_md.exists(): + existing_content = hermes_md.read_text() + if "CoDA Orchestrator Role" not in existing_content: + hermes_md.write_text(existing_content + CODA_ORCHESTRATOR_INSTRUCTIONS) + print("CoDA orchestrator instructions appended to HERMES.md") + else: + print("CoDA orchestrator instructions already present in HERMES.md") + # 6. Create projects directory (parity with other agents) projects_dir = home / "projects" projects_dir.mkdir(exist_ok=True) diff --git a/setup_mlflow.py b/setup/setup_mlflow.py similarity index 100% rename from setup_mlflow.py rename to setup/setup_mlflow.py diff --git a/setup_opencode.py b/setup/setup_opencode.py similarity index 100% rename from setup_opencode.py rename to setup/setup_opencode.py diff --git a/setup_proxy.py b/setup/setup_proxy.py similarity index 100% rename from setup_proxy.py rename to setup/setup_proxy.py diff --git a/static/index.html b/static/index.html index 9f517a6..f5b0f2a 100644 --- a/static/index.html +++ b/static/index.html @@ -1010,7 +1010,10 @@

General

return; } - socket = io({ transports: ['websocket', 'polling'] }); + // Start with polling (HTTP) so Databricks proxy identity headers are present + // for auth, then upgrade to WebSocket transparently. Direct WebSocket-first + // fails because the proxy doesn't inject X-Forwarded-Email on WS upgrade. + socket = io({ transports: ['polling', 'websocket'] }); socket.on('connect', () => { // Check actual transport — Socket.IO reports connected=true even on long-polling @@ -1353,6 +1356,116 @@

General

return sessionId; } + // ── Deep-link helpers ───────────────────────────────────────────── + + async function _doReplay(term, sessionId, content) { + // Chunk the write to avoid main-thread jank on multi-MB transcripts. + const CHUNK = 64 * 1024; + for (let i = 0; i < content.length; i += CHUNK) { + term.write(content.slice(i, i + CHUNK)); + await new Promise(r => requestAnimationFrame(r)); + } + // Mount a static banner above the pane. + _showReplayBanner(term, sessionId); + // NOTE: do NOT wire term.onData → terminal_input; do NOT include in heartbeat + // session_ids list; do NOT emit join_session. + return sessionId; + } + + function _showReplayBanner(term, sessionId) { + const pane = getAllPanes().find(p => p.sessionId === sessionId); + if (!pane || !pane.element) return; + const banner = document.createElement('div'); + banner.className = 'replay-banner'; + banner.textContent = 'Task completed — viewing replay'; + banner.style.cssText = 'padding:4px 8px;background:#333;color:#aaa;font-size:12px;text-align:center;'; + pane.element.insertBefore(banner, pane.element.firstChild); + } + + function _renderExpiredPage(sessionId) { + // Use DOM construction instead of innerHTML interpolation to prevent XSS + // via crafted ?session= values. textContent escapes everything. + document.body.innerHTML = ''; // clear + + const wrap = document.createElement('div'); + wrap.style.cssText = 'font-family:monospace;padding:40px;text-align:center;color:#ccc;'; + + const heading = document.createElement('h2'); + heading.textContent = 'Session expired'; + wrap.appendChild(heading); + + const intro = document.createElement('p'); + intro.appendChild(document.createTextNode('Session ')); + const code = document.createElement('code'); + code.textContent = sessionId; // textContent escapes <>"'`& + intro.appendChild(code); + intro.appendChild(document.createTextNode(' is gone, and no replay is available.')); + wrap.appendChild(intro); + + const explain = document.createElement('p'); + explain.textContent = 'The transcript may have aged out after the 24-hour retention window.'; + wrap.appendChild(explain); + + const link = document.createElement('a'); + link.href = '/'; + link.style.color = '#6cf'; + link.textContent = '← Back to terminal'; + const linkPara = document.createElement('p'); + linkPara.appendChild(link); + wrap.appendChild(linkPara); + + document.body.appendChild(wrap); + } + + async function _initFromQueryString() { + const params = new URLSearchParams(location.search); + const sessionId = params.get('session'); + if (!sessionId) return false; + + try { + const resp = await fetch('/api/session/attach', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ session_id: sessionId }) + }); + + if (resp.status === 404) { + _renderExpiredPage(sessionId); + return true; // handled, skip picker + } + + const data = await resp.json(); + + // Create a tab that skips the session picker and uses our known session id. + const tab = await createTab({ deepLinkSessionId: sessionId }); + if (!tab || tab.panes.length === 0) return false; + const pane = tab.panes[0]; + const term = pane.term; + + if (data.replay) { + // Replay pane: static, read-only. createPane skipped onData/join_session + // wiring because of deepLinkSessionId, so we leave it that way. Keystrokes + // are ignored; nothing to clean up. + const content = (data.output || []).join(''); + await _doReplay(term, sessionId, content); + } else { + // Live pane: createPane skipped the default wiring, so we own it here. + await _doAttach(term, sessionId); + term.onData(d => sendInput(d, pane.sessionId)); + if (wsConnected && socket) { + socket.emit('join_session', { session_id: sessionId }); + } else { + pollWorker.postMessage({ type: 'start_poll', paneId: pane.id, sessionId: sessionId }); + } + } + + return true; // handled, skip picker + } catch (err) { + console.error('deep-link attach failed:', err); + return false; + } + } + function _formatAge(timestamp) { const seconds = Math.floor((Date.now() / 1000) - timestamp); if (seconds < 60) return 'just now'; @@ -1672,6 +1785,10 @@

General

await waitForSetup(); } var { sid, reattached } = await getOrPromptSession(term, tab.label, opts.skipPrompt); + } else if (opts.deepLinkSessionId) { + // Deep-link boot — session id is already known; skip picker entirely. + var sid = opts.deepLinkSessionId; + var reattached = true; } else if (!opts.newSession) { // PAT is valid, initial page load — check for existing sessions first. const setupResp2 = await fetch('/api/setup-status'); @@ -1715,13 +1832,19 @@

General

const pane = { id, element, term, fitAddon, searchAddon, sessionId: sid, batchWrite: createWriteBatcher(term) }; - term.onData(data => sendInput(data, pane.sessionId)); - // Join WebSocket room if connected; otherwise start HTTP polling (AC-11, AC-16) - if (wsConnected && socket) { - socket.emit('join_session', { session_id: sid }); - } else { - pollWorker.postMessage({ type: 'start_poll', paneId: id, sessionId: sid }); + // Deep-link panes own their own input wiring + transport joins from + // _initFromQueryString (so replay mode can stay read-only and live mode + // doesn't double-emit join_session). Skip the default wiring here. + if (!opts.deepLinkSessionId) { + term.onData(data => sendInput(data, pane.sessionId)); + + // Join WebSocket room if connected; otherwise start HTTP polling (AC-11, AC-16) + if (wsConnected && socket) { + socket.emit('join_session', { session_id: sid }); + } else { + pollWorker.postMessage({ type: 'start_poll', paneId: id, sessionId: sid }); + } } // Click to focus @@ -1806,6 +1929,13 @@

General

p.term.dispose(); }); + // If the tab contained a deep-linked pane, drop ?session= from the URL. + const _ctParams = new URLSearchParams(location.search); + const _ctSid = _ctParams.get('session'); + if (_ctSid && tab.panes.some(p => p.sessionId === _ctSid)) { + history.replaceState({}, '', '/'); + } + // Remove DOM tab.paneContainer.remove(); @@ -1959,7 +2089,17 @@

General

const ap = tab.panes.find(p => p.id === tab.activePaneId) || tab.panes[0]; if (!ap) return; + // Capture before cleanupPane() nulls pane.sessionId. + const _apSessionId = ap.sessionId; cleanupPane(ap); + + // If this pane was opened via ?session=, drop the query param so a + // refresh doesn't re-attach to a stale id. + const _cpParams = new URLSearchParams(location.search); + if (_apSessionId && _cpParams.get('session') === _apSessionId) { + history.replaceState({}, '', '/'); + } + ap.term.dispose(); ap.element.remove(); @@ -2263,7 +2403,11 @@

General

// The element is kept in the DOM for error reporting (see catch below). status.style.display = 'none'; - await createTab(); + // ── Deep-link: ?session= takes priority over the session picker ── + const deepLinkHandled = await _initFromQueryString(); + if (!deepLinkHandled) { + await createTab(); + } updateSessionBadge(); let resizeTimer; diff --git a/tests/test_app_url_middleware.py b/tests/test_app_url_middleware.py new file mode 100644 index 0000000..46ee7df --- /dev/null +++ b/tests/test_app_url_middleware.py @@ -0,0 +1,71 @@ +"""Tests for AppUrlCaptureMiddleware — populates url_builder._app_url_cache.""" +import asyncio +import importlib + +import pytest + +from coda_mcp import url_builder + + +@pytest.fixture(autouse=True) +def _reset_cache(): + importlib.reload(url_builder) + yield + + +async def _fake_app(scope, receive, send): + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"", "more_body": False}) + + +def _make_scope(headers: list[tuple[bytes, bytes]]): + return { + "type": "http", + "asgi": {"version": "3.0"}, + "method": "POST", + "path": "/mcp", + "headers": headers, + } + + +async def _drive(middleware, scope): + sent = [] + async def send(msg): sent.append(msg) + async def receive(): return {"type": "http.request", "body": b"", "more_body": False} + await middleware(scope, receive, send) + + +def test_middleware_captures_x_forwarded_host(): + from coda_mcp.mcp_asgi import AppUrlCaptureMiddleware + mw = AppUrlCaptureMiddleware(_fake_app) + scope = _make_scope([(b"x-forwarded-host", b"app.databricksapps.com")]) + asyncio.run(_drive(mw, scope)) + assert url_builder._app_url_cache == "app.databricksapps.com" + + +def test_middleware_falls_back_to_host_when_no_xforwarded(): + from coda_mcp.mcp_asgi import AppUrlCaptureMiddleware + mw = AppUrlCaptureMiddleware(_fake_app) + scope = _make_scope([(b"host", b"localhost:8000")]) + asyncio.run(_drive(mw, scope)) + assert url_builder._app_url_cache == "localhost:8000" + + +def test_middleware_skips_non_http_scope(): + from coda_mcp.mcp_asgi import AppUrlCaptureMiddleware + mw = AppUrlCaptureMiddleware(_fake_app) + scope = {"type": "lifespan"} + async def receive(): return {"type": "lifespan.startup"} + sent = [] + async def send(msg): sent.append(msg) + # Must not crash. Cache stays None. + asyncio.run(mw(scope, receive, send)) + assert url_builder._app_url_cache is None + + +def test_middleware_no_op_when_no_host_header(): + from coda_mcp.mcp_asgi import AppUrlCaptureMiddleware + mw = AppUrlCaptureMiddleware(_fake_app) + scope = _make_scope([]) + asyncio.run(_drive(mw, scope)) + assert url_builder._app_url_cache is None diff --git a/tests/test_coda_bridge.py b/tests/test_coda_bridge.py new file mode 100644 index 0000000..8d1e39f --- /dev/null +++ b/tests/test_coda_bridge.py @@ -0,0 +1,122 @@ +"""Unit tests for the stdio→HTTP MCP bridge (tools/coda-bridge.py). + +The bridge sits between a local MCP client (Claude Code's OAuth flow) and a +remote deployed CoDA app. It must: + 1. Mint a Databricks access token via the CLI and inject it as Bearer auth + 2. Forward the JSON-RPC payload unchanged to the configured APP_URL + 3. Surface server errors without dropping them + 4. Refuse to run without an APP_URL (operator misconfiguration) +""" +import importlib.util +import json +import os +import sys +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[1] +BRIDGE_PATH = REPO_ROOT / "tools" / "coda-bridge.py" + + +def _load_bridge(): + spec = importlib.util.spec_from_file_location("coda_bridge", BRIDGE_PATH) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +@pytest.fixture +def bridge(monkeypatch, tmp_path): + monkeypatch.setenv("CODA_MCP_URL", "https://fake-app.databricksapps.com/mcp") + monkeypatch.setenv("DATABRICKS_PROFILE", "test") + monkeypatch.setenv("HOME", str(tmp_path)) + return _load_bridge() + + +def test_bridge_loads_with_app_url(bridge): + assert bridge is not None + assert callable(getattr(bridge, "_forward", None)) or callable( + getattr(bridge, "forward", None) + ), "bridge must expose a forward function" + + +def test_forward_injects_authorization_header(bridge): + forward = getattr(bridge, "_forward", None) or getattr(bridge, "forward", None) + if forward is None: + pytest.skip("bridge implementation does not expose a forward entrypoint") + + fake_resp = MagicMock() + fake_resp.status = 200 + fake_resp.headers = {} + fake_resp.read.return_value = b'{"jsonrpc":"2.0","id":1,"result":{}}' + fake_resp.__enter__ = lambda s: s + fake_resp.__exit__ = MagicMock(return_value=False) + + fake_proc = MagicMock( + returncode=0, + stdout=json.dumps({"access_token": "tok-from-cli"}), + stderr="", + ) + + with patch("subprocess.run", return_value=fake_proc), \ + patch("urllib.request.urlopen", return_value=fake_resp) as mock_open: + forward(json.dumps({"jsonrpc": "2.0", "id": 1, "method": "ping", "params": {}})) + + sent_req = mock_open.call_args[0][0] + headers_lower = {k.lower(): v for k, v in sent_req.headers.items()} + assert "authorization" in headers_lower, "Bearer token MUST be injected" + assert "tok-from-cli" in headers_lower["authorization"], ( + "Authorization header should contain the token from `databricks auth token`" + ) + + +def test_forward_returns_server_response_body(bridge): + forward = getattr(bridge, "_forward", None) or getattr(bridge, "forward", None) + if forward is None: + pytest.skip("bridge implementation does not expose a forward entrypoint") + + server_payload = b'{"jsonrpc":"2.0","id":42,"result":{"ok":true}}' + fake_resp = MagicMock() + fake_resp.status = 200 + fake_resp.headers = {} + fake_resp.read.return_value = server_payload + fake_resp.__enter__ = lambda s: s + fake_resp.__exit__ = MagicMock(return_value=False) + + fake_proc = MagicMock( + returncode=0, + stdout=json.dumps({"access_token": "tok"}), + stderr="", + ) + + with patch("subprocess.run", return_value=fake_proc), \ + patch("urllib.request.urlopen", return_value=fake_resp): + result = forward( + json.dumps({"jsonrpc": "2.0", "id": 42, "method": "tools/list", "params": {}}) + ) + + if result is None: + pytest.skip("bridge writes directly to stdout — capture via capsys in a follow-up") + if isinstance(result, (bytes, bytearray)): + result = result.decode() + assert "ok" in result and "true" in result.lower(), ( + f"forward should surface the server response body; got {result!r}" + ) + + +def test_missing_app_url_is_handled(monkeypatch, tmp_path): + monkeypatch.delenv("CODA_MCP_URL", raising=False) + monkeypatch.delenv("APP_URL", raising=False) + monkeypatch.setenv("HOME", str(tmp_path)) + sys.modules.pop("coda_bridge", None) + with pytest.raises((SystemExit, ValueError, RuntimeError, KeyError)): + spec = importlib.util.spec_from_file_location("coda_bridge", BRIDGE_PATH) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + # If import-time guard is absent, the forward call itself should refuse. + forward = getattr(mod, "_forward", None) or getattr(mod, "forward", None) + if forward: + forward(json.dumps({"jsonrpc": "2.0", "id": 1, "method": "ping", "params": {}})) diff --git a/tests/test_content_filter_proxy.py b/tests/test_content_filter_proxy.py new file mode 100644 index 0000000..4aad029 --- /dev/null +++ b/tests/test_content_filter_proxy.py @@ -0,0 +1,556 @@ +"""Tests for content_filter_proxy — request/response sanitization for OpenCode.""" + +import json +import time + +import pytest +from unittest import mock + + +# --------------------------------------------------------------------------- +# strip_unsupported_schema_keys +# --------------------------------------------------------------------------- + +class TestStripUnsupportedSchemaKeys: + def test_strips_top_level_keys(self): + from content_filter_proxy import strip_unsupported_schema_keys + obj = {"type": "object", "$schema": "http://...", "additionalProperties": False, "title": "Foo"} + result = strip_unsupported_schema_keys(obj) + assert result == {"type": "object", "title": "Foo"} + + def test_strips_nested_keys(self): + from content_filter_proxy import strip_unsupported_schema_keys + obj = { + "type": "object", + "properties": { + "name": {"type": "string", "$ref": "#/defs/Name", "$comment": "ignore"}, + }, + } + result = strip_unsupported_schema_keys(obj) + assert result == { + "type": "object", + "properties": { + "name": {"type": "string"}, + }, + } + + def test_strips_inside_lists(self): + from content_filter_proxy import strip_unsupported_schema_keys + obj = [{"$id": "x", "type": "string"}, {"type": "int"}] + result = strip_unsupported_schema_keys(obj) + assert result == [{"type": "string"}, {"type": "int"}] + + def test_passes_through_primitives(self): + from content_filter_proxy import strip_unsupported_schema_keys + assert strip_unsupported_schema_keys("hello") == "hello" + assert strip_unsupported_schema_keys(42) == 42 + assert strip_unsupported_schema_keys(None) is None + + +# --------------------------------------------------------------------------- +# sanitize_tool_schemas +# --------------------------------------------------------------------------- + +class TestSanitizeToolSchemas: + def test_cleans_tool_parameters(self): + from content_filter_proxy import sanitize_tool_schemas + data = { + "tools": [ + {"function": {"name": "foo", "parameters": {"$schema": "x", "type": "object"}}}, + ], + } + result = sanitize_tool_schemas(data) + assert result["tools"][0]["function"]["parameters"] == {"type": "object"} + + def test_strips_top_level_request_keys(self): + from content_filter_proxy import sanitize_tool_schemas + data = { + "tools": [{"function": {"name": "foo", "parameters": {"type": "object"}}}], + "stream_options": {"include_usage": True}, + "$schema": "x", + } + result = sanitize_tool_schemas(data) + assert "stream_options" not in result + assert "$schema" not in result + + def test_no_tools_is_noop(self): + from content_filter_proxy import sanitize_tool_schemas + data = {"messages": [{"role": "user", "content": "hi"}]} + result = sanitize_tool_schemas(data) + assert result == data + + +# --------------------------------------------------------------------------- +# _extract_tool_ids_from_message +# --------------------------------------------------------------------------- + +class TestExtractToolIds: + def test_anthropic_format(self): + from content_filter_proxy import _extract_tool_ids_from_message + msg = { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "tu_1", "name": "bash"}, + {"type": "text", "text": "running..."}, + {"type": "tool_use", "id": "tu_2", "name": "read"}, + ], + } + assert _extract_tool_ids_from_message(msg) == {"tu_1", "tu_2"} + + def test_openai_format(self): + from content_filter_proxy import _extract_tool_ids_from_message + msg = { + "role": "assistant", + "tool_calls": [ + {"id": "tc_1", "function": {"name": "bash"}}, + {"id": "tc_2", "function": {"name": "read"}}, + ], + } + assert _extract_tool_ids_from_message(msg) == {"tc_1", "tc_2"} + + def test_no_tools(self): + from content_filter_proxy import _extract_tool_ids_from_message + msg = {"role": "assistant", "content": "hello"} + assert _extract_tool_ids_from_message(msg) == set() + + +# --------------------------------------------------------------------------- +# _extract_tool_refs_from_message +# --------------------------------------------------------------------------- + +class TestExtractToolRefs: + def test_anthropic_tool_result(self): + from content_filter_proxy import _extract_tool_refs_from_message + msg = { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "tu_1", "content": "ok"}, + ], + } + assert _extract_tool_refs_from_message(msg) == {"tu_1"} + + def test_openai_tool_message(self): + from content_filter_proxy import _extract_tool_refs_from_message + msg = {"role": "tool", "tool_call_id": "tc_1", "content": "result"} + assert _extract_tool_refs_from_message(msg) == {"tc_1"} + + def test_no_refs(self): + from content_filter_proxy import _extract_tool_refs_from_message + msg = {"role": "user", "content": "hi"} + assert _extract_tool_refs_from_message(msg) == set() + + +# --------------------------------------------------------------------------- +# sanitize_messages — the big one +# --------------------------------------------------------------------------- + +class TestSanitizeMessages: + def test_strips_empty_text_blocks(self): + from content_filter_proxy import sanitize_messages + messages = [ + {"role": "user", "content": [ + {"type": "text", "text": "hello"}, + {"type": "text", "text": ""}, + {"type": "text", "text": " "}, + ]}, + ] + result = sanitize_messages(messages) + assert len(result) == 1 + assert len(result[0]["content"]) == 1 + assert result[0]["content"][0]["text"] == "hello" + + def test_strips_orphaned_tool_result_anthropic(self): + """tool_result referencing a tool_use ID that doesn't exist in prev assistant msg.""" + from content_filter_proxy import sanitize_messages + messages = [ + {"role": "assistant", "content": [ + {"type": "tool_use", "id": "tu_1", "name": "bash"}, + ]}, + {"role": "user", "content": [ + {"type": "tool_result", "tool_use_id": "tu_1", "content": "ok"}, + {"type": "tool_result", "tool_use_id": "tu_ORPHAN", "content": "stale"}, + ]}, + ] + result = sanitize_messages(messages) + assert len(result) == 2 + # Only tu_1 should survive + user_blocks = result[1]["content"] + assert len(user_blocks) == 1 + assert user_blocks[0]["tool_use_id"] == "tu_1" + + def test_strips_orphaned_openai_tool_message(self): + from content_filter_proxy import sanitize_messages + messages = [ + {"role": "assistant", "tool_calls": [{"id": "tc_1", "function": {"name": "bash"}}]}, + {"role": "tool", "tool_call_id": "tc_1", "content": "ok"}, + {"role": "tool", "tool_call_id": "tc_ORPHAN", "content": "stale"}, + ] + result = sanitize_messages(messages) + assert len(result) == 2 + assert result[1]["role"] == "tool" + assert result[1]["tool_call_id"] == "tc_1" + + def test_cascading_orphan_removal(self): + """Dropping one message can make the next one orphaned too — multi-pass.""" + from content_filter_proxy import sanitize_messages + messages = [ + # assistant with tool_use tu_A + {"role": "assistant", "content": [{"type": "tool_use", "id": "tu_A", "name": "bash"}]}, + # user responds to tu_A + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "tu_A", "content": "ok"}]}, + # assistant with tool_use tu_B (referencing something dropped) + {"role": "assistant", "content": [{"type": "tool_use", "id": "tu_B", "name": "read"}]}, + # user responds to tu_B AND orphan tu_C (no matching tool_use) + {"role": "user", "content": [ + {"type": "tool_result", "tool_use_id": "tu_B", "content": "ok"}, + {"type": "tool_result", "tool_use_id": "tu_C", "content": "orphan"}, + ]}, + ] + result = sanitize_messages(messages) + # tu_C should be stripped, tu_A and tu_B should survive + assert len(result) == 4 + last_user_blocks = result[3]["content"] + assert len(last_user_blocks) == 1 + assert last_user_blocks[0]["tool_use_id"] == "tu_B" + + def test_drops_empty_user_message_after_filter(self): + """If all content blocks are stripped, the user message is dropped entirely.""" + from content_filter_proxy import sanitize_messages + messages = [ + {"role": "assistant", "content": [{"type": "tool_use", "id": "tu_1", "name": "bash"}]}, + {"role": "user", "content": [ + {"type": "tool_result", "tool_use_id": "tu_ORPHAN", "content": "stale"}, + ]}, + ] + result = sanitize_messages(messages) + # The user message should be dropped (all blocks were orphaned) + assert len(result) == 1 + assert result[0]["role"] == "assistant" + + def test_keeps_empty_assistant_message(self): + """Empty assistant messages are kept (not dropped) to preserve alternation.""" + from content_filter_proxy import sanitize_messages + messages = [ + {"role": "assistant", "content": [{"type": "text", "text": ""}]}, + ] + result = sanitize_messages(messages) + assert len(result) == 1 + assert result[0]["role"] == "assistant" + + def test_replaces_null_assistant_content(self): + from content_filter_proxy import sanitize_messages + messages = [ + {"role": "assistant", "content": None}, + ] + result = sanitize_messages(messages) + assert result[0]["content"] == "." + + def test_replaces_empty_string_assistant(self): + from content_filter_proxy import sanitize_messages + messages = [ + {"role": "assistant", "content": " "}, + ] + result = sanitize_messages(messages) + assert result[0]["content"] == "." + + def test_strips_empty_string_user(self): + from content_filter_proxy import sanitize_messages + messages = [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi"}, + {"role": "user", "content": ""}, + ] + result = sanitize_messages(messages) + assert len(result) == 2 # empty user dropped + + def test_passthrough_non_list(self): + from content_filter_proxy import sanitize_messages + assert sanitize_messages("not a list") == "not a list" + assert sanitize_messages(None) is None + + def test_preserves_non_dict_blocks(self): + """Non-dict items in content list are preserved as-is.""" + from content_filter_proxy import sanitize_messages + messages = [ + {"role": "user", "content": ["plain string", {"type": "text", "text": "hi"}]}, + ] + result = sanitize_messages(messages) + assert len(result[0]["content"]) == 2 + + def test_null_assistant_with_tool_calls_not_replaced(self): + """Assistant msg with null content but tool_calls should NOT get placeholder.""" + from content_filter_proxy import sanitize_messages + messages = [ + {"role": "assistant", "content": None, "tool_calls": [{"id": "tc_1"}]}, + ] + result = sanitize_messages(messages) + assert result[0]["content"] is None # preserved because tool_calls exist + + +# --------------------------------------------------------------------------- +# remap_tool_call +# --------------------------------------------------------------------------- + +class TestRemapToolCall: + def test_remaps_databricks_tool_call(self): + from content_filter_proxy import remap_tool_call + tc = { + "id": "tc_1", + "function": { + "name": "databricks-tool-call", + "arguments": json.dumps({"name": "execute_sql", "query": "SELECT 1"}), + }, + } + result = remap_tool_call(tc) + assert result["function"]["name"] == "execute_sql" + args = json.loads(result["function"]["arguments"]) + assert "name" not in args + assert args["query"] == "SELECT 1" + + def test_passthrough_normal_tool(self): + from content_filter_proxy import remap_tool_call + tc = {"id": "tc_1", "function": {"name": "bash", "arguments": '{"cmd": "ls"}'}} + result = remap_tool_call(tc) + assert result["function"]["name"] == "bash" + + def test_handles_invalid_json_args(self): + from content_filter_proxy import remap_tool_call + tc = {"id": "tc_1", "function": {"name": "databricks-tool-call", "arguments": "not json"}} + result = remap_tool_call(tc) + assert result["function"]["name"] == "databricks-tool-call" # unchanged + + +# --------------------------------------------------------------------------- +# fix_response_data +# --------------------------------------------------------------------------- + +class TestFixResponseData: + def test_remaps_tool_calls_in_message(self): + from content_filter_proxy import fix_response_data + data = { + "choices": [{ + "message": { + "tool_calls": [{ + "id": "tc_1", + "function": { + "name": "databricks-tool-call", + "arguments": json.dumps({"name": "run_sql", "q": "SELECT 1"}), + }, + }], + }, + "finish_reason": "stop", + }], + } + result = fix_response_data(data) + assert result["choices"][0]["message"]["tool_calls"][0]["function"]["name"] == "run_sql" + assert result["choices"][0]["finish_reason"] == "tool_calls" + + def test_fixes_streaming_delta(self): + from content_filter_proxy import fix_response_data + data = { + "choices": [{ + "delta": { + "tool_calls": [{ + "id": "tc_1", + "function": { + "name": "databricks-tool-call", + "arguments": json.dumps({"name": "run_sql"}), + }, + }], + }, + "finish_reason": "stop", + }], + } + result = fix_response_data(data) + assert result["choices"][0]["delta"]["tool_calls"][0]["function"]["name"] == "run_sql" + assert result["choices"][0]["finish_reason"] == "tool_calls" + + def test_noop_on_non_dict(self): + from content_filter_proxy import fix_response_data + assert fix_response_data("string") == "string" + assert fix_response_data(None) is None + + def test_no_choices_is_noop(self): + from content_filter_proxy import fix_response_data + data = {"id": "resp_1"} + assert fix_response_data(data) == data + + +# --------------------------------------------------------------------------- +# SSEProcessor +# --------------------------------------------------------------------------- + +class TestSSEProcessor: + def test_passthrough_non_data_lines(self): + from content_filter_proxy import SSEProcessor + proc = SSEProcessor() + assert proc.process_line("event: message") == ["event: message"] + assert proc.process_line(": comment") == [": comment"] + + def test_passthrough_done_signal(self): + from content_filter_proxy import SSEProcessor + proc = SSEProcessor() + result = proc.process_line("data: [DONE]") + assert "data: [DONE]" in result + + def test_passthrough_normal_tool(self): + from content_filter_proxy import SSEProcessor + proc = SSEProcessor() + event = { + "choices": [{ + "delta": {"tool_calls": [{"index": 0, "function": {"name": "bash"}}]}, + "finish_reason": None, + }], + } + result = proc.process_line(f"data: {json.dumps(event)}") + assert len(result) == 1 + assert "bash" in result[0] + + def test_buffers_databricks_tool_call(self): + """First chunk with databricks-tool-call name should be buffered.""" + from content_filter_proxy import SSEProcessor + proc = SSEProcessor() + event = { + "choices": [{ + "delta": { + "tool_calls": [{ + "index": 0, + "function": {"name": "databricks-tool-call", "arguments": ""}, + }], + }, + "finish_reason": None, + }], + } + result = proc.process_line(f"data: {json.dumps(event)}") + assert result == [] # buffered, not sent + + def test_resolves_name_from_args(self): + """Once args JSON is complete, name is resolved and buffered events flushed.""" + from content_filter_proxy import SSEProcessor + proc = SSEProcessor() + # First chunk — name is databricks-tool-call + event1 = { + "choices": [{ + "delta": { + "tool_calls": [{ + "index": 0, + "function": {"name": "databricks-tool-call", "arguments": ""}, + }], + }, + "finish_reason": None, + }], + } + proc.process_line(f"data: {json.dumps(event1)}") + + # Second chunk — args with real name + event2 = { + "choices": [{ + "delta": { + "tool_calls": [{ + "index": 0, + "function": {"arguments": json.dumps({"name": "execute_sql", "query": "SELECT 1"})}, + }], + }, + "finish_reason": None, + }], + } + result = proc.process_line(f"data: {json.dumps(event2)}") + # Should flush buffered events + current event + assert len(result) >= 1 + # The resolved name should appear in flushed output + combined = " ".join(result) + assert "execute_sql" in combined + + def test_flush_remaining(self): + from content_filter_proxy import SSEProcessor + proc = SSEProcessor() + # Buffer a databricks-tool-call but never resolve it + event = { + "choices": [{ + "delta": { + "tool_calls": [{ + "index": 0, + "function": {"name": "databricks-tool-call", "arguments": '{"partial'}, + }], + }, + "finish_reason": None, + }], + } + proc.process_line(f"data: {json.dumps(event)}") + remaining = proc.flush_remaining() + assert len(remaining) >= 1 # buffered lines flushed as-is + + def test_fixes_finish_reason_on_stop(self): + """finish_reason 'stop' with active tool state should become 'tool_calls'.""" + from content_filter_proxy import SSEProcessor + proc = SSEProcessor() + # Seed tool state + proc._tool_state[0] = {"args_buffer": "", "resolved_name": "bash", "buffered_lines": []} + event = { + "choices": [{"delta": {}, "finish_reason": "stop"}], + } + result = proc.process_line(f"data: {json.dumps(event)}") + parsed = json.loads(result[0][6:]) # strip "data: " + assert parsed["choices"][0]["finish_reason"] == "tool_calls" + + def test_invalid_json_passthrough(self): + from content_filter_proxy import SSEProcessor + proc = SSEProcessor() + result = proc.process_line("data: {invalid json}") + assert result == ["data: {invalid json}"] + + +# --------------------------------------------------------------------------- +# _get_fresh_token +# --------------------------------------------------------------------------- + +class TestGetFreshToken: + def setup_method(self): + """Reset token cache before each test.""" + from content_filter_proxy import _TOKEN_CACHE + _TOKEN_CACHE["token"] = None + _TOKEN_CACHE["read_at"] = 0.0 + + def test_reads_from_databrickscfg(self, tmp_path): + from content_filter_proxy import _get_fresh_token, _TOKEN_CACHE + cfg = tmp_path / ".databrickscfg" + cfg.write_text("[DEFAULT]\nhost = https://test.cloud.databricks.com\ntoken = dapi_test123\n") + with mock.patch("content_filter_proxy._DATABRICKSCFG_PATH", str(cfg)): + token = _get_fresh_token() + assert token == "dapi_test123" + assert _TOKEN_CACHE["token"] == "dapi_test123" + + def test_returns_cached_within_ttl(self, tmp_path): + from content_filter_proxy import _get_fresh_token, _TOKEN_CACHE + _TOKEN_CACHE["token"] = "cached_token" + _TOKEN_CACHE["read_at"] = time.time() # just now + # Even with a bad path, should return cached + with mock.patch("content_filter_proxy._DATABRICKSCFG_PATH", "/nonexistent"): + token = _get_fresh_token() + assert token == "cached_token" + + def test_refreshes_after_ttl(self, tmp_path): + from content_filter_proxy import _get_fresh_token, _TOKEN_CACHE + _TOKEN_CACHE["token"] = "old_token" + _TOKEN_CACHE["read_at"] = time.time() - 60 # expired + cfg = tmp_path / ".databrickscfg" + cfg.write_text("[DEFAULT]\nhost = https://test.cloud.databricks.com\ntoken = new_token\n") + with mock.patch("content_filter_proxy._DATABRICKSCFG_PATH", str(cfg)): + token = _get_fresh_token() + assert token == "new_token" + + def test_returns_stale_on_read_error(self, tmp_path): + from content_filter_proxy import _get_fresh_token, _TOKEN_CACHE + _TOKEN_CACHE["token"] = "stale_token" + _TOKEN_CACHE["read_at"] = 0.0 # force re-read + with mock.patch("content_filter_proxy._DATABRICKSCFG_PATH", "/nonexistent"): + token = _get_fresh_token() + assert token == "stale_token" + + def test_returns_none_when_no_cache_and_no_file(self): + from content_filter_proxy import _get_fresh_token, _TOKEN_CACHE + _TOKEN_CACHE["token"] = None + _TOKEN_CACHE["read_at"] = 0.0 + with mock.patch("content_filter_proxy._DATABRICKSCFG_PATH", "/nonexistent"): + token = _get_fresh_token() + assert token is None diff --git a/tests/test_gateway_discovery.py b/tests/test_gateway_discovery.py index 698445a..92ca725 100644 --- a/tests/test_gateway_discovery.py +++ b/tests/test_gateway_discovery.py @@ -132,7 +132,7 @@ def test_workspace_id_whitespace_stripped(self, mock_probe): # Integration tests — verify endpoint URLs constructed by setup scripts # --------------------------------------------------------------------------- -SETUP_DIR = Path(__file__).parent.parent +SETUP_DIR = Path(__file__).parent.parent / "setup" class TestEndpointConstruction: @@ -146,9 +146,11 @@ def _run_setup(self, script_name, tmp_path, env_overrides=None): "DATABRICKS_TOKEN": "dapi_test_token", "DATABRICKS_WORKSPACE_ID": "6280049833385130", "PATH": os.environ.get("PATH", ""), - "PYTHONPATH": str(SETUP_DIR), + "PYTHONPATH": str(SETUP_DIR.parent), # Pre-resolve gateway so subprocess skips the network probe "_GATEWAY_RESOLVED": "", + # Skip CLI install (curl | bash) — tests only verify config files + "SKIP_CLAUDE_INSTALL": "1", } # Ensure DATABRICKS_GATEWAY_HOST is NOT set (test auto-discovery) env.pop("DATABRICKS_GATEWAY_HOST", None) @@ -175,15 +177,15 @@ def test_setup_claude_falls_back_when_gateway_unreachable(self, tmp_path): # Gateway is unreachable from test env, so should fall back import json settings_path = tmp_path / ".claude" / "settings.json" - if settings_path.exists(): - settings = json.loads(settings_path.read_text()) - base_url = settings.get("env", {}).get("ANTHROPIC_BASE_URL", "") - assert base_url.endswith("/anthropic") - # Either gateway or serving-endpoints is valid - assert ( - "ai-gateway.cloud.databricks.com" in base_url - or "serving-endpoints/anthropic" in base_url - ) + assert settings_path.exists(), "settings.json was not written" + settings = json.loads(settings_path.read_text()) + base_url = settings.get("env", {}).get("ANTHROPIC_BASE_URL", "") + assert base_url.endswith("/anthropic") + # Either gateway or serving-endpoints is valid + assert ( + "ai-gateway.cloud.databricks.com" in base_url + or "serving-endpoints/anthropic" in base_url + ) def test_setup_claude_explicit_override(self, tmp_path): """setup_claude.py should prefer explicit DATABRICKS_GATEWAY_HOST.""" @@ -196,10 +198,10 @@ def test_setup_claude_explicit_override(self, tmp_path): import json settings_path = tmp_path / ".claude" / "settings.json" - if settings_path.exists(): - settings = json.loads(settings_path.read_text()) - base_url = settings.get("env", {}).get("ANTHROPIC_BASE_URL", "") - assert "custom.gateway.example.com" in base_url + assert settings_path.exists(), "settings.json was not written" + settings = json.loads(settings_path.read_text()) + base_url = settings.get("env", {}).get("ANTHROPIC_BASE_URL", "") + assert "custom.gateway.example.com" in base_url def test_setup_claude_fallback_no_gateway(self, tmp_path): """setup_claude.py falls back to DATABRICKS_HOST when no gateway available.""" @@ -210,10 +212,10 @@ def test_setup_claude_fallback_no_gateway(self, tmp_path): import json settings_path = tmp_path / ".claude" / "settings.json" - if settings_path.exists(): - settings = json.loads(settings_path.read_text()) - base_url = settings.get("env", {}).get("ANTHROPIC_BASE_URL", "") - assert "test.cloud.databricks.com/serving-endpoints/anthropic" in base_url + assert settings_path.exists(), "settings.json was not written" + settings = json.loads(settings_path.read_text()) + base_url = settings.get("env", {}).get("ANTHROPIC_BASE_URL", "") + assert "test.cloud.databricks.com/serving-endpoints/anthropic" in base_url @mock.patch("utils._probe_gateway", return_value=True) def test_codex_gateway_url_construction(self, mock_probe): diff --git a/tests/test_mcp_endpoint.py b/tests/test_mcp_endpoint.py new file mode 100644 index 0000000..da3e230 --- /dev/null +++ b/tests/test_mcp_endpoint.py @@ -0,0 +1,102 @@ +"""Unit tests for the Flask Blueprint fallback at coda_mcp.mcp_endpoint. + +Production traffic flows through coda_mcp.mcp_asgi (uvicorn + native MCP SDK). +This blueprint is the WSGI-only fallback. These tests pin the JSON-RPC contract +so the two paths stay in lockstep. +""" +import json + +import pytest + + +@pytest.fixture +def client(): + from app import app as flask_app + + return flask_app.test_client() + + +def _rpc(method, params=None, rpc_id=1): + return {"jsonrpc": "2.0", "id": rpc_id, "method": method, "params": params or {}} + + +def test_initialize_returns_server_info(client): + r = client.post("/mcp", json=_rpc("initialize", {"protocolVersion": "2025-03-26"})) + assert r.status_code == 200 + body = r.get_json() + assert body["jsonrpc"] == "2.0" + assert body["result"]["serverInfo"]["name"] == "coda" + assert "capabilities" in body["result"] + + +def test_tools_list_returns_three_v2_tools(client): + r = client.post("/mcp", json=_rpc("tools/list", {}, rpc_id=2)) + assert r.status_code == 200 + tools = r.get_json()["result"]["tools"] + names = {t["name"] for t in tools} + assert names == {"coda_run", "coda_inbox", "coda_get_result"}, ( + f"Tool surface drifted from the v2 contract (docs/mcp-v2-background-execution.md). Got: {names}" + ) + + +def test_tools_list_each_tool_has_description_and_schema(client): + r = client.post("/mcp", json=_rpc("tools/list", {}, rpc_id=3)) + for t in r.get_json()["result"]["tools"]: + assert t.get("description"), f"tool {t['name']} missing description (MCP requires it)" + assert isinstance(t.get("inputSchema"), dict), f"tool {t['name']} missing inputSchema" + + +def test_cors_preflight_returns_204(client): + r = client.options( + "/mcp", + headers={ + "Origin": "https://test.cloud.databricks.com", + "Access-Control-Request-Method": "POST", + }, + ) + assert r.status_code == 204 + assert "Access-Control-Allow-Origin" in r.headers + + +def test_ping_returns_empty_result(client): + r = client.post("/mcp", json=_rpc("ping", {}, rpc_id=4)) + assert r.status_code == 200 + body = r.get_json() + assert body["result"] == {} + assert "error" not in body + + +def test_unknown_method_returns_method_not_found(client): + r = client.post("/mcp", json=_rpc("does/not/exist", {}, rpc_id=5)) + body = r.get_json() + assert body.get("error", {}).get("code") == -32601, ( + f"Expected JSON-RPC method-not-found (-32601); got {body}" + ) + + +def test_unknown_tool_returns_jsonrpc_error(client): + r = client.post( + "/mcp", + json=_rpc("tools/call", {"name": "not_a_real_tool", "arguments": {}}, rpc_id=6), + ) + body = r.get_json() + assert "error" in body or ( + "result" in body and body["result"].get("isError") is True + ), f"Calling an unknown tool should error; got {body}" + + +def test_jsonrpc_id_is_echoed(client): + for rpc_id in (7, "string-id", 0): + r = client.post("/mcp", json=_rpc("ping", {}, rpc_id=rpc_id)) + assert r.get_json()["id"] == rpc_id + + +def test_post_with_non_json_body_does_not_crash(client): + r = client.post( + "/mcp", + data="not json at all", + headers={"Content-Type": "application/json"}, + ) + assert r.status_code in (200, 400) + if r.status_code == 200: + assert "error" in r.get_json() diff --git a/tests/test_mcp_integration.py b/tests/test_mcp_integration.py new file mode 100644 index 0000000..48a3173 --- /dev/null +++ b/tests/test_mcp_integration.py @@ -0,0 +1,433 @@ +"""End-to-end MCP integration tests — v2 background execution + inbox API. + +Exercises the full flow: coda_run -> coda_inbox -> coda_get_result. +No real PTY — app hooks are mocked. +""" + +import json +import os +import time +from unittest.mock import MagicMock + +import pytest + + +# ── helpers ────────────────────────────────────────────────────────── + + +def _parse(result: str) -> dict: + """Parse JSON string returned by MCP tools.""" + return json.loads(result) + + +# ── fixture ────────────────────────────────────────────────────────── + + +@pytest.fixture(autouse=True) +def isolated_env(tmp_path): + """Redirect state to tmp and mock PTY hooks.""" + from coda_mcp import task_manager as tm + from coda_mcp import mcp_server as ms + + original_dir = tm.SESSIONS_DIR + tm.SESSIONS_DIR = str(tmp_path / "sessions") + + mock_send = MagicMock() + mock_close = MagicMock() + ms.set_app_hooks( + create_session_fn=lambda label, **kwargs: f"pty-mock-{label}", + send_input_fn=mock_send, + close_session_fn=mock_close, + ) + + yield {"tmp": tmp_path, "mock_send": mock_send, "mock_close": mock_close} + + tm.SESSIONS_DIR = original_dir + ms.set_app_hooks(None, None, None) + + +# ── 1. Happy-path: fire-and-forget → inbox → result ───────────────── + + +class TestFullMcpFlow: + @pytest.mark.asyncio + async def test_full_background_flow(self, isolated_env): + """Happy path: run (fire-and-forget) → inbox → result.""" + from coda_mcp import mcp_server as ms + from coda_mcp import task_manager as tm + + # Step 1: submit task (returns immediately) + with MagicMock() as mock_thread: + from coda_mcp import mcp_server + with pytest.MonkeyPatch.context() as mp: + mp.setattr("coda_mcp.mcp_server.threading", mock_thread) + raw = await ms.coda_run( + prompt="create a sales pipeline", + email="alice@test.com", + context='{"tables": ["sales.transactions"]}', + ) + + task = _parse(raw) + assert task["status"] == "running" + task_id = task["task_id"] + session_id = task["session_id"] + assert task_id.startswith("task-") + assert session_id.startswith("sess-") + + # Step 2: inbox shows running task + raw = await ms.coda_inbox() + inbox = _parse(raw) + assert len(inbox["tasks"]) == 1 + assert inbox["tasks"][0]["task_id"] == task_id + assert inbox["tasks"][0]["status"] == "running" + assert inbox["counts"]["running"] == 1 + + # Step 3: simulate agent writing result.json + tdir = tm._task_dir(session_id, task_id) + result_path = os.path.join(tdir, "result.json") + with open(result_path, "w") as f: + json.dump({ + "status": "completed", + "summary": "Created sales pipeline with 3 stages", + "files_changed": ["pipeline.py", "config.yaml"], + "artifacts": ["/workspace/pipeline.py"], + "errors": [], + }, f) + + # Step 4: complete_task (simulating what _watch_task does) + tm.complete_task(session_id, task_id) + + # Step 5: inbox shows completed + raw = await ms.coda_inbox() + inbox = _parse(raw) + assert len(inbox["tasks"]) == 1 + assert inbox["tasks"][0]["status"] == "completed" + assert inbox["tasks"][0]["summary"] == "Created sales pipeline with 3 stages" + assert inbox["counts"]["completed"] == 1 + + # Step 6: get full result + raw = await ms.coda_get_result(task_id=task_id, session_id=session_id) + result = _parse(raw) + assert result["task_id"] == task_id + assert result["summary"] == "Created sales pipeline with 3 stages" + assert result["files_changed"] == ["pipeline.py", "config.yaml"] + + # Step 7: session was auto-closed + session = tm._read_session(session_id) + assert session["status"] == "closed" + + +# ── 2. Task chaining with previous_session_id ─────────────────────── + + +class TestTaskChaining: + @pytest.mark.asyncio + async def test_chained_task_references_prior_session(self, isolated_env): + """A chained task includes prior session context in prompt.""" + from coda_mcp import mcp_server as ms + from coda_mcp import task_manager as tm + + # First task + raw = await ms.coda_run( + prompt="build pipeline", + email="bob@test.com", + ) + first = _parse(raw) + first_sid = first["session_id"] + first_tid = first["task_id"] + + # Complete first task + tdir = tm._task_dir(first_sid, first_tid) + with open(os.path.join(tdir, "result.json"), "w") as f: + json.dump({ + "status": "completed", + "summary": "Built pipeline.py", + "files_changed": ["pipeline.py"], + }, f) + tm.complete_task(first_sid, first_tid) + + # Second task chained to first + raw = await ms.coda_run( + prompt="add tests for the pipeline", + email="bob@test.com", + previous_session_id=first_sid, + ) + second = _parse(raw) + second_sid = second["session_id"] + second_tid = second["task_id"] + + # Verify prompt references prior session + prompt_path = os.path.join( + tm._task_dir(second_sid, second_tid), "prompt.txt" + ) + with open(prompt_path) as f: + prompt_text = f.read() + assert f"PRIOR SESSION: {first_sid}" in prompt_text + + # Verify meta.json has previous_session_id + meta_path = os.path.join( + tm._task_dir(second_sid, second_tid), "meta.json" + ) + with open(meta_path) as f: + meta = json.load(f) + assert meta["previous_session_id"] == first_sid + + # Verify inbox shows chaining + raw = await ms.coda_inbox() + inbox = _parse(raw) + running_tasks = [t for t in inbox["tasks"] if t["status"] == "running"] + assert len(running_tasks) == 1 + assert running_tasks[0]["previous_session_id"] == first_sid + + +# ── 3. Concurrency limit ──────────────────────────────────────────── + + +class TestConcurrencyLimit: + @pytest.mark.asyncio + async def test_exceeding_limit_returns_error(self, isolated_env): + """Exceeding MAX_CONCURRENT_TASKS returns a clear error.""" + from coda_mcp import mcp_server as ms + from unittest.mock import patch + + with patch("coda_mcp.task_manager.MAX_CONCURRENT_TASKS", 1): + r1 = await ms.coda_run(prompt="task1", email="a@b.com") + assert _parse(r1)["status"] == "running" + + r2 = await ms.coda_run(prompt="task2", email="a@b.com") + d2 = _parse(r2) + assert d2["status"] == "error" + assert "concurrency" in d2["error"].lower() + + +# ── 4. Yolo permissions → --yolo flag ─────────────────────────────── + + +class TestYoloPermissions: + @pytest.mark.asyncio + async def test_yolo_permissions(self, isolated_env): + """permissions='yolo' causes the PTY command to include --yolo.""" + from coda_mcp import mcp_server as ms + + mock_send = isolated_env["mock_send"] + + with MagicMock() as mock_thread: + from coda_mcp import mcp_server + with pytest.MonkeyPatch.context() as mp: + mp.setattr("coda_mcp.mcp_server.threading", mock_thread) + await ms.coda_run( + prompt="deploy everything", + email="dave@test.com", + permissions="yolo", + ) + + mock_send.assert_called_once() + cmd = mock_send.call_args[0][1] + assert "--yolo" in cmd + + +# ── 5. Session auto-close on completion ────────────────────────────── + + +class TestAutoClose: + @pytest.mark.asyncio + async def test_session_auto_closes(self, isolated_env): + """Session is auto-closed when task completes.""" + from coda_mcp import mcp_server as ms + from coda_mcp import task_manager as tm + + raw = await ms.coda_run(prompt="quick job", email="a@b.com") + d = _parse(raw) + + # Session should be busy + session = tm._read_session(d["session_id"]) + assert session["status"] == "busy" + + # Complete the task + tdir = tm._task_dir(d["session_id"], d["task_id"]) + with open(os.path.join(tdir, "result.json"), "w") as f: + json.dump({"status": "completed", "summary": "done"}, f) + tm.complete_task(d["session_id"], d["task_id"]) + + # Session should now be closed + session = tm._read_session(d["session_id"]) + assert session["status"] == "closed" + assert "closed_at" in session + + +# ── 6. Cleanup expired tasks ──────────────────────────────────────── + + +class TestCleanup: + @pytest.mark.asyncio + async def test_cleanup_removes_expired(self, isolated_env): + """cleanup_expired_tasks removes old closed sessions.""" + from coda_mcp import mcp_server as ms + from coda_mcp import task_manager as tm + from unittest.mock import patch + + raw = await ms.coda_run(prompt="old task", email="a@b.com") + d = _parse(raw) + + # Complete and close + tdir = tm._task_dir(d["session_id"], d["task_id"]) + with open(os.path.join(tdir, "result.json"), "w") as f: + json.dump({"status": "completed", "summary": "done"}, f) + tm.complete_task(d["session_id"], d["task_id"]) + + # Backdate closed_at to expire it + session = tm._read_session(d["session_id"]) + session["closed_at"] = time.time() - 90000 # 25 hours ago + tm._write_json(tm._session_file(d["session_id"]), session) + + # Cleanup should remove it + removed = tm.cleanup_expired_tasks() + assert removed == 1 + + # Inbox should be empty now + raw = await ms.coda_inbox() + inbox = _parse(raw) + assert len(inbox["tasks"]) == 0 + + +# ── 7. E2E: grace period + transcript replay ──────────────────────── +# Import the PTY-availability guard from test_transcript. +# The test below requires a real PTY to be usable. + + +def _pty_is_usable() -> bool: + import os + if not hasattr(os, "openpty"): + return False + try: + master, slave = os.openpty() + os.close(master) + os.close(slave) + return True + except OSError: + return False + + +_pty_skip = pytest.mark.skipif(not _pty_is_usable(), reason="pty.openpty() not available") + + +@_pty_skip +def test_end_to_end_grace_and_replay(tmp_path, monkeypatch): + """Stub hermes via direct file I/O, then exercise the full coda_run flow. + + Wires up real Flask PTY hooks (not mocks) to verify the complete pipeline: + create PTY → send input → result.json written → grace period → PTY closed + → transcript persists → find_task_dir_by_pty_session resolves correctly. + """ + import asyncio + import json + import time + from pathlib import Path + + from coda_mcp import mcp_server, task_manager, url_builder + + # Override the autouse isolated_env fixture's SESSIONS_DIR patch with our own. + monkeypatch.setattr(task_manager, "SESSIONS_DIR", str(tmp_path / "sessions")) + monkeypatch.setattr(url_builder, "_app_url_cache", "app.example.com") + # Shrink grace period so the test runs fast. + monkeypatch.setattr(mcp_server, "GRACE_PERIOD_S", 2) + + from app import ( + mcp_create_pty_session, + mcp_send_input, + mcp_close_pty_session, + _mark_grace_for_session, + _bump_session_last_poll, + sessions, + ) + + # Disable the watcher thread so the test's manual orchestration + # (write result.json + call _schedule_deferred_close) is the sole driver. + # The watcher otherwise races with the manual orchestration on a + # 5-second poll cycle and produces SessionNotFoundError. + monkeypatch.setattr(mcp_server, "_watch_task", lambda *a, **kw: None) + + mcp_server.set_app_hooks( + mcp_create_pty_session, + mcp_send_input, + mcp_close_pty_session, + _mark_grace_for_session, + _bump_session_last_poll, + ) + + # Initialize cleanup-referenced names BEFORE the try so an early failure + # (e.g., coda_run or _read_session raising) doesn't shadow the original + # exception with an UnboundLocalError in the finally block. + pty_id = None + sess_id = None + task_id = None + + try: + # --- Step 1: Submit a fake task ------------------------------------------ + result_json = asyncio.run(mcp_server.coda_run( + prompt="test", + email="u@x", + timeout_s=5, + )) + result = json.loads(result_json) + assert result["status"] == "running", f"Unexpected status: {result}" + sess_id = result["session_id"] + task_id = result["task_id"] + pty_id = task_manager._read_session(sess_id)["pty_session_id"] + + # --- Step 2: viewer_url contains the pty_id ------------------------------ + assert pty_id in result["viewer_url"] + + # --- Step 3: Simulate hermes writing to the PTY -------------------------- + mcp_send_input(pty_id, "echo HELLO_FROM_HERMES\n") + time.sleep(0.5) + + # --- Step 4: Simulate hermes completion by writing result.json ----------- + tdir = task_manager._task_dir(sess_id, task_id) + Path(tdir).joinpath("result.json").write_text(json.dumps({ + "status": "completed", + "summary": "stub", + "files_changed": [], + "artifacts": {}, + "errors": [], + })) + + # --- Step 5: Trigger deferred close (watcher normally does this) --------- + # complete_task first (watcher calls this before _schedule_deferred_close) + task_manager.complete_task(sess_id, task_id) + mcp_server._schedule_deferred_close(sess_id) + + # --- Step 6: PTY still alive immediately after grace scheduling ---------- + assert pty_id in sessions, "PTY should still be in sessions during grace" + assert sessions[pty_id]["grace"] is True, "PTY should be marked grace" + + # --- Step 7: Wait past GRACE_PERIOD_S (2 s) + small margin -------------- + time.sleep(2.5) + + # --- Step 8: PTY now gone ------------------------------------------------ + assert pty_id not in sessions, "PTY should have been closed after grace" + + # --- Step 9: Transcript file exists and contains echoed output ----------- + transcript = Path(tdir) / "transcript.log" + assert transcript.exists(), f"transcript.log missing at {transcript}" + assert b"HELLO_FROM_HERMES" in transcript.read_bytes(), \ + "Echoed string not found in transcript" + + # --- Step 10: find_task_dir_by_pty_session resolves to the right dir ----- + found = task_manager.find_task_dir_by_pty_session(pty_id) + assert found == str(tdir), f"Expected {tdir!r}, got {found!r}" + + finally: + # Re-install mock hooks so the autouse fixture's teardown is consistent. + from unittest.mock import MagicMock + mcp_server.set_app_hooks( + create_session_fn=lambda label, **kwargs: f"pty-mock-{label}", + send_input_fn=MagicMock(), + close_session_fn=MagicMock(), + ) + # Best-effort PTY cleanup if the test failed before the Timer fired. + if pty_id and pty_id in sessions: + try: + mcp_close_pty_session(pty_id) + except Exception: + pass diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py new file mode 100644 index 0000000..7c1e287 --- /dev/null +++ b/tests/test_mcp_server.py @@ -0,0 +1,508 @@ +"""Tests for mcp_server — v2 background execution + inbox API.""" + +import json +import os +from unittest import mock + +import pytest + + +# ── helpers ────────────────────────────────────────────────────────── + + +@pytest.fixture(autouse=True) +def _reset_hooks(): + """Clear app hooks before/after each test.""" + from coda_mcp import mcp_server + + mcp_server._app_create_session = None + mcp_server._app_send_input = None + mcp_server._app_close_session = None + mcp_server._app_mark_grace = None + mcp_server._app_bump_poll = None + yield + mcp_server._app_create_session = None + mcp_server._app_send_input = None + mcp_server._app_close_session = None + mcp_server._app_mark_grace = None + mcp_server._app_bump_poll = None + + +@pytest.fixture(autouse=True) +def _isolated_sessions(tmp_path): + """Point task_manager.SESSIONS_DIR at a temp dir.""" + sessions_dir = str(tmp_path / ".coda" / "sessions") + with mock.patch("coda_mcp.task_manager.SESSIONS_DIR", sessions_dir): + yield sessions_dir + + +def _parse(result: str) -> dict: + """Parse JSON string returned by MCP tools.""" + return json.loads(result) + + +# ── Tool registration ──────────────────────────────────────────────── + + +class TestToolRegistration: + def test_three_tools_registered(self): + from coda_mcp import mcp_server + + tool_mgr = mcp_server.mcp._tool_manager + tool_names = set(tool_mgr._tools.keys()) + expected = {"coda_run", "coda_inbox", "coda_get_result"} + assert expected == tool_names, f"Expected {expected}, got {tool_names}" + + def test_tool_count_is_three(self): + from coda_mcp import mcp_server + + tool_mgr = mcp_server.mcp._tool_manager + assert len(tool_mgr._tools) == 3 + + +# ── coda_run ───────────────────────────────────────────────────────── + + +class TestCodaRun: + @pytest.mark.asyncio + async def test_creates_task_disk_only(self): + """Without app hooks, creates session+task on disk, returns immediately.""" + from coda_mcp import mcp_server + + result = await mcp_server.coda_run( + prompt="fix the bug", + email="a@b.com", + ) + data = _parse(result) + assert data["status"] == "running" + assert data["task_id"].startswith("task-") + assert data["session_id"].startswith("sess-") + + @pytest.mark.asyncio + async def test_auto_creates_session(self): + """coda_run auto-creates a session — no separate create_session needed.""" + from coda_mcp import mcp_server + from coda_mcp import task_manager + + result = await mcp_server.coda_run( + prompt="build pipeline", + email="a@b.com", + ) + data = _parse(result) + session = task_manager._read_session(data["session_id"]) + assert session["email"] == "a@b.com" + assert session["status"] == "busy" # task is running + + @pytest.mark.asyncio + async def test_sends_to_pty_when_hooks_set(self): + """With hooks, creates PTY and sends hermes command.""" + from coda_mcp import mcp_server + + mock_create = mock.Mock(return_value="pty-xyz") + mock_send = mock.Mock() + mcp_server.set_app_hooks( + create_session_fn=mock_create, + send_input_fn=mock_send, + close_session_fn=mock.Mock(), + ) + + with mock.patch("coda_mcp.mcp_server.threading"): + result = await mcp_server.coda_run( + prompt="fix the bug", + email="a@b.com", + ) + + data = _parse(result) + assert data["status"] == "running" + mock_create.assert_called_once() + call_kwargs = mock_create.call_args.kwargs + assert call_kwargs["label"] == "hermes-mcp" + assert "transcript_path" in call_kwargs + mock_send.assert_called_once() + assert "hermes" in mock_send.call_args[0][1] + + @pytest.mark.asyncio + async def test_yolo_permission(self): + """permissions='yolo' produces --yolo flag in PTY command.""" + from coda_mcp import mcp_server + + mock_send = mock.Mock() + mcp_server.set_app_hooks( + create_session_fn=mock.Mock(return_value="pty-1"), + send_input_fn=mock_send, + close_session_fn=mock.Mock(), + ) + + with mock.patch("coda_mcp.mcp_server.threading"): + await mcp_server.coda_run( + prompt="go fast", + email="a@b.com", + permissions="yolo", + ) + + cmd = mock_send.call_args[0][1] + assert "--yolo" in cmd + + @pytest.mark.asyncio + async def test_previous_session_id_in_prompt(self): + """previous_session_id appears in the wrapped prompt.""" + from coda_mcp import mcp_server + from coda_mcp import task_manager + + # Create a "prior" session with a completed task + prior = task_manager.create_session("a@b.com", "u1") + prior_sid = prior["session_id"] + + result = await mcp_server.coda_run( + prompt="add tests", + email="a@b.com", + previous_session_id=prior_sid, + ) + data = _parse(result) + + # Read the prompt.txt and verify prior session reference + tdir = task_manager._task_dir(data["session_id"], data["task_id"]) + with open(os.path.join(tdir, "prompt.txt")) as f: + prompt_text = f.read() + + assert f"PRIOR SESSION: {prior_sid}" in prompt_text + + @pytest.mark.asyncio + async def test_meta_json_written(self): + """coda_run writes meta.json with task metadata.""" + from coda_mcp import mcp_server + from coda_mcp import task_manager + + result = await mcp_server.coda_run( + prompt="build a dashboard for sales", + email="alice@test.com", + previous_session_id="sess-old", + ) + data = _parse(result) + + meta_path = os.path.join( + task_manager._task_dir(data["session_id"], data["task_id"]), + "meta.json", + ) + with open(meta_path) as f: + meta = json.load(f) + + assert meta["email"] == "alice@test.com" + assert meta["previous_session_id"] == "sess-old" + assert meta["prompt_summary"] == "build a dashboard for sales" + assert "created_at" in meta + + @pytest.mark.asyncio + async def test_concurrency_limit(self): + """Exceeding MAX_CONCURRENT_TASKS returns an error.""" + from coda_mcp import mcp_server + + with mock.patch("coda_mcp.task_manager.MAX_CONCURRENT_TASKS", 1): + # First task succeeds + r1 = await mcp_server.coda_run(prompt="task1", email="a@b.com") + assert _parse(r1)["status"] == "running" + + # Second task should fail (1 already running) + r2 = await mcp_server.coda_run(prompt="task2", email="a@b.com") + d2 = _parse(r2) + assert d2["status"] == "error" + assert "concurrency" in d2["error"].lower() + + +# ── coda_inbox ─────────────────────────────────────────────────────── + + +class TestCodaInbox: + @pytest.mark.asyncio + async def test_empty_inbox(self): + """No tasks → empty inbox.""" + from coda_mcp import mcp_server + + result = await mcp_server.coda_inbox() + data = _parse(result) + assert data["tasks"] == [] + assert data["counts"] == {"running": 0, "completed": 0, "failed": 0} + + @pytest.mark.asyncio + async def test_running_task_in_inbox(self): + """A running task shows up in the inbox.""" + from coda_mcp import mcp_server + + await mcp_server.coda_run(prompt="build pipeline", email="a@b.com") + + result = await mcp_server.coda_inbox() + data = _parse(result) + assert len(data["tasks"]) == 1 + assert data["tasks"][0]["status"] == "running" + assert data["tasks"][0]["prompt_summary"] == "build pipeline" + assert data["counts"]["running"] == 1 + + @pytest.mark.asyncio + async def test_completed_task_in_inbox(self): + """A completed task shows summary in inbox.""" + from coda_mcp import mcp_server + from coda_mcp import task_manager + + r = await mcp_server.coda_run(prompt="fix bug", email="a@b.com") + d = _parse(r) + + # Simulate agent writing result.json + tdir = task_manager._task_dir(d["session_id"], d["task_id"]) + result_path = os.path.join(tdir, "result.json") + with open(result_path, "w") as f: + json.dump({ + "status": "completed", + "summary": "Fixed the login bug", + "files_changed": ["auth.py"], + "artifacts": [], + "errors": [], + }, f) + + result = await mcp_server.coda_inbox() + data = _parse(result) + assert len(data["tasks"]) == 1 + assert data["tasks"][0]["status"] == "completed" + assert data["tasks"][0]["summary"] == "Fixed the login bug" + + @pytest.mark.asyncio + async def test_status_filter(self): + """Filtering inbox by status works.""" + from coda_mcp import mcp_server + from coda_mcp import task_manager + + # Create two tasks — one running, one completed + r1 = await mcp_server.coda_run(prompt="task1", email="a@b.com") + d1 = _parse(r1) + + r2 = await mcp_server.coda_run(prompt="task2", email="a@b.com") + d2 = _parse(r2) + + # Complete task2 + tdir = task_manager._task_dir(d2["session_id"], d2["task_id"]) + with open(os.path.join(tdir, "result.json"), "w") as f: + json.dump({"status": "completed", "summary": "done"}, f) + + # Filter running only + result = await mcp_server.coda_inbox(status="running") + data = _parse(result) + assert len(data["tasks"]) == 1 + assert data["tasks"][0]["task_id"] == d1["task_id"] + + @pytest.mark.asyncio + async def test_multiple_tasks_sorted_recent_first(self): + """Inbox returns tasks sorted most recent first.""" + from coda_mcp import mcp_server + + r1 = await mcp_server.coda_run(prompt="first", email="a@b.com") + r2 = await mcp_server.coda_run(prompt="second", email="a@b.com") + + result = await mcp_server.coda_inbox() + data = _parse(result) + assert len(data["tasks"]) == 2 + # Most recent first + assert data["tasks"][0]["prompt_summary"] == "second" + assert data["tasks"][1]["prompt_summary"] == "first" + + +# ── coda_get_result ────────────────────────────────────────────────── + + +class TestCodaGetResult: + @pytest.mark.asyncio + async def test_returns_result(self): + from coda_mcp import mcp_server + from coda_mcp import task_manager + + r = await mcp_server.coda_run(prompt="go", email="a@b.com") + d = _parse(r) + + # Simulate agent writing result.json + tdir = task_manager._task_dir(d["session_id"], d["task_id"]) + with open(os.path.join(tdir, "result.json"), "w") as f: + json.dump({ + "summary": "Fixed the bug", + "files_changed": ["app.py"], + "artifacts": [], + "errors": [], + }, f) + + result = await mcp_server.coda_get_result( + task_id=d["task_id"], session_id=d["session_id"] + ) + data = _parse(result) + assert data["task_id"] == d["task_id"] + assert data["session_id"] == d["session_id"] + assert data["summary"] == "Fixed the bug" + + @pytest.mark.asyncio + async def test_no_result_yet(self): + from coda_mcp import mcp_server + + r = await mcp_server.coda_run(prompt="go", email="a@b.com") + d = _parse(r) + + result = await mcp_server.coda_get_result( + task_id=d["task_id"], session_id=d["session_id"] + ) + data = _parse(result) + assert data["status"] == "running" + assert "not yet available" in data["message"] + + +# ── Deferred close (Timer) ─────────────────────────────────────────── + + +import threading +from unittest import mock + +from coda_mcp import mcp_server, task_manager + + +def test_set_app_hooks_accepts_grace_and_bump_hooks(): + create = mock.MagicMock() + send = mock.MagicMock() + close = mock.MagicMock() + mark_grace = mock.MagicMock() + bump_poll = mock.MagicMock() + mcp_server.set_app_hooks(create, send, close, mark_grace, bump_poll) + assert mcp_server._app_mark_grace is mark_grace + assert mcp_server._app_bump_poll is bump_poll + + +def test_watch_task_schedules_timer_on_completion(tmp_path, monkeypatch): + monkeypatch.setattr(task_manager, "SESSIONS_DIR", str(tmp_path)) + # Create a session + task with a faked result.json + s = task_manager.create_session("u@x", "uid", label="t") + sid = s["session_id"] + task_manager._update_session_field(sid, "pty_session_id", "pty-abc") + t = task_manager.create_task(sid, "do thing", "u@x") + tid = t["task_id"] + tdir = task_manager._task_dir(sid, tid) + task_manager._write_json(tdir + "/result.json", {"status": "completed"}) + + mark = mock.MagicMock() + bump = mock.MagicMock() + closer = mock.MagicMock() + mcp_server.set_app_hooks(mock.MagicMock(), mock.MagicMock(), closer, mark, bump) + + timer_created = [] + real_timer = threading.Timer + + def fake_timer(seconds, fn, args=None, kwargs=None): + timer_created.append((seconds, fn, args)) + t = real_timer(seconds, fn, args=args, kwargs=kwargs) + return t + + monkeypatch.setattr(mcp_server.threading, "Timer", fake_timer) + + # Use a very short watch interval and ensure no real Timer fires + monkeypatch.setattr(mcp_server, "GRACE_PERIOD_S", 0.05) + + # Run one iteration manually + mcp_server._watch_task(sid, tid, timeout_s=10) + + # Timer should be scheduled for GRACE_PERIOD_S seconds with closer + pty_session_id + assert len(timer_created) == 1 + delay, fn, args = timer_created[0] + assert delay == 0.05 + assert fn is closer + assert args == ("pty-abc",) + + # _mark_grace and _bump_session_last_poll should have been called + mark.assert_called_once_with("pty-abc") + bump.assert_called_once_with("pty-abc", 0.05) + + +# ── viewer_url + transcript_path wiring ───────────────────────────── + + +import asyncio +import json +import os +from unittest import mock + +from coda_mcp import mcp_server, task_manager, url_builder + + +def _run(coro): + return asyncio.get_event_loop().run_until_complete(coro) if not asyncio.iscoroutine(coro) else asyncio.run(coro) + + +def test_coda_run_includes_viewer_url_when_builder_returns_one(tmp_path, monkeypatch): + monkeypatch.setattr(task_manager, "SESSIONS_DIR", str(tmp_path)) + monkeypatch.setattr(url_builder, "_app_url_cache", "app.example.com") + + create = mock.MagicMock(return_value="pty-abc") + send = mock.MagicMock() + closer = mock.MagicMock() + mcp_server.set_app_hooks(create, send, closer, mock.MagicMock(), mock.MagicMock()) + + result_json = asyncio.run(mcp_server.coda_run(prompt="do it", email="u@x")) + result = json.loads(result_json) + assert result["status"] == "running" + assert "?session=pty-abc" in result["viewer_url"] + assert result["viewer_url"].startswith("https://app.example.com") + + +def test_coda_run_omits_viewer_url_when_builder_returns_none(tmp_path, monkeypatch): + monkeypatch.setattr(task_manager, "SESSIONS_DIR", str(tmp_path)) + monkeypatch.setattr(url_builder, "_app_url_cache", None) + monkeypatch.delenv("CODA_APP_URL", raising=False) + + create = mock.MagicMock(return_value="pty-abc") + mcp_server.set_app_hooks(create, mock.MagicMock(), mock.MagicMock(), mock.MagicMock(), mock.MagicMock()) + + result_json = asyncio.run(mcp_server.coda_run(prompt="do it", email="u@x")) + result = json.loads(result_json) + # viewer_url present but None when builder returns None + assert result.get("viewer_url") is None + + +def test_coda_run_passes_transcript_path_to_create_session(tmp_path, monkeypatch): + monkeypatch.setattr(task_manager, "SESSIONS_DIR", str(tmp_path)) + create = mock.MagicMock(return_value="pty-abc") + mcp_server.set_app_hooks(create, mock.MagicMock(), mock.MagicMock(), mock.MagicMock(), mock.MagicMock()) + + asyncio.run(mcp_server.coda_run(prompt="do it", email="u@x")) + # create_session was called with transcript_path=... pointing into ~/.coda/sessions//tasks//transcript.log + kwargs = create.call_args.kwargs + assert "transcript_path" in kwargs + assert kwargs["transcript_path"].endswith("transcript.log") + assert "tasks" in kwargs["transcript_path"] + + +def test_coda_inbox_decorates_each_task_with_viewer_url(tmp_path, monkeypatch): + monkeypatch.setattr(task_manager, "SESSIONS_DIR", str(tmp_path)) + monkeypatch.setattr(url_builder, "_app_url_cache", "app.example.com") + + # Seed one session with one task and a pty_session_id + s = task_manager.create_session("u@x", "uid", label="t") + sid = s["session_id"] + task_manager._update_session_field(sid, "pty_session_id", "pty-xyz") + task_manager.create_task(sid, "prompt", "u@x") + + result_json = asyncio.run(mcp_server.coda_inbox()) + result = json.loads(result_json) + assert len(result["tasks"]) == 1 + assert "viewer_url" in result["tasks"][0] + assert "?session=pty-xyz" in result["tasks"][0]["viewer_url"] + + +def test_coda_get_result_includes_viewer_url(tmp_path, monkeypatch): + monkeypatch.setattr(task_manager, "SESSIONS_DIR", str(tmp_path)) + monkeypatch.setattr(url_builder, "_app_url_cache", "app.example.com") + + s = task_manager.create_session("u@x", "uid", label="t") + sid = s["session_id"] + task_manager._update_session_field(sid, "pty_session_id", "pty-xyz") + t = task_manager.create_task(sid, "prompt", "u@x") + tid = t["task_id"] + tdir = task_manager._task_dir(sid, tid) + task_manager._write_json(tdir + "/result.json", { + "status": "completed", "summary": "ok", + }) + + result_json = asyncio.run(mcp_server.coda_get_result(tid, sid)) + result = json.loads(result_json) + assert "viewer_url" in result + assert "?session=pty-xyz" in result["viewer_url"] diff --git a/tests/test_mlflow_tracing.py b/tests/test_mlflow_tracing.py index fb6e975..c72113f 100644 --- a/tests/test_mlflow_tracing.py +++ b/tests/test_mlflow_tracing.py @@ -14,7 +14,7 @@ # Helpers # --------------------------------------------------------------------------- -SETUP_MLFLOW = Path(__file__).parent.parent / "setup_mlflow.py" +SETUP_MLFLOW = Path(__file__).parent.parent / "setup" / "setup_mlflow.py" def run_setup_mlflow(tmp_path, env_overrides=None): diff --git a/tests/test_npm_version_pinning.py b/tests/test_npm_version_pinning.py index ee128dd..97d298b 100644 --- a/tests/test_npm_version_pinning.py +++ b/tests/test_npm_version_pinning.py @@ -322,8 +322,12 @@ class TestNpmVersionLive: """Run against real npm registry to verify the function works end-to-end.""" @pytest.mark.skipif( - not __import__("shutil").which("npm"), - reason="npm not installed" + not __import__("shutil").which("npm") or + __import__("subprocess").run( + ["npm", "view", "npm", "version"], + capture_output=True, timeout=15 + ).returncode != 0, + reason="npm not installed or not functional" ) def test_resolves_real_package(self): get_npm_version = _get_npm_version() diff --git a/tests/test_replay_attach.py b/tests/test_replay_attach.py new file mode 100644 index 0000000..482fae5 --- /dev/null +++ b/tests/test_replay_attach.py @@ -0,0 +1,53 @@ +"""Tests for /api/session/attach replay fallback.""" +import json +import os +from pathlib import Path + +import pytest + +from coda_mcp import task_manager + + +@pytest.fixture +def client(tmp_path, monkeypatch): + monkeypatch.setattr(task_manager, "SESSIONS_DIR", str(tmp_path)) + monkeypatch.setenv("MAX_CONCURRENT_SESSIONS", "5") + import app as app_module + # Set app_owner so check_authorization returns (True, None) for requests + # with no user header (same pattern used by test_session_detach.py) + app_module.app_owner = "test@example.com" + with app_module.app.test_client() as c: + yield c, tmp_path + + +def _seed_transcript(sessions_root: Path, pty_id: str, content: bytes) -> None: + sess_id = "sess-test" + task_id = "task-test" + sdir = sessions_root / sess_id + tdir = sdir / "tasks" / task_id + tdir.mkdir(parents=True) + (sdir / "session.json").write_text(json.dumps({ + "session_id": sess_id, + "pty_session_id": pty_id, + "current_task": None, + "completed_tasks": [task_id], + "status": "closed", + })) + (tdir / "transcript.log").write_bytes(content) + + +def test_attach_returns_replay_when_pty_gone_and_transcript_exists(client): + c, root = client + _seed_transcript(root, "pty-gone", b"hello\r\nworld\r\n") + resp = c.post("/api/session/attach", json={"session_id": "pty-gone"}) + assert resp.status_code == 200 + data = resp.get_json() + assert data["replay"] is True + assert data["output"] == ["hello\r\nworld\r\n"] + assert data["label"] == "hermes-mcp (replay)" + + +def test_attach_404_when_pty_gone_and_no_transcript(client): + c, root = client + resp = c.post("/api/session/attach", json={"session_id": "pty-nope"}) + assert resp.status_code == 404 diff --git a/tests/test_run_step.py b/tests/test_run_step.py new file mode 100644 index 0000000..af09733 --- /dev/null +++ b/tests/test_run_step.py @@ -0,0 +1,170 @@ +"""Tests for _run_step and _configure_all_cli_auth — env setup for subprocesses.""" + +import os +import subprocess +from unittest import mock + +import pytest + + +# We need to test _run_step from app.py. It calls subprocess.run, so we mock that. +# The function also updates setup_state, so we mock that too. + + +@pytest.fixture +def patch_app_globals(): + """Patch app.py globals needed by _run_step.""" + with mock.patch("app._update_step"): + yield + + +class TestRunStepEnvStripping: + """Verify _run_step strips OAuth credentials from subprocess env.""" + + def test_strips_databricks_client_id(self, patch_app_globals): + from app import _run_step + with mock.patch.dict(os.environ, { + "DATABRICKS_CLIENT_ID": "sp-client-id", + "DATABRICKS_CLIENT_SECRET": "sp-client-secret", + "HOME": "/tmp/test-home", + }), mock.patch("subprocess.run") as mock_run: + mock_run.return_value = subprocess.CompletedResult = mock.MagicMock( + returncode=0, stdout="ok", stderr="" + ) + _run_step("test-step", "echo hello") + + call_env = mock_run.call_args.kwargs.get("env", {}) + assert "DATABRICKS_CLIENT_ID" not in call_env + assert "DATABRICKS_CLIENT_SECRET" not in call_env + + def test_preserves_other_env_vars(self, patch_app_globals): + from app import _run_step + with mock.patch.dict(os.environ, { + "HOME": "/tmp/test-home", + "MY_CUSTOM_VAR": "keep-this", + "DATABRICKS_CLIENT_ID": "remove-this", + }), mock.patch("subprocess.run") as mock_run: + mock_run.return_value = mock.MagicMock(returncode=0, stdout="ok", stderr="") + _run_step("test-step", "echo hello") + + call_env = mock_run.call_args.kwargs.get("env", {}) + assert call_env.get("MY_CUSTOM_VAR") == "keep-this" + + +class TestRunStepPythonpath: + """Verify _run_step injects PYTHONPATH for setup script imports.""" + + def test_sets_pythonpath_to_app_dir(self, patch_app_globals): + from app import _run_step + with mock.patch.dict(os.environ, {"HOME": "/tmp/test-home"}), \ + mock.patch("subprocess.run") as mock_run: + mock_run.return_value = mock.MagicMock(returncode=0, stdout="ok", stderr="") + _run_step("test-step", "echo hello") + + call_env = mock_run.call_args.kwargs.get("env", {}) + # PYTHONPATH should contain the app directory (dirname of app.py) + assert "PYTHONPATH" in call_env + assert call_env["PYTHONPATH"] # non-empty + + def test_prepends_to_existing_pythonpath(self, patch_app_globals): + from app import _run_step + with mock.patch.dict(os.environ, { + "HOME": "/tmp/test-home", + "PYTHONPATH": "/existing/path", + }), mock.patch("subprocess.run") as mock_run: + mock_run.return_value = mock.MagicMock(returncode=0, stdout="ok", stderr="") + _run_step("test-step", "echo hello") + + call_env = mock_run.call_args.kwargs.get("env", {}) + assert "/existing/path" in call_env["PYTHONPATH"] + + +class TestRunStepPath: + """Verify _run_step adds ~/.local/bin to PATH.""" + + def test_adds_local_bin_to_path(self, patch_app_globals): + from app import _run_step + with mock.patch.dict(os.environ, { + "HOME": "/tmp/test-home", + "PATH": "/usr/bin", + }), mock.patch("subprocess.run") as mock_run: + mock_run.return_value = mock.MagicMock(returncode=0, stdout="ok", stderr="") + _run_step("test-step", "echo hello") + + call_env = mock_run.call_args.kwargs.get("env", {}) + assert "/tmp/test-home/.local/bin" in call_env["PATH"] + + def test_skips_if_already_in_path(self, patch_app_globals): + from app import _run_step + with mock.patch.dict(os.environ, { + "HOME": "/tmp/test-home", + "PATH": "/tmp/test-home/.local/bin:/usr/bin", + }), mock.patch("subprocess.run") as mock_run: + mock_run.return_value = mock.MagicMock(returncode=0, stdout="ok", stderr="") + _run_step("test-step", "echo hello") + + call_env = mock_run.call_args.kwargs.get("env", {}) + # Should not duplicate + assert call_env["PATH"].count(".local/bin") == 1 + + def test_defaults_home_when_empty(self, patch_app_globals): + """When HOME is empty or '/', should default to /app/python/source_code.""" + from app import _run_step + with mock.patch.dict(os.environ, {"HOME": ""}, clear=False), \ + mock.patch("subprocess.run") as mock_run: + mock_run.return_value = mock.MagicMock(returncode=0, stdout="ok", stderr="") + _run_step("test-step", "echo hello") + + call_env = mock_run.call_args.kwargs.get("env", {}) + assert "/app/python/source_code" in call_env.get("HOME", "") + + +# --------------------------------------------------------------------------- +# _configure_all_cli_auth — PAT reconfiguration path +# --------------------------------------------------------------------------- + +class TestConfigureAllCliAuth: + """Verify _configure_all_cli_auth injects PYTHONPATH for setup script imports. + + This is a separate code path from _run_step — it runs setup scripts via + subprocess.run after PAT rotation. Without PYTHONPATH, the scripts can't + `from utils import ...` since they live in setup/ subdirectory. + """ + + def _call_configure(self, mock_run, tmp_path, token="dapi_test"): + """Helper to call _configure_all_cli_auth with all dependencies mocked.""" + from app import _configure_all_cli_auth + # Create .claude dir so settings.json write succeeds + (tmp_path / ".claude").mkdir(exist_ok=True) + with mock.patch("utils.resolve_and_cache_gateway"), \ + mock.patch("app.get_gateway_host", return_value=None), \ + mock.patch("app.ensure_https", return_value="https://test.databricks.com"), \ + mock.patch("app.pat_rotator"), \ + mock.patch.dict(os.environ, {"HOME": str(tmp_path)}): + _configure_all_cli_auth(token) + + def test_injects_pythonpath(self, tmp_path): + with mock.patch("subprocess.run") as mock_run: + mock_run.return_value = mock.MagicMock(returncode=0, stdout="", stderr="") + self._call_configure(mock_run, tmp_path) + + # Find a subprocess call that runs a setup script + setup_calls = [c for c in mock_run.call_args_list + if any("setup/" in str(a) for a in c[0][0])] + assert len(setup_calls) > 0, "Expected subprocess calls for setup scripts" + + for call in setup_calls: + call_env = call.kwargs.get("env") or call[1].get("env", {}) + assert "PYTHONPATH" in call_env, f"PYTHONPATH missing from env for {call[0][0]}" + assert call_env["PYTHONPATH"], "PYTHONPATH should not be empty" + + def test_passes_token_in_env(self, tmp_path): + with mock.patch("subprocess.run") as mock_run: + mock_run.return_value = mock.MagicMock(returncode=0, stdout="", stderr="") + self._call_configure(mock_run, tmp_path, token="dapi_mytoken") + + setup_calls = [c for c in mock_run.call_args_list + if any("setup/" in str(a) for a in c[0][0])] + for call in setup_calls: + call_env = call.kwargs.get("env") or call[1].get("env", {}) + assert call_env.get("DATABRICKS_TOKEN") == "dapi_mytoken" diff --git a/tests/test_session_detach.py b/tests/test_session_detach.py index c381a40..6e3b60f 100644 --- a/tests/test_session_detach.py +++ b/tests/test_session_detach.py @@ -7,7 +7,6 @@ import os import subprocess -import sys import threading import time from collections import deque @@ -40,42 +39,23 @@ def test_detects_child_process_name(self): """When a shell has a child process, return the child's name.""" app_mod = _get_app() - # Launch a shell (bash) with a child process (sleep) - shell = subprocess.Popen( - ["bash", "-c", "sleep 300"], - stdin=subprocess.PIPE, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - # Give the child time to spawn - time.sleep(0.5) - - try: - result = app_mod._get_session_process(shell.pid) - assert result == "sleep", f"Expected 'sleep', got '{result}'" - finally: - shell.kill() - shell.wait() + # Mock pgrep returning a child PID, then ps resolving it to "sleep" + pgrep_result = mock.Mock(returncode=0, stdout="12345\n") + ps_result = mock.Mock(returncode=0, stdout="sleep\n") + with mock.patch("subprocess.run", side_effect=[pgrep_result, ps_result]): + result = app_mod._get_session_process(100) + assert result == "sleep", f"Expected 'sleep', got '{result}'" def test_returns_parent_process_name_when_no_children(self): """When a shell has no foreground children, return the shell name.""" app_mod = _get_app() - # Launch a bare shell that just sleeps via bash built-in wait - # Use cat which will block on stdin with no children of its own - proc = subprocess.Popen( - ["cat"], - stdin=subprocess.PIPE, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - - try: - result = app_mod._get_session_process(proc.pid) - assert result == "cat", f"Expected 'cat', got '{result}'" - finally: - proc.kill() - proc.wait() + # Mock pgrep finding no children (exit 1), then ps resolving the process itself + pgrep_result = mock.Mock(returncode=1, stdout="") + ps_result = mock.Mock(returncode=0, stdout="cat\n") + with mock.patch("subprocess.run", side_effect=[pgrep_result, ps_result]): + result = app_mod._get_session_process(100) + assert result == "cat", f"Expected 'cat', got '{result}'" def test_returns_unknown_for_dead_pid(self): """Return 'unknown' when the PID does not exist.""" @@ -230,28 +210,31 @@ def setup_app(self): app_module.sessions.clear() def test_exited_session_removed_from_dict(self): - import pty - master_fd, slave_fd = pty.openpty() + fake_master = 50 + # Use a completed process so waitpid works proc = subprocess.Popen( - ["bash", "-c", "echo hello && exit 0"], - stdin=slave_fd, stdout=slave_fd, stderr=slave_fd, - preexec_fn=os.setsid + ["bash", "-c", "exit 0"], + stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, ) - os.close(slave_fd) + proc.wait() session_id = "sess-eof-test" with self.app_module.sessions_lock: self.app_module.sessions[session_id] = { "pid": proc.pid, - "master_fd": master_fd, + "master_fd": fake_master, "output_buffer": deque(maxlen=1000), "lock": threading.Lock(), "last_poll_time": time.time(), "created_at": time.time(), } - # read_pty_output should detect EOF and call terminate_session - self.app_module.read_pty_output(session_id, master_fd) + # Simulate EOF: select says readable, os.read returns empty bytes + with mock.patch("select.select", return_value=([fake_master], [], [])), \ + mock.patch("os.read", return_value=b""), \ + mock.patch("os.close"), \ + mock.patch("os.kill"): + self.app_module.read_pty_output(session_id, fake_master) with self.app_module.sessions_lock: assert session_id not in self.app_module.sessions diff --git a/tests/test_sync_to_workspace.py b/tests/test_sync_to_workspace.py new file mode 100644 index 0000000..6faedf4 --- /dev/null +++ b/tests/test_sync_to_workspace.py @@ -0,0 +1,181 @@ +"""Tests for sync_to_workspace — path-escape guard and workspace sync.""" + +import subprocess +from pathlib import Path +from unittest import mock + +import pytest + + +# --------------------------------------------------------------------------- +# _read_databrickscfg +# --------------------------------------------------------------------------- + +class TestReadDatabrickscfg: + def test_reads_host_and_token(self, tmp_path): + cfg = tmp_path / ".databrickscfg" + cfg.write_text("[DEFAULT]\nhost = https://test.cloud.databricks.com\ntoken = dapi_abc123\n") + with mock.patch("sync_to_workspace.Path.home", return_value=tmp_path): + from sync_to_workspace import _read_databrickscfg + host, token = _read_databrickscfg() + assert host == "https://test.cloud.databricks.com" + assert token == "dapi_abc123" + + def test_returns_none_when_missing(self, tmp_path): + with mock.patch("sync_to_workspace.Path.home", return_value=tmp_path): + from sync_to_workspace import _read_databrickscfg + host, token = _read_databrickscfg() + assert host is None + assert token is None + + def test_returns_none_for_missing_keys(self, tmp_path): + cfg = tmp_path / ".databrickscfg" + cfg.write_text("[DEFAULT]\n# empty section\n") + with mock.patch("sync_to_workspace.Path.home", return_value=tmp_path): + from sync_to_workspace import _read_databrickscfg + host, token = _read_databrickscfg() + assert host is None + assert token is None + + +# --------------------------------------------------------------------------- +# get_user_email +# --------------------------------------------------------------------------- + +class TestGetUserEmail: + def test_raises_when_no_config(self, tmp_path): + from sync_to_workspace import get_user_email + with mock.patch("sync_to_workspace._read_databrickscfg", return_value=(None, None)): + with pytest.raises(RuntimeError, match="missing host or token"): + get_user_email() + + def test_raises_when_no_token(self): + from sync_to_workspace import get_user_email + with mock.patch("sync_to_workspace._read_databrickscfg", return_value=("https://host", None)): + with pytest.raises(RuntimeError, match="missing host or token"): + get_user_email() + + def test_returns_email(self): + from sync_to_workspace import get_user_email + mock_user = mock.MagicMock() + mock_user.user_name = "test@example.com" + mock_client = mock.MagicMock() + mock_client.current_user.me.return_value = mock_user + with mock.patch("sync_to_workspace._read_databrickscfg", return_value=("https://host", "tok")): + with mock.patch("sync_to_workspace.WorkspaceClient", return_value=mock_client): + email = get_user_email() + assert email == "test@example.com" + + +# --------------------------------------------------------------------------- +# sync_project — path-escape guard +# --------------------------------------------------------------------------- + +class TestSyncProject: + def test_rejects_path_outside_projects_dir(self, tmp_path, capsys): + from sync_to_workspace import sync_project + # Create a path outside ~/projects/ + outside = tmp_path / "evil-repo" + outside.mkdir() + with mock.patch("sync_to_workspace.Path.home", return_value=tmp_path): + sync_project(outside) + captured = capsys.readouterr() + assert "SKIP" in captured.err + assert "outside" in captured.err + + def test_accepts_path_inside_projects_dir(self, tmp_path): + from sync_to_workspace import sync_project + projects = tmp_path / "projects" + projects.mkdir() + repo = projects / "my-repo" + repo.mkdir() + + mock_user = mock.MagicMock() + mock_user.user_name = "test@example.com" + mock_client = mock.MagicMock() + mock_client.current_user.me.return_value = mock_user + + with mock.patch("sync_to_workspace.Path.home", return_value=tmp_path), \ + mock.patch("sync_to_workspace._read_databrickscfg", return_value=("https://host", "tok")), \ + mock.patch("sync_to_workspace.WorkspaceClient", return_value=mock_client), \ + mock.patch("sync_to_workspace.subprocess.run") as mock_run: + mock_run.return_value = subprocess.CompletedProcess([], 0, stdout="", stderr="") + sync_project(repo) + + mock_run.assert_called_once() + args = mock_run.call_args + assert "databricks" in args[0][0][0] + assert "sync" in args[0][0][1] + + def test_strips_oauth_env_from_subprocess(self, tmp_path): + """Verify OAuth credentials are stripped so CLI falls through to ~/.databrickscfg.""" + from sync_to_workspace import sync_project + projects = tmp_path / "projects" + projects.mkdir() + repo = projects / "my-repo" + repo.mkdir() + + mock_user = mock.MagicMock() + mock_user.user_name = "test@example.com" + mock_client = mock.MagicMock() + mock_client.current_user.me.return_value = mock_user + + with mock.patch("sync_to_workspace.Path.home", return_value=tmp_path), \ + mock.patch("sync_to_workspace._read_databrickscfg", return_value=("https://host", "tok")), \ + mock.patch("sync_to_workspace.WorkspaceClient", return_value=mock_client), \ + mock.patch("sync_to_workspace.subprocess.run") as mock_run, \ + mock.patch.dict("os.environ", { + "DATABRICKS_CLIENT_ID": "sp-id", + "DATABRICKS_CLIENT_SECRET": "sp-secret", + "DATABRICKS_HOST": "https://host", + "DATABRICKS_TOKEN": "dapi_tok", + }): + mock_run.return_value = subprocess.CompletedProcess([], 0, stdout="", stderr="") + sync_project(repo) + + call_env = mock_run.call_args[1].get("env") or mock_run.call_args.kwargs.get("env", {}) + assert "DATABRICKS_CLIENT_ID" not in call_env + assert "DATABRICKS_CLIENT_SECRET" not in call_env + assert "DATABRICKS_HOST" not in call_env + assert "DATABRICKS_TOKEN" not in call_env + + def test_logs_error_on_failure(self, tmp_path, capsys): + from sync_to_workspace import sync_project + projects = tmp_path / "projects" + projects.mkdir() + repo = projects / "my-repo" + repo.mkdir() + + with mock.patch("sync_to_workspace.Path.home", return_value=tmp_path), \ + mock.patch("sync_to_workspace.get_user_email", side_effect=Exception("auth failed")): + sync_project(repo) + + captured = capsys.readouterr() + assert "Sync failed" in captured.err + # Error should be logged to file + error_log = tmp_path / ".sync-errors.log" + assert error_log.exists() + assert "auth failed" in error_log.read_text() + + def test_sync_failure_warns(self, tmp_path, capsys): + """Non-zero return code from databricks sync should print warning.""" + from sync_to_workspace import sync_project + projects = tmp_path / "projects" + projects.mkdir() + repo = projects / "my-repo" + repo.mkdir() + + mock_user = mock.MagicMock() + mock_user.user_name = "test@example.com" + mock_client = mock.MagicMock() + mock_client.current_user.me.return_value = mock_user + + with mock.patch("sync_to_workspace.Path.home", return_value=tmp_path), \ + mock.patch("sync_to_workspace._read_databrickscfg", return_value=("https://host", "tok")), \ + mock.patch("sync_to_workspace.WorkspaceClient", return_value=mock_client), \ + mock.patch("sync_to_workspace.subprocess.run") as mock_run: + mock_run.return_value = subprocess.CompletedProcess([], 1, stdout="", stderr="permission denied") + sync_project(repo) + + captured = capsys.readouterr() + assert "Sync warning" in captured.err diff --git a/tests/test_task_manager.py b/tests/test_task_manager.py new file mode 100644 index 0000000..c0e0d7a --- /dev/null +++ b/tests/test_task_manager.py @@ -0,0 +1,541 @@ +"""Tests for task_manager — disk-based MCP session/task state.""" + +import json +import os +import time +from unittest import mock + +import pytest + + +@pytest.fixture(autouse=True) +def isolated_sessions(tmp_path): + """Point task_manager.SESSIONS_DIR at a temp dir.""" + sessions_dir = str(tmp_path / ".coda" / "sessions") + with mock.patch("coda_mcp.task_manager.SESSIONS_DIR", sessions_dir): + yield sessions_dir + + +# ── helpers ────────────────────────────────────────────────────────── + + +def _read_json(path): + with open(path) as f: + return json.load(f) + + +def _read_text(path): + with open(path) as f: + return f.read() + + +def _read_jsonl(path): + lines = [] + with open(path) as f: + for line in f: + line = line.strip() + if line: + lines.append(json.loads(line)) + return lines + + +# ── Session lifecycle ──────────────────────────────────────────────── + + +class TestCreateSession: + def test_returns_session_id_and_status(self): + from coda_mcp import task_manager + + result = task_manager.create_session("a@b.com", "u1", "my-label") + assert result["status"] == "ready" + assert result["session_id"].startswith("sess-") + assert len(result["session_id"]) == 5 + 12 # "sess-" + 12 hex + + def test_creates_session_json_on_disk(self, isolated_sessions): + from coda_mcp import task_manager + + result = task_manager.create_session("a@b.com", "u1", "my-label") + sid = result["session_id"] + path = os.path.join(isolated_sessions, sid, "session.json") + assert os.path.isfile(path) + data = _read_json(path) + assert data["email"] == "a@b.com" + assert data["user_id"] == "u1" + assert data["label"] == "my-label" + assert data["status"] == "ready" + assert data["current_task"] is None + assert data["completed_tasks"] == [] + assert "created_at" in data + + def test_unique_ids(self): + from coda_mcp import task_manager + + ids = {task_manager.create_session("a@b.com", "u1")["session_id"] for _ in range(20)} + assert len(ids) == 20 + + +class TestCloseSession: + def test_marks_session_closed(self, isolated_sessions): + from coda_mcp import task_manager + + sid = task_manager.create_session("a@b.com", "u1")["session_id"] + task_manager.close_session(sid) + data = _read_json(os.path.join(isolated_sessions, sid, "session.json")) + assert data["status"] == "closed" + + def test_close_nonexistent_raises(self): + from coda_mcp import task_manager + + with pytest.raises(task_manager.SessionNotFoundError): + task_manager.close_session("sess-doesnotexist") + + +class TestReadSession: + def test_read_existing(self, isolated_sessions): + from coda_mcp import task_manager + + sid = task_manager.create_session("a@b.com", "u1", "lbl")["session_id"] + data = task_manager._read_session(sid) + assert data["email"] == "a@b.com" + + def test_read_nonexistent_raises(self): + from coda_mcp import task_manager + + with pytest.raises(task_manager.SessionNotFoundError): + task_manager._read_session("sess-000000000000") + + +class TestUpdateSessionField: + def test_updates_single_field(self, isolated_sessions): + from coda_mcp import task_manager + + sid = task_manager.create_session("a@b.com", "u1")["session_id"] + task_manager._update_session_field(sid, "status", "busy") + data = task_manager._read_session(sid) + assert data["status"] == "busy" + + def test_preserves_other_fields(self, isolated_sessions): + from coda_mcp import task_manager + + sid = task_manager.create_session("a@b.com", "u1", "lbl")["session_id"] + task_manager._update_session_field(sid, "status", "busy") + data = task_manager._read_session(sid) + assert data["email"] == "a@b.com" + assert data["label"] == "lbl" + + +# ── Task lifecycle ─────────────────────────────────────────────────── + + +class TestCreateTask: + def test_returns_task_id_and_running(self): + from coda_mcp import task_manager + + sid = task_manager.create_session("a@b.com", "u1")["session_id"] + result = task_manager.create_task(sid, "do something", "a@b.com") + assert result["status"] == "running" + assert result["task_id"].startswith("task-") + assert len(result["task_id"]) == 5 + 8 # "task-" + 8 hex + + def test_creates_task_directory_with_files(self, isolated_sessions): + from coda_mcp import task_manager + + sid = task_manager.create_session("a@b.com", "u1")["session_id"] + tid = task_manager.create_task(sid, "do something", "a@b.com")["task_id"] + task_dir = task_manager._task_dir(sid, tid) + assert os.path.isdir(task_dir) + assert os.path.isfile(os.path.join(task_dir, "prompt.txt")) + assert os.path.isfile(os.path.join(task_dir, "status.jsonl")) + + def test_prompt_txt_contains_wrapped_prompt(self, isolated_sessions): + from coda_mcp import task_manager + + sid = task_manager.create_session("a@b.com", "u1")["session_id"] + tid = task_manager.create_task(sid, "fix the bug", "a@b.com")["task_id"] + prompt = _read_text(os.path.join(task_manager._task_dir(sid, tid), "prompt.txt")) + assert "---CODA-TASK---" in prompt + assert "fix the bug" in prompt + + def test_session_marked_busy(self, isolated_sessions): + from coda_mcp import task_manager + + sid = task_manager.create_session("a@b.com", "u1")["session_id"] + task_manager.create_task(sid, "do it", "a@b.com") + data = task_manager._read_session(sid) + assert data["status"] == "busy" + + def test_session_current_task_set(self, isolated_sessions): + from coda_mcp import task_manager + + sid = task_manager.create_session("a@b.com", "u1")["session_id"] + tid = task_manager.create_task(sid, "do it", "a@b.com")["task_id"] + data = task_manager._read_session(sid) + assert data["current_task"] == tid + + def test_busy_session_raises(self): + from coda_mcp import task_manager + + sid = task_manager.create_session("a@b.com", "u1")["session_id"] + task_manager.create_task(sid, "first", "a@b.com") + with pytest.raises(task_manager.SessionBusyError): + task_manager.create_task(sid, "second", "a@b.com") + + def test_nonexistent_session_raises(self): + from coda_mcp import task_manager + + with pytest.raises(task_manager.SessionNotFoundError): + task_manager.create_task("sess-doesnotexist", "p", "e@x.com") + + def test_status_jsonl_has_initial_entry(self, isolated_sessions): + from coda_mcp import task_manager + + sid = task_manager.create_session("a@b.com", "u1")["session_id"] + tid = task_manager.create_task(sid, "go", "a@b.com")["task_id"] + entries = _read_jsonl( + os.path.join(task_manager._task_dir(sid, tid), "status.jsonl") + ) + assert len(entries) == 1 + assert entries[0]["status"] == "running" + + def test_optional_params_stored(self, isolated_sessions): + from coda_mcp import task_manager + + sid = task_manager.create_session("a@b.com", "u1")["session_id"] + tid = task_manager.create_task( + sid, "go", "a@b.com", + context={"repo": "myrepo"}, + context_hint="look at utils.py", + timeout_s=120, + permissions=["read", "write"], + )["task_id"] + prompt = _read_text(os.path.join(task_manager._task_dir(sid, tid), "prompt.txt")) + assert "myrepo" in prompt + assert "utils.py" in prompt + + +class TestTaskDir: + def test_returns_correct_path(self, isolated_sessions): + from coda_mcp import task_manager + + path = task_manager._task_dir("sess-aabbccddee01", "task-11223344") + expected = os.path.join( + isolated_sessions, "sess-aabbccddee01", "tasks", "task-11223344" + ) + assert path == expected + + +# ── Task status / result ───────────────────────────────────────────── + + +class TestGetTaskStatus: + def test_returns_latest_status(self, isolated_sessions): + from coda_mcp import task_manager + + sid = task_manager.create_session("a@b.com", "u1")["session_id"] + tid = task_manager.create_task(sid, "go", "a@b.com")["task_id"] + status = task_manager.get_task_status(tid, sid) + assert status["status"] == "running" + + def test_reads_appended_lines(self, isolated_sessions): + from coda_mcp import task_manager + + sid = task_manager.create_session("a@b.com", "u1")["session_id"] + tid = task_manager.create_task(sid, "go", "a@b.com")["task_id"] + # simulate agent appending progress + status_path = os.path.join(task_manager._task_dir(sid, tid), "status.jsonl") + with open(status_path, "a") as f: + f.write(json.dumps({"status": "progress", "pct": 50, "ts": time.time()}) + "\n") + status = task_manager.get_task_status(tid, sid) + assert status["status"] == "progress" + assert status["pct"] == 50 + + def test_missing_task_returns_not_found(self, isolated_sessions): + from coda_mcp import task_manager + + sid = task_manager.create_session("a@b.com", "u1")["session_id"] + status = task_manager.get_task_status("task-nonexist", sid) + assert status["status"] == "not_found" + + +class TestGetTaskResult: + def test_returns_result_when_present(self, isolated_sessions): + from coda_mcp import task_manager + + sid = task_manager.create_session("a@b.com", "u1")["session_id"] + tid = task_manager.create_task(sid, "go", "a@b.com")["task_id"] + # simulate agent writing result + result_path = os.path.join(task_manager._task_dir(sid, tid), "result.json") + with open(result_path, "w") as f: + json.dump({"answer": 42}, f) + result = task_manager.get_task_result(tid, sid) + assert result["answer"] == 42 + + def test_returns_none_when_absent(self, isolated_sessions): + from coda_mcp import task_manager + + sid = task_manager.create_session("a@b.com", "u1")["session_id"] + tid = task_manager.create_task(sid, "go", "a@b.com")["task_id"] + result = task_manager.get_task_result(tid, sid) + assert result is None + + def test_missing_task_returns_none(self, isolated_sessions): + from coda_mcp import task_manager + + sid = task_manager.create_session("a@b.com", "u1")["session_id"] + result = task_manager.get_task_result("task-nonexist", sid) + assert result is None + + +# ── Complete task ───────────────────────────────────────────────────── + + +class TestCompleteTask: + def test_marks_session_closed(self, isolated_sessions): + """v2: sessions are ephemeral — complete_task auto-closes the session.""" + from coda_mcp import task_manager + + sid = task_manager.create_session("a@b.com", "u1")["session_id"] + tid = task_manager.create_task(sid, "go", "a@b.com")["task_id"] + task_manager.complete_task(sid, tid) + data = task_manager._read_session(sid) + assert data["status"] == "closed" + assert "closed_at" in data + + def test_appends_to_completed_tasks(self, isolated_sessions): + from coda_mcp import task_manager + + sid = task_manager.create_session("a@b.com", "u1")["session_id"] + tid = task_manager.create_task(sid, "go", "a@b.com")["task_id"] + task_manager.complete_task(sid, tid) + data = task_manager._read_session(sid) + assert tid in data["completed_tasks"] + + def test_closed_session_rejects_new_task(self, isolated_sessions): + """v2: ephemeral sessions — new tasks need new sessions.""" + from coda_mcp import task_manager + + sid = task_manager.create_session("a@b.com", "u1")["session_id"] + tid1 = task_manager.create_task(sid, "first", "a@b.com")["task_id"] + task_manager.complete_task(sid, tid1) + with pytest.raises(task_manager.SessionNotFoundError): + task_manager.create_task(sid, "second", "a@b.com") + + def test_appends_done_to_status_jsonl(self, isolated_sessions): + from coda_mcp import task_manager + + sid = task_manager.create_session("a@b.com", "u1")["session_id"] + tid = task_manager.create_task(sid, "go", "a@b.com")["task_id"] + task_manager.complete_task(sid, tid) + entries = _read_jsonl( + os.path.join(task_manager._task_dir(sid, tid), "status.jsonl") + ) + assert entries[-1]["status"] == "done" + + def test_nonexistent_session_raises(self): + from coda_mcp import task_manager + + with pytest.raises(task_manager.SessionNotFoundError): + task_manager.complete_task("sess-doesnotexist", "task-00000000") + + +# ── Prompt wrapping ────────────────────────────────────────────────── + + +class TestWrapPrompt: + def test_contains_marker(self): + from coda_mcp import task_manager + + wrapped = task_manager.wrap_prompt( + task_id="task-aabbccdd", + session_id="sess-112233445566", + email="a@b.com", + prompt="fix the bug", + context=None, + results_dir="/tmp/r", + context_hint=None, + ) + assert "---CODA-TASK---" in wrapped + assert "fix the bug" in wrapped + assert "task-aabbccdd" in wrapped + assert "sess-112233445566" in wrapped + assert "a@b.com" in wrapped + assert "/tmp/r" in wrapped + + def test_includes_context_when_provided(self): + from coda_mcp import task_manager + + wrapped = task_manager.wrap_prompt( + task_id="task-aabbccdd", + session_id="sess-112233445566", + email="a@b.com", + prompt="go", + context={"repo": "myrepo", "branch": "main"}, + results_dir="/tmp/r", + context_hint=None, + ) + assert "myrepo" in wrapped + assert "main" in wrapped + + def test_includes_context_hint(self): + from coda_mcp import task_manager + + wrapped = task_manager.wrap_prompt( + task_id="task-aabbccdd", + session_id="sess-112233445566", + email="a@b.com", + prompt="go", + context=None, + results_dir="/tmp/r", + context_hint="look at utils.py first", + ) + assert "look at utils.py first" in wrapped + + def test_no_context_still_valid(self): + from coda_mcp import task_manager + + wrapped = task_manager.wrap_prompt( + task_id="task-aabbccdd", + session_id="sess-112233445566", + email="a@b.com", + prompt="hello", + context=None, + results_dir="/tmp/r", + context_hint=None, + ) + assert "---CODA-TASK---" in wrapped + assert "hello" in wrapped + + +# ── Edge cases ──────────────────────────────────────────────────────── + + +class TestEdgeCases: + def test_closed_session_rejects_task(self, isolated_sessions): + from coda_mcp import task_manager + + sid = task_manager.create_session("a@b.com", "u1")["session_id"] + task_manager.close_session(sid) + with pytest.raises(task_manager.SessionNotFoundError): + task_manager.create_task(sid, "go", "a@b.com") + + def test_multiple_tasks_across_sessions(self, isolated_sessions): + """v2: each task gets its own ephemeral session; all appear in list_all_tasks.""" + from coda_mcp import task_manager + + tids = [] + for i in range(3): + sid = task_manager.create_session("a@b.com", "u1")["session_id"] + tid = task_manager.create_task(sid, f"task {i}", "a@b.com")["task_id"] + task_manager.complete_task(sid, tid) + tids.append(tid) + # Each session auto-closes + data = task_manager._read_session(sid) + assert data["status"] == "closed" + + all_tasks = task_manager.list_all_tasks() + all_tids = [t["task_id"] for t in all_tasks] + for tid in tids: + assert tid in all_tids + + def test_corrupt_session_json_raises(self, isolated_sessions): + from coda_mcp import task_manager + + sid = task_manager.create_session("a@b.com", "u1")["session_id"] + path = os.path.join(isolated_sessions, sid, "session.json") + with open(path, "w") as f: + f.write("{bad json") + with pytest.raises(task_manager.SessionNotFoundError): + task_manager._read_session(sid) + + +# ── find_task_dir_by_pty_session ───────────────────────────────────── + + +@pytest.fixture +def sessions_root(tmp_path, monkeypatch): + from coda_mcp import task_manager + monkeypatch.setattr(task_manager, "SESSIONS_DIR", str(tmp_path)) + # Reset the lookup cache between tests + task_manager._pty_lookup_cache.clear() + return tmp_path + + +def _make_session_dir(root, sess_id, pty_id, current_task=None, completed=None): + sdir = root / sess_id + (sdir / "tasks").mkdir(parents=True) + data = { + "session_id": sess_id, + "pty_session_id": pty_id, + "current_task": current_task, + "completed_tasks": completed or [], + "status": "ready", + } + (sdir / "session.json").write_text(json.dumps(data)) + return sdir + + +def test_find_task_dir_hits_current_task(sessions_root): + from coda_mcp import task_manager + + _make_session_dir(sessions_root, "sess-A", "pty-1", current_task="task-X") + result = task_manager.find_task_dir_by_pty_session("pty-1") + assert result == str(sessions_root / "sess-A" / "tasks" / "task-X") + + +def test_find_task_dir_falls_back_to_last_completed(sessions_root): + from coda_mcp import task_manager + + _make_session_dir( + sessions_root, "sess-A", "pty-1", + current_task=None, + completed=["task-old", "task-recent"], + ) + result = task_manager.find_task_dir_by_pty_session("pty-1") + assert result == str(sessions_root / "sess-A" / "tasks" / "task-recent") + + +def test_find_task_dir_returns_none_when_no_match(sessions_root): + from coda_mcp import task_manager + + _make_session_dir(sessions_root, "sess-A", "pty-1", current_task="task-X") + assert task_manager.find_task_dir_by_pty_session("pty-NONEXIST") is None + + +def test_find_task_dir_ignores_corrupt_session_json(sessions_root): + from coda_mcp import task_manager + + sdir = sessions_root / "sess-bad" + sdir.mkdir() + (sdir / "session.json").write_text("not json {{{") + _make_session_dir(sessions_root, "sess-good", "pty-1", current_task="task-X") + assert task_manager.find_task_dir_by_pty_session("pty-1") == \ + str(sessions_root / "sess-good" / "tasks" / "task-X") + + +def test_find_task_dir_cache_hits_within_ttl(sessions_root): + from coda_mcp import task_manager + + _make_session_dir(sessions_root, "sess-A", "pty-1", current_task="task-X") + task_manager.find_task_dir_by_pty_session("pty-1") + # Remove session.json — cache should still return the hit + (sessions_root / "sess-A" / "session.json").unlink() + assert task_manager.find_task_dir_by_pty_session("pty-1") == \ + str(sessions_root / "sess-A" / "tasks" / "task-X") + + +def test_find_task_dir_cache_expires(sessions_root, monkeypatch): + from coda_mcp import task_manager + + monkeypatch.setattr(task_manager, "_PTY_LOOKUP_TTL", 0.01) + _make_session_dir(sessions_root, "sess-A", "pty-1", current_task="task-X") + task_manager.find_task_dir_by_pty_session("pty-1") + (sessions_root / "sess-A" / "session.json").unlink() + time.sleep(0.02) + assert task_manager.find_task_dir_by_pty_session("pty-1") is None + + +def test_find_task_dir_no_sessions_dir(sessions_root, monkeypatch): + from coda_mcp import task_manager + + monkeypatch.setattr(task_manager, "SESSIONS_DIR", "/nonexistent/path/that/does/not/exist") + assert task_manager.find_task_dir_by_pty_session("pty-1") is None diff --git a/tests/test_transcript.py b/tests/test_transcript.py new file mode 100644 index 0000000..c993a0a --- /dev/null +++ b/tests/test_transcript.py @@ -0,0 +1,176 @@ +"""Unit tests for the transcript tee in read_pty_output. + +These tests exercise the tee logic directly by simulating output dispatch into +a synthesized session dict and a real on-disk transcript file. The full PTY +read loop is not exercised here — see test_mcp_integration.py for E2E. +""" +import os +import stat +import threading +from pathlib import Path + +import pytest + +# The three tests that hit mcp_create_pty_session call pty.openpty(), which +# fails in headless CI containers without TTY allocators. Mark those tests +# explicitly so existing fixture-based tests (test_tee_*) keep running. +def _pty_is_usable() -> bool: + if not hasattr(os, "openpty"): + return False + try: + master, slave = os.openpty() + os.close(master) + os.close(slave) + return True + except OSError: + return False + + +_pty_available = _pty_is_usable() +_pty_skip = pytest.mark.skipif(not _pty_available, reason="pty.openpty() not available") + + +@pytest.fixture +def session_dict(tmp_path): + """Build a minimally valid sessions[pty_id] entry with a real transcript handle.""" + transcript = tmp_path / "transcript.log" + fh = open(transcript, "ab", buffering=0) + os.fchmod(fh.fileno(), 0o600) + return { + "transcript_path": str(transcript), + "transcript_fh": fh, + "transcript_bytes": 0, + "lock": threading.Lock(), + } + + +def _write_chunk(session, output: bytes, cap: int = 10 * 1024 * 1024) -> None: + """Mirror the tee logic from read_pty_output for unit testing.""" + from app import _tee_transcript_chunk + _tee_transcript_chunk(session, output, cap=cap) + + +def test_tee_writes_bytes_and_flushes(session_dict): + _write_chunk(session_dict, b"hello world\n") + assert session_dict["transcript_bytes"] == 12 + assert Path(session_dict["transcript_path"]).read_bytes() == b"hello world\n" + + +def test_tee_chmod_is_0600(session_dict): + mode = stat.S_IMODE(os.stat(session_dict["transcript_path"]).st_mode) + assert mode == 0o600 + + +def test_tee_truncation_at_cap(session_dict): + cap = 16 + _write_chunk(session_dict, b"AAAAAAAAAA", cap=cap) + _write_chunk(session_dict, b"BBBBBBBBBBBBBBBBBBBB", cap=cap) + body = Path(session_dict["transcript_path"]).read_bytes() + # 10 A's, then 6 B's, then truncation marker. + assert body.startswith(b"AAAAAAAAAABBBBBB") + assert b"[transcript truncated at" in body + # Handle is closed after marker + assert session_dict["transcript_fh"] is None + + +def test_tee_no_op_when_fh_is_none(session_dict): + session_dict["transcript_fh"] = None + _write_chunk(session_dict, b"should not write") + assert Path(session_dict["transcript_path"]).read_bytes() == b"" + + +def test_tee_handles_write_error(session_dict, monkeypatch): + # Close the handle out from under the tee — write() will ValueError. + session_dict["transcript_fh"].close() + _write_chunk(session_dict, b"this will fail") + # Handle replaced with None; no crash. + assert session_dict["transcript_fh"] is None + + +@_pty_skip +def test_mcp_create_pty_session_opens_transcript_when_path_given(tmp_path, monkeypatch): + monkeypatch.setattr("app.MAX_CONCURRENT_SESSIONS", 5) + transcript = tmp_path / "transcript.log" + from app import mcp_create_pty_session, sessions, mcp_close_pty_session + sid = mcp_create_pty_session(label="test", transcript_path=str(transcript)) + try: + assert transcript.exists() + mode = stat.S_IMODE(os.stat(transcript).st_mode) + assert mode == 0o600 + sess = sessions[sid] + assert sess["transcript_path"] == str(transcript) + assert sess["transcript_fh"] is not None + assert sess["transcript_bytes"] == 0 + finally: + mcp_close_pty_session(sid) + + +@_pty_skip +def test_mcp_create_pty_session_no_transcript_when_path_none(monkeypatch): + monkeypatch.setattr("app.MAX_CONCURRENT_SESSIONS", 5) + from app import mcp_create_pty_session, sessions, mcp_close_pty_session + sid = mcp_create_pty_session(label="test") + try: + sess = sessions[sid] + assert sess.get("transcript_fh") is None + assert sess.get("transcript_path") is None + finally: + mcp_close_pty_session(sid) + + +@_pty_skip +def test_terminate_session_closes_transcript_handle(tmp_path, monkeypatch): + monkeypatch.setattr("app.MAX_CONCURRENT_SESSIONS", 5) + transcript = tmp_path / "transcript.log" + from app import mcp_create_pty_session, sessions, mcp_close_pty_session + sid = mcp_create_pty_session(label="test", transcript_path=str(transcript)) + fh = sessions[sid]["transcript_fh"] + mcp_close_pty_session(sid) + assert fh.closed + # Session removed from dict + assert sid not in sessions + + +@_pty_skip +def test_grace_period_pty_does_not_count_toward_max(monkeypatch): + monkeypatch.setattr("app.MAX_CONCURRENT_SESSIONS", 2) + from app import mcp_create_pty_session, mcp_close_pty_session, sessions, _mark_grace_for_session + + sid1 = mcp_create_pty_session(label="t1") + sid2 = mcp_create_pty_session(label="t2") + try: + # At cap. A third creation should raise. + with pytest.raises(RuntimeError, match="Maximum"): + mcp_create_pty_session(label="t3") + # Mark one as grace; now we should have headroom. + _mark_grace_for_session(sid1) + assert sessions[sid1]["grace"] is True + sid3 = mcp_create_pty_session(label="t3") + mcp_close_pty_session(sid3) + finally: + for s in [sid1, sid2]: + try: mcp_close_pty_session(s) + except Exception: pass + + +@_pty_skip +def test_bump_session_last_poll_advances_clock(monkeypatch): + monkeypatch.setattr("app.MAX_CONCURRENT_SESSIONS", 5) + from app import mcp_create_pty_session, mcp_close_pty_session, sessions, _bump_session_last_poll + sid = mcp_create_pty_session(label="t") + try: + baseline = sessions[sid]["last_poll_time"] + _bump_session_last_poll(sid, 300) + assert sessions[sid]["last_poll_time"] >= baseline + 299 + finally: + mcp_close_pty_session(sid) + + +def test_mark_grace_on_missing_session_is_noop(): + from app import _mark_grace_for_session + _mark_grace_for_session("nonexistent-pty-id") # must not raise + + +def test_bump_session_last_poll_missing_is_noop(): + from app import _bump_session_last_poll + _bump_session_last_poll("nonexistent-pty-id", 100) # must not raise diff --git a/tests/test_url_builder.py b/tests/test_url_builder.py new file mode 100644 index 0000000..907287e --- /dev/null +++ b/tests/test_url_builder.py @@ -0,0 +1,82 @@ +"""Tests for url_builder module — base URL resolution for viewer_url.""" +import os +import importlib +from unittest import mock + +import pytest + + +@pytest.fixture(autouse=True) +def _reset_module(): + """Re-import url_builder fresh for each test (module-level cache).""" + from coda_mcp import url_builder + importlib.reload(url_builder) + yield + + +def test_returns_none_when_neither_env_nor_cache(): + from coda_mcp import url_builder + assert url_builder.build_viewer_url("pty-1") is None + + +def test_env_override_wins(): + from coda_mcp import url_builder + with mock.patch.dict(os.environ, {"CODA_APP_URL": "https://override.example.com"}): + assert url_builder.build_viewer_url("pty-1") == \ + "https://override.example.com/?session=pty-1" + + +def test_env_override_strips_trailing_slash(): + from coda_mcp import url_builder + with mock.patch.dict(os.environ, {"CODA_APP_URL": "https://override.example.com/"}): + assert url_builder.build_viewer_url("pty-1") == \ + "https://override.example.com/?session=pty-1" + + +def test_header_capture_used_when_no_env(): + from coda_mcp import url_builder + url_builder.capture_from_headers("app.databricksapps.com") + assert url_builder.build_viewer_url("pty-1") == \ + "https://app.databricksapps.com/?session=pty-1" + + +def test_env_overrides_header_capture(): + from coda_mcp import url_builder + url_builder.capture_from_headers("captured.example.com") + with mock.patch.dict(os.environ, {"CODA_APP_URL": "https://override.example.com"}): + assert url_builder.build_viewer_url("pty-1") == \ + "https://override.example.com/?session=pty-1" + + +def test_header_capture_overwrites_previous(): + from coda_mcp import url_builder + url_builder.capture_from_headers("first.example.com") + url_builder.capture_from_headers("second.example.com") + assert "second.example.com" in url_builder.build_viewer_url("pty-1") + + +def test_capture_empty_string_does_not_overwrite(): + from coda_mcp import url_builder + url_builder.capture_from_headers("good.example.com") + url_builder.capture_from_headers("") + assert "good.example.com" in url_builder.build_viewer_url("pty-1") + + +def test_capture_none_does_not_crash(): + from coda_mcp import url_builder + url_builder.capture_from_headers(None) + assert url_builder.build_viewer_url("pty-1") is None + + +def test_capture_strips_scheme_prefix(): + from coda_mcp import url_builder + url_builder.capture_from_headers("https://app.example.com") + assert url_builder._app_url_cache == "app.example.com" + assert url_builder.build_viewer_url("pty-1") == "https://app.example.com/?session=pty-1" + + +def test_capture_strips_http_scheme_prefix(): + from coda_mcp import url_builder + url_builder.capture_from_headers("http://app.example.com/") + # http stripped, trailing slash stripped + assert url_builder._app_url_cache == "app.example.com" diff --git a/tools/coda-bridge.py b/tools/coda-bridge.py new file mode 100644 index 0000000..c67b54c --- /dev/null +++ b/tools/coda-bridge.py @@ -0,0 +1,118 @@ +#!/usr/bin/env python3 +"""Stdio-to-HTTP MCP bridge with Databricks OAuth token injection. + +Proxies MCP JSON-RPC (stdio) to a Databricks App (Streamable HTTP), +injecting fresh OAuth tokens via `databricks auth token`. + +Config via environment variables (set in Claude Code settings.json): + + CODA_MCP_URL — App MCP endpoint URL + DATABRICKS_PROFILE — Databricks CLI profile for auth +""" + +import json +import os +import subprocess +import sys +import time +import urllib.request +import urllib.error + +APP_URL = os.environ.get("CODA_MCP_URL", "") +PROFILE = os.environ.get("DATABRICKS_PROFILE", "DEFAULT") +TOKEN_TTL = 1800 # cache 30 min (tokens last 60) + +_cache = {"token": None, "expires_at": 0.0} +_session_id = None + + +def _log(msg): + print(f"[coda-bridge] {msg}", file=sys.stderr, flush=True) + + +def _get_token(force=False): + now = time.time() + if not force and _cache["token"] and now < _cache["expires_at"]: + return _cache["token"] + result = subprocess.run( + ["databricks", "auth", "token", "-p", PROFILE], + capture_output=True, text=True, timeout=15, + ) + if result.returncode != 0: + raise RuntimeError(f"databricks auth token failed: {result.stderr.strip()}") + data = json.loads(result.stdout) + _cache["token"] = data["access_token"] + _cache["expires_at"] = now + TOKEN_TTL + _log("OAuth token refreshed") + return _cache["token"] + + +def _forward(line): + global _session_id + token = _get_token() + + headers = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "Authorization": f"Bearer {token}", + } + if _session_id: + headers["Mcp-Session-Id"] = _session_id + + req = urllib.request.Request(APP_URL, data=line.encode(), headers=headers, method="POST") + try: + with urllib.request.urlopen(req, timeout=300) as resp: + sid = resp.headers.get("Mcp-Session-Id") + if sid: + _session_id = sid + body = resp.read().decode() + if body.strip(): + sys.stdout.write(body.rstrip("\n") + "\n") + sys.stdout.flush() + except urllib.error.HTTPError as e: + if e.code in (302, 401, 403): + _log(f"Auth failed ({e.code}), forcing token refresh") + token = _get_token(force=True) + headers["Authorization"] = f"Bearer {token}" + retry = urllib.request.Request(APP_URL, data=line.encode(), headers=headers, method="POST") + with urllib.request.urlopen(retry, timeout=300) as resp: + sid = resp.headers.get("Mcp-Session-Id") + if sid: + _session_id = sid + body = resp.read().decode() + if body.strip(): + sys.stdout.write(body.rstrip("\n") + "\n") + sys.stdout.flush() + else: + raise + + +def main(): + if not APP_URL: + _log("FATAL: CODA_MCP_URL not set") + sys.exit(1) + _log(f"Proxying to {APP_URL} (profile={PROFILE})") + for line in sys.stdin: + line = line.strip() + if not line: + continue + try: + _forward(line) + except Exception as e: + _log(f"Error: {e}") + try: + msg_id = json.loads(line).get("id") + except Exception: + msg_id = None + if msg_id is not None: + err = json.dumps({ + "jsonrpc": "2.0", + "id": msg_id, + "error": {"code": -32000, "message": str(e)}, + }) + sys.stdout.write(err + "\n") + sys.stdout.flush() + + +if __name__ == "__main__": + main()