openswarm/backend/OLDapps/agents/manager/agent_manager_ops.py

"""Complex agent-manager operations extracted for the 250-line limit.

Each function is a standalone async operation that receives the sessions
dict (and other dependencies) explicitly.
"""

from __future__ import annotations

import asyncio
import logging
from datetime import datetime
from uuid import uuid4

from backend.apps.agents.models import (
    AgentSession, Message, MessageBranch,
)
from backend.apps.agents.manager.ws_manager import ws_manager
from backend.apps.agents.manager.session_store import (
    save_session, load_session_data, delete_session_file,
    build_search_text, copy_session_messages,
)
from backend.apps.analytics.collector import record as _analytics
from backend.apps.agents.execution.agent_loop import run_agent_loop

from backend.apps.agents.execution.agent_loop import run_agent_loop

from backend.apps.agents.execution.agent_mock import fire_session_completed



logger = logging.getLogger(__name__)


async def edit_message_op(
    sessions: dict, tasks: dict,
    session_id: str, message_id: str, new_content: str,
):
    session = sessions.get(session_id)
    if not session:
        raise ValueError(f"Session {session_id} not found")
    existing = tasks.get(session_id)
    if existing and not existing.done():
        existing.cancel()
        try:
            await existing
        except asyncio.CancelledError:
            pass

    target_msg = next((m for m in session.messages if m.id == message_id), None)
    if not target_msg or target_msg.role != "user":
        raise ValueError("Can only edit user messages")

    fork_point_id = message_id
    fork_parent_branch = target_msg.branch_id
    msg_branch = session.branches.get(target_msg.branch_id)
    if msg_branch and msg_branch.fork_point_message_id:
        branch_user_msgs = [m for m in session.messages if m.branch_id == target_msg.branch_id and m.role == "user"]
        if branch_user_msgs and branch_user_msgs[0].id == message_id:
            fork_point_id = msg_branch.fork_point_message_id
            fork_parent_branch = msg_branch.parent_branch_id or "main"

    new_branch_id = uuid4().hex
    new_branch = MessageBranch(id=new_branch_id, parent_branch_id=fork_parent_branch, fork_point_message_id=fork_point_id)
    session.branches[new_branch_id] = new_branch
    session.active_branch_id = new_branch_id
    _analytics("feature.used", {
        "feature": "message.branched",
        "branch_depth": len([b for b in session.branches.values() if b.parent_branch_id]),
        "total_branches_in_session": len(session.branches),
        "messages_before_fork": len([m for m in session.messages if m.branch_id == fork_parent_branch]),
    }, session_id=session_id, dashboard_id=session.dashboard_id)

    edited_msg = Message(
        role="user", content=new_content, branch_id=new_branch_id,
        parent_id=target_msg.parent_id, images=target_msg.images,
        context_paths=target_msg.context_paths, forced_tools=target_msg.forced_tools,
        attached_skills=target_msg.attached_skills,
    )
    session.messages.append(edited_msg)
    await ws_manager.emit_message(session_id, edited_msg)
    await ws_manager.emit_branch_created(session_id, new_branch, new_branch_id)
    session.sdk_session_id = None
    session.status = "running"
    await ws_manager.emit_status(session_id, "running", session)
    task = asyncio.create_task(run_agent_loop(
        sessions, session_id, new_content,
        images=target_msg.images, context_paths=target_msg.context_paths,
        forced_tools=target_msg.forced_tools, attached_skills=target_msg.attached_skills,
    ))
    tasks[session_id] = task


async def close_session_op(
    sessions: dict, tasks: dict,
    session_id: str,
):
    children = [s for s in sessions.values() if s.parent_session_id == session_id and s.mode == "browser-agent"]
    from backend.apps.agents.manager.agent_manager import agent_manager
    # NOTE: this is a circular dependency, must be fixed soon by fixing the ai slop
    for child in children:
        await agent_manager.stop_agent(child.id)
    task = tasks.get(session_id)
    if task and not task.done():
        task.cancel()
        try:
            await task
        except asyncio.CancelledError:
            pass
    session = sessions.get(session_id)
    if not session:
        raise ValueError(f"Session {session_id} not found")
    if session.status in ("running", "waiting_approval"):
        session.status = "stopped"
    session.closed_at = datetime.now()
    for req in list(session.pending_approvals):
        ws_manager.resolve_approval(req.id, {"behavior": "deny", "message": "Session closed"})
    session.pending_approvals = []
    if hasattr(session, '_cancel_event'):
        session._cancel_event.set()
    fire_session_completed(session, sessions)
    doc_data = session.model_dump(mode="json")
    doc_data["search_text"] = build_search_text(session)
    save_session(session_id, doc_data)
    await ws_manager.emit_closed(session_id, session)
    sessions.pop(session_id, None)
    tasks.pop(session_id, None)
    logger.info(f"Session {session_id} closed and persisted")


