feat: add atomic generic draft step insertion

src/wf_api/__init__.py

@@ -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

@@ -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:
+class RouteSource:
-    """A reference to a specific outcome of a draft step."""
+    """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

@@ -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"])
+            RouteSource(step_id=b["step_id"], outcome=b["outcome"]) for b in branches
-            for b in branches
         ]
         return await self.draft_authoring.handle_draft(
             workspace_id=workspace_id,

src/wf_api/surface.py

@@ -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

@@ -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"),
+            RouteSource(step_id="lookup", outcome="error"),
-            DraftOutcomeRef(step_id="transform", outcome="error"),
+            RouteSource(step_id="transform", outcome="error"),
         ],
         target="__end__",
     )