code-muse 0.0.1__py3-none-any.whl

code_muse/plugins/claude_code_oauth/config.py

@@ -0,0 +1,52 @@
+from pathlib import Path
+from typing import Any
+
+from code_muse import config
+
+# Claude Code OAuth configuration
+CLAUDE_CODE_OAUTH_CONFIG: dict[str, Any] = {
+    # OAuth endpoints inferred from official Claude Code OAuth flow
+    "auth_url": "https://claude.ai/oauth/authorize",
+    "token_url": "https://console.anthropic.com/v1/oauth/token",
+    "api_base_url": "https://api.anthropic.com",
+    # OAuth client configuration observed in Claude Code CLI flow
+    "client_id": "9d1c250a-e61b-44d9-88ed-5944d1962f5e",
+    "scope": "org:create_api_key user:profile user:inference",
+    # Callback handling (we host a localhost callback to capture the redirect)
+    "redirect_host": "http://localhost",
+    "redirect_path": "callback",
+    "callback_port_range": (8765, 8795),
+    "callback_timeout": 180,
+    # Console redirect fallback (for manual flows, if needed)
+    "console_redirect_uri": "https://console.anthropic.com/oauth/code/callback",
+    # Local configuration (uses XDG_DATA_HOME)
+    "token_storage": None,  # Set dynamically in get_token_storage_path()
+    # Model configuration
+    "prefix": "claude-code-",
+    "default_context_length": 200000,
+    "long_context_length": 1000000,
+    "long_context_models": ["claude-opus-4-6", "claude-opus-4-7"],
+    "api_key_env_var": "CLAUDE_CODE_ACCESS_TOKEN",
+    "anthropic_version": "2023-06-01",
+}
+
+
+def get_token_storage_path() -> Path:
+    """Get the 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 / "claude_code_oauth.json"
+
+
+def get_config_dir() -> Path:
+    """Get the Muse configuration directory (uses XDG_CONFIG_HOME)."""
+    config_dir = Path(config.CONFIG_DIR)
+    config_dir.mkdir(parents=True, exist_ok=True, mode=0o700)
+    return config_dir
+
+
+def get_claude_models_path() -> Path:
+    """Get the path to the dedicated claude_models.json file (uses XDG_DATA_HOME)."""
+    data_dir = Path(config.DATA_DIR)
+    data_dir.mkdir(parents=True, exist_ok=True, mode=0o700)
+    return data_dir / "claude_models.json"

code_muse/plugins/claude_code_oauth/fast_mode.py

@@ -0,0 +1,124 @@
+"""Fast mode support for Claude Code OAuth.
+
+Fast mode adds two wire-level changes when enabled via the per-model
+``fast`` setting:
+
+1. ``speed: "fast"`` is injected into the ``/v1/messages`` request payload.
+2. ``fast-mode-2026-02-01`` is appended to the ``anthropic-beta`` header.
+
+The setting is read at request time (not model construction time) so that
+toggling takes effect on the next call without a full agent reload for the
+payload side. The beta header must still be present on the httpx client at
+construction time, so we set that up once and rely on the existing
+interleaved-thinking pattern for consistency.
+"""
+
+from collections.abc import Callable, MutableMapping
+from typing import Any
+
+# Anthropic beta marker for fast mode. Keep as a module-level constant so
+# there's one source of truth for both header setup and tests.
+FAST_MODE_BETA = "fast-mode-2026-02-01"
+
+# The per-model setting key. Stored via code_muse.config.set_model_setting
+# as "true"/"false" strings and read back via get_all_model_settings.
+#
+# We intentionally DO NOT use get_effective_model_settings here: that helper
+# filters values through model_supports_setting(), whose default allowlist for
+# claude-* models is {temperature, extended_thinking, budget_tokens, effort}.
+# ``fast`` is a plugin-owned setting, not a core setting, so it would be
+# silently dropped. get_all_model_settings returns the raw per-model config,
+# which is exactly what a plugin should own and read.
+FAST_SETTING_KEY = "fast"
+
+
+def is_fast_mode_enabled(model_name: str) -> bool:
+    """Return True if the ``fast`` setting is on for ``model_name``."""
+    # Imported lazily to avoid circular imports at module load.
+    from code_muse.config import get_all_model_settings
+
+    settings = get_all_model_settings(model_name)
+    return bool(settings.get(FAST_SETTING_KEY, False))
+
+
+def ensure_fast_beta_header(headers: MutableMapping[str, str], enabled: bool) -> None:
+    """Add or remove ``fast-mode-2026-02-01`` from the anthropic-beta header.
+
+    Mutates ``headers`` in place. Mirrors the interleaved-thinking handling so
+    beta markers stay comma-separated and de-duplicated.
+    """
+    existing = headers.get("anthropic-beta")
+    parts = [p.strip() for p in existing.split(",") if p.strip()] if existing else []
+
+    has_marker = FAST_MODE_BETA in parts
+    if enabled and not has_marker:
+        parts.append(FAST_MODE_BETA)
+    elif not enabled and has_marker:
+        parts = [p for p in parts if p != FAST_MODE_BETA]
+
+    if parts:
+        headers["anthropic-beta"] = ",".join(parts)
+    elif "anthropic-beta" in headers:
+        del headers["anthropic-beta"]
+
+
+def _inject_speed_in_payload(payload: dict[str, Any]) -> None:
+    """Inject ``speed: "fast"`` into a /v1/messages payload dict in place."""
+    if not isinstance(payload, dict):
+        return
+    # Respect an explicit speed already present (don't clobber caller intent).
+    if "speed" not in payload:
+        payload["speed"] = "fast"
+
+
+def make_fast_mode_wrapper(
+    original_create: Callable[..., Any], model_name: str
+) -> Callable[..., Any]:
+    """Wrap ``messages.create`` to inject ``speed: fast`` when enabled.
+
+    The setting is re-read on every call so toggling via ``/claude-code-fast``
+    takes effect on the very next request (no agent reload needed for the
+    payload side of the equation).
+    """
+
+    async def wrapped_create(*args: Any, **kwargs: Any):
+        if is_fast_mode_enabled(model_name):
+            if kwargs:
+                _inject_speed_in_payload(kwargs)
+            elif args:
+                maybe_payload = args[-1]
+                if isinstance(maybe_payload, dict):
+                    _inject_speed_in_payload(maybe_payload)
+
+        return await original_create(*args, **kwargs)
+
+    return wrapped_create
+
+
+def patch_anthropic_client_fast_mode(client: Any, model_name: str) -> None:
+    """Monkey-patch ``messages.create`` (and ``beta.messages.create``).
+
+    Must be called *after* ``patch_anthropic_client_messages`` so our wrapper
+    sits outside the cache-control injector — order of wrappers doesn't matter
+    for correctness here, but keeping fast-mode outermost makes it easy to
+    reason about.
+    """
+    try:
+        messages_obj = getattr(client, "messages", None)
+        if messages_obj is not None and hasattr(messages_obj, "create"):
+            messages_obj.create = make_fast_mode_wrapper(
+                messages_obj.create, model_name
+            )
+    except Exception:  # pragma: no cover - defensive
+        pass
+
+    try:
+        beta_obj = getattr(client, "beta", None)
+        if beta_obj is not None:
+            beta_messages_obj = getattr(beta_obj, "messages", None)
+            if beta_messages_obj is not None and hasattr(beta_messages_obj, "create"):
+                beta_messages_obj.create = make_fast_mode_wrapper(
+                    beta_messages_obj.create, model_name
+                )
+    except Exception:  # pragma: no cover - defensive
+        pass

code_muse/plugins/claude_code_oauth/prompt_handler.py

@@ -0,0 +1,63 @@
+"""Prompt preparation logic for Claude Code OAuth models.
+
+This module owns everything that's special about preparing prompts for
+``claude-code-*`` models:
+
+- The fixed system-instruction string Anthropic's Claude Code CLI expects.
+- The ``is_claude_code_model`` predicate.
+- A callback wired into the ``prepare_model_prompt`` hook which runs inside
+  ``code_muse.model_utils.prepare_prompt_for_model``.
+
+Keeping this here (rather than in core ``model_utils``) keeps the core
+model-agnostic and lets the claude-code behavior ship/fail as a single
+plugin — exactly per the Contributing guide.
+"""
+
+from typing import Any
+
+# The instruction override Anthropic's Claude Code CLI expects as the system
+# prompt. When using OAuth-authenticated claude-code models, the *real* system
+# prompt has to move into the user turn instead (see the callback below).
+CLAUDE_CODE_INSTRUCTIONS = "You are Claude Code, Anthropic's official CLI for Claude."
+
+# Prefix used by every claude-code-* model (as registered by this plugin).
+_CLAUDE_CODE_PREFIX = "claude-code"
+
+
+def is_claude_code_model(model_name: str) -> bool:
+    """Return True if ``model_name`` is a claude-code OAuth model."""
+    return model_name.startswith(_CLAUDE_CODE_PREFIX)
+
+
+def get_claude_code_instructions() -> str:
+    """Return the fixed Claude Code system instruction string."""
+    return CLAUDE_CODE_INSTRUCTIONS
+
+
+def prepare_claude_code_prompt(
+    model_name: str,
+    system_prompt: str,
+    user_prompt: str,
+    prepend_system_to_user: bool = True,
+) -> dict[str, Any | None]:
+    """Callback for the ``prepare_model_prompt`` hook.
+
+    For claude-code models, swap in the Anthropic-expected instruction string
+    and (optionally) fold the caller's system prompt into the user message.
+
+    Returns ``None`` for non-claude-code models so other handlers / the
+    default passthrough can take over.
+    """
+    if not is_claude_code_model(model_name):
+        return None
+
+    modified_prompt = user_prompt
+    if prepend_system_to_user and system_prompt:
+        modified_prompt = f"{system_prompt}\n\n{user_prompt}"
+
+    return {
+        "handled": True,
+        "instructions": CLAUDE_CODE_INSTRUCTIONS,
+        "user_prompt": modified_prompt,
+        "is_claude_code": True,
+    }