Make Cancelled PO status sticky via atomic write

The PO merge path read status with a get_item (_is_cancelled) and then
wrote with an unconditional update_item. Two defects followed from this:

- Race (Issue A): a cancellation landing between the read and the write
  was silently un-cancelled by a revision carrying a non-cancelled
  po_status — a TOCTOU on a table with concurrent email processing.
- Over-broad strip (Issue B): save_revision dropped po_status whenever
  the PO was Cancelled, so legitimate status updates on non-cancelled
  POs and status-less revisions were affected rather than only the true
  un-cancel transition.

Enforce the invariant server-side instead. "Cancelled" is a sticky,
authoritative status: once set, later new_po/revision emails may enrich
other fields but must never move it to a non-cancelled status. When the
payload carries a non-cancelled po_status, _merge_update issues the
update_item guarded by ConditionExpression "attribute_not_exists(po_status)
OR po_status <> :marker", evaluated atomically at write time, so a
cancellation that lands first always wins. On ConditionalCheckFailedException
the same fields are re-written without po_status/cancelled_at, enriching the
record while Cancelled sticks. Payloads with no status change, or an already
-Cancelled status, take a plain merge — the status is only ever suppressed on
a real un-cancel. This removes the non-atomic get_item from the write path;
_is_cancelled is deleted. Key schema and attribute names are unchanged, so the
cross-stack purchase-orders contract (read-only by seahaven-slack-bot) holds.

Add moto-backed tests covering un-cancel suppression with field enrichment,
status-less merge onto a Cancelled PO, legitimate status updates on
non-cancelled POs, new_po backfill of a Cancelled skeleton, fresh
create/merge, and authoritative save_cancellation.

Refs: #97

lambdas/po/email_processor/handler.py

@@ -27,6 +27,11 @@ dynamodb = boto3.resource("dynamodb")
 PO_TABLE = os.environ.get("PO_TABLE", "purchase-orders")
 ANTHROPIC_API_KEY_SECRET_ARN = os.environ.get("ANTHROPIC_API_KEY_SECRET_ARN")
 
+# "Cancelled" is a sticky, authoritative status: once a PO reaches it, a later
+# new_po/revision may enrich other fields but must never move it back to a
+# non-cancelled status.
+CANCELLED_STATUS = "Cancelled"
+
 EXTRACTION_PROMPT = """\
 You are an email parser for a purchase order ingest pipeline.
 The emails are Coupa procurement platform notifications containing purchase order
@@ -320,12 +325,18 @@ def enrich_parsed(parsed: dict, s3_key: str, email_subject: str):
     return parsed
 
 
-def _merge_update(po_number: str, fields: dict):
-    """Apply a merge (SET-only) update of the given fields onto a PO record.
+def _write_fields(po_number: str, fields: dict, *, guard_cancelled: bool):
+    """SET the given non-null fields on a PO record via update_item.
 
     Only the fields supplied are written; absent fields are left untouched, so a
     partial payload can never delete data that an earlier email established. The
     record is created if it does not exist (DynamoDB update_item upsert).
+
+    When ``guard_cancelled`` is True the write carries a ConditionExpression that
+    only permits it while the record is not already Cancelled. The condition is
+    evaluated atomically by DynamoDB at write time, so a cancellation that lands
+    first always wins — there is no read-then-write TOCTOU window. A failed guard
+    raises ConditionalCheckFailedException for the caller to handle.
     """
     table = dynamodb.Table(PO_TABLE)
 
@@ -344,19 +355,58 @@ def _merge_update(po_number: str, fields: dict):
     if not set_parts:
         return
 
