more schema validation tests

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