diff --git a/CHANGELOG.md b/CHANGELOG.md index 6c5aa5fd..7ec70e3a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,12 @@ GitHub Releases page; `0.8.0` is the new starting line. ## Unreleased +- Add xAI Grok OAuth login (browser loopback and device-code). +- Add GitHub Copilot device-code OAuth login for individual github.com accounts. +- Add DigitalOcean Gradient AI browser OAuth login with dynamically discovered Inference Routers; router-discovery failures (unauthorized, outage, malformed, empty) are now reported distinctly instead of silently yielding no models. +- Add Snowflake Cortex account-scoped browser OAuth login (`pythinker login --snowflake`); the account identifier is validated before any request, and models register without becoming the default until Cortex chat support ships. +- Discover each provider's models from a cached, typed, provider-neutral models.dev catalog with curated fallbacks, so xAI, GitHub Copilot, and Snowflake pick up new models automatically; the catalog is now portable across platforms (no longer Unix-only) and distinguishes fresh, cached, stale, disabled, and unavailable results. +- Persist provider login and logout atomically so a failed save never leaves orphaned credentials or a half-applied configuration; a re-login whose save fails now restores the previous credential instead of deleting it, and provider persistence failures are logged. - Fix queued follow-up input showing a bordered ghost; pressing Enter during an active turn now shows one intentional queued row. ## 0.60.0 (2026-07-18) diff --git a/docs/en/release-notes/changelog.md b/docs/en/release-notes/changelog.md index 380fea1f..20eccaeb 100644 --- a/docs/en/release-notes/changelog.md +++ b/docs/en/release-notes/changelog.md @@ -17,6 +17,12 @@ GitHub Releases page; `0.8.0` is the new starting line. ## Unreleased +- Add xAI Grok OAuth login (browser loopback and device-code). +- Add GitHub Copilot device-code OAuth login for individual github.com accounts. +- Add DigitalOcean Gradient AI browser OAuth login with dynamically discovered Inference Routers; router-discovery failures (unauthorized, outage, malformed, empty) are now reported distinctly instead of silently yielding no models. +- Add Snowflake Cortex account-scoped browser OAuth login (`pythinker login --snowflake`); the account identifier is validated before any request, and models register without becoming the default until Cortex chat support ships. +- Discover each provider's models from a cached, typed, provider-neutral models.dev catalog with curated fallbacks, so xAI, GitHub Copilot, and Snowflake pick up new models automatically; the catalog is now portable across platforms (no longer Unix-only) and distinguishes fresh, cached, stale, disabled, and unavailable results. +- Persist provider login and logout atomically so a failed save never leaves orphaned credentials or a half-applied configuration; a re-login whose save fails now restores the previous credential instead of deleting it, and provider persistence failures are logged. - Fix queued follow-up input showing a bordered ghost; pressing Enter during an active turn now shows one intentional queued row. ## 0.60.0 (2026-07-18) diff --git a/docs/superpowers/specs/2026-07-18-pr-215-auth-review-fixes-design.md b/docs/superpowers/specs/2026-07-18-pr-215-auth-review-fixes-design.md new file mode 100644 index 00000000..c8b4cee3 --- /dev/null +++ b/docs/superpowers/specs/2026-07-18-pr-215-auth-review-fixes-design.md @@ -0,0 +1,69 @@ +# PR #215 auth review fixes design + +## Goal + +Resolve every validated review defect on PR #215 without broad auth or configuration refactoring. Preserve existing provider behavior except where it is unsafe, misleading, or violates the repository's failure-truthfulness and async-runtime contracts. + +## Scope + +1. Prevent nested OAuth references such as `oauth/snowflake-cortex/` from sharing credential or lock files with flat references such as `oauth/xai`. +2. Make the new login/logout persistence unit: + - serialized across processes at the shared config-file boundary; + - executed outside the asyncio event loop; + - safe against partial config-file writes; + - explicit when rollback or credential deletion fails; + - unable to report logout success while credentials remain undeleted. +3. Validate the implicit-flow `state` before accepting either success or error payloads. +4. Keep DigitalOcean's intentional degraded credential persistence, but never claim that an unrelated or empty default model is a configured DigitalOcean model. +5. Log curated-model fallback both when catalog status is non-authoritative and when an authoritative catalog yields no usable provider models. +6. Preserve the existing OpenCode Go unavailable-catalog regression test and resolve its stale review thread with evidence. + +Out of scope: provider protocol redesigns, replacing DigitalOcean's provider-required implicit grant, Snowflake Cortex chat transforms, global config merge semantics unrelated to these auth transactions, and new dependencies. + +## Design + +### Credential filenames + +Keep the legacy filename for ordinary one-segment OAuth keys so existing credentials remain readable. Encode the complete relative key for multi-segment keys into a deterministic, filesystem-safe filename, and apply the same mapping to lock files. This makes the mapping injective without changing released flat-key paths. + +### Persistence transaction + +Expose async persistence helpers to provider login/logout callers. Each helper offloads one complete synchronous transaction to a worker thread. The transaction acquires a bounded, fail-closed inter-process lock derived from the default config path before it snapshots credentials/config, mutates state, writes, or rolls back. + +Config writes use a temporary file in the target directory, flush and fsync it, apply private permissions, and atomically replace the destination. A failed write therefore leaves the previous config file intact. + +Login order remains token write then config write. On failure, restore the previous token (or remove the new one), restore the in-memory snapshot, and re-raise. If rollback also fails, raise an explicit persistence error carrying both failure contexts and log the rollback failure without secrets. + +Logout writes the config removal first, then deletes credentials. If deletion fails, restore and persist the previous config before raising. It must never emit success after failed credential deletion. + +The config-scoped lock covers snapshot through rollback, so a failed concurrent login cannot restore a token observed before another successful transaction. + +### OAuth callback validation + +For implicit callbacks, parse `state` and compare it with the expected value before interpreting `error`. RFC 6749 requires the original state on both successful and error responses. A wrong-state error is rejected as `OAuthStateMismatch`, not accepted as a user denial. + +### Provider messaging and fallback visibility + +DigitalOcean continues to persist a valid OAuth token when router discovery is empty or unavailable, matching the approved PR scope. Its terminal success text distinguishes “router configured” from “credentials saved with no routers configured.” + +Copilot, xAI, and Snowflake log when they use curated models because the catalog is non-authoritative or because authoritative conversion is empty. Logs include status/source only, never credentials. + +## Tests + +Use red-green TDD for each behavior: + +- nested Snowflake/xAI credential and lock paths differ; +- concurrent failed login cannot overwrite a successful token; +- failed config save preserves prior file bytes; +- rollback/delete failures are surfaced and logout does not report success; +- async callers do not run persistence I/O on the event-loop thread; +- wrong-state implicit error yields `OAuthStateMismatch`; +- DigitalOcean degraded success does not name an unrelated model; +- empty authoritative catalogs emit fallback logs; +- unavailable OpenCode Go catalog coverage remains green. + +After focused tests, run `make check-pythinker-code`, `make test-pythinker-code`, and `git diff --check`. + +## GitHub review completion + +Push the verified commit to `feat/auth-login-providers`. Reply inside each CodeRabbit thread with the specific fix and test evidence. For the already-satisfied OpenCode Go coverage thread, cite the existing unavailable-catalog test. Resolve threads only after the pushed head and checks reflect the fixes. diff --git a/src/pythinker_code/auth/__init__.py b/src/pythinker_code/auth/__init__.py index cdc5b77c..e0429b99 100644 --- a/src/pythinker_code/auth/__init__.py +++ b/src/pythinker_code/auth/__init__.py @@ -9,10 +9,14 @@ MOONSHOT_PLATFORM_ID = "moonshot" KIMI_PLATFORM_ID = "kimi" DEEPSEEK_PLATFORM_ID = "deepseek" +DIGITALOCEAN_PLATFORM_ID = "digitalocean" +GITHUB_COPILOT_PLATFORM_ID = "copilot" ANTHROPIC_PLATFORM_ID = "anthropic" OPENROUTER_PLATFORM_ID = "openrouter" +SNOWFLAKE_CORTEX_PLATFORM_ID = "snowflake-cortex" LM_STUDIO_PLATFORM_ID = "lm-studio" OLLAMA_PLATFORM_ID = "ollama" +XAI_PLATFORM_ID = "xai" ZAI_CODING_PLATFORM_ID = "z-ai-coding" ZAI_API_PLATFORM_ID = "z-ai-api" @@ -20,6 +24,8 @@ "ALIBABA_PLATFORM_ID", "ANTHROPIC_PLATFORM_ID", "DEEPSEEK_PLATFORM_ID", + "DIGITALOCEAN_PLATFORM_ID", + "GITHUB_COPILOT_PLATFORM_ID", "KIMI_PLATFORM_ID", "LM_STUDIO_PLATFORM_ID", "MINIMAX_PLATFORM_ID", @@ -30,6 +36,8 @@ "OPENCODE_GO_PLATFORM_ID", "OPENROUTER_PLATFORM_ID", "PYTHINKER_CODE_PLATFORM_ID", + "SNOWFLAKE_CORTEX_PLATFORM_ID", + "XAI_PLATFORM_ID", "ZAI_API_PLATFORM_ID", "ZAI_CODING_PLATFORM_ID", ] diff --git a/src/pythinker_code/auth/copilot.py b/src/pythinker_code/auth/copilot.py new file mode 100644 index 00000000..d792b7fa --- /dev/null +++ b/src/pythinker_code/auth/copilot.py @@ -0,0 +1,305 @@ +from __future__ import annotations + +from collections.abc import AsyncIterator +from dataclasses import dataclass +from typing import Any, cast + +import aiohttp +from pydantic import SecretStr + +from pythinker_code.auth import GITHUB_COPILOT_PLATFORM_ID +from pythinker_code.auth.models_dev import ( + CatalogModel, + build_catalog_models, + get_models_dev_catalog, +) +from pythinker_code.auth.oauth import ( + OAuthError, + OAuthEvent, + OAuthToken, + OAuthUnauthorized, + persist_login, + persist_logout, +) +from pythinker_code.auth.oauth_flows import poll_device_token, request_device_code +from pythinker_code.auth.platforms import managed_model_key, managed_provider_key +from pythinker_code.config import Config, LLMModel, LLMProvider, OAuthRef +from pythinker_code.thinking import apply_login_thinking_defaults +from pythinker_code.utils.aiohttp import new_client_session +from pythinker_code.utils.logging import logger + +GITHUB_COPILOT_CLIENT_ID = "Iv1.b507a08c87ecfe98" +GITHUB_COPILOT_SCOPE = "read:user" +GITHUB_DEVICE_CODE_ENDPOINT = "https://github.com/login/device/code" +GITHUB_TOKEN_ENDPOINT = "https://github.com/login/oauth/access_token" +GITHUB_COPILOT_TOKEN_ENDPOINT = "https://api.github.com/copilot_internal/v2/token" +GITHUB_COPILOT_OAUTH_KEY = "oauth/github-copilot" +GITHUB_COPILOT_PROVIDER_KEY = managed_provider_key(GITHUB_COPILOT_PLATFORM_ID) +COPILOT_BASE_URL = "https://api.githubcopilot.com" +GITHUB_COPILOT_MODELS_DEV_PROVIDER_ID = "github-copilot" +GITHUB_COPILOT_DEFAULT_CONTEXT = 128_000 + +GITHUB_JSON_HEADERS = {"Accept": "application/json"} +_EDITOR_VERSION = "vscode/1.99.0" +_EDITOR_PLUGIN_VERSION = "copilot-chat/0.26.7" +_USER_AGENT = "GitHubCopilotChat/0.26.7" +_GITHUB_API_VERSION = "2025-04-01" + + +def build_copilot_headers() -> dict[str, str]: + return { + "Copilot-Integration-Id": "vscode-chat", + "Editor-Version": _EDITOR_VERSION, + "Editor-Plugin-Version": _EDITOR_PLUGIN_VERSION, + "User-Agent": _USER_AGENT, + "X-GitHub-Api-Version": _GITHUB_API_VERSION, + "Openai-Intent": "conversation-panel", + } + + +def _build_exchange_headers(github_token: str) -> dict[str, str]: + return { + "Authorization": f"token {github_token}", + "Editor-Version": _EDITOR_VERSION, + "Editor-Plugin-Version": _EDITOR_PLUGIN_VERSION, + "User-Agent": _USER_AGENT, + "X-GitHub-Api-Version": _GITHUB_API_VERSION, + } + + +@dataclass(frozen=True, slots=True) +class GitHubCopilotModel: + model_id: str + display_name: str + max_context_size: int + + @property + def alias(self) -> str: + return managed_model_key(GITHUB_COPILOT_PLATFORM_ID, self.model_id) + + +GITHUB_COPILOT_MODELS: tuple[GitHubCopilotModel, ...] = ( + GitHubCopilotModel("gpt-4.1", "GPT-4.1", 1_000_000), + GitHubCopilotModel("gpt-4o", "GPT-4o", 128_000), + GitHubCopilotModel("o4-mini", "o4-mini", 200_000), +) + + +def _copilot_oauth_ref() -> OAuthRef: + return OAuthRef(storage="file", key=GITHUB_COPILOT_OAUTH_KEY) + + +def _apply_copilot_config( + config: Config, models: tuple[GitHubCopilotModel, ...] = GITHUB_COPILOT_MODELS +) -> None: + oauth_ref = _copilot_oauth_ref() + config.providers[GITHUB_COPILOT_PROVIDER_KEY] = LLMProvider( + type="openai_legacy", + base_url=COPILOT_BASE_URL, + api_key=SecretStr(""), + oauth=oauth_ref, + custom_headers=build_copilot_headers(), + ) + + for alias, model in list(config.models.items()): + if model.provider == GITHUB_COPILOT_PROVIDER_KEY: + del config.models[alias] + + for model in models: + config.models[model.alias] = LLMModel( + provider=GITHUB_COPILOT_PROVIDER_KEY, + model=model.model_id, + max_context_size=model.max_context_size, + display_name=model.display_name, + ) + + if models: + config.default_model = models[0].alias + apply_login_thinking_defaults(config, thinking=False, effort="off") + + +def _catalog_models_to_copilot(built: tuple[CatalogModel, ...]) -> tuple[GitHubCopilotModel, ...]: + return tuple( + GitHubCopilotModel(model.model_id, model.display_name, model.max_context_size) + for model in built + ) + + +async def _discover_copilot_models() -> tuple[GitHubCopilotModel, ...]: + """Resolve the live Copilot model list from the shared models.dev catalog. + + Falls back to the curated list when the catalog is unavailable or degraded. + """ + result = await get_models_dev_catalog() + if not result.is_authoritative: + logger.debug( + "models.dev catalog not authoritative (status={status}, source={source}); " + "using curated GitHub Copilot models.", + status=result.status, + source=result.source, + ) + return GITHUB_COPILOT_MODELS + built = build_catalog_models( + result.catalog, + GITHUB_COPILOT_MODELS_DEV_PROVIDER_ID, + default_context=GITHUB_COPILOT_DEFAULT_CONTEXT, + ) + converted = _catalog_models_to_copilot(built) + if converted: + return converted + logger.debug( + "models.dev catalog contained no usable GitHub Copilot models " + "(status={status}, source={source}); using curated models.", + status=result.status, + source=result.source, + ) + return GITHUB_COPILOT_MODELS + + +async def refresh_copilot_token(github_token: str) -> OAuthToken: + try: + async with ( + new_client_session() as session, + session.get( + GITHUB_COPILOT_TOKEN_ENDPOINT, + headers=_build_exchange_headers(github_token), + ) as response, + ): + status = response.status + try: + payload_any: Any = await response.json(content_type=None) + except ValueError: + payload_any = None + except (aiohttp.ClientError, TimeoutError, OSError) as exc: + raise OAuthError("GitHub Copilot token exchange request failed.") from exc + + if status in {401, 403}: + raise OAuthUnauthorized("GitHub Copilot token exchange was unauthorized.") + if status != 200: + raise OAuthError(f"GitHub Copilot token exchange failed (HTTP {status}).") + if not isinstance(payload_any, dict): + raise OAuthError("GitHub Copilot token exchange returned an invalid response.") + + payload = cast(dict[str, Any], payload_any) + bearer = payload.get("token") + expires_at = payload.get("expires_at") + refresh_in = payload.get("refresh_in") + if ( + not isinstance(bearer, str) + or not bearer + or isinstance(expires_at, bool) + or not isinstance(expires_at, int) + or expires_at <= 0 + or isinstance(refresh_in, bool) + or not isinstance(refresh_in, int) + or refresh_in <= 0 + ): + raise OAuthError("GitHub Copilot token exchange returned an incomplete response.") + + return OAuthToken( + access_token=bearer, + refresh_token=github_token, + expires_at=float(expires_at), + scope=GITHUB_COPILOT_SCOPE, + token_type="Bearer", + expires_in=float(refresh_in), + ) + + +async def login_copilot(config: Config, *, open_browser: bool = True) -> AsyncIterator[OAuthEvent]: + if not config.is_from_default_location: + yield OAuthEvent( + "error", + "Login requires the default config file; restart without --config/--config-file.", + ) + return + + try: + device_code = await request_device_code( + device_authorization_endpoint=GITHUB_DEVICE_CODE_ENDPOINT, + client_id=GITHUB_COPILOT_CLIENT_ID, + scope=GITHUB_COPILOT_SCOPE, + headers=GITHUB_JSON_HEADERS, + ) + except OAuthError as exc: + yield OAuthEvent("error", f"Failed to start GitHub Copilot login: {exc}") + return + + yield OAuthEvent( + "verification_url", + f"Open {device_code.verification_uri} and enter code {device_code.user_code}.", + data={ + "verification_url": device_code.verification_uri, + "user_code": device_code.user_code, + }, + ) + if open_browser: + try: + from pythinker_code.utils.term import open_url_in_browser + + open_url_in_browser(device_code.verification_uri) + except Exception as exc: + logger.warning("Failed to open browser: {error}", error=exc) + yield OAuthEvent("waiting", "Waiting for GitHub Copilot authorization...") + + try: + payload = await poll_device_token( + token_endpoint=GITHUB_TOKEN_ENDPOINT, + client_id=GITHUB_COPILOT_CLIENT_ID, + device_code=device_code, + headers=GITHUB_JSON_HEADERS, + ) + github_token = payload.get("access_token") + if not isinstance(github_token, str) or not github_token: + raise OAuthError("GitHub device token response was incomplete.") + copilot_token = await refresh_copilot_token(github_token) + except OAuthError as exc: + yield OAuthEvent("error", f"GitHub Copilot login failed: {exc}") + return + + token = OAuthToken( + access_token=copilot_token.access_token, + refresh_token=github_token, + expires_at=copilot_token.expires_at, + scope=copilot_token.scope, + token_type=copilot_token.token_type, + expires_in=copilot_token.expires_in, + ) + models = await _discover_copilot_models() + try: + await persist_login( + config, _copilot_oauth_ref(), token, lambda cfg: _apply_copilot_config(cfg, models) + ) + except Exception as exc: + logger.warning("Failed to persist GitHub Copilot login: {exc}", exc=exc) + yield OAuthEvent("error", "Failed to save GitHub Copilot login.") + return + yield OAuthEvent( + "success", + f"GitHub Copilot configured with model {config.default_model}.", + ) + + +async def logout_copilot(config: Config) -> AsyncIterator[OAuthEvent]: + if not config.is_from_default_location: + yield OAuthEvent( + "error", + "Logout requires the default config file; restart without --config/--config-file.", + ) + return + + def _remove(cfg: Config) -> None: + cfg.providers.pop(GITHUB_COPILOT_PROVIDER_KEY, None) + for alias, model in list(cfg.models.items()): + if model.provider == GITHUB_COPILOT_PROVIDER_KEY: + del cfg.models[alias] + if cfg.default_model not in cfg.models: + cfg.default_model = next(iter(cfg.models), "") + + try: + await persist_logout(config, _copilot_oauth_ref(), _remove) + except Exception as exc: + logger.warning("Failed to persist GitHub Copilot logout: {exc}", exc=exc) + yield OAuthEvent("error", "Failed to log out of GitHub Copilot.") + return + yield OAuthEvent("success", "Logged out of GitHub Copilot successfully.") diff --git a/src/pythinker_code/auth/digitalocean.py b/src/pythinker_code/auth/digitalocean.py new file mode 100644 index 00000000..8f949da8 --- /dev/null +++ b/src/pythinker_code/auth/digitalocean.py @@ -0,0 +1,227 @@ +from __future__ import annotations + +from collections.abc import AsyncIterator +from dataclasses import dataclass +from enum import Enum +from typing import Any, cast + +import aiohttp +from pydantic import SecretStr + +from pythinker_code.auth import DIGITALOCEAN_PLATFORM_ID +from pythinker_code.auth.oauth import OAuthError, OAuthEvent, persist_config_change +from pythinker_code.auth.oauth_flows import run_loopback_implicit_flow +from pythinker_code.auth.platforms import managed_model_key, managed_provider_key +from pythinker_code.config import Config, LLMModel, LLMProvider +from pythinker_code.thinking import apply_login_thinking_defaults +from pythinker_code.utils.aiohttp import new_client_session +from pythinker_code.utils.logging import logger + +DIGITALOCEAN_CLIENT_ID = "b1a6c5158156caac821fd1b30253ca8acb52454a48fa744420e41889cb589f82" +DIGITALOCEAN_AUTHORIZE_URL = "https://cloud.digitalocean.com/v1/oauth/authorize" +DIGITALOCEAN_SCOPE = "genai:read inference:query" +DIGITALOCEAN_BASE_URL = "https://inference.do-ai.run/v1" +DIGITALOCEAN_ROUTERS_URL = "https://api.digitalocean.com/v2/gen-ai/models/routers" +DIGITALOCEAN_REDIRECT_PORT = 1456 +DIGITALOCEAN_CALLBACK_PATH = "/auth/callback" +DIGITALOCEAN_TOKEN_PATH = "/auth/token" +DIGITALOCEAN_PROVIDER_KEY = managed_provider_key(DIGITALOCEAN_PLATFORM_ID) +DIGITALOCEAN_DEFAULT_CONTEXT = 128_000 + + +class RouterDiscovery(str, Enum): + """Outcome of a DigitalOcean inference-router discovery call.""" + + OK = "ok" # routers were discovered + EMPTY = "empty" # the account has no routers (valid, but empty) + UNAUTHORIZED = "unauthorized" # the token could not list routers + UNAVAILABLE = "unavailable" # timeout / outage / non-2xx response + MALFORMED = "malformed" # the response body was not the expected shape + + +@dataclass(frozen=True, slots=True) +class RouterCatalog: + """Discovered routers paired with the outcome that produced them.""" + + status: RouterDiscovery + names: tuple[str, ...] + + +def _skip_browser_open(_url: str) -> None: + return None + + +def _apply_digitalocean_config( + config: Config, + api_key: SecretStr, + router_names: tuple[str, ...], +) -> None: + config.providers[DIGITALOCEAN_PROVIDER_KEY] = LLMProvider( + type="openai_legacy", + base_url=DIGITALOCEAN_BASE_URL, + api_key=api_key, + ) + + for alias, model in list(config.models.items()): + if model.provider == DIGITALOCEAN_PROVIDER_KEY: + del config.models[alias] + + for name in router_names: + model_id = f"router:{name}" + alias = managed_model_key(DIGITALOCEAN_PLATFORM_ID, model_id) + config.models[alias] = LLMModel( + provider=DIGITALOCEAN_PROVIDER_KEY, + model=model_id, + max_context_size=DIGITALOCEAN_DEFAULT_CONTEXT, + display_name=name, + ) + + if router_names: + config.default_model = managed_model_key( + DIGITALOCEAN_PLATFORM_ID, + f"router:{router_names[0]}", + ) + elif config.default_model not in config.models: + config.default_model = next(iter(config.models), "") + apply_login_thinking_defaults(config, thinking=False, effort="off") + + +def _parse_router_names(payload: object) -> RouterCatalog: + if not isinstance(payload, dict): + return RouterCatalog(RouterDiscovery.MALFORMED, ()) + raw = cast(dict[str, Any], payload).get("model_routers") + if not isinstance(raw, list): + return RouterCatalog(RouterDiscovery.MALFORMED, ()) + + names: list[str] = [] + for item in cast(list[Any], raw): + if isinstance(item, dict): + name = cast(dict[str, Any], item).get("name") + if isinstance(name, str) and name: + names.append(name) + if not names: + return RouterCatalog(RouterDiscovery.EMPTY, ()) + return RouterCatalog(RouterDiscovery.OK, tuple(names)) + + +async def _fetch_router_catalog(access_token: str) -> RouterCatalog: + """Discover the account's inference routers, preserving distinct outcomes. + + Authentication failures, dependency outages, malformed responses, and a + valid-but-empty catalog are each reported separately so the login flow can + surface an accurate message instead of collapsing everything into "no + routers". + """ + try: + async with ( + new_client_session() as session, + session.get( + DIGITALOCEAN_ROUTERS_URL, + headers={ + "Authorization": f"Bearer {access_token}", + "Accept": "application/json", + }, + ) as response, + ): + status = response.status + if status in (401, 403): + return RouterCatalog(RouterDiscovery.UNAUTHORIZED, ()) + if not 200 <= status < 300: + return RouterCatalog(RouterDiscovery.UNAVAILABLE, ()) + try: + payload: Any = await response.json(content_type=None) + except (ValueError, aiohttp.ClientError): + return RouterCatalog(RouterDiscovery.MALFORMED, ()) + except (TimeoutError, aiohttp.ClientError, OSError): + return RouterCatalog(RouterDiscovery.UNAVAILABLE, ()) + + return _parse_router_names(payload) + + +def _router_status_message(status: RouterDiscovery) -> str | None: + if status is RouterDiscovery.OK: + return None + if status is RouterDiscovery.EMPTY: + return "DigitalOcean returned no inference routers; sign-in saved with no models." + if status is RouterDiscovery.UNAUTHORIZED: + return ( + "DigitalOcean did not authorize inference-router discovery; sign-in saved. " + "Re-run login once the account has inference access." + ) + return ( + "DigitalOcean Inference Routers were unavailable; sign-in saved. " + "Re-run login to load routers." + ) + + +async def login_digitalocean( + config: Config, *, open_browser: bool = True +) -> AsyncIterator[OAuthEvent]: + if not config.is_from_default_location: + yield OAuthEvent( + "error", + "Login requires the default config file; restart without --config/--config-file.", + ) + return + + yield OAuthEvent("waiting", "Waiting for DigitalOcean browser authorization...") + try: + auth = await run_loopback_implicit_flow( + authorize_endpoint=DIGITALOCEAN_AUTHORIZE_URL, + client_id=DIGITALOCEAN_CLIENT_ID, + scope=DIGITALOCEAN_SCOPE, + callback_path=DIGITALOCEAN_CALLBACK_PATH, + token_path=DIGITALOCEAN_TOKEN_PATH, + port=DIGITALOCEAN_REDIRECT_PORT, + browser_open=None if open_browser else _skip_browser_open, + ) + except OAuthError as exc: + yield OAuthEvent("error", f"DigitalOcean browser login failed: {exc}") + return + + catalog = await _fetch_router_catalog(auth.access_token) + try: + await persist_config_change( + config, + lambda cfg: _apply_digitalocean_config( + cfg, SecretStr(auth.access_token), catalog.names + ), + ) + except Exception as exc: + logger.warning("Failed to persist DigitalOcean login: {exc}", exc=exc) + yield OAuthEvent("error", "Failed to save DigitalOcean login.") + return + + message = _router_status_message(catalog.status) + if message: + yield OAuthEvent("info", message) + if catalog.names: + success_message = f"DigitalOcean configured with model {config.default_model}." + else: + success_message = "DigitalOcean credentials saved; no inference routers are configured." + yield OAuthEvent("success", success_message) + + +async def logout_digitalocean(config: Config) -> AsyncIterator[OAuthEvent]: + if not config.is_from_default_location: + yield OAuthEvent( + "error", + "Logout requires the default config file; restart without --config/--config-file.", + ) + return + + def _remove(cfg: Config) -> None: + cfg.providers.pop(DIGITALOCEAN_PROVIDER_KEY, None) + for alias, model in list(cfg.models.items()): + if model.provider == DIGITALOCEAN_PROVIDER_KEY: + del cfg.models[alias] + if cfg.default_model not in cfg.models: + cfg.default_model = next(iter(cfg.models), "") + + try: + await persist_config_change(config, _remove) + except Exception as exc: + logger.warning("Failed to persist DigitalOcean logout: {exc}", exc=exc) + yield OAuthEvent("error", "Failed to log out of DigitalOcean.") + return + yield OAuthEvent("success", "Logged out of DigitalOcean successfully.") diff --git a/src/pythinker_code/auth/models_dev.py b/src/pythinker_code/auth/models_dev.py new file mode 100644 index 00000000..a2a992d1 --- /dev/null +++ b/src/pythinker_code/auth/models_dev.py @@ -0,0 +1,454 @@ +from __future__ import annotations + +import asyncio +import contextlib +import json +import os +import tempfile +import time +from collections.abc import Mapping +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from typing import IO, Any, cast + +import httpx + +from pythinker_code.share import get_share_dir +from pythinker_code.utils.logging import logger + +# Provider-agnostic model catalog backed by the public models.dev API. It +# supplies dynamic per-model metadata (display name, context window, reasoning +# and modality flags) for any provider, with on-disk caching and fail-open +# behavior so a slow or unreachable endpoint never blocks startup. Callers get +# a CatalogResult that names the outcome (fresh, cached, stale, disabled, or +# unavailable) so degraded data is never mistaken for authoritative success. + +MODELS_DEV_BASE_URL = "https://models.dev" +MODELS_DEV_CACHE_TTL_SECONDS = 5 * 60 +MODELS_DEV_TIMEOUT_SECONDS = 10.0 +MODELS_DEV_USER_AGENT = "pythinker-code models.dev catalog" +_CACHE_FILENAME = "models_dev.json" +_LOCK_FILENAME = "models_dev.lock" +_FETCH_ATTEMPTS = 2 +_RETRY_BACKOFF_SECONDS = 0.25 +_LOCK_TIMEOUT_SECONDS = 10.0 +_LOCK_POLL_INTERVAL_SECONDS = 0.05 + +# models.dev has no explicit "chat" flag; a text-only output modality is the +# discriminator that excludes image/video/audio/embedding generators. A few +# text-output-but-non-chat models (rerankers, embeddings, moderation) slip past +# that test, so they are additionally excluded by id substring. +_NON_CHAT_ID_SUBSTRINGS = ( + "embedding", + "embed", + "reranker", + "rerank", + "moderation", + "guard", + "whisper", + "tts", + "stt", + "ocr", +) + + +class CatalogStatus(str, Enum): + """Outcome of a catalog load, so callers/logs can tell states apart.""" + + OK = "ok" # fresh data fetched from the network this call + CACHED = "cached" # served from a fresh on-disk cache + STALE = "stale" # served from a stale cache (fetch failed or disabled) + DISABLED = "disabled" # fetching disabled and no cache exists + UNAVAILABLE = "unavailable" # no data at all (fetch failed, no cache) + + +@dataclass(frozen=True, slots=True) +class ModelsDevModel: + provider_id: str + model_id: str + display_name: str + context_length: int | None + reasoning: bool + supports_image_input: bool + supports_video_input: bool + output_modalities: tuple[str, ...] = () + npm: str | None = None + + @property + def is_chat_model(self) -> bool: + """True when the model produces only text output (usable as a chat model).""" + return bool(self.output_modalities) and all( + modality == "text" for modality in self.output_modalities + ) + + +@dataclass(frozen=True, slots=True) +class ModelsDevProvider: + provider_id: str + display_name: str + env: tuple[str, ...] + npm: str | None + api: str | None + models: Mapping[str, ModelsDevModel] + + +ModelsDevCatalog = dict[str, ModelsDevProvider] + + +@dataclass(frozen=True, slots=True) +class CatalogResult: + """A catalog load paired with its outcome and origin.""" + + catalog: ModelsDevCatalog + status: CatalogStatus + source: str # "network" | "cache" | "none" + + @property + def is_authoritative(self) -> bool: + """True only for a fresh network fetch or a fresh cache hit.""" + return self.status in (CatalogStatus.OK, CatalogStatus.CACHED) + + +@dataclass(frozen=True, slots=True) +class CatalogModel: + """A provider-neutral chat model resolved from the catalog.""" + + model_id: str + display_name: str + max_context_size: int + reasoning: bool + + +def _optional_string(value: object) -> str | None: + return value if isinstance(value, str) and value else None + + +def _string_tuple(value: object) -> tuple[str, ...]: + if not isinstance(value, list): + return () + return tuple(item for item in cast(list[object], value) if isinstance(item, str)) + + +def _parse_catalog(payload: object) -> ModelsDevCatalog: + if not isinstance(payload, dict): + raise ValueError("models.dev catalog must be a JSON object") + + catalog: ModelsDevCatalog = {} + for provider_id, raw_provider in cast(dict[str, Any], payload).items(): + if not isinstance(raw_provider, dict): + continue + provider = cast(dict[str, Any], raw_provider) + raw_models = provider.get("models") + if not isinstance(raw_models, dict): + continue + + provider_npm = _optional_string(provider.get("npm")) + models: dict[str, ModelsDevModel] = {} + for model_id, raw_model in cast(dict[str, Any], raw_models).items(): + if not model_id or not isinstance(raw_model, dict): + continue + model = cast(dict[str, Any], raw_model) + display_name = _optional_string(model.get("name")) or model_id + raw_limit = model.get("limit") + raw_context = ( + cast(dict[str, Any], raw_limit).get("context") + if isinstance(raw_limit, dict) + else None + ) + context_length = ( + raw_context + if isinstance(raw_context, int) + and not isinstance(raw_context, bool) + and raw_context > 0 + else None + ) + raw_modalities = model.get("modalities") + modalities = ( + cast(dict[str, Any], raw_modalities) if isinstance(raw_modalities, dict) else {} + ) + input_modalities = set(_string_tuple(modalities.get("input"))) + output_modalities = _string_tuple(modalities.get("output")) + raw_model_provider = model.get("provider") + model_npm = ( + _optional_string(cast(dict[str, Any], raw_model_provider).get("npm")) + if isinstance(raw_model_provider, dict) + else None + ) + models[model_id] = ModelsDevModel( + provider_id=provider_id, + model_id=model_id, + display_name=display_name, + context_length=context_length, + reasoning=model.get("reasoning") is True, + supports_image_input="image" in input_modalities, + supports_video_input="video" in input_modalities, + output_modalities=output_modalities, + npm=model_npm or provider_npm, + ) + + env = _string_tuple(provider.get("env")) + catalog[provider_id] = ModelsDevProvider( + provider_id=provider_id, + display_name=_optional_string(provider.get("name")) or provider_id, + env=env, + npm=provider_npm, + api=_optional_string(provider.get("api")), + models=models, + ) + return catalog + + +def parse_models_dev_catalog(payload: object) -> ModelsDevCatalog: + """Normalize a models.dev JSON object, returning an empty catalog if invalid.""" + try: + return _parse_catalog(payload) + except (TypeError, ValueError): + return {} + + +def get_provider_models( + catalog: Mapping[str, ModelsDevProvider], provider_id: str +) -> Mapping[str, ModelsDevModel]: + """Return normalized models for one provider, or an empty mapping.""" + provider = catalog.get(provider_id) + return provider.models if provider is not None else {} + + +def _is_non_chat_id(model_id: str) -> bool: + normalized = model_id.lower() + return any(token in normalized for token in _NON_CHAT_ID_SUBSTRINGS) + + +def chat_models( + catalog: Mapping[str, ModelsDevProvider], provider_id: str +) -> dict[str, ModelsDevModel]: + """Return the provider's chat-capable models (text output, non-embedding).""" + return { + model_id: model + for model_id, model in get_provider_models(catalog, provider_id).items() + if model.is_chat_model and not _is_non_chat_id(model_id) + } + + +def build_catalog_models( + catalog: Mapping[str, ModelsDevProvider], + provider_id: str, + *, + default_context: int, +) -> tuple[CatalogModel, ...]: + """Resolve a provider's chat models into provider-neutral entries. + + Filters to text-output chat models, applies ``default_context`` when the + catalog has no context window, and returns a deterministically ordered + tuple so provider config stays stable across runs. + """ + resolved = [ + CatalogModel( + model_id=model_id, + display_name=model.display_name, + max_context_size=model.context_length or default_context, + reasoning=model.reasoning, + ) + for model_id, model in sorted(chat_models(catalog, provider_id).items()) + ] + return tuple(resolved) + + +def _cache_is_fresh(path: Path) -> bool: + try: + return time.time() - path.stat().st_mtime < MODELS_DEV_CACHE_TTL_SECONDS + except OSError: + return False + + +def _read_cache(path: Path) -> ModelsDevCatalog | None: + try: + with path.open(encoding="utf-8") as cache_file: + payload = json.load(cache_file) + return _parse_catalog(payload) + except (OSError, TypeError, ValueError): + return None + + +def _write_cache(path: Path, payload: object) -> None: + temporary_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + delete=False, + ) as temporary_file: + temporary_path = Path(temporary_file.name) + json.dump(payload, temporary_file, ensure_ascii=False, separators=(",", ":")) + temporary_file.flush() + os.fsync(temporary_file.fileno()) + os.replace(temporary_path, path) + finally: + if temporary_path is not None and temporary_path.exists(): + temporary_path.unlink(missing_ok=True) + + +def _fetch_disabled() -> bool: + value = os.getenv("PYTHINKER_DISABLE_MODELS_FETCH", "").strip().lower() + return value not in {"", "0", "false", "no", "off"} + + +def _catalog_url() -> str: + base_url = os.getenv("PYTHINKER_MODELS_URL", MODELS_DEV_BASE_URL).rstrip("/") + return f"{base_url}/api.json" + + +async def _fetch_catalog_payload() -> object: + last_error: Exception | None = None + async with httpx.AsyncClient( + headers={"User-Agent": MODELS_DEV_USER_AGENT}, + timeout=MODELS_DEV_TIMEOUT_SECONDS, + follow_redirects=True, + ) as client: + for attempt in range(_FETCH_ATTEMPTS): + try: + response = await client.get(_catalog_url()) + response.raise_for_status() + payload = response.json() + _parse_catalog(payload) + return payload + except (httpx.HTTPError, TypeError, ValueError) as exc: + last_error = exc + if attempt + 1 < _FETCH_ATTEMPTS: + await asyncio.sleep(_RETRY_BACKOFF_SECONDS * (attempt + 1)) + if last_error is None: + raise RuntimeError("models.dev fetch failed without an error") + raise last_error + + +def _try_acquire_lock(lock_file: IO[str]) -> bool: + """Attempt a non-blocking exclusive lock; True if acquired, False if held.""" + if os.name == "nt": + import msvcrt + + try: + lock_file.seek(0) + msvcrt.locking(lock_file.fileno(), msvcrt.LK_NBLCK, 1) + return True + except OSError: + return False + + import fcntl + + try: + fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + return True + except OSError: + return False + + +def _release_lock(lock_file: IO[str]) -> None: + try: + if os.name == "nt": + import msvcrt + + lock_file.seek(0) + with contextlib.suppress(OSError): + msvcrt.locking(lock_file.fileno(), msvcrt.LK_UNLCK, 1) + else: + import fcntl + + fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) + finally: + lock_file.close() + + +async def _acquire_lock_bounded(path: Path) -> IO[str] | None: + """Acquire the cross-process lock with a bounded, cancellable wait. + + Returns the locked handle, or None if the deadline elapsed or the lock file + could not be opened. The lock is polled non-blockingly with an ``asyncio`` + sleep between attempts, so cancellation is honored promptly and the handle + is always closed on exit — no descriptor or cross-process lock is orphaned. + """ + loop = asyncio.get_running_loop() + deadline = loop.time() + _LOCK_TIMEOUT_SECONDS + try: + lock_file = path.open("a", encoding="utf-8") + except OSError: + return None + try: + while True: + if _try_acquire_lock(lock_file): + return lock_file + if loop.time() >= deadline: + lock_file.close() + return None + await asyncio.sleep(_LOCK_POLL_INTERVAL_SECONDS) + except BaseException: + lock_file.close() + raise + + +def _fallback_result(cached: ModelsDevCatalog | None) -> CatalogResult: + if cached is not None: + return CatalogResult(cached, CatalogStatus.STALE, "cache") + return CatalogResult({}, CatalogStatus.UNAVAILABLE, "none") + + +async def _load_catalog() -> CatalogResult: + share_dir = get_share_dir() + cache_path = share_dir / _CACHE_FILENAME + cached = _read_cache(cache_path) + if cached is not None and _cache_is_fresh(cache_path): + return CatalogResult(cached, CatalogStatus.CACHED, "cache") + if _fetch_disabled(): + if cached is not None: + return CatalogResult(cached, CatalogStatus.STALE, "cache") + return CatalogResult({}, CatalogStatus.DISABLED, "none") + + lock_file = await _acquire_lock_bounded(share_dir / _LOCK_FILENAME) + if lock_file is None: + logger.debug("models.dev catalog lock unavailable; using fallback data") + return _fallback_result(cached) + try: + locked_cache = _read_cache(cache_path) + if locked_cache is not None and _cache_is_fresh(cache_path): + return CatalogResult(locked_cache, CatalogStatus.CACHED, "cache") + fallback = locked_cache if locked_cache is not None else cached + try: + payload = await _fetch_catalog_payload() + catalog = _parse_catalog(payload) + except Exception as exc: + logger.debug("models.dev catalog fetch failed: {error}", error=exc) + return _fallback_result(fallback) + try: + _write_cache(cache_path, payload) + except OSError as exc: + logger.debug("models.dev catalog cache write failed: {error}", error=exc) + return CatalogResult(catalog, CatalogStatus.OK, "network") + finally: + _release_lock(lock_file) + + +async def get_models_dev_catalog() -> CatalogResult: + """Return the catalog with its status/source, failing open to empty data.""" + try: + return await _load_catalog() + except Exception as exc: + logger.debug("models.dev catalog unavailable: {error}", error=exc) + return CatalogResult({}, CatalogStatus.UNAVAILABLE, "none") + + +__all__ = [ + "CatalogModel", + "CatalogResult", + "CatalogStatus", + "ModelsDevCatalog", + "ModelsDevModel", + "ModelsDevProvider", + "build_catalog_models", + "chat_models", + "get_models_dev_catalog", + "get_provider_models", + "parse_models_dev_catalog", +] diff --git a/src/pythinker_code/auth/oauth.py b/src/pythinker_code/auth/oauth.py index 6ed958a1..c5552f9c 100644 --- a/src/pythinker_code/auth/oauth.py +++ b/src/pythinker_code/auth/oauth.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import base64 import json import os import platform @@ -10,7 +11,7 @@ import tempfile import time import uuid -from collections.abc import AsyncGenerator, AsyncIterator +from collections.abc import AsyncGenerator, AsyncIterator, Callable from contextlib import asynccontextmanager, suppress from dataclasses import dataclass from pathlib import Path @@ -35,12 +36,14 @@ OAuthRef, PythinkerAIFetchConfig, PythinkerAISearchConfig, + get_config_file, save_config, ) from pythinker_code.constant import VERSION from pythinker_code.share import get_share_dir from pythinker_code.thinking import apply_login_thinking_defaults from pythinker_code.utils.aiohttp import new_client_session +from pythinker_code.utils.io import file_lock from pythinker_code.utils.logging import logger if TYPE_CHECKING: @@ -74,6 +77,10 @@ class OAuthUnauthorized(OAuthError): """OAuth credentials rejected.""" +class OAuthPersistenceError(OAuthError): + """OAuth credentials and configuration could not be persisted consistently.""" + + class _RetryableRefreshError(OAuthError): """Transient HTTP error during token refresh (5xx / 429).""" @@ -297,14 +304,22 @@ def _credentials_dir() -> Path: return path +def _credential_file_stem(key: str) -> str: + relative_key = key.removeprefix("oauth/") + if "/" not in relative_key: + return relative_key or key + encoded = base64.urlsafe_b64encode(relative_key.encode(encoding="utf-8")).decode( + encoding="utf-8" + ) + return f"v2-{encoded.rstrip('=')}" + + def _credentials_path(key: str) -> Path: - name = key.removeprefix("oauth/").split("/")[-1] or key - return _credentials_dir() / f"{name}.json" + return _credentials_dir() / f"{_credential_file_stem(key)}.json" def _credentials_lock_path(key: str) -> Path: - name = key.removeprefix("oauth/").split("/")[-1] or key - return _credentials_dir() / f"{name}.lock" + return _credentials_dir() / f"{_credential_file_stem(key)}.lock" class _CrossProcessLock: @@ -488,6 +503,114 @@ def delete_tokens(ref: OAuthRef) -> None: _delete_from_file(ref.key) +def restore_config_state(config: Config, snapshot: Config) -> None: + """Restore every field of ``config`` in place from a prior deep-copy snapshot.""" + for field_name in type(config).model_fields: + setattr(config, field_name, getattr(snapshot, field_name)) + + +def _persist_login_sync( + config: Config, + ref: OAuthRef, + token: OAuthToken, + apply_config: Callable[[Config], None], +) -> None: + with file_lock(get_config_file()): + snapshot = config.model_copy(deep=True) + previous_token = load_tokens(ref) + save_tokens(ref, token) + try: + apply_config(config) + save_config(config) + except BaseException as persistence_error: + restore_config_state(config, snapshot) + try: + if previous_token is not None: + save_tokens(ref, previous_token) + else: + delete_tokens(ref) + except Exception as rollback_error: + logger.error( + "Failed to roll back OAuth login persistence: {error}", + error=rollback_error, + ) + raise OAuthPersistenceError( + "OAuth login persistence failed and credential rollback also failed." + ) from persistence_error + raise + + +async def persist_login( + config: Config, + ref: OAuthRef, + token: OAuthToken, + apply_config: Callable[[Config], None], +) -> None: + """Persist OAuth credentials and config together outside the event loop.""" + await asyncio.to_thread(_persist_login_sync, config, ref, token, apply_config) + + +def _persist_logout_sync( + config: Config, + ref: OAuthRef, + remove_config: Callable[[Config], None], +) -> None: + with file_lock(get_config_file()): + snapshot = config.model_copy(deep=True) + remove_config(config) + try: + save_config(config) + except BaseException: + restore_config_state(config, snapshot) + raise + try: + delete_tokens(ref) + except Exception as delete_error: + restore_config_state(config, snapshot) + try: + save_config(config) + except Exception as rollback_error: + logger.error( + "Failed to restore configuration after OAuth credential deletion failed: " + "{error}", + error=rollback_error, + ) + raise OAuthPersistenceError( + "OAuth logout failed and configuration rollback also failed." + ) from delete_error + raise OAuthPersistenceError( + "OAuth credential deletion failed; logout was rolled back." + ) from delete_error + + +async def persist_logout( + config: Config, + ref: OAuthRef, + remove_config: Callable[[Config], None], +) -> None: + """Persist OAuth logout outside the event loop, config removal first.""" + await asyncio.to_thread(_persist_logout_sync, config, ref, remove_config) + + +def _persist_config_change_sync( + config: Config, + apply_config: Callable[[Config], None], +) -> None: + with file_lock(get_config_file()): + snapshot = config.model_copy(deep=True) + try: + apply_config(config) + save_config(config) + except BaseException: + restore_config_state(config, snapshot) + raise + + +async def persist_config_change(config: Config, apply_config: Callable[[Config], None]) -> None: + """Persist a config-only change atomically outside the event loop.""" + await asyncio.to_thread(_persist_config_change_sync, config, apply_config) + + async def request_device_authorization() -> DeviceAuthorization: async with ( new_client_session() as session, @@ -1152,10 +1275,26 @@ async def _refresh_tokens( xlock.release() async def _refresh_token_for_ref(self, ref: OAuthRef, refresh_token_value: str) -> OAuthToken: + from pythinker_code.auth.copilot import ( + GITHUB_COPILOT_OAUTH_KEY, + refresh_copilot_token, + ) + + if ref.key == GITHUB_COPILOT_OAUTH_KEY: + return await refresh_copilot_token(refresh_token_value) if ref.key == "oauth/openai-chatgpt": from pythinker_code.auth.openai import refresh_openai_chatgpt_token return await refresh_openai_chatgpt_token(refresh_token_value) + if ref.key == "oauth/xai": + from pythinker_code.auth.xai import refresh_xai_token + + return await refresh_xai_token(refresh_token_value) + if ref.key.startswith("oauth/snowflake-cortex/"): + from pythinker_code.auth.snowflake import refresh_snowflake_cortex_token + + account = ref.key.removeprefix("oauth/snowflake-cortex/") + return await refresh_snowflake_cortex_token(account, refresh_token_value) return await refresh_token(refresh_token_value) def _apply_access_token( diff --git a/src/pythinker_code/auth/oauth_flows.py b/src/pythinker_code/auth/oauth_flows.py new file mode 100644 index 00000000..c6c77d6d --- /dev/null +++ b/src/pythinker_code/auth/oauth_flows.py @@ -0,0 +1,721 @@ +"""Provider-neutral helpers for OAuth device-code and loopback-PKCE flows.""" + +from __future__ import annotations + +import asyncio +import base64 +import hashlib +import json +import secrets +import time +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import Any, NamedTuple, cast +from urllib.parse import parse_qs, urlencode, urlsplit + +import aiohttp + +from pythinker_code.auth.oauth import OAuthDeviceExpired, OAuthError +from pythinker_code.utils.aiohttp import new_client_session + +_DEVICE_CODE_GRANT = "urn:ietf:params:oauth:grant-type:device_code" +_DEFAULT_DEVICE_INTERVAL = 5 +_SLOW_DOWN_INCREMENT = 5 +_IMPLICIT_BOOTSTRAP_HTML = """ +

