Extract OAuth flow logic into reusable components for proxy use cases

Summary

Refactor OAuth implementation so the flow logic and state machine are usable by server-side proxy services, not just client-side browser flows.

Problem

The SDK's OAuth implementation is designed for local client-side flows (opening a browser locally). The business logic is embedded inside an httpx auth module, making it hard to reuse for other scenarios.

While individual helper functions have been extracted (PKCE utilities, token exchange, discovery), the core state machine that orchestrates the OAuth flow is not reusable. Proxy services that need to perform OAuth on behalf of users currently have to reimplement significant portions of the flow themselves — and when the SDK updates its OAuth logic, those reimplementations can fall out of sync.

Goal

Make the OAuth portions of the SDK compatible with proxy/gateway services that currently use custom workarounds
When an issue is fixed in the SDK, updating the SDK version should fix it everywhere — no custom OAuth reimplementations needed
Keep existing client-side flows working

Design Requirements (from maintainer discussion, Feb 2026)

Modularization into zones: Break the monolithic OAuth flow into modular, overridable pieces:

Discovery — obtaining and potentially customizing discovery URLs
Client Registration — dynamic client registration
Interactive Flow — authorization URL generation, redirect handling
Token Fetching — code exchange, refresh, new token extensions (XA, WIF)
Token Storage — pluggable storage (already exists)

Key requirements:

Each zone should operate as a pure function requiring minimal state
Every HTTP request in the flow must be interceptable — allow injection of a custom HTTP client/fetch interface (httpx client in Python, fetch in TypeScript) for custom headers, metrics, response handling
Support an "Auth Required" state as an SDK primitive — when a server responds with 401/403 mid-flow, the SDK should capture discovery metadata, scope, and WWW-Authenticate info and surface it so the calling application can handle it (rather than assuming auth happens upfront)
The flow must be resumable — a caller should be able to pick up an auth flow at any point (e.g., after a redirect returns on a different machine/request)
Support bypassing discovery when configuration is provided directly (important for enterprise environments with broken discovery)
Support new token-getting extensions (XA, WIF) that don't require interactive flows

Next steps:

Draft code sketches (potentially TypeScript first) to validate the modular function approach
Cross-SDK coordination — this applies to both Python and TypeScript SDKs

"""
OAuth V2 Sketch - Action/Effect Pattern with Strategy Configuration

This separates:
- WHAT to do (Actions - pure data)
- HOW to execute (Adapters - httpx, fetch, etc.)
- WHEN to transition (Flow - pure state machine)
- CUSTOMIZATION (Config - strategies)
"""

from __future__ import annotations

import secrets
from dataclasses import dataclass, field
from enum import Enum, auto
from typing import Callable, Protocol
from urllib.parse import urlencode, urljoin, urlparse

import httpx
from pydantic import AnyUrl

# =============================================================================
# MODELS (existing, mostly unchanged)
# =============================================================================


@dataclass
class OAuthToken:
    access_token: str
    token_type: str = "Bearer"
    refresh_token: str | None = None
    expires_in: int | None = None
    scope: str | None = None


@dataclass
class ProtectedResourceMetadata:
    resource: str
    authorization_servers: list[str]
    scopes_supported: list[str] | None = None


@dataclass
class OAuthMetadata:
    issuer: str
    authorization_endpoint: str | None = None
    token_endpoint: str | None = None
    registration_endpoint: str | None = None
    scopes_supported: list[str] | None = None
    token_endpoint_auth_methods_supported: list[str] | None = None
    code_challenge_methods_supported: list[str] | None = None
    client_id_metadata_document_supported: bool | None = None


@dataclass
class OAuthClientMetadata:
    redirect_uris: list[str]
    client_name: str | None = None
    scope: str | None = None
    token_endpoint_auth_method: str | None = None


@dataclass
class OAuthClientInfo:
    client_id: str
    client_secret: str | None = None
    token_endpoint_auth_method: str = "none"
    redirect_uris: list[str] | None = None


# =============================================================================
# ACTIONS - Pure data describing what HTTP call to make
# =============================================================================


@dataclass
class HttpRequest:
    """Generic HTTP request descriptor."""

    method: str
    url: str
    headers: dict[str, str] = field(default_factory=dict)
    body: bytes | None = None
    json: dict | None = None
    form: dict[str, str] | None = None


@dataclass
class DiscoverPRM:
    """Discover Protected Resource Metadata."""

    urls: list[str]
    current_index: int = 0

    @property
    def current_url(self) -> str | None:
        if self.current_index < len(self.urls):
            return self.urls[self.current_index]
        return None

    def to_request(self) -> HttpRequest | None:
        if url := self.current_url:
            return HttpRequest(method="GET", url=url)
        return None


@dataclass
class DiscoverOASM:
    """Discover OAuth Authorization Server Metadata."""

    urls: list[str]
    current_index: int = 0

    @property
    def current_url(self) -> str | None:
        if self.current_index < len(self.urls):
            return self.urls[self.current_index]
        return None

    def to_request(self) -> HttpRequest | None:
        if url := self.current_url:
            return HttpRequest(method="GET", url=url)
        return None


@dataclass
class RegisterClient:
    """Dynamic Client Registration."""

    endpoint: str
    client_metadata: dict

    def to_request(self) -> HttpRequest:
        return HttpRequest(
            method="POST",
            url=self.endpoint,
            headers={"Content-Type": "application/json"},
            json=self.client_metadata,
        )


