mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-25 05:56:18 +00:00
fix(middlewares): end length-capped turns cleanly, prevent todo re-engagement, annotate write_file budget (#5569)
* fix(middlewares): end length-capped turns cleanly, prevent todo re-engagement, annotate write_file budget When a model hits its per-response output cap (finish_reason=length) while emitting a write_file tool call, ModelLengthFinishReasonMiddleware suppresses the truncated call and stamps model_length_termination. TodoMiddleware must not re-engage (jump_to=model) on such a capped turn -- doing so re-emits the same oversized call into the same cap, producing up to 3 futile responses with junk fragments instead of a clean truncation notice. Changes: - TodoMiddleware.after_model: skip completion reminder jump when additional_kwargs.model_length_termination is present (follows the existing deerflow_error_fallback precedent). - ModelLengthFinishReasonMiddleware: always append the length notice when tool calls were suppressed, even when partial text survived (collapses the visible-content ternary). Fixes a latent bug in append_visible_text that silently dropped string content. - tools.get_available_tools: annotate write_file's model-visible description with the model's configured max_tokens output budget. Guarded extraction safely handles missing or non-numeric tokens, and the tool is cloned via model_copy to keep module-level singletons immutable across assemblies and prevent guidance leakage to unbudgeted models. - release_policy_parameters() updated for both middlewares. - AGENTS.md chain entries (#20, #35) and module docstrings updated within AG002 guidance limits. - Tests: 8 new/focused unit tests + 1 updated pin + 1 real create_agent() integration test reproducing the incident (thread b1723286). * fix(tools): use effective model cap for write_file guidance --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
Dan Caldr
GitHub
parent
c668716737
commit
19266a5eac
@ -1640,6 +1640,8 @@ an explicit **Load full file** action before fetching the remainder or mounting
|
||||
the full code editor. Active HTML, XHTML, and SVG artifacts remain forced
|
||||
downloads at the Gateway boundary.
|
||||
|
||||
The `write_file` guidance reflects the active model's output token limit, including custom-agent and thinking-mode overrides. For longer documents, the agent is guided to write sections with `append=True`; models without a known limit receive no numeric budget hint.
|
||||
|
||||
With `AioSandboxProvider`, shell execution runs inside isolated containers. With `LocalSandboxProvider`, file tools still map to per-thread directories on the host, but host `bash` is disabled by default because it is not a secure isolation boundary. Re-enable host bash only for fully trusted local workflows. Host bash commands have a wall-clock timeout, and long-lived processes should be started in the background with output redirected to a workspace log. On Windows, Git Bash/MSYS argument-conversion exclusions are limited to safe non-root virtual path prefixes, so host-native CLI launchers retain their normal MSYS compatibility. When the local sandbox falls back to PowerShell, it captures output as UTF-8 so CJK text does not depend on the Gateway host locale.
|
||||
|
||||
Docker AIO sandboxes default to their existing open egress behavior for
|
||||
|
||||
@ -1057,7 +1057,8 @@ def _assemble_lead_agent(config: RunnableConfig, *, app_config: AppConfig) -> Le
|
||||
enabled=skill_search_enabled,
|
||||
container_base_path=container_base_path,
|
||||
)
|
||||
raw_tools = get_available_tools(model_name=model_name, subagent_enabled=subagent_enabled, app_config=resolved_app_config) + [setup_agent]
|
||||
chat_model = create_chat_model(name=model_name, thinking_enabled=thinking_enabled, app_config=resolved_app_config, attach_tracing=False)
|
||||
raw_tools = get_available_tools(model_name=model_name, subagent_enabled=subagent_enabled, app_config=resolved_app_config, chat_model=chat_model) + [setup_agent]
|
||||
configured_tools = raw_tools
|
||||
configured_tools = [tool for tool in configured_tools if tool.name not in interaction_policy.disabled_tool_names]
|
||||
authorization_candidates = [*configured_tools]
|
||||
@ -1111,7 +1112,7 @@ def _assemble_lead_agent(config: RunnableConfig, *, app_config: AppConfig) -> Le
|
||||
memory_enabled=memory_enabled,
|
||||
)
|
||||
graph = create_agent(
|
||||
model=create_chat_model(name=model_name, thinking_enabled=thinking_enabled, app_config=resolved_app_config, attach_tracing=False),
|
||||
model=chat_model,
|
||||
tools=final_tools,
|
||||
middleware=normalize_middleware_state_schemas(middlewares, mode),
|
||||
system_prompt=system_prompt,
|
||||
@ -1177,7 +1178,8 @@ def _assemble_lead_agent(config: RunnableConfig, *, app_config: AppConfig) -> Le
|
||||
channel_name = cfg.get("channel_name")
|
||||
is_webhook_channel = channel_name in _WEBHOOK_CHANNELS
|
||||
extra_tools = [update_agent] if agent_name and not is_webhook_channel else []
|
||||
# Default lead agent (unchanged behavior)
|
||||
# Resolve the model once so tool guidance uses the same effective settings.
|
||||
chat_model = create_chat_model(name=model_name, thinking_enabled=thinking_enabled, reasoning_effort=reasoning_effort, app_config=resolved_app_config, attach_tracing=False, model_overrides=agent_model_overrides)
|
||||
raw_tools = get_available_tools(
|
||||
model_name=model_name,
|
||||
groups=agent_config.tool_groups if agent_config else None,
|
||||
@ -1185,6 +1187,7 @@ def _assemble_lead_agent(config: RunnableConfig, *, app_config: AppConfig) -> Le
|
||||
subagent_enabled=subagent_enabled,
|
||||
include_conversation_reader=callable(cfg.get(CONVERSATION_READER_CONTEXT_KEY)) and not bool(cfg.get("is_subagent")),
|
||||
app_config=resolved_app_config,
|
||||
chat_model=chat_model,
|
||||
)
|
||||
configured_tools = raw_tools + extra_tools
|
||||
configured_tools = [tool for tool in configured_tools if tool.name not in interaction_policy.disabled_tool_names]
|
||||
@ -1241,7 +1244,7 @@ def _assemble_lead_agent(config: RunnableConfig, *, app_config: AppConfig) -> Le
|
||||
memory_enabled=memory_enabled,
|
||||
)
|
||||
graph = create_agent(
|
||||
model=create_chat_model(name=model_name, thinking_enabled=thinking_enabled, reasoning_effort=reasoning_effort, app_config=resolved_app_config, attach_tracing=False, model_overrides=agent_model_overrides),
|
||||
model=chat_model,
|
||||
tools=final_tools,
|
||||
middleware=normalize_middleware_state_schemas(middlewares, mode),
|
||||
system_prompt=system_prompt,
|
||||
|
||||
@ -103,7 +103,7 @@ Before changing a later authorization phase, read the [authorization RFC](../../
|
||||
17. **SkillToolPolicyMiddleware** - Applies `allowed-tools` only after real activation; passive enabled skills and a custom agent's configured skill allowlist do not clamp the lead toolset. A run-scoped slash activation is authoritative and suppresses `skill_context` as a policy source, so reading another skill cannot widen the explicit skill's tools; without slash activation, skills captured after configured `read_file` loads retain the existing union semantics. The middleware filters model-visible schemas and blocks unauthorized execution, resolving canonical paths against the live enabled/agent-allowed registry on every model call, then stores a versioned, JSON-safe, middleware-token-bound decision signed by policy source plus active paths in run context for the resulting tool calls to reuse. The next model call always refreshes it, and malformed, foreign, stale, or unmatched decisions fall back to live resolution. `tool_search` and `describe_skill` remain framework-safe discovery tools under a restrictive policy; they may reveal or promote metadata, but a deferred business tool must still be declared by the active policy before its schema or execution can survive the policy middleware. The decision's owner token is authorization-sensitive, so its reserved context key is owned by `runtime.secret_context` and included in `REDACTED_CONTEXT_KEYS` for observable and persisted context copies. Registry load failures and a non-empty active set with no authorized skill fail closed to framework-safe tools; an individual stale path is skipped only when at least one valid active skill remains. This is best-effort behavioral scoping rather than a hard security boundary: alternate loads such as `bash cat` are not captured, and bounded autonomous `skill_context` can evict old entries. `task` is not framework-exempt, so a restricted skill cannot delegate around its policy. The middleware must remain immediately after `SkillActivationMiddleware` (which publishes the slash source through `runtime.secret_context`'s public path helpers authenticated by a required token shared only within the assembled middleware chain) and immediately before `DurableContextMiddleware`; assembly and compiled-graph tests pin ordering, token sharing, schema filtering, and execution blocking.
|
||||
18. **DurableContextMiddleware** - Captures `task` delegations into `ThreadState.delegations` (including in-progress dispatches and terminal result summaries) and loaded skill-file references (name/path/description, parsed in-memory - not the body) into `ThreadState.skill_context` before summarization can compact the paired tool-call/result messages, then projects durable context into each model request. Static authority rules are injected as a `SystemMessage`; untrusted field values (`summary_text`, delegation results, skill descriptions) are injected separately as a hidden `HumanMessage` data block so compressed history, delegated work, and which skills are active stay visible without being stored as `messages` or promoted to system-role instructions. `build_subagent_runtime_middlewares` also attaches this middleware immediately before subagent summarization so a compacted `summary_text` is projected ahead of a preserved assistant/tool tail instead of leaving strict providers with an assistant-first request.
|
||||
19. **SummarizationMiddleware** - *(optional, if enabled)* Compacts near token limits; memory flush follows runtime policy, while manual compaction trusts the checkpoint-bound agent, so opt-out never captures removed turns. It preserves the latest real user request by ID and tagged DynamicContext reminders while allowing stale ID-swap peers into the summary. Moving the cutoff backward can retain old AI/tool turns and make first-turn compaction a no-op.
|
||||
20. **TodoListMiddleware** - *(optional, if `is_plan_mode`)* Task tracking with the `write_todos` tool
|
||||
20. **TodoListMiddleware** - *(optional, if `is_plan_mode`)* Task tracking with `write_todos`. Skips reminders on `model_length_termination` so capped turns end cleanly.
|
||||
21. **TokenUsageMiddleware** - *(optional, if `token_usage.enabled`)* Records token usage metrics; subagent usage is read from terminal `ToolMessage.additional_kwargs` in the current run and merged back into the dispatching AIMessage by message position. The same state update marks the ToolMessage with `subagent_token_usage_attributed=true`, so checkpoint replay or middleware re-entry cannot add the cumulative snapshot twice; missing/malformed usage or a result with no matching dispatch remains unmarked and retryable.
|
||||
22. **TitleMiddleware** - Auto-generates the thread title after the first complete exchange and normalizes structured message content before prompting the title model. If a first-turn run is interrupted before this middleware can write a title, `runtime/runs/worker.py` keeps the run in a finalizing state, persists a local fallback title from the latest checkpoint or original run input, and then syncs it to `threads_meta.display_name`. Replacement runs admitted by `multitask_strategy="interrupt"` / `"rollback"` wait for older same-thread finalization before entering the graph; the interrupted run only skips the fallback title write once a later run has started and may have advanced the checkpoint.
|
||||
23. **MemoryMiddleware** - Queues conversations for async memory update (filters to user + final AI responses); captures the runtime-resolved user so standalone LangGraph Server reads and writes stay in the same bucket
|
||||
@ -123,6 +123,6 @@ Before changing a later authorization phase, read the [authorization RFC](../../
|
||||
32. **Custom middlewares** - *(optional)* Any `custom_middlewares` passed to `build_middlewares` are injected here, before config-declared extensions and the terminal-response/safety/clarification tail
|
||||
33. **Configured extension middlewares** - `extensions.middlewares` in `config.yaml` or `extensions_config.json` optionally accepts `module.path:ClassName` strings or `{class, kwargs}` objects. `deerflow.reflection.resolve_class` loads `AgentMiddleware` classes; import, class, and constructor errors fail agent creation. `kwargs` must be JSON-compatible; YAML dates/timestamps become ISO strings. Order: built-ins/custom and loop/token guards → extensions → terminal-response/safety/clarification tail. Subagents share the list before their safety tail; separate lead/subagent lists are unsupported. Trusted operator config only: paths instantiate arbitrary code. Gateway skill/MCP toggles preserve it in raw JSON; adding an API write path requires explicit trust-boundary review.
|
||||
34. **TerminalResponseMiddleware** - When a provider returns an empty terminal `AIMessage` after tool execution, injects a hidden recovery prompt and retries the model once; a second empty response is replaced in checkpoint state by a visible error fallback marked for the run worker, so the run finishes as an error instead of a silent success
|
||||
35. **ModelLengthFinishReasonMiddleware** - Records `stop_reason=model_length_capped` when provider-specific length detectors match a terminal `AIMessage` without tool-call intent (`finish_reason=length` / `MAX_TOKENS`, or `stop_reason=max_tokens`), preserving the original assistant content and never reparsing textual tool-call-like envelopes
|
||||
35. **ModelLengthFinishReasonMiddleware** - Match stamps `stop_reason=model_length_capped` and ends tool loop. Suppresses calls, appends notice even with partial text, and stamps `model_length_termination` so downstream guards stand down. Preserves content blocks.
|
||||
36. **SafetyFinishReasonMiddleware** - *(optional, if `safety_finish_reason.enabled`)* Suppresses tool execution when the provider safety-terminated the response (e.g. `finish_reason=content_filter`); registered after terminal-response/custom/configured middlewares so LangChain's reverse-order `after_model` dispatch runs it first
|
||||
37. **ClarificationMiddleware** - Intercepts `ask_clarification`, writes a readable `ToolMessage.content` fallback plus a structured `ToolMessage.artifact.human_input` payload, and interrupts via `Command(goto=END)` (must be last). `after_model` drops same-turn sibling tool calls so they cannot run before the user answers; a malformed `ask_clarification` parked on `invalid_tool_calls` is the same stop signal. `disable_clarification` runs keep the siblings. Payloads are versioned — legacy `free_text`/`choice_with_other` stay `version: 1`; the v2 `form` mode (from `fields`) is `version: 2` so older frontends reject it and fall back to plain text. Field normalization is deterministic and lives in the middleware (it short-circuits before tool execution, so tool-arg typing gives no runtime validation), and it is atomic: any structurally broken entry — non-dict, bad/duplicate name, a name colliding with a JS `Object.prototype` member (`__proto__`/`constructor`), or exceeding the caps (16 fields / 24 options per field / 200 chars per text / `MAX_FORM_SERIALIZED_BYTES` = 16KB UTF-8, the per-item caps alone admitting forms whose IM text fallback overruns channel limits) — degrades the whole form to the legacy option/free-text modes, so a card never renders "complete" while missing a field. Benign issues degrade locally (unknown types — incl. unhashable JSON like `type: []`, which must not raise from the membership probe — and option-less selects become `text`); options are trimmed/deduped with blanks dropped (form- and top-level) since the frontend rejects blank labels. XML-to-dict option payloads are recursively flattened from dict/list containers in source order, scalar leaves kept, residual XML tags stripped before that trimming. Checkboxes are booleans defaulting to "no"; `required` on one means consent semantics. The response protocol is unchanged (v1 `text`/`option`): form cards submit a text summary as `response_kind: "text"`, so journal persistence needs no new allowlist entries. Because this middleware can short-circuit before `on_tool_end`, `RunJournal` does a root-run reconciliation for `ToolMessage`s whose `tool_call_id` came from the current run, so cards survive checkpoint compaction. That reconciliation is **not** `ask_clarification`-only — any middleware that answers a tool call has the same gap, and a result the user saw must not vanish on reload (#4666 — `ReadBeforeWriteMiddleware` blocked-write errors reached the UI but not the event store). It is bounded by three conditions, not a name allowlist: the message is user-visible, the call belongs to this run's **lead agent** (`_remember_current_run_tool_calls` records lead-agent calls only; subagent results stay in `subagent.step`), and it is not already persisted. Human Input Card replies are `hide_from_ui` `HumanMessage`s with `additional_kwargs.human_input_response`; `RunJournal` persists only allowlisted hidden sources (currently `ask_clarification`) as `llm.human.input`.
|
||||
|
||||
@ -4,9 +4,12 @@ Background — see issue bytedance/deer-flow#4271.
|
||||
|
||||
Some providers stop generation because the output budget is exhausted and
|
||||
surface that through ``finish_reason='length'`` while still returning assistant
|
||||
content. DeerFlow preserves visible content, adds a deterministic notice when
|
||||
no visible answer was produced, and drops tool calls that may have been
|
||||
truncated at the output boundary before they can execute.
|
||||
content. DeerFlow preserves visible content, adds a deterministic notice
|
||||
whenever tool calls were suppressed (even when partial text survived), and
|
||||
drops tool calls that may have been truncated at the output boundary before
|
||||
they can execute. The stamped ``model_length_termination`` marker tells
|
||||
downstream guards (notably ``TodoMiddleware``) to stand down rather than
|
||||
re-engage a capped turn.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@ -85,7 +88,7 @@ class ModelLengthFinishReasonMiddleware(AgentMiddleware[AgentState]):
|
||||
|
||||
return {
|
||||
"suppress_truncated_tool_calls": True,
|
||||
"empty_content_fallback_hash": canonical_hash(_MODEL_LENGTH_CAPPED_CONTENT),
|
||||
"length_notice_hash": canonical_hash(_MODEL_LENGTH_CAPPED_CONTENT),
|
||||
}
|
||||
|
||||
def _detect(self, message: AIMessage) -> ModelLengthTermination | None:
|
||||
@ -153,7 +156,7 @@ class ModelLengthFinishReasonMiddleware(AgentMiddleware[AgentState]):
|
||||
}
|
||||
replacement = last.model_copy(
|
||||
update={
|
||||
"content": (cleaned_content if contains_visible_content else append_visible_text(content_source, _MODEL_LENGTH_CAPPED_CONTENT)),
|
||||
"content": append_visible_text(content_source, _MODEL_LENGTH_CAPPED_CONTENT),
|
||||
"tool_calls": [],
|
||||
"invalid_tool_calls": [],
|
||||
"additional_kwargs": additional_kwargs,
|
||||
|
||||
@ -48,6 +48,8 @@ def append_visible_text(message: AIMessage, text: str) -> Any:
|
||||
"""Append a visible text block without dropping existing content blocks."""
|
||||
if isinstance(message.content, list):
|
||||
return [*message.content, {"type": "text", "text": text}]
|
||||
if isinstance(message.content, str) and message.content.strip():
|
||||
return f"{message.content}\n\n{text}"
|
||||
return text
|
||||
|
||||
|
||||
|
||||
@ -11,6 +11,10 @@ there are still incomplete todo items. When the model produces a final response
|
||||
for the next model request and jumps back to the model node to force continued
|
||||
engagement. The completion reminder is injected via ``wrap_model_call`` instead
|
||||
of being persisted into graph state as a normal user-visible message.
|
||||
|
||||
The completion guard defers to ``model_length_termination``: when a length-capped
|
||||
turn has already been terminalized by ``ModelLengthFinishReasonMiddleware``,
|
||||
re-engaging would only re-emit the same oversized tool call into the same cap.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@ -119,6 +123,7 @@ class TodoMiddleware(TodoListMiddleware):
|
||||
"system_prompt_hash": canonical_hash(self.system_prompt),
|
||||
"tool_description_hash": canonical_hash(self.tool_description),
|
||||
"state_channel": "todos",
|
||||
"skip_completion_reminder_on_length_cap": True,
|
||||
}
|
||||
|
||||
@override
|
||||
@ -300,6 +305,13 @@ class TodoMiddleware(TodoListMiddleware):
|
||||
if (last_ai.additional_kwargs or {}).get("deerflow_error_fallback"):
|
||||
return None
|
||||
|
||||
# A length-capped turn was already terminalized by
|
||||
# ModelLengthFinishReasonMiddleware (tool calls suppressed, notice
|
||||
# appended); re-engaging would only re-emit the same oversized tool
|
||||
# call into the same cap.
|
||||
if (last_ai.additional_kwargs or {}).get("model_length_termination"):
|
||||
return None
|
||||
|
||||
# 3. Allow exit when all todos are completed or there are no todos.
|
||||
todos: list[Todo] = state.get("todos") or [] # type: ignore[assignment]
|
||||
if not todos or all(t.get("status") == "completed" for t in todos):
|
||||
|
||||
@ -11,7 +11,13 @@ The Gateway sizes pages to the `CONVERSATION_TOOL_NAME` tool-output budget so
|
||||
results stay inline. Cut messages carry a `message_seq`/`offset` continuation that
|
||||
the same host reader serves; keep reading guidance separate from permission enforcement.
|
||||
|
||||
Lead and bootstrap assembly pass the constructed `chat_model` to tool assembly.
|
||||
The cloned `write_file` budget hint uses that instance's effective `max_tokens`,
|
||||
including custom-agent and thinking-mode overrides; an absent cap omits the hint.
|
||||
Only standalone tool discovery without a model falls back to the base profile.
|
||||
|
||||
`get_available_tools(groups, include_mcp, model_name, subagent_enabled)` assembles:
|
||||
|
||||
1. **Config-defined tools** - Resolved from `config.yaml` via `resolve_variable()`
|
||||
2. **MCP tools** - From enabled MCP servers (lazy initialized, cached with resolved-path + content-signature invalidation)
|
||||
3. **Built-in tools**:
|
||||
|
||||
@ -1,7 +1,10 @@
|
||||
import copy
|
||||
import logging
|
||||
import threading
|
||||
|
||||
from langchain.tools import BaseTool
|
||||
from langchain_core.language_models import BaseChatModel
|
||||
from pydantic import BaseModel
|
||||
|
||||
from deerflow.config import get_app_config
|
||||
from deerflow.config.app_config import AppConfig
|
||||
@ -70,6 +73,34 @@ def _ensure_sync_invocable_tool(tool: BaseTool) -> BaseTool:
|
||||
return tool
|
||||
|
||||
|
||||
def _extract_max_tokens(model_config: object | None) -> int | None:
|
||||
"""Safely extract a positive integer max_tokens from a model config object.
|
||||
|
||||
Handles ModelConfig (where max_tokens may be stored as an extra dynamic field),
|
||||
dicts, SimpleNamespace, or test stubs. Rejects booleans, mocks, non-numeric
|
||||
values, negative numbers, zero, and None.
|
||||
"""
|
||||
if model_config is None:
|
||||
return None
|
||||
raw = model_config.get("max_tokens") if isinstance(model_config, dict) else getattr(model_config, "max_tokens", None)
|
||||
if isinstance(raw, bool) or not isinstance(raw, (int, float, str)):
|
||||
return None
|
||||
try:
|
||||
val = int(raw)
|
||||
return val if val > 0 else None
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
|
||||
|
||||
def _clone_tool_with_description(tool: BaseTool, description: str) -> BaseTool:
|
||||
"""Return a copy of tool with an updated description, leaving the original intact."""
|
||||
if isinstance(tool, BaseModel):
|
||||
return tool.model_copy(update={"description": description})
|
||||
cloned = copy.copy(tool)
|
||||
cloned.description = description
|
||||
return cloned
|
||||
|
||||
|
||||
def get_available_tools(
|
||||
groups: list[str] | None = None,
|
||||
include_mcp: bool = True,
|
||||
@ -80,6 +111,7 @@ def get_available_tools(
|
||||
include_upload_tool: bool = True,
|
||||
include_conversation_reader: bool = False,
|
||||
app_config: AppConfig | None = None,
|
||||
chat_model: BaseChatModel | None = None,
|
||||
) -> list[BaseTool]:
|
||||
"""Get all available tools from config.
|
||||
|
||||
@ -90,6 +122,9 @@ def get_available_tools(
|
||||
groups: Optional list of tool groups to filter by.
|
||||
include_mcp: Whether to include tools from MCP servers (default: True).
|
||||
model_name: Optional model name to determine if vision tools should be included.
|
||||
chat_model: Constructed model whose effective output cap supplies write_file
|
||||
guidance. When supplied, an absent cap omits the hint; only callers
|
||||
without a model fall back to the configured profile.
|
||||
subagent_enabled: Whether to include subagent tools (task, task_status).
|
||||
include_upload_tool: Whether to include ``list_uploaded_files`` (default: True).
|
||||
Ordinary task subagents enable it only after snapshotting the
|
||||
@ -164,6 +199,30 @@ def get_available_tools(
|
||||
builtin_tools.append(view_image_tool)
|
||||
logger.info(f"Including view_image_tool for model '{model_name}' (supports_vision=True)")
|
||||
|
||||
# Annotate write_file with the constructed model's effective output budget so the
|
||||
# model does not assume the 80 KB streaming ceiling is the practical limit
|
||||
# for a single completion. The tool is cloned to avoid mutating the
|
||||
# module-level singleton in-place across assemblies or leaking guidance to
|
||||
# models configured without max_tokens.
|
||||
max_tokens = _extract_max_tokens(chat_model if chat_model is not None else model_config)
|
||||
if max_tokens is not None:
|
||||
safe_chars = int(max_tokens * 3 * 0.7)
|
||||
budget_note = (
|
||||
f"\n\nPER-RESPONSE BUDGET: your output limit is {max_tokens} tokens "
|
||||
f"(≈{safe_chars} chars). Single non-append writes above this will be truncated. "
|
||||
"For larger documents, write the first section now, "
|
||||
"then use append=True for subsequent sections."
|
||||
)
|
||||
loaded_tools = [
|
||||
_clone_tool_with_description(
|
||||
tool,
|
||||
f"{getattr(tool, 'description', '') or ''}{budget_note}",
|
||||
)
|
||||
if tool.name == "write_file" and hasattr(tool, "description") and "PER-RESPONSE BUDGET:" not in (getattr(tool, "description", "") or "")
|
||||
else tool
|
||||
for tool in loaded_tools
|
||||
]
|
||||
|
||||
# Get cached MCP tools if enabled
|
||||
# NOTE: We use ExtensionsConfig.from_file() instead of config.extensions
|
||||
# to always read the latest configuration from disk. This ensures that changes
|
||||
|
||||
@ -7,7 +7,7 @@ import inspect
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock, create_autospec
|
||||
from unittest.mock import ANY, AsyncMock, MagicMock, create_autospec
|
||||
|
||||
import pytest
|
||||
from langchain.agents import create_agent
|
||||
@ -659,7 +659,7 @@ def test_make_lead_agent_reads_runtime_options_from_context(monkeypatch):
|
||||
"reasoning_effort": "high",
|
||||
"app_config": app_config,
|
||||
}
|
||||
get_available_tools.assert_called_once_with(model_name="context-model", groups=None, subagent_enabled=True, mcp_plugins=None, include_conversation_reader=False, app_config=app_config)
|
||||
get_available_tools.assert_called_once_with(model_name="context-model", groups=None, subagent_enabled=True, mcp_plugins=None, include_conversation_reader=False, app_config=app_config, chat_model=result["model"])
|
||||
assert result["model"] is not None
|
||||
|
||||
|
||||
@ -1534,6 +1534,7 @@ def test_empty_allowed_subagents_disables_requested_delegation(monkeypatch, mcp_
|
||||
subagent_enabled=False,
|
||||
include_conversation_reader=False,
|
||||
app_config=app_config,
|
||||
chat_model=ANY,
|
||||
)
|
||||
assert config["context"]["subagent_enabled"] is False
|
||||
assert config["configurable"]["subagent_enabled"] is False
|
||||
|
||||
@ -150,7 +150,8 @@ def test_finish_reason_length_drops_potentially_truncated_tool_calls():
|
||||
replacement = result["messages"][0]
|
||||
assert replacement.tool_calls == []
|
||||
assert replacement.invalid_tool_calls == []
|
||||
assert replacement.content == [{"type": "text", "text": "partial answer"}]
|
||||
assert replacement.content[0] == {"type": "text", "text": "partial answer"}
|
||||
assert "output limit" in replacement.content[-1]["text"]
|
||||
assert "tool_calls" not in replacement.additional_kwargs
|
||||
assert replacement.additional_kwargs["model_length_termination"]["suppressed_tool_call_count"] == 1
|
||||
assert replacement.additional_kwargs["model_length_termination"]["suppressed_tool_call_names"] == ["write_file"]
|
||||
@ -220,6 +221,34 @@ def test_anthropic_thinking_is_preserved_when_native_tool_use_is_removed():
|
||||
assert "output limit" in content[-1]["text"]
|
||||
|
||||
|
||||
def test_suppressed_tool_call_with_string_fragment_appends_notice():
|
||||
"""Reproduces the incident: string content fragment + suppressed tool call
|
||||
at length cap -> notice is appended alongside the fragment."""
|
||||
mw = ModelLengthFinishReasonMiddleware()
|
||||
runtime = _runtime()
|
||||
msg = AIMessage(
|
||||
content="nit",
|
||||
tool_calls=[
|
||||
{
|
||||
"name": "write_file",
|
||||
"id": "call_truncated",
|
||||
"args": {"path": "/mnt/user-data/outputs/report.md", "content": "# Deep Research\n| ext4 | jbd2 | Every 5s |"},
|
||||
}
|
||||
],
|
||||
response_metadata={"finish_reason": "length", "model_name": "deepseek-v4-pro"},
|
||||
)
|
||||
|
||||
result = mw._apply({"messages": [msg]}, runtime)
|
||||
|
||||
assert result is not None
|
||||
replacement = result["messages"][0]
|
||||
assert replacement.tool_calls == []
|
||||
assert "output limit" in replacement.content
|
||||
assert "nit" in replacement.content
|
||||
assert replacement.additional_kwargs["model_length_termination"]["suppressed_tool_call_names"] == ["write_file"]
|
||||
assert runtime.context["stop_reason"] == MODEL_LENGTH_CAPPED_STOP_REASON
|
||||
|
||||
|
||||
def test_finish_reason_length_suppresses_complete_tool_call_as_safety_policy():
|
||||
"""Even parsed arguments cannot be proven complete after a length cap."""
|
||||
mw = ModelLengthFinishReasonMiddleware()
|
||||
|
||||
@ -10,6 +10,9 @@ from langchain_core.language_models.fake_chat_models import FakeMessagesListChat
|
||||
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
|
||||
from pydantic import PrivateAttr
|
||||
|
||||
from deerflow.agents.middlewares.model_length_finish_reason_middleware import (
|
||||
ModelLengthFinishReasonMiddleware,
|
||||
)
|
||||
from deerflow.agents.middlewares.todo_middleware import (
|
||||
TodoMiddleware,
|
||||
_format_todos,
|
||||
@ -453,6 +456,41 @@ class TestAfterModel:
|
||||
}
|
||||
assert mw.after_model(state, _make_runtime()) is None
|
||||
|
||||
def test_does_not_reengage_when_model_length_capped_marker_present(self):
|
||||
mw = TodoMiddleware()
|
||||
state = {
|
||||
"messages": [
|
||||
AIMessage(
|
||||
content="nit",
|
||||
tool_calls=[],
|
||||
additional_kwargs={
|
||||
"model_length_termination": {
|
||||
"detector": "openai_compatible_length",
|
||||
"suppressed_tool_call_count": 1,
|
||||
"suppressed_tool_call_names": ["write_file"],
|
||||
}
|
||||
},
|
||||
)
|
||||
],
|
||||
"todos": _incomplete_todos(),
|
||||
}
|
||||
assert mw.after_model(state, _make_runtime()) is None
|
||||
|
||||
def test_pure_text_length_cap_without_marker_still_reengages(self):
|
||||
mw = TodoMiddleware()
|
||||
state = {
|
||||
"messages": [
|
||||
AIMessage(
|
||||
content="partial answer",
|
||||
response_metadata={"finish_reason": "length"},
|
||||
)
|
||||
],
|
||||
"todos": _incomplete_todos(),
|
||||
}
|
||||
result = mw.after_model(state, _make_runtime())
|
||||
assert result is not None
|
||||
assert result["jump_to"] == "model"
|
||||
|
||||
|
||||
class TestAafterModel:
|
||||
def test_delegates_to_sync(self):
|
||||
@ -617,6 +655,44 @@ class TestTodoMiddlewareAgentGraphIntegration:
|
||||
assert mw._pending_completion_reminders == {}
|
||||
assert mw._completion_reminder_counts == {}
|
||||
|
||||
def test_length_capped_write_file_does_not_reengage_todos(self):
|
||||
"""Reproduces the incident (thread b1723286): model emits a write_file
|
||||
call with finish_reason=length, ModelLength suppresses it, and
|
||||
TodoMiddleware must NOT re-engage via jump_to=model."""
|
||||
todo_mw = TodoMiddleware()
|
||||
model = _CapturingFakeMessagesListChatModel(
|
||||
responses=[
|
||||
AIMessage(
|
||||
content="nit",
|
||||
tool_calls=[{"name": "write_file", "id": "call_1", "args": {"path": "/mnt/user-data/outputs/report.md", "content": "# truncated report\n| ext4 | jbd2"}}],
|
||||
response_metadata={"finish_reason": "length", "model_name": "deepseek-v4-pro"},
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
graph = create_agent(
|
||||
model=model,
|
||||
tools=[],
|
||||
middleware=[todo_mw, ModelLengthFinishReasonMiddleware()],
|
||||
state_schema=ThreadState,
|
||||
)
|
||||
|
||||
result = graph.invoke(
|
||||
{"messages": [("user", "write the report")], "todos": [{"content": "Phase 9: Synthesis and final report writing", "status": "in_progress"}]},
|
||||
context={"thread_id": "cap-thread", "run_id": "cap-run"},
|
||||
)
|
||||
|
||||
assert len(model.seen_messages) == 1
|
||||
reminders_by_call = [_todo_completion_reminders(messages) for messages in model.seen_messages]
|
||||
assert all(len(r) == 0 for r in reminders_by_call)
|
||||
|
||||
final_ai = result["messages"][-1]
|
||||
assert final_ai.additional_kwargs.get("model_length_termination")
|
||||
assert "output limit" in str(final_ai.content)
|
||||
assert "nit" in str(final_ai.content)
|
||||
assert result["todos"] == [{"content": "Phase 9: Synthesis and final report writing", "status": "in_progress"}]
|
||||
assert todo_mw._pending_completion_reminders == {}
|
||||
|
||||
|
||||
class TestRunScopedReminderCleanup:
|
||||
def test_before_agent_clears_stale_count_without_pending_reminder(self):
|
||||
|
||||
208
backend/tests/test_write_file_tool_budget_annotation.py
Normal file
208
backend/tests/test_write_file_tool_budget_annotation.py
Normal file
@ -0,0 +1,208 @@
|
||||
"""Tests for write_file tool budget annotation and assembly isolation.
|
||||
|
||||
Verifies fixes for reviewer findings on PR #5569:
|
||||
- [P1] Guarded max_tokens extraction avoiding AttributeError on ModelConfig without max_tokens
|
||||
- [P2] Cloning write_file tool to preserve module-level singleton immutability across assemblies
|
||||
and prevent cross-model guidance leakage.
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
|
||||
from deerflow.config.app_config import AppConfig, ModelConfig, SandboxConfig, ToolConfig
|
||||
from deerflow.sandbox.tools import write_file_tool
|
||||
from deerflow.tools.tools import (
|
||||
_clone_tool_with_description,
|
||||
_extract_max_tokens,
|
||||
get_available_tools,
|
||||
)
|
||||
|
||||
|
||||
def test_extract_max_tokens_various_inputs():
|
||||
"""Verify _extract_max_tokens safely handles all expected and edge-case inputs."""
|
||||
# None and empty
|
||||
assert _extract_max_tokens(None) is None
|
||||
|
||||
# ModelConfig with and without max_tokens
|
||||
mc_without = ModelConfig(name="test", model="m", use="u")
|
||||
assert _extract_max_tokens(mc_without) is None
|
||||
|
||||
mc_with = ModelConfig(name="test", model="m", use="u", max_tokens=4096)
|
||||
assert _extract_max_tokens(mc_with) == 4096
|
||||
|
||||
mc_zero = ModelConfig(name="test", model="m", use="u", max_tokens=0)
|
||||
assert _extract_max_tokens(mc_zero) is None
|
||||
|
||||
mc_neg = ModelConfig(name="test", model="m", use="u", max_tokens=-100)
|
||||
assert _extract_max_tokens(mc_neg) is None
|
||||
|
||||
# Dictionaries
|
||||
assert _extract_max_tokens({}) is None
|
||||
assert _extract_max_tokens({"max_tokens": 2048}) == 2048
|
||||
assert _extract_max_tokens({"max_tokens": "8192"}) == 8192
|
||||
assert _extract_max_tokens({"max_tokens": None}) is None
|
||||
assert _extract_max_tokens({"max_tokens": 0}) is None
|
||||
|
||||
# SimpleNamespace
|
||||
assert _extract_max_tokens(SimpleNamespace()) is None
|
||||
assert _extract_max_tokens(SimpleNamespace(max_tokens=1024)) == 1024
|
||||
|
||||
# Booleans (must NOT be treated as 1 or 0)
|
||||
assert _extract_max_tokens({"max_tokens": True}) is None
|
||||
assert _extract_max_tokens({"max_tokens": False}) is None
|
||||
assert _extract_max_tokens(SimpleNamespace(max_tokens=True)) is None
|
||||
|
||||
# Floats
|
||||
assert _extract_max_tokens({"max_tokens": 4096.0}) == 4096
|
||||
|
||||
# Unparseable strings
|
||||
assert _extract_max_tokens({"max_tokens": "unlimited"}) is None
|
||||
|
||||
# MagicMock (in Python unittest.mock, int(MagicMock()) defaults to 1; must be rejected)
|
||||
mock_without = MagicMock(spec=[])
|
||||
assert _extract_max_tokens(mock_without) is None
|
||||
|
||||
mock_with = MagicMock()
|
||||
mock_with.max_tokens = 8000
|
||||
assert _extract_max_tokens(mock_with) == 8000
|
||||
|
||||
|
||||
def test_clone_tool_with_description_preserves_singleton():
|
||||
"""Verify _clone_tool_with_description returns an isolated copy and keeps original unchanged."""
|
||||
original_desc = write_file_tool.description
|
||||
assert "CUSTOM_TEST_BUDGET" not in original_desc
|
||||
|
||||
cloned = _clone_tool_with_description(write_file_tool, original_desc + "\n\nCUSTOM_TEST_BUDGET")
|
||||
|
||||
assert "CUSTOM_TEST_BUDGET" in cloned.description
|
||||
assert write_file_tool.description == original_desc
|
||||
assert cloned.name == write_file_tool.name
|
||||
assert cloned.func is write_file_tool.func
|
||||
assert cloned.coroutine is write_file_tool.coroutine
|
||||
assert cloned.args_schema is write_file_tool.args_schema
|
||||
assert isinstance(cloned, BaseModel)
|
||||
|
||||
|
||||
def _build_minimal_app_config(models: list[ModelConfig]) -> AppConfig:
|
||||
return AppConfig(
|
||||
models=models,
|
||||
sandbox=SandboxConfig(use="deerflow.sandbox.local:LocalSandboxProvider"),
|
||||
tools=[
|
||||
ToolConfig(name="write_file", group="file:write", use="deerflow.sandbox.tools:write_file_tool"),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_get_available_tools_with_model_config_lacking_max_tokens():
|
||||
"""Verify get_available_tools does not raise AttributeError when max_tokens is omitted."""
|
||||
model_without_max_tokens = ModelConfig(name="capped-model", model="m", use="u")
|
||||
config = _build_minimal_app_config([model_without_max_tokens])
|
||||
|
||||
tools = get_available_tools(model_name="capped-model", app_config=config, include_mcp=False)
|
||||
write_tool = next((t for t in tools if t.name == "write_file"), None)
|
||||
assert write_tool is not None
|
||||
assert "PER-RESPONSE BUDGET:" not in write_tool.description
|
||||
|
||||
|
||||
def test_write_file_singleton_remains_unmutated_across_assemblies():
|
||||
"""Verify write_file_tool process singleton is never mutated during tool assembly."""
|
||||
baseline_desc = write_file_tool.description
|
||||
|
||||
model_with_budget = ModelConfig(name="budget-model", model="m", use="u", max_tokens=4096)
|
||||
config = _build_minimal_app_config([model_with_budget])
|
||||
|
||||
tools = get_available_tools(model_name="budget-model", app_config=config, include_mcp=False)
|
||||
assembled_write_file = next(t for t in tools if t.name == "write_file")
|
||||
|
||||
assert "PER-RESPONSE BUDGET: your output limit is 4096 tokens" in assembled_write_file.description
|
||||
# The process-wide singleton must remain pristine
|
||||
assert write_file_tool.description == baseline_desc
|
||||
assert "PER-RESPONSE BUDGET:" not in write_file_tool.description
|
||||
|
||||
|
||||
def test_repeated_assembly_cross_model_isolation():
|
||||
"""Verify consecutive tool assemblies for different models do not leak guidance or duplicate notes."""
|
||||
baseline_desc = write_file_tool.description
|
||||
|
||||
model_4k = ModelConfig(name="model-4k", model="m", use="u", max_tokens=4096)
|
||||
model_32k = ModelConfig(name="model-32k", model="m", use="u", max_tokens=32768)
|
||||
model_none = ModelConfig(name="model-none", model="m", use="u")
|
||||
config = _build_minimal_app_config([model_4k, model_32k, model_none])
|
||||
|
||||
# Assembly 1: 4K model
|
||||
tools_4k = get_available_tools(model_name="model-4k", app_config=config, include_mcp=False)
|
||||
wf_4k = next(t for t in tools_4k if t.name == "write_file")
|
||||
assert "your output limit is 4096 tokens" in wf_4k.description
|
||||
assert "32768" not in wf_4k.description
|
||||
assert write_file_tool.description == baseline_desc
|
||||
|
||||
# Assembly 2: 32K model (must not carry 4096 note)
|
||||
tools_32k = get_available_tools(model_name="model-32k", app_config=config, include_mcp=False)
|
||||
wf_32k = next(t for t in tools_32k if t.name == "write_file")
|
||||
assert "your output limit is 32768 tokens" in wf_32k.description
|
||||
assert "4096" not in wf_32k.description
|
||||
assert wf_32k.description.count("PER-RESPONSE BUDGET:") == 1
|
||||
assert write_file_tool.description == baseline_desc
|
||||
|
||||
# Assembly 3: model without max_tokens (must have NO budget note at all)
|
||||
tools_none = get_available_tools(model_name="model-none", app_config=config, include_mcp=False)
|
||||
wf_none = next(t for t in tools_none if t.name == "write_file")
|
||||
assert "PER-RESPONSE BUDGET:" not in wf_none.description
|
||||
assert wf_none.description == baseline_desc
|
||||
assert write_file_tool.description == baseline_desc
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("profile_overrides", "agent_settings", "thinking_enabled", "bootstrap", "expected"),
|
||||
[
|
||||
({}, {"max_tokens": 1024}, False, False, 1024),
|
||||
({"when_thinking_enabled": {"max_tokens": 1024}}, {}, True, False, 1024),
|
||||
({"when_thinking_disabled": {"max_tokens": 1024}}, {}, False, False, 1024),
|
||||
({"when_thinking_enabled": {"max_tokens": 1024}}, {"max_tokens": 2048}, True, False, 1024),
|
||||
({"when_thinking_disabled": {"max_tokens": None}}, {}, False, False, None),
|
||||
({"when_thinking_enabled": {"max_tokens": 1024}}, {}, True, True, 1024),
|
||||
({"when_thinking_disabled": {"max_tokens": 1024}}, {}, False, True, 1024),
|
||||
],
|
||||
ids=["custom-agent", "thinking-on", "thinking-off", "thinking-over-agent", "uncapped", "bootstrap-thinking-on", "bootstrap-thinking-off"],
|
||||
)
|
||||
def test_lead_write_file_budget_matches_constructed_model(monkeypatch, profile_overrides, agent_settings, thinking_enabled, bootstrap, expected):
|
||||
"""Exercise real model and tool assembly, including override precedence."""
|
||||
from deerflow.agents.lead_agent import agent as lead_agent_module
|
||||
from deerflow.config.agents_config import AgentConfig
|
||||
from deerflow.config.extensions_config import ExtensionsConfig
|
||||
|
||||
model = ModelConfig(
|
||||
name="budget-model",
|
||||
model="budget-model",
|
||||
use="langchain_openai:ChatOpenAI",
|
||||
api_key="test-key",
|
||||
max_tokens=32768,
|
||||
supports_thinking=True,
|
||||
**profile_overrides,
|
||||
)
|
||||
app_config = _build_minimal_app_config([model])
|
||||
agent_config = AgentConfig(name="researcher", model="budget-model", model_settings=agent_settings)
|
||||
monkeypatch.setattr(lead_agent_module, "load_agent_config", lambda *args, **kwargs: agent_config)
|
||||
monkeypatch.setattr(lead_agent_module, "_load_enabled_available_skills", lambda *args, **kwargs: [])
|
||||
monkeypatch.setattr(lead_agent_module, "build_middlewares", lambda *args, **kwargs: [])
|
||||
monkeypatch.setattr(lead_agent_module, "apply_prompt_template", lambda **kwargs: "system prompt")
|
||||
monkeypatch.setattr(lead_agent_module, "create_agent", lambda **kwargs: kwargs)
|
||||
monkeypatch.setattr(lead_agent_module, "build_tracing_callbacks", lambda: [])
|
||||
monkeypatch.setattr(ExtensionsConfig, "from_file", lambda *args, **kwargs: ExtensionsConfig())
|
||||
|
||||
graph = lead_agent_module._make_lead_agent(
|
||||
{"context": {"agent_name": "researcher", "thinking_enabled": thinking_enabled, "is_bootstrap": bootstrap}},
|
||||
app_config=app_config,
|
||||
)
|
||||
|
||||
assert graph["model"].max_tokens == expected
|
||||
write_tool = next(tool for tool in graph["tools"] if tool.name == "write_file")
|
||||
if expected is None:
|
||||
assert "PER-RESPONSE BUDGET:" not in write_tool.description
|
||||
else:
|
||||
assert f"your output limit is {expected} tokens" in write_tool.description
|
||||
assert "32768 tokens" not in write_tool.description
|
||||
assert "PER-RESPONSE BUDGET:" not in write_file_tool.description
|
||||
Loading…
x
Reference in New Issue
Block a user