feat: add draft state schema projection helper

docs/current_roadmap.md

@@ -63,6 +63,9 @@ clear operator feedback before adding more architecture.
 - Completed: `wf schema` now lists workflow document/component models, emits
   compact JSON outlines for agent discovery, and emits valid self-contained
   JSON Schema with `--verbose`.
+- Completed: `wf draft add-state-from-output` projects capability output
+  property schemas into draft state schemas, preserving `$defs` / `definitions`
+  for schema refs and reducing brittle whole-`state_schema` patches.
 - Keep status read-only; do not mutate registry, auth, config, or stores.
 
 ## Priority 2: Durable Run/Resume Hardening

docs/wf_cli.md

@@ -297,6 +297,7 @@ wf draft set-route concat_ws --revision 2 --step call --outcome ok --to __end__
 wf draft set-input concat_ws --revision 3 --step call --map input.items=items --map input.separator=separator
 wf draft set-output concat_ws --revision 4 --step call --map value=state.value
 wf draft set-input concat_ws --revision 5 --step call --merge --map input.limit=limit
+wf draft add-state-from-output concat_ws --revision 5 --step call --output value --state state.value
 ```
 
 `set-input` maps graph source paths to node-local input fields:
@@ -310,6 +311,13 @@ Use repeated `--map` flags in one command when you know the complete map. Use
 `--merge` when adding or updating one entry across a later revision while
 preserving existing bindings.
 
+Before mapping a step output to a new state field, the state schema must declare
+that root field. `add-state-from-output` copies the selected step capability's
+top-level output property schema into `state_schema.properties`, including local
+`$defs` / `definitions` blocks needed by `$ref` schemas. It only declares the
+state field; still run `set-output` or `draft patch` to write values into that
+field, then run `wf draft validate`.
+
 Validate:
 
 ```bash

skills/wf-cli/SKILL.md

@@ -44,6 +44,7 @@ wf draft set-input <workspace_id> --revision <n> --step <step_id> --map input.te
 wf draft set-input <workspace_id> --revision <n> --step <step_id> --merge --map input.other=other
 wf draft set-output <workspace_id> --revision <n> --step <step_id> --map text=state.text
 wf draft set-output <workspace_id> --revision <n> --step <step_id> --merge --map other=state.other
+wf draft add-state-from-output <workspace_id> --revision <n> --step <step_id> --output <field> --state state.<field>
 wf draft validate <workspace_id>
 wf draft save <workspace_id> --artifact <artifact_id> --version <n> --title <title>
 
@@ -73,6 +74,10 @@ For `draft set-input` and `draft set-output`, repeated `--map` flags in one
 command define the complete replacement map. If you split map edits across
 multiple commands, pass `--merge` or the later command replaces the earlier map.
 
+If mapping `LOCAL_SOURCE=state.new_field`, declare the state field first with
+`draft add-state-from-output` when the schema should match a capability output
+field. Do not hand-copy `$defs` unless the helper cannot express the shape.
+
 ## Rules
 
 - Use explicit `--config <path>` for examples, challenge workspaces, and

skills/wf-workflow/references/draft-workspaces.md

@@ -73,6 +73,7 @@ Prefer focused helpers over JSON Patch for common edits:
 - `set_draft_route`
 - `set_step_input_map`
 - `set_step_output_map`
+- `add_state_schema_from_output`
 
 CLI equivalents:
 
@@ -83,6 +84,7 @@ wf draft set-input <workspace_id> --revision <n> --step <step_id> --map input.te
 wf draft set-input <workspace_id> --revision <n> --step <step_id> --merge --map input.other=other
 wf draft set-output <workspace_id> --revision <n> --step <step_id> --map text=state.text
 wf draft set-output <workspace_id> --revision <n> --step <step_id> --merge --map other=state.other
