fix: smooth draft schema projection ux

src/wf_api/draft_authoring.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from typing import Any
 
@@ -64,6 +64,18 @@ def _local_parts(path: str) -> tuple[str, ...]:
     return LocalPath.parse(path.removeprefix("local.")).parts
 
 
+def _schema_path_exists(schema: Mapping[str, Any], parts: Sequence[str]) -> bool:
+    current: Any = schema
+    for part in parts:
+        if not isinstance(current, Mapping):
+            return False
+        properties = current.get("properties")
+        if not isinstance(properties, Mapping) or part not in properties:
+            return False
+        current = properties[part]
+    return True
+
+
 class WorkflowDraftAuthoringApi:
     """Capability-aware semantic edits over revisioned workflow drafts."""
 
@@ -343,6 +355,7 @@ class WorkflowDraftAuthoringApi:
                 source_schema=output_schema,
                 source_field=local_field,
                 target_parts=target_parts,
+                allow_existing_equivalent=True,
             )
             output_map = {
                 **self.drafts._step_output_map(
@@ -434,7 +447,40 @@ class WorkflowDraftAuthoringApi:
 
         input_map = input_map or {}
         bind_outputs = bind_outputs or {}
+        projected_input_schema = workspace.draft.get("input_schema", {})
+        if not isinstance(projected_input_schema, dict):
+            raise ValueError("draft input_schema must be an object")
         projected_state_schema = state_schema
+        input_schema = (
+            spec.input_schema_contract or spec.input_model.model_json_schema()
+        )
+        for graph_path, local_path in input_map.items():
+            try:
+                source_root, source_parts = _graph_parts(graph_path)
+                local_parts = LocalPath.parse(local_path).parts
+            except ValueError:
+                continue
+            if source_root not in {"input", "state"} or len(local_parts) != 1:
+                continue
+            schema_key = "input_schema" if source_root == "input" else "state_schema"
+            target_schema = (
+                projected_input_schema
+                if source_root == "input"
+                else projected_state_schema
+            )
+            if _schema_path_exists(target_schema, source_parts):
+                continue
+            projected = project_property_to_schema_path(
+                target_schema=target_schema,
+                source_schema=input_schema,
+                source_field=local_parts[0],
+                target_parts=source_parts,
+                allow_existing_equivalent=True,
+            )
+            if schema_key == "input_schema":
+                projected_input_schema = projected
+            else:
+                projected_state_schema = projected
         for output_field, path in bind_outputs.items():
             sf = state_root_field(path)
             projected_state_schema = project_output_property_to_state_schema(
@@ -442,6 +488,7 @@ class WorkflowDraftAuthoringApi:
                 output_schema=output_schema,
                 output_field=output_field,
                 state_field=sf,
+                allow_existing_equivalent=True,
             )
 
         patch: list[dict[str, Any]] = [
@@ -460,6 +507,15 @@ class WorkflowDraftAuthoringApi:
                 "value": step_routes,
             },
         ]
+        if projected_input_schema != workspace.draft.get("input_schema", {}):
+            patch.insert(
+                0,
+                {
+                    "op": "replace",
+                    "path": "/input_schema",
+                    "value": projected_input_schema,
+                },
+            )
         if projected_state_schema != state_schema:
             patch.insert(
                 0,

src/wf_api/drafts.py

@@ -25,6 +25,7 @@ from wf_core.models.steps import (
     InputValueBinding,
     OutputBinding,
 )
+from wf_core.paths import GraphSourcePath, parse_toml_path_segments
 
 from .capability_requirements import (
     required_capabilities_for_plan,
@@ -43,6 +44,7 @@ from .draft_payloads import (
     output_bindings_payload as _draft_output_bindings_payload,
 )
 from .operation_context import WorkflowOperationContext
+from .schema_projection import project_property_to_schema_path
 
 
 class WorkflowDraftApi:
@@ -339,18 +341,88 @@ class WorkflowDraftApi:
                 {"path": source, "target": target}
                 for source, target in output_map.items()
             ]
+        workspace = self._draft_store().get_workspace(workspace_id)
+        output_schema = self._workflow_output_schema_for_bindings(
+            draft=workspace.draft,
+            output_bindings=output_bindings,
+        )
+        patch = [
+            {
+                "op": "replace",
+                "path": "/output",
+                "value": output_bindings,
+            }
+        ]
+        if output_schema is not workspace.draft.get("output_schema"):
+            patch.insert(
+                0,
+                {
+                    "op": "replace",
+                    "path": "/output_schema",
+                    "value": output_schema,
+                },
+            )
         return await self.patch_draft_workspace(
             workspace_id=workspace_id,
             revision=revision,
-            patch=[
-                {
-                    "op": "replace",
-                    "path": "/output",
-                    "value": output_bindings,
-                }
-            ],
+            patch=patch,
         )
 
+    def _workflow_output_schema_for_bindings(
+        self,
+        *,
+        draft: dict[str, Any],
+        output_bindings: Sequence[Mapping[str, Any]],
+    ) -> dict[str, Any]:
+        """Project missing top-level output fields from input/state schemas.
+
+        This is intentionally conservative: only single-field ``input.x`` and
+        ``state.x`` sources can be copied unambiguously. More complex sources
+        still fall through to existing validation diagnostics instead of
+        guessing a schema.
+        """
+        output_schema = draft.get("output_schema", {})
+        if not isinstance(output_schema, dict):
+            raise ValueError("draft output_schema must be an object")
+        projected = output_schema
+        changed = False
+        for binding in output_bindings:
+            source = binding.get("path")
+            target = binding.get("target")
+            if not isinstance(source, str) or not isinstance(target, str):
+                continue
+            source_schema = _workflow_source_schema(draft, source)
+            if source_schema is None:
+                continue
+            try:
+                target_parts = parse_toml_path_segments(target)
+            except ValueError:
+                continue
+            if _schema_path_exists(projected, target_parts):
+                continue
+            try:
+                source_path = GraphSourcePath.parse(source)
+            except ValueError:
+                continue
+            if len(source_path.parts) != 1:
+                continue
+            try:
+                updated = project_property_to_schema_path(
+                    target_schema=projected,
+                    source_schema=source_schema,
+                    source_field=source_path.parts[0],
+                    target_parts=target_parts,
+                    allow_existing_equivalent=True,
+                )
+            except ValueError as exc:
+                if str(exc).startswith("source field "):
+                    continue
+                raise
+            if updated != projected:
+                changed = True
+                projected = updated
+        return projected if changed else output_schema
+
     def _step_input_maps(
         self,
         *,
@@ -367,6 +439,35 @@ class WorkflowDraftApi:
         return _output_map_from_payload(step.get("output", []))
 
 
+def _workflow_source_schema(
+    draft: Mapping[str, Any],
+    source_path: str,
+) -> dict[str, Any] | None:
+    try:
+        parsed = GraphSourcePath.parse(source_path)
+    except ValueError:
+        return None
+    if parsed.root == "input":
+        schema = draft.get("input_schema")
+    elif parsed.root == "state":
+        schema = draft.get("state_schema")
+    else:
+        return None
+    return schema if isinstance(schema, dict) else None
+
+
+def _schema_path_exists(schema: Mapping[str, Any], parts: Sequence[str]) -> bool:
+    current: Any = schema
+    for part in parts:
+        if not isinstance(current, Mapping):
+            return False
+        properties = current.get("properties")
+        if not isinstance(properties, Mapping) or part not in properties:
+            return False
+        current = properties[part]
+    return True
+
+
 def _draft_input_maps(
     *,
     input: Sequence[InputBinding] | None,

src/wf_api/schema_projection.py

@@ -75,6 +75,7 @@ def project_output_property_to_state_schema(
     output_schema: JsonObject,
     output_field: str,
     state_field: str,
+    allow_existing_equivalent: bool = False,
 ) -> JsonObject:
     """Root state projection convenience wrapper.
 
@@ -86,6 +87,7 @@ def project_output_property_to_state_schema(
             source_schema=output_schema,
             source_field=output_field,
             target_parts=(state_field,),
+            allow_existing_equivalent=allow_existing_equivalent,
         )
     except ValueError as exc:
         msg = str(exc)

src/wf_cli/commands/drafts.py

@@ -364,6 +364,9 @@ def set_workflow_output(
     This edits WorkflowDraft.output (top-level workflow output). Use
     wf draft set-output for step-level output bindings.
 
+    For single-field input/state sources, missing output_schema fields are
+    projected automatically from the source schema.
+
     Repeat --map for multiple mappings:
     --map state.markdown=markdown --map state.title=title