On-Behalf-Of Authentication

import re
from fastmcp import Client
from fastmcp.client.transports import StreamableHttpTransport

def get_auth_error_data(exception: Exception) -> str | None:
    """Traverse exception chain to find auth error with URL."""
    current = exception
    while current is not None:
        try:
            data = current.error.data
            if data and "Please visit:" in str(data):
                return str(data)
        except AttributeError:
            pass
        current = current.__cause__
    return None


def extract_auth_urls(error_data: str) -> list[str]:
    """Extract authorization URLs from MCP Gateway error."""
    match = re.search(r"Please visit:\s*(.+)$", error_data)
    if match:
        # URLs are separated by " , " (space-comma-space)
        return [url.strip() for url in match.group(1).split(" , ")]
    return []


async def call_mcp_with_oauth(
    mcp_url: str,
    auth_token: str,  # User's TrueFoundry PAT or External Identity JWT
    tool_name: str,
    tool_args: dict
):
    transport = StreamableHttpTransport(
        url=mcp_url,
        headers={"Authorization": f"Bearer {auth_token}"}
    )
    
    try:
        # Note: OAuth errors occur during connection, not tool calls
        async with Client(transport) as client:
            result = await client.call_tool(tool_name, tool_args)
            return {"success": True, "result": result}
            
    except Exception as e:
        error_data = get_auth_error_data(e)
        if error_data:
            auth_urls = extract_auth_urls(error_data)
            if auth_urls:
                return {
                    "success": False,
                    "auth_required": True,
                    "authorization_urls": auth_urls
                }
        raise


# Usage
result = await call_mcp_with_oauth(mcp_url, user_token, "get_repositories", {})

if result.get("auth_required"):
    # In a web app: redirect user to the authorization URL
    # In a CLI: print the URL for user to visit
    for url in result["authorization_urls"]:
        print(f"Please authorize access: {url}")
    
    # After user completes OAuth, retry the request
    # The MCP Gateway stores the OAuth token and future requests will succeed