"""Comprehensive tests for ToolOutputBudgetMiddleware. Covers: pass-through, disk externalization, fallback truncation, UTF-8 boundaries, Command results, model-request history patching, superseded write_file payload elision (issue #5328), config variations, exempt tools, per-tool overrides, bash exit-marker preservation after budget rewrites, edge cases, and both sync/async code paths. """ from __future__ import annotations import contextlib import csv import hashlib import importlib.util import io import json import os import pathlib import re import sys import tempfile from types import SimpleNamespace import pytest from langchain.agents.middleware.types import ModelRequest from langchain_core.messages import AIMessage, HumanMessage, ToolMessage from langgraph.types import Command from deerflow.agents.middlewares.tool_output_budget_middleware import ( _BASH_EXIT_MARKER_TAIL_RE, TOOL_OUTPUT_BLOB_KEY, ToolOutputBudgetMiddleware, _build_fallback, _build_preview, _effective_trigger, _externalize, _keep_exit_marker_last, _message_text, _needs_budget, _patch_model_messages, _sanitize_tool_name, _snap_start_to_line_boundary, _snap_to_line_boundary, _tool_message_over_budget, ) from deerflow.agents.middlewares.tool_output_synopsis import build_tool_output_synopsis from deerflow.config.app_config import AppConfig from deerflow.config.sandbox_config import SandboxConfig from deerflow.config.tool_output_config import ToolOutputConfig from deerflow.storage import BlobReadError, BlobRef, BlobWriteError # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _lines_then_long_line(total: int, newline_ratio: float = 0.6) -> str: """Content that is line-oriented for the first *newline_ratio*, then one unbroken line. Mirrors real bash/web_fetch output that logs progress lines and then dumps a single-line artifact (minified JSON, base64 blob). The last newline lands in the second half of the content, which is what exercises line snapping around the tail offset. """ head_len = int(total * newline_ratio) lines = "".join(f"[info] step {i} ok\n" for i in range(head_len // 18 + 1))[:head_len] lines = lines[:-1] + "\n" if not lines.endswith("\n") else lines return lines + "A" * (total - len(lines)) def _make_request(tool_name: str = "remote_executor", tool_call_id: str = "tc-1", outputs_path: str | None = None) -> SimpleNamespace: thread_data = {"outputs_path": outputs_path} if outputs_path else None state = {"thread_data": thread_data} if thread_data else {} runtime = SimpleNamespace(state=state, context={"thread_id": "thread-1"}, config={}) return SimpleNamespace( tool_call={"name": tool_name, "id": tool_call_id}, runtime=runtime, ) @contextlib.contextmanager def _unwritable_outputs_path(): """Yield an ``outputs_path`` that ``os.makedirs`` cannot create, on any platform. The parent component is a regular file, so creating a directory below it fails with an ``OSError`` subclass everywhere (``NotADirectoryError`` on POSIX, ``FileNotFoundError`` on Windows) and nothing is written outside the temporary directory. This deliberately avoids expressing "unwritable" as a magic absolute path. ``/nonexistent/...`` was creatable by root in the CI container, and its replacement ``/dev/null/...`` relies on ``/dev/null`` being a character device, which is only true on POSIX -- on Windows it is an ordinary relative path that ``os.makedirs`` happily creates at the drive root. """ with tempfile.TemporaryDirectory() as tmpdir: blocker = pathlib.Path(tmpdir) / "not-a-directory" blocker.touch() yield os.path.join(blocker, "outputs") def _tm(content: str = "ok", name: str = "tool", tool_call_id: str = "tc-1") -> ToolMessage: return ToolMessage(content=content, name=name, tool_call_id=tool_call_id) def _pytest_like_output(exit_line: str = "Exit Code: 1") -> str: """Pytest-shaped bash output in the default 12k–20k budget window.""" return "============================= test session starts =============================\n" + ("collected 12 items\n" + "." * 12 + "\n") * 500 + f"12 passed in 1.23s\n{exit_line}" def _bash_tm(content: str, name: str = "bash") -> ToolMessage: return ToolMessage( content=content, name=name, tool_call_id="tc-1", id="tool-msg-1", artifact={"k": "v"}, additional_kwargs={"deerflow_tool_meta": {"status": "success", "source": "normalized"}}, ) # =========================================================================== # Unit tests for helper functions # =========================================================================== class TestMessageText: def test_string_content(self): assert _message_text("hello") == "hello" def test_none_content(self): assert _message_text(None) is None def test_list_of_strings(self): assert _message_text(["a", "b"]) == "a\nb" def test_list_of_text_dicts(self): assert _message_text([{"text": "x"}, {"text": "y"}]) == "x\ny" def test_list_with_image_returns_none(self): assert _message_text([{"type": "image", "data": "..."}]) is None def test_empty_list(self): assert _message_text([]) is None def test_non_string_non_list(self): assert _message_text(42) is None def test_json_block_is_rendered(self): assert _message_text([{"type": "json", "json": {"a": 1}}]) == '{"a": 1}' def test_json_block_mixed_with_text(self): assert _message_text([{"text": "rows:"}, {"type": "json", "json": [1, 2]}]) == "rows:\n[1, 2]" def test_json_block_non_serializable_falls_back_to_str(self): assert _message_text([{"type": "json", "json": {"bad": {1}}}]) == "{'bad': {1}}" def test_json_block_circular_falls_back_to_str(self): payload: dict = {} payload["self"] = payload result = _message_text([{"type": "json", "json": payload}]) # repr of a recursive dict differs across versions ("..." vs "{...}") assert result is not None and result.startswith("{'self': ") and "..." in result def test_json_block_without_json_key_returns_none(self): assert _message_text([{"type": "json"}]) is None class TestStructuredJsonMindIEBudget: def test_mixed_json_media_externalizes_without_losing_media(self, tmp_path): from deerflow.models.mindie_provider import _fix_messages payload = {"rows": ["x" * 10_000], "answer": "TAIL_SENTINEL"} media = {"type": "image", "base64": "BINARY_SENTINEL", "mime_type": "image/png"} message = ToolMessage( content=[{"type": "json", "json": payload}, media, {"type": "json"}], name="query_rows", tool_call_id="call_rows", artifact={"source": "rows"}, ) middleware = ToolOutputBudgetMiddleware(config=ToolOutputConfig(externalize_min_chars=100, preview_head_chars=20, preview_tail_chars=10)) result = middleware.wrap_tool_call(_make_request(tool_name="query_rows", outputs_path=str(tmp_path)), lambda _: message) assert result is not message assert isinstance(result.content, list) assert media in result.content assert result.artifact == message.artifact files = list((tmp_path / ".tool-results").iterdir()) assert len(files) == 1 assert json.loads(files[0].read_text(encoding="utf-8")) == payload visible = _fix_messages([result])[0].content assert "Full query_rows output saved to" in visible assert "BINARY_SENTINEL" not in visible assert "x" * 10_000 not in visible assert message.content[0]["json"] == payload @pytest.mark.anyio async def test_mixed_json_media_history_is_budgeted_before_mindie(self): from deerflow.models.mindie_provider import _fix_messages payload = {"rows": "x" * 2000, "answer": "TAIL_SENTINEL"} media = {"type": "image", "base64": "BINARY_SENTINEL", "mime_type": "image/png"} message = ToolMessage(content=[{"type": "json", "json": payload}, media], name="query_rows", tool_call_id="call_history") request = ModelRequest(model=None, messages=[message], tools=[], state={}) middleware = ToolOutputBudgetMiddleware(config=ToolOutputConfig(externalize_min_chars=100, fallback_max_chars=300, fallback_head_chars=80, fallback_tail_chars=40)) captured = {} async def handler(req): captured["request"] = req return [] await middleware.awrap_model_call(request, handler) forwarded = captured["request"] assert forwarded is not request assert media in forwarded.messages[0].content visible = _fix_messages(forwarded.messages)[0].content assert len(visible) <= 333 # Configured text limit plus XML framing. assert "TAIL_SENTINEL" in visible assert "BINARY_SENTINEL" not in visible assert message.content[0]["json"] == payload @pytest.mark.parametrize("structured", [False, True], ids=["plain-text", "json"]) @pytest.mark.parametrize( "config,tool_name", [ (ToolOutputConfig(enabled=False), "query_rows"), (ToolOutputConfig(externalize_min_chars=60_000, fallback_max_chars=60_000), "query_rows"), (ToolOutputConfig(), "read_file"), ], ids=["disabled", "increased-limits", "exempt-read"], ) def test_configured_passthrough_survives_provider_normalization(self, config, tool_name, structured): from deerflow.models.mindie_provider import _fix_messages payload = "x" * 35_000 + "TAIL_SENTINEL" content = [{"type": "json", "json": {"rows": payload}}] if structured else payload message = ToolMessage(content=content, name=tool_name, tool_call_id="call_passthrough") middleware = ToolOutputBudgetMiddleware(config=config) result = middleware.wrap_tool_call(_make_request(tool_name=tool_name), lambda _: message) assert result is message assert payload in _fix_messages([result])[0].content class TestSnapToLineBoundary: def test_snaps_to_newline(self): text = "line1\nline2\nline3" pos = 14 # inside "line3" result = _snap_to_line_boundary(text, pos) assert text[result - 1] == "\n" def test_no_snap_when_no_newline_in_range(self): text = "abcdefghij" assert _snap_to_line_boundary(text, 8) == 8 def test_zero_pos(self): assert _snap_to_line_boundary("abc", 0) == 0 def test_pos_beyond_length(self): assert _snap_to_line_boundary("abc", 10) == 10 class TestSnapStartToLineBoundary: def test_snaps_forward_to_newline(self): text = "line1\nline2\nline3" result = _snap_start_to_line_boundary(text, 2) # inside "line1" assert text[result - 1] == "\n" assert result >= 2 def test_never_moves_backwards(self): text = "aaaa\n" + "b" * 20 for pos in range(1, len(text)): assert _snap_start_to_line_boundary(text, pos) >= pos def test_no_snap_when_no_newline_in_range(self): assert _snap_start_to_line_boundary("abcdefghij", 2) == 2 def test_zero_pos(self): assert _snap_start_to_line_boundary("a\nbc", 0) == 0 def test_pos_beyond_length(self): assert _snap_start_to_line_boundary("abc", 10) == 10 class TestExternalize: def test_writes_file_and_returns_virtual_path(self): with tempfile.TemporaryDirectory() as tmpdir: path = _externalize( "full content here", tool_name="bash", tool_call_id="tc-1", outputs_path=tmpdir, storage_subdir=".tool-results", ) assert path is not None assert path.startswith("/mnt/user-data/outputs/.tool-results/bash-") assert path.endswith(".log") # Verify actual file on disk storage_dir = os.path.join(tmpdir, ".tool-results") files = os.listdir(storage_dir) assert len(files) == 1 with open(os.path.join(storage_dir, files[0]), encoding="utf-8") as f: assert f.read() == "full content here" def test_returns_none_on_invalid_path(self): with _unwritable_outputs_path() as outputs_path: path = _externalize( "data", tool_name="test", tool_call_id="tc-1", outputs_path=outputs_path, storage_subdir=".tool-results", ) assert path is None def test_txt_extension_for_unknown_tool(self): with tempfile.TemporaryDirectory() as tmpdir: path = _externalize( "data", tool_name="unknown_tool", tool_call_id="tc-1", outputs_path=tmpdir, storage_subdir=".tool-results", ) assert path is not None assert path.endswith(".txt") class TestSanitizeToolName: def test_strips_path_separators(self): assert _sanitize_tool_name("../../etc/passwd") == "passwd" def test_strips_backslashes(self): result = _sanitize_tool_name("..\\..\\windows\\system32") assert ".." not in result assert "/" not in result def test_normal_name_unchanged(self): assert _sanitize_tool_name("bash") == "bash" def test_empty_becomes_unknown(self): assert _sanitize_tool_name("") == "unknown" def test_dots_only_becomes_unknown(self): assert _sanitize_tool_name("..") == "unknown" class TestExternalizePathTraversal: def test_traversal_tool_name_is_sanitized(self): with tempfile.TemporaryDirectory() as tmpdir: path = _externalize( "data", tool_name="../../etc/passwd", tool_call_id="tc-1", outputs_path=tmpdir, storage_subdir=".tool-results", ) assert path is not None assert "passwd-" in path assert "../" not in path def test_absolute_storage_subdir_rejected(self): path = _externalize( "data", tool_name="tool", tool_call_id="tc-1", outputs_path="/tmp", storage_subdir="/etc/evil", ) assert path is None def test_traversal_storage_subdir_rejected(self): path = _externalize( "data", tool_name="tool", tool_call_id="tc-1", outputs_path="/tmp", storage_subdir="../../../etc", ) assert path is None class TestStorageSubdirConfig: """storage_subdir must be a single path segment. The workspace-changes scanner prunes by directory name during os.walk, which yields one-segment dirnames — a nested value like ``cache/tool-results`` would silently never match the exclusion and its files would be counted as produced artifacts again. Rejecting it at config time keeps the dir-name-based exclusion sound. """ def test_default_is_single_segment(self): assert ToolOutputConfig().storage_subdir == ".tool-results" def test_single_segment_custom_accepted(self): assert ToolOutputConfig(storage_subdir="tool-output-cache").storage_subdir == "tool-output-cache" def test_nested_storage_subdir_rejected(self): with pytest.raises(ValueError): ToolOutputConfig(storage_subdir="cache/tool-results") def test_windows_separator_storage_subdir_rejected(self): with pytest.raises(ValueError): ToolOutputConfig(storage_subdir="cache\\tool-results") def test_absolute_storage_subdir_rejected(self): with pytest.raises(ValueError): ToolOutputConfig(storage_subdir="/abs/path") def test_dot_and_dotdot_storage_subdir_rejected(self): with pytest.raises(ValueError): ToolOutputConfig(storage_subdir=".") with pytest.raises(ValueError): ToolOutputConfig(storage_subdir="..") def test_empty_storage_subdir_rejected(self): with pytest.raises(ValueError): ToolOutputConfig(storage_subdir="") class TestNeedsBudget: def test_small_output_does_not_need_budget(self): config = ToolOutputConfig(externalize_min_chars=1000) msg = _tm("small", name="tool") assert _needs_budget(msg, config) is False def test_large_output_needs_budget(self): config = ToolOutputConfig(externalize_min_chars=50) msg = _tm("x" * 100, name="tool") assert _needs_budget(msg, config) is True def test_exempt_tool_does_not_need_budget(self): config = ToolOutputConfig(externalize_min_chars=10) msg = _tm("x" * 100, name="read_file") assert _needs_budget(msg, config) is False def test_multimodal_does_not_need_budget(self): config = ToolOutputConfig(externalize_min_chars=10) msg = ToolMessage(content=[{"type": "image", "data": "x" * 100}], name="tool", tool_call_id="tc-1") assert _needs_budget(msg, config) is False def test_structured_json_output_needs_budget(self): config = ToolOutputConfig(externalize_min_chars=50) msg = ToolMessage(content=[{"type": "json", "json": {"rows": ["x" * 100]}}], name="query_rows", tool_call_id="tc-1") assert _needs_budget(msg, config) is True class TestBuildPreview: def test_contains_typed_summary_and_reference(self): content = "HEAD_" + "x" * 5000 + "_TAIL" preview = _build_preview( content, tool_name="bash", virtual_path="/mnt/test/bash-abc.log", head_chars=100, tail_chars=50, ) assert preview.startswith("[Full bash output saved to /mnt/test/bash-abc.log") assert "Preview kind: text" in preview assert "Text output" in preview assert "/mnt/test/bash-abc.log" in preview assert "read_file" in preview assert "start_line and end_line" in preview def test_reports_total_chars(self): content = "a" * 10000 preview = _build_preview( content, tool_name="web_search", virtual_path="/mnt/test/file.txt", head_chars=200, tail_chars=100, ) assert "10000 chars" in preview def test_json_preview_includes_structure_and_raw_sample(self): content = '{"meta":{"source":"unit"},"items":[{"id":1,"name":"alpha"},{"id":2,"name":"beta"}],"payload":"' + "x" * 5000 + '","tail_marker":"SHOULD_NOT_NEED_TAIL"}' preview = _build_preview( content, tool_name="mcp_json", virtual_path="/mnt/test/result.json", head_chars=80, tail_chars=40, ) assert "Preview kind: json" in preview assert "JSON object with 4 top-level keys" in preview assert "Top-level keys: meta, items, payload, tail_marker" in preview assert "items: array length 2" in preview assert '$.meta.source: "unit"' in preview assert not preview.startswith('{"meta"') # The synopsis no longer hides the raw head/tail bytes; the model # gets the typed synopsis AND inline raw samples so it can read # the file with a tighter start_line range. assert "Raw sample (head + tail" in preview # payload segment dominates the document, so even with head_chars=80 # the raw head sample is almost entirely 'x' characters. assert preview.count("x") >= 50 def test_json_preview_reports_nested_paths_and_line_hints(self): content = json.dumps( { "data": { "items": [{"id": idx, "name": f"item-{idx}"} for idx in range(47)], "next_cursor": "cursor-2", }, "meta": {"source": "unit"}, }, indent=2, ) preview = _build_preview( content, tool_name="api_tool", virtual_path="/mnt/test/api.json", head_chars=80, tail_chars=40, ) assert "$.data: object keys 2; keys items, next_cursor" in preview assert "$.data.items: array length 47; first item object" in preview # Location hints were removed: they are wrong when a key string also # appears as a value earlier in the document, or when the same key # recurs at multiple depths. assert "line " not in preview.split("Access:")[0] assert "byte offset " not in preview def test_json_paths_are_emitted_without_line_hints(self): content = "\n\n" + json.dumps({"data": {"items": [1, 2, 3]}}, indent=2) preview = _build_preview( content, tool_name="api_tool", virtual_path="/mnt/test/api.json", head_chars=80, tail_chars=40, ) assert "$.data: object keys 1; keys items" in preview assert "line " not in preview.split("Access:")[0] assert "byte offset " not in preview def test_table_preview_extracts_columns(self): content = "name,score\n" + "\n".join(f"Ada{i},{90 + i}" for i in range(10)) + "\n" preview = _build_preview( content, tool_name="csv_tool", virtual_path="/mnt/test/table.csv", head_chars=80, tail_chars=40, ) assert "Preview kind: csv" in preview assert "CSV table with 10 data rows and 2 columns" in preview assert "columns: name, score" in preview assert "first data row: name=Ada0 | score=90" in preview class TestToolOutputSynopsis: def test_code_synopsis_extracts_imports_and_symbols(self): content = "import os\nfrom pathlib import Path\n\nclass Runner:\n pass\n\ndef main():\n return Path(os.getcwd())\n" synopsis = build_tool_output_synopsis(content, tool_name="python") assert synopsis.kind == "code" assert "line count" in synopsis.structure[0] assert any("imports: os, pathlib" in item for item in synopsis.structure) assert "class Runner" in synopsis.notable_items assert "def main" in synopsis.notable_items def test_yaml_synopsis_extracts_top_level_keys(self): content = "name: deer\nsettings:\n enabled: true\n retries: 3\nitems:\n - alpha\n" synopsis = build_tool_output_synopsis(content, tool_name="config") assert synopsis.kind == "yaml" assert "Top-level keys: name, settings, items" in synopsis.summary assert "settings: object" in synopsis.structure assert "items: array" in synopsis.structure def test_xml_synopsis_extracts_root_and_children(self): content = '' synopsis = build_tool_output_synopsis(content, tool_name="xml") assert synopsis.kind == "xml" assert "XML document with root tag feed." in synopsis.summary assert "root tag: feed" in synopsis.structure assert "entry: 2" in synopsis.structure # ------------------------------------------------------------------ # Regression tests for the @willem-bd review of PR #3377. # Each test pins one of the eight findings so a future change cannot # silently regress the fix. # ------------------------------------------------------------------ def test_review_5_log_lines_are_not_misclassified_as_yaml(self): # 200 lines of log output shaped like "LEVEL: message". The previous # _looks_yaml counted 2 'key:' lines and accepted it; _try_yaml then # produced a "YAML object with 3 top-level keys: INFO, ERROR, WARN" # summary that hid every line, count, and middle-of-log signal. content = "INFO: starting service\nERROR: failed to connect\nWARN: retrying\nINFO: connected\n" * 200 synopsis = build_tool_output_synopsis(content, tool_name="bash") assert synopsis.kind == "text", f"expected text, got {synopsis.kind!r}: {synopsis.summary}" def test_review_6_json_paths_are_emitted_without_byte_offset(self): # The previous _json_path_location anchored at the first textual # occurrence of the key string, which is wrong when the key also # appears as a value earlier in the document. content = '{"label": "items", "items": {"id": 1, "name": "foo"}}' preview = _build_preview( content, tool_name="api_tool", virtual_path="/mnt/test/api.json", head_chars=200, tail_chars=200, ) assert "byte offset" not in preview # The path itself is still useful navigation. assert "$.items" in preview def test_review_7_scalar_examples_respects_depth_cap(self): # build_tool_output_synopsis used to recurse without a depth cap # in _scalar_examples, which could trigger RecursionError on # deeply nested JSON. The cap is now mirrored from # _JSON_STRUCTURE_DEPTH. deep = {"k": 1} for _ in range(500): deep = {"k": deep} # Should not raise. synopsis = build_tool_output_synopsis(json.dumps(deep)) assert synopsis.kind == "json" def test_review_8_csv_first_row_quoted_cells_round_trip(self): # delimiter.join(rows[1]) silently re-split cells containing the # delimiter inside a quoted cell, misleading the model about # column count. header = "name,description,score" rows = [ 'Ada,"a fine, brilliant logician",98', 'Grace,"a creator, of compilers",99', 'Alan,"a pioneer, of computing",95', 'Kurt,"a poet, of logic",91', 'Ada2,"another, fine mind",97', 'Grace2,"yet another, creator",93', ] content = header + "\n" + "\n".join(rows) + "\n" synopsis = build_tool_output_synopsis(content, tool_name="csv_tool") assert synopsis.kind == "csv" first_row = next((line for line in synopsis.structure if line.startswith("first data row:")), "") # All three columns must be present, and the quoted cell must # round-trip without losing the embedded comma. assert "name=Ada" in first_row assert "score=98" in first_row assert "a fine, brilliant logician" in first_row # The re-joined comma-broken row is the failure mode we are guarding. assert "Ada,a fine, brilliant" not in first_row @pytest.mark.parametrize("delimiter,kind", [(",", "csv"), ("\t", "tsv")]) @pytest.mark.parametrize("newline", ["\n", "\r\n"]) @pytest.mark.parametrize("count", [6, 60]) def test_table_synopsis_counts_logical_records(self, delimiter, kind, newline, count): """Count logical records even beyond the physical-line recognition sample.""" buffer = io.StringIO(newline="") writer = csv.writer(buffer, delimiter=delimiter, lineterminator=newline) writer.writerow(["id", "description", "score"]) for index in range(count): writer.writerow([index, f"first{delimiter}part{newline}second part", 90]) content = buffer.getvalue() synopsis = build_tool_output_synopsis(content) assert synopsis.kind == kind assert synopsis.summary == [f"{kind.upper()} table with {count} data rows and 3 columns."] preview = _build_preview(content, tool_name="bash", virtual_path="/mnt/test/table", head_chars=100, tail_chars=100) assert synopsis.summary[0] in preview @pytest.mark.parametrize("delimiter,kind", [(",", "csv"), ("\t", "tsv")]) def test_table_synopsis_ignores_blank_records(self, delimiter, kind): """Ignore blank records while preserving blank lines inside quoted fields.""" buffer = io.StringIO(newline="") writer = csv.writer(buffer, delimiter=delimiter) writer.writerow(["id", "description"]) for index in range(6): writer.writerow([index, "first\n\nlast"]) writer.writerow([]) writer.writerow([" ", ""]) synopsis = build_tool_output_synopsis(buffer.getvalue()) assert synopsis.kind == kind assert synopsis.summary == [f"{kind.upper()} table with 6 data rows and 2 columns."] @pytest.mark.parametrize("delimiter,kind", [(",", "csv"), ("\t", "tsv")]) @pytest.mark.parametrize("suffix", ['"unterminated', '"closed"invalid', "field_limit"]) def test_table_synopsis_does_not_invent_count_after_parse_failure(self, delimiter, kind, suffix): """Do not report an exact total when parsing fails beyond the sample.""" content = f"id{delimiter}description\n" + "".join(f"{index}{delimiter}ok\n" for index in range(60)) original_limit = csv.field_size_limit() if suffix == "field_limit": suffix = "x" * (original_limit + 1) synopsis = build_tool_output_synopsis(content + f"61{delimiter}{suffix}") assert synopsis.kind == kind assert synopsis.summary == [f"{kind.upper()} table with an undetermined number of data rows and 2 columns."] assert csv.field_size_limit() == original_limit def test_table_synopsis_preserves_oversized_input_guard(self): """Skip structured parsing when the input exceeds the byte budget.""" content = "id,description\n" + "1,ok\n" * 6 + "x" * 5_000_000 synopsis = build_tool_output_synopsis(content) assert synopsis.kind == "unknown" assert "Parsing skipped due to size limit" in synopsis.summary[0] def test_review_9_tsv_detector_rejects_tab_indented_bash(self): # Tab-indented output (ls -l, tree, indented logs) used to be # accepted as TSV because _try_table only checked that the # delimiter is present and rows agree on width. row = "drwxr-xr-x 2 user group 64 Jun 24 17:00 dir" bash_out = "ls -l output:\n\ttotal 0\n" + "\n".join(f"\t{row}{i}" for i in range(1, 6)) + "\n" synopsis = build_tool_output_synopsis(bash_out, tool_name="bash") assert synopsis.kind == "text", f"expected text, got {synopsis.kind!r}: {synopsis.summary}" def test_review_10_preview_includes_raw_head_and_tail_sample(self): # Default behavior change in the PR removed the inline raw bytes # for non-binary previews. The fix restores them so the model # can see the actual first/last KB without a follow-up read_file. content = "log line 1\n" * 200 preview = _build_preview( content, tool_name="bash", virtual_path="/mnt/test/run.log", head_chars=400, tail_chars=400, ) assert "Raw sample (head + tail" in preview # head_chars=400 should capture the first 80 'log line 1' lines # verbatim; tail_chars=400 should capture the last 80. assert preview.count("log line 1") >= 70 # line snapping may lose a few def test_review_11_short_text_does_not_duplicate_excerpts(self): # For inputs shorter than 2 * _TEXT_EXCERPT_CHARS, the previous # opener/closer slices overlapped and the model saw the same # body twice. build_tool_output_synopsis is reachable directly # from tests and other callers that pass small inputs. short = "hello world " * 30 # ~360 chars synopsis = build_tool_output_synopsis(short) opener_line = next((ln for ln in synopsis.summary if ln.startswith("Opening excerpt: ")), "") # Closer is now suppressed entirely for short inputs. assert all(not ln.startswith("Closing excerpt: ") for ln in synopsis.summary), f"unexpected closer for short input: {synopsis.summary}" assert opener_line, "opening excerpt should still be present" def test_review_12_preview_head_tail_chars_are_operational(self): # preview_head_chars / preview_tail_chars were silently no-op # for every non-binary kind. The fix plumbs them through # render_tool_output_preview as an explicit 'Raw sample' section. content = "alpha " * 1000 # 6000 chars preview = _build_preview( content, tool_name="bash", virtual_path="/mnt/test/run.log", head_chars=300, tail_chars=300, ) # The head sample should contain 'alpha' more times than the # tail (or split-count), proving head_chars=300 took effect. # The full document has 1000 'alpha' tokens; without head_chars # we'd see fewer than 50 in the head sample. assert preview.count("alpha") >= 50 class TestBuildFallback: def test_short_content_unchanged(self): assert _build_fallback("short", tool_name="t", max_chars=100, head_chars=50, tail_chars=50) == "short" def test_zero_max_disables(self): content = "a" * 1000 assert _build_fallback(content, tool_name="t", max_chars=0, head_chars=50, tail_chars=50) == content def test_truncates_long_content(self): content = "H" * 5000 + "M" * 20000 + "T" * 5000 result = _build_fallback(content, tool_name="bash", max_chars=12000, head_chars=6000, tail_chars=3000) assert len(result) < len(content) assert "omitted from bash output" in result assert "Persistent storage unavailable" in result def test_preserves_head_and_tail(self): content = "HEADSTART" + "x" * 50000 + "TAILEND" result = _build_fallback(content, tool_name="t", max_chars=20000, head_chars=10000, tail_chars=5000) assert result.startswith("HEADSTART") assert "TAILEND" in result def test_result_never_exceeds_max_chars(self): """The marker itself has non-zero length; total must still respect max_chars.""" for max_chars in [200, 500, 1000, 5000, 20000]: content = "x" * 50000 result = _build_fallback(content, tool_name="long_tool_name", max_chars=max_chars, head_chars=max_chars // 2, tail_chars=max_chars // 4) assert len(result) <= max_chars, f"max_chars={max_chars}: got {len(result)}" def test_result_never_exceeds_max_chars_with_newlines(self): """Same guarantee as above, on content that actually exercises line snapping. ``test_result_never_exceeds_max_chars`` passes newline-free content, so the tail offset is never snapped. Real bash/web_fetch output has newlines. """ for total in [50_000, 200_000, 1_000_000]: content = _lines_then_long_line(total) result = _build_fallback(content, tool_name="bash", max_chars=30_000, head_chars=8_000, tail_chars=3_000) assert len(result) <= 30_000, f"total={total}: got {len(result)}" def test_fallback_forward_snaps_tail_onto_line_boundary(self): """The tail must begin *after* the newline, never before it. The bound test above never moves the tail offset: its content has no newline inside the snap window, so it would pass even with the snap removed. Placing a newline in the window pins the direction instead — a backward snap leaves the tail starting mid-line. """ total, newline_pos = 100_000, 98_000 # window is [97_000, 98_500) content = "A" * newline_pos + "\n" + "B" * (total - newline_pos - 1) result = _build_fallback(content, tool_name="bash", max_chars=30_000, head_chars=8_000, tail_chars=3_000) assert len(result) <= 30_000 tail = result.rsplit("]\n\n", 1)[1] assert tail.startswith("B"), f"tail begins mid-line: {tail[:20]!r}" def test_very_small_max_chars_does_not_crash(self): content = "x" * 1000 result = _build_fallback(content, tool_name="t", max_chars=50, head_chars=20, tail_chars=10) assert len(result) <= 50 # =========================================================================== # Bash exit-marker preservation (harvest reads Exit Code: N from the end) # =========================================================================== class TestBashExitMarkerPreservation: """Budget rewrites must keep a trailing bash exit marker last. The subagent executor harvests status with ``Exit Code: N\\s*$``. Preview construction always ends with an Access footer, so without this the marker is no longer last and a failed command is harvested as success. """ def test_keep_exit_marker_last_appends_for_bash(self): original = "out\nExit Code: 1" rewritten = "preview\nAccess:\n- Use read_file" assert _keep_exit_marker_last(original, rewritten, tool_name="bash") == "preview\nAccess:\n- Use read_file\nExit Code: 1" def test_keep_exit_marker_last_is_noop_for_other_tools(self): original = "out\nExit Code: 1" rewritten = "preview\nAccess:\n- Use read_file" assert _keep_exit_marker_last(original, rewritten, tool_name="web_fetch") == rewritten def test_keep_exit_marker_last_is_noop_when_already_last(self): original = "out\nExit Code: 0" rewritten = "preview\nExit Code: 0" assert _keep_exit_marker_last(original, rewritten, tool_name="bash") == rewritten def test_wrap_tool_call_keeps_failed_exit_marker_last(self): content = _pytest_like_output("Exit Code: 1") assert 12_000 < len(content) <= 20_000 with tempfile.TemporaryDirectory() as tmpdir: mw = ToolOutputBudgetMiddleware(config=ToolOutputConfig()) msg = _bash_tm(content) result = mw.wrap_tool_call(_make_request(tool_name="bash", outputs_path=tmpdir), lambda _: msg) assert isinstance(result, ToolMessage) assert result is not msg assert str(result.content).rstrip().endswith("Exit Code: 1") assert "Access:" in result.content assert "read_file" in result.content assert result.additional_kwargs["deerflow_tool_meta"] == {"status": "success", "source": "normalized"} assert result.artifact == {"k": "v"} assert result.id == "tool-msg-1" @pytest.mark.anyio async def test_awrap_tool_call_keeps_failed_exit_marker_last(self): content = _pytest_like_output("Exit Code: 1") with tempfile.TemporaryDirectory() as tmpdir: mw = ToolOutputBudgetMiddleware(config=ToolOutputConfig()) msg = _bash_tm(content) async def handler(_): return msg result = await mw.awrap_tool_call(_make_request(tool_name="bash", outputs_path=tmpdir), handler) assert str(result.content).rstrip().endswith("Exit Code: 1") assert "Access:" in result.content assert result.artifact == {"k": "v"} assert result.id == "tool-msg-1" def test_wrap_tool_call_keeps_successful_exit_marker_last(self): content = _pytest_like_output("Exit Code: 0") with tempfile.TemporaryDirectory() as tmpdir: mw = ToolOutputBudgetMiddleware(config=ToolOutputConfig()) result = mw.wrap_tool_call(_make_request(tool_name="bash", outputs_path=tmpdir), lambda _: _bash_tm(content)) assert str(result.content).rstrip().endswith("Exit Code: 0") assert "Access:" in result.content def test_wrap_tool_call_without_marker_ends_at_access_footer(self): content = _pytest_like_output("12 failed in 1.23s") assert not content.rstrip().endswith("Exit Code: 1") with tempfile.TemporaryDirectory() as tmpdir: mw = ToolOutputBudgetMiddleware(config=ToolOutputConfig()) result = mw.wrap_tool_call(_make_request(tool_name="bash", outputs_path=tmpdir), lambda _: _bash_tm(content)) assert "Access:" in result.content assert "read_file" in result.content assert not str(result.content).rstrip().endswith("Exit Code: 1") assert "Exit Code:" not in str(result.content).rsplit("Access:", 1)[-1] def test_non_bash_tool_does_not_promote_exit_marker(self): content = _pytest_like_output("Exit Code: 1") with tempfile.TemporaryDirectory() as tmpdir: mw = ToolOutputBudgetMiddleware(config=ToolOutputConfig()) result = mw.wrap_tool_call( _make_request(tool_name="web_fetch", outputs_path=tmpdir), lambda _: _bash_tm(content, name="web_fetch"), ) assert "Access:" in result.content assert not str(result.content).rstrip().endswith("Exit Code: 1") def test_bash_tool_alias_also_preserves_marker(self): content = _pytest_like_output("Exit Code: 1") with tempfile.TemporaryDirectory() as tmpdir: mw = ToolOutputBudgetMiddleware(config=ToolOutputConfig()) result = mw.wrap_tool_call( _make_request(tool_name="bash_tool", outputs_path=tmpdir), lambda _: _bash_tm(content, name="bash_tool"), ) assert str(result.content).rstrip().endswith("Exit Code: 1") def test_fallback_tiny_budget_keeps_marker_and_stays_within_max_chars(self): content = _pytest_like_output("Exit Code: 1") result = _build_fallback(content, tool_name="bash", max_chars=80, head_chars=8000, tail_chars=3000) assert len(result) <= 80 assert result.rstrip().endswith("Exit Code: 1") def test_fallback_default_budget_leaves_payload_unchanged(self): content = _pytest_like_output("Exit Code: 1") result = _build_fallback(content, tool_name="bash", max_chars=30_000, head_chars=8000, tail_chars=3000) assert result == content assert result.rstrip().endswith("Exit Code: 1") def test_fallback_never_exceeds_max_chars_with_exit_marker(self): content = _pytest_like_output("Exit Code: 1") marker_line = "\nExit Code: 1" for max_chars in [10, len(marker_line) - 1, len(marker_line), 20, 80, 200, 500, 1000, 5000, 20000]: result = _build_fallback(content, tool_name="bash", max_chars=max_chars, head_chars=max_chars // 2, tail_chars=max_chars // 4) assert len(result) <= max_chars, f"max_chars={max_chars}: got {len(result)}" if max_chars >= len(marker_line): assert result.rstrip().endswith("Exit Code: 1"), f"max_chars={max_chars}" else: assert not result.rstrip().endswith("Exit Code: 1") def test_exit_marker_regex_matches_sandbox_truncation_tail_shapes(self): from deerflow.sandbox.tools import _BASH_EXIT_MARKER_TAIL_RE as sandbox_re for content in ("out\nExit Code: 1", "out\nExit Code: -9 \n", "out\nCommand exited with code 3", "Command exited with code 3"): ours = _BASH_EXIT_MARKER_TAIL_RE.search(content) theirs = sandbox_re.search(content) assert ours is not None and theirs is not None assert ours.start() == theirs.start(), content def test_budgeted_failed_pytest_is_harvested_as_error_not_pass(self): """End-to-end: wrap_tool_call rewrite → harvest status=error → leaf holds=False. ``tests/conftest.py`` mocks ``deerflow.subagents.executor``, so the production harvest helpers are loaded under a unique module name. """ path = pathlib.Path(__file__).parents[1] / "packages/harness/deerflow/subagents/executor.py" spec = importlib.util.spec_from_file_location("_budget_exit_marker_executor", path) assert spec is not None and spec.loader is not None module = importlib.util.module_from_spec(spec) try: sys.modules[spec.name] = module spec.loader.exec_module(module) from deerflow.subagents import acceptance_checks content = _pytest_like_output("Exit Code: 1") with tempfile.TemporaryDirectory() as tmpdir: mw = ToolOutputBudgetMiddleware(config=ToolOutputConfig()) result = mw.wrap_tool_call(_make_request(tool_name="bash", outputs_path=tmpdir), lambda _: _bash_tm(content)) assert str(result.content).rstrip().endswith("Exit Code: 1") cmd = "pytest -q" ai = AIMessage( content="", tool_calls=[{"name": "bash", "args": {"command": cmd}, "id": "tc-1", "type": "tool_call"}], ) state = {"messages": [HumanMessage(content="task"), ai, result]} executions = module._harvest_bash_executions(state) assert executions, "harvest returned no bash executions" for entry in executions: entry["shell_persistent"] = False latest = executions[-1] assert latest["status"] == "error" assert latest["status_marker"] == "Exit Code: 1" leaf = acceptance_checks._check_tests_passed_leaf(cmd, executions) assert leaf["checked"] is True assert leaf["holds"] is False finally: sys.modules.pop(spec.name, None) # Absent when exec_module failed part-way; don't mask that error. shutdown = getattr(module, "_shutdown_isolated_subagent_loop", None) if shutdown is not None: shutdown() # =========================================================================== # Middleware integration tests — wrap_tool_call # =========================================================================== class TestWrapToolCallPassThrough: def test_small_output_passes_through(self): mw = ToolOutputBudgetMiddleware(config=ToolOutputConfig(externalize_min_chars=1000)) msg = _tm("small output", name="bash") result = mw.wrap_tool_call(_make_request(), lambda _: msg) assert result is msg def test_disabled_middleware_passes_through(self): mw = ToolOutputBudgetMiddleware(config=ToolOutputConfig(enabled=False, externalize_min_chars=10, fallback_max_chars=20)) msg = _tm("x" * 50000, name="bash") result = mw.wrap_tool_call(_make_request(), lambda _: msg) assert result is msg class TestWrapToolCallExternalize: def test_oversized_output_externalized_to_disk(self): with tempfile.TemporaryDirectory() as tmpdir: config = ToolOutputConfig(externalize_min_chars=100, preview_head_chars=50, preview_tail_chars=30) mw = ToolOutputBudgetMiddleware(config=config) content = "x" * 500 msg = _tm(content, name="remote_executor") req = _make_request(outputs_path=tmpdir) result = mw.wrap_tool_call(req, lambda _: msg) assert isinstance(result, ToolMessage) assert result is not msg assert "Full remote_executor output saved to" in result.content assert "read_file" in result.content assert result.tool_call_id == "tc-1" assert TOOL_OUTPUT_BLOB_KEY not in result.additional_kwargs # Verify file was written storage_dir = os.path.join(tmpdir, ".tool-results") assert os.path.isdir(storage_dir) files = os.listdir(storage_dir) assert len(files) == 1 with open(os.path.join(storage_dir, files[0]), encoding="utf-8") as f: assert f.read() == content class TestToolOutputBlobPersistence: def test_host_externalization_stamps_durable_blob_ref(self, monkeypatch, tmp_path): from deerflow.agents.middlewares import tool_output_budget_middleware as mod content = "跨节点工具输出" * 80 calls: list[dict] = [] class RecordingStore: def put_bytes(self, data: bytes, **kwargs) -> BlobRef: calls.append({"data": data, **kwargs}) return BlobRef( sha256=hashlib.sha256(data).hexdigest(), size=len(data), kind=kwargs["kind"], content_type=kwargs["content_type"], ) monkeypatch.setattr(mod, "get_blob_store_if_enabled", lambda: RecordingStore()) config = ToolOutputConfig(externalize_min_chars=50, preview_head_chars=20, preview_tail_chars=10) result = ToolOutputBudgetMiddleware(config=config).wrap_tool_call( _make_request(outputs_path=str(tmp_path)), lambda _: _tm(content, name="remote_executor"), ) payload = result.additional_kwargs[TOOL_OUTPUT_BLOB_KEY] assert payload["version"] == 1 assert payload["virtual_path"].startswith("/mnt/user-data/outputs/.tool-results/") assert payload["storage_subdir"] == ".tool-results" assert payload["encoding"] == "utf-8" assert payload["ref"] == { "sha256": hashlib.sha256(content.encode()).hexdigest(), "size": len(content.encode()), "kind": "tool-output", "content_type": "text/plain; charset=utf-8", } assert calls == [ { "data": content.encode(), "kind": "tool-output", "content_type": "text/plain; charset=utf-8", "thread_id": "thread-1", } ] def test_host_externalized_file_keeps_its_blob_ref_bytes_under_windows_newlines(self, monkeypatch, tmp_path): # Windows text mode wrote "\n" as "\r\n", so the file no longer matched # the ref stamped from content.encode() and the next model call on this # Gateway deleted it as a mismatch when no blob store was configured. from deerflow.agents.middlewares import tool_output_budget_middleware as mod content = "first line\nsecond line\n" * 20 builtin_open = open def windows_open(file, mode="r", *args, **kwargs): if "b" not in mode and any(flag in mode for flag in "wax+"): kwargs.setdefault("newline", "\r\n") return builtin_open(file, mode, *args, **kwargs) class RecordingStore: def put_bytes(self, data: bytes, **kwargs) -> BlobRef: return BlobRef(sha256=hashlib.sha256(data).hexdigest(), size=len(data), kind=kwargs["kind"], content_type=kwargs["content_type"]) monkeypatch.setattr(mod, "open", windows_open, raising=False) monkeypatch.setattr(mod, "get_blob_store_if_enabled", lambda: RecordingStore()) mw = ToolOutputBudgetMiddleware(config=ToolOutputConfig(externalize_min_chars=50, preview_head_chars=20, preview_tail_chars=10)) result = mw.wrap_tool_call(_make_request(outputs_path=str(tmp_path)), lambda _: _tm(content, name="remote_executor")) saved = tmp_path / ".tool-results" / os.path.basename(result.additional_kwargs[TOOL_OUTPUT_BLOB_KEY]["virtual_path"]) assert saved.read_bytes() == content.encode() monkeypatch.setattr(mod, "get_blob_store_if_enabled", lambda: None) request = ModelRequest(model=None, messages=[result], tools=[], state={"thread_data": {"outputs_path": str(tmp_path)}}) mw.wrap_model_call(request, lambda prepared: []) assert saved.read_bytes() == content.encode() def test_configured_blob_write_failure_falls_back_inline(self, monkeypatch, tmp_path): from deerflow.agents.middlewares import tool_output_budget_middleware as mod class FailingStore: def put_bytes(self, data: bytes, **kwargs) -> BlobRef: raise BlobWriteError("backend detail") monkeypatch.setattr(mod, "get_blob_store_if_enabled", lambda: FailingStore()) config = ToolOutputConfig( externalize_min_chars=50, fallback_max_chars=200, fallback_head_chars=80, fallback_tail_chars=40, ) result = ToolOutputBudgetMiddleware(config=config).wrap_tool_call( _make_request(outputs_path=str(tmp_path)), lambda _: _tm("x" * 500, name="remote_executor"), ) assert "Persistent storage unavailable" in result.content assert TOOL_OUTPUT_BLOB_KEY not in result.additional_kwargs published = list((tmp_path / ".tool-results").glob("remote_executor-*")) assert len(published) == 1 assert published[0].read_text(encoding="utf-8") == "x" * 500 def test_blob_write_failure_stays_bounded_when_regular_fallback_is_disabled(self, monkeypatch, tmp_path): from deerflow.agents.middlewares import tool_output_budget_middleware as mod class FailingStore: def put_bytes(self, data: bytes, **kwargs) -> BlobRef: raise BlobWriteError("backend detail") monkeypatch.setattr(mod, "get_blob_store_if_enabled", lambda: FailingStore()) monkeypatch.setattr(mod, "_DURABLE_FAILURE_FALLBACK_MAX_CHARS", 80) config = ToolOutputConfig( externalize_min_chars=10, fallback_max_chars=0, fallback_head_chars=20, fallback_tail_chars=10, ) result = ToolOutputBudgetMiddleware(config=config).wrap_tool_call( _make_request(outputs_path=str(tmp_path)), lambda _: _tm("x" * 500, name="remote_executor"), ) assert len(result.content) == 80 assert result.additional_kwargs["deerflow_tool_transforms"][-1]["kind"] == "truncated" assert TOOL_OUTPUT_BLOB_KEY not in result.additional_kwargs def test_blob_producer_cap_falls_back_without_writing(self, monkeypatch, tmp_path): from deerflow.agents.middlewares import tool_output_budget_middleware as mod class StoreMustNotBeWritten: def put_bytes(self, data: bytes, **kwargs) -> BlobRef: raise AssertionError("oversized content must not reach the blob backend") monkeypatch.setattr(mod, "get_blob_store_if_enabled", lambda: StoreMustNotBeWritten()) monkeypatch.setattr(mod, "_MAX_TOOL_OUTPUT_BLOB_BYTES", 32) config = ToolOutputConfig( externalize_min_chars=10, fallback_max_chars=80, fallback_head_chars=20, fallback_tail_chars=10, ) result = ToolOutputBudgetMiddleware(config=config).wrap_tool_call( _make_request(outputs_path=str(tmp_path)), lambda _: _tm("x" * 100, name="remote_executor"), ) assert len(result.content) == 80 assert result.additional_kwargs["deerflow_tool_transforms"][-1]["kind"] == "truncated" assert TOOL_OUTPUT_BLOB_KEY not in result.additional_kwargs assert not (tmp_path / ".tool-results").exists() def test_blob_producer_cap_stays_bounded_when_regular_fallback_is_disabled(self, monkeypatch, tmp_path): from deerflow.agents.middlewares import tool_output_budget_middleware as mod class StoreMustNotBeWritten: def put_bytes(self, data: bytes, **kwargs) -> BlobRef: raise AssertionError("oversized content must not reach the blob backend") monkeypatch.setattr(mod, "get_blob_store_if_enabled", lambda: StoreMustNotBeWritten()) monkeypatch.setattr(mod, "_MAX_TOOL_OUTPUT_BLOB_BYTES", 32) monkeypatch.setattr(mod, "_DURABLE_FAILURE_FALLBACK_MAX_CHARS", 80) config = ToolOutputConfig( externalize_min_chars=10, fallback_max_chars=0, fallback_head_chars=20, fallback_tail_chars=10, ) result = ToolOutputBudgetMiddleware(config=config).wrap_tool_call( _make_request(outputs_path=str(tmp_path)), lambda _: _tm("x" * 100, name="remote_executor"), ) assert len(result.content) == 80 assert result.additional_kwargs["deerflow_tool_transforms"][-1]["kind"] == "truncated" assert TOOL_OUTPUT_BLOB_KEY not in result.additional_kwargs assert not (tmp_path / ".tool-results").exists() def test_preview_contains_typed_summary(self): with tempfile.TemporaryDirectory() as tmpdir: config = ToolOutputConfig(externalize_min_chars=50, preview_head_chars=20, preview_tail_chars=10) mw = ToolOutputBudgetMiddleware(config=config) content = "HEADPART_" + "m" * 200 + "_TAILPART" msg = _tm(content, name="web_search") req = _make_request(outputs_path=tmpdir) result = mw.wrap_tool_call(req, lambda _: msg) assert result.content.startswith("[Full web_search output saved to") assert "Preview kind: text" in result.content assert "Text output" in result.content assert "HEADPART_" in result.content class TestWrapToolCallFallback: def test_fallback_when_no_outputs_path(self): config = ToolOutputConfig( externalize_min_chars=50, fallback_max_chars=200, fallback_head_chars=80, fallback_tail_chars=40, ) mw = ToolOutputBudgetMiddleware(config=config) content = "x" * 500 msg = _tm(content, name="mcp_tool") req = _make_request(outputs_path=None) result = mw.wrap_tool_call(req, lambda _: msg) assert isinstance(result, ToolMessage) assert result is not msg assert "omitted from mcp_tool output" in result.content assert "Persistent storage unavailable" in result.content assert len(result.content) < len(content) def test_fallback_when_disk_write_fails(self): config = ToolOutputConfig( externalize_min_chars=50, fallback_max_chars=200, fallback_head_chars=80, fallback_tail_chars=40, ) mw = ToolOutputBudgetMiddleware(config=config) content = "x" * 500 msg = _tm(content, name="tool") with _unwritable_outputs_path() as outputs_path: req = _make_request(outputs_path=outputs_path) result = mw.wrap_tool_call(req, lambda _: msg) assert isinstance(result, ToolMessage) assert "omitted from tool output" in result.content def test_structured_json_output_budgeted_before_provider(self): # Review feedback on #6208: the MindIE adapter serializes {"type": "json"} # blocks to text, so a large structured result must hit the budget here, # not sail through to provider normalization at full size. config = ToolOutputConfig( externalize_min_chars=50, fallback_max_chars=200, fallback_head_chars=80, fallback_tail_chars=40, ) mw = ToolOutputBudgetMiddleware(config=config) msg = ToolMessage( content=[{"type": "json", "json": {"rows": ["x" * 500]}}], name="query_rows", tool_call_id="tc-1", ) req = _make_request(outputs_path=None) result = mw.wrap_tool_call(req, lambda _: msg) assert isinstance(result, ToolMessage) assert result is not msg assert isinstance(result.content, str) assert "omitted from query_rows output" in result.content assert len(result.content) <= 200 class TestWrapToolCallExemption: def test_read_file_exempt(self): config = ToolOutputConfig(externalize_min_chars=10, fallback_max_chars=50) mw = ToolOutputBudgetMiddleware(config=config) content = "x" * 100 msg = _tm(content, name="read_file") result = mw.wrap_tool_call(_make_request(tool_name="read_file"), lambda _: msg) assert result is msg def test_read_file_tool_exempt(self): config = ToolOutputConfig(externalize_min_chars=10, fallback_max_chars=50) mw = ToolOutputBudgetMiddleware(config=config) content = "x" * 100 msg = _tm(content, name="read_file_tool") result = mw.wrap_tool_call(_make_request(tool_name="read_file_tool"), lambda _: msg) assert result is msg def test_custom_exempt_tool(self): config = ToolOutputConfig(externalize_min_chars=10, fallback_max_chars=50, exempt_tools=["my_tool"]) mw = ToolOutputBudgetMiddleware(config=config) content = "x" * 100 msg = _tm(content, name="my_tool") result = mw.wrap_tool_call(_make_request(tool_name="my_tool"), lambda _: msg) assert result is msg class TestWrapToolCallPerToolOverride: def test_per_tool_threshold(self): with tempfile.TemporaryDirectory() as tmpdir: config = ToolOutputConfig( externalize_min_chars=50000, # global: high tool_overrides={"sensitive_tool": 100}, # override: low ) mw = ToolOutputBudgetMiddleware(config=config) content = "x" * 500 msg = _tm(content, name="sensitive_tool") req = _make_request(tool_name="sensitive_tool", outputs_path=tmpdir) result = mw.wrap_tool_call(req, lambda _: msg) assert result is not msg assert "Full sensitive_tool output saved to" in result.content def test_per_tool_zero_disables_externalization(self): config = ToolOutputConfig( externalize_min_chars=50, tool_overrides={"bash": 0}, fallback_max_chars=200, fallback_head_chars=80, fallback_tail_chars=40, ) mw = ToolOutputBudgetMiddleware(config=config) content = "x" * 500 msg = _tm(content, name="bash") # Even with outputs_path, externalization disabled for bash req = _make_request(tool_name="bash", outputs_path="/tmp/test") result = mw.wrap_tool_call(req, lambda _: msg) assert isinstance(result, ToolMessage) # Should use fallback instead of externalization assert "Persistent storage unavailable" in result.content or "omitted" in result.content class TestWrapToolCallCommand: def test_command_messages_are_patched(self): with tempfile.TemporaryDirectory() as tmpdir: config = ToolOutputConfig(externalize_min_chars=50, preview_head_chars=20, preview_tail_chars=10) mw = ToolOutputBudgetMiddleware(config=config) tool_msg = _tm("x" * 200, name="present_files") command = Command(update={"messages": [tool_msg], "artifacts": ["/mnt/report.html"]}) req = _make_request(tool_name="present_files", outputs_path=tmpdir) result = mw.wrap_tool_call(req, lambda _: command) assert isinstance(result, Command) assert result is not command assert result.update["artifacts"] == ["/mnt/report.html"] new_msg = result.update["messages"][0] assert isinstance(new_msg, ToolMessage) assert "Full present_files output saved to" in new_msg.content def test_command_without_messages_unchanged(self): config = ToolOutputConfig(externalize_min_chars=10) mw = ToolOutputBudgetMiddleware(config=config) command = Command(update={"key": "value"}) result = mw.wrap_tool_call(_make_request(), lambda _: command) assert result is command class TestWrapToolCallEdgeCases: def test_none_content_passes_through(self): config = ToolOutputConfig(externalize_min_chars=10, fallback_max_chars=20) mw = ToolOutputBudgetMiddleware(config=config) msg = ToolMessage(content=None, name="tool", tool_call_id="tc-1") result = mw.wrap_tool_call(_make_request(), lambda _: msg) assert result is msg def test_empty_string_passes_through(self): config = ToolOutputConfig(externalize_min_chars=10, fallback_max_chars=20) mw = ToolOutputBudgetMiddleware(config=config) msg = _tm("", name="tool") result = mw.wrap_tool_call(_make_request(), lambda _: msg) assert result is msg def test_multimodal_content_skipped(self): config = ToolOutputConfig(externalize_min_chars=10, fallback_max_chars=20) mw = ToolOutputBudgetMiddleware(config=config) content = [{"type": "image", "data": "x" * 100}] msg = ToolMessage(content=content, name="tool", tool_call_id="tc-1") result = mw.wrap_tool_call(_make_request(), lambda _: msg) assert result is msg def test_exactly_at_threshold_passes_through(self): config = ToolOutputConfig(externalize_min_chars=100, fallback_max_chars=100) mw = ToolOutputBudgetMiddleware(config=config) msg = _tm("x" * 100, name="tool") result = mw.wrap_tool_call(_make_request(), lambda _: msg) assert result is msg def test_one_char_over_threshold_triggers(self): with tempfile.TemporaryDirectory() as tmpdir: config = ToolOutputConfig(externalize_min_chars=100) mw = ToolOutputBudgetMiddleware(config=config) msg = _tm("x" * 101, name="tool") req = _make_request(outputs_path=tmpdir) result = mw.wrap_tool_call(req, lambda _: msg) assert result is not msg def test_chinese_content_preserved(self): with tempfile.TemporaryDirectory() as tmpdir: config = ToolOutputConfig(externalize_min_chars=50, preview_head_chars=20, preview_tail_chars=10) mw = ToolOutputBudgetMiddleware(config=config) content = "你好世界" * 50 msg = _tm(content, name="tool") req = _make_request(outputs_path=tmpdir) result = mw.wrap_tool_call(req, lambda _: msg) assert isinstance(result, ToolMessage) # File should contain the full Chinese content storage_dir = os.path.join(tmpdir, ".tool-results") files = os.listdir(storage_dir) with open(os.path.join(storage_dir, files[0]), encoding="utf-8") as f: assert f.read() == content def test_no_runtime_state_uses_fallback(self): config = ToolOutputConfig( externalize_min_chars=50, fallback_max_chars=500, fallback_head_chars=100, fallback_tail_chars=50, ) mw = ToolOutputBudgetMiddleware(config=config) content = "x" * 1000 msg = _tm(content, name="tool") req = SimpleNamespace( tool_call={"name": "tool", "id": "tc-1"}, runtime=None, ) result = mw.wrap_tool_call(req, lambda _: msg) assert isinstance(result, ToolMessage) assert "omitted" in result.content assert len(result.content) <= 500 # =========================================================================== # MCP content_and_artifact format tests # =========================================================================== class TestMCPContentAndArtifact: """MCP tools return content as list of content blocks, not plain strings.""" def test_text_content_blocks_are_budgeted(self): with tempfile.TemporaryDirectory() as tmpdir: config = ToolOutputConfig(externalize_min_chars=50, preview_head_chars=20, preview_tail_chars=10) mw = ToolOutputBudgetMiddleware(config=config) content = [{"type": "text", "text": "x" * 200}] msg = ToolMessage(content=content, name="mcp_tool", tool_call_id="tc-mcp") req = _make_request(tool_name="mcp_tool", outputs_path=tmpdir) result = mw.wrap_tool_call(req, lambda _: msg) assert result is not msg assert isinstance(result.content, str) assert "Full mcp_tool output saved to" in result.content assert result.tool_call_id == "tc-mcp" def test_multiple_text_blocks_joined_and_budgeted(self): config = ToolOutputConfig(externalize_min_chars=50, fallback_max_chars=500, fallback_head_chars=100, fallback_tail_chars=50) mw = ToolOutputBudgetMiddleware(config=config) content = [{"type": "text", "text": "a" * 300}, {"type": "text", "text": "b" * 300}] msg = ToolMessage(content=content, name="mcp_tool", tool_call_id="tc-mcp2") req = _make_request(tool_name="mcp_tool") result = mw.wrap_tool_call(req, lambda _: msg) assert result is not msg assert "omitted" in result.content def test_image_content_blocks_are_skipped(self): config = ToolOutputConfig(externalize_min_chars=10, fallback_max_chars=20) mw = ToolOutputBudgetMiddleware(config=config) content = [{"type": "image", "data": "base64data" * 100}] msg = ToolMessage(content=content, name="mcp_tool", tool_call_id="tc-img") req = _make_request(tool_name="mcp_tool") result = mw.wrap_tool_call(req, lambda _: msg) assert result is msg def test_mixed_text_and_image_blocks_are_skipped(self): config = ToolOutputConfig(externalize_min_chars=10) mw = ToolOutputBudgetMiddleware(config=config) content = [{"type": "text", "text": "x" * 100}, {"type": "image", "data": "base64"}] msg = ToolMessage(content=content, name="mcp_tool", tool_call_id="tc-mix") req = _make_request(tool_name="mcp_tool") result = mw.wrap_tool_call(req, lambda _: msg) assert result is msg def test_small_text_blocks_pass_through(self): config = ToolOutputConfig(externalize_min_chars=1000) mw = ToolOutputBudgetMiddleware(config=config) content = [{"type": "text", "text": "small result"}] msg = ToolMessage(content=content, name="mcp_tool", tool_call_id="tc-sm") req = _make_request(tool_name="mcp_tool") result = mw.wrap_tool_call(req, lambda _: msg) assert result is msg # =========================================================================== # Async path tests # =========================================================================== class TestAsyncPaths: @pytest.mark.anyio async def test_async_tool_call_externalized(self): with tempfile.TemporaryDirectory() as tmpdir: config = ToolOutputConfig(externalize_min_chars=50, preview_head_chars=20, preview_tail_chars=10) mw = ToolOutputBudgetMiddleware(config=config) content = "x" * 200 msg = _tm(content, name="async_tool") req = _make_request(tool_name="async_tool", outputs_path=tmpdir) async def handler(_): return msg result = await mw.awrap_tool_call(req, handler) assert isinstance(result, ToolMessage) assert result is not msg assert "Full async_tool output saved to" in result.content @pytest.mark.anyio async def test_async_model_call_patches_history(self): config = ToolOutputConfig(fallback_max_chars=500, fallback_head_chars=100, fallback_tail_chars=50) mw = ToolOutputBudgetMiddleware(config=config) oversized = _tm("h" * 1000, name="tool", tool_call_id="tc-h") request = ModelRequest(model=None, messages=[oversized], tools=[], state={}) captured: dict[str, ModelRequest] = {} async def handler(req): captured["request"] = req return [] await mw.awrap_model_call(request, handler) forwarded = captured["request"] assert forwarded is not request msg = forwarded.messages[0] assert isinstance(msg, ToolMessage) assert "omitted" in msg.content # =========================================================================== # wrap_model_call — historical message patching # =========================================================================== class TestWrapModelCall: def test_oversized_historical_messages_truncated(self): config = ToolOutputConfig(fallback_max_chars=500, fallback_head_chars=100, fallback_tail_chars=50) mw = ToolOutputBudgetMiddleware(config=config) oversized = _tm("q" * 1000, name="tool", tool_call_id="tc-q") request = ModelRequest(model=None, messages=[oversized], tools=[], state={}) captured: dict[str, ModelRequest] = {} def handler(req): captured["request"] = req return [] mw.wrap_model_call(request, handler) forwarded = captured["request"] assert forwarded is not request msg = forwarded.messages[0] assert isinstance(msg, ToolMessage) assert "omitted" in msg.content assert len(msg.content) < len(oversized.content) + 150 def test_small_historical_messages_unchanged(self): config = ToolOutputConfig(fallback_max_chars=1000) mw = ToolOutputBudgetMiddleware(config=config) small = _tm("small", name="tool") request = ModelRequest(model=None, messages=[small], tools=[], state={}) captured: dict[str, ModelRequest] = {} def handler(req): captured["request"] = req return [] mw.wrap_model_call(request, handler) assert captured["request"] is request def test_exempt_tools_in_history_unchanged(self): config = ToolOutputConfig(fallback_max_chars=50) mw = ToolOutputBudgetMiddleware(config=config) read_msg = _tm("x" * 200, name="read_file", tool_call_id="tc-r") request = ModelRequest(model=None, messages=[read_msg], tools=[], state={}) captured: dict[str, ModelRequest] = {} def handler(req): captured["request"] = req return [] mw.wrap_model_call(request, handler) assert captured["request"] is request def test_non_tool_messages_preserved(self): config = ToolOutputConfig(fallback_max_chars=500, fallback_head_chars=100, fallback_tail_chars=50) mw = ToolOutputBudgetMiddleware(config=config) human = HumanMessage(content="x" * 200) ai = AIMessage(content="y" * 200) oversized_tool = _tm("z" * 1000, name="tool") request = ModelRequest(model=None, messages=[human, ai, oversized_tool], tools=[], state={}) captured: dict[str, ModelRequest] = {} def handler(req): captured["request"] = req return [] mw.wrap_model_call(request, handler) msgs = captured["request"].messages assert msgs[0] is human assert msgs[1] is ai assert isinstance(msgs[2], ToolMessage) assert "omitted" in msgs[2].content class TestToolOutputBlobRestore: @staticmethod def _request(tmp_path, data: bytes, *, virtual_path: str | None = None) -> tuple[ModelRequest, BlobRef]: ref = BlobRef( sha256=hashlib.sha256(data).hexdigest(), size=len(data), kind="tool-output", content_type="text/plain; charset=utf-8", ) path = virtual_path or "/mnt/user-data/outputs/.tool-results/remote_executor-tc-1.txt" message = ToolMessage( content=f"[Full remote_executor output saved to {path}]", name="remote_executor", tool_call_id="tc-1", additional_kwargs={ TOOL_OUTPUT_BLOB_KEY: { "version": 1, "ref": ref.model_dump(mode="json", exclude_none=True), "virtual_path": path, "storage_subdir": ".tool-results", "encoding": "utf-8", } }, ) request = ModelRequest( model=None, messages=[message], tools=[], state={"thread_data": {"outputs_path": str(tmp_path)}}, ) return request, ref def test_model_call_restores_exact_bytes_on_another_gateway(self, monkeypatch, tmp_path): from deerflow.agents.middlewares import tool_output_budget_middleware as mod data = "跨节点完整输出\nsecond line".encode() request, ref = self._request(tmp_path, data) reads: list[BlobRef] = [] class SharedStore: def get_bytes(self, requested: BlobRef) -> bytes: reads.append(requested) return data monkeypatch.setattr(mod, "get_blob_store_if_enabled", lambda: SharedStore()) destination = tmp_path / ".tool-results" / "remote_executor-tc-1.txt" def handler(prepared): assert prepared is request assert destination.read_bytes() == data return [] ToolOutputBudgetMiddleware().wrap_model_call(request, handler) assert reads == [ref] def test_digest_matching_local_file_avoids_blob_store(self, monkeypatch, tmp_path): from deerflow.agents.middlewares import tool_output_budget_middleware as mod data = b"already materialized" request, _ = self._request(tmp_path, data) destination = tmp_path / ".tool-results" / "remote_executor-tc-1.txt" destination.parent.mkdir() destination.write_bytes(data) def store_must_not_be_resolved(): raise AssertionError("matching local bytes should stay on the fast path") monkeypatch.setattr(mod, "get_blob_store_if_enabled", store_must_not_be_resolved) ToolOutputBudgetMiddleware().wrap_model_call(request, lambda prepared: []) assert destination.read_bytes() == data def test_wrong_size_local_file_skips_read_before_restore(self, monkeypatch, tmp_path): from deerflow.agents.middlewares import tool_output_budget_middleware as mod data = b"restored shared output" request, _ = self._request(tmp_path, data) destination = tmp_path / ".tool-results" / "remote_executor-tc-1.txt" destination.parent.mkdir() destination.write_bytes(b"wrong size") builtin_open = open def fail_if_stale_file_is_read(file, mode="r", *args, **kwargs): if os.fspath(file) == os.fspath(destination) and mode == "rb": raise AssertionError("wrong-size local file should not be read") return builtin_open(file, mode, *args, **kwargs) class SharedStore: def get_bytes(self, requested: BlobRef) -> bytes: return data monkeypatch.setattr(mod, "open", fail_if_stale_file_is_read, raising=False) monkeypatch.setattr(mod, "get_blob_store_if_enabled", lambda: SharedStore()) ToolOutputBudgetMiddleware().wrap_model_call(request, lambda prepared: []) assert destination.read_bytes() == data def test_invalid_virtual_path_is_ignored_without_blob_read(self, monkeypatch, tmp_path): from deerflow.agents.middlewares import tool_output_budget_middleware as mod request, _ = self._request( tmp_path, b"must-not-escape", virtual_path="/mnt/user-data/outputs/.tool-results/../escape.txt", ) class StoreMustNotBeRead: def get_bytes(self, requested: BlobRef) -> bytes: raise AssertionError(f"untrusted path triggered a blob read: {requested}") monkeypatch.setattr(mod, "get_blob_store_if_enabled", lambda: StoreMustNotBeRead()) ToolOutputBudgetMiddleware().wrap_model_call(request, lambda prepared: []) assert not (tmp_path / "escape.txt").exists() def test_blob_read_failure_does_not_block_model_call(self, monkeypatch, tmp_path): from deerflow.agents.middlewares import tool_output_budget_middleware as mod request, _ = self._request(tmp_path, b"missing") class FailingStore: def get_bytes(self, requested: BlobRef) -> bytes: raise BlobReadError("temporarily unavailable") monkeypatch.setattr(mod, "get_blob_store_if_enabled", lambda: FailingStore()) called = [] destination = tmp_path / ".tool-results" / "remote_executor-tc-1.txt" destination.parent.mkdir() destination.write_bytes(b"corrupt local copy") ToolOutputBudgetMiddleware().wrap_model_call(request, lambda prepared: called.append(prepared) or []) assert called == [request] assert not destination.exists() @pytest.mark.anyio async def test_async_model_call_restores_before_handler(self, monkeypatch, tmp_path): from deerflow.agents.middlewares import tool_output_budget_middleware as mod data = b"async shared output" request, _ = self._request(tmp_path, data) class SharedStore: def get_bytes(self, requested: BlobRef) -> bytes: return data monkeypatch.setattr(mod, "get_blob_store_if_enabled", lambda: SharedStore()) destination = tmp_path / ".tool-results" / "remote_executor-tc-1.txt" async def handler(prepared): assert destination.read_bytes() == data return [] await ToolOutputBudgetMiddleware().awrap_model_call(request, handler) # =========================================================================== # Config integration # =========================================================================== class TestFromAppConfig: def test_from_app_config_with_tool_output(self): config = AppConfig( sandbox=SandboxConfig(use="test"), tool_output={"externalize_min_chars": 5000, "preview_head_chars": 500}, ) mw = ToolOutputBudgetMiddleware.from_app_config(config) assert mw._config.externalize_min_chars == 5000 assert mw._config.preview_head_chars == 500 def test_from_app_config_defaults(self): config = AppConfig(sandbox=SandboxConfig(use="test")) mw = ToolOutputBudgetMiddleware.from_app_config(config) assert mw._config.externalize_min_chars == 12000 class TestPatchModelMessages: def test_returns_none_when_no_changes(self): config = ToolOutputConfig(fallback_max_chars=1000) messages = [_tm("short", name="tool")] assert _patch_model_messages(messages, config) is None def test_patches_oversized_messages(self): config = ToolOutputConfig(fallback_max_chars=500, fallback_head_chars=100, fallback_tail_chars=50) messages = [_tm("x" * 1000, name="tool")] result = _patch_model_messages(messages, config) assert result is not None assert len(result) == 1 assert "omitted" in result[0].content def test_patches_oversized_structured_json_history(self): config = ToolOutputConfig(fallback_max_chars=500, fallback_head_chars=100, fallback_tail_chars=50) messages = [ToolMessage(content=[{"type": "json", "json": {"rows": ["x" * 1000]}}], name="query_rows", tool_call_id="tc-1")] result = _patch_model_messages(messages, config) assert result is not None assert isinstance(result[0].content, str) assert "omitted" in result[0].content # =========================================================================== # Pre-scan helpers (_effective_trigger / _tool_message_over_budget / _needs_budget) # These guard the fast-path optimization — a false negative here is a real bug # (budgeting silently skipped), so per-tool overrides must be honored. # =========================================================================== class TestPreScanHelpers: def test_effective_trigger_uses_global_externalize(self): config = ToolOutputConfig(externalize_min_chars=12000, fallback_max_chars=30000) # smallest of the two thresholds wins assert _effective_trigger("any_tool", config) == 12000 def test_effective_trigger_respects_per_tool_override(self): config = ToolOutputConfig(externalize_min_chars=50000, fallback_max_chars=0, tool_overrides={"sensitive": 100}) assert _effective_trigger("sensitive", config) == 100 # other tools fall back to the (high) global assert _effective_trigger("other", config) == 50000 def test_effective_trigger_per_tool_zero_falls_to_fallback(self): config = ToolOutputConfig(externalize_min_chars=50, tool_overrides={"bash": 0}, fallback_max_chars=200) # externalize disabled for bash → only fallback can trigger assert _effective_trigger("bash", config) == 200 def test_effective_trigger_returns_negative_when_fully_disabled(self): config = ToolOutputConfig(externalize_min_chars=0, fallback_max_chars=0) assert _effective_trigger("any", config) == -1 def test_pre_scan_does_not_short_circuit_per_tool_override(self): """Regression: pre-scan must honor per-tool overrides, not just global threshold.""" config = ToolOutputConfig(externalize_min_chars=50000, fallback_max_chars=0, tool_overrides={"sensitive": 100}) msg = _tm("x" * 500, name="sensitive") # 500 < global 50000 but > per-tool 100 → must still be flagged assert _tool_message_over_budget(msg, config) is True assert _needs_budget(msg, config) is True def test_exempt_tool_never_over_budget(self): config = ToolOutputConfig(externalize_min_chars=10, fallback_max_chars=20, exempt_tools=["read_file"]) msg = _tm("x" * 1000, name="read_file") assert _tool_message_over_budget(msg, config) is False def test_model_call_pre_scan_skips_when_nothing_oversized(self): """_patch_model_messages returns None (no list rebuild) when all messages are small.""" config = ToolOutputConfig(externalize_min_chars=12000, fallback_max_chars=30000) messages = [_tm("small", name="tool"), HumanMessage(content="hi"), _tm("also small", name="bash")] assert _patch_model_messages(messages, config) is None # =========================================================================== # Middleware ordering in the chain # =========================================================================== class TestMiddlewareChainIntegration: def test_budget_middleware_is_first_in_chain(self): from deerflow.agents.middlewares.tool_error_handling_middleware import build_subagent_runtime_middlewares app_config = AppConfig(sandbox=SandboxConfig(use="test")) middlewares = build_subagent_runtime_middlewares(app_config=app_config, lazy_init=False) # InputSanitizationMiddleware is the outermost wrap_model_call wrapper; # KnowledgeScopeMiddleware cleans model input immediately inside it; # ToolOutputBudgetMiddleware remains immediately inside the scope guard. from deerflow.agents.middlewares.input_sanitization_middleware import InputSanitizationMiddleware from deerflow.agents.middlewares.knowledge_scope_middleware import KnowledgeScopeMiddleware assert isinstance(middlewares[0], InputSanitizationMiddleware) assert isinstance(middlewares[1], KnowledgeScopeMiddleware) assert isinstance(middlewares[2], ToolOutputBudgetMiddleware) def test_budget_middleware_in_lead_chain(self): from deerflow.agents.middlewares.tool_error_handling_middleware import build_lead_runtime_middlewares app_config = AppConfig(sandbox=SandboxConfig(use="test")) middlewares = build_lead_runtime_middlewares(app_config=app_config, lazy_init=False) from deerflow.agents.middlewares.input_sanitization_middleware import InputSanitizationMiddleware from deerflow.agents.middlewares.knowledge_scope_middleware import KnowledgeScopeMiddleware assert isinstance(middlewares[0], InputSanitizationMiddleware) assert isinstance(middlewares[1], KnowledgeScopeMiddleware) assert isinstance(middlewares[2], ToolOutputBudgetMiddleware) # =========================================================================== # Config version bump # =========================================================================== class TestConfigVersion: def test_config_version_bumped(self): import yaml example_path = os.path.join(os.path.dirname(__file__), "..", "..", "config.example.yaml") if os.path.exists(example_path): with open(example_path, encoding="utf-8") as f: data = yaml.safe_load(f) assert data.get("config_version", 0) >= 11 def test_config_example_has_tool_output_section(self): import yaml example_path = os.path.join(os.path.dirname(__file__), "..", "..", "config.example.yaml") if os.path.exists(example_path): with open(example_path, encoding="utf-8") as f: data = yaml.safe_load(f) assert "tool_output" in data tool_output = data["tool_output"] assert tool_output["enabled"] is True assert tool_output["externalize_min_chars"] == 12000 assert "read_file" in tool_output["exempt_tools"] # =========================================================================== # externalize into sandbox for non-mounted (remote) sandboxes # =========================================================================== class _FakeSandbox: """In-memory stand-in for a Sandbox. Records calls and supports failure injection.""" def __init__(self, *, write_ok: bool = True, check_result: str | None = None) -> None: self.commands: list[str] = [] self.writes: list[tuple[str, str]] = [] self.files: dict[str, str] = {} self._write_ok = write_ok self._check_result = check_result def execute_command( self, command: str, env: dict[str, str] | None = None, timeout: float | None = None, ) -> str: del env, timeout self.commands.append(command) if self._check_result is not None: return self._check_result # Simulate shell execution of: # test -f && test "$(wc -c < )" -eq && echo OK || echo MISSING match = re.search(r'test -f (\S+) && test "\$\(wc -c < \S+\)" -eq (\d+)', command) if match: path, expected_bytes = match.group(1), int(match.group(2)) if path in self.files: actual_bytes = len(self.files[path].encode("utf-8")) if actual_bytes == expected_bytes: return "OK" return "MISSING" return "" def write_file(self, path: str, content: str, append: bool = False) -> None: if not self._write_ok: raise RuntimeError("simulated write failure") self.writes.append((path, content)) if append and path in self.files: self.files[path] += content else: self.files[path] = content class _FakeProvider: """Minimal SandboxProvider stand-in for monkeypatching get_sandbox_provider.""" def __init__(self, *, uses_thread_data_mounts: bool, sandbox: _FakeSandbox | None = None) -> None: self.uses_thread_data_mounts = uses_thread_data_mounts self._sandbox = sandbox def get(self, sandbox_id: str): return self._sandbox class TestExternalizeToSandbox: def test_writes_and_returns_virtual_path(self): from deerflow.agents.middlewares.tool_output_budget_middleware import ( _externalize_to_sandbox, ) sb = _FakeSandbox() result = _externalize_to_sandbox( "x" * 100, tool_name="bash", tool_call_id="tc-1", storage_subdir=".tool-results", sandbox=sb, ) assert result is not None assert result.startswith("/mnt/user-data/outputs/.tool-results/bash-") assert result.endswith(".log") assert any(c.startswith("mkdir -p ") for c in sb.commands) assert any("wc -c" in c for c in sb.commands) assert any("-eq 100" in c for c in sb.commands) assert sb.writes and sb.writes[0][0] == result assert sb.writes[0][1] == "x" * 100 def test_returns_none_when_write_raises(self): from deerflow.agents.middlewares.tool_output_budget_middleware import ( _externalize_to_sandbox, ) result = _externalize_to_sandbox( "x" * 100, tool_name="web_fetch", tool_call_id="tc-2", storage_subdir=".tool-results", sandbox=_FakeSandbox(write_ok=False), ) assert result is None def test_returns_none_when_validation_fails(self): from deerflow.agents.middlewares.tool_output_budget_middleware import ( _externalize_to_sandbox, ) result = _externalize_to_sandbox( "x" * 100, tool_name="bash", tool_call_id="tc-3", storage_subdir=".tool-results", sandbox=_FakeSandbox(check_result="MISSING"), ) assert result is None def test_returns_none_when_byte_size_is_mismatched(self): """A truncated write (fewer bytes than expected) fails validation and returns None.""" from deerflow.agents.middlewares.tool_output_budget_middleware import ( _externalize_to_sandbox, ) class _TruncatingSandbox(_FakeSandbox): def write_file(self, path: str, content: str, append: bool = False) -> None: # Simulate a truncated write (e.g. disk full / broken pipe) where only half lands super().write_file(path, content[: len(content) // 2], append=append) sb = _TruncatingSandbox() result = _externalize_to_sandbox( "x" * 100, tool_name="bash", tool_call_id="tc-3-truncated", storage_subdir=".tool-results", sandbox=sb, ) assert result is None assert any("-eq 100" in c for c in sb.commands) # Confirm the file was actually written with half size, triggering the real byte mismatch assert sb.writes and len(sb.writes[0][1]) == 50 assert any(len(content) == 50 for content in sb.files.values()) def test_rejects_unsafe_storage_subdir(self): from deerflow.agents.middlewares.tool_output_budget_middleware import ( _externalize_to_sandbox, ) sb = _FakeSandbox() assert ( _externalize_to_sandbox( "x" * 100, tool_name="bash", tool_call_id="tc-4", storage_subdir="../escape", sandbox=sb, ) is None ) assert ( _externalize_to_sandbox( "x" * 100, tool_name="bash", tool_call_id="tc-5", storage_subdir="/abs/path", sandbox=sb, ) is None ) # Sandbox must not be touched when the subdir is rejected up-front. assert sb.commands == [] assert sb.writes == [] def test_default_extension_for_unknown_tool(self): from deerflow.agents.middlewares.tool_output_budget_middleware import ( _externalize_to_sandbox, ) result = _externalize_to_sandbox( "data", tool_name="unknown_tool", tool_call_id="tc-6", storage_subdir=".tool-results", sandbox=_FakeSandbox(), ) assert result is not None and result.endswith(".txt") class TestBudgetContentSandboxDispatch: """_budget_content must branch on uses_thread_data_mounts (issue #3416).""" def test_mounted_sandbox_uses_host_disk(self, monkeypatch, tmp_path): from deerflow.agents.middlewares import tool_output_budget_middleware as mod sb = _FakeSandbox() monkeypatch.setattr( mod, "get_sandbox_provider", lambda: _FakeProvider(uses_thread_data_mounts=True, sandbox=sb), ) config = ToolOutputConfig(externalize_min_chars=50, preview_head_chars=20, preview_tail_chars=10) result = mod._budget_content( "x" * 500, tool_name="remote_executor", tool_call_id="tc-m", outputs_path=str(tmp_path), config=config, sandbox=sb, ) assert result is not None assert "Full remote_executor output saved to /mnt/user-data/outputs/" in result[0] assert result[1] == "externalized" # Mounted path must NOT touch the sandbox. assert sb.commands == [] assert sb.writes == [] # And the host file must exist. storage_dir = tmp_path / ".tool-results" assert storage_dir.is_dir() assert len(list(storage_dir.iterdir())) == 1 def test_non_mounted_sandbox_writes_to_sandbox(self, monkeypatch, tmp_path): from deerflow.agents.middlewares import tool_output_budget_middleware as mod sb = _FakeSandbox() blob_calls = {"count": 0} def blob_store_must_not_be_resolved(): blob_calls["count"] += 1 raise AssertionError("sandbox-resident output is outside the host blob migration") monkeypatch.setattr( mod, "get_sandbox_provider", lambda: _FakeProvider(uses_thread_data_mounts=False, sandbox=sb), ) monkeypatch.setattr(mod, "get_blob_store_if_enabled", blob_store_must_not_be_resolved) config = ToolOutputConfig(externalize_min_chars=50, preview_head_chars=20, preview_tail_chars=10) result = mod._budget_content( "x" * 500, tool_name="remote_executor", tool_call_id="tc-n", outputs_path=str(tmp_path), # present, but ignored on non-mounted path config=config, sandbox=sb, ) assert result is not None assert "Full remote_executor output saved to /mnt/user-data/outputs/" in result[0] assert result[1] == "externalized" # Non-mounted path MUST write into the sandbox. assert sb.writes and sb.writes[0][1] == "x" * 500 # And MUST NOT touch the host. assert not (tmp_path / ".tool-results").exists() assert blob_calls["count"] == 0 def test_non_mounted_without_sandbox_falls_back(self, monkeypatch): from deerflow.agents.middlewares import tool_output_budget_middleware as mod monkeypatch.setattr( mod, "get_sandbox_provider", lambda: _FakeProvider(uses_thread_data_mounts=False, sandbox=None), ) config = ToolOutputConfig( externalize_min_chars=50, fallback_max_chars=500, fallback_head_chars=100, fallback_tail_chars=50, ) result = mod._budget_content( "x" * 5000, tool_name="web_search", tool_call_id="tc-fb", outputs_path=None, config=config, sandbox=None, ) assert result is not None assert "Persistent storage unavailable" in result[0] assert result[1] == "truncated" class TestResolveSandbox: def test_returns_none_when_no_state(self): from deerflow.agents.middlewares.tool_output_budget_middleware import _resolve_sandbox req = SimpleNamespace(runtime=None) assert _resolve_sandbox(req) is None def test_returns_none_when_state_has_no_sandbox(self): from deerflow.agents.middlewares.tool_output_budget_middleware import _resolve_sandbox req = SimpleNamespace(runtime=SimpleNamespace(state={})) assert _resolve_sandbox(req) is None def test_returns_none_when_sandbox_id_missing(self): from deerflow.agents.middlewares.tool_output_budget_middleware import _resolve_sandbox req = SimpleNamespace(runtime=SimpleNamespace(state={"sandbox": {}})) assert _resolve_sandbox(req) is None def test_returns_sandbox_from_provider(self, monkeypatch): from deerflow.agents.middlewares import tool_output_budget_middleware as mod sb = _FakeSandbox() monkeypatch.setattr( mod, "get_sandbox_provider", lambda: _FakeProvider(uses_thread_data_mounts=False, sandbox=sb), ) req = SimpleNamespace(runtime=SimpleNamespace(state={"sandbox": {"sandbox_id": "sb-1"}})) assert mod._resolve_sandbox(req) is sb def test_returns_sandbox_from_provider_when_overwrite_wrapped(self, monkeypatch): from langgraph.types import Overwrite from deerflow.agents.middlewares import tool_output_budget_middleware as mod sb = _FakeSandbox() monkeypatch.setattr( mod, "get_sandbox_provider", lambda: _FakeProvider(uses_thread_data_mounts=False, sandbox=sb), ) # Fork-restored state delivers sandbox wrapped in Overwrite req = SimpleNamespace(runtime=SimpleNamespace(state={"sandbox": Overwrite({"sandbox_id": "sb-fork"})})) assert mod._resolve_sandbox(req) is sb def test_returns_none_on_provider_exception(self, monkeypatch): from deerflow.agents.middlewares import tool_output_budget_middleware as mod class _Boom: def get(self, sandbox_id): raise RuntimeError("boom") monkeypatch.setattr(mod, "get_sandbox_provider", lambda: _Boom()) req = SimpleNamespace(runtime=SimpleNamespace(state={"sandbox": {"sandbox_id": "sb-x"}})) assert mod._resolve_sandbox(req) is None class TestWrapToolCallSandboxIntegration: """End-to-end via wrap_tool_call for the non-mounted path (issue #3416).""" def test_oversized_output_lands_in_sandbox_not_host(self, monkeypatch, tmp_path): from deerflow.agents.middlewares import tool_output_budget_middleware as mod sb = _FakeSandbox() monkeypatch.setattr( mod, "get_sandbox_provider", lambda: _FakeProvider(uses_thread_data_mounts=False, sandbox=sb), ) config = ToolOutputConfig(externalize_min_chars=50, preview_head_chars=20, preview_tail_chars=10) mw = ToolOutputBudgetMiddleware(config=config) content = "x" * 500 msg = _tm(content, name="remote_executor") # Request carries BOTH outputs_path (host) AND a sandbox_id; the # non-mounted branch must ignore outputs_path and write into sandbox. req = SimpleNamespace( tool_call={"name": "remote_executor", "id": "tc-1"}, runtime=SimpleNamespace( state={ "thread_data": {"outputs_path": str(tmp_path)}, "sandbox": {"sandbox_id": "sb-1"}, } ), ) result = mw.wrap_tool_call(req, lambda _: msg) assert isinstance(result, ToolMessage) assert "Full remote_executor output saved to /mnt/user-data/outputs/" in result.content assert sb.writes and sb.writes[0][1] == content # Host disk must not have been written. assert not (tmp_path / ".tool-results").exists() class TestBudgetContentNoSandboxNoProviderCall: """Without a sandbox, _budget_content must NOT call get_sandbox_provider. This is the legacy host-disk path (and the CI-without-config.yaml path): touching the provider would raise and force inline fallback, regressing issue #3416's fix and breaking environments that never opt into sandbox. """ def test_no_provider_call_when_sandbox_absent(self, monkeypatch, tmp_path): from deerflow.agents.middlewares import tool_output_budget_middleware as mod called = {"n": 0} def boom(): called["n"] += 1 raise RuntimeError("provider must not be called on the legacy path") monkeypatch.setattr(mod, "get_sandbox_provider", boom) config = ToolOutputConfig(externalize_min_chars=50, preview_head_chars=20, preview_tail_chars=10) result = mod._budget_content( "x" * 500, tool_name="remote_executor", tool_call_id="tc-legacy", outputs_path=str(tmp_path), config=config, sandbox=None, ) assert result is not None assert "Full remote_executor output saved to /mnt/user-data/outputs/" in result[0] assert result[1] == "externalized" assert called["n"] == 0 assert (tmp_path / ".tool-results").is_dir() # =========================================================================== # Superseded write payload elision (issue #5328, step 2) # =========================================================================== def _meta_result(name: str, tool_call_id: str, content: str = "OK", *, status: str | None = "success") -> ToolMessage: """A ToolMessage stamped the way ToolErrorHandlingMiddleware stamps it; ``status=None`` leaves it unstamped.""" msg = ToolMessage(content=content, name=name, tool_call_id=tool_call_id, status="error" if status == "error" else "success") if status is not None: msg.additional_kwargs["deerflow_tool_meta"] = { "status": status, "error_type": None, "recoverable_by_model": True, "recommended_next_action": "continue", "source": "content_analysis", } return msg def _write(tool_call_id: str, path: str, content: str, *, append: bool = False, status: str | None = "success") -> tuple[AIMessage, ToolMessage]: args = {"description": "d", "path": path, "content": content, "append": append} ai = AIMessage(content="", tool_calls=[{"name": "write_file", "id": tool_call_id, "args": args}]) return ai, _meta_result("write_file", tool_call_id, "Error: boom" if status == "error" else "OK", status=status) def _read(tool_call_id: str, path: str, *, status: str | None = "success") -> tuple[AIMessage, ToolMessage]: ai = AIMessage(content="", tool_calls=[{"name": "read_file", "id": tool_call_id, "args": {"path": path}}]) return ai, _meta_result("read_file", tool_call_id, "Error: File not found" if status == "error" else "file text", status=status) def _str_replace(tool_call_id: str, path: str, *, new_str: str = "n", status: str | None = "success") -> tuple[AIMessage, ToolMessage]: ai = AIMessage(content="", tool_calls=[{"name": "str_replace", "id": tool_call_id, "args": {"path": path, "old_str": "o", "new_str": new_str}}]) return ai, _meta_result("str_replace", tool_call_id, "Error: boom" if status == "error" else "OK", status=status) class TestSupersededWriteElision: """Model-bound requests drop the content of successful write_file calls superseded by a later read or write of the same path.""" PATH = "/mnt/user-data/outputs/report.md" OTHER = "/mnt/user-data/outputs/other.md" @staticmethod def _middleware(**overrides) -> ToolOutputBudgetMiddleware: return ToolOutputBudgetMiddleware(config=ToolOutputConfig(**overrides)) @staticmethod def _model_request(messages) -> ModelRequest: return ModelRequest(model=None, messages=list(messages), tools=[], state={"messages": list(messages)}) def _forward(self, mw: ToolOutputBudgetMiddleware, messages) -> tuple[ModelRequest, ModelRequest]: """Run ``wrap_model_call`` and return ``(original request, request the handler received)``.""" captured: dict[str, ModelRequest] = {} def handler(req): captured["request"] = req return AIMessage(content="ok") request = self._model_request(messages) mw.wrap_model_call(request, handler) return request, captured["request"] @staticmethod def _content(forwarded: ModelRequest, index: int) -> str: return forwarded.messages[index].tool_calls[0]["args"]["content"] # -- policy ------------------------------------------------------------ def test_superseded_write_is_elided_and_newest_write_is_kept(self): mw = self._middleware() payload = "x" * 5000 human = HumanMessage(content="go") w1, r1 = _write("call-1", self.PATH, payload) rd, rr = _read("call-2", self.PATH) w2, r2 = _write("call-3", self.PATH, "y" * 5000, append=True) request, forwarded = self._forward(mw, [human, w1, r1, rd, rr, w2, r2]) assert forwarded is not request elided = forwarded.messages[1].tool_calls[0]["args"] assert elided["content"].startswith("[content elided: 5000 chars") assert "read_file" in elided["content"] assert payload not in elided["content"] assert elided["path"] == self.PATH assert elided["description"] == "d" assert elided["append"] is False # The newest successful write stays visible (keep_recent_writes=1). assert forwarded.messages[5] is w2 # Untouched neighbours pass through by identity; state and stored history keep the original. assert forwarded.messages[0] is human assert forwarded.messages[2] is r1 assert forwarded.messages[3] is rd assert request.messages[1] is w1 assert request.state["messages"][1] is w1 assert w1.tool_calls[0]["args"]["content"] == payload def test_write_without_a_later_touch_of_the_path_is_kept(self): mw = self._middleware() w1, r1 = _write("call-1", self.PATH, "x" * 5000) w2, r2 = _write("call-2", self.OTHER, "y" * 5000) w3, r3 = _write("call-3", "/mnt/user-data/outputs/third.md", "z" * 5000) request, forwarded = self._forward(mw, [w1, r1, w2, r2, w3, r3]) assert forwarded is request def test_newest_write_is_kept_even_when_superseded(self): w1, r1 = _write("call-1", self.PATH, "x" * 5000) rd, rr = _read("call-2", self.PATH) request, forwarded = self._forward(self._middleware(), [w1, r1, rd, rr]) assert forwarded is request _request, forwarded = self._forward(self._middleware(keep_recent_writes=0), [w1, r1, rd, rr]) assert self._content(forwarded, 0).startswith("[content elided: 5000 chars") def test_keep_recent_counts_successful_writes_across_paths(self): w1, r1 = _write("call-1", self.PATH, "x" * 5000) rd, rr = _read("call-2", self.PATH) w2, r2 = _write("call-3", self.OTHER, "tiny") w3, r3 = _write("call-4", "/mnt/user-data/outputs/third.md", "tiny") history = [w1, r1, rd, rr, w2, r2, w3, r3] _request, forwarded = self._forward(self._middleware(keep_recent_writes=2), history) assert self._content(forwarded, 0).startswith("[content elided") request, forwarded = self._forward(self._middleware(keep_recent_writes=3), history) assert forwarded is request def test_later_successful_write_of_the_same_path_supersedes(self): mw = self._middleware() w1, r1 = _write("call-1", self.PATH, "x" * 5000) w2, r2 = _write("call-2", self.PATH, "y" * 5000) w3, r3 = _write("call-3", self.OTHER, "z" * 5000) _request, forwarded = self._forward(mw, [w1, r1, w2, r2, w3, r3]) assert self._content(forwarded, 0).startswith("[content elided: 5000 chars") assert forwarded.messages[2] is w2 # not superseded, and older than the newest write assert forwarded.messages[4] is w3 def test_later_successful_str_replace_supersedes(self): mw = self._middleware() w1, r1 = _write("call-1", self.PATH, "x" * 5000) sr, srr = _str_replace("call-2", self.PATH) w2, r2 = _write("call-3", self.OTHER, "tiny") _request, forwarded = self._forward(mw, [w1, r1, sr, srr, w2, r2]) assert self._content(forwarded, 0).startswith("[content elided") assert forwarded.messages[2] is sr @pytest.mark.parametrize("status", ["error", "partial_success", None], ids=["error", "partial", "unstamped"]) def test_later_write_that_did_not_succeed_does_not_supersede(self, status): mw = self._middleware() w1, r1 = _write("call-1", self.PATH, "x" * 5000) w2, r2 = _write("call-2", self.PATH, "y" * 5000, status=status) w3, r3 = _write("call-3", self.OTHER, "tiny") request, forwarded = self._forward(mw, [w1, r1, w2, r2, w3, r3]) assert forwarded is request def test_gate_blocked_later_write_does_not_supersede(self): from deerflow.agents.middlewares.read_before_write_middleware import WRITE_BLOCK_KEY mw = self._middleware() w1, r1 = _write("call-1", self.PATH, "x" * 5000) w2, r2 = _write("call-2", self.PATH, "y" * 5000, status="error") r2.additional_kwargs[WRITE_BLOCK_KEY] = {"path": self.PATH, "tool": "write_file"} w3, r3 = _write("call-3", self.OTHER, "tiny") request, forwarded = self._forward(mw, [w1, r1, w2, r2, w3, r3]) assert forwarded is request @pytest.mark.parametrize("status", ["error", None], ids=["error", "unstamped"]) def test_read_that_did_not_succeed_does_not_supersede(self, status): mw = self._middleware() w1, r1 = _write("call-1", self.PATH, "x" * 5000) rd, rr = _read("call-2", self.PATH, status=status) w2, r2 = _write("call-3", self.OTHER, "tiny") request, forwarded = self._forward(mw, [w1, r1, rd, rr, w2, r2]) assert forwarded is request def test_partial_read_still_supersedes(self): """A truncated or ranged read still showed the model the on-disk file.""" mw = self._middleware() w1, r1 = _write("call-1", self.PATH, "x" * 5000) rd, rr = _read("call-2", self.PATH, status="partial_success") w2, r2 = _write("call-3", self.OTHER, "tiny") _request, forwarded = self._forward(mw, [w1, r1, rd, rr, w2, r2]) assert self._content(forwarded, 0).startswith("[content elided") @pytest.mark.parametrize("status", ["error", "partial_success", None], ids=["error", "partial", "unstamped"]) def test_write_that_did_not_succeed_is_never_a_candidate(self, status): mw = self._middleware() w1, r1 = _write("call-1", self.PATH, "x" * 5000, status=status) rd, rr = _read("call-2", self.PATH) w2, r2 = _write("call-3", self.OTHER, "tiny") request, forwarded = self._forward(mw, [w1, r1, rd, rr, w2, r2]) assert forwarded is request def test_unanswered_write_is_never_a_candidate(self): mw = self._middleware() w1, _unused = _write("call-1", self.PATH, "x" * 5000) rd, rr = _read("call-2", self.PATH) w2, r2 = _write("call-3", self.OTHER, "tiny") request, forwarded = self._forward(mw, [w1, rd, rr, w2, r2]) assert forwarded is request def test_same_turn_read_does_not_supersede(self): """Parallel calls in one AIMessage run in no fixed order, so the read may predate the write.""" mw = self._middleware() payload = "x" * 5000 ai = AIMessage( content="", tool_calls=[ {"name": "write_file", "id": "call-1", "args": {"description": "d", "path": self.PATH, "content": payload}}, {"name": "read_file", "id": "call-2", "args": {"path": self.PATH}}, ], ) results = [_meta_result("write_file", "call-1"), _meta_result("read_file", "call-2", "file text")] w2, r2 = _write("call-3", self.OTHER, "tiny") request, forwarded = self._forward(mw, [ai, *results, w2, r2]) assert forwarded is request def test_paths_are_normalized_before_matching(self): mw = self._middleware() w1, r1 = _write("call-1", "/mnt/user-data/outputs/./report.md", "x" * 5000) rd, rr = _read("call-2", "/mnt/user-data/outputs/sub/../report.md") w2, r2 = _write("call-3", self.OTHER, "tiny") _request, forwarded = self._forward(mw, [w1, r1, rd, rr, w2, r2]) assert self._content(forwarded, 0).startswith("[content elided") def test_different_path_does_not_supersede(self): mw = self._middleware() w1, r1 = _write("call-1", self.PATH, "x" * 5000) rd, rr = _read("call-2", self.OTHER) w2, r2 = _write("call-3", self.OTHER, "tiny") request, forwarded = self._forward(mw, [w1, r1, rd, rr, w2, r2]) assert forwarded is request def test_str_replace_payloads_are_never_elided(self): mw = self._middleware() sr, srr = _str_replace("call-1", self.PATH, new_str="n" * 5000) rd, rr = _read("call-2", self.PATH) w2, r2 = _write("call-3", self.OTHER, "tiny") request, forwarded = self._forward(mw, [sr, srr, rd, rr, w2, r2]) assert forwarded is request # -- thresholds and config --------------------------------------------- def test_content_below_min_chars_stays_visible(self): mw = self._middleware() w1, r1 = _write("call-1", self.PATH, "short " * 20) rd, rr = _read("call-2", self.PATH) w2, r2 = _write("call-3", self.OTHER, "tiny") request, forwarded = self._forward(mw, [w1, r1, rd, rr, w2, r2]) assert forwarded is request def test_min_chars_zero_elides_any_non_empty_content(self): mw = self._middleware(superseded_write_min_chars=0) w1, r1 = _write("call-1", self.PATH, "v1") rd, rr = _read("call-2", self.PATH) w2, r2 = _write("call-3", self.OTHER, "") rd2, rr2 = _read("call-4", self.OTHER) w3, r3 = _write("call-5", "/mnt/user-data/outputs/third.md", "tiny") _request, forwarded = self._forward(mw, [w1, r1, rd, rr, w2, r2, rd2, rr2, w3, r3]) assert self._content(forwarded, 0).startswith("[content elided: 2 chars") assert forwarded.messages[4] is w2 def test_non_string_content_is_left_alone(self): mw = self._middleware(superseded_write_min_chars=0) ai = AIMessage(content="", tool_calls=[{"name": "write_file", "id": "call-1", "args": {"path": self.PATH, "content": ["not", "a", "string"]}}]) rd, rr = _read("call-2", self.PATH) w2, r2 = _write("call-3", self.OTHER, "tiny") request, forwarded = self._forward(mw, [ai, _meta_result("write_file", "call-1"), rd, rr, w2, r2]) assert forwarded is request def test_disabled_by_config_passes_request_through(self): mw = self._middleware(elide_superseded_writes=False) w1, r1 = _write("call-1", self.PATH, "x" * 5000) rd, rr = _read("call-2", self.PATH) w2, r2 = _write("call-3", self.OTHER, "tiny") request, forwarded = self._forward(mw, [w1, r1, rd, rr, w2, r2]) assert forwarded is request def test_middleware_disabled_passes_request_through(self): mw = self._middleware(enabled=False) w1, r1 = _write("call-1", self.PATH, "x" * 5000) rd, rr = _read("call-2", self.PATH) w2, r2 = _write("call-3", self.OTHER, "tiny") request, forwarded = self._forward(mw, [w1, r1, rd, rr, w2, r2]) assert forwarded is request def test_config_defaults(self): config = ToolOutputConfig() assert config.elide_superseded_writes is True assert config.superseded_write_min_chars == 2000 assert config.keep_recent_writes == 1 def test_release_policy_declares_config(self): params = self._middleware(keep_recent_writes=3, superseded_write_min_chars=123).release_policy_parameters() assert params["config"]["elide_superseded_writes"] is True assert params["config"]["superseded_write_min_chars"] == 123 assert params["config"]["keep_recent_writes"] == 3 def test_from_app_config_passes_the_keys(self): config = AppConfig(sandbox=SandboxConfig(use="test"), tool_output={"superseded_write_min_chars": 10, "keep_recent_writes": 0}) mw = ToolOutputBudgetMiddleware.from_app_config(config) assert mw._config.superseded_write_min_chars == 10 assert mw._config.keep_recent_writes == 0 def test_config_example_documents_the_keys(self): import yaml example_path = os.path.join(os.path.dirname(__file__), "..", "..", "config.example.yaml") with open(example_path, encoding="utf-8") as f: data = yaml.safe_load(f) tool_output = data["tool_output"] assert tool_output["elide_superseded_writes"] is True assert tool_output["superseded_write_min_chars"] == 2000 assert tool_output["keep_recent_writes"] == 1 # New user-settable keys are a schema change: the outdated-config warning must fire. assert data["config_version"] >= 42 # -- surfaces, pairing, determinism, composition ------------------------ def test_rewrites_every_provider_surface_together(self): mw = self._middleware() payload = "y" * 5000 args = {"description": "d", "path": self.PATH, "content": payload} ai = AIMessage( content=[ {"type": "text", "text": "writing"}, {"type": "tool_use", "id": "call-1", "name": "write_file", "input": dict(args), "partial_json": json.dumps(args)}, ], tool_calls=[{"name": "write_file", "id": "call-1", "args": dict(args)}], additional_kwargs={"tool_calls": [{"id": "call-1", "type": "function", "function": {"name": "write_file", "arguments": json.dumps(args)}}]}, ) rd, rr = _read("call-2", self.PATH) w2, r2 = _write("call-3", self.OTHER, "tiny") _request, forwarded = self._forward(mw, [ai, _meta_result("write_file", "call-1"), rd, rr, w2, r2]) rewritten = forwarded.messages[0] structured = rewritten.tool_calls[0]["args"] assert structured["content"].startswith("[content elided: 5000 chars") raw = json.loads(rewritten.additional_kwargs["tool_calls"][0]["function"]["arguments"]) assert raw == structured block = rewritten.content[1] assert block["input"] == structured assert "partial_json" not in block assert rewritten.content[0] == {"type": "text", "text": "writing"} assert payload not in json.dumps(rewritten.model_dump(), ensure_ascii=False) # Original objects are untouched. assert ai.content[1]["input"]["content"] == payload assert payload in ai.additional_kwargs["tool_calls"][0]["function"]["arguments"] def test_reused_call_ids_pair_per_occurrence(self): """A failed write and a later successful one may share a tool-call id; the failed one must not inherit success.""" mw = self._middleware() w1, r1 = _write("call-1", self.PATH, "x" * 5000, status="error") rd, rr = _read("call-2", self.PATH) w2, r2 = _write("call-1", self.OTHER, "y" * 5000) w3, r3 = _write("call-3", "/mnt/user-data/outputs/third.md", "tiny") request, forwarded = self._forward(mw, [w1, r1, rd, rr, w2, r2, w3, r3]) assert forwarded is request def test_elision_is_deterministic_across_model_calls(self): mw = self._middleware() w1, r1 = _write("call-1", self.PATH, "x" * 5000) rd, rr = _read("call-2", self.PATH) w2, r2 = _write("call-3", self.OTHER, "tiny") history = [w1, r1, rd, rr, w2, r2] _request, first = self._forward(mw, history) _request, second = self._forward(mw, history) assert first.messages[0].tool_calls == second.messages[0].tool_calls def test_elision_is_monotonic_as_history_grows(self): """Once a write is elided, appending more history never brings its content back.""" mw = self._middleware() w1, r1 = _write("call-1", self.PATH, "x" * 5000) rd, rr = _read("call-2", self.PATH) w2, r2 = _write("call-3", self.OTHER, "tiny") history = [w1, r1, rd, rr, w2, r2] _request, before = self._forward(mw, history) assert self._content(before, 0).startswith("[content elided") w3, r3 = _write("call-4", "/mnt/user-data/outputs/third.md", "z" * 5000) _request, after = self._forward(mw, [*history, HumanMessage(content="more"), w3, r3]) assert after.messages[0].tool_calls == before.messages[0].tool_calls def test_rewritten_history_drops_openai_response_chain_ids(self): mw = self._middleware() w1, r1 = _write("call-1", self.PATH, "x" * 5000) w1.response_metadata = {"id": "resp_write", "output_version": "responses/v1"} rd, rr = _read("call-2", self.PATH) w2, r2 = _write("call-3", self.OTHER, "tiny") w2.response_metadata = {"id": "resp_latest"} _request, forwarded = self._forward(mw, [w1, r1, rd, rr, w2, r2]) assert "id" not in forwarded.messages[0].response_metadata assert "id" not in forwarded.messages[4].response_metadata def test_applies_alongside_historical_output_truncation(self): mw = self._middleware(fallback_max_chars=500, fallback_head_chars=100, fallback_tail_chars=50) w1, r1 = _write("call-1", self.PATH, "x" * 5000) rd, rr = _read("call-2", self.PATH) oversized = _tm("q" * 1000, name="tool", tool_call_id="tc-q") w2, r2 = _write("call-3", self.OTHER, "tiny") _request, forwarded = self._forward(mw, [w1, r1, rd, rr, oversized, w2, r2]) assert self._content(forwarded, 0).startswith("[content elided") assert "omitted" in forwarded.messages[4].content assert forwarded.messages[5] is w2 def test_async_model_call_elides(self): import asyncio mw = self._middleware() w1, r1 = _write("call-1", self.PATH, "x" * 5000) rd, rr = _read("call-2", self.PATH) w2, r2 = _write("call-3", self.OTHER, "tiny") request = self._model_request([w1, r1, rd, rr, w2, r2]) seen: dict[str, ModelRequest] = {} async def handler(req): seen["request"] = req return AIMessage(content="ok") asyncio.run(mw.awrap_model_call(request, handler)) assert seen["request"] is not request assert self._content(seen["request"], 0).startswith("[content elided") def test_no_write_calls_in_history_is_a_cheap_no_op(self): mw = self._middleware() history = [HumanMessage(content="go"), AIMessage(content="", tool_calls=[{"name": "bash", "id": "call-1", "args": {"command": "ls"}}]), _meta_result("bash", "call-1", "files")] request, forwarded = self._forward(mw, history) assert forwarded is request def test_chat_completions_payload_uses_the_placeholder(self): """End to end against the OpenAI chat-completions message converter.""" from langchain_openai.chat_models.base import _convert_message_to_dict mw = self._middleware() payload = "x" * 5000 w1, r1 = _write("call-1", self.PATH, payload) rd, rr = _read("call-2", self.PATH) w2, r2 = _write("call-3", self.OTHER, "tiny") _request, forwarded = self._forward(mw, [w1, r1, rd, rr, w2, r2]) wire = json.loads(_convert_message_to_dict(forwarded.messages[0])["tool_calls"][0]["function"]["arguments"]) assert wire["content"].startswith("[content elided: 5000 chars") assert payload not in json.dumps(wire) def test_unanswered_write_with_a_reused_id_is_never_treated_as_successful(self): """Review on #5374: an interrupted write must not inherit the success of a later call that reused its id.""" mw = self._middleware() draft = "d" * 5000 interrupted, _never_delivered = _write("reused", self.PATH, draft) rd, rr = _read("call-2", self.PATH) later, later_ok = _write("reused", self.OTHER, "n" * 5000) last, last_ok = _write("call-4", "/mnt/user-data/outputs/last.md", "l" * 5000) request, forwarded = self._forward(mw, [interrupted, rd, rr, later, later_ok, last, last_ok]) assert forwarded is request assert interrupted.tool_calls[0]["args"]["content"] == draft def test_duplicate_ids_in_one_turn_are_never_rewritten(self): """Review on #5374: a failed sibling sharing the id of a superseded successful write must not be rewritten into it.""" mw = self._middleware() turn = AIMessage( content="", tool_calls=[ {"name": "write_file", "id": "dup", "args": {"path": self.PATH, "content": "a" * 5000}}, {"name": "write_file", "id": "dup", "args": {"path": self.OTHER, "content": "b" * 5000}}, ], ) results = [_meta_result("write_file", "dup"), _meta_result("write_file", "dup", "Error: boom", status="error")] rd, rr = _read("call-2", self.PATH) last, last_ok = _write("call-3", "/mnt/user-data/outputs/last.md", "l" * 5000) request, forwarded = self._forward(mw, [turn, *results, rd, rr, last, last_ok]) assert forwarded is request assert [call["args"]["path"] for call in turn.tool_calls] == [self.PATH, self.OTHER] assert turn.tool_calls[1]["args"]["content"] == "b" * 5000 def test_unhashable_sibling_id_does_not_crash_the_model_call(self): """Review on #5374 (round 3): a malformed sibling id next to an elision candidate must be skipped, not hashed.""" mw = self._middleware(keep_recent_writes=0) payload = "x" * 5000 turn, ok = _write("call-1", self.PATH, payload) turn.tool_calls.append({"name": "bash", "id": ["not", "a", "string"], "args": {"command": "ls"}}) rd, rr = _read("call-2", self.PATH) _request, forwarded = self._forward(mw, [turn, ok, rd, rr]) assert self._content(forwarded, 0).startswith("[content elided: 5000 chars") assert forwarded.messages[0].tool_calls[1] == turn.tool_calls[1]