From f28e6ed911f0f92194fdcd11fb6281b1497de4b1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mi=C5=82osz=20Sobczyk?= Date: Tue, 14 Jul 2026 06:24:37 -0700 Subject: [PATCH] feat: add a Vertex AI-backed `SessionService` Implement `SessionService` against the managed Vertex AI Session Service so agent sessions and events persist in Vertex AI. This ports the Java ADK `VertexAiSessionService` and preserves its observable behavior, so callers can switch session backends without code changes. - Transport: use a small custom OkHttp-based client instead of a generated SDK client. The service needs only a few REST calls (create/get/list/delete session, list/append events) plus long-running-operation polling on create, so a dedicated client keeps the dependency surface small and the request/response mapping explicit. - Wire format: map payloads with `@Serializable` DTOs and mapper extensions rather than serializing the domain `Event`/`Content` directly, keeping the domain model decoupled from the wire. The mappers flatten streaming/turn signaling onto `Event` and encode the `State.REMOVED` sentinel as JSON null, which is the state-delta representation the service expects. - Encode byte fields (`Part.thoughtSignature`, `Blob.data`) as base64 strings to match the proto3-JSON representation of `bytes`. The kotlinx default serializes a `ByteArray` as an array of ints, which the service rejects ("Proto field is not repeating, cannot start list."); that silently dropped model events carrying a thought signature and corrupted session history. - Follow the Java ADK's semantics: update the in-memory session before the remote append, persist every event (including partial ones), and return list results in server order. Append failures are logged and skipped so one bad event does not abort a run, while read errors (5xx) are surfaced instead of masked as empty results. - Add unit tests for the service (mocked client) and the mappers (including a base64 wire-format regression test), plus transport tests driving a real client against a `MockWebServer`. PiperOrigin-RevId: 947644188 --- core/build.gradle.kts | 1 + .../Base64ByteArraySerializer.kt | 50 +++ .../kotlin/com/google/adk/kt/types/Blob.kt | 3 +- .../kotlin/com/google/adk/kt/types/Part.kt | 3 +- .../google/adk/kt/sessions/HttpApiClient.kt | 94 ++++++ .../google/adk/kt/sessions/VertexAiClient.kt | 247 +++++++++++++++ .../adk/kt/sessions/VertexAiSessionService.kt | 145 +++++++++ .../adk/kt/sessions/dto/EventActionsDto.kt | 39 +++ .../adk/kt/sessions/dto/EventMetadataDto.kt | 40 +++ .../adk/kt/sessions/dto/OperationDto.kt | 26 ++ .../adk/kt/sessions/dto/ResponsesDto.kt | 41 +++ .../google/adk/kt/sessions/dto/SessionDto.kt | 34 ++ .../adk/kt/sessions/dto/SessionEventDto.kt | 40 +++ .../adk/kt/sessions/dto/SessionMappers.kt | 178 +++++++++++ .../adk/kt/sessions/dto/TimestampDto.kt | 86 +++++ .../adk/kt/sessions/VertexAiClientTest.kt | 195 ++++++++++++ .../kt/sessions/VertexAiSessionServiceTest.kt | 293 ++++++++++++++++++ .../adk/kt/sessions/dto/SessionMappersTest.kt | 208 +++++++++++++ gradle/libs.versions.toml | 1 + 19 files changed, 1722 insertions(+), 2 deletions(-) create mode 100644 core/src/commonMain/kotlin/com/google/adk/kt/serialization/Base64ByteArraySerializer.kt create mode 100644 core/src/jvmMain/kotlin/com/google/adk/kt/sessions/HttpApiClient.kt create mode 100644 core/src/jvmMain/kotlin/com/google/adk/kt/sessions/VertexAiClient.kt create mode 100644 core/src/jvmMain/kotlin/com/google/adk/kt/sessions/VertexAiSessionService.kt create mode 100644 core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/EventActionsDto.kt create mode 100644 core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/EventMetadataDto.kt create mode 100644 core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/OperationDto.kt create mode 100644 core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/ResponsesDto.kt create mode 100644 core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/SessionDto.kt create mode 100644 core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/SessionEventDto.kt create mode 100644 core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/SessionMappers.kt create mode 100644 core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/TimestampDto.kt create mode 100644 core/src/jvmTest/kotlin/com/google/adk/kt/sessions/VertexAiClientTest.kt create mode 100644 core/src/jvmTest/kotlin/com/google/adk/kt/sessions/VertexAiSessionServiceTest.kt create mode 100644 core/src/jvmTest/kotlin/com/google/adk/kt/sessions/dto/SessionMappersTest.kt diff --git a/core/build.gradle.kts b/core/build.gradle.kts index a5676585..7e27a8d3 100644 --- a/core/build.gradle.kts +++ b/core/build.gradle.kts @@ -78,6 +78,7 @@ kotlin { implementation(libs.google.gson) implementation(libs.google.cloud.storage) implementation(libs.mcp) + implementation(libs.okhttp) implementation(libs.kotlinx.coroutines.reactor) implementation(libs.slf4j.api) implementation(libs.google.flogger.extensions) diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/serialization/Base64ByteArraySerializer.kt b/core/src/commonMain/kotlin/com/google/adk/kt/serialization/Base64ByteArraySerializer.kt new file mode 100644 index 00000000..fa144e27 --- /dev/null +++ b/core/src/commonMain/kotlin/com/google/adk/kt/serialization/Base64ByteArraySerializer.kt @@ -0,0 +1,50 @@ +/* + * 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 + * + * https://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.kt.serialization + +import kotlin.io.encoding.Base64 +import kotlin.io.encoding.ExperimentalEncodingApi +import kotlinx.serialization.KSerializer +import kotlinx.serialization.descriptors.PrimitiveKind +import kotlinx.serialization.descriptors.PrimitiveSerialDescriptor +import kotlinx.serialization.descriptors.SerialDescriptor +import kotlinx.serialization.encoding.Decoder +import kotlinx.serialization.encoding.Encoder + +/** + * Serializes a [ByteArray] as a base64 string. + * + * This matches the proto3-JSON representation of `bytes` fields, which is what the Vertex AI APIs + * (including the managed Session Service) expect. The kotlinx default encodes a `ByteArray` as a + * JSON array of numbers, which those APIs reject (`"Proto field is not repeating, cannot start + * list."`). Apply it with `@Serializable(with = Base64ByteArraySerializer::class)` on `ByteArray` + * fields that are persisted or sent over such wires (e.g. `Part.thoughtSignature`, `Blob.data`). + */ +@OptIn(ExperimentalEncodingApi::class) +object Base64ByteArraySerializer : KSerializer { + override val descriptor: SerialDescriptor = + PrimitiveSerialDescriptor( + "com.google.adk.kt.serialization.Base64ByteArray", + PrimitiveKind.STRING, + ) + + override fun serialize(encoder: Encoder, value: ByteArray) { + encoder.encodeString(Base64.encode(value)) + } + + override fun deserialize(decoder: Decoder): ByteArray = Base64.decode(decoder.decodeString()) +} diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/types/Blob.kt b/core/src/commonMain/kotlin/com/google/adk/kt/types/Blob.kt index e94c2146..11b7c0d7 100644 --- a/core/src/commonMain/kotlin/com/google/adk/kt/types/Blob.kt +++ b/core/src/commonMain/kotlin/com/google/adk/kt/types/Blob.kt @@ -16,6 +16,7 @@ package com.google.adk.kt.types +import com.google.adk.kt.serialization.Base64ByteArraySerializer import kotlinx.serialization.Serializable /** Represents binary data. */ @@ -23,7 +24,7 @@ import kotlinx.serialization.Serializable data class Blob( val mimeType: String? = null, val displayName: String? = null, - val data: ByteArray? = null, + @Serializable(with = Base64ByteArraySerializer::class) val data: ByteArray? = null, ) { override fun equals(other: Any?): Boolean { if (this === other) return true diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/types/Part.kt b/core/src/commonMain/kotlin/com/google/adk/kt/types/Part.kt index dac6442a..8acb5aef 100644 --- a/core/src/commonMain/kotlin/com/google/adk/kt/types/Part.kt +++ b/core/src/commonMain/kotlin/com/google/adk/kt/types/Part.kt @@ -16,6 +16,7 @@ package com.google.adk.kt.types import com.google.adk.kt.annotations.FrameworkInternalApi +import com.google.adk.kt.serialization.Base64ByteArraySerializer import kotlinx.serialization.Contextual import kotlinx.serialization.Serializable import kotlinx.serialization.Transient @@ -49,7 +50,7 @@ constructor( /** Indicates whether the part represents the model's thought process. */ val thought: Boolean? = null, /** An opaque signature for the thought. */ - val thoughtSignature: ByteArray? = null, + @Serializable(with = Base64ByteArraySerializer::class) val thoughtSignature: ByteArray? = null, /** Metadata for a video part (segment offsets and frame rate). */ val videoMetadata: VideoMetadata? = null, /** Arbitrary key-value metadata associated with this part. The map must be JSON serializable. */ diff --git a/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/HttpApiClient.kt b/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/HttpApiClient.kt new file mode 100644 index 00000000..ad6f053e --- /dev/null +++ b/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/HttpApiClient.kt @@ -0,0 +1,94 @@ +/* + * 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.kt.sessions + +import com.google.auth.oauth2.GoogleCredentials +import okhttp3.MediaType.Companion.toMediaType +import okhttp3.OkHttpClient +import okhttp3.Request +import okhttp3.RequestBody.Companion.toRequestBody +import okhttp3.Response + +/** + * A simplified HTTP client for interacting with the Vertex AI Session Service. + * + * @property project The Google Cloud Project ID. + * @property location The Google Cloud Location (region) where the service resides. The special + * value `"global"` selects the global, non-region-prefixed endpoint. + * @property httpClient The underlying [OkHttpClient] instance to use. + * @property credentials The Google credentials used for authentication. + * @property baseUrlOverride Optional override for the API base URL (useful for testing). When set, + * it must already include the API version segment (e.g. `http://localhost:8080/v1beta1`). + */ +class HttpApiClient( + private val project: String, + private val location: String, + private val httpClient: OkHttpClient = OkHttpClient(), + private val credentials: GoogleCredentials = defaultCredentials(), + private val baseUrlOverride: String? = null, +) { + private val baseUrl: String = + baseUrlOverride + ?: if (location == "global") { + "https://aiplatform.googleapis.com/$API_VERSION" + } else { + "https://$location-aiplatform.googleapis.com/$API_VERSION" + } + + /** + * Sends an authenticated HTTP request to the Vertex AI Session Service. + * + * This method refreshes the credentials if they are expired, injects the necessary + * `Authorization` and `Content-Type` headers (plus `x-goog-user-project` when the credentials + * declare a quota project), and executes the request synchronously. + * + * @param httpMethod The HTTP method to use (`GET`, `POST`, or `DELETE`). + * @param path The relative API path (e.g. `reasoningEngines/engine-id/sessions`). + * @param requestJson The JSON payload for `POST` requests. Ignored for `GET` and `DELETE`. + * @return The [Response] object from OkHttp. The caller is responsible for closing it. + * @throws IllegalArgumentException if the provided [httpMethod] is not supported. + */ + fun request(httpMethod: String, path: String, requestJson: String): Response { + val url = "$baseUrl/projects/$project/locations/$location/$path" + val requestBuilder = Request.Builder().url(url) + + credentials.refreshIfExpired() + requestBuilder.header("Authorization", "Bearer ${credentials.accessToken.tokenValue}") + requestBuilder.header("Content-Type", "application/json") + credentials.quotaProjectId?.let { requestBuilder.header("x-goog-user-project", it) } + + when (httpMethod.uppercase()) { + "POST" -> + requestBuilder.post( + requestJson.toRequestBody("application/json; charset=utf-8".toMediaType()) + ) + "GET" -> requestBuilder.get() + "DELETE" -> requestBuilder.delete() + else -> throw IllegalArgumentException("Unsupported HTTP method: $httpMethod") + } + + return httpClient.newCall(requestBuilder.build()).execute() + } + + companion object { + private const val API_VERSION = "v1beta1" + + private fun defaultCredentials(): GoogleCredentials = + GoogleCredentials.getApplicationDefault() + .createScoped("https://www.googleapis.com/auth/cloud-platform") + } +} diff --git a/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/VertexAiClient.kt b/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/VertexAiClient.kt new file mode 100644 index 00000000..4b5d579e --- /dev/null +++ b/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/VertexAiClient.kt @@ -0,0 +1,247 @@ +/* + * 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. + */ + +@file:OptIn(FrameworkInternalApi::class) + +package com.google.adk.kt.sessions + +import com.google.adk.kt.annotations.FrameworkInternalApi +import com.google.adk.kt.logging.LoggerFactory +import com.google.adk.kt.serialization.adkJson +import com.google.adk.kt.sessions.dto.CreateSessionRequestDto +import com.google.adk.kt.sessions.dto.ListEventsResponseDto +import com.google.adk.kt.sessions.dto.ListSessionsResponseDto +import com.google.adk.kt.sessions.dto.OperationDto +import com.google.adk.kt.sessions.dto.SessionDto +import com.google.adk.kt.sessions.dto.SessionEventDto +import java.io.IOException +import java.net.URLEncoder +import java.nio.charset.StandardCharsets +import java.util.concurrent.TimeoutException +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.delay +import kotlinx.coroutines.withContext +import kotlinx.serialization.json.JsonElement +import kotlinx.serialization.json.JsonObject +import kotlinx.serialization.json.JsonPrimitive +import okhttp3.Response + +/** + * Client for interacting with the Vertex AI Session Service API. + * + * Session creation is a long-running operation (LRO): [createSession] polls the returned operation + * until it is `done` and then fetches the materialized session. + * + * @property apiClient The authenticated HTTP transport used for all requests. + */ +internal open class VertexAiClient(private val apiClient: HttpApiClient) { + + private val logger = LoggerFactory.getLogger(VertexAiClient::class) + + /** + * Creates a new session for the given reasoning engine. + * + * This is a long-running operation: the method polls the operation returned by the create call + * until completion and then fetches the materialized session. + * + * @param reasoningEngineId The ID of the reasoning engine. + * @param userId The ID of the user owning the session. + * @param state Optional initial state map to inject into the session. + * @return The created [SessionDto], or `null` if creation failed. + */ + open suspend fun createSession( + reasoningEngineId: String, + userId: String, + state: Map?, + ): SessionDto? { + val requestBody = + adkJson.encodeToString( + CreateSessionRequestDto.serializer(), + CreateSessionRequestDto( + userId = userId, + sessionState = state?.let { stateMapToJsonElement(it) }, + ), + ) + val path = "reasoningEngines/$reasoningEngineId/sessions" + val response = performApiRequest("POST", path, requestBody) + val createResponse = + decodeResponse(response, "createSession", path) ?: return null + + // The create call returns an LRO whose `name` encodes both the session id and the operation id, + // e.g. ".../sessions/{sessionId}/operations/{operationId}". + val operationName = createResponse.name ?: return null + val parts = operationName.split("/") + if (parts.size < 3) { + return null + } + val sessionId = parts[parts.size - 3] + val operationId = parts.last() + + pollOperation(operationId) + return getSession(reasoningEngineId, sessionId) + } + + /** Retrieves an existing session. */ + open suspend fun getSession(reasoningEngineId: String, sessionId: String): SessionDto? { + val path = "reasoningEngines/$reasoningEngineId/sessions/$sessionId" + val response = performApiRequest("GET", path, "") + return decodeResponse(response, "getSession", path) + } + + /** Lists sessions associated with a reasoning engine, filtered by user ID. */ + open suspend fun listSessions( + reasoningEngineId: String, + userId: String, + ): ListSessionsResponseDto? { + val path = "reasoningEngines/$reasoningEngineId/sessions?filter=user_id=$userId" + val response = performApiRequest("GET", path, "") + return decodeResponse(response, "listSessions", path) + } + + /** + * Lists the events associated with a session. + * + * @param filter Optional server-side filter expression (e.g. `timestamp>="..."`). It is URL form + * escaped before being appended to the request. + */ + open suspend fun listEvents( + reasoningEngineId: String, + sessionId: String, + filter: String? = null, + ): ListEventsResponseDto? { + var path = "reasoningEngines/$reasoningEngineId/sessions/$sessionId/events" + if (filter != null) { + path += "?filter=" + URLEncoder.encode(filter, StandardCharsets.UTF_8.name()) + } + val response = performApiRequest("GET", path, "") + return decodeResponse(response, "listEvents", path) + } + + /** Deletes an existing session. */ + open suspend fun deleteSession(reasoningEngineId: String, sessionId: String) { + performApiRequest("DELETE", "reasoningEngines/$reasoningEngineId/sessions/$sessionId", "") + .close() + } + + /** Appends a new event (e.g. user message or model response) to a session. */ + open suspend fun appendEvent( + reasoningEngineId: String, + sessionId: String, + event: SessionEventDto, + ) { + val body = adkJson.encodeToString(SessionEventDto.serializer(), event) + val path = "reasoningEngines/$reasoningEngineId/sessions/$sessionId:appendEvent" + val response = performApiRequest("POST", path, body) + response.use { resp -> + val responseBody = resp.body?.string() + if (!resp.isSuccessful) { + // Match Java ADK: log and continue. The Runner keeps going with a partial history; the log + // trail names the offending field or condition so the user can diagnose. + logger.warn { + "appendEvent failed: HTTP ${resp.code} ${resp.message} path=$path\n" + + " request=${truncate(body)}\n" + + " response=${truncate(responseBody)}" + } + } + } + } + + private suspend fun performApiRequest(method: String, path: String, body: String): Response = + withContext(Dispatchers.IO) { apiClient.request(method, path, body) } + + /** + * Reads and decodes an API response into [T]. + * + * - `2xx` with a non-empty body → decode. + * - `2xx` with empty body → `null`. + * - `404` → `null` (legitimate "not found"). + * - Other `4xx` → `null` with a warning; caller sees "not found" but the log records why. + * - `5xx` → throws [IOException]. + */ + private inline fun decodeResponse( + response: Response, + opName: String, + path: String, + ): T? { + response.use { resp -> + val bodyString = resp.body?.string() + if (!resp.isSuccessful) { + when { + resp.code == 404 -> return null + resp.code in 500..599 -> + throw IOException( + "$opName failed: HTTP ${resp.code} ${resp.message} path=$path" + + " response=${truncate(bodyString)}" + ) + else -> { + logger.warn { + "$opName failed: HTTP ${resp.code} ${resp.message} path=$path" + + " response=${truncate(bodyString)}" + } + return null + } + } + } + if (bodyString.isNullOrEmpty()) return null + return adkJson.decodeFromString(bodyString) + } + } + + private fun truncate(s: String?, max: Int = 2000): String { + if (s == null) return "" + return if (s.length <= max) s else s.take(max) + "…(${s.length - max} more chars truncated)" + } + + /** + * Polls a long-running operation until it reports `done` or the retry budget is exhausted. + * + * @throws TimeoutException if the operation does not complete within [MAX_RETRY_ATTEMPTS]. + */ + private suspend fun pollOperation(operationId: String) { + val path = "operations/$operationId" + for (attempt in 0 until MAX_RETRY_ATTEMPTS) { + val op = + decodeResponse(performApiRequest("GET", path, ""), "pollOperation", path) + if (op?.done == true) { + return + } + delay(POLL_INTERVAL_MILLIS) + } + throw TimeoutException("Operation $operationId did not complete in time.") + } + + /** + * Converts a caller-supplied state map of arbitrary values into a [JsonObject] suitable for + * embedding in a [CreateSessionRequestDto]. Values that aren't JSON-native fall back to their + * `toString()` representation. + */ + private fun stateMapToJsonElement(state: Map): JsonElement { + val entries: Map = state.mapValues { (_, value) -> + when (value) { + is JsonElement -> value + is Boolean -> JsonPrimitive(value) + is Number -> JsonPrimitive(value) + else -> JsonPrimitive(value.toString()) + } + } + return JsonObject(entries) + } + + private companion object { + private const val MAX_RETRY_ATTEMPTS = 5 + private const val POLL_INTERVAL_MILLIS = 1000L + } +} diff --git a/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/VertexAiSessionService.kt b/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/VertexAiSessionService.kt new file mode 100644 index 00000000..8949601c --- /dev/null +++ b/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/VertexAiSessionService.kt @@ -0,0 +1,145 @@ +/* + * 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.kt.sessions + +import com.google.adk.kt.events.Event +import com.google.adk.kt.sessions.dto.toAdk +import com.google.adk.kt.sessions.dto.toDto +import java.io.IOException + +/** + * A [SessionService] backed by the managed Vertex AI Session Service. + * + * This is a Kotlin port of the Java ADK `com.google.adk.sessions.VertexAiSessionService`. It uses a + * custom [VertexAiClient] for transport. The wire format (defined by `session.proto`) is modeled + * with a small set of `@Serializable` DTOs in the `dto` sub-package; this class only translates + * between those DTOs and the ADK domain types via the mapper extensions in + * `com.google.adk.kt.sessions.dto.SessionMappers`. + * + * The application name ([SessionKey.appName]) must be either a bare numeric reasoning engine id or + * a full `projects/{project}/locations/{location}/reasoningEngines/{id}` resource name. + * + * @property client The [VertexAiClient] used to talk to the Vertex AI Session API. + */ +class VertexAiSessionService internal constructor(private val client: VertexAiClient) : + SessionService { + + constructor(apiClient: HttpApiClient) : this(VertexAiClient(apiClient)) + + override suspend fun createSession(key: SessionKey, state: Map?): Session { + val reasoningEngineId = parseReasoningEngineId(key.appName) + val sessionDto = + client.createSession(reasoningEngineId, key.userId, state) + ?: throw IOException("Failed to create session") + return sessionDto.toAdk(key.appName, key.userId, key.id) + } + + override suspend fun getSession(key: SessionKey, config: GetSessionConfig?): Session? { + val sessionId = requireNotNull(key.id) { "SessionKey.id is required for getSession." } + val reasoningEngineId = parseReasoningEngineId(key.appName) + val sessionDto = client.getSession(reasoningEngineId, sessionId) ?: return null + val session = sessionDto.toAdk(key.appName, key.userId, sessionId) + + val eventsResponse = + client.listEvents(reasoningEngineId, sessionId, afterTimestampFilter(config)) + val events = eventsResponse?.sessionEvents?.map { it.toAdk() } ?: emptyList() + session.events.addAll(filterEvents(events, config)) + return session + } + + override suspend fun listSessions(appName: String, userId: String): ListSessionsResponse { + val reasoningEngineId = parseReasoningEngineId(appName) + val response = client.listSessions(reasoningEngineId, userId) ?: return ListSessionsResponse() + val sessions = + response.sessions?.map { it.toAdk(appName, userId, fallbackId = null) } ?: emptyList() + return ListSessionsResponse(sessions) + } + + override suspend fun deleteSession(key: SessionKey) { + val sessionId = requireNotNull(key.id) { "SessionKey.id is required for deleteSession." } + val reasoningEngineId = parseReasoningEngineId(key.appName) + client.deleteSession(reasoningEngineId, sessionId) + } + + override suspend fun listEvents(key: SessionKey): ListEventsResponse { + val sessionId = requireNotNull(key.id) { "SessionKey.id is required for listEvents." } + val reasoningEngineId = parseReasoningEngineId(key.appName) + val response = client.listEvents(reasoningEngineId, sessionId) ?: return ListEventsResponse() + val events = response.sessionEvents?.map { it.toAdk() } ?: emptyList() + return ListEventsResponse(events) + } + + override suspend fun appendEvent(session: Session, event: Event): Event { + val reasoningEngineId = parseReasoningEngineId(session.key.appName) + val sessionId = requireNotNull(session.key.id) { "Session.key.id is required for appendEvent." } + val appended = super.appendEvent(session, event) + client.appendEvent(reasoningEngineId, sessionId, appended.toDto()) + return appended + } + + /** + * Trims the events to [GetSessionConfig.numRecentEvents] after sorting by timestamp. The + * [GetSessionConfig.afterTimestamp] filter is applied server-side (see [afterTimestampFilter]) so + * it is not re-applied here, matching the Java implementation. + */ + private fun filterEvents(events: List, config: GetSessionConfig?): List { + val sorted = events.sortedBy { it.timestamp } + val numRecentEvents = config?.numRecentEvents ?: return sorted + return if (sorted.size > numRecentEvents) { + sorted.subList(sorted.size - numRecentEvents, sorted.size) + } else { + sorted + } + } + + companion object { + private val APP_NAME_PATTERN = + Regex("^projects/([a-zA-Z0-9-_]+)/locations/([a-zA-Z0-9-_]+)/reasoningEngines/(\\d+)$") + + /** + * Builds the inclusive server-side `timestamp>=` filter for [GetSessionConfig.afterTimestamp]. + * The filter is only applied when [GetSessionConfig.numRecentEvents] is not set, matching the + * precedence in [filterEvents]. + */ + private fun afterTimestampFilter(config: GetSessionConfig?): String? { + if (config?.numRecentEvents == null && config?.afterTimestamp != null) { + return "timestamp>=\"${config.afterTimestamp}\"" + } + return null + } + + /** + * Extracts the reasoning engine id from a bare numeric id or a full ReasoningEngine resource + * name. + * + * @throws IllegalArgumentException if [appName] is neither a numeric id nor a valid resource + * name. + */ + internal fun parseReasoningEngineId(appName: String): String { + if (appName.matches(Regex("\\d+"))) { + return appName + } + val match = + APP_NAME_PATTERN.matchEntire(appName) + ?: throw IllegalArgumentException( + "App name $appName is not valid. It should either be the full ReasoningEngine resource" + + " name, or the reasoning engine id." + ) + return match.groupValues[3] + } + } +} diff --git a/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/EventActionsDto.kt b/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/EventActionsDto.kt new file mode 100644 index 00000000..7bf380ae --- /dev/null +++ b/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/EventActionsDto.kt @@ -0,0 +1,39 @@ +/* + * 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.kt.sessions.dto + +import kotlinx.serialization.Serializable +import kotlinx.serialization.json.JsonElement + +/** + * Wire-level representation of `EventActions` from `session.proto`. + * + * [stateDelta] is a raw [JsonElement] so the mapper can control the encoding of the + * [com.google.adk.kt.sessions.State.REMOVED] sentinel (emitted as JSON `null`, decoded back). Both + * [transferAgent] (canonical wire name) and [transferToAgent] (accepted for compatibility) are + * modeled: the mapper reads either and always writes [transferAgent]. + */ +@Serializable +internal data class EventActionsDto( + val skipSummarization: Boolean? = null, + val stateDelta: JsonElement? = null, + val artifactDelta: Map? = null, + val transferAgent: String? = null, + val transferToAgent: String? = null, + val escalate: Boolean? = null, + val endOfAgent: Boolean? = null, +) diff --git a/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/EventMetadataDto.kt b/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/EventMetadataDto.kt new file mode 100644 index 00000000..56890021 --- /dev/null +++ b/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/EventMetadataDto.kt @@ -0,0 +1,40 @@ +/* + * 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.kt.sessions.dto + +import kotlinx.serialization.Serializable +import kotlinx.serialization.json.JsonElement + +/** + * Wire-level representation of `EventMetadata` from `session.proto`. + * + * The streaming/turn signaling fields ([partial], [turnComplete], [interrupted], [branch], + * [longRunningToolIds]) live under this nested object on the wire even though the ADK domain + * [com.google.adk.kt.events.Event] carries them flat at the top level. Structured sub-objects + * ([groundingMetadata], [usageMetadata]) are kept as [JsonElement] to decouple the DTO from the ADK + * model shape; the mapper decodes them into the corresponding ADK types. + */ +@Serializable +internal data class EventMetadataDto( + val partial: Boolean? = null, + val turnComplete: Boolean? = null, + val interrupted: Boolean? = null, + val branch: String? = null, + val longRunningToolIds: List? = null, + val groundingMetadata: JsonElement? = null, + val usageMetadata: JsonElement? = null, +) diff --git a/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/OperationDto.kt b/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/OperationDto.kt new file mode 100644 index 00000000..f4c1a72b --- /dev/null +++ b/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/OperationDto.kt @@ -0,0 +1,26 @@ +/* + * 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.kt.sessions.dto + +import kotlinx.serialization.Serializable + +/** + * Minimal wire representation of a `google.longrunning.Operation`. Only the fields inspected by the + * LRO polling loop in [com.google.adk.kt.sessions.VertexAiClient] are modeled; unknown fields are + * silently ignored by the shared JSON reader. + */ +@Serializable internal data class OperationDto(val name: String? = null, val done: Boolean? = null) diff --git a/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/ResponsesDto.kt b/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/ResponsesDto.kt new file mode 100644 index 00000000..55f721ac --- /dev/null +++ b/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/ResponsesDto.kt @@ -0,0 +1,41 @@ +/* + * 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.kt.sessions.dto + +import kotlinx.serialization.Serializable +import kotlinx.serialization.json.JsonElement + +/** Wire response for `reasoningEngines/{id}/sessions` (list). */ +@Serializable +internal data class ListSessionsResponseDto( + val sessions: List? = null, + val nextPageToken: String? = null, +) + +/** Wire response for `reasoningEngines/{id}/sessions/{sid}/events` (list). */ +@Serializable +internal data class ListEventsResponseDto( + val sessionEvents: List? = null, + val nextPageToken: String? = null, +) + +/** Wire request body for `POST reasoningEngines/{id}/sessions` (create). */ +@Serializable +internal data class CreateSessionRequestDto( + val userId: String, + val sessionState: JsonElement? = null, +) diff --git a/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/SessionDto.kt b/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/SessionDto.kt new file mode 100644 index 00000000..1e44e196 --- /dev/null +++ b/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/SessionDto.kt @@ -0,0 +1,34 @@ +/* + * 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.kt.sessions.dto + +import kotlinx.serialization.Serializable +import kotlinx.serialization.json.JsonElement + +/** + * Wire-level representation of the `Session` resource from `session.proto`. + * + * Only the fields consumed by the Kotlin ADK are modeled; other fields returned by the API + * (`labels`, `createTime`, `expireTime`/`ttl`, `agents`, ...) are silently ignored by the shared + * JSON reader. + */ +@Serializable +internal data class SessionDto( + val name: String? = null, + val updateTime: String? = null, + val sessionState: JsonElement? = null, +) diff --git a/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/SessionEventDto.kt b/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/SessionEventDto.kt new file mode 100644 index 00000000..b83bded2 --- /dev/null +++ b/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/SessionEventDto.kt @@ -0,0 +1,40 @@ +/* + * 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.kt.sessions.dto + +import kotlinx.serialization.Serializable +import kotlinx.serialization.json.JsonElement + +/** + * Wire-level representation of `SessionEvent` from `session.proto`. + * + * [content] is left as a [JsonElement] and decoded to [com.google.adk.kt.types.Content] by the + * mapper. [name] is populated only on read (the resource-name path whose last segment is the local + * event id); it is never emitted on write. + */ +@Serializable +internal data class SessionEventDto( + val name: String? = null, + val author: String? = null, + val invocationId: String? = null, + val timestamp: TimestampDto? = null, + val errorCode: String? = null, + val errorMessage: String? = null, + val content: JsonElement? = null, + val actions: EventActionsDto? = null, + val eventMetadata: EventMetadataDto? = null, +) diff --git a/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/SessionMappers.kt b/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/SessionMappers.kt new file mode 100644 index 00000000..1b182f52 --- /dev/null +++ b/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/SessionMappers.kt @@ -0,0 +1,178 @@ +/* + * 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. + */ + +@file:OptIn(FrameworkInternalApi::class) + +package com.google.adk.kt.sessions.dto + +import com.google.adk.kt.annotations.FrameworkInternalApi +import com.google.adk.kt.events.Event +import com.google.adk.kt.events.EventActions +import com.google.adk.kt.serialization.adkJson +import com.google.adk.kt.sessions.Session +import com.google.adk.kt.sessions.SessionKey +import com.google.adk.kt.sessions.State +import com.google.adk.kt.types.Content +import com.google.adk.kt.types.GroundingMetadata +import com.google.adk.kt.types.UsageMetadata +import kotlin.time.Instant +import kotlinx.serialization.json.JsonArray +import kotlinx.serialization.json.JsonElement +import kotlinx.serialization.json.JsonNull +import kotlinx.serialization.json.JsonObject +import kotlinx.serialization.json.JsonPrimitive +import kotlinx.serialization.json.boolean +import kotlinx.serialization.json.booleanOrNull +import kotlinx.serialization.json.decodeFromJsonElement +import kotlinx.serialization.json.encodeToJsonElement +import kotlinx.serialization.json.int +import kotlinx.serialization.json.intOrNull +import kotlinx.serialization.json.longOrNull + +/** Mappers between the wire-level DTOs in this package and the ADK domain types. */ +internal fun Event.toDto(): SessionEventDto { + val metadata = + EventMetadataDto( + partial = partial, + turnComplete = turnComplete, + interrupted = interrupted, + branch = branch, + longRunningToolIds = longRunningToolIds.takeIf { it.isNotEmpty() }?.toList(), + groundingMetadata = groundingMetadata?.let { adkJson.encodeToJsonElement(it) }, + usageMetadata = usageMetadata?.let { adkJson.encodeToJsonElement(it) }, + ) + val actionsDto = + EventActionsDto( + skipSummarization = actions.skipSummarization.takeIf { it }, + stateDelta = stateDeltaToDto(actions.stateDelta), + artifactDelta = actions.artifactDelta.takeIf { it.isNotEmpty() }?.toMap(), + transferAgent = actions.transferToAgent, + escalate = actions.escalate.takeIf { it }, + endOfAgent = actions.endOfAgent.takeIf { it }, + ) + return SessionEventDto( + author = author, + invocationId = invocationId, + timestamp = TimestampDto.fromEpochMillis(timestamp), + errorCode = errorCode, + errorMessage = errorMessage, + content = content?.let { adkJson.encodeToJsonElement(it) }, + actions = actionsDto, + eventMetadata = metadata, + ) +} + +internal fun SessionEventDto.toAdk(): Event { + val id = name?.substringAfterLast('/') ?: "" + val metadata = eventMetadata + return Event( + id = id, + invocationId = invocationId, + author = author ?: "", + content = content?.let { adkJson.decodeFromJsonElement(it) }, + actions = actions?.toAdk() ?: EventActions(), + longRunningToolIds = metadata?.longRunningToolIds?.toSet() ?: emptySet(), + partial = metadata?.partial ?: false, + turnComplete = metadata?.turnComplete ?: false, + errorCode = errorCode, + errorMessage = errorMessage, + interrupted = metadata?.interrupted ?: false, + branch = metadata?.branch, + groundingMetadata = + metadata?.groundingMetadata?.let { adkJson.decodeFromJsonElement(it) }, + usageMetadata = + metadata?.usageMetadata?.let { adkJson.decodeFromJsonElement(it) }, + timestamp = timestamp?.toEpochMillis() ?: 0L, + ) +} + +internal fun SessionDto.toAdk(appName: String, userId: String, fallbackId: String?): Session { + val sessionId = + name?.substringAfterLast('/') + ?: fallbackId + ?: error("Session response is missing a name and no fallback session id was provided.") + val lastUpdateTime = + updateTime?.let { Instant.fromEpochMilliseconds(java.time.Instant.parse(it).toEpochMilli()) } + ?: Instant.fromEpochMilliseconds(0) + val initialState: Map = + (sessionState as? JsonObject)?.let { jsonObj -> + jsonObj.mapValues { (_, v) -> jsonElementToAny(v) ?: State.REMOVED } + } ?: emptyMap() + return Session( + key = SessionKey(appName, userId, sessionId), + state = State(initialState), + lastUpdateTime = lastUpdateTime, + ) +} + +private fun EventActionsDto.toAdk(): EventActions { + val actions = EventActions() + actions.skipSummarization = skipSummarization ?: false + (stateDelta as? JsonObject)?.let { delta -> + for ((key, value) in delta) { + actions.stateDelta[key] = + if (value is JsonNull) State.REMOVED else (jsonElementToAny(value) ?: State.REMOVED) + } + } + artifactDelta?.let { actions.artifactDelta.putAll(it) } + actions.transferToAgent = transferAgent ?: transferToAgent + actions.escalate = escalate ?: false + actions.endOfAgent = endOfAgent ?: false + return actions +} + +/** + * Serializes the in-memory state delta into a [JsonObject] where [State.REMOVED] is emitted as JSON + * `null` (matches the Java ADK). + */ +private fun stateDeltaToDto(stateDelta: Map): JsonElement { + val entries = stateDelta.mapValues { (_, value) -> + if (value === State.REMOVED) JsonNull else anyToJsonElement(value) + } + return JsonObject(entries) +} + +private fun anyToJsonElement(value: Any): JsonElement = + when (value) { + is JsonElement -> value + is Boolean -> JsonPrimitive(value) + is Number -> JsonPrimitive(value) + is String -> JsonPrimitive(value) + is Map<*, *> -> + JsonObject( + value.entries.associate { (k, v) -> + k.toString() to (if (v == null) JsonNull else anyToJsonElement(v)) + } + ) + is List<*> -> JsonArray(value.map { if (it == null) JsonNull else anyToJsonElement(it) }) + else -> JsonPrimitive(value.toString()) + } + +private fun jsonElementToAny(element: JsonElement): Any? = + when (element) { + is JsonNull -> null + is JsonObject -> + element.mapValues { (_, v) -> if (v is JsonNull) State.REMOVED else jsonElementToAny(v) } + is JsonArray -> element.map { jsonElementToAny(it) } + is JsonPrimitive -> + when { + element.isString -> element.content + element.booleanOrNull != null -> element.boolean + element.longOrNull != null -> element.longOrNull!! + element.intOrNull != null -> element.int + else -> element.content + } + } diff --git a/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/TimestampDto.kt b/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/TimestampDto.kt new file mode 100644 index 00000000..8d5080b9 --- /dev/null +++ b/core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/TimestampDto.kt @@ -0,0 +1,86 @@ +/* + * 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.kt.sessions.dto + +import kotlinx.serialization.KSerializer +import kotlinx.serialization.Serializable +import kotlinx.serialization.descriptors.SerialDescriptor +import kotlinx.serialization.descriptors.buildClassSerialDescriptor +import kotlinx.serialization.encoding.Decoder +import kotlinx.serialization.encoding.Encoder +import kotlinx.serialization.json.JsonDecoder +import kotlinx.serialization.json.JsonEncoder +import kotlinx.serialization.json.JsonObject +import kotlinx.serialization.json.JsonPrimitive +import kotlinx.serialization.json.contentOrNull +import kotlinx.serialization.json.jsonPrimitive +import kotlinx.serialization.json.longOrNull + +/** + * Wire-level representation of `google.protobuf.Timestamp` used by the Vertex AI Session Service. + * + * The service accepts either an ISO-8601 string or a `{seconds, nanos}` object for event + * timestamps. This DTO always emits the object form (matching the Java ADK's writer) but its + * [TimestampDtoSerializer] transparently reads both shapes. + */ +@Serializable(with = TimestampDtoSerializer::class) +internal data class TimestampDto(val seconds: Long = 0, val nanos: Long = 0) { + fun toEpochMillis(): Long = seconds * 1000 + nanos / 1_000_000 + + companion object { + fun fromEpochMillis(epochMillis: Long): TimestampDto = + TimestampDto(seconds = epochMillis / 1000, nanos = (epochMillis % 1000) * 1_000_000) + } +} + +/** + * Serializer for [TimestampDto] that emits the `{seconds, nanos}` object form and accepts either + * that shape or an ISO-8601 string on read. + */ +internal object TimestampDtoSerializer : KSerializer { + override val descriptor: SerialDescriptor = + buildClassSerialDescriptor("com.google.adk.kt.sessions.dto.TimestampDto") + + override fun serialize(encoder: Encoder, value: TimestampDto) { + val jsonEncoder = + encoder as? JsonEncoder ?: error("TimestampDtoSerializer supports JSON formats only") + jsonEncoder.encodeJsonElement( + JsonObject( + mapOf("seconds" to JsonPrimitive(value.seconds), "nanos" to JsonPrimitive(value.nanos)) + ) + ) + } + + override fun deserialize(decoder: Decoder): TimestampDto { + val jsonDecoder = + decoder as? JsonDecoder ?: error("TimestampDtoSerializer supports JSON formats only") + val element = jsonDecoder.decodeJsonElement() + return when (element) { + is JsonObject -> { + val seconds = element["seconds"]?.jsonPrimitive?.longOrNull ?: 0L + val nanos = element["nanos"]?.jsonPrimitive?.longOrNull ?: 0L + TimestampDto(seconds, nanos) + } + is JsonPrimitive -> { + val text = element.contentOrNull ?: return TimestampDto() + val instant = java.time.Instant.parse(text) + TimestampDto(instant.epochSecond, instant.nano.toLong()) + } + else -> TimestampDto() + } + } +} diff --git a/core/src/jvmTest/kotlin/com/google/adk/kt/sessions/VertexAiClientTest.kt b/core/src/jvmTest/kotlin/com/google/adk/kt/sessions/VertexAiClientTest.kt new file mode 100644 index 00000000..67fd0761 --- /dev/null +++ b/core/src/jvmTest/kotlin/com/google/adk/kt/sessions/VertexAiClientTest.kt @@ -0,0 +1,195 @@ +/* + * 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.kt.sessions + +import com.google.adk.kt.sessions.dto.SessionEventDto +import com.google.adk.kt.sessions.dto.TimestampDto +import com.google.auth.oauth2.AccessToken +import com.google.auth.oauth2.GoogleCredentials +import com.google.common.truth.Truth.assertThat +import java.io.IOException +import java.time.Instant +import java.time.temporal.ChronoUnit +import java.util.Date +import kotlin.test.assertFailsWith +import kotlinx.coroutines.runBlocking +import okhttp3.mockwebserver.MockResponse +import okhttp3.mockwebserver.MockWebServer +import org.junit.After +import org.junit.Before +import org.junit.Test +import org.junit.runner.RunWith +import org.junit.runners.JUnit4 + +/** + * Transport-level tests for [VertexAiClient]. These drive a real [HttpApiClient] (with a fake, non + * expiring credential) against a [MockWebServer], so they exercise URL construction, header + * injection, JSON (de)serialization, the create-session long-running-operation poll, and the error + * handling contract end to end. + */ +@RunWith(JUnit4::class) +class VertexAiClientTest { + + private lateinit var server: MockWebServer + private lateinit var client: VertexAiClient + + @Before + fun setUp() { + server = MockWebServer() + server.start() + val apiClient = + HttpApiClient( + project = "test-project", + location = "test-location", + credentials = fakeCredentials(), + baseUrlOverride = server.url("/v1beta1").toString(), + ) + client = VertexAiClient(apiClient) + } + + @After + fun tearDown() { + server.shutdown() + } + + @Test + fun createSession_pollsOperationThenFetchesSession() { + // 1) create returns an LRO whose name encodes the session id and operation id. + server.enqueue( + jsonResponse("""{"name":"reasoningEngines/123/sessions/sess-1/operations/op-1"}""") + ) + // 2) the operation poll reports completion immediately. + server.enqueue(jsonResponse("""{"name":"operations/op-1","done":true}""")) + // 3) the materialized session is fetched. + server.enqueue( + jsonResponse( + """{"name":"reasoningEngines/123/sessions/sess-1","updateTime":"2024-12-12T12:12:12Z"}""" + ) + ) + + val session = runBlocking { client.createSession("123", "user", mapOf("k" to "v")) } + + assertThat(session).isNotNull() + assertThat(session!!.name).isEqualTo("reasoningEngines/123/sessions/sess-1") + + val createRequest = server.takeRequest() + assertThat(createRequest.method).isEqualTo("POST") + assertThat(createRequest.path) + .isEqualTo( + "/v1beta1/projects/test-project/locations/test-location/reasoningEngines/123/sessions" + ) + assertThat(createRequest.getHeader("Authorization")).isEqualTo("Bearer fake-token") + assertThat(createRequest.getHeader("Content-Type")).contains("application/json") + val body = createRequest.body.readUtf8() + assertThat(body).contains("\"userId\":\"user\"") + assertThat(body).contains("\"k\":\"v\"") + + assertThat(server.takeRequest().path).endsWith("/operations/op-1") + assertThat(server.takeRequest().path).endsWith("/reasoningEngines/123/sessions/sess-1") + } + + @Test + fun getSession_notFound_returnsNull() = runBlocking { + server.enqueue(MockResponse().setResponseCode(404)) + + assertThat(client.getSession("123", "missing")).isNull() + } + + @Test + fun getSession_serverError_throwsIoException() { + server.enqueue(MockResponse().setResponseCode(500).setBody("boom")) + + assertFailsWith { runBlocking { client.getSession("123", "s1") } } + } + + @Test + fun listSessions_parsesSessionsAndSendsUserFilter() { + server.enqueue( + jsonResponse( + """ + {"sessions":[ + {"name":"reasoningEngines/123/sessions/1"}, + {"name":"reasoningEngines/123/sessions/2"} + ]} + """ + .trimIndent() + ) + ) + + val response = runBlocking { client.listSessions("123", "user") } + + assertThat(response!!.sessions!!.map { it.name!!.substringAfterLast('/') }) + .containsExactly("1", "2") + .inOrder() + assertThat(server.takeRequest().path).contains("filter=user_id=user") + } + + @Test + fun listEvents_withFilter_urlEncodesFilter() { + server.enqueue(jsonResponse("""{"sessionEvents":[]}""")) + + val unused = runBlocking { + client.listEvents("123", "s1", "timestamp>=\"2024-12-12T12:00:10Z\"") + } + + val path = server.takeRequest().path!! + // The filter operator and quotes are URL-escaped (>= -> %3E%3D, " -> %22), not sent raw. + assertThat(path).contains("filter=timestamp%3E%3D%22") + assertThat(path).doesNotContain("timestamp>=") + } + + @Test + fun deleteSession_sendsDeleteToSessionPath() { + server.enqueue(MockResponse().setResponseCode(200)) + + runBlocking { client.deleteSession("123", "s1") } + + val request = server.takeRequest() + assertThat(request.method).isEqualTo("DELETE") + assertThat(request.path).endsWith("/reasoningEngines/123/sessions/s1") + } + + @Test + fun appendEvent_serverError_isSwallowed() { + // The Java ADK contract is to log and continue rather than fail the run. + server.enqueue(MockResponse().setResponseCode(400).setBody("bad request")) + + runBlocking { + client.appendEvent( + "123", + "s1", + SessionEventDto(author = "user", timestamp = TimestampDto.fromEpochMillis(1000)), + ) + } + + val request = server.takeRequest() + assertThat(request.method).isEqualTo("POST") + assertThat(request.path).endsWith("/reasoningEngines/123/sessions/s1:appendEvent") + } + + private companion object { + fun jsonResponse(body: String): MockResponse = + MockResponse().setHeader("Content-Type", "application/json").setBody(body) + + fun fakeCredentials(): GoogleCredentials = + GoogleCredentials.newBuilder() + .setAccessToken( + AccessToken("fake-token", Date(Instant.now().plus(1, ChronoUnit.DAYS).toEpochMilli())) + ) + .build() + } +} diff --git a/core/src/jvmTest/kotlin/com/google/adk/kt/sessions/VertexAiSessionServiceTest.kt b/core/src/jvmTest/kotlin/com/google/adk/kt/sessions/VertexAiSessionServiceTest.kt new file mode 100644 index 00000000..2b37046e --- /dev/null +++ b/core/src/jvmTest/kotlin/com/google/adk/kt/sessions/VertexAiSessionServiceTest.kt @@ -0,0 +1,293 @@ +/* + * 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.kt.sessions + +import com.google.adk.kt.events.Event +import com.google.adk.kt.events.EventActions +import com.google.adk.kt.sessions.dto.ListEventsResponseDto +import com.google.adk.kt.sessions.dto.ListSessionsResponseDto +import com.google.adk.kt.sessions.dto.SessionDto +import com.google.adk.kt.sessions.dto.SessionEventDto +import com.google.adk.kt.sessions.dto.TimestampDto +import com.google.adk.kt.types.Content +import com.google.adk.kt.types.Part +import com.google.common.truth.Truth.assertThat +import java.io.IOException +import kotlin.test.assertFailsWith +import kotlin.time.Instant +import kotlinx.coroutines.test.runTest +import kotlinx.serialization.json.JsonObject +import kotlinx.serialization.json.JsonPrimitive +import org.junit.Test +import org.junit.runner.RunWith +import org.junit.runners.JUnit4 +import org.mockito.kotlin.any +import org.mockito.kotlin.anyOrNull +import org.mockito.kotlin.doReturn +import org.mockito.kotlin.eq +import org.mockito.kotlin.isNull +import org.mockito.kotlin.mock +import org.mockito.kotlin.verifyBlocking + +/** + * Unit tests for [VertexAiSessionService]. + * + * These exercise the service-level logic (app-name parsing, DTO-to-domain mapping, event + * filtering/sorting and the append-event write-through) against a mocked [VertexAiClient], so no + * HTTP transport is involved. The transport itself is covered by `VertexAiClientTest`. + */ +@RunWith(JUnit4::class) +class VertexAiSessionServiceTest { + + @Test + fun parseReasoningEngineId_numericId_returnedVerbatim() { + assertThat(VertexAiSessionService.parseReasoningEngineId("123")).isEqualTo("123") + } + + @Test + fun parseReasoningEngineId_fullResourceName_extractsId() { + assertThat( + VertexAiSessionService.parseReasoningEngineId( + "projects/my-project/locations/us-central1/reasoningEngines/456" + ) + ) + .isEqualTo("456") + } + + @Test + fun parseReasoningEngineId_invalid_throws() { + assertFailsWith { + VertexAiSessionService.parseReasoningEngineId("not-a-valid-app-name") + } + } + + @Test + fun createSession_mapsClientResponse() = runTest { + val client = + mock { + onBlocking { createSession(eq("123"), eq("user"), anyOrNull()) } doReturn + SessionDto( + name = "reasoningEngines/123/sessions/session-1", + updateTime = "2024-12-12T12:12:12Z", + sessionState = JsonObject(mapOf("k" to JsonPrimitive("v"))), + ) + } + val service = VertexAiSessionService(client) + + val session = service.createSession(SessionKey("123", "user", id = null), mapOf("k" to "v")) + + assertThat(session.key.id).isEqualTo("session-1") + assertThat(session.key.appName).isEqualTo("123") + assertThat(session.key.userId).isEqualTo("user") + assertThat(session.state["k"]).isEqualTo("v") + } + + @Test + fun createSession_clientReturnsNull_throws() = runTest { + val service = VertexAiSessionService(mock()) + + assertFailsWith { + service.createSession(SessionKey("123", "user", id = null), state = null) + } + } + + @Test + fun getSession_notFound_returnsNull() = runTest { + // Unstubbed getSession returns null by default -> service short-circuits to null. + val service = VertexAiSessionService(mock()) + + assertThat(service.getSession(SessionKey("123", "user", "missing"))).isNull() + } + + @Test + fun getSession_numRecentEvents_returnsMostRecentSorted() = runTest { + val client = + mock { + onBlocking { getSession(eq("123"), eq("s1")) } doReturn + SessionDto(name = "reasoningEngines/123/sessions/s1", updateTime = "2024-12-12T12:00:30Z") + onBlocking { listEvents(eq("123"), eq("s1"), anyOrNull()) } doReturn + ListEventsResponseDto( + sessionEvents = listOf(eventDto("e3", 3000), eventDto("e1", 1000), eventDto("e2", 2000)) + ) + } + val service = VertexAiSessionService(client) + + val session = + service.getSession(SessionKey("123", "user", "s1"), GetSessionConfig(numRecentEvents = 2)) + + assertThat(session!!.events.map { it.id }).containsExactly("e2", "e3").inOrder() + } + + @Test + fun getSession_afterTimestamp_passesInclusiveServerFilter() = runTest { + val client = + mock { + onBlocking { getSession(eq("123"), eq("s1")) } doReturn + SessionDto(name = "reasoningEngines/123/sessions/s1") + } + val service = VertexAiSessionService(client) + val threshold = Instant.parse("2024-12-12T12:00:10Z") + + val unused = + service.getSession( + SessionKey("123", "user", "s1"), + GetSessionConfig(afterTimestamp = threshold), + ) + + verifyBlocking(client) { listEvents(eq("123"), eq("s1"), eq("timestamp>=\"$threshold\"")) } + } + + @Test + fun getSession_numRecentEventsTakesPrecedence_noServerFilter() = runTest { + val client = + mock { + onBlocking { getSession(eq("123"), eq("s1")) } doReturn + SessionDto(name = "reasoningEngines/123/sessions/s1") + } + val service = VertexAiSessionService(client) + + val unused = + service.getSession( + SessionKey("123", "user", "s1"), + GetSessionConfig( + numRecentEvents = 2, + afterTimestamp = Instant.parse("2024-12-12T12:00:10Z"), + ), + ) + + // When numRecentEvents is set, the afterTimestamp filter is not sent to the server. + verifyBlocking(client) { listEvents(eq("123"), eq("s1"), isNull()) } + } + + @Test + fun getSession_missingId_throws() = runTest { + val service = VertexAiSessionService(mock()) + + assertFailsWith { + service.getSession(SessionKey("123", "user", id = null)) + } + } + + @Test + fun listSessions_mapsClientResponse() = runTest { + val client = + mock { + onBlocking { listSessions(eq("123"), eq("user")) } doReturn + ListSessionsResponseDto( + sessions = + listOf( + SessionDto(name = "reasoningEngines/123/sessions/1"), + SessionDto(name = "reasoningEngines/123/sessions/2"), + ) + ) + } + val service = VertexAiSessionService(client) + + val response = service.listSessions("123", "user") + + assertThat(response.sessions.map { it.key.id }).containsExactly("1", "2").inOrder() + } + + @Test + fun listSessions_clientReturnsNull_returnsEmpty() = runTest { + val service = VertexAiSessionService(mock()) + + assertThat(service.listSessions("123", "user").sessions).isEmpty() + } + + @Test + fun listEvents_returnsEventsInServerOrder() = runTest { + val client = + mock { + onBlocking { listEvents(eq("123"), eq("s1"), anyOrNull()) } doReturn + ListEventsResponseDto( + sessionEvents = listOf(eventDto("e2", 2000), eventDto("e3", 3000), eventDto("e1", 1000)) + ) + } + val service = VertexAiSessionService(client) + + val response = service.listEvents(SessionKey("123", "user", "s1")) + + // listEvents does not re-sort; it returns events in the order the server sent them. + assertThat(response.events.map { it.id }).containsExactly("e2", "e3", "e1").inOrder() + } + + @Test + fun deleteSession_delegatesToClient() = runTest { + val client = mock() + val service = VertexAiSessionService(client) + + service.deleteSession(SessionKey("123", "user", "s1")) + + verifyBlocking(client) { deleteSession(eq("123"), eq("s1")) } + } + + @Test + fun deleteSession_missingId_throws() = runTest { + val service = VertexAiSessionService(mock()) + + assertFailsWith { + service.deleteSession(SessionKey("123", "user", id = null)) + } + } + + @Test + fun appendEvent_nonPartial_writesThroughAndUpdatesSession() = runTest { + val client = mock() + val service = VertexAiSessionService(client) + val session = Session(SessionKey("123", "user", "s1")) + val event = + Event( + author = "user", + timestamp = 1000L, + content = Content(role = "user", parts = listOf(Part(text = "hi"))), + actions = EventActions(stateDelta = mutableMapOf("k" to "v")), + ) + + val returned = service.appendEvent(session, event) + + assertThat(returned).isEqualTo(event) + verifyBlocking(client) { appendEvent(eq("123"), eq("s1"), any()) } + // super.appendEvent keeps the in-memory session in sync. + assertThat(session.events).contains(event) + assertThat(session.state["k"]).isEqualTo("v") + } + + @Test + fun appendEvent_partial_persistsRemotelyButNotInMemory() = runTest { + val client = mock() + val service = VertexAiSessionService(client) + val session = Session(SessionKey("123", "user", "s1")) + val event = Event(author = "user", partial = true, timestamp = 1000L) + + val unused = service.appendEvent(session, event) + + // Matching the Java/Python ADK: a partial event is still written to the remote service, but the + // base implementation does not add it to the in-memory session. + verifyBlocking(client) { appendEvent(eq("123"), eq("s1"), any()) } + assertThat(session.events).isEmpty() + } + + private companion object { + fun eventDto(id: String, epochMillis: Long): SessionEventDto = + SessionEventDto( + name = "reasoningEngines/123/sessions/s1/events/$id", + author = "agent", + timestamp = TimestampDto.fromEpochMillis(epochMillis), + ) + } +} diff --git a/core/src/jvmTest/kotlin/com/google/adk/kt/sessions/dto/SessionMappersTest.kt b/core/src/jvmTest/kotlin/com/google/adk/kt/sessions/dto/SessionMappersTest.kt new file mode 100644 index 00000000..07ebcfae --- /dev/null +++ b/core/src/jvmTest/kotlin/com/google/adk/kt/sessions/dto/SessionMappersTest.kt @@ -0,0 +1,208 @@ +/* + * 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.kt.sessions.dto + +import com.google.adk.kt.events.Event +import com.google.adk.kt.events.EventActions +import com.google.adk.kt.sessions.State +import com.google.adk.kt.types.Blob +import com.google.adk.kt.types.Content +import com.google.adk.kt.types.FunctionCall +import com.google.adk.kt.types.Part +import com.google.common.truth.Truth.assertThat +import kotlinx.serialization.json.JsonNull +import kotlinx.serialization.json.JsonObject +import kotlinx.serialization.json.JsonPrimitive +import kotlinx.serialization.json.jsonArray +import kotlinx.serialization.json.jsonObject +import kotlinx.serialization.json.jsonPrimitive +import org.junit.Test +import org.junit.runner.RunWith +import org.junit.runners.JUnit4 + +/** + * Unit tests for the wire-DTO <-> ADK-domain mappers in [SessionMappers]. + * + * These pin down the flattening of streaming/turn signaling onto [Event], the [State.REMOVED] + * sentinel encoding (JSON `null` on the wire), the `transferAgent`/`transferToAgent` compatibility + * read, and the session name/fallback-id and timestamp handling. + */ +@RunWith(JUnit4::class) +class SessionMappersTest { + + @Test + fun eventToDto_mapsCoreFieldsAndMetadata() { + val event = + Event( + id = "e1", + invocationId = "inv1", + author = "user", + timestamp = 1_734_005_532_123L, + turnComplete = true, + longRunningToolIds = setOf("tool1"), + actions = EventActions(transferToAgent = "agent").apply { stateDelta["k"] = "v" }, + ) + + val dto = event.toDto() + + assertThat(dto.author).isEqualTo("user") + assertThat(dto.invocationId).isEqualTo("inv1") + assertThat(dto.timestamp!!.toEpochMillis()).isEqualTo(1_734_005_532_123L) + assertThat(dto.eventMetadata!!.turnComplete).isTrue() + assertThat(dto.eventMetadata.longRunningToolIds).containsExactly("tool1") + assertThat(dto.actions!!.transferAgent).isEqualTo("agent") + assertThat((dto.actions.stateDelta as JsonObject)["k"]).isEqualTo(JsonPrimitive("v")) + } + + @Test + fun eventToDto_stateRemoved_isEncodedAsJsonNull() { + val event = + Event(author = "user", actions = EventActions().apply { stateDelta["gone"] = State.REMOVED }) + + val delta = event.toDto().actions!!.stateDelta as JsonObject + + assertThat(delta["gone"]).isEqualTo(JsonNull) + } + + @Test + fun sessionEventDtoToAdk_mapsFieldsAndCompatTransferAgent() { + val dto = + SessionEventDto( + name = "reasoningEngines/123/sessions/s1/events/e9", + author = "agent", + invocationId = "inv9", + timestamp = TimestampDto.fromEpochMillis(2000L), + // Only the compatibility field `transferToAgent` is set; the mapper must still read it. + actions = + EventActionsDto( + transferToAgent = "compatAgent", + stateDelta = JsonObject(mapOf("a" to JsonNull)), + ), + eventMetadata = EventMetadataDto(partial = true, branch = "b1"), + ) + + val event = dto.toAdk() + + assertThat(event.id).isEqualTo("e9") + assertThat(event.author).isEqualTo("agent") + assertThat(event.invocationId).isEqualTo("inv9") + assertThat(event.timestamp).isEqualTo(2000L) + assertThat(event.partial).isTrue() + assertThat(event.branch).isEqualTo("b1") + assertThat(event.actions.transferToAgent).isEqualTo("compatAgent") + // JSON null in the wire state delta decodes back to the REMOVED sentinel. + assertThat(event.actions.stateDelta["a"]).isEqualTo(State.REMOVED) + } + + @Test + fun event_roundTrip_preservesContentAndFields() { + val original = + Event( + author = "user", + invocationId = "inv-rt", + timestamp = 1000L, + content = Content(role = "user", parts = listOf(Part(text = "round-trip"))), + ) + + val restored = original.toDto().toAdk() + + assertThat(restored.author).isEqualTo("user") + assertThat(restored.invocationId).isEqualTo("inv-rt") + assertThat(restored.timestamp).isEqualTo(1000L) + assertThat(restored.content?.parts?.single()?.text).isEqualTo("round-trip") + } + + @Test + fun sessionDtoToAdk_mapsNameStateAndUpdateTime() { + val dto = + SessionDto( + name = "reasoningEngines/123/sessions/sX", + updateTime = "2024-12-12T12:12:12Z", + sessionState = JsonObject(mapOf("k" to JsonPrimitive("v"))), + ) + + val session = dto.toAdk(appName = "123", userId = "user", fallbackId = null) + + assertThat(session.key.id).isEqualTo("sX") + assertThat(session.key.appName).isEqualTo("123") + assertThat(session.key.userId).isEqualTo("user") + assertThat(session.state["k"]).isEqualTo("v") + assertThat(session.lastUpdateTime.toEpochMilliseconds()) + .isEqualTo(java.time.Instant.parse("2024-12-12T12:12:12Z").toEpochMilli()) + } + + @Test + fun sessionDtoToAdk_missingName_usesFallbackIdAndEpochTime() { + val dto = SessionDto(name = null, updateTime = null, sessionState = null) + + val session = dto.toAdk(appName = "123", userId = "user", fallbackId = "fallback-1") + + assertThat(session.key.id).isEqualTo("fallback-1") + assertThat(session.state).isEmpty() + assertThat(session.lastUpdateTime.toEpochMilliseconds()).isEqualTo(0L) + } + + @Test + fun eventToDto_richContent_serializesBytesAsBase64AndRoundTrips() { + // The event content on the wire is a google.genai Content. Parts must use the SDK's camelCase + // field names, and bytes fields (thoughtSignature, inlineData.data) must be base64 STRINGS, not + // JSON int arrays - Vertex rejects int arrays for `bytes` proto fields ("Proto field is not + // repeating, cannot start list."), which silently dropped every model (thinking) event. + val signature = byteArrayOf(1, -113, 61, 107) + val content = + Content( + role = "model", + parts = + listOf( + Part( + functionCall = FunctionCall(name = "get_weather", args = mapOf("city" to "SF")), + thoughtSignature = signature, + ), + Part(inlineData = Blob(mimeType = "image/png", data = byteArrayOf(1, 2, 3))), + ), + ) + val event = Event(author = "model", timestamp = 1000L, content = content) + + val contentJson = event.toDto().content!!.jsonObject + + assertThat(contentJson["role"]!!.jsonPrimitive.content).isEqualTo("model") + val parts = contentJson["parts"]!!.jsonArray + assertThat(parts).hasSize(2) + val functionCallPart = parts[0].jsonObject + assertThat(functionCallPart["functionCall"]!!.jsonObject["name"]!!.jsonPrimitive.content) + .isEqualTo("get_weather") + assertThat( + functionCallPart["functionCall"]!! + .jsonObject["args"]!! + .jsonObject["city"]!! + .jsonPrimitive + .content + ) + .isEqualTo("SF") + // Bytes must be base64 strings. If serialized as an int array these `.jsonPrimitive.isString` + // accesses fail, which is the exact defect that made Vertex 400 and drop model events. + assertThat(functionCallPart["thoughtSignature"]!!.jsonPrimitive.isString).isTrue() + val inlineData = parts[1].jsonObject["inlineData"]!!.jsonObject + assertThat(inlineData["mimeType"]!!.jsonPrimitive.content).isEqualTo("image/png") + assertThat(inlineData["data"]!!.jsonPrimitive.isString).isTrue() + + val restored = event.toDto().toAdk().content!!.parts + assertThat(restored[0].functionCall?.name).isEqualTo("get_weather") + assertThat(restored[0].thoughtSignature?.toList()).isEqualTo(signature.toList()) + assertThat(restored[1].inlineData?.data?.toList()).isEqualTo(byteArrayOf(1, 2, 3).toList()) + } +} diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index dc79402a..e9ee7439 100644 --- a/gradle/libs.versions.toml +++ b/gradle/libs.versions.toml @@ -102,6 +102,7 @@ kxml2 = { module = "net.sf.kxml:kxml2", version.ref = "kxml2" } mcp = { module = "io.modelcontextprotocol.sdk:mcp", version.ref = "mcp" } mockito-android = { module = "org.mockito:mockito-android", version.ref = "mockito-android" } mockito-kotlin = { module = "org.mockito.kotlin:mockito-kotlin", version.ref = "mockito-kotlin" } +okhttp = { module = "com.squareup.okhttp3:okhttp", version.ref = "okhttp" } okhttp-mockwebserver = { module = "com.squareup.okhttp3:mockwebserver", version.ref = "okhttp" } opentelemetry-api = { module = "io.opentelemetry:opentelemetry-api", version.ref = "opentelemetry" } opentelemetry-sdk = { module = "io.opentelemetry:opentelemetry-sdk", version.ref = "opentelemetry" }