Merge pull request #2413 from Molecule-AI/fix/external-runtime-universal-mcp

feat(workspace-runtime): expose universal MCP server to runtime=external operators

canvas/src/components/ExternalConnectModal.tsx

@@ -32,6 +32,14 @@ export interface ExternalConnectionInfo {
   // haven't shipped molecule-core PR #2304 yet (older response payload
   // omits the field; tab is hidden if empty).
   claude_code_channel_snippet?: string;
+  // Universal MCP snippet — runtime-agnostic outbound tool path via
+  // the `molecule-mcp` console script in the
+  // molecule-ai-workspace-runtime PyPI wheel. Works with any MCP-aware
+  // agent runtime (Claude Code, hermes, codex, third-party). Outbound-
+  // only: pair with claude_code_channel or python tabs for heartbeat
+  // + inbound. Optional for backward compat with platforms that
+  // haven't shipped PR #2413 yet.
+  universal_mcp_snippet?: string;
 }
 
 interface Props {
@@ -39,7 +47,7 @@ interface Props {
   onClose: () => void;
 }
 
-type Tab = "python" | "curl" | "claude" | "fields";
+type Tab = "python" | "curl" | "claude" | "mcp" | "fields";
 
 export function ExternalConnectModal({ info, onClose }: Props) {
   // Default to Claude Code when the platform offers it — that's the
@@ -89,6 +97,17 @@ export function ExternalConnectModal({ info, onClose }: Props) {
     'MOLECULE_WORKSPACE_TOKENS=<paste auth_token from create response>',
     `MOLECULE_WORKSPACE_TOKENS=${info.auth_token}`,
   );
+  // Universal MCP snippet uses MOLECULE_WORKSPACE_TOKEN as the env-var
+  // name passed through to molecule-mcp via `claude mcp add ... -- env
+  // MOLECULE_WORKSPACE_TOKEN=...`. The placeholder must match the
+  // template's literal — pre-2026-04-30 polish this looked for
+  // WORKSPACE_AUTH_TOKEN (carryover from the curl tab), which silently
+  // skipped the substitution and left "<paste from create response>"
+  // visible in the operator's clipboard.
+  const filledUniversalMcp = info.universal_mcp_snippet?.replace(
+    'MOLECULE_WORKSPACE_TOKEN="<paste from create response>"',
+    `MOLECULE_WORKSPACE_TOKEN="${info.auth_token}"`,
+  );
 
   return (
     <Dialog.Root open onOpenChange={(o) => !o && onClose()}>
@@ -110,11 +129,19 @@ export function ExternalConnectModal({ info, onClose }: Props) {
             aria-label="Connection snippet format"
             className="mt-4 flex gap-1 border-b border-zinc-800"
           >
-            {(
-              filledChannel
-                ? (["claude", "python", "curl", "fields"] as Tab[])
-                : (["python", "curl", "fields"] as Tab[])
-            ).map((t) => (
+            {(() => {
+              // Build the tab order dynamically. Claude Code first
+              // (when offered) since it's the simplest setup; Python
+              // SDK second (full register+heartbeat+inbound); Universal
+              // MCP third (any MCP-aware runtime, outbound-only); curl
+              // for one-shot register; Fields for raw values.
+              const tabs: Tab[] = [];
+              if (filledChannel) tabs.push("claude");
+              tabs.push("python");
+              if (filledUniversalMcp) tabs.push("mcp");
+              tabs.push("curl", "fields");
+              return tabs;
+            })().map((t) => (
               <button
                 key={t}
                 type="button"
@@ -131,6 +158,8 @@ export function ExternalConnectModal({ info, onClose }: Props) {
                   ? "Claude Code"
                   : t === "python"
                   ? "Python SDK"
+                  : t === "mcp"
+                  ? "Universal MCP"
                   : t === "curl"
                   ? "curl"
                   : "Fields"}
@@ -167,6 +196,15 @@ export function ExternalConnectModal({ info, onClose }: Props) {
                 onCopy={() => copy(filledCurl, "curl")}
               />
             )}
+            {tab === "mcp" && filledUniversalMcp && (
+              <SnippetBlock
+                value={filledUniversalMcp}
+                label="Universal MCP — standalone register + heartbeat + tools for any MCP-aware runtime (Claude Code, hermes, codex). Pair with Python or Claude Code tab if you need inbound A2A delivery."
+                copyKey="mcp"
+                copied={copiedKey === "mcp"}
+                onCopy={() => copy(filledUniversalMcp, "mcp")}
+              />
+            )}
             {tab === "fields" && (
               <div className="space-y-2">
                 <Field label="workspace_id" value={info.workspace_id} onCopy={() => copy(info.workspace_id, "wsid")} copied={copiedKey === "wsid"} />

scripts/build_runtime_package.py

@@ -68,6 +68,7 @@ TOP_LEVEL_MODULES = {
     "internal_chat_uploads",
     "internal_file_read",
     "main",
+    "mcp_cli",
     "molecule_ai_status",
     "platform_auth",
     "platform_inbound_auth",
@@ -217,6 +218,7 @@ dependencies = [
 
 [project.scripts]
 molecule-runtime = "molecule_runtime.main:main_sync"
+molecule-mcp = "molecule_runtime.mcp_cli:main"
 
 [tool.setuptools.packages.find]
 where = ["."]
@@ -240,6 +242,31 @@ directory** by the `publish-runtime` GitHub Actions workflow on every
 `runtime-v*` tag push. **Do not edit this package directly** — edit
 `workspace/` in the monorepo.
 
+## External-runtime MCP server (`molecule-mcp`)
+
+Operators running an agent outside the platform's container fleet
+(any runtime that supports MCP stdio — Claude Code, hermes, codex,
+etc.) can install this wheel and run the universal MCP server
+locally:
+
+```sh
+pip install molecule-ai-workspace-runtime
+WORKSPACE_ID=<uuid> \\
+  PLATFORM_URL=https://<tenant>.staging.moleculesai.app \\
+  MOLECULE_WORKSPACE_TOKEN=<bearer> \\
+  molecule-mcp
+```
+
+That exposes the same 8 platform tools (`delegate_task`, `list_peers`,
+`send_message_to_user`, `commit_memory`, etc.) that container-bound
+runtimes already get via the workspace's auto-spawned MCP. Register
+the binary in your agent's MCP config (e.g. Claude Code's
+`claude mcp add molecule -- molecule-mcp` with the env above).
+
+The token comes from the canvas → Tokens tab. Restarting an external
+workspace from the canvas no longer revokes the token (PR #2412), so
+operator tokens persist across status nudges.
+
 See [`docs/workspace-runtime-package.md`](https://github.com/Molecule-AI/molecule-core/blob/main/docs/workspace-runtime-package.md)
 for the publish flow and architecture.
 """

scripts/wheel_smoke.py

@@ -32,6 +32,17 @@ def smoke_imports_and_invariants() -> None:
     from molecule_runtime.builtin_tools import memory  # noqa: F401
     from molecule_runtime.adapters import get_adapter, BaseAdapter, AdapterConfig
 
+    # cli_main + mcp_cli.main are the molecule-mcp console-script entry
+    # points — the external-runtime universal MCP path. Same regression
+    # class as the 0.1.16 main_sync incident: a silent rename or missed
+    # rewrite here would break every external operator's MCP install on
+    # the next wheel publish. Pin both names because pyproject points
+    # at mcp_cli.main, which then imports a2a_mcp_server.cli_main.
+    from molecule_runtime.a2a_mcp_server import cli_main  # noqa: F401
+    from molecule_runtime.mcp_cli import main as mcp_cli_main  # noqa: F401
+    assert callable(cli_main), "a2a_mcp_server.cli_main must be callable"
+    assert callable(mcp_cli_main), "mcp_cli.main must be callable"
+
     assert a2a_client._A2A_ERROR_PREFIX, "a2a_client missing error sentinel"
     assert callable(get_adapter), "adapters.get_adapter must be callable"
     assert hasattr(BaseAdapter, "name"), "BaseAdapter interface broken"

workspace-server/internal/handlers/external_connection.go

@@ -53,8 +53,14 @@ const externalCurlTemplate = `# Replace AGENT_URL with YOUR agent's public HTTPS
 export WORKSPACE_AUTH_TOKEN="<paste from create response>"
 export AGENT_URL="https://your-agent.example.com"
 
+# NOTE on the "Origin" header below: hosted SaaS tenants run behind an
+# edge WAF that requires same-origin requests. Without "Origin", paths
+# like /workspaces/* silently 404 (rewritten to the canvas Next.js).
+# /registry/register is currently allowed without Origin, but setting
+# it preemptively keeps your snippet working if the WAF rules expand.
 curl -fsS -X POST "{{PLATFORM_URL}}/registry/register" \
   -H "Authorization: Bearer $WORKSPACE_AUTH_TOKEN" \
+  -H "Origin: {{PLATFORM_URL}}" \
   -H "Content-Type: application/json" \
   -d '{
     "id": "{{WORKSPACE_ID}}",
@@ -98,6 +104,51 @@ claude --channels plugin:molecule@Molecule-AI/molecule-mcp-claude-channel
 # pairing flow, push-mode upgrade, and v0.2 roadmap.
 `
 
+// externalUniversalMcpTemplate — runtime-agnostic standalone path.
+// Ships as the `molecule-mcp` console script in the
+// molecule-ai-workspace-runtime PyPI wheel (workspace/mcp_cli.py).
+// Any MCP-aware runtime (Claude Code, hermes, codex, third-party)
+// registers it once and gets the same 8 universal tools that
+// container-bound runtimes use today: delegate_task, list_peers,
+// send_message_to_user, commit_memory, etc.
+//
+// Standalone: the binary itself handles register-on-startup +
+// continuous heartbeats (daemon thread, 20s cadence). No separate
+// SDK or channel process needed to keep the workspace online. The
+// only thing it does NOT yet do is poll inbound A2A messages — for
+// runtimes that need their agent to react to canvas messages or
+// peer-initiated tasks, pair with the Claude Code channel tab
+// (poll-based inbound delivery into a Claude Code session) or the
+// Python SDK tab (push-mode inbound + heartbeat).
+//
+// Origin/WAF: handled automatically by platform_auth.auth_headers()
+// in the wheel — operator doesn't need to configure anything.
+const externalUniversalMcpTemplate = `# Universal MCP — standalone register + heartbeat + outbound platform tools
+# for any MCP-aware runtime (Claude Code, hermes, codex, etc.).
+# Pair with the Claude Code or Python SDK tab if your runtime needs
+# inbound A2A delivery (canvas messages → agent conversation turns).
+
+# 1. Install the workspace runtime wheel:
+pip install molecule-ai-workspace-runtime
+
+# 2. Wire molecule-mcp into your agent's MCP config. Claude Code:
+claude mcp add molecule -s user -- env \
+  WORKSPACE_ID={{WORKSPACE_ID}} \
+  PLATFORM_URL={{PLATFORM_URL}} \
+  MOLECULE_WORKSPACE_TOKEN="<paste from create response>" \
+  molecule-mcp
+
+# molecule-mcp registers the workspace + heartbeats every 20s in a
+# daemon thread, then runs the MCP stdio loop. Same env-var contract
+# works with hermes-agent, codex, or any MCP stdio runtime. Tools
+# exposed: delegate_task, delegate_task_async, check_task_status,
+# list_peers, get_workspace_info, send_message_to_user,
+# commit_memory, recall_memory.
+#
+# Origin/WAF handling is built into the wheel — no manual headers
+# needed when calling tools through the MCP server.
+`
+
 // externalPythonTemplate uses molecule-sdk-python's RemoteAgentClient +
 // A2AServer (PR #13 in that repo). Until the SDK cuts a v0.y release
 // to PyPI the snippet pins git+main.

workspace-server/internal/handlers/registry.go

@@ -720,6 +720,34 @@ func (h *RegistryHandler) evaluateStatus(c *gin.Context, payload models.Heartbea
 		})
 	}
 
+	// Auto-recovery from awaiting_agent: external workspaces are flipped
+	// to 'awaiting_agent' by registry/healthsweep when their heartbeat
+	// goes stale (>staleAfter). When the operator's poller comes back —
+	// for example when their laptop wakes from sleep — the heartbeat
+	// resumes but does NOT re-register. Without this branch the
+	// workspace would stay 'awaiting_agent' forever (visible as OFFLINE
+	// in the canvas with a "Restart" CTA) even though the agent is
+	// actively heartbeating.
+	//
+	// Discovered while smoke-testing the universal MCP path against a
+	// freshly-registered external workspace: register set status=online
+	// + sent one heartbeat → healthsweep then flipped back to
+	// awaiting_agent because the smoke didn't loop. The molecule-mcp
+	// console script's built-in heartbeat thread (PR #2413) drives
+	// continuous heartbeats now, but without THIS branch those
+	// heartbeats can't lift the workspace out of awaiting_agent on
+	// their own.
+	if currentStatus == "awaiting_agent" {
+		if _, err := db.DB.ExecContext(ctx, `UPDATE workspaces SET status = $1, updated_at = now() WHERE id = $2 AND status = 'awaiting_agent'`, models.StatusOnline, payload.WorkspaceID); err != nil {
+			log.Printf("Heartbeat: failed to recover %s from awaiting_agent: %v", payload.WorkspaceID, err)
+		} else {
+			log.Printf("Heartbeat: transitioned %s from awaiting_agent to online (heartbeat received)", payload.WorkspaceID)
+		}
+		h.broadcaster.RecordAndBroadcast(ctx, "WORKSPACE_ONLINE", payload.WorkspaceID, map[string]interface{}{
+			"recovered_from": currentStatus,
+		})
+	}
+
 	// #1870 Phase 1: drain one queued A2A request if the target reports
 	// spare capacity. The heartbeat's active_tasks field reflects what the
 	// workspace runtime is ACTUALLY running right now, independent of

workspace-server/internal/handlers/registry_test.go

@@ -193,6 +193,58 @@ func TestHeartbeatHandler_ProvisioningToOnline(t *testing.T) {
 	}
 }
 
+// ==================== Heartbeat — awaiting_agent → online recovery ====================
+// External workspaces flip to 'awaiting_agent' via healthsweep when their
+// heartbeat goes stale. When the operator's poller comes back, heartbeat
+// must lift the workspace out of awaiting_agent the same way it does for
+// 'offline' and 'provisioning'. Without this branch, an external workspace
+// stays OFFLINE in the canvas forever despite active heartbeats.
+
+func TestHeartbeatHandler_AwaitingAgentToOnline(t *testing.T) {
+	mock := setupTestDB(t)
+	setupTestRedis(t)
+	broadcaster := newTestBroadcaster()
+	handler := NewRegistryHandler(broadcaster)
+
+	mock.ExpectQuery("SELECT COALESCE\\(current_task").
+		WithArgs("ws-external").
+		WillReturnRows(sqlmock.NewRows([]string{"current_task"}).AddRow(""))
+
+	mock.ExpectExec("UPDATE workspaces SET").
+		WithArgs("ws-external", 0.0, "", 0, 60, "").
+		WillReturnResult(sqlmock.NewResult(0, 1))
+
+	mock.ExpectQuery("SELECT status FROM workspaces WHERE id =").
+		WithArgs("ws-external").
+		WillReturnRows(sqlmock.NewRows([]string{"status"}).AddRow("awaiting_agent"))
+
+	// The new branch — UPDATE ... WHERE status = 'awaiting_agent'
+	mock.ExpectExec("UPDATE workspaces SET status =").
+		WithArgs(models.StatusOnline, "ws-external").
+		WillReturnResult(sqlmock.NewResult(0, 1))
+
+	// Broadcast WORKSPACE_ONLINE
+	mock.ExpectExec("INSERT INTO structure_events").
+		WillReturnResult(sqlmock.NewResult(0, 1))
+
+	w := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(w)
+
+	body := `{"workspace_id":"ws-external","error_rate":0.0,"sample_error":"","active_tasks":0,"uptime_seconds":60}`
+	c.Request = httptest.NewRequest("POST", "/registry/heartbeat", bytes.NewBufferString(body))
+	c.Request.Header.Set("Content-Type", "application/json")
+
+	handler.Heartbeat(c)
+
+	if w.Code != http.StatusOK {
+		t.Errorf("expected status 200, got %d: %s", w.Code, w.Body.String())
+	}
+
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Errorf("unmet sqlmock expectations: %v", err)
+	}
+}
+
 func TestHeartbeatHandler_BadJSON(t *testing.T) {
 	setupTestDB(t)
 	setupTestRedis(t)

workspace-server/internal/handlers/workspace.go

@@ -412,6 +412,17 @@ func (h *WorkspaceHandler) Create(c *gin.Context) {
 						"{{PLATFORM_URL}}", platformURL),
 					"{{WORKSPACE_ID}}", id,
 				),
+				// Universal MCP snippet — runtime-agnostic outbound
+				// tool path via the molecule-mcp console script. Same
+				// 8 platform tools any MCP-aware runtime can register
+				// (Claude Code, hermes, codex, etc.). Outbound-only:
+				// the snippet calls out that heartbeat/inbound need
+				// pairing with the SDK or channel tab.
+				"universal_mcp_snippet": strings.ReplaceAll(
+					strings.ReplaceAll(externalUniversalMcpTemplate,
+						"{{PLATFORM_URL}}", platformURL),
+					"{{WORKSPACE_ID}}", id,
+				),
 			}
 		}
 		c.JSON(http.StatusCreated, resp)

workspace/a2a_mcp_server.py

@@ -201,5 +201,23 @@ async def main():  # pragma: no cover
             break
 
 
-if __name__ == "__main__":  # pragma: no cover
+def cli_main() -> None:  # pragma: no cover
+    """Synchronous wrapper around the async MCP stdio loop.
+
+    Called by ``mcp_cli.main`` (the ``molecule-mcp`` console-script
+    entry point in scripts/build_runtime_package.py) AFTER env
+    validation and the standalone register + heartbeat thread setup.
+    Direct callers (in-container code that already validated env and
+    runs heartbeat.py separately) can also invoke this — it's the
+    smallest possible "run the MCP stdio JSON-RPC loop" surface.
+
+    Wheel-smoke gates in scripts/wheel_smoke.py pin the importability
+    of this name (alongside ``mcp_cli.main``) so a silent rename can't
+    break every external-runtime operator's MCP install — the 0.1.16
+    ``main_sync`` rename incident is the cautionary precedent.
+    """
     asyncio.run(main())
+
+
+if __name__ == "__main__":  # pragma: no cover
+    cli_main()

workspace/mcp_cli.py

@@ -0,0 +1,302 @@
+"""Console-script entry point for the ``molecule-mcp`` universal MCP server.
+
+Validates required environment BEFORE importing the heavy
+``a2a_mcp_server`` module — that module triggers a ``RuntimeError`` at
+import time when ``WORKSPACE_ID`` is unset (a2a_client.py:22), and
+console-script entry-point shims surface it as an ugly traceback. This
+wrapper catches the missing-env case early and prints actionable help
+to stderr so an operator running ``molecule-mcp`` for the first time
+gets the right pointer in the first 3 lines of output instead of a
+20-line traceback.
+
+Standalone-runtime contract: this wrapper is responsible for keeping
+the workspace ALIVE on the platform side, not just exposing tools.
+Concretely it:
+    1. Calls ``POST /registry/register`` once at startup (idempotent —
+       the upsert flips status awaiting_agent → online for an external
+       workspace whose token matches).
+    2. Spawns a daemon heartbeat thread that POSTs to
+       ``POST /registry/heartbeat`` every 20s. Without continuous
+       heartbeats the platform's healthsweep flips the workspace back
+       to awaiting_agent (visible as OFFLINE in the canvas with a
+       "Restart" CTA) within 60-90s.
+    3. Runs the MCP stdio loop in the foreground.
+
+Why threads + sync requests: the MCP stdio server is async. The
+heartbeat work is fire-and-forget HTTP. A daemon thread is the
+lowest-friction integration — no asyncio bridging, dies automatically
+when the main process exits, and ``requests`` is already a transitive
+dependency via ``a2a-sdk``.
+
+In-container usage (``python -m molecule_runtime.a2a_mcp_server`` or
+direct import) bypasses this wrapper — the workspace runtime has its
+own heartbeat loop in ``heartbeat.py`` so we don't double-heartbeat.
+"""
+from __future__ import annotations
+
+import logging
+import os
+import sys
+import threading
+import time
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+# Heartbeat cadence. Must be tighter than healthsweep's stale window
+# (currently 60-90s — see registry/healthsweep.go) by a comfortable
+# margin so a single missed heartbeat doesn't flip awaiting_agent.
+# 20s gives the operator's network 3 attempts within the budget; long
+# enough that it doesn't spam, short enough to recover quickly after
+# laptop sleep.
+HEARTBEAT_INTERVAL_SECONDS = 20.0
+
+
+def _platform_register(platform_url: str, workspace_id: str, token: str) -> None:
+    """One-shot register at startup; fails fast on auth errors.
+
+    Lifts the workspace from ``awaiting_agent`` to ``online`` for
+    operators who never ran the curl-register snippet. Safe to call
+    repeatedly: the platform's register handler is an upsert that
+    just refreshes ``url``, ``agent_card``, and ``status``.
+
+    Failure model (post-review):
+        - 401 / 403  → ``sys.exit(3)`` immediately. The operator's
+          token is wrong; silently looping in a broken state would
+          make this hard to diagnose because the MCP tools would 401
+          on every call too. Hard-fail is the kindest option.
+        - Other 4xx/5xx → log a warning + continue. The heartbeat
+          thread will surface persistent failures; transient platform
+          blips shouldn't abort the MCP loop.
+        - Network / transport errors → log + continue. Same reasoning.
+
+    Origin header is required by the SaaS edge WAF; without it
+    /registry/register currently still works (it's on the WAF
+    allowlist), but the heartbeat path needs Origin and we want one
+    consistent header set across both calls.
+    """
+    try:
+        import httpx
+    except ImportError:
+        # httpx is a transitive dep via a2a-sdk; if missing, the MCP
+        # server won't import either. Let the caller's later import
+        # surface the real error.
+        return
+
+    payload = {
+        "id": workspace_id,
+        "url": "",
+        "agent_card": {"name": f"molecule-mcp-{workspace_id[:8]}", "skills": []},
+        "delivery_mode": "poll",
+    }
+    headers = {
+        "Authorization": f"Bearer {token}",
+        "Origin": platform_url,
+        "Content-Type": "application/json",
+    }
+    try:
+        with httpx.Client(timeout=10.0) as client:
+            resp = client.post(
+                f"{platform_url}/registry/register",
+                json=payload,
+                headers=headers,
+            )
+        if resp.status_code in (401, 403):
+            print(
+                f"molecule-mcp: register rejected with HTTP {resp.status_code} — "
+                f"the token in MOLECULE_WORKSPACE_TOKEN is invalid for workspace "
+                f"{workspace_id}. Regenerate from the canvas → Tokens tab.",
+                file=sys.stderr,
+            )
+            sys.exit(3)
+        if resp.status_code >= 400:
+            logger.warning(
+                "molecule-mcp: register POST returned HTTP %d: %s",
+                resp.status_code,
+                (resp.text or "")[:200],
+            )
+        else:
+            logger.info(
+                "molecule-mcp: registered workspace %s with platform",
+                workspace_id,
+            )
+    except SystemExit:
+        raise
+    except Exception as exc:  # noqa: BLE001
+        logger.warning("molecule-mcp: register POST failed: %s", exc)
+
+
+def _heartbeat_loop(
+    platform_url: str,
+    workspace_id: str,
+    token: str,
+    interval: float = HEARTBEAT_INTERVAL_SECONDS,
+) -> None:
+    """Daemon thread body: POST /registry/heartbeat every ``interval``s.
+
+    Failures are logged at WARNING and the loop continues. The thread
+    exits when the main process does (daemon=True). Each iteration
+    rebuilds the payload + headers — cheap and ensures token rotation
+    via env var (rare but possible) is picked up on the next tick.
+    """
+    try:
+        import httpx
+    except ImportError:
+        return
+
+    start_time = time.time()
+    while True:
+        body = {
+            "workspace_id": workspace_id,
+            "error_rate": 0.0,
+            "sample_error": "",
+            "active_tasks": 0,
+            "uptime_seconds": int(time.time() - start_time),
+        }
+        headers = {
+            "Authorization": f"Bearer {token}",
+            "Origin": platform_url,
+            "Content-Type": "application/json",
+        }
+        try:
+            with httpx.Client(timeout=10.0) as client:
+                resp = client.post(
+                    f"{platform_url}/registry/heartbeat",
+                    json=body,
+                    headers=headers,
+                )
+            if resp.status_code >= 400:
+                logger.warning(
+                    "molecule-mcp: heartbeat HTTP %d: %s",
+                    resp.status_code,
+                    (resp.text or "")[:200],
+                )
+        except Exception as exc:  # noqa: BLE001
+            logger.warning("molecule-mcp: heartbeat failed: %s", exc)
+        time.sleep(interval)
+
+
+def _start_heartbeat_thread(
+    platform_url: str,
+    workspace_id: str,
+    token: str,
+) -> threading.Thread:
+    """Start the heartbeat daemon thread. Returns the Thread handle.
+
+    The MCP stdio loop runs in the foreground (asyncio); this thread
+    runs alongside it. ``daemon=True`` so when the operator hits
+    Ctrl-C / closes the runtime, the heartbeat dies with it instead
+    of leaking and writing to a stale workspace.
+    """
+    t = threading.Thread(
+        target=_heartbeat_loop,
+        args=(platform_url, workspace_id, token),
+        name="molecule-mcp-heartbeat",
+        daemon=True,
+    )
+    t.start()
+    return t
+
+
+def _print_missing_env_help(missing: list[str], have_token_file: bool) -> None:
+    print("molecule-mcp: missing required environment.\n", file=sys.stderr)
+    print("Set the following before running molecule-mcp:", file=sys.stderr)
+    print("  WORKSPACE_ID                — your workspace UUID (from canvas)", file=sys.stderr)
+    print(
+        "  PLATFORM_URL                — base URL of your Molecule platform "
+        "(e.g. https://your-tenant.staging.moleculesai.app)",
+        file=sys.stderr,
+    )
+    if not have_token_file:
+        print(
+            "  MOLECULE_WORKSPACE_TOKEN    — bearer token for this workspace "
+            "(canvas → Tokens tab)",
+            file=sys.stderr,
+        )
+    print("", file=sys.stderr)
+    print(f"Currently missing: {', '.join(missing)}", file=sys.stderr)
+
+
+def main() -> None:
+    """Entry point for the ``molecule-mcp`` console script.
+
+    Returns nothing — calls ``sys.exit`` on validation failure or on
+    normal completion of the underlying MCP server loop.
+    """
+    missing: list[str] = []
+    if not os.environ.get("WORKSPACE_ID", "").strip():
+        missing.append("WORKSPACE_ID")
+    if not os.environ.get("PLATFORM_URL", "").strip():
+        missing.append("PLATFORM_URL")
+    # Token can come from env OR file — only flag when both are absent.
+    # Mirrors platform_auth.get_token's resolution order (file-first,
+    # env-fallback).
+    configs_dir = Path(os.environ.get("CONFIGS_DIR", "/configs"))
+    has_token_file = (configs_dir / ".auth_token").is_file()
+    has_token_env = bool(os.environ.get("MOLECULE_WORKSPACE_TOKEN", "").strip())
+    if not has_token_file and not has_token_env:
+        missing.append("MOLECULE_WORKSPACE_TOKEN (or CONFIGS_DIR/.auth_token)")
+
+    if missing:
+        _print_missing_env_help(missing, have_token_file=has_token_file)
+        sys.exit(2)
+
+    # Resolve the effective token: env wins (operator override), then
+    # the on-disk file (in-container default). Mirrors
+    # platform_auth.get_token's resolution order so we don't
+    # double-implement.
+    token = (
+        os.environ.get("MOLECULE_WORKSPACE_TOKEN", "").strip()
+        or _read_token_file()
+    )
+    workspace_id = os.environ["WORKSPACE_ID"].strip()
+    platform_url = os.environ["PLATFORM_URL"].strip().rstrip("/")
+
+    # Configure logging so the operator sees register/heartbeat status
+    # without needing to set up logging themselves. WARNING by default
+    # keeps the steady-state quiet (only failures); MOLECULE_MCP_VERBOSE=1
+    # surfaces register-success + per-tick heartbeat info for debugging.
+    log_level = (
+        logging.INFO
+        if os.environ.get("MOLECULE_MCP_VERBOSE", "").strip()
+        else logging.WARNING
+    )
+    logging.basicConfig(level=log_level, format="[molecule-mcp] %(message)s")
+
+    # Standalone-mode register + heartbeat. Skipped via env var so an
+    # in-container caller (which has its own heartbeat loop) can reuse
+    # this entry point without double-heartbeating. The wheel's main
+    # console-script path always runs them; the
+    # MOLECULE_MCP_DISABLE_HEARTBEAT escape hatch exists for tests +
+    # the rare embedded use-case.
+    if not os.environ.get("MOLECULE_MCP_DISABLE_HEARTBEAT", "").strip():
+        _platform_register(platform_url, workspace_id, token)
+        _start_heartbeat_thread(platform_url, workspace_id, token)
+
+    # Env is valid — safe to import the heavy module now. Importing
+    # earlier would trigger a2a_client.py:22's module-level RuntimeError
+    # before our friendly help reaches the user.
+    from a2a_mcp_server import cli_main
+    cli_main()
+
+
+def _read_token_file() -> str:
+    """Read the token from ${CONFIGS_DIR}/.auth_token if present.
+
+    Mirrors platform_auth._token_file but without importing the heavy
+    module here (that import triggers a2a_client's WORKSPACE_ID guard
+    which is fine after env validation, but cheaper to inline a 4-line
+    file read than pull in the whole stack just for the path).
+    """
+    configs_dir = Path(os.environ.get("CONFIGS_DIR", "/configs"))
+    path = configs_dir / ".auth_token"
+    if not path.is_file():
+        return ""
+    try:
+        return path.read_text().strip()
+    except OSError:
+        return ""
+
+
+if __name__ == "__main__":  # pragma: no cover
+    main()

workspace/platform_auth.py

@@ -39,22 +39,42 @@ def _token_file() -> Path:
 
 
 def get_token() -> str | None:
-    """Return the cached token, reading it from disk on first call."""
+    """Return the cached token, reading it from disk on first call.
+
+    Resolution order:
+        1. In-process cache (hot path)
+        2. ``${CONFIGS_DIR}/.auth_token`` file (in-container default —
+           the platform writes this on provision and rotates it on
+           restart)
+        3. ``MOLECULE_WORKSPACE_TOKEN`` env var (external-runtime path —
+           operators running the universal MCP server outside a
+           container have no /configs volume to populate, so they pass
+           the token via env)
+
+    File-first preserves in-container behavior unchanged: containers
+    always have /configs/.auth_token on disk, env-var fallback only
+    fires when there's no file. This is additive — no existing caller
+    sees a behavior change.
+    """
     global _cached_token
     if _cached_token is not None:
         return _cached_token
     path = _token_file()
-    if not path.exists():
-        return None
-    try:
-        tok = path.read_text().strip()
-    except OSError as exc:
-        logger.warning("platform_auth: failed to read %s: %s", path, exc)
-        return None
-    if not tok:
-        return None
-    _cached_token = tok
-    return tok
+    if path.exists():
+        try:
+            tok = path.read_text().strip()
+        except OSError as exc:
+            logger.warning("platform_auth: failed to read %s: %s", path, exc)
+            tok = ""
+        if tok:
+            _cached_token = tok
+            return tok
+    # File missing or empty — fall back to env (external-runtime path).
+    env_tok = os.environ.get("MOLECULE_WORKSPACE_TOKEN", "").strip()
+    if env_tok:
+        _cached_token = env_tok
+        return env_tok
+    return None
 
 
 def save_token(token: str) -> None:
@@ -91,11 +111,26 @@ def auth_headers() -> dict[str, str]:
     """Return a header dict to merge into httpx calls. Empty if no token
     is available yet — callers send the request as-is and the platform's
     heartbeat handler grandfathers pre-token workspaces through until
-    their next /registry/register issues one."""
+    their next /registry/register issues one.
+
+    Always sets ``Origin`` to ``PLATFORM_URL`` when that env var is set.
+    On hosted SaaS deployments the tenant's edge WAF requires a same-
+    origin header — without it ``/workspaces/*`` and ``/registry/*/peers``
+    requests get silently rewritten to the canvas Next.js app, which has
+    no such routes and returns an empty 404. Inside-container calls are
+    unaffected (Docker-internal PLATFORM_URLs aren't behind the WAF).
+    Discovered while smoke-testing the molecule-mcp external-runtime
+    path against a live tenant — every tool call returned "not found"
+    because the WAF was eating them.
+    """
+    headers: dict[str, str] = {}
+    platform_url = os.environ.get("PLATFORM_URL", "").strip()
+    if platform_url:
+        headers["Origin"] = platform_url
     tok = get_token()
-    if not tok:
-        return {}
-    return {"Authorization": f"Bearer {tok}"}
+    if tok:
+        headers["Authorization"] = f"Bearer {tok}"
+    return headers
 
 
 def self_source_headers(workspace_id: str) -> dict[str, str]:

workspace/tests/test_mcp_cli.py

@@ -0,0 +1,492 @@
+"""Tests for workspace/mcp_cli.py — the molecule-mcp console-script
+entry-point validator.
+
+The wrapper exists to surface a friendly missing-env error before
+a2a_client.py:22's module-level RuntimeError fires. Regressions here
+ship a poor first-run UX to every external-runtime operator.
+"""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import pytest
+
+import mcp_cli
+
+
+@pytest.fixture(autouse=True)
+def _isolate(monkeypatch, tmp_path):
+    """Each test starts with no Molecule env vars set + a fresh
+    CONFIGS_DIR pointing at an empty tmpdir. The heartbeat thread is
+    disabled by default so happy-path tests don't spawn a background
+    POST loop against a fake URL — individual tests opt back in via
+    monkeypatch.delenv when they want to assert heartbeat behavior."""
+    for var in ("WORKSPACE_ID", "PLATFORM_URL", "MOLECULE_WORKSPACE_TOKEN"):
+        monkeypatch.delenv(var, raising=False)
+    monkeypatch.setenv("CONFIGS_DIR", str(tmp_path))
+    monkeypatch.setenv("MOLECULE_MCP_DISABLE_HEARTBEAT", "1")
+    yield
+
+
+def _run_main_capturing_exit(capsys) -> tuple[int, str]:
+    """Call mcp_cli.main and return (exit_code, stderr).
+
+    main() is supposed to sys.exit on missing env. Any non-exit return
+    means it tried to run the real MCP loop, which we don't want in a
+    unit test (and which would also fail because we never set the
+    mandatory env).
+    """
+    with pytest.raises(SystemExit) as exc_info:
+        mcp_cli.main()
+    captured = capsys.readouterr()
+    code = exc_info.value.code if isinstance(exc_info.value.code, int) else 1
+    return code, captured.err
+
+
+def test_missing_workspace_id_exits_with_message(capsys):
+    code, err = _run_main_capturing_exit(capsys)
+    assert code == 2, f"expected exit code 2, got {code}"
+    assert "WORKSPACE_ID" in err
+    assert "PLATFORM_URL" in err  # also missing
+    assert "MOLECULE_WORKSPACE_TOKEN" in err  # also missing
+
+
+def test_only_workspace_id_missing(capsys, monkeypatch):
+    monkeypatch.setenv("PLATFORM_URL", "http://localhost:8080")
+    monkeypatch.setenv("MOLECULE_WORKSPACE_TOKEN", "tok")
+    code, err = _run_main_capturing_exit(capsys)
+    assert code == 2
+    # Only WORKSPACE_ID should appear in the "currently missing" list.
+    assert "Currently missing: WORKSPACE_ID" in err
+
+
+def test_only_platform_url_missing(capsys, monkeypatch):
+    monkeypatch.setenv("WORKSPACE_ID", "00000000-0000-0000-0000-000000000000")
+    monkeypatch.setenv("MOLECULE_WORKSPACE_TOKEN", "tok")
+    code, err = _run_main_capturing_exit(capsys)
+    assert code == 2
+    assert "Currently missing: PLATFORM_URL" in err
+
+
+def test_only_token_missing(capsys, monkeypatch):
+    monkeypatch.setenv("WORKSPACE_ID", "00000000-0000-0000-0000-000000000000")
+    monkeypatch.setenv("PLATFORM_URL", "http://localhost:8080")
+    code, err = _run_main_capturing_exit(capsys)
+    assert code == 2
+    assert "MOLECULE_WORKSPACE_TOKEN" in err
+
+
+def test_token_file_satisfies_token_requirement(capsys, monkeypatch, tmp_path):
+    """Token from CONFIGS_DIR/.auth_token must be accepted (in-container
+    path)."""
+    (tmp_path / ".auth_token").write_text("file-token")
+    monkeypatch.setenv("WORKSPACE_ID", "00000000-0000-0000-0000-000000000000")
+    monkeypatch.setenv("PLATFORM_URL", "http://localhost:8080")
+    # No MOLECULE_WORKSPACE_TOKEN — but file exists. Validation should
+    # pass; we then short-circuit before importing the heavy module by
+    # patching the import to a no-op spy.
+
+    spy_called: dict[str, bool] = {"called": False}
+
+    def fake_cli_main():
+        spy_called["called"] = True
+
+    # Patch the heavy import to avoid actually running the MCP server.
+    # mcp_cli does the import lazily inside main(), so we monkeypatch
+    # sys.modules to inject a fake a2a_mcp_server.
+    import types
+    fake_module = types.ModuleType("a2a_mcp_server")
+    fake_module.cli_main = fake_cli_main
+    monkeypatch.setitem(sys.modules, "a2a_mcp_server", fake_module)
+
+    mcp_cli.main()  # should NOT exit
+    assert spy_called["called"], "expected cli_main to be invoked when env+file are valid"
+
+
+def test_env_token_satisfies_token_requirement(capsys, monkeypatch):
+    """Token from env must be accepted (external-runtime path)."""
+    monkeypatch.setenv("WORKSPACE_ID", "00000000-0000-0000-0000-000000000000")
+    monkeypatch.setenv("PLATFORM_URL", "http://localhost:8080")
+    monkeypatch.setenv("MOLECULE_WORKSPACE_TOKEN", "env-token")
+
+    spy_called: dict[str, bool] = {"called": False}
+
+    def fake_cli_main():
+        spy_called["called"] = True
+
+    import types
+    fake_module = types.ModuleType("a2a_mcp_server")
+    fake_module.cli_main = fake_cli_main
+    monkeypatch.setitem(sys.modules, "a2a_mcp_server", fake_module)
+
+    mcp_cli.main()
+    assert spy_called["called"]
+
+
+def test_whitespace_only_env_treated_as_missing(capsys, monkeypatch):
+    """An accidentally-empty env var (WORKSPACE_ID="   ") must NOT be
+    considered set — otherwise the error would surface deep inside an
+    HTTP call instead of in this validator."""
+    monkeypatch.setenv("WORKSPACE_ID", "   ")
+    monkeypatch.setenv("PLATFORM_URL", "http://localhost:8080")
+    monkeypatch.setenv("MOLECULE_WORKSPACE_TOKEN", "tok")
+    code, err = _run_main_capturing_exit(capsys)
+    assert code == 2
+    assert "WORKSPACE_ID" in err
+
+
+def test_help_lists_canvas_tokens_tab_pointer(capsys):
+    """Operator must know WHERE to get a token. The help mentions the
+    canvas Tokens tab so they can self-recover without asking on
+    Slack."""
+    code, err = _run_main_capturing_exit(capsys)
+    assert code == 2
+    assert "Tokens tab" in err or "canvas" in err.lower()
+
+
+# ==================== Standalone register + heartbeat ====================
+# molecule-mcp must be a single-process standalone runtime: it registers
+# the workspace at startup AND continuously heartbeats so the platform
+# healthsweep doesn't flip status back to awaiting_agent. Without these,
+# the operator sees "OFFLINE — Restart" in the canvas within ~60s of
+# launching the agent, which was the bug that motivated this PR.
+
+
+def test_register_called_at_startup(monkeypatch):
+    """When env is valid and heartbeat enabled, register fires once
+    before the MCP loop starts."""
+    monkeypatch.setenv("WORKSPACE_ID", "00000000-0000-0000-0000-000000000000")
+    monkeypatch.setenv("PLATFORM_URL", "https://test.moleculesai.app")
+    monkeypatch.setenv("MOLECULE_WORKSPACE_TOKEN", "tok")
+    monkeypatch.delenv("MOLECULE_MCP_DISABLE_HEARTBEAT", raising=False)
+
+    register_calls: list[tuple[str, str, str]] = []
+
+    def fake_register(platform_url, workspace_id, token):
+        register_calls.append((platform_url, workspace_id, token))
+
+    def fake_start_thread(*_args, **_kwargs):
+        # Return a dummy thread-shaped object so the caller's reference
+        # is harmless. Real thread spawning is asserted separately.
+        class _Stub:
+            def join(self): pass
+        return _Stub()
+
+    monkeypatch.setattr(mcp_cli, "_platform_register", fake_register)
+    monkeypatch.setattr(mcp_cli, "_start_heartbeat_thread", fake_start_thread)
+
+    spy_called: dict[str, bool] = {"called": False}
+
+    def fake_cli_main():
+        spy_called["called"] = True
+
+    import types
+    fake_module = types.ModuleType("a2a_mcp_server")
+    fake_module.cli_main = fake_cli_main
+    monkeypatch.setitem(sys.modules, "a2a_mcp_server", fake_module)
+
+    mcp_cli.main()
+
+    assert register_calls == [
+        ("https://test.moleculesai.app", "00000000-0000-0000-0000-000000000000", "tok"),
+    ]
+    assert spy_called["called"], "MCP loop must run AFTER register"
+
+
+def test_heartbeat_thread_started(monkeypatch):
+    """The heartbeat daemon thread must start before the MCP loop runs."""
+    monkeypatch.setenv("WORKSPACE_ID", "00000000-0000-0000-0000-000000000000")
+    monkeypatch.setenv("PLATFORM_URL", "https://test.moleculesai.app")
+    monkeypatch.setenv("MOLECULE_WORKSPACE_TOKEN", "tok")
+    monkeypatch.delenv("MOLECULE_MCP_DISABLE_HEARTBEAT", raising=False)
+
+    monkeypatch.setattr(mcp_cli, "_platform_register", lambda *a, **k: None)
+
+    thread_started: dict[str, bool] = {"started": False}
+
+    def fake_start_thread(platform_url, workspace_id, token):
+        thread_started["started"] = True
+        thread_started["args"] = (platform_url, workspace_id, token)
+        class _Stub:
+            def join(self): pass
+        return _Stub()
+
+    monkeypatch.setattr(mcp_cli, "_start_heartbeat_thread", fake_start_thread)
+
+    import types
+    fake_module = types.ModuleType("a2a_mcp_server")
+    fake_module.cli_main = lambda: None
+    monkeypatch.setitem(sys.modules, "a2a_mcp_server", fake_module)
+
+    mcp_cli.main()
+
+    assert thread_started["started"], "heartbeat thread must be spawned"
+    assert thread_started["args"][1] == "00000000-0000-0000-0000-000000000000"
+    assert thread_started["args"][2] == "tok"
+
+
+def test_heartbeat_disable_env_skips_both(monkeypatch):
+    """MOLECULE_MCP_DISABLE_HEARTBEAT=1 (the test fixture default + the
+    in-container escape hatch) must skip BOTH register and heartbeat,
+    so the in-container heartbeat loop in heartbeat.py doesn't compete
+    with this thread."""
+    monkeypatch.setenv("WORKSPACE_ID", "00000000-0000-0000-0000-000000000000")
+    monkeypatch.setenv("PLATFORM_URL", "https://test.moleculesai.app")
+    monkeypatch.setenv("MOLECULE_WORKSPACE_TOKEN", "tok")
+    # MOLECULE_MCP_DISABLE_HEARTBEAT=1 is set by the autouse fixture.
+
+    register_called: dict[str, bool] = {"called": False}
+    thread_started: dict[str, bool] = {"started": False}
+
+    monkeypatch.setattr(
+        mcp_cli, "_platform_register",
+        lambda *a, **k: register_called.update(called=True),
+    )
+    monkeypatch.setattr(
+        mcp_cli, "_start_heartbeat_thread",
+        lambda *a, **k: thread_started.update(started=True),
+    )
+
+    import types
+    fake_module = types.ModuleType("a2a_mcp_server")
+    fake_module.cli_main = lambda: None
+    monkeypatch.setitem(sys.modules, "a2a_mcp_server", fake_module)
+
+    mcp_cli.main()
+
+    assert register_called["called"] is False, "disable env must skip register"
+    assert thread_started["started"] is False, "disable env must skip heartbeat thread"
+
+
+def test_token_resolved_from_env_when_no_file(monkeypatch):
+    """Operator without a /configs volume — token comes from env var."""
+    monkeypatch.setenv("WORKSPACE_ID", "00000000-0000-0000-0000-000000000000")
+    monkeypatch.setenv("PLATFORM_URL", "https://test.moleculesai.app")
+    monkeypatch.setenv("MOLECULE_WORKSPACE_TOKEN", "env-token")
+    monkeypatch.delenv("MOLECULE_MCP_DISABLE_HEARTBEAT", raising=False)
+
+    captured_token: dict[str, str] = {}
+
+    def fake_register(platform_url, workspace_id, token):
+        captured_token["t"] = token
+
+    monkeypatch.setattr(mcp_cli, "_platform_register", fake_register)
+    monkeypatch.setattr(mcp_cli, "_start_heartbeat_thread", lambda *a, **k: None)
+
+    import types
+    fake_module = types.ModuleType("a2a_mcp_server")
+    fake_module.cli_main = lambda: None
+    monkeypatch.setitem(sys.modules, "a2a_mcp_server", fake_module)
+
+    mcp_cli.main()
+
+    assert captured_token["t"] == "env-token"
+
+
+def test_token_resolved_from_file_when_no_env(monkeypatch, tmp_path):
+    """In-container parity: token comes from /configs/.auth_token when
+    env is unset. Mirrors platform_auth.get_token resolution order."""
+    (tmp_path / ".auth_token").write_text("file-token")
+    monkeypatch.setenv("WORKSPACE_ID", "00000000-0000-0000-0000-000000000000")
+    monkeypatch.setenv("PLATFORM_URL", "https://test.moleculesai.app")
+    monkeypatch.delenv("MOLECULE_WORKSPACE_TOKEN", raising=False)
+    monkeypatch.delenv("MOLECULE_MCP_DISABLE_HEARTBEAT", raising=False)
+
+    captured_token: dict[str, str] = {}
+
+    def fake_register(platform_url, workspace_id, token):
+        captured_token["t"] = token
+
+    monkeypatch.setattr(mcp_cli, "_platform_register", fake_register)
+    monkeypatch.setattr(mcp_cli, "_start_heartbeat_thread", lambda *a, **k: None)
+
+    import types
+    fake_module = types.ModuleType("a2a_mcp_server")
+    fake_module.cli_main = lambda: None
+    monkeypatch.setitem(sys.modules, "a2a_mcp_server", fake_module)
+
+    mcp_cli.main()
+
+    assert captured_token["t"] == "file-token"
+
+
+def test_register_401_exits_with_actionable_error(monkeypatch, capsys):
+    """Bad token at startup must hard-fail. Otherwise the operator
+    sees no error in their MCP client (which spawns the binary in a
+    subprocess), the heartbeat thread silently 401's forever, and
+    every tool call also 401's — needle-in-haystack debugging.
+    Hard-exiting prints a clear pointer to the canvas Tokens tab."""
+
+    class FakeResp:
+        status_code = 401
+        text = "invalid workspace auth token"
+
+    class FakeClient:
+        def __init__(self, **_kwargs): pass
+        def __enter__(self): return self
+        def __exit__(self, *_a): return False
+        def post(self, *_a, **_kw): return FakeResp()
+
+    import types
+    fake_httpx = types.ModuleType("httpx")
+    fake_httpx.Client = FakeClient
+    monkeypatch.setitem(sys.modules, "httpx", fake_httpx)
+
+    with pytest.raises(SystemExit) as exc_info:
+        mcp_cli._platform_register(
+            "https://test.moleculesai.app",
+            "ws-bad-token",
+            "wrong-token",
+        )
+    assert exc_info.value.code == 3
+    err = capsys.readouterr().err
+    assert "401" in err
+    assert "ws-bad-token" in err
+    assert "Tokens tab" in err or "canvas" in err.lower()
+
+
+def test_register_403_also_exits(monkeypatch, capsys):
+    """403 is the C18 hijack-prevention rejection — same operator
+    action (regenerate token) as 401."""
+
+    class FakeResp:
+        status_code = 403
+        text = "C18: live tokens exist; bearer didn't match"
+
+    class FakeClient:
+        def __init__(self, **_kwargs): pass
+        def __enter__(self): return self
+        def __exit__(self, *_a): return False
+        def post(self, *_a, **_kw): return FakeResp()
+
+    import types
+    fake_httpx = types.ModuleType("httpx")
+    fake_httpx.Client = FakeClient
+    monkeypatch.setitem(sys.modules, "httpx", fake_httpx)
+
+    with pytest.raises(SystemExit) as exc_info:
+        mcp_cli._platform_register(
+            "https://test.moleculesai.app",
+            "ws-hijack",
+            "stolen-token",
+        )
+    assert exc_info.value.code == 3
+
+
+def test_register_500_does_not_exit(monkeypatch):
+    """Transient platform errors (500, 503) must NOT hard-fail —
+    those clear on retry and the heartbeat thread will surface
+    persistent failures via warning logs."""
+
+    class FakeResp:
+        status_code = 503
+        text = "service unavailable"
+
+    class FakeClient:
+        def __init__(self, **_kwargs): pass
+        def __enter__(self): return self
+        def __exit__(self, *_a): return False
+        def post(self, *_a, **_kw): return FakeResp()
+
+    import types
+    fake_httpx = types.ModuleType("httpx")
+    fake_httpx.Client = FakeClient
+    monkeypatch.setitem(sys.modules, "httpx", fake_httpx)
+
+    # Should return cleanly, no SystemExit raised
+    mcp_cli._platform_register(
+        "https://test.moleculesai.app",
+        "ws-ok",
+        "tok",
+    )
+
+
+def test_register_payload_shape(monkeypatch):
+    """The register POST body must use the field names the workspace-
+    server expects (id/url/agent_card/delivery_mode), and must include
+    the Origin header for the SaaS edge WAF."""
+    captured: dict[str, object] = {}
+
+    class FakeResp:
+        status_code = 200
+        text = ""
+
+    class FakeClient:
+        def __init__(self, **_kwargs): pass
+        def __enter__(self): return self
+        def __exit__(self, *_a): return False
+        def post(self, url, json=None, headers=None):
+            captured["url"] = url
+            captured["json"] = json
+            captured["headers"] = headers
+            return FakeResp()
+
+    import types
+    fake_httpx = types.ModuleType("httpx")
+    fake_httpx.Client = FakeClient
+    monkeypatch.setitem(sys.modules, "httpx", fake_httpx)
+
+    mcp_cli._platform_register(
+        "https://test.moleculesai.app",
+        "ws-abc",
+        "tok",
+    )
+
+    assert captured["url"] == "https://test.moleculesai.app/registry/register"
+    body = captured["json"]
+    assert body["id"] == "ws-abc"
+    assert body["delivery_mode"] == "poll"
+    assert body["url"] == ""
+    assert "agent_card" in body
+    headers = captured["headers"]
+    assert headers["Authorization"] == "Bearer tok"
+    assert headers["Origin"] == "https://test.moleculesai.app"
+
+
+def test_heartbeat_loop_posts_to_correct_endpoint(monkeypatch):
+    """Heartbeat thread must POST to /registry/heartbeat with the
+    workspace_id + Origin/Authorization headers."""
+    captured: dict[str, object] = {}
+
+    class FakeResp:
+        status_code = 200
+        text = ""
+
+    class FakeClient:
+        def __init__(self, **_kwargs): pass
+        def __enter__(self): return self
+        def __exit__(self, *_a): return False
+        def post(self, url, json=None, headers=None):
+            captured["url"] = url
+            captured["json"] = json
+            captured["headers"] = headers
+            return FakeResp()
+
+    import types
+    fake_httpx = types.ModuleType("httpx")
+    fake_httpx.Client = FakeClient
+    monkeypatch.setitem(sys.modules, "httpx", fake_httpx)
+
+    # Patch sleep so the loop exits after one tick (raise to break out).
+    sleep_calls: list[float] = []
+
+    def fake_sleep(seconds):
+        sleep_calls.append(seconds)
+        raise SystemExit  # break out of the infinite loop
+
+    monkeypatch.setattr("time.sleep", fake_sleep)
+
+    with pytest.raises(SystemExit):
+        mcp_cli._heartbeat_loop(
+            "https://test.moleculesai.app",
+            "ws-abc",
+            "tok",
+            interval=20.0,
+        )
+
+    assert captured["url"] == "https://test.moleculesai.app/registry/heartbeat"
+    assert captured["json"]["workspace_id"] == "ws-abc"
+    assert captured["headers"]["Authorization"] == "Bearer tok"
+    assert captured["headers"]["Origin"] == "https://test.moleculesai.app"
+    assert sleep_calls == [20.0], "heartbeat must sleep the configured interval"

workspace/tests/test_platform_auth.py

@@ -65,15 +65,36 @@ def test_save_token_rotation_overwrites(tmp_path):
     assert platform_auth.get_token() == "token-v2"
 
 
-def test_auth_headers_when_no_token_is_empty():
+def test_auth_headers_when_no_token_and_no_platform_is_empty(monkeypatch):
+    monkeypatch.delenv("PLATFORM_URL", raising=False)
     assert platform_auth.auth_headers() == {}
 
 
-def test_auth_headers_format():
+def test_auth_headers_when_no_token_includes_origin(monkeypatch):
+    """Origin must be set even without a token — the WAF gates ALL
+    requests to /workspaces and /registry, including pre-token bootstrap
+    register calls. Without Origin those would silently 404 from Next.js."""
+    monkeypatch.setenv("PLATFORM_URL", "https://tenant.moleculesai.app")
+    assert platform_auth.auth_headers() == {"Origin": "https://tenant.moleculesai.app"}
+
+
+def test_auth_headers_format(monkeypatch):
+    monkeypatch.delenv("PLATFORM_URL", raising=False)
     platform_auth.save_token("hello-world")
     assert platform_auth.auth_headers() == {"Authorization": "Bearer hello-world"}
 
 
+def test_auth_headers_includes_origin_when_platform_url_set(monkeypatch):
+    """Both Authorization and Origin land on the same dict so the
+    SaaS edge WAF accepts every workspace-runtime request."""
+    monkeypatch.setenv("PLATFORM_URL", "https://hongmingwang.moleculesai.app")
+    platform_auth.save_token("tok")
+    assert platform_auth.auth_headers() == {
+        "Authorization": "Bearer tok",
+        "Origin": "https://hongmingwang.moleculesai.app",
+    }
+
+
 def test_get_token_caches_after_first_disk_read(tmp_path, monkeypatch):
     path = tmp_path / ".auth_token"
     path.write_text("disk-token")
@@ -119,3 +140,66 @@ def test_default_configs_dir_fallback(tmp_path, monkeypatch):
     # We expect _token_file() to resolve under /configs when env is unset.
     path = platform_auth._token_file()
     assert str(path).startswith("/configs")
+
+
+# ==================== MOLECULE_WORKSPACE_TOKEN env-var fallback ====================
+# External-runtime path: operators running the universal MCP server outside
+# a container have no /configs volume. They pass the token via env. The
+# fallback must NOT override the file when both are present (in-container
+# rotation must keep working) and MUST surface env when the file is absent.
+
+
+def test_get_token_uses_env_when_file_absent(tmp_path, monkeypatch):
+    monkeypatch.setenv("MOLECULE_WORKSPACE_TOKEN", "env-token-xyz")
+    assert not (tmp_path / ".auth_token").exists()
+    assert platform_auth.get_token() == "env-token-xyz"
+
+
+def test_get_token_file_takes_priority_over_env(tmp_path, monkeypatch):
+    """In-container rotation must keep working — file overrides env."""
+    (tmp_path / ".auth_token").write_text("file-token")
+    monkeypatch.setenv("MOLECULE_WORKSPACE_TOKEN", "env-token-should-be-ignored")
+    assert platform_auth.get_token() == "file-token"
+
+
+def test_get_token_falls_back_to_env_when_file_empty(tmp_path, monkeypatch):
+    """Empty file is equivalent to absent — env still fires."""
+    (tmp_path / ".auth_token").write_text("")
+    monkeypatch.setenv("MOLECULE_WORKSPACE_TOKEN", "env-token-fallback")
+    assert platform_auth.get_token() == "env-token-fallback"
+
+
+def test_get_token_strips_env_whitespace(tmp_path, monkeypatch):
+    monkeypatch.setenv("MOLECULE_WORKSPACE_TOKEN", "  padded-env-token  \n")
+    assert platform_auth.get_token() == "padded-env-token"
+
+
+def test_get_token_ignores_empty_env(tmp_path, monkeypatch):
+    """Empty string env var is the same as unset — no false positive."""
+    monkeypatch.setenv("MOLECULE_WORKSPACE_TOKEN", "")
+    assert platform_auth.get_token() is None
+
+
+def test_get_token_ignores_whitespace_only_env(tmp_path, monkeypatch):
+    monkeypatch.setenv("MOLECULE_WORKSPACE_TOKEN", "   \n\n   ")
+    assert platform_auth.get_token() is None
+
+
+def test_env_token_caches_like_file_token(tmp_path, monkeypatch):
+    """Once env-token is read, mutating env shouldn't affect cached value."""
+    monkeypatch.setenv("MOLECULE_WORKSPACE_TOKEN", "first-env-token")
+    assert platform_auth.get_token() == "first-env-token"
+    monkeypatch.setenv("MOLECULE_WORKSPACE_TOKEN", "second-env-token")
+    # Cache returns first value
+    assert platform_auth.get_token() == "first-env-token"
+    # clear_cache forces re-read of env
+    platform_auth.clear_cache()
+    assert platform_auth.get_token() == "second-env-token"
+
+
+def test_auth_headers_works_with_env_token(tmp_path, monkeypatch):
+    """Header construction must use the env-fallback token, not silently
+    return {} when no file exists."""
+    monkeypatch.delenv("PLATFORM_URL", raising=False)
+    monkeypatch.setenv("MOLECULE_WORKSPACE_TOKEN", "external-bearer")
+    assert platform_auth.auth_headers() == {"Authorization": "Bearer external-bearer"}