diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/client/ClaimedWorkItem.java b/jamjet-agent/src/main/java/dev/jamjet/agent/client/ClaimedWorkItem.java index c642c04..07fb5b1 100644 --- a/jamjet-agent/src/main/java/dev/jamjet/agent/client/ClaimedWorkItem.java +++ b/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. * + *

{@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. + * *

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 payload, int attempt, - Long leaseFence -) {} + Long leaseFence, + String idempotencyKey +) { + /** + * The pre-{@code idempotencyKey} constructor, kept so adding the component stays + * source- and binary-compatible. + * + *

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 payload, + int attempt, + Long leaseFence) { + this(id, executionId, nodeId, queueType, payload, attempt, leaseFence, null); + } +} diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/client/JamjetEngineClient.java b/jamjet-agent/src/main/java/dev/jamjet/agent/client/JamjetEngineClient.java index 95f0afe..b514ffe 100644 --- a/jamjet-agent/src/main/java/dev/jamjet/agent/client/JamjetEngineClient.java +++ b/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}. + * + *

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 statePatch, + long durationMs, + String genAiModel, + String finishReason, + Long leaseFence, + String idempotencyKey) { Map 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?}}. + * + *

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. + * + *

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 body = new LinkedHashMap<>(); + body.put("error", error); + if (leaseFence != null) { + body.put("lease_fence", leaseFence); + } + execute(buildPost("/work-items/" + itemId + "/fail", body)); } /** diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/worker/JavaToolWorker.java b/jamjet-agent/src/main/java/dev/jamjet/agent/worker/JavaToolWorker.java index fef8c06..3ed734b 100644 --- a/jamjet-agent/src/main/java/dev/jamjet/agent/worker/JavaToolWorker.java +++ b/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 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. + * + *

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 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 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 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); } } diff --git a/jamjet-agent/src/test/java/dev/jamjet/agent/client/JamjetEngineClientTest.java b/jamjet-agent/src/test/java/dev/jamjet/agent/client/JamjetEngineClientTest.java index 7124e88..bc5e9f5 100644 --- a/jamjet-agent/src/test/java/dev/jamjet/agent/client/JamjetEngineClientTest.java +++ b/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); + } + } diff --git a/jamjet-agent/src/test/java/dev/jamjet/agent/worker/JavaToolWorkerTest.java b/jamjet-agent/src/test/java/dev/jamjet/agent/worker/JavaToolWorkerTest.java index 1b151ca..356043b 100644 --- a/jamjet-agent/src/test/java/dev/jamjet/agent/worker/JavaToolWorkerTest.java +++ b/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 toolCall(String id, String name, Map input, long fence) { + return claimBody(dispatchClass, dispatchMethod, input, fence, null); + } + + private static String claimBody(String dispatchClass, String dispatchMethod, + Map input, long fence, String idempotencyKey) { Map 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 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 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 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 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 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()))); + } + }