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" }