fix(worker): treat 409 complete-conflict as lost-lease + atomic memory settle + honest absent-fence test

Author: sunilp

runtime/state/src/memory.rs

@@ -446,16 +446,19 @@ impl StateBackend for InMemoryBackend {
         //   2. A stale fence after a reclaim — `reclaim_expired_leases` clears
         //      `worker_id` but LEAVES `lease_fence` set, so the old fence lingers.
         // Either case now returns `Ok(false)`. A mismatched fence, an absent
-        // item, or one already settled (removed) also returns `false`. The
-        // get-then-remove is two DashMap ops rather than one atomic step; that is
-        // acceptable for the in-memory dev/test backend (the SQLite backends are
-        // the production path where atomicity is guaranteed by the single UPDATE).
-        match self.work_items.get(&item_id) {
-            Some(entry) if entry.worker_id.is_some() && entry.lease_fence == lease_fence => {}
-            _ => return Ok(false),
-        }
-        self.work_items.remove(&item_id);
-        Ok(true)
+        // item, or one already settled (removed) also returns `false`.
+        //
+        // `remove_if` makes the guard check and the delete a SINGLE atomic step
+        // (it holds the shard lock across the predicate and removal), closing the
+        // get-then-remove TOCTOU where two concurrent settles could both observe a
+        // claimed+matching item and both remove it. It returns `Some` iff the
+        // predicate held and the entry was removed, so exactly one racer settles.
+        Ok(self
+            .work_items
+            .remove_if(&item_id, |_k, entry| {
+                entry.worker_id.is_some() && entry.lease_fence == lease_fence
+            })
+            .is_some())
     }
 
     async fn commit_turn(

runtime/state/tests/memory_backend.rs

@@ -345,6 +345,59 @@ async fn fenced_complete_settles_claimed_item_with_matching_fence() {
     );
 }
 