@dataclass
class AuthorizeUser:
    """User authorization - NOT an HTTP request, but a user interaction."""

    authorization_url: str
    state: str
    code_verifier: str  # Keep for token exchange


@dataclass
class ExchangeToken:
    """Exchange authorization code for token."""

    endpoint: str
    code: str
    code_verifier: str
    redirect_uri: str
    client_id: str
    client_secret: str | None = None
    auth_method: str = "none"
    resource: str | None = None

    def to_request(self) -> HttpRequest:
        data = {
            "grant_type": "authorization_code",
            "code": self.code,
            "redirect_uri": self.redirect_uri,
            "client_id": self.client_id,
            "code_verifier": self.code_verifier,
        }
        if self.resource:
            data["resource"] = self.resource

        headers = {"Content-Type": "application/x-www-form-urlencoded"}

        # Handle client authentication
        if self.auth_method == "client_secret_post" and self.client_secret:
            data["client_secret"] = self.client_secret
        elif self.auth_method == "client_secret_basic" and self.client_secret:
            import base64

            credentials = base64.b64encode(f"{self.client_id}:{self.client_secret}".encode()).decode()
            headers["Authorization"] = f"Basic {credentials}"

        return HttpRequest(method="POST", url=self.endpoint, headers=headers, form=data)


@dataclass
class RefreshToken:
    """Refresh an access token."""

    endpoint: str
    refresh_token: str
    client_id: str
    client_secret: str | None = None
    auth_method: str = "none"
    resource: str | None = None

    def to_request(self) -> HttpRequest:
        data = {
            "grant_type": "refresh_token",
            "refresh_token": self.refresh_token,
            "client_id": self.client_id,
        }
        if self.resource:
            data["resource"] = self.resource

        headers = {"Content-Type": "application/x-www-form-urlencoded"}

        if self.auth_method == "client_secret_post" and self.client_secret:
            data["client_secret"] = self.client_secret
        elif self.auth_method == "client_secret_basic" and self.client_secret:
            import base64

            credentials = base64.b64encode(f"{self.client_id}:{self.client_secret}".encode()).decode()
            headers["Authorization"] = f"Basic {credentials}"

        return HttpRequest(method="POST", url=self.endpoint, headers=headers, form=data)


# Union of all actions
OAuthAction = DiscoverPRM | DiscoverOASM | RegisterClient | AuthorizeUser | ExchangeToken | RefreshToken


@dataclass
class OAuthComplete:
    """Flow completed successfully."""

    token: OAuthToken


@dataclass
class OAuthFailed:
    """Flow failed."""

    error: str
    details: str | None = None


OAuthResult = OAuthAction | OAuthComplete | OAuthFailed


# =============================================================================
# ACTION RESULTS - What comes back from executing an action
# =============================================================================


@dataclass
class HttpResponse:
    """Generic HTTP response."""

    status_code: int
    body: bytes
    headers: dict[str, str] = field(default_factory=dict)

    def json(self) -> dict:
        import json

        return json.loads(self.body)


@dataclass
class PRMResult:
    """Result of PRM discovery attempt."""

    found: bool
    metadata: ProtectedResourceMetadata | None = None


@dataclass
class OASMResult:
    """Result of OASM discovery attempt."""

    found: bool
    metadata: OAuthMetadata | None = None
    should_stop: bool = False  # True if got non-4xx error


@dataclass
class RegistrationResult:
    """Result of client registration."""

    success: bool
    client_info: OAuthClientInfo | None = None
    error: str | None = None


@dataclass
class AuthorizationResult:
    """Result of user authorization (callback)."""

    code: str
    state: str


@dataclass
class TokenResult:
    """Result of token exchange or refresh."""

    success: bool
    token: OAuthToken | None = None
    error: str | None = None


# =============================================================================
# CONFIGURATION - Strategies for customization
# =============================================================================


def default_prm_urls(server_url: str, www_auth_url: str | None) -> list[str]:
    """Default PRM discovery URL builder (SEP-985)."""
    urls: list[str] = []

    # Priority 1: WWW-Authenticate header
    if www_auth_url:
        urls.append(www_auth_url)

    # Priority 2-3: Well-known URIs
    parsed = urlparse(server_url)
    base_url = f"{parsed.scheme}://{parsed.netloc}"

    # Path-based
    if parsed.path and parsed.path != "/":
        urls.append(urljoin(base_url, f"/.well-known/oauth-protected-resource{parsed.path}"))

    # Root-based
    urls.append(urljoin(base_url, "/.well-known/oauth-protected-resource"))

    return urls


def default_oasm_urls(auth_server_url: str | None, server_url: str) -> list[str]:
    """Default OASM discovery URL builder (RFC 8414 + OIDC)."""
    # No PRM found - legacy March 2025 spec fallback
    if not auth_server_url:
        parsed = urlparse(server_url)
        return [f"{parsed.scheme}://{parsed.netloc}/.well-known/oauth-authorization-server"]

    urls: list[str] = []
    parsed = urlparse(auth_server_url)
    base_url = f"{parsed.scheme}://{parsed.netloc}"

    if parsed.path and parsed.path != "/":
        path = parsed.path.rstrip("/")
        # RFC 8414 path-aware
        urls.append(urljoin(base_url, f"/.well-known/oauth-authorization-server{path}"))
        urls.append(urljoin(base_url, f"/.well-known/openid-configuration{path}"))
        # OIDC 1.0 appended
        urls.append(urljoin(base_url, f"{path}/.well-known/openid-configuration"))
    else:
        urls.append(urljoin(base_url, "/.well-known/oauth-authorization-server"))
        urls.append(urljoin(base_url, "/.well-known/openid-configuration"))

    return urls


