mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-09-06 01:37:43 +02:00
[eric] agents: split session-control ops out of MessagingMixin (turn-producing vs control), drop dead analytics stubs
This commit is contained in:
ciregenz
@@ -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 {
|
||||
|
||||
@@ -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"),
|
||||
})
|
||||
|
||||
@@ -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"),
|
||||
})
|
||||
Reference in New Issue
Block a user