Skip to content

Commit b701a8f

Browse files
zzstoatzzclaude
andauthored
fix: get_identity and review_rate_limits broken in Cloud OAuth mode (#200)
* fix: get_identity 403 in Cloud OAuth mode OAuth access tokens are workspace-scoped and cannot call account-level endpoints like /me/ or /accounts/{id}. get_identity only short-circuited to the OAuth grant description when workspace_id was None, so passing a workspace_id fell through to the full Cloud identity path and failed with 403 on /accounts/{account_id}. Always describe the OAuth grant when an OAuth token is present, and drop the now-unreachable service-account fallback. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix: review_rate_limits unusable in Cloud OAuth mode The tool errored with 'workspace_id is required when using Prefect Cloud OAuth mode' but its input schema had no workspace_id parameter, making it impossible to call. Thread workspace_id from the tool through get_rate_limits to both client factories. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix: replace retired claude-opus-4-1 eval judge model with claude-opus-5 claude-opus-4-1-20250805 was retired on 2026-08-05, so every eval that reaches the judge now fails with a 404 from the Anthropic API. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * style: ruff format after eval model swap Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
1 parent 2940a3e commit b701a8f

7 files changed

Lines changed: 85 additions & 43 deletions

File tree

‎evals/README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ MODEL_PROVIDER=openai just evals
3636
SIMPLE_AGENT_MODEL=openai:gpt-4o REASONING_AGENT_MODEL=openai:gpt-4.1 just evals
3737

3838
# override just one model
39-
REASONING_AGENT_MODEL=anthropic:claude-opus-4-1-20250805 just evals
39+
REASONING_AGENT_MODEL=anthropic:claude-opus-5 just evals
4040
```
4141

4242
Provider defaults:

‎evals/conftest.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
- Default for anthropic: "anthropic:claude-sonnet-4-5"
1414
- Default for openai: "openai:gpt-4.1"
1515
16-
EVALUATOR_MODEL: Override default evaluator model (default: "anthropic:claude-opus-4-1-20250805")
16+
EVALUATOR_MODEL: Override default evaluator model (default: "anthropic:claude-opus-5")
1717
"""
1818

1919
import os
@@ -221,9 +221,7 @@ async def _evaluate(evaluation_prompt: str, agent_response: str) -> None:
221221
Raises:
222222
AssertionError: If evaluation fails, with explanation
223223
"""
224-
evaluator_model = os.getenv(
225-
"EVALUATOR_MODEL", "anthropic:claude-opus-4-1-20250805"
226-
)
224+
evaluator_model = os.getenv("EVALUATOR_MODEL", "anthropic:claude-opus-5")
227225
evaluator = Agent(
228226
name="Response Evaluator",
229227
model=evaluator_model,

‎src/prefect_mcp_server/_prefect_client/identity.py‎

Lines changed: 8 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,6 @@
22

33
from uuid import UUID
44

5-
from prefect.client.cloud import CloudUnauthorizedError
6-
75
from prefect_mcp_server import cloud_oauth
86
from prefect_mcp_server._prefect_client.client import (
97
get_prefect_client,
@@ -51,8 +49,13 @@ async def get_identity(workspace_id: UUID | None = None) -> IdentityResult:
5149
"""Get identity and connection information for the current Prefect instance."""
5250
try:
5351
access_token = cloud_oauth.current_oauth_access_token()
54-
if workspace_id is None and cloud_oauth.settings.enabled and access_token:
55-
identity = await _get_cloud_oauth_identity(access_token)
52+
if cloud_oauth.settings.enabled and access_token:
53+
# OAuth tokens are workspace-scoped and cannot call account-level
54+
# endpoints like /me/ or /accounts/{id}, so always describe the
55+
# grant instead of fetching full Cloud identity.
56+
identity = await _get_cloud_oauth_identity(
57+
access_token, workspace_id=workspace_id
58+
)
5659
return {
5760
"success": True,
5861
"identity": identity,
@@ -73,19 +76,7 @@ async def get_identity(workspace_id: UUID | None = None) -> IdentityResult:
7376
workspace_id=workspace_id
7477
) as cloud_client:
7578
# Get user info from /me/ endpoint
76-
try:
77-
me_data = await cloud_client.get("/me/")
78-
except CloudUnauthorizedError:
79-
if cloud_oauth.settings.enabled and access_token:
80-
identity = await _get_cloud_oauth_identity(
81-
access_token, workspace_id=workspace_id
82-
)
83-
return {
84-
"success": True,
85-
"identity": identity,
86-
"error": None,
87-
}
88-
raise
79+
me_data = await cloud_client.get("/me/")
8980

9081
user_info: UserInfo = {
9182
"id": str(me_data.get("id")) if me_data.get("id") else None,

‎src/prefect_mcp_server/_prefect_client/rate_limits.py‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
"""Rate limits inspection for Prefect Cloud."""
22

33
from datetime import datetime, timedelta, timezone
4+
from uuid import UUID
45

56
from prefect_mcp_server._prefect_client.client import (
67
get_prefect_client,
@@ -37,20 +38,23 @@
3738
async def get_rate_limits(
3839
since: datetime | None = None,
3940
until: datetime | None = None,
41+
workspace_id: UUID | None = None,
4042
) -> RateLimitsResult:
4143
"""Get rate limit usage for the Cloud account across all common operation groups.
4244
4345
Args:
4446
since: Start time for usage data (defaults to 3 days ago)
4547
until: End time for usage data (defaults to 1 minute ago)
48+
workspace_id: Workspace to resolve the account from (required in
49+
Prefect Cloud OAuth mode)
4650
4751
Returns:
4852
RateLimitsResult with grouped throttling periods showing which operation
4953
groups (runs, deployments, writing-logs, etc.) were throttled during each
5054
consecutive stretch of throttling
5155
"""
5256
try:
53-
async with get_prefect_client() as client:
57+
async with get_prefect_client(workspace_id=workspace_id) as client:
5458
api_url = str(client.api_url)
5559

5660
# Extract account_id from Cloud API URL
@@ -65,7 +69,9 @@ async def get_rate_limits(
6569
if until is None:
6670
until = datetime.now(timezone.utc) - timedelta(minutes=1)
6771

68-
async with get_prefect_cloud_client() as cloud_client:
72+
async with get_prefect_cloud_client(
73+
workspace_id=workspace_id
74+
) as cloud_client:
6975
# Query all keys in one request
7076
params = {
7177
"since": since.isoformat(),

‎src/prefect_mcp_server/server.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -500,6 +500,7 @@ async def get_object_schema(
500500

501501

502502
async def review_rate_limits(
503+
workspace_id: WorkspaceId | None = None,
503504
since: Annotated[
504505
datetime | None,
505506
Field(
@@ -539,7 +540,9 @@ async def review_rate_limits(
539540
- Check recent throttling: review_rate_limits()
540541
- Custom time range: review_rate_limits(since="2025-09-30T00:00:00Z", until="2025-10-01T00:00:00Z")
541542
"""
542-
return await _prefect_client.get_rate_limits(since=since, until=until)
543+
return await _prefect_client.get_rate_limits(
544+
since=since, until=until, workspace_id=workspace_id
545+
)
543546

544547

545548
CORE_TOOLS = (

‎tests/test_cloud_oauth.py‎

Lines changed: 5 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77
import pytest
88
from fastmcp import Client
99
from fastmcp.server.auth.providers.jwt import JWTVerifier
10-
from prefect.client.cloud import CloudUnauthorizedError
1110
from starlette.applications import Starlette
1211
from starlette.testclient import TestClient
1312

@@ -130,25 +129,14 @@ async def test_get_identity_describes_oauth_grant_without_workspace() -> None:
130129
)
131130

132131

133-
async def test_get_identity_describes_service_account_oauth_grant_with_workspace() -> (
134-
None
135-
):
132+
async def test_get_identity_with_workspace_never_calls_account_endpoints() -> None:
133+
"""OAuth tokens are workspace-scoped; account-level endpoints return 403."""
136134
workspace = cloud_oauth.WorkspaceRef(
137135
account_id=ACCOUNT_ID,
138136
account_handle="acme",
139137
workspace_id=WORKSPACE_ID,
140138
workspace_handle="prod",
141139
)
142-
mock_client = AsyncMock()
143-
mock_client.api_url = (
144-
f"https://api.prefect.cloud/api/accounts/{ACCOUNT_ID}/workspaces/{WORKSPACE_ID}"
145-
)
146-
mock_cloud_client = AsyncMock()
147-
mock_cloud_client.get = AsyncMock(
148-
side_effect=CloudUnauthorizedError(
149-
"Only users (not service accounts) can access this endpoint."
150-
)
151-
)
152140

153141
with (
154142
patch(
@@ -172,12 +160,11 @@ async def test_get_identity_describes_service_account_oauth_grant_with_workspace
172160
AsyncMock(return_value=[workspace]),
173161
) as mock_list_authorized_workspaces,
174162
):
175-
mock_get_client.return_value.__aenter__.return_value = mock_client
176-
mock_get_cloud_client.return_value.__aenter__.return_value = mock_cloud_client
177-
mock_get_cloud_client.return_value.__aexit__.return_value = None
178-
179163
result = await get_identity(workspace_id=WORKSPACE_ID)
180164

165+
mock_get_client.assert_not_called()
166+
mock_get_cloud_client.assert_not_called()
167+
181168
assert result["success"] is True
182169
assert result["identity"] == {
183170
"api_url": "https://api.prefect.cloud",

‎tests/test_rate_limits.py‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,12 @@
22

33
from datetime import datetime, timezone
44
from unittest.mock import AsyncMock, patch
5+
from uuid import UUID
6+
7+
from fastmcp import Client
58

69
from prefect_mcp_server._prefect_client.rate_limits import get_rate_limits
10+
from prefect_mcp_server.server import build_prefect_mcp_server
711

812

913
async def test_get_rate_limits_empty_response() -> None:
@@ -816,3 +820,56 @@ async def test_get_rate_limits_handles_errors() -> None:
816820
assert result["account_id"] is None
817821
assert result["summary"] is None
818822
assert result["throttling_periods"] is None
823+
824+
825+
async def test_get_rate_limits_passes_workspace_id_to_clients() -> None:
826+
"""workspace_id must reach both clients so OAuth mode can scope the request."""
827+
workspace_id = UUID("bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb")
828+
829+
mock_client = AsyncMock()
830+
mock_client.api_url = (
831+
f"https://api.prefect.cloud/api/accounts/abc-123/workspaces/{workspace_id}"
832+
)
833+
834+
mock_cloud_client = AsyncMock()
835+
mock_cloud_client.get = AsyncMock(
836+
return_value={
837+
"account": "abc-123",
838+
"since": "2025-09-30T00:00:00Z",
839+
"until": "2025-10-01T00:00:00Z",
840+
"minutes": [],
841+
"keys": {},
842+
}
843+
)
844+
845+
with (
846+
patch(
847+
"prefect_mcp_server._prefect_client.rate_limits.get_prefect_client"
848+
) as mock_get_client,
849+
patch(
850+
"prefect_mcp_server._prefect_client.rate_limits.get_prefect_cloud_client"
851+
) as mock_get_cloud_client,
852+
):
853+
mock_get_client.return_value.__aenter__.return_value = mock_client
854+
mock_get_cloud_client.return_value.__aenter__.return_value = mock_cloud_client
855+
mock_get_cloud_client.return_value.__aexit__.return_value = None
856+
857+
result = await get_rate_limits(workspace_id=workspace_id)
858+
859+
assert result["success"] is True
860+
mock_get_client.assert_called_once_with(workspace_id=workspace_id)
861+
mock_get_cloud_client.assert_called_once_with(workspace_id=workspace_id)
862+
863+
864+
async def test_review_rate_limits_tool_accepts_workspace_id() -> None:
865+
"""The MCP tool schema must expose workspace_id for OAuth mode."""
866+
server = build_prefect_mcp_server(
867+
include_docs_proxy=False,
868+
include_cloud_tools=True,
869+
)
870+
871+
async with Client(server) as client:
872+
tools = {tool.name: tool for tool in await client.list_tools()}
873+
874+
schema = tools["review_rate_limits"].inputSchema
875+
assert "workspace_id" in schema["properties"]

0 commit comments

Comments
 (0)