351 lines
13 KiB
Python
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
|