Merge branch 'main' into fix/memory-leak-delete-session

Author: YuqiGuo105

core/src/main/java/com/google/adk/agents/InvocationContext.java

@@ -43,18 +43,18 @@ public class InvocationContext {
   private final BaseArtifactService artifactService;
   private final BaseMemoryService memoryService;
   private final Plugin pluginManager;
-  private final Optional<LiveRequestQueue> liveRequestQueue;
+  @Nullable private final LiveRequestQueue liveRequestQueue;
   private final Map<String, ActiveStreamingTool> activeStreamingTools;
   private final String invocationId;
   private final Session session;
-  private final Optional<Content> userContent;
+  @Nullable private final Content userContent;
   private final RunConfig runConfig;
   @Nullable private final EventsCompactionConfig eventsCompactionConfig;
   @Nullable private final ContextCacheConfig contextCacheConfig;
   private final InvocationCostManager invocationCostManager;
   private final Map<String, Object> callbackContextData;
 
-  private Optional<String> branch;
+  @Nullable private String branch;
   private BaseAgent agent;
   private boolean endInvocation;
 
@@ -153,10 +153,10 @@ public InvocationContext(
               + ".invocationId(invocationId)"
               + ".agent(agent)"
               + ".session(session)"
-              + ".userContent(Optional.ofNullable(userContent))"
+              + ".userContent(userContent)"
               + ".runConfig(runConfig)"
               + ".build()",
-      imports = {"com.google.adk.agents.InvocationContext", "java.util.Optional"})
+      imports = {"com.google.adk.agents.InvocationContext"})
   @Deprecated(forRemoval = true)
   public static InvocationContext create(
       BaseSessionService sessionService,
@@ -172,7 +172,7 @@ public static InvocationContext create(
         .invocationId(invocationId)
         .agent(agent)
         .session(session)
-        .userContent(Optional.ofNullable(userContent))
+        .userContent(userContent)
         .runConfig(runConfig)
         .build();
   }
@@ -245,7 +245,7 @@ public Map<String, ActiveStreamingTool> activeStreamingTools() {
 
   /** Returns the queue for managing live requests, if available for this invocation. */
   public Optional<LiveRequestQueue> liveRequestQueue() {
-    return liveRequestQueue;
+    return Optional.ofNullable(liveRequestQueue);
   }
 
   /** Returns the unique ID for this invocation. */
@@ -258,15 +258,15 @@ public String invocationId() {
    * history.
    */
   public void branch(@Nullable String branch) {
-    this.branch = Optional.ofNullable(branch);
+    this.branch = branch;
   }
 
   /**
    * Returns the branch ID for the current invocation, if one is set. A branch represents a fork in
    * the conversation history.
    */
   public Optional<String> branch() {
-    return branch;
+    return Optional.ofNullable(branch);
   }
 
   /** Returns the agent being invoked. */
@@ -291,7 +291,7 @@ public Session session() {
 
   /** Returns the user content that triggered this invocation, if any. */
   public Optional<Content> userContent() {
-    return userContent;
+    return Optional.ofNullable(userContent);
   }
 
   /** Returns the configuration for the current agent run. */
@@ -416,13 +416,13 @@ private Builder(InvocationContext context) {
     private BaseArtifactService artifactService;
     private BaseMemoryService memoryService;
     private Plugin pluginManager = new PluginManager();
-    private Optional<LiveRequestQueue> liveRequestQueue = Optional.empty();
+    @Nullable private LiveRequestQueue liveRequestQueue = null;
     private Map<String, ActiveStreamingTool> activeStreamingTools = new ConcurrentHashMap<>();
-    private Optional<String> branch = Optional.empty();
+    @Nullable private String branch = null;
     private String invocationId = newInvocationContextId();
     private BaseAgent agent;
     private Session session;
-    private Optional<Content> userContent = Optional.empty();
+    @Nullable private Content userContent = null;
     private RunConfig runConfig = RunConfig.builder().build();
     private boolean endInvocation = false;
     @Nullable private EventsCompactionConfig eventsCompactionConfig;
@@ -489,7 +489,7 @@ public Builder pluginManager(Plugin pluginManager) {
     @Deprecated(forRemoval = true)
     @CanIgnoreReturnValue
     public Builder liveRequestQueue(Optional<LiveRequestQueue> liveRequestQueue) {
-      this.liveRequestQueue = liveRequestQueue;
+      this.liveRequestQueue = liveRequestQueue.orElse(null);
       return this;
     }
 
@@ -501,7 +501,7 @@ public Builder liveRequestQueue(Optional<LiveRequestQueue> liveRequestQueue) {
      */
     @CanIgnoreReturnValue
     public Builder liveRequestQueue(@Nullable LiveRequestQueue liveRequestQueue) {
-      this.liveRequestQueue = Optional.ofNullable(liveRequestQueue);
+      this.liveRequestQueue = liveRequestQueue;
       return this;
     }
 
@@ -516,7 +516,7 @@ public Builder liveRequestQueue(@Nullable LiveRequestQueue liveRequestQueue) {
     @Deprecated(forRemoval = true)
     @CanIgnoreReturnValue
     public Builder branch(Optional<String> branch) {
-      this.branch = branch;
+      this.branch = branch.orElse(null);
       return this;
     }
 
@@ -527,8 +527,8 @@ public Builder branch(Optional<String> branch) {
      * @return this builder instance for chaining.
      */
     @CanIgnoreReturnValue
-    public Builder branch(String branch) {
-      this.branch = Optional.of(branch);
+    public Builder branch(@Nullable String branch) {
+      this.branch = branch;
       return this;
     }
 
@@ -569,14 +569,12 @@ public Builder session(Session session) {
     }
 
     /**
-     * Sets the user content that triggered this invocation.
-     *
-     * @param userContent the user content that triggered this invocation.
-     * @return this builder instance for chaining.
+     * @deprecated Use {@link #userContent(Content)} instead.
      */
     @CanIgnoreReturnValue
+    @Deprecated
     public Builder userContent(Optional<Content> userContent) {
-      this.userContent = userContent;
+      this.userContent = userContent.orElse(null);
       return this;
     }
 
@@ -587,8 +585,8 @@ public Builder userContent(Optional<Content> userContent) {
      * @return this builder instance for chaining.
      */
     @CanIgnoreReturnValue
-    public Builder userContent(Content userContent) {
-      this.userContent = Optional.of(userContent);
+    public Builder userContent(@Nullable Content userContent) {
+      this.userContent = userContent;
       return this;
     }
 

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);
+  }
 }

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.
    *