Finishing sign-in…

+""" + + +class OAuthAccessDenied(OAuthError): + """The resource owner denied an OAuth authorization request.""" + + +class OAuthStateMismatch(OAuthError): + """The loopback callback did not contain the expected OAuth state.""" + + +@dataclass(frozen=True, slots=True) +class PkceCodes: + code_verifier: str + code_challenge: str + + +@dataclass(frozen=True, slots=True) +class DeviceCode: + user_code: str + verification_uri: str + device_code: str + interval: int + expires_in: int + verification_uri_complete: str | None = None + + +class LoopbackAuthorization(NamedTuple): + """Successful loopback result, with named fields and tuple unpacking.""" + + authorization_code: str + code_verifier: str + redirect_uri: str + + +class ImplicitAuthorization(NamedTuple): + """Successful OAuth implicit-flow result. + + ``expires_in`` is ``None`` when the provider omitted the value or returned a + malformed/nonpositive one: implicit tokens carry no refresh token, so a + fabricated lifetime would be a lie that could later feed a validity check. + Callers must treat ``None`` as "unknown", never as a trusted duration. + """ + + access_token: str + expires_in: int | None + state: str + + +def _base64url(data: bytes) -> str: + return base64.urlsafe_b64encode(data).decode(encoding="utf-8", errors="replace").rstrip("=") + + +def generate_pkce() -> PkceCodes: + """Generate an RFC 7636 verifier and its S256 challenge.""" + verifier = _base64url(secrets.token_bytes(32)) + verifier_bytes = verifier.encode(encoding="utf-8") + challenge = _base64url(hashlib.sha256(verifier_bytes).digest()) + return PkceCodes(code_verifier=verifier, code_challenge=challenge) + + +def generate_state() -> str: + """Generate a cryptographically random OAuth state value.""" + return secrets.token_urlsafe(32) + + +async def _post_form( + endpoint: str, + data: Mapping[str, str], + *, + operation: str, + headers: Mapping[str, str] | None = None, +) -> tuple[int, dict[str, Any]]: + try: + async with ( + new_client_session() as session, + session.post(endpoint, data=dict(data), headers=headers) as response, + ): + status = response.status + payload_any: Any = await response.json(content_type=None) + except (aiohttp.ClientError, TimeoutError, OSError, ValueError) as exc: + raise OAuthError(f"{operation} request failed.") from exc + + if not isinstance(payload_any, dict): + raise OAuthError(f"{operation} returned an invalid response.") + return status, cast(dict[str, Any], payload_any) + + +async def request_device_code( + *, + device_authorization_endpoint: str, + client_id: str, + scope: str | Sequence[str] | None = None, + extra_params: Mapping[str, str] | None = None, + headers: Mapping[str, str] | None = None, +) -> DeviceCode: + """Request an RFC 8628 device code without beginning token polling.""" + data = dict(extra_params or {}) + data["client_id"] = client_id + if scope: + data["scope"] = scope if isinstance(scope, str) else " ".join(scope) + + status, payload = await _post_form( + device_authorization_endpoint, + data, + operation="Device authorization", + headers=headers, + ) + if not 200 <= status < 300: + raise OAuthError(f"Device authorization failed (HTTP {status}).") + + required = ("user_code", "verification_uri", "device_code", "expires_in") + if any(not payload.get(field) for field in required): + raise OAuthError("Device authorization response was incomplete.") + try: + interval = int(payload.get("interval") or _DEFAULT_DEVICE_INTERVAL) + expires_in = int(payload["expires_in"]) + except (TypeError, ValueError) as exc: + raise OAuthError("Device authorization response contained invalid timing values.") from exc + if interval <= 0 or expires_in <= 0: + raise OAuthError("Device authorization response contained invalid timing values.") + + complete = payload.get("verification_uri_complete") + return DeviceCode( + user_code=str(payload["user_code"]), + verification_uri=str(payload["verification_uri"]), + verification_uri_complete=str(complete) if complete else None, + device_code=str(payload["device_code"]), + interval=interval, + expires_in=expires_in, + ) + + +async def poll_device_token( + *, + token_endpoint: str, + client_id: str, + device_code: DeviceCode, + extra_params: Mapping[str, str] | None = None, + deadline: float | None = None, + headers: Mapping[str, str] | None = None, +) -> dict[str, Any]: + """Poll an RFC 8628 token endpoint until authorization succeeds or terminates. + + ``deadline`` is an absolute ``time.monotonic()`` value. When omitted, the + device authorization's ``expires_in`` value defines the deadline. + """ + effective_deadline = deadline + if effective_deadline is None: + effective_deadline = time.monotonic() + device_code.expires_in + interval = float(device_code.interval) + + data = dict(extra_params or {}) + data.update( + { + "client_id": client_id, + "device_code": device_code.device_code, + "grant_type": _DEVICE_CODE_GRANT, + } + ) + + while True: + remaining = effective_deadline - time.monotonic() + if remaining <= 0: + raise OAuthDeviceExpired("Device authorization expired before completion.") + await asyncio.sleep(min(interval, remaining)) + if time.monotonic() >= effective_deadline: + raise OAuthDeviceExpired("Device authorization expired before completion.") + + status, payload = await _post_form( + token_endpoint, + data, + operation="Device token polling", + headers=headers, + ) + if time.monotonic() >= effective_deadline: + raise OAuthDeviceExpired("Device authorization expired before completion.") + error = str(payload.get("error") or "") + if 200 <= status < 300 and not error: + access_token = payload.get("access_token") + if not isinstance(access_token, str) or not access_token: + raise OAuthError("Device token polling returned an incomplete response.") + return payload + if error == "authorization_pending": + continue + if error == "slow_down": + interval += _SLOW_DOWN_INCREMENT + continue + if error == "expired_token": + raise OAuthDeviceExpired("Device authorization expired.") + if error == "access_denied": + raise OAuthAccessDenied("Device authorization was denied.") + raise OAuthError(f"Device token polling failed (HTTP {status}).") + + +def _authorize_url( + *, + authorize_endpoint: str, + client_id: str, + redirect_uri: str, + scope: str | Sequence[str], + pkce: PkceCodes, + state: str, + extra_params: Mapping[str, str] | None, +) -> str: + params = dict(extra_params or {}) + params.update( + { + "client_id": client_id, + "code_challenge": pkce.code_challenge, + "code_challenge_method": "S256", + "redirect_uri": redirect_uri, + "response_type": "code", + "scope": scope if isinstance(scope, str) else " ".join(scope), + "state": state, + } + ) + separator = "&" if "?" in authorize_endpoint else "?" + return f"{authorize_endpoint}{separator}{urlencode(params)}" + + +def _open_browser(url: str) -> None: + from pythinker_code.utils.term import open_url_in_browser + + open_url_in_browser(url) + + +async def _write_callback_response( + writer: asyncio.StreamWriter, *, status: str, message: str +) -> None: + body = bytes(message, encoding="utf-8") + headers = bytes( + f"HTTP/1.1 {status}\r\n" + "Content-Type: text/plain; charset=utf-8\r\n" + f"Content-Length: {len(body)}\r\n" + "Connection: close\r\n\r\n", + encoding="utf-8", + ) + writer.write(headers + body) + await writer.drain() + + +async def _write_http_response( + writer: asyncio.StreamWriter, + *, + status: str, + body: bytes, + content_type: str, +) -> None: + headers = bytes( + f"HTTP/1.1 {status}\r\n" + f"Content-Type: {content_type}\r\n" + f"Content-Length: {len(body)}\r\n" + "Connection: close\r\n\r\n", + encoding="utf-8", + ) + writer.write(headers + body) + await writer.drain() + + +async def _handle_loopback_callback( + reader: asyncio.StreamReader, + writer: asyncio.StreamWriter, + *, + redirect_path: str, + expected_state: str, + result: asyncio.Future[str], +) -> None: + try: + line = await reader.readline() + parts = line.decode(encoding="utf-8", errors="replace").strip().split() + if len(parts) < 2 or parts[0] != "GET": + await _write_callback_response(writer, status="404 Not Found", message="Not found.") + return + + parsed = urlsplit(parts[1]) + if parsed.path != redirect_path: + await _write_callback_response(writer, status="404 Not Found", message="Not found.") + return + + params = parse_qs(parsed.query) + state = params.get("state", [None])[0] + if state != expected_state: + await _write_callback_response( + writer, + status="400 Bad Request", + message="OAuth callback state did not match.", + ) + if not result.done(): + result.set_exception(OAuthStateMismatch("OAuth callback state did not match.")) + return + + error = params.get("error", [None])[0] + if error: + await _write_callback_response( + writer, + status="400 Bad Request", + message="OAuth authorization was not completed.", + ) + if not result.done(): + exc: OAuthError + if error == "access_denied": + exc = OAuthAccessDenied("OAuth authorization was denied.") + else: + exc = OAuthError("OAuth authorization callback reported an error.") + result.set_exception(exc) + return + + code = params.get("code", [None])[0] + if not code: + await _write_callback_response( + writer, + status="400 Bad Request", + message="OAuth callback did not include an authorization code.", + ) + if not result.done(): + result.set_exception( + OAuthError("OAuth callback did not include an authorization code.") + ) + return + + await _write_callback_response( + writer, + status="200 OK", + message="Authorization complete. You can close this window.", + ) + if not result.done(): + result.set_result(code) + except asyncio.CancelledError: + raise + except (OSError, ValueError): + if not result.done(): + result.set_exception(OAuthError("Failed to process the OAuth callback.")) + finally: + writer.close() + await writer.wait_closed() + + +async def _handle_implicit_loopback_callback( + reader: asyncio.StreamReader, + writer: asyncio.StreamWriter, + *, + callback_path: str, + token_path: str, + expected_state: str, + result: asyncio.Future[ImplicitAuthorization], +) -> None: + try: + line = await reader.readline() + parts = line.decode(encoding="utf-8", errors="replace").strip().split() + if len(parts) < 2: + await _write_http_response( + writer, + status="404 Not Found", + body=bytes("Not found.", encoding="utf-8"), + content_type="text/plain; charset=utf-8", + ) + return + + method = parts[0] + path = urlsplit(parts[1]).path + if method == "GET" and path == callback_path: + html = _IMPLICIT_BOOTSTRAP_HTML.replace("__TOKEN_PATH__", token_path) + await _write_http_response( + writer, + status="200 OK", + body=bytes(html, encoding="utf-8"), + content_type="text/html; charset=utf-8", + ) + return + + if method != "POST" or path != token_path: + await _write_http_response( + writer, + status="404 Not Found", + body=bytes("Not found.", encoding="utf-8"), + content_type="text/plain; charset=utf-8", + ) + return + + content_length = 0 + while True: + header = await reader.readline() + if header in (b"\r\n", b""): + break + header_text = header.decode(encoding="utf-8", errors="replace") + name, separator, value = header_text.partition(":") + if separator and name.strip().lower() == "content-length": + content_length = int(value.strip()) + + try: + body = await reader.readexactly(content_length) + except asyncio.IncompleteReadError: + await _write_http_response( + writer, + status="400 Bad Request", + body=bytes('{"ok": false}', encoding="utf-8"), + content_type="application/json", + ) + if not result.done(): + result.set_exception(OAuthError("OAuth callback body was truncated.")) + return + try: + payload_any: Any = json.loads(body.decode(encoding="utf-8")) + if not isinstance(payload_any, dict): + raise ValueError("OAuth callback body must be a JSON object.") + payload = cast(dict[str, Any], payload_any) + except ValueError: + await _write_http_response( + writer, + status="400 Bad Request", + body=bytes('{"ok": false}', encoding="utf-8"), + content_type="application/json", + ) + if not result.done(): + result.set_exception(OAuthError("OAuth callback contained invalid JSON.")) + return + + state = str(payload.get("state") or "") + if state != expected_state: + await _write_http_response( + writer, + status="400 Bad Request", + body=bytes('{"ok": false}', encoding="utf-8"), + content_type="application/json", + ) + if not result.done(): + result.set_exception(OAuthStateMismatch("OAuth callback state did not match.")) + return + + error = str(payload.get("error") or "") + if error: + await _write_http_response( + writer, + status="400 Bad Request", + body=bytes('{"ok": false}', encoding="utf-8"), + content_type="application/json", + ) + if not result.done(): + exc: OAuthError + if error == "access_denied": + exc = OAuthAccessDenied("OAuth authorization was denied.") + else: + exc = OAuthError("OAuth authorization callback reported an error.") + result.set_exception(exc) + return + + access_token = payload.get("access_token") + if not isinstance(access_token, str) or not access_token: + await _write_http_response( + writer, + status="400 Bad Request", + body=bytes('{"ok": false}', encoding="utf-8"), + content_type="application/json", + ) + if not result.done(): + result.set_exception(OAuthError("OAuth callback did not include an access token.")) + return + + raw_expires_in = payload.get("expires_in") + expires_in: int | None + try: + expires_in = int(raw_expires_in) if raw_expires_in not in (None, "") else None + except (TypeError, ValueError): + expires_in = None + if expires_in is not None and expires_in <= 0: + expires_in = None + + await _write_http_response( + writer, + status="200 OK", + body=bytes('{"ok": true}', encoding="utf-8"), + content_type="application/json", + ) + if not result.done(): + result.set_result( + ImplicitAuthorization( + access_token=access_token, + expires_in=expires_in, + state=state, + ) + ) + except asyncio.CancelledError: + raise + except (OSError, ValueError): + if not result.done(): + result.set_exception(OAuthError("Failed to process the OAuth callback.")) + finally: + writer.close() + await writer.wait_closed() + + +def _server_port(server: asyncio.Server) -> int: + sockets = server.sockets + if not sockets: + raise OAuthError("OAuth callback server did not expose a listening socket.") + address = sockets[0].getsockname() + if not isinstance(address, tuple): + raise OAuthError("OAuth callback server returned an invalid address.") + parts = cast("tuple[object, ...]", address) + if len(parts) < 2 or not isinstance(parts[1], int): + raise OAuthError("OAuth callback server returned an invalid address.") + return parts[1] + + +async def run_loopback_pkce_flow( + *, + authorize_endpoint: str, + client_id: str, + scope: str | Sequence[str], + redirect_path: str, + port: int = 0, + timeout: float = 15 * 60, + extra_authorize_params: Mapping[str, str] | None = None, + browser_open: Callable[[str], object] | None = None, +) -> LoopbackAuthorization: + """Run an OAuth authorization-code flow using a 127.0.0.1 PKCE callback.""" + if not redirect_path.startswith("/"): + raise ValueError("redirect_path must start with '/'.") + if not 0 <= port <= 65535: + raise ValueError("port must be between 0 and 65535.") + if timeout <= 0: + raise ValueError("timeout must be positive.") + + pkce = generate_pkce() + state = generate_state() + result: asyncio.Future[str] = asyncio.get_running_loop().create_future() + callback_tasks: set[asyncio.Task[None]] = set() + + def on_client_connected(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + task = asyncio.create_task( + _handle_loopback_callback( + reader, + writer, + redirect_path=redirect_path, + expected_state=state, + result=result, + ) + ) + callback_tasks.add(task) + task.add_done_callback(callback_tasks.discard) + + try: + server = await asyncio.start_server(on_client_connected, "127.0.0.1", port) + except OSError as exc: + raise OAuthError("Failed to start the OAuth callback server on 127.0.0.1.") from exc + + try: + bound_port = _server_port(server) + redirect_uri = f"http://127.0.0.1:{bound_port}{redirect_path}" + authorize_url = _authorize_url( + authorize_endpoint=authorize_endpoint, + client_id=client_id, + redirect_uri=redirect_uri, + scope=scope, + pkce=pkce, + state=state, + extra_params=extra_authorize_params, + ) + try: + (browser_open or _open_browser)(authorize_url) + except Exception as exc: + raise OAuthError("Failed to open a browser for OAuth authorization.") from exc + + try: + authorization_code = await asyncio.wait_for(result, timeout=timeout) + except TimeoutError as exc: + raise OAuthError("Timed out waiting for the OAuth authorization callback.") from exc + return LoopbackAuthorization( + authorization_code=authorization_code, + code_verifier=pkce.code_verifier, + redirect_uri=redirect_uri, + ) + finally: + server.close() + for task in callback_tasks: + task.cancel() + if callback_tasks: + await asyncio.gather(*callback_tasks, return_exceptions=True) + await server.wait_closed() + + +async def run_loopback_implicit_flow( + *, + authorize_endpoint: str, + client_id: str, + scope: str | Sequence[str], + callback_path: str, + token_path: str, + port: int, + timeout: float = 5 * 60, + redirect_host: str = "localhost", + extra_authorize_params: Mapping[str, str] | None = None, + browser_open: Callable[[str], object] | None = None, +) -> ImplicitAuthorization: + """Run an OAuth implicit flow using a pinned loopback callback server.""" + if not callback_path.startswith("/"): + raise ValueError("callback_path must start with '/'.") + if not token_path.startswith("/"): + raise ValueError("token_path must start with '/'.") + if redirect_host not in {"localhost", "127.0.0.1"}: + raise ValueError("redirect_host must resolve to the loopback interface.") + if not 1 <= port <= 65535: + raise ValueError("port must be between 1 and 65535.") + if timeout <= 0: + raise ValueError("timeout must be positive.") + + state = generate_state() + result: asyncio.Future[ImplicitAuthorization] = asyncio.get_running_loop().create_future() + callback_tasks: set[asyncio.Task[None]] = set() + + def on_client_connected(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + task = asyncio.create_task( + _handle_implicit_loopback_callback( + reader, + writer, + callback_path=callback_path, + token_path=token_path, + expected_state=state, + result=result, + ) + ) + callback_tasks.add(task) + task.add_done_callback(callback_tasks.discard) + + try: + server = await asyncio.start_server(on_client_connected, redirect_host, port) + except OSError as exc: + raise OAuthError(f"Failed to start the OAuth callback server on {redirect_host}.") from exc + + try: + params = dict(extra_authorize_params or {}) + params.update( + { + "response_type": "token", + "client_id": client_id, + "redirect_uri": f"http://{redirect_host}:{port}{callback_path}", + "scope": scope if isinstance(scope, str) else " ".join(scope), + "state": state, + } + ) + separator = "&" if "?" in authorize_endpoint else "?" + authorize_url = f"{authorize_endpoint}{separator}{urlencode(params)}" + try: + (browser_open or _open_browser)(authorize_url) + except Exception as exc: + raise OAuthError("Failed to open a browser for OAuth authorization.") from exc + + try: + return await asyncio.wait_for(result, timeout=timeout) + except TimeoutError as exc: + raise OAuthError("Timed out waiting for the OAuth authorization callback.") from exc + finally: + server.close() + for task in callback_tasks: + task.cancel() + if callback_tasks: + await asyncio.gather(*callback_tasks, return_exceptions=True) + await server.wait_closed() + + +__all__ = [ + "DeviceCode", + "ImplicitAuthorization", + "LoopbackAuthorization", + "OAuthAccessDenied", + "OAuthStateMismatch", + "PkceCodes", + "generate_pkce", + "generate_state", + "poll_device_token", + "request_device_code", + "run_loopback_implicit_flow", + "run_loopback_pkce_flow", +] diff --git a/src/pythinker_code/auth/opencode_go.py b/src/pythinker_code/auth/opencode_go.py index 2168e995..f7ef6902 100644 --- a/src/pythinker_code/auth/opencode_go.py +++ b/src/pythinker_code/auth/opencode_go.py @@ -9,6 +9,13 @@ from pydantic import SecretStr from pythinker_code.auth import OPENCODE_GO_PLATFORM_ID +from pythinker_code.auth.models_dev import ( + ModelsDevCatalog, + ModelsDevProvider, + get_models_dev_catalog, + get_provider_models, + parse_models_dev_catalog, +) from pythinker_code.auth.oauth import OAuthEvent from pythinker_code.config import Config, LLMModel, LLMProvider, save_config from pythinker_code.llm import ModelCapability @@ -30,15 +37,7 @@ OPENCODE_GO_DEFAULT_MODEL_ALIAS = "opencode-go/kimi-k2.6" OPENCODE_GO_DEFAULT_CONTEXT = 262_000 -# models.dev is OpenCode's own source of truth for model metadata (context -# window, display name). The Go /models endpoint returns ids only, so we -# enrich ids not in the curated catalog below from this catalog. -MODELS_DEV_API_URL = "https://models.dev/api.json" MODELS_DEV_PROVIDER_ID = "opencode-go" -# The models.dev fetch is best-effort enrichment, so it must not stall login on -# the 120s default. A tight cap means a slow/partial endpoint degrades quickly -# to the curated catalog instead of holding the user for up to two minutes. -MODELS_DEV_TIMEOUT = aiohttp.ClientTimeout(total=15, sock_connect=8, sock_read=10) @dataclass(frozen=True, slots=True) @@ -210,41 +209,30 @@ def _extract_model_ids(data: object) -> list[str]: return ids -def _parse_models_dev_metadata(data: object) -> dict[str, _ModelsDevMeta]: - """Extract display name, context, and API shape per opencode-go model id.""" - if not isinstance(data, dict): - return {} - provider = cast(dict[str, Any], data).get(MODELS_DEV_PROVIDER_ID) - if not isinstance(provider, dict): - return {} - models = cast(dict[str, Any], provider).get("models") - if not isinstance(models, dict): - return {} - default_npm = cast(dict[str, Any], provider).get("npm") +def _metadata_from_catalog(catalog: ModelsDevCatalog) -> dict[str, _ModelsDevMeta]: result: dict[str, _ModelsDevMeta] = {} - for model_id, entry in cast(dict[str, Any], models).items(): - if not isinstance(entry, dict): - continue - entry_d = cast(dict[str, Any], entry) - name = entry_d.get("name") - display_name = name if isinstance(name, str) and name else None - limit = entry_d.get("limit") - context = cast(dict[str, Any], limit).get("context") if isinstance(limit, dict) else None - max_context = context if isinstance(context, int) and context > 0 else None - model_provider = entry_d.get("provider") - npm = ( - cast(dict[str, Any], model_provider).get("npm") - if isinstance(model_provider, dict) - else None - ) - effective_npm = npm or default_npm - is_anthropic = ( - effective_npm == MODELS_DEV_ANTHROPIC_NPM if isinstance(effective_npm, str) else None + for model_id, model in get_provider_models(catalog, MODELS_DEV_PROVIDER_ID).items(): + is_anthropic = model.npm == MODELS_DEV_ANTHROPIC_NPM if model.npm is not None else None + result[model_id] = _ModelsDevMeta( + display_name=model.display_name, + max_context=model.context_length, + is_anthropic=is_anthropic, ) - result[model_id] = _ModelsDevMeta(display_name, max_context, is_anthropic) return result +def _parse_models_dev_metadata(data: object) -> dict[str, _ModelsDevMeta]: + """Compatibility wrapper over the shared models.dev catalog parser.""" + raw_catalog = cast(dict[object, object], data) if isinstance(data, dict) else None + if raw_catalog is not None and all( + isinstance(provider, ModelsDevProvider) for provider in raw_catalog.values() + ): + catalog = cast(ModelsDevCatalog, raw_catalog) + else: + catalog = parse_models_dev_catalog(cast(object, data)) + return _metadata_from_catalog(catalog) + + def _build_models( model_ids: list[str], metadata: dict[str, _ModelsDevMeta], @@ -282,18 +270,9 @@ def _build_models( async def _fetch_models_dev_metadata() -> dict[str, _ModelsDevMeta]: - """Best-effort metadata fetch. Returns {} on any failure so login still - succeeds (falling back to the curated catalog) when models.dev is - unreachable.""" - try: - async with ( - new_client_session(timeout=MODELS_DEV_TIMEOUT) as session, - session.get(MODELS_DEV_API_URL, raise_for_status=True) as response, - ): - payload = await response.json(content_type=None) - except (TimeoutError, aiohttp.ClientError, ValueError): - return {} - return _parse_models_dev_metadata(payload) + """Load best-effort OpenCode metadata from the shared models.dev catalog.""" + result = await get_models_dev_catalog() + return _parse_models_dev_metadata(result.catalog) async def _discover_opencode_go_models(api_key: str) -> tuple[OpenCodeGoModel, ...]: @@ -311,9 +290,8 @@ async def _discover_opencode_go_models(api_key: str) -> tuple[OpenCodeGoModel, . if not model_ids: return () - # models.dev is the authority for API shape + context; fetch it on every - # login (best-effort) so the live list self-corrects even when our curated - # catalog drifts. Falls back to the catalog when models.dev is unreachable. + # models.dev is the authority for API shape + context. Its shared loader + # refreshes stale data on demand and falls back to disk when unavailable. metadata = await _fetch_models_dev_metadata() return _build_models(model_ids, metadata) diff --git a/src/pythinker_code/auth/platforms.py b/src/pythinker_code/auth/platforms.py index ad22780b..9172b04a 100644 --- a/src/pythinker_code/auth/platforms.py +++ b/src/pythinker_code/auth/platforms.py @@ -7,11 +7,15 @@ from pydantic import BaseModel from pythinker_code.auth import ( + DIGITALOCEAN_PLATFORM_ID, + GITHUB_COPILOT_PLATFORM_ID, LM_STUDIO_PLATFORM_ID, OLLAMA_PLATFORM_ID, OPENAI_API_PLATFORM_ID, OPENAI_CHATGPT_PLATFORM_ID, PYTHINKER_CODE_PLATFORM_ID, + SNOWFLAKE_CORTEX_PLATFORM_ID, + XAI_PLATFORM_ID, ) from pythinker_code.config import Config, LLMModel, load_config, save_config from pythinker_code.llm import ModelCapability @@ -96,6 +100,27 @@ def _ollama_base_url() -> str: name="OpenAI ChatGPT Codex", base_url="https://chatgpt.com/backend-api/codex", ), + Platform( + id=GITHUB_COPILOT_PLATFORM_ID, + name="GitHub Copilot", + base_url="https://api.githubcopilot.com", + ), + Platform( + id=DIGITALOCEAN_PLATFORM_ID, + name="DigitalOcean", + base_url="https://inference.do-ai.run/v1", + ), + Platform( + id=XAI_PLATFORM_ID, + name="xAI Grok", + base_url="https://api.x.ai/v1", + ), + Platform( + id=SNOWFLAKE_CORTEX_PLATFORM_ID, + name="Snowflake Cortex", + # Display-only; inference and OAuth endpoints are account-scoped. + base_url="https://app.snowflake.com", + ), Platform( id="pythinker_ai-cn", name="Pythinker AI Open Platform (pythinker-ai.cn)", @@ -274,7 +299,15 @@ async def refresh_managed_models(config: Config) -> bool: # the wire-shape suffix (`managed:minimax-anthropic`). if ( provider_key in OPENCODE_GO_PROVIDER_KEYS - or provider_key in {MINIMAX_ANTHROPIC_PROVIDER_KEY, KIMI_PROVIDER_KEY} + or provider_key + in { + managed_provider_key(DIGITALOCEAN_PLATFORM_ID), + managed_provider_key(GITHUB_COPILOT_PLATFORM_ID), + managed_provider_key(SNOWFLAKE_CORTEX_PLATFORM_ID), + managed_provider_key(XAI_PLATFORM_ID), + MINIMAX_ANTHROPIC_PROVIDER_KEY, + KIMI_PROVIDER_KEY, + } or provider_key in z_ai_provider_keys ): continue diff --git a/src/pythinker_code/auth/snowflake.py b/src/pythinker_code/auth/snowflake.py new file mode 100644 index 00000000..2d9db74d --- /dev/null +++ b/src/pythinker_code/auth/snowflake.py @@ -0,0 +1,372 @@ +from __future__ import annotations + +import base64 +import re +from collections.abc import AsyncIterator +from dataclasses import dataclass +from typing import Any, cast +from urllib.parse import quote + +import aiohttp +from pydantic import SecretStr + +from pythinker_code.auth import SNOWFLAKE_CORTEX_PLATFORM_ID +from pythinker_code.auth.models_dev import ( + CatalogModel, + build_catalog_models, + get_models_dev_catalog, +) +from pythinker_code.auth.oauth import ( + OAuthError, + OAuthEvent, + OAuthToken, + OAuthUnauthorized, + persist_config_change, + persist_login, + persist_logout, +) +from pythinker_code.auth.oauth_flows import run_loopback_pkce_flow +from pythinker_code.auth.platforms import managed_model_key, managed_provider_key +from pythinker_code.config import Config, LLMModel, LLMProvider, OAuthRef +from pythinker_code.utils.aiohttp import new_client_session +from pythinker_code.utils.logging import logger + +SNOWFLAKE_CLIENT_ID = "LOCAL_APPLICATION" +SNOWFLAKE_REDIRECT_PATH = "/" +SNOWFLAKE_PROVIDER_KEY = managed_provider_key(SNOWFLAKE_CORTEX_PLATFORM_ID) +SNOWFLAKE_OAUTH_KEY_PREFIX = "oauth/snowflake-cortex/" +SNOWFLAKE_JSON_HEADERS = {"Accept": "application/json"} +SNOWFLAKE_MODELS_DEV_PROVIDER_ID = "snowflake-cortex" +SNOWFLAKE_DEFAULT_CONTEXT = 128_000 +_ROLE_SIMPLE = re.compile(r"^[-_A-Za-z0-9]+$") +# A Snowflake account locator is a plain host label: alphanumerics plus the +# org-account separators '.', '-', '_'. Anything carrying an authority, path, +# port, userinfo, query, or fragment delimiter is rejected so it can never +# redirect the OAuth/token host. +_ACCOUNT_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$") +_DEFAULT_EXPIRES_IN = 600 # Snowflake access tokens are short-lived; PENDING-LIVE + + +def _skip_browser_open(_url: str) -> None: + return None + + +def normalize_account(raw: str) -> str: + """Normalize and validate a Snowflake account locator. + + Strips an optional scheme and the ``.snowflakecomputing.com`` suffix, then + rejects anything that is not a plain account locator. Raises ``ValueError`` + for empty or malformed input before any URL is built. + """ + account = raw.strip() + account = re.sub(r"^https?://", "", account, flags=re.IGNORECASE) + account = re.sub(r"\.snowflakecomputing\.com/?$", "", account, flags=re.IGNORECASE) + account = account.rstrip("/") + if not account: + raise ValueError("Snowflake account identifier is required.") + if not _ACCOUNT_RE.match(account): + raise ValueError("Snowflake account identifier is invalid.") + return account + + +def _authorize_url(account: str) -> str: + return f"https://{account}.snowflakecomputing.com/oauth/authorize" + + +def _token_url(account: str) -> str: + return f"https://{account}.snowflakecomputing.com/oauth/token-request" + + +def _cortex_base_url(account: str) -> str: + return f"https://{account}.snowflakecomputing.com/api/v2/cortex/v1" + + +def _oauth_ref(account: str) -> OAuthRef: + return OAuthRef(storage="file", key=f"{SNOWFLAKE_OAUTH_KEY_PREFIX}{account}") + + +def _scope(role: str | None) -> str: + if not role: + return "refresh_token" + if _ROLE_SIMPLE.match(role): + return f"refresh_token session:role:{role}" + return f"refresh_token session:role-encoded:{quote(role, safe='')}" + + +def _basic_header() -> str: + raw = f"{SNOWFLAKE_CLIENT_ID}:{SNOWFLAKE_CLIENT_ID}" + encoded = base64.b64encode(raw.encode(encoding="utf-8")).decode(encoding="utf-8") + return f"Basic {encoded}" + + +def _headers() -> dict[str, str]: + return { + "Content-Type": "application/x-www-form-urlencoded", + "Accept": "application/json", + "Authorization": _basic_header(), + } + + +def _inject_default_expiry(payload: dict[str, Any]) -> dict[str, Any]: + if not payload.get("expires_in"): + payload["expires_in"] = _DEFAULT_EXPIRES_IN + return payload + + +@dataclass(frozen=True, slots=True) +class SnowflakeModel: + model_id: str + display_name: str + max_context_size: int # PENDING-LIVE: Cortex does not publish context windows + + @property + def alias(self) -> str: + return managed_model_key(SNOWFLAKE_CORTEX_PLATFORM_ID, self.model_id) + + +SNOWFLAKE_MODELS: tuple[SnowflakeModel, ...] = ( + SnowflakeModel("claude-sonnet-4-5", "Claude Sonnet 4.5 (Cortex)", 200_000), + SnowflakeModel("openai-gpt-5", "OpenAI GPT-5 (Cortex)", 128_000), + SnowflakeModel("llama3.1-405b", "Llama 3.1 405B (Cortex)", 128_000), +) + + +async def _post_form( + endpoint: str, + data: dict[str, str], + *, + operation: str, + headers: dict[str, str], +) -> tuple[int, dict[str, Any]]: + try: + async with ( + new_client_session() as session, + session.post(endpoint, data=dict(data), headers=headers) as response, + ): + status = response.status + payload_any: Any = await response.json(content_type=None) + except (aiohttp.ClientError, TimeoutError, OSError, ValueError) as exc: + raise OAuthError(f"{operation} request failed.") from exc + + if not isinstance(payload_any, dict): + raise OAuthError(f"{operation} returned an invalid response.") + return status, cast(dict[str, Any], payload_any) + + +async def _post_token(account: str, data: dict[str, str], *, operation: str) -> dict[str, Any]: + status, payload = await _post_form( + _token_url(account), + data, + operation=operation, + headers=_headers(), + ) + # A 400 invalid_grant means the refresh token was rejected; surface it as + # unauthorized so the refresh path suppresses the token instead of retrying. + if status in {401, 403} or str(payload.get("error") or "") == "invalid_grant": + raise OAuthUnauthorized(f"{operation} was unauthorized.") + if not 200 <= status < 300: + raise OAuthError(f"{operation} failed (HTTP {status}).") + return _inject_default_expiry(payload) + + +async def _exchange_code_for_tokens( + account: str, + code: str, + code_verifier: str, + redirect_uri: str, +) -> dict[str, Any]: + return await _post_token( + account, + { + "grant_type": "authorization_code", + "code": code, + "redirect_uri": redirect_uri, + "client_id": SNOWFLAKE_CLIENT_ID, + "code_verifier": code_verifier, + }, + operation="Snowflake authorization code exchange", + ) + + +async def refresh_snowflake_cortex_token(account: str, refresh_token: str) -> OAuthToken: + payload = await _post_token( + account, + { + "grant_type": "refresh_token", + "refresh_token": refresh_token, + "client_id": SNOWFLAKE_CLIENT_ID, + }, + operation="Snowflake token refresh", + ) + return OAuthToken.from_response(payload) + + +def _apply_snowflake_config( + config: Config, + account: str, + models: tuple[SnowflakeModel, ...] = SNOWFLAKE_MODELS, +) -> None: + # Register the provider and its models, but do NOT make Snowflake the + # default model: its Cortex chat adapter (request/response transforms) is + # not implemented yet, so a Snowflake model cannot serve chat. Selecting it + # as default would claim a capability that does not exist. The existing + # default is left untouched. + config.providers[SNOWFLAKE_PROVIDER_KEY] = LLMProvider( + type="openai_legacy", + base_url=_cortex_base_url(account), + api_key=SecretStr(""), + oauth=_oauth_ref(account), + ) + + for alias, model in list(config.models.items()): + if model.provider == SNOWFLAKE_PROVIDER_KEY: + del config.models[alias] + + for model in models: + config.models[model.alias] = LLMModel( + provider=SNOWFLAKE_PROVIDER_KEY, + model=model.model_id, + max_context_size=model.max_context_size, + display_name=model.display_name, + ) + + +def _catalog_models_to_snowflake(built: tuple[CatalogModel, ...]) -> tuple[SnowflakeModel, ...]: + return tuple( + SnowflakeModel(model.model_id, model.display_name, model.max_context_size) + for model in built + ) + + +async def _discover_snowflake_models() -> tuple[SnowflakeModel, ...]: + """Resolve the Snowflake Cortex model list from the shared models.dev catalog. + + Falls back to the curated list when the catalog is unavailable or degraded. + """ + result = await get_models_dev_catalog() + if not result.is_authoritative: + logger.debug( + "models.dev catalog not authoritative (status={status}, source={source}); " + "using curated Snowflake Cortex models.", + status=result.status, + source=result.source, + ) + return SNOWFLAKE_MODELS + built = build_catalog_models( + result.catalog, + SNOWFLAKE_MODELS_DEV_PROVIDER_ID, + default_context=SNOWFLAKE_DEFAULT_CONTEXT, + ) + converted = _catalog_models_to_snowflake(built) + if converted: + return converted + logger.debug( + "models.dev catalog contained no usable Snowflake Cortex models " + "(status={status}, source={source}); using curated models.", + status=result.status, + source=result.source, + ) + return SNOWFLAKE_MODELS + + +async def login_snowflake( + config: Config, + account: str, + role: str | None = None, + *, + open_browser: bool = True, +) -> AsyncIterator[OAuthEvent]: + if not config.is_from_default_location: + yield OAuthEvent( + "error", + "Login requires the default config file; restart without --config/--config-file.", + ) + return + + try: + account = normalize_account(account) + except ValueError as exc: + yield OAuthEvent("error", str(exc)) + return + + yield OAuthEvent("waiting", "Waiting for Snowflake browser authorization...") + browser_open = None if open_browser else _skip_browser_open + try: + auth = await run_loopback_pkce_flow( + authorize_endpoint=_authorize_url(account), + client_id=SNOWFLAKE_CLIENT_ID, + scope=_scope(role), + redirect_path=SNOWFLAKE_REDIRECT_PATH, + port=0, + extra_authorize_params=None, + browser_open=browser_open, + ) + payload = await _exchange_code_for_tokens( + account, + auth.authorization_code, + auth.code_verifier, + auth.redirect_uri, + ) + except OAuthError as exc: + yield OAuthEvent("error", f"Snowflake Cortex browser login failed: {exc}") + return + + token = OAuthToken.from_response(payload) + if not token.refresh_token: + yield OAuthEvent( + "error", + "Snowflake did not return a refresh token; " + "ensure the integration issues refresh tokens.", + ) + return + + models = await _discover_snowflake_models() + try: + await persist_login( + config, + _oauth_ref(account), + token, + lambda cfg: _apply_snowflake_config(cfg, account, models), + ) + except Exception as exc: + logger.warning("Failed to persist Snowflake Cortex login: {exc}", exc=exc) + yield OAuthEvent("error", "Failed to save Snowflake Cortex login.") + return + + yield OAuthEvent( + "success", + f"Snowflake Cortex credentials saved for account '{account}' " + f"({len(models)} models registered). Chat support is pending the Cortex " + "adapter; select a Snowflake model manually once it ships.", + ) + + +async def logout_snowflake(config: Config) -> AsyncIterator[OAuthEvent]: + if not config.is_from_default_location: + yield OAuthEvent( + "error", + "Logout requires the default config file; restart without --config/--config-file.", + ) + return + + provider = config.providers.get(SNOWFLAKE_PROVIDER_KEY) + ref = provider.oauth if provider is not None and provider.oauth is not None else None + + def _remove(cfg: Config) -> None: + cfg.providers.pop(SNOWFLAKE_PROVIDER_KEY, None) + for alias, model in list(cfg.models.items()): + if model.provider == SNOWFLAKE_PROVIDER_KEY: + del cfg.models[alias] + if cfg.default_model not in cfg.models: + cfg.default_model = next(iter(cfg.models), "") + + try: + if ref is not None: + await persist_logout(config, ref, _remove) + else: + await persist_config_change(config, _remove) + except Exception as exc: + logger.warning("Failed to persist Snowflake Cortex logout: {exc}", exc=exc) + yield OAuthEvent("error", "Failed to log out of Snowflake Cortex.") + return + yield OAuthEvent("success", "Logged out of Snowflake Cortex successfully.") diff --git a/src/pythinker_code/auth/xai.py b/src/pythinker_code/auth/xai.py new file mode 100644 index 00000000..e30023f2 --- /dev/null +++ b/src/pythinker_code/auth/xai.py @@ -0,0 +1,340 @@ +from __future__ import annotations + +from collections.abc import AsyncIterator +from dataclasses import dataclass +from typing import Any, cast + +import aiohttp +from pydantic import SecretStr + +from pythinker_code.auth import XAI_PLATFORM_ID +from pythinker_code.auth.models_dev import ( + CatalogModel, + build_catalog_models, + get_models_dev_catalog, +) +from pythinker_code.auth.oauth import ( + OAuthError, + OAuthEvent, + OAuthToken, + OAuthUnauthorized, + persist_login, + persist_logout, +) +from pythinker_code.auth.oauth_flows import ( + generate_state, + poll_device_token, + request_device_code, + run_loopback_pkce_flow, +) +from pythinker_code.auth.platforms import managed_model_key, managed_provider_key +from pythinker_code.config import Config, LLMModel, LLMProvider, OAuthRef +from pythinker_code.thinking import apply_login_thinking_defaults +from pythinker_code.utils.aiohttp import new_client_session +from pythinker_code.utils.logging import logger + +XAI_CLIENT_ID = "b1a00492-073a-47ea-816f-4c329264a828" +XAI_SCOPE = "openid profile email offline_access grok-cli:access api:access" +XAI_AUTHORIZE_URL = "https://auth.x.ai/oauth2/authorize" +XAI_TOKEN_URL = "https://auth.x.ai/oauth2/token" +XAI_DEVICE_CODE_URL = "https://auth.x.ai/oauth2/device/code" +XAI_BASE_URL = "https://api.x.ai/v1" +XAI_REDIRECT_PORT = 56121 +XAI_REDIRECT_PATH = "/callback" +XAI_OAUTH_KEY = "oauth/xai" +XAI_PROVIDER_KEY = managed_provider_key(XAI_PLATFORM_ID) +XAI_JSON_HEADERS = {"Accept": "application/json"} +XAI_MODELS_DEV_PROVIDER_ID = "xai" +XAI_DEFAULT_CONTEXT = 131_072 + + +@dataclass(frozen=True, slots=True) +class XAIModel: + model_id: str + display_name: str + max_context_size: int + + @property + def alias(self) -> str: + return managed_model_key(XAI_PLATFORM_ID, self.model_id) + + +XAI_MODELS: tuple[XAIModel, ...] = ( + XAIModel("grok-4", "Grok 4", 256_000), + XAIModel("grok-3", "Grok 3", 131_072), + XAIModel("grok-3-mini", "Grok 3 Mini", 131_072), +) + + +def _skip_browser_open(_url: str) -> None: + return None + + +def _xai_oauth_ref() -> OAuthRef: + return OAuthRef(storage="file", key=XAI_OAUTH_KEY) + + +def _apply_xai_config(config: Config, models: tuple[XAIModel, ...] = XAI_MODELS) -> None: + config.providers[XAI_PROVIDER_KEY] = LLMProvider( + type="openai_legacy", + base_url=XAI_BASE_URL, + api_key=SecretStr(""), + oauth=_xai_oauth_ref(), + ) + + for alias, model in list(config.models.items()): + if model.provider == XAI_PROVIDER_KEY: + del config.models[alias] + + for model in models: + config.models[model.alias] = LLMModel( + provider=XAI_PROVIDER_KEY, + model=model.model_id, + max_context_size=model.max_context_size, + display_name=model.display_name, + ) + + if models: + config.default_model = models[0].alias + apply_login_thinking_defaults(config, thinking=False, effort="off") + + +def _catalog_models_to_xai(built: tuple[CatalogModel, ...]) -> tuple[XAIModel, ...]: + return tuple( + XAIModel(model.model_id, model.display_name, model.max_context_size) for model in built + ) + + +async def _discover_xai_models() -> tuple[XAIModel, ...]: + """Resolve the live xAI model list from the shared models.dev catalog. + + Falls back to the curated list when the catalog is unavailable or degraded, + so login never depends on a reachable catalog. + """ + result = await get_models_dev_catalog() + if not result.is_authoritative: + logger.debug( + "models.dev catalog not authoritative (status={status}, source={source}); " + "using curated xAI models.", + status=result.status, + source=result.source, + ) + return XAI_MODELS + built = build_catalog_models( + result.catalog, XAI_MODELS_DEV_PROVIDER_ID, default_context=XAI_DEFAULT_CONTEXT + ) + converted = _catalog_models_to_xai(built) + if converted: + return converted + logger.debug( + "models.dev catalog contained no usable xAI models " + "(status={status}, source={source}); using curated models.", + status=result.status, + source=result.source, + ) + return XAI_MODELS + + +def _error_description(payload: object) -> str: + if not isinstance(payload, dict): + return "" + value = cast(dict[str, Any], payload).get("error_description") + return str(value) if value else "" + + +def _error_code(payload: object) -> str: + if not isinstance(payload, dict): + return "" + value = cast(dict[str, Any], payload).get("error") + return str(value) if value else "" + + +async def _post_token(data: dict[str, str], *, operation: str) -> dict[str, Any]: + try: + async with ( + new_client_session() as session, + session.post(XAI_TOKEN_URL, data=data) as response, + ): + status = response.status + try: + payload_any: Any = await response.json(content_type=None) + except ValueError: + payload_any = None + except (aiohttp.ClientError, TimeoutError, OSError) as exc: + raise OAuthError(f"{operation} request failed.") from exc + + # A 400 invalid_grant means the refresh/authorization token was rejected; + # surface it as unauthorized so the refresh path can suppress the token + # instead of retrying a doomed grant. + if status in {401, 403} or _error_code(payload_any) == "invalid_grant": + raise OAuthUnauthorized(_error_description(payload_any) or f"{operation} was unauthorized.") + if status != 200: + raise OAuthError(_error_description(payload_any) or f"{operation} failed (HTTP {status}).") + if not isinstance(payload_any, dict): + raise OAuthError(f"{operation} returned an invalid response.") + payload = cast(dict[str, Any], payload_any) + return payload + + +async def _exchange_code_for_tokens( + code: str, code_verifier: str, redirect_uri: str +) -> dict[str, Any]: + return await _post_token( + { + "grant_type": "authorization_code", + "code": code, + "redirect_uri": redirect_uri, + "client_id": XAI_CLIENT_ID, + "code_verifier": code_verifier, + }, + operation="xAI authorization code exchange", + ) + + +async def refresh_xai_token(refresh_token: str) -> OAuthToken: + payload = await _post_token( + { + "grant_type": "refresh_token", + "refresh_token": refresh_token, + "client_id": XAI_CLIENT_ID, + }, + operation="xAI token refresh", + ) + return OAuthToken.from_response(payload) + + +async def login_xai_browser( + config: Config, *, open_browser: bool = True +) -> AsyncIterator[OAuthEvent]: + if not config.is_from_default_location: + yield OAuthEvent( + "error", + "Login requires the default config file; restart without --config/--config-file.", + ) + return + + yield OAuthEvent("waiting", "Waiting for xAI Grok browser authorization...") + nonce = generate_state() + browser_open = None if open_browser else _skip_browser_open + try: + auth = await run_loopback_pkce_flow( + authorize_endpoint=XAI_AUTHORIZE_URL, + client_id=XAI_CLIENT_ID, + scope=XAI_SCOPE, + redirect_path=XAI_REDIRECT_PATH, + port=XAI_REDIRECT_PORT, + extra_authorize_params={"plan": "generic", "nonce": nonce}, + browser_open=browser_open, + ) + payload = await _exchange_code_for_tokens( + auth.authorization_code, + auth.code_verifier, + auth.redirect_uri, + ) + except OAuthError as exc: + yield OAuthEvent("error", f"xAI Grok browser login failed: {exc}") + return + + token = OAuthToken.from_response(payload) + if not token.refresh_token: + yield OAuthEvent( + "error", + "xAI Grok did not return a refresh token; the login was not saved.", + ) + return + + models = await _discover_xai_models() + try: + await persist_login( + config, _xai_oauth_ref(), token, lambda cfg: _apply_xai_config(cfg, models) + ) + except Exception as exc: + logger.warning("Failed to persist xAI Grok login: {exc}", exc=exc) + yield OAuthEvent("error", "Failed to save xAI Grok login.") + return + yield OAuthEvent("success", f"xAI Grok configured with model {config.default_model}.") + + +async def login_xai_headless(config: Config) -> AsyncIterator[OAuthEvent]: + if not config.is_from_default_location: + yield OAuthEvent( + "error", + "Login requires the default config file; restart without --config/--config-file.", + ) + return + + try: + device_code = await request_device_code( + device_authorization_endpoint=XAI_DEVICE_CODE_URL, + client_id=XAI_CLIENT_ID, + scope=XAI_SCOPE, + headers=XAI_JSON_HEADERS, + ) + except OAuthError as exc: + yield OAuthEvent("error", f"Failed to start xAI Grok device login: {exc}") + return + + yield OAuthEvent( + "verification_url", + f"Open {device_code.verification_uri} and enter code {device_code.user_code}.", + data={ + "verification_url": device_code.verification_uri, + "user_code": device_code.user_code, + }, + ) + yield OAuthEvent("waiting", "Waiting for xAI Grok device authorization...") + + try: + payload = await poll_device_token( + token_endpoint=XAI_TOKEN_URL, + client_id=XAI_CLIENT_ID, + device_code=device_code, + headers=XAI_JSON_HEADERS, + ) + token = OAuthToken.from_response(payload) + except OAuthError as exc: + yield OAuthEvent("error", f"xAI Grok device login failed: {exc}") + return + + if not token.refresh_token: + yield OAuthEvent( + "error", + "xAI Grok did not return a refresh token; the login was not saved.", + ) + return + + models = await _discover_xai_models() + try: + await persist_login( + config, _xai_oauth_ref(), token, lambda cfg: _apply_xai_config(cfg, models) + ) + except Exception as exc: + logger.warning("Failed to persist xAI Grok login: {exc}", exc=exc) + yield OAuthEvent("error", "Failed to save xAI Grok login.") + return + yield OAuthEvent("success", f"xAI Grok configured with model {config.default_model}.") + + +async def logout_xai(config: Config) -> AsyncIterator[OAuthEvent]: + if not config.is_from_default_location: + yield OAuthEvent( + "error", + "Logout requires the default config file; restart without --config/--config-file.", + ) + return + + def _remove(cfg: Config) -> None: + cfg.providers.pop(XAI_PROVIDER_KEY, None) + for alias, model in list(cfg.models.items()): + if model.provider == XAI_PROVIDER_KEY: + del cfg.models[alias] + if cfg.default_model not in cfg.models: + cfg.default_model = next(iter(cfg.models), "") + + try: + await persist_logout(config, _xai_oauth_ref(), _remove) + except Exception as exc: + logger.warning("Failed to persist xAI Grok logout: {exc}", exc=exc) + yield OAuthEvent("error", "Failed to log out of xAI Grok.") + return + yield OAuthEvent("success", "Logged out of xAI Grok successfully.") diff --git a/src/pythinker_code/cli/__init__.py b/src/pythinker_code/cli/__init__.py index 97bd875a..cd9fd93a 100644 --- a/src/pythinker_code/cli/__init__.py +++ b/src/pythinker_code/cli/__init__.py @@ -158,6 +158,60 @@ def logout_openai(*args: Any, **kwargs: Any) -> Any: return impl(*args, **kwargs) +def login_copilot(*args: Any, **kwargs: Any) -> Any: + from pythinker_code.auth.copilot import login_copilot as impl + + return impl(*args, **kwargs) + + +def logout_copilot(*args: Any, **kwargs: Any) -> Any: + from pythinker_code.auth.copilot import logout_copilot as impl + + return impl(*args, **kwargs) + + +def login_digitalocean(*args: Any, **kwargs: Any) -> Any: + from pythinker_code.auth.digitalocean import login_digitalocean as impl + + return impl(*args, **kwargs) + + +def logout_digitalocean(*args: Any, **kwargs: Any) -> Any: + from pythinker_code.auth.digitalocean import logout_digitalocean as impl + + return impl(*args, **kwargs) + + +def login_snowflake(*args: Any, **kwargs: Any) -> Any: + from pythinker_code.auth.snowflake import login_snowflake as impl + + return impl(*args, **kwargs) + + +def logout_snowflake(*args: Any, **kwargs: Any) -> Any: + from pythinker_code.auth.snowflake import logout_snowflake as impl + + return impl(*args, **kwargs) + + +def login_xai_browser(*args: Any, **kwargs: Any) -> Any: + from pythinker_code.auth.xai import login_xai_browser as impl + + return impl(*args, **kwargs) + + +def login_xai_headless(*args: Any, **kwargs: Any) -> Any: + from pythinker_code.auth.xai import login_xai_headless as impl + + return impl(*args, **kwargs) + + +def logout_xai(*args: Any, **kwargs: Any) -> Any: + from pythinker_code.auth.xai import logout_xai as impl + + return impl(*args, **kwargs) + + def login_opencode_go_api_key(*args: Any, **kwargs: Any) -> Any: from pythinker_code.auth.opencode_go import login_opencode_go_api_key as impl @@ -1467,6 +1521,23 @@ def login( False, "--headless", help="Use OpenAI ChatGPT device-code login." ), api_key: bool = typer.Option(False, "--api-key", help="Configure OpenAI with an API key."), + copilot: bool = typer.Option( + False, "--copilot", help="Login with GitHub Copilot (device code)." + ), + digitalocean: bool = typer.Option( + False, "--digitalocean", help="Login with DigitalOcean (browser)." + ), + snowflake: bool = typer.Option( + False, "--snowflake", help="Login with Snowflake Cortex (browser)." + ), + account: str = typer.Option( + "", "--account", help="Snowflake account identifier (for --snowflake)." + ), + role: str = typer.Option("", "--role", help="Snowflake role (optional, for --snowflake)."), + xai: bool = typer.Option(False, "--xai", help="Login with xAI Grok (browser)."), + xai_device: bool = typer.Option( + False, "--xai-device", help="Login with xAI Grok (device code)." + ), opencode_go: bool = typer.Option( False, "--opencode-go", help="Configure OpenCode Go with an API key." ), @@ -1498,19 +1569,28 @@ def login( help="Override the default base URL for --lm-studio or --ollama.", ), ) -> None: - """Login with OpenAI, OpenCode Go, MiniMax, DeepSeek, Anthropic, or local providers.""" + """Login with OpenAI, GitHub Copilot, DigitalOcean, Snowflake, xAI, or API keys.""" import asyncio from rich.console import Console from rich.status import Status async def _run() -> bool: + if (account.strip() or role.strip()) and not snowflake: + typer.echo("--account and --role require --snowflake.", err=True) + return False + selected_modes = sum( bool(value) for value in ( browser, headless, api_key, + copilot, + digitalocean, + snowflake, + xai, + xai_device, opencode_go, minimax, deepseek, @@ -1525,8 +1605,9 @@ async def _run() -> bool: if selected_modes > 1: typer.echo( "Choose only one of --browser, --headless, --api-key, " - "--opencode-go, --minimax, --deepseek, --z-ai-coding, --z-ai-api, " - "--anthropic, --openrouter, --lm-studio, or --ollama.", + "--copilot, --digitalocean, --snowflake, --xai, --xai-device, --opencode-go, " + "--minimax, --deepseek, --z-ai-coding, --z-ai-api, --anthropic, --openrouter, " + "--lm-studio, or --ollama.", err=True, ) return False @@ -1570,6 +1651,24 @@ async def _run() -> bool: elif minimax: key = typer.prompt("MiniMax API key", hide_input=True).strip() events = login_minimax_api_key(config, key) + elif digitalocean: + events = login_digitalocean(config) + elif snowflake: + account_value = account.strip() + if not account_value: + account_value = typer.prompt("Snowflake account identifier").strip() + role_value = role.strip() + if not role_value: + role_value = typer.prompt( + "Snowflake role (optional)", default="", show_default=False + ).strip() + events = login_snowflake(config, account_value, role_value or None) + elif xai_device: + events = login_xai_headless(config) + elif xai: + events = login_xai_browser(config) + elif copilot: + events = login_copilot(config) elif opencode_go: key = typer.prompt("OpenCode Go API key", hide_input=True).strip() events = login_opencode_go_api_key(config, key) @@ -1606,7 +1705,7 @@ async def _run() -> bool: async for event in events: if event.type == "waiting": if status is None: - status = console.status("Waiting for OpenAI authorization.") + status = console.status("Waiting for authorization.") status.start() continue if status is not None: @@ -1639,6 +1738,10 @@ def logout( "--json", help="Emit OAuth events as JSON lines.", ), + copilot: bool = typer.Option(False, "--copilot", help="Logout from GitHub Copilot."), + digitalocean: bool = typer.Option(False, "--digitalocean", help="Logout from DigitalOcean."), + snowflake: bool = typer.Option(False, "--snowflake", help="Logout from Snowflake Cortex."), + xai: bool = typer.Option(False, "--xai", help="Logout from xAI Grok."), opencode_go: bool = typer.Option(False, "--opencode-go", help="Logout from OpenCode Go."), minimax: bool = typer.Option(False, "--minimax", help="Logout from MiniMax."), deepseek: bool = typer.Option(False, "--deepseek", help="Logout from DeepSeek."), @@ -1651,7 +1754,7 @@ def logout( ), ollama: bool = typer.Option(False, "--ollama", help="Logout from Ollama."), ) -> None: - """Logout from OpenAI, OpenCode Go, MiniMax, DeepSeek, Anthropic, or local providers.""" + """Logout from OpenAI, GitHub Copilot, DigitalOcean, Snowflake, xAI, or API keys.""" import asyncio from rich.console import Console @@ -1659,6 +1762,10 @@ def logout( async def _run() -> bool: ok = True selected_modes = ( + copilot, + digitalocean, + snowflake, + xai, opencode_go, minimax, deepseek, @@ -1671,8 +1778,8 @@ async def _run() -> bool: ) if sum(bool(v) for v in selected_modes) > 1: typer.echo( - "Choose only one of --opencode-go, --minimax, --deepseek, " - "--z-ai-coding, --z-ai-api, --anthropic, --openrouter, " + "Choose only one of --copilot, --digitalocean, --snowflake, --xai, --opencode-go, " + "--minimax, --deepseek, --z-ai-coding, --z-ai-api, --anthropic, --openrouter, " "--lm-studio, or --ollama.", err=True, ) @@ -1691,6 +1798,14 @@ async def _run() -> bool: events = logout_deepseek(config) elif minimax: events = logout_minimax(config) + elif digitalocean: + events = logout_digitalocean(config) + elif snowflake: + events = logout_snowflake(config) + elif xai: + events = logout_xai(config) + elif copilot: + events = logout_copilot(config) elif opencode_go: events = logout_opencode_go(config) elif lm_studio: diff --git a/src/pythinker_code/config.py b/src/pythinker_code/config.py index 57f9f108..60147b5a 100644 --- a/src/pythinker_code/config.py +++ b/src/pythinker_code/config.py @@ -3,6 +3,7 @@ import contextlib import json import os +import tempfile from enum import StrEnum from pathlib import Path from types import UnionType @@ -1537,15 +1538,34 @@ def save_config(config: Config, config_file: Path | None = None): """ config_file = config_file or get_config_file() logger.debug("Saving config to file: {file}", file=config_file) - config_file.parent.mkdir(parents=True, exist_ok=True) + target = config_file.resolve(strict=False) if config_file.is_symlink() else config_file + target.parent.mkdir(parents=True, exist_ok=True) config_data = config.model_dump(mode="json", exclude_none=True) - with open(config_file, "w", encoding="utf-8") as f: - if config_file.suffix.lower() == ".json": - f.write(json.dumps(config_data, ensure_ascii=False, indent=2)) - else: - f.write(tomlkit.dumps(config_data)) # type: ignore[reportUnknownMemberType] - with contextlib.suppress(OSError): - os.chmod(config_file, 0o600) + if config_file.suffix.lower() == ".json": + serialized = json.dumps(config_data, ensure_ascii=False, indent=2) + else: + serialized = tomlkit.dumps(config_data) # type: ignore[reportUnknownMemberType] + + temporary_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=target.parent, + prefix=f".{target.name}.", + suffix=".tmp", + delete=False, + ) as temporary_file: + temporary_path = Path(temporary_file.name) + temporary_file.write(serialized) + temporary_file.flush() + os.fsync(temporary_file.fileno()) + with contextlib.suppress(OSError): + os.chmod(temporary_path, 0o600) + os.replace(temporary_path, target) + finally: + if temporary_path is not None and temporary_path.exists(): + temporary_path.unlink(missing_ok=True) class MigrationError(Exception): diff --git a/src/pythinker_code/tools/lsp/tool.py b/src/pythinker_code/tools/lsp/tool.py index 5ab467e7..6e28bae8 100644 --- a/src/pythinker_code/tools/lsp/tool.py +++ b/src/pythinker_code/tools/lsp/tool.py @@ -174,10 +174,10 @@ def _recommendation_hint(self, absolute_path: str) -> str | None: # plugin once per extension per session. Gated by recommendation_disabled / # recommendation_never inside get_matching_lsp_plugins. The reference's # persisted >=5 ignored-count auto-disable is intentionally NOT wired here: - # it requires incremental writes to the shared global config, and the - # current save_config() rewrites the whole file with no lock/atomic rename - # (multi-instance clobber risk). Deferred until a safe global-write path - # exists; the disabled/never flags still apply. + # it requires incremental writes to the shared global config. Atomic + # replacement prevents torn files, but save_config() still has no global + # read-modify-write lock (multi-instance clobber risk). Deferred until a + # safe global-update path exists; the disabled/never flags still apply. ext = Path(absolute_path).suffix.lower() if not ext or ext in self._recommended_exts: return None diff --git a/src/pythinker_code/ui/shell/oauth.py b/src/pythinker_code/ui/shell/oauth.py index a5233207..6a33f0f4 100644 --- a/src/pythinker_code/ui/shell/oauth.py +++ b/src/pythinker_code/ui/shell/oauth.py @@ -11,6 +11,8 @@ ALIBABA_PLATFORM_ID, ANTHROPIC_PLATFORM_ID, DEEPSEEK_PLATFORM_ID, + DIGITALOCEAN_PLATFORM_ID, + GITHUB_COPILOT_PLATFORM_ID, KIMI_PLATFORM_ID, LM_STUDIO_PLATFORM_ID, MINIMAX_PLATFORM_ID, @@ -20,6 +22,8 @@ OPENAI_CHATGPT_PLATFORM_ID, OPENCODE_GO_PLATFORM_ID, OPENROUTER_PLATFORM_ID, + SNOWFLAKE_CORTEX_PLATFORM_ID, + XAI_PLATFORM_ID, ZAI_API_PLATFORM_ID, ZAI_CODING_PLATFORM_ID, ) @@ -33,11 +37,21 @@ login_anthropic_api_key, logout_anthropic, ) +from pythinker_code.auth.copilot import ( + GITHUB_COPILOT_PROVIDER_KEY, + login_copilot, + logout_copilot, +) from pythinker_code.auth.deepseek import ( DEEPSEEK_PROVIDER_KEY, login_deepseek_api_key, logout_deepseek, ) +from pythinker_code.auth.digitalocean import ( + DIGITALOCEAN_PROVIDER_KEY, + login_digitalocean, + logout_digitalocean, +) from pythinker_code.auth.kimi import ( KIMI_PROVIDER_KEY, login_kimi_api_key, @@ -82,6 +96,17 @@ logout_openrouter, ) from pythinker_code.auth.platforms import managed_provider_key +from pythinker_code.auth.snowflake import ( + SNOWFLAKE_PROVIDER_KEY, + login_snowflake, + logout_snowflake, +) +from pythinker_code.auth.xai import ( + XAI_PROVIDER_KEY, + login_xai_browser, + login_xai_headless, + logout_xai, +) from pythinker_code.auth.z_ai import ( ZAI_API_ROUTE, ZAI_CODING_ROUTE, @@ -114,7 +139,7 @@ async def _render_oauth_events(events: AsyncIterator[OAuthEvent]) -> bool: async for event in events: if event.type == "waiting": if status is None: - status = console.status(f"[{_t.info}]Waiting for OpenAI authorization.[/]") + status = console.status(f"[{_t.info}]Waiting for authorization.[/]") status.start() continue if status is not None: @@ -159,6 +184,11 @@ async def _prompt_text(label: str) -> str | None: _SELECTOR_PROVIDER_ENTRIES: list[OAuthProviderEntry] = [ OAuthProviderEntry(id="browser", name="OpenAI ChatGPT (browser)", auth_type="oauth"), OAuthProviderEntry(id="headless", name="OpenAI ChatGPT (device code)", auth_type="oauth"), + OAuthProviderEntry(id="copilot", name="GitHub Copilot", auth_type="oauth"), + OAuthProviderEntry(id="digitalocean", name="DigitalOcean", auth_type="oauth"), + OAuthProviderEntry(id="snowflake", name="Snowflake Cortex", auth_type="oauth"), + OAuthProviderEntry(id="xai", name="xAI Grok (browser)", auth_type="oauth"), + OAuthProviderEntry(id="xai-device", name="xAI Grok (device code)", auth_type="oauth"), OAuthProviderEntry(id="api-key", name="OpenAI API key", auth_type="api_key"), OAuthProviderEntry(id="opencode-go", name="OpenCode Go", auth_type="api_key"), OAuthProviderEntry(id="minimax", name="MiniMax", auth_type="api_key"), @@ -187,6 +217,11 @@ async def _prompt_text(label: str) -> str | None: managed_provider_key(OPENAI_API_PLATFORM_ID), managed_provider_key(OPENAI_CHATGPT_PLATFORM_ID), ), + "copilot": (GITHUB_COPILOT_PROVIDER_KEY,), + "digitalocean": (DIGITALOCEAN_PROVIDER_KEY,), + "snowflake": (SNOWFLAKE_PROVIDER_KEY,), + "xai": (XAI_PROVIDER_KEY,), + "xai-device": (XAI_PROVIDER_KEY,), "opencode-go": (OPENCODE_GO_OPENAI_PROVIDER_KEY, OPENCODE_GO_ANTHROPIC_PROVIDER_KEY), "minimax": (MINIMAX_ANTHROPIC_PROVIDER_KEY,), "deepseek": (DEEPSEEK_PROVIDER_KEY,), @@ -205,6 +240,10 @@ async def _prompt_text(label: str) -> str | None: # (a single OpenAI entry that clears both OpenAI credentials). _LOGOUT_PROVIDER_ENTRIES: list[OAuthProviderEntry] = [ OAuthProviderEntry(id="openai", name="OpenAI", auth_type="oauth"), + OAuthProviderEntry(id="copilot", name="GitHub Copilot", auth_type="oauth"), + OAuthProviderEntry(id="digitalocean", name="DigitalOcean", auth_type="oauth"), + OAuthProviderEntry(id="snowflake", name="Snowflake Cortex", auth_type="oauth"), + OAuthProviderEntry(id="xai", name="xAI Grok", auth_type="oauth"), OAuthProviderEntry(id="opencode-go", name="OpenCode Go", auth_type="api_key"), OAuthProviderEntry(id="minimax", name="MiniMax", auth_type="api_key"), OAuthProviderEntry(id="deepseek", name="DeepSeek", auth_type="api_key"), @@ -239,7 +278,7 @@ def current_model_key(soul: PythinkerSoul) -> str | None: @registry.command(aliases=["setup"]) async def login(app: Shell, args: str) -> None: - """Login with OpenAI, OpenCode Go, MiniMax, DeepSeek, Anthropic, or local providers.""" + """Login with OpenAI, GitHub Copilot, DigitalOcean, Snowflake, xAI, or API keys.""" soul = ensure_pythinker_soul(app) if soul is None: return @@ -262,6 +301,26 @@ async def login(app: Shell, args: str) -> None: elif mode in ("headless", "device", "device-code"): ok = await _render_oauth_events(login_openai_headless(soul.runtime.config)) provider = "openai-chatgpt" + elif mode in ("copilot", "github-copilot"): + ok = await _render_oauth_events(login_copilot(soul.runtime.config)) + provider = GITHUB_COPILOT_PLATFORM_ID + elif mode == "digitalocean": + ok = await _render_oauth_events(login_digitalocean(soul.runtime.config)) + provider = DIGITALOCEAN_PLATFORM_ID + elif mode == "snowflake": + account = await _prompt_text("Snowflake account identifier") + if not account: + console.print(f"[{_t.error}]Snowflake account identifier is required.[/]") + return + role = await _prompt_text("Snowflake role (optional)") + ok = await _render_oauth_events(login_snowflake(soul.runtime.config, account, role or None)) + provider = SNOWFLAKE_CORTEX_PLATFORM_ID + elif mode == "xai": + ok = await _render_oauth_events(login_xai_browser(soul.runtime.config)) + provider = XAI_PLATFORM_ID + elif mode in ("xai-device", "xai-headless"): + ok = await _render_oauth_events(login_xai_headless(soul.runtime.config)) + provider = XAI_PLATFORM_ID elif mode in ("api-key", "apikey", "api"): api_key = await _prompt_api_key("OpenAI") if not api_key: @@ -356,8 +415,10 @@ async def login(app: Shell, args: str) -> None: else: console.print( f"[{_t.error}]Usage: /login " - "[browser|headless|api-key|opencode-go|minimax|deepseek|z-ai-coding|z-ai-api|" - "moonshot|kimi|alibaba|anthropic|openrouter|lm-studio|ollama][/]" + "[browser|headless|copilot|digitalocean|snowflake|xai|xai-device|api-key|" + "opencode-go|" + "minimax|deepseek|z-ai-coding|z-ai-api|moonshot|kimi|alibaba|anthropic|" + "openrouter|lm-studio|ollama][/]" ) return if not ok: @@ -372,7 +433,7 @@ async def login(app: Shell, args: str) -> None: @registry.command async def logout(app: Shell, args: str) -> None: - """Logout from OpenAI, OpenCode Go, MiniMax, DeepSeek, Anthropic, or local providers.""" + """Logout from OpenAI, GitHub Copilot, DigitalOcean, Snowflake, xAI, or API keys.""" soul = ensure_pythinker_soul(app) if soul is None: return @@ -405,6 +466,14 @@ async def logout(app: Shell, args: str) -> None: if mode == "openai": ok = await _render_oauth_events(logout_openai(config)) + elif mode in ("copilot", "github-copilot"): + ok = await _render_oauth_events(logout_copilot(config)) + elif mode == "digitalocean": + ok = await _render_oauth_events(logout_digitalocean(config)) + elif mode == "snowflake": + ok = await _render_oauth_events(logout_snowflake(config)) + elif mode == "xai": + ok = await _render_oauth_events(logout_xai(config)) elif mode == "openrouter": ok = await _render_oauth_events(logout_openrouter(config)) elif mode == "anthropic": @@ -438,8 +507,10 @@ async def logout(app: Shell, args: str) -> None: else: console.print( f"[{_t.error}]Usage: /logout " - "[openai|opencode-go|minimax|deepseek|z-ai-coding|z-ai-api|moonshot|kimi|" - "alibaba|anthropic|openrouter|lm-studio|ollama|github-feedback][/]" + "[openai|copilot|digitalocean|snowflake|xai|opencode-go|minimax|deepseek|" + "z-ai-coding|" + "z-ai-api|moonshot|kimi|alibaba|anthropic|openrouter|lm-studio|ollama|" + "github-feedback][/]" ) return if not ok: diff --git a/src/pythinker_code/utils/io.py b/src/pythinker_code/utils/io.py index 3170ff07..893b38ff 100644 --- a/src/pythinker_code/utils/io.py +++ b/src/pythinker_code/utils/io.py @@ -17,18 +17,35 @@ def file_lock(path: Path) -> Generator[None]: that both load before either saves drop each other's changes. Wrap the whole load → mutate → save in this lock to serialize concurrent writers. The lock file (``.lock``) is kept on disk — unlinking would split the lock across - inodes. On platforms without ``fcntl`` (Windows), this is a no-op. Blocking - (flock + small JSON I/O) — call via ``asyncio.to_thread`` from event-loop code. + inodes. This call blocks; event-loop callers must run it via ``asyncio.to_thread``. """ lock_file = path.with_name(path.name + ".lock") lock_file.parent.mkdir(parents=True, exist_ok=True) - fh = lock_file.open("a+", encoding="utf-8") + fh = lock_file.open("a+b") try: - try: - import fcntl - except ImportError: - yield + if os.name == "nt": + import msvcrt + import time + + if os.fstat(fh.fileno()).st_size == 0: + fh.write(b"\0") + fh.flush() + while True: + try: + fh.seek(0) + msvcrt.locking(fh.fileno(), msvcrt.LK_NBLCK, 1) + break + except OSError: + time.sleep(0.05) + try: + yield + finally: + with contextlib.suppress(OSError): + fh.seek(0) + msvcrt.locking(fh.fileno(), msvcrt.LK_UNLCK, 1) else: + import fcntl + fcntl.flock(fh.fileno(), fcntl.LOCK_EX) try: yield diff --git a/tasks/lessons.md b/tasks/lessons.md index ab80cb24..b1b12d44 100644 --- a/tasks/lessons.md +++ b/tasks/lessons.md @@ -3,6 +3,13 @@ Repo-specific rules accumulated from corrections and post-session reviews. Format: trigger → rule. +## User-directed execution + +- **When the user explicitly requests direct implementation and says not to create plans**, skip + optional spec and planning checkpoints for a bounded edit; inspect enough to preserve safety, + then implement and verify immediately. This never waives mandatory approval for destructive + actions, dependencies, telemetry, production changes, or other tracked safety gates. + ## Subagent orchestration - **When testing that subagent preparation failed before a prompt snapshot was written**, assert diff --git a/tasks/todo.md b/tasks/todo.md index 5c7adffb..1d6bf9a6 100644 --- a/tasks/todo.md +++ b/tasks/todo.md @@ -2,6 +2,140 @@ ## Active +### Generic auth + API-key login providers, dynamic catalog, effort mapping (2026-07-18) + +Branch: `feat/auth-login-providers`. Framing is generic (add auth + API-key login providers); +`blackbox/opencode/AUTH_PROVIDERS.md` is only the reference source, not user-facing branding. + +**Source:** `blackbox/opencode/AUTH_PROVIDERS.md` + opencode source (**MIT**; pythinker is +Apache-2.0 — compatible; ported logic carries an attribution notice). +**Delivery:** plan-first, then **sequential** delegation to Codex `gpt-5.6-sol` (high). NOT parallel — +every workstream mutates the same core files (`config.py ProviderType`, `llm.py create_llm`, +`platforms.py PLATFORMS`/`refresh_managed_models`, both UI menus, both CLI dispatchers). +**Effort ceiling:** `max` (no `ultra`) — pythinker's `ThinkingEffort` union already covers it, so the +effort workstream needs **no union change / no new-enum snapshot churn**. + +Decisions locked (AskUserQuestion): all four workstreams; plan-first; cap at `max`. + +**Headline open decision (user's call): registry-first vs additive-first.** +Recon: adding one provider today touches **8–12 hand-edited enumeration files** (no registry). +- additive-first = cheapest per phase, lowest risk, but re-pays the 8–12-file tax per provider and + that code is throwaway if a registry lands later. +- registry-first = build the abstraction up front so new providers are cheap/non-throwaway, but it is + the highest-churn refactor of working, public-compat code. +Discriminator = how many providers wanted. 4 OAuth only → additive; full roster → registry-first. + +**Size (honest):** multi-thousand-LOC across ~15–25 files; new `ProviderType` values, config keys, +CLI flags, persisted shape → tests + docs + CI snapshots (config-dump / pyinstaller / wire) each. + +Phases (each = one verified Codex delegation, sequential): +- [x] P1 — Dynamic models.dev catalog (DONE, green: `make check-pythinker-code` + focused tests 342 + passed): new `auth/models_dev.py` (httpx fetch, 5-min TTL, `fcntl.flock`, atomic write, + env override/disable, fail-open), `opencode_go.py` consumes it, tests + changelog added. + Codex `gpt-5.6-sol` candidate (commit `ee6d8384`) + architect compat-fix (`e6adfead`). + Generalize `auth/opencode_go.py` fetch into a provider-agnostic + module (`GET https://models.dev/api.json`, env override + disable flag, disk cache, 5-min TTL + + 60-min bg refresh, **stdlib `fcntl.flock`**, atomic temp+rename, **fail-open** cached→static, + never block startup). Wire one provider through it. +- [ ] P2 — Effort/family mapping: port opencode `variants()` tier-selection into the effort layer, + **capped at max**, with attribution. Extend `openai_gpt_reasoning_levels` → per-family table. +- [~] P3 — OAuth providers (sub-phased, sequential; each provider touches shared enumeration files): + - [x] P3a — shared `auth/oauth_flows.py` (device-code + loopback-PKCE) + tests. DONE, green (349 + passed). Codex candidate `cc59239a` + architect fix `79b64d42`. + - [x] P3b — GitHub Copilot (device-code; github token → copilot bearer exchange). DONE, + committed `5bd283e4`. Codex candidate `9ab0a577` (runId 9fe426d4), correctness-approved + after the token-preservation fix (F-001/F-002); clean-room green: `All checks passed` + (ruff+format+pyright+ty) + `458 passed`. Two-refresh regression test present. **PENDING + LIVE VERIFICATION**: no offline gate exercises client_id/exchange/headers/URL — needs a + real `pythinker login --copilot` + one chat call before "truly done". **Design + FINAL (primary-source verified 2026-07-18):** client_id `Iv1.b507a08c87ecfe98` + scope + `read:user` (the exchange-proven pair from ericc-ch/copilot-api; NOT opencode's + `Ov23li…` which is proven only with the no-exchange direct-token path). Device: + POST github.com/login/device/code + poll .../login/oauth/access_token (both need + `Accept: application/json` → thread a `headers` param through oauth_flows + `_post_form`/`request_device_code`/`poll_device_token`). Exchange: + GET api.github.com/copilot_internal/v2/token, hdrs `Authorization: token `, + `Editor-Version: vscode/`, `Editor-Plugin-Version: copilot-chat/0.26.7`, + `User-Agent: GitHubCopilotChat/0.26.7`, `X-GitHub-Api-Version: 2025-04-01` + → `{token, expires_at, refresh_in}`. Store OAuthToken(access=bearer, + refresh=gh_token, expires_at); `_refresh_token_for_ref` re-runs exchange from gh_token. + Provider: type `openai_legacy`, base_url `https://api.githubcopilot.com` (**NO /v1** — + SDK appends /chat/completions to root), oauth ref `oauth/github-copilot`, + custom_headers = the copilot chat headers (Copilot-Integration-Id: vscode-chat + editor + hdrs + Openai-Intent: conversation-panel; **NO Authorization** — bearer flows via + resolve_api_key→api_key). Skip-guard `managed:copilot` in refresh_managed_models. + **Scope: github.com individual only.** Business/Enterprise (endpoints.api routing, + api.business/individual.*) DEFERRED — do not claim exchange fixes Business (opencode + #23540) while hardcoding the individual host. **Acceptance:** offline gates green ≠ done; + none of client_id/exchange/headers/URL are gate-exercisable → requires live + `pythinker login --copilot` + one real chat call before marking done. + - [x] P3c — xAI/Grok (browser loopback + device-code). DONE, committed `b586080c`. Codex + candidate `1419210a` (runId f12a8003) — verification-failed on strict pyright + (`.get()` on an isinstance-narrowed bare `dict` → reportUnknownMemberType/Argument); + salvaged via `git checkout -- .`, added a cast'd `_error_description` + helper, re-ran gates (`All checks passed` + `2741 passed`). Two OAuth methods only + (loopback port 56121 + `plan=generic`/OIDC `nonce`; device-code); no API-key method; + openai_legacy → api.x.ai/v1; rotating refresh persisted by OAuthManager. **PENDING LIVE + VERIFICATION** (login flow untested against real auth.x.ai). + - [x] P3d — DigitalOcean. DONE. Lane A `run_loopback_implicit_flow` helper committed `48a1fdbb` + (salvaged after base-changed abort); Lane B provider+wiring+tests committed `962ef990` + (salvaged after producer verify-fail, fixed reportUnnecessaryIsInstance/Cast + + ruff-format nit). Gates green locally: make check-pythinker-code + pytest tests/auth + tests/ui_and_conv tests/cli. **PENDING LIVE VERIFY** (implicit + browser-JS + real DO acct). + - [~] P3d(orig) — DigitalOcean. **SCOPE = FULL ROBUST BUILD (user-confirmed 2026-07-18).** OAuth + IMPLICIT flow (response_type=token; token in URL fragment) → needs a NEW reusable + `run_loopback_implicit_flow` helper in oauth_flows.py that serves an HTML-bootstrap page + (inline JS reads location.hash, POSTs {access_token,expires_in,state} to a pinned-port + /auth/token) — P3a's authorization-code loopback does NOT cover this. Token stored AS AN + API KEY (no refresh, ~30d; re-login on expiry) → NO oauth ref / NO _refresh_token_for_ref + branch. Provider openai_legacy → base_url https://inference.do-ai.run/v1, api_key=token. + Dynamic model catalog: GET https://api.digitalocean.com/v2/gen-ai/models/routers (Bearer) + → model_routers[].name → seed 'router:' models at login (login still succeeds if + fetch fails; seed none + warn). Verified constants: client_id + b1a6c5158156caac821fd1b30253ca8acb52454a48fa744420e41889cb589f82; authorize + https://cloud.digitalocean.com/v1/oauth/authorize; redirect http://localhost:1456/auth/callback; + scope 'genai:read inference:query'. Recon ad3fed51 DONE (advisor-confirmed); split into + two serialized lanes to fit the 30-min Codex cap: + - Lane A (DISPATCHED, task kz00tnywe, Codex gpt-5.6-sol/high): add + `run_loopback_implicit_flow` + `ImplicitAuthorization` + content-type writer to + oauth_flows.py + tests. Self-contained, no src caller yet. + - Lane B (after A integrates): new auth/digitalocean.py (DeepSeek storage template: + LLMProvider openai_legacy, api_key=SecretStr(token), NO oauth) + __init__/platforms/ + shell/cli wiring (mirror xai) + tests. platforms.py skip-guard for managed:digitalocean. + Empty-routers guard: still persist key + save_config, seed 0 models, guard default_model. + **PENDING LIVE VERIFY** (implicit + browser-JS + real DO account — largely untestable offline). + - [x] P3e — Snowflake Cortex. DONE, committed `1847ab9d` (salvaged after producer verify-fail: + fixed reportPrivateUsage on oauth_flows._post_form → local _post_form; ruff-format nit). + Gates green: make check-pythinker-code + pytest tests/auth tests/ui_and_conv tests/cli + (497 passed). One shared-oauth.py change = the account-parsing refresh branch. + **PENDING LIVE VERIFY** (real Snowflake acct + browser). Known live-inference gap: cortexFetch + transforms not replicated by openai_legacy — follow-up, does not block login. + - [~] P3e(orig) — Snowflake Cortex. **SCOPE = FULL ROBUST ACCOUNT-SCOPED BUILD (user-confirmed + 2026-07-18).** Materially the most complex P3 provider (NOT "simpler like xai"): + account-scoped OAuth + inference base_url (account + optional role PROMPTED at login), + role-dependent scope (`refresh_token session:role:`), HTTP Basic client creds + (base64 `LOCAL_APPLICATION:LOCAL_APPLICATION`), loopback-PKCE (reuses P3a + run_loopback_pkce_flow, redirect_path "/"). Endpoints: + https://.snowflakecomputing.com/oauth/{authorize,token-request}. **Account-aware + refresh: encode account into OAuthRef key `oauth/snowflake-cortex/` and parse it + in oauth.py refresh dispatch → the ONE shared-oauth.py change.** Provider openai_legacy, + base_url per-account = https://.snowflakecomputing.com/api/v2/cortex (VERIFY exact + path). KNOWN LIVE-INFERENCE GAP: cortexFetch transforms (max_tokens→max_completion_tokens, + 400 "conversation complete"→stop, streaming role:""→"assistant") NOT replicated by + openai_legacy — does not block login; chat may need follow-up. Recon afd25065 mapping exact + wiring + primary-source constant verification. **PENDING LIVE VERIFY** (real Snowflake acct). +- [ ] P4 — API-key providers: batch the models.dev env-keyed providers through the P1 catalog. +- [ ] P0/P5 — Registry refactor (only if registry-first chosen; else optional last). + +Non-negotiables per Codex spec: full `make check-pythinker-code && make test-pythinker-code` gate; +`## Unreleased` changelog line; deliberate snapshot updates; models.dev fail-open + no new dep; +attribution on ported files; root-cause robust design (no workarounds). + +Out of scope (logged): `ultra` effort level (dropped, cap at max); any external endpoint beyond +models.dev without approval. + +Review: _pending first delegation._ + ### Implementer-agent deepening: lighter/smarter/more robust (2026-07-17) Architecture review found: every implementer spawn carries ~7,300 words of prompt (root diff --git a/tests/auth/test_copilot_auth.py b/tests/auth/test_copilot_auth.py new file mode 100644 index 00000000..68454382 --- /dev/null +++ b/tests/auth/test_copilot_auth.py @@ -0,0 +1,406 @@ +from __future__ import annotations + +from collections.abc import Mapping +from types import SimpleNamespace +from typing import Any + +import aiohttp +import pytest +from pydantic import SecretStr + +from pythinker_code.auth.models_dev import CatalogResult, CatalogStatus +from pythinker_code.auth.oauth import ( + OAuthError, + OAuthManager, + OAuthToken, + OAuthUnauthorized, + load_tokens, + save_tokens, +) +from pythinker_code.auth.oauth_flows import DeviceCode +from pythinker_code.config import Config, LLMModel, LLMProvider, OAuthRef + + +def _mock_unavailable_catalog(monkeypatch: pytest.MonkeyPatch) -> None: + from pythinker_code.auth import copilot + + async def _fake() -> CatalogResult: + return CatalogResult({}, CatalogStatus.UNAVAILABLE, "none") + + monkeypatch.setattr(copilot, "get_models_dev_catalog", _fake) + + +@pytest.mark.asyncio +async def test_discover_copilot_models_logs_empty_authoritative_fallback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from pythinker_code.auth import copilot + + messages: list[str] = [] + + async def empty_catalog() -> CatalogResult: + return CatalogResult({}, CatalogStatus.OK, "network") + + monkeypatch.setattr(copilot, "get_models_dev_catalog", empty_catalog) + monkeypatch.setattr( + copilot, + "logger", + SimpleNamespace(debug=lambda message, **_kwargs: messages.append(message)), + ) + + assert await copilot._discover_copilot_models() == copilot.GITHUB_COPILOT_MODELS + assert any("no usable" in message for message in messages) + + +@pytest.mark.asyncio +async def test_login_copilot_saves_two_tokens_provider_and_models( + monkeypatch: pytest.MonkeyPatch, + tmp_path, +) -> None: + from pythinker_code.auth import copilot + + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + config = Config(is_from_default_location=True) + device_calls: list[dict[str, Any]] = [] + poll_calls: list[dict[str, Any]] = [] + + async def fake_request_device_code(**kwargs: Any) -> DeviceCode: + device_calls.append(kwargs) + return DeviceCode( + user_code="ABCD-EFGH", + verification_uri="https://github.com/login/device", + device_code="device-secret", + interval=5, + expires_in=900, + ) + + async def fake_poll_device_token(**kwargs: Any) -> dict[str, Any]: + poll_calls.append(kwargs) + return {"access_token": "github-oauth-token", "token_type": "bearer"} + + async def fake_refresh_copilot_token(github_token: str) -> OAuthToken: + assert github_token == "github-oauth-token" + return OAuthToken( + access_token="copilot-bearer", + refresh_token="", + expires_at=2_000_000_000, + scope="read:user", + token_type="Bearer", + expires_in=1500, + ) + + monkeypatch.setattr(copilot, "request_device_code", fake_request_device_code) + monkeypatch.setattr(copilot, "poll_device_token", fake_poll_device_token) + monkeypatch.setattr(copilot, "refresh_copilot_token", fake_refresh_copilot_token) + _mock_unavailable_catalog(monkeypatch) + + events = [event async for event in copilot.login_copilot(config, open_browser=False)] + + assert [event.type for event in events] == ["verification_url", "waiting", "success"] + assert device_calls[0]["headers"] == {"Accept": "application/json"} + assert poll_calls[0]["headers"] == {"Accept": "application/json"} + oauth_ref = OAuthRef(storage="file", key="oauth/github-copilot") + stored = load_tokens(oauth_ref) + assert stored is not None + assert stored.access_token == "copilot-bearer" + assert stored.refresh_token == "github-oauth-token" + + provider = config.providers["managed:copilot"] + assert provider.type == "openai_legacy" + assert provider.base_url == "https://api.githubcopilot.com" + assert provider.oauth == oauth_ref + assert provider.custom_headers is not None + assert provider.custom_headers["Copilot-Integration-Id"] == "vscode-chat" + assert provider.custom_headers["X-GitHub-Api-Version"] == "2025-04-01" + assert "Authorization" not in provider.custom_headers + assert {model.provider for model in config.models.values()} == {"managed:copilot"} + + +@pytest.mark.asyncio +async def test_copilot_persistence_errors_hide_internal_details( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from pythinker_code.auth import copilot + + diagnostic = "permission denied at /private/credentials/github-copilot.json" + config = Config(is_from_default_location=True) + + async def fake_request_device_code(**_kwargs: Any) -> DeviceCode: + return DeviceCode( + user_code="ABCD-EFGH", + verification_uri="https://github.com/login/device", + device_code="device-secret", + interval=5, + expires_in=900, + ) + + async def fake_poll_device_token(**_kwargs: Any) -> dict[str, Any]: + return {"access_token": "github-oauth-token"} + + async def fake_refresh_copilot_token(_github_token: str) -> OAuthToken: + return OAuthToken.from_response( + { + "access_token": "copilot-bearer", + "refresh_token": "github-oauth-token", + "expires_in": 1500, + } + ) + + async def fail_persistence(*_args: object, **_kwargs: object) -> None: + raise OSError(diagnostic) + + monkeypatch.setattr(copilot, "request_device_code", fake_request_device_code) + monkeypatch.setattr(copilot, "poll_device_token", fake_poll_device_token) + monkeypatch.setattr(copilot, "refresh_copilot_token", fake_refresh_copilot_token) + monkeypatch.setattr(copilot, "persist_login", fail_persistence) + monkeypatch.setattr(copilot, "persist_logout", fail_persistence) + _mock_unavailable_catalog(monkeypatch) + + login_events = [event async for event in copilot.login_copilot(config, open_browser=False)] + logout_events = [event async for event in copilot.logout_copilot(config)] + + assert login_events[-1].message == "Failed to save GitHub Copilot login." + assert logout_events[-1].message == "Failed to log out of GitHub Copilot." + assert diagnostic not in login_events[-1].json + assert diagnostic not in logout_events[-1].json + + +class _ExchangeResponse: + def __init__(self, status: int, payload: object) -> None: + self.status = status + self.payload = payload + + async def __aenter__(self) -> _ExchangeResponse: + return self + + async def __aexit__(self, exc_type: object, exc: object, traceback: object) -> None: + return None + + async def json(self, *, content_type: object = None) -> object: + return self.payload + + +class _ExchangeSession: + def __init__( + self, + response: _ExchangeResponse, + calls: list[tuple[str, Mapping[str, str]]], + ) -> None: + self.response = response + self.calls = calls + + async def __aenter__(self) -> _ExchangeSession: + return self + + async def __aexit__(self, exc_type: object, exc: object, traceback: object) -> None: + return None + + def get(self, endpoint: str, *, headers: Mapping[str, str]) -> _ExchangeResponse: + self.calls.append((endpoint, headers)) + return self.response + + +@pytest.mark.asyncio +async def test_refresh_copilot_token_exchanges_github_token( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from pythinker_code.auth import copilot + + calls: list[tuple[str, Mapping[str, str]]] = [] + response = _ExchangeResponse( + 200, + {"token": "copilot-bearer", "expires_at": 2_000_000_000, "refresh_in": 1500}, + ) + monkeypatch.setattr( + copilot, + "new_client_session", + lambda: _ExchangeSession(response, calls), + ) + + token = await copilot.refresh_copilot_token("github-oauth-token") + + assert token.access_token == "copilot-bearer" + assert token.refresh_token == "github-oauth-token" + assert token.expires_at == 2_000_000_000 + assert token.expires_in == 1500 + assert calls[0][0] == "https://api.github.com/copilot_internal/v2/token" + assert calls[0][1]["Authorization"] == "token github-oauth-token" + assert calls[0][1]["Editor-Version"] == "vscode/1.99.0" + assert calls[0][1]["Editor-Plugin-Version"] == "copilot-chat/0.26.7" + assert calls[0][1]["User-Agent"] == "GitHubCopilotChat/0.26.7" + assert calls[0][1]["X-GitHub-Api-Version"] == "2025-04-01" + + +@pytest.mark.asyncio +async def test_oauth_manager_preserves_github_token_across_copilot_refreshes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from pythinker_code.auth import copilot + + responses = [ + _ExchangeResponse( + 200, + {"token": "copilot-bearer-1", "expires_at": 2_000_000_000, "refresh_in": 1500}, + ), + _ExchangeResponse( + 200, + {"token": "copilot-bearer-2", "expires_at": 2_000_001_500, "refresh_in": 1500}, + ), + ] + calls: list[tuple[str, Mapping[str, str]]] = [] + monkeypatch.setattr( + copilot, + "new_client_session", + lambda: _ExchangeSession(responses.pop(0), calls), + ) + manager = OAuthManager(Config()) + oauth_ref = OAuthRef(storage="file", key=copilot.GITHUB_COPILOT_OAUTH_KEY) + + first = await manager._refresh_token_for_ref(oauth_ref, "github-oauth-token") + second = await manager._refresh_token_for_ref(oauth_ref, first.refresh_token) + + assert first.access_token == "copilot-bearer-1" + assert first.refresh_token == "github-oauth-token" + assert second.access_token == "copilot-bearer-2" + assert second.refresh_token == "github-oauth-token" + assert [headers["Authorization"] for _, headers in calls] == [ + "token github-oauth-token", + "token github-oauth-token", + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [401, 403]) +async def test_refresh_copilot_token_raises_unauthorized( + monkeypatch: pytest.MonkeyPatch, + status: int, +) -> None: + from pythinker_code.auth import copilot + + response = _ExchangeResponse(status, {"message": "Bad credentials"}) + monkeypatch.setattr( + copilot, + "new_client_session", + lambda: _ExchangeSession(response, []), + ) + + with pytest.raises(OAuthUnauthorized): + await copilot.refresh_copilot_token("revoked-github-token") + + +class _RaisingSession: + def __init__(self, exc: BaseException) -> None: + self.exc = exc + + async def __aenter__(self) -> _RaisingSession: + return self + + async def __aexit__(self, exc_type: object, exc: object, traceback: object) -> None: + return None + + def get(self, endpoint: str, *, headers: Mapping[str, str]) -> object: + raise self.exc + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "response", + [ + _ExchangeResponse(200, ["not", "a", "dict"]), # malformed body + _ExchangeResponse(200, {"expires_at": 2_000_000_000, "refresh_in": 1500}), # no token + _ExchangeResponse(200, {"token": "", "expires_at": 2_000_000_000, "refresh_in": 1500}), + _ExchangeResponse(500, {"message": "server error"}), # non-auth HTTP failure + ], +) +async def test_refresh_copilot_token_rejects_malformed_responses( + monkeypatch: pytest.MonkeyPatch, + response: _ExchangeResponse, +) -> None: + from pythinker_code.auth import copilot + + monkeypatch.setattr(copilot, "new_client_session", lambda: _ExchangeSession(response, [])) + + with pytest.raises(OAuthError): + await copilot.refresh_copilot_token("github-oauth-token") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("exc", [aiohttp.ClientError(), TimeoutError(), OSError()]) +async def test_refresh_copilot_token_wraps_transport_failures( + monkeypatch: pytest.MonkeyPatch, + exc: BaseException, +) -> None: + from pythinker_code.auth import copilot + + monkeypatch.setattr(copilot, "new_client_session", lambda: _RaisingSession(exc)) + + with pytest.raises(OAuthError): + await copilot.refresh_copilot_token("github-oauth-token") + + +@pytest.mark.asyncio +async def test_logout_copilot_removes_tokens_provider_models_and_repairs_default( + monkeypatch: pytest.MonkeyPatch, + tmp_path, +) -> None: + from pythinker_code.auth.copilot import ( + GITHUB_COPILOT_OAUTH_KEY, + GITHUB_COPILOT_PROVIDER_KEY, + logout_copilot, + ) + + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + oauth_ref = OAuthRef(storage="file", key=GITHUB_COPILOT_OAUTH_KEY) + save_tokens( + oauth_ref, + OAuthToken( + access_token="copilot-bearer", + refresh_token="github-oauth-token", + expires_at=2_000_000_000, + scope="read:user", + token_type="Bearer", + expires_in=1500, + ), + ) + config = Config( + is_from_default_location=True, + default_model="copilot/gpt-4.1", + providers={ + GITHUB_COPILOT_PROVIDER_KEY: LLMProvider( + type="openai_legacy", + base_url="https://api.githubcopilot.com", + api_key=SecretStr(""), + oauth=oauth_ref, + ), + "fallback": LLMProvider( + type="openai_legacy", + base_url="https://example.test/v1", + api_key=SecretStr("test"), + ), + }, + models={ + "copilot/gpt-4.1": LLMModel( + provider=GITHUB_COPILOT_PROVIDER_KEY, + model="gpt-4.1", + max_context_size=1_000_000, + ), + "copilot/gpt-4o": LLMModel( + provider=GITHUB_COPILOT_PROVIDER_KEY, + model="gpt-4o", + max_context_size=128_000, + ), + "fallback/model": LLMModel( + provider="fallback", + model="model", + max_context_size=32_000, + ), + }, + ) + + events = [event async for event in logout_copilot(config)] + + assert [event.type for event in events] == ["success"] + assert load_tokens(oauth_ref) is None + assert GITHUB_COPILOT_PROVIDER_KEY not in config.providers + assert all(model.provider != GITHUB_COPILOT_PROVIDER_KEY for model in config.models.values()) + assert config.default_model == "fallback/model" diff --git a/tests/auth/test_digitalocean_auth.py b/tests/auth/test_digitalocean_auth.py new file mode 100644 index 00000000..0cb14db2 --- /dev/null +++ b/tests/auth/test_digitalocean_auth.py @@ -0,0 +1,288 @@ +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import aiohttp +import pytest +from pydantic import SecretStr + +from pythinker_code.config import Config, LLMModel, LLMProvider + + +class _RouterResponse: + def __init__(self, status: int, payload: object) -> None: + self.status = status + self.payload = payload + + async def __aenter__(self) -> _RouterResponse: + return self + + async def __aexit__(self, exc_type: object, exc: object, traceback: object) -> None: + return None + + async def json(self, *, content_type: object = None) -> object: + if isinstance(self.payload, Exception): + raise self.payload + return self.payload + + +class _RouterSession: + def __init__( + self, response: _RouterResponse, calls: list[tuple[str, Mapping[str, str]]] + ) -> None: + self.response = response + self.calls = calls + + async def __aenter__(self) -> _RouterSession: + return self + + async def __aexit__(self, exc_type: object, exc: object, traceback: object) -> None: + return None + + def get(self, url: str, *, headers: Mapping[str, str]) -> _RouterResponse: + self.calls.append((url, dict(headers))) + return self.response + + +class _RaisingRouterSession: + def __init__(self, exc: BaseException) -> None: + self.exc = exc + + async def __aenter__(self) -> _RaisingRouterSession: + return self + + async def __aexit__(self, exc_type: object, exc: object, traceback: object) -> None: + return None + + def get(self, url: str, *, headers: Mapping[str, str]) -> object: + raise self.exc + + +def test_apply_digitalocean_config_writes_provider_models_and_default() -> None: + from pythinker_code.auth.digitalocean import ( + DIGITALOCEAN_BASE_URL, + DIGITALOCEAN_PROVIDER_KEY, + _apply_digitalocean_config, + ) + + config = Config(is_from_default_location=True) + + _apply_digitalocean_config(config, SecretStr("do-test"), ("production", "staging")) + + assert set(config.providers) == {DIGITALOCEAN_PROVIDER_KEY} + provider = config.providers[DIGITALOCEAN_PROVIDER_KEY] + assert provider.type == "openai_legacy" + assert provider.base_url == DIGITALOCEAN_BASE_URL + assert provider.api_key.get_secret_value() == "do-test" + assert provider.oauth is None + assert config.models["digitalocean/router:production"].provider == DIGITALOCEAN_PROVIDER_KEY + assert config.models["digitalocean/router:production"].model == "router:production" + assert config.models["digitalocean/router:production"].max_context_size == 128_000 + assert config.models["digitalocean/router:production"].display_name == "production" + assert config.models["digitalocean/router:staging"].model == "router:staging" + assert config.default_model == "digitalocean/router:production" + + +def test_apply_digitalocean_config_handles_empty_router_catalog() -> None: + from pythinker_code.auth.digitalocean import ( + DIGITALOCEAN_PROVIDER_KEY, + _apply_digitalocean_config, + ) + + fallback_alias = "openai/existing" + config = Config( + is_from_default_location=True, + default_model=fallback_alias, + providers={ + "managed:openai": LLMProvider( + type="openai_legacy", + base_url="https://example.test/v1", + api_key=SecretStr("existing"), + ) + }, + models={ + fallback_alias: LLMModel( + provider="managed:openai", + model="existing", + max_context_size=1, + ) + }, + ) + + _apply_digitalocean_config(config, SecretStr("tok"), ()) + + assert DIGITALOCEAN_PROVIDER_KEY in config.providers + assert not any(model.provider == DIGITALOCEAN_PROVIDER_KEY for model in config.models.values()) + assert config.default_model == fallback_alias + assert config.default_model in config.models + + +@pytest.mark.asyncio +async def test_fetch_router_catalog_returns_ok_with_names( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from pythinker_code.auth import digitalocean as do + + calls: list[tuple[str, Mapping[str, str]]] = [] + response = _RouterResponse( + 200, + {"model_routers": [{"name": "primary"}, {"name": "fallback"}, {"no_name": 1}]}, + ) + monkeypatch.setattr(do, "new_client_session", lambda: _RouterSession(response, calls)) + + catalog = await do._fetch_router_catalog("tok") + + assert catalog.status is do.RouterDiscovery.OK + assert catalog.names == ("primary", "fallback") + assert calls[0][0] == do.DIGITALOCEAN_ROUTERS_URL + assert calls[0][1]["Authorization"] == "Bearer tok" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("response", "expected"), + [ + (_RouterResponse(200, {"model_routers": []}), "EMPTY"), + (_RouterResponse(200, {"model_routers": [{"no_name": 1}]}), "EMPTY"), + (_RouterResponse(401, {"id": "unauthorized"}), "UNAUTHORIZED"), + (_RouterResponse(403, {"id": "forbidden"}), "UNAUTHORIZED"), + (_RouterResponse(500, {"id": "server_error"}), "UNAVAILABLE"), + (_RouterResponse(200, ["not", "a", "dict"]), "MALFORMED"), + (_RouterResponse(200, {"model_routers": "nope"}), "MALFORMED"), + (_RouterResponse(200, ValueError("bad json")), "MALFORMED"), + ], +) +async def test_fetch_router_catalog_distinguishes_outcomes( + monkeypatch: pytest.MonkeyPatch, + response: _RouterResponse, + expected: str, +) -> None: + from pythinker_code.auth import digitalocean as do + + monkeypatch.setattr(do, "new_client_session", lambda: _RouterSession(response, [])) + + catalog = await do._fetch_router_catalog("tok") + + assert catalog.status.name == expected + assert catalog.names == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize("exc", [aiohttp.ClientError(), TimeoutError(), OSError()]) +async def test_fetch_router_catalog_treats_transport_errors_as_unavailable( + monkeypatch: pytest.MonkeyPatch, + exc: BaseException, +) -> None: + from pythinker_code.auth import digitalocean as do + + monkeypatch.setattr(do, "new_client_session", lambda: _RaisingRouterSession(exc)) + + catalog = await do._fetch_router_catalog("tok") + + assert catalog.status is do.RouterDiscovery.UNAVAILABLE + + +@pytest.mark.asyncio +async def test_login_digitalocean_saves_discovered_routers_without_leaking_token( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from pythinker_code.auth import digitalocean as do + from pythinker_code.auth.oauth_flows import ImplicitAuthorization + + access_token = "secret-digitalocean-token" + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + config = Config(is_from_default_location=True) + + async def fake_implicit_flow(**kwargs: Any) -> ImplicitAuthorization: + return ImplicitAuthorization(access_token, None, "state") + + response = _RouterResponse(200, {"model_routers": [{"name": "primary"}, {"name": "fallback"}]}) + monkeypatch.setattr(do, "run_loopback_implicit_flow", fake_implicit_flow) + monkeypatch.setattr(do, "new_client_session", lambda: _RouterSession(response, [])) + + events = [event async for event in do.login_digitalocean(config, open_browser=False)] + + assert [event.type for event in events] == ["waiting", "success"] + assert config.default_model == "digitalocean/router:primary" + provider = config.providers[do.DIGITALOCEAN_PROVIDER_KEY] + assert provider.api_key.get_secret_value() == access_token + assert access_token not in "\n".join(f"{event!r}\n{event.json}" for event in events) + assert (tmp_path / "config.toml").exists() + + +@pytest.mark.asyncio +async def test_digitalocean_persistence_errors_hide_internal_details( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from pythinker_code.auth import digitalocean as do + from pythinker_code.auth.oauth_flows import ImplicitAuthorization + + diagnostic = "permission denied at /private/config.toml" + config = Config(is_from_default_location=True) + + async def fake_implicit_flow(**_kwargs: Any) -> ImplicitAuthorization: + return ImplicitAuthorization("access-token", None, "state") + + async def fake_catalog(_access_token: str) -> do.RouterCatalog: + return do.RouterCatalog(do.RouterDiscovery.OK, ("primary",)) + + async def fail_persistence(*_args: object, **_kwargs: object) -> None: + raise OSError(diagnostic) + + monkeypatch.setattr(do, "run_loopback_implicit_flow", fake_implicit_flow) + monkeypatch.setattr(do, "_fetch_router_catalog", fake_catalog) + monkeypatch.setattr(do, "persist_config_change", fail_persistence) + + login_events = [event async for event in do.login_digitalocean(config)] + logout_events = [event async for event in do.logout_digitalocean(config)] + + assert login_events[-1].message == "Failed to save DigitalOcean login." + assert logout_events[-1].message == "Failed to log out of DigitalOcean." + assert diagnostic not in login_events[-1].json + assert diagnostic not in logout_events[-1].json + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "response", + [ + _RouterResponse(200, {"model_routers": []}), # empty + _RouterResponse(401, {"id": "unauthorized"}), # unauthorized + _RouterResponse(500, {"id": "server_error"}), # outage + ], +) +async def test_login_digitalocean_reports_router_discovery_failures( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + response: _RouterResponse, +) -> None: + from pythinker_code.auth import digitalocean as do + from pythinker_code.auth.oauth_flows import ImplicitAuthorization + + access_token = "secret-empty-router-token" + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + config = Config(is_from_default_location=True) + + async def fake_implicit_flow(**kwargs: Any) -> ImplicitAuthorization: + return ImplicitAuthorization(access_token, None, "state") + + monkeypatch.setattr(do, "run_loopback_implicit_flow", fake_implicit_flow) + monkeypatch.setattr(do, "new_client_session", lambda: _RouterSession(response, [])) + + events = [event async for event in do.login_digitalocean(config)] + + # Sign-in is saved but the degraded discovery is surfaced as an info event + # before success, never silently swallowed. + assert [event.type for event in events] == ["waiting", "info", "success"] + assert do.DIGITALOCEAN_PROVIDER_KEY in config.providers + assert not any( + model.provider == do.DIGITALOCEAN_PROVIDER_KEY for model in config.models.values() + ) + assert access_token not in "\n".join(f"{event!r}\n{event.json}" for event in events) + assert "configured with model" not in events[-1].message + assert events[-1].message == ( + "DigitalOcean credentials saved; no inference routers are configured." + ) diff --git a/tests/auth/test_models_dev.py b/tests/auth/test_models_dev.py new file mode 100644 index 00000000..4ab1159e --- /dev/null +++ b/tests/auth/test_models_dev.py @@ -0,0 +1,175 @@ +from __future__ import annotations + +import json +import os +import time + +import pytest + +from pythinker_code.auth import models_dev + + +def _catalog_payload() -> dict[str, object]: + return { + "example": { + "name": "Example Provider", + "env": ["EXAMPLE_API_KEY"], + "npm": "@ai-sdk/example", + "api": "https://example.test/v1", + "models": { + "reasoner": { + "id": "reasoner", + "name": "Example Reasoner", + "reasoning": True, + "limit": {"context": 128_000, "output": 8_192}, + "modalities": {"input": ["text", "image", "video"], "output": ["text"]}, + } + }, + } + } + + +def _write_cache(tmp_path, payload: object) -> None: + (tmp_path / "models_dev.json").write_text( + json.dumps(payload), + encoding="utf-8", + ) + + +def test_parse_models_dev_catalog_normalizes_model_metadata(): + catalog = models_dev.parse_models_dev_catalog(_catalog_payload()) + + provider = catalog["example"] + model = models_dev.get_provider_models(catalog, "example")["reasoner"] + assert provider.display_name == "Example Provider" + assert provider.env == ("EXAMPLE_API_KEY",) + assert model.display_name == "Example Reasoner" + assert model.context_length == 128_000 + assert model.reasoning is True + assert model.supports_image_input is True + assert model.supports_video_input is True + assert model.npm == "@ai-sdk/example" + + +@pytest.mark.asyncio +async def test_fresh_cache_is_reused_without_network(monkeypatch, tmp_path): + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + monkeypatch.delenv("PYTHINKER_DISABLE_MODELS_FETCH", raising=False) + _write_cache(tmp_path, _catalog_payload()) + + async def unexpected_fetch(): + pytest.fail("fresh cache must not trigger a network fetch") + + monkeypatch.setattr(models_dev, "_fetch_catalog_payload", unexpected_fetch) + + result = await models_dev.get_models_dev_catalog() + + assert result.status is models_dev.CatalogStatus.CACHED + assert result.source == "cache" + assert result.catalog["example"].models["reasoner"].context_length == 128_000 + + +@pytest.mark.asyncio +async def test_network_failure_fails_open_to_stale_cache(monkeypatch, tmp_path): + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + monkeypatch.delenv("PYTHINKER_DISABLE_MODELS_FETCH", raising=False) + _write_cache(tmp_path, _catalog_payload()) + stale_time = time.time() - models_dev.MODELS_DEV_CACHE_TTL_SECONDS - 1 + os.utime(tmp_path / "models_dev.json", (stale_time, stale_time)) + + async def failed_fetch(): + raise RuntimeError("network unavailable") + + monkeypatch.setattr(models_dev, "_fetch_catalog_payload", failed_fetch) + + result = await models_dev.get_models_dev_catalog() + + # A stale-cache fallback must announce that it is degraded, not masquerade + # as an authoritative fresh load. + assert result.status is models_dev.CatalogStatus.STALE + assert result.is_authoritative is False + assert result.catalog["example"].models["reasoner"].display_name == "Example Reasoner" + + +@pytest.mark.asyncio +async def test_disabled_fetch_without_cache_returns_empty(monkeypatch, tmp_path): + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + monkeypatch.setenv("PYTHINKER_DISABLE_MODELS_FETCH", "yes") + + async def unexpected_fetch(): + pytest.fail("disabled model fetching must not access the network") + + monkeypatch.setattr(models_dev, "_fetch_catalog_payload", unexpected_fetch) + + result = await models_dev.get_models_dev_catalog() + + assert result.status is models_dev.CatalogStatus.DISABLED + assert result.catalog == {} + + +def _chat_catalog_payload() -> dict[str, object]: + return { + "demo": { + "name": "Demo", + "models": { + "chat-small": { + "id": "chat-small", + "name": "Chat Small", + "limit": {"context": 64_000}, + "modalities": {"input": ["text"], "output": ["text"]}, + }, + "image-gen": { + "id": "image-gen", + "name": "Image Gen", + "modalities": {"input": ["text"], "output": ["image"]}, + }, + "text-embedding-3": { + "id": "text-embedding-3", + "name": "Embedding", + "modalities": {"input": ["text"], "output": ["text"]}, + }, + }, + } + } + + +def test_chat_models_excludes_non_chat_output_and_embeddings(): + catalog = models_dev.parse_models_dev_catalog(_chat_catalog_payload()) + + ids = set(models_dev.chat_models(catalog, "demo")) + + # Image generator excluded by output modality; embedding excluded by id. + assert ids == {"chat-small"} + + +def test_build_catalog_models_applies_default_context_and_ordering(): + catalog = models_dev.parse_models_dev_catalog(_chat_catalog_payload()) + + built = models_dev.build_catalog_models(catalog, "demo", default_context=100_000) + + assert [model.model_id for model in built] == ["chat-small"] + assert built[0].max_context_size == 64_000 + + empty = models_dev.build_catalog_models(catalog, "missing", default_context=100_000) + assert empty == () + + +def test_build_catalog_models_uses_default_context_when_limit_missing(): + payload = { + "demo": { + "name": "Demo", + "models": { + "chat-nolimit": { + "id": "chat-nolimit", + "name": "Chat No Limit", + "modalities": {"input": ["text"], "output": ["text"]}, + }, + }, + } + } + catalog = models_dev.parse_models_dev_catalog(payload) + + built = models_dev.build_catalog_models(catalog, "demo", default_context=100_000) + + assert [model.model_id for model in built] == ["chat-nolimit"] + assert built[0].max_context_size == 100_000 diff --git a/tests/auth/test_oauth_flows.py b/tests/auth/test_oauth_flows.py new file mode 100644 index 00000000..1ee66f56 --- /dev/null +++ b/tests/auth/test_oauth_flows.py @@ -0,0 +1,618 @@ +from __future__ import annotations + +import asyncio +import base64 +import hashlib +import json +import re +from collections.abc import Callable, Mapping +from typing import Any, cast +from urllib.parse import parse_qs, urlsplit + +import pytest + +from pythinker_code.auth.oauth import OAuthDeviceExpired, OAuthError +from pythinker_code.auth.oauth_flows import ( + DeviceCode, + ImplicitAuthorization, + OAuthAccessDenied, + OAuthStateMismatch, + _handle_implicit_loopback_callback, + generate_pkce, + generate_state, + poll_device_token, + request_device_code, + run_loopback_implicit_flow, + run_loopback_pkce_flow, +) + + +class _FakeResponse: + def __init__(self, status: int, payload: dict[str, Any]) -> None: + self.status = status + self.payload = payload + + async def __aenter__(self) -> _FakeResponse: + return self + + async def __aexit__(self, exc_type: object, exc: object, traceback: object) -> None: + return None + + async def json(self, *, content_type: object = None) -> dict[str, Any]: + return self.payload + + +class _FakeSession: + def __init__( + self, + responses: list[_FakeResponse], + calls: list[tuple[str, dict[str, str], Mapping[str, str] | None]], + ) -> None: + self.responses = responses + self.calls = calls + + async def __aenter__(self) -> _FakeSession: + return self + + async def __aexit__(self, exc_type: object, exc: object, traceback: object) -> None: + return None + + def post( + self, + endpoint: str, + *, + data: dict[str, str], + headers: Mapping[str, str] | None = None, + ) -> _FakeResponse: + self.calls.append((endpoint, data, headers)) + return self.responses.pop(0) + + +def _mock_http( + monkeypatch: pytest.MonkeyPatch, + *responses: tuple[int, dict[str, Any]], +) -> list[tuple[str, dict[str, str], Mapping[str, str] | None]]: + queued = [_FakeResponse(status, payload) for status, payload in responses] + calls: list[tuple[str, dict[str, str], Mapping[str, str] | None]] = [] + monkeypatch.setattr( + "pythinker_code.auth.oauth_flows.new_client_session", + lambda: _FakeSession(queued, calls), + ) + return calls + + +def test_generate_pkce_and_state_have_oauth_safe_shape() -> None: + first = generate_pkce() + second = generate_pkce() + state = generate_state() + + assert re.fullmatch(r"[A-Za-z0-9_-]{43}", first.code_verifier) + assert re.fullmatch(r"[A-Za-z0-9_-]{43}", first.code_challenge) + assert re.fullmatch(r"[A-Za-z0-9_-]{43}", state) + verifier_bytes = first.code_verifier.encode(encoding="utf-8") + expected_challenge = ( + base64.urlsafe_b64encode(hashlib.sha256(verifier_bytes).digest()) + .decode(encoding="ascii", errors="replace") + .rstrip("=") + ) + assert first.code_challenge == expected_challenge + assert first != second + + +@pytest.mark.asyncio +async def test_request_device_code_parses_response(monkeypatch: pytest.MonkeyPatch) -> None: + calls = _mock_http( + monkeypatch, + ( + 200, + { + "device_code": "device-secret", + "user_code": "ABCD-EFGH", + "verification_uri": "https://login.example/device", + "verification_uri_complete": "https://login.example/device?user_code=ABCD-EFGH", + "expires_in": 600, + "interval": 3, + }, + ), + ) + + authorization = await request_device_code( + device_authorization_endpoint="https://login.example/oauth/device", + client_id="client-id", + scope=["openid", "profile"], + extra_params={"audience": "example-api"}, + headers={"Accept": "application/json"}, + ) + + assert authorization == DeviceCode( + user_code="ABCD-EFGH", + verification_uri="https://login.example/device", + verification_uri_complete="https://login.example/device?user_code=ABCD-EFGH", + device_code="device-secret", + interval=3, + expires_in=600, + ) + assert calls == [ + ( + "https://login.example/oauth/device", + { + "audience": "example-api", + "client_id": "client-id", + "scope": "openid profile", + }, + {"Accept": "application/json"}, + ) + ] + + +@pytest.mark.asyncio +async def test_poll_device_token_handles_pending_slow_down_then_success( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = _mock_http( + monkeypatch, + (400, {"error": "authorization_pending"}), + (400, {"error": "slow_down"}), + (200, {"access_token": "access-secret", "token_type": "Bearer"}), + ) + sleeps: list[float] = [] + + async def fake_sleep(delay: float) -> None: + sleeps.append(delay) + + monkeypatch.setattr("pythinker_code.auth.oauth_flows.asyncio.sleep", fake_sleep) + + payload = await poll_device_token( + token_endpoint="https://login.example/oauth/token", + client_id="client-id", + device_code=_device_code(interval=2), + headers={"Accept": "application/json"}, + ) + + assert payload == {"access_token": "access-secret", "token_type": "Bearer"} + assert sleeps == [2, 2, 7] + assert len(calls) == 3 + assert calls[0][1] == { + "client_id": "client-id", + "device_code": "device-secret", + "grant_type": "urn:ietf:params:oauth:grant-type:device_code", + } + assert all(call[2] == {"Accept": "application/json"} for call in calls) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("error", "expected_error"), + [ + ("expired_token", OAuthDeviceExpired), + ("access_denied", OAuthAccessDenied), + ], +) +async def test_poll_device_token_raises_typed_terminal_errors( + monkeypatch: pytest.MonkeyPatch, + error: str, + expected_error: type[Exception], +) -> None: + _mock_http(monkeypatch, (400, {"error": error})) + + async def fake_sleep(delay: float) -> None: + assert delay == 1 + + monkeypatch.setattr("pythinker_code.auth.oauth_flows.asyncio.sleep", fake_sleep) + + with pytest.raises(expected_error): + await poll_device_token( + token_endpoint="https://login.example/oauth/token", + client_id="client-id", + device_code=_device_code(interval=1), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("payload", [{}, {"token_type": "Bearer"}, {"access_token": ""}]) +async def test_poll_device_token_rejects_success_without_access_token( + monkeypatch: pytest.MonkeyPatch, + payload: dict[str, object], +) -> None: + # A 2xx response carrying no usable access token must fail closed, not be + # returned as success. + _mock_http(monkeypatch, (200, payload)) + + async def fake_sleep(delay: float) -> None: + assert delay == 1 + + monkeypatch.setattr("pythinker_code.auth.oauth_flows.asyncio.sleep", fake_sleep) + + with pytest.raises(OAuthError): + await poll_device_token( + token_endpoint="https://login.example/oauth/token", + client_id="client-id", + device_code=_device_code(interval=1), + ) + + +def _device_code(*, interval: int) -> DeviceCode: + return DeviceCode( + user_code="ABCD-EFGH", + verification_uri="https://login.example/device", + verification_uri_complete=None, + device_code="device-secret", + interval=interval, + expires_in=600, + ) + + +class _FakeSocket: + def __init__(self, port: int) -> None: + self.port = port + + def getsockname(self) -> tuple[str, int]: + return ("127.0.0.1", self.port) + + +class _FakeServer: + def __init__(self, port: int) -> None: + self.sockets = [_FakeSocket(port)] + self.closed = False + self.waited_closed = False + + def close(self) -> None: + self.closed = True + + async def wait_closed(self) -> None: + self.waited_closed = True + + +class _FakeWriter: + def __init__(self) -> None: + self.buffer = bytearray() + self.closed = False + + def write(self, data: bytes) -> None: + self.buffer.extend(data) + + async def drain(self) -> None: + return None + + def close(self) -> None: + self.closed = True + + async def wait_closed(self) -> None: + return None + + +def _reader(request_target: str) -> asyncio.StreamReader: + reader = asyncio.StreamReader() + request = f"GET {request_target} HTTP/1.1\r\n\r\n" + reader.feed_data(bytes(request, encoding="utf-8")) + reader.feed_eof() + return reader + + +def _request_reader(request: str, body: bytes = b"") -> asyncio.StreamReader: + reader = asyncio.StreamReader() + reader.feed_data(bytes(request, encoding="utf-8") + body) + reader.feed_eof() + return reader + + +def _implicit_post_reader(payload: dict[str, object]) -> asyncio.StreamReader: + body_text = json.dumps(payload) + body = bytes(body_text, encoding="utf-8") + request = ( + f"POST /oauth/token HTTP/1.1\r\nHost: localhost\r\ncOnTeNt-LeNgTh: {len(body)}\r\n\r\n" + ) + return _request_reader(request, body) + + +async def _drive_implicit_handler( + payload: dict[str, object], + *, + expected_state: str = "expected-state", +) -> tuple[asyncio.Future[ImplicitAuthorization], _FakeWriter]: + result: asyncio.Future[ImplicitAuthorization] = asyncio.get_running_loop().create_future() + writer = _FakeWriter() + await _handle_implicit_loopback_callback( + _implicit_post_reader(payload), + cast("asyncio.StreamWriter", writer), + callback_path="/oauth/callback", + token_path="/oauth/token", + expected_state=expected_state, + result=result, + ) + return result, writer + + +def _mock_loopback_server( + monkeypatch: pytest.MonkeyPatch, +) -> tuple[ + _FakeServer, + dict[str, Callable[[asyncio.StreamReader, asyncio.StreamWriter], None]], +]: + server = _FakeServer(port=43123) + captured: dict[str, Callable[[asyncio.StreamReader, asyncio.StreamWriter], None]] = {} + + async def fake_start_server( + handler: Callable[[asyncio.StreamReader, asyncio.StreamWriter], None], + host: str, + port: int, + ) -> _FakeServer: + assert host == "127.0.0.1" + assert port == 0 + captured["handler"] = handler + return server + + monkeypatch.setattr("pythinker_code.auth.oauth_flows.asyncio.start_server", fake_start_server) + return server, captured + + +@pytest.mark.asyncio +async def test_loopback_flow_rejects_state_mismatch(monkeypatch: pytest.MonkeyPatch) -> None: + server, captured = _mock_loopback_server(monkeypatch) + writer = _FakeWriter() + + def open_browser(url: str) -> None: + handler = captured["handler"] + handler( + _reader("/oauth/callback?code=authorization-code&state=wrong-state"), + cast("asyncio.StreamWriter", writer), + ) + + with pytest.raises(OAuthStateMismatch): + await run_loopback_pkce_flow( + authorize_endpoint="https://login.example/oauth/authorize", + client_id="client-id", + scope="openid profile", + redirect_path="/oauth/callback", + browser_open=open_browser, + ) + + assert bytes(writer.buffer).startswith(b"HTTP/1.1 400 Bad Request") + assert server.closed + assert server.waited_closed + + +@pytest.mark.asyncio +async def test_loopback_flow_captures_code_with_ephemeral_port( + monkeypatch: pytest.MonkeyPatch, +) -> None: + server, captured = _mock_loopback_server(monkeypatch) + opened_urls: list[str] = [] + writer = _FakeWriter() + + def open_browser(url: str) -> None: + opened_urls.append(url) + params = parse_qs(urlsplit(url).query) + handler = captured["handler"] + handler( + _reader(f"/oauth/callback?code=authorization-code&state={params['state'][0]}"), + cast("asyncio.StreamWriter", writer), + ) + + result = await run_loopback_pkce_flow( + authorize_endpoint="https://login.example/oauth/authorize", + client_id="client-id", + scope=["openid", "profile"], + redirect_path="/oauth/callback", + extra_authorize_params={"prompt": "login"}, + browser_open=open_browser, + ) + + assert result.authorization_code == "authorization-code" + assert re.fullmatch(r"[A-Za-z0-9_-]{43}", result.code_verifier) + assert result.redirect_uri == "http://127.0.0.1:43123/oauth/callback" + params = parse_qs(urlsplit(opened_urls[0]).query) + assert params["redirect_uri"] == [result.redirect_uri] + assert params["scope"] == ["openid profile"] + assert params["prompt"] == ["login"] + assert params["code_challenge_method"] == ["S256"] + assert bytes(writer.buffer).startswith(b"HTTP/1.1 200 OK") + assert server.closed + assert server.waited_closed + + +@pytest.mark.asyncio +async def test_implicit_callback_serves_bootstrap_without_resolving_result() -> None: + result: asyncio.Future[ImplicitAuthorization] = asyncio.get_running_loop().create_future() + writer = _FakeWriter() + + await _handle_implicit_loopback_callback( + _request_reader("GET /oauth/callback HTTP/1.1\r\n\r\n"), + cast("asyncio.StreamWriter", writer), + callback_path="/oauth/callback", + token_path="/oauth/token", + expected_state="expected-state", + result=result, + ) + + response = bytes(writer.buffer) + assert response.startswith(b"HTTP/1.1 200 OK") + assert b"Content-Type: text/html; charset=utf-8" in response + assert b'fetch("/oauth/token"' in response + assert b"window.location.hash" in response + assert not result.done() + assert writer.closed + + +@pytest.mark.asyncio +async def test_implicit_callback_resolves_access_token() -> None: + result, writer = await _drive_implicit_handler( + { + "access_token": "access-secret", + "expires_in": "3600", + "state": "expected-state", + } + ) + + assert await result == ImplicitAuthorization( + access_token="access-secret", + expires_in=3600, + state="expected-state", + ) + assert bytes(writer.buffer).startswith(b"HTTP/1.1 200 OK") + assert b"Content-Type: application/json" in writer.buffer + + +@pytest.mark.asyncio +async def test_implicit_callback_rejects_state_mismatch() -> None: + result, writer = await _drive_implicit_handler( + { + "access_token": "access-secret", + "expires_in": 3600, + "state": "wrong-state", + } + ) + + with pytest.raises(OAuthStateMismatch): + _ = await result + assert bytes(writer.buffer).startswith(b"HTTP/1.1 400 Bad Request") + + +@pytest.mark.asyncio +async def test_implicit_callback_maps_access_denied_error() -> None: + result, writer = await _drive_implicit_handler( + {"error": "access_denied", "state": "expected-state"} + ) + + with pytest.raises(OAuthAccessDenied): + _ = await result + assert bytes(writer.buffer).startswith(b"HTTP/1.1 400 Bad Request") + + +@pytest.mark.asyncio +async def test_implicit_callback_rejects_wrong_state_before_error() -> None: + result, writer = await _drive_implicit_handler( + {"error": "access_denied", "state": "wrong-state"} + ) + + with pytest.raises(OAuthStateMismatch): + _ = await result + assert bytes(writer.buffer).startswith(b"HTTP/1.1 400 Bad Request") + + +@pytest.mark.asyncio +async def test_implicit_callback_requires_access_token() -> None: + result, writer = await _drive_implicit_handler({"access_token": "", "state": "expected-state"}) + + with pytest.raises(OAuthError): + _ = await result + assert bytes(writer.buffer).startswith(b"HTTP/1.1 400 Bad Request") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("expires_in", [None, "not-a-number", 0, -5]) +async def test_implicit_callback_marks_invalid_expiry_unknown(expires_in: object) -> None: + payload: dict[str, object] = { + "access_token": "access-secret", + "state": "expected-state", + } + if expires_in is not None: + payload["expires_in"] = expires_in + + result, _writer = await _drive_implicit_handler(payload) + + # Malformed/missing/nonpositive expiry is reported as unknown (None), never + # fabricated into a trusted lifetime. + assert (await result).expires_in is None + + +@pytest.mark.asyncio +async def test_implicit_callback_rejects_truncated_body() -> None: + # Content-Length promises more bytes than the client actually sends, so + # readexactly() raises IncompleteReadError. The handler must fail closed + # with a 400 rather than leaving the result pending until timeout. + result: asyncio.Future[ImplicitAuthorization] = asyncio.get_running_loop().create_future() + writer = _FakeWriter() + request = "POST /oauth/token HTTP/1.1\r\nHost: localhost\r\nContent-Length: 4096\r\n\r\n" + reader = _request_reader(request, b'{"access_token": "x"}') + + await _handle_implicit_loopback_callback( + reader, + cast("asyncio.StreamWriter", writer), + callback_path="/oauth/callback", + token_path="/oauth/token", + expected_state="expected-state", + result=result, + ) + + with pytest.raises(OAuthError): + _ = await result + assert bytes(writer.buffer).startswith(b"HTTP/1.1 400 Bad Request") + + +@pytest.mark.asyncio +async def test_implicit_flow_uses_pinned_localhost_redirect( + monkeypatch: pytest.MonkeyPatch, +) -> None: + server = _FakeServer(port=9999) + captured: dict[str, Callable[[asyncio.StreamReader, asyncio.StreamWriter], None]] = {} + opened_urls: list[str] = [] + writer = _FakeWriter() + + async def fake_start_server( + handler: Callable[[asyncio.StreamReader, asyncio.StreamWriter], None], + host: str, + port: int, + ) -> _FakeServer: + assert host == "localhost" + assert port == 43124 + captured["handler"] = handler + return server + + def open_browser(url: str) -> None: + opened_urls.append(url) + params = parse_qs(urlsplit(url).query) + captured["handler"]( + _implicit_post_reader( + { + "access_token": "access-secret", + "expires_in": 7200, + "state": params["state"][0], + } + ), + cast("asyncio.StreamWriter", writer), + ) + + monkeypatch.setattr("pythinker_code.auth.oauth_flows.asyncio.start_server", fake_start_server) + + result = await run_loopback_implicit_flow( + authorize_endpoint="https://login.example/oauth/authorize?audience=example", + client_id="client-id", + scope=["openid", "profile"], + callback_path="/oauth/callback", + token_path="/oauth/token", + port=43124, + extra_authorize_params={"prompt": "login"}, + browser_open=open_browser, + ) + + assert result.access_token == "access-secret" + params = parse_qs(urlsplit(opened_urls[0]).query) + assert params["response_type"] == ["token"] + assert params["redirect_uri"] == ["http://localhost:43124/oauth/callback"] + assert params["scope"] == ["openid profile"] + assert params["prompt"] == ["login"] + assert server.closed + assert server.waited_closed + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("redirect_host", "port"), + [("0.0.0.0", 43124), ("example.com", 43124), ("localhost", 0)], +) +async def test_implicit_flow_rejects_non_loopback_or_zero_port( + redirect_host: str, port: int +) -> None: + # A non-loopback bind would expose the bearer-token callback; port 0 yields + # an unreachable ":0" redirect URI. Both must fail closed before binding. + with pytest.raises(ValueError): + await run_loopback_implicit_flow( + authorize_endpoint="https://login.example/oauth/authorize", + client_id="client-id", + scope="openid", + callback_path="/oauth/callback", + token_path="/oauth/token", + port=port, + redirect_host=redirect_host, + ) diff --git a/tests/auth/test_oauth_persist.py b/tests/auth/test_oauth_persist.py new file mode 100644 index 00000000..9ec1bfa3 --- /dev/null +++ b/tests/auth/test_oauth_persist.py @@ -0,0 +1,278 @@ +"""Atomicity guarantees for the shared OAuth persistence helpers. + +These pin the rollback contract of ``persist_login`` — in particular that a +re-login whose config save fails must not destroy a still-valid existing +credential (only a fresh login with no prior token deletes on failure). +""" + +from __future__ import annotations + +import asyncio +import threading +from types import SimpleNamespace + +import pytest + +from pythinker_code.auth import oauth +from pythinker_code.auth.oauth import OAuthToken +from pythinker_code.config import Config, OAuthRef + + +def _token(access: str, refresh: str) -> OAuthToken: + return OAuthToken( + access_token=access, + refresh_token=refresh, + expires_at=9_999_999_999.0, + scope="test-scope", + token_type="Bearer", + ) + + +@pytest.mark.asyncio +async def test_persist_login_restores_previous_token_on_config_save_failure( + monkeypatch: pytest.MonkeyPatch, tmp_path +) -> None: + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + ref = OAuthRef(storage="file", key="oauth/test-provider") + oauth.save_tokens(ref, _token("old-access", "old-refresh")) + + def boom(_config: Config) -> None: + raise OSError("disk full") + + monkeypatch.setattr(oauth, "save_config", boom) + + config = Config(is_from_default_location=True) + + def apply_config(cfg: Config) -> None: + cfg.default_model = "provider/new-model" + + with pytest.raises(OSError): + await oauth.persist_login(config, ref, _token("new-access", "new-refresh"), apply_config) + + # The previously valid credential must survive an unrelated config-save + # failure — the re-login attempt failed, but the working token is intact. + restored = oauth.load_tokens(ref) + assert restored is not None + assert restored.access_token == "old-access" + assert restored.refresh_token == "old-refresh" + # In-memory config is rolled back to before the mutation. + assert config.default_model == "" + + +@pytest.mark.asyncio +async def test_persist_login_deletes_token_on_fresh_login_failure( + monkeypatch: pytest.MonkeyPatch, tmp_path +) -> None: + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + ref = OAuthRef(storage="file", key="oauth/test-provider") + + def boom(_config: Config) -> None: + raise OSError("disk full") + + monkeypatch.setattr(oauth, "save_config", boom) + + config = Config(is_from_default_location=True) + + with pytest.raises(OSError): + await oauth.persist_login( + config, ref, _token("new-access", "new-refresh"), lambda cfg: None + ) + + # No prior credential existed, so the failed login leaves no orphan token. + assert oauth.load_tokens(ref) is None + + +@pytest.mark.asyncio +async def test_persist_login_persists_both_on_success( + monkeypatch: pytest.MonkeyPatch, tmp_path +) -> None: + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + ref = OAuthRef(storage="file", key="oauth/test-provider") + + saved: list[Config] = [] + monkeypatch.setattr(oauth, "save_config", lambda cfg: saved.append(cfg)) + + config = Config(is_from_default_location=True) + await oauth.persist_login( + config, + ref, + _token("new-access", "new-refresh"), + lambda cfg: setattr(cfg, "default_model", "provider/new-model"), + ) + + assert saved == [config] + assert config.default_model == "provider/new-model" + stored = oauth.load_tokens(ref) + assert stored is not None + assert stored.access_token == "new-access" + + +@pytest.mark.asyncio +async def test_persist_login_runs_complete_unit_off_event_loop( + monkeypatch: pytest.MonkeyPatch, tmp_path +) -> None: + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + ref = OAuthRef(storage="file", key="oauth/thread-probe") + event_loop_thread = threading.get_ident() + observed_threads: list[int] = [] + + def record_save(_config: Config) -> None: + observed_threads.append(threading.get_ident()) + + monkeypatch.setattr(oauth, "save_config", record_save) + + await oauth.persist_login( + Config(is_from_default_location=True), + ref, + _token("access", "refresh"), + lambda _cfg: observed_threads.append(threading.get_ident()), + ) + + assert observed_threads + assert all(thread_id != event_loop_thread for thread_id in observed_threads) + + +@pytest.mark.asyncio +async def test_failed_concurrent_login_cannot_clobber_successful_token( + monkeypatch: pytest.MonkeyPatch, tmp_path +) -> None: + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + ref = OAuthRef(storage="file", key="oauth/concurrent-provider") + oauth.save_tokens(ref, _token("old-access", "old-refresh")) + + failed_apply_entered = threading.Event() + allow_failed_save = threading.Event() + successful_apply_entered = threading.Event() + + def save_config(config: Config) -> None: + if config.default_model == "provider/failed": + raise OSError("disk full") + + monkeypatch.setattr(oauth, "save_config", save_config) + + def apply_failed(config: Config) -> None: + config.default_model = "provider/failed" + failed_apply_entered.set() + if not allow_failed_save.wait(timeout=5): + raise TimeoutError("failed-login test barrier timed out") + + def apply_successful(config: Config) -> None: + config.default_model = "provider/successful" + successful_apply_entered.set() + + failed_task = asyncio.create_task( + oauth.persist_login( + Config(is_from_default_location=True), + ref, + _token("failed-access", "failed-refresh"), + apply_failed, + ) + ) + assert await asyncio.to_thread(failed_apply_entered.wait, 5) + + successful_task = asyncio.create_task( + oauth.persist_login( + Config(is_from_default_location=True), + ref, + _token("successful-access", "successful-refresh"), + apply_successful, + ) + ) + _ = await asyncio.to_thread(successful_apply_entered.wait, 1) + allow_failed_save.set() + + results = await asyncio.gather(failed_task, successful_task, return_exceptions=True) + assert any(isinstance(result, OSError) for result in results) + assert any(result is None for result in results) + + stored = oauth.load_tokens(ref) + assert stored is not None + assert stored.access_token == "successful-access" + assert stored.refresh_token == "successful-refresh" + + +@pytest.mark.asyncio +async def test_persist_login_surfaces_credential_rollback_failure( + monkeypatch: pytest.MonkeyPatch, tmp_path +) -> None: + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + ref = OAuthRef(storage="file", key="oauth/rollback-provider") + oauth.save_tokens(ref, _token("old-access", "old-refresh")) + + original_save_tokens = oauth.save_tokens + save_calls = 0 + + def fail_rollback(target_ref: OAuthRef, token: OAuthToken) -> OAuthRef: + nonlocal save_calls + save_calls += 1 + if save_calls == 2: + raise OSError("rollback denied") + return original_save_tokens(target_ref, token) + + def fail_config_save(_config: Config) -> None: + raise OSError("disk full") + + errors: list[str] = [] + monkeypatch.setattr(oauth, "save_tokens", fail_rollback) + monkeypatch.setattr(oauth, "save_config", fail_config_save) + monkeypatch.setattr( + oauth, + "logger", + SimpleNamespace(error=lambda message, **_kwargs: errors.append(message)), + ) + + with pytest.raises( + oauth.OAuthPersistenceError, + match="credential rollback also failed", + ): + await oauth.persist_login( + Config(is_from_default_location=True), + ref, + _token("new-access", "new-refresh"), + lambda _config: None, + ) + + assert errors == ["Failed to roll back OAuth login persistence: {error}"] + + +@pytest.mark.asyncio +async def test_persist_logout_surfaces_delete_failure_and_restores_config( + monkeypatch: pytest.MonkeyPatch, tmp_path +) -> None: + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + ref = OAuthRef(storage="file", key="oauth/logout-provider") + oauth.save_tokens(ref, _token("access", "refresh")) + config = Config(is_from_default_location=True) + config.default_model = "provider/model" + saved_defaults: list[str] = [] + + monkeypatch.setattr(oauth, "save_config", lambda cfg: saved_defaults.append(cfg.default_model)) + + def fail_delete(_ref: OAuthRef) -> None: + raise OSError("delete denied") + + monkeypatch.setattr(oauth, "delete_tokens", fail_delete) + + with pytest.raises(oauth.OAuthPersistenceError, match="logout was rolled back"): + await oauth.persist_logout( + config, + ref, + lambda cfg: setattr(cfg, "default_model", ""), + ) + + assert config.default_model == "provider/model" + assert saved_defaults == ["", "provider/model"] + + +def test_nested_oauth_ref_has_distinct_credential_and_lock_paths( + monkeypatch: pytest.MonkeyPatch, tmp_path +) -> None: + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + + flat_key = "oauth/xai" + nested_key = "oauth/snowflake-cortex/xai" + + assert oauth._credentials_path(flat_key).name == "xai.json" + assert oauth._credentials_lock_path(flat_key).name == "xai.lock" + assert oauth._credentials_path(nested_key) != oauth._credentials_path(flat_key) + assert oauth._credentials_lock_path(nested_key) != oauth._credentials_lock_path(flat_key) diff --git a/tests/auth/test_opencode_go_auth.py b/tests/auth/test_opencode_go_auth.py index 0568c759..9cee8ea5 100644 --- a/tests/auth/test_opencode_go_auth.py +++ b/tests/auth/test_opencode_go_auth.py @@ -230,25 +230,70 @@ async def fake_discover(api_key): @pytest.mark.asyncio -async def test_fetch_models_dev_metadata_uses_short_best_effort_timeout(monkeypatch): - """The best-effort enrichment fetch must use a tight timeout so a slow - models.dev cannot block login for up to the 120s default.""" +async def test_fetch_models_dev_metadata_degrades_when_catalog_unavailable(monkeypatch): + """Best-effort enrichment degrades to empty metadata (the curated catalog) + when the shared models.dev catalog is unavailable, so login never blocks on + it. The fetch timeout and fail-open behavior now live in and are tested by + ``auth/models_dev.py``; this layer only delegates to it.""" from pythinker_code.auth import opencode_go + from pythinker_code.auth.models_dev import CatalogResult, CatalogStatus - captured: dict[str, aiohttp.ClientTimeout | None] = {} + async def empty_catalog(): + return CatalogResult({}, CatalogStatus.UNAVAILABLE, "none") - def fake_session(*, timeout=None): - captured["timeout"] = timeout - raise aiohttp.ClientError("unreachable") - - monkeypatch.setattr(opencode_go, "new_client_session", fake_session) + monkeypatch.setattr(opencode_go, "get_models_dev_catalog", empty_catalog) result = await opencode_go._fetch_models_dev_metadata() assert result == {} # degrades gracefully to the curated catalog - assert captured["timeout"] is opencode_go.MODELS_DEV_TIMEOUT - assert opencode_go.MODELS_DEV_TIMEOUT.total is not None - assert opencode_go.MODELS_DEV_TIMEOUT.total <= 15 + + +@pytest.mark.asyncio +async def test_fetch_models_dev_metadata_uses_authoritative_catalog(monkeypatch): + """A successful (authoritative) catalog load yields populated metadata rather + than being dropped: the normalized ModelsDevProvider objects flow straight + through to _ModelsDevMeta.""" + from pythinker_code.auth import opencode_go + from pythinker_code.auth.models_dev import ( + CatalogResult, + CatalogStatus, + parse_models_dev_catalog, + ) + + catalog = parse_models_dev_catalog( + { + opencode_go.MODELS_DEV_PROVIDER_ID: { + "name": "OpenCode Go", + "models": { + "kimi-k2.6": { + "id": "kimi-k2.6", + "name": "Kimi K2.6", + "limit": {"context": 262_144}, + "provider": {"npm": "@ai-sdk/moonshot"}, + }, + "claude-sonnet": { + "id": "claude-sonnet", + "name": "Claude Sonnet", + "limit": {"context": 200_000}, + "provider": {"npm": opencode_go.MODELS_DEV_ANTHROPIC_NPM}, + }, + }, + } + } + ) + + async def ok_catalog(): + return CatalogResult(catalog, CatalogStatus.OK, "models.dev") + + monkeypatch.setattr(opencode_go, "get_models_dev_catalog", ok_catalog) + + result = await opencode_go._fetch_models_dev_metadata() + + assert set(result) == {"kimi-k2.6", "claude-sonnet"} + assert result["kimi-k2.6"].display_name == "Kimi K2.6" + assert result["kimi-k2.6"].max_context == 262_144 + assert result["kimi-k2.6"].is_anthropic is False + assert result["claude-sonnet"].is_anthropic is True @pytest.mark.parametrize( diff --git a/tests/auth/test_platforms.py b/tests/auth/test_platforms.py index 0adb0f19..25084f79 100644 --- a/tests/auth/test_platforms.py +++ b/tests/auth/test_platforms.py @@ -42,6 +42,111 @@ def _make_config_with_model( ) +def test_digitalocean_platform_is_registered() -> None: + from pythinker_code.auth import DIGITALOCEAN_PLATFORM_ID + from pythinker_code.auth.platforms import get_platform_by_id, managed_provider_key + + platform = get_platform_by_id(DIGITALOCEAN_PLATFORM_ID) + + assert platform is not None + assert platform.name == "DigitalOcean" + assert platform.base_url == "https://inference.do-ai.run/v1" + assert managed_provider_key(platform.id) == "managed:digitalocean" + + +def test_snowflake_platform_is_registered() -> None: + from pythinker_code.auth import SNOWFLAKE_CORTEX_PLATFORM_ID + from pythinker_code.auth.platforms import get_platform_by_id, managed_provider_key + + platform = get_platform_by_id(SNOWFLAKE_CORTEX_PLATFORM_ID) + + assert platform is not None + assert platform.name == "Snowflake Cortex" + assert platform.base_url == "https://app.snowflake.com" + assert managed_provider_key(platform.id) == "managed:snowflake-cortex" + + +@pytest.mark.asyncio +async def test_refresh_managed_models_skips_snowflake_curated_catalog( + monkeypatch: pytest.MonkeyPatch, +) -> None: + provider_key = "managed:snowflake-cortex" + alias = "snowflake-cortex/claude-sonnet-4-5" + config = Config( + is_from_default_location=True, + default_model=alias, + providers={ + provider_key: LLMProvider( + type="openai_legacy", + base_url="https://myorg-acct.snowflakecomputing.com/api/v2/cortex/v1", + api_key=SecretStr(""), + oauth=OAuthRef(storage="file", key="oauth/snowflake-cortex/myorg-acct"), + ) + }, + models={ + alias: LLMModel( + provider=provider_key, + model="claude-sonnet-4-5", + max_context_size=200_000, + ) + }, + ) + called = False + + async def should_not_list_models(*args: Any, **kwargs: Any) -> list[ModelInfo]: + nonlocal called + called = True + return [] + + monkeypatch.setattr("pythinker_code.auth.platforms.list_models", should_not_list_models) + + changed = await refresh_managed_models(config) + + assert called is False + assert changed is False + assert alias in config.models + + +@pytest.mark.asyncio +async def test_refresh_managed_models_skips_digitalocean_router_catalog( + monkeypatch: pytest.MonkeyPatch, +) -> None: + provider_key = "managed:digitalocean" + alias = "digitalocean/router:production" + config = Config( + is_from_default_location=True, + default_model=alias, + providers={ + provider_key: LLMProvider( + type="openai_legacy", + base_url="https://inference.do-ai.run/v1", + api_key=SecretStr("token"), + ) + }, + models={ + alias: LLMModel( + provider=provider_key, + model="router:production", + max_context_size=128_000, + ) + }, + ) + called = False + + async def should_not_list_models(*args: Any, **kwargs: Any) -> list[ModelInfo]: + nonlocal called + called = True + return [] + + monkeypatch.setattr("pythinker_code.auth.platforms.list_models", should_not_list_models) + + changed = await refresh_managed_models(config) + + assert called is False + assert changed is False + assert alias in config.models + + # ── ModelInfo / _list_models: display_name parsing ───────────────── diff --git a/tests/auth/test_snowflake_auth.py b/tests/auth/test_snowflake_auth.py new file mode 100644 index 00000000..caa2a421 --- /dev/null +++ b/tests/auth/test_snowflake_auth.py @@ -0,0 +1,479 @@ +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest +from pydantic import SecretStr + +from pythinker_code.auth.models_dev import CatalogResult, CatalogStatus +from pythinker_code.auth.oauth import ( + OAuthError, + OAuthManager, + OAuthToken, + OAuthUnauthorized, + _credentials_path, + load_tokens, + save_tokens, +) +from pythinker_code.auth.oauth_flows import LoopbackAuthorization +from pythinker_code.config import Config, LLMModel, LLMProvider, OAuthRef + + +def _mock_unavailable_catalog(monkeypatch: pytest.MonkeyPatch) -> None: + from pythinker_code.auth import snowflake + + async def _fake() -> CatalogResult: + return CatalogResult({}, CatalogStatus.UNAVAILABLE, "none") + + monkeypatch.setattr(snowflake, "get_models_dev_catalog", _fake) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [CatalogStatus.OK, CatalogStatus.UNAVAILABLE]) +async def test_discover_snowflake_models_logs_catalog_fallback( + monkeypatch: pytest.MonkeyPatch, + status: CatalogStatus, +) -> None: + from pythinker_code.auth import snowflake + + messages: list[str] = [] + + async def empty_catalog() -> CatalogResult: + source = "network" if status is CatalogStatus.OK else "none" + return CatalogResult({}, status, source) + + monkeypatch.setattr(snowflake, "get_models_dev_catalog", empty_catalog) + monkeypatch.setattr( + snowflake, + "logger", + SimpleNamespace(debug=lambda message, **_kwargs: messages.append(message)), + ) + + assert await snowflake._discover_snowflake_models() == snowflake.SNOWFLAKE_MODELS + assert messages + if status is CatalogStatus.OK: + assert any("no usable" in message for message in messages) + else: + assert any("not authoritative" in message for message in messages) + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + ("https://MYORG-acct.snowflakecomputing.com/", "MYORG-acct"), + ("MYORG-acct", "MYORG-acct"), + (" http://MYORG-acct.snowflakecomputing.com ", "MYORG-acct"), + ("MYORG-acct///", "MYORG-acct"), + ], +) +def test_normalize_account(raw: str, expected: str) -> None: + from pythinker_code.auth.snowflake import normalize_account + + assert normalize_account(raw) == expected + + +@pytest.mark.parametrize( + "raw", + [ + "", + " ", + "acct/../evil", + "acct/path", + "user@acct", + "acct:443", + "acct?query=1", + "acct#fragment", + "acct evil", + "http://acct.snowflakecomputing.com:8443/", + "//evil.example.com", + ], +) +def test_normalize_account_rejects_hostile_input(raw: str) -> None: + from pythinker_code.auth.snowflake import normalize_account + + # Authority/path/port/userinfo/query/fragment payloads must be rejected + # before any URL is built, so they can never redirect the OAuth host. + with pytest.raises(ValueError): + normalize_account(raw) + + +@pytest.mark.parametrize( + ("role", "expected"), + [ + (None, "refresh_token"), + ("DATA_ENGINEER", "refresh_token session:role:DATA_ENGINEER"), + ("Data Science/Admin", "refresh_token session:role-encoded:Data%20Science%2FAdmin"), + ], +) +def test_scope(role: str | None, expected: str) -> None: + from pythinker_code.auth.snowflake import _scope + + assert _scope(role) == expected + + +class _FormResponse: + def __init__(self, status: int, payload: object) -> None: + self.status = status + self.payload = payload + + async def __aenter__(self) -> _FormResponse: + return self + + async def __aexit__(self, exc_type: object, exc: object, traceback: object) -> None: + return None + + async def json(self, *, content_type: object = None) -> object: + if isinstance(self.payload, ValueError): + raise self.payload + return self.payload + + +class _FormSession: + def __init__(self, response: _FormResponse, calls: list[dict[str, Any]]) -> None: + self.response = response + self.calls = calls + + async def __aenter__(self) -> _FormSession: + return self + + async def __aexit__(self, exc_type: object, exc: object, traceback: object) -> None: + return None + + def post( + self, endpoint: str, *, data: Mapping[str, str], headers: Mapping[str, str] + ) -> _FormResponse: + self.calls.append({"endpoint": endpoint, "data": dict(data), "headers": dict(headers)}) + return self.response + + +@pytest.mark.asyncio +async def test_refresh_uses_account_endpoint_basic_auth_and_default_expiry( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from pythinker_code.auth import snowflake + + calls: list[dict[str, Any]] = [] + response = _FormResponse(200, {"access_token": "access", "refresh_token": "refresh"}) + monkeypatch.setattr(snowflake, "new_client_session", lambda: _FormSession(response, calls)) + + token = await snowflake.refresh_snowflake_cortex_token("myorg-acct", "old-refresh") + + assert token.expires_in == 600 + assert calls == [ + { + "endpoint": "https://myorg-acct.snowflakecomputing.com/oauth/token-request", + "data": { + "grant_type": "refresh_token", + "refresh_token": "old-refresh", + "client_id": "LOCAL_APPLICATION", + }, + "headers": { + "Content-Type": "application/x-www-form-urlencoded", + "Accept": "application/json", + "Authorization": "Basic TE9DQUxfQVBQTElDQVRJT046TE9DQUxfQVBQTElDQVRJT04=", + }, + } + ] + + +@pytest.mark.asyncio +async def test_refresh_maps_invalid_grant_to_unauthorized( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from pythinker_code.auth import snowflake + + response = _FormResponse(400, {"error": "invalid_grant"}) + monkeypatch.setattr(snowflake, "new_client_session", lambda: _FormSession(response, [])) + + with pytest.raises(OAuthUnauthorized): + await snowflake.refresh_snowflake_cortex_token("myorg-acct", "revoked") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "response", + [ + _FormResponse(401, {"error": "unauthorized"}), + _FormResponse(500, {"error": "server_error"}), + _FormResponse(200, ValueError("bad json")), + _FormResponse(200, ["not", "a", "dict"]), + ], +) +async def test_refresh_rejects_error_and_malformed_responses( + monkeypatch: pytest.MonkeyPatch, + response: _FormResponse, +) -> None: + from pythinker_code.auth import snowflake + + monkeypatch.setattr(snowflake, "new_client_session", lambda: _FormSession(response, [])) + + with pytest.raises(OAuthError): + await snowflake.refresh_snowflake_cortex_token("myorg-acct", "old-refresh") + + +@pytest.mark.asyncio +async def test_login_snowflake_saves_account_scoped_token_provider_and_models( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from pythinker_code.auth import snowflake + + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + config = Config(is_from_default_location=True) + loopback_calls: list[dict[str, Any]] = [] + + async def fake_loopback(**kwargs: Any) -> LoopbackAuthorization: + loopback_calls.append(kwargs) + return LoopbackAuthorization("auth-code", "verifier", "http://127.0.0.1:49231/") + + async def fake_exchange( + account: str, + code: str, + code_verifier: str, + redirect_uri: str, + ) -> dict[str, Any]: + assert account == "MYORG-acct" + assert code == "auth-code" + assert code_verifier == "verifier" + assert redirect_uri == "http://127.0.0.1:49231/" + return snowflake._inject_default_expiry( + { + "access_token": "snowflake-access-secret", + "refresh_token": "snowflake-refresh-secret", + "token_type": "Bearer", + } + ) + + monkeypatch.setattr(snowflake, "run_loopback_pkce_flow", fake_loopback) + monkeypatch.setattr(snowflake, "_exchange_code_for_tokens", fake_exchange) + _mock_unavailable_catalog(monkeypatch) + + events = [ + event + async for event in snowflake.login_snowflake( + config, + " https://MYORG-acct.snowflakecomputing.com/ ", + "DATA_ENGINEER", + ) + ] + + assert [event.type for event in events] == ["waiting", "success"] + call = loopback_calls[0] + assert call["authorize_endpoint"] == ( + "https://MYORG-acct.snowflakecomputing.com/oauth/authorize" + ) + assert call["client_id"] == "LOCAL_APPLICATION" + assert call["scope"] == "refresh_token session:role:DATA_ENGINEER" + assert call["redirect_path"] == "/" + assert call["port"] == 0 + assert call["extra_authorize_params"] is None + + oauth_ref = OAuthRef(storage="file", key="oauth/snowflake-cortex/MYORG-acct") + stored = load_tokens(oauth_ref) + assert stored is not None + assert stored.access_token == "snowflake-access-secret" + assert stored.refresh_token == "snowflake-refresh-secret" + assert stored.expires_in == 600 + assert _credentials_path(oauth_ref.key).is_file() + + provider = config.providers["managed:snowflake-cortex"] + assert provider.type == "openai_legacy" + assert provider.base_url == "https://MYORG-acct.snowflakecomputing.com/api/v2/cortex/v1" + assert provider.api_key.get_secret_value() == "" + assert provider.oauth == oauth_ref + assert {model.provider for model in config.models.values()} == {"managed:snowflake-cortex"} + # Snowflake registers its models but must NOT become the default: its Cortex + # chat adapter is not implemented, so it cannot serve chat yet. + assert config.default_model == "" + assert not config.default_model.startswith("snowflake-cortex/") + rendered_events = "\n".join(event.json for event in events) + assert "snowflake-access-secret" not in rendered_events + assert "snowflake-refresh-secret" not in rendered_events + + +@pytest.mark.asyncio +async def test_snowflake_persistence_errors_hide_internal_details( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from pythinker_code.auth import snowflake + + diagnostic = "permission denied at /private/credentials/snowflake.json" + config = Config(is_from_default_location=True) + + async def fake_loopback(**_kwargs: Any) -> LoopbackAuthorization: + return LoopbackAuthorization("auth-code", "verifier", "http://127.0.0.1:49231/") + + async def fake_exchange( + _account: str, + _code: str, + _code_verifier: str, + _redirect_uri: str, + ) -> dict[str, Any]: + return { + "access_token": "snowflake-access", + "refresh_token": "snowflake-refresh", + "expires_in": 600, + } + + async def fake_models() -> tuple[snowflake.SnowflakeModel, ...]: + return snowflake.SNOWFLAKE_MODELS + + async def fail_persistence(*_args: object, **_kwargs: object) -> None: + raise OSError(diagnostic) + + monkeypatch.setattr(snowflake, "run_loopback_pkce_flow", fake_loopback) + monkeypatch.setattr(snowflake, "_exchange_code_for_tokens", fake_exchange) + monkeypatch.setattr(snowflake, "_discover_snowflake_models", fake_models) + monkeypatch.setattr(snowflake, "persist_login", fail_persistence) + monkeypatch.setattr(snowflake, "persist_config_change", fail_persistence) + + login_events = [event async for event in snowflake.login_snowflake(config, "myorg-acct")] + logout_events = [event async for event in snowflake.logout_snowflake(config)] + + assert login_events[-1].message == "Failed to save Snowflake Cortex login." + assert logout_events[-1].message == "Failed to log out of Snowflake Cortex." + assert diagnostic not in login_events[-1].json + assert diagnostic not in logout_events[-1].json + + +@pytest.mark.asyncio +async def test_login_snowflake_fails_when_refresh_token_is_missing( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from pythinker_code.auth import snowflake + + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + config = Config(is_from_default_location=True) + + async def fake_loopback(**_kwargs: Any) -> LoopbackAuthorization: + return LoopbackAuthorization("auth-code", "verifier", "http://127.0.0.1:49231/") + + async def fake_exchange( + _account: str, + _code: str, + _code_verifier: str, + _redirect_uri: str, + ) -> dict[str, Any]: + return {"access_token": "access-only", "expires_in": 600} + + monkeypatch.setattr(snowflake, "run_loopback_pkce_flow", fake_loopback) + monkeypatch.setattr(snowflake, "_exchange_code_for_tokens", fake_exchange) + + events = [event async for event in snowflake.login_snowflake(config, "myorg-acct")] + + assert [event.type for event in events] == ["waiting", "error"] + assert "refresh token" in events[-1].message + assert "managed:snowflake-cortex" not in config.providers + assert load_tokens(OAuthRef(storage="file", key="oauth/snowflake-cortex/myorg-acct")) is None + + +@pytest.mark.asyncio +async def test_login_snowflake_rejects_empty_account( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from pythinker_code.auth.snowflake import login_snowflake + + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + config = Config(is_from_default_location=True) + + events = [event async for event in login_snowflake(config, " ")] + + assert [event.type for event in events] == ["error"] + assert events[0].message == "Snowflake account identifier is required." + assert "managed:snowflake-cortex" not in config.providers + assert not (tmp_path / "credentials").exists() + + +@pytest.mark.asyncio +async def test_logout_snowflake_removes_account_token_provider_models_and_repairs_default( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from pythinker_code.auth.snowflake import SNOWFLAKE_PROVIDER_KEY, logout_snowflake + + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + oauth_ref = OAuthRef(storage="file", key="oauth/snowflake-cortex/myorg-acct") + save_tokens( + oauth_ref, + OAuthToken.from_response( + { + "access_token": "snowflake-access", + "refresh_token": "snowflake-refresh", + "expires_in": 600, + } + ), + ) + config = Config( + is_from_default_location=True, + default_model="snowflake-cortex/claude-sonnet-4-5", + providers={ + SNOWFLAKE_PROVIDER_KEY: LLMProvider( + type="openai_legacy", + base_url="https://myorg-acct.snowflakecomputing.com/api/v2/cortex/v1", + api_key=SecretStr(""), + oauth=oauth_ref, + ), + "fallback": LLMProvider( + type="openai_legacy", + base_url="https://example.test/v1", + api_key=SecretStr("test"), + ), + }, + models={ + "snowflake-cortex/claude-sonnet-4-5": LLMModel( + provider=SNOWFLAKE_PROVIDER_KEY, + model="claude-sonnet-4-5", + max_context_size=200_000, + ), + "fallback/model": LLMModel( + provider="fallback", + model="model", + max_context_size=32_000, + ), + }, + ) + + events = [event async for event in logout_snowflake(config)] + + assert [event.type for event in events] == ["success"] + assert load_tokens(oauth_ref) is None + assert not _credentials_path(oauth_ref.key).exists() + assert SNOWFLAKE_PROVIDER_KEY not in config.providers + assert all(model.provider != SNOWFLAKE_PROVIDER_KEY for model in config.models.values()) + assert config.default_model == "fallback/model" + + +@pytest.mark.asyncio +async def test_oauth_manager_routes_snowflake_refresh_with_account( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from pythinker_code.auth import snowflake + + calls: list[tuple[str, str]] = [] + + async def fake_refresh(account: str, refresh_token: str) -> OAuthToken: + calls.append((account, refresh_token)) + return OAuthToken.from_response( + { + "access_token": "new-access", + "refresh_token": "rotated-refresh", + "expires_in": 600, + } + ) + + monkeypatch.setattr(snowflake, "refresh_snowflake_cortex_token", fake_refresh) + manager = OAuthManager(Config()) + + token = await manager._refresh_token_for_ref( + OAuthRef(storage="file", key="oauth/snowflake-cortex/myorg-acct"), + "old-refresh", + ) + + assert calls == [("myorg-acct", "old-refresh")] + assert token.access_token == "new-access" + assert token.refresh_token == "rotated-refresh" diff --git a/tests/auth/test_xai_auth.py b/tests/auth/test_xai_auth.py new file mode 100644 index 00000000..61d5cb2d --- /dev/null +++ b/tests/auth/test_xai_auth.py @@ -0,0 +1,531 @@ +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest +from pydantic import SecretStr + +from pythinker_code.auth.models_dev import CatalogResult, CatalogStatus +from pythinker_code.auth.oauth import ( + OAuthError, + OAuthManager, + OAuthToken, + OAuthUnauthorized, + load_tokens, + save_tokens, +) +from pythinker_code.auth.oauth_flows import DeviceCode, LoopbackAuthorization +from pythinker_code.config import Config, LLMModel, LLMProvider, OAuthRef + + +def _mock_unavailable_catalog(monkeypatch: pytest.MonkeyPatch) -> None: + """Force provider logins to fall back to their curated model list.""" + from pythinker_code.auth import xai + + async def _fake() -> CatalogResult: + return CatalogResult({}, CatalogStatus.UNAVAILABLE, "none") + + monkeypatch.setattr(xai, "get_models_dev_catalog", _fake) + + +@pytest.mark.asyncio +async def test_discover_xai_models_logs_empty_authoritative_fallback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from pythinker_code.auth import xai + + messages: list[str] = [] + + async def empty_catalog() -> CatalogResult: + return CatalogResult({}, CatalogStatus.OK, "network") + + monkeypatch.setattr(xai, "get_models_dev_catalog", empty_catalog) + monkeypatch.setattr( + xai, + "logger", + SimpleNamespace(debug=lambda message, **_kwargs: messages.append(message)), + ) + + assert await xai._discover_xai_models() == xai.XAI_MODELS + assert any("no usable" in message for message in messages) + + +@pytest.mark.asyncio +async def test_login_xai_browser_saves_token_provider_and_models( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from pythinker_code.auth import xai + + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + config = Config(is_from_default_location=True) + loopback_calls: list[dict[str, Any]] = [] + + async def fake_loopback(**kwargs: Any) -> LoopbackAuthorization: + loopback_calls.append(kwargs) + return LoopbackAuthorization("auth-code", "verifier", "http://127.0.0.1:56121/callback") + + async def fake_exchange(code: str, code_verifier: str, redirect_uri: str) -> dict[str, Any]: + assert code == "auth-code" + assert code_verifier == "verifier" + assert redirect_uri == "http://127.0.0.1:56121/callback" + return { + "access_token": "xai-access", + "refresh_token": "xai-refresh", + "expires_in": 3600, + "token_type": "Bearer", + } + + monkeypatch.setattr(xai, "run_loopback_pkce_flow", fake_loopback) + monkeypatch.setattr(xai, "_exchange_code_for_tokens", fake_exchange) + _mock_unavailable_catalog(monkeypatch) + + events = [event async for event in xai.login_xai_browser(config)] + + assert [event.type for event in events] == ["waiting", "success"] + call = loopback_calls[0] + assert call["port"] == 56121 + assert call["redirect_path"] == "/callback" + assert call["extra_authorize_params"]["plan"] == "generic" + assert call["extra_authorize_params"]["nonce"] + + oauth_ref = OAuthRef(storage="file", key="oauth/xai") + stored = load_tokens(oauth_ref) + assert stored is not None + assert stored.access_token == "xai-access" + assert stored.refresh_token == "xai-refresh" + + provider = config.providers["managed:xai"] + assert provider.type == "openai_legacy" + assert provider.base_url == "https://api.x.ai/v1" + assert provider.oauth == oauth_ref + assert not provider.custom_headers + assert {model.provider for model in config.models.values()} == {"managed:xai"} + assert config.default_model == "xai/grok-4" + + +@pytest.mark.asyncio +async def test_xai_persistence_errors_hide_internal_details( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from pythinker_code.auth import xai + + diagnostic = "permission denied at /private/credentials/xai.json" + config = Config(is_from_default_location=True) + + async def fake_loopback(**_kwargs: Any) -> LoopbackAuthorization: + return LoopbackAuthorization("auth-code", "verifier", "http://127.0.0.1:56121/callback") + + async def fake_exchange( + _code: str, + _code_verifier: str, + _redirect_uri: str, + ) -> dict[str, Any]: + return { + "access_token": "xai-access", + "refresh_token": "xai-refresh", + "expires_in": 3600, + } + + async def fake_request_device_code(**_kwargs: Any) -> DeviceCode: + return DeviceCode( + user_code="GROK-CODE", + verification_uri="https://auth.x.ai/activate", + device_code="device-secret", + interval=5, + expires_in=900, + ) + + async def fake_poll_device_token(**_kwargs: Any) -> dict[str, Any]: + return { + "access_token": "xai-access", + "refresh_token": "xai-refresh", + "expires_in": 3600, + } + + async def fake_models() -> tuple[xai.XAIModel, ...]: + return xai.XAI_MODELS + + async def fail_persistence(*_args: object, **_kwargs: object) -> None: + raise OSError(diagnostic) + + monkeypatch.setattr(xai, "run_loopback_pkce_flow", fake_loopback) + monkeypatch.setattr(xai, "_exchange_code_for_tokens", fake_exchange) + monkeypatch.setattr(xai, "request_device_code", fake_request_device_code) + monkeypatch.setattr(xai, "poll_device_token", fake_poll_device_token) + monkeypatch.setattr(xai, "_discover_xai_models", fake_models) + monkeypatch.setattr(xai, "persist_login", fail_persistence) + monkeypatch.setattr(xai, "persist_logout", fail_persistence) + + browser_events = [event async for event in xai.login_xai_browser(config)] + headless_events = [event async for event in xai.login_xai_headless(config)] + logout_events = [event async for event in xai.logout_xai(config)] + + assert browser_events[-1].message == "Failed to save xAI Grok login." + assert headless_events[-1].message == "Failed to save xAI Grok login." + assert logout_events[-1].message == "Failed to log out of xAI Grok." + assert diagnostic not in browser_events[-1].json + assert diagnostic not in headless_events[-1].json + assert diagnostic not in logout_events[-1].json + + +@pytest.mark.asyncio +async def test_login_xai_headless_uses_device_flow( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from pythinker_code.auth import xai + + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + config = Config(is_from_default_location=True) + device_calls: list[dict[str, Any]] = [] + poll_calls: list[dict[str, Any]] = [] + + async def fake_request_device_code(**kwargs: Any) -> DeviceCode: + device_calls.append(kwargs) + return DeviceCode( + user_code="GROK-CODE", + verification_uri="https://auth.x.ai/activate", + device_code="device-secret", + interval=5, + expires_in=900, + ) + + async def fake_poll_device_token(**kwargs: Any) -> dict[str, Any]: + poll_calls.append(kwargs) + return { + "access_token": "xai-access", + "refresh_token": "xai-refresh", + "expires_in": 3600, + } + + monkeypatch.setattr(xai, "request_device_code", fake_request_device_code) + monkeypatch.setattr(xai, "poll_device_token", fake_poll_device_token) + _mock_unavailable_catalog(monkeypatch) + + events = [event async for event in xai.login_xai_headless(config)] + + assert [event.type for event in events] == ["verification_url", "waiting", "success"] + assert device_calls[0]["device_authorization_endpoint"] == xai.XAI_DEVICE_CODE_URL + assert device_calls[0]["headers"] == {"Accept": "application/json"} + assert poll_calls[0]["token_endpoint"] == xai.XAI_TOKEN_URL + assert poll_calls[0]["headers"] == {"Accept": "application/json"} + stored = load_tokens(OAuthRef(storage="file", key="oauth/xai")) + assert stored is not None + assert stored.refresh_token == "xai-refresh" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("refresh_token", [None, ""]) +async def test_login_xai_browser_rejects_missing_refresh_token( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + refresh_token: object, +) -> None: + from pythinker_code.auth import xai + + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + config = Config(is_from_default_location=True) + + async def fake_loopback(**kwargs: Any) -> LoopbackAuthorization: + return LoopbackAuthorization("auth-code", "verifier", "http://127.0.0.1:56121/callback") + + async def fake_exchange(code: str, code_verifier: str, redirect_uri: str) -> dict[str, Any]: + payload: dict[str, Any] = {"access_token": "xai-access", "expires_in": 3600} + if refresh_token is not None: + payload["refresh_token"] = refresh_token + return payload + + monkeypatch.setattr(xai, "run_loopback_pkce_flow", fake_loopback) + monkeypatch.setattr(xai, "_exchange_code_for_tokens", fake_exchange) + _mock_unavailable_catalog(monkeypatch) + + events = [event async for event in xai.login_xai_browser(config)] + + # A login without a refresh token must fail closed: no credentials, no + # provider, no models persisted. + assert events[-1].type == "error" + assert load_tokens(OAuthRef(storage="file", key="oauth/xai")) is None + assert "managed:xai" not in config.providers + assert config.models == {} + + +@pytest.mark.asyncio +async def test_login_xai_headless_rejects_missing_refresh_token( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from pythinker_code.auth import xai + + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + config = Config(is_from_default_location=True) + + async def fake_request_device_code(**kwargs: Any) -> DeviceCode: + return DeviceCode( + user_code="GROK-CODE", + verification_uri="https://auth.x.ai/activate", + device_code="device-secret", + interval=5, + expires_in=900, + ) + + async def fake_poll_device_token(**kwargs: Any) -> dict[str, Any]: + return {"access_token": "xai-access", "expires_in": 3600} + + monkeypatch.setattr(xai, "request_device_code", fake_request_device_code) + monkeypatch.setattr(xai, "poll_device_token", fake_poll_device_token) + _mock_unavailable_catalog(monkeypatch) + + events = [event async for event in xai.login_xai_headless(config)] + + assert events[-1].type == "error" + assert load_tokens(OAuthRef(storage="file", key="oauth/xai")) is None + assert "managed:xai" not in config.providers + + +@pytest.mark.asyncio +async def test_login_xai_browser_registers_catalog_models( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from pythinker_code.auth import xai + from pythinker_code.auth.models_dev import parse_models_dev_catalog + + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + config = Config(is_from_default_location=True) + catalog = parse_models_dev_catalog( + { + "xai": { + "name": "xAI", + "models": { + "grok-9": { + "id": "grok-9", + "name": "Grok 9", + "limit": {"context": 512_000}, + "modalities": {"input": ["text"], "output": ["text"]}, + }, + "grok-imagine-video": { + "id": "grok-imagine-video", + "name": "Grok Imagine Video", + "modalities": {"input": ["text"], "output": ["video"]}, + }, + }, + } + } + ) + + async def fake_loopback(**kwargs: Any) -> LoopbackAuthorization: + return LoopbackAuthorization("auth-code", "verifier", "http://127.0.0.1:56121/callback") + + async def fake_exchange(code: str, code_verifier: str, redirect_uri: str) -> dict[str, Any]: + return {"access_token": "xai-access", "refresh_token": "xai-refresh", "expires_in": 3600} + + async def fake_catalog() -> CatalogResult: + return CatalogResult(catalog, CatalogStatus.OK, "network") + + monkeypatch.setattr(xai, "run_loopback_pkce_flow", fake_loopback) + monkeypatch.setattr(xai, "_exchange_code_for_tokens", fake_exchange) + monkeypatch.setattr(xai, "get_models_dev_catalog", fake_catalog) + + events = [event async for event in xai.login_xai_browser(config)] + + assert events[-1].type == "success" + # The image/video model is filtered out; only the chat model is registered. + assert set(config.models) == {"xai/grok-9"} + assert config.models["xai/grok-9"].max_context_size == 512_000 + + +class _TokenResponse: + def __init__(self, status: int, payload: object) -> None: + self.status = status + self.payload = payload + + async def __aenter__(self) -> _TokenResponse: + return self + + async def __aexit__(self, exc_type: object, exc: object, traceback: object) -> None: + return None + + async def json(self, *, content_type: object = None) -> object: + return self.payload + + +class _TokenSession: + def __init__( + self, + response: _TokenResponse, + calls: list[tuple[str, Mapping[str, str]]], + ) -> None: + self.response = response + self.calls = calls + + async def __aenter__(self) -> _TokenSession: + return self + + async def __aexit__(self, exc_type: object, exc: object, traceback: object) -> None: + return None + + def post(self, endpoint: str, *, data: Mapping[str, str]) -> _TokenResponse: + self.calls.append((endpoint, data)) + return self.response + + +@pytest.mark.asyncio +async def test_refresh_xai_token_returns_rotated_token(monkeypatch: pytest.MonkeyPatch) -> None: + from pythinker_code.auth import xai + + calls: list[tuple[str, Mapping[str, str]]] = [] + response = _TokenResponse( + 200, + { + "access_token": "new-access", + "refresh_token": "rotated-refresh", + "expires_in": 3600, + "token_type": "Bearer", + }, + ) + monkeypatch.setattr(xai, "new_client_session", lambda: _TokenSession(response, calls)) + + token = await xai.refresh_xai_token("old-refresh") + + assert token.access_token == "new-access" + assert token.refresh_token == "rotated-refresh" + assert calls == [ + ( + "https://auth.x.ai/oauth2/token", + { + "grant_type": "refresh_token", + "refresh_token": "old-refresh", + "client_id": "b1a00492-073a-47ea-816f-4c329264a828", + }, + ) + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [401, 403]) +async def test_refresh_xai_token_raises_unauthorized( + monkeypatch: pytest.MonkeyPatch, + status: int, +) -> None: + from pythinker_code.auth import xai + + response = _TokenResponse(status, {"error_description": "revoked"}) + monkeypatch.setattr(xai, "new_client_session", lambda: _TokenSession(response, [])) + + with pytest.raises(OAuthUnauthorized): + await xai.refresh_xai_token("revoked-refresh") + + +@pytest.mark.asyncio +async def test_refresh_xai_token_treats_invalid_grant_as_unauthorized( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from pythinker_code.auth import xai + + # A 400 invalid_grant means the refresh token was rejected and must be + # surfaced as unauthorized so the token is suppressed, not retried. + response = _TokenResponse(400, {"error": "invalid_grant"}) + monkeypatch.setattr(xai, "new_client_session", lambda: _TokenSession(response, [])) + + with pytest.raises(OAuthUnauthorized): + await xai.refresh_xai_token("revoked-refresh") + + +@pytest.mark.asyncio +async def test_refresh_xai_token_raises_error_on_malformed_response( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from pythinker_code.auth import xai + + response = _TokenResponse(200, ["not", "a", "dict"]) + monkeypatch.setattr(xai, "new_client_session", lambda: _TokenSession(response, [])) + + with pytest.raises(OAuthError): + await xai.refresh_xai_token("some-refresh") + + +@pytest.mark.asyncio +async def test_oauth_manager_routes_xai_refresh(monkeypatch: pytest.MonkeyPatch) -> None: + from pythinker_code.auth import xai + + async def fake_refresh(refresh_token: str) -> OAuthToken: + assert refresh_token == "old-refresh" + return OAuthToken.from_response( + { + "access_token": "new-access", + "refresh_token": "rotated-refresh", + "expires_in": 3600, + } + ) + + monkeypatch.setattr(xai, "refresh_xai_token", fake_refresh) + manager = OAuthManager(Config()) + + token = await manager._refresh_token_for_ref( + OAuthRef(storage="file", key="oauth/xai"), "old-refresh" + ) + + assert token.access_token == "new-access" + assert token.refresh_token == "rotated-refresh" + + +@pytest.mark.asyncio +async def test_logout_xai_removes_tokens_provider_models_and_repairs_default( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from pythinker_code.auth.xai import XAI_OAUTH_KEY, XAI_PROVIDER_KEY, logout_xai + + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + oauth_ref = OAuthRef(storage="file", key=XAI_OAUTH_KEY) + save_tokens( + oauth_ref, + OAuthToken.from_response( + { + "access_token": "xai-access", + "refresh_token": "xai-refresh", + "expires_in": 3600, + } + ), + ) + config = Config( + is_from_default_location=True, + default_model="xai/grok-4", + providers={ + XAI_PROVIDER_KEY: LLMProvider( + type="openai_legacy", + base_url="https://api.x.ai/v1", + api_key=SecretStr(""), + oauth=oauth_ref, + ), + "fallback": LLMProvider( + type="openai_legacy", + base_url="https://example.test/v1", + api_key=SecretStr("test"), + ), + }, + models={ + "xai/grok-4": LLMModel( + provider=XAI_PROVIDER_KEY, + model="grok-4", + max_context_size=256_000, + ), + "fallback/model": LLMModel( + provider="fallback", + model="model", + max_context_size=32_000, + ), + }, + ) + + events = [event async for event in logout_xai(config)] + + assert [event.type for event in events] == ["success"] + assert load_tokens(oauth_ref) is None + assert XAI_PROVIDER_KEY not in config.providers + assert all(model.provider != XAI_PROVIDER_KEY for model in config.models.values()) + assert config.default_model == "fallback/model" diff --git a/tests/cli/test_openai_login_cli.py b/tests/cli/test_openai_login_cli.py index 4fb4faef..3926e286 100644 --- a/tests/cli/test_openai_login_cli.py +++ b/tests/cli/test_openai_login_cli.py @@ -47,6 +47,58 @@ def test_cli_login_headless_routes_to_openai_headless(monkeypatch): assert headless.called +def test_cli_login_copilot_routes_to_copilot(monkeypatch): + login = Mock(side_effect=_success_event) + monkeypatch.setattr("pythinker_code.cli.login_copilot", login, raising=False) + monkeypatch.setattr( + "pythinker_code.cli.load_config", + lambda: Config(is_from_default_location=True), + raising=False, + ) + + result = runner.invoke(cli, ["login", "--copilot"]) + + assert result.exit_code == 0 + assert login.called + + +def test_cli_login_xai_routes_to_browser(monkeypatch): + login = Mock(side_effect=_success_event) + monkeypatch.setattr("pythinker_code.cli.login_xai_browser", login, raising=False) + monkeypatch.setattr( + "pythinker_code.cli.load_config", + lambda: Config(is_from_default_location=True), + raising=False, + ) + + result = runner.invoke(cli, ["login", "--xai"]) + + assert result.exit_code == 0 + assert login.called + + +def test_cli_login_xai_device_routes_to_headless(monkeypatch): + login = Mock(side_effect=_success_event) + monkeypatch.setattr("pythinker_code.cli.login_xai_headless", login, raising=False) + monkeypatch.setattr( + "pythinker_code.cli.load_config", + lambda: Config(is_from_default_location=True), + raising=False, + ) + + result = runner.invoke(cli, ["login", "--xai-device"]) + + assert result.exit_code == 0 + assert login.called + + +def test_cli_login_rejects_multiple_xai_modes(monkeypatch): + result = runner.invoke(cli, ["login", "--xai", "--xai-device"]) + + assert result.exit_code == 1 + assert "Choose only one" in result.output + + def test_cli_login_api_key_routes_to_openai_api_key(monkeypatch): api_key = Mock(side_effect=_success_event) monkeypatch.setattr("pythinker_code.cli.login_openai_api_key", api_key, raising=False) @@ -77,6 +129,43 @@ def test_cli_logout_routes_to_openai_logout(monkeypatch): assert logout.called +def test_cli_logout_copilot_routes_to_copilot(monkeypatch): + logout = Mock(side_effect=_success_event) + monkeypatch.setattr("pythinker_code.cli.logout_copilot", logout, raising=False) + monkeypatch.setattr( + "pythinker_code.cli.load_config", + lambda: Config(is_from_default_location=True), + raising=False, + ) + + result = runner.invoke(cli, ["logout", "--copilot"]) + + assert result.exit_code == 0 + assert logout.called + + +def test_cli_logout_xai_routes_to_xai(monkeypatch): + logout = Mock(side_effect=_success_event) + monkeypatch.setattr("pythinker_code.cli.logout_xai", logout, raising=False) + monkeypatch.setattr( + "pythinker_code.cli.load_config", + lambda: Config(is_from_default_location=True), + raising=False, + ) + + result = runner.invoke(cli, ["logout", "--xai"]) + + assert result.exit_code == 0 + assert logout.called + + +def test_cli_logout_rejects_xai_with_other_mode(monkeypatch): + result = runner.invoke(cli, ["logout", "--xai", "--copilot"]) + + assert result.exit_code == 1 + assert "Choose only one" in result.output + + def test_cli_login_opencode_go_routes_to_opencode_go(monkeypatch): login = Mock(side_effect=_success_event) monkeypatch.setattr("pythinker_code.cli.login_opencode_go_api_key", login, raising=False) diff --git a/tests/core/test_config_atomic_save.py b/tests/core/test_config_atomic_save.py new file mode 100644 index 00000000..60c3bba2 --- /dev/null +++ b/tests/core/test_config_atomic_save.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from pythinker_code import config as config_module +from pythinker_code.config import Config, save_config + + +def test_save_config_preserves_previous_file_when_replace_fails( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + config_path = tmp_path / "config.toml" + previous = 'default_model = "existing/model"\n' + config_path.write_text(previous, encoding="utf-8") + + def fail_replace(_source: object, _destination: object) -> None: + raise OSError("simulated replace failure") + + monkeypatch.setattr(config_module.os, "replace", fail_replace) + + with pytest.raises(OSError, match="simulated replace failure"): + save_config(Config(), config_path) + + assert config_path.read_text(encoding="utf-8") == previous + assert list(tmp_path.glob(".config.toml.*.tmp")) == [] + + +def test_save_config_preserves_symlink_and_uses_link_format(tmp_path: Path) -> None: + target = tmp_path / "config-target" + config_path = tmp_path / "config.json" + config_path.symlink_to(target) + + save_config(Config(), config_path) + + assert config_path.is_symlink() + assert json.loads(target.read_text(encoding="utf-8"))["default_model"] == "" diff --git a/tests/ui_and_conv/test_shell_slash_commands.py b/tests/ui_and_conv/test_shell_slash_commands.py index 5842a5c2..abe0dd23 100644 --- a/tests/ui_and_conv/test_shell_slash_commands.py +++ b/tests/ui_and_conv/test_shell_slash_commands.py @@ -2,7 +2,7 @@ from __future__ import annotations -from collections.abc import Awaitable +from collections.abc import AsyncIterator, Awaitable from contextlib import nullcontext from pathlib import Path from types import SimpleNamespace @@ -14,6 +14,7 @@ from pythinker_core.message import Message from pythinker_host.path import HostPath +from pythinker_code.auth.oauth import OAuthEvent from pythinker_code.cli import Reload from pythinker_code.config import Config, LLMModel, LLMProvider from pythinker_code.session import Session @@ -30,6 +31,10 @@ async def _invoke_slash_command(command: SlashCommand[ShellSlashCmdFunc], shell: await ret +async def _oauth_success_event(*args: Any, **kwargs: Any) -> AsyncIterator[OAuthEvent]: + yield OAuthEvent("success", "ok") + + # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- @@ -92,6 +97,195 @@ def test_shell_slash_aliases_are_registered() -> None: assert command.name == canonical +async def test_shell_login_copilot_routes_to_copilot(monkeypatch: pytest.MonkeyPatch) -> None: + from pythinker_code.ui.shell import oauth as shell_oauth + + login = Mock(side_effect=_oauth_success_event) + config = Config(is_from_default_location=True) + app = SimpleNamespace(soul=SimpleNamespace(runtime=SimpleNamespace(config=config))) + monkeypatch.setattr(shell_oauth, "ensure_pythinker_soul", lambda _app: _app.soul) + monkeypatch.setattr(shell_oauth, "login_copilot", login) + + async def no_sleep(_delay: float) -> None: + return None + + monkeypatch.setattr(shell_oauth.asyncio, "sleep", no_sleep) + monkeypatch.setattr(shell_oauth.console, "clear", Mock()) + + with pytest.raises(Reload): + await cast(Any, shell_oauth.login)(app, "copilot") + + assert login.called + + +async def test_shell_logout_copilot_routes_to_copilot(monkeypatch: pytest.MonkeyPatch) -> None: + from pythinker_code.ui.shell import oauth as shell_oauth + + logout = Mock(side_effect=_oauth_success_event) + config = Config(is_from_default_location=True) + app = SimpleNamespace(soul=SimpleNamespace(runtime=SimpleNamespace(config=config))) + monkeypatch.setattr(shell_oauth, "ensure_pythinker_soul", lambda _app: _app.soul) + monkeypatch.setattr(shell_oauth, "logout_copilot", logout) + + async def no_sleep(_delay: float) -> None: + return None + + monkeypatch.setattr(shell_oauth.asyncio, "sleep", no_sleep) + monkeypatch.setattr(shell_oauth.console, "clear", Mock()) + + with pytest.raises(Reload): + await cast(Any, shell_oauth.logout)(app, "copilot") + + assert logout.called + + +async def test_shell_login_digitalocean_routes_to_digitalocean( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from pythinker_code.ui.shell import oauth as shell_oauth + + login = Mock(side_effect=_oauth_success_event) + config = Config(is_from_default_location=True) + app = SimpleNamespace(soul=SimpleNamespace(runtime=SimpleNamespace(config=config))) + monkeypatch.setattr(shell_oauth, "ensure_pythinker_soul", lambda _app: _app.soul) + monkeypatch.setattr(shell_oauth, "login_digitalocean", login) + + async def no_sleep(_delay: float) -> None: + return None + + monkeypatch.setattr(shell_oauth.asyncio, "sleep", no_sleep) + monkeypatch.setattr(shell_oauth.console, "clear", Mock()) + + with pytest.raises(Reload): + await cast(Any, shell_oauth.login)(app, "digitalocean") + + assert login.called + + +async def test_shell_logout_digitalocean_routes_to_digitalocean( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from pythinker_code.ui.shell import oauth as shell_oauth + + logout = Mock(side_effect=_oauth_success_event) + config = Config(is_from_default_location=True) + app = SimpleNamespace(soul=SimpleNamespace(runtime=SimpleNamespace(config=config))) + monkeypatch.setattr(shell_oauth, "ensure_pythinker_soul", lambda _app: _app.soul) + monkeypatch.setattr(shell_oauth, "logout_digitalocean", logout) + + async def no_sleep(_delay: float) -> None: + return None + + monkeypatch.setattr(shell_oauth.asyncio, "sleep", no_sleep) + monkeypatch.setattr(shell_oauth.console, "clear", Mock()) + + with pytest.raises(Reload): + await cast(Any, shell_oauth.logout)(app, "digitalocean") + + assert logout.called + + +@pytest.mark.parametrize( + ("mode", "function_name"), + [("xai", "login_xai_browser"), ("xai-device", "login_xai_headless")], +) +async def test_shell_login_xai_routes_to_xai_flow( + monkeypatch: pytest.MonkeyPatch, + mode: str, + function_name: str, +) -> None: + from pythinker_code.ui.shell import oauth as shell_oauth + + login = Mock(side_effect=_oauth_success_event) + config = Config(is_from_default_location=True) + app = SimpleNamespace(soul=SimpleNamespace(runtime=SimpleNamespace(config=config))) + monkeypatch.setattr(shell_oauth, "ensure_pythinker_soul", lambda _app: _app.soul) + monkeypatch.setattr(shell_oauth, function_name, login) + + async def no_sleep(_delay: float) -> None: + return None + + monkeypatch.setattr(shell_oauth.asyncio, "sleep", no_sleep) + monkeypatch.setattr(shell_oauth.console, "clear", Mock()) + + with pytest.raises(Reload): + await cast(Any, shell_oauth.login)(app, mode) + + assert login.called + + +async def test_shell_logout_xai_routes_to_xai(monkeypatch: pytest.MonkeyPatch) -> None: + from pythinker_code.ui.shell import oauth as shell_oauth + + logout = Mock(side_effect=_oauth_success_event) + config = Config(is_from_default_location=True) + app = SimpleNamespace(soul=SimpleNamespace(runtime=SimpleNamespace(config=config))) + monkeypatch.setattr(shell_oauth, "ensure_pythinker_soul", lambda _app: _app.soul) + monkeypatch.setattr(shell_oauth, "logout_xai", logout) + + async def no_sleep(_delay: float) -> None: + return None + + monkeypatch.setattr(shell_oauth.asyncio, "sleep", no_sleep) + monkeypatch.setattr(shell_oauth.console, "clear", Mock()) + + with pytest.raises(Reload): + await cast(Any, shell_oauth.logout)(app, "xai") + + assert logout.called + + +async def test_shell_login_snowflake_prompts_and_routes_to_snowflake( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from pythinker_code.ui.shell import oauth as shell_oauth + + login = Mock(side_effect=_oauth_success_event) + config = Config(is_from_default_location=True) + app = SimpleNamespace(soul=SimpleNamespace(runtime=SimpleNamespace(config=config))) + prompt_values = iter(("myorg-acct", "DATA_ENGINEER")) + + async def fake_prompt_text(_label: str) -> str: + return next(prompt_values) + + async def no_sleep(_delay: float) -> None: + return None + + monkeypatch.setattr(shell_oauth, "ensure_pythinker_soul", lambda _app: _app.soul) + monkeypatch.setattr(shell_oauth, "login_snowflake", login) + monkeypatch.setattr(shell_oauth, "_prompt_text", fake_prompt_text) + monkeypatch.setattr(shell_oauth.asyncio, "sleep", no_sleep) + monkeypatch.setattr(shell_oauth.console, "clear", Mock()) + + with pytest.raises(Reload): + await cast(Any, shell_oauth.login)(app, "snowflake") + + login.assert_called_once_with(config, "myorg-acct", "DATA_ENGINEER") + + +async def test_shell_logout_snowflake_routes_to_snowflake( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from pythinker_code.ui.shell import oauth as shell_oauth + + logout = Mock(side_effect=_oauth_success_event) + config = Config(is_from_default_location=True) + app = SimpleNamespace(soul=SimpleNamespace(runtime=SimpleNamespace(config=config))) + + async def no_sleep(_delay: float) -> None: + return None + + monkeypatch.setattr(shell_oauth, "ensure_pythinker_soul", lambda _app: _app.soul) + monkeypatch.setattr(shell_oauth, "logout_snowflake", logout) + monkeypatch.setattr(shell_oauth.asyncio, "sleep", no_sleep) + monkeypatch.setattr(shell_oauth.console, "clear", Mock()) + + with pytest.raises(Reload): + await cast(Any, shell_oauth.logout)(app, "snowflake") + + logout.assert_called_once_with(config) + + async def test_model_switch_starts_fresh_session(monkeypatch: pytest.MonkeyPatch) -> None: """Changing models should reload into a new session so old context is not reused.""" from pythinker_code.soul.pythinkersoul import PythinkerSoul