fix: Forward X-Org-Id during automation auth

openhands/automation/auth.py

@@ -39,6 +39,7 @@
 # Auth cache - initialized lazily to use config values
 _auth_cache: TTLCache[str, "AuthenticatedUser"] | None = None
 SESSION_COOKIE_NAME = "keycloak_auth"
+X_ORG_ID_HEADER = "X-Org-Id"
 # Keep parity with OpenHands' cookie chunking helper: 8 * 3000 bytes is
 # comfortably above expected session token sizes while staying bounded.
 MAX_SESSION_COOKIE_CHUNKS = 8
@@ -185,9 +186,37 @@ def clear_auth_cache() -> None:
     _get_auth_cache().clear()
 
 
-def _credential_cache_key(credential: str) -> str:
-    """Hash a credential for use as a cache key (never store raw credential)."""
-    return hashlib.sha256(credential.encode()).hexdigest()
+def _credential_cache_key(
+    credential: str, auth_method: AuthMethod, x_org_id: str | None
+) -> str:
+    """Hash auth scope for use as a cache key (never store raw credentials)."""
+    cache_material = "\0".join((auth_method.value, x_org_id or "", credential))
+    return hashlib.sha256(cache_material.encode()).hexdigest()
+
+
+def _request_has_header(request: Request, name: str) -> bool:
+    try:
+        return name in request.headers
+    except TypeError:
+        return False
+
+
+def _extract_x_org_id(request: Request) -> str | None:
+    """Extract and normalize the optional organization scope header."""
+    if not _request_has_header(request, X_ORG_ID_HEADER):
+        return None
+
+    header_value = request.headers.get(X_ORG_ID_HEADER, "").strip()
+    if not header_value:
+        return None
+
+    try:
+        return str(uuid.UUID(header_value))
+    except ValueError as exc:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="Invalid X-Org-Id header (must be a UUID)",
+        ) from exc
 
 
 def _is_rate_limited(response: httpx.Response) -> bool:
