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,