🧪 Update test files to improve the test coverage

Author: Jasonxia007

backend/services/agent_service.py

@@ -94,7 +94,7 @@ async def _stream_agent_chunks(
     local_messages = []
     captured_final_answer = None
     try:
-        async for chunk in agent_run(agent_run_info, memory_ctx):
+        async for chunk in agent_run(agent_run_info):
             local_messages.append(chunk)
             # Try to capture the final answer as it streams by in order to start memory addition
             try:

sdk/nexent/core/agents/run_agent.py

@@ -4,15 +4,15 @@
 
 from smolagents import ToolCollection
 
-from .agent_model import AgentRunInfo, MemoryContext
+from .agent_model import AgentRunInfo
 from .nexent_agent import NexentAgent, ProcessType
 
 
 logger = logging.getLogger("run_agent")
 logger.setLevel(logging.DEBUG)
 
 
-def agent_run_thread(agent_run_info: AgentRunInfo, memory_context: MemoryContext):
+def agent_run_thread(agent_run_info: AgentRunInfo):
     try:
         mcp_host = agent_run_info.mcp_host
         if mcp_host is None or len(mcp_host)==0:
@@ -50,10 +50,10 @@ def agent_run_thread(agent_run_info: AgentRunInfo, memory_context: MemoryContext
         raise ValueError(f"Error in agent_run_thread: {e}")
 
 
-async def agent_run(agent_run_info: AgentRunInfo, memory_context: MemoryContext):
+async def agent_run(agent_run_info: AgentRunInfo):
     observer = agent_run_info.observer
 
-    thread_agent = Thread(target=agent_run_thread, args=(agent_run_info, memory_context))
+    thread_agent = Thread(target=agent_run_thread, args=agent_run_info)
     thread_agent.start()
 
     while thread_agent.is_alive():

sdk/nexent/memory/memory_service.py

@@ -24,14 +24,14 @@
 
 def _filter_by_memory_level(memory_level: str, raw_results: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
     """
-    根据 memory_level 过滤搜索或列表结果，仅供模块内部使用。
+    Filter search or list results by memory_level
 
-    参数:
+    args:
         memory_level: "tenant" | "user" | "agent" | "user_agent"
-        raw_results:   需要过滤的结果列表
+        raw_results:   The list of results to filter
 
-    返回:
-        过滤后的结果列表
+    return:
+        The filtered list of results
     """
     if memory_level in {"tenant", "user"}:
         return [r for r in raw_results if not r.get("agent_id")]

sdk/pyproject.toml

@@ -54,7 +54,7 @@ dependencies = [
     "linkup-sdk",
     "paramiko>=3.4.0",
     "linkup-sdk",
-    "mem0ai>=0.1.116"
+    "mem0ai>=0.1.117"
 ]
 
 [tool.uv]

test/backend/agents/test_create_agent_info.py

@@ -1,6 +1,8 @@
 from fastapi import Request
 from consts.model import ExportAndImportAgentInfo, ExportAndImportDataFormat, MCPInfo, AgentRequest
 import sys
+import asyncio
+import json
 from unittest.mock import patch, MagicMock, mock_open, call, Mock, AsyncMock
 
 import pytest
@@ -2355,6 +2357,306 @@ def fake_unregister(conv_id):
 
 
 @pytest.mark.asyncio
+async def test__stream_agent_chunks_emits_error_chunk_on_run_failure(monkeypatch):
+    """When agent_run raises, an error SSE chunk should be emitted and run unregistered."""
+    agent_request = AgentRequest(
+        agent_id=1,
+        conversation_id=1001,
+        query="trigger error",
+        history=[],
+        minio_files=[],
+        is_debug=True,  # avoid persisting messages to focus on error path
+    )
+
+    async def failing_agent_run(*_, **__):
+        raise Exception("oops")
+
+    monkeypatch.setattr(
+        "backend.services.agent_service.agent_run", failing_agent_run, raising=False
+    )
+
+    called = {"unregistered": None}
+
+    def fake_unregister(conv_id):
+        called["unregistered"] = conv_id
+
+    monkeypatch.setattr(
+        "backend.services.agent_service.agent_run_manager.unregister_agent_run",
+        fake_unregister,
+        raising=False,
+    )
+
+    # Collect streamed chunks
+    collected = []
+    async for out in agent_service._stream_agent_chunks(
+        agent_request, "u", "t", MagicMock(), MagicMock()
+    ):
+        collected.append(out)
+
+    # Expect a single error payload chunk and unregister called
+    assert collected and collected[0].startswith(
+        "data: {") and "\"type\": \"error\"" in collected[0]
+    assert called["unregistered"] == 1001
+
+
+@pytest.mark.asyncio
+async def test__stream_agent_chunks_captures_final_answer_and_adds_memory(monkeypatch):
+    """Final answer should be captured and appended to memory via add_memory_in_levels."""
+    agent_request = AgentRequest(
+        agent_id=3,
+        conversation_id=3003,
+        query="hello",
+        history=[],
+        minio_files=[],
+        is_debug=False,
+    )
+
+    async def yield_final_answer(*_, **__):
+        yield json.dumps({"type": "token", "content": "hi"}, ensure_ascii=False)
+        yield json.dumps({"type": "final_answer", "content": "bye"}, ensure_ascii=False)
+
+    monkeypatch.setattr(
+        "backend.services.agent_service.agent_run", yield_final_answer, raising=False
+    )
+
+    add_calls = {"args": None, "called": False}
+
+    async def fake_add_memory_in_levels(**kwargs):
+        add_calls["args"] = kwargs
+        add_calls["called"] = True
+        return {"results": [{"ok": True}]}
+
+    monkeypatch.setattr(
+        "backend.services.agent_service.add_memory_in_levels",
+        fake_add_memory_in_levels,
+        raising=False,
+    )
+
+    # Memory context with switch ON
+    memory_ctx = MagicMock()
+    memory_ctx.user_config = MagicMock(
+        memory_switch=True,
+        agent_share_option="always",
+        disable_agent_ids=[],
+        disable_user_agent_ids=[],
+    )
+    memory_ctx.memory_config = {"cfg": 1}
+    memory_ctx.tenant_id = "t"
+    memory_ctx.user_id = "u"
+    memory_ctx.agent_id = 3
+
+    # Capture and await scheduled background task
+    task_holder = {"task": None}
+    orig_create_task = asyncio.create_task
+
+    def capture_task(coro):
+        t = orig_create_task(coro)
+        task_holder["task"] = t
+        return t
+
+    monkeypatch.setattr(asyncio, "create_task", capture_task)
+
+    # Run stream
+    collected = []
+    async for out in agent_service._stream_agent_chunks(
+        agent_request, "u", "t", MagicMock(query="hello"), memory_ctx
+    ):
+        collected.append(out)
+
+    # Ensure background task completed
+    if task_holder["task"] is not None:
+        await task_holder["task"]
+
+    assert add_calls["called"] is True
+    assert add_calls["args"]["messages"] == [
+        {"role": "user", "content": "hello"},
+        {"role": "assistant", "content": "bye"},
+    ]
+    assert set(add_calls["args"]["memory_levels"]) == {"agent", "user_agent"}
+    assert add_calls["args"]["memory_config"] == {"cfg": 1}
+    assert add_calls["args"]["tenant_id"] == "t"
+    assert add_calls["args"]["user_id"] == "u"
+    assert add_calls["args"]["agent_id"] == 3
+
+
+@pytest.mark.asyncio
+async def test__stream_agent_chunks_skips_memory_when_switch_off(monkeypatch):
+    """When memory switch is off, background memory addition exits early."""
+    agent_request = AgentRequest(
+        agent_id=4,
+        conversation_id=4004,
+        query="q",
+        history=[],
+        minio_files=[],
+        is_debug=False,
+    )
+
+    async def yield_one(*_, **__):
+        yield json.dumps({"type": "final_answer", "content": "ans"}, ensure_ascii=False)
+
+    monkeypatch.setattr(
+        "backend.services.agent_service.agent_run", yield_one, raising=False
+    )
+
+    called = {"count": 0}
+
+    async def track_add(**kwargs):
+        called["count"] += 1
+        return {"results": []}
+
+    monkeypatch.setattr(
+        "backend.services.agent_service.add_memory_in_levels", track_add, raising=False
+    )
+
+    memory_ctx = MagicMock()
+    memory_ctx.user_config = MagicMock(memory_switch=False)
+
+    async for _ in agent_service._stream_agent_chunks(
+        agent_request, "u", "t", MagicMock(query="q"), memory_ctx
+    ):
+        pass
+
+    await asyncio.sleep(0)
+    assert called["count"] == 0
+
+
+@pytest.mark.asyncio
+async def test__stream_agent_chunks_background_add_exception(monkeypatch):
+    """Exceptions in background memory addition should be caught and not crash the stream."""
+    agent_request = AgentRequest(
+        agent_id=5,
+        conversation_id=5005,
+        query="q",
+        history=[],
+        minio_files=[],
+        is_debug=False,
+    )
+
+    async def yield_final(*_, **__):
+        yield json.dumps({"type": "final_answer", "content": "A"}, ensure_ascii=False)
+
+    monkeypatch.setattr(
+        "backend.services.agent_service.agent_run", yield_final, raising=False
+    )
+
+    async def raise_in_add(**kwargs):
+        raise RuntimeError("mem add fail")
+
+    monkeypatch.setattr(
+        "backend.services.agent_service.add_memory_in_levels", raise_in_add, raising=False
+    )
+
+    memory_ctx = MagicMock()
+    memory_ctx.user_config = MagicMock(
+        memory_switch=True,
+        agent_share_option="always",
+        disable_agent_ids=[],
+        disable_user_agent_ids=[],
+    )
+
+    # Capture and await scheduled background task
+    task_holder = {"task": None}
+    orig_create_task = asyncio.create_task
+
+    def capture_task(coro):
+        t = orig_create_task(coro)
+        task_holder["task"] = t
+        return t
+
+    monkeypatch.setattr(asyncio, "create_task", capture_task)
+
+    async for _ in agent_service._stream_agent_chunks(
+        agent_request, "u", "t", MagicMock(query="q"), memory_ctx
+    ):
+        pass
+
+    # Let background exception be handled by awaiting the task
+    if task_holder["task"] is not None:
+        await task_holder["task"]
+
+
+@pytest.mark.asyncio
+async def test__stream_agent_chunks_schedule_task_failure(monkeypatch):
+    """Scheduling background task failure should be caught and logged."""
+    agent_request = AgentRequest(
+        agent_id=6,
+        conversation_id=6006,
+        query="q",
+        history=[],
+        minio_files=[],
+        is_debug=False,
+    )
+
+    async def yield_final(*_, **__):
+        yield json.dumps({"type": "final_answer", "content": "A"}, ensure_ascii=False)
+
+    monkeypatch.setattr(
+        "backend.services.agent_service.agent_run", yield_final, raising=False
+    )
+
+    # Force asyncio.create_task to fail
+    def fail_create_task(*_, **__):
+        raise RuntimeError("schedule fail")
+
+    monkeypatch.setattr("asyncio.create_task", fail_create_task)
+
+    memory_ctx = MagicMock()
+    memory_ctx.user_config = MagicMock(
+        memory_switch=True,
+        agent_share_option="always",
+        disable_agent_ids=[],
+        disable_user_agent_ids=[],
+    )
+
+    collected = []
+    async for out in agent_service._stream_agent_chunks(
+        agent_request, "u", "t", MagicMock(query="q"), memory_ctx
+    ):
+        collected.append(out)
+
+    assert collected  # Stream still produced data without crashing
+
+
+def test_insert_related_agent_impl_failure_returns_400():
+    """When insertion fails, should return 400 JSONResponse."""
+    with patch(
+        "backend.services.agent_service.query_sub_agents_id_list", return_value=[]
+    ) as _, patch(
+        "backend.services.agent_service.insert_related_agent", return_value=False
+    ) as __:
+        resp = insert_related_agent_impl(
+            parent_agent_id=1, child_agent_id=2, tenant_id="t")
+        assert resp.status_code == 400
+
+
+@pytest.mark.asyncio
+async def test_generate_stream_with_memory_unexpected_exception_emits_error(monkeypatch):
+    """Generic exceptions should emit an error SSE chunk and stop."""
+    agent_request = AgentRequest(
+        agent_id=9,
+        conversation_id=9009,
+        query="q",
+        history=[],
+        minio_files=[],
+        is_debug=False,
+    )
+
+    # Cause an unexpected error inside the try block
+    monkeypatch.setattr(
+        "backend.services.agent_service.build_memory_context",
+        MagicMock(side_effect=Exception("unexpected")),
+        raising=False,
+    )
+
+    out = []
+    async for d in agent_service.generate_stream_with_memory(
+        agent_request, user_id="u", tenant_id="t"
+    ):
+        out.append(d)
+
+    assert out and out[0].startswith(
+        "data: {") and "\"type\": \"error\"" in out[0]
+
 async def test_generate_stream_no_memory_registers_and_streams(monkeypatch):
     """generate_stream_no_memory should prepare run info, register it and stream data without memory tokens."""
     # Prepare AgentRequest & Request