mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-11 14:38:38 +00:00
feat(agents): elide blocked write payloads from model-bound requests (#5329)
* feat(agents): elide blocked write payloads from model-bound requests A write_file / str_replace call rejected by the read-before-write gate never runs, yet its payload (up to 80 KB for a non-append write, unbounded for append) stayed verbatim in every later model request: nothing in the chain rewrites AIMessage tool-call arguments, and ToolOutputBudgetMiddleware only budgets ToolMessage output. The gate demands a re-read plus a fresh call, so the model re-emits the content anyway and the original is pure dead weight. - ReadBeforeWriteMiddleware stamps `deerflow_write_block` on the blocked ToolMessage and, in wrap_model_call, replaces the paired call's payload fields (content / old_str / new_str) with a short deterministic placeholder in the model-bound request only. state["messages"], receipts, loop detection, and the run journal keep the original arguments; nothing is externalized to disk, since a file reference to content the model must re-derive after reading the target would only invite bypassing the gate. - New `tool_call_args` helper rewrites every provider surface together (structured tool_calls, raw additional_kwargs.tool_calls, tool_use content blocks, tool_call_chunks) so strict providers never see them disagree; the gate only supplies the policy (which calls, what placeholder). - `read_before_write.elide_blocked_payloads` (default on) and `read_before_write.elide_min_chars` (default 2000) configure it; the runtime builder passes the config through and the middleware declares it via release_policy_parameters. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * docs(agents): condense middleware guide entry 11 to fit the guidance budget The agent-guidance CI check failed: the effective AGENTS.md chain for agents/middlewares was 99673 bytes against a 98304-byte hard limit. The chain already sat at 98459 on main, so the ReadBeforeWrite entry could not grow. Rewrite entry 11 so it states the same facts (gate, lock scope, fail-open, authorization scope, blocked-payload elision, shared tool_call_args helper) in 1229 bytes instead of 2640; the chain is now 98262 bytes. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * fix(config): bump config_version for the read_before_write elision keys Review follow-ups on #5329: - `read_before_write.elide_blocked_payloads` / `elide_min_chars` are new user-settable YAML keys, i.e. a config schema change, so bump `config_version` 40 -> 41 in config.example.yaml; without it an existing config.yaml gets no outdated-config warning and `make config-upgrade` has nothing to signal. - Say in the `elide_min_chars` description (and the example comment) that the threshold and the placeholder's size figure are Python character counts, not tokens: the same value spans roughly 3-4x in real context cost between ASCII and CJK text. - The builder wiring test now asserts only the wired `elide_min_chars` value instead of the whole `ReadBeforeWriteConfig` dump, so future knobs do not have to edit an unrelated test. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * chore(helm): bump chart config_version to 41 validate-chart's config_version drift check failed after config.example.yaml moved to 41 in ef9ee267. Bare bump of the chart's embedded `config:` block and the README example; the chart does not mirror the read_before_write section, so no field changes are needed. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * fix(agents): rewrite Responses and v1 content-block arguments too Review finding on #5329 (P2): the content rewriter only handled Anthropic `tool_use` blocks. With `use_responses_api=true` and `output_version='responses/v1'`, AIMessage.content carries `function_call` blocks whose `arguments` still hold the full write payload, and langchain_openai's Responses input builder emits that block instead of the rewritten structured call whose `call_id` it already carries. Standard `v1` `tool_call` blocks likewise keep `extras.arguments`, which the v1->Responses translator prefers over the structured args. So the blocked payload was still sent on every later Responses API request. `tool_call_args` now rewrites every content dialect that carries its own copy of the arguments: Anthropic `tool_use` (input, drop partial_json), Responses `function_call` (arguments, matched by call_id, `fc_...` item id and status preserved), and v1 `tool_call` / `tool_call_chunk` (args plus `extras.arguments`). Tests assert against the real adapter serializers: `_construct_responses_api_input` for responses/v1, v1, and v0 messages, `_convert_message_to_dict` for chat completions, and Anthropic `_format_messages` for native and v1 content, plus an end-to-end probe through the gate's wrap_model_call. * fix(agents): pair blocked writes per call occurrence and defeat Responses chaining Two review findings on #5329: - Tool-call ids may repeat across assistant turns (DanglingToolCallMiddleware pairs them with per-id queues). The gate matched blocked results against a history-wide id set, so a successful write sharing an id with a later (or earlier) blocked one also lost its payload and was labelled as blocked. `_blocked_call_occurrences` now pairs ToolMessages with call occurrences the same FIFO-per-id way and the selector keys on (message, call id). - With `use_previous_response_id`, the OpenAI adapter sends only the messages after the last AIMessage carrying a `resp_` response id and lets the server rebuild the rest from its stored copy, which still holds the original arguments and cannot be edited; every later response chains back to it. `rewrite_messages_tool_call_args` now drops every `resp_` id from the model-bound copy whenever it rewrote anything, so the adapter replays the full rewritten history (the `use_previous_response_id=False` request shape). OpenAI bills chained input tokens as input either way, so replay costs no more; the state keeps its ids. Tests cover success-before-block and block-before-success histories through the Chat Completions serializer, and chaining through `ChatOpenAI._get_request_payload` with `use_previous_response_id=True`: unrewritten history chains and omits the call, rewritten history is replayed with the placeholder and no `previous_response_id`. --------- Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
Nan Gao
GitHub
parent
f09e824460
commit
3f0b6ecc81
10
CHANGELOG.md
10
CHANGELOG.md
@ -214,6 +214,16 @@ This section accumulates work toward the **2.1.0** milestone
|
||||
subagent's graph state, making `list_uploaded_files` eligible for normal
|
||||
tool-policy filtering (durable `batch_task` workers keep it disabled).
|
||||
([#5170])
|
||||
- **agents:** The read-before-write gate now elides the dead payload of a
|
||||
blocked `write_file` / `str_replace` call (`content`, `old_str`, `new_str`)
|
||||
from model-bound requests. A blocked call never ran and must be re-issued
|
||||
after a re-read, so the original arguments only cost context; stored
|
||||
history, receipts, and the run journal keep them. Blocked results are
|
||||
paired with call occurrences (tool-call ids may repeat across turns), and
|
||||
a request whose history was rewritten drops OpenAI `resp_` response ids so
|
||||
`use_previous_response_id` chaining cannot resume the original server-side
|
||||
history. Controlled by `read_before_write.elide_blocked_payloads` (default
|
||||
on) and `read_before_write.elide_min_chars` (default 2000).
|
||||
|
||||
#### Memory
|
||||
|
||||
|
||||
@ -65,7 +65,7 @@ alongside every behaviour-affecting field.
|
||||
forgeries). Consumers pop it; the publisher and the consumer share only that
|
||||
contract module.
|
||||
10. **SandboxAuditMiddleware** - Audits sandboxed shell/file operations before tool execution; command classification is **defense-in-depth and audit, not a security boundary** (the sandbox is the isolation boundary). Command substitution is judged by *position*, not the presence of `$(`: **command position** (`$(curl url)`, `` `curl url` ``, the word after `|`/`&&`/`;`, an `eval`/`source` argument) executes fetched content and is blocked; **value position** (`x=$(curl url)`, `echo $(curl url)`, an argument, a `for` word list) only captures output and passes (#4611). So `_HIGH_RISK_COMMAND_POSITION_PATTERNS` is matched anchored against each sub-command from `_split_compound_command(split_pipes=True)`, never the whole string; pipe-spanning rules (`| sh`, `base64 -d | ...`) still use `_classify_command`'s whole-command Pass 1. `_COMMAND_POSITION_PREFIX` extends the anchor over leading assignments and exec wrappers (`FOO=1 $(curl url)`, `env`/`command`/`builtin`/`exec`/`nohup`/`time`/`sudo`/`doas`); its assignment branch requires whitespace before the substitution, which keeps `x=$(curl url)` in value position. Two contexts are deliberately **position-blind** (matched whole-command in Pass 1, since they execute their input anywhere, e.g. `xargs sh -c "$(curl url)"`): an `eval`/`source` argument, and an interpreter **code-string flag** — `-c` (shells, `python`), `-e` (`perl`/`ruby`/`node`), `-p` (`perl`/`node`), `-r` (`php`) — plus the here-string (`<<<`) reaching the same place via stdin. All three substitution spellings (`$(`, `<(`, `` ` ``) share one `_RISKY_SUBSTITUTION` opener. An unquoted newline splits like `;` (else `echo hi\n$(curl url)` evades the anchored rules). Heredoc bodies are data: `_split_compound_command` records headers (`<<EOF`, `<<-EOF`, `<<'EOF'`) and consumes their bodies verbatim, so a body line starting `$(curl url)` isn't promoted to a command position; `<<<` (here-string, needs look-ahead + look-behind) and a `<<` inside `$(( ))`/`(( ))` (bit shift, arithmetic depth tracked with the quote flags) must not open one. This is a heuristic, not shell parsing — an unterminated body consumes the rest of the string, an unclosed `((` only disables heredoc detection, and the failure direction is always toward *more* command positions, not fewer. Known gaps: process substitution outside `eval`/`source` (`. <(curl u)`) is undetected, and two-step forms (`x=$(curl u); eval "$x"`) need dataflow analysis. No config gate — appended unconditionally in `_build_runtime_middlewares`, for both lead and subagents.
|
||||
11. **ReadBeforeWriteMiddleware** - *(optional, if `read_before_write.enabled`, default on)* Outermost write gate (issue #3857): `read_file` stamps a content hash onto its ToolMessage; `write_file` (append/overwrite-existing) and `str_replace` are blocked unless the newest mark for that path matches the file's current hash. Sits outside ToolProgressMiddleware and ToolErrorHandlingMiddleware so a blocked write returns immediately without consuming a ToolProgress slot. Blocked results call `normalize_tool_result` directly to stamp `deerflow_tool_meta` (`recoverable_by_model=True`) before returning, keeping the result well-formed for any outer consumer. Marks live on messages, so summarization dropping the read result invalidates the gate automatically; writes never refresh marks, forcing a re-read between consecutive edits. Gate check + tool execution are serialized per (thread, path) so same-turn parallel writes cannot reuse one stale mark; on sandboxes whose `read_file` reports failures as `"Error: ..."` strings instead of raising (AIO/E2B), uninspectable targets fail open (creation proceeds, no mark stamped). The middleware also owns the sandbox authorization scope for these composed calls: pre-write inspection, the tool body, and post-read hashing share one sync/async provider decision, while `SandboxAuthorizationError` bypasses the generic inspection fail-open paths and becomes an error ToolMessage.
|
||||
11. **ReadBeforeWriteMiddleware** - *(optional, `read_before_write.enabled`, default on)* Outermost write gate (#3857): `read_file` stamps a content hash on its ToolMessage; `write_file` (existing file, incl. append) and `str_replace` are blocked unless the newest mark for the path matches its current hash. Sits outside ToolProgress/ToolErrorHandling (a block consumes no ToolProgress slot); blocked results self-stamp `deerflow_tool_meta` and carry `deerflow_write_block` (`{path, tool}`). Marks live on messages, so summarization dropping the read invalidates the gate; writes never refresh marks. Gate check + execution are serialized per (thread, path); `"Error: ..."`-string sandboxes (AIO/E2B) fail open. It owns the composed call's sandbox authorization scope; `SandboxAuthorizationError` becomes an error ToolMessage. Its `wrap_model_call` swaps blocked calls' dead payload (`content`, `old_str`/`new_str`) for a deterministic placeholder in the model-bound request only (`elide_blocked_payloads`, `elide_min_chars`); state, receipts, and the journal keep the originals. Policy stays in the gate; the shared `tool_call_args` helper rewrites every arg surface together and every model-bound arg rewrite must use it.
|
||||
12. **ToolProgressMiddleware** - *(optional, if `tool_progress.enabled`)* State-machine-based stagnation guard (RFC #3177). Outer wrapper around ToolErrorHandlingMiddleware so its `wrap_tool_call` receives results already stamped with `deerflow_tool_meta`. Tracks per-(thread, tool) consecutive "no-new-info" calls across three error categories: (a) `recoverable_by_model=True` (no_results, not_found, permission, Jaccard-duplicate success): ACTIVE → WARNED (terminal — hint re-injected on each subsequent problem); (b) `recoverable_by_model=False, action≠stop` (rate_limited, transient): ACTIVE → WARNED → BLOCKED after `warn_escalation_count` more problems; (c) `recoverable_by_model=False, action=stop` (auth, config, internal): immediately BLOCKED on first occurrence. **Division of labor with LoopDetectionMiddleware:** ToolProgressMiddleware is a result-quality guard — fires after tool execution and blocks specific tools that stop producing new information; LoopDetectionMiddleware is a call-pattern guard — fires after the model responds and hard-stops the whole turn when the model repeatedly issues identical tool_calls. Both can inject HumanMessage hints in the same model call without conflict; neither reads the other's internal state.
|
||||
13. **ToolReceiptMiddleware + ToolErrorHandlingMiddleware** - `ToolReceiptMiddleware` is *(optional, if `verification.receipts_enabled`, default on)*. It is the **outermost `wrap_tool_call` layer** — registered ahead of entries 9-12 — because Guardrail/SandboxAudit/ReadBeforeWrite/ToolProgress can short-circuit a call with their own ToolMessage (and SandboxAudit rebuilds medium-risk results); an inner receipt layer would silently gap the ledger on those results (ordering constraints in `deerflow.extensions.ordering`). Normal results still carry the `deerflow_tool_meta` status ToolErrorHandlingMiddleware stamps on the inner return path; short-circuit messages self-stamp meta or fall back to `message.status`. It stamps deterministic provenance (tool name, status, args/output hashes, byte count, timestamp) onto direct `ToolMessage` results and every matching `ToolMessage` carried in `Command.update.messages`, including delegated `task`, `present_file`, `view_image`, and `tool_search` results; before model calls it derives a hidden receipt ledger (display ids r1..rN) from message state, and when the 2,000-character budget is exceeded the newest receipts are retained in chronological order with their original ids plus an older-receipts omission marker. Rendering returns both the text and its retained receipt subset; every response that received a ledger carries only that exact server-owned subset, never omitted receipts. Snapshot validation accepts a strictly consecutive positive original-id range (for example `r24`–`r30`) rather than requiring `r1`, so subagent terminal citation verification resolves ids against evidence present in the citing turn even when later summarization drops and renumbers tool messages. Model-generated citation IDs are digit-bounded before integer conversion; oversized IDs are ignored as malformed input rather than raising through task write-back. Citation parsing deduplicates exact `(id, anchor)` pairs, not IDs alone, so repeated identical references stay compact while every distinct anchor claim is verified. Gateway strips delegated receipts/verdicts from external messages. `ToolErrorHandlingMiddleware` receives `AppConfig`, converts tool exceptions into error `ToolMessage`s so the run can continue instead of aborting, stamps every result with `deerflow_tool_meta` (status / error_type / recoverable_by_model / recommended_next_action / source) via `tool_result_meta.normalize_tool_result`, stamps structured metadata for task exception wrappers, and stamps skill-read metadata for downstream durable-context capture. Task tool result text is generated from the same status/result/error inputs as the structured metadata so callers do not hand-write a second protocol string.
|
||||
|
||||
|
||||
@ -23,6 +23,15 @@ Design invariants:
|
||||
binary content, or sandboxes like AIO/E2B that report read failures as
|
||||
``"Error: ..."`` strings instead of raising), it lets the tool run and
|
||||
produce its own error.
|
||||
- Blocked payloads are dead weight: the call never ran, and the gate demands
|
||||
a re-read plus a fresh call, so the model re-emits the content anyway. The
|
||||
blocked ToolMessage carries ``WRITE_BLOCK_KEY`` and ``wrap_model_call``
|
||||
replaces the paired call's payload arguments (``content``, ``old_str``,
|
||||
``new_str``) with a short deterministic placeholder in the *model-bound
|
||||
request only*. ``state["messages"]``, tool receipts, and the run journal
|
||||
keep the original arguments, and nothing is externalized to disk: handing
|
||||
the model a file reference to content it must re-derive after reading the
|
||||
target would only invite bypassing the gate through ``bash``.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
@ -31,15 +40,19 @@ import logging
|
||||
import posixpath
|
||||
import threading
|
||||
import weakref
|
||||
from collections import defaultdict, deque
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any, override
|
||||
|
||||
from langchain.agents.middleware import AgentMiddleware
|
||||
from langchain_core.messages import ToolMessage
|
||||
from langchain.agents.middleware.types import ModelCallResult, ModelRequest, ModelResponse
|
||||
from langchain_core.messages import AIMessage, ToolMessage
|
||||
from langgraph.prebuilt.tool_node import ToolCallRequest
|
||||
from langgraph.types import Command
|
||||
|
||||
from deerflow.agents.middlewares.tool_call_args import rewrite_messages_tool_call_args
|
||||
from deerflow.agents.middlewares.tool_result_meta import normalize_tool_result, stamp_exception_meta
|
||||
from deerflow.config.read_before_write_config import ReadBeforeWriteConfig
|
||||
from deerflow.sandbox.exceptions import SandboxAuthorizationError
|
||||
from deerflow.sandbox.tools import (
|
||||
read_current_file_content,
|
||||
@ -50,9 +63,20 @@ from deerflow.sandbox.tools import (
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
READ_MARK_KEY = "deerflow_read_mark"
|
||||
#: Stamped on the error ToolMessage of a gate-blocked call: ``{"path", "tool"}``.
|
||||
WRITE_BLOCK_KEY = "deerflow_write_block"
|
||||
|
||||
_READ_TOOLS = frozenset({"read_file"})
|
||||
_GATED_WRITE_TOOLS = frozenset({"write_file", "str_replace"})
|
||||
# Payload arguments per gated tool — the bulk of a write call. Everything else
|
||||
# (path, description, flags) stays visible after a block.
|
||||
_PAYLOAD_FIELDS: dict[str, tuple[str, ...]] = {
|
||||
"write_file": ("content",),
|
||||
"str_replace": ("old_str", "new_str"),
|
||||
}
|
||||
# Deterministic for a given payload so repeated model calls keep the same
|
||||
# request prefix (prompt caching) instead of drifting.
|
||||
_ELIDED_PAYLOAD_TEMPLATE = "[payload elided: {chars} chars; this {tool_name} call was blocked by the read-before-write gate and nothing was written]"
|
||||
|
||||
# AIO/E2B-style sandboxes convert read failures (including missing files)
|
||||
# into "Error: ..." strings instead of raising. Content with this prefix is
|
||||
@ -95,9 +119,18 @@ def _content_hash(content: str) -> str:
|
||||
class ReadBeforeWriteMiddleware(AgentMiddleware):
|
||||
"""Version gate: block writes to existing files not read at their current version."""
|
||||
|
||||
def __init__(self, content_reader: Callable[[Any, str], str] | None = None) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
content_reader: Callable[[Any, str], str] | None = None,
|
||||
*,
|
||||
config: ReadBeforeWriteConfig | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self._content_reader = content_reader or read_current_file_content
|
||||
self._config = config if config is not None else ReadBeforeWriteConfig()
|
||||
|
||||
def release_policy_parameters(self) -> dict[str, object]:
|
||||
return {"config": self._config.model_dump(mode="python")}
|
||||
|
||||
@override
|
||||
def wrap_tool_call(
|
||||
@ -248,6 +281,7 @@ class ReadBeforeWriteMiddleware(AgentMiddleware):
|
||||
tool_call_id=str(tool_call.get("id", "")),
|
||||
name=tool_name,
|
||||
status="error",
|
||||
additional_kwargs={WRITE_BLOCK_KEY: {"path": norm_path, "tool": tool_name}},
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
@ -272,6 +306,36 @@ class ReadBeforeWriteMiddleware(AgentMiddleware):
|
||||
return mark_hash if isinstance(mark_hash, str) else None
|
||||
return None
|
||||
|
||||
# -- model-bound payload elision --------------------------------------
|
||||
|
||||
@override
|
||||
def wrap_model_call(
|
||||
self,
|
||||
request: ModelRequest,
|
||||
handler: Callable[[ModelRequest], ModelResponse],
|
||||
) -> ModelCallResult:
|
||||
return handler(self._elide_blocked_payloads(request))
|
||||
|
||||
@override
|
||||
async def awrap_model_call(
|
||||
self,
|
||||
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._elide_blocked_payloads(request))
|
||||
|
||||
def _elide_blocked_payloads(self, request: ModelRequest) -> ModelRequest:
|
||||
if not self._config.elide_blocked_payloads:
|
||||
return request
|
||||
messages = getattr(request, "messages", None)
|
||||
if not isinstance(messages, list):
|
||||
return request
|
||||
patched = elide_blocked_write_payloads(messages, min_chars=self._config.elide_min_chars)
|
||||
if patched is None:
|
||||
return request
|
||||
return request.override(messages=patched)
|
||||
|
||||
# -- mark stamping ---------------------------------------------------
|
||||
|
||||
def _attach_read_mark(self, request: ToolCallRequest, result: ToolMessage | Command) -> None:
|
||||
@ -305,3 +369,75 @@ class ReadBeforeWriteMiddleware(AgentMiddleware):
|
||||
if candidates:
|
||||
return candidates[-1]
|
||||
return None
|
||||
|
||||
|
||||
# -- blocked payload elision (policy) -----------------------------------------
|
||||
|
||||
|
||||
def elide_blocked_write_payloads(messages: list[Any], *, min_chars: int) -> list[Any] | None:
|
||||
"""Return ``messages`` with gate-blocked write payloads replaced by placeholders, or ``None`` if unchanged.
|
||||
|
||||
Only the policy lives here: a call qualifies when a ``WRITE_BLOCK_KEY``
|
||||
ToolMessage answered it, and its payload fields become
|
||||
``_ELIDED_PAYLOAD_TEMPLATE``. The surface-by-surface rewrite (structured
|
||||
``tool_calls``, raw provider payload, ``tool_use`` blocks, chunk args) is
|
||||
``tool_call_args.rewrite_messages_tool_call_args``, which never mutates the
|
||||
input and passes untouched messages through by identity, so the stored
|
||||
history keeps the original arguments and the output is identical across
|
||||
model calls.
|
||||
"""
|
||||
blocked = _blocked_call_occurrences(messages)
|
||||
if not blocked:
|
||||
return None
|
||||
|
||||
def replacement_for(message: AIMessage, tool_call: dict[str, Any]) -> dict[str, Any] | None:
|
||||
if (id(message), tool_call.get("id")) not in blocked:
|
||||
return None
|
||||
args = tool_call.get("args")
|
||||
return _elide_args(args, str(tool_call.get("name")), min_chars) if isinstance(args, dict) else None
|
||||
|
||||
return rewrite_messages_tool_call_args(messages, replacement_for)
|
||||
|
||||
|
||||
def _blocked_call_occurrences(messages: list[Any]) -> set[tuple[int, str]]:
|
||||
"""Return ``(id(ai_message), call_id)`` for every call occurrence answered by a gate-blocked result.
|
||||
|
||||
Tool-call ids may repeat across assistant turns, so a history-wide id set
|
||||
would also hit an earlier (or later) *successful* call with the same id and
|
||||
mislabel it as blocked. Results are paired with call occurrences the way
|
||||
``DanglingToolCallMiddleware`` does: ToolMessages queue per id in history
|
||||
order and each AIMessage call consumes the next one for its id.
|
||||
"""
|
||||
results_by_id: dict[str, deque[ToolMessage]] = defaultdict(deque)
|
||||
for message in messages:
|
||||
if isinstance(message, ToolMessage) and isinstance(message.tool_call_id, str) and message.tool_call_id:
|
||||
results_by_id[message.tool_call_id].append(message)
|
||||
|
||||
blocked: set[tuple[int, str]] = set()
|
||||
for message in messages:
|
||||
if not isinstance(message, AIMessage):
|
||||
continue
|
||||
for tool_call in message.tool_calls or ():
|
||||
call_id = tool_call.get("id") if isinstance(tool_call, dict) else None
|
||||
if not isinstance(call_id, str) or not call_id:
|
||||
continue
|
||||
queue = results_by_id.get(call_id)
|
||||
result = queue.popleft() if queue else None
|
||||
if result is not None and isinstance((result.additional_kwargs or {}).get(WRITE_BLOCK_KEY), dict):
|
||||
blocked.add((id(message), call_id))
|
||||
return blocked
|
||||
|
||||
|
||||
def _elide_args(args: dict[str, Any], tool_name: str, min_chars: int) -> dict[str, Any] | None:
|
||||
fields = _PAYLOAD_FIELDS.get(tool_name)
|
||||
if not fields:
|
||||
return None
|
||||
elided: dict[str, Any] | None = None
|
||||
for field in fields:
|
||||
value = args.get(field)
|
||||
if not isinstance(value, str) or not value or len(value) < min_chars:
|
||||
continue
|
||||
if elided is None:
|
||||
elided = dict(args)
|
||||
elided[field] = _ELIDED_PAYLOAD_TEMPLATE.format(chars=len(value), tool_name=tool_name)
|
||||
return elided
|
||||
|
||||
@ -0,0 +1,199 @@
|
||||
"""Rewrite AIMessage tool-call arguments on every provider surface at once.
|
||||
|
||||
Middlewares that shrink or replace a historical tool call's arguments in the
|
||||
*model-bound request* (never in graph state) share one hazard: a LangChain
|
||||
``AIMessage`` carries the same arguments on up to four surfaces, and provider
|
||||
adapters do not all read the same one —
|
||||
|
||||
- ``tool_calls``: the structured list most adapters prefer;
|
||||
- ``additional_kwargs["tool_calls"]``: the raw provider payload (OpenAI
|
||||
``function.arguments`` JSON string) some adapters fall back to;
|
||||
- ``content`` blocks that carry their own copy of the arguments: Anthropic
|
||||
``tool_use`` (``input`` + ``partial_json``), OpenAI Responses
|
||||
``function_call`` (``arguments`` string, matched by ``call_id``; the
|
||||
``fc_…`` item id is preserved), and LangChain standard-content
|
||||
``tool_call`` / ``tool_call_chunk`` (``args`` plus ``extras.arguments``);
|
||||
- ``tool_call_chunks`` on an ``AIMessageChunk``.
|
||||
|
||||
Rewriting only one surface leaves the original payload reachable through the
|
||||
others and can hand a strict provider a request whose surfaces disagree. The
|
||||
content surfaces matter most: ``langchain_openai``'s Responses input builder
|
||||
emits a content ``function_call`` block *instead of* the structured call
|
||||
whose ``call_id`` it already carries, and prefers ``extras.arguments`` over
|
||||
the structured args when translating a v1 ``tool_call`` block, so a rewrite
|
||||
that touched ``tool_calls`` alone would still send the original payload.
|
||||
:func:`rewrite_tool_call_args` rewrites them together and returns a
|
||||
``model_copy`` (or the same object when nothing matched), so callers never
|
||||
mutate state and the result is identical across model calls. Policy — which
|
||||
calls, and what replaces their arguments — stays with the caller; see
|
||||
``read_before_write_middleware.elide_blocked_write_payloads`` for one.
|
||||
|
||||
A rewrite also invalidates server-side continuation. With
|
||||
``use_previous_response_id`` the OpenAI adapter sends only the messages after
|
||||
the last AIMessage carrying a ``resp_…`` ``response_metadata["id"]`` and lets
|
||||
the server rebuild the rest from *its* stored copy of the conversation, which
|
||||
still holds the original arguments; stored responses cannot be edited, and
|
||||
every response produced after the rewritten call chains back to that history.
|
||||
So whenever anything was rewritten, :func:`rewrite_messages_tool_call_args`
|
||||
drops every ``resp_`` id from the model-bound copy and the adapter falls back
|
||||
to replaying the full rewritten history (the same request shape as
|
||||
``use_previous_response_id=False``; per OpenAI's docs chained input tokens are
|
||||
billed either way, so replay costs no more).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from typing import Any
|
||||
|
||||
from langchain_core.messages import AIMessage
|
||||
|
||||
#: Replacement args keyed by tool-call id.
|
||||
ArgsReplacements = Mapping[str, dict[str, Any]]
|
||||
#: ``(message, tool_call) -> new_args`` or ``None`` to leave the call alone.
|
||||
ReplacementSelector = Callable[[AIMessage, dict[str, Any]], dict[str, Any] | None]
|
||||
|
||||
|
||||
def rewrite_messages_tool_call_args(messages: list[Any], replacement_for: ReplacementSelector) -> list[Any] | None:
|
||||
"""Apply ``replacement_for(message, tool_call)`` to every AIMessage tool call in ``messages``.
|
||||
|
||||
Returns a new list with the rewritten AIMessages, or ``None`` when no call
|
||||
was replaced. Untouched messages pass through by identity, except that once
|
||||
anything was rewritten every AIMessage loses its ``resp_`` response id (see
|
||||
the module docstring: the server-side history behind that id still holds
|
||||
the original arguments). Only calls with a non-empty string id are offered
|
||||
to the selector, since nothing else can be matched across surfaces.
|
||||
"""
|
||||
updated: list[Any] = []
|
||||
changed = False
|
||||
for message in messages:
|
||||
patched = message
|
||||
if isinstance(message, AIMessage) and message.tool_calls:
|
||||
replacements: dict[str, dict[str, Any]] = {}
|
||||
for tool_call in message.tool_calls:
|
||||
if not isinstance(tool_call, dict):
|
||||
continue
|
||||
call_id = tool_call.get("id")
|
||||
if not isinstance(call_id, str) or not call_id:
|
||||
continue
|
||||
new_args = replacement_for(message, tool_call)
|
||||
if new_args is not None:
|
||||
replacements[call_id] = new_args
|
||||
if replacements:
|
||||
patched = rewrite_tool_call_args(message, replacements)
|
||||
if patched is not message:
|
||||
changed = True
|
||||
updated.append(patched)
|
||||
if not changed:
|
||||
return None
|
||||
return [_without_response_chain_id(message) for message in updated]
|
||||
|
||||
|
||||
def _without_response_chain_id(message: Any) -> Any:
|
||||
"""Drop an OpenAI ``resp_`` response id so the adapter replays history instead of chaining to it."""
|
||||
if not isinstance(message, AIMessage):
|
||||
return message
|
||||
response_metadata = message.response_metadata or {}
|
||||
response_id = response_metadata.get("id")
|
||||
if not (isinstance(response_id, str) and response_id.startswith("resp_")):
|
||||
return message
|
||||
return message.model_copy(update={"response_metadata": {key: value for key, value in response_metadata.items() if key != "id"}})
|
||||
|
||||
|
||||
def rewrite_tool_call_args(message: AIMessage, replacements: ArgsReplacements) -> AIMessage:
|
||||
"""Return ``message`` with the args of every tool call in ``replacements`` (by id) rewritten on all surfaces.
|
||||
|
||||
``message`` is never mutated; the same object comes back when no id matches.
|
||||
"""
|
||||
if not replacements:
|
||||
return message
|
||||
update: dict[str, Any] = {}
|
||||
|
||||
tool_calls = message.tool_calls or []
|
||||
rewritten_calls = [dict(tool_call, args=new_args) if isinstance(tool_call, dict) and (new_args := _replacement_for_id(tool_call.get("id"), replacements)) is not None else tool_call for tool_call in tool_calls]
|
||||
if _any_replaced(rewritten_calls, tool_calls):
|
||||
update["tool_calls"] = rewritten_calls
|
||||
|
||||
tool_call_chunks = getattr(message, "tool_call_chunks", None)
|
||||
if isinstance(tool_call_chunks, list):
|
||||
rewritten_chunks = [dict(chunk, args=_serialize(new_args)) if isinstance(chunk, dict) and (new_args := _replacement_for_id(chunk.get("id"), replacements)) is not None else chunk for chunk in tool_call_chunks]
|
||||
if _any_replaced(rewritten_chunks, tool_call_chunks):
|
||||
update["tool_call_chunks"] = rewritten_chunks
|
||||
|
||||
additional_kwargs = message.additional_kwargs or {}
|
||||
raw_tool_calls = additional_kwargs.get("tool_calls")
|
||||
if isinstance(raw_tool_calls, list):
|
||||
rewritten_raw = [_rewrite_raw_tool_call(entry, replacements) for entry in raw_tool_calls]
|
||||
if _any_replaced(rewritten_raw, raw_tool_calls):
|
||||
update["additional_kwargs"] = {**additional_kwargs, "tool_calls": rewritten_raw}
|
||||
|
||||
if isinstance(message.content, list):
|
||||
rewritten_content = [_rewrite_content_block(block, replacements) for block in message.content]
|
||||
if _any_replaced(rewritten_content, message.content):
|
||||
update["content"] = rewritten_content
|
||||
|
||||
return message.model_copy(update=update) if update else message
|
||||
|
||||
|
||||
def _replacement_for_id(identifier: Any, replacements: ArgsReplacements) -> dict[str, Any] | None:
|
||||
"""Non-string ids (malformed provider payloads) never match, and never raise from a membership probe."""
|
||||
return replacements.get(identifier) if isinstance(identifier, str) else None
|
||||
|
||||
|
||||
def _any_replaced(rewritten: Sequence[Any], original: Sequence[Any]) -> bool:
|
||||
return any(new is not old for new, old in zip(rewritten, original, strict=True))
|
||||
|
||||
|
||||
def _serialize(args: dict[str, Any]) -> str:
|
||||
return json.dumps(args, ensure_ascii=False)
|
||||
|
||||
|
||||
def _rewrite_raw_tool_call(entry: Any, replacements: ArgsReplacements) -> Any:
|
||||
"""Rewrite one raw provider tool-call payload (OpenAI ``function.arguments`` JSON string, or flattened variants)."""
|
||||
if not isinstance(entry, dict):
|
||||
return entry
|
||||
new_args = _replacement_for_id(entry.get("id"), replacements)
|
||||
if new_args is None:
|
||||
return entry
|
||||
function = entry.get("function")
|
||||
if isinstance(function, dict):
|
||||
return {**entry, "function": {**function, "arguments": _serialize(new_args)}}
|
||||
if isinstance(entry.get("arguments"), str):
|
||||
return {**entry, "arguments": _serialize(new_args)}
|
||||
if isinstance(entry.get("args"), dict):
|
||||
return {**entry, "args": new_args}
|
||||
return entry
|
||||
|
||||
|
||||
def _rewrite_content_block(block: Any, replacements: ArgsReplacements) -> Any:
|
||||
"""Rewrite one content block that carries tool-call arguments; anything else passes through by identity."""
|
||||
if not isinstance(block, dict):
|
||||
return block
|
||||
block_type = block.get("type")
|
||||
if block_type == "tool_use":
|
||||
# Anthropic: ``partial_json`` is dropped so it cannot leak the old payload.
|
||||
new_args = _replacement_for_id(block.get("id"), replacements)
|
||||
if new_args is None:
|
||||
return block
|
||||
rewritten = {key: value for key, value in block.items() if key != "partial_json"}
|
||||
rewritten["input"] = new_args
|
||||
return rewritten
|
||||
if block_type == "function_call":
|
||||
# OpenAI Responses (``responses/v1``): matched by ``call_id``; the ``fc_…`` item id and status are kept.
|
||||
new_args = _replacement_for_id(block.get("call_id"), replacements)
|
||||
if new_args is None:
|
||||
return block
|
||||
return {**block, "arguments": _serialize(new_args)}
|
||||
if block_type in ("tool_call", "tool_call_chunk"):
|
||||
# LangChain standard content (``v1``): ``args`` is a dict on tool_call and a JSON string on
|
||||
# tool_call_chunk; ``extras.arguments`` (raw provider string) wins in the Responses translator.
|
||||
new_args = _replacement_for_id(block.get("id"), replacements)
|
||||
if new_args is None:
|
||||
return block
|
||||
rewritten = {**block, "args": new_args if block_type == "tool_call" else _serialize(new_args)}
|
||||
extras = block.get("extras")
|
||||
if isinstance(extras, dict) and "arguments" in extras:
|
||||
rewritten["extras"] = {**extras, "arguments": _serialize(new_args)}
|
||||
return rewritten
|
||||
return block
|
||||
@ -288,11 +288,13 @@ def _build_runtime_middlewares(
|
||||
# the model hasn't read in their current version. It must sit outside ToolProgress
|
||||
# and ToolErrorHandling so that a blocked write returns immediately without consuming
|
||||
# a ToolProgress slot. The middleware stamps deerflow_tool_meta on the blocked
|
||||
# ToolMessage itself so downstream callers receive a well-formed result.
|
||||
# ToolMessage itself so downstream callers receive a well-formed result, and its
|
||||
# wrap_model_call elides the dead payload of blocked calls from model-bound
|
||||
# requests (config-gated, state untouched).
|
||||
if app_config.read_before_write.enabled:
|
||||
from deerflow.agents.middlewares.read_before_write_middleware import ReadBeforeWriteMiddleware
|
||||
|
||||
tail.append(ReadBeforeWriteMiddleware())
|
||||
tail.append(ReadBeforeWriteMiddleware(config=app_config.read_before_write))
|
||||
|
||||
# ToolProgressMiddleware must be outer (lower index) so its wrap_tool_call handler
|
||||
# chain includes ToolErrorHandlingMiddleware (inner), which stamps deerflow_tool_meta
|
||||
|
||||
@ -16,3 +16,22 @@ class ReadBeforeWriteConfig(BaseModel):
|
||||
default=True,
|
||||
description="Whether to block writes to existing files that were not read at their current version",
|
||||
)
|
||||
elide_blocked_payloads: bool = Field(
|
||||
default=True,
|
||||
description=(
|
||||
"Replace the payload arguments of gate-blocked calls (write_file content, str_replace old_str/new_str) "
|
||||
"with a short placeholder in model-bound requests. A blocked call never ran and must be re-issued after a "
|
||||
"re-read, so its payload is dead weight in every later model call. Only the request copy changes: stored "
|
||||
"message history, receipts, and the run journal keep the original arguments."
|
||||
),
|
||||
)
|
||||
elide_min_chars: int = Field(
|
||||
default=2000,
|
||||
ge=0,
|
||||
description=(
|
||||
"Elide only payload fields at least this many characters long; shorter payloads stay visible so the model "
|
||||
"can reuse them after re-reading. 0 elides every non-empty payload. This is a Python character count, not a "
|
||||
"token count: the same value spans roughly 3-4x in real context cost between ASCII and CJK text, and the "
|
||||
"placeholder's elided-size figure is the same character count."
|
||||
),
|
||||
)
|
||||
|
||||
@ -13,11 +13,11 @@ def _sha(text: str) -> str:
|
||||
return hashlib.sha256(text.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
def _make_request(name, args, messages=()):
|
||||
def _make_request(name, args, messages=(), tool_call_id="call-1"):
|
||||
runtime = MagicMock()
|
||||
runtime.context = {"thread_id": "t-test"}
|
||||
return ToolCallRequest(
|
||||
tool_call={"name": name, "args": args, "id": "call-1"},
|
||||
tool_call={"name": name, "args": args, "id": tool_call_id},
|
||||
tool=None,
|
||||
state={"messages": list(messages)},
|
||||
runtime=runtime,
|
||||
@ -437,3 +437,358 @@ class TestSamePathSerialization:
|
||||
mark = read_result.additional_kwargs.get("deerflow_read_mark")
|
||||
assert mark is not None
|
||||
assert mark["hash"] == _sha(read_result.content)
|
||||
|
||||
|
||||
class TestBlockedPayloadElision:
|
||||
"""Model-bound requests drop the dead payload of gate-blocked writes; state stays intact."""
|
||||
|
||||
PATH = "/mnt/user-data/outputs/report.md"
|
||||
|
||||
@staticmethod
|
||||
def _config(**overrides):
|
||||
from deerflow.config.read_before_write_config import ReadBeforeWriteConfig
|
||||
|
||||
return ReadBeforeWriteConfig(**overrides)
|
||||
|
||||
def _middleware(self, files=None, **config_overrides):
|
||||
from deerflow.agents.middlewares.read_before_write_middleware import ReadBeforeWriteMiddleware
|
||||
|
||||
files = {self.PATH: "v1"} if files is None else files
|
||||
|
||||
def reader(_runtime, path):
|
||||
normalized = posixpath.normpath(path)
|
||||
if normalized not in files:
|
||||
raise FileNotFoundError(path)
|
||||
return files[normalized]
|
||||
|
||||
return ReadBeforeWriteMiddleware(content_reader=reader, config=self._config(**config_overrides))
|
||||
|
||||
@staticmethod
|
||||
def _model_request(messages):
|
||||
from langchain.agents.middleware.types import ModelRequest
|
||||
|
||||
return ModelRequest(model=None, messages=list(messages), tools=[], state={"messages": list(messages)}, runtime=MagicMock())
|
||||
|
||||
def _blocked_turn(self, mw, name, args, tool_call_id="call-1"):
|
||||
"""Run a gated call against an unread file; return ``(AIMessage, blocked ToolMessage)``."""
|
||||
ai = AIMessage(content="", tool_calls=[{"name": name, "id": tool_call_id, "args": dict(args)}])
|
||||
request = _make_request(name, dict(args), [HumanMessage(content="go"), ai], tool_call_id=tool_call_id)
|
||||
blocked = mw.wrap_tool_call(request, MagicMock(side_effect=AssertionError("handler must not run when blocked")))
|
||||
assert blocked.status == "error"
|
||||
return ai, blocked
|
||||
|
||||
@staticmethod
|
||||
def _captured(handler):
|
||||
return handler.call_args[0][0]
|
||||
|
||||
def test_blocked_result_carries_write_block_marker(self):
|
||||
from deerflow.agents.middlewares.read_before_write_middleware import WRITE_BLOCK_KEY
|
||||
|
||||
mw = self._middleware()
|
||||
_ai, blocked = self._blocked_turn(mw, "write_file", {"description": "d", "path": self.PATH, "content": "v2"})
|
||||
assert blocked.additional_kwargs[WRITE_BLOCK_KEY] == {"path": self.PATH, "tool": "write_file"}
|
||||
|
||||
def test_allowed_write_result_has_no_marker(self):
|
||||
from deerflow.agents.middlewares.read_before_write_middleware import WRITE_BLOCK_KEY
|
||||
|
||||
mw = self._middleware()
|
||||
messages = [_read_marked_message(self.PATH, "v1")]
|
||||
request = _make_request("write_file", {"description": "d", "path": self.PATH, "content": "v2"}, messages)
|
||||
handler = MagicMock(return_value=ToolMessage(content="OK", tool_call_id="call-1", name="write_file"))
|
||||
result = mw.wrap_tool_call(request, handler)
|
||||
assert WRITE_BLOCK_KEY not in result.additional_kwargs
|
||||
|
||||
def test_elides_blocked_write_file_content_in_model_request(self):
|
||||
mw = self._middleware()
|
||||
payload = "x" * 5000
|
||||
ai, blocked = self._blocked_turn(mw, "write_file", {"description": "d", "path": self.PATH, "content": payload})
|
||||
human = HumanMessage(content="go")
|
||||
request = self._model_request([human, ai, blocked])
|
||||
handler = MagicMock(return_value=AIMessage(content="ok"))
|
||||
|
||||
mw.wrap_model_call(request, handler)
|
||||
|
||||
captured = self._captured(handler)
|
||||
assert captured is not request
|
||||
rewritten = captured.messages[1]
|
||||
assert rewritten is not ai
|
||||
args = rewritten.tool_calls[0]["args"]
|
||||
assert args["path"] == self.PATH
|
||||
assert args["description"] == "d"
|
||||
assert args["content"].startswith("[payload elided: 5000 chars")
|
||||
assert "read-before-write" in args["content"]
|
||||
assert payload not in args["content"]
|
||||
# Untouched neighbours are passed through by identity; the stored history is never rewritten.
|
||||
assert captured.messages[0] is human
|
||||
assert captured.messages[2] is blocked
|
||||
assert request.messages[1] is ai
|
||||
assert request.state["messages"][1] is ai
|
||||
assert ai.tool_calls[0]["args"]["content"] == payload
|
||||
|
||||
def test_successful_write_payload_is_left_alone(self):
|
||||
mw = self._middleware()
|
||||
payload = "x" * 5000
|
||||
ai = AIMessage(content="", tool_calls=[{"name": "write_file", "id": "call-1", "args": {"description": "d", "path": self.PATH, "content": payload}}])
|
||||
ok = ToolMessage(content="OK", tool_call_id="call-1", name="write_file")
|
||||
request = self._model_request([HumanMessage(content="go"), ai, ok])
|
||||
handler = MagicMock(return_value=AIMessage(content="ok"))
|
||||
|
||||
mw.wrap_model_call(request, handler)
|
||||
|
||||
assert self._captured(handler) is request
|
||||
assert ai.tool_calls[0]["args"]["content"] == payload
|
||||
|
||||
def test_only_the_blocked_call_is_elided_when_ids_differ(self):
|
||||
mw = self._middleware()
|
||||
payload = "x" * 5000
|
||||
ai, blocked = self._blocked_turn(mw, "write_file", {"description": "d", "path": self.PATH, "content": payload}, tool_call_id="call-blocked")
|
||||
other = AIMessage(content="", tool_calls=[{"name": "write_file", "id": "call-ok", "args": {"description": "d", "path": "/mnt/user-data/outputs/new.md", "content": payload}}])
|
||||
ok = ToolMessage(content="OK", tool_call_id="call-ok", name="write_file")
|
||||
request = self._model_request([HumanMessage(content="go"), other, ok, ai, blocked])
|
||||
handler = MagicMock(return_value=AIMessage(content="ok"))
|
||||
|
||||
mw.wrap_model_call(request, handler)
|
||||
|
||||
captured = self._captured(handler)
|
||||
assert captured.messages[1] is other
|
||||
assert captured.messages[3].tool_calls[0]["args"]["content"].startswith("[payload elided")
|
||||
|
||||
def test_rewrites_raw_tool_calls_and_tool_use_blocks_consistently(self):
|
||||
import json
|
||||
|
||||
mw = self._middleware()
|
||||
payload = "y" * 5000
|
||||
args = {"description": "d", "path": self.PATH, "content": payload}
|
||||
ai = AIMessage(
|
||||
content=[
|
||||
{"type": "text", "text": "writing"},
|
||||
{"type": "tool_use", "id": "call-1", "name": "write_file", "input": dict(args), "partial_json": json.dumps(args)},
|
||||
],
|
||||
tool_calls=[{"name": "write_file", "id": "call-1", "args": dict(args)}],
|
||||
additional_kwargs={"tool_calls": [{"id": "call-1", "type": "function", "function": {"name": "write_file", "arguments": json.dumps(args)}}]},
|
||||
)
|
||||
request = _make_request("write_file", dict(args), [HumanMessage(content="go"), ai])
|
||||
blocked = mw.wrap_tool_call(request, MagicMock())
|
||||
model_request = self._model_request([HumanMessage(content="go"), ai, blocked])
|
||||
handler = MagicMock(return_value=AIMessage(content="ok"))
|
||||
|
||||
mw.wrap_model_call(model_request, handler)
|
||||
|
||||
rewritten = self._captured(handler).messages[1]
|
||||
structured = rewritten.tool_calls[0]["args"]
|
||||
assert structured["content"].startswith("[payload elided")
|
||||
raw = json.loads(rewritten.additional_kwargs["tool_calls"][0]["function"]["arguments"])
|
||||
assert raw == structured
|
||||
assert rewritten.additional_kwargs["tool_calls"][0]["function"]["name"] == "write_file"
|
||||
block = rewritten.content[1]
|
||||
assert block["input"] == structured
|
||||
assert "partial_json" not in block
|
||||
assert rewritten.content[0] == {"type": "text", "text": "writing"}
|
||||
# Serialized payload must be gone from every surface the provider adapters read.
|
||||
assert payload not in json.dumps(rewritten.model_dump(), ensure_ascii=False)
|
||||
# Original objects are untouched.
|
||||
assert ai.content[1]["input"]["content"] == payload
|
||||
assert payload in ai.additional_kwargs["tool_calls"][0]["function"]["arguments"]
|
||||
|
||||
def test_str_replace_elides_old_and_new_str(self):
|
||||
mw = self._middleware()
|
||||
old_str, new_str = "a" * 3000, "b" * 4000
|
||||
ai, blocked = self._blocked_turn(mw, "str_replace", {"description": "d", "path": self.PATH, "old_str": old_str, "new_str": new_str})
|
||||
request = self._model_request([HumanMessage(content="go"), ai, blocked])
|
||||
handler = MagicMock(return_value=AIMessage(content="ok"))
|
||||
|
||||
mw.wrap_model_call(request, handler)
|
||||
|
||||
args = self._captured(handler).messages[1].tool_calls[0]["args"]
|
||||
assert args["old_str"].startswith("[payload elided: 3000 chars")
|
||||
assert args["new_str"].startswith("[payload elided: 4000 chars")
|
||||
assert "str_replace" in args["new_str"]
|
||||
assert args["path"] == self.PATH
|
||||
|
||||
def test_payload_below_min_chars_stays_visible(self):
|
||||
mw = self._middleware()
|
||||
ai, blocked = self._blocked_turn(mw, "write_file", {"description": "d", "path": self.PATH, "content": "short " * 20})
|
||||
request = self._model_request([HumanMessage(content="go"), ai, blocked])
|
||||
handler = MagicMock(return_value=AIMessage(content="ok"))
|
||||
|
||||
mw.wrap_model_call(request, handler)
|
||||
|
||||
assert self._captured(handler) is request
|
||||
|
||||
def test_mixed_fields_only_elide_those_over_threshold(self):
|
||||
mw = self._middleware(elide_min_chars=1000)
|
||||
ai, blocked = self._blocked_turn(mw, "str_replace", {"description": "d", "path": self.PATH, "old_str": "tiny", "new_str": "n" * 1000})
|
||||
request = self._model_request([HumanMessage(content="go"), ai, blocked])
|
||||
handler = MagicMock(return_value=AIMessage(content="ok"))
|
||||
|
||||
mw.wrap_model_call(request, handler)
|
||||
|
||||
args = self._captured(handler).messages[1].tool_calls[0]["args"]
|
||||
assert args["old_str"] == "tiny"
|
||||
assert args["new_str"].startswith("[payload elided: 1000 chars")
|
||||
|
||||
def test_min_chars_zero_elides_any_non_empty_payload(self):
|
||||
mw = self._middleware(elide_min_chars=0)
|
||||
ai, blocked = self._blocked_turn(mw, "write_file", {"description": "d", "path": self.PATH, "content": "v2"})
|
||||
empty_ai, empty_blocked = self._blocked_turn(mw, "write_file", {"description": "d", "path": self.PATH, "content": ""}, tool_call_id="call-2")
|
||||
request = self._model_request([HumanMessage(content="go"), ai, blocked, empty_ai, empty_blocked])
|
||||
handler = MagicMock(return_value=AIMessage(content="ok"))
|
||||
|
||||
mw.wrap_model_call(request, handler)
|
||||
|
||||
captured = self._captured(handler)
|
||||
assert captured.messages[1].tool_calls[0]["args"]["content"].startswith("[payload elided: 2 chars")
|
||||
assert captured.messages[3] is empty_ai
|
||||
|
||||
def test_disabled_by_config_passes_request_through(self):
|
||||
mw = self._middleware(elide_blocked_payloads=False)
|
||||
ai, blocked = self._blocked_turn(mw, "write_file", {"description": "d", "path": self.PATH, "content": "x" * 5000})
|
||||
request = self._model_request([HumanMessage(content="go"), ai, blocked])
|
||||
handler = MagicMock(return_value=AIMessage(content="ok"))
|
||||
|
||||
mw.wrap_model_call(request, handler)
|
||||
|
||||
assert self._captured(handler) is request
|
||||
|
||||
def test_elision_is_deterministic_across_model_calls(self):
|
||||
mw = self._middleware()
|
||||
ai, blocked = self._blocked_turn(mw, "write_file", {"description": "d", "path": self.PATH, "content": "x" * 5000})
|
||||
first, second = MagicMock(return_value=AIMessage(content="ok")), MagicMock(return_value=AIMessage(content="ok"))
|
||||
|
||||
mw.wrap_model_call(self._model_request([HumanMessage(content="go"), ai, blocked]), first)
|
||||
mw.wrap_model_call(self._model_request([HumanMessage(content="go"), ai, blocked]), second)
|
||||
|
||||
assert self._captured(first).messages[1].tool_calls == self._captured(second).messages[1].tool_calls
|
||||
|
||||
def test_async_model_call_elides(self):
|
||||
import asyncio
|
||||
|
||||
mw = self._middleware()
|
||||
ai, blocked = self._blocked_turn(mw, "write_file", {"description": "d", "path": self.PATH, "content": "x" * 5000})
|
||||
request = self._model_request([HumanMessage(content="go"), ai, blocked])
|
||||
seen = {}
|
||||
|
||||
async def handler(model_request):
|
||||
seen["request"] = model_request
|
||||
return AIMessage(content="ok")
|
||||
|
||||
asyncio.run(mw.awrap_model_call(request, handler))
|
||||
|
||||
assert seen["request"] is not request
|
||||
assert seen["request"].messages[1].tool_calls[0]["args"]["content"].startswith("[payload elided")
|
||||
|
||||
def test_release_policy_declares_config(self):
|
||||
mw = self._middleware(elide_min_chars=123)
|
||||
params = mw.release_policy_parameters()
|
||||
assert params["config"]["enabled"] is True
|
||||
assert params["config"]["elide_blocked_payloads"] is True
|
||||
assert params["config"]["elide_min_chars"] == 123
|
||||
|
||||
def test_malformed_unhashable_ids_do_not_break_elision(self):
|
||||
import json
|
||||
|
||||
mw = self._middleware()
|
||||
payload = "z" * 5000
|
||||
args = {"description": "d", "path": self.PATH, "content": payload}
|
||||
ai, blocked = self._blocked_turn(mw, "write_file", args)
|
||||
# A provider payload with a list-typed id must be skipped, not raise from a membership probe.
|
||||
weird = AIMessage(
|
||||
content=[{"type": "tool_use", "id": ["not", "a", "string"], "name": "write_file", "input": dict(args)}],
|
||||
tool_calls=[{"name": "write_file", "id": "call-1", "args": dict(args)}],
|
||||
additional_kwargs={"tool_calls": [{"id": ["not", "a", "string"], "type": "function", "function": {"name": "write_file", "arguments": json.dumps(args)}}]},
|
||||
)
|
||||
request = self._model_request([HumanMessage(content="go"), weird, blocked])
|
||||
handler = MagicMock(return_value=AIMessage(content="ok"))
|
||||
|
||||
mw.wrap_model_call(request, handler)
|
||||
|
||||
rewritten = self._captured(handler).messages[1]
|
||||
assert rewritten.tool_calls[0]["args"]["content"].startswith("[payload elided")
|
||||
assert rewritten.content[0]["input"]["content"] == payload
|
||||
assert payload in rewritten.additional_kwargs["tool_calls"][0]["function"]["arguments"]
|
||||
|
||||
def test_responses_api_request_input_never_carries_the_blocked_payload(self):
|
||||
"""End to end against the real OpenAI Responses input builder (reviewer probe on #5329)."""
|
||||
import json
|
||||
|
||||
from langchain_openai.chat_models.base import _construct_responses_api_input
|
||||
|
||||
mw = self._middleware()
|
||||
payload = "r" * 5000
|
||||
args = {"description": "d", "path": self.PATH, "content": payload}
|
||||
ai = AIMessage(
|
||||
content=[{"type": "function_call", "id": "fc_1", "call_id": "call-1", "name": "write_file", "arguments": json.dumps(args), "status": "completed"}],
|
||||
tool_calls=[{"name": "write_file", "id": "call-1", "args": dict(args)}],
|
||||
response_metadata={"output_version": "responses/v1"},
|
||||
)
|
||||
blocked = mw.wrap_tool_call(_make_request("write_file", dict(args), [HumanMessage(content="go"), ai]), MagicMock())
|
||||
request = self._model_request([HumanMessage(content="go"), ai, blocked])
|
||||
handler = MagicMock(return_value=AIMessage(content="ok"))
|
||||
|
||||
mw.wrap_model_call(request, handler)
|
||||
|
||||
items = _construct_responses_api_input(self._captured(handler).messages[1:2])
|
||||
calls = [item for item in items if item.get("type") == "function_call"]
|
||||
assert len(calls) == 1
|
||||
assert calls[0]["id"] == "fc_1"
|
||||
assert json.loads(calls[0]["arguments"])["content"].startswith("[payload elided: 5000 chars")
|
||||
assert payload not in json.dumps(items, ensure_ascii=False)
|
||||
|
||||
def _successful_write(self, tool_call_id, path="/mnt/user-data/outputs/other.md", payload="s" * 5000):
|
||||
ai = AIMessage(content="", tool_calls=[{"name": "write_file", "id": tool_call_id, "args": {"description": "d", "path": path, "content": payload}}])
|
||||
return ai, ToolMessage(content="OK", tool_call_id=tool_call_id, name="write_file")
|
||||
|
||||
@pytest.mark.parametrize("success_first", [True, False], ids=["success-before-block", "block-before-success"])
|
||||
def test_reused_call_id_only_elides_the_blocked_occurrence(self, success_first):
|
||||
"""Tool-call ids repeat across turns; pairing is per occurrence, not per id (review on #5329)."""
|
||||
import json
|
||||
|
||||
from langchain_openai.chat_models.base import _convert_message_to_dict
|
||||
|
||||
mw = self._middleware()
|
||||
ok_ai, ok_tool = self._successful_write("call-1")
|
||||
blocked_ai, blocked = self._blocked_turn(mw, "write_file", {"description": "d", "path": self.PATH, "content": "b" * 5000}, tool_call_id="call-1")
|
||||
turns = [ok_ai, ok_tool, blocked_ai, blocked] if success_first else [blocked_ai, blocked, ok_ai, ok_tool]
|
||||
request = self._model_request([HumanMessage(content="go"), *turns])
|
||||
handler = MagicMock(return_value=AIMessage(content="ok"))
|
||||
|
||||
mw.wrap_model_call(request, handler)
|
||||
|
||||
captured = self._captured(handler).messages
|
||||
ok_index, blocked_index = (1, 3) if success_first else (3, 1)
|
||||
assert captured[ok_index] is ok_ai
|
||||
assert captured[blocked_index].tool_calls[0]["args"]["content"].startswith("[payload elided: 5000 chars")
|
||||
ok_wire = json.loads(_convert_message_to_dict(captured[ok_index])["tool_calls"][0]["function"]["arguments"])
|
||||
blocked_wire = json.loads(_convert_message_to_dict(captured[blocked_index])["tool_calls"][0]["function"]["arguments"])
|
||||
assert ok_wire["content"] == "s" * 5000
|
||||
assert blocked_wire["content"].startswith("[payload elided")
|
||||
|
||||
def test_chained_responses_request_replays_the_rewritten_history(self):
|
||||
"""With use_previous_response_id the adapter must not chain past the elided call (review on #5329)."""
|
||||
import json
|
||||
|
||||
from langchain_openai import ChatOpenAI
|
||||
|
||||
mw = self._middleware()
|
||||
payload = "c" * 5000
|
||||
args = {"description": "d", "path": self.PATH, "content": payload}
|
||||
ai = AIMessage(
|
||||
content=[{"type": "function_call", "id": "fc_1", "call_id": "call-1", "name": "write_file", "arguments": json.dumps(args), "status": "completed"}],
|
||||
tool_calls=[{"name": "write_file", "id": "call-1", "args": dict(args)}],
|
||||
response_metadata={"id": "resp_blocked", "output_version": "responses/v1"},
|
||||
)
|
||||
blocked = mw.wrap_tool_call(_make_request("write_file", dict(args), [HumanMessage(content="go"), ai]), MagicMock())
|
||||
request = self._model_request([HumanMessage(content="go"), ai, blocked])
|
||||
handler = MagicMock(return_value=AIMessage(content="ok"))
|
||||
model = ChatOpenAI(model="gpt-4.1", api_key="test-key", use_responses_api=True, use_previous_response_id=True)
|
||||
|
||||
mw.wrap_model_call(request, handler)
|
||||
|
||||
leaked = model._get_request_payload(request.messages)
|
||||
assert leaked["previous_response_id"] == "resp_blocked"
|
||||
sent = model._get_request_payload(self._captured(handler).messages)
|
||||
assert "previous_response_id" not in sent
|
||||
calls = [item for item in sent["input"] if item.get("type") == "function_call"]
|
||||
assert len(calls) == 1
|
||||
assert json.loads(calls[0]["arguments"])["content"].startswith("[payload elided: 5000 chars")
|
||||
assert payload not in json.dumps(sent, ensure_ascii=False)
|
||||
|
||||
387
backend/tests/test_tool_call_args.py
Normal file
387
backend/tests/test_tool_call_args.py
Normal file
@ -0,0 +1,387 @@
|
||||
"""Tests for the shared model-bound tool-call argument rewriter (``tool_call_args``)."""
|
||||
|
||||
import json
|
||||
|
||||
from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, ToolMessage
|
||||
|
||||
from deerflow.agents.middlewares.tool_call_args import rewrite_messages_tool_call_args, rewrite_tool_call_args
|
||||
|
||||
ARGS = {"path": "/mnt/user-data/outputs/report.md", "content": "x" * 50}
|
||||
NEW_ARGS = {"path": "/mnt/user-data/outputs/report.md", "content": "[elided]"}
|
||||
|
||||
|
||||
def _full_surface_message(call_id="call-1"):
|
||||
"""An AIMessage carrying the same call on every surface a provider adapter may read."""
|
||||
return AIMessage(
|
||||
content=[
|
||||
{"type": "text", "text": "writing"},
|
||||
{"type": "tool_use", "id": call_id, "name": "write_file", "input": dict(ARGS), "partial_json": json.dumps(ARGS)},
|
||||
],
|
||||
tool_calls=[{"name": "write_file", "id": call_id, "args": dict(ARGS)}],
|
||||
additional_kwargs={"tool_calls": [{"id": call_id, "type": "function", "function": {"name": "write_file", "arguments": json.dumps(ARGS)}}]},
|
||||
)
|
||||
|
||||
|
||||
class TestRewriteToolCallArgs:
|
||||
def test_no_matching_id_returns_same_object(self):
|
||||
message = _full_surface_message()
|
||||
assert rewrite_tool_call_args(message, {"other": NEW_ARGS}) is message
|
||||
assert rewrite_tool_call_args(message, {}) is message
|
||||
|
||||
def test_rewrites_every_surface_together(self):
|
||||
message = _full_surface_message()
|
||||
|
||||
rewritten = rewrite_tool_call_args(message, {"call-1": NEW_ARGS})
|
||||
|
||||
assert rewritten is not message
|
||||
assert rewritten.tool_calls[0]["args"] == NEW_ARGS
|
||||
assert rewritten.tool_calls[0]["name"] == "write_file"
|
||||
raw = rewritten.additional_kwargs["tool_calls"][0]
|
||||
assert json.loads(raw["function"]["arguments"]) == NEW_ARGS
|
||||
assert raw["function"]["name"] == "write_file"
|
||||
assert rewritten.content[0] == {"type": "text", "text": "writing"}
|
||||
assert rewritten.content[1] == {"type": "tool_use", "id": "call-1", "name": "write_file", "input": NEW_ARGS}
|
||||
assert "x" * 50 not in json.dumps(rewritten.model_dump(), ensure_ascii=False)
|
||||
|
||||
def test_original_message_is_never_mutated(self):
|
||||
message = _full_surface_message()
|
||||
|
||||
rewrite_tool_call_args(message, {"call-1": NEW_ARGS})
|
||||
|
||||
assert message.tool_calls[0]["args"] == ARGS
|
||||
assert message.content[1]["input"] == ARGS
|
||||
assert "partial_json" in message.content[1]
|
||||
assert json.loads(message.additional_kwargs["tool_calls"][0]["function"]["arguments"]) == ARGS
|
||||
|
||||
def test_untouched_sibling_calls_keep_identity(self):
|
||||
other = {"name": "bash", "id": "call-2", "args": {"command": "ls"}}
|
||||
message = AIMessage(content="", tool_calls=[{"name": "write_file", "id": "call-1", "args": dict(ARGS)}, other])
|
||||
|
||||
rewritten = rewrite_tool_call_args(message, {"call-1": NEW_ARGS})
|
||||
|
||||
# AIMessage validation copies tool-call dicts at construction, so identity is against the message's own list.
|
||||
assert rewritten.tool_calls[1] is message.tool_calls[1]
|
||||
assert rewritten.tool_calls[1]["args"] == other["args"]
|
||||
assert rewritten.tool_calls[0]["args"] == NEW_ARGS
|
||||
|
||||
def test_rewrites_chunk_surfaces(self):
|
||||
chunk = AIMessageChunk(content="", tool_call_chunks=[{"name": "write_file", "args": json.dumps(ARGS), "id": "call-1", "index": 0}])
|
||||
assert chunk.tool_calls[0]["args"] == ARGS
|
||||
|
||||
rewritten = rewrite_tool_call_args(chunk, {"call-1": NEW_ARGS})
|
||||
|
||||
assert rewritten.tool_calls[0]["args"] == NEW_ARGS
|
||||
assert json.loads(rewritten.tool_call_chunks[0]["args"]) == NEW_ARGS
|
||||
assert rewritten.tool_call_chunks[0]["index"] == 0
|
||||
assert chunk.tool_call_chunks[0]["args"] == json.dumps(ARGS)
|
||||
|
||||
def test_flattened_raw_provider_variants(self):
|
||||
message = AIMessage(
|
||||
content="",
|
||||
tool_calls=[{"name": "write_file", "id": "call-1", "args": dict(ARGS)}, {"name": "write_file", "id": "call-2", "args": dict(ARGS)}],
|
||||
additional_kwargs={
|
||||
"tool_calls": [
|
||||
{"id": "call-1", "name": "write_file", "arguments": json.dumps(ARGS)},
|
||||
{"id": "call-2", "name": "write_file", "args": dict(ARGS)},
|
||||
{"id": "call-3", "name": "write_file"},
|
||||
"not-a-dict",
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
rewritten = rewrite_tool_call_args(message, {"call-1": NEW_ARGS, "call-2": NEW_ARGS, "call-3": NEW_ARGS})
|
||||
|
||||
raw = rewritten.additional_kwargs["tool_calls"]
|
||||
assert json.loads(raw[0]["arguments"]) == NEW_ARGS
|
||||
assert raw[1]["args"] == NEW_ARGS
|
||||
assert raw[2] is message.additional_kwargs["tool_calls"][2]
|
||||
assert raw[3] == "not-a-dict"
|
||||
|
||||
def test_non_string_ids_never_match(self):
|
||||
message = AIMessage(
|
||||
content=[{"type": "tool_use", "id": ["list", "id"], "name": "write_file", "input": dict(ARGS)}],
|
||||
tool_calls=[{"name": "write_file", "id": "call-1", "args": dict(ARGS)}],
|
||||
additional_kwargs={"tool_calls": [{"id": {"dict": "id"}, "type": "function", "function": {"name": "write_file", "arguments": json.dumps(ARGS)}}]},
|
||||
)
|
||||
|
||||
rewritten = rewrite_tool_call_args(message, {"call-1": NEW_ARGS})
|
||||
|
||||
assert rewritten.tool_calls[0]["args"] == NEW_ARGS
|
||||
assert rewritten.content[0]["input"] == ARGS
|
||||
assert json.loads(rewritten.additional_kwargs["tool_calls"][0]["function"]["arguments"]) == ARGS
|
||||
|
||||
def test_result_is_deterministic(self):
|
||||
message = _full_surface_message()
|
||||
first = rewrite_tool_call_args(message, {"call-1": NEW_ARGS})
|
||||
second = rewrite_tool_call_args(message, {"call-1": NEW_ARGS})
|
||||
assert first.model_dump() == second.model_dump()
|
||||
|
||||
|
||||
class TestRewriteMessagesToolCallArgs:
|
||||
def test_returns_none_when_nothing_replaced(self):
|
||||
messages = [HumanMessage(content="go"), _full_surface_message(), ToolMessage(content="ok", tool_call_id="call-1", name="write_file")]
|
||||
assert rewrite_messages_tool_call_args(messages, lambda _message, _tool_call: None) is None
|
||||
assert rewrite_messages_tool_call_args([], lambda _message, _tool_call: NEW_ARGS) is None
|
||||
|
||||
def test_selector_sees_message_and_call_and_untouched_messages_keep_identity(self):
|
||||
human = HumanMessage(content="go")
|
||||
target = _full_surface_message("call-1")
|
||||
other = AIMessage(content="", tool_calls=[{"name": "bash", "id": "call-2", "args": {"command": "ls"}}])
|
||||
tool = ToolMessage(content="ok", tool_call_id="call-1", name="write_file")
|
||||
seen = []
|
||||
|
||||
def replacement_for(message, tool_call):
|
||||
seen.append((message, tool_call["id"]))
|
||||
return NEW_ARGS if tool_call["name"] == "write_file" else None
|
||||
|
||||
rewritten = rewrite_messages_tool_call_args([human, target, other, tool], replacement_for)
|
||||
|
||||
assert seen == [(target, "call-1"), (other, "call-2")]
|
||||
assert rewritten[0] is human
|
||||
assert rewritten[1] is not target
|
||||
assert rewritten[1].tool_calls[0]["args"] == NEW_ARGS
|
||||
assert rewritten[2] is other
|
||||
assert rewritten[3] is tool
|
||||
assert target.tool_calls[0]["args"] == ARGS
|
||||
|
||||
def test_calls_without_a_string_id_are_not_offered(self):
|
||||
message = AIMessage(content="", tool_calls=[{"name": "write_file", "id": None, "args": dict(ARGS)}])
|
||||
offered = []
|
||||
|
||||
assert rewrite_messages_tool_call_args([message], lambda _m, tc: offered.append(tc) or NEW_ARGS) is None
|
||||
assert offered == []
|
||||
|
||||
|
||||
RESPONSES_V1_BLOCK = {"type": "function_call", "id": "fc_1", "call_id": "call-1", "name": "write_file", "arguments": json.dumps(ARGS), "status": "completed"}
|
||||
V1_BLOCK = {"type": "tool_call", "id": "call-1", "name": "write_file", "args": dict(ARGS), "extras": {"item_id": "fc_1", "arguments": json.dumps(ARGS), "status": "completed"}}
|
||||
V1_CHUNK_BLOCK = {"type": "tool_call_chunk", "id": "call-1", "name": "write_file", "args": json.dumps(ARGS), "index": 0, "extras": {"item_id": "fc_1"}}
|
||||
|
||||
|
||||
def _responses_v1_message():
|
||||
return AIMessage(content=[{"type": "text", "text": "writing"}, dict(RESPONSES_V1_BLOCK)], tool_calls=[{"name": "write_file", "id": "call-1", "args": dict(ARGS)}], response_metadata={"output_version": "responses/v1"})
|
||||
|
||||
|
||||
def _v1_message():
|
||||
return AIMessage(content=[{"type": "text", "text": "writing"}, {**V1_BLOCK, "extras": dict(V1_BLOCK["extras"])}], tool_calls=[{"name": "write_file", "id": "call-1", "args": dict(ARGS)}], response_metadata={"output_version": "v1"})
|
||||
|
||||
|
||||
class TestContentBlockVariants:
|
||||
"""Every content-block dialect that carries its own copy of the arguments is rewritten, ids preserved."""
|
||||
|
||||
def test_responses_function_call_block_matched_by_call_id_keeps_item_id(self):
|
||||
message = _responses_v1_message()
|
||||
|
||||
rewritten = rewrite_tool_call_args(message, {"call-1": NEW_ARGS})
|
||||
|
||||
block = rewritten.content[1]
|
||||
assert json.loads(block["arguments"]) == NEW_ARGS
|
||||
assert block["id"] == "fc_1"
|
||||
assert block["call_id"] == "call-1"
|
||||
assert block["status"] == "completed"
|
||||
assert rewritten.content[0] is message.content[0]
|
||||
assert json.loads(message.content[1]["arguments"]) == ARGS
|
||||
|
||||
def test_responses_function_call_block_ignores_item_id_as_match_key(self):
|
||||
message = _responses_v1_message()
|
||||
assert rewrite_tool_call_args(message, {"fc_1": NEW_ARGS}) is message
|
||||
|
||||
def test_v1_tool_call_block_rewrites_args_and_extras_arguments(self):
|
||||
message = _v1_message()
|
||||
|
||||
rewritten = rewrite_tool_call_args(message, {"call-1": NEW_ARGS})
|
||||
|
||||
block = rewritten.content[1]
|
||||
assert block["args"] == NEW_ARGS
|
||||
assert json.loads(block["extras"]["arguments"]) == NEW_ARGS
|
||||
assert block["extras"]["item_id"] == "fc_1"
|
||||
assert block["extras"]["status"] == "completed"
|
||||
assert message.content[1]["args"] == ARGS
|
||||
assert json.loads(message.content[1]["extras"]["arguments"]) == ARGS
|
||||
|
||||
def test_v1_tool_call_block_without_extras_arguments_gets_no_extras_entry(self):
|
||||
block = {"type": "tool_call", "id": "call-1", "name": "write_file", "args": dict(ARGS), "extras": {"item_id": "fc_1"}}
|
||||
message = AIMessage(content=[block], tool_calls=[{"name": "write_file", "id": "call-1", "args": dict(ARGS)}])
|
||||
|
||||
rewritten = rewrite_tool_call_args(message, {"call-1": NEW_ARGS})
|
||||
|
||||
assert rewritten.content[0]["args"] == NEW_ARGS
|
||||
assert rewritten.content[0]["extras"] == {"item_id": "fc_1"}
|
||||
|
||||
def test_v1_tool_call_chunk_block_rewrites_serialized_args(self):
|
||||
chunk = AIMessageChunk(content=[dict(V1_CHUNK_BLOCK)], tool_call_chunks=[{"name": "write_file", "args": json.dumps(ARGS), "id": "call-1", "index": 0}])
|
||||
|
||||
rewritten = rewrite_tool_call_args(chunk, {"call-1": NEW_ARGS})
|
||||
|
||||
assert json.loads(rewritten.content[0]["args"]) == NEW_ARGS
|
||||
assert rewritten.content[0]["extras"] == {"item_id": "fc_1"}
|
||||
assert rewritten.content[0]["index"] == 0
|
||||
assert json.loads(rewritten.tool_call_chunks[0]["args"]) == NEW_ARGS
|
||||
assert json.loads(chunk.content[0]["args"]) == ARGS
|
||||
|
||||
def test_unrelated_block_types_pass_through_by_identity(self):
|
||||
reasoning = {"type": "reasoning", "id": "rs_1", "summary": []}
|
||||
message = AIMessage(content=[reasoning, dict(RESPONSES_V1_BLOCK)], tool_calls=[{"name": "write_file", "id": "call-1", "args": dict(ARGS)}])
|
||||
|
||||
rewritten = rewrite_tool_call_args(message, {"call-1": NEW_ARGS})
|
||||
|
||||
assert rewritten.content[0] is message.content[0]
|
||||
|
||||
|
||||
class TestProviderSerializers:
|
||||
"""Lock the rewrite against the real adapter request builders: the payload must not reach the wire."""
|
||||
|
||||
PAYLOAD = ARGS["content"]
|
||||
|
||||
@staticmethod
|
||||
def _responses_input(message):
|
||||
from langchain_openai.chat_models.base import _construct_responses_api_input
|
||||
|
||||
return _construct_responses_api_input([message])
|
||||
|
||||
def _function_calls(self, message):
|
||||
items = self._responses_input(message)
|
||||
assert self.PAYLOAD not in json.dumps(items, ensure_ascii=False)
|
||||
return [item for item in items if item.get("type") == "function_call"]
|
||||
|
||||
def test_responses_v1_content_sends_rewritten_arguments_once(self):
|
||||
calls = self._function_calls(rewrite_tool_call_args(_responses_v1_message(), {"call-1": NEW_ARGS}))
|
||||
|
||||
assert len(calls) == 1
|
||||
assert json.loads(calls[0]["arguments"]) == NEW_ARGS
|
||||
assert calls[0]["call_id"] == "call-1"
|
||||
assert calls[0]["id"] == "fc_1"
|
||||
|
||||
def test_v1_content_sends_rewritten_arguments_once(self):
|
||||
calls = self._function_calls(rewrite_tool_call_args(_v1_message(), {"call-1": NEW_ARGS}))
|
||||
|
||||
assert len(calls) == 1
|
||||
assert json.loads(calls[0]["arguments"]) == NEW_ARGS
|
||||
assert calls[0]["call_id"] == "call-1"
|
||||
assert calls[0]["id"] == "fc_1"
|
||||
|
||||
def test_v0_responses_message_sends_rewritten_arguments_with_item_id(self):
|
||||
message = AIMessage(
|
||||
content=[{"type": "text", "text": "writing"}],
|
||||
tool_calls=[{"name": "write_file", "id": "call-1", "args": dict(ARGS)}],
|
||||
additional_kwargs={"__openai_function_call_ids__": {"call-1": "fc_1"}},
|
||||
)
|
||||
|
||||
calls = self._function_calls(rewrite_tool_call_args(message, {"call-1": NEW_ARGS}))
|
||||
|
||||
assert len(calls) == 1
|
||||
assert json.loads(calls[0]["arguments"]) == NEW_ARGS
|
||||
assert calls[0]["id"] == "fc_1"
|
||||
|
||||
def test_unrewritten_responses_message_still_carries_payload(self):
|
||||
"""Sanity check that the probe can see the payload at all."""
|
||||
items = self._responses_input(_responses_v1_message())
|
||||
assert self.PAYLOAD in json.dumps(items, ensure_ascii=False)
|
||||
|
||||
def test_chat_completions_payload_uses_rewritten_arguments(self):
|
||||
from langchain_openai.chat_models.base import _convert_message_to_dict
|
||||
|
||||
payload = _convert_message_to_dict(rewrite_tool_call_args(_full_surface_message(), {"call-1": NEW_ARGS}))
|
||||
|
||||
assert json.loads(payload["tool_calls"][0]["function"]["arguments"]) == NEW_ARGS
|
||||
assert self.PAYLOAD not in json.dumps(payload, ensure_ascii=False)
|
||||
|
||||
def test_anthropic_native_tool_use_payload_uses_rewritten_input(self):
|
||||
from langchain_anthropic.chat_models import _format_messages
|
||||
|
||||
_system, formatted = _format_messages([rewrite_tool_call_args(_full_surface_message(), {"call-1": NEW_ARGS})])
|
||||
|
||||
tool_use = [block for block in formatted[0]["content"] if block["type"] == "tool_use"]
|
||||
assert len(tool_use) == 1
|
||||
assert tool_use[0]["input"] == NEW_ARGS
|
||||
assert self.PAYLOAD not in json.dumps(formatted, ensure_ascii=False)
|
||||
|
||||
def test_anthropic_v1_content_payload_uses_rewritten_input(self):
|
||||
from langchain_anthropic._compat import _convert_from_v1_to_anthropic
|
||||
from langchain_anthropic.chat_models import _format_messages
|
||||
|
||||
rewritten = rewrite_tool_call_args(_v1_message(), {"call-1": NEW_ARGS})
|
||||
# Mirrors ChatAnthropic._get_request_payload's v1 translation step.
|
||||
tcs = [{"type": "tool_call", "name": tc["name"], "args": tc["args"], "id": tc.get("id")} for tc in rewritten.tool_calls]
|
||||
translated = rewritten.model_copy(update={"content": _convert_from_v1_to_anthropic(rewritten.content, tcs, "anthropic")})
|
||||
|
||||
_system, formatted = _format_messages([translated])
|
||||
|
||||
tool_use = [block for block in formatted[0]["content"] if block["type"] == "tool_use"]
|
||||
assert len(tool_use) == 1
|
||||
assert tool_use[0]["input"] == NEW_ARGS
|
||||
assert self.PAYLOAD not in json.dumps(formatted, ensure_ascii=False)
|
||||
|
||||
|
||||
class TestResponseChainInvalidation:
|
||||
"""A rewritten history must be replayed, never chained to the original server-side copy."""
|
||||
|
||||
PAYLOAD = ARGS["content"]
|
||||
|
||||
@staticmethod
|
||||
def _chained_history(rewritten_call=True):
|
||||
first = AIMessage(content=[{"type": "text", "text": "earlier"}], response_metadata={"id": "resp_a", "model_name": "gpt-x"})
|
||||
call = _responses_v1_message()
|
||||
call = call.model_copy(update={"response_metadata": {**call.response_metadata, "id": "resp_b", "model_name": "gpt-x"}})
|
||||
tool = ToolMessage(content="Error: blocked", tool_call_id="call-1", name="write_file", status="error")
|
||||
later = AIMessage(content=[{"type": "text", "text": "later"}], response_metadata={"id": "resp_c"})
|
||||
return [HumanMessage(content="go"), first, call, tool, later]
|
||||
|
||||
def test_rewrite_drops_resp_ids_from_every_ai_message(self):
|
||||
messages = self._chained_history()
|
||||
|
||||
rewritten = rewrite_messages_tool_call_args(messages, lambda _m, tc: NEW_ARGS if tc["id"] == "call-1" else None)
|
||||
|
||||
assert [type(m) for m in rewritten] == [type(m) for m in messages]
|
||||
for index in (1, 2, 4):
|
||||
assert "id" not in rewritten[index].response_metadata
|
||||
assert rewritten[index] is not messages[index]
|
||||
assert rewritten[1].response_metadata == {"model_name": "gpt-x"}
|
||||
assert rewritten[2].tool_calls[0]["args"] == NEW_ARGS
|
||||
assert rewritten[0] is messages[0]
|
||||
assert rewritten[3] is messages[3]
|
||||
# Stored history keeps its chain ids and arguments.
|
||||
assert messages[1].response_metadata["id"] == "resp_a"
|
||||
assert messages[2].response_metadata["id"] == "resp_b"
|
||||
assert messages[4].response_metadata["id"] == "resp_c"
|
||||
assert messages[2].tool_calls[0]["args"] == ARGS
|
||||
|
||||
def test_non_resp_ids_are_left_alone(self):
|
||||
anthropic_style = AIMessage(content="earlier", response_metadata={"id": "msg_01", "model": "claude"})
|
||||
messages = [anthropic_style, _full_surface_message(), ToolMessage(content="ok", tool_call_id="call-1", name="write_file")]
|
||||
|
||||
rewritten = rewrite_messages_tool_call_args(messages, lambda _m, tc: NEW_ARGS)
|
||||
|
||||
assert rewritten[0] is anthropic_style
|
||||
assert rewritten[0].response_metadata["id"] == "msg_01"
|
||||
|
||||
def test_no_rewrite_keeps_chain_ids(self):
|
||||
messages = self._chained_history()
|
||||
assert rewrite_messages_tool_call_args(messages, lambda _m, _tc: None) is None
|
||||
assert messages[2].response_metadata["id"] == "resp_b"
|
||||
|
||||
@staticmethod
|
||||
def _chained_model():
|
||||
from langchain_openai import ChatOpenAI
|
||||
|
||||
return ChatOpenAI(model="gpt-4.1", api_key="test-key", use_responses_api=True, use_previous_response_id=True)
|
||||
|
||||
def test_unrewritten_history_chains_and_never_sends_the_call(self):
|
||||
"""Documents the leak: with chaining on, the adapter sends only the tail after the last resp_ id."""
|
||||
payload = self._chained_model()._get_request_payload(self._chained_history()[:4])
|
||||
|
||||
assert payload["previous_response_id"] == "resp_b"
|
||||
assert [item["type"] for item in payload["input"]] == ["function_call_output"]
|
||||
|
||||
def test_rewritten_history_is_replayed_with_rewritten_arguments(self):
|
||||
messages = self._chained_history()[:4]
|
||||
rewritten = rewrite_messages_tool_call_args(messages, lambda _m, tc: NEW_ARGS if tc["id"] == "call-1" else None)
|
||||
|
||||
payload = self._chained_model()._get_request_payload(rewritten)
|
||||
|
||||
assert "previous_response_id" not in payload
|
||||
calls = [item for item in payload["input"] if item.get("type") == "function_call"]
|
||||
assert len(calls) == 1
|
||||
assert json.loads(calls[0]["arguments"]) == NEW_ARGS
|
||||
assert calls[0]["id"] == "fc_1"
|
||||
assert any(item.get("type") == "function_call_output" for item in payload["input"])
|
||||
assert self.PAYLOAD not in json.dumps(payload, ensure_ascii=False)
|
||||
@ -1245,3 +1245,17 @@ def test_subagent_summarization_fires_mid_run_and_produces_usable_result(monkeyp
|
||||
ai_finals = [m for m in final_messages if isinstance(m, AIMessage)]
|
||||
assert ai_finals, "the run must produce a final AIMessage after compaction"
|
||||
assert ai_finals[-1].content == "final answer after compaction"
|
||||
|
||||
|
||||
def test_build_lead_runtime_middlewares_passes_read_before_write_config():
|
||||
"""The gate's model-bound payload elision is configured from app_config.read_before_write."""
|
||||
from deerflow.agents.middlewares.read_before_write_middleware import ReadBeforeWriteMiddleware
|
||||
from deerflow.config.read_before_write_config import ReadBeforeWriteConfig
|
||||
|
||||
app_config = _make_app_config().model_copy(update={"read_before_write": ReadBeforeWriteConfig(elide_min_chars=321)})
|
||||
middlewares = build_lead_runtime_middlewares(app_config=app_config)
|
||||
|
||||
gates = [m for m in middlewares if isinstance(m, ReadBeforeWriteMiddleware)]
|
||||
assert len(gates) == 1
|
||||
# Only the wired value is under test; the full policy identity is covered by the middleware's own tests.
|
||||
assert gates[0].release_policy_parameters()["config"]["elide_min_chars"] == 321
|
||||
|
||||
@ -20,7 +20,7 @@
|
||||
# ============================================================================
|
||||
# Bump this number when the config schema changes.
|
||||
# Run `make config-upgrade` to merge new fields into your local config.yaml.
|
||||
config_version: 40
|
||||
config_version: 41
|
||||
|
||||
# ============================================================================
|
||||
# Logging
|
||||
@ -1290,6 +1290,15 @@ loop_detection:
|
||||
|
||||
read_before_write:
|
||||
enabled: true
|
||||
# A blocked call never ran and must be re-issued after a re-read, so its
|
||||
# payload (write_file content, str_replace old_str/new_str) is dead weight in
|
||||
# every later model call. Replace it with a short placeholder in model-bound
|
||||
# requests; stored history, receipts, and the run journal keep the original.
|
||||
elide_blocked_payloads: true
|
||||
# Only elide payload fields at least this many characters long (a character
|
||||
# count, not tokens: CJK text costs ~3-4x more per character than ASCII).
|
||||
# 0 elides every non-empty payload.
|
||||
elide_min_chars: 2000
|
||||
|
||||
# ============================================================================
|
||||
# Provider Safety Termination Configuration
|
||||
|
||||
@ -131,7 +131,7 @@ they resolve from the `secrets` map):
|
||||
|
||||
```yaml
|
||||
config: |
|
||||
config_version: 40
|
||||
config_version: 41
|
||||
models:
|
||||
- name: gpt-4
|
||||
use: langchain_openai:ChatOpenAI
|
||||
|
||||
@ -249,7 +249,7 @@ ingress:
|
||||
# -- DeerFlow config.yaml content. Secrets MUST stay as $VAR references — never
|
||||
# inline literal secret values here. The default enables provisioner sandbox.
|
||||
config: |
|
||||
config_version: 40
|
||||
config_version: 41
|
||||
log_level: info
|
||||
|
||||
models: []
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user