From d62d9e786cf3b47e70e65932c0ddb019d66e8038 Mon Sep 17 00:00:00 2001 From: PeaceMaker-best Date: Fri, 2 Oct 2026 19:33:25 +0800 Subject: [PATCH] feat(storage): restore tool outputs from shared blobs (#6159) Checkpoint durable refs for host-externalized tool results and materialize verified bytes before model calls on another Gateway. Refs: #4189 Signed-off-by: PeaceMaker-best <221849497+PeaceMaker-best@users.noreply.github.com> Co-authored-by: PeaceMaker-best <221849497+PeaceMaker-best@users.noreply.github.com> --- backend/app/gateway/services.py | 2 + backend/docs/blob-storage.md | 34 +- .../deerflow/agents/middlewares/AGENTS.md | 2 +- .../middlewares/TOOL_OUTPUT_PERSISTENCE.md | 7 + .../tool_output_budget_middleware.py | 348 ++++++++++++++++-- .../harness/deerflow/storage/AGENTS.md | 16 +- backend/tests/test_gateway_services.py | 35 ++ .../test_tool_output_budget_middleware.py | 321 +++++++++++++++- 8 files changed, 716 insertions(+), 49 deletions(-) diff --git a/backend/app/gateway/services.py b/backend/app/gateway/services.py index 05e69043c6..f1d76bd172 100644 --- a/backend/app/gateway/services.py +++ b/backend/app/gateway/services.py @@ -43,6 +43,7 @@ from deerflow.agents.middlewares.dynamic_context_middleware import _DYNAMIC_CONT from deerflow.agents.middlewares.input_sanitization_middleware import frame_untrusted_text from deerflow.agents.middlewares.message_utils import _SUMMARY_MESSAGE_NAME, is_genuine_user_message from deerflow.agents.middlewares.skill_usage import SKILL_USAGE_KEY, SKILL_USAGES_KEY +from deerflow.agents.middlewares.tool_output_budget_middleware import TOOL_OUTPUT_BLOB_KEY from deerflow.agents.middlewares.tool_receipt import TOOL_RECEIPT_KEY, TOOL_RECEIPT_LEDGER_KEY from deerflow.agents.middlewares.tool_transform_meta import TOOL_TRANSFORMS_KEY from deerflow.agents.middlewares.view_image_middleware import _IMAGE_CONTEXT_MESSAGE_MARKER_KEY @@ -147,6 +148,7 @@ _SERVER_OWNED_MESSAGE_METADATA_KEYS = ( KNOWLEDGE_SCOPE_RUNTIME_KEY, TOOL_RECEIPT_KEY, TOOL_RECEIPT_LEDGER_KEY, + TOOL_OUTPUT_BLOB_KEY, TOOL_TRANSFORMS_KEY, SKILL_USAGE_KEY, SKILL_USAGES_KEY, diff --git a/backend/docs/blob-storage.md b/backend/docs/blob-storage.md index 33c66aace2..6b01371fe7 100644 --- a/backend/docs/blob-storage.md +++ b/backend/docs/blob-storage.md @@ -5,7 +5,7 @@ Resolves the multi-instance half of [#4189](https://github.com/bytedance/deer-fl | Producer | Write | Reads | |---|---|---| | Viewed images | `view_image_tool` → `ViewedImageData.actual_path` (`deerflow/agents/thread_state.py:52`) | `ViewImageMiddleware._read_image_as_data_url`, gateway artifact routes, IM channels, `present_file_tool` | -| Externalized tool results | `ToolOutputBudgetMiddleware._externalize` → virtual path under the thread `outputs` tree | model `read_file` via the thread-data mount | +| Externalized tool results | `ToolOutputBudgetMiddleware` → server-owned `tool-output` ref beside the virtual path | pre-model restoration into the current thread-data mount, then model `read_file` | On a single gateway both are correct. Behind a load balancer, the instance handling the read is frequently not the instance that wrote the file. The blob store replaces **"where on this machine"** with **"which content"**, so every instance that can reach the backing store resolves the same bytes. @@ -33,7 +33,7 @@ Reads **verify the digest**. A content-addressed store that silently returns wro ``` - Writes go to a unique temp file in the destination directory then `os.replace` — atomic within a volume, so a concurrent reader never sees a partial file and two instances racing the same blob converge on identical content. -- Single-put size cap: `_MAX_BLOB_BYTES` (64 MiB) in `local_fs_store.py` rejects larger puts with `BlobWriteError`. This is a backend-level defense-in-depth limit, **not** part of the `BlobStore` contract — other backends set their own policy — but the externalized-tool-results follow-up should assume it and stream or split oversized payloads before they reach `put_bytes`. +- Single-put size cap: `_MAX_BLOB_BYTES` (64 MiB) in `local_fs_store.py` rejects larger puts with `BlobWriteError`. This is a backend-level defense-in-depth limit, **not** part of the `BlobStore` contract — other backends set their own policy. `ToolOutputBudgetMiddleware` applies the same producer cap before `put_bytes`; larger results use an inline fallback capped at 30,000 characters (or a lower configured fallback) instead of creating a reference the default backend cannot satisfy. - The sidecar is GC metadata, **not a read dependency and not a reference count**: losing it must not make content unreadable, and `writer_thread_id` is advisory provenance that cannot stand in for liveness (see the deletion rule above). This backend already delivers multi-instance resolution when `root` points at a shared volume (NFS / EFS / a `ReadWriteMany` PVC). The S3/MinIO backend is a later, optional extra implementing the same contract — it is deliberately not part of this change, because it needs a dependency decision (`[tool.uv.sources]`, optional extra) that belongs in its own PR. @@ -52,14 +52,27 @@ Fail-fast on an unresolvable backend (`ValueError`), mirroring `MemoryConfig.man The factory checks the effective backend and backend configuration on each access. After a `config.yaml` edit, new accesses use the new store; disabling storage makes the optional accessor return `None` and the required accessor raise `BlobNotConfiguredError`. A store already handed to an in-flight caller remains usable and is closed by `reset_blob_store()` rather than during the switch. When changing a storage root in a multi-instance deployment, coordinate the rollout and keep previously written blobs available until existing references have been migrated. -## What this PR deliberately does not do +## Producer migration status -**No producer is migrated.** `blob_storage.enabled` defaults to `false`, no existing call site changed, and the diff is purely additive — a deployment that never sets the key behaves exactly as before. +`blob_storage.enabled` defaults to `false`, so a deployment that never sets the +key behaves exactly as before. -Migration happens in two follow-up PRs, each independently revertible: +The producer migrations are independently revertible: -1. **Viewed images.** `ViewedImageData` gains an optional `blob_ref` (`actual_path` is kept); `view_image_tool` writes the blob when the store is enabled; `ViewImageMiddleware._read_image_as_data_url` resolves blob-first, path-second. The gateway artifact routes and IM channels keep their local-path reads until they can be exercised against a multi-instance deployment. -2. **Externalized tool results.** `ToolOutputBudgetMiddleware._externalize` records a ref alongside the virtual path; the sandbox variant (`_externalize_to_sandbox`, issue #3416) stays as-is — sandbox-resident content is a different failure mode one layer down. +1. **Viewed images — pending.** `ViewedImageData` will gain an optional + `blob_ref` (`actual_path` is kept); `view_image_tool` will write the blob + when the store is enabled; `ViewImageMiddleware._read_image_as_data_url` + will resolve blob-first, path-second. +2. **Externalized tool results — migrated.** Host externalization stores the + exact UTF-8 bytes as `kind="tool-output"` and checkpoints a versioned, + server-owned ref beside the virtual `read_file` path. Before a model call, + another Gateway validates the ref and confined path, reuses a matching + local file, or atomically restores the exact bytes from the shared store. + Blob write failures and outputs over the 64 MiB producer cap fall back to + inline content capped at 30,000 characters (or a lower configured fallback) + rather than advertising a non-durable path. The + non-mounted sandbox variant (`_externalize_to_sandbox`, issue #3416) stays + as-is — sandbox-resident content is a different failure mode one layer down. ## Interaction with checkpoint retention (#5255) @@ -92,9 +105,10 @@ What this means for the two migration PRs: upload the same bytes), so the sweep must check liveness against the checkpoint references before unlinking; `kind` and `writer_thread_id` bound the search. -- **Externalized tool results** — per-thread by construction, so a thread-scoped - sweep is safe while that remains true; `kind` is what makes the "while" - checkable. +- **Externalized tool results** — references are thread-scoped, but identical + bytes from different threads still deduplicate to one blob. A sweep must + therefore confirm that no surviving checkpoint ref names the digest; `kind` + and advisory writer provenance only bound that check. This is also the seam [#5188](https://github.com/bytedance/deer-flow/issues/5188) needs: a thread-scoped blob sweep keyed by thread incarnation, rather than by the reusable thread id. diff --git a/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md b/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md index 78742ed9fb..9e5674d6a0 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md +++ b/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md @@ -45,7 +45,7 @@ strict providers reject. 1. **InputSanitizationMiddleware** - First, so it is the outermost `wrap_model_call` wrapper; every inner middleware (including LLM retries) sees sanitized messages. `additional_kwargs.original_user_content` is server-owned provenance: Gateway strips caller-supplied values for non-internal run requests, trusted IM calls may carry the string they captured before adding transport/file context, and the middleware replaces any non-string value before wrapping. Uploads and sanitization retain first-writer-wins only for validated strings. Caller markers are marked `untrusted_input`, never stripped; scope is every turn. **KnowledgeScopeMiddleware** follows input sanitization: it exposes only Gateway-admitted execution scope, removes scope/display data from model messages, and blocks `knowledge_search` when disabled without reading storage or RAGFlow. -2. **ToolOutputBudgetMiddleware** - Caps model-bound outputs per app config. Externalizes oversized results under thread outputs in `tool_output.storage_subdir` (default `.tool-results`, `TOOL_RESULTS_DIRNAME`), leaving a typed synopsis + `read_file` reference. Host permissions, publication and crash cleanup: [contract](TOOL_OUTPUT_PERSISTENCE.md). These files bypass workspace-change scans and delivery verification. `wrap_model_call` elides successful `write_file` content only in model requests after a later successful same-path `read_file`/`write_file`/`str_replace` (#5328); disk is the reference. Keeps `keep_recent_writes` newest writes. Shared `tool_call_args` helpers pair calls (`pair_tool_call_results`) and rewrite args. Controls: `elide_superseded_writes`, `superseded_write_min_chars`. +2. **ToolOutputBudgetMiddleware** - Caps model-bound outputs per app config. Externalizes oversized results under thread outputs in `tool_output.storage_subdir` (default `.tool-results`, `TOOL_RESULTS_DIRNAME`), leaving a typed synopsis + `read_file` reference. Host permissions, publication and crash cleanup: [contract](TOOL_OUTPUT_PERSISTENCE.md). Blob refs restore across hosts; failures use ≤30,000 inline chars. These files bypass workspace-change scans and delivery verification. `wrap_model_call` elides successful `write_file` content only in model requests after a later successful same-path `read_file`/`write_file`/`str_replace` (#5328); disk is the reference. Keeps `keep_recent_writes` newest writes. Shared `tool_call_args` helpers pair calls (`pair_tool_call_results`) and rewrite args. Controls: `elide_superseded_writes`, `superseded_write_min_chars`. 3. **ToolResultSanitizationMiddleware** - Neutralizes framework/injection tags (e.g. ``) and boundary markers in *remote-content* tool results (`web_fetch`/`web_search`/`image_search`/`web_capture`) so attacker-controlled fetched pages cannot forge trusted framework context. Mirrors `InputSanitizationMiddleware`'s user-input guardrail for the other untrusted-content entry point; sits inner of `ToolOutputBudgetMiddleware` (neutralizes the raw output, then the budget truncates). Local tool output (bash/read_file) is left untouched. Scope is a name-based allowlist for the first-party web tools, plus every MCP-sourced tool via its `deerflow_mcp` metadata tag, so an MCP server naming its fetcher `fetch_url` is still covered Result-rewriting middlewares between the raw callable boundary and the diff --git a/backend/packages/harness/deerflow/agents/middlewares/TOOL_OUTPUT_PERSISTENCE.md b/backend/packages/harness/deerflow/agents/middlewares/TOOL_OUTPUT_PERSISTENCE.md index cf231a7b57..fc20974554 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/TOOL_OUTPUT_PERSISTENCE.md +++ b/backend/packages/harness/deerflow/agents/middlewares/TOOL_OUTPUT_PERSISTENCE.md @@ -12,6 +12,13 @@ creation succeeds: a collision or creation failure must not remove another writer's pending file. An observed `OSError` cleans only this invocation's temp and leaves previously published content intact. +When blob storage is enabled, blob publication follows host publication. If the +blob write fails, that message uses the safety-limited inline fallback and does +not advertise the host path. It also must not unlink the deterministic final +path: a concurrent publisher may already have replaced it and durably +checkpointed that content. Remove published files only as part of inactive +thread-data maintenance. + ## Unclean shutdown SIGKILL, OOM termination, and host failure bypass exception cleanup and can leave diff --git a/backend/packages/harness/deerflow/agents/middlewares/tool_output_budget_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/tool_output_budget_middleware.py index f38179b199..ac5c38a9f4 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/tool_output_budget_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/tool_output_budget_middleware.py @@ -29,10 +29,11 @@ import logging import os import posixpath import shlex +import tempfile import uuid from collections.abc import Awaitable, Callable, Mapping from dataclasses import replace as dc_replace -from typing import TYPE_CHECKING, Any, override +from typing import TYPE_CHECKING, Any, NamedTuple, override from langchain.agents import AgentState from langchain.agents.middleware import AgentMiddleware @@ -51,7 +52,9 @@ from deerflow.community.ragflow.sources import budget_source_artifact from deerflow.config.summarization_config import DEFAULT_SKILL_FILE_READ_TOOL_NAMES from deerflow.config.tool_output_config import ToolOutputConfig from deerflow.constants import DEFAULT_SKILLS_CONTAINER_PATH +from deerflow.sandbox.lease import run_sync_lifecycle_operation from deerflow.sandbox.sandbox_provider import get_sandbox_provider +from deerflow.storage import BlobRef, BlobStoreError, get_blob_store_if_enabled if TYPE_CHECKING: from deerflow.sandbox.sandbox import Sandbox @@ -64,6 +67,25 @@ logger = logging.getLogger(__name__) # ``read_file`` tool can read it back (issue #3416). _VIRTUAL_OUTPUTS_BASE = "/mnt/user-data/outputs" +# Checkpointed capability added only by this middleware. The Gateway strips +# client-supplied copies before they can enter state. +TOOL_OUTPUT_BLOB_KEY = "deerflow_tool_output_blob" +_TOOL_OUTPUT_BLOB_VERSION = 1 +_TOOL_OUTPUT_BLOB_KIND = "tool-output" +_TOOL_OUTPUT_CONTENT_TYPE = "text/plain; charset=utf-8" +# Matches the default local_fs backend object limit. A producer must not create +# a reference that another default-configured Gateway cannot read. +_MAX_TOOL_OUTPUT_BLOB_BYTES = 64 * 1024 * 1024 +# A configured shared store must never fail back to an unbounded model payload, +# even when the operator disabled the ordinary disk-unavailable fallback. +_DURABLE_FAILURE_FALLBACK_MAX_CHARS = 30_000 + + +class _BudgetedContent(NamedTuple): + replacement: str + transform_kind: str + blob_metadata: dict[str, Any] | None = None + def _default_config() -> ToolOutputConfig: return ToolOutputConfig() @@ -221,6 +243,44 @@ def _externalize( return f"{_VIRTUAL_OUTPUTS_BASE}/{storage_subdir}/{filename}" +def _host_path_for_externalized_output( + virtual_path: str, + *, + outputs_path: str, + storage_subdir: str, +) -> str | None: + """Map one middleware-owned virtual path into the current thread root. + + Blob metadata is durable capability data. Even though external callers are + stripped at the Gateway, validate it again before a store read or host write + so a malformed legacy checkpoint cannot escape the tool-results directory. + """ + if not isinstance(virtual_path, str) or not virtual_path: + return None + if not storage_subdir or storage_subdir in {".", ".."} or os.path.isabs(storage_subdir) or "/" in storage_subdir or "\\" in storage_subdir: + return None + virtual_dir = f"{_VIRTUAL_OUTPUTS_BASE}/{storage_subdir}" + if posixpath.normpath(virtual_path) != virtual_path or posixpath.dirname(virtual_path) != virtual_dir: + return None + filename = posixpath.basename(virtual_path) + if not filename or filename in {".", ".."}: + return None + + outputs_root = os.path.realpath(os.path.abspath(outputs_path)) + storage_root = os.path.realpath(os.path.join(outputs_root, storage_subdir)) + candidate = os.path.realpath(os.path.join(storage_root, filename)) + try: + if os.path.commonpath([outputs_root, storage_root]) != outputs_root: + return None + if os.path.commonpath([storage_root, candidate]) != storage_root: + return None + except ValueError: + return None + if os.path.dirname(candidate) != storage_root: + return None + return candidate + + def _externalize_to_sandbox( content: str, *, @@ -358,6 +418,20 @@ def _resolve_outputs_path(request: ToolCallRequest) -> str | None: return outputs_path if isinstance(outputs_path, str) else None +def _resolve_thread_id(request: ToolCallRequest) -> str | None: + """Resolve advisory blob provenance from runtime context or config.""" + runtime = getattr(request, "runtime", None) + context = getattr(runtime, "context", None) + if isinstance(context, Mapping): + thread_id = context.get("thread_id") + if isinstance(thread_id, str) and thread_id: + return thread_id + runtime_config = getattr(runtime, "config", None) + configurable = runtime_config.get("configurable") if isinstance(runtime_config, Mapping) else None + thread_id = configurable.get("thread_id") if isinstance(configurable, Mapping) else None + return thread_id if isinstance(thread_id, str) and thread_id else None + + def _resolve_sandbox(request: ToolCallRequest) -> Sandbox | None: """Resolve the active sandbox for the current tool call, or ``None``. @@ -395,11 +469,12 @@ def _budget_content( outputs_path: str | None, config: ToolOutputConfig, sandbox: Sandbox | None = None, -) -> tuple[str, str] | None: + thread_id: str | None = None, +) -> _BudgetedContent | None: """Apply budget to *content* and name the applied transform. - Returns ``(replacement, transform_kind)`` — ``"externalized"`` or - ``"truncated"`` — or ``None`` if no change was needed. + Returns the replacement, transform kind, and optional durable blob + metadata, or ``None`` if no change was needed. """ threshold = config.tool_overrides.get(tool_name, config.externalize_min_chars) if threshold <= 0 and config.fallback_max_chars <= 0: @@ -409,6 +484,9 @@ def _budget_content( if threshold > 0 and len(content) > threshold: virtual_path: str | None = None + blob_metadata: dict[str, Any] | None = None + host_outputs_path: str | None = None + durable_fallback_required = False # Decide persistence target based on what's available, without touching # the sandbox provider unless a sandbox was actually resolved for this # call. This keeps the legacy host-disk path provider-free, so callers @@ -426,13 +504,7 @@ def _budget_content( # equivalent. Preserve the original behavior to avoid extra # sandbox round-trips. if outputs_path: - virtual_path = _externalize( - content, - tool_name=tool_name, - tool_call_id=tool_call_id, - outputs_path=outputs_path, - storage_subdir=config.storage_subdir, - ) + host_outputs_path = outputs_path else: virtual_path = _externalize_to_sandbox( content, @@ -444,13 +516,53 @@ def _budget_content( elif outputs_path: # No sandbox in this call (legacy / non-sandbox tools): write to # host outputs path directly, no provider needed. - virtual_path = _externalize( - content, - tool_name=tool_name, - tool_call_id=tool_call_id, - outputs_path=outputs_path, - storage_subdir=config.storage_subdir, - ) + host_outputs_path = outputs_path + + if host_outputs_path is not None: + blob_store = get_blob_store_if_enabled() + blob_bytes = content.encode("utf-8") if blob_store is not None else None + if blob_bytes is not None and len(blob_bytes) > _MAX_TOOL_OUTPUT_BLOB_BYTES: + durable_fallback_required = True + logger.warning( + "Tool output is too large for durable blob externalization: %d bytes > %d", + len(blob_bytes), + _MAX_TOOL_OUTPUT_BLOB_BYTES, + ) + else: + virtual_path = _externalize( + content, + tool_name=tool_name, + tool_call_id=tool_call_id, + outputs_path=host_outputs_path, + storage_subdir=config.storage_subdir, + ) + if virtual_path is not None and blob_store is not None and blob_bytes is not None: + try: + ref = blob_store.put_bytes( + blob_bytes, + kind=_TOOL_OUTPUT_BLOB_KIND, + content_type=_TOOL_OUTPUT_CONTENT_TYPE, + thread_id=thread_id, + ) + except BlobStoreError: + durable_fallback_required = True + logger.warning( + "Failed to persist externalized %s output in shared blob storage", + tool_name, + exc_info=True, + ) + # A concurrent publisher may already have replaced the + # deterministic host path. Do not unlink a file that + # may belong to another successful blob checkpoint. + virtual_path = None + else: + blob_metadata = { + "version": _TOOL_OUTPUT_BLOB_VERSION, + "ref": ref.model_dump(mode="json", exclude_none=True), + "virtual_path": virtual_path, + "storage_subdir": config.storage_subdir, + "encoding": "utf-8", + } if virtual_path is not None: logger.info( "Externalized %s output (%d chars) to %s", @@ -458,7 +570,7 @@ def _budget_content( len(content), virtual_path, ) - return ( + return _BudgetedContent( _build_preview( content, tool_name=tool_name, @@ -467,20 +579,29 @@ def _budget_content( tail_chars=config.preview_tail_chars, ), "externalized", + blob_metadata, ) - if config.fallback_max_chars > 0 and len(content) > config.fallback_max_chars: + if durable_fallback_required: + configured_limit = config.fallback_max_chars if config.fallback_max_chars > 0 else _DURABLE_FAILURE_FALLBACK_MAX_CHARS + fallback_max_chars = min(configured_limit, _DURABLE_FAILURE_FALLBACK_MAX_CHARS) + else: + fallback_max_chars = config.fallback_max_chars + else: + fallback_max_chars = config.fallback_max_chars + + if fallback_max_chars > 0 and len(content) > fallback_max_chars: logger.warning( "Fallback-truncating %s output: %d chars → %d max", tool_name, len(content), - config.fallback_max_chars, + fallback_max_chars, ) - return ( + return _BudgetedContent( _build_fallback( content, tool_name=tool_name, - max_chars=config.fallback_max_chars, + max_chars=fallback_max_chars, head_chars=config.fallback_head_chars, tail_chars=config.fallback_tail_chars, ), @@ -500,6 +621,7 @@ def _patch_tool_message( config: ToolOutputConfig, outputs_path: str | None, sandbox: Sandbox | None = None, + thread_id: str | None = None, ) -> ToolMessage: """Apply budget to a single ToolMessage. Returns the original if unchanged.""" tool_name = msg.name or "unknown" @@ -517,6 +639,7 @@ def _patch_tool_message( outputs_path=outputs_path, config=config, sandbox=sandbox, + thread_id=thread_id, ) update: dict[str, Any] = {} trigger = _effective_trigger(tool_name, config) @@ -527,14 +650,18 @@ def _patch_tool_message( replacement, update["artifact"] = citation_result transform_kind = "truncated" elif budgeted is not None: - replacement, transform_kind = budgeted + replacement = budgeted.replacement + transform_kind = budgeted.transform_kind else: return msg update["content"] = replacement if getattr(msg, "response_metadata", None): update["response_metadata"] = dict(msg.response_metadata) new_kwargs = dict(getattr(msg, "additional_kwargs", None) or {}) - if citation_result is not None and budgeted is not None and budgeted[1] == "externalized": + new_kwargs.pop(TOOL_OUTPUT_BLOB_KEY, None) + if budgeted is not None and budgeted.blob_metadata is not None: + new_kwargs[TOOL_OUTPUT_BLOB_KEY] = budgeted.blob_metadata + if citation_result is not None and budgeted is not None and budgeted.transform_kind == "externalized": append_tool_transform(new_kwargs, "externalized", by="ToolOutputBudgetMiddleware") append_tool_transform(new_kwargs, transform_kind, by="ToolOutputBudgetMiddleware") update["additional_kwargs"] = new_kwargs @@ -587,10 +714,11 @@ def _patch_result( config: ToolOutputConfig, outputs_path: str | None, sandbox: Sandbox | None = None, + thread_id: str | None = None, ) -> ToolMessage | Command: """Apply budget to a tool call result (ToolMessage or Command).""" if isinstance(result, ToolMessage): - return _patch_tool_message(result, config, outputs_path, sandbox) + return _patch_tool_message(result, config, outputs_path, sandbox, thread_id) update = getattr(result, "update", None) if not isinstance(update, dict): @@ -598,7 +726,7 @@ def _patch_result( messages = update.get("messages") if isinstance(messages, ToolMessage): - patched = _patch_tool_message(messages, config, outputs_path, sandbox) + patched = _patch_tool_message(messages, config, outputs_path, sandbox, thread_id) return result if patched is messages else dc_replace(result, update={**update, "messages": patched}) if not isinstance(messages, (list, tuple)): return result @@ -607,7 +735,7 @@ def _patch_result( changed = False for msg in messages: if isinstance(msg, ToolMessage): - patched = _patch_tool_message(msg, config, outputs_path, sandbox) + patched = _patch_tool_message(msg, config, outputs_path, sandbox, thread_id) if patched is not msg: changed = True new_messages.append(patched) @@ -701,6 +829,135 @@ def _patch_model_messages(messages: list[Any], config: ToolOutputConfig) -> list return updated if changed else None +def _model_outputs_path(request: ModelRequest) -> str | None: + state = getattr(request, "state", None) + if not isinstance(state, Mapping): + return None + thread_data = state.get("thread_data") + if not isinstance(thread_data, Mapping): + return None + outputs_path = thread_data.get("outputs_path") + return outputs_path if isinstance(outputs_path, str) and outputs_path else None + + +def _has_tool_output_blob_refs(messages: list[Any]) -> bool: + """Cheap pre-scan used to keep normal async model calls on the event loop.""" + return any(isinstance(message, ToolMessage) and isinstance(getattr(message, "additional_kwargs", None), dict) and TOOL_OUTPUT_BLOB_KEY in message.additional_kwargs for message in messages) + + +def _write_restored_output(filepath: str, data: bytes) -> bool: + """Atomically publish restored bytes without exposing a partial file.""" + directory = os.path.dirname(filepath) + tmp_path: str | None = None + try: + os.makedirs(directory, exist_ok=True) + fd, tmp_path = tempfile.mkstemp(prefix=f".{os.path.basename(filepath)}.", suffix=".tmp", dir=directory) + with os.fdopen(fd, "wb") as handle: + handle.write(data) + handle.flush() + os.fsync(handle.fileno()) + os.replace(tmp_path, filepath) + return True + except OSError: + logger.warning("Failed to restore externalized tool output on this Gateway", exc_info=True) + return False + finally: + if tmp_path is not None: + try: + os.unlink(tmp_path) + except FileNotFoundError: + pass + except OSError: + logger.warning("Failed to clean up tool-output restore temp file", exc_info=True) + + +def _discard_mismatched_output(filepath: str) -> None: + """Do not leave known-wrong bytes behind a trusted read_file reference.""" + try: + os.unlink(filepath) + except FileNotFoundError: + pass + except OSError: + logger.warning("Failed to remove mismatched externalized tool output", exc_info=True) + + +def _restore_tool_output_blobs(messages: list[Any], *, outputs_path: str) -> None: + """Materialize durable host-externalized results under this Gateway. + + Existing digest-matching files are reused without touching the blob store. + Invalid metadata and unavailable blobs are ignored so a transient storage + outage does not prevent the model from using the inline preview. + """ + store_resolved = False + blob_store = None + for message in messages: + if not isinstance(message, ToolMessage): + continue + payload = (message.additional_kwargs or {}).get(TOOL_OUTPUT_BLOB_KEY) + if not isinstance(payload, Mapping): + continue + if payload.get("version") != _TOOL_OUTPUT_BLOB_VERSION or payload.get("encoding") != "utf-8": + continue + ref_data = payload.get("ref") + virtual_path = payload.get("virtual_path") + storage_subdir = payload.get("storage_subdir") + if not isinstance(ref_data, Mapping) or not isinstance(virtual_path, str) or not isinstance(storage_subdir, str): + continue + try: + ref = BlobRef.model_validate(dict(ref_data)) + except (TypeError, ValueError): + continue + if ref.kind != _TOOL_OUTPUT_BLOB_KIND or ref.content_type != _TOOL_OUTPUT_CONTENT_TYPE: + continue + filepath = _host_path_for_externalized_output( + virtual_path, + outputs_path=outputs_path, + storage_subdir=storage_subdir, + ) + if filepath is None: + continue + + local_mismatch = False + try: + if os.stat(filepath).st_size != ref.size: + local_mismatch = True + else: + with open(filepath, "rb") as handle: + if ref.matches(handle.read()): + continue + local_mismatch = True + except FileNotFoundError: + pass + except OSError: + logger.warning("Failed to inspect local externalized tool output", exc_info=True) + + if not store_resolved: + blob_store = get_blob_store_if_enabled() + store_resolved = True + if blob_store is None: + if local_mismatch: + _discard_mismatched_output(filepath) + continue + try: + data = blob_store.get_bytes(ref) + except BlobStoreError: + logger.warning( + "Failed to resolve externalized tool output blob %s", + ref.sha256[:12], + exc_info=True, + ) + if local_mismatch: + _discard_mismatched_output(filepath) + continue + if not ref.matches(data): + logger.warning("Blob backend returned mismatched tool output %s", ref.sha256[:12]) + if local_mismatch: + _discard_mismatched_output(filepath) + continue + if not _write_restored_output(filepath, data) and local_mismatch: + _discard_mismatched_output(filepath) + + # --------------------------------------------------------------------------- # Superseded write payload elision (issue #5328, step 2) # --------------------------------------------------------------------------- @@ -856,7 +1113,13 @@ class ToolOutputBudgetMiddleware(AgentMiddleware[AgentState]): if self._config.enabled and _needs_budget(result, self._config): outputs_path = _resolve_outputs_path(request) sandbox = _resolve_sandbox(request) - result = _patch_result(result, self._config, outputs_path, sandbox) + result = _patch_result( + result, + self._config, + outputs_path, + sandbox, + _resolve_thread_id(request), + ) return _record_visible_skill_usage(result, request, skill_read_tool_names=self._skill_read_tool_names, skills_root=self._skills_root) @override @@ -872,7 +1135,14 @@ class ToolOutputBudgetMiddleware(AgentMiddleware[AgentState]): # in-memory sandbox registry, so it is safe to call on the event # loop. The actual sandbox I/O happens in the worker thread. sandbox = _resolve_sandbox(request) - result = await asyncio.to_thread(_patch_result, result, self._config, outputs_path, sandbox) + result = await asyncio.to_thread( + _patch_result, + result, + self._config, + outputs_path, + sandbox, + _resolve_thread_id(request), + ) return _record_visible_skill_usage(result, request, skill_read_tool_names=self._skill_read_tool_names, skills_root=self._skills_root) # -- model call hooks (historical context budgeting) ------------------- @@ -891,16 +1161,26 @@ class ToolOutputBudgetMiddleware(AgentMiddleware[AgentState]): request: ModelRequest, handler: Callable[[ModelRequest], Awaitable[ModelResponse]], ) -> ModelCallResult: - # Pure in-memory rewrite: no sandbox or file I/O, so it stays on the loop. - return await handler(self._budget_model_request(request)) + messages = getattr(request, "messages", None) + outputs_path = _model_outputs_path(request) + if isinstance(messages, list) and outputs_path is not None and _has_tool_output_blob_refs(messages): + # A blob read and atomic host write may outlive cancellation. Drain + # the worker before the request boundary releases its resources. + prepared = await run_sync_lifecycle_operation(self._budget_model_request, request) + else: + prepared = self._budget_model_request(request) + return await handler(prepared) def _budget_model_request(self, request: ModelRequest) -> ModelRequest: """Truncate oversized historical tool output and elide superseded write payloads in the request copy only.""" - if not self._config.enabled: - return request original = getattr(request, "messages", None) if not isinstance(original, list): return request + outputs_path = _model_outputs_path(request) + if outputs_path is not None and _has_tool_output_blob_refs(original): + _restore_tool_output_blobs(original, outputs_path=outputs_path) + if not self._config.enabled: + return request messages = original patched = _patch_model_messages(messages, self._config) if patched is not None: diff --git a/backend/packages/harness/deerflow/storage/AGENTS.md b/backend/packages/harness/deerflow/storage/AGENTS.md index 5476fc7759..f19d5b6b23 100644 --- a/backend/packages/harness/deerflow/storage/AGENTS.md +++ b/backend/packages/harness/deerflow/storage/AGENTS.md @@ -8,9 +8,9 @@ a reader can verify what it got back. Contract: `storage/contract.py` Operator-facing behaviour and the producer migration plan: `docs/blob-storage.md`. Factory: `storage/manager.py`. -**Disabled by default**: `blob_storage.enabled` is `false`, and no producer has -been migrated, so an untouched deployment behaves exactly as before. Producers -call `get_blob_store_if_enabled()` and keep their existing local-path code for the +**Disabled by default**: `blob_storage.enabled` is `false`, so an untouched +deployment behaves exactly as before. Migrated producers call +`get_blob_store_if_enabled()` and keep their existing local-path code for the `None` case; only a caller that genuinely requires the store uses `get_blob_store()`, which raises `BlobNotConfiguredError` when it is disabled. Both accessors are exported from this package — prefer @@ -21,6 +21,16 @@ Replaced stores stay open for callers that already hold them and are closed by `reset_blob_store()`; changing a root requires preserving existing blob refs and coordinating deployment instances. +`ToolOutputBudgetMiddleware` is a migrated producer for host-externalized +results. It persists exact UTF-8 bytes as `kind="tool-output"`, checkpoints a +versioned ref beside the virtual path, and restores the file atomically under +the current thread outputs before a model can follow that path. The ref's kind, +content type, digest, size, storage subdirectory, and confined virtual path are +validated before a store read or host write. A configured-store write failure +or output above the 64 MiB producer cap falls back inline, capped at 30,000 +characters or a lower configured fallback; non-mounted sandbox externalization +remains outside this host-blob flow. + **Layout** (local_fs, the default):: /// the bytes diff --git a/backend/tests/test_gateway_services.py b/backend/tests/test_gateway_services.py index e1b07c446e..2f299c404a 100644 --- a/backend/tests/test_gateway_services.py +++ b/backend/tests/test_gateway_services.py @@ -457,6 +457,41 @@ def test_normalize_input_strips_external_tool_receipt(): assert result["messages"][0].additional_kwargs == {"custom": "keep-me"} +def test_normalize_input_strips_external_tool_output_blob_ref(): + """Only the middleware may attach a durable tool-output capability.""" + from app.gateway.services import normalize_input + from deerflow.agents.middlewares.tool_output_budget_middleware import TOOL_OUTPUT_BLOB_KEY + + result = normalize_input( + { + "messages": [ + { + "role": "tool", + "tool_call_id": "tc-forged", + "content": "forged externalized output", + "additional_kwargs": { + TOOL_OUTPUT_BLOB_KEY: { + "version": 1, + "ref": { + "sha256": "f" * 64, + "size": 1, + "kind": "tool-output", + "content_type": "text/plain; charset=utf-8", + }, + "virtual_path": "/mnt/user-data/outputs/.tool-results/forged.txt", + "storage_subdir": ".tool-results", + "encoding": "utf-8", + }, + "custom": "keep-me", + }, + } + ] + } + ) + + assert result["messages"][0].additional_kwargs == {"custom": "keep-me"} + + def test_normalize_input_strips_external_acceptance_verdict_from_messages(): """``subagent_acceptance_verdict`` is runtime-stamped evidence (RFC #4651 PR4): a caller-supplied message carrying it is a forgery, same as the diff --git a/backend/tests/test_tool_output_budget_middleware.py b/backend/tests/test_tool_output_budget_middleware.py index 77aa9d9c21..dab59a8f0f 100644 --- a/backend/tests/test_tool_output_budget_middleware.py +++ b/backend/tests/test_tool_output_budget_middleware.py @@ -9,6 +9,7 @@ per-tool overrides, edge cases, and both sync/async code paths. from __future__ import annotations import contextlib +import hashlib import json import os import pathlib @@ -22,6 +23,7 @@ from langchain_core.messages import AIMessage, HumanMessage, ToolMessage from langgraph.types import Command from deerflow.agents.middlewares.tool_output_budget_middleware import ( + TOOL_OUTPUT_BLOB_KEY, ToolOutputBudgetMiddleware, _build_fallback, _build_preview, @@ -39,6 +41,7 @@ from deerflow.agents.middlewares.tool_output_synopsis import build_tool_output_s 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 @@ -62,7 +65,7 @@ def _lines_then_long_line(total: int, newline_ratio: float = 0.6) -> str: 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) + runtime = SimpleNamespace(state=state, context={"thread_id": "thread-1"}, config={}) return SimpleNamespace( tool_call={"name": tool_name, "id": tool_call_id}, runtime=runtime, @@ -673,6 +676,7 @@ class TestWrapToolCallExternalize: 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") @@ -682,6 +686,153 @@ class TestWrapToolCallExternalize: 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_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) @@ -1114,6 +1265,166 @@ class TestWrapModelCall: 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 # =========================================================================== @@ -1478,11 +1789,18 @@ class TestBudgetContentSandboxDispatch: 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, @@ -1499,6 +1817,7 @@ class TestBudgetContentSandboxDispatch: 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