LCORE-1246: Add OAuth authentication method for MCP servers#1156
Open
jrobertboos wants to merge 1 commit intolightspeed-core:mainfrom
Open
LCORE-1246: Add OAuth authentication method for MCP servers#1156jrobertboos wants to merge 1 commit intolightspeed-core:mainfrom
jrobertboos wants to merge 1 commit intolightspeed-core:mainfrom
Conversation
- Introduced a new OAuth authentication method for MCP servers, allowing clients to use the "oauth" keyword in authorization headers. - Updated documentation to reflect the addition of OAuth as a supported authentication method. - Enhanced error handling in the tools endpoint to return a 401 status with a WWW-Authenticate header when OAuth is required but no token is provided. - Modified relevant tests to cover the new OAuth functionality and ensure proper error responses.
Contributor
WalkthroughIntroduces OAuth as a new MCP server authentication method. Adds OAuth constant, updates authorization header resolution to handle oauth keyword, implements 401/WWW-Authenticate response handling in tools endpoint, and extends documentation and tests to cover OAuth flows. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant ToolsEndpoint as Tools Endpoint
participant MCP as MCP Server
Client->>ToolsEndpoint: GET /tools
ToolsEndpoint->>MCP: Request tools (no token)
MCP-->>ToolsEndpoint: 401 Unauthorized + WWW-Authenticate header
ToolsEndpoint->>ToolsEndpoint: Construct HTTPException with header
ToolsEndpoint-->>Client: 401 + WWW-Authenticate header
Client->>Client: Parse OAuth challenge
Client->>ToolsEndpoint: GET /tools + Authorization header
ToolsEndpoint->>MCP: Request tools (with token)
MCP-->>ToolsEndpoint: 200 OK + tools
ToolsEndpoint-->>Client: 200 OK + tools
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Tip Issue Planner is now in beta. Read the docs and try it out! Share your feedback on Discord. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 Fix all issues with AI agents
In `@docs/config.md`:
- Line 374: In the authorization_headers description update the phrase "client
provided token" to the hyphenated form "client-provided token" (the clause that
explains using the string 'client' instead of a file path); edit the docs entry
for authorization_headers to use "client-provided" wherever it appears to
correct the compound adjective usage.
In `@src/app/endpoints/tools.py`:
- Line 5: The module imports the synchronous requests library (requests) but
defines async endpoint handlers, which will block the event loop; replace the
requests import with an async HTTP client (e.g., aiohttp.ClientSession or
httpx.AsyncClient), update any call sites that use requests.get/post/etc. to use
the async client with await (for example create or accept an
aiohttp.ClientSession and call await session.get(...)), and ensure proper
session lifecycle (use async context manager or dependency-injected session) so
all network I/O in your async endpoint functions remains non-blocking.
- Around line 93-104: The block catching AuthenticationError is performing a
blocking synchronous probe and not guarding against probe failures or missing
headers; update the probe to use an async HTTP client (e.g., aiohttp) and await
the request instead of requests.get, wrap the probe in try/except to catch
connection, timeout, and client errors (e.g., aiohttp.ClientError,
asyncio.TimeoutError) and fall back to a safe UnauthorizedResponse when the
probe fails, and when constructing the HTTPException use
resp.headers.get("WWW-Authenticate") with a sensible default (or omit the header
if absent) to avoid KeyError; apply the same async/non-blocking error-handling
pattern to the get_mcp_tools function in src/utils/responses.py.
In `@src/models/config.py`:
- Around line 512-514: The concatenated description string for the special cases
is missing a space between the end of the "client" sentence and the start of "3.
Usage of the oauth..." sentence; update the string fragments (the literal
containing "To specify this use a string 'client' instead of the file path." or
the immediately following fragment starting "3. Usage of the oauth token in the
header.") so there is a space either appended to the first fragment or prefixed
to the second fragment to ensure the final description reads "...file path. 3.
Usage of the oauth..." (refer to the config description string that includes the
"'client'" and "'oauth'" special-case text).
In `@src/utils/responses.py`:
- Around line 389-405: The probe using requests.get in the get_mcp_tools logic
blocks the async event loop and lacks error handling for connection/timeouts and
missing WWW-Authenticate header; replace the synchronous probe with an async
HTTP call (e.g., httpx.AsyncClient) or run requests.get in an executor, wrap the
call in try/except to catch request exceptions (ConnectionError, Timeout,
HTTPError) and handle them gracefully, and when building the HTTPException use
resp.headers.get("WWW-Authenticate") (only include the header if present) so you
don't KeyError; update the code paths in get_mcp_tools and callers
(prepare_tools) accordingly to await the async probe or handle executor futures.
🧹 Nitpick comments (6)
README.md (1)
462-462: Consider whether OAuth servers should appear in the/v1/mcp-auth/client-optionsendpoint response.This line states that only
"client"servers are returned by the discovery endpoint. Since"oauth"servers also require client-provided tokens (viaMCP-HEADERS), clients would benefit from discovering them here too. If intentionally excluded, a brief note explaining why would help.tests/unit/utils/test_mcp_auth_headers.py (1)
96-112: Consider adding"oauth"to the mixed-types test.The mixed-types test covers
client, file, andkubernetesbut notoauth. Adding it here would validate that all four header types work together.Suggested addition
headers = { "Authorization": "client", "X-API-Key": str(secret_file), "X-K8s-Token": "kubernetes", + "X-OAuth-Token": "oauth", } result = resolve_authorization_headers(headers) # Special keywords should be preserved, file should be resolved assert result["Authorization"] == "client" assert result["X-API-Key"] == "file-secret" assert result["X-K8s-Token"] == "kubernetes" + assert result["X-OAuth-Token"] == "oauth"tests/unit/utils/test_responses.py (1)
513-537: Test covers the happy path well, but consider adding edge-case tests.The test correctly validates the OAuth 401 flow. However, given the issues flagged in the production code, consider adding tests for:
- MCP server probe returning a response without
WWW-Authenticateheader (currently would causeKeyError)requests.getraising aConnectionErrorduring the probeThese would serve as regression tests once the production code is hardened.
src/utils/mcp_auth_headers.py (1)
19-22: Docstring updated correctly for OAuth.The parameter documentation now covers all three special keywords. Consider also adding an OAuth example in the
Examplessection (lines 28-39) for completeness — entirely optional.src/models/config.py (1)
535-551: Docstring doesn't mention'oauth'as a preserved special value.The
resolve_auth_headersdocstring (Line 542) states only'kubernetes'and'client'are preserved for later substitution. Now that'oauth'is also a pass-through keyword, consider updating the docstring to include it.src/utils/responses.py (1)
354-361: OAuth token resolution is an exact duplicate of theMCP_AUTH_CLIENTcase.The
MCP_AUTH_OAUTHbranch (lines 354-361) is identical toMCP_AUTH_CLIENT(lines 346-353). Consider merging them to avoid duplication.Proposed fix
- case constants.MCP_AUTH_CLIENT: - # use client provided token - if mcp_headers is None: - return None - c_headers = mcp_headers.get(mcp_server.name, None) - if c_headers is None: - return None - return c_headers.get(header, None) - case constants.MCP_AUTH_OAUTH: - # use oauth token - if mcp_headers is None: - return None - c_headers = mcp_headers.get(mcp_server.name, None) - if c_headers is None: - return None - return c_headers.get(header, None) + case constants.MCP_AUTH_CLIENT | constants.MCP_AUTH_OAUTH: + # use client or oauth provided token + if mcp_headers is None: + return None + c_headers = mcp_headers.get(mcp_server.name, None) + if c_headers is None: + return None + return c_headers.get(header, None)
| | provider_id | string | MCP provider identification | | ||
| | url | string | URL of the MCP server | | ||
| | authorization_headers | object | Headers to send to the MCP server. The map contains the header name and the path to a file containing the header value (secret). There are 2 special cases: 1. Usage of the kubernetes token in the header. To specify this use a string 'kubernetes' instead of the file path. 2. Usage of the client provided token in the header. To specify this use a string 'client' instead of the file path. | | ||
| | authorization_headers | object | Headers to send to the MCP server. The map contains the header name and the path to a file containing the header value (secret). There are 3 special cases: 1. Usage of the kubernetes token in the header — use the string 'kubernetes' instead of the file path. 2. Usage of the client provided token in the header — use the string 'client' instead of the file path. 3. Usage of OAuth token (resolved at request time or 401 with WWW-Authenticate) — use the string 'oauth' instead of the file path. | |
Contributor
There was a problem hiding this comment.
Minor grammar: "client provided" → "client-provided".
Static analysis flagged this. When used as a compound adjective before a noun, it should be hyphenated.
Suggested fix
-| authorization_headers | object | Headers to send to the MCP server. The map contains the header name and the path to a file containing the header value (secret). There are 3 special cases: 1. Usage of the kubernetes token in the header — use the string 'kubernetes' instead of the file path. 2. Usage of the client provided token in the header — use the string 'client' instead of the file path. 3. Usage of OAuth token (resolved at request time or 401 with WWW-Authenticate) — use the string 'oauth' instead of the file path. |
+| authorization_headers | object | Headers to send to the MCP server. The map contains the header name and the path to a file containing the header value (secret). There are 3 special cases: 1. Usage of the kubernetes token in the header — use the string 'kubernetes' instead of the file path. 2. Usage of the client-provided token in the header — use the string 'client' instead of the file path. 3. Usage of OAuth token (resolved at request time or 401 with WWW-Authenticate) — use the string 'oauth' instead of the file path. |📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| | authorization_headers | object | Headers to send to the MCP server. The map contains the header name and the path to a file containing the header value (secret). There are 3 special cases: 1. Usage of the kubernetes token in the header — use the string 'kubernetes' instead of the file path. 2. Usage of the client provided token in the header — use the string 'client' instead of the file path. 3. Usage of OAuth token (resolved at request time or 401 with WWW-Authenticate) — use the string 'oauth' instead of the file path. | | |
| | authorization_headers | object | Headers to send to the MCP server. The map contains the header name and the path to a file containing the header value (secret). There are 3 special cases: 1. Usage of the kubernetes token in the header — use the string 'kubernetes' instead of the file path. 2. Usage of the client-provided token in the header — use the string 'client' instead of the file path. 3. Usage of OAuth token (resolved at request time or 401 with WWW-Authenticate) — use the string 'oauth' instead of the file path. | |
🧰 Tools
🪛 LanguageTool
[grammar] ~374-~374: Use a hyphen to join words.
Context: ...of the file path. 2. Usage of the client provided token in the header — use the s...
(QB_NEW_EN_HYPHEN)
🤖 Prompt for AI Agents
In `@docs/config.md` at line 374, In the authorization_headers description update
the phrase "client provided token" to the hyphenated form "client-provided
token" (the clause that explains using the string 'client' instead of a file
path); edit the docs entry for authorization_headers to use "client-provided"
wherever it appears to correct the compound adjective usage.
|
|
||
| import logging | ||
| from typing import Annotated, Any | ||
| import requests |
Contributor
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Synchronous requests import in an async endpoint module.
This module defines async def endpoint handlers. Using the synchronous requests library blocks the event loop. Since the project already depends on aiohttp (per the library context), consider using that instead, or httpx.AsyncClient.
🤖 Prompt for AI Agents
In `@src/app/endpoints/tools.py` at line 5, The module imports the synchronous
requests library (requests) but defines async endpoint handlers, which will
block the event loop; replace the requests import with an async HTTP client
(e.g., aiohttp.ClientSession or httpx.AsyncClient), update any call sites that
use requests.get/post/etc. to use the async client with await (for example
create or accept an aiohttp.ClientSession and call await session.get(...)), and
ensure proper session lifecycle (use async context manager or
dependency-injected session) so all network I/O in your async endpoint functions
remains non-blocking.
Comment on lines
+93
to
+104
| except AuthenticationError as e: | ||
| logger.error("Authentication error: %s", e) | ||
| if toolgroup.mcp_endpoint: | ||
| resp = requests.get(toolgroup.mcp_endpoint.uri, timeout=10) | ||
| cause = f"MCP server at {toolgroup.mcp_endpoint.uri} requires OAuth" | ||
| error_response = UnauthorizedResponse(cause=cause) | ||
| raise HTTPException( | ||
| **error_response.model_dump(), | ||
| headers={"WWW-Authenticate": resp.headers["WWW-Authenticate"]}, | ||
| ) from e | ||
| error_response = UnauthorizedResponse(cause=str(e)) | ||
| raise HTTPException(**error_response.model_dump()) from e |
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Step 1: Check the function context around lines 93-104 in tools.py
echo "=== Function context around lines 93-104 in src/app/endpoints/tools.py ==="
sed -n '70,110p' src/app/endpoints/tools.py | cat -n
echo ""
echo "=== Check function signature ==="
sed -n '70,95p' src/app/endpoints/tools.py | grep -E "^(async )?def " | head -1Repository: lightspeed-core/lightspeed-stack
Length of output: 2401
🏁 Script executed:
#!/bin/bash
# Step 2: Check project dependencies for async HTTP clients
echo "=== Check for async HTTP client dependencies ==="
if [ -f "pyproject.toml" ]; then
grep -E "aiohttp|httpx" pyproject.toml
elif [ -f "requirements.txt" ]; then
grep -E "aiohttp|httpx" requirements.txt
else
find . -name "pyproject.toml" -o -name "requirements.txt" -o -name "setup.py" | head -3
fi
echo ""
echo "=== Check imports in tools.py ==="
head -30 src/app/endpoints/tools.py | cat -nRepository: lightspeed-core/lightspeed-stack
Length of output: 1347
🏁 Script executed:
#!/bin/bash
# Step 3: Search for requests.get usage in src/**/*.py
echo "=== All requests.get() calls in src/ ==="
rg -n "requests\." --type=py src/ | head -20Repository: lightspeed-core/lightspeed-stack
Length of output: 1610
🏁 Script executed:
#!/bin/bash
# Step 4: Check src/utils/responses.py for get_mcp_tools pattern
echo "=== Check get_mcp_tools function in src/utils/responses.py ==="
rg -n -B2 -A8 "def get_mcp_tools" src/utils/responses.pyRepository: lightspeed-core/lightspeed-stack
Length of output: 531
🏁 Script executed:
#!/bin/bash
# Check if get_mcp_tools is async and full context around line 398
echo "=== Check if get_mcp_tools is async ==="
sed -n '318,325p' src/utils/responses.py | cat -n
echo ""
echo "=== Full context around line 398 in responses.py ==="
sed -n '390,410p' src/utils/responses.py | cat -n
echo ""
echo "=== Check which function contains line 398 ==="
sed -n '380,400p' src/utils/responses.py | grep -E "^(async )?def " | tail -1Repository: lightspeed-core/lightspeed-stack
Length of output: 1712
🏁 Script executed:
#!/bin/bash
# Check if there are any await calls or async patterns near line 398
echo "=== Context around requests.get in responses.py with function boundaries ==="
sed -n '1,400p' src/utils/responses.py | tail -100 | head -50Repository: lightspeed-core/lightspeed-stack
Length of output: 1901
🏁 Script executed:
#!/bin/bash
# Check where get_mcp_tools is called from
echo "=== Where is get_mcp_tools called? ==="
rg -n "get_mcp_tools" src/ --type=py
echo ""
echo "=== Check if those call sites are in async functions ==="
# Check the call in responses.py
rg -n -B5 "get_mcp_tools" src/ --type=py | grep -E "(async def|def )" | head -5Repository: lightspeed-core/lightspeed-stack
Length of output: 661
🏁 Script executed:
#!/bin/bash
# Verify the full handler function signature in tools.py to confirm it's async
echo "=== Find the full async endpoint function in tools.py ==="
rg -n "^async def\|^@router" src/app/endpoints/tools.py | head -10
echo ""
echo "=== Get the specific endpoint that contains the blocking code ==="
sed -n '60,95p' src/app/endpoints/tools.py | cat -n | tail -20Repository: lightspeed-core/lightspeed-stack
Length of output: 1126
🏁 Script executed:
#!/bin/bash
# Find the complete function definition that contains the blocking code
echo "=== Find the router decorator and function def in tools.py ==="
sed -n '35,65p' src/app/endpoints/tools.py | cat -n
echo ""
echo "=== Also check for `@router.get` or similar ==="
rg -n "@router" src/app/endpoints/tools.py | head -5Repository: lightspeed-core/lightspeed-stack
Length of output: 1539
🏁 Script executed:
#!/bin/bash
# Check the call site in responses.py line 188 - is that function async?
echo "=== Check function containing line 188 in responses.py ==="
sed -n '170,195p' src/utils/responses.py | cat -n
echo ""
echo "=== Check the rlsapi_v1.py call site line 268 ==="
sed -n '250,275p' src/app/endpoints/rlsapi_v1.py | cat -nRepository: lightspeed-core/lightspeed-stack
Length of output: 2547
Blocking I/O in async handler and unhandled failure modes in the OAuth probe.
Three critical issues in this block:
-
Blocking call on the event loop:
requests.get()is synchronous and will block the entire async event loop for up to 10 seconds. Use an async HTTP client (aiohttp,httpx.AsyncClient) or run it in an executor. (aiohttp is already available in dependencies.) -
KeyErrorifWWW-Authenticateis missing:resp.headers["WWW-Authenticate"](line 101) will raiseKeyErrorif the MCP server doesn't return that header, resulting in an unhandled 500. Use.get()with a fallback. -
No error handling on the probe request:
requests.get()can raiseConnectionError,Timeout, etc. If the MCP server is unreachable, this surfaces as an unhandled 500 instead of a meaningful error.
The same pattern exists in src/utils/responses.py:398 (get_mcp_tools) and should be fixed there too.
Proposed fix (minimal — still uses requests; async migration recommended separately)
except AuthenticationError as e:
logger.error("Authentication error: %s", e)
if toolgroup.mcp_endpoint:
- resp = requests.get(toolgroup.mcp_endpoint.uri, timeout=10)
- cause = f"MCP server at {toolgroup.mcp_endpoint.uri} requires OAuth"
- error_response = UnauthorizedResponse(cause=cause)
- raise HTTPException(
- **error_response.model_dump(),
- headers={"WWW-Authenticate": resp.headers["WWW-Authenticate"]},
- ) from e
+ www_authenticate = ""
+ try:
+ resp = requests.get(toolgroup.mcp_endpoint.uri, timeout=10)
+ www_authenticate = resp.headers.get("WWW-Authenticate", "")
+ except requests.RequestException:
+ logger.warning(
+ "Failed to probe MCP server %s for WWW-Authenticate header",
+ toolgroup.mcp_endpoint.uri,
+ )
+ cause = f"MCP server at {toolgroup.mcp_endpoint.uri} requires OAuth"
+ error_response = UnauthorizedResponse(cause=cause)
+ headers = {"WWW-Authenticate": www_authenticate} if www_authenticate else {}
+ raise HTTPException(
+ **error_response.model_dump(),
+ headers=headers,
+ ) from e
error_response = UnauthorizedResponse(cause=str(e))
raise HTTPException(**error_response.model_dump()) from ePer coding guidelines, "Use async def for I/O operations and external API calls".
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| except AuthenticationError as e: | |
| logger.error("Authentication error: %s", e) | |
| if toolgroup.mcp_endpoint: | |
| resp = requests.get(toolgroup.mcp_endpoint.uri, timeout=10) | |
| cause = f"MCP server at {toolgroup.mcp_endpoint.uri} requires OAuth" | |
| error_response = UnauthorizedResponse(cause=cause) | |
| raise HTTPException( | |
| **error_response.model_dump(), | |
| headers={"WWW-Authenticate": resp.headers["WWW-Authenticate"]}, | |
| ) from e | |
| error_response = UnauthorizedResponse(cause=str(e)) | |
| raise HTTPException(**error_response.model_dump()) from e | |
| except AuthenticationError as e: | |
| logger.error("Authentication error: %s", e) | |
| if toolgroup.mcp_endpoint: | |
| www_authenticate = "" | |
| try: | |
| resp = requests.get(toolgroup.mcp_endpoint.uri, timeout=10) | |
| www_authenticate = resp.headers.get("WWW-Authenticate", "") | |
| except requests.RequestException: | |
| logger.warning( | |
| "Failed to probe MCP server %s for WWW-Authenticate header", | |
| toolgroup.mcp_endpoint.uri, | |
| ) | |
| cause = f"MCP server at {toolgroup.mcp_endpoint.uri} requires OAuth" | |
| error_response = UnauthorizedResponse(cause=cause) | |
| headers = {"WWW-Authenticate": www_authenticate} if www_authenticate else {} | |
| raise HTTPException( | |
| **error_response.model_dump(), | |
| headers=headers, | |
| ) from e | |
| error_response = UnauthorizedResponse(cause=str(e)) | |
| raise HTTPException(**error_response.model_dump()) from e |
🤖 Prompt for AI Agents
In `@src/app/endpoints/tools.py` around lines 93 - 104, The block catching
AuthenticationError is performing a blocking synchronous probe and not guarding
against probe failures or missing headers; update the probe to use an async HTTP
client (e.g., aiohttp) and await the request instead of requests.get, wrap the
probe in try/except to catch connection, timeout, and client errors (e.g.,
aiohttp.ClientError, asyncio.TimeoutError) and fall back to a safe
UnauthorizedResponse when the probe fails, and when constructing the
HTTPException use resp.headers.get("WWW-Authenticate") with a sensible default
(or omit the header if absent) to avoid KeyError; apply the same
async/non-blocking error-handling pattern to the get_mcp_tools function in
src/utils/responses.py.
Comment on lines
512
to
+514
| "To specify this use a string 'client' instead of the file path." | ||
| "3. Usage of the oauth token in the header. " | ||
| "To specify this use a string 'oauth' instead of the file path. " |
Contributor
There was a problem hiding this comment.
Missing space between special cases 2 and 3 in the description string.
Line 512 ends without a trailing space, so the concatenated text reads "...file path.3. Usage of the oauth...". Add a trailing space (or leading space on Line 513) to match the formatting of case 1→2.
Proposed fix
- "To specify this use a string 'client' instead of the file path."
- "3. Usage of the oauth token in the header. "
+ "To specify this use a string 'client' instead of the file path. "
+ "3. Usage of the oauth token in the header. "📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| "To specify this use a string 'client' instead of the file path." | |
| "3. Usage of the oauth token in the header. " | |
| "To specify this use a string 'oauth' instead of the file path. " | |
| "To specify this use a string 'client' instead of the file path. " | |
| "3. Usage of the oauth token in the header. " | |
| "To specify this use a string 'oauth' instead of the file path. " |
🤖 Prompt for AI Agents
In `@src/models/config.py` around lines 512 - 514, The concatenated description
string for the special cases is missing a space between the end of the "client"
sentence and the start of "3. Usage of the oauth..." sentence; update the string
fragments (the literal containing "To specify this use a string 'client' instead
of the file path." or the immediately following fragment starting "3. Usage of
the oauth token in the header.") so there is a space either appended to the
first fragment or prefixed to the second fragment to ensure the final
description reads "...file path. 3. Usage of the oauth..." (refer to the config
description string that includes the "'client'" and "'oauth'" special-case
text).
Comment on lines
+389
to
+405
| # If OAuth was required and no headers passed, probe endpoint and forward | ||
| # 401 with WWW-Authenticate so the client can perform OAuth | ||
| uses_oauth = ( | ||
| constants.MCP_AUTH_OAUTH | ||
| in mcp_server.resolved_authorization_headers.values() | ||
| ) | ||
| if uses_oauth and ( | ||
| mcp_headers is None or not mcp_headers.get(mcp_server.name) | ||
| ): | ||
| resp = requests.get(mcp_server.url, timeout=10) | ||
| error_response = UnauthorizedResponse( | ||
| cause=f"MCP server at {mcp_server.url} requires OAuth authentication", | ||
| ) | ||
| raise HTTPException( | ||
| **error_response.model_dump(), | ||
| headers={"WWW-Authenticate": resp.headers["WWW-Authenticate"]}, | ||
| ) |
Contributor
There was a problem hiding this comment.
Missing error handling: requests.get() can fail, and WWW-Authenticate header may be absent.
Several unhandled failure modes:
requests.get(mcp_server.url, timeout=10)(Line 398) — can raiserequests.ConnectionError,requests.Timeout, etc. if the MCP server is unreachable. This will surface as an unhandled 500.resp.headers["WWW-Authenticate"](Line 404) — raisesKeyErrorif the MCP server doesn't include this header (e.g., returns 200, or a 401 without it).- Blocking I/O in async context —
get_mcp_toolsis called from the asyncprepare_tools(line 188). The synchronousrequests.get()blocks the event loop. Consider usinghttpx.AsyncClientor running in an executor. The same issue exists insrc/app/endpoints/tools.pyline ~98.
At minimum, wrap the probe in a try/except and use .get() for the header:
Proposed fix for error handling
if uses_oauth and (
mcp_headers is None or not mcp_headers.get(mcp_server.name)
):
- resp = requests.get(mcp_server.url, timeout=10)
- error_response = UnauthorizedResponse(
- cause=f"MCP server at {mcp_server.url} requires OAuth authentication",
- )
- raise HTTPException(
- **error_response.model_dump(),
- headers={"WWW-Authenticate": resp.headers["WWW-Authenticate"]},
- )
+ error_response = UnauthorizedResponse(
+ cause=f"MCP server at {mcp_server.url} requires OAuth authentication",
+ )
+ headers: dict[str, str] = {}
+ try:
+ resp = requests.get(mcp_server.url, timeout=10)
+ www_auth = resp.headers.get("WWW-Authenticate")
+ if www_auth:
+ headers["WWW-Authenticate"] = www_auth
+ except requests.RequestException:
+ logger.warning(
+ "Failed to probe MCP server %s for WWW-Authenticate header",
+ mcp_server.url,
+ )
+ raise HTTPException(
+ **error_response.model_dump(),
+ headers=headers if headers else None,
+ )As per coding guidelines: Use async def for I/O operations and external API calls. The synchronous requests.get call in this code path blocks the event loop. Consider using an async HTTP client for this probe request, or at least running it in an executor.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| # If OAuth was required and no headers passed, probe endpoint and forward | |
| # 401 with WWW-Authenticate so the client can perform OAuth | |
| uses_oauth = ( | |
| constants.MCP_AUTH_OAUTH | |
| in mcp_server.resolved_authorization_headers.values() | |
| ) | |
| if uses_oauth and ( | |
| mcp_headers is None or not mcp_headers.get(mcp_server.name) | |
| ): | |
| resp = requests.get(mcp_server.url, timeout=10) | |
| error_response = UnauthorizedResponse( | |
| cause=f"MCP server at {mcp_server.url} requires OAuth authentication", | |
| ) | |
| raise HTTPException( | |
| **error_response.model_dump(), | |
| headers={"WWW-Authenticate": resp.headers["WWW-Authenticate"]}, | |
| ) | |
| # If OAuth was required and no headers passed, probe endpoint and forward | |
| # 401 with WWW-Authenticate so the client can perform OAuth | |
| uses_oauth = ( | |
| constants.MCP_AUTH_OAUTH | |
| in mcp_server.resolved_authorization_headers.values() | |
| ) | |
| if uses_oauth and ( | |
| mcp_headers is None or not mcp_headers.get(mcp_server.name) | |
| ): | |
| error_response = UnauthorizedResponse( | |
| cause=f"MCP server at {mcp_server.url} requires OAuth authentication", | |
| ) | |
| headers: dict[str, str] = {} | |
| try: | |
| resp = requests.get(mcp_server.url, timeout=10) | |
| www_auth = resp.headers.get("WWW-Authenticate") | |
| if www_auth: | |
| headers["WWW-Authenticate"] = www_auth | |
| except requests.RequestException: | |
| logger.warning( | |
| "Failed to probe MCP server %s for WWW-Authenticate header", | |
| mcp_server.url, | |
| ) | |
| raise HTTPException( | |
| **error_response.model_dump(), | |
| headers=headers if headers else None, | |
| ) |
🤖 Prompt for AI Agents
In `@src/utils/responses.py` around lines 389 - 405, The probe using requests.get
in the get_mcp_tools logic blocks the async event loop and lacks error handling
for connection/timeouts and missing WWW-Authenticate header; replace the
synchronous probe with an async HTTP call (e.g., httpx.AsyncClient) or run
requests.get in an executor, wrap the call in try/except to catch request
exceptions (ConnectionError, Timeout, HTTPError) and handle them gracefully, and
when building the HTTPException use resp.headers.get("WWW-Authenticate") (only
include the header if present) so you don't KeyError; update the code paths in
get_mcp_tools and callers (prepare_tools) accordingly to await the async probe
or handle executor futures.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
MCP servers can use the "oauth" keyword in authorization_headers. If no token is provided, the service returns 401 with a WWW-Authenticate header (from the MCP server) so clients can run OAuth and retry with a token in MCP-HEADERS.
Type of change
Tools used to create PR
Identify any AI code assistants used in this PR (for transparency and review context)
Related Tickets & Documents
Checklist before requesting a review
Testing
Summary by CodeRabbit
Release Notes
New Features
Documentation