feat: add draft state schema projection helper

src/wf_api/drafts.py

@@ -38,6 +38,7 @@ from .constants import (
     RUNTIME_ERROR_CAPABILITY,
 )
 from .operation_context import WorkflowOperationContext
+from .schema_projection import project_output_property_to_state_schema
 
 
 class WorkflowDraftApi:
@@ -254,6 +255,48 @@ class WorkflowDraftApi:
             ],
         )
 
+    async def add_state_schema_from_output(
+        self,
+        *,
+        workspace_id: str,
+        revision: int,
+        step_id: str,
+        output_field: str,
+        state_path: str,
+    ) -> dict[str, Any]:
+        workspace = self._draft_store().get_workspace(workspace_id)
+        step = _draft_step(workspace.draft, step_id)
+        capability_name = step.get("use")
+        if not isinstance(capability_name, str) or not capability_name:
+            raise ValueError(
+                f"draft step {step_id!r} does not declare a capability use"
+            )
+        state_field = _state_root_field(state_path)
+        spec = self.context.specs.get_qualified_spec(capability_name)
+        output_schema = (
+            spec.output_schema_contract or spec.output_model.model_json_schema()
+        )
+        state_schema = workspace.draft.get("state_schema", {})
+        if not isinstance(state_schema, dict):
+            raise ValueError("draft state_schema must be an object")
+        projected = project_output_property_to_state_schema(
+            state_schema=state_schema,
+            output_schema=output_schema,
+            output_field=output_field,
+            state_field=state_field,
+        )
+        return await self.patch_draft_workspace(
+            workspace_id=workspace_id,
+            revision=revision,
+            patch=[
+                {
+                    "op": "replace",
+                    "path": "/state_schema",
+                    "value": projected,
+                }
+            ],
+        )
+
     def _step_input_maps(
         self,
         *,
@@ -485,3 +528,10 @@ def _state_path_payload(value: str) -> dict[str, str | list[str]]:
 def _escape_json_pointer(value: str) -> str:
     """Escape one JSON Pointer path segment for generated JSON Patch helpers."""
     return value.replace("~", "~0").replace("/", "~1")
+
+
+def _state_root_field(value: str) -> str:
+    path = StatePath.parse(value)
+    if len(path.parts) != 1:
+        raise ValueError("state_path must name one root field, such as state.after")
+    return path.parts[0]

src/wf_api/schema_projection.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+from copy import deepcopy
+from typing import Any
+
+from jsonschema import Draft202012Validator, SchemaError
+
+JsonObject = dict[str, Any]
+
+
+def project_output_property_to_state_schema(
+    *,
+    state_schema: JsonObject,
+    output_schema: JsonObject,
+    output_field: str,
+    state_field: str,
+) -> JsonObject:
+    """Project one capability output property schema into workflow state schema.
+
+    Capability output schemas may use local references such as
+    ``{"$ref": "#/$defs/Snapshot"}``. Copying only the property schema would
+    create dangling references, so this helper also merges local definition
+    blocks and rejects conflicting definition names.
+    """
+    _check_schema("state_schema", state_schema)
+    _check_schema("output_schema", output_schema)
+    output_properties = output_schema.get("properties")
+    if not isinstance(output_properties, dict) or output_field not in output_properties:
+        raise ValueError(f"output field {output_field!r} is not declared")
+    output_property = output_properties[output_field]
+    if not isinstance(output_property, dict):
+        raise ValueError(f"output field {output_field!r} is not a JSON Schema object")
+
+    projected = deepcopy(state_schema)
+    projected.setdefault("type", "object")
+    properties = projected.setdefault("properties", {})
+    if not isinstance(properties, dict):
+        raise ValueError("state_schema.properties must be an object")
+    properties[state_field] = deepcopy(output_property)
+
+    _merge_definition_block(projected, output_schema, "$defs")
+    _merge_definition_block(projected, output_schema, "definitions")
+    _check_schema("projected state_schema", projected)
+    return projected
+
+
+def _check_schema(name: str, schema: JsonObject) -> None:
+    try:
+        Draft202012Validator.check_schema(schema)
+    except SchemaError as exc:
+        raise ValueError(f"{name} is not valid JSON Schema: {exc.message}") from exc
+
+
+def _merge_definition_block(
+    target_schema: JsonObject,
+    source_schema: JsonObject,
+    key: str,
+) -> None:
+    source_defs = source_schema.get(key)
+    if source_defs is None:
+        return
+    if not isinstance(source_defs, dict):
+        raise ValueError(f"output_schema.{key} must be an object")
+    target_defs = target_schema.setdefault(key, {})
+    if not isinstance(target_defs, dict):
+        raise ValueError(f"state_schema.{key} must be an object")
+    for name, definition in source_defs.items():
+        if name in target_defs and target_defs[name] != definition:
+            raise ValueError(f"conflicting {key}.{name}")
+        target_defs[name] = deepcopy(definition)

src/wf_api/service.py

@@ -362,6 +362,23 @@ class WorkflowApi:
             merge=merge,
         )
 
