code-puppy 0.0.613__tar.gz → 0.0.614__tar.gz

{code_puppy-0.0.613 → code_puppy-0.0.614}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: code-puppy
-Version: 0.0.613
+Version: 0.0.614
 Summary: Code generation agent
 Project-URL: repository, https://github.com/mpfaffenberger/code_puppy
 Project-URL: HomePage, https://github.com/mpfaffenberger/code_puppy

code_puppy-0.0.614/code_puppy/plugins/grok_oauth/__init__.py

@@ -0,0 +1,6 @@
+"""Grok (x.ai) OAuth plugin for Code Puppy.
+
+Authenticates against xAI's OAuth service (the same flow used by the
+official Grok CLI and the pi-xai-oauth package) and registers Grok
+models that run on the resulting access token.
+"""

code_puppy-0.0.614/code_puppy/plugins/grok_oauth/config.py

@@ -0,0 +1,49 @@
+"""Configuration for the Grok (x.ai) OAuth plugin.
+
+Endpoints, client id, and scopes mirror the official Grok CLI OAuth flow
+(as implemented by https://pi.dev/packages/pi-xai-oauth).
+"""
+
+from pathlib import Path
+from typing import Any, Dict
+
+from code_puppy import config
+
+GROK_OAUTH_CONFIG: Dict[str, Any] = {
+    # OIDC issuer + discovery. Authorization/token endpoints are resolved
+    # dynamically from the discovery document and validated against *.x.ai.
+    "issuer": "https://auth.x.ai",
+    "discovery_url": "https://auth.x.ai/.well-known/openid-configuration",
+    # Public OAuth client used by the official Grok CLI.
+    "client_id": "b1a00492-073a-47ea-816f-4c329264a828",
+    "scope": "openid profile email offline_access grok-cli:access api:access",
+    # Localhost callback. xAI allows an ephemeral-port fallback when the
+    # preferred port is busy.
+    "redirect_host": "127.0.0.1",
+    "redirect_port": 56121,
+    "redirect_path": "/callback",
+    "callback_timeout": 180,
+    # Refresh this many seconds before the token actually expires.
+    "refresh_skew_seconds": 120,
+    # Grok models are served from the regular xAI API with the OAuth
+    # access token as the bearer key.
+    "api_base_url": "https://api.x.ai/v1",
+    # Registered model names are prefixed to avoid catalogue collisions.
+    "prefix": "xai-",
+    "default_model": "grok-4.5",
+}
+
+# Grok models made available after authentication.
+GROK_MODELS: Dict[str, Dict[str, Any]] = {
+    "grok-4.5": {
+        "context_length": 500_000,
+        "supported_settings": ["reasoning_effort"],
+    },
+}
+
+
+def get_token_storage_path() -> Path:
+    """Path for storing OAuth tokens (uses XDG_DATA_HOME)."""
+    data_dir = Path(config.DATA_DIR)
+    data_dir.mkdir(parents=True, exist_ok=True, mode=0o700)
+    return data_dir / "grok_oauth.json"

code_puppy-0.0.614/code_puppy/plugins/grok_oauth/oauth_flow.py

