Merge pull request 'feat(molecule_agent): add org_id and origin kwargs to RemoteAgentClient' (#7) from feat/client-multi-tenant-headers into main

molecule_agent/README.md

@@ -219,30 +219,24 @@ reference it directly.
     - `Origin: <PLATFORM_URL>` — `/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 <token>` 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": "<your-org-uuid>",
-      "Origin": "https://<your-tenant>.moleculesai.app",
-  })
   client = RemoteAgentClient(
       workspace_id="…",
       platform_url="https://<your-tenant>.moleculesai.app",
-      session=session,
+      org_id="<your-org-uuid>",          # sets X-Molecule-Org-Id
+      origin="https://<your-tenant>.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,

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: <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: <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,

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"})