Skip to content

Commit 6b77251

Browse files
google-genai-botcopybara-github
authored andcommitted
feat: Add JSON cycle detection
This chanve enhances the JsonFormatter within the ADK core to detect and handle circular references in JSON structures during smart truncation, preventing infinite recursion. PiperOrigin-RevId: 904554270
1 parent 52323b4 commit 6b77251

2 files changed

Lines changed: 100 additions & 28 deletions

File tree

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

Lines changed: 63 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616

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

19+
import static java.util.Collections.newSetFromMap;
20+
1921
import com.fasterxml.jackson.annotation.JsonProperty;
2022
import com.fasterxml.jackson.databind.JsonNode;
2123
import com.fasterxml.jackson.databind.ObjectMapper;
@@ -30,14 +32,17 @@
3032
import com.google.genai.types.FunctionCall;
3133
import com.google.genai.types.Part;
3234
import java.util.ArrayList;
35+
import java.util.IdentityHashMap;
3336
import java.util.List;
3437
import java.util.Map;
3538
import java.util.Optional;
3639
import java.util.Set;
40+
import java.util.logging.Logger;
3741
import org.jspecify.annotations.Nullable;
3842

3943
/** Utility for parsing, formatting and truncating content for BigQuery logging. */
4044
final class JsonFormatter {
45+
private static final Logger logger = Logger.getLogger(JsonFormatter.class.getName());
4146
private static final ObjectMapper mapper = new ObjectMapper().findAndRegisterModules();
4247

4348
@AutoValue
@@ -304,11 +309,15 @@ static TruncationResult smartTruncate(Object obj, int maxLength) {
304309
if (obj == null) {
305310
return TruncationResult.create(mapper.nullNode(), false);
306311
}
312+
if (obj instanceof JsonNode jsonNode) {
313+
return recursiveSmartTruncate(jsonNode, maxLength, newSetFromMap(new IdentityHashMap<>()));
314+
}
307315
try {
308-
return recursiveSmartTruncate(mapper.valueToTree(obj), maxLength);
309-
} catch (IllegalArgumentException e) {
316+
return recursiveSmartTruncate(
317+
mapper.valueToTree(obj), maxLength, newSetFromMap(new IdentityHashMap<>()));
318+
} catch (IllegalArgumentException | StackOverflowError e) {
310319
// Fallback for types that mapper can't handle directly as a tree
311-
return truncateWithStatus(String.valueOf(obj), maxLength);
320+
return truncateWithStatus(safeToString(obj), maxLength);
312321
}
313322
}
314323

@@ -318,39 +327,65 @@ static JsonNode convertToJsonNode(Object obj) {
318327
}
319328
try {
320329
return mapper.valueToTree(obj);
321-
} catch (IllegalArgumentException e) {
330+
} catch (IllegalArgumentException | StackOverflowError e) {
322331
// Fallback for types that mapper can't handle directly as a tree
323-
return mapper.valueToTree(String.valueOf(obj));
332+
return mapper.valueToTree(safeToString(obj));
324333
}
325334
}
326335