@@ -0,0 +1,213 @@
+"""Grok (x.ai) OAuth flow matching the official Grok CLI / pi-xai-oauth."""
+
+from __future__ import annotations
+
+import secrets
+import threading
+import time
+import urllib.parse
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from typing import Any
+
+from code_puppy.messaging import emit_error, emit_info, emit_success, emit_warning
+
+from ..oauth_puppy_html import oauth_failure_html, oauth_success_html
+from .config import GROK_OAUTH_CONFIG
+from .utils import (
+    exchange_token,
+    fetch_discovery,
+    generate_pkce_pair,
+    load_stored_tokens,
+    save_tokens,
+    tokens_from_payload,
+)
+
+_ALLOWED_CORS_ORIGINS = ("https://accounts.x.ai", "https://auth.x.ai")
+
+
+class _OAuthServer(HTTPServer):
+    """Localhost callback server holding the state for one OAuth attempt."""
+
+    def __init__(self) -> None:
+        host = GROK_OAUTH_CONFIG["redirect_host"]
+        preferred_port = GROK_OAUTH_CONFIG["redirect_port"]
+        try:
+            super().__init__((host, preferred_port), _CallbackHandler)
+        except OSError:
+            # Preferred port busy — xAI accepts ephemeral localhost ports.
+            super().__init__((host, 0), _CallbackHandler)
+
+        self.exit_code = 1
+        discovery = fetch_discovery()
+        self.authorization_endpoint = discovery["authorization_endpoint"]
+        self.token_endpoint = discovery["token_endpoint"]
+        self.state = secrets.token_hex(16)
+        self.nonce = secrets.token_hex(16)
+        self.code_verifier, self.code_challenge = generate_pkce_pair()
+        port = self.server_address[1]
+        self.redirect_uri = f"http://{host}:{port}{GROK_OAUTH_CONFIG['redirect_path']}"
+
+    def auth_url(self) -> str:
+        params = {
+            "response_type": "code",
+            "client_id": GROK_OAUTH_CONFIG["client_id"],
+            "redirect_uri": self.redirect_uri,
+            "scope": GROK_OAUTH_CONFIG["scope"],
+            "code_challenge": self.code_challenge,
+            "code_challenge_method": "S256",
+            "state": self.state,
+            "nonce": self.nonce,
+        }
+        return f"{self.authorization_endpoint}?{urllib.parse.urlencode(params)}"
+
+    def exchange_code(self, code: str) -> bool:
+        payload = exchange_token(
+            self.token_endpoint,
+            {
+                "grant_type": "authorization_code",
+                "code": code,
+                "redirect_uri": self.redirect_uri,
+                "client_id": GROK_OAUTH_CONFIG["client_id"],
+                "code_verifier": self.code_verifier,
+            },
+        )
+        return save_tokens(tokens_from_payload(payload, self.token_endpoint))
+
+
+class _CallbackHandler(BaseHTTPRequestHandler):
+    server: "_OAuthServer"
+
+    def do_OPTIONS(self) -> None:  # noqa: N802
+        self.send_response(204)
+        self._send_cors_headers()
+        self.end_headers()
+
+    def do_GET(self) -> None:  # noqa: N802
+        parsed = urllib.parse.urlparse(self.path)
+        if parsed.path != GROK_OAUTH_CONFIG["redirect_path"]:
+            self._send_failure(404, "Callback endpoint not found for the puppy parade.")
+            return
+
+        params = urllib.parse.parse_qs(parsed.query)
+        error = params.get("error", [None])[0]
+        if error:
+            description = params.get("error_description", [error])[0]
+            self._send_failure(400, f"xAI authorization failed: {description}")
+            self._shutdown_after_delay()
+            return
+
+        if params.get("state", [None])[0] != self.server.state:
+            self._send_failure(
+                400, "OAuth state mismatch — a squirrel may have tampered with it."
+            )
+            return
+
+        code = params.get("code", [None])[0]
+        if not code:
+            self._send_failure(400, "Missing auth code — the token treat rolled away.")
+            self._shutdown_after_delay()
+            return
+
+        try:
+            saved = self.server.exchange_code(code)
+        except Exception as exc:  # noqa: BLE001
+            self._send_failure(500, f"Token exchange failed: {exc}")
+            self._shutdown_after_delay()
+            return
+
+        if saved:
+            self.server.exit_code = 0
+            self._send_html(
+                oauth_success_html(
+                    "Grok",
+                    "You can now close this window and return to Code Puppy.",
+                )
+            )
+        else:
+            self._send_failure(
+                500, "Unable to persist auth file — a puppy probably chewed it."
+            )
+        self._shutdown_after_delay()
+
+    def log_message(self, fmt: str, *args: Any) -> None:  # noqa: A003
+        pass
+
+    def _send_cors_headers(self) -> None:
+        origin = self.headers.get("Origin")
+        if origin not in _ALLOWED_CORS_ORIGINS:
+            return
+        self.send_header("Access-Control-Allow-Origin", origin)
+        self.send_header("Access-Control-Allow-Methods", "GET, OPTIONS")
+        self.send_header("Access-Control-Allow-Headers", "Content-Type")
+        self.send_header("Access-Control-Allow-Private-Network", "true")
+        self.send_header("Vary", "Origin")
+
+    def _send_html(self, body: str, status: int = 200) -> None:
+        encoded = body.encode("utf-8")
+        self.send_response(status)
+        self._send_cors_headers()
+        self.send_header("Content-Type", "text/html; charset=utf-8")
+        self.send_header("Content-Length", str(len(encoded)))
+        self.end_headers()
+        self.wfile.write(encoded)
+
+    def _send_failure(self, status: int, reason: str) -> None:
+        self._send_html(oauth_failure_html("Grok", reason), status)
+
+    def _shutdown_after_delay(self, seconds: float = 2.0) -> None:
+        def _later() -> None:
+            time.sleep(seconds)
+            self.server.shutdown()
+
+        threading.Thread(target=_later, daemon=True).start()
+
+
+def run_oauth_flow() -> bool:
+    """Run the browser OAuth flow. Returns True when tokens were saved."""
+    existing_tokens = load_stored_tokens()
+    if existing_tokens and existing_tokens.get("access_token"):
+        emit_warning("Existing Grok tokens will be overwritten.")
+
+    try:
+        server = _OAuthServer()
+    except Exception as exc:  # noqa: BLE001
+        emit_error(f"Could not start Grok OAuth flow: {exc}")
+        return False
+
+    auth_url = server.auth_url()
+    emit_info(f"Open this URL in your browser: {auth_url}")
+
+    server_thread = threading.Thread(target=server.serve_forever, daemon=True)
+    server_thread.start()
+
+    from code_puppy.tools.common import should_suppress_browser
+
+    if should_suppress_browser():
+        emit_info(f"[HEADLESS MODE] Would normally open: {auth_url}")
+    else:
+        try:
+            import webbrowser
+
+            if not webbrowser.open(auth_url):
+                emit_warning("Please open the URL manually — the browser stayed home.")
+        except Exception as exc:  # noqa: BLE001
+            emit_warning(f"Could not open browser automatically: {exc}")
+
+    emit_info("Waiting for xAI authentication callback…")
+
+    elapsed = 0.0
+    timeout = GROK_OAUTH_CONFIG["callback_timeout"]
+    interval = 0.25
+    while elapsed < timeout and server.exit_code != 0:
+        time.sleep(interval)
+        elapsed += interval
+
+    server.shutdown()
+    server_thread.join(timeout=5)
+
+    if server.exit_code != 0:
+        emit_error("Grok authentication failed or timed out.")
+        return False
+
+    emit_success("Grok OAuth complete — access token saved.")
+    return True