def default_base_url(server_url: str) -> str:
    """Extract base URL (scheme + netloc) from server URL."""
    parsed = urlparse(server_url)
    return f"{parsed.scheme}://{parsed.netloc}"


def default_scope_selector(
    www_auth_scope: str | None,
    prm: ProtectedResourceMetadata | None,
    oasm: OAuthMetadata | None,
) -> str | None:
    """Default scope selection strategy."""
    if www_auth_scope:
        return www_auth_scope
    if prm and prm.scopes_supported:
        return " ".join(prm.scopes_supported)
    if oasm and oasm.scopes_supported:
        return " ".join(oasm.scopes_supported)
    return None


@dataclass
class OAuthFlowConfig:
    """Configuration for OAuth flow behavior."""

    # Server and client info
    server_url: str
    client_metadata: OAuthClientMetadata
    client_metadata_url: str | None = None  # For CIMD

    # Customizable strategies
    prm_url_builder: Callable[[str, str | None], list[str]] = default_prm_urls
    oasm_url_builder: Callable[[str | None, str], list[str]] = default_oasm_urls
    base_url_extractor: Callable[[str], str] = default_base_url
    scope_selector: Callable[
        [str | None, ProtectedResourceMetadata | None, OAuthMetadata | None],
        str | None,
    ] = default_scope_selector

    # Feature flags
    include_resource_param: bool = True  # RFC 8707


# =============================================================================
# FLOW STATE - Internal state of the OAuth flow
# =============================================================================


class FlowPhase(Enum):
    """Current phase of the OAuth flow."""

    INITIAL = auto()
    DISCOVERING_PRM = auto()
    DISCOVERING_OASM = auto()
    REGISTERING = auto()
    AWAITING_USER_AUTH = auto()
    EXCHANGING_TOKEN = auto()
    REFRESHING = auto()
    COMPLETE = auto()
    FAILED = auto()


@dataclass
class FlowState:
    """Internal state of the OAuth flow."""

    phase: FlowPhase = FlowPhase.INITIAL

    # Discovery results
    www_auth_scope: str | None = None
    www_auth_resource_metadata_url: str | None = None
    prm: ProtectedResourceMetadata | None = None
    oasm: OAuthMetadata | None = None
    auth_server_url: str | None = None

    # Client info
    client_info: OAuthClientInfo | None = None

    # Authorization state
    pending_auth: AuthorizeUser | None = None

    # Token
    token: OAuthToken | None = None

    # Error
    error: str | None = None


# =============================================================================
# OAUTH FLOW - Pure state machine, no I/O
# =============================================================================


