From ad4ab50c741d7ec1a45d6cbf2d60647c06422c0e Mon Sep 17 00:00:00 2001 From: Mateusz Krawiec Date: Fri, 13 Nov 2026 07:24:36 -0800 Subject: [PATCH] fix: validate free-form schemas instead of crashing PiperOrigin-RevId: 1021017109 --- .../main/java/com/google/adk/SchemaUtils.java | 62 +++--- .../adk/tools/SetModelResponseTool.java | 8 +- .../java/com/google/adk/SchemaUtilsTest.java | 184 ++++++++++++++++++ .../com/google/adk/agents/LlmAgentTest.java | 37 ++++ .../adk/flows/llmflows/OutputSchemaTest.java | 54 +++++ .../adk/tools/SetModelResponseToolTest.java | 19 ++ 6 files changed, 335 insertions(+), 29 deletions(-) diff --git a/core/src/main/java/com/google/adk/SchemaUtils.java b/core/src/main/java/com/google/adk/SchemaUtils.java index 0df7ee775..9ba1a3588 100644 --- a/core/src/main/java/com/google/adk/SchemaUtils.java +++ b/core/src/main/java/com/google/adk/SchemaUtils.java @@ -41,11 +41,23 @@ private SchemaUtils() {} // Private constructor for utility class */ @SuppressWarnings("unchecked") // For tool parameter type casting. private static Boolean matchType(Object value, Schema schema, Boolean isInput) { + // An untyped schema with properties or required keys is an object, one with items is an array, + // and one with none of them accepts any value, null included. + Type.Known type = + schema + .type() + .map(Type::knownEnum) + .orElse( + schema.properties().isPresent() || schema.required().isPresent() + ? Type.Known.OBJECT + : schema.items().isPresent() ? Type.Known.ARRAY : null); if (value == null) { - return schema.nullable().orElse(false); + return type == null || schema.nullable().orElse(false); + } + if (type == null) { + return true; } // Based on types from https://cloud.google.com/vertex-ai/docs/reference/rest/v1/Schema - Type.Known type = schema.type().get().knownEnum(); switch (type) { case STRING: return value instanceof String; @@ -57,6 +69,9 @@ private static Boolean matchType(Object value, Schema schema, Boolean isInput) { return value instanceof Number; case ARRAY: if (value instanceof List) { + if (schema.items().isEmpty()) { + return true; + } for (Object element : (List) value) { if (!matchType(element, schema.items().get(), isInput)) { return false; @@ -91,29 +106,14 @@ private static Boolean matchType(Object value, Schema schema, Boolean isInput) { */ public static void validateMapOnSchema(Map args, Schema schema, Boolean isInput) { Preconditions.checkNotNull(isInput, "IsInput cannot be null"); - Map properties = schema.properties().get(); - for (Entry arg : args.entrySet()) { - // Check if the argument is in the schema. - if (!properties.containsKey(arg.getKey())) { - if (isInput) { - throw new IllegalArgumentException( - "Input arg: " + arg.getKey() + " does not match agent input schema: " + schema); - } else { - throw new IllegalArgumentException( - "Output arg: " + arg.getKey() + " does not match agent output schema: " + schema); - } - } - // Check if the argument type matches the schema type. - if (!matchType(arg.getValue(), properties.get(arg.getKey()), isInput)) { - if (isInput) { - throw new IllegalArgumentException( - "Input arg: " + arg.getKey() + " does not match agent input schema: " + schema); - } else { - throw new IllegalArgumentException( - "Output arg: " + arg.getKey() + " does not match agent output schema: " + schema); - } - } + if (schema.type().map(Type::knownEnum).orElse(Type.Known.OBJECT) != Type.Known.OBJECT) { + throw new IllegalArgumentException( + (isInput ? "Input" : "Output") + " schema does not describe an object: " + schema); } + // A schema without properties leaves the keys open, as for a free-form map. + schema + .properties() + .ifPresent(properties -> validateProperties(args, properties, schema, isInput)); // Check if all required arguments are present. if (schema.required().isPresent()) { for (String required : schema.required().get()) { @@ -128,6 +128,20 @@ public static void validateMapOnSchema(Map args, Schema schema, } } + /** Checks that every key in {@code args} is one of {@code properties} and matches its schema. */ + private static void validateProperties( + Map args, Map properties, Schema schema, boolean isInput) { + for (Entry arg : args.entrySet()) { + if (!properties.containsKey(arg.getKey()) + || !matchType(arg.getValue(), properties.get(arg.getKey()), isInput)) { + throw new IllegalArgumentException( + isInput + ? "Input arg: " + arg.getKey() + " does not match agent input schema: " + schema + : "Output arg: " + arg.getKey() + " does not match agent output schema: " + schema); + } + } + } + /** * Validates an output string against a schema. * diff --git a/core/src/main/java/com/google/adk/tools/SetModelResponseTool.java b/core/src/main/java/com/google/adk/tools/SetModelResponseTool.java index 9cab1749a..af5368859 100644 --- a/core/src/main/java/com/google/adk/tools/SetModelResponseTool.java +++ b/core/src/main/java/com/google/adk/tools/SetModelResponseTool.java @@ -87,8 +87,8 @@ public Single> runAsync(Map args, ToolContex /** * Drops nulls like Python's {@code model_dump(exclude_none=True)}, but keeps the null of a - * property that {@code schema} marks required and nullable, which the schema demands. A value - * without a schema loses all its nulls. + * property that {@code schema} marks required, which validation already accepted. A value without + * a schema loses all its nulls. */ private static Map excludeNullFields( Map values, @Nullable Schema schema) { @@ -102,9 +102,7 @@ private static Map excludeNullFields( Object value = entry.getValue(); if (value != null) { result.put(entry.getKey(), excludeNullsFromValue(value, property)); - } else if (property != null - && property.nullable().orElse(false) - && required.contains(entry.getKey())) { + } else if (required.contains(entry.getKey())) { result.put(entry.getKey(), null); } } diff --git a/core/src/test/java/com/google/adk/SchemaUtilsTest.java b/core/src/test/java/com/google/adk/SchemaUtilsTest.java index ac5df39a8..7e91eba17 100644 --- a/core/src/test/java/com/google/adk/SchemaUtilsTest.java +++ b/core/src/test/java/com/google/adk/SchemaUtilsTest.java @@ -18,6 +18,7 @@ import static org.junit.Assert.assertThrows; +import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.genai.types.Schema; import java.util.HashMap; @@ -98,4 +99,187 @@ public void validateMapOnSchema_integerField_allowsIntegerAndLong() { args.put("intField", 123); SchemaUtils.validateMapOnSchema(args, schema, /* isInput= */ true); } + + @Test + public void validateMapOnSchema_untypedProperty_acceptsAnyValue() { + Schema schema = + Schema.builder() + .type("OBJECT") + .properties(ImmutableMap.of("anything", Schema.builder().build())) + .build(); + + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("anything", ImmutableMap.of("nested", 1)), schema, /* isInput= */ false); + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("anything", "text"), schema, /* isInput= */ false); + } + + @Test + public void validateMapOnSchema_objectWithoutProperties_acceptsAnyKeys() { + Schema schema = + Schema.builder() + .type("OBJECT") + .properties(ImmutableMap.of("metadata", Schema.builder().type("OBJECT").build())) + .build(); + + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("metadata", ImmutableMap.of("source", "web", "score", 0.5)), + schema, + /* isInput= */ false); + } + + @Test + public void validateMapOnSchema_topLevelObjectWithoutProperties_stillEnforcesRequired() { + Schema schema = Schema.builder().type("OBJECT").required(ImmutableList.of("id")).build(); + + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("id", "1", "extra", true), schema, /* isInput= */ true); + assertThrows( + IllegalArgumentException.class, + () -> + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("extra", true), schema, /* isInput= */ true)); + } + + @Test + public void validateMapOnSchema_arrayWithoutItems_acceptsAnyElements() { + Schema schema = + Schema.builder() + .type("OBJECT") + .properties(ImmutableMap.of("tags", Schema.builder().type("ARRAY").build())) + .build(); + + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("tags", ImmutableList.of("a", 1, ImmutableMap.of())), + schema, + /* isInput= */ false); + } + + @Test + public void validateMapOnSchema_arrayWithoutItems_stillRejectsANonList() { + Schema schema = + Schema.builder() + .type("OBJECT") + .properties(ImmutableMap.of("tags", Schema.builder().type("ARRAY").build())) + .build(); + + assertThrows( + IllegalArgumentException.class, + () -> + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("tags", "a"), schema, /* isInput= */ false)); + } + + @Test + public void validateMapOnSchema_nonObjectSchema_throwsException() { + Schema schema = Schema.builder().type("ARRAY").build(); + + assertThrows( + IllegalArgumentException.class, + () -> + SchemaUtils.validateMapOnSchema(ImmutableMap.of("a", 1), schema, /* isInput= */ false)); + } + + @Test + public void validateMapOnSchema_untypedProperty_acceptsNull() { + Schema schema = + Schema.builder() + .type("OBJECT") + .properties(ImmutableMap.of("anything", Schema.builder().build())) + .build(); + Map args = new HashMap<>(); + args.put("anything", null); + + SchemaUtils.validateMapOnSchema(args, schema, /* isInput= */ false); + } + + @Test + public void validateMapOnSchema_untypedNonNullableProperty_acceptsNull() { + Schema schema = + Schema.builder() + .type("OBJECT") + .properties(ImmutableMap.of("anything", Schema.builder().nullable(false).build())) + .build(); + Map args = new HashMap<>(); + args.put("anything", null); + + SchemaUtils.validateMapOnSchema(args, schema, /* isInput= */ false); + } + + @Test + public void validateMapOnSchema_untypedPropertyWithRequired_isAnObject() { + Schema schema = + Schema.builder() + .type("OBJECT") + .properties( + ImmutableMap.of("item", Schema.builder().required(ImmutableList.of("id")).build())) + .build(); + + assertThrows( + IllegalArgumentException.class, + () -> + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("item", 5), schema, /* isInput= */ false)); + assertThrows( + IllegalArgumentException.class, + () -> + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("item", ImmutableMap.of()), schema, /* isInput= */ false)); + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("item", ImmutableMap.of("id", "x")), schema, /* isInput= */ false); + } + + @Test + public void validateMapOnSchema_untypedPropertyWithItems_isAnArray() { + Schema schema = + Schema.builder() + .type("OBJECT") + .properties( + ImmutableMap.of( + "tags", + Schema.builder().items(Schema.builder().type("STRING").build()).build())) + .build(); + + assertThrows( + IllegalArgumentException.class, + () -> + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("tags", 5), schema, /* isInput= */ false)); + assertThrows( + IllegalArgumentException.class, + () -> + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("tags", ImmutableList.of(1)), schema, /* isInput= */ false)); + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("tags", ImmutableList.of("a")), schema, /* isInput= */ false); + } + + @Test + public void validateMapOnSchema_untypedNestedSchemaWithProperties_validatesAsObject() { + Schema nested = + Schema.builder() + .properties(ImmutableMap.of("a", Schema.builder().type("STRING").build())) + .required(ImmutableList.of("a")) + .build(); + Schema schema = + Schema.builder().type("OBJECT").properties(ImmutableMap.of("nested", nested)).build(); + + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("nested", ImmutableMap.of("a", "x")), schema, /* isInput= */ false); + assertThrows( + IllegalArgumentException.class, + () -> + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("nested", 5), schema, /* isInput= */ false)); + assertThrows( + IllegalArgumentException.class, + () -> + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("nested", ImmutableMap.of()), schema, /* isInput= */ false)); + Map nullNested = new HashMap<>(); + nullNested.put("nested", null); + assertThrows( + IllegalArgumentException.class, + () -> SchemaUtils.validateMapOnSchema(nullNested, schema, /* isInput= */ false)); + } } diff --git a/core/src/test/java/com/google/adk/agents/LlmAgentTest.java b/core/src/test/java/com/google/adk/agents/LlmAgentTest.java index c32eafc4d..2f10651ba 100644 --- a/core/src/test/java/com/google/adk/agents/LlmAgentTest.java +++ b/core/src/test/java/com/google/adk/agents/LlmAgentTest.java @@ -71,7 +71,9 @@ import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.Maybe; import io.reactivex.rxjava3.core.Single; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicBoolean; import org.junit.After; @@ -142,6 +144,41 @@ public void testRun_withOutputKey_savesState() { assertThat(events.get(0).actions().stateDelta()).containsEntry("myOutput", "Saved output"); } + @Test + public void testRun_withOutputKeyAndFreeFormOutputSchema_savesParsedState() { + Schema outputSchema = + Schema.builder() + .type("OBJECT") + .properties( + ImmutableMap.of( + "city", + Schema.builder().type("STRING").build(), + "details", + Schema.builder().type("OBJECT").build(), + "tags", + Schema.builder().type("ARRAY").build(), + "extra", + Schema.builder().build())) + .build(); + TestLlm testLlm = + createTestLlm( + createTextLlmResponse( + "{\"city\":\"Paris\",\"details\":{\"population\":2100000}," + + "\"tags\":[\"a\",1],\"extra\":null}")); + LlmAgent agent = + createTestAgentBuilder(testLlm).outputSchema(outputSchema).outputKey("answer").build(); + InvocationContext invocationContext = createInvocationContext(agent); + + List events = agent.runAsync(invocationContext).toList().blockingGet(); + + Map expected = new HashMap<>(); + expected.put("city", "Paris"); + expected.put("details", ImmutableMap.of("population", 2100000)); + expected.put("tags", ImmutableList.of("a", 1)); + expected.put("extra", null); + assertThat(getOnlyElement(events).actions().stateDelta()).containsEntry("answer", expected); + } + @Test public void testRun_withOutputKey_savesMultiPartState() { Content modelContent = Content.fromParts(Part.fromText("Part 1."), Part.fromText(" Part 2.")); diff --git a/core/src/test/java/com/google/adk/flows/llmflows/OutputSchemaTest.java b/core/src/test/java/com/google/adk/flows/llmflows/OutputSchemaTest.java index 8b9a55263..c08946d22 100644 --- a/core/src/test/java/com/google/adk/flows/llmflows/OutputSchemaTest.java +++ b/core/src/test/java/com/google/adk/flows/llmflows/OutputSchemaTest.java @@ -24,6 +24,7 @@ import com.google.adk.agents.InvocationContext; import com.google.adk.agents.LlmAgent; +import com.google.adk.agents.RunConfig; import com.google.adk.events.Event; import com.google.adk.events.EventActions; import com.google.adk.flows.llmflows.RequestProcessor.RequestProcessingResult; @@ -613,4 +614,57 @@ private LlmRequest basicThenOutputSchema(BaseLlm model) { .blockingGet() .updatedRequest(); } + + @Test + public void setModelResponseWithARequiredUntypedNull_keepsItInTheOutputKey() { + Schema schema = + Schema.builder() + .type("OBJECT") + .properties(ImmutableMap.of("extra", Schema.builder().build())) + .required(ImmutableList.of("extra")) + .build(); + Map args = new LinkedHashMap<>(); + args.put("extra", null); + TestLlm scriptedLlm = + createTestLlm( + createLlmResponse( + Content.fromParts(Part.fromFunctionCall(SetModelResponseTool.NAME, args)))); + BaseLlm gemini2NamedLlm = + new BaseLlm("gemini-2.0-flash") { + @Override + public Flowable generateContent(LlmRequest llmRequest, boolean stream) { + return scriptedLlm.generateContent(llmRequest, stream); + } + + @Override + public BaseLlmConnection connect(LlmRequest llmRequest) { + return scriptedLlm.connect(llmRequest); + } + }; + LlmAgent agent = + LlmAgent.builder() + .name("agent") + .model(gemini2NamedLlm) + .outputSchema(schema) + .tools(ImmutableList.of(new TestTool())) + .outputKey("answer") + .build(); + InMemoryRunner runner = new InMemoryRunner(agent, "test-app"); + Session session = runner.sessionService().createSession("test-app", "user").blockingGet(); + + List events = + runner + .runAsync( + "user", + session.id(), + Content.fromParts(Part.fromText("hello")), + // A rejected answer would make the model retry; stop instead of looping. + RunConfig.builder().maxLlmCalls(2).build()) + .toList() + .blockingGet(); + + Event finalEvent = events.get(events.size() - 1); + assertThat(finalEvent.content().get().parts().get().get(0).text()).hasValue("{\"extra\":null}"); + assertThat(finalEvent.actions().stateDelta()).containsEntry("answer", args); + } } diff --git a/core/src/test/java/com/google/adk/tools/SetModelResponseToolTest.java b/core/src/test/java/com/google/adk/tools/SetModelResponseToolTest.java index 959c1c6dc..9335fe475 100644 --- a/core/src/test/java/com/google/adk/tools/SetModelResponseToolTest.java +++ b/core/src/test/java/com/google/adk/tools/SetModelResponseToolTest.java @@ -151,6 +151,25 @@ public void runAsync_requiredNullableNull_keepsTheNull() { SchemaUtils.validateMapOnSchema(result, outputSchema, /* isInput= */ false); } + @Test + public void runAsync_requiredUntypedNull_keepsTheNull() { + Schema outputSchema = + Schema.builder() + .type("OBJECT") + .properties(ImmutableMap.of("extra", Schema.builder().build())) + .required(ImmutableList.of("extra")) + .build(); + Map args = new LinkedHashMap<>(); + args.put("extra", null); + + ToolContext toolContext = createToolContext(); + Map result = + new SetModelResponseTool(outputSchema).runAsync(args, toolContext).blockingGet(); + + assertThat(result).isEqualTo(args); + assertThat(toolContext.actions().setModelResponse()).hasValue(args); + } + @Test public void runAsync_requiredAndOptionalNullableNulls_keepsOnlyTheRequiredOne() { Schema outputSchema =