mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-08-20 11:42:23 +02:00
349 lines
24 KiB
Python
349 lines
24 KiB
Python
import asyncio
|
|
import logging
|
|
import time
|
|
import os
|
|
from contextlib import asynccontextmanager
|
|
from typing import AsyncIterator, Dict, List, Optional
|
|
from typeguard import typechecked
|
|
|
|
from backend.apps.agents.core.models import (
|
|
AgentSession, Message,
|
|
)
|
|
from backend.apps.agents.core.ws_manager import ws_manager
|
|
from backend.apps.settings.settings import load_settings
|
|
from backend.apps.tools_lib.tools_lib import load_builtin_permissions
|
|
# SESSIONS_DIR is re-exported on purpose: session_store reads agent_manager.SESSIONS_DIR at call time (dodging a circular import), and the disk-resilience test monkeypatches it here.
|
|
from backend.config.paths import SESSIONS_DIR as SESSIONS_DIR
|
|
from backend.apps.agents.manager.session.session_store import (
|
|
save_session,
|
|
load_session_data as load_session_data,
|
|
)
|
|
from backend.apps.agents.manager.streaming.state import ThinkingState, TurnState
|
|
from backend.apps.agents.manager.streaming.PartialReply import PartialReply
|
|
from backend.apps.agents.manager.session.SessionLifecycle import SessionLifecycle
|
|
from backend.apps.agents.manager.SpawnAgentRun import SpawnAgentRun
|
|
from backend.apps.agents.manager.session.SessionPersistence import SessionPersistence
|
|
from backend.apps.agents.manager.Messaging import Messaging, QueuedMessage
|
|
from backend.apps.agents.manager.SessionControl import SessionControl
|
|
from backend.apps.agents.manager.AgentLaunch import AgentLaunch
|
|
from backend.apps.agents.manager.MockAgent import MockAgent
|
|
from backend.apps.agents.manager.RunSupport import RunSupport
|
|
from backend.apps.agents.manager.run.handle_run_error import handle_run_error
|
|
from backend.apps.agents.manager.run.TurnRunner import TurnRunner
|
|
from backend.apps.agents.manager.run.client_pool import ClientHandle
|
|
from backend.apps.agents.manager.streaming.HookContext import HookContext
|
|
from backend.apps.agents.manager.run.RunOptions import RunOptions
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
os.environ.setdefault("CLAUDE_CODE_STREAM_CLOSE_TIMEOUT", "3600000")
|
|
|
|
# Cap concurrent ROOT agent turns so firing 30 agents at once doesn't spawn 30 CLIs in the same instant; the overflow queues (agents are model/IO-bound, so they're waiting anyway). Env-tunable, 0/blank disables the gate.
|
|
MAX_CONCURRENT_TURNS = int(os.environ.get("OSW_MAX_CONCURRENT_TURNS", "8") or "0")
|
|
|
|
|
|
class AgentManager(SessionLifecycle, SessionPersistence, Messaging, SessionControl, AgentLaunch, SpawnAgentRun, MockAgent, TurnRunner, RunOptions, RunSupport):
|
|
@typechecked
|
|
def __init__(self):
|
|
self.sessions: Dict[str, AgentSession] = {}
|
|
from backend.apps.agents.core.flight_recorder import set_sessions_provider
|
|
set_sessions_provider(lambda: self.sessions)
|
|
self.tasks: Dict[str, asyncio.Task] = {}
|
|
# Live mirror of the in-flight streamed assistant text per session, so a stop can persist the partial reply instantly instead of waiting out the multi-second SDK teardown the cancel handler sits behind.
|
|
self.live_partial: Dict[str, PartialReply] = {}
|
|
# Per-session cancel signal: the loop stashes its asyncio.Event here so a stop/close can set it. Lives on the manager, not the AgentSession model, so it stays out of serialization (an Event can't be model_dump'd).
|
|
self.cancel_events: Dict[str, asyncio.Event] = {}
|
|
# Persistent-client pool (lever A, flag-gated): one live CLI per session, reused across turns.
|
|
self.client_pool: Dict[str, ClientHandle] = {}
|
|
# Per-SESSION hook context + stderr buffer, updated in place each turn: a persistent client's hooks/stderr callback were bound at connect, so they must read stable objects, not per-turn rebuilds.
|
|
self.hook_ctxs: Dict[str, HookContext] = {}
|
|
self.stderr_buffers: Dict[str, List[str]] = {}
|
|
# Messages typed while a turn was live, replayed in order when it ends (see Messaging).
|
|
self.pending_messages: Dict[str, List[QueuedMessage]] = {}
|
|
# Admission gate: one shared semaphore caps concurrent ROOT turns (children bypass). (Re)created per running loop by get_turn_admission so it never binds to a dead loop across a uvicorn reload or a test's asyncio.run.
|
|
self.p_turn_admission_sema: Optional[asyncio.Semaphore] = None
|
|
self.p_turn_admission_loop: Optional[asyncio.AbstractEventLoop] = None
|
|
|
|
|
|
@typechecked
|
|
def get_turn_admission(self) -> asyncio.Semaphore:
|
|
"""The shared admission semaphore for the CURRENT loop; rebuilt if the loop changed so a
|
|
reload/test-run can never await a semaphore bound to a dead loop."""
|
|
loop = asyncio.get_running_loop()
|
|
if self.p_turn_admission_sema is None or self.p_turn_admission_loop is not loop:
|
|
self.p_turn_admission_sema = asyncio.Semaphore(MAX_CONCURRENT_TURNS)
|
|
self.p_turn_admission_loop = loop
|
|
return self.p_turn_admission_sema
|
|
|
|
@asynccontextmanager
|
|
async def turn_admission_slot(self, session: AgentSession, session_id: str) -> AsyncIterator[None]:
|
|
"""Hold one concurrency slot for the duration of a ROOT turn. Overflow turns queue on the
|
|
semaphore (emitting agent:queued, then agent:admitted when they start). Two bypasses, both
|
|
load-bearing: (1) MAX_CONCURRENT_TURNS<=0 disables the gate entirely (kill switch); (2) a
|
|
CHILD turn (parent_session_id set) is NEVER gated, because a parent holds its own slot while
|
|
awaiting a delegated child, so gating children would deadlock the pool. `async with` release
|
|
is cancellation-safe: a stop while queued never acquired, so it can't over-release."""
|
|
if MAX_CONCURRENT_TURNS <= 0 or session.parent_session_id is not None:
|
|
yield
|
|
return
|
|
sema = self.get_turn_admission()
|
|
was_queued = sema.locked()
|
|
if was_queued:
|
|
try:
|
|
await ws_manager.send_to_session(session_id, "agent:queued", {"session_id": session_id})
|
|
except Exception:
|
|
pass
|
|
async with sema:
|
|
if was_queued:
|
|
try:
|
|
await ws_manager.send_to_session(session_id, "agent:admitted", {"session_id": session_id})
|
|
except Exception:
|
|
pass
|
|
yield
|
|
|
|
@typechecked
|
|
async def prewarm_client(self, session_id: str) -> None:
|
|
"""Spawn the session's CLI in the seconds between create and the first message, so the first
|
|
turn's acquire is a pool hit instead of a 0.6-1.6s cold connect. Best-effort: any failure
|
|
just means the first turn pays the connect it always paid. Kill switch OSW_PREWARM_CLI=0."""
|
|
if os.environ.get("OSW_PREWARM_CLI", "1") == "0":
|
|
return
|
|
session = self.sessions.get(session_id)
|
|
if not session or session.messages:
|
|
return
|
|
try:
|
|
import claude_agent_sdk # noqa: F401
|
|
except ImportError:
|
|
return
|
|
try:
|
|
from backend.apps.agents.providers.registry import (
|
|
resolve_model_id_for_sdk as p_resolve,
|
|
get_api_type as p_api_of,
|
|
)
|
|
p_router_model_id = p_resolve(session.model, load_settings())
|
|
p_api_type = p_api_of(session.model)
|
|
builtin_perms = load_builtin_permissions()
|
|
# Representative-LENGTH prompt: thinking derives from prompt length (<50 chars forces it
|
|
# off), so an empty prewarm prompt would boot a different thinking config than a typical
|
|
# first message and fingerprint-miss into a respawn. 50+ chars matches the common case.
|
|
p_representative = "prewarm placeholder prompt of representative length for boot"
|
|
(options, options_kwargs, _, _, _) = await self.build_agent_options(
|
|
session, session_id, p_representative, "", builtin_perms,
|
|
None, None, None, False, p_router_model_id, p_api_type)
|
|
from claude_agent_sdk import ClaudeSDKClient
|
|
from backend.apps.agents.manager.run.client_pool import acquire_client, boot_fingerprint
|
|
|
|
async def p_connect():
|
|
p_client = ClaudeSDKClient(options=options)
|
|
logger.info(f"[SPAWN-PHASE] prewarm-connect start session={session_id[:8]} t={time.monotonic():.3f}")
|
|
await p_client.connect()
|
|
logger.info(f"[SPAWN-PHASE] prewarm-connect done session={session_id[:8]} t={time.monotonic():.3f}")
|
|
return p_client
|
|
|
|
fp = boot_fingerprint(options_kwargs, session)
|
|
await acquire_client(self.client_pool, session_id, fp, p_connect)
|
|
# Deleted mid-connect: the late-arriving client just pooled into a dead session; nothing else will ever dispose it.
|
|
if session_id not in self.sessions:
|
|
from backend.apps.agents.manager.run.client_pool import dispose_client
|
|
await dispose_client(self.client_pool, session_id)
|
|
except Exception:
|
|
logger.info("[client-pool] prewarm skipped for %s", session_id[:8], exc_info=True)
|
|
|
|
@typechecked
|
|
async def run_agent_loop(self, session_id: str, prompt: str, images: Optional[List] = None, context_paths: Optional[List] = None, forced_tools: Optional[List[str]] = None, attached_skills: Optional[List] = None, fork_session: bool = False, selected_browser_ids: Optional[List[str]] = None, selected_app_output_ids: Optional[List[str]] = None, selected_setting_ids: Optional[List[str]] = None, context_valve_retry: bool = False):
|
|
"""Run the Claude Agent SDK query loop for a session."""
|
|
session = self.sessions.get(session_id)
|
|
if not session:
|
|
return
|
|
|
|
from backend.apps.agents.providers.registry import get_api_type as p_get_api_type
|
|
p_api = p_get_api_type(session.model)
|
|
prompt_content = self.build_prompt_content(
|
|
prompt, images, context_paths, forced_tools, attached_skills,
|
|
api_type=p_api, model=session.model,
|
|
)
|
|
|
|
try:
|
|
# SDK presence check: fall to mock mode here, before the options build, so a missing SDK is a clean mock run, not an error card. The real use is in run_options / turn_runner (lazy-imported there).
|
|
import claude_agent_sdk # noqa: F401
|
|
except ImportError:
|
|
logger.warning("claude_agent_sdk not installed, running in mock mode")
|
|
await self.run_mock_agent(session_id, prompt)
|
|
return
|
|
|
|
session.status = "running"
|
|
|
|
# Resolve the model id now so every closure (approval hook, tool executed handler, etc.) has both the short name and the 9Router-prefixed id available without re-resolving. The short name is what the user sees; the router id is what 9Router reports its per-model counters under.
|
|
from backend.apps.agents.providers.registry import (
|
|
resolve_model_id_for_sdk as p_resolve_model_id_early,
|
|
get_api_type as p_get_api_type_early,
|
|
)
|
|
p_router_model_id = p_resolve_model_id_early(session.model, load_settings())
|
|
p_api_type_for_session = p_get_api_type_early(session.model)
|
|
|
|
builtin_perms = load_builtin_permissions()
|
|
|
|
# Builtins default to always_allow (frictionless); path_gate still force-prompts on catastrophic patterns (rm -rf), OS-scheduling, and sensitive paths, so poisoned-email -> destructive-command is still caught. Flip Bash to "ask" in the UI for a prompt on every command. Bind turn + stderr first: build_agent_options can raise early (no provider) and the except hands both to handle_run_error.
|
|
turn = TurnState()
|
|
p_stderr_buffer: List[str] = []
|
|
# Read BEFORE build_agent_options consumes these flags: a fresh-session/fork request must force the persistent client to respawn (same branch id would otherwise fingerprint-match a client still holding the old transcript).
|
|
p_force_respawn = bool(session.needs_fresh_session or session.needs_fork or fork_session)
|
|
try:
|
|
logger.info(f"[SPAWN-PHASE] run-loop start session={session_id[:8]} t={time.monotonic():.3f}")
|
|
(options, options_kwargs, prompt_content, p_stderr_buffer,
|
|
global_settings) = await self.build_agent_options(
|
|
session, session_id, prompt, prompt_content, builtin_perms,
|
|
selected_browser_ids, selected_app_output_ids, selected_setting_ids,
|
|
fork_session, p_router_model_id, p_api_type_for_session)
|
|
resolved_model = p_router_model_id
|
|
api_type = p_api_type_for_session
|
|
|
|
thinking = ThinkingState()
|
|
# Gate the CLI turn (spawn + stream) behind the admission slot so a burst can't run every turn at once; the slot is held ONLY for run_turn_with_retry, so the context-valve retry below re-acquires cleanly instead of nesting.
|
|
logger.info(f"[SPAWN-PHASE] admission-wait session={session_id[:8]} t={time.monotonic():.3f}")
|
|
async with self.turn_admission_slot(session, session_id):
|
|
logger.info(f"[SPAWN-PHASE] admitted session={session_id[:8]} t={time.monotonic():.3f}")
|
|
await self.run_turn_with_retry(
|
|
session, session_id, prompt_content, options, options_kwargs,
|
|
turn, thinking, p_stderr_buffer, resolved_model, api_type, global_settings,
|
|
force_respawn=p_force_respawn,
|
|
)
|
|
session.status = "completed"
|
|
|
|
# Silent-quit seal: a turn that ran tools and ended with no visible answer gets ONE hidden continue nudge (dispatched by the auto-continuation block below); a second silent quit in the same ask surfaces as-is rather than looping.
|
|
try:
|
|
from backend.apps.agents.manager.run.empty_finish import maybe_nudge_empty_finish
|
|
maybe_nudge_empty_finish(session, session_id)
|
|
except Exception:
|
|
logger.exception("empty-finish detection failed; continuing")
|
|
|
|
# Auto-continuation hook (Phase 3). If MCPActivate (or any analogous flow) flagged pending_continuation during this turn, kick off a follow-up turn immediately with the captured prompt. We dispatch as a fire-and-forget task so the current run_agent_loop frame can unwind cleanly before the next turn's options + history rebuild kicks in. The follow-up is `hidden=True` so it doesn't add a user bubble to the visible chat; the model sees it as a synthetic prompt to keep working.
|
|
try:
|
|
if getattr(session, "pending_continuation", False):
|
|
p_continuation_prompt = session.pending_continuation_prompt or "Continue."
|
|
session.pending_continuation = False
|
|
session.pending_continuation_prompt = None
|
|
asyncio.create_task(self.send_message(
|
|
session_id,
|
|
p_continuation_prompt,
|
|
hidden=True,
|
|
))
|
|
logger.info(f"Auto-continuing session {session_id} with hidden prompt")
|
|
except Exception:
|
|
logger.exception("auto-continuation dispatch failed")
|
|
except asyncio.CancelledError:
|
|
# Only act if we're still the session's live task. A user stop pops this task (stop_agent already finalized status + partial), and a follow-up message may have started a newer turn; either way this dying task must NOT clobber the live status or pop the new turn's in-flight partial mirror.
|
|
if self.tasks.get(session_id) is asyncio.current_task():
|
|
session.status = "stopped"
|
|
# A cancelled turn desyncs the CLI's resume transcript from session.messages (the SDK never recorded the interrupted turn), so force the next turn to rebuild history from session.messages, else resume/follow-ups replay a transcript with no trace of the stopped reply ("nothing to continue").
|
|
session.needs_fresh_session = True
|
|
# Persist whatever streamed before the cancel (edit / branch switch paths; the user-stop path already did this in stop_agent).
|
|
await self.commit_partial_now(session)
|
|
turn.stream_text_msg_id = None
|
|
turn.stream_text_accum = ""
|
|
except Exception as e:
|
|
from backend.apps.agents.core.error_classify import is_context_overflow_error, is_context_pressure_death
|
|
p_stderr_tail = "\n".join(p_stderr_buffer[-50:])
|
|
p_overflow = is_context_overflow_error(e, extra_text=p_stderr_tail)
|
|
if not context_valve_retry and (p_overflow or is_context_pressure_death(e, turn.compact_boundaries, extra_text=p_stderr_tail)):
|
|
# Pressure-release valve, two entry shapes: the CLI compacted this turn and still died (autocompact thrash), or the provider rejected the query outright as over the context window. Either way the CLI's resume transcript is beyond saving, but ours isn't: rebuild from the local mirror via the proven fresh-session recap path and transparently re-run the turn ONCE.
|
|
logger.warning(
|
|
f"Agent {session_id}: {'context overflow' if p_overflow else 'context-pressure death'} after "
|
|
f"{turn.compact_boundaries} compact boundaries; one fresh-session recap retry"
|
|
)
|
|
# The recap rebuild trims at compacted_through_msg_id; an overflow can hit before the proactive threshold ever fired, so force a cutoff or the rebuilt prompt is full history again.
|
|
self.maybe_compact(session, force=True)
|
|
session.needs_fresh_session = True
|
|
if turn.stream_text_msg_id:
|
|
await ws_manager.send_to_session(session_id, "agent:stream_end", {
|
|
"session_id": session_id,
|
|
"message_id": turn.stream_text_msg_id,
|
|
})
|
|
for p_tool_msg_id in turn.stream_tool_msg_ids_ordered:
|
|
await ws_manager.send_to_session(session_id, "agent:stream_end", {
|
|
"session_id": session_id,
|
|
"message_id": p_tool_msg_id,
|
|
})
|
|
self.live_partial.pop(session_id, None)
|
|
# Tell the user we self-healed instead of retrying in silence: the frontend renders this as a muted transient pill (same language as the rate-limit pill), not an error card.
|
|
try:
|
|
await ws_manager.send_to_session(session_id, "agent:context_recovered", {
|
|
"session_id": session_id,
|
|
})
|
|
except Exception:
|
|
logger.debug("context_recovered broadcast failed", exc_info=True)
|
|
try:
|
|
from backend.apps.service.client import submit_diagnostic
|
|
from backend.apps.agents.core.redact_for_telemetry import redact_for_telemetry
|
|
from backend.apps.agents.core import flight_recorder as p_fr
|
|
submit_diagnostic({
|
|
"kind": "context_pressure_valve",
|
|
"trigger": "overflow" if p_overflow else "pressure_death",
|
|
"flight": p_fr.build_envelope(session_id, "context_pressure_valve", "overflow" if p_overflow else "pressure_death", session.model, "stream", turn.compact_boundaries),
|
|
"session_id": session_id,
|
|
"model": session.model,
|
|
"compact_boundaries": turn.compact_boundaries,
|
|
"error_preview": redact_for_telemetry(str(e), limit=300),
|
|
})
|
|
except Exception:
|
|
logger.debug("submit_diagnostic context_pressure_valve failed", exc_info=True)
|
|
await self.run_agent_loop(
|
|
session_id, prompt, images, context_paths, forced_tools,
|
|
attached_skills, fork_session, selected_browser_ids,
|
|
selected_app_output_ids, selected_setting_ids,
|
|
context_valve_retry=True,
|
|
)
|
|
return
|
|
await handle_run_error(e, session, session_id, turn, p_stderr_buffer)
|
|
except BaseException as e:
|
|
# Catch BaseExceptionGroup from anyio task groups (e.g. concurrent CLI crash + pending approval cancellation) so it doesn't escape and kill the uvicorn process.
|
|
logger.exception(f"Agent {session_id} fatal error: {e}")
|
|
# A group's str() names the group, not the cause; unwrap to the real member so a wrapped 429/auth error still gets its friendly card + retry-pill semantics instead of a raw group dump.
|
|
from backend.apps.agents.core.first_real_exception import first_real_exception
|
|
p_real = first_real_exception(e)
|
|
if p_real is not None:
|
|
await handle_run_error(p_real, session, session_id, turn, p_stderr_buffer)
|
|
else:
|
|
session.status = "error"
|
|
error_msg = Message(role="system", content=f"Error: {str(e)}", branch_id=session.active_branch_id)
|
|
session.messages.append(error_msg)
|
|
await ws_manager.send_to_session(session_id, "agent:message", {
|
|
"session_id": session_id,
|
|
"message": error_msg.model_dump(mode="json"),
|
|
})
|
|
finally:
|
|
# Only the session's live task finalizes. A stopped task (popped by stop_agent, which already finalized status + saved) or one superseded by a newer turn must not pop the new turn's partial mirror, broadcast a stale terminal status, or overwrite the snapshot the live turn is writing.
|
|
p_is_live_task = self.tasks.get(session_id) is asyncio.current_task()
|
|
if p_is_live_task:
|
|
self.live_partial.pop(session_id, None)
|
|
if session_id in self.sessions and p_is_live_task:
|
|
# For canvas-launched App Builder sessions, the workspace folder IS the session_id (see launch_agent), so meta.json lives at outputs_workspace/<session_id>/meta.json. Read it and propagate name/description into the Output row before the terminal status fires; without this, the row stays "Untitled App" forever because no React component polls the file on the canvas path. Best-effort, only acts when the row's name is still the default placeholder.
|
|
if session.mode == "view-builder":
|
|
try:
|
|
from backend.apps.outputs.outputs import sync_output_from_meta_json
|
|
from backend.apps.outputs.workspace_io import load_all as load_outputs
|
|
if sync_output_from_meta_json(session_id, fallback_name=session.name):
|
|
# Broadcast the renamed row so the sidebar flips from "Untitled App" to the real name without waiting for the next mount.
|
|
try:
|
|
matching = [o for o in load_outputs() if o.workspace_id == session_id]
|
|
if matching:
|
|
await ws_manager.broadcast_global("agent:output_upserted", {
|
|
"output": matching[0].model_dump(mode="json"),
|
|
})
|
|
except Exception:
|
|
logger.exception("post-sync output_upserted broadcast failed")
|
|
except Exception:
|
|
logger.exception("post-session meta sync failed")
|
|
await ws_manager.send_to_session(session_id, "agent:status", {
|
|
"session_id": session_id,
|
|
"status": session.status,
|
|
"session": session.model_dump(mode="json"),
|
|
})
|
|
try:
|
|
save_session(session_id, session.model_dump(mode="json"))
|
|
except Exception as e:
|
|
logger.warning(f"Failed to snapshot session {session_id}: {e}")
|
|
|
|
|
|
agent_manager = AgentManager()
|