feat: Adding a SessionKey for typeSafety

google-genai-bot · copybara-github · commit d899f6f4ad52 · 2026-03-02T06:40:32.000-08:00
PiperOrigin-RevId: 877372854
diff --git a/core/src/main/java/com/google/adk/artifacts/BaseArtifactService.java b/core/src/main/java/com/google/adk/artifacts/BaseArtifactService.java
@@ -16,6 +16,7 @@
 
 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;
@@ -39,6 +40,12 @@ public interface BaseArtifactService {
   Single<Integer> saveArtifact(
       String appName, String userId, String sessionId, String filename, Part artifact);
 
+  /** Saves an artifact. */
+  default Single<Integer> 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.
    *
@@ -58,18 +65,35 @@ default Single<Part> saveAndReloadArtifact(
         .flatMap(version -> loadArtifact(appName, userId, sessionId, filename, version).toSingle());
   }
 
+  /** Saves an artifact and returns it with fileData if available. */
+  default Single<Part> 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<Part> 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<Part> 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<Part> loadArtifact(
       String appName, String userId, String sessionId, String filename, int version) {
     return loadArtifact(appName, userId, sessionId, filename, Optional.of(version));
   }
 
+  default Maybe<Part> 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.
@@ -88,6 +112,10 @@ Maybe<Part> loadArtifact(
    */
   Single<ListArtifactsResponse> listArtifactKeys(String appName, String userId, String sessionId);
 
+  default Single<ListArtifactsResponse> listArtifactKeys(SessionKey sessionKey) {
+    return listArtifactKeys(sessionKey.appName(), sessionKey.userId(), sessionKey.id());
+  }
+
   /**
    * Deletes an artifact.
    *
@@ -98,6 +126,10 @@ Maybe<Part> loadArtifact(
    */
   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.
    *
@@ -109,4 +141,8 @@ Maybe<Part> loadArtifact(
    */
   Single<ImmutableList<Integer>> listVersions(
       String appName, String userId, String sessionId, String filename);
+
+  default Single<ImmutableList<Integer>> 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
@@ -35,6 +35,7 @@
 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;
@@ -383,6 +384,25 @@ public Flowable<Event> runAsync(
         .flatMapPublisher(session -> this.runAsyncImpl(session, newMessage, runConfig, stateDelta));
   }
 
+  /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */
+  public Flowable<Event> runAsync(
+      SessionKey sessionKey,
+      Content newMessage,
+      RunConfig runConfig,
+      @Nullable Map<String, Object> stateDelta) {
+    return runAsync(sessionKey.userId(), sessionKey.id(), newMessage, runConfig, stateDelta);
+  }
+
+  /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */
+  public Flowable<Event> runAsync(SessionKey sessionKey, Content newMessage, RunConfig runConfig) {
+    return runAsync(sessionKey, newMessage, runConfig, /* stateDelta= */ null);
+  }
+
+  /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */
+  public Flowable<Event> runAsync(SessionKey sessionKey, Content newMessage) {
+    return runAsync(sessionKey, newMessage, RunConfig.builder().build());
+  }
+
   /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */
   public Flowable<Event> runAsync(String userId, String sessionId, Content newMessage) {
     return runAsync(userId, sessionId, newMessage, RunConfig.builder().build());
@@ -671,6 +691,17 @@ public Flowable<Event> 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<Event> 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
@@ -78,6 +78,18 @@ default Single<Session> 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<Session> createSession(
+      SessionKey sessionKey, @Nullable Map<String, Object> 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.
@@ -94,6 +106,14 @@ default Single<Session> 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<Session> createSession(SessionKey sessionKey) {
+    return createSession(sessionKey.appName(), sessionKey.userId(), null, sessionKey.id());
+  }
+
   /**
    * Retrieves a specific session, optionally filtering the events included.
    *
@@ -110,6 +130,12 @@ default Single<Session> createSession(String appName, String userId) {
   Maybe<Session> getSession(
       String appName, String userId, String sessionId, Optional<GetSessionConfig> config);
 
+  /** Retrieves a specific session, optionally filtering the events included. */
+  default Maybe<Session> 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.
    *
@@ -123,6 +149,11 @@ Maybe<Session> getSession(
    */
   Single<ListSessionsResponse> listSessions(String appName, String userId);
 
+  /** Lists sessions associated with a specific application and user. */
+  default Single<ListSessionsResponse> listSessions(SessionKey sessionKey) {
+    return listSessions(sessionKey.appName(), sessionKey.userId());
+  }
+
   /**
    * Deletes a specific session.
    *
@@ -134,6 +165,11 @@ Maybe<Session> getSession(
    */
   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.
    *
@@ -147,6 +183,11 @@ Maybe<Session> getSession(
    */
   Single<ListEventsResponse> listEvents(String appName, String userId, String sessionId);
 
+  /** Lists the events within a specific session. */
+  default Single<ListEventsResponse> 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.
@@ -190,20 +231,18 @@ default Single<Event> appendEvent(Session session, Event event) {
     EventActions actions = event.actions();
     if (actions != null) {
       Map<String, Object> stateDelta = actions.stateDelta();
-      if (stateDelta != null && !stateDelta.isEmpty()) {
-        Map<String, Object> 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);
-                  }
+      Map<String, Object> 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);
                 }
-              });
-        }
+              }
+            });
       }
     }
 
diff --git a/core/src/main/java/com/google/adk/sessions/Session.java b/core/src/main/java/com/google/adk/sessions/Session.java
@@ -49,6 +49,11 @@ 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;
@@ -62,6 +67,13 @@ 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() {}
 
@@ -72,6 +84,15 @@ 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;
@@ -130,6 +151,11 @@ 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
@@ -0,0 +1,82 @@
+/*
+ * 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
