runtime-sdk 0.4.49__py3-none-win_amd64.whl

runtime_sdk/providers.py

@@ -0,0 +1,422 @@
+from __future__ import annotations
+
+import contextlib
+import json
+import os
+import re
+import select
+import shutil
+import subprocess
+import sys
+import tempfile
+import time
+from pathlib import Path
+from typing import Any
+
+from runtime_sdk import output as cli_output, terminal
+from runtime_sdk.auth import authenticated_client
+from runtime_sdk.client import RuntimeAPIError
+from runtime_sdk.config import RuntimeConfig, RuntimeConfigError
+
+
+def _read_provider_auth_file(path: str | Path) -> str:
+    try:
+        return Path(path).read_text(encoding="utf-8")
+    except OSError as exc:
+        raise RuntimeConfigError(f"failed to read auth file: {exc}") from exc
+
+
+def _codex_auth_path(codex_home: Path) -> Path:
+    return codex_home / "auth.json"
+
+
+def _local_codex_auth_path() -> Path:
+    codex_home = os.environ.get("CODEX_HOME", "").strip()
+    return _codex_auth_path(
+        Path(codex_home).expanduser() if codex_home else Path.home() / ".codex"
+    )
+
+
+def _claude_auth_env(home: Path) -> dict[str, str]:
+    env = os.environ.copy()
+    env["HOME"] = str(home)
+    env["USERPROFILE"] = str(home)
+    env["XDG_CONFIG_HOME"] = str(home / ".config")
+    env.pop("ANTHROPIC_API_KEY", None)
+    env.pop("ANTHROPIC_AUTH_TOKEN", None)
+    env.pop("CLAUDE_CODE_OAUTH_TOKEN", None)
+    return env
+
+
+def _run_codex_login(codex_bin: str, codex_home: Path, *, device: bool) -> int:
+    cmd = [
+        codex_bin,
+        "login",
+        "-c",
+        'cli_auth_credentials_store="file"',
+    ]
+    if device:
+        cmd.append("--device-auth")
+    env = os.environ.copy()
+    env["CODEX_HOME"] = str(codex_home)
+    return subprocess.run(cmd, env=env).returncode
+
+
+def _run_claude_setup_token(
+    claude_bin: str, home: Path, *, timeout_seconds: int = 300
+) -> str:
+    if sys.platform == "win32":
+        return _run_claude_setup_token_windows(claude_bin, home)
+
+    import pty
+
+    sys.stderr.write(
+        "Creating Claude Runtime proxy token. Complete the browser flow if prompted.\n"
+    )
+    sys.stderr.flush()
+    master_fd, slave_fd = pty.openpty()
+    proc = subprocess.Popen(
+        [claude_bin, "setup-token"],
+        stdin=slave_fd,
+        stdout=slave_fd,
+        stderr=slave_fd,
+        env=_claude_auth_env(home),
+        close_fds=True,
+    )
+    os.close(slave_fd)
+    chunks: list[bytes] = []
+    displayed = 0
+    token_output_hidden = False
+    deadline = time.monotonic() + timeout_seconds
+    try:
+        while proc.poll() is None:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                proc.terminate()
+                try:
+                    proc.wait(timeout=5)
+                except subprocess.TimeoutExpired:
+                    proc.kill()
+                    proc.wait()
+                raise RuntimeConfigError("claude setup-token timed out")
+            readers = [master_fd]
+            if cli_output.interactive():
+                readers.append(sys.stdin.fileno())
+            readable, _, _ = select.select(readers, [], [], min(0.2, remaining))
+            for fd in readable:
+                if fd == master_fd:
+                    try:
+                        data = os.read(master_fd, 4096)
+                    except OSError:
+                        data = b""
+                    if data:
+                        chunks.append(data)
+                        displayed, token_output_hidden = _write_claude_setup_output(
+                            b"".join(chunks).decode("utf-8", errors="ignore"),
+                            displayed,
+                            token_output_hidden,
+                        )
+                else:
+                    data = os.read(fd, 1024)
+                    if data:
+                        filtered = _filter_claude_setup_input(data)
+                        if filtered:
+                            os.write(master_fd, filtered)
+        while True:
+            try:
+                data = os.read(master_fd, 4096)
+            except OSError:
+                break
+            if not data:
+                break
+            chunks.append(data)
+            displayed, token_output_hidden = _write_claude_setup_output(
+                b"".join(chunks).decode("utf-8", errors="ignore"),
+                displayed,
+                token_output_hidden,
+            )
+    finally:
+        with contextlib.suppress(OSError):
+            os.close(master_fd)
+    status = proc.wait()
+    if status != 0:
+        raise RuntimeConfigError(f"claude setup-token failed with exit code {status}")
+    token = _extract_claude_oauth_token(
+        b"".join(chunks).decode("utf-8", errors="ignore")
+    )
+    if not token:
+        raise RuntimeConfigError(
+            "claude setup-token completed but did not print an OAuth token"
+        )
+    return token
+
+
+def _run_claude_setup_token_windows(claude_bin: str, home: Path) -> str:
+    sys.stderr.write(
+        "Creating Claude Runtime proxy token. Complete the browser flow, then copy the token shown by Claude.\n"
+    )
+    sys.stderr.flush()
+    status = subprocess.run(
+        [claude_bin, "setup-token"], env=_claude_auth_env(home)
+    ).returncode
+    if status != 0:
+        raise RuntimeConfigError(f"claude setup-token failed with exit code {status}")
+    token = cli_output.prompt_password("Paste Claude OAuth token") or ""
+    if not re.fullmatch(r"sk-ant-[A-Za-z0-9_-]{20,}", token):
+        raise RuntimeConfigError("Claude did not provide a valid OAuth token")
+    return token
+
+
+def _write_claude_setup_output(
+    output: str, displayed: int, token_output_hidden: bool
+) -> tuple[int, bool]:
+    marker = re.search(r"Your OAuth token", output, flags=re.I)
+    if marker:
+        if marker.start() > displayed:
+            sys.stderr.write(
+                terminal.strip_controls(output[displayed : marker.start()])
+            )
+            sys.stderr.flush()
+        if not token_output_hidden:
+            sys.stderr.write("\nClaude OAuth token captured and hidden.\n")
+            sys.stderr.flush()
+        return len(output), True
+    if len(output) > displayed:
+        chunk = terminal.strip_controls(output[displayed:])
+        chunk = re.sub(r"sk-ant-[A-Za-z0-9_-]{12,}", "[hidden]", chunk)
+        sys.stderr.write(chunk)
+        sys.stderr.flush()
+    return len(output), token_output_hidden
+
+
+def _filter_claude_setup_input(data: bytes) -> bytes:
+    cleaned = terminal.strip_controls(data.decode("utf-8", errors="ignore"))
+    return cleaned.encode()
+
+
+def _extract_claude_oauth_token(output: str) -> str:
+    output = terminal.strip_controls(output)
+    section = output
+    marker = re.search(r"Your OAuth token.*?:", output, flags=re.I | re.S)
+    if marker:
+        section = output[marker.end() :]
+    section = re.split(r"Store this token", section, maxsplit=1, flags=re.I)[0]
+    compact = re.sub(r"\s+", "", section)
+    match = re.search(r"sk-ant-[A-Za-z0-9_-]{20,}", compact)
+    return match.group(0) if match else ""
+
+
+def _claude_oauth_token_json(token: str) -> str:
+    return json.dumps(
+        {"type": "claude-code-oauth-token", "token": token}, separators=(",", ":")
+    )
+
+
+def _grok_auth_path() -> Path:
+    grok_home = os.environ.get("GROK_HOME", "").strip()
+    return (
+        Path(grok_home).expanduser() / "auth.json"
+        if grok_home
+        else Path.home() / ".grok" / "auth.json"
+    )
+
+
+def _grok_oauth_token_json(auth_json: str) -> str:
+    try:
+        payload = json.loads(auth_json)
+    except json.JSONDecodeError as exc:
+        raise RuntimeConfigError("Grok auth.json is not valid JSON") from exc
+    if not isinstance(payload, dict):
+        raise RuntimeConfigError("Grok auth.json must contain an account login")
+
+    candidates: list[tuple[str, str]] = []
+    for value in payload.values():
+        if not isinstance(value, dict):
+            continue
+        token = str(value.get("key") or "").strip()
+        if token:
+            candidates.append((str(value.get("expires_at") or ""), token))
+    if not candidates:
+        raise RuntimeConfigError("Grok auth.json does not contain a login token")
+
+    token = max(candidates, key=lambda candidate: candidate[0])[1]
+    return json.dumps(
+        {"type": "grok-oauth-token", "token": token}, separators=(",", ":")
+    )
+
+
+def _xai_api_key_json(api_key: str) -> str:
+    api_key = api_key.strip()
+    if not api_key:
+        raise RuntimeConfigError("xAI API key is required")
+    return json.dumps(
+        {"type": "xai-api-key", "api_key": api_key}, separators=(",", ":")
+    )
+
+
+def _xai_api_key(api_key: str | None) -> str | None:
+    resolved = (api_key or "").strip()
+    if resolved:
+        return resolved
+    resolved = os.environ.get("XAI_API_KEY", "").strip()
+    if resolved:
+        return resolved
+    if not cli_output.interactive():
+        return None
+    return cli_output.prompt_password("xAI API key")
+
+
+def _find_provider(credentials: dict[str, Any], provider: str) -> dict[str, Any] | None:
+    for credential in credentials.get("provider_credentials") or []:
+        if credential.get("provider") == provider:
+            return credential
+    return None
+
+
+def handle_provider_status(config: RuntimeConfig, provider: str) -> int:
+    if (err := cli_output.require_auth(config)) is not None:
+        return err
+    client = authenticated_client(config)
+    credential = _find_provider(client.list_provider_credentials(), provider)
+    return cli_output.report_success(
+        {
+            "provider": provider,
+            "saved": credential is not None,
+            "provider_credential": credential,
+        }
+    )
+
+
+def handle_provider_logout(config: RuntimeConfig, provider: str) -> int:
+    if (err := cli_output.require_auth(config)) is not None:
+        return err
+    client = authenticated_client(config)
+    try:
+        client.delete_provider_credential(provider)
+    except RuntimeAPIError as exc:
+        if exc.status_code != 404:
+            raise
+    return cli_output.report_success({"provider": provider, "saved": False})
+
+
+def handle_codex_status(config: RuntimeConfig) -> int:
+    return handle_provider_status(config, "codex")
+
+
+def handle_codex_login(config: RuntimeConfig, *, device: bool = False) -> int:
+    if (err := cli_output.require_auth(config)) is not None:
+        return err
+    local_auth_path = _local_codex_auth_path()
+    if local_auth_path.exists():
+        auth_json = _read_provider_auth_file(local_auth_path)
+    else:
+        codex_bin = shutil.which("codex")
+        if not codex_bin:
+            return cli_output.report_error("codex CLI is not installed")
+
+        with tempfile.TemporaryDirectory(prefix="runtime-codex-") as codex_home:
+            home = Path(codex_home)
+            status = _run_codex_login(codex_bin, home, device=device)
+            if status != 0:
+                return cli_output.report_error(
+                    f"codex login failed with exit code {status}"
+                )
+            auth_path = _codex_auth_path(home)
+            if not auth_path.exists():
+                return cli_output.report_error(
+                    "codex login completed but did not write auth.json"
+                )
+            auth_json = _read_provider_auth_file(auth_path)
+
+    client = authenticated_client(config)
+    credential = client.put_provider_credential("codex", auth_json)
+    return cli_output.report_success(
+        {"provider": "codex", "saved": True, "provider_credential": credential}
+    )
+
+
+def handle_codex_logout(config: RuntimeConfig) -> int:
+    return handle_provider_logout(config, "codex")
+
+
+def handle_claude_status(config: RuntimeConfig) -> int:
+    return handle_provider_status(config, "claude")
+
+
+def handle_claude_login(config: RuntimeConfig) -> int:
+    if (err := cli_output.require_auth(config)) is not None:
+        return err
+    claude_bin = shutil.which("claude")
+    if not claude_bin:
+        return cli_output.report_error("claude CLI is not installed")
+
+    with tempfile.TemporaryDirectory(prefix="runtime-claude-") as claude_home:
+        home = Path(claude_home)
+        token = _run_claude_setup_token(claude_bin, home)
+        auth_json = _claude_oauth_token_json(token)
+
+    client = authenticated_client(config)
+    credential = client.put_provider_credential("claude", auth_json)
+    return cli_output.report_success(
+        {"provider": "claude", "saved": True, "provider_credential": credential}
+    )
+
+
+def handle_claude_logout(config: RuntimeConfig) -> int:
+    return handle_provider_logout(config, "claude")
+
+
+def handle_grok_status(config: RuntimeConfig) -> int:
+    return handle_provider_status(config, "grok")
+
+
+def handle_grok_login(config: RuntimeConfig) -> int:
+    if (err := cli_output.require_auth(config)) is not None:
+        return err
+    auth_path = _grok_auth_path()
+    if not auth_path.exists():
+        grok_bin = shutil.which("grok")
+        if not grok_bin:
+            return cli_output.report_error("grok CLI is not installed")
+        status = subprocess.run([grok_bin, "login"]).returncode
+        if status != 0:
+            return cli_output.report_error(f"grok login failed with exit code {status}")
+    if not auth_path.exists():
+        return cli_output.report_error(
+            "grok login completed but did not write auth.json"
+        )
+
+    auth_json = _grok_oauth_token_json(_read_provider_auth_file(auth_path))
+    client = authenticated_client(config)
+    credential = client.put_provider_credential("grok", auth_json)
+    return cli_output.report_success(
+        {"provider": "grok", "saved": True, "provider_credential": credential}
+    )
+
+
+def handle_grok_logout(config: RuntimeConfig) -> int:
+    return handle_provider_logout(config, "grok")
+
+
+def handle_xai_status(config: RuntimeConfig) -> int:
+    return handle_provider_status(config, "xai")
+
+
+def handle_xai_login(config: RuntimeConfig, *, api_key: str | None = None) -> int:
+    if (err := cli_output.require_auth(config)) is not None:
+        return err
+    resolved_key = _xai_api_key(api_key)
+    if not resolved_key:
+        return cli_output.report_error(
+            "xAI API key is required; pass --api-key or set XAI_API_KEY"
+        )
+    auth_json = _xai_api_key_json(resolved_key)
+    client = authenticated_client(config)
+    credential = client.put_provider_credential("xai", auth_json)
+    return cli_output.report_success(
+        {"provider": "xai", "saved": True, "provider_credential": credential}
+    )
+
+
+def handle_xai_logout(config: RuntimeConfig) -> int:
+    return handle_provider_logout(config, "xai")