code-muse 0.0.1__py3-none-any.whl

code_muse/plugins/claude_code_oauth/test_fast_mode.py

@@ -0,0 +1,165 @@
+"""Tests for the fast_mode submodule.
+
+These tests don't touch the network or config file - they exercise pure
+helpers for beta-header merging and payload injection.
+"""
+
+from unittest.mock import patch
+
+import pytest
+
+from code_muse.plugins.claude_code_oauth.fast_mode import (
+    FAST_MODE_BETA,
+    FAST_SETTING_KEY,
+    ensure_fast_beta_header,
+    is_fast_mode_enabled,
+    make_fast_mode_wrapper,
+)
+
+
+class TestEnsureFastBetaHeader:
+    def test_adds_marker_when_enabled_and_absent(self):
+        headers = {}
+        ensure_fast_beta_header(headers, enabled=True)
+        assert headers["anthropic-beta"] == FAST_MODE_BETA
+
+    def test_appends_marker_to_existing(self):
+        headers = {"anthropic-beta": "oauth-2025-04-20,interleaved-thinking-2025-05-14"}
+        ensure_fast_beta_header(headers, enabled=True)
+        parts = headers["anthropic-beta"].split(",")
+        assert FAST_MODE_BETA in parts
+        assert "oauth-2025-04-20" in parts
+        assert "interleaved-thinking-2025-05-14" in parts
+
+    def test_no_duplicate_when_already_present(self):
+        headers = {"anthropic-beta": f"oauth-2025-04-20,{FAST_MODE_BETA}"}
+        ensure_fast_beta_header(headers, enabled=True)
+        assert headers["anthropic-beta"].count(FAST_MODE_BETA) == 1
+
+    def test_removes_marker_when_disabled(self):
+        headers = {"anthropic-beta": f"oauth-2025-04-20,{FAST_MODE_BETA}"}
+        ensure_fast_beta_header(headers, enabled=False)
+        assert FAST_MODE_BETA not in headers["anthropic-beta"]
+        assert "oauth-2025-04-20" in headers["anthropic-beta"]
+
+    def test_drops_header_when_only_marker_removed(self):
+        headers = {"anthropic-beta": FAST_MODE_BETA}
+        ensure_fast_beta_header(headers, enabled=False)
+        assert "anthropic-beta" not in headers
+
+    def test_noop_when_disabled_and_absent(self):
+        headers = {"anthropic-beta": "oauth-2025-04-20"}
+        ensure_fast_beta_header(headers, enabled=False)
+        assert headers == {"anthropic-beta": "oauth-2025-04-20"}
+
+
+class TestIsFastModeEnabled:
+    def test_returns_false_when_setting_absent(self):
+        with patch(
+            "code_muse.config.get_all_model_settings",
+            return_value={},
+        ):
+            assert is_fast_mode_enabled("claude-code-foo") is False
+
+    def test_returns_true_when_setting_true(self):
+        with patch(
+            "code_muse.config.get_all_model_settings",
+            return_value={FAST_SETTING_KEY: True},
+        ):
+            assert is_fast_mode_enabled("claude-code-foo") is True
+
+    def test_returns_false_when_setting_false(self):
+        with patch(
+            "code_muse.config.get_all_model_settings",
+            return_value={FAST_SETTING_KEY: False},
+        ):
+            assert is_fast_mode_enabled("claude-code-foo") is False
+
+
+class TestMakeFastModeWrapper:
+    @pytest.mark.asyncio
+    async def test_injects_speed_when_enabled(self):
+        captured = {}
+
+        async def fake_create(**kwargs):
+            captured.update(kwargs)
+            return "ok"
+
+        wrapped = make_fast_mode_wrapper(fake_create, "claude-code-foo")
+
+        with patch(
+            "code_muse.config.get_all_model_settings",
+            return_value={FAST_SETTING_KEY: True},
+        ):
+            await wrapped(model="foo", messages=[])
+
+        assert captured.get("speed") == "fast"
+
+    @pytest.mark.asyncio
+    async def test_does_not_inject_when_disabled(self):
+        captured = {}
+
+        async def fake_create(**kwargs):
+            captured.update(kwargs)
+            return "ok"
+
+        wrapped = make_fast_mode_wrapper(fake_create, "claude-code-foo")
+
+        with patch(
+            "code_muse.config.get_all_model_settings",
+            return_value={FAST_SETTING_KEY: False},
+        ):
+            await wrapped(model="foo", messages=[])
+
+        assert "speed" not in captured
+
+    @pytest.mark.asyncio
+    async def test_does_not_clobber_explicit_speed(self):
+        captured = {}
+
+        async def fake_create(**kwargs):
+            captured.update(kwargs)
+            return "ok"
+
+        wrapped = make_fast_mode_wrapper(fake_create, "claude-code-foo")
+
+        with patch(
+            "code_muse.config.get_all_model_settings",
+            return_value={FAST_SETTING_KEY: True},
+        ):
+            await wrapped(model="foo", messages=[], speed="slow")
+
+        # Caller-provided speed wins
+        assert captured.get("speed") == "slow"
+
+
+class TestIsFastModeEnabledRealConfigRoundTrip:
+    """End-to-end regression test: set_model_setting -> is_fast_mode_enabled.
+
+    This would have caught the original bug where get_effective_model_settings
+    silently filtered the 'fast' key out because it wasn't in the core
+    supported_settings allowlist for claude-* models.
+    """
+
+    def test_real_roundtrip_with_claude_code_model(self, tmp_path, monkeypatch):
+        # Point config at a throwaway file so we don't clobber user config.
+        cfg = tmp_path / "muse.cfg"
+        monkeypatch.setattr("code_muse.config.CONFIG_FILE", str(cfg))
+
+        from code_muse.config import set_model_setting
+
+        model = "claude-code-claude-opus-4-7"
+
+        # Initially off.
+        assert is_fast_mode_enabled(model) is False
+
+        # Flip it on via the same path /claude-code-fast uses.
+        set_model_setting(model, FAST_SETTING_KEY, "true")
+        assert is_fast_mode_enabled(model) is True, (
+            "Regression: fast setting written to config but read as False. "
+            "Likely the reader is filtering through the supported_settings allowlist."
+        )
+
+        # Flip it back off.
+        set_model_setting(model, FAST_SETTING_KEY, "false")
+        assert is_fast_mode_enabled(model) is False

