diff --git a/core/src/main/java/com/google/adk/sessions/SessionJsonConverter.java b/core/src/main/java/com/google/adk/sessions/SessionJsonConverter.java index 1e6c5327a..3107fa25c 100644 --- a/core/src/main/java/com/google/adk/sessions/SessionJsonConverter.java +++ b/core/src/main/java/com/google/adk/sessions/SessionJsonConverter.java @@ -16,26 +16,48 @@ package com.google.adk.sessions; +import static com.google.common.collect.ImmutableList.toImmutableList; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.core.JsonParser; import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.DeserializationContext; +import com.fasterxml.jackson.databind.JsonMappingException; import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.deser.std.StdDeserializer; +import com.fasterxml.jackson.databind.module.SimpleModule; import com.google.adk.JsonBaseModel; import com.google.adk.events.Event; import com.google.adk.events.EventActions; +import com.google.adk.events.EventCompaction; import com.google.adk.events.ToolConfirmation; +import com.google.common.base.CaseFormat; import com.google.common.base.Splitter; +import com.google.common.base.Throwables; +import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import com.google.common.collect.Iterables; +import com.google.common.collect.Maps; import com.google.genai.types.Content; +import com.google.genai.types.CustomMetadata; import com.google.genai.types.FinishReason; import com.google.genai.types.GenerateContentResponseUsageMetadata; import com.google.genai.types.GroundingMetadata; +import com.google.genai.types.StringList; +import java.io.IOException; import java.io.UncheckedIOException; import java.time.Instant; +import java.util.ArrayList; +import java.util.Base64; import java.util.Collection; import java.util.HashMap; import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.Optional; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; @@ -43,11 +65,47 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -/** Handles JSON serialization and deserialization for session-related objects. */ +/** + * Static utilities for converting session events to and from Vertex AI Sessions API JSON. + * + *

Events are stored as Python ADK stores them: the typed fields carry what the API keeps, and + * {@code rawEvent} carries a copy of the event in the shape Python ADK reads. + */ final class SessionJsonConverter { private static final ObjectMapper objectMapper = JsonBaseModel.getMapper(); + + /** + * Maps event JSON stored in API {@code Struct}s as Python ADK does: null fields, but not null map + * values, are left out on write, and the URL-safe base64 Python ADK writes is accepted on read. + */ + private static final ObjectMapper storedJsonMapper = + objectMapper + .copy() + .setDefaultPropertyInclusion( + JsonInclude.Value.construct(JsonInclude.Include.NON_NULL, JsonInclude.Include.ALWAYS)) + .registerModule( + new SimpleModule().addDeserializer(byte[].class, new LenientBase64Deserializer())); + private static final Logger logger = LoggerFactory.getLogger(SessionJsonConverter.class); + // customMetadata keys under which Python ADK stores fields the API does not keep. + private static final String USAGE_METADATA_KEY = "_usage_metadata"; + private static final String COMPACTION_KEY = "_compaction"; + + /** + * {@code EventActions} keys Python ADK rejects ({@code setModelResponse} only before 2.x), left + * out of {@code rawEvent} so that it can read the event. + */ + private static final ImmutableSet PYTHON_UNSUPPORTED_ACTION_KEYS = + ImmutableSet.of("deletedArtifactIds", "setModelResponse"); + + /** + * Keys holding user data ({@code FunctionCall.args}, {@code FunctionResponse.response} and {@code + * Part.partMetadata}), whose nested keys {@link #snakeToCamelKeys} leaves as written. + */ + private static final ImmutableSet USER_DATA_KEYS = + ImmutableSet.of("args", "response", "part_metadata", "partMetadata"); + private SessionJsonConverter() {} /** @@ -68,6 +126,7 @@ static String convertEventToJson(Event event) { * @throws UncheckedIOException if serialization fails. */ static String convertEventToJson(Event event, boolean useIsoString) { + EventActions actions = event.actions(); Map metadataJson = new HashMap<>(); event.partial().ifPresent(v -> metadataJson.put("partial", v)); event.turnComplete().ifPresent(v -> metadataJson.put("turnComplete", v)); @@ -75,7 +134,18 @@ static String convertEventToJson(Event event, boolean useIsoString) { event.branch().ifPresent(v -> metadataJson.put("branch", v)); event.longRunningToolIds().ifPresent(v -> putIfNotEmpty(metadataJson, "longRunningToolIds", v)); event.groundingMetadata().ifPresent(v -> metadataJson.put("groundingMetadata", v)); - event.usageMetadata().ifPresent(v -> metadataJson.put("usageMetadata", v)); + Map customMetadataJson = + event + .customMetadata() + .map(SessionJsonConverter::customMetadataToJson) + .orElseGet(LinkedHashMap::new); + event.usageMetadata().ifPresent(v -> customMetadataJson.put(USAGE_METADATA_KEY, v)); + if (actions != null) { + actions + .compaction() + .ifPresent(v -> customMetadataJson.put(COMPACTION_KEY, compactionToJson(v))); + } + putIfNotEmpty(metadataJson, "customMetadata", customMetadataJson); Map eventJson = new HashMap<>(); eventJson.put("author", event.author()); eventJson.put("invocationId", event.invocationId()); @@ -95,36 +165,80 @@ static String convertEventToJson(Event event, boolean useIsoString) { event.errorMessage().ifPresent(errorMessage -> eventJson.put("errorMessage", errorMessage)); eventJson.put("eventMetadata", metadataJson); - if (event.actions() != null) { + if (actions != null) { + // The API keeps only these actions fields and drops the rest; rawEvent carries them. Map actionsJson = new HashMap<>(); - EventActions actions = event.actions(); actions.skipSummarization().ifPresent(v -> actionsJson.put("skipSummarization", v)); actionsJson.put("stateDelta", stateDeltaToJson(actions.stateDelta())); putIfNotEmpty(actionsJson, "artifactDelta", actions.artifactDelta()); - actions - .transferToAgent() - .ifPresent( - v -> { - actionsJson.put("transferAgent", v); - }); + actions.transferToAgent().ifPresent(v -> actionsJson.put("transferAgent", v)); actions.escalate().ifPresent(v -> actionsJson.put("escalate", v)); - if (actions.endOfAgent()) { - actionsJson.put("endOfAgent", actions.endOfAgent()); - } putIfNotEmpty(actionsJson, "requestedAuthConfigs", actions.requestedAuthConfigs()); - putIfNotEmpty( - actionsJson, "requestedToolConfirmations", actions.requestedToolConfirmations()); - actions.agentState().ifPresent(v -> actionsJson.put("agentState", v)); eventJson.put("actions", actionsJson); } event.content().ifPresent(c -> eventJson.put("content", SessionUtils.encodeContent(c))); try { + eventJson.put("rawEvent", rawEventToJson(event)); return objectMapper.writeValueAsString(eventJson); } catch (JsonProcessingException e) { throw new UncheckedIOException(e); } } + /** Serializes the event in the shape Python ADK writes to and reads from {@code rawEvent}. */ + private static Map rawEventToJson(Event event) throws JsonProcessingException { + Map rawEvent = + storedJsonMapper.readValue( + storedJsonMapper.writeValueAsString(event), + new TypeReference>() {}); + rawEvent.put("timestamp", millisToSeconds(event.timestamp())); + event.customMetadata().ifPresent(v -> rawEvent.put("customMetadata", customMetadataToJson(v))); + EventActions actions = event.actions(); + if (rawEvent.get("actions") instanceof Map rawActions) { + @SuppressWarnings("unchecked") // Jackson maps a JSON object to Map. + Map actionsJson = (Map) rawActions; + actionsJson.keySet().removeAll(PYTHON_UNSUPPORTED_ACTION_KEYS); + actionsJson.put("stateDelta", stateDeltaToJson(actions.stateDelta())); + actions.compaction().ifPresent(v -> actionsJson.put("compaction", compactionToJson(v))); + } + return rawEvent; + } + + /** Maps {@code customMetadata} entries to the key-value object Python ADK stores. */ + private static Map customMetadataToJson(List entries) { + Map json = new LinkedHashMap<>(); + for (CustomMetadata entry : entries) { + entry.key().ifPresent(key -> json.put(key, customMetadataValue(entry))); + } + return json; + } + + private static @Nullable Object customMetadataValue(CustomMetadata entry) { + return entry + .stringValue() + .map(value -> value) + .or(entry::numericValue) + .or(() -> entry.stringListValue().flatMap(StringList::values)) + .orElse(null); + } + + /** Serializes a compaction with the epoch-second timestamps Python ADK uses. */ + private static ImmutableMap compactionToJson(EventCompaction compaction) { + return ImmutableMap.of( + "startTimestamp", millisToSeconds(compaction.startTimestamp()), + "endTimestamp", millisToSeconds(compaction.endTimestamp()), + "compactedContent", compaction.compactedContent()); + } + + private static double millisToSeconds(long millis) { + return millis / 1000.0; + } + + /** Floors to millis, as event timestamps are; rounding to micros first cancels float error. */ + private static long secondsToMillis(Number seconds) { + return Math.floorDiv(Math.round(seconds.doubleValue() * 1_000_000), 1000); + } + /** * Converts a raw value to a {@link Content} object. * @@ -152,22 +266,162 @@ private static Content convertMapToContent(Object rawContentValue) { } } - /** Returns {@code key} from {@code actions}, falling back to the copy under {@code rawEvent}. */ - private static @Nullable Object checkpointValue( - @Nullable Map actions, @Nullable Map rawActions, String key) { - if (actions != null && actions.get(key) != null) { - return actions.get(key); - } - return rawActions == null ? null : rawActions.get(key); - } - /** * Converts raw API event data into an {@link Event} object. * + *

A readable {@code rawEvent} takes precedence over the typed fields, which remain the + * fallback when {@code rawEvent} is missing or unreadable. + * * @return parsed {@link Event}. */ - @SuppressWarnings("unchecked") // Parsing raw Map from JSON following a known schema. static Event fromApiEvent(Map apiEvent) { + if (apiEvent.get("rawEvent") instanceof Map rawEvent && !rawEvent.isEmpty()) { + long timestampMillis = convertToInstant(apiEvent.get("timestamp")).toEpochMilli(); + try { + return fromRawEvent(apiEvent, rawEvent, timestampMillis); + } catch (IllegalArgumentException e) { + logger.warn( + "Ignoring unreadable rawEvent ({}), using the typed event fields", + Throwables.getRootCause(e).getClass().getSimpleName()); + } + } + return fromTypedFields(apiEvent); + } + + private static Event fromRawEvent( + Map apiEvent, Map rawEvent, long timestampMillis) { + Map eventJson = new HashMap<>(rawEvent); + eventJson.put("invocationId", apiEvent.get("invocationId")); + eventJson.put("author", apiEvent.get("author")); + eventJson.put("timestamp", timestampMillis); + // As in Python ADK 2.x, keep the stored id so a reloaded event matches the streamed one. + if (!(eventJson.get("id") instanceof String id) || id.isEmpty()) { + eventJson.put("id", eventIdFromName(apiEvent)); + } + if (eventJson.get("customMetadata") instanceof Map customMetadata) { + eventJson.put("customMetadata", customMetadataFromJson(customMetadata)); + } + if (eventJson.get("actions") instanceof Map actions + && actions.containsKey("compaction")) { + Map actionsJson = new HashMap<>(actions); + actionsJson.put("compaction", compactionWithMillis(actionsJson.get("compaction"))); + eventJson.put("actions", actionsJson); + } + return storedJsonMapper.convertValue(eventJson, Event.class); + } + + /** Reads the key-value {@code customMetadata} object, skipping the keys Python ADK reserves. */ + private static @Nullable ImmutableList customMetadataFromJson(Map json) { + ImmutableList entries = + json.entrySet().stream() + .filter( + e -> !USAGE_METADATA_KEY.equals(e.getKey()) && !COMPACTION_KEY.equals(e.getKey())) + .map(e -> customMetadataEntry(String.valueOf(e.getKey()), e.getValue())) + .collect(toImmutableList()); + return entries.isEmpty() ? null : entries; + } + + private static CustomMetadata customMetadataEntry(String key, @Nullable Object value) { + CustomMetadata.Builder entry = CustomMetadata.builder().key(key); + if (value instanceof String string) { + entry.stringValue(string); + } else if (value instanceof Number number) { + entry.numericValue(number.floatValue()); + } else if (value instanceof List list) { + entry.stringListValue( + StringList.builder() + .values( + list.stream() + .map( + item -> + item instanceof String string + ? string + : JsonBaseModel.toJsonString(item)) + .collect(toImmutableList())) + .build()); + } else if (value != null) { + // CustomMetadata has no boolean or object value, so keep their JSON text. + entry.stringValue(JsonBaseModel.toJsonString(value)); + } + return entry.build(); + } + + /** Converts a compaction's epoch-second timestamps, as Python ADK writes them, to millis. */ + private static @Nullable Object compactionWithMillis(@Nullable Object compaction) { + if (!(compaction instanceof Map json) + || !(json.get("startTimestamp") instanceof Number start) + || !(json.get("endTimestamp") instanceof Number end)) { + return compaction; + } + Map converted = new HashMap<>(json); + converted.put("startTimestamp", secondsToMillis(start)); + converted.put("endTimestamp", secondsToMillis(end)); + return converted; + } + + /** Camel-cases the snake_case keys Python ADK writes in {@code customMetadata} values. */ + private static @Nullable Object snakeToCamelKeys(@Nullable Object value) { + if (value instanceof Map map) { + Map converted = new LinkedHashMap<>(); + for (Map.Entry entry : map.entrySet()) { + String key = String.valueOf(entry.getKey()); + Object nested = entry.getValue(); + converted.put( + snakeToCamel(key), USER_DATA_KEYS.contains(key) ? nested : snakeToCamelKeys(nested)); + } + return converted; + } + if (value instanceof List list) { + List converted = new ArrayList<>(); + for (Object nested : list) { + converted.add(snakeToCamelKeys(nested)); + } + return converted; + } + return value; + } + + private static String snakeToCamel(String key) { + // CaseFormat would lowercase a key that is already camelCase. + return key.indexOf('_') < 0 ? key : CaseFormat.LOWER_UNDERSCORE.to(CaseFormat.LOWER_CAMEL, key); + } + + /** Returns the value of {@code key} in the typed actions, or else in {@code rawEvent.actions}. */ + private static @Nullable Object actionsValue( + @Nullable Map actions, Map rawActions, String key) { + Object value = actions == null ? null : actions.get(key); + return value != null ? value : rawActions.get(key); + } + + private static String eventIdFromName(Map apiEvent) { + return Iterables.getLast(Splitter.on('/').split(apiEvent.get("name").toString())); + } + + /** Converts a stored value, or returns null so that one bad value cannot fail a load. */ + private static @Nullable T convertStoredValue( + @Nullable Object value, TypeReference type, String field) { + if (value == null) { + return null; + } + try { + return storedJsonMapper.convertValue(value, type); + } catch (IllegalArgumentException e) { + logger.warn( + "Ignoring unreadable '{}' ({})", + field, + Throwables.getRootCause(e).getClass().getSimpleName()); + return null; + } + } + + @SuppressWarnings("unchecked") // Parsing raw Map from JSON following a known schema. + private static Event fromTypedFields(Map apiEvent) { + Map eventMetadata = (Map) apiEvent.get("eventMetadata"); + Map customMetadata = + eventMetadata != null && eventMetadata.get("customMetadata") instanceof Map json + ? (Map) json + : ImmutableMap.of(); + EventActions.Builder eventActionsBuilder = EventActions.builder(); Map actionsMap = (Map) apiEvent.get("actions"); if (actionsMap != null) { @@ -194,41 +448,44 @@ static Event fromApiEvent(Map apiEvent) { Optional.ofNullable(actionsMap.get("requestedAuthConfigs")) .map(SessionJsonConverter::asConcurrentMapOfConcurrentMaps) .orElse(new ConcurrentHashMap<>())); - eventActionsBuilder.requestedToolConfirmations( - Optional.ofNullable(actionsMap.get("requestedToolConfirmations")) - .map(SessionJsonConverter::asConcurrentMapOfToolConfirmations) - .orElse(new ConcurrentHashMap<>())); } - - // Python's Vertex session service leaves agentState and endOfAgent out of "actions" and keeps - // them only under "rawEvent", so a session it wrote would otherwise load with no checkpoints. - Map rawActions = null; - Object rawEvent = apiEvent.get("rawEvent"); - if (rawEvent instanceof Map) { - Object actions = ((Map) rawEvent).get("actions"); - if (actions instanceof Map) { - rawActions = (Map) actions; - } + // The API drops these typed fields; an unreadable rawEvent may still carry them. + Map rawActions = + apiEvent.get("rawEvent") instanceof Map rawEvent + && rawEvent.get("actions") instanceof Map actions + ? actions + : ImmutableMap.of(); + Map toolConfirmations = + convertStoredValue( + actionsValue(actionsMap, rawActions, "requestedToolConfirmations"), + new TypeReference>() {}, + "requestedToolConfirmations"); + if (toolConfirmations != null) { + eventActionsBuilder.requestedToolConfirmations( + Maps.filterValues(toolConfirmations, Objects::nonNull)); } - // Resolved per key, not as a whole map: a rawEvent carrying only one of the two must not - // shadow the other where "actions" does carry it. - Object endOfAgent = checkpointValue(actionsMap, rawActions, "endOfAgent"); + Object endOfAgent = actionsValue(actionsMap, rawActions, "endOfAgent"); if (endOfAgent instanceof Boolean value) { eventActionsBuilder.endOfAgent(value); } else if (endOfAgent != null) { logger.warn("Ignoring 'endOfAgent' of unexpected type {}", endOfAgent.getClass().getName()); } - Object agentState = checkpointValue(actionsMap, rawActions, "agentState"); - if (agentState instanceof Map) { - eventActionsBuilder.agentState((Map) agentState); + Object agentState = actionsValue(actionsMap, rawActions, "agentState"); + if (agentState instanceof Map state) { + eventActionsBuilder.agentState((Map) state); } else if (agentState != null) { // Drop a foreign agentState rather than fail the whole session load on a bad cast. logger.warn("Ignoring 'agentState' of unexpected type {}", agentState.getClass().getName()); } + eventActionsBuilder.compaction( + convertStoredValue( + compactionWithMillis(snakeToCamelKeys(customMetadata.get(COMPACTION_KEY))), + new TypeReference() {}, + COMPACTION_KEY)); Event event = Event.builder() - .id((String) Iterables.getLast(Splitter.on('/').split(apiEvent.get("name").toString()))) + .id(eventIdFromName(apiEvent)) .invocationId((String) apiEvent.get("invocationId")) .author((String) apiEvent.get("author")) .actions(eventActionsBuilder.build()) @@ -247,7 +504,6 @@ static Event fromApiEvent(Map apiEvent) { .map(value -> (String) value) .orElse(null)) .build(); - Map eventMetadata = (Map) apiEvent.get("eventMetadata"); if (eventMetadata != null) { List longRunningToolIdsList = (List) eventMetadata.get("longRunningToolIds"); @@ -257,12 +513,14 @@ static Event fromApiEvent(Map apiEvent) { groundingMetadata = objectMapper.convertValue(rawGroundingMetadata, GroundingMetadata.class); } - GenerateContentResponseUsageMetadata usageMetadata = null; - Object rawUsageMetadata = eventMetadata.get("usageMetadata"); - if (rawUsageMetadata != null) { - usageMetadata = - objectMapper.convertValue(rawUsageMetadata, GenerateContentResponseUsageMetadata.class); - } + Object storedUsageMetadata = customMetadata.get(USAGE_METADATA_KEY); + GenerateContentResponseUsageMetadata usageMetadata = + convertStoredValue( + storedUsageMetadata != null + ? snakeToCamelKeys(storedUsageMetadata) + : eventMetadata.get("usageMetadata"), + new TypeReference() {}, + "usageMetadata"); event = event.toBuilder() @@ -274,6 +532,7 @@ static Event fromApiEvent(Map apiEvent) { .branch((String) eventMetadata.get("branch")) .groundingMetadata(groundingMetadata) .usageMetadata(usageMetadata) + .customMetadata(customMetadataFromJson(customMetadata)) .longRunningToolIds( longRunningToolIdsList != null ? new HashSet<>(longRunningToolIdsList) : null) .build(); @@ -367,20 +626,6 @@ private static ConcurrentMap convertToArtifactDeltaMap(Object a ConcurrentHashMap::putAll); } - @SuppressWarnings("unchecked") // Parsing raw Map from JSON following a known schema. - private static ConcurrentMap asConcurrentMapOfToolConfirmations( - Object value) { - return ((Map) value) - .entrySet().stream() - .collect( - ConcurrentHashMap::new, - (map, entry) -> - map.put( - entry.getKey(), - objectMapper.convertValue(entry.getValue(), ToolConfirmation.class)), - ConcurrentHashMap::putAll); - } - private static void putIfNotEmpty(Map map, String key, Map values) { if (values != null && !values.isEmpty()) { map.put(key, values); @@ -393,4 +638,27 @@ private static void putIfNotEmpty( map.put(key, values); } } + + /** A {@code byte[]} deserializer for standard base64 and the URL-safe base64 of Python ADK. */ + private static final class LenientBase64Deserializer extends StdDeserializer { + LenientBase64Deserializer() { + super(byte[].class); + } + + @Override + public byte[] deserialize(JsonParser parser, DeserializationContext context) + throws IOException { + String text = parser.getValueAsString(); + if (text == null) { + return (byte[]) context.handleUnexpectedToken(byte[].class, parser); + } + try { + // Python ADK writes the URL-safe alphabet; map it to the standard one. + return Base64.getDecoder().decode(text.replace('-', '+').replace('_', '/')); + } catch (IllegalArgumentException e) { + // Leaves the value out of the message: it is event content. + throw JsonMappingException.from(parser, "Invalid base64 value", e); + } + } + } } diff --git a/core/src/test/java/com/google/adk/sessions/MockApiAnswer.java b/core/src/test/java/com/google/adk/sessions/MockApiAnswer.java index 496737553..d93eb02f6 100644 --- a/core/src/test/java/com/google/adk/sessions/MockApiAnswer.java +++ b/core/src/test/java/com/google/adk/sessions/MockApiAnswer.java @@ -22,6 +22,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.adk.JsonBaseModel; import com.google.adk.events.Event; +import com.google.common.collect.ImmutableSet; import java.io.IOException; import java.net.URLDecoder; import java.time.Instant; @@ -57,6 +58,27 @@ class MockApiAnswer implements Answer { private static final MediaType JSON_MEDIA_TYPE = MediaType.parse("application/json; charset=utf-8"); + // Like the real API, appendEvent keeps only these fields and drops any others without an error. + private static final ImmutableSet STORED_ACTIONS_FIELDS = + ImmutableSet.of( + "skipSummarization", + "stateDelta", + "artifactDelta", + "escalate", + "requestedAuthConfigs", + "transferAgent"); + private static final ImmutableSet STORED_EVENT_METADATA_FIELDS = + ImmutableSet.of( + "groundingMetadata", + "partial", + "turnComplete", + "interrupted", + "longRunningToolIds", + "branch", + "customMetadata", + "inputTranscription", + "outputTranscription"); + /** The id this fake mints when the caller requests none. */ static final String GENERATED_SESSION_ID = "4"; @@ -244,6 +266,10 @@ private ApiResponse handleAppendEvent(String path, InvocationOnMock invocation) eventDataString, new TypeReference>>() {})); } + extractObjectMap(newEventData, "actions") + .ifPresent(actions -> actions.keySet().retainAll(STORED_ACTIONS_FIELDS)); + extractObjectMap(newEventData, "eventMetadata") + .ifPresent(metadata -> metadata.keySet().retainAll(STORED_EVENT_METADATA_FIELDS)); newEventData.put( "name", path.replaceFirst(":appendEvent$", "/events/" + Event.generateEventId())); diff --git a/core/src/test/java/com/google/adk/sessions/SessionJsonConverterTest.java b/core/src/test/java/com/google/adk/sessions/SessionJsonConverterTest.java index d52ef93da..ddfde5d11 100644 --- a/core/src/test/java/com/google/adk/sessions/SessionJsonConverterTest.java +++ b/core/src/test/java/com/google/adk/sessions/SessionJsonConverterTest.java @@ -20,23 +20,30 @@ import static org.junit.Assert.assertThrows; import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; import com.google.adk.JsonBaseModel; import com.google.adk.events.Event; import com.google.adk.events.EventActions; +import com.google.adk.events.EventCompaction; import com.google.adk.events.ToolConfirmation; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; +import com.google.genai.types.Blob; import com.google.genai.types.Content; +import com.google.genai.types.CustomMetadata; import com.google.genai.types.FinishReason; import com.google.genai.types.GenerateContentResponseUsageMetadata; import com.google.genai.types.GroundingMetadata; import com.google.genai.types.Part; +import com.google.genai.types.StringList; +import java.io.UncheckedIOException; import java.time.Instant; import java.util.Collections; import java.util.HashMap; +import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; @@ -47,6 +54,12 @@ @RunWith(JUnit4.class) public final class SessionJsonConverterTest { private static final ObjectMapper objectMapper = JsonBaseModel.getMapper(); + private static final EventCompaction COMPACTION = + EventCompaction.builder() + .startTimestamp(1_700_000_000_100L) + .endTimestamp(1_700_000_000_200L) + .compactedContent(Content.fromParts(Part.fromText("summary"))) + .build(); @Test public void convertEventToJson_fullEvent_success() throws JsonProcessingException { @@ -225,18 +238,33 @@ public void convertEventToJson_complexActions_success() throws JsonProcessingExc JsonNode eventMetadata = jsonNode.get("eventMetadata"); assertThat(eventMetadata.get("longRunningToolIds").get(0).asText()).isEqualTo("tool-id-1"); - assertThat(eventMetadata.has("usageMetadata")).isTrue(); assertThat(eventMetadata.has("groundingMetadata")).isTrue(); + // The API drops the typed usageMetadata field, so it is stored in customMetadata. + assertThat(eventMetadata.has("usageMetadata")).isFalse(); + assertThat( + eventMetadata + .get("customMetadata") + .get("_usage_metadata") + .get("promptTokenCount") + .asInt()) + .isEqualTo(10); JsonNode actionsNode = jsonNode.get("actions"); assertThat(actionsNode.get("requestedAuthConfigs").get("auth1").get("param1").asText()) .isEqualTo("value1"); - assertThat(actionsNode.get("requestedToolConfirmations").get("tool1").get("hint").asText()) + // The API drops these too; only rawEvent carries them. + assertThat(actionsNode.has("requestedToolConfirmations")).isFalse(); + assertThat(actionsNode.has("endOfAgent")).isFalse(); + + JsonNode rawActions = jsonNode.get("rawEvent").get("actions"); + assertThat(rawActions.get("requestedToolConfirmations").get("tool1").get("hint").asText()) .isEqualTo("hint1"); assertThat( - actionsNode.get("requestedToolConfirmations").get("tool1").get("confirmed").asBoolean()) + rawActions.get("requestedToolConfirmations").get("tool1").get("confirmed").asBoolean()) .isTrue(); - assertThat(actionsNode.get("endOfAgent").asBoolean()).isTrue(); + assertThat(rawActions.get("endOfAgent").asBoolean()).isTrue(); + assertThat(jsonNode.get("rawEvent").get("usageMetadata").get("promptTokenCount").asInt()) + .isEqualTo(10); } @Test @@ -291,32 +319,46 @@ public void convertEventToJson_agentState_success() throws JsonProcessingExcepti .build(); String json = SessionJsonConverter.convertEventToJson(event, true); - JsonNode actionsNode = objectMapper.readTree(json).get("actions"); + JsonNode jsonNode = objectMapper.readTree(json); + JsonNode rawActions = jsonNode.get("rawEvent").get("actions"); - assertThat(actionsNode.get("agentState").get("current_sub_agent").asText()) - .isEqualTo("b_agent"); - assertThat(actionsNode.get("agentState").get("times_looped").asInt()).isEqualTo(2); + assertThat(jsonNode.get("actions").has("agentState")).isFalse(); + assertThat(rawActions.get("agentState").get("current_sub_agent").asText()).isEqualTo("b_agent"); + assertThat(rawActions.get("agentState").get("times_looped").asInt()).isEqualTo(2); } - // rawEvent supplies only the keys "actions" lacks: a rawEvent carrying endOfAgent alone must not - // shadow an agentState that actions does carry. + // As in Python ADK, a readable rawEvent replaces the typed fields as a whole. @Test - public void fromApiEvent_rawEventMissingAgentState_doesNotShadowActions() { + public void fromApiEvent_rawEventPresent_takesPrecedenceOverTypedFields() { Map apiEvent = new HashMap<>(); apiEvent.put("name", "sessions/123/events/456"); apiEvent.put("invocationId", "inv-1"); apiEvent.put("author", "agent"); apiEvent.put("timestamp", "2023-01-01T00:00:00.123Z"); - Map actions = new HashMap<>(); - actions.put("agentState", ImmutableMap.of("current_sub_agent", "b_agent")); - apiEvent.put("actions", actions); - apiEvent.put("rawEvent", ImmutableMap.of("actions", ImmutableMap.of("endOfAgent", true))); + apiEvent.put( + "content", ImmutableMap.of("parts", ImmutableList.of(ImmutableMap.of("text", "a")))); + apiEvent.put("actions", new HashMap<>(ImmutableMap.of("transferAgent", "typed-agent"))); + apiEvent.put("eventMetadata", new HashMap<>(ImmutableMap.of("branch", "typed-branch"))); + apiEvent.put( + "rawEvent", + ImmutableMap.of( + "content", + ImmutableMap.of("parts", ImmutableList.of(ImmutableMap.of("text", "b"))), + "actions", + ImmutableMap.of("transferToAgent", "raw-agent"), + "branch", + "raw-branch")); Event event = SessionJsonConverter.fromApiEvent(apiEvent); - assertThat(event.actions().agentState()).isPresent(); - assertThat(event.actions().agentState().get()).containsEntry("current_sub_agent", "b_agent"); - assertThat(event.actions().endOfAgent()).isTrue(); + assertThat(event.id()).isEqualTo("456"); + assertThat(event.content().get().text()).isEqualTo("b"); + assertThat(event.actions().transferToAgent()).hasValue("raw-agent"); + assertThat(event.branch()).hasValue("raw-branch"); + assertThat(event.invocationId()).isEqualTo("inv-1"); + assertThat(event.author()).isEqualTo("agent"); + assertThat(event.timestamp()) + .isEqualTo(Instant.parse("2023-01-01T00:00:00.123Z").toEpochMilli()); } // Python's Vertex session service keeps agentState and endOfAgent only under rawEvent, so a @@ -540,4 +582,691 @@ public void fromApiEvent_withNullStateDeltaValue_success() { assertThat(eventActions.stateDelta()).containsEntry("key1", "value1"); assertThat(eventActions.stateDelta()).containsEntry("key2", State.REMOVED); } + + @Test + public void convertEventToJson_customMetadata_writtenUnderPythonKeys() + throws JsonProcessingException { + Event event = + Event.builder() + .author("agent") + .invocationId("inv-1") + .timestamp(1_700_000_000_123L) + .customMetadata( + ImmutableList.of( + CustomMetadata.builder().key("str").stringValue("v").build(), + CustomMetadata.builder().key("num").numericValue(1.5f).build(), + CustomMetadata.builder() + .key("list") + .stringListValue( + StringList.builder().values(ImmutableList.of("a", "b")).build()) + .build())) + .usageMetadata( + GenerateContentResponseUsageMetadata.builder().totalTokenCount(12).build()) + .actions(EventActions.builder().compaction(COMPACTION).build()) + .build(); + + JsonNode customMetadata = + objectMapper + .readTree(SessionJsonConverter.convertEventToJson(event)) + .get("eventMetadata") + .get("customMetadata"); + + assertThat(customMetadata.get("str").asText()).isEqualTo("v"); + assertThat(customMetadata.get("num").asDouble()).isEqualTo(1.5); + assertThat(customMetadata.get("list").get(1).asText()).isEqualTo("b"); + assertThat(customMetadata.get("_usage_metadata").get("totalTokenCount").asInt()).isEqualTo(12); + JsonNode compaction = customMetadata.get("_compaction"); + assertThat(compaction.get("startTimestamp").asDouble()).isWithin(1e-6).of(1_700_000_000.1); + assertThat(compaction.get("endTimestamp").asDouble()).isWithin(1e-6).of(1_700_000_000.2); + assertThat(compaction.get("compactedContent").get("parts").get(0).get("text").asText()) + .isEqualTo("summary"); + } + + @Test + public void convertEventToJson_rawEvent_hasShapePythonAdkReads() throws JsonProcessingException { + EventActions actions = + EventActions.builder() + .stateDelta( + new ConcurrentHashMap<>(ImmutableMap.of("kept", "v", "removed", State.REMOVED))) + .compaction(COMPACTION) + .build(); + Event event = + Event.builder() + .id("event-1") + .author("agent") + .invocationId("inv-1") + .timestamp(1_700_000_000_123L) + .customMetadata( + ImmutableList.of(CustomMetadata.builder().key("k").stringValue("v").build())) + .actions(actions) + .build(); + + JsonNode rawEvent = + objectMapper.readTree(SessionJsonConverter.convertEventToJson(event)).get("rawEvent"); + + assertThat(rawEvent.get("id").asText()).isEqualTo("event-1"); + assertThat(rawEvent.get("timestamp").asDouble()).isWithin(1e-6).of(1_700_000_000.123); + assertThat(rawEvent.get("customMetadata").get("k").asText()).isEqualTo("v"); + JsonNode rawActions = rawEvent.get("actions"); + assertThat(rawActions.get("stateDelta").get("kept").asText()).isEqualTo("v"); + assertThat(rawActions.get("stateDelta").get("removed").isNull()).isTrue(); + assertThat(rawActions.get("compaction").get("startTimestamp").asDouble()) + .isWithin(1e-6) + .of(1_700_000_000.1); + } + + // Python ADK fails to load the whole session on any of these. + @Test + public void convertEventToJson_rawEvent_leavesOutFieldsPythonAdkRejects() + throws JsonProcessingException { + EventActions actions = + EventActions.builder() + .requestedToolConfirmations( + new ConcurrentHashMap<>( + ImmutableMap.of("call-1", ToolConfirmation.builder().hint(null).build()))) + .deletedArtifactIds(ImmutableSet.of("artifact")) + .setModelResponse(ImmutableMap.of("k", "v")) + .build(); + Event event = + Event.builder() + .author("agent") + .invocationId("inv-1") + .timestamp(1L) + .actions(actions) + .build(); + + JsonNode rawEvent = + objectMapper.readTree(SessionJsonConverter.convertEventToJson(event)).get("rawEvent"); + + assertThat(rawEvent.has("id")).isFalse(); + JsonNode rawActions = rawEvent.get("actions"); + assertThat(rawActions.get("requestedToolConfirmations").get("call-1").has("hint")).isFalse(); + assertThat(rawActions.has("deletedArtifactIds")).isFalse(); + assertThat(rawActions.has("setModelResponse")).isFalse(); + } + + @Test + public void convertEventToJson_rawEvent_keepsNullMapValues() throws JsonProcessingException { + Map agentState = new HashMap<>(); + agentState.put("unset", null); + Event event = + Event.builder() + .author("agent") + .invocationId("inv-1") + .timestamp(1L) + .actions(EventActions.builder().agentState(agentState).build()) + .build(); + + JsonNode rawEvent = + objectMapper.readTree(SessionJsonConverter.convertEventToJson(event)).get("rawEvent"); + + assertThat(rawEvent.get("actions").get("agentState").get("unset").isNull()).isTrue(); + } + + @Test + public void convertEventToJson_thenFromApiEvent_roundTripsWholeEvent() throws Exception { + Event event = + Event.builder() + .id("event-1") + .author("agent") + .invocationId("inv-1") + .timestamp(1_700_000_000_123L) + .content(Content.fromParts(Part.fromText("hi"))) + .customMetadata( + ImmutableList.of(CustomMetadata.builder().key("k").stringValue("v").build())) + .usageMetadata( + GenerateContentResponseUsageMetadata.builder() + .promptTokenCount(10) + .totalTokenCount(12) + .build()) + .finishReason(new FinishReason("STOP")) + .modelVersion("model-1") + .turnComplete(true) + .actions( + EventActions.builder() + .stateDelta( + new ConcurrentHashMap<>(ImmutableMap.of("k", "v", "gone", State.REMOVED))) + .transferToAgent("other") + .requestedToolConfirmations( + new ConcurrentHashMap<>( + ImmutableMap.of( + "call-1", + ToolConfirmation.builder().hint("ok?").confirmed(true).build()))) + .endOfAgent(true) + .agentState(ImmutableMap.of("step", "2")) + .compaction(COMPACTION) + .build()) + .build(); + + Event reloaded = SessionJsonConverter.fromApiEvent(asStoredEvent(event)); + + assertThat(reloaded).isEqualTo(event); + } + + @Test + public void convertEventToJson_thenFromApiEvent_roundTripsInlineBytes() throws Exception { + byte[] data = {(byte) 0xfb, (byte) 0xff, 0x01}; + Event event = + Event.builder() + .author("user") + .invocationId("inv-1") + .timestamp(1_700_000_000_123L) + .content(Content.fromParts(Part.fromBytes(data, "image/png"))) + .build(); + + Blob reloaded = + SessionJsonConverter.fromApiEvent(asStoredEvent(event)) + .content() + .get() + .parts() + .get() + .get(0) + .inlineData() + .get(); + + assertThat(reloaded.data().get()).isEqualTo(data); + // Only rawEvent keeps the MIME type, so this proves the bytes came from it. + assertThat(reloaded.mimeType()).hasValue("image/png"); + } + + @Test + public void fromApiEvent_pythonWrittenRawEvent_readsFieldsJavaModels() { + // As Python ADK dumps it: float seconds, URL-safe base64 and keys Java does not model. + Map rawEvent = new HashMap<>(); + rawEvent.put("id", "python-event-id"); + rawEvent.put("invocationId", "inv-1"); + rawEvent.put("author", "agent"); + rawEvent.put("timestamp", 1_700_000_000.5); + rawEvent.put("nodeInfo", ImmutableMap.of("path", "")); + rawEvent.put( + "content", + ImmutableMap.of( + "role", + "model", + "parts", + ImmutableList.of(ImmutableMap.of("text", "hi", "thoughtSignature", "-_8B")))); + rawEvent.put("usageMetadata", ImmutableMap.of("promptTokenCount", 10, "totalTokenCount", 12)); + rawEvent.put( + "customMetadata", + ImmutableMap.of( + "flag", + true, + "k", + "v", + "n", + 2.5, + "tags", + ImmutableList.of("a", ImmutableMap.of("b", 1)))); + rawEvent.put( + "actions", + ImmutableMap.of( + "requestedToolConfirmations", + ImmutableMap.of("call-1", ImmutableMap.of("hint", "ok?", "confirmed", false)), + "compaction", + ImmutableMap.of( + "startTimestamp", + 1_700_000_000.1, + "endTimestamp", + 1_700_000_000.2, + "compactedContent", + ImmutableMap.of( + "role", "model", "parts", ImmutableList.of(ImmutableMap.of("text", "sum")))))); + Map apiEvent = new HashMap<>(); + apiEvent.put("name", "sessions/123/events/server-id"); + apiEvent.put("invocationId", "inv-1"); + apiEvent.put("author", "agent"); + apiEvent.put("timestamp", "2023-11-14T22:13:20.500Z"); + apiEvent.put("rawEvent", rawEvent); + + Event event = SessionJsonConverter.fromApiEvent(apiEvent); + + assertThat(event.id()).isEqualTo("python-event-id"); + assertThat(event.timestamp()).isEqualTo(1_700_000_000_500L); + assertThat(event.content().get().parts().get().get(0).thoughtSignature().get()) + .isEqualTo(new byte[] {(byte) 0xfb, (byte) 0xff, 0x01}); + assertThat(event.usageMetadata().get().totalTokenCount()).hasValue(12); + assertThat(event.customMetadata().get()) + .containsExactly( + CustomMetadata.builder().key("flag").stringValue("true").build(), + CustomMetadata.builder().key("k").stringValue("v").build(), + CustomMetadata.builder().key("n").numericValue(2.5f).build(), + CustomMetadata.builder() + .key("tags") + .stringListValue( + StringList.builder().values(ImmutableList.of("a", "{\"b\":1}")).build()) + .build()); + assertThat(event.actions().requestedToolConfirmations().get("call-1").hint()).isEqualTo("ok?"); + assertThat(event.actions().compaction().get().startTimestamp()).isEqualTo(1_700_000_000_100L); + assertThat(event.actions().compaction().get().endTimestamp()).isEqualTo(1_700_000_000_200L); + } + + // Python timestamps carry microseconds, which event timestamps drop by flooring to millis. + @Test + public void fromApiEvent_pythonCompactionWithMicroseconds_floorsToMillis() { + Map apiEvent = new HashMap<>(); + apiEvent.put("name", "sessions/123/events/456"); + apiEvent.put("invocationId", "inv-1"); + apiEvent.put("author", "user"); + apiEvent.put("timestamp", "2023-01-01T00:00:00Z"); + apiEvent.put( + "rawEvent", + ImmutableMap.of( + "actions", + ImmutableMap.of( + "compaction", + ImmutableMap.of( + "startTimestamp", + 1_700_000_000.1236, + "endTimestamp", + 1_700_000_000.9999, + "compactedContent", + ImmutableMap.of("parts", ImmutableList.of(ImmutableMap.of("text", "s"))))))); + + EventCompaction compaction = + SessionJsonConverter.fromApiEvent(apiEvent).actions().compaction().get(); + + assertThat(compaction.startTimestamp()).isEqualTo(1_700_000_000_123L); + assertThat(compaction.endTimestamp()).isEqualTo(1_700_000_000_999L); + } + + // Python ADK stores an event without raw_event when its Vertex AI SDK rejects that field. + @Test + public void fromApiEvent_pythonEventWithoutRawEvent_readsCustomMetadataKeys() { + Map customMetadata = new HashMap<>(); + customMetadata.put("k", "v"); + customMetadata.put( + "_usage_metadata", + ImmutableMap.of( + "prompt_token_count", + 10, + "total_token_count", + 12, + "prompt_tokens_details", + ImmutableList.of(ImmutableMap.of("modality", "TEXT", "token_count", 10)))); + customMetadata.put( + "_compaction", + ImmutableMap.of( + "start_timestamp", + 1_700_000_000.1, + "end_timestamp", + 1_700_000_000.2, + "compacted_content", + ImmutableMap.of( + "role", "model", "parts", ImmutableList.of(ImmutableMap.of("text", "summary"))))); + Map apiEvent = new HashMap<>(); + apiEvent.put("name", "sessions/123/events/456"); + apiEvent.put("invocationId", "inv-1"); + apiEvent.put("author", "agent"); + apiEvent.put("timestamp", "2023-01-01T00:00:00Z"); + apiEvent.put("eventMetadata", new HashMap<>(ImmutableMap.of("customMetadata", customMetadata))); + + Event event = SessionJsonConverter.fromApiEvent(apiEvent); + + assertThat(event.id()).isEqualTo("456"); + assertThat(event.usageMetadata().get().promptTokenCount()).hasValue(10); + assertThat(event.usageMetadata().get().promptTokensDetails().get().get(0).tokenCount()) + .hasValue(10); + assertThat(event.actions().compaction().get().startTimestamp()).isEqualTo(1_700_000_000_100L); + assertThat(event.actions().compaction().get().compactedContent().text()).isEqualTo("summary"); + assertThat(event.customMetadata().get()) + .containsExactly(CustomMetadata.builder().key("k").stringValue("v").build()); + } + + @Test + public void fromApiEvent_pythonCompactionWithUserData_keepsItsKeys() { + Map functionCall = + ImmutableMap.of("name", "lookup", "args", ImmutableMap.of("user_id", "u1")); + Map functionResponse = + ImmutableMap.of("name", "lookup", "response", ImmutableMap.of("user_name", "Ann")); + Map compaction = + ImmutableMap.of( + "start_timestamp", + 1_700_000_000.1, + "end_timestamp", + 1_700_000_000.2, + "compacted_content", + ImmutableMap.of( + "role", + "model", + "parts", + ImmutableList.of( + ImmutableMap.of("function_call", functionCall), + ImmutableMap.of("function_response", functionResponse), + ImmutableMap.of( + "text", "summary", "part_metadata", ImmutableMap.of("user_tag", 1))))); + Map apiEvent = new HashMap<>(); + apiEvent.put("name", "sessions/123/events/456"); + apiEvent.put("invocationId", "inv-1"); + apiEvent.put("author", "agent"); + apiEvent.put("timestamp", "2023-01-01T00:00:00Z"); + apiEvent.put( + "eventMetadata", + new HashMap<>( + ImmutableMap.of("customMetadata", ImmutableMap.of("_compaction", compaction)))); + + List parts = + SessionJsonConverter.fromApiEvent(apiEvent) + .actions() + .compaction() + .get() + .compactedContent() + .parts() + .get(); + + assertThat(parts.get(0).functionCall().get().args().get()).containsExactly("user_id", "u1"); + assertThat(parts.get(1).functionResponse().get().response().get()) + .containsExactly("user_name", "Ann"); + assertThat(parts.get(2).partMetadata().get()).containsExactly("user_tag", 1); + } + + // Java writes _compaction with camelCase keys. + @Test + public void fromApiEvent_javaCompactionWithPartMetadata_keepsItsKeys() { + Map compaction = + ImmutableMap.of( + "startTimestamp", + 1_700_000_000.1, + "endTimestamp", + 1_700_000_000.2, + "compactedContent", + ImmutableMap.of( + "parts", + ImmutableList.of( + ImmutableMap.of( + "text", "summary", "partMetadata", ImmutableMap.of("user_tag", 1))))); + Map apiEvent = new HashMap<>(); + apiEvent.put("name", "sessions/123/events/456"); + apiEvent.put("invocationId", "inv-1"); + apiEvent.put("author", "agent"); + apiEvent.put("timestamp", "2023-01-01T00:00:00Z"); + apiEvent.put( + "eventMetadata", + new HashMap<>( + ImmutableMap.of("customMetadata", ImmutableMap.of("_compaction", compaction)))); + + Part part = + SessionJsonConverter.fromApiEvent(apiEvent) + .actions() + .compaction() + .get() + .compactedContent() + .parts() + .get() + .get(0); + + assertThat(part.partMetadata().get()).containsExactly("user_tag", 1); + } + + @Test + public void fromApiEvent_unreadableRawEvent_fallsBackToTypedFields() { + Map apiEvent = new HashMap<>(); + apiEvent.put("name", "sessions/123/events/456"); + apiEvent.put("invocationId", "inv-1"); + apiEvent.put("author", "agent"); + apiEvent.put("timestamp", "2023-01-01T00:00:00Z"); + apiEvent.put("actions", new HashMap<>(ImmutableMap.of("transferAgent", "typed-agent"))); + apiEvent.put("rawEvent", ImmutableMap.of("id", "raw-id", "actions", "not-an-object")); + + Event event = SessionJsonConverter.fromApiEvent(apiEvent); + + assertThat(event.id()).isEqualTo("456"); + assertThat(event.actions().transferToAgent()).hasValue("typed-agent"); + } + + // As in Python ADK, an empty rawEvent counts as missing. + @Test + public void fromApiEvent_emptyRawEvent_usesTypedFields() { + Map apiEvent = new HashMap<>(); + apiEvent.put("name", "sessions/123/events/456"); + apiEvent.put("invocationId", "inv-1"); + apiEvent.put("author", "agent"); + apiEvent.put("timestamp", "2023-01-01T00:00:00Z"); + apiEvent.put( + "content", ImmutableMap.of("parts", ImmutableList.of(ImmutableMap.of("text", "typed")))); + apiEvent.put("actions", new HashMap<>(ImmutableMap.of("transferAgent", "typed-agent"))); + apiEvent.put("rawEvent", ImmutableMap.of()); + + Event event = SessionJsonConverter.fromApiEvent(apiEvent); + + assertThat(event.content().get().text()).isEqualTo("typed"); + assertThat(event.actions().transferToAgent()).hasValue("typed-agent"); + } + + // As in Python ADK, an empty stored id falls back to the resource name. + @Test + public void fromApiEvent_rawEventEmptyId_usesResourceId() { + Map apiEvent = new HashMap<>(); + apiEvent.put("name", "sessions/123/events/456"); + apiEvent.put("invocationId", "inv-1"); + apiEvent.put("author", "agent"); + apiEvent.put("timestamp", "2023-01-01T00:00:00Z"); + apiEvent.put("rawEvent", ImmutableMap.of("id", "", "branch", "raw-branch")); + + Event event = SessionJsonConverter.fromApiEvent(apiEvent); + + assertThat(event.id()).isEqualTo("456"); + // Only rawEvent carries the branch, so this proves rawEvent was read. + assertThat(event.branch()).hasValue("raw-branch"); + } + + @Test + public void fromApiEvent_rawEventWithInvalidBase64_fallsBackToTypedFields() { + Map apiEvent = new HashMap<>(); + apiEvent.put("name", "sessions/123/events/456"); + apiEvent.put("invocationId", "inv-1"); + apiEvent.put("author", "user"); + apiEvent.put("timestamp", "2023-01-01T00:00:00Z"); + apiEvent.put( + "content", ImmutableMap.of("parts", ImmutableList.of(ImmutableMap.of("text", "a")))); + apiEvent.put( + "rawEvent", + ImmutableMap.of( + "id", + "raw-id", + "content", + ImmutableMap.of( + "parts", + ImmutableList.of(ImmutableMap.of("inlineData", ImmutableMap.of("data", "!!!!")))))); + + Event event = SessionJsonConverter.fromApiEvent(apiEvent); + + assertThat(event.id()).isEqualTo("456"); + assertThat(event.content().get().text()).isEqualTo("a"); + } + + @Test + public void fromApiEvent_rawEventWithNonStringBytes_fallsBackToTypedFields() { + Map apiEvent = new HashMap<>(); + apiEvent.put("name", "sessions/123/events/456"); + apiEvent.put("invocationId", "inv-1"); + apiEvent.put("author", "user"); + apiEvent.put("timestamp", "2023-01-01T00:00:00Z"); + apiEvent.put( + "content", ImmutableMap.of("parts", ImmutableList.of(ImmutableMap.of("text", "a")))); + apiEvent.put( + "rawEvent", + ImmutableMap.of( + "id", + "raw-id", + "content", + ImmutableMap.of( + "parts", + ImmutableList.of( + ImmutableMap.of( + "inlineData", ImmutableMap.of("data", ImmutableList.of(1, 2))))))); + + Event event = SessionJsonConverter.fromApiEvent(apiEvent); + + assertThat(event.id()).isEqualTo("456"); + assertThat(event.content().get().text()).isEqualTo("a"); + } + + @Test + public void fromApiEvent_unreadableRawEvent_keepsFieldsTheApiDropsFromRawActions() { + Map apiEvent = new HashMap<>(); + apiEvent.put("name", "sessions/123/events/456"); + apiEvent.put("invocationId", "inv-1"); + apiEvent.put("author", "agent"); + apiEvent.put("timestamp", "2023-01-01T00:00:00Z"); + // Typed actions as the API returns them, without the fields it drops. + apiEvent.put("actions", new HashMap<>(ImmutableMap.of("stateDelta", new HashMap<>()))); + apiEvent.put( + "rawEvent", + ImmutableMap.of( + "id", + "raw-id", + "content", + ImmutableMap.of( + "parts", + ImmutableList.of(ImmutableMap.of("inlineData", ImmutableMap.of("data", "!!!!")))), + "actions", + ImmutableMap.of( + "agentState", + ImmutableMap.of("step", "2"), + "endOfAgent", + true, + "requestedToolConfirmations", + ImmutableMap.of("call-1", ImmutableMap.of("hint", "ok?", "confirmed", false))))); + + Event event = SessionJsonConverter.fromApiEvent(apiEvent); + + // The resource-name id shows that the typed-field fallback ran. + assertThat(event.id()).isEqualTo("456"); + assertThat(event.actions().agentState().get()).containsEntry("step", "2"); + assertThat(event.actions().endOfAgent()).isTrue(); + assertThat(event.actions().requestedToolConfirmations().get("call-1").hint()).isEqualTo("ok?"); + } + + @Test + public void fromApiEvent_unreadableRawEvent_prefersTypedActionsPerKey() { + Map apiEvent = new HashMap<>(); + apiEvent.put("name", "sessions/123/events/456"); + apiEvent.put("invocationId", "inv-1"); + apiEvent.put("author", "agent"); + apiEvent.put("timestamp", "2023-01-01T00:00:00Z"); + apiEvent.put( + "actions", new HashMap<>(ImmutableMap.of("agentState", ImmutableMap.of("step", "typed")))); + apiEvent.put( + "rawEvent", + ImmutableMap.of( + "content", + ImmutableMap.of( + "parts", + ImmutableList.of(ImmutableMap.of("inlineData", ImmutableMap.of("data", "!!!!")))), + "actions", + ImmutableMap.of("agentState", ImmutableMap.of("step", "raw"), "endOfAgent", true))); + + Event event = SessionJsonConverter.fromApiEvent(apiEvent); + + assertThat(event.actions().agentState().get()).containsEntry("step", "typed"); + assertThat(event.actions().endOfAgent()).isTrue(); + } + + @Test + public void fromApiEvent_unreadableToolConfirmations_areSkipped() { + Map toolConfirmations = new HashMap<>(); + toolConfirmations.put("call-1", "not-an-object"); + Event event = fromUnreadableRawEventWithActions(toolConfirmations); + + assertThat(event.actions().requestedToolConfirmations()).isEmpty(); + assertThat(event.actions().agentState()).isPresent(); + } + + @Test + public void fromApiEvent_nullToolConfirmation_isSkipped() { + Map toolConfirmations = new HashMap<>(); + toolConfirmations.put("call-1", null); + toolConfirmations.put("call-2", ImmutableMap.of("hint", "ok?", "confirmed", false)); + Event event = fromUnreadableRawEventWithActions(toolConfirmations); + + assertThat(event.actions().requestedToolConfirmations().keySet()).containsExactly("call-2"); + } + + /** Reads an event whose rawEvent content is unreadable but whose actions carry these values. */ + private static Event fromUnreadableRawEventWithActions(Map toolConfirmations) { + Map apiEvent = new HashMap<>(); + apiEvent.put("name", "sessions/123/events/456"); + apiEvent.put("invocationId", "inv-1"); + apiEvent.put("author", "agent"); + apiEvent.put("timestamp", "2023-01-01T00:00:00Z"); + apiEvent.put( + "rawEvent", + ImmutableMap.of( + "content", + ImmutableMap.of( + "parts", + ImmutableList.of(ImmutableMap.of("inlineData", ImmutableMap.of("data", "!!!!")))), + "actions", + ImmutableMap.of( + "agentState", + ImmutableMap.of("step", "2"), + "requestedToolConfirmations", + toolConfirmations))); + return SessionJsonConverter.fromApiEvent(apiEvent); + } + + @Test + public void fromApiEvent_unreadableStoredUsageMetadata_isSkipped() { + Map apiEvent = new HashMap<>(); + apiEvent.put("name", "sessions/123/events/456"); + apiEvent.put("invocationId", "inv-1"); + apiEvent.put("author", "agent"); + apiEvent.put("timestamp", "2023-01-01T00:00:00Z"); + apiEvent.put( + "eventMetadata", + new HashMap<>( + ImmutableMap.of( + "customMetadata", + ImmutableMap.of("_usage_metadata", ImmutableMap.of("total_token_count", "many"))))); + + Event event = SessionJsonConverter.fromApiEvent(apiEvent); + + assertThat(event.usageMetadata()).isEmpty(); + } + + // Java writes camelCase keys under _usage_metadata; snake_case conversion must leave them as is. + @Test + public void fromApiEvent_camelCaseUsageMetadataWithoutRawEvent_readsIt() { + Map apiEvent = new HashMap<>(); + apiEvent.put("name", "sessions/123/events/456"); + apiEvent.put("invocationId", "inv-1"); + apiEvent.put("author", "agent"); + apiEvent.put("timestamp", "2023-01-01T00:00:00Z"); + apiEvent.put( + "eventMetadata", + new HashMap<>( + ImmutableMap.of( + "customMetadata", + ImmutableMap.of("_usage_metadata", ImmutableMap.of("promptTokenCount", 10))))); + + Event event = SessionJsonConverter.fromApiEvent(apiEvent); + + assertThat(event.usageMetadata().get().promptTokenCount()).hasValue(10); + } + + @Test + public void convertEventToJson_unserializableValue_throwsUncheckedIOException() { + Event event = + Event.builder() + .author("agent") + .invocationId("inv-1") + .timestamp(1L) + .actions( + EventActions.builder() + .stateDelta(new ConcurrentHashMap<>(ImmutableMap.of("k", new Object()))) + .build()) + .build(); + + assertThrows(UncheckedIOException.class, () -> SessionJsonConverter.convertEventToJson(event)); + } + + /** Returns the event as the API stores it, under a server-assigned name. */ + private static Map asStoredEvent(Event event) throws JsonProcessingException { + Map apiEvent = + objectMapper.readValue( + SessionJsonConverter.convertEventToJson(event), + new TypeReference>() {}); + apiEvent.put("name", "sessions/123/events/server-id"); + return apiEvent; + } } diff --git a/core/src/test/java/com/google/adk/sessions/VertexAiSessionServiceTest.java b/core/src/test/java/com/google/adk/sessions/VertexAiSessionServiceTest.java index 1cf7c247d..caab80806 100644 --- a/core/src/test/java/com/google/adk/sessions/VertexAiSessionServiceTest.java +++ b/core/src/test/java/com/google/adk/sessions/VertexAiSessionServiceTest.java @@ -31,10 +31,13 @@ import com.google.adk.JsonBaseModel; import com.google.adk.events.Event; import com.google.adk.events.EventActions; +import com.google.adk.events.ToolConfirmation; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.genai.types.Content; +import com.google.genai.types.CustomMetadata; +import com.google.genai.types.GenerateContentResponseUsageMetadata; import com.google.genai.types.Part; import io.reactivex.rxjava3.core.Completable; import io.reactivex.rxjava3.core.Single; @@ -463,6 +466,50 @@ public void appendEvent_success() { .isEqualTo(Instant.parse("2024-12-12T12:12:12.123456Z").toEpochMilli()); } + @Test + public void appendEvent_thenListEvents_keepsEventFields() { + Session session = vertexAiSessionService.createSession("987", "userA").blockingGet(); + Event event = + Event.builder() + .id("client-event-id") + .invocationId("inv-1") + .author("agent") + .timestamp(Instant.parse("2024-12-12T12:12:12.123Z").toEpochMilli()) + .usageMetadata( + GenerateContentResponseUsageMetadata.builder().totalTokenCount(12).build()) + .customMetadata( + ImmutableList.of(CustomMetadata.builder().key("k").stringValue("v").build())) + .actions( + EventActions.builder() + .transferToAgent("other") + .requestedToolConfirmations( + new ConcurrentHashMap<>( + ImmutableMap.of( + "call-1", ToolConfirmation.builder().hint("ok?").build()))) + .endOfAgent(true) + .agentState(ImmutableMap.of("step", "2")) + .build()) + .build(); + var unused = vertexAiSessionService.appendEvent(session, event).blockingGet(); + + Event reloaded = + vertexAiSessionService + .listEvents(session.appName(), session.userId(), session.id()) + .blockingGet() + .events() + .get(0); + + assertThat(reloaded.id()).isEqualTo("client-event-id"); + assertThat(reloaded.usageMetadata().get().totalTokenCount()).hasValue(12); + assertThat(reloaded.customMetadata().get()) + .containsExactly(CustomMetadata.builder().key("k").stringValue("v").build()); + assertThat(reloaded.actions().transferToAgent()).hasValue("other"); + assertThat(reloaded.actions().requestedToolConfirmations().get("call-1").hint()) + .isEqualTo("ok?"); + assertThat(reloaded.actions().endOfAgent()).isTrue(); + assertThat(reloaded.actions().agentState().get()).containsEntry("step", "2"); + } + @Test public void listSessions_empty() { assertThat(vertexAiSessionService.listSessions("789", "user1").blockingGet().sessions()) diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/JavaAdkToKt.kt b/tokt/src/main/kotlin/com/google/adk/tokt/JavaAdkToKt.kt index c1959bb86..6c0171f40 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/JavaAdkToKt.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/JavaAdkToKt.kt @@ -141,7 +141,7 @@ object JavaAdkToKt { * Java view of a Kotlin service is unwrapped to that Kotlin service, with no `dispatcher` hop. * Resumption works across the adapter (`EventActions.agentState` crosses); * `rewindBeforeInvocationId` crosses too, but a rewind drops the rewound turns only if the Java - * service stores it, which ADK Java's `VertexAiSessionService` does not. + * service stores that field. */ @JvmStatic @JvmOverloads