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

pi_ai/auth/oauth/kimi_coding.py

@@ -0,0 +1,283 @@
+"""Kimi Code (subscription) OAuth flow.
+
+RFC 8628 device authorization grant against https://auth.kimi.com with JSON
+responses. The access token authenticates requests to https://api.kimi.com/coding
+as an ``Authorization: Bearer`` header.
+"""
+
+from __future__ import annotations
+
+import math
+import re
+from dataclasses import dataclass
+from typing import Any
+
+import httpx
+
+from ...abort import AbortError, AbortSignal
+from ...utils.misc import get_provider_env_value
+from ..types import AuthEvent, Credential, LoginOptions, ModelAuth, OAuthAuth
+from ..types import ProviderAuthInteraction as Interaction
+from . import _common
+from ._common import OAuthError, form_encode, json_stringify, parse_url
+from .device_code import (
+    DeviceCodeComplete,
+    DeviceCodeFailed,
+    DeviceCodePending,
+    DeviceCodePollResult,
+    DeviceCodeSlowDown,
+    poll_oauth_device_code_flow,
+)
+
+__all__ = ["KimiCodingOAuth", "kimi_coding_oauth"]
+
+CLIENT_ID = "17e5f671-d194-4dfb-9706-5516cb48c098"
+DEFAULT_OAUTH_HOST = "https://auth.kimi.com"
+DEVICE_CODE_TIMEOUT_SECONDS = 15 * 60
+DEFAULT_POLL_INTERVAL_SECONDS = 5
+REQUEST_TIMEOUT_MS = 30 * 1000
+REFRESH_MAX_RETRIES = 3
+
+_FORM_HEADERS = {"Content-Type": "application/x-www-form-urlencoded", "Accept": "application/json"}
+
+
+@dataclass(slots=True)
+class _DeviceAuthorization:
+    device_code: str
+    user_code: str
+    verification_uri: str
+    verification_uri_complete: str
+    interval_seconds: float
+    expires_in_seconds: float
+
+
+@dataclass(slots=True)
+class _TokenResponse:
+    access: str
+    refresh: str
+    expires: int
+
+
+def _oauth_host() -> str:
+    override = get_provider_env_value("KIMI_CODE_OAUTH_HOST") or get_provider_env_value("KIMI_OAUTH_HOST")
+    return re.sub(r"/+$", "", override or DEFAULT_OAUTH_HOST)
+
+
+def _read_json(response: httpx.Response) -> Any:
+    """The JSON body when it is an object or array, otherwise ``None``."""
+    try:
+        value = _common.parse_json(response)
+    except ValueError:
+        return None
+    return value if isinstance(value, (dict, list)) else None
+
+
+def _field(data: Any, key: str) -> Any:
+    return data.get(key) if isinstance(data, dict) else None
+
+
+def _trusted_http_url(value: Any) -> str | None:
+    """The verification URI is opened in the user's browser; only http(s) URLs are trusted."""
+    if not isinstance(value, str) or not value:
+        return None
+    url = parse_url(value)
+    if url is None or url.protocol not in ("https:", "http:"):
+        return None
+    return url.href
+
+
+def _positive(value: Any) -> bool:
+    return _common.is_number(value) and math.isfinite(value) and value > 0
+
+
+async def _start_device_authorization(oauth_host: str, signal: AbortSignal) -> _DeviceAuthorization:
+    response = await _common.fetch(
+        f"{oauth_host}/api/oauth/device_authorization",
+        method="POST",
+        headers=_FORM_HEADERS,
+        body=form_encode({"client_id": CLIENT_ID}),
+        signal=signal,
+        timeout_ms=REQUEST_TIMEOUT_MS,
+    )
+
+    if not response.is_success:
+        text = response.text
+        raise OAuthError(
+            f"Kimi Code device authorization failed with status {response.status_code}{f': {text}' if text else ''}"
+        )
+
+    data = _read_json(response)
+    device_code = _field(data, "device_code")
+    user_code = _field(data, "user_code")
+    verification_uri = _field(data, "verification_uri")
+    verification_uri_complete = _field(data, "verification_uri_complete")
+    if (
+        not isinstance(device_code, str)
+        or not isinstance(user_code, str)
+        or not isinstance(verification_uri, str)
+        or not isinstance(verification_uri_complete, str)
+        or not _trusted_http_url(verification_uri_complete)
+        or not _trusted_http_url(verification_uri)
+    ):
+        raise OAuthError(f"Invalid Kimi Code device authorization response: {json_stringify(data)}")
+
+    interval = _field(data, "interval")
+    expires_in = _field(data, "expires_in")
+    return _DeviceAuthorization(
+        device_code=device_code,
+        user_code=user_code,
+        verification_uri=verification_uri,
+        verification_uri_complete=verification_uri_complete,
+        interval_seconds=interval if _positive(interval) else DEFAULT_POLL_INTERVAL_SECONDS,
+        expires_in_seconds=expires_in if _positive(expires_in) else DEVICE_CODE_TIMEOUT_SECONDS,
+    )
+
+
+def _parse_token_response(data: Any, operation: str) -> _TokenResponse:
+    access_token = _field(data, "access_token")
+    refresh_token = _field(data, "refresh_token")
+    expires_in = _field(data, "expires_in")
+    if (
+        not isinstance(access_token, str)
+        or not access_token
+        or not isinstance(refresh_token, str)
+        or not refresh_token
+        or not _positive(expires_in)
+    ):
+        raise OAuthError(f"Kimi Code token {operation} response missing fields: {json_stringify(data)}")
+    return _TokenResponse(access_token, refresh_token, _common.now_ms() + _common.expires_in_ms(expires_in))
+
+
+async def _poll_for_token(oauth_host: str, device: _DeviceAuthorization, signal: AbortSignal) -> _TokenResponse:
+    async def poll() -> DeviceCodePollResult:
+        response = await _common.fetch(
+            f"{oauth_host}/api/oauth/token",
+            method="POST",
+            headers=_FORM_HEADERS,
+            body=form_encode(
+                {
+                    "client_id": CLIENT_ID,
+                    "device_code": device.device_code,
+                    "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
+                }
+            ),
+            signal=signal,
+            timeout_ms=REQUEST_TIMEOUT_MS,
+        )
+
+        if response.status_code >= 500:
+            text = response.text
+            return DeviceCodeFailed(
+                f"Kimi Code device token request failed with status {response.status_code}{f': {text}' if text else ''}"
+            )
+
+        data = _read_json(response)
+        if response.is_success and isinstance(_field(data, "access_token"), str):
+            try:
+                return DeviceCodeComplete(_parse_token_response(data, "poll"))
+            except Exception as error:
+                return DeviceCodeFailed(str(error))
+
+        error = _field(data, "error")
+        error_description = _field(data, "error_description")
+        description = f": {error_description}" if isinstance(error_description, str) else ""
+        if error == "authorization_pending":
+            return DeviceCodePending()
+        if error == "slow_down":
+            interval = _field(data, "interval")
+            return DeviceCodeSlowDown(interval if _common.is_number(interval) and interval > 0 else None)
+        if error == "expired_token":
+            return DeviceCodeFailed("Kimi Code device authorization expired. Please restart login.")
+        if error == "access_denied":
+            return DeviceCodeFailed("Kimi Code login was denied.")
+        return DeviceCodeFailed(
+            f"Kimi Code device token request failed (status {response.status_code})"
+            f"{f': {error}{description}' if isinstance(error, str) else ''}"
+        )
+
+    return await poll_oauth_device_code_flow(
+        poll=poll,
+        signal=signal,
+        interval_seconds=device.interval_seconds,
+        expires_in_seconds=device.expires_in_seconds,
+        wait_before_first_poll=True,
+    )
+
+
+def _is_retryable_refresh_failure(response: httpx.Response) -> bool:
+    return response.status_code == 429 or response.status_code >= 500
+
+
+async def _refresh_token(oauth_host: str, refresh_token: str, signal: AbortSignal) -> _TokenResponse:
+    last_error: BaseException | None = None
+    for attempt in range(REFRESH_MAX_RETRIES + 1):
+        if attempt > 0:
+            await _common.sleep(2 ** (attempt - 1), signal)
+        if signal.aborted:
+            raise AbortError("Kimi Code token refresh aborted")
+
+        try:
+            response = await _common.fetch(
+                f"{oauth_host}/api/oauth/token",
+                method="POST",
+                headers=_FORM_HEADERS,
+                body=form_encode(
+                    {"client_id": CLIENT_ID, "grant_type": "refresh_token", "refresh_token": refresh_token}
+                ),
+                signal=signal,
+                timeout_ms=REQUEST_TIMEOUT_MS,
+            )
+        except Exception as error:
+            last_error = error
+            continue
+
+        data = _read_json(response)
+        if response.is_success:
+            return _parse_token_response(data, "refresh")
+
+        # Unauthorized: the stored credential is dead; Models clears it and prompts re-login.
+        if response.status_code in (401, 403) or _field(data, "error") == "invalid_grant":
+            error_description = _field(data, "error_description")
+            description = f": {error_description}" if isinstance(error_description, str) else ""
+            raise OAuthError(f"Kimi Code token refresh unauthorized (status {response.status_code}){description}")
+
+        if _is_retryable_refresh_failure(response) and attempt < REFRESH_MAX_RETRIES:
+            last_error = OAuthError(f"Kimi Code token refresh failed with status {response.status_code}")
+            continue
+
+        raise OAuthError(f"Kimi Code token refresh failed with status {response.status_code}: {json_stringify(data)}")
+
+    raise last_error if last_error is not None else OAuthError("Kimi Code token refresh failed")
+
+
+def _credential(token: _TokenResponse) -> Credential:
+    return {"type": "oauth", "access": token.access, "refresh": token.refresh, "expires": token.expires}
+
+
+class KimiCodingOAuth(OAuthAuth):
+    name = "Kimi Code (subscription)"
+    is_subscription = True
+    login_label = "Sign in with Kimi Code"
+
+    async def login(self, interaction: Interaction, options: LoginOptions | None = None) -> Credential:
+        oauth_host = _oauth_host()
+        device = await _start_device_authorization(oauth_host, interaction.signal)
+        interaction.notify(
+            AuthEvent(
+                type="device_code",
+                user_code=device.user_code,
+                verification_uri=device.verification_uri_complete,
+                interval_seconds=device.interval_seconds,
+                expires_in_seconds=device.expires_in_seconds,
+            )
+        )
+        return _credential(await _poll_for_token(oauth_host, device, interaction.signal))
+
+    async def refresh(self, credential: Credential, signal: AbortSignal) -> Credential:
+        return _credential(await _refresh_token(_oauth_host(), credential["refresh"], signal))
+
+    async def to_auth(self, credential: Credential) -> ModelAuth:
+        return ModelAuth(headers={"Authorization": f"Bearer {credential['access']}"})
+
+
+kimi_coding_oauth = KimiCodingOAuth()

