From 7141a8818d8820398d9ecdd39e0f5927aff83a4f Mon Sep 17 00:00:00 2001 From: lda Date: Sat, 5 Sep 2026 21:00:00 +0700 Subject: [PATCH] feat: persist and inspect run step budgets --- src/wf_api/models/runs.py | 3 + src/wf_api/operation_context.py | 3 +- src/wf_api/run_lifecycle.py | 24 +- src/wf_api/runs.py | 34 ++- src/wf_api/service.py | 2 + src/wf_api/surface.py | 1 + src/wf_artifacts/__init__.py | 2 + src/wf_artifacts/runs/__init__.py | 2 + src/wf_artifacts/runs/models.py | 19 +- src/wf_core/runtime/engine.py | 10 +- src/wf_mcp/broker/service/core.py | 3 + .../service/workflow_operation_context.py | 3 + src/wf_mcp/broker/service/workflow_runtime.py | 3 + src/wf_server/context.py | 3 + tests/core/test_run_step_budget.py | 99 ++++++ tests/wf_api/test_run_lifecycle.py | 282 ++++++++++++++++++ tests/wf_api/test_runs.py | 222 ++++++++++++++ 17 files changed, 707 insertions(+), 8 deletions(-) create mode 100644 tests/wf_api/test_run_lifecycle.py create mode 100644 tests/wf_api/test_runs.py diff --git a/src/wf_api/models/runs.py b/src/wf_api/models/runs.py index 87b2d06c..7fe2160e 100644 --- a/src/wf_api/models/runs.py +++ b/src/wf_api/models/runs.py @@ -83,6 +83,9 @@ class RunResultBase(ArtifactVersionPayload, GuidedResultPayload): error: str | None output: JsonObject | None trace_count: int + max_steps: int + steps_executed: int + steps_remaining: int class RunResult(RunResultBase): diff --git a/src/wf_api/operation_context.py b/src/wf_api/operation_context.py index c22baf58..19dff342 100644 --- a/src/wf_api/operation_context.py +++ b/src/wf_api/operation_context.py @@ -13,7 +13,7 @@ from wf_artifacts import ( WorkflowDeployment, ) from wf_authoring import NodeSpec -from wf_core import RunState +from wf_core import RunLimits, RunState from wf_platform import CapabilitySource from .models import RawWorkflowPlan @@ -61,6 +61,7 @@ class WorkflowRuntimeRunner(Protocol): deployment: WorkflowDeployment | None = None, artifact: WorkflowArtifact | None = None, saved_subgraph_tree: SavedSubgraphTree | None = None, + limits: RunLimits | None = None, ) -> RunState: """Execute one raw workflow plan and return its run state.""" ... diff --git a/src/wf_api/run_lifecycle.py b/src/wf_api/run_lifecycle.py index 711c61fd..5a362548 100644 --- a/src/wf_api/run_lifecycle.py +++ b/src/wf_api/run_lifecycle.py @@ -25,6 +25,7 @@ from wf_core import ( RunStatus, dump_run_state, load_run_state, + load_run_state_with_upgrade, ) @@ -100,10 +101,29 @@ def persist_stopped_run( def restore_interrupted_run( store: RunStore, run_id: str ) -> tuple[WorkflowRunRecord, RunState]: - """Load a persisted interrupted run and its latest typed runtime state.""" - record, run = load_stored_run(store, run_id) + """Load a persisted interrupted run, persisting a v1 upgrade first. + + A pre-budget (v1) checkpoint receives its one-time defaults and is + rewritten as a new v2 interrupted checkpoint under the same run id and + pinned environment *before* the run is returned, so resume dispatch + never runs on unmigrated state and a failed upgrade fails resume before + any handler runs. Ordinary inspection uses :func:`load_stored_run`, + which decodes v1 prospectively without mutating the store. + """ + record = store.get_run(run_id) if record.status is not StoredRunStatus.INTERRUPTED: raise ValueError(f"workflow run {run_id!r} is not interrupted") + checkpoint = store.get_latest_checkpoint(run_id) + run, upgraded = load_run_state_with_upgrade( + checkpoint.state.model_dump(mode="json") + ) + if upgraded: + record = persist_stopped_run( + store=store, + environment=record.environment, + run=run, + run_id=run_id, + ) return record, run diff --git a/src/wf_api/runs.py b/src/wf_api/runs.py index 6f427cd2..46978b7c 100644 --- a/src/wf_api/runs.py +++ b/src/wf_api/runs.py @@ -11,7 +11,7 @@ from wf_artifacts import ( WorkflowDeployment, WorkflowRunRecord, ) -from wf_core import RunState +from wf_core import RunLimits, RunState from .artifact_plans import raw_plan_from_artifact from .deployments import WorkflowDeploymentApi, _available_sources @@ -80,6 +80,7 @@ class WorkflowRunApi: deployment_id: str, workflow_input: dict[str, Any], trace_range: TraceRangeLike | None = None, + max_steps: int | None = None, ) -> RunResult: trace_values = _trace_range_values(trace_range) deployment, artifact, diagnostics, tree = ( @@ -94,12 +95,16 @@ class WorkflowRunApi: ) plan = raw_plan_from_artifact(artifact) + limits = ( + RunLimits(max_steps=max_steps) if max_steps is not None else RunLimits() + ) run = await self.context.runtime.run_workflow_from_plan( plan, workflow_input, deployment=deployment, artifact=artifact, saved_subgraph_tree=tree, + limits=limits, ) record = persist_stopped_run( store=self._run_store(), @@ -121,6 +126,9 @@ class WorkflowRunApi: error=run.error, output=run.output, trace_count=len(run.trace), + max_steps=run.limits.max_steps, + steps_executed=run.steps_executed, + steps_remaining=run.steps_remaining, **_trace_slice_fields(run, trace_values), ) @@ -176,6 +184,9 @@ class WorkflowRunApi: output=stopped_run.output, diagnostics=diagnostics, trace_count=len(stopped_run.trace), + max_steps=stopped_run.limits.max_steps, + steps_executed=stopped_run.steps_executed, + steps_remaining=stopped_run.steps_remaining, ) plan = raw_plan_from_artifact(environment.root_artifact) tree = saved_subgraph_tree_from_snapshots(environment.child_artifacts) @@ -205,6 +216,9 @@ class WorkflowRunApi: error=run.error, output=run.output, trace_count=len(run.trace), + max_steps=run.limits.max_steps, + steps_executed=run.steps_executed, + steps_remaining=run.steps_remaining, **_trace_slice_fields(run, trace_values), ) @@ -261,6 +275,9 @@ class WorkflowRunApi: output=run.output, diagnostics=record.diagnostics, trace_count=len(run.trace), + max_steps=run.limits.max_steps, + steps_executed=run.steps_executed, + steps_remaining=run.steps_remaining, ) async def read_run_trace( @@ -281,6 +298,9 @@ class WorkflowRunApi: resume_readiness=record.resume_readiness.value, diagnostics=record.diagnostics, trace_count=len(run.trace), + max_steps=run.limits.max_steps, + steps_executed=run.steps_executed, + steps_remaining=run.steps_remaining, **_trace_slice_fields(run, trace_values), ) # A concrete trace range makes _run_payload include the four trace @@ -366,7 +386,16 @@ def _run_payload( trace_start: int | None = None, trace_limit: int | None = None, trace_truncated: bool = False, + max_steps: int | None = None, + steps_executed: int = 0, + steps_remaining: int | None = None, ) -> RunResult: + effective_max = max_steps if max_steps is not None else RunLimits().max_steps + effective_remaining = ( + steps_remaining + if steps_remaining is not None + else max(effective_max - steps_executed, 0) + ) payload = { "deployment_id": deployment.id, "artifact_id": artifact.id, @@ -382,6 +411,9 @@ def _run_payload( diagnostic.model_dump(mode="json") for diagnostic in diagnostics or [] ], "trace_count": trace_count, + "max_steps": effective_max, + "steps_executed": steps_executed, + "steps_remaining": effective_remaining, "next_actions": NextActions.from_run_result( run_id=run_id, status=status, diff --git a/src/wf_api/service.py b/src/wf_api/service.py index 74780849..30ff67b9 100644 --- a/src/wf_api/service.py +++ b/src/wf_api/service.py @@ -1048,11 +1048,13 @@ class WorkflowApi: deployment_id: str, workflow_input: dict[str, Any], trace_range: TraceRangeLike | None = None, + max_steps: int | None = None, ) -> RunResult: return await self.runs.run_deployment( deployment_id=deployment_id, workflow_input=workflow_input, trace_range=trace_range, + max_steps=max_steps, ) async def resume_run( diff --git a/src/wf_api/surface.py b/src/wf_api/surface.py index 59732c22..d8c3a12f 100644 --- a/src/wf_api/surface.py +++ b/src/wf_api/surface.py @@ -517,6 +517,7 @@ class WorkflowRunSurface(Protocol): deployment_id: str, workflow_input: dict[str, Any], trace_range: TraceRangeLike | None = None, + max_steps: int | None = None, ) -> RunResult: ... async def resume_run( diff --git a/src/wf_artifacts/__init__.py b/src/wf_artifacts/__init__.py index 7e9d98b6..6773d9bd 100644 --- a/src/wf_artifacts/__init__.py +++ b/src/wf_artifacts/__init__.py @@ -49,6 +49,7 @@ from .runs import ( RunCheckpoint, RunStore, StoredRunStatus, + VersionedCheckpointState, WorkflowRunRecord, ensure_run_id, ) @@ -75,6 +76,7 @@ __all__ = [ "RunStore", "SourceBinding", "StoredRunStatus", + "VersionedCheckpointState", "WorkflowArtifact", "WorkflowArtifactCatalogEntry", "WorkflowArtifactStore", diff --git a/src/wf_artifacts/runs/__init__.py b/src/wf_artifacts/runs/__init__.py index 11f0188a..172efe7c 100644 --- a/src/wf_artifacts/runs/__init__.py +++ b/src/wf_artifacts/runs/__init__.py @@ -4,6 +4,7 @@ from .models import ( ResumeReadiness, RunCheckpoint, StoredRunStatus, + VersionedCheckpointState, WorkflowRunRecord, ensure_run_id, ) @@ -17,6 +18,7 @@ __all__ = [ "RunCheckpoint", "RunStore", "StoredRunStatus", + "VersionedCheckpointState", "WorkflowRunRecord", "ensure_run_id", ] diff --git a/src/wf_artifacts/runs/models.py b/src/wf_artifacts/runs/models.py index 22f53671..49eb5238 100644 --- a/src/wf_artifacts/runs/models.py +++ b/src/wf_artifacts/runs/models.py @@ -3,6 +3,7 @@ from __future__ import annotations import re from datetime import datetime from enum import StrEnum +from typing import Any, Literal from pydantic import BaseModel, ConfigDict, Field @@ -72,6 +73,22 @@ class WorkflowRunRecord(BaseModel): updated_at: datetime +class VersionedCheckpointState(BaseModel): + """Lenient read envelope for stopped-run checkpoints. + + Writes always produce version 2 via ``wf_core.dump_run_state``; reads + accept version 1 so pre-budget checkpoints reach + ``load_run_state_with_upgrade`` instead of failing checkpoint validation + with a ``version == 2`` literal error first. The inner state stays an + untyped dict because core owns strict budget validation there. + """ + + model_config = ConfigDict(extra="forbid") + + version: Literal[1, 2] = 2 + state: dict[str, Any] + + class RunCheckpoint(BaseModel): """One stopped-state snapshot persisted at an external run boundary.""" @@ -81,5 +98,5 @@ class RunCheckpoint(BaseModel): run_id: str = Field(pattern=RUN_ID_PATTERN) sequence: int = Field(ge=1) reason: CheckpointReason - state: PersistedRunState + state: PersistedRunState | VersionedCheckpointState created_at: datetime diff --git a/src/wf_core/runtime/engine.py b/src/wf_core/runtime/engine.py index d4b54b71..17869080 100644 --- a/src/wf_core/runtime/engine.py +++ b/src/wf_core/runtime/engine.py @@ -6,6 +6,7 @@ from typing import Any from wf_core.errors import WorkflowExecutionError from wf_core.models.workflow import Workflow from wf_core.run_state import ROOT_SCOPE_ID, RunState, RunStatus +from wf_core.runtime.limits import RunLimits from wf_core.runtime.ops.flow import finalize_run from wf_core.runtime.ops.merges import ReducerDefinition from wf_core.runtime.ops.nodes import AsyncNodeHandler, NodeHandler @@ -25,9 +26,10 @@ def execute_workflow( *, reducers: Mapping[str, ReducerDefinition] | None = None, subgraphs: Mapping[str, PreparedSubgraph[NodeHandler]] | None = None, + limits: RunLimits | None = None, ) -> RunState: """Create a run and execute a workflow synchronously until it stops.""" - run = create_run_state(workflow, workflow_input) + run = create_run_state(workflow, workflow_input, limits=limits) try: prepare_new_run(workflow, workflow_input, run) @@ -52,9 +54,10 @@ async def execute_workflow_async( reducers: Mapping[str, ReducerDefinition] | None = None, subgraphs: Mapping[str, PreparedSubgraph[AsyncNodeHandler]] | None = None, platform: object | None = None, + limits: RunLimits | None = None, ) -> RunState: """Create a run and execute a workflow asynchronously until it stops.""" - run = create_run_state(workflow, workflow_input) + run = create_run_state(workflow, workflow_input, limits=limits) try: prepare_new_run(workflow, workflow_input, run) @@ -80,9 +83,10 @@ async def execute_workflow_result_async( reducers: Mapping[str, ReducerDefinition] | None = None, subgraphs: Mapping[str, PreparedSubgraph[AsyncNodeHandler]] | None = None, platform: object | None = None, + limits: RunLimits | None = None, ) -> RunState: """Execute asynchronously and return failed state instead of raising failures.""" - run = create_run_state(workflow, workflow_input) + run = create_run_state(workflow, workflow_input, limits=limits) try: prepare_new_run(workflow, workflow_input, run) diff --git a/src/wf_mcp/broker/service/core.py b/src/wf_mcp/broker/service/core.py index 4ba728b7..3ea98e9d 100644 --- a/src/wf_mcp/broker/service/core.py +++ b/src/wf_mcp/broker/service/core.py @@ -15,6 +15,7 @@ from wf_artifacts import ( ) from wf_authoring import NodeSpec from wf_core import ( + RunLimits, RunState, Workflow, ) @@ -343,6 +344,7 @@ class WfMcpService: deployment: WorkflowDeployment | None = None, artifact: WorkflowArtifact | None = None, saved_subgraph_tree: SavedSubgraphTree | None = None, + limits: RunLimits | None = None, ): return await self.workflow_runtime.run_workflow_from_plan( plan, @@ -350,6 +352,7 @@ class WfMcpService: deployment=deployment, artifact=artifact, saved_subgraph_tree=saved_subgraph_tree, + limits=limits, ) async def resume_workflow_from_plan( diff --git a/src/wf_mcp/broker/service/workflow_operation_context.py b/src/wf_mcp/broker/service/workflow_operation_context.py index 5be52ef2..edb400d1 100644 --- a/src/wf_mcp/broker/service/workflow_operation_context.py +++ b/src/wf_mcp/broker/service/workflow_operation_context.py @@ -13,6 +13,7 @@ from wf_api.operation_context import ( ) from wf_artifacts import DependencyDiagnostic, WorkflowArtifact, WorkflowDeployment from wf_authoring import NodeSpec +from wf_core import RunLimits from .core import WfMcpService from .events import BrokerEventRecorder @@ -70,6 +71,7 @@ class WfMcpWorkflowRuntimeRunner(WorkflowRuntimeRunner): deployment=None, artifact=None, saved_subgraph_tree=None, + limits: RunLimits | None = None, ): return await self.runtime.run_workflow_from_plan( plan, @@ -77,6 +79,7 @@ class WfMcpWorkflowRuntimeRunner(WorkflowRuntimeRunner): deployment=deployment, artifact=artifact, saved_subgraph_tree=saved_subgraph_tree, + limits=limits, ) async def resume_workflow_from_plan( diff --git a/src/wf_mcp/broker/service/workflow_runtime.py b/src/wf_mcp/broker/service/workflow_runtime.py index 6573c34f..705ef5b3 100644 --- a/src/wf_mcp/broker/service/workflow_runtime.py +++ b/src/wf_mcp/broker/service/workflow_runtime.py @@ -16,6 +16,7 @@ from wf_artifacts import WorkflowArtifact, WorkflowArtifactStore, WorkflowDeploy from wf_authoring import NodeSpec from wf_core import ( NodeUse, + RunLimits, RunState, RunStatus, Workflow, @@ -163,6 +164,7 @@ class WorkflowRuntimeService: deployment: WorkflowDeployment | None = None, artifact: WorkflowArtifact | None = None, saved_subgraph_tree: SavedSubgraphTree | None = None, + limits: RunLimits | None = None, ) -> RunState: self.emit_event( make_event( @@ -186,6 +188,7 @@ class WorkflowRuntimeService: reducers=reducers, subgraphs=prepared_subgraphs, platform=platform_context, + limits=limits, ) self.emit_event( make_event( diff --git a/src/wf_server/context.py b/src/wf_server/context.py index 0496fa05..486a0559 100644 --- a/src/wf_server/context.py +++ b/src/wf_server/context.py @@ -33,6 +33,7 @@ from wf_artifacts import WorkflowArtifact, WorkflowDeployment from wf_authoring import NodeSpec from wf_core import ( NodeUse, + RunLimits, RunState, Workflow, execute_workflow_result_async, @@ -224,6 +225,7 @@ class LocalWorkflowRuntimeRunner(WorkflowRuntimeRunner): deployment: WorkflowDeployment | None = None, artifact: WorkflowArtifact | None = None, saved_subgraph_tree: SavedSubgraphTree | None = None, + limits: RunLimits | None = None, ) -> RunState: workflow, registry, reducers, prepared_subgraphs, platform_context = ( self.prepare_workflow_runtime( @@ -240,6 +242,7 @@ class LocalWorkflowRuntimeRunner(WorkflowRuntimeRunner): reducers=reducers, subgraphs=prepared_subgraphs, platform=platform_context, + limits=limits, ) async def resume_workflow_from_plan( diff --git a/tests/core/test_run_step_budget.py b/tests/core/test_run_step_budget.py index 96f83d95..ca170505 100644 --- a/tests/core/test_run_step_budget.py +++ b/tests/core/test_run_step_budget.py @@ -1,5 +1,7 @@ from __future__ import annotations +import inspect + import pytest from wf_core import ( @@ -22,8 +24,12 @@ from wf_core import ( WorkflowExecutionError, dump_run_state, execute_workflow, + execute_workflow_async, + execute_workflow_result_async, load_run_state, resume_workflow, + resume_workflow_async, + resume_workflow_result_async, step_workflow, ) from wf_core.errors import WorkflowStepLimitExceeded @@ -34,6 +40,7 @@ from wf_core.runtime.limits import ( remaining_step_attempts, ) from wf_core.runtime.ops.runs import create_run_state +from wf_core.runtime.preparation import prepare_resume def _minimal_workflow(name: str = "budget") -> Workflow: @@ -994,3 +1001,95 @@ def test_v2_missing_interrupt_step_number_is_corrupt() -> None: with pytest.raises(ValueError): load_run_state_with_upgrade(stored) + + +# --- Task 4: engine-level limits seam --- + + +def test_execute_workflow_accepts_explicit_limits() -> None: + workflow = _chain_workflow() + + run = execute_workflow( + workflow, + {}, + {"da": _ok_handler, "db": _ok_handler}, + limits=RunLimits(max_steps=10), + ) + + assert run.status == RunStatus.COMPLETED + assert run.limits.max_steps == 10 + assert run.steps_executed == 2 + assert run.steps_remaining == 8 + + +def test_execute_workflow_defaults_to_ten_thousand() -> None: + workflow = _chain_workflow() + + run = execute_workflow(workflow, {}, {"da": _ok_handler, "db": _ok_handler}) + + assert run.limits.max_steps == 10_000 + assert run.steps_executed == 2 + + +def test_execute_workflow_enforces_limits() -> None: + workflow = _cyclic_workflow() + + def ok_handler(_payload: dict, _context: object) -> dict: + return {"outcome": "ok", "output": {}} + + with pytest.raises(WorkflowStepLimitExceeded): + execute_workflow( + workflow, + {}, + {"da": ok_handler, "db": ok_handler}, + limits=RunLimits(max_steps=2), + ) + + +async def test_execute_workflow_async_accepts_explicit_limits() -> None: + workflow = _chain_workflow() + + async def ok_async(_payload: dict, _context: object) -> dict: + return {"outcome": "ok", "output": {}} + + run = await execute_workflow_async( + workflow, + {}, + {"da": ok_async, "db": ok_async}, + limits=RunLimits(max_steps=10), + ) + + assert run.status == RunStatus.COMPLETED + assert run.limits.max_steps == 10 + assert run.steps_executed == 2 + assert run.steps_remaining == 8 + + +async def test_execute_workflow_result_async_reports_exhaustion() -> None: + workflow = _cyclic_workflow() + + async def ok_async(_payload: dict, _context: object) -> dict: + return {"outcome": "ok", "output": {}} + + run = await execute_workflow_result_async( + workflow, + {}, + {"da": ok_async, "db": ok_async}, + limits=RunLimits(max_steps=2), + ) + + assert run.status == RunStatus.FAILED + assert run.limits.max_steps == 2 + assert run.steps_executed == 2 + assert "step budget" in (run.error or "") + + +def test_resume_entry_points_accept_no_replacement_limits() -> None: + """Ordinary resume reuses persisted limits; it never takes new ones.""" + for entry in ( + resume_workflow, + resume_workflow_async, + resume_workflow_result_async, + prepare_resume, + ): + assert "limits" not in inspect.signature(entry).parameters diff --git a/tests/wf_api/test_run_lifecycle.py b/tests/wf_api/test_run_lifecycle.py new file mode 100644 index 00000000..190b133f --- /dev/null +++ b/tests/wf_api/test_run_lifecycle.py @@ -0,0 +1,282 @@ +"""Stopped-run step budget migration tests (Task 4). + +Pins the v1-to-v2 upgrade contract: a pre-budget interrupted checkpoint is +rewritten as v2 and persisted *before* resume dispatch (the runtime must +observe the upgraded checkpoint when it is called), while ordinary +inspection decodes v1 prospectively without touching the store. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +from wf_api.models import RawWorkflowPlan +from wf_api.operation_context import WorkflowOperationContext +from wf_api.run_lifecycle import ( + create_pinned_environment, + persist_stopped_run, + restore_interrupted_run, +) +from wf_api.runs import WorkflowRunApi +from wf_api.saved_subgraphs import SavedSubgraphTree +from wf_artifacts import ( + FileRunStore, + WorkflowArtifact, + WorkflowDeployment, +) +from wf_authoring import NodeSpec +from wf_core import InterruptRequest, RunState, RunStatus +from wf_platform import CapabilitySource + + +class DummyEvents: + def record_event(self, event: object) -> None: + pass + + def record_workflow_event( + self, + event_type: str, + *, + capability_id: str, + payload: dict[str, Any], + ) -> None: + pass + + +class EmptySpecProvider: + @property + def capability_sources(self) -> dict[str, CapabilitySource]: + return {} + + def get_qualified_spec(self, qualified_name: str) -> NodeSpec[Any, Any]: + raise KeyError(f"unknown capability {qualified_name!r}") + + +class UpgradeAssertingRuntime: + """Resume-only fake proving the v1 upgrade persists before dispatch. + + When the API calls resume, the upgraded v2 checkpoint must already be + the latest persisted checkpoint; otherwise resume dispatched work on top + of unmigrated state. + """ + + def __init__(self, store: FileRunStore, run_id: str) -> None: + self.store = store + self.run_id = run_id + self.resume_calls = 0 + + async def run_workflow_from_plan( + self, + plan: RawWorkflowPlan, + workflow_input: dict[str, Any], + deployment: WorkflowDeployment | None = None, + artifact: WorkflowArtifact | None = None, + saved_subgraph_tree: SavedSubgraphTree | None = None, + limits: Any | None = None, + ) -> RunState: + raise AssertionError("test must not start new workflow runs") + + async def resume_workflow_from_plan( + self, + plan: RawWorkflowPlan, + run: RunState, + *, + resume_payload: dict[str, Any], + resume_outcome: str, + deployment: WorkflowDeployment | None = None, + artifact: WorkflowArtifact | None = None, + saved_subgraph_tree: SavedSubgraphTree | None = None, + ) -> RunState: + self.resume_calls += 1 + latest = self.store.get_latest_checkpoint(self.run_id) + upgraded = latest.state.model_dump(mode="json") + assert upgraded["version"] == 2 + assert upgraded["state"]["limits"] == {"max_steps": 10_000} + assert upgraded["state"]["steps_executed"] == 0 + assert run.limits.max_steps == 10_000 + assert run.steps_executed == 0 + return RunState( + workflow_name=plan.name, + status=RunStatus.COMPLETED, + workflow_input=run.workflow_input, + state={"answer": resume_payload["answer"]}, + outcome=resume_outcome, + output={"answer": resume_payload["answer"]}, + ) + + +def _artifact() -> WorkflowArtifact: + return WorkflowArtifact( + id="pause", + version=1, + title="Pause", + input_schema={"type": "object", "properties": {}}, + output_schema={"type": "object", "properties": {}}, + outcomes=("ok", "submitted"), + plan={ + "name": "pause", + "input_schema": {"type": "object", "properties": {}}, + "state_schema": {"type": "object", "properties": {}}, + "output_schema": {"type": "object", "properties": {}}, + "outcomes": ["ok", "submitted"], + "start": "end_submitted", + "nodes": [{"id": "end_submitted", "type": "end", "outcome": "submitted"}], + "edges": [], + }, + ) + + +def _deployment(artifact: WorkflowArtifact) -> WorkflowDeployment: + return WorkflowDeployment( + id="pause.default", + artifact_id=artifact.id, + artifact_version=artifact.version, + bindings=[], + ) + + +def _seed_interrupted_run(store: FileRunStore) -> str: + artifact = _artifact() + interrupted = RunState( + workflow_name="pause", + status=RunStatus.INTERRUPTED, + workflow_input={"question": "continue?"}, + state={}, + interrupt=InterruptRequest( + id="interrupt:approval", + frame_id="root", + node_id="approval", + kind="approval", + payload={"question": "continue?"}, + ), + ) + record = persist_stopped_run( + store=store, + environment=create_pinned_environment( + deployment=_deployment(artifact), + artifact=artifact, + tree=SavedSubgraphTree(artifacts_by_ref={}, diagnostics=[]), + ), + run=interrupted, + ) + return record.id + + +def _checkpoint_path(store: FileRunStore, run_id: str, sequence: int) -> Path: + return store.runs_dir / run_id / "checkpoints" / f"{sequence:06d}.json" + + +def _read_raw_checkpoint( + store: FileRunStore, run_id: str, sequence: int +) -> dict[str, Any]: + return json.loads(_checkpoint_path(store, run_id, sequence).read_text("utf-8")) + + +def _downgrade_latest_checkpoint_to_v1(store: FileRunStore, run_id: str) -> None: + """Rewrite the latest checkpoint file as a pre-budget v1 envelope. + + Writes raw JSON directly so the file matches what a legacy (pre-budget) + writer left on disk, bypassing current model validation. + """ + path = _checkpoint_path(store, run_id, 1) + payload: dict[str, Any] = json.loads(path.read_text("utf-8")) + inner = payload["state"]["state"] + inner.pop("limits", None) + inner.pop("steps_executed", None) + for frame in inner.get("frames", {}).values(): + frame.pop("step_number", None) + for entry in inner.get("trace", []): + entry.pop("step_number", None) + if inner.get("interrupt") is not None: + inner["interrupt"].pop("step_number", None) + payload["state"] = {"version": 1, "state": inner} + path.write_text(json.dumps(payload, indent=2), encoding="utf-8") + + +def _api(store: FileRunStore, runtime: UpgradeAssertingRuntime) -> WorkflowRunApi: + return WorkflowRunApi( + WorkflowOperationContext( + artifact_store=None, + draft_workspace_store=None, + run_store=store, + events=DummyEvents(), + specs=EmptySpecProvider(), + runtime=runtime, + live_sources=None, + ) + ) + + +def test_restore_interrupted_run_persists_v1_upgrade_before_returning( + tmp_path: Path, +) -> None: + store = FileRunStore(tmp_path / "runs") + run_id = _seed_interrupted_run(store) + _downgrade_latest_checkpoint_to_v1(store, run_id) + + record, run = restore_interrupted_run(store, run_id) + + assert record.id == run_id + assert run.limits.max_steps == 10_000 + assert run.steps_executed == 0 + assert run.steps_remaining == 10_000 + assert [item.sequence for item in store.list_checkpoints(run_id)] == [1, 2] + upgraded = _read_raw_checkpoint(store, run_id, 2) + assert upgraded["state"]["version"] == 2 + assert upgraded["reason"] == "interrupted" + + +def test_restore_interrupted_run_leaves_v2_checkpoints_untouched( + tmp_path: Path, +) -> None: + store = FileRunStore(tmp_path / "runs") + run_id = _seed_interrupted_run(store) + + record, run = restore_interrupted_run(store, run_id) + + assert record.id == run_id + assert run.steps_executed == 0 + assert [item.sequence for item in store.list_checkpoints(run_id)] == [1] + + +async def test_resume_persists_v1_upgrade_before_runtime_dispatch( + tmp_path: Path, +) -> None: + store = FileRunStore(tmp_path / "runs") + run_id = _seed_interrupted_run(store) + _downgrade_latest_checkpoint_to_v1(store, run_id) + runtime = UpgradeAssertingRuntime(store, run_id) + api = _api(store, runtime) + + result = await api.resume_run(run_id=run_id, resume_payload={"answer": "yes"}) + + assert runtime.resume_calls == 1 + assert result["status"] == "completed" + assert result["max_steps"] == 10_000 + assert result["steps_executed"] == 0 + assert result["steps_remaining"] == 10_000 + assert [item.sequence for item in store.list_checkpoints(run_id)] == [1, 2, 3] + upgraded = _read_raw_checkpoint(store, run_id, 2) + assert upgraded["state"]["version"] == 2 + assert upgraded["state"]["state"]["steps_executed"] == 0 + + +async def test_inspect_decodes_v1_without_mutation(tmp_path: Path) -> None: + store = FileRunStore(tmp_path / "runs") + run_id = _seed_interrupted_run(store) + _downgrade_latest_checkpoint_to_v1(store, run_id) + runtime = UpgradeAssertingRuntime(store, run_id) + api = _api(store, runtime) + + summary = await api.inspect_run(run_id=run_id) + + assert summary["status"] == "interrupted" + assert summary["max_steps"] == 10_000 + assert summary["steps_executed"] == 0 + assert summary["steps_remaining"] == 10_000 + assert runtime.resume_calls == 0 + assert [item.sequence for item in store.list_checkpoints(run_id)] == [1] + untouched = _read_raw_checkpoint(store, run_id, 1) + assert untouched["state"]["version"] == 1 diff --git a/tests/wf_api/test_runs.py b/tests/wf_api/test_runs.py new file mode 100644 index 00000000..7783c118 --- /dev/null +++ b/tests/wf_api/test_runs.py @@ -0,0 +1,222 @@ +"""Run step budget creation and inspection tests (Task 4). + +Pins the API surface for persisted step budgets: optional ``max_steps`` on +run creation only, effective ``max_steps``/``steps_executed``/ +``steps_remaining`` on every run result, counter preservation across resume, +and no replacement budget on resume. +""" + +from __future__ import annotations + +import inspect +from pathlib import Path +from typing import Any + +import pytest + +from tests.wf_mcp.test_support import echo_tool +from tests.wf_mcp.workflow_surface.conftest import echo_artifact +from wf_api.runs import WorkflowRunApi +from wf_artifacts import ( + FileRunStore, + FileWorkflowArtifactStore, + WorkflowArtifact, + WorkflowDeployment, +) +from wf_mcp.broker import WfMcpService +from wf_mcp.broker.service.workflow_operation_context import context_from_service +from wf_mcp.models import ConnectionConfig +from wf_mcp.storage import FileStore + + +def _echo_service(root: Path) -> WfMcpService: + artifact_store = FileWorkflowArtifactStore(root) + artifact_store.save_artifact(echo_artifact()) + artifact_store.save_deployment( + WorkflowDeployment( + id="echo.personal", + artifact_id="echo", + artifact_version=1, + bindings=[{"logical_source": "demo", "concrete_source": "demo.personal"}], + ) + ) + service = WfMcpService( + store=FileStore(root / "mcp"), + artifact_store=artifact_store, + run_store=FileRunStore(root / "mcp"), + ) + service.register_connection( + ConnectionConfig(id="demo.personal", server="demo", account="personal") + ) + service.register_specs("demo.personal", echo_tool) + return service + + +def _interrupt_artifact() -> WorkflowArtifact: + return WorkflowArtifact( + id="approval", + version=1, + title="Approval", + input_schema={ + "type": "object", + "properties": {"message": {"type": "string"}}, + "required": ["message"], + }, + output_schema={"type": "object", "properties": {}}, + outcomes=("submitted",), + plan={ + "name": "approval", + "input_schema": { + "type": "object", + "properties": {"message": {"type": "string"}}, + "required": ["message"], + }, + "state_schema": {"fields": {}}, + "output_schema": {"type": "object", "properties": {}}, + "outcomes": ["submitted"], + "start": "approval", + "nodes": [ + { + "id": "approval", + "type": "interrupt", + "kind": "approval", + "request": [ + { + "path": {"root": "input", "parts": ["message"]}, + "target": {"root": "local", "parts": ["message"]}, + } + ], + "resume": [], + "outcomes": ["submitted"], + "resume_schema": { + "type": "object", + "properties": {"approved": {"type": "boolean"}}, + "required": ["approved"], + "additionalProperties": False, + }, + }, + {"id": "end_submitted", "type": "end", "outcome": "submitted"}, + ], + "edges": [ + {"from": "approval", "outcome": "submitted", "to": "end_submitted"} + ], + }, + ) + + +def _interrupt_service(root: Path) -> WfMcpService: + artifact_store = FileWorkflowArtifactStore(root) + artifact_store.save_artifact(_interrupt_artifact()) + artifact_store.save_deployment( + WorkflowDeployment( + id="approval.default", + artifact_id="approval", + artifact_version=1, + bindings=[], + ) + ) + return WfMcpService( + store=FileStore(root / "mcp"), + artifact_store=artifact_store, + run_store=FileRunStore(root / "mcp"), + ) + + +async def test_run_deployment_reports_default_budget(tmp_path: Path) -> None: + api = WorkflowRunApi(context_from_service(_echo_service(tmp_path / "default"))) + + result = await api.run_deployment( + deployment_id="echo.personal", + workflow_input={"text": "hello"}, + ) + + assert result["status"] == "completed" + assert result["max_steps"] == 10_000 + assert result["steps_executed"] == 1 + assert result["steps_remaining"] == 10_000 - result["steps_executed"] + + +async def test_run_deployment_accepts_requested_max_steps(tmp_path: Path) -> None: + api = WorkflowRunApi(context_from_service(_echo_service(tmp_path / "requested"))) + + result = await api.run_deployment( + deployment_id="echo.personal", + workflow_input={"text": "hello"}, + max_steps=5, + ) + + assert result["status"] == "completed" + assert result["max_steps"] == 5 + assert result["steps_executed"] == 1 + assert result["steps_remaining"] == 4 + + +async def test_run_deployment_rejects_non_positive_max_steps(tmp_path: Path) -> None: + api = WorkflowRunApi(context_from_service(_echo_service(tmp_path / "invalid"))) + + with pytest.raises(ValueError, match="positive"): + await api.run_deployment( + deployment_id="echo.personal", + workflow_input={"text": "hello"}, + max_steps=0, + ) + + +async def test_inspect_run_reports_effective_budget(tmp_path: Path) -> None: + api = WorkflowRunApi(context_from_service(_echo_service(tmp_path / "inspect"))) + started = await api.run_deployment( + deployment_id="echo.personal", + workflow_input={"text": "hello"}, + max_steps=7, + ) + run_id = started["run_id"] + assert isinstance(run_id, str) + + summary = await api.inspect_run(run_id=run_id) + + assert summary["max_steps"] == 7 + assert summary["steps_executed"] == started["steps_executed"] + assert summary["steps_remaining"] == 7 - started["steps_executed"] + + +async def test_interrupted_resume_preserves_budget_counter(tmp_path: Path) -> None: + api = WorkflowRunApi(context_from_service(_interrupt_service(tmp_path / "resume"))) + started = await api.run_deployment( + deployment_id="approval.default", + workflow_input={"message": "approve?"}, + max_steps=9, + ) + run_id = started["run_id"] + assert isinstance(run_id, str) + + assert started["status"] == "interrupted" + assert started["max_steps"] == 9 + assert started["steps_executed"] == 1 + assert started["steps_remaining"] == 8 + + resumed = await api.resume_run( + run_id=run_id, + resume_payload={"approved": True}, + resume_outcome="submitted", + ) + + assert resumed["status"] == "completed" + assert resumed["max_steps"] == 9 + assert resumed["steps_executed"] == started["steps_executed"] + 1 + assert resumed["steps_remaining"] == 9 - resumed["steps_executed"] + + +async def test_resume_run_accepts_no_replacement_limit(tmp_path: Path) -> None: + api = WorkflowRunApi(context_from_service(_echo_service(tmp_path / "resume_sig"))) + parameters = inspect.signature(WorkflowRunApi.resume_run).parameters + + assert "max_steps" not in parameters + assert "limits" not in parameters + + extra: dict[str, Any] = {"max_steps": 5} + with pytest.raises(TypeError): + await api.resume_run( + run_id="missing", + resume_payload={}, + **extra, + )