more schema validation tests

This commit is contained in:
lda
2026-05-08 23:42:10 +07:00 Verified
parent ca14d726d0
commit 4d8e1e29ee
10 changed files with 257 additions and 26 deletions
+8 -2
View File
@@ -17,13 +17,19 @@ class NodeCatalogEntry:
@classmethod @classmethod
def from_spec(cls, spec: NodeSpec[Any, Any]) -> "NodeCatalogEntry": 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( return cls(
name=spec.name, name=spec.name,
display_name=None, display_name=None,
description=spec.description, description=spec.description,
outcomes=spec.outcomes, outcomes=spec.outcomes,
input_schema=spec.input_model.model_json_schema(), input_schema=input_schema,
output_schema=spec.output_model.model_json_schema(), output_schema=output_schema,
) )
+9 -2
View File
@@ -1,10 +1,17 @@
from __future__ import annotations from __future__ import annotations
from typing import Any
from pydantic import BaseModel from pydantic import BaseModel
from wf_core import SchemaRef from wf_core import SchemaRef
def schema_ref_for(model_type: type[BaseModel]) -> SchemaRef: def schema_ref_for(
"""Build a core schema reference from a pydantic model class.""" 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()) return SchemaRef.model_validate(model_type.model_json_schema())
+10 -2
View File
@@ -60,6 +60,8 @@ class NodeSpec(Generic[InputT, OutputT]):
description: str | None = None description: str | None = None
is_async: bool = False is_async: bool = False
accepts_context: bool = True accepts_context: bool = True
input_schema_contract: dict[str, Any] | None = None
output_schema_contract: dict[str, Any] | None = None
def __call__( def __call__(
self, self,
@@ -75,8 +77,14 @@ class NodeSpec(Generic[InputT, OutputT]):
def to_node_def(self) -> NodeDef: def to_node_def(self) -> NodeDef:
return NodeDef( return NodeDef(
name=self.name, name=self.name,
input_schema=schema_ref_for(self.input_model), input_schema=schema_ref_for(
output_schema=schema_ref_for(self.output_model), self.input_model,
self.input_schema_contract,
),
output_schema=schema_ref_for(
self.output_model,
self.output_schema_contract,
),
outcomes=list(self.outcomes), outcomes=list(self.outcomes),
) )
+3
View File
@@ -17,6 +17,9 @@ def qualify_spec(connection_id: str, spec: NodeSpec[Any, Any]) -> NodeSpec[Any,
fn=spec.fn, fn=spec.fn,
description=spec.description, description=spec.description,
is_async=spec.is_async, is_async=spec.is_async,
accepts_context=spec.accepts_context,
input_schema_contract=spec.input_schema_contract,
output_schema_contract=spec.output_schema_contract,
) )
+12 -2
View File
@@ -4,10 +4,10 @@ from dataclasses import dataclass, field
from typing import Any from typing import Any
## are these our own redef of MCP VERY MUCH NICE VERY MUCH READY structs?
## These are unfortunately boundary. Hence, we need good typecheck on these, and since mcp lib is good stuff, carry those over. Could be pro
@dataclass(slots=True) @dataclass(slots=True)
class DiscoveredTool: class DiscoveredTool:
"""Tool snapshot after converting from an upstream MCP SDK tool."""
name: str name: str
title: str | None title: str | None
description: str | None description: str | None
@@ -19,6 +19,8 @@ class DiscoveredTool:
@dataclass(slots=True) @dataclass(slots=True)
class DiscoveredResource: class DiscoveredResource:
"""Resource snapshot after converting from an upstream MCP SDK resource."""
uri: str uri: str
name: str name: str
title: str | None title: str | None
@@ -29,6 +31,8 @@ class DiscoveredResource:
@dataclass(slots=True) @dataclass(slots=True)
class DiscoveredPrompt: class DiscoveredPrompt:
"""Prompt snapshot after converting from an upstream MCP SDK prompt."""
name: str name: str
title: str | None title: str | None
description: str | None description: str | None
@@ -38,6 +42,8 @@ class DiscoveredPrompt:
@dataclass(slots=True) @dataclass(slots=True)
class CatalogNodeEntry: class CatalogNodeEntry:
"""Namespaced tool entry stored in the broker catalog snapshot."""
qualified_name: str qualified_name: str
connection_id: str connection_id: str
local_name: str local_name: str
@@ -50,6 +56,8 @@ class CatalogNodeEntry:
@dataclass(slots=True) @dataclass(slots=True)
class CatalogResourceEntry: class CatalogResourceEntry:
"""Namespaced resource entry stored in the broker catalog snapshot."""
qualified_name: str qualified_name: str
connection_id: str connection_id: str
local_name: str local_name: str
@@ -62,6 +70,8 @@ class CatalogResourceEntry:
@dataclass(slots=True) @dataclass(slots=True)
class CatalogPromptEntry: class CatalogPromptEntry:
"""Namespaced prompt entry stored in the broker catalog snapshot."""
qualified_name: str qualified_name: str
connection_id: str connection_id: str
local_name: str local_name: str
+74 -5
View File
@@ -1,27 +1,94 @@
from __future__ import annotations from __future__ import annotations
from collections.abc import Callable from collections.abc import Callable
from typing import Any, cast from types import NoneType, UnionType
from typing import Any, Union, cast, get_args, get_origin
from pydantic import BaseModel, ConfigDict, Field, create_model from pydantic import BaseModel, ConfigDict, Field, create_model
from wf_authoring import NodeReturn, NodeSpec from wf_authoring import NodeReturn, NodeSpec
from wf_core import RuntimeContext from wf_core import RuntimeContext
from wf_mcp.broker.events import McpEvent, make_event
from ..capabilities import DiscoveredTool from ..capabilities import DiscoveredTool
from ..models import AuthRecord, ConnectionConfig from ..models import AuthRecord, ConnectionConfig
from ..sdk import BackendAdapter from ..sdk import BackendAdapter
from wf_mcp.broker.events import McpEvent, make_event
_JSON_TYPE_MAP: dict[str, object] = {
"string": str,
"integer": int,
"number": float,
"boolean": bool,
"object": dict[str, Any],
}
def _python_type_from_schema(schema: object) -> object:
"""Map a small JSON Schema subset into a Pydantic field annotation."""
if not isinstance(schema, dict):
return Any
if "enum" in schema:
return Any
schema_type = schema.get("type")
if isinstance(schema_type, list):
non_null_types = [item for item in schema_type if item != "null"]
if len(non_null_types) == 1:
return _optional_type(
_python_type_from_schema({**schema, "type": non_null_types[0]})
)
return Any
if schema_type == "array":
item_type = _python_type_from_schema(schema.get("items", {}))
return list[item_type] if isinstance(item_type, type) else list[Any]
if not isinstance(schema_type, str):
return Any
return _JSON_TYPE_MAP.get(schema_type, Any)
def _optional_type(annotation: object) -> object:
"""Return an optional version of a supported runtime annotation."""
if annotation is Any:
return Any
origin = get_origin(annotation)
if origin in {Union, UnionType} and NoneType in get_args(annotation):
return annotation
return annotation | None if isinstance(annotation, type) else Any
def _field_default(
field_name: str,
property_schema: object,
required: set[str],
) -> object:
"""Return the Pydantic field default for a JSON Schema property."""
if isinstance(property_schema, dict) and "default" in property_schema:
return property_schema["default"]
return ... if field_name in required else None
def _model_from_schema(name: str, schema: dict[str, Any]) -> type[BaseModel]: def _model_from_schema(name: str, schema: dict[str, Any]) -> type[BaseModel]:
"""Create a loose Pydantic adapter model for an MCP JSON Schema object."""
properties = cast(dict[str, Any], schema.get("properties", {})) properties = cast(dict[str, Any], schema.get("properties", {}))
required = set(cast(list[str], schema.get("required", []))) required = set(cast(list[str], schema.get("required", [])))
field_defs: dict[str, tuple[object, object]] = {} field_defs: dict[str, tuple[object, object]] = {}
for field_name in properties: for field_name, property_schema in properties.items():
default = ... if field_name in required else None annotation = _python_type_from_schema(property_schema)
field_defs[field_name] = (Any, Field(default=default)) default = _field_default(field_name, property_schema, required)
description = (
property_schema.get("description")
if isinstance(property_schema, dict)
else None
)
field_defs[field_name] = (
annotation,
Field(default=default, description=description),
)
raw_field_defs = cast(dict[str, Any], field_defs) raw_field_defs = cast(dict[str, Any], field_defs)
model = create_model( model = create_model(
@@ -93,4 +160,6 @@ def wrap_discovered_tool(
fn=invoke_tool, fn=invoke_tool,
description=tool.description, description=tool.description,
is_async=True, is_async=True,
input_schema_contract=tool.input_schema,
output_schema_contract=tool.output_schema,
) )
+39
View File
@@ -1,5 +1,6 @@
from __future__ import annotations from __future__ import annotations
import pytest
from pydantic import BaseModel from pydantic import BaseModel
from wf_core import ( from wf_core import (
@@ -7,6 +8,7 @@ from wf_core import (
FrameStatus, FrameStatus,
RuntimeContext, RuntimeContext,
RunStatus, RunStatus,
WorkflowExecutionError,
execute_workflow, execute_workflow,
resume_workflow, resume_workflow,
step_workflow, step_workflow,
@@ -286,6 +288,43 @@ def test_non_interrupt_path_skips_email() -> None:
assert run.interrupt is None assert run.interrupt is None
def test_workflow_input_schema_rejects_wrong_type() -> None:
workflow = build_demo_workflow()
registry = build_demo_registry()
with pytest.raises(
WorkflowExecutionError,
match=r"workflow input\['should_email'\].*not of type 'boolean'",
):
execute_workflow(
workflow,
{"folder_id": "demo-folder", "should_email": "false"},
registry,
)
def test_node_output_schema_rejects_wrong_type() -> None:
workflow = build_demo_workflow()
registry = build_demo_registry()
def bad_list_files(
payload: dict[str, object], ctx: RuntimeContext
) -> dict[str, object]:
return {"outcome": "ok", "output": {"documents": "not-a-list"}}
registry["drive_list_files"] = bad_list_files
with pytest.raises(
WorkflowExecutionError,
match=r"node output for list_files\['documents'\].*not of type 'array'",
):
execute_workflow(
workflow,
{"folder_id": "demo-folder", "should_email": False},
registry,
)
def test_stepwise_execution_reaches_interrupt() -> None: def test_stepwise_execution_reaches_interrupt() -> None:
workflow = build_demo_workflow() workflow = build_demo_workflow()
registry = build_demo_registry() registry = build_demo_registry()
+30 -2
View File
@@ -1,8 +1,8 @@
from __future__ import annotations from __future__ import annotations
from pydantic import BaseModel from pydantic import BaseModel, Field
from wf_authoring import NodeReturn, Nothing, build_registry, node, outcome from wf_authoring import NodeCatalog, NodeReturn, Nothing, build_registry, node, outcome
from wf_core import RuntimeContext from wf_core import RuntimeContext
@@ -38,6 +38,14 @@ class InferredAsyncOutput(BaseModel):
echoed: str echoed: str
class DocumentedInput(BaseModel):
message: str = Field(description="Message to echo")
class DocumentedOutput(BaseModel):
echoed: str = Field(description="Echoed message")
@node() @node()
def inferred_echo( def inferred_echo(
payload: InferredEchoInput, payload: InferredEchoInput,
@@ -157,3 +165,23 @@ def test_node_decorator_infers_nodereturn_output_model() -> None:
def test_node_decorator_detects_async_automatically() -> None: def test_node_decorator_detects_async_automatically() -> None:
assert inferred_async_echo.is_async is True assert inferred_async_echo.is_async is True
def test_pydantic_field_descriptions_survive_node_catalog_schema() -> None:
@node(description="Echoes a documented message.")
def documented_echo(payload: DocumentedInput) -> DocumentedOutput:
return DocumentedOutput(echoed=payload.message)
entry = NodeCatalog.from_specs(documented_echo).entries()[0]
assert entry.description == "Echoes a documented message."
assert entry.input_schema["type"] == "object"
assert (
entry.input_schema["properties"]["message"]["description"]
== "Message to echo"
)
assert entry.output_schema["type"] == "object"
assert (
entry.output_schema["properties"]["echoed"]["description"]
== "Echoed message"
)
+57 -6
View File
@@ -117,17 +117,23 @@ def test_service_refreshes_catalog_from_adapter() -> None:
"description": "Echo text back", "description": "Echo text back",
"outcomes": ["ok"], "outcomes": ["ok"],
"input_schema": { "input_schema": {
"additionalProperties": True, "properties": {
"properties": {"text": {"title": "Text"}}, "text": {
"description": "Text to echo",
"type": "string",
}
},
"required": ["text"], "required": ["text"],
"title": "demo.personal_echo_tool_Input",
"type": "object", "type": "object",
}, },
"output_schema": { "output_schema": {
"additionalProperties": True, "properties": {
"properties": {"echoed": {"title": "Echoed"}}, "echoed": {
"description": "Echoed text",
"type": "string",
}
},
"required": ["echoed"], "required": ["echoed"],
"title": "demo.personal_echo_tool_Output",
"type": "object", "type": "object",
}, },
} }
@@ -178,6 +184,51 @@ def test_service_refreshes_catalog_from_adapter() -> None:
assert "catalog_refresh_completed" in event_kinds assert "catalog_refresh_completed" in event_kinds
def test_service_catalog_preserves_json_schema_description_metadata() -> None:
service = WfMcpService(store=FileStore(local_temp_root() / "schema_doc_store"))
service.register_connection(
ConnectionConfig(id="demo.personal", server="demo", account="personal")
)
service.register_adapter("demo", FakeAdapter())
asyncio.run(service.refresh_connection_catalog("demo.personal"))
node = service.get_catalog().as_payload()["nodes"][0]
assert node["input_schema"]["type"] == "object"
assert isinstance(node["input_schema"]["properties"], dict)
assert (
node["input_schema"]["properties"]["text"]["description"]
== "Text to echo"
)
assert node["output_schema"]["type"] == "object"
assert isinstance(node["output_schema"]["properties"], dict)
def test_service_wrapped_tool_adapter_model_validates_simple_schema_types() -> None:
service = WfMcpService(store=FileStore(local_temp_root() / "schema_model_store"))
service.register_connection(
ConnectionConfig(id="demo.personal", server="demo", account="personal")
)
service.register_adapter("demo", FakeAdapter())
asyncio.run(service.refresh_connection_catalog("demo.personal"))
spec = service.specs_by_connection["demo.personal"]["demo.personal.echo_tool"]
parsed = spec.input_model.model_validate({"text": "hello"})
assert parsed.model_dump() == {"text": "hello"}
assert (
spec.input_model.model_json_schema()["properties"]["text"]["description"]
== "Text to echo"
)
try:
spec.input_model.model_validate({"text": 123})
except ValueError as exc:
assert "text" in str(exc)
else:
raise AssertionError("expected generated adapter model to reject non-string text")
def test_service_records_tool_call_events() -> None: def test_service_records_tool_call_events() -> None:
service = WfMcpService(store=FileStore(local_temp_root() / "event_store")) service = WfMcpService(store=FileStore(local_temp_root() / "event_store"))
service.register_connection( service.register_connection(
+15 -5
View File
@@ -5,7 +5,7 @@ import sys
from pathlib import Path from pathlib import Path
from typing import Any from typing import Any
from pydantic import BaseModel from pydantic import BaseModel, Field
from wf_authoring import NodeReturn, node from wf_authoring import NodeReturn, node
from wf_core import RuntimeContext from wf_core import RuntimeContext
@@ -15,11 +15,11 @@ from wf_mcp.sdk import ToolCallResult
class EchoInput(BaseModel): class EchoInput(BaseModel):
text: str text: str = Field(description="Text to echo")
class EchoOutput(BaseModel): class EchoOutput(BaseModel):
echoed: str echoed: str = Field(description="Echoed text")
class FinalizeInput(BaseModel): class FinalizeInput(BaseModel):
@@ -104,12 +104,22 @@ class FakeAdapter:
description="Echo text back", description="Echo text back",
input_schema={ input_schema={
"type": "object", "type": "object",
"properties": {"text": {"type": "string"}}, "properties": {
"text": {
"type": "string",
"description": "Text to echo",
}
},
"required": ["text"], "required": ["text"],
}, },
output_schema={ output_schema={
"type": "object", "type": "object",
"properties": {"echoed": {"type": "string"}}, "properties": {
"echoed": {
"type": "string",
"description": "Echoed text",
}
},
"required": ["echoed"], "required": ["echoed"],
}, },
) )