fix(worker): echo the lease fence on fail and the idempotency key on complete

jamjet-agent/src/main/java/dev/jamjet/agent/client/ClaimedWorkItem.java

@@ -12,6 +12,12 @@
  * completion: a stale fence yields HTTP 409. It is {@code null} only for legacy
  * unfenced claims.
  *
+ * <p>{@code idempotencyKey} (added by engine PR #128) is the key the engine derived
+ * for THIS node occurrence. Echo it on {@code complete} and the engine records the
+ * result against it, so a re-run replays instead of firing the tool a second time.
+ * Without it nothing lands in {@code tool_effects} and every replay re-fires.
+ * {@code null} against an engine that predates the field.
+ *
  * <p>For a {@code java_tool} item the {@code payload} carries the enriched
  * {@code class} / {@code method} / {@code input} the durable tool-worker dispatches on.
  */
@@ -22,5 +28,24 @@ public record ClaimedWorkItem(
         String queueType,
         Map<String, Object> payload,
         int attempt,
-        Long leaseFence
-) {}
+        Long leaseFence,
+        String idempotencyKey
+) {
+    /**
+     * The pre-{@code idempotencyKey} constructor, kept so adding the component stays
+     * source- and binary-compatible.
+     *
+     * <p>A record's canonical constructor is public API: widening it would break any
+     * caller that builds a {@code ClaimedWorkItem} directly, even though every accessor
+     * still resolves.
+     */
+    public ClaimedWorkItem(String id,
+                           String executionId,
+                           String nodeId,
+                           String queueType,
+                           Map<String, Object> payload,
+                           int attempt,
+                           Long leaseFence) {
+        this(id, executionId, nodeId, queueType, payload, attempt, leaseFence, null);
+    }
+}

jamjet-agent/src/main/java/dev/jamjet/agent/client/JamjetEngineClient.java