+/// Atomicity: many concurrent fenced completes of the SAME claimed item with the
+/// SAME matching fence must yield EXACTLY ONE `true`. `complete_work_item_fenced`
+/// uses DashMap's `remove_if`, so the claimed+fence-match check and the delete are
+/// one atomic step — the get-then-remove TOCTOU where two racers both observe the
+/// item and both remove it (double-settle) cannot happen.
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn fenced_complete_is_atomic_under_concurrency() {
+    use std::sync::Arc;
+    use tokio::sync::Barrier;
+
+    let backend = Arc::new(InMemoryBackend::new());
+    let exec_id = ExecutionId::new();
+    let item_id = backend
+        .enqueue_work_item(pending_item(&exec_id))
+        .await
+        .unwrap();
+    let claimed = backend
+        .claim_work_item("w1", &["default"])
+        .await
+        .unwrap()
+        .unwrap();
+    let fence = claimed.lease_fence;
+    assert!(fence > 0);
+
+    // Release all racers simultaneously (the barrier maximizes the overlap on the
+    // single contended item) and have each present the matching fence.
+    const RACERS: usize = 16;
+    let barrier = Arc::new(Barrier::new(RACERS));
+    let mut handles = Vec::with_capacity(RACERS);
+    for _ in 0..RACERS {
+        let backend = Arc::clone(&backend);
+        let barrier = Arc::clone(&barrier);
+        handles.push(tokio::spawn(async move {
+            barrier.wait().await;
+            backend
+                .complete_work_item_fenced(item_id, fence)
+                .await
+                .unwrap()
+        }));
+    }
+
+    let mut settled = 0usize;
+    for h in handles {
+        if h.await.unwrap() {
+            settled += 1;
+        }
+    }
+    assert_eq!(
+        settled, 1,
+        "exactly one concurrent fenced complete may settle the item (atomic check-and-delete)"
+    );
+}
+
 #[tokio::test]
 async fn test_patch_append_array() {
     let backend = InMemoryBackend::new();

sdk/python/jamjet/cli/main.py

@@ -1850,20 +1850,36 @@ async def _heartbeat_loop() -> None:
                     gen_ai_model_field = output.get("gen_ai_model") or None
                     finish_reason_field = output.get("finish_reason") or None
 
-                await client.complete_work_item(
-                    item_id,
-                    exec_id,
-                    node_id,
-                    output,
-                    state_patch,
-                    duration_ms,
-                    gen_ai_model=gen_ai_model_field,
-                    finish_reason=finish_reason_field,
-                    # Echo the lease fence from the claim so the runtime can fence
-                    # the completion (reject a stale/reclaimed lease). Omitted when
-                    # 0/absent, keeping the unfenced fallback backward-compatible.
-                    lease_fence=lease_fence,
-                )
+                try:
+                    await client.complete_work_item(
+                        item_id,
+                        exec_id,
+                        node_id,
+                        output,
+                        state_patch,
+                        duration_ms,
+                        gen_ai_model=gen_ai_model_field,
+                        finish_reason=finish_reason_field,
+                        # Echo the lease fence from the claim so the runtime can fence
+                        # the completion (reject a stale/reclaimed lease). Omitted when
+                        # 0/absent, keeping the unfenced fallback backward-compatible.
+                        lease_fence=lease_fence,
+                    )
+                except Exception as complete_exc:
+                    # A 409 means our echoed fence no longer matches: the lease was
+                    # reclaimed and a NEW worker owns this item now. This (stale)
+                    # worker must NOT fail_work_item — that would clobber the
+                    # reclaimed work the new claimant is running and defeat the fence.
+                    # Treat the lost lease as a no-op and move on. httpx surfaces the
+                    # rejection as HTTPStatusError carrying .response.status_code.
+                    status = getattr(getattr(complete_exc, "response", None), "status_code", None)
+                    if status == 409:
+                        console.print(f"[yellow]Completion rejected; lease lost[/yellow] id={item_id}")
+                        if once:
+                            return
+                        continue
+                    # Any other completion error keeps the existing fail behavior.
+                    raise
                 console.print(f"[green]Completed[/green] id={item_id} node=[bold]{node_id}[/bold] {duration_ms}ms")
             except Exception as exc:
                 console.print(f"[red]Failed[/red] node={node_id}: {exc}")

sdk/python/tests/test_worker.py

@@ -13,6 +13,8 @@
 
 from __future__ import annotations
 
+import httpx
+
 from jamjet.cli.main import _worker_loop  # noqa: E402
 
 # ── Test handler functions ────────────────────────────────────────────────────
@@ -97,6 +99,30 @@ async def heartbeat_work_item(self, item_id: str, worker_id: str, lease_fence: i
         self.heartbeat_calls.append({"item_id": item_id, "lease_fence": lease_fence})
 
 
+def _http_status_error(status: int) -> httpx.HTTPStatusError:
+    """Build the exact exception JamjetClient.complete_work_item raises on a non-2xx
+    response: httpx's HTTPStatusError carrying `.response.status_code` (the engine
+    returns 409 when the echoed lease fence no longer matches)."""
+    req = httpx.Request("POST", "http://runtime/work-items/wi/complete")
+    resp = httpx.Response(status, request=req)
+    return httpx.HTTPStatusError(f"HTTP {status}", request=req, response=resp)
+
+
+class _Conflict409Client(_StubClient):
+    """Stub whose complete_work_item raises a 409 — i.e. the lease was reclaimed and
+    a NEW worker now owns the item; this (stale) worker's settle is rejected."""
+
+    async def complete_work_item(self, *args: object, **kwargs: object) -> None:
+        raise _http_status_error(409)
+
+
+class _ServerError500Client(_StubClient):
+    """Stub whose complete_work_item raises a NON-409 error (transient server fault)."""
+
+    async def complete_work_item(self, *args: object, **kwargs: object) -> None:
+        raise _http_status_error(500)
+
+
 # ── Fixtures / constants ──────────────────────────────────────────────────────
 
 _ADD_ITEM: dict = {
@@ -275,17 +301,46 @@ async def test_worker_echoes_lease_fence_on_complete() -> None:
 
 
 async def test_worker_absent_fence_forwarded_as_zero() -> None:
-    """An item with no real fence (lease_fence=0) forwards 0; the client omits it,
-    keeping the unfenced fallback backward-compatible."""
-    stub = _StubClient(claimed_item=_ADD_ITEM)
+    """When the claim response OMITS the lease_fence key ENTIRELY (not even an
+    explicit 0), the worker must still forward 0 — the client then drops it, keeping
+    the unfenced fallback backward-compatible.
+
+    Regression: this previously claimed an item built from _ADD_ITEM, which already
+    carries an explicit lease_fence=0, so it only exercised the explicit-0 path and
+    never the missing-key path it is named for. Build an item without the key."""
+    no_fence_item = {k: v for k, v in _ADD_ITEM.items() if k != "lease_fence"}
+    assert "lease_fence" not in no_fence_item, "the key must be truly absent from the claim"
+    stub = _StubClient(claimed_item=no_fence_item)
     await _worker_loop(stub, "test-worker", ["python_tool"], once=True)
 
     assert len(stub.complete_calls) == 1
     assert stub.complete_calls[0]["lease_fence"] == 0, (
-        "absent fence is forwarded as 0 and dropped by the client (unfenced path)"
+        "an absent fence is forwarded as 0 and dropped by the client (unfenced path)"
     )
 
 
+async def test_worker_409_complete_is_lost_lease_noop() -> None:
+    """A 409 from complete_work_item means our lease fence no longer matches: the
+    item was reclaimed and a NEW worker owns it now. The stale worker must treat the
+    rejection as a lost-lease no-op — it must NOT call fail_work_item, because that
+    would clobber the reclaimed work the new claimant is running. With --once it
+    returns cleanly."""
+    stub = _Conflict409Client(claimed_item=_FENCED_ITEM)
+    await _worker_loop(stub, "test-worker", ["python_tool"], once=True)
+
+    assert len(stub.fail_calls) == 0, "a 409 lost-lease must NOT call fail_work_item"
+
+
+async def test_worker_non_409_complete_error_still_fails() -> None:
+    """A NON-409 completion error keeps the existing behavior: it falls through to
+    fail_work_item (the lease is still ours; the settle just did not land)."""
+    stub = _ServerError500Client(claimed_item=_FENCED_ITEM)
+    await _worker_loop(stub, "test-worker", ["python_tool"], once=True)
+
+    assert len(stub.fail_calls) == 1, "a non-409 completion error must still call fail_work_item"
+    assert stub.fail_calls[0]["item_id"] == "wi-006"
+
+
 async def test_worker_dispatch_tool_calls_return_reaches_state() -> None:
     """2j-4 G2: the agent loop's dispatch_tool_calls runs on the worker and its
     {"messages": [...]} return is posted as the state_patch — so the accumulated