more schema validation tests

src/wf_authoring/catalog.py

@@ -17,13 +17,19 @@ class NodeCatalogEntry:
 
     @classmethod
     def from_spec(cls, spec: NodeSpec[Any, Any]) -> "NodeCatalogEntry":
+        input_schema = (
+            spec.input_schema_contract or spec.input_model.model_json_schema()
+        )
+        output_schema = (
+            spec.output_schema_contract or spec.output_model.model_json_schema()
+        )
         return cls(
             name=spec.name,
             display_name=None,
             description=spec.description,
             outcomes=spec.outcomes,
-            input_schema=spec.input_model.model_json_schema(),
-            output_schema=spec.output_model.model_json_schema(),
+            input_schema=input_schema,
+            output_schema=output_schema,
         )
 
 

src/wf_authoring/nodes/schema.py

@@ -1,10 +1,17 @@
 from __future__ import annotations
 
+from typing import Any
+
 from pydantic import BaseModel
 
 from wf_core import SchemaRef
 
 
-def schema_ref_for(model_type: type[BaseModel]) -> SchemaRef:
-    """Build a core schema reference from a pydantic model class."""
+def schema_ref_for(
+    model_type: type[BaseModel],
+    schema_override: dict[str, Any] | None = None,
+) -> SchemaRef:
+    """Build a core schema reference from a Pydantic model or schema override."""
+    if schema_override is not None:
+        return SchemaRef.model_validate(schema_override)
     return SchemaRef.model_validate(model_type.model_json_schema())

src/wf_authoring/nodes/spec.py

@@ -60,6 +60,8 @@ class NodeSpec(Generic[InputT, OutputT]):
     description: str | None = None
     is_async: bool = False
     accepts_context: bool = True
+    input_schema_contract: dict[str, Any] | None = None
+    output_schema_contract: dict[str, Any] | None = None
 
     def __call__(
         self,
@@ -75,8 +77,14 @@ class NodeSpec(Generic[InputT, OutputT]):
     def to_node_def(self) -> NodeDef:
         return NodeDef(
             name=self.name,
-            input_schema=schema_ref_for(self.input_model),
-            output_schema=schema_ref_for(self.output_model),
+            input_schema=schema_ref_for(
+                self.input_model,
+                self.input_schema_contract,
+            ),
+            output_schema=schema_ref_for(
+                self.output_model,
+                self.output_schema_contract,
+            ),
             outcomes=list(self.outcomes),
         )