axion-code 1.0.0__py3-none-any.whl

axion/runtime/memory.py

@@ -0,0 +1,196 @@
+"""Persistent memory system.
+
+Maps to: rust/crates/runtime/src/memory.rs
+
+Provides a file-backed memory store where entries are saved as .md files
+with YAML frontmatter.  An index file (MEMORY.md) tracks all entries.
+"""
+
+from __future__ import annotations
+
+import enum
+import logging
+import re
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_MEMORY_DIR = Path.home() / ".axion" / "memory"
+
+# Reuse the lightweight frontmatter parser from skills
+_FRONTMATTER_RE = re.compile(
+    r"\A---\s*\n(.*?)\n---\s*\n(.*)",
+    re.DOTALL,
+)
+
+
+class MemoryType(enum.Enum):
+    USER = "user"
+    FEEDBACK = "feedback"
+    PROJECT = "project"
+    REFERENCE = "reference"
+
+
+@dataclass
+class MemoryEntry:
+    """A single memory entry."""
+
+    name: str
+    description: str
+    type: MemoryType
+    content: str
+    created_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
+    updated_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
+
+
+# ---------------------------------------------------------------------------
+# Frontmatter helpers
+# ---------------------------------------------------------------------------
+
+
+def _parse_frontmatter(text: str) -> tuple[dict[str, str], str]:
+    match = _FRONTMATTER_RE.match(text)
+    if not match:
+        return {}, text
+    meta: dict[str, str] = {}
+    for line in match.group(1).splitlines():
+        line = line.strip()
+        if not line or line.startswith("#"):
+            continue
+        if ":" in line:
+            key, _, value = line.partition(":")
+            meta[key.strip()] = value.strip()
+    return meta, match.group(2)
+
+
+def _build_frontmatter(entry: MemoryEntry) -> str:
+    lines = [
+        "---",
+        f"name: {entry.name}",
+        f"description: {entry.description}",
+        f"type: {entry.type.value}",
+        f"created_at: {entry.created_at}",
+        f"updated_at: {entry.updated_at}",
+        "---",
+        "",
+        entry.content,
+    ]
+    return "\n".join(lines)
+
+
+# ---------------------------------------------------------------------------
+# MemoryStore
+# ---------------------------------------------------------------------------
+
+
+class MemoryStore:
+    """File-backed memory store.
+
+    Each entry is stored as a .md file under *memory_dir*.  An index
+    file (MEMORY.md) provides a human-readable listing.
+    """
+
+    def __init__(self, memory_dir: Path | None = None) -> None:
+        self.memory_dir = memory_dir or _DEFAULT_MEMORY_DIR
+
+    def _ensure_dir(self) -> None:
+        self.memory_dir.mkdir(parents=True, exist_ok=True)
+
+    def _entry_path(self, name: str) -> Path:
+        safe_name = re.sub(r"[^\w\-.]", "_", name)
+        return self.memory_dir / f"{safe_name}.md"
+
+    # -- CRUD ---------------------------------------------------------------
+
+    def save(self, entry: MemoryEntry) -> Path:
+        """Write an entry as a .md file with YAML frontmatter."""
+        self._ensure_dir()
+        entry.updated_at = datetime.now(timezone.utc).isoformat()
+        path = self._entry_path(entry.name)
+        path.write_text(_build_frontmatter(entry), encoding="utf-8")
+        logger.debug("Saved memory entry '%s' to %s", entry.name, path)
+        return path
+
+    def load(self, name: str) -> MemoryEntry | None:
+        """Read a single memory file by name."""
+        path = self._entry_path(name)
+        if not path.is_file():
+            return None
+        return self._read_entry(path)
+
+    def load_all(self) -> list[MemoryEntry]:
+        """Read all memory files in the store."""
+        if not self.memory_dir.is_dir():
+            return []
+        entries: list[MemoryEntry] = []
+        for path in sorted(self.memory_dir.glob("*.md")):
+            if path.name == "MEMORY.md":
+                continue
+            entry = self._read_entry(path)
+            if entry is not None:
+                entries.append(entry)
+        return entries
+
+    def remove(self, name: str) -> bool:
+        """Delete a memory file. Returns True if the file existed."""
+        path = self._entry_path(name)
+        if path.is_file():
+            path.unlink()
+            logger.debug("Removed memory entry '%s'", name)
+            return True
+        return False
+
+    # -- Index --------------------------------------------------------------
+
+    def load_index(self) -> str | None:
+        """Read the MEMORY.md index file."""
+        index_path = self.memory_dir / "MEMORY.md"
+        if index_path.is_file():
+            return index_path.read_text(encoding="utf-8")
+        return None
+
+    def save_index(self, entries: list[MemoryEntry] | None = None) -> Path:
+        """Write a MEMORY.md index listing all entries."""
+        self._ensure_dir()
+        if entries is None:
+            entries = self.load_all()
+
+        lines = ["# Memory Index", ""]
+        for entry in entries:
+            lines.append(
+                f"- **{entry.name}** ({entry.type.value}): {entry.description}"
+            )
+        lines.append("")
+
+        index_path = self.memory_dir / "MEMORY.md"
+        index_path.write_text("\n".join(lines), encoding="utf-8")
+        logger.debug("Saved memory index with %d entries", len(entries))
+        return index_path
+
+    # -- Internal -----------------------------------------------------------
+
+    @staticmethod
+    def _read_entry(path: Path) -> MemoryEntry | None:
+        try:
+            text = path.read_text(encoding="utf-8")
+        except OSError as exc:
+            logger.warning("Failed to read memory file %s: %s", path, exc)
+            return None
+
+        meta, content = _parse_frontmatter(text)
+
+        try:
+            mem_type = MemoryType(meta.get("type", "reference"))
+        except ValueError:
+            mem_type = MemoryType.REFERENCE
+
+        return MemoryEntry(
+            name=meta.get("name", path.stem),
+            description=meta.get("description", ""),
+            type=mem_type,
+            content=content.strip(),
+            created_at=meta.get("created_at", ""),
+            updated_at=meta.get("updated_at", ""),
+        )

