diff --git a/backend/app/gateway/routers/thread_runs.py b/backend/app/gateway/routers/thread_runs.py index a97854e3d..a4de3b8ab 100644 --- a/backend/app/gateway/routers/thread_runs.py +++ b/backend/app/gateway/routers/thread_runs.py @@ -126,6 +126,7 @@ class RunResponse(BaseModel): subagent_tokens: int = 0 middleware_tokens: int = 0 message_count: int = 0 + stop_reason: str | None = None class ThreadTokenUsageModelBreakdown(BaseModel): @@ -230,6 +231,7 @@ def _record_to_response(record: RunRecord) -> RunResponse: subagent_tokens=record.subagent_tokens, middleware_tokens=record.middleware_tokens, message_count=record.message_count, + stop_reason=record.stop_reason, ) diff --git a/backend/packages/harness/deerflow/agents/middlewares/loop_detection_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/loop_detection_middleware.py index fccf053a1..459429131 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/loop_detection_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/loop_detection_middleware.py @@ -592,6 +592,11 @@ class LoopDetectionMiddleware(AgentMiddleware[AgentState]): run_id = self._get_run_id(runtime) with self._lock: self._stop_reason[run_id] = "loop_capped" + # Also write to runtime.context so the lead worker can read it + # without needing a reference to this middleware instance (#4176). + ctx = getattr(runtime, "context", None) + if isinstance(ctx, dict): + ctx["stop_reason"] = "loop_capped" # Strip tool_calls from the last AIMessage to force text output. # Once tool_calls are stripped, the AIMessage no longer requires # matching ToolMessage responses, so mutating it in place here 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 8fd733c23..b215b9712 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 @@ -283,6 +283,11 @@ class SafetyFinishReasonMiddleware(AgentMiddleware[AgentState]): if termination is None: return None + # Stamp stop_reason so the worker can surface this capped completion + # alongside loop_capped / token_capped (#4176). + ctx = getattr(runtime, "context", None) + if isinstance(ctx, dict): + ctx["stop_reason"] = "safety_capped" patched = self._build_suppressed_message(last, termination) thread_id = None diff --git a/backend/packages/harness/deerflow/agents/middlewares/subagent_limit_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/subagent_limit_middleware.py index 59f8671f9..39e8e87d2 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/subagent_limit_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/subagent_limit_middleware.py @@ -154,6 +154,12 @@ class SubagentLimitMiddleware(AgentMiddleware[AgentState]): prior_delegation_count, ) + # Stamp stop_reason when the total per-run cap is exhausted so the + # worker surfaces this capped completion alongside loop_capped / + # token_capped / safety_capped (#4176). + if remaining_total == 0 and isinstance(getattr(runtime, "context", None), dict): + runtime.context["stop_reason"] = "subagent_limit_capped" + # Replace the AIMessage with truncated tool_calls (same id triggers replacement) content = _append_text(last_msg.content, _TOTAL_LIMIT_STOP_MSG) if remaining_total == 0 else None updated_msg = clone_ai_message_with_tool_calls(last_msg, truncated_tool_calls, content=content) diff --git a/backend/packages/harness/deerflow/agents/middlewares/token_budget_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/token_budget_middleware.py index 6799adf08..7518ba5b7 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/token_budget_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/token_budget_middleware.py @@ -251,6 +251,11 @@ class TokenBudgetMiddleware(AgentMiddleware[AgentState]): # returns (the hard stop itself does not raise). See # ``consume_stop_reason``. self._stop_reason[run_id] = "token_capped" + # Also write to runtime.context so the lead worker can read it + # without needing a reference to this middleware instance (#4176). + ctx = getattr(runtime, "context", None) + if isinstance(ctx, dict): + ctx["stop_reason"] = "token_capped" stop_text = _BUDGET_EXCEEDED_MSG.format(reason=trigger_reason, used=trigger_used, budget=trigger_budget) return self._build_hard_stop_update(last_msg, stop_text) diff --git a/backend/packages/harness/deerflow/persistence/migrations/versions/0005_run_stop_reason.py b/backend/packages/harness/deerflow/persistence/migrations/versions/0005_run_stop_reason.py new file mode 100644 index 000000000..e300351c8 --- /dev/null +++ b/backend/packages/harness/deerflow/persistence/migrations/versions/0005_run_stop_reason.py @@ -0,0 +1,28 @@ +"""run stop_reason + +Revision ID: 0005_run_stop_reason +Revises: 0004_run_ownership +Create Date: 2026-07-15 +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op + +revision: str = "0005_run_stop_reason" +down_revision: str | Sequence[str] | None = "0004_run_ownership" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + from deerflow.persistence.migrations._helpers import safe_add_column + + safe_add_column("runs", sa.Column("stop_reason", sa.String(50), nullable=True)) + + +def downgrade() -> None: + op.drop_column("runs", "stop_reason") diff --git a/backend/packages/harness/deerflow/persistence/run/model.py b/backend/packages/harness/deerflow/persistence/run/model.py index 19b73e032..723a32370 100644 --- a/backend/packages/harness/deerflow/persistence/run/model.py +++ b/backend/packages/harness/deerflow/persistence/run/model.py @@ -25,6 +25,7 @@ class RunRow(Base): metadata_json: Mapped[dict] = mapped_column(JSON, default=dict) kwargs_json: Mapped[dict] = mapped_column(JSON, default=dict) error: Mapped[str | None] = mapped_column(Text) + stop_reason: Mapped[str | None] = mapped_column(String(50)) # Convenience fields (for listing pages without querying RunEventStore) message_count: Mapped[int] = mapped_column(default=0) diff --git a/backend/packages/harness/deerflow/persistence/run/sql.py b/backend/packages/harness/deerflow/persistence/run/sql.py index 11d44af1c..c84e56b3f 100644 --- a/backend/packages/harness/deerflow/persistence/run/sql.py +++ b/backend/packages/harness/deerflow/persistence/run/sql.py @@ -96,6 +96,7 @@ class RunRepository(RunStore): metadata=None, kwargs=None, error=None, + stop_reason: str | None = None, created_at=None, follow_up_to_run_id=None, owner_worker_id: str | None = None, @@ -121,6 +122,7 @@ class RunRepository(RunStore): "metadata_json": self._safe_json(metadata) or {}, "kwargs_json": self._safe_json(kwargs) or {}, "error": error, + "stop_reason": stop_reason, "follow_up_to_run_id": follow_up_to_run_id, "owner_worker_id": owner_worker_id, "lease_expires_at": lease_dt, @@ -203,10 +205,12 @@ class RunRepository(RunStore): result = await session.execute(stmt) return {row.run_id: self._row_to_dict(row) for row in result.scalars()} - async def update_status(self, run_id, status, *, error=None) -> bool: + async def update_status(self, run_id, status, *, error=None, stop_reason=None) -> bool: values: dict[str, Any] = {"status": status, "updated_at": datetime.now(UTC)} if error is not None: values["error"] = error + if stop_reason is not None: + values["stop_reason"] = stop_reason # Guard: only transition rows that are still active. ``interrupted`` is # included because the rollback path goes ``running → interrupted`` # (cancel acknowledged) then ``interrupted → error`` (task finalize). diff --git a/backend/packages/harness/deerflow/runtime/runs/manager.py b/backend/packages/harness/deerflow/runtime/runs/manager.py index 0248b5024..588498a94 100644 --- a/backend/packages/harness/deerflow/runtime/runs/manager.py +++ b/backend/packages/harness/deerflow/runtime/runs/manager.py @@ -181,6 +181,7 @@ class RunRecord: finalizing: bool = False owner_worker_id: str | None = None lease_expires_at: str | None = None + stop_reason: str | None = None class RunManager: @@ -244,7 +245,7 @@ class RunManager: return [record for run_id in run_ids if (record := self._runs.get(run_id)) is not None] @staticmethod - def _store_put_payload(record: RunRecord, *, error: str | None = None) -> dict[str, Any]: + def _store_put_payload(record: RunRecord, *, error: str | None = None, stop_reason: str | None = None) -> dict[str, Any]: payload = { "thread_id": record.thread_id, "assistant_id": record.assistant_id, @@ -260,6 +261,8 @@ class RunManager: } if record.user_id is not None: payload["user_id"] = record.user_id + if record.stop_reason is not None: + payload["stop_reason"] = record.stop_reason return payload async def _call_store_with_retry( @@ -331,16 +334,16 @@ class RunManager: self._store_put_payload(record, error=error), ) - async def _persist_status(self, record: RunRecord, status: RunStatus, *, error: str | None = None) -> bool: + async def _persist_status(self, record: RunRecord, status: RunStatus, *, error: str | None = None, stop_reason: str | None = None) -> bool: """Best-effort persist a status transition to the backing store.""" if self._store is None: return True - row_recovery_payload = self._store_put_payload(record, error=error) + row_recovery_payload = self._store_put_payload(record, error=error, stop_reason=stop_reason) try: updated = await self._call_store_with_retry( "update_status", record.run_id, - lambda: self._store.update_status(record.run_id, status.value, error=error), + lambda: self._store.update_status(record.run_id, status.value, error=error, stop_reason=stop_reason), ) if updated is False: # ``update_status`` is now guarded by ``status IN ('pending','running')``. @@ -407,6 +410,7 @@ class RunManager: first_human_message=row.get("first_human_message"), owner_worker_id=row.get("owner_worker_id"), lease_expires_at=row.get("lease_expires_at"), + stop_reason=row.get("stop_reason"), ) async def update_run_completion(self, run_id: str, **kwargs) -> None: @@ -656,7 +660,7 @@ class RunManager: logger.warning("Failed to map store row for run %s", run_id, exc_info=True) return records_by_id - async def set_status(self, run_id: str, status: RunStatus, *, error: str | None = None) -> None: + async def set_status(self, run_id: str, status: RunStatus, *, error: str | None = None, stop_reason: str | None = None) -> None: """Transition a run to a new status.""" async with self._lock: record = self._runs.get(run_id) @@ -667,7 +671,9 @@ class RunManager: record.updated_at = _now_iso() if error is not None: record.error = error - await self._persist_status(record, status, error=error) + if stop_reason is not None: + record.stop_reason = stop_reason + await self._persist_status(record, status, error=error, stop_reason=stop_reason) logger.info("Run %s -> %s", run_id, status.value) async def set_finalizing(self, run_id: str, finalizing: bool) -> None: diff --git a/backend/packages/harness/deerflow/runtime/runs/store/base.py b/backend/packages/harness/deerflow/runtime/runs/store/base.py index f8420f8e3..cd6e9eae0 100644 --- a/backend/packages/harness/deerflow/runtime/runs/store/base.py +++ b/backend/packages/harness/deerflow/runtime/runs/store/base.py @@ -29,6 +29,7 @@ class RunStore(abc.ABC): metadata: dict[str, Any] | None = None, kwargs: dict[str, Any] | None = None, error: str | None = None, + stop_reason: str | None = None, created_at: str | None = None, owner_worker_id: str | None = None, lease_expires_at: str | None = None, @@ -84,6 +85,7 @@ class RunStore(abc.ABC): status: str, *, error: str | None = None, + stop_reason: str | None = None, ) -> bool | None: """Update a run status. diff --git a/backend/packages/harness/deerflow/runtime/runs/store/memory.py b/backend/packages/harness/deerflow/runtime/runs/store/memory.py index 97fc49e9f..a91a958f4 100644 --- a/backend/packages/harness/deerflow/runtime/runs/store/memory.py +++ b/backend/packages/harness/deerflow/runtime/runs/store/memory.py @@ -45,6 +45,7 @@ class MemoryRunStore(RunStore): metadata=None, kwargs=None, error=None, + stop_reason=None, created_at=None, owner_worker_id=None, lease_expires_at=None, @@ -61,6 +62,7 @@ class MemoryRunStore(RunStore): "metadata": metadata or {}, "kwargs": kwargs or {}, "error": error, + "stop_reason": stop_reason, "created_at": created_at or now, "updated_at": now, "owner_worker_id": owner_worker_id, @@ -105,7 +107,7 @@ class MemoryRunStore(RunStore): thread_run_ids = self._runs_by_thread.get(thread_id) or () return {run_id: run for run_id in thread_run_ids if run_id in run_ids and (run := self._runs.get(run_id)) is not None and (user_id is None or run.get("user_id") == user_id)} - async def update_status(self, run_id, status, *, error=None): + async def update_status(self, run_id, status, *, error=None, stop_reason=None): run = self._runs.get(run_id) if run is None: return False @@ -116,6 +118,8 @@ class MemoryRunStore(RunStore): run["status"] = status if error is not None: run["error"] = error + if stop_reason is not None: + run["stop_reason"] = stop_reason run["updated_at"] = datetime.now(UTC).isoformat() return True diff --git a/backend/packages/harness/deerflow/runtime/runs/worker.py b/backend/packages/harness/deerflow/runtime/runs/worker.py index 4f23065d4..f8dc08a53 100644 --- a/backend/packages/harness/deerflow/runtime/runs/worker.py +++ b/backend/packages/harness/deerflow/runtime/runs/worker.py @@ -517,6 +517,12 @@ async def run_agent( await subagent_events.add(chunk) # 7. Stream the requested turn, then optionally continue hidden goal turns. + # Clear any stale stop_reason before the first (user-visible) turn only. + # Continuation turns preserve a cap reason from the user turn: a run that + # hits a cap during the user turn IS capped even if hidden goal-evaluator + # turns complete cleanly afterward (#4176 review). + if isinstance(runtime.context, dict): + runtime.context.pop("stop_reason", None) await _stream_once(graph_input, initial_runnable_config) while not record.abort_event.is_set() and not llm_error_fallback_message and (journal is None or not journal.had_llm_error_fallback): continuation_input = await _prepare_goal_continuation_input( @@ -562,7 +568,22 @@ async def run_agent( error_msg = error_msg or "LLM provider failed after retries" await run_manager.set_status(run_id, RunStatus.error, error=error_msg) else: - await run_manager.set_status(run_id, RunStatus.success) + runtime_context = runtime.context if isinstance(runtime.context, dict) else None + # Guard middlewares that hard-stop a run by stripping tool_calls + # stamp stop_reason into runtime.context so the worker can surface + # it on the run record: + # loop_detection -> "loop_capped" + # token_budget -> "token_capped" + # safety_finish_reason -> "safety_capped" + # subagent_limit -> "subagent_limit_capped" + # + # If more guards grow stop_reason semantics, consider a publish/ + # collect pattern (e.g. each guard middleware publishes its cap + # reason to a dedicated runtime.context channel, and the worker + # collects the most severe / first / all reasons) instead of each + # guard writing directly to the same key. + stop_reason = runtime_context.get("stop_reason") if runtime_context is not None else None + await run_manager.set_status(run_id, RunStatus.success, stop_reason=stop_reason) except asyncio.CancelledError: await run_manager.set_finalizing(run_id, True) diff --git a/backend/tests/test_loop_detection_stop_reason.py b/backend/tests/test_loop_detection_stop_reason.py new file mode 100644 index 000000000..ac66042e8 --- /dev/null +++ b/backend/tests/test_loop_detection_stop_reason.py @@ -0,0 +1,280 @@ +"""Integration tests: verify that guard middlewares write ``stop_reason`` to +``runtime.context`` and the worker surfaces it on the run record (#4176). + +The lead worker calls ``agent.astream()``. During streaming, guard +middlewares (loop detection, token budget) may detect a cap and write +``stop_reason`` into ``runtime.context``. After streaming completes, the +worker reads ``runtime.context["stop_reason"]`` and persists it. + +The key invariant: the middleware's ``runtime.context`` IS the worker's +``runtime.context`` — LangGraph surfaces the same dict — so the worker +sees whatever the middleware wrote. + +These tests exercise that invariant end-to-end, using real middleware +instances (not hand-written simulations of the write) driven inside +``astream`` to prove the full middleware→context→worker pipeline. +""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import AsyncMock + +import pytest +from langchain_core.messages import AIMessage + +from deerflow.runtime import RunContext, RunManager, RunStatus +from deerflow.runtime.runs.worker import run_agent + + +@pytest.mark.asyncio +async def test_worker_surfaces_stop_reason_from_loop_detection(): + """The worker persists ``stop_reason=loop_capped`` when the real + LoopDetectionMiddleware triggers a hard stop during streaming.""" + from deerflow.agents.middlewares.loop_detection_middleware import LoopDetectionMiddleware + + run_manager = RunManager() + record = await run_manager.create("thread-1") + + mw = LoopDetectionMiddleware(warn_threshold=1, hard_limit=3, window_size=5) + 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, "LangGraph Runtime must be in configurable" + captured_runtime[0] = runtime + + # Drive the real middleware to a hard stop with repeated identical + # tool calls. With hard_limit=3, the 3rd identical call fires the + # hard stop, triggering the runtime.context write. + tool_calls = [{"name": "bash", "args": {"command": "ls"}, "id": "c1", "type": "tool_call"}] + for _ in range(2): + mw._apply({"messages": [AIMessage(content="", tool_calls=tool_calls)]}, runtime) + # 3rd call — hard stop fires here. + mw._apply({"messages": [AIMessage(content="", tool_calls=tool_calls)]}, runtime) + + yield {"messages": [AIMessage(content="Done.")]} + + bridge = AsyncMock() + bridge.publish = AsyncMock() + bridge.publish_end = AsyncMock() + bridge.cleanup = AsyncMock() + + def factory(*, config): + return DummyAgent() + + await run_agent( + bridge, + run_manager, + record, + ctx=RunContext(checkpointer=None), + agent_factory=factory, + graph_input={"messages": []}, + config={}, + ) + + # Prove runtime object identity: the DummyAgent captured the same Runtime + # object the worker reads from; if LangGraph's merge created a copy, the + # worker would see a different context dict and stop_reason would be None. + assert captured_runtime[0] is not None, "DummyAgent never captured runtime" + runtime_ctx = captured_runtime[0].context + assert isinstance(runtime_ctx, dict) + assert runtime_ctx.get("stop_reason") == "loop_capped", "The runtime the DummyAgent wrote to is the same one the worker read from" + + fetched = await run_manager.get(record.run_id) + assert fetched is not None + assert fetched.status == RunStatus.success + assert fetched.stop_reason == "loop_capped" + + +@pytest.mark.asyncio +async def test_worker_surfaces_stop_reason_from_token_budget(): + """The worker persists ``stop_reason=token_capped`` when the real + TokenBudgetMiddleware triggers a hard stop during streaming.""" + from deerflow.agents.middlewares.token_budget_middleware import TokenBudgetMiddleware + from deerflow.config.token_budget_config import TokenBudgetConfig + + run_manager = RunManager() + record = await run_manager.create("thread-1") + # Use a moderate budget with hard_stop_threshold=0.0 so even + # modest usage triggers the hard stop immediately. + config = TokenBudgetConfig( + enabled=True, + max_tokens=1000, + hard_stop_threshold=0.0, + warn_threshold=0.0, + ) + mw = TokenBudgetMiddleware(config=config) + 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, "LangGraph Runtime must be in configurable" + captured_runtime[0] = runtime + + # Feed a single AIMessage with token usage that exceeds the budget. + msg = AIMessage( + id="msg-budget", + content="hello", + usage_metadata={"input_tokens": 100, "output_tokens": 50, "total_tokens": 150}, + ) + mw._apply({"messages": [msg]}, runtime) + + yield {"messages": [AIMessage(content="Budget exceeded, wrapping up.")]} + + bridge = AsyncMock() + bridge.publish = AsyncMock() + bridge.publish_end = AsyncMock() + bridge.cleanup = AsyncMock() + + def factory(*, config): + return DummyAgent() + + await run_agent( + bridge, + run_manager, + record, + ctx=RunContext(checkpointer=None), + agent_factory=factory, + graph_input={"messages": []}, + config={}, + ) + # Prove runtime object identity (same rationale as the loop-detection test). + assert captured_runtime[0] is not None, "DummyAgent never captured runtime" + runtime_ctx = captured_runtime[0].context + assert isinstance(runtime_ctx, dict) + assert runtime_ctx.get("stop_reason") == "token_capped" + + fetched = await run_manager.get(record.run_id) + assert fetched is not None + assert fetched.status == RunStatus.success + assert fetched.stop_reason == "token_capped" + + +@pytest.mark.asyncio +async def test_worker_surfaces_stop_reason_from_safety_finish_reason(): + """The worker persists ``stop_reason=safety_capped`` when the real + SafetyFinishReasonMiddleware strips tool_calls on a safety termination.""" + from unittest.mock import MagicMock + + from deerflow.agents.middlewares.safety_finish_reason_middleware import SafetyFinishReasonMiddleware + from deerflow.agents.middlewares.safety_termination_detectors import SafetyTermination + + # A detector that always fires, simulating any provider safety signal. + always_detector = MagicMock() + always_detector.name = "test-always-fire" + always_detector.detect.return_value = SafetyTermination( + detector="test-always-fire", + reason_field="finish_reason", + reason_value="content_filter", + ) + mw = SafetyFinishReasonMiddleware(detectors=[always_detector]) + 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 + + # Feed an AIMessage with tool_calls so the middleware triggers. + msg = AIMessage( + content="I can't do that.", + tool_calls=[{"name": "bash", "args": {}, "id": "c1", "type": "tool_call"}], + response_metadata={"finish_reason": "content_filter"}, + ) + mw._apply({"messages": [msg]}, runtime) + + yield {"messages": [AIMessage(content="Safety filter triggered.")]} + + 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") == "safety_capped" + + fetched = await run_manager.get(record.run_id) + assert fetched is not None + assert fetched.status == RunStatus.success + assert fetched.stop_reason == "safety_capped" + + +@pytest.mark.asyncio +async def test_worker_surfaces_stop_reason_from_subagent_limit(): + """The worker persists ``stop_reason=subagent_limit_capped`` when the + real SubagentLimitMiddleware hits the total per-run cap.""" + from deerflow.agents.middlewares.subagent_limit_middleware import SubagentLimitMiddleware + + # max_total=1: first delegation exhausts the cap. + mw = SubagentLimitMiddleware(max_concurrent=3, max_total=1) + 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 + + run_id = runtime.context.get("run_id") + # Simulate one prior delegation so remaining_total = 0. + state: dict[str, Any] = { + "messages": [ + AIMessage( + content="Delegating...", + tool_calls=[{"name": "task", "args": {"subagent_type": "general-purpose"}, "id": "c1", "type": "tool_call"}], + ) + ], + "delegations": [{"id": "prior-delegation", "run_id": run_id}], + } + mw._truncate_task_calls(state, runtime) + + yield {"messages": [AIMessage(content="Subagent limit reached.")]} + + 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") == "subagent_limit_capped" + + fetched = await run_manager.get(record.run_id) + assert fetched is not None + assert fetched.status == RunStatus.success + assert fetched.stop_reason == "subagent_limit_capped" diff --git a/backend/tests/test_migration_0004_run_ownership_dedupe.py b/backend/tests/test_migration_0004_run_ownership_dedupe.py index db6aa01eb..03d948050 100644 --- a/backend/tests/test_migration_0004_run_ownership_dedupe.py +++ b/backend/tests/test_migration_0004_run_ownership_dedupe.py @@ -156,7 +156,7 @@ async def test_migration_dedupes_duplicate_active_rows_before_unique_index(tmp_p with sqlite3.connect(db_path) as raw: version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone() - assert version_row[0] == "0004_run_ownership" + assert version_row[0] == "0005_run_stop_reason" # Sanity: the invariant the index enforces is now true — at most one # active row per thread. diff --git a/backend/tests/test_multi_worker_run_ownership.py b/backend/tests/test_multi_worker_run_ownership.py index 6f2e985b0..aadd99456 100644 --- a/backend/tests/test_multi_worker_run_ownership.py +++ b/backend/tests/test_multi_worker_run_ownership.py @@ -1349,7 +1349,7 @@ async def test_cancel_returns_taken_over_when_peer_claims_during_local_cancel(): # the race: in-memory cancel succeeds, but store write is blocked. original = store.update_status - async def race_update(run_id, status, *, error=None): + async def race_update(run_id, status, *, error=None, stop_reason=None): # Simulate peer takeover: flip to error before our write lands run = store._runs.get(run_id) if run and run["status"] == "running" and status == "interrupted": @@ -1357,7 +1357,7 @@ async def test_cancel_returns_taken_over_when_peer_claims_during_local_cancel(): run["error"] = "peer takeover" run["updated_at"] = datetime.now(UTC).isoformat() return False # our write was blocked - return await original(run_id, status, error=error) + return await original(run_id, status, error=error, stop_reason=stop_reason) store.update_status = race_update diff --git a/backend/tests/test_persistence_bootstrap.py b/backend/tests/test_persistence_bootstrap.py index 257208c3b..11cfe926c 100644 --- a/backend/tests/test_persistence_bootstrap.py +++ b/backend/tests/test_persistence_bootstrap.py @@ -48,7 +48,7 @@ from deerflow.persistence.migrations._helpers import _normalize_default asyncio_test = pytest.mark.asyncio -HEAD = "0004_run_ownership" +HEAD = "0005_run_stop_reason" BASELINE = "0001_baseline" diff --git a/backend/tests/test_persistence_bootstrap_concurrency.py b/backend/tests/test_persistence_bootstrap_concurrency.py index de41df28f..045c8a9bb 100644 --- a/backend/tests/test_persistence_bootstrap_concurrency.py +++ b/backend/tests/test_persistence_bootstrap_concurrency.py @@ -28,7 +28,7 @@ from deerflow.persistence.bootstrap import bootstrap_schema pytestmark = pytest.mark.asyncio -HEAD = "0004_run_ownership" +HEAD = "0005_run_stop_reason" def _url(tmp_path: Path) -> str: diff --git a/backend/tests/test_persistence_bootstrap_regression.py b/backend/tests/test_persistence_bootstrap_regression.py index 16683db58..6981fb27e 100644 --- a/backend/tests/test_persistence_bootstrap_regression.py +++ b/backend/tests/test_persistence_bootstrap_regression.py @@ -76,7 +76,7 @@ async def test_legacy_database_recovers_token_usage_column(tmp_path: Path) -> No cols = {row[1] for row in raw.execute("PRAGMA table_info(runs)").fetchall()} assert "token_usage_by_model" in cols version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone() - assert version_row[0] == "0004_run_ownership" + assert version_row[0] == "0005_run_stop_reason" # And the read path that originally 500'd must now succeed. sf = get_session_factory() @@ -116,6 +116,6 @@ async def test_legacy_database_with_manual_alter_still_bootstraps(tmp_path: Path # No duplicate column -- list, not set, to catch dupes. assert cols.count("token_usage_by_model") == 1 version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone() - assert version_row[0] == "0004_run_ownership" + assert version_row[0] == "0005_run_stop_reason" finally: await close_engine() diff --git a/backend/tests/test_run_manager.py b/backend/tests/test_run_manager.py index 33d213003..5a61e3fa9 100644 --- a/backend/tests/test_run_manager.py +++ b/backend/tests/test_run_manager.py @@ -29,19 +29,19 @@ class FlakyStatusRunStore(MemoryRunStore): self.status_failures = status_failures self.status_update_attempts = 0 - async def update_status(self, run_id, status, *, error=None): + async def update_status(self, run_id, status, *, error=None, stop_reason=None): self.status_update_attempts += 1 if self.status_failures > 0: self.status_failures -= 1 raise sqlite3.OperationalError("database is locked") - return await super().update_status(run_id, status, error=error) + return await super().update_status(run_id, status, error=error, stop_reason=stop_reason) class MissingRowStatusRunStore(MemoryRunStore): """Memory run store that reports a missing row for status updates.""" - async def update_status(self, run_id, status, *, error=None): - await super().update_status(run_id, status, error=error) + async def update_status(self, run_id, status, *, error=None, stop_reason=None): + await super().update_status(run_id, status, error=error, stop_reason=stop_reason) return False @@ -52,7 +52,7 @@ class PermanentStatusRunStore(MemoryRunStore): super().__init__() self.status_update_attempts = 0 - async def update_status(self, run_id, status, *, error=None): + async def update_status(self, run_id, status, *, error=None, stop_reason=None): self.status_update_attempts += 1 raise SQLAlchemyDatabaseError( "UPDATE runs SET status = :status WHERE run_id = :run_id", @@ -68,7 +68,7 @@ class FailingStatusRunStore(MemoryRunStore): super().__init__() self.status_update_attempts = 0 - async def update_status(self, run_id, status, *, error=None): + async def update_status(self, run_id, status, *, error=None, stop_reason=None): self.status_update_attempts += 1 raise sqlite3.OperationalError("database is locked")