From c9d83431e0dcd49454590be96982b165f6c1f604 Mon Sep 17 00:00:00 2001
From: Chaim <chaim@marcus-law.co.il>
Date: Wed, 24 Jun 2026 12:58:03 +0000
Subject: [PATCH] fix(arguments): validate claim_ids + per-row savepoint in
 argument_aggregator
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

aggregate_claims_to_arguments crashed with "current transaction is aborted,
commands ignored until end of transaction block" on large cases (confirmed on
1027-04-26, 195 claims; reported via CMP-186).

Root cause: the proposition INSERT (legal_argument_propositions) was wrapped in
a broad except Exception with no savepoint. When the LLM echoes a
syntactically-valid-but-nonexistent claim_id, the FK violation
(legal_argument_propositions_claim_id_fkey) puts the asyncpg transaction into
aborted state. The except caught only that INSERT's error but never issued
ROLLBACK TO SAVEPOINT, so the next statement (the following argument's INSERT
into legal_arguments) raised InFailedSQLTransactionError uncaught and crashed
the whole call. With many claims the bad-UUID probability is high -> consistent
failure.

Fix:
- Validate each claim_id against the known set of claim ids fetched for the
  case before INSERT, so a hallucinated id never reaches the DB (G1: fix at
  source). Malformed UUIDs are already dropped in _normalize_argument; this
  catches the valid-but-nonexistent ones.
- Wrap the INSERT in a per-row savepoint (async with conn.transaction()) as
  defense-in-depth, so any unexpected constraint failure rolls back to the
  savepoint instead of poisoning the surrounding transaction.

Invariants: G1 (normalize at source, not symptom-catch on read). No parallel
path (G2). No silent swallow — skips are logged.

TaskMaster: legal-ai #156

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 .../legal_mcp/services/argument_aggregator.py | 40 ++++++++++++++-----
 1 file changed, 31 insertions(+), 9 deletions(-)

diff --git a/mcp-server/src/legal_mcp/services/argument_aggregator.py b/mcp-server/src/legal_mcp/services/argument_aggregator.py
index 25c8b4f..a3bf220 100644
--- a/mcp-server/src/legal_mcp/services/argument_aggregator.py
+++ b/mcp-server/src/legal_mcp/services/argument_aggregator.py
@@ -230,6 +230,15 @@ async def aggregate_claims_to_arguments(
             party = "unknown"
         by_party.setdefault(party, []).append(dict(r))
 
+    # Valid claim_ids for this case == the ids of the claims we just fetched.
+    # The LLM is asked to echo back supporting claim_ids, but it may hallucinate
+    # a syntactically-valid-but-nonexistent UUID (malformed ones are already
+    # dropped in ``_normalize_argument``). Validating against this known set at
+    # source keeps a doomed INSERT — which would poison the surrounding asyncpg
+    # transaction (FK violation -> "current transaction is aborted") — out of
+    # the transaction entirely (G1: fix at source, not symptom).
+    valid_claim_ids: set[UUID] = {r["id"] for r in rows}
+
     party_counts: dict[str, int] = {}
     inserted = 0
     errors: list[str] = []
@@ -275,17 +284,30 @@ async def aggregate_claims_to_arguments(
                         arg["priority"],
                     )
                     for cid in arg["claim_ids"]:
-                        try:
-                            await conn.execute(
-                                """INSERT INTO legal_argument_propositions
-                                    (argument_id, claim_id)
-                                   VALUES ($1, $2)
-                                   ON CONFLICT DO NOTHING""",
-                                arg_id, cid,
+                        if cid not in valid_claim_ids:
+                            # Hallucinated claim_id that doesn't belong to this
+                            # case. Skip it rather than letting the FK violation
+                            # abort the whole transaction.
+                            logger.warning(
+                                "argument_aggregator: skipped unknown claim_id %s for arg %s",
+                                cid, arg_id,
                             )
+                            continue
+                        try:
+                            # Per-row savepoint: even after the validation above,
+                            # wrap the INSERT so any unexpected constraint failure
+                            # rolls back to the savepoint instead of poisoning the
+                            # surrounding transaction (asyncpg nests transaction()
+                            # as SAVEPOINT when already inside one).
+                            async with conn.transaction():
+                                await conn.execute(
+                                    """INSERT INTO legal_argument_propositions
+                                        (argument_id, claim_id)
+                                       VALUES ($1, $2)
+                                       ON CONFLICT DO NOTHING""",
+                                    arg_id, cid,
+                                )
                         except Exception as e:  # noqa: BLE001
-                            # Likely FK violation if the LLM hallucinated
-                            # a claim_id. Log and continue.
                             logger.warning(
                                 "argument_aggregator: skipped bad claim_id %s for arg %s: %s",
                                 cid, arg_id, e,