From d677358801f3cb030555afc541727e085cfcb8f9 Mon Sep 17 00:00:00 2001
From: Adam Moussa <166072409+amoussa1229@users.noreply.github.com>
Date: Tue, 4 Aug 2026 19:39:20 -0400
Subject: [PATCH] fix(wo): reject non-str AI free-text fields (#158)

---
 lambdas/wo/email_processor/extraction.py      |  4 +-
 lambdas/wo/email_processor/template_parser.py | 23 +++++++++
 .../tests/test_validation_gate.py             | 48 +++++++++++++++++++
 3 files changed, 74 insertions(+), 1 deletion(-)

diff --git a/lambdas/wo/email_processor/extraction.py b/lambdas/wo/email_processor/extraction.py
index 9bafe0e..0e1027a 100644
--- a/lambdas/wo/email_processor/extraction.py
+++ b/lambdas/wo/email_processor/extraction.py
@@ -9,6 +9,7 @@ client is constructed (moto-before-handler invariant).
 import json
 import os
 import re
+from decimal import Decimal
 
 import boto3
 from prompts import EXTRACTION_PROMPT
@@ -85,4 +86,5 @@ def extract_with_bedrock(email_data: dict) -> dict:
     if json_match:
         response_text = json_match.group(1)
 
-    return json.loads(response_text.strip())
+    # parse_float=Decimal is CRITICAL: DynamoDB rejects Python floats.
+    return json.loads(response_text.strip(), parse_float=Decimal)
diff --git a/lambdas/wo/email_processor/template_parser.py b/lambdas/wo/email_processor/template_parser.py
index 4df3518..048aee5 100644
--- a/lambdas/wo/email_processor/template_parser.py
+++ b/lambdas/wo/email_processor/template_parser.py
@@ -53,6 +53,21 @@ VALID_STATUSES = {
     "unknown",
 }
 
+# Free-text scalar fields that must be None or str (blocks LLM-emitted maps/
+# lists from landing as DynamoDB Map/List attribute pollution, and floats that
+# would crash update_item). email_type, status, work_order_id, site_code, and
+# date fields are validated separately.
+_AI_FREE_TEXT_STR_FIELDS = (
+    "description",
+    "building",
+    "address",
+    "severity",
+    "priority",
+    "assigned_to",
+    "commenter",
+    "comment_text",
+)
+
 # Subject classifiers.
 _T1_SUBJECT = re.compile(r"^AMAZON UPDATE WO DETAILS\s+(?P<wo>\S+)\s*$")
 _T2_SUBJECT = re.compile(
@@ -480,6 +495,14 @@ def validate_ai_fallback(candidate):
         except ValueError:
             return False, "creation_time_unparseable"
 
+    # Free-text scalars must be None or str. A prompt-injected float reaches
+    # DynamoDB as Python float and crashes update_item; a dict/list lands as
+    # a Map/List attribute and pollutes downstream readers (PO parity).
+    for field in _AI_FREE_TEXT_STR_FIELDS:
+        val = candidate.get(field)
+        if val is not None and not isinstance(val, str):
+            return False, "invalid_field_type"
+
     return True, "ok"
 
 
diff --git a/lambdas/wo/email_processor/tests/test_validation_gate.py b/lambdas/wo/email_processor/tests/test_validation_gate.py
index 6b873f5..01c3f5f 100644
--- a/lambdas/wo/email_processor/tests/test_validation_gate.py
+++ b/lambdas/wo/email_processor/tests/test_validation_gate.py
@@ -374,3 +374,51 @@ def test_ai_fallback_all_valid_email_types():
         cand["email_type"] = et
         ok, reason = validate_ai_fallback(cand)
         assert ok, f"email_type={et} should pass"
+
+
+def test_ai_fallback_free_text_float_rejected():
+    # Prompt-injected JSON float in a free-text slot must fail closed before
+    # update_item (boto3 rejects Python floats -> async retries / DLQ).
+    cand = _ai_candidate()
+    cand["work_order_id"] = "12345"
+    cand["email_type"] = "update"
+    cand["severity"] = 1.5
+    ok, reason = validate_ai_fallback(cand)
+    assert not ok and reason == "invalid_field_type"
+
+
+def test_ai_fallback_free_text_dict_rejected():
+    # LLM-emitted maps in scalar slots must be rejected (DynamoDB Map pollution).
+    cand = _ai_candidate()
+    cand["work_order_id"] = "12345"
+    cand["email_type"] = "update"
+    cand["assigned_to"] = {"a": 1}
+    ok, reason = validate_ai_fallback(cand)
+    assert not ok and reason == "invalid_field_type"
+
+
+def test_ai_fallback_free_text_list_rejected():
+    # LLM-emitted lists in scalar slots must be rejected (DynamoDB List pollution).
+    cand = _ai_candidate()
+    cand["work_order_id"] = "12345"
+    cand["email_type"] = "update"
+    cand["comment_text"] = ["x"]
+    ok, reason = validate_ai_fallback(cand)
+    assert not ok and reason == "invalid_field_type"
+
+
+def test_ai_fallback_free_text_none_and_str_ok():
+    # None and str remain valid for every free-text field.
+    cand = _ai_candidate()
+    cand["work_order_id"] = "12345"
+    cand["email_type"] = "update"
+    ok, reason = validate_ai_fallback(cand)
+    assert ok and reason == "ok"
+
+    for field in template_parser._AI_FREE_TEXT_STR_FIELDS:
+        cand = _ai_candidate()
+        cand["work_order_id"] = "12345"
+        cand["email_type"] = "update"
+        cand[field] = "ok"
+        ok, reason = validate_ai_fallback(cand)
+        assert ok and reason == "ok", field