lda.chat/lda-wf
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: expose run step budgets to clients

Browse Source
This commit is contained in:
lda
2026-09-05 21:26:31 +07:00 Verified
parent 7141a8818d
commit 715035dce8
18 changed files with 458 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
contracts/workflow-api.manifest.json
+42 -1
View File
@@ -3478,6 +3478,9 @@
}
]
},
"max_steps": {
"type": "integer"
},
"next_actions": {
"$ref": "#/components/schemas/NextActionsPayload"
},
@@ -3524,6 +3527,12 @@
"status": {
"$ref": "#/components/schemas/RunStatus"
},
"steps_executed": {
"type": "integer"
},
"steps_remaining": {
"type": "integer"
},
"trace": {
"items": {
"$ref": "#/components/schemas/TraceEntryPayload"
@@ -3556,7 +3565,10 @@
"outcome",
"error",
"output",
"trace_count"
"trace_count",
"max_steps",
"steps_executed",
"steps_remaining"
],
"type": "object"
},
@@ -3644,6 +3656,9 @@
}
]
},
"max_steps": {
"type": "integer"
},
"next_actions": {
"$ref": "#/components/schemas/NextActionsPayload"
},
@@ -3690,6 +3705,12 @@
"status": {
"$ref": "#/components/schemas/RunStatus"
},
"steps_executed": {
"type": "integer"
},
"steps_remaining": {
"type": "integer"
},
"trace": {
"items": {
"$ref": "#/components/schemas/TraceEntryPayload"
@@ -3723,6 +3744,9 @@
"error",
"output",
"trace_count",
"max_steps",
"steps_executed",
"steps_remaining",
"trace",
"trace_start",
"trace_limit",
@@ -8760,6 +8784,23 @@
],
"default": null
}
},
{
"name": "max_steps",
"required": false,
"schema": {
"anyOf": [
{
"minimum": 1,
"type": "integer"
},
{
"type": "null"
}
],
"default": null,
"description": "Optional run step budget. The server default applies when omitted; resume never accepts a replacement."
}
}
],
"result": {
src/wf_cli/commands/runs.py
+9
View File
@@ -70,6 +70,14 @@ def start_run(
"--trace-limit", min=1, max=100, help="Optional trace slice limit."
),
] = None,
max_steps: Annotated[
int | None,
typer.Option(
"--max-steps",
min=1,
help="Optional run step budget; the server default applies when omitted.",
),
] = None,
) -> None:
"""Start one workflow deployment."""
try:
@@ -84,6 +92,7 @@ def start_run(
deployment_id=deployment_id,
workflow_input=workflow_input,
trace_range=trace_range,
max_steps=max_steps,
),
)
emit_json(payload)
src/wf_client/_http_port.py
+2
View File
@@ -259,6 +259,7 @@ class PublicErrorWorkflowClientPort:
deployment_id: str,
workflow_input: dict[str, Any],
trace_range: TraceRangeLike | None = None,
max_steps: int | None = None,
) -> RunResult:
return await self._invoke(
"workflow.runs.start",
@@ -266,6 +267,7 @@ class PublicErrorWorkflowClientPort:
deployment_id=deployment_id,
workflow_input=workflow_input,
trace_range=trace_range,
max_steps=max_steps,
)
async def inspect_run(self, *, run_id: str) -> RunResult:
src/wf_client/codec.py
+12 -1
View File
@@ -51,6 +51,9 @@ class _DecodedRunFields:
error: str | None
output: dict[str, Any] | None
trace_count: int
max_steps: int
steps_executed: int
steps_remaining: int
diagnostics: tuple[DependencyDiagnostic, ...]
next_actions: dict[str, Any]
trace: tuple[dict[str, Any], ...] | None = None
@@ -255,6 +258,9 @@ def _decode_run_fields(
error=wire["error"],
output=dict(wire["output"]) if wire["output"] is not None else None,
trace_count=wire["trace_count"],
max_steps=wire["max_steps"],
steps_executed=wire["steps_executed"],
steps_remaining=wire["steps_remaining"],
diagnostics=diagnostics,
next_actions=dict(wire["next_actions"]),
trace=trace,
@@ -269,7 +275,12 @@ def decode_run_result(
*,
operation: str = "workflow.runs.inspect",
) -> DecodedRunResult:
"""Validate and decode a start/inspect/resume run response."""
"""Validate and decode a start/inspect/resume run response.
``RunResult`` declares the ``max_steps``/``steps_executed``/
``steps_remaining`` trio as required, so a response missing any of them
fails validation here instead of reaching client reconstruction.
"""
wire = _validate(payload, RunResult, operation)
fields = _decode_run_fields(
wire,
src/wf_client/deployments.py
+16 -6
View File
@@ -179,17 +179,27 @@ class Deployment:
diagnostics=diagnostics,
)
async def run(self, workflow_input: Mapping[str, Any]) -> Run:
async def run(
self,
workflow_input: Mapping[str, Any],
*,
max_steps: int | None = None,
) -> Run:
from .codec import decode_run_result
from .runs import _run_from_decoded
decoded = decode_run_result(
await self._port.run_deployment(
if max_steps is None:
raw = await self._port.run_deployment(
deployment_id=self.deployment_id,
workflow_input=dict(workflow_input),
),
operation="workflow.runs.start",
)
)
else:
raw = await self._port.run_deployment(
deployment_id=self.deployment_id,
workflow_input=dict(workflow_input),
max_steps=max_steps,
)
decoded = decode_run_result(raw, operation="workflow.runs.start")
if decoded.deployment_id != self.deployment_id:
raise InvalidResponse(
operation="workflow.runs.start",
src/wf_client/protocols.py
+1
View File
@@ -121,6 +121,7 @@ class WorkflowClientPort(Protocol):
deployment_id: str,
workflow_input: dict[str, Any],
trace_range: TraceRangeLike | None = None,
max_steps: int | None = None,
) -> RunResult: ...
async def list_runs(
src/wf_client/runs.py
+19 -1
View File
@@ -119,12 +119,21 @@ def _run_from_decoded(
interrupt=_interrupt(decoded.interrupt, operation=operation),
diagnostics=decoded.diagnostics,
trace_count=decoded.trace_count,
max_steps=decoded.max_steps,
steps_executed=decoded.steps_executed,
steps_remaining=decoded.steps_remaining,
)
@dataclass(frozen=True, slots=True, init=False)
class Run:
"""Immutable client snapshot of one durable deployment run."""
"""Immutable client snapshot of one durable deployment run.
``max_steps``, ``steps_executed``, and ``steps_remaining`` are the
server-effective budget values returned with every run response. The
server substitutes its default limit when creation omits one, so
refresh/resume reconstruction simply preserves whatever was returned.
"""
_port: WorkflowClientPort = field(repr=False, compare=False)
run_id: str
@@ -135,6 +144,9 @@ class Run:
_interrupt: InterruptRequest | None = field(repr=False)
_diagnostics: tuple[DependencyDiagnostic, ...] = field(repr=False)
trace_count: int
max_steps: int
steps_executed: int
steps_remaining: int
def __init__(
self,
@@ -148,6 +160,9 @@ class Run:
interrupt: InterruptRequest | None,
diagnostics: tuple[DependencyDiagnostic, ...],
trace_count: int,
max_steps: int,
steps_executed: int,
steps_remaining: int,
) -> None:
object.__setattr__(self, "_port", _port)
object.__setattr__(self, "run_id", run_id)
@@ -162,6 +177,9 @@ class Run:
tuple(item.model_copy(deep=True) for item in diagnostics),
)
object.__setattr__(self, "trace_count", trace_count)
object.__setattr__(self, "max_steps", max_steps)
object.__setattr__(self, "steps_executed", steps_executed)
object.__setattr__(self, "steps_remaining", steps_remaining)
@property
def output(self) -> dict[str, Any] | None:
src/wf_transport_rpc_http/client/runs.py
+9 -8
View File
@@ -36,17 +36,18 @@ class RpcRunClientMixin:
deployment_id: str,
workflow_input: dict[str, Any],
trace_range: TraceRangeLike | None = None,
max_steps: int | None = None,
) -> RunResult:
params: dict[str, Any] = {
"deployment_id": deployment_id,
"workflow_input": workflow_input,
"trace_range": _trace_range_payload(trace_range),
}
if max_steps is not None:
params["max_steps"] = max_steps
return cast(
RunResult,
await self._call(
"workflow.runs.start",
{
"deployment_id": deployment_id,
"workflow_input": workflow_input,
"trace_range": _trace_range_payload(trace_range),
},
),
await self._call("workflow.runs.start", params),
)
async def resume_run(
src/wf_transport_rpc_http/methods/runs.py
+1
View File
@@ -52,6 +52,7 @@ def register_methods(
if params.trace_range is not None
else None
),
max_steps=params.max_steps,
)
except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
raise_workflow_rpc_error(exc)
src/wf_transport_rpc_http/models.py
+8
View File
@@ -435,6 +435,14 @@ class StartRunParams(RpcParamsModel):
deployment_id: str = Field(min_length=1)
workflow_input: dict[str, Any] = Field(default_factory=dict)
trace_range: TraceRangeParams | None = None
max_steps: int | None = Field(
default=None,
ge=1,
description=(
"Optional run step budget. The server default applies when omitted; "
"resume never accepts a replacement."
),
)
class InspectRunParams(RpcParamsModel):
tests/wf_cli/test_app.py
+50
View File
@@ -3240,3 +3240,53 @@ def test_wf_draft_import_rejects_bad_files_before_loading_context(
assert result.exit_code == 2
assert expected_error in " ".join(result.output.split())
assert "context loaded" not in result.output
def test_wf_run_start_rejects_zero_max_steps_before_loading_context(
monkeypatch,
) -> None:
monkeypatch.setattr(
"wf_cli.commands.runs.load_cli_context_from_typer",
lambda _ctx: (_ for _ in ()).throw(AssertionError("context loaded")),
)
result = runner.invoke(
app,
["run", "start", "report.production", "--input", "{}", "--max-steps", "0"],
)
assert result.exit_code != 0
assert "context loaded" not in result.output
def test_wf_run_start_threads_max_steps_to_handler(monkeypatch) -> None:
calls: list[dict[str, Any]] = []
class FakeHandlers:
async def run_deployment(self, **kwargs: Any) -> dict[str, Any]:
calls.append(kwargs)
return {
"deployment_id": "report.production",
"status": "completed",
"run_id": "run-1",
}
context = SimpleNamespace(handlers=FakeHandlers(), verbose=False)
monkeypatch.setattr(
"wf_cli.commands.runs.load_cli_context_from_typer",
lambda _ctx: context,
)
budgeted = runner.invoke(
app,
["run", "start", "report.production", "--input", "{}", "--max-steps", "7"],
)
defaulted = runner.invoke(
app,
["run", "start", "report.production", "--input", "{}"],
)
assert budgeted.exit_code == 0, budgeted.output
assert defaulted.exit_code == 0, defaulted.output
assert calls[0]["max_steps"] == 7
assert calls[1]["max_steps"] is None
tests/wf_client/test_codec.py
+3
View File
@@ -124,6 +124,9 @@ def _run_payload(*, trace: list[dict[str, Any]] | None = None) -> dict[str, Any]
"error": None,
"output": {"result": "hello"},
"trace_count": 0 if trace is None else len(trace),
"max_steps": 10_000,
"steps_executed": 1,
"steps_remaining": 9_999,
"diagnostics": [],
"next_actions": {
"can_continue": False,
tests/wf_client/test_deployments.py
+28
View File
@@ -77,6 +77,9 @@ class _FakePort:
"error": None,
"output": {"result": "done"},
"trace_count": 0,
"max_steps": 10_000,
"steps_executed": 1,
"steps_remaining": 9_999,
"diagnostics": [],
"next_actions": {
"can_continue": False,
@@ -365,3 +368,28 @@ async def test_deployment_snapshot_defensively_copies_model_and_diagnostics() ->
assert deployment.diagnostics[0].message == "original"
await deployment.run({})
assert port.calls[-1][1]["deployment_id"] == "report.production"
@pytest.mark.asyncio
async def test_deployment_run_omits_max_steps_when_not_supplied() -> None:
artifact = _artifact()
deployment = await artifact.deploy("report.production")
port = cast(_FakePort, deployment._port)
await deployment.run({"topic": "workflow"})
run_calls = [params for name, params in port.calls if name == "run_deployment"]
assert len(run_calls) == 1
assert "max_steps" not in run_calls[0]
@pytest.mark.asyncio
async def test_deployment_run_threads_max_steps_when_supplied() -> None:
artifact = _artifact()
deployment = await artifact.deploy("report.production")
port = cast(_FakePort, deployment._port)
port.run_result.update({"max_steps": 7, "steps_executed": 2, "steps_remaining": 5})
run = await deployment.run({"topic": "workflow"}, max_steps=7)
run_calls = [params for name, params in port.calls if name == "run_deployment"]
assert run_calls[-1]["max_steps"] == 7
assert run.max_steps == 7
assert run.steps_executed == 2
assert run.steps_remaining == 5
tests/wf_client/test_repr.py
+3
View File
@@ -73,6 +73,9 @@ def test_rich_representations_bound_large_values_and_redact_secret_like_fields()
interrupt=None,
diagnostics=(),
trace_count=1000,
max_steps=10_000,
steps_executed=3,
steps_remaining=9_997,
)
rendered = repr(run)
tests/wf_client/test_runs.py
+67
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import inspect
from typing import Any, cast
import pytest
@@ -38,6 +39,9 @@ def _payload(
"error": None,
"output": None if status == "interrupted" else {"result": "done"},
"trace_count": 1,
"max_steps": 10_000,
"steps_executed": 1,
"steps_remaining": 9_999,
"diagnostics": [],
"next_actions": {
"can_continue": status == "interrupted",
@@ -253,3 +257,66 @@ async def test_run_snapshot_defensively_copies_nested_public_values() -> None:
assert run.diagnostics[0].message == "original"
await run.resume({"approved": True})
assert port.calls[-1][1]["run_id"] == "run-1"
def _budgeted_payload(
*, status: str = "interrupted", max_steps: int = 7
) -> dict[str, Any]:
payload = _payload(status=status)
payload.update(
{
"max_steps": max_steps,
"steps_executed": 2,
"steps_remaining": max_steps - 2,
}
)
return payload
@pytest.mark.asyncio
async def test_run_decoder_requires_budget_fields() -> None:
for field in ("max_steps", "steps_executed", "steps_remaining"):
payload = _budgeted_payload()
del payload[field]
with pytest.raises(InvalidResponse, match="workflow.runs.inspect"):
Run.from_payload(cast(WorkflowClientPort, _Port()), payload)
@pytest.mark.asyncio
async def test_run_exposes_step_budget() -> None:
run = Run.from_payload(
cast(WorkflowClientPort, _Port()), _budgeted_payload(max_steps=7)
)
assert run.max_steps == 7
assert run.steps_executed == 2
assert run.steps_remaining == 5
@pytest.mark.asyncio
async def test_refresh_preserves_step_budget() -> None:
port = _Port()
port.resume_payload = _budgeted_payload(status="completed", max_steps=9)
run = Run.from_payload(
cast(WorkflowClientPort, port), _budgeted_payload(max_steps=7)
)
refreshed = await run.refresh()
assert refreshed.max_steps == 9
assert refreshed.steps_executed == 2
assert refreshed.steps_remaining == 7
@pytest.mark.asyncio
async def test_resume_preserves_step_budget() -> None:
port = _Port()
port.resume_payload = _budgeted_payload(status="completed", max_steps=9)
run = Run.from_payload(
cast(WorkflowClientPort, port), _budgeted_payload(max_steps=7)
)
resumed = await run.resume({"approved": True})
assert resumed.max_steps == 9
assert resumed.steps_executed == 2
assert resumed.steps_remaining == 7
def test_resume_signature_excludes_max_steps() -> None:
assert "max_steps" not in inspect.signature(Run.resume).parameters
tests/wf_transport_rpc_http/test_app.py
+92
View File
@@ -2233,3 +2233,95 @@ async def test_rpc_diagnoses_source(tmp_path) -> None:
assert payload["result"]["source_id"] == "wf.std"
assert payload["result"]["status"] == "unknown"
async def _seed_step_budget_deployment(server: Any) -> str:
await server.api.create_artifact_from_plan(
artifact_id="budget_constant",
version=1,
title="Budget Constant",
plan=_constant_plan(),
outcomes=["ok"],
source_bindings={},
)
await server.api.save_deployment(
{
"id": "budget_constant.default",
"artifact_id": "budget_constant",
"artifact_version": 1,
"bindings": {},
}
)
return "budget_constant.default"
async def test_rpc_runs_start_applies_optional_step_budget(tmp_path) -> None:
server = build_local_static_workflow_server(tmp_path / "store")
deployment_id = await _seed_step_budget_deployment(server)
app = create_rpc_app(server)
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
default = await _rpc(
client,
"workflow.runs.start",
{"deployment_id": deployment_id, "workflow_input": {}},
)
budgeted = await _rpc(
client,
"workflow.runs.start",
{
"deployment_id": deployment_id,
"workflow_input": {},
"max_steps": 5,
},
)
assert default["result"]["max_steps"] == 10_000
assert default["result"]["steps_executed"] >= 1
assert (
default["result"]["steps_remaining"]
== 10_000 - default["result"]["steps_executed"]
)
assert budgeted["result"]["max_steps"] == 5
assert budgeted["result"]["steps_executed"] >= 1
assert (
budgeted["result"]["steps_remaining"]
== 5 - budgeted["result"]["steps_executed"]
)
async def test_rpc_runs_start_rejects_non_positive_step_budget(tmp_path) -> None:
server = build_local_static_workflow_server(tmp_path / "store")
deployment_id = await _seed_step_budget_deployment(server)
app = create_rpc_app(server)
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
rejected = await _rpc(
client,
"workflow.runs.start",
{
"deployment_id": deployment_id,
"workflow_input": {},
"max_steps": 0,
},
)
assert rejected["error"]["code"] == -32602
async def test_rpc_runs_resume_rejects_step_budget_replacement(tmp_path) -> None:
server = build_local_static_workflow_server(tmp_path / "store")
app = create_rpc_app(server)
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
rejected = await _rpc(
client,
"workflow.runs.resume",
{
"run_id": "missing-run",
"resume_payload": {},
"max_steps": 3,
},
)
assert rejected["error"]["code"] == -32602
tests/wf_transport_rpc_http/test_client.py
+48
View File
@@ -28,6 +28,7 @@ from wf_server import build_local_static_workflow_server
from wf_transport_rpc_http import RpcWorkflowApiClient, create_rpc_app
from wf_transport_rpc_http.client.base import RpcProtocolError
from wf_transport_rpc_http.client.drafts import RpcDraftClientMixin
from wf_transport_rpc_http.client.runs import RpcRunClientMixin
from wf_transport_rpc_http.client.sources import RpcSourceAdminClientMixin
@@ -1342,3 +1343,50 @@ async def test_rpc_client_diagnoses_source(tmp_path) -> None:
assert payload == {"source_id": "demo.personal", "status": "ok"}
assert calls == [("workflow.sources.diagnose", {"source_id": "demo.personal"})]
async def test_rpc_run_client_omits_max_steps_unless_supplied() -> None:
calls: list[tuple[str, dict[str, object]]] = []
class Client(RpcRunClientMixin):
async def _call(self, method: str, params: dict[str, object]):
calls.append((method, params))
return {"run_id": "run-1"}
client = Client()
await client.run_deployment(deployment_id="report.default", workflow_input={})
assert calls[-1] == (
"workflow.runs.start",
{
"deployment_id": "report.default",
"workflow_input": {},
"trace_range": None,
},
)
async def test_rpc_run_client_threads_max_steps_when_supplied() -> None:
calls: list[tuple[str, dict[str, object]]] = []
class Client(RpcRunClientMixin):
async def _call(self, method: str, params: dict[str, object]):
calls.append((method, params))
return {"run_id": "run-1"}
client = Client()
await client.run_deployment(
deployment_id="report.default",
workflow_input={},
max_steps=5,
)
assert calls[-1] == (
"workflow.runs.start",
{
"deployment_id": "report.default",
"workflow_input": {},
"trace_range": None,
"max_steps": 5,
},
)
web/packages/rpc/src/generated/workflow-contract.ts
+48 -1
View File
@@ -895,6 +895,10 @@ export interface WorkflowContractMap {
[k: string]: unknown;
};
trace_range?: TraceRangeParams | null;
/**
* Optional run step budget. The server default applies when omitted; resume never accepts a replacement.
*/
max_steps?: number | null;
};
result: RunResult;
};
@@ -2399,12 +2403,15 @@ export interface RunResult {
diagnostics: DependencyDiagnosticPayload[];
error: string | null;
interrupt: InterruptPayload | null;
max_steps: number;
next_actions: NextActionsPayload;
outcome: string | null;
output: JsonObject | null;
resume_readiness: ResumeReadiness | null;
run_id: string | null;
status: RunStatus;
steps_executed: number;
steps_remaining: number;
trace?: TraceEntryPayload[];
trace_count: number;
trace_limit?: number;
@@ -2532,12 +2539,15 @@ export interface RunTraceResult {
diagnostics: DependencyDiagnosticPayload[];
error: string | null;
interrupt: InterruptPayload | null;
max_steps: number;
next_actions: NextActionsPayload;
outcome: string | null;
output: JsonObject | null;
resume_readiness: ResumeReadiness | null;
run_id: string | null;
status: RunStatus;
steps_executed: number;
steps_remaining: number;
trace: TraceEntryPayload[];
trace_count: number;
trace_limit: number;
@@ -4473,6 +4483,9 @@ export const workflowRuntimeContract = {
}
]
},
"max_steps": {
"type": "integer"
},
"next_actions": {
"$ref": "#/components/schemas/NextActionsPayload"
},
@@ -4519,6 +4532,12 @@ export const workflowRuntimeContract = {
"status": {
"$ref": "#/components/schemas/RunStatus"
},
"steps_executed": {
"type": "integer"
},
"steps_remaining": {
"type": "integer"
},
"trace": {
"items": {
"$ref": "#/components/schemas/TraceEntryPayload"
@@ -4551,7 +4570,10 @@ export const workflowRuntimeContract = {
"outcome",
"error",
"output",
"trace_count"
"trace_count",
"max_steps",
"steps_executed",
"steps_remaining"
],
"type": "object"
},
@@ -4639,6 +4661,9 @@ export const workflowRuntimeContract = {
}
]
},
"max_steps": {
"type": "integer"
},
"next_actions": {
"$ref": "#/components/schemas/NextActionsPayload"
},
@@ -4685,6 +4710,12 @@ export const workflowRuntimeContract = {
"status": {
"$ref": "#/components/schemas/RunStatus"
},
"steps_executed": {
"type": "integer"
},
"steps_remaining": {
"type": "integer"
},
"trace": {
"items": {
"$ref": "#/components/schemas/TraceEntryPayload"
@@ -4718,6 +4749,9 @@ export const workflowRuntimeContract = {
"error",
"output",
"trace_count",
"max_steps",
"steps_executed",
"steps_remaining",
"trace",
"trace_start",
"trace_limit",
@@ -6571,6 +6605,19 @@ export const workflowRuntimeContract = {
}
],
"default": null
},
"max_steps": {
"anyOf": [
{
"minimum": 1,
"type": "integer"
},
{
"type": "null"
}
],
"default": null,
"description": "Optional run step budget. The server default applies when omitted; resume never accepts a replacement."
}
},
"required": [