Skip to content

feat: Add JSON cycle detection #1156

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
+101 −29
Open

feat: Add JSON cycle detection #1156

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
93 changes: 64 additions & 29 deletions core/src/main/java/com/google/adk/plugins/agentanalytics/JsonFormatter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,20 +16,26 @@

package com.google.adk.plugins.agentanalytics;

import static java.util.Collections.newSetFromMap;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.google.auto.value.AutoValue;
import com.google.common.base.Utf8;
import java.util.IdentityHashMap;
import java.util.Map;
import java.util.Set;
import java.util.logging.Logger;
import org.jspecify.annotations.Nullable;

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

@AutoValue
abstract static class TruncationResult {
Expand All @@ -48,10 +54,14 @@ static TruncationResult smartTruncate(Object obj, int maxLength) {
return TruncationResult.create(mapper.nullNode(), false);
}
try {
return recursiveSmartTruncate(mapper.valueToTree(obj), maxLength);
} catch (IllegalArgumentException e) {
if (obj instanceof JsonNode jsonNode) {
return recursiveSmartTruncate(jsonNode, maxLength, newSetFromMap(new IdentityHashMap<>()));
}
return recursiveSmartTruncate(
mapper.valueToTree(obj), maxLength, newSetFromMap(new IdentityHashMap<>()));
} catch (IllegalArgumentException | StackOverflowError e) {
// Fallback for types that mapper can't handle directly as a tree
return truncateWithStatus(String.valueOf(obj), maxLength);
return truncateWithStatus(safeToString(obj), maxLength);
}
}

Expand All @@ -61,39 +71,64 @@ static JsonNode convertToJsonNode(Object obj) {
}
try {
return mapper.valueToTree(obj);
} catch (IllegalArgumentException e) {
// Fallback for types that mapper can't handle directly as a tree
return mapper.valueToTree(String.valueOf(obj));
} catch (IllegalArgumentException | StackOverflowError e) {
// Fallback for types that mapper can't handle directly as a tree.
return mapper.valueToTree(safeToString(obj));
}
}

private static TruncationResult recursiveSmartTruncate(JsonNode node, int maxLength) {
boolean isTruncated = false;
if (node.isTextual()) {
String text = node.asText();
if (Utf8.encodedLength(text) > maxLength) {
return TruncationResult.create(mapper.valueToTree(truncate(text, maxLength)), true);
static String safeToString(Object obj) {
try {
return String.valueOf(obj);
} catch (StackOverflowError e) {
logger.warning("StackOverflowError when converting object to string");
return "[STACK OVERFLOW ERROR CONVERTING TO STRING]";
} catch (RuntimeException e) {
logger.warning("RuntimeException when converting object to string");
return "[ERROR CONVERTING TO STRING]";
}
}

private static TruncationResult recursiveSmartTruncate(
JsonNode node, int maxLength, Set<JsonNode> visited) {
if (node.isContainerNode()) {
if (visited.contains(node)) {
return TruncationResult.create(mapper.valueToTree(CYCLE_DETECTED_MESSAGE), true);
}
return TruncationResult.create(node, false);
} else if (node.isObject()) {
ObjectNode newNode = mapper.createObjectNode();
Set<Map.Entry<String, JsonNode>> properties = node.properties();
for (Map.Entry<String, JsonNode> entry : properties) {
TruncationResult res = recursiveSmartTruncate(entry.getValue(), maxLength);
newNode.set(entry.getKey(), res.node());
isTruncated = isTruncated || res.isTruncated();
visited.add(node);
}
try {
boolean isTruncated = false;
if (node.isTextual()) {
String text = node.asText();
if (Utf8.encodedLength(text) > maxLength) {
return TruncationResult.create(mapper.valueToTree(truncate(text, maxLength)), true);
}
return TruncationResult.create(node, false);
} else if (node.isObject()) {
ObjectNode newNode = mapper.createObjectNode();
Set<Map.Entry<String, JsonNode>> properties = node.properties();
for (Map.Entry<String, JsonNode> entry : properties) {
TruncationResult res = recursiveSmartTruncate(entry.getValue(), maxLength, visited);
newNode.set(entry.getKey(), res.node());
isTruncated = isTruncated || res.isTruncated();
}
return TruncationResult.create(newNode, isTruncated);
} else if (node.isArray()) {
ArrayNode newNode = mapper.createArrayNode();
for (JsonNode element : node) {
TruncationResult res = recursiveSmartTruncate(element, maxLength, visited);
newNode.add(res.node());
isTruncated = isTruncated || res.isTruncated();
}
return TruncationResult.create(newNode, isTruncated);
}
return TruncationResult.create(newNode, isTruncated);
} else if (node.isArray()) {
ArrayNode newNode = mapper.createArrayNode();
for (JsonNode element : node) {
TruncationResult res = recursiveSmartTruncate(element, maxLength);
newNode.add(res.node());
isTruncated = isTruncated || res.isTruncated();
return TruncationResult.create(node, false);
} finally {
if (node.isContainerNode()) {
visited.remove(node);
}
return TruncationResult.create(newNode, isTruncated);
}
return TruncationResult.create(node, false);
}

static TruncationResult truncateWithStatus(String s, int maxLength) {
Expand Down
37 changes: 37 additions & 0 deletions core/src/test/java/com/google/adk/plugins/agentanalytics/JsonFormatterTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,9 @@
import static org.junit.Assert.assertTrue;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.google.adk.models.LlmRequest;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
Expand Down Expand Up @@ -202,4 +204,39 @@ public void parse_multibyteContent_truncatesBasedOnBytes() throws Exception {
assertTrue(result.isTruncated());
assertEquals("こん...[truncated]", result.content().get("text_summary").asText());
}

@Test
public void smartTruncate_withCycle_detectsCycle() {
ObjectMapper mapper = new ObjectMapper();
ObjectNode node = mapper.createObjectNode();
node.set("child", node);

// Verify that smartTruncate handles circular JsonNode structures by detecting the cycle.
JsonFormatter.TruncationResult result = JsonFormatter.smartTruncate(node, 100);

assertTrue(result.isTruncated());
assertEquals("[cycle detected]", result.node().get("child").asText());
}

@Test
public void smartTruncate_withToStringStackOverflow_handlesGracefully() {
Object recursiveObj =
new Object() {
@Override
public String toString() {
return String.valueOf(this);
}
};

// Verify both direct safeToString and via smartTruncate
assertEquals(
"[STACK OVERFLOW ERROR CONVERTING TO STRING]", JsonFormatter.safeToString(recursiveObj));

JsonFormatter.TruncationResult result = JsonFormatter.smartTruncate(recursiveObj, 100);
assertTrue(result.node().isTextual());
// Note: This expectation depends on whether mapper.valueToTree(recursiveObj)
// throws IllegalArgumentException or StackOverflowError.
// If it throws StackOverflowError and we don't catch it in smartTruncate, this test will fail.
assertEquals("[STACK OVERFLOW ERROR CONVERTING TO STRING]", result.node().asText());
}
}