diff --git a/backend/apps/agents/__init__.py b/backend/apps/agents/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/backend/apps/agents/agents.py b/backend/apps/agents/agents.py index 6a8fdd9b..d1f1e2c7 100644 --- a/backend/apps/agents/agents.py +++ b/backend/apps/agents/agents.py @@ -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 diff --git a/backend/apps/agents/utils/__init__.py b/backend/apps/agents/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/backend/apps/agents/utils/agent_utils/__init__.py b/backend/apps/agents/utils/agent_utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/backend/apps/agents/compose_system_prompt.py b/backend/apps/agents/utils/agent_utils/compose_system_prompt.py similarity index 100% rename from backend/apps/agents/compose_system_prompt.py rename to backend/apps/agents/utils/agent_utils/compose_system_prompt.py diff --git a/backend/apps/agents/create_sdk_hooks.py b/backend/apps/agents/utils/agent_utils/create_sdk_hooks.py similarity index 100% rename from backend/apps/agents/create_sdk_hooks.py rename to backend/apps/agents/utils/agent_utils/create_sdk_hooks.py diff --git a/backend/apps/agents/utils/comms/FutureBridge.py b/backend/apps/agents/utils/comms/FutureBridge.py new file mode 100644 index 00000000..4fba30f6 --- /dev/null +++ b/backend/apps/agents/utils/comms/FutureBridge.py @@ -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() diff --git a/backend/apps/agents/utils/comms/__init__.py b/backend/apps/agents/utils/comms/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/backend/apps/agents/utils/comms/ws.py b/backend/apps/agents/utils/comms/ws.py new file mode 100644 index 00000000..bc5249dd --- /dev/null +++ b/backend/apps/agents/utils/comms/ws.py @@ -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 diff --git a/backend/apps/agents/utils/comms_utils/__init__.py b/backend/apps/agents/utils/comms_utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/backend/apps/agents/utils/comms_utils/build_search_text.py b/backend/apps/agents/utils/comms_utils/build_search_text.py new file mode 100644 index 00000000..84aba573 --- /dev/null +++ b/backend/apps/agents/utils/comms_utils/build_search_text.py @@ -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] diff --git a/backend/apps/agents/utils/comms_utils/make_session_emitter.py b/backend/apps/agents/utils/comms_utils/make_session_emitter.py new file mode 100644 index 00000000..3b1582f9 --- /dev/null +++ b/backend/apps/agents/utils/comms_utils/make_session_emitter.py @@ -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 \ No newline at end of file diff --git a/backend/apps/agents/utils/comms_utils/send_browser_command.py b/backend/apps/agents/utils/comms_utils/send_browser_command.py new file mode 100644 index 00000000..87026c8d --- /dev/null +++ b/backend/apps/agents/utils/comms_utils/send_browser_command.py @@ -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, + ) \ No newline at end of file diff --git a/backend/apps/agents/ws.py b/backend/apps/agents/ws.py deleted file mode 100644 index 603aedf4..00000000 --- a/backend/apps/agents/ws.py +++ /dev/null @@ -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() diff --git a/backend/apps/dashboards/dashboards.py b/backend/apps/dashboards/dashboards.py index a9a42f7a..621f7ff7 100644 --- a/backend/apps/dashboards/dashboards.py +++ b/backend/apps/dashboards/dashboards.py @@ -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 diff --git a/backend/apps/health/__init__.py b/backend/apps/health/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/backend/apps/settings/__init__.py b/backend/apps/settings/__init__.py new file mode 100644 index 00000000..e69de29b