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
This commit is contained in:
Taylor Mullen
2025-05-26 14:17:56 -07:00
committed by N. Taylor Mullen
parent 02503a3248
commit 480549e02e
7 changed files with 339 additions and 19 deletions

View File

@@ -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) {