From a3bba8a3f333b362aa32ff6e0f2fa59ce1140379 Mon Sep 17 00:00:00 2001 From: Molecule AI SDK-Dev Date: Sun, 10 May 2026 13:33:17 +0000 Subject: [PATCH] feat(molecule_agent): add org_id and origin kwargs to RemoteAgentClient Adds optional org_id and origin constructor kwargs that inject X-Molecule-Org-Id and Origin headers on every request, enabling SDK use against multi-tenant SaaS deployments (*.moleculesai.app) without needing a pre-configured requests.Session. Co-Authored-By: Claude Opus 4.7 --- molecule_agent/README.md | 22 +++++------- molecule_agent/client.py | 22 ++++++++++-- tests/test_remote_agent.py | 72 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 99 insertions(+), 17 deletions(-) diff --git a/molecule_agent/README.md b/molecule_agent/README.md index 01b6233..bff6f93 100644 --- a/molecule_agent/README.md +++ b/molecule_agent/README.md @@ -219,30 +219,24 @@ reference it directly. - `Origin: ` — `/workspaces/*` and `/registry/*/peers` silently rewrite to Next.js without it (returns an empty 404, easy to misdiagnose as auth) - `RemoteAgentClient` does not set either header today — it ships only - `Authorization: Bearer ` and per-call `X-Workspace-ID` / - `X-Source-Workspace-Id`. Workaround: pass a pre-configured - `requests.Session` to the constructor with the headers set globally: + A follow-up PR will accept `org_id` and `origin` constructor kwargs and + inject the headers automatically. + +- **Tenant + Origin headers (resolved).** + `RemoteAgentClient` now accepts `org_id` and `origin` constructor kwargs and + injects them automatically on every request: ```python - import requests from molecule_agent import RemoteAgentClient - session = requests.Session() - session.headers.update({ - "X-Molecule-Org-Id": "", - "Origin": "https://.moleculesai.app", - }) client = RemoteAgentClient( workspace_id="…", platform_url="https://.moleculesai.app", - session=session, + org_id="", # sets X-Molecule-Org-Id + origin="https://.moleculesai.app", # sets Origin ) ``` - A follow-up PR will accept `org_id` and `origin` constructor kwargs and - inject the headers automatically. - ## Design choices - **Blocking (`requests`), not async.** Drops into any runtime — script, diff --git a/molecule_agent/client.py b/molecule_agent/client.py index 6f7c330..c51b8ff 100644 --- a/molecule_agent/client.py +++ b/molecule_agent/client.py @@ -268,6 +268,13 @@ class RemoteAgentClient: 0700 permissions if missing. heartbeat_interval: Seconds between heartbeats in the run loop. state_poll_interval: Seconds between state polls in the run loop. + org_id: Optional tenant UUID for multi-tenant SaaS deployments + (``*.moleculesai.app``). When set the client sends + ``X-Molecule-Org-Id: `` on every request so the WAF + can route to the correct tenant. + origin: Optional origin string for multi-tenant SaaS deployments. + When set the client sends ``Origin: `` on every + request, preventing silent Next.js rewrites on SaaS edges. """ def __init__( @@ -281,6 +288,8 @@ class RemoteAgentClient: state_poll_interval: float = DEFAULT_STATE_POLL_INTERVAL, url_cache_ttl: float = DEFAULT_URL_CACHE_TTL, session: requests.Session | None = None, + org_id: str = "", + origin: str = "", ) -> None: self.workspace_id = workspace_id self.platform_url = platform_url.rstrip("/") @@ -289,6 +298,8 @@ class RemoteAgentClient: self.heartbeat_interval = heartbeat_interval self.state_poll_interval = state_poll_interval self.url_cache_ttl = url_cache_ttl + self._org_id = org_id + self._origin = origin # Phase 30.6 — sibling URL cache keyed by workspace id. Values are # (url, expires_at_unix_seconds). Process-memory only; we re-fetch # on restart because agent lifetimes are short enough that @@ -381,9 +392,14 @@ class RemoteAgentClient: def _auth_headers(self) -> dict[str, str]: tok = self.load_token() - if not tok: - return {} - return {"Authorization": f"Bearer {tok}"} + headers: dict[str, str] = {} + if tok: + headers["Authorization"] = f"Bearer {tok}" + if self._org_id: + headers["X-Molecule-Org-Id"] = self._org_id + if self._origin: + headers["Origin"] = self._origin + return headers def _get_with_retry( self, diff --git a/tests/test_remote_agent.py b/tests/test_remote_agent.py index ac60271..3c52aee 100644 --- a/tests/test_remote_agent.py +++ b/tests/test_remote_agent.py @@ -750,6 +750,78 @@ def test_delegate_sends_bearer_and_workspace_headers(client: RemoteAgentClient): assert kwargs["headers"]["X-Workspace-ID"] == "ws-abc-123" +def test_auth_headers_injects_org_id_and_origin(): + """org_id and origin kwargs are injected into every request headers.""" + session = MagicMock() + session.post.return_value = FakeResponse(200, {}) + client = RemoteAgentClient( + workspace_id="ws-test", + platform_url="https://platform.example.com", + org_id="org-uuid-123", + origin="https://tenant.moleculesai.app", + session=session, + ) + client.save_token("tok") + client.delegate(task="x", target_id="peer") + hdrs = session.post.call_args[1]["headers"] + assert hdrs["Authorization"] == "Bearer tok" + assert hdrs["X-Molecule-Org-Id"] == "org-uuid-123" + assert hdrs["Origin"] == "https://tenant.moleculesai.app" + + +def test_auth_headers_org_id_only(): + """origin can be omitted when only org_id is needed.""" + session = MagicMock() + session.post.return_value = FakeResponse(200, {}) + client = RemoteAgentClient( + workspace_id="ws-test", + platform_url="https://platform.example.com", + org_id="org-uuid-456", + session=session, + ) + client.save_token("tok") + client.poll_state() + hdrs = session.get.call_args[1]["headers"] + assert hdrs["Authorization"] == "Bearer tok" + assert hdrs["X-Molecule-Org-Id"] == "org-uuid-456" + assert "Origin" not in hdrs + + +def test_auth_headers_origin_only(): + """org_id can be omitted when only origin is needed.""" + session = MagicMock() + session.post.return_value = FakeResponse(200, {}) + client = RemoteAgentClient( + workspace_id="ws-test", + platform_url="https://platform.example.com", + origin="https://other-tenant.moleculesai.app", + session=session, + ) + client.save_token("tok") + client.poll_state() + hdrs = session.get.call_args[1]["headers"] + assert hdrs["Authorization"] == "Bearer tok" + assert "X-Molecule-Org-Id" not in hdrs + assert hdrs["Origin"] == "https://other-tenant.moleculesai.app" + + +def test_auth_headers_no_extra_when_unset(): + """When neither org_id nor origin is set, headers contain only auth.""" + session = MagicMock() + session.post.return_value = FakeResponse(200, {}) + client = RemoteAgentClient( + workspace_id="ws-test", + platform_url="https://platform.example.com", + session=session, + ) + client.save_token("tok") + client.poll_state() + hdrs = session.get.call_args[1]["headers"] + assert hdrs["Authorization"] == "Bearer tok" + assert "X-Molecule-Org-Id" not in hdrs + assert "Origin" not in hdrs + + def test_delegate_raises_on_http_error(client: RemoteAgentClient): client.save_token("tok") client._session.post.return_value = FakeResponse(500, {"error": "boom"})