mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
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>
This commit is contained in:
parent
2d192dc245
commit
d62d9e786c
@ -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,
|
||||
|
||||
@ -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.
|
||||
|
||||
|
||||
@ -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. `<system-reminder>`) 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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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)::
|
||||
|
||||
<root>/<kind>/<sha256[:2]>/<sha256> the bytes
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user