fix(sql): validate checkpoint run ownership

sql/absurd.sql

@@ -1470,6 +1470,7 @@ create function absurd.set_task_checkpoint_state (
 as $$
 declare
   v_now timestamptz := absurd.current_time();
+  v_run_task_id uuid;
   v_new_attempt integer;
   v_existing_attempt integer;
   v_existing_owner uuid;
@@ -1481,20 +1482,25 @@ begin
   end if;
 
   execute format(
-    'select r.attempt, r.state, t.state
+    'select r.task_id, r.attempt, r.state, t.state
        from absurd.%I r
        join absurd.%I t on t.task_id = r.task_id
-      where r.run_id = $1',
+      where r.run_id = $1
+      for update of r',
     'r_' || p_queue_name,
     't_' || p_queue_name
   )
-  into v_new_attempt, v_run_state, v_task_state
+  into v_run_task_id, v_new_attempt, v_run_state, v_task_state
   using p_owner_run;
 
-  if v_new_attempt is null then
+  if v_run_task_id is null then
     raise exception 'Run "%" not found for checkpoint', p_owner_run;
   end if;
 
+  if v_run_task_id <> p_task_id then
+    raise exception 'Run "%" does not belong to task "%" in queue "%"', p_owner_run, p_task_id, p_queue_name;
+  end if;
+
   if v_task_state = 'cancelled' then
     raise exception sqlstate 'AB001' using message = 'Task has been cancelled';
   end if;
@@ -1503,6 +1509,13 @@ begin
     raise exception sqlstate 'AB002' using message = format('Run "%s" has already failed in queue "%s"', p_owner_run, p_queue_name);
   end if;
 
+  if v_task_state in ('completed', 'failed') or v_run_state in ('completed', 'cancelled') then
+    raise exception 'Cannot checkpoint terminal run "%" or task "%" in queue "%"',
+      p_owner_run,
+      p_task_id,
+      p_queue_name;
+  end if;
+
   -- Extend the claim if requested
   if p_extend_claim_by is not null and p_extend_claim_by > 0 then
     execute format(

tests/test_checkpoints.py

@@ -115,6 +115,61 @@ def test_checkpoint_write_on_failed_run_raises_ab002(client):
     client.conn.execute(f"release savepoint {savepoint}")
 
 
+def test_checkpoint_write_requires_run_task_match(client):
+    queue = "checkpoint-run-match"
+    client.create_queue(queue)
+
+    first = client.spawn_task(queue, "first", {"value": 1})
+    second = client.spawn_task(queue, "second", {"value": 2})
+    claims = client.claim_tasks(queue, worker="worker-1", claim_timeout=60, qty=2)
+    claims_by_task = {claim["task_id"]: claim for claim in claims}
+    first_run = claims_by_task[first.task_id]["run_id"]
+    second_run = claims_by_task[second.task_id]["run_id"]
+
+    client.set_task_checkpoint_state(
+        queue, second.task_id, "step", {"value": "original"}, second_run, 60
+    )
+    first_run_before = client.get_run(queue, first_run)
+
+    savepoint = "checkpoint_run_task_match"
+    client.conn.execute(f"savepoint {savepoint}")
+    with pytest.raises(Exception, match="does not belong"):
+        client.set_task_checkpoint_state(
+            queue, second.task_id, "step", {"value": "wrong"}, first_run, 600
+        )
+    client.conn.execute(f"rollback to savepoint {savepoint}")
+    client.conn.execute(f"release savepoint {savepoint}")
+
+    assert client.get_checkpoint(queue, second.task_id, "step")["state"] == {
+        "value": "original"
+    }
+    assert client.get_run(queue, first_run) == first_run_before
+
+
+def test_checkpoint_write_rejects_terminal_task(client):
+    queue = "checkpoint-terminal"
+    client.create_queue(queue)
+
+    spawn = client.spawn_task(queue, "task", {"value": 1})
+    claim = client.claim_tasks(queue)[0]
+    client.complete_run(queue, claim["run_id"], {"value": "done"})
+
+    savepoint = "checkpoint_terminal"
+    client.conn.execute(f"savepoint {savepoint}")
+    with pytest.raises(Exception, match="Cannot checkpoint terminal"):
+        client.set_task_checkpoint_state(
+            queue,
+            spawn.task_id,
+            "too-late",
+            {"value": "orphan"},
+            claim["run_id"],
+        )
+    client.conn.execute(f"rollback to savepoint {savepoint}")
+    client.conn.execute(f"release savepoint {savepoint}")
+
+    assert client.get_checkpoint(queue, spawn.task_id, "too-late") is None
+
+
 def test_checkpoint_preloading_survives_retry(client):
     queue = "checkpoint-retry"
     client.create_queue(queue)