pi_ai/auth/oauth/load.py

@@ -0,0 +1,64 @@
+"""Loaders for the built-in OAuth flows.
+
+Provider factories pass these to :func:`pi_ai.lazy_oauth` so flow modules (local
+callback servers, PKCE helpers) import only when a login or refresh happens.
+"""
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+
+from ..types import OAuthAuth
+
+__all__ = [
+    "load_anthropic_oauth",
+    "load_openai_codex_oauth",
+    "load_openai_chatgpt_oauth",
+    "load_github_copilot_oauth",
+    "load_openrouter_oauth",
+    "load_kimi_coding_oauth",
+    "load_meta_oauth",
+    "load_xai_oauth",
+    "load_radius_oauth",
+]
+
+
+def _load(module: str, attribute: str) -> Any:
+    return getattr(importlib.import_module(f"pi_ai.auth.oauth.{module}"), attribute)
+
+
+def load_anthropic_oauth() -> OAuthAuth:
+    return _load("anthropic", "anthropic_oauth")
+
+
+def load_openai_codex_oauth() -> OAuthAuth:
+    return _load("openai_codex", "openai_codex_oauth")
+
+
+def load_openai_chatgpt_oauth() -> OAuthAuth:
+    return _load("openai_chatgpt", "openai_chatgpt_oauth")
+
+
+def load_github_copilot_oauth() -> OAuthAuth:
+    return _load("github_copilot", "github_copilot_oauth")
+
+
+def load_openrouter_oauth() -> OAuthAuth:
+    return _load("openrouter", "openrouter_oauth")
+
+
+def load_kimi_coding_oauth() -> OAuthAuth:
+    return _load("kimi_coding", "kimi_coding_oauth")
+
+
+def load_meta_oauth() -> OAuthAuth:
+    return _load("meta", "meta_oauth")
+
+
+def load_xai_oauth() -> OAuthAuth:
+    return _load("xai", "xai_oauth")
+
+
+def load_radius_oauth(*, name: str, gateway: str) -> OAuthAuth:
+    return _load("radius", "create_radius_oauth")(name=name, gateway=gateway)

