From a0fd4686a7030031d12bae097503658546bc4ae5 Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Mon, 2 Mar 2026 07:19:00 -0800 Subject: [PATCH] feat: Adding a SessionKey for typeSafety PiperOrigin-RevId: 877386034 --- .../adk/artifacts/BaseArtifactService.java | 36 ------ .../java/com/google/adk/runner/Runner.java | 31 ------ .../adk/sessions/BaseSessionService.java | 65 +++-------- .../java/com/google/adk/sessions/Session.java | 26 ----- .../com/google/adk/sessions/SessionKey.java | 82 -------------- .../com/google/adk/runner/RunnerTest.java | 103 ------------------ 6 files changed, 13 insertions(+), 330 deletions(-) delete mode 100644 core/src/main/java/com/google/adk/sessions/SessionKey.java diff --git a/core/src/main/java/com/google/adk/artifacts/BaseArtifactService.java b/core/src/main/java/com/google/adk/artifacts/BaseArtifactService.java index a9bb6ba4d..32ef9ff4d 100644 --- a/core/src/main/java/com/google/adk/artifacts/BaseArtifactService.java +++ b/core/src/main/java/com/google/adk/artifacts/BaseArtifactService.java @@ -16,7 +16,6 @@ package com.google.adk.artifacts; -import com.google.adk.sessions.SessionKey; import com.google.common.collect.ImmutableList; import com.google.genai.types.Part; import io.reactivex.rxjava3.core.Completable; @@ -40,12 +39,6 @@ public interface BaseArtifactService { Single saveArtifact( String appName, String userId, String sessionId, String filename, Part artifact); - /** Saves an artifact. */ - default Single saveArtifact(SessionKey sessionKey, String filename, Part artifact) { - return saveArtifact( - sessionKey.appName(), sessionKey.userId(), sessionKey.id(), filename, artifact); - } - /** * Saves an artifact and returns it with fileData if available. * @@ -65,35 +58,18 @@ default Single saveAndReloadArtifact( .flatMap(version -> loadArtifact(appName, userId, sessionId, filename, version).toSingle()); } - /** Saves an artifact and returns it with fileData if available. */ - default Single saveAndReloadArtifact( - SessionKey sessionKey, String filename, Part artifact) { - return saveAndReloadArtifact( - sessionKey.appName(), sessionKey.userId(), sessionKey.id(), filename, artifact); - } - /** Loads the latest version of an artifact from the service. */ default Maybe loadArtifact( String appName, String userId, String sessionId, String filename) { return loadArtifact(appName, userId, sessionId, filename, Optional.empty()); } - /** Loads the latest version of an artifact from the service. */ - default Maybe loadArtifact(SessionKey sessionKey, String filename) { - return loadArtifact(sessionKey.appName(), sessionKey.userId(), sessionKey.id(), filename); - } - /** Loads a specific version of an artifact from the service. */ default Maybe loadArtifact( String appName, String userId, String sessionId, String filename, int version) { return loadArtifact(appName, userId, sessionId, filename, Optional.of(version)); } - default Maybe loadArtifact(SessionKey sessionKey, String filename, int version) { - return loadArtifact( - sessionKey.appName(), sessionKey.userId(), sessionKey.id(), filename, version); - } - /** * @deprecated Use {@link #loadArtifact(String, String, String, String)} or {@link * #loadArtifact(String, String, String, String, int)} instead. @@ -112,10 +88,6 @@ Maybe loadArtifact( */ Single listArtifactKeys(String appName, String userId, String sessionId); - default Single listArtifactKeys(SessionKey sessionKey) { - return listArtifactKeys(sessionKey.appName(), sessionKey.userId(), sessionKey.id()); - } - /** * Deletes an artifact. * @@ -126,10 +98,6 @@ default Single listArtifactKeys(SessionKey sessionKey) { */ Completable deleteArtifact(String appName, String userId, String sessionId, String filename); - default Completable deleteArtifact(SessionKey sessionKey, String filename) { - return deleteArtifact(sessionKey.appName(), sessionKey.userId(), sessionKey.id(), filename); - } - /** * Lists all the versions (as revision IDs) of an artifact. * @@ -141,8 +109,4 @@ default Completable deleteArtifact(SessionKey sessionKey, String filename) { */ Single> listVersions( String appName, String userId, String sessionId, String filename); - - default Single> listVersions(SessionKey sessionKey, String filename) { - return listVersions(sessionKey.appName(), sessionKey.userId(), sessionKey.id(), filename); - } } diff --git a/core/src/main/java/com/google/adk/runner/Runner.java b/core/src/main/java/com/google/adk/runner/Runner.java index ea7cb80f6..0ddfdaea1 100644 --- a/core/src/main/java/com/google/adk/runner/Runner.java +++ b/core/src/main/java/com/google/adk/runner/Runner.java @@ -35,7 +35,6 @@ import com.google.adk.sessions.BaseSessionService; import com.google.adk.sessions.InMemorySessionService; import com.google.adk.sessions.Session; -import com.google.adk.sessions.SessionKey; import com.google.adk.summarizer.EventsCompactionConfig; import com.google.adk.summarizer.LlmEventSummarizer; import com.google.adk.summarizer.SlidingWindowEventCompactor; @@ -384,25 +383,6 @@ public Flowable runAsync( .flatMapPublisher(session -> this.runAsyncImpl(session, newMessage, runConfig, stateDelta)); } - /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */ - public Flowable runAsync( - SessionKey sessionKey, - Content newMessage, - RunConfig runConfig, - @Nullable Map stateDelta) { - return runAsync(sessionKey.userId(), sessionKey.id(), newMessage, runConfig, stateDelta); - } - - /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */ - public Flowable runAsync(SessionKey sessionKey, Content newMessage, RunConfig runConfig) { - return runAsync(sessionKey, newMessage, runConfig, /* stateDelta= */ null); - } - - /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */ - public Flowable runAsync(SessionKey sessionKey, Content newMessage) { - return runAsync(sessionKey, newMessage, RunConfig.builder().build()); - } - /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */ public Flowable runAsync(String userId, String sessionId, Content newMessage) { return runAsync(userId, sessionId, newMessage, RunConfig.builder().build()); @@ -691,17 +671,6 @@ public Flowable runLive( .flatMapPublisher(session -> this.runLive(session, liveRequestQueue, runConfig)); } - /** - * Retrieves the session and runs the agent in live mode. - * - * @return stream of events from the agent. - * @throws IllegalArgumentException if the session is not found. - */ - public Flowable runLive( - SessionKey sessionKey, LiveRequestQueue liveRequestQueue, RunConfig runConfig) { - return runLive(sessionKey.userId(), sessionKey.id(), liveRequestQueue, runConfig); - } - /** * Runs the agent asynchronously with a default user ID. * diff --git a/core/src/main/java/com/google/adk/sessions/BaseSessionService.java b/core/src/main/java/com/google/adk/sessions/BaseSessionService.java index 7a0885544..94e8cd7ba 100644 --- a/core/src/main/java/com/google/adk/sessions/BaseSessionService.java +++ b/core/src/main/java/com/google/adk/sessions/BaseSessionService.java @@ -78,18 +78,6 @@ default Single createSession( return createSession(appName, userId, ensureConcurrentMap(state), sessionId); } - /** - * Creates a new session with the specified parameters. - * - * @param sessionKey The session key containing appName, userId and sessionId. - * @param state An optional map representing the initial state of the session. Can be null or - * empty. - */ - default Single createSession( - SessionKey sessionKey, @Nullable Map state) { - return createSession(sessionKey.appName(), sessionKey.userId(), state, sessionKey.id()); - } - /** * Creates a new session with the specified application name and user ID, using a default state * (null) and allowing the service to generate a unique session ID. @@ -106,14 +94,6 @@ default Single createSession(String appName, String userId) { return createSession(appName, userId, null, null); } - /** - * Creates a new session with the specified application name and user ID, using a default state - * (null) and allowing the service to generate a unique session ID. - */ - default Single createSession(SessionKey sessionKey) { - return createSession(sessionKey.appName(), sessionKey.userId(), null, sessionKey.id()); - } - /** * Retrieves a specific session, optionally filtering the events included. * @@ -130,12 +110,6 @@ default Single createSession(SessionKey sessionKey) { Maybe getSession( String appName, String userId, String sessionId, Optional config); - /** Retrieves a specific session, optionally filtering the events included. */ - default Maybe getSession(SessionKey sessionKey, @Nullable GetSessionConfig config) { - return getSession( - sessionKey.appName(), sessionKey.userId(), sessionKey.id(), Optional.ofNullable(config)); - } - /** * Lists sessions associated with a specific application and user. * @@ -149,11 +123,6 @@ default Maybe getSession(SessionKey sessionKey, @Nullable GetSessionCon */ Single listSessions(String appName, String userId); - /** Lists sessions associated with a specific application and user. */ - default Single listSessions(SessionKey sessionKey) { - return listSessions(sessionKey.appName(), sessionKey.userId()); - } - /** * Deletes a specific session. * @@ -165,11 +134,6 @@ default Single listSessions(SessionKey sessionKey) { */ Completable deleteSession(String appName, String userId, String sessionId); - /** Deletes a specific session. */ - default Completable deleteSession(SessionKey sessionKey) { - return deleteSession(sessionKey.appName(), sessionKey.userId(), sessionKey.id()); - } - /** * Lists the events within a specific session. Supports pagination via the response object. * @@ -183,11 +147,6 @@ default Completable deleteSession(SessionKey sessionKey) { */ Single listEvents(String appName, String userId, String sessionId); - /** Lists the events within a specific session. */ - default Single listEvents(SessionKey sessionKey) { - return listEvents(sessionKey.appName(), sessionKey.userId(), sessionKey.id()); - } - /** * Closes a session. This is currently a placeholder and may involve finalizing session state or * performing cleanup actions in future implementations. The default implementation does nothing. @@ -231,18 +190,20 @@ default Single appendEvent(Session session, Event event) { EventActions actions = event.actions(); if (actions != null) { Map stateDelta = actions.stateDelta(); - Map sessionState = session.state(); - if (stateDelta != null && !stateDelta.isEmpty() && sessionState != null) { - stateDelta.forEach( - (key, value) -> { - if (!key.startsWith(State.TEMP_PREFIX)) { - if (value == State.REMOVED) { - sessionState.remove(key); - } else { - sessionState.put(key, value); + if (stateDelta != null && !stateDelta.isEmpty()) { + Map sessionState = session.state(); + if (sessionState != null) { + stateDelta.forEach( + (key, value) -> { + if (!key.startsWith(State.TEMP_PREFIX)) { + if (value == State.REMOVED) { + sessionState.remove(key); + } else { + sessionState.put(key, value); + } } - } - }); + }); + } } } diff --git a/core/src/main/java/com/google/adk/sessions/Session.java b/core/src/main/java/com/google/adk/sessions/Session.java index f8376589a..877a95220 100644 --- a/core/src/main/java/com/google/adk/sessions/Session.java +++ b/core/src/main/java/com/google/adk/sessions/Session.java @@ -49,11 +49,6 @@ public static Builder builder(String id) { return new Builder(id); } - /** Creates a new {@link Builder} with the given session key. */ - public static Builder builder(SessionKey sessionKey) { - return new Builder(sessionKey); - } - /** Builder for {@link Session}. */ public static final class Builder { private String id; @@ -67,13 +62,6 @@ public Builder(String id) { this.id = id; } - /** Creates a new {@link Builder} with the given session key. */ - public Builder(SessionKey sessionKey) { - this.id = sessionKey.id(); - this.appName = sessionKey.appName(); - this.userId = sessionKey.userId(); - } - @JsonCreator private Builder() {} @@ -84,15 +72,6 @@ public Builder id(String id) { return this; } - /** Sets the session key. */ - @CanIgnoreReturnValue - public Builder sessionKey(SessionKey sessionKey) { - this.id = sessionKey.id(); - this.appName = sessionKey.appName(); - this.userId = sessionKey.userId(); - return this; - } - @CanIgnoreReturnValue public Builder state(State state) { this.state = state; @@ -151,11 +130,6 @@ public Session build() { } } - /** Returns the session key. */ - public SessionKey sessionKey() { - return new SessionKey(appName, userId, id); - } - @JsonProperty("id") public String id() { return id; diff --git a/core/src/main/java/com/google/adk/sessions/SessionKey.java b/core/src/main/java/com/google/adk/sessions/SessionKey.java deleted file mode 100644 index db26b5a3a..000000000 --- a/core/src/main/java/com/google/adk/sessions/SessionKey.java +++ /dev/null @@ -1,82 +0,0 @@ -/* - * Copyright 2025 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.sessions; - -import com.fasterxml.jackson.annotation.JsonCreator; -import com.fasterxml.jackson.annotation.JsonProperty; -import com.google.adk.JsonBaseModel; -import java.util.Objects; - -/** Key for a session, composed of appName, userId and session id. */ -public final class SessionKey extends JsonBaseModel { - private final String appName; - private final String userId; - private final String id; - - @JsonCreator - public SessionKey( - @JsonProperty("appName") String appName, - @JsonProperty("userId") String userId, - @JsonProperty("id") String id) { - this.appName = appName; - this.userId = userId; - this.id = id; - } - - @JsonProperty("appName") - public String appName() { - return appName; - } - - @JsonProperty("userId") - public String userId() { - return userId; - } - - @JsonProperty("id") - public String id() { - return id; - } - - @Override - public boolean equals(Object o) { - if (this == o) { - return true; - } - if (o == null || getClass() != o.getClass()) { - return false; - } - SessionKey that = (SessionKey) o; - return Objects.equals(appName, that.appName) - && Objects.equals(userId, that.userId) - && Objects.equals(id, that.id); - } - - @Override - public int hashCode() { - return Objects.hash(appName, userId, id); - } - - @Override - public String toString() { - return toJson(); - } - - public static SessionKey fromJson(String json) { - return fromJsonString(json, SessionKey.class); - } -} diff --git a/core/src/test/java/com/google/adk/runner/RunnerTest.java b/core/src/test/java/com/google/adk/runner/RunnerTest.java index 42452c6a0..421b79abb 100644 --- a/core/src/test/java/com/google/adk/runner/RunnerTest.java +++ b/core/src/test/java/com/google/adk/runner/RunnerTest.java @@ -41,7 +41,6 @@ import com.google.adk.models.LlmResponse; import com.google.adk.plugins.BasePlugin; import com.google.adk.sessions.Session; -import com.google.adk.sessions.SessionKey; import com.google.adk.summarizer.EventsCompactionConfig; import com.google.adk.telemetry.Tracing; import com.google.adk.testing.TestLlm; @@ -579,14 +578,6 @@ public void onEventCallback_success() { verify(plugin).onEventCallback(any(), any()); } - @Test - public void runAsync_withSessionKey_success() { - var events = - runner.runAsync(session.sessionKey(), createContent("from user")).toList().blockingGet(); - - assertThat(simplifyEvents(events)).containsExactly("test agent: from llm"); - } - @Test public void runAsync_withStateDelta_mergesStateIntoSession() { ImmutableMap stateDelta = ImmutableMap.of("key1", "value1", "key2", 42); @@ -614,32 +605,6 @@ public void runAsync_withStateDelta_mergesStateIntoSession() { assertThat(finalSession.state()).containsAtLeastEntriesIn(stateDelta); } - @Test - public void runAsync_withSessionKeyAndStateDelta_mergesStateIntoSession() { - ImmutableMap stateDelta = ImmutableMap.of("key1", "value1", "key2", 42); - - var events = - runner - .runAsync( - session.sessionKey(), - createContent("test message"), - RunConfig.builder().build(), - stateDelta) - .toList() - .blockingGet(); - - // Verify agent runs successfully - assertThat(simplifyEvents(events)).containsExactly("test agent: from llm"); - - // Verify state was merged into session - Session finalSession = - runner - .sessionService() - .getSession("test", "user", session.id(), Optional.empty()) - .blockingGet(); - assertThat(finalSession.state()).containsAtLeastEntriesIn(stateDelta); - } - @Test public void runAsync_withEmptyStateDelta_doesNotModifySession() { ImmutableMap emptyStateDelta = ImmutableMap.of(); @@ -875,20 +840,6 @@ public void runLive_success() throws Exception { assertThat(simplifyEvents(testSubscriber.values())).containsExactly("test agent: from llm"); } - @Test - public void runLive_withSessionKey_success() throws Exception { - LiveRequestQueue liveRequestQueue = new LiveRequestQueue(); - TestSubscriber testSubscriber = - runner.runLive(session.sessionKey(), liveRequestQueue, RunConfig.builder().build()).test(); - - liveRequestQueue.content(createContent("from user")); - liveRequestQueue.close(); - - testSubscriber.await(); - testSubscriber.assertComplete(); - assertThat(simplifyEvents(testSubscriber.values())).containsExactly("test agent: from llm"); - } - @Test public void runLive_withToolExecution() throws Exception { LlmAgent agentWithTool = @@ -997,18 +948,6 @@ public void runAsync_withoutSessionAndAutoCreateSessionTrue_createsSession() { .isNotNull(); } - @Test - public void runAsync_withoutSessionAndAutoCreateSessionTrue_withSessionKey_createsSession() { - RunConfig runConfig = RunConfig.builder().setAutoCreateSession(true).build(); - SessionKey sessionKey = new SessionKey("test", "user", UUID.randomUUID().toString()); - - var events = - runner.runAsync(sessionKey, createContent("from user"), runConfig).toList().blockingGet(); - - assertThat(simplifyEvents(events)).containsExactly("test agent: from llm"); - assertThat(runner.sessionService().getSession(sessionKey, null).blockingGet()).isNotNull(); - } - @Test public void runAsync_withoutSessionAndAutoCreateSessionFalse_throwsException() { RunConfig runConfig = RunConfig.builder().setAutoCreateSession(false).build(); @@ -1020,17 +959,6 @@ public void runAsync_withoutSessionAndAutoCreateSessionFalse_throwsException() { .assertError(IllegalArgumentException.class); } - @Test - public void runAsync_withoutSessionAndAutoCreateSessionFalse_withSessionKey_throwsException() { - RunConfig runConfig = RunConfig.builder().setAutoCreateSession(false).build(); - SessionKey sessionKey = new SessionKey("test", "user", UUID.randomUUID().toString()); - - runner - .runAsync(sessionKey, createContent("from user"), runConfig) - .test() - .assertError(IllegalArgumentException.class); - } - @Test public void runLive_withoutSessionAndAutoCreateSessionTrue_createsSession() throws Exception { RunConfig runConfig = RunConfig.builder().setAutoCreateSession(true).build(); @@ -1054,25 +982,6 @@ public void runLive_withoutSessionAndAutoCreateSessionTrue_createsSession() thro .isNotNull(); } - @Test - public void runLive_withoutSessionAndAutoCreateSessionTrue_withSessionKey_createsSession() - throws Exception { - RunConfig runConfig = RunConfig.builder().setAutoCreateSession(true).build(); - SessionKey sessionKey = new SessionKey("test", "user", UUID.randomUUID().toString()); - LiveRequestQueue liveRequestQueue = new LiveRequestQueue(); - - TestSubscriber testSubscriber = - runner.runLive(sessionKey, liveRequestQueue, runConfig).test(); - - liveRequestQueue.content(createContent("from user")); - liveRequestQueue.close(); - - testSubscriber.await(); - testSubscriber.assertComplete(); - assertThat(simplifyEvents(testSubscriber.values())).containsExactly("test agent: from llm"); - assertThat(runner.sessionService().getSession(sessionKey, null).blockingGet()).isNotNull(); - } - @Test public void runLive_withoutSessionAndAutoCreateSessionFalse_throwsException() { RunConfig runConfig = RunConfig.builder().setAutoCreateSession(false).build(); @@ -1085,18 +994,6 @@ public void runLive_withoutSessionAndAutoCreateSessionFalse_throwsException() { .assertError(IllegalArgumentException.class); } - @Test - public void runLive_withoutSessionAndAutoCreateSessionFalse_withSessionKey_throwsException() { - RunConfig runConfig = RunConfig.builder().setAutoCreateSession(false).build(); - SessionKey sessionKey = new SessionKey("test", "user", UUID.randomUUID().toString()); - LiveRequestQueue liveRequestQueue = new LiveRequestQueue(); - - runner - .runLive(sessionKey, liveRequestQueue, runConfig) - .test() - .assertError(IllegalArgumentException.class); - } - @Test public void runAsync_withToolConfirmation() { TestLlm testLlm =