diff --git a/src/prefect_mcp_server/_prefect_client/identity.py b/src/prefect_mcp_server/_prefect_client/identity.py index 3070526..089013b 100644 --- a/src/prefect_mcp_server/_prefect_client/identity.py +++ b/src/prefect_mcp_server/_prefect_client/identity.py @@ -2,8 +2,6 @@ from uuid import UUID -from prefect.client.cloud import CloudUnauthorizedError - from prefect_mcp_server import cloud_oauth from prefect_mcp_server._prefect_client.client import ( get_prefect_client, @@ -51,8 +49,13 @@ async def get_identity(workspace_id: UUID | None = None) -> IdentityResult: """Get identity and connection information for the current Prefect instance.""" try: access_token = cloud_oauth.current_oauth_access_token() - if workspace_id is None and cloud_oauth.settings.enabled and access_token: - identity = await _get_cloud_oauth_identity(access_token) + if cloud_oauth.settings.enabled and access_token: + # OAuth tokens are workspace-scoped and cannot call account-level + # endpoints like /me/ or /accounts/{id}, so always describe the + # grant instead of fetching full Cloud identity. + identity = await _get_cloud_oauth_identity( + access_token, workspace_id=workspace_id + ) return { "success": True, "identity": identity, @@ -73,19 +76,7 @@ async def get_identity(workspace_id: UUID | None = None) -> IdentityResult: workspace_id=workspace_id ) as cloud_client: # Get user info from /me/ endpoint - try: - me_data = await cloud_client.get("/me/") - except CloudUnauthorizedError: - if cloud_oauth.settings.enabled and access_token: - identity = await _get_cloud_oauth_identity( - access_token, workspace_id=workspace_id - ) - return { - "success": True, - "identity": identity, - "error": None, - } - raise + me_data = await cloud_client.get("/me/") user_info: UserInfo = { "id": str(me_data.get("id")) if me_data.get("id") else None, diff --git a/tests/test_cloud_oauth.py b/tests/test_cloud_oauth.py index b96a2e5..23110a1 100644 --- a/tests/test_cloud_oauth.py +++ b/tests/test_cloud_oauth.py @@ -7,7 +7,6 @@ import pytest from fastmcp import Client from fastmcp.server.auth.providers.jwt import JWTVerifier -from prefect.client.cloud import CloudUnauthorizedError from starlette.applications import Starlette from starlette.testclient import TestClient @@ -130,25 +129,14 @@ async def test_get_identity_describes_oauth_grant_without_workspace() -> None: ) -async def test_get_identity_describes_service_account_oauth_grant_with_workspace() -> ( - None -): +async def test_get_identity_with_workspace_never_calls_account_endpoints() -> None: + """OAuth tokens are workspace-scoped; account-level endpoints return 403.""" workspace = cloud_oauth.WorkspaceRef( account_id=ACCOUNT_ID, account_handle="acme", workspace_id=WORKSPACE_ID, workspace_handle="prod", ) - mock_client = AsyncMock() - mock_client.api_url = ( - f"https://api.prefect.cloud/api/accounts/{ACCOUNT_ID}/workspaces/{WORKSPACE_ID}" - ) - mock_cloud_client = AsyncMock() - mock_cloud_client.get = AsyncMock( - side_effect=CloudUnauthorizedError( - "Only users (not service accounts) can access this endpoint." - ) - ) with ( patch( @@ -172,12 +160,11 @@ async def test_get_identity_describes_service_account_oauth_grant_with_workspace AsyncMock(return_value=[workspace]), ) as mock_list_authorized_workspaces, ): - mock_get_client.return_value.__aenter__.return_value = mock_client - mock_get_cloud_client.return_value.__aenter__.return_value = mock_cloud_client - mock_get_cloud_client.return_value.__aexit__.return_value = None - result = await get_identity(workspace_id=WORKSPACE_ID) + mock_get_client.assert_not_called() + mock_get_cloud_client.assert_not_called() + assert result["success"] is True assert result["identity"] == { "api_url": "https://api.prefect.cloud",