class OAuthFlow:
    """
    Pure state machine for OAuth flow.

    No I/O - just returns actions and processes results.
    Adapters (httpx, fetch, etc.) execute the actions.
    """

    def __init__(self, config: OAuthFlowConfig, initial_client_info: OAuthClientInfo | None = None):
        self.config = config
        self.state = FlowState(client_info=initial_client_info)
        self._current_action: OAuthAction | None = None

    @property
    def phase(self) -> FlowPhase:
        return self.state.phase

    @property
    def token(self) -> OAuthToken | None:
        return self.state.token

    @property
    def is_complete(self) -> bool:
        return self.state.phase == FlowPhase.COMPLETE

    @property
    def is_failed(self) -> bool:
        return self.state.phase == FlowPhase.FAILED

    # -------------------------------------------------------------------------
    # Public API
    # -------------------------------------------------------------------------

    def start(
        self,
        www_auth_scope: str | None = None,
        www_auth_resource_metadata_url: str | None = None,
    ) -> OAuthResult:
        """
        Start the OAuth flow (typically after receiving 401).

        Args:
            www_auth_scope: Scope from WWW-Authenticate header
            www_auth_resource_metadata_url: resource_metadata from WWW-Authenticate

        Returns:
            First action to execute, or OAuthComplete if already authenticated
        """
        self.state.www_auth_scope = www_auth_scope
        self.state.www_auth_resource_metadata_url = www_auth_resource_metadata_url

        # Start with PRM discovery
        self.state.phase = FlowPhase.DISCOVERING_PRM
        urls = self.config.prm_url_builder(self.config.server_url, www_auth_resource_metadata_url)

        action = DiscoverPRM(urls=urls)
        self._current_action = action
        return action

    def continue_with_prm_result(self, result: PRMResult) -> OAuthResult:
        """Process PRM discovery result."""
        assert self.state.phase == FlowPhase.DISCOVERING_PRM
        assert isinstance(self._current_action, DiscoverPRM)

        if result.found and result.metadata:
            # Success - store PRM and move to OASM discovery
            self.state.prm = result.metadata
            self.state.auth_server_url = result.metadata.authorization_servers[0]
            return self._start_oasm_discovery()
        else:
            # Try next URL
            action = DiscoverPRM(
                urls=self._current_action.urls,
                current_index=self._current_action.current_index + 1,
            )
            if action.current_url is None:
                # All URLs exhausted - continue without PRM (March 2025 fallback)
                return self._start_oasm_discovery()
            self._current_action = action
            return action

    def continue_with_oasm_result(self, result: OASMResult) -> OAuthResult:
        """Process OASM discovery result."""
        assert self.state.phase == FlowPhase.DISCOVERING_OASM
        assert isinstance(self._current_action, DiscoverOASM)

        if result.should_stop:
            # Non-4xx error, stop trying
            return self._start_registration()

        if result.found and result.metadata:
            # Success - store OASM and move to registration
            self.state.oasm = result.metadata
            return self._start_registration()
        else:
            # Try next URL
            action = DiscoverOASM(
                urls=self._current_action.urls,
                current_index=self._current_action.current_index + 1,
            )
            if action.current_url is None:
                # All URLs exhausted - continue without OASM (use defaults)
                return self._start_registration()
            self._current_action = action
            return action

    def continue_with_registration_result(self, result: RegistrationResult) -> OAuthResult:
        """Process client registration result."""
        assert self.state.phase == FlowPhase.REGISTERING

        if not result.success:
            self.state.phase = FlowPhase.FAILED
            self.state.error = result.error or "Registration failed"
            return OAuthFailed(error=self.state.error)

        self.state.client_info = result.client_info
        return self._start_authorization()

    def continue_with_authorization_result(self, result: AuthorizationResult) -> OAuthResult:
        """Process user authorization callback result."""
        assert self.state.phase == FlowPhase.AWAITING_USER_AUTH
        assert self.state.pending_auth is not None

        # Validate state
        if result.state != self.state.pending_auth.state:
            self.state.phase = FlowPhase.FAILED
            self.state.error = "State mismatch"
            return OAuthFailed(error="State mismatch")

        # Build token exchange action
        self.state.phase = FlowPhase.EXCHANGING_TOKEN
        return self._build_token_exchange(result.code)

    def continue_with_token_result(self, result: TokenResult) -> OAuthResult:
        """Process token exchange/refresh result."""
        assert self.state.phase in (FlowPhase.EXCHANGING_TOKEN, FlowPhase.REFRESHING)

        if not result.success:
            self.state.phase = FlowPhase.FAILED
            self.state.error = result.error or "Token exchange failed"
            return OAuthFailed(error=self.state.error)

        self.state.token = result.token
        self.state.phase = FlowPhase.COMPLETE
        return OAuthComplete(token=result.token)  # type: ignore

    # -------------------------------------------------------------------------
    # Refresh flow (separate entry point)
    # -------------------------------------------------------------------------

    def start_refresh(self, refresh_token: str) -> OAuthResult:
        """Start a token refresh flow."""
        if not self.state.client_info:
            return OAuthFailed(error="No client info for refresh")

        self.state.phase = FlowPhase.REFRESHING

        endpoint = self._get_token_endpoint()
        resource = self._get_resource_url() if self.config.include_resource_param else None

        action = RefreshToken(
            endpoint=endpoint,
            refresh_token=refresh_token,
            client_id=self.state.client_info.client_id,
            client_secret=self.state.client_info.client_secret,
            auth_method=self.state.client_info.token_endpoint_auth_method,
            resource=resource,
        )
        self._current_action = action
        return action

    # -------------------------------------------------------------------------
    # Internal helpers
    # -------------------------------------------------------------------------

    def _start_oasm_discovery(self) -> OAuthResult:
        """Transition to OASM discovery phase."""
        self.state.phase = FlowPhase.DISCOVERING_OASM
        urls = self.config.oasm_url_builder(self.state.auth_server_url, self.config.server_url)

        action = DiscoverOASM(urls=urls)
        self._current_action = action
        return action

    def _start_registration(self) -> OAuthResult:
        """Transition to registration phase."""
        # Check if we already have client info
        if self.state.client_info:
            return self._start_authorization()

        # Check for CIMD (URL-based client ID)
        if self._should_use_cimd():
            self.state.client_info = OAuthClientInfo(
                client_id=self.config.client_metadata_url,  # type: ignore
                token_endpoint_auth_method="none",
                redirect_uris=self.config.client_metadata.redirect_uris,
            )
            return self._start_authorization()

        # Dynamic Client Registration
        self.state.phase = FlowPhase.REGISTERING
        endpoint = self._get_registration_endpoint()

        action = RegisterClient(
            endpoint=endpoint,
            client_metadata={
                "redirect_uris": self.config.client_metadata.redirect_uris,
                "client_name": self.config.client_metadata.client_name,
                "token_endpoint_auth_method": self.config.client_metadata.token_endpoint_auth_method,
            },
        )
        self._current_action = action
        return action

    def _should_use_cimd(self) -> bool:
        """Check if we should use Client ID Metadata Document."""
        if not self.config.client_metadata_url:
            return False
        if not self.state.oasm:
            return False
        return self.state.oasm.client_id_metadata_document_supported is True

    def _start_authorization(self) -> OAuthResult:
        """Transition to authorization phase."""
        assert self.state.client_info is not None

        self.state.phase = FlowPhase.AWAITING_USER_AUTH

        # Select scopes
        scope = self.config.scope_selector(
            self.state.www_auth_scope,
            self.state.prm,
            self.state.oasm,
        )

        # Generate PKCE
        code_verifier = secrets.token_urlsafe(96)
        import base64
        import hashlib

        code_challenge = base64.urlsafe_b64encode(hashlib.sha256(code_verifier.encode()).digest()).decode().rstrip("=")

        # Build authorization URL
        auth_endpoint = self._get_authorization_endpoint()
        state = secrets.token_urlsafe(32)

        params = {
            "response_type": "code",
            "client_id": self.state.client_info.client_id,
            "redirect_uri": self.config.client_metadata.redirect_uris[0],
            "state": state,
            "code_challenge": code_challenge,
            "code_challenge_method": "S256",
        }
        if scope:
            params["scope"] = scope
        if self.config.include_resource_param:
            resource = self._get_resource_url()
            if resource:
                params["resource"] = resource

        authorization_url = f"{auth_endpoint}?{urlencode(params)}"

        action = AuthorizeUser(
            authorization_url=authorization_url,
            state=state,
            code_verifier=code_verifier,
        )
        self.state.pending_auth = action
        self._current_action = action
        return action

    def _build_token_exchange(self, code: str) -> ExchangeToken:
        """Build token exchange action."""
        assert self.state.client_info is not None
        assert self.state.pending_auth is not None

        endpoint = self._get_token_endpoint()
        resource = self._get_resource_url() if self.config.include_resource_param else None

        action = ExchangeToken(
            endpoint=endpoint,
            code=code,
            code_verifier=self.state.pending_auth.code_verifier,
            redirect_uri=self.config.client_metadata.redirect_uris[0],
            client_id=self.state.client_info.client_id,
            client_secret=self.state.client_info.client_secret,
            auth_method=self.state.client_info.token_endpoint_auth_method,
            resource=resource,
        )
        self._current_action = action
        return action

    def _get_authorization_endpoint(self) -> str:
        """Get authorization endpoint with fallback."""
        if self.state.oasm and self.state.oasm.authorization_endpoint:
            return self.state.oasm.authorization_endpoint
        base = self.config.base_url_extractor(self.config.server_url)
        return urljoin(base, "/authorize")

    def _get_token_endpoint(self) -> str:
        """Get token endpoint with fallback."""
        if self.state.oasm and self.state.oasm.token_endpoint:
            return self.state.oasm.token_endpoint
        base = self.config.base_url_extractor(self.config.server_url)
        return urljoin(base, "/token")

    def _get_registration_endpoint(self) -> str:
        """Get registration endpoint with fallback."""
        if self.state.oasm and self.state.oasm.registration_endpoint:
            return self.state.oasm.registration_endpoint
        base = self.config.base_url_extractor(self.config.server_url)
        return urljoin(base, "/register")

    def _get_resource_url(self) -> str | None:
        """Get resource URL for RFC 8707."""
        if self.state.prm:
            return self.state.prm.resource
        # Canonical form of server URL
        parsed = urlparse(self.config.server_url)
        return f"{parsed.scheme}://{parsed.netloc}{parsed.path}".rstrip("/") or None