@@ -170,6 +170,28 @@ public void completeWorkItem(String itemId,
                                  String genAiModel,
                                  String finishReason,
                                  Long leaseFence) {
+        completeWorkItem(itemId, executionId, nodeId, output, statePatch, durationMs,
+                genAiModel, finishReason, leaseFence, null);
+    }
+
+    /**
+     * As {@link #completeWorkItem(String, String, String, Object, Map, long, String, String, Long)},
+     * additionally echoing the claim's {@code idempotency_key}.
+     *
+     * <p>The engine records the result against that key, so a re-run replays it instead
+     * of firing the tool again. Omitting it records no effect at all and every replay
+     * re-fires — the behaviour before engine PR #128.
+     */
+    public void completeWorkItem(String itemId,
+                                 String executionId,
+                                 String nodeId,
+                                 Object output,
+                                 Map<String, Object> statePatch,
+                                 long durationMs,
+                                 String genAiModel,
+                                 String finishReason,
+                                 Long leaseFence,
+                                 String idempotencyKey) {
         Map<String, Object> body = new LinkedHashMap<>();
         body.put("output", output);
         body.put("state_patch", statePatch == null ? Map.of() : statePatch);
@@ -189,12 +211,47 @@ public void completeWorkItem(String itemId,
         if (leaseFence != null) {
             body.put("lease_fence", leaseFence);
         }
+        // Sent whenever present, not whenever non-blank: the worker echoes what the
+        // claim handed it, and the engine rejects a malformed key (HTTP 400). Dropping
+        // it here would turn a bad key into "no key", which fails silently.
+        if (idempotencyKey != null) {
+            body.put("idempotency_key", idempotencyKey);
+        }
         execute(buildPost("/work-items/" + itemId + "/complete", body));
     }
 
-    /** {@code POST /work-items/{id}/fail}. Body {@code {error}}. */
+    /**
+     * {@code POST /work-items/{id}/fail} without a fence — the legacy unfenced path.
+     *
+     * @deprecated the engine settles the item but emits no {@code NodeFailed}, so the
+     *     scheduler fold keeps the node scheduled and the execution never reaches a
+     *     terminal state. Use {@link #failWorkItem(String, String, Long)} with the
+     *     claim's fence to get retry semantics.
+     */
+    @Deprecated
     public void failWorkItem(String itemId, String error) {
-        execute(buildPost("/work-items/" + itemId + "/fail", Map.of("error", error)));
+        failWorkItem(itemId, error, null);
+    }
+
+    /**
+     * {@code POST /work-items/{id}/fail}. Body {@code {error, lease_fence?}}.
+     *
+     * <p>With the fence the engine emits {@code NodeFailed} (or {@code RetryScheduled})
+     * so the node is rescheduled or dead-lettered. Without it the item is settled but
+     * the fold keeps the node scheduled and the workflow is stranded — every Java tool
+     * failure did this before engine PR #118.
+     *
+     * <p>A stale fence yields HTTP 409, surfaced as a {@link JamjetHttpException} with
+     * {@link JamjetHttpException#isConflict()} true: the lease was reclaimed and another
+     * worker owns the item, so the caller must treat it as a no-op rather than an error.
+     */
+    public void failWorkItem(String itemId, String error, Long leaseFence) {
+        Map<String, Object> body = new LinkedHashMap<>();
+        body.put("error", error);
+        if (leaseFence != null) {
+            body.put("lease_fence", leaseFence);
+        }
+        execute(buildPost("/work-items/" + itemId + "/fail", body));
     }
 
     /**

jamjet-agent/src/main/java/dev/jamjet/agent/worker/JavaToolWorker.java

@@ -176,18 +176,20 @@ private ItemResult process(ClaimedWorkItem item) {
         // RCE GATE (dispatch coordinate): the only java_fn coordinate the Agent builder
         // emits is the fixed ToolDispatcher. A payload naming any other class/method is
         // a forged/unknown node: fail cleanly, NEVER Class.forName a payload string.
+        // Read the fence before the gate below: a rejection is still a settle, and an
+        // unfenced settle emits no NodeFailed, so the node stays scheduled forever.
+        long fence = item.leaseFence() == null ? 0L : item.leaseFence();
+
         String cls = String.valueOf(payload.get("class"));
         String method = String.valueOf(payload.get("method"));
         if (!ToolDispatcher.DISPATCH_CLASS.equals(cls) || !ToolDispatcher.DISPATCH_METHOD.equals(method)) {
             String msg = "unsupported java_fn dispatch coordinate: " + cls + "#" + method
                     + " (only " + ToolDispatcher.DISPATCH_CLASS + "#" + ToolDispatcher.DISPATCH_METHOD + " is callable)";
             LOG.log(Level.WARNING, () -> "rejecting work item " + item.id() + ": " + msg);
-            client.failWorkItem(item.id(), msg);
-            return ItemResult.FAILED;
+            return fail(item, msg, fence);
         }
 
         Map<String, Object> input = asMap(payload.get("input"));
-        long fence = item.leaseFence() == null ? 0L : item.leaseFence();
         AtomicBoolean leaseLost = new AtomicBoolean(false);
 
         long startNanos = System.nanoTime();
@@ -220,8 +222,7 @@ private ItemResult process(ClaimedWorkItem item) {
                 Throwable cause = (e instanceof ExecutionException) ? e.getCause() : e;
                 String err = cause == null ? String.valueOf(e) : String.valueOf(cause.getMessage() != null ? cause.getMessage() : cause);
                 LOG.log(Level.WARNING, () -> "tool dispatch failed for item " + item.id() + ": " + err);
-                client.failWorkItem(item.id(), err);
-                return ItemResult.FAILED;
+                return fail(item, err, fence);
             }
 
             // Dispatch succeeded. If the lease was lost during/just after it, do NOT
@@ -238,6 +239,32 @@ private ItemResult process(ClaimedWorkItem item) {
         }
     }
 
+
+    /**
+     * Fail the item, threading the lease fence; a 409 is a lost-lease no-op.
+     *
+     * <p>Mirrors {@link #complete}: the fence makes the engine emit {@code NodeFailed},
+     * so the node is retried or dead-lettered instead of staying scheduled forever.
+     * A 409 means the lease was reclaimed and a NEW worker owns the item — reporting
+     * our failure would kill work that worker is running, so return quietly.
+     */
+    private ItemResult fail(ClaimedWorkItem item, String error, long fence) {
+        try {
+            client.failWorkItem(item.id(), error, fence == 0L ? null : fence);
+            return ItemResult.FAILED;
+        } catch (JamjetHttpException e) {
+            if (e.isConflict()) {
+                LOG.log(Level.INFO, () -> "failure rejected (409); lease lost for item " + item.id());
+                return ItemResult.LOST_LEASE;
+            }
+            // Anything else: we could not settle either, so the answer is the same —
+            // leave the item for lease expiry rather than pretending it failed cleanly.
+            // The two arms differ only in log level; both mean "the reclaimer has it".
+            LOG.log(Level.WARNING, () -> "could not report failure for item " + item.id() + ": " + e.getMessage());
+            return ItemResult.LOST_LEASE;
+        }
+    }
+
     /** Settle the item, threading the lease fence; a 409 is a lost-lease no-op. */
     private ItemResult complete(ClaimedWorkItem item, Map<String, Object> output, long durationMs, long fence) {
         // The dispatcher return ({"messages": [...]}) is BOTH the node output and the
@@ -251,7 +278,10 @@ private ItemResult complete(ClaimedWorkItem item, Map<String, Object> output, lo
                     item.id(), item.executionId(), item.nodeId(),
                     output, output, durationMs, genAiModel, finishReason,
                     // Echo the claim's fence so the engine fences this completion.
-                    fence == 0L ? null : fence);
+                    fence == 0L ? null : fence,
+                    // Echo the claim's key so the engine records the effect against it
+                    // and a re-run replays instead of firing the tool a second time.
+                    item.idempotencyKey());
             LOG.log(Level.DEBUG, () -> "completed item " + item.id() + " in " + durationMs + "ms");
             return ItemResult.COMPLETED;
         } catch (JamjetHttpException e) {
@@ -264,8 +294,7 @@ private ItemResult complete(ClaimedWorkItem item, Map<String, Object> output, lo
             }
             // Any other completion error keeps the fail behavior (mirror Python re-raise).
             LOG.log(Level.WARNING, () -> "completion failed for item " + item.id() + ": " + e.getMessage());
-            client.failWorkItem(item.id(), "complete failed: " + e.getMessage());
-            return ItemResult.FAILED;
+            return fail(item, "complete failed: " + e.getMessage(), fence);
         }
     }
 