+wf draft add-state-from-output <workspace_id> --revision <n> --step <step_id> --output <field> --state state.<field>
 ```
 
 `set-input` direction: `input.text=text` means graph source `input.text` maps to
@@ -95,6 +97,10 @@ Without `--merge`, `set-input` and `set-output` replace the whole map for that
 step. Use repeated `--map` flags in one command for a complete replacement. Use
 `--merge` only when adding/updating entries over multiple revisions.
 
+Use `add-state-from-output` when the target state field should reuse a capability
+output schema. This prevents dangling `$ref` values by copying local `$defs` /
+`definitions` with the selected property schema.
+
 Use JSON Patch for structural edits the helpers do not cover.
 
 For larger patches, write a JSON Patch array to a file and pass it with

skills/wf-workflow/references/workflow-lifecycle.md

@@ -18,6 +18,8 @@ validated, runnable deployment.
      for common edits.
    - `set-input` and `set-output` replace full maps by default; pass `--merge`
      only when adding or updating one entry across a later revision.
+   - Before output-mapping into a new state field, declare it with
+     `add-state-from-output` when it should mirror a capability output property.
    - Use JSON Patch only for general structural edits.
 6. Save an artifact.
    - Draft artifact:

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)
 

tests/wf_api/test_drafts_service.py

@@ -4,10 +4,12 @@ from pathlib import Path
 from typing import Any
 
 import pytest
+from pydantic import BaseModel
 
 from tests.wf_mcp.test_support import echo_tool
 from wf_api.drafts import WorkflowDraftApi
 from wf_artifacts import FileDraftWorkspaceStore, FileWorkflowArtifactStore
+from wf_authoring import node
 from wf_mcp.broker import WfMcpService
 from wf_mcp.broker.service.workflow_operation_context import context_from_service
 from wf_mcp.models import ConnectionConfig
@@ -51,6 +53,23 @@ def _echo_draft() -> dict[str, Any]:
     }
 
 
+class _Snapshot(BaseModel):
+    clicked: bool
+
+
+class _SnapshotOutput(BaseModel):
+    after: _Snapshot
+
+
+class _SnapshotInput(BaseModel):
+    pass
+
+
+@node(name="snapshot_tool")
+def _snapshot_tool(payload: _SnapshotInput) -> _SnapshotOutput:
+    return _SnapshotOutput(after=_Snapshot(clicked=True))
+
+
 def _draft_api(
     artifact_store: FileWorkflowArtifactStore,
     *,
@@ -375,3 +394,92 @@ async def test_delegation_smoke_validate_draft_equivalence(tmp_path: Path) -> No
     assert (
         handler_result["compiled_plan"]["nodes"] == api_result["compiled_plan"]["nodes"]
     )
+
+
+@pytest.mark.asyncio
+async def test_add_state_schema_from_output_copies_output_property_defs(
+    tmp_path: Path,
+) -> None:
+    artifact_store = FileWorkflowArtifactStore(tmp_path / "drafts_state_from_output")
+    api, service = _draft_api(artifact_store)
+    service.register_connection(
+        ConnectionConfig(id="demo.personal", server="demo", account="personal")
+    )
+    service.register_specs("demo.personal", _snapshot_tool)
+    await api.create_draft_workspace(
+        workspace_id="snapshot_ws",
+        draft={
+            "name": "snapshot",
+            "input_schema": {"type": "object", "properties": {}},
+            "state_schema": {"type": "object", "properties": {}},
+            "output_schema": {"type": "object", "properties": {}},
+            "start": "snap",
+            "steps": {
+                "snap": {
+                    "use": "demo.personal.snapshot_tool",
+                    "input": [],
+                    "output": [],
+                }
+            },
+            "routes": {"snap": {"ok": "__end__"}},
+        },
+    )
+
+    updated = await api.add_state_schema_from_output(
+        workspace_id="snapshot_ws",
+        revision=1,
+        step_id="snap",
+        output_field="after",
+        state_path="state.after",
+    )
+    fetched = await api.get_draft_workspace(
+        workspace_id="snapshot_ws",
+        include_draft=True,
+    )
+
+    state_schema = fetched["draft"]["state_schema"]
+    assert updated["revision"] == 2
+    assert state_schema["properties"]["after"]["$ref"] == "#/$defs/_Snapshot"
+    assert state_schema["$defs"]["_Snapshot"]["properties"]["clicked"] == {
+        "title": "Clicked",
+        "type": "boolean",
+    }
+
+
+@pytest.mark.asyncio
+async def test_add_state_schema_from_output_rejects_nested_state_path(
+    tmp_path: Path,
+) -> None:
+    artifact_store = FileWorkflowArtifactStore(tmp_path / "drafts_nested_state_output")
+    api, service = _draft_api(artifact_store)
+    service.register_connection(
+        ConnectionConfig(id="demo.personal", server="demo", account="personal")
+    )
+    service.register_specs("demo.personal", _snapshot_tool)
+    await api.create_draft_workspace(
+        workspace_id="snapshot_ws",
+        draft={
+            "name": "snapshot",
+            "input_schema": {"type": "object", "properties": {}},
+            "state_schema": {"type": "object", "properties": {}},
+            "output_schema": {"type": "object", "properties": {}},
+            "start": "snap",
+            "steps": {
+                "snap": {
+                    "use": "demo.personal.snapshot_tool",
+                    "input": [],
+                    "output": [],
+                }
+            },
+            "routes": {"snap": {"ok": "__end__"}},
+        },
+    )
+
+    with pytest.raises(ValueError, match="state_path must name one root field"):
+        await api.add_state_schema_from_output(
+            workspace_id="snapshot_ws",
+            revision=1,
+            step_id="snap",
+            output_field="after",
+            state_path="state.after.clicked",
+        )

tests/wf_api/test_schema_projection.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+import pytest
+
+from wf_api.schema_projection import project_output_property_to_state_schema
+
+
+def test_project_output_property_copies_schema_and_defs() -> None:
+    state_schema = {
+        "type": "object",
+        "properties": {"before": {"type": "object"}},
+    }
+    output_schema = {
+        "type": "object",
+        "properties": {
+            "after": {"$ref": "#/$defs/Snapshot"},
+        },
+        "$defs": {
+            "Snapshot": {
+                "type": "object",
+                "properties": {"clicked": {"type": "boolean"}},
+                "required": ["clicked"],
+            }
+        },
+    }
+
+    projected = project_output_property_to_state_schema(
+        state_schema=state_schema,
+        output_schema=output_schema,
+        output_field="after",
+        state_field="after",
+    )
+
+    assert projected["properties"]["before"] == {"type": "object"}
+    assert projected["properties"]["after"] == {"$ref": "#/$defs/Snapshot"}
+    assert projected["$defs"]["Snapshot"]["properties"]["clicked"] == {
+        "type": "boolean"
+    }
+    assert "after" not in state_schema["properties"]
+
+
+def test_project_output_property_rejects_missing_output_field() -> None:
+    with pytest.raises(ValueError, match="output field 'after'"):
+        project_output_property_to_state_schema(
+            state_schema={"type": "object", "properties": {}},
+            output_schema={"type": "object", "properties": {}},
+            output_field="after",
+            state_field="after",
+        )
+
+
+def test_project_output_property_rejects_conflicting_defs() -> None:
+    with pytest.raises(ValueError, match=r"conflicting \$defs.Snapshot"):
+        project_output_property_to_state_schema(
+            state_schema={
+                "type": "object",
+                "properties": {},
+                "$defs": {"Snapshot": {"type": "string"}},
+            },
+            output_schema={
+                "type": "object",
+                "properties": {"after": {"$ref": "#/$defs/Snapshot"}},
+                "$defs": {"Snapshot": {"type": "object"}},
+            },
+            output_field="after",
+            state_field="after",
+        )
+
+
+def test_project_output_property_rejects_invalid_output_schema() -> None:
+    with pytest.raises(ValueError, match="output_schema is not valid JSON Schema"):
+        project_output_property_to_state_schema(
+            state_schema={"type": "object", "properties": {}},
+            output_schema={
+                "type": "object",
+                "properties": {"after": {"type": "definitely-not-jsonschema"}},
+            },
+            output_field="after",
+            state_field="after",
+        )

tests/wf_cli/test_app.py

@@ -142,3 +142,13 @@ def test_wf_draft_map_help_explains_replace_merge_and_validate() -> None:
     assert "replaces the full output map" in output_help
     assert "Use --merge only" in output_help
     assert "draft validate" in output_help
+
+
+def test_wf_draft_add_state_from_output_help_explains_schema_copy() -> None:
+    result = runner.invoke(app, ["draft", "add-state-from-output", "--help"])
+
+    assert result.exit_code == 0
+    help_text = " ".join(result.output.split())
+    assert "capability output field schema" in help_text
+    assert "$defs" in help_text
+    assert "draft validate" in help_text

tests/wf_cli/test_remote_target.py

@@ -1059,6 +1059,23 @@ def test_wf_draft_focused_edit_commands_use_rpc_target(monkeypatch, tmp_path) ->
             "--merge",
         ],
     )
+    state_added = runner.invoke(
+        app,
+        [
+            *base_args,
+            "draft",
+            "add-state-from-output",
+            "focused_ws",
+            "--revision",
+            "7",
+            "--step",
+            "call",
+            "--output",
+            "value",
+            "--state",
+            "state.extra_value",
+        ],
+    )
     inspected = runner.invoke(
         app,
         [*base_args, "draft", "inspect", "focused_ws", "--include-draft"],
@@ -1070,6 +1087,7 @@ def test_wf_draft_focused_edit_commands_use_rpc_target(monkeypatch, tmp_path) ->
     assert output_mapped.exit_code == 0, output_mapped.output
     assert input_merged.exit_code == 0, input_merged.output
     assert output_merged.exit_code == 0, output_merged.output
+    assert state_added.exit_code == 0, state_added.output
     assert inspected.exit_code == 0, inspected.output
     payload = json.loads(inspected.output)
     draft = payload["draft"]
@@ -1095,6 +1113,7 @@ def test_wf_draft_focused_edit_commands_use_rpc_target(monkeypatch, tmp_path) ->
             "target": {"root": "state", "parts": ["extra"]},
         },
     ]
+    assert "extra_value" in draft["state_schema"]["properties"]
 
 
 def test_wf_deploy_create_alias_saves_deployment(monkeypatch, tmp_path) -> None:

tests/wf_mcp/server/test_config.py

@@ -56,6 +56,7 @@ def test_server_exposes_upstream_admin_and_workflow_tools() -> None:
             assert "wf.workflow.set_draft_route" in names
             assert "wf.workflow.set_step_input_map" in names
             assert "wf.workflow.set_step_output_map" in names
+            assert "wf.workflow.add_state_from_output" in names
             assert "wf.workflow.create_minimal_draft_workspace" in names
             assert "wf.workflow.create_draft_workspace_from_capability" in names
             assert "wf.workflow.create_artifact_from_workspace" in names
@@ -112,6 +113,12 @@ def test_server_exposes_upstream_admin_and_workflow_tools() -> None:
             ].inputSchema
             set_input_request = set_input_schema["properties"]["request"]
             assert "merge" in set_input_request["properties"]
+            state_from_output_schema = tools_by_name[
+                "wf.workflow.add_state_from_output"
+            ].inputSchema
+            state_request = state_from_output_schema["properties"]["request"]
+            assert "output_field" in state_request["properties"]
+            assert "state_path" in state_request["properties"]
             from_capability_output = tools_by_name[
                 "wf.workflow.create_draft_workspace_from_capability"
             ].outputSchema

tests/wf_transport_rpc_http/test_app.py

@@ -722,6 +722,17 @@ async def test_rpc_draft_workspace_focused_edit_methods(tmp_path) -> None:
                 "merge": True,
             },
         )
+        state_added = await _rpc(
+            client,
+            "workflow.draft_workspaces.add_state_from_output",
+            {
+                "workspace_id": "focused_ws",
+                "revision": 7,
+                "step_id": "call",
+                "output_field": "value",
+                "state_path": "state.extra_value",
+            },
+        )
         fetched = await _rpc(
             client,
             "workflow.draft_workspaces.get",
@@ -734,6 +745,7 @@ async def test_rpc_draft_workspace_focused_edit_methods(tmp_path) -> None:
     assert output_mapped["result"]["revision"] == 5
     assert input_merged["result"]["revision"] == 6
     assert output_merged["result"]["revision"] == 7
+    assert state_added["result"]["revision"] == 8
     draft = fetched["result"]["draft"]
     assert draft["name"] == "focused_renamed"
     assert draft["routes"]["call"]["ok"] == "__end__"
@@ -757,6 +769,10 @@ async def test_rpc_draft_workspace_focused_edit_methods(tmp_path) -> None:
             "target": {"root": "state", "parts": ["extra"]},
         },
     ]
+    assert (
+        draft["state_schema"]["properties"]["extra_value"]
+        == draft["state_schema"]["properties"]["value"]
+    )
 
 
 async def test_rpc_diagnoses_source(tmp_path) -> None:

tests/wf_transport_rpc_http/test_client.py

@@ -501,6 +501,13 @@ async def test_rpc_client_draft_workspace_focused_edit_methods(tmp_path) -> None
             output_map={"extra": "state.extra"},
             merge=True,
         )
+        state_added = await client.add_state_schema_from_output(
+            workspace_id="client_focused_ws",
+            revision=7,
+            step_id="call",
+            output_field="value",
+            state_path="state.extra_value",
+        )
 
     assert named["revision"] == 2
     assert routed["revision"] == 3
@@ -508,6 +515,7 @@ async def test_rpc_client_draft_workspace_focused_edit_methods(tmp_path) -> None
     assert output_mapped["revision"] == 5
     assert input_merged["revision"] == 6
     assert output_merged["revision"] == 7
+    assert state_added["revision"] == 8
 
 
 async def test_rpc_client_diagnoses_source(tmp_path) -> None: