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

feat: add atomic generic draft step insertion

Browse Source
This commit is contained in:
lda
2026-07-20 08:49:21 +07:00 Verified
parent 7e8ce5330f
commit 3119276014
5 changed files with 460 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/wf_api/__init__.py
+2
View File
@@ -17,6 +17,7 @@ from .constants import (
RUNTIME_ERROR_CAPABILITY,
)
from .deployments import WorkflowDeploymentApi
from .draft_authoring import RouteSource
from .drafts import WorkflowDraftApi
from .durable_context import durable_workflow_api, require_workflow_stores
from .listing import matches_query, paged_list_payload
@@ -88,6 +89,7 @@ __all__ = [
"OutcomeCandidateKind",
"RUNTIME_ERROR_CAPABILITY",
"RawWorkflowPlan",
"RouteSource",
"RuntimeDependencies",
"TraceRange",
"WorkflowAdminApi",
src/wf_api/draft_authoring.py
+119 -4
View File
@@ -8,6 +8,18 @@ from wf_artifacts.draft_workspaces.models import (
WorkflowDraftWorkspace,
summarize_draft_workspace,
)
from wf_artifacts.drafts.models import (
DraftChooseStep,
DraftEndStep,
DraftForeachStep,
DraftInterruptStep,
DraftJoinStep,
DraftMatchStep,
DraftStep,
DraftSubgraphStep,
DraftUseStep,
DraftWhenStep,
)
from wf_core.models.steps import (
InputBinding,
OutputBinding,
@@ -120,6 +132,109 @@ class WorkflowDraftAuthoringApi:
outcomes = getattr(spec, "outcomes", None)
return tuple(outcomes) if outcomes is not None else None
def _draft_step_route_outcomes(self, step: DraftStep) -> set[str] | None:
"""Return top-level route outcomes, or ``None`` for non-routable steps."""
if isinstance(step, DraftUseStep):
return set(self._outcomes_for_capability(step.use) or (DEFAULT_OK_OUTCOME,))
if isinstance(step, DraftForeachStep):
outcomes = {"loop", "done"}
if step.foreach.item_error.action in {"skip", "collect"}:
outcomes.add("completed_with_errors")
return outcomes
if isinstance(step, DraftInterruptStep):
return set(step.interrupt.outcomes)
if isinstance(step, DraftJoinStep):
return {"done"}
if isinstance(step, DraftSubgraphStep):
return set(step.subgraph.outcomes)
if isinstance(
step, (DraftEndStep, DraftWhenStep, DraftChooseStep, DraftMatchStep)
):
return None
raise TypeError(f"unsupported draft step {type(step)!r}")
async def add_step(
self,
*,
workspace_id: str,
revision: int,
step_id: str,
step: DraftStep,
incoming: RouteSource | None = None,
routes: dict[str, str] | None = None,
) -> dict[str, Any]:
"""Add one typed draft step and optional route edits in one revision."""
workspace = self.drafts._draft_store().get_workspace(workspace_id)
steps = workspace.draft.get("steps")
if not isinstance(steps, dict):
raise ValueError("draft steps must be an object")
draft_routes = workspace.draft.get("routes")
if not isinstance(draft_routes, dict):
raise ValueError("draft routes must be an object")
if step_id in steps:
raise ValueError(f"draft step {step_id!r} already exists")
route_outcomes = self._draft_step_route_outcomes(step)
if routes is not None:
if route_outcomes is None:
raise ValueError(f"routes are not allowed for draft step {step_id!r}")
unknown_outcomes = set(routes) - route_outcomes
if unknown_outcomes:
raise ValueError(
f"unknown route outcome(s) for draft step {step_id!r}: "
f"{sorted(unknown_outcomes)!r}"
)
if incoming is not None and incoming.step_id not in steps:
raise ValueError(f"unknown incoming source step {incoming.step_id!r}")
patch: list[dict[str, Any]] = [
{
"op": "add",
"path": f"/steps/{escape_json_pointer(step_id)}",
"value": step.model_dump(mode="json", by_alias=True),
}
]
if routes is not None:
patch.append(
{
"op": "add",
"path": f"/routes/{escape_json_pointer(step_id)}",
"value": routes,
}
)
if incoming is not None:
source_routes = draft_routes.get(incoming.step_id)
if source_routes is None:
# JSON Patch cannot add a nested outcome until its parent exists.
patch.append(
{
"op": "add",
"path": f"/routes/{escape_json_pointer(incoming.step_id)}",
"value": {incoming.outcome: step_id},
}
)
else:
if not isinstance(source_routes, dict):
raise ValueError(
f"routes for step {incoming.step_id!r} must be an object"
)
patch.append(
{
"op": "add",
"path": (
f"/routes/{escape_json_pointer(incoming.step_id)}/"
f"{escape_json_pointer(incoming.outcome)}"
),
"value": step_id,
}
)
return await self.drafts.patch_draft_workspace(
workspace_id=workspace_id,
revision=revision,
patch=patch,
)
async def create_minimal_draft_workspace(
self,
*,
@@ -600,7 +715,7 @@ class WorkflowDraftAuthoringApi:
*,
workspace_id: str,
revision: int,
branches: Sequence[DraftOutcomeRef],
branches: Sequence[RouteSource],
target: str,
) -> dict[str, Any]:
"""Update the target for multiple (step, outcome) pairs atomically."""
@@ -801,8 +916,8 @@ class WorkflowDraftAuthoringApi:
@dataclass(frozen=True)
class DraftOutcomeRef:
"""A reference to a specific outcome of a draft step."""
class RouteSource:
"""One source step/outcome pair used for atomic route edits."""
step_id: str
outcome: str
outcome: str = DEFAULT_OK_OUTCOME
src/wf_api/service.py
+22 -3
View File
@@ -4,11 +4,12 @@ from collections.abc import Sequence
from typing import Any
from wf_artifacts import ArtifactKind
from wf_artifacts.drafts.models import DraftStep
from .artifacts import WorkflowArtifactApi
from .capabilities import WorkflowCapabilityApi
from .deployments import WorkflowDeploymentApi
from .draft_authoring import DraftOutcomeRef, WorkflowDraftAuthoringApi
from .draft_authoring import RouteSource, WorkflowDraftAuthoringApi
from .drafts import WorkflowDraftApi
from .models import RawWorkflowPlan
from .operation_context import WorkflowOperationContext
@@ -428,6 +429,25 @@ class WorkflowApi:
bind_outputs=bind_outputs,
)
async def add_step(
self,
*,
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.draft_authoring.add_step(
workspace_id=workspace_id,
revision=revision,
step_id=step_id,
step=step,
incoming=incoming,
routes=routes,
)
async def branch_draft(
self,
*,
@@ -452,8 +472,7 @@ class WorkflowApi:
target: str,
) -> dict[str, Any]:
refs = [
DraftOutcomeRef(step_id=b["step_id"], outcome=b["outcome"])
for b in branches
RouteSource(step_id=b["step_id"], outcome=b["outcome"]) for b in branches
]
return await self.draft_authoring.handle_draft(
workspace_id=workspace_id,
src/wf_api/surface.py
+13
View File
@@ -4,7 +4,9 @@ from collections.abc import Mapping, Sequence
from typing import Any, Protocol
from wf_artifacts import ArtifactKind
from wf_artifacts.drafts.models import DraftStep
from .draft_authoring import RouteSource
from .runs import TraceRangeLike
@@ -147,6 +149,17 @@ class WorkflowDraftSurface(Protocol):
bind_outputs: dict[str, str] | None = None,
) -> dict[str, Any]: ...
async def add_step(
self,
*,
workspace_id: str,
revision: int,
step_id: str,
step: DraftStep,
incoming: RouteSource | None = None,
routes: dict[str, str] | None = None,
) -> dict[str, Any]: ...
async def branch_draft(
self,
*,
tests/wf_api/test_drafts_service.py
+304 -4
View File
@@ -4,13 +4,14 @@ from pathlib import Path
from typing import Any
import pytest
from pydantic import BaseModel
from pydantic import BaseModel, TypeAdapter
from tests.wf_mcp.test_support import echo_tool
from wf_api.draft_authoring import DraftOutcomeRef, WorkflowDraftAuthoringApi
from wf_api.draft_authoring import RouteSource, WorkflowDraftAuthoringApi
from wf_api.drafts import WorkflowDraftApi
from wf_api.service import WorkflowApi
from wf_artifacts import FileDraftWorkspaceStore, FileWorkflowArtifactStore
from wf_artifacts.drafts.models import DraftStep
from wf_authoring import node
from wf_mcp.broker import WfMcpService
from wf_mcp.broker.service.workflow_operation_context import context_from_service
@@ -811,6 +812,305 @@ async def test_add_step_from_capability_rejects_existing_step_id(
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("step_name", "step_payload"),
[
(
"use",
{
"use": "demo.personal.echo_tool",
"input": [],
"output": [],
},
),
(
"foreach",
{"foreach": {"over": "state.items", "as": "item"}},
),
(
"interrupt",
{"interrupt": {"kind": "review", "outcomes": ["submitted"]}},
),
("join", {"join": {}}),
("end", {"end": {"outcome": "ok"}}),
(
"when",
{
"when": {
"if": {"op": "exists", "path": "state.ready"},
"then": "echo",
}
},
),
(
"choose",
{
"choose": {
"clauses": [
{
"if": {"op": "exists", "path": "state.ready"},
"then": "echo",
}
]
}
},
),
(
"match",
{
"match": {
"value": "state.status",
"cases": [{"equals": "ready", "then": "echo"}],
}
},
),
(
"subgraph",
{
"subgraph": {
"workflow": {"name": "child"},
"outcomes": ["ok"],
}
},
),
],
)
async def test_add_step_accepts_every_typed_draft_step(
tmp_path: Path,
step_name: str,
step_payload: dict[str, Any],
) -> None:
artifact_store = FileWorkflowArtifactStore(tmp_path / f"draft_add_{step_name}")
draft_api, _service, authoring = _draft_api(artifact_store, register_echo=True)
api = WorkflowApi(authoring.context)
await draft_api.create_draft_workspace(workspace_id="draft_ws", draft=_echo_draft())
step = TypeAdapter(DraftStep).validate_python(step_payload)
result = await api.add_step(
workspace_id="draft_ws",
revision=1,
step_id="new_step",
step=step,
)
assert result["revision"] == 2
workspace = await draft_api.get_draft_workspace(
workspace_id="draft_ws", include_draft=True
)
assert workspace["draft"]["steps"]["new_step"] == step.model_dump(
mode="json", by_alias=True
)
@pytest.mark.asyncio
async def test_add_step_routes_incoming_and_outgoing_edges_atomically(
tmp_path: Path,
) -> None:
artifact_store = FileWorkflowArtifactStore(tmp_path / "draft_add_routes")
draft_api, _service, authoring = _draft_api(artifact_store, register_echo=True)
api = WorkflowApi(authoring.context)
await draft_api.create_draft_workspace(workspace_id="draft_ws", draft=_echo_draft())
step = TypeAdapter(DraftStep).validate_python(
{"use": "demo.personal.echo_tool", "input": [], "output": []}
)
result = await api.add_step(
workspace_id="draft_ws",
revision=1,
step_id="new_step",
step=step,
incoming=RouteSource("echo", "ok"),
routes={"ok": "__end__"},
)
assert result["revision"] == 2
workspace = await draft_api.get_draft_workspace(
workspace_id="draft_ws", include_draft=True
)
assert workspace["draft"]["routes"]["echo"]["ok"] == "new_step"
assert workspace["draft"]["routes"]["new_step"] == {"ok": "__end__"}
@pytest.mark.asyncio
async def test_add_step_adds_missing_incoming_route_parent_atomically(
tmp_path: Path,
) -> None:
artifact_store = FileWorkflowArtifactStore(tmp_path / "draft_add_missing_parent")
draft_api, _service, authoring = _draft_api(artifact_store, register_echo=True)
api = WorkflowApi(authoring.context)
draft = _echo_draft()
draft["routes"] = {}
await draft_api.create_draft_workspace(workspace_id="draft_ws", draft=draft)
step = TypeAdapter(DraftStep).validate_python(
{"use": "demo.personal.echo_tool", "input": [], "output": []}
)
result = await api.add_step(
workspace_id="draft_ws",
revision=1,
step_id="new_step",
step=step,
incoming=RouteSource("echo", "ok"),
)
assert result["revision"] == 2
workspace = await draft_api.get_draft_workspace(
workspace_id="draft_ws", include_draft=True
)
assert workspace["draft"]["routes"]["echo"] == {"ok": "new_step"}
async def _assert_add_step_rejected_without_mutation(
api: WorkflowApi,
draft_api: WorkflowDraftApi,
*,
step_id: str = "new_step",
incoming: RouteSource | None = None,
routes: dict[str, str] | None = None,
message: str,
) -> None:
before = await draft_api.get_draft_workspace(
workspace_id="draft_ws", include_draft=True
)
step = TypeAdapter(DraftStep).validate_python(
{"use": "demo.personal.echo_tool", "input": [], "output": []}
)
with pytest.raises(ValueError, match=message):
await api.add_step(
workspace_id="draft_ws",
revision=1,
step_id=step_id,
step=step,
incoming=incoming,
routes=routes,
)
after = await draft_api.get_draft_workspace(
workspace_id="draft_ws", include_draft=True
)
assert after["revision"] == before["revision"]
assert after["draft"] == before["draft"]
@pytest.mark.asyncio
async def test_add_step_rejects_invalid_routing_inputs_atomically(
tmp_path: Path,
) -> None:
artifact_store = FileWorkflowArtifactStore(tmp_path / "draft_add_errors")
draft_api, _service, authoring = _draft_api(artifact_store, register_echo=True)
api = WorkflowApi(authoring.context)
await draft_api.create_draft_workspace(workspace_id="draft_ws", draft=_echo_draft())
await _assert_add_step_rejected_without_mutation(
api,
draft_api,
incoming=RouteSource("missing", "ok"),
message="unknown incoming source",
)
await _assert_add_step_rejected_without_mutation(
api,
draft_api,
step_id="echo",
routes={"ok": "__end__"},
message="already exists",
)
await _assert_add_step_rejected_without_mutation(
api,
draft_api,
routes={"typo": "__end__"},
message="unknown route outcome",
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"step_payload",
[
{"end": {"outcome": "ok"}},
{
"when": {
"if": {"op": "exists", "path": "state.ready"},
"then": "echo",
}
},
{
"choose": {
"clauses": [
{
"if": {"op": "exists", "path": "state.ready"},
"then": "echo",
}
]
}
},
{
"match": {
"value": "state.status",
"cases": [{"equals": "ready", "then": "echo"}],
}
},
],
)
async def test_add_step_rejects_routes_for_non_routable_steps_atomically(
tmp_path: Path,
step_payload: dict[str, Any],
) -> None:
artifact_store = FileWorkflowArtifactStore(tmp_path / "draft_add_forbidden_routes")
draft_api, _service, authoring = _draft_api(artifact_store, register_echo=True)
api = WorkflowApi(authoring.context)
await draft_api.create_draft_workspace(workspace_id="draft_ws", draft=_echo_draft())
before = await draft_api.get_draft_workspace(
workspace_id="draft_ws", include_draft=True
)
step = TypeAdapter(DraftStep).validate_python(step_payload)
with pytest.raises(ValueError, match="routes are not allowed"):
await api.add_step(
workspace_id="draft_ws",
revision=1,
step_id="new_step",
step=step,
routes={"ok": "__end__"},
)
after = await draft_api.get_draft_workspace(
workspace_id="draft_ws", include_draft=True
)
assert after["revision"] == before["revision"]
assert after["draft"] == before["draft"]
@pytest.mark.asyncio
async def test_add_step_accepts_incomplete_declared_route_subset(
tmp_path: Path,
) -> None:
artifact_store = FileWorkflowArtifactStore(tmp_path / "draft_add_partial_routes")
draft_api, service, authoring = _draft_api(artifact_store, register_echo=True)
service.register_specs("demo.personal", echo_tool, _snapshot_tool)
api = WorkflowApi(authoring.context)
await draft_api.create_draft_workspace(workspace_id="draft_ws", draft=_echo_draft())
step = TypeAdapter(DraftStep).validate_python(
{"use": "demo.personal.snapshot_tool", "input": [], "output": []}
)
result = await api.add_step(
workspace_id="draft_ws",
revision=1,
step_id="new_step",
step=step,
routes={"ok": "__end__"},
)
assert result["revision"] == 2
workspace = await draft_api.get_draft_workspace(
workspace_id="draft_ws", include_draft=True
)
assert workspace["draft"]["routes"]["new_step"] == {"ok": "__end__"}
@pytest.mark.asyncio
async def test_branch_draft_updates_routes_atomically(tmp_path: Path) -> None:
artifact_store = FileWorkflowArtifactStore(tmp_path / "drafts_branch")
@@ -894,8 +1194,8 @@ async def test_handle_draft_updates_multiple_source_outcomes(tmp_path: Path) ->
workspace_id="handling",
revision=1,
branches=[
DraftOutcomeRef(step_id="lookup", outcome="error"),
DraftOutcomeRef(step_id="transform", outcome="error"),
RouteSource(step_id="lookup", outcome="error"),
RouteSource(step_id="transform", outcome="error"),
],
target="__end__",
)