pi-ai-python 1.3.0__py3-none-any.whl

pi_ai/auth/oauth/anthropic.py

@@ -0,0 +1,283 @@
+"""Anthropic OAuth flow (Claude Pro/Max).
+
+Browser login runs a loopback callback server; copy-code login is for headless use.
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+import traceback
+from typing import Any
+
+from ...abort import AbortSignal
+from ...utils.misc import get_provider_env_value
+from ..types import AuthEvent, AuthPrompt, AuthSelectOption, Credential, LoginOptions, ModelAuth, OAuthAuth
+from ..types import ProviderAuthInteraction as Interaction
+from . import _common
+from ._common import OAuthError, form_encode, json_stringify, parse_authorization_input
+from .callback_server import (
+    CallbackResult,
+    OAuthCallbackServer,
+    start_oauth_callback_server,
+    wait_for_callback_or_manual_input,
+)
+from .pkce import generate_pkce
+
+__all__ = ["AnthropicOAuth", "anthropic_oauth"]
+
+CLIENT_ID = base64.b64decode("OWQxYzI1MGEtZTYxYi00NGQ5LTg4ZWQtNTk0NGQxOTYyZjVl").decode()
+AUTHORIZE_URL = "https://claude.ai/oauth/authorize"
+TOKEN_URL = "https://platform.claude.com/v1/oauth/token"
+# Preferred so the port can be forwarded into containers or over SSH. Anthropic accepts any loopback port,
+# so login falls back to a free port when this one cannot be bound (#10571).
+CALLBACK_PORT = 53692
+CALLBACK_PATH = "/callback"
+REDIRECT_URI = f"http://localhost:{CALLBACK_PORT}{CALLBACK_PATH}"
+COPY_CODE_REDIRECT_URI = "https://platform.claude.com/oauth/code/callback"
+ANTHROPIC_BROWSER_LOGIN_METHOD = "browser"
+ANTHROPIC_COPY_CODE_LOGIN_METHOD = "copy_code"
+SCOPES = "org:create_api_key user:profile user:inference user:sessions:claude_code user:mcp_servers user:file_upload"
+_EXPIRY_SKEW_MS = 5 * 60 * 1000
+_REQUEST_TIMEOUT_MS = 30_000
+
+
+def _callback_host() -> str:
+    return get_provider_env_value("PI_OAUTH_CALLBACK_HOST") or "127.0.0.1"
+
+
+def _format_error_details(error: BaseException, depth: int = 0) -> str:
+    details = [f"{type(error).__name__}: {error}"]
+    code = getattr(error, "code", None)
+    if code:
+        details.append(f"code={code}")
+    errno = getattr(error, "errno", None)
+    if errno is not None:
+        details.append(f"errno={errno}")
+    cause = error.__cause__ if error.__cause__ is not None or error.__suppress_context__ else error.__context__
+    if cause is not None and depth < 8:
+        details.append(f"cause={_format_error_details(cause, depth + 1)}")
+    if error.__traceback__ is not None:
+        details.append(f"stack={''.join(traceback.format_tb(error.__traceback__)).rstrip()}")
+    return "; ".join(details)
+
+
+async def _post_json(url: str, body: dict[str, Any], signal: AbortSignal) -> str:
+    response = await _common.fetch(
+        url,
+        method="POST",
+        headers={"Content-Type": "application/json", "Accept": "application/json"},
+        body=json_stringify(body),
+        signal=signal,
+        timeout_ms=_REQUEST_TIMEOUT_MS,
+    )
+    response_body = response.text
+    if not response.is_success:
+        raise OAuthError(f"HTTP request failed. status={response.status_code}; url={url}; body={response_body}")
+    return response_body
+
+
+def _credential_from_token(data: Any) -> Credential:
+    token = data if isinstance(data, dict) else {}
+    credential: Credential = {"type": "oauth"}
+    if token.get("refresh_token") is not None:
+        credential["refresh"] = token["refresh_token"]
+    if token.get("access_token") is not None:
+        credential["access"] = token["access_token"]
+    credential["expires"] = _common.now_ms() + _common.expires_in_ms(token.get("expires_in")) - _EXPIRY_SKEW_MS
+    return credential
+
+
+async def _exchange_authorization_code(
+    code: str, state: str, verifier: str, redirect_uri: str, signal: AbortSignal
+) -> Credential:
+    try:
+        response_body = await _post_json(
+            TOKEN_URL,
+            {
+                "grant_type": "authorization_code",
+                "client_id": CLIENT_ID,
+                "code": code,
+                "state": state,
+                "redirect_uri": redirect_uri,
+                "code_verifier": verifier,
+            },
+            signal,
+        )
+    except Exception as error:
+        raise OAuthError(
+            f"Token exchange request failed. url={TOKEN_URL}; redirect_uri={redirect_uri}; "
+            f"response_type=authorization_code; details={_format_error_details(error)}"
+        ) from error
+
+    try:
+        token_data = json.loads(response_body)
+    except ValueError as error:
+        raise OAuthError(
+            f"Token exchange returned invalid JSON. url={TOKEN_URL}; body={response_body}; "
+            f"details={_format_error_details(error)}"
+        ) from error
+
+    return _credential_from_token(token_data)
+
+
+def _authorize_url(redirect_uri: str, challenge: str, verifier: str) -> str:
+    params = {
+        "code": "true",
+        "client_id": CLIENT_ID,
+        "response_type": "code",
+        "redirect_uri": redirect_uri,
+        "scope": SCOPES,
+        "code_challenge": challenge,
+        "code_challenge_method": "S256",
+        "state": verifier,
+    }
+    return f"{AUTHORIZE_URL}?{form_encode(params)}"
+
+
+async def _login_anthropic(interaction: Interaction) -> Credential:
+    verifier, challenge = generate_pkce()
+
+    async def identity(code: str) -> str:
+        return code
+
+    async def start_callback_server(port: int) -> OAuthCallbackServer[str]:
+        return await start_oauth_callback_server(
+            provider_name="Anthropic",
+            host=_callback_host(),
+            port=port,
+            path=CALLBACK_PATH,
+            redirect_host="localhost",
+            state=verifier,
+            complete=identity,
+            signal=interaction.signal,
+        )
+
+    # Without a callback server, login continues with the pasted redirect URL.
+    callback: OAuthCallbackServer[str] | None = None
+    for port in (CALLBACK_PORT, 0):
+        try:
+            callback = await start_callback_server(port)
+            break
+        except Exception:
+            continue
+    redirect_uri = callback.redirect_uri if callback is not None else REDIRECT_URI
+
+    try:
+        interaction.notify(
+            AuthEvent(
+                type="auth_url",
+                url=_authorize_url(redirect_uri, challenge, verifier),
+                instructions=(
+                    "Complete login in your browser. If the browser is on another machine, "
+                    "paste the final redirect URL here."
+                ),
+            )
+        )
+
+        result = await wait_for_callback_or_manual_input(
+            interaction,
+            callback,
+            message="Complete login in your browser, or paste the authorization code / redirect URL here:",
+            placeholder=redirect_uri,
+        )
+        code: str | None
+        state = verifier
+        if isinstance(result, CallbackResult):
+            code = result.value
+        else:
+            code, parsed_state = parse_authorization_input(result.input)
+            if parsed_state and parsed_state != verifier:
+                raise OAuthError("OAuth state mismatch")
+            state = parsed_state if parsed_state is not None else verifier
+
+        if not code:
+            raise OAuthError("Missing authorization code")
+        interaction.notify(AuthEvent(type="progress", message="Exchanging authorization code for tokens..."))
+        return await _exchange_authorization_code(code, state, verifier, redirect_uri, interaction.signal)
+    finally:
+        if callback is not None:
+            callback.close()
+
+
+async def _login_anthropic_copy_code(interaction: Interaction) -> Credential:
+    verifier, challenge = generate_pkce()
+    interaction.notify(
+        AuthEvent(
+            type="auth_url",
+            url=_authorize_url(COPY_CODE_REDIRECT_URI, challenge, verifier),
+            instructions="Complete login in your browser, then copy the code Anthropic shows and paste it here.",
+        )
+    )
+
+    entered = await interaction.prompt(
+        AuthPrompt(
+            type="manual_code",
+            message="Paste the code Anthropic shows after you sign in:",
+            placeholder="code#state",
+            signal=interaction.signal,
+        )
+    )
+    code, state = parse_authorization_input(entered)
+    if state and state != verifier:
+        raise OAuthError("OAuth state mismatch")
+    if not code:
+        raise OAuthError("Missing authorization code")
+    interaction.notify(AuthEvent(type="progress", message="Exchanging authorization code for tokens..."))
+    return await _exchange_authorization_code(
+        code, state if state is not None else verifier, verifier, COPY_CODE_REDIRECT_URI, interaction.signal
+    )
+
+
+async def _refresh_anthropic_token(refresh_token: str, signal: AbortSignal) -> Credential:
+    try:
+        response_body = await _post_json(
+            TOKEN_URL,
+            {"grant_type": "refresh_token", "client_id": CLIENT_ID, "refresh_token": refresh_token},
+            signal,
+        )
+    except Exception as error:
+        raise OAuthError(
+            f"Anthropic token refresh request failed. url={TOKEN_URL}; details={_format_error_details(error)}"
+        ) from error
+
+    try:
+        data = json.loads(response_body)
+    except ValueError as error:
+        raise OAuthError(
+            f"Anthropic token refresh returned invalid JSON. url={TOKEN_URL}; body={response_body}; "
+            f"details={_format_error_details(error)}"
+        ) from error
+
+    return _credential_from_token(data)
+
+
+class AnthropicOAuth(OAuthAuth):
+    name = "Anthropic (Claude Pro/Max)"
+    is_subscription = True
+
+    async def login(self, interaction: Interaction, options: LoginOptions | None = None) -> Credential:
+        method = await interaction.prompt(
+            AuthPrompt(
+                type="select",
+                message="Select Anthropic login method:",
+                options=[
+                    AuthSelectOption(ANTHROPIC_BROWSER_LOGIN_METHOD, "Browser login (default)"),
+                    AuthSelectOption(ANTHROPIC_COPY_CODE_LOGIN_METHOD, "Copy code login (headless)"),
+                ],
+            )
+        )
+        if method == ANTHROPIC_COPY_CODE_LOGIN_METHOD:
+            return await _login_anthropic_copy_code(interaction)
+        if method != ANTHROPIC_BROWSER_LOGIN_METHOD:
+            raise OAuthError(f"Unknown Anthropic login method: {method}")
+        return await _login_anthropic(interaction)
+
+    async def refresh(self, credential: Credential, signal: AbortSignal) -> Credential:
+        return await _refresh_anthropic_token(credential["refresh"], signal)
+
+    async def to_auth(self, credential: Credential) -> ModelAuth:
+        return ModelAuth(api_key=credential["access"])
+
+
+anthropic_oauth = AnthropicOAuth()

pi_ai/auth/oauth/callback_server.py

@@ -0,0 +1,339 @@
+"""Loopback OAuth redirect handler shared by the browser sign-in flows.
+
+The server is a minimal asyncio HTTP/1.1 listener (one request per connection)
+running on the caller's event loop, so it is cancellable together with the login.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import os
+import socket
+from collections.abc import Awaitable, Callable, Mapping
+from dataclasses import dataclass
+from http import HTTPStatus
+from typing import Any, ClassVar, Generic, TypeVar
+
+from ...abort import AbortController, AbortSignal
+from ...utils.oauth_page import oauth_error_html, oauth_success_html
+from ..types import AuthPrompt, ProviderAuthInteraction
+from ._common import OAuthError, ParsedUrl, login_cancelled, parse_url
+
+__all__ = [
+    "LoopbackHttpServer",
+    "OAuthCallbackServer",
+    "CallbackResult",
+    "ManualInput",
+    "start_oauth_callback_server",
+    "wait_for_callback_or_manual_input",
+]
+
+T = TypeVar("T")
+
+#: ``(method, url) -> (status, html)`` where ``url`` is the request target resolved against the server origin.
+RequestHandler = Callable[[str, ParsedUrl | None], Awaitable[tuple[int, str]]]
+
+_MAX_HEAD_BYTES = 64 * 1024
+
+
+class _Connection:
+    __slots__ = ("active", "writer")
+
+    def __init__(self, writer: asyncio.StreamWriter) -> None:
+        self.writer = writer
+        self.active = False
+
+
+class LoopbackHttpServer:
+    """Tiny HTTP/1.1 server that answers every request with an HTML page and closes the connection."""
+
+    def __init__(self, handler: RequestHandler, *, headers: Mapping[str, str] | None = None) -> None:
+        self._handler = handler
+        self._headers = dict(headers or {"content-type": "text/html; charset=utf-8"})
+        self._server: asyncio.Server | None = None
+        self._connections: set[_Connection] = set()
+        self.port = 0
+        self._origin = "http://localhost"
+
+    async def listen(self, host: str, port: int) -> None:
+        """Bind and start serving. Raises ``OSError`` (e.g. ``EADDRINUSE``) when the port is taken."""
+        loop = asyncio.get_running_loop()
+        # Bind the first resolved address only (like Node's `server.listen(port, host)`), so a
+        # free port picked with `port=0` is the same for every client of this server.
+        infos = await loop.getaddrinfo(host, port, type=socket.SOCK_STREAM, flags=socket.AI_PASSIVE)
+        family, sock_type, proto, _, address = infos[0]
+        sock = socket.socket(family, sock_type, proto)
+        try:
+            if os.name != "nt":
+                sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            sock.bind(address)
+            sock.listen(511)
+            sock.setblocking(False)
+            self._server = await asyncio.start_server(self._serve, sock=sock, limit=_MAX_HEAD_BYTES)
+        except BaseException:
+            sock.close()
+            raise
+        self.port = sock.getsockname()[1]
+        self._origin = f"http://{f'[{host}]' if ':' in host else host}:{self.port}"
+
+    def close(self) -> None:
+        """Stop accepting connections and drop idle ones; requests being handled complete."""
+        if self._server is not None:
+            self._server.close()
+        for connection in list(self._connections):
+            if not connection.active:
+                self._drop(connection)
+
+    def close_all_connections(self) -> None:
+        for connection in list(self._connections):
+            self._drop(connection)
+
+    def _drop(self, connection: _Connection) -> None:
+        self._connections.discard(connection)
+        with contextlib.suppress(Exception):
+            connection.writer.close()
+
+    async def _serve(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
+        connection = _Connection(writer)
+        self._connections.add(connection)
+        try:
+            try:
+                head = await reader.readuntil(b"\r\n\r\n")
+            except (asyncio.IncompleteReadError, asyncio.LimitOverrunError, ConnectionError):
+                return
+            connection.active = True
+            request_line = head.split(b"\r\n", 1)[0].decode("latin-1")
+            parts = request_line.split(" ")
+            if len(parts) != 3 or not parts[2].startswith("HTTP/"):
+                status, html = 400, oauth_error_html("Bad request.")
+            else:
+                method, target = parts[0], parts[1]
+                try:
+                    status, html = await self._handler(method, parse_url(target, self._origin))
+                except Exception:
+                    status, html = 500, oauth_error_html("Internal error while processing the callback.")
+            body = html.encode("utf-8")
+            try:
+                reason = HTTPStatus(status).phrase
+            except ValueError:
+                reason = ""
+            lines = [f"HTTP/1.1 {status} {reason}"]
+            lines.extend(f"{name}: {value}" for name, value in self._headers.items())
+            lines.append(f"content-length: {len(body)}")
+            lines.append("connection: close")
+            writer.write(("\r\n".join(lines) + "\r\n\r\n").encode("latin-1") + body)
+            await writer.drain()
+        except (ConnectionError, RuntimeError):
+            pass
+        finally:
+            self._connections.discard(connection)
+            with contextlib.suppress(Exception):
+                writer.close()
+
+
+class OAuthCallbackServer(Generic[T]):
+    """A running loopback redirect handler. Create with :func:`start_oauth_callback_server`."""
+
+    def __init__(
+        self,
+        *,
+        provider_name: str,
+        path: str,
+        state: str | None,
+        complete: Callable[[str], Awaitable[T]],
+    ) -> None:
+        self.redirect_uri = ""
+        self._provider_name = provider_name
+        self._path = path
+        self._state = state
+        self._complete = complete
+        self._loop = asyncio.get_running_loop()
+        self._future: asyncio.Future[T | None] = self._loop.create_future()
+        # A cancelled or closed wait may never be observed.
+        self._future.add_done_callback(_consume_future)
+        self._claimed = False
+        self._settled = False
+        self._timer: asyncio.TimerHandle | None = None
+        self._remove_abort: Callable[[], None] | None = None
+        self._http = LoopbackHttpServer(
+            self._handle, headers={"content-type": "text/html; charset=utf-8", "cache-control": "no-store"}
+        )
+
+    def _finish(self, *, value: T | None = None, error: BaseException | None = None) -> None:
+        if self._settled:
+            return
+        self._settled = True
+        if self._timer is not None:
+            self._timer.cancel()
+        if self._remove_abort is not None:
+            self._remove_abort()
+        if self._future.done():
+            return
+        if error is not None:
+            self._future.set_exception(error)
+        else:
+            self._future.set_result(value)
+
+    def _on_abort(self) -> None:
+        def finish() -> None:
+            self._finish(error=login_cancelled())
+
+        try:
+            running = asyncio.get_running_loop()
+        except RuntimeError:
+            running = None
+        if running is self._loop:
+            finish()
+        elif not self._loop.is_closed():
+            self._loop.call_soon_threadsafe(finish)
+
+    async def _handle(self, method: str, url: ParsedUrl | None) -> tuple[int, str]:
+        if method != "GET" or url is None or url.pathname != self._path:
+            return 404, oauth_error_html("Callback route not found.")
+        if self._state is not None and url.param("state") != self._state:
+            return 400, oauth_error_html("State mismatch.")
+        if self._claimed or self._settled:
+            return 409, oauth_error_html("This sign-in has already been handled.")
+        error = url.param("error")
+        if error:
+            description = url.param("error_description")
+            if description is None:
+                description = error
+            self._finish(error=OAuthError(f"{self._provider_name} authorization failed: {description}"))
+            return 400, oauth_error_html(f"{self._provider_name} authorization failed.", description)
+        code = url.param("code")
+        if not code:
+            return 400, oauth_error_html("Missing authorization code.")
+        self._claimed = True
+        try:
+            value = await self._complete(code)
+        except Exception as failure:
+            self._finish(error=failure)
+            return 502, oauth_error_html(f"{self._provider_name} sign-in failed.", str(failure))
+        self._finish(value=value)
+        return 200, oauth_success_html(f"Signed in to {self._provider_name}. You may now close this page.")
+
+    async def wait(self) -> T | None:
+        """The result of ``complete``, or ``None`` after :meth:`cancel`.
+
+        Raises when the provider redirects with an error, ``complete`` fails, the signal
+        aborts, or the timeout elapses.
+        """
+        return await asyncio.shield(self._future)
+
+    def cancel(self) -> None:
+        """Stop waiting for the browser unless a callback is already being completed."""
+        if not self._claimed:
+            self._finish(value=None)
+
+    def close(self) -> None:
+        self._finish(error=OAuthError("OAuth callback server closed"))
+        self._http.close()
+
+
+async def start_oauth_callback_server(
+    *,
+    provider_name: str,
+    host: str,
+    port: int,
+    path: str,
+    complete: Callable[[str], Awaitable[T]],
+    redirect_host: str | None = None,
+    state: str | None = None,
+    signal: AbortSignal | None = None,
+    timeout_ms: float | None = None,
+) -> OAuthCallbackServer[T]:
+    """Listen for the OAuth redirect on ``host:port`` (``0`` picks a free port).
+
+    ``provider_name`` is shown on the browser page. ``redirect_host`` replaces ``host``
+    in ``redirect_uri`` (for example ``localhost``). ``state`` is the expected ``state``
+    parameter (omit when the provider sends none). ``complete`` finishes the sign-in with
+    the received code before the browser page is sent, so the page can show exchange
+    failures; pass ``lambda code: code`` (async) to exchange the code later.
+    """
+    if signal is not None and signal.aborted:
+        raise login_cancelled()
+    server: OAuthCallbackServer[T] = OAuthCallbackServer(
+        provider_name=provider_name, path=path, state=state, complete=complete
+    )
+    await server._http.listen(host, port)
+    if signal is not None:
+        server._remove_abort = signal.add_listener(server._on_abort)
+    if timeout_ms is not None:
+        server._timer = server._loop.call_later(
+            timeout_ms / 1000, lambda: server._finish(error=OAuthError(f"{provider_name} sign-in timed out"))
+        )
+    effective_host = redirect_host if redirect_host is not None else host
+    formatted = f"[{effective_host}]" if ":" in effective_host else effective_host
+    server.redirect_uri = f"http://{formatted}:{server._http.port}{path}"
+    return server
+
+
+@dataclass(frozen=True, slots=True)
+class CallbackResult:
+    """The browser callback completed with ``value``."""
+
+    value: Any
+    type: ClassVar[str] = "callback"
+
+
+@dataclass(frozen=True, slots=True)
+class ManualInput:
+    """The user pasted ``input`` (a code or redirect URL)."""
+
+    input: str
+    type: ClassVar[str] = "manual"
+
+
+async def wait_for_callback_or_manual_input(
+    interaction: ProviderAuthInteraction,
+    callback: OAuthCallbackServer[Any] | None,
+    *,
+    message: str,
+    placeholder: str,
+) -> CallbackResult | ManualInput:
+    """Wait for the browser callback, or for the user to paste the code or redirect URL.
+
+    The paste prompt covers browsers that cannot reach the loopback server (for example
+    over SSH). Without a callback server only the manual prompt is used. Once one side
+    wins, the manual prompt's signal aborts so UIs can dismiss it.
+    """
+    manual_abort = AbortController()
+    manual_errors: list[BaseException] = []
+
+    async def manual() -> str | None:
+        try:
+            value = await interaction.prompt(
+                AuthPrompt(type="manual_code", message=message, placeholder=placeholder, signal=manual_abort.signal)
+            )
+        except Exception as error:
+            manual_errors.append(error)
+            if callback is not None:
+                callback.cancel()
+            return None
+        if callback is not None:
+            callback.cancel()
+        return value
+
+    manual_task = asyncio.ensure_future(manual())
+    try:
+        value = await callback.wait() if callback is not None else None
+        if manual_errors:
+            raise manual_errors[0]
+        if value is not None:
+            return CallbackResult(value)
+        entered = await manual_task
+        if manual_errors:
+            raise manual_errors[0]
+        return ManualInput(entered or "")
+    finally:
+        manual_abort.abort()
+        if not manual_task.done():
+            # Prompts that ignore their signal would otherwise stay pending forever.
+            manual_task.cancel()
+
+
+def _consume_future(future: asyncio.Future[Any]) -> None:
+    if not future.cancelled():
+        future.exception()