diff --git a/backend/AGENTS.md b/backend/AGENTS.md index 22f08a28a..35241fcb3 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -361,8 +361,9 @@ Before changing a later authorization phase, read the [authorization RFC](../doc 30. **Custom middlewares** - *(optional)* Any `custom_middlewares` passed to `build_middlewares` are injected here, before config-declared extensions and the terminal-response/safety/clarification tail 31. **Configured extension middlewares** - *(optional, if `extensions.middlewares` is set in `config.yaml` or `extensions_config.json`)* Zero-argument `AgentMiddleware` classes loaded from `module.path:ClassName` entries via `deerflow.reflection.resolve_class`. Missing packages, invalid classes, and broken modules fail loudly at agent creation. These run after built-ins/programmatic custom middleware and after the lead/subagent loop/token guards, but before the terminal-response/safety/clarification tail; subagents receive the same configured extension middleware class list before their safety tail. Treat these files as trusted operator config because middleware paths instantiate arbitrary code. Gateway skill/MCP toggle endpoints preserve this field through `to_file_dict()` but must not add a write path for `extensions.middlewares` without an explicit trust-boundary review. Lead-only vs subagent-only middleware lists and per-context constructor parameters are not expressible in this MVP. 32. **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 -33. **SafetyFinishReasonMiddleware** - *(optional, if `safety_finish_reason.enabled`)* Repairs AIMessages the provider safety-terminated (e.g. `finish_reason=content_filter`): strips truncated tool calls so they are not executed (#3028), and — when the response is otherwise blank (no tool calls, no visible content) — backfills a user-facing explanation so the empty message is not persisted and then rejected by strict OpenAI-compatible providers on the next request (`message ... with role 'assistant' must not be empty`), which would otherwise strand the whole thread (#4393). A safety-terminated response that still carries visible text is left untouched. Registered after terminal-response/custom/configured middlewares so LangChain's reverse-order `after_model` dispatch runs it first -34. **ClarificationMiddleware** - Intercepts `ask_clarification` tool calls, writes a readable `ToolMessage.content` fallback plus structured `ToolMessage.artifact.human_input` request payload, and interrupts via `Command(goto=END)` (must be last). Because this middleware can short-circuit tool execution before LangChain emits `on_tool_end`, `RunJournal` performs a root-run final reconciliation for allowlisted clarification `ToolMessage`s whose `tool_call_id` was produced by the current run, so human-input request cards remain recoverable from `run_events` after checkpoint compaction. Human Input Card replies are submitted as `hide_from_ui` `HumanMessage`s with `additional_kwargs.human_input_response`; `RunJournal` persists only allowlisted hidden response sources (currently `ask_clarification`) as `llm.human.input`, which preserves answered-card state after compaction without exposing generic internal hidden context. +33. **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 +34. **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 +35. **ClarificationMiddleware** - Intercepts `ask_clarification` tool calls, writes a readable `ToolMessage.content` fallback plus structured `ToolMessage.artifact.human_input` request payload, and interrupts via `Command(goto=END)` (must be last). Because this middleware can short-circuit tool execution before LangChain emits `on_tool_end`, `RunJournal` performs a root-run final reconciliation for allowlisted clarification `ToolMessage`s whose `tool_call_id` was produced by the current run, so human-input request cards remain recoverable from `run_events` after checkpoint compaction. Human Input Card replies are submitted as `hide_from_ui` `HumanMessage`s with `additional_kwargs.human_input_response`; `RunJournal` persists only allowlisted hidden response sources (currently `ask_clarification`) as `llm.human.input`, which preserves answered-card state after compaction without exposing generic internal hidden context. ### Configuration System diff --git a/backend/packages/harness/deerflow/agents/lead_agent/agent.py b/backend/packages/harness/deerflow/agents/lead_agent/agent.py index 189c87cee..b20c2e337 100644 --- a/backend/packages/harness/deerflow/agents/lead_agent/agent.py +++ b/backend/packages/harness/deerflow/agents/lead_agent/agent.py @@ -36,6 +36,7 @@ from deerflow.agents.middlewares.clarification_middleware import ClarificationMi from deerflow.agents.middlewares.configured_extensions import load_configured_extension_middlewares from deerflow.agents.middlewares.loop_detection_middleware import LoopDetectionMiddleware from deerflow.agents.middlewares.memory_middleware import MemoryMiddleware +from deerflow.agents.middlewares.model_length_finish_reason_middleware import ModelLengthFinishReasonMiddleware from deerflow.agents.middlewares.safety_finish_reason_middleware import SafetyFinishReasonMiddleware from deerflow.agents.middlewares.subagent_limit_middleware import SubagentLimitMiddleware from deerflow.agents.middlewares.summarization_middleware import DeerFlowSummarizationMiddleware, create_summarization_middleware @@ -452,6 +453,12 @@ def build_middlewares( # allowing LangChain's no-tool-call router to end a silent successful run. middlewares.append(TerminalResponseMiddleware()) + # A provider may also cap the final assistant response at the model output + # limit. Preserve the assistant content unchanged, but stamp a run-level + # stop_reason so Gateway consumers can tell a length-capped completion from + # a clean one. + middlewares.append(ModelLengthFinishReasonMiddleware()) + # SafetyFinishReasonMiddleware — suppress tool execution when the provider # safety-terminated the response. Registered after the terminal-response # and custom/configured middlewares so LangChain's reverse-order after_model diff --git a/backend/packages/harness/deerflow/agents/middlewares/model_length_finish_reason_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/model_length_finish_reason_middleware.py new file mode 100644 index 000000000..49650e469 --- /dev/null +++ b/backend/packages/harness/deerflow/agents/middlewares/model_length_finish_reason_middleware.py @@ -0,0 +1,130 @@ +"""Surface provider length-capped model responses as run stop reasons. + +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 should preserve that content for audit, but it should not +silently treat the run as an uncapped clean completion when the provider has +explicitly signaled truncation. + +This middleware keeps that boundary narrow: +- it only marks a run-level stop reason when the final AIMessage is capped + by a provider length signal and still has visible content; +- it never rewrites the assistant content or reparses XML-like text into a + tool call; +- it ignores any response that still carries tool-call intent, malformed + tool-call metadata, or no visible content, so only terminal assistant + responses with visible content can be marked capped. + +""" + +from __future__ import annotations + +import logging +from typing import Any, override + +from langchain.agents import AgentState +from langchain.agents.middleware import AgentMiddleware +from langchain_core.messages import AIMessage +from langgraph.runtime import Runtime + +from deerflow.agents.middlewares.model_length_termination_detectors import ( + ModelLengthTermination, + ModelLengthTerminationDetector, + default_detectors, +) + +MODEL_LENGTH_CAPPED_STOP_REASON = "model_length_capped" +logger = logging.getLogger(__name__) + + +def _has_tool_call_intent_or_error(message: AIMessage) -> bool: + if message.tool_calls or getattr(message, "invalid_tool_calls", None): + return True + additional_kwargs = message.additional_kwargs or {} + return bool(additional_kwargs.get("tool_calls") or additional_kwargs.get("function_call")) + + +def _has_visible_content(message: AIMessage) -> bool: + content = message.content + if isinstance(content, str): + return bool(content.strip()) + if isinstance(content, list): + for block in content: + if isinstance(block, str) and block.strip(): + return True + if isinstance(block, dict) and block.get("type") in {"text", "output_text"}: + text = block.get("text") + if isinstance(text, str) and text.strip(): + return True + return False + + +class ModelLengthFinishReasonMiddleware(AgentMiddleware[AgentState]): + """Record provider length caps for terminal assistant responses with content. + + If the last AIMessage still carries tool-call intent, this middleware + leaves it alone and lets the normal tool-handling path decide what to do. + """ + + def __init__(self, detectors: list[ModelLengthTerminationDetector] | None = None) -> None: + super().__init__() + self._detectors: list[ModelLengthTerminationDetector] = list(detectors) if detectors else default_detectors() + + def _detect(self, message: AIMessage) -> ModelLengthTermination | None: + for detector in self._detectors: + try: + hit = detector.detect(message) + except Exception: # noqa: BLE001 - provider detectors must not break a run + logger.exception("ModelLengthTerminationDetector %r raised; treating as no-match", getattr(detector, "name", type(detector).__name__)) + continue + if hit is not None: + return hit + return None + + def _apply(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None: + messages = list(state.get("messages") or []) + if not messages or not isinstance(messages[-1], AIMessage): + return None + + last = messages[-1] + if _has_tool_call_intent_or_error(last): + return None + if not _has_visible_content(last): + return None + + termination = self._detect(last) + if termination is None: + return None + + ctx = getattr(runtime, "context", None) + thread_id = ctx.get("thread_id") if isinstance(ctx, dict) else None + run_id = ctx.get("run_id") if isinstance(ctx, dict) else None + stamped_stop_reason = False + if isinstance(ctx, dict): + # Preserve any earlier cap reason carried across hidden continuation turns. + if "stop_reason" not in ctx: + ctx["stop_reason"] = MODEL_LENGTH_CAPPED_STOP_REASON + stamped_stop_reason = True + logger.info( + "Provider model length cap detected", + extra={ + "thread_id": thread_id, + "run_id": run_id, + "message_id": getattr(last, "id", None), + "detector": termination.detector, + "reason_field": termination.reason_field, + "reason_value": termination.reason_value, + "stamped_stop_reason": stamped_stop_reason, + }, + ) + return None + + @override + def after_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None: + return self._apply(state, runtime) + + @override + async def aafter_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None: + return self._apply(state, runtime) diff --git a/backend/packages/harness/deerflow/agents/middlewares/model_length_termination_detectors.py b/backend/packages/harness/deerflow/agents/middlewares/model_length_termination_detectors.py new file mode 100644 index 000000000..8158f5a99 --- /dev/null +++ b/backend/packages/harness/deerflow/agents/middlewares/model_length_termination_detectors.py @@ -0,0 +1,126 @@ +"""Detectors for provider-side model length termination signals. + +Different providers report "the response hit the output-token limit" through +different fields and values. Keep those provider details here so +``ModelLengthFinishReasonMiddleware`` can stay focused on when to mark a run, +not on which provider emitted which spelling. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Protocol, runtime_checkable + +from langchain_core.messages import AIMessage + + +@dataclass(frozen=True) +class ModelLengthTermination: + """A detected model-output length cap.""" + + detector: str + reason_field: str + reason_value: str + extras: dict[str, Any] = field(default_factory=dict) + + +@runtime_checkable +class ModelLengthTerminationDetector(Protocol): + """Strategy interface for provider length-cap detection.""" + + name: str + + def detect(self, message: AIMessage) -> ModelLengthTermination | None: + """Return a hit when *message* indicates output length truncation.""" + ... + + +def _get_metadata_value(message: AIMessage, field_name: str) -> str | None: + """Read a string metadata value from common LangChain provider fields.""" + for container_name in ("response_metadata", "additional_kwargs"): + container = getattr(message, container_name, None) or {} + if not isinstance(container, dict): + continue + value = container.get(field_name) + if isinstance(value, str) and value: + return value + return None + + +class OpenAICompatibleLengthDetector: + """OpenAI-compatible ``finish_reason == "length"`` signal.""" + + name = "openai_compatible_length" + + def __init__(self, finish_reasons: list[str] | tuple[str, ...] | None = None) -> None: + configured = finish_reasons if finish_reasons is not None else ("length",) + self._finish_reasons: frozenset[str] = frozenset(r.lower() for r in configured) + + def detect(self, message: AIMessage) -> ModelLengthTermination | None: + value = _get_metadata_value(message, "finish_reason") + if value is None or value.lower() not in self._finish_reasons: + return None + return ModelLengthTermination( + detector=self.name, + reason_field="finish_reason", + reason_value=value, + ) + + +class AnthropicMaxTokensDetector: + """Anthropic ``stop_reason == "max_tokens"`` signal.""" + + name = "anthropic_max_tokens" + + def __init__(self, stop_reasons: list[str] | tuple[str, ...] | None = None) -> None: + configured = stop_reasons if stop_reasons is not None else ("max_tokens",) + self._stop_reasons: frozenset[str] = frozenset(r.lower() for r in configured) + + def detect(self, message: AIMessage) -> ModelLengthTermination | None: + value = _get_metadata_value(message, "stop_reason") + if value is None or value.lower() not in self._stop_reasons: + return None + return ModelLengthTermination( + detector=self.name, + reason_field="stop_reason", + reason_value=value, + ) + + +class GeminiMaxTokensDetector: + """Gemini / Vertex AI ``finish_reason == "MAX_TOKENS"`` signal.""" + + name = "gemini_max_tokens" + + def __init__(self, finish_reasons: list[str] | tuple[str, ...] | None = None) -> None: + configured = finish_reasons if finish_reasons is not None else ("MAX_TOKENS",) + self._finish_reasons: frozenset[str] = frozenset(r.upper() for r in configured) + + def detect(self, message: AIMessage) -> ModelLengthTermination | None: + value = _get_metadata_value(message, "finish_reason") + if value is None or value.upper() not in self._finish_reasons: + return None + return ModelLengthTermination( + detector=self.name, + reason_field="finish_reason", + reason_value=value, + ) + + +def default_detectors() -> list[ModelLengthTerminationDetector]: + """Built-in detector set used for provider length-cap signals.""" + return [ + OpenAICompatibleLengthDetector(), + AnthropicMaxTokensDetector(), + GeminiMaxTokensDetector(), + ] + + +__all__ = [ + "AnthropicMaxTokensDetector", + "GeminiMaxTokensDetector", + "ModelLengthTermination", + "ModelLengthTerminationDetector", + "OpenAICompatibleLengthDetector", + "default_detectors", +] diff --git a/backend/packages/harness/deerflow/runtime/runs/worker.py b/backend/packages/harness/deerflow/runtime/runs/worker.py index 32da098cc..7faa1794b 100644 --- a/backend/packages/harness/deerflow/runtime/runs/worker.py +++ b/backend/packages/harness/deerflow/runtime/runs/worker.py @@ -761,6 +761,7 @@ async def run_agent( # token_budget -> "token_capped" # safety_finish_reason -> "safety_capped" # subagent_limit -> "subagent_limit_capped" + # model_length_finish_reason -> "model_length_capped" # # If more guards grow stop_reason semantics, consider a publish/ # collect pattern (e.g. each guard middleware publishes its cap diff --git a/backend/tests/_replay_fixture.py b/backend/tests/_replay_fixture.py index 5886c3f15..a0220cbfb 100644 --- a/backend/tests/_replay_fixture.py +++ b/backend/tests/_replay_fixture.py @@ -156,7 +156,10 @@ def drive_gateway(app, *, prompt: str, context: dict) -> list[dict]: body = { "assistant_id": "lead_agent", "input": {"messages": [{"role": "user", "content": prompt}]}, - "config": {"recursion_limit": 50}, + # Keep replay close to the Gateway default. A tighter limit can + # produce false golden drift when protocol-neutral middleware adds + # graph steps without changing the streamed SSE contract. + "config": {"recursion_limit": 100}, "context": context, "stream_mode": ["values"], } diff --git a/backend/tests/test_lead_agent_model_resolution.py b/backend/tests/test_lead_agent_model_resolution.py index 2c6ad3939..c8dc0f7c0 100644 --- a/backend/tests/test_lead_agent_model_resolution.py +++ b/backend/tests/test_lead_agent_model_resolution.py @@ -545,15 +545,18 @@ def test_build_middlewares_uses_resolved_model_name_for_vision(monkeypatch): assert any(isinstance(m, lead_agent_module.ViewImageMiddleware) for m in middlewares) # verify the custom middleware is injected correctly. # With this test's default safety config enabled, the tail order is: - # ..., custom, TerminalResponseMiddleware, SafetyFinishReasonMiddleware, - # ClarificationMiddleware, so the custom mock sits at index [-4]. - assert len(middlewares) > 0 and isinstance(middlewares[-4], MagicMock) + # ..., custom, TerminalResponseMiddleware, ModelLengthFinishReasonMiddleware, + # SafetyFinishReasonMiddleware, ClarificationMiddleware, so the custom mock + # sits at index [-5]. + assert len(middlewares) > 0 and isinstance(middlewares[-5], MagicMock) from deerflow.agents.middlewares.clarification_middleware import ClarificationMiddleware + from deerflow.agents.middlewares.model_length_finish_reason_middleware import ModelLengthFinishReasonMiddleware from deerflow.agents.middlewares.safety_finish_reason_middleware import SafetyFinishReasonMiddleware from deerflow.agents.middlewares.terminal_response_middleware import TerminalResponseMiddleware - assert isinstance(middlewares[-3], TerminalResponseMiddleware) + assert isinstance(middlewares[-4], TerminalResponseMiddleware) + assert isinstance(middlewares[-3], ModelLengthFinishReasonMiddleware) assert isinstance(middlewares[-2], SafetyFinishReasonMiddleware) assert isinstance(middlewares[-1], ClarificationMiddleware) @@ -842,10 +845,11 @@ def test_build_middlewares_injects_configured_extension_middlewares(monkeypatch) ) middleware_types = [type(m).__name__ for m in middlewares] - assert middleware_types[-5:] == [ + assert middleware_types[-6:] == [ "ConfiguredGuardMiddleware", "ConfiguredAuditMiddleware", "TerminalResponseMiddleware", + "ModelLengthFinishReasonMiddleware", "SafetyFinishReasonMiddleware", "ClarificationMiddleware", ] diff --git a/backend/tests/test_loop_detection_stop_reason.py b/backend/tests/test_loop_detection_stop_reason.py index ac66042e8..b1e909093 100644 --- a/backend/tests/test_loop_detection_stop_reason.py +++ b/backend/tests/test_loop_detection_stop_reason.py @@ -221,6 +221,59 @@ async def test_worker_surfaces_stop_reason_from_safety_finish_reason(): assert fetched.stop_reason == "safety_capped" +@pytest.mark.asyncio +async def test_worker_surfaces_stop_reason_from_model_length_finish_reason(): + """The worker persists ``stop_reason=model_length_capped`` when the + provider caps a terminal assistant response with ``finish_reason=length``.""" + from deerflow.agents.middlewares.model_length_finish_reason_middleware import ModelLengthFinishReasonMiddleware + + mw = ModelLengthFinishReasonMiddleware() + captured_runtime: list[Any] = [None] + + class DummyAgent: + metadata: dict[str, Any] = {"model_name": "test-model"} + + async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False): + runtime = ((config or {}).get("configurable") or {}).get("__pregel_runtime") + assert runtime is not None + captured_runtime[0] = runtime + + msg = AIMessage( + content=('/mnt/user-data/outputs/report.md# partial'), + tool_calls=[], + invalid_tool_calls=[], + response_metadata={"finish_reason": "length"}, + ) + mw._apply({"messages": [msg]}, runtime) + + yield {"messages": [msg]} + + run_manager = RunManager() + record = await run_manager.create("thread-1") + bridge = AsyncMock() + bridge.publish = AsyncMock() + bridge.publish_end = AsyncMock() + bridge.cleanup = AsyncMock() + + await run_agent( + bridge, + run_manager, + record, + ctx=RunContext(checkpointer=None), + agent_factory=lambda *, config: DummyAgent(), + graph_input={"messages": []}, + config={}, + ) + + assert captured_runtime[0] is not None + assert captured_runtime[0].context.get("stop_reason") == "model_length_capped" + + fetched = await run_manager.get(record.run_id) + assert fetched is not None + assert fetched.status == RunStatus.success + assert fetched.stop_reason == "model_length_capped" + + @pytest.mark.asyncio async def test_worker_surfaces_stop_reason_from_subagent_limit(): """The worker persists ``stop_reason=subagent_limit_capped`` when the diff --git a/backend/tests/test_model_length_finish_reason_middleware.py b/backend/tests/test_model_length_finish_reason_middleware.py new file mode 100644 index 000000000..eedaedeb4 --- /dev/null +++ b/backend/tests/test_model_length_finish_reason_middleware.py @@ -0,0 +1,161 @@ +"""Unit tests for ModelLengthFinishReasonMiddleware.""" + +import logging +from unittest.mock import MagicMock + +from langchain_core.messages import AIMessage, HumanMessage + +from deerflow.agents.middlewares.model_length_finish_reason_middleware import ( + MODEL_LENGTH_CAPPED_STOP_REASON, + ModelLengthFinishReasonMiddleware, +) + +_MW_LOGGER = "deerflow.agents.middlewares.model_length_finish_reason_middleware" + + +def _runtime(run_id: str = "run-1"): + runtime = MagicMock() + runtime.context = {"thread_id": "thread-1", "run_id": run_id} + return runtime + + +def test_finish_reason_length_records_stop_reason_without_rewriting_textual_tool_call(): + mw = ModelLengthFinishReasonMiddleware() + runtime = _runtime() + msg = AIMessage( + content=('/mnt/user-data/outputs/report.md# partial'), + tool_calls=[], + invalid_tool_calls=[], + response_metadata={"finish_reason": "length"}, + ) + + result = mw._apply({"messages": [msg]}, runtime) + + assert result is None + assert runtime.context["stop_reason"] == MODEL_LENGTH_CAPPED_STOP_REASON + assert msg.content.startswith('') + assert msg.tool_calls == [] + assert msg.invalid_tool_calls == [] + + +def test_finish_reason_stop_with_tool_call_example_in_prose_passes_through(): + mw = ModelLengthFinishReasonMiddleware() + runtime = _runtime() + msg = AIMessage( + content=('Here is an example, not a real tool call:\n\n```xml\n\n```'), + response_metadata={"finish_reason": "stop"}, + ) + + assert mw._apply({"messages": [msg]}, runtime) is None + assert "stop_reason" not in runtime.context + + +def test_additional_kwargs_finish_reason_length_records_stop_reason(): + mw = ModelLengthFinishReasonMiddleware() + runtime = _runtime() + msg = AIMessage( + content="partial", + additional_kwargs={"finish_reason": "length"}, + ) + + assert mw._apply({"messages": [msg]}, runtime) is None + assert runtime.context["stop_reason"] == MODEL_LENGTH_CAPPED_STOP_REASON + + +def test_gemini_max_tokens_finish_reason_records_stop_reason(): + mw = ModelLengthFinishReasonMiddleware() + runtime = _runtime() + msg = AIMessage( + content="partial", + response_metadata={"finish_reason": "MAX_TOKENS"}, + ) + + assert mw._apply({"messages": [msg]}, runtime) is None + assert runtime.context["stop_reason"] == MODEL_LENGTH_CAPPED_STOP_REASON + + +def test_anthropic_max_tokens_stop_reason_records_stop_reason(): + mw = ModelLengthFinishReasonMiddleware() + runtime = _runtime() + msg = AIMessage( + content="partial", + response_metadata={"stop_reason": "max_tokens"}, + ) + + assert mw._apply({"messages": [msg]}, runtime) is None + assert runtime.context["stop_reason"] == MODEL_LENGTH_CAPPED_STOP_REASON + + +def test_length_cap_detection_logs_observability_fields(caplog): + mw = ModelLengthFinishReasonMiddleware() + runtime = _runtime(run_id="run-observe") + msg = AIMessage( + id="msg-1", + content="partial", + response_metadata={"finish_reason": "length"}, + ) + + with caplog.at_level(logging.INFO, logger=_MW_LOGGER): + assert mw._apply({"messages": [msg]}, runtime) is None + + records = [record for record in caplog.records if record.message == "Provider model length cap detected"] + assert len(records) == 1 + record = records[0] + assert record.thread_id == "thread-1" + assert record.run_id == "run-observe" + assert record.message_id == "msg-1" + assert record.detector == "openai_compatible_length" + assert record.reason_field == "finish_reason" + assert record.reason_value == "length" + assert record.stamped_stop_reason is True + + +def test_finish_reason_length_with_tool_calls_passes_through(): + mw = ModelLengthFinishReasonMiddleware() + runtime = _runtime() + msg = AIMessage( + content="", + tool_calls=[ + { + "id": "call_write_1", + "name": "write_file", + "args": {"path": "/mnt/user-data/outputs/report.md", "content": "# partial"}, + } + ], + response_metadata={"finish_reason": "length"}, + ) + + assert mw._apply({"messages": [msg]}, runtime) is None + assert "stop_reason" not in runtime.context + + +def test_empty_finish_reason_length_passes_through_for_terminal_response_recovery(): + mw = ModelLengthFinishReasonMiddleware() + runtime = _runtime() + msg = AIMessage(content="", response_metadata={"finish_reason": "length"}) + + assert mw._apply({"messages": [msg]}, runtime) is None + assert "stop_reason" not in runtime.context + + +def test_existing_stop_reason_is_not_overwritten(caplog): + mw = ModelLengthFinishReasonMiddleware() + runtime = _runtime() + runtime.context["stop_reason"] = "token_capped" + msg = AIMessage(content="partial", response_metadata={"finish_reason": "length"}) + + with caplog.at_level(logging.INFO, logger=_MW_LOGGER): + assert mw._apply({"messages": [msg]}, runtime) is None + + assert runtime.context["stop_reason"] == "token_capped" + records = [record for record in caplog.records if record.message == "Provider model length cap detected"] + assert len(records) == 1 + assert records[0].stamped_stop_reason is False + + +def test_non_ai_last_message_passes_through(): + mw = ModelLengthFinishReasonMiddleware() + runtime = _runtime() + + assert mw._apply({"messages": [HumanMessage(content="hello")]}, runtime) is None + assert "stop_reason" not in runtime.context diff --git a/backend/tests/test_model_length_termination_detectors.py b/backend/tests/test_model_length_termination_detectors.py new file mode 100644 index 000000000..61eb0ea3f --- /dev/null +++ b/backend/tests/test_model_length_termination_detectors.py @@ -0,0 +1,62 @@ +"""Unit tests for provider length-termination detectors.""" + +from langchain_core.messages import AIMessage + +from deerflow.agents.middlewares.model_length_termination_detectors import ( + AnthropicMaxTokensDetector, + GeminiMaxTokensDetector, + OpenAICompatibleLengthDetector, + default_detectors, +) + + +def test_openai_compatible_length_detector_matches_finish_reason_length(): + hit = OpenAICompatibleLengthDetector().detect( + AIMessage( + content="partial", + response_metadata={"finish_reason": "length"}, + ) + ) + + assert hit is not None + assert hit.detector == "openai_compatible_length" + assert hit.reason_field == "finish_reason" + assert hit.reason_value == "length" + + +def test_gemini_max_tokens_detector_matches_uppercase_finish_reason(): + hit = GeminiMaxTokensDetector().detect( + AIMessage( + content="partial", + response_metadata={"finish_reason": "MAX_TOKENS"}, + ) + ) + + assert hit is not None + assert hit.detector == "gemini_max_tokens" + assert hit.reason_field == "finish_reason" + assert hit.reason_value == "MAX_TOKENS" + + +def test_anthropic_max_tokens_detector_matches_stop_reason(): + hit = AnthropicMaxTokensDetector().detect( + AIMessage( + content="partial", + response_metadata={"stop_reason": "max_tokens"}, + ) + ) + + assert hit is not None + assert hit.detector == "anthropic_max_tokens" + assert hit.reason_field == "stop_reason" + assert hit.reason_value == "max_tokens" + + +def test_default_detectors_cover_openai_anthropic_and_gemini(): + names = [detector.name for detector in default_detectors()] + + assert names == [ + "openai_compatible_length", + "anthropic_max_tokens", + "gemini_max_tokens", + ]