feat(molecule_agent): add org_id and origin kwargs to RemoteAgentClient
Some checks failed
Test / test (3.13) (pull_request) Failing after 1s
Test / test (3.12) (pull_request) Failing after 3s
Test / test (3.11) (pull_request) Failing after 4s

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 <noreply@anthropic.com>
This commit is contained in:
Molecule AI · sdk-dev 2026-05-10 13:33:17 +00:00
parent e8b9d42fe6
commit a3bba8a3f3
3 changed files with 99 additions and 17 deletions

View File

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

View File

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

View File

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