jamjet-agent/src/test/java/dev/jamjet/agent/client/JamjetEngineClientTest.java

@@ -220,4 +220,22 @@ void staleFenceCompleteSurfacesAsConflict() {
         assertThat(ex.isConflict()).isTrue();
         assertThat(ex.body()).contains("stale");
     }
+
+    @Test
+    void thePreIdempotencyKeyConstructorStillCompiles() {
+        // A record's canonical constructor is public API. Adding `idempotencyKey` as a
+        // component would have widened it and broken every caller that builds one
+        // directly — accessors resolving is not the same as the constructor resolving.
+        //
+        // This call IS the compatibility promise: delete the 7-arg overload and this
+        // test stops compiling.
+        ClaimedWorkItem legacy = new ClaimedWorkItem(
+                "wi_1", "ex_1", "n1", "java_tool", Map.of(), 1, 5L);
+
+        assertThat(legacy.idempotencyKey())
+                .as("an item built the old way has no key, and must not invent one")
+                .isNull();
+        assertThat(legacy.leaseFence()).isEqualTo(5L);
+    }
+
 }

jamjet-agent/src/test/java/dev/jamjet/agent/worker/JavaToolWorkerTest.java

@@ -22,6 +22,7 @@
 import java.util.function.BooleanSupplier;
 
 import static com.github.tomakehurst.wiremock.client.WireMock.aResponse;
+import static com.github.tomakehurst.wiremock.client.WireMock.absent;
 import static com.github.tomakehurst.wiremock.client.WireMock.containing;
 import static com.github.tomakehurst.wiremock.client.WireMock.equalTo;
 import static com.github.tomakehurst.wiremock.client.WireMock.matchingJsonPath;
