axion-code 1.0.0__py3-none-any.whl

axion/runtime/claude_subscription.py

@@ -0,0 +1,300 @@
+"""Claude Pro/Max subscription OAuth — bypass API billing using your subscription.
+
+This is the same OAuth flow Claude Code uses. When authenticated via subscription,
+requests are billed against your Claude Pro/Max plan instead of pay-per-token API.
+
+Flow (paste-style, like Claude Code):
+  1. Open https://claude.ai/oauth/authorize?...&redirect_uri=https://platform.claude.com/oauth/code/success?app=claude-code
+  2. User logs in and authorizes
+  3. claude.ai redirects to platform.claude.com/oauth/code/success which displays the code
+  4. User copies the code and pastes it into the CLI
+  5. We exchange the code for an access token at console.anthropic.com
+  6. Use the token as `Authorization: Bearer <token>` with `anthropic-beta: oauth-2025-04-20`
+
+Tokens are saved to ~/.axion/credentials/anthropic-oauth.json and auto-refreshed.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+import urllib.parse
+from dataclasses import dataclass
+
+from axion.runtime.oauth import (
+    OAuthTokenSet,
+    PkceCodePair,
+    clear_oauth_credentials,
+    generate_pkce_pair,
+    generate_state,
+    load_oauth_credentials,
+    open_browser,
+    save_oauth_credentials,
+)
+
+logger = logging.getLogger(__name__)
+
+# Claude Code's well-known OAuth client ID for subscription auth
+CLAUDE_CODE_CLIENT_ID = "9d1c250a-e61b-44d9-88ed-5944d1962f5e"
+
+# Subscription OAuth endpoints
+CLAUDE_AUTHORIZE_URL = "https://claude.ai/oauth/authorize"
+CLAUDE_TOKEN_URL = "https://console.anthropic.com/v1/oauth/token"
+
+# Registered redirect URI for Claude Code's OAuth client.
+# When `code=true` is set in the authorize URL, claude.ai displays the code
+# on https://platform.claude.com/oauth/code/success?app=claude-code instead
+# of actually redirecting — but the OAuth `redirect_uri` parameter must still
+# match what's registered for this client.
+CLAUDE_REDIRECT_URI = "https://console.anthropic.com/oauth/code/callback"
+
+# OAuth requires this beta header on Messages API requests
+SUBSCRIPTION_BETA_HEADER = "oauth-2025-04-20"
+
+# Scopes required for subscription-based inference
+SUBSCRIPTION_SCOPES = ["org:create_api_key", "user:profile", "user:inference"]
+
+# Provider key for credential storage
+SUBSCRIPTION_PROVIDER = "anthropic-oauth"
+
+
+@dataclass
+class SubscriptionAuthResult:
+    """Result of a subscription OAuth login attempt."""
+
+    success: bool
+    token_set: OAuthTokenSet | None = None
+    error: str | None = None
+
+
+def build_subscription_authorize_url(pkce: PkceCodePair, state: str) -> str:
+    """Build the claude.ai authorize URL for subscription auth."""
+    params = {
+        "code": "true",  # Tells claude.ai this is for paste-style flow
+        "client_id": CLAUDE_CODE_CLIENT_ID,
+        "response_type": "code",
+        "redirect_uri": CLAUDE_REDIRECT_URI,
+        "scope": " ".join(SUBSCRIPTION_SCOPES),
+        "state": state,
+        "code_challenge": pkce.code_challenge,
+        "code_challenge_method": "S256",
+    }
+    return f"{CLAUDE_AUTHORIZE_URL}?{urllib.parse.urlencode(params)}"
+
+
+def parse_pasted_code(pasted: str) -> tuple[str, str | None]:
+    """Parse the code from whatever the user pasted.
+
+    Handles all common paste formats:
+      - `code` — just the code
+      - `code#state` — Claude's success page format ("Paste this into Claude Code")
+      - `https://platform.claude.com/oauth/code/callback?code=ABC&state=XYZ` — full URL
+      - `?code=ABC&state=XYZ` — query string only
+      - URLs with the authorize endpoint (user pasted wrong thing) — return empty
+    """
+    pasted = pasted.strip()
+
+    if not pasted:
+        return "", None
+
+    # If they pasted the AUTHORIZE URL by mistake, reject it
+    if "/oauth/authorize" in pasted:
+        return "", None
+
+    # If they pasted a full URL or query string, extract code & state
+    if pasted.startswith(("http://", "https://", "?")):
+        from urllib.parse import parse_qs, urlparse
+
+        if pasted.startswith("?"):
+            # Just a query string
+            qs = pasted[1:]
+        else:
+            qs = urlparse(pasted).query
+
+        params = parse_qs(qs)
+        code = (params.get("code") or [""])[0].strip()
+        state = (params.get("state") or [None])[0]
+        if state:
+            state = state.strip() or None
+        return code, state
+
+    # Claude's success page format: code#state
+    if "#" in pasted:
+        code, _, state = pasted.partition("#")
+        return code.strip(), state.strip() or None
+
+    # Plain code
+    return pasted, None
+
+
+async def exchange_subscription_code(
+    code: str,
+    code_verifier: str,
+    state: str,
+) -> OAuthTokenSet:
+    """Exchange an authorization code for subscription tokens."""
+    import httpx
+
+    payload = {
+        "grant_type": "authorization_code",
+        "code": code,
+        "code_verifier": code_verifier,
+        "client_id": CLAUDE_CODE_CLIENT_ID,
+        "redirect_uri": CLAUDE_REDIRECT_URI,
+        "state": state,
+    }
+
+    async with httpx.AsyncClient(timeout=30.0) as client:
+        response = await client.post(
+            CLAUDE_TOKEN_URL,
+            json=payload,
+            headers={"Content-Type": "application/json"},
+        )
+        if response.status_code != 200:
+            raise RuntimeError(
+                f"Token exchange failed ({response.status_code}): {response.text[:500]}"
+            )
+        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 SUBSCRIPTION_SCOPES,
+    )
+
+
+async def refresh_subscription_token(refresh_token_str: str) -> OAuthTokenSet:
+    """Refresh an expired subscription access token."""
+    import httpx
+
+    payload = {
+        "grant_type": "refresh_token",
+        "refresh_token": refresh_token_str,
+        "client_id": CLAUDE_CODE_CLIENT_ID,
+    }
+
+    async with httpx.AsyncClient(timeout=30.0) as client:
+        response = await client.post(
+            CLAUDE_TOKEN_URL,
+            json=payload,
+            headers={"Content-Type": "application/json"},
+        )
+        if response.status_code != 200:
+            raise RuntimeError(
+                f"Token refresh failed ({response.status_code}): {response.text[:500]}"
+            )
+        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 SUBSCRIPTION_SCOPES,
+    )
+
+
+async def begin_subscription_login(
+    *, open_browser_automatically: bool = True
+) -> tuple[str, PkceCodePair, str]:
+    """Start the subscription login flow.
+
+    Returns (auth_url, pkce_pair, state). Caller should:
+      1. Tell the user to open auth_url (we try to open the browser)
+      2. After authorizing, the user will see a code on the success page
+      3. Caller prompts the user to paste the code
+      4. Caller calls complete_subscription_login(code, pkce, state)
+    """
+    pkce = generate_pkce_pair()
+    state = generate_state()
+    auth_url = build_subscription_authorize_url(pkce, state)
+
+    if open_browser_automatically:
+        open_browser(auth_url)
+
+    return auth_url, pkce, state
+
+
+async def complete_subscription_login(
+    pasted: str,
+    pkce: PkceCodePair,
+    expected_state: str,
+) -> SubscriptionAuthResult:
+    """Finish subscription login using the pasted code from the success page."""
+    # Detect if user pasted the authorize URL by mistake
+    if "/oauth/authorize" in pasted:
+        return SubscriptionAuthResult(
+            success=False,
+            error="That's the authorize URL, not the code. After logging in, the success "
+                  "page shows an 'Authentication Code'. Copy that and paste it here.",
+        )
+
+    code, pasted_state = parse_pasted_code(pasted)
+
+    if not code:
+        return SubscriptionAuthResult(
+            success=False,
+            error="No code found in what you pasted. Copy the value from the "
+                  "'Authentication Code' box on the success page.",
+        )
+
+    # If the user pasted code#state, verify state for CSRF protection
+    if pasted_state and pasted_state != expected_state:
+        return SubscriptionAuthResult(
+            success=False,
+            error="State mismatch — the code didn't come from the login you started. "
+                  "Try again from scratch.",
+        )
+
+    try:
+        token_set = await exchange_subscription_code(
+            code=code,
+            code_verifier=pkce.code_verifier,
+            state=expected_state,
+        )
+    except Exception as exc:
+        return SubscriptionAuthResult(
+            success=False,
+            error=f"Token exchange failed: {exc}",
+        )
+
+    save_oauth_credentials(SUBSCRIPTION_PROVIDER, token_set)
+    return SubscriptionAuthResult(success=True, token_set=token_set)
+
+
+async def get_valid_subscription_token() -> str | None:
+    """Get a valid subscription access token, refreshing if needed.
+
+    Returns None if no subscription credentials are saved.
+    """
+    creds = load_oauth_credentials(SUBSCRIPTION_PROVIDER)
+    if creds is None:
+        return None
+
+    # If expired, refresh
+    if creds.is_expired() and creds.refresh_token:
+        try:
+            new_creds = await refresh_subscription_token(creds.refresh_token)
+            save_oauth_credentials(SUBSCRIPTION_PROVIDER, new_creds)
+            return new_creds.access_token
+        except Exception as exc:
+            logger.warning("Subscription token refresh failed: %s", exc)
+            return None
+
+    return creds.access_token
+
+
+def has_subscription_credentials() -> bool:
+    """Check if subscription credentials are saved (without validating)."""
+    return load_oauth_credentials(SUBSCRIPTION_PROVIDER) is not None
+
+
+def logout_subscription() -> None:
+    """Remove subscription credentials."""
+    clear_oauth_credentials(SUBSCRIPTION_PROVIDER)

axion/runtime/compact.py

@@ -0,0 +1,233 @@
+"""Session compaction.
+
+Maps to: rust/crates/runtime/src/compact.rs
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from axion.runtime.session import (
+    ConversationMessage,
+    MessageRole,
+    Session,
+    SessionCompaction,
+    TextBlock,
+)
+
+if TYPE_CHECKING:
+    from axion.api.client import ProviderClient
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class CompactionConfig:
+    """Configuration for session compaction."""
+
+    max_tokens: int = 100_000
+    preserve_recent_messages: int = 4
+    summary_max_tokens: int = 2000
+
+
+@dataclass
+class CompactionResult:
+    """Result of a compaction operation."""
+
+    removed_count: int
+    summary: str
+    estimated_tokens_before: int
+    estimated_tokens_after: int
+
+
+def estimate_session_tokens(session: Session) -> int:
+    """Rough estimate of token count in a session.
+
+    Uses ~4 chars per token heuristic.
+    """
+    total_chars = 0
+    for msg in session.messages:
+        for block in msg.blocks:
+            if isinstance(block, TextBlock):
+                total_chars += len(block.text)
+            else:
+                total_chars += 100  # Rough estimate for non-text blocks
+    return total_chars // 4
+
+
+def compact_session(
+    session: Session,
+    config: CompactionConfig | None = None,
+) -> CompactionResult | None:
+    """Compact a session by summarizing old messages.
+
+    Preserves the most recent messages and replaces older ones
+    with a summary message.
+    """
+    cfg = config or CompactionConfig()
+
+    estimated_tokens = estimate_session_tokens(session)
+    if estimated_tokens < cfg.max_tokens:
+        return None
+
+    total_messages = len(session.messages)
+    if total_messages <= cfg.preserve_recent_messages:
+        return None
+
+    # Split into old and recent
+    split_idx = total_messages - cfg.preserve_recent_messages
+    old_messages = session.messages[:split_idx]
+    recent_messages = session.messages[split_idx:]
+
+    # Build summary from old messages
+    summary_parts: list[str] = []
+    for msg in old_messages:
+        for block in msg.blocks:
+            if isinstance(block, TextBlock):
+                # Truncate long texts
+                text = block.text[:200] + "..." if len(block.text) > 200 else block.text
+                summary_parts.append(f"[{msg.role.value}] {text}")
+
+    summary = "\n".join(summary_parts)
+    if len(summary) > cfg.summary_max_tokens * 4:
+        summary = summary[: cfg.summary_max_tokens * 4] + "\n... (summary truncated)"
+
+    # Replace old messages with summary
+    summary_message = ConversationMessage(
+        role=MessageRole.USER,
+        blocks=[TextBlock(text=f"[Session compacted. Summary of {split_idx} earlier messages:]\n{summary}")],
+    )
+
+    session.messages = [summary_message] + recent_messages
+
+    # Update compaction metadata
+    compaction_count = (session.compaction.count + 1) if session.compaction else 1
+    session.compaction = SessionCompaction(
+        count=compaction_count,
+        removed_message_count=split_idx,
+        summary=summary[:500],
+    )
+
+    tokens_after = estimate_session_tokens(session)
+
+    return CompactionResult(
+        removed_count=split_idx,
+        summary=summary[:200],
+        estimated_tokens_before=estimated_tokens,
+        estimated_tokens_after=tokens_after,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Model-based intelligent compaction
+# ---------------------------------------------------------------------------
+
+
+async def model_compact_session(
+    session: Session,
+    provider: "ProviderClient",
+    model: str,
+    config: CompactionConfig | None = None,
+) -> CompactionResult | None:
+    """Use a model to produce an intelligent summary for compaction.
+
+    1. Builds a summary request asking the model to summarize old messages.
+    2. Sends it to the API.
+    3. Replaces old messages with the model-generated summary.
+
+    Falls back to the heuristic ``compact_session`` if the API call fails.
+    """
+    from axion.api.client import max_tokens_for_model, resolve_model_alias
+    from axion.api.types import InputMessage, MessageRequest, TextInputBlock
+
+    cfg = config or CompactionConfig()
+
+    estimated_tokens = estimate_session_tokens(session)
+    if estimated_tokens < cfg.max_tokens:
+        return None
+
+    total_messages = len(session.messages)
+    if total_messages <= cfg.preserve_recent_messages:
+        return None
+
+    split_idx = total_messages - cfg.preserve_recent_messages
+    old_messages = session.messages[:split_idx]
+    recent_messages = session.messages[split_idx:]
+
+    # Build text representation of old messages for the summariser
+    old_text_parts: list[str] = []
+    for msg in old_messages:
+        for block in msg.blocks:
+            if isinstance(block, TextBlock):
+                text = block.text[:500] if len(block.text) > 500 else block.text
+                old_text_parts.append(f"[{msg.role.value}] {text}")
+
+    old_text = "\n".join(old_text_parts)
+
+    # Build a summarisation request
+    summary_prompt = (
+        "Summarise the following conversation history concisely. "
+        "Preserve key decisions, tool outputs, file paths, and any context "
+        "the assistant will need to continue the conversation.\n\n"
+        f"{old_text}"
+    )
+
+    resolved_model = resolve_model_alias(model)
+    request = MessageRequest(
+        model=resolved_model,
+        max_tokens=min(cfg.summary_max_tokens, max_tokens_for_model(resolved_model)),
+        messages=[
+            InputMessage(
+                role="user",
+                content=[TextInputBlock(text=summary_prompt)],
+            )
+        ],
+        system="You are a conversation summariser. Produce a concise summary.",
+        stream=False,
+    )
+
+    try:
+        response = await provider.send_message(request)
+        # Extract text from the response
+        summary_text = ""
+        if response.content:
+            for block in response.content:
+                if hasattr(block, "text"):
+                    summary_text += block.text
+        if not summary_text:
+            raise ValueError("Empty summary from model")
+    except Exception as exc:
+        logger.warning(
+            "Model-based compaction failed, falling back to heuristic: %s", exc
+        )
+        return compact_session(session, cfg)
+
+    # Replace old messages with model-generated summary
+    summary_message = ConversationMessage(
+        role=MessageRole.USER,
+        blocks=[
+            TextBlock(
+                text=f"[Session compacted by model. Summary of {split_idx} earlier messages:]\n{summary_text}"
+            )
+        ],
+    )
+
+    session.messages = [summary_message] + recent_messages
+
+    compaction_count = (session.compaction.count + 1) if session.compaction else 1
+    session.compaction = SessionCompaction(
+        count=compaction_count,
+        removed_message_count=split_idx,
+        summary=summary_text[:500],
+    )
+
+    tokens_after = estimate_session_tokens(session)
+
+    return CompactionResult(
+        removed_count=split_idx,
+        summary=summary_text[:200],
+        estimated_tokens_before=estimated_tokens,
+        estimated_tokens_after=tokens_after,
+    )