feat: default init agent optimze (#163)

Author: yhjun1026

packages/derisk-app/src/derisk_app/app.py

@@ -7,6 +7,19 @@
 from fastapi.openapi.docs import get_swagger_ui_html
 from fastapi.staticfiles import StaticFiles
 
+
+class WebSocketAwareStaticFiles(StaticFiles):
+    """StaticFiles that gracefully handles WebSocket connections."""
+
+    async def __call__(self, scope, receive, send):
+        if scope["type"] == "websocket":
+            # WebSocket connections should not reach static files
+            # Let them fall through to 404 or other handlers
+            await send({"type": "websocket.close", "code": 1000})
+            return
+        await super().__call__(scope, receive, send)
+
+
 from derisk._version import version
 from derisk.component import SystemApp
 from derisk.configs.model_config import (
@@ -131,14 +144,18 @@ def mount_static_files(app: FastAPI, param: ApplicationConfig):
     os.makedirs(STATIC_MESSAGE_IMG_PATH, exist_ok=True)
     app.mount(
         "/images",
-        StaticFiles(directory=STATIC_MESSAGE_IMG_PATH, html=True),
+        WebSocketAwareStaticFiles(directory=STATIC_MESSAGE_IMG_PATH, html=True),
         name="static2",
     )
-    app.mount("/", StaticFiles(directory=static_file_path, html=True), name="static")
+    app.mount(
+        "/",
+        WebSocketAwareStaticFiles(directory=static_file_path, html=True),
+        name="static",
+    )
 
     app.mount(
         "/swagger_static",
-        StaticFiles(directory=static_file_path),
+        WebSocketAwareStaticFiles(directory=static_file_path),
         name="swagger_static",
     )
 

packages/derisk-core/src/derisk/agent/core/memory/gpts/gpts_memory.py

@@ -56,23 +56,89 @@
 # 消息通道迭代器
 # --------------------------
 class QueueIterator:
-    def __init__(self, queue: asyncio.Queue):
+    """异步队列迭代器，支持超时机制防止无限等待。
+
+    当队列长时间没有新消息时，会定期检查是否有错误发生，
+    避免因后台任务卡死或异常导致的无限等待。
+    """
+
+    DEFAULT_TIMEOUT = 30.0
+    MAX_TIMEOUT_COUNT = 3  # 最大连续超时次数，超过后抛出异常
+
+    def __init__(self, queue: asyncio.Queue, timeout: Optional[float] = None):
         self.queue = queue
+        self.timeout = timeout if timeout is not None else self.DEFAULT_TIMEOUT
+        self._error: Optional[Exception] = None
+        self._stopped = False
+        self._timeout_count = 0  # 连续超时计数器
+
+    def set_error(self, error: Exception):
+        """设置错误状态，将在下次迭代时抛出。"""
+        self._error = error
+        try:
+            self.queue.put_nowait(None)
+        except asyncio.QueueFull:
+            pass
+
+    def stop(self):
+        """停止迭代器。"""
+        self._stopped = True
+        try:
+            self.queue.put_nowait("[DONE]")
+        except asyncio.QueueFull:
+            pass
 
     def __aiter__(self):
         return self
 
     async def __anext__(self):
         start = time.perf_counter()
-        item = await self.queue.get()
-        if item == "[DONE]":
-            self.queue.task_done()
-            raise StopAsyncIteration
-        logger.debug(f"Queue wait: {(time.perf_counter() - start) * 1000:.2f}ms")
-        try:
-            return item
-        finally:
-            self.queue.task_done()
+
+        while True:
+            if self._error:
+                raise self._error
+
+            if self._stopped:
+                raise StopAsyncIteration
+
+            try:
+                item = await asyncio.wait_for(self.queue.get(), timeout=self.timeout)
+                self._timeout_count = 0  # 成功获取消息，重置超时计数
+            except asyncio.TimeoutError:
+                self._timeout_count += 1
+                wait_time = time.perf_counter() - start
+
+                if self._timeout_count >= self.MAX_TIMEOUT_COUNT:
+                    logger.error(
+                        f"Queue timeout exceeded max retries ({self.MAX_TIMEOUT_COUNT}), "
+                        f"total wait: {wait_time:.2f}s. Terminating to prevent infinite wait."
+                    )
+                    raise TimeoutError(
+                        f"对话响应超时，已等待 {wait_time:.1f} 秒。"
+                        f"请检查后端服务状态或稍后重试。"
+                    )
+
+                logger.warning(
+                    f"Queue wait timeout ({self._timeout_count}/{self.MAX_TIMEOUT_COUNT}) "
+                    f"after {wait_time:.2f}s, continuing to wait... (queue size: {self.queue.qsize()})"
+                )
+                continue
+
+            if item == "[DONE]":
+                self.queue.task_done()
+                raise StopAsyncIteration
+
+            if item is None:
+                if self._error:
+                    raise self._error
+                self.queue.task_done()
+                continue
+
+            logger.debug(f"Queue wait: {(time.perf_counter() - start) * 1000:.2f}ms")
+            try:
+                return item
+            finally:
+                self.queue.task_done()
 
 
 class AgentTaskType(Enum):
@@ -616,9 +682,11 @@ async def _merge_messages_async(
     # --------------------------
     # 外部核心方法区
     # --------------------------
-    async def queue_iterator(self, conv_id: str) -> Optional[QueueIterator]:
+    async def queue_iterator(
+        self, conv_id: str, timeout: Optional[float] = None
+    ) -> Optional[QueueIterator]:
         cache = await self._get_cache(conv_id)
-        return QueueIterator(cache.channel) if cache else None
+        return QueueIterator(cache.channel, timeout=timeout) if cache else None
 
     async def init(
         self,

packages/derisk-core/src/derisk/agent/core_v2/enhanced_agent.py

@@ -21,6 +21,7 @@
 
 from .improved_compaction import ImprovedSessionCompaction, CompactionConfig
 from .llm_utils import call_llm, LLMCaller
+from .tools_v2 import ToolRegistry, ToolResult
 
 logger = logging.getLogger(__name__)
 
@@ -154,88 +155,6 @@ def _match_pattern(self, tool_name: str, pattern: str) -> bool:
         return fnmatch.fnmatch(tool_name, pattern)
 
 
-class ToolRegistry:
-    """工具注册表"""
-
-    def __init__(self):
-        self._tools: Dict[str, Any] = {}
-
-    def register(self, tool: Any) -> "ToolRegistry":
-        # 优先使用 metadata.name，其次使用 name 属性
-        if hasattr(tool, "metadata") and hasattr(tool.metadata, "name"):
-            name = tool.metadata.name
-        elif hasattr(tool, "name"):
-            name = tool.name
-        else:
-            name = str(tool)
-        self._tools[name] = tool
-        logger.debug(f"[ToolRegistry] 注册工具: {name}")
-        return self
-
-    def get(self, name: str) -> Optional[Any]:
-        return self._tools.get(name)
-
-    async def execute(
-        self, name: str, args: Dict[str, Any], context: Optional[Dict[str, Any]] = None
-    ) -> "ToolResult":
-        """执行工具"""
-        from .tools_v2 import ToolResult
-
-        tool = self._tools.get(name)
-        if not tool:
-            return ToolResult(success=False, output="", error=f"工具不存在: {name}")
-
-        try:
-            if hasattr(tool, "execute"):
-                result = await tool.execute(args, context)
-                return result
-            elif callable(tool):
-                result = tool(**args)
-                if hasattr(result, "__await__"):
-                    result = await result
-                if isinstance(result, ToolResult):
-                    return result
-                return ToolResult(
-                    success=True,
-                    output=str(result) if result else "",
-                )
-            else:
-                return ToolResult(
-                    success=False, output="", error=f"工具不可执行: {name}"
-                )
-        except Exception as e:
-            logger.exception(f"[ToolRegistry] 工具执行异常: {name}")
-            return ToolResult(success=False, output="", error=str(e))
-
-    def list_tools(self) -> List[str]:
-        return list(self._tools.keys())
-
-    def list_all(self) -> List[Any]:
-        """列出所有工具对象"""
-        return list(self._tools.values())
-
-    def list_names(self) -> List[str]:
-        """列出所有工具名称"""
-        return list(self._tools.keys())
-
-    def get_openai_tools(self) -> List[Dict[str, Any]]:
-        result = []
-        for name, tool in self._tools.items():
-            if hasattr(tool, "get_openai_spec"):
-                result.append(tool.get_openai_spec())
-            else:
-                result.append(
-                    {
-                        "type": "function",
-                        "function": {
-                            "name": name,
-                            "description": getattr(tool, "description", ""),
-                        },
-                    }
-                )
-        return result
-
-
 @dataclass
 class SubagentSession:
     """子代理会话"""
@@ -1119,8 +1038,9 @@ def _build_llm_messages(self) -> List:
             else:
                 messages.append(SystemMessage(content=msg.content))
 
-        if self.tools.list_tools():
-            tools_desc = "Available tools: " + ", ".join(self.tools.list_tools())
+        if self.tools.list_all():
+            tool_names = [t.metadata.name for t in self.tools.list_all()]
+            tools_desc = "Available tools: " + ", ".join(tool_names)
             messages.append(SystemMessage(content=tools_desc))
 
         return messages

packages/derisk-core/src/derisk/agent/core_v2/integration/runtime.py

@@ -745,18 +745,28 @@ async def _push_user_message(self, conv_id: str, message: str):
 
         # 保存到 GptsMemory (gpts_messages 表)
         if self.gpts_memory:
-            user_msg = type(
-                "GptsMessage",
-                (),
-                {
-                    "message_id": str(uuid.uuid4().hex),
-                    "conv_id": conv_id,
-                    "sender": "user",
-                    "receiver": "assistant",
-                    "content": message,
-                    "rounds": 0,
-                },
-            )()
+            # Find the session to get additional info
+            session = None
+            for s in self._sessions.values():
+                if s.conv_id == conv_id:
+                    session = s
+                    break
+
+            # Create a proper GptsMessage with all required attributes
+            user_msg = GptsMessage(
+                conv_id=conv_id,
+                conv_session_id=session.session_id if session else conv_id,
+                message_id=str(uuid.uuid4().hex),
+                sender="user",
+                sender_name="user",
+                receiver="assistant",
+                receiver_name="assistant",
+                role="user",
+                content=message,
+                rounds=0,
+                app_code=session.agent_name if session else None,
+                app_name=session.agent_name if session else None,
+            )
             await self.gpts_memory.append_message(conv_id, user_msg, save_db=True)
 
         # 同时保存到 StorageConversation (ChatHistoryMessageEntity 表)
@@ -865,19 +875,21 @@ async def _push_stream_chunk(self, conv_id: str, chunk: V2StreamChunk):
 
             # 保存到 GptsMemory (gpts_messages 表)
             if self.gpts_memory and session.accumulated_content:
-                assistant_msg = type(
-                    "GptsMessage",
-                    (),
-                    {
-                        "message_id": session.current_message_id
-                        or str(uuid.uuid4().hex),
-                        "conv_id": conv_id,
-                        "sender": session.agent_name,
-                        "receiver": "user",
-                        "content": vis_final_content,
-                        "rounds": 0,
-                    },
-                )()
+                # Create a proper GptsMessage with all required attributes
+                assistant_msg = GptsMessage(
+                    conv_id=conv_id,
+                    conv_session_id=session.session_id,
+                    message_id=session.current_message_id or str(uuid.uuid4().hex),
+                    sender=session.agent_name,
+                    sender_name=session.agent_name,
+                    receiver="user",
+                    receiver_name="user",
+                    role="assistant",
+                    content=vis_final_content,
+                    rounds=0,
+                    app_code=session.agent_name,
+                    app_name=session.agent_name,
+                )
                 await self.gpts_memory.append_message(
                     conv_id, assistant_msg, save_db=True
                 )

packages/derisk-core/src/derisk/agent/core_v2/tools_v2/__init__.py

@@ -35,9 +35,21 @@
 from derisk.agent.tools.registry import (
     ToolRegistry,
     tool_registry,
-    register_builtin_tools,
 )
 
+
+def register_builtin_tools(registry: ToolRegistry) -> None:
+    """
+    注册内置工具到注册表（兼容层）
+
+    此函数接受 registry 参数，以兼容旧版 API。
+    内部调用 derisk.agent.tools.builtin.register_all
+    """
+    from derisk.agent.tools.builtin import register_all
+
+    register_all(registry)
+
+
 from derisk.agent.tools.decorators import tool
 
 # 从内置工具模块导入具体工具
@@ -167,7 +179,7 @@ def register_all_tools(
     import logging
 
     logger = logging.getLogger(__name__)
-    logger.info(f"[Tools] 已注册所有工具，共 {len(registry.list_names())} 个")
+    logger.info(f"[Tools] 已注册所有工具，共 {len(registry.list_all())} 个")
 
     return registry