Merge remote-tracking branch 'origin/main' into feat-mcp_client_timeout

Author: wenhaozhao

core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java

@@ -39,7 +39,6 @@
 import com.google.adk.tools.ToolContext;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
-import com.google.genai.types.FunctionCall;
 import com.google.genai.types.FunctionResponse;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.StatusCode;
@@ -621,11 +620,10 @@ private Event buildModelResponseEvent(
 
     Event event = eventBuilder.build();
 
-    ImmutableList<FunctionCall> functionCalls = event.functionCalls();
-    if (!functionCalls.isEmpty()) {
+    if (!event.functionCalls().isEmpty()) {
       Functions.populateClientFunctionCallId(event);
       Set<String> longRunningToolIds =
-          Functions.getLongRunningFunctionCalls(functionCalls, llmRequest.tools());
+          Functions.getLongRunningFunctionCalls(event.functionCalls(), llmRequest.tools());
       if (!longRunningToolIds.isEmpty()) {
         event.setLongRunningToolIds(Optional.of(longRunningToolIds));
       }

core/src/main/java/com/google/adk/runner/Runner.java

@@ -38,7 +38,6 @@
 import io.reactivex.rxjava3.core.Flowable;
 import io.reactivex.rxjava3.core.Maybe;
 import io.reactivex.rxjava3.core.Single;
-import io.reactivex.rxjava3.core.Completable;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -84,7 +83,7 @@ public BaseSessionService sessionService() {
    *
    * @throws IllegalArgumentException if message has no parts.
    */
-  private Completable appendNewMessageToSession(
+  private void appendNewMessageToSession(
       Session session,
       Content newMessage,
       InvocationContext invocationContext,
@@ -124,7 +123,7 @@ private Completable appendNewMessageToSession(
             .author("user")
             .content(Optional.of(newMessage))
             .build();
-    return this.sessionService.appendEvent(session, event).ignoreElement();
+    this.sessionService.appendEvent(session, event);
   }
 
   /**
@@ -191,20 +190,14 @@ public Flowable<Event> runAsync(Session session, Content newMessage, RunConfig r
                         newMessage,
                         runConfig);
 
-                Completable newMessageCompletable = Completable.complete();
                 if (newMessage != null) {
-                  newMessageCompletable =
-                      appendNewMessageToSession(
-                          sess,
-                          newMessage,
-                          invocationContext,
-                          runConfig.saveInputBlobsAsArtifacts());
+                  appendNewMessageToSession(
+                      sess, newMessage, invocationContext, runConfig.saveInputBlobsAsArtifacts());
                 }
 
                 invocationContext.agent(this.findAgentToRun(sess, rootAgent));
                 Flowable<Event> events = invocationContext.agent().runAsync(invocationContext);
-                return newMessageCompletable.andThen(
-                    events.concatMapSingle(event -> this.sessionService.appendEvent(sess, event)));
+                return events.doOnNext(event -> this.sessionService.appendEvent(sess, event));
               })
           .doOnError(
               throwable -> {

core/src/main/java/com/google/adk/tools/LongRunningFunctionTool.java

@@ -35,7 +35,22 @@ public static LongRunningFunctionTool create(Class<?> cls, String methodName) {
         String.format("Method %s not found in class %s.", methodName, cls.getName()));
   }
 
+  public static LongRunningFunctionTool create(Object instance, String methodName) {
+    Class<?> cls = instance.getClass();
+    for (Method method : cls.getMethods()) {
+      if (method.getName().equals(methodName)) {
+        return new LongRunningFunctionTool(instance, method);
+      }
+    }
+    throw new IllegalArgumentException(
+        String.format("Method %s not found in class %s.", methodName, cls.getName()));
+  }
+
   private LongRunningFunctionTool(Method func) {
     super(null, func, /* isLongRunning= */ true);
   }
+
+  private LongRunningFunctionTool(Object instance, Method func) {
+    super(instance, func, /* isLongRunning= */ true);
+  }
 }

core/src/main/java/com/google/adk/tools/applicationintegrationtoolset/ApplicationIntegrationToolset.java

@@ -5,24 +5,28 @@
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.google.adk.agents.ReadonlyContext;
 import com.google.adk.tools.BaseTool;
+import com.google.adk.tools.BaseToolset;
 import com.google.adk.tools.applicationintegrationtoolset.IntegrationConnectorTool.DefaultHttpExecutor;
 import com.google.adk.tools.applicationintegrationtoolset.IntegrationConnectorTool.HttpExecutor;
+import io.reactivex.rxjava3.core.Flowable;
 import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import org.jspecify.annotations.Nullable;
 
 /** Application Integration Toolset */
-public class ApplicationIntegrationToolset {
+public class ApplicationIntegrationToolset implements BaseToolset {
   String project;
   String location;
   @Nullable String integration;
   @Nullable List<String> triggers;
   @Nullable String connection;
   @Nullable Map<String, List<String>> entityOperations;
   @Nullable List<String> actions;
+  String serviceAccountJson;
   @Nullable String toolNamePrefix;
   @Nullable String toolInstructions;
   public static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
@@ -35,16 +39,16 @@ public class ApplicationIntegrationToolset {
    *
    * <p>integrationTool = new ApplicationIntegrationToolset( project="test-project",
    * location="us-central1", integration="test-integration",
-   * triggers=ImmutableList.of("api_trigger/test_trigger",
-   * "api_trigger/test_trigger_2"),connection=null,enitityOperations=null,actions=null,toolNamePrefix="test-integration-tool",toolInstructions="This
+   * triggers=ImmutableList.of("api_trigger/test_trigger", "api_trigger/test_trigger_2",
+   * serviceAccountJson="{....}"),connection=null,enitityOperations=null,actions=null,toolNamePrefix="test-integration-tool",toolInstructions="This
    * tool is used to get response from test-integration.");
    *
    * <p>connectionTool = new ApplicationIntegrationToolset( project="test-project",
    * location="us-central1", integration=null, triggers=null, connection="test-connection",
    * entityOperations=ImmutableMap.of("Entity1", ImmutableList.of("LIST", "GET", "UPDATE")),
    * "Entity2", ImmutableList.of()), actions=ImmutableList.of("ExecuteCustomQuery"),
-   * toolNamePrefix="test-tool", toolInstructions="This tool is used to list, get and update issues
-   * in Jira.");
+   * serviceAccountJson="{....}", toolNamePrefix="test-tool", toolInstructions="This tool is used to
+   * list, get and update issues in Jira.");
    *
    * @param project The GCP project ID.
    * @param location The GCP location of integration.
@@ -53,6 +57,9 @@ public class ApplicationIntegrationToolset {
    * @param connection(Optional) The connection name.
    * @param entityOperations(Optional) The entity operations.
    * @param actions(Optional) The actions.
+   * @param serviceAccountJson(Optional) The service account configuration as a dictionary. Required
+   *     if not using default service credential. Used for fetching the Application Integration or
+   *     Integration Connector resource.
    * @param toolNamePrefix(Optional) The tool name prefix.
    * @param toolInstructions(Optional) The tool instructions.
    */
@@ -64,6 +71,7 @@ public ApplicationIntegrationToolset(
       String connection,
       Map<String, List<String>> entityOperations,
       List<String> actions,
+      String serviceAccountJson,
       String toolNamePrefix,
       String toolInstructions) {
     this(
@@ -74,9 +82,10 @@ public ApplicationIntegrationToolset(
         connection,
         entityOperations,
         actions,
+        serviceAccountJson,
         toolNamePrefix,
         toolInstructions,
-        new DefaultHttpExecutor());
+        new DefaultHttpExecutor().createExecutor(serviceAccountJson));
   }
 
   ApplicationIntegrationToolset(
@@ -87,6 +96,7 @@ public ApplicationIntegrationToolset(
       String connection,
       Map<String, List<String>> entityOperations,
       List<String> actions,
+      String serviceAccountJson,
       String toolNamePrefix,
       String toolInstructions,
       HttpExecutor httpExecutor) {
@@ -97,6 +107,7 @@ public ApplicationIntegrationToolset(
     this.connection = connection;
     this.entityOperations = entityOperations;
     this.actions = actions;
+    this.serviceAccountJson = serviceAccountJson;
     this.toolNamePrefix = toolNamePrefix;
     this.toolInstructions = toolInstructions;
     this.httpExecutor = httpExecutor;
@@ -121,7 +132,7 @@ List<String> getPathUrl(String openApiSchemaString) throws Exception {
     return pathUrls;
   }
 
-  public List<BaseTool> getTools() throws Exception {
+  private List<BaseTool> getAllTools() throws Exception {
     String openApiSchemaString = null;
     List<BaseTool> tools = new ArrayList<>();
     if (!isNullOrEmpty(this.integration)) {
@@ -142,7 +153,15 @@ public List<BaseTool> getTools() throws Exception {
         if (toolName != null) {
           tools.add(
               new IntegrationConnectorTool(
-                  openApiSchemaString, pathUrl, toolName, "", null, null, null, this.httpExecutor));
+                  openApiSchemaString,
+                  pathUrl,
+                  toolName,
+                  toolInstructions,
+                  null,
+                  null,
+                  null,
+                  this.serviceAccountJson,
+                  this.httpExecutor));
         }
       }
     } else if (!isNullOrEmpty(this.connection)
@@ -182,6 +201,7 @@ public List<BaseTool> getTools() throws Exception {
                   connectionDetails.name,
                   connectionDetails.serviceName,
                   connectionDetails.host,
+                  this.serviceAccountJson,
                   this.httpExecutor));
         }
       }
@@ -193,4 +213,19 @@ public List<BaseTool> getTools() throws Exception {
 
     return tools;
   }
+
+  @Override
+  public Flowable<BaseTool> getTools(@Nullable ReadonlyContext readonlyContext) {
+    try {
+      List<BaseTool> allTools = getAllTools();
+      return Flowable.fromIterable(allTools);
+    } catch (Exception e) {
+      return Flowable.error(e);
+    }
+  }
+
+  @Override
+  public void close() throws Exception {
+    // Nothing to close.
+  }
 }