From 58fa2998a26897ec68fb1eebc3fcb81021a9c702 Mon Sep 17 00:00:00 2001 From: konlue Date: Sat, 7 Nov 2026 00:18:48 +0800 Subject: [PATCH 1/4] fix(workflow): validate declared input and output schemas at publish and run boundaries (issue #274) - WorkflowDefinitionValidator: parse non-blank inputSchema/outputSchema with the networknt draft 2020-12 registry; blank rejection and no-schema behavior unchanged - WorkflowAgent: compile declared schemas at construction; validate request inputs before the first step (no tool or agent side effects); validate the output snapshot on the shared completion path before delivering COMPLETED; failures carry AgentResult.OutputValidationIssue and terminate the stream with WorkflowValidationException - ExecutionResult.continueWith: keep the suspension-point snapshot so a suspended outcome can be built (pre-existing NPE, untested path) - boundaries: resumes do not re-validate inputs; subflow input schemas are a publish-time contract, not an extra gate inside the parent run - tests: WorkflowSchemaContractTest, 10 cases (red 5F+2E before, green after) --- argi-orchestration/argi-workflow-core/pom.xml | 4 + .../agentic/ai/workflow/WorkflowAgent.java | 132 ++++++-- .../workflow/WorkflowDefinitionValidator.java | 39 ++- .../workflow/WorkflowValidationException.java | 39 +++ .../workflow/WorkflowSchemaContractTest.java | 299 ++++++++++++++++++ 5 files changed, 473 insertions(+), 40 deletions(-) create mode 100644 argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowValidationException.java create mode 100644 argi-orchestration/argi-workflow-core/src/test/java/io/github/agentic/ai/workflow/WorkflowSchemaContractTest.java diff --git a/argi-orchestration/argi-workflow-core/pom.xml b/argi-orchestration/argi-workflow-core/pom.xml index 3764bf7bb..7f1556249 100644 --- a/argi-orchestration/argi-workflow-core/pom.xml +++ b/argi-orchestration/argi-workflow-core/pom.xml @@ -43,6 +43,10 @@ com.fasterxml.jackson.core jackson-databind + + com.networknt + json-schema-validator + org.junit.jupiter junit-jupiter diff --git a/argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowAgent.java b/argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowAgent.java index 42af2fc0a..5c78d9cca 100644 --- a/argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowAgent.java +++ b/argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowAgent.java @@ -16,11 +16,16 @@ package io.github.agentic.ai.workflow; import com.fasterxml.jackson.databind.ObjectMapper; +import com.networknt.schema.Error; +import com.networknt.schema.Schema; +import com.networknt.schema.SchemaRegistry; +import com.networknt.schema.SpecificationVersion; import io.github.agentic.ai.agent.Agent; import io.github.agentic.ai.agent.AgentEvent; import io.github.agentic.ai.agent.AgentExecutionContext; import io.github.agentic.ai.agent.AgentOutcome; import io.github.agentic.ai.agent.AgentRequest; +import io.github.agentic.ai.agent.AgentResult; import io.github.agentic.ai.api.chat.ChatMessage; import io.github.agentic.ai.api.chat.ContentPart; import io.github.agentic.ai.api.chat.ToolCall; @@ -44,13 +49,31 @@ import java.util.concurrent.Flow; import java.util.concurrent.Semaphore; import java.util.function.Function; +import tools.jackson.databind.JsonNode; +import tools.jackson.databind.json.JsonMapper; /** * Independent workflow executor. It interprets the closed workflow DSL directly and * never creates or references StateGraph, CompiledGraph or Loop runtime classes. + * + *

+ * When the definition declares an input schema, the request inputs are validated + * against it before the first step runs — before any tool or agent side effect. When + * it declares an output schema, the completed output snapshot is validated before the + * COMPLETED outcome is delivered; a violating output terminates the stream with a + * {@link WorkflowValidationException} instead. Inputs are validated on the fresh run + * only: a resume request continues from the state owned by its continuation and often + * carries no inputs, so re-validating them there would make suspended workflows + * impossible to complete. Subflows share the parent run's inputs, so a nested + * definition's input schema is a publish-time contract for its own callers and is not + * applied again at the subflow entry. */ public final class WorkflowAgent implements Agent { + private static final JsonMapper JSON = JsonMapper.builder().build(); + + private static final SchemaRegistry SCHEMAS = SchemaRegistry.withDefaultDialect(SpecificationVersion.DRAFT_2020_12); + private final WorkflowDefinition definition; private final Function agentResolver; @@ -65,6 +88,10 @@ public final class WorkflowAgent implements Agent { private final Map continuations = new ConcurrentHashMap<>(); + private final Schema inputSchema; + + private final Schema outputSchema; + private WorkflowAgent(Builder builder) { this.definition = Objects.requireNonNull(builder.definition, "definition must not be null"); this.definition.validate(); @@ -73,6 +100,8 @@ private WorkflowAgent(Builder builder) { this.registry = Objects.requireNonNull(builder.registry, "registry must not be null"); this.mapper = Objects.requireNonNull(builder.mapper, "mapper must not be null"); this.subflowMaxDepth = builder.subflowMaxDepth; + this.inputSchema = compileSchema(definition.inputSchema()); + this.outputSchema = compileSchema(definition.outputSchema()); } public static Builder builder(WorkflowDefinition definition) { @@ -90,6 +119,34 @@ public WorkflowDefinition definition() { return definition; } + private static Schema compileSchema(String schemaText) { + return schemaText == null ? null : SCHEMAS.getSchema(schemaText); + } + + private static void validate(Map value, Schema schema, String boundary) { + if (schema == null) { + return; + } + JsonNode node; + try { + node = JSON.valueToTree(value); + } + catch (RuntimeException exception) { + throw new WorkflowValidationException("Workflow " + boundary + " is not representable as JSON", + List.of(new AgentResult.OutputValidationIssue(AgentResult.ERROR_CODE_SCHEMA_MALFORMED, "$", + "Workflow " + boundary + " is not representable as JSON: " + exception.getMessage()))); + } + List issues = new ArrayList<>(); + for (Error error : schema.validate(node)) { + issues.add(new AgentResult.OutputValidationIssue(AgentResult.ERROR_CODE_SCHEMA_VIOLATION, + error.getInstanceLocation().toString(), error.getMessage())); + } + if (!issues.isEmpty()) { + throw new WorkflowValidationException( + "Workflow " + boundary + " does not match the declared " + boundary + "Schema", issues); + } + } + public static final class Builder { private final WorkflowDefinition definition; @@ -204,44 +261,51 @@ public synchronized void cancel() { } } - private void runSafely() { - try { - ExecutionResult result = run(); - if (!cancelled) { - if (result.suspension != null) { - String token = register(result.continuation); - AgentOutcome outcome = new AgentOutcome(AgentOutcome.Status.SUSPENDED, result.data.messages, - result.data.state, result.data.usage(), - new AgentOutcome.Suspension(token, result.suspension.prompts())); - emit(new AgentEvent.Finished(outcome)); - } - else { - AgentOutcome outcome = AgentOutcome.completed(result.data.messages, result.data.state, - result.data.usage()); - emit(new AgentEvent.Finished(outcome)); - } - subscriber.onComplete(); + private void runSafely() { + try { + ExecutionResult result = run(); + if (!cancelled) { + if (result.suspension != null) { + String token = register(result.continuation); + AgentOutcome outcome = new AgentOutcome(AgentOutcome.Status.SUSPENDED, result.data.messages, + result.data.state, result.data.usage(), + new AgentOutcome.Suspension(token, result.suspension.prompts())); + emit(new AgentEvent.Finished(outcome)); } - } - catch (Throwable error) { - if (!cancelled) { - subscriber.onError(error); + else { + // The output boundary: a completed delivery must match the declared + // output schema, whatever path (fresh run or resume) produced it. + validate(result.data.state, outputSchema, "output"); + AgentOutcome outcome = AgentOutcome.completed(result.data.messages, result.data.state, + result.data.usage()); + emit(new AgentEvent.Finished(outcome)); } + subscriber.onComplete(); } } - - private ExecutionResult run() { - context.control().assertActive(); - if (request.resumeToken() != null) { - StoredContinuation stored = continuations.remove(request.resumeToken()); - if (stored == null || !stored.matches(context)) { - throw new IllegalArgumentException("Unknown or unauthorized workflow resume token"); - } - return stored.continuation.resume(request, context); + catch (Throwable error) { + if (!cancelled) { + subscriber.onError(error); } - RunData data = new RunData(request); - return executeSteps(definition.steps(), 0, data, context); } + } + + private ExecutionResult run() { + context.control().assertActive(); + if (request.resumeToken() != null) { + StoredContinuation stored = continuations.remove(request.resumeToken()); + if (stored == null || !stored.matches(context)) { + throw new IllegalArgumentException("Unknown or unauthorized workflow resume token"); + } + // Inputs were validated on the fresh run that produced this continuation; + // a resume request carries the continuation's state forward. + return stored.continuation.resume(request, context); + } + RunData data = new RunData(request); + // The input boundary: refuse before any step, tool or agent side effect. + validate(data.inputs, inputSchema, "input"); + return executeSteps(definition.steps(), 0, data, context); + } private String register(Continuation continuation) { String token = definition.id() + ":" + definition.version() + ":" + UUID.randomUUID(); @@ -654,7 +718,9 @@ private static ExecutionResult suspended(RunData data, Map promp } private ExecutionResult continueWith(Continuation next) { - return new ExecutionResult(null, suspension, (request, context) -> { + // The wrapper keeps the suspension-point snapshot: the suspended outcome is + // built from this data before the continuation ever runs. + return new ExecutionResult(data, suspension, (request, context) -> { ExecutionResult result = continuation.resume(request, context); return result.suspension == null ? next.resume(request, context) : result; }); diff --git a/argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowDefinitionValidator.java b/argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowDefinitionValidator.java index b460af17b..d5aad20c1 100644 --- a/argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowDefinitionValidator.java +++ b/argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowDefinitionValidator.java @@ -15,6 +15,9 @@ */ package io.github.agentic.ai.workflow; +import com.networknt.schema.SchemaRegistry; +import com.networknt.schema.SpecificationVersion; + import java.util.ArrayList; import java.util.HashSet; import java.util.LinkedHashSet; @@ -23,19 +26,25 @@ import java.util.Set; import java.util.function.Function; -/** Structural validation performed before a definition can be published. */ +/** + * Structural validation performed before a definition can be published. + * + *

+ * Declared input and output schemas are checked for syntax against the supported + * JSON Schema dialect (draft 2020-12, the same validator the agent layer uses for + * structured outputs). A definition that declares no schema is unaffected; a blank + * schema is still rejected. + */ public final class WorkflowDefinitionValidator { + private static final SchemaRegistry SCHEMAS = SchemaRegistry.withDefaultDialect(SpecificationVersion.DRAFT_2020_12); + private WorkflowDefinitionValidator() { } public static void validate(WorkflowDefinition definition) { - if (definition.inputSchema() != null && definition.inputSchema().isBlank()) { - throw new IllegalArgumentException("inputSchema must not be blank"); - } - if (definition.outputSchema() != null && definition.outputSchema().isBlank()) { - throw new IllegalArgumentException("outputSchema must not be blank"); - } + validateSchema("inputSchema", definition.inputSchema()); + validateSchema("outputSchema", definition.outputSchema()); Set ids = new HashSet<>(); definition.steps().forEach(step -> validateStep(step, ids)); } @@ -108,6 +117,22 @@ else if (step instanceof WorkflowStep.Loop loop) { } } + private static void validateSchema(String field, String schema) { + if (schema == null) { + return; + } + if (schema.isBlank()) { + throw new IllegalArgumentException(field + " must not be blank"); + } + try { + SCHEMAS.getSchema(schema); + } + catch (RuntimeException exception) { + throw new IllegalArgumentException( + field + " is not a valid JSON Schema (draft 2020-12): " + exception.getMessage(), exception); + } + } + private static void validateStep(WorkflowStep step, Set ids) { if (!ids.add(step.id())) { throw new IllegalArgumentException("duplicate workflow step id: " + step.id()); diff --git a/argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowValidationException.java b/argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowValidationException.java new file mode 100644 index 000000000..65cb74ac5 --- /dev/null +++ b/argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowValidationException.java @@ -0,0 +1,39 @@ +/* + * Copyright 2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.github.agentic.ai.workflow; + +import io.github.agentic.ai.agent.AgentResult; +import java.util.List; + +/** + * A workflow input or output does not match the schema declared by the definition. + * Carries one structured issue per violating location, so the caller can point at + * the field and the error type instead of parsing a message. + */ +public class WorkflowValidationException extends WorkflowExecutionException { + + private final List issues; + + public WorkflowValidationException(String message, List issues) { + super(message); + this.issues = List.copyOf(issues); + } + + public List issues() { + return issues; + } + +} diff --git a/argi-orchestration/argi-workflow-core/src/test/java/io/github/agentic/ai/workflow/WorkflowSchemaContractTest.java b/argi-orchestration/argi-workflow-core/src/test/java/io/github/agentic/ai/workflow/WorkflowSchemaContractTest.java new file mode 100644 index 000000000..d6567bfdb --- /dev/null +++ b/argi-orchestration/argi-workflow-core/src/test/java/io/github/agentic/ai/workflow/WorkflowSchemaContractTest.java @@ -0,0 +1,299 @@ +/* + * Copyright 2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.github.agentic.ai.workflow; + +import io.github.agentic.ai.agent.Agent; +import io.github.agentic.ai.agent.AgentEvent; +import io.github.agentic.ai.agent.AgentExecutionContext; +import io.github.agentic.ai.agent.AgentOutcome; +import io.github.agentic.ai.agent.AgentRequest; +import io.github.agentic.ai.agent.AgentResult; +import io.github.agentic.ai.api.chat.ChatMessage; +import io.github.agentic.ai.api.chat.Usage; +import io.github.agentic.ai.api.tool.Tool; +import io.github.agentic.ai.api.tool.ToolContext; +import io.github.agentic.ai.api.tool.ToolRegistry; +import io.github.agentic.ai.api.tool.ToolResult; +import io.github.agentic.ai.runtime.execution.DelegationScope; +import io.github.agentic.ai.runtime.execution.ExecutionBudget; +import io.github.agentic.ai.runtime.execution.PrincipalReference; +import io.github.agentic.ai.runtime.execution.TenantReference; +import io.github.agentic.ai.runtime.execution.TrustedExecutionContext; +import java.time.Clock; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionStage; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Flow; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Contract tests for the input/output schema boundary of a workflow definition + * (issue #274): declared schemas must be validated at publish time, inputs must + * be validated before any step or tool side effect, and a final output that does + * not match the declared schema must fail the completion instead of delivering + * COMPLETED. + */ +class WorkflowSchemaContractTest { + + private static final String MALFORMED_SCHEMA = "{\"type\":\"objec"; + + private static final String INPUT_SCHEMA = + "{\"type\":\"object\",\"properties\":{\"question\":{\"type\":\"string\"}},\"required\":[\"question\"]}"; + + // The step outputKey stores the agent outputs map, so the output schema describes + // result.value. + private static final String OUTPUT_SCHEMA = + "{\"type\":\"object\",\"properties\":{\"result\":{\"type\":\"object\"," + + "\"properties\":{\"value\":{\"type\":\"string\"}},\"required\":[\"value\"]}}," + + "\"required\":[\"result\"]}"; + + @Test + void publishRejectsMalformedSchemas() { + WorkflowRegistry registry = new WorkflowRegistry(); + IllegalArgumentException input = assertThrows(IllegalArgumentException.class, + () -> registry.publish(definition(MALFORMED_SCHEMA, null))); + assertTrue(input.getMessage().contains("inputSchema")); + IllegalArgumentException output = assertThrows(IllegalArgumentException.class, + () -> registry.publish(definition(null, MALFORMED_SCHEMA))); + assertTrue(output.getMessage().contains("outputSchema")); + } + + @Test + void blankSchemasStayRejected() { + WorkflowRegistry registry = new WorkflowRegistry(); + assertThrows(IllegalArgumentException.class, () -> registry.publish(definition(" ", null))); + assertThrows(IllegalArgumentException.class, () -> registry.publish(definition(null, " "))); + } + + @Test + void undeclaredSchemasKeepExistingBehavior() throws Exception { + AtomicInteger executions = new AtomicInteger(); + WorkflowAgent workflow = WorkflowAgent.builder(definition(null, null)) + .tools(new ToolRegistry(List.of(countingTool("record", executions)))).build(); + try (AgentExecutionContext context = context()) { + AgentOutcome outcome = workflow + .invoke(new AgentRequest(List.of(ChatMessage.user("hello")), Map.of(), null), context) + .get(3, TimeUnit.SECONDS); + assertEquals(AgentOutcome.Status.COMPLETED, outcome.status()); + assertEquals(1, executions.get()); + } + } + + @Test + void missingRequiredInputIsRejectedBeforeAnySideEffect() { + AtomicInteger executions = new AtomicInteger(); + WorkflowAgent workflow = WorkflowAgent.builder(definition(INPUT_SCHEMA, null)) + .tools(new ToolRegistry(List.of(countingTool("record", executions)))).build(); + try (AgentExecutionContext context = context()) { + WorkflowValidationException failure = validationFailure(() -> workflow + .invoke(new AgentRequest(List.of(ChatMessage.user("hello")), Map.of(), null), context)); + assertEquals(AgentResult.ERROR_CODE_SCHEMA_VIOLATION, failure.issues().get(0).errorCode()); + assertEquals(0, executions.get()); + } + } + + @Test + void typeMismatchReportsTheFieldPath() { + AtomicInteger executions = new AtomicInteger(); + WorkflowAgent workflow = WorkflowAgent.builder(definition(INPUT_SCHEMA, null)) + .tools(new ToolRegistry(List.of(countingTool("record", executions)))).build(); + try (AgentExecutionContext context = context()) { + WorkflowValidationException failure = validationFailure(() -> workflow.invoke( + new AgentRequest(List.of(ChatMessage.user("hello")), Map.of("question", 123), null), context)); + assertTrue(failure.issues().get(0).fieldPath().contains("question")); + assertEquals(0, executions.get()); + } + } + + @Test + void validInputAndOutputComplete() throws Exception { + Agent agent = agent(request -> AgentOutcome.completed(List.of(ChatMessage.assistant("done")), + Map.of("value", "analysis complete"), new Usage(1L, 1L, 2L))); + WorkflowStep step = new WorkflowStep.Agent("answer", "ref", Map.of(), "result", null); + WorkflowDefinition definition = new WorkflowDefinition("schema", "1", List.of(step), INPUT_SCHEMA, + OUTPUT_SCHEMA, Map.of()); + WorkflowAgent workflow = WorkflowAgent.builder(definition).agents(Map.of("ref", agent)).build(); + try (AgentExecutionContext context = context()) { + AgentOutcome outcome = workflow + .invoke(new AgentRequest(List.of(ChatMessage.user("hello")), Map.of("question", "why"), null), + context) + .get(3, TimeUnit.SECONDS); + assertEquals(AgentOutcome.Status.COMPLETED, outcome.status()); + assertEquals(Map.of("value", "analysis complete"), outcome.outputs().get("result")); + } + } + + @Test + void outputViolationFailsInsteadOfCompleting() { + Agent agent = agent(request -> AgentOutcome.completed(List.of(ChatMessage.assistant("done")), + Map.of("value", 42), new Usage(1L, 1L, 2L))); + WorkflowStep step = new WorkflowStep.Agent("answer", "ref", Map.of(), "result", null); + WorkflowDefinition definition = new WorkflowDefinition("schema", "1", List.of(step), INPUT_SCHEMA, + OUTPUT_SCHEMA, Map.of()); + WorkflowAgent workflow = WorkflowAgent.builder(definition).agents(Map.of("ref", agent)).build(); + try (AgentExecutionContext context = context()) { + WorkflowValidationException failure = validationFailure(() -> workflow.invoke( + new AgentRequest(List.of(ChatMessage.user("hello")), Map.of("question", "why"), null), context)); + assertTrue(failure.issues().stream().anyMatch(issue -> issue.fieldPath().contains("result"))); + } + } + + @Test + void resumeSkipsInputRevalidationButStillAcceptsValidOutputs() throws Exception { + WorkflowAgent workflow = suspendingWorkflow("done"); + try (AgentExecutionContext first = context()) { + AgentOutcome suspended = workflow + .invoke(new AgentRequest(List.of(ChatMessage.user("hello")), Map.of("question", "why"), null), + first) + .get(3, TimeUnit.SECONDS); + assertEquals(AgentOutcome.Status.SUSPENDED, suspended.status()); + String token = suspended.suspension().resumeToken(); + try (AgentExecutionContext resumed = context()) { + AgentOutcome outcome = workflow + .invoke(new AgentRequest(List.of(), Map.of(), token), resumed) + .get(3, TimeUnit.SECONDS); + // The resume request carries no inputs at all: the schema was enforced + // on the fresh run, the continuation owns the state afterwards. + assertEquals(AgentOutcome.Status.COMPLETED, outcome.status()); + } + } + } + + @Test + void resumeOutputViolationStillFails() throws Exception { + WorkflowAgent workflow = suspendingWorkflow(42); + try (AgentExecutionContext first = context()) { + AgentOutcome suspended = workflow + .invoke(new AgentRequest(List.of(ChatMessage.user("hello")), Map.of("question", "why"), null), + first) + .get(3, TimeUnit.SECONDS); + String token = suspended.suspension().resumeToken(); + try (AgentExecutionContext resumed = context()) { + WorkflowValidationException failure = validationFailure( + () -> workflow.invoke(new AgentRequest(List.of(), Map.of(), token), resumed)); + assertTrue(failure.issues().stream().anyMatch(issue -> issue.fieldPath().contains("result"))); + } + } + } + + @Test + void subflowSchemasAreValidatedAtPublishAndDoNotGateParentInputs() throws Exception { + WorkflowRegistry registry = new WorkflowRegistry(); + WorkflowDefinition broken = new WorkflowDefinition("nested", "1", List.of(new WorkflowStep.Tool("step", + "record", Map.of())), MALFORMED_SCHEMA, null, Map.of()); + assertThrows(IllegalArgumentException.class, () -> registry.publish(broken)); + + WorkflowDefinition valid = new WorkflowDefinition("nested", "1", List.of(new WorkflowStep.Tool("step", + "record", Map.of())), INPUT_SCHEMA, null, Map.of()); + registry.publish(valid); + WorkflowDefinition parent = new WorkflowDefinition("parent", "1", + List.of(new WorkflowStep.Subflow("call", "nested", "1")), null, null, Map.of()); + registry.publish(parent); + AtomicInteger executions = new AtomicInteger(); + WorkflowAgent workflow = WorkflowAgent.builder(parent).registry(registry) + .tools(new ToolRegistry(List.of(countingTool("record", executions)))).build(); + try (AgentExecutionContext context = context()) { + AgentOutcome outcome = workflow + .invoke(new AgentRequest(List.of(ChatMessage.user("hello")), Map.of("question", "why"), null), + context) + .get(3, TimeUnit.SECONDS); + // The subflow shares the parent inputs; the nested inputSchema is a publish- + // time contract for its own callers, not an extra gate inside the parent run. + assertEquals(AgentOutcome.Status.COMPLETED, outcome.status()); + assertEquals(1, executions.get()); + } + } + + private WorkflowAgent suspendingWorkflow(Object resumedValue) { + AtomicReference secondCallValue = new AtomicReference<>(resumedValue); + Agent agent = agent(request -> { + if (request.resumeToken() == null) { + return new AgentOutcome(AgentOutcome.Status.SUSPENDED, List.of(ChatMessage.assistant("need input")), + Map.of(), new Usage(1L, 1L, 2L), new AgentOutcome.Suspension("agent-token", + Map.of("prompt", "provide the missing confirmation"))); + } + return AgentOutcome.completed(List.of(ChatMessage.assistant("done")), + Map.of("value", secondCallValue.get()), new Usage(1L, 1L, 2L)); + }); + WorkflowStep step = new WorkflowStep.Agent("answer", "ref", Map.of(), "result", null); + WorkflowDefinition definition = new WorkflowDefinition("schema", "1", List.of(step), INPUT_SCHEMA, + OUTPUT_SCHEMA, Map.of()); + return WorkflowAgent.builder(definition).agents(Map.of("ref", agent)).build(); + } + + private WorkflowDefinition definition(String inputSchema, String outputSchema) { + return new WorkflowDefinition("schema", "1", + List.of(new WorkflowStep.Tool("record", "record", Map.of(), "result", null)), inputSchema, + outputSchema, Map.of()); + } + + private static Tool countingTool(String name, AtomicInteger executions) { + return new Tool() { + public io.github.agentic.ai.api.chat.ToolDefinition definition() { + return new io.github.agentic.ai.api.chat.ToolDefinition(name, name, "{}"); + } + + public CompletionStage execute(io.github.agentic.ai.api.chat.ToolCall call, + ToolContext context) { + executions.incrementAndGet(); + return CompletableFuture.completedFuture(ToolResult.text("recorded")); + } + }; + } + + private static WorkflowValidationException validationFailure( + java.util.function.Supplier> invocation) { + ExecutionException failure = assertThrows(ExecutionException.class, + () -> invocation.get().get(3, TimeUnit.SECONDS)); + return assertInstanceOf(WorkflowValidationException.class, failure.getCause()); + } + + private static Agent agent(java.util.function.Function function) { + return (request, context) -> subscriber -> { + subscriber.onSubscribe(new Flow.Subscription() { + public void request(long demand) { + if (demand > 0) { + subscriber.onNext(new AgentEvent.Finished(function.apply(request))); + subscriber.onComplete(); + } + } + + public void cancel() { + } + }); + }; + } + + private static AgentExecutionContext context() { + TrustedExecutionContext trusted = TrustedExecutionContext.builder() + .tenant(new TenantReference("test", "tenant")) + .principal(new PrincipalReference("test", "principal")) + .runId("run").correlationId("correlation").budget(ExecutionBudget.unlimited()) + .delegationScope(DelegationScope.none()).build(); + return new AgentExecutionContext(trusted, Clock.systemUTC()); + } + +} From 940cdcd06b0d114b91dc3298ad10ba495e911d68 Mon Sep 17 00:00:00 2001 From: konlue Date: Sun, 8 Nov 2026 10:55:29 +0800 Subject: [PATCH 2/4] fix(workflow): resume the suspended continuation exactly once Self-review of the #274 branch found that resuming a suspended workflow invoked the wrapped agent continuation twice: continueWith already resumed the continuation, and the composed step lambdas of executeSteps and Loop resumed it again. A real agent would reject the second call with the same resume token. The composed step now receives the resumed result through a ContinuedStep callback, and the resume tests assert a single invocation. Also fixes the indentation of the runSafely/run block (format:apply). --- .../agentic/ai/workflow/WorkflowAgent.java | 102 +++++++++--------- .../workflow/WorkflowSchemaContractTest.java | 13 ++- 2 files changed, 63 insertions(+), 52 deletions(-) diff --git a/argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowAgent.java b/argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowAgent.java index 5c78d9cca..c56acd323 100644 --- a/argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowAgent.java +++ b/argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowAgent.java @@ -261,51 +261,51 @@ public synchronized void cancel() { } } - private void runSafely() { - try { - ExecutionResult result = run(); - if (!cancelled) { - if (result.suspension != null) { - String token = register(result.continuation); - AgentOutcome outcome = new AgentOutcome(AgentOutcome.Status.SUSPENDED, result.data.messages, - result.data.state, result.data.usage(), - new AgentOutcome.Suspension(token, result.suspension.prompts())); - emit(new AgentEvent.Finished(outcome)); - } - else { - // The output boundary: a completed delivery must match the declared - // output schema, whatever path (fresh run or resume) produced it. - validate(result.data.state, outputSchema, "output"); - AgentOutcome outcome = AgentOutcome.completed(result.data.messages, result.data.state, - result.data.usage()); - emit(new AgentEvent.Finished(outcome)); + private void runSafely() { + try { + ExecutionResult result = run(); + if (!cancelled) { + if (result.suspension != null) { + String token = register(result.continuation); + AgentOutcome outcome = new AgentOutcome(AgentOutcome.Status.SUSPENDED, result.data.messages, + result.data.state, result.data.usage(), + new AgentOutcome.Suspension(token, result.suspension.prompts())); + emit(new AgentEvent.Finished(outcome)); + } + else { + // The output boundary: a completed delivery must match the declared + // output schema, whatever path (fresh run or resume) produced it. + validate(result.data.state, outputSchema, "output"); + AgentOutcome outcome = AgentOutcome.completed(result.data.messages, result.data.state, + result.data.usage()); + emit(new AgentEvent.Finished(outcome)); + } + subscriber.onComplete(); } - subscriber.onComplete(); } - } - catch (Throwable error) { - if (!cancelled) { - subscriber.onError(error); + catch (Throwable error) { + if (!cancelled) { + subscriber.onError(error); + } } } - } - private ExecutionResult run() { - context.control().assertActive(); - if (request.resumeToken() != null) { - StoredContinuation stored = continuations.remove(request.resumeToken()); - if (stored == null || !stored.matches(context)) { - throw new IllegalArgumentException("Unknown or unauthorized workflow resume token"); + private ExecutionResult run() { + context.control().assertActive(); + if (request.resumeToken() != null) { + StoredContinuation stored = continuations.remove(request.resumeToken()); + if (stored == null || !stored.matches(context)) { + throw new IllegalArgumentException("Unknown or unauthorized workflow resume token"); + } + // Inputs were validated on the fresh run that produced this continuation; + // a resume request carries the continuation's state forward. + return stored.continuation.resume(request, context); } - // Inputs were validated on the fresh run that produced this continuation; - // a resume request carries the continuation's state forward. - return stored.continuation.resume(request, context); - } - RunData data = new RunData(request); - // The input boundary: refuse before any step, tool or agent side effect. - validate(data.inputs, inputSchema, "input"); - return executeSteps(definition.steps(), 0, data, context); - } + RunData data = new RunData(request); + // The input boundary: refuse before any step, tool or agent side effect. + validate(data.inputs, inputSchema, "input"); + return executeSteps(definition.steps(), 0, data, context); + } private String register(Continuation continuation) { String token = definition.id() + ":" + definition.version() + ":" + UUID.randomUUID(); @@ -330,10 +330,8 @@ private ExecutionResult executeSteps(List steps, int index, RunDat if (result.suspension == null) { return executeSteps(steps, index + 1, data, context); } - return result.continueWith((request, resumeContext) -> { - ExecutionResult resumed = result.continuation.resume(request, resumeContext); - return resumed.suspension == null ? executeSteps(steps, index + 1, data, resumeContext) : resumed; - }); + return result.continueWith((resumed, request, resumeContext) -> + resumed.suspension == null ? executeSteps(steps, index + 1, data, resumeContext) : resumed); } private ExecutionResult executeStep(WorkflowStep step, RunData data, AgentExecutionContext context) { @@ -363,8 +361,7 @@ private ExecutionResult executeStep(WorkflowStep step, RunData data, AgentExecut ExecutionResult result = executeStep(loop.body(), data, context); if (result.suspension != null) { int nextIteration = iteration + 1; - return result.continueWith((request, resumeContext) -> { - ExecutionResult resumed = result.continuation.resume(request, resumeContext); + return result.continueWith((resumed, request, resumeContext) -> { if (resumed.suspension != null) { return resumed; } @@ -705,6 +702,12 @@ private interface Continuation { ExecutionResult resume(AgentRequest request, AgentExecutionContext context); } + /** The work that continues after the wrapped continuation has resumed exactly once. */ + @FunctionalInterface + private interface ContinuedStep { + ExecutionResult after(ExecutionResult resumed, AgentRequest request, AgentExecutionContext context); + } + private record SuspensionData(Map prompts) { } @@ -717,12 +720,13 @@ private static ExecutionResult suspended(RunData data, Map promp return new ExecutionResult(data, new SuspensionData(Map.copyOf(prompts)), continuation); } - private ExecutionResult continueWith(Continuation next) { + private ExecutionResult continueWith(ContinuedStep next) { // The wrapper keeps the suspension-point snapshot: the suspended outcome is - // built from this data before the continuation ever runs. + // built from this data before the continuation ever runs. The continuation + // itself runs exactly once; the step after the resume receives its result. return new ExecutionResult(data, suspension, (request, context) -> { - ExecutionResult result = continuation.resume(request, context); - return result.suspension == null ? next.resume(request, context) : result; + ExecutionResult resumed = continuation.resume(request, context); + return next.after(resumed, request, context); }); } } diff --git a/argi-orchestration/argi-workflow-core/src/test/java/io/github/agentic/ai/workflow/WorkflowSchemaContractTest.java b/argi-orchestration/argi-workflow-core/src/test/java/io/github/agentic/ai/workflow/WorkflowSchemaContractTest.java index d6567bfdb..162741209 100644 --- a/argi-orchestration/argi-workflow-core/src/test/java/io/github/agentic/ai/workflow/WorkflowSchemaContractTest.java +++ b/argi-orchestration/argi-workflow-core/src/test/java/io/github/agentic/ai/workflow/WorkflowSchemaContractTest.java @@ -163,7 +163,8 @@ void outputViolationFailsInsteadOfCompleting() { @Test void resumeSkipsInputRevalidationButStillAcceptsValidOutputs() throws Exception { - WorkflowAgent workflow = suspendingWorkflow("done"); + AtomicInteger resumedInvocations = new AtomicInteger(); + WorkflowAgent workflow = suspendingWorkflow("done", resumedInvocations); try (AgentExecutionContext first = context()) { AgentOutcome suspended = workflow .invoke(new AgentRequest(List.of(ChatMessage.user("hello")), Map.of("question", "why"), null), @@ -178,13 +179,16 @@ void resumeSkipsInputRevalidationButStillAcceptsValidOutputs() throws Exception // The resume request carries no inputs at all: the schema was enforced // on the fresh run, the continuation owns the state afterwards. assertEquals(AgentOutcome.Status.COMPLETED, outcome.status()); + // The continuation must resume the agent exactly once: a second + // invocation with the same token would be rejected by a real agent. + assertEquals(1, resumedInvocations.get()); } } } @Test void resumeOutputViolationStillFails() throws Exception { - WorkflowAgent workflow = suspendingWorkflow(42); + WorkflowAgent workflow = suspendingWorkflow(42, new AtomicInteger()); try (AgentExecutionContext first = context()) { AgentOutcome suspended = workflow .invoke(new AgentRequest(List.of(ChatMessage.user("hello")), Map.of("question", "why"), null), @@ -227,9 +231,12 @@ void subflowSchemasAreValidatedAtPublishAndDoNotGateParentInputs() throws Except } } - private WorkflowAgent suspendingWorkflow(Object resumedValue) { + private WorkflowAgent suspendingWorkflow(Object resumedValue, AtomicInteger resumedInvocations) { AtomicReference secondCallValue = new AtomicReference<>(resumedValue); Agent agent = agent(request -> { + if (request.resumeToken() != null) { + resumedInvocations.incrementAndGet(); + } if (request.resumeToken() == null) { return new AgentOutcome(AgentOutcome.Status.SUSPENDED, List.of(ChatMessage.assistant("need input")), Map.of(), new Usage(1L, 1L, 2L), new AgentOutcome.Suspension("agent-token", From 25868988aedfcdd1c43eafbf604b214b581ee432 Mon Sep 17 00:00:00 2001 From: konlue <2362586639@qq.com> Date: Mon, 9 Nov 2026 15:27:20 +0800 Subject: [PATCH 3/4] fix(workflow): reject schemas that violate the draft 2020-12 meta-schema at publish time (issue #274) --- .../workflow/WorkflowDefinitionValidator.java | 20 +++++++++++++++++ .../workflow/WorkflowSchemaContractTest.java | 22 ++++++++++++++++++- 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowDefinitionValidator.java b/argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowDefinitionValidator.java index d5aad20c1..45ccd41a2 100644 --- a/argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowDefinitionValidator.java +++ b/argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowDefinitionValidator.java @@ -15,6 +15,10 @@ */ package io.github.agentic.ai.workflow; +import com.networknt.schema.Error; +import com.networknt.schema.InputFormat; +import com.networknt.schema.Schema; +import com.networknt.schema.SchemaLocation; import com.networknt.schema.SchemaRegistry; import com.networknt.schema.SpecificationVersion; @@ -39,6 +43,9 @@ public final class WorkflowDefinitionValidator { private static final SchemaRegistry SCHEMAS = SchemaRegistry.withDefaultDialect(SpecificationVersion.DRAFT_2020_12); + private static final Schema META_SCHEMA = SCHEMAS + .getSchema(SchemaLocation.of(SpecificationVersion.DRAFT_2020_12.getDialectId())); + private WorkflowDefinitionValidator() { } @@ -124,6 +131,19 @@ private static void validateSchema(String field, String schema) { if (schema.isBlank()) { throw new IllegalArgumentException(field + " must not be blank"); } + // Compiling the document is not enough: a non-object document or an invalid + // keyword value would be accepted here and silently ignored at run time. + List errors; + try { + errors = META_SCHEMA.validate(schema, InputFormat.JSON); + } + catch (RuntimeException exception) { + throw new IllegalArgumentException(field + " is not valid JSON: " + exception.getMessage(), exception); + } + if (!errors.isEmpty()) { + throw new IllegalArgumentException( + field + " is not a valid JSON Schema (draft 2020-12): " + errors.get(0).getMessage()); + } try { SCHEMAS.getSchema(schema); } diff --git a/argi-orchestration/argi-workflow-core/src/test/java/io/github/agentic/ai/workflow/WorkflowSchemaContractTest.java b/argi-orchestration/argi-workflow-core/src/test/java/io/github/agentic/ai/workflow/WorkflowSchemaContractTest.java index 162741209..c5a8def4a 100644 --- a/argi-orchestration/argi-workflow-core/src/test/java/io/github/agentic/ai/workflow/WorkflowSchemaContractTest.java +++ b/argi-orchestration/argi-workflow-core/src/test/java/io/github/agentic/ai/workflow/WorkflowSchemaContractTest.java @@ -43,6 +43,8 @@ import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertInstanceOf; @@ -58,7 +60,7 @@ */ class WorkflowSchemaContractTest { - private static final String MALFORMED_SCHEMA = "{\"type\":\"objec"; + private static final String MALFORMED_SCHEMA = "{\"type\":\"object\""; private static final String INPUT_SCHEMA = "{\"type\":\"object\",\"properties\":{\"question\":{\"type\":\"string\"}},\"required\":[\"question\"]}"; @@ -81,6 +83,24 @@ void publishRejectsMalformedSchemas() { assertTrue(output.getMessage().contains("outputSchema")); } + @ParameterizedTest + @ValueSource(strings = { "{\"required\":\"question\"}", "[]", "123", "\"schema\"", "null" }) + void publishRejectsDocumentsThatAreNotJsonSchemas(String schema) { + WorkflowRegistry registry = new WorkflowRegistry(); + IllegalArgumentException input = assertThrows(IllegalArgumentException.class, + () -> registry.publish(definition(schema, null))); + assertTrue(input.getMessage().contains("inputSchema")); + IllegalArgumentException output = assertThrows(IllegalArgumentException.class, + () -> registry.publish(definition(null, schema))); + assertTrue(output.getMessage().contains("outputSchema")); + } + + @Test + void publishAcceptsObjectAndBooleanSchemas() { + new WorkflowRegistry().publish(definition("{\"type\":\"object\"}", "true")); + new WorkflowRegistry().publish(definition("false", "{}")); + } + @Test void blankSchemasStayRejected() { WorkflowRegistry registry = new WorkflowRegistry(); From 95e258683dc835d92844d203498347e40e9ee93a Mon Sep 17 00:00:00 2001 From: konlue Date: Tue, 10 Nov 2026 23:40:31 +0800 Subject: [PATCH 4/4] fix(workflow): preserve the remaining work across repeated suspensions (issue #274 review) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Resuming a suspended workflow step that suspends again returned the new suspension with a bare continuation: the composed step that continues the remaining steps was dropped, so the run completed without ever executing what followed. executeSteps and Loop now recompose their continue-after step onto every repeated suspension (continueAfter / continueLoopAfter), and two regression tests drive a child that suspends twice — the following step and the remaining loop iterations both run after the second resume. --- .../agentic/ai/workflow/WorkflowAgent.java | 55 ++++++--- .../workflow/WorkflowSchemaContractTest.java | 109 ++++++++++++++++++ 2 files changed, 147 insertions(+), 17 deletions(-) diff --git a/argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowAgent.java b/argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowAgent.java index c56acd323..647d6e37c 100644 --- a/argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowAgent.java +++ b/argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowAgent.java @@ -330,8 +330,19 @@ private ExecutionResult executeSteps(List steps, int index, RunDat if (result.suspension == null) { return executeSteps(steps, index + 1, data, context); } - return result.continueWith((resumed, request, resumeContext) -> - resumed.suspension == null ? executeSteps(steps, index + 1, data, resumeContext) : resumed); + return continueAfter(result, steps, index, data); + } + + /** + * Wraps a suspension so the remaining steps survive any number of further + * suspensions: every time the wrapped continuation suspends again, the + * continue-after step is composed onto the new suspension (issue #274 review). + */ + private ExecutionResult continueAfter(ExecutionResult suspended, List steps, int index, + RunData data) { + return suspended.continueWith((resumed, request, resumeContext) -> + resumed.suspension == null ? executeSteps(steps, index + 1, data, resumeContext) + : continueAfter(resumed, steps, index, data)); } private ExecutionResult executeStep(WorkflowStep step, RunData data, AgentExecutionContext context) { @@ -361,21 +372,7 @@ private ExecutionResult executeStep(WorkflowStep step, RunData data, AgentExecut ExecutionResult result = executeStep(loop.body(), data, context); if (result.suspension != null) { int nextIteration = iteration + 1; - return result.continueWith((resumed, request, resumeContext) -> { - if (resumed.suspension != null) { - return resumed; - } - for (int next = nextIteration; next < loop.maxIterations(); next++) { - if (!truthy(resolve(loop.condition(), data))) { - return ExecutionResult.completed(data); - } - ExecutionResult nested = executeStep(loop.body(), data, resumeContext); - if (nested.suspension != null) { - return nested; - } - } - return ExecutionResult.completed(data); - }); + return continueLoopAfter(result, loop, nextIteration, data); } } throw new WorkflowExecutionException("Workflow loop exceeded maxIterations: " + loop.id()); @@ -427,6 +424,30 @@ private ExecutionResult executeSubflow(WorkflowStep.Subflow subflow, RunData dat } } + /** + * The loop counterpart of {@link #continueAfter}: the remaining iterations survive + * any number of further suspensions, both of the wrapped continuation and of the + * body invocations inside the remaining iterations. + */ + private ExecutionResult continueLoopAfter(ExecutionResult suspended, WorkflowStep.Loop loop, int nextIteration, + RunData data) { + return suspended.continueWith((resumed, request, resumeContext) -> { + if (resumed.suspension != null) { + return continueLoopAfter(resumed, loop, nextIteration, data); + } + for (int next = nextIteration; next < loop.maxIterations(); next++) { + if (!truthy(resolve(loop.condition(), data))) { + return ExecutionResult.completed(data); + } + ExecutionResult nested = executeStep(loop.body(), data, resumeContext); + if (nested.suspension != null) { + return continueLoopAfter(nested, loop, next + 1, data); + } + } + return ExecutionResult.completed(data); + }); + } + private ExecutionResult executeAgent(WorkflowStep.Agent step, RunData data, AgentExecutionContext context) { Agent agent = Objects.requireNonNull(agentResolver.apply(step.agentRef()), "agent resolver returned no agent for " + step.agentRef()); diff --git a/argi-orchestration/argi-workflow-core/src/test/java/io/github/agentic/ai/workflow/WorkflowSchemaContractTest.java b/argi-orchestration/argi-workflow-core/src/test/java/io/github/agentic/ai/workflow/WorkflowSchemaContractTest.java index c5a8def4a..e5b934b12 100644 --- a/argi-orchestration/argi-workflow-core/src/test/java/io/github/agentic/ai/workflow/WorkflowSchemaContractTest.java +++ b/argi-orchestration/argi-workflow-core/src/test/java/io/github/agentic/ai/workflow/WorkflowSchemaContractTest.java @@ -206,6 +206,101 @@ void resumeSkipsInputRevalidationButStillAcceptsValidOutputs() throws Exception } } + @Test + void aTwiceSuspendedChildStillRunsTheFollowingSteps() throws Exception { + AtomicInteger invocations = new AtomicInteger(); + AtomicInteger resumedInvocations = new AtomicInteger(); + Agent twice = agent(request -> { + int call = invocations.incrementAndGet(); + if (request.resumeToken() != null) { + resumedInvocations.incrementAndGet(); + } + if (call <= 2) { + return new AgentOutcome(AgentOutcome.Status.SUSPENDED, List.of(ChatMessage.assistant("need more")), + Map.of(), new Usage(1L, 1L, 2L), new AgentOutcome.Suspension("agent-token-" + call, + Map.of("prompt", "provide confirmation " + call))); + } + return AgentOutcome.completed(List.of(ChatMessage.assistant("done")), + Map.of("value", "finally done"), new Usage(1L, 1L, 2L)); + }); + AtomicInteger toolExecutions = new AtomicInteger(); + WorkflowDefinition definition = new WorkflowDefinition("twice", "1", + List.of(new WorkflowStep.Agent("answer", "ref"), + new WorkflowStep.Tool("mark", "mark", Map.of(), "marked", null))); + WorkflowAgent workflow = WorkflowAgent.builder(definition).agents(Map.of("ref", twice)) + .tools(new ToolRegistry(List.of(countingTool(toolExecutions)))).build(); + + try (AgentExecutionContext first = context()) { + AgentOutcome suspendedOnce = workflow + .invoke(new AgentRequest(List.of(ChatMessage.user("go")), Map.of("question", "why"), null), first) + .get(3, TimeUnit.SECONDS); + assertEquals(AgentOutcome.Status.SUSPENDED, suspendedOnce.status()); + String firstToken = suspendedOnce.suspension().resumeToken(); + try (AgentExecutionContext second = context()) { + AgentOutcome suspendedTwice = workflow + .invoke(new AgentRequest(List.of(), Map.of(), firstToken), second) + .get(3, TimeUnit.SECONDS); + assertEquals(AgentOutcome.Status.SUSPENDED, suspendedTwice.status()); + String secondToken = suspendedTwice.suspension().resumeToken(); + try (AgentExecutionContext third = context()) { + AgentOutcome outcome = workflow + .invoke(new AgentRequest(List.of(), Map.of(), secondToken), third) + .get(3, TimeUnit.SECONDS); + assertEquals(AgentOutcome.Status.COMPLETED, outcome.status()); + // The step after the twice-suspended child must still run. + assertEquals(1, toolExecutions.get()); + assertEquals(2, resumedInvocations.get()); + } + } + } + } + + @Test + void aLoopBodySuspendingTwiceStillCompletesTheRemainingIterations() throws Exception { + AtomicInteger invocations = new AtomicInteger(); + Agent twice = agent(request -> { + int call = invocations.incrementAndGet(); + if (call <= 2) { + return new AgentOutcome(AgentOutcome.Status.SUSPENDED, List.of(ChatMessage.assistant("need more")), + Map.of(), new Usage(1L, 1L, 2L), new AgentOutcome.Suspension("agent-token-" + call, + Map.of("prompt", "provide confirmation " + call))); + } + // Every later invocation completes with a falsy value so the loop leaves + // through its condition once the twice-suspended first iteration is done. + return AgentOutcome.completed(List.of(ChatMessage.assistant("done")), + Map.of("value", false), new Usage(1L, 1L, 2L)); + }); + WorkflowDefinition definition = new WorkflowDefinition("looped", "1", + List.of(new WorkflowStep.Loop("looped", WorkflowValue.state("keep"), + new WorkflowStep.Agent("body", "ref"), 3))); + WorkflowAgent workflow = WorkflowAgent.builder(definition).agents(Map.of("ref", twice)).build(); + try (AgentExecutionContext first = context()) { + AgentOutcome suspendedOnce = workflow + .invoke(new AgentRequest(List.of(ChatMessage.user("go")), Map.of("keep", true), null), first) + .get(5, TimeUnit.SECONDS); + assertEquals(AgentOutcome.Status.SUSPENDED, suspendedOnce.status()); + String firstToken = suspendedOnce.suspension().resumeToken(); + try (AgentExecutionContext second = context()) { + AgentOutcome suspendedTwice = workflow + .invoke(new AgentRequest(List.of(), Map.of(), firstToken), second) + .get(5, TimeUnit.SECONDS); + assertEquals(AgentOutcome.Status.SUSPENDED, suspendedTwice.status()); + String secondToken = suspendedTwice.suspension().resumeToken(); + try (AgentExecutionContext third = context()) { + AgentOutcome outcome = workflow + .invoke(new AgentRequest(List.of(), Map.of(), secondToken), third) + .get(5, TimeUnit.SECONDS); + assertEquals(AgentOutcome.Status.COMPLETED, outcome.status()); + // Iteration 0 suspends twice (invocations 1-2) and completes on + // its third resume; iterations 1 and 2 still run (invocations 4-5) + // before the loop leaves through maxIterations. With the dropped- + // continuation bug this run ended after invocation 3. + assertEquals(5, invocations.get()); + } + } + } + } + @Test void resumeOutputViolationStillFails() throws Exception { WorkflowAgent workflow = suspendingWorkflow(42, new AtomicInteger()); @@ -298,6 +393,20 @@ private static WorkflowValidationException validationFailure( return assertInstanceOf(WorkflowValidationException.class, failure.getCause()); } + private static Tool countingTool(AtomicInteger executions) { + return new Tool() { + public io.github.agentic.ai.api.chat.ToolDefinition definition() { + return new io.github.agentic.ai.api.chat.ToolDefinition("mark", "mark", "{}"); + } + + public CompletionStage execute(io.github.agentic.ai.api.chat.ToolCall call, + ToolContext toolContext) { + executions.incrementAndGet(); + return CompletableFuture.completedFuture(ToolResult.text("marked")); + } + }; + } + private static Agent agent(java.util.function.Function function) { return (request, context) -> subscriber -> { subscriber.onSubscribe(new Flow.Subscription() {