@google/adk 2.3.0 → 2.5.0

package/dist/esm/agents/agent_engine_client.js

@@ -0,0 +1,382 @@
+import {createRequire as topLevelCreateRequire} from 'module';
+const require = topLevelCreateRequire(import.meta.url);
+var __create = Object.create;
+var __defProp = Object.defineProperty;
+var __getOwnPropDesc = Object.getOwnPropertyDescriptor;
+var __getOwnPropNames = Object.getOwnPropertyNames;
+var __getProtoOf = Object.getPrototypeOf;
+var __hasOwnProp = Object.prototype.hasOwnProperty;
+var __name = (target, value) => __defProp(target, "name", { value, configurable: true });
+var __copyProps = (to, from, except, desc) => {
+  if (from && typeof from === "object" || typeof from === "function") {
+    for (let key of __getOwnPropNames(from))
+      if (!__hasOwnProp.call(to, key) && key !== except)
+        __defProp(to, key, { get: () => from[key], enumerable: !(desc = __getOwnPropDesc(from, key)) || desc.enumerable });
+  }
+  return to;
+};
+var __toESM = (mod, isNodeMode, target) => (target = mod != null ? __create(__getProtoOf(mod)) : {}, __copyProps(
+  // If the importer is in node compatibility mode or this is not an ESM
+  // file that has been converted to a CommonJS file using a Babel-
+  // compatible transform (i.e. "__esModule" has not been set), then set
+  // "default" to the CommonJS "module.exports" for node compatibility.
+  isNodeMode || !mod || !mod.__esModule ? __defProp(target, "default", { value: mod, enumerable: true }) : target,
+  mod
+));
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import { getErrorMessage } from "../utils/error_utils.js";
+import { loadOptionalPeer } from "../utils/optional_peer.js";
+const _AgentExecutionError = class _AgentExecutionError extends Error {
+  constructor(message, cause) {
+    super(message);
+    this.cause = cause;
+    this.name = "AgentExecutionError";
+  }
+};
+__name(_AgentExecutionError, "AgentExecutionError");
+let AgentExecutionError = _AgentExecutionError;
+const WIRE_KEY_ALIASES = [
+  ["userId", "user_id"],
+  ["sessionId", "session_id"],
+  ["expireTime", "expire_time"],
+  ["runConfig", "run_config"],
+  ["stateDelta", "state_delta"]
+];
+function toWireInput(config) {
+  var _a;
+  const input = { ...config };
+  delete input["classMethod"];
+  for (const [camel, snake] of WIRE_KEY_ALIASES) {
+    input[snake] = (_a = input[snake]) != null ? _a : input[camel];
+    delete input[camel];
+  }
+  for (const key of Object.keys(input)) {
+    if (input[key] === void 0) {
+      delete input[key];
+    }
+  }
+  return input;
+}
+__name(toWireInput, "toWireInput");
+function buildInputStruct(sdkHelpers, input) {
+  var _a;
+  const value = sdkHelpers.toValue(input);
+  return (_a = value == null ? void 0 : value.structValue) != null ? _a : void 0;
+}
+__name(buildInputStruct, "buildInputStruct");
+function parseOutput(sdkHelpers, output) {
+  return sdkHelpers.fromValue(output);
+}
+__name(parseOutput, "parseOutput");
+const IGNORED_SSE_FIELDS = /* @__PURE__ */ new Set([
+  "event",
+  "id",
+  "retry"
+]);
+const DONE_SENTINEL = "[DONE]";
+function parseFragment(payload) {
+  try {
+    return JSON.parse(payload);
+  } catch (e) {
+    throw new AgentExecutionError(
+      `Failed to parse stream fragment: ${payload}`,
+      e
+    );
+  }
+}
+__name(parseFragment, "parseFragment");
+function* dispatchEvent(dataLines) {
+  if (dataLines.length === 0) {
+    return;
+  }
+  const lines = dataLines.splice(0);
+  const payload = lines.join("\n").trim();
+  if (payload === "" || payload === DONE_SENTINEL) {
+    return;
+  }
+  let parsed;
+  try {
+    parsed = parseFragment(payload);
+  } catch (e) {
+    if (lines.length === 1) {
+      throw e;
+    }
+    const values = [];
+    try {
+      for (const line of lines) {
+        const trimmed = line.trim();
+        if (trimmed !== "" && trimmed !== DONE_SENTINEL) {
+          values.push(JSON.parse(trimmed));
+        }
+      }
+    } catch {
+      throw e;
+    }
+    yield* values;
+    return;
+  }
+  yield parsed;
+}
+__name(dispatchEvent, "dispatchEvent");
+function* processLine(rawLine, dataLines) {
+  const line = (rawLine.endsWith("\r") ? rawLine.slice(0, -1) : rawLine).replace(/^[ \t]+/, "");
+  if (line === "") {
+    yield* dispatchEvent(dataLines);
+    return;
+  }
+  if (line.startsWith(":")) {
+    return;
+  }
+  const colon = line.indexOf(":");
+  const field = colon === -1 ? line : line.slice(0, colon);
+  if (IGNORED_SSE_FIELDS.has(field)) {
+    return;
+  }
+  if (field === "data") {
+    if (dataLines.length === 1) {
+      const buffered = dataLines[0].trim();
+      if (buffered === DONE_SENTINEL) {
+        yield* dispatchEvent(dataLines);
+      } else if (buffered !== "") {
+        try {
+          JSON.parse(buffered);
+          yield* dispatchEvent(dataLines);
+        } catch {
+        }
+      }
+    }
+    const value = colon === -1 ? "" : line.slice(colon + 1);
+    dataLines.push(value.startsWith(" ") ? value.slice(1) : value);
+    return;
+  }
+  yield* dispatchEvent(dataLines);
+  const trimmed = line.trim();
+  if (trimmed === "" || trimmed === DONE_SENTINEL) {
+    return;
+  }
+  yield parseFragment(trimmed);
+}
+__name(processLine, "processLine");
+async function* parseStream(stream) {
+  var _a;
+  const decoder = new TextDecoder("utf-8");
+  const dataLines = [];
+  let buffer = "";
+  for await (const chunk of stream) {
+    const chunkData = chunk == null ? void 0 : chunk.data;
+    if (chunkData === void 0) {
+      if (typeof chunk === "object" && chunk !== null) {
+        yield chunk;
+      }
+      continue;
+    }
+    buffer += chunkData instanceof Uint8Array ? decoder.decode(chunkData, { stream: true }) : String(chunkData);
+    const lines = buffer.split("\n");
+    buffer = (_a = lines.pop()) != null ? _a : "";
+    for (const line of lines) {
+      yield* processLine(line, dataLines);
+    }
+  }
+  buffer += decoder.decode();
+  for (const line of buffer.split("\n")) {
+    yield* processLine(line, dataLines);
+  }
+  yield* dispatchEvent(dataLines);
+}
+__name(parseStream, "parseStream");
+const _AgentEngineClient = class _AgentEngineClient {
+  constructor(config) {
+    var _a;
+    this.apiEndpoint = (_a = config.apiEndpoint) != null ? _a : config.location === "global" ? "aiplatform.googleapis.com" : `${config.location}-aiplatform.googleapis.com`;
+    this.reasoningEnginePath = `projects/${config.project}/locations/${config.location}/reasoningEngines/${config.reasoningEngineId}`;
+  }
+  /**
+   * Creates a remote session.
+   *
+   * Calls the `async_create_session` class method unless `config.classMethod`
+   * names another one.
+   *
+   * @returns The session payload returned by the Agent Engine. When
+   *     `config.sessionId` is omitted the engine generates one, and this is the
+   *     only place the caller can read it back from.
+   */
+  async createSession(config) {
+    return this.callQueryMethod({
+      config,
+      defaultClassMethod: "async_create_session",
+      errorPrefix: "Failed to create session"
+    });
+  }
+  /**
+   * Reads a remote session.
+   *
+   * Calls the `async_get_session` class method unless `config.classMethod`
+   * names another one.
+   *
+   * @throws {@link AgentExecutionError} if the user id or the session id is
+   *     missing, or if the call fails.
+   */
+  async getSession(config) {
+    var _a, _b;
+    const userId = (_a = config.user_id) != null ? _a : config.userId;
+    const sessionId = (_b = config.session_id) != null ? _b : config.sessionId;
+    if (!userId || !sessionId) {
+      throw new AgentExecutionError(
+        "Failed to get session: user_id and session_id are required"
+      );
+    }
+    return this.callQueryMethod({
+      config,
+      defaultClassMethod: "async_get_session",
+      errorPrefix: "Failed to get session"
+    });
+  }
+  /**
+   * Lists the remote sessions of a user.
+   *
+   * Calls the `async_list_sessions` class method unless `config.classMethod`
+   * names another one.
+   */
+  async listSessions(config) {
+    return this.callQueryMethod({
+      config,
+      defaultClassMethod: "async_list_sessions",
+      errorPrefix: "Failed to list sessions"
+    });
+  }
+  /**
+   * Deletes a remote session.
+   *
+   * Calls the `async_delete_session` class method unless `config.classMethod`
+   * names another one.
+   */
+  async deleteSession(config) {
+    await this.callQueryMethod({
+      config,
+      defaultClassMethod: "async_delete_session",
+      errorPrefix: "Failed to delete session"
+    });
+  }
+  /**
+   * Sends a query and waits for the whole response.
+   *
+   * Calls the `query` class method unless `config.classMethod` names another
+   * one.
+   */
+  async query(config) {
+    return this.callQueryMethod({
+      config,
+      defaultClassMethod: "query",
+      errorPrefix: "Failed to execute query"
+    });
+  }
+  /**
+   * Sends a query and yields each streamed response value.
+   *
+   * Calls the `async_stream_query` class method unless `config.classMethod`
+   * names another one. The stream can be server-sent events or
+   * newline-delimited JSON.
+   */
+  async *streamQuery(config) {
+    var _a;
+    try {
+      const { helpers: sdkHelpers, client } = await this.getClient();
+      const stream = client.streamQueryReasoningEngine({
+        name: this.reasoningEnginePath,
+        classMethod: (_a = config.classMethod) != null ? _a : "async_stream_query",
+        input: buildInputStruct(sdkHelpers, toWireInput(config))
+      });
+      yield* parseStream(stream);
+    } catch (err) {
+      if (err instanceof AgentExecutionError) {
+        throw err;
+      }
+      throw new AgentExecutionError(
+        `Failed to execute stream query: ${getErrorMessage(err)}`,
+        err
+      );
+    }
+  }
+  /**
+   * Closes the underlying gRPC client.
+   *
+   * Does nothing if no call has created the client yet, and is safe to call
+   * more than once. A later call creates a new client.
+   */
+  async close() {
+    const pending = this.clientPromise;
+    if (!pending) {
+      return;
+    }
+    this.clientPromise = void 0;
+    let loaded;
+    try {
+      loaded = await pending;
+    } catch {
+      return;
+    }
+    await loaded.client.close();
+  }
+  async callQueryMethod({
+    config,
+    defaultClassMethod,
+    errorPrefix
+  }) {
+    var _a;
+    try {
+      const { helpers: sdkHelpers, client } = await this.getClient();
+      const [response] = await client.queryReasoningEngine({
+        name: this.reasoningEnginePath,
+        classMethod: (_a = config.classMethod) != null ? _a : defaultClassMethod,
+        input: buildInputStruct(sdkHelpers, toWireInput(config))
+      });
+      return response.output ? parseOutput(sdkHelpers, response.output) : void 0;
+    } catch (err) {
+      throw new AgentExecutionError(
+        `${errorPrefix}: ${getErrorMessage(err)}`,
+        err
+      );
+    }
+  }
+  /**
+   * Resolves the SDK helpers and the gRPC client, loading the optional peer on
+   * first use.
+   *
+   * A created client is kept. A failure is not kept, so the next call tries
+   * again, for example once the peer is installed.
+   */
+  getClient() {
+    if (!this.clientPromise) {
+      const promise = this.createClient();
+      this.clientPromise = promise;
+      promise.catch(() => {
+        if (this.clientPromise === promise) {
+          this.clientPromise = void 0;
+        }
+      });
+    }
+    return this.clientPromise;
+  }
+  async createClient() {
+    const aiplatform = await loadOptionalPeer(
+      { packageName: "@google-cloud/aiplatform", feature: "AgentEngineClient" },
+      () => Promise.resolve().then(() => __toESM(require("@google-cloud/aiplatform"), 1))
+    );
+    return {
+      helpers: aiplatform.helpers,
+      client: new aiplatform.v1.ReasoningEngineExecutionServiceClient({
+        apiEndpoint: this.apiEndpoint
+      })
+    };
+  }
+};
+__name(_AgentEngineClient, "AgentEngineClient");
+let AgentEngineClient = _AgentEngineClient;
+export {
+  AgentEngineClient,
+  AgentExecutionError
+};

