feat: expose run step budgets to clients
This commit is contained in:
@@ -3478,6 +3478,9 @@
|
|||||||
}
|
}
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
|
"max_steps": {
|
||||||
|
"type": "integer"
|
||||||
|
},
|
||||||
"next_actions": {
|
"next_actions": {
|
||||||
"$ref": "#/components/schemas/NextActionsPayload"
|
"$ref": "#/components/schemas/NextActionsPayload"
|
||||||
},
|
},
|
||||||
@@ -3524,6 +3527,12 @@
|
|||||||
"status": {
|
"status": {
|
||||||
"$ref": "#/components/schemas/RunStatus"
|
"$ref": "#/components/schemas/RunStatus"
|
||||||
},
|
},
|
||||||
|
"steps_executed": {
|
||||||
|
"type": "integer"
|
||||||
|
},
|
||||||
|
"steps_remaining": {
|
||||||
|
"type": "integer"
|
||||||
|
},
|
||||||
"trace": {
|
"trace": {
|
||||||
"items": {
|
"items": {
|
||||||
"$ref": "#/components/schemas/TraceEntryPayload"
|
"$ref": "#/components/schemas/TraceEntryPayload"
|
||||||
@@ -3556,7 +3565,10 @@
|
|||||||
"outcome",
|
"outcome",
|
||||||
"error",
|
"error",
|
||||||
"output",
|
"output",
|
||||||
"trace_count"
|
"trace_count",
|
||||||
|
"max_steps",
|
||||||
|
"steps_executed",
|
||||||
|
"steps_remaining"
|
||||||
],
|
],
|
||||||
"type": "object"
|
"type": "object"
|
||||||
},
|
},
|
||||||
@@ -3644,6 +3656,9 @@
|
|||||||
}
|
}
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
|
"max_steps": {
|
||||||
|
"type": "integer"
|
||||||
|
},
|
||||||
"next_actions": {
|
"next_actions": {
|
||||||
"$ref": "#/components/schemas/NextActionsPayload"
|
"$ref": "#/components/schemas/NextActionsPayload"
|
||||||
},
|
},
|
||||||
@@ -3690,6 +3705,12 @@
|
|||||||
"status": {
|
"status": {
|
||||||
"$ref": "#/components/schemas/RunStatus"
|
"$ref": "#/components/schemas/RunStatus"
|
||||||
},
|
},
|
||||||
|
"steps_executed": {
|
||||||
|
"type": "integer"
|
||||||
|
},
|
||||||
|
"steps_remaining": {
|
||||||
|
"type": "integer"
|
||||||
|
},
|
||||||
"trace": {
|
"trace": {
|
||||||
"items": {
|
"items": {
|
||||||
"$ref": "#/components/schemas/TraceEntryPayload"
|
"$ref": "#/components/schemas/TraceEntryPayload"
|
||||||
@@ -3723,6 +3744,9 @@
|
|||||||
"error",
|
"error",
|
||||||
"output",
|
"output",
|
||||||
"trace_count",
|
"trace_count",
|
||||||
|
"max_steps",
|
||||||
|
"steps_executed",
|
||||||
|
"steps_remaining",
|
||||||
"trace",
|
"trace",
|
||||||
"trace_start",
|
"trace_start",
|
||||||
"trace_limit",
|
"trace_limit",
|
||||||
@@ -8760,6 +8784,23 @@
|
|||||||
],
|
],
|
||||||
"default": null
|
"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": {
|
"result": {
|
||||||
|
|||||||
@@ -70,6 +70,14 @@ def start_run(
|
|||||||
"--trace-limit", min=1, max=100, help="Optional trace slice limit."
|
"--trace-limit", min=1, max=100, help="Optional trace slice limit."
|
||||||
),
|
),
|
||||||
] = None,
|
] = 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:
|
) -> None:
|
||||||
"""Start one workflow deployment."""
|
"""Start one workflow deployment."""
|
||||||
try:
|
try:
|
||||||
@@ -84,6 +92,7 @@ def start_run(
|
|||||||
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,
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
emit_json(payload)
|
emit_json(payload)
|
||||||
|
|||||||
@@ -259,6 +259,7 @@ class PublicErrorWorkflowClientPort:
|
|||||||
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._invoke(
|
return await self._invoke(
|
||||||
"workflow.runs.start",
|
"workflow.runs.start",
|
||||||
@@ -266,6 +267,7 @@ class PublicErrorWorkflowClientPort:
|
|||||||
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 inspect_run(self, *, run_id: str) -> RunResult:
|
async def inspect_run(self, *, run_id: str) -> RunResult:
|
||||||
|
|||||||
+12
-1
@@ -51,6 +51,9 @@ class _DecodedRunFields:
|
|||||||
error: str | None
|
error: str | None
|
||||||
output: dict[str, Any] | None
|
output: dict[str, Any] | None
|
||||||
trace_count: int
|
trace_count: int
|
||||||
|
max_steps: int
|
||||||
|
steps_executed: int
|
||||||
|
steps_remaining: int
|
||||||
diagnostics: tuple[DependencyDiagnostic, ...]
|
diagnostics: tuple[DependencyDiagnostic, ...]
|
||||||
next_actions: dict[str, Any]
|
next_actions: dict[str, Any]
|
||||||
trace: tuple[dict[str, Any], ...] | None = None
|
trace: tuple[dict[str, Any], ...] | None = None
|
||||||
@@ -255,6 +258,9 @@ def _decode_run_fields(
|
|||||||
error=wire["error"],
|
error=wire["error"],
|
||||||
output=dict(wire["output"]) if wire["output"] is not None else None,
|
output=dict(wire["output"]) if wire["output"] is not None else None,
|
||||||
trace_count=wire["trace_count"],
|
trace_count=wire["trace_count"],
|
||||||
|
max_steps=wire["max_steps"],
|
||||||
|
steps_executed=wire["steps_executed"],
|
||||||
|
steps_remaining=wire["steps_remaining"],
|
||||||
diagnostics=diagnostics,
|
diagnostics=diagnostics,
|
||||||
next_actions=dict(wire["next_actions"]),
|
next_actions=dict(wire["next_actions"]),
|
||||||
trace=trace,
|
trace=trace,
|
||||||
@@ -269,7 +275,12 @@ def decode_run_result(
|
|||||||
*,
|
*,
|
||||||
operation: str = "workflow.runs.inspect",
|
operation: str = "workflow.runs.inspect",
|
||||||
) -> DecodedRunResult:
|
) -> 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)
|
wire = _validate(payload, RunResult, operation)
|
||||||
fields = _decode_run_fields(
|
fields = _decode_run_fields(
|
||||||
wire,
|
wire,
|
||||||
|
|||||||
@@ -179,17 +179,27 @@ class Deployment:
|
|||||||
diagnostics=diagnostics,
|
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 .codec import decode_run_result
|
||||||
from .runs import _run_from_decoded
|
from .runs import _run_from_decoded
|
||||||
|
|
||||||
decoded = decode_run_result(
|
if max_steps is None:
|
||||||
await self._port.run_deployment(
|
raw = await self._port.run_deployment(
|
||||||
deployment_id=self.deployment_id,
|
deployment_id=self.deployment_id,
|
||||||
workflow_input=dict(workflow_input),
|
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:
|
if decoded.deployment_id != self.deployment_id:
|
||||||
raise InvalidResponse(
|
raise InvalidResponse(
|
||||||
operation="workflow.runs.start",
|
operation="workflow.runs.start",
|
||||||
|
|||||||
@@ -121,6 +121,7 @@ class WorkflowClientPort(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 list_runs(
|
async def list_runs(
|
||||||
|
|||||||
+19
-1
@@ -119,12 +119,21 @@ def _run_from_decoded(
|
|||||||
interrupt=_interrupt(decoded.interrupt, operation=operation),
|
interrupt=_interrupt(decoded.interrupt, operation=operation),
|
||||||
diagnostics=decoded.diagnostics,
|
diagnostics=decoded.diagnostics,
|
||||||
trace_count=decoded.trace_count,
|
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)
|
@dataclass(frozen=True, slots=True, init=False)
|
||||||
class Run:
|
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)
|
_port: WorkflowClientPort = field(repr=False, compare=False)
|
||||||
run_id: str
|
run_id: str
|
||||||
@@ -135,6 +144,9 @@ class Run:
|
|||||||
_interrupt: InterruptRequest | None = field(repr=False)
|
_interrupt: InterruptRequest | None = field(repr=False)
|
||||||
_diagnostics: tuple[DependencyDiagnostic, ...] = field(repr=False)
|
_diagnostics: tuple[DependencyDiagnostic, ...] = field(repr=False)
|
||||||
trace_count: int
|
trace_count: int
|
||||||
|
max_steps: int
|
||||||
|
steps_executed: int
|
||||||
|
steps_remaining: int
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
@@ -148,6 +160,9 @@ class Run:
|
|||||||
interrupt: InterruptRequest | None,
|
interrupt: InterruptRequest | None,
|
||||||
diagnostics: tuple[DependencyDiagnostic, ...],
|
diagnostics: tuple[DependencyDiagnostic, ...],
|
||||||
trace_count: int,
|
trace_count: int,
|
||||||
|
max_steps: int,
|
||||||
|
steps_executed: int,
|
||||||
|
steps_remaining: int,
|
||||||
) -> None:
|
) -> None:
|
||||||
object.__setattr__(self, "_port", _port)
|
object.__setattr__(self, "_port", _port)
|
||||||
object.__setattr__(self, "run_id", run_id)
|
object.__setattr__(self, "run_id", run_id)
|
||||||
@@ -162,6 +177,9 @@ class Run:
|
|||||||
tuple(item.model_copy(deep=True) for item in diagnostics),
|
tuple(item.model_copy(deep=True) for item in diagnostics),
|
||||||
)
|
)
|
||||||
object.__setattr__(self, "trace_count", trace_count)
|
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
|
@property
|
||||||
def output(self) -> dict[str, Any] | None:
|
def output(self) -> dict[str, Any] | None:
|
||||||
|
|||||||
@@ -36,17 +36,18 @@ class RpcRunClientMixin:
|
|||||||
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:
|
||||||
|
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(
|
return cast(
|
||||||
RunResult,
|
RunResult,
|
||||||
await self._call(
|
await self._call("workflow.runs.start", params),
|
||||||
"workflow.runs.start",
|
|
||||||
{
|
|
||||||
"deployment_id": deployment_id,
|
|
||||||
"workflow_input": workflow_input,
|
|
||||||
"trace_range": _trace_range_payload(trace_range),
|
|
||||||
},
|
|
||||||
),
|
|
||||||
)
|
)
|
||||||
|
|
||||||
async def resume_run(
|
async def resume_run(
|
||||||
|
|||||||
@@ -52,6 +52,7 @@ def register_methods(
|
|||||||
if params.trace_range is not None
|
if params.trace_range is not None
|
||||||
else None
|
else None
|
||||||
),
|
),
|
||||||
|
max_steps=params.max_steps,
|
||||||
)
|
)
|
||||||
except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
|
except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
|
||||||
raise_workflow_rpc_error(exc)
|
raise_workflow_rpc_error(exc)
|
||||||
|
|||||||
@@ -435,6 +435,14 @@ class StartRunParams(RpcParamsModel):
|
|||||||
deployment_id: str = Field(min_length=1)
|
deployment_id: str = Field(min_length=1)
|
||||||
workflow_input: dict[str, Any] = Field(default_factory=dict)
|
workflow_input: dict[str, Any] = Field(default_factory=dict)
|
||||||
trace_range: TraceRangeParams | None = None
|
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):
|
class InspectRunParams(RpcParamsModel):
|
||||||
|
|||||||
@@ -3240,3 +3240,53 @@ def test_wf_draft_import_rejects_bad_files_before_loading_context(
|
|||||||
assert result.exit_code == 2
|
assert result.exit_code == 2
|
||||||
assert expected_error in " ".join(result.output.split())
|
assert expected_error in " ".join(result.output.split())
|
||||||
assert "context loaded" not in result.output
|
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
|
||||||
|
|||||||
@@ -124,6 +124,9 @@ def _run_payload(*, trace: list[dict[str, Any]] | None = None) -> dict[str, Any]
|
|||||||
"error": None,
|
"error": None,
|
||||||
"output": {"result": "hello"},
|
"output": {"result": "hello"},
|
||||||
"trace_count": 0 if trace is None else len(trace),
|
"trace_count": 0 if trace is None else len(trace),
|
||||||
|
"max_steps": 10_000,
|
||||||
|
"steps_executed": 1,
|
||||||
|
"steps_remaining": 9_999,
|
||||||
"diagnostics": [],
|
"diagnostics": [],
|
||||||
"next_actions": {
|
"next_actions": {
|
||||||
"can_continue": False,
|
"can_continue": False,
|
||||||
|
|||||||
@@ -77,6 +77,9 @@ class _FakePort:
|
|||||||
"error": None,
|
"error": None,
|
||||||
"output": {"result": "done"},
|
"output": {"result": "done"},
|
||||||
"trace_count": 0,
|
"trace_count": 0,
|
||||||
|
"max_steps": 10_000,
|
||||||
|
"steps_executed": 1,
|
||||||
|
"steps_remaining": 9_999,
|
||||||
"diagnostics": [],
|
"diagnostics": [],
|
||||||
"next_actions": {
|
"next_actions": {
|
||||||
"can_continue": False,
|
"can_continue": False,
|
||||||
@@ -365,3 +368,28 @@ async def test_deployment_snapshot_defensively_copies_model_and_diagnostics() ->
|
|||||||
assert deployment.diagnostics[0].message == "original"
|
assert deployment.diagnostics[0].message == "original"
|
||||||
await deployment.run({})
|
await deployment.run({})
|
||||||
assert port.calls[-1][1]["deployment_id"] == "report.production"
|
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
|
||||||
|
|||||||
@@ -73,6 +73,9 @@ def test_rich_representations_bound_large_values_and_redact_secret_like_fields()
|
|||||||
interrupt=None,
|
interrupt=None,
|
||||||
diagnostics=(),
|
diagnostics=(),
|
||||||
trace_count=1000,
|
trace_count=1000,
|
||||||
|
max_steps=10_000,
|
||||||
|
steps_executed=3,
|
||||||
|
steps_remaining=9_997,
|
||||||
)
|
)
|
||||||
|
|
||||||
rendered = repr(run)
|
rendered = repr(run)
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import inspect
|
||||||
from typing import Any, cast
|
from typing import Any, cast
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
@@ -38,6 +39,9 @@ def _payload(
|
|||||||
"error": None,
|
"error": None,
|
||||||
"output": None if status == "interrupted" else {"result": "done"},
|
"output": None if status == "interrupted" else {"result": "done"},
|
||||||
"trace_count": 1,
|
"trace_count": 1,
|
||||||
|
"max_steps": 10_000,
|
||||||
|
"steps_executed": 1,
|
||||||
|
"steps_remaining": 9_999,
|
||||||
"diagnostics": [],
|
"diagnostics": [],
|
||||||
"next_actions": {
|
"next_actions": {
|
||||||
"can_continue": status == "interrupted",
|
"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"
|
assert run.diagnostics[0].message == "original"
|
||||||
await run.resume({"approved": True})
|
await run.resume({"approved": True})
|
||||||
assert port.calls[-1][1]["run_id"] == "run-1"
|
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
|
||||||
|
|||||||
@@ -2233,3 +2233,95 @@ async def test_rpc_diagnoses_source(tmp_path) -> None:
|
|||||||
|
|
||||||
assert payload["result"]["source_id"] == "wf.std"
|
assert payload["result"]["source_id"] == "wf.std"
|
||||||
assert payload["result"]["status"] == "unknown"
|
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
|
||||||
|
|||||||
@@ -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 import RpcWorkflowApiClient, create_rpc_app
|
||||||
from wf_transport_rpc_http.client.base import RpcProtocolError
|
from wf_transport_rpc_http.client.base import RpcProtocolError
|
||||||
from wf_transport_rpc_http.client.drafts import RpcDraftClientMixin
|
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
|
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 payload == {"source_id": "demo.personal", "status": "ok"}
|
||||||
assert calls == [("workflow.sources.diagnose", {"source_id": "demo.personal"})]
|
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,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|||||||
@@ -895,6 +895,10 @@ export interface WorkflowContractMap {
|
|||||||
[k: string]: unknown;
|
[k: string]: unknown;
|
||||||
};
|
};
|
||||||
trace_range?: TraceRangeParams | null;
|
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;
|
result: RunResult;
|
||||||
};
|
};
|
||||||
@@ -2399,12 +2403,15 @@ export interface RunResult {
|
|||||||
diagnostics: DependencyDiagnosticPayload[];
|
diagnostics: DependencyDiagnosticPayload[];
|
||||||
error: string | null;
|
error: string | null;
|
||||||
interrupt: InterruptPayload | null;
|
interrupt: InterruptPayload | null;
|
||||||
|
max_steps: number;
|
||||||
next_actions: NextActionsPayload;
|
next_actions: NextActionsPayload;
|
||||||
outcome: string | null;
|
outcome: string | null;
|
||||||
output: JsonObject | null;
|
output: JsonObject | null;
|
||||||
resume_readiness: ResumeReadiness | null;
|
resume_readiness: ResumeReadiness | null;
|
||||||
run_id: string | null;
|
run_id: string | null;
|
||||||
status: RunStatus;
|
status: RunStatus;
|
||||||
|
steps_executed: number;
|
||||||
|
steps_remaining: number;
|
||||||
trace?: TraceEntryPayload[];
|
trace?: TraceEntryPayload[];
|
||||||
trace_count: number;
|
trace_count: number;
|
||||||
trace_limit?: number;
|
trace_limit?: number;
|
||||||
@@ -2532,12 +2539,15 @@ export interface RunTraceResult {
|
|||||||
diagnostics: DependencyDiagnosticPayload[];
|
diagnostics: DependencyDiagnosticPayload[];
|
||||||
error: string | null;
|
error: string | null;
|
||||||
interrupt: InterruptPayload | null;
|
interrupt: InterruptPayload | null;
|
||||||
|
max_steps: number;
|
||||||
next_actions: NextActionsPayload;
|
next_actions: NextActionsPayload;
|
||||||
outcome: string | null;
|
outcome: string | null;
|
||||||
output: JsonObject | null;
|
output: JsonObject | null;
|
||||||
resume_readiness: ResumeReadiness | null;
|
resume_readiness: ResumeReadiness | null;
|
||||||
run_id: string | null;
|
run_id: string | null;
|
||||||
status: RunStatus;
|
status: RunStatus;
|
||||||
|
steps_executed: number;
|
||||||
|
steps_remaining: number;
|
||||||
trace: TraceEntryPayload[];
|
trace: TraceEntryPayload[];
|
||||||
trace_count: number;
|
trace_count: number;
|
||||||
trace_limit: number;
|
trace_limit: number;
|
||||||
@@ -4473,6 +4483,9 @@ export const workflowRuntimeContract = {
|
|||||||
}
|
}
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
|
"max_steps": {
|
||||||
|
"type": "integer"
|
||||||
|
},
|
||||||
"next_actions": {
|
"next_actions": {
|
||||||
"$ref": "#/components/schemas/NextActionsPayload"
|
"$ref": "#/components/schemas/NextActionsPayload"
|
||||||
},
|
},
|
||||||
@@ -4519,6 +4532,12 @@ export const workflowRuntimeContract = {
|
|||||||
"status": {
|
"status": {
|
||||||
"$ref": "#/components/schemas/RunStatus"
|
"$ref": "#/components/schemas/RunStatus"
|
||||||
},
|
},
|
||||||
|
"steps_executed": {
|
||||||
|
"type": "integer"
|
||||||
|
},
|
||||||
|
"steps_remaining": {
|
||||||
|
"type": "integer"
|
||||||
|
},
|
||||||
"trace": {
|
"trace": {
|
||||||
"items": {
|
"items": {
|
||||||
"$ref": "#/components/schemas/TraceEntryPayload"
|
"$ref": "#/components/schemas/TraceEntryPayload"
|
||||||
@@ -4551,7 +4570,10 @@ export const workflowRuntimeContract = {
|
|||||||
"outcome",
|
"outcome",
|
||||||
"error",
|
"error",
|
||||||
"output",
|
"output",
|
||||||
"trace_count"
|
"trace_count",
|
||||||
|
"max_steps",
|
||||||
|
"steps_executed",
|
||||||
|
"steps_remaining"
|
||||||
],
|
],
|
||||||
"type": "object"
|
"type": "object"
|
||||||
},
|
},
|
||||||
@@ -4639,6 +4661,9 @@ export const workflowRuntimeContract = {
|
|||||||
}
|
}
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
|
"max_steps": {
|
||||||
|
"type": "integer"
|
||||||
|
},
|
||||||
"next_actions": {
|
"next_actions": {
|
||||||
"$ref": "#/components/schemas/NextActionsPayload"
|
"$ref": "#/components/schemas/NextActionsPayload"
|
||||||
},
|
},
|
||||||
@@ -4685,6 +4710,12 @@ export const workflowRuntimeContract = {
|
|||||||
"status": {
|
"status": {
|
||||||
"$ref": "#/components/schemas/RunStatus"
|
"$ref": "#/components/schemas/RunStatus"
|
||||||
},
|
},
|
||||||
|
"steps_executed": {
|
||||||
|
"type": "integer"
|
||||||
|
},
|
||||||
|
"steps_remaining": {
|
||||||
|
"type": "integer"
|
||||||
|
},
|
||||||
"trace": {
|
"trace": {
|
||||||
"items": {
|
"items": {
|
||||||
"$ref": "#/components/schemas/TraceEntryPayload"
|
"$ref": "#/components/schemas/TraceEntryPayload"
|
||||||
@@ -4718,6 +4749,9 @@ export const workflowRuntimeContract = {
|
|||||||
"error",
|
"error",
|
||||||
"output",
|
"output",
|
||||||
"trace_count",
|
"trace_count",
|
||||||
|
"max_steps",
|
||||||
|
"steps_executed",
|
||||||
|
"steps_remaining",
|
||||||
"trace",
|
"trace",
|
||||||
"trace_start",
|
"trace_start",
|
||||||
"trace_limit",
|
"trace_limit",
|
||||||
@@ -6571,6 +6605,19 @@ export const workflowRuntimeContract = {
|
|||||||
}
|
}
|
||||||
],
|
],
|
||||||
"default": null
|
"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": [
|
"required": [
|
||||||
|
|||||||
Reference in New Issue
Block a user