"""WebSocket chat endpoint for real-time agent conversations.""" import json import uuid from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query from loguru import logger from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import selectinload from app.core.security import decode_access_token from app.core.permissions import check_agent_access, is_agent_expired from app.database import async_session from app.models.agent import Agent from app.models.audit import ChatMessage from app.models.llm import LLMModel from app.models.user import User router = APIRouter(tags=["websocket"]) class ConnectionManager: """Manage WebSocket connections per agent.""" def __init__(self): # agent_id_str -> list of (WebSocket, session_id_str | None) self.active_connections: dict[str, list[tuple]] = {} async def connect(self, agent_id: str, websocket: WebSocket, session_id: str = None): await websocket.accept() if agent_id not in self.active_connections: self.active_connections[agent_id] = [] self.active_connections[agent_id].append((websocket, session_id)) def disconnect(self, agent_id: str, websocket: WebSocket): if agent_id in self.active_connections: self.active_connections[agent_id] = [ (ws, sid) for ws, sid in self.active_connections[agent_id] if ws != websocket ] async def send_message(self, agent_id: str, message: dict): if agent_id in self.active_connections: for ws, _sid in self.active_connections[agent_id]: try: await ws.send_json(message) except Exception: pass async def send_to_session(self, agent_id: str, session_id: str, message: dict): """Send message only to WebSocket connections matching the given session_id.""" if agent_id in self.active_connections: for ws, sid in self.active_connections[agent_id]: if sid == session_id: try: await ws.send_json(message) except Exception: pass def get_active_session_ids(self, agent_id: str) -> list[str]: """Return distinct session IDs for all active WS connections of an agent.""" if agent_id not in self.active_connections: return [] return list(set(sid for _ws, sid in self.active_connections[agent_id] if sid)) manager = ConnectionManager() from fastapi import Depends from app.core.security import get_current_user from app.database import get_db from app.models.user import User @router.get("/api/chat/{agent_id}/history") async def get_chat_history( agent_id: uuid.UUID, current_user: User = Depends(get_current_user), db: AsyncSession = Depends(get_db), ): """Return web chat message history for this user + agent.""" conv_id = f"web_{current_user.id}" result = await db.execute( select(ChatMessage) .where(ChatMessage.agent_id == agent_id, ChatMessage.conversation_id == conv_id) .order_by(ChatMessage.created_at.asc()) .limit(200) ) messages = result.scalars().all() out = [] for m in messages: entry: dict = {"role": m.role, "content": m.content, "created_at": m.created_at.isoformat() if m.created_at else None} if getattr(m, 'thinking', None): entry["thinking"] = m.thinking if m.role == "tool_call": # Parse JSON-encoded tool call data try: import json data = json.loads(m.content) entry["content"] = "" entry["toolName"] = data.get("name", "") entry["toolArgs"] = data.get("args") entry["toolStatus"] = data.get("status", "done") entry["toolResult"] = data.get("result", "") except Exception: pass out.append(entry) return out async def call_llm( model: LLMModel, messages: list[dict], agent_name: str, role_description: str, agent_id=None, user_id=None, session_id: str = "", on_chunk=None, on_tool_call=None, on_thinking=None, supports_vision=False, max_tool_rounds_override: int | None = None, ) -> str: """Call LLM via unified client with function-calling tool loop. Args: on_chunk: Optional async callback(text: str) for streaming chunks to client. on_thinking: Optional async callback(text: str) for reasoning/thinking content. on_tool_call: Optional async callback(dict) for tool call status updates. """ from app.services.agent_tools import AGENT_TOOLS, execute_tool, get_agent_tools_for_llm from app.services.llm_utils import create_llm_client, get_max_tokens, LLMMessage, LLMError # ── Token limit check & config ── _max_tool_rounds = 50 # default if agent_id: try: from app.models.agent import Agent as AgentModel async with async_session() as _db: _ar = await _db.execute(select(AgentModel).where(AgentModel.id == agent_id)) _agent = _ar.scalar_one_or_none() if _agent: _max_tool_rounds = _agent.max_tool_rounds or 50 if max_tool_rounds_override and max_tool_rounds_override < _max_tool_rounds: _max_tool_rounds = max_tool_rounds_override if _agent.max_tokens_per_day and _agent.tokens_used_today >= _agent.max_tokens_per_day: return f"⚠️ Daily token usage has reached the limit ({_agent.tokens_used_today:,}/{_agent.max_tokens_per_day:,}). Please try again tomorrow or ask admin to increase the limit." if _agent.max_tokens_per_month and _agent.tokens_used_month >= _agent.max_tokens_per_month: return f"⚠️ Monthly token usage has reached the limit ({_agent.tokens_used_month:,}/{_agent.max_tokens_per_month:,}). Please ask admin to increase the limit." except Exception: pass if max_tool_rounds_override and max_tool_rounds_override < _max_tool_rounds: _max_tool_rounds = max_tool_rounds_override # Build rich prompt with soul, memory, skills, relationships from app.services.agent_context import build_agent_context # Look up current user's display name so the agent knows who it's talking to _current_user_name = None if user_id: try: from app.models.user import User as _UserModel async with async_session() as _udb: _ur = await _udb.execute(select(_UserModel).where(_UserModel.id == user_id)) _u = _ur.scalar_one_or_none() if _u: _current_user_name = _u.display_name or _u.username except Exception: pass static_prompt, dynamic_prompt = await build_agent_context(agent_id, agent_name, role_description, current_user_name=_current_user_name) # Load tools dynamically from DB tools_for_llm = await get_agent_tools_for_llm(agent_id) if agent_id else AGENT_TOOLS # Convert messages to LLMMessage format api_messages = [LLMMessage(role="system", content=static_prompt, dynamic_content=dynamic_prompt)] for msg in messages: api_messages.append(LLMMessage( role=msg.get("role", "user"), content=msg.get("content"), tool_calls=msg.get("tool_calls"), tool_call_id=msg.get("tool_call_id"), )) # ── Vision format conversion ── # If the model supports vision, convert image markers in user messages # to OpenAI Vision API format: content becomes an array of parts. if supports_vision: import re as _re_v for i, msg in enumerate(api_messages): if msg.role != "user" or not msg.content or not isinstance(msg.content, str): continue content_str = msg.content # Find [image_data:data:image/...;base64,...] markers pattern = r'\[image_data:(data:image/[^;]+;base64,[A-Za-z0-9+/=]+)\]' images = _re_v.findall(pattern, content_str) if not images: continue # Build content array text = _re_v.sub(pattern, '', content_str).strip() parts = [] for img_url in images: parts.append({"type": "image_url", "image_url": {"url": img_url}}) if text: parts.append({"type": "text", "text": text}) # Replace the message content with the array format api_messages[i] = LLMMessage( role=msg.role, content=parts, # type: ignore # This is valid for vision models ) else: # Strip base64 image markers for non-vision models to avoid wasting tokens import re as _re_strip _img_pattern = r'\[image_data:data:image/[^;]+;base64,[A-Za-z0-9+/=]+\]' for i, msg in enumerate(api_messages): if msg.role != "user" or not isinstance(msg.content, str): continue if "[image_data:" in msg.content: _n_imgs = len(_re_strip.findall(_img_pattern, msg.content)) cleaned = _re_strip.sub(_img_pattern, '', msg.content).strip() if _n_imgs > 0: cleaned += f"\n[用户发送了 {_n_imgs} 张图片，但当前模型不支持视觉，无法查看图片内容]" api_messages[i] = LLMMessage( role=msg.role, content=cleaned, ) # Create the unified LLM client try: client = create_llm_client( provider=model.provider, api_key=model.api_key_encrypted, model=model.model, base_url=model.base_url, timeout=float(getattr(model, 'request_timeout', None) or 120.0), ) except Exception as e: return f"[Error] Failed to create LLM client: {e}" max_tokens = get_max_tokens(model.provider, model.model, getattr(model, 'max_output_tokens', None)) # ── Per-round token accumulator ── from app.services.token_tracker import record_token_usage, extract_usage_tokens, estimate_tokens_from_chars _accumulated_tokens = 0 # Tool-calling loop (configurable per agent, default 50) for round_i in range(_max_tool_rounds): # ── Dynamic tool-call limit warning (Aware engine) ── # Don't tell the agent about limits at the start — only warn when approaching. # This prevents models from rushing to complete tasks prematurely. _warn_threshold_80 = int(_max_tool_rounds * 0.8) _warn_threshold_96 = _max_tool_rounds - 2 if round_i == _warn_threshold_80: api_messages.append(LLMMessage( role="user", content=( f"⚠️ 你已使用 {round_i}/{_max_tool_rounds} 轮工具调用。" "如果当前任务尚未完成，请尽快保存进度到 focus.md，" "并使用 set_trigger 设置续接触发器，在剩余轮次中做好收尾。" ), )) elif round_i == _warn_threshold_96: api_messages.append(LLMMessage( role="user", content=f"🚨 仅剩 2 轮工具调用。请立即保存进度到 focus.md 并设置续接触发器。", )) try: # Use streaming API for real-time responses response = await client.stream( messages=api_messages, tools=tools_for_llm if tools_for_llm else None, temperature=model.temperature, max_tokens=max_tokens, on_chunk=on_chunk, on_thinking=on_thinking, ) except LLMError as e: # Record accumulated tokens before returning error logger.error( f"[LLM] LLMError provider={getattr(model, 'provider', '?')} " f"model={getattr(model, 'model', '?')} round={round_i + 1}: {e}" ) if agent_id and _accumulated_tokens > 0: await record_token_usage(agent_id, _accumulated_tokens) return f"[LLM Error] {e}" except Exception as e: logger.error( f"[LLM] Unexpected error provider={getattr(model, 'provider', '?')} " f"model={getattr(model, 'model', '?')} round={round_i + 1}: " f"{type(e).__name__}: {str(e)[:300]}" ) if agent_id and _accumulated_tokens > 0: await record_token_usage(agent_id, _accumulated_tokens) return f"[LLM call error] {type(e).__name__}: {str(e)[:200]}" # ── Track tokens for this round ── logger.debug(f"[LLM] stream() returned: {len(response.content or '')} chars, finish={response.finish_reason}, tools={len(response.tool_calls or [])}") real_tokens = extract_usage_tokens(response.usage) if real_tokens: _accumulated_tokens += real_tokens else: round_chars = sum(len(m.content or '') if isinstance(m.content, str) else 0 for m in api_messages) + len(response.content or '') _accumulated_tokens += estimate_tokens_from_chars(round_chars) # If no tool calls, return the final content if not response.tool_calls: if agent_id and _accumulated_tokens > 0: await record_token_usage(agent_id, _accumulated_tokens) await client.close() return response.content or "[LLM returned empty content]" # Execute tool calls logger.info(f"[LLM] Round {round_i+1}: {len(response.tool_calls)} tool call(s), finish_reason={response.finish_reason}") # Add assistant message with tool calls api_messages.append(LLMMessage( role="assistant", content=response.content or None, tool_calls=[{ "id": tc["id"], "type": "function", "function": tc["function"], } for tc in response.tool_calls], reasoning_content=response.reasoning_content, )) full_reasoning_content = response.reasoning_content or "" # Tools that require arguments — if LLM sends empty args, skip and ask to retry _TOOLS_REQUIRING_ARGS = {"write_file", "read_file", "delete_file", "read_document", "send_message_to_agent", "send_feishu_message", "send_email"} for tc in response.tool_calls: fn = tc["function"] tool_name = fn["name"] raw_args = fn.get("arguments", "{}") logger.info(f"[LLM] Raw arguments for {tool_name} (len={len(raw_args)}): {repr(raw_args[:300])}") try: args = json.loads(raw_args) if raw_args else {} except json.JSONDecodeError: args = {} # Guard: if a tool that requires arguments received empty args, # return an error to LLM instead of executing (Claude sometimes # emits tool_use blocks with no input_json_delta events) if not args and tool_name in _TOOLS_REQUIRING_ARGS: logger.warning(f"[LLM] Empty arguments for {tool_name}, asking LLM to retry") api_messages.append(LLMMessage( role="tool", content=f"Error: {tool_name} was called with empty arguments. You must provide the required parameters. Please retry with the correct arguments.", tool_call_id=tc.get("id", ""), )) continue logger.info(f"[LLM] Calling tool: {tool_name}({args})") # Notify client about tool call (in-progress) if on_tool_call: try: await on_tool_call({ "name": tool_name, "args": args, "status": "running", "reasoning_content": full_reasoning_content }) except Exception: pass result = await execute_tool( tool_name, args, agent_id=agent_id, user_id=user_id or agent_id, session_id=session_id, ) logger.debug(f"[LLM] Tool result: {result[:100]}") # Notify client about tool call result if on_tool_call: try: await on_tool_call({ "name": tool_name, "args": args, "status": "done", "result": result, "reasoning_content": full_reasoning_content }) except Exception as _cb_err: logger.warning(f"[LLM] on_tool_call callback error: {_cb_err}") # ── Vision injection for screenshot tools ── # If the model supports vision, try to inject the actual screenshot # image into the tool result so the LLM can SEE what's on screen. # Without this, the LLM only gets text like "Screenshot saved to ..." # and blindly guesses the page content. tool_content: str | list = str(result) if supports_vision and agent_id: try: from app.services.vision_inject import try_inject_screenshot_vision from app.services.agent_tools import WORKSPACE_ROOT ws_path = WORKSPACE_ROOT / str(agent_id) vision_content = try_inject_screenshot_vision(tool_name, str(result), ws_path) if vision_content: tool_content = vision_content logger.info(f"[LLM] Injected screenshot vision for {tool_name}") except Exception as e: logger.warning(f"[LLM] Vision injection failed for {tool_name}: {e}") api_messages.append(LLMMessage( role="tool", tool_call_id=tc["id"], content=tool_content, )) # Record tokens even on "too many rounds" exit if agent_id and _accumulated_tokens > 0: await record_token_usage(agent_id, _accumulated_tokens) await client.close() return "[Error] Too many tool call rounds" @router.websocket("/ws/chat/{agent_id}") async def websocket_chat( websocket: WebSocket, agent_id: uuid.UUID, token: str = Query(...), session_id: str = Query(None), ): """WebSocket endpoint for real-time chat with an agent. Flow: 1. Client connects with JWT token + optional session_id as query params 2. Server accepts immediately so browser onopen fires quickly 3. Server authenticates and checks agent access 4. If session_id provided, uses it; otherwise finds/creates the user's latest session 5. Client sends messages as JSON: {"content": "..."} 6. Server calls the agent's configured LLM and sends response back 7. Messages are persisted to chat_messages table under the session """ # Accept immediately so browser sees onopen without waiting for DB setup await websocket.accept() # Authenticate try: payload = decode_access_token(token) user_id = uuid.UUID(payload["sub"]) except Exception: await websocket.send_json({"type": "error", "content": "Authentication failed"}) await websocket.close(code=4001) return # Verify access and load agent + model agent_name = "" agent_type = "" # Track agent type for OpenClaw routing role_description = "" welcome_message = "" llm_model = None fallback_llm_model = None history_messages = [] try: async with async_session() as db: logger.info(f"[WS] Looking up user {user_id}") result = await db.execute(select(User).where(User.id == user_id)) user = result.scalar_one_or_none() if not user: logger.info("[WS] User not found") await websocket.send_json({"type": "error", "content": "User not found"}) await websocket.close(code=4001) return logger.info(f"[WS] Checking agent access for {agent_id}") agent, _ = await check_agent_access(db, user, agent_id) # Check agent expiry if is_agent_expired(agent): await websocket.send_json({"type": "error", "content": "This Agent has expired and is off duty. Please contact your admin to extend its service."}) await websocket.close(code=4003) return agent_name = agent.name agent_type = agent.agent_type or "" role_description = agent.role_description or "" welcome_message = agent.welcome_message or "" from app.models.agent import DEFAULT_CONTEXT_WINDOW_SIZE ctx_size = agent.context_window_size or DEFAULT_CONTEXT_WINDOW_SIZE logger.info(f"[WS] Agent: {agent_name}, type: {agent_type}, model_id: {agent.primary_model_id}, ctx: {ctx_size}") # Load the agent's primary model if agent.primary_model_id: model_result = await db.execute( select(LLMModel).where(LLMModel.id == agent.primary_model_id) ) llm_model = model_result.scalar_one_or_none() # Treat disabled models as unavailable at runtime if llm_model and not llm_model.enabled: logger.info(f"[WS] Primary model {llm_model.model} is disabled, skipping") llm_model = None else: logger.info(f"[WS] Primary model loaded: {llm_model.model if llm_model else 'None'}") # Load fallback model if agent.fallback_model_id: fb_result = await db.execute( select(LLMModel).where(LLMModel.id == agent.fallback_model_id) ) fallback_llm_model = fb_result.scalar_one_or_none() # Treat disabled fallback models as unavailable if fallback_llm_model and not fallback_llm_model.enabled: logger.info(f"[WS] Fallback model {fallback_llm_model.model} is disabled, skipping") fallback_llm_model = None elif fallback_llm_model: logger.info(f"[WS] Fallback model loaded: {fallback_llm_model.model}") # Config-level fallback: primary missing -> use fallback if not llm_model and fallback_llm_model: llm_model = fallback_llm_model fallback_llm_model = None # No further fallback available logger.info(f"[WS] Primary model unavailable, using fallback: {llm_model.model}") # Resolve or create chat session from app.models.chat_session import ChatSession from sqlalchemy import select as _sel from datetime import datetime as _dt, timezone as _tz conv_id = session_id if conv_id: # Validate the session belongs to this agent and to this user (no hijacking others' sessions). try: _sid = uuid.UUID(conv_id) except (ValueError, TypeError): conv_id = None _existing = None else: _sr = await db.execute( _sel(ChatSession).where( ChatSession.id == _sid, ChatSession.agent_id == agent_id, ) ) _existing = _sr.scalar_one_or_none() if not _existing: conv_id = None elif _existing.source_channel != "agent" and str(_existing.user_id) != str(user_id): await websocket.send_json({"type": "error", "content": "Not authorized for this session"}) await websocket.close(code=4003) return if not conv_id: # Find most recent session for this user+agent _sr = await db.execute( _sel(ChatSession) .where(ChatSession.agent_id == agent_id, ChatSession.user_id == user_id) .order_by(ChatSession.last_message_at.desc().nulls_last(), ChatSession.created_at.desc()) .limit(1) ) _latest = _sr.scalar_one_or_none() if _latest: conv_id = str(_latest.id) else: # Create a default session now = _dt.now(_tz.utc) _new_session = ChatSession( agent_id=agent_id, user_id=user_id, title=f"Session {now.strftime('%m-%d %H:%M')}", source_channel="web", created_at=now, ) db.add(_new_session) await db.commit() await db.refresh(_new_session) conv_id = str(_new_session.id) logger.info(f"[WS] Created default session {conv_id}") try: history_result = await db.execute( select(ChatMessage) .where(ChatMessage.agent_id == agent_id, ChatMessage.conversation_id == conv_id) .order_by(ChatMessage.created_at.desc()) .limit(ctx_size) ) history_messages = list(reversed(history_result.scalars().all())) logger.info(f"[WS] Loaded {len(history_messages)} history messages for session {conv_id}") except Exception as e: logger.warning(f"[WS] History load failed (non-fatal): {e}") except Exception as e: logger.error(f"[WS] Setup error: {type(e).__name__}: {e}") import traceback traceback.print_exc() await websocket.send_json({"type": "error", "content": "Setup failed"}) await websocket.close(code=4002) # Config error — client should NOT retry return agent_id_str = str(agent_id) if agent_id_str not in manager.active_connections: manager.active_connections[agent_id_str] = [] manager.active_connections[agent_id_str].append((websocket, conv_id)) logger.info(f"[WS] Ready! Agent={agent_name}") # Send session_id to frontend so Take Control can reference the correct session await websocket.send_json({"type": "connected", "session_id": conv_id}) # Build conversation context from history # IMPORTANT: Include tool_call messages so the LLM maintains tool-calling behavior. # Without them, Claude sees user→assistant-text patterns and learns to skip tools. conversation: list[dict] = [] for msg in history_messages: if msg.role == "tool_call": # Convert stored tool_call JSON into OpenAI-format assistant+tool pair try: import json as _j_hist tc_data = _j_hist.loads(msg.content) tc_name = tc_data.get("name", "unknown") tc_args = tc_data.get("args", {}) tc_result = tc_data.get("result", "") tc_id = f"call_{msg.id}" # synthetic tool_call_id # Assistant message with tool_calls array asst_msg = { "role": "assistant", "content": None, "tool_calls": [{ "id": tc_id, "type": "function", "function": {"name": tc_name, "arguments": _j_hist.dumps(tc_args, ensure_ascii=False)}, }], } if tc_data.get("reasoning_content"): asst_msg["reasoning_content"] = tc_data["reasoning_content"] conversation.append(asst_msg) # Tool result message. # Sanitize any stale [ImageID: ...] markers left by the ephemeral # screenshot cache — those images are gone from memory and would # confuse the LLM if sent as-is. from app.services.vision_inject import sanitize_history_tool_result sanitized_result = sanitize_history_tool_result(str(tc_result)) conversation.append({ "role": "tool", "tool_call_id": tc_id, "content": sanitized_result[:500], }) except Exception: continue # Skip malformed tool_call records else: entry = {"role": msg.role, "content": msg.content} if hasattr(msg, 'thinking') and msg.thinking: entry["thinking"] = msg.thinking conversation.append(entry) try: # Send welcome message on new session (no history) if welcome_message and not history_messages: await websocket.send_json({"type": "done", "role": "assistant", "content": welcome_message}) while True: logger.info(f"[WS] Waiting for message from {agent_name}...") data = await websocket.receive_json() # Set a unique trace ID for this specific message processing from app.core.logging_config import set_trace_id import uuid as _trace_uuid trace_id = str(_trace_uuid.uuid4())[:12] set_trace_id(trace_id) content = data.get("content", "") display_content = data.get("display_content", "") # User-facing display text file_name = data.get("file_name", "") # Original file name for attachment display logger.info(f"[WS] Received: {content[:50]}") if not content: continue # ── Quota checks ── try: from app.services.quota_guard import ( check_conversation_quota, increment_conversation_usage, check_agent_expired, check_agent_llm_quota, increment_agent_llm_usage, QuotaExceeded, AgentExpired, ) await check_conversation_quota(user_id) await check_agent_expired(agent_id) except QuotaExceeded as qe: await websocket.send_json({"type": "done", "role": "assistant", "content": f"⚠️ {qe.message}"}) continue except AgentExpired as ae: await websocket.send_json({"type": "done", "role": "assistant", "content": f"⚠️ {ae.message}"}) continue # Add user message to conversation (full LLM context) conversation.append({"role": "user", "content": content}) # Save user message to DB. # # Strategy: # - If the LLM content contains [image_data:...] markers (i.e. the user # attached an image and the model supports vision), persist the FULL # content including the base64 marker. This makes history self-contained # so subsequent turns can forward the image to the LLM without any # disk-based rehydration step. # - For all other messages (text, non-image files) use display_content for # cleaner history (avoids e.g. the raw file-text blob appearing in chat). # # The call_llm() path already strips [image_data:] for non-vision models # (websocket.py ~line 210), so no extra handling is needed at read time. HAS_IMAGE_MARKER = "[image_data:" in content if HAS_IMAGE_MARKER: # Preserve the full LLM content (includes base64) for multi-turn context. # Prefix with [file:name] for the UI history parser if a file name exists. saved_content = f"[file:{file_name}]\n{content}" if file_name else content else: saved_content = display_content if display_content else content if file_name: saved_content = f"[file:{file_name}]\n{saved_content}" async with async_session() as db: user_msg = ChatMessage( agent_id=agent_id, user_id=user_id, role="user", content=saved_content, conversation_id=conv_id, ) db.add(user_msg) # Update session last_message_at + auto-title on first message from app.models.chat_session import ChatSession as _CS from datetime import datetime as _dt2, timezone as _tz2 _now = _dt2.now(_tz2.utc) _sess_r = await db.execute( select(_CS).where(_CS.id == uuid.UUID(conv_id)) ) _sess = _sess_r.scalar_one_or_none() if _sess: _sess.last_message_at = _now if not history_messages and _sess.title.startswith("Session "): # Always use display_content for title (never expose raw base64) title_src = display_content if display_content else content # Clean up common prefixes from image/file messages clean_title = title_src.replace("[图片] ", "📷 ").replace("[image_data:", "").strip() if file_name and not clean_title: clean_title = f"📎 {file_name}" _sess.title = clean_title[:40] if clean_title else content[:40] await db.commit() logger.info("[WS] User message saved") # ── OpenClaw routing: insert into gateway_messages instead of LLM ── if agent_type == "openclaw": from app.models.gateway_message import GatewayMessage as GwMsg async with async_session() as db: gw_msg = GwMsg( agent_id=agent_id, sender_user_id=user_id, conversation_id=conv_id, content=content, status="pending", ) db.add(gw_msg) await db.commit() logger.info("[WS] OpenClaw: message queued for gateway poll") await websocket.send_json({ "type": "done", "role": "assistant", "content": "Message forwarded to OpenClaw agent. Waiting for response..." }) continue # Detect task creation intent import re task_match = re.search( r'(?:创建|新建|添加|建一个|帮我建|create|add)(?:一个|a )?(?:任务|待办|todo|task)[，,：：:\\s]*(.+)', content, re.IGNORECASE ) # Track thinking content for storage (initialize before condition) thinking_content = [] # Call LLM with streaming if llm_model: try: logger.info(f"[WS] Calling LLM {llm_model.model} (streaming)...") # Accumulate partial content for abort handling partial_chunks: list[str] = [] async def stream_to_ws(text: str): """Send each chunk to client in real-time.""" partial_chunks.append(text) await websocket.send_json({"type": "chunk", "content": text}) # Track which agentbay live URLs have been sent to avoid redundant pushes _sent_live_envs: set[str] = set() async def tool_call_to_ws(data: dict): """Send tool call info to client and persist completed ones.""" # ── AgentBay live preview: embed screenshot URL in tool_call message ── # We embed live preview data directly in the tool_call payload # because separate WebSocket messages get silently dropped by nginx. if data.get("status") == "done": try: from app.services.agentbay_live import detect_agentbay_env, get_desktop_screenshot, get_browser_snapshot import re as _re_live tool_name = data.get("name", "") env = detect_agentbay_env(tool_name) if env: tool_result = data.get("result", "") or "" if env == "desktop": b64_url = await get_desktop_screenshot(agent_id, session_id=conv_id) if b64_url: data["live_preview"] = {"env": env, "screenshot_url": b64_url} logger.info(f"[WS][LivePreview] Embedded {env} base64 in tool_call") elif env == "browser": b64_url = await get_browser_snapshot(agent_id, session_id=conv_id) if b64_url: data["live_preview"] = {"env": env, "screenshot_url": b64_url} logger.info(f"[WS][LivePreview] Embedded {env} base64 in tool_call") elif env == "code": data["live_preview"] = {"env": "code", "output": tool_result[:5000]} except Exception as _lp_err: logger.warning(f"[WS][LivePreview] Embed failed: {_lp_err}") await websocket.send_json({"type": "tool_call", **data}) # Save completed tool calls to DB so they persist in chat history if data.get("status") == "done": try: import json as _json_tc async with async_session() as _tc_db: tc_msg = ChatMessage( agent_id=agent_id, user_id=user_id, role="tool_call", content=_json_tc.dumps({ "name": data.get("name", ""), "args": data.get("args"), "status": "done", "result": (data.get("result") or "")[:500], "reasoning_content": data.get("reasoning_content"), }), conversation_id=conv_id, ) _tc_db.add(tc_msg) await _tc_db.commit() except Exception as _tc_err: logger.warning(f"[WS] Failed to save tool_call: {_tc_err}") # Track thinking content for storage thinking_content = [] async def thinking_to_ws(text: str): """Send thinking chunks to client for collapsible display.""" thinking_content.append(text) await websocket.send_json({"type": "thinking", "content": text}) import asyncio as _aio # Run call_llm as a cancellable task llm_task = _aio.create_task(call_llm( llm_model, conversation[-ctx_size:], agent_name, role_description, agent_id=agent_id, user_id=user_id, session_id=conv_id, on_chunk=stream_to_ws, on_tool_call=tool_call_to_ws, on_thinking=thinking_to_ws, supports_vision=getattr(llm_model, 'supports_vision', False), )) # Listen for abort while LLM is running aborted = False queued_messages: list[dict] = [] while not llm_task.done(): try: msg = await _aio.wait_for( websocket.receive_json(), timeout=0.5 ) if msg.get("type") == "abort": logger.info(f"[WS] Abort received, cancelling LLM task") llm_task.cancel() aborted = True break else: # Queue non-abort messages for later queued_messages.append(msg) except _aio.TimeoutError: continue except WebSocketDisconnect: llm_task.cancel() raise if aborted: # Wait for task to finish cancelling try: await llm_task except (_aio.CancelledError, Exception): pass partial_text = "".join(partial_chunks).strip() if partial_text: assistant_response = partial_text + "\n\n*[Generation stopped]*" else: assistant_response = "*[Generation stopped]*" logger.info(f"[WS] LLM aborted, partial: {assistant_response[:80]}") else: assistant_response = await llm_task logger.info(f"[WS] LLM response: {assistant_response[:80]}") # Update last_active_at from datetime import datetime, timezone as tz async with async_session() as _db: from app.models.agent import Agent as AgentModel _ar = await _db.execute(select(AgentModel).where(AgentModel.id == agent_id)) _agent = _ar.scalar_one_or_none() if _agent: _agent.last_active_at = datetime.now(tz.utc) await _db.commit() # Increment quota usage try: await increment_conversation_usage(user_id) await increment_agent_llm_usage(agent_id) except Exception: pass # Log activity from app.services.activity_logger import log_activity await log_activity(agent_id, "chat_reply", f"Replied to web chat: {assistant_response[:80]}", detail={"channel": "web", "user_text": content[:200], "reply": assistant_response[:500]}) except WebSocketDisconnect: raise except Exception as e: logger.error(f"[WS] LLM error: {e}") import traceback traceback.print_exc() # Runtime fallback: primary model failed -> retry with fallback model if fallback_llm_model: logger.info(f"[WS] Primary model failed, retrying with fallback: {fallback_llm_model.model}") try: await websocket.send_json({"type": "info", "content": f"Primary model error, switching to fallback model ({fallback_llm_model.model})..."}) assistant_response = await call_llm( fallback_llm_model, conversation[-ctx_size:], agent_name, role_description, agent_id=agent_id, user_id=user_id, session_id=conv_id, on_chunk=stream_to_ws, on_tool_call=tool_call_to_ws, on_thinking=thinking_to_ws, supports_vision=getattr(fallback_llm_model, 'supports_vision', False), ) logger.info(f"[WS] Fallback LLM response: {assistant_response[:80]}") except Exception as e2: logger.error(f"[WS] Fallback LLM also failed: {e2}") traceback.print_exc() assistant_response = f"[LLM call error] Primary: {str(e)[:100]} | Fallback: {str(e2)[:100]}" else: assistant_response = f"[LLM call error] {str(e)[:200]}" else: assistant_response = f"⚠️ {agent_name} has no LLM model configured. Please select a model in the agent's Settings tab." # If task creation detected, create a real Task record if task_match: task_title = task_match.group(1).strip() if task_title: try: from app.models.task import Task from app.services.task_executor import execute_task import asyncio as _asyncio async with async_session() as db: task = Task( agent_id=agent_id, title=task_title, created_by=user_id, status="pending", priority="medium", ) db.add(task) await db.commit() await db.refresh(task) task_id = task.id _asyncio.create_task(execute_task(task_id, agent_id)) assistant_response += f"\n\n📋 Task synced to task board: [{task_title}]" logger.info(f"[WS] Created task: {task_title}") except Exception as e: logger.error(f"[WS] Failed to create task: {e}") # Add assistant response to conversation conversation.append({"role": "assistant", "content": assistant_response}) # Save assistant message async with async_session() as db: asst_msg = ChatMessage( agent_id=agent_id, user_id=user_id, role="assistant", content=assistant_response, thinking=''.join(thinking_content) if thinking_content else None, conversation_id=conv_id, ) db.add(asst_msg) await db.commit() logger.info("[WS] Assistant message saved") # Send done signal with final content (for non-streaming clients) await websocket.send_json({ "type": "done", "role": "assistant", "content": assistant_response, }) logger.info("[WS] Response done sent to client") except WebSocketDisconnect: logger.info(f"[WS] Client disconnected: {agent_name}") manager.disconnect(agent_id_str, websocket) except Exception as e: logger.error(f"[WS] Error in message loop: {type(e).__name__}: {e}") import traceback traceback.print_exc() manager.disconnect(agent_id_str, websocket) try: await websocket.close(code=1011) except Exception: pass