# =============================================================================
# RESULT PARSERS - Parse HTTP responses into result types
# =============================================================================


def parse_prm_response(response: HttpResponse) -> PRMResult:
    """Parse PRM discovery response."""
    if response.status_code == 200:
        try:
            data = response.json()
            prm = ProtectedResourceMetadata(
                resource=data["resource"],
                authorization_servers=data["authorization_servers"],
                scopes_supported=data.get("scopes_supported"),
            )
            return PRMResult(found=True, metadata=prm)
        except (KeyError, ValueError):
            return PRMResult(found=False)
    return PRMResult(found=False)


def parse_oasm_response(response: HttpResponse) -> OASMResult:
    """Parse OASM discovery response."""
    if response.status_code == 200:
        try:
            data = response.json()
            oasm = OAuthMetadata(
                issuer=data.get("issuer", ""),
                authorization_endpoint=data.get("authorization_endpoint"),
                token_endpoint=data.get("token_endpoint"),
                registration_endpoint=data.get("registration_endpoint"),
                scopes_supported=data.get("scopes_supported"),
                token_endpoint_auth_methods_supported=data.get("token_endpoint_auth_methods_supported"),
                code_challenge_methods_supported=data.get("code_challenge_methods_supported"),
                client_id_metadata_document_supported=data.get("client_id_metadata_document_supported"),
            )
            return OASMResult(found=True, metadata=oasm)
        except (KeyError, ValueError):
            return OASMResult(found=False)
    elif response.status_code >= 400 and response.status_code < 500:
        return OASMResult(found=False)  # Keep trying
    else:
        return OASMResult(found=False, should_stop=True)  # Non-4xx, stop


def parse_registration_response(response: HttpResponse) -> RegistrationResult:
    """Parse registration response."""
    if response.status_code in (200, 201):
        try:
            data = response.json()
            client_info = OAuthClientInfo(
                client_id=data["client_id"],
                client_secret=data.get("client_secret"),
                token_endpoint_auth_method=data.get("token_endpoint_auth_method", "none"),
                redirect_uris=data.get("redirect_uris"),
            )
            return RegistrationResult(success=True, client_info=client_info)
        except (KeyError, ValueError) as e:
            return RegistrationResult(success=False, error=str(e))
    return RegistrationResult(success=False, error=f"HTTP {response.status_code}")


def parse_token_response(response: HttpResponse) -> TokenResult:
    """Parse token response."""
    if response.status_code == 200:
        try:
            data = response.json()
            token = OAuthToken(
                access_token=data["access_token"],
                token_type=data.get("token_type", "Bearer"),
                refresh_token=data.get("refresh_token"),
                expires_in=data.get("expires_in"),
                scope=data.get("scope"),
            )
            return TokenResult(success=True, token=token)
        except (KeyError, ValueError) as e:
            return TokenResult(success=False, error=str(e))
    return TokenResult(success=False, error=f"HTTP {response.status_code}")


# =============================================================================
# ADAPTERS - Execute actions with specific HTTP clients
# =============================================================================


