vtx-coding-agent 0.1.1__py3-none-any.whl

vtx/llm/oauth/copilot.py

@@ -0,0 +1,358 @@
+"""
+GitHub Copilot OAuth device flow.
+
+Implements the device code flow to authenticate with GitHub and
+exchange for a Copilot token that can be used with the Copilot API.
+"""
+
+import asyncio
+import json
+from base64 import b64decode
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import aiohttp
+
+from vtx import get_config_dir
+
+# GitHub OAuth client ID (same as VS Code Copilot extension)
+_CLIENT_ID = b64decode("SXYxLmI1MDdhMDhjODdlY2ZlOTg=").decode()
+
+# Required headers for Copilot API
+COPILOT_HEADERS = {
+    "User-Agent": "GitHubCopilotChat/0.35.0",
+    "Editor-Version": "vscode/1.107.0",
+    "Editor-Plugin-Version": "copilot-chat/0.35.0",
+    "Copilot-Integration-Id": "vscode-chat",
+}
+
+
+@dataclass
+class CopilotCredentials:
+    github_token: str  # Long-lived GitHub OAuth token (refresh token)
+    copilot_token: str  # Short-lived Copilot API token (access token)
+    expires_at: int  # Unix timestamp (milliseconds) when copilot_token expires
+    enterprise_domain: str | None = None  # For GitHub Enterprise
+
+
+@dataclass
+class DeviceCodeResponse:
+    device_code: str
+    user_code: str
+    verification_uri: str
+    interval: int
+    expires_in: int
+
+
+def get_copilot_auth_path() -> Path:
+    return get_config_dir() / "copilot_auth.json"
+
+
+def load_credentials() -> CopilotCredentials | None:
+    path = get_copilot_auth_path()
+    if not path.exists():
+        return None
+
+    try:
+        data = json.loads(path.read_text())
+        return CopilotCredentials(
+            github_token=data["github_token"],
+            copilot_token=data["copilot_token"],
+            expires_at=data["expires_at"],
+            enterprise_domain=data.get("enterprise_domain"),
+        )
+    except (json.JSONDecodeError, KeyError):
+        return None
+
+
+def save_credentials(creds: CopilotCredentials) -> None:
+    path = get_copilot_auth_path()
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    data = {
+        "github_token": creds.github_token,
+        "copilot_token": creds.copilot_token,
+        "expires_at": creds.expires_at,
+    }
+    if creds.enterprise_domain:
+        data["enterprise_domain"] = creds.enterprise_domain
+
+    path.write_text(json.dumps(data, indent=2))
+    path.chmod(0o600)
+
+
+def clear_credentials() -> None:
+    path = get_copilot_auth_path()
+    if path.exists():
+        path.unlink()
+
+
+def is_copilot_logged_in() -> bool:
+    return load_credentials() is not None
+
+
+def _get_urls(domain: str) -> dict[str, str]:
+    return {
+        "device_code": f"https://{domain}/login/device/code",
+        "access_token": f"https://{domain}/login/oauth/access_token",
+        "copilot_token": f"https://api.{domain}/copilot_internal/v2/token",
+    }
+
+
+def get_base_url_from_token(token: str, enterprise_domain: str | None = None) -> str:
+    """
+    Extract API base URL from Copilot token.
+
+    Token format: tid=...;exp=...;proxy-ep=proxy.individual.githubcopilot.com;...
+    Returns API URL like https://api.individual.githubcopilot.com
+    """
+    import re
+
+    match = re.search(r"proxy-ep=([^;]+)", token)
+    if match:
+        proxy_host = match.group(1)
+        api_host = proxy_host.replace("proxy.", "api.", 1)
+        return f"https://{api_host}"
+
+    # Fallback
+    if enterprise_domain:
+        return f"https://copilot-api.{enterprise_domain}"
+    return "https://api.individual.githubcopilot.com"
+
+
+async def start_device_flow(domain: str = "github.com") -> DeviceCodeResponse:
+    urls = _get_urls(domain)
+
+    async with (
+        aiohttp.ClientSession() as session,
+        session.post(
+            urls["device_code"],
+            headers={
+                "Accept": "application/json",
+                "Content-Type": "application/json",
+                "User-Agent": "GitHubCopilotChat/0.35.0",
+            },
+            json={"client_id": _CLIENT_ID, "scope": "read:user"},
+        ) as response,
+    ):
+        response.raise_for_status()
+        data = await response.json()
+
+    return DeviceCodeResponse(
+        device_code=data["device_code"],
+        user_code=data["user_code"],
+        verification_uri=data["verification_uri"],
+        interval=data["interval"],
+        expires_in=data["expires_in"],
+    )
+
+
+async def poll_for_github_token(
+    device_code: str,
+    interval: int,
+    expires_in: int,
+    domain: str = "github.com",
+    on_poll: Any | None = None,
+) -> str:
+    """
+    Poll GitHub for the access token after user authorizes.
+
+    Returns the GitHub OAuth access token.
+    Raises TimeoutError if the flow expires.
+    """
+    import time
+
+    urls = _get_urls(domain)
+    deadline = time.time() + expires_in
+    poll_interval = max(1, interval)
+
+    async with aiohttp.ClientSession() as session:
+        while time.time() < deadline:
+            if on_poll:
+                on_poll()
+
+            async with session.post(
+                urls["access_token"],
+                headers={
+                    "Accept": "application/json",
+                    "Content-Type": "application/json",
+                    "User-Agent": "GitHubCopilotChat/0.35.0",
+                },
+                json={
+                    "client_id": _CLIENT_ID,
+                    "device_code": device_code,
+                    "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
+                },
+            ) as response:
+                data = await response.json()
+
+            if "access_token" in data:
+                return data["access_token"]
+
+            error = data.get("error")
+            if error == "authorization_pending":
+                await asyncio.sleep(poll_interval)
+                continue
+            elif error == "slow_down":
+                poll_interval += 5
+                await asyncio.sleep(poll_interval)
+                continue
+            elif error == "expired_token":
+                raise TimeoutError("Device code expired")
+            else:
+                raise RuntimeError(f"OAuth error: {error}")
+
+    raise TimeoutError("Device code flow timed out")
+
+
+async def exchange_for_copilot_token(
+    github_token: str, domain: str = "github.com"
+) -> tuple[str, int]:
+    """
+    Exchange GitHub OAuth token for Copilot API token.
+
+    Returns (copilot_token, expires_at_ms).
+    """
+    urls = _get_urls(domain)
+
+    async with (
+        aiohttp.ClientSession() as session,
+        session.get(
+            urls["copilot_token"],
+            headers={
+                "Accept": "application/json",
+                "Authorization": f"Bearer {github_token}",
+                **COPILOT_HEADERS,
+            },
+        ) as response,
+    ):
+        if response.status == 401:
+            raise RuntimeError(
+                "GitHub Copilot subscription not found. "
+                "Make sure you have an active Copilot subscription."
+            )
+        response.raise_for_status()
+        data = await response.json()
+
+    token = data["token"]
+    # expires_at is in seconds, convert to milliseconds with 5min buffer
+    expires_at = data["expires_at"] * 1000 - 5 * 60 * 1000
+
+    return token, expires_at
+
+
+async def refresh_copilot_token(creds: CopilotCredentials) -> CopilotCredentials:
+    domain = creds.enterprise_domain or "github.com"
+    copilot_token, expires_at = await exchange_for_copilot_token(creds.github_token, domain)
+
+    new_creds = CopilotCredentials(
+        github_token=creds.github_token,
+        copilot_token=copilot_token,
+        expires_at=expires_at,
+        enterprise_domain=creds.enterprise_domain,
+    )
+    save_credentials(new_creds)
+    return new_creds
+
+
+async def get_valid_token() -> str | None:
+    """
+    Get a valid Copilot API token, refreshing if needed.
+
+    Returns None if not logged in.
+    """
+    import time
+
+    creds = load_credentials()
+    if not creds:
+        return None
+
+    # Check if token needs refresh (with 1 minute buffer)
+    if time.time() * 1000 >= creds.expires_at - 60_000:
+        try:
+            creds = await refresh_copilot_token(creds)
+        except Exception:
+            # Token refresh failed, need to re-login
+            return None
+
+    return creds.copilot_token
+
+
+async def _enable_copilot_model(
+    token: str, model_id: str, enterprise_domain: str | None = None
+) -> bool:
+    base_url = get_base_url_from_token(token, enterprise_domain)
+    url = f"{base_url}/models/{model_id}/policy"
+
+    try:
+        async with (
+            aiohttp.ClientSession() as session,
+            session.post(
+                url,
+                headers={
+                    "Content-Type": "application/json",
+                    "Authorization": f"Bearer {token}",
+                    **COPILOT_HEADERS,
+                    "openai-intent": "chat-policy",
+                    "x-interaction-type": "chat-policy",
+                },
+                json={"state": "enabled"},
+            ) as response,
+        ):
+            return response.status < 400
+    except Exception:
+        return False
+
+
+async def enable_all_copilot_models(token: str, enterprise_domain: str | None = None) -> None:
+    from ..models import get_all_models
+
+    copilot_models = [m for m in get_all_models() if m.provider == "github-copilot"]
+    tasks = [_enable_copilot_model(token, model.id, enterprise_domain) for model in copilot_models]
+    await asyncio.gather(*tasks, return_exceptions=True)
+
+
+async def login(
+    on_user_code: Any | None = None, enterprise_domain: str | None = None
+) -> CopilotCredentials:
+    """
+    Perform the full Copilot login flow.
+
+    Args:
+        on_user_code: Callback with (verification_uri, user_code) when user action needed
+        enterprise_domain: Optional GitHub Enterprise domain
+
+    Returns:
+        CopilotCredentials that are saved and ready to use
+    """
+    domain = enterprise_domain or "github.com"
+
+    # Start device flow
+    device = await start_device_flow(domain)
+
+    # Notify caller about user action needed
+    if on_user_code:
+        on_user_code(device.verification_uri, device.user_code)
+
+    # Poll for GitHub token
+    github_token = await poll_for_github_token(
+        device.device_code, device.interval, device.expires_in, domain
+    )
+
+    # Exchange for Copilot token
+    copilot_token, expires_at = await exchange_for_copilot_token(github_token, domain)
+
+    # Save and return credentials
+    creds = CopilotCredentials(
+        github_token=github_token,
+        copilot_token=copilot_token,
+        expires_at=expires_at,
+        enterprise_domain=enterprise_domain,
+    )
+    save_credentials(creds)
+
+    # Enable all Copilot models (some require policy acceptance)
+    await enable_all_copilot_models(copilot_token, enterprise_domain)
+
+    return creds

