code-muse 0.0.1__py3-none-any.whl

code_muse/agents/agent_manager.py

@@ -0,0 +1,740 @@
+"""Agent manager for handling different agent configurations."""
+
+import importlib
+import json
+import os
+import pkgutil
+import re
+import threading
+import uuid
+from pathlib import Path
+
+from pydantic_ai.messages import ModelMessage
+
+from code_muse.agents.base_agent import BaseAgent
+from code_muse.agents.json_agent import JSONAgent, discover_json_agents
+from code_muse.callbacks import on_agent_reload, on_register_agents
+from code_muse.messaging import emit_success, emit_warning
+
+# Registry of available agents (Python classes and JSON file paths)
+_AGENT_REGISTRY: dict[str, type[BaseAgent] | str] = {}
+_AGENT_HISTORIES: dict[str, list[ModelMessage]] = {}
+_CURRENT_AGENT: BaseAgent | None = None
+
+# Terminal session-based agent selection
+_SESSION_AGENTS_CACHE: dict[str, str] = {}
+_SESSION_FILE_LOADED: bool = False
+# FREE-THREADED: _SESSION_LOCK guards sync-only session cache access.
+_SESSION_LOCK = threading.Lock()
+
+
+# Session persistence file path
+def _get_session_file_path() -> Path:
+    """Get the path to the terminal sessions file."""
+    from ..config import STATE_DIR
+
+    return Path(STATE_DIR) / "terminal_sessions.json"
+
+
+def get_terminal_session_id() -> str:
+    """Get a unique identifier for the current terminal session.
+
+    Uses parent process ID (PPID) as the session identifier.
+    This works across all platforms and provides session isolation.
+
+    Returns:
+        str: Unique session identifier (e.g., "session_12345")
+    """
+    try:
+        ppid = os.getppid()
+        return f"session_{ppid}"
+    except OSError, AttributeError:
+        # Fallback to current process ID if PPID unavailable
+        return f"fallback_{os.getpid()}"
+
+
+def _is_process_alive(pid: int) -> bool:
+    """Check if a process with the given PID is still alive, cross-platform.
+
+    Args:
+        pid: Process ID to check
+
+    Returns:
+        bool: True if process likely exists, False otherwise
+    """
+    try:
+        if os.name == "nt":
+            # Windows: use OpenProcess to probe liveness safely
+            import ctypes
+            from ctypes import wintypes
+
+            PROCESS_QUERY_LIMITED_INFORMATION = 0x1000
+            kernel32 = ctypes.windll.kernel32  # type: ignore[attr-defined]
+            kernel32.OpenProcess.argtypes = [
+                wintypes.DWORD,
+                wintypes.BOOL,
+                wintypes.DWORD,
+            ]
+            kernel32.OpenProcess.restype = wintypes.HANDLE
+            handle = kernel32.OpenProcess(
+                PROCESS_QUERY_LIMITED_INFORMATION, False, int(pid)
+            )
+            if handle:
+                kernel32.CloseHandle(handle)
+                return True
+            # If access denied, process likely exists but we can't query it
+            last_error = kernel32.GetLastError()
+            # ERROR_ACCESS_DENIED = 5
+            return last_error == 5
+        else:
+            # Unix-like: signal 0 does not deliver a signal but checks existence
+            os.kill(int(pid), 0)
+            return True
+    except PermissionError:
+        # No permission to signal -> process exists
+        return True
+    except OSError:
+        # Process does not exist
+        return False
+    except ValueError:
+        # Invalid signal or pid format
+        return False
+    except Exception:
+        # Be conservative – don't crash session cleanup due to platform quirks
+        return True
+
+
+def _cleanup_dead_sessions(sessions: dict[str, str]) -> dict[str, str]:
+    """Remove sessions for processes that no longer exist.
+
+    Args:
+        sessions: Dictionary of session_id -> agent_name
+
+    Returns:
+        dict: Cleaned sessions dictionary
+    """
+    cleaned = {}
+    for session_id, agent_name in sessions.items():
+        if session_id.startswith("session_"):
+            try:
+                pid_str = session_id.replace("session_", "")
+                pid = int(pid_str)
+                if _is_process_alive(pid):
+                    cleaned[session_id] = agent_name
+                # else: skip dead session
+            except ValueError, TypeError:
+                # Invalid session ID format, keep it anyway
+                cleaned[session_id] = agent_name
+        else:
+            # Non-standard session ID (like "fallback_"), keep it
+            cleaned[session_id] = agent_name
+    return cleaned
+
+
+def _load_session_data() -> dict[str, str]:
+    """Load terminal session data from the JSON file.
+
+    Returns:
+        dict: Session ID to agent name mapping
+    """
+    session_file = _get_session_file_path()
+    try:
+        if session_file.exists():
+            with open(session_file, encoding="utf-8") as f:
+                data = json.load(f)
+                # Clean up dead sessions while loading
+                return _cleanup_dead_sessions(data)
+        return {}
+    except json.JSONDecodeError, OSError:
+        # File corrupted or permission issues, start fresh
+        return {}
+
+
+def _save_session_data(sessions: dict[str, str]) -> None:
+    """Save terminal session data to the JSON file.
+
+    Args:
+        sessions: Session ID to agent name mapping
+    """
+    session_file = _get_session_file_path()
+    try:
+        # Ensure the config directory exists
+        session_file.parent.mkdir(parents=True, exist_ok=True)
+
+        # Clean up dead sessions before saving
+        cleaned_sessions = _cleanup_dead_sessions(sessions)
+
+        # Write to file atomically (write to temp file, then rename)
+        temp_file = session_file.with_suffix(".tmp")
+        with open(temp_file, "w", encoding="utf-8") as f:
+            json.dump(cleaned_sessions, f, indent=2)
+
+        # Atomic rename (works on all platforms)
+        temp_file.replace(session_file)
+
+    except OSError:
+        # File permission issues, etc. - just continue without persistence
+        pass
+
+
+def _ensure_session_cache_loaded() -> None:
+    """Ensure the session cache is loaded from disk."""
+    global _SESSION_AGENTS_CACHE, _SESSION_FILE_LOADED
+    with _SESSION_LOCK:
+        if not _SESSION_FILE_LOADED:
+            _SESSION_AGENTS_CACHE.update(_load_session_data())
+            _SESSION_FILE_LOADED = True
+
+
+def _discover_agents(message_group_id: str | None = None):
+    """Dynamically discover all agent classes and JSON agents."""
+    # Always clear the registry to force refresh
+    _AGENT_REGISTRY.clear()
+
+    # 1. Discover Python agent classes in the agents package
+    import code_muse.agents as agents_package
+
+    # Iterate through all modules in the agents package
+    for _, modname, _ in pkgutil.iter_modules(agents_package.__path__):
+        if modname.startswith("_") or modname in [
+            "base_agent",
+            "json_agent",
+            "agent_manager",
+        ]:
+            continue
+
+        try:
+            # Import the module
+            module = importlib.import_module(f"code_muse.agents.{modname}")
+
+            # Look for BaseAgent subclasses
+            for attr_name in dir(module):
+                attr = getattr(module, attr_name)
+                if (
+                    isinstance(attr, type)
+                    and issubclass(attr, BaseAgent)
+                    and attr not in [BaseAgent, JSONAgent]
+                ):
+                    # Create an instance to get the name
+                    agent_instance = attr()
+                    _AGENT_REGISTRY[agent_instance.name] = attr
+
+        except Exception as e:
+            # Skip problematic modules
+            emit_warning(
+                f"Warning: Could not load agent module {modname}: {e}",
+                message_group=message_group_id,
+            )
+            continue
+
+    # 1b. Discover agents in sub-packages (like 'pack')
+    for _, subpkg_name, ispkg in pkgutil.iter_modules(agents_package.__path__):
+        if not ispkg or subpkg_name.startswith("_"):
+            continue
+
+        try:
+            # Import the sub-package
+            subpkg = importlib.import_module(f"code_muse.agents.{subpkg_name}")
+
+            # Iterate through modules in the sub-package
+            if not hasattr(subpkg, "__path__"):
+                continue
+
+            for _, modname, _ in pkgutil.iter_modules(subpkg.__path__):
+                if modname.startswith("_"):
+                    continue
+
+                try:
+                    # Import the submodule
+                    module = importlib.import_module(
+                        f"code_muse.agents.{subpkg_name}.{modname}"
+                    )
+
+                    # Look for BaseAgent subclasses
+                    for attr_name in dir(module):
+                        attr = getattr(module, attr_name)
+                        if (
+                            isinstance(attr, type)
+                            and issubclass(attr, BaseAgent)
+                            and attr not in [BaseAgent, JSONAgent]
+                        ):
+                            # Create an instance to get the name
+                            agent_instance = attr()
+                            _AGENT_REGISTRY[agent_instance.name] = attr
+
+                except Exception as e:
+                    emit_warning(
+                        f"Warning: Could not load agent {subpkg_name}.{modname}: {e}",
+                        message_group=message_group_id,
+                    )
+                    continue
+
+        except Exception as e:
+            emit_warning(
+                f"Warning: Could not load agent sub-package {subpkg_name}: {e}",
+                message_group=message_group_id,
+            )
+            continue
+
+    # 2. Discover JSON agents in user directory
+    try:
+        json_agents = discover_json_agents()
+
+        # Add JSON agents to registry (store file path instead of class)
+        # Python (builtin) agents take precedence over JSON agents.
+        for agent_name, json_path in json_agents.items():
+            if agent_name in _AGENT_REGISTRY:
+                emit_warning(
+                    f"JSON agent '{agent_name}' skipped: builtin Python agent with the same name takes precedence.",
+                    message_group=message_group_id,
+                )
+                continue
+            _AGENT_REGISTRY[agent_name] = json_path
+
+    except Exception as e:
+        emit_warning(
+            f"Warning: Could not discover JSON agents: {e}",
+            message_group=message_group_id,
+        )
+
+    # 3. Discover agents registered by plugins
+    try:
+        results = on_register_agents()
+        for result in results:
+            if result is None:
+                continue
+            # Each result should be a list of agent definitions
+            agents_list = result if isinstance(result, list) else [result]
+            for agent_def in agents_list:
+                if not isinstance(agent_def, dict) or "name" not in agent_def:
+                    continue
+
+                agent_name = agent_def["name"]
+
+                # Support both class-based and JSON path-based registration
+                if "class" in agent_def:
+                    agent_class = agent_def["class"]
+                    if isinstance(agent_class, type) and issubclass(
+                        agent_class, BaseAgent
+                    ):
+                        _AGENT_REGISTRY[agent_name] = agent_class
+                elif "json_path" in agent_def:
+                    json_path = agent_def["json_path"]
+                    if isinstance(json_path, str):
+                        _AGENT_REGISTRY[agent_name] = json_path
+
+    except Exception as e:
+        emit_warning(
+            f"Warning: Could not load plugin agents: {e}",
+            message_group=message_group_id,
+        )
+
+
+def get_available_agents() -> dict[str, str]:
+    """Get a dictionary of available agents with their display names.
+
+    Returns:
+        Dict mapping agent names to display names.
+    """
+    from ..config import (
+        PACK_AGENT_NAMES,
+        UC_AGENT_NAMES,
+        get_pack_agents_enabled,
+        get_universal_constructor_enabled,
+    )
+
+    # Generate a message group ID for this operation
+    message_group_id = str(uuid.uuid4())
+    _discover_agents(message_group_id=message_group_id)
+
+    # Check if pack agents are enabled
+    pack_agents_enabled = get_pack_agents_enabled()
+
+    # Check if UC is enabled
+    uc_enabled = get_universal_constructor_enabled()
+
+    agents = {}
+    for name, agent_ref in _AGENT_REGISTRY.items():
+        # Filter out pack agents if disabled
+        if not pack_agents_enabled and name in PACK_AGENT_NAMES:
+            continue
+
+        # Filter out UC-dependent agents if UC is disabled
+        if not uc_enabled and name in UC_AGENT_NAMES:
+            continue
+
+        try:
+            if isinstance(agent_ref, str):  # JSON agent (file path)
+                agent_instance = JSONAgent(agent_ref)
+            else:  # Python agent (class)
+                agent_instance = agent_ref()
+            agents[name] = agent_instance.display_name
+        except Exception:
+            agents[name] = name.title()  # Fallback
+
+    return agents
+
+
+def get_current_agent_name() -> str:
+    """Get the name of the currently active agent for this terminal session.
+
+    Returns:
+        The name of the current agent for this session.
+        Priority: session agent > config default > 'muse'.
+    """
+    _ensure_session_cache_loaded()
+    session_id = get_terminal_session_id()
+
+    # First check for session-specific agent
+    with _SESSION_LOCK:
+        session_agent = _SESSION_AGENTS_CACHE.get(session_id)
+    if session_agent:
+        return session_agent
+
+    # Fall back to config default
+    from ..config import get_default_agent
+
+    return get_default_agent()
+
+
+def set_current_agent(agent_name: str) -> bool:
+    """Set the current agent by name.
+
+    Args:
+        agent_name: The name of the agent to set as current.
+
+    Returns:
+        True if the agent was set successfully, False if agent not found.
+    """
+    global _CURRENT_AGENT
+    curr_agent = get_current_agent()
+    if curr_agent is not None:
+        # Store a shallow copy so future mutations don't affect saved history
+        _AGENT_HISTORIES[curr_agent.name] = list(curr_agent.get_message_history())
+    # Generate a message group ID for agent switching
+    message_group_id = str(uuid.uuid4())
+    _discover_agents(message_group_id=message_group_id)
+
+    # Save current agent's history before switching
+
+    # Clear the cached config when switching agents
+    agent_obj = load_agent(agent_name)
+    _CURRENT_AGENT = agent_obj
+
+    # Update session-based agent selection and persist to disk
+    _ensure_session_cache_loaded()
+    session_id = get_terminal_session_id()
+    with _SESSION_LOCK:
+        _SESSION_AGENTS_CACHE[session_id] = agent_name
+        cache_snapshot = dict(_SESSION_AGENTS_CACHE)
+    _save_session_data(cache_snapshot)
+    if agent_obj.name in _AGENT_HISTORIES:
+        # Restore a copy to avoid sharing the same list instance
+        agent_obj.set_message_history(list(_AGENT_HISTORIES[agent_obj.name]))
+    on_agent_reload(agent_obj.id, agent_name)
+    return True
+
+
+def get_current_agent() -> BaseAgent:
+    """Get the current agent configuration.
+
+    Returns:
+        The current agent configuration instance.
+    """
+    global _CURRENT_AGENT
+
+    if _CURRENT_AGENT is None:
+        agent_name = get_current_agent_name()
+        _CURRENT_AGENT = load_agent(agent_name)
+
+    return _CURRENT_AGENT
+
+
+def load_agent(agent_name: str) -> BaseAgent:
+    """Load an agent configuration by name.
+
+    Args:
+        agent_name: The name of the agent to load.
+
+    Returns:
+        The agent configuration instance.
+
+    Raises:
+        ValueError: If the agent is not found.
+    """
+    # Generate a message group ID for agent loading
+    message_group_id = str(uuid.uuid4())
+    _discover_agents(message_group_id=message_group_id)
+
+    if agent_name not in _AGENT_REGISTRY:
+        # Fallback to muse if agent not found
+        if "muse" in _AGENT_REGISTRY:
+            agent_name = "muse"
+        else:
+            raise ValueError(
+                f"Agent '{agent_name}' not found and no fallback available"
+            )
+
+    agent_ref = _AGENT_REGISTRY[agent_name]
+    if isinstance(agent_ref, str):  # JSON agent (file path)
+        return JSONAgent(agent_ref)
+    else:  # Python agent (class)
+        return agent_ref()
+
+
+def get_agent_descriptions() -> dict[str, str]:
+    """Get descriptions for all available agents.
+
+    Returns:
+        Dict mapping agent names to their descriptions.
+    """
+    from ..config import (
+        PACK_AGENT_NAMES,
+        UC_AGENT_NAMES,
+        get_pack_agents_enabled,
+        get_universal_constructor_enabled,
+    )
+
+    # Generate a message group ID for this operation
+    message_group_id = str(uuid.uuid4())
+    _discover_agents(message_group_id=message_group_id)
+
+    # Check if pack agents are enabled
+    pack_agents_enabled = get_pack_agents_enabled()
+
+    # Check if UC is enabled
+    uc_enabled = get_universal_constructor_enabled()
+
+    descriptions = {}
+    for name, agent_ref in _AGENT_REGISTRY.items():
+        # Filter out pack agents if disabled
+        if not pack_agents_enabled and name in PACK_AGENT_NAMES:
+            continue
+
+        # Filter out UC-dependent agents if UC is disabled
+        if not uc_enabled and name in UC_AGENT_NAMES:
+            continue
+
+        try:
+            if isinstance(agent_ref, str):  # JSON agent (file path)
+                agent_instance = JSONAgent(agent_ref)
+            else:  # Python agent (class)
+                agent_instance = agent_ref()
+            descriptions[name] = agent_instance.description
+        except Exception:
+            descriptions[name] = "No description available"
+
+    return descriptions
+
+
+def refresh_agents():
+    """Refresh the agent discovery to pick up newly created agents.
+
+    This clears the agent registry cache and forces a rediscovery of all agents.
+    """
+    # Generate a message group ID for agent refreshing
+    message_group_id = str(uuid.uuid4())
+    _discover_agents(message_group_id=message_group_id)
+
+
+_CLONE_NAME_PATTERN = re.compile(r"^(?P<base>.+)-clone-(?P<index>\d+)$")
+_CLONE_DISPLAY_PATTERN = re.compile(r"\s*\(Clone\s+\d+\)$", re.IGNORECASE)
+
+
+def _strip_clone_suffix(agent_name: str) -> str:
+    """Strip a trailing -clone-N suffix from a name if present."""
+    match = _CLONE_NAME_PATTERN.match(agent_name)
+    return match.group("base") if match else agent_name
+
+
+def _strip_clone_display_suffix(display_name: str) -> str:
+    """Remove a trailing "(Clone N)" suffix from display names."""
+    cleaned = _CLONE_DISPLAY_PATTERN.sub("", display_name).strip()
+    return cleaned or display_name
+
+
+def is_clone_agent_name(agent_name: str) -> bool:
+    """Return True if the agent name looks like a clone."""
+    return bool(_CLONE_NAME_PATTERN.match(agent_name))
+
+
+def _default_display_name(agent_name: str) -> str:
+    """Build a default display name from an agent name."""
+    title = agent_name.title()
+    return f"{title} 🤖"
+
+
+def _build_clone_display_name(display_name: str, clone_index: int) -> str:
+    """Build a clone display name based on the source display name."""
+    base_name = _strip_clone_display_suffix(display_name)
+    return f"{base_name} (Clone {clone_index})"
+
+
+def _filter_available_tools(tool_names: list[str]) -> list[str]:
+    """Filter a tool list to only available tool names."""
+    from code_muse.tools import get_available_tool_names
+
+    available_tools = set(get_available_tool_names())
+    return [tool for tool in tool_names if tool in available_tools]
+
+
+def _next_clone_index(
+    base_name: str, existing_names: set[str], agents_dir: Path
+) -> int:
+    """Compute the next clone index for a base name."""
+    clone_pattern = re.compile(rf"^{re.escape(base_name)}-clone-(\\d+)$")
+    indices = []
+    for name in existing_names:
+        match = clone_pattern.match(name)
+        if match:
+            indices.append(int(match.group(1)))
+
+    next_index = max(indices, default=0) + 1
+    while True:
+        clone_name = f"{base_name}-clone-{next_index}"
+        clone_path = agents_dir / f"{clone_name}.json"
+        if clone_name not in existing_names and not clone_path.exists():
+            return next_index
+        next_index += 1
+
+
+def clone_agent(agent_name: str) -> str | None:
+    """Clone an agent definition into the user agents directory.
+
+    Args:
+        agent_name: Source agent name to clone.
+
+    Returns:
+        The cloned agent name, or None if cloning failed.
+    """
+    # Generate a message group ID for agent cloning
+    message_group_id = str(uuid.uuid4())
+    _discover_agents(message_group_id=message_group_id)
+
+    agent_ref = _AGENT_REGISTRY.get(agent_name)
+    if agent_ref is None:
+        emit_warning(f"Agent '{agent_name}' not found for cloning.")
+        return None
+
+    from ..config import get_agent_pinned_model, get_user_agents_directory
+
+    agents_dir = Path(get_user_agents_directory())
+    base_name = _strip_clone_suffix(agent_name)
+    existing_names = set(_AGENT_REGISTRY.keys())
+    clone_index = _next_clone_index(base_name, existing_names, agents_dir)
+    clone_name = f"{base_name}-clone-{clone_index}"
+    clone_path = agents_dir / f"{clone_name}.json"
+
+    try:
+        if isinstance(agent_ref, str):
+            with open(agent_ref, encoding="utf-8") as f:
+                source_config = json.load(f)
+
+            source_display_name = source_config.get("display_name")
+            if not source_display_name:
+                source_display_name = _default_display_name(base_name)
+
+            clone_config = dict(source_config)
+            clone_config["name"] = clone_name
+            clone_config["display_name"] = _build_clone_display_name(
+                source_display_name, clone_index
+            )
+
+            tools = source_config.get("tools", [])
+            clone_config["tools"] = (
+                _filter_available_tools(tools) if isinstance(tools, list) else []
+            )
+
+            if not clone_config.get("model"):
+                clone_config.pop("model", None)
+        else:
+            agent_instance = agent_ref()
+            clone_config = {
+                "name": clone_name,
+                "display_name": _build_clone_display_name(
+                    agent_instance.display_name, clone_index
+                ),
+                "description": agent_instance.description,
+                "system_prompt": agent_instance.get_full_system_prompt(),
+                "tools": _filter_available_tools(agent_instance.get_available_tools()),
+            }
+
+            user_prompt = agent_instance.get_user_prompt()
+            if user_prompt is not None:
+                clone_config["user_prompt"] = user_prompt
+
+            tools_config = agent_instance.get_tools_config()
+            if tools_config is not None:
+                clone_config["tools_config"] = tools_config
+
+            pinned_model = get_agent_pinned_model(agent_instance.name)
+            if pinned_model:
+                clone_config["model"] = pinned_model
+    except Exception as exc:
+        emit_warning(f"Failed to build clone for '{agent_name}': {exc}")
+        return None
+
+    if clone_path.exists():
+        emit_warning(f"Clone target '{clone_name}' already exists.")
+        return None
+
+    try:
+        with open(clone_path, "w", encoding="utf-8") as f:
+            json.dump(clone_config, f, indent=2, ensure_ascii=False)
+        emit_success(f"Cloned '{agent_name}' to '{clone_name}'.")
+        return clone_name
+    except Exception as exc:
+        emit_warning(f"Failed to write clone file '{clone_path}': {exc}")
+        return None
+
+
+def delete_clone_agent(agent_name: str) -> bool:
+    """Delete a cloned JSON agent definition.
+
+    Args:
+        agent_name: Clone agent name to delete.
+
+    Returns:
+        True if the clone was deleted, False otherwise.
+    """
+    message_group_id = str(uuid.uuid4())
+    _discover_agents(message_group_id=message_group_id)
+
+    if not is_clone_agent_name(agent_name):
+        emit_warning(f"Agent '{agent_name}' is not a clone.")
+        return False
+
+    if get_current_agent_name() == agent_name:
+        emit_warning("Cannot delete the active agent. Switch agents first.")
+        return False
+
+    agent_ref = _AGENT_REGISTRY.get(agent_name)
+    if agent_ref is None:
+        emit_warning(f"Clone '{agent_name}' not found.")
+        return False
+
+    if not isinstance(agent_ref, str):
+        emit_warning(f"Clone '{agent_name}' is not a JSON agent.")
+        return False
+
+    clone_path = Path(agent_ref)
+    if not clone_path.exists():
+        emit_warning(f"Clone file for '{agent_name}' does not exist.")
+        return False
+
+    from ..config import get_user_agents_directory
+
+    agents_dir = Path(get_user_agents_directory()).resolve()
+    if clone_path.resolve().parent != agents_dir:
+        emit_warning(f"Refusing to delete non-user clone '{agent_name}'.")
+        return False
+
+    try:
+        clone_path.unlink()
+        emit_success(f"Deleted clone '{agent_name}'.")
+        _AGENT_REGISTRY.pop(agent_name, None)
+        _AGENT_HISTORIES.pop(agent_name, None)
+        return True
+    except Exception as exc:
+        emit_warning(f"Failed to delete clone '{agent_name}': {exc}")
+        return False