async def execute_flow_generic(
    flow: OAuthFlow,
    fetch: Callable[[HttpRequest], HttpResponse],
    authorize: Callable[[str], AuthorizationResult],
    www_auth_scope: str | None = None,
    www_auth_resource_metadata_url: str | None = None,
) -> OAuthToken:
    """
    Execute OAuth flow with any HTTP client.

    Args:
        flow: The OAuth flow state machine
        fetch: Function to execute HTTP requests (sync or async)
        authorize: Function to handle user authorization (opens browser, waits for callback)
        www_auth_scope: Scope from initial 401's WWW-Authenticate header
        www_auth_resource_metadata_url: resource_metadata from WWW-Authenticate

    Returns:
        OAuthToken on success

    Raises:
        Exception on failure
    """
    result = flow.start(www_auth_scope, www_auth_resource_metadata_url)

    while True:
        match result:
            case OAuthComplete(token=token):
                return token

            case OAuthFailed(error=error):
                raise Exception(f"OAuth flow failed: {error}")

            case DiscoverPRM() as action:
                request = action.to_request()
                if request is None:
                    # No more URLs to try
                    result = flow.continue_with_prm_result(PRMResult(found=False))
                else:
                    response = fetch(request)
                    prm_result = parse_prm_response(response)
                    result = flow.continue_with_prm_result(prm_result)

            case DiscoverOASM() as action:
                request = action.to_request()
                if request is None:
                    result = flow.continue_with_oasm_result(OASMResult(found=False))
                else:
                    response = fetch(request)
                    oasm_result = parse_oasm_response(response)
                    result = flow.continue_with_oasm_result(oasm_result)

            case RegisterClient() as action:
                response = fetch(action.to_request())
                reg_result = parse_registration_response(response)
                result = flow.continue_with_registration_result(reg_result)

            case AuthorizeUser(authorization_url=url):
                # Hand off to user authorization handler
                auth_result = authorize(url)
                result = flow.continue_with_authorization_result(auth_result)

            case ExchangeToken() as action:
                response = fetch(action.to_request())
                token_result = parse_token_response(response)
                result = flow.continue_with_token_result(token_result)

            case RefreshToken() as action:
                response = fetch(action.to_request())
                token_result = parse_token_response(response)
                result = flow.continue_with_token_result(token_result)


class HttpxOAuthProvider(httpx.Auth):
    """
    httpx Auth adapter for OAuthFlow.

    Thin wrapper that executes actions using httpx.
    """

    requires_response_body = True

    def __init__(
        self,
        config: OAuthFlowConfig,
        redirect_handler: Callable[[str], None],
        callback_handler: Callable[[], tuple[str, str]],  # Returns (code, state)
    ):
        self.config = config
        self.redirect_handler = redirect_handler
        self.callback_handler = callback_handler
        self.flow: OAuthFlow | None = None
        self.token: OAuthToken | None = None

    async def async_auth_flow(self, request: httpx.Request):
        """httpx auth flow integration."""
        # Add existing token if valid
        if self.token and self.token.access_token:
            request.headers["Authorization"] = f"Bearer {self.token.access_token}"

        response = yield request

        if response.status_code == 401:
            # Extract WWW-Authenticate info
            www_auth = response.headers.get("WWW-Authenticate", "")
            scope = self._extract_scope(www_auth)
            resource_metadata = self._extract_resource_metadata(www_auth)

            # Run OAuth flow
            self.flow = OAuthFlow(self.config)
            result = self.flow.start(scope, resource_metadata)

            while not isinstance(result, (OAuthComplete, OAuthFailed)):
                match result:
                    case DiscoverPRM() | DiscoverOASM() | RegisterClient() | ExchangeToken() | RefreshToken() as action:
                        http_request = action.to_request()
                        if http_request:
                            httpx_request = self._to_httpx_request(http_request)
                            httpx_response = yield httpx_request
                            result = self._handle_http_response(result, httpx_response)
                        else:
                            # No more URLs for discovery
                            result = self._handle_exhausted(result)

                    case AuthorizeUser(authorization_url=url, state=state):
                        # User authorization
                        await self.redirect_handler(url)
                        code, returned_state = await self.callback_handler()
                        result = self.flow.continue_with_authorization_result(
                            AuthorizationResult(code=code, state=returned_state)
                        )

            if isinstance(result, OAuthComplete):
                self.token = result.token
                request.headers["Authorization"] = f"Bearer {self.token.access_token}"
                yield request
            else:
                raise Exception(f"OAuth failed: {result.error}")

    def _to_httpx_request(self, req: HttpRequest) -> httpx.Request:
        """Convert HttpRequest to httpx.Request."""
        if req.form:
            return httpx.Request(req.method, req.url, headers=req.headers, data=req.form)
        elif req.json:
            return httpx.Request(req.method, req.url, headers=req.headers, json=req.json)
        else:
            return httpx.Request(req.method, req.url, headers=req.headers, content=req.body)

    def _to_http_response(self, resp: httpx.Response) -> HttpResponse:
        """Convert httpx.Response to HttpResponse."""
        return HttpResponse(
            status_code=resp.status_code,
            body=resp.content,
            headers=dict(resp.headers),
        )

    def _handle_http_response(self, action: OAuthAction, response: httpx.Response) -> OAuthResult:
        """Route HTTP response to appropriate handler."""
        http_resp = self._to_http_response(response)
        match action:
            case DiscoverPRM():
                return self.flow.continue_with_prm_result(parse_prm_response(http_resp))  # type: ignore
            case DiscoverOASM():
                return self.flow.continue_with_oasm_result(parse_oasm_response(http_resp))  # type: ignore
            case RegisterClient():
                return self.flow.continue_with_registration_result(parse_registration_response(http_resp))  # type: ignore
            case ExchangeToken() | RefreshToken():
                return self.flow.continue_with_token_result(parse_token_response(http_resp))  # type: ignore
            case _:
                raise ValueError(f"Unexpected action type: {type(action)}")

    def _handle_exhausted(self, action: OAuthAction) -> OAuthResult:
        """Handle when discovery URLs are exhausted."""
        match action:
            case DiscoverPRM():
                return self.flow.continue_with_prm_result(PRMResult(found=False))  # type: ignore
            case DiscoverOASM():
                return self.flow.continue_with_oasm_result(OASMResult(found=False))  # type: ignore
            case _:
                raise ValueError(f"Unexpected exhausted action: {type(action)}")

    def _extract_scope(self, www_auth: str) -> str | None:
        """Extract scope from WWW-Authenticate header."""
        import re

        match = re.search(r'scope="([^"]+)"', www_auth)
        return match.group(1) if match else None

    def _extract_resource_metadata(self, www_auth: str) -> str | None:
        """Extract resource_metadata from WWW-Authenticate header."""
        import re

        match = re.search(r'resource_metadata="([^"]+)"', www_auth)
        return match.group(1) if match else None