async def resume_session_op(sessions: dict, session_id: str) -> AgentSession:
    if session_id in sessions:
        return sessions[session_id]
    data = load_session_data(session_id)
    if data is None:
        raise ValueError(f"Session {session_id} not found in history")
    session = AgentSession(**data)
    hours_since = 0
    if data.get("closed_at"):
        try:
            closed = datetime.fromisoformat(data["closed_at"][:19])
            hours_since = round((datetime.now() - closed).total_seconds() / 3600, 1)
        except Exception:
            pass
    _analytics("session.resumed", {
        "hours_since_closed": hours_since,
        "original_message_count": len(data.get("messages", [])),
        "original_cost_usd": data.get("cost_usd", 0), "model": session.model,
    }, session_id=session_id, dashboard_id=session.dashboard_id)
    session.closed_at = None
    sessions[session_id] = session
    delete_session_file(session_id)
    await ws_manager.emit_status(session_id, session.status, session)
    logger.info(f"Session {session_id} resumed from history")
    return session


async def duplicate_session_op(
    sessions: dict, session_id: str,
    dashboard_id: str | None = None, up_to_message_id: str | None = None,
) -> AgentSession:
    source = sessions.get(session_id)
    if not source:
        data = load_session_data(session_id)
        if data is None:
            raise ValueError(f"Session {session_id} not found")
        source = AgentSession(**data)
    new_messages, new_branches, _ = copy_session_messages(source, up_to_message_id)
    new_session = AgentSession(
        id=uuid4().hex, name=f"{source.name} (copy)", status="stopped",
        model=source.model, mode=source.mode, system_prompt=source.system_prompt,
        allowed_tools=list(source.allowed_tools), max_turns=source.max_turns,
        cwd=source.cwd, created_at=datetime.now(), messages=new_messages,
        branches=new_branches, active_branch_id=source.active_branch_id,
        tool_group_meta=dict(source.tool_group_meta),
        dashboard_id=dashboard_id or source.dashboard_id,
    )
    sessions[new_session.id] = new_session
    await ws_manager.emit_status(new_session.id, new_session.status, new_session)
    return new_session


async def invoke_agent_op(
    sessions: dict, source_session_id: str, message: str,
    parent_session_id: str | None = None, dashboard_id: str | None = None,
) -> dict:
    source = sessions.get(source_session_id)
    if not source:
        data = load_session_data(source_session_id)
        if data is None:
            raise ValueError(f"Session {source_session_id} not found")
        source = AgentSession(**data)
    source_name = source.name
    new_messages, new_branches, _ = copy_session_messages(source)
    fork = AgentSession(
        id=uuid4().hex, name=f"{source_name} (invoked)", status="running",
        model=source.model, mode="invoked-agent", sdk_session_id=source.sdk_session_id,
        system_prompt=source.system_prompt, allowed_tools=list(source.allowed_tools),
        max_turns=source.max_turns or 25, cwd=source.cwd, created_at=datetime.now(),
        messages=new_messages, branches=new_branches,
        active_branch_id=source.active_branch_id,
        tool_group_meta=dict(source.tool_group_meta),
        dashboard_id=dashboard_id or source.dashboard_id,
        parent_session_id=parent_session_id,
    )
    sessions[fork.id] = fork
    await ws_manager.broadcast_global("agent:status", {
        "session_id": fork.id, "status": fork.status,
        "session": fork.model_dump(mode="json"),
    })
    user_msg = Message(role="user", content=message, branch_id=fork.active_branch_id)
    fork.messages.append(user_msg)
    await ws_manager.emit_message(fork.id, user_msg)
    await run_agent_loop(sessions, fork.id, message, fork_session=True)
    last_assistant = None
    for msg in reversed(fork.messages):
        if msg.role == "assistant":
            content = msg.content
            if isinstance(content, str):
                last_assistant = content
            elif isinstance(content, list):
                texts = [b.get("text", "") for b in content if isinstance(b, dict) and b.get("type") == "text"]
                last_assistant = "\n".join(texts)
            else:
                last_assistant = str(content)
            break
    return {
        "forked_session_id": fork.id, "source_name": source_name,
        "response": last_assistant or "No response from invoked agent.",
        "cost_usd": fork.cost_usd,
    }