code_muse/plugins/claude_code_oauth/test_plugin.py

@@ -0,0 +1,283 @@
+#!/usr/bin/env python3
+"""Manual sanity checks for the Claude Code OAuth plugin."""
+
+import os
+import sys
+from pathlib import Path
+
+# Ensure project root on path
+PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent.parent
+sys.path.insert(0, str(PROJECT_ROOT))
+
+# Switch to project root for predictable relative paths
+os.chdir(PROJECT_ROOT)
+
+
+def test_plugin_imports() -> bool:
+    """Verify the plugin modules import correctly."""
+    print("\n=== Testing Plugin Imports ===")
+
+    try:
+        from code_muse.plugins.claude_code_oauth.config import (
+            CLAUDE_CODE_OAUTH_CONFIG,
+            get_token_storage_path,
+        )
+
+        print("✅ Config import successful")
+        print(f"✅ Token storage path: {get_token_storage_path()}")
+        print(f"✅ Known auth URL: {CLAUDE_CODE_OAUTH_CONFIG['auth_url']}")
+    except Exception as exc:  # pragma: no cover - manual harness
+        print(f"❌ Config import failed: {exc}")
+        return False
+
+    try:
+        from code_muse.plugins.claude_code_oauth.utils import (
+            add_models_to_extra_config,
+            build_authorization_url,
+            exchange_code_for_tokens,
+            fetch_claude_code_models,
+            load_claude_models,
+            load_stored_tokens,
+            parse_authorization_code,
+            prepare_oauth_context,
+            remove_claude_code_models,
+            save_claude_models,
+            save_tokens,
+        )
+
+        _ = (
+            add_models_to_extra_config,
+            build_authorization_url,
+            exchange_code_for_tokens,
+            fetch_claude_code_models,
+            load_claude_models,
+            load_stored_tokens,
+            parse_authorization_code,
+            prepare_oauth_context,
+            remove_claude_code_models,
+            save_claude_models,
+            save_tokens,
+        )
+        print("✅ Utils import successful")
+    except Exception as exc:  # pragma: no cover - manual harness
+        print(f"❌ Utils import failed: {exc}")
+        return False
+
+    try:
+        from code_muse.plugins.claude_code_oauth.register_callbacks import (
+            _custom_help,
+            _handle_custom_command,
+        )
+
+        commands = _custom_help()
+        print("✅ Callback registration import successful")
+        for name, description in commands:
+            print(f"  /{name} - {description}")
+        # Ensure handler callable exists
+        _ = _handle_custom_command
+    except Exception as exc:  # pragma: no cover - manual harness
+        print(f"❌ Callback import failed: {exc}")
+        return False
+
+    return True
+
+
+def test_oauth_helpers() -> bool:
+    """Exercise helper functions without performing network requests."""
+    print("\n=== Testing OAuth Helper Functions ===")
+
+    try:
+        from urllib.parse import parse_qs, urlparse
+
+        from code_muse.plugins.claude_code_oauth.utils import (
+            assign_redirect_uri,
+            build_authorization_url,
+            parse_authorization_code,
+            prepare_oauth_context,
+        )
+
+        context = prepare_oauth_context()
+        assert context.state, "Expected non-empty OAuth state"
+        assert context.code_verifier, "Expected PKCE code verifier"
+        assert context.code_challenge, "Expected PKCE code challenge"
+
+        assign_redirect_uri(context, 8765)
+        auth_url = build_authorization_url(context)
+        parsed = urlparse(auth_url)
+        params = parse_qs(parsed.query)
+        print(f"✅ Authorization URL: {auth_url}")
+        assert parsed.scheme == "https", "Authorization URL must use https"
+        assert params.get("client_id", [None])[0], "client_id missing"
+        assert params.get("code_challenge_method", [None])[0] == "S256"
+        assert params.get("state", [None])[0] == context.state
+        assert params.get("code_challenge", [None])[0] == context.code_challenge
+
+        sample_code = f"MYCODE#{context.state}"
+        parsed_code, parsed_state = parse_authorization_code(sample_code)
+        assert parsed_code == "MYCODE", "Code parsing failed"
+        assert parsed_state == context.state, "State parsing failed"
+        print("✅ parse_authorization_code handled state suffix correctly")
+
+        parsed_code, parsed_state = parse_authorization_code("SINGLECODE")
+        assert parsed_code == "SINGLECODE" and parsed_state is None
+        print("✅ parse_authorization_code handled bare code correctly")
+
+        return True
+
+    except AssertionError as exc:
+        print(f"❌ Assertion failed: {exc}")
+        return False
+    except Exception as exc:  # pragma: no cover - manual harness
+        print(f"❌ OAuth helper test crashed: {exc}")
+        import traceback
+
+        traceback.print_exc()
+        return False
+
+
+def test_file_operations() -> bool:
+    """Ensure token/model storage helpers behave sanely."""
+    print("\n=== Testing File Operations ===")
+
+    try:
+        from code_muse.plugins.claude_code_oauth.config import (
+            get_claude_models_path,
+            get_token_storage_path,
+        )
+        from code_muse.plugins.claude_code_oauth.utils import (
+            load_claude_models,
+            load_stored_tokens,
+        )
+
+        tokens = load_stored_tokens()
+        print(f"✅ Token load result: {'present' if tokens else 'none'}")
+
+        models = load_claude_models()
+        print(f"✅ Loaded {len(models)} Claude models")
+        for name, config in models.items():
+            print(f"  - {name}: {config.get('type', 'unknown type')}")
+
+        token_path = get_token_storage_path()
+        models_path = get_claude_models_path()
+        token_path.parent.mkdir(parents=True, exist_ok=True)
+        models_path.parent.mkdir(parents=True, exist_ok=True)
+        print(f"✅ Token path: {token_path}")
+        print(f"✅ Models path: {models_path}")
+
+        return True
+
+    except Exception as exc:  # pragma: no cover - manual harness
+        print(f"❌ File operations test failed: {exc}")
+        import traceback
+
+        traceback.print_exc()
+        return False
+
+
+def test_command_handlers() -> bool:
+    """Smoke-test command handler routing without simulating authentication."""
+    print("\n=== Testing Command Handlers ===")
+
+    from code_muse.plugins.claude_code_oauth.register_callbacks import (
+        _handle_custom_command,
+    )
+
+    unknown = _handle_custom_command("/bogus", "bogus")
+    print(f"✅ Unknown command returned: {unknown}")
+
+    partial = _handle_custom_command("/claude-code", "claude-code")
+    print(f"✅ Partial command returned: {partial}")
+
+    # Do not invoke the real auth command here because it prompts for input.
+    return True
+
+
+def test_configuration() -> bool:
+    """Validate configuration keys and basic formats."""
+    print("\n=== Testing Configuration ===")
+
+    try:
+        from code_muse.plugins.claude_code_oauth.config import CLAUDE_CODE_OAUTH_CONFIG
+
+        required_keys = [
+            "auth_url",
+            "token_url",
+            "api_base_url",
+            "client_id",
+            "scope",
+            "redirect_host",
+            "redirect_path",
+            "callback_port_range",
+            "callback_timeout",
+            "token_storage",
+            "prefix",
+            "default_context_length",
+            "api_key_env_var",
+        ]
+
+        missing = [key for key in required_keys if key not in CLAUDE_CODE_OAUTH_CONFIG]
+        if missing:
+            print(f"❌ Missing configuration keys: {missing}")
+            return False
+
+        for key in required_keys:
+            value = CLAUDE_CODE_OAUTH_CONFIG[key]
+            print(f"✅ {key}: {value}")
+
+        for url_key in ["auth_url", "token_url", "api_base_url"]:
+            url = CLAUDE_CODE_OAUTH_CONFIG[url_key]
+            if not str(url).startswith("https://"):
+                print(f"❌ URL must use HTTPS: {url_key} -> {url}")
+                return False
+            print(f"✅ {url_key} uses HTTPS")
+
+        return True
+
+    except Exception as exc:  # pragma: no cover - manual harness
+        print(f"❌ Configuration test crashed: {exc}")
+        import traceback
+
+        traceback.print_exc()
+        return False
+
+
+def main() -> bool:
+    """Run all manual checks."""
+    print("Claude Code OAuth Plugin Test Suite")
+    print("=" * 40)
+
+    tests = [
+        test_plugin_imports,
+        test_oauth_helpers,
+        test_file_operations,
+        test_command_handlers,
+        test_configuration,
+    ]
+
+    passed = 0
+    for test in tests:
+        try:
+            if test():
+                passed += 1
+            else:
+                print("\n❌ Test failed")
+        except Exception as exc:  # pragma: no cover - manual harness
+            print(f"\n❌ Test crashed: {exc}")
+
+    print("\n=== Test Results ===")
+    print(f"Passed: {passed}/{len(tests)}")
+
+    if passed == len(tests):
+        print("✅ All sanity checks passed!")
+        print("Next steps:")
+        print("1. Restart Muse if it was running")
+        print("2. Run /claude-code-auth")
+        print("3. Paste the Claude Console authorization code when prompted")
+        return True
+
+    print("❌ Some checks failed. Investigate before using the plugin.")
+    return False
+
+
+if __name__ == "__main__":
+    sys.exit(0 if main() else 1)

