code-muse 0.0.1__py3-none-any.whl

code_muse/plugins/ollama_setup/register_callbacks.py

@@ -0,0 +1,410 @@
+"""Ollama cloud model setup — /ollama-setup command.
+
+Pulls an Ollama `:cloud` model and registers it in extra_models.json so the
+model is immediately available for use in Muse.
+
+Cloud models supported (the Ollama "Recommended Models" cloud tier):
+    kimi-k2.6:cloud, kimi-k2.5:cloud, glm-5:cloud, glm-5.1:cloud, minimax-m2.7:cloud, qwen3.5:cloud
+"""
+
+import json
+import logging
+import shutil
+import subprocess
+from pathlib import Path
+from typing import Any
+
+from code_muse.callbacks import register_callback
+from code_muse.config import EXTRA_MODELS_FILE
+from code_muse.messaging import emit_error, emit_info, emit_success, emit_warning
+
+logger = logging.getLogger(__name__)
+
+# ── Cloud model catalogue ───────────────────────────────────────────────────
+# The `:cloud` models from the Ollama recommended list.
+# Each entry maps the ollama tag → extra_models.json config metadata.
+
+CLOUD_MODELS: dict[str, dict[str, Any]] = {
+    "kimi-k2.6:cloud": {
+        "context_length": 262144,  # 256k
+        "description": "Moonshot Kimi K2.6 (cloud)",
+    },
+    "kimi-k2.5:cloud": {
+        "context_length": 262144,  # 256k
+        "description": "Moonshot Kimi K2.5 (cloud)",
+    },
+    "glm-5:cloud": {
+        "context_length": 204800,  # 200k
+        "description": "ZhipuAI GLM-5 (cloud)",
+    },
+    "glm-5.1:cloud": {
+        "context_length": 204800,  # 200k
+        "description": "ZhipuAI GLM-5.1 (cloud)",
+    },
+    "minimax-m2.7:cloud": {
+        "context_length": 199680,  # 195k
+        "description": "MiniMax M2.7 (cloud)",
+    },
+    "qwen3.5:cloud": {
+        "context_length": 204800,  # 200k
+        "description": "Alibaba Qwen 3.5 (cloud)",
+    },
+}
+
+OLLAMA_ENDPOINT = "http://localhost:11434/v1"
+OLLAMA_API_KEY = "ollama"  # Ollama's local API doesn't need a real key
+
+
+# ── Helpers ─────────────────────────────────────────────────────────────────
+
+
+def _model_key(model_tag: str) -> str:
+    """Derive a unique extra_models.json key from an ollama tag.
+
+    e.g. ``glm-5:cloud`` → ``ollama-glm-5-cloud``
+    """
+    return "ollama-" + model_tag.replace(":", "-").replace("/", "-")
+
+
+def _ollama_available() -> bool:
+    """Return True if the ``ollama`` CLI is on PATH."""
+    return shutil.which("ollama") is not None
+
+
+def _pull_model(model_tag: str) -> bool:
+    """Run ``ollama pull <model_tag>``, streaming output to the terminal.
+
+    Returns True on success, False on failure.
+    """
+    emit_info(f"[Run] Pulling {model_tag} via ollama …")
+    try:
+        result = subprocess.run(
+            ["ollama", "pull", model_tag],
+            capture_output=True,
+            text=True,
+            timeout=600,
+        )
+        if result.returncode != 0:
+            stderr = result.stderr.strip()
+            emit_error(f"ollama pull failed (exit {result.returncode}): {stderr}")
+            return False
+
+        stdout = result.stdout.strip()
+        if stdout:
+            emit_info(stdout)
+        return True
+    except subprocess.TimeoutExpired:
+        emit_error("ollama pull timed out after 10 minutes")
+        return False
+    except Exception as exc:
+        emit_error(f"Failed to run ollama pull: {exc}")
+        return False
+
+
+def _register_model(model_tag: str) -> bool:
+    """Write (or update) the model entry in extra_models.json.
+
+    Uses atomic-write via a temp file, same pattern as add_model_menu.py.
+    Returns True on success.
+    """
+    meta = CLOUD_MODELS[model_tag]
+    key = _model_key(model_tag)
+
+    extra_path = Path(EXTRA_MODELS_FILE)
+    extra_models: dict[str, Any] = {}
+
+    if extra_path.exists():
+        try:
+            with open(extra_path, encoding="utf-8") as fh:
+                extra_models = json.load(fh)
+                if not isinstance(extra_models, dict):
+                    emit_error("extra_models.json must be a dict, not a list")
+                    return False
+        except json.JSONDecodeError as exc:
+            emit_error(f"Corrupt extra_models.json: {exc}")
+            return False
+
+    if key in extra_models:
+        emit_info(f"Model {key} already registered — updating entry")
+
+    extra_models[key] = {
+        "type": "custom_openai",
+        "name": model_tag,
+        "custom_endpoint": {
+            "url": OLLAMA_ENDPOINT,
+            "api_key": OLLAMA_API_KEY,
+        },
+        "context_length": meta["context_length"],
+        "supported_settings": ["temperature", "top_p"],
+    }
+
+    extra_path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = extra_path.with_suffix(".tmp")
+    try:
+        with open(tmp, "w", encoding="utf-8") as fh:
+            json.dump(extra_models, fh, indent=4, ensure_ascii=False)
+        tmp.replace(extra_path)
+    except Exception as exc:
+        emit_error(f"Failed to write extra_models.json: {exc}")
+        return False
+
+    emit_success(f"✅ Registered {key} in extra_models.json")
+    return True
+
+
+def _test_model_auth(model_tag: str) -> tuple[bool, str]:
+    """Test if the model requires authentication.
+
+    Returns (authorized: bool, message: str).
+    If unauthorized, message contains guidance for the user.
+    """
+    import json as json_mod
+    import urllib.request
+
+    test_payload = json_mod.dumps(
+        {
+            "model": model_tag,
+            "messages": [{"role": "user", "content": "hi"}],
+            "stream": False,
+        }
+    ).encode("utf-8")
+
+    req = urllib.request.Request(
+        f"{OLLAMA_ENDPOINT}/chat/completions",
+        data=test_payload,
+        headers={
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {OLLAMA_API_KEY}",
+        },
+        method="POST",
+    )
+
+    try:
+        with urllib.request.urlopen(req, timeout=10):
+            # 200 = good to go
+            return (True, "")
+    except urllib.error.HTTPError as e:
+        body = e.read().decode("utf-8", errors="replace")
+        # Check for auth errors
+        if e.code == 401 or "unauthorized" in body.lower():
+            return (False, _auth_help_message(model_tag))
+        if e.code == 500 and (
+            "internal service error" in body.lower() or "unauthorized" in body.lower()
+        ):
+            # Cloud models often return 500 instead of 401
+            return (False, _auth_help_message(model_tag))
+        return (False, f"HTTP {e.code}: {body[:200]}")
+    except Exception as e:
+        return (False, f"Connection test failed: {e}")
+
+
+def _auth_help_message(model_tag: str) -> str:
+    """Generate a helpful auth guidance message."""
+    return f"""🔐 Authentication Required for {model_tag}
+
+This is a cloud-hosted model (756B+ parameters) that requires Ollama Cloud access.
+
+To use this model:
+1. Run: ollama login
+2. Visit https://ollama.com to create an account
+3. Subscribe to Ollama Cloud (cloud models aren't free)
+4. Accept the terms for {model_tag.split(":")[0]}
+
+After logging in, you can use the model with:
+/model {_model_key(model_tag)}
+"""
+
+
+def _run_ollama_login() -> bool:
+    """Run ``ollama login`` interactively in the user's terminal.
+
+    This will open a browser for OAuth flow. Returns True if login succeeded.
+    """
+    emit_info(
+        "🔐 Launching 'ollama login' — this will open your browser for authentication..."
+    )
+    emit_info("   (Press Ctrl+C here if you want to cancel and do it manually later)")
+    emit_info("")
+
+    try:
+        # Run ollama login with stdout/stderr connected to terminal
+        # so the user can see the browser-open message and any instructions
+        result = subprocess.run(
+            ["ollama", "login"],
+            capture_output=False,  # Let it flow to terminal
+            text=True,
+            timeout=120,  # 2 minutes — OAuth can take a while
+        )
+
+        if result.returncode == 0:
+            emit_success("✅ Login completed!")
+            return True
+        else:
+            emit_error(f"ollama login exited with code {result.returncode}")
+            return False
+    except subprocess.TimeoutExpired:
+        emit_error("ollama login timed out after 2 minutes")
+        return False
+    except Exception as exc:
+        emit_error(f"Failed to run ollama login: {exc}")
+        return False
+
+
+def _start_ollama_serve() -> subprocess.Popen | None:
+    """Start ``ollama serve`` as a background subprocess if not already running.
+
+    Returns the Popen handle, or None if already running / on failure.
+    """
+    # Quick liveness check — hit the version endpoint
+    try:
+        result = subprocess.run(
+            ["ollama", "list"],
+            capture_output=True,
+            text=True,
+            timeout=5,
+        )
+        if result.returncode == 0:
+            emit_info("🟢 Ollama is already running")
+            return None
+    except Exception:
+        pass
+
+    emit_info("🚀 Starting ollama serve in the background …")
+    try:
+        proc = subprocess.Popen(
+            ["ollama", "serve"],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        )
+        emit_success(f"ollama serve started (pid {proc.pid})")
+        return proc
+    except Exception as exc:
+        emit_error(f"Could not start ollama serve: {exc}")
+        return None
+
+
+# ── Command handler ─────────────────────────────────────────────────────────
+
+
+def _handle_ollama_setup(command: str, name: str) -> Any:
+    """Handle ``/ollama-setup [model]``."""
+    if name != "ollama-setup":
+        return None  # Not our command
+
+    parts = command.split(maxsplit=1)
+    model_tag = parts[1].strip() if len(parts) > 1 else ""
+
+    # No argument → show available models
+    if not model_tag:
+        lines = ["🦙 Available Ollama cloud models:\n"]
+        for tag, meta in CLOUD_MODELS.items():
+            lines.append(f"  • {tag:25s} — {meta['description']}")
+        lines.append("\nUsage: /ollama-setup <model>")
+        emit_info("\n".join(lines))
+        return True
+
+    # Fuzzy-ish match: allow partial names like "glm" → "glm-5:cloud"
+    matched = _resolve_model(model_tag)
+    if matched is None:
+        emit_error(f"Unknown model '{model_tag}'. Available: {', '.join(CLOUD_MODELS)}")
+        return True
+
+    # Pre-flight: is ollama installed?
+    if not _ollama_available():
+        emit_error(
+            "ollama CLI not found on PATH. "
+            "Install it from https://ollama.com and try again."
+        )
+        return True
+
+    # 1. Ensure ollama serve is up
+    _start_ollama_serve()
+
+    # 2. Pull the model
+    if not _pull_model(matched):
+        return True
+
+    # 3. Register in extra_models.json
+    _register_model(matched)
+
+    # 4. Test auth / accessibility
+    emit_info(f"🔍 Testing {matched} accessibility...")
+    authorized, msg = _test_model_auth(matched)
+
+    if authorized:
+        emit_success(
+            f"🎉 Done! Switch to the model with:  /model {_model_key(matched)}"
+        )
+    else:
+        emit_warning(msg)
+        emit_info("")
+        emit_info(
+            "💡 Want me to run 'ollama login' for you? (This will open your browser)"
+        )
+        emit_info(
+            "   Type 'yes' to proceed, or anything else to skip and do it manually later."
+        )
+
+        # We can't do interactive input here (custom commands shouldn't block)
+        # So we just run it — worst case it fails and they do it manually
+        login_ok = _run_ollama_login()
+
+        if login_ok:
+            # Re-test auth after login
+            emit_info("🔍 Re-testing accessibility after login...")
+            authorized, msg = _test_model_auth(matched)
+
+            if authorized:
+                emit_success(
+                    f"🎉 All set! Switch to the model with:  /model {_model_key(matched)}"
+                )
+            else:
+                emit_warning("Still not authorized after login.")
+                emit_info(
+                    f"📋 Model registered as '{_model_key(matched)}'. "
+                    f"You may need to subscribe to Ollama Cloud at https://ollama.com, "
+                    f"then use /model {_model_key(matched)}"
+                )
+        else:
+            emit_info(
+                f"📋 Model registered as '{_model_key(matched)}' but needs auth. "
+                f"Run 'ollama login' manually, then use /model {_model_key(matched)}"
+            )
+
+    return True
+
+
+def _resolve_model(user_input: str) -> str | None:
+    """Resolve user input to an exact cloud model tag.
+
+    Supports exact match and unambiguous prefix matching.
+    """
+    lowered = user_input.lower().strip()
+
+    # Exact match first
+    for tag in CLOUD_MODELS:
+        if tag.lower() == lowered:
+            return tag
+
+    # Prefix / substring match
+    candidates = [t for t in CLOUD_MODELS if lowered in t.lower()]
+    if len(candidates) == 1:
+        return candidates[0]
+
+    return None
+
+
+# ── Help entry ──────────────────────────────────────────────────────────────
+
+
+def _custom_help():
+    model_names = ", ".join(CLOUD_MODELS)
+    return [
+        ("ollama-setup", f"Pull & register an Ollama cloud model ({model_names})"),
+    ]
+
+
+# ── Register ────────────────────────────────────────────────────────────────
+
+register_callback("custom_command", _handle_ollama_setup)
+register_callback("custom_command_help", _custom_help)

