diff --git a/core/src/main/java/com/google/adk/workflow/FunctionNode.java b/core/src/main/java/com/google/adk/workflow/FunctionNode.java index ef14143cd..e26a84003 100644 --- a/core/src/main/java/com/google/adk/workflow/FunctionNode.java +++ b/core/src/main/java/com/google/adk/workflow/FunctionNode.java @@ -81,8 +81,11 @@ public Builder toBuilder() { /** * Returns the schema that describes the input port: the one set on the builder or, if none, one - * inferred from the input type. An {@code Object}, {@code Void} or other type with no fixed JSON - * shape infers none. + * inferred from the input type, where an {@code Object}, {@code Void} or other type with no fixed + * JSON shape infers none. An edge into an inferred record or POJO fails at build if the + * producer's output schema lists a key the type does not, while a builder schema checks only the + * keys it lists and requires. On a {@code Content} port a builder schema describes the message + * text, and an edge is checked against the {@code Content} schema instead. */ @Override public Optional inputSchema() { @@ -104,6 +107,21 @@ boolean inputSchemaInferred() { return super.inputSchema().isEmpty() && inputCodec.schema() != null; } + /** Returns the schema an edge into this node must fit, the inferred one on a Content port. */ + Optional edgeInputSchema() { + return inputCodec.takesContent() ? Optional.ofNullable(inputCodec.schema()) : inputSchema(); + } + + /** Whether a builder schema would let this port take keys its type does not list. */ + boolean acceptsExtraKeysWithASchema() { + return inputSchemaInferred() && !inputCodec.takesContent(); + } + + /** Whether {@link #edgeInputSchema} was inferred, so its objects with properties are closed. */ + boolean edgeInputSchemaInferred() { + return inputCodec.takesContent() || inputSchemaInferred(); + } + @Override public Flowable runNode(Context context, @Nullable Object nodeInput) { return Flowable.defer( @@ -227,7 +245,10 @@ public Builder config(NodeConfig config) { return this; } - /** Sets the schema that describes the node's input, in place of one inferred from its type. */ + /** + * Sets the schema that describes the node's input, in place of one inferred from its type. On a + * {@code Content} port, the schema describes the message's text. + */ @CanIgnoreReturnValue public Builder inputSchema(@Nullable Schema inputSchema) { this.inputSchema = inputSchema; diff --git a/core/src/main/java/com/google/adk/workflow/Graph.java b/core/src/main/java/com/google/adk/workflow/Graph.java index 270bd62e1..1f91fcf04 100644 --- a/core/src/main/java/com/google/adk/workflow/Graph.java +++ b/core/src/main/java/com/google/adk/workflow/Graph.java @@ -24,8 +24,10 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Sets; +import com.google.genai.types.Schema; import java.util.ArrayList; import java.util.List; +import java.util.Optional; import java.util.Set; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -80,6 +82,7 @@ static Graph of(List edges) { NodeNames.validate(node.name()); } GraphValidation.validate(nodes, edgeList); + checkEdgeSchemas(edgeList); ImmutableMap nodesByName = nodes.stream().collect(toImmutableMap(Node::name, node -> node)); ImmutableSet sourceNames = @@ -92,6 +95,41 @@ static Graph of(List edges) { return new Graph(nodes, edgeList, nodesByName, terminalNodeNames); } + /** + * Throws unless the output schema of each edge's source fits the input schema of its target, when + * both are set. Each object that lists properties in a schema the target inferred is closed. + */ + private static void checkEdgeSchemas(List edges) { + for (Edge edge : edges) { + FunctionNode function = edge.to() instanceof FunctionNode node ? node : null; + Optional output = edge.from().outputSchema(); + Optional input = + function != null ? function.edgeInputSchema() : edge.to().inputSchema(); + if (output.isEmpty() || input.isEmpty()) { + continue; + } + boolean closed = function != null && function.edgeInputSchemaInferred(); + if (!SchemaFit.fits(output.get(), input.get(), closed)) { + throw new WorkflowConfigurationException( + "Schema mismatch on edge " + + edge.from().name() + + " -> " + + edge.to().name() + + ". Output schema " + + output.get().toJson() + + " does not fit input schema " + + input.get().toJson() + + "." + + (function != null + && function.acceptsExtraKeysWithASchema() + && SchemaFit.fits(output.get(), input.get(), /* closed= */ false) + ? " It was inferred from the port's type and takes no other keys; set" + + " inputSchema on the builder or read the objects as Maps to accept them." + : "")); + } + } + } + ImmutableList nodes() { return nodes; } diff --git a/core/src/main/java/com/google/adk/workflow/Node.java b/core/src/main/java/com/google/adk/workflow/Node.java index b81fdfd47..752df2f55 100644 --- a/core/src/main/java/com/google/adk/workflow/Node.java +++ b/core/src/main/java/com/google/adk/workflow/Node.java @@ -62,7 +62,13 @@ default NodeConfig config() { return NodeConfig.builder().build(); } - /** Returns the schema the node's input must match. */ + /** + * Returns the schema the node's input must match. An edge into the node fails at build when its + * source's output schema lists a key or type this one rejects, so an open producer such as a map + * builds, and a {@link FunctionNode} on a {@code Content} port checks the edge against the {@code + * Content} schema. A node that requires all predecessors applies it to each predecessor's output, + * not the joined map. + */ @Experimental default Optional inputSchema() { return Optional.empty(); diff --git a/core/src/main/java/com/google/adk/workflow/PortCodec.java b/core/src/main/java/com/google/adk/workflow/PortCodec.java index 57e08a367..ff75efad9 100644 --- a/core/src/main/java/com/google/adk/workflow/PortCodec.java +++ b/core/src/main/java/com/google/adk/workflow/PortCodec.java @@ -499,6 +499,11 @@ private static String capitalized(String what) { return schema.get(); } + /** Whether this port reads a {@code Content}. */ + boolean takesContent() { + return type.getRawClass() == Content.class; + } + /** * Returns {@code value} read as this port's type; {@code what} names the port in errors. * diff --git a/core/src/main/java/com/google/adk/workflow/SchemaFit.java b/core/src/main/java/com/google/adk/workflow/SchemaFit.java new file mode 100644 index 000000000..a351948ed --- /dev/null +++ b/core/src/main/java/com/google/adk/workflow/SchemaFit.java @@ -0,0 +1,200 @@ +/* + * Copyright 2026 Google LLC + * + * 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 + * + * http://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 com.google.adk.workflow; + +import static com.google.common.collect.ImmutableList.toImmutableList; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.genai.types.Schema; +import com.google.genai.types.Type; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import org.jspecify.annotations.Nullable; + +/** + * Decides at build time whether an output schema fits an input schema by types, enums, keys and + * nullability; it ignores formats, patterns and bounds. An output with no type or enum fits any + * input that accepts its nulls, and an object output that lists no keys fits any object input. + */ +final class SchemaFit { + + /** + * Whether {@code source} fits {@code target} by types, enums, keys and nullability. With {@code + * closed}, each object in {@code target} that lists properties takes no other key, as in a schema + * inferred from a type. + */ + static boolean fits(Schema source, Schema target, boolean closed) { + if (source.equals(target)) { + return true; + } + if (source.anyOf().isPresent()) { + return unionFits(source, source.anyOf().get(), target, closed); + } + Type.Known type = declaredType(source); + if (isNullable(source) || type == Type.Known.NULL) { + return acceptsNull(target) + && (type == Type.Known.NULL + || fits(source.toBuilder().clearNullable().build(), target, closed)); + } + if (type == null && source.enum_().isPresent()) { + // An enum lists strings, as a typed one does. + return fits(source.toBuilder().type(Type.Known.STRING).build(), target, closed); + } + if (target.anyOf().isPresent()) { + return fits(source, target.toBuilder().clearAnyOf().build(), closed) + && target.anyOf().get().stream().anyMatch(branch -> fits(source, branch, closed)); + } + if (type == null) { + return true; + } + if (target.enum_().isPresent() + && !(source.enum_().isPresent() + && target.enum_().get().containsAll(source.enum_().get()))) { + return false; + } + Type.Known targetType = declaredType(target); + if (targetType != null + && targetType != type + && !(type == Type.Known.INTEGER && targetType == Type.Known.NUMBER)) { + return false; + } + return switch (type) { + case ARRAY -> + target.items().isEmpty() + || source.items().isEmpty() + || fits(source.items().get(), target.items().get(), closed); + case OBJECT -> objectFits(source, target, closed); + default -> true; + }; + } + + /** Whether a union fits: its own part and each of its {@code branches} must. */ + private static boolean unionFits( + Schema source, List branches, Schema target, boolean closed) { + boolean takesNull = acceptsNull(source); + if (takesNull && !acceptsNull(target)) { + return false; + } + // A union that rejects null admits only its branches' other values. + List values = + takesNull + ? branches + : branches.stream() + .filter(branch -> declaredType(branch) != Type.Known.NULL) + .map(branch -> branch.toBuilder().clearNullable().build()) + .collect(toImmutableList()); + // An untyped union's own keys, items or enum take the one type its branches share, if any. + List types = + values.stream() + .filter(value -> declaredType(value) != Type.Known.NULL) + .map(SchemaFit::impliedType) + .distinct() + .collect(Collectors.toCollection(ArrayList::new)); + Type.Known type = declaredType(source); + Schema.Builder own = source.toBuilder().clearAnyOf(); + if (type == null + && (impliedType(source) != null || source.enum_().isPresent()) + && types.size() == 1 + && types.get(0) != null) { + own.type(types.get(0)); + } + // An untyped branch takes its union's type. + return fits(own.build(), target, closed) + && values.stream() + .allMatch( + value -> + fits( + type != null && declaredType(value) == null + ? value.toBuilder().type(type).build() + : value, + target, + closed)); + } + + /** Whether {@code source}, an object, fits {@code target}'s keys. */ + private static boolean objectFits(Schema source, Schema target, boolean closed) { + // An object that constrains no keys, such as a map, passes: the build cannot know its keys. + if (!declaresKeys(source)) { + return true; + } + Map given = source.properties().orElse(ImmutableMap.of()); + Map wanted = target.properties().orElse(ImmutableMap.of()); + List required = source.required().orElse(ImmutableList.of()); + List targetRequired = target.required().orElse(ImmutableList.of()); + // The source must require each key the target requires, and declare it if the target does. + if (!required.containsAll(targetRequired) + || targetRequired.stream() + .anyMatch(key -> wanted.containsKey(key) && !given.containsKey(key))) { + return false; + } + boolean closedTarget = closed && target.properties().isPresent(); + if (closedTarget && !wanted.keySet().containsAll(required)) { + return false; + } + return given.entrySet().stream() + .allMatch( + property -> + wanted.containsKey(property.getKey()) + ? fits(property.getValue(), wanted.get(property.getKey()), closed) + : !closedTarget); + } + + /** + * Whether {@code schema} takes null: if nullable, typed NULL, or untyped with no enum, unless it + * has anyOf branches and none of them takes null. + */ + private static boolean acceptsNull(Schema schema) { + return isNullable(schema) + || declaredType(schema) == Type.Known.NULL + || (declaredType(schema) == null + && schema.enum_().isEmpty() + && schema + .anyOf() + .map(branches -> branches.stream().anyMatch(SchemaFit::acceptsNull)) + .orElse(true)); + } + + private static boolean isNullable(Schema schema) { + return schema.nullable().orElse(false); + } + + private static boolean declaresKeys(Schema schema) { + return schema.properties().isPresent() || schema.required().isPresent(); + } + + /** Returns the type {@code schema} declares, or else the shape its keys or items imply. */ + private static Type.@Nullable Known impliedType(Schema schema) { + Type.Known type = declaredType(schema); + if (type != null) { + return type; + } + if (declaresKeys(schema)) { + return Type.Known.OBJECT; + } + return schema.items().isPresent() ? Type.Known.ARRAY : null; + } + + private static Type.@Nullable Known declaredType(Schema schema) { + Type.Known type = schema.type().map(Type::knownEnum).orElse(null); + return type == Type.Known.TYPE_UNSPECIFIED ? null : type; + } + + private SchemaFit() {} +} diff --git a/core/src/main/java/com/google/adk/workflow/Workflow.java b/core/src/main/java/com/google/adk/workflow/Workflow.java index a2a8fdcdb..4484a0df5 100644 --- a/core/src/main/java/com/google/adk/workflow/Workflow.java +++ b/core/src/main/java/com/google/adk/workflow/Workflow.java @@ -210,6 +210,8 @@ public Builder stateSchema(@Nullable Schema stateSchema) { * empty or contains {@code /}, {@code @}, or {@code .} * @throws GraphValidationException if the graph is malformed, for example if {@link Start} * cannot reach one of its nodes + * @throws WorkflowConfigurationException if the output schema of an edge's source admits a + * value that the input schema of its target does not */ public Workflow build() { checkState(name != null, "Workflow.Builder requires name to be set."); diff --git a/core/src/test/java/com/google/adk/workflow/FunctionNodeTest.java b/core/src/test/java/com/google/adk/workflow/FunctionNodeTest.java index 3d82ceb6f..805c00bdb 100644 --- a/core/src/test/java/com/google/adk/workflow/FunctionNodeTest.java +++ b/core/src/test/java/com/google/adk/workflow/FunctionNodeTest.java @@ -491,6 +491,82 @@ public String note() { } } + record FullOrder(String id, double total, String note) {} + + record BasicOrder(String id, double total) {} + + record OrderWithOptionalNote(String id, double total, @Nullable String note) {} + + /** A record that writes a total it derives, and so does not read. */ + record Invoice(String id, int quantity) { + public int getTotal() { + return quantity * 10; + } + } + + record Creds(String user, String password) {} + + /** A class with a property it only reads, by a setter. */ + static final class Notepad { + public String id; + + public void setNote(String note) {} + } + + record FlatOut(String label, String id, int quantity) {} + + record Qty(int qty) {} + + @JsonInclude(JsonInclude.Include.NON_DEFAULT) + record Count(int count) {} + + @JsonInclude(JsonInclude.Include.NON_EMPTY) + record Tags(String text, List tags) {} + + /** A record whose nested record NON_DEFAULT still writes. */ + @JsonInclude(JsonInclude.Include.NON_DEFAULT) + record Packed(int count, Order order) {} + + /** A record whose property-level inclusion overrides its class's. */ + @JsonInclude(JsonInclude.Include.NON_EMPTY) + record Overridden(@JsonInclude(JsonInclude.Include.ALWAYS) String text, String note) {} + + /** A class whose NON_DEFAULT omits any value equal to its field's initializer. */ + @JsonInclude(JsonInclude.Include.NON_DEFAULT) + static final class Defaults { + public Status status = Status.OPEN; + public Order order = new Order("a1", 1); + } + + record View(Status status, Order order) {} + + @JsonInclude(JsonInclude.Include.NON_EMPTY) + record Labeled(@JsonProperty(required = true) String label, int n) {} + + /** A class that writes its order unwrapped and has two setters for it. */ + static final class OverFlat { + public String label; + + @JsonUnwrapped + public Order getOrder() { + return new Order("a1", 1); + } + + public void setOrder(Order order) {} + + public void setOrder(Map order) {} + } + + /** A class that reads its unwrapped order through a setter. */ + static final class SetterFlat { + public String label; + + @JsonUnwrapped + public void setOrder(Order order) {} + } + + record Authored(String author) {} + record TreeNode(String label, List children) {} /** Holds a declaration annotation named like the {@code @Nullable} of other packages. */ @@ -943,8 +1019,9 @@ public void of_bodySettingStateAndReturningAnEvent_mergesBothStateDeltas() { @Test public void of_inputThatDoesNotFitThePort_failsTheNodeWithoutQuotingTheValue() { - FunctionNode secret = - FunctionNodes.of("secret", Content.class, String.class, (ctx, message) -> "s3cr3t"); + // Arrange: an Object output has no schema, so the build leaves this edge to run time. + FunctionNode secret = + FunctionNodes.of("secret", Content.class, Object.class, (ctx, message) -> "s3cr3t"); FunctionNode describe = FunctionNodes.of("describe", Order.class, String.class, (ctx, order) -> order.id()); @@ -957,8 +1034,9 @@ public void of_inputThatDoesNotFitThePort_failsTheNodeWithoutQuotingTheValue() { @Test public void of_valueThatIsNotAContent_failsAContentPortWithoutQuotingTheValue() { - FunctionNode secret = - FunctionNodes.of("secret", Content.class, String.class, (ctx, message) -> "s3cr3t"); + // Arrange: an Object output has no schema, so the build leaves this edge to run time. + FunctionNode secret = + FunctionNodes.of("secret", Content.class, Object.class, (ctx, message) -> "s3cr3t"); FunctionNode read = FunctionNodes.of("read", Content.class, String.class, (ctx, content) -> content.text()); @@ -1364,6 +1442,431 @@ public void of_eventAndObjectPorts_inferASchemaOnlyForAnEventInput() { assertThat(object).isEmpty(); } + @Test + public void build_recordIntoAClosedRecordThatLacksOneOfItsKeys_failsNamingBothSchemasAndTheFix() { + // Arrange: the inferred BasicOrder schema takes no note. + FunctionNode full = fullOrder(); + FunctionNode basic = + FunctionNodes.of("basic", BasicOrder.class, String.class, (ctx, order) -> order.id()); + + // Act + WorkflowConfigurationException error = + assertThrows(WorkflowConfigurationException.class, () -> chain(full, basic)); + + // Assert + assertThat(error) + .hasMessageThat() + .isEqualTo( + "Schema mismatch on edge full -> basic. Output schema " + + full.outputSchema().get().toJson() + + " does not fit input schema " + + basic.inputSchema().get().toJson() + + ". It was inferred from the port's type and takes no other keys; set inputSchema" + + " on the builder or read the objects as Maps to accept them."); + } + + @Test + public void build_recordIntoAClosedRecordWithAnOptionalKeyItLacks_buildsAndRuns() { + // Arrange: the inferred OrderWithOptionalNote schema lists a note it does not require. + FunctionNode basic = + FunctionNodes.of( + "basic", Content.class, BasicOrder.class, (ctx, m) -> new BasicOrder("a1", 9.5)); + FunctionNode noted = + FunctionNodes.of( + "noted", + OrderWithOptionalNote.class, + String.class, + (ctx, order) -> order.id() + ":" + order.note()); + + // Act + RunResult result = run(chain(basic, noted)); + + // Assert + result.assertCompleted(); + assertThat(result.outputAt("wf@1/noted@1")).isEqualTo("a1:null"); + } + + @Test + public void build_listOfRecordsIntoAListOfClosedRecords_hintsAtMapsThatThenBuild() { + // Arrange: the inferred BasicOrder items take no note; map items take any key. + FunctionNode> fulls = fullOrders(); + FunctionNode, Integer> basics = + FunctionNodes.of( + "basics", + new TypeReference>() {}, + Integer.class, + (ctx, l) -> l.size()); + FunctionNode>, Object> maps = + FunctionNodes.of( + "maps", + new TypeReference>>() {}, + Object.class, + (ctx, l) -> l.get(0).get("note")); + + // Act + WorkflowConfigurationException error = + assertThrows(WorkflowConfigurationException.class, () -> chain(fulls, basics)); + RunResult result = run(chain(fulls, maps)); + + // Assert + assertThat(error) + .hasMessageThat() + .endsWith("set inputSchema on the builder or read the objects as Maps to accept them."); + result.assertCompleted(); + assertThat(result.outputAt("wf@1/maps@1")).isEqualTo("rush"); + } + + @Test + public void build_recordIntoAHandWrittenSchema_admitsKeysItDoesNotList() { + // Arrange: a schema set on the builder is open, unlike an inferred one. + FunctionNode basic = + FunctionNodes.builder(BasicOrder.class, String.class) + .name("basic") + .inputSchema(object(ImmutableMap.of("id", STRING), "id")) + .body((ctx, order) -> order.id()) + .build(); + + // Act + RunResult result = run(chain(fullOrder(), basic)); + + // Assert + result.assertCompleted(); + assertThat(result.outputAt("wf@1/basic@1")).isEqualTo("a1"); + } + + @Test + public void build_integerIntoANullableIntegerSchema_buildsAndRuns() { + // Arrange + FunctionNode three = + FunctionNodes.of("three", Content.class, Integer.class, (ctx, message) -> 3); + FunctionNode increment = + FunctionNodes.builder(Integer.class, Integer.class) + .name("increment") + .inputSchema(nullable(INTEGER)) + .body((ctx, n) -> n + 1) + .build(); + + // Act + RunResult result = run(chain(three, increment)); + + // Assert + result.assertCompleted(); + assertThat(result.outputAt("wf@1/increment@1")).isEqualTo(4); + } + + @Test + public void build_recordIntoAContentPort_failsWithoutSuggestingABuilderSchema() { + // Arrange: a builder schema on a Content port describes the text, so it cannot help here. + FunctionNode full = fullOrder(); + FunctionNode read = + FunctionNodes.of("read", Content.class, String.class, (ctx, content) -> content.text()); + + // Act + WorkflowConfigurationException error = + assertThrows(WorkflowConfigurationException.class, () -> chain(full, read)); + + // Assert + assertThat(error).hasMessageThat().startsWith("Schema mismatch on edge full -> read."); + assertThat(error).hasMessageThat().doesNotContain("inputSchema on the builder"); + } + + @Test + public void build_contentIntoAStringPort_fails() { + // Arrange: a returned Content is the output in its JSON form, an object. + FunctionNode reply = reply(); + FunctionNode shout = + FunctionNodes.of("shout", String.class, String.class, (ctx, text) -> text); + + // Act + WorkflowConfigurationException error = + assertThrows(WorkflowConfigurationException.class, () -> chain(reply, shout)); + + // Assert: the keys are not the cause, so the message suggests no way to accept them. + assertThat(error).hasMessageThat().startsWith("Schema mismatch on edge reply -> shout."); + assertThat(error).hasMessageThat().doesNotContain("takes no other keys"); + } + + @Test + public void build_recordIntoAnEventPort_fails() { + // Arrange: an Event port takes only an event's keys, and a FullOrder has others. + FunctionNode full = fullOrder(); + FunctionNode author = + FunctionNodes.of("author", Event.class, String.class, (ctx, event) -> event.author()); + + // Act + WorkflowConfigurationException error = + assertThrows(WorkflowConfigurationException.class, () -> chain(full, author)); + + // Assert + assertThat(error).hasMessageThat().startsWith("Schema mismatch on edge full -> author."); + } + + @Test + public void build_contentIntoContentOrObjectPorts_buildsAndRuns() { + // Arrange + FunctionNode read = + FunctionNodes.of("read", Content.class, String.class, (ctx, content) -> content.text()); + FunctionNode any = + FunctionNodes.of("any", Object.class, String.class, (ctx, value) -> "ok"); + + // Act + RunResult intoContent = run(chain(reply(), read)); + RunResult intoObject = run(chain(reply(), any)); + + // Assert + intoContent.assertCompleted(); + intoObject.assertCompleted(); + assertThat(intoContent.outputAt("wf@1/read@1")).isEqualTo("hi"); + assertThat(intoObject.outputAt("wf@1/any@1")).isEqualTo("ok"); + } + + @Test + public void build_contentPortWithAnExplicitSchema_checksTheEdgeAgainstItsContentSchema() { + // Arrange: an explicit schema on a Content port describes the message's text. + FunctionNode text = + FunctionNodes.builder(Content.class, String.class) + .name("text") + .inputSchema(STRING) + .body((ctx, content) -> content.text()) + .build(); + FunctionNode full = fullOrder(); + + // Act + RunResult result = run(chain(reply(), text)); + WorkflowConfigurationException error = + assertThrows(WorkflowConfigurationException.class, () -> chain(full, text)); + + // Assert + result.assertCompleted(); + assertThat(result.outputAt("wf@1/text@1")).isEqualTo("hi"); + assertThat(error).hasMessageThat().startsWith("Schema mismatch on edge full -> text."); + } + + @Test + public void build_edgeIntoANodeThatIsNotAFunctionNode_checksItsOpenSchema() { + // Arrange: such a node infers nothing, so its schema takes keys it does not list. + FunctionNode text = + FunctionNodes.of("text", Content.class, String.class, (ctx, message) -> message.text()); + JoinNode wantsInteger = JoinNode.builder().name("join").inputSchema(INTEGER).build(); + JoinNode wantsId = + JoinNode.builder() + .name("join") + .inputSchema(object(ImmutableMap.of("id", STRING), "id")) + .build(); + + // Act + WorkflowConfigurationException error = + assertThrows(WorkflowConfigurationException.class, () -> chain(text, wantsInteger)); + Workflow workflow = chain(fullOrder(), wantsId); + + // Assert + assertThat(error).hasMessageThat().startsWith("Schema mismatch on edge text -> join."); + assertThat(workflow.graph()).isPresent(); + } + + @Test + public void build_recordsIntoPortsThatReadOtherKeysThanTheyWrite_buildAndRun() { + // Arrange: Invoice writes a total it does not read; Login reads a password it does not write. + FunctionNode order = + FunctionNodes.of("order", Content.class, Order.class, (ctx, m) -> new Order("a1", 2)); + FunctionNode invoice = + FunctionNodes.of( + "invoice", Invoice.class, String.class, (ctx, i) -> i.id() + "=" + i.getTotal()); + FunctionNode creds = + FunctionNodes.of("creds", Content.class, Creds.class, (ctx, m) -> new Creds("ann", "pw")); + FunctionNode login = + FunctionNodes.of( + "login", Login.class, String.class, (ctx, l) -> l.user() + ":" + l.password().length()); + + // Act + RunResult invoiced = run(chain(order, invoice)); + RunResult loggedIn = run(chain(creds, login)); + + // Assert + invoiced.assertCompleted(); + loggedIn.assertCompleted(); + assertThat(invoiced.outputAt("wf@1/invoice@1")).isEqualTo("a1=20"); + assertThat(loggedIn.outputAt("wf@1/login@1")).isEqualTo("ann:2"); + } + + @Test + public void build_contentIntoASchemaWithAnOptionalRole_buildsAndRuns() { + // Arrange: a Content omits an unset role rather than writing null. + FunctionNode model = + FunctionNodes.of( + "model", + Content.class, + Content.class, + (ctx, m) -> Content.builder().role("model").parts(Part.fromText("hi")).build()); + FunctionNode, Object> role = + FunctionNodes.builder(new TypeReference>() {}, Object.class) + .name("role") + .inputSchema(object(ImmutableMap.of("role", STRING))) + .body((ctx, message) -> message.get("role")) + .build(); + + // Act + RunResult result = run(chain(model, role)); + + // Assert + result.assertCompleted(); + assertThat(result.outputAt("wf@1/role@1")).isEqualTo("model"); + } + + @Test + public void build_recordWithAnEventKeyIntoAnEventPort_buildsAndRuns() { + // Arrange: an Event port reads each key as optional, as Event.fromJson does. + FunctionNode authored = + FunctionNodes.of( + "authored", Content.class, Authored.class, (ctx, m) -> new Authored("ann")); + FunctionNode author = + FunctionNodes.of("author", Event.class, String.class, (ctx, event) -> event.author()); + + // Act + RunResult result = run(chain(authored, author)); + + // Assert + result.assertCompleted(); + assertThat(result.outputAt("wf@1/author@1")).isEqualTo("ann"); + } + + @Test + public void schemas_propertiesOnlyReadOrOnlyWritten_appearOnlyInThatPortsSchema() { + // Act + Optional notepadIn = inputSchema(Notepad.class); + Optional notepadOut = outputSchema(Notepad.class); + Optional invoiceIn = inputSchema(Invoice.class); + Optional invoiceOut = outputSchema(Invoice.class); + Optional loginIn = inputSchema(Login.class); + Optional loginOut = outputSchema(Login.class); + + // Assert: Jackson leaves a field or setter property it does not read at its default. + assertThat(notepadIn).hasValue(object(ImmutableMap.of("id", STRING, "note", STRING))); + assertThat(notepadOut).hasValue(object(ImmutableMap.of("id", STRING), "id")); + assertThat(invoiceIn).hasValue(ORDER); + assertThat(invoiceOut.orElseThrow().required()) + .hasValue(ImmutableList.of("id", "quantity", "total")); + assertThat(loginIn) + .hasValue(object(ImmutableMap.of("user", STRING, "password", STRING), "user")); + assertThat(loginOut) + .hasValue(object(ImmutableMap.of("user", STRING, "role", STRING), "user", "role")); + } + + @Test + public void build_sameTypeUnderClassInclusionOrFlatKeysIntoUnwrapped_builds() { + // Arrange: NON_NULL omits only nullable keys, and an unwrapped port infers no schema. + FunctionNode counts = + FunctionNodes.of( + "counts", + Content.class, + Counted.class, + (ctx, m) -> new Counted(1, "a", null, Optional.empty())); + FunctionNode count = + FunctionNodes.of("count", Counted.class, Integer.class, (ctx, c) -> c.count()); + FunctionNode flatOut = + FunctionNodes.of( + "flat_out", Content.class, FlatOut.class, (ctx, m) -> new FlatOut("l", "a1", 2)); + FunctionNode flat = + FunctionNodes.of("flat", Flat.class, String.class, (ctx, f) -> f.label()); + FunctionNode setterFlat = + FunctionNodes.of("setter_flat", SetterFlat.class, String.class, (ctx, f) -> f.label); + + // Act + Workflow counted = chain(counts, count); + Workflow flattened = chain(flatOut, flat); + Workflow setterFlattened = chain(flatOut, setterFlat); + + // Assert + assertThat(counted.graph()).isPresent(); + assertThat(flattened.graph()).isPresent(); + assertThat(setterFlattened.graph()).isPresent(); + } + + @Test + public void build_sameTypeUnderClassNonDefaultOrNonEmpty_builds() { + // Arrange: such a key may be omitted on write, and an absent one reads as its default. + FunctionNode makeCount = + FunctionNodes.of("make_count", Content.class, Count.class, (ctx, m) -> new Count(0)); + FunctionNode readCount = + FunctionNodes.of("read_count", Count.class, Integer.class, (ctx, c) -> c.count()); + FunctionNode makeTags = + FunctionNodes.of( + "make_tags", Content.class, Tags.class, (ctx, m) -> new Tags("", ImmutableList.of())); + FunctionNode readTags = + FunctionNodes.of("read_tags", Tags.class, String.class, (ctx, t) -> t.text()); + + // Act + Workflow counts = chain(makeCount, readCount); + Workflow tags = chain(makeTags, readTags); + + // Assert + assertThat(counts.graph()).isPresent(); + assertThat(tags.graph()).isPresent(); + } + + @Test + public void build_producerThatMayOmitAKeyTheTargetNeeds_fails() { + // Arrange: Defaults omits a value equal to its initializer, and Labeled an empty label. + FunctionNode defaults = + FunctionNodes.of("defaults", Content.class, Defaults.class, (ctx, m) -> new Defaults()); + FunctionNode view = + FunctionNodes.of("view", View.class, Status.class, (ctx, v) -> v.status()); + FunctionNode makeLabeled = + FunctionNodes.of( + "make_labeled", Content.class, Labeled.class, (ctx, m) -> new Labeled("", 1)); + FunctionNode readLabeled = + FunctionNodes.of("read_labeled", Labeled.class, String.class, (ctx, l) -> l.label()); + + // Act + WorkflowConfigurationException intoView = + assertThrows(WorkflowConfigurationException.class, () -> chain(defaults, view)); + WorkflowConfigurationException intoLabeled = + assertThrows(WorkflowConfigurationException.class, () -> chain(makeLabeled, readLabeled)); + + // Assert + assertThat(intoView).hasMessageThat().startsWith("Schema mismatch on edge defaults -> view."); + assertThat(intoLabeled) + .hasMessageThat() + .startsWith("Schema mismatch on edge make_labeled -> read_labeled."); + } + + @Test + public void outputSchema_unwrappedGetterBesideOverloadedSetters_inferNone() { + // Act + Optional schema = outputSchema(OverFlat.class); + + // Assert: Jackson cannot pick a setter, but the getter still writes the order's keys. + assertThat(schema).isEmpty(); + } + + @Test + public void outputSchema_nonDefaultNestedRecordOrPropertyLevelAlways_staysRequired() { + // Act + Optional packed = outputSchema(Packed.class); + Optional overridden = outputSchema(Overridden.class); + + // Assert: NON_DEFAULT always writes a non-null record. + assertThat(packed).hasValue(object(ImmutableMap.of("count", INTEGER, "order", ORDER), "order")); + assertThat(overridden) + .hasValue(object(ImmutableMap.of("text", STRING, "note", STRING), "text")); + } + + @Test + public void build_recordIntoAnAliasOfItsKey_buildsAndRuns() { + // Arrange: Aliased reads its quantity from a qty key too. + FunctionNode qty = + FunctionNodes.of("qty", Content.class, Qty.class, (ctx, m) -> new Qty(3)); + FunctionNode aliased = + FunctionNodes.of("aliased", Aliased.class, String.class, (ctx, a) -> "x" + a.quantity()); + + // Act + RunResult result = run(chain(qty, aliased)); + + // Assert + result.assertCompleted(); + assertThat(result.outputAt("wf@1/aliased@1")).isEqualTo("x3"); + } + @Test public void builder_withoutABody_failsAtBuild() { FunctionNode.Builder builder = @@ -1437,9 +1940,9 @@ public void of_mapOfRecordsPort_readsEachValueAsARecord() { @Test public void of_numberIntoAListPort_failsTheNode() { - // Arrange: a number feeds a list port. - FunctionNode number = - FunctionNodes.of("number", Content.class, Integer.class, (ctx, message) -> 123); + // Arrange: a number feeds a list port, by an Object output the build does not check. + FunctionNode number = + FunctionNodes.of("number", Content.class, Object.class, (ctx, message) -> 123); FunctionNode, String> first = FunctionNodes.of( "first", @@ -2650,6 +3153,26 @@ private static Schema object(ImmutableMap properties, String... return (required.length == 0 ? object : object.required(required)).build(); } + /** Returns a node named {@code reply} that returns a Content saying "hi". */ + private static FunctionNode reply() { + return FunctionNodes.of( + "reply", Content.class, Content.class, (ctx, m) -> Content.fromParts(Part.fromText("hi"))); + } + + /** Returns a node named {@code full} that returns a {@link FullOrder} with a note. */ + private static FunctionNode fullOrder() { + return FunctionNodes.of( + "full", Content.class, FullOrder.class, (ctx, m) -> new FullOrder("a1", 9.5, "rush")); + } + + private static FunctionNode> fullOrders() { + return FunctionNodes.of( + "fulls", + Content.class, + new TypeReference>() {}, + (ctx, m) -> ImmutableList.of(new FullOrder("a1", 9.5, "rush"))); + } + /** Returns a workflow named {@code wf} that runs {@code nodes} in sequence from START. */ private static Workflow chain(Node... nodes) { List edges = new ArrayList<>(); diff --git a/core/src/test/java/com/google/adk/workflow/SchemaFitTest.java b/core/src/test/java/com/google/adk/workflow/SchemaFitTest.java new file mode 100644 index 000000000..07f524ccc --- /dev/null +++ b/core/src/test/java/com/google/adk/workflow/SchemaFitTest.java @@ -0,0 +1,245 @@ +/* + * Copyright 2026 Google LLC + * + * 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 + * + * http://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 com.google.adk.workflow; + +import static com.google.common.collect.ImmutableMap.toImmutableMap; +import static com.google.common.collect.ImmutableSet.toImmutableSet; +import static com.google.common.truth.Truth.assertThat; + +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import com.google.genai.types.Schema; +import com.google.genai.types.Type; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests {@link SchemaFit}, which checks the schemas of a workflow edge when it is built. */ +@RunWith(JUnit4.class) +public final class SchemaFitTest { + + private static final Schema NULL = typed(Type.Known.NULL); + private static final Schema INTEGER = typed(Type.Known.INTEGER); + private static final Schema NUMBER = typed(Type.Known.NUMBER); + private static final Schema STRING = typed(Type.Known.STRING); + private static final Schema INTEGER_OR_STRING = Schema.builder().anyOf(INTEGER, STRING).build(); + private static final Schema RED = Schema.builder().type(Type.Known.STRING).enum_("red").build(); + private static final Schema COLORS = + Schema.builder().type(Type.Known.STRING).enum_("red", "blue").build(); + private static final Schema ID_ONLY = object(ImmutableMap.of("id", STRING), "id"); + private static final Schema ID_AND_NOTE = + object(ImmutableMap.of("id", STRING, "note", STRING), "id"); + + @Test + public void fits_everySourceShapeIntoEveryTargetShape_matchesTheAssignabilityTable() { + // Arrange: "closed object" is "open object" as a node infers it, so it is only a target. + ImmutableMap shapes = + ImmutableMap.builder() + .put("integer", INTEGER) + .put("number", NUMBER) + .put("string", STRING) + .put("untyped", Schema.builder().build()) + .put("integer or string", INTEGER_OR_STRING) + .put("red", RED) + .put("nullable integer", INTEGER.toBuilder().nullable(true).build()) + .put("map", typed(Type.Known.OBJECT)) + .put("open object", ID_ONLY) + .put("closed object", ID_ONLY) + .put("id and note", ID_AND_NOTE) + .put("integer list", array(INTEGER)) + .buildOrThrow(); + ImmutableSet objects = + ImmutableSet.of("untyped", "map", "open object", "closed object", "id and note"); + ImmutableMap> expected = + ImmutableMap.>builder() + .put( + "integer", + ImmutableSet.of( + "integer", "number", "untyped", "integer or string", "nullable integer")) + .put("number", ImmutableSet.of("number", "untyped")) + .put("string", ImmutableSet.of("string", "untyped", "integer or string")) + .put("untyped", shapes.keySet()) + .put("integer or string", ImmutableSet.of("untyped", "integer or string")) + .put("red", ImmutableSet.of("string", "untyped", "integer or string", "red")) + .put("nullable integer", ImmutableSet.of("untyped", "nullable integer")) + .put("map", objects) + .put("open object", objects) + .put("id and note", ImmutableSet.of("untyped", "map", "open object", "id and note")) + .put("integer list", ImmutableSet.of("untyped", "integer list")) + .buildOrThrow(); + + // Act + ImmutableMap> actual = + shapes.keySet().stream() + .filter(source -> !source.equals("closed object")) + .collect( + toImmutableMap( + source -> source, + source -> + shapes.keySet().stream() + .filter( + target -> + SchemaFit.fits( + shapes.get(source), + shapes.get(target), + target.equals("closed object"))) + .collect(toImmutableSet()))); + + // Assert + assertThat(actual).isEqualTo(expected); + } + + @Test + public void fits_nullableSourceOrUnion_needsATargetThatTakesNull() { + // Arrange + Schema integerOrNull = Schema.builder().anyOf(INTEGER, NULL).build(); + Schema nullableIntegerOrString = INTEGER_OR_STRING.toBuilder().nullable(true).build(); + + // Act / Assert + assertThat(fits(INTEGER.toBuilder().nullable(true).build(), integerOrNull)).isTrue(); + assertThat(fits(nullableIntegerOrString, INTEGER_OR_STRING)).isFalse(); + assertThat(fits(NULL, nullableIntegerOrString)).isTrue(); + assertThat(fits(NULL, INTEGER_OR_STRING)).isFalse(); + } + + @Test + public void fits_enums_stayWithinTheTargetEnum() { + // Arrange: an untyped enum lists strings, as a typed one does. + Schema untypedRed = Schema.builder().enum_("red").build(); + + // Act / Assert + assertThat(fits(RED, COLORS)).isTrue(); + assertThat(fits(COLORS, RED)).isFalse(); + assertThat(fits(STRING, RED)).isFalse(); + assertThat(fits(untypedRed, STRING)).isTrue(); + assertThat(fits(untypedRed, INTEGER)).isFalse(); + assertThat(fits(Schema.builder().enum_("green").build(), COLORS)).isFalse(); + } + + @Test + public void fits_objectsAndArrays_checkRequiredDeclaredAndClosedKeysAtEachLevel() { + // Arrange + Schema integerId = object(ImmutableMap.of("id", INTEGER), "id"); + Schema requiresId = Schema.builder().type(Type.Known.OBJECT).required("id").build(); + Schema noteOnly = object(ImmutableMap.of("note", STRING)); + Schema holdsIdAndNote = object(ImmutableMap.of("order", ID_AND_NOTE), "order"); + Schema holdsIdOnly = object(ImmutableMap.of("order", ID_ONLY), "order"); + + // Act / Assert + assertThat(fits(integerId, ID_ONLY)).isFalse(); + assertThat(fits(ID_ONLY, ID_AND_NOTE.toBuilder().required("id", "note").build())).isFalse(); + assertThat(fits(requiresId, ID_ONLY)).isFalse(); + assertThat(SchemaFit.fits(requiresId, noteOnly, /* closed= */ true)).isFalse(); + assertThat(fits(holdsIdAndNote, holdsIdOnly)).isTrue(); + assertThat(SchemaFit.fits(holdsIdAndNote, holdsIdOnly, /* closed= */ true)).isFalse(); + // A map lists no properties, so it stays open inside a closed target. + assertThat(SchemaFit.fits(ID_ONLY, typed(Type.Known.OBJECT), /* closed= */ true)).isTrue(); + assertThat(fits(array(STRING), array(INTEGER))).isFalse(); + assertThat(SchemaFit.fits(array(ID_AND_NOTE), array(ID_ONLY), /* closed= */ true)).isFalse(); + assertThat(fits(typed(Type.Known.ARRAY), array(INTEGER))).isTrue(); + } + + @Test + public void fits_closedTarget_admitsAnotherObjectThatUsesOnlyItsKeys() { + // Arrange: each target lists a key the source does not have, or does not require one it has. + Schema basicOrder = object(ImmutableMap.of("id", STRING, "total", NUMBER), "id", "total"); + Schema orderWithOptionalNote = + object(ImmutableMap.of("id", STRING, "total", NUMBER, "note", STRING), "id", "total"); + Schema requiresIdAndNote = ID_AND_NOTE.toBuilder().required("id", "note").build(); + Schema integerId = object(ImmutableMap.of("id", INTEGER), "id"); + Schema numberId = object(ImmutableMap.of("id", NUMBER), "id"); + + // Act / Assert + assertThat(SchemaFit.fits(basicOrder, orderWithOptionalNote, /* closed= */ true)).isTrue(); + assertThat(SchemaFit.fits(ID_ONLY, ID_AND_NOTE, /* closed= */ true)).isTrue(); + assertThat(SchemaFit.fits(requiresIdAndNote, ID_AND_NOTE, /* closed= */ true)).isTrue(); + assertThat(SchemaFit.fits(integerId, numberId, /* closed= */ true)).isTrue(); + } + + @Test + public void fits_unions_needEverySourceBranchAndOneTargetBranch() { + // Arrange + Schema green = Schema.builder().type(Type.Known.STRING).enum_("green").build(); + Schema requiresId = Schema.builder().required("id").build(); + Schema requiresNote = object(ImmutableMap.of("note", STRING), "note"); + Schema stringRed = Schema.builder().type(Type.Known.STRING).anyOf(RED).build(); + + // Act / Assert + assertThat(fits(INTEGER_OR_STRING, INTEGER)).isFalse(); + assertThat(fits(INTEGER, Schema.builder().type(Type.Known.STRING).anyOf(INTEGER).build())) + .isFalse(); + assertThat(fits(Schema.builder().anyOf(RED).build(), COLORS)).isTrue(); + assertThat(fits(Schema.builder().anyOf(RED, green).build(), COLORS)).isFalse(); + assertThat(fits(Schema.builder().enum_("green").anyOf(COLORS).build(), COLORS)).isFalse(); + // An untyped union's own enum takes its branches' type. + assertThat(fits(Schema.builder().enum_("1").anyOf(INTEGER).build(), INTEGER)).isTrue(); + // An untyped branch takes its union's type, so its keys count. + assertThat( + fits(Schema.builder().type(Type.Known.OBJECT).anyOf(requiresId).build(), requiresNote)) + .isFalse(); + // A union that rejects null drops its NULL branch. + assertThat(fits(Schema.builder().type(Type.Known.STRING).anyOf(STRING, NULL).build(), STRING)) + .isTrue(); + // A NULL branch does not count toward the type an untyped union's own enum takes. + assertThat( + fits( + Schema.builder().enum_("1").nullable(true).anyOf(INTEGER, NULL).build(), + INTEGER.toBuilder().nullable(true).build())) + .isTrue(); + // Untyped branches imply a type by their keys or items, so the union's own part is checked. + assertThat( + fits( + Schema.builder() + .required("id") + .anyOf(Schema.builder().properties(ImmutableMap.of("id", STRING)).build()) + .build(), + requiresNote.toBuilder().nullable(true).build())) + .isFalse(); + assertThat( + fits( + Schema.builder() + .items(INTEGER) + .anyOf(Schema.builder().items(INTEGER).build()) + .build(), + array(STRING).toBuilder().nullable(true).build())) + .isFalse(); + // A closed target stays closed for each branch of a union. + assertThat( + SchemaFit.fits( + Schema.builder().anyOf(ID_AND_NOTE).build(), ID_ONLY, /* closed= */ true)) + .isFalse(); + // Only identity lets this union fit itself: its own part fits none of its branches. + assertThat(fits(stringRed, stringRed)).isTrue(); + } + + private static boolean fits(Schema source, Schema target) { + return SchemaFit.fits(source, target, /* closed= */ false); + } + + private static Schema typed(Type.Known type) { + return Schema.builder().type(type).build(); + } + + private static Schema array(Schema items) { + return Schema.builder().type(Type.Known.ARRAY).items(items).build(); + } + + private static Schema object(ImmutableMap properties, String... required) { + Schema.Builder object = Schema.builder().type(Type.Known.OBJECT).properties(properties); + return (required.length == 0 ? object : object.required(required)).build(); + } +}