@@ -97,6 +98,11 @@ private static Map<String, Object> toolCall(String id, String name, Map<String,
     /** Serialize a {@code {claimed, work_item}} claim response with the given payload + fence. */
     private static String claimBody(String dispatchClass, String dispatchMethod,
                                     Map<String, Object> input, long fence) {
+        return claimBody(dispatchClass, dispatchMethod, input, fence, null);
+    }
+
+    private static String claimBody(String dispatchClass, String dispatchMethod,
+                                    Map<String, Object> input, long fence, String idempotencyKey) {
         Map<String, Object> payload = new LinkedHashMap<>();
         payload.put("class", dispatchClass);
         payload.put("method", dispatchMethod);
@@ -110,6 +116,9 @@ private static String claimBody(String dispatchClass, String dispatchMethod,
         wi.put("payload", payload);
         wi.put("attempt", 1);
         wi.put("lease_fence", fence);
+        if (idempotencyKey != null) {
+            wi.put("idempotency_key", idempotencyKey);
+        }
 
         Map<String, Object> resp = new LinkedHashMap<>();
         resp.put("claimed", true);
@@ -347,4 +356,96 @@ void emptyQueueReturnsEmpty() {
             assertThat(worker.runOnce()).isEqualTo(JavaToolWorker.ItemResult.EMPTY);
         }
     }
+
+    @Test
+    @Timeout(15)
+    void aFailureReportsTheFenceSoTheNodeIsRescheduled() {
+        // Without the fence the engine settles the item but emits no NodeFailed, so the
+        // scheduler fold keeps the node `scheduled` and the execution never terminates.
+        // Every Java tool failure stranded its workflow this way.
+        Map<String, Object> input = Map.of("last_model_tool_calls", List.of());
+        wm.stubFor(post(urlEqualTo("/work-items/claim"))
+                .willReturn(okJson(claimBody("java.lang.Runtime", "exec", input, 7))));
+        wm.stubFor(post(urlEqualTo("/work-items/wi_1/fail")).willReturn(ok()));
+
+        try (var client = new JamjetEngineClient(wm.baseUrl());
+             var worker = new JavaToolWorker(client, "w1",
+                     ToolRegistry.of(new TestTools.WebSearchTool()),
+                     Duration.ofSeconds(2), Duration.ofMillis(10))) {
+            assertThat(worker.runOnce()).isEqualTo(JavaToolWorker.ItemResult.FAILED);
+        }
+
+        wm.verify(postRequestedFor(urlEqualTo("/work-items/wi_1/fail"))
+                .withRequestBody(matchingJsonPath("$.lease_fence", equalTo("7"))));
+    }
+
+    @Test
+    @Timeout(15)
+    void aConflictOnFailDoesNotEscapeAndKillTheWorker() {
+        // 409 means the lease was reclaimed and a NEW worker owns the item. Reporting
+        // our failure would kill work that worker is running, so it must be a quiet
+        // no-op rather than an exception out of runOnce().
+        //
+        // This pins the escape, not the 409/other distinction: a failure we could not
+        // report is LOST_LEASE either way, because in both cases we did not settle the
+        // item and the reclaimer must have it. Only the log level differs.
+        Map<String, Object> input = Map.of("last_model_tool_calls", List.of());
+        wm.stubFor(post(urlEqualTo("/work-items/claim"))
+                .willReturn(okJson(claimBody("java.lang.Runtime", "exec", input, 7))));
+        wm.stubFor(post(urlEqualTo("/work-items/wi_1/fail"))
+                .willReturn(aResponse().withStatus(409).withBody("{\"reason\":\"stale or invalid lease fence\"}")));
+
+        try (var client = new JamjetEngineClient(wm.baseUrl());
+             var worker = new JavaToolWorker(client, "w1",
+                     ToolRegistry.of(new TestTools.WebSearchTool()),
+                     Duration.ofSeconds(2), Duration.ofMillis(10))) {
+            assertThat(worker.runOnce()).isEqualTo(JavaToolWorker.ItemResult.LOST_LEASE);
+        }
+    }
+
+    @Test
+    @Timeout(15)
+    void theClaimsIdempotencyKeyIsEchoedOnComplete() {
+        // The engine records the result against this key, so a re-run replays it instead
+        // of firing the tool a second time. Omit it and nothing lands in tool_effects.
+        Map<String, Object> input = Map.of(
+                "last_model_tool_calls", List.of(toolCall("tc1", "web_search", Map.of("query", "x"))));
+        wm.stubFor(post(urlEqualTo("/work-items/claim")).willReturn(okJson(
+                claimBody(ToolDispatcher.DISPATCH_CLASS, ToolDispatcher.DISPATCH_METHOD, input, 9, "key-abc"))));
+        wm.stubFor(post(urlEqualTo("/work-items/wi_1/complete")).willReturn(ok()));
+        wm.stubFor(post(urlEqualTo("/work-items/wi_1/heartbeat")).willReturn(ok()));
+
+        try (var client = new JamjetEngineClient(wm.baseUrl());
+             var worker = new JavaToolWorker(client, "w1",
+                     ToolRegistry.of(new TestTools.WebSearchTool()),
+                     Duration.ofSeconds(2), Duration.ofMillis(10))) {
+            assertThat(worker.runOnce()).isEqualTo(JavaToolWorker.ItemResult.COMPLETED);
+        }
+
+        wm.verify(postRequestedFor(urlEqualTo("/work-items/wi_1/complete"))
+                .withRequestBody(matchingJsonPath("$.idempotency_key", equalTo("key-abc"))));
+    }
+
+    @Test
+    @Timeout(15)
+    void aClaimWithoutAKeyStillCompletes() {
+        // Against an engine that predates the field the key is absent. That must record
+        // no effect, not break the completion.
+        Map<String, Object> input = Map.of(
+                "last_model_tool_calls", List.of(toolCall("tc1", "web_search", Map.of("query", "x"))));
+        wm.stubFor(post(urlEqualTo("/work-items/claim")).willReturn(okJson(validClaim(input, 9))));
+        wm.stubFor(post(urlEqualTo("/work-items/wi_1/complete")).willReturn(ok()));
+        wm.stubFor(post(urlEqualTo("/work-items/wi_1/heartbeat")).willReturn(ok()));
+
+        try (var client = new JamjetEngineClient(wm.baseUrl());
+             var worker = new JavaToolWorker(client, "w1",
+                     ToolRegistry.of(new TestTools.WebSearchTool()),
+                     Duration.ofSeconds(2), Duration.ofMillis(10))) {
+            assertThat(worker.runOnce()).isEqualTo(JavaToolWorker.ItemResult.COMPLETED);
+        }
+
+        wm.verify(postRequestedFor(urlEqualTo("/work-items/wi_1/complete"))
+                .withRequestBody(matchingJsonPath("$.idempotency_key", absent())));
+    }
+
 }