Authentication and Security

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, Virtual Account token, or IdP JWT
    tool_name: str,
    tool_args: dict
):
    """Call an MCP tool, handling OAuth authorization if needed."""
    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


# Example usage
async def main():
    mcp_url = "https://<control-plane-url>/api/llm/mcp/<integration-id>/server"
    user_token = "user-tfy-token-or-idp-jwt"
    
    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
    else:
        print(f"Success: {result['result']}")

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

# Store Virtual Account token securely (e.g., environment variable)
VA_TOKEN = os.environ["VIRTUAL_ACCOUNT_TOKEN"]

async def call_mcp_tool(tool_name: str, tool_args: dict):
    """Call an MCP tool using Virtual Account credentials."""
    mcp_url = "https://<control-plane-url>/api/llm/mcp/<integration-id>/server"
    
    transport = StreamableHttpTransport(
        url=mcp_url,
        headers={"Authorization": f"Bearer {VA_TOKEN}"}
    )
    
    async with Client(transport) as client:
        result = await client.call_tool(tool_name, tool_args)
        return result

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

async def call_passthrough_mcp(
    mcp_url: str,
    user_jwt: str,  # JWT from your IdP
    tool_name: str,
    tool_args: dict
):
    """Call an MCP tool with token passthrough."""
    transport = StreamableHttpTransport(
        url=mcp_url,
        headers={"Authorization": f"Bearer {user_jwt}"}
    )
    
    async with Client(transport) as client:
        result = await client.call_tool(tool_name, tool_args)
        return result

Does the MCP server need per-user data from third-party services?
├── Yes → Use OAuth2 (Per-User Credentials)
│         └── Are users TrueFoundry platform users?
│             ├── Yes → Use TrueFoundry PAT
│             └── No → Set up External Identity
│
└── No → Does the MCP server validate tokens from your IdP?
         ├── Yes → Use Token Passthrough
         └── No → Does it require any auth?
                  ├── Yes → Use Static Header Auth (Shared Credentials)
                  └── No → Use No Auth