fix(core): retry invalid structured model responses

Author: BlueCatPro

core/src/main/java/com/google/adk/events/EventActions.java

@@ -45,6 +45,7 @@ public class EventActions extends JsonBaseModel {
   private ConcurrentMap<String, ToolConfirmation> requestedToolConfirmations;
   private boolean endOfAgent;
   private @Nullable EventCompaction compaction;
+  private @Nullable Object setModelResponse;
 
   /** Default constructor for Jackson. */
   public EventActions() {
@@ -67,6 +68,7 @@ private EventActions(Builder builder) {
     this.requestedToolConfirmations = builder.requestedToolConfirmations;
     this.endOfAgent = builder.endOfAgent;
     this.compaction = builder.compaction;
+    this.setModelResponse = builder.setModelResponse;
   }
 
   @JsonProperty("skipSummarization")
@@ -201,6 +203,19 @@ public void setCompaction(@Nullable EventCompaction compaction) {
     this.compaction = compaction;
   }
 
+  /**
+   * The successfully validated structured response set by the {@code set_model_response} tool.
+   * Empty when the tool was not called or its arguments failed output-schema validation.
+   */
+  @JsonProperty("setModelResponse")
+  public Optional<Object> setModelResponse() {
+    return Optional.ofNullable(setModelResponse);
+  }
+
+  public void setSetModelResponse(@Nullable Object setModelResponse) {
+    this.setModelResponse = setModelResponse;
+  }
+
   public static Builder builder() {
     return new Builder();
   }
@@ -226,7 +241,8 @@ public boolean equals(Object o) {
         && Objects.equals(requestedAuthConfigs, that.requestedAuthConfigs)
         && Objects.equals(requestedToolConfirmations, that.requestedToolConfirmations)
         && (endOfAgent == that.endOfAgent)
-        && Objects.equals(compaction, that.compaction);
+        && Objects.equals(compaction, that.compaction)
+        && Objects.equals(setModelResponse, that.setModelResponse);
   }
 
   @Override
@@ -241,7 +257,8 @@ public int hashCode() {
         requestedAuthConfigs,
         requestedToolConfirmations,
         endOfAgent,
-        compaction);
+        compaction,
+        setModelResponse);
   }
 
   /** Builder for {@link EventActions}. */