code_muse/plugins/plan_command/register_callbacks.py

@@ -0,0 +1,206 @@
+"""Plan Command Plugin — /plan slash command.
+
+Invokes the planning-agent to break down a coding task into
+clear, actionable steps and displays the plan to the user.
+"""
+
+import logging
+from contextlib import AsyncExitStack
+from dataclasses import dataclass
+from functools import partial
+
+from code_muse.callbacks import on_agent_run_context, register_callback
+from code_muse.messaging import emit_info, emit_success, emit_warning
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Lightweight standalone agent invocation (mirrors invoke_agent tool logic)
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class _AgentResult:
+    """Simple container for an invoked agent's response."""
+
+    response_text: str | None = None
+    error: str | None = None
+
+
+async def _invoke_agent(agent_name: str, prompt: str) -> _AgentResult:
+    """Invoke a named agent with a prompt and return its text response.
+
+    This is a lightweight standalone wrapper that mirrors the core
+    ``invoke_agent`` tool logic but doesn't require a Pydantic-AI
+    ``RunContext`` — making it callable from plugin callbacks.
+    """
+    from pydantic_ai import Agent, UsageLimits
+
+    from code_muse.agents._compaction import make_history_processor
+    from code_muse.agents.agent_manager import load_agent
+    from code_muse.agents.subagent_stream_handler import subagent_stream_handler
+    from code_muse.config import get_message_limit
+    from code_muse.model_factory import ModelFactory, make_model_settings
+    from code_muse.model_utils import prepare_prompt_for_model
+    from code_muse.tools import register_tools_for_agent
+    from code_muse.tools.common import generate_group_id
+    from code_muse.tools.subagent_context import subagent_context
+
+    group_id = generate_group_id("invoke_agent", agent_name)
+
+    try:
+        agent_config = load_agent(agent_name)
+
+        model_name = agent_config.get_model_name()
+        models_config = ModelFactory.load_config()
+
+        if model_name not in models_config:
+            return _AgentResult(
+                error=f"Model '{model_name}' not found in configuration"
+            )
+
+        model = ModelFactory.get_model(model_name, models_config)
+
+        instructions = agent_config.get_full_system_prompt()
+
+        # Add AGENTS.md content
+        from code_muse.agents._builder import load_muse_rules
+
+        agent_rules = load_muse_rules()
+        if agent_rules:
+            instructions += f"\n\n{agent_rules}"
+
+        # Apply plugin prompt additions
+        from code_muse import callbacks as cb
+
+        prompt_additions = cb.on_load_prompt()
+        if prompt_additions:
+            instructions += "\n" + "\n".join(prompt_additions)
+
+        # Prepare prompt for model (handles claude-code models etc.)
+        prepared = prepare_prompt_for_model(
+            model_name,
+            instructions,
+            prompt,
+            prepend_system_to_user=True,
+        )
+        instructions = prepared.instructions
+        prompt = prepared.user_prompt
+
+        model_settings = make_model_settings(model_name)
+
+        temp_agent = Agent(
+            model=model,
+            instructions=instructions,
+            output_type=str,
+            retries=3,
+            toolsets=[],
+            history_processors=[make_history_processor(agent_config)],
+            model_settings=model_settings,
+        )
+
+        # Register the tools the agent needs
+        agent_tools = agent_config.get_available_tools()
+        register_tools_for_agent(temp_agent, agent_tools, model_name=model_name)
+
+        # Use subagent_stream_handler for clean output
+        stream_handler = partial(
+            subagent_stream_handler, session_id=f"{agent_name}-plan-session"
+        )
+
+        with subagent_context(agent_name):
+            run_ctxs = on_agent_run_context(agent_config, temp_agent, group_id)
+            async with AsyncExitStack() as stack:
+                for cm in run_ctxs:
+                    await stack.enter_async_context(cm)
+                result = await temp_agent.run(
+                    prompt,
+                    message_history=[],
+                    usage_limits=UsageLimits(request_limit=get_message_limit()),
+                    event_stream_handler=stream_handler,
+                )
+
+        return _AgentResult(response_text=result.output)
+
+    except Exception as exc:
+        logger.exception("invoke_agent failed for %s", agent_name)
+        return _AgentResult(error=str(exc))
+
+
+# ---------------------------------------------------------------------------
+# Slash-command help
+# ---------------------------------------------------------------------------
+
+
+def _custom_help() -> list[tuple[str, str]]:
+    """Provide help entry for /help display."""
+    return [
+        ("plan", "Invoke the Planning Agent to create an execution roadmap"),
+    ]
+
+
+# ---------------------------------------------------------------------------
+# Slash-command handler
+# ---------------------------------------------------------------------------
+
+
+async def _handle_custom_command(command: str, name: str) -> bool | None:
+    """Handle the /plan slash command.
+
+    Usage:
+        /plan <description of what to plan>
+
+    Invokes the planning-agent with the user's request and displays
+    the resulting plan.
+
+    Returns:
+        True if handled.
+        None if the command name doesn't match.
+    """
+    if name != "plan":
+        return None
+
+    # Extract the planning request (everything after "/plan")
+    parts = command.split(maxsplit=1)
+    plan_request = parts[1].strip() if len(parts) > 1 else ""
+
+    if not plan_request:
+        emit_warning(
+            "Usage: /plan <description of what to plan>\n"
+            "Example: /plan Add user authentication with OAuth2"
+        )
+        return True
+
+    emit_info("📋 Planning Agent is creating your roadmap…")
+
+    result = await _invoke_agent(
+        agent_name="planning-agent",
+        prompt=(
+            f"Create a detailed execution plan for the following task:\n\n"
+            f"{plan_request}\n\n"
+            f"Include: project structure analysis, dependencies, "
+            f"execution steps, and validation strategy."
+        ),
+    )
+
+    # Display the plan
+    if result.response_text:
+        emit_success("📋 Planning Agent — Execution Roadmap")
+        emit_info(result.response_text)
+    elif result.error:
+        emit_warning(f"Planning Agent error: {result.error}")
+    else:
+        emit_warning("Planning Agent returned no output.")
+
+    return True  # Handled — no additional model invocation needed
+
+
+# ---------------------------------------------------------------------------
+# Register callbacks
+# ---------------------------------------------------------------------------
+
+register_callback("custom_command_help", _custom_help)
+register_callback("custom_command", _handle_custom_command)
+
+logger.debug("Plan Command plugin callbacks registered")

