Merge pull request #44 from rongxinzy/techdebt/runtime-bridge-extraction

refactor: extract runtime bridge and harden base agent

Author: hawkli-1994

swarmmind/agents/base.py

@@ -3,7 +3,7 @@
 from abc import ABC, abstractmethod
 
 from swarmmind.layered_memory import LayeredMemory
-from swarmmind.models import MemoryContext, MemoryScope
+from swarmmind.models import MemoryContext, MemoryLayer, MemoryScope
 from swarmmind.repositories.action_proposal import ActionProposalRepository
 from swarmmind.repositories.agent import AgentRepository
 
@@ -12,7 +12,6 @@ class AgentError(Exception):
     """Base exception for SwarmMind agent errors."""
 
     pass
-    pass
 
 
 class BaseAgent(ABC):
@@ -45,8 +44,6 @@ def _resolve_write_scope(self, ctx: MemoryContext | None) -> MemoryScope:
         but since ctx always provides user_id as a fallback, we fall through
         to the most specific available scope (never L4 for regular agents).
         """
-        from swarmmind.models import MemoryLayer
-
         if ctx is None:
             # No context — use a default user scope (L4, but agent will be denied
             # unless they are soul_writer; this is intentional guardrail)

swarmmind/agents/general_agent.py

@@ -5,7 +5,6 @@
 
 from __future__ import annotations
 
-import asyncio
 import json
 import logging
 import uuid
@@ -22,6 +21,7 @@
 from swarmmind.prompting import rewrite_swarmmind_identity_prompt
 from swarmmind.runtime import RuntimeExecutionError, ensure_default_runtime_instance
 from swarmmind.runtime.models import RuntimeInstance
+from swarmmind.services.runtime_bridge import iter_async_generator_in_thread
 
 logger = logging.getLogger(__name__)
 
@@ -524,72 +524,12 @@ def stream_events(
         with its own event loop. This isolates the runtime from any existing loop
         in the caller while preserving the synchronous generator API.
         """
-        import queue
-        import threading
-
-        result_queue: queue.Queue = queue.Queue()
-        exception_container = []
-        stop_event = threading.Event()
-
-        async def _run_async_stream():
-            """Run the async stream and put events into the queue."""
-            try:
-                async for event in self._astream_events(goal, ctx=ctx, runtime_options=runtime_options):
-                    result_queue.put(("event", event))
-            except Exception as e:
-                exception_container.append(e)
-            finally:
-                result_queue.put(("done", None))
-                stop_event.set()
-
-        def _run_in_thread():
-            """Run the async code in a new event loop in a separate thread.
-
-            NOTE: We create a new event loop in this thread to isolate the
-            DeerFlow agent execution from any existing event loop. This is
-            necessary because subagents also create their own event loops,
-            and we need to avoid httpx client binding conflicts.
-            """
-            # Set the event loop policy to create new loops per-thread
-            loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(loop)
-            try:
-                loop.run_until_complete(_run_async_stream())
-            finally:
-                # Clean up any remaining tasks
-                try:
-                    pending = asyncio.all_tasks(loop)
-                    if pending:
-                        for task in pending:
-                            task.cancel()
-                        loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
-                except Exception:  # nosec: B110 - cleanup code, safe to ignore
-                    pass
-                loop.close()
-
-        # Start the async execution in a separate thread
-        thread = threading.Thread(target=_run_in_thread, name="deerflow-stream")
-        thread.start()
-
-        # Yield events as they become available
-        try:
-            while not stop_event.is_set() or not result_queue.empty():
-                try:
-                    item_type, event = result_queue.get(timeout=0.1)
-                    if item_type == "done":
-                        break
-                    if item_type == "event":
-                        yield event
-                except queue.Empty:
-                    continue
-        finally:
-            thread.join(timeout=5.0)
-            if thread.is_alive():
-                logger.warning("Stream thread did not terminate within timeout")
-
-        # Re-raise any exception from the async execution
-        if exception_container:
-            raise exception_container[0]
+        yield from iter_async_generator_in_thread(
+            lambda: self._astream_events(goal, ctx=ctx, runtime_options=runtime_options),
+            thread_name="deerflow-stream",
+            join_timeout=5.0,
+            bridge_logger=logger,
+        )
 
         return self._last_final_text, self._last_tool_results
 