# =============================================================================
# TESTS - Example tests showing how easy this is to test
# =============================================================================


def test_prm_discovery_success():
    """Test PRM discovery with successful first URL."""
    config = OAuthFlowConfig(
        server_url="https://api.example.com/v1/mcp",
        client_metadata=OAuthClientMetadata(redirect_uris=["http://localhost:3000/callback"]),
    )
    flow = OAuthFlow(config)

    # Start flow
    action = flow.start()
    assert isinstance(action, DiscoverPRM)
    assert len(action.urls) == 2  # path-based + root-based

    # First URL succeeds
    result = flow.continue_with_prm_result(
        PRMResult(
            found=True,
            metadata=ProtectedResourceMetadata(
                resource="https://api.example.com/v1/mcp",
                authorization_servers=["https://auth.example.com"],
            ),
        )
    )

    # Should move to OASM discovery
    assert isinstance(result, DiscoverOASM)
    assert flow.state.prm is not None
    assert flow.state.auth_server_url == "https://auth.example.com"


def test_prm_discovery_fallback():
    """Test PRM discovery with fallback to second URL."""
    config = OAuthFlowConfig(
        server_url="https://api.example.com/v1/mcp",
        client_metadata=OAuthClientMetadata(redirect_uris=["http://localhost:3000/callback"]),
    )
    flow = OAuthFlow(config)

    # Start flow
    action = flow.start()
    assert isinstance(action, DiscoverPRM)
    assert action.current_index == 0

    # First URL fails
    result = flow.continue_with_prm_result(PRMResult(found=False))
    assert isinstance(result, DiscoverPRM)
    assert result.current_index == 1  # Moved to next URL

    # Second URL also fails
    result = flow.continue_with_prm_result(PRMResult(found=False))

    # Should move to OASM with legacy fallback (no auth_server_url)
    assert isinstance(result, DiscoverOASM)
    assert flow.state.auth_server_url is None
    # Legacy fallback: only root .well-known URL
    assert result.urls == ["https://api.example.com/.well-known/oauth-authorization-server"]


def test_march_2025_spec_fallback():
    """Test complete fallback to March 2025 spec defaults."""
    config = OAuthFlowConfig(
        server_url="https://api.example.com/v1/mcp",
        client_metadata=OAuthClientMetadata(redirect_uris=["http://localhost:3000/callback"]),
    )
    flow = OAuthFlow(config, initial_client_info=OAuthClientInfo(client_id="test-client"))

    # Start and exhaust PRM discovery
    action = flow.start()
    result = flow.continue_with_prm_result(PRMResult(found=False))
    result = flow.continue_with_prm_result(PRMResult(found=False))

    # Now in OASM discovery with legacy URL
    assert isinstance(result, DiscoverOASM)

    # OASM also fails
    result = flow.continue_with_oasm_result(OASMResult(found=False))

    # Should go to authorization with default endpoints
    assert isinstance(result, AuthorizeUser)

    # Verify fallback endpoints are used
    assert "/authorize?" in result.authorization_url
    assert "api.example.com" in result.authorization_url


def test_cimd_skips_registration():
    """Test that CIMD (URL-based client ID) skips registration."""
    config = OAuthFlowConfig(
        server_url="https://api.example.com/mcp",
        client_metadata=OAuthClientMetadata(redirect_uris=["http://localhost:3000/callback"]),
        client_metadata_url="https://myapp.example.com/oauth/metadata.json",
    )
    flow = OAuthFlow(config)

    # Start and complete PRM discovery
    action = flow.start()
    result = flow.continue_with_prm_result(
        PRMResult(
            found=True,
            metadata=ProtectedResourceMetadata(
                resource="https://api.example.com/mcp",
                authorization_servers=["https://auth.example.com"],
            ),
        )
    )

    # Complete OASM discovery with CIMD support
    assert isinstance(result, DiscoverOASM)
    result = flow.continue_with_oasm_result(
        OASMResult(
            found=True,
            metadata=OAuthMetadata(
                issuer="https://auth.example.com",
                authorization_endpoint="https://auth.example.com/authorize",
                token_endpoint="https://auth.example.com/token",
                client_id_metadata_document_supported=True,  # Server supports CIMD
            ),
        )
    )

    # Should skip registration and go straight to authorization
    assert isinstance(result, AuthorizeUser)
    assert flow.state.client_info is not None
    assert flow.state.client_info.client_id == "https://myapp.example.com/oauth/metadata.json"


