Feat: add BaseSessionService.flush() and drain it from the Runner (adk-python parity)

core/src/runner/runner.ts

@@ -454,6 +454,7 @@ export class Runner {
       span.end();
       const toolsets = getAllToolsets(this.agent);
       await Promise.allSettled(toolsets.map((t) => t.close()));
+      await this.sessionService.flush();
     }
   }
 

core/src/sessions/base_session_service.ts

@@ -183,6 +183,17 @@ export abstract class BaseSessionService {
     return event;
   }
 
+  /**
+   * Flushes any buffered events.
+   *
+   * A session service that batches writes should drain its buffer here. The
+   * `Runner` calls this when an invocation finishes. Implementations that
+   * write synchronously need no override; the default is a no-op.
+   *
+   * @return A promise that resolves when buffered events have been written.
+   */
+  async flush(): Promise<void> {}
+
   /**
    * Updates the session state based on the event.
    *

core/test/runner/runner_test.ts

@@ -6,6 +6,7 @@
 
 import {
   App,
+  AppendEventRequest,
   BaseAgent,
   BasePlugin,
   createEvent,
@@ -1252,3 +1253,138 @@ describe('Runner artifact saving (`saveInputBlobsAsArtifacts`)', () => {
     ]);
   });
 });
+
+/** A session service that only writes events out when it is drained. */
+class BufferingSessionService extends InMemorySessionService {
+  readonly pending: Event[] = [];
+  readonly written: Event[] = [];
+
+  override async appendEvent(request: AppendEventRequest): Promise<Event> {
+    const event = await super.appendEvent(request);
+    this.pending.push(event);
+    return event;
+  }
+
+  override async flush(): Promise<void> {
+    this.written.push(...this.pending);
+    this.pending.length = 0;
+  }
+}
+
+describe('Runner session service flush', () => {
+  let sessionService: InMemorySessionService;
+  let runner: Runner;
+
+  beforeEach(() => {
+    sessionService = new InMemorySessionService();
+    runner = new Runner({
+      appName: TEST_APP_ID,
+      agent: new MockLlmAgent('test_agent'),
+      sessionService,
+    });
+  });
+
+  async function runToCompletion(sessionId: string): Promise<void> {
+    for await (const _ of runner.runAsync({
+      userId: TEST_USER_ID,
+      sessionId,
+      newMessage: {role: 'user', parts: [{text: TEST_MESSAGE}]},
+    })) {
+      // Consume stream
+    }
+  }
+
+  it('drains the session service once per completed invocation', async () => {
+    const session = await sessionService.createSession({
+      appName: TEST_APP_ID,
+      userId: TEST_USER_ID,
+      sessionId: TEST_SESSION_ID,
+    });
+    const flushSpy = vi.spyOn(sessionService, 'flush');
+
+    await runToCompletion(session.id);
+
+    expect(flushSpy).toHaveBeenCalledTimes(1);
+  });
+
+  it('drains the session service after the last appendEvent', async () => {
+    const session = await sessionService.createSession({
+      appName: TEST_APP_ID,
+      userId: TEST_USER_ID,
+      sessionId: TEST_SESSION_ID,
+    });
+    const appendEventSpy = vi.spyOn(sessionService, 'appendEvent');
+    const flushSpy = vi.spyOn(sessionService, 'flush');
+
+    await runToCompletion(session.id);
+
+    expect(appendEventSpy.mock.invocationCallOrder.length).toBeGreaterThan(0);
+    expect(flushSpy.mock.invocationCallOrder[0]).toBeGreaterThan(
+      Math.max(...appendEventSpy.mock.invocationCallOrder),
+    );
+  });
+
+  it('drains the session service once per invocation, not once per runner', async () => {
+    const session = await sessionService.createSession({
+      appName: TEST_APP_ID,
+      userId: TEST_USER_ID,
+      sessionId: TEST_SESSION_ID,
+    });
+    const flushSpy = vi.spyOn(sessionService, 'flush');
+
+    await runToCompletion(session.id);
+    await runToCompletion(session.id);
+
+    expect(flushSpy).toHaveBeenCalledTimes(2);
+  });
+
+  it('drains the session service when the invocation throws', async () => {
+    const flushSpy = vi.spyOn(sessionService, 'flush');
+
+    await expect(runToCompletion('non_existent_session_id')).rejects.toThrow(
+      'Session not found: non_existent_session_id',
+    );
+
+    expect(flushSpy).toHaveBeenCalledTimes(1);
+  });
+
+  it('drains the session service before runEphemeral deletes the session', async () => {
+    const flushSpy = vi.spyOn(sessionService, 'flush');
+    const deleteSessionSpy = vi.spyOn(sessionService, 'deleteSession');
+
+    for await (const _ of runner.runEphemeral({
+      userId: TEST_USER_ID,
+      newMessage: {role: 'user', parts: [{text: TEST_MESSAGE}]},
+    })) {
+      // Consume stream
+    }
+
+    expect(flushSpy).toHaveBeenCalledTimes(1);
+    expect(deleteSessionSpy).toHaveBeenCalledTimes(1);
+    expect(flushSpy.mock.invocationCallOrder[0]).toBeLessThan(
+      deleteSessionSpy.mock.invocationCallOrder[0],
+    );
+  });
+
+  it('drains a buffering session service by the end of the invocation', async () => {
+    const bufferingService = new BufferingSessionService();
+    const bufferingRunner = new Runner({
+      appName: TEST_APP_ID,
+      agent: new MockLlmAgent('test_agent'),
+      sessionService: bufferingService,
+    });
+
+    for await (const _ of bufferingRunner.runEphemeral({
+      userId: TEST_USER_ID,
+      newMessage: {role: 'user', parts: [{text: TEST_MESSAGE}]},
+    })) {
+      // Consume stream
+    }
+
+    expect(bufferingService.pending).toEqual([]);
+    expect(bufferingService.written.map((e) => e.author)).toEqual([
+      'user',
+      'test_agent',
+    ]);
+  });
+});

