@google/adk 2.4.0 → 2.5.0

package/dist/esm/runner/runner.js

@@ -9,7 +9,6 @@ var __name = (target, value) => __defProp(target, "name", { value, configurable:
  */
 var _a;
 import { createPartFromText, Modality } from "@google/genai";
-import { context, trace } from "@opentelemetry/api";
 import { isBaseAgent } from "../agents/base_agent.js";
 import { reservedFunctionCallName } from "../agents/framework_function_calls.js";
 import {
@@ -39,18 +38,25 @@ import {
   withoutInternalMetadata
 } from "../events/internal_metadata.js";
 import { PluginManager } from "../plugins/plugin_manager.js";
-import {
-  runAsyncGeneratorWithOtelContext,
-  tracer
-} from "../telemetry/tracing.js";
+import { recordInvocation } from "../telemetry/api/instrumentation.js";
 import { isBaseToolset } from "../tools/base_toolset.js";
+import {
+  formatError,
+  isInvocationCancellation,
+  toError
+} from "../utils/error_utils.js";
 import { logger } from "../utils/logger.js";
 import { isGemini2OrAbove } from "../utils/model_name.js";
 import { stringifyWithRedactedInlineData } from "../utils/redact_inline_data.js";
+import { isParallelWorker, isWorkflow } from "../workflow/index.js";
 import {
   asRunnableRoot,
   runNodeAsInvocation
 } from "../workflow/run_node_as_invocation.js";
+const TOOLSET_CLEANUP_TIMEOUT_MS = 1e4;
+const TOOLSET_CLEANUP_TIMEOUT_SYMBOL = /* @__PURE__ */ Symbol(
+  "google.adk.runner.toolsetCleanupTimeout"
+);
 const RUNNER_SIGNATURE_SYMBOL = /* @__PURE__ */ Symbol.for("google.adk.runner");
 function isRunner(obj) {
   return typeof obj === "object" && obj !== null && RUNNER_SIGNATURE_SYMBOL in obj && obj[RUNNER_SIGNATURE_SYMBOL] === true;
@@ -150,59 +156,57 @@ const _Runner = class _Runner {
       newMessage.role = "user";
     }
     rejectReservedFunctionCalls(newMessage);
-    const span = tracer.startSpan("invocation");
-    const ctx = trace.setSpan(context.active(), span);
+    const runnerCtx = this;
     try {
-      yield* runAsyncGeneratorWithOtelContext(
-        ctx,
-        this,
+      yield* recordInvocation(
+        { entrypoint: this.agent, sessionId, runConfig },
         async function* () {
-          var _a2, _b, _c, _d, _e, _f, _g, _h, _i, _j, _k, _l;
-          const session = await this.resolveSession(userId, sessionId);
+          var _a2, _b, _c, _d, _e;
+          const session = await runnerCtx.resolveSession(userId, sessionId);
           if ((_a2 = params.abortSignal) == null ? void 0 : _a2.aborted) {
             return;
           }
           if (!session) {
-            if (!this.appName) {
+            if (!runnerCtx.appName) {
               throw new Error(
                 `Session lookup failed: appName must be provided in runner constructor (or via app.name)`
               );
             }
             throw new Error(
-              `Session not found: ${sessionId} (appName=${this.appName}, userId=${userId})`
+              `Session not found: ${sessionId} (appName=${runnerCtx.appName}, userId=${userId})`
             );
           }
-          if (runConfig.supportCfc && isLlmAgent(this.agent)) {
-            const modelName = this.agent.canonicalModel.model;
+          if (runConfig.supportCfc && isLlmAgent(runnerCtx.agent)) {
+            const modelName = runnerCtx.agent.canonicalModel.model;
             if (!isGemini2OrAbove(modelName)) {
               throw new Error(
-                `CFC is not supported for model: ${modelName} in agent: ${this.agent.name}`
+                `CFC is not supported for model: ${modelName} in agent: ${runnerCtx.agent.name}`
               );
             }
-            if (!isBuiltInCodeExecutor(this.agent.codeExecutor)) {
-              this.agent.codeExecutor = new BuiltInCodeExecutor();
+            if (!isBuiltInCodeExecutor(runnerCtx.agent.codeExecutor)) {
+              runnerCtx.agent.codeExecutor = new BuiltInCodeExecutor();
             }
           }
           const invocationContext = new InvocationContext({
-            artifactService: this.artifactService ? isSessionArtifactService(this.artifactService) ? this.artifactService : new ScopedArtifactService(
-              this.artifactService,
-              this.appName,
+            artifactService: runnerCtx.artifactService ? isSessionArtifactService(runnerCtx.artifactService) ? runnerCtx.artifactService : new ScopedArtifactService(
+              runnerCtx.artifactService,
+              runnerCtx.appName,
               userId,
               sessionId
             ) : void 0,
-            sessionService: this.sessionService,
-            memoryService: this.memoryService,
-            credentialService: this.credentialService,
+            sessionService: runnerCtx.sessionService,
+            memoryService: runnerCtx.memoryService,
+            credentialService: runnerCtx.credentialService,
             invocationId: newInvocationContextId(),
-            agent: isBaseAgent(this.agent) ? this.agent : void 0,
+            agent: isBaseAgent(runnerCtx.agent) ? runnerCtx.agent : void 0,
             session,
             userContent: newMessage,
             runConfig,
             a2aMetadata: runConfig.a2aMetadata,
-            pluginManager: this.pluginManager,
+            pluginManager: runnerCtx.pluginManager,
             abortSignal: params.abortSignal
           });
-          const pluginUserMessage = await this.pluginManager.runOnUserMessageCallback({
+          const pluginUserMessage = await runnerCtx.pluginManager.runOnUserMessageCallback({
             userMessage: newMessage,
             invocationContext
           });
@@ -217,7 +221,7 @@ const _Runner = class _Runner {
               throw new Error("No parts in the newMessage.");
             }
             if (runConfig.saveInputBlobsAsArtifacts) {
-              newMessage = await this.saveArtifacts(
+              newMessage = await runnerCtx.saveArtifacts(
                 invocationContext.artifactService,
                 invocationContext.invocationId,
                 newMessage
@@ -226,7 +230,7 @@ const _Runner = class _Runner {
                 return;
               }
             }
-            await this.sessionService.appendEvent({
+            await runnerCtx.sessionService.appendEvent({
               session,
               event: createEvent({
                 invocationId: invocationContext.invocationId,
@@ -241,83 +245,137 @@ const _Runner = class _Runner {
               return;
             }
           }
-          if (isBaseAgent(this.agent)) {
-            invocationContext.agent = this.determineAgentForResumption(
+          if (isBaseAgent(runnerCtx.agent)) {
+            invocationContext.agent = runnerCtx.determineAgentForResumption(
               session,
-              this.agent
+              runnerCtx.agent
             );
           }
           if (newMessage) {
-            const beforeRunCallbackResponse = await this.pluginManager.runBeforeRunCallback({
-              invocationContext
-            });
-            if ((_f = params.abortSignal) == null ? void 0 : _f.aborted) {
-              return;
-            }
-            if (beforeRunCallbackResponse) {
-              const earlyExitEvent = createEvent({
-                invocationId: invocationContext.invocationId,
-                author: "model",
-                content: beforeRunCallbackResponse
-              });
-              await this.sessionService.appendEvent({
-                session,
-                event: earlyExitEvent
-              });
-              if ((_g = params.abortSignal) == null ? void 0 : _g.aborted) {
-                return;
-              }
-              yield earlyExitEvent;
-            } else {
-              let abortSealed = false;
-              const sealAbortedInvocation = /* @__PURE__ */ __name(() => {
-                abortSealed = true;
-                return this.synthesizeAbortEvents(invocationContext);
-              }, "sealAbortedInvocation");
-              try {
-                for await (const event of this.runRoot(invocationContext)) {
-                  if ((_h = params.abortSignal) == null ? void 0 : _h.aborted) {
-                    break;
-                  }
-                  const outputEvent = await this.processAndAppendEvent(
-                    invocationContext,
-                    event
-                  );
-                  if ((_i = params.abortSignal) == null ? void 0 : _i.aborted) {
-                    break;
-                  }
-                  yield outputEvent;
-                }
-                if ((_j = params.abortSignal) == null ? void 0 : _j.aborted) {
-                  for (const abortEvent of await sealAbortedInvocation()) {
-                    yield abortEvent;
-                  }
-                  return;
-                }
-              } finally {
-                if (((_k = params.abortSignal) == null ? void 0 : _k.aborted) && !abortSealed) {
-                  try {
-                    await sealAbortedInvocation();
-                  } catch (e) {
-                    logger.error(
-                      `Failed to seal aborted invocation ${invocationContext.invocationId}.`,
-                      e
-                    );
-                  }
-                }
-              }
-              await this.pluginManager.runAfterRunCallback({ invocationContext });
-              if ((_l = params.abortSignal) == null ? void 0 : _l.aborted) {
-                return;
-              }
-            }
+            yield* runnerCtx.runWithPluginLifecycle(
+              invocationContext,
+              () => runnerCtx.runAsyncWithPlugins(invocationContext)
+            );
           }
         }
       );
     } finally {
-      span.end();
-      const toolsets = isBaseAgent(this.agent) ? getAllToolsets(this.agent) : [];
-      await Promise.allSettled(toolsets.map((t) => t.close()));
+      await this.close();
+    }
+  }
+  /**
+   * Closes the runner and cleans up all toolsets reachable from the root agent
+   * or workflow.
+   */
+  async close() {
+    await cleanupToolsets(collectToolsets(this.agent));
+  }
+  async *runAsyncWithPlugins(invocationContext) {
+    var _a2, _b, _c, _d, _e, _f;
+    const beforeRunCallbackResponse = await this.pluginManager.runBeforeRunCallback({
+      invocationContext
+    });
+    if ((_a2 = invocationContext.abortSignal) == null ? void 0 : _a2.aborted) {
+      return;
+    }
+    if (beforeRunCallbackResponse) {
+      const earlyExitEvent = createEvent({
+        invocationId: invocationContext.invocationId,
+        author: "model",
+        content: beforeRunCallbackResponse
+      });
+      await this.sessionService.appendEvent({
+        session: invocationContext.session,
+        event: earlyExitEvent
+      });
+      if ((_b = invocationContext.abortSignal) == null ? void 0 : _b.aborted) {
+        return;
+      }
+      yield earlyExitEvent;
+    } else {
+      let abortSealed = false;
+      const sealAbortedInvocation = /* @__PURE__ */ __name(() => {
+        abortSealed = true;
+        return this.synthesizeAbortEvents(invocationContext);
+      }, "sealAbortedInvocation");
+      try {
+        for await (const event of this.runRoot(invocationContext)) {
+          if ((_c = invocationContext.abortSignal) == null ? void 0 : _c.aborted) {
+            break;
+          }
+          const outputEvent = await this.processAndAppendEvent(
+            invocationContext,
+            event
+          );
+          if ((_d = invocationContext.abortSignal) == null ? void 0 : _d.aborted) {
+            break;
+          }
+          yield outputEvent;
+        }
+        if ((_e = invocationContext.abortSignal) == null ? void 0 : _e.aborted) {
+          for (const abortEvent of await sealAbortedInvocation()) {
+            yield abortEvent;
+          }
+          return;
+        }
+      } finally {
+        if (((_f = invocationContext.abortSignal) == null ? void 0 : _f.aborted) && !abortSealed) {
+          try {
+            await sealAbortedInvocation();
+          } catch (e) {
+            logger.error(
+              `Failed to seal aborted invocation ${invocationContext.invocationId}.`,
+              e
+            );
+          }
+        }
+      }
+    }
+  }
+  /** Finalizes a started run, including when its consumer closes the iterator. */
+  async *runWithPluginLifecycle(invocationContext, execute) {
+    let failed = false;
+    let completed = false;
+    try {
+      yield* execute();
+      completed = true;
+    } catch (error) {
+      if (!isInvocationCancellation(error, invocationContext.abortSignal)) {
+        failed = true;
+        await this.notifyRunError(invocationContext, error);
+      }
+      throw error;
+    } finally {
+      if (!failed) {
+        await this.finishRun(invocationContext, !completed);
+      }
+    }
+  }
+  async finishRun(invocationContext, closingEarly) {
+    try {
+      await this.pluginManager.runAfterRunCallback({ invocationContext });
+    } catch (error) {
+      await this.notifyRunError(invocationContext, error);
+      if (!closingEarly) {
+        throw error;
+      }
+      logger.error(
+        `afterRunCallback failed while closing invocation ${invocationContext.invocationId}.`,
+        error
+      );
+    }
+  }
+  async notifyRunError(invocationContext, error) {
+    try {
+      await this.pluginManager.runOnRunErrorCallback({
+        invocationContext,
+        error: toError(error)
+      });
+    } catch (callbackError) {
+      logger.error(
+        "Run error notification failed; preserving the original error.",
+        callbackError
+      );
     }
   }
   /**
@@ -508,6 +566,7 @@ const _Runner = class _Runner {
       throw new Error("runLive is only supported for agents.");
     }
     const agent = this.agent;
+    const { sessionId } = params;
     const runConfig = createRunConfig(params.runConfig);
     if (!((_a2 = runConfig.responseModalities) == null ? void 0 : _a2.length)) {
       runConfig.responseModalities = [Modality.AUDIO];
@@ -518,92 +577,92 @@ const _Runner = class _Runner {
       }
       (_d = runConfig.inputAudioTranscription) != null ? _d : runConfig.inputAudioTranscription = {};
     }
-    const span = tracer.startSpan("invocation");
-    const ctx = trace.setSpan(context.active(), span);
-    try {
-      yield* runAsyncGeneratorWithOtelContext(
-        ctx,
-        this,
-        async function* () {
-          var _a3, _b2, _c2, _d2;
-          const session = await this.sessionService.getOrCreateSession({
-            appName: this.appName,
-            userId: params.userId,
-            sessionId: params.sessionId
-          });
-          if ((_a3 = params.abortSignal) == null ? void 0 : _a3.aborted) {
-            return;
-          }
-          const invocationContext = new InvocationContext({
-            artifactService: this.artifactService ? isSessionArtifactService(this.artifactService) ? this.artifactService : new ScopedArtifactService(
-              this.artifactService,
-              this.appName,
-              params.userId,
-              params.sessionId
-            ) : void 0,
-            sessionService: this.sessionService,
-            memoryService: this.memoryService,
-            credentialService: this.credentialService,
-            invocationId: newInvocationContextId(),
-            agent,
-            session,
-            runConfig,
-            a2aMetadata: runConfig.a2aMetadata,
-            pluginManager: this.pluginManager,
-            liveRequestQueue: params.liveRequestQueue,
-            abortSignal: params.abortSignal,
-            liveSessionResumptionHandle: params.liveSessionResumptionHandle
-          });
-          invocationContext.agent = this.determineAgentForResumption(
-            session,
-            agent
-          );
-          const beforeRunCallbackResponse = await this.pluginManager.runBeforeRunCallback({
-            invocationContext
-          });
-          if ((_b2 = params.abortSignal) == null ? void 0 : _b2.aborted) {
-            return;
-          }
-          if (beforeRunCallbackResponse) {
-            const earlyExitEvent = createEvent({
-              invocationId: invocationContext.invocationId,
-              author: "model",
-              content: beforeRunCallbackResponse
-            });
-            await this.sessionService.appendEvent({
-              session,
-              event: earlyExitEvent
-            });
-            yield earlyExitEvent;
-            return;
-          }
-          for await (const event of requireAgent(invocationContext).runLive(
-            invocationContext
-          )) {
-            if ((_c2 = params.abortSignal) == null ? void 0 : _c2.aborted) {
-              return;
-            }
-            const modifiedEvent = await this.pluginManager.runOnEventCallback({
-              invocationContext,
-              event
-            });
-            if ((_d2 = params.abortSignal) == null ? void 0 : _d2.aborted) {
-              return;
-            }
-            const eventToProcess = modifiedEvent ? withInternalMetadataOf(event, modifiedEvent) : event;
-            if (!eventToProcess.partial && !isLiveModelMediaEventWithInlineData(eventToProcess)) {
-              await this.sessionService.appendEvent({
-                session,
-                event: eventToProcess
-              });
-            }
-            yield eventToProcess;
-          }
-          await this.pluginManager.runAfterRunCallback({ invocationContext });
+    const runnerCtx = this;
+    yield* recordInvocation(
+      { entrypoint: agent, sessionId, runConfig },
+      async function* () {
+        var _a3;
+        const session = await runnerCtx.sessionService.getOrCreateSession({
+          appName: runnerCtx.appName,
+          userId: params.userId,
+          sessionId: params.sessionId
+        });
+        if ((_a3 = params.abortSignal) == null ? void 0 : _a3.aborted) {
+          return;
         }
-      );
-    } finally {
-      span.end();
+        const invocationContext = new InvocationContext({
+          artifactService: runnerCtx.artifactService ? isSessionArtifactService(runnerCtx.artifactService) ? runnerCtx.artifactService : new ScopedArtifactService(
+            runnerCtx.artifactService,
+            runnerCtx.appName,
+            params.userId,
+            params.sessionId
+          ) : void 0,
+          sessionService: runnerCtx.sessionService,
+          memoryService: runnerCtx.memoryService,
+          credentialService: runnerCtx.credentialService,
+          invocationId: newInvocationContextId(),
+          agent,
+          session,
+          runConfig,
+          a2aMetadata: runConfig.a2aMetadata,
+          pluginManager: runnerCtx.pluginManager,
+          liveRequestQueue: params.liveRequestQueue,
+          abortSignal: params.abortSignal,
+          liveSessionResumptionHandle: params.liveSessionResumptionHandle
+        });
+        invocationContext.agent = runnerCtx.determineAgentForResumption(
+          session,
+          agent
+        );
+        yield* runnerCtx.runWithPluginLifecycle(
+          invocationContext,
+          () => runnerCtx.runLiveWithPlugins(invocationContext)
+        );
+      }
+    );
+  }
+  async *runLiveWithPlugins(invocationContext) {
+    var _a2, _b, _c;
+    const beforeRunCallbackResponse = await this.pluginManager.runBeforeRunCallback({
+      invocationContext
+    });
+    if ((_a2 = invocationContext.abortSignal) == null ? void 0 : _a2.aborted) {
+      return;
+    }
+    if (beforeRunCallbackResponse) {
+      const earlyExitEvent = createEvent({
+        invocationId: invocationContext.invocationId,
+        author: "model",
+        content: beforeRunCallbackResponse
+      });
+      await this.sessionService.appendEvent({
+        session: invocationContext.session,
+        event: earlyExitEvent
+      });
+      yield earlyExitEvent;
+      return;
+    }
+    for await (const event of requireAgent(invocationContext).runLive(
+      invocationContext
+    )) {
+      if ((_b = invocationContext.abortSignal) == null ? void 0 : _b.aborted) {
+        return;
+      }
+      const modifiedEvent = await this.pluginManager.runOnEventCallback({
+        invocationContext,
+        event
+      });
+      if ((_c = invocationContext.abortSignal) == null ? void 0 : _c.aborted) {
+        return;
+      }
+      const eventToProcess = modifiedEvent ? withInternalMetadataOf(event, modifiedEvent) : event;
+      if (!eventToProcess.partial && !isLiveModelMediaEventWithInlineData(eventToProcess)) {
+        await this.sessionService.appendEvent({
+          session: invocationContext.session,
+          event: eventToProcess
+        });
+      }
+      yield eventToProcess;
     }
   }
 };
@@ -706,28 +765,87 @@ function findEventByLastFunctionResponseId(events) {
   return (_a2 = findMatchingFunctionCall(events)) != null ? _a2 : null;
 }
 __name(findEventByLastFunctionResponseId, "findEventByLastFunctionResponseId");
-function getAllToolsets(agent) {
-  const toolsets = [];
+function collectToolsets(root) {
+  const toolsets = /* @__PURE__ */ new Set();
   const visited = /* @__PURE__ */ new Set();
   function traverse(curr) {
-    if (visited.has(curr)) return;
+    if (visited.has(curr)) {
+      return;
+    }
     visited.add(curr);
     if (isLlmAgent(curr)) {
       for (const tool of curr.tools) {
         if (isBaseToolset(tool)) {
-          toolsets.push(tool);
+          toolsets.add(tool);
         }
       }
     }
-    for (const sub of curr.subAgents) {
-      traverse(sub);
+    if (isBaseAgent(curr)) {
+      for (const sub of curr.subAgents) {
+        traverse(sub);
+      }
+    }
+    if (isWorkflow(curr) && curr.graph) {
+      for (const node of curr.graph.nodes) {
+        traverse(node);
+      }
+    }
+    if (isParallelWorker(curr)) {
+      traverse(curr.inner);
     }
   }
   __name(traverse, "traverse");
-  traverse(agent);
+  traverse(root);
   return toolsets;
 }
-__name(getAllToolsets, "getAllToolsets");
+__name(collectToolsets, "collectToolsets");
+function formatToolsetCloseError(err) {
+  const formatted = formatError(err);
+  if (formatted !== "[object Object]") {
+    return formatted;
+  }
+  try {
+    return JSON.stringify(err);
+  } catch {
+    return formatted;
+  }
+}
+__name(formatToolsetCloseError, "formatToolsetCloseError");
+async function cleanupToolsets(toolsets) {
+  await Promise.allSettled(
+    [...toolsets].map(async (toolset) => {
+      var _a2;
+      const toolsetName = ((_a2 = toolset.constructor) == null ? void 0 : _a2.name) || "BaseToolset";
+      let timeoutId;
+      try {
+        logger.debug(`Closing toolset: ${toolsetName}`);
+        const result = await Promise.race([
+          toolset.close().then(() => void 0),
+          new Promise((resolve) => {
+            timeoutId = setTimeout(
+              () => resolve(TOOLSET_CLEANUP_TIMEOUT_SYMBOL),
+              TOOLSET_CLEANUP_TIMEOUT_MS
+            );
+          })
+        ]);
+        if (result === TOOLSET_CLEANUP_TIMEOUT_SYMBOL) {
+          logger.warn(`Toolset ${toolsetName} cleanup timed out`);
+        } else {
+          logger.debug(`Successfully closed toolset: ${toolsetName}`);
+        }
+      } catch (e) {
+        logger.error(
+          `Error closing toolset ${toolsetName}: ${formatToolsetCloseError(e)}`
+        );
+      } finally {
+        if (timeoutId !== void 0) {
+          clearTimeout(timeoutId);
+        }
+      }
+    })
+  );
+}
+__name(cleanupToolsets, "cleanupToolsets");
 function rejectReservedFunctionCalls(newMessage) {
   const reserved = reservedFunctionCallName(newMessage);
   if (reserved) {

package/dist/esm/sessions/base_session_service.js

@@ -41,8 +41,8 @@ const _BaseSessionService = class _BaseSessionService {
     if (event.partial) {
       return event;
     }
-    event = trimTempDeltaState(event);
     this.updateSessionState({ session, event });
+    event = trimTempDeltaState(event);
     upsertSessionEvent(session, event);
     return event;
   }
@@ -56,9 +56,6 @@ const _BaseSessionService = class _BaseSessionService {
       return;
     }
     for (const [key, value] of Object.entries(event.actions.stateDelta)) {
-      if (key.startsWith(State.TEMP_PREFIX)) {
-        continue;
-      }
       if (!shouldApplyDeltaWrite(session.state, event.actions.stateDelta, key)) {
         continue;
       }