vtx/llm/oauth/dynamic.py

@@ -0,0 +1,236 @@
+"""
+API-key storage for dynamic OpenAI-compatible providers.
+
+The dynamic providers (``airouter``, ``opencode``, ``kilo``, ``tokenrouter``) do
+not need an OAuth flow — they just need an API key. Users can set one of three
+ways, in priority order:
+
+1. The provider's ``<NAME>_API_KEY`` environment variable (e.g. ``KILO_API_KEY``).
+2. The encrypted-on-disk key file at the configured location (mode 0600),
+   written by the in-app ``/login`` command.
+3. None — for providers that support a free tier (airouter, kilo), vtx will
+   fall back to a placeholder key.
+
+This module owns path #2: it reads/writes the key file and exposes a small
+helper, :func:`get_dynamic_api_key`, that already implements the env-var-first
+priority so the rest of vtx does not have to.
+
+The storage location and format can be configured via the vtx-api-key-storage skill.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import json
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+from vtx import get_config_dir
+from vtx.llm.dynamic_models import DYNAMIC_PROVIDERS
+from vtx.llm.provider_catalog import get as get_provider_info
+
+# Default configuration
+AUTH_FILENAME = "dynamic_auth.json"
+Vtx_STORAGE_DIR = Path.home() / "vtx"
+
+
+@dataclass
+class DynamicProviderStatus:
+    """Status of a dynamic provider's credentials."""
+
+    provider: str
+    env_var: str | None
+    has_env_key: bool
+    has_stored_key: bool
+    api_key_optional: bool
+
+    @property
+    def is_configured(self) -> bool:
+        """True if we have any way to authenticate (key or no-auth provider)."""
+        return self.has_env_key or self.has_stored_key or self.api_key_optional
+
+
+def get_dynamic_auth_path() -> Path:
+    """Get the path to the API key storage file.
+
+    This function checks for the new YAML format first, then falls back
+    to the JSON format for backward compatibility.
+    """
+    # Check for new YAML format
+    yaml_path = Vtx_STORAGE_DIR / "dynamic_auth.yml"
+    if yaml_path.exists():
+        return yaml_path
+
+    # Check for JSON format in both old and new locations
+    # First check XDG_CONFIG_HOME/vtx/
+    xdg_config_dir = os.environ.get("XDG_CONFIG_HOME")
+    if xdg_config_dir:
+        xdg_path = Path(xdg_config_dir) / "vtx" / AUTH_FILENAME
+        if xdg_path.exists():
+            return xdg_path
+
+    # Then check ~/.vtx/ for backward compatibility
+    return get_config_dir() / AUTH_FILENAME
+
+
+def _read_all() -> dict[str, str]:
+    path = get_dynamic_auth_path()
+    if not path.exists():
+        return {}
+
+    try:
+        content = path.read_text(encoding="utf-8")
+
+        # Determine format based on file extension
+        if path.suffix.lower() == ".yml" or path.suffix.lower() == ".yaml":
+            import yaml
+
+            data = yaml.safe_load(content) or {}
+        elif path.suffix.lower() == ".json":
+            data = json.loads(content)
+        else:
+            # Default to JSON for backward compatibility
+            data = json.loads(content)
+
+    except (OSError, json.JSONDecodeError, ImportError):
+        return {}
+    except Exception:
+        return {}
+
+    if not isinstance(data, dict):
+        return {}
+    # Only keep str→str entries; ignore anything weird.
+    return {k: v for k, v in data.items() if isinstance(k, str) and isinstance(v, str)}
+
+    if not isinstance(data, dict):
+        return {}
+    # Only keep str→str entries; ignore anything weird.
+    return {k: v for k, v in data.items() if isinstance(k, str) and isinstance(v, str)}
+
+
+def _write_all(keys: dict[str, str]) -> None:
+    path = get_dynamic_auth_path()
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    # Determine format based on file extension
+    if path.suffix.lower() == ".yml" or path.suffix.lower() == ".yaml":
+        tmp = path.with_suffix(".yml.tmp")
+        try:
+            import yaml
+
+            tmp.write_text(yaml.dump(keys, default_flow_style=False), encoding="utf-8")
+        except ImportError:
+            # Fallback to JSON if yaml not available
+            tmp = path.with_suffix(".json.tmp")
+            tmp.write_text(json.dumps(keys, indent=2), encoding="utf-8")
+    else:
+        # Default to JSON
+        tmp = path.with_suffix(".json.tmp")
+        tmp.write_text(json.dumps(keys, indent=2), encoding="utf-8")
+
+    with contextlib.suppress(OSError):
+        # Non-POSIX filesystems (e.g. Windows) don't support chmod; ignore.
+        os.chmod(tmp, 0o600)
+    tmp.replace(path)
+    with contextlib.suppress(OSError):
+        os.chmod(path, 0o600)
+
+
+def load_api_key(provider: str) -> str | None:
+    """Return the API key stored on disk for a provider, if any."""
+    return _read_all().get(provider)
+
+
+def save_api_key(provider: str, key: str) -> None:
+    """Persist an API key for a provider."""
+    key = key.strip()
+    if not key:
+        raise ValueError("API key must not be empty")
+    if provider not in DYNAMIC_PROVIDERS and get_provider_info(provider) is None:
+        raise ValueError(f"Unknown provider: {provider}")
+    keys = _read_all()
+    keys[provider] = key
+    _write_all(keys)
+
+
+def clear_api_key(provider: str) -> bool:
+    """Remove a stored API key. Returns True if one was removed."""
+    keys = _read_all()
+    if provider not in keys:
+        return False
+    del keys[provider]
+    _write_all(keys)
+    return True
+
+
+def has_api_key(provider: str) -> bool:
+    """True if a stored key exists for the provider."""
+    return provider in _read_all()
+
+
+def _env_var_for(provider: str) -> str | None:
+    config = DYNAMIC_PROVIDERS.get(provider)
+    if config is not None:
+        return config.env_var
+    p = get_provider_info(provider)
+    return p.api_key_env if p else None
+
+
+def get_dynamic_api_key(provider: str) -> str | None:
+    """Return the best available API key for a dynamic provider.
+
+    Priority: ``<NAME>_API_KEY`` env var > stored ``dynamic_auth.json`` entry.
+    Returns ``None`` if neither is set.
+    """
+    env_var = _env_var_for(provider)
+    if env_var:
+        env_value = os.environ.get(env_var)
+        if env_value and env_value.strip():
+            return env_value.strip()
+    return load_api_key(provider)
+
+
+def get_provider_status(provider: str) -> DynamicProviderStatus | None:
+    """Return credential status for a provider, or ``None`` if unknown.
+
+    Works for both the built-in ``DYNAMIC_PROVIDERS`` (airouter, opencode,
+    kilo, tokenrouter) and any provider defined in ``provider.yaml``.
+    """
+    config = DYNAMIC_PROVIDERS.get(provider)
+    if config is not None:
+        env_var = config.env_var
+        has_env = bool(env_var and os.environ.get(env_var, "").strip())
+        return DynamicProviderStatus(
+            provider=provider,
+            env_var=env_var,
+            has_env_key=has_env,
+            has_stored_key=has_api_key(provider),
+            api_key_optional=config.api_key_optional,
+        )
+
+    p = get_provider_info(provider)
+    if p is None or not p.base_url:
+        return None
+    env_var = p.api_key_env
+    has_env = bool(env_var and os.environ.get(env_var, "").strip())
+    return DynamicProviderStatus(
+        provider=provider,
+        env_var=env_var,
+        has_env_key=has_env,
+        has_stored_key=has_api_key(provider),
+        api_key_optional=p.api_key_optional,
+    )
+
+
+__all__ = [
+    "AUTH_FILENAME",
+    "DynamicProviderStatus",
+    "clear_api_key",
+    "get_dynamic_api_key",
+    "get_dynamic_auth_path",
+    "get_provider_status",
+    "has_api_key",
+    "load_api_key",
+    "save_api_key",
+]