mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-08-09 22:38:46 +00:00
2730ee1f7b
When _process_queue found another worker already processing, it called _schedule_timer(0), spawning a fresh Timer thread immediately and looping tightly (spawn -> busy -> reschedule -> spawn) until the active worker finished. Replace this with a _reprocess_pending flag: a concurrent caller sets the flag and returns, and the active worker reschedules exactly once in its finally block when work remains. Reset the flag in clear().
458 lines
16 KiB
Python
458 lines
16 KiB
Python
import threading
|
|
import time
|
|
from unittest.mock import MagicMock, call, patch
|
|
|
|
from deerflow.agents.memory.queue import ConversationContext, MemoryUpdateQueue
|
|
from deerflow.config.memory_config import MemoryConfig
|
|
from deerflow.trace_context import get_current_trace_id, request_trace_context
|
|
|
|
|
|
def _memory_config(**overrides: object) -> MemoryConfig:
|
|
config = MemoryConfig()
|
|
for key, value in overrides.items():
|
|
setattr(config, key, value)
|
|
return config
|
|
|
|
|
|
def test_queue_add_preserves_existing_correction_flag_for_same_thread() -> None:
|
|
queue = MemoryUpdateQueue()
|
|
|
|
with (
|
|
patch("deerflow.agents.memory.queue.get_memory_config", return_value=_memory_config(enabled=True)),
|
|
patch.object(queue, "_reset_timer"),
|
|
):
|
|
queue.add(thread_id="thread-1", messages=["first"], correction_detected=True)
|
|
queue.add(thread_id="thread-1", messages=["second"], correction_detected=False)
|
|
|
|
assert len(queue._queue) == 1
|
|
assert queue._queue[0].messages == ["second"]
|
|
assert queue._queue[0].correction_detected is True
|
|
|
|
|
|
def test_process_queue_forwards_correction_flag_to_updater() -> None:
|
|
queue = MemoryUpdateQueue()
|
|
queue._queue = [
|
|
ConversationContext(
|
|
thread_id="thread-1",
|
|
messages=["conversation"],
|
|
agent_name="lead_agent",
|
|
correction_detected=True,
|
|
)
|
|
]
|
|
mock_updater = MagicMock()
|
|
mock_updater.update_memory.return_value = True
|
|
|
|
with patch("deerflow.agents.memory.updater.MemoryUpdater", return_value=mock_updater):
|
|
queue._process_queue()
|
|
|
|
mock_updater.update_memory.assert_called_once_with(
|
|
messages=["conversation"],
|
|
thread_id="thread-1",
|
|
agent_name="lead_agent",
|
|
correction_detected=True,
|
|
reinforcement_detected=False,
|
|
user_id=None,
|
|
deerflow_trace_id=None,
|
|
)
|
|
|
|
|
|
def test_queue_add_preserves_existing_reinforcement_flag_for_same_thread() -> None:
|
|
queue = MemoryUpdateQueue()
|
|
|
|
with (
|
|
patch("deerflow.agents.memory.queue.get_memory_config", return_value=_memory_config(enabled=True)),
|
|
patch.object(queue, "_reset_timer"),
|
|
):
|
|
queue.add(thread_id="thread-1", messages=["first"], reinforcement_detected=True)
|
|
queue.add(thread_id="thread-1", messages=["second"], reinforcement_detected=False)
|
|
|
|
assert len(queue._queue) == 1
|
|
assert queue._queue[0].messages == ["second"]
|
|
assert queue._queue[0].reinforcement_detected is True
|
|
|
|
|
|
def test_process_queue_forwards_reinforcement_flag_to_updater() -> None:
|
|
queue = MemoryUpdateQueue()
|
|
queue._queue = [
|
|
ConversationContext(
|
|
thread_id="thread-1",
|
|
messages=["conversation"],
|
|
agent_name="lead_agent",
|
|
reinforcement_detected=True,
|
|
)
|
|
]
|
|
mock_updater = MagicMock()
|
|
mock_updater.update_memory.return_value = True
|
|
|
|
with patch("deerflow.agents.memory.updater.MemoryUpdater", return_value=mock_updater):
|
|
queue._process_queue()
|
|
|
|
mock_updater.update_memory.assert_called_once_with(
|
|
messages=["conversation"],
|
|
thread_id="thread-1",
|
|
agent_name="lead_agent",
|
|
correction_detected=False,
|
|
reinforcement_detected=True,
|
|
user_id=None,
|
|
deerflow_trace_id=None,
|
|
)
|
|
|
|
|
|
def test_flush_nowait_cancels_existing_timer_and_starts_immediate_timer() -> None:
|
|
queue = MemoryUpdateQueue()
|
|
existing_timer = MagicMock()
|
|
queue._timer = existing_timer
|
|
created_timer = MagicMock()
|
|
|
|
with patch("deerflow.agents.memory.queue.threading.Timer", return_value=created_timer) as timer_cls:
|
|
queue.flush_nowait()
|
|
|
|
existing_timer.cancel.assert_called_once_with()
|
|
timer_cls.assert_called_once_with(0, queue._process_queue)
|
|
assert created_timer.daemon is True
|
|
created_timer.start.assert_called_once_with()
|
|
assert queue._timer is created_timer
|
|
|
|
|
|
def test_add_nowait_cancels_existing_timer_and_starts_immediate_timer() -> None:
|
|
queue = MemoryUpdateQueue()
|
|
existing_timer = MagicMock()
|
|
queue._timer = existing_timer
|
|
created_timer = MagicMock()
|
|
|
|
with (
|
|
patch("deerflow.agents.memory.queue.get_memory_config", return_value=_memory_config(enabled=True)),
|
|
patch("deerflow.agents.memory.queue.threading.Timer", return_value=created_timer) as timer_cls,
|
|
):
|
|
queue.add_nowait(thread_id="thread-1", messages=["conversation"], agent_name="lead-agent")
|
|
|
|
existing_timer.cancel.assert_called_once_with()
|
|
timer_cls.assert_called_once_with(0, queue._process_queue)
|
|
assert queue.pending_count == 1
|
|
assert queue._queue[0].agent_name == "lead-agent"
|
|
assert created_timer.daemon is True
|
|
created_timer.start.assert_called_once_with()
|
|
|
|
|
|
def test_process_queue_defers_reprocess_when_already_processing() -> None:
|
|
"""When a timer fires while a worker is active, ``_process_queue`` must set the
|
|
deferred-rerun flag instead of spinning up a tight 0-delay Timer chain.
|
|
|
|
The old behavior re-scheduled a 0-delay Timer on every re-entry while busy,
|
|
burning a fresh thread each time. The fix defers a single re-run via
|
|
``_reprocess_pending`` that the finishing worker honors once.
|
|
"""
|
|
queue = MemoryUpdateQueue()
|
|
queue._processing = True
|
|
|
|
with patch("deerflow.agents.memory.queue.threading.Timer") as timer_cls:
|
|
queue._process_queue()
|
|
|
|
timer_cls.assert_not_called()
|
|
assert queue._reprocess_pending is True
|
|
|
|
|
|
def test_finishing_worker_reschedules_once_when_reprocess_pending() -> None:
|
|
"""A worker that finishes with ``_reprocess_pending`` set and work still queued
|
|
schedules exactly one follow-up run (not a per-arrival timer spin)."""
|
|
queue = MemoryUpdateQueue()
|
|
queue._queue = [ConversationContext(thread_id="thread-1", messages=["first"], agent_name="lead_agent")]
|
|
queue._reprocess_pending = True
|
|
created_timer = MagicMock()
|
|
mock_updater = MagicMock()
|
|
|
|
def _enqueue_more_while_processing(**_kwargs) -> bool:
|
|
# Simulate a new update arriving mid-processing so the finally block sees
|
|
# remaining work and reschedules exactly once.
|
|
queue._queue.append(ConversationContext(thread_id="thread-2", messages=["second"], agent_name="lead_agent"))
|
|
return True
|
|
|
|
mock_updater.update_memory.side_effect = _enqueue_more_while_processing
|
|
|
|
with (
|
|
patch("deerflow.agents.memory.updater.MemoryUpdater", return_value=mock_updater),
|
|
patch("deerflow.agents.memory.queue.threading.Timer", return_value=created_timer) as timer_cls,
|
|
):
|
|
queue._process_queue()
|
|
|
|
timer_cls.assert_called_once_with(0, queue._process_queue)
|
|
assert created_timer.daemon is True
|
|
created_timer.start.assert_called_once_with()
|
|
assert queue._reprocess_pending is False
|
|
|
|
|
|
def test_finishing_worker_does_not_reschedule_when_no_work_remains() -> None:
|
|
"""The deferred re-run is cleared even when nothing is left to process, so a
|
|
stray flag never leaves a dangling ``_reprocess_pending``."""
|
|
queue = MemoryUpdateQueue()
|
|
queue._queue = [ConversationContext(thread_id="thread-1", messages=["only"], agent_name="lead_agent")]
|
|
queue._reprocess_pending = True
|
|
mock_updater = MagicMock()
|
|
mock_updater.update_memory.return_value = True
|
|
|
|
with (
|
|
patch("deerflow.agents.memory.updater.MemoryUpdater", return_value=mock_updater),
|
|
patch("deerflow.agents.memory.queue.threading.Timer") as timer_cls,
|
|
):
|
|
queue._process_queue()
|
|
|
|
timer_cls.assert_not_called()
|
|
assert queue._reprocess_pending is False
|
|
|
|
|
|
def test_flush_nowait_is_non_blocking() -> None:
|
|
queue = MemoryUpdateQueue()
|
|
started = threading.Event()
|
|
finished = threading.Event()
|
|
|
|
def _slow_process_queue() -> None:
|
|
started.set()
|
|
time.sleep(0.2)
|
|
finished.set()
|
|
|
|
queue._process_queue = _slow_process_queue
|
|
|
|
start = time.perf_counter()
|
|
queue.flush_nowait()
|
|
elapsed = time.perf_counter() - start
|
|
|
|
assert started.wait(0.1) is True
|
|
assert elapsed < 0.1
|
|
assert finished.is_set() is False
|
|
assert finished.wait(1.0) is True
|
|
|
|
|
|
def test_queue_keeps_updates_for_different_agents_in_same_thread() -> None:
|
|
queue = MemoryUpdateQueue()
|
|
|
|
with (
|
|
patch("deerflow.agents.memory.queue.get_memory_config", return_value=_memory_config(enabled=True)),
|
|
patch.object(queue, "_reset_timer"),
|
|
):
|
|
queue.add(thread_id="thread-1", messages=["agent-a"], agent_name="agent-a")
|
|
queue.add(thread_id="thread-1", messages=["agent-b"], agent_name="agent-b")
|
|
|
|
assert queue.pending_count == 2
|
|
assert [context.agent_name for context in queue._queue] == ["agent-a", "agent-b"]
|
|
|
|
|
|
def test_queue_still_coalesces_updates_for_same_agent_in_same_thread() -> None:
|
|
queue = MemoryUpdateQueue()
|
|
|
|
with (
|
|
patch("deerflow.agents.memory.queue.get_memory_config", return_value=_memory_config(enabled=True)),
|
|
patch.object(queue, "_reset_timer"),
|
|
):
|
|
queue.add(
|
|
thread_id="thread-1",
|
|
messages=["first"],
|
|
agent_name="agent-a",
|
|
correction_detected=True,
|
|
)
|
|
queue.add(
|
|
thread_id="thread-1",
|
|
messages=["second"],
|
|
agent_name="agent-a",
|
|
correction_detected=False,
|
|
)
|
|
|
|
assert queue.pending_count == 1
|
|
assert queue._queue[0].agent_name == "agent-a"
|
|
assert queue._queue[0].messages == ["second"]
|
|
assert queue._queue[0].correction_detected is True
|
|
|
|
|
|
def test_process_queue_updates_different_agents_in_same_thread_separately() -> None:
|
|
queue = MemoryUpdateQueue()
|
|
|
|
with (
|
|
patch("deerflow.agents.memory.queue.get_memory_config", return_value=_memory_config(enabled=True)),
|
|
patch.object(queue, "_reset_timer"),
|
|
):
|
|
queue.add(thread_id="thread-1", messages=["agent-a"], agent_name="agent-a")
|
|
queue.add(thread_id="thread-1", messages=["agent-b"], agent_name="agent-b")
|
|
|
|
mock_updater = MagicMock()
|
|
mock_updater.update_memory.return_value = True
|
|
|
|
with (
|
|
patch("deerflow.agents.memory.updater.MemoryUpdater", return_value=mock_updater),
|
|
patch("deerflow.agents.memory.queue.time.sleep"),
|
|
):
|
|
queue.flush()
|
|
|
|
assert mock_updater.update_memory.call_count == 2
|
|
mock_updater.update_memory.assert_has_calls(
|
|
[
|
|
call(
|
|
messages=["agent-a"],
|
|
thread_id="thread-1",
|
|
agent_name="agent-a",
|
|
correction_detected=False,
|
|
reinforcement_detected=False,
|
|
user_id=None,
|
|
deerflow_trace_id=None,
|
|
),
|
|
call(
|
|
messages=["agent-b"],
|
|
thread_id="thread-1",
|
|
agent_name="agent-b",
|
|
correction_detected=False,
|
|
reinforcement_detected=False,
|
|
user_id=None,
|
|
deerflow_trace_id=None,
|
|
),
|
|
]
|
|
)
|
|
|
|
|
|
def test_process_queue_forwards_deerflow_trace_id_to_updater() -> None:
|
|
queue = MemoryUpdateQueue()
|
|
queue._queue = [
|
|
ConversationContext(
|
|
thread_id="thread-1",
|
|
messages=["conversation"],
|
|
agent_name="lead_agent",
|
|
deerflow_trace_id="trace-memory-1",
|
|
)
|
|
]
|
|
mock_updater = MagicMock()
|
|
mock_updater.update_memory.return_value = True
|
|
|
|
with patch("deerflow.agents.memory.updater.MemoryUpdater", return_value=mock_updater):
|
|
queue._process_queue()
|
|
|
|
mock_updater.update_memory.assert_called_once_with(
|
|
messages=["conversation"],
|
|
thread_id="thread-1",
|
|
agent_name="lead_agent",
|
|
correction_detected=False,
|
|
reinforcement_detected=False,
|
|
user_id=None,
|
|
deerflow_trace_id="trace-memory-1",
|
|
)
|
|
|
|
|
|
class TestProcessQueueBindsTraceContextVar:
|
|
"""Regression: ``_process_queue`` runs in a Timer thread where the request
|
|
trace ContextVar is unbound. The per-context iteration must bind
|
|
``ConversationContext.deerflow_trace_id`` into the ContextVar so
|
|
``TraceContextFilter`` (which only reads the ContextVar) attaches the correct
|
|
``trace_id`` to log records emitted from ``queue.py`` itself (``"Updating
|
|
memory for thread ..."``, ``"Memory updated successfully..."``, exception
|
|
logs) — not just from the deep memory-updater stack.
|
|
"""
|
|
|
|
@staticmethod
|
|
def _run_process_queue_in_fresh_thread(queue: MemoryUpdateQueue, mock_updater: MagicMock) -> None:
|
|
def _target() -> None:
|
|
with patch("deerflow.agents.memory.updater.MemoryUpdater", return_value=mock_updater):
|
|
queue._process_queue()
|
|
|
|
thread = threading.Thread(target=_target)
|
|
thread.start()
|
|
thread.join()
|
|
|
|
def test_process_queue_binds_deerflow_trace_id_during_iteration(self) -> None:
|
|
queue = MemoryUpdateQueue()
|
|
queue._queue = [
|
|
ConversationContext(
|
|
thread_id="thread-1",
|
|
messages=["conversation"],
|
|
agent_name="lead_agent",
|
|
deerflow_trace_id="trace-queue-abc",
|
|
)
|
|
]
|
|
captured: list[str | None] = []
|
|
mock_updater = MagicMock()
|
|
|
|
def _capture(**_kwargs) -> bool:
|
|
captured.append(get_current_trace_id())
|
|
return True
|
|
|
|
mock_updater.update_memory.side_effect = _capture
|
|
|
|
self._run_process_queue_in_fresh_thread(queue, mock_updater)
|
|
|
|
assert captured == ["trace-queue-abc"]
|
|
|
|
def test_process_queue_binds_distinct_ids_per_context(self) -> None:
|
|
"""Each queued context must be scoped independently — a per-iteration bind,
|
|
not a batch-level one — so id A's logs don't bleed into id B's iteration."""
|
|
queue = MemoryUpdateQueue()
|
|
queue._queue = [
|
|
ConversationContext(
|
|
thread_id="thread-1",
|
|
messages=["conv-a"],
|
|
agent_name="agent-a",
|
|
deerflow_trace_id="trace-a",
|
|
),
|
|
ConversationContext(
|
|
thread_id="thread-2",
|
|
messages=["conv-b"],
|
|
agent_name="agent-b",
|
|
deerflow_trace_id="trace-b",
|
|
),
|
|
]
|
|
captured: list[str | None] = []
|
|
mock_updater = MagicMock()
|
|
|
|
def _capture(**_kwargs) -> bool:
|
|
captured.append(get_current_trace_id())
|
|
return True
|
|
|
|
mock_updater.update_memory.side_effect = _capture
|
|
|
|
with (
|
|
patch("deerflow.agents.memory.updater.MemoryUpdater", return_value=mock_updater),
|
|
patch("deerflow.agents.memory.queue.time.sleep"),
|
|
):
|
|
queue._process_queue()
|
|
|
|
assert captured == ["trace-a", "trace-b"]
|
|
|
|
def test_process_queue_leaves_contextvar_unbound_when_no_trace_id(self) -> None:
|
|
"""A queued context without ``deerflow_trace_id`` must not fabricate one;
|
|
the ContextVar stays unbound and log records fall through to '-'."""
|
|
queue = MemoryUpdateQueue()
|
|
queue._queue = [
|
|
ConversationContext(
|
|
thread_id="thread-1",
|
|
messages=["conversation"],
|
|
agent_name="lead_agent",
|
|
deerflow_trace_id=None,
|
|
)
|
|
]
|
|
captured: list[str | None] = []
|
|
mock_updater = MagicMock()
|
|
|
|
def _capture(**_kwargs) -> bool:
|
|
captured.append(get_current_trace_id())
|
|
return True
|
|
|
|
mock_updater.update_memory.side_effect = _capture
|
|
|
|
self._run_process_queue_in_fresh_thread(queue, mock_updater)
|
|
|
|
assert captured == [None]
|
|
|
|
def test_process_queue_restores_outer_contextvar_after_return(self) -> None:
|
|
queue = MemoryUpdateQueue()
|
|
queue._queue = [
|
|
ConversationContext(
|
|
thread_id="thread-1",
|
|
messages=["conversation"],
|
|
agent_name="lead_agent",
|
|
deerflow_trace_id="trace-inner",
|
|
)
|
|
]
|
|
mock_updater = MagicMock()
|
|
mock_updater.update_memory.return_value = True
|
|
|
|
with (
|
|
request_trace_context("trace-outer"),
|
|
patch("deerflow.agents.memory.updater.MemoryUpdater", return_value=mock_updater),
|
|
):
|
|
queue._process_queue()
|
|
assert get_current_trace_id() == "trace-outer"
|