mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-05-14 20:53:41 +00:00
eab7ae3d62
* feat: real-time subagent token usage display in header and per-turn Backend: - Persist subagent token usage to AIMessage.usage_metadata via TokenUsageMiddleware, so accumulateUsage() naturally includes subagent tokens without frontend state management - Cache subagent usage by tool_call_id in task_tool, write back to the dispatching AIMessage on next model response - Emit subagent token usage on all terminal task events (task_completed, task_failed, task_cancelled, task_timed_out) - Report subagent usage to parent RunJournal for API totals - Search backward from ToolMessage to find dispatching AIMessage for correct multi-tool-call attribution Frontend: - Remove subagentUsage state, custom event handling, and prop threading — subagent tokens are now embedded in message metadata - Simplify selectHeaderTokenUsage (no subagentUsage parameter) - Per-turn inline badges show turn-specific usage via message accumulation - Remove isLoading guard from MessageTokenUsageList for dynamic updates during streaming * fix: prevent header token double counting from baseline reset race onFinish, onError, and thread-switch useEffect all reset pendingUsageBaselineMessageIdsRef to an empty Set. If thread.isLoading is still true on the next render, all messages pass the getMessagesAfterBaseline filter and their tokens are added to backendUsage (which already includes them), causing the header to display up to 2× the actual token count. Capture current message IDs instead of using an empty Set so that getMessagesAfterBaseline correctly returns no pending messages even if thread.isLoading lags behind the stream end. * fix: write back subagent tokens for all concurrent task tool calls TokenUsageMiddleware only processed messages[-2], so when a single model response dispatched multiple task tool calls only the last ToolMessage had its cached subagent usage written back to the dispatch AIMessage.usage_metadata. Earlier tasks' usage stayed in _subagent_usage_cache indefinitely (leak) and never appeared in the per-turn inline token display. Walk backward through all consecutive ToolMessages before the new AIMessage, and accumulate updates targeting the same dispatch message into one state update so overlapping writes don't clobber each other. * fix: clean up subagent usage cache entry on task cancellation When a task_tool invocation is cancelled via CancelledError, any cached subagent usage entry leaked because the TokenUsageMiddleware writeback path never fires after cancellation. Pop the cache entry before re-raising to prevent unbounded growth of the module-level _subagent_usage_cache dict. * fix: address token usage review feedback * fix: handle missing config for subagent usage cache --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
41 lines
1.5 KiB
Python
41 lines
1.5 KiB
Python
"""Tests for user_id propagation through memory queue."""
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from deerflow.agents.memory.queue import ConversationContext, MemoryUpdateQueue
|
|
from deerflow.config.memory_config import MemoryConfig
|
|
|
|
|
|
def test_conversation_context_has_user_id():
|
|
ctx = ConversationContext(thread_id="t1", messages=[], user_id="alice")
|
|
assert ctx.user_id == "alice"
|
|
|
|
|
|
def test_conversation_context_user_id_default_none():
|
|
ctx = ConversationContext(thread_id="t1", messages=[])
|
|
assert ctx.user_id is None
|
|
|
|
|
|
def test_queue_add_stores_user_id():
|
|
q = MemoryUpdateQueue()
|
|
with patch("deerflow.agents.memory.queue.get_memory_config", return_value=MemoryConfig(enabled=True)), patch.object(q, "_reset_timer"):
|
|
q.add(thread_id="t1", messages=["msg"], user_id="alice")
|
|
assert len(q._queue) == 1
|
|
assert q._queue[0].user_id == "alice"
|
|
q.clear()
|
|
|
|
|
|
def test_queue_process_passes_user_id_to_updater():
|
|
q = MemoryUpdateQueue()
|
|
with patch("deerflow.agents.memory.queue.get_memory_config", return_value=MemoryConfig(enabled=True)), patch.object(q, "_reset_timer"):
|
|
q.add(thread_id="t1", messages=["msg"], user_id="alice")
|
|
|
|
mock_updater = MagicMock()
|
|
mock_updater.update_memory.return_value = True
|
|
with patch("deerflow.agents.memory.updater.MemoryUpdater", return_value=mock_updater):
|
|
q._process_queue()
|
|
|
|
mock_updater.update_memory.assert_called_once()
|
|
call_kwargs = mock_updater.update_memory.call_args.kwargs
|
|
assert call_kwargs["user_id"] == "alice"
|