diff --git a/core/src/sessions/vertex_ai_session_service.ts b/core/src/sessions/vertex_ai_session_service.ts index af9f7bb56..7a0f0c3f7 100644 --- a/core/src/sessions/vertex_ai_session_service.ts +++ b/core/src/sessions/vertex_ai_session_service.ts @@ -7,6 +7,7 @@ import {Client} from '@google-cloud/vertexai/build/src/genai/client.js'; import {Sessions} from '@google-cloud/vertexai/build/src/genai/sessions.js'; import { + EventActions as ApiEventActions, AppendAgentEngineSessionEventConfig, AppendAgentEngineSessionEventRequestParameters, EventMetadata, @@ -40,6 +41,7 @@ import {createSession, Session} from './session.js'; const DEFAULT_MAX_ATTEMPTS = 30; const GRPC_NOT_FOUND = 5; const HTTP_NOT_FOUND = 404; +const HTTP_BAD_REQUEST = 400; /** * Checks if the given URI is a Vertex AI session service URI. @@ -431,12 +433,16 @@ export class VertexAiSessionService extends BaseSessionService { } const config = partialCopy(event, [ - 'content', - 'actions', 'errorCode', 'errorMessage', ]); + const content = event.content && dropUnsupportedPartFields(event.content); + config.content = content; + config.actions = event.actions + ? toApiEventActions(event.actions) + : undefined; + config.eventMetadata = { ...partialCopy(event, [ 'partial', @@ -450,7 +456,7 @@ export class VertexAiSessionService extends BaseSessionService { Object.keys(customMetadata).length > 0 ? customMetadata : undefined, }; - config.rawEvent = JSON.parse(JSON.stringify(event)) as Record< + config.rawEvent = JSON.parse(JSON.stringify({...event, content})) as Record< string, unknown >; @@ -466,24 +472,66 @@ export class VertexAiSessionService extends BaseSessionService { try { await this.sessions.events.append(params); } catch (error) { + if (!isInvalidArgumentError(error)) { + throw error; + } logger.warn( - 'Failed to append event with rawEvent, falling back...', + 'appendEvent was rejected with rawEvent; retrying without it.', error, ); delete config.rawEvent; - await this.sessions.events.append({ - name: `reasoningEngines/${reasoningEngineId}/sessions/${session.id}`, - author: event.author || 'user', - invocationId: event.invocationId || `inv-${Date.now()}`, - timestamp: new Date(event.timestamp).toISOString(), - config, - }); + await this.sessions.events.append(params); } return event; } } +/** + * Returns a copy of `content` without Part fields the Agent Engine Sessions + * API rejects. + * + * `partMetadata` is a Gemini Developer API-only field; the Sessions API fails + * appendEvent with 400 INVALID_ARGUMENT ("Unknown name \"part_metadata\""). + */ +function dropUnsupportedPartFields(content: Content): Content { + if (!content.parts) { + return content; + } + return { + ...content, + parts: content.parts.map((part) => { + const copy = {...part}; + delete copy.partMetadata; + return copy; + }), + }; +} + +/** + * Maps ADK `EventActions` onto the Sessions API wire shape. ADK's + * `transferToAgent` is the API's `transferAgent` (adk-python writes the same + * field as `transfer_agent`); every other field keeps its name, including + * `requestedToolConfirmations`, which the SDK type omits but `_fromApiEvent` + * reads back. + */ +function toApiEventActions(actions: EventActions): ApiEventActions { + const {transferToAgent, ...rest} = actions; + return {...rest, transferAgent: transferToAgent}; +} + +/** + * True when the service rejected the request payload itself, which is what an + * API that does not know `rawEvent` returns. Any other failure must propagate: + * the event may already be persisted, so retrying would append it twice. + * + * Matched structurally on the `ApiError`'s `status`, for the reason given in + * getSession's catch. + */ +function isInvalidArgumentError(error: unknown): boolean { + return (error as {status?: number} | null)?.status === HTTP_BAD_REQUEST; +} + interface ExtendedEventActions extends EventActions { compaction?: { startTime: number; @@ -554,7 +602,12 @@ function _fromApiEvent(apiEventObj: VertexAiSessionEvent): Event { 'requestedToolConfirmations' ] as Record) || {}, skipSummarization: actions['skipSummarization'] as boolean | undefined, - transferToAgent: actions['transferAgent'] as string | undefined, + // Earlier adk-js versions copied `event.actions` onto the request + // verbatim, so sessions they wrote store ADK's own `transferToAgent` key. + transferToAgent: (actions['transferAgent'] ?? + (actions as Record)['transferToAgent']) as + | string + | undefined, escalate: actions['escalate'] as boolean | undefined, compaction: compactionData || undefined, }; @@ -574,6 +627,7 @@ function _fromApiEvent(apiEventObj: VertexAiSessionEvent): Event { turnComplete: eventMetadata['turnComplete'] as boolean | undefined, interrupted: eventMetadata['interrupted'] as boolean | undefined, branch: eventMetadata['branch'] as string | undefined, + groundingMetadata: eventMetadata.groundingMetadata, customMetadata, longRunningToolIds: eventMetadata['longRunningToolIds'] as | string[] diff --git a/core/test/sessions/vertex_ai_session_service_test.ts b/core/test/sessions/vertex_ai_session_service_test.ts index c0e834607..e33d08d0a 100644 --- a/core/test/sessions/vertex_ai_session_service_test.ts +++ b/core/test/sessions/vertex_ai_session_service_test.ts @@ -5,10 +5,22 @@ */ import {Sessions} from '@google-cloud/vertexai/build/src/genai/sessions.js'; -import {createEvent, State, VertexAiSessionService} from '@google/adk'; +import { + AppendAgentEngineSessionEventConfig, + AppendAgentEngineSessionEventRequestParameters, + SessionEvent as VertexAiSessionEvent, +} from '@google-cloud/vertexai/build/src/genai/types.js'; +import { + createEvent, + createEventActions, + createSession, + State, + ToolConfirmation, + VertexAiSessionService, +} from '@google/adk'; import {Session} from '@google/adk/sessions/session.js'; import {ApiError} from '@google/genai'; -import {beforeEach, describe, expect, it, vi} from 'vitest'; +import {beforeEach, describe, expect, it, Mock, vi} from 'vitest'; // Mock the unreleased nodejs-vertexai package so the import resolves vi.mock('nodejs-vertexai', () => ({ @@ -62,6 +74,10 @@ describe('quoteFilterLiteral', () => { }); }); +type AppendFn = ( + params: AppendAgentEngineSessionEventRequestParameters, +) => Promise; + describe('VertexAiSessionService', () => { let service: VertexAiSessionService; interface MockSessions { @@ -72,7 +88,7 @@ describe('VertexAiSessionService', () => { delete: ReturnType; events: { listInternal: ReturnType; - append: ReturnType; + append: Mock; }; } let mockClient: MockSessions; @@ -109,7 +125,7 @@ describe('VertexAiSessionService', () => { delete: vi.fn().mockResolvedValue({}), events: { listInternal: vi.fn().mockResolvedValue({sessionEvents: []}), - append: vi.fn().mockResolvedValue({}), + append: vi.fn().mockResolvedValue({}), }, }; @@ -118,6 +134,21 @@ describe('VertexAiSessionService', () => { }); }); + /** A fresh session to append events to. */ + function appendSession(): Session { + return createSession({ + id: 'append-session', + appName: '12345', + userId: 'testUser', + lastUpdateTime: Date.now(), + }); + } + + /** The request config captured by the first appendEvent call. */ + function appendedConfig(): AppendAgentEngineSessionEventConfig { + return mockClient.events.append.mock.calls[0][0].config!; + } + it('can initialize without passing a client explicitly', () => { const defaultService = new VertexAiSessionService({ projectId: 'test-project', @@ -1182,5 +1213,186 @@ describe('VertexAiSessionService', () => { }), ); }); + + it('sends transferToAgent as transferAgent and keeps every other action', async () => { + const authConfig = { + authScheme: {type: 'apiKey', name: 'x-api-key', in: 'header'}, + credentialKey: 'key1', + } as const; + const actions = createEventActions({ + stateDelta: {counter: 1}, + artifactDelta: {'report.pdf': 2}, + requestedAuthConfigs: {call1: authConfig}, + requestedToolConfirmations: { + call2: new ToolConfirmation({hint: 'confirm?', confirmed: true}), + }, + skipSummarization: true, + escalate: true, + transferToAgent: 'specialist', + }); + const event = createEvent({timestamp: 1620000000000, actions}); + + await service.appendEvent({session: appendSession(), event}); + + expect(appendedConfig().actions).toEqual({ + stateDelta: {counter: 1}, + artifactDelta: {'report.pdf': 2}, + requestedAuthConfigs: {call1: authConfig}, + requestedToolConfirmations: { + call2: new ToolConfirmation({hint: 'confirm?', confirmed: true}), + }, + skipSummarization: true, + escalate: true, + transferAgent: 'specialist', + }); + }); + + it('strips partMetadata from content and rawEvent without mutating the event', async () => { + const event = createEvent({ + timestamp: 1620000000000, + content: { + role: 'user', + parts: [ + {text: 'hello', partMetadata: {source: 'portal'}}, + {text: 'world', partMetadata: {source: 'portal'}}, + ], + }, + }); + + await service.appendEvent({session: appendSession(), event}); + + const config = appendedConfig(); + expect(config.content).toEqual({ + role: 'user', + parts: [{text: 'hello'}, {text: 'world'}], + }); + expect(config.rawEvent?.content).toEqual({ + role: 'user', + parts: [{text: 'hello'}, {text: 'world'}], + }); + expect(event.content?.parts?.[0].partMetadata).toEqual({ + source: 'portal', + }); + }); + + it('handles content without parts', async () => { + const event = createEvent({ + timestamp: 1620000000000, + content: {role: 'user'}, + }); + + await service.appendEvent({session: appendSession(), event}); + + expect(appendedConfig().content).toEqual({role: 'user'}); + }); + + it('handles an event without content', async () => { + const event = createEvent({timestamp: 1620000000000}); + delete event.content; + + await service.appendEvent({session: appendSession(), event}); + + const config = appendedConfig(); + expect(config.content).toBeUndefined(); + expect(config.rawEvent).not.toHaveProperty('content'); + }); + + it('retries without rawEvent when the API rejects it with 400', async () => { + const loggerSpy = vi.spyOn(logger, 'warn').mockImplementation(() => {}); + // The retry reuses the request object, so record what each attempt + // actually carried instead of inspecting it afterwards. + const sentRawEvent: boolean[] = []; + mockClient.events.append + .mockImplementationOnce(async (params) => { + sentRawEvent.push(params.config?.rawEvent !== undefined); + throw new ApiError({message: 'Unknown name', status: 400}); + }) + .mockImplementationOnce(async (params) => { + sentRawEvent.push(params.config?.rawEvent !== undefined); + return {}; + }); + const event = createEvent({ + timestamp: 1620000000000, + content: {role: 'user', parts: [{text: 'hello'}]}, + }); + + await service.appendEvent({session: appendSession(), event}); + + expect(sentRawEvent).toEqual([true, false]); + // Reusing the request is what keeps the retry's invocation id and + // timestamp identical to the first attempt's. + const [first, second] = mockClient.events.append.mock.calls; + expect(second[0]).toBe(first[0]); + loggerSpy.mockRestore(); + }); + + it.each([ + ['a server error', new ApiError({message: 'try later', status: 503})], + ['a network error', new Error('socket hang up')], + ])('rethrows %s without re-appending', async (_label, failure) => { + mockClient.events.append.mockRejectedValueOnce(failure); + const event = createEvent({ + timestamp: 1620000000000, + content: {role: 'user', parts: [{text: 'hello'}]}, + }); + + await expect( + service.appendEvent({session: appendSession(), event}), + ).rejects.toBe(failure); + expect(mockClient.events.append).toHaveBeenCalledTimes(1); + }); + }); + + describe('legacy read path', () => { + /** Reads back a single legacy (rawEvent-less) event via getSession. */ + async function readLegacyEvent(apiEvent: VertexAiSessionEvent) { + mockClient.events.listInternal.mockResolvedValue({ + sessionEvents: [apiEvent], + }); + const session = await service.getSession({ + appName: '12345', + userId: 'testUser', + sessionId: 'my-session-id', + }); + expect(session?.events).toHaveLength(1); + return session!.events[0]; + } + + it('restores groundingMetadata from eventMetadata', async () => { + const groundingMetadata = { + webSearchQueries: ['adk'], + groundingChunks: [{web: {uri: 'https://example.com', title: 't'}}], + }; + + const event = await readLegacyEvent({ + name: 'reasoningEngines/12345/sessions/s/events/e1', + author: 'model', + eventMetadata: {groundingMetadata}, + }); + + expect(event.groundingMetadata).toEqual(groundingMetadata); + }); + + it('restores transferToAgent from actions.transferAgent', async () => { + const event = await readLegacyEvent({ + name: 'reasoningEngines/12345/sessions/s/events/e1', + author: 'model', + actions: {transferAgent: 'specialist'}, + }); + + expect(event.actions.transferToAgent).toBe('specialist'); + }); + + it('restores transferToAgent from a legacy transferToAgent key', async () => { + const event = await readLegacyEvent({ + name: 'reasoningEngines/12345/sessions/s/events/e1', + author: 'model', + actions: { + transferToAgent: 'legacy-specialist', + } as VertexAiSessionEvent['actions'], + }); + + expect(event.actions.transferToAgent).toBe('legacy-specialist'); + }); }); });