+    async def add_state_schema_from_output(
+        self,
+        *,
+        workspace_id: str,
+        revision: int,
+        step_id: str,
+        output_field: str,
+        state_path: str,
+    ) -> dict[str, Any]:
+        return await self.drafts.add_state_schema_from_output(
+            workspace_id=workspace_id,
+            revision=revision,
+            step_id=step_id,
+            output_field=output_field,
+            state_path=state_path,
+        )
+
     async def create_minimal_draft_workspace(
         self,
         *,

src/wf_api/surface.py

@@ -114,6 +114,16 @@ class WorkflowDraftSurface(Protocol):
         merge: bool = False,
     ) -> dict[str, Any]: ...
 
+    async def add_state_schema_from_output(
+        self,
+        *,
+        workspace_id: str,
+        revision: int,
+        step_id: str,
+        output_field: str,
+        state_path: str,
+    ) -> dict[str, Any]: ...
+
     async def validate_draft_workspace(
         self,
         *,

src/wf_cli/commands/drafts.py

@@ -277,6 +277,47 @@ def set_step_output_map(
     )
 
 
+@app.command("add-state-from-output")
+def add_state_from_output(
+    ctx: typer.Context,
+    workspace_id: Annotated[str, typer.Argument(help="Draft workspace id.")],
+    revision: Annotated[
+        int, typer.Option("--revision", min=1, help="Expected workspace revision.")
+    ],
+    step_id: Annotated[str, typer.Option("--step", help="Draft step id.")],
+    output_field: Annotated[
+        str,
+        typer.Option("--output", help="Top-level capability output field."),
+    ],
+    state_path: Annotated[
+        str,
+        typer.Option("--state", help="Root state path, for example state.after."),
+    ],
+) -> None:
+    """Copy one capability output field schema into draft state_schema.
+
+    Use this before mapping a step output into a new state field. The command
+    reads the selected draft step's capability output schema, copies the
+    requested output property schema, and preserves local $defs/definitions so
+    JSON Schema refs remain valid.
+
+    Run `wf draft validate <workspace_id>` after adding state schema fields.
+    """
+    context = load_cli_context(ctx)
+    emit_json(
+        run_cli_operation(
+            context,
+            context.handlers.add_state_schema_from_output(
+                workspace_id=workspace_id,
+                revision=revision,
+                step_id=step_id,
+                output_field=output_field,
+                state_path=state_path,
+            ),
+        )
+    )
+
+
 @app.command("validate")
 def validate_draft(
     ctx: typer.Context,

src/wf_mcp/workflow_surface/models.py

@@ -237,6 +237,20 @@ class SetStepOutputMapRequest(BaseModel):
     )
 
 
+class AddStateFromOutputRequest(BaseModel):
+    """Typed MCP request for declaring a state field from a step output schema."""
+
+    workspace_id: WorkspaceId
+    revision: int = Field(ge=1, description="Expected current workspace revision.")
+    step_id: str = Field(description="Draft step id whose capability output is used.")
+    output_field: str = Field(
+        description="Top-level output field to copy, for example after."
+    )
+    state_path: str = Field(
+        description="Root state path to declare, for example state.after."
+    )
+
+
 class DeleteDraftWorkspaceRequest(BaseModel):
     """Typed MCP request payload for deleting one draft workspace."""
 

src/wf_mcp/workflow_surface/tools.py

@@ -12,6 +12,7 @@ from wf_mcp.broker.service import WfMcpService
 from wf_mcp.broker.service.workflow_operation_context import context_from_service
 
 from .models import (
+    AddStateFromOutputRequest,
     CallCapabilityResult,
     CreateArtifactFromWorkspaceRequest,
     CreateDraftWorkspaceFromCapabilityRequest,
@@ -439,6 +440,27 @@ def register_workflow_tools(server: FastMCP[Any], service: WfMcpService) -> None
             )
         )
 
