"""Thin coordinator for agent sessions. Heavy logic lives in sibling modules: - agent_manager_ops – edit, close, resume, duplicate, invoke, LLM metadata - agent_loop – the SDK query loop, streaming, mock agent - agent_mock – session-completed analytics - prompt_builder – system-prompt composition & context injection - mcp_builder – MCP server construction & tool-policy helpers - session_store – on-disk persistence, history, message copying """ import asyncio import logging import os from datetime import datetime from typing import Optional from uuid import uuid4 from pydantic import InstanceOf from backend.apps.agents.models import AgentConfig, AgentSession, Message from backend.apps.agents.manager.ws_manager import ws_manager from backend.apps.agents.execution.prompt_builder import resolve_mode from backend.apps.agents.execution.mcp_builder import get_all_tool_names from backend.apps.agents.manager.session_store import ( get_history, reconcile_on_startup, get_browser_agent_children, ) from backend.apps.agents.execution.agent_loop import run_agent_loop from backend.apps.agents.manager.agent_manager_ops import ( edit_message_op, close_session_op, resume_session_op, duplicate_session_op, invoke_agent_op ) from backend.apps.agents.manager.agent_manager_meta import ( generate_title_op, generate_group_meta_op, persist_all_sessions_op, restore_all_sessions_op, delete_session_op ) from backend.apps.settings.settings import load_settings from backend.apps.analytics.collector import record as _analytics logger = logging.getLogger(__name__) os.environ.setdefault("CLAUDE_CODE_STREAM_CLOSE_TIMEOUT", "3600000") class AgentManager: def __init__(self): self.sessions: dict[str, AgentSession] = {} self.tasks: dict[str, InstanceOf[asyncio.Task]] = {} async def launch_agent(self, config: AgentConfig) -> AgentSession: session_id = uuid4().hex mode_tools, _, mode_folder = resolve_mode(config.mode, get_all_tool_names) global_settings = load_settings() effective_cwd = ( config.target_directory or mode_folder or global_settings.default_folder or os.path.expanduser("~") ) if config.mode in ("view-builder", "skill-builder") and not config.target_directory: effective_cwd = os.path.join(effective_cwd, session_id) os.makedirs(effective_cwd, exist_ok=True) session = AgentSession( id=session_id, name=config.name, provider=getattr(config, "provider", "anthropic"), model=config.model, mode=config.mode, system_prompt=config.system_prompt, allowed_tools=mode_tools, max_turns=config.max_turns, cwd=effective_cwd, dashboard_id=config.dashboard_id, ) self.sessions[session_id] = session _analytics("session.started", { "model": session.model, "provider": session.provider, "mode": session.mode, "tool_count": len(mode_tools), }, session_id=session_id, dashboard_id=config.dashboard_id) await ws_manager.emit_status(session_id, "running", session) return session async def send_message( self, session_id: str, prompt: str, mode: str | None = None, model: str | None = None, provider: str | None = None, images: list | None = None, context_paths: list | None = None, forced_tools: list[str] | None = None, attached_skills: list | None = None, hidden: bool = False, selected_browser_ids: list[str] | None = None, ): session = self.sessions.get(session_id) if not session: raise ValueError(f"Session {session_id} not found") existing = self.tasks.get(session_id) if existing and not existing.done(): return session_changed = False if model and model != session.model: _analytics("model.switched", { "from_model": session.model, "to_model": model, "from_provider": session.provider, "to_provider": provider or session.provider, "message_number": len([m for m in session.messages if m.role == "user"]), "cost_so_far": session.cost_usd, }, session_id=session_id, dashboard_id=session.dashboard_id) session.model = model session_changed = True if mode and mode != session.mode: _analytics("feature.used", {"feature": "mode.switched", "from_mode": session.mode, "to_mode": mode}, session_id=session_id, dashboard_id=session.dashboard_id) session.mode = mode mode_tools, _, _ = resolve_mode(mode, get_all_tool_names) session.allowed_tools = mode_tools session_changed = True if session_changed: await ws_manager.emit_status(session_id, session.status, session) skill_meta = [{"id": s["id"], "name": s["name"]} for s in (attached_skills or [])] or None image_meta = [{"data": img["data"], "media_type": img.get("media_type", "image/png")} for img in (images or [])] or None user_msg = Message( role="user", content=prompt, branch_id=session.active_branch_id, context_paths=context_paths or None, attached_skills=skill_meta, forced_tools=forced_tools or None, images=image_meta, hidden=hidden, ) session.messages.append(user_msg) await ws_manager.emit_message(session_id, user_msg) if context_paths or attached_skills or images or forced_tools: _analytics("context.attached", { "file_count": len([c for c in (context_paths or []) if c.get("type") == "file"]), "directory_count": len([c for c in (context_paths or []) if c.get("type") == "directory"]), "skill_count": len(attached_skills or []), "image_count": len(images or []), "has_forced_tools": bool(forced_tools), }, session_id=session_id, dashboard_id=session.dashboard_id) for skill in (attached_skills or []): _analytics("feature.used", {"feature": "skill.used", "skill_name": skill.get("name", "")}, session_id=session_id, dashboard_id=session.dashboard_id) is_first = sum(1 for m in session.messages if m.role == "user") == 1 if is_first: _analytics("session.first_message", { "message_length": len(prompt), "has_code_block": "```" in prompt, "has_url": "http://" in prompt or "https://" in prompt, "model": session.model, "mode": session.mode, }, session_id=session_id, dashboard_id=session.dashboard_id) session.status = "running" await ws_manager.emit_status(session_id, "running", session) task = asyncio.create_task(run_agent_loop( self.sessions, session_id, prompt, images=images, context_paths=context_paths, forced_tools=forced_tools, attached_skills=attached_skills, selected_browser_ids=selected_browser_ids, )) self.tasks[session_id] = task async def stop_agent(self, session_id: str): task = self.tasks.get(session_id) if task and not task.done(): task.cancel() try: await task except asyncio.CancelledError: pass session = self.sessions.get(session_id) if session: for req in list(session.pending_approvals): ws_manager.resolve_approval(req.id, {"behavior": "deny", "message": "Agent stopped"}) session.pending_approvals = [] if hasattr(session, '_cancel_event'): session._cancel_event.set() session.status = "stopped" if not session.closed_at: session.closed_at = datetime.now() await ws_manager.emit_status(session_id, "stopped", session) children = [s for s in self.sessions.values() if s.parent_session_id == session_id and s.mode == "browser-agent"] for child in children: await self.stop_agent(child.id) def handle_approval(self, request_id: str, decision: dict): ws_manager.resolve_approval(request_id, decision) async def edit_message(self, session_id: str, message_id: str, new_content: str): await edit_message_op(self.sessions, self.tasks, session_id, message_id, new_content) async def switch_branch(self, session_id: str, branch_id: str): session = self.sessions.get(session_id) if not session: raise ValueError(f"Session {session_id} not found") if branch_id not in session.branches: raise ValueError(f"Branch {branch_id} not found") session.active_branch_id = branch_id await ws_manager.emit_branch_switched(session_id, branch_id) async def generate_title(self, session_id: str, first_prompt: str) -> str: return await generate_title_op(self.sessions, session_id, first_prompt) async def generate_group_meta(self, session_id: str, group_id: str, tool_calls: list[dict], results_summary: list[str] | None = None, is_refinement: bool = False) -> dict: return await generate_group_meta_op(self.sessions, session_id, group_id, tool_calls, results_summary, is_refinement) async def update_session(self, session_id: str, **fields): session = self.sessions.get(session_id) if not session: raise ValueError(f"Session {session_id} not found") for key, value in fields.items(): if key in {"system_prompt", "name"}: setattr(session, key, value) await ws_manager.emit_status(session_id, session.status, session) async def close_session(self, session_id: str) -> None: await close_session_op(self.sessions, self.tasks, session_id) async def delete_session(self, session_id: str) -> None: await delete_session_op(self, session_id) async def resume_session(self, session_id: str) -> AgentSession: return await resume_session_op(self.sessions, session_id) async def duplicate_session(self, session_id: str, dashboard_id: str | None = None, up_to_message_id: str | None = None) -> AgentSession: return await duplicate_session_op(self.sessions, session_id, dashboard_id, up_to_message_id) async def invoke_agent(self, source_session_id: str, message: str, parent_session_id: str | None = None, dashboard_id: str | None = None) -> dict: return await invoke_agent_op(self.sessions, source_session_id, message, parent_session_id, dashboard_id) def get_all_sessions(self, dashboard_id: str | None = None) -> list[AgentSession]: if dashboard_id: return [s for s in self.sessions.values() if s.dashboard_id == dashboard_id] return list(self.sessions.values()) def get_session(self, session_id: str) -> Optional[AgentSession]: return self.sessions.get(session_id) def get_history(self, q: str = "", limit: int = 20, offset: int = 0, dashboard_id: str | None = None) -> dict: return get_history(q=q, limit=limit, offset=offset, dashboard_id=dashboard_id) async def reconcile_on_startup(self) -> None: return await reconcile_on_startup() async def persist_all_sessions(self) -> None: await persist_all_sessions_op(self.sessions, self.tasks) async def restore_all_sessions(self) -> None: await restore_all_sessions_op(self.sessions) def get_browser_agent_children(self, parent_session_id: str) -> list[dict]: return get_browser_agent_children(self.sessions, parent_session_id) agent_manager = AgentManager()