code-puppy-core-plugins 0.0.2__py3-none-any.whl

code_puppy_core_plugins/frontend_emitter/session_context.py

@@ -0,0 +1,50 @@
+"""Session-id ContextVar used by the frontend_emitter plugin.
+
+This module defines a single ``contextvars.ContextVar`` that carries
+the *current agent invocation's* session identifier across the asyncio
+call graph without requiring every callable to take an explicit kwarg.
+
+The primary use case is letting embedders (e.g. a WebSocket backend
+hosting multiple concurrent agent sessions) tag every event emitted
+during an agent run with a session identifier, so subscribers can
+filter events by session without monkey-patching the emitter.
+
+Usage::
+
+    from code_puppy_core_plugins.frontend_emitter.session_context import (
+        current_emitter_session_id,
+    )
+
+    token = current_emitter_session_id.set("my-session-id")
+    try:
+        await some_agent_work()  # any emit_event() inside picks up the sid
+    finally:
+        current_emitter_session_id.reset(token)
+
+Because ``ContextVar`` values are copied into each new asyncio Task at
+creation time, sessions are naturally isolated across concurrent tasks
+without explicit propagation.
+
+Scope note
+----------
+This ContextVar is intentionally scoped to the ``frontend_emitter``
+plugin.  It is not a general-purpose ``code_puppy.context.session_id``
+because the only consumer is the emitter's session-channel fan-out
+logic.  Other subsystems that need session-level metadata should
+introduce their own ContextVar with a clearly-scoped name.
+"""
+
+from __future__ import annotations
+
+from contextvars import ContextVar
+from typing import Optional
+
+# session_id of the currently running agent invocation, if any. Set by embedders
+# before agent work; read by emit_event()/register_callbacks to tag events. ``None``
+# is the legitimate "no session context" value, not an error.
+current_emitter_session_id: ContextVar[Optional[str]] = ContextVar(
+    "code_puppy_frontend_emitter_session_id", default=None
+)
+
+
+__all__ = ["current_emitter_session_id"]

code_puppy_core_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_core_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_core_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_core_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_core_plugins/grok_oauth/utils.py

