Skip to content

Commit 1631a8f

Browse files
committed
Add regression tests for refreshed access tokens (#1250)
1 parent 077f7ad commit 1631a8f

1 file changed

Lines changed: 231 additions & 0 deletions

File tree

Lines changed: 231 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,231 @@
1+
"""Regression tests for issue #1250 - refreshed access tokens stay invisible to tools.
2+
3+
A client that refreshes its access token mid-session sends the new token on
4+
the next request's `Authorization` header, but a tool calling
5+
`get_access_token()` kept receiving the token the session was opened with:
6+
the auth context middleware sets the contextvar on each HTTP request's task,
7+
while handlers run in the session's long-lived task.
8+
9+
These tests drive the server over raw HTTP in process because the property
10+
under test - which credential a *changed* `Authorization` header maps to
11+
inside an already-open session - cannot be produced by the typed client API,
12+
whose auth layer owns the header. Everything else (the session manager, the
13+
middleware stack, the tool handler) is the real stack.
14+
"""
15+
16+
import json
17+
from collections.abc import AsyncGenerator
18+
from contextlib import asynccontextmanager
19+
from typing import Any
20+
21+
import anyio
22+
import httpx
23+
import pytest
24+
from starlette.applications import Starlette
25+
from starlette.middleware import Middleware
26+
from starlette.middleware.authentication import AuthenticationMiddleware
27+
from starlette.routing import Mount
28+
29+
from mcp.server.auth.middleware.auth_context import AuthContextMiddleware, get_access_token
30+
from mcp.server.auth.middleware.bearer_auth import BearerAuthBackend
31+
from mcp.server.auth.provider import AccessToken, TokenVerifier
32+
from mcp.server.lowlevel import Server
33+
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
34+
from mcp.types import TextContent, Tool
35+
36+
pytestmark = pytest.mark.anyio
37+
38+
PROTOCOL_VERSION = "2025-06-18"
39+
BASE_URL = "http://127.0.0.1:8000"
40+
41+
42+
class AcceptAllTokenVerifier(TokenVerifier):
43+
"""Verifies any bearer token as the same client; the token string is its identity.
44+
45+
Models a refresh faithfully: the refreshed token belongs to the same
46+
client (and principal) as the one it replaces, only the token value differs.
47+
"""
48+
49+
async def verify_token(self, token: str) -> AccessToken | None:
50+
return AccessToken(token=token, client_id="test-client", scopes=["mcp"])
51+
52+
53+
def build_app(server: Server) -> tuple[Starlette, StreamableHTTPSessionManager]:
54+
session_manager = StreamableHTTPSessionManager(app=server, stateless=False, json_response=False)
55+
56+
@asynccontextmanager
57+
async def lifespan(app: Starlette) -> AsyncGenerator[None, None]:
58+
async with session_manager.run():
59+
yield
60+
61+
app = Starlette(
62+
routes=[Mount("/mcp", app=session_manager.handle_request)],
63+
middleware=[
64+
Middleware(AuthenticationMiddleware, backend=BearerAuthBackend(AcceptAllTokenVerifier())),
65+
Middleware(AuthContextMiddleware),
66+
],
67+
lifespan=lifespan,
68+
)
69+
return app, session_manager
70+
71+
72+
def make_server(call_tool: Any) -> Server:
73+
server = Server("token-echo")
74+
75+
@server.list_tools()
76+
async def list_tools() -> list[Tool]:
77+
return [
78+
Tool(
79+
name="echo_token",
80+
description="Return the token get_access_token() sees",
81+
inputSchema={"type": "object"},
82+
)
83+
]
84+
85+
server.call_tool()(call_tool)
86+
return server
87+
88+
89+
async def post_message(
90+
http: httpx.AsyncClient,
91+
payload: dict[str, Any],
92+
*,
93+
token: str,
94+
session_id: str | None = None,
95+
) -> httpx.Response:
96+
headers = {
97+
"Authorization": f"Bearer {token}",
98+
"Accept": "application/json, text/event-stream",
99+
"Content-Type": "application/json",
100+
}
101+
if session_id is not None:
102+
headers["mcp-session-id"] = session_id
103+
headers["mcp-protocol-version"] = PROTOCOL_VERSION
104+
return await http.post("/mcp/", headers=headers, json=payload)
105+
106+
107+
def response_for(response: httpx.Response, request_id: int) -> dict[str, Any]:
108+
"""Extract the JSON-RPC response for `request_id` from an SSE response body."""
109+
for line in response.text.splitlines():
110+
if line.startswith("data:"):
111+
message = json.loads(line.removeprefix("data:").strip())
112+
if message.get("id") == request_id:
113+
return message
114+
raise AssertionError(f"no response for request {request_id} in body: {response.text!r}")
115+
116+
117+
async def open_session(http: httpx.AsyncClient, token: str) -> str:
118+
"""Run the initialize handshake with `token` and return the session id."""
119+
response = await post_message(
120+
http,
121+
{
122+
"jsonrpc": "2.0",
123+
"id": 1,
124+
"method": "initialize",
125+
"params": {
126+
"protocolVersion": PROTOCOL_VERSION,
127+
"capabilities": {},
128+
"clientInfo": {"name": "issue-1250-regression", "version": "0"},
129+
},
130+
},
131+
token=token,
132+
)
133+
assert response.status_code == 200
134+
session_id = response.headers["mcp-session-id"]
135+
initialized = await post_message(
136+
http,
137+
{"jsonrpc": "2.0", "method": "notifications/initialized"},
138+
token=token,
139+
session_id=session_id,
140+
)
141+
assert initialized.status_code == 202
142+
return session_id
143+
144+
145+
async def call_echo_token(http: httpx.AsyncClient, session_id: str, token: str, request_id: int) -> str:
146+
"""Call the echo tool presenting `token`; return the token the tool observed."""
147+
response = await post_message(
148+
http,
149+
{
150+
"jsonrpc": "2.0",
151+
"id": request_id,
152+
"method": "tools/call",
153+
"params": {"name": "echo_token", "arguments": {}},
154+
},
155+
token=token,
156+
session_id=session_id,
157+
)
158+
assert response.status_code == 200
159+
message = response_for(response, request_id)
160+
content = message["result"]["content"]
161+
assert content[0]["type"] == "text"
162+
return content[0]["text"]
163+
164+
165+
async def test_a_refreshed_token_is_what_tools_see_in_the_same_session() -> None:
166+
"""After a refresh, the next call in the same session observes the new token.
167+
168+
Steps:
169+
1. Open a session presenting `token-a`; the tool sees `token-a`.
170+
2. Present the refreshed `token-b` on the next call in the same session;
171+
the tool must see `token-b`, not the session-opening `token-a`.
172+
"""
173+
174+
async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
175+
assert name == "echo_token"
176+
token = get_access_token()
177+
return [TextContent(type="text", text=token.token if token else "no-token")]
178+
179+
server = make_server(call_tool)
180+
app, session_manager = build_app(server)
181+
182+
with anyio.fail_after(5):
183+
async with session_manager.run():
184+
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url=BASE_URL) as http:
185+
session_id = await open_session(http, "token-a")
186+
assert await call_echo_token(http, session_id, "token-a", 2) == "token-a"
187+
assert await call_echo_token(http, session_id, "token-b", 3) == "token-b"
188+
189+
190+
async def test_parallel_sessions_each_see_their_own_current_token() -> None:
191+
"""Interleaved calls in two sessions never observe each other's tokens.
192+
193+
Both calls are held in flight at the same time (the handler parks on a
194+
barrier until both have read their token), so the scheduler cannot
195+
serialize them into looking correct.
196+
"""
197+
in_flight = 0
198+
both_in_flight = anyio.Event()
199+
release = anyio.Event()
200+
201+
async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
202+
nonlocal in_flight
203+
assert name == "echo_token"
204+
token = get_access_token()
205+
seen = token.token if token else "no-token"
206+
in_flight += 1
207+
if in_flight == 2:
208+
both_in_flight.set()
209+
await release.wait()
210+
return [TextContent(type="text", text=seen)]
211+
212+
server = make_server(call_tool)
213+
app, session_manager = build_app(server)
214+
215+
results: dict[str, str] = {}
216+
217+
async def call(http: httpx.AsyncClient, session_id: str, token: str, request_id: int) -> None:
218+
results[token] = await call_echo_token(http, session_id, token, request_id)
219+
220+
with anyio.fail_after(5):
221+
async with session_manager.run():
222+
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url=BASE_URL) as http:
223+
session_one = await open_session(http, "token-a")
224+
session_two = await open_session(http, "token-c")
225+
async with anyio.create_task_group() as tg:
226+
# Session one presents its refreshed token; session two its own.
227+
tg.start_soon(call, http, session_one, "token-b", 2)
228+
tg.start_soon(call, http, session_two, "token-c", 3)
229+
await both_in_flight.wait()
230+
release.set()
231+
assert results == {"token-b": "token-b", "token-c": "token-c"}

0 commit comments

Comments
 (0)