WEBDIENSTE/openswarm
1
0
Fork 0
mirror of https://github.com/openswarm-ai/openswarm.git synced 2026-09-11 12:17:45 +02:00
Code Issues Packages Projects Releases Wiki Activity

[Haik]: ckpt on agents.py abstraction, looking good so far but we can do more

Browse Source
This commit is contained in:
haikdc
2026-04-05 08:37:38 -07:00
parent eca8cccf1a
commit 1cd5eaa496
17 changed files with 195 additions and 198 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/agents/__init__.py
View File
backend/apps/agents/agents.py
+17 -78
View File
@@ -14,7 +14,6 @@ from datetime import datetime
from typing import Optional, List, Dict
from uuid import uuid4
from typeguard import typechecked
from fastapi import HTTPException
from pydantic import BaseModel
@@ -23,14 +22,14 @@ from backend.config.paths import DB_ROOT
from backend.core.Agent.Agent import Agent
from backend.core.db.PydanticStore import PydanticStore
from backend.core.shared_structs.agent.Message.Message import UserMessage
from backend.core.events.events import (
AnyEvent, AgentStatusEvent, AgentClosedEvent, BranchSwitchedEvent,
ApprovalRequestEvent, EventCallback,
)
from backend.apps.agents import ws
from backend.apps.agents.compose_system_prompt import compose_system_prompt
from backend.core.events.events import AgentStatusEvent, AgentClosedEvent, BranchSwitchedEvent
from backend.apps.agents.utils.comms.FutureBridge import APPROVAL_BRIDGE
from backend.apps.agents.utils.agent_utils.compose_system_prompt import compose_system_prompt
from backend.core.tools.make_builtin_toolkit.make_builtin_toolkit import make_builtin_toolkit
from backend.apps.agents.create_sdk_hooks import create_sdk_hooks
from backend.apps.agents.utils.agent_utils.create_sdk_hooks import create_sdk_hooks
from backend.apps.agents.utils.comms_utils.make_session_emitter import make_session_emitter
from backend.apps.agents.utils.comms_utils.send_browser_command import send_browser_command
from backend.apps.agents.utils.comms_utils.build_search_text import build_search_text
from claude_agent_sdk import ClaudeAgentOptions
from claude_agent_sdk.types import HookMatcher, McpServerConfig
from backend.core.tools.shared_structs.Toolkit import Toolkit
@@ -46,57 +45,6 @@ AGENT_STORE: PydanticStore[Agent] = PydanticStore[Agent](
SESSIONS: dict[str, Agent] = {}
@typechecked
def p_make_session_emitter(session_id: str) -> EventCallback:
"""Create an event callback that routes typed events to the WS connection pool.
ApprovalRequestEvents are special-cased: instead of just broadcasting,
the emitter routes through the APPROVAL_BRIDGE and resolves the
embedded future with the user's decision.
"""
async def emit(event: AnyEvent) -> None:
if isinstance(event, ApprovalRequestEvent):
if not ws.has_global_connections():
if not event.future.done():
event.future.set_result({"behavior": "deny", "message": "No dashboard connected for approval."})
return
result = await ws.APPROVAL_BRIDGE.request(
request_id=event.request_id,
send_fn=lambda: ws.send_to_session(session_id, event.event, {
"request_id": event.request_id,
"session_id": event.session_id,
"tool_name": event.tool_name,
"tool_input": event.tool_input,
}),
timeout=600.0,
)
if not event.future.done():
event.future.set_result(result)
return
await ws.send_to_session(session_id, event.event, event.model_dump(mode="json"))
return emit
async def p_send_browser_command(
action: str, browser_id: str, tab_id: str, params: dict,
) -> dict:
"""BrowserCommandFn implementation that routes through the browser FutureBridge."""
request_id: str = uuid4().hex
if not ws.has_global_connections():
return {"error": "No dashboard connected. Open the dashboard to use browser tools."}
return await ws.BROWSER_BRIDGE.request(
request_id=request_id,
send_fn=lambda: ws.broadcast_global("browser:command", {
"request_id": request_id,
"action": action,
"browser_id": browser_id,
"tab_id": tab_id,
"params": params,
}),
timeout=30.0,
)
def get_agent(session_id: str) -> Agent:
agent: Optional[Agent] = SESSIONS.get(session_id)
if not agent:
@@ -113,8 +61,8 @@ async def agents_lifespan():
for stored in AGENT_STORE.load_all():
try:
stored.status = "stopped"
stored.on_event = p_make_session_emitter(stored.session_id)
toolkit: Toolkit = make_builtin_toolkit(stored, SESSIONS, p_send_browser_command)
stored.on_event = make_session_emitter(stored.session_id)
toolkit: Toolkit = make_builtin_toolkit(stored, SESSIONS, send_browser_command)
stored.toolkit = toolkit
SESSIONS[stored.session_id] = stored
except Exception as e:
@@ -163,10 +111,10 @@ async def launch(body: LaunchBody) -> dict:
status="stopped",
config=ClaudeAgentOptions(max_turns=body.max_turns),
)
agent.on_event = p_make_session_emitter(agent.session_id)
agent.on_event = make_session_emitter(agent.session_id)
SESSIONS[agent.session_id] = agent
toolkit: Toolkit = make_builtin_toolkit(agent, SESSIONS, p_send_browser_command)
toolkit: Toolkit = make_builtin_toolkit(agent, SESSIONS, send_browser_command)
agent.toolkit = toolkit
mcp_servers: Dict[str, McpServerConfig] = toolkit.collect_mcp_servers()
allowed_tools, disallowed_tools = toolkit.collect_tool_permissions()
@@ -270,7 +218,7 @@ class ApprovalBody(BaseModel):
@agents.router.post("/approval")
async def handle_approval(body: ApprovalBody) -> dict:
ws.APPROVAL_BRIDGE.resolve(body.request_id, {
APPROVAL_BRIDGE.resolve(body.request_id, {
"behavior": body.behavior,
"message": body.message,
"updated_input": body.updated_input,
@@ -339,8 +287,8 @@ async def resume_session(session_id: str) -> dict:
if not agent:
raise HTTPException(status_code=404, detail="Session not found in history")
agent.status = "stopped"
agent.on_event = p_make_session_emitter(agent.session_id)
toolkit: Toolkit = make_builtin_toolkit(agent, SESSIONS, p_send_browser_command)
agent.on_event = make_session_emitter(agent.session_id)
toolkit: Toolkit = make_builtin_toolkit(agent, SESSIONS, send_browser_command)
agent.toolkit = toolkit
SESSIONS[agent.session_id] = agent
AGENT_STORE.delete(session_id)
@@ -366,8 +314,8 @@ async def duplicate_session(session_id: str, body: dict = {}) -> dict:
clone.lock = asyncio.Lock()
clone.pending_approvals = []
clone.sub_agents = []
clone.on_event = p_make_session_emitter(clone.session_id)
toolkit: Toolkit = make_builtin_toolkit(clone, SESSIONS, p_send_browser_command)
clone.on_event = make_session_emitter(clone.session_id)
toolkit: Toolkit = make_builtin_toolkit(clone, SESSIONS, send_browser_command)
clone.toolkit = toolkit
SESSIONS[clone.session_id] = clone
await clone.emit(AgentStatusEvent(
@@ -377,15 +325,6 @@ async def duplicate_session(session_id: str, body: dict = {}) -> dict:
return {"session": clone.snapshot().model_dump(mode="json")}
@typechecked
def p_build_search_text(agent: Agent, max_len: int = 5000) -> str:
parts: List[str] = []
for msg in agent.messages.messages:
if msg.role in ("user", "assistant") and isinstance(msg.content, str):
parts.append(msg.content)
return " ".join(parts)[:max_len]
@agents.router.get("/history")
async def get_history(q: str = "", limit: int = 20, offset: int = 0, dashboard_id: str = "") -> dict:
all_agents: List[Agent] = AGENT_STORE.load_all()
@@ -398,7 +337,7 @@ async def get_history(q: str = "", limit: int = 20, offset: int = 0, dashboard_i
history: List[dict] = []
for agent in all_agents:
if q_lower:
search_text: str = p_build_search_text(agent).lower()
search_text: str = build_search_text(agent).lower()
if q_lower not in search_text:
continue
msgs = agent.messages.messages
backend/apps/agents/utils/__init__.py
View File
backend/apps/agents/utils/agent_utils/__init__.py
View File
backend/apps/agents/compose_system_prompt.py → backend/apps/agents/utils/agent_utils/compose_system_prompt.py
View File
backend/apps/agents/create_sdk_hooks.py → backend/apps/agents/utils/agent_utils/create_sdk_hooks.py
View File
backend/apps/agents/utils/comms/FutureBridge.py
+42
View File
@@ -0,0 +1,42 @@
import asyncio
from typing import Callable, Awaitable, Dict
from pydantic import BaseModel, Field
from typeguard import typechecked
class FutureBridge(BaseModel):
"""Async request/response bridge over WebSocket.
Pattern: create a Future, send a question to the frontend,
block until the frontend responds (or timeout).
"""
p_pending: Dict[str, asyncio.Future] = Field(default_factory=dict)
# TODO: add better type specing for the output of this function
@typechecked
async def request(
self,
request_id: str,
send_fn: Callable[[], Awaitable[None]],
timeout: float,
) -> dict:
future: asyncio.Future = asyncio.get_event_loop().create_future()
self.p_pending[request_id] = future
await send_fn()
try:
return await asyncio.wait_for(future, timeout=timeout)
except asyncio.TimeoutError:
print(f"[FutureBridge.request] Request {request_id} timed out after {timeout}s")
return {"error": "Timed out"}
finally:
self.p_pending.pop(request_id, None)
# TODO: add better type specing for the input of this function
@typechecked
def resolve(self, request_id: str, result: dict) -> None:
future = self.p_pending.get(request_id)
if future and not future.done():
future.set_result(result)
APPROVAL_BRIDGE = FutureBridge()
BROWSER_BRIDGE = FutureBridge()
backend/apps/agents/utils/comms/__init__.py
View File
backend/apps/agents/utils/comms/ws.py
+63
View File
@@ -0,0 +1,63 @@
import json
from fastapi import WebSocket
from typing import Dict, List
from typeguard import typechecked
P_SESSION_CONNECTIONS: Dict[str, List[WebSocket]] = {}
P_GLOBAL_CONNECTIONS: List[WebSocket] = []
@typechecked
async def connect_session(session_id: str, ws: WebSocket) -> None:
await ws.accept()
P_SESSION_CONNECTIONS.setdefault(session_id, []).append(ws)
@typechecked
async def connect_global(ws: WebSocket) -> None:
await ws.accept()
P_GLOBAL_CONNECTIONS.append(ws)
@typechecked
def disconnect_session(session_id: str, ws: WebSocket) -> None:
conns = P_SESSION_CONNECTIONS.get(session_id)
if not conns:
return
conns[:] = [c for c in conns if c is not ws]
if not conns:
del P_SESSION_CONNECTIONS[session_id]
@typechecked
def has_global_connections() -> bool:
return len(P_GLOBAL_CONNECTIONS) > 0
@typechecked
def disconnect_global(ws: WebSocket) -> None:
P_GLOBAL_CONNECTIONS[:] = [c for c in P_GLOBAL_CONNECTIONS if c is not ws]
@typechecked
async def send_to_session(session_id: str, event: str, data: dict) -> None:
payload = json.dumps({"event": event, "session_id": session_id, "data": data})
for ws in P_SESSION_CONNECTIONS.get(session_id, []):
try:
await ws.send_text(payload)
except Exception:
pass
for ws in P_GLOBAL_CONNECTIONS:
try:
await ws.send_text(payload)
except Exception:
pass
@typechecked
async def broadcast_global(event: str, data: dict) -> None:
payload = json.dumps({"event": event, "data": data})
for ws in P_GLOBAL_CONNECTIONS:
try:
await ws.send_text(payload)
except Exception:
pass
backend/apps/agents/utils/comms_utils/__init__.py
View File
backend/apps/agents/utils/comms_utils/build_search_text.py
+12
View File
@@ -0,0 +1,12 @@
from typeguard import typechecked
from backend.core.Agent.Agent import Agent
from typing import List
@typechecked
def build_search_text(agent: Agent, max_len: int = 5000) -> str:
parts: List[str] = []
for msg in agent.messages.messages:
if msg.role in ("user", "assistant") and isinstance(msg.content, str):
parts.append(msg.content)
return " ".join(parts)[:max_len]
backend/apps/agents/utils/comms_utils/make_session_emitter.py
+34
View File
@@ -0,0 +1,34 @@
from typeguard import typechecked
from backend.core.events.events import AnyEvent, ApprovalRequestEvent, EventCallback
from backend.apps.agents.utils.comms.ws import send_to_session, has_global_connections
from backend.apps.agents.utils.comms.FutureBridge import APPROVAL_BRIDGE
@typechecked
def make_session_emitter(session_id: str) -> EventCallback:
"""Create an event callback that routes typed events to the WS connection pool.
ApprovalRequestEvents are special-cased: instead of just broadcasting,
the emitter routes through the APPROVAL_BRIDGE and resolves the
embedded future with the user's decision.
"""
async def emit(event: AnyEvent) -> None:
if isinstance(event, ApprovalRequestEvent):
if not has_global_connections():
if not event.future.done():
event.future.set_result({"behavior": "deny", "message": "No dashboard connected for approval."})
return
result = await APPROVAL_BRIDGE.request(
request_id=event.request_id,
send_fn=lambda: send_to_session(session_id, event.event, {
"request_id": event.request_id,
"session_id": event.session_id,
"tool_name": event.tool_name,
"tool_input": event.tool_input,
}),
timeout=600.0,
)
if not event.future.done():
event.future.set_result(result)
return
await send_to_session(session_id, event.event, event.model_dump(mode="json"))
return emit
backend/apps/agents/utils/comms_utils/send_browser_command.py
+26
View File
@@ -0,0 +1,26 @@
from typeguard import typechecked
from backend.apps.agents.utils.comms.ws import has_global_connections, broadcast_global
from backend.apps.agents.utils.comms.FutureBridge import BROWSER_BRIDGE
from uuid import uuid4
# TODO: add better type specing for the output of this function
@typechecked
async def send_browser_command(
action: str, browser_id: str, tab_id: str, params: dict,
) -> dict:
"""BrowserCommandFn implementation that routes through the browser FutureBridge."""
request_id: str = uuid4().hex
if not has_global_connections():
return {"error": "No dashboard connected. Open the dashboard to use browser tools."}
return await BROWSER_BRIDGE.request(
request_id=request_id,
send_fn=lambda: broadcast_global("browser:command", {
"request_id": request_id,
"action": action,
"browser_id": browser_id,
"tab_id": tab_id,
"params": params,
}),
timeout=30.0,
)
backend/apps/agents/ws.py
-117
View File
@@ -1,117 +0,0 @@
"""WebSocket connection pool and async future bridges.
Three concerns, one module:
1. Connection pool — holds WebSocket objects, delivers JSON payloads.
2. FutureBridge — generic async request/response over WebSocket.
3. Bridge instances — approval_bridge and browser_bridge.
"""
import asyncio
import json
from typing import Callable, Awaitable
from fastapi import WebSocket
from pydantic import BaseModel, Field
# ---------------------------------------------------------------------------
# 1. Connection pool
# ---------------------------------------------------------------------------
P_SESSION_CONNECTIONS: dict[str, list[WebSocket]] = {}
P_GLOBAL_CONNECTIONS: list[WebSocket] = []
async def connect_session(session_id: str, ws: WebSocket) -> None:
await ws.accept()
P_SESSION_CONNECTIONS.setdefault(session_id, []).append(ws)
async def connect_global(ws: WebSocket) -> None:
await ws.accept()
P_GLOBAL_CONNECTIONS.append(ws)
def disconnect_session(session_id: str, ws: WebSocket) -> None:
conns = P_SESSION_CONNECTIONS.get(session_id)
if not conns:
return
conns[:] = [c for c in conns if c is not ws]
if not conns:
del P_SESSION_CONNECTIONS[session_id]
def has_global_connections() -> bool:
return len(P_GLOBAL_CONNECTIONS) > 0
def disconnect_global(ws: WebSocket) -> None:
P_GLOBAL_CONNECTIONS[:] = [c for c in P_GLOBAL_CONNECTIONS if c is not ws]
async def send_to_session(session_id: str, event: str, data: dict) -> None:
payload = json.dumps({"event": event, "session_id": session_id, "data": data})
for ws in P_SESSION_CONNECTIONS.get(session_id, []):
try:
await ws.send_text(payload)
except Exception:
pass
for ws in P_GLOBAL_CONNECTIONS:
try:
await ws.send_text(payload)
except Exception:
pass
async def broadcast_global(event: str, data: dict) -> None:
payload = json.dumps({"event": event, "data": data})
for ws in P_GLOBAL_CONNECTIONS:
try:
await ws.send_text(payload)
except Exception:
pass
# ---------------------------------------------------------------------------
# 2. FutureBridge
# ---------------------------------------------------------------------------
class FutureBridge(BaseModel):
"""Async request/response bridge over WebSocket.
Pattern: create a Future, send a question to the frontend,
block until the frontend responds (or timeout).
"""
p_pending: dict[str, asyncio.Future] = Field(default_factory=dict)
async def request(
self,
request_id: str,
send_fn: Callable[[], Awaitable[None]],
timeout: float,
) -> dict:
future: asyncio.Future = asyncio.get_event_loop().create_future()
self.p_pending[request_id] = future
await send_fn()
try:
return await asyncio.wait_for(future, timeout=timeout)
except asyncio.TimeoutError:
print(f"[FutureBridge.request] Request {request_id} timed out after {timeout}s")
return {"error": "Timed out"}
finally:
self.p_pending.pop(request_id, None)
def resolve(self, request_id: str, result: dict) -> None:
future = self.p_pending.get(request_id)
if future and not future.done():
future.set_result(result)
# ---------------------------------------------------------------------------
# 3. Bridge instances
# ---------------------------------------------------------------------------
APPROVAL_BRIDGE = FutureBridge()
BROWSER_BRIDGE = FutureBridge()
backend/apps/dashboards/dashboards.py
+1 -3
View File
@@ -8,10 +8,8 @@ from uuid import uuid4
from backend.config.Apps import SubApp
# from backend.apps.common.json_store import JsonStore
from backend.core.db.PydanticStore import PydanticStore
from backend.apps.dashboards.models import (
from backend.apps.dashboards.Dashboard import (
Dashboard,
DashboardCreate,
DashboardUpdate,
DashboardLayout,
)
from backend.apps.common.llm_helpers import _resolve_model as _rm
backend/apps/health/__init__.py
View File
backend/apps/settings/__init__.py
View File