Forward the workspace header through the wallet payment bridge (#1221)

strix/interface/cloud/billing.py

@@ -284,6 +284,7 @@ def _run_wallet_client(
         with wallet_payment_bridge(
             upstream_url=upstream_url,
             api_token=http.api_token(token),
+            workspace_id=http.expected_workspace_id(token_override=token is not None),
             expected_body=body_json.encode(),
             timeout=getattr(args, "timeout", None),
             response_observer=upstream_responses.append,

strix/interface/cloud/http.py

@@ -159,7 +159,7 @@ def request(
     headers = {
         "Authorization": f"Bearer {api_token(token)}",
     }
-    workspace_id = _expected_workspace_id(token_override=token is not None)
+    workspace_id = expected_workspace_id(token_override=token is not None)
     if workspace_id:
         headers["X-Strix-Workspace"] = workspace_id
     if idempotency_key is not None:
@@ -185,7 +185,7 @@ def request(
     return response
 
 
-def _expected_workspace_id(*, token_override: bool) -> str | None:
+def expected_workspace_id(*, token_override: bool) -> str | None:
     """Pin every request in this process to the workspace selected at startup."""
     if _workspace_id_override:
         return _workspace_id_override

strix/interface/cloud/payment_proxy.py

@@ -45,6 +45,7 @@ _HOP_BY_HOP_HEADERS = frozenset(
 class _BridgeState:
     upstream_url: str
     authorization: str
+    workspace_id: str | None
     expected_body: bytes
     path: str
     timeout: float
@@ -112,6 +113,7 @@ def _forward_request_headers(handler: BaseHTTPRequestHandler) -> dict[str, str]:
         "x-forwarded-proto",
         "x-real-ip",
         "x-strix-authorization",
+        "x-strix-workspace",
         "x-vercel-forwarded-for",
     }
     return {name: value for name, value in handler.headers.items() if name.lower() not in blocked}
@@ -163,6 +165,8 @@ def _make_handler(state: _BridgeState) -> type[BaseHTTPRequestHandler]:
 
             headers = _forward_request_headers(self)
             headers["X-Strix-Authorization"] = state.authorization
+            if state.workspace_id:
+                headers["X-Strix-Workspace"] = state.workspace_id
             try:
                 response = requests.request(
                     "POST",
@@ -243,6 +247,7 @@ def wallet_payment_bridge(
     *,
     upstream_url: str,
     api_token: str,
+    workspace_id: str | None = None,
     expected_body: bytes,
     timeout: float | None = None,
     response_observer: Callable[[WalletUpstreamResponse], None] | None = None,
@@ -258,6 +263,7 @@ def wallet_payment_bridge(
     state = _BridgeState(
         upstream_url=upstream_url,
         authorization=f"Bearer {api_token}",
+        workspace_id=workspace_id,
         expected_body=expected_body,
         path=path,
         timeout=timeout or _DEFAULT_REQUEST_TIMEOUT_S,

tests/test_cloud_payment_proxy.py

@@ -84,6 +84,7 @@ def test_bridge_forwards_only_the_approved_request_and_protected_headers(
     with payment_proxy.wallet_payment_bridge(
         upstream_url="https://app.example.test/api/v1/billing/topup",
         api_token="strix-secret",  # noqa: S106
+        workspace_id="org_trusted",
         expected_body=b'{"credits":5}',
         response_observer=observed.append,
     ) as wallet_url:
@@ -94,6 +95,7 @@ def test_bridge_forwards_only_the_approved_request_and_protected_headers(
                 "Authorization": "Payment wallet-proof",
                 "Proxy-Authorization": "Basic drop-me",
                 "X-Strix-Authorization": "Bearer attacker",
+                "X-Strix-Workspace": "org_attacker",
             },
         )
 
@@ -101,6 +103,7 @@ def test_bridge_forwards_only_the_approved_request_and_protected_headers(
         headers = captured[0]["headers"]
         assert headers["Authorization"] == "Payment wallet-proof"
         assert headers["X-Strix-Authorization"] == "Bearer strix-secret"
+        assert headers["X-Strix-Workspace"] == "org_trusted"
         assert "Proxy-Authorization" not in headers
         assert not any(name.lower() in {"host", "content-length"} for name in headers)
         assert observed == [