327-
private static TruncationResult recursiveSmartTruncate(JsonNode node, int maxLength) {
328-
boolean isTruncated = false;
329-
if (node.isTextual()) {
330-
String text = node.asText();
331-
if (text.length() > maxLength) {
332-
return TruncationResult.create(mapper.valueToTree(truncate(text, maxLength)), true);
336+
static String safeToString(Object obj) {
337+
try {
338+
return String.valueOf(obj);
339+
} catch (StackOverflowError e) {
340+
logger.warning("StackOverflowError when converting object to string");
341+
return "[STACK OVERFLOW ERROR CONVERTING TO STRING]";
342+
} catch (RuntimeException e) {
343+
logger.warning("RuntimeException when converting object to string");
344+
return "[ERROR CONVERTING TO STRING]";
345+
}
346+
}
347+
348+
private static TruncationResult recursiveSmartTruncate(
349+
JsonNode node, int maxLength, Set<JsonNode> visited) {
350+
if (node.isContainerNode()) {
351+
if (visited.contains(node)) {
352+
return TruncationResult.create(mapper.valueToTree("[CYCLE DETECTED]"), true);
333353
}
334-
return TruncationResult.create(node, false);
335-
} else if (node.isObject()) {
336-
ObjectNode newNode = mapper.createObjectNode();
337-
Set<Map.Entry<String, JsonNode>> properties = node.properties();
338-
for (Map.Entry<String, JsonNode> entry : properties) {
339-
TruncationResult res = recursiveSmartTruncate(entry.getValue(), maxLength);
340-
newNode.set(entry.getKey(), res.node());
341-
isTruncated = isTruncated || res.isTruncated();
354+
visited.add(node);
355+
}
356+
357+
try {
358+
boolean isTruncated = false;
359+
if (node.isTextual()) {
360+
String text = node.asText();
361+
if (text.length() > maxLength) {
362+
return TruncationResult.create(mapper.valueToTree(truncate(text, maxLength)), true);
363+
}
364+
return TruncationResult.create(node, false);
365+
} else if (node.isObject()) {
366+
ObjectNode newNode = mapper.createObjectNode();
367+
Set<Map.Entry<String, JsonNode>> properties = node.properties();
368+
for (Map.Entry<String, JsonNode> entry : properties) {
369+
TruncationResult res = recursiveSmartTruncate(entry.getValue(), maxLength, visited);
370+
newNode.set(entry.getKey(), res.node());
371+
isTruncated = isTruncated || res.isTruncated();
372+
}
373+
return TruncationResult.create(newNode, isTruncated);
374+
} else if (node.isArray()) {
375+
ArrayNode newNode = mapper.createArrayNode();
376+
for (JsonNode element : node) {
377+
TruncationResult res = recursiveSmartTruncate(element, maxLength, visited);
378+
newNode.add(res.node());
379+
isTruncated = isTruncated || res.isTruncated();
380+
}
381+
return TruncationResult.create(newNode, isTruncated);
342382
}
343-
return TruncationResult.create(newNode, isTruncated);
344-
} else if (node.isArray()) {
345-
ArrayNode newNode = mapper.createArrayNode();
346-
for (JsonNode element : node) {
347-
TruncationResult res = recursiveSmartTruncate(element, maxLength);
348-
newNode.add(res.node());
349-
isTruncated = isTruncated || res.isTruncated();
383+
return TruncationResult.create(node, false);
384+
} finally {
385+
if (node.isContainerNode()) {
386+
visited.remove(node);
350387
}
351-
return TruncationResult.create(newNode, isTruncated);
352388
}
353-
return TruncationResult.create(node, false);
354389
}
355390

356391
private static TruncationResult truncateWithStatus(String s, int maxLength) {

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

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,9 @@
2121
import static org.junit.Assert.assertTrue;
2222

2323
import com.fasterxml.jackson.databind.JsonNode;
24+
import com.fasterxml.jackson.databind.ObjectMapper;
2425
import com.fasterxml.jackson.databind.node.ArrayNode;
26+
import com.fasterxml.jackson.databind.node.ObjectNode;
2527
import com.google.adk.models.LlmRequest;
2628
import com.google.common.collect.ImmutableList;
2729
import com.google.common.collect.ImmutableMap;
@@ -142,4 +144,39 @@ public void parse_list_truncatesElements() {
142144
assertEquals("short", arrayNode.get(0).asText());
143145
assertEquals("this is a ...[truncated]", arrayNode.get(1).asText());
144146
}
147+
148+
@Test
149+
public void smartTruncate_withCycle_detectsCycle() {
150+
ObjectMapper mapper = new ObjectMapper();
151+
ObjectNode node = mapper.createObjectNode();
152+
node.set("child", node);
153+
154+
// Verify that smartTruncate handles circular JsonNode structures by detecting the cycle.
155+
JsonFormatter.TruncationResult result = JsonFormatter.smartTruncate(node, 100);
156+
157+
assertTrue(result.isTruncated());
158+
assertEquals("[CYCLE DETECTED]", result.node().get("child").asText());
159+
}
160+
161+
@Test
162+
public void smartTruncate_withToStringStackOverflow_handlesGracefully() {
163+
Object recursiveObj =
164+
new Object() {
165+
@Override
166+
public String toString() {
167+
return String.valueOf(this);
168+
}
169+
};
170+
171+
// Verify both direct safeToString and via smartTruncate
172+
assertEquals(
173+
"[STACK OVERFLOW ERROR CONVERTING TO STRING]", JsonFormatter.safeToString(recursiveObj));
174+
175+
JsonFormatter.TruncationResult result = JsonFormatter.smartTruncate(recursiveObj, 100);
176+
assertTrue(result.node().isTextual());
177+
// Note: This expectation depends on whether mapper.valueToTree(recursiveObj)
178+
// throws IllegalArgumentException or StackOverflowError.
179+
// If it throws StackOverflowError and we don't catch it in smartTruncate, this test will fail.
180+
assertEquals("[STACK OVERFLOW ERROR CONVERTING TO STRING]", result.node().asText());
181+
}
145182
}

0 commit comments

Comments
 (0)