ajarbot/agent_registry.py

"""Agent Registry - Thread-safe global singleton for MCP tool access.

MCP tools are module-level functions that cannot access the Agent instance
directly. This registry provides a thread-safe bridge so that tools like
delegate_task can call Agent.delegate() without circular imports.

Usage:
    # At bot startup (bot_runner.py):
    from agent_registry import register_agent
    agent = Agent(...)
    register_agent(agent)

    # In MCP tools (mcp_tools.py):
    from agent_registry import get_agent
    agent = get_agent()
    if agent:
        result = agent.delegate(task, specialist_prompt)
"""

import threading
from typing import Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from agent import Agent

# Module-level singleton state
_agent: Optional['Agent'] = None
_lock = threading.Lock()


def register_agent(agent: 'Agent') -> None:
    """Register the main Agent instance for MCP tool access.

    Must be called exactly once at bot startup, after Agent is initialized.
    Thread-safe.

    Args:
        agent: The main Agent instance (not a sub-agent).

    Raises:
        ValueError: If agent is None or is a sub-agent.
    """
    global _agent

    if agent is None:
        raise ValueError("Cannot register None as the main agent")
    if getattr(agent, 'is_sub_agent', False):
        raise ValueError("Cannot register a sub-agent as the main agent")

    with _lock:
        _agent = agent
        print(f"[AgentRegistry] Main agent registered (provider={agent.llm.provider})")


def get_agent() -> Optional['Agent']:
    """Get the registered main Agent instance.

    Thread-safe. Returns None if no agent has been registered yet.

    Returns:
        The main Agent instance, or None.
    """
    with _lock:
        return _agent


def clear_agent() -> None:
    """Clear the registered agent (for testing or shutdown).

    Thread-safe.
    """
    global _agent

    with _lock:
        _agent = None
        print("[AgentRegistry] Agent registry cleared")