From cd7d91f35ebd241c4d35c7da73360998a3ff65cc Mon Sep 17 00:00:00 2001 From: Michael Bleigh Date: Mon, 30 Mar 2026 12:35:31 -0700 Subject: [PATCH] refactor(cli): remove config dependency from useAgentStream --- packages/cli/src/ui/AppContainer.tsx | 1 + .../cli/src/ui/hooks/useAgentStream.test.tsx | 8 +------ packages/cli/src/ui/hooks/useAgentStream.ts | 9 ++++--- .../src/agent/legacy-agent-session.test.ts | 18 +++++++------- .../core/src/agent/legacy-agent-session.ts | 24 ++++++++++++------- 5 files changed, 32 insertions(+), 28 deletions(-) diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 2d622470c1..7de80e130c 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -1108,6 +1108,7 @@ Logging in with Google... Restarting Gemini CLI to continue. addItem: historyManager.addItem, onCancelSubmit, isShellFocused: embeddedShellFocused, + logger, }) : // eslint-disable-next-line react-hooks/rules-of-hooks useGeminiStream( diff --git a/packages/cli/src/ui/hooks/useAgentStream.test.tsx b/packages/cli/src/ui/hooks/useAgentStream.test.tsx index c825ccaa17..53bb512504 100644 --- a/packages/cli/src/ui/hooks/useAgentStream.test.tsx +++ b/packages/cli/src/ui/hooks/useAgentStream.test.tsx @@ -17,12 +17,6 @@ const mockLegacyAgentProtocol = vi.hoisted(() => ({ abort: vi.fn().mockResolvedValue(undefined), })); -vi.mock('./useLogger.js', () => ({ - useLogger: vi.fn().mockReturnValue({ - logMessage: vi.fn().mockResolvedValue(undefined), - }), -})); - vi.mock('../contexts/SessionContext.js', async (importOriginal) => { const actual = await importOriginal>(); return { @@ -74,7 +68,7 @@ describe('useAgentStream', () => { }); expect(mockLegacyAgentProtocol.send).toHaveBeenCalledWith({ - message: [{ type: 'text', text: 'hello' }], + message: { content: [{ type: 'text', text: 'hello' }] }, }); expect(mockAddItem).toHaveBeenCalledWith( expect.objectContaining({ type: MessageType.USER, text: 'hello' }), diff --git a/packages/cli/src/ui/hooks/useAgentStream.ts b/packages/cli/src/ui/hooks/useAgentStream.ts index 44f70652b7..3600337169 100644 --- a/packages/cli/src/ui/hooks/useAgentStream.ts +++ b/packages/cli/src/ui/hooks/useAgentStream.ts @@ -20,6 +20,7 @@ import { type RetryAttemptPayload, type AgentEvent, type AgentProtocol, + type Logger, } from '@google/gemini-cli-core'; import type { HistoryItemWithoutId, @@ -32,10 +33,8 @@ import { findLastSafeSplitPoint } from '../utils/markdownUtilities.js'; import { getToolGroupBorderAppearance } from '../utils/borderStyles.js'; import { type BackgroundShell } from './shellCommandProcessor.js'; import type { UseHistoryManagerReturn } from './useHistoryManager.js'; -import { useLogger } from './useLogger.js'; import { useSessionStats } from '../contexts/SessionContext.js'; import { useStateAndRef } from './useStateAndRef.js'; -import { useConfig } from '../contexts/ConfigContext.js'; import { type MinimalTrackedToolCall } from './useTurnActivityMonitor.js'; export interface UseAgentStreamOptions { @@ -43,6 +42,7 @@ export interface UseAgentStreamOptions { addItem: UseHistoryManagerReturn['addItem']; onCancelSubmit: (shouldRestorePrompt?: boolean) => void; isShellFocused?: boolean; + logger?: Logger | null; } /** @@ -54,8 +54,8 @@ export const useAgentStream = ({ addItem, onCancelSubmit, isShellFocused, + logger, }: UseAgentStreamOptions) => { - const config = useConfig(); const [initError] = useState(null); const [retryStatus] = useState(null); const [streamingState, setStreamingState] = useState( @@ -78,7 +78,6 @@ export const useAgentStream = ({ useStateAndRef(true); const { startNewPrompt } = useSessionStats(); - const logger = useLogger(config?.storage); const activePtyId = undefined; const backgroundShellCount = 0; @@ -320,7 +319,7 @@ export const useAgentStream = ({ try { const { streamId } = await agent.send({ - message: parts, + message: { content: parts }, }); currentStreamIdRef.current = streamId; } catch (err) { diff --git a/packages/core/src/agent/legacy-agent-session.test.ts b/packages/core/src/agent/legacy-agent-session.test.ts index 93eb8acb2b..1de5d90e20 100644 --- a/packages/core/src/agent/legacy-agent-session.test.ts +++ b/packages/core/src/agent/legacy-agent-session.test.ts @@ -17,6 +17,9 @@ import type { ToolCallRequestInfo, } from '../scheduler/types.js'; import { CoreToolCallStatus } from '../scheduler/types.js'; +import type { GeminiClient } from '../core/client.js'; +import type { Scheduler } from '../scheduler/scheduler.js'; +import type { Config } from '../config/config.js'; // --------------------------------------------------------------------------- // Mock helpers @@ -24,7 +27,7 @@ import { CoreToolCallStatus } from '../scheduler/types.js'; function createMockDeps( overrides?: Partial, -): LegacyAgentSessionDeps { +): Required { const mockClient = { sendMessageStream: vi.fn(), getChat: vi.fn().mockReturnValue({ @@ -48,15 +51,14 @@ function createMockDeps( }; return { - client: mockClient as unknown as LegacyAgentSessionDeps['client'], - - scheduler: mockScheduler as unknown as LegacyAgentSessionDeps['scheduler'], - - config: mockConfig as unknown as LegacyAgentSessionDeps['config'], + client: mockClient as unknown as GeminiClient, + scheduler: mockScheduler as unknown as Scheduler, + config: mockConfig as unknown as Config, promptId: 'test-prompt', streamId: 'test-stream', + getPreferredEditor: vi.fn().mockReturnValue(undefined), ...overrides, - }; + } as Required; } async function* makeStream( @@ -134,7 +136,7 @@ async function collectEvents( // --------------------------------------------------------------------------- describe('LegacyAgentSession', () => { - let deps: LegacyAgentSessionDeps; + let deps: Required; beforeEach(() => { deps = createMockDeps(); diff --git a/packages/core/src/agent/legacy-agent-session.ts b/packages/core/src/agent/legacy-agent-session.ts index 667c85f5ed..88a5e897dd 100644 --- a/packages/core/src/agent/legacy-agent-session.ts +++ b/packages/core/src/agent/legacy-agent-session.ts @@ -14,10 +14,11 @@ import type { Part } from '@google/genai'; import type { GeminiClient } from '../core/client.js'; import type { Config } from '../config/config.js'; import type { ToolCallRequestInfo } from '../scheduler/types.js'; -import type { Scheduler } from '../scheduler/scheduler.js'; +import { Scheduler } from '../scheduler/scheduler.js'; import { recordToolCallInteractions } from '../code_assist/telemetry.js'; import { ToolErrorType, isFatalToolError } from '../tools/tool-error.js'; import { debugLogger } from '../utils/debugLogger.js'; +import type { EditorType } from '../utils/editor.js'; import { buildToolResponseData, contentPartsToGeminiParts, @@ -45,14 +46,15 @@ function isAbortLikeError(err: unknown): boolean { } export interface LegacyAgentSessionDeps { - client: GeminiClient; - scheduler: Scheduler; config: Config; - promptId: string; + client?: GeminiClient; + scheduler?: Scheduler; + promptId?: string; streamId?: string; + getPreferredEditor?: () => EditorType | undefined; } -class LegacyAgentProtocol implements AgentProtocol { +export class LegacyAgentProtocol implements AgentProtocol { private _events: AgentEvent[] = []; private _subscribers = new Set<(event: AgentEvent) => void>(); private _translationState: TranslationState; @@ -69,10 +71,16 @@ class LegacyAgentProtocol implements AgentProtocol { constructor(deps: LegacyAgentSessionDeps) { this._translationState = createTranslationState(deps.streamId); this._nextStreamIdOverride = deps.streamId; - this._client = deps.client; - this._scheduler = deps.scheduler; this._config = deps.config; - this._promptId = deps.promptId; + this._client = deps.client ?? deps.config.getGeminiClient(); + this._promptId = deps.promptId ?? deps.config.promptId ?? ''; + this._scheduler = + deps.scheduler ?? + new Scheduler({ + context: deps.config, + schedulerId: 'legacy-agent-scheduler', + getPreferredEditor: deps.getPreferredEditor ?? (() => undefined), + }); } get events(): readonly AgentEvent[] {