feat: expose generic draft steps over rpc

This commit is contained in:
lda
2026-07-20 16:09:26 +07:00 Verified
parent dda013dd12
commit 1b5bbd2944
5 changed files with 368 additions and 1 deletions
@@ -3,6 +3,9 @@ from __future__ import annotations
from collections.abc import Sequence
from typing import Any, Literal
from wf_api.surface import RouteSource
from wf_artifacts.drafts.models import DraftStep
from .base import RpcCaller
@@ -207,6 +210,32 @@ class RpcDraftClientMixin:
},
)
async def add_step(
self: RpcCaller,
*,
workspace_id: str,
revision: int,
step_id: str,
step: DraftStep,
incoming: RouteSource | None = None,
routes: dict[str, str] | None = None,
) -> dict[str, Any]:
return await self._call(
"workflow.draft_workspaces.add_step",
{
"workspace_id": workspace_id,
"revision": revision,
"step_id": step_id,
"step": step.model_dump(mode="json", by_alias=True),
"incoming": (
None
if incoming is None
else {"step_id": incoming.step_id, "outcome": incoming.outcome}
),
"routes": routes,
},
)
async def branch_draft(
self: RpcCaller,
*,
@@ -4,10 +4,12 @@ from typing import Any
import fastapi_jsonrpc as jsonrpc
from wf_api.surface import RouteSource
from wf_server import WorkflowServer
from ..errors import WorkflowRpcError, raise_workflow_rpc_error
from ..models import (
AddDraftStepParams,
AddStepFromCapabilityParams,
BindDraftParams,
BranchDraftParams,
@@ -243,6 +245,33 @@ def register_methods(
except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
raise_workflow_rpc_error(exc)
@entrypoint.method(
name="workflow.draft_workspaces.add_step",
errors=[WorkflowRpcError],
)
async def workflow_draft_workspaces_add_step(
params: AddDraftStepParams = RpcParams(),
) -> dict[str, Any]:
try:
incoming = (
None
if params.incoming is None
else RouteSource(
step_id=params.incoming.step_id,
outcome=params.incoming.outcome,
)
)
return await server.api.add_step(
workspace_id=params.workspace_id,
revision=params.revision,
step_id=params.step_id,
step=params.step,
incoming=incoming,
routes=params.routes,
)
except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
raise_workflow_rpc_error(exc)
@entrypoint.method(
name="workflow.draft_workspaces.branch", errors=[WorkflowRpcError]
)
+15
View File
@@ -5,6 +5,7 @@ from typing import Any, Literal
from pydantic import BaseModel, ConfigDict, Field
from wf_api.models import TraceRange
from wf_artifacts.drafts.models import DraftStep
class RpcParamsModel(BaseModel):
@@ -100,6 +101,20 @@ class ListDraftWorkspacesParams(RpcParamsModel):
pass
class RouteSourceParams(RpcParamsModel):
step_id: str = Field(min_length=1)
outcome: str = Field(default="ok", min_length=1)
class AddDraftStepParams(RpcParamsModel):
workspace_id: str = Field(min_length=1)
revision: int = Field(ge=1)
step_id: str = Field(min_length=1)
step: DraftStep
incoming: RouteSourceParams | None = None
routes: dict[str, str] | None = None
class GetDraftWorkspaceParams(RpcParamsModel):
workspace_id: str = Field(min_length=1)
include_draft: bool = False
+177
View File
@@ -3,6 +3,8 @@ from __future__ import annotations
from typing import Any
import httpx
import pytest
from pydantic import ValidationError
from wf_api.models import RawWorkflowPlan
from wf_config import WorkflowConfigFile
@@ -10,6 +12,7 @@ from wf_core import END
from wf_server import build_local_static_workflow_server
from wf_server.config import build_workflow_server_from_workflow_config
from wf_transport_rpc_http.app import create_rpc_app
from wf_transport_rpc_http.models import AddDraftStepParams
async def _rpc(
@@ -886,6 +889,180 @@ async def test_rpc_draft_workspace_add_step_from_capability(tmp_path) -> None:
assert result["status"] == "valid"
def test_add_draft_step_params_preserve_typed_step_json() -> None:
foreach = AddDraftStepParams.model_validate(
{
"workspace_id": "ws",
"revision": 1,
"step_id": "each",
"step": {"foreach": {"over": "state.items", "as": "item"}},
"incoming": {"step_id": "call"},
}
)
foreach_dump = foreach.model_dump(mode="json", by_alias=True)
assert foreach_dump["step"]["foreach"]["as"] == "item"
assert "as_" not in foreach_dump["step"]["foreach"]
assert foreach_dump["incoming"] == {"step_id": "call", "outcome": "ok"}
when = AddDraftStepParams.model_validate(
{
"workspace_id": "ws",
"revision": 1,
"step_id": "decide",
"step": {
"when": {
"if": {"op": "exists", "path": "state.ready"},
"then": "next",
}
},
}
)
assert when.model_dump(mode="json", by_alias=True)["step"]["when"]["if"] == {
"op": "exists",
"path": "state.ready",
}
interrupt = AddDraftStepParams.model_validate(
{
"workspace_id": "ws",
"revision": 1,
"step_id": "review",
"step": {
"interrupt": {
"kind": "issue_review",
"request_schema": {
"type": "object",
"properties": {"issues": {"type": "array"}},
},
"resume_schema": {
"type": "object",
"properties": {"selected": {"type": "array"}},
},
},
},
}
)
interrupt_dump = interrupt.model_dump(mode="json", by_alias=True)
assert interrupt_dump["step"]["interrupt"]["request_schema"]["type"] == "object"
assert interrupt_dump["step"]["interrupt"]["resume_schema"]["type"] == "object"
subgraph = AddDraftStepParams.model_validate(
{
"workspace_id": "ws",
"revision": 1,
"step_id": "child",
"step": {
"subgraph": {"workflow": {"artifact_id": "child", "version": 2}}
},
}
)
assert subgraph.model_dump(mode="json", by_alias=True)["step"]["subgraph"][
"workflow"
] == {"artifact_id": "child", "version": 2}
def test_add_draft_step_params_reject_invalid_kind_and_route_source() -> None:
with pytest.raises(ValidationError):
AddDraftStepParams.model_validate(
{
"workspace_id": "ws",
"revision": 1,
"step_id": "bad",
"step": {"unknown": {}},
}
)
with pytest.raises(ValidationError):
AddDraftStepParams.model_validate(
{
"workspace_id": "ws",
"revision": 1,
"step_id": "bad",
"step": {"use": "demo.echo", "join": {}},
}
)
for incoming in ({"step_id": "", "outcome": "ok"}, {"step_id": "call", "outcome": ""}):
with pytest.raises(ValidationError):
AddDraftStepParams.model_validate(
{
"workspace_id": "ws",
"revision": 1,
"step_id": "new",
"step": {"join": {}},
"incoming": incoming,
}
)
async def test_rpc_draft_workspace_add_typed_step_round_trip(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:
created = await _rpc(
client,
"workflow.draft_workspaces.create_from_capability",
{
"workspace_id": "typed_step_ws",
"capability_name": "wf.std.constant",
"name": "typed_step",
},
)
added = await _rpc(
client,
"workflow.draft_workspaces.add_step",
{
"workspace_id": "typed_step_ws",
"revision": created["result"]["revision"],
"step_id": "review",
"step": {
"interrupt": {
"kind": "issue_review",
"request_schema": {
"type": "object",
"properties": {"issues": {"type": "array"}},
},
"resume_schema": {
"type": "object",
"properties": {"selected": {"type": "array"}},
},
"outcomes": ["submitted", "cancelled"],
}
},
"incoming": {"step_id": "call", "outcome": "ok"},
"routes": {"submitted": "__end__", "cancelled": "__end__"},
},
)
malformed = await _rpc(
client,
"workflow.draft_workspaces.add_step",
{
"workspace_id": "typed_step_ws",
"revision": added["result"]["revision"],
"step_id": "bad",
"step": {"unknown": {}},
},
)
fetched = await _rpc(
client,
"workflow.draft_workspaces.get",
{"workspace_id": "typed_step_ws", "include_draft": True},
)
assert added["result"]["revision"] == created["result"]["revision"] + 1
assert added["result"]["status"] in {"valid", "invalid"}
assert "error" in malformed
assert fetched["result"]["revision"] == added["result"]["revision"]
assert "bad" not in fetched["result"]["draft"]["steps"]
review = fetched["result"]["draft"]["steps"]["review"]["interrupt"]
assert review["request_schema"]["type"] == "object"
assert review["resume_schema"]["properties"]["selected"] == {
"type": "array"
}
async def test_rpc_diagnoses_source(tmp_path) -> None:
server = build_local_static_workflow_server(tmp_path / "store")
app = create_rpc_app(server)
+118 -1
View File
@@ -3,9 +3,12 @@ from __future__ import annotations
from typing import Any
import httpx
import pytest
from pydantic import TypeAdapter
from wf_api.models import RawWorkflowPlan, TraceRange
from wf_api.surface import WorkflowDraftSurface
from wf_api.surface import RouteSource, WorkflowDraftSurface
from wf_artifacts.drafts.models import DraftStep
from wf_core import END
from wf_server import build_local_static_workflow_server
from wf_transport_rpc_http import RpcWorkflowApiClient, create_rpc_app
@@ -625,6 +628,120 @@ async def test_rpc_client_draft_workspace_add_step_from_capability(tmp_path) ->
assert result["status"] == "valid"
@pytest.mark.parametrize(
("step_id", "step"),
[
("use", TypeAdapter(DraftStep).validate_python({"use": "demo.echo"})),
(
"foreach",
TypeAdapter(DraftStep).validate_python(
{"foreach": {"over": "state.items", "as": "item"}}
),
),
(
"interrupt",
TypeAdapter(DraftStep).validate_python(
{
"interrupt": {
"kind": "approval",
"request_schema": {"type": "object"},
"resume_schema": {"type": "object"},
}
}
),
),
("join", TypeAdapter(DraftStep).validate_python({"join": {}})),
("end", TypeAdapter(DraftStep).validate_python({"end": {}})),
(
"when",
TypeAdapter(DraftStep).validate_python(
{
"when": {
"if": {"op": "exists", "path": "state.ready"},
"then": "next",
}
}
),
),
(
"choose",
TypeAdapter(DraftStep).validate_python(
{
"choose": {
"clauses": [
{"if": {"op": "exists", "path": "state.ready"}, "then": "next"}
]
}
}
),
),
(
"match",
TypeAdapter(DraftStep).validate_python(
{
"match": {
"value": "state.status",
"cases": [{"equals": "ready", "then": "next"}],
}
}
),
),
(
"subgraph",
TypeAdapter(DraftStep).validate_python(
{
"subgraph": {"workflow": {"artifact_id": "child", "version": 2}}
}
),
),
],
)
async def test_rpc_client_add_step_preserves_all_typed_variants(
step_id: str, step: DraftStep
) -> None:
calls: list[dict[str, Any]] = []
class Client(RpcDraftClientMixin):
async def _call(self, method: str, params: dict[str, object]):
calls.append({"method": method, "params": params})
return {"revision": 2}
result = await Client().add_step(
workspace_id="ws",
revision=1,
step_id=step_id,
step=step,
incoming=RouteSource(step_id="lookup"),
routes=None,
)
assert result == {"revision": 2}
request = calls[0]
assert request["method"] == "workflow.draft_workspaces.add_step"
assert request["params"]["step"] == step.model_dump(mode="json", by_alias=True)
assert request["params"]["incoming"] == {
"step_id": "lookup",
"outcome": "ok",
}
if step_id == "when":
assert request["params"]["step"]["when"]["if"]["op"] == "exists"
if step_id == "foreach":
assert request["params"]["step"]["foreach"]["as"] == "item"
assert "as_" not in request["params"]["step"]["foreach"]
if step_id == "interrupt":
assert request["params"]["step"]["interrupt"]["request_schema"][
"type"
] == "object"
assert request["params"]["step"]["interrupt"]["resume_schema"][
"type"
] == "object"
if step_id == "subgraph":
assert request["params"]["step"]["subgraph"]["workflow"] == {
"artifact_id": "child",
"version": 2,
}
async def test_rpc_client_diagnoses_source(tmp_path) -> None:
calls: list[tuple[str, dict[str, object]]] = []