code_muse/plugins/plan_mode/__init__.py

@@ -0,0 +1,37 @@
+"""Plan Mode plugin for Muse.
+
+Provides tools and hooks to enter/exit a planning-only mode where
+destructive tools (write, replace, delete, shell) are blocked while
+research tools (read, list, grep, ask) remain available.
+"""
+
+from code_muse.plugins.plan_mode.mode_cycling import cycle_mode
+from code_muse.plugins.plan_mode.plan_generation import generate_plan_md, save_plan
+from code_muse.plugins.plan_mode.plan_hooks import plan_mode_pre_tool_call_hook
+from code_muse.plugins.plan_mode.plan_mode_tools import (
+    PlanModeState,
+    get_current_mode,
+    get_plan_goal,
+    register_approve_plan,
+    register_enter_plan_mode,
+    register_exit_plan_mode,
+    register_get_plan_mode,
+    register_open_plan_in_editor,
+    set_plan_mode,
+)
+
+__all__ = [
+    "PlanModeState",
+    "get_current_mode",
+    "get_plan_goal",
+    "set_plan_mode",
+    "register_enter_plan_mode",
+    "register_exit_plan_mode",
+    "register_get_plan_mode",
+    "register_approve_plan",
+    "register_open_plan_in_editor",
+    "generate_plan_md",
+    "save_plan",
+    "cycle_mode",
+    "plan_mode_pre_tool_call_hook",
+]

