MCP Enum Validation Fails Due to Dynamic Enum Class Creation

[maxjeblick]

Disclaimer:
Bug report is generated by opus 4.5, after inspecting and trying to fix an error I encountered with kaggle mcp tools.
I haven't checked the faithfulness of the entire bug report.

Bug Report: MCP Enum Validation Fails Due to Dynamic Enum Class Creation

Version

https://github.com/NVIDIA/NeMo-Agent-Toolkit.git@release/1.4

Installation Method

• PyPi
• Source

Describe the bug

When using MCP tools with enum parameters (e.g., Kaggle's search_datasets with sortBy), validation fails because NAT's model_from_mcp_schema() function creates new Python Enum classes on each invocation.

Python enums with the same name and values but from different class definitions are not equal:

Enum1 = Enum("SortbyEnum", {"Relevance": "Relevance"})
Enum2 = Enum("SortbyEnum", {"Relevance": "Relevance"})
assert Enum1.Relevance != Enum2.Relevance  # True! Different classes

Flow causing the bug:

1. model_from_mcp_schema() creates SortbyEnum1 when building tool schema
2. LangChain creates enum instance SortbyEnum1.Relevance from LLM tool call
3. Inside mcp_session_tool_function._response_fn(), session_tool.input_schema.model_validate(kwargs) is called
4. client.get_tool(tool.name) internally calls model_from_mcp_schema() again, creating SortbyEnum2
5. Pydantic validation fails: SortbyEnum1.Relevance is not a valid SortbyEnum2

Error message:

1 validation error for SearchDatasetsInputSchema
request.sortBy
  Input should be 'Hottest', 'Votes', 'Updated', 'Active', 'Published', 'Relevance', 'LastViewed', 'Usability', 'DownloadCount' or 'NotebookCount' [type=enum, input_value=<SortbyEnum.Relevance: 'Relevance'>, input_type=SortbyEnum]

Minimum reproducible example

"""
Minimal reproduction of NAT MCP enum validation bug.

Run with: python test_enum_bug.py
Requires: pip install pydantic
"""
from enum import Enum
from typing import Any

from pydantic import BaseModel, Field, ValidationError, create_model


def model_from_mcp_schema_simplified(name: str, enum_values: list[str]) -> type[BaseModel]:
    """Simplified version of NAT's model_from_mcp_schema that shows the bug."""
    # This is what NAT does - creates a NEW enum class each time
    enum_name = f"{name.capitalize()}Enum"
    enum_type = Enum(enum_name, {item: item for item in enum_values})
    
    return create_model(
        f"{name}InputSchema",
        sortBy=(enum_type, Field(default=None)),
    )


def test_enum_validation_bug():
    """Demonstrate the enum validation bug."""
    enum_values = ["Relevance", "Votes", "DownloadCount"]
    
    # Step 1: NAT creates schema for tool (first call to model_from_mcp_schema)
    Schema1 = model_from_mcp_schema_simplified("search", enum_values)
    
    # Step 2: LangChain creates an instance with enum value
    # (LangChain uses Schema1 to parse LLM output)
    instance1 = Schema1(sortBy=Schema1.model_fields["sortBy"].annotation.Relevance)
    print(f"Created instance with: {instance1.sortBy}")
    print(f"  Type: {type(instance1.sortBy)}")
    print(f"  Value: {instance1.sortBy.value}")
    
    # Step 3: NAT re-validates with a NEW schema (second call to model_from_mcp_schema)
    # This happens in mcp_session_tool_function when calling session_tool.input_schema.model_validate()
    Schema2 = model_from_mcp_schema_simplified("search", enum_values)
    
    # Step 4: Validation fails because Schema1's enum != Schema2's enum
    try:
        # Try to validate the instance from Schema1 against Schema2
        Schema2.model_validate({"sortBy": instance1.sortBy})
        print("✓ Validation passed (unexpected)")
    except ValidationError as e:
        print(f"\n✗ Validation FAILED (this is the bug):")
        print(f"  {e.errors()[0]['msg']}")
        print(f"\nRoot cause:")
        print(f"  Schema1 enum class: {Schema1.model_fields['sortBy'].annotation}")
        print(f"  Schema2 enum class: {Schema2.model_fields['sortBy'].annotation}")
        print(f"  Same class? {Schema1.model_fields['sortBy'].annotation is Schema2.model_fields['sortBy'].annotation}")
    
    # Workaround: Convert enum to string before re-validation
    print("\n--- Workaround: Convert enum to string ---")
    string_value = instance1.sortBy.value  # "Relevance"
    try:
        validated = Schema2.model_validate({"sortBy": string_value})
        print(f"✓ Validation passed with string: {validated.sortBy}")
    except ValidationError as e:
        print(f"✗ Still failed: {e}")


def test_suggested_fix():
    """Demonstrate the suggested fix: cache enum classes."""
    enum_values = ["Relevance", "Votes", "DownloadCount"]
    
    # Cache for enum classes (keyed by frozen set of values)
    _enum_cache: dict[tuple[str, frozenset[str]], type[Enum]] = {}
    
    def get_or_create_enum(name: str, values: list[str]) -> type[Enum]:
        """Get cached enum or create new one."""
        cache_key = (name, frozenset(values))
        if cache_key not in _enum_cache:
            enum_name = f"{name.capitalize()}Enum"
            _enum_cache[cache_key] = Enum(enum_name, {item: item for item in values})
        return _enum_cache[cache_key]
    
    def model_from_mcp_schema_fixed(name: str, enum_values: list[str]) -> type[BaseModel]:
        """Fixed version that caches enum classes."""
        enum_type = get_or_create_enum(name, enum_values)
        return create_model(
            f"{name}InputSchema",
            sortBy=(enum_type, Field(default=None)),
        )
    
    # Now both schemas use the SAME enum class
    Schema1 = model_from_mcp_schema_fixed("search", enum_values)
    Schema2 = model_from_mcp_schema_fixed("search", enum_values)
    
    print(f"Same enum class? {Schema1.model_fields['sortBy'].annotation is Schema2.model_fields['sortBy'].annotation}")
    
    instance1 = Schema1(sortBy=Schema1.model_fields["sortBy"].annotation.Relevance)
    
    try:
        validated = Schema2.model_validate({"sortBy": instance1.sortBy})
        print(f"✓ Validation passed with cached enum: {validated.sortBy}")
    except ValidationError as e:
        print(f"✗ Failed: {e}")


if __name__ == "__main__":
    print("=" * 60)
    print("TEST 1: Demonstrating the bug")
    print("=" * 60)
    test_enum_validation_bug()
    
    print("\n" + "=" * 60)
    print("TEST 2: Suggested fix with enum caching")
    print("=" * 60)
    test_suggested_fix()

Output:

============================================================
TEST 1: Demonstrating the bug
============================================================
Created instance with: SearchEnum.Relevance
  Type: <enum 'SearchEnum'>
  Value: Relevance

✗ Validation FAILED (this is the bug):
  Input should be 'Relevance', 'Votes' or 'DownloadCount'

Root cause:
  Schema1 enum class: <enum 'SearchEnum'>
  Schema2 enum class: <enum 'SearchEnum'>
  Same class? False

--- Workaround: Convert enum to string ---
✓ Validation passed with string: SearchEnum.Relevance

============================================================
TEST 2: Suggested fix with enum caching
============================================================
Same enum class? True
✓ Validation passed with cached enum: SearchEnum.Relevance

Relevant log output

Click here to see error details

From NAT workflow trajectory when calling Kaggle MCP tool:
tool_outputs: "1 validation error for SearchDatasetsInputSchema

request.sortBy

Input should be 'Hottest', 'Votes', 'Updated', 'Active', 'Published', 'Relevance', 'LastViewed', 'Usability', 'DownloadCount' or 'NotebookCount' [type=enum, input_value=<SortbyEnum.Relevance: 'Relevance'>, input_type=SortbyEnum]

For further information visit https://errors.pydantic.dev/2.12/v/enum"

Other/Misc.

Location of bug: nat/plugins/mcp/utils.py in model_from_mcp_schema() function, lines 96-98:

enum_name = f"{name.capitalize()}Enum"
enum_type: Any = Enum(enum_name, {item: item for item in non_null_vals})

Suggested fix: Add enum class caching at module level:

# Module-level cache for enum classes
_enum_class_cache: dict[tuple[str, frozenset[str]], type[Enum]] = {}

def _get_or_create_enum(name: str, values: list[str]) -> type[Enum]:
    """Get cached enum class or create new one."""
    cache_key = (name.capitalize() + "Enum", frozenset(values))
    if cache_key not in _enum_class_cache:
        _enum_class_cache[cache_key] = Enum(cache_key[0], {item: item for item in values})
    return _enum_class_cache[cache_key]

Then replace line 98 with:

enum_type: Any = _get_or_create_enum(name, non_null_vals)

Alternative fix: Instead of caching, configure Pydantic to use use_enum_values=True in model config, which would store the string value instead of the enum instance. However, this changes the API contract.

Current workaround: Convert all enum instances to their string .value before calling tool.ainvoke(). However, this workaround is incomplete because NAT internally re-validates with new enum classes in mcp_session_tool_function._response_fn() at line 498 in client_impl.py.

Code of Conduct

• I agree to follow the NeMo Agent toolkit Code of Conduct
• I have searched the open bugs and have found no duplicates for this bug report