feat(sessions): support session expiry in VertexAiSessionService

A Vertex-only `createSession` overload now takes `ttl` or `expireTime`. At most
one may be set, since the wire field is a single choice. Without one the backend
applies its own default, which is long; with one, the minimum it accepts is 24
hours. The `SessionService` interface is unchanged.

PiperOrigin-RevId: 970431402

Author: kvmilos

core/src/jvmMain/kotlin/com/google/adk/kt/sessions/VertexAiSessionService.kt

@@ -25,6 +25,8 @@ import com.google.auth.oauth2.GoogleCredentials
 import io.ktor.client.HttpClient
 import io.ktor.client.engine.java.Java
 import kotlin.jvm.JvmStatic
+import kotlin.time.Duration
+import kotlin.time.Instant
 
 /**
  * A [SessionService] backed by the managed Vertex AI Session Service.
@@ -86,8 +88,36 @@ internal constructor(
     reasoningEngineId,
   )
 
-  override suspend fun createSession(key: SessionKey, state: Map<String, Any>?): Session {
-    val sessionDto = client.createSession(engine, key.userId, state).getOrThrow()
+  override suspend fun createSession(key: SessionKey, state: Map<String, Any>?): Session =
+    createSession(key, state, ttl = null, expireTime = null)
+
+  /**
+   * Creates a session that expires, addressing the reasoning engine fixed at construction.
+   *
+   * At most one of [ttl] and [expireTime] may be set, because the backend models them as a single
+   * choice; setting both is rejected. The backend also requires the expiry to be at least 24 hours
+   * out, and applies its own default when neither is given.
+   *
+   * @param key The composite identifier of the session; [SessionKey.appName] is only a label.
+   * @param state An optional map representing the initial state of the session.
+   * @param ttl How long the session lives, measured from creation. Sub-second precision is dropped.
+   * @param expireTime The absolute instant at which the session expires.
+   * @return The newly created [Session].
+   */
+  suspend fun createSession(
+    key: SessionKey,
+    state: Map<String, Any>? = null,
+    ttl: Duration? = null,
+    expireTime: Instant? = null,
+  ): Session {
+    require(ttl == null || expireTime == null) {
+      "Cannot specify both ttl and expireTime simultaneously."
+    }
+    // The wire format is whole seconds, so a sub-second ttl would silently travel as "0s".
+    require(ttl == null || ttl.inWholeSeconds > 0) {
+      "ttl must be at least one second, but was $ttl."
+    }
+    val sessionDto = client.createSession(engine, key.userId, state, ttl, expireTime).getOrThrow()
     return sessionDto.toAdk(key.appName, key.userId, key.id)
   }
 

core/src/jvmMain/kotlin/com/google/adk/kt/sessions/VertexAiSessionsClient.kt

@@ -38,6 +38,8 @@ import java.io.IOException
 import java.net.URLEncoder
 import java.nio.charset.StandardCharsets
 import java.util.concurrent.TimeoutException
+import kotlin.time.Duration
+import kotlin.time.Instant
 import kotlinx.coroutines.delay
 import kotlinx.serialization.DeserializationStrategy
 