code_puppy-0.0.614/code_puppy/plugins/grok_oauth/register_callbacks.py

@@ -0,0 +1,132 @@
+"""Grok (x.ai) OAuth plugin callbacks.
+
+Provides OAuth authentication for Grok models (SuperGrok / Grok CLI flow)
+and registers the 'grok_oauth' model type handler plus the Grok model
+catalogue (headlined by grok-4.5).
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Tuple
+
+from code_puppy.callbacks import register_callback
+from code_puppy.messaging import emit_info, emit_success, emit_warning
+from code_puppy.model_switching import set_model_and_reload_agent
+
+from .config import GROK_MODELS, GROK_OAUTH_CONFIG, get_token_storage_path
+from .oauth_flow import run_oauth_flow
+from .utils import get_valid_access_token, load_stored_tokens
+
+_PREFIX = GROK_OAUTH_CONFIG["prefix"]
+_DEFAULT_MODEL = f"{_PREFIX}{GROK_OAUTH_CONFIG['default_model']}"
+
+
+def _custom_help() -> List[Tuple[str, str]]:
+    return [
+        (
+            "grok-auth",
+            "Authenticate with x.ai (Grok) via OAuth and register Grok models",
+        ),
+        ("grok-status", "Check Grok OAuth authentication status"),
+        ("grok-logout", "Remove Grok OAuth tokens and registered models"),
+    ]
+
+
+def _handle_grok_status() -> None:
+    tokens = load_stored_tokens()
+    if tokens and tokens.get("access_token"):
+        emit_success(" Grok OAuth: Authenticated")
+        emit_info(f" Available Grok models: {', '.join(_grok_model_names())}")
+    else:
+        emit_warning(" Grok OAuth: Not authenticated")
+        emit_info(" Run /grok-auth to launch the browser sign-in flow.")
+
+
+def _handle_grok_logout() -> None:
+    token_path = get_token_storage_path()
+    if token_path.exists():
+        token_path.unlink()
+        emit_info("Removed Grok OAuth tokens")
+    emit_success("Grok logout complete")
+
+
+def _handle_custom_command(command: str, name: str) -> Optional[bool]:
+    if not name:
+        return None
+
+    if name == "grok-auth":
+        if run_oauth_flow():
+            emit_info(f"Registering Grok models: {', '.join(_grok_model_names())}")
+            set_model_and_reload_agent(_DEFAULT_MODEL)
+        return True
+
+    if name == "grok-status":
+        _handle_grok_status()
+        return True
+
+    if name == "grok-logout":
+        _handle_grok_logout()
+        return True
+
+    return None
+
+
+def _grok_model_names() -> List[str]:
+    return [f"{_PREFIX}{model_id}" for model_id in GROK_MODELS]
+
+
+def _load_grok_models() -> Dict[str, Any]:
+    """Inject Grok models into the catalogue once authenticated."""
+    tokens = load_stored_tokens()
+    if not (tokens and tokens.get("access_token")):
+        return {}
+
+    return {
+        f"{_PREFIX}{model_id}": {
+            "type": "grok_oauth",
+            "name": model_id,
+            "context_length": spec["context_length"],
+            "supported_settings": spec.get("supported_settings", []),
+            "oauth_source": "grok-oauth-plugin",
+        }
+        for model_id, spec in GROK_MODELS.items()
+    }
+
+
+def _create_grok_oauth_model(model_name: str, model_config: Dict, config: Dict) -> Any:
+    """Create a Grok OAuth model instance.
+
+    This handler is registered via the 'register_model_type' callback to
+    handle models with type='grok_oauth'. Grok models speak the OpenAI
+    Responses API on api.x.ai with the OAuth access token as bearer key.
+    """
+    from pydantic_ai.models.openai import OpenAIResponsesModel
+    from pydantic_ai.providers.openai import OpenAIProvider
+
+    from code_puppy.http_utils import create_async_client
+
+    access_token = get_valid_access_token()
+    if not access_token:
+        emit_warning(
+            f"Failed to get valid Grok OAuth token; skipping model "
+            f"'{model_config.get('name')}'. Run /grok-auth to authenticate."
+        )
+        return None
+
+    provider = OpenAIProvider(
+        api_key=access_token,
+        base_url=GROK_OAUTH_CONFIG["api_base_url"],
+        http_client=create_async_client(model_name=model_name),
+    )
+    return OpenAIResponsesModel(model_name=model_config["name"], provider=provider)
+
+
+def _register_model_types() -> List[Dict[str, Any]]:
+    """Register the grok_oauth model type handler."""
+    return [{"type": "grok_oauth", "handler": _create_grok_oauth_model}]
+
+
+register_callback("custom_command_help", _custom_help)
+register_callback("custom_command", _handle_custom_command)
+register_callback("register_model_type", _register_model_types)
+register_callback("load_models_config", _load_grok_models)

code_puppy-0.0.614/code_puppy/plugins/grok_oauth/test_plugin.py

@@ -0,0 +1,160 @@
+"""Tests for the Grok (x.ai) OAuth plugin."""
+
+from __future__ import annotations
+
+import time
+from unittest.mock import patch
+
+import pytest
+
+from code_puppy.plugins.grok_oauth import config as grok_config
+from code_puppy.plugins.grok_oauth import register_callbacks as rc
+from code_puppy.plugins.grok_oauth import utils
+from code_puppy.plugins.grok_oauth.config import GROK_OAUTH_CONFIG
+
+
+@pytest.fixture
+def token_path(tmp_path):
+    path = tmp_path / "grok_oauth.json"
+    with (
+        patch.object(grok_config, "get_token_storage_path", return_value=path),
+        patch.object(utils, "get_token_storage_path", return_value=path),
+        patch.object(rc, "get_token_storage_path", return_value=path),
+    ):
+        yield path
+
+
+def _save_valid_tokens(**overrides):
+    tokens = {
+        "access_token": "fake-token",
+        "refresh_token": "fake-refresh",
+        "expires_at": time.time() + 3600,
+        "token_endpoint": "https://auth.x.ai/oauth2/token",
+    }
+    tokens.update(overrides)
+    assert utils.save_tokens(tokens)
+    return tokens
+
+
+def test_config_values():
+    assert GROK_OAUTH_CONFIG["issuer"] == "https://auth.x.ai"
+    assert "grok-cli:access" in GROK_OAUTH_CONFIG["scope"]
+    assert "api:access" in GROK_OAUTH_CONFIG["scope"]
+    assert GROK_OAUTH_CONFIG["api_base_url"] == "https://api.x.ai/v1"
+    assert GROK_OAUTH_CONFIG["default_model"] in grok_config.GROK_MODELS
+
+
+def test_pkce_pair_shape():
+    verifier, challenge = utils.generate_pkce_pair()
+    assert len(verifier) >= 43
+    assert "=" not in verifier and "=" not in challenge
+    # Deterministic: same verifier always hashes to the same challenge
+    import base64
+    import hashlib
+
+    digest = hashlib.sha256(verifier.encode()).digest()
+    expected = base64.urlsafe_b64encode(digest).decode().rstrip("=")
+    assert challenge == expected
+
+
+@pytest.mark.parametrize(
+    "url",
+    [
+        "http://auth.x.ai/oauth2/token",  # not https
+        "https://evil.com/oauth2/token",  # wrong host
+        "https://notx.ai/oauth2/token",  # suffix trickery
+        "https://x.ai.evil.com/token",  # prefix trickery
+    ],
+)
+def test_validate_xai_endpoint_rejects(url):
+    with pytest.raises(ValueError):
+        utils.validate_xai_endpoint(url)
+
+
+def test_validate_xai_endpoint_accepts():
+    assert utils.validate_xai_endpoint("https://auth.x.ai/oauth2/token")
+    assert utils.validate_xai_endpoint("https://x.ai/oauth2/token")
+
+
+def test_save_and_load_tokens(token_path):
+    _save_valid_tokens()
+    loaded = utils.load_stored_tokens()
+    assert loaded["access_token"] == "fake-token"
+    assert oct(token_path.stat().st_mode & 0o777) == oct(0o600)
+
+
+def test_get_valid_access_token_fresh(token_path):
+    _save_valid_tokens()
+    assert utils.get_valid_access_token() == "fake-token"
+
+
+def test_get_valid_access_token_refreshes_expired(token_path):
+    _save_valid_tokens(access_token="stale", expires_at=time.time() - 10)
+    payload = {
+        "access_token": "fresh-token",
+        "refresh_token": "new-refresh",
+        "expires_in": 3600,
+    }
+    with patch.object(utils, "exchange_token", return_value=payload) as exchange:
+        assert utils.get_valid_access_token() == "fresh-token"
+
+    body = exchange.call_args[0][1]
+    assert body["grant_type"] == "refresh_token"
+    assert body["client_id"] == GROK_OAUTH_CONFIG["client_id"]
+    assert utils.load_stored_tokens()["refresh_token"] == "new-refresh"
+
+
+def test_get_valid_access_token_unauthenticated(token_path):
+    assert utils.get_valid_access_token() is None
+
+
+def test_load_grok_models_requires_auth(token_path):
+    assert rc._load_grok_models() == {}
+
+
+def test_load_grok_models_when_authenticated(token_path):
+    _save_valid_tokens()
+    models = rc._load_grok_models()
+    assert "xai-grok-4.5" in models
+    entry = models["xai-grok-4.5"]
+    assert entry["type"] == "grok_oauth"
+    assert entry["name"] == "grok-4.5"
+    assert entry["context_length"] == 500_000
+
+
+def test_model_type_registration():
+    handlers = rc._register_model_types()
+    assert handlers == [{"type": "grok_oauth", "handler": rc._create_grok_oauth_model}]
+
+
+def test_create_model_returns_none_without_auth(token_path):
+    model = rc._create_grok_oauth_model("xai-grok-4.5", {"name": "grok-4.5"}, {})
+    assert model is None
+
+
+def test_create_model_when_authenticated(token_path):
+    _save_valid_tokens()
+    model = rc._create_grok_oauth_model("xai-grok-4.5", {"name": "grok-4.5"}, {})
+    assert type(model).__name__ == "OpenAIResponsesModel"
+    assert model.model_name == "grok-4.5"
+
+
+def test_custom_command_ignores_unknown():
+    assert rc._handle_custom_command("/other", "not-ours") is None
+    assert rc._handle_custom_command("/other", "") is None
+
+
+def test_custom_command_handles_status(token_path):
+    assert rc._handle_custom_command("/grok-status", "grok-status") is True
+
+
+def test_custom_command_logout(token_path):
+    _save_valid_tokens()
+    assert rc._handle_custom_command("/grok-logout", "grok-logout") is True
+    assert not token_path.exists()
+
+
+def test_custom_help_entries():
+    entries = rc._custom_help()
+    names = [name for name, _ in entries]
+    assert names == ["grok-auth", "grok-status", "grok-logout"]