Fix SSN and EIN classification in TaxCalc validation

finbot/mcp/servers/taxcalc/server.py

@@ -11,9 +11,14 @@
 from fastmcp import FastMCP
 
 from finbot.core.auth.session import SessionContext
+from finbot.ctf.detectors.primitives.pii import _PATTERN_CATALOG
 
 logger = logging.getLogger(__name__)
 
+# Full US SSN pattern from pii.py (rejects area 000, 666, 900-999, group 00, serial 0000)
+_US_SSN_PATTERN: str = next(p.regex for p in _PATTERN_CATALOG if p.name == "us_ssn")
+_US_SSN_REGEX: re.Pattern[str] = re.compile(_US_SSN_PATTERN)
+
 DEFAULT_CONFIG: dict[str, Any] = {
     "default_jurisdiction": "US-CA",
     "tax_rates": {
@@ -150,32 +155,48 @@ def validate_tax_id(tax_id: str, country: str = "US") -> dict[str, Any]:
         Checks if the provided tax ID matches expected formats for the given country.
         Note: This validates format only, not whether the ID is registered with tax authorities.
         """
-        tax_id_clean = re.sub(r"[\s\-]", "", tax_id)
+        tax_id_trimmed = tax_id.strip()
+        tax_id_clean = re.sub(r"[\s\-]", "", tax_id_trimmed)
 
         if country == "US":
-            # EIN format: XX-XXXXXXX (9 digits)
-            if re.match(r"^\d{9}$", tax_id_clean):
+            # Normalize accepted space-separated SSN input (XXX XX XXXX) to hyphenated form (XXX-XX-XXXX)
+            tax_id_hyphenated = re.sub(r"\s+", "-", tax_id_trimmed)
+
+            # SSN format: XXX-XX-XXXX (3-2-4 grouping, validated against US SSN rules from pii.py)
+            if re.match(r"^\d{3}-\d{2}-\d{4}$", tax_id_hyphenated):
+                if _US_SSN_REGEX.fullmatch(tax_id_hyphenated):
+                    return {
+                        "tax_id": tax_id,
+                        "country": country,
+                        "format_valid": True,
+                        "id_type": "SSN",
+                        "formatted": f"{tax_id_clean[:3]}-{tax_id_clean[3:5]}-{tax_id_clean[5:]}",
+                    }
+                area = tax_id_clean[:3]
+                if area == "000" or area == "666" or area.startswith("9"):
+                    error_msg = "Invalid SSN: area number cannot be 000, 666, or 900-999"
+                else:
+                    error_msg = "Invalid SSN: invalid group or serial number"
                 return {
                     "tax_id": tax_id,
                     "country": country,
-                    "format_valid": True,
-                    "id_type": "EIN",
-                    "formatted": f"{tax_id_clean[:2]}-{tax_id_clean[2:]}",
+                    "format_valid": False,
+                    "error": error_msg,
                 }
-            # SSN format: XXX-XX-XXXX (9 digits, different grouping)
-            if re.match(r"^\d{9}$", tax_id_clean) and tax_id_clean[:3] != "00":
+            # EIN format: XX-XXXXXXX (9 digits, 2-7 grouping or unformatted digits)
+            if re.match(r"^\d{9}$", tax_id_clean):
                 return {
                     "tax_id": tax_id,
                     "country": country,
                     "format_valid": True,
-                    "id_type": "SSN",
-                    "formatted": f"{tax_id_clean[:3]}-{tax_id_clean[3:5]}-{tax_id_clean[5:]}",
+                    "id_type": "EIN",
+                    "formatted": f"{tax_id_clean[:2]}-{tax_id_clean[2:]}",
                 }
             return {
                 "tax_id": tax_id,
                 "country": country,
                 "format_valid": False,
-                "error": "US tax IDs must be 9 digits (EIN: XX-XXXXXXX)",
+                "error": "US tax IDs must be 9 digits (EIN: XX-XXXXXXX or SSN: XXX-XX-XXXX)",
             }
 
         return {

tests/unit/mcp/__init__.py

@@ -0,0 +1 @@
+"""Unit tests for MCP servers"""

tests/unit/mcp/test_taxcalc_server.py

@@ -0,0 +1,293 @@
+"""
+Unit tests for the TaxCalc MCP Server -- tax ID validation.
+"""
+
+from typing import Any
+
+import pytest
+
+from finbot.mcp.servers.taxcalc.server import create_taxcalc_server
+
+
+@pytest.fixture
+def taxcalc_server() -> Any:
+    """Create a TaxCalc MCP server instance."""
+    return create_taxcalc_server(session_context=None)  # type: ignore[arg-type]
+
+
+# ============================================================================
+# EIN Validation Tests
+# ============================================================================
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_validate_tax_id_valid_ein_hyphenated(taxcalc_server: Any) -> None:
+    """EIN formatted with standard hyphen (XX-XXXXXXX) is valid."""
+    res = await taxcalc_server.call_tool("validate_tax_id", {"tax_id": "12-3456789"})
+    data = res.structured_content
+    assert data["format_valid"] is True
+    assert data["id_type"] == "EIN"
+    assert data["formatted"] == "12-3456789"
+    assert data["country"] == "US"
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_validate_tax_id_valid_ein_raw_digits(taxcalc_server: Any) -> None:
+    """Raw 9-digit string without hyphens defaults to EIN."""
+    res = await taxcalc_server.call_tool("validate_tax_id", {"tax_id": "123456789"})
+    data = res.structured_content
+    assert data["format_valid"] is True
+    assert data["id_type"] == "EIN"
+    assert data["formatted"] == "12-3456789"
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_validate_tax_id_valid_ein_spaced(taxcalc_server: Any) -> None:
+    """EIN with spaces instead of hyphen is valid and formatted correctly."""
+    res = await taxcalc_server.call_tool("validate_tax_id", {"tax_id": "98 7654321"})
+    data = res.structured_content
+    assert data["format_valid"] is True
+    assert data["id_type"] == "EIN"
+    assert data["formatted"] == "98-7654321"
+
+
+# ============================================================================
+# SSN Validation Tests (Explicit 3-2-4 Grouping)
+# ============================================================================
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_validate_tax_id_valid_ssn_hyphenated(taxcalc_server: Any) -> None:
+    """SSN formatted with standard hyphens (XXX-XX-XXXX) is identified as SSN."""
+    res = await taxcalc_server.call_tool("validate_tax_id", {"tax_id": "123-45-6789"})
+    data = res.structured_content
+    assert data["format_valid"] is True
+    assert data["id_type"] == "SSN"
+    assert data["formatted"] == "123-45-6789"
+    assert data["country"] == "US"
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_validate_tax_id_valid_ssn_spaced(taxcalc_server: Any) -> None:
+    """SSN formatted with spaces (XXX XX XXXX) is identified as SSN."""
+    res = await taxcalc_server.call_tool("validate_tax_id", {"tax_id": "123 45 6789"})
+    data = res.structured_content
+    assert data["format_valid"] is True
+    assert data["id_type"] == "SSN"
+    assert data["formatted"] == "123-45-6789"
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_validate_tax_id_invalid_ssn_area_zero(taxcalc_server: Any) -> None:
+    """SSN with area number 000 is rejected as invalid."""
+    res = await taxcalc_server.call_tool("validate_tax_id", {"tax_id": "000-12-3456"})
+    data = res.structured_content
+    assert data["format_valid"] is False
+    assert "Invalid SSN" in data["error"]
+    assert "000" in data["error"]
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_validate_tax_id_invalid_ssn_spaced_area_zero(taxcalc_server: Any) -> None:
+    """Spaced SSN with area number 000 is rejected as invalid."""
+    res = await taxcalc_server.call_tool("validate_tax_id", {"tax_id": "000 12 3456"})
+    data = res.structured_content
+    assert data["format_valid"] is False
+    assert "Invalid SSN" in data["error"]
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_validate_tax_id_invalid_ssn_area_666(taxcalc_server: Any) -> None:
+    """SSN with area number 666 is rejected as invalid."""
+    res = await taxcalc_server.call_tool("validate_tax_id", {"tax_id": "666-12-3456"})
+    data = res.structured_content
+    assert data["format_valid"] is False
+    assert "Invalid SSN" in data["error"]
+    assert "666" in data["error"]
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_validate_tax_id_invalid_ssn_spaced_area_666(taxcalc_server: Any) -> None:
+    """Spaced SSN with area number 666 is rejected as invalid."""
+    res = await taxcalc_server.call_tool("validate_tax_id", {"tax_id": "666 12 3456"})
+    data = res.structured_content
+    assert data["format_valid"] is False
+    assert "Invalid SSN" in data["error"]
+    assert "666" in data["error"]
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_validate_tax_id_invalid_ssn_area_900(taxcalc_server: Any) -> None:
+    """SSN with area number 900 is rejected as invalid."""
+    res = await taxcalc_server.call_tool("validate_tax_id", {"tax_id": "900-12-3456"})
+    data = res.structured_content
+    assert data["format_valid"] is False
+    assert "Invalid SSN" in data["error"]
+    assert "900-999" in data["error"]
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_validate_tax_id_invalid_ssn_area_999(taxcalc_server: Any) -> None:
+    """SSN with area number 999 is rejected as invalid."""
+    res = await taxcalc_server.call_tool("validate_tax_id", {"tax_id": "999-12-3456"})
+    data = res.structured_content
+    assert data["format_valid"] is False
+    assert "Invalid SSN" in data["error"]
+    assert "900-999" in data["error"]
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_validate_tax_id_invalid_ssn_spaced_area_900(taxcalc_server: Any) -> None:
+    """Spaced SSN with area number in 900-999 range is rejected as invalid."""
+    res = await taxcalc_server.call_tool("validate_tax_id", {"tax_id": "950 12 3456"})
+    data = res.structured_content
+    assert data["format_valid"] is False
+    assert "Invalid SSN" in data["error"]
+    assert "900-999" in data["error"]
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_validate_tax_id_invalid_ssn_group_zero(taxcalc_server: Any) -> None:
+    """SSN with group number 00 is rejected as invalid."""
+    res = await taxcalc_server.call_tool("validate_tax_id", {"tax_id": "123-00-6789"})
+    data = res.structured_content
+    assert data["format_valid"] is False
+    assert "Invalid SSN" in data["error"]
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_validate_tax_id_invalid_ssn_serial_zero(taxcalc_server: Any) -> None:
+    """SSN with serial number 0000 is rejected as invalid."""
+    res = await taxcalc_server.call_tool("validate_tax_id", {"tax_id": "123-45-0000"})
+    data = res.structured_content
+    assert data["format_valid"] is False
+    assert "Invalid SSN" in data["error"]
+
+
+# ============================================================================
+# Whitespace Handling Tests
+# ============================================================================
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_validate_tax_id_ssn_with_surrounding_whitespace(taxcalc_server: Any) -> None:
+    """SSN with leading and trailing whitespace is trimmed and recognized as SSN."""
+    res = await taxcalc_server.call_tool("validate_tax_id", {"tax_id": "  123-45-6789  "})
+    data = res.structured_content
+    assert data["format_valid"] is True
+    assert data["id_type"] == "SSN"
+    assert data["formatted"] == "123-45-6789"
+    assert data["country"] == "US"
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_validate_tax_id_ein_with_surrounding_whitespace(taxcalc_server: Any) -> None:
+    """EIN with leading and trailing whitespace is trimmed and recognized as EIN."""
+    res = await taxcalc_server.call_tool("validate_tax_id", {"tax_id": "  12-3456789  "})
+    data = res.structured_content
+    assert data["format_valid"] is True
+    assert data["id_type"] == "EIN"
+    assert data["formatted"] == "12-3456789"
+    assert data["country"] == "US"
+
+
+# ============================================================================
+# Invalid Lengths & Malformed Identifiers
+# ============================================================================
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_validate_tax_id_too_short(taxcalc_server: Any) -> None:
+    """Tax ID with fewer than 9 digits is rejected."""
+    res = await taxcalc_server.call_tool("validate_tax_id", {"tax_id": "12345"})
+    data = res.structured_content
+    assert data["format_valid"] is False
+    assert "must be 9 digits" in data["error"]
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_validate_tax_id_too_long(taxcalc_server: Any) -> None:
+    """Tax ID with more than 9 digits is rejected."""
+    res = await taxcalc_server.call_tool("validate_tax_id", {"tax_id": "12345678901"})
+    data = res.structured_content
+    assert data["format_valid"] is False
+    assert "must be 9 digits" in data["error"]
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_validate_tax_id_empty_string(taxcalc_server: Any) -> None:
+    """Empty string tax ID is rejected."""
+    res = await taxcalc_server.call_tool("validate_tax_id", {"tax_id": ""})
+    data = res.structured_content
+    assert data["format_valid"] is False
+    assert "must be 9 digits" in data["error"]
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_validate_tax_id_non_numeric_characters(taxcalc_server: Any) -> None:
+    """Tax ID containing letters or special characters is rejected."""
+    res = await taxcalc_server.call_tool("validate_tax_id", {"tax_id": "12-ABC5678"})
+    data = res.structured_content
+    assert data["format_valid"] is False
+    assert "must be 9 digits" in data["error"]
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_validate_tax_id_ssn_embedded_in_extra_characters(
+    taxcalc_server: Any,
+) -> None:
+    """A valid SSN embedded within extra characters or text is rejected."""
+    cases = [
+        "prefix 123-45-6789",
+        "123-45-6789 suffix",
+        "prefix 123-45-6789 suffix",
+        "123-45-6789-extra",
+        "extra-123-45-6789",
+        "123 45 6789 extra",
+        "prefix 123 45 6789",
+        "Tax ID: 123-45-6789",
+        "A123-45-6789",
+        "123-45-6789B",
+    ]
+    for case in cases:
+        res = await taxcalc_server.call_tool("validate_tax_id", {"tax_id": case})
+        data = res.structured_content
+        assert data["format_valid"] is False, f"Expected {case!r} to be rejected"
+
+
+# ============================================================================
+# Country Handling
+# ============================================================================
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_validate_tax_id_unsupported_country(taxcalc_server: Any) -> None:
+    """Countries other than 'US' return an unsupported error."""
+    res = await taxcalc_server.call_tool(
+        "validate_tax_id", {"tax_id": "12-3456789", "country": "CA"}
+    )
+    data = res.structured_content
+    assert data["format_valid"] is False
+    assert "Tax ID validation not supported for country: CA" in data["error"]