fix(durable): enforce lifecycle progress invariants

pg_llm_batch/db.py

@@ -384,6 +384,14 @@ def _normalize_remote_batch_snapshot(
     )
     counts_value = provider_batch.get("request_counts")
     request_counts = counts_value if isinstance(counts_value, Mapping) else {}
+    total_requests = _provider_count(request_counts.get("total"))
+    completed_requests = _provider_count(request_counts.get("completed"))
+    failed_requests = _provider_count(request_counts.get("failed"))
+    if (
+        "total" in request_counts
+        and completed_requests + failed_requests > total_requests
+    ):
+        raise ValueError("request_counts progress is inconsistent")
     provider_metadata, metadata_json = _provider_metadata(
         provider_batch.get("metadata")
     )
@@ -400,9 +408,9 @@ def _normalize_remote_batch_snapshot(
         "batch_status": batch_status,
         "output_file_id": output_file_id,
         "error_file_id": error_file_id,
-        "total_requests": _provider_count(request_counts.get("total")),
-        "completed_requests": _provider_count(request_counts.get("completed")),
-        "failed_requests": _provider_count(request_counts.get("failed")),
+        "total_requests": total_requests,
+        "completed_requests": completed_requests,
+        "failed_requests": failed_requests,
         "provider_metadata": provider_metadata,
         "observed_at": observed,
         "terminal_at": terminal_at,
@@ -496,6 +504,16 @@ def _persist_remote_batch_state(
             ),
             updated_at = EXCLUDED.updated_at
         WHERE EXCLUDED.observation_order > llm_remote_batch_jobs.observation_order
+          AND GREATEST(
+              llm_remote_batch_jobs.completed_requests,
+              EXCLUDED.completed_requests
+          ) + GREATEST(
+              llm_remote_batch_jobs.failed_requests,
+              EXCLUDED.failed_requests
+          ) <= GREATEST(
+              llm_remote_batch_jobs.total_requests,
+              EXCLUDED.total_requests
+          )
           AND (
               llm_remote_batch_jobs.batch_status NOT IN (
                   'completed', 'failed', 'expired', 'cancelled'

tests/test_remote_batch_progress_invariant.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Regression tests for durable lifecycle progress invariants."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import pytest
+
+from pg_llm_batch import db
+
+
+class _Cursor:
+    """Capture lifecycle SQL without contacting PostgreSQL."""
+
+    def __init__(self, driver: "_Psycopg") -> None:
+        self.driver = driver
+
+    def __enter__(self) -> "_Cursor":
+        return self
+
+    def __exit__(self, *_exc: Any) -> None:
+        return None
+
+    def execute(self, sql: str, params: Any = None) -> None:
+        """Record one SQL statement and its bound parameters."""
+        self.driver.executions.append((sql, params))
+
+
+class _Connection:
+    """Expose the minimal connection surface used by lifecycle persistence."""
+
+    def __init__(self, driver: "_Psycopg") -> None:
+        self.driver = driver
+
+    def __enter__(self) -> "_Connection":
+        return self
+
+    def __exit__(self, *_exc: Any) -> None:
+        return None
+
+    def cursor(self) -> _Cursor:
+        """Return one SQL-capturing cursor."""
+        return _Cursor(self.driver)
+
+    def commit(self) -> None:
+        """Record an explicit transaction commit."""
+        self.driver.commits += 1
+
+
+class _Psycopg:
+    """Minimal deterministic psycopg replacement for boundary tests."""
+
+    def __init__(self) -> None:
+        self.executions: list[tuple[str, Any]] = []
+        self.connections: list[str] = []
+        self.commits = 0
+
+    def connect(self, dsn: str) -> _Connection:
+        """Record the selected target before returning a fake connection."""
+        self.connections.append(dsn)
+        return _Connection(self)
+
+
+def _compact_sql(sql: str) -> str:
+    """Normalize SQL whitespace for semantic contract assertions."""
+    return " ".join(sql.split())
+
+
+def test_persistence_rejects_impossible_same_observation_before_database(
+    monkeypatch: Any,
+) -> None:
+    """Completed plus failed requests cannot exceed the observed total."""
+    driver = _Psycopg()
+    monkeypatch.setattr(db, "psycopg", driver)
+
+    with pytest.raises(ValueError, match="request_counts progress is inconsistent"):
+        db.persist_remote_batch_state(
+            "postgresql://example",
+            "primary",
+            {
+                "id": "batch-1",
+                "status": "in_progress",
+                "request_counts": {"total": 1, "completed": 1, "failed": 1},
+            },
+            observation_order=1,
+        )
+
+    assert driver.connections == []
+    assert driver.executions == []
+    assert driver.commits == 0
+
+
+def test_upsert_guards_combined_monotonic_progress_invariant(monkeypatch: Any) -> None:
+    """Independent monotonic counter updates cannot create impossible progress."""
+    driver = _Psycopg()
+    monkeypatch.setattr(db, "psycopg", driver)
+
+    db.persist_remote_batch_state(
+        "postgresql://example",
+        "primary",
+        {
+            "id": "batch-1",
+            "status": "in_progress",
+            "request_counts": {"total": 10, "completed": 9, "failed": 0},
+        },
+        observation_order=1,
+    )
+
+    sql = _compact_sql(driver.executions[1][0])
+    expected_guard = (
+        "GREATEST( llm_remote_batch_jobs.completed_requests, "
+        "EXCLUDED.completed_requests ) + GREATEST( "
+        "llm_remote_batch_jobs.failed_requests, EXCLUDED.failed_requests ) "
+        "<= GREATEST( llm_remote_batch_jobs.total_requests, "
+        "EXCLUDED.total_requests )"
+    )
+    assert expected_guard in sql