From 909c35c2d50310fda336a1abec44bce75ad7e914 Mon Sep 17 00:00:00 2001 From: Mahima Shanware Date: Mon, 30 Mar 2026 21:56:59 +0000 Subject: [PATCH] fix(core): address PR feedback regarding hooks, unused promises, and UI double-render --- .../cli/src/ui/components/MainContent.tsx | 50 +++++++++---------- packages/core/src/core/geminiChat.ts | 26 ++++------ packages/core/src/hooks/hookEventHandler.ts | 4 ++ packages/core/src/hooks/hookSystem.ts | 9 +++- packages/core/src/hooks/types.ts | 2 + 5 files changed, 49 insertions(+), 42 deletions(-) diff --git a/packages/cli/src/ui/components/MainContent.tsx b/packages/cli/src/ui/components/MainContent.tsx index 6f0af57b7c..084949ff77 100644 --- a/packages/cli/src/ui/components/MainContent.tsx +++ b/packages/cli/src/ui/components/MainContent.tsx @@ -242,6 +242,28 @@ export const MainContent = () => { ], ); + const btwDisplayNode = useMemo( + () => + uiState.btwState.isActive ? ( + + ) : null, + [ + uiState.btwState.isActive, + uiState.btwState.query, + uiState.btwState.response, + uiState.btwState.isStreaming, + uiState.btwState.error, + uiState.terminalWidth, + ], + ); + const virtualizedData = useMemo(() => { const data: Array< | { type: 'header' } @@ -281,16 +303,7 @@ export const MainContent = () => { // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion return (item as any).element; } else if (item.type === 'btw') { - return ( - - ); + return <>{btwDisplayNode}; } else { return pendingItems; } @@ -299,11 +312,7 @@ export const MainContent = () => { showHeaderDetails, version, pendingItems, - uiState.btwState.query, - uiState.btwState.response, - uiState.btwState.isStreaming, - uiState.btwState.error, - uiState.terminalWidth, + btwDisplayNode, ], ); @@ -392,16 +401,7 @@ export const MainContent = () => { {(item) => item} {pendingItems} - {uiState.btwState.isActive && ( - - )} + {btwDisplayNode} ); }; diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index c736165230..ec19bdfa49 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -494,12 +494,6 @@ export class GeminiChat { signal: AbortSignal, role: LlmRole, ): AsyncGenerator { - let streamDoneResolver: () => void; - // We don't await this.sendPromise because /btw is safe to run concurrently. - void new Promise((resolve) => { - streamDoneResolver = resolve; - }); - const requestContents = [ ...this.getHistory(true), createUserContent(message), @@ -541,8 +535,6 @@ export class GeminiChat { } else { throw error; } - } finally { - streamDoneResolver!(); } }; @@ -636,12 +628,15 @@ export class GeminiChat { : [...requestContents]; const hookSystem = this.context.config.getHookSystem(); - if (hookSystem && !isBtw) { - const beforeModelResult = await hookSystem.fireBeforeModelEvent({ - model: modelToUse, - config, - contents: contentsToUse, - }); + if (hookSystem) { + const beforeModelResult = await hookSystem.fireBeforeModelEvent( + { + model: modelToUse, + config, + contents: contentsToUse, + }, + isBtw, + ); if (beforeModelResult.stopped) { throw new AgentExecutionStoppedError( @@ -987,10 +982,11 @@ export class GeminiChat { } const hookSystem = this.context.config.getHookSystem(); - if (originalRequest && chunk && hookSystem && !isBtw) { + if (originalRequest && chunk && hookSystem) { const hookResult = await hookSystem.fireAfterModelEvent( originalRequest, chunk, + isBtw, ); if (hookResult.stopped) { diff --git a/packages/core/src/hooks/hookEventHandler.ts b/packages/core/src/hooks/hookEventHandler.ts index 24dd77d76e..1c67352588 100644 --- a/packages/core/src/hooks/hookEventHandler.ts +++ b/packages/core/src/hooks/hookEventHandler.ts @@ -220,10 +220,12 @@ export class HookEventHandler { */ async fireBeforeModelEvent( llmRequest: GenerateContentParameters, + isBtw?: boolean, ): Promise { const input: BeforeModelInput = { ...this.createBaseInput(HookEventName.BeforeModel), llm_request: defaultHookTranslator.toHookLLMRequest(llmRequest), + isBtw, }; return this.executeHooks( @@ -241,11 +243,13 @@ export class HookEventHandler { async fireAfterModelEvent( llmRequest: GenerateContentParameters, llmResponse: GenerateContentResponse, + isBtw?: boolean, ): Promise { const input: AfterModelInput = { ...this.createBaseInput(HookEventName.AfterModel), llm_request: defaultHookTranslator.toHookLLMRequest(llmRequest), llm_response: defaultHookTranslator.toHookLLMResponse(llmResponse), + isBtw, }; return this.executeHooks( diff --git a/packages/core/src/hooks/hookSystem.ts b/packages/core/src/hooks/hookSystem.ts index 0b943256f3..e2a2f626b4 100644 --- a/packages/core/src/hooks/hookSystem.ts +++ b/packages/core/src/hooks/hookSystem.ts @@ -260,10 +260,13 @@ export class HookSystem { async fireBeforeModelEvent( llmRequest: GenerateContentParameters, + isBtw?: boolean, ): Promise { try { - const result = - await this.hookEventHandler.fireBeforeModelEvent(llmRequest); + const result = await this.hookEventHandler.fireBeforeModelEvent( + llmRequest, + isBtw, + ); const hookOutput = result.finalOutput; if (hookOutput?.shouldStopExecution()) { @@ -310,11 +313,13 @@ export class HookSystem { async fireAfterModelEvent( originalRequest: GenerateContentParameters, chunk: GenerateContentResponse, + isBtw?: boolean, ): Promise { try { const result = await this.hookEventHandler.fireAfterModelEvent( originalRequest, chunk, + isBtw, ); const hookOutput = result.finalOutput; diff --git a/packages/core/src/hooks/types.ts b/packages/core/src/hooks/types.ts index 418dcde03e..3dbe969562 100644 --- a/packages/core/src/hooks/types.ts +++ b/packages/core/src/hooks/types.ts @@ -673,6 +673,7 @@ export interface PreCompressOutput { */ export interface BeforeModelInput extends HookInput { llm_request: LLMRequest; + isBtw?: boolean; } /** @@ -692,6 +693,7 @@ export interface BeforeModelOutput extends HookOutput { export interface AfterModelInput extends HookInput { llm_request: LLMRequest; llm_response: LLMResponse; + isBtw?: boolean; } /**