@@ -256,6 +273,7 @@ public static class Builder {
     private ConcurrentMap<String, ToolConfirmation> requestedToolConfirmations;
     private boolean endOfAgent = false;
     private @Nullable EventCompaction compaction;
+    private @Nullable Object setModelResponse;
 
     public Builder() {
       this.stateDelta = new ConcurrentHashMap<>();
@@ -277,6 +295,7 @@ private Builder(EventActions eventActions) {
           new ConcurrentHashMap<>(eventActions.requestedToolConfirmations());
       this.endOfAgent = eventActions.endOfAgent;
       this.compaction = eventActions.compaction;
+      this.setModelResponse = eventActions.setModelResponse;
     }
 
     @CanIgnoreReturnValue
@@ -383,6 +402,13 @@ public Builder compaction(@Nullable EventCompaction value) {
       return this;
     }
 
+    @CanIgnoreReturnValue
+    @JsonProperty("setModelResponse")
+    public Builder setModelResponse(@Nullable Object value) {
+      this.setModelResponse = value;
+      return this;
+    }
+
     @CanIgnoreReturnValue
     public Builder merge(EventActions other) {
       other.skipSummarization().ifPresent(this::skipSummarization);
@@ -395,6 +421,7 @@ public Builder merge(EventActions other) {
       this.requestedToolConfirmations.putAll(other.requestedToolConfirmations());
       this.endOfAgent = this.endOfAgent || other.endOfAgent();
       other.compaction().ifPresent(this::compaction);
+      other.setModelResponse().ifPresent(this::setModelResponse);
       return this;
     }
 

core/src/main/java/com/google/adk/flows/llmflows/OutputSchema.java

@@ -78,18 +78,24 @@ public Single<RequestProcessingResult> processRequest(
   }
 
   /**
-   * Check if function response contains set_model_response and extract JSON.
+   * Extracts a successfully validated {@code set_model_response} result as JSON.
+   *
+   * <p>Only a result that passed output-schema validation (recorded on the event actions by {@link
+   * SetModelResponseTool}) is returned. Validation feedback sent back to the model is never
+   * promoted to the final structured response.
    *
    * @param functionResponseEvent The function response event to check.
-   * @return JSON response string if set_model_response was called, Optional.empty() otherwise.
+   * @return JSON response string if set_model_response succeeded, Optional.empty() otherwise.
    */
   public static Optional<String> getStructuredModelResponse(Event functionResponseEvent) {
     for (FunctionResponse funcResponse : functionResponseEvent.functionResponses()) {
       if (Objects.equals(funcResponse.name().orElse(""), SetModelResponseTool.NAME)) {
-        Object response = funcResponse.response();
-        // The tool returns the args map directly.
+        Optional<Object> validatedResponse = functionResponseEvent.actions().setModelResponse();
+        if (validatedResponse.isEmpty()) {
+          return Optional.empty();
+        }
         try {
-          return Optional.of(JsonBaseModel.getMapper().writeValueAsString(response));
+          return Optional.of(JsonBaseModel.getMapper().writeValueAsString(validatedResponse.get()));
         } catch (JsonProcessingException e) {
           logger.error("Failed to serialize set_model_response result", e);
           return Optional.empty();

core/src/main/java/com/google/adk/tools/SetModelResponseTool.java

@@ -17,9 +17,14 @@
 package com.google.adk.tools;
 
 import com.google.adk.SchemaUtils;
+import com.google.common.collect.ImmutableMap;
 import com.google.genai.types.FunctionDeclaration;
 import com.google.genai.types.Schema;
+import com.google.genai.types.Type;
 import io.reactivex.rxjava3.core.Single;
+import java.util.ArrayList;
+import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 
@@ -33,6 +38,12 @@
 public class SetModelResponseTool extends BaseTool {
   public static final String NAME = "set_model_response";
 
+  // Prefix of the SchemaUtils validation message after which the full schema is appended. Used to
+  // strip the schema dump from feedback on a best-effort basis; if SchemaUtils changes its wording
+  // the feedback simply stays unstripped. runAsync_unknownArg_feedbackOmitsSchemaDump pins the
+  // current format.
+  private static final String OUTPUT_SCHEMA_DUMP_MARKER = " does not match agent output schema: ";
+
   private final Schema outputSchema;
 
   public SetModelResponseTool(Schema outputSchema) {
@@ -56,12 +67,66 @@ public Optional<FunctionDeclaration> declaration() {
 
   @Override
   public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContext toolContext) {
-    // This tool is a marker for the final response, it doesn't do anything but return its arguments
-    // which will be captured as the final result.
+    // Record validated responses on the event actions; return validation feedback so the model can
+    // retry.
     return Single.fromCallable(
         () -> {
-          SchemaUtils.validateMapOnSchema(args, outputSchema, /* isInput= */ false);
-          return args;
+          try {
+            SchemaUtils.validateMapOnSchema(args, outputSchema, /* isInput= */ false);
+          } catch (IllegalArgumentException e) {
+            return ImmutableMap.of(
+                "error",
+                "Validation Error found:\n"
+                    + sanitizeValidationMessage(e.getMessage())
+                    + "\nRecall the set_model_response function correctly, fix the errors, and"
+                    + " call it again with all required fields using the correct types.");
+          }
+          // Match Python's model_dump(exclude_none=True) for Java's map-shaped response.
+          Map<String, Object> validatedResponse = excludeNullFields(args, outputSchema);
+          toolContext.actions().setSetModelResponse(validatedResponse);
+          return validatedResponse;
         });
   }
+
+  private static Map<String, Object> excludeNullFields(Map<String, Object> values, Schema schema) {
+    Map<String, Schema> properties = schema.properties().get();
+    Map<String, Object> result = new LinkedHashMap<>();
+    for (Map.Entry<String, Object> entry : values.entrySet()) {
+      Object value = entry.getValue();
+      if (value != null) {
+        result.put(entry.getKey(), excludeNullFields(value, properties.get(entry.getKey())));
+      }
+    }
+    return result;
+  }
+
+  @SuppressWarnings("unchecked")
+  private static Object excludeNullFields(Object value, Schema schema) {
+    Type.Known type = schema.type().get().knownEnum();
+    if (type == Type.Known.OBJECT) {
+      return excludeNullFields((Map<String, Object>) value, schema);
+    }
+    if (type == Type.Known.ARRAY) {
+      Schema itemSchema = schema.items().get();
+      List<Object> result = new ArrayList<>();
+      for (Object item : (List<?>) value) {
+        result.add(item == null ? null : excludeNullFields(item, itemSchema));
+      }
+      return result;
+    }
+    return value;
+  }
+
+  private static String sanitizeValidationMessage(String message) {
+    if (message == null) {
+      return "Arguments do not match the output schema.";
+    }
+    // The model already knows the schema from the tool declaration, so the appended schema dump is
+    // redundant in feedback.
+    int schemaDumpIndex = message.indexOf(OUTPUT_SCHEMA_DUMP_MARKER);
+    if (schemaDumpIndex >= 0) {
+      message = message.substring(0, schemaDumpIndex) + " does not match agent output schema.";
+    }
+    return message;
+  }
 }

core/src/test/java/com/google/adk/events/EventActionsTest.java

@@ -89,6 +89,7 @@ public void merge_mergesAllFields() {
             .requestedToolConfirmations(
                 new ConcurrentHashMap<>(ImmutableMap.of("tool2", TOOL_CONFIRMATION)))
             .endOfAgent(true)
+            .setModelResponse(ImmutableMap.of("field1", "value1"))
             .build();
 
     EventActions merged = eventActions1.toBuilder().merge(eventActions2).build();
@@ -109,6 +110,7 @@ public void merge_mergesAllFields() {
         .containsExactly("tool1", TOOL_CONFIRMATION, "tool2", TOOL_CONFIRMATION);
     assertThat(merged.endOfAgent()).isTrue();
     assertThat(merged.compaction()).hasValue(COMPACTION);
+    assertThat(merged.setModelResponse()).hasValue(ImmutableMap.of("field1", "value1"));
   }
 
   @Test
@@ -177,13 +179,15 @@ public void jsonSerialization_works() throws Exception {
         EventActions.builder()
             .deletedArtifactIds(ImmutableSet.of("d1", "d2"))
             .stateDelta(new ConcurrentHashMap<>(ImmutableMap.of("k", "v")))
+            .setModelResponse(ImmutableMap.of("field1", "value1"))
             .build();
 
     String json = eventActions.toJson();
     EventActions deserialized = EventActions.fromJsonString(json, EventActions.class);
 
     assertThat(deserialized).isEqualTo(eventActions);
     assertThat(deserialized.deletedArtifactIds()).containsExactly("d1", "d2");
+    assertThat(deserialized.setModelResponse()).hasValue(ImmutableMap.of("field1", "value1"));
   }
 
   @Test