From 480549e02ed4ae01c7df2abbd98bb0eb5b23bdd5 Mon Sep 17 00:00:00 2001
From: Taylor Mullen <ntaylormullen@google.com>
Date: Mon, 26 May 2025 14:17:56 -0700
Subject: Refactor(chat): Introduce custom Chat class for future modifications

- Copied the `Chat` class from `@google/genai` into `packages/server/src/core/geminiChat.ts`.
- This change is in preparation for future modifications to the chat handling logic.
- Updated relevant files to use the new `GeminiChat` class.

Part of https://github.com/google-gemini/gemini-cli/issues/551
---
 packages/cli/src/ui/hooks/useGeminiStream.ts | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

(limited to 'packages/cli/src')

diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts
index 76d29189..f369d796 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.ts
+++ b/packages/cli/src/ui/hooks/useGeminiStream.ts
@@ -23,7 +23,7 @@ import {
   ToolResultDisplay,
   ToolCallRequestInfo,
 } from '@gemini-code/server';
-import { type Chat, type PartListUnion, type Part } from '@google/genai';
+import { type PartListUnion, type Part } from '@google/genai';
 import {
   StreamingState,
   ToolCallStatus,
@@ -39,6 +39,7 @@ import { useStateAndRef } from './useStateAndRef.js';
 import { UseHistoryManagerReturn } from './useHistoryManager.js';
 import { useLogger } from './useLogger.js';
 import { useToolScheduler, mapToDisplay } from './useToolScheduler.js';
+import { GeminiChat } from '@gemini-code/server/src/core/geminiChat.js';
 
 enum StreamProcessingStatus {
   Completed,
@@ -63,7 +64,7 @@ export const useGeminiStream = (
 ) => {
   const [initError, setInitError] = useState<string | null>(null);
   const abortControllerRef = useRef<AbortController | null>(null);
-  const chatSessionRef = useRef<Chat | null>(null);
+  const chatSessionRef = useRef<GeminiChat | null>(null);
   const geminiClientRef = useRef<GeminiClient | null>(null);
   const [isResponding, setIsResponding] = useState<boolean>(false);
   const [pendingHistoryItemRef, setPendingHistoryItem] =
@@ -235,7 +236,7 @@ export const useGeminiStream = (
 
   const ensureChatSession = useCallback(async (): Promise<{
     client: GeminiClient | null;
-    chat: Chat | null;
+    chat: GeminiChat | null;
   }> => {
     const currentClient = geminiClientRef.current;
     if (!currentClient) {
-- 
cgit v1.2.3