if __name__ == "__main__":
    # Run tests
    test_prm_discovery_success()
    print("✓ test_prm_discovery_success")

    test_prm_discovery_fallback()
    print("✓ test_prm_discovery_fallback")

    test_march_2025_spec_fallback()
    print("✓ test_march_2025_spec_fallback")

    test_cimd_skips_registration()
    print("✓ test_cimd_skips_registration")

    print("\nAll tests passed!")

Note: This issue now consolidates #1724 (Support for synchronous auth flow with separate re-entry). The action/effect state machine pattern proposed here directly enables the synchronous re-entry use case requested by the Google ADK team - the "awaiting user auth" phase naturally supports pausing and resuming the flow with an authorization code.

Thanks for putting this together Max. I read through the code quickly and may be misunderstanding, but it appears to create a state machine internally within a single auth flow rather than allowing for state machine behavior between separate invocations of the auth flow.

Here's some pseudo-code of how I was imagining this would look. Very high-level, but let me know what you think.

New state machine provider

class StateMachineOAuthProvider(httpx.Auth):
   def __init__(
      authorization_code: str | None,
      code_verifier: str | None,
      redirect_callback: Callable[[str, str, str], None],  # args are auth URL, state and verifier
   ): ...

   async def async_auth_flow(self, request: httpx.Request):
      if (self._authorization_code && self._code_verifier)
         # exchange code for access token and append token to request
      
      response = yield request
      if response == 401:
         # perform WWW-auth, PRM, ASM discovery and construct auth URL
         redirect_callback(auth_url, state, verifier)
         return 401
      
      return response

Generic usage example

def redirect_callback(auth_url, state, verifier):
   self._auth_url = auth_url
   self._state = state
   self._verifier = verifier

def execute_request(request):
   provider = StateMachineOAuthProvider(None, None, redirect_callback)
   response = execute_with_provider(request, provider)

   if response == 401:
      auth_code, state = # handle auth synchronously in any fashion
      if (self._state != state) raise error
   else:
      return response

   provider = StateMachineOAuthProvider(auth_code, self._verifier, redirect_callback)
   return execute_with_provider(request, provider)

Re-implementation of existing async provider to use state machine provider

class OAuthClientProvider(httpx.Auth):
   def __init__(
       redirect_handler: Callable[[str], Awaitable[None]] | None = None,
       callback_handler: Callable[[], Awaitable[tuple[str, str | None]]] | None = None,
   ): ...

   def internal_callback(auth_url, state, verifier):
      self._auth_url = auth_url
      self._state = state
      self._verifier = verifier

   async def async_auth_flow(self, request: httpx.Request):
      internal_provider = StateMachineOAuthProvider(None, None, internal_callback)
      response = internal_provider.async_auth_flow(request)
      
      if response == 401:
         await redirect_handler(self._auth_url)
         auth_code, state = await callback_handler()
         if (self._state != state) raise error
      else:
         return response

      internal_provider = StateMachineOAuthProvider(auth_code, self._verifier, internal_callback)
      return internal_provider.async_auth_flow(request)

@jake-mayer

@jake-mayer - I like your proposed solution - having a way to async continue the auth flow would be helpful.

For reference, this would be helpful for me as my MCP client is being hosted on a serverless architecture (invoked by an endpoint); the auth flow would need to be split in two separate flows (each on different endpoints)

FlowA: Discovery, storing of state, returning auth_url
FlowB: handle callback, fetch state, store tokens

Yea fully agree with the above! We definitely want the process to be resumable at any point by a different machine if needed. This is a requirement for v2.

Thanks Max. Do we have an idea of when v2 will be out? We'd be happy to contrib this from the Google side earlier if that seems reasonable to you.

Current estimate is Q1, but it's a bit uncertain as I'm currently thinking we wait until the big upcoming changes around transport in the next spec release are finalised. They're going to be significant not just in the protocol but also the SDK, so IMO makes sense to wait for that till we officially release it.

For now we'll be doing releases of v2 of the SDK in prerelease, but each version will almost definitely have breaking changes.

@pcarleton

Related: #1240 proposes using Authlib as the OAuth implementation, which could be relevant as an implementation approach for this rework. cc @pcarleton

@maxisbey

@maxisbey is the authlib rework you mentioned something that still help? I can get back working on it again

Navigation Menu

Repository navigation

Extract OAuth flow logic into reusable components for proxy use cases #1743

Summary

Problem

Goal

Design Requirements (from maintainer discussion, Feb 2026)

Related

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Extract OAuth flow logic into reusable components for proxy use cases #1743

Description

Summary

Problem

Goal

Design Requirements (from maintainer discussion, Feb 2026)

Related

Activity

maxisbey commented on Dec 8, 2025

maxisbey commented on Dec 8, 2025

jake-mayer commented on Dec 12, 2025

ryanschulz46 commented on Dec 16, 2025

maxisbey commented on Dec 17, 2025

jake-mayer commented on Dec 18, 2025

maxisbey commented on Jan 5, 2026

maxisbey commented on Feb 19, 2026

yannj-fr commented on Apr 30, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions