deer-flow/backend/tests/test_artifact_registry.py
李靖 10bd065615
feat(agents): persist tool artifacts as resolvable handles (#4676) (#4929)
* feat(agents): persist tool artifacts as resolvable handles (#4676)

MCP and sandbox tools return file paths, URLs, and task ids that are
lost when context compaction summarizes the conversation. This adds a
per-thread artifact handle registry so the model can reference
artifacts by short deterministic handles (art_xxxxxxxx) that survive
compaction and resolve to real references at the tool-call boundary.

- ThreadState.tool_artifacts: reducer channel capped at 100 entries,
  same-handle latest-wins, survives summarization compaction
- artifact_registry: deterministic handle derivation from
  (thread_id, tool_call_id, call_index); extraction precedence
  structured_content -> file/image blocks -> conservative text scan;
  HTML-escaped model-context rendering with char budget
- ArtifactCaptureMiddleware: before_model hook captures lightweight
  metadata (never content) and tracks handle consumption
- ArtifactResolutionMiddleware: rewrites handles in tool-call args to
  real references before execution; unknown handles untouched
- DurableContextMiddleware: projects the registry into the hidden
  durable-context data block ([available]/[consumed] status)
- config: new tool_artifacts section (all switches default on)
- frontend: artifact badges on tool-call steps keyed by tool_call_id

Fixes #4676

* fix(agents): harden artifact registry after review (#4676)

Review findings on #4929:

- capture/consumption updates emitted in the same before_model call are
  now concatenated instead of dict-merged; dict-merging clobbered the
  shared tool_artifacts key and permanently dropped fresh captures
- handles embed a per-result ordinal so multiple references inside one
  tool result never collide (two paths in one text block previously
  collapsed into one handle)
- structured_content extraction pulls concrete refs from known keys
  (file/path/url/task_id/..., string or list-of-string values) instead
  of storing a truncated JSON preview, falls back to an untruncated
  whole-object data entry, and no longer skips file/image content blocks
- plain-string tool results (e.g. bash/write_file output) are scanned
  for sandbox paths / remote URLs too, gated by detect_refs_in_text
- wire max_entries through configure_tool_artifact_max_entries at
  assembly (clamped 10..1000) and gate detect_refs_in_text scans
- resolution requires both enabled and resolve_handles_in_args,
  matching its docs
- add real ToolCallRequest.override coverage and the spec 14.4
  integration cycle (capture -> inject -> resolve -> consumption)

* fix(agents): address second-round artifact registry review (#4676)

Follow-up findings on #4929:

- drop generic `output`/`outputs` from structured ref keys: prose under
  those keys no longer becomes a bogus file entry (falls back to the
  whole-object data entry; path-shaped refs are still caught by the
  free-text scan)
- strip closing quotes, brackets, and backticks from detected refs so
  markdown- and JSON-wrapped paths resolve to clean real references
- remove the process-wide cap setter: the reducer enforces only an
  absolute ceiling (1000) and ArtifactCaptureMiddleware applies the
  configured max_entries per agent when emitting updates, so agents
  with different configs in one process never share or flap a cap
- skip extraction for tool messages whose tool_call_id already has
  registry entries and memoize known-empty results (bounded FIFO), so
  steady-state capture cost drops to the new message tail instead of a
  full-history rescan per model call
- correct the inverted ordering claim: capture is position-independent
  in the wrap chain because it is a before_model hook reading state
  messages; fixed the assembly comment, AGENTS.md item 15, and the
  ordering test docstring

CI parity verified locally: ruff check + format --check, full offline
suite, blocking-io suite, collect-only, uv lock --check,
check_agent_guidance, frontend format/lint/typecheck/build/tests.

* fix(agents): sliding-window cap and URI gates for artifact registry (#4676)

Third-round review findings on #4929:

- max_entries is now a sliding retention window: fresh captures are
  always registered, and when the projection exceeds the configured
  cap the middleware appends a trailing {"op": "trim_to", "keep": N}
  directive that the reducer applies by evicting the oldest entries.
  The previous remaining-budget emission froze the registry at the cap,
  permanently dropping post-compaction artifacts — the exact failure
  mode #4676 targets — and silently disabled the round-2 rescan
  optimizations at the point threads grew longest
- _track_consumption memoizes settled tool-call args in a bounded FIFO
  alongside the existing empty-result memo, so historical AIMessage
  args are not regex-rescanned every model call
- file/image content-block URLs must be http(s) or path-shaped;
  data:/blob: URIs (MCP embedded resources carrying base64 payloads)
  can no longer enter thread state, tool args, or crowd out the render
  budget; protocol-relative //host forms are rejected as well
- structured-content ref values go through the same URI-shape gate as
  content blocks, closing the embedded-payload hole via known keys such
  as url/urls
- tool_artifacts.enabled now gates the whole middleware including
  consumption tracking via an early return in before_model

Docs updated to describe the sliding window, URI gates, enabled scope,
and both bounded memos.

* fix(agents): final evictions and pending-aware consumption for artifact registry (#4676)

Fourth-round review findings on #4929:

- the sliding window interacted with the seen_call_ids skip to
  resurrect evicted entries every round: once an entry was trimmed, its
  tool_call_id left the seen set while its ToolMessage stayed in
  context, so the next model call re-extracted it, re-registered it as
  newest, and evicted something else — a never-converging rotation.
  The empty-result memo is now a unified _processed_results memo that
  records every extraction regardless of outcome, making evictions
  final; steady state over a saturated thread is quiescent (pinned by
  a multi-round test asserting no updates and zero extractions)
- the consumption scan built its handle map from pre-capture state, so
  a call referencing a handle captured in the same before_model call
  missed, got settled as quiet, and its consumption mark was lost
  forever. _track_consumption now receives the pending capture entries
  and resolves against state + pending; a call settles as quiet only
  after every referenced handle resolved, so transient misses retry
  next round instead of permanently suppressing the mark

AGENTS.md entry 15 updated to describe both memos and the
eviction-is-final semantics.

* fix(artifacts): address occurrence identity, PII and URL review findings

* fix(artifacts): enforce resolved policy checks and bounded registry handling

---------

Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
2026-09-27 11:09:41 +08:00

435 lines
16 KiB
Python

"""Tests for the persistent artifact handle registry (issue #4676)."""
import json
from langchain_core.messages import ToolMessage
from deerflow.agents.thread_state import merge_tool_artifacts
from deerflow.tools.artifact_registry import (
_detect_refs_in_text,
extract_artifacts_from_result,
generate_handle,
)
def _entry(handle: str, i: int) -> dict:
return {
"handle": handle,
"tool_name": "t",
"tool_call_id": f"call-{i}",
"call_index": 0,
"artifact_type": "file",
"display_name": f"{i}.txt",
"real_ref": f"/tmp/{i}.txt",
"created_at": "2026-08-19T00:00:00Z",
}
def test_generate_handle_deterministic():
assert generate_handle("thread-1", "call-1", 0) == generate_handle("thread-1", "call-1", 0)
def test_generate_handle_unique():
assert generate_handle("thread-1", "call-1", 0) != generate_handle("thread-1", "call-1", 1)
assert generate_handle("thread-1", "call-1", 0) != generate_handle("thread-1", "call-2", 0)
assert generate_handle("thread-1", "call-1", 0) != generate_handle("thread-2", "call-1", 0)
def test_generate_handle_format():
handle = generate_handle("thread-1", "call-1", 0)
assert handle.startswith("art_")
assert len(handle) == len("art_") + 8
def test_extract_from_file_block():
result = ToolMessage(
content=[
{
"type": "file",
"source": {"type": "url", "url": "/mnt/user-data/outputs/report.html", "mime_type": "text/html"},
}
],
tool_call_id="call-1",
name="mcp_server_analyze",
)
entries = extract_artifacts_from_result(result, thread_id="thread-1")
assert len(entries) == 1
entry = entries[0]
assert entry["artifact_type"] == "file"
assert entry["real_ref"] == "/mnt/user-data/outputs/report.html"
assert entry["display_name"] == "report.html"
assert entry["tool_name"] == "mcp_server_analyze"
assert entry["tool_call_id"] == "call-1"
assert entry["mime_type"] == "text/html"
def test_extract_from_image_block():
result = ToolMessage(
content=[
{
"type": "image",
"source": {"type": "url", "url": "/mnt/user-data/outputs/chart.png", "mime_type": "image/png"},
}
],
tool_call_id="call-2",
name="mcp_chart_gen",
)
entries = extract_artifacts_from_result(result, thread_id="thread-1")
assert len(entries) == 1
assert entries[0]["artifact_type"] == "image"
assert entries[0]["real_ref"] == "/mnt/user-data/outputs/chart.png"
def test_file_block_rejects_data_and_blob_urls():
"""Embedded-resource URIs must never enter the registry or tool args."""
result = ToolMessage(
content=[
{"type": "file", "source": {"type": "url", "url": "data:text/plain;base64,QUFB" + "A" * 500}},
{"type": "image", "source": {"type": "url", "url": "blob:https://example.com/uuid"}},
{"type": "file", "source": {"type": "url", "url": "//cdn.example.com/pixel.gif"}},
{"type": "file", "source": {"type": "url", "url": "/mnt/user-data/outputs/real.txt"}},
],
tool_call_id="call-2b",
name="mcp_embedded",
)
entries = extract_artifacts_from_result(result, thread_id="thread-1")
refs = {entry["real_ref"] for entry in entries}
assert refs == {"/mnt/user-data/outputs/real.txt"}
def test_structured_url_key_rejects_non_http_uris():
"""Structured refs get the same URI-shape gate as content blocks."""
result = ToolMessage(
content=[{"type": "text", "text": "done"}],
tool_call_id="call-4g",
name="mcp_embedded",
artifact={
"structured_content": {
"url": "data:text/html;base64,PGh0bWw+" + "Q" * 300,
"file": "/mnt/user-data/outputs/keep.md",
"task_id": "job-11",
}
},
)
entries = extract_artifacts_from_result(result, thread_id="thread-1")
refs = {entry["real_ref"] for entry in entries}
assert refs == {"/mnt/user-data/outputs/keep.md", "job-11"}
types = {entry["real_ref"]: entry["artifact_type"] for entry in entries}
assert types["job-11"] == "task"
def test_extract_from_text_with_path():
result = ToolMessage(
content=[{"type": "text", "text": "Report saved to /mnt/user-data/outputs/report.html"}],
tool_call_id="call-3",
name="mcp_server_analyze",
)
entries = extract_artifacts_from_result(result, thread_id="thread-1")
assert len(entries) == 1
assert entries[0]["artifact_type"] == "file"
assert entries[0]["real_ref"] == "/mnt/user-data/outputs/report.html"
def test_extract_from_structured_content():
result = ToolMessage(
content=[{"type": "text", "text": "task submitted"}],
tool_call_id="call-4",
name="mcp_task_submit",
artifact={"structured_content": {"task_id": "remote-task-42", "status": "working"}},
)
entries = extract_artifacts_from_result(result, thread_id="thread-1")
assert len(entries) == 1
assert entries[0]["artifact_type"] == "task"
assert entries[0]["real_ref"] == "remote-task-42"
assert entries[0]["display_name"] == "remote-task-42"
def test_extract_from_structured_file_key_is_concrete_ref():
result = ToolMessage(
content=[{"type": "text", "text": "saved"}],
tool_call_id="call-4b",
name="mcp_writer",
artifact={"structured_content": {"file": "/mnt/user-data/outputs/report.md", "mime_type": "text/markdown"}},
)
entries = extract_artifacts_from_result(result, thread_id="thread-1")
assert len(entries) == 1
assert entries[0]["artifact_type"] == "file"
assert entries[0]["real_ref"] == "/mnt/user-data/outputs/report.md"
def test_extract_from_structured_list_valued_keys():
result = ToolMessage(
content=[{"type": "text", "text": "saved"}],
tool_call_id="call-4e",
name="mcp_batch_writer",
artifact={"structured_content": {"files": ["/mnt/user-data/outputs/a.txt", "/mnt/user-data/outputs/b.txt"], "task_ids": ["x"]}},
)
entries = extract_artifacts_from_result(result, thread_id="thread-1")
refs = {entry["real_ref"] for entry in entries}
assert refs == {"/mnt/user-data/outputs/a.txt", "/mnt/user-data/outputs/b.txt"}
handles = {entry["handle"] for entry in entries}
assert len(handles) == 2
def test_structured_generic_output_key_not_treated_as_ref():
"""Prose under generic result keys must not become a bogus file entry."""
result = ToolMessage(
content=[{"type": "text", "text": "done"}],
tool_call_id="call-4f",
name="mcp_analyst",
artifact={"structured_content": {"output": "Analysis complete; revenue up 12% quarter over quarter."}},
)
entries = extract_artifacts_from_result(result, thread_id="thread-1")
assert len(entries) == 1
assert entries[0]["artifact_type"] == "data"
assert "Analysis complete" in entries[0]["real_ref"]
def test_structured_fallback_whole_object_not_truncated():
structured = {"custom_payload": {"nested": ["x" * 200] * 10}}
result = ToolMessage(
content=[{"type": "text", "text": "done"}],
tool_call_id="call-4c",
name="mcp_exotic",
artifact={"structured_content": structured},
)
entries = extract_artifacts_from_result(result, thread_id="thread-1")
assert len(entries) == 1
assert entries[0]["artifact_type"] == "data"
assert entries[0]["real_ref"] == json.dumps(structured, ensure_ascii=False)
def test_structured_content_does_not_skip_file_blocks():
result = ToolMessage(
content=[
{
"type": "file",
"source": {"type": "url", "url": "/mnt/user-data/outputs/chart.png", "mime_type": "image/png"},
}
],
tool_call_id="call-4d",
name="mcp_mixed",
artifact={"structured_content": {"task_id": "job-7"}},
)
entries = extract_artifacts_from_result(result, thread_id="thread-1")
refs = {entry["real_ref"] for entry in entries}
assert "job-7" in refs
assert "/mnt/user-data/outputs/chart.png" in refs
def test_two_refs_in_one_text_block_get_distinct_handles():
result = ToolMessage(
content=[{"type": "text", "text": "Saved /mnt/user-data/outputs/a.txt and /mnt/user-data/outputs/b.csv"}],
tool_call_id="call-3b",
name="mcp_server_analyze",
)
entries = extract_artifacts_from_result(result, thread_id="thread-1")
assert len(entries) == 2
handles = {entry["handle"] for entry in entries}
assert len(handles) == 2
assert {entry["real_ref"] for entry in entries} == {"/mnt/user-data/outputs/a.txt", "/mnt/user-data/outputs/b.csv"}
def test_string_content_with_path_captured():
result = ToolMessage(
content="Wrote /mnt/user-data/outputs/notes.md successfully",
tool_call_id="call-6b",
name="write_file",
)
entries = extract_artifacts_from_result(result, thread_id="thread-1")
assert len(entries) == 1
assert entries[0]["real_ref"] == "/mnt/user-data/outputs/notes.md"
assert entries[0]["artifact_type"] == "file"
def test_detect_refs_in_text_flag_disables_text_scans():
result = ToolMessage(
content=[{"type": "text", "text": "Report saved to /mnt/user-data/outputs/report.html"}],
tool_call_id="call-3c",
name="mcp_server_analyze",
)
entries = extract_artifacts_from_result(result, thread_id="thread-1", detect_refs_in_text=False)
assert entries == []
def test_no_extraction_from_error_result():
result = ToolMessage(
content="Error: Tool failed",
tool_call_id="call-5",
name="mcp_server_analyze",
status="error",
)
entries = extract_artifacts_from_result(result, thread_id="thread-1")
assert entries == []
def test_no_extraction_from_plain_text_without_refs():
result = ToolMessage(content="All done", tool_call_id="call-6", name="mcp_server_analyze")
entries = extract_artifacts_from_result(result, thread_id="thread-1")
assert entries == []
def test_detect_refs_in_text_paths():
refs = _detect_refs_in_text("See /mnt/user-data/outputs/a.txt and /mnt/user-data/outputs/b.csv here.")
paths = [r["ref"] for r in refs]
assert "/mnt/user-data/outputs/a.txt" in paths
assert "/mnt/user-data/outputs/b.csv" in paths
def test_detect_refs_in_text_urls():
refs = _detect_refs_in_text("Download https://example.com/files/report.pdf now.")
assert len(refs) == 1
assert refs[0]["ref"] == "https://example.com/files/report.pdf"
def test_detect_refs_in_text_noise():
refs = _detect_refs_in_text("No references here, just text.")
assert refs == []
def test_detect_refs_in_text_strips_quotes_backticks_brackets():
text = 'Saved to `/mnt/user-data/outputs/report.md` and {"path": "/mnt/user-data/a.txt"} plus [/mnt/user-data/c.csv]'
values = [ref["ref"] for ref in _detect_refs_in_text(text)]
assert "/mnt/user-data/outputs/report.md" in values
assert "/mnt/user-data/a.txt" in values
assert "/mnt/user-data/c.csv" in values
def test_string_content_with_quoted_path_captured_cleanly():
result = ToolMessage(
content='Wrote "/mnt/user-data/outputs/notes.md" successfully',
tool_call_id="call-6c",
name="write_file",
)
entries = extract_artifacts_from_result(result, thread_id="thread-1")
assert len(entries) == 1
assert entries[0]["real_ref"] == "/mnt/user-data/outputs/notes.md"
def test_merge_tool_artifacts_append():
existing = [
{
"handle": "art_00000001",
"tool_name": "t1",
"tool_call_id": "call-1",
"call_index": 0,
"artifact_type": "file",
"display_name": "a.txt",
"real_ref": "/tmp/a.txt",
"created_at": "2026-08-19T00:00:00Z",
}
]
new = [
{
"handle": "art_00000002",
"tool_name": "t2",
"tool_call_id": "call-2",
"call_index": 0,
"artifact_type": "file",
"display_name": "b.txt",
"real_ref": "/tmp/b.txt",
"created_at": "2026-08-19T00:00:01Z",
}
]
merged = merge_tool_artifacts(existing, new)
assert len(merged) == 2
assert merged[0]["handle"] == "art_00000001"
assert merged[1]["handle"] == "art_00000002"
def test_merge_tool_artifacts_dedup_same_handle_latest_wins():
existing = [
{
"handle": "art_00000001",
"tool_name": "t1",
"tool_call_id": "call-1",
"call_index": 0,
"artifact_type": "file",
"display_name": "a.txt",
"real_ref": "/tmp/a.txt",
"created_at": "2026-08-19T00:00:00Z",
}
]
new = [
{
"handle": "art_00000001",
"tool_name": "t1",
"tool_call_id": "call-1",
"call_index": 0,
"artifact_type": "file",
"display_name": "a.txt",
"real_ref": "/tmp/a-new.txt",
"consumed_by": ["call-9"],
"created_at": "2026-08-19T00:00:01Z",
}
]
merged = merge_tool_artifacts(existing, new)
assert len(merged) == 1
assert merged[0]["real_ref"] == "/tmp/a-new.txt"
assert merged[0]["consumed_by"] == ["call-9"]
def test_merge_tool_artifacts_empty_new_preserves_existing():
existing = [
{
"handle": "art_00000001",
"tool_name": "t1",
"tool_call_id": "call-1",
"call_index": 0,
"artifact_type": "file",
"display_name": "a.txt",
"real_ref": "/tmp/a.txt",
"created_at": "2026-08-19T00:00:00Z",
}
]
merged = merge_tool_artifacts(existing, None)
assert merged == existing
merged2 = merge_tool_artifacts(existing, [])
assert merged2 == existing
def test_merge_tool_artifacts_trim_directive_slides_window():
"""A trailing trim directive makes the configured cap a sliding window."""
existing = [_entry(f"art_{i:08x}", i) for i in range(20)]
fresh = [_entry(f"art_new{i:02x}", 100 + i) for i in range(2)]
update = [*fresh, {"op": "trim_to", "keep": 20}]
merged = merge_tool_artifacts(existing, update)
assert len(merged) == 20
assert merged[0]["handle"] == "art_00000002", "oldest two must be evicted"
handles = {entry["handle"] for entry in merged}
assert {f"art_new{i:02x}" for i in range(2)} <= handles, "fresh entries must survive"
assert "art_00000000" not in handles and "art_00000001" not in handles
def test_merge_tool_artifacts_no_trim_without_directive():
existing = [_entry(f"art_{i:08x}", i) for i in range(5)]
fresh = [_entry(f"art_new{i:02x}", 50 + i) for i in range(2)]
merged = merge_tool_artifacts(existing, fresh)
assert len(merged) == 7
def test_merge_tool_artifacts_trim_clamped_to_ceiling():
entries = [_entry(f"art_{i:08x}", i) for i in range(1200)]
update = [*entries, {"op": "trim_to", "keep": 5000}]
merged = merge_tool_artifacts(None, update)
assert len(merged) == 1000
assert merged[0]["handle"] == f"art_{200:08x}"
def test_merge_tool_artifacts_cap():
"""Reducer enforces the absolute ceiling only; configured caps live in the middleware."""
entries = [_entry(f"art_{i:08x}", i) for i in range(1500)]
merged = merge_tool_artifacts(None, entries)
assert len(merged) == 1000
assert merged[0]["handle"] == f"art_{500:08x}"
assert merged[-1]["handle"] == f"art_{1499:08x}"
def test_merge_tool_artifacts_consumption_update_does_not_grow_count():
entries = [_entry(f"art_{i:08x}", i) for i in range(1000)]
consumed = {**entries[0], "consumed_by": ["call-x"]}
merged = merge_tool_artifacts(entries, [consumed])
assert len(merged) == 1000