Skip to content

Commit 1685a4e

Browse files
google-genai-botcopybara-github
authored andcommitted
feat: Add JSON cycle detection
This change enhances the JsonFormatter within the ADK core to detect and handle circular references in JSON structures during smart truncation, preventing infinite recursion. PiperOrigin-RevId: 923585268
1 parent 51c9d1a commit 1685a4e

2 files changed

Lines changed: 74 additions & 27 deletions

File tree

core/src/main/java/com/google/adk/plugins/agentanalytics/JsonFormatter.java

Lines changed: 59 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -16,20 +16,26 @@
1616

1717
package com.google.adk.plugins.agentanalytics;
1818

19+
import static java.util.Collections.newSetFromMap;
20+
1921
import com.fasterxml.jackson.databind.JsonNode;
2022
import com.fasterxml.jackson.databind.ObjectMapper;
2123
import com.fasterxml.jackson.databind.node.ArrayNode;
2224
import com.fasterxml.jackson.databind.node.ObjectNode;
2325
import com.google.auto.value.AutoValue;
2426
import com.google.common.base.Utf8;
27+
import java.util.IdentityHashMap;
2528
import java.util.Map;
2629
import java.util.Set;
30+
import java.util.logging.Logger;
2731
import org.jspecify.annotations.Nullable;
2832

