Merge pull request #116 from Serverless-Devs/dev-0609

Enhance MemoryConversation with async initialization and fire-and-for…

Author: OhYee

agentrun/memory_collection/memory_conversation.py

@@ -4,6 +4,7 @@
 
 """
 
+import asyncio
 import json
 import os
 from typing import (
@@ -76,6 +77,7 @@ def __init__(
         # 延迟初始化
         self._memory_store = None
         self._ots_client = None
+        self._init_lock = asyncio.Lock()
 
     @staticmethod
     def _default_user_id_extractor(req: Any) -> str:
@@ -153,10 +155,18 @@ def _default_agent_id_extractor(req: Any) -> str:
         return "default_agent"
 
     async def _get_memory_store(self):
-        """获取或创建 AsyncMemoryStore 实例"""
+        """获取或创建 AsyncMemoryStore 实例（双检锁，并发安全）"""
         if self._memory_store is not None:
             return self._memory_store
 
+        async with self._init_lock:
+            # 拿到锁后再检查一次，防止并发请求重复初始化
+            if self._memory_store is not None:
+                return self._memory_store
+            return await self._init_memory_store()
+
+    async def _init_memory_store(self):
+        """内部初始化方法，由 _get_memory_store 在持锁状态下调用"""
         try:
             # 导入依赖
             from tablestore_for_agent_memory.base.base_memory_store import (
@@ -228,7 +238,7 @@ async def _get_memory_store(self):
             )
             await self._memory_store.init_table()
             await self._memory_store.init_search_index()
-            logger.info(f"Tables and indexes initialized successfully")
+            logger.info("Tables and indexes initialized successfully")
         except Exception as e:
             # 如果表已存在，会抛出异常，这是正常的
             logger.info(
@@ -384,10 +394,13 @@ async def wrap_invoke_agent(
             metadata={"agent_id": agent_id},
         )
 
-        try:
-            await memory_store.put_session(session)
-        except Exception as e:
-            logger.error(f"Failed to save session: {e}", exc_info=True)
+        async def _put_session_bg():
+            try:
+                await memory_store.put_session(session)
+            except Exception as e:
+                logger.error(f"Failed to save session: {e}", exc_info=True)
+
+        asyncio.create_task(_put_session_bg())
 
         # 构建输入消息列表（包含所有历史消息）
         input_messages = []
@@ -465,57 +478,62 @@ async def wrap_invoke_agent(
                 yield event
 
             # 保存完整的对话轮次（输入 + 输出）
-            # 只有当有文本内容或工具调用时才保存
+            # 使用 fire-and-forget 避免阻塞流式响应关闭
             if agent_response_content or tool_calls or tool_results:
-                try:
-                    # 构建助手响应消息
-                    assistant_message: Dict[str, Any] = {
-                        "role": "assistant",
-                    }
-
-                    # 添加文本内容（如果有）
-                    if agent_response_content:
-                        assistant_message["content"] = agent_response_content
-                    else:
-                        # OpenAI 格式要求：如果有 tool_calls，content 可以为 null
-                        assistant_message["content"] = None
-
-                    # 添加工具调用（如果有）
-                    if tool_calls:
-                        assistant_message["tool_calls"] = list(
-                            tool_calls.values()
+                # 构建助手响应消息
+                assistant_message: Dict[str, Any] = {
+                    "role": "assistant",
+                }
+
+                if agent_response_content:
+                    assistant_message["content"] = agent_response_content
+                else:
+                    assistant_message["content"] = None
+
+                if tool_calls:
+                    assistant_message["tool_calls"] = list(tool_calls.values())
+
+                output_messages = input_messages + [assistant_message]
+
+                if tool_results:
+                    output_messages.extend(tool_results)
+
+                conversation_message = Message(
+                    session_id=session_id,
+                    message_id=f"msg_{uuid.uuid4().hex[:16]}",
+                    content=json.dumps(output_messages, ensure_ascii=False),
+                )
+
+                async def _save_conversation_bg(
+                    ms=memory_store,
+                    msg=conversation_message,
+                    sess=session,
+                    n_msgs=len(output_messages),
+                    text_len=len(agent_response_content),
+                    n_tc=len(tool_calls),
+                    n_tr=len(tool_results),
+                ):
+                    try:
+                        await ms.put_message(msg)
+                        sess.update_time = microseconds_timestamp()
+                        await ms.update_session(sess)
+                        logger.debug(
+                            "Saved conversation: %d messages,"
+                            " text length: %d chars,"
+                            " tool_calls: %d, tool_results: %d",
+                            n_msgs,
+                            text_len,
+                            n_tc,
+                            n_tr,
+                        )
+                    except Exception as e:
+                        logger.error(
+                            "Failed to save conversation: %s",
+                            e,
+                            exc_info=True,
                         )
 
-                    # 构建完整的消息列表
-                    output_messages = input_messages + [assistant_message]
-
-                    # 添加工具执行结果（如果有）
-                    if tool_results:
-                        output_messages.extend(tool_results)
-
-                    # 将完整的对话历史存储为一条消息
-                    # content 字段存储 JSON 格式的消息列表
-                    conversation_message = Message(
-                        session_id=session_id,
-                        message_id=f"msg_{uuid.uuid4().hex[:16]}",
-                        content=json.dumps(output_messages, ensure_ascii=False),
-                    )
-                    await memory_store.put_message(conversation_message)
-
-                    # 更新 Session 时间
-                    session.update_time = microseconds_timestamp()
-                    await memory_store.update_session(session)
-
-                    logger.debug(
-                        f"Saved conversation: {len(output_messages)} messages,"
-                        f" text length: {len(agent_response_content)} chars,"
-                        f" tool_calls: {len(tool_calls)}, tool_results:"
-                        f" {len(tool_results)}"
-                    )
-                except Exception as e:
-                    logger.error(
-                        f"Failed to save conversation: {e}", exc_info=True
-                    )
+                asyncio.create_task(_save_conversation_bg())
 
         except Exception as e:
             logger.error(f"Error in agent handler: {e}", exc_info=True)

tests/unittests/memory_collection/test_memory_conversation.py

@@ -1,5 +1,6 @@
 """Tests for AgentRun Memory Conversation / AgentRun 记忆对话测试"""
 
+import asyncio
 from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
 import pytest
@@ -8,6 +9,11 @@
 from agentrun.server.model import AgentRequest, Message, MessageRole
 
 
+async def _flush_bg_tasks():
+    """Let fire-and-forget background tasks complete before assertions."""
+    await asyncio.sleep(0.05)
+
+
 @pytest.fixture
 def mock_memory_collection():
     """Mock MemoryCollection"""
@@ -185,6 +191,9 @@ async def mock_agent(request: AgentRequest):
         # Verify results
         assert results == ["Hello", ", ", "world!"]
 
+        # Wait for fire-and-forget background tasks to complete
+        await _flush_bg_tasks()
+
         # Verify memory store calls
         assert mock_memory_store.put_session.called
         assert mock_memory_store.put_message.called
@@ -252,6 +261,9 @@ async def mock_agent(request: AgentRequest):
         async for event in memory.wrap_invoke_agent(request, mock_agent):
             results.append(event)
 
+        # Wait for fire-and-forget background tasks to complete
+        await _flush_bg_tasks()
+
         # Verify agent still responds
         assert results == ["Still works!"]
 
@@ -339,6 +351,9 @@ async def mock_agent(request: AgentRequest):
         assert results[0] == "Let me search for that..."
         assert results[3] == "Based on the search, it's sunny today."
 
+        # Wait for fire-and-forget background tasks to complete
+        await _flush_bg_tasks()
+
         # Verify message was saved with tool calls
         assert mock_memory_store.put_message.called
         saved_message = mock_memory_store.put_message.call_args[0][0]
@@ -437,6 +452,9 @@ async def mock_agent(request: AgentRequest):
         # Verify all events were passed through
         assert len(results) == 4
 
+        # Wait for fire-and-forget background tasks to complete
+        await _flush_bg_tasks()
+
         # Verify message was saved with accumulated tool call
         assert mock_memory_store.put_message.called
         saved_message = mock_memory_store.put_message.call_args[0][0]