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 ? (
+        <BtwDisplay
+          key="btw-display"
+          query={uiState.btwState.query}
+          response={uiState.btwState.response}
+          isStreaming={uiState.btwState.isStreaming}
+          error={uiState.btwState.error}
+          terminalWidth={uiState.terminalWidth}
+        />
+      ) : 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 (
-          <BtwDisplay
-            key="btw-display"
-            query={uiState.btwState.query}
-            response={uiState.btwState.response}
-            isStreaming={uiState.btwState.isStreaming}
-            error={uiState.btwState.error}
-            terminalWidth={uiState.terminalWidth}
-          />
-        );
+        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}
       </Static>
       {pendingItems}
-      {uiState.btwState.isActive && (
-        <BtwDisplay
-          key="btw-display"
-          query={uiState.btwState.query}
-          response={uiState.btwState.response}
-          isStreaming={uiState.btwState.isStreaming}
-          error={uiState.btwState.error}
-          terminalWidth={uiState.terminalWidth}
-        />
-      )}
+      {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<StreamEvent> {
-    let streamDoneResolver: () => void;
-    // We don't await this.sendPromise because /btw is safe to run concurrently.
-    void new Promise<void>((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<AggregatedHookResult> {
     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<AggregatedHookResult> {
     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<BeforeModelHookResult> {
     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<AfterModelHookResult> {
     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;
 }
 
 /**