swarmmind/services/runtime_bridge.py

@@ -0,0 +1,75 @@
+"""Helpers for bridging async runtime streams into sync iterators."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import queue
+import threading
+from collections.abc import AsyncGenerator, Callable, Generator
+from typing import TypeVar
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+def iter_async_generator_in_thread[T](
+    async_factory: Callable[[], AsyncGenerator[T, None]],
+    *,
+    thread_name: str = "runtime-stream",
+    join_timeout: float = 5.0,
+    bridge_logger: logging.Logger = logger,
+) -> Generator[T, None, None]:
+    """Run an async generator inside a worker thread and yield its items synchronously."""
+    result_queue: queue.Queue[tuple[str, T | None]] = queue.Queue()
+    exception_container: list[BaseException] = []
+    stop_event = threading.Event()
+
+    async def _drain_async_generator() -> None:
+        try:
+            async for item in async_factory():
+                result_queue.put(("event", item))
+        except BaseException as exc:
+            exception_container.append(exc)
+        finally:
+            result_queue.put(("done", None))
+            stop_event.set()
+
+    def _run_in_thread() -> None:
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        try:
+            loop.run_until_complete(_drain_async_generator())
+        finally:
+            try:
+                pending = asyncio.all_tasks(loop)
+                if pending:
+                    for task in pending:
+                        task.cancel()
+                    loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
+            except Exception:  # nosec: B110 - cleanup code, safe to ignore
+                pass
+            loop.close()
+
+    thread = threading.Thread(target=_run_in_thread, name=thread_name, daemon=True)
+    thread.start()
+
+    try:
+        while not stop_event.is_set() or not result_queue.empty():
+            try:
+                item_type, item = result_queue.get(timeout=0.1)
+            except queue.Empty:
+                continue
+
+            if item_type == "done":
+                break
+            if item_type == "event" and item is not None:
+                yield item
+    finally:
+        thread.join(timeout=join_timeout)
+        if thread.is_alive():
+            bridge_logger.warning("Async bridge thread %s did not terminate within timeout", thread_name)
+
+    if exception_container:
+        raise exception_container[0]

tests/test_base_agent.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+import pytest
+
+from swarmmind.agents.base import AgentError, BaseAgent
+from swarmmind.models import MemoryContext, MemoryLayer
+
+
+class _FakeLayeredMemory:
+    def __init__(self, agent_id: str) -> None:
+        self.agent_id = agent_id
+
+
+class _DummyAgent(BaseAgent):
+    @property
+    def domain_tags(self) -> list[str]:
+        return [self.domain]
+
+
+def test_base_agent_loads_system_prompt_from_repository(monkeypatch: pytest.MonkeyPatch) -> None:
+    calls: list[str] = []
+
+    class _FakeAgentRepository:
+        def get_system_prompt(self, agent_id: str) -> str | None:
+            calls.append(agent_id)
+            return "system prompt from repository"
+
+    monkeypatch.setattr("swarmmind.agents.base.LayeredMemory", _FakeLayeredMemory)
+    monkeypatch.setattr("swarmmind.agents.base.AgentRepository", _FakeAgentRepository)
+
+    agent = _DummyAgent(agent_id="general", domain="general")
+
+    assert calls == ["general"]
+    assert agent._system_prompt == "system prompt from repository"
+    assert agent.memory.agent_id == "general"
+
+
+def test_base_agent_raises_for_missing_agent(monkeypatch: pytest.MonkeyPatch) -> None:
+    class _FakeAgentRepository:
+        def get_system_prompt(self, agent_id: str) -> str | None:
+            return None
+
+    monkeypatch.setattr("swarmmind.agents.base.LayeredMemory", _FakeLayeredMemory)
+    monkeypatch.setattr("swarmmind.agents.base.AgentRepository", _FakeAgentRepository)
+
+    with pytest.raises(AgentError, match="Agent missing not found in database\\."):
+        _DummyAgent(agent_id="missing", domain="general")
+
+
+def test_create_rejected_proposal_delegates_to_repository(monkeypatch: pytest.MonkeyPatch) -> None:
+    rejected_calls: list[tuple[str, str]] = []
+
+    class _FakeAgentRepository:
+        def get_system_prompt(self, agent_id: str) -> str | None:
+            return "prompt"
+
+    class _FakeActionProposalRepository:
+        def reject(self, proposal_id: str, description: str) -> None:
+            rejected_calls.append((proposal_id, description))
+
+    monkeypatch.setattr("swarmmind.agents.base.LayeredMemory", _FakeLayeredMemory)
+    monkeypatch.setattr("swarmmind.agents.base.AgentRepository", _FakeAgentRepository)
+    monkeypatch.setattr("swarmmind.agents.base.ActionProposalRepository", _FakeActionProposalRepository)
+
+    agent = _DummyAgent(agent_id="general", domain="general")
+    agent._create_rejected_proposal("proposal-123", "runtime failed")
+
+    assert rejected_calls == [("proposal-123", "runtime failed")]
+
+
+@pytest.mark.parametrize(
+    ("ctx", "expected_layer", "expected_scope_id"),
+    [
+        (
+            MemoryContext(user_id="user-1", project_id="project-1", team_id="team-1", session_id="session-1"),
+            MemoryLayer.TMP,
+            "session-1",
+        ),
+        (
+            MemoryContext(user_id="user-1", project_id="project-1", team_id="team-1"),
+            MemoryLayer.TEAM,
+            "team-1",
+        ),
+        (
+            MemoryContext(user_id="user-1", project_id="project-1"),
+            MemoryLayer.PROJECT,
+            "project-1",
+        ),
+        (
+            MemoryContext(user_id="user-1"),
+            MemoryLayer.USER_SOUL,
+            "user-1",
+        ),
+    ],
+)
+def test_resolve_write_scope_uses_most_specific_context(
+    monkeypatch: pytest.MonkeyPatch,
+    ctx: MemoryContext,
+    expected_layer: MemoryLayer,
+    expected_scope_id: str,
+) -> None:
+    class _FakeAgentRepository:
+        def get_system_prompt(self, agent_id: str) -> str | None:
+            return "prompt"
+
+    monkeypatch.setattr("swarmmind.agents.base.LayeredMemory", _FakeLayeredMemory)
+    monkeypatch.setattr("swarmmind.agents.base.AgentRepository", _FakeAgentRepository)
+
+    agent = _DummyAgent(agent_id="general", domain="general")
+    scope = agent._resolve_write_scope(ctx)
+
+    assert scope.layer == expected_layer
+    assert scope.scope_id == expected_scope_id

tests/test_runtime_bridge.py

@@ -0,0 +1,28 @@
+"""Tests for async-to-sync runtime bridge helpers."""
+
+from __future__ import annotations
+
+from swarmmind.services.runtime_bridge import iter_async_generator_in_thread
+
+
+def test_iter_async_generator_in_thread_yields_items_in_order() -> None:
+    async def factory():
+        yield "first"
+        yield "second"
+
+    assert list(iter_async_generator_in_thread(factory, thread_name="bridge-test")) == ["first", "second"]
+
+
+def test_iter_async_generator_in_thread_reraises_async_errors() -> None:
+    async def factory():
+        raise RuntimeError("bridge failed")
+        yield  # pragma: no cover
+
+    iterator = iter_async_generator_in_thread(factory, thread_name="bridge-test")
+
+    try:
+        next(iterator)
+    except RuntimeError as exc:
+        assert str(exc) == "bridge failed"
+    else:  # pragma: no cover
+        raise AssertionError("RuntimeError was not re-raised")