@@ -101,17 +103,28 @@ internal open class VertexAiSessionsClient(
    * @param engine The reasoning engine that owns the session.
    * @param userId The ID of the user owning the session.
    * @param state Optional initial state map to inject into the session.
+   * @param ttl Optional lifetime of the session, measured from creation.
+   * @param expireTime Optional absolute expiry of the session. The wire field is a oneof, so the
+   *   caller must not set both this and [ttl].
    * @return The created [SessionDto], or a [Result.failure] describing why creation failed.
    */
   open suspend fun createSession(
     engine: ReasoningEngineRef,
     userId: String,
     state: Map<String, Any>?,
+    ttl: Duration? = null,
+    expireTime: Instant? = null,
   ): Result<SessionDto> {
     val requestBody =
       adkJson.encodeToString(
         CreateSessionRequestDto.serializer(),
-        CreateSessionRequestDto(userId = userId, sessionState = state?.let { anyToJsonElement(it) }),
+        CreateSessionRequestDto(
+          userId = userId,
+          sessionState = state?.let { anyToJsonElement(it) },
+          // proto3 JSON writes a Duration as seconds with an "s" suffix.
+          ttl = ttl?.let { "${it.inWholeSeconds}s" },
+          expireTime = expireTime?.toString(),
+        ),
       )
     val createResponse =
       postAndDecode(

core/src/jvmMain/kotlin/com/google/adk/kt/sessions/dto/ResponsesDto.kt

@@ -38,4 +38,7 @@ internal data class ListEventsResponseDto(
 internal data class CreateSessionRequestDto(
   val userId: String,
   val sessionState: JsonElement? = null,
+  // `ttl`/`expire_time` are a proto oneof; at most one is set, in its JSON string form.
+  val ttl: String? = null,
+  val expireTime: String? = null,
 )

core/src/jvmTest/kotlin/com/google/adk/kt/sessions/VertexAiSessionServiceTest.kt

@@ -28,7 +28,12 @@ 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.Duration
+import kotlin.time.Duration.Companion.hours
+import kotlin.time.Duration.Companion.milliseconds
+import kotlin.time.Duration.Companion.seconds
 import kotlin.time.Instant
+import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.runTest
 import kotlinx.serialization.json.JsonObject
 import kotlinx.serialization.json.JsonPrimitive
@@ -62,19 +67,24 @@ class VertexAiSessionServiceTest {
       reasoningEngineId = ENGINE_ID,
     )
 
+  /** A client that accepts any create call, for asserting which expiration was forwarded. */
+  private fun expiringSessionClient() =
+    mock<VertexAiSessionsClient> {
+      onBlocking { createSession(any(), any(), anyOrNull(), anyOrNull(), anyOrNull()) } doReturn
+        Result.success(SessionDto(name = "reasoningEngines/123/sessions/s"))
+    }
+
   @Test
   fun addressesConfiguredEngineRegardlessOfAppName() = runTest {
-    val client =
-      mock<VertexAiSessionsClient> {
-        onBlocking { createSession(any(), any(), anyOrNull()) } doReturn
-          Result.success(SessionDto(name = "reasoningEngines/123/sessions/s"))
-      }
+    val client = expiringSessionClient()
 
     // The app name is only a label; the service always addresses the engine set at construction.
     val unused =
       service(client).createSession(SessionKey("any-label", "user", id = null), state = null)
 
-    verifyBlocking(client) { createSession(eq(ENGINE), eq("user"), anyOrNull()) }
+    verifyBlocking(client) {
+      createSession(eq(ENGINE), eq("user"), anyOrNull(), anyOrNull(), anyOrNull())
+    }
   }
 
   @Test
@@ -137,7 +147,9 @@ class VertexAiSessionServiceTest {
   fun createSession_mapsClientResponse() = runTest {
     val client =
       mock<VertexAiSessionsClient> {
-        onBlocking { createSession(eq(ENGINE), eq("user"), anyOrNull()) } doReturn
+        onBlocking {
+          createSession(eq(ENGINE), eq("user"), anyOrNull(), anyOrNull(), anyOrNull())
+        } doReturn
           Result.success(
             SessionDto(
               name = "reasoningEngines/123/sessions/session-1",
@@ -160,7 +172,7 @@ class VertexAiSessionServiceTest {
   fun createSession_clientFails_propagates() = runTest {
     val client =
       mock<VertexAiSessionsClient> {
-        onBlocking { createSession(any(), any(), anyOrNull()) } doReturn
+        onBlocking { createSession(any(), any(), anyOrNull(), anyOrNull(), anyOrNull()) } doReturn
           Result.failure(IOException("boom"))
       }
 
@@ -169,6 +181,85 @@ class VertexAiSessionServiceTest {
     }
   }
 
+  @Test
+  fun createSession_ttl_forwardsTtlOnly() {
+    val client = expiringSessionClient()
+
+    runBlocking {
+      val unused =
+        service(client).createSession(SessionKey("123", "user", id = null), ttl = 24.hours)
+    }
+
+    verifyBlocking(client) {
+      createSession(eq(ENGINE), eq("user"), anyOrNull(), eq(24.hours), eq(null))
+    }
+  }
+
+  @Test
+  fun createSession_expireTime_forwardsExpireTimeOnly() {
+    val client = expiringSessionClient()
+    val expiry = Instant.parse("2026-10-01T00:00:00Z")
+
+    runBlocking {
+      val unused =
+        service(client).createSession(SessionKey("123", "user", id = null), expireTime = expiry)
+    }
+
+    verifyBlocking(client) {
+      createSession(eq(ENGINE), eq("user"), anyOrNull(), eq(null), eq(expiry))
+    }
+  }
+
+  @Test
+  fun createSession_noExpiration_forwardsNeither() {
+    val client = expiringSessionClient()
+
+    // The SessionService overload must not invent an expiration of its own.
+    runBlocking {
+      val unused = service(client).createSession(SessionKey("123", "user", id = null))
+    }
+
+    verifyBlocking(client) {
+      createSession(eq(ENGINE), eq("user"), anyOrNull(), eq(null), eq(null))
+    }
+  }
+
+  @Test
+  fun createSession_ttlAndExpireTime_throwsWithoutCallingBackend() {
+    val client = expiringSessionClient()
+
+    assertFailsWith<IllegalArgumentException> {
+      runBlocking {
+        service(client)
+          .createSession(
+            SessionKey("123", "user", id = null),
+            ttl = 24.hours,
+            expireTime = Instant.parse("2026-10-01T00:00:00Z"),
+          )
+      }
+    }
+    verifyBlocking(client, never()) {
+      createSession(any(), any(), anyOrNull(), anyOrNull(), anyOrNull())
+    }
+  }
+
+  @Test
+  fun createSession_ttlBelowOneSecond_throwsWithoutCallingBackend() {
+    val client = expiringSessionClient()
+
+    // 500ms is positive but truncates to "0s" on the wire, so it must be rejected too.
+    for (bad in listOf(Duration.ZERO, (-1).seconds, 500.milliseconds)) {
+      assertFailsWith<IllegalArgumentException> {
+        runBlocking {
+          service(client).createSession(SessionKey("123", "user", id = null), ttl = bad)
+        }
+      }
+    }
+    verifyBlocking(client, never()) {
+      createSession(any(), any(), anyOrNull(), anyOrNull(), anyOrNull())
+    }
+  }
+
   @Test
   fun getSession_notFound_returnsNull() = runTest {
     val client =

core/src/jvmTest/kotlin/com/google/adk/kt/sessions/VertexAiSessionsClientTest.kt

@@ -24,9 +24,12 @@ 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.time.Clock
+import kotlin.time.Duration.Companion.days
+import kotlin.time.Duration.Companion.hours
+import kotlin.time.Duration.Companion.milliseconds
+import kotlin.time.Instant
 import kotlinx.coroutines.runBlocking
 import kotlinx.serialization.json.JsonObject
 import kotlinx.serialization.json.JsonPrimitive
@@ -108,9 +111,7 @@ class VertexAiSessionsClientTest {
 
   @Test
   fun createSession_nestedState_serializedAsJsonNotToString() {
-    server.enqueue(jsonResponse("""{"name":"reasoningEngines/123/sessions/s/operations/op"}"""))
-    server.enqueue(jsonResponse("""{"name":"operations/op","done":true}"""))
-    server.enqueue(jsonResponse("""{"name":"reasoningEngines/123/sessions/s"}"""))
+    enqueueCreateSessionExchange()
 
     val unused = runBlocking {
       client.createSession(ENGINE, "user", mapOf("nested" to mapOf("a" to listOf(1L, 2L))))
@@ -121,6 +122,54 @@ class VertexAiSessionsClientTest {
     assertThat(body).contains("\"nested\":{\"a\":[1,2]}")
   }
 
+  @Test
+  fun createSession_ttl_sendsProtoDurationAndNoExpireTime() {
+    enqueueCreateSessionExchange()
+
+    val unused = runBlocking { client.createSession(ENGINE, "user", null, ttl = 24.hours) }
+
+    val body = server.takeRequest().body?.utf8()
+    assertThat(body).contains("\"ttl\":\"86400s\"")
+    // `ttl` and `expireTime` are a wire oneof, so the unset arm must not be sent at all.
+    assertThat(body).doesNotContain("expireTime")
+  }
+
+  @Test
+  fun createSession_subSecondTtl_truncatesToWholeSeconds() {
+    enqueueCreateSessionExchange()
+
+    val unused = runBlocking {
+      client.createSession(ENGINE, "user", null, ttl = 24.hours + 900.milliseconds)
+    }
+
+    assertThat(server.takeRequest().body?.utf8()).contains("\"ttl\":\"86400s\"")
+  }
+
+  @Test
+  fun createSession_expireTime_sendsRfc3339AndNoTtl() {
+    enqueueCreateSessionExchange()
+
+    val unused = runBlocking {
+      client.createSession(ENGINE, "user", null, expireTime = EXPIRE_TIME)
+    }
+
+    val body = server.takeRequest().body?.utf8()
+    assertThat(body).contains("\"expireTime\":\"2026-10-01T00:00:00Z\"")
+    assertThat(body).doesNotContain("ttl")
+  }
+
+  @Test
+  fun createSession_noExpiration_sendsNeitherField() {
+    enqueueCreateSessionExchange()
+
+    val unused = runBlocking { client.createSession(ENGINE, "user", null) }
+
+    // An unset arm is omitted entirely rather than sent as an explicit JSON null.
+    val body = server.takeRequest().body?.utf8()
+    assertThat(body).doesNotContain("ttl")
+    assertThat(body).doesNotContain("expireTime")
+  }
+
   @Test
   fun getSession_notFound_returnsNullSuccess() = runBlocking {
     server.enqueue(MockResponse(code = 404))
@@ -274,17 +323,26 @@ class VertexAiSessionsClientTest {
       .endsWith("/reasoningEngines/123/sessions/s1:appendEvent")
   }
 
+  /** Enqueues the create response, the completed operation poll, and the materialized session. */
+  private fun enqueueCreateSessionExchange() {
+    server.enqueue(jsonResponse("""{"name":"reasoningEngines/123/sessions/s/operations/op"}"""))
+    server.enqueue(jsonResponse("""{"name":"operations/op","done":true}"""))
+    server.enqueue(jsonResponse("""{"name":"reasoningEngines/123/sessions/s"}"""))
+  }
+
   private companion object {
     val ENGINE =
       ReasoningEngineRef(project = "test-project", location = "test-location", id = "123")
 
+    val EXPIRE_TIME = Instant.parse("2026-10-01T00:00:00Z")
+
     fun jsonResponse(body: String): MockResponse =
       MockResponse(headers = Headers.headersOf("Content-Type", "application/json"), body = body)
 
     fun fakeCredentials(): GoogleCredentials =
       GoogleCredentials.newBuilder()
         .setAccessToken(
-          AccessToken("fake-token", Date(Instant.now().plus(1, ChronoUnit.DAYS).toEpochMilli()))
+          AccessToken("fake-token", Date((Clock.System.now() + 1.days).toEpochMilliseconds()))
         )
         .build()
   }