axion/runtime/oauth.py

@@ -0,0 +1,329 @@
+"""OAuth PKCE flow with HTTP callback server, token refresh, and browser launch.
+
+Maps to: rust/crates/runtime/src/oauth.rs
+"""
+
+from __future__ import annotations
+
+import base64
+import hashlib
+import http.server
+import json
+import logging
+import os
+import platform
+import secrets
+import subprocess
+import threading
+import time
+import urllib.parse
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_OAUTH_CALLBACK_PORT = 4545
+DEFAULT_CALLBACK_PATH = "/oauth/callback"
+
+
+@dataclass
+class OAuthTokenSet:
+    access_token: str
+    refresh_token: str | None = None
+    expires_at: int | None = None
+    scopes: list[str] = field(default_factory=list)
+
+    def is_expired(self) -> bool:
+        if self.expires_at is None:
+            return False
+        return int(time.time()) >= self.expires_at
+
+
+@dataclass
+class OAuthConfig:
+    client_id: str = ""
+    authorize_url: str = "https://console.anthropic.com/oauth/authorize"
+    token_url: str = "https://console.anthropic.com/oauth/token"
+    callback_port: int = DEFAULT_OAUTH_CALLBACK_PORT
+    scopes: list[str] = field(default_factory=lambda: ["user:inference"])
+
+
+@dataclass
+class PkceCodePair:
+    code_verifier: str
+    code_challenge: str
+
+
+@dataclass
+class OAuthCallbackParams:
+    code: str | None = None
+    state: str | None = None
+    error: str | None = None
+    error_description: str | None = None
+
+
+# ---------------------------------------------------------------------------
+# PKCE helpers
+# ---------------------------------------------------------------------------
+
+def generate_pkce_pair() -> PkceCodePair:
+    """Generate a PKCE code verifier and S256 challenge."""
+    verifier = secrets.token_urlsafe(64)
+    digest = hashlib.sha256(verifier.encode("ascii")).digest()
+    challenge = base64.urlsafe_b64encode(digest).rstrip(b"=").decode("ascii")
+    return PkceCodePair(code_verifier=verifier, code_challenge=challenge)
+
+
+def generate_state() -> str:
+    """Generate a random OAuth state parameter."""
+    return secrets.token_urlsafe(32)
+
+
+# ---------------------------------------------------------------------------
+# Credential persistence
+# ---------------------------------------------------------------------------
+
+def _credentials_path(provider: str) -> Path:
+    return Path.home() / ".axion" / "credentials" / f"{provider}.json"
+
+
+def save_oauth_credentials(provider: str, token_set: OAuthTokenSet) -> None:
+    """Save OAuth credentials to disk."""
+    path = _credentials_path(provider)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    data = {
+        "access_token": token_set.access_token,
+        "refresh_token": token_set.refresh_token,
+        "expires_at": token_set.expires_at,
+        "scopes": token_set.scopes,
+    }
+    path.write_text(json.dumps(data, indent=2), encoding="utf-8")
+    try:
+        os.chmod(path, 0o600)
+    except OSError:
+        pass
+
+
+def load_oauth_credentials(provider: str) -> OAuthTokenSet | None:
+    """Load OAuth credentials from disk."""
+    path = _credentials_path(provider)
+    if not path.exists():
+        return None
+    try:
+        data = json.loads(path.read_text(encoding="utf-8"))
+        return OAuthTokenSet(
+            access_token=data["access_token"],
+            refresh_token=data.get("refresh_token"),
+            expires_at=data.get("expires_at"),
+            scopes=data.get("scopes", []),
+        )
+    except (json.JSONDecodeError, KeyError, OSError):
+        return None
+
+
+def clear_oauth_credentials(provider: str) -> None:
+    """Remove OAuth credentials from disk."""
+    path = _credentials_path(provider)
+    if path.exists():
+        path.unlink()
+
+
+# ---------------------------------------------------------------------------
+# Browser launch (platform-specific)
+# ---------------------------------------------------------------------------
+
+def open_browser(url: str) -> bool:
+    """Open a URL in the default browser. Returns True on success.
+
+    On Windows, `cmd /C start` interprets `&` in URLs as command separators,
+    truncating OAuth URLs with multiple query params. We use os.startfile
+    (Windows) and webbrowser.open (cross-platform) which handle this correctly.
+    """
+    system = platform.system().lower()
+    try:
+        if system == "windows":
+            # os.startfile preserves & in URLs; cmd start does not
+            os.startfile(url)
+            return True
+        elif system == "darwin":
+            subprocess.Popen(["open", url], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+            return True
+        elif system == "linux":
+            subprocess.Popen(
+                ["xdg-open", url],
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+            )
+            return True
+    except (FileNotFoundError, OSError, AttributeError) as exc:
+        logger.warning("Native browser launcher failed: %s, falling back to webbrowser module", exc)
+
+    # Fallback: Python stdlib webbrowser (handles all platforms safely)
+    try:
+        import webbrowser
+        return webbrowser.open(url)
+    except Exception as exc:
+        logger.warning("Failed to open browser: %s", exc)
+    return False
+
+
+# ---------------------------------------------------------------------------
+# OAuth callback HTTP server
+# ---------------------------------------------------------------------------
+
+class _OAuthCallbackHandler(http.server.BaseHTTPRequestHandler):
+    """HTTP handler that captures the OAuth callback parameters."""
+
+    callback_result: OAuthCallbackParams | None = None
+
+    def do_GET(self) -> None:
+        parsed = urllib.parse.urlparse(self.path)
+        params = urllib.parse.parse_qs(parsed.query)
+
+        result = OAuthCallbackParams(
+            code=params.get("code", [None])[0],
+            state=params.get("state", [None])[0],
+            error=params.get("error", [None])[0],
+            error_description=params.get("error_description", [None])[0],
+        )
+        _OAuthCallbackHandler.callback_result = result
+
+        # Send response
+        if result.error:
+            body = f"<h1>OAuth Error</h1><p>{result.error}: {result.error_description}</p>"
+            self.send_response(400)
+        else:
+            body = "<h1>Success!</h1><p>You can close this window and return to Axion Code.</p>"
+            self.send_response(200)
+
+        self.send_header("Content-Type", "text/html")
+        self.end_headers()
+        self.wfile.write(body.encode())
+
+    def log_message(self, format: str, *args: Any) -> None:
+        # Suppress default logging
+        pass
+
+
+def wait_for_oauth_callback(
+    port: int = DEFAULT_OAUTH_CALLBACK_PORT,
+    timeout: float = 120.0,
+) -> OAuthCallbackParams | None:
+    """Start a local HTTP server and wait for the OAuth callback.
+
+    Returns the callback parameters or None on timeout.
+    """
+    _OAuthCallbackHandler.callback_result = None
+
+    server = http.server.HTTPServer(("127.0.0.1", port), _OAuthCallbackHandler)
+    server.timeout = timeout
+
+    # Run in a thread with timeout
+    result: OAuthCallbackParams | None = None
+
+    def serve() -> None:
+        nonlocal result
+        server.handle_request()  # Handle exactly one request
+        result = _OAuthCallbackHandler.callback_result
+
+    thread = threading.Thread(target=serve, daemon=True)
+    thread.start()
+    thread.join(timeout=timeout)
+
+    server.server_close()
+    return result
+
+
+# ---------------------------------------------------------------------------
+# Token exchange and refresh
+# ---------------------------------------------------------------------------
+
+async def exchange_authorization_code(
+    token_url: str,
+    code: str,
+    code_verifier: str,
+    client_id: str,
+    redirect_uri: str,
+) -> OAuthTokenSet:
+    """Exchange an authorization code for tokens."""
+    import httpx
+
+    async with httpx.AsyncClient() as client:
+        response = await client.post(
+            token_url,
+            data={
+                "grant_type": "authorization_code",
+                "code": code,
+                "code_verifier": code_verifier,
+                "client_id": client_id,
+                "redirect_uri": redirect_uri,
+            },
+        )
+        response.raise_for_status()
+        data = response.json()
+
+    expires_in = data.get("expires_in")
+    expires_at = int(time.time()) + expires_in if expires_in else None
+
+    return OAuthTokenSet(
+        access_token=data["access_token"],
+        refresh_token=data.get("refresh_token"),
+        expires_at=expires_at,
+        scopes=data.get("scope", "").split() if data.get("scope") else [],
+    )
+
+
+async def refresh_token(
+    token_url: str,
+    refresh_token_str: str,
+    client_id: str,
+) -> OAuthTokenSet:
+    """Refresh an expired OAuth token."""
+    import httpx
+
+    async with httpx.AsyncClient() as client:
+        response = await client.post(
+            token_url,
+            data={
+                "grant_type": "refresh_token",
+                "refresh_token": refresh_token_str,
+                "client_id": client_id,
+            },
+        )
+        response.raise_for_status()
+        data = response.json()
+
+    expires_in = data.get("expires_in")
+    expires_at = int(time.time()) + expires_in if expires_in else None
+
+    return OAuthTokenSet(
+        access_token=data["access_token"],
+        refresh_token=data.get("refresh_token", refresh_token_str),
+        expires_at=expires_at,
+        scopes=data.get("scope", "").split() if data.get("scope") else [],
+    )
+
+
+# ---------------------------------------------------------------------------
+# Full login flow
+# ---------------------------------------------------------------------------
+
+def build_authorization_url(
+    config: OAuthConfig,
+    pkce: PkceCodePair,
+    state: str,
+) -> str:
+    """Build the full OAuth authorization URL."""
+    params = {
+        "client_id": config.client_id,
+        "response_type": "code",
+        "redirect_uri": f"http://127.0.0.1:{config.callback_port}{DEFAULT_CALLBACK_PATH}",
+        "state": state,
+        "code_challenge": pkce.code_challenge,
+        "code_challenge_method": "S256",
+    }
+    if config.scopes:
+        params["scope"] = " ".join(config.scopes)
+
+    return f"{config.authorize_url}?{urllib.parse.urlencode(params)}"