diff --git a/backend/apps/agents/agent_manager.py b/backend/apps/agents/agent_manager.py index 325b80ee..90d13f8e 100644 --- a/backend/apps/agents/agent_manager.py +++ b/backend/apps/agents/agent_manager.py @@ -51,6 +51,7 @@ from backend.apps.agents.manager.builtin_mcp_servers import register_builtin_mcp from backend.apps.agents.manager.provider_env import configure_provider_env from backend.apps.agents.manager.session.SessionLifecycleMixin import SessionLifecycleMixin from backend.apps.agents.manager.MessagingMixin import MessagingMixin +from backend.apps.agents.manager.SessionControlMixin import SessionControlMixin from backend.apps.agents.manager.AgentLaunchMixin import AgentLaunchMixin from backend.apps.agents.manager.RunSupportMixin import RunSupportMixin from backend.apps.agents.manager.permissions import gate_hooks @@ -70,7 +71,7 @@ logger = logging.getLogger(__name__) os.environ.setdefault("CLAUDE_CODE_STREAM_CLOSE_TIMEOUT", "3600000") -class AgentManager(SessionLifecycleMixin, MessagingMixin, AgentLaunchMixin, RunSupportMixin): +class AgentManager(SessionLifecycleMixin, MessagingMixin, SessionControlMixin, AgentLaunchMixin, RunSupportMixin): @typechecked def __init__(self): self.sessions: Dict[str, AgentSession] = {} @@ -592,7 +593,7 @@ class AgentManager(SessionLifecycleMixin, MessagingMixin, AgentLaunchMixin, RunS logger.info(f"[MCP-DEBUG] Creating ClaudeAgentOptions short={session.model} resolved={resolved_model} api_type={api_type}") options = ClaudeAgentOptions(**options_kwargs) - logger.info(f"[MCP-DEBUG] ClaudeAgentOptions created. Starting query...") + logger.info("[MCP-DEBUG] ClaudeAgentOptions created. Starting query...") async def prompt_stream(): yield { diff --git a/backend/apps/agents/manager/MessagingMixin.py b/backend/apps/agents/manager/MessagingMixin.py index 2482ccfe..a82f6762 100644 --- a/backend/apps/agents/manager/MessagingMixin.py +++ b/backend/apps/agents/manager/MessagingMixin.py @@ -1,23 +1,19 @@ -"""User-facing message operations for AgentManager (send / stop / edit / branch / approve / -update), split into a mixin to keep the manager file under the size ceiling. Pure relocation: -self.p_run_agent_loop / self.sessions / self.stop_agent all resolve across the MRO as before.""" +"""Turn-producing message operations for AgentManager (send + edit), the ones that append a +user Message and spawn the agent loop. Session-control ops (stop / approve / branch / update) +live in SessionControlMixin. Pure relocation: self.* resolves across the MRO as before.""" import asyncio import logging -from typing import Dict, List, Optional +from typing import List, Optional from typeguard import typechecked -from datetime import datetime from uuid import uuid4 from backend.apps.agents.core.models import AgentSession, Message, MessageBranch from backend.apps.agents.core.ws_manager import ws_manager from backend.apps.settings.settings import load_settings from backend.apps.agents.manager import browser_dispatch -from backend.apps.agents.manager.session.session_store import ( - load_session_data, - save_session, -) +from backend.apps.agents.manager.session.session_store import load_session_data from backend.apps.agents.manager.session.apply_context_window import apply_context_window from backend.apps.agents.manager.prompt.tool_catalog import get_all_tool_names from backend.apps.agents.manager.prompt.prompt_context import resolve_mode @@ -124,18 +120,7 @@ class MessagingMixin: except Exception: pass - # Track context attachment patterns - if context_paths or attached_skills or images or forced_tools: - pass - - # Track skill usage - for skill in (attached_skills or []): - pass - - # Track first message sophistication is_first_message = sum(1 for m in session.messages if m.role == "user") == 1 - if is_first_message: - pass session.status = "running" await ws_manager.send_to_session(session_id, "agent:status", { @@ -172,65 +157,6 @@ class MessagingMixin: task = asyncio.create_task(self.p_run_agent_loop(session_id, prompt, images=images, context_paths=context_paths, forced_tools=forced_tools, attached_skills=attached_skills, selected_browser_ids=selected_browser_ids, selected_app_output_ids=selected_app_output_ids, selected_setting_ids=selected_setting_ids)) self.tasks[session_id] = task - @typechecked - async def stop_agent(self, session_id: str): - """Stop a running agent and all its browser-agent children.""" - # Stop children first so browser agents get cancelled before parent - 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) - - session = self.sessions.get(session_id) - if session: - # Set cancel event BEFORE cancelling the task so in-flight - # browser agent loops see it immediately - if hasattr(session, '_cancel_event'): - session._cancel_event.set() - - for req in list(session.pending_approvals): - ws_manager.resolve_approval(req.id, {"behavior": "deny", "message": "Agent stopped"}) - session.pending_approvals = [] - - session.status = "stopped" - session.needs_fresh_session = True - if not session.closed_at: - session.closed_at = datetime.now() - # Persist the partial reply NOW, before tearing down the SDK. The - # cancel handler also does this, but it sits behind the generator's - # teardown, which can take several seconds; doing it here means the - # streamed text stays put the instant Stop is pressed instead of - # blinking out and reappearing once teardown finishes. - await self.p_commit_partial_now(session) - await ws_manager.send_to_session(session_id, "agent:status", { - "session_id": session_id, - "status": "stopped", - "session": session.model_dump(mode="json"), - }) - # Snapshot now: the cancelled task's finally skips the save (it's no - # longer the live task once we pop it below), so persist the partial - # here or it'd live only in memory until the next turn / shutdown. - try: - save_session(session_id, session.model_dump(mode="json")) - except Exception: - pass - - # Drop the task from the registry immediately so a follow-up message - # isn't rejected as "still running" while the cancelled task slowly - # tears down (that window was eating user messages). Drain it in the - # background; we've already captured the partial above. - task = self.tasks.pop(session_id, None) - if task and not task.done(): - task.cancel() - asyncio.create_task(self.p_drain_task(task)) - - @typechecked - def handle_approval(self, request_id: str, decision: Dict): - """Resolve a pending HITL approval.""" - ws_manager.resolve_approval(request_id, decision) - @typechecked async def edit_message(self, session_id: str, message_id: str, new_content: str): """Edit a prior user message, creating a new branch (fork).""" @@ -318,37 +244,3 @@ class MessagingMixin: )) self.tasks[session_id] = task - @typechecked - 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 - session.needs_fresh_session = True - await ws_manager.send_to_session(session_id, "agent:branch_switched", { - "session_id": session_id, - "active_branch_id": branch_id, - }) - - @typechecked - async def update_session(self, session_id: str, **fields): - """Update mutable session fields (system_prompt, name).""" - session = self.sessions.get(session_id) - if not session: - raise ValueError(f"Session {session_id} not found") - - allowed = {"system_prompt", "name", "thinking_level"} - for key, value in fields.items(): - if key in allowed: - # Defend against bad thinking_level values - if key == "thinking_level" and value not in ("off", "low", "medium", "high", "auto"): - continue - setattr(session, key, value) - - await ws_manager.send_to_session(session_id, "agent:status", { - "session_id": session_id, - "status": session.status, - "session": session.model_dump(mode="json"), - }) diff --git a/backend/apps/agents/manager/SessionControlMixin.py b/backend/apps/agents/manager/SessionControlMixin.py new file mode 100644 index 00000000..f12d8ac8 --- /dev/null +++ b/backend/apps/agents/manager/SessionControlMixin.py @@ -0,0 +1,111 @@ +"""Session-control operations for AgentManager (stop / approve / switch-branch / update), +split from MessagingMixin so each file stays one responsibility: these control or mutate a +session WITHOUT producing a new agent turn. Pure relocation, self.* resolves across the MRO.""" + +import asyncio +import logging +from datetime import datetime +from typing import Dict + +from typeguard import typechecked + +from backend.apps.agents.core.ws_manager import ws_manager +from backend.apps.agents.manager.session.session_store import save_session + +logger = logging.getLogger(__name__) + + +class SessionControlMixin: + @typechecked + async def stop_agent(self, session_id: str): + """Stop a running agent and all its browser-agent children.""" + # Stop children first so browser agents get cancelled before parent + 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) + + session = self.sessions.get(session_id) + if session: + # Set cancel event BEFORE cancelling the task so in-flight + # browser agent loops see it immediately + if hasattr(session, '_cancel_event'): + session._cancel_event.set() + + for req in list(session.pending_approvals): + ws_manager.resolve_approval(req.id, {"behavior": "deny", "message": "Agent stopped"}) + session.pending_approvals = [] + + session.status = "stopped" + session.needs_fresh_session = True + if not session.closed_at: + session.closed_at = datetime.now() + # Persist the partial reply NOW, before tearing down the SDK. The + # cancel handler also does this, but it sits behind the generator's + # teardown, which can take several seconds; doing it here means the + # streamed text stays put the instant Stop is pressed instead of + # blinking out and reappearing once teardown finishes. + await self.p_commit_partial_now(session) + await ws_manager.send_to_session(session_id, "agent:status", { + "session_id": session_id, + "status": "stopped", + "session": session.model_dump(mode="json"), + }) + # Snapshot now: the cancelled task's finally skips the save (it's no + # longer the live task once we pop it below), so persist the partial + # here or it'd live only in memory until the next turn / shutdown. + try: + save_session(session_id, session.model_dump(mode="json")) + except Exception: + pass + + # Drop the task from the registry immediately so a follow-up message + # isn't rejected as "still running" while the cancelled task slowly + # tears down (that window was eating user messages). Drain it in the + # background; we've already captured the partial above. + task = self.tasks.pop(session_id, None) + if task and not task.done(): + task.cancel() + asyncio.create_task(self.p_drain_task(task)) + + @typechecked + def handle_approval(self, request_id: str, decision: Dict): + """Resolve a pending HITL approval.""" + ws_manager.resolve_approval(request_id, decision) + + @typechecked + 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 + session.needs_fresh_session = True + await ws_manager.send_to_session(session_id, "agent:branch_switched", { + "session_id": session_id, + "active_branch_id": branch_id, + }) + + @typechecked + async def update_session(self, session_id: str, **fields): + """Update mutable session fields (system_prompt, name).""" + session = self.sessions.get(session_id) + if not session: + raise ValueError(f"Session {session_id} not found") + + allowed = {"system_prompt", "name", "thinking_level"} + for key, value in fields.items(): + if key in allowed: + # Defend against bad thinking_level values + if key == "thinking_level" and value not in ("off", "low", "medium", "high", "auto"): + continue + setattr(session, key, value) + + await ws_manager.send_to_session(session_id, "agent:status", { + "session_id": session_id, + "status": session.status, + "session": session.model_dump(mode="json"), + })