@google/adk 0.2.5 → 0.4.0

package/dist/esm/runner/runner.js

@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import { createPartFromText } from "@google/genai";
-import { trace } from "@opentelemetry/api";
+import { context, trace } from "@opentelemetry/api";
 import {
   InvocationContext,
   newInvocationContextId
@@ -18,6 +18,10 @@ import {
 import { createEvent, getFunctionCalls } from "../events/event.js";
 import { createEventActions } from "../events/event_actions.js";
 import { PluginManager } from "../plugins/plugin_manager.js";
+import {
+  runAsyncGeneratorWithOtelContext,
+  tracer
+} from "../telemetry/tracing.js";
 import { logger } from "../utils/logger.js";
 import { isGemini2OrAbove } from "../utils/model_name.js";
 class Runner {
@@ -31,6 +35,37 @@ class Runner {
     this.memoryService = input.memoryService;
     this.credentialService = input.credentialService;
   }
+  /**
+   * Runs the agent with a new, ephemeral session.
+   *
+   * @param params.userId The user ID of the session.
+   * @param params.newMessage A new message to append to the session.
+   * @param params.stateDelta An optional state delta to apply to the session.
+   * @param params.runConfig The run config for the agent.
+   * @yields The Events generated by the agent.
+   */
+  async *runEphemeral(params) {
+    const session = await this.sessionService.createSession({
+      appName: this.appName,
+      userId: params.userId
+    });
+    const sessionId = session.id;
+    try {
+      yield* this.runAsync({
+        userId: params.userId,
+        sessionId,
+        newMessage: params.newMessage,
+        stateDelta: params.stateDelta,
+        runConfig: params.runConfig
+      });
+    } finally {
+      await this.sessionService.deleteSession({
+        appName: this.appName,
+        userId: params.userId,
+        sessionId
+      });
+    }
+  }
   /**
    * Runs the agent with the given message, and returns an async generator of
    * events.
@@ -44,109 +79,123 @@ class Runner {
    */
   // TODO - b/425992518: user, sessionId, and runConfig can be internalized.
   async *runAsync(params) {
-    var _a;
     const { userId, sessionId, stateDelta } = params;
     const runConfig = createRunConfig(params.runConfig);
     let newMessage = params.newMessage;
-    const span = trace.getTracer("gcp.vertex.agent").startSpan("invocation");
+    const span = tracer.startSpan("invocation");
+    const ctx = trace.setSpan(context.active(), span);
     try {
-      const session = await this.sessionService.getSession({
-        appName: this.appName,
-        userId,
-        sessionId
-      });
-      if (!session) {
-        if (!this.appName) {
-          throw new Error(
-            `Session lookup failed: appName must be provided in runner constructor`
-          );
-        }
-        throw new Error(`Session not found: ${sessionId}`);
-      }
-      if (runConfig.supportCfc && isLlmAgent(this.agent)) {
-        const modelName = this.agent.canonicalModel.model;
-        if (!isGemini2OrAbove(modelName)) {
-          throw new Error(
-            `CFC is not supported for model: ${modelName} in agent: ${this.agent.name}`
-          );
-        }
-        if (!isBuiltInCodeExecutor(this.agent.codeExecutor)) {
-          this.agent.codeExecutor = new BuiltInCodeExecutor();
-        }
-      }
-      const invocationContext = new InvocationContext({
-        artifactService: this.artifactService,
-        sessionService: this.sessionService,
-        memoryService: this.memoryService,
-        credentialService: this.credentialService,
-        invocationId: newInvocationContextId(),
-        agent: this.agent,
-        session,
-        userContent: newMessage,
-        runConfig,
-        pluginManager: this.pluginManager
-      });
-      const pluginUserMessage = await this.pluginManager.runOnUserMessageCallback({
-        userMessage: newMessage,
-        invocationContext
-      });
-      if (pluginUserMessage) {
-        newMessage = pluginUserMessage;
-      }
-      if (newMessage) {
-        if (!((_a = newMessage.parts) == null ? void 0 : _a.length)) {
-          throw new Error("No parts in the newMessage.");
-        }
-        if (runConfig.saveInputBlobsAsArtifacts) {
-          await this.saveArtifacts(
-            invocationContext.invocationId,
-            session.userId,
-            session.id,
-            newMessage
-          );
-        }
-        await this.sessionService.appendEvent({
-          session,
-          event: createEvent({
-            invocationId: invocationContext.invocationId,
-            author: "user",
-            actions: stateDelta ? createEventActions({ stateDelta }) : void 0,
-            content: newMessage
-          })
-        });
-      }
-      invocationContext.agent = this.determineAgentForResumption(
-        session,
-        this.agent
-      );
-      const beforeRunCallbackResponse = await this.pluginManager.runBeforeRunCallback({ invocationContext });
-      if (beforeRunCallbackResponse) {
-        const earlyExitEvent = createEvent({
-          invocationId: invocationContext.invocationId,
-          author: "model",
-          content: beforeRunCallbackResponse
-        });
-        await this.sessionService.appendEvent({ session, event: earlyExitEvent });
-        yield earlyExitEvent;
-      } else {
-        for await (const event of invocationContext.agent.runAsync(
-          invocationContext
-        )) {
-          if (!event.partial) {
-            await this.sessionService.appendEvent({ session, event });
+      yield* runAsyncGeneratorWithOtelContext(
+        ctx,
+        this,
+        async function* () {
+          var _a;
+          const session = await this.sessionService.getSession({
+            appName: this.appName,
+            userId,
+            sessionId
+          });
+          if (!session) {
+            if (!this.appName) {
+              throw new Error(
+                `Session lookup failed: appName must be provided in runner constructor`
+              );
+            }
+            throw new Error(`Session not found: ${sessionId}`);
+          }
+          if (runConfig.supportCfc && isLlmAgent(this.agent)) {
+            const modelName = this.agent.canonicalModel.model;
+            if (!isGemini2OrAbove(modelName)) {
+              throw new Error(
+                `CFC is not supported for model: ${modelName} in agent: ${this.agent.name}`
+              );
+            }
+            if (!isBuiltInCodeExecutor(this.agent.codeExecutor)) {
+              this.agent.codeExecutor = new BuiltInCodeExecutor();
+            }
           }
-          const modifiedEvent = await this.pluginManager.runOnEventCallback({
-            invocationContext,
-            event
+          const invocationContext = new InvocationContext({
+            artifactService: this.artifactService,
+            sessionService: this.sessionService,
+            memoryService: this.memoryService,
+            credentialService: this.credentialService,
+            invocationId: newInvocationContextId(),
+            agent: this.agent,
+            session,
+            userContent: newMessage,
+            runConfig,
+            pluginManager: this.pluginManager
+          });
+          const pluginUserMessage = await this.pluginManager.runOnUserMessageCallback({
+            userMessage: newMessage,
+            invocationContext
           });
-          if (modifiedEvent) {
-            yield modifiedEvent;
-          } else {
-            yield event;
+          if (pluginUserMessage) {
+            newMessage = pluginUserMessage;
+          }
+          if (newMessage) {
+            if (!((_a = newMessage.parts) == null ? void 0 : _a.length)) {
+              throw new Error("No parts in the newMessage.");
+            }
+            if (runConfig.saveInputBlobsAsArtifacts) {
+              await this.saveArtifacts(
+                invocationContext.invocationId,
+                session.userId,
+                session.id,
+                newMessage
+              );
+            }
+            await this.sessionService.appendEvent({
+              session,
+              event: createEvent({
+                invocationId: invocationContext.invocationId,
+                author: "user",
+                actions: stateDelta ? createEventActions({ stateDelta }) : void 0,
+                content: newMessage
+              })
+            });
+          }
+          invocationContext.agent = this.determineAgentForResumption(
+            session,
+            this.agent
+          );
+          if (newMessage) {
+            const beforeRunCallbackResponse = await this.pluginManager.runBeforeRunCallback({
+              invocationContext
+            });
+            if (beforeRunCallbackResponse) {
+              const earlyExitEvent = createEvent({
+                invocationId: invocationContext.invocationId,
+                author: "model",
+                content: beforeRunCallbackResponse
+              });
+              await this.sessionService.appendEvent({
+                session,
+                event: earlyExitEvent
+              });
+              yield earlyExitEvent;
+            } else {
+              for await (const event of invocationContext.agent.runAsync(
+                invocationContext
+              )) {
+                if (!event.partial) {
+                  await this.sessionService.appendEvent({ session, event });
+                }
+                const modifiedEvent = await this.pluginManager.runOnEventCallback({
+                  invocationContext,
+                  event
+                });
+                if (modifiedEvent) {
+                  yield modifiedEvent;
+                } else {
+                  yield event;
+                }
+              }
+              await this.pluginManager.runAfterRunCallback({ invocationContext });
+            }
           }
         }
-      }
-      await this.pluginManager.runAfterRunCallback({ invocationContext });
+      );
     } finally {
       span.end();
     }

package/dist/esm/sessions/base_session_service.js

@@ -3,8 +3,29 @@
  * Copyright 2025 Google LLC
  * SPDX-License-Identifier: Apache-2.0
  */
+import { cloneDeep } from "lodash-es";
 import { State } from "./state.js";
 class BaseSessionService {
+  /**
+   * Gets a session or creates one if it doesn't exist.
+   *
+   * @param request The request to get or create a session.
+   * @return A promise that resolves to the session instance.
+   */
+  async getOrCreateSession(request) {
+    if (!request.sessionId) {
+      return this.createSession(request);
+    }
+    const session = await this.getSession({
+      appName: request.appName,
+      userId: request.userId,
+      sessionId: request.sessionId
+    });
+    if (session) {
+      return session;
+    }
+    return this.createSession(request);
+  }
   /**
    * Appends an event to a session.
    *
@@ -15,6 +36,7 @@ class BaseSessionService {
     if (event.partial) {
       return event;
     }
+    event = trimTempDeltaState(event);
     this.updateSessionState({ session, event });
     session.events.push(event);
     return event;
@@ -36,6 +58,32 @@ class BaseSessionService {
     }
   }
 }
+function trimTempDeltaState(event) {
+  if (!event.actions || !event.actions.stateDelta) {
+    return event;
+  }
+  const stateDelta = event.actions.stateDelta;
+  const filteredStateDelta = {};
+  for (const [key, value] of Object.entries(stateDelta)) {
+    if (!key.startsWith(State.TEMP_PREFIX)) {
+      filteredStateDelta[key] = value;
+    }
+  }
+  event.actions.stateDelta = filteredStateDelta;
+  return event;
+}
+function mergeStates(appState = {}, userState = {}, sessionState = {}) {
+  const merged = cloneDeep(sessionState);
+  for (const [k, v] of Object.entries(appState)) {
+    merged[State.APP_PREFIX + k] = v;
+  }
+  for (const [k, v] of Object.entries(userState)) {
+    merged[State.USER_PREFIX + k] = v;
+  }
+  return merged;
+}
 export {
-  BaseSessionService
+  BaseSessionService,
+  mergeStates,
+  trimTempDeltaState
 };

package/dist/esm/sessions/database_session_service.js

@@ -0,0 +1,350 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import {
+  LockMode,
+  MikroORM
+} from "@mikro-orm/core";
+import { randomUUID } from "../utils/env_aware_utils.js";
+import {
+  BaseSessionService,
+  mergeStates,
+  trimTempDeltaState
+} from "./base_session_service.js";
+import {
+  ensureDatabaseCreated,
+  getConnectionOptionsFromUri,
+  validateDatabaseSchemaVersion
+} from "./db/operations.js";
+import {
+  ENTITIES,
+  StorageAppState,
+  StorageEvent,
+  StorageSession,
+  StorageUserState
+} from "./db/schema.js";
+import { createSession } from "./session.js";
+import { State } from "./state.js";
+function isDatabaseConnectionString(uri) {
+  if (!uri) {
+    return false;
+  }
+  return uri.startsWith("postgres://") || uri.startsWith("postgresql://") || uri.startsWith("mysql://") || uri.startsWith("mariadb://") || uri.startsWith("mssql://") || uri.startsWith("sqlite://");
+}
+class DatabaseSessionService extends BaseSessionService {
+  constructor(connectionStringOrOptions) {
+    super();
+    this.initialized = false;
+    if (typeof connectionStringOrOptions === "string") {
+      this.options = getConnectionOptionsFromUri(connectionStringOrOptions);
+    } else {
+      if (!connectionStringOrOptions.driver) {
+        throw new Error("Driver is required when passing options object.");
+      }
+      this.options = {
+        ...connectionStringOrOptions,
+        entities: ENTITIES
+      };
+    }
+  }
+  async init() {
+    if (this.initialized) {
+      return;
+    }
+    this.orm = await MikroORM.init(this.options);
+    await ensureDatabaseCreated(this.orm);
+    await validateDatabaseSchemaVersion(this.orm);
+    this.initialized = true;
+  }
+  async createSession({
+    appName,
+    userId,
+    state,
+    sessionId
+  }) {
+    await this.init();
+    const em = this.orm.em.fork();
+    const id = sessionId || randomUUID();
+    const now = /* @__PURE__ */ new Date();
+    const existing = await em.findOne(StorageSession, {
+      id,
+      appName,
+      userId
+    });
+    if (existing) {
+      throw new Error(`Session with id ${id} already exists.`);
+    }
+    let appStateModel = await em.findOne(StorageAppState, { appName });
+    if (!appStateModel) {
+      appStateModel = em.create(StorageAppState, {
+        appName,
+        state: {},
+        updateTime: now
+      });
+      em.persist(appStateModel);
+    }
+    let userStateModel = await em.findOne(StorageUserState, { appName, userId });
+    if (!userStateModel) {
+      userStateModel = em.create(StorageUserState, {
+        appName,
+        userId,
+        state: {}
+      });
+      em.persist(userStateModel);
+    }
+    const appStateDelta = {};
+    const userStateDelta = {};
+    const sessionState = {};
+    if (state) {
+      for (const [key, value] of Object.entries(state)) {
+        if (key.startsWith(State.APP_PREFIX)) {
+          appStateDelta[key.replace(State.APP_PREFIX, "")] = value;
+        } else if (key.startsWith(State.USER_PREFIX)) {
+          userStateDelta[key.replace(State.USER_PREFIX, "")] = value;
+        } else {
+          sessionState[key] = value;
+        }
+      }
+    }
+    if (Object.keys(appStateDelta).length > 0) {
+      appStateModel.state = { ...appStateModel.state, ...appStateDelta };
+    }
+    if (Object.keys(userStateDelta).length > 0) {
+      userStateModel.state = { ...userStateModel.state, ...userStateDelta };
+    }
+    const storageSession = em.create(StorageSession, {
+      id,
+      appName,
+      userId,
+      state: sessionState,
+      createTime: now,
+      updateTime: now
+    });
+    em.persist(storageSession);
+    await em.flush();
+    const mergedState = mergeStates(
+      appStateModel.state,
+      userStateModel.state,
+      sessionState
+    );
+    return createSession({
+      id,
+      appName,
+      userId,
+      state: mergedState,
+      events: [],
+      lastUpdateTime: storageSession.createTime.getTime()
+    });
+  }
+  async getSession({
+    appName,
+    userId,
+    sessionId,
+    config
+  }) {
+    await this.init();
+    const em = this.orm.em.fork();
+    const storageSession = await em.findOne(StorageSession, {
+      appName,
+      userId,
+      id: sessionId
+    });
+    if (!storageSession) {
+      return void 0;
+    }
+    const eventWhere = {
+      appName,
+      userId,
+      sessionId
+    };
+    if (config == null ? void 0 : config.afterTimestamp) {
+      eventWhere.timestamp = { $gt: new Date(config.afterTimestamp) };
+    }
+    const storageEvents = await em.find(StorageEvent, eventWhere, {
+      orderBy: { timestamp: "DESC" },
+      limit: config == null ? void 0 : config.numRecentEvents
+    });
+    storageEvents.reverse();
+    const appStateModel = await em.findOne(StorageAppState, { appName });
+    const userStateModel = await em.findOne(StorageUserState, {
+      appName,
+      userId
+    });
+    const mergedState = mergeStates(
+      (appStateModel == null ? void 0 : appStateModel.state) || {},
+      (userStateModel == null ? void 0 : userStateModel.state) || {},
+      storageSession.state
+    );
+    return createSession({
+      id: sessionId,
+      appName,
+      userId,
+      state: mergedState,
+      events: storageEvents.map((se) => se.eventData),
+      lastUpdateTime: storageSession.updateTime.getTime()
+    });
+  }
+  async listSessions({
+    appName,
+    userId
+  }) {
+    await this.init();
+    const em = this.orm.em.fork();
+    const where = { appName };
+    if (userId) {
+      where.userId = userId;
+    }
+    const storageSessions = await em.find(StorageSession, where);
+    const appStateModel = await em.findOne(StorageAppState, { appName });
+    const appState = (appStateModel == null ? void 0 : appStateModel.state) || {};
+    const userStateMap = {};
+    if (userId) {
+      const u = await em.findOne(StorageUserState, { appName, userId });
+      if (u) userStateMap[userId] = u.state;
+    } else {
+      const allUserStates = await em.find(StorageUserState, { appName });
+      for (const u of allUserStates) {
+        userStateMap[u.userId] = u.state;
+      }
+    }
+    const sessions = storageSessions.map((ss) => {
+      const uState = userStateMap[ss.userId] || {};
+      const merged = mergeStates(appState, uState, ss.state);
+      return createSession({
+        id: ss.id,
+        appName: ss.appName,
+        userId: ss.userId,
+        state: merged,
+        events: [],
+        lastUpdateTime: ss.updateTime.getTime()
+      });
+    });
+    return { sessions };
+  }
+  async deleteSession({
+    appName,
+    userId,
+    sessionId
+  }) {
+    await this.init();
+    const em = this.orm.em.fork();
+    await em.nativeDelete(StorageSession, { appName, userId, id: sessionId });
+    await em.nativeDelete(StorageEvent, { appName, userId, sessionId });
+  }
+  async appendEvent({
+    session,
+    event
+  }) {
+    await this.init();
+    const em = this.orm.em.fork();
+    if (event.partial) {
+      return event;
+    }
+    const trimmedEvent = trimTempDeltaState(event);
+    await em.transactional(async (txEm) => {
+      const storageSession = await txEm.findOne(
+        StorageSession,
+        {
+          appName: session.appName,
+          userId: session.userId,
+          id: session.id
+        },
+        { lockMode: LockMode.PESSIMISTIC_WRITE }
+      );
+      if (!storageSession) {
+        throw new Error(`Session ${session.id} not found for appendEvent`);
+      }
+      let appStateModel = await txEm.findOne(StorageAppState, {
+        appName: session.appName
+      });
+      if (!appStateModel) {
+        appStateModel = txEm.create(StorageAppState, {
+          appName: session.appName,
+          state: {},
+          updateTime: /* @__PURE__ */ new Date()
+        });
+        txEm.persist(appStateModel);
+      }
+      let userStateModel = await txEm.findOne(StorageUserState, {
+        appName: session.appName,
+        userId: session.userId
+      });
+      if (!userStateModel) {
+        userStateModel = txEm.create(StorageUserState, {
+          appName: session.appName,
+          userId: session.userId,
+          state: {}
+        });
+        txEm.persist(userStateModel);
+      }
+      if (storageSession.updateTime.getTime() > session.lastUpdateTime) {
+        const events = await txEm.find(
+          StorageEvent,
+          {
+            appName: session.appName,
+            userId: session.userId,
+            sessionId: session.id
+          },
+          { orderBy: { timestamp: "ASC" } }
+        );
+        const mergedState = mergeStates(
+          appStateModel.state,
+          userStateModel.state,
+          storageSession.state
+        );
+        session.state = mergedState;
+        session.events = events.map((e) => e.eventData);
+      }
+      if (event.actions && event.actions.stateDelta) {
+        const appDelta = {};
+        const userDelta = {};
+        const sessionDelta = {};
+        for (const [key, value] of Object.entries(event.actions.stateDelta)) {
+          if (key.startsWith(State.APP_PREFIX)) {
+            appDelta[key.replace(State.APP_PREFIX, "")] = value;
+          } else if (key.startsWith(State.USER_PREFIX)) {
+            userDelta[key.replace(State.USER_PREFIX, "")] = value;
+          } else {
+            sessionDelta[key] = value;
+          }
+        }
+        if (Object.keys(appDelta).length > 0) {
+          appStateModel.state = { ...appStateModel.state, ...appDelta };
+        }
+        if (Object.keys(userDelta).length > 0) {
+          userStateModel.state = { ...userStateModel.state, ...userDelta };
+        }
+        if (Object.keys(sessionDelta).length > 0) {
+          storageSession.state = { ...storageSession.state, ...sessionDelta };
+        }
+      }
+      const newStorageEvent = txEm.create(StorageEvent, {
+        id: trimmedEvent.id,
+        appName: session.appName,
+        userId: session.userId,
+        sessionId: session.id,
+        invocationId: trimmedEvent.invocationId,
+        timestamp: new Date(trimmedEvent.timestamp),
+        eventData: trimmedEvent
+      });
+      txEm.persist(newStorageEvent);
+      await txEm.commit();
+      storageSession.updateTime = new Date(event.timestamp);
+      const newMergedState = mergeStates(
+        appStateModel.state,
+        userStateModel.state,
+        storageSession.state
+      );
+      session.state = newMergedState;
+      session.events.push(event);
+      session.lastUpdateTime = storageSession.updateTime.getTime();
+    });
+    return event;
+  }
+}
+export {
+  DatabaseSessionService,
+  isDatabaseConnectionString
+};