From a5cca4e7a2a791e7b19f3c900bfe8e52166bb0ff Mon Sep 17 00:00:00 2001 From: rohanbuilds-ai Date: Sun, 11 Oct 2026 12:50:08 +0530 Subject: [PATCH 1/2] fix(taxcalc): distinguish formatted SSNs from EINs --- finbot/mcp/servers/taxcalc/server.py | 28 ++-- tests/unit/mcp/__init__.py | 1 + tests/unit/mcp/test_taxcalc_server.py | 194 ++++++++++++++++++++++++++ 3 files changed, 213 insertions(+), 10 deletions(-) create mode 100644 tests/unit/mcp/__init__.py create mode 100644 tests/unit/mcp/test_taxcalc_server.py diff --git a/finbot/mcp/servers/taxcalc/server.py b/finbot/mcp/servers/taxcalc/server.py index 4aa48e56..782e8ced 100644 --- a/finbot/mcp/servers/taxcalc/server.py +++ b/finbot/mcp/servers/taxcalc/server.py @@ -150,32 +150,40 @@ 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): + # SSN format: XXX-XX-XXXX (3-2-4 grouping, area number cannot be 000) + if re.match(r"^\d{3}[\s\-]\d{2}[\s\-]\d{4}$", tax_id_trimmed): + if tax_id_clean.startswith("000"): + return { + "tax_id": tax_id, + "country": country, + "format_valid": False, + "error": "Invalid SSN: area number cannot be 000", + } return { "tax_id": tax_id, "country": country, "format_valid": True, - "id_type": "EIN", - "formatted": f"{tax_id_clean[:2]}-{tax_id_clean[2:]}", + "id_type": "SSN", + "formatted": f"{tax_id_clean[:3]}-{tax_id_clean[3:5]}-{tax_id_clean[5:]}", } - # 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 { diff --git a/tests/unit/mcp/__init__.py b/tests/unit/mcp/__init__.py new file mode 100644 index 00000000..f580e070 --- /dev/null +++ b/tests/unit/mcp/__init__.py @@ -0,0 +1 @@ +"""Unit tests for MCP servers""" diff --git a/tests/unit/mcp/test_taxcalc_server.py b/tests/unit/mcp/test_taxcalc_server.py new file mode 100644 index 00000000..26722caf --- /dev/null +++ b/tests/unit/mcp/test_taxcalc_server.py @@ -0,0 +1,194 @@ +""" +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"] + + +# ============================================================================ +# 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"] + + +# ============================================================================ +# 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"] From cb9c7bb76c45072591c84196a71a5b4513cb29c8 Mon Sep 17 00:00:00 2001 From: rohanbuilds-ai Date: Mon, 12 Oct 2026 14:27:17 +0530 Subject: [PATCH 2/2] fix(taxcalc): enforce complete SSN validation --- finbot/mcp/servers/taxcalc/server.py | 29 +++++--- tests/unit/mcp/test_taxcalc_server.py | 99 +++++++++++++++++++++++++++ 2 files changed, 120 insertions(+), 8 deletions(-) diff --git a/finbot/mcp/servers/taxcalc/server.py b/finbot/mcp/servers/taxcalc/server.py index 782e8ced..4d86c7b0 100644 --- a/finbot/mcp/servers/taxcalc/server.py +++ b/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": { @@ -154,21 +159,29 @@ def validate_tax_id(tax_id: str, country: str = "US") -> dict[str, Any]: tax_id_clean = re.sub(r"[\s\-]", "", tax_id_trimmed) if country == "US": - # SSN format: XXX-XX-XXXX (3-2-4 grouping, area number cannot be 000) - if re.match(r"^\d{3}[\s\-]\d{2}[\s\-]\d{4}$", tax_id_trimmed): - if tax_id_clean.startswith("000"): + # 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": False, - "error": "Invalid SSN: area number cannot be 000", + "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": "SSN", - "formatted": f"{tax_id_clean[:3]}-{tax_id_clean[3:5]}-{tax_id_clean[5:]}", + "format_valid": False, + "error": error_msg, } # EIN format: XX-XXXXXXX (9 digits, 2-7 grouping or unformatted digits) if re.match(r"^\d{9}$", tax_id_clean): diff --git a/tests/unit/mcp/test_taxcalc_server.py b/tests/unit/mcp/test_taxcalc_server.py index 26722caf..c9016430 100644 --- a/tests/unit/mcp/test_taxcalc_server.py +++ b/tests/unit/mcp/test_taxcalc_server.py @@ -103,6 +103,81 @@ async def test_validate_tax_id_invalid_ssn_spaced_area_zero(taxcalc_server: Any) 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 # ============================================================================ @@ -177,6 +252,30 @@ async def test_validate_tax_id_non_numeric_characters(taxcalc_server: Any) -> No 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 # ============================================================================