"""Agent sessions sub-app. Endpoints operate directly on a module-level sessions dict and the Agent class. No manager layer — Agent already encapsulates its own runtime state. The ws module is used ONLY in this file — the Agent class and its internals communicate via the on_event callback, never importing ws directly. """ import asyncio import os from contextlib import asynccontextmanager from datetime import datetime from typing import Optional, List, Dict from uuid import uuid4 from fastapi import HTTPException, WebSocket, WebSocketDisconnect, Body import json from pydantic import BaseModel from backend.config.Apps import SubApp from backend.config.paths import DB_ROOT from backend.core.Agent.Agent import Agent from backend.core.db.PydanticStore import PydanticStore from backend.core.shared_structs.agent.Message.Message import UserMessage from backend.core.events.events import AgentStatusEvent, AgentClosedEvent, BranchSwitchedEvent from backend.apps.agents.agent_utils.create_sdk_hooks import create_sdk_hooks from backend.apps.agents.agent_utils.build_search_text import build_search_text from backend.apps.agents.ResolvedModeConfig.ResolvedModeConfig import ResolvedModeConfig from backend.apps.agents.COMMS_MANAGER.COMMS_MANAGER import COMMS_MANAGER from backend.apps.settings.settings import load_settings from backend.core.llm.resolve_sdk_env import resolve_sdk_env from backend.ports import NINE_ROUTER_PORT from swarm_debug import debug from claude_agent_sdk import ClaudeAgentOptions from claude_agent_sdk.types import HookMatcher, McpServerConfig from backend.core.tools.shared_structs.Toolkit import Toolkit from backend.apps.agents.agent_utils.build_agent_toolkit import build_agent_toolkit AGENT_STORE: PydanticStore[Agent] = PydanticStore[Agent]( model_cls=Agent, data_dir=os.path.join(DB_ROOT, "sessions"), id_field="session_id", dump_mode="json", not_found_detail="Session not found in history", ) # NOTE: Essentially the SESSIONS is a cache for active agents. SESSIONS: dict[str, Agent] = {} def _persist_agent(agent: Agent) -> None: """Called when an agent reaches a terminal state (completed/error).""" debug("auto-saving session %s (status=%s)", agent.session_id, agent.status) try: AGENT_STORE.save(agent) except Exception as e: debug("auto-save failed for session %s: %s", agent.session_id, e) def get_agent(session_id: str) -> Agent: agent: Optional[Agent] = SESSIONS.get(session_id) if not agent: raise HTTPException(status_code=404, detail="Session not found") return agent # --------------------------------------------------------------------------- # Lifespan # --------------------------------------------------------------------------- @asynccontextmanager async def agents_lifespan(): for stored in AGENT_STORE.load_all(): try: stored.status = "stopped" stored.on_event = COMMS_MANAGER.make_session_emitter(stored.session_id) stored.on_done = _persist_agent stored.toolkit = await build_agent_toolkit( agent=stored, sessions=SESSIONS, comms_manager=COMMS_MANAGER, ) SESSIONS[stored.session_id] = stored except Exception as e: debug(f"[agents lifespan] Skipping corrupt session {stored.session_id}: {e}") yield debug("agents_lifespan: shutting down — %s sessions to save", len(SESSIONS)) for agent in list[Agent](SESSIONS.values()): debug("agents_lifespan: stopping agent %s", agent.session_id) await agent.stop_agent() debug("agents_lifespan: saving agent %s", agent.session_id) AGENT_STORE.save(agent) debug("agents_lifespan: saved agent %s", agent.session_id) SESSIONS.clear() debug("agents_lifespan: shutdown complete") agents = SubApp("agents", agents_lifespan) # --------------------------------------------------------------------------- # WebSocket # --------------------------------------------------------------------------- # TODO: type spec this more @agents.router.websocket("/ws") async def websocket_dashboard(websocket: WebSocket): await COMMS_MANAGER.broadcaster.connect(websocket) try: while True: raw: str = await websocket.receive_text() msg: dict = json.loads(raw) event: str = msg.get("event", "") payload: dict = msg.get("data", {}) if event == "agent:approval_response": await COMMS_MANAGER.resolve_approval( request_id=payload.get("request_id", ""), behavior=payload.get("behavior", "deny"), message=payload.get("message"), updated_input=payload.get("updated_input"), ) elif event == "browser:result": COMMS_MANAGER.browser_bridge.resolve( payload.get("request_id", ""), payload, ) except WebSocketDisconnect: COMMS_MANAGER.broadcaster.disconnect(websocket) # --------------------------------------------------------------------------- # Session CRUD # --------------------------------------------------------------------------- @agents.router.get("/get_all_sessions") async def get_all_sessions(dashboard_id: str = "") -> dict: result: List[Agent] = list(SESSIONS.values()) if dashboard_id: result = [a for a in result if a.dashboard_id == dashboard_id] return {"sessions": [a.model_dump(mode="json") for a in result]} @agents.router.get("/get_session") async def get_session(session_id: str = Body()) -> dict: return get_agent(session_id).model_dump(mode="json") @agents.router.post("/launch_agent") async def launch_agent( model: str = Body(), mode: str = Body(), system_prompt: str = Body(), max_turns: int = Body(), dashboard_id: Optional[str] = Body(default=None), ) -> dict: agent: Agent = Agent( model=model, mode=mode, status="stopped", dashboard_id=dashboard_id, config=ClaudeAgentOptions(max_turns=max_turns), ) agent.on_event = COMMS_MANAGER.make_session_emitter(agent.session_id) agent.on_done = _persist_agent SESSIONS[agent.session_id] = agent toolkit: Toolkit = await build_agent_toolkit( agent=agent, sessions=SESSIONS, comms_manager=COMMS_MANAGER, ) agent.toolkit = toolkit mcp_servers: Dict[str, McpServerConfig] = toolkit.collect_mcp_servers() resolved_mode_config: ResolvedModeConfig = await ResolvedModeConfig.create( mode_id=mode, session_prompt=system_prompt or None, toolkit=toolkit, ) can_use_tool, pre_tool_hook, post_tool_hook = create_sdk_hooks(agent) settings = load_settings() env = resolve_sdk_env( api_key=settings.anthropic_api_key, nine_router_port=NINE_ROUTER_PORT if not settings.anthropic_api_key else None, ) NINE_ROUTER_MODEL_MAP = { "sonnet": "cc/claude-sonnet-4-6", "opus": "cc/claude-opus-4-6", "haiku": "cc/claude-haiku-4-5-20251001", } resolved_model = NINE_ROUTER_MODEL_MAP.get(model, f"cc/{model}") if not settings.anthropic_api_key else model agent.config = ClaudeAgentOptions( env=env, model=resolved_model, system_prompt=resolved_mode_config.system_prompt, max_turns=max_turns, cwd=resolved_mode_config.cwd, mcp_servers=mcp_servers if mcp_servers else None, allowed_tools=resolved_mode_config.allowed_tools, disallowed_tools=resolved_mode_config.disallowed_tools, permission_mode="default", can_use_tool=can_use_tool, hooks={ "PreToolUse": [HookMatcher(matcher=None, hooks=[pre_tool_hook])], "PostToolUse": [HookMatcher(matcher=None, hooks=[post_tool_hook])], }, ) await agent.emit(AgentStatusEvent( session_id=agent.session_id, status="stopped", session=agent.snapshot(), )) return {"session_id": agent.session_id, "session": agent.snapshot().model_dump(mode="json")} @agents.router.patch("/update_system_prompt") async def update_system_prompt( session_id: str = Body(), system_prompt: Optional[str] = Body(default=None), ) -> dict: agent: Agent = get_agent(session_id) if system_prompt is not None: agent.config.system_prompt = system_prompt await agent.emit(AgentStatusEvent( session_id=session_id, status=agent.status, session=agent.snapshot(), )) return {"ok": True} @agents.router.delete("/delete_session") async def delete_session(session_id: str = Body()) -> dict: agent: Optional[Agent] = SESSIONS.pop(session_id, None) if agent is not None: await agent.stop_agent() AGENT_STORE.delete(session_id) return {"ok": True} # --------------------------------------------------------------------------- # Agent lifecycle # --------------------------------------------------------------------------- class MessageBody(BaseModel): prompt: str mode: Optional[str] = None model: Optional[str] = None images: Optional[List[str]] = None image_media_types: Optional[List[str]] = None context_paths: Optional[List[dict]] = None forced_tools: Optional[List[str]] = None attached_skills: Optional[List[dict]] = None hidden: bool = False @agents.router.post("/send_message") async def send_message( session_id: str = Body(), prompt: str = Body(), mode: Optional[str] = Body(default=None), model: Optional[str] = Body(default=None), images: Optional[List[str]] = Body(default=None), image_media_types: Optional[List[str]] = Body(default=None), context_paths: Optional[List[dict]] = Body(default=None), forced_tools: Optional[List[str]] = Body(default=None), attached_skills: Optional[List[dict]] = Body(default=None), hidden: bool = Body(default=False), ) -> dict: agent: Agent = get_agent(session_id) mode_changed: bool = bool(mode and mode != agent.mode) model_changed: bool = bool(model and model != agent.model) if mode_changed: agent.mode = mode # type: ignore[assignment] if model_changed: agent.model = model # type: ignore[assignment] if mode_changed or model_changed: resolved_mode_config: ResolvedModeConfig = await ResolvedModeConfig.create( mode_id=agent.mode, session_prompt=agent.config.system_prompt, toolkit=agent.toolkit, ) agent.config.system_prompt = resolved_mode_config.system_prompt agent.config.model = agent.model agent.config.allowed_tools = resolved_mode_config.allowed_tools agent.config.disallowed_tools = resolved_mode_config.disallowed_tools if resolved_mode_config.cwd: agent.config.cwd = resolved_mode_config.cwd msg: UserMessage = UserMessage( content=prompt, branch_id=agent.branch_id, images=images or [], image_media_types=image_media_types or [], context_paths=context_paths or [], attached_skills=attached_skills or [], forced_tools=forced_tools or [], hidden=hidden, ) await agent.send_message(msg) return {"ok": True} @agents.router.post("/stop_agent") async def stop_agent(session_id: str = Body()) -> dict: agent: Agent = get_agent(session_id) await agent.stop_agent() return {"ok": True} @agents.router.post("/handle_approval") async def handle_approval( request_id: str = Body(), behavior: str = Body(), message: str = Body(default=""), updated_input: Optional[dict] = Body(default=None), ) -> dict: await COMMS_MANAGER.resolve_approval( request_id=request_id, behavior=behavior, message=message, updated_input=updated_input, ) return {"ok": True} # --------------------------------------------------------------------------- # Branching # --------------------------------------------------------------------------- @agents.router.post("/edit_message") async def edit_message( session_id: str = Body(), message_id: str = Body(), content: str = Body(), ) -> dict: agent: Agent = get_agent(session_id) await agent.stop_agent() fork: Agent = agent.branch(message_id) SESSIONS[fork.session_id] = fork edited_msg: UserMessage = UserMessage(content=content, branch_id=fork.branch_id) await fork.send_message(edited_msg) return {"ok": True, "branch_id": fork.branch_id, "session_id": fork.session_id} @agents.router.post("/switch_branch") async def switch_branch( session_id: str = Body(), branch_id: str = Body(), ) -> dict: agent: Agent = get_agent(session_id) agent.branch_id = branch_id await agent.emit(BranchSwitchedEvent( session_id=session_id, active_branch_id=branch_id, )) return {"ok": True} # --------------------------------------------------------------------------- # Persistence # --------------------------------------------------------------------------- @agents.router.post("/close_session") async def close_session(session_id: str = Body()) -> dict: agent: Optional[Agent] = SESSIONS.pop(session_id, None) if not agent: raise HTTPException(status_code=404, detail="Session not found") await agent.stop_agent() AGENT_STORE.save(agent) msgs = agent.messages.messages closed_at: str = msgs[-1].timestamp.isoformat() if msgs else datetime.now().isoformat() await agent.emit(AgentClosedEvent( session_id=session_id, status=agent.status, closed_at=closed_at, )) return {"ok": True} @agents.router.post("/resume_session") async def resume_session(session_id: str = Body()) -> dict: if session_id in SESSIONS: return {"session": SESSIONS[session_id].model_dump(mode="json")} agent: Optional[Agent] = AGENT_STORE.load_or_none(session_id) if not agent: raise HTTPException(status_code=404, detail="Session not found in history") agent.status = "stopped" agent.on_event = COMMS_MANAGER.make_session_emitter(agent.session_id) agent.on_done = _persist_agent agent.toolkit = await build_agent_toolkit( agent=agent, sessions=SESSIONS, comms_manager=COMMS_MANAGER, ) SESSIONS[agent.session_id] = agent AGENT_STORE.delete(session_id) await agent.emit(AgentStatusEvent( session_id=session_id, status=agent.status, session=agent.snapshot(), )) return {"session": agent.snapshot().model_dump(mode="json")} @agents.router.post("/duplicate_session") async def duplicate_session(session_id: str = Body()) -> dict: source: Optional[Agent] = SESSIONS.get(session_id) if source is None: source = AGENT_STORE.load_or_none(session_id) if not source: raise HTTPException(status_code=404, detail="Session not found") clone: Agent = source.model_copy(deep=True) clone.session_id = uuid4().hex clone.status = "stopped" clone.task = None clone.lock = asyncio.Lock() clone.pending_approvals = [] clone.sub_agents = [] clone.on_event = COMMS_MANAGER.make_session_emitter(clone.session_id) clone.on_done = _persist_agent clone.toolkit = await build_agent_toolkit( agent=clone, sessions=SESSIONS, comms_manager=COMMS_MANAGER, ) SESSIONS[clone.session_id] = clone await clone.emit(AgentStatusEvent( session_id=clone.session_id, status=clone.status, session=clone.snapshot(), )) return {"session": clone.snapshot().model_dump(mode="json")} @agents.router.get("/get_history") async def get_history( q: str = Body(default=""), limit: int = Body(default=20), offset: int = Body(default=0), ) -> dict: all_agents: List[Agent] = AGENT_STORE.load_all() all_agents.sort( key=lambda a: a.messages.messages[-1].timestamp if a.messages.messages else datetime.min, reverse=True, ) q_lower: str = q.strip().lower() history: List[dict] = [] for agent in all_agents: if q_lower: search_text: str = build_search_text(agent).lower() if q_lower not in search_text: continue msgs = agent.messages.messages closed_at: str = msgs[-1].timestamp.isoformat() if msgs else "" history.append({ "id": agent.session_id, "status": agent.status, "model": agent.model, "mode": agent.mode, "closed_at": closed_at, }) total: int = len(history) page: List[dict] = history[offset : offset + limit] return {"sessions": page, "total": total, "has_more": offset + limit < total}