Files
bytedance--deer-flow/backend/tests/test_subagent_step_events.py
2026-07-13 11:59:58 +08:00

351 lines
13 KiB
Python

"""Tests for the pure subagent step-payload builder (issue #3779).
``build_subagent_step`` turns a captured subagent message dict (the
``model_dump()`` of an AIMessage or ToolMessage) into the compact,
serializable step payload that is both streamed (``task_running``) and
persisted (``subagent.step`` run events). It is a pure function so it can
be unit-tested without the executor/graph.
"""
from __future__ import annotations
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
from deerflow.subagents.step_events import (
SUBAGENT_EVENT_CATEGORY,
SUBAGENT_STEP_MAX_CHARS,
build_subagent_step,
capture_new_step_messages,
capture_step_message,
subagent_run_event,
truncate_step_text,
)
def test_ai_message_becomes_ai_step_with_tool_calls():
message = {
"type": "ai",
"id": "ai-1",
"content": "Let me search the web.",
"tool_calls": [
{"name": "web_search", "args": {"query": "deerflow"}, "id": "call_1", "type": "tool_call"},
],
}
step = build_subagent_step(message, task_id="call_task", message_index=1)
assert step["task_id"] == "call_task"
assert step["message_index"] == 1
assert step["kind"] == "ai"
assert step["text"] == "Let me search the web."
assert step["truncated"] is False
assert step["tool_calls"] == [{"name": "web_search", "args": {"query": "deerflow"}}]
assert "tool_name" not in step
def test_tool_message_becomes_tool_step_with_output():
message = {
"type": "tool",
"id": "tool-1",
"name": "web_search",
"tool_call_id": "call_1",
"content": "Result: DeerFlow is a LangGraph super-agent.",
}
step = build_subagent_step(message, task_id="call_task", message_index=2)
assert step["kind"] == "tool"
assert step["tool_name"] == "web_search"
assert step["text"] == "Result: DeerFlow is a LangGraph super-agent."
assert step["truncated"] is False
assert "tool_calls" not in step
def test_long_tool_output_is_truncated_and_flagged():
big = "x" * (SUBAGENT_STEP_MAX_CHARS + 500)
message = {"type": "tool", "name": "read_file", "content": big}
step = build_subagent_step(message, task_id="t", message_index=3, max_chars=SUBAGENT_STEP_MAX_CHARS)
assert step["truncated"] is True
assert len(step["text"]) == SUBAGENT_STEP_MAX_CHARS
def test_list_content_blocks_are_flattened_to_text():
message = {
"type": "ai",
"content": [
{"type": "text", "text": "first"},
{"type": "text", "text": "second"},
],
"tool_calls": [],
}
step = build_subagent_step(message, task_id="t", message_index=1)
assert "first" in step["text"]
assert "second" in step["text"]
assert step["tool_calls"] == []
def test_ai_text_is_also_truncated():
big = "y" * (SUBAGENT_STEP_MAX_CHARS + 10)
message = {"type": "ai", "content": big, "tool_calls": []}
step = build_subagent_step(message, task_id="t", message_index=1, max_chars=SUBAGENT_STEP_MAX_CHARS)
assert step["truncated"] is True
assert len(step["text"]) == SUBAGENT_STEP_MAX_CHARS
def test_truncate_step_text_helper():
assert truncate_step_text("abc", 10) == ("abc", False)
assert truncate_step_text("abcdef", 3) == ("abc", True)
def test_capture_ai_message_appends_dict():
captured: list[dict] = []
seen: set[str] = set()
appended = capture_step_message(AIMessage(content="hi", id="ai-1"), captured, seen)
assert appended is True
assert len(captured) == 1
assert captured[0]["type"] == "ai"
def test_capture_tool_message_is_now_captured():
# Regression for #3779: tool outputs (ToolMessage) used to be dropped,
# so "what each step produced" never reached the UI/store.
captured: list[dict] = []
seen: set[str] = set()
appended = capture_step_message(
ToolMessage(content="search results", tool_call_id="call_1", name="web_search", id="tool-1"),
captured,
seen,
)
assert appended is True
assert captured[0]["type"] == "tool"
assert captured[0]["name"] == "web_search"
def test_capture_dedupes_by_id():
captured: list[dict] = []
seen: set[str] = set()
msg = AIMessage(content="hi", id="ai-1")
assert capture_step_message(msg, captured, seen) is True
assert capture_step_message(msg, captured, seen) is False
assert len(captured) == 1
def test_capture_ignores_human_message():
captured: list[dict] = []
seen: set[str] = set()
appended = capture_step_message(HumanMessage(content="user input", id="h-1"), captured, seen)
assert appended is False
assert captured == []
def test_none_content_flattens_to_empty_string():
# A tool-call-only AI turn can carry content=None; it must render as "" (not
# the literal "None"), matching the shared message_content_to_text guard.
message = {"type": "ai", "content": None, "tool_calls": []}
step = build_subagent_step(message, task_id="t", message_index=1)
assert step["text"] == ""
def test_ai_step_caps_large_tool_call_args():
# Regression for #3779: build_subagent_step capped `text` but copied
# `tool_calls[].args` verbatim, so a write_file/bash call carrying a big
# payload produced an unbounded persisted row. Args must now be capped too.
big_payload = "F" * (SUBAGENT_STEP_MAX_CHARS + 4096)
message = {
"type": "ai",
"content": "writing the file",
"tool_calls": [
{"name": "write_file", "args": {"path": "/mnt/out.txt", "content": big_payload}},
],
}
step = build_subagent_step(message, task_id="t", message_index=1, max_chars=SUBAGENT_STEP_MAX_CHARS)
call = step["tool_calls"][0]
assert call["name"] == "write_file"
assert call["args_truncated"] is True
# The serialized args are bounded by the same cap the text field uses.
assert isinstance(call["args"], str)
assert len(call["args"]) == SUBAGENT_STEP_MAX_CHARS
def test_ai_step_keeps_small_tool_call_args_structured():
message = {
"type": "ai",
"content": "searching",
"tool_calls": [{"name": "web_search", "args": {"query": "deerflow"}}],
}
step = build_subagent_step(message, task_id="t", message_index=1)
call = step["tool_calls"][0]
assert call["args"] == {"query": "deerflow"}
assert "args_truncated" not in call
def test_capture_new_step_messages_captures_full_multi_tool_tail():
# Regression for #3779: a single super-step can append several ToolMessages
# (one per tool call in a multi-tool turn). Capturing only messages[-1]
# dropped all but the last; the tail walk must capture every new message.
captured: list[dict] = []
seen: set[str] = set()
# Chunk 1: human + one AIMessage requesting 3 tool calls.
chunk1 = [
HumanMessage(content="do work", id="h-1"),
AIMessage(content="running tools", id="ai-1"),
]
processed = capture_new_step_messages(chunk1, captured, seen, 0)
assert processed == 2
assert [c["id"] for c in captured] == ["ai-1"]
# Chunk 2: values-mode re-yields the whole history plus 3 new ToolMessages
# appended in one super-step.
chunk2 = chunk1 + [
ToolMessage(content="r1", tool_call_id="c1", name="web_search", id="tool-1"),
ToolMessage(content="r2", tool_call_id="c2", name="read_file", id="tool-2"),
ToolMessage(content="r3", tool_call_id="c3", name="web_search", id="tool-3"),
]
processed = capture_new_step_messages(chunk2, captured, seen, processed)
assert processed == 5
# All three tool outputs survive, not just the last.
assert [c["id"] for c in captured] == ["ai-1", "tool-1", "tool-2", "tool-3"]
def test_capture_new_step_messages_is_noop_on_values_reyield():
# stream_mode="values" re-yields the same trailing message with unchanged
# length; re-processing must not duplicate captures.
captured: list[dict] = []
seen: set[str] = set()
messages = [AIMessage(content="hi", id="ai-1")]
processed = capture_new_step_messages(messages, captured, seen, 0)
assert processed == 1
# Same list handed back (no growth) — cursor already at the end.
processed = capture_new_step_messages(messages, captured, seen, processed)
assert processed == 1
assert len(captured) == 1
def test_capture_new_step_messages_handles_history_contraction():
# Regression for #3875 Phase 3: DeerFlowSummarizationMiddleware rewrites the
# messages channel via RemoveMessage(id=REMOVE_ALL_MESSAGES), which shrinks
# len(messages) below the cursor we were tracking. Without a contraction
# reset, every step appended AFTER the compaction is dropped until total
# overtakes the stale cursor.
#
# Faithful to the real middleware: compaction puts the summary into a
# SEPARATE ``summary_text`` state key — the messages channel after
# compaction holds only the preserved recent tail (already-seen messages),
# NOT a synthetic summary AIMessage. So the contraction chunk is the
# already-seen tail, deduped by id; the real regression coverage is that
# POST-compaction growth is still captured.
captured: list[dict] = []
seen: set[str] = set()
# Pre-compaction: a normal growing turn captures 3 steps (cursor → 4).
ai1 = AIMessage(content="searching", id="ai-1")
tool1 = ToolMessage(content="r1", tool_call_id="c1", name="web_search", id="tool-1")
ai2 = AIMessage(content="done turn", id="ai-2")
before = [HumanMessage(content="do research", id="h-1"), ai1, tool1, ai2]
processed = capture_new_step_messages(before, captured, seen, 0)
assert processed == 4
assert [c["id"] for c in captured] == ["ai-1", "tool-1", "ai-2"]
# Compaction rewrites the channel to just the preserved tail (ai2) —
# length drops from 4 to 1, below the cursor. ai2 is already seen, so the
# dedup makes it a no-op; no new step is emitted. (The summary itself lives
# in summary_text and is never a capturable AIMessage — see step_events.py
# INVARIANT.)
compacted = [ai2]
processed = capture_new_step_messages(compacted, captured, seen, processed)
assert processed == 1
assert [c["id"] for c in captured] == ["ai-1", "tool-1", "ai-2"] # unchanged
# Post-compaction growth: a new turn appends after the preserved tail. This
# is the bug the fix targets — without the reset, processed_count stays at
# 4, total (3) never exceeds it, and tool-2/ai-3 are silently dropped.
tool2 = ToolMessage(content="r2", tool_call_id="c2", name="read_file", id="tool-2")
ai3 = AIMessage(content="final answer", id="ai-3")
after = [ai2, tool2, ai3]
processed = capture_new_step_messages(after, captured, seen, processed)
assert processed == 3
assert [c["id"] for c in captured] == ["ai-1", "tool-1", "ai-2", "tool-2", "ai-3"]
def test_run_event_for_task_started():
record = subagent_run_event({"type": "task_started", "task_id": "call_1", "description": "research X"})
assert record["event_type"] == "subagent.start"
assert record["category"] == SUBAGENT_EVENT_CATEGORY
assert record["metadata"]["task_id"] == "call_1"
assert record["content"]["description"] == "research X"
def test_run_event_for_task_running_carries_step_payload():
chunk = {
"type": "task_running",
"task_id": "call_1",
"message": {"type": "tool", "name": "web_search", "content": "results"},
"message_index": 2,
}
record = subagent_run_event(chunk)
assert record["event_type"] == "subagent.step"
assert record["category"] == SUBAGENT_EVENT_CATEGORY
assert record["metadata"] == {"task_id": "call_1", "message_index": 2}
assert record["content"] == build_subagent_step(chunk["message"], task_id="call_1", message_index=2)
def test_run_event_for_terminal_status():
record = subagent_run_event(
{
"type": "task_completed",
"task_id": "call_1",
"result": "done",
"model_name": "claude-3-7-sonnet",
"usage": {"input_tokens": 100, "output_tokens": 20, "total_tokens": 120},
}
)
assert record["event_type"] == "subagent.end"
assert record["content"]["status"] == "completed"
assert record["content"]["result"] == "done"
assert record["content"]["model_name"] == "claude-3-7-sonnet"
assert record["content"]["usage"]["total_tokens"] == 120
failed = subagent_run_event({"type": "task_failed", "task_id": "call_1", "error": "boom"})
assert failed["content"]["status"] == "failed"
assert failed["content"]["error"] == "boom"
def test_run_event_terminal_result_is_truncated():
big = "z" * (SUBAGENT_STEP_MAX_CHARS + 100)
record = subagent_run_event({"type": "task_completed", "task_id": "c1", "result": big})
assert len(record["content"]["result"]) == SUBAGENT_STEP_MAX_CHARS
assert record["content"]["result_truncated"] is True
def test_run_event_ignores_non_task_chunks():
assert subagent_run_event({"type": "something_else"}) is None
assert subagent_run_event({"no_type": True}) is None
assert subagent_run_event("not-a-dict") is None