diff --git a/core/src/runner/runner.ts b/core/src/runner/runner.ts index ba43f2a93..f1d3a7de2 100644 --- a/core/src/runner/runner.ts +++ b/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(); } } diff --git a/core/src/sessions/base_session_service.ts b/core/src/sessions/base_session_service.ts index 459fa17b2..69862a185 100644 --- a/core/src/sessions/base_session_service.ts +++ b/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 {} + /** * Updates the session state based on the event. * diff --git a/core/test/runner/runner_test.ts b/core/test/runner/runner_test.ts index dddee9b04..b1fb33047 100644 --- a/core/test/runner/runner_test.ts +++ b/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 { + const event = await super.appendEvent(request); + this.pending.push(event); + return event; + } + + override async flush(): Promise { + 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 { + 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', + ]); + }); +}); diff --git a/core/test/sessions/base_session_service_test.ts b/core/test/sessions/base_session_service_test.ts new file mode 100644 index 000000000..667259da2 --- /dev/null +++ b/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 { + return createSession({ + id: request.sessionId ?? 'session-id', + appName: request.appName, + userId: request.userId, + }); + } + + override async getSession( + _request: GetSessionRequest, + ): Promise { + return undefined; + } + + override async listSessions( + _request: ListSessionsRequest, + ): Promise { + return {sessions: [], page: 1, limit: 0, totalItems: 0, totalPages: 0}; + } + + override async deleteSession(_request: DeleteSessionRequest): Promise {} +} + +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'}]); + }); +});