feat(agent): add streaming delivery boundary with bounded format repair (T39)

argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/AgentResult.java

@@ -56,6 +56,9 @@ public record AgentResult(ChatMessage message, Status status, Usage usage, Objec
 	/** Error code for a terminal message that only reports a model failure. */
 	public static final String ERROR_CODE_MODEL_FAILED = "MODEL_FAILED";
 
+	/** Error code marking that the single bounded repair attempt still failed validation. */
+	public static final String ERROR_CODE_REPAIRED_STILL_INVALID = "REPAIRED_STILL_INVALID";
+
 	/**
 	 * Delivery status of the run.
 	 */

argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/ReactAgent.java

@@ -20,6 +20,8 @@
 import io.github.agentic.ai.api.chat.ContextAuthorityKeys;
 import io.github.agentic.ai.api.chat.MessageProvenance;
 import io.github.agentic.ai.api.chat.Usage;
+import io.github.agentic.ai.api.chat.ChatClient;
+import io.github.agentic.ai.api.chat.ChatResponse;
 
 import io.github.agentic.ai.graph.CompileConfig;
 import io.github.agentic.ai.graph.CompiledGraph;
@@ -65,6 +67,8 @@
 import io.github.agentic.ai.graph.internal.node.Node;
 import io.github.agentic.ai.graph.internal.node.ResumableSubGraphAction;
 import io.github.agentic.ai.graph.serializer.AgentInstructionMessage;
+import io.github.agentic.ai.graph.streaming.OutputType;
+import io.github.agentic.ai.graph.streaming.StreamingOutput;
 import io.github.agentic.ai.graph.serializer.StateSerializer;
 import io.github.agentic.ai.graph.serializer.plain_text.jackson.ArgiJacksonStateSerializer;
 import io.github.agentic.ai.graph.state.strategy.AppendStrategy;
@@ -89,6 +93,7 @@
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.function.Supplier;
@@ -111,6 +116,10 @@ public class ReactAgent extends BaseAgent {
 
 	private static final String MODEL_ERROR_MESSAGE_PREFIX = "Exception: ";
 
+	private static final String FINAL_RESPONSE_METADATA_KEY = "argiFinalResponse";
+
+	private final ChatClient chatClient;
+
 	private volatile StructuredOutputFinalizer structuredOutputFinalizer;
 
 	private final ConcurrentMap<String, Map<String, Object>> threadIdStateMap;
@@ -151,6 +160,7 @@ public ReactAgent(AgentLlmNode llmNode, AgentToolNode toolNode, CompileConfig co
 		this.inputType = builder.inputType;
 		this.outputSchema = builder.outputSchema;
 		this.outputType = builder.outputType;
+		this.chatClient = builder.chatClient;
 
 		// Set state serializer from builder, or use default
 		// Default to Jackson serializer for better compatibility and features
@@ -338,6 +348,139 @@ public AgentResult callResult(Map<String, Object> inputs, RunnableConfig config)
 		return doResultInvoke(inputs, config);
 	}
 
+	/**
+	 * Streams the agent run as provisional {@link AgentResult} chunks followed by the
+	 * final validated delivery. Chunks with a {@code null} status are provisional and
+	 * explicitly unvalidated; the last element of the stream carries the delivery
+	 * outcome of the terminal message (one of the four terminal statuses), finalized
+	 * through the same boundary as {@link #callResult(String)}. A run that ends without
+	 * a terminal response is reported as INTERRUPTED, never as completed.
+	 * @param message the user input message
+	 * @return the provisional chunk stream terminated by the validated result
+	 */
+	public Flux<AgentResult> streamResult(String message) {
+		return doStreamResult(message, null);
+	}
+
+	/**
+	 * Streams the agent run with a runtime configuration, mirroring
+	 * {@link #streamResult(String)}.
+	 * @param message the user input message
+	 * @param config runtime configuration controlling execution behavior
+	 * @return the provisional chunk stream terminated by the validated result
+	 */
+	public Flux<AgentResult> streamResult(String message, RunnableConfig config) {
+		return doStreamResult(message, config);
+	}
+
+	/**
+	 * Streams the agent run with a {@link ChatMessage} input, mirroring
+	 * {@link #streamResult(String)}.
+	 * @param message the input message
+	 * @return the provisional chunk stream terminated by the validated result
+	 */
+	public Flux<AgentResult> streamResult(ChatMessage message) {
+		return doStreamResult(message, null);
+	}
+
+	/**
+	 * Streams the agent run with a {@link ChatMessage} input and runtime configuration,
+	 * mirroring {@link #streamResult(String)}.
+	 * @param message the input message
+	 * @param config runtime configuration controlling execution behavior
+	 * @return the provisional chunk stream terminated by the validated result
+	 */
+	public Flux<AgentResult> streamResult(ChatMessage message, RunnableConfig config) {
+		return doStreamResult(message, config);
+	}
+
+	/**
+	 * Streams the agent run with a message list input, mirroring
+	 * {@link #streamResult(String)}.
+	 * @param messages the input messages
+	 * @return the provisional chunk stream terminated by the validated result
+	 */
+	public Flux<AgentResult> streamResult(List<ChatMessage> messages) {
+		return doStreamResult(messages, null);
+	}
+
+	/**
+	 * Streams the agent run with a message list input and runtime configuration,
+	 * mirroring {@link #streamResult(String)}.
+	 * @param messages the input messages
+	 * @param config runtime configuration controlling execution behavior
+	 * @return the provisional chunk stream terminated by the validated result
+	 */
+	public Flux<AgentResult> streamResult(List<ChatMessage> messages, RunnableConfig config) {
+		return doStreamResult(messages, config);
+	}
+
+	/**
+	 * Streams the agent run with an inputs map, mirroring {@link #streamResult(String)}.
+	 * @param inputs the input map (reserved: messages, input; other keys as state)
+	 * @return the provisional chunk stream terminated by the validated result
+	 */
+	public Flux<AgentResult> streamResult(Map<String, Object> inputs) {
+		return doStreamResult(inputs, null);
+	}
+
+	/**
+	 * Streams the agent run with an inputs map and runtime configuration, mirroring
+	 * {@link #streamResult(String)}.
+	 * @param inputs the input map (reserved: messages, input; other keys as state)
+	 * @param config runtime configuration controlling execution behavior
+	 * @return the provisional chunk stream terminated by the validated result
+	 */
+	public Flux<AgentResult> streamResult(Map<String, Object> inputs, RunnableConfig config) {
+		return doStreamResult(inputs, config);
+	}
+
+	private Flux<AgentResult> doStreamResult(Object message, RunnableConfig config) {
+		return doStream(buildMessageInput(message), config).transform(this::finalizeStream);
+	}
+
+	private Flux<AgentResult> doStreamResult(Map<String, Object> inputs, RunnableConfig config) {
+		return doStream(normalizeMessageInput(inputs), config).transform(this::finalizeStream);
+	}
+
+	/**
+	 * Consumes the node-output stream: text-delta chunks are re-published as provisional
+	 * results (status {@code null}); the final aggregated response (marked by the
+	 * existing {@code argiFinalResponse} protocol convention) is buffered and finalized
+	 * once the stream completes, so unverified chunks are never reported as a delivery.
+	 */
+	private Flux<AgentResult> finalizeStream(Flux<NodeOutput> outputs) {
+		AtomicReference<ChatResponse> finalResponse = new AtomicReference<>();
+		return outputs.concatMap(output -> {
+			if (!(output instanceof StreamingOutput<?> streamed)) {
+				return Flux.empty();
+			}
+			if (streamed.getOriginData() instanceof ChatResponse response
+					&& Boolean.TRUE.equals(response.metadata().get(FINAL_RESPONSE_METADATA_KEY))) {
+				finalResponse.set(response);
+				return Flux.empty();
+			}
+			if (streamed.getOutputType() == OutputType.AGENT_MODEL_STREAMING) {
+				return Flux.fromIterable(streamed.messages())
+					.map(message -> new AgentResult(message, null, Usage.unknown(), null, List.of()));
+			}
+			return Flux.empty();
+		}).concatWith(Flux.defer(() -> {
+			ChatResponse response = finalResponse.get();
+			if (response == null) {
+				return Flux.just(AgentResult.interrupted(Usage.unknown()));
+			}
+			ChatMessage terminal = response.messages().stream()
+				.filter(message -> message.role() == ChatRole.ASSISTANT)
+				.reduce((first, second) -> second)
+				.orElse(null);
+			if (terminal == null) {
+				return Flux.just(AgentResult.interrupted(response.usage()));
+			}
+			return Flux.just(deliver(terminal, response.usage()));
+		}));
+	}
+
 	private AgentResult doResultInvoke(Object message, RunnableConfig config) {
 		try {
 			return finalizeResult(doInvokeAndGetOutput(buildMessageInput(message), config));
@@ -373,6 +516,11 @@ private AgentResult finalizeResult(Optional<NodeOutput> output) {
 			// interruption): report it as unfinished, never as completed.
 			return AgentResult.interrupted(usage);
 		}
+		return deliver(terminal, usage);
+	}
+
+	/** Maps a terminal message onto the delivery boundary shared by call and stream paths. */
+	private AgentResult deliver(ChatMessage terminal, Usage usage) {
 		if (terminal.text() != null && terminal.text().startsWith(MODEL_ERROR_MESSAGE_PREFIX)) {
 			// throwOnModelError=false surfaces model failures as assistant text.
 			return AgentResult.error(terminal, usage,
@@ -387,12 +535,15 @@ private AgentResult finalizeResult(Optional<NodeOutput> output) {
 	 * the JSON Schema validator classes, which are not guaranteed to be on every
 	 * consumer's classpath (the studio module, for one, resolves an older
 	 * json-schema-validator for its MCP client); agents that never finalize a
-	 * schema-constrained output must not trigger that load.
+	 * schema-constrained output must not trigger that load. When the agent has a chat
+	 * client, the finalizer gains the bounded repair capability.
 	 */
 	private StructuredOutputFinalizer finalizer() {
 		StructuredOutputFinalizer finalizer = this.structuredOutputFinalizer;
 		if (finalizer == null) {
-			finalizer = new StructuredOutputFinalizer();
+			finalizer = this.chatClient != null
+					? new StructuredOutputFinalizer(this.chatClient)
+					: new StructuredOutputFinalizer();
 			this.structuredOutputFinalizer = finalizer;
 		}
 		return finalizer;

argi-agent-framework/src/main/java/io/github/agentic/ai/graph/agent/result/StructuredOutputFinalizer.java

@@ -17,14 +17,22 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
 
 import com.networknt.schema.Error;
 import com.networknt.schema.Schema;
 import com.networknt.schema.SchemaRegistry;
 import com.networknt.schema.SpecificationVersion;
 import org.springframework.util.StringUtils;
 
+import io.github.agentic.ai.api.chat.ChatClient;
 import io.github.agentic.ai.api.chat.ChatMessage;
+import io.github.agentic.ai.api.chat.ChatRequest;
+import io.github.agentic.ai.api.chat.ChatResponse;
+import io.github.agentic.ai.api.chat.ChatRole;
+import io.github.agentic.ai.api.chat.GenerationOptions;
+import io.github.agentic.ai.api.chat.ToolChoice;
 import io.github.agentic.ai.api.chat.Usage;
 import io.github.agentic.ai.graph.agent.AgentResult;
 import tools.jackson.core.StreamReadFeature;
@@ -37,7 +45,7 @@
  * and schema-validating the configured output schema. This is the single delivery
  * boundary shared by every termination path (normal final message, returnDirect
  * synthesis, hook {@code JumpTo.end}): the terminal message itself is carried through
- * unchanged and never regenerated.
+ * unchanged and never regenerated by the validation itself.
  *
  * <p>
  * Prompt-level schema constraints do not count as validation: without this component
@@ -52,6 +60,14 @@
  * schema violations) becomes {@link AgentResult.Status#VALIDATION_FAILED} with issue
  * entries that do not echo the full output content.
  * </p>
+ *
+ * <p>
+ * When constructed with a {@link ChatClient}, a failed validation gets exactly one
+ * bounded repair attempt: a single-turn model re-ask that carries no tools (and never
+ * touches the graph), after which the repaired output is re-validated without a second
+ * repair. A failed repair is marked with {@link #ERROR_CODE_REPAIRED_STILL_INVALID};
+ * a failing repair client leaves the original failure untouched.
+ * </p>
  */
 public final class StructuredOutputFinalizer {
 
@@ -64,6 +80,23 @@ public final class StructuredOutputFinalizer {
 
 	private static final String ROOT_PATH = "$";
 
+	private final ChatClient repairClient;
+
+	/** Creates a pure validating finalizer without repair capability. */
+	public StructuredOutputFinalizer() {
+		this(null);
+	}
+
+	/**
+	 * Creates a finalizer that performs one bounded repair re-ask through the given
+	 * client when validation fails.
+	 * @param repairClient the client used for the single repair re-ask; the request
+	 * carries no tools and never touches the graph
+	 */
+	public StructuredOutputFinalizer(ChatClient repairClient) {
+		this.repairClient = repairClient;
+	}
+
 	/**
 	 * Finalizes the terminal message of a run against the configured output schema.
 	 * @param terminal the terminal assistant message; must not be {@code null}
@@ -87,6 +120,28 @@ public AgentResult finalizeOutput(ChatMessage terminal, String outputSchema, Usa
 							ROOT_PATH, "Configured output schema is not a valid JSON Schema")));
 		}
 
+		AgentResult first = validateAgainst(terminal, schema, usage);
+		if (first.status() != AgentResult.Status.VALIDATION_FAILED || this.repairClient == null) {
+			return first;
+		}
+
+		ChatMessage repaired = attemptRepair(terminal, outputSchema, first.issues());
+		if (repaired == null) {
+			// The repair client failed: keep the original failure untouched.
+			return first;
+		}
+		AgentResult second = validateAgainst(repaired, schema, usage);
+		if (second.status() != AgentResult.Status.VALIDATION_FAILED) {
+			return second;
+		}
+		List<AgentResult.OutputValidationIssue> marked = new ArrayList<>();
+		marked.add(new AgentResult.OutputValidationIssue(AgentResult.ERROR_CODE_REPAIRED_STILL_INVALID, ROOT_PATH,
+				"One bounded repair attempt did not produce a schema-valid output"));
+		marked.addAll(second.issues());
+		return AgentResult.validationFailed(repaired, usage, List.copyOf(marked));
+	}
+
+	private AgentResult validateAgainst(ChatMessage terminal, Schema schema, Usage usage) {
 		JsonNode root;
 		try {
 			root = MAPPER.readTree(terminal.text());
@@ -110,4 +165,33 @@ public AgentResult finalizeOutput(ChatMessage terminal, String outputSchema, Usa
 		return AgentResult.completed(terminal, usage, root);
 	}
 
+	/**
+	 * One bounded repair re-ask: single-turn, no tools, no graph. Returns the repaired
+	 * assistant message, or {@code null} when the repair client itself failed.
+	 */
+	private ChatMessage attemptRepair(ChatMessage terminal, String outputSchema,
+			List<AgentResult.OutputValidationIssue> issues) {
+		try {
+			String issueSummary = issues.stream()
+				.map(issue -> issue.errorCode()
+						+ (StringUtils.hasLength(issue.fieldPath()) && !ROOT_PATH.equals(issue.fieldPath())
+								? " at " + issue.fieldPath() : ""))
+				.collect(Collectors.joining("; "));
+			ChatMessage system = ChatMessage
+				.system("Return only valid JSON matching this JSON Schema without any surrounding text: "
+						+ outputSchema);
+			ChatMessage user = ChatMessage.user("The previous response failed schema validation (" + issueSummary
+					+ "). Return the corrected JSON only.\nPrevious response:\n" + terminal.text());
+			ChatResponse response = this.repairClient.callBlocking(new ChatRequest(List.of(system, user), List.of(),
+					GenerationOptions.builder().toolChoice(ToolChoice.none()).build(), Map.of()));
+			return response.messages().stream()
+				.filter(message -> message.role() == ChatRole.ASSISTANT)
+				.findFirst()
+				.orElse(null);
+		}
+		catch (RuntimeException ex) {
+			return null;
+		}
+	}
+
 }

argi-agent-framework/src/test/java/io/github/agentic/ai/graph/agent/ReactAgentCallResultTest.java

@@ -73,8 +73,11 @@ void invalidOutputFailsValidation() {
 		AgentResult result = agent.callResult("give me the status");
 
 		assertThat(result.status()).isEqualTo(AgentResult.Status.VALIDATION_FAILED);
-		assertThat(result.issues()).hasSize(1);
-		assertThat(result.issues().get(0).errorCode()).isEqualTo(AgentResult.ERROR_CODE_INVALID_JSON);
+		// The stub answers the bounded repair with the same invalid text, so the
+		// REPAIRED_STILL_INVALID marker is prepended to the original issue.
+		assertThat(result.issues()).hasSize(2);
+		assertThat(result.issues().get(0).errorCode()).isEqualTo(AgentResult.ERROR_CODE_REPAIRED_STILL_INVALID);
+		assertThat(result.issues().get(1).errorCode()).isEqualTo(AgentResult.ERROR_CODE_INVALID_JSON);
 	}
 
 	@Test
@@ -85,8 +88,9 @@ void schemaViolationReported() {
 		AgentResult result = agent.callResult("give me the status");
 
 		assertThat(result.status()).isEqualTo(AgentResult.Status.VALIDATION_FAILED);
-		assertThat(result.issues()).isNotEmpty();
-		assertThat(result.issues()).allMatch(issue -> AgentResult.ERROR_CODE_SCHEMA_VIOLATION.equals(issue.errorCode()));
+		// Repair re-asks with the same payload: the marker precedes the actual violations.
+		assertThat(result.issues()).anyMatch(issue -> AgentResult.ERROR_CODE_REPAIRED_STILL_INVALID.equals(issue.errorCode()));
+		assertThat(result.issues()).anyMatch(issue -> AgentResult.ERROR_CODE_SCHEMA_VIOLATION.equals(issue.errorCode()));
 	}
 
 	@Test