fix bugs, new type

src/wf_core/models/schemas.py

@@ -262,6 +262,7 @@ def _iter_state_field_declarations(
             for key, value in resolved_schema.items()
             if key not in {"reducer", "trace"}
         }
+        _attach_root_schema_context(validation_schema, root_schema)
         yield (
             path,
             StateFieldDecl.model_validate(
@@ -392,6 +393,24 @@ def _resolve_local_ref(
     return resolved if isinstance(resolved, Mapping) else property_schema
 
 
+def _attach_root_schema_context(
+    field_schema: dict[str, Any],
+    root_schema: Mapping[str, Any],
+) -> None:
+    """Keep local JSON Schema refs valid after extracting one state field.
+
+    Runtime state writes validate one declared state path at a time. If a field
+    schema contains a Pydantic-style local ref such as ``#/$defs/Thing``, the
+    extracted subschema still needs the root ``$defs`` table to resolve it.
+    """
+    definitions = root_schema.get("$defs")
+    if isinstance(definitions, Mapping):
+        field_schema.setdefault("$defs", dict(definitions))
+    schema_dialect = root_schema.get("$schema")
+    if isinstance(schema_dialect, str):
+        field_schema.setdefault("$schema", schema_dialect)
+
+
 class NodeDef(BaseModel):
     """Reusable node contract referenced by one or more node uses."""
 

src/wf_mcp/workflow_surface/handlers.py

@@ -42,6 +42,7 @@ from .constants import (
     DEFAULT_OK_OUTCOME,
     RUNTIME_ERROR_CAPABILITY,
 )
+from .refs import WorkflowSurfaceCapabilityId
 
 if TYPE_CHECKING:
     from ..broker.service import WfMcpService
@@ -159,12 +160,22 @@ class WorkflowSurfaceHandlers:
         qualified_name: str,
     ) -> WorkflowArtifact | None:
         """Resolve a saved node-like wrapper artifact from its stable capability name."""
-        parsed = _parse_artifact_capability_id(qualified_name)
-        if parsed is None or self.service.artifact_store is None:
-            return None
-        artifact_id, version = parsed
         try:
-            artifact = self.service.artifact_store.get_artifact(artifact_id, version)
+            capability_id = WorkflowSurfaceCapabilityId.parse(qualified_name)
+        except ValueError:
+            return None
+        if (
+            not capability_id.is_wrapper_artifact
+            or self.service.artifact_store is None
+            or capability_id.artifact_id is None
+            or capability_id.artifact_version is None
+        ):
+            return None
+        try:
+            artifact = self.service.artifact_store.get_artifact(
+                capability_id.artifact_id,
+                capability_id.artifact_version,
+            )
         except KeyError:
             return None
         if artifact.kind != "wrapper":
@@ -984,9 +995,12 @@ def _source_id_for_capability(
 def _capability_name(qualified_name: str) -> str | None:
     """Return the local name of one qualified capability ref if it is valid."""
     try:
-        return CapabilityRef.parse(qualified_name).name
+        parsed = WorkflowSurfaceCapabilityId.parse(qualified_name)
     except ValueError:
         return None
+    if parsed.is_wrapper_artifact:
+        return None
+    return parsed.live_name
 
 
 def _artifact_capability_id(artifact: WorkflowArtifact) -> str:
@@ -999,15 +1013,6 @@ def _artifact_capability_id(artifact: WorkflowArtifact) -> str:
     )
 
 
-def _parse_artifact_capability_id(qualified_name: str) -> tuple[str, int] | None:
-    """Parse the stable `workflow.<artifact_id>.v<version>` capability name."""
-    try:
-        ref = WorkflowCapabilityRef.parse(qualified_name)
-    except ValueError:
-        return None
-    return ref.artifact_id, ref.version
-
-
 def _raw_plan_from_artifact(artifact: WorkflowArtifact) -> RawWorkflowPlan:
     """Validate the stored plan shape expected by the broker workflow runner."""
     return RawWorkflowPlan.model_validate(

src/wf_mcp/workflow_surface/refs.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from wf_artifacts import WorkflowCapabilityRef
+from wf_platform import CapabilityRef
+
+
+@dataclass(frozen=True, slots=True)
+class WorkflowSurfaceCapabilityId:
+    """Typed internal form for workflow-surface capability names.
+
+    MCP tools still accept and return plain strings. This type is an internal
+    boundary so handlers can distinguish live source capabilities from saved
+    wrapper artifacts without repeating ad hoc string parsing.
+    """
+
+    qualified_name: str
+    source_id: str
+    live_name: str | None = None
+    artifact_id: str | None = None
+    artifact_version: int | None = None
+
+    @classmethod
+    def parse(cls, value: str) -> WorkflowSurfaceCapabilityId:
+        """Parse one workflow-facing capability id into its internal kind."""
+        try:
+            artifact_ref = WorkflowCapabilityRef.parse(value)
+        except ValueError:
+            capability_ref = CapabilityRef.parse(value)
+            return cls(
+                qualified_name=str(capability_ref),
+                source_id=str(capability_ref.source),
+                live_name=capability_ref.name,
+            )
+        return cls(
+            qualified_name=str(artifact_ref),
+            source_id="workflow",
+            artifact_id=artifact_ref.artifact_id,
+            artifact_version=artifact_ref.version,
+        )
+
+    @property
+    def is_wrapper_artifact(self) -> bool:
+        """Return whether this id targets a saved workflow wrapper artifact."""
+        return self.artifact_id is not None

tests/authoring/test_ops.py

@@ -29,7 +29,6 @@ from wf_core import (
     RunStatus,
     RuntimeContext,
     SchemaRef,
-    StateField,
     StateSchema,
     execute_workflow,
 )
@@ -40,10 +39,13 @@ def _build_first_workflow(use_safe_first: bool = False):
     builder = WorkflowBuilder(
         name="first_demo",
         input_schema=SchemaRef(type="object"),
-        state_schema=StateSchema.from_field_map(
+        state_schema=StateSchema.model_validate(
             {
-                "items": StateField(type="array"),
-                "item": StateField(type="object"),
+                "type": "object",
+                "properties": {
+                    "items": {"type": "array"},
+                    "item": {"type": ["string", "null"]},
+                },
             }
         ),
         output_schema=SchemaRef(type="object"),
@@ -63,11 +65,14 @@ def _build_first_maybe_workflow():
     builder = WorkflowBuilder(
         name="first_maybe_demo",
         input_schema=SchemaRef(type="object"),
-        state_schema=StateSchema.from_field_map(
+        state_schema=StateSchema.model_validate(
             {
-                "items": StateField(type="array"),
-                "item": StateField(type="object"),
-                "missing": StateField(type="boolean"),
+                "type": "object",
+                "properties": {
+                    "items": {"type": "array"},
+                    "item": {"type": ["string", "null"]},
+                    "missing": {"type": "boolean"},
+                },
             }
         ),
         output_schema=SchemaRef(type="object"),

tests/core/test_schema_validation.py

@@ -169,3 +169,34 @@ def test_state_schema_dump_is_valid_json_schema_with_reducer_keyword() -> None:
     assert dumped["properties"]["count"]["description"] == "Running count"
     assert dumped["properties"]["count"]["reducer"] == "wf.std.add"
     Draft202012Validator.check_schema(dumped)
+
+
+def test_state_field_validation_schema_preserves_root_defs_for_local_refs() -> None:
+    from wf_core import StateSchema
+
+    schema = StateSchema.model_validate(
+        {
+            "type": "object",
+            "$defs": {
+                "PoolByCategory": {
+                    "type": "object",
+                    "properties": {"category": {"type": "string"}},
+                    "required": ["category"],
+                }
+            },
+            "properties": {
+                "current_pools": {
+                    "type": "array",
+                    "items": {"$ref": "#/$defs/PoolByCategory"},
+                }
+            },
+        }
+    )
+
+    field_schema = schema.field_map()["current_pools"].validation_schema
+
+    validate_payload_against_schema(
+        field_schema,
+        [{"category": "1"}],
+        "state write state.current_pools",
+    )

tests/wf_mcp/test_workflow_surface_refs.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+from wf_mcp.workflow_surface.refs import WorkflowSurfaceCapabilityId
+
+
+def test_workflow_surface_capability_id_parses_live_capability_ref() -> None:
+    capability = WorkflowSurfaceCapabilityId.parse("demo.personal.echo_tool")
+
+    assert capability.qualified_name == "demo.personal.echo_tool"
+    assert capability.source_id == "demo.personal"
+    assert capability.live_name == "echo_tool"
+    assert capability.is_wrapper_artifact is False
+
+
+def test_workflow_surface_capability_id_parses_saved_wrapper_ref() -> None:
+    capability = WorkflowSurfaceCapabilityId.parse("workflow.echo_wrapper.v2")
+
+    assert capability.qualified_name == "workflow.echo_wrapper.v2"
+    assert capability.source_id == "workflow"
+    assert capability.artifact_id == "echo_wrapper"
+    assert capability.artifact_version == 2
+    assert capability.is_wrapper_artifact is True