From b228b5cf72e3003a4d148f9ff05e8c25956ce0b3 Mon Sep 17 00:00:00 2001 From: haikdc Date: Wed, 1 Apr 2026 01:59:41 -0700 Subject: [PATCH] [Haik]: ckpt, started working on an abstraction for Agent. Starting to understand the slop, honestly im not mad, not even disapointed, im just kinda sad tbh. Ima go to sleep and hopefully wake up less sad :( I think tmr ima just pull an all nigher to deslop the backend entirely. --- backend/apps/agents/README.md | 225 +++++++++++++ backend/apps/agents/browser/README.md | 240 ++++++++++++++ backend/apps/agents/execution/README.md | 253 ++++++++++++++ backend/apps/agents/manager/HaikFix/Agent.py | 197 +++++++++++ .../agents/manager/HaikFix/PromptChunks.py | 39 +++ .../apps/agents/manager/HaikFix/agent_loop.py | 106 ++++++ .../agents/manager/HaikFix/helpers/Message.py | 53 +++ .../helpers/handle_assistant_message.py | 32 ++ .../HaikFix/helpers/handle_stream_event.py | 0 backend/apps/agents/manager/README.md | 313 ++++++++++++++++++ 10 files changed, 1458 insertions(+) create mode 100644 backend/apps/agents/README.md create mode 100644 backend/apps/agents/browser/README.md create mode 100644 backend/apps/agents/execution/README.md create mode 100644 backend/apps/agents/manager/HaikFix/Agent.py create mode 100644 backend/apps/agents/manager/HaikFix/PromptChunks.py create mode 100644 backend/apps/agents/manager/HaikFix/agent_loop.py create mode 100644 backend/apps/agents/manager/HaikFix/helpers/Message.py create mode 100644 backend/apps/agents/manager/HaikFix/helpers/handle_assistant_message.py create mode 100644 backend/apps/agents/manager/HaikFix/helpers/handle_stream_event.py create mode 100644 backend/apps/agents/manager/README.md diff --git a/backend/apps/agents/README.md b/backend/apps/agents/README.md new file mode 100644 index 00000000..91ba20ad --- /dev/null +++ b/backend/apps/agents/README.md @@ -0,0 +1,225 @@ +# Agents Package + +The `agents/` package is the core of OpenSwarm's AI agent system. It manages the full lifecycle of Claude-powered agent sessions — from launching and configuring agents, through real-time streaming conversations, to persistence and history. + +## Architecture Overview + +``` +┌─────────────────────────────────────────────────────────────────┐ +│ Frontend (Browser) │ +│ REST API calls ↕ WebSocket events ↕ │ +├─────────────────────────────────────────────────────────────────┤ +│ │ +│ agents.py ─── REST endpoints ──┐ │ +│ ws_routes.py ─ WS dispatch ────┤ │ +│ ▼ │ +│ ┌─────────────────┐ │ +│ │ AgentManager │ (singleton facade) │ +│ │ agent_manager │ │ +│ └────────┬────────┘ │ +│ │ │ +│ ┌──────────────────┼──────────────────┐ │ +│ ▼ ▼ ▼ │ +│ ┌────────────────┐ ┌──────────────┐ ┌─────────────────┐ │ +│ │ manager/ │ │ execution/ │ │ browser/ │ │ +│ │ │ │ │ │ │ │ +│ │ Session store │ │ Agent loop │ │ Browser agent │ │ +│ │ WS manager │ │ SDK hooks │ │ runner + tools │ │ +│ │ Operations │ │ Prompts │ │ MCP server │ │ +│ │ Meta/LLM calls │ │ MCP config │ │ │ │ +│ │ Persistence │ │ Approval │ │ │ │ +│ └────────────────┘ └──────────────┘ └─────────────────┘ │ +│ │ +│ models.py ─── Shared Pydantic data models │ +└─────────────────────────────────────────────────────────────────┘ +``` + +**Data flow for a typical user message:** + +1. Frontend sends a REST `POST /sessions/{id}/message` or a WebSocket `agent:send_message` event +2. `agents.py` / `ws_routes.py` delegates to `agent_manager.send_message()` +3. `AgentManager` creates a `Message`, emits it via WebSocket, spawns `run_agent_loop()` as an async task +4. `run_agent_loop()` (in `execution/`) builds the prompt, configures MCP servers, creates SDK hooks, then streams the Claude Agent SDK `query()` call +5. Streaming events (text deltas, tool calls, results) are emitted in real-time via `ws_manager` +6. Tool calls go through the permission/approval system (`agent_hooks.py` → `approval.py`) +7. On completion, the session is persisted to disk and analytics are recorded + +## Directory Structure + +``` +agents/ +├── README.md # This file +├── __init__.py # Empty package marker +├── agents.py # FastAPI sub-app + all REST endpoints +├── models.py # Pydantic models (AgentSession, Message, etc.) +├── ws_routes.py # WebSocket event dispatch +│ +├── execution/ # Agent runtime engine +│ ├── README.md # Detailed docs for execution/ +│ ├── __init__.py +│ ├── agent_loop.py # Main Claude SDK query loop + streaming +│ ├── agent_hooks.py # SDK permission/lifecycle hook factories +│ ├── agent_options.py # ClaudeAgentOptions builder +│ ├── agent_mock.py # Session-completed analytics +│ ├── approval.py # Human-in-the-loop approval flow +│ ├── mcp_builder.py # MCP server config + tool policies +│ ├── prompt_builder.py # Prompt composition helpers +│ ├── prompt_context.py # Context builders (tools, browser, files) +│ └── invoke_agent_mcp_server.py # Stdio MCP server for InvokeAgent +│ +├── manager/ # Session management + WebSocket infra +│ ├── README.md # Detailed docs for manager/ +│ ├── agent_manager.py # Central AgentManager singleton +│ ├── agent_manager_ops.py # Complex ops (edit, close, resume, etc.) +│ ├── agent_manager_meta.py # LLM metadata, persistence, deletion +│ ├── session_store.py # On-disk JSON persistence + history +│ └── ws_manager.py # WebSocket ConnectionManager singleton +│ +└── browser/ # Browser automation sub-agents + ├── README.md # Detailed docs for browser/ + ├── __init__.py # Re-exports run_browser_agent(s) + ├── schemas.py # Browser tool definitions + system prompt + ├── executor.py # Tool execution bridge to frontend + ├── runner.py # Core browser agent loop + ├── browser_agent_mcp_schemas.py # MCP delegation tool schemas + └── browser_agent_mcp_server.py # Stdio MCP server for browser delegation +``` + +## Top-Level Files + +### `agents.py` — REST API Surface + +The FastAPI sub-application. Defines ~20 REST endpoints that form the entire HTTP API for agent management. Every endpoint delegates to the `agent_manager` singleton. + +| Method | Route | Purpose | +|--------|-------|---------| +| GET | `/sessions` | List active sessions (optionally by dashboard) | +| GET | `/sessions/{id}` | Get a single session | +| POST | `/launch` | Launch a new agent from an `AgentConfig` | +| POST | `/sessions/{id}/message` | Send a user message (with optional mode/model/images/tools) | +| POST | `/sessions/{id}/stop` | Stop a running agent | +| POST | `/approval` | Handle tool approval decision | +| POST | `/sessions/{id}/edit_message` | Edit a message (triggers branching) | +| POST | `/sessions/{id}/switch_branch` | Switch active conversation branch | +| POST | `/sessions/{id}/generate-title` | AI-generate a session title | +| POST | `/sessions/{id}/generate-group-meta` | AI-generate tool group name + icon | +| PATCH | `/sessions/{id}` | Partial update (name, system prompt) | +| POST | `/sessions/{id}/duplicate` | Deep-copy a session | +| POST | `/sessions/{id}/close` | Close and persist a session | +| DELETE | `/sessions/{id}` | Permanently delete a session | +| GET | `/history` | Search/paginate closed session history | +| GET | `/sessions/{id}/browser-agents` | Get child browser-agent sessions | +| POST | `/sessions/{id}/resume` | Resume a closed session | +| POST | `/browser-agent/run` | Run browser sub-agents | +| POST | `/invoke-agent/run` | Fork and invoke an agent session | + +Also defines a **lifespan** context manager that on startup reconciles stale sessions and restores persisted ones, and on shutdown stops all agents and persists state. + +### `models.py` — Shared Data Models + +Pydantic models used across the entire package: + +| Model | Purpose | +|-------|---------| +| `AgentConfig` | Launch configuration (model, mode, tools, system prompt, target directory, dashboard) | +| `AgentSession` | Full session state — status, messages, branches, cost, tokens, approvals, metadata | +| `Message` | Conversation message with role, content, branching info, attachments | +| `MessageBranch` | Branch metadata (parent branch, fork point) | +| `ApprovalRequest` | Pending tool approval sent to user | +| `ApprovalResponse` | User's allow/deny decision | +| `ToolGroupMeta` | AI-generated name + SVG icon for tool call groups | + +**Defaults:** Model is `"sonnet"`, provider is `"anthropic"`, mode is `"agent"`, default tools are `[Read, Edit, Write, Bash, Glob, Grep, AskUserQuestion]`. + +**Session status flow:** +``` +launched → running → completed + → stopped (user cancelled) + → error + → waiting_approval → running (after decision) +``` + +### `ws_routes.py` — WebSocket Dispatch + +Thin event router that handles WebSocket messages from the frontend. Two handlers: + +- **`handle_session_message`** — Per-session events: `agent:send_message`, `agent:approval_response`, `agent:edit_message`, `agent:stop` +- **`handle_dashboard_message`** — Dashboard-level events: `agent:approval_response`, `browser:result` + +Each event is dispatched to the appropriate `agent_manager` or `ws_manager` method. + +## Key Concepts + +### Conversation Branching + +When a user edits a message, the system creates a new `MessageBranch` forking from the edit point. Messages are linked via `parent_id` and `branch_id`. The active branch can be switched to navigate between conversation paths. + +### Human-in-the-Loop (HITL) Approval + +Tools can have three permission policies: `always_allow`, `deny`, or `ask`. When a tool with `ask` policy is invoked, the system sends an approval request to the frontend via WebSocket, waits for the user's decision (with a 10-minute timeout), and then allows or denies the tool execution. + +### MCP (Model Context Protocol) Servers + +The agent system uses MCP servers to extend tool capabilities: +- **User tools** — External MCP servers configured by the user (with OAuth2 support) +- **Browser agent MCP** — Stdio subprocess exposing `CreateBrowserAgent`, `BrowserAgent`, `BrowserAgents` +- **Invoke agent MCP** — Stdio subprocess exposing `InvokeAgent` for cross-session invocation + +### Session Persistence + +Sessions are persisted as JSON files via `SessionStore`. On shutdown, all active sessions are saved. On startup, persisted sessions are restored to memory and the disk files are removed. Closed sessions remain on disk for history/search. + +### Browser Sub-Agents + +Browser agents are autonomous agents that control browser tabs in the frontend via a WebSocket bridge. They can screenshot, click, type, scroll, navigate, and evaluate JavaScript. The main agent can delegate browser tasks via MCP tools. + +## Dependency Graph + +``` +agents.py ──────────────────────► agent_manager (singleton) +ws_routes.py ───────────────────► agent_manager, ws_manager + +agent_manager + ├── execution/agent_loop.py (run_agent_loop) + ├── execution/prompt_builder.py (resolve_mode) + ├── execution/mcp_builder.py (get_all_tool_names) + ├── manager/ws_manager.py (emit events) + ├── manager/session_store.py (persistence) + ├── manager/agent_manager_ops.py (edit, close, resume, duplicate, invoke) + └── manager/agent_manager_meta.py (title gen, group meta, persist/restore) + +execution/agent_loop.py + ├── prompt_builder.py (build_prompt_content) + ├── agent_hooks.py (create_sdk_hooks) + ├── agent_options.py (build_agent_options) + └── Claude Agent SDK (query, streaming) + +browser/runner.py + ├── browser/executor.py (execute_browser_tool) + ├── browser/schemas.py (tool defs, system prompt) + ├── Anthropic API (direct, not SDK) + └── ws_manager (real-time comms) +``` + +## Design Notes + +- **250-line file limit** — Files are deliberately kept under ~250 lines. Complex logic is split across multiple files (e.g., `agent_manager.py` delegates to `agent_manager_ops.py` and `agent_manager_meta.py`). +- **Singleton pattern** — `agent_manager` and `ws_manager` are module-level singletons, imported directly by consumers. +- **Stateless functions** — Most logic is in standalone functions that receive data as parameters rather than relying on class state, making testing easier. +- **Separation of concerns** — Prompt building, MCP configuration, hook creation, and the query loop are each in their own module within `execution/`. +- **Two browser tool layers** — `browser/schemas.py` defines the low-level tools the browser agent uses internally (Screenshot, Click, etc.), while `browser_agent_mcp_schemas.py` defines the high-level delegation tools the main agent uses to spawn browser agents. + +## External Dependencies + +| Dependency | Used For | +|------------|----------| +| `claude_agent_sdk` | Agent query loop, streaming, tool hooks | +| Anthropic API | Browser agent loop (direct API calls) | +| FastAPI | REST endpoints, WebSocket handling | +| Pydantic | Data models and validation | +| PIL (optional) | Screenshot compression in browser MCP server | + +See the sub-package READMEs for detailed per-file documentation: +- [execution/README.md](execution/README.md) — Agent runtime engine +- [manager/README.md](manager/README.md) — Session management and WebSocket infrastructure +- [browser/README.md](browser/README.md) — Browser automation sub-agents diff --git a/backend/apps/agents/browser/README.md b/backend/apps/agents/browser/README.md new file mode 100644 index 00000000..f19c9f11 --- /dev/null +++ b/backend/apps/agents/browser/README.md @@ -0,0 +1,240 @@ +# browser/ — Browser Automation Sub-Agents + +This package implements autonomous browser agents that can control browser tabs in the frontend. The main agent can delegate browser tasks (navigate, click, type, screenshot, etc.) and these sub-agents execute them independently. + +## Architecture + +There are **two layers** of tools here, which is important to understand: + +``` +┌──────────────────────────────────────────────────────────────────┐ +│ Main Agent (Claude SDK) │ +│ │ +│ Uses DELEGATION tools (MCP): │ +│ CreateBrowserAgent — spin up a new browser + assign a task │ +│ BrowserAgent — assign a task to an existing browser │ +│ BrowserAgents — parallel tasks on multiple browsers │ +│ │ +│ These are defined in browser_agent_mcp_schemas.py │ +│ and served by browser_agent_mcp_server.py (stdio subprocess) │ +└───────────────────────────────────┬──────────────────────────────┘ + │ HTTP POST to /browser-agent/run + ▼ +┌──────────────────────────────────────────────────────────────────┐ +│ Browser Agent (Anthropic API) │ +│ │ +│ Uses EXECUTION tools (direct): │ +│ BrowserScreenshot — capture current page │ +│ BrowserGetText — get visible text content │ +│ BrowserNavigate — go to a URL │ +│ BrowserClick — click an element by CSS selector │ +│ BrowserType — type text into an element │ +│ BrowserEvaluate — run JavaScript on the page │ +│ BrowserGetElements — query elements by selector │ +│ BrowserScroll — scroll the page │ +│ BrowserWait — wait for a specified duration │ +│ │ +│ These are defined in schemas.py │ +│ and executed by executor.py (via WebSocket to frontend) │ +└───────────────────────────────────┬──────────────────────────────┘ + │ ws_manager.send_browser_command + ▼ +┌──────────────────────────────────────────────────────────────────┐ +│ Frontend Browser Iframe │ +│ │ +│ Receives WebSocket commands, executes in the actual browser, │ +│ and returns results (screenshots, text, element lists) │ +└──────────────────────────────────────────────────────────────────┘ +``` + +## Files + +### `schemas.py` — Browser Tool Definitions + System Prompt + +Pure data file with no imports. Defines everything the browser agent needs to operate. + +**`BROWSER_TOOLS_SCHEMA`** — List of 9 Anthropic-compatible tool definitions: + +| Tool | Parameters | Description | +|------|-----------|-------------| +| `BrowserScreenshot` | (none) | Capture a screenshot of the current page | +| `BrowserGetText` | (none) | Get all visible text content from the page | +| `BrowserNavigate` | `url` | Navigate to a URL | +| `BrowserClick` | `selector` | Click an element by CSS selector | +| `BrowserType` | `selector`, `text` | Type text into an input element | +| `BrowserEvaluate` | `expression` | Execute JavaScript and return the result | +| `BrowserGetElements` | `selector` | Query DOM elements by CSS selector | +| `BrowserScroll` | `direction` (up/down), `amount` (pixels) | Scroll the page | +| `BrowserWait` | `duration` (ms) | Wait for a specified duration | + +**`ACTION_MAP`** — Maps tool names to short action strings for the WebSocket protocol: +``` +BrowserScreenshot → screenshot BrowserClick → click +BrowserGetText → get_text BrowserType → type +BrowserNavigate → navigate BrowserEvaluate → evaluate +BrowserGetElements→ get_elements BrowserScroll → scroll +BrowserWait → wait +``` + +**`SYSTEM_PROMPT`** — Multi-paragraph instructions for the browser agent, including: +- Always screenshot first to see the current state +- Wait 2-3 seconds after navigation before screenshots +- Use `BrowserGetElements` before clicking to find correct selectors +- Don't get stuck in loops — try alternative approaches +- Provide clear summaries of what was accomplished + +**`MAX_TURNS`** — `25` (maximum LLM turns per browser agent run) + +--- + +### `executor.py` — Tool Execution Bridge + +Bridges between the browser agent's tool calls and the actual browser in the frontend. + +**`execute_browser_tool(tool_name, tool_input, browser_id, tab_id="")`** (async) +1. Looks up the action string from `ACTION_MAP` +2. Sends the command to the frontend via `ws_manager.send_browser_command()` +3. Waits up to 30 seconds for the frontend to return a result +4. Returns the raw result dict + +**`_format_tool_result(result, tool_name)`** +- Converts raw browser results into Anthropic content blocks +- Special case for `BrowserScreenshot`: returns an image content block with base64 PNG +- Other tools: returns text content blocks + +**`_request_browser_approval(session, tool_name, tool_input)`** (async) +- Wraps the generic `request_approval()` from `execution/approval.py` +- Uses browser-specific defaults: 300s timeout, analytics tracking disabled + +--- + +### `runner.py` — Core Browser Agent Loop + +The main engine that runs browser agents. Uses the Anthropic API directly (not the Claude Agent SDK). + +**`run_browser_agent(task, browser_id, model, dashboard_id?, tab_id?, pre_selected?, initial_url?, parent_session_id?)`** (async) + +Full lifecycle of a single browser agent: + +1. **Setup** — Creates an `AgentSession` in `"browser-agent"` mode with the parent session ID +2. **Initial navigation** — If `initial_url` is provided, navigates and takes an initial screenshot +3. **Agent loop** (up to `MAX_TURNS`): + a. Calls the Anthropic API with the conversation history + browser tools + b. For each tool call in the response: + - Checks builtin permissions for approval requirements + - Requests approval if needed (via `_request_browser_approval`) + - Executes the tool via `execute_browser_tool` + - Formats the result and appends to conversation + - Logs the action for the action log + c. If no tool calls → agent is done (the response is the summary) + d. If cancelled → stop early +4. **Completion** — Takes a final screenshot, sets status, emits via WebSocket +5. **Returns** `{session_id, browser_id, summary, action_log, final_screenshot}` + +**Error handling:** +- API errors → logged, session status set to `error` +- Cancellation → session status set to `stopped` +- Always emits final status via WebSocket + +**`_create_browser_card(dashboard_id, url, parent_session_id?)`** (async) +- Creates a new browser card on the dashboard +- Adds a `BrowserTab` with the given URL +- Positions the card in the layout +- Persists the dashboard and broadcasts `dashboard:browser_card_added` +- Returns the new `browser_id` + +**`run_browser_agents(tasks, model, dashboard_id?, pre_selected_browser_ids?, parent_session_id?)`** (async) +- Runs multiple browser agents in parallel via `asyncio.gather` +- For tasks without a `browser_id`: calls `_create_browser_card` first +- For tasks with `pre_selected_browser_ids`: assigns available pre-selected IDs +- Records `browser_agent.batch_completed` analytics +- Returns list of result dicts + +--- + +### `browser_agent_mcp_schemas.py` — MCP Delegation Tool Schemas + +Pure data file. Defines the 3 high-level tools the main agent uses to delegate browser work. + +| Tool | Parameters | Description | +|------|-----------|-------------| +| `CreateBrowserAgent` | `task`, `initial_url` | Create a new browser card and run a task on it | +| `BrowserAgent` | `browser_id`, `task`, `tab_id?` | Run a task on an existing browser card | +| `BrowserAgents` | `tasks[]` (each with `browser_id?`, `task`, `tab_id?`, `initial_url?`) | Run multiple browser tasks in parallel | + +These are the tools listed by the MCP server when the SDK calls `tools/list`. + +--- + +### `browser_agent_mcp_server.py` — Stdio MCP Server + +A standalone script launched as a subprocess by the Claude Agent SDK. It implements the JSON-RPC MCP protocol and proxies browser agent requests to the OpenSwarm backend via HTTP. + +**Lifecycle:** +1. The Claude Agent SDK starts this process with stdin/stdout pipes +2. It receives `initialize` → responds with server info and capabilities +3. It receives `tools/list` → returns the 3 delegation tools from `browser_agent_mcp_schemas.py` +4. It receives `tools/call` → dispatches to `handle_tool_call` +5. `handle_tool_call` POSTs to `http://127.0.0.1:{port}/api/agents/browser-agent/run` +6. The backend runs the browser agents and returns results +7. Results are formatted into MCP content blocks and returned to the SDK + +**Environment variables:** +| Variable | Default | Purpose | +|----------|---------|---------| +| `OPENSWARM_PORT` | `8325` | Backend server port | +| `OPENSWARM_AGENT_MODEL` | `sonnet` | Model for browser agents | +| `OPENSWARM_DASHBOARD_ID` | — | Dashboard to create browser cards on | +| `OPENSWARM_PRE_SELECTED_BROWSER_IDS` | — | Comma-separated pre-selected browser IDs | +| `OPENSWARM_PARENT_SESSION_ID` | — | Parent session for child tracking | + +**Screenshot compression:** +- If a screenshot's base64 exceeds 400KB, it's compressed via PIL (if available) +- Resized to max 1024px wide, converted to JPEG at quality 60 +- Falls back gracefully if PIL isn't installed + +**No direct Python imports from the backend** — communicates purely via HTTP. This isolation is necessary because it runs as a separate subprocess. + +--- + +### `__init__.py` — Package Exports + +Re-exports `run_browser_agent` and `run_browser_agents` from `runner.py` for convenient importing: + +```python +from backend.apps.agents.browser import run_browser_agent, run_browser_agents +``` + +## Key Concepts + +### Browser Cards + +Browser cards are UI elements in the dashboard that contain an embedded browser. Each card has: +- A unique `browser_id` +- One or more `BrowserTab` instances (each with a URL) +- A position in the dashboard layout + +Browser agents are always associated with a specific browser card. + +### Parent-Child Sessions + +When the main agent spawns browser agents, the browser agent sessions are linked to the parent via `parent_session_id`. This enables: +- Querying all browser agents for a given session +- Stopping all children when the parent stops +- Tracking sub-agent costs and analytics + +### Two API Paths + +There are two ways browser agents get triggered: + +1. **Via MCP** (agent-initiated): The main agent calls `CreateBrowserAgent` → MCP server → HTTP → `run_browser_agents` → `runner.py` +2. **Via REST** (user-initiated): Direct POST to `/browser-agent/run` → `run_browser_agents` → `runner.py` + +Both paths end up in the same `run_browser_agents` function. + +### Tool Approval in Browser Context + +Browser agents share the same HITL approval system as the main agent, but with: +- Shorter timeout (300s vs 600s) +- Analytics tracking disabled (to avoid double-counting) +- Permission checks against the same builtin permission policies diff --git a/backend/apps/agents/execution/README.md b/backend/apps/agents/execution/README.md new file mode 100644 index 00000000..5461a7aa --- /dev/null +++ b/backend/apps/agents/execution/README.md @@ -0,0 +1,253 @@ +# execution/ — Agent Runtime Engine + +This package contains the runtime core of the agent system: the Claude Agent SDK query loop, prompt assembly, MCP server configuration, SDK hooks, and the human-in-the-loop approval flow. + +## How It All Fits Together + +``` +run_agent_loop() ← entry point (called by AgentManager) + │ + ├─ build_prompt_content() ← prompt_builder.py + │ └─ resolve_context_paths() ← prompt_context.py + │ + ├─ create_sdk_hooks() ← agent_hooks.py + │ └─ request_approval() ← approval.py + │ └─ get_effective_policy() ← mcp_builder.py + │ + ├─ build_agent_options() ← agent_options.py + │ ├─ resolve_mode() ← prompt_builder.py + │ ├─ compose_system_prompt() ← prompt_builder.py + │ ├─ build_connected_tools_context()← prompt_context.py + │ ├─ build_browser_context() ← prompt_context.py + │ └─ build_mcp_servers() ← mcp_builder.py + │ + └─ Claude Agent SDK query() ← streaming loop + ├─ StreamEvent → ws_manager (real-time token streaming) + ├─ AssistantMessage → Messages (text + tool calls) + └─ ResultMessage → cost/tokens (final accounting) +``` + +## Files + +### `agent_loop.py` — Main Query Loop + +The top-level entry point for running an agent. Orchestrates everything else. + +**`run_agent_loop(sessions, session_id, prompt, ...)`** (async) + +1. Builds the user prompt via `build_prompt_content()` — resolves context paths, forced tools, attached skills, and images +2. Creates SDK hooks via `create_sdk_hooks()` — permission checking, tool approval, result formatting +3. Builds the full options dict via `build_agent_options()` — system prompt, MCP servers, tool permissions, API config +4. Creates `ClaudeAgentOptions` and calls `query()` to start the streaming agent loop +5. Iterates the async stream, dispatching to three handlers: + +| Event Type | Handler | What It Does | +|------------|---------|--------------| +| `StreamEvent` | `_handle_stream_event()` | Real-time text/tool streaming deltas → WebSocket | +| `AssistantMessage` | `_handle_assistant_message()` | Extracts text + tool_use blocks, creates Messages, emits via WS | +| `ResultMessage` | `_handle_result_message()` | Captures session ID, cost, token usage | + +6. On completion: sets status to `completed`, persists session, fires analytics +7. On cancellation: sets status to `stopped` +8. On error: sets status to `error`, creates an error message + +--- + +### `agent_hooks.py` — SDK Hook Factories + +Creates the three hook functions the Claude Agent SDK needs for tool execution control. + +**`create_sdk_hooks(session, session_id, sessions, builtin_perms, ...)`** + +Returns a tuple of `(can_use_tool, pre_tool_hook, post_tool_hook)`: + +**`can_use_tool(tool_name, input_data)`** +- Looks up the effective permission policy for the tool +- `always_allow` → auto-approve +- `deny` → auto-reject +- `ask` → triggers HITL approval (except for `AskUserQuestion` which is always allowed) + +**`pre_tool_hook(input_data, tool_use_id)`** +- Enforces `deny` policy by returning a denial result +- For `ask` policy, calls `request_approval()` and blocks until user decides +- Tracks tool start time for duration analytics + +**`post_tool_hook(input_data, tool_use_id)`** +- Calculates elapsed execution time +- Records `tool.executed` analytics (tool name, MCP server, duration, content length) +- Normalizes response content from the SDK +- Creates `tool_result` Message objects and appends to session +- Special handling for `Agent` tool results — creates sub-agent sessions via `_build_sub_agent_session()` +- Emits everything via WebSocket + +**`_build_sub_agent_session(input_data, raw_response, content, session, ...)`** +- Parses the result of an `Agent` tool call +- Creates a child `AgentSession` with the sub-agent's messages +- Stores it in the sessions dict and broadcasts via WebSocket + +--- + +### `agent_options.py` — ClaudeAgentOptions Builder + +Assembles the complete configuration dict passed to the Claude Agent SDK. + +**`build_agent_options(session, builtin_perms, hooks..., fork_session?, selected_browser_ids?)`** (async) + +Builds a kwargs dict containing: + +| Key | Source | Description | +|-----|--------|-------------| +| `system_prompt` | `compose_system_prompt()` | Global + mode + session + tool context + browser context | +| `mcp_servers` | `build_mcp_servers()` | User MCP tools (with OAuth2 refresh) | +| + browser MCP | stdio subprocess | `openswarm-browser-agent` (if not fully denied) | +| + invoke MCP | stdio subprocess | `openswarm-invoke-agent` (if not fully denied) | +| `allowed_tools` | `_compute_tool_permissions()` | Tools with `always_allow` or `ask` policy | +| `disallowed_tools` | `_compute_tool_permissions()` | Tools with `deny` policy | +| `model` | `resolve_model_id()` | Resolved from session model + provider | +| `api_key` or proxy | settings | Direct Anthropic key or 9Router proxy URL | +| `session_id` | session | For SDK resume/fork support | + +**`_compute_tool_permissions(session, builtin_perms, mcp_servers, ...)`** + +Maps permission policies to the SDK's `mcp____` naming convention: +- Builtins → direct policy lookup +- Browser/invoke MCP tools → mapped back to their builtin equivalents +- User MCP tools → per-server per-tool permission lookup + +Special handling: +- `VIEW_BUILDER_SKILL` is injected when mode is `"view-builder"` +- API key vs. 9Router proxy configuration (with `bare` mode and `cc/` model prefix) +- Resume vs. fork session behavior + +--- + +### `prompt_builder.py` — Prompt Composition + +Stateless helpers for building the system prompt and user prompt content. + +**`resolve_mode(mode_id, get_all_tool_names_fn)`** +- Loads mode definition (from the modes package) +- Returns `(allowed_tools, system_prompt, default_folder)` +- Falls back to all tools if mode not found + +**`compose_system_prompt(default_prompt, mode_prompt, session_prompt, tools_ctx?, browser_ctx?)`** +- Joins non-empty prompt fragments with `\n\n` +- Returns `None` if all fragments are empty + +**`resolve_forced_tools(forced_tools, load_all_tools_fn)`** +- Builds XML `` block describing user-selected tools +- Includes tool descriptions, MCP server names, and connected account emails + +**`resolve_attached_skills(attached_skills)`** +- Formats skill attachments as `[Using skill: name]\n\ncontent` + +**`build_prompt_content(prompt, images?, context_paths?, forced_tools?, attached_skills?, ...)`** +- Orchestrator that calls all the above +- If images are present: returns a multimodal content list with base64 image blocks +- Otherwise: returns a plain string + +--- + +### `prompt_context.py` — Context Building + +Generates XML context blocks that get injected into the system or user prompt. + +**`build_connected_tools_context(allowed_tools, ...)`** +- Generates `` XML listing all MCP servers, their status, connected accounts, and available tools +- Also lists installed-but-not-connected tools + +**`build_browser_context(dashboard_id, selected_browser_ids?)`** +- Generates `` XML explaining the browser delegation tools +- Lists user-selected browser cards with IDs, titles, and current URLs + +**`get_pre_selected_browser_ids(dashboard_id)`** +- Returns browser card IDs from the dashboard layout + +**`resolve_context_paths(context_paths)`** +- For each path: reads file contents (up to 512KB) or builds directory tree (depth 4) +- Wraps in `` or `` XML tags + +**`build_dir_tree(root, max_depth=4, prefix="")`** +- Recursive directory listing, skipping dotfiles + +--- + +### `mcp_builder.py` — MCP Server Config + Tool Policies + +Manages MCP server configuration and resolves tool permission policies. + +**Constants:** +- `FULL_TOOLS` — canonical list of ~22 built-in tool names (Read, Edit, Write, Bash, Glob, Grep, AskUserQuestion, WebSearch, WebFetch, NotebookEdit, TodoWrite, EnterPlanMode, ExitPlanMode, EnterWorktree, TaskOutput, TaskStop, CronCreate, CronList, CronDelete, RenderOutput, InvokeAgent, Agent) + +**`build_mcp_servers(allowed_tools)`** (async) +- Iterates enabled MCP tools, filters by allowed list +- Skips fully-denied tools +- Refreshes OAuth2 tokens for Google tools +- Derives MCP server configs +- Returns `{server_name: config}` dict + +**`get_effective_policy(tool_name, builtin_perms)`** +- Resolves policy for any tool name: + - Builtins → direct lookup in `builtin_perms` + - `mcp__openswarm-browser-agent__X` → maps to the builtin browser tool equivalent + - `mcp__openswarm-invoke-agent__X` → maps to the builtin InvokeAgent equivalent + - Other MCP tools → per-server per-tool permission lookup +- Default policy: `"ask"` + +**`get_all_tool_names()`** +- Returns `FULL_TOOLS` (minus denied builtins) + `mcp:` for enabled/connected MCP tools + +Helper functions: `_get_denied_tool_names()`, `_get_all_known_tool_names()`, `_is_fully_denied()` + +--- + +### `approval.py` — Human-in-the-Loop Approval + +Shared approval flow used by both the main agent and browser sub-agents. + +**`request_approval(session, tool_name, tool_input, timeout?, track_analytics?)`** (async) + +1. Creates an `ApprovalRequest` with a unique ID +2. Adds it to the session's `pending_approvals` +3. Sets session status to `waiting_approval`, emits via WebSocket +4. Calls `ws_manager.send_approval_request()` — creates an `asyncio.Future` and waits +5. User's decision resolves the Future (or timeout triggers auto-deny) +6. Records `approval.requested` and `approval.resolved` analytics (with latency) +7. Cleans up, restores status to `running` +8. Returns `{"behavior": "allow"|"deny", "message": ..., "updated_input": ...}` + +--- + +### `agent_mock.py` — Session Completion Analytics + +**`fire_session_completed(session, sessions_dict)`** + +Fires a comprehensive `session.completed` analytics event with: +- Model, provider, mode +- Total cost (USD), token usage (input + output) +- Message count, duration (seconds) +- Final status, tool usage counts +- Session title, first user message +- Sub-agent count and IDs +- Branch count + +--- + +### `invoke_agent_mcp_server.py` — InvokeAgent MCP Server + +A standalone stdio MCP server launched as a subprocess by the Claude Agent SDK. + +**How it works:** +1. The SDK starts this as a child process +2. It reads JSON-RPC messages from stdin +3. For `tools/list`: returns the `InvokeAgent` tool schema +4. For `tools/call`: POSTs to `http://127.0.0.1:{port}/api/agents/invoke-agent/run` +5. The backend forks the source session, runs the agent loop, and returns the response +6. The MCP server formats and returns the result + +**Environment variables it reads:** +- `OPENSWARM_PORT` — backend port (default 8325) +- `OPENSWARM_PARENT_SESSION_ID` — parent session for tracking +- `OPENSWARM_DASHBOARD_ID` — dashboard context + +**No internal Python imports** — communicates with the backend purely via HTTP. This is necessary because it runs as a separate subprocess. diff --git a/backend/apps/agents/manager/HaikFix/Agent.py b/backend/apps/agents/manager/HaikFix/Agent.py new file mode 100644 index 00000000..1507596b --- /dev/null +++ b/backend/apps/agents/manager/HaikFix/Agent.py @@ -0,0 +1,197 @@ +"""Thin coordinator for agent sessions. + +Heavy logic lives in sibling modules: +- agent_manager_ops – edit, close, resume, duplicate, invoke, LLM metadata +- agent_loop – the SDK query loop, streaming, mock agent +- agent_mock – session-completed analytics +- prompt_builder – system-prompt composition & context injection +- mcp_builder – MCP server construction & tool-policy helpers +- session_store – on-disk persistence, history, message copying +""" + +from __future__ import annotations + +import asyncio +import logging +import os +from datetime import datetime +from uuid import uuid4 + +from claude_agent_sdk import ClaudeAgentOptions +from pydantic import BaseModel, InstanceOf +from typing import List, Literal, Optional +from typeguard import typechecked + +from backend.apps.agents.models import AgentSession, Message +from backend.apps.agents.manager.AgentConfig import AgentConfig +from backend.apps.agents.manager.ws_manager import ws_manager +from backend.apps.agents.execution.prompt_builder import resolve_mode +from backend.apps.agents.execution.mcp_builder import get_all_tool_names +from backend.apps.agents.execution.agent_loop import run_agent_loop +from backend.apps.settings.settings import load_settings + + +logger = logging.getLogger(__name__) + +os.environ.setdefault("CLAUDE_CODE_STREAM_CLOSE_TIMEOUT", "3600000") + +class ContextPath(BaseModel): + path: str + type: Literal["file", "directory"] + +class Skill(BaseModel): + name: str + content: str + + +class Agent(BaseModel): + model: str + mode: str + session_id: str + status: Literal["running", "waiting_approval", "completed", "error", "stopped"] + lock: InstanceOf[asyncio.Lock] + config: ClaudeAgentOptions + task: Optional[asyncio.Task] = None + + @typechecked + def __init__( + self, + model: str, + mode: str, + tools: List[str], + effective_cwd: str, + config: ClaudeAgentOptions, + ) -> None: + id: str = uuid4().hex + lock: asyncio.Lock = asyncio.Lock() + super().__init__( + model=model, + mode=mode, + tools=tools, + effective_cwd=effective_cwd, + status="running", + id=id, + task=None, + lock=lock, + config=config, + ) + + async def launch_agent(self, config: AgentConfig) -> AgentSession: + session_id = uuid4().hex + mode_tools, _, mode_folder = resolve_mode(config.mode, get_all_tool_names) + global_settings = load_settings() + effective_cwd = ( + config.target_directory or mode_folder + or global_settings.default_folder or os.path.expanduser("~") + ) + if config.mode in ("view-builder", "skill-builder") and not config.target_directory: + effective_cwd = os.path.join(effective_cwd, session_id) + os.makedirs(effective_cwd, exist_ok=True) + session = AgentSession( + id=session_id, name=config.name, + provider=getattr(config, "provider", "anthropic"), + model=config.model, mode=config.mode, + system_prompt=config.system_prompt, allowed_tools=mode_tools, + max_turns=config.max_turns, cwd=effective_cwd, + dashboard_id=config.dashboard_id, + ) + await ws_manager.emit_status(session_id, "running", session) + return session + + async def send_message( + self, + prompt: str, + images: Optional[list] = None, + ): + async with self.lock: + if self.task is not None and not self.task.done(): + print("[Agent.send_message] Agent is already running") + return + + skill_meta = [{"id": s["id"], "name": s["name"]} for s in (attached_skills or [])] or None + image_meta = [{"data": img["data"], "media_type": img.get("media_type", "image/png")} for img in (images or [])] or None + user_msg = Message( + role="user", content=prompt, branch_id=session.active_branch_id, + context_paths=context_paths or None, attached_skills=skill_meta, + forced_tools=forced_tools or None, images=image_meta, hidden=hidden, + ) + session.messages.append(user_msg) + await ws_manager.emit_message(session_id, user_msg) + + session.status = "running" + await ws_manager.emit_status(session_id, "running", session) + task = asyncio.create_task(run_agent_loop( + self.sessions, session_id, prompt, images=images, + context_paths=context_paths, forced_tools=forced_tools, + attached_skills=attached_skills, selected_browser_ids=selected_browser_ids, + )) + self.tasks[session_id] = task + + async def send_message_old( + self, session_id: str, prompt: str, + mode: str | None = None, model: str | None = None, + provider: str | None = None, images: list | None = None, + context_paths: list | None = None, forced_tools: list[str] | None = None, + attached_skills: list | None = None, hidden: bool = False, + selected_browser_ids: list[str] | None = None, + ): + session = self.sessions.get(session_id) + if not session: + raise ValueError(f"Session {session_id} not found") + existing = self.tasks.get(session_id) + if existing and not existing.done(): + return + + session_changed = False + if model and model != session.model: + session.model = model + session_changed = True + if mode and mode != session.mode: + session.mode = mode + mode_tools, _, _ = resolve_mode(mode, get_all_tool_names) + session.allowed_tools = mode_tools + session_changed = True + if session_changed: + await ws_manager.emit_status(session_id, session.status, session) + + skill_meta = [{"id": s["id"], "name": s["name"]} for s in (attached_skills or [])] or None + image_meta = [{"data": img["data"], "media_type": img.get("media_type", "image/png")} for img in (images or [])] or None + user_msg = Message( + role="user", content=prompt, branch_id=session.active_branch_id, + context_paths=context_paths or None, attached_skills=skill_meta, + forced_tools=forced_tools or None, images=image_meta, hidden=hidden, + ) + session.messages.append(user_msg) + await ws_manager.emit_message(session_id, user_msg) + + session.status = "running" + await ws_manager.emit_status(session_id, "running", session) + task = asyncio.create_task(run_agent_loop( + self.sessions, session_id, prompt, images=images, + context_paths=context_paths, forced_tools=forced_tools, + attached_skills=attached_skills, selected_browser_ids=selected_browser_ids, + )) + self.tasks[session_id] = task + + async def stop_agent(self, session_id: str): + task = self.tasks.get(session_id) + if task and not task.done(): + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + session = self.sessions.get(session_id) + if session: + for req in list(session.pending_approvals): + ws_manager.resolve_approval(req.id, {"behavior": "deny", "message": "Agent stopped"}) + session.pending_approvals = [] + if hasattr(session, '_cancel_event'): + session._cancel_event.set() + session.status = "stopped" + if not session.closed_at: + session.closed_at = datetime.now() + await ws_manager.emit_status(session_id, "stopped", session) + 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) diff --git a/backend/apps/agents/manager/HaikFix/PromptChunks.py b/backend/apps/agents/manager/HaikFix/PromptChunks.py new file mode 100644 index 00000000..07a5ff26 --- /dev/null +++ b/backend/apps/agents/manager/HaikFix/PromptChunks.py @@ -0,0 +1,39 @@ +from typing import Dict, Literal +from pydantic import BaseModel +from typeguard import typechecked + +TextChunkDict = Dict[Literal["type", "text"], str] +class TextChunk(BaseModel): + type: str = "text" + text: str + + @typechecked + def __init__(self, text: str) -> None: + self.text = text + + @typechecked + def to_dict(self) -> TextChunkDict: + return { + "type": self.type, + "text": self.text, + } + + +ImageChunkDict = Dict[Literal["type", "data", "media_type"], str] +class ImageChunk(BaseModel): + type: str = "base64" + data: str + media_type: str + + @typechecked + def __init__(self, data: str, media_type: str) -> None: + self.data = data + self.media_type = media_type + + @typechecked + def to_dict(self) -> ImageChunkDict: + return { + "type": self.type, + "data": self.data, + "media_type": self.media_type, + } \ No newline at end of file diff --git a/backend/apps/agents/manager/HaikFix/agent_loop.py b/backend/apps/agents/manager/HaikFix/agent_loop.py new file mode 100644 index 00000000..e7d4f90d --- /dev/null +++ b/backend/apps/agents/manager/HaikFix/agent_loop.py @@ -0,0 +1,106 @@ +"""Main agent loop — orchestrates the Claude Agent SDK query loop. + +Heavy logic is delegated to sibling modules: +- agent_mock – mock-agent fallback, streaming helpers, session analytics +- agent_hooks – SDK hook factories (approval, permissions, post-tool) +- agent_options – MCP server construction & ClaudeAgentOptions building +""" + +from __future__ import annotations + +import asyncio +import logging +from uuid import uuid4 + +from typeguard import typechecked + +from backend.apps.agents.models import AgentSession, Message +from backend.apps.agents.manager.ws_manager import ws_manager +from backend.apps.agents.manager.session_store import save_session +from backend.apps.agents.execution.prompt_builder import build_prompt_content +from backend.apps.tools_lib.tools_lib import ( + _load_all as load_all_tools, + load_builtin_permissions, +) +from backend.apps.analytics.collector import record as _analytics +from backend.apps.agents.execution.agent_hooks import create_sdk_hooks + +from claude_agent_sdk import ( + query, ClaudeAgentOptions, AssistantMessage, ResultMessage, +) +from claude_agent_sdk.types import ( + PermissionResultAllow, PermissionResultDeny, + TextBlock, ToolUseBlock, StreamEvent, SystemMessage, +) +from backend.apps.agents.execution.agent_options import build_agent_options + +from backend.apps.agents.manager.HaikFix.PromptChunks import ImageChunk, ImageChunkDict, TextChunk, TextChunkDict +from typing import List, Dict, Literal, Any, Union, Optional + +logger = logging.getLogger(__name__) + + +@typechecked +def build_image_prompt_content(prompt: str, images: List[ImageChunk]) -> List[TextChunk | ImageChunk]: + content: List[Union[ImageChunkDict, TextChunkDict]] = [TextChunk(text=prompt).to_dict()] + for img in images: + content.append(img.to_dict()) + return content + + +PromptMsgDict = Dict[ + Literal["type", "message"], + Dict[ + Literal["role", "content"], + List[ + Union[ImageChunkDict, TextChunkDict] + ] + ] + ] + +@typechecked +def build_prompt_msg(prompt: str, images: Optional[List[ImageChunk]]) -> PromptMsgDict: + content = build_image_prompt_content(prompt, images) + return { + "type": "user", + "message": { + "role": "user", + "content": content + } + } + +async def run_agent_loop( + prompt: str, + images: list | None = None, + options: ClaudeAgentOptions | None = None, +): + """Run the Claude Agent SDK query loop for a session.""" + + prompt_msg = build_prompt_msg(prompt, images) + + async def prompt_stream(): + yield prompt_msg + + stream_text_msg_id = None + stream_tool_msg_ids_ordered: list[str] = [] + stream_block_index_map: dict[int, str] = {} + _turn_number = 0 + _first_event = True + + async for message in query(prompt=prompt_stream(), options=options): + + if isinstance(message, StreamEvent): + stream_text_msg_id = await _handle_stream_event( + session_id, message.event, + stream_text_msg_id, stream_tool_msg_ids_ordered, stream_block_index_map, + ) + + elif isinstance(message, AssistantMessage): + stream_text_msg_id, stream_tool_msg_ids_ordered, stream_block_index_map = ( + await _handle_assistant_message( + session, session_id, message, stream_text_msg_id, + stream_tool_msg_ids_ordered, _turn_number, + TextBlock, ToolUseBlock, + ) + ) + _turn_number += 1 \ No newline at end of file diff --git a/backend/apps/agents/manager/HaikFix/helpers/Message.py b/backend/apps/agents/manager/HaikFix/helpers/Message.py new file mode 100644 index 00000000..556f515b --- /dev/null +++ b/backend/apps/agents/manager/HaikFix/helpers/Message.py @@ -0,0 +1,53 @@ +from pydantic import BaseModel, Field +from typing import Optional, Literal, Union, List, Dict +from datetime import datetime +from uuid import uuid4 + + +######################################################## +# Message Content Types +######################################################## + +class ToolCallContent(BaseModel): + id: str + tool: str + input: dict + +class ToolResultContent(BaseModel): + text: str + tool_name: Optional[str] = None + elapsed_ms: Optional[float] = None + sub_session_id: Optional[str] = None + +MessageContent = Union[str, ToolCallContent, ToolResultContent] + + +######################################################## +# Additional Message Types +######################################################## + +class ContextPath(BaseModel): + path: str + type: Literal["file", "directory"] + +class SkillMeta(BaseModel): + id: str + name: str + # NOTE: content is omitted in backend to save space + +class ImageMeta(BaseModel): + data: str # base64-encoded + media_type: str = "image/png" + +class Message(BaseModel): + id: str = Field(default_factory=lambda: uuid4().hex) + role: Literal["user", "assistant", "tool_call", "tool_result", "system"] + content: MessageContent + timestamp: datetime = Field(default_factory=datetime.now) + branch_id: str = "main" + parent_id: Optional[str] = None + context_paths: Optional[List[ContextPath]] = None + attached_skills: Optional[List[SkillMeta]] = None + forced_tools: Optional[List[str]] = None + images: Optional[List[ImageMeta]] = None + hidden: bool = False diff --git a/backend/apps/agents/manager/HaikFix/helpers/handle_assistant_message.py b/backend/apps/agents/manager/HaikFix/helpers/handle_assistant_message.py new file mode 100644 index 00000000..5258570b --- /dev/null +++ b/backend/apps/agents/manager/HaikFix/helpers/handle_assistant_message.py @@ -0,0 +1,32 @@ +from backend.apps.agents.manager.ws_manager import ws_manager +from claude_agent_sdk.types import ( + PermissionResultAllow, PermissionResultDeny, + TextBlock, ToolUseBlock, StreamEvent, SystemMessage, +) + +async def handle_assistant_message( + session, session_id, message, stream_text_msg_id, + stream_tool_ids +): + content_parts = [] + tool_uses = [] + for block in message.content: + if isinstance(block, TextBlock): + content_parts.append(block.text) + elif isinstance(block, ToolUseBlock): + tool_uses.append({"id": block.id, "tool": block.name, "input": block.input}) + + if content_parts: + asst_msg = Message( + id=stream_text_msg_id or uuid4().hex, + role="assistant", content="\n".join(content_parts), + branch_id=session.active_branch_id, + ) + session.messages.append(asst_msg) + await ws_manager.emit_message(session_id, asst_msg) + + for i, tu in enumerate(tool_uses): + mid = stream_tool_ids[i] if i < len(stream_tool_ids) else uuid4().hex + tool_msg = Message(id=mid, role="tool_call", content=tu, branch_id=session.active_branch_id) + session.messages.append(tool_msg) + await ws_manager.emit_message(session_id, tool_msg) \ No newline at end of file diff --git a/backend/apps/agents/manager/HaikFix/helpers/handle_stream_event.py b/backend/apps/agents/manager/HaikFix/helpers/handle_stream_event.py new file mode 100644 index 00000000..e69de29b diff --git a/backend/apps/agents/manager/README.md b/backend/apps/agents/manager/README.md new file mode 100644 index 00000000..cb86da36 --- /dev/null +++ b/backend/apps/agents/manager/README.md @@ -0,0 +1,313 @@ +# manager/ — Session Management & WebSocket Infrastructure + +This package handles the full lifecycle of agent sessions — creating, running, stopping, editing, branching, persisting, restoring, duplicating, and deleting — plus the WebSocket infrastructure that powers real-time communication with the frontend. + +## Architecture + +``` + ┌──────────────────────────┐ + │ agent_manager.py │ ← singleton facade + │ (AgentManager) │ + └────────────┬─────────────┘ + │ + ┌──────────────────┼──────────────────┐ + ▼ ▼ ▼ + agent_manager_ops.py agent_manager_meta.py ws_manager.py + (edit, close, resume, (title gen, group meta, (WebSocket connections, + duplicate, invoke) persist, restore, HITL futures, + delete) browser bridge) + │ │ + └────────┬─────────┘ + ▼ + session_store.py + (JSON persistence, + history, search) +``` + +**Design principle:** `agent_manager.py` is a thin coordinator — it holds the state dicts and delegates all complex logic to sibling modules. This keeps every file under ~250 lines. + +## Files + +### `agent_manager.py` — Central AgentManager Singleton + +The single entry point consumed by all API routes and WebSocket handlers. Holds two core dicts: + +- **`sessions: dict[str, AgentSession]`** — all active in-memory sessions +- **`tasks: dict[str, asyncio.Task]`** — running agent loop tasks + +**Every method either handles simple logic directly or delegates to a sibling module.** + +| Method | Delegates To | Purpose | +|--------|-------------|---------| +| `launch_agent(config)` | — | Creates session, resolves mode/tools, records analytics, emits WS status | +| `send_message(session_id, prompt, ...)` | `run_agent_loop` | Validates session, handles model/mode switching, creates Message, spawns agent loop task | +| `stop_agent(session_id)` | — | Cancels task, resolves approvals, stops browser children, sets status to `stopped` | +| `handle_approval(request_id, decision)` | `ws_manager` | Resolves pending approval Future | +| `edit_message(...)` | `agent_manager_ops` | Triggers branching and re-run | +| `switch_branch(session_id, branch_id)` | — | Sets `active_branch_id`, emits WS event | +| `generate_title(...)` | `agent_manager_meta` | LLM-powered title generation | +| `generate_group_meta(...)` | `agent_manager_meta` | LLM-powered tool group naming + SVG icon | +| `update_session(session_id, **fields)` | — | Updates `system_prompt` or `name`, emits WS status | +| `close_session(session_id)` | `agent_manager_ops` | Stops children, persists, fires analytics | +| `delete_session(session_id)` | `agent_manager_meta` | Permanent deletion from memory and disk | +| `resume_session(session_id)` | `agent_manager_ops` | Loads from disk, restores to memory | +| `duplicate_session(...)` | `agent_manager_ops` | Deep-copies messages and branches | +| `invoke_agent(...)` | `agent_manager_ops` | Forks session, runs agent loop synchronously | +| `get_all_sessions(dashboard_id?)` | — | Filters in-memory sessions | +| `get_session(session_id)` | — | Dict lookup | +| `get_history(...)` | `session_store` | Paginated, filterable session history | +| `reconcile_on_startup()` | `session_store` | Marks stale running sessions as stopped | +| `persist_all_sessions()` | `agent_manager_meta` | Shutdown persistence | +| `restore_all_sessions()` | `agent_manager_meta` | Startup restore | +| `get_browser_agent_children(...)` | `session_store` | Finds child browser sessions | + +**Exported as:** `agent_manager = AgentManager()` (module-level singleton) + +**Environment:** +- Sets `CLAUDE_CODE_STREAM_CLOSE_TIMEOUT` to 1 hour (3,600,000ms) + +--- + +### `agent_manager_ops.py` — Complex Session Operations + +Implements operations that involve multiple steps (cancellation, branching, persistence, re-execution). + +**`edit_message_op(sessions, tasks, session_id, message_id, new_content)`** +1. Cancels any running agent loop task for the session +2. Creates a new `MessageBranch` forking from the edited message's position +3. Appends a new `Message` with the edited content on the new branch +4. Resets `sdk_session_id` (forces a fresh SDK session) +5. Spawns a new `run_agent_loop` with the edited content +6. Records `session.branched` analytics + +**`close_session_op(sessions, tasks, session_id)`** +1. Stops all child browser-agent sessions (has a known circular import workaround) +2. Cancels the running task +3. Resolves any pending approvals with denial +4. Fires `session.completed` analytics +5. Persists the session to disk via `save_session` +6. Removes from in-memory dicts + +**`resume_session_op(sessions, session_id)`** +1. Checks if session is already in memory (returns it directly if so) +2. Loads from disk via `load_session_data` +3. Records `session.resumed` analytics (with hours since closed) +4. Clears `closed_at`, sets status back to `stopped` +5. Deletes the on-disk file (session is now in memory) +6. Emits WS status + +**`duplicate_session_op(sessions, session_id, dashboard_id?, up_to_message_id?)`** +1. Deep-copies all messages and branches via `copy_session_messages` +2. Creates a new `AgentSession` with `"(copy)"` suffix +3. Emits WS status for the new session + +**`invoke_agent_op(sessions, source_session_id, message, parent_session_id?, dashboard_id?)`** +1. Forks the source session (copies messages/branches) +2. Creates a new session in `"invoked-agent"` mode +3. Appends the new user message +4. Runs the agent loop synchronously (awaits completion) +5. Returns the last assistant response text + cost + +**Known issue:** `close_session_op` has a circular import from `agent_manager` (flagged with a TODO comment in the code). + +--- + +### `agent_manager_meta.py` — LLM Metadata, Persistence, Deletion + +Handles LLM-powered metadata generation and the full persistence lifecycle. + +**`generate_title_op(sessions, session_id, first_prompt)`** +- Calls `quick_llm_call` with a prompt asking for a 3-6 word session title +- Falls back to truncating the first prompt on failure +- Emits `agent:name_updated` via WebSocket + +**`generate_group_meta_op(sessions, session_id, group_id, tool_calls, ...)`** +- Calls `quick_llm_json` asking for a 2-5 word name and 24x24 SVG icon for a group of tool calls +- Stores result as `ToolGroupMeta` on the session +- Supports refinement (regeneration) via `is_refinement` flag +- Emits `agent:group_meta_updated` via WebSocket + +**`persist_all_sessions_op(sessions, tasks)`** — Shutdown hook +1. Iterates all in-memory sessions +2. Stops running sessions, resolves pending approvals +3. Fires `session.completed` analytics for each +4. Serializes to JSON and saves to disk +5. Clears both `sessions` and `tasks` dicts + +**`restore_all_sessions_op(sessions)`** — Startup hook +1. Loads all session data from disk +2. Skips closed or corrupt sessions +3. Resets `"running"` status to `"stopped"` (since the agent loop is no longer active) +4. Clears any stale pending approvals +5. Adds to in-memory `sessions` dict +6. Deletes the disk file (session is now managed in memory) + +**`delete_session_op(manager, session_id)`** +1. Stops child browser-agent sessions +2. Cancels the running task +3. Removes from in-memory dicts +4. Deletes the on-disk file + +--- + +### `ws_manager.py` — WebSocket ConnectionManager + +Manages all WebSocket connections and provides Future-based async bridges for HITL approval and browser commands. + +**Zero internal dependencies** — only uses `fastapi.WebSocket` and stdlib. This makes it the lowest-level component in the dependency graph. + +#### Connection Management + +| Method | Purpose | +|--------|---------| +| `connect_session(session_id, ws)` | Accept and register a per-session WS connection | +| `connect_global(ws)` | Accept and register a dashboard-level WS connection | +| `disconnect_session(session_id, ws)` | Remove a session connection | +| `disconnect_global(ws)` | Remove a global connection | + +#### Message Sending + +**`send_to_session(session_id, event, data)`** +- Broadcasts to ALL connections for that session AND all global connections +- This ensures dashboard-level listeners always see session updates + +**`broadcast_global(event, data)`** +- Sends only to global (dashboard) connections + +#### HITL Approval Bridge + +**`send_approval_request(session_id, request_id, tool_name, tool_input, timeout=600)`** +- Creates an `asyncio.Future` +- Sends the approval request to the frontend +- Awaits the Future with a 10-minute timeout +- Returns the user's decision (or auto-deny on timeout) + +**`resolve_approval(request_id, decision)`** +- Sets the result on the pending Future, unblocking the waiting agent + +#### Browser Command Bridge + +**`send_browser_command(request_id, action, browser_id, params, tab_id?)`** +- Sends a browser command to the frontend via global WS connections +- Waits up to 30 seconds for the frontend to return a result +- Returns error if no dashboard is connected + +**`resolve_browser_command(request_id, result)`** +- Sets the result on the pending Future + +#### Typed Event Emitters + +14 convenience methods that wrap `send_to_session` with specific event types: + +| Emitter | Event Name | Data | +|---------|-----------|------| +| `emit_status` | `agent:status` | Status string + optional full session | +| `emit_message` | `agent:message` | Message object | +| `emit_cost_update` | `agent:cost_update` | Cost in USD | +| `emit_stream_start` | `agent:stream_start` | Message ID, role, optional tool name | +| `emit_stream_delta` | `agent:stream_delta` | Message ID + text delta | +| `emit_stream_end` | `agent:stream_end` | Message ID | +| `emit_branch_created` | `agent:branch_created` | Branch object + active branch ID | +| `emit_branch_switched` | `agent:branch_switched` | Active branch ID | +| `emit_name_updated` | `agent:name_updated` | New session name | +| `emit_group_meta_updated` | `agent:group_meta_updated` | Group ID, name, SVG, is_refined | +| `emit_closed` | `agent:closed` | Full session object | + +**Exported as:** `ws_manager = ConnectionManager()` (module-level singleton) + +--- + +### `session_store.py` — On-Disk Persistence & History + +Wraps the generic `SessionStore` (from `backend.apps.common.json_store`) with agent-specific logic. + +#### Persistence (re-exported from `SessionStore`) + +| Name | Purpose | +|------|---------| +| `save_session(id, data)` | Save session JSON to `SESSIONS_DIR/{id}.json` | +| `load_session_data(id)` | Load a session's JSON from disk | +| `delete_session_file(id)` | Delete a session file | +| `load_all_session_data()` | Load all session files from disk | + +#### Agent-Specific Functions + +**`build_search_text(session, max_len=5000)`** +- Builds a search-indexing string from session name + all user/assistant message text +- Truncated to `max_len` characters +- Used by `get_history` for text search + +**`get_history(q?, limit=20, offset=0, dashboard_id?)`** +- Loads all sessions from disk +- Sorts by `closed_at` descending (most recent first) +- Applies text search filter (case-insensitive against `build_search_text`) +- Applies optional `dashboard_id` filter +- Returns `{sessions: [...], total: N, has_more: bool}` + +**`reconcile_on_startup()`** +- Iterates all on-disk sessions +- Sets any `"running"` or `"waiting_approval"` status to `"stopped"` +- Handles crashes/restarts gracefully + +**`get_browser_agent_children(sessions, parent_session_id)`** +- Finds all browser-agent sessions belonging to a parent +- Checks both in-memory sessions and on-disk data +- Deduplicates by session ID +- Returns list of session summary dicts + +**`copy_session_messages(source, up_to_message_id?)`** +- Deep-copies all messages and branches from a source session +- Generates fresh UUIDs for each message +- Re-maps `parent_id` references to new IDs +- Updates branch `fork_point_message_id` to new IDs +- Returns `(new_messages, new_branches, old_to_new_id_map)` +- Used by duplicate and invoke operations + +## WebSocket Event Flow + +``` +Frontend ws_manager Agent System + │ │ │ + │── connect_session ────────►│ │ + │ │ │ + │ │◄── emit_status("running") ──│ (agent starts) + │◄── agent:status ──────────│ │ + │ │ │ + │ │◄── emit_stream_start ───────│ (LLM streaming) + │◄── agent:stream_start ────│ │ + │ │◄── emit_stream_delta ───────│ + │◄── agent:stream_delta ────│ (repeated) │ + │ │◄── emit_stream_end ─────────│ + │◄── agent:stream_end ──────│ │ + │ │ │ + │ │◄── emit_message ────────────│ (tool call) + │◄── agent:message ─────────│ │ + │ │ │ + │ │◄── send_approval_request ───│ (HITL needed) + │◄── agent:approval_request │ │ + │ │ (user decides) │ + │── approval_response ──────►│ │ + │ │── resolve_approval ─────────►│ (unblocks agent) + │ │ │ + │ │◄── emit_cost_update ────────│ (completion) + │◄── agent:cost_update ─────│ │ + │ │◄── emit_status("completed")─│ + │◄── agent:status ──────────│ │ +``` + +## Persistence Lifecycle + +``` +Server Startup: + reconcile_on_startup() → Fix stale "running" statuses on disk + restore_all_sessions_op() → Load disk sessions into memory, delete disk files + +During Operation: + close_session_op() → Stop, persist to disk, remove from memory + resume_session_op() → Load from disk to memory, delete disk file + +Server Shutdown: + persist_all_sessions_op() → Stop all, save all to disk, clear memory +``` + +Sessions live in memory while active. They move to disk when closed. They move back to memory when resumed. On server restart, disk sessions are loaded back into memory.