feat: persist and inspect run step budgets

This commit is contained in:
lda
2026-09-05 21:00:00 +07:00 Verified
parent 344902c17e
commit 7141a8818d
17 changed files with 707 additions and 8 deletions
+3
View File
@@ -83,6 +83,9 @@ class RunResultBase(ArtifactVersionPayload, GuidedResultPayload):
error: str | None error: str | None
output: JsonObject | None output: JsonObject | None
trace_count: int trace_count: int
max_steps: int
steps_executed: int
steps_remaining: int
class RunResult(RunResultBase): class RunResult(RunResultBase):
+2 -1
View File
@@ -13,7 +13,7 @@ from wf_artifacts import (
WorkflowDeployment, WorkflowDeployment,
) )
from wf_authoring import NodeSpec from wf_authoring import NodeSpec
from wf_core import RunState from wf_core import RunLimits, RunState
from wf_platform import CapabilitySource from wf_platform import CapabilitySource
from .models import RawWorkflowPlan from .models import RawWorkflowPlan
@@ -61,6 +61,7 @@ class WorkflowRuntimeRunner(Protocol):
deployment: WorkflowDeployment | None = None, deployment: WorkflowDeployment | None = None,
artifact: WorkflowArtifact | None = None, artifact: WorkflowArtifact | None = None,
saved_subgraph_tree: SavedSubgraphTree | None = None, saved_subgraph_tree: SavedSubgraphTree | None = None,
limits: RunLimits | None = None,
) -> RunState: ) -> RunState:
"""Execute one raw workflow plan and return its run state.""" """Execute one raw workflow plan and return its run state."""
... ...
+22 -2
View File
@@ -25,6 +25,7 @@ from wf_core import (
RunStatus, RunStatus,
dump_run_state, dump_run_state,
load_run_state, load_run_state,
load_run_state_with_upgrade,
) )
@@ -100,10 +101,29 @@ def persist_stopped_run(
def restore_interrupted_run( def restore_interrupted_run(
store: RunStore, run_id: str store: RunStore, run_id: str
) -> tuple[WorkflowRunRecord, RunState]: ) -> tuple[WorkflowRunRecord, RunState]:
"""Load a persisted interrupted run and its latest typed runtime state.""" """Load a persisted interrupted run, persisting a v1 upgrade first.
record, run = load_stored_run(store, run_id)
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: if record.status is not StoredRunStatus.INTERRUPTED:
raise ValueError(f"workflow run {run_id!r} is not 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 return record, run
+33 -1
View File
@@ -11,7 +11,7 @@ from wf_artifacts import (
WorkflowDeployment, WorkflowDeployment,
WorkflowRunRecord, WorkflowRunRecord,
) )
from wf_core import RunState from wf_core import RunLimits, RunState
from .artifact_plans import raw_plan_from_artifact from .artifact_plans import raw_plan_from_artifact
from .deployments import WorkflowDeploymentApi, _available_sources from .deployments import WorkflowDeploymentApi, _available_sources
@@ -80,6 +80,7 @@ class WorkflowRunApi:
deployment_id: str, deployment_id: str,
workflow_input: dict[str, Any], workflow_input: dict[str, Any],
trace_range: TraceRangeLike | None = None, trace_range: TraceRangeLike | None = None,
max_steps: int | None = None,
) -> RunResult: ) -> RunResult:
trace_values = _trace_range_values(trace_range) trace_values = _trace_range_values(trace_range)
deployment, artifact, diagnostics, tree = ( deployment, artifact, diagnostics, tree = (
@@ -94,12 +95,16 @@ class WorkflowRunApi:
) )
plan = raw_plan_from_artifact(artifact) 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( run = await self.context.runtime.run_workflow_from_plan(
plan, plan,
workflow_input, workflow_input,
deployment=deployment, deployment=deployment,
artifact=artifact, artifact=artifact,
saved_subgraph_tree=tree, saved_subgraph_tree=tree,
limits=limits,
) )
record = persist_stopped_run( record = persist_stopped_run(
store=self._run_store(), store=self._run_store(),
@@ -121,6 +126,9 @@ class WorkflowRunApi:
error=run.error, error=run.error,
output=run.output, output=run.output,
trace_count=len(run.trace), 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), **_trace_slice_fields(run, trace_values),
) )
@@ -176,6 +184,9 @@ class WorkflowRunApi:
output=stopped_run.output, output=stopped_run.output,
diagnostics=diagnostics, diagnostics=diagnostics,
trace_count=len(stopped_run.trace), 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) plan = raw_plan_from_artifact(environment.root_artifact)
tree = saved_subgraph_tree_from_snapshots(environment.child_artifacts) tree = saved_subgraph_tree_from_snapshots(environment.child_artifacts)
@@ -205,6 +216,9 @@ class WorkflowRunApi:
error=run.error, error=run.error,
output=run.output, output=run.output,
trace_count=len(run.trace), 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), **_trace_slice_fields(run, trace_values),
) )
@@ -261,6 +275,9 @@ class WorkflowRunApi:
output=run.output, output=run.output,
diagnostics=record.diagnostics, diagnostics=record.diagnostics,
trace_count=len(run.trace), 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( async def read_run_trace(
@@ -281,6 +298,9 @@ class WorkflowRunApi:
resume_readiness=record.resume_readiness.value, resume_readiness=record.resume_readiness.value,
diagnostics=record.diagnostics, diagnostics=record.diagnostics,
trace_count=len(run.trace), 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), **_trace_slice_fields(run, trace_values),
) )
# A concrete trace range makes _run_payload include the four trace # A concrete trace range makes _run_payload include the four trace
@@ -366,7 +386,16 @@ def _run_payload(
trace_start: int | None = None, trace_start: int | None = None,
trace_limit: int | None = None, trace_limit: int | None = None,
trace_truncated: bool = False, trace_truncated: bool = False,
max_steps: int | None = None,
steps_executed: int = 0,
steps_remaining: int | None = None,
) -> RunResult: ) -> 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 = { payload = {
"deployment_id": deployment.id, "deployment_id": deployment.id,
"artifact_id": artifact.id, "artifact_id": artifact.id,
@@ -382,6 +411,9 @@ def _run_payload(
diagnostic.model_dump(mode="json") for diagnostic in diagnostics or [] diagnostic.model_dump(mode="json") for diagnostic in diagnostics or []
], ],
"trace_count": trace_count, "trace_count": trace_count,
"max_steps": effective_max,
"steps_executed": steps_executed,
"steps_remaining": effective_remaining,
"next_actions": NextActions.from_run_result( "next_actions": NextActions.from_run_result(
run_id=run_id, run_id=run_id,
status=status, status=status,
+2
View File
@@ -1048,11 +1048,13 @@ class WorkflowApi:
deployment_id: str, deployment_id: str,
workflow_input: dict[str, Any], workflow_input: dict[str, Any],
trace_range: TraceRangeLike | None = None, trace_range: TraceRangeLike | None = None,
max_steps: int | None = None,
) -> RunResult: ) -> RunResult:
return await self.runs.run_deployment( return await self.runs.run_deployment(
deployment_id=deployment_id, deployment_id=deployment_id,
workflow_input=workflow_input, workflow_input=workflow_input,
trace_range=trace_range, trace_range=trace_range,
max_steps=max_steps,
) )
async def resume_run( async def resume_run(
+1
View File
@@ -517,6 +517,7 @@ class WorkflowRunSurface(Protocol):
deployment_id: str, deployment_id: str,
workflow_input: dict[str, Any], workflow_input: dict[str, Any],
trace_range: TraceRangeLike | None = None, trace_range: TraceRangeLike | None = None,
max_steps: int | None = None,
) -> RunResult: ... ) -> RunResult: ...
async def resume_run( async def resume_run(
+2
View File
@@ -49,6 +49,7 @@ from .runs import (
RunCheckpoint, RunCheckpoint,
RunStore, RunStore,
StoredRunStatus, StoredRunStatus,
VersionedCheckpointState,
WorkflowRunRecord, WorkflowRunRecord,
ensure_run_id, ensure_run_id,
) )
@@ -75,6 +76,7 @@ __all__ = [
"RunStore", "RunStore",
"SourceBinding", "SourceBinding",
"StoredRunStatus", "StoredRunStatus",
"VersionedCheckpointState",
"WorkflowArtifact", "WorkflowArtifact",
"WorkflowArtifactCatalogEntry", "WorkflowArtifactCatalogEntry",
"WorkflowArtifactStore", "WorkflowArtifactStore",
+2
View File
@@ -4,6 +4,7 @@ from .models import (
ResumeReadiness, ResumeReadiness,
RunCheckpoint, RunCheckpoint,
StoredRunStatus, StoredRunStatus,
VersionedCheckpointState,
WorkflowRunRecord, WorkflowRunRecord,
ensure_run_id, ensure_run_id,
) )
@@ -17,6 +18,7 @@ __all__ = [
"RunCheckpoint", "RunCheckpoint",
"RunStore", "RunStore",
"StoredRunStatus", "StoredRunStatus",
"VersionedCheckpointState",
"WorkflowRunRecord", "WorkflowRunRecord",
"ensure_run_id", "ensure_run_id",
] ]
+18 -1
View File
@@ -3,6 +3,7 @@ from __future__ import annotations
import re import re
from datetime import datetime from datetime import datetime
from enum import StrEnum from enum import StrEnum
from typing import Any, Literal
from pydantic import BaseModel, ConfigDict, Field from pydantic import BaseModel, ConfigDict, Field
@@ -72,6 +73,22 @@ class WorkflowRunRecord(BaseModel):
updated_at: datetime 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): class RunCheckpoint(BaseModel):
"""One stopped-state snapshot persisted at an external run boundary.""" """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) run_id: str = Field(pattern=RUN_ID_PATTERN)
sequence: int = Field(ge=1) sequence: int = Field(ge=1)
reason: CheckpointReason reason: CheckpointReason
state: PersistedRunState state: PersistedRunState | VersionedCheckpointState
created_at: datetime created_at: datetime
+7 -3
View File
@@ -6,6 +6,7 @@ from typing import Any
from wf_core.errors import WorkflowExecutionError from wf_core.errors import WorkflowExecutionError
from wf_core.models.workflow import Workflow from wf_core.models.workflow import Workflow
from wf_core.run_state import ROOT_SCOPE_ID, RunState, RunStatus 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.flow import finalize_run
from wf_core.runtime.ops.merges import ReducerDefinition from wf_core.runtime.ops.merges import ReducerDefinition
from wf_core.runtime.ops.nodes import AsyncNodeHandler, NodeHandler from wf_core.runtime.ops.nodes import AsyncNodeHandler, NodeHandler
@@ -25,9 +26,10 @@ def execute_workflow(
*, *,
reducers: Mapping[str, ReducerDefinition] | None = None, reducers: Mapping[str, ReducerDefinition] | None = None,
subgraphs: Mapping[str, PreparedSubgraph[NodeHandler]] | None = None, subgraphs: Mapping[str, PreparedSubgraph[NodeHandler]] | None = None,
limits: RunLimits | None = None,
) -> RunState: ) -> RunState:
"""Create a run and execute a workflow synchronously until it stops.""" """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: try:
prepare_new_run(workflow, workflow_input, run) prepare_new_run(workflow, workflow_input, run)
@@ -52,9 +54,10 @@ async def execute_workflow_async(
reducers: Mapping[str, ReducerDefinition] | None = None, reducers: Mapping[str, ReducerDefinition] | None = None,
subgraphs: Mapping[str, PreparedSubgraph[AsyncNodeHandler]] | None = None, subgraphs: Mapping[str, PreparedSubgraph[AsyncNodeHandler]] | None = None,
platform: object | None = None, platform: object | None = None,
limits: RunLimits | None = None,
) -> RunState: ) -> RunState:
"""Create a run and execute a workflow asynchronously until it stops.""" """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: try:
prepare_new_run(workflow, workflow_input, run) prepare_new_run(workflow, workflow_input, run)
@@ -80,9 +83,10 @@ async def execute_workflow_result_async(
reducers: Mapping[str, ReducerDefinition] | None = None, reducers: Mapping[str, ReducerDefinition] | None = None,
subgraphs: Mapping[str, PreparedSubgraph[AsyncNodeHandler]] | None = None, subgraphs: Mapping[str, PreparedSubgraph[AsyncNodeHandler]] | None = None,
platform: object | None = None, platform: object | None = None,
limits: RunLimits | None = None,
) -> RunState: ) -> RunState:
"""Execute asynchronously and return failed state instead of raising failures.""" """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: try:
prepare_new_run(workflow, workflow_input, run) prepare_new_run(workflow, workflow_input, run)
+3
View File
@@ -15,6 +15,7 @@ from wf_artifacts import (
) )
from wf_authoring import NodeSpec from wf_authoring import NodeSpec
from wf_core import ( from wf_core import (
RunLimits,
RunState, RunState,
Workflow, Workflow,
) )
@@ -343,6 +344,7 @@ class WfMcpService:
deployment: WorkflowDeployment | None = None, deployment: WorkflowDeployment | None = None,
artifact: WorkflowArtifact | None = None, artifact: WorkflowArtifact | None = None,
saved_subgraph_tree: SavedSubgraphTree | None = None, saved_subgraph_tree: SavedSubgraphTree | None = None,
limits: RunLimits | None = None,
): ):
return await self.workflow_runtime.run_workflow_from_plan( return await self.workflow_runtime.run_workflow_from_plan(
plan, plan,
@@ -350,6 +352,7 @@ class WfMcpService:
deployment=deployment, deployment=deployment,
artifact=artifact, artifact=artifact,
saved_subgraph_tree=saved_subgraph_tree, saved_subgraph_tree=saved_subgraph_tree,
limits=limits,
) )
async def resume_workflow_from_plan( async def resume_workflow_from_plan(
@@ -13,6 +13,7 @@ from wf_api.operation_context import (
) )
from wf_artifacts import DependencyDiagnostic, WorkflowArtifact, WorkflowDeployment from wf_artifacts import DependencyDiagnostic, WorkflowArtifact, WorkflowDeployment
from wf_authoring import NodeSpec from wf_authoring import NodeSpec
from wf_core import RunLimits
from .core import WfMcpService from .core import WfMcpService
from .events import BrokerEventRecorder from .events import BrokerEventRecorder
@@ -70,6 +71,7 @@ class WfMcpWorkflowRuntimeRunner(WorkflowRuntimeRunner):
deployment=None, deployment=None,
artifact=None, artifact=None,
saved_subgraph_tree=None, saved_subgraph_tree=None,
limits: RunLimits | None = None,
): ):
return await self.runtime.run_workflow_from_plan( return await self.runtime.run_workflow_from_plan(
plan, plan,
@@ -77,6 +79,7 @@ class WfMcpWorkflowRuntimeRunner(WorkflowRuntimeRunner):
deployment=deployment, deployment=deployment,
artifact=artifact, artifact=artifact,
saved_subgraph_tree=saved_subgraph_tree, saved_subgraph_tree=saved_subgraph_tree,
limits=limits,
) )
async def resume_workflow_from_plan( async def resume_workflow_from_plan(
@@ -16,6 +16,7 @@ from wf_artifacts import WorkflowArtifact, WorkflowArtifactStore, WorkflowDeploy
from wf_authoring import NodeSpec from wf_authoring import NodeSpec
from wf_core import ( from wf_core import (
NodeUse, NodeUse,
RunLimits,
RunState, RunState,
RunStatus, RunStatus,
Workflow, Workflow,
@@ -163,6 +164,7 @@ class WorkflowRuntimeService:
deployment: WorkflowDeployment | None = None, deployment: WorkflowDeployment | None = None,
artifact: WorkflowArtifact | None = None, artifact: WorkflowArtifact | None = None,
saved_subgraph_tree: SavedSubgraphTree | None = None, saved_subgraph_tree: SavedSubgraphTree | None = None,
limits: RunLimits | None = None,
) -> RunState: ) -> RunState:
self.emit_event( self.emit_event(
make_event( make_event(
@@ -186,6 +188,7 @@ class WorkflowRuntimeService:
reducers=reducers, reducers=reducers,
subgraphs=prepared_subgraphs, subgraphs=prepared_subgraphs,
platform=platform_context, platform=platform_context,
limits=limits,
) )
self.emit_event( self.emit_event(
make_event( make_event(
+3
View File
@@ -33,6 +33,7 @@ from wf_artifacts import WorkflowArtifact, WorkflowDeployment
from wf_authoring import NodeSpec from wf_authoring import NodeSpec
from wf_core import ( from wf_core import (
NodeUse, NodeUse,
RunLimits,
RunState, RunState,
Workflow, Workflow,
execute_workflow_result_async, execute_workflow_result_async,
@@ -224,6 +225,7 @@ class LocalWorkflowRuntimeRunner(WorkflowRuntimeRunner):
deployment: WorkflowDeployment | None = None, deployment: WorkflowDeployment | None = None,
artifact: WorkflowArtifact | None = None, artifact: WorkflowArtifact | None = None,
saved_subgraph_tree: SavedSubgraphTree | None = None, saved_subgraph_tree: SavedSubgraphTree | None = None,
limits: RunLimits | None = None,
) -> RunState: ) -> RunState:
workflow, registry, reducers, prepared_subgraphs, platform_context = ( workflow, registry, reducers, prepared_subgraphs, platform_context = (
self.prepare_workflow_runtime( self.prepare_workflow_runtime(
@@ -240,6 +242,7 @@ class LocalWorkflowRuntimeRunner(WorkflowRuntimeRunner):
reducers=reducers, reducers=reducers,
subgraphs=prepared_subgraphs, subgraphs=prepared_subgraphs,
platform=platform_context, platform=platform_context,
limits=limits,
) )
async def resume_workflow_from_plan( async def resume_workflow_from_plan(
+99
View File
@@ -1,5 +1,7 @@
from __future__ import annotations from __future__ import annotations
import inspect
import pytest import pytest
from wf_core import ( from wf_core import (
@@ -22,8 +24,12 @@ from wf_core import (
WorkflowExecutionError, WorkflowExecutionError,
dump_run_state, dump_run_state,
execute_workflow, execute_workflow,
execute_workflow_async,
execute_workflow_result_async,
load_run_state, load_run_state,
resume_workflow, resume_workflow,
resume_workflow_async,
resume_workflow_result_async,
step_workflow, step_workflow,
) )
from wf_core.errors import WorkflowStepLimitExceeded from wf_core.errors import WorkflowStepLimitExceeded
@@ -34,6 +40,7 @@ from wf_core.runtime.limits import (
remaining_step_attempts, remaining_step_attempts,
) )
from wf_core.runtime.ops.runs import create_run_state 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: 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): with pytest.raises(ValueError):
load_run_state_with_upgrade(stored) 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
+282
View File
@@ -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
+222
View File
@@ -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,
)