Implementing support for recitations events in responses from A2A Server (#12067)

Co-authored-by: Alisa Novikova <alisanovikova@google.com>

packages/a2a-server/src/agent/task.test.ts

@@ -4,11 +4,24 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import {
+  describe,
+  it,
+  expect,
+  vi,
+  beforeEach,
+  afterEach,
+  type Mock,
+} from 'vitest';
 import { Task } from './task.js';
-import type { Config, ToolCallRequestInfo } from '@google/gemini-cli-core';
+import {
+  GeminiEventType,
+  type Config,
+  type ToolCallRequestInfo,
+} from '@google/gemini-cli-core';
 import { createMockConfig } from '../utils/testing_utils.js';
 import type { ExecutionEventBus } from '@a2a-js/sdk/server';
+import { CoderAgentEvent } from '../types.js';
 import type { ToolCall } from '@google/gemini-cli-core';
 
 describe('Task', () => {
@@ -94,6 +107,50 @@ describe('Task', () => {
         }),
       );
     });
+
+    it('should handle Citation event and publish to event bus', async () => {
+      const mockConfig = createMockConfig();
+      const mockEventBus: ExecutionEventBus = {
+        publish: vi.fn(),
+        on: vi.fn(),
+        off: vi.fn(),
+        once: vi.fn(),
+        removeAllListeners: vi.fn(),
+        finished: vi.fn(),
+      };
+
+      // @ts-expect-error - Calling private constructor for test purposes.
+      const task = new Task(
+        'task-id',
+        'context-id',
+        mockConfig as Config,
+        mockEventBus,
+      );
+
+      const citationText = 'Source: example.com';
+      const citationEvent = {
+        type: GeminiEventType.Citation,
+        value: citationText,
+      };
+
+      await task.acceptAgentMessage(citationEvent);
+
+      expect(mockEventBus.publish).toHaveBeenCalledOnce();
+      const publishedEvent = (mockEventBus.publish as Mock).mock.calls[0][0];
+
+      expect(publishedEvent.kind).toBe('status-update');
+      expect(publishedEvent.taskId).toBe('task-id');
+      expect(publishedEvent.metadata.coderAgent.kind).toBe(
+        CoderAgentEvent.CitationEvent,
+      );
+      expect(publishedEvent.status.message).toBeDefined();
+      expect(publishedEvent.status.message.parts).toEqual([
+        {
+          kind: 'text',
+          text: citationText,
+        },
+      ]);
+    });
   });
 
   describe('_schedulerToolCallsUpdate', () => {

packages/a2a-server/src/agent/task.ts

@@ -49,6 +49,7 @@ import type {
   TaskMetadata,
   Thought,
   ThoughtSummary,
+  Citation,
 } from '../types.js';
 import type { PartUnion, Part as genAiPart } from '@google/genai';
 
@@ -638,6 +639,10 @@ export class Task {
         logger.info('[Task] Sending agent thought...');
         this._sendThought(event.value, traceId);
         break;
+      case GeminiEventType.Citation:
+        logger.info('[Task] Received citation from LLM stream.');
+        this._sendCitation(event.value);
+        break;
       case GeminiEventType.ChatCompressed:
         break;
       case GeminiEventType.Finished:
@@ -979,4 +984,18 @@ export class Task {
       ),
     );
   }
+
+  _sendCitation(citation: string) {
+    if (!citation || citation.trim() === '') {
+      return;
+    }
+    logger.info('[Task] Sending citation to event bus.');
+    const message = this._createTextMessage(citation);
+    const citationEvent: Citation = {
+      kind: CoderAgentEvent.CitationEvent,
+    };
+    this.eventBus?.publish(
+      this._createStatusUpdateEvent(this.taskState, citationEvent, message),
+    );
+  }
 }

packages/a2a-server/src/types.ts

@@ -37,6 +37,10 @@ export enum CoderAgentEvent {
    * An event that contains a thought from the agent.
    */
   ThoughtEvent = 'thought',
+  /**
+   * An event that contains citation from the agent.
+   */
+  CitationEvent = 'citation',
 }
 
 export interface AgentSettings {
@@ -64,6 +68,10 @@ export interface Thought {
   kind: CoderAgentEvent.ThoughtEvent;
 }
 
+export interface Citation {
+  kind: CoderAgentEvent.CitationEvent;
+}
+
 export type ThoughtSummary = {
   subject: string;
   description: string;
@@ -80,7 +88,8 @@ export type CoderAgentMessage =
   | ToolCallUpdate
   | TextContent
   | StateChange
-  | Thought;
+  | Thought
+  | Citation;
 
 export interface TaskMetadata {
   id: string;