+    @server.tool(
+        name="wf.workflow.add_state_from_output",
+        title="Add State From Output",
+        description=(
+            "Declare one root state field by copying a draft step capability output "
+            "field schema, including local $defs/definitions when present."
+        ),
+    )
+    async def add_state_from_output(
+        request: AddStateFromOutputRequest,
+    ) -> DraftWorkspaceResult:
+        return DraftWorkspaceResult.model_validate(
+            await handlers.add_state_schema_from_output(
+                workspace_id=request.workspace_id,
+                revision=request.revision,
+                step_id=request.step_id,
+                output_field=request.output_field,
+                state_path=request.state_path,
+            )
+        )
+
     @server.tool(
         name="wf.workflow.create_minimal_draft_workspace",
         title="Create Minimal Draft Workspace",

src/wf_transport_rpc_http/client/drafts.py

@@ -141,6 +141,26 @@ class RpcDraftClientMixin:
             },
         )
 
+    async def add_state_schema_from_output(
+        self: RpcCaller,
+        *,
+        workspace_id: str,
+        revision: int,
+        step_id: str,
+        output_field: str,
+        state_path: str,
+    ) -> dict[str, Any]:
+        return await self._call(
+            "workflow.draft_workspaces.add_state_from_output",
+            {
+                "workspace_id": workspace_id,
+                "revision": revision,
+                "step_id": step_id,
+                "output_field": output_field,
+                "state_path": state_path,
+            },
+        )
+
     async def validate_draft_workspace(
         self: RpcCaller,
         *,

src/wf_transport_rpc_http/methods/drafts.py

@@ -8,6 +8,7 @@ from wf_server import WorkflowServer
 
 from ..errors import WorkflowRpcError, raise_workflow_rpc_error
 from ..models import (
+    AddStateFromOutputParams,
     CreateArtifactFromWorkspaceParams,
     CreateDraftFromCapabilityParams,
     CreateWrapperFromWorkspaceParams,
@@ -177,6 +178,24 @@ def register_methods(
         except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
             raise_workflow_rpc_error(exc)
 
+    @entrypoint.method(
+        name="workflow.draft_workspaces.add_state_from_output",
+        errors=[WorkflowRpcError],
+    )
+    async def workflow_draft_workspaces_add_state_from_output(
+        params: AddStateFromOutputParams = RpcParams(),
+    ) -> dict[str, Any]:
+        try:
+            return await server.api.add_state_schema_from_output(
+                workspace_id=params.workspace_id,
+                revision=params.revision,
+                step_id=params.step_id,
+                output_field=params.output_field,
+                state_path=params.state_path,
+            )
+        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
+            raise_workflow_rpc_error(exc)
+
     @entrypoint.method(
         name="workflow.draft_workspaces.validate", errors=[WorkflowRpcError]
     )

src/wf_transport_rpc_http/models.py

@@ -141,6 +141,14 @@ class SetStepOutputMapParams(RpcParamsModel):
     merge: bool = False
 
 
+class AddStateFromOutputParams(RpcParamsModel):
+    workspace_id: str = Field(min_length=1)
+    revision: int = Field(ge=1)
+    step_id: str = Field(min_length=1)
+    output_field: str = Field(min_length=1)
+    state_path: str = Field(min_length=1)
+
+
 class ValidateDraftWorkspaceParams(RpcParamsModel):
     workspace_id: str = Field(min_length=1)