core/test/sessions/base_session_service_test.ts

@@ -0,0 +1,83 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {
+  BaseSessionService,
+  CreateSessionRequest,
+  DeleteSessionRequest,
+  GetSessionRequest,
+  InMemorySessionService,
+  ListSessionsRequest,
+  ListSessionsResponse,
+  Session,
+  createEvent,
+  createSession,
+} from '@google/adk';
+import {describe, expect, it} from 'vitest';
+
+const TEST_APP_NAME = 'test-app';
+const TEST_USER_ID = 'test-user';
+
+/** The smallest subclass that satisfies the abstract members. */
+class MinimalSessionService extends BaseSessionService {
+  override async createSession(
+    request: CreateSessionRequest,
+  ): Promise<Session> {
+    return createSession({
+      id: request.sessionId ?? 'session-id',
+      appName: request.appName,
+      userId: request.userId,
+    });
+  }
+
+  override async getSession(
+    _request: GetSessionRequest,
+  ): Promise<Session | undefined> {
+    return undefined;
+  }
+
+  override async listSessions(
+    _request: ListSessionsRequest,
+  ): Promise<ListSessionsResponse> {
+    return {sessions: [], page: 1, limit: 0, totalItems: 0, totalPages: 0};
+  }
+
+  override async deleteSession(_request: DeleteSessionRequest): Promise<void> {}
+}
+
+describe('BaseSessionService.flush', () => {
+  it('resolves to undefined by default', async () => {
+    const service = new MinimalSessionService();
+
+    await expect(service.flush()).resolves.toBeUndefined();
+  });
+
+  it('is inherited by a shipped service and leaves its events untouched', async () => {
+    const service = new InMemorySessionService();
+    const session = await service.createSession({
+      appName: TEST_APP_NAME,
+      userId: TEST_USER_ID,
+    });
+    await service.appendEvent({
+      session,
+      event: createEvent({
+        invocationId: 'invocation-id',
+        author: 'user',
+        content: {role: 'user', parts: [{text: 'Hello'}]},
+      }),
+    });
+
+    await expect(service.flush()).resolves.toBeUndefined();
+
+    const stored = await service.getSession({
+      appName: TEST_APP_NAME,
+      userId: TEST_USER_ID,
+      sessionId: session.id,
+    });
+    expect(stored?.events).toHaveLength(1);
+    expect(stored?.events[0].content?.parts).toEqual([{text: 'Hello'}]);
+  });
+});