@@ -446,9 +475,10 @@ async def authenticate_request(
 
     # --- Resolve credential (API key or cookie) ---
     credential, auth_method = _extract_credential(request)
+    x_org_id = _extract_x_org_id(request)
 
     # --- Cache lookup ---
-    cache_key = _credential_cache_key(credential)
+    cache_key = _credential_cache_key(credential, auth_method, x_org_id)
     auth_cache = _get_auth_cache()
     cached_user = auth_cache.get(cache_key)
     if cached_user is not None:
@@ -462,6 +492,8 @@ async def authenticate_request(
         if auth_method == AuthMethod.API_KEY
         else {"Cookie": f"{SESSION_COOKIE_NAME}={credential}"}
     )
+    if x_org_id:
+        outbound_headers[X_ORG_ID_HEADER] = x_org_id
 
     try:
         resp = await _make_auth_request_with_retry(

tests/test_auth.py

@@ -74,6 +74,11 @@ def _get(name: str, default: str = "") -> str:
     return _get
 
 
+def _set_mock_headers(request, headers_dict: dict[str, str]):
+    request.headers.get.side_effect = _make_header_getter(headers_dict)
+    request.headers.__contains__.side_effect = lambda name: name in headers_dict
+
+
 class TestAuthentication:
     """Tests for authenticate_request function with API key auth.
 
@@ -105,6 +110,67 @@ async def test_authenticate_valid_api_key(self, mock_request, mock_http_client):
         assert result.auth_method == AuthMethod.API_KEY
         assert result.api_key == "valid-api-key"
 
+    async def test_authenticate_forwards_x_org_id_with_api_key(
+        self, mock_request, mock_http_client
+    ):
+        """Bearer auth forwards the selected org scope to OpenHands /users/me."""
+        _set_mock_headers(
+            mock_request,
+            {
+                "Authorization": "Bearer valid-api-key",
+                "X-Org-Id": str(TEST_ORG_ID),
+            },
+        )
+
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = MOCK_USERS_ME_RESPONSE
+        mock_http_client.get = AsyncMock(return_value=mock_response)
+
+        result = await authenticate_request(mock_request, client=mock_http_client)
+
+        assert result.org_id == TEST_ORG_ID
+        headers = mock_http_client.get.call_args[1]["headers"]
+        assert headers["Authorization"] == "Bearer valid-api-key"
+        assert headers["X-Org-Id"] == str(TEST_ORG_ID)
+
+    async def test_authenticate_forwards_x_org_id_with_cookie(
+        self, mock_request, mock_http_client
+    ):
+        """Cookie auth forwards the selected org scope to OpenHands /users/me."""
+        _set_mock_headers(mock_request, {"X-Org-Id": str(TEST_ORG_ID)})
+        mock_request.cookies = {"keycloak_auth": "valid-cookie-value"}
+
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = MOCK_USERS_ME_RESPONSE
+        mock_http_client.get = AsyncMock(return_value=mock_response)
+
+        result = await authenticate_request(mock_request, client=mock_http_client)
+
+        assert result.org_id == TEST_ORG_ID
+        headers = mock_http_client.get.call_args[1]["headers"]
+        assert headers["Cookie"] == "keycloak_auth=valid-cookie-value"
+        assert headers["X-Org-Id"] == str(TEST_ORG_ID)
+
+    async def test_authenticate_rejects_invalid_x_org_id(
+        self, mock_request, mock_http_client
+    ):
+        """Invalid org scope headers fail as client errors before auth forwarding."""
+        _set_mock_headers(
+            mock_request,
+            {
+                "Authorization": "Bearer valid-api-key",
+                "X-Org-Id": "not-a-uuid",
+            },
+        )
+
+        with pytest.raises(HTTPException) as exc_info:
+            await authenticate_request(mock_request, client=mock_http_client)
+
+        assert exc_info.value.status_code == 400
+        mock_http_client.get.assert_not_called()
+
     async def test_authenticate_extracts_model_profile_metadata(
         self, mock_request, mock_http_client
     ):
@@ -779,7 +845,6 @@ async def test_different_keys_cached_separately(
             "role": "member",
             "permissions": [],
         }
-
         mock_http_client.get = AsyncMock(side_effect=[mock_response1, mock_response2])
 
         # First key
@@ -794,6 +859,53 @@ async def test_different_keys_cached_separately(
         assert result1.user_id == TEST_USER_ID
         assert result2.user_id == user2_id
 
+    async def test_auth_cache_is_scoped_by_x_org_id(
+        self, mock_request, mock_http_client
+    ):
+        """The same credential can authenticate separately for different orgs."""
+        other_org_id = uuid.UUID("99999999-9999-4999-8999-999999999999")
+
+        response_for_default_org = MagicMock()
+        response_for_default_org.status_code = 200
+        response_for_default_org.json.return_value = MOCK_USERS_ME_RESPONSE
+
+        response_for_other_org = MagicMock()
+        response_for_other_org.status_code = 200
+        response_for_other_org.json.return_value = {
+            **MOCK_USERS_ME_RESPONSE,
+            "org_id": str(other_org_id),
+        }
+        mock_http_client.get = AsyncMock(
+            side_effect=[response_for_default_org, response_for_other_org]
+        )
+
+        _set_mock_headers(
+            mock_request,
+            {
+                "Authorization": "Bearer shared-key",
+                "X-Org-Id": str(TEST_ORG_ID),
+            },
+        )
+        result1 = await authenticate_request(mock_request, client=mock_http_client)
+
+        _set_mock_headers(
+            mock_request,
+            {
+                "Authorization": "Bearer shared-key",
+                "X-Org-Id": str(other_org_id),
+            },
+        )
+        result2 = await authenticate_request(mock_request, client=mock_http_client)
+
+        assert mock_http_client.get.call_count == 2
+        assert result1.org_id == TEST_ORG_ID
+        assert result2.org_id == other_org_id
+        forwarded_orgs = [
+            call[1]["headers"]["X-Org-Id"]
+            for call in mock_http_client.get.call_args_list
+        ]
+        assert forwarded_orgs == [str(TEST_ORG_ID), str(other_org_id)]
+
     async def test_cookie_and_api_key_cached_separately(
         self, mock_request, mock_http_client
     ):