Bug: `Optional[int]` type hint incorrectly generates `string` type in JSON Schema

[SeanWei-lab]

Describe the bug

Description

When using Optional[int] type hints in MCP tool functions, the awslabs.mcp_lambda_handler package incorrectly generates JSON Schema with type: "string" instead of type: "integer".

Environment

• Package: awslabs.mcp_lambda_handler
• Python version: 3.x
• Type hint: Optional[int] (which is Union[int, None])

Expected Behavior

When a function parameter is annotated with Optional[int], the generated JSON Schema should be:

{
  "type": "integer",
  "description": "..."
}

Or with nullable support:

{
  "type": ["integer", "null"],
  "description": "..."
}

Current Behavior

The generated JSON Schema incorrectly uses string type:

{
  "type": "string",
  "description": "..."
}

Reproduction Steps

from typing import Optional
from awslabs.mcp_lambda_handler import MCPLambdaHandler

mcp = MCPLambdaHandler(
    name="test-mcp",
    version="1.0.0",
)

@mcp.tool()
def test_int_param(days: Optional[int] = 30, status: Optional[str] = "new") -> str:
    """Test integer parameter
    
    Args:
        days: Number of days, default 30
        status: Status, default "new"
    
    Returns:
        Test result
    """
    return f"days={days} (type: {type(days).__name__}), status={status}"

# Check generated schema
import json
print(json.dumps(mcp.tools, indent=2))

Output:

{
  "test_int_param": {
    "name": "test_int_param",
    "description": "Test integer parameter",
    "inputSchema": {
      "type": "object",
      "properties": {
        "days": {
          "type": "string",  // ❌ Should be "integer"
          "description": "Number of days, default 30"
        },
        "status": {
          "type": "string",  // ✅ Correct
          "description": "Status, default \"new\""
        }
      },
      "required": ["days", "status"]
    }
  }
}

Function annotations (correct):

{
  'days': typing.Optional[int],
  'status': typing.Optional[str],
  'return': <class 'str'>
}

Possible Solution

Add handling for Union types before the default case:

from typing import Union, get_origin, get_args

def get_type_schema(type_hint: Any) -> Dict[str, Any]:
    # Handle basic types
    if type_hint is int:
        return {'type': 'integer'}
    elif type_hint is float:
        return {'type': 'number'}
    elif type_hint is bool:
        return {'type': 'boolean'}
    elif type_hint is str:
        return {'type': 'string'}
    
    # Get origin type
    origin = get_origin(type_hint)
    
    # Handle Union types (including Optional)
    if origin is Union:
        args = get_args(type_hint)
        # Filter out NoneType
        non_none_args = [arg for arg in args if arg is not type(None)]
        
        if len(non_none_args) == 1:
            # This is Optional[T], return schema for T
            return get_type_schema(non_none_args[0])
        else:
            # Multiple non-None types in Union
            # Could use anyOf, but for simplicity, use the first type
            return get_type_schema(non_none_args[0])
    
    if origin is None:
        return {'type': 'string'}  # Default for unknown types
    
    # ... rest of the code ...

Additional Information/Context

This issue was discovered when integrating with MCP clients that validate parameter types against the JSON Schema. The clients correctly send integer values, but the schema indicates they should be strings, causing confusion and requiring workarounds.

OS

Lambda

Server

lambda-mcp-server

Server Version

0.1.9

Region experiencing the issue

cn-north-1

Other information

No response

Service quota

• I have reviewed the service quotas for this construct