openswarm/backend/apps/HaikFix/agents.py

"""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.

ws_manager is used ONLY in this file — the Agent class and its internals
communicate via the on_event callback, never importing ws_manager directly.
"""

from contextlib import asynccontextmanager
from datetime import datetime
from typeguard import typechecked
from uuid import uuid4
import asyncio

from fastapi import HTTPException
from pydantic import BaseModel
from typing import Optional, List

from backend.config.Apps import SubApp
from backend.apps.HaikFix.Agent.Agent import Agent
from backend.apps.HaikFix.Agent.shared_structs.Message.Message import UserMessage
from backend.apps.HaikFix.Agent.shared_structs.events import (
    AnyEvent, AgentStatusEvent, AgentClosedEvent, BranchSwitchedEvent,
    EventCallback,
)
from backend.apps.HaikFix import session_store
from backend.apps.agents.manager.ws_manager import ws_manager
from claude_agent_sdk import ClaudeAgentOptions


SESSIONS: dict[str, Agent] = {}


@typechecked
def p_make_session_emitter(session_id: str) -> EventCallback:
    """Create an event callback that routes typed events to ws_manager for a session."""
    async def emit(event: AnyEvent) -> None:
        await ws_manager.send_to_session(session_id, event.event, event.model_dump(mode="json"))
    return emit

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():
    await session_store.reconcile_on_startup()
    for sid, data in session_store.load_all():
        if data.get("closed_at") is not None:
            continue
        try:
            data.pop("task", None)
            data.pop("lock", None)
            agent: Agent = Agent(**data)
            agent.status = "stopped"
            agent.on_event = p_make_session_emitter(agent.session_id)
            SESSIONS[agent.session_id] = agent
            session_store.delete(sid)
        except Exception as e:
            print(f"[agents lifespan] Skipping corrupt session {sid}: {e}")
    yield
    for agent in list[Agent](SESSIONS.values()):
        await agent.stop_agent()
        data: dict = agent.model_dump(mode="json")
        data["search_text"] = session_store.build_search_text(data)
        session_store.save(agent.session_id, data)
    SESSIONS.clear()


agents = SubApp("agents", agents_lifespan)


# ---------------------------------------------------------------------------
# Session CRUD
# ---------------------------------------------------------------------------

@agents.router.get("/SESSIONS")
async def list_sessions(dashboard_id: str = "") -> dict:
    result: List[Agent] = list[Agent](SESSIONS.values())
    if dashboard_id:
        result: List[Agent] = [a for a in result if getattr(a, "dashboard_id", None) == dashboard_id]
    return {"SESSIONS": [a.model_dump(mode="json") for a in result]}


@agents.router.get("/sessions/{session_id}")
async def get_session(session_id: str) -> dict:
    return get_agent(session_id).model_dump(mode="json")


class LaunchBody(BaseModel):
    name: str = "New Agent"
    model: str = "sonnet"
    mode: str = "agent"
    system_prompt: str = ""
    max_turns: int = 200
    target_directory: Optional[str] = None
    dashboard_id: Optional[str] = None

@agents.router.post("/launch")
async def launch(body: LaunchBody) -> dict:
    # TODO: build ClaudeAgentOptions from body once prompt/options builder exists
    agent: Agent = Agent(
        model=body.model,
        mode=body.mode,
        status="running",
        config=ClaudeAgentOptions(
            system_prompt=body.system_prompt,
            max_turns=body.max_turns,
        ),
    )
    agent.on_event = p_make_session_emitter(agent.session_id)
    SESSIONS[agent.session_id] = agent
    await agent._emit(AgentStatusEvent(
        session_id=agent.session_id, status="running",
        session=agent.snapshot(),
    ))
    return {"session_id": agent.session_id, "session": agent.snapshot().model_dump(mode="json")}


class UpdateBody(BaseModel):
    name: Optional[str] = None
    system_prompt: Optional[str] = None

@agents.router.patch("/SESSIONS/{session_id}")
async def update_session(session_id: str, body: UpdateBody) -> dict:
    agent: Agent = get_agent(session_id)
    if body.name is not None:
        agent.name = body.name
    if body.system_prompt is not None:
        agent.config.system_prompt = body.system_prompt
    await agent._emit(AgentStatusEvent(
        session_id=session_id, status=agent.status,
        session=agent.snapshot(),
    ))
    return {"ok": True}