2933
/** Utility for parsing, formatting and truncating content for BigQuery logging. */
3034
final class JsonFormatter {
35+
private static final Logger logger = Logger.getLogger(JsonFormatter.class.getName());
3136
static final ObjectMapper mapper = new ObjectMapper().findAndRegisterModules();
3237
static final String TRUNCATION_SUFFIX = "...[truncated]";
38+
static final String CYCLE_DETECTED_MESSAGE = "[cycle detected]";
3339

3440
@AutoValue
3541
abstract static class TruncationResult {
@@ -48,10 +54,14 @@ static TruncationResult smartTruncate(Object obj, int maxLength) {
4854
return TruncationResult.create(mapper.nullNode(), false);
4955
}
5056
try {
51-
return recursiveSmartTruncate(mapper.valueToTree(obj), maxLength);
57+
if (obj instanceof JsonNode jsonNode) {
58+
return recursiveSmartTruncate(jsonNode, maxLength, newSetFromMap(new IdentityHashMap<>()));
59+
}
60+
return recursiveSmartTruncate(
61+
mapper.valueToTree(obj), maxLength, newSetFromMap(new IdentityHashMap<>()));
5262
} catch (IllegalArgumentException e) {
5363
// Fallback for types that mapper can't handle directly as a tree
54-
return truncateWithStatus(String.valueOf(obj), maxLength);
64+
return truncateWithStatus(safeToString(obj), maxLength);
5565
}
5666
}
5767

@@ -62,38 +72,60 @@ static JsonNode convertToJsonNode(Object obj) {
6272
try {
6373
return mapper.valueToTree(obj);
6474
} catch (IllegalArgumentException e) {
65-
// Fallback for types that mapper can't handle directly as a tree
66-
return mapper.valueToTree(String.valueOf(obj));
75+
// Fallback for types that mapper can't handle directly as a tree.
76+
return mapper.valueToTree(safeToString(obj));
6777
}
6878
}
6979

70-
private static TruncationResult recursiveSmartTruncate(JsonNode node, int maxLength) {
71-
boolean isTruncated = false;
72-
if (node.isTextual()) {
73-
String text = node.asText();
74-
if (Utf8.encodedLength(text) > maxLength) {
75-
return TruncationResult.create(mapper.valueToTree(truncate(text, maxLength)), true);
80+
static String safeToString(Object obj) {
81+
try {
82+
return String.valueOf(obj);
83+
} catch (RuntimeException e) {
84+
logger.warning("RuntimeException when converting object to string");
85+
return "[ERROR CONVERTING TO STRING]";
86+
}
87+
}
88+
89+
private static TruncationResult recursiveSmartTruncate(
90+
JsonNode node, int maxLength, Set<JsonNode> visited) {
91+
if (node.isContainerNode()) {
92+
if (visited.contains(node)) {
93+
return TruncationResult.create(mapper.valueToTree(CYCLE_DETECTED_MESSAGE), true);
7694
}
77-
return TruncationResult.create(node, false);
78-
} else if (node.isObject()) {
79-
ObjectNode newNode = mapper.createObjectNode();
80-
Set<Map.Entry<String, JsonNode>> properties = node.properties();
81-
for (Map.Entry<String, JsonNode> entry : properties) {
82-
TruncationResult res = recursiveSmartTruncate(entry.getValue(), maxLength);
83-
newNode.set(entry.getKey(), res.node());
84-
isTruncated = isTruncated || res.isTruncated();
95+
visited.add(node);
96+
}
97+
try {
98+
boolean isTruncated = false;
99+
if (node.isTextual()) {
100+
String text = node.asText();
101+
if (Utf8.encodedLength(text) > maxLength) {
102+
return TruncationResult.create(mapper.valueToTree(truncate(text, maxLength)), true);
103+
}
104+
return TruncationResult.create(node, false);
105+
} else if (node.isObject()) {
106+
ObjectNode newNode = mapper.createObjectNode();
107+
Set<Map.Entry<String, JsonNode>> properties = node.properties();
108+
for (Map.Entry<String, JsonNode> entry : properties) {
109+
TruncationResult res = recursiveSmartTruncate(entry.getValue(), maxLength, visited);
110+
newNode.set(entry.getKey(), res.node());
111+
isTruncated = isTruncated || res.isTruncated();
112+
}
113+
return TruncationResult.create(newNode, isTruncated);
114+
} else if (node.isArray()) {
115+
ArrayNode newNode = mapper.createArrayNode();
116+
for (JsonNode element : node) {
117+
TruncationResult res = recursiveSmartTruncate(element, maxLength, visited);
118+
newNode.add(res.node());
119+
isTruncated = isTruncated || res.isTruncated();
120+
}
121+
return TruncationResult.create(newNode, isTruncated);
85122
}
86-
return TruncationResult.create(newNode, isTruncated);
87-
} else if (node.isArray()) {
88-
ArrayNode newNode = mapper.createArrayNode();
89-
for (JsonNode element : node) {
90-
TruncationResult res = recursiveSmartTruncate(element, maxLength);
91-
newNode.add(res.node());
92-
isTruncated = isTruncated || res.isTruncated();
123+
return TruncationResult.create(node, false);
124+
} finally {
125+
if (node.isContainerNode()) {
126+
visited.remove(node);
93127
}
94-
return TruncationResult.create(newNode, isTruncated);
95128
}
96-
return TruncationResult.create(node, false);
97129
}
98130

99131
static TruncationResult truncateWithStatus(String s, int maxLength) {

core/src/test/java/com/google/adk/plugins/agentanalytics/JsonFormatterTest.java

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,9 @@
2727
import static org.mockito.Mockito.when;
2828

2929
import com.fasterxml.jackson.databind.JsonNode;
30+
import com.fasterxml.jackson.databind.ObjectMapper;
3031
import com.fasterxml.jackson.databind.node.ArrayNode;
32+
import com.fasterxml.jackson.databind.node.ObjectNode;
3133
import com.google.adk.models.LlmRequest;
3234
import com.google.common.collect.ImmutableList;
3335
import com.google.common.collect.ImmutableMap;
@@ -257,4 +259,17 @@ public void parse_multibyteContent_truncatesBasedOnBytes() throws Exception {
257259
assertTrue(result.isTruncated());
258260
assertEquals("こん...[truncated]", result.content().get("text_summary").asText());
259261
}
262+
263+
@Test
264+
public void smartTruncate_withCycle_detectsCycle() {
265+
ObjectMapper mapper = new ObjectMapper();
266+
ObjectNode node = mapper.createObjectNode();
267+
node.set("child", node);
268+
269+
// Verify that smartTruncate handles circular JsonNode structures by detecting the cycle.
270+
JsonFormatter.TruncationResult result = JsonFormatter.smartTruncate(node, 100);
271+
272+
assertTrue(result.isTruncated());
273+
assertEquals("[cycle detected]", result.node().get("child").asText());
274+
}
260275
}

0 commit comments

Comments
 (0)