code_muse/plugins/plan_mode/mode_cycling.py

@@ -0,0 +1,40 @@
+"""Mode cycling logic and UI emission.
+
+Provides a ``cycle_mode()`` function that rotates through the three
+plan-mode states.  A ``/mode`` slash command exposes this in the UI.
+
+TODO: Shift+Tab key-binding integration requires future prompt_toolkit
+hooking (not implemented here to avoid core changes).
+"""
+
+from code_muse.messaging import emit_info
+from code_muse.plugins.plan_mode.plan_mode_tools import (
+    PlanModeState,
+    get_current_mode,
+    set_plan_mode,
+)
+
+# Cyclic ordering: DEFAULT → AUTO_EDIT → PLAN → DEFAULT
+_MODE_CYCLE = [
+    PlanModeState.DEFAULT,
+    PlanModeState.AUTO_EDIT,
+    PlanModeState.PLAN,
+]
+
+
+def cycle_mode() -> PlanModeState:
+    """Advance to the next mode in the cycle and emit a UI indicator.
+
+    Returns:
+        The new active mode.
+    """
+    current = get_current_mode()
+    try:
+        idx = _MODE_CYCLE.index(current)
+    except ValueError:
+        idx = 0
+    next_idx = (idx + 1) % len(_MODE_CYCLE)
+    new_mode = _MODE_CYCLE[next_idx]
+    set_plan_mode(new_mode)
+    emit_info(f"🔁 Mode cycled to: {new_mode.value}")
+    return new_mode