@agents.router.delete("/sessions/{session_id}")
async def delete_session(session_id: str) -> dict:
    agent: Optional[Agent] = SESSIONS.pop(session_id, None)
    if agent is not None:
        await agent.stop_agent()
    session_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("/SESSIONS/{session_id}/message")
async def send_message(session_id: str, body: MessageBody) -> dict:
    agent: Agent = get_agent(session_id)
    if body.mode and body.mode != agent.mode:
        agent.mode = body.mode
    if body.model and body.model != agent.model:
        agent.model = body.model

    msg: UserMessage = UserMessage(
        content=body.prompt,
        branch_id=agent.branch_id,
        images=body.images or [],
        image_media_types=body.image_media_types or [],
        context_paths=body.context_paths or [],
        attached_skills=body.attached_skills or [],
        forced_tools=body.forced_tools or [],
        hidden=body.hidden,
    )
    await agent.send_message(msg)
    return {"ok": True}


@agents.router.post("/sessions/{session_id}/stop")
async def stop_agent(session_id: str) -> dict:
    agent: Agent = get_agent(session_id)
    await agent.stop_agent()
    return {"ok": True}


class ApprovalBody(BaseModel):
    request_id: str
    behavior: str
    message: str = ""
    updated_input: Optional[dict] = None

@agents.router.post("/approval")
async def handle_approval(body: ApprovalBody) -> dict:
    ws_manager.resolve_approval(body.request_id, {
        "behavior": body.behavior,
        "message": body.message,
        "updated_input": body.updated_input,
    })
    return {"ok": True}


# ---------------------------------------------------------------------------
# Branching
# ---------------------------------------------------------------------------

class EditMessageBody(BaseModel):
    message_id: str
    content: str

@agents.router.post("/sessions/{session_id}/edit_message")
async def edit_message(session_id: str, body: EditMessageBody) -> dict:
    agent: Agent = get_agent(session_id)
    await agent.stop_agent()
    fork: Agent = agent.branch(body.message_id)
    SESSIONS[fork.session_id] = fork

    edited_msg = UserMessage(content=body.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}


class SwitchBranchBody(BaseModel):
    branch_id: str

@agents.router.post("/SESSIONS/{session_id}/switch_branch")
async def switch_branch(session_id: str, body: SwitchBranchBody) -> dict:
    agent: Agent = get_agent(session_id)
    agent.branch_id = body.branch_id
    await agent.emit(BranchSwitchedEvent(
        session_id=session_id, active_branch_id=body.branch_id,
    ))
    return {"ok": True}


# ---------------------------------------------------------------------------
# Persistence
# ---------------------------------------------------------------------------

@agents.router.post("/sessions/{session_id}/close")
async def close_session(session_id: str) -> 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()
    closed_at: str = datetime.now().isoformat()
    data: dict = agent.model_dump(mode="json")
    data["search_text"] = session_store.build_search_text(data)
    data["closed_at"] = closed_at
    session_store.save(session_id, data)
    await agent.emit(AgentClosedEvent(
        session_id=session_id, status=agent.status,
        closed_at=closed_at,
    ))
    return {"ok": True}


@agents.router.post("/SESSIONS/{session_id}/resume")
async def resume_session(session_id: str) -> dict:
    if session_id in SESSIONS:
        return {"session": SESSIONS[session_id].model_dump(mode="json")}
    data: Optional[dict] = session_store.load(session_id)
    if not data:
        raise HTTPException(status_code=404, detail="Session not found in history")
    data.pop("task", None)
    data.pop("lock", None)
    data.pop("search_text", None)
    data.pop("closed_at", None)
    agent: Agent = Agent(**data)
    agent.status = "stopped"
    agent.on_event = p_make_session_emitter(agent.session_id)
    SESSIONS[agent.session_id] = agent
    session_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("/SESSIONS/{session_id}/duplicate")
async def duplicate_session(session_id: str, body: dict = {}) -> dict:
    source: Optional[Agent] = SESSIONS.get(session_id)
    if source is None:
        data: Optional[dict] = session_store.load(session_id)
        if not data:
            raise HTTPException(status_code=404, detail="Session not found")
        data.pop("task", None)
        data.pop("lock", None)
        source = Agent(**data)

    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 = p_make_session_emitter(clone.session_id)
    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("/history")
async def get_history(q: str = "", limit: int = 20, offset: int = 0, dashboard_id: str = "") -> dict:
    return session_store.get_history(
        q=q, limit=limit, offset=offset,
        dashboard_id=dashboard_id or None,
    )