From d8e085630ac04eef9bdd6c95c84d4bae5320617b Mon Sep 17 00:00:00 2001 From: Evan <1178327328@qq.com> Date: Sun, 8 Nov 2026 06:19:56 -0500 Subject: [PATCH] fix(agents): keep todo reminders in Goal turns after an earlier provider stop (#8004) * fix(agents): keep todo reminders in Goal turns after an earlier provider stop TodoMiddleware skipped completion reminders whenever runtime.context held safety_capped or model_length_capped. The worker keeps that key across the run's hidden Goal continuation turns so the run record reports the first cap, and an unmet Goal continues after these provider stops (it stands down only on loop and token caps). After one capped response, every continuation turn therefore lost its todo reminders and could stop with unfinished items. Loop and token stops stay run-level and are still read from the context. Provider safety and length stops now apply to the response they ended, identified by the safety_termination / model_length_termination marker that those guards stamp on it. The check still runs before the parent hook, so a capped response cannot be resumed by the parallel-write retry. * refactor(agents): import provider-stop markers from the guards that stamp them TodoMiddleware matched the safety and length markers by string literals that duplicated the guards' own. Each guard now exports its marker and stop reason as module constants, used at its stamp and check sites and imported by TodoMiddleware, so a rename in a guard cannot silently disable suppression. Comment that completion_reminder_hard_stop_reasons is the inventory surface and only its run-level entries are read from runtime.context, and run the worker-level Goal continuation test for the safety guard too, which overwrites stop_reason instead of keeping the first one. --------- Co-authored-by: Totoro --- README.md | 2 + .../deerflow/agents/middlewares/AGENTS.md | 2 +- .../model_length_finish_reason_middleware.py | 6 +- .../safety_finish_reason_middleware.py | 10 +- .../agents/middlewares/todo_middleware.py | 44 +++-- backend/tests/test_todo_goal_continuation.py | 154 ++++++++++++++++++ backend/tests/test_todo_middleware.py | 24 ++- 7 files changed, 219 insertions(+), 23 deletions(-) create mode 100644 backend/tests/test_todo_goal_continuation.py diff --git a/README.md b/README.md index 4340b4d1de..e7f1d6027a 100644 --- a/README.md +++ b/README.md @@ -2746,6 +2746,8 @@ When loop detection, the run token budget, or a provider safety/output-length limit stops a run, unfinished todos are retained for a later user run; todo completion reminders do not restart model calls or tool execution after that stop. Safety-stopped partial text is preserved; it does not trigger a Todo retry. +A provider safety or output-length stop applies to the response it ended: if an +unmet goal continues the run, later continuation turns still get todo reminders. Optional `pii_redaction.enabled` redacts detected identifiers in user messages, remote tool results, compaction input, reinjected summaries, configured diff --git a/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md b/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md index 97c32f0daa..c55621e2d6 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md +++ b/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md @@ -150,7 +150,7 @@ Before changing a later authorization phase, read the [authorization RFC](../../ 18. **SkillActivationMiddleware** - Detects strict `/skill-name task` syntax on the latest real user message, resolves only enabled and runtime-allowed skills, injects the `SKILL.md` body as hidden current-turn context, and records a `middleware:skill_activation` audit event; a retried activation call replays the recorded reminder without re-auditing. Skill `allowed-tools` is compatibility metadata only: activation never filters, grants, or denies tools — Agent assembly/execution authorization owns the permission ceiling. 19. **DurableContextMiddleware** - Before compaction, captures task dispatch/result summaries in `ThreadState.delegations` and loaded skill name/path/description (never bodies) in `ThreadState.skill_context`; projects them per request. Static authority uses `SystemMessage`; untrusted `summary_text`, delegation results and skill descriptions use a separate hidden `HumanMessage` data block, never persisted as messages or promoted to system instructions. Subagents attach it immediately before summarization: summary precedes preserved AI/tool history for strict providers. See [tool artifacts](TOOL_ARTIFACTS.md) for projection, retry limits and delegation scope. 20. **SummarizationMiddleware** - *(optional)* Compacts near token limits. Fraction trigger/keep profiles belong to the active run model; an explicit summary model only generates. Memory flush/manual compaction follow runtime/checkpoint agent policy. Preserve the latest user request and DynamicContext reminders; a backward cutoff may retain old AI/tool turns or make first-turn compaction a no-op. -21. **TodoListMiddleware** - *(optional, if `is_plan_mode`)* Task tracking with `write_todos`. Before the parent hook, completion reminders respect loop/token/safety/model-length `stop_reason` values; provider termination markers remain fallbacks. Unfinished todos survive for later runs. A raising model call requeues its reminder, uncounted, for the retry. Both hooks read stored todos through `_usable_todos`: thread-state updates and direct integrations can Overwrite that channel, and a non-mapping entry would otherwise raise on every model call. +21. **TodoListMiddleware** - *(optional, if `is_plan_mode`)* Task tracking with `write_todos`. Before the parent hook, completion reminders respect run-level loop/token `stop_reason` values and the current response's safety/length markers; context keeps an earlier provider stop across Goal turns. Unfinished todos survive for later runs. A raising model call requeues its reminder, uncounted, for the retry. Both hooks read stored todos through `_usable_todos`: thread-state updates and direct integrations can Overwrite that channel, and a non-mapping entry would otherwise raise on every model call. 22. **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. TODO attribution skips non-string statuses before membership; preserve args/TODO state for normal tool validation. Attribution normalizes the stored `todos` channel with `_normalize_todos` before reading entries, the same tolerance the todo middleware applies. 23. **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. 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 index a287e4161d..b05c75cf7b 100644 --- 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 @@ -42,6 +42,8 @@ from deerflow.agents.middlewares.model_response import append_visible_text, has_ from deerflow.agents.middlewares.tool_call_metadata import clone_ai_message_with_tool_calls MODEL_LENGTH_CAPPED_STOP_REASON = "model_length_capped" +# Stamped on the AIMessage this guard rewrote; consumers import it, not the literal. +MODEL_LENGTH_TERMINATION_MARKER = "model_length_termination" _MODEL_LENGTH_CAPPED_CONTENT = "The model reached its output limit before producing a complete final response. Please continue the conversation to resume." logger = logging.getLogger(__name__) @@ -158,7 +160,7 @@ class ModelLengthFinishReasonMiddleware(AgentMiddleware[AgentState]): content_source = clone_ai_message_with_tool_calls(last.model_copy(update={"invalid_tool_calls": []}), []) contains_tool_call = has_tool_call_intent(last) or content_source.content is not last.content contains_visible_content = has_visible_content(content_source) - if not contains_tool_call and contains_visible_content and (last.additional_kwargs or {}).get("model_length_termination"): + if not contains_tool_call and contains_visible_content and (last.additional_kwargs or {}).get(MODEL_LENGTH_TERMINATION_MARKER): return None additional_kwargs = dict(last.additional_kwargs or {}) @@ -167,7 +169,7 @@ class ModelLengthFinishReasonMiddleware(AgentMiddleware[AgentState]): # Tool arguments may be incomplete at a max-token boundary. additional_kwargs.pop("tool_calls", None) additional_kwargs.pop("function_call", None) - additional_kwargs["model_length_termination"] = { + additional_kwargs[MODEL_LENGTH_TERMINATION_MARKER] = { "detector": termination.detector, "reason_field": termination.reason_field, "reason_value": termination.reason_value, diff --git a/backend/packages/harness/deerflow/agents/middlewares/safety_finish_reason_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/safety_finish_reason_middleware.py index 807e2a6509..dec34fb0e1 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/safety_finish_reason_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/safety_finish_reason_middleware.py @@ -71,6 +71,10 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +SAFETY_CAPPED_STOP_REASON = "safety_capped" +# Stamped on the AIMessage this guard rewrote; consumers import it, not the literal. +SAFETY_TERMINATION_MARKER = "safety_termination" + _USER_FACING_MESSAGE = ( "The model provider stopped this response with a safety-related signal " @@ -218,7 +222,7 @@ class SafetyFinishReasonMiddleware(AgentMiddleware[AgentState]): # made a shallow copy, but downstream model_copy still references # it). Then stamp the observability record. kwargs = dict(getattr(cleared, "additional_kwargs", None) or {}) - kwargs["safety_termination"] = { + kwargs[SAFETY_TERMINATION_MARKER] = { "detector": termination.detector, "reason_field": termination.reason_field, "reason_value": termination.reason_value, @@ -392,10 +396,10 @@ class SafetyFinishReasonMiddleware(AgentMiddleware[AgentState]): # alongside loop_capped / token_capped (#4176). ctx = getattr(runtime, "context", None) if isinstance(ctx, dict): - ctx["stop_reason"] = "safety_capped" + ctx["stop_reason"] = SAFETY_CAPPED_STOP_REASON # Re-observing a terminal message must not replace its original # suppression record or emit another event. Still stamp this runtime. - if not tool_calls and not content_is_blank and (last.additional_kwargs or {}).get("safety_termination"): + if not tool_calls and not content_is_blank and (last.additional_kwargs or {}).get(SAFETY_TERMINATION_MARKER): return None patched = self._build_suppressed_message(last, termination) diff --git a/backend/packages/harness/deerflow/agents/middlewares/todo_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/todo_middleware.py index d2d17b0887..95cf9a82dd 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/todo_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/todo_middleware.py @@ -15,8 +15,10 @@ 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. -Run-level loop, token, provider-safety and model-length stops also suppress -completion reminders, leaving unfinished todos available for a later user run. +Run-level loop and token stops, and a provider safety or length stop of the +current response, also suppress completion reminders, leaving unfinished todos +available for a later user run. An earlier response's provider stop does not +suppress reminders in later Goal continuation turns of the same run. """ from __future__ import annotations @@ -30,12 +32,21 @@ from langchain.agents.middleware.types import ModelCallResult, ModelRequest, Mod from langchain_core.messages import AIMessage, HumanMessage from langgraph.runtime import Runtime +from deerflow.agents.middlewares.model_length_finish_reason_middleware import MODEL_LENGTH_CAPPED_STOP_REASON, MODEL_LENGTH_TERMINATION_MARKER +from deerflow.agents.middlewares.safety_finish_reason_middleware import SAFETY_CAPPED_STOP_REASON, SAFETY_TERMINATION_MARKER from deerflow.agents.thread_state import ThreadState TODO_REMINDER_MESSAGE_NAME = "todo_reminder" -# Runtime Goal continuation has a narrower policy; keep this local to Todo. -# Source-inventory and graph tests pin these values to the actual guard stamps. -_COMPLETION_REMINDER_HARD_STOP_REASONS = ("loop_capped", "token_capped", "safety_capped", "model_length_capped") +# Source-inventory and graph tests pin these values to the actual guard stamps; +# the provider-stop vocabulary is imported from the guards that write it. +# Loop and token stops cap the whole run: their budgets stay exhausted and a +# Goal stands down on them, so runtime.context["stop_reason"] is authoritative. +_COMPLETION_REMINDER_RUN_STOP_REASONS = ("loop_capped", "token_capped") +# Provider safety and length stops end one response, which carries the marker. +# runtime.context keeps the first cap across hidden Goal turns, and a Goal +# continues after these, so the context alone does not mean the current +# response was capped. +_COMPLETION_REMINDER_RESPONSE_STOP_MARKERS = {SAFETY_CAPPED_STOP_REASON: SAFETY_TERMINATION_MARKER, MODEL_LENGTH_CAPPED_STOP_REASON: MODEL_LENGTH_TERMINATION_MARKER} def _todos_in_messages(messages: list[Any]) -> bool: @@ -160,7 +171,11 @@ class TodoMiddleware(TodoListMiddleware): "state_channel": "todos", "skip_completion_reminder_on_length_cap": True, "skip_completion_reminder_on_safety_cap": True, - "completion_reminder_hard_stop_reasons": list(_COMPLETION_REMINDER_HARD_STOP_REASONS), + # Every stop_reason that suppresses reminders, for the stop-reason + # inventory test. Only the run-level reasons are read from + # runtime.context; provider stops are matched by response marker. + "completion_reminder_hard_stop_reasons": [*_COMPLETION_REMINDER_RUN_STOP_REASONS, *_COMPLETION_REMINDER_RESPONSE_STOP_MARKERS], + "completion_reminder_response_stop_markers": dict(_COMPLETION_REMINDER_RESPONSE_STOP_MARKERS), } @override @@ -358,7 +373,12 @@ class TodoMiddleware(TodoListMiddleware): # calls and stamp loop/token/safety/model-length stops. Check before # the parent hook so its parallel-write retry cannot override a stop. context = getattr(runtime, "context", None) - if isinstance(context, dict) and context.get("stop_reason") in _COMPLETION_REMINDER_HARD_STOP_REASONS: + if isinstance(context, dict) and context.get("stop_reason") in _COMPLETION_REMINDER_RUN_STOP_REASONS: + return None + messages = state.get("messages") or [] + last_ai = next((m for m in reversed(messages) if isinstance(m, AIMessage)), None) + termination_metadata = (last_ai.additional_kwargs or {}) if last_ai else {} + if any(termination_metadata.get(marker) for marker in _COMPLETION_REMINDER_RESPONSE_STOP_MARKERS.values()): return None # 1. Preserve base class logic (parallel write_todos detection). @@ -369,18 +389,10 @@ class TodoMiddleware(TodoListMiddleware): # 2. Only intervene when the agent wants to exit cleanly. Tool-call # intent or tool-call parse errors should be handled by the tool path # instead of being masked by todo reminders. - messages = state.get("messages") or [] - last_ai = next((m for m in reversed(messages) if isinstance(m, AIMessage)), None) if not last_ai or _has_tool_call_intent_or_error(last_ai): return None - if (last_ai.additional_kwargs or {}).get("deerflow_error_fallback"): - return None - - # Provider caps also mark the message: preserve this fallback when - # runtime context is unavailable or another reason already occupies it. - termination_metadata = last_ai.additional_kwargs or {} - if termination_metadata.get("model_length_termination") or termination_metadata.get("safety_termination"): + if termination_metadata.get("deerflow_error_fallback"): return None # 3. Allow exit when all todos are completed or there are no todos. diff --git a/backend/tests/test_todo_goal_continuation.py b/backend/tests/test_todo_goal_continuation.py new file mode 100644 index 0000000000..ed6fb3ebf1 --- /dev/null +++ b/backend/tests/test_todo_goal_continuation.py @@ -0,0 +1,154 @@ +"""Todo completion reminders in Goal continuation turns after an earlier response cap. + +The worker keeps the first cap reason in ``runtime.context["stop_reason"]`` across +hidden Goal turns so the run record reports it, and an unmet Goal keeps going after +a provider length or safety cap. Those turns are new model responses: a cap that +ended an earlier response must not switch off Todo reminders for them. +""" + +import json +import os +import socket +import tempfile +from pathlib import Path +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, patch +from uuid import uuid4 + +import pytest +from langchain.agents import create_agent +from langchain_core.language_models.chat_models import BaseChatModel +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage +from langchain_core.outputs import ChatGeneration, ChatResult +from langgraph.checkpoint.memory import InMemorySaver +from pydantic import ConfigDict + +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.todo_middleware import TodoMiddleware +from deerflow.agents.thread_state import ThreadState +from deerflow.config.app_config import AppConfig +from deerflow.config.paths import Paths +from deerflow.config.sandbox_config import SandboxConfig +from deerflow.runtime.events.store.memory import MemoryRunEventStore +from deerflow.runtime.goal import build_goal_state, read_thread_goal, write_thread_goal +from deerflow.runtime.runs import worker +from deerflow.runtime.runs.manager import RunManager +from deerflow.runtime.runs.store.memory import MemoryRunStore +from deerflow.runtime.user_context import reset_current_user, set_current_user + +_UNFINISHED_TODOS = [ + {"content": "Draft the report", "status": "completed"}, + {"content": "Add the appendix", "status": "in_progress"}, +] + + +class _ScriptedModel(BaseChatModel): + """Return scripted responses in order and record every request.""" + + responses: list[AIMessage] + requests: list[list[BaseMessage]] + model_config = ConfigDict(arbitrary_types_allowed=True) + + @property + def _llm_type(self) -> str: + return "offline-todo-goal-continuation" + + def bind_tools(self, tools: Any, **kwargs: Any) -> "_ScriptedModel": + return self + + def _generate(self, messages: Any, stop: Any = None, run_manager: Any = None, **kwargs: Any) -> ChatResult: + self.requests.append(list(messages)) + if not self.responses: + raise AssertionError("The lead model was called more often than scripted") + return ChatResult(generations=[ChatGeneration(message=self.responses.pop(0))]) + + +class _UnmetEvaluator: + async def ainvoke(self, messages: Any, config: Any = None) -> AIMessage: + return AIMessage(content=json.dumps({"satisfied": False, "blocker": "goal_not_met_yet", "reason": "The appendix is missing.", "evidence_summary": "Only the draft exists.", "relied_on_assumption": False})) + + +def _has_completion_reminder(request: list[BaseMessage]) -> bool: + return any(isinstance(message, HumanMessage) and message.name == "todo_completion_reminder" for message in request) + + +async def _run_goal_with_capped_user_turn(model: _ScriptedModel, guard: Any): + saver = InMemorySaver() + manager = RunManager(store=MemoryRunStore()) + graph = create_agent( + model=model, + tools=[], + middleware=[TodoMiddleware(), guard], + state_schema=ThreadState, + context_schema=dict, + checkpointer=saver, + ) + cfg = AppConfig(sandbox=SandboxConfig(use="deerflow.sandbox.local:LocalSandboxProvider")) + cfg.title.enabled = cfg.memory.enabled = cfg.memory.injection_enabled = False + bridge = SimpleNamespace(publish=AsyncMock(), publish_end=AsyncMock(), cleanup=AsyncMock()) + token = set_current_user(SimpleNamespace(id="todo-goal-continuation", role="admin")) + try: + with tempfile.TemporaryDirectory(prefix="todo-goal-continuation-") as folder: + with ( + patch("deerflow.config.paths._paths", Paths(base_dir=Path(folder))), + patch.dict(os.environ, {"DEER_FLOW_HOME": folder, "LANGSMITH_TRACING": "false", "LANGCHAIN_TRACING_V2": "false"}), + patch.object(worker, "create_goal_evaluator_model", lambda **kwargs: _UnmetEvaluator()), + ): + record = await manager.create(f"todo-goal-{uuid4().hex}", user_id="todo-goal-continuation") + await write_thread_goal(saver, record.thread_id, build_goal_state("Write the report with its appendix.", max_continuations=1), create_if_missing=True) + await worker.run_agent( + bridge, + manager, + record, + ctx=worker.RunContext(checkpointer=saver, event_store=MemoryRunEventStore(), app_config=cfg), + agent_factory=lambda config: graph, + graph_input={"messages": [HumanMessage(content="Write the report with its appendix.")], "todos": _UNFINISHED_TODOS}, + config={"configurable": {"thread_id": record.thread_id}, "recursion_limit": 100}, + stream_modes=["values"], + ) + goal = await read_thread_goal(saver, record.thread_id) + return record, goal + finally: + reset_current_user(token) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("guard", "finish_reason", "stop_reason"), + [ + # ModelLength keeps the first cap reason already in the context. + pytest.param(ModelLengthFinishReasonMiddleware, "length", "model_length_capped", id="length"), + # Safety overwrites the context's stop_reason unconditionally. + pytest.param(SafetyFinishReasonMiddleware, "content_filter", "safety_capped", id="safety"), + ], +) +async def test_goal_continuation_after_a_provider_stop_still_gets_todo_reminders(monkeypatch, guard, finish_reason, stop_reason): + def deny_network(*args, **kwargs): + raise AssertionError("This controlled regression must not use the network") + + monkeypatch.setattr(socket.socket, "connect", deny_network) + model = _ScriptedModel( + responses=[ + # User turn: the provider stops the answer early. + AIMessage(content="Here is the first half of the report", response_metadata={"finish_reason": finish_reason}), + # Hidden Goal turn: a complete response that tries to stop with the appendix unfinished. + AIMessage(content="The report is done."), + AIMessage(content="Adding the appendix now."), + AIMessage(content="Appendix added."), + ], + requests=[], + ) + + record, goal = await _run_goal_with_capped_user_turn(model, guard()) + + assert record.error is None + # The run still reports the stop from the user turn. + assert record.stop_reason == stop_reason + assert goal is not None and goal["continuation_count"] == 1 + # The stopped user turn ends without a reminder ... + assert not _has_completion_reminder(model.requests[0]) + # ... but the continuation's premature exit is reminded twice, up to the usual cap. + assert len(model.requests) == 4 + assert [_has_completion_reminder(request) for request in model.requests[1:]] == [False, True, True] diff --git a/backend/tests/test_todo_middleware.py b/backend/tests/test_todo_middleware.py index d412dc4cd2..19bb16dbb2 100644 --- a/backend/tests/test_todo_middleware.py +++ b/backend/tests/test_todo_middleware.py @@ -222,6 +222,12 @@ def _ai_no_tool_calls(): return AIMessage(content="I'm done!") +def _ai_stopped_by(stop_reason): + """The response a guard leaves behind: provider stops also mark the message.""" + marker = {"safety_capped": "safety_termination", "model_length_capped": "model_length_termination"}.get(stop_reason) + return AIMessage(content="I'm done!", additional_kwargs={marker: {"suppressed_tool_call_count": 0}} if marker else {}) + + def _ai_with_invalid_tool_calls(): return AIMessage( content="", @@ -357,7 +363,7 @@ class TestAfterModel: mw = TodoMiddleware() runtime = _make_runtime() runtime.context["stop_reason"] = stop_reason - state = {"messages": [_ai_no_tool_calls()], "todos": _incomplete_todos()} + state = {"messages": [_ai_stopped_by(stop_reason)], "todos": _incomplete_todos()} result = asyncio.run(mw.aafter_model(state, runtime)) if invoke_async else mw.after_model(state, runtime) @@ -383,9 +389,25 @@ class TestAfterModel: base_hook.assert_called_once() base_hook.reset_mock() runtime.context["stop_reason"] = stop_reason + state["messages"] = [_ai_stopped_by(stop_reason)] assert after_model() is None base_hook.assert_not_called() + @pytest.mark.parametrize("invoke_async", [False, True], ids=["sync", "async"]) + @pytest.mark.parametrize("stop_reason", ["safety_capped", "model_length_capped"]) + def test_earlier_provider_stop_in_context_keeps_completion_reminders(self, invoke_async, stop_reason): + # Goal continuation turns share the run's context, which keeps the first + # cap for the run record. A complete later response is not capped. + mw = TodoMiddleware() + runtime = _make_runtime() + runtime.context["stop_reason"] = stop_reason + state = {"messages": [_ai_no_tool_calls()], "todos": _incomplete_todos()} + + result = asyncio.run(mw.aafter_model(state, runtime)) if invoke_async else mw.after_model(state, runtime) + + assert result == {"jump_to": "model"} + assert mw._completion_reminder_count_for_runtime(runtime) == 1 + @pytest.mark.parametrize("stop_reason", [None, "custom_stop", "subagent_limit_capped", [], {}]) def test_unrecognized_stop_reason_keeps_completion_reminders(self, stop_reason): mw = TodoMiddleware()