fix(workflow): validate declared input and output schemas at publish and run boundaries (issue #274)

argi-orchestration/argi-workflow-core/pom.xml

@@ -43,6 +43,10 @@
             <groupId>com.fasterxml.jackson.core</groupId>
             <artifactId>jackson-databind</artifactId>
         </dependency>
+        <dependency>
+            <groupId>com.networknt</groupId>
+            <artifactId>json-schema-validator</artifactId>
+        </dependency>
         <dependency>
             <groupId>org.junit.jupiter</groupId>
             <artifactId>junit-jupiter</artifactId>

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.
+ *
+ * <p>
+ * 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<String, Agent> agentResolver;
@@ -65,6 +88,10 @@ public final class WorkflowAgent implements Agent {
 
 	private final Map<String, StoredContinuation> 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<String, Object> 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<AgentResult.OutputValidationIssue> 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;
@@ -216,6 +273,9 @@ private void runSafely() {
 						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));
@@ -237,9 +297,13 @@ private ExecutionResult run() {
 				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);
 		}
 
@@ -266,10 +330,19 @@ private ExecutionResult executeSteps(List<WorkflowStep> 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 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<WorkflowStep> 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) {
@@ -299,22 +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((request, resumeContext) -> {
-						ExecutionResult resumed = result.continuation.resume(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());
@@ -366,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());
@@ -641,6 +723,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<String, Object> prompts) {
 	}
 
@@ -653,10 +741,13 @@ private static ExecutionResult suspended(RunData data, Map<String, Object> promp
 			return new ExecutionResult(data, new SuspensionData(Map.copyOf(prompts)), continuation);
 		}
 
-		private ExecutionResult continueWith(Continuation next) {
-			return new ExecutionResult(null, suspension, (request, context) -> {
-				ExecutionResult result = continuation.resume(request, context);
-				return result.suspension == null ? next.resume(request, context) : result;
+		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. The continuation
+			// itself runs exactly once; the step after the resume receives its result.
+			return new ExecutionResult(data, suspension, (request, context) -> {
+				ExecutionResult resumed = continuation.resume(request, context);
+				return next.after(resumed, request, context);
 			});
 		}
 	}

argi-orchestration/argi-workflow-core/src/main/java/io/github/agentic/ai/workflow/WorkflowDefinitionValidator.java

@@ -15,6 +15,13 @@
  */
 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;
+
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.LinkedHashSet;
@@ -23,19 +30,28 @@
 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.
+ *
+ * <p>
+ * 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 static final Schema META_SCHEMA = SCHEMAS
+		.getSchema(SchemaLocation.of(SpecificationVersion.DRAFT_2020_12.getDialectId()));
+
 	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<String> ids = new HashSet<>();
 		definition.steps().forEach(step -> validateStep(step, ids));
 	}
@@ -108,6 +124,35 @@ 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");
+		}
+		// 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<Error> 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);
+		}
+		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<String> ids) {
 		if (!ids.add(step.id())) {
 			throw new IllegalArgumentException("duplicate workflow step id: " + step.id());

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<AgentResult.OutputValidationIssue> issues;
+
+	public WorkflowValidationException(String message, List<AgentResult.OutputValidationIssue> issues) {
+		super(message);
+		this.issues = List.copyOf(issues);
+	}
+
+	public List<AgentResult.OutputValidationIssue> issues() {
+		return issues;
+	}
+
+}