feat: make workflow builder edits lossless

src/wf_authoring/builder/__init__.py

@@ -1,4 +1,5 @@
 from .core import WorkflowBuilder
+from .mapping import auto_input_map_from_schema, auto_output_map_from_schema
 from .refs import BranchRef, BranchResult, DecisionResult, HandleResult, StepRef
 
 __all__ = [
@@ -8,4 +9,6 @@ __all__ = [
     "HandleResult",
     "StepRef",
     "WorkflowBuilder",
+    "auto_input_map_from_schema",
+    "auto_output_map_from_schema",
 ]

src/wf_authoring/builder/core.py

@@ -15,6 +15,7 @@ from wf_core import (
     ForeachItemErrorPolicy,
     ForeachNode,
     InterruptNode,
+    NodeDef,
     NodeHandler,
     NodeUse,
     PreparedSubgraph,
@@ -22,6 +23,7 @@ from wf_core import (
     SchemaRef,
     StateSchema,
     SubgraphNode,
+    ValidationReport,
     Workflow,
     WorkflowRef,
     execute_workflow,
@@ -31,6 +33,7 @@ from wf_core.errors import WorkflowExecutionError
 from wf_core.models.conditions import BinaryCondition, ExistsCondition, PathOperand
 from wf_core.models.conditions import Condition as CoreCondition
 from wf_core.models.steps import (
+    InputBinding,
     InputPathBinding,
     InputValueBinding,
     OutputBinding,
@@ -54,7 +57,9 @@ from .mapping import (
     OutputBindingArg,
     StepInputBindingArg,
     auto_input_map,
+    auto_input_map_from_schema,
     auto_output_map,
+    auto_output_map_from_schema,
     coerce_path,
     normalize_input_mapping,
     normalize_input_values,
@@ -136,6 +141,11 @@ def _canonical_output_bindings(
     ]
 
 
+def _node_defs_compatible(left: NodeDef, right: NodeDef) -> bool:
+    """Compare contracts by their serialized canonical content."""
+    return left.model_dump(mode="json") == right.model_dump(mode="json")
+
+
 def _reject_mixed_binding_styles(
     *,
     input: object | None,
@@ -213,6 +223,8 @@ class WorkflowBuilder:
     node_specs: dict[str, NodeSpec[Any, Any]] = field(default_factory=dict)
     nodes: list[Step] = field(default_factory=list)
     edges: list[Edge] = field(default_factory=list)
+    workflow_output: list[InputBinding] = field(default_factory=list)
+    seeded_node_defs: dict[str, NodeDef] = field(default_factory=dict, repr=False)
     prepared_subgraphs: dict[str, PreparedSubgraph[NodeHandler]] = field(
         default_factory=dict
     )
@@ -392,6 +404,134 @@ class WorkflowBuilder:
         self.nodes.append(node)
         return node
 
+    def use_contract(
+        self,
+        node_def: NodeDef,
+        *,
+        id: str | None = None,
+        input: Sequence[StepInputBindingArg] | None = None,
+        output: Sequence[OutputBindingArg] | None = None,
+        desc: str | None = None,
+    ) -> NodeUse:
+        """Use a schema-backed external node contract without a local handler."""
+        existing = self.seeded_node_defs.get(node_def.name)
+        if existing is not None and not _node_defs_compatible(existing, node_def):
+            raise ValueError(
+                f"incompatible duplicate node definition {node_def.name!r}"
+            )
+        self.seeded_node_defs[node_def.name] = node_def.model_copy(deep=True)
+        normalized_input_schema = cast(SchemaRef, self.input_schema)
+        normalized_state_schema = cast(StateSchema, self.state_schema)
+        node_input = (
+            normalize_step_input_bindings(input)
+            if input is not None
+            else _canonical_input_bindings(
+                normalize_input_mapping(
+                    auto_input_map_from_schema(
+                        node_def.input_schema,
+                        input_schema=normalized_input_schema,
+                        state_schema=normalized_state_schema,
+                    )
+                ),
+                {},
+            )
+        )
+        node_output = (
+            normalize_output_bindings(output)
+            if output is not None
+            else _canonical_output_bindings(
+                normalize_output_mapping(
+                    auto_output_map_from_schema(
+                        node_def.output_schema,
+                        state_schema=normalized_state_schema,
+                    )
+                )
+            )
+        )
+        return self.use_ref(
+            node_def.name,
+            id=id,
+            input=node_input,
+            output=node_output,
+            desc=desc,
+        )
+
+    @classmethod
+    def from_workflow(cls, workflow: Workflow) -> WorkflowBuilder:
+        """Create an independently editable builder from a canonical workflow."""
+        return cls(
+            name=workflow.name,
+            input_schema=workflow.input_schema.model_copy(deep=True),
+            state_schema=workflow.state_schema.model_copy(deep=True),
+            output_schema=workflow.output_schema.model_copy(deep=True),
+            outcomes=tuple(workflow.outcomes),
+            start=workflow.start,
+            nodes=[node.model_copy(deep=True) for node in workflow.nodes],
+            edges=[edge.model_copy(deep=True) for edge in workflow.edges],
+            workflow_output=[
+                binding.model_copy(deep=True) for binding in workflow.output
+            ],
+            seeded_node_defs={
+                node_def.name: node_def.model_copy(deep=True)
+                for node_def in workflow.node_defs
+            },
+        )
+
+    def set_output(self, bindings: Sequence[StepInputBindingArg]) -> None:
+        """Replace final workflow output projection bindings."""
+        self.workflow_output = cast(
+            list[InputBinding], normalize_step_input_bindings(bindings)
+        )
+
+    def set_route(self, source: StepRef, outcome: str, target: StepRef) -> None:
+        """Replace the unique route for one source/outcome pair."""
+        source_id = step_id(source)
+        target_id = step_id(target)
+        self.edges = [
+            edge
+            for edge in self.edges
+            if not (edge.from_ == source_id and edge.outcome == outcome)
+        ]
+        self.edges.append(
+            Edge.model_validate(
+                {"from": source_id, "outcome": outcome, "to": target_id}
+            )
+        )
+
+    def remove_route(self, source: StepRef, outcome: str) -> None:
+        """Remove one route, raising when the requested route is absent."""
+        source_id = step_id(source)
+        matching = [
+            edge
+            for edge in self.edges
+            if edge.from_ == source_id and edge.outcome == outcome
+        ]
+        if not matching:
+            raise ValueError(
+                f"route from step {source_id!r} with outcome {outcome!r} not found"
+            )
+        self.edges = [
+            edge
+            for edge in self.edges
+            if not (edge.from_ == source_id and edge.outcome == outcome)
+        ]
+
+    def remove_step(self, step: StepRef) -> None:
+        """Remove an unreferenced step, rejecting dangling graph references."""
+        step_id_value = step_id(step)
+        if not any(node.id == step_id_value for node in self.nodes):
+            raise ValueError(f"step {step_id_value!r} not found")
+        if self.start == step_id_value:
+            raise ValueError(
+                f"step {step_id_value!r} is still referenced as workflow start"
+            )
+        if any(
+            edge.from_ == step_id_value or edge.to == step_id_value
+            for edge in self.edges
+        ):
+            raise ValueError(f"step {step_id_value!r} is still referenced by route")
+        self.nodes = [node for node in self.nodes if node.id != step_id_value]
+
     def subgraph(
         self,
         *,
@@ -869,21 +1009,45 @@ class WorkflowBuilder:
             )
         return self.match(value, cases, id=id, default=default)
 
+    def _build_workflow(self, *, start: str) -> Workflow:
+        """Build a canonical workflow snapshot from current builder state."""
+        node_defs = [
+            node_def.model_copy(deep=True)
+            for node_def in self.seeded_node_defs.values()
+        ]
+        by_name = {node_def.name: node_def for node_def in node_defs}
+        for spec in self.node_specs.values():
+            node_def = spec.to_node_def()
+            existing = by_name.get(node_def.name)
+            if existing is not None:
+                if not _node_defs_compatible(existing, node_def):
+                    raise ValueError(
+                        f"incompatible duplicate node definition {node_def.name!r}"
+                    )
+                continue
+            by_name[node_def.name] = node_def
+            node_defs.append(node_def)
+        return Workflow(
+            name=self.name,
+            input_schema=cast(SchemaRef, self.input_schema).model_copy(deep=True),
+            state_schema=cast(StateSchema, self.state_schema).model_copy(deep=True),
+            output_schema=cast(SchemaRef, self.output_schema).model_copy(deep=True),
+            outcomes=list(self.outcomes),
+            node_defs=node_defs,
+            start=start,
+            output=[binding.model_copy(deep=True) for binding in self.workflow_output],
+            nodes=[node.model_copy(deep=True) for node in self.nodes],
+            edges=[edge.model_copy(deep=True) for edge in self.edges],
+        )
+
+    def validate_structure(self) -> ValidationReport:
+        """Return structural issues, including an unset workflow start."""
+        return self._build_workflow(start=self.start or "").validate_structure()
+
     def compile(self) -> Workflow:
         if self.start is None:
             raise WorkflowExecutionError(
                 "workflow builder requires an explicit start; "
                 "call set_entry_point(...) or pass start=..."
             )
-        node_defs = [spec.to_node_def() for spec in self.node_specs.values()]
+        return self._build_workflow(start=self.start)
-        return Workflow(
-            name=self.name,
-            input_schema=cast(SchemaRef, self.input_schema),
-            state_schema=cast(StateSchema, self.state_schema),
-            output_schema=cast(SchemaRef, self.output_schema),
-            outcomes=list(self.outcomes),
-            node_defs=node_defs,
-            start=self.start,
-            nodes=self.nodes,
-            edges=self.edges,
-        )

src/wf_authoring/builder/mapping.py

@@ -143,11 +143,25 @@ def auto_input_map(
     state_schema: StateSchema,
 ) -> dict[str, str]:
     """Map node input fields from state first, then workflow input."""
+    return auto_input_map_from_schema(
+        spec.to_node_def().input_schema,
+        input_schema=input_schema,
+        state_schema=state_schema,
+    )
+
+
+def auto_input_map_from_schema(
+    capability_input_schema: SchemaRef,
+    *,
+    input_schema: SchemaRef,
+    state_schema: StateSchema,
+) -> dict[str, str]:
+    """Map schema-declared capability inputs from state or workflow input."""
     return {
         _auto_source_path(
             field, input_schema=input_schema, state_schema=state_schema
         ): field
-        for field in spec.input_model.model_json_schema().get("properties", {})
+        for field in capability_input_schema.properties
     }
 
 
@@ -157,10 +171,22 @@ def auto_output_map(
     state_schema: StateSchema,
 ) -> dict[str, str]:
     """Map node output fields back into matching state fields."""
+    return auto_output_map_from_schema(
+        spec.to_node_def().output_schema,
+        state_schema=state_schema,
+    )
+
+
+def auto_output_map_from_schema(
+    capability_output_schema: SchemaRef,
+    *,
+    state_schema: StateSchema,
+) -> dict[str, str]:
+    """Map schema-declared capability outputs into matching state fields."""
     state_fields = state_schema.field_map()
     return {
         field: f"state.{field}"
-        for field in spec.output_model.model_json_schema().get("properties", {})
+        for field in capability_output_schema.properties
         if field in state_fields
     }
 

tests/authoring/test_builder.py

@@ -21,7 +21,15 @@ from wf_authoring import (
     state_path,
 )
 from wf_authoring.builder.mapping import normalize_input_mapping
-from wf_core import END, EndNode, RunStatus, WorkflowExecutionError
+from wf_core import (
+    END,
+    EndNode,
+    NodeDef,
+    RunStatus,
+    ValidationIssueCode,
+    Workflow,
+    WorkflowExecutionError,
+)
 from wf_core.models.steps import (
     InputExpressionBinding,
     InputPathBinding,
@@ -31,6 +39,163 @@ from wf_core.paths import GraphSourcePath, LocalPath, StatePath
 from wf_platform import CapabilityRef
 
 
+def _editable_three_step_builder() -> WorkflowBuilder:
+    builder = WorkflowBuilder(
+        name="editable_three_step",
+        input_schema={"type": "object", "properties": {}},
+        state_schema={"type": "object", "properties": {}},
+        output_schema={"type": "object", "properties": {}},
+        start="first",
+    )
+    builder.use_ref("demo.first", id="first")
+    builder.use_ref("demo.second", id="second")
+    builder.use_ref("demo.third", id="third")
+    builder.connect("first", "ok", "second")
+    builder.connect("second", "ok", "third")
+    builder.connect("third", "ok", END)
+    return builder
+
+
+def test_builder_round_trip_preserves_complete_workflow() -> None:
+    original = Workflow.model_validate(
+        {
+            "name": "round_trip",
+            "input_schema": {
+                "type": "object",
+                "properties": {"topic": {"type": "string"}},
+            },
+            "state_schema": {
+                "type": "object",
+                "properties": {"result": {"type": "string"}},
+            },
+            "output_schema": {
+                "type": "object",
+                "properties": {"result": {"type": "string"}},
+            },
+            "output": [{"path": "state.result", "target": "result"}],
+            "node_defs": [
+                {
+                    "name": "app.default.search",
+                    "input_schema": {"type": "object", "properties": {}},
+                    "output_schema": {
+                        "type": "object",
+                        "properties": {"result": {"type": "string"}},
+                    },
+                    "outcomes": ["ok"],
+                }
+            ],
+            "outcomes": ["ok"],
+            "start": "search",
+            "nodes": [
+                {
+                    "id": "search",
+                    "type": "node",
+                    "node": "app.default.search",
+                    "input": [],
+                    "output": [{"source": "result", "target": "state.result"}],
+                },
+                {"id": "end_ok", "type": "end", "outcome": "ok"},
+            ],
+            "edges": [{"from": "search", "outcome": "ok", "to": "end_ok"}],
+        }
+    )
+
+    builder = WorkflowBuilder.from_workflow(original)
+    rebuilt = builder.compile()
+
+    assert rebuilt.model_dump(mode="json", by_alias=True) == original.model_dump(
+        mode="json", by_alias=True
+    )
+    builder.set_output([{"value": "changed", "target": "result"}])
+    original_output = original.output[0]
+    assert isinstance(original_output, InputPathBinding)
+    assert str(original_output.path) == "state.result"
+
+
+def test_set_route_replaces_unique_source_outcome_edge() -> None:
+    builder = _editable_three_step_builder()
+
+    builder.set_route("first", "ok", "third")
+
+    matching = [
+        edge for edge in builder.edges if edge.from_ == "first" and edge.outcome == "ok"
+    ]
+    assert [(edge.from_, edge.outcome, edge.to) for edge in matching] == [
+        ("first", "ok", "third")
+    ]
+
+
+def test_remove_route_removes_only_requested_source_outcome_pair() -> None:
+    builder = _editable_three_step_builder()
+    builder.connect("first", "error", "third")
+
+    builder.remove_route("first", "ok")
+
+    assert [(edge.from_, edge.outcome) for edge in builder.edges] == [
+        ("second", "ok"),
+        ("third", "ok"),
+        ("first", "error"),
+    ]
+
+
+def test_remove_step_rejects_referenced_step() -> None:
+    builder = _editable_three_step_builder()
+
+    with pytest.raises(ValueError, match="still referenced by route"):
+        builder.remove_step("second")
+
+
+def test_use_contract_registers_schema_only_external_node() -> None:
+    builder = WorkflowBuilder(
+        name="contract_builder",
+        input_schema={"type": "object", "properties": {"topic": {"type": "string"}}},
+        state_schema={
+            "type": "object",
+            "properties": {"result": {"type": "string"}},
+        },
+        output_schema={"type": "object", "properties": {"result": {"type": "string"}}},
+    )
+    contract = NodeDef.model_validate(
+        {
+            "name": "app.remote.search",
+            "input_schema": {
+                "type": "object",
+                "properties": {"topic": {"type": "string"}},
+            },
+            "output_schema": {
+                "type": "object",
+                "properties": {"result": {"type": "string"}},
+            },
+            "outcomes": ["ok"],
+        }
+    )
+
+    step = builder.use_contract(contract, id="search")
+    builder.set_entry_point(step)
+    builder.connect(step, "ok", END)
+
+    assert step.node == contract.name
+    assert contract.name in {node_def.name for node_def in builder.compile().node_defs}
+    assert builder.node_specs == {}
+    assert builder.registry() == {}
+    assert isinstance(step.input[0], InputPathBinding)
+    assert step.input[0].path == GraphSourcePath.input("topic")
+    assert step.output[0].target == StatePath.of("result")
+
+
+def test_validate_structure_reports_unknown_start_when_unset() -> None:
+    builder = WorkflowBuilder(
+        name="missing_start",
+        input_schema={},
+        state_schema={"type": "object"},
+        output_schema={},
+    )
+
+    report = builder.validate_structure()
+
+    assert report.errors[0].code == ValidationIssueCode.UNKNOWN_START
+
+
 def test_builder_auto_binds_matching_node_inputs_and_outputs_to_state() -> None:
     builder = WorkflowBuilder(
         name="auto_bind_demo",