pi_ai/auth/oauth/meta.py

@@ -0,0 +1,238 @@
+"""Meta Model API OAuth flow.
+
+RFC 8628 device authorization grant against https://auth.meta.com (JSON
+responses). Meta splits identity from API access: the resulting identity token
+is not accepted for inference, so it is exchanged for a Model API key via the
+Muse Code key-mint endpoint (minted keys live about a day). The identity token
+is stored as ``refresh`` and the minted key as ``access``, so the standard OAuth
+scheduler re-mints the key when it expires with no bespoke renewal machinery.
+The identity token itself is not renewable (auth.meta.com answers
+``grant_type=refresh_token`` with 404 and issues no refresh_token), so a 401/403
+from mint means the session is dead and the user must sign in again.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Any
+
+import httpx
+
+from ...abort import AbortSignal
+from ..types import AuthEvent, Credential, LoginOptions, ModelAuth, OAuthAuth
+from ..types import ProviderAuthInteraction as Interaction
+from . import _common
+from ._common import OAuthError, form_encode, json_stringify, parse_url
+from .device_code import (
+    DeviceCodeComplete,
+    DeviceCodeFailed,
+    DeviceCodePending,
+    DeviceCodePollResult,
+    DeviceCodeSlowDown,
+    poll_oauth_device_code_flow,
+)
+
+__all__ = ["MetaOAuth", "meta_oauth"]
+
+# Muse Code CLI client id.
+CLIENT_ID = "1031625952748946"
+AUTH_HOST = "https://auth.meta.com"
+DEVICE_AUTHORIZATION_URL = f"{AUTH_HOST}/oidc/device/authorization/"
+DEVICE_TOKEN_URL = f"{AUTH_HOST}/oidc/device/token/"
+API_KEY_MINT_URL = "https://api.meta.ai/muse-code/key"
+API_KEY_LIFETIME_MS = 24 * 60 * 60 * 1000
+REQUEST_TIMEOUT_MS = 30 * 1000
+
+_FORM_HEADERS = {"Content-Type": "application/x-www-form-urlencoded", "Accept": "application/json"}
+
+
+@dataclass(slots=True)
+class _DeviceAuthorization:
+    device_code: str
+    user_code: str
+    verification_uri: str
+    interval_seconds: float | None
+    expires_in_seconds: float | None
+
+
+def _read_json(response: httpx.Response) -> Any:
+    try:
+        value = _common.parse_json(response)
+    except ValueError:
+        return None
+    return value if isinstance(value, (dict, list)) else None
+
+
+def _field(data: Any, key: str) -> Any:
+    return data.get(key) if isinstance(data, dict) else None
+
+
+def _error_detail(data: Any) -> str:
+    for key in ("error_description", "detail", "message", "error"):
+        value = _field(data, key)
+        if isinstance(value, str) and value.strip():
+            return f": {value.strip()}"
+    return ""
+
+
+def _trusted_http_url(value: Any) -> str | None:
+    """The verification URI is opened in the user's browser; only http(s) URLs are trusted."""
+    if not isinstance(value, str) or not value:
+        return None
+    url = parse_url(value)
+    if url is None or url.protocol not in ("https:", "http:"):
+        return None
+    return url.href
+
+
+def _positive_number(value: Any) -> float | None:
+    return value if _common.is_number(value) and math.isfinite(value) and value > 0 else None
+
+
+async def _start_device_authorization(signal: AbortSignal) -> _DeviceAuthorization:
+    response = await _common.fetch(
+        DEVICE_AUTHORIZATION_URL,
+        method="POST",
+        headers=_FORM_HEADERS,
+        body=form_encode({"client_id": CLIENT_ID}),
+        signal=signal,
+        timeout_ms=REQUEST_TIMEOUT_MS,
+    )
+    data = _read_json(response)
+    if not response.is_success:
+        raise OAuthError(f"Meta device authorization failed with status {response.status_code}{_error_detail(data)}")
+    device_code = _field(data, "device_code")
+    user_code = _field(data, "user_code")
+    verification_uri = _trusted_http_url(_field(data, "verification_uri_complete")) or _trusted_http_url(
+        _field(data, "verification_uri")
+    )
+    if not isinstance(device_code, str) or not device_code or not isinstance(user_code, str) or not user_code:
+        raise OAuthError(f"Invalid Meta device authorization response: {json_stringify(data)}")
+    if not verification_uri:
+        raise OAuthError(f"Invalid Meta device authorization response: {json_stringify(data)}")
+    return _DeviceAuthorization(
+        device_code=device_code,
+        user_code=user_code,
+        verification_uri=verification_uri,
+        interval_seconds=_positive_number(_field(data, "interval")),
+        expires_in_seconds=_positive_number(_field(data, "expires_in")),
+    )
+
+
+async def _poll_for_identity_token(device: _DeviceAuthorization, signal: AbortSignal) -> str:
+    async def poll() -> DeviceCodePollResult:
+        response = await _common.fetch(
+            DEVICE_TOKEN_URL,
+            method="POST",
+            headers=_FORM_HEADERS,
+            body=form_encode(
+                {
+                    "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
+                    "device_code": device.device_code,
+                    "client_id": CLIENT_ID,
+                }
+            ),
+            signal=signal,
+            timeout_ms=REQUEST_TIMEOUT_MS,
+        )
+        data = _read_json(response)
+        access_token = _field(data, "access_token")
+        if response.is_success and isinstance(access_token, str) and access_token:
+            return DeviceCodeComplete(access_token)
+        error = _field(data, "error")
+        if error == "authorization_pending":
+            return DeviceCodePending()
+        if error == "slow_down":
+            return DeviceCodeSlowDown(_positive_number(_field(data, "interval")))
+        if error == "access_denied":
+            return DeviceCodeFailed("Meta login was denied.")
+        if error == "expired_token":
+            return DeviceCodeFailed("Meta device authorization expired. Please restart login.")
+        return DeviceCodeFailed(
+            f"Meta device token request failed with status {response.status_code}{_error_detail(data)}"
+        )
+
+    return await poll_oauth_device_code_flow(
+        poll=poll,
+        signal=signal,
+        interval_seconds=device.interval_seconds,
+        expires_in_seconds=device.expires_in_seconds,
+        wait_before_first_poll=True,
+    )
+
+
+async def _mint_api_key(identity_token: str, signal: AbortSignal) -> Credential:
+    """Exchange an identity token for a Model API key. Keys are valid for about a day."""
+    response = await _common.fetch(
+        API_KEY_MINT_URL,
+        method="POST",
+        headers={
+            "Accept": "application/json",
+            "Authorization": f"Bearer {identity_token}",
+            "Content-Type": "application/json",
+            "x-api-version": "1.0.0",
+        },
+        body="{}",
+        signal=signal,
+        timeout_ms=REQUEST_TIMEOUT_MS,
+    )
+    data = _read_json(response)
+    if response.status_code in (401, 403):
+        # Identity token is not renewable (see module docstring); only a fresh device flow helps.
+        raise OAuthError(
+            f"Meta session expired (status {response.status_code}). Run `/login meta` to sign in again."
+            f"{_error_detail(data)}"
+        )
+    if not response.is_success:
+        raise OAuthError(f"Meta API key mint failed with status {response.status_code}{_error_detail(data)}")
+    api_key = _field(data, "api_key")
+    if not isinstance(api_key, str) or not api_key:
+        action_url = _trusted_http_url(_field(data, "action_url"))
+        raise OAuthError(f"Meta did not issue an API key.{f' Complete setup at {action_url}' if action_url else ''}")
+    return {
+        "type": "oauth",
+        "refresh": identity_token,
+        "access": api_key,
+        "expires": _common.now_ms() + API_KEY_LIFETIME_MS,
+    }
+
+
+async def _login_meta(interaction: Interaction) -> Credential:
+    try:
+        device = await _start_device_authorization(interaction.signal)
+        interaction.notify(
+            AuthEvent(
+                type="device_code",
+                user_code=device.user_code,
+                verification_uri=device.verification_uri,
+                interval_seconds=device.interval_seconds,
+                expires_in_seconds=device.expires_in_seconds,
+            )
+        )
+        identity_token = await _poll_for_identity_token(device, interaction.signal)
+        interaction.notify(AuthEvent(type="progress", message="Enabling Meta Model API access..."))
+        return await _mint_api_key(identity_token, interaction.signal)
+    except Exception as error:
+        # An in-flight request raises an AbortError on abort; the login UI matches on this message.
+        if interaction.signal.aborted:
+            raise _common.login_cancelled() from error
+        raise
+
+
+class MetaOAuth(OAuthAuth):
+    name = "Meta (Muse subscription)"
+    is_subscription = True
+    login_label = "Sign in with Meta"
+
+    async def login(self, interaction: Interaction, options: LoginOptions | None = None) -> Credential:
+        return await _login_meta(interaction)
+
+    async def refresh(self, credential: Credential, signal: AbortSignal) -> Credential:
+        return await _mint_api_key(credential["refresh"], signal)
+
+    async def to_auth(self, credential: Credential) -> ModelAuth:
+        return ModelAuth(api_key=credential["access"])
+
+
+meta_oauth = MetaOAuth()