From 1cb01a0093f2218b2f23569bf76b89a5d393f68b Mon Sep 17 00:00:00 2001 From: sdev Date: Sun, 28 Jun 2026 23:23:39 +0530 Subject: [PATCH 1/7] feat(agent): bare-route engine client with work-item methods + WireMock round-trip New jamjet-agent module with JamjetEngineClient: JDK java.net.http on a virtual-thread executor, JamjetJson snake_case bodies, BARE engine routes (no /api/v1). Mirrors the canonical Python jamjet.client shapes. - createWorkflow/startExecution/getExecution/listEvents - claimWorkItem/completeWorkItem/failWorkItem/heartbeatWorkItem - lease_fence threaded on heartbeat and complete (engine PR #108); a stale fence on complete surfaces as JamjetHttpException with isConflict()==true so a reclaimed worker can detect a lost lease - startExecution uses the workflow_version body key (matches client.py) - complete includes lease_fence only when non-null Hermetic WireMock test proves the wire shape: create -> start -> poll to Completed, plus the work-item path (claim returns lease_fence, complete and heartbeat post it, null fence is omitted, stubbed 409 surfaces as conflict). --- jamjet-agent/pom.xml | 65 ++++ .../jamjet/agent/client/ClaimedWorkItem.java | 26 ++ .../agent/client/CreateWorkflowResult.java | 8 + .../jamjet/agent/client/ExecutionState.java | 16 + .../agent/client/JamjetEngineClient.java | 313 ++++++++++++++++++ .../agent/client/JamjetHttpException.java | 49 +++ .../agent/client/StartExecutionResult.java | 7 + .../agent/client/JamjetEngineClientTest.java | 207 ++++++++++++ pom.xml | 1 + 9 files changed, 692 insertions(+) create mode 100644 jamjet-agent/pom.xml create mode 100644 jamjet-agent/src/main/java/dev/jamjet/agent/client/ClaimedWorkItem.java create mode 100644 jamjet-agent/src/main/java/dev/jamjet/agent/client/CreateWorkflowResult.java create mode 100644 jamjet-agent/src/main/java/dev/jamjet/agent/client/ExecutionState.java create mode 100644 jamjet-agent/src/main/java/dev/jamjet/agent/client/JamjetEngineClient.java create mode 100644 jamjet-agent/src/main/java/dev/jamjet/agent/client/JamjetHttpException.java create mode 100644 jamjet-agent/src/main/java/dev/jamjet/agent/client/StartExecutionResult.java create mode 100644 jamjet-agent/src/test/java/dev/jamjet/agent/client/JamjetEngineClientTest.java diff --git a/jamjet-agent/pom.xml b/jamjet-agent/pom.xml new file mode 100644 index 0000000..b17f213 --- /dev/null +++ b/jamjet-agent/pom.xml @@ -0,0 +1,65 @@ + + + 4.0.0 + + + dev.jamjet + jamjet-runtime-java-parent + 0.3.1 + + + jamjet-agent + jar + + JamJet Agent (Java) + + Idiomatic Java agent authoring on the governed durable JamJet engine: a thin + Agent builder that compiles to the shared agent-loop WorkflowIr plus a durable + Java tool-worker. Plain Java 21 (virtual threads), framework-free. This module + owns the bare-route engine HTTP client (workflows, executions, work-items). + + + + + + dev.jamjet + jamjet-runtime-core + ${project.version} + + + + + com.fasterxml.jackson.core + jackson-databind + + + + + + + org.junit.jupiter + junit-jupiter + test + + + org.assertj + assertj-core + test + + + + org.wiremock + wiremock-standalone + 3.9.2 + test + + + + ch.qos.logback + logback-classic + test + + + 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 new file mode 100644 index 0000000..c642c04 --- /dev/null +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/client/ClaimedWorkItem.java @@ -0,0 +1,26 @@ +package dev.jamjet.agent.client; + +import java.util.Map; + +/** + * A work item claimed from the engine via {@code POST /work-items/claim}. + * Mirrors the {@code work_item} object the Python client returns + * ({@code {id, execution_id, node_id, queue_type, payload, attempt, lease_fence}}). + * + *

{@code leaseFence} (added by engine PR #108) MUST be threaded back on both + * {@code heartbeat} and {@code complete} so a reclaimed worker cannot duplicate a + * completion: a stale fence yields HTTP 409. It is {@code null} only for legacy + * unfenced claims. + * + *

For a {@code java_tool} item the {@code payload} carries the enriched + * {@code class} / {@code method} / {@code input} the durable tool-worker dispatches on. + */ +public record ClaimedWorkItem( + String id, + String executionId, + String nodeId, + String queueType, + Map payload, + int attempt, + Long leaseFence +) {} diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/client/CreateWorkflowResult.java b/jamjet-agent/src/main/java/dev/jamjet/agent/client/CreateWorkflowResult.java new file mode 100644 index 0000000..d335eda --- /dev/null +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/client/CreateWorkflowResult.java @@ -0,0 +1,8 @@ +package dev.jamjet.agent.client; + +/** + * Response of {@code POST /workflows}. Mirrors the Python client's + * {@code create_workflow} return ({@code {"workflow_id": ...}}); {@code version} + * is populated when the engine echoes it. + */ +public record CreateWorkflowResult(String workflowId, String version) {} diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/client/ExecutionState.java b/jamjet-agent/src/main/java/dev/jamjet/agent/client/ExecutionState.java new file mode 100644 index 0000000..7d2dfab --- /dev/null +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/client/ExecutionState.java @@ -0,0 +1,16 @@ +package dev.jamjet.agent.client; + +import java.util.Map; + +/** + * Response of {@code GET /executions/{id}} — the polled execution snapshot. + * Mirrors the Python client's {@code get_execution} shape ({@code {status, + * current_state, ...}}). The terminal status is {@code "Completed"} (or + * {@code "Failed"} / {@code "Cancelled"}); {@code currentState} carries the + * accumulated execution state. + */ +public record ExecutionState( + String executionId, + String status, + Map currentState +) {} 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 new file mode 100644 index 0000000..60b6480 --- /dev/null +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/client/JamjetEngineClient.java @@ -0,0 +1,313 @@ +package dev.jamjet.agent.client; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import dev.jamjet.runtime.core.JamjetJson; +import dev.jamjet.runtime.core.ir.WorkflowIr; + +import java.io.IOException; +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.net.http.HttpTimeoutException; +import java.time.Duration; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.Executors; + +/** + * Plain-Java HTTP client for the JamJet engine REST API — the load-bearing + * Java→engine transport for the Track-9 Java ADK. + * + *

Built on the JDK {@link java.net.http.HttpClient} over a virtual-thread + * executor (no third-party HTTP dependency), serializing every body through + * {@link JamjetJson#shared()} (snake_case, {@code NON_NULL}) so Java JSON is + * wire-compatible with the Rust engine. It mirrors the canonical Python + * {@code jamjet.client.JamjetClient}: the engine serves bare + * routes ({@code /workflows}, {@code /executions}, {@code /work-items/...}), + * not {@code /api/v1/...}. + * + *

Unlike jamjet-spring's {@code JamjetRuntimeClient} (which posts to the wrong + * {@code /api/v1} routes and lacks the worker protocol), this client uses bare + * routes and adds the four work-item methods the durable Java tool-worker needs: + * {@link #claimWorkItem}, {@link #completeWorkItem}, {@link #failWorkItem}, + * {@link #heartbeatWorkItem}. + * + *

Lease fencing. The claim response carries a {@code lease_fence} + * (engine PR #108). A worker must thread it back on both {@code heartbeat} and + * {@code complete}; a stale fence on complete is rejected with HTTP 409, which this + * client surfaces as a {@link JamjetHttpException} where {@link JamjetHttpException#isConflict()} + * is true, so the worker can treat it as a lost lease rather than a generic error. + * + *

Request bodies are hand-built with literal snake_case keys (faithful to the + * Python client). This is deliberate: {@code JamjetJson}'s SNAKE_CASE strategy + * renames record/bean property names but NOT {@code Map} keys, so building bodies + * as maps with camelCase keys would emit the wrong wire shape. + */ +public final class JamjetEngineClient implements AutoCloseable { + + private static final Duration CONNECT_TIMEOUT = Duration.ofSeconds(10); + private static final Duration REQUEST_TIMEOUT = Duration.ofSeconds(30); + + private final String baseUrl; + private final String bearerToken; + private final String tenantId; + private final HttpClient http; + private final ObjectMapper json; + + /** Anonymous client (no auth headers) against {@code baseUrl}. */ + public JamjetEngineClient(String baseUrl) { + this(baseUrl, null, null); + } + + /** + * @param baseUrl engine base URL, e.g. {@code http://localhost:7700} + * @param bearerToken optional; sent as {@code Authorization: Bearer ...} when non-blank + * @param tenantId optional; sent as {@code X-Tenant-Id} when non-blank + */ + public JamjetEngineClient(String baseUrl, String bearerToken, String tenantId) { + this.baseUrl = stripTrailingSlash(baseUrl); + this.bearerToken = bearerToken; + this.tenantId = tenantId; + this.http = HttpClient.newBuilder() + .executor(Executors.newVirtualThreadPerTaskExecutor()) + .connectTimeout(CONNECT_TIMEOUT) + .build(); + this.json = JamjetJson.shared(); + } + + // -- Workflows -------------------------------------------------------------- + + /** {@code POST /workflows} with a typed IR. Body {@code {"ir": ir}}. */ + public CreateWorkflowResult createWorkflow(WorkflowIr ir) { + return postJson("/workflows", Map.of("ir", ir), CreateWorkflowResult.class); + } + + /** {@code POST /workflows} with a raw IR map. Body {@code {"ir": ir}}. */ + public CreateWorkflowResult createWorkflow(Map ir) { + return postJson("/workflows", Map.of("ir", ir), CreateWorkflowResult.class); + } + + // -- Executions ------------------------------------------------------------- + + /** + * {@code POST /executions}. Body {@code {workflow_id, input, workflow_version?}} + * (note: the version key is {@code workflow_version}, matching the Python client). + */ + public StartExecutionResult startExecution(String workflowId, + Map initialInput, + String workflowVersion) { + Map body = new LinkedHashMap<>(); + body.put("workflow_id", workflowId); + body.put("input", initialInput == null ? Map.of() : initialInput); + if (workflowVersion != null && !workflowVersion.isBlank()) { + body.put("workflow_version", workflowVersion); + } + return postJson("/executions", body, StartExecutionResult.class); + } + + /** {@code GET /executions/{id}} — the current execution snapshot for polling. */ + public ExecutionState getExecution(String executionId) { + return getJson("/executions/" + executionId, ExecutionState.class); + } + + /** {@code GET /executions/{id}/events} — the event stream (response {@code {"events": [...]}}). */ + @SuppressWarnings("unchecked") + public List> listEvents(String executionId) { + String body = execute(buildGet("/executions/" + executionId + "/events")); + Map resp = (Map) parse(body, Map.class); + if (resp == null) { + return List.of(); + } + Object events = resp.get("events"); + if (events instanceof List list) { + return list.stream().map(e -> (Map) e).toList(); + } + return List.of(); + } + + // -- Work items (worker protocol) ------------------------------------------ + + /** + * {@code POST /work-items/claim}. Body {@code {worker_id, queue_types}}. + * Returns the claimed item, or {@link Optional#empty()} when the queue is empty + * (engine responds {@code {"claimed": false}}). + */ + public Optional claimWorkItem(String workerId, List queueTypes) { + Map body = new LinkedHashMap<>(); + body.put("worker_id", workerId); + body.put("queue_types", queueTypes); + JsonNode root = readTree(execute(buildPost("/work-items/claim", body))); + if (root == null || !root.path("claimed").asBoolean(false)) { + return Optional.empty(); + } + JsonNode item = root.get("work_item"); + if (item == null || item.isNull()) { + return Optional.empty(); + } + return Optional.of(json.convertValue(item, ClaimedWorkItem.class)); + } + + /** + * {@code POST /work-items/{id}/complete} — mark the item done. + * Body {@code {output, state_patch, duration_ms, execution_id?, node_id?, + * gen_ai_model?, finish_reason?, lease_fence?}}. + * + *

{@code leaseFence} is included in the body only when non-null; when present + * the engine fences the completion, so a stale fence (the item was reclaimed) + * yields HTTP 409, surfaced as a {@link JamjetHttpException} with + * {@link JamjetHttpException#isConflict()} true. + */ + public void completeWorkItem(String itemId, + String executionId, + String nodeId, + Object output, + Map statePatch, + long durationMs, + String genAiModel, + String finishReason, + Long leaseFence) { + Map body = new LinkedHashMap<>(); + body.put("output", output); + body.put("state_patch", statePatch == null ? Map.of() : statePatch); + body.put("duration_ms", durationMs); + if (executionId != null) { + body.put("execution_id", executionId); + } + if (nodeId != null) { + body.put("node_id", nodeId); + } + if (genAiModel != null) { + body.put("gen_ai_model", genAiModel); + } + if (finishReason != null) { + body.put("finish_reason", finishReason); + } + if (leaseFence != null) { + body.put("lease_fence", leaseFence); + } + execute(buildPost("/work-items/" + itemId + "/complete", body)); + } + + /** {@code POST /work-items/{id}/fail}. Body {@code {error}}. */ + public void failWorkItem(String itemId, String error) { + execute(buildPost("/work-items/" + itemId + "/fail", Map.of("error", error))); + } + + /** + * {@code POST /work-items/{id}/heartbeat} — renew the lease. + * Body {@code {worker_id, lease_fence}}; the fence must match the claim's value + * or the engine rejects the renewal. + */ + public void heartbeatWorkItem(String itemId, String workerId, long leaseFence) { + Map body = new LinkedHashMap<>(); + body.put("worker_id", workerId); + body.put("lease_fence", leaseFence); + execute(buildPost("/work-items/" + itemId + "/heartbeat", body)); + } + + // -- Internal HTTP plumbing ------------------------------------------------- + + private T postJson(String path, Object body, Class type) { + return parse(execute(buildPost(path, body)), type); + } + + private T getJson(String path, Class type) { + return parse(execute(buildGet(path)), type); + } + + private HttpRequest buildPost(String path, Object body) { + byte[] payload; + try { + payload = json.writeValueAsBytes(body); + } catch (IOException e) { + throw new JamjetHttpException("Failed to serialize request body for " + path, e); + } + return withHeaders(HttpRequest.newBuilder(uri(path))) + .timeout(REQUEST_TIMEOUT) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofByteArray(payload)) + .build(); + } + + private HttpRequest buildGet(String path) { + return withHeaders(HttpRequest.newBuilder(uri(path))) + .timeout(REQUEST_TIMEOUT) + .GET() + .build(); + } + + private HttpRequest.Builder withHeaders(HttpRequest.Builder builder) { + builder.header("Accept", "application/json"); + if (bearerToken != null && !bearerToken.isBlank()) { + builder.header("Authorization", "Bearer " + bearerToken); + } + if (tenantId != null && !tenantId.isBlank()) { + builder.header("X-Tenant-Id", tenantId); + } + return builder; + } + + private URI uri(String path) { + return URI.create(baseUrl + path); + } + + /** Send, throwing {@link JamjetHttpException} on a non-2xx (status + body preserved). */ + private String execute(HttpRequest request) { + HttpResponse resp; + try { + resp = http.send(request, HttpResponse.BodyHandlers.ofString()); + } catch (HttpTimeoutException e) { + throw new JamjetHttpException("Request timed out: " + request.uri(), e); + } catch (IOException e) { + throw new JamjetHttpException("HTTP request failed: " + request.uri(), e); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new JamjetHttpException("HTTP request interrupted: " + request.uri(), e); + } + int status = resp.statusCode(); + String body = resp.body(); + if (status < 200 || status >= 300) { + throw new JamjetHttpException(status, body); + } + return body; + } + + private T parse(String body, Class type) { + if (body == null || body.isBlank()) { + return null; + } + try { + return json.readValue(body, type); + } catch (IOException e) { + throw new JamjetHttpException("Failed to parse response body: " + e.getMessage(), e); + } + } + + private JsonNode readTree(String body) { + if (body == null || body.isBlank()) { + return null; + } + try { + return json.readTree(body); + } catch (IOException e) { + throw new JamjetHttpException("Failed to parse response body: " + e.getMessage(), e); + } + } + + private static String stripTrailingSlash(String url) { + if (url == null) { + throw new IllegalArgumentException("baseUrl must not be null"); + } + return url.endsWith("/") ? url.substring(0, url.length() - 1) : url; + } + + /** HttpClient is JVM-managed and the virtual-thread executor has no pool to drain. */ + @Override + public void close() { + // no-op + } +} diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/client/JamjetHttpException.java b/jamjet-agent/src/main/java/dev/jamjet/agent/client/JamjetHttpException.java new file mode 100644 index 0000000..987b7da --- /dev/null +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/client/JamjetHttpException.java @@ -0,0 +1,49 @@ +package dev.jamjet.agent.client; + +/** + * Raised when an engine call fails — either a non-2xx HTTP response (carrying the + * status code + raw body) or a transport-level error (carrying a cause). + * + *

The status code is surfaced so callers can branch on it. In particular a + * {@code 409 Conflict} on {@code POST /work-items/{id}/complete} means a stale + * lease fence (the item was reclaimed): the durable tool-worker must treat this + * as a lost lease, not a generic failure. Use {@link #isConflict()} for that + * branch. + */ +public final class JamjetHttpException extends RuntimeException { + + /** -1 when the failure is transport-level (no HTTP status was received). */ + private final int statusCode; + + /** The raw response body for a non-2xx; {@code null} for transport errors. */ + private final String body; + + /** A non-2xx HTTP response. */ + public JamjetHttpException(int statusCode, String body) { + super("engine returned HTTP " + statusCode + (body == null || body.isBlank() ? "" : ": " + body)); + this.statusCode = statusCode; + this.body = body; + } + + /** A transport-level failure (timeout, connection refused, serialization, ...). */ + public JamjetHttpException(String message, Throwable cause) { + super(message, cause); + this.statusCode = -1; + this.body = null; + } + + /** The HTTP status code, or -1 for a transport-level error. */ + public int statusCode() { + return statusCode; + } + + /** The raw response body for a non-2xx, or {@code null} for a transport error. */ + public String body() { + return body; + } + + /** True iff this is a {@code 409 Conflict} (e.g. a stale lease fence on complete). */ + public boolean isConflict() { + return statusCode == 409; + } +} diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/client/StartExecutionResult.java b/jamjet-agent/src/main/java/dev/jamjet/agent/client/StartExecutionResult.java new file mode 100644 index 0000000..65c99a9 --- /dev/null +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/client/StartExecutionResult.java @@ -0,0 +1,7 @@ +package dev.jamjet.agent.client; + +/** + * Response of {@code POST /executions}. Mirrors the Python client's + * {@code start_execution} return ({@code {"execution_id": ..., "status": ...}}). + */ +public record StartExecutionResult(String executionId, String status) {} 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 new file mode 100644 index 0000000..e4550c5 --- /dev/null +++ b/jamjet-agent/src/test/java/dev/jamjet/agent/client/JamjetEngineClientTest.java @@ -0,0 +1,207 @@ +package dev.jamjet.agent.client; + +import com.github.tomakehurst.wiremock.WireMockServer; +import com.github.tomakehurst.wiremock.verification.LoggedRequest; +import dev.jamjet.runtime.core.ir.NodeDef; +import dev.jamjet.runtime.core.ir.NodeKind; +import dev.jamjet.runtime.core.ir.WorkflowIr; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import static com.github.tomakehurst.wiremock.client.WireMock.aResponse; +import static com.github.tomakehurst.wiremock.client.WireMock.equalTo; +import static com.github.tomakehurst.wiremock.client.WireMock.get; +import static com.github.tomakehurst.wiremock.client.WireMock.matchingJsonPath; +import static com.github.tomakehurst.wiremock.client.WireMock.ok; +import static com.github.tomakehurst.wiremock.client.WireMock.okJson; +import static com.github.tomakehurst.wiremock.client.WireMock.post; +import static com.github.tomakehurst.wiremock.client.WireMock.postRequestedFor; +import static com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo; +import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.options; +import static com.github.tomakehurst.wiremock.stubbing.Scenario.STARTED; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.catchThrowableOfType; + +/** + * Hermetic round-trip proof for {@link JamjetEngineClient} against a WireMock stub + * of the engine's BARE routes. Asserts both the typed results AND the exact request + * bodies/headers WireMock received, so the test proves the wire shape (snake_case + * keys, the lease fence, bare paths), not just that the code runs. + */ +class JamjetEngineClientTest { + + WireMockServer wm; + + @BeforeEach + void up() { + wm = new WireMockServer(options().dynamicPort()); + wm.start(); + } + + @AfterEach + void down() { + wm.stop(); + } + + /** A minimal agent-loop-shaped IR: a single Model start node (never a tool node). */ + private static WorkflowIr demoIr() { + var model = new NodeKind.Model("gpt-4o", null, null, "you are helpful"); + var node = new NodeDef("model_0", model, null, null, null, null, null, null); + return new WorkflowIr( + "demo", "v1", "demo agent", null, null, + "model_0", // startNode = the Model node (M2) + Map.of("model_0", node), + List.of(), + null, null, null, null, null, null, null, + null, null, null, null, null); + } + + @Test + void createStartPollToCompleted() { + wm.stubFor(post(urlEqualTo("/workflows")) + .willReturn(okJson("{\"workflow_id\":\"wf_123\",\"version\":\"v1\"}"))); + wm.stubFor(post(urlEqualTo("/executions")) + .willReturn(okJson("{\"execution_id\":\"ex_123\",\"status\":\"Running\"}"))); + // First poll -> Running, second poll -> Completed (WireMock scenario transition). + wm.stubFor(get(urlEqualTo("/executions/ex_123")).inScenario("poll") + .whenScenarioStateIs(STARTED) + .willReturn(okJson("{\"execution_id\":\"ex_123\",\"status\":\"Running\"}")) + .willSetStateTo("done")); + wm.stubFor(get(urlEqualTo("/executions/ex_123")).inScenario("poll") + .whenScenarioStateIs("done") + .willReturn(okJson( + "{\"execution_id\":\"ex_123\",\"status\":\"Completed\"," + + "\"current_state\":{\"answer\":\"42\"}}"))); + + var client = new JamjetEngineClient(wm.baseUrl(), "tkn", "t1"); + + CreateWorkflowResult wf = client.createWorkflow(demoIr()); + assertThat(wf.workflowId()).isEqualTo("wf_123"); + assertThat(wf.version()).isEqualTo("v1"); + + StartExecutionResult exec = client.startExecution("wf_123", Map.of("prompt", "hello"), "v1"); + assertThat(exec.executionId()).isEqualTo("ex_123"); + + ExecutionState state = null; + for (int i = 0; i < 10; i++) { + state = client.getExecution("ex_123"); + if ("Completed".equals(state.status())) { + break; + } + } + assertThat(state).isNotNull(); + assertThat(state.status()).isEqualTo("Completed"); + assertThat(state.currentState()).containsEntry("answer", "42"); + + // Wire shape: bare path, auth + tenant headers, typed IR serialized snake_case. + wm.verify(postRequestedFor(urlEqualTo("/workflows")) + .withHeader("Authorization", equalTo("Bearer tkn")) + .withHeader("X-Tenant-Id", equalTo("t1")) + .withRequestBody(matchingJsonPath("$.ir.workflow_id", equalTo("demo"))) + .withRequestBody(matchingJsonPath("$.ir.start_node", equalTo("model_0")))); + // Wire shape: the version key is workflow_version (not version), input nested. + wm.verify(postRequestedFor(urlEqualTo("/executions")) + .withRequestBody(matchingJsonPath("$.workflow_id", equalTo("wf_123"))) + .withRequestBody(matchingJsonPath("$.workflow_version", equalTo("v1"))) + .withRequestBody(matchingJsonPath("$.input.prompt", equalTo("hello")))); + } + + @Test + void claimThreadsLeaseFenceThroughHeartbeatAndComplete() { + wm.stubFor(post(urlEqualTo("/work-items/claim")).willReturn(okJson( + "{\"claimed\":true,\"work_item\":{" + + "\"id\":\"wi_1\",\"execution_id\":\"ex_123\",\"node_id\":\"tool_0\"," + + "\"queue_type\":\"java_tool\"," + + "\"payload\":{\"class\":\"C\",\"method\":\"m\",\"input\":{\"x\":1}}," + + "\"attempt\":1,\"lease_fence\":7}}"))); + wm.stubFor(post(urlEqualTo("/work-items/wi_1/heartbeat")).willReturn(ok())); + wm.stubFor(post(urlEqualTo("/work-items/wi_1/complete")).willReturn(ok())); + wm.stubFor(post(urlEqualTo("/work-items/wi_1/fail")).willReturn(ok())); + + var client = new JamjetEngineClient(wm.baseUrl()); + + Optional claimed = client.claimWorkItem("w1", List.of("java_tool")); + assertThat(claimed).isPresent(); + ClaimedWorkItem wi = claimed.get(); + assertThat(wi.id()).isEqualTo("wi_1"); + assertThat(wi.executionId()).isEqualTo("ex_123"); + assertThat(wi.nodeId()).isEqualTo("tool_0"); + assertThat(wi.queueType()).isEqualTo("java_tool"); + assertThat(wi.attempt()).isEqualTo(1); + assertThat(wi.leaseFence()).isEqualTo(7L); + assertThat(wi.payload()).containsEntry("method", "m"); + + client.heartbeatWorkItem(wi.id(), "w1", wi.leaseFence()); + client.completeWorkItem(wi.id(), wi.executionId(), wi.nodeId(), + Map.of("result", "ok"), Map.of("done", true), 12L, null, "stop", wi.leaseFence()); + client.failWorkItem(wi.id(), "boom"); + + wm.verify(postRequestedFor(urlEqualTo("/work-items/claim")) + .withRequestBody(matchingJsonPath("$.worker_id", equalTo("w1"))) + .withRequestBody(matchingJsonPath("$.queue_types[0]", equalTo("java_tool")))); + // The fence is echoed on heartbeat... + wm.verify(postRequestedFor(urlEqualTo("/work-items/wi_1/heartbeat")) + .withRequestBody(matchingJsonPath("$.worker_id", equalTo("w1"))) + .withRequestBody(matchingJsonPath("$.lease_fence", equalTo("7")))); + // ...and in the complete body (the fenced completion path). + wm.verify(postRequestedFor(urlEqualTo("/work-items/wi_1/complete")) + .withRequestBody(matchingJsonPath("$.lease_fence", equalTo("7"))) + .withRequestBody(matchingJsonPath("$.execution_id", equalTo("ex_123"))) + .withRequestBody(matchingJsonPath("$.node_id", equalTo("tool_0"))) + .withRequestBody(matchingJsonPath("$.duration_ms", equalTo("12"))) + .withRequestBody(matchingJsonPath("$.finish_reason", equalTo("stop"))) + .withRequestBody(matchingJsonPath("$.output.result", equalTo("ok"))) + .withRequestBody(matchingJsonPath("$.state_patch.done", equalTo("true")))); + wm.verify(postRequestedFor(urlEqualTo("/work-items/wi_1/fail")) + .withRequestBody(matchingJsonPath("$.error", equalTo("boom")))); + } + + @Test + void emptyQueueClaimReturnsEmpty() { + wm.stubFor(post(urlEqualTo("/work-items/claim")) + .willReturn(okJson("{\"claimed\":false}"))); + + var client = new JamjetEngineClient(wm.baseUrl()); + + assertThat(client.claimWorkItem("w1", List.of("java_tool"))).isEmpty(); + } + + @Test + void completeOmitsLeaseFenceWhenNull() { + wm.stubFor(post(urlEqualTo("/work-items/wi_nf/complete")).willReturn(ok())); + + var client = new JamjetEngineClient(wm.baseUrl()); + client.completeWorkItem("wi_nf", "ex_1", "n1", + Map.of("r", "ok"), Map.of(), 1L, null, null, null); + + List reqs = wm.findAll(postRequestedFor(urlEqualTo("/work-items/wi_nf/complete"))); + assertThat(reqs).hasSize(1); + String body = reqs.get(0).getBodyAsString(); + assertThat(body).doesNotContain("lease_fence"); + assertThat(body).contains("output"); + assertThat(body).contains("duration_ms"); + } + + @Test + void staleFenceCompleteSurfacesAsConflict() { + wm.stubFor(post(urlEqualTo("/work-items/wi_stale/complete")) + .willReturn(aResponse().withStatus(409).withBody("{\"error\":\"stale lease fence\"}"))); + + var client = new JamjetEngineClient(wm.baseUrl()); + + JamjetHttpException ex = catchThrowableOfType( + JamjetHttpException.class, + () -> client.completeWorkItem("wi_stale", "ex_1", "n1", + Map.of("r", "ok"), Map.of(), 1L, null, null, 1L)); + + assertThat(ex).isNotNull(); + assertThat(ex.statusCode()).isEqualTo(409); + assertThat(ex.isConflict()).isTrue(); + assertThat(ex.body()).contains("stale"); + } +} diff --git a/pom.xml b/pom.xml index ee1a7a5..7fb4548 100644 --- a/pom.xml +++ b/pom.xml @@ -50,6 +50,7 @@ jamjet-spring-boot-starter jamjet-cloud-sdk jamjet-cloud-spring-boot-starter + jamjet-agent From ca34ead647000302efc1bb046de9ac8cc08fe9e5 Mon Sep 17 00:00:00 2001 From: sdev Date: Sun, 28 Jun 2026 23:56:56 +0530 Subject: [PATCH 2/7] feat(agent): @Tool + registry + Agent builder compiling to agent-loop IR with JavaFn nodes Add NodeKind.JavaFn (snake_case java_fn: class_name/method/output_schema) + QueueType.JAVA_TOOL (java_tool), matching the Phase A Rust engine, plus a tools field on NodeKind.Model carrying OpenAI schemas. Stop serializing the computed queueType()/isDurable() helpers (@JsonIgnore) so a node matches the Rust shape exactly. Build the dev.jamjet.agent module: a @Tool method annotation, a ToolRegistry that extracts OpenAI-format schemas keyed by class+method (the single source for authoring and B-3 dispatch), and an Agent builder whose compileToIr() emits the static-unroll agent-loop WorkflowIr (Model -> tool gate -> JavaFn dispatch -> next Model, startNode always the first Model node) with governance folded into policy/token_budget/cost_budget_usd/data_policy. A golden-file parity test asserts the Java IR matches the Python compile_agent_to_ir output structurally (node ids, edges, gate branches, model tool schemas, governance, turn count), with the documented tool-node difference (java_fn vs python_fn) asserted explicitly. --- .../src/main/java/dev/jamjet/agent/Agent.java | 259 ++++++++++++++ .../dev/jamjet/agent/AgentIrCompiler.java | 305 +++++++++++++++++ .../main/java/dev/jamjet/agent/Budget.java | 34 ++ .../src/main/java/dev/jamjet/agent/Tool.java | 49 +++ .../jamjet/agent/tools/RegisteredTool.java | 39 +++ .../dev/jamjet/agent/tools/ToolRegistry.java | 209 ++++++++++++ .../dev/jamjet/agent/AgentIrCompilerTest.java | 167 +++++++++ .../dev/jamjet/agent/AgentIrParityTest.java | 291 ++++++++++++++++ .../test/java/dev/jamjet/agent/TestTools.java | 44 +++ .../jamjet/agent/tools/ToolRegistryTest.java | 112 ++++++ .../golden/agent_loop_ir.python.json | 322 ++++++++++++++++++ .../dev/jamjet/runtime/core/QueueType.java | 1 + .../dev/jamjet/runtime/core/ir/NodeKind.java | 40 ++- .../core/ir/NodeKindSerializationTest.java | 72 ++++ 14 files changed, 1942 insertions(+), 2 deletions(-) create mode 100644 jamjet-agent/src/main/java/dev/jamjet/agent/Agent.java create mode 100644 jamjet-agent/src/main/java/dev/jamjet/agent/AgentIrCompiler.java create mode 100644 jamjet-agent/src/main/java/dev/jamjet/agent/Budget.java create mode 100644 jamjet-agent/src/main/java/dev/jamjet/agent/Tool.java create mode 100644 jamjet-agent/src/main/java/dev/jamjet/agent/tools/RegisteredTool.java create mode 100644 jamjet-agent/src/main/java/dev/jamjet/agent/tools/ToolRegistry.java create mode 100644 jamjet-agent/src/test/java/dev/jamjet/agent/AgentIrCompilerTest.java create mode 100644 jamjet-agent/src/test/java/dev/jamjet/agent/AgentIrParityTest.java create mode 100644 jamjet-agent/src/test/java/dev/jamjet/agent/TestTools.java create mode 100644 jamjet-agent/src/test/java/dev/jamjet/agent/tools/ToolRegistryTest.java create mode 100644 jamjet-agent/src/test/resources/golden/agent_loop_ir.python.json diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/Agent.java b/jamjet-agent/src/main/java/dev/jamjet/agent/Agent.java new file mode 100644 index 0000000..f6615d2 --- /dev/null +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/Agent.java @@ -0,0 +1,259 @@ +package dev.jamjet.agent; + +import dev.jamjet.agent.tools.ToolRegistry; +import dev.jamjet.runtime.core.ir.PolicySetIr; +import dev.jamjet.runtime.core.ir.WorkflowIr; + +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; + +/** + * Idiomatic Java authoring for a governed, durable JamJet agent — the Java analog + * of the Python {@code jamjet.Agent}. A {@code model} + {@code @Tool} methods + + * {@code instructions} + governance knobs, compiled by {@link #compileToIr()} to + * the same agent-loop {@link WorkflowIr} the Python ADK emits, so a Java + * agent and a Python agent run on the identical Rust engine (the Java tool nodes + * being {@code java_fn} where Python's are {@code python_fn}). + * + *

Construct via the {@link #builder(String)} fluent builder: + * 

{@code
+ * Agent agent = Agent.builder("research_agent")
+ *         .model("anthropic/claude-sonnet-4-6")
+ *         .instructions("You are a helpful research assistant.")
+ *         .tools(new WebSearchTools(), new MathTools())
+ *         .policy(new PolicySetIr(List.of("delete_db"), List.of(), List.of()))
+ *         .approvalRequired(List.of("delete_*"))
+ *         .budget(new Budget(100_000, 2.5))
+ *         .build();
+ *
+ * WorkflowIr ir = agent.compileToIr();   // the durable agent-loop IR
+ * }
+ * + *

The model side is free: a Java agent emits only {@code Model} nodes; the Rust + * engine routes model calls through the governed Python model-seam sidecar. No + * Java model code is needed. + */ +public final class Agent { + + /** Default static-unroll bound for the agent loop (mirrors the Python default). */ + public static final int DEFAULT_MAX_TURNS = 8; + + private final String name; + private final String model; + private final String instructions; + private final String strategy; + private final ToolRegistry registry; + private final PolicySetIr policy; + private final boolean approvalAll; + private final List approvalGlobs; + private final Budget budget; + private final boolean pii; + private final int timeoutSeconds; + + private Agent(Builder b) { + this.name = b.name; + this.model = b.model; + this.instructions = b.instructions; + this.strategy = b.strategy; + this.registry = b.registry; + this.policy = b.policy; + this.approvalAll = b.approvalAll; + this.approvalGlobs = List.copyOf(b.approvalGlobs); + this.budget = b.budget; + this.pii = b.pii; + this.timeoutSeconds = b.timeoutSeconds; + } + + /** Start building an agent with the given logical name (used as the workflow id). */ + public static Builder builder(String name) { + return new Builder(name); + } + + /** + * Compile this agent to the durable agent-loop {@link WorkflowIr} with the + * default unroll bound ({@link #DEFAULT_MAX_TURNS}). + */ + public WorkflowIr compileToIr() { + return compileToIr(DEFAULT_MAX_TURNS); + } + + /** + * Compile this agent to the durable agent-loop {@link WorkflowIr}, statically + * unrolling up to {@code maxTurns} {@code model -> tools} turns plus a final + * tool-less answer turn. The start node is always the first Model node. + */ + public WorkflowIr compileToIr(int maxTurns) { + return AgentIrCompiler.compile(this, maxTurns); + } + + // -- accessors (read by AgentIrCompiler) ------------------------------------ + + public String name() { + return name; + } + + public String model() { + return model; + } + + public String instructions() { + return instructions; + } + + public String strategy() { + return strategy; + } + + public ToolRegistry registry() { + return registry; + } + + /** The inline policy, or {@code null} when none is set. */ + public PolicySetIr policy() { + return policy; + } + + /** {@code true} when every tool call requires human approval ({@code approvalRequired(true)}). */ + public boolean approvalAll() { + return approvalAll; + } + + /** The per-tool approval globs (possibly empty). */ + public List approvalGlobs() { + return approvalGlobs; + } + + /** The per-run budget, or {@code null} when uncapped. */ + public Budget budget() { + return budget; + } + + /** Whether PII governance is on (emits the default {@code data_policy} IR). */ + public boolean pii() { + return pii; + } + + /** The workflow timeout in seconds (compiled into {@code timeouts.workflow_timeout}). */ + public int timeoutSeconds() { + return timeoutSeconds; + } + + @Override + public String toString() { + return "Agent(name=" + name + ", model=" + model + + ", tools=" + registry.tools().stream().map(t -> t.name()).toList() + + ", strategy=" + strategy + ")"; + } + + /** Fluent builder for {@link Agent}. */ + public static final class Builder { + private final String name; + private String model; + private String instructions = ""; + private String strategy = "react"; + private ToolRegistry registry = new ToolRegistry(); + private PolicySetIr policy; + private boolean approvalAll; + private List approvalGlobs = List.of(); + private Budget budget; + private boolean pii = true; + private int timeoutSeconds = 300; + + private Builder(String name) { + this.name = Objects.requireNonNull(name, "agent name must not be null"); + } + + /** The model reference, e.g. {@code "anthropic/claude-sonnet-4-6"}. */ + public Builder model(String model) { + this.model = model; + return this; + } + + /** The system instructions for the agent. */ + public Builder instructions(String instructions) { + this.instructions = instructions == null ? "" : instructions; + return this; + } + + /** + * The reasoning strategy. v1 always compiles the durable IR as the + * react-style {@code model -> tools -> model} loop regardless of this + * value (parity with the Python durable path); the field is carried for + * forward compatibility. + */ + public Builder strategy(String strategy) { + this.strategy = strategy == null ? "react" : strategy; + return this; + } + + /** The tool-holder instances whose {@code @Tool} methods the agent may call. */ + public Builder tools(Object... holders) { + this.registry = ToolRegistry.of(holders); + return this; + } + + /** The tool-holder instances whose {@code @Tool} methods the agent may call. */ + public Builder tools(List holders) { + this.registry = ToolRegistry.of(holders); + return this; + } + + /** Use a pre-built tool registry. */ + public Builder registry(ToolRegistry registry) { + this.registry = registry == null ? new ToolRegistry() : registry; + return this; + } + + /** An inline policy ({@code blocked_tools} / {@code require_approval_for} / {@code model_allowlist}). */ + public Builder policy(PolicySetIr policy) { + this.policy = policy; + return this; + } + + /** + * Require human approval for tools. {@code true} requires approval for + * every tool call ({@code require_approval_for = ["*"]}); {@code false} + * clears the gate. + */ + public Builder approvalRequired(boolean all) { + this.approvalAll = all; + this.approvalGlobs = List.of(); + return this; + } + + /** + * Require human approval for the given tool-name globs (e.g. + * {@code ["delete_*", "send_*"]}). Unioned into {@code require_approval_for}. + */ + public Builder approvalRequired(List globs) { + this.approvalAll = false; + this.approvalGlobs = globs == null ? List.of() : new ArrayList<>(globs); + return this; + } + + /** The per-run budget (token and/or cost cap). */ + public Builder budget(Budget budget) { + this.budget = budget; + return this; + } + + /** Toggle PII governance (on by default). */ + public Builder pii(boolean pii) { + this.pii = pii; + return this; + } + + /** The workflow timeout in seconds (default 300). */ + public Builder timeoutSeconds(int timeoutSeconds) { + this.timeoutSeconds = timeoutSeconds; + return this; + } + + /** Build the immutable {@link Agent}. */ + public Agent build() { + Objects.requireNonNull(model, "agent model must be set"); + return new Agent(this); + } + } +} diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/AgentIrCompiler.java b/jamjet-agent/src/main/java/dev/jamjet/agent/AgentIrCompiler.java new file mode 100644 index 0000000..0bf09bd --- /dev/null +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/AgentIrCompiler.java @@ -0,0 +1,305 @@ +package dev.jamjet.agent; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.SerializationFeature; +import dev.jamjet.runtime.core.JamjetJson; +import dev.jamjet.runtime.core.TimeoutConfig; +import dev.jamjet.runtime.core.ir.ConditionalBranch; +import dev.jamjet.runtime.core.ir.DataPolicyIr; +import dev.jamjet.runtime.core.ir.EdgeDef; +import dev.jamjet.runtime.core.ir.NodeDef; +import dev.jamjet.runtime.core.ir.NodeKind; +import dev.jamjet.runtime.core.ir.PolicySetIr; +import dev.jamjet.runtime.core.ir.TokenBudgetIr; +import dev.jamjet.runtime.core.ir.WorkflowIr; + +import java.security.MessageDigest; +import java.time.Duration; +import java.util.ArrayList; +import java.util.HexFormat; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Locale; +import java.util.Map; + +/** + * Compiles a {@link Agent} into a durable agent-loop {@link WorkflowIr}, mirroring + * the Python {@code jamjet.compiler.agent_ir.compile_agent_to_ir} so a Java agent + * and a Python agent produce the same graph for the Rust engine. + * + *

Because the model picks tools dynamically, the loop is statically + * unrolled into {@code maxTurns} turns, each turn being three nodes: + *

    + *
  1. {@code __model_{t}__} — a {@link NodeKind.Model} carrying the agent's + * OpenAI tool schemas;
    2. + *
  2. {@code __tool_gate_{t}__} — a {@link NodeKind.Condition} on + * {@code state.last_model_finish_reason == "tool_calls"} (true → the + * turn's tool-dispatch node, false → the terminal {@code end});
    3. + *
  3. {@code __tools_{t}__} — a {@link NodeKind.JavaFn} pointing at the fixed + * Java tool dispatcher (the analog of Python's + * {@code jamjet.agents.tool_runtime:dispatch_tool_calls}); it routes to the + * next turn's model node.
    4. + *
+ * A final tool-less {@code __model_{maxTurns}__} produces the answer and routes to + * {@code end}, so the terminal is always reached via a model turn. + * + *

M2: {@code start_node} is always {@code __model_0__} (the first Model + * node), never a tool node. + */ +final class AgentIrCompiler { + + /** + * The fixed Java tool-dispatch coordinate every {@code __tools_{t}__} node + * carries — the Java analog of Python's + * {@code jamjet.agents.tool_runtime:dispatch_tool_calls}. The Phase B-3 durable + * worker resolves this dispatcher, which reads {@code last_model_tool_calls} + * from state and fans out to the requested {@code @Tool} methods via the + * {@link dev.jamjet.agent.tools.ToolRegistry}. (A single dispatcher per turn, + * not one node per dynamic call, mirrors the Python loop exactly.) + */ + static final String DISPATCH_CLASS = "dev.jamjet.agent.tools.ToolDispatcher"; + static final String DISPATCH_METHOD = "dispatchToolCalls"; + + /** The condition the tool gate branches on (matches the Model executor's recorded finish reason). */ + static final String TOOL_CALLS_EXPR = "state.last_model_finish_reason == \"tool_calls\""; + + /** Graph terminal sentinel (edges target it; it is not a node). */ + static final String END = "end"; + + private static final String MODEL_RETRY_POLICY = "llm_default"; + private static final String TOOLS_RETRY_POLICY = "no_retry"; + private static final String VERSION_BASE = "0.1.0"; + private static final int HEARTBEAT_INTERVAL_SECS = 30; + + /** The five built-in PII detectors, matching the Rust PiiRedactor + the Python default. */ + private static final List DEFAULT_PII_DETECTORS = + List.of("email", "ssn", "credit_card", "phone", "ip_address"); + + /** + * Canonical (sorted-key) snake_case mapper used only to derive the immutable + * content-version hash. Sorting map entries makes the hash independent of + * {@code Map.of}/{@code LinkedHashMap} iteration order, so the same agent + * always yields the same version across runs and JVMs. + */ + private static final ObjectMapper CANON = + JamjetJson.copy().enable(SerializationFeature.ORDER_MAP_ENTRIES_BY_KEYS); + + private AgentIrCompiler() {} + + static WorkflowIr compile(Agent agent, int maxTurns) { + if (maxTurns < 1) { + throw new IllegalArgumentException("maxTurns must be >= 1"); + } + + String modelRef = litellmModel(agent.model()); + List> toolSchemas = agent.registry().openAiToolSchemas(); + String systemPrompt = blankToNull(agent.instructions()); + + Map nodes = new LinkedHashMap<>(); + List edges = new ArrayList<>(); + + for (int t = 0; t < maxTurns; t++) { + String modelId = "__model_" + t + "__"; + String gateId = "__tool_gate_" + t + "__"; + String toolsId = "__tools_" + t + "__"; + + // Model node: carries the tool schemas, reads messages from state. + nodes.put(modelId, node( + modelId, + new NodeKind.Model(modelRef, "", "", systemPrompt, toolSchemas), + MODEL_RETRY_POLICY, + "agent turn " + t + ": model call", + Map.of("jamjet.agent.loop", "model", "jamjet.agent.turn", String.valueOf(t)))); + edges.add(new EdgeDef(modelId, gateId, null)); + + // Condition gate: tool_calls -> dispatch, else -> terminal answer. + nodes.put(gateId, node( + gateId, + new NodeKind.Condition(List.of( + new ConditionalBranch(TOOL_CALLS_EXPR, toolsId), + new ConditionalBranch(null, END))), + null, + "agent turn " + t + ": tool-call gate", + Map.of("jamjet.agent.loop", "gate", "jamjet.agent.turn", String.valueOf(t)))); + edges.add(new EdgeDef(gateId, toolsId, TOOL_CALLS_EXPR)); + edges.add(new EdgeDef(gateId, END, null)); + + // Tool-dispatch node: a JavaFn pointing at the fixed dispatcher. + // no_retry: the dispatch runs user @Tool methods (possible + // non-idempotent writes), so an already-succeeded dispatch must not + // re-run on retry (mirrors the Python tool node). + nodes.put(toolsId, node( + toolsId, + new NodeKind.JavaFn(DISPATCH_CLASS, DISPATCH_METHOD, ""), + TOOLS_RETRY_POLICY, + "agent turn " + t + ": dispatch tool calls", + Map.of("jamjet.agent.loop", "tools", "jamjet.agent.turn", String.valueOf(t)))); + // Always route forward to the NEXT model node (the last turn's dispatch + // flows into the final model node below); a tool node never targets end. + edges.add(new EdgeDef(toolsId, "__model_" + (t + 1) + "__", null)); + } + + // Final model node: no tool schemas, so the model must return a text + // answer; routes straight to end. Guarantees the terminal is reached via + // a model turn that saw the tool results. + String finalId = "__model_" + maxTurns + "__"; + nodes.put(finalId, node( + finalId, + new NodeKind.Model(modelRef, "", "", systemPrompt, List.of()), + MODEL_RETRY_POLICY, + "agent turn " + maxTurns + ": final answer (no tools)", + Map.of( + "jamjet.agent.loop", "model", + "jamjet.agent.turn", String.valueOf(maxTurns), + "jamjet.agent.final", "true"))); + edges.add(new EdgeDef(finalId, END, null)); + + Gov gov = compileGovernance(agent); + + TimeoutConfig timeouts = new TimeoutConfig( + null, + Duration.ofSeconds(agent.timeoutSeconds()), + Duration.ofSeconds(HEARTBEAT_INTERVAL_SECS), + null); + + Map labels = new LinkedHashMap<>(); + labels.put("jamjet.agent.id", agent.name()); + labels.put("jamjet.agent.loop", "true"); + labels.put("jamjet.agent.max_turns", String.valueOf(maxTurns)); + + String description = blankToNull(agent.instructions()) != null + ? agent.instructions() + : "Durable agent: " + agent.name(); + + // Build once with a placeholder version, content-hash it, rebuild with the + // real version so a changed agent (tools / instructions / governance / + // maxTurns) never reuses an immutably-cached graph (mirrors agent_ir.py). + WorkflowIr draft = assemble("", agent.name(), description, nodes, edges, timeouts, labels, gov); + return assemble(contentVersion(draft), agent.name(), description, nodes, edges, timeouts, labels, gov); + } + + // -- governance (mirrors agent_ir._compile_governance_ir) ------------------- + + private record Gov(PolicySetIr policy, TokenBudgetIr tokenBudget, Double costBudgetUsd, DataPolicyIr dataPolicy) {} + + private static Gov compileGovernance(Agent agent) { + Double costBudgetUsd = null; + TokenBudgetIr tokenBudget = null; + Budget budget = agent.budget(); + if (budget != null) { + if (budget.costUsd() != null) { + costBudgetUsd = budget.costUsd(); + } + if (budget.tokens() != null) { + // total_tokens only: the single most useful combined input+output cap. + tokenBudget = new TokenBudgetIr(null, null, budget.tokens()); + } + } + PolicySetIr policy = compilePolicy(agent); + DataPolicyIr dataPolicy = agent.pii() + ? new DataPolicyIr(List.of(), DEFAULT_PII_DETECTORS, "mask", false, true, null) + : null; + return new Gov(policy, tokenBudget, costBudgetUsd, dataPolicy); + } + + /** Mirrors {@code agent_ir._compile_agent_policy_ir}: returns {@code null} when no rules are needed. */ + private static PolicySetIr compilePolicy(Agent agent) { + boolean hasApproval = agent.approvalAll() || !agent.approvalGlobs().isEmpty(); + boolean hasPolicy = agent.policy() != null; + if (!hasPolicy && !hasApproval) { + return null; + } + + List blocked = new ArrayList<>(); + List require = new ArrayList<>(); + List allowlist = new ArrayList<>(); + if (agent.policy() != null) { + blocked.addAll(agent.policy().blockedTools()); + require.addAll(agent.policy().requireApprovalFor()); + allowlist.addAll(agent.policy().modelAllowlist()); + } + + if (agent.approvalAll()) { + require = new ArrayList<>(List.of("*")); // every tool requires approval + } else if (!agent.approvalGlobs().isEmpty()) { + // Union the caller globs with any policy require_approval_for, + // preserving declaration order and deduplicating. + LinkedHashSet merged = new LinkedHashSet<>(require); + merged.addAll(agent.approvalGlobs()); + require = new ArrayList<>(merged); + } + + return new PolicySetIr(blocked, require, allowlist); + } + + // -- helpers ---------------------------------------------------------------- + + private static NodeDef node(String id, NodeKind kind, String retryPolicy, String description, + Map labels) { + return new NodeDef(id, kind, retryPolicy, null, description, labels, null, null); + } + + private static WorkflowIr assemble(String version, String name, String description, + Map nodes, List edges, + TimeoutConfig timeouts, Map labels, Gov gov) { + return new WorkflowIr( + name, // workflowId + version, + name, // name + description, + "", // stateSchema + "__model_0__", // startNode — always the first Model node (M2) + nodes, + edges, + Map.of(), // retryPolicies + timeouts, + Map.of(), // models + Map.of(), // tools + Map.of(), // mcpServers + Map.of(), // remoteAgents + labels, + gov.policy(), + gov.tokenBudget(), + gov.costBudgetUsd(), + null, // onBudgetExceeded + gov.dataPolicy()); + } + + /** + * The litellm model string, mirroring Python {@code parse_model_ref}: a + * {@code provider/rest} string lowercases the provider; a bare string passes + * through unchanged. + */ + static String litellmModel(String model) { + String raw = model.strip(); + int slash = raw.indexOf('/'); + if (slash >= 0) { + String provider = raw.substring(0, slash).strip().toLowerCase(Locale.ROOT); + String rest = raw.substring(slash + 1); + return provider + "/" + rest; + } + return raw; + } + + /** + * The immutable content-version: {@code "0.1.0+" + sha256(canonical_ir)[:12]}. + * The hash differs from Python's (the Java IR has {@code java_fn} tool nodes + * and no python_fn input block), but its purpose — a fresh cache key whenever + * the agent's compiled content changes — is identical. + */ + static String contentVersion(WorkflowIr ir) { + try { + byte[] canonical = CANON.writeValueAsBytes(ir); + byte[] digest = MessageDigest.getInstance("SHA-256").digest(canonical); + String hex = HexFormat.of().formatHex(digest); + return VERSION_BASE + "+" + hex.substring(0, 12); + } catch (Exception e) { + throw new IllegalStateException("failed to content-version agent IR", e); + } + } + + private static String blankToNull(String s) { + return (s == null || s.isBlank()) ? null : s; + } +} diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/Budget.java b/jamjet-agent/src/main/java/dev/jamjet/agent/Budget.java new file mode 100644 index 0000000..0d61684 --- /dev/null +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/Budget.java @@ -0,0 +1,34 @@ +package dev.jamjet.agent; + +/** + * A per-run spending cap for a governed agent — the Java analog of the Python + * {@code jamjet.agents.governance.Budget}. Either or both fields may be set. + * + *

At compile time {@code costUsd} maps to the IR's {@code cost_budget_usd} and + * {@code tokens} maps to {@code token_budget.total_tokens} (a combined input+output + * cap); the Rust engine enforces both fail-closed. A {@code null} field is uncapped. + * + * @param tokens total token cap (input + output combined), or {@code null} for uncapped + * @param costUsd wall-cost cap in US dollars, or {@code null} for uncapped + */ +public record Budget(Integer tokens, Double costUsd) { + + public Budget { + if (tokens != null && tokens <= 0) { + throw new IllegalArgumentException("Budget.tokens must be positive"); + } + if (costUsd != null && costUsd <= 0) { + throw new IllegalArgumentException("Budget.costUsd must be positive"); + } + } + + /** A token-only budget. */ + public static Budget ofTokens(int tokens) { + return new Budget(tokens, null); + } + + /** A cost-only budget. */ + public static Budget ofCostUsd(double costUsd) { + return new Budget(null, costUsd); + } +} diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/Tool.java b/jamjet-agent/src/main/java/dev/jamjet/agent/Tool.java new file mode 100644 index 0000000..32f9245 --- /dev/null +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/Tool.java @@ -0,0 +1,49 @@ +package dev.jamjet.agent; + +import java.lang.annotation.ElementType; +import java.lang.annotation.Retention; +import java.lang.annotation.RetentionPolicy; +import java.lang.annotation.Target; + +/** + * Marks a method as a JamJet agent tool — the Java authoring analog of the + * Python {@code @jamjet.tool} decorator. + * + *

A {@code @Tool} method on a tool-holder object is discovered by + * {@link dev.jamjet.agent.tools.ToolRegistry}, which derives an OpenAI-format + * function schema from the method's name + parameters and offers it to the model + * in the agent loop's Model nodes. When the model requests the tool, the durable + * Java tool-worker (Phase B-3) resolves it back to this method by + * class name + method name via the same registry and invokes it + * reflectively — only declared {@code @Tool} methods are ever callable, never an + * arbitrary class named in a payload. + * + *

This is a net-new authoring annotation, deliberately distinct from the + * unrelated {@code dev.jamjet.cloud.agentboundary.Tool} (receipt metadata) and + * {@code org.springframework.ai.tool.annotation.Tool} (Spring AI) types. + * + *

Usage: + * 

{@code
+ * class WeatherTools {
+ *     @Tool(description = "Look up the current weather for a city.")
+ *     public String getWeather(String city) { ... }
+ * }
+ * }
+ */ +@Target(ElementType.METHOD) +@Retention(RetentionPolicy.RUNTIME) +public @interface Tool { + + /** + * The tool name offered to the model. Defaults to the method name when empty. + * Set this when the idiomatic Java method name (camelCase) should be exposed + * to the model under a different identifier (e.g. {@code "web_search"}). + */ + String name() default ""; + + /** + * A human-readable description of what the tool does, surfaced to the model + * in the function schema. Empty by default. + */ + String description() default ""; +} diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/tools/RegisteredTool.java b/jamjet-agent/src/main/java/dev/jamjet/agent/tools/RegisteredTool.java new file mode 100644 index 0000000..5490a2d --- /dev/null +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/tools/RegisteredTool.java @@ -0,0 +1,39 @@ +package dev.jamjet.agent.tools; + +import java.lang.reflect.Method; +import java.util.Map; + +/** + * One discovered {@link dev.jamjet.agent.Tool @Tool} method, the single source of + * truth for both authoring (the OpenAI schema offered to the model) and durable + * dispatch (the reflective handle the Phase B-3 worker resolves by class+method). + * + * @param name the tool name offered to the model ({@code @Tool.name} or the method name) + * @param description the {@code @Tool.description} (may be empty) + * @param className the fully-qualified declaring class name (the JavaFn dispatch coordinate) + * @param methodName the method name (the JavaFn dispatch coordinate) + * @param method the reflective handle (used by the B-3 worker to invoke; never serialized) + * @param instance the tool-holder instance the method is invoked on (never serialized) + * @param inputSchema the JSON-schema {@code {type, properties, required}} for the parameters + * @param openAiSchema the full OpenAI function wrapper {@code {type:function, function:{...}}} + */ +public record RegisteredTool( + String name, + String description, + String className, + String methodName, + Method method, + Object instance, + Map inputSchema, + Map openAiSchema +) { + /** A stable registry key: {@code className#methodName} (the B-3 dispatch coordinate). */ + public String key() { + return key(className, methodName); + } + + /** Build the {@code className#methodName} registry key. */ + public static String key(String className, String methodName) { + return className + "#" + methodName; + } +} diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/tools/ToolRegistry.java b/jamjet-agent/src/main/java/dev/jamjet/agent/tools/ToolRegistry.java new file mode 100644 index 0000000..be34123 --- /dev/null +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/tools/ToolRegistry.java @@ -0,0 +1,209 @@ +package dev.jamjet.agent.tools; + +import dev.jamjet.agent.Tool; + +import java.lang.reflect.Method; +import java.lang.reflect.Parameter; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +/** + * Discovers {@link Tool @Tool} methods on tool-holder objects and is the single + * source of truth for both halves of the Java agent loop: + * + *
    + *
  • Authoring — {@link #openAiToolSchemas()} returns the OpenAI-format + * function schemas (in deterministic order) that the {@code Agent} builder + * threads into every Model node's {@code tools}, mirroring the Python + * {@code agent_ir._tool_schema}.
    • + *
  • Dispatch — {@link #byClassAndMethod} lets the Phase B-3 durable + * worker resolve a {@code JavaFn} node's {@code class_name}+{@code method} + * back to the exact {@link RegisteredTool} (and its reflective handle), + * gating invocation to declared {@code @Tool} methods only (no arbitrary + * {@code Class.forName} from a payload field).
    • + *
+ * + *

Tool order is deterministic: holders are kept in registration order and, + * within a holder, {@code @Tool} methods are ordered by name (since + * {@link Class#getDeclaredMethods()} is unordered). This keeps the emitted IR + * stable across runs and JVMs, which the content-version cache key relies on. + * + *

The OpenAI schema mirrors the Python {@code @tool} decorator exactly, + * including its bare-type-string property values ({@code {"query": "string"}}, + * not {@code {"query": {"type": "string"}}}) so the schemas a Java agent offers + * the model are byte-identical to the Python equivalent. + */ +public final class ToolRegistry { + + private final List tools = new ArrayList<>(); + private final Map byKey = new LinkedHashMap<>(); + private final Map byName = new LinkedHashMap<>(); + + /** An empty registry. */ + public ToolRegistry() {} + + /** + * Build a registry from the given tool-holder instances, registered in order. + * Each holder is scanned for {@code @Tool} methods. + */ + public static ToolRegistry of(Object... holders) { + ToolRegistry registry = new ToolRegistry(); + for (Object holder : holders) { + registry.register(holder); + } + return registry; + } + + /** + * Build a registry from the given tool-holder instances, registered in order. + */ + public static ToolRegistry of(List holders) { + ToolRegistry registry = new ToolRegistry(); + for (Object holder : holders) { + registry.register(holder); + } + return registry; + } + + /** + * Scan {@code holder}'s class for {@code @Tool} methods and add them, in + * method-name order. Returns {@code this} for chaining. + * + * @throws IllegalArgumentException if two registered tools share a name + */ + public ToolRegistry register(Object holder) { + if (holder == null) { + throw new IllegalArgumentException("tool holder must not be null"); + } + Class clazz = holder.getClass(); + + List annotated = new ArrayList<>(); + for (Method m : clazz.getDeclaredMethods()) { + if (m.isAnnotationPresent(Tool.class)) { + annotated.add(m); + } + } + // getDeclaredMethods() has no defined order — sort by name for determinism. + annotated.sort(Comparator.comparing(Method::getName)); + + for (Method m : annotated) { + Tool ann = m.getAnnotation(Tool.class); + String name = ann.name().isBlank() ? m.getName() : ann.name(); + String description = ann.description(); + + Map inputSchema = buildInputSchema(m); + Map openAi = openAiSchema(name, description, inputSchema); + + RegisteredTool rt = new RegisteredTool( + name, description, clazz.getName(), m.getName(), m, holder, inputSchema, openAi); + + if (byName.containsKey(name)) { + throw new IllegalArgumentException( + "Duplicate @Tool name '" + name + "' (from " + clazz.getName() + "#" + m.getName() + + "); tool names must be unique within an agent."); + } + tools.add(rt); + byKey.put(rt.key(), rt); + byName.put(name, rt); + } + return this; + } + + /** The registered tools in deterministic order. */ + public List tools() { + return Collections.unmodifiableList(tools); + } + + /** {@code true} if no tools are registered. */ + public boolean isEmpty() { + return tools.isEmpty(); + } + + /** + * The OpenAI-format function schemas for every registered tool, in order — + * the value the {@code Agent} builder puts in each Model node's {@code tools}. + */ + public List> openAiToolSchemas() { + List> out = new ArrayList<>(tools.size()); + for (RegisteredTool t : tools) { + out.add(t.openAiSchema()); + } + return out; + } + + /** + * Resolve a tool by its dispatch coordinate ({@code className}+{@code method}). + * This is the B-3 worker's lookup: a registry hit means the method is a + * declared {@code @Tool}, so reflective invocation is safe. + */ + public RegisteredTool byClassAndMethod(String className, String methodName) { + return byKey.get(RegisteredTool.key(className, methodName)); + } + + /** Resolve a tool by the name the model used in its {@code tool_call}. */ + public RegisteredTool byName(String name) { + return byName.get(name); + } + + // -- schema construction (mirrors Python jamjet.tools.decorators) ----------- + + private static Map buildInputSchema(Method m) { + LinkedHashMap properties = new LinkedHashMap<>(); + List required = new ArrayList<>(); + for (Parameter p : m.getParameters()) { + String paramName = p.getName(); + properties.put(paramName, jsonType(p.getType())); + // Java reflection exposes no notion of an optional/defaulted parameter, + // so every @Tool parameter is required (parity with the Python schema + // where a param without a default is required). + required.add(paramName); + } + LinkedHashMap schema = new LinkedHashMap<>(); + schema.put("type", "object"); + schema.put("properties", properties); + schema.put("required", required); + return schema; + } + + private static Map openAiSchema(String name, String description, Map inputSchema) { + LinkedHashMap function = new LinkedHashMap<>(); + function.put("name", name); + function.put("description", description); + function.put("parameters", inputSchema); + + LinkedHashMap schema = new LinkedHashMap<>(); + schema.put("type", "function"); + schema.put("function", function); + return schema; + } + + /** + * Map a Java parameter type to the bare JSON-schema type string the Python + * {@code _type_to_schema} emits ({@code "string"}/{@code "integer"}/ + * {@code "number"}/{@code "boolean"}/{@code "object"}). Integer-family and + * floating-family Java types widen to {@code "integer"}/{@code "number"}. + */ + private static String jsonType(Class t) { + if (t == String.class || t == CharSequence.class || t == char.class || t == Character.class) { + return "string"; + } + if (t == int.class || t == Integer.class + || t == long.class || t == Long.class + || t == short.class || t == Short.class + || t == byte.class || t == Byte.class) { + return "integer"; + } + if (t == float.class || t == Float.class || t == double.class || t == Double.class) { + return "number"; + } + if (t == boolean.class || t == Boolean.class) { + return "boolean"; + } + // Anything else (records, POJOs, collections, ...) is an object. + return "object"; + } +} diff --git a/jamjet-agent/src/test/java/dev/jamjet/agent/AgentIrCompilerTest.java b/jamjet-agent/src/test/java/dev/jamjet/agent/AgentIrCompilerTest.java new file mode 100644 index 0000000..981632f --- /dev/null +++ b/jamjet-agent/src/test/java/dev/jamjet/agent/AgentIrCompilerTest.java @@ -0,0 +1,167 @@ +package dev.jamjet.agent; + +import dev.jamjet.runtime.core.JamjetJson; +import dev.jamjet.runtime.core.ir.NodeKind; +import dev.jamjet.runtime.core.ir.PolicySetIr; +import dev.jamjet.runtime.core.ir.WorkflowIr; +import org.junit.jupiter.api.Test; + +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Structural unit tests for {@link AgentIrCompiler} that do not require the Python + * golden: the static unroll shape, the M2 start-node invariant, the JavaFn tool + * nodes, governance omission, model-ref parsing, and content-version behaviour. + */ +class AgentIrCompilerTest { + + private static Agent agent(int timeout, String instructions) { + return Agent.builder("a1") + .model("anthropic/claude-sonnet-4-6") + .instructions(instructions) + .tools(new TestTools.WebSearchTool()) + .timeoutSeconds(timeout) + .build(); + } + + @Test + void rejectsNonPositiveMaxTurns() { + assertThatThrownBy(() -> agent(300, "hi").compileToIr(0)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("maxTurns"); + } + + @Test + void startNodeIsAlwaysFirstModelNode_M2() { + WorkflowIr ir = agent(300, "hi").compileToIr(3); + assertThat(ir.startNode()).isEqualTo("__model_0__"); + // The start node must be a Model node, never a tool node. + assertThat(ir.node(ir.startNode()).kind()).isInstanceOf(NodeKind.Model.class); + } + + @Test + void unrollEmitsThreeNodesPerTurnPlusFinalModel() { + WorkflowIr ir = agent(300, "hi").compileToIr(4); + // 4 turns * (model + gate + tools) + 1 final model = 13 nodes. + assertThat(ir.nodes()).hasSize(4 * 3 + 1); + assertThat(ir.node("__model_4__")).isNotNull(); // final model + assertThat(ir.node("__tools_3__").kind()).isInstanceOf(NodeKind.JavaFn.class); + assertThat(ir.node("__tools_4__")).isNull(); // no tools after the final model + } + + @Test + void toolNodesAreJavaFnPointingAtTheFixedDispatcher() { + WorkflowIr ir = agent(300, "hi").compileToIr(2); + NodeKind.JavaFn tools = (NodeKind.JavaFn) ir.node("__tools_0__").kind(); + assertThat(tools.className()).isEqualTo("dev.jamjet.agent.tools.ToolDispatcher"); + assertThat(tools.method()).isEqualTo("dispatchToolCalls"); + assertThat(tools.outputSchema()).isEmpty(); + assertThat(tools.queueType().getValue()).isEqualTo("java_tool"); + assertThat(ir.node("__tools_0__").retryPolicy()).isEqualTo("no_retry"); + } + + @Test + void perTurnModelsCarryToolSchemasFinalModelDoesNot() { + WorkflowIr ir = agent(300, "hi").compileToIr(2); + NodeKind.Model turn0 = (NodeKind.Model) ir.node("__model_0__").kind(); + NodeKind.Model finalModel = (NodeKind.Model) ir.node("__model_2__").kind(); + assertThat(turn0.tools()).hasSize(1); + assertThat(finalModel.tools()).isEmpty(); + assertThat(ir.node("__model_0__").retryPolicy()).isEqualTo("llm_default"); + } + + @Test + void blankInstructionsYieldNullSystemPromptAndDefaultDescription() { + WorkflowIr ir = Agent.builder("a1").model("gpt-4o").tools().build().compileToIr(1); + NodeKind.Model m = (NodeKind.Model) ir.node("__model_0__").kind(); + assertThat(m.systemPrompt()).isNull(); + assertThat(ir.description()).isEqualTo("Durable agent: a1"); + } + + @Test + void litellmModelMirrorsPythonParseModelRef() { + // provider/rest lowercases the provider; bare strings pass through. + assertThat(AgentIrCompiler.litellmModel("anthropic/claude-sonnet-4-6")) + .isEqualTo("anthropic/claude-sonnet-4-6"); + assertThat(AgentIrCompiler.litellmModel("Anthropic/Claude-Opus")).isEqualTo("anthropic/Claude-Opus"); + assertThat(AgentIrCompiler.litellmModel("gpt-4o")).isEqualTo("gpt-4o"); + assertThat(AgentIrCompiler.litellmModel(" openai/gpt-4o ")).isEqualTo("openai/gpt-4o"); + } + + @Test + void defaultMaxTurnsIsEight() { + WorkflowIr ir = agent(300, "hi").compileToIr(); + assertThat(ir.node("__model_8__")).isNotNull(); // final answer node at maxTurns=8 + assertThat(ir.node("__tools_7__")).isNotNull(); + assertThat(ir.node("__tools_8__")).isNull(); + } + + @Test + void governanceFieldsOmittedWhenUnset() throws Exception { + // No policy, no budget, pii off -> none of the governance IR fields emit. + Agent bare = Agent.builder("a1").model("gpt-4o").tools(new TestTools.WebSearchTool()) + .pii(false).build(); + WorkflowIr ir = bare.compileToIr(1); + assertThat(ir.policy()).isNull(); + assertThat(ir.tokenBudget()).isNull(); + assertThat(ir.costBudgetUsd()).isNull(); + assertThat(ir.dataPolicy()).isNull(); + + String json = JamjetJson.shared().writeValueAsString(ir); + assertThat(json).doesNotContain("\"policy\""); + assertThat(json).doesNotContain("\"cost_budget_usd\""); + assertThat(json).doesNotContain("\"token_budget\""); + assertThat(json).doesNotContain("\"data_policy\""); + } + + @Test + void approvalAllSetsWildcardRequireApproval() { + Agent a = Agent.builder("a1").model("gpt-4o").tools(new TestTools.WebSearchTool()) + .approvalRequired(true).build(); + PolicySetIr policy = a.compileToIr(1).policy(); + assertThat(policy).isNotNull(); + assertThat(policy.requireApprovalFor()).containsExactly("*"); + } + + @Test + void approvalGlobsUnionWithInlinePolicyPreservingOrder() { + Agent a = Agent.builder("a1").model("gpt-4o").tools(new TestTools.WebSearchTool()) + .policy(new PolicySetIr(List.of(), List.of("wire_*"), List.of())) + .approvalRequired(List.of("delete_*", "wire_*")) // wire_* is a dup + .build(); + PolicySetIr policy = a.compileToIr(1).policy(); + // Union preserves the policy's order first, then new globs; dedups wire_*. + assertThat(policy.requireApprovalFor()).containsExactly("wire_*", "delete_*"); + } + + @Test + void emittedIrRoundTripsThroughTypedPort() throws Exception { + // The emitted agent IR must survive a serialize -> deserialize through the + // same typed WorkflowIr port the engine uses, with JavaFn tool nodes and + // tool-carrying Model nodes intact. + WorkflowIr ir = agent(300, "hi").compileToIr(2); + String json = JamjetJson.shared().writeValueAsString(ir); + WorkflowIr back = WorkflowIr.fromJson(json); + + assertThat(back.startNode()).isEqualTo("__model_0__"); + assertThat(back.nodes()).hasSize(7); + assertThat(back.node("__tools_0__").kind()).isInstanceOf(NodeKind.JavaFn.class); + NodeKind.Model m = (NodeKind.Model) back.node("__model_0__").kind(); + assertThat(m.tools()).hasSize(1); + } + + @Test + void contentVersionIsDeterministicAndContentSensitive() { + String v1 = agent(300, "hello").compileToIr(2).version(); + String v1Again = agent(300, "hello").compileToIr(2).version(); + String vDifferentInstructions = agent(300, "different").compileToIr(2).version(); + String vDifferentTurns = agent(300, "hello").compileToIr(3).version(); + + assertThat(v1).startsWith("0.1.0+").isEqualTo(v1Again); // same agent -> same version + assertThat(vDifferentInstructions).isNotEqualTo(v1); // changed content -> new key + assertThat(vDifferentTurns).isNotEqualTo(v1); // changed unroll -> new key + } +} diff --git a/jamjet-agent/src/test/java/dev/jamjet/agent/AgentIrParityTest.java b/jamjet-agent/src/test/java/dev/jamjet/agent/AgentIrParityTest.java new file mode 100644 index 0000000..b912345 --- /dev/null +++ b/jamjet-agent/src/test/java/dev/jamjet/agent/AgentIrParityTest.java @@ -0,0 +1,291 @@ +package dev.jamjet.agent; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +import dev.jamjet.runtime.core.JamjetJson; +import dev.jamjet.runtime.core.ir.PolicySetIr; +import dev.jamjet.runtime.core.ir.WorkflowIr; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +import java.io.InputStream; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.TreeSet; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * GOLDEN-FILE PARITY: the Java {@code Agent} builder must emit the SAME agent-loop + * graph the Python ADK emits, so both run on the identical Rust engine. + * + *

The golden {@code golden/agent_loop_ir.python.json} is the output of the + * canonical Python {@code jamjet.compiler.agent_ir.compile_agent_to_ir} for a small + * fixture agent (model + two @tool functions + instructions + inline policy + + * approval globs + token/cost budget + pii) at {@code max_turns=2}. This test builds + * the structurally identical Java {@link Agent} and asserts every load-bearing field + * matches: node ids, edges, the gate condition, the Model nodes' tool schemas, + * governance (policy / budget / data_policy), turn count, and {@code start_node}. + * + *

Documented, expected differences (asserted explicitly below)

+ *
    + *
  • Tool node kind: Python emits {@code python_fn{module,function,input}} + * pointing at {@code jamjet.agents.tool_runtime:dispatch_tool_calls}; Java emits + * {@code java_fn{class_name,method}} pointing at the fixed Java dispatcher. The + * Rust {@code JavaFn} struct has no {@code input} field (the scheduler enriches + * the payload from state), so the Java tool node carries no {@code input} block.
    • + *
  • Version: both are {@code "0.1.0+"+sha256(content)[:12]}, but the hash + * differs because the Java IR's content differs (java_fn vs python_fn).
    • + *
  • Per-node {@code description}: cosmetic human text the engine ignores; + * Java uses house punctuation. Excluded from the structural comparison.
    • + *
  • Gate {@code expression}: Python adds a redundant {@code expression} + * alongside {@code branches}; Java emits only {@code branches} (the engine reads + * branches and ignores expression).
    • + *
  • Explicit nulls: the snake_case {@code JamjetJson} mapper omits null + * fields (NON_NULL) where Python writes {@code "condition": null} etc.; both + * trees are null-stripped before comparison.
    • + *
+ */ +class AgentIrParityTest { + + private static final ObjectMapper M = JamjetJson.shared(); + + private static final List EXPECTED_NODE_IDS = List.of( + "__model_0__", "__tool_gate_0__", "__tools_0__", + "__model_1__", "__tool_gate_1__", "__tools_1__", + "__model_2__"); + + private static JsonNode py; // the Python golden IR + private static JsonNode jv; // the Java-emitted IR + + @BeforeAll + static void compileBoth() throws Exception { + py = loadGolden(); + jv = M.readTree(M.writeValueAsBytes(buildJavaAgent().compileToIr(2))); + } + + /** The structurally identical Java agent matching the Python golden fixture. */ + private static Agent buildJavaAgent() { + return Agent.builder("research_agent") + .model("anthropic/claude-sonnet-4-6") + .instructions("You are a helpful research assistant.") + .tools(new TestTools.WebSearchTool(), new TestTools.MathTool()) + .policy(new PolicySetIr( + List.of("delete_db"), // blocked_tools + List.of(), // require_approval_for (from approval globs) + List.of("anthropic/claude-sonnet-4-6"))) // model_allowlist + .approvalRequired(List.of("delete_*")) + .budget(new Budget(100_000, 2.5)) + .pii(true) + .build(); + } + + // -- structural parity (must match) ----------------------------------------- + + @Test + void sameTopLevelIdentityAndStartNode() { + for (String f : List.of("workflow_id", "name", "description", "state_schema", "start_node")) { + assertThat(jv.get(f)).as(f).isEqualTo(py.get(f)); + } + // M2: the loop ALWAYS starts on the first Model node, never a tool node. + assertThat(jv.get("start_node").asText()).isEqualTo("__model_0__"); + } + + @Test + void sameNodeIdSet() { + Set pyIds = nodeIds(py); + Set jvIds = nodeIds(jv); + assertThat(jvIds).isEqualTo(pyIds); + assertThat(jvIds).containsExactlyInAnyOrderElementsOf(EXPECTED_NODE_IDS); + } + + @Test + void sameTurnCount() { + // 2 model->tools turns + 1 final model node; max_turns label = "2". + assertThat(jv.get("labels").get("jamjet.agent.max_turns").asText()).isEqualTo("2"); + assertThat(jv.get("labels").get("jamjet.agent.max_turns").asText()) + .isEqualTo(py.get("labels").get("jamjet.agent.max_turns").asText()); + long jvToolNodes = nodeIds(jv).stream().filter(id -> id.startsWith("__tools_")).count(); + assertThat(jvToolNodes).isEqualTo(2); + } + + @Test + void sameEdgesGraph() { + assertThat(edgeSet(jv)).isEqualTo(edgeSet(py)); + // 9 edges: per turn {model->gate, gate->tools, gate->end, tools->model}, + final model->end. + assertThat(jv.get("edges")).hasSize(9); + } + + @Test + void sameModelNodesIncludingToolSchemas() { + // The Model node kinds must be byte-identical: model_ref, prompt_ref, + // output_schema, system_prompt, AND the OpenAI tool schemas offered to the + // model. This is what makes "the same graph the engine runs" true. + for (String id : List.of("__model_0__", "__model_1__", "__model_2__")) { + JsonNode pk = canonicalize(py.get("nodes").get(id).get("kind")); + JsonNode jk = canonicalize(jv.get("nodes").get(id).get("kind")); + assertThat(jk).as("model kind %s", id).isEqualTo(pk); + } + // The per-turn models carry both tools; the final model carries none. + assertThat(jv.get("nodes").get("__model_0__").get("kind").get("tools")).hasSize(2); + assertThat(jv.get("nodes").get("__model_2__").get("kind").get("tools")).hasSize(0); + assertThat(jv.get("nodes").get("__model_0__").get("kind").get("model_ref").asText()) + .isEqualTo("anthropic/claude-sonnet-4-6"); + } + + @Test + void sameGateBranches() { + for (String id : List.of("__tool_gate_0__", "__tool_gate_1__")) { + JsonNode pb = canonicalize(py.get("nodes").get(id).get("kind").get("branches")); + JsonNode jb = canonicalize(jv.get("nodes").get(id).get("kind").get("branches")); + assertThat(jb).as("gate branches %s", id).isEqualTo(pb); + } + // The condition the gate branches on. + assertThat(jv.get("nodes").get("__tool_gate_0__").get("kind") + .get("branches").get(0).get("condition").asText()) + .isEqualTo("state.last_model_finish_reason == \"tool_calls\""); + } + + @Test + void sameGovernanceFields() { + for (String f : List.of("policy", "cost_budget_usd", "token_budget", "data_policy")) { + assertThat(canonicalize(jv.get(f))).as(f).isEqualTo(canonicalize(py.get(f))); + } + // Spot-check the merged require_approval_for (the ["delete_*"] approval glob). + assertThat(jv.get("policy").get("require_approval_for").get(0).asText()).isEqualTo("delete_*"); + assertThat(jv.get("token_budget").get("total_tokens").asInt()).isEqualTo(100_000); + assertThat(jv.get("cost_budget_usd").asDouble()).isEqualTo(2.5); + assertThat(jv.get("data_policy").get("pii_detectors")).hasSize(5); + } + + @Test + void sameTimeoutsLabelsAndEmptyMaps() { + assertThat(canonicalize(jv.get("timeouts"))).isEqualTo(canonicalize(py.get("timeouts"))); + assertThat(canonicalize(jv.get("labels"))).isEqualTo(canonicalize(py.get("labels"))); + for (String f : List.of("retry_policies", "models", "tools", "mcp_servers", "remote_agents")) { + assertThat(jv.get(f)).as(f).isEqualTo(py.get(f)); + } + } + + // -- documented differences (asserted explicitly) --------------------------- + + @Test + void toolNodesAreJavaFnNotPythonFn() { + for (String id : List.of("__tools_0__", "__tools_1__")) { + JsonNode pyTool = py.get("nodes").get(id); + JsonNode jvTool = jv.get("nodes").get(id); + + // THE key difference: python_fn -> java_fn, with the Java dispatch ref. + assertThat(pyTool.get("kind").get("type").asText()).isEqualTo("python_fn"); + assertThat(jvTool.get("kind").get("type").asText()).isEqualTo("java_fn"); + + assertThat(jvTool.get("kind").get("class_name").asText()) + .isEqualTo("dev.jamjet.agent.tools.ToolDispatcher"); + assertThat(jvTool.get("kind").get("method").asText()).isEqualTo("dispatchToolCalls"); + assertThat(jvTool.get("kind").get("output_schema").asText()).isEmpty(); + + // The Rust JavaFn struct has no `input` field; the scheduler enriches the + // payload from state. Python's python_fn carries a redundant `input` block. + assertThat(jvTool.get("kind").has("input")).isFalse(); + assertThat(pyTool.get("kind").has("input")).isTrue(); + + // Everything else about the tool node matches: id, retry policy, labels, + // and the surrounding edges (verified in sameEdgesGraph). + assertThat(jvTool.get("id").asText()).isEqualTo(pyTool.get("id").asText()); + assertThat(jvTool.get("retry_policy").asText()) + .isEqualTo(pyTool.get("retry_policy").asText()) + .isEqualTo("no_retry"); + assertThat(canonicalize(jvTool.get("labels"))).isEqualTo(canonicalize(pyTool.get("labels"))); + } + } + + @Test + void versionDiffersButBothAreContentHashed() { + assertThat(py.get("version").asText()).startsWith("0.1.0+"); + assertThat(jv.get("version").asText()).startsWith("0.1.0+"); + // The java_fn content yields a different hash than the python_fn content. + assertThat(jv.get("version").asText()).isNotEqualTo(py.get("version").asText()); + } + + // -- backstop: deep-equal of the whole graph after normalization ------------ + + @Test + void deepEqualAfterNormalization() { + // The strongest statement: once the documented differences are normalized + // away (version, per-node description, gate expression, the tool-node kind), + // the ENTIRE Java IR tree deep-equals the Python IR tree. + assertThat(normalize(jv)).isEqualTo(normalize(py)); + } + + // -- helpers ---------------------------------------------------------------- + + private static JsonNode loadGolden() throws Exception { + try (InputStream in = AgentIrParityTest.class.getResourceAsStream("/golden/agent_loop_ir.python.json")) { + assertThat(in).as("golden resource present").isNotNull(); + return M.readTree(in); + } + } + + private static Set nodeIds(JsonNode ir) { + Set ids = new TreeSet<>(); + ir.get("nodes").fieldNames().forEachRemaining(ids::add); + return ids; + } + + private static Set edgeSet(JsonNode ir) { + Set edges = new HashSet<>(); + canonicalize(ir.get("edges")).forEach(edges::add); + return edges; + } + + /** Recursively strip null-valued object fields (harmonize NON_NULL vs explicit null). */ + private static JsonNode canonicalize(JsonNode node) { + if (node == null) { + return null; + } + if (node.isObject()) { + ObjectNode out = M.createObjectNode(); + node.fields().forEachRemaining(e -> { + if (!e.getValue().isNull()) { + out.set(e.getKey(), canonicalize(e.getValue())); + } + }); + return out; + } + if (node.isArray()) { + ArrayNode out = M.createArrayNode(); + node.forEach(c -> out.add(canonicalize(c))); + return out; + } + return node; + } + + /** + * Null-strip the whole IR and remove the documented-divergent fields so the + * remaining graph (identity / model nodes incl. tool schemas / gate branches / + * governance / edges / timeouts / labels) can be deep-compared. The tool-node + * {@code kind} legitimately differs (java_fn vs python_fn) and is masked with a + * sentinel so the rest of each tool node (id / retry_policy / labels) is still + * compared. + */ + private static JsonNode normalize(JsonNode root) { + ObjectNode r = (ObjectNode) canonicalize(root); + r.remove("version"); + ObjectNode nodes = (ObjectNode) r.get("nodes"); + nodes.fields().forEachRemaining(e -> { + ObjectNode n = (ObjectNode) e.getValue(); + n.remove("description"); // cosmetic; Java uses house punctuation + ObjectNode kind = (ObjectNode) n.get("kind"); + if ("condition".equals(kind.path("type").asText())) { + kind.remove("expression"); // Python-only redundant metadata + } + if (e.getKey().startsWith("__tools_")) { + n.set("kind", M.createObjectNode().put("type", "__tool_dispatch__")); + } + }); + return r; + } +} diff --git a/jamjet-agent/src/test/java/dev/jamjet/agent/TestTools.java b/jamjet-agent/src/test/java/dev/jamjet/agent/TestTools.java new file mode 100644 index 0000000..452b427 --- /dev/null +++ b/jamjet-agent/src/test/java/dev/jamjet/agent/TestTools.java @@ -0,0 +1,44 @@ +package dev.jamjet.agent; + +/** + * Tool-holder fixtures shared by the agent tests. The two tools mirror the Python + * golden fixture used to generate {@code golden/agent_loop_ir.python.json}: + * {@code web_search(query: str) -> str} and {@code add_numbers(a: int, b: int) -> str}. + */ +public final class TestTools { + + private TestTools() {} + + /** One @Tool method: {@code web_search(query)}. */ + public static final class WebSearchTool { + @Tool(name = "web_search", description = "Search the web for a query.") + public String webSearch(String query) { + return "results for " + query; + } + } + + /** One @Tool method: {@code add_numbers(a, b)}. */ + public static final class MathTool { + @Tool(name = "add_numbers", description = "Add two integers.") + public String addNumbers(int a, int b) { + return String.valueOf(a + b); + } + } + + /** Two @Tool methods in one holder, plus a non-tool method, for ordering tests. */ + public static final class MultiTool { + @Tool(description = "Z tool.") + public String zebra(String s) { + return s; + } + + @Tool(description = "A tool.") + public boolean alpha(boolean flag) { + return flag; + } + + public String notATool(String ignored) { + return ignored; + } + } +} diff --git a/jamjet-agent/src/test/java/dev/jamjet/agent/tools/ToolRegistryTest.java b/jamjet-agent/src/test/java/dev/jamjet/agent/tools/ToolRegistryTest.java new file mode 100644 index 0000000..0a4c487 --- /dev/null +++ b/jamjet-agent/src/test/java/dev/jamjet/agent/tools/ToolRegistryTest.java @@ -0,0 +1,112 @@ +package dev.jamjet.agent.tools; + +import dev.jamjet.agent.TestTools; +import org.junit.jupiter.api.Test; + +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Proves the registry extracts OpenAI-format schemas from {@code @Tool} methods + * exactly as the Python {@code @tool} decorator does (bare-type-string property + * values, all params required), keys tools by class+method for B-3 dispatch, and + * gates invocation to declared tools only. + */ +class ToolRegistryTest { + + @Test + void extractsOpenAiSchemaForStringParamTool() { + ToolRegistry registry = ToolRegistry.of(new TestTools.WebSearchTool()); + + assertThat(registry.tools()).hasSize(1); + RegisteredTool t = registry.tools().get(0); + + assertThat(t.name()).isEqualTo("web_search"); + assertThat(t.description()).isEqualTo("Search the web for a query."); + assertThat(t.className()).isEqualTo(TestTools.WebSearchTool.class.getName()); + assertThat(t.methodName()).isEqualTo("webSearch"); + + // The OpenAI function wrapper, mirroring agent_ir._tool_schema. + assertThat(t.openAiSchema()).isEqualTo(Map.of( + "type", "function", + "function", Map.of( + "name", "web_search", + "description", "Search the web for a query.", + "parameters", Map.of( + "type", "object", + // Bare type string, NOT {"type":"string"} — matches Python. + "properties", Map.of("query", "string"), + "required", List.of("query"))))); + } + + @Test + void mapsPrimitiveParamTypesToBareJsonSchemaStrings() { + ToolRegistry registry = ToolRegistry.of(new TestTools.MathTool()); + RegisteredTool t = registry.tools().get(0); + + @SuppressWarnings("unchecked") + Map params = (Map) t.inputSchema(); + assertThat(params).containsEntry("type", "object"); + assertThat(params.get("properties")).isEqualTo(Map.of("a", "integer", "b", "integer")); + // Every Java @Tool parameter is required (no optional-param notion). + assertThat(params.get("required")).isEqualTo(List.of("a", "b")); + } + + @Test + void resolvesByClassAndMethodForWorkerDispatch() { + TestTools.WebSearchTool holder = new TestTools.WebSearchTool(); + ToolRegistry registry = ToolRegistry.of(holder); + + RegisteredTool byKey = registry.byClassAndMethod( + TestTools.WebSearchTool.class.getName(), "webSearch"); + assertThat(byKey).isNotNull(); + assertThat(byKey.name()).isEqualTo("web_search"); + // The reflective handle + instance the B-3 worker invokes. + assertThat(byKey.method().getName()).isEqualTo("webSearch"); + assertThat(byKey.instance()).isSameAs(holder); + + // The model uses the tool NAME in its tool_call; that resolves too. + assertThat(registry.byName("web_search")).isSameAs(byKey); + } + + @Test + void unregisteredClassMethodResolvesToNull() { + ToolRegistry registry = ToolRegistry.of(new TestTools.WebSearchTool()); + // No registry entry -> the worker must fail, never Class.forName an + // arbitrary class named in a payload. + assertThat(registry.byClassAndMethod("com.evil.Rce", "exploit")).isNull(); + assertThat(registry.byName("nonexistent")).isNull(); + } + + @Test + void preservesHolderOrderAndSortsMethodsWithinHolderByName() { + // Holder order preserved across holders; @Tool methods within a holder + // are name-sorted for determinism (getDeclaredMethods is unordered). + ToolRegistry registry = ToolRegistry.of( + new TestTools.WebSearchTool(), new TestTools.MultiTool()); + + assertThat(registry.tools().stream().map(RegisteredTool::name).toList()) + .containsExactly("web_search", "alpha", "zebra"); + + // Non-@Tool methods are not registered. + assertThat(registry.byName("notATool")).isNull(); + } + + @Test + void rejectsDuplicateToolNames() { + assertThatThrownBy(() -> ToolRegistry.of( + new TestTools.WebSearchTool(), new TestTools.WebSearchTool())) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("Duplicate @Tool name 'web_search'"); + } + + @Test + void emptyRegistryHasNoSchemas() { + ToolRegistry registry = new ToolRegistry(); + assertThat(registry.isEmpty()).isTrue(); + assertThat(registry.openAiToolSchemas()).isEmpty(); + } +} diff --git a/jamjet-agent/src/test/resources/golden/agent_loop_ir.python.json b/jamjet-agent/src/test/resources/golden/agent_loop_ir.python.json new file mode 100644 index 0000000..fa7b0da --- /dev/null +++ b/jamjet-agent/src/test/resources/golden/agent_loop_ir.python.json @@ -0,0 +1,322 @@ +{ + "workflow_id": "research_agent", + "version": "0.1.0+8476bdc8e9fb", + "name": "research_agent", + "description": "You are a helpful research assistant.", + "state_schema": "", + "start_node": "__model_0__", + "nodes": { + "__model_0__": { + "id": "__model_0__", + "kind": { + "type": "model", + "model_ref": "anthropic/claude-sonnet-4-6", + "prompt_ref": "", + "output_schema": "", + "system_prompt": "You are a helpful research assistant.", + "tools": [ + { + "type": "function", + "function": { + "name": "web_search", + "description": "Search the web for a query.", + "parameters": { + "type": "object", + "properties": { + "query": "string" + }, + "required": [ + "query" + ] + } + } + }, + { + "type": "function", + "function": { + "name": "add_numbers", + "description": "Add two integers.", + "parameters": { + "type": "object", + "properties": { + "a": "integer", + "b": "integer" + }, + "required": [ + "a", + "b" + ] + } + } + } + ] + }, + "retry_policy": "llm_default", + "node_timeout_secs": null, + "description": "agent turn 0 — model call", + "labels": { + "jamjet.agent.loop": "model", + "jamjet.agent.turn": "0" + } + }, + "__tool_gate_0__": { + "id": "__tool_gate_0__", + "kind": { + "type": "condition", + "branches": [ + { + "condition": "state.last_model_finish_reason == \"tool_calls\"", + "target": "__tools_0__" + }, + { + "condition": null, + "target": "end" + } + ], + "expression": "state.last_model_finish_reason == \"tool_calls\"" + }, + "retry_policy": null, + "node_timeout_secs": null, + "description": "agent turn 0 — tool-call gate", + "labels": { + "jamjet.agent.loop": "gate", + "jamjet.agent.turn": "0" + } + }, + "__tools_0__": { + "id": "__tools_0__", + "kind": { + "type": "python_fn", + "module": "jamjet.agents.tool_runtime", + "function": "dispatch_tool_calls", + "output_schema": "", + "input": { + "messages": "$state.messages", + "assistant_content": "$state.last_model_output", + "tool_calls": "$state.last_model_tool_calls", + "tools": { + "web_search": "__main__:web_search", + "add_numbers": "__main__:add_numbers" + } + } + }, + "retry_policy": "no_retry", + "node_timeout_secs": null, + "description": "agent turn 0 — dispatch tool calls", + "labels": { + "jamjet.agent.loop": "tools", + "jamjet.agent.turn": "0" + } + }, + "__model_1__": { + "id": "__model_1__", + "kind": { + "type": "model", + "model_ref": "anthropic/claude-sonnet-4-6", + "prompt_ref": "", + "output_schema": "", + "system_prompt": "You are a helpful research assistant.", + "tools": [ + { + "type": "function", + "function": { + "name": "web_search", + "description": "Search the web for a query.", + "parameters": { + "type": "object", + "properties": { + "query": "string" + }, + "required": [ + "query" + ] + } + } + }, + { + "type": "function", + "function": { + "name": "add_numbers", + "description": "Add two integers.", + "parameters": { + "type": "object", + "properties": { + "a": "integer", + "b": "integer" + }, + "required": [ + "a", + "b" + ] + } + } + } + ] + }, + "retry_policy": "llm_default", + "node_timeout_secs": null, + "description": "agent turn 1 — model call", + "labels": { + "jamjet.agent.loop": "model", + "jamjet.agent.turn": "1" + } + }, + "__tool_gate_1__": { + "id": "__tool_gate_1__", + "kind": { + "type": "condition", + "branches": [ + { + "condition": "state.last_model_finish_reason == \"tool_calls\"", + "target": "__tools_1__" + }, + { + "condition": null, + "target": "end" + } + ], + "expression": "state.last_model_finish_reason == \"tool_calls\"" + }, + "retry_policy": null, + "node_timeout_secs": null, + "description": "agent turn 1 — tool-call gate", + "labels": { + "jamjet.agent.loop": "gate", + "jamjet.agent.turn": "1" + } + }, + "__tools_1__": { + "id": "__tools_1__", + "kind": { + "type": "python_fn", + "module": "jamjet.agents.tool_runtime", + "function": "dispatch_tool_calls", + "output_schema": "", + "input": { + "messages": "$state.messages", + "assistant_content": "$state.last_model_output", + "tool_calls": "$state.last_model_tool_calls", + "tools": { + "web_search": "__main__:web_search", + "add_numbers": "__main__:add_numbers" + } + } + }, + "retry_policy": "no_retry", + "node_timeout_secs": null, + "description": "agent turn 1 — dispatch tool calls", + "labels": { + "jamjet.agent.loop": "tools", + "jamjet.agent.turn": "1" + } + }, + "__model_2__": { + "id": "__model_2__", + "kind": { + "type": "model", + "model_ref": "anthropic/claude-sonnet-4-6", + "prompt_ref": "", + "output_schema": "", + "system_prompt": "You are a helpful research assistant.", + "tools": [] + }, + "retry_policy": "llm_default", + "node_timeout_secs": null, + "description": "agent turn 2 — final answer (no tools)", + "labels": { + "jamjet.agent.loop": "model", + "jamjet.agent.turn": "2", + "jamjet.agent.final": "true" + } + } + }, + "edges": [ + { + "from": "__model_0__", + "to": "__tool_gate_0__", + "condition": null + }, + { + "from": "__tool_gate_0__", + "to": "__tools_0__", + "condition": "state.last_model_finish_reason == \"tool_calls\"" + }, + { + "from": "__tool_gate_0__", + "to": "end", + "condition": null + }, + { + "from": "__tools_0__", + "to": "__model_1__", + "condition": null + }, + { + "from": "__model_1__", + "to": "__tool_gate_1__", + "condition": null + }, + { + "from": "__tool_gate_1__", + "to": "__tools_1__", + "condition": "state.last_model_finish_reason == \"tool_calls\"" + }, + { + "from": "__tool_gate_1__", + "to": "end", + "condition": null + }, + { + "from": "__tools_1__", + "to": "__model_2__", + "condition": null + }, + { + "from": "__model_2__", + "to": "end", + "condition": null + } + ], + "retry_policies": {}, + "timeouts": { + "workflow_timeout": 300, + "heartbeat_interval": 30 + }, + "models": {}, + "tools": {}, + "mcp_servers": {}, + "remote_agents": {}, + "labels": { + "jamjet.agent.id": "research_agent", + "jamjet.agent.loop": "true", + "jamjet.agent.max_turns": "2" + }, + "cost_budget_usd": 2.5, + "token_budget": { + "total_tokens": 100000 + }, + "policy": { + "blocked_tools": [ + "delete_db" + ], + "require_approval_for": [ + "delete_*" + ], + "model_allowlist": [ + "anthropic/claude-sonnet-4-6" + ] + }, + "data_policy": { + "pii_fields": [], + "pii_detectors": [ + "email", + "ssn", + "credit_card", + "phone", + "ip_address" + ], + "redaction_mode": "mask", + "retain_prompts": false, + "retain_outputs": true + } +} diff --git a/jamjet-runtime-core/src/main/java/dev/jamjet/runtime/core/QueueType.java b/jamjet-runtime-core/src/main/java/dev/jamjet/runtime/core/QueueType.java index d2129bd..008afd5 100644 --- a/jamjet-runtime-core/src/main/java/dev/jamjet/runtime/core/QueueType.java +++ b/jamjet-runtime-core/src/main/java/dev/jamjet/runtime/core/QueueType.java @@ -7,6 +7,7 @@ public enum QueueType { MODEL("model"), TOOL("tool"), PYTHON_TOOL("python_tool"), + JAVA_TOOL("java_tool"), RETRIEVAL("retrieval"), PRIVILEGED("privileged"), GENERAL("general"); diff --git a/jamjet-runtime-core/src/main/java/dev/jamjet/runtime/core/ir/NodeKind.java b/jamjet-runtime-core/src/main/java/dev/jamjet/runtime/core/ir/NodeKind.java index 5a59ec6..c65fd7e 100644 --- a/jamjet-runtime-core/src/main/java/dev/jamjet/runtime/core/ir/NodeKind.java +++ b/jamjet-runtime-core/src/main/java/dev/jamjet/runtime/core/ir/NodeKind.java @@ -1,5 +1,6 @@ package dev.jamjet.runtime.core.ir; +import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.annotation.JsonSubTypes; import com.fasterxml.jackson.annotation.JsonTypeInfo; import dev.jamjet.runtime.core.QueueType; @@ -12,6 +13,7 @@ @JsonSubTypes.Type(value = NodeKind.Model.class, name = "model"), @JsonSubTypes.Type(value = NodeKind.Tool.class, name = "tool"), @JsonSubTypes.Type(value = NodeKind.PythonFn.class, name = "python_fn"), + @JsonSubTypes.Type(value = NodeKind.JavaFn.class, name = "java_fn"), @JsonSubTypes.Type(value = NodeKind.Condition.class, name = "condition"), @JsonSubTypes.Type(value = NodeKind.Parallel.class, name = "parallel"), @JsonSubTypes.Type(value = NodeKind.Join.class, name = "join"), @@ -30,11 +32,16 @@ }) public sealed interface NodeKind { + // Computed routing/durability helpers — NOT IR fields. @JsonIgnore keeps them + // off the wire so a serialized node matches the Rust NodeKind shape exactly + // (the Rust structs have no `queue_type`/`durable` field on a node kind). + @JsonIgnore default QueueType queueType() { return switch (this) { case Model m -> QueueType.MODEL; case Tool t -> QueueType.TOOL; case PythonFn p -> QueueType.PYTHON_TOOL; + case JavaFn jf -> QueueType.JAVA_TOOL; case MemoryRetrieval r -> QueueType.RETRIEVAL; case Finalizer f -> QueueType.TOOL; case McpTool m -> QueueType.TOOL; @@ -43,6 +50,7 @@ default QueueType queueType() { }; } + @JsonIgnore default boolean isDurable() { return !(this instanceof Condition); } @@ -51,8 +59,22 @@ record Model( String modelRef, String promptRef, String outputSchema, - String systemPrompt - ) implements NodeKind {} + String systemPrompt, + List> tools + ) implements NodeKind { + public Model { + // OpenAI-format tool/function schemas offered to the model for this + // call (mirrors the Rust Model.tools / Python agent_ir _model_kind). + // Empty (never null) means no tools are offered; it serializes as + // "tools":[] matching the Rust #[serde(default)] Vec field. + tools = tools == null ? List.of() : List.copyOf(tools); + } + + /** Back-compat constructor: a Model node offering no tools (plain text completion). */ + public Model(String modelRef, String promptRef, String outputSchema, String systemPrompt) { + this(modelRef, promptRef, outputSchema, systemPrompt, List.of()); + } + } record Tool( String toolRef, @@ -70,6 +92,20 @@ record PythonFn( String outputSchema ) implements NodeKind {} + /** + * Arbitrary Java method executed by an external durable Java tool-worker — + * the Java analog of {@link PythonFn}. Serializes with the snake_case type + * tag {@code "java_fn"} and fields {@code class_name}/{@code method}/ + * {@code output_schema}, exactly matching the Rust engine's + * {@code JavaFn { class_name, method, output_schema }} (Phase A). Routes to + * the {@link QueueType#JAVA_TOOL} queue, which the Java tool-worker drains. + */ + record JavaFn( + String className, + String method, + String outputSchema + ) implements NodeKind {} + record Condition( List branches ) implements NodeKind { diff --git a/jamjet-runtime-core/src/test/java/dev/jamjet/runtime/core/ir/NodeKindSerializationTest.java b/jamjet-runtime-core/src/test/java/dev/jamjet/runtime/core/ir/NodeKindSerializationTest.java index 04de3bd..7229a23 100644 --- a/jamjet-runtime-core/src/test/java/dev/jamjet/runtime/core/ir/NodeKindSerializationTest.java +++ b/jamjet-runtime-core/src/test/java/dev/jamjet/runtime/core/ir/NodeKindSerializationTest.java @@ -29,6 +29,77 @@ void modelNodeRoundTrip() throws JsonProcessingException { assertThat(deserialized).isEqualTo(original); } + @Test + void javaFnNodeRoundTrip() throws JsonProcessingException { + NodeKind original = new NodeKind.JavaFn("com.example.AgentTools", "dispatch", ""); + String json = mapper.writeValueAsString(original); + + // Exactly the Rust engine's Phase-A wire shape: + // JavaFn { class_name, method, output_schema } with the "java_fn" tag. + assertThat(json).contains("\"type\":\"java_fn\""); + assertThat(json).contains("\"class_name\":\"com.example.AgentTools\""); + assertThat(json).contains("\"method\":\"dispatch\""); + assertThat(json).contains("\"output_schema\":\"\""); + // It must NOT leak the camelCase Java property name. + assertThat(json).doesNotContain("className"); + + NodeKind deserialized = mapper.readValue(json, NodeKind.class); + assertThat(deserialized).isEqualTo(original); + assertThat(deserialized).isInstanceOf(NodeKind.JavaFn.class); + } + + @Test + void modelNodeCarriesToolSchemas() throws JsonProcessingException { + var schema = Map.of( + "type", "function", + "function", Map.of( + "name", "web_search", + "description", "Search the web.", + "parameters", Map.of( + "type", "object", + "properties", Map.of("query", "string"), + "required", List.of("query")))); + NodeKind original = new NodeKind.Model("gpt4", "", "", "be helpful", List.of(schema)); + String json = mapper.writeValueAsString(original); + + assertThat(json).contains("\"type\":\"model\""); + assertThat(json).contains("\"tools\":["); + assertThat(json).contains("\"name\":\"web_search\""); + + NodeKind deserialized = mapper.readValue(json, NodeKind.class); + assertThat(deserialized).isEqualTo(original); + assertThat(((NodeKind.Model) deserialized).tools()).hasSize(1); + } + + @Test + void modelNodeWithoutToolsEmitsEmptyToolsArray() throws JsonProcessingException { + // The 4-arg back-compat constructor offers no tools; it serializes as + // "tools":[] (never null), matching the Rust #[serde(default)] Vec and + // the Python final-answer node's "tools": []. + NodeKind original = new NodeKind.Model("gpt4", "", "", null); + String json = mapper.writeValueAsString(original); + assertThat(json).contains("\"tools\":[]"); + + NodeKind deserialized = mapper.readValue(json, NodeKind.class); + assertThat(deserialized).isEqualTo(original); + assertThat(((NodeKind.Model) deserialized).tools()).isEmpty(); + } + + @Test + void computedHelpersNeverSerialize() throws JsonProcessingException { + // queueType()/isDurable() are computed routing helpers, not IR fields: + // they must not leak onto the wire (the Rust NodeKind has no such field). + for (NodeKind kind : List.of( + new NodeKind.Model("m", null, null, null), + new NodeKind.JavaFn("C", "m", ""), + new NodeKind.Condition(List.of()))) { + String json = mapper.writeValueAsString(kind); + assertThat(json).doesNotContain("durable"); + assertThat(json).doesNotContain("queue_type"); + assertThat(json).doesNotContain("queueType"); + } + } + @Test void toolNodeRoundTrip() throws JsonProcessingException { NodeKind original = new NodeKind.Tool("search_tool", Map.of("query", "$.input.query"), "{}"); @@ -83,6 +154,7 @@ void queueTypeMapping() { assertThat(new NodeKind.Model("m", null, null, null).queueType()).isEqualTo(QueueType.MODEL); assertThat(new NodeKind.Tool("t", null, null).queueType()).isEqualTo(QueueType.TOOL); assertThat(new NodeKind.PythonFn("m", "f", null).queueType()).isEqualTo(QueueType.PYTHON_TOOL); + assertThat(new NodeKind.JavaFn("C", "m", null).queueType()).isEqualTo(QueueType.JAVA_TOOL); assertThat(new NodeKind.MemoryRetrieval("c", "q", null).queueType()).isEqualTo(QueueType.RETRIEVAL); assertThat(new NodeKind.Finalizer("t", FinalizerTrigger.ALWAYS).queueType()).isEqualTo(QueueType.TOOL); assertThat(new NodeKind.McpTool("s", "t", null, null).queueType()).isEqualTo(QueueType.TOOL); From 121300bf6cd9112299d93c272bfe4ffd797bf0e9 Mon Sep 17 00:00:00 2001 From: sdev Date: Mon, 29 Jun 2026 00:25:59 +0530 Subject: [PATCH 3/7] feat(agent): durable java_tool worker (fence + heartbeat + registry-gated dispatch) B-3 of Track 9 Phase B. Adds the ToolDispatcher the JavaFn nodes point at and the durable worker that drains the java_tool queue, mirroring the hardened Python worker and dispatch_tool_calls. ToolDispatcher.dispatchToolCalls reads last_model_tool_calls / last_model_output (with the tool_calls / assistant_content fallbacks) from the work item input, resolves each call by name through the ToolRegistry only, coerces JSON args to the typed method params via Jackson, invokes reflectively, and returns the updated message list as {messages:[...]} (the node output and the state_patch). An unknown tool yields a clean tool error to the model, never an arbitrary invocation. JavaToolWorker claims java_tool items and settles them exactly-once: - completes WITH the claim lease_fence (the engine fences /complete; PR #108) - treats a 409 on complete as a lost lease no-op, never failing the item - runs a heartbeat that, on a definitive engine rejection, aborts the in-flight tool and skips completion (M3); the engine surfaces a lost heartbeat lease as 500 - gates dispatch to the one fixed ToolDispatcher coordinate, never Class.forName on a payload field (the RCE gate) Tests cover the happy path (complete with fence), fence-lost mid-flight (no complete), 409 = lost lease (no fail), a forged dispatch coordinate (clean fail, not RCE), and arg coercion. mvn -pl jamjet-agent -am verify green (50 tests). --- .../jamjet/agent/tools/ToolDispatcher.java | 286 +++++++++++++++ .../agent/tools/ToolInvocationException.java | 28 ++ .../jamjet/agent/worker/JavaToolWorker.java | 327 ++++++++++++++++++ .../agent/tools/ToolDispatcherTest.java | 186 ++++++++++ .../agent/worker/JavaToolWorkerTest.java | 302 ++++++++++++++++ 5 files changed, 1129 insertions(+) create mode 100644 jamjet-agent/src/main/java/dev/jamjet/agent/tools/ToolDispatcher.java create mode 100644 jamjet-agent/src/main/java/dev/jamjet/agent/tools/ToolInvocationException.java create mode 100644 jamjet-agent/src/main/java/dev/jamjet/agent/worker/JavaToolWorker.java create mode 100644 jamjet-agent/src/test/java/dev/jamjet/agent/tools/ToolDispatcherTest.java create mode 100644 jamjet-agent/src/test/java/dev/jamjet/agent/worker/JavaToolWorkerTest.java diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/tools/ToolDispatcher.java b/jamjet-agent/src/main/java/dev/jamjet/agent/tools/ToolDispatcher.java new file mode 100644 index 0000000..577bf1b --- /dev/null +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/tools/ToolDispatcher.java @@ -0,0 +1,286 @@ +package dev.jamjet.agent.tools; + +import com.fasterxml.jackson.databind.ObjectMapper; +import dev.jamjet.runtime.core.JamjetJson; + +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; +import java.lang.reflect.Parameter; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +/** + * The Java analog of Python's {@code jamjet.agents.tool_runtime.dispatch_tool_calls}: + * runs one model turn's requested tool calls and returns the updated message list. + * + *

This is the method every {@code __tools_{t}__} {@code JavaFn} node points at + * ({@code dev.jamjet.agent.tools.ToolDispatcher#dispatchToolCalls}). On the durable + * engine the Phase B-3 {@code JavaToolWorker} resolves that fixed coordinate to a + * {@code ToolDispatcher} bound to the agent's {@link ToolRegistry} and invokes + * {@link #dispatchToolCalls(Map)} with the work item's {@code input} (the full + * accumulated workflow state the scheduler enriched in — see the Rust scheduler's + * {@code JavaFn} arm). The return replaces {@code state["messages"]} for the next + * model turn, exactly as the Python dispatcher does. + * + *

Input / output contract (mirrors {@code dispatch_tool_calls})

+ * Reads from the state map, with the same fallbacks so it is unit-testable without + * the engine: + *
    + *
  • {@code messages} — the running message list (default {@code []}).
    • + *
  • {@code tool_calls} then {@code last_model_tool_calls} — the requested calls + * {@code [{id, name, arguments}, ...]} (the Model executor records the latter).
    • + *
  • {@code assistant_content} then {@code last_model_output} — the assistant text.
    • + *
+ * Returns {@code {"messages": }}: the original messages, plus the + * assistant message that requested the calls (OpenAI shape), plus one {@code role: + * tool} message per call. The worker uses this both as the node {@code output} and as + * the {@code state_patch} (the engine merge-patches it into state, replacing the + * top-level {@code messages} key) — identical to the Python {@code python_tool} worker. + * + *

Tool resolution is registry-gated (the RCE gate)

+ * Each call is resolved by name through {@link ToolRegistry#byName(String)} + * only. The registry's name map is populated exclusively from + * {@code @Tool}-annotated methods scanned off the agent's tool holders, so the only + * reflective handles that exist are those pre-captured {@link RegisteredTool#method()} + * objects. A model-supplied (hence untrusted) tool name that is not a declared + * {@code @Tool} simply misses the map: there is no code path that + * turns a name into a {@code Class.forName(...)} or an arbitrary {@code Method}, so a + * hallucinated or malicious tool name can never invoke arbitrary code. An unknown tool + * yields a clean {@code role: tool} error message surfaced to the model (so a single + * hallucinated name does not abort the durable run), never an exception that crashes + * the worker and never an arbitrary invocation. A registered tool that throws + * at runtime propagates (the worker fails the item), mirroring the Python dispatcher. + */ +public final class ToolDispatcher { + + /** The fixed dispatch class coordinate the {@code JavaFn} tool nodes carry. */ + public static final String DISPATCH_CLASS = "dev.jamjet.agent.tools.ToolDispatcher"; + + /** The fixed dispatch method coordinate the {@code JavaFn} tool nodes carry. */ + public static final String DISPATCH_METHOD = "dispatchToolCalls"; + + private final ToolRegistry registry; + private final ObjectMapper json; + + /** Bind a dispatcher to the agent's tool registry (the only callable tool set). */ + public ToolDispatcher(ToolRegistry registry) { + if (registry == null) { + throw new IllegalArgumentException("registry must not be null"); + } + this.registry = registry; + this.json = JamjetJson.shared(); + } + + /** + * Run the tool calls from one model turn and return the updated message list as + * {@code {"messages": }}. + * + * @param input the accumulated workflow state (the work item's enriched {@code input}) + * @throws ToolInvocationException if a registered tool throws, or the + * dispatch thread is interrupted (lease lost) between calls + */ + public Map dispatchToolCalls(Map input) { + Map state = input == null ? Map.of() : input; + + List messages = new ArrayList<>(asList(state.get("messages"))); + List> toolCalls = readToolCalls(state); + String assistantContent = readAssistantContent(state); + + // 1. The assistant message that requested these calls (OpenAI shape, so the + // history replays cleanly into the next model turn). + messages.add(assistantMessage(assistantContent, toolCalls)); + + // 2. Execute each requested tool; append its result as a `role: tool` message. + for (Map call : toolCalls) { + // Cooperative abort checkpoint: if the worker cancelled us because the + // lease was lost (M3), stop before running another (possibly + // side-effecting) tool. The worker's lease-lost gate is the real + // safety net; this just avoids extra work after an abort. + if (Thread.currentThread().isInterrupted()) { + throw new ToolInvocationException( + "", new InterruptedException("tool dispatch aborted (lease lost)")); + } + Object callId = call.get("id"); + String name = stringOrNull(call.get("name")); + Map arguments = coerceArguments(call.get("arguments")); + + String content = invokeTool(name, arguments); + messages.add(toolMessage(callId, name, content)); + } + + // 3. Return the FULL updated list — replaces state["messages"]. + Map out = new LinkedHashMap<>(); + out.put("messages", messages); + return out; + } + + // -- tool invocation (registry-gated) --------------------------------------- + + /** + * Resolve and invoke one tool by name. A registry miss (unknown/unregistered + * name) returns a clean error string for a {@code role: tool} message rather than + * throwing — this is the safe, model-recoverable path and never touches an + * arbitrary class. A registered tool that throws propagates as a + * {@link ToolInvocationException} (the worker fails the item). + */ + private String invokeTool(String name, Map arguments) { + RegisteredTool tool = name == null ? null : registry.byName(name); + if (tool == null) { + // RCE GATE: an unknown name is not a declared @Tool. Do NOT reflect on it; + // surface a clean tool error to the model. + return "ERROR: tool '" + name + "' is not a registered @Tool"; + } + + Method method = tool.method(); + Object[] args = coerceToParameters(method, arguments); + try { + method.setAccessible(true); + Object result = method.invoke(tool.instance(), args); + return stringify(result); + } catch (InvocationTargetException e) { + Throwable cause = e.getCause() != null ? e.getCause() : e; + throw new ToolInvocationException(name, cause); + } catch (ReflectiveOperationException | IllegalArgumentException e) { + throw new ToolInvocationException(name, e); + } + } + + /** + * Coerce the JSON argument map to the method's typed positional parameters via + * Jackson, matching each parameter by its (compiled {@code -parameters}) name. + * A missing argument coerces to {@code null} (or the primitive default). + */ + private Object[] coerceToParameters(Method method, Map arguments) { + Parameter[] params = method.getParameters(); + Object[] out = new Object[params.length]; + for (int i = 0; i < params.length; i++) { + Parameter p = params[i]; + Object raw = arguments.get(p.getName()); + out[i] = json.convertValue(raw, p.getType()); + } + return out; + } + + // -- message shaping (mirrors tool_runtime helpers) ------------------------- + + private Map assistantMessage(String content, List> toolCalls) { + List calls = new ArrayList<>(toolCalls.size()); + for (Map call : toolCalls) { + calls.add(assistantToolCall(call)); + } + Map m = new LinkedHashMap<>(); + m.put("role", "assistant"); + m.put("content", content == null ? "" : content); + m.put("tool_calls", calls); + return m; + } + + /** Render a {@code {id, name, arguments}} call into the OpenAI assistant shape. */ + private Map assistantToolCall(Map call) { + Object rawArgs = call.get("arguments"); + String argString; + if (rawArgs instanceof String s) { + argString = s; + } else { + try { + argString = json.writeValueAsString(rawArgs == null ? Map.of() : rawArgs); + } catch (Exception e) { + argString = "{}"; + } + } + Map function = new LinkedHashMap<>(); + function.put("name", call.get("name")); + function.put("arguments", argString); + + Map m = new LinkedHashMap<>(); + m.put("id", call.get("id")); + m.put("type", "function"); + m.put("function", function); + return m; + } + + private Map toolMessage(Object callId, String name, String content) { + Map m = new LinkedHashMap<>(); + m.put("role", "tool"); + m.put("tool_call_id", callId); + m.put("name", name); + m.put("content", content); + return m; + } + + // -- state reading (with Python fallbacks) ---------------------------------- + + @SuppressWarnings("unchecked") + private List> readToolCalls(Map state) { + Object calls = state.get("tool_calls"); + if (calls == null) { + calls = state.get("last_model_tool_calls"); + } + if (calls instanceof List list) { + List> out = new ArrayList<>(list.size()); + for (Object e : list) { + if (e instanceof Map m) { + out.add((Map) m); + } + } + return out; + } + return List.of(); + } + + private String readAssistantContent(Map state) { + Object content = state.get("assistant_content"); + if (content == null) { + content = state.get("last_model_output"); + } + return content == null ? null : String.valueOf(content); + } + + /** Tool arguments arrive as a JSON object or a JSON string — normalise to a map. */ + @SuppressWarnings("unchecked") + private Map coerceArguments(Object arguments) { + if (arguments instanceof String s) { + if (s.isBlank()) { + return Map.of(); + } + try { + return json.readValue(s, Map.class); + } catch (Exception e) { + return Map.of(); + } + } + if (arguments instanceof Map m) { + return (Map) m; + } + return Map.of(); + } + + /** Coerce a tool result to the string content a {@code role: tool} message needs. */ + private String stringify(Object result) { + if (result == null) { + return "null"; + } + if (result instanceof String s) { + return s; + } + try { + return json.writeValueAsString(result); + } catch (Exception e) { + return String.valueOf(result); + } + } + + @SuppressWarnings("unchecked") + private static List asList(Object value) { + if (value instanceof List list) { + return (List) list; + } + return List.of(); + } + + private static String stringOrNull(Object value) { + return value == null ? null : String.valueOf(value); + } +} diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/tools/ToolInvocationException.java b/jamjet-agent/src/main/java/dev/jamjet/agent/tools/ToolInvocationException.java new file mode 100644 index 0000000..52bff34 --- /dev/null +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/tools/ToolInvocationException.java @@ -0,0 +1,28 @@ +package dev.jamjet.agent.tools; + +/** + * Raised by {@link ToolDispatcher} when a registered {@code @Tool} method + * throws during reflective invocation, or when the dispatch thread is interrupted + * (the worker aborted it because the lease was lost). It carries the offending tool + * name plus the underlying cause so the durable worker can fail the work item with a + * faithful message (mirroring the Python dispatcher, which lets a tool exception + * propagate out of {@code dispatch_tool_calls}). + * + *

An unknown/unregistered tool name does NOT raise this — it is surfaced + * to the model as a clean {@code role: tool} error message instead, so a hallucinated + * name cannot fail the run (and, critically, never triggers an arbitrary invocation). + */ +public final class ToolInvocationException extends RuntimeException { + + private final String toolName; + + public ToolInvocationException(String toolName, Throwable cause) { + super("tool '" + toolName + "' failed: " + (cause == null ? "unknown error" : cause.toString()), cause); + this.toolName = toolName; + } + + /** The name of the tool whose invocation failed. */ + public String toolName() { + return toolName; + } +} 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 new file mode 100644 index 0000000..2af8e6f --- /dev/null +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/worker/JavaToolWorker.java @@ -0,0 +1,327 @@ +package dev.jamjet.agent.worker; + +import dev.jamjet.agent.client.ClaimedWorkItem; +import dev.jamjet.agent.client.JamjetEngineClient; +import dev.jamjet.agent.client.JamjetHttpException; +import dev.jamjet.agent.tools.ToolDispatcher; +import dev.jamjet.agent.tools.ToolRegistry; + +import java.lang.System.Logger; +import java.lang.System.Logger.Level; +import java.time.Duration; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.CancellationException; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; + +/** + * The durable Java tool-worker: drains the {@code java_tool} queue over the + * {@link JamjetEngineClient}, dispatching each claimed agent-loop tool turn to the + * registry-bound {@link ToolDispatcher} and settling the work item back to the engine + * exactly-once. It is the Java mirror of the hardened Python {@code jamjet worker} + * loop ({@code jamjet.cli.main._worker_loop}), framework-free plain Java 21. + * + *

Lease fencing — the enforcement axis (M3 + PR #108)

+ * A reclaimed / zombie worker must never complete an item the engine has handed to a + * new claimant. Two independent guards enforce this: + *
    + *
  1. Proactive heartbeat abort (M3). While a tool runs, a heartbeat task + * renews the lease with the claim's {@code lease_fence}. If a heartbeat is + * rejected by the engine (any definitive HTTP error response, {@code status >= + * 400}), the lease is treated as lost: the in-flight dispatch is cancelled + * ({@code Future.cancel(true)} interrupts the tool thread) and the worker does + * NOT complete. The engine surfaces a stale heartbeat fence as HTTP 500 + * ({@code FenceLost -> Internal}), so we key the abort on a definitive HTTP + * error rather than a specific code; a pure transport error (no HTTP status) is + * treated as a transient blip and does not abort (matching the Python worker's + * tolerance), because the complete-time fence below is the authoritative + * backstop.
    2. + *
  2. Authoritative complete-time fence (the backstop). The completion echoes + * the same {@code lease_fence}; the engine's {@code /complete} fences on it and + * rejects a stale/reclaimed lease with HTTP 409, emitting NO + * {@code NodeCompleted}. The worker treats a 409 on complete as a lost-lease + * no-op: it does NOT {@code fail} the item (failing would clobber the work the + * new claimant is running). This mirrors the hardened Python worker exactly and + * guarantees correctness even if the heartbeat abort has not fired yet.
    3. + *
+ * + *

Registry-gated dispatch — the RCE gate

+ * The worker never reflects on an attacker-influenceable string. Its reflective + * surface is a closed set: (a) the single fixed {@link ToolDispatcher} — the + * worker verifies the work item's {@code class}/{@code method} payload equals the + * known dispatch coordinate and fails cleanly otherwise, so a forged {@code java_fn} + * node naming an arbitrary class can never be {@code Class.forName}'d; and (b) the + * declared {@code @Tool} methods the dispatcher resolves by name through the + * {@link ToolRegistry}. No payload or model field becomes a class/method to load. + */ +public final class JavaToolWorker implements AutoCloseable { + + /** The single queue this worker drains. */ + public static final String JAVA_TOOL_QUEUE = "java_tool"; + + private static final Logger LOG = System.getLogger(JavaToolWorker.class.getName()); + private static final Duration DEFAULT_HEARTBEAT = Duration.ofSeconds(10); + private static final Duration DEFAULT_POLL_BACKOFF = Duration.ofSeconds(2); + + /** The outcome of processing a single claim. */ + public enum ItemResult { + /** The queue was empty; nothing was claimed. */ + EMPTY, + /** The tool ran and the completion was accepted by the engine. */ + COMPLETED, + /** A genuine tool / dispatch / completion error; the item was failed. */ + FAILED, + /** The lease was lost (heartbeat abort or a 409 on complete); a no-op, NOT a failure. */ + LOST_LEASE + } + + private final JamjetEngineClient client; + private final String workerId; + private final ToolDispatcher dispatcher; + private final long heartbeatMillis; + private final long pollBackoffMillis; + + private final ScheduledExecutorService heartbeats; + private final ExecutorService dispatchPool; + private volatile boolean stopped; + + /** Build a worker with default heartbeat (10s) and poll backoff (2s). */ + public JavaToolWorker(JamjetEngineClient client, String workerId, ToolRegistry registry) { + this(client, workerId, registry, DEFAULT_HEARTBEAT, DEFAULT_POLL_BACKOFF); + } + + /** + * @param client the engine client (work-item protocol) + * @param workerId this worker's stable id (echoed on claim + heartbeat) + * @param registry the agent's tool registry — the ONLY callable tools + * @param heartbeatInterval lease-renewal interval while a tool runs + * @param pollBackoff sleep between empty claims in {@link #run()} + */ + public JavaToolWorker(JamjetEngineClient client, String workerId, ToolRegistry registry, + Duration heartbeatInterval, Duration pollBackoff) { + this.client = Objects.requireNonNull(client, "client"); + this.workerId = Objects.requireNonNull(workerId, "workerId"); + this.dispatcher = new ToolDispatcher(Objects.requireNonNull(registry, "registry")); + this.heartbeatMillis = heartbeatInterval.toMillis(); + this.pollBackoffMillis = pollBackoff.toMillis(); + this.heartbeats = Executors.newSingleThreadScheduledExecutor(r -> { + Thread t = new Thread(r, "jamjet-java-tool-heartbeat"); + t.setDaemon(true); + return t; + }); + this.dispatchPool = Executors.newThreadPerTaskExecutor(Thread.ofVirtual().name("jamjet-java-tool-", 0).factory()); + } + + /** Request a clean stop of {@link #run()} (takes effect at the next loop boundary). */ + public void stop() { + this.stopped = true; + } + + /** + * Claim and process exactly one work item, then return its outcome. Returns + * {@link ItemResult#EMPTY} when the queue is empty. Mirrors the Python worker's + * {@code --once} mode. + */ + public ItemResult runOnce() { + Optional claimed = client.claimWorkItem(workerId, List.of(JAVA_TOOL_QUEUE)); + if (claimed.isEmpty()) { + return ItemResult.EMPTY; + } + return process(claimed.get()); + } + + /** + * Continuously claim and process {@code java_tool} items until {@link #stop()} is + * called or the thread is interrupted, backing off when the queue is empty. A poll + * error is logged and retried after the backoff (never fatal). + */ + public void run() { + while (!stopped && !Thread.currentThread().isInterrupted()) { + try { + Optional claimed = client.claimWorkItem(workerId, List.of(JAVA_TOOL_QUEUE)); + if (claimed.isEmpty()) { + if (sleep(pollBackoffMillis)) { + return; // interrupted + } + continue; + } + process(claimed.get()); + } catch (Exception pollErr) { + LOG.log(Level.WARNING, () -> "poll error: " + pollErr); + if (sleep(pollBackoffMillis)) { + return; + } + } + } + } + + // -- per-item processing ---------------------------------------------------- + + private ItemResult process(ClaimedWorkItem item) { + Map payload = item.payload() == null ? Map.of() : item.payload(); + + // 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. + 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; + } + + Map input = asMap(payload.get("input")); + long fence = item.leaseFence() == null ? 0L : item.leaseFence(); + AtomicBoolean leaseLost = new AtomicBoolean(false); + + long startNanos = System.nanoTime(); + Future> dispatchFuture = dispatchPool.submit(() -> dispatcher.dispatchToolCalls(input)); + ScheduledFuture heartbeat = scheduleHeartbeat(item, fence, leaseLost, dispatchFuture); + try { + Map output; + try { + output = dispatchFuture.get(); + } catch (CancellationException | InterruptedException | ExecutionException e) { + if (Thread.interrupted()) { + // restore the flag for the caller; we're exiting this item anyway + Thread.currentThread().interrupt(); + } + // Lease-lost gate FIRST: if the heartbeat aborted us, this is a no-op, + // never a failure (the new claimant owns the item now). + if (leaseLost.get()) { + LOG.log(Level.INFO, () -> "lease lost mid-dispatch; aborting item " + item.id() + " (not completing)"); + return ItemResult.LOST_LEASE; + } + // A genuine tool / dispatch failure -> fail the item (mirror Python). + 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; + } + + // Dispatch succeeded. If the lease was lost during/just after it, do NOT + // complete (proactive M3 gate; the 409 backstop covers the residual race). + if (leaseLost.get()) { + LOG.log(Level.INFO, () -> "lease lost; not completing item " + item.id()); + return ItemResult.LOST_LEASE; + } + + long durationMs = (System.nanoTime() - startNanos) / 1_000_000L; + return complete(item, output, durationMs, fence); + } finally { + heartbeat.cancel(true); + } + } + + /** 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 + // state_patch: the engine merge-patches it into state (top-level keys replaced), + // replacing state["messages"] for the next model turn. Identical to the Python + // python_tool worker (return-as-state_patch). + String genAiModel = stringOrNull(output.get("gen_ai_model")); + String finishReason = stringOrNull(output.get("finish_reason")); + try { + client.completeWorkItem( + 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); + LOG.log(Level.DEBUG, () -> "completed item " + item.id() + " in " + durationMs + "ms"); + return ItemResult.COMPLETED; + } catch (JamjetHttpException e) { + if (e.isConflict()) { + // 409: our fence no longer matches -> the lease was reclaimed and a NEW + // worker owns this item. Do NOT fail (that would clobber the reclaimed + // work). Treat the lost lease as a no-op. Mirrors the Python worker. + LOG.log(Level.INFO, () -> "completion rejected (409); lease lost for item " + item.id()); + return ItemResult.LOST_LEASE; + } + // 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; + } + } + + /** + * Schedule the lease-renewal heartbeat. On a definitive engine rejection + * ({@code status >= 400}) it flags the lease lost and cancels the in-flight + * dispatch (M3); a transport-level blip is logged and tolerated. + */ + private ScheduledFuture scheduleHeartbeat(ClaimedWorkItem item, long fence, + AtomicBoolean leaseLost, Future dispatchFuture) { + return heartbeats.scheduleAtFixedRate(() -> { + if (leaseLost.get()) { + return; + } + try { + client.heartbeatWorkItem(item.id(), workerId, fence); + } catch (JamjetHttpException e) { + if (e.statusCode() >= 400) { + // Definitive engine rejection: the lease is gone (FenceLost -> 500, + // or a 4xx). Abort the in-flight tool and stop renewing (M3). + LOG.log(Level.WARNING, () -> "heartbeat rejected for item " + item.id() + + " (HTTP " + e.statusCode() + "); lease lost -> aborting dispatch"); + leaseLost.set(true); + dispatchFuture.cancel(true); + } else { + // Transport-level blip (no HTTP status): tolerate; the complete-time + // fence is the authoritative backstop. + LOG.log(Level.WARNING, () -> "heartbeat transport error for item " + item.id() + ": " + e.getMessage()); + } + } catch (Exception e) { + LOG.log(Level.WARNING, () -> "heartbeat error for item " + item.id() + ": " + e); + } + }, heartbeatMillis, heartbeatMillis, TimeUnit.MILLISECONDS); + } + + // -- helpers ---------------------------------------------------------------- + + @SuppressWarnings("unchecked") + private static Map asMap(Object value) { + if (value instanceof Map m) { + return (Map) m; + } + return Map.of(); + } + + private static String stringOrNull(Object value) { + if (value == null) { + return null; + } + String s = String.valueOf(value); + return s.isBlank() ? null : s; + } + + /** Sleep, returning true if interrupted (so the caller can exit cleanly). */ + private static boolean sleep(long millis) { + try { + Thread.sleep(millis); + return false; + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return true; + } + } + + @Override + public void close() { + stopped = true; + heartbeats.shutdownNow(); + dispatchPool.shutdownNow(); + } +} diff --git a/jamjet-agent/src/test/java/dev/jamjet/agent/tools/ToolDispatcherTest.java b/jamjet-agent/src/test/java/dev/jamjet/agent/tools/ToolDispatcherTest.java new file mode 100644 index 0000000..cc4911c --- /dev/null +++ b/jamjet-agent/src/test/java/dev/jamjet/agent/tools/ToolDispatcherTest.java @@ -0,0 +1,186 @@ +package dev.jamjet.agent.tools; + +import dev.jamjet.agent.TestTools; +import dev.jamjet.agent.Tool; +import org.junit.jupiter.api.Test; + +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Unit tests for {@link ToolDispatcher} — the registry-gated tool-call dispatcher + * the {@code JavaFn} nodes point at. Mirrors the Python {@code dispatch_tool_calls} + * INPUT/OUTPUT contract: reads {@code last_model_tool_calls} / {@code last_model_output} + * (with the {@code tool_calls} / {@code assistant_content} fallbacks) from state, and + * returns {@code {"messages": }} (assistant message + one tool message + * per call). Includes the adversarial cases: arg coercion, unknown-tool clean error, and + * a registered tool that throws. + */ +class ToolDispatcherTest { + + /** A tool that always throws, to prove a registered failure propagates. */ + public static final class BoomTool { + @Tool(name = "boom", description = "always throws") + public String boom(String why) { + throw new IllegalStateException("kaboom: " + why); + } + } + + @SuppressWarnings("unchecked") + private static List> messagesOf(Map out) { + assertThat(out).containsKey("messages"); + return (List>) out.get("messages"); + } + + private static Map toolCall(String id, String name, Map args) { + Map c = new LinkedHashMap<>(); + c.put("id", id); + c.put("name", name); + c.put("arguments", args); + return c; + } + + @Test + void dispatchesRegisteredToolAndAppendsAssistantThenToolMessage() { + var dispatcher = new ToolDispatcher(ToolRegistry.of(new TestTools.WebSearchTool())); + + Map state = new LinkedHashMap<>(); + state.put("last_model_output", "let me search"); + state.put("last_model_tool_calls", + List.of(toolCall("tc1", "web_search", Map.of("query", "jamjet")))); + + List> messages = messagesOf(dispatcher.dispatchToolCalls(state)); + + // [0] the assistant message that requested the call (OpenAI shape). + assertThat(messages).hasSize(2); + Map assistant = messages.get(0); + assertThat(assistant).containsEntry("role", "assistant").containsEntry("content", "let me search"); + @SuppressWarnings("unchecked") + List> calls = (List>) assistant.get("tool_calls"); + assertThat(calls).hasSize(1); + assertThat(calls.get(0)).containsEntry("id", "tc1").containsEntry("type", "function"); + @SuppressWarnings("unchecked") + Map fn = (Map) calls.get(0).get("function"); + assertThat(fn).containsEntry("name", "web_search"); + assertThat((String) fn.get("arguments")).contains("jamjet"); // arguments are a JSON string + + // [1] the tool result message. + Map toolMsg = messages.get(1); + assertThat(toolMsg) + .containsEntry("role", "tool") + .containsEntry("tool_call_id", "tc1") + .containsEntry("name", "web_search") + .containsEntry("content", "results for jamjet"); + } + + @Test + void coercesJsonArgsToTypedMethodParameters() { + // add_numbers(int a, int b) with JSON numbers -> "5". + var dispatcher = new ToolDispatcher(ToolRegistry.of(new TestTools.MathTool())); + + Map state = Map.of( + "last_model_tool_calls", + List.of(toolCall("tc1", "add_numbers", Map.of("a", 2, "b", 3)))); + + List> messages = messagesOf(dispatcher.dispatchToolCalls(state)); + assertThat(messages.get(1)).containsEntry("content", "5"); + } + + @Test + void argumentsSuppliedAsJsonStringAreParsed() { + var dispatcher = new ToolDispatcher(ToolRegistry.of(new TestTools.WebSearchTool())); + + Map call = new LinkedHashMap<>(); + call.put("id", "tc1"); + call.put("name", "web_search"); + call.put("arguments", "{\"query\":\"from-string\"}"); // a JSON STRING, not an object + + Map state = Map.of("last_model_tool_calls", List.of(call)); + List> messages = messagesOf(dispatcher.dispatchToolCalls(state)); + assertThat(messages.get(1)).containsEntry("content", "results for from-string"); + } + + @Test + void unknownToolYieldsCleanErrorMessageNotExceptionAndNoInvocation() { + // RCE gate at the name-resolution level: an unregistered name is never turned + // into a class/method; it surfaces a clean tool error to the model. + var dispatcher = new ToolDispatcher(ToolRegistry.of(new TestTools.WebSearchTool())); + + Map state = Map.of( + "last_model_tool_calls", + List.of(toolCall("tc1", "java.lang.Runtime", Map.of("cmd", "rm -rf /")))); + + Map out = dispatcher.dispatchToolCalls(state); // must NOT throw + List> messages = messagesOf(out); + assertThat(messages.get(1)) + .containsEntry("role", "tool") + .containsEntry("tool_call_id", "tc1"); + assertThat((String) messages.get(1).get("content")) + .contains("not a registered @Tool") + .contains("java.lang.Runtime"); + } + + @Test + void registeredToolThatThrowsPropagatesAsToolInvocationException() { + var dispatcher = new ToolDispatcher(ToolRegistry.of(new BoomTool())); + + Map state = Map.of( + "last_model_tool_calls", + List.of(toolCall("tc1", "boom", Map.of("why", "test")))); + + assertThatThrownBy(() -> dispatcher.dispatchToolCalls(state)) + .isInstanceOf(ToolInvocationException.class) + .hasMessageContaining("boom") + .hasRootCauseInstanceOf(IllegalStateException.class); + } + + @Test + void preservesExistingMessagesAndOrdersMultipleToolCalls() { + var dispatcher = new ToolDispatcher( + ToolRegistry.of(new TestTools.WebSearchTool(), new TestTools.MathTool())); + + Map prior = Map.of("role", "user", "content", "hi"); + Map state = new LinkedHashMap<>(); + state.put("messages", List.of(prior)); + state.put("last_model_tool_calls", List.of( + toolCall("a", "web_search", Map.of("query", "q")), + toolCall("b", "add_numbers", Map.of("a", 10, "b", 1)))); + + List> messages = messagesOf(dispatcher.dispatchToolCalls(state)); + // prior user msg + 1 assistant + 2 tool msgs, in order. + assertThat(messages).hasSize(4); + assertThat(messages.get(0)).containsEntry("role", "user"); + assertThat(messages.get(1)).containsEntry("role", "assistant"); + assertThat(messages.get(2)).containsEntry("tool_call_id", "a").containsEntry("content", "results for q"); + assertThat(messages.get(3)).containsEntry("tool_call_id", "b").containsEntry("content", "11"); + } + + @Test + void prefersExplicitToolCallsAndAssistantContentFallbackKeys() { + var dispatcher = new ToolDispatcher(ToolRegistry.of(new TestTools.WebSearchTool())); + + // Both keys present: tool_calls / assistant_content take precedence. + Map state = new LinkedHashMap<>(); + state.put("assistant_content", "preferred"); + state.put("last_model_output", "ignored"); + state.put("tool_calls", List.of(toolCall("tc1", "web_search", Map.of("query", "x")))); + state.put("last_model_tool_calls", List.of()); // would be empty if used + + List> messages = messagesOf(dispatcher.dispatchToolCalls(state)); + assertThat(messages).hasSize(2); + assertThat(messages.get(0)).containsEntry("content", "preferred"); + assertThat(messages.get(1)).containsEntry("content", "results for x"); + } + + @Test + void emptyInputYieldsOnlyTheAssistantMessage() { + var dispatcher = new ToolDispatcher(ToolRegistry.of(new TestTools.WebSearchTool())); + List> messages = messagesOf(dispatcher.dispatchToolCalls(Map.of())); + assertThat(messages).hasSize(1); + assertThat(messages.get(0)).containsEntry("role", "assistant").containsEntry("content", ""); + } +} 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 new file mode 100644 index 0000000..6b5b31a --- /dev/null +++ b/jamjet-agent/src/test/java/dev/jamjet/agent/worker/JavaToolWorkerTest.java @@ -0,0 +1,302 @@ +package dev.jamjet.agent.worker; + +import com.github.tomakehurst.wiremock.WireMockServer; +import dev.jamjet.agent.TestTools; +import dev.jamjet.agent.Tool; +import dev.jamjet.agent.client.JamjetEngineClient; +import dev.jamjet.agent.tools.ToolDispatcher; +import dev.jamjet.agent.tools.ToolRegistry; +import dev.jamjet.runtime.core.JamjetJson; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +import java.time.Duration; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.function.BooleanSupplier; + +import static com.github.tomakehurst.wiremock.client.WireMock.aResponse; +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; +import static com.github.tomakehurst.wiremock.client.WireMock.ok; +import static com.github.tomakehurst.wiremock.client.WireMock.okJson; +import static com.github.tomakehurst.wiremock.client.WireMock.post; +import static com.github.tomakehurst.wiremock.client.WireMock.postRequestedFor; +import static com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo; +import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.options; +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Adversarial integration tests for {@link JavaToolWorker} against a WireMock stub of + * the engine's bare work-item routes. This is the enforcement surface (a stale/zombie + * worker must not complete a reclaimed item), so every guard is exercised end-to-end + * over real HTTP: + * + *
    + *
  • Happy path — claim → dispatch a registered {@code @Tool} + * → complete WITH the lease fence (asserted in the request body).
    • + *
  • Fence-lost mid-flight (M3) — a heartbeat is rejected while a tool + * runs → the worker aborts the tool and NEVER calls complete.
    • + *
  • 409 on complete = lost lease — the worker does NOT fail the item.
    • + *
  • Unknown dispatch coordinate = clean fail, not RCE — a forged + * {@code java_fn} payload class is failed cleanly, never {@code Class.forName}'d.
    • + *
  • Arg coercion — JSON args coerce to typed params end-to-end.
    • + *
+ */ +class JavaToolWorkerTest { + + WireMockServer wm; + + @BeforeEach + void up() { + wm = new WireMockServer(options().dynamicPort()); + wm.start(); + } + + @AfterEach + void down() { + wm.stop(); + } + + // -- fixtures --------------------------------------------------------------- + + /** A tool that blocks until interrupted, so the M3 heartbeat abort can hit it mid-flight. */ + public static final class SlowTool { + final CountDownLatch started = new CountDownLatch(1); + volatile boolean interrupted = false; + + @Tool(name = "slow_tool", description = "blocks until interrupted") + public String slow(String ignored) { + started.countDown(); + try { + Thread.sleep(5_000); + } catch (InterruptedException e) { + interrupted = true; + Thread.currentThread().interrupt(); + throw new RuntimeException("interrupted while running slow_tool"); + } + return "should never be reached"; + } + } + + private static Map toolCall(String id, String name, Map args) { + Map c = new LinkedHashMap<>(); + c.put("id", id); + c.put("name", name); + c.put("arguments", args); + return c; + } + + /** Serialize a {@code {claimed, work_item}} claim response with the given payload + fence. */ + private static String claimBody(String dispatchClass, String dispatchMethod, + Map input, long fence) { + Map payload = new LinkedHashMap<>(); + payload.put("class", dispatchClass); + payload.put("method", dispatchMethod); + payload.put("input", input); + + Map wi = new LinkedHashMap<>(); + wi.put("id", "wi_1"); + wi.put("execution_id", "ex_1"); + wi.put("node_id", "__tools_0__"); + wi.put("queue_type", "java_tool"); + wi.put("payload", payload); + wi.put("attempt", 1); + wi.put("lease_fence", fence); + + Map resp = new LinkedHashMap<>(); + resp.put("claimed", true); + resp.put("work_item", wi); + try { + return JamjetJson.shared().writeValueAsString(resp); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + /** A claim carrying the FIXED, valid dispatch coordinate the Agent builder emits. */ + private static String validClaim(Map input, long fence) { + return claimBody(ToolDispatcher.DISPATCH_CLASS, ToolDispatcher.DISPATCH_METHOD, input, fence); + } + + private static boolean waitFor(BooleanSupplier cond, long timeoutMillis) { + long deadline = System.currentTimeMillis() + timeoutMillis; + while (System.currentTimeMillis() < deadline) { + if (cond.getAsBoolean()) { + return true; + } + try { + Thread.sleep(5); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return cond.getAsBoolean(); + } + } + return cond.getAsBoolean(); + } + + // -- tests ------------------------------------------------------------------ + + @Test + @Timeout(15) + void happyPathDispatchesToolAndCompletesWithTheFence() { + Map input = Map.of( + "last_model_output", "searching", + "last_model_tool_calls", List.of(toolCall("tc1", "web_search", Map.of("query", "jamjet")))); + + wm.stubFor(post(urlEqualTo("/work-items/claim")).willReturn(okJson(validClaim(input, 7)))); + wm.stubFor(post(urlEqualTo("/work-items/wi_1/complete")).willReturn(ok())); + wm.stubFor(post(urlEqualTo("/work-items/wi_1/heartbeat")).willReturn(ok())); + 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))) { + + JavaToolWorker.ItemResult result = worker.runOnce(); + assertThat(result).isEqualTo(JavaToolWorker.ItemResult.COMPLETED); + } + + // The completion MUST echo the claim's lease fence (the fenced completion path) + // and carry the dispatcher's {messages:[assistant, tool]} as output + state_patch. + wm.verify(postRequestedFor(urlEqualTo("/work-items/wi_1/complete")) + .withRequestBody(matchingJsonPath("$.lease_fence", equalTo("7"))) + .withRequestBody(matchingJsonPath("$.execution_id", equalTo("ex_1"))) + .withRequestBody(matchingJsonPath("$.node_id", equalTo("__tools_0__"))) + .withRequestBody(matchingJsonPath("$.output.messages[0].role", equalTo("assistant"))) + .withRequestBody(matchingJsonPath("$.output.messages[1].role", equalTo("tool"))) + .withRequestBody(matchingJsonPath("$.output.messages[1].content", equalTo("results for jamjet"))) + .withRequestBody(matchingJsonPath("$.state_patch.messages[1].content", equalTo("results for jamjet")))); + // No failure on the happy path. + wm.verify(0, postRequestedFor(urlEqualTo("/work-items/wi_1/fail"))); + } + + @Test + @Timeout(15) + void argsCoerceToTypedParametersEndToEnd() { + Map input = Map.of( + "last_model_tool_calls", List.of(toolCall("tc1", "add_numbers", Map.of("a", 2, "b", 3)))); + + wm.stubFor(post(urlEqualTo("/work-items/claim")).willReturn(okJson(validClaim(input, 5)))); + 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.MathTool()), + Duration.ofSeconds(2), Duration.ofMillis(10))) { + assertThat(worker.runOnce()).isEqualTo(JavaToolWorker.ItemResult.COMPLETED); + } + + // add_numbers(int,int) with JSON {a:2,b:3} -> "5". + wm.verify(postRequestedFor(urlEqualTo("/work-items/wi_1/complete")) + .withRequestBody(matchingJsonPath("$.lease_fence", equalTo("5"))) + .withRequestBody(matchingJsonPath("$.output.messages[1].content", equalTo("5")))); + } + + @Test + @Timeout(15) + void fenceLostMidFlightAbortsAndNeverCompletes_M3() { + Map input = Map.of( + "last_model_tool_calls", List.of(toolCall("tc1", "slow_tool", Map.of("ignored", "x")))); + + // Heartbeat is rejected: the engine surfaces a stale fence on renew as HTTP 500 + // (FenceLost -> Internal). The worker must abort the in-flight tool and NOT complete. + wm.stubFor(post(urlEqualTo("/work-items/claim")).willReturn(okJson(validClaim(input, 7)))); + wm.stubFor(post(urlEqualTo("/work-items/wi_1/heartbeat")) + .willReturn(aResponse().withStatus(500).withBody("{\"error\":\"FenceLost(wi_1)\"}"))); + wm.stubFor(post(urlEqualTo("/work-items/wi_1/complete")).willReturn(ok())); + wm.stubFor(post(urlEqualTo("/work-items/wi_1/fail")).willReturn(ok())); + + SlowTool slow = new SlowTool(); + try (var client = new JamjetEngineClient(wm.baseUrl()); + var worker = new JavaToolWorker(client, "w1", + ToolRegistry.of(slow), + Duration.ofMillis(20), Duration.ofMillis(10))) { + + JavaToolWorker.ItemResult result = worker.runOnce(); + + // Lost lease is a no-op, NOT a completion and NOT a failure. + assertThat(result).isEqualTo(JavaToolWorker.ItemResult.LOST_LEASE); + } + + // The in-flight tool was actually signalled to abort (M3). + assertThat(waitFor(() -> slow.interrupted, 2_000)).as("slow tool was interrupted").isTrue(); + // A heartbeat fired (and was rejected). + wm.verify(postRequestedFor(urlEqualTo("/work-items/wi_1/heartbeat"))); + // CRITICAL: the reclaimed item was NEVER completed and NEVER failed by this worker. + wm.verify(0, postRequestedFor(urlEqualTo("/work-items/wi_1/complete"))); + wm.verify(0, postRequestedFor(urlEqualTo("/work-items/wi_1/fail"))); + } + + @Test + @Timeout(15) + void conflictOnCompleteIsLostLeaseNotFailure() { + 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)))); + // The completion is fence-rejected: a reclaimed lease yields HTTP 409. + wm.stubFor(post(urlEqualTo("/work-items/wi_1/complete")) + .willReturn(aResponse().withStatus(409).withBody("{\"reason\":\"stale or invalid lease fence\"}"))); + wm.stubFor(post(urlEqualTo("/work-items/wi_1/heartbeat")).willReturn(ok())); + 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.LOST_LEASE); + } + + // It attempted the (fenced) complete... + wm.verify(postRequestedFor(urlEqualTo("/work-items/wi_1/complete")) + .withRequestBody(matchingJsonPath("$.lease_fence", equalTo("9")))); + // ...but a 409 must NOT escalate to a fail (that would clobber the new claimant). + wm.verify(0, postRequestedFor(urlEqualTo("/work-items/wi_1/fail"))); + } + + @Test + @Timeout(15) + void unknownDispatchCoordinateFailsCleanlyAndIsNotRce() { + // A FORGED java_fn payload names an arbitrary class/method. The worker must + // fail it cleanly and NEVER Class.forName / invoke it. + Map input = Map.of("last_model_tool_calls", List.of()); + String forged = claimBody("java.lang.Runtime", "exec", input, 3); + + wm.stubFor(post(urlEqualTo("/work-items/claim")).willReturn(okJson(forged))); + wm.stubFor(post(urlEqualTo("/work-items/wi_1/fail")).willReturn(ok())); + wm.stubFor(post(urlEqualTo("/work-items/wi_1/complete")).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); + } + + // Failed with a clear message naming the rejected coordinate; the arbitrary + // class was NEVER invoked or completed. + wm.verify(postRequestedFor(urlEqualTo("/work-items/wi_1/fail")) + .withRequestBody(matchingJsonPath("$.error", containing("unsupported java_fn dispatch coordinate")))); + wm.verify(0, postRequestedFor(urlEqualTo("/work-items/wi_1/complete"))); + } + + @Test + @Timeout(15) + void emptyQueueReturnsEmpty() { + wm.stubFor(post(urlEqualTo("/work-items/claim")).willReturn(okJson("{\"claimed\":false}"))); + 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.EMPTY); + } + } +} From 005b1bf84677c45d2f3549995ebca4fc7a6fa0d0 Mon Sep 17 00:00:00 2001 From: sdev Date: Mon, 29 Jun 2026 00:58:32 +0530 Subject: [PATCH 4/7] feat(agent): Agent.runDurable + durable result extraction (B-4) Add the durable run path to the Java Agent builder, mirroring the Python Agent.run_durable: compile -> create_workflow -> start_execution (seeded with the system+user messages and the tool-resolver map) -> poll get_execution to a terminal state -> extract the final assistant text + tool-call trace into an AgentResult. - AgentResult (output + toolCalls + raw terminal state) and a nested ToolCall record, parallel to the Python result shape. - RunOptions (maxTurns / runtimeUrl / auth / pollInterval / timeout) with defaults and withers; timeout derives from the agent when unset. - AgentRunException (non-completed terminal: failed / cancelled / limit_exceeded) and AgentRunTimeoutException, mirroring the Python RuntimeError / TimeoutError contract. - DurableRunner orchestrates create/start/poll/extract. Terminal statuses are the snake_case the engine serializes (completed / failed / cancelled / limit_exceeded). The answer is current_state.last_model_output (written inline by the Rust Model executor's state_patch), falling back to the last assistant message; tool calls reconstruct from current_state.messages. The existing ExecutionState.currentState carries both, so no record change. - AgentIrCompiler.buildInitialState / toolsMap seed the run input. - DRY (B-3 review): AgentIrCompiler.DISPATCH_CLASS/METHOD now reference the ToolDispatcher constants the worker's RCE gate checks, so the coordinate the builder emits and the coordinate the worker accepts are one literal. Tests: WireMock-driven create/start/poll/extract, the last-assistant fallback, terminal failed/limit_exceeded surfacing, the poll timeout, and the client-owning vs client-injecting overloads. The parity golden stays green after the DRY fix. --- .../src/main/java/dev/jamjet/agent/Agent.java | 62 ++++ .../dev/jamjet/agent/AgentIrCompiler.java | 56 +++- .../java/dev/jamjet/agent/AgentResult.java | 63 +++++ .../dev/jamjet/agent/AgentRunException.java | 29 ++ .../agent/AgentRunTimeoutException.java | 18 ++ .../java/dev/jamjet/agent/DurableRunner.java | 266 ++++++++++++++++++ .../java/dev/jamjet/agent/RunOptions.java | 81 ++++++ .../dev/jamjet/agent/AgentRunDurableTest.java | 243 ++++++++++++++++ 8 files changed, 816 insertions(+), 2 deletions(-) create mode 100644 jamjet-agent/src/main/java/dev/jamjet/agent/AgentResult.java create mode 100644 jamjet-agent/src/main/java/dev/jamjet/agent/AgentRunException.java create mode 100644 jamjet-agent/src/main/java/dev/jamjet/agent/AgentRunTimeoutException.java create mode 100644 jamjet-agent/src/main/java/dev/jamjet/agent/DurableRunner.java create mode 100644 jamjet-agent/src/main/java/dev/jamjet/agent/RunOptions.java create mode 100644 jamjet-agent/src/test/java/dev/jamjet/agent/AgentRunDurableTest.java diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/Agent.java b/jamjet-agent/src/main/java/dev/jamjet/agent/Agent.java index f6615d2..7393d8d 100644 --- a/jamjet-agent/src/main/java/dev/jamjet/agent/Agent.java +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/Agent.java @@ -1,5 +1,6 @@ package dev.jamjet.agent; +import dev.jamjet.agent.client.JamjetEngineClient; import dev.jamjet.agent.tools.ToolRegistry; import dev.jamjet.runtime.core.ir.PolicySetIr; import dev.jamjet.runtime.core.ir.WorkflowIr; @@ -87,6 +88,67 @@ public WorkflowIr compileToIr(int maxTurns) { return AgentIrCompiler.compile(this, maxTurns); } + // -- durable run (B-4) ------------------------------------------------------ + + /** + * Run this agent durably on the JamJet engine with the default {@link RunOptions} + * (local runtime, {@link #DEFAULT_MAX_TURNS} turns), returning the final assistant + * text + tool-call trace as an {@link AgentResult}. The Java mirror of the Python + * {@code Agent.run_durable}. + * + *

This compiles the agent to the agent-loop {@link WorkflowIr}, registers it + * ({@code POST /workflows}), starts an execution seeded with the system+user + * {@code messages} ({@code POST /executions}), polls {@code GET /executions/{id}} + * to a terminal state, and extracts the answer from the terminal + * {@code current_state.last_model_output} (falling back to the last assistant + * message). Every model call and tool dispatch runs through the durable engine, so + * the run is event-sourced, replayable, idempotent, and governed (budget / policy / + * PII enforced fail-closed from the compiled IR). + * + *

Required running services (mirrors the Python {@code run_durable})

+ * A durable run is NOT self-contained — three services must be running: + *
    + *
  1. the JamJet engine at {@link RunOptions#runtimeUrl()} (the + * {@code jamjet-server} that owns the {@code java_tool} queue);
    2. + *
  2. the model sidecar ({@code JAMJET_MODEL_SEAM_URL}) — the engine routes + * every governed model call through it (no Java model code);
    3. + *
  3. a {@link dev.jamjet.agent.worker.JavaToolWorker} draining the + * {@code java_tool} queue with THIS agent's tool registry, so the + * {@code @Tool} methods execute durably exactly-once. Run it in a separate + * thread/process; {@code runDurable} does not start one.
    4. + *
+ * + * @throws AgentRunException if the run reaches a non-{@code completed} terminal + * state ({@code failed} / {@code cancelled} / + * {@code limit_exceeded}) + * @throws AgentRunTimeoutException if no terminal state is reached before the deadline + */ + public AgentResult runDurable(String prompt) { + return runDurable(prompt, RunOptions.defaults()); + } + + /** + * Run this agent durably with the given {@link RunOptions}, building (and closing) a + * {@link JamjetEngineClient} for {@link RunOptions#runtimeUrl()}. See + * {@link #runDurable(String)} for the running-services contract. + */ + public AgentResult runDurable(String prompt, RunOptions options) { + try (JamjetEngineClient client = + new JamjetEngineClient(options.runtimeUrl(), options.bearerToken(), options.tenantId())) { + return runDurable(prompt, client, options); + } + } + + /** + * Run this agent durably over a caller-provided {@link JamjetEngineClient}. The + * caller owns the client's lifecycle (this overload does NOT close it), so a run + * and a {@link dev.jamjet.agent.worker.JavaToolWorker} can share one client against + * the same engine. See {@link #runDurable(String)} for the running-services contract. + */ + public AgentResult runDurable(String prompt, JamjetEngineClient client, RunOptions options) { + return DurableRunner.run(this, prompt, client, options); + } + // -- accessors (read by AgentIrCompiler) ------------------------------------ public String name() { diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/AgentIrCompiler.java b/jamjet-agent/src/main/java/dev/jamjet/agent/AgentIrCompiler.java index 0bf09bd..25f248a 100644 --- a/jamjet-agent/src/main/java/dev/jamjet/agent/AgentIrCompiler.java +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/AgentIrCompiler.java @@ -2,6 +2,8 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.SerializationFeature; +import dev.jamjet.agent.tools.RegisteredTool; +import dev.jamjet.agent.tools.ToolDispatcher; import dev.jamjet.runtime.core.JamjetJson; import dev.jamjet.runtime.core.TimeoutConfig; import dev.jamjet.runtime.core.ir.ConditionalBranch; @@ -57,9 +59,15 @@ final class AgentIrCompiler { * from state and fans out to the requested {@code @Tool} methods via the * {@link dev.jamjet.agent.tools.ToolRegistry}. (A single dispatcher per turn, * not one node per dynamic call, mirrors the Python loop exactly.) + * + *

Single source of truth. These reference {@link ToolDispatcher}'s own + * constants — the SAME constants the {@link dev.jamjet.agent.worker.JavaToolWorker} + * RCE gate checks a claimed {@code java_fn} payload against. So the coordinate the + * builder emits and the coordinate the worker accepts are one + * literal and can never silently diverge (a B-3-review DRY fix). */ - static final String DISPATCH_CLASS = "dev.jamjet.agent.tools.ToolDispatcher"; - static final String DISPATCH_METHOD = "dispatchToolCalls"; + static final String DISPATCH_CLASS = ToolDispatcher.DISPATCH_CLASS; + static final String DISPATCH_METHOD = ToolDispatcher.DISPATCH_METHOD; /** The condition the tool gate branches on (matches the Model executor's recorded finish reason). */ static final String TOOL_CALLS_EXPR = "state.last_model_finish_reason == \"tool_calls\""; @@ -282,6 +290,50 @@ static String litellmModel(String model) { return raw; } + // -- durable-run initial state (mirrors agent_ir.build_initial_state) -------- + + /** + * The execution {@code initial_input} for a compiled agent-loop IR, mirroring the + * Python {@code build_initial_state}: it seeds the running {@code messages} + * (system + user prompt) plus the {@code {name: "class#method"}} tool-resolver map + * into workflow state, so the first Model node reads the messages and every + * {@code java_fn} tool node finds them. The Rust engine copies this verbatim into + * {@code current_state} at {@code start_execution} time. + * + *

The prompt is per-run and private, so it is seeded HERE (not embedded in the + * workflow definition) — the compiled IR's {@code description} never carries it. + */ + static Map buildInitialState(Agent agent, String prompt) { + String system = blankToNull(agent.instructions()) != null + ? agent.instructions() + : "You are a helpful assistant."; + + List> messages = List.of( + Map.of("role", "system", "content", system), + Map.of("role", "user", "content", prompt == null ? "" : prompt)); + + Map out = new LinkedHashMap<>(); + out.put("messages", messages); + out.put("tools", toolsMap(agent)); + return out; + } + + /** + * {@code {tool_name: "class#method"}} — the Java analog of Python's {@code _tools_map} + * ({@code {name: "module:qualname"}}). Carried in the seeded state for shape-parity + * with the Python durable run; the Java {@code JavaToolWorker} resolves tools by name + * through its own {@link dev.jamjet.agent.tools.ToolRegistry}, so it does not consume + * this map (the dispatch coordinate is the same {@code class#method}, see + * {@link RegisteredTool#key()}). + */ + static Map toolsMap(Agent agent) { + Map map = new LinkedHashMap<>(); + for (RegisteredTool t : agent.registry().tools()) { + map.put(t.name(), t.key()); + } + return map; + } + /** * The immutable content-version: {@code "0.1.0+" + sha256(canonical_ir)[:12]}. * The hash differs from Python's (the Java IR has {@code java_fn} tool nodes diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/AgentResult.java b/jamjet-agent/src/main/java/dev/jamjet/agent/AgentResult.java new file mode 100644 index 0000000..49167d0 --- /dev/null +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/AgentResult.java @@ -0,0 +1,63 @@ +package dev.jamjet.agent; + +import dev.jamjet.agent.client.ExecutionState; + +import java.util.List; +import java.util.Map; + +/** + * The result of a durable agent run ({@link Agent#runDurable(String)}) — the Java + * analog of the Python {@code jamjet.agents.agent.AgentResult}. It carries the SAME + * load-bearing fields the Python durable path extracts: + * + *

    + *
  • {@link #output()} — the final assistant text. On the durable engine this is + * the terminal {@code current_state.last_model_output} (written inline by the + * Rust Model executor's {@code state_patch}); it falls back to the last + * assistant message's content in {@code current_state.messages}.
    • + *
  • {@link #toolCalls()} — the tool-call trace, reconstructed from the + * accumulated {@code messages} (each {@code role: tool} result paired with the + * arguments of the assistant {@code tool_call} that requested it, matched by + * {@code tool_call_id}), mirroring the Python {@code _tool_calls_from_messages}.
    • + *
  • {@link #terminalState()} — the raw terminal execution snapshot + * ({@code GET /executions/{id}}), so callers can inspect the full + * {@code current_state} / status without re-fetching.
    • + *
+ * + *

Governance artifacts (budget enforcement, policy gates, PII redaction, the audit + * log) ride in the durable engine itself — they are compiled into the {@link Agent}'s + * {@code WorkflowIr} (see {@link Agent#compileToIr()}) and enforced fail-closed by the + * Rust engine, not re-derived here. The terminal state reflects their outcome (e.g. a + * {@code limit_exceeded} status surfaces as an {@link AgentRunException}). + */ +public record AgentResult( + String output, + List toolCalls, + ExecutionState terminalState +) { + public AgentResult { + output = output == null ? "" : output; + toolCalls = toolCalls == null ? List.of() : List.copyOf(toolCalls); + } + + /** + * One reconstructed tool call from the durable run's message history — the Java + * analog of the Python {@code ToolCallRecord} dict ({@code tool} / {@code input} / + * {@code output}). Per-call durations are not carried in the message history, so + * (unlike the in-process path) no duration is reconstructed here. + * + * @param tool the tool name the model called + * @param input the parsed arguments the model passed (possibly empty) + * @param output the tool's returned content (the {@code role: tool} message body) + */ + public record ToolCall(String tool, Map input, String output) { + public ToolCall { + input = input == null ? Map.of() : Map.copyOf(input); + } + } + + @Override + public String toString() { + return "AgentResult(output=" + output + ", toolCalls=" + toolCalls.size() + ")"; + } +} diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/AgentRunException.java b/jamjet-agent/src/main/java/dev/jamjet/agent/AgentRunException.java new file mode 100644 index 0000000..5776ecc --- /dev/null +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/AgentRunException.java @@ -0,0 +1,29 @@ +package dev.jamjet.agent; + +/** + * Raised when a durable agent run ({@link Agent#runDurable(String)}) reaches a + * terminal state that is not {@code completed} — i.e. {@code failed}, + * {@code cancelled}, or {@code limit_exceeded}. This mirrors the Python + * {@code run_durable}, which raises {@code RuntimeError} for a non-{@code completed} + * terminal state, so a budget breach or a policy denial surfaces loudly rather than + * returning a hollow result. + * + *

{@link #status()} carries the terminal {@link dev.jamjet.runtime.core.workflow} + * status string the engine reported ({@code failed} / {@code cancelled} / + * {@code limit_exceeded}), so callers can branch on the failure mode (e.g. treat a + * {@code limit_exceeded} budget breach differently from a hard {@code failed}). + */ +public class AgentRunException extends RuntimeException { + + private final String status; + + public AgentRunException(String status, String message) { + super(message); + this.status = status; + } + + /** The terminal status the engine reported ({@code failed} / {@code cancelled} / {@code limit_exceeded}). */ + public String status() { + return status; + } +} diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/AgentRunTimeoutException.java b/jamjet-agent/src/main/java/dev/jamjet/agent/AgentRunTimeoutException.java new file mode 100644 index 0000000..ddf9c7e --- /dev/null +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/AgentRunTimeoutException.java @@ -0,0 +1,18 @@ +package dev.jamjet.agent; + +/** + * Raised when a durable agent run ({@link Agent#runDurable(String)}) does not reach a + * terminal state within the poll deadline ({@link RunOptions#timeout()}, or the + * agent's {@link Agent#timeoutSeconds()} when unset). Mirrors the Python + * {@code run_durable}, which raises {@code TimeoutError} when no terminal state is + * reached within {@code self.limits.timeout_seconds}. + * + *

{@link #status()} carries the last status observed before the deadline + * (typically {@code running}), to aid diagnosis of a stalled run. + */ +public final class AgentRunTimeoutException extends AgentRunException { + + public AgentRunTimeoutException(String lastStatus, String message) { + super(lastStatus, message); + } +} diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/DurableRunner.java b/jamjet-agent/src/main/java/dev/jamjet/agent/DurableRunner.java new file mode 100644 index 0000000..94d42bf --- /dev/null +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/DurableRunner.java @@ -0,0 +1,266 @@ +package dev.jamjet.agent; + +import com.fasterxml.jackson.databind.ObjectMapper; +import dev.jamjet.agent.client.ExecutionState; +import dev.jamjet.agent.client.JamjetEngineClient; +import dev.jamjet.agent.client.StartExecutionResult; +import dev.jamjet.runtime.core.JamjetJson; +import dev.jamjet.runtime.core.ir.WorkflowIr; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; + +/** + * Drives a compiled {@link Agent} through one durable run on the JamJet engine — + * the mechanics behind {@link Agent#runDurable(String)}, mirroring the Python + * {@code Agent.run_durable}: compile → {@code create_workflow} → + * {@code start_execution} (seeded with the system+user {@code messages} and the + * tool-resolver map) → poll {@code get_execution} to a terminal state → + * extract the final assistant text + tool-call trace into an {@link AgentResult}. + * + *

Unlike the in-process path, every model call and tool dispatch runs through the + * durable event-sourced engine, so the run gets the event log, replay, idempotency, + * and fail-closed governance enforcement (budget / policy / PII) compiled into the IR. + * + *

Where the terminal result lives

+ * The Rust Model executor writes the final model turn's content inline into + * {@code current_state.last_model_output} (via its {@code state_patch}; never spilled + * to the artifact store), so the canonical answer is {@code last_model_output}. The + * fallback is the last assistant message's content in the accumulated + * {@code current_state.messages} (which the Java tool dispatcher replaces each turn). + * The terminal is always reached via a model turn (the final tool-less + * {@code __model_{maxTurns}__} node), so the answer is a real model turn that saw the + * tool results — exactly as in the Python contract. + */ +final class DurableRunner { + + /** + * Terminal {@code WorkflowStatus} strings, snake_case as the Rust engine + * serializes them ({@code runtime/core/src/workflow.rs}: {@code rename_all = + * "snake_case"}), matching the Python {@code _TERMINAL_STATUSES}. + */ + private static final Set TERMINAL_STATUSES = + Set.of("completed", "failed", "cancelled", "limit_exceeded"); + + /** The single status the run treats as success; every other terminal is an error. */ + private static final String COMPLETED = "completed"; + + private static final ObjectMapper JSON = JamjetJson.shared(); + + private DurableRunner() {} + + /** + * Run {@code agent} on {@code prompt} over {@code client}, returning the extracted + * {@link AgentResult}. The caller owns {@code client}'s lifecycle (this method does + * NOT close it), so a worker and a run can share one client. + * + * @throws AgentRunException if the run reaches a non-{@code completed} terminal state + * @throws AgentRunTimeoutException if no terminal state is reached before the deadline + */ + static AgentResult run(Agent agent, String prompt, JamjetEngineClient client, RunOptions options) { + WorkflowIr ir = agent.compileToIr(options.maxTurns()); + Map initialInput = AgentIrCompiler.buildInitialState(agent, prompt); + + // Register the compiled IR, then start an execution seeded with the running + // messages + tool-resolver map (build_initial_state). The version is the + // content hash so a changed agent never reuses an immutably-cached graph. + client.createWorkflow(ir); + StartExecutionResult started = client.startExecution(ir.workflowId(), initialInput, ir.version()); + String execId = started == null ? null : started.executionId(); + if (execId == null || execId.isBlank()) { + throw new AgentRunException("unknown", + "start_execution returned no execution_id (started=" + started + ")"); + } + + Duration timeout = options.timeout() != null + ? options.timeout() + : Duration.ofSeconds(agent.timeoutSeconds()); + ExecutionState terminal = pollToTerminal(client, execId, options.pollInterval(), timeout); + return extractResult(terminal); + } + + /** Poll {@code get_execution} until a terminal status, mirroring Python {@code _poll_to_terminal}. */ + private static ExecutionState pollToTerminal(JamjetEngineClient client, String execId, + Duration pollInterval, Duration timeout) { + // Monotonic deadline: time spent inside slow get_execution() calls counts + // toward the timeout, so a stalled run can't blow past it. + long deadlineNanos = System.nanoTime() + timeout.toNanos(); + String lastStatus = "unknown"; + while (true) { + ExecutionState execution = client.getExecution(execId); + lastStatus = execution == null ? "unknown" : execution.status(); + if (isTerminal(lastStatus)) { + return execution; + } + if (System.nanoTime() >= deadlineNanos) { + throw new AgentRunTimeoutException(lastStatus, + "durable agent run " + execId + " did not reach a terminal state within " + + timeout.toSeconds() + "s (last status: " + lastStatus + ")"); + } + if (sleep(pollInterval.toMillis())) { + throw new AgentRunTimeoutException(lastStatus, + "durable agent run " + execId + " interrupted while polling (last status: " + lastStatus + ")"); + } + } + } + + /** + * Build an {@link AgentResult} from a terminal execution snapshot, mirroring the + * Python {@code _extract_result}: the answer is {@code current_state.last_model_output} + * (the inline final model turn), falling back to the last assistant message; the + * tool-call trace is reconstructed from {@code current_state.messages}. + * + * @throws AgentRunException if the terminal status is not {@code completed} + */ + static AgentResult extractResult(ExecutionState execution) { + String status = normalize(execution == null ? null : execution.status()); + if (!COMPLETED.equals(status)) { + String detail = failureDetail(execution, status); + throw new AgentRunException(status, + "durable agent run ended in non-completed state: " + detail); + } + + Map state = execution.currentState() == null ? Map.of() : execution.currentState(); + List messages = asList(state.get("messages")); + + Object output = state.get("last_model_output"); + if (output == null) { + output = lastAssistantContent(messages); + } + String outputStr = output == null ? "" : String.valueOf(output); + + return new AgentResult(outputStr, toolCallsFromMessages(messages), execution); + } + + // -- result reconstruction (mirrors agent.py helpers) ----------------------- + + /** The content of the last assistant message that carried text. */ + private static String lastAssistantContent(List messages) { + for (int i = messages.size() - 1; i >= 0; i--) { + Map msg = asMap(messages.get(i)); + if ("assistant".equals(msg.get("role"))) { + Object content = msg.get("content"); + if (content != null && !String.valueOf(content).isEmpty()) { + return String.valueOf(content); + } + } + } + return null; + } + + /** + * Reconstruct the tool-call trace from the message history, mirroring the Python + * {@code _tool_calls_from_messages}: pair each {@code role: tool} result with the + * arguments of the assistant {@code tool_call} that requested it (matched by + * {@code tool_call_id}). + */ + private static List toolCallsFromMessages(List messages) { + Map> argsById = new LinkedHashMap<>(); + for (Object raw : messages) { + Map msg = asMap(raw); + if (!"assistant".equals(msg.get("role"))) { + continue; + } + for (Object callRaw : asList(msg.get("tool_calls"))) { + Map call = asMap(callRaw); + Map function = asMap(call.get("function")); + Map parsed = parseArguments(function.get("arguments")); + Object id = call.get("id"); + if (id != null) { + argsById.put(id, parsed); + } + } + } + + List calls = new ArrayList<>(); + for (Object raw : messages) { + Map msg = asMap(raw); + if (!"tool".equals(msg.get("role"))) { + continue; + } + Object name = msg.get("name"); + Map input = argsById.getOrDefault(msg.get("tool_call_id"), Map.of()); + Object content = msg.get("content"); + calls.add(new AgentResult.ToolCall( + name == null ? null : String.valueOf(name), + input, + content == null ? null : String.valueOf(content))); + } + return calls; + } + + /** Tool-call arguments arrive as a JSON object or a JSON string — normalise to a map. */ + @SuppressWarnings("unchecked") + private static Map parseArguments(Object arguments) { + if (arguments instanceof Map m) { + return (Map) m; + } + if (arguments instanceof String s && !s.isBlank()) { + try { + return JSON.readValue(s, Map.class); + } catch (Exception e) { + return Map.of(); + } + } + return Map.of(); + } + + // -- status / detail -------------------------------------------------------- + + private static boolean isTerminal(String status) { + return TERMINAL_STATUSES.contains(normalize(status)); + } + + /** Lower-case the status so a snake_case engine ({@code completed}) and any capitalized stub agree. */ + private static String normalize(String status) { + return status == null ? "" : status.toLowerCase(Locale.ROOT); + } + + /** Best-effort failure detail for a non-completed terminal: a {@code detail}/{@code error} in state, else the status. */ + private static String failureDetail(ExecutionState execution, String status) { + if (execution != null && execution.currentState() != null) { + Object detail = execution.currentState().get("detail"); + if (detail == null) { + detail = execution.currentState().get("error"); + } + if (detail != null && !String.valueOf(detail).isBlank()) { + return status + " (" + detail + ")"; + } + } + return status; + } + + // -- coercion helpers ------------------------------------------------------- + + @SuppressWarnings("unchecked") + private static List asList(Object value) { + if (value instanceof List list) { + return (List) list; + } + return List.of(); + } + + @SuppressWarnings("unchecked") + private static Map asMap(Object value) { + if (value instanceof Map m) { + return (Map) m; + } + return Map.of(); + } + + /** Sleep, returning true if interrupted (so the poll loop can exit cleanly). */ + private static boolean sleep(long millis) { + try { + Thread.sleep(millis); + return false; + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return true; + } + } +} diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/RunOptions.java b/jamjet-agent/src/main/java/dev/jamjet/agent/RunOptions.java new file mode 100644 index 0000000..01d28d7 --- /dev/null +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/RunOptions.java @@ -0,0 +1,81 @@ +package dev.jamjet.agent; + +import java.time.Duration; + +/** + * Tuning knobs for a durable agent run ({@link Agent#runDurable(String, RunOptions)}), + * mirroring the keyword arguments of the Python {@code Agent.run_durable} + * ({@code max_turns}, {@code runtime_url}). All fields have sensible defaults via + * {@link #defaults()}; override individually with the {@code with*} methods. + * + * @param maxTurns static-unroll bound for the compiled agent loop (default + * {@link Agent#DEFAULT_MAX_TURNS}); matches the Python default. + * @param runtimeUrl the JamJet engine base URL the run targets (default + * {@value #DEFAULT_RUNTIME_URL}); a bare run never targets a + * remote URL, mirroring the Python dev default. + * @param bearerToken optional {@code Authorization: Bearer} token, or {@code null}. + * @param tenantId optional {@code X-Tenant-Id}, or {@code null}. + * @param pollInterval how often to poll {@code GET /executions/{id}} while waiting for + * a terminal state (default {@value #DEFAULT_POLL_INTERVAL_MS} ms, + * matching the Python {@code _POLL_INTERVAL_SECONDS}). + * @param timeout overall poll deadline, or {@code null} to derive it from the + * agent's {@link Agent#timeoutSeconds()} (the Python behaviour: + * {@code self.limits.timeout_seconds}). + */ +public record RunOptions( + int maxTurns, + String runtimeUrl, + String bearerToken, + String tenantId, + Duration pollInterval, + Duration timeout +) { + /** The Python dev default ({@code http://127.0.0.1:7700}); a bare run never goes remote. */ + public static final String DEFAULT_RUNTIME_URL = "http://127.0.0.1:7700"; + + /** The default poll interval in milliseconds (mirrors Python's 0.5s). */ + public static final long DEFAULT_POLL_INTERVAL_MS = 500; + + public RunOptions { + if (maxTurns < 1) { + throw new IllegalArgumentException("maxTurns must be >= 1"); + } + if (runtimeUrl == null || runtimeUrl.isBlank()) { + throw new IllegalArgumentException("runtimeUrl must not be blank"); + } + if (pollInterval == null || pollInterval.isNegative() || pollInterval.isZero()) { + throw new IllegalArgumentException("pollInterval must be positive"); + } + } + + /** The default options: {@code maxTurns=8}, local runtime, no auth, 500ms poll, agent-derived timeout. */ + public static RunOptions defaults() { + return new RunOptions( + Agent.DEFAULT_MAX_TURNS, + DEFAULT_RUNTIME_URL, + null, + null, + Duration.ofMillis(DEFAULT_POLL_INTERVAL_MS), + null); + } + + public RunOptions withMaxTurns(int maxTurns) { + return new RunOptions(maxTurns, runtimeUrl, bearerToken, tenantId, pollInterval, timeout); + } + + public RunOptions withRuntimeUrl(String runtimeUrl) { + return new RunOptions(maxTurns, runtimeUrl, bearerToken, tenantId, pollInterval, timeout); + } + + public RunOptions withAuth(String bearerToken, String tenantId) { + return new RunOptions(maxTurns, runtimeUrl, bearerToken, tenantId, pollInterval, timeout); + } + + public RunOptions withPollInterval(Duration pollInterval) { + return new RunOptions(maxTurns, runtimeUrl, bearerToken, tenantId, pollInterval, timeout); + } + + public RunOptions withTimeout(Duration timeout) { + return new RunOptions(maxTurns, runtimeUrl, bearerToken, tenantId, pollInterval, timeout); + } +} diff --git a/jamjet-agent/src/test/java/dev/jamjet/agent/AgentRunDurableTest.java b/jamjet-agent/src/test/java/dev/jamjet/agent/AgentRunDurableTest.java new file mode 100644 index 0000000..6b951d2 --- /dev/null +++ b/jamjet-agent/src/test/java/dev/jamjet/agent/AgentRunDurableTest.java @@ -0,0 +1,243 @@ +package dev.jamjet.agent; + +import com.github.tomakehurst.wiremock.WireMockServer; +import dev.jamjet.agent.client.JamjetEngineClient; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +import java.time.Duration; + +import static com.github.tomakehurst.wiremock.client.WireMock.equalTo; +import static com.github.tomakehurst.wiremock.client.WireMock.get; +import static com.github.tomakehurst.wiremock.client.WireMock.matchingJsonPath; +import static com.github.tomakehurst.wiremock.client.WireMock.okJson; +import static com.github.tomakehurst.wiremock.client.WireMock.post; +import static com.github.tomakehurst.wiremock.client.WireMock.postRequestedFor; +import static com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo; +import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.options; +import static com.github.tomakehurst.wiremock.stubbing.Scenario.STARTED; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Hermetic integration test for {@link Agent#runDurable} against a WireMock stub of the + * engine's bare routes ({@code /workflows}, {@code /executions}, {@code /executions/{id}}). + * Proves the B-4 durable-run contract end-to-end at the transport boundary: + * + *
    + *
  • create → start → poll → extract — the workflow is created, + * the execution is started with the seeded {@code messages} + tool-resolver map and + * the content-hashed {@code workflow_version}, the run polls to a terminal state, and + * the final assistant text + tool-call trace are extracted into the + * {@link AgentResult}.
    • + *
  • terminal failure surfaces — a {@code failed} / {@code limit_exceeded} + * terminal raises a clear {@link AgentRunException} (mirrors the Python + * {@code RuntimeError}); never a hollow result.
    • + *
  • timeout surfaces — a run that never terminates raises + * {@link AgentRunTimeoutException} (mirrors the Python {@code TimeoutError}).
    • + *
  • last-assistant fallback — when {@code last_model_output} is absent the + * answer falls back to the last assistant message's content.
    • + *
+ * + *

The status strings are the real snake_case the Rust engine emits + * ({@code running} / {@code completed} / {@code failed} / {@code limit_exceeded}; + * {@code runtime/core/src/workflow.rs}), not capitalized stubs. + */ +class AgentRunDurableTest { + + WireMockServer wm; + + @BeforeEach + void up() { + wm = new WireMockServer(options().dynamicPort()); + wm.start(); + } + + @AfterEach + void down() { + wm.stop(); + } + + /** A small governed agent with two @Tool methods (mirrors the golden fixture). */ + private static Agent demoAgent() { + return Agent.builder("research_agent") + .model("anthropic/claude-sonnet-4-6") + .instructions("You are a helpful research assistant.") + .tools(new TestTools.WebSearchTool(), new TestTools.MathTool()) + .budget(new Budget(100_000, 2.5)) + .build(); + } + + /** Run options pinned at the WireMock base URL with a fast poll + bounded timeout. */ + private RunOptions fastOptions() { + return RunOptions.defaults() + .withMaxTurns(2) + .withRuntimeUrl(wm.baseUrl()) + .withPollInterval(Duration.ofMillis(15)) + .withTimeout(Duration.ofSeconds(5)); + } + + @Test + @Timeout(20) + void createStartPollExtractFinalAssistantText() { + Agent agent = demoAgent(); + // The version runDurable will send on start is the content hash of the SAME IR. + String expectedVersion = agent.compileToIr(2).version(); + + wm.stubFor(post(urlEqualTo("/workflows")) + .willReturn(okJson("{\"workflow_id\":\"research_agent\",\"version\":\"" + expectedVersion + "\"}"))); + wm.stubFor(post(urlEqualTo("/executions")) + .willReturn(okJson("{\"execution_id\":\"ex_99\",\"status\":\"running\"}"))); + + // First poll -> running, second poll -> completed (prove the loop actually polls). + wm.stubFor(get(urlEqualTo("/executions/ex_99")).inScenario("poll") + .whenScenarioStateIs(STARTED) + .willReturn(okJson("{\"execution_id\":\"ex_99\",\"status\":\"running\"}")) + .willSetStateTo("done")); + wm.stubFor(get(urlEqualTo("/executions/ex_99")).inScenario("poll") + .whenScenarioStateIs("done") + .willReturn(okJson(completedBody()))); + + AgentResult result; + try (var client = new JamjetEngineClient(wm.baseUrl())) { + result = agent.runDurable("add 2 and 3", client, fastOptions()); + } + + // The answer is the terminal current_state.last_model_output (the inline final turn). + assertThat(result.output()).isEqualTo("The sum of 2 and 3 is 5."); + // The tool-call trace is reconstructed from the accumulated messages. + assertThat(result.toolCalls()).hasSize(1); + AgentResult.ToolCall call = result.toolCalls().get(0); + assertThat(call.tool()).isEqualTo("add_numbers"); + assertThat(call.input()).containsEntry("a", 2).containsEntry("b", 3); + assertThat(call.output()).isEqualTo("5"); + // The raw terminal snapshot is carried through. + assertThat(result.terminalState().status()).isEqualTo("completed"); + + // CREATE wire shape: the compiled agent-loop IR, java_fn tool node, M2 start node. + wm.verify(postRequestedFor(urlEqualTo("/workflows")) + .withRequestBody(matchingJsonPath("$.ir.workflow_id", equalTo("research_agent"))) + .withRequestBody(matchingJsonPath("$.ir.start_node", equalTo("__model_0__"))) + .withRequestBody(matchingJsonPath("$.ir.nodes.__tools_0__.kind.type", equalTo("java_fn"))) + .withRequestBody(matchingJsonPath("$.ir.nodes.__tools_0__.kind.class_name", + equalTo("dev.jamjet.agent.tools.ToolDispatcher")))); + // START wire shape: workflow_version = the content hash; the seeded user prompt + + // tool-resolver map are the execution input (build_initial_state). + wm.verify(postRequestedFor(urlEqualTo("/executions")) + .withRequestBody(matchingJsonPath("$.workflow_id", equalTo("research_agent"))) + .withRequestBody(matchingJsonPath("$.workflow_version", equalTo(expectedVersion))) + .withRequestBody(matchingJsonPath("$.input.messages[0].role", equalTo("system"))) + .withRequestBody(matchingJsonPath("$.input.messages[1].role", equalTo("user"))) + .withRequestBody(matchingJsonPath("$.input.messages[1].content", equalTo("add 2 and 3"))) + .withRequestBody(matchingJsonPath("$.input.tools.add_numbers"))); + } + + @Test + @Timeout(20) + void terminalFailedSurfacesAsAgentRunException() { + stubCreateAndStart(); + wm.stubFor(get(urlEqualTo("/executions/ex_99")) + .willReturn(okJson("{\"execution_id\":\"ex_99\",\"status\":\"failed\"," + + "\"current_state\":{\"error\":\"model adapter exploded\"}}"))); + + try (var client = new JamjetEngineClient(wm.baseUrl())) { + assertThatThrownBy(() -> demoAgent().runDurable("hi", client, fastOptions())) + .isInstanceOf(AgentRunException.class) + .hasMessageContaining("non-completed") + .hasMessageContaining("model adapter exploded") + .extracting(e -> ((AgentRunException) e).status()) + .isEqualTo("failed"); + } + } + + @Test + @Timeout(20) + void terminalLimitExceededSurfacesAsAgentRunException() { + // A budget breach is a terminal limit_exceeded — it must surface loudly, not as a result. + stubCreateAndStart(); + wm.stubFor(get(urlEqualTo("/executions/ex_99")) + .willReturn(okJson("{\"execution_id\":\"ex_99\",\"status\":\"limit_exceeded\"}"))); + + try (var client = new JamjetEngineClient(wm.baseUrl())) { + assertThatThrownBy(() -> demoAgent().runDurable("hi", client, fastOptions())) + .isInstanceOf(AgentRunException.class) + .extracting(e -> ((AgentRunException) e).status()) + .isEqualTo("limit_exceeded"); + } + } + + @Test + @Timeout(20) + void neverTerminalRaisesTimeout() { + stubCreateAndStart(); + wm.stubFor(get(urlEqualTo("/executions/ex_99")) + .willReturn(okJson("{\"execution_id\":\"ex_99\",\"status\":\"running\"}"))); + + RunOptions tinyTimeout = fastOptions().withTimeout(Duration.ofMillis(120)); + try (var client = new JamjetEngineClient(wm.baseUrl())) { + assertThatThrownBy(() -> demoAgent().runDurable("hi", client, tinyTimeout)) + .isInstanceOf(AgentRunTimeoutException.class) + .extracting(e -> ((AgentRunException) e).status()) + .isEqualTo("running"); + } + } + + @Test + @Timeout(20) + void twoArgOverloadBuildsAndClosesItsOwnClient() { + // The common public path: runDurable(prompt, options) builds + closes its own + // JamjetEngineClient for options.runtimeUrl(). Point it at WireMock end-to-end. + stubCreateAndStart(); + wm.stubFor(get(urlEqualTo("/executions/ex_99")) + .willReturn(okJson("{\"execution_id\":\"ex_99\",\"status\":\"completed\"," + + "\"current_state\":{\"last_model_output\":\"done\"}}"))); + + AgentResult result = demoAgent().runDurable("hi", fastOptions()); + + assertThat(result.output()).isEqualTo("done"); + wm.verify(postRequestedFor(urlEqualTo("/workflows"))); + wm.verify(postRequestedFor(urlEqualTo("/executions"))); + } + + @Test + @Timeout(20) + void extractsFromLastAssistantMessageWhenNoLastModelOutput() { + stubCreateAndStart(); + // No last_model_output; the answer must fall back to the last assistant message. + wm.stubFor(get(urlEqualTo("/executions/ex_99")) + .willReturn(okJson("{\"execution_id\":\"ex_99\",\"status\":\"completed\"," + + "\"current_state\":{\"messages\":[" + + "{\"role\":\"user\",\"content\":\"hi\"}," + + "{\"role\":\"assistant\",\"content\":\"fallback answer\"}]}}"))); + + try (var client = new JamjetEngineClient(wm.baseUrl())) { + AgentResult result = demoAgent().runDurable("hi", client, fastOptions()); + assertThat(result.output()).isEqualTo("fallback answer"); + assertThat(result.toolCalls()).isEmpty(); + } + } + + // -- helpers ---------------------------------------------------------------- + + private void stubCreateAndStart() { + wm.stubFor(post(urlEqualTo("/workflows")) + .willReturn(okJson("{\"workflow_id\":\"research_agent\",\"version\":\"v1\"}"))); + wm.stubFor(post(urlEqualTo("/executions")) + .willReturn(okJson("{\"execution_id\":\"ex_99\",\"status\":\"running\"}"))); + } + + /** A terminal completed snapshot: an inline last_model_output + a tool-interleaved thread. */ + private static String completedBody() { + return "{\"execution_id\":\"ex_99\",\"status\":\"completed\",\"current_state\":{" + + "\"last_model_output\":\"The sum of 2 and 3 is 5.\"," + + "\"messages\":[" + + "{\"role\":\"system\",\"content\":\"You are a helpful research assistant.\"}," + + "{\"role\":\"user\",\"content\":\"add 2 and 3\"}," + + "{\"role\":\"assistant\",\"content\":\"\",\"tool_calls\":[{\"id\":\"tc1\",\"type\":\"function\"," + + "\"function\":{\"name\":\"add_numbers\",\"arguments\":\"{\\\"a\\\":2,\\\"b\\\":3}\"}}]}," + + "{\"role\":\"tool\",\"tool_call_id\":\"tc1\",\"name\":\"add_numbers\",\"content\":\"5\"}" + + "]}}"; + } +} From 4bed136b62e2dc691c8118dc360771e5944cda83 Mon Sep 17 00:00:00 2001 From: sdev Date: Mon, 29 Jun 2026 00:58:48 +0530 Subject: [PATCH 5/7] test(agent): worked governed durable agent example + e2e (B-5) An idiomatic Java example under jamjet-agent: a calculator Agent with a @Tool plus governance defaults (token+cost budget, inline policy with a blocked tool + model allowlist, an approval gate, PII), authored in clean Java and runnable durable+governed. - CalculatorTools: one @Tool method (calculate). - CalculatorAgentExample: buildAgent() is the whole authoring surface; main runs it durably with a background JavaToolWorker, documenting the run requirement (engine + model sidecar + a java_tool worker). Real shipped API only, no fictional methods. - CalculatorAgentExampleTest: hermetic end-to-end against a WireMock stub of the bare engine routes. Builds the agent, asserts governance rode into the IR, dispatches the @Tool durably through the worker (the tool actually ran: calculate(5,3,add) -> 8 in the completion body), and runDurable extracts the final assistant text + tool call. The result shape parallels the Python run_durable result. Also asserts the worker accepts EXACTLY the java_fn coordinate the compiler emits (the claim payload uses the IR's emitted coordinate), proving the DRY single-source dispatch coordinate end-to-end. - README: idiomatic authoring + the durable-run requirement. --- .../examples/CalculatorAgentExample.java | 86 ++++++++ .../examples/CalculatorAgentExampleTest.java | 188 ++++++++++++++++++ .../agent/examples/CalculatorTools.java | 37 ++++ .../java/dev/jamjet/agent/examples/README.md | 75 +++++++ 4 files changed, 386 insertions(+) create mode 100644 jamjet-agent/src/test/java/dev/jamjet/agent/examples/CalculatorAgentExample.java create mode 100644 jamjet-agent/src/test/java/dev/jamjet/agent/examples/CalculatorAgentExampleTest.java create mode 100644 jamjet-agent/src/test/java/dev/jamjet/agent/examples/CalculatorTools.java create mode 100644 jamjet-agent/src/test/java/dev/jamjet/agent/examples/README.md diff --git a/jamjet-agent/src/test/java/dev/jamjet/agent/examples/CalculatorAgentExample.java b/jamjet-agent/src/test/java/dev/jamjet/agent/examples/CalculatorAgentExample.java new file mode 100644 index 0000000..8086c0b --- /dev/null +++ b/jamjet-agent/src/test/java/dev/jamjet/agent/examples/CalculatorAgentExample.java @@ -0,0 +1,86 @@ +package dev.jamjet.agent.examples; + +import dev.jamjet.agent.Agent; +import dev.jamjet.agent.AgentResult; +import dev.jamjet.agent.Budget; +import dev.jamjet.agent.RunOptions; +import dev.jamjet.agent.client.JamjetEngineClient; +import dev.jamjet.agent.worker.JavaToolWorker; +import dev.jamjet.runtime.core.ir.PolicySetIr; + +import java.util.List; + +/** + * End-to-end example: a governed, durable calculator agent authored in idiomatic Java. + * + *

{@link #buildAgent()} is the whole authoring surface — a {@link Agent} with a model, + * instructions, one {@code @Tool} ({@link CalculatorTools#calculate}), and governance + * defaults (a token+cost {@link Budget}, an inline {@link PolicySetIr} with a blocked + * tool + a model allowlist, and an approval gate). It compiles to the SAME agent-loop + * {@code WorkflowIr} the Python ADK emits, so it runs on the identical Rust engine; the + * governance knobs ride in the IR and the engine enforces them fail-closed. + * + *

Running it for real ({@link #main})

+ * A durable run needs three services live (mirrors the Python {@code agent.run_durable}): + *
    + *
  1. the JamJet engine ({@code jamjet-server}) at {@link RunOptions#runtimeUrl()} + * — it owns the {@code java_tool} queue;
    2. + *
  2. the model sidecar ({@code JAMJET_MODEL_SEAM_URL}) — the engine routes every + * governed model call through it;
    3. + *
  3. a {@link JavaToolWorker} draining the {@code java_tool} queue with this + * agent's tool registry, so {@code calculate} executes durably exactly-once.
    4. + *
+ * {@code main} starts the worker on a background thread and shares ONE + * {@link JamjetEngineClient} between the worker and the run. (The hermetic + * {@code CalculatorAgentExampleTest} drives this same agent against a WireMock stub of + * the engine, since a live engine + sidecar is not available in CI.) + * + *

This example uses only the real shipped API — no fictional methods. + */ +public final class CalculatorAgentExample { + + private CalculatorAgentExample() {} + + /** The complete authoring surface: a governed, tool-using durable agent in idiomatic Java. */ + public static Agent buildAgent() { + return Agent.builder("calculator_agent") + .model("anthropic/claude-sonnet-4-6") + .instructions("You are a precise calculator. Use the calculate tool for every arithmetic step.") + .tools(new CalculatorTools()) + // -- governance: compiled into the IR, enforced fail-closed by the engine -- + .budget(new Budget(50_000, 0.50)) // cap tokens AND cost + .policy(new PolicySetIr( + List.of("delete_*"), // blocked_tools + List.of(), // require_approval_for (from globs below) + List.of("anthropic/claude-sonnet-4-6"))) // model_allowlist + .approvalRequired(List.of("wire_*")) // HITL gate on any wire_* tool + .pii(true) // PII redaction metadata in the IR + .build(); + } + + /** + * Run the agent durably against a local engine, with a background + * {@link JavaToolWorker} executing the {@code @Tool} call. Requires the three + * services above to be running; intended as a copy-paste starting point, not a CI test. + */ + public static void main(String[] args) { + Agent agent = buildAgent(); + RunOptions options = RunOptions.defaults(); // local engine at http://127.0.0.1:7700 + + try (JamjetEngineClient client = new JamjetEngineClient(options.runtimeUrl()); + JavaToolWorker worker = new JavaToolWorker(client, "calculator-worker-1", agent.registry())) { + + // Drain the java_tool queue on a background thread for the duration of the run. + Thread workerThread = new Thread(worker::run, "calculator-java-tool-worker"); + workerThread.setDaemon(true); + workerThread.start(); + + AgentResult result = agent.runDurable("What is (5 + 3) * 2?", client, options); + + System.out.println("answer: " + result.output()); + System.out.println("toolCalls: " + result.toolCalls()); + + worker.stop(); + } + } +} diff --git a/jamjet-agent/src/test/java/dev/jamjet/agent/examples/CalculatorAgentExampleTest.java b/jamjet-agent/src/test/java/dev/jamjet/agent/examples/CalculatorAgentExampleTest.java new file mode 100644 index 0000000..d88eb56 --- /dev/null +++ b/jamjet-agent/src/test/java/dev/jamjet/agent/examples/CalculatorAgentExampleTest.java @@ -0,0 +1,188 @@ +package dev.jamjet.agent.examples; + +import com.github.tomakehurst.wiremock.WireMockServer; +import dev.jamjet.agent.Agent; +import dev.jamjet.agent.AgentResult; +import dev.jamjet.agent.RunOptions; +import dev.jamjet.agent.client.JamjetEngineClient; +import dev.jamjet.agent.tools.ToolDispatcher; +import dev.jamjet.agent.worker.JavaToolWorker; +import dev.jamjet.runtime.core.JamjetJson; +import dev.jamjet.runtime.core.ir.NodeKind; +import dev.jamjet.runtime.core.ir.WorkflowIr; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +import java.time.Duration; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +import static com.github.tomakehurst.wiremock.client.WireMock.equalTo; +import static com.github.tomakehurst.wiremock.client.WireMock.get; +import static com.github.tomakehurst.wiremock.client.WireMock.matchingJsonPath; +import static com.github.tomakehurst.wiremock.client.WireMock.ok; +import static com.github.tomakehurst.wiremock.client.WireMock.okJson; +import static com.github.tomakehurst.wiremock.client.WireMock.post; +import static com.github.tomakehurst.wiremock.client.WireMock.postRequestedFor; +import static com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo; +import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.options; +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Hermetic END-TO-END proof for the {@link CalculatorAgentExample}: the idiomatic Java + * agent builds, compiles to the agent-loop IR, has its {@code @Tool} executed durably by + * a {@link JavaToolWorker}, and {@code runDurable} returns the expected governed result — + * all against a WireMock stub of the engine's bare routes (a live Rust engine + sidecar + * is not available in CI). + * + *

Parity with the Python equivalent. The asserted result shape matches the + * Python {@code agent.run_durable} result: the {@link AgentResult#output()} is the final + * assistant text, and the tool actually ran (proven twice — the worker's completion body + * carries the tool's output, and the reconstructed {@link AgentResult#toolCalls()} carries + * the call + its result). + * + *

Because WireMock cannot run the engine's scheduler, the two halves of the durable + * loop are stubbed independently and asserted together: (1) the {@link JavaToolWorker} + * draining a stubbed {@code java_tool} claim dispatches the real {@code @Tool}; (2) + * {@code runDurable} extracts the final answer from the terminal {@code current_state}. + */ +class CalculatorAgentExampleTest { + + WireMockServer wm; + + @BeforeEach + void up() { + wm = new WireMockServer(options().dynamicPort()); + wm.start(); + } + + @AfterEach + void down() { + wm.stop(); + } + + @Test + @Timeout(20) + void buildsCompilesDispatchesToolDurablyAndReturnsGovernedResult() { + Agent agent = CalculatorAgentExample.buildAgent(); + + // -- (a) the example compiles to a governed agent-loop IR ------------------- + WorkflowIr ir = agent.compileToIr(2); + assertThat(ir.startNode()).isEqualTo("__model_0__"); // M2: start is a Model node + assertThat(ir.costBudgetUsd()).isEqualTo(0.50); // governance rode into the IR + assertThat(ir.tokenBudget().totalTokens()).isEqualTo(50_000); + assertThat(ir.policy().blockedTools()).contains("delete_*"); + assertThat(ir.policy().requireApprovalFor()).contains("wire_*"); + assertThat(ir.dataPolicy().piiDetectors()).hasSize(5); + + // -- (b) DRY: the worker accepts EXACTLY the coordinate the compiler emits --- + NodeKind.JavaFn toolNode = (NodeKind.JavaFn) ir.node("__tools_0__").kind(); + // The emitted JavaFn coordinate is the same literal the worker's RCE gate checks. + assertThat(toolNode.className()).isEqualTo(ToolDispatcher.DISPATCH_CLASS); + assertThat(toolNode.method()).isEqualTo(ToolDispatcher.DISPATCH_METHOD); + + // -- stubs: the bare engine routes for BOTH the worker and the run ---------- + Map dispatchInput = Map.of( + "last_model_tool_calls", + List.of(toolCall("tc1", "calculate", Map.of("a", 5, "b", 3, "op", "add")))); + // The claim carries the coordinate the COMPILER emitted (not a hardcoded literal), + // so a green dispatch proves the worker accepts the builder's emitted coordinate. + wm.stubFor(post(urlEqualTo("/work-items/claim")) + .willReturn(okJson(claim(toolNode.className(), toolNode.method(), dispatchInput, 7L)))); + wm.stubFor(post(urlEqualTo("/work-items/wi_1/complete")).willReturn(ok())); + wm.stubFor(post(urlEqualTo("/work-items/wi_1/heartbeat")).willReturn(ok())); + + wm.stubFor(post(urlEqualTo("/workflows")) + .willReturn(okJson("{\"workflow_id\":\"calculator_agent\",\"version\":\"" + ir.version() + "\"}"))); + wm.stubFor(post(urlEqualTo("/executions")) + .willReturn(okJson("{\"execution_id\":\"ex_1\",\"status\":\"running\"}"))); + wm.stubFor(get(urlEqualTo("/executions/ex_1")) + .willReturn(okJson(completedBody()))); + + RunOptions runOptions = RunOptions.defaults() + .withMaxTurns(2) + .withRuntimeUrl(wm.baseUrl()) + .withPollInterval(Duration.ofMillis(15)) + .withTimeout(Duration.ofSeconds(5)); + + AgentResult result; + try (JamjetEngineClient client = new JamjetEngineClient(wm.baseUrl()); + JavaToolWorker worker = new JavaToolWorker(client, "calculator-worker-1", agent.registry(), + Duration.ofSeconds(2), Duration.ofMillis(10))) { + + // -- (c) the @Tool runs DURABLY: the worker dispatches the claimed java_tool item. + JavaToolWorker.ItemResult dispatched = worker.runOnce(); + assertThat(dispatched).isEqualTo(JavaToolWorker.ItemResult.COMPLETED); + + // -- (d) the durable run extracts the final governed answer ------------- + result = agent.runDurable("What is (5 + 3) * 2?", client, runOptions); + } + + // The tool actually ran: calculate(5, 3, "add") -> "8" in the completion body. + wm.verify(postRequestedFor(urlEqualTo("/work-items/wi_1/complete")) + .withRequestBody(matchingJsonPath("$.lease_fence", equalTo("7"))) + .withRequestBody(matchingJsonPath("$.output.messages[1].role", equalTo("tool"))) + .withRequestBody(matchingJsonPath("$.output.messages[1].content", equalTo("8")))); + + // The result PARALLELS the Python run_durable result: final assistant text + the tool call. + assertThat(result.output()).isEqualTo("(5 + 3) * 2 = 16"); + assertThat(result.terminalState().status()).isEqualTo("completed"); + assertThat(result.toolCalls()).hasSize(1); + AgentResult.ToolCall call = result.toolCalls().get(0); + assertThat(call.tool()).isEqualTo("calculate"); + assertThat(call.input()).containsEntry("op", "add"); + assertThat(call.output()).isEqualTo("8"); + } + + // -- helpers ---------------------------------------------------------------- + + private static Map toolCall(String id, String name, Map args) { + Map c = new LinkedHashMap<>(); + c.put("id", id); + c.put("name", name); + c.put("arguments", args); + return c; + } + + /** A {@code {claimed, work_item}} claim response carrying the given dispatch coordinate + input. */ + private static String claim(String dispatchClass, String dispatchMethod, Map input, long fence) { + Map payload = new LinkedHashMap<>(); + payload.put("class", dispatchClass); + payload.put("method", dispatchMethod); + payload.put("input", input); + + Map wi = new LinkedHashMap<>(); + wi.put("id", "wi_1"); + wi.put("execution_id", "ex_1"); + wi.put("node_id", "__tools_0__"); + wi.put("queue_type", "java_tool"); + wi.put("payload", payload); + wi.put("attempt", 1); + wi.put("lease_fence", fence); + + Map resp = new LinkedHashMap<>(); + resp.put("claimed", true); + resp.put("work_item", wi); + try { + return JamjetJson.shared().writeValueAsString(resp); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + /** A terminal completed snapshot: inline last_model_output + a calculate tool turn. */ + private static String completedBody() { + return "{\"execution_id\":\"ex_1\",\"status\":\"completed\",\"current_state\":{" + + "\"last_model_output\":\"(5 + 3) * 2 = 16\"," + + "\"messages\":[" + + "{\"role\":\"system\",\"content\":\"You are a precise calculator.\"}," + + "{\"role\":\"user\",\"content\":\"What is (5 + 3) * 2?\"}," + + "{\"role\":\"assistant\",\"content\":\"\",\"tool_calls\":[{\"id\":\"tc1\",\"type\":\"function\"," + + "\"function\":{\"name\":\"calculate\",\"arguments\":\"{\\\"a\\\":5,\\\"b\\\":3,\\\"op\\\":\\\"add\\\"}\"}}]}," + + "{\"role\":\"tool\",\"tool_call_id\":\"tc1\",\"name\":\"calculate\",\"content\":\"8\"}" + + "]}}"; + } +} diff --git a/jamjet-agent/src/test/java/dev/jamjet/agent/examples/CalculatorTools.java b/jamjet-agent/src/test/java/dev/jamjet/agent/examples/CalculatorTools.java new file mode 100644 index 0000000..7e19597 --- /dev/null +++ b/jamjet-agent/src/test/java/dev/jamjet/agent/examples/CalculatorTools.java @@ -0,0 +1,37 @@ +package dev.jamjet.agent.examples; + +import dev.jamjet.agent.Tool; + +/** + * A tiny tool-holder for the {@link CalculatorAgentExample}: one {@code @Tool} method + * the model can call to do exact arithmetic. Idiomatic Java — a plain class with a + * {@code @Tool}-annotated method; the {@code Agent} builder derives the OpenAI function + * schema from the method's name + typed parameters, and the durable + * {@link dev.jamjet.agent.worker.JavaToolWorker} invokes it reflectively when the model + * calls it. + */ +public final class CalculatorTools { + + /** + * Evaluate a binary arithmetic operation. The {@code op} is one of + * {@code add} / {@code subtract} / {@code multiply} / {@code divide}. + * + *

The JSON arguments the model sends ({@code {"a": 5, "b": 3, "op": "add"}}) are + * coerced to these typed parameters by name (Jackson), so the method stays plain Java. + */ + @Tool(name = "calculate", description = "Evaluate a binary arithmetic operation: add, subtract, multiply, or divide.") + public String calculate(double a, double b, String op) { + double result = switch (op) { + case "add" -> a + b; + case "subtract" -> a - b; + case "multiply" -> a * b; + case "divide" -> a / b; + default -> throw new IllegalArgumentException("unknown op: " + op); + }; + // Render whole-number results without a trailing ".0" for clean tool output. + if (result == Math.rint(result) && !Double.isInfinite(result)) { + return String.valueOf((long) result); + } + return String.valueOf(result); + } +} diff --git a/jamjet-agent/src/test/java/dev/jamjet/agent/examples/README.md b/jamjet-agent/src/test/java/dev/jamjet/agent/examples/README.md new file mode 100644 index 0000000..3a38650 --- /dev/null +++ b/jamjet-agent/src/test/java/dev/jamjet/agent/examples/README.md @@ -0,0 +1,75 @@ +# Example: a governed, durable calculator agent (Java) + +A worked end-to-end example of authoring a JamJet agent in idiomatic Java with +`dev.jamjet.agent.Agent`, running it durably and governed on the Rust engine. Every +call here is the real shipped `jamjet-agent` API. + +- `CalculatorTools` — a plain class with one `@Tool` method (`calculate`). +- `CalculatorAgentExample` — `buildAgent()` is the whole authoring surface; `main` + runs it durably against a local engine. +- `CalculatorAgentExampleTest` — drives the same agent end-to-end against a WireMock + stub of the engine (a live engine + sidecar is not available in CI). + +## Authoring + +```java +Agent agent = Agent.builder("calculator_agent") + .model("anthropic/claude-sonnet-4-6") + .instructions("You are a precise calculator. Use the calculate tool for every arithmetic step.") + .tools(new CalculatorTools()) + // governance: compiled into the IR, enforced fail-closed by the engine + .budget(new Budget(50_000, 0.50)) // token + cost cap + .policy(new PolicySetIr( + List.of("delete_*"), // blocked_tools + List.of(), // require_approval_for + List.of("anthropic/claude-sonnet-4-6"))) // model_allowlist + .approvalRequired(List.of("wire_*")) // HITL gate on wire_* tools + .pii(true) // PII redaction metadata + .build(); +``` + +`agent.compileToIr()` produces the same agent-loop `WorkflowIr` the Python ADK emits +(the Java tool nodes are `java_fn` where Python's are `python_fn`), so a Java agent and +a Python agent run on the identical engine. The governance knobs ride in the IR; the +engine enforces them, not the Java SDK. + +## Running it durably + +```java +RunOptions options = RunOptions.defaults(); // local engine at http://127.0.0.1:7700 + +try (JamjetEngineClient client = new JamjetEngineClient(options.runtimeUrl()); + JavaToolWorker worker = new JavaToolWorker(client, "calculator-worker-1", agent.registry())) { + + Thread workerThread = new Thread(worker::run, "calculator-java-tool-worker"); + workerThread.setDaemon(true); + workerThread.start(); + + AgentResult result = agent.runDurable("What is (5 + 3) * 2?", client, options); + System.out.println(result.output()); // the final assistant text + System.out.println(result.toolCalls()); // the calculate call + its result + + worker.stop(); +} +``` + +`runDurable` compiles the IR, registers it (`POST /workflows`), starts an execution +seeded with the system+user messages (`POST /executions`), polls +`GET /executions/{id}` to a terminal state, and extracts the answer from +`current_state.last_model_output` (falling back to the last assistant message). A +non-`completed` terminal (`failed` / `cancelled` / `limit_exceeded`, e.g. a budget +breach) raises `AgentRunException`; a run that never terminates raises +`AgentRunTimeoutException`. + +## Required running services + +A durable run is not self-contained. Three services must be live (the same contract as +the Python `agent.run_durable`): + +1. **the JamJet engine** (`jamjet-server`) at `RunOptions.runtimeUrl()`, which owns the + `java_tool` queue; +2. **the model sidecar** (`JAMJET_MODEL_SEAM_URL`), through which the engine routes + every governed model call (there is no Java model code); +3. **a `JavaToolWorker`** draining the `java_tool` queue with this agent's tool + registry, so the `@Tool` methods execute durably exactly-once. Run it in a separate + thread or process; `runDurable` does not start one. From 1e6b3764d4ea58c4d29b39b4a80dc11eb4a88716 Mon Sep 17 00:00:00 2001 From: sdev Date: Mon, 29 Jun 2026 01:13:07 +0530 Subject: [PATCH 6/7] docs(agent): honest PII scope on runDurable + snake_case terminal status --- jamjet-agent/src/main/java/dev/jamjet/agent/Agent.java | 6 ++++-- .../main/java/dev/jamjet/agent/client/ExecutionState.java | 7 ++++--- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/Agent.java b/jamjet-agent/src/main/java/dev/jamjet/agent/Agent.java index 7393d8d..e3f8b89 100644 --- a/jamjet-agent/src/main/java/dev/jamjet/agent/Agent.java +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/Agent.java @@ -102,8 +102,10 @@ public WorkflowIr compileToIr(int maxTurns) { * to a terminal state, and extracts the answer from the terminal * {@code current_state.last_model_output} (falling back to the last assistant * message). Every model call and tool dispatch runs through the durable engine, so - * the run is event-sourced, replayable, idempotent, and governed (budget / policy / - * PII enforced fail-closed from the compiled IR). + * the run is event-sourced, replayable, idempotent, and governed: budget and policy + * (the model allowlist plus approval gates) are enforced fail-closed by the engine, + * and PII redaction is applied at the model-seam sidecar (the {@code data_policy} IR + * signals it; redaction is the sidecar's job, not an IR-level guarantee). * *

Required running services (mirrors the Python {@code run_durable})

* A durable run is NOT self-contained — three services must be running: diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/client/ExecutionState.java b/jamjet-agent/src/main/java/dev/jamjet/agent/client/ExecutionState.java index 7d2dfab..7c27b3b 100644 --- a/jamjet-agent/src/main/java/dev/jamjet/agent/client/ExecutionState.java +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/client/ExecutionState.java @@ -5,9 +5,10 @@ /** * Response of {@code GET /executions/{id}} — the polled execution snapshot. * Mirrors the Python client's {@code get_execution} shape ({@code {status, - * current_state, ...}}). The terminal status is {@code "Completed"} (or - * {@code "Failed"} / {@code "Cancelled"}); {@code currentState} carries the - * accumulated execution state. + * current_state, ...}}). The engine serializes status in snake_case, so a + * terminal status is {@code "completed"} (or {@code "failed"} / + * {@code "cancelled"} / {@code "limit_exceeded"}); {@code currentState} carries + * the accumulated execution state. */ public record ExecutionState( String executionId, From 74be915807c5a05943e2f4337d26db18ca105448 Mon Sep 17 00:00:00 2001 From: sdev Date: Mon, 29 Jun 2026 01:49:00 +0530 Subject: [PATCH 7/7] fix(agent): address CodeRabbit review on Track 9 Phase B Correctness (with tests): - ToolRegistry: reject overloaded @Tool methods that collide on the class+method dispatch key, which would silently corrupt byClassAndMethod - JavaToolWorker: an interrupt mid-dispatch is a reclaimable no-op, not a tool failure; restore the interrupt and stop without failing the item - ToolDispatcher: a tool call missing a required arg yields a clean tool error surfaced to the model, not a null-coerced invocation Input validation: - Agent.Builder: reject non-positive timeoutSeconds and a blank model - Budget: reject an entirely-empty budget and a non-finite costUsd - RunOptions: reject a zero or negative timeout Hygiene: - NodeKind.Model: deep-copy and freeze tool schemas, not just the outer list - JamjetEngineClient.heartbeatWorkItem: accept a nullable Long lease fence and include it in the body only when present - ToolInvocationException: keep the surfaced message generic and retain the cause as the exception cause, never embedding raw cause.toString - CalculatorTools example: BigDecimal arithmetic with a clean divide-by-zero error and whole-number rendering without a trailing .0 Adds regression coverage for older Model JSON with no tools field. Full jamjet-agent and jamjet-runtime-core suites green (Java 21, 163 tests). --- .../src/main/java/dev/jamjet/agent/Agent.java | 8 ++- .../main/java/dev/jamjet/agent/Budget.java | 17 +++++- .../java/dev/jamjet/agent/RunOptions.java | 5 ++ .../agent/client/JamjetEngineClient.java | 12 +++- .../jamjet/agent/tools/ToolDispatcher.java | 27 ++++++++- .../agent/tools/ToolInvocationException.java | 7 ++- .../dev/jamjet/agent/tools/ToolRegistry.java | 16 +++++- .../jamjet/agent/worker/JavaToolWorker.java | 34 +++++++---- .../agent/AgentBuilderValidationTest.java | 46 +++++++++++++++ .../java/dev/jamjet/agent/BudgetTest.java | 47 +++++++++++++++ .../agent/RunOptionsValidationTest.java | 42 ++++++++++++++ .../agent/client/JamjetEngineClientTest.java | 16 ++++++ .../agent/examples/CalculatorTools.java | 40 +++++++++---- .../agent/examples/CalculatorToolsTest.java | 43 ++++++++++++++ .../agent/tools/ToolDispatcherTest.java | 50 ++++++++++++++++ .../jamjet/agent/tools/ToolRegistryTest.java | 26 +++++++++ .../agent/worker/JavaToolWorkerTest.java | 48 ++++++++++++++++ .../dev/jamjet/runtime/core/ir/NodeKind.java | 47 ++++++++++++++- .../core/ir/NodeKindSerializationTest.java | 57 +++++++++++++++++++ 19 files changed, 556 insertions(+), 32 deletions(-) create mode 100644 jamjet-agent/src/test/java/dev/jamjet/agent/AgentBuilderValidationTest.java create mode 100644 jamjet-agent/src/test/java/dev/jamjet/agent/BudgetTest.java create mode 100644 jamjet-agent/src/test/java/dev/jamjet/agent/RunOptionsValidationTest.java create mode 100644 jamjet-agent/src/test/java/dev/jamjet/agent/examples/CalculatorToolsTest.java diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/Agent.java b/jamjet-agent/src/main/java/dev/jamjet/agent/Agent.java index e3f8b89..c585344 100644 --- a/jamjet-agent/src/main/java/dev/jamjet/agent/Agent.java +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/Agent.java @@ -308,8 +308,11 @@ public Builder pii(boolean pii) { return this; } - /** The workflow timeout in seconds (default 300). */ + /** The workflow timeout in seconds (default 300; must be positive). */ public Builder timeoutSeconds(int timeoutSeconds) { + if (timeoutSeconds <= 0) { + throw new IllegalArgumentException("timeoutSeconds must be positive (got " + timeoutSeconds + ")"); + } this.timeoutSeconds = timeoutSeconds; return this; } @@ -317,6 +320,9 @@ public Builder timeoutSeconds(int timeoutSeconds) { /** Build the immutable {@link Agent}. */ public Agent build() { Objects.requireNonNull(model, "agent model must be set"); + if (model.isBlank()) { + throw new IllegalArgumentException("agent model must not be blank"); + } return new Agent(this); } } diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/Budget.java b/jamjet-agent/src/main/java/dev/jamjet/agent/Budget.java index 0d61684..6d721e4 100644 --- a/jamjet-agent/src/main/java/dev/jamjet/agent/Budget.java +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/Budget.java @@ -2,7 +2,8 @@ /** * A per-run spending cap for a governed agent — the Java analog of the Python - * {@code jamjet.agents.governance.Budget}. Either or both fields may be set. + * {@code jamjet.agents.governance.Budget}. At least one field must be set (an entirely + * empty budget is rejected); set either or both. * *

At compile time {@code costUsd} maps to the IR's {@code cost_budget_usd} and * {@code tokens} maps to {@code token_budget.total_tokens} (a combined input+output @@ -14,11 +15,21 @@ public record Budget(Integer tokens, Double costUsd) { public Budget { + if (tokens == null && costUsd == null) { + throw new IllegalArgumentException("Budget must cap tokens and/or cost (both were null)"); + } if (tokens != null && tokens <= 0) { throw new IllegalArgumentException("Budget.tokens must be positive"); } - if (costUsd != null && costUsd <= 0) { - throw new IllegalArgumentException("Budget.costUsd must be positive"); + if (costUsd != null) { + // Guard non-finite first: NaN/Infinity slip past the `<= 0` check (every NaN + // comparison is false), so they would otherwise be accepted as a "cap". + if (!Double.isFinite(costUsd)) { + throw new IllegalArgumentException("Budget.costUsd must be a finite number"); + } + if (costUsd <= 0) { + throw new IllegalArgumentException("Budget.costUsd must be positive"); + } } } diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/RunOptions.java b/jamjet-agent/src/main/java/dev/jamjet/agent/RunOptions.java index 01d28d7..69b4546 100644 --- a/jamjet-agent/src/main/java/dev/jamjet/agent/RunOptions.java +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/RunOptions.java @@ -46,6 +46,11 @@ public record RunOptions( if (pollInterval == null || pollInterval.isNegative() || pollInterval.isZero()) { throw new IllegalArgumentException("pollInterval must be positive"); } + // timeout is nullable (null = derive from the agent's timeoutSeconds); when set it + // must be a positive duration, never zero/negative. + if (timeout != null && (timeout.isNegative() || timeout.isZero())) { + throw new IllegalArgumentException("timeout must be positive when set"); + } } /** The default options: {@code maxTurns=8}, local runtime, no auth, 500ms poll, agent-derived timeout. */ 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 60b6480..95f0afe 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 @@ -199,13 +199,19 @@ public void failWorkItem(String itemId, String error) { /** * {@code POST /work-items/{id}/heartbeat} — renew the lease. - * Body {@code {worker_id, lease_fence}}; the fence must match the claim's value + * Body {@code {worker_id, lease_fence?}}; the fence must match the claim's value * or the engine rejects the renewal. + * + *

{@code leaseFence} is nullable — consistent with {@link #completeWorkItem} and + * {@link ClaimedWorkItem#leaseFence()} — and included in the body only when non-null. + * For a real fenced claim it is always present; null is only the legacy unfenced case. */ - public void heartbeatWorkItem(String itemId, String workerId, long leaseFence) { + public void heartbeatWorkItem(String itemId, String workerId, Long leaseFence) { Map body = new LinkedHashMap<>(); body.put("worker_id", workerId); - body.put("lease_fence", leaseFence); + if (leaseFence != null) { + body.put("lease_fence", leaseFence); + } execute(buildPost("/work-items/" + itemId + "/heartbeat", body)); } diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/tools/ToolDispatcher.java b/jamjet-agent/src/main/java/dev/jamjet/agent/tools/ToolDispatcher.java index 577bf1b..783a8a9 100644 --- a/jamjet-agent/src/main/java/dev/jamjet/agent/tools/ToolDispatcher.java +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/tools/ToolDispatcher.java @@ -134,6 +134,14 @@ private String invokeTool(String name, Map arguments) { } Method method = tool.method(); + // Required-arg gate: every @Tool parameter is required (ToolRegistry.buildInputSchema + // marks them all). A call that omits one must surface a clean role:tool error to the + // model, NOT invoke the tool with a null-coerced (or primitive-defaulted) argument. + List missing = missingRequiredArgs(method, arguments); + if (!missing.isEmpty()) { + return "ERROR: tool '" + name + "' missing required argument(s): " + String.join(", ", missing); + } + Object[] args = coerceToParameters(method, arguments); try { method.setAccessible(true); @@ -147,10 +155,27 @@ private String invokeTool(String name, Map arguments) { } } + /** + * The names of the tool's required parameters that are absent from {@code arguments}. + * Every {@code @Tool} parameter is required (parity with the emitted schema), so any + * parameter whose name is not a key is missing. A present-but-{@code null} value is the + * model's explicit choice and is left to the tool. + */ + private static List missingRequiredArgs(Method method, Map arguments) { + List missing = new ArrayList<>(); + for (Parameter p : method.getParameters()) { + if (!arguments.containsKey(p.getName())) { + missing.add(p.getName()); + } + } + return missing; + } + /** * Coerce the JSON argument map to the method's typed positional parameters via * Jackson, matching each parameter by its (compiled {@code -parameters}) name. - * A missing argument coerces to {@code null} (or the primitive default). + * Callers must have already rejected missing required args (see + * {@link #missingRequiredArgs}); a present-but-null value coerces to {@code null}. */ private Object[] coerceToParameters(Method method, Map arguments) { Parameter[] params = method.getParameters(); diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/tools/ToolInvocationException.java b/jamjet-agent/src/main/java/dev/jamjet/agent/tools/ToolInvocationException.java index 52bff34..5694e74 100644 --- a/jamjet-agent/src/main/java/dev/jamjet/agent/tools/ToolInvocationException.java +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/tools/ToolInvocationException.java @@ -11,13 +11,18 @@ *

An unknown/unregistered tool name does NOT raise this — it is surfaced * to the model as a clean {@code role: tool} error message instead, so a hallucinated * name cannot fail the run (and, critically, never triggers an arbitrary invocation). + * + *

The surfaced message stays generic (the tool name only). The raw cause is retained + * as the exception {@code cause} (for stack traces) but deliberately NOT embedded in the + * message text, because that message is logged and persisted to the engine via + * {@code failWorkItem}, where a raw {@code cause.toString()} could leak sensitive detail. */ public final class ToolInvocationException extends RuntimeException { private final String toolName; public ToolInvocationException(String toolName, Throwable cause) { - super("tool '" + toolName + "' failed: " + (cause == null ? "unknown error" : cause.toString()), cause); + super("tool '" + toolName + "' failed", cause); this.toolName = toolName; } diff --git a/jamjet-agent/src/main/java/dev/jamjet/agent/tools/ToolRegistry.java b/jamjet-agent/src/main/java/dev/jamjet/agent/tools/ToolRegistry.java index be34123..9676de0 100644 --- a/jamjet-agent/src/main/java/dev/jamjet/agent/tools/ToolRegistry.java +++ b/jamjet-agent/src/main/java/dev/jamjet/agent/tools/ToolRegistry.java @@ -74,6 +74,9 @@ public static ToolRegistry of(List holders) { * method-name order. Returns {@code this} for chaining. * * @throws IllegalArgumentException if two registered tools share a name + * @throws IllegalStateException if two {@code @Tool} methods overload the same + * class+method dispatch key (which would corrupt + * {@link #byClassAndMethod} dispatch) */ public ToolRegistry register(Object holder) { if (holder == null) { @@ -101,13 +104,24 @@ public ToolRegistry register(Object holder) { RegisteredTool rt = new RegisteredTool( name, description, clazz.getName(), m.getName(), m, holder, inputSchema, openAi); + String key = rt.key(); if (byName.containsKey(name)) { throw new IllegalArgumentException( "Duplicate @Tool name '" + name + "' (from " + clazz.getName() + "#" + m.getName() + "); tool names must be unique within an agent."); } + // Overloaded @Tool methods (same declaring class + method name, but distinct + // @Tool names) collide on the class+method dispatch key. Silently letting the + // second overwrite the first would corrupt byClassAndMethod, so reject it. + if (byKey.containsKey(key)) { + throw new IllegalStateException( + "Overloaded @Tool method '" + clazz.getName() + "#" + m.getName() + + "': another @Tool already maps to the same class+method dispatch key '" + + key + "'. Tool dispatch keys on class+method, so overloads collide; " + + "give each @Tool a distinct method name."); + } tools.add(rt); - byKey.put(rt.key(), rt); + byKey.put(key, rt); byName.put(name, rt); } return this; 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 2af8e6f..fef8c06 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 @@ -80,7 +80,11 @@ public enum ItemResult { COMPLETED, /** A genuine tool / dispatch / completion error; the item was failed. */ FAILED, - /** The lease was lost (heartbeat abort or a 409 on complete); a no-op, NOT a failure. */ + /** + * The item was abandoned without failing it: the lease was lost (heartbeat abort or + * a 409 on complete), or the worker was interrupted mid-dispatch (shutdown). A no-op, + * NOT a failure — the engine reclaims the item on lease expiry. + */ LOST_LEASE } @@ -193,18 +197,26 @@ private ItemResult process(ClaimedWorkItem item) { Map output; try { output = dispatchFuture.get(); - } catch (CancellationException | InterruptedException | ExecutionException e) { - if (Thread.interrupted()) { - // restore the flag for the caller; we're exiting this item anyway - Thread.currentThread().interrupt(); - } - // Lease-lost gate FIRST: if the heartbeat aborted us, this is a no-op, - // never a failure (the new claimant owns the item now). + } catch (InterruptedException e) { + // Worker shutdown / cancellation (the WORKER thread was interrupted while + // waiting) is NOT a tool failure. Restore the interrupt and stop cleanly so + // the engine reclaims the item on lease expiry — never fail a reclaimable + // item (the same invariant as a lost lease). Best-effort: stop the in-flight + // tool so a side-effecting call doesn't keep running after we abandon it. + Thread.currentThread().interrupt(); + dispatchFuture.cancel(true); + LOG.log(Level.INFO, () -> "worker interrupted mid-dispatch; abandoning item " + + item.id() + " for reclaim (not failing)"); + return ItemResult.LOST_LEASE; + } catch (CancellationException | ExecutionException e) { + // Lease-lost gate FIRST: if the heartbeat aborted us (Future.cancel), this is + // a no-op, never a failure (the new claimant owns the item now). if (leaseLost.get()) { LOG.log(Level.INFO, () -> "lease lost mid-dispatch; aborting item " + item.id() + " (not completing)"); return ItemResult.LOST_LEASE; } - // A genuine tool / dispatch failure -> fail the item (mirror Python). + // A genuine tool / dispatch failure -> fail the item (mirror Python). An + // ExecutionException here wraps the real tool error. 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); @@ -269,7 +281,9 @@ private ScheduledFuture scheduleHeartbeat(ClaimedWorkItem item, long fence, return; } try { - client.heartbeatWorkItem(item.id(), workerId, fence); + // Echo the claim's fence (null only for a legacy unfenced claim), matching + // the nullable contract the complete path uses. + client.heartbeatWorkItem(item.id(), workerId, fence == 0L ? null : fence); } catch (JamjetHttpException e) { if (e.statusCode() >= 400) { // Definitive engine rejection: the lease is gone (FenceLost -> 500, diff --git a/jamjet-agent/src/test/java/dev/jamjet/agent/AgentBuilderValidationTest.java b/jamjet-agent/src/test/java/dev/jamjet/agent/AgentBuilderValidationTest.java new file mode 100644 index 0000000..7d88bb1 --- /dev/null +++ b/jamjet-agent/src/test/java/dev/jamjet/agent/AgentBuilderValidationTest.java @@ -0,0 +1,46 @@ +package dev.jamjet.agent; + +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Input-validation guards on the {@link Agent.Builder}: a non-positive timeout and a + * blank/empty model are rejected eagerly rather than producing a nonsensical IR. + */ +class AgentBuilderValidationTest { + + @Test + void rejectsNonPositiveTimeoutSeconds() { + assertThatThrownBy(() -> Agent.builder("a").timeoutSeconds(0)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("timeoutSeconds"); + assertThatThrownBy(() -> Agent.builder("a").timeoutSeconds(-1)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("timeoutSeconds"); + } + + @Test + void acceptsPositiveTimeoutSeconds() { + Agent agent = Agent.builder("a").model("anthropic/claude-sonnet-4-6") + .timeoutSeconds(120).build(); + assertThat(agent.timeoutSeconds()).isEqualTo(120); + } + + @Test + void rejectsBlankModel() { + assertThatThrownBy(() -> Agent.builder("a").model(" ").build()) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("model"); + assertThatThrownBy(() -> Agent.builder("a").model("").build()) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("model"); + } + + @Test + void rejectsNullModel() { + assertThatThrownBy(() -> Agent.builder("a").build()) + .isInstanceOf(NullPointerException.class); + } +} diff --git a/jamjet-agent/src/test/java/dev/jamjet/agent/BudgetTest.java b/jamjet-agent/src/test/java/dev/jamjet/agent/BudgetTest.java new file mode 100644 index 0000000..abf3eb2 --- /dev/null +++ b/jamjet-agent/src/test/java/dev/jamjet/agent/BudgetTest.java @@ -0,0 +1,47 @@ +package dev.jamjet.agent; + +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatCode; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Validation guards on {@link Budget}: besides the existing non-positive checks, an + * entirely-empty budget (no tokens AND no cost) and a non-finite {@code costUsd} + * (NaN / Infinity) are rejected — a NaN would otherwise slip past {@code costUsd <= 0}. + */ +class BudgetTest { + + @Test + void rejectsEntirelyEmptyBudget() { + assertThatThrownBy(() -> new Budget(null, null)) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + void rejectsNonFiniteCostUsd() { + assertThatThrownBy(() -> new Budget(null, Double.NaN)) + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> new Budget(null, Double.POSITIVE_INFINITY)) + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> new Budget(1_000, Double.NaN)) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + void stillRejectsNonPositiveValues() { + assertThatThrownBy(() -> new Budget(0, null)).isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> new Budget(null, 0.0)).isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> new Budget(-1, null)).isInstanceOf(IllegalArgumentException.class); + } + + @Test + void acceptsValidBudgets() { + assertThatCode(() -> new Budget(100_000, null)).doesNotThrowAnyException(); + assertThatCode(() -> new Budget(null, 2.5)).doesNotThrowAnyException(); + assertThatCode(() -> new Budget(100_000, 2.5)).doesNotThrowAnyException(); + assertThat(Budget.ofTokens(50_000).tokens()).isEqualTo(50_000); + assertThat(Budget.ofCostUsd(0.5).costUsd()).isEqualTo(0.5); + } +} diff --git a/jamjet-agent/src/test/java/dev/jamjet/agent/RunOptionsValidationTest.java b/jamjet-agent/src/test/java/dev/jamjet/agent/RunOptionsValidationTest.java new file mode 100644 index 0000000..d525fb0 --- /dev/null +++ b/jamjet-agent/src/test/java/dev/jamjet/agent/RunOptionsValidationTest.java @@ -0,0 +1,42 @@ +package dev.jamjet.agent; + +import org.junit.jupiter.api.Test; + +import java.time.Duration; + +import static org.assertj.core.api.Assertions.assertThatCode; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Validation guards on {@link RunOptions}: the compact constructor already rejects a + * non-positive {@code pollInterval}; it must equally reject a zero/negative + * {@code timeout} (a {@code null} timeout is still allowed — it derives from the agent). + */ +class RunOptionsValidationTest { + + @Test + void rejectsZeroTimeout() { + assertThatThrownBy(() -> RunOptions.defaults().withTimeout(Duration.ZERO)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("timeout"); + } + + @Test + void rejectsNegativeTimeout() { + assertThatThrownBy(() -> RunOptions.defaults().withTimeout(Duration.ofSeconds(-1))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("timeout"); + } + + @Test + void allowsNullTimeout() { + // null timeout = derive the deadline from the agent's timeoutSeconds. + assertThatCode(() -> RunOptions.defaults().withTimeout(null)).doesNotThrowAnyException(); + } + + @Test + void allowsPositiveTimeout() { + assertThatCode(() -> RunOptions.defaults().withTimeout(Duration.ofSeconds(5))) + .doesNotThrowAnyException(); + } +} 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 e4550c5..7124e88 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 @@ -187,6 +187,22 @@ void completeOmitsLeaseFenceWhenNull() { assertThat(body).contains("duration_ms"); } + @Test + void heartbeatOmitsLeaseFenceWhenNull() { + wm.stubFor(post(urlEqualTo("/work-items/wi_nf/heartbeat")).willReturn(ok())); + + var client = new JamjetEngineClient(wm.baseUrl()); + // A null fence (legacy unfenced claim) must be omitted from the body, mirroring the + // nullable contract of completeWorkItem. + client.heartbeatWorkItem("wi_nf", "w1", null); + + List reqs = wm.findAll(postRequestedFor(urlEqualTo("/work-items/wi_nf/heartbeat"))); + assertThat(reqs).hasSize(1); + String body = reqs.get(0).getBodyAsString(); + assertThat(body).doesNotContain("lease_fence"); + assertThat(body).contains("worker_id"); + } + @Test void staleFenceCompleteSurfacesAsConflict() { wm.stubFor(post(urlEqualTo("/work-items/wi_stale/complete")) diff --git a/jamjet-agent/src/test/java/dev/jamjet/agent/examples/CalculatorTools.java b/jamjet-agent/src/test/java/dev/jamjet/agent/examples/CalculatorTools.java index 7e19597..59f179c 100644 --- a/jamjet-agent/src/test/java/dev/jamjet/agent/examples/CalculatorTools.java +++ b/jamjet-agent/src/test/java/dev/jamjet/agent/examples/CalculatorTools.java @@ -2,6 +2,9 @@ import dev.jamjet.agent.Tool; +import java.math.BigDecimal; +import java.math.MathContext; + /** * A tiny tool-holder for the {@link CalculatorAgentExample}: one {@code @Tool} method * the model can call to do exact arithmetic. Idiomatic Java — a plain class with a @@ -21,17 +24,32 @@ public final class CalculatorTools { */ @Tool(name = "calculate", description = "Evaluate a binary arithmetic operation: add, subtract, multiply, or divide.") public String calculate(double a, double b, String op) { - double result = switch (op) { - case "add" -> a + b; - case "subtract" -> a - b; - case "multiply" -> a * b; - case "divide" -> a / b; - default -> throw new IllegalArgumentException("unknown op: " + op); - }; - // Render whole-number results without a trailing ".0" for clean tool output. - if (result == Math.rint(result) && !Double.isInfinite(result)) { - return String.valueOf((long) result); + // Exact decimal arithmetic via BigDecimal: avoids binary-float noise (0.1 + 0.2) + // and lets divide-by-zero be a clean error rather than producing "Infinity". + // valueOf(double) uses the canonical short decimal string, so "0.1" stays "0.1". + BigDecimal x = BigDecimal.valueOf(a); + BigDecimal y = BigDecimal.valueOf(b); + BigDecimal result; + switch (op) { + case "add" -> result = x.add(y); + case "subtract" -> result = x.subtract(y); + case "multiply" -> result = x.multiply(y); + case "divide" -> { + if (y.signum() == 0) { + return "ERROR: division by zero"; + } + result = x.divide(y, MathContext.DECIMAL64); + } + default -> { + return "ERROR: unknown op '" + op + "' (expected add, subtract, multiply, or divide)"; + } + } + // Render whole-number results without a trailing ".0"; otherwise strip trailing + // zeros and render in plain (non-scientific) decimal notation. + result = result.stripTrailingZeros(); + if (result.scale() <= 0) { + return result.toBigInteger().toString(); } - return String.valueOf(result); + return result.toPlainString(); } } diff --git a/jamjet-agent/src/test/java/dev/jamjet/agent/examples/CalculatorToolsTest.java b/jamjet-agent/src/test/java/dev/jamjet/agent/examples/CalculatorToolsTest.java new file mode 100644 index 0000000..44e66f1 --- /dev/null +++ b/jamjet-agent/src/test/java/dev/jamjet/agent/examples/CalculatorToolsTest.java @@ -0,0 +1,43 @@ +package dev.jamjet.agent.examples; + +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Unit tests for the example {@link CalculatorTools}: exact decimal arithmetic via + * BigDecimal, a clean error (not {@code Infinity}) on divide-by-zero, and whole-number + * results rendered without a trailing {@code .0}. + */ +class CalculatorToolsTest { + + private final CalculatorTools calc = new CalculatorTools(); + + @Test + void wholeNumberResultsHaveNoTrailingDecimal() { + assertThat(calc.calculate(5, 3, "add")).isEqualTo("8"); + assertThat(calc.calculate(5, 3, "subtract")).isEqualTo("2"); + assertThat(calc.calculate(5, 3, "multiply")).isEqualTo("15"); + assertThat(calc.calculate(10, 2, "divide")).isEqualTo("5"); + } + + @Test + void decimalDivisionIsExactNotFloatingPointNoise() { + assertThat(calc.calculate(10, 4, "divide")).isEqualTo("2.5"); + // 0.1 + 0.2 is 0.3 in decimal arithmetic (not 0.30000000000000004). + assertThat(calc.calculate(0.1, 0.2, "add")).isEqualTo("0.3"); + } + + @Test + void divisionByZeroReturnsCleanErrorNotInfinity() { + String out = calc.calculate(5, 0, "divide"); + assertThat(out).doesNotContain("Infinity"); + assertThat(out.toLowerCase()).contains("zero"); + } + + @Test + void unknownOperationReturnsCleanError() { + String out = calc.calculate(1, 2, "modulo"); + assertThat(out.toLowerCase()).contains("modulo"); + } +} diff --git a/jamjet-agent/src/test/java/dev/jamjet/agent/tools/ToolDispatcherTest.java b/jamjet-agent/src/test/java/dev/jamjet/agent/tools/ToolDispatcherTest.java index cc4911c..a5c9377 100644 --- a/jamjet-agent/src/test/java/dev/jamjet/agent/tools/ToolDispatcherTest.java +++ b/jamjet-agent/src/test/java/dev/jamjet/agent/tools/ToolDispatcherTest.java @@ -30,6 +30,17 @@ public String boom(String why) { } } + /** A tool that records whether it was invoked, to prove a null-coerced call never happens. */ + public static final class GreetTool { + volatile boolean invoked = false; + + @Tool(name = "greet", description = "greets a person") + public String greet(String name) { + invoked = true; + return "hello " + name; + } + } + @SuppressWarnings("unchecked") private static List> messagesOf(Map out) { assertThat(out).containsKey("messages"); @@ -124,6 +135,45 @@ void unknownToolYieldsCleanErrorMessageNotExceptionAndNoInvocation() { .contains("java.lang.Runtime"); } + @Test + void missingRequiredArgYieldsCleanToolErrorNotANullCall() { + // Every @Tool param is required (buildInputSchema marks them all). A call missing + // a required arg must surface a clean role:tool error to the model, NOT invoke the + // tool with a null-coerced argument. + GreetTool greet = new GreetTool(); + var dispatcher = new ToolDispatcher(ToolRegistry.of(greet)); + + Map state = Map.of( + "last_model_tool_calls", + List.of(toolCall("tc1", "greet", Map.of()))); // no "name" arg + + Map out = dispatcher.dispatchToolCalls(state); // must NOT throw + List> messages = messagesOf(out); + assertThat(messages.get(1)) + .containsEntry("role", "tool") + .containsEntry("tool_call_id", "tc1"); + assertThat((String) messages.get(1).get("content")) + .contains("missing required argument") + .contains("name"); + // CRITICAL: the tool was never invoked with a null arg. + assertThat(greet.invoked).isFalse(); + } + + @Test + void toolInvocationExceptionMessageDoesNotLeakCauseDetail() { + // The surfaced message is logged + persisted to the engine via failWorkItem, so it + // must stay generic (tool name only) — never embed the raw cause.toString(). + ToolInvocationException ex = new ToolInvocationException( + "wire_money", new IllegalStateException("secret-account-9988")); + assertThat(ex.getMessage()) + .contains("wire_money") + .doesNotContain("secret-account-9988") + .doesNotContain("IllegalStateException"); + // The cause is retained as the exception cause (just not in the message text). + assertThat(ex.getCause()).isInstanceOf(IllegalStateException.class); + assertThat(ex.toolName()).isEqualTo("wire_money"); + } + @Test void registeredToolThatThrowsPropagatesAsToolInvocationException() { var dispatcher = new ToolDispatcher(ToolRegistry.of(new BoomTool())); diff --git a/jamjet-agent/src/test/java/dev/jamjet/agent/tools/ToolRegistryTest.java b/jamjet-agent/src/test/java/dev/jamjet/agent/tools/ToolRegistryTest.java index 0a4c487..38bb8a2 100644 --- a/jamjet-agent/src/test/java/dev/jamjet/agent/tools/ToolRegistryTest.java +++ b/jamjet-agent/src/test/java/dev/jamjet/agent/tools/ToolRegistryTest.java @@ -1,6 +1,7 @@ package dev.jamjet.agent.tools; import dev.jamjet.agent.TestTools; +import dev.jamjet.agent.Tool; import org.junit.jupiter.api.Test; import java.util.List; @@ -103,6 +104,31 @@ void rejectsDuplicateToolNames() { .hasMessageContaining("Duplicate @Tool name 'web_search'"); } + /** + * Two {@code @Tool} overloads (same Java method name) collide on the class+method + * dispatch key, even with distinct {@code @Tool.name}s. The second would silently + * overwrite the first in {@code byClassAndMethod} and corrupt dispatch, so + * registration must reject it loudly. + */ + public static final class OverloadedTools { + @Tool(name = "compute_one", description = "first overload") + public String compute(String a) { + return a; + } + + @Tool(name = "compute_two", description = "second overload") + public String compute(int a) { + return String.valueOf(a); + } + } + + @Test + void rejectsOverloadedToolMethodsSameClassAndMethod() { + assertThatThrownBy(() -> ToolRegistry.of(new OverloadedTools())) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("compute"); + } + @Test void emptyRegistryHasNoSchemas() { ToolRegistry registry = new ToolRegistry(); 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 6b5b31a..1b151ca 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 @@ -17,6 +17,8 @@ import java.util.List; import java.util.Map; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.BooleanSupplier; import static com.github.tomakehurst.wiremock.client.WireMock.aResponse; @@ -235,6 +237,52 @@ void fenceLostMidFlightAbortsAndNeverCompletes_M3() { wm.verify(0, postRequestedFor(urlEqualTo("/work-items/wi_1/fail"))); } + @Test + @Timeout(15) + void interruptedDispatchIsNotAToolFailureAndExitsCleanly() throws InterruptedException { + // The WORKER thread is interrupted (shutdown / cancellation) while it blocks on + // dispatchFuture.get(). That is NOT a tool failure: the item must be left for the + // engine to reclaim on lease expiry, never failed (which would clobber a reclaimer). + Map input = Map.of( + "last_model_tool_calls", List.of(toolCall("tc1", "slow_tool", Map.of("ignored", "x")))); + + wm.stubFor(post(urlEqualTo("/work-items/claim")).willReturn(okJson(validClaim(input, 7)))); + wm.stubFor(post(urlEqualTo("/work-items/wi_1/complete")).willReturn(ok())); + // Heartbeat interval is far longer than the test: the ONLY signal is the interrupt, + // so this isolates the interrupt path from the M3 heartbeat-abort path. + wm.stubFor(post(urlEqualTo("/work-items/wi_1/heartbeat")).willReturn(ok())); + wm.stubFor(post(urlEqualTo("/work-items/wi_1/fail")).willReturn(ok())); + + SlowTool slow = new SlowTool(); + AtomicReference resultRef = new AtomicReference<>(); + CountDownLatch done = new CountDownLatch(1); + + try (var client = new JamjetEngineClient(wm.baseUrl()); + var worker = new JavaToolWorker(client, "w1", + ToolRegistry.of(slow), + Duration.ofSeconds(30), Duration.ofMillis(10))) { + + Thread runner = new Thread(() -> { + try { + resultRef.set(worker.runOnce()); + } finally { + done.countDown(); + } + }, "worker-runner"); + runner.start(); + + // Once the tool is actually executing, interrupt the worker thread mid-get(). + assertThat(slow.started.await(5, TimeUnit.SECONDS)).isTrue(); + runner.interrupt(); + assertThat(done.await(5, TimeUnit.SECONDS)).isTrue(); + } + + // A clean stop (no-op), NOT a failure: the item was neither failed nor completed. + assertThat(resultRef.get()).isEqualTo(JavaToolWorker.ItemResult.LOST_LEASE); + wm.verify(0, postRequestedFor(urlEqualTo("/work-items/wi_1/fail"))); + wm.verify(0, postRequestedFor(urlEqualTo("/work-items/wi_1/complete"))); + } + @Test @Timeout(15) void conflictOnCompleteIsLostLeaseNotFailure() { diff --git a/jamjet-runtime-core/src/main/java/dev/jamjet/runtime/core/ir/NodeKind.java b/jamjet-runtime-core/src/main/java/dev/jamjet/runtime/core/ir/NodeKind.java index c65fd7e..844eb56 100644 --- a/jamjet-runtime-core/src/main/java/dev/jamjet/runtime/core/ir/NodeKind.java +++ b/jamjet-runtime-core/src/main/java/dev/jamjet/runtime/core/ir/NodeKind.java @@ -5,6 +5,9 @@ import com.fasterxml.jackson.annotation.JsonTypeInfo; import dev.jamjet.runtime.core.QueueType; +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; @@ -67,13 +70,55 @@ record Model( // call (mirrors the Rust Model.tools / Python agent_ir _model_kind). // Empty (never null) means no tools are offered; it serializes as // "tools":[] matching the Rust #[serde(default)] Vec field. - tools = tools == null ? List.of() : List.copyOf(tools); + // + // Deep-copy + freeze ALL THE WAY DOWN: List.copyOf only freezes the outer + // list, leaving nested schema maps/lists mutable so a caller could mutate a + // node's tool schema post-construction. The serde shape is unchanged. + tools = deepImmutableTools(tools); } /** Back-compat constructor: a Model node offering no tools (plain text completion). */ public Model(String modelRef, String promptRef, String outputSchema, String systemPrompt) { this(modelRef, promptRef, outputSchema, systemPrompt, List.of()); } + + private static List> deepImmutableTools(List> tools) { + if (tools == null || tools.isEmpty()) { + return List.of(); + } + List> out = new ArrayList<>(tools.size()); + for (Map tool : tools) { + out.add(deepImmutableMap(tool)); + } + return Collections.unmodifiableList(out); + } + + @SuppressWarnings("unchecked") + private static Map deepImmutableMap(Map map) { + if (map == null || map.isEmpty()) { + return Map.of(); + } + LinkedHashMap copy = new LinkedHashMap<>(map.size()); + for (Map.Entry e : map.entrySet()) { + copy.put(e.getKey(), deepImmutableValue(e.getValue())); + } + return Collections.unmodifiableMap(copy); + } + + @SuppressWarnings("unchecked") + private static Object deepImmutableValue(Object value) { + if (value instanceof Map m) { + return deepImmutableMap((Map) m); + } + if (value instanceof List list) { + List out = new ArrayList<>(list.size()); + for (Object e : list) { + out.add(deepImmutableValue(e)); + } + return Collections.unmodifiableList(out); + } + return value; + } } record Tool( diff --git a/jamjet-runtime-core/src/test/java/dev/jamjet/runtime/core/ir/NodeKindSerializationTest.java b/jamjet-runtime-core/src/test/java/dev/jamjet/runtime/core/ir/NodeKindSerializationTest.java index 7229a23..734f2ee 100644 --- a/jamjet-runtime-core/src/test/java/dev/jamjet/runtime/core/ir/NodeKindSerializationTest.java +++ b/jamjet-runtime-core/src/test/java/dev/jamjet/runtime/core/ir/NodeKindSerializationTest.java @@ -6,10 +6,13 @@ import dev.jamjet.runtime.core.QueueType; import org.junit.jupiter.api.Test; +import java.util.ArrayList; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; class NodeKindSerializationTest { @@ -85,6 +88,60 @@ void modelNodeWithoutToolsEmitsEmptyToolsArray() throws JsonProcessingException assertThat(((NodeKind.Model) deserialized).tools()).isEmpty(); } + @Test + void modelToolSchemasAreDeeplyImmutable() { + // A Model node freezes its tool schemas all the way down: mutating the SOURCE + // structures after construction must not leak in, and the node's nested maps are + // themselves unmodifiable (the outer-list-only copy of List.copyOf was not enough). + Map nestedProps = new LinkedHashMap<>(); + nestedProps.put("query", "string"); + Map params = new LinkedHashMap<>(); + params.put("type", "object"); + params.put("properties", nestedProps); + Map tool = new LinkedHashMap<>(); + tool.put("name", "web_search"); + tool.put("parameters", params); + List> mutableTools = new ArrayList<>(); + mutableTools.add(tool); + + NodeKind.Model model = new NodeKind.Model("gpt4", "", "", "be helpful", mutableTools); + + // Mutating the source structures must NOT affect the node (deep copy). + nestedProps.put("injected", "boom"); + mutableTools.add(Map.of("name", "evil")); + + assertThat(model.tools()).hasSize(1); + @SuppressWarnings("unchecked") + Map p = (Map) model.tools().get(0).get("parameters"); + @SuppressWarnings("unchecked") + Map props = (Map) p.get("properties"); + assertThat(props).doesNotContainKey("injected"); + + // The node's nested maps are unmodifiable. + assertThatThrownBy(() -> model.tools().get(0).put("x", "y")) + .isInstanceOf(UnsupportedOperationException.class); + assertThatThrownBy(() -> props.put("x", "y")) + .isInstanceOf(UnsupportedOperationException.class); + } + + @Test + void deserializesOlderModelJsonWithoutToolsFieldToEmpty() throws JsonProcessingException { + // Back-compat: older Model JSON had NO `tools` field at all. Jackson must default + // it to an empty list (compact ctor maps null -> List.of()), matching the Rust + // #[serde(default)] Vec, and round-trip to the canonical "tools":[] shape. + String olderJson = "{\"type\":\"model\",\"model_ref\":\"gpt4\"," + + "\"prompt_ref\":\"summarize\",\"output_schema\":\"{}\"," + + "\"system_prompt\":\"You are helpful\"}"; + + NodeKind deserialized = mapper.readValue(olderJson, NodeKind.class); + assertThat(deserialized).isInstanceOf(NodeKind.Model.class); + assertThat(((NodeKind.Model) deserialized).tools()).isEmpty(); + + String reserialized = mapper.writeValueAsString(deserialized); + assertThat(reserialized).contains("\"tools\":[]"); + assertThat(mapper.readValue(reserialized, NodeKind.class)).isEqualTo(deserialized); + } + @Test void computedHelpersNeverSerialize() throws JsonProcessingException { // queueType()/isDurable() are computed routing helpers, not IR fields: