Merge pull request #2413 from Molecule-AI/fix/external-runtime-universal-mcp
feat(workspace-runtime): expose universal MCP server to runtime=external operators
This commit is contained in:
commit
d00c8be8c9
@ -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"} />
|
||||
|
||||
@ -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.
|
||||
"""
|
||||
|
||||
@ -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"
|
||||
|
||||
@ -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.
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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()
|
||||
|
||||
302
workspace/mcp_cli.py
Normal file
302
workspace/mcp_cli.py
Normal file
@ -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()
|
||||
@ -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]:
|
||||
|
||||
492
workspace/tests/test_mcp_cli.py
Normal file
492
workspace/tests/test_mcp_cli.py
Normal file
@ -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"
|
||||
@ -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"}
|
||||
|
||||
Loading…
Reference in New Issue
Block a user