fix: preserve scalar enum and const constraints in MCP schemas

atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py

@@ -1,15 +1,29 @@
 """Module for transforming JSON schemas to Pydantic models."""
 
 import logging
-from typing import Any, Dict, List, Optional, Type, Tuple, Literal, Union, cast
+from typing import Any, Dict, List, Optional, Type, Tuple, Literal, Union, Annotated, cast
 
 from atomic_agents.connectors.mcp.mcp_definition_service import MCPAttributeType
-from pydantic import Field, create_model
+from jsonschema import Draft202012Validator, validators
+from pydantic import BeforeValidator, Field, create_model
+from referencing import Registry
+from referencing.exceptions import Unresolvable
 
 from atomic_agents.base.base_io_schema import BaseIOSchema
 
 logger = logging.getLogger(__name__)
 
+
+def _validate_ref(validator: Any, reference: str, value: Any, schema: Any) -> Any:
+    try:
+        yield from Draft202012Validator.VALIDATORS["$ref"](validator, reference, value, schema)
+    except Unresolvable:
+        # Match the transformer's Any fallback only for the unresolved reference.
+        return
+
+
+_ScalarChoiceValidator = validators.extend(Draft202012Validator, {"$ref": _validate_ref})
+
 # JSON type mapping
 JSON_TYPE_MAP = {
     "string": str,
@@ -25,6 +39,42 @@
 class SchemaTransformer:
     """Class for transforming JSON schemas to Pydantic models."""
 
+    @staticmethod
+    def _matches_scalar_type(value: Any, json_types: List[str]) -> bool:
+        """Match JSON scalar types without treating Python bools as numbers."""
+        if isinstance(value, bool):
+            return "boolean" in json_types
+        if isinstance(value, (int, float)):
+            return "number" in json_types or ("integer" in json_types and (isinstance(value, int) or value.is_integer()))
+        return (value is None and "null" in json_types) or (isinstance(value, str) and "string" in json_types)
+
+    @staticmethod
+    def _scalar_choice_type(prop_schema: Dict[str, Any], choices: List[Any], root_schema: Dict[str, Any]) -> Type:
+        """Keep only choices satisfying all advertised constraints, including compositions."""
+        declared_type = prop_schema.get("type")
+        json_types = declared_type if isinstance(declared_type, list) else [declared_type]
+        # The finite set can be intersected once, without coercing values through union types.
+        # Bind local references to the original root, and never retrieve remote schemas.
+        validator = _ScalarChoiceValidator(root_schema, registry=Registry()).evolve(schema=prop_schema)
+        choices = [value for value in choices if validator.is_valid(value)]
+        if not choices:
+            raise ValueError("Scalar choices do not match the declared JSON Schema constraints")
+
+        def validate_choice(value: Any) -> Any:
+            if declared_type is not None and not SchemaTransformer._matches_scalar_type(value, json_types):
+                raise ValueError("Value does not match the declared JSON Schema type")
+            for choice in choices:
+                if isinstance(value, bool) == isinstance(choice, bool) and value == choice:
+                    # Give Literal the matching choice so 1.0 cannot resolve to True.
+                    return choice
+            raise ValueError("Value does not match the JSON Schema choices")
+
+        choice_type = Annotated[Literal[tuple(choices)], BeforeValidator(validate_choice)]
+        if declared_type is not None:
+            # Keep type metadata inside the choice branch when Optional wraps it.
+            choice_type = Annotated[choice_type, Field(json_schema_extra={"type": declared_type})]
+        return choice_type
+
     @staticmethod
     def _resolve_ref(ref_path: str, root_schema: Dict[str, Any], model_cache: Dict[str, Type]) -> Type:
         """Resolve a $ref to its Pydantic-compatible value type."""
@@ -100,6 +150,13 @@ def json_to_pydantic_field(
         # Handle $ref
         if "$ref" in prop_schema:
             python_type = SchemaTransformer._resolve_ref(prop_schema["$ref"], root_schema, model_cache)
+        # Scalar value constraints must also appear in the model-facing schema.
+        elif "const" in prop_schema and isinstance(prop_schema["const"], (str, int, float, bool, type(None))):
+            python_type = SchemaTransformer._scalar_choice_type(prop_schema, [prop_schema["const"]], root_schema)
+        elif prop_schema.get("enum") and all(
+            isinstance(value, (str, int, float, bool, type(None))) for value in prop_schema["enum"]
+        ):
+            python_type = SchemaTransformer._scalar_choice_type(prop_schema, prop_schema["enum"], root_schema)
         # Handle oneOf/anyOf (unions)
         elif "oneOf" in prop_schema or "anyOf" in prop_schema:
             union_schemas = prop_schema.get("oneOf", prop_schema.get("anyOf", []))

atomic-agents/tests/connectors/mcp/test_schema_transformer.py

@@ -1,6 +1,7 @@
 import logging
 
 import pytest
+from jsonschema import Draft202012Validator
 from typing import Any, Dict, List, Optional, Union
 
 from pydantic import ValidationError
@@ -75,6 +76,269 @@ def test_no_type(self):
 
 
 class TestCreateModelFromSchema:
+    @pytest.mark.parametrize("required", [True, False])
+    @pytest.mark.parametrize(
+        "prop_schema",
+        [
+            {"enum": [1, "bad"], "anyOf": [{"type": "integer"}]},
+            {"enum": [1, 2.5, "bad"], "oneOf": [{"type": "integer"}, {"type": "number"}]},
+            {"enum": [1, 2, 3, "bad"], "allOf": [{"type": "number"}, {"minimum": 2}]},
+            {"const": 1, "anyOf": [{"type": "integer"}]},
+            {"const": 2.5, "oneOf": [{"type": "integer"}, {"type": "number"}]},
+            {"const": 3, "allOf": [{"type": "number"}, {"minimum": 2}]},
+            {"enum": [True, 1, "bad"], "anyOf": [{"type": "boolean"}, {"type": "integer"}]},
+            {"enum": [1, 2], "const": 2},
+            {"enum": [1, "bad"], "anyOf": [{"$ref": "#/$defs/Integer"}]},
+            {
+                "enum": [0, 1, 2, 3],
+                "anyOf": [{"type": "integer"}, {"type": "string"}],
+                "oneOf": [{"const": 1}, {"const": 2}],
+                "allOf": [{"minimum": 2}],
+            },
+            {"enum": [1, "bad"], "anyOf": [False, {"type": "integer"}]},
+            {"enum": ["brief", None, 1], "anyOf": [{"type": "string"}, {"type": "null"}]},
+            {"const": None, "anyOf": [{"type": "string"}, {"type": "null"}]},
+        ],
+    )
+    def test_scalar_choices_intersect_compositions_in_validation_and_emitted_schema(self, prop_schema, required):
+        schema = {
+            "type": "object",
+            "properties": {"value": prop_schema},
+            "required": ["value"] if required else [],
+            "$defs": {"Integer": {"type": "integer"}},
+        }
+        model = SchemaTransformer.create_model_from_schema(schema, "ComposedChoiceInput", "choose")
+        advertised = Draft202012Validator(schema)
+        emitted = Draft202012Validator(model.model_json_schema())
+        for value in [None, True, False, 0, 1, 1.0, 2, 2.5, 3, "brief", "bad", "unknown", [], {}]:
+            payload = {"tool_name": "choose", "value": value}
+            # Optional fields retain the transformer's existing local None allowance.
+            allowed = advertised.is_valid(payload) or (not required and value is None)
+            assert emitted.is_valid(payload) == allowed
+            if allowed:
+                assert model.model_validate(payload).value == value
+            else:
+                with pytest.raises(ValidationError):
+                    model.model_validate(payload)
+
+    @pytest.mark.parametrize(
+        "prop_schema",
+        [
+            {"const": "bad", "anyOf": [{"type": "integer"}]},
+            {"const": 1, "oneOf": [{"type": "integer"}, {"type": "number"}]},
+            {"enum": [1], "allOf": [{"minimum": 2}]},
+        ],
+    )
+    def test_composition_with_no_admissible_scalar_choice_rejects_definition(self, prop_schema):
+        schema = {"type": "object", "properties": {"value": prop_schema}, "required": ["value"]}
+        with pytest.raises(ValueError, match="choices"):
+            SchemaTransformer.create_model_from_schema(schema, "ImpossibleChoiceInput", "choose")
+
+    @pytest.mark.parametrize("ref_first", [True, False])
+    def test_scalar_choice_composition_unresolved_reference_keeps_known_sibling_constraint(self, ref_first):
+        branch = {"$ref": "https://example.invalid/schema.json", "type": "integer"}
+        if not ref_first:
+            branch = dict(reversed(branch.items()))
+        schema = {
+            "type": "object",
+            "properties": {"value": {"enum": [1, "bad"], "anyOf": [branch]}},
+            "required": ["value"],
+        }
+        model = SchemaTransformer.create_model_from_schema(schema, "UnresolvedChoiceInput", "choose")
+        assert model(tool_name="choose", value=1).value == 1
+        with pytest.raises(ValidationError):
+            model(tool_name="choose", value="bad")
+        assert not Draft202012Validator(model.model_json_schema()).is_valid({"tool_name": "choose", "value": "bad"})
+
+    def test_scalar_choice_composition_unresolved_reference_alone_preserves_fallback(self):
+        schema = {
+            "type": "object",
+            "properties": {"value": {"enum": [1, "brief"], "anyOf": [{"$ref": "#/$defs/Missing"}]}},
+            "required": ["value"],
+        }
+        model = SchemaTransformer.create_model_from_schema(schema, "UnresolvedChoiceInput", "choose")
+        assert model(tool_name="choose", value=1).value == 1
+        assert model(tool_name="choose", value="brief").value == "brief"
+        with pytest.raises(ValidationError):
+            model(tool_name="choose", value="other")
+
+    @pytest.mark.parametrize(
+        "prop_schema, invalid",
+        [
+            ({"type": ["boolean", "integer"], "enum": [True]}, 1),
+            ({"type": ["boolean", "integer"], "enum": [1]}, True),
+            ({"type": ["boolean", "number"], "const": False}, 0.0),
+            ({"type": ["boolean", "number"], "const": 0}, False),
+            ({"enum": [True]}, 1.0),
+            ({"enum": [0]}, False),
+            ({"const": True}, 1),
+            ({"const": 1}, True),
+        ],
+    )
+    def test_scalar_choices_preserve_json_identity(self, prop_schema, invalid):
+        schema = {"type": "object", "properties": {"value": prop_schema}, "required": ["value"]}
+        model = SchemaTransformer.create_model_from_schema(schema, "ChoiceInput", "choose")
+        payload = {"tool_name": "choose", "value": invalid}
+
+        assert not Draft202012Validator(schema).is_valid(payload)
+        with pytest.raises(ValidationError):
+            model.model_validate(payload)
+        assert not Draft202012Validator(model.model_json_schema()).is_valid(payload)
+
+    @pytest.mark.parametrize("choices", [[True, 1], [1, True], [False, 0], [0, False]])
+    def test_mixed_boolean_and_numeric_choices_keep_the_matching_choice(self, choices):
+        schema = {"type": "object", "properties": {"value": {"enum": choices}}, "required": ["value"]}
+        model = SchemaTransformer.create_model_from_schema(schema, "MixedChoiceInput", "choose")
+        numeric = next(value for value in choices if not isinstance(value, bool))
+
+        assert type(model(tool_name="choose", value=float(numeric)).value) is int
+        boolean = next(value for value in choices if isinstance(value, bool))
+        assert model(tool_name="choose", value=boolean).value is boolean
+
+    @pytest.mark.parametrize("prop_schema", [{"type": "number", "enum": [1]}, {"const": 1.0}, {"enum": [0, 1]}])
+    def test_scalar_choices_retain_json_numeric_equality(self, prop_schema):
+        schema = {"type": "object", "properties": {"value": prop_schema}, "required": ["value"]}
+        model = SchemaTransformer.create_model_from_schema(schema, "NumericChoiceInput", "choose")
+        for value in (1, 1.0):
+            payload = {"tool_name": "choose", "value": value}
+            assert Draft202012Validator(schema).is_valid(payload)
+            assert model.model_validate(payload).value == value
+            assert Draft202012Validator(model.model_json_schema()).is_valid(payload)
+
+    @pytest.mark.parametrize(
+        "prop_schema",
+        [
+            {"type": "string", "enum": ["brief", "detailed"]},
+            {"type": "string", "const": "brief"},
+            {"type": ["string", "null"], "enum": ["brief", None]},
+            {"type": "number", "enum": [1, 2]},
+        ],
+    )
+    def test_optional_scalar_choice_schema_matches_none_validation(self, prop_schema):
+        schema = {"type": "object", "properties": {"value": prop_schema}}
+        model = SchemaTransformer.create_model_from_schema(schema, "OptionalChoiceInput", "choose")
+        payload = {"tool_name": "choose", "value": None}
+
+        assert model.model_validate(payload).value is None
+        emitted = model.model_json_schema()
+        assert Draft202012Validator(emitted).is_valid(payload)
+        assert Draft202012Validator(emitted).is_valid({"tool_name": "choose"})
+        assert not Draft202012Validator(emitted).is_valid({"tool_name": "choose", "value": "unknown"})
+
+    @pytest.mark.parametrize(
+        "prop_schema, valid, invalid, emitted_type, choices",
+        [
+            ({"type": "string", "enum": ["brief", None]}, "brief", None, "string", ["brief"]),
+            ({"type": "integer", "enum": [1, "two"]}, 1, "two", "integer", [1]),
+            ({"type": "integer", "enum": [1, True, 2.5]}, 1, True, "integer", [1]),
+            ({"type": "number", "enum": [1, 2.5, False]}, 2.5, False, "number", [1, 2.5]),
+            ({"type": "boolean", "enum": [False, 0, "false"]}, False, 0, "boolean", [False]),
+            ({"type": "integer", "enum": [1.0, 2.5]}, 1, 2.5, "integer", [1.0]),
+            ({"type": "boolean", "const": False}, False, 0, "boolean", [False]),
+        ],
+    )
+    def test_typed_scalar_choices(self, prop_schema, valid, invalid, emitted_type, choices):
+        """Choice constraints intersect the declared JSON type, including bool/number boundaries."""
+        schema = {"type": "object", "properties": {"mode": prop_schema}, "required": ["mode"]}
+        model = SchemaTransformer.create_model_from_schema(schema, "TypedModeInput", "configure")
+        assert model(tool_name="configure", mode=valid).mode == valid
+        with pytest.raises(ValueError):
+            model(tool_name="configure", mode=invalid)
+        emitted = model.model_json_schema()["properties"]["mode"]
+        assert emitted["type"] == emitted_type
+        assert emitted.get("enum", [emitted.get("const")]) == choices
+
+    @pytest.mark.parametrize("prop_schema", [{"type": "string", "const": None}, {"type": "integer", "enum": ["two"]}])
+    def test_incompatible_scalar_choices_reject_definition(self, prop_schema):
+        """An impossible type/choice combination never becomes an unrestricted input field."""
+        schema = {"type": "object", "properties": {"mode": prop_schema}, "required": ["mode"]}
+        with pytest.raises(ValueError, match="choices.*declared"):
+            SchemaTransformer.create_model_from_schema(schema, "InvalidModeInput", "configure")
+
+    def test_nullable_type_array_with_mixed_enum(self):
+        schema = {
+            "type": "object",
+            "properties": {"mode": {"type": ["string", "null"], "enum": ["brief", None, 1]}},
+            "required": ["mode"],
+        }
+        model = SchemaTransformer.create_model_from_schema(schema, "NullableModeInput", "configure")
+        assert model(tool_name="configure", mode="brief").mode == "brief"
+        assert model(tool_name="configure", mode=None).mode is None
+        with pytest.raises(ValueError):
+            model(tool_name="configure", mode=1)
+        assert model.model_json_schema()["properties"]["mode"]["enum"] == ["brief", None]
+
+    @pytest.mark.parametrize(
+        "prop_schema, valid, invalid",
+        [
+            ({"type": "string", "enum": ["brief", "detailed"]}, "brief", "unknown"),
+            ({"type": "integer", "enum": [1, 2]}, 1, 3),
+            ({"enum": ["brief", None]}, None, "unknown"),
+            ({"type": "string", "const": "brief"}, "brief", "detailed"),
+            ({"type": "integer", "const": 0}, 0, 1),
+            ({"type": "boolean", "const": False}, False, True),
+            ({"const": None}, None, "unknown"),
+        ],
+    )
+    def test_scalar_value_constraints(self, prop_schema, valid, invalid):
+        """MCP parameter choices constrain both validation and the model-facing schema."""
+        schema = {"type": "object", "properties": {"mode": prop_schema}, "required": ["mode"]}
+        model = SchemaTransformer.create_model_from_schema(schema, "ModeInput", "configure")
+
+        assert model(tool_name="configure", mode=valid).mode == valid
+        with pytest.raises(ValueError):
+            model(tool_name="configure", mode=invalid)
+        emitted = model.model_json_schema()["properties"]["mode"]
+        choices = prop_schema.get("enum", [prop_schema.get("const")])
+        if len(choices) == 1:
+            assert emitted["const"] == choices[0]
+        else:
+            assert emitted["enum"] == choices
+
+    def test_array_item_enum(self):
+        """Enum constraints survive recursive conversion of tool parameter arrays."""
+        schema = {
+            "type": "object",
+            "properties": {"modes": {"type": "array", "items": {"type": "string", "enum": ["brief", "detailed"]}}},
+            "required": ["modes"],
+        }
+        model = SchemaTransformer.create_model_from_schema(schema, "ModesInput", "configure")
+
+        assert model(tool_name="configure", modes=["brief", "detailed"]).modes == ["brief", "detailed"]
+        with pytest.raises(ValueError):
+            model(tool_name="configure", modes=["unknown"])
+        assert model.model_json_schema()["properties"]["modes"]["items"]["enum"] == ["brief", "detailed"]
+
+    @pytest.mark.parametrize("definitions_key", ["$defs", "definitions"])
+    def test_referenced_enum(self, definitions_key):
+        """Enum definitions emitted by MCP servers remain scalar parameter choices."""
+        schema = {
+            "type": "object",
+            "properties": {"mode": {"$ref": f"#/{definitions_key}/Mode"}},
+            "required": ["mode"],
+            definitions_key: {"Mode": {"type": "string", "enum": ["brief", "detailed"]}},
+        }
+        model = SchemaTransformer.create_model_from_schema(schema, "ModeInput", "configure")
+
+        assert model(tool_name="configure", mode="brief").mode == "brief"
+        with pytest.raises(ValueError):
+            model(tool_name="configure", mode="unknown")
+        assert model.model_json_schema()["properties"]["mode"]["enum"] == ["brief", "detailed"]
+
+    def test_nullable_enum_with_default(self):
+        schema = {
+            "type": "object",
+            "properties": {
+                "mode": {"anyOf": [{"type": "string", "enum": ["brief", "detailed"]}, {"type": "null"}], "default": "brief"}
+            },
+        }
+        model = SchemaTransformer.create_model_from_schema(schema, "ModeInput", "configure")
+
+        assert model(tool_name="configure").mode == "brief"
+        assert model(tool_name="configure", mode=None).mode is None
+        with pytest.raises(ValueError):
+            model(tool_name="configure", mode="unknown")
+
     @pytest.mark.parametrize("invalid_field", ["first", "second"])
     def test_repeated_reference_reuses_model_and_validates_both_fields(self, invalid_field: str) -> None:
         reference = {"$ref": "#/$defs/Item"}

pyproject.toml

@@ -33,6 +33,8 @@ dependencies = [
     "pyyaml>=6.0.2,<7.0.0",
     "requests>=2.32.3,<3.0.0",
     "mcp[cli]>=1.8.0,<2",
+    "jsonschema>=4.20.0,<5",
+    "referencing>=0.28.4,<1",
     "litellm>=1.50.0,<2.0.0",
 ]