code_muse/plugins/claude_code_oauth/token_refresh_heartbeat.py

@@ -0,0 +1,237 @@
+"""Token refresh heartbeat for long-running Claude Code OAuth sessions.
+
+This module provides a background task that periodically checks and refreshes
+Claude Code OAuth tokens during long-running agentic operations. This ensures
+that tokens don't expire during extended streaming responses or tool processing.
+
+Usage:
+    async with token_refresh_heartbeat_context():
+        # Long running agent operation
+        await agent.run(...)
+"""
+
+import asyncio
+import logging
+import time
+from contextlib import asynccontextmanager, suppress
+
+logger = logging.getLogger(__name__)
+
+# Heartbeat interval in seconds - check token every 2 minutes
+# This is frequent enough to catch expiring tokens before they cause issues
+# but not so frequent as to spam the token endpoint
+HEARTBEAT_INTERVAL_SECONDS = 120
+
+# Minimum time between refresh attempts to avoid hammering the endpoint
+MIN_REFRESH_INTERVAL_SECONDS = 60
+
+# Global tracking of last refresh time to coordinate across heartbeats
+_last_refresh_time: float = 0.0
+
+
+class TokenRefreshHeartbeat:
+    """Background task that periodically refreshes Claude Code OAuth tokens.
+
+    This runs as an asyncio task during agent operations and checks if the
+    token needs refreshing at regular intervals.
+    """
+
+    def __init__(
+        self,
+        interval: float = HEARTBEAT_INTERVAL_SECONDS,
+        min_refresh_interval: float = MIN_REFRESH_INTERVAL_SECONDS,
+    ):
+        self._interval = interval
+        self._min_refresh_interval = min_refresh_interval
+        self._task: asyncio.Task | None = None
+        self._stop_event = asyncio.Event()
+        self._lock = asyncio.Lock()
+        self._refresh_count = 0
+
+    async def start(self) -> None:
+        """Start the heartbeat background task."""
+        if self._task is not None:
+            logger.debug("Heartbeat already running")
+            return
+
+        self._stop_event.clear()
+        self._task = asyncio.create_task(self._heartbeat_loop())
+        logger.debug("Token refresh heartbeat started")
+
+    async def stop(self) -> None:
+        """Stop the heartbeat background task."""
+        if self._task is None:
+            return
+
+        self._stop_event.set()
+        self._task.cancel()
+
+        with suppress(asyncio.CancelledError):
+            await self._task
+
+        self._task = None
+        logger.debug(
+            "Token refresh heartbeat stopped (refreshed %d times)",
+            self._refresh_count,
+        )
+
+    async def _heartbeat_loop(self) -> None:
+        """Main heartbeat loop that periodically checks token status."""
+        global _last_refresh_time
+
+        while not self._stop_event.is_set():
+            try:
+                # Wait for the interval or until stopped
+                try:
+                    await asyncio.wait_for(
+                        self._stop_event.wait(), timeout=self._interval
+                    )
+                    # If we got here, stop event was set
+                    break
+                except TimeoutError:
+                    # Normal timeout - time to check token
+                    pass
+
+                # Check if we should attempt refresh
+                async with self._lock:
+                    now = time.time()
+                    if now - _last_refresh_time < self._min_refresh_interval:
+                        logger.debug(
+                            "Skipping refresh - last refresh was %.1f seconds ago",
+                            now - _last_refresh_time,
+                        )
+                        continue
+
+                    # Attempt the refresh
+                    refreshed = await self._attempt_refresh()
+                    if refreshed:
+                        _last_refresh_time = now
+                        self._refresh_count += 1
+
+            except asyncio.CancelledError:
+                break
+            except Exception as exc:
+                logger.debug("Error in heartbeat loop: %s", exc)
+                # Continue running - don't let errors kill the heartbeat
+                await asyncio.sleep(5)  # Brief pause before retrying
+
+    async def _attempt_refresh(self) -> bool:
+        """Attempt to refresh the token if needed.
+
+        Returns:
+            True if a refresh was performed, False otherwise.
+        """
+        try:
+            # Import here to avoid circular imports
+            from .utils import (
+                is_token_expired,
+                load_stored_tokens,
+                refresh_access_token,
+            )
+
+            tokens = await asyncio.to_thread(load_stored_tokens)
+            if not tokens:
+                logger.debug("No stored tokens found")
+                return False
+
+            if not is_token_expired(tokens):
+                logger.debug("Token not yet expired, skipping refresh")
+                return False
+
+            # Token is expiring soon, refresh it
+            logger.info("Heartbeat: Token expiring soon, refreshing proactively")
+            refreshed_token = await asyncio.to_thread(refresh_access_token, force=False)
+
+            if refreshed_token:
+                logger.info("Heartbeat: Successfully refreshed token")
+                return True
+            else:
+                logger.warning("Heartbeat: Token refresh returned None")
+                return False
+
+        except Exception as exc:
+            logger.error("Heartbeat: Error during token refresh: %s", exc)
+            return False
+
+    @property
+    def refresh_count(self) -> int:
+        """Get the number of successful refreshes performed by this heartbeat."""
+        return self._refresh_count
+
+    @property
+    def is_running(self) -> bool:
+        """Check if the heartbeat is currently running."""
+        return self._task is not None and not self._task.done()
+
+
+# Global heartbeat instance for the current session
+_current_heartbeat: TokenRefreshHeartbeat | None = None
+
+
+@asynccontextmanager
+async def token_refresh_heartbeat_context(
+    interval: float = HEARTBEAT_INTERVAL_SECONDS,
+):
+    """Context manager that runs token refresh heartbeat during its scope.
+
+    Use this around long-running agent operations to ensure tokens stay fresh.
+
+    Args:
+        interval: Seconds between heartbeat checks. Default is 2 minutes.
+
+    Example:
+        async with token_refresh_heartbeat_context():
+            result = await agent.run(prompt)
+    """
+    global _current_heartbeat
+
+    heartbeat = TokenRefreshHeartbeat(interval=interval)
+
+    try:
+        await heartbeat.start()
+        _current_heartbeat = heartbeat
+        yield heartbeat
+    finally:
+        await heartbeat.stop()
+        _current_heartbeat = None
+
+
+def is_heartbeat_running() -> bool:
+    """Check if a token refresh heartbeat is currently active."""
+    return _current_heartbeat is not None and _current_heartbeat.is_running
+
+
+def get_current_heartbeat() -> TokenRefreshHeartbeat | None:
+    """Get the currently running heartbeat instance, if any."""
+    return _current_heartbeat
+
+
+async def force_token_refresh() -> bool:
+    """Force an immediate token refresh.
+
+    This can be called from anywhere to trigger a token refresh,
+    regardless of whether a heartbeat is running.
+
+    Returns:
+        True if refresh was successful, False otherwise.
+    """
+    global _last_refresh_time
+
+    try:
+        from .utils import refresh_access_token
+
+        logger.info("Forcing token refresh")
+        # TODO: PEP 734 async bridge — convert refresh_access_token to async
+        refreshed_token = await asyncio.to_thread(refresh_access_token, force=True)
+
+        if refreshed_token:
+            _last_refresh_time = time.time()
+            logger.info("Force refresh successful")
+            return True
+        else:
+            logger.warning("Force refresh returned None")
+            return False
+
+    except Exception as exc:
+        logger.error("Force refresh error: %s", exc)
+        return False