-    table.update_item(
-        Key={"po_number": po_number},
-        UpdateExpression="SET " + ", ".join(set_parts),
-        ExpressionAttributeNames=attr_names,
-        ExpressionAttributeValues=attr_values,
-    )
+    params = {
+        "Key": {"po_number": po_number},
+        "UpdateExpression": "SET " + ", ".join(set_parts),
+        "ExpressionAttributeNames": attr_names,
+        "ExpressionAttributeValues": attr_values,
+    }
+    if guard_cancelled:
+        params["ExpressionAttributeValues"][":__cancelled_marker"] = CANCELLED_STATUS
+        params["ConditionExpression"] = (
+            "attribute_not_exists(po_status) OR po_status <> :__cancelled_marker"
+        )
+
+    table.update_item(**params)
 
 
-def _is_cancelled(po_number: str) -> bool:
-    """Return True if the PO already exists with a Cancelled status."""
-    table = dynamodb.Table(PO_TABLE)
-    existing = table.get_item(Key={"po_number": po_number}).get("Item")
-    return bool(existing) and existing.get("po_status") == "Cancelled"
+def _merge_update(po_number: str, fields: dict):
+    """Merge (SET-only) the given fields onto a PO record, keeping Cancelled sticky.
+
+    Only the fields supplied are written; absent fields are left untouched. The
+    record is created if it does not exist (DynamoDB update_item upsert).
+
+    "Cancelled" is a sticky, authoritative status. When the incoming payload
+    carries a non-cancelled ``po_status``, the write is guarded by a
+    ConditionExpression so the status is only applied while the record is not
+    already Cancelled — enforced atomically at write time, eliminating the
+    read-then-write TOCTOU where a concurrently-landing cancellation could be
+    silently un-cancelled. If the guard fails (the PO is already Cancelled), the
+    same fields are re-written WITHOUT po_status/cancelled_at and
+    unconditionally, so the other fields still merge while the Cancelled status
+    stays intact.
+
+    A payload with no ``po_status``, or one whose status is already "Cancelled",
+    needs no guard — a plain merge is correct. This is what keeps legitimate
+    status updates (non-cancelled PO) and status-less revisions from ever being
+    dropped: the status is only ever suppressed on a true un-cancel transition.
+    """
+    incoming_status = fields.get("po_status")
+    if incoming_status is None or incoming_status == CANCELLED_STATUS:
+        _write_fields(po_number, fields, guard_cancelled=False)
+        return
+
+    try:
+        _write_fields(po_number, fields, guard_cancelled=True)
+    except dynamodb.meta.client.exceptions.ConditionalCheckFailedException:
+        logger.info(
+            f"PO {po_number} is Cancelled; suppressing incoming "
+            f"po_status={incoming_status!r} and merging remaining fields"
+        )
+        enrich_fields = {
+            k: v for k, v in fields.items() if k not in ("po_status", "cancelled_at")
+        }
+        _write_fields(po_number, enrich_fields, guard_cancelled=False)
 
 
 def save_new_po(parsed: dict):
@@ -364,21 +414,13 @@ def save_new_po(parsed: dict):
 
     Uses a merge update rather than a conditional put so that an out-of-order
     cancellation (which leaves a Cancelled skeleton) is filled in with the full
-    PO data instead of the new_po being silently dropped. The cancellation marker
-    (po_status=Cancelled, cancelled_at) is preserved; new_po data backfills the
-    remaining fields.
+    PO data instead of the new_po being silently dropped. "Cancelled" is a sticky
+    status enforced atomically inside _merge_update: if the PO was already
+    cancelled, the new_po backfills its remaining fields (supplier, line_items,
+    amounts) but never un-cancels it.
     """
     po_number = parsed["po_number"]
     fields = {k: v for k, v in parsed.items() if v is not None}
-    if _is_cancelled(po_number):
-        # Preserve the cancellation: don't overwrite po_status/cancelled_at.
-        fields.pop("po_status", None)
-        fields.pop("cancelled_at", None)
-        logger.info(
-            f"PO {po_number} was cancelled before new_po arrived; "
-            f"backfilling data and keeping Cancelled status"
-        )
-
     _merge_update(po_number, fields)
     logger.info(f"Created/merged PO {po_number}")
 
@@ -389,12 +431,15 @@ def save_revision(parsed: dict):
     A revision email often omits unchanged sections (line_items, supplier). The
     previous full-overwrite put_item permanently dropped those. This SETs only the
     fields present in the revision, leaving everything else intact.
+
+    "Cancelled" is a sticky status: a revision may enrich a cancelled PO's fields
+    but must never move it to a non-cancelled status. That invariant is enforced
+    atomically inside _merge_update and applies ONLY to the un-cancel transition —
+    a revision that carries no status change, or one targeting a non-cancelled PO,
+    updates po_status normally.
     """
     po_number = parsed["po_number"]
     fields = {k: v for k, v in parsed.items() if v is not None}
-    if _is_cancelled(po_number):
-        # A revision must not silently un-cancel a PO.
-        fields.pop("po_status", None)
     _merge_update(po_number, fields)
     logger.info(f"Revised PO {po_number}")
 

tests/requirements.txt

@@ -0,0 +1 @@
+moto>=5.0.0

tests/test_po_merge.py

