Show citations at the end of each turn (#7350)

2026-07-12 11:00:51 -07:00 · 2025-08-28 16:42:54 -07:00
parent e02cb6c6a3
commit 95db24f234
8 changed files with 262 additions and 24 deletions
@@ -445,6 +445,170 @@ describe('Turn', () => {
      ]);
    });

+    it('should yield citation and finished events when response has citationMetadata', async () => {
+      const mockResponseStream = (async function* () {
+        yield {
+          candidates: [
+            {
+              content: { parts: [{ text: 'Some text.' }] },
+              citationMetadata: {
+                citations: [
+                  {
+                    uri: 'https://example.com/source1',
+                    title: 'Source 1 Title',
+                  },
+                ],
+              },
+              finishReason: 'STOP',
+            },
+          ],
+        } as unknown as GenerateContentResponse;
+      })();
+      mockSendMessageStream.mockResolvedValue(mockResponseStream);
+
+      const events = [];
+      for await (const event of turn.run(
+        [{ text: 'Test citations' }],
+        new AbortController().signal,
+      )) {
+        events.push(event);
+      }
+
+      expect(events).toEqual([
+        { type: GeminiEventType.Content, value: 'Some text.' },
+        {
+          type: GeminiEventType.Citation,
+          value: 'Citations:\n(Source 1 Title) https://example.com/source1',
+        },
+        { type: GeminiEventType.Finished, value: 'STOP' },
+      ]);
+    });
+
+    it('should yield a single citation event for multiple citations in one response', async () => {
+      const mockResponseStream = (async function* () {
+        yield {
+          candidates: [
+            {
+              content: { parts: [{ text: 'Some text.' }] },
+              citationMetadata: {
+                citations: [
+                  {
+                    uri: 'https://example.com/source2',
+                    title: 'Title2',
+                  },
+                  {
+                    uri: 'https://example.com/source1',
+                    title: 'Title1',
+                  },
+                ],
+              },
+              finishReason: 'STOP',
+            },
+          ],
+        } as unknown as GenerateContentResponse;
+      })();
+      mockSendMessageStream.mockResolvedValue(mockResponseStream);
+
+      const events = [];
+      for await (const event of turn.run(
+        [{ text: 'test' }],
+        new AbortController().signal,
+      )) {
+        events.push(event);
+      }
+
+      expect(events).toEqual([
+        { type: GeminiEventType.Content, value: 'Some text.' },
+        {
+          type: GeminiEventType.Citation,
+          value:
+            'Citations:\n(Title1) https://example.com/source1\n(Title2) https://example.com/source2',
+        },
+        { type: GeminiEventType.Finished, value: 'STOP' },
+      ]);
+    });
+
+    it('should not yield citation event if there is no finish reason', async () => {
+      const mockResponseStream = (async function* () {
+        yield {
+          candidates: [
+            {
+              content: { parts: [{ text: 'Some text.' }] },
+              citationMetadata: {
+                citations: [
+                  {
+                    uri: 'https://example.com/source1',
+                    title: 'Source 1 Title',
+                  },
+                ],
+              },
+              // No finishReason
+            },
+          ],
+        } as unknown as GenerateContentResponse;
+      })();
+      mockSendMessageStream.mockResolvedValue(mockResponseStream);
+
+      const events = [];
+      for await (const event of turn.run(
+        [{ text: 'test' }],
+        new AbortController().signal,
+      )) {
+        events.push(event);
+      }
+
+      expect(events).toEqual([
+        { type: GeminiEventType.Content, value: 'Some text.' },
+      ]);
+      // No Citation or Finished event
+      expect(events.some((e) => e.type === GeminiEventType.Citation)).toBe(
+        false,
+      );
+    });
+
+    it('should ignore citations without a URI', async () => {
+      const mockResponseStream = (async function* () {
+        yield {
+          candidates: [
+            {
+              content: { parts: [{ text: 'Some text.' }] },
+              citationMetadata: {
+                citations: [
+                  {
+                    uri: 'https://example.com/source1',
+                    title: 'Good Source',
+                  },
+                  {
+                    // uri is undefined
+                    title: 'Bad Source',
+                  },
+                ],
+              },
+              finishReason: 'STOP',
+            },
+          ],
+        } as unknown as GenerateContentResponse;
+      })();
+      mockSendMessageStream.mockResolvedValue(mockResponseStream);
+
+      const events = [];
+      for await (const event of turn.run(
+        [{ text: 'test' }],
+        new AbortController().signal,
+      )) {
+        events.push(event);
+      }
+
+      expect(events).toEqual([
+        { type: GeminiEventType.Content, value: 'Some text.' },
+        {
+          type: GeminiEventType.Citation,
+          value: 'Citations:\n(Good Source) https://example.com/source1',
+        },
+        { type: GeminiEventType.Finished, value: 'STOP' },
+      ]);
+    });
+
    it('should not crash when cancelled request has malformed error', async () => {
      const abortController = new AbortController();

@@ -54,6 +54,7 @@ export enum GeminiEventType {
  MaxSessionTurns = 'max_session_turns',
  Finished = 'finished',
  LoopDetected = 'loop_detected',
+  Citation = 'citation',
 }

 export interface StructuredError {
@@ -163,34 +164,37 @@ export type ServerGeminiLoopDetectedEvent = {
  type: GeminiEventType.LoopDetected;
 };

+export type ServerGeminiCitationEvent = {
+  type: GeminiEventType.Citation;
+  value: string;
+};
+
 // The original union type, now composed of the individual types
 export type ServerGeminiStreamEvent =
+  | ServerGeminiChatCompressedEvent
+  | ServerGeminiCitationEvent
  | ServerGeminiContentEvent
+  | ServerGeminiErrorEvent
+  | ServerGeminiFinishedEvent
+  | ServerGeminiLoopDetectedEvent
+  | ServerGeminiMaxSessionTurnsEvent
+  | ServerGeminiThoughtEvent
+  | ServerGeminiToolCallConfirmationEvent
  | ServerGeminiToolCallRequestEvent
  | ServerGeminiToolCallResponseEvent
-  | ServerGeminiToolCallConfirmationEvent
-  | ServerGeminiUserCancelledEvent
-  | ServerGeminiErrorEvent
-  | ServerGeminiChatCompressedEvent
-  | ServerGeminiThoughtEvent
-  | ServerGeminiMaxSessionTurnsEvent
-  | ServerGeminiFinishedEvent
-  | ServerGeminiLoopDetectedEvent;
+  | ServerGeminiUserCancelledEvent;

 // A turn manages the agentic loop turn within the server context.
 export class Turn {
-  readonly pendingToolCalls: ToolCallRequestInfo[];
-  private debugResponses: GenerateContentResponse[];
-  finishReason: FinishReason | undefined;
+  readonly pendingToolCalls: ToolCallRequestInfo[] = [];
+  private debugResponses: GenerateContentResponse[] = [];
+  private pendingCitations = new Set<string>();
+  finishReason: FinishReason | undefined = undefined;

  constructor(
    private readonly chat: GeminiChat,
    private readonly prompt_id: string,
-  ) {
-    this.pendingToolCalls = [];
-    this.debugResponses = [];
-    this.finishReason = undefined;
-  }
+  ) {}
  // The run method yields simpler events suitable for server logic
  async *run(
    req: PartListUnion,
@@ -251,10 +255,22 @@ export class Turn {
          }
        }

+        for (const citation of getCitations(resp)) {
+          this.pendingCitations.add(citation);
+        }
+
        // Check if response was truncated or stopped for various reasons
        const finishReason = resp.candidates?.[0]?.finishReason;

        if (finishReason) {
+          if (this.pendingCitations.size > 0) {
+            yield {
+              type: GeminiEventType.Citation,
+              value: `Citations:\n${[...this.pendingCitations].sort().join('\n')}`,
+            };
+            this.pendingCitations.clear();
+          }
+
          this.finishReason = finishReason;
          yield {
            type: GeminiEventType.Finished,
@@ -325,3 +341,14 @@ export class Turn {
    return this.debugResponses;
  }
 }
+
+function getCitations(resp: GenerateContentResponse): string[] {
+  return (resp.candidates?.[0]?.citationMetadata?.citations ?? [])
+    .filter((citation) => citation.uri !== undefined)
+    .map((citation) => {
+      if (citation.title) {
+        return `(${citation.title}) ${citation.uri}`;
+      }
+      return citation.uri!;
+    });
+}