fix: read run stop_reason from runtime context (#4188)

* fix: read run stop_reason from runtime context

* fix: address review feedback for #4188 stop_reason integration

   - migration 0005: use safe_add_column for consistency and drift detection
   - worker: clear runtime.context stop_reason at start of each _stream_once
     turn so a clean continuation doesn't inherit a prior cap reason
   - tests: replace circular unit test with real middleware integration
     tests that exercise LoopDetectionMiddleware._apply and
     TokenBudgetMiddleware._apply through the worker, proving the full
     middleware → runtime.context → persist pipeline

* fix(test): resume conftest

* fix: stamp stop_reason in all guard middlewares, fix clearing semantics
This commit is contained in:
Vanzeren 2026-07-16 08:19:52 +08:00 committed by GitHub
parent de55982c5a
commit 1769b2de0d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
19 changed files with 391 additions and 22 deletions

View File

@ -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,
)

View File

@ -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

View File

@ -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

View File

@ -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)

View File

@ -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)

View File

@ -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")

View File

@ -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)

View File

@ -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).

View File

@ -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:

View File

@ -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.

View File

@ -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

View File

@ -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)

View File

@ -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 middlewarecontextworker 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"

View File

@ -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.

View File

@ -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

View File

@ -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"

View File

@ -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:

View File

@ -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()

View File

@ -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")