@@ -0,0 +1,163 @@
+"""Tests for the sticky-Cancelled merge semantics in the PO email processor.
+
+These exercise the atomic conditional write in ``_merge_update`` that makes
+"Cancelled" a sticky, authoritative status: once a PO is Cancelled, a later
+new_po or revision may enrich its other fields but can never move it back to a
+non-cancelled status. The invariant is enforced server-side with a DynamoDB
+ConditionExpression, so it holds even under concurrent email processing — there
+is no read-then-write TOCTOU window.
+
+A real in-memory purchase-orders table is stood up with moto so the tests
+exercise the actual DynamoDB conditional-write and ConditionalCheckFailedException
+fallback path, not a stub.
+"""
+
+import boto3
+import pytest
+from moto import mock_aws
+
+TABLE_NAME = "purchase-orders"
+
+
+@pytest.fixture
+def po_table():
+    """Stand up a mock purchase-orders table matching the real key schema."""
+    with mock_aws():
+        resource = boto3.resource("dynamodb", region_name="us-east-1")
+        table = resource.create_table(
+            TableName=TABLE_NAME,
+            KeySchema=[{"AttributeName": "po_number", "KeyType": "HASH"}],
+            AttributeDefinitions=[
+                {"AttributeName": "po_number", "AttributeType": "S"},
+            ],
+            BillingMode="PAY_PER_REQUEST",
+        )
+        table.wait_until_exists()
+        yield table
+
+
+def _item(table, po_number):
+    return table.get_item(Key={"po_number": po_number}).get("Item")
+
+
+def test_revision_cannot_uncancel_but_enriches(po_handler, po_table):
+    """(a) Revision carrying a live status onto a Cancelled PO keeps it Cancelled
+    while still writing the other revision fields (no resurrection)."""
+    po_table.put_item(
+        Item={
+            "po_number": "PO-A",
+            "po_status": "Cancelled",
+            "cancelled_at": "2026-07-01T00:00:00",
+        }
+    )
+
+    po_handler.save_revision(
+        {
+            "po_number": "PO-A",
+            "po_status": "Issued",
+            "total_amount": 500,
+            "supplier": {"name": "Acme"},
+        }
+    )
+
+    item = _item(po_table, "PO-A")
+    assert item["po_status"] == "Cancelled"  # not resurrected
+    assert item["cancelled_at"] == "2026-07-01T00:00:00"  # marker preserved
+    assert item["total_amount"] == 500  # other fields still enriched
+    assert item["supplier"] == {"name": "Acme"}
+
+
+def test_revision_without_status_merges_on_cancelled(po_handler, po_table):
+    """(b) Revision with NO po_status onto a Cancelled PO merges cleanly, no
+    exception, status stays Cancelled."""
+    po_table.put_item(Item={"po_number": "PO-B", "po_status": "Cancelled"})
+
+    po_handler.save_revision({"po_number": "PO-B", "total_amount": 250})
+
+    item = _item(po_table, "PO-B")
+    assert item["po_status"] == "Cancelled"
+    assert item["total_amount"] == 250
+
+
+def test_revision_updates_status_on_non_cancelled(po_handler, po_table):
+    """(c) Issue B: a legitimate status update on a NON-cancelled PO is applied,
+    never dropped."""
+    po_table.put_item(Item={"po_number": "PO-C", "po_status": "Issued"})
+
+    po_handler.save_revision(
+        {"po_number": "PO-C", "po_status": "Revised", "total_amount": 999}
+    )
+
+    item = _item(po_table, "PO-C")
+    assert item["po_status"] == "Revised"
+    assert item["total_amount"] == 999
+
+
+def test_new_po_backfills_cancelled_skeleton(po_handler, po_table):
+    """(d) new_po onto a Cancelled skeleton backfills its fields but keeps the
+    Cancelled status and cancellation marker."""
+    po_table.put_item(
+        Item={
+            "po_number": "PO-D",
+            "po_status": "Cancelled",
+            "cancelled_at": "2026-07-02T00:00:00",
+        }
+    )
+
+    po_handler.save_new_po(
+        {
+            "po_number": "PO-D",
+            "po_status": "Issued",
+            "supplier": {"name": "Beta"},
+            "total_amount": 1200,
+        }
+    )
+
+    item = _item(po_table, "PO-D")
+    assert item["po_status"] == "Cancelled"
+    assert item["cancelled_at"] == "2026-07-02T00:00:00"
+    assert item["supplier"] == {"name": "Beta"}
+    assert item["total_amount"] == 1200
+
+
+def test_new_po_creates_fresh(po_handler, po_table):
+    """(e) new_po on a fresh PO creates it normally with its status."""
+    po_handler.save_new_po(
+        {"po_number": "PO-E", "po_status": "Issued", "total_amount": 42}
+    )
+
+    item = _item(po_table, "PO-E")
+    assert item["po_status"] == "Issued"
+    assert item["total_amount"] == 42
+
+
+def test_revision_creates_fresh_merge(po_handler, po_table):
+    """(e) revision on a fresh PO merges/creates normally with its status."""
+    po_handler.save_revision(
+        {"po_number": "PO-F", "po_status": "Revised", "total_amount": 77}
+    )
+
+    item = _item(po_table, "PO-F")
+    assert item["po_status"] == "Revised"
+    assert item["total_amount"] == 77
+
+
+def test_save_cancellation_is_authoritative(po_handler, po_table):
+    """Cancellation always wins: save_cancellation writes Cancelled over a live
+    status while leaving previously-established data intact."""
+    po_table.put_item(
+        Item={"po_number": "PO-G", "po_status": "Issued", "total_amount": 300}
+    )
+
+    po_handler.save_cancellation(
+        {
+            "po_number": "PO-G",
+            "processed_at": "2026-07-03T00:00:00",
+            "raw_s3_key": "s3://bucket/key",
+        }
+    )
+
+    item = _item(po_table, "PO-G")
+    assert item["po_status"] == "Cancelled"
+    assert item["cancelled_at"] == "2026-07-03T00:00:00"
+    assert item["total_amount"] == 300  # pre-existing data retained