@@ -0,0 +1,195 @@
+"""Utility helpers for the Grok (x.ai) OAuth plugin."""
+
+from __future__ import annotations
+
+import base64
+import datetime
+import hashlib
+import json
+import logging
+import secrets
+import time
+from typing import Any, Dict, Optional, Tuple
+from urllib.parse import urlparse
+
+import requests
+
+from .config import GROK_OAUTH_CONFIG, get_token_storage_path
+
+logger = logging.getLogger(__name__)
+
+_discovery_cache: Optional[Dict[str, str]] = None
+
+
+def _urlsafe_b64encode(data: bytes) -> str:
+    return base64.urlsafe_b64encode(data).decode("utf-8").rstrip("=")
+
+
+def generate_pkce_pair() -> Tuple[str, str]:
+    """Return a (code_verifier, code_challenge) PKCE S256 pair."""
+    code_verifier = _urlsafe_b64encode(secrets.token_bytes(32))
+    digest = hashlib.sha256(code_verifier.encode("utf-8")).digest()
+    return code_verifier, _urlsafe_b64encode(digest)
+
+
+def validate_xai_endpoint(url: str) -> str:
+    """Ensure a discovered endpoint is HTTPS on x.ai (or a subdomain)."""
+    parsed = urlparse(url)
+    host = (parsed.hostname or "").lower()
+    if parsed.scheme != "https" or (host != "x.ai" and not host.endswith(".x.ai")):
+        raise ValueError(f"xAI OAuth discovery returned an unexpected endpoint: {url}")
+    return url
+
+
+def fetch_discovery() -> Dict[str, str]:
+    """Fetch (and cache) the xAI OIDC discovery document."""
+    global _discovery_cache
+    if _discovery_cache:
+        return _discovery_cache
+
+    response = requests.get(
+        GROK_OAUTH_CONFIG["discovery_url"],
+        headers={"Accept": "application/json"},
+        timeout=30,
+    )
+    response.raise_for_status()
+    data = response.json()
+
+    authorization_endpoint = data.get("authorization_endpoint")
+    token_endpoint = data.get("token_endpoint")
+    if not authorization_endpoint or not token_endpoint:
+        raise ValueError(
+            "xAI OAuth discovery response did not include authorization/token endpoints"
+        )
+
+    _discovery_cache = {
+        "authorization_endpoint": validate_xai_endpoint(authorization_endpoint),
+        "token_endpoint": validate_xai_endpoint(token_endpoint),
+    }
+    return _discovery_cache
+
+
+def _utc_now_iso() -> str:
+    return (
+        datetime.datetime.now(datetime.timezone.utc).isoformat().replace("+00:00", "Z")
+    )
+
+
+def tokens_from_payload(
+    payload: Dict[str, Any],
+    token_endpoint: str,
+    fallback_refresh: str = "",
+) -> Dict[str, Any]:
+    """Normalize an OAuth token response into our stored-token shape."""
+    access_token = payload.get("access_token")
+    if not access_token:
+        raise ValueError("xAI token response did not include an access token")
+
+    expires_in = payload.get("expires_in") or 3600
+    return {
+        "access_token": access_token,
+        "refresh_token": payload.get("refresh_token") or fallback_refresh,
+        "id_token": payload.get("id_token", ""),
+        "token_type": payload.get("token_type", "Bearer"),
+        "expires_at": time.time() + float(expires_in),
+        "token_endpoint": token_endpoint,
+        "last_refresh": _utc_now_iso(),
+    }
+
+
+def load_stored_tokens() -> Optional[Dict[str, Any]]:
+    try:
+        token_path = get_token_storage_path()
+        if token_path.exists():
+            with open(token_path, "r", encoding="utf-8") as handle:
+                return json.load(handle)
+    except Exception as exc:
+        logger.error("Failed to load Grok OAuth tokens: %s", exc)
+    return None
+
+
+def save_tokens(tokens: Dict[str, Any]) -> bool:
+    if tokens is None:
+        raise TypeError("tokens cannot be None")
+    try:
+        token_path = get_token_storage_path()
+        with open(token_path, "w", encoding="utf-8") as handle:
+            json.dump(tokens, handle, indent=2)
+        token_path.chmod(0o600)
+        return True
+    except Exception as exc:
+        logger.error("Failed to save Grok OAuth tokens: %s", exc)
+    return False
+
+
+def exchange_token(token_endpoint: str, body: Dict[str, str]) -> Dict[str, Any]:
+    """POST a form-encoded token request and return the JSON payload."""
+    response = requests.post(
+        token_endpoint,
+        data=body,
+        headers={
+            "Accept": "application/json",
+            "Content-Type": "application/x-www-form-urlencoded",
+        },
+        timeout=30,
+    )
+    if response.status_code != 200:
+        raise RuntimeError(
+            f"xAI token request failed: {response.status_code} {response.text}"
+        )
+    return response.json()
+
+
+def refresh_access_token() -> Optional[str]:
+    """Refresh the access token using the stored refresh token."""
+    tokens = load_stored_tokens()
+    if not tokens:
+        return None
+
+    refresh_token = tokens.get("refresh_token")
+    if not refresh_token:
+        logger.debug("No refresh_token available for Grok OAuth")
+        return None
+
+    try:
+        token_endpoint = tokens.get("token_endpoint") or ""
+        token_endpoint = (
+            validate_xai_endpoint(token_endpoint)
+            if token_endpoint
+            else fetch_discovery()["token_endpoint"]
+        )
+        payload = exchange_token(
+            token_endpoint,
+            {
+                "grant_type": "refresh_token",
+                "refresh_token": refresh_token,
+                "client_id": GROK_OAUTH_CONFIG["client_id"],
+            },
+        )
+        new_tokens = tokens_from_payload(payload, token_endpoint, refresh_token)
+        if save_tokens(new_tokens):
+            logger.info("Successfully refreshed Grok OAuth token")
+            return new_tokens["access_token"]
+    except Exception as exc:
+        logger.error("Grok OAuth token refresh failed: %s", exc)
+    return None
+
+
+def get_valid_access_token() -> Optional[str]:
+    """Return a valid access token, refreshing if it is (nearly) expired."""
+    tokens = load_stored_tokens()
+    if not tokens:
+        logger.debug("No stored Grok OAuth tokens found")
+        return None
+
+    access_token = tokens.get("access_token")
+    if not access_token:
+        return None
+
+    expires_at = tokens.get("expires_at")
+    skew = GROK_OAUTH_CONFIG["refresh_skew_seconds"]
+    if isinstance(expires_at, (int, float)) and time.time() > expires_at - skew:
+        logger.info("Grok OAuth token expired, attempting refresh")
+        return refresh_access_token()
+
+    return access_token