package/dist/esm/agents/base_agent.js

@@ -8,13 +8,10 @@ var __name = (target, value) => __defProp(target, "name", { value, configurable:
  * SPDX-License-Identifier: Apache-2.0
  */
 var _a, _b;
-import { context, trace } from "@opentelemetry/api";
 import { createEvent } from "../events/event.js";
-import {
-  runAsyncGeneratorWithOtelContext,
-  traceAgentInvocation,
-  tracer
-} from "../telemetry/tracing.js";
+import { recordAgentInvocation } from "../telemetry/api/instrumentation.js";
+import { isInvocationCancellation, toError } from "../utils/error_utils.js";
+import { logger } from "../utils/logger.js";
 import { BaseNode } from "../workflow/base_node.js";
 import { Context } from "./context.js";
 import { InvocationContext } from "./invocation_context.js";
@@ -98,38 +95,38 @@ const _BaseAgent = class _BaseAgent extends (_b = BaseNode, _a = BASE_AGENT_SIGN
    * @returns An AsyncGenerator that yields the events generated by the agent.
    */
   async *runAsync(parentContext) {
-    const span = tracer.startSpan(`invoke_agent ${this.name}`);
-    const ctx = trace.setSpan(context.active(), span);
-    try {
-      yield* runAsyncGeneratorWithOtelContext(
-        ctx,
-        this,
-        async function* () {
-          var _a2, _b2;
-          const context2 = this.createInvocationContext(parentContext);
-          const beforeAgentCallbackEvent = await this.handleBeforeAgentCallback(context2);
+    const agentCtx = this;
+    yield* recordAgentInvocation(
+      { agent: this, invocationContext: parentContext },
+      async function* () {
+        var _a2, _b2;
+        const context = agentCtx.createInvocationContext(parentContext);
+        try {
+          const beforeAgentCallbackEvent = await agentCtx.handleBeforeAgentCallback(context);
           if (beforeAgentCallbackEvent) {
             yield beforeAgentCallbackEvent;
           }
-          if (context2.endInvocation || ((_a2 = parentContext.abortSignal) == null ? void 0 : _a2.aborted)) {
+          if (context.endInvocation || ((_a2 = parentContext.abortSignal) == null ? void 0 : _a2.aborted)) {
             return;
           }
-          traceAgentInvocation({ agent: this, invocationContext: context2 });
-          for await (const event of this.runAsyncImpl(context2)) {
+          for await (const event of agentCtx.runAsyncImpl(context)) {
             yield event;
           }
-          if (context2.endInvocation || ((_b2 = parentContext.abortSignal) == null ? void 0 : _b2.aborted)) {
+          if (context.endInvocation || ((_b2 = parentContext.abortSignal) == null ? void 0 : _b2.aborted)) {
             return;
           }
-          const afterAgentCallbackEvent = await this.handleAfterAgentCallback(context2);
+          const afterAgentCallbackEvent = await agentCtx.handleAfterAgentCallback(context);
           if (afterAgentCallbackEvent) {
             yield afterAgentCallbackEvent;
           }
+        } catch (error) {
+          if (!isInvocationCancellation(error, context.abortSignal)) {
+            await agentCtx.handleAgentErrorCallback(context, error);
+          }
+          throw error;
         }
-      );
-    } finally {
-      span.end();
-    }
+      }
+    );
   }
   /**
    * Runs this agent as a workflow node.
@@ -164,36 +161,52 @@ const _BaseAgent = class _BaseAgent extends (_b = BaseNode, _a = BASE_AGENT_SIGN
    * @returns An AsyncGenerator that yields the events generated by the agent.
    */
   async *runLive(parentContext) {
-    const span = tracer.startSpan(`invoke_agent ${this.name}`);
-    const ctx = trace.setSpan(context.active(), span);
-    try {
-      yield* runAsyncGeneratorWithOtelContext(
-        ctx,
-        this,
-        async function* () {
-          var _a2, _b2;
-          const context2 = this.createInvocationContext(parentContext);
-          const beforeAgentCallbackEvent = await this.handleBeforeAgentCallback(context2);
+    const agentCtx = this;
+    yield* recordAgentInvocation(
+      { agent: this, invocationContext: parentContext },
+      async function* () {
+        var _a2, _b2;
+        const context = agentCtx.createInvocationContext(parentContext);
+        try {
+          const beforeAgentCallbackEvent = await agentCtx.handleBeforeAgentCallback(context);
           if (beforeAgentCallbackEvent) {
             yield beforeAgentCallbackEvent;
           }
-          if (context2.endInvocation || ((_a2 = parentContext.abortSignal) == null ? void 0 : _a2.aborted)) {
+          if (context.endInvocation || ((_a2 = parentContext.abortSignal) == null ? void 0 : _a2.aborted)) {
             return;
           }
-          for await (const event of this.runLiveImpl(context2)) {
+          for await (const event of agentCtx.runLiveImpl(context)) {
             yield event;
           }
-          if (context2.endInvocation || ((_b2 = parentContext.abortSignal) == null ? void 0 : _b2.aborted)) {
+          if (context.endInvocation || ((_b2 = parentContext.abortSignal) == null ? void 0 : _b2.aborted)) {
             return;
           }
-          const afterAgentCallbackEvent = await this.handleAfterAgentCallback(context2);
+          const afterAgentCallbackEvent = await agentCtx.handleAfterAgentCallback(context);
           if (afterAgentCallbackEvent) {
             yield afterAgentCallbackEvent;
           }
+        } catch (error) {
+          if (!isInvocationCancellation(error, context.abortSignal)) {
+            await agentCtx.handleAgentErrorCallback(context, error);
+          }
+          throw error;
         }
+      }
+    );
+  }
+  /** Reports execution failures without replacing the original thrown value. */
+  async handleAgentErrorCallback(invocationContext, error) {
+    try {
+      await invocationContext.pluginManager.runOnAgentErrorCallback({
+        agent: this,
+        callbackContext: new Context({ invocationContext }),
+        error: toError(error)
+      });
+    } catch (callbackError) {
+      logger.error(
+        "Agent error notification failed; preserving the original error.",
+        callbackError
       );
-    } finally {
-      span.end();
     }
   }
   /**
@@ -235,6 +248,22 @@ const _BaseAgent = class _BaseAgent extends (_b = BaseNode, _a = BASE_AGENT_SIGN
       agent: this
     });
   }
+  /**
+   * The resolved beforeAgentCallback field as a list of SingleAgentCallback.
+   *
+   * This method is only for use by Agent Development Kit.
+   */
+  get canonicalBeforeAgentCallbacks() {
+    return this.beforeAgentCallback;
+  }
+  /**
+   * The resolved afterAgentCallback field as a list of SingleAgentCallback.
+   *
+   * This method is only for use by Agent Development Kit.
+   */
+  get canonicalAfterAgentCallbacks() {
+    return this.afterAgentCallback;
+  }
   /**
    * Runs the before agent callback if it exists.
    *
@@ -243,27 +272,40 @@ const _BaseAgent = class _BaseAgent extends (_b = BaseNode, _a = BASE_AGENT_SIGN
    *     generated.
    */
   async handleBeforeAgentCallback(invocationContext) {
-    var _a2;
-    if (this.beforeAgentCallback.length === 0) {
+    var _a2, _b2, _c;
+    const hasPlugins = Boolean((_a2 = invocationContext.pluginManager) == null ? void 0 : _a2.hasPlugins);
+    if (!hasPlugins && this.canonicalBeforeAgentCallbacks.length === 0) {
       return void 0;
     }
     const callbackContext = new Context({ invocationContext });
-    for (const callback of this.beforeAgentCallback) {
-      const content = await callback(callbackContext);
-      if ((_a2 = invocationContext.abortSignal) == null ? void 0 : _a2.aborted) {
-        return;
-      }
-      if (content) {
-        invocationContext.endInvocation = true;
-        return createEvent({
-          invocationId: invocationContext.invocationId,
-          author: this.name,
-          branch: invocationContext.branch,
-          content,
-          actions: callbackContext.eventActions
-        });
+    let content = hasPlugins ? await invocationContext.pluginManager.runBeforeAgentCallback({
+      agent: this,
+      callbackContext
+    }) : void 0;
+    if ((_b2 = invocationContext.abortSignal) == null ? void 0 : _b2.aborted) {
+      return;
+    }
+    if (!content) {
+      for (const callback of this.canonicalBeforeAgentCallbacks) {
+        content = await callback(callbackContext);
+        if ((_c = invocationContext.abortSignal) == null ? void 0 : _c.aborted) {
+          return;
+        }
+        if (content) {
+          break;
+        }
       }
     }
+    if (content) {
+      invocationContext.endInvocation = true;
+      return createEvent({
+        invocationId: invocationContext.invocationId,
+        author: this.name,
+        branch: invocationContext.branch,
+        content,
+        actions: callbackContext.eventActions
+      });
+    }
     if (callbackContext.state.hasDelta()) {
       return createEvent({
         invocationId: invocationContext.invocationId,
@@ -282,26 +324,39 @@ const _BaseAgent = class _BaseAgent extends (_b = BaseNode, _a = BASE_AGENT_SIGN
    *     generated.
    */
   async handleAfterAgentCallback(invocationContext) {
-    var _a2;
-    if (this.afterAgentCallback.length === 0) {
+    var _a2, _b2, _c;
+    const hasPlugins = Boolean((_a2 = invocationContext.pluginManager) == null ? void 0 : _a2.hasPlugins);
+    if (!hasPlugins && this.canonicalAfterAgentCallbacks.length === 0) {
       return void 0;
     }
     const callbackContext = new Context({ invocationContext });
-    for (const callback of this.afterAgentCallback) {
-      const content = await callback(callbackContext);
-      if ((_a2 = invocationContext.abortSignal) == null ? void 0 : _a2.aborted) {
-        return;
-      }
-      if (content) {
-        return createEvent({
-          invocationId: invocationContext.invocationId,
-          author: this.name,
-          branch: invocationContext.branch,
-          content,
-          actions: callbackContext.eventActions
-        });
+    let content = hasPlugins ? await invocationContext.pluginManager.runAfterAgentCallback({
+      agent: this,
+      callbackContext
+    }) : void 0;
+    if ((_b2 = invocationContext.abortSignal) == null ? void 0 : _b2.aborted) {
+      return;
+    }
+    if (!content) {
+      for (const callback of this.canonicalAfterAgentCallbacks) {
+        content = await callback(callbackContext);
+        if ((_c = invocationContext.abortSignal) == null ? void 0 : _c.aborted) {
+          return;
+        }
+        if (content) {
+          break;
+        }
       }
     }
+    if (content) {
+      return createEvent({
+        invocationId: invocationContext.invocationId,
+        author: this.name,
+        branch: invocationContext.branch,
+        content,
+        actions: callbackContext.eventActions
+      });
+    }
     if (callbackContext.state.hasDelta()) {
       return createEvent({
         invocationId: invocationContext.invocationId,