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