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,
         *,