From e699238d83e654c1278bf03c73aeca604114021e Mon Sep 17 00:00:00 2001 From: Alex Yanchenko Date: Tue, 7 Jul 2026 23:34:20 +0200 Subject: [PATCH 01/10] feat(hooks): add MessageDisplay hook for mid-turn streaming Fires repeatedly as the assistant reply streams, before Stop (which only fires once at the end of the turn). Fire-and-forget, cumulative text payload, debounced (~200ms) except for the unconditional final firing. Fires from the single streaming loop in client.ts shared by the terminal UI and ACP paths. Fixes #6488 --- docs/users/features/hooks.md | 20 ++++ packages/acp-bridge/src/status.ts | 3 + packages/cli/src/config/settingsSchema.ts | 12 ++ .../cli/src/ui/components/hooks/constants.ts | 16 +++ packages/core/src/config/config.ts | 16 +++ packages/core/src/core/client.test.ts | 67 +++++++++++ packages/core/src/core/client.ts | 90 ++++++++++++++ .../src/core/message-display-buffer.test.ts | 110 ++++++++++++++++++ .../core/src/core/message-display-buffer.ts | 81 +++++++++++++ .../core/src/hooks/hookAggregator.test.ts | 28 +++++ .../core/src/hooks/hookEventHandler.test.ts | 75 ++++++++++++ packages/core/src/hooks/hookEventHandler.ts | 28 +++++ packages/core/src/hooks/hookPlanner.test.ts | 3 + packages/core/src/hooks/hookPlanner.ts | 1 + packages/core/src/hooks/hookSystem.test.ts | 52 +++++++++ packages/core/src/hooks/hookSystem.ts | 17 +++ packages/core/src/hooks/types.ts | 31 +++++ .../schemas/settings.schema.json | 103 ++++++++++++++++ 18 files changed, 753 insertions(+) create mode 100644 packages/core/src/core/message-display-buffer.test.ts create mode 100644 packages/core/src/core/message-display-buffer.ts diff --git a/docs/users/features/hooks.md b/docs/users/features/hooks.md index 8709b22d105..f7c7d74a005 100644 --- a/docs/users/features/hooks.md +++ b/docs/users/features/hooks.md @@ -243,6 +243,7 @@ Hooks fire at specific points during a Qwen Code session. Different events suppo | `UserPromptSubmit` | After user submits prompt | None (always fires) | | `SessionStart` | When session starts or resumes | Source (`startup`, `resume`, `clear`, `compact`) | | `SessionEnd` | When session ends | Reason (`clear`, `logout`, `prompt_input_exit`, etc.) | +| `MessageDisplay` | Repeatedly, as the reply streams | None (always fires) | | `Stop` | When Claude prepares to conclude response | None (always fires) | | `SubagentStart` | When subagent starts | Agent type (`Bash`, `Explorer`, `Plan`, etc.) | | `SubagentStop` | When subagent stops | Agent type | @@ -267,6 +268,7 @@ Hooks fire at specific points during a Qwen Code session. Different events suppo | Todo Events | `TodoCreated`, `TodoCompleted` | ❌ No | N/A | | Prompt Events | `UserPromptSubmit` | ❌ No | N/A | | Stop Events | `Stop` | ❌ No | N/A | +| Message Display | `MessageDisplay` | ❌ No | N/A | **Matcher Syntax:** @@ -563,6 +565,24 @@ The `permissionDecision` value controls whether the tool runs: - Standard hook output fields (typically not used for blocking) +#### MessageDisplay + +**Purpose**: Fires repeatedly as the assistant's reply streams — before `Stop`, which fires once at the end of the turn. Useful for live narration, incremental logging, or any consumer that wants to react to the reply as it's written rather than after the fact. This is a **fire-and-forget** event - hook output and exit codes are ignored. + +**Event-specific fields**: + +```json +{ + "message_id": "stable id for the whole streamed message", + "displayed_text": "the CUMULATIVE text streamed so far for this message (not a delta)", + "is_final": "true on the last firing for this message, false otherwise" +} +``` + +`displayed_text` is cumulative rather than a delta so hook scripts never need to reassemble chunks themselves — each firing carries the full text so far. Firing is debounced (at most every ~200ms) except for the final firing (`is_final: true`), which always fires immediately once the message ends, so the reply's tail is never dropped waiting on the debounce window. + +**Note**: Fires in both the terminal UI and ACP (IDE/editor) sessions — they share the same underlying streaming event loop. + #### Stop **Purpose**: Executed before Qwen concludes its response to provide final feedback or summaries. diff --git a/packages/acp-bridge/src/status.ts b/packages/acp-bridge/src/status.ts index 9cce58cd147..fc899caf6fe 100644 --- a/packages/acp-bridge/src/status.ts +++ b/packages/acp-bridge/src/status.ts @@ -936,6 +936,9 @@ export const IDLE_HOOK_EVENTS: Record = { description: 'When a new session is started', matcherKind: 'sessionTrigger', }, + MessageDisplay: { + description: 'Repeatedly, as the assistant reply streams', + }, Stop: { description: 'Right before Qwen Code concludes its response' }, SubagentStart: { description: 'When a subagent is started', diff --git a/packages/cli/src/config/settingsSchema.ts b/packages/cli/src/config/settingsSchema.ts index 1d0c971984f..e79ae739b50 100644 --- a/packages/cli/src/config/settingsSchema.ts +++ b/packages/cli/src/config/settingsSchema.ts @@ -2889,6 +2889,18 @@ const SETTINGS_SCHEMA = { mergeStrategy: MergeStrategy.CONCAT, items: HOOK_DEFINITION_ITEMS, }, + MessageDisplay: { + type: 'array', + label: 'Message Display Hooks', + category: 'Advanced', + requiresRestart: false, + default: [], + description: + 'Hooks that execute repeatedly as the assistant reply streams, before the After Agent (Stop) hooks fire.', + showInDialog: false, + mergeStrategy: MergeStrategy.CONCAT, + items: HOOK_DEFINITION_ITEMS, + }, Notification: { type: 'array', label: 'Notification Hooks', diff --git a/packages/cli/src/ui/components/hooks/constants.ts b/packages/cli/src/ui/components/hooks/constants.ts index 673d136c9d2..c37174769b5 100644 --- a/packages/cli/src/ui/components/hooks/constants.ts +++ b/packages/cli/src/ui/components/hooks/constants.ts @@ -48,6 +48,16 @@ export function getHookExitCodes(eventName: string): HookExitCode[] { { code: 2, description: t('show stderr to model immediately') }, { code: 'Other', description: t('show stderr to user only') }, ], + [HookEventName.MessageDisplay]: [ + { + code: 0, + description: t('fire-and-forget; exit status is ignored'), + }, + { + code: 'Other', + description: t('fire-and-forget; exit status is ignored'), + }, + ], [HookEventName.Notification]: [ { code: 0, description: t('stdout/stderr not shown') }, { code: 'Other', description: t('show stderr to user only') }, @@ -171,6 +181,9 @@ export function getHookShortDescription(eventName: string): string { 'When a slash command expands into a prompt', ), [HookEventName.SessionStart]: t('When a new session is started'), + [HookEventName.MessageDisplay]: t( + 'Repeatedly, as the assistant reply streams', + ), [HookEventName.Stop]: t('Right before Qwen Code concludes its response'), [HookEventName.SubagentStart]: t( 'When a subagent (Agent tool call) is started', @@ -229,6 +242,9 @@ export function getHookDescription(eventName: string): string { [HookEventName.SessionStart]: t( 'Input to command is JSON with session start source.', ), + [HookEventName.MessageDisplay]: t( + 'Input to command is JSON with message_id, displayed_text (cumulative text streamed so far), and is_final. Fire-and-forget: output and exit status are ignored.', + ), [HookEventName.SessionEnd]: t( 'Input to command is JSON with session end reason.', ), diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 7dafb838533..3702cc88b9f 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -2154,6 +2154,22 @@ export class Config { stopHookCount = stopResult.allOutputs.length; break; } + case 'MessageDisplay': { + const messageDisplayResult = + await hookSystem.fireMessageDisplayEvent( + (input['message_id'] as string) || '', + (input['displayed_text'] as string) || '', + (input['is_final'] as boolean) || false, + signal, + ); + result = messageDisplayResult.finalOutput + ? createHookOutput( + 'MessageDisplay', + messageDisplayResult.finalOutput, + ) + : undefined; + break; + } case 'PreToolUse': { result = await hookSystem.firePreToolUseEvent( (input['tool_name'] as string) || '', diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 574a1968d36..82c21544a12 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -7544,6 +7544,73 @@ Other open files: expect(mockMessageBus.request).not.toHaveBeenCalled(); }); + it('should skip messageBus.request for MessageDisplay when hasHooksForEvent returns false', async () => { + const mockMessageBus = { + request: vi.fn(), + response: vi.fn(), + }; + vi.mocked(mockConfig.getDisableAllHooks).mockReturnValue(false); + vi.mocked(mockConfig.getMessageBus).mockReturnValue( + mockMessageBus as unknown as ReturnType, + ); + vi.mocked(mockConfig.hasHooksForEvent).mockReturnValue(false); + + const stream = client.sendMessageStream( + [{ text: 'Hi' }], + new AbortController().signal, + 'prompt-hooks-message-display-off', + ); + for await (const _ of stream) { + // consume stream + } + + expect(mockMessageBus.request).not.toHaveBeenCalled(); + }); + + it('fires MessageDisplay with the cumulative streamed text, exactly once, when is_final on turn end', async () => { + const mockMessageBus = { + request: vi.fn().mockResolvedValue({}), + response: vi.fn(), + }; + vi.mocked(mockConfig.getDisableAllHooks).mockReturnValue(false); + vi.mocked(mockConfig.getMessageBus).mockReturnValue( + mockMessageBus as unknown as ReturnType, + ); + vi.mocked(mockConfig.hasHooksForEvent).mockImplementation( + (event: string) => event === 'MessageDisplay', + ); + mockTurnRunFn.mockReturnValue( + (async function* () { + yield { type: GeminiEventType.Content, value: 'Hello, ' }; + yield { type: GeminiEventType.Content, value: 'world.' }; + })(), + ); + + const stream = client.sendMessageStream( + [{ text: 'Hi' }], + new AbortController().signal, + 'prompt-message-display', + ); + for await (const _ of stream) { + // consume stream + } + + // A fast-running test never crosses the debounce window between the two + // Content chunks, so the only firing is the unconditional final flush — + // this also pins that mid-stream chunks don't each spawn their own call. + expect(mockMessageBus.request).toHaveBeenCalledTimes(1); + const [request] = mockMessageBus.request.mock.calls[0]; + expect(request).toMatchObject({ + eventName: 'MessageDisplay', + input: { + displayed_text: 'Hello, world.', + is_final: true, + }, + }); + expect(request.input.message_id).toEqual(expect.any(String)); + expect(request.input.message_id.length).toBeGreaterThan(0); + }); + it('ends the Stop hook loop when the blocking cap is reached', async () => { const mockMessageBus = { request: vi.fn().mockResolvedValue({ diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 4f8d90234a1..0aa44ca4410 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -135,6 +135,13 @@ import { import { partToString } from '../utils/partUtils.js'; import { createHookOutput, SessionStartSource } from '../hooks/types.js'; import fsPromises from 'node:fs/promises'; +import { randomUUID } from 'node:crypto'; +import { + createInitialMessageDisplayState, + stepMessageDisplay, + MESSAGE_DISPLAY_DEBOUNCE_MS, + type MessageDisplayState, +} from './message-display-buffer.js'; // IDE integration import { ideContextStore } from '../ide/ideContext.js'; @@ -1221,6 +1228,41 @@ export class GeminiClient { } } + /** + * Fire one MessageDisplay batch through MessageBus, WITHOUT being awaited by the + * caller — the streaming loop calls this and immediately continues to the next + * event, so a slow or hung hook command never stalls the reply's display. Errors + * are caught and logged here since there is no caller left to observe them. + */ + private fireMessageDisplayHook( + messageBus: ReturnType, + messageId: string, + displayedText: string, + isFinal: boolean, + signal: AbortSignal, + ): void { + if (!messageBus) { + return; + } + messageBus + .request( + { + type: MessageBusType.HOOK_EXECUTION_REQUEST, + eventName: 'MessageDisplay', + input: { + message_id: messageId, + displayed_text: displayedText, + is_final: isFinal, + }, + signal, + }, + MessageBusType.HOOK_EXECUTION_RESPONSE, + ) + .catch((err) => { + this.config.getDebugLogger().warn(`MessageDisplay hook failed: ${err}`); + }); + } + async startChat( extraHistory?: Content[], sessionStartSource = extraHistory @@ -2366,9 +2408,40 @@ export class GeminiClient { }; }; + // MessageDisplay hook: fires repeatedly as this turn's reply streams (before + // Stop, which fires once at the end). One id/buffer per turn.run() call — + // recursion into sendMessageStream (hook-forced continuations) naturally + // gets its own id since these locals are re-declared on each invocation. + const messageDisplayEnabled = + hooksEnabled && + !!messageBus && + this.config.hasHooksForEvent('MessageDisplay'); + const messageDisplayId = messageDisplayEnabled ? randomUUID() : ''; + let messageDisplayState: MessageDisplayState = + createInitialMessageDisplayState(Date.now()); + const resultStream = turn.run(model, requestToSend, signal); let didUpdateIdeContextState = false; for await (const event of resultStream) { + if (messageDisplayEnabled && event.type === GeminiEventType.Content) { + const step = stepMessageDisplay( + messageDisplayState, + event.value, + Date.now(), + MESSAGE_DISPLAY_DEBOUNCE_MS, + false, + ); + messageDisplayState = step.next; + if (step.flush) { + this.fireMessageDisplayHook( + messageBus, + messageDisplayId, + step.flush.displayedText, + step.flush.isFinal, + signal, + ); + } + } if (shouldUpdateIdeContextState && !didUpdateIdeContextState) { this.lastSentIdeContext = nextIdeContext; this.forceFullIdeContext = false; @@ -2499,6 +2572,23 @@ export class GeminiClient { } } + // Final MessageDisplay flush: this turn.run() stream is exhausted, so this + // message is done, regardless of whether pending tool calls will trigger a + // continuation (that continuation is its own message.run() call and gets its + // own message_id — see the const declarations above the loop). Unconditional, + // same as the debounced mid-stream flushes: not gated on !turn.pendingToolCalls + // the way the Stop hook below is, since a message boundary and a Stop-worthy + // end-of-turn are different things here. + if (messageDisplayEnabled) { + this.fireMessageDisplayHook( + messageBus, + messageDisplayId, + messageDisplayState.displayedText, + true, + signal, + ); + } + // Track API completion time for thinking block idle cleanup this.lastApiCompletionTimestamp = Date.now(); diff --git a/packages/core/src/core/message-display-buffer.test.ts b/packages/core/src/core/message-display-buffer.test.ts new file mode 100644 index 00000000000..21c1e7fc1f2 --- /dev/null +++ b/packages/core/src/core/message-display-buffer.test.ts @@ -0,0 +1,110 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect } from 'vitest'; +import { + createInitialMessageDisplayState, + stepMessageDisplay, + MESSAGE_DISPLAY_DEBOUNCE_MS, +} from './message-display-buffer.js'; + +describe('messageDisplayBuffer', () => { + describe('createInitialMessageDisplayState', () => { + it('starts empty with the given clock reading', () => { + const state = createInitialMessageDisplayState(1000); + expect(state).toEqual({ + displayedText: '', + lastFlushMs: 1000, + lastFlushedText: '', + }); + }); + }); + + describe('stepMessageDisplay', () => { + it('does not flush a chunk that arrives before the debounce window elapses', () => { + const initial = createInitialMessageDisplayState(0); + const step = stepMessageDisplay(initial, 'Hello', 50, 200, false); + expect(step.flush).toBeUndefined(); + expect(step.next.displayedText).toBe('Hello'); + }); + + it('flushes once the debounce window has elapsed and there is new text', () => { + const initial = createInitialMessageDisplayState(0); + const step = stepMessageDisplay(initial, 'Hello', 200, 200, false); + expect(step.flush).toEqual({ displayedText: 'Hello', isFinal: false }); + expect(step.next.lastFlushMs).toBe(200); + expect(step.next.lastFlushedText).toBe('Hello'); + }); + + it('accumulates chunks across non-flushing steps', () => { + let state = createInitialMessageDisplayState(0); + state = stepMessageDisplay(state, 'Hel', 10, 200, false).next; + state = stepMessageDisplay(state, 'lo', 20, 200, false).next; + expect(state.displayedText).toBe('Hello'); + }); + + it('does not flush when due by time but there is no new text since the last flush', () => { + let state = createInitialMessageDisplayState(0); + const first = stepMessageDisplay(state, 'Hello', 200, 200, false); + expect(first.flush).toBeDefined(); + state = first.next; + + // No new chunk arrived, but plenty of time has passed. + const second = stepMessageDisplay(state, '', 500, 200, false); + expect(second.flush).toBeUndefined(); + }); + + it('always flushes on isFinal, even with an empty chunk and within the debounce window', () => { + const initial = createInitialMessageDisplayState(0); + const step = stepMessageDisplay(initial, '', 5, 200, true); + expect(step.flush).toEqual({ displayedText: '', isFinal: true }); + }); + + it('isFinal flush carries the full cumulative text, including the final chunk', () => { + let state = createInitialMessageDisplayState(0); + state = stepMessageDisplay(state, 'Hello, ', 10, 200, false).next; + const final = stepMessageDisplay(state, 'world.', 15, 200, true); + expect(final.flush).toEqual({ + displayedText: 'Hello, world.', + isFinal: true, + }); + }); + + it('resets the debounce clock after each flush, independent of the caller-supplied window', () => { + let state = createInitialMessageDisplayState(0); + const first = stepMessageDisplay( + state, + 'a', + MESSAGE_DISPLAY_DEBOUNCE_MS, + MESSAGE_DISPLAY_DEBOUNCE_MS, + false, + ); + expect(first.flush).toBeDefined(); + state = first.next; + + // Immediately after the flush, a new chunk should NOT flush yet. + const tooSoon = stepMessageDisplay( + state, + 'b', + MESSAGE_DISPLAY_DEBOUNCE_MS + 1, + MESSAGE_DISPLAY_DEBOUNCE_MS, + false, + ); + expect(tooSoon.flush).toBeUndefined(); + + // Once the window elapses again, it flushes with everything accumulated. + state = tooSoon.next; + const later = stepMessageDisplay( + state, + 'c', + 2 * MESSAGE_DISPLAY_DEBOUNCE_MS + 1, + MESSAGE_DISPLAY_DEBOUNCE_MS, + false, + ); + expect(later.flush).toEqual({ displayedText: 'abc', isFinal: false }); + }); + }); +}); diff --git a/packages/core/src/core/message-display-buffer.ts b/packages/core/src/core/message-display-buffer.ts new file mode 100644 index 00000000000..354706a0eee --- /dev/null +++ b/packages/core/src/core/message-display-buffer.ts @@ -0,0 +1,81 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Debounce window for the MessageDisplay hook. Unlike Claude Code — whose + * MessageDisplay hook spawns a separate OS process per streamed batch and + * therefore needs cross-process locking to reassemble order — Qwen Code's + * streaming loop is single and sequential, so no such reassembly is needed. + * This constant just bounds how often a `command` hook process gets spawned + * per streamed reply. + */ +export const MESSAGE_DISPLAY_DEBOUNCE_MS = 200; + +/** + * Per-message accumulation state for the MessageDisplay hook, threaded + * through repeated calls to {@link stepMessageDisplay} as Content events + * arrive from the model. Kept as plain data (no clock/IO) so the whole + * decision is unit-testable without mocking timers. + */ +export interface MessageDisplayState { + /** Cumulative text streamed so far (all Content chunks appended in order). */ + displayedText: string; + /** Wall-clock time (ms) the last flush fired, or the state's creation time if none yet. */ + lastFlushMs: number; + /** The `displayedText` value as of the last flush, to detect "nothing new to say". */ + lastFlushedText: string; +} + +export function createInitialMessageDisplayState( + nowMs: number, +): MessageDisplayState { + return { displayedText: '', lastFlushMs: nowMs, lastFlushedText: '' }; +} + +/** What a batch produced: the updated state, plus a flush payload if one is due. */ +export interface MessageDisplayStep { + next: MessageDisplayState; + flush?: { displayedText: string; isFinal: boolean }; +} + +/** + * Decide what one streamed chunk does to the MessageDisplay accumulator, + * PURELY (no IO, no real timer) — the seam this feature's unit tests drive. + * + * A flush fires when either: + * - `isFinal` is true (the caller is closing out this message — always + * flushes, even with an empty `chunk`, so the reply's tail is never + * dropped waiting on the debounce window), or + * - there is new text since the last flush AND at least `debounceMs` has + * elapsed since then. + * Otherwise the chunk is folded into `displayedText` with no flush — the + * caller fires nothing this batch. + */ +export function stepMessageDisplay( + prev: MessageDisplayState, + chunk: string, + nowMs: number, + debounceMs: number, + isFinal: boolean, +): MessageDisplayStep { + const displayedText = prev.displayedText + chunk; + const hasNewText = displayedText !== prev.lastFlushedText; + const dueByTime = nowMs - prev.lastFlushMs >= debounceMs; + const shouldFlush = isFinal || (hasNewText && dueByTime); + + if (!shouldFlush) { + return { next: { ...prev, displayedText } }; + } + + return { + next: { + displayedText, + lastFlushMs: nowMs, + lastFlushedText: displayedText, + }, + flush: { displayedText, isFinal }, + }; +} diff --git a/packages/core/src/hooks/hookAggregator.test.ts b/packages/core/src/hooks/hookAggregator.test.ts index ece02f8c80b..33b7664d3e1 100644 --- a/packages/core/src/hooks/hookAggregator.test.ts +++ b/packages/core/src/hooks/hookAggregator.test.ts @@ -580,6 +580,34 @@ describe('HookAggregator', () => { result.finalOutput?.hookSpecificOutput?.['otherField'], ).toBeUndefined(); }); + + it('falls through to mergeSimple/DefaultHookOutput for MessageDisplay (no control-effect merge logic)', () => { + // MessageDisplay is fire-and-forget with no control effects, so it deliberately + // has no case in aggregateResults's switch and no case in createSpecificHookOutput + // — this pins that it lands in the same default path as Notification/PostCompact, + // not the OR-logic (mergeWithOrLogic) path used by control-affecting events. + const outputs: HookOutput[] = [ + { hookSpecificOutput: { additionalContext: 'a' } }, + { hookSpecificOutput: { additionalContext: 'b' } }, + ]; + const results: HookExecutionResult[] = outputs.map((output) => ({ + hookConfig: { type: HookType.Command, command: 'echo test' }, + eventName: HookEventName.MessageDisplay, + success: true, + output, + duration: 10, + })); + + const result = aggregator.aggregateResults( + results, + HookEventName.MessageDisplay, + ); + + expect( + result.finalOutput?.hookSpecificOutput?.['additionalContext'], + ).toBe('a\nb'); + expect(result.finalOutput?.constructor.name).toBe('DefaultHookOutput'); + }); }); describe('createSpecificHookOutput', () => { diff --git a/packages/core/src/hooks/hookEventHandler.test.ts b/packages/core/src/hooks/hookEventHandler.test.ts index 9f86b8ecb66..fa92bf15177 100644 --- a/packages/core/src/hooks/hookEventHandler.test.ts +++ b/packages/core/src/hooks/hookEventHandler.test.ts @@ -412,6 +412,81 @@ describe('HookEventHandler', () => { }); }); + describe('fireMessageDisplayEvent', () => { + it('should execute hooks for MessageDisplay event', async () => { + const mockPlan = createMockExecutionPlan([]); + const mockAggregated = createMockAggregatedResult(true); + + vi.mocked(mockHookPlanner.createExecutionPlan).mockReturnValue(mockPlan); + vi.mocked(mockHookRunner.executeHooksParallel).mockResolvedValue([]); + vi.mocked(mockHookAggregator.aggregateResults).mockReturnValue( + mockAggregated, + ); + + const result = await hookEventHandler.fireMessageDisplayEvent( + 'msg-1', + 'Hello', + false, + ); + + expect(mockHookPlanner.createExecutionPlan).toHaveBeenCalledWith( + HookEventName.MessageDisplay, + undefined, + ); + expect(result.success).toBe(true); + }); + + it('should include message_id, displayed_text, and is_final in hook input', async () => { + const mockPlan = createMockExecutionPlan([ + { + type: HookType.Command, + command: 'echo test', + source: HooksConfigSource.Project, + }, + ]); + vi.mocked(mockHookPlanner.createExecutionPlan).mockReturnValue(mockPlan); + vi.mocked(mockHookRunner.executeHooksParallel).mockResolvedValue([]); + vi.mocked(mockHookAggregator.aggregateResults).mockReturnValue( + createMockAggregatedResult(true), + ); + + await hookEventHandler.fireMessageDisplayEvent( + 'msg-42', + 'Hello, world.', + true, + ); + + const mockCalls = (mockHookRunner.executeHooksParallel as Mock).mock + .calls; + const input = mockCalls[0][2] as { + message_id: string; + displayed_text: string; + is_final: boolean; + }; + expect(input.message_id).toBe('msg-42'); + expect(input.displayed_text).toBe('Hello, world.'); + expect(input.is_final).toBe(true); + }); + + it('should handle missing finalOutput gracefully', async () => { + const mockPlan = createMockExecutionPlan([]); + vi.mocked(mockHookPlanner.createExecutionPlan).mockReturnValue(mockPlan); + vi.mocked(mockHookRunner.executeHooksParallel).mockResolvedValue([]); + vi.mocked(mockHookAggregator.aggregateResults).mockReturnValue( + createMockAggregatedResult(true, undefined), + ); + + const result = await hookEventHandler.fireMessageDisplayEvent( + 'msg-1', + '', + false, + ); + + expect(result.success).toBe(true); + expect(result.finalOutput).toBeUndefined(); + }); + }); + describe('fireSessionStartEvent', () => { it('should execute hooks for SessionStart event', async () => { const mockPlan = createMockExecutionPlan([]); diff --git a/packages/core/src/hooks/hookEventHandler.ts b/packages/core/src/hooks/hookEventHandler.ts index 17f7c097f4c..120cd0188f1 100644 --- a/packages/core/src/hooks/hookEventHandler.ts +++ b/packages/core/src/hooks/hookEventHandler.ts @@ -18,6 +18,7 @@ import type { UserPromptSubmitInput, UserPromptExpansionInput, StopInput, + MessageDisplayInput, ContextUsageData, SessionStartInput, SessionEndInput, @@ -202,6 +203,33 @@ export class HookEventHandler { return this.executeHooks(HookEventName.Stop, input, undefined, signal); } + /** + * Fire a MessageDisplay event + * Called repeatedly as the assistant's reply streams (before Stop). Fire-and-forget: + * callers should not await this on the critical streaming path — see client.ts, which + * fires it without blocking the next chunk's display. + */ + async fireMessageDisplayEvent( + messageId: string, + displayedText: string, + isFinal: boolean, + signal?: AbortSignal, + ): Promise { + const input: MessageDisplayInput = { + ...this.createBaseInput(HookEventName.MessageDisplay), + message_id: messageId, + displayed_text: displayedText, + is_final: isFinal, + }; + + return this.executeHooks( + HookEventName.MessageDisplay, + input, + undefined, + signal, + ); + } + /** * Fire a SessionStart event * Called when a new session starts or resumes diff --git a/packages/core/src/hooks/hookPlanner.test.ts b/packages/core/src/hooks/hookPlanner.test.ts index fbc5d2b1f35..8fccbdbb0d5 100644 --- a/packages/core/src/hooks/hookPlanner.test.ts +++ b/packages/core/src/hooks/hookPlanner.test.ts @@ -107,6 +107,9 @@ describe('HookPlanner', () => { undefined, ); expect(getHookMatcherTarget(HookEventName.PostToolBatch)).toBe(undefined); + expect(getHookMatcherTarget(HookEventName.MessageDisplay)).toBe( + undefined, + ); }); }); diff --git a/packages/core/src/hooks/hookPlanner.ts b/packages/core/src/hooks/hookPlanner.ts index eb5fa368fd6..66d8bb7a655 100644 --- a/packages/core/src/hooks/hookPlanner.ts +++ b/packages/core/src/hooks/hookPlanner.ts @@ -69,6 +69,7 @@ export function getHookMatcherTarget( case HookEventName.UserPromptSubmit: case HookEventName.Stop: + case HookEventName.MessageDisplay: case HookEventName.PostToolBatch: case HookEventName.TodoCreated: case HookEventName.TodoCompleted: diff --git a/packages/core/src/hooks/hookSystem.test.ts b/packages/core/src/hooks/hookSystem.test.ts index 0ccfb09e82c..ae27996d060 100644 --- a/packages/core/src/hooks/hookSystem.test.ts +++ b/packages/core/src/hooks/hookSystem.test.ts @@ -90,6 +90,7 @@ describe('HookSystem', () => { fireInstructionsLoadedEvent: vi.fn(), fireUserPromptExpansionEvent: vi.fn(), fireStopEvent: vi.fn(), + fireMessageDisplayEvent: vi.fn(), fireSessionStartEvent: vi.fn(), fireSessionEndEvent: vi.fn(), firePreToolUseEvent: vi.fn(), @@ -373,6 +374,57 @@ describe('HookSystem', () => { }); }); + describe('fireMessageDisplayEvent', () => { + it('should fire message display event and return AggregatedHookResult', async () => { + const mockResult = { + success: true, + allOutputs: [], + errors: [], + totalDuration: 5, + finalOutput: undefined, + }; + vi.mocked(mockHookEventHandler.fireMessageDisplayEvent).mockResolvedValue( + mockResult, + ); + + const result = await hookSystem.fireMessageDisplayEvent( + 'msg-1', + 'Hello', + false, + ); + + expect(mockHookEventHandler.fireMessageDisplayEvent).toHaveBeenCalledWith( + 'msg-1', + 'Hello', + false, + undefined, + ); + expect(result).toEqual(mockResult); + }); + + it('should return AggregatedHookResult even when no final output', async () => { + const mockResult = { + success: true, + allOutputs: [], + errors: [], + totalDuration: 0, + finalOutput: undefined, + }; + vi.mocked(mockHookEventHandler.fireMessageDisplayEvent).mockResolvedValue( + mockResult, + ); + + const result = await hookSystem.fireMessageDisplayEvent( + 'msg-1', + 'Hello, world.', + true, + ); + + expect(result).toEqual(mockResult); + expect(result.finalOutput).toBeUndefined(); + }); + }); + describe('fireUserPromptSubmitEvent', () => { it('should fire UserPromptSubmit event and return output', async () => { const mockResult = { diff --git a/packages/core/src/hooks/hookSystem.ts b/packages/core/src/hooks/hookSystem.ts index dbf4e5940c8..a7ab0c366c0 100644 --- a/packages/core/src/hooks/hookSystem.ts +++ b/packages/core/src/hooks/hookSystem.ts @@ -215,6 +215,23 @@ export class HookSystem { ); } + /** + * Fire a MessageDisplay event - called repeatedly as the assistant's reply streams + */ + async fireMessageDisplayEvent( + messageId: string, + displayedText: string, + isFinal: boolean, + signal?: AbortSignal, + ): Promise { + return this.hookEventHandler.fireMessageDisplayEvent( + messageId, + displayedText, + isFinal, + signal, + ); + } + async fireSessionStartEvent( source: SessionStartSource, model: string, diff --git a/packages/core/src/hooks/types.ts b/packages/core/src/hooks/types.ts index da76ed2e178..559b88944ff 100644 --- a/packages/core/src/hooks/types.ts +++ b/packages/core/src/hooks/types.ts @@ -39,6 +39,8 @@ export enum HookEventName { SessionStart = 'SessionStart', // Stop - Right before Claude concludes its response Stop = 'Stop', + // MessageDisplay - Fires repeatedly as the assistant's reply streams, before Stop + MessageDisplay = 'MessageDisplay', // SubagentStart - When a subagent (Task tool call) is started SubagentStart = 'SubagentStart', // SubagentStop - Right before a subagent (Task tool call) concludes its response @@ -941,6 +943,35 @@ export interface StopOutput extends HookOutput { }; } +/** + * MessageDisplay hook input + * + * Fires repeatedly as the assistant's reply streams (before `Stop`, which fires + * once at the end of the turn). `message_id` is stable for the whole streamed + * message; `displayed_text` is CUMULATIVE (the full text so far, not a delta), + * so hook authors never need to reassemble chunks themselves. `is_final` is + * true on the last firing for this message, so a hook script knows to flush + * (e.g. speak the tail of a buffered reply) rather than wait for more text + * that will never arrive. + */ +export interface MessageDisplayInput extends HookInput { + message_id: string; + displayed_text: string; + is_final: boolean; +} + +/** + * MessageDisplay hook output + * + * Fire-and-forget, no control effects (no blocking/permission semantics) — + * purely observational, like `Notification`/`PostCompact`. + */ +export interface MessageDisplayOutput extends HookOutput { + hookSpecificOutput?: { + hookEventName: 'MessageDisplay'; + }; +} + /** * SessionStart source types */ diff --git a/packages/vscode-ide-companion/schemas/settings.schema.json b/packages/vscode-ide-companion/schemas/settings.schema.json index 26c7ce8124b..4f2ab80f5af 100644 --- a/packages/vscode-ide-companion/schemas/settings.schema.json +++ b/packages/vscode-ide-companion/schemas/settings.schema.json @@ -1645,6 +1645,109 @@ ] } }, + "MessageDisplay": { + "description": "Hooks that execute repeatedly as the assistant reply streams, before the After Agent (Stop) hooks fire.", + "type": "array", + "items": { + "description": "A hook definition with an optional matcher and a list of hook configurations.", + "type": "object", + "properties": { + "matcher": { + "description": "An optional matcher pattern to filter when this hook definition applies.", + "type": "string" + }, + "sequential": { + "description": "Whether the hooks should be executed sequentially instead of in parallel.", + "type": "boolean" + }, + "hooks": { + "description": "The list of hook configurations to execute.", + "type": "array", + "items": { + "description": "A hook configuration entry that defines a hook to execute.", + "type": "object", + "properties": { + "type": { + "description": "The type of hook. Note: \"function\" type is only available via SDK registration, not settings.json.", + "type": "string", + "enum": [ + "command", + "http" + ] + }, + "command": { + "description": "The command to execute when the hook is triggered. Required for \"command\" type.", + "type": "string" + }, + "url": { + "description": "The URL to send the POST request to. Required for \"http\" type.", + "type": "string" + }, + "headers": { + "description": "HTTP headers to include in the request. Supports env var interpolation ($VAR, ${VAR}).", + "type": "object", + "additionalProperties": { + "type": "string" + } + }, + "allowedEnvVars": { + "description": "List of environment variables allowed for interpolation in headers and URL.", + "type": "array", + "items": { + "type": "string" + } + }, + "name": { + "description": "An optional name for the hook.", + "type": "string" + }, + "description": { + "description": "An optional description of what the hook does.", + "type": "string" + }, + "timeout": { + "description": "Timeout in seconds for the hook execution.", + "type": "number" + }, + "env": { + "description": "Environment variables to set when executing the hook command.", + "type": "object", + "additionalProperties": { + "type": "string" + } + }, + "async": { + "description": "Whether to execute the hook asynchronously (non-blocking, for \"command\" type only).", + "type": "boolean" + }, + "once": { + "description": "Whether to execute the hook only once per session (for \"http\" type).", + "type": "boolean" + }, + "statusMessage": { + "description": "A message to display while the hook is executing.", + "type": "string" + }, + "shell": { + "description": "The shell to use for command execution.", + "type": "string", + "enum": [ + "bash", + "powershell" + ] + } + }, + "required": [ + "type" + ] + } + } + }, + "required": [ + "hooks" + ] + } + }, "Notification": { "description": "Hooks that execute when notifications are sent.", "type": "array", From f5c76f4c8a9866a380f89506ad3a63a2cbbae663 Mon Sep 17 00:00:00 2001 From: Alex Yanchenko Date: Wed, 8 Jul 2026 02:53:11 +0200 Subject: [PATCH 02/10] fix(hooks): address MessageDisplay review feedback - Chain fire-and-forget MessageDisplay requests per message_id instead of firing them fully unbounded, so a slow hook command can't pile up concurrent processes. - Gate the final flush on non-empty displayed_text and !signal.aborted, matching the adjacent Stop hook's guard. - Document why the final flush intentionally re-sends the last debounced text (is_final itself is new information). - Simplify the debounce constant's JSDoc to drop the competitor comparison. - Add tests for the mid-stream debounced flush and the rejected-request warn path. --- packages/core/src/core/client.test.ts | 109 ++++++++++++++++++ packages/core/src/core/client.ts | 65 ++++++++--- .../core/src/core/message-display-buffer.ts | 8 +- 3 files changed, 159 insertions(+), 23 deletions(-) diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 82c21544a12..98eee881fee 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -23,6 +23,7 @@ import { tmpdir } from 'node:os'; import { join } from 'node:path'; import type { Content, GenerateContentResponse, Part } from '@google/genai'; import { GeminiClient, SendMessageType } from './client.js'; +import { MESSAGE_DISPLAY_DEBOUNCE_MS } from './message-display-buffer.js'; import { getRecentGitStatus } from '../utils/gitUtils.js'; import { AuthType, @@ -7611,6 +7612,114 @@ Other open files: expect(request.input.message_id.length).toBeGreaterThan(0); }); + it('fires a debounced mid-stream flush once the debounce window elapses, then a separate final flush', async () => { + vi.useFakeTimers(); + const mockMessageBus = { + request: vi.fn().mockResolvedValue({}), + response: vi.fn(), + }; + vi.mocked(mockConfig.getDisableAllHooks).mockReturnValue(false); + vi.mocked(mockConfig.getMessageBus).mockReturnValue( + mockMessageBus as unknown as ReturnType, + ); + vi.mocked(mockConfig.hasHooksForEvent).mockImplementation( + (event: string) => event === 'MessageDisplay', + ); + + let releaseSecondChunk!: () => void; + const secondChunkGate = new Promise((resolve) => { + releaseSecondChunk = resolve; + }); + mockTurnRunFn.mockReturnValue( + (async function* () { + yield { type: GeminiEventType.Content, value: 'Hello, ' }; + await secondChunkGate; + yield { type: GeminiEventType.Content, value: 'world.' }; + })(), + ); + + const stream = client.sendMessageStream( + [{ text: 'Hi' }], + new AbortController().signal, + 'prompt-message-display-debounced', + ); + const consumed = (async () => { + for await (const _ of stream) { + // consume stream + } + })(); + + // Let the first chunk get processed. It arrives in the same instant the + // debounce state was created, so it does not clear the debounce window + // by itself. + await vi.advanceTimersByTimeAsync(0); + expect(mockMessageBus.request).not.toHaveBeenCalled(); + + // Cross the debounce window, then let the second chunk arrive — this + // should fire a mid-stream flush (is_final: false) on its own, distinct + // from the unconditional final flush that fires once the stream ends. + await vi.advanceTimersByTimeAsync(MESSAGE_DISPLAY_DEBOUNCE_MS); + releaseSecondChunk(); + await consumed; + + expect(mockMessageBus.request).toHaveBeenCalledTimes(2); + const [midStreamCall, finalCall] = mockMessageBus.request.mock.calls; + expect(midStreamCall[0]).toMatchObject({ + eventName: 'MessageDisplay', + input: { displayed_text: 'Hello, world.', is_final: false }, + }); + expect(finalCall[0]).toMatchObject({ + eventName: 'MessageDisplay', + input: { displayed_text: 'Hello, world.', is_final: true }, + }); + }); + + it('logs and swallows a rejected MessageDisplay hook request', async () => { + const debugLogger = { + isEnabled: vi.fn().mockReturnValue(true), + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }; + const mockMessageBus = { + request: vi.fn().mockRejectedValue(new Error('hook process failed')), + response: vi.fn(), + }; + vi.mocked(mockConfig.getDisableAllHooks).mockReturnValue(false); + vi.mocked(mockConfig.getMessageBus).mockReturnValue( + mockMessageBus as unknown as ReturnType, + ); + vi.mocked(mockConfig.hasHooksForEvent).mockImplementation( + (event: string) => event === 'MessageDisplay', + ); + vi.mocked(mockConfig.getDebugLogger).mockReturnValue(debugLogger); + mockTurnRunFn.mockReturnValue( + (async function* () { + yield { type: GeminiEventType.Content, value: 'Hello, world.' }; + })(), + ); + + const stream = client.sendMessageStream( + [{ text: 'Hi' }], + new AbortController().signal, + 'prompt-message-display-rejected', + ); + for await (const _ of stream) { + // consume stream + } + // fireMessageDisplayHook is fire-and-forget; give its chained promise a + // tick to settle before asserting on the logger. + await Promise.resolve(); + await Promise.resolve(); + + expect(debugLogger.warn).toHaveBeenCalledWith( + expect.stringContaining( + 'MessageDisplay hook failed: Error: hook process failed', + ), + ); + }); + it('ends the Stop hook loop when the blocking cap is reached', async () => { const mockMessageBus = { request: vi.fn().mockResolvedValue({ diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 0aa44ca4410..36da7cf4763 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -229,6 +229,10 @@ export class GeminiClient { private cachedGitStatus: string | null | undefined; private readonly surfacedRelevantAutoMemoryPaths = new Set(); private shutdownRequested = false; + // Per-message-id chain of in-flight MessageDisplay hook requests, so a slow + // hook command doesn't let concurrent processes for the same message pile + // up unboundedly — see fireMessageDisplayHook. + private readonly messageDisplayChains = new Map>(); private readonly loopDetector: LoopDetectionService; private lastPromptId: string | undefined = undefined; @@ -1233,6 +1237,13 @@ export class GeminiClient { * caller — the streaming loop calls this and immediately continues to the next * event, so a slow or hung hook command never stalls the reply's display. Errors * are caught and logged here since there is no caller left to observe them. + * + * Batches for the same messageId are chained (not run concurrently): each waits + * for the previous batch's hook process to finish before its own request goes + * out. This bounds concurrent hook processes per message to one, so a slow or + * hung `command` hook can't let them pile up, while still preserving arrival + * order (important since `displayed_text` is cumulative — an out-of-order + * delivery would let an older, shorter payload land after a newer one). */ private fireMessageDisplayHook( messageBus: ReturnType, @@ -1244,23 +1255,33 @@ export class GeminiClient { if (!messageBus) { return; } - messageBus - .request( - { - type: MessageBusType.HOOK_EXECUTION_REQUEST, - eventName: 'MessageDisplay', - input: { - message_id: messageId, - displayed_text: displayedText, - is_final: isFinal, + const prior = this.messageDisplayChains.get(messageId) ?? Promise.resolve(); + const next = prior + .then(() => + messageBus.request( + { + type: MessageBusType.HOOK_EXECUTION_REQUEST, + eventName: 'MessageDisplay', + input: { + message_id: messageId, + displayed_text: displayedText, + is_final: isFinal, + }, + signal, }, - signal, - }, - MessageBusType.HOOK_EXECUTION_RESPONSE, + MessageBusType.HOOK_EXECUTION_RESPONSE, + ), ) + .then(() => undefined) .catch((err) => { this.config.getDebugLogger().warn(`MessageDisplay hook failed: ${err}`); + }) + .finally(() => { + if (this.messageDisplayChains.get(messageId) === next) { + this.messageDisplayChains.delete(messageId); + } }); + this.messageDisplayChains.set(messageId, next); } async startChat( @@ -2575,11 +2596,21 @@ export class GeminiClient { // Final MessageDisplay flush: this turn.run() stream is exhausted, so this // message is done, regardless of whether pending tool calls will trigger a // continuation (that continuation is its own message.run() call and gets its - // own message_id — see the const declarations above the loop). Unconditional, - // same as the debounced mid-stream flushes: not gated on !turn.pendingToolCalls - // the way the Stop hook below is, since a message boundary and a Stop-worthy - // end-of-turn are different things here. - if (messageDisplayEnabled) { + // own message_id — see the const declarations above the loop). Not gated on + // !turn.pendingToolCalls the way the Stop hook below is, since a message + // boundary and a Stop-worthy end-of-turn are different things here. It does + // re-send the same cumulative text as the last debounced flush when nothing + // changed since then — that's intentional, not a missed dedup: is_final is + // itself new information (it tells subscribers this message is done), so the + // event still needs to fire even when displayedText didn't change. Gated on + // there being any text at all so tool-call-only turns (no Content events ever + // arrived) don't fire a vacuous empty-text final event, and on !signal.aborted + // to match the Stop hook's cancellation guard below. + if ( + messageDisplayEnabled && + messageDisplayState.displayedText !== '' && + !signal.aborted + ) { this.fireMessageDisplayHook( messageBus, messageDisplayId, diff --git a/packages/core/src/core/message-display-buffer.ts b/packages/core/src/core/message-display-buffer.ts index 354706a0eee..2711fbe0c73 100644 --- a/packages/core/src/core/message-display-buffer.ts +++ b/packages/core/src/core/message-display-buffer.ts @@ -5,12 +5,8 @@ */ /** - * Debounce window for the MessageDisplay hook. Unlike Claude Code — whose - * MessageDisplay hook spawns a separate OS process per streamed batch and - * therefore needs cross-process locking to reassemble order — Qwen Code's - * streaming loop is single and sequential, so no such reassembly is needed. - * This constant just bounds how often a `command` hook process gets spawned - * per streamed reply. + * Debounce window for the MessageDisplay hook: bounds how often a `command` + * hook process gets spawned per streamed reply. */ export const MESSAGE_DISPLAY_DEBOUNCE_MS = 200; From 85e6f31b356711e3e455dc9918df536ee7145e78 Mon Sep 17 00:00:00 2001 From: Alex Yanchenko Date: Wed, 8 Jul 2026 02:56:05 +0200 Subject: [PATCH 03/10] test(hooks): drain microtasks before asserting on chained MessageDisplay calls MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit fireMessageDisplayHook now chains per-message_id through a promise (see previous commit), so the final flush's actual messageBus.request() call lands a few microtask ticks after the generator itself finishes — the mid-stream-flush test needs to let that chain settle before asserting. --- packages/core/src/core/client.test.ts | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 98eee881fee..177ffb096c7 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -7661,6 +7661,13 @@ Other open files: await vi.advanceTimersByTimeAsync(MESSAGE_DISPLAY_DEBOUNCE_MS); releaseSecondChunk(); await consumed; + // fireMessageDisplayHook chains requests for the same message_id through + // a promise so a slow hook can't run concurrently with itself — that + // chain settles a few microtask ticks after the generator itself + // finishes, since it's deliberately not awaited by the caller. + for (let i = 0; i < 10; i++) { + await Promise.resolve(); + } expect(mockMessageBus.request).toHaveBeenCalledTimes(2); const [midStreamCall, finalCall] = mockMessageBus.request.mock.calls; From dd355fa22b17de58a253b49555abde54b8e8666b Mon Sep 17 00:00:00 2001 From: Alex Yanchenko Date: Wed, 8 Jul 2026 04:10:43 +0200 Subject: [PATCH 04/10] fix(hooks): flush MessageDisplay is_final on every for-await exit path The three early `return turn` paths inside the streaming loop (always-on loop-detection safety, heuristic loop detection, and the stream Error event) exited before the final MessageDisplay flush, which only sat after the loop ended normally. Hook scripts relying on is_final: true to know when to flush never received it when a turn ended via loop detection or an API error. Extracts the flush into a shared closure and calls it from all four exits (the three early returns plus the normal fall-through), instead of only the one at the bottom of the loop. Adds regression tests for all three previously missed exits, plus the two guard-coverage tests requested in review (abort suppresses the flush, a tool-call-only turn with no Content events does not fire a vacuous empty-text event). Addresses the outstanding critical review comment and the follow-up test coverage suggestion on PR #6489. --- packages/core/src/core/client.test.ts | 204 ++++++++++++++++++++++++++ packages/core/src/core/client.ts | 59 ++++---- 2 files changed, 237 insertions(+), 26 deletions(-) diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 177ffb096c7..95642048c89 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -7727,6 +7727,210 @@ Other open files: ); }); + it('fires the final MessageDisplay flush when the always-on loop-detection safety trips mid-stream', async () => { + const mockMessageBus = { + request: vi.fn().mockResolvedValue({}), + response: vi.fn(), + }; + vi.mocked(mockConfig.getDisableAllHooks).mockReturnValue(false); + vi.mocked(mockConfig.getMessageBus).mockReturnValue( + mockMessageBus as unknown as ReturnType, + ); + vi.mocked(mockConfig.hasHooksForEvent).mockImplementation( + (event: string) => event === 'MessageDisplay', + ); + + const loopDetector = client['loopDetector']; + vi.spyOn(loopDetector, 'checkAlwaysOnSafeties').mockReturnValue(true); + vi.spyOn(loopDetector, 'getLastLoopType').mockReturnValue(null); + + mockTurnRunFn.mockReturnValue( + (async function* () { + yield { type: GeminiEventType.Content, value: 'Hello, world.' }; + })(), + ); + + const stream = client.sendMessageStream( + [{ text: 'trigger the always-on safety' }], + new AbortController().signal, + 'prompt-message-display-always-on-loop', + ); + for await (const _ of stream) { + // consume stream + } + + // Regression: this early `return turn` used to exit the method before the + // final-flush block that sat only after the `for await` loop, so hook + // scripts relying on `is_final: true` never saw the turn end. + const finalCall = mockMessageBus.request.mock.calls.find( + ([request]) => + request.eventName === 'MessageDisplay' && request.input?.is_final, + ); + expect(finalCall).toBeDefined(); + expect(finalCall![0].input.displayed_text).toBe('Hello, world.'); + }); + + it('fires the final MessageDisplay flush when heuristic loop detection trips mid-stream', async () => { + const mockMessageBus = { + request: vi.fn().mockResolvedValue({}), + response: vi.fn(), + }; + vi.mocked(mockConfig.getDisableAllHooks).mockReturnValue(false); + vi.mocked(mockConfig.getMessageBus).mockReturnValue( + mockMessageBus as unknown as ReturnType, + ); + vi.mocked(mockConfig.hasHooksForEvent).mockImplementation( + (event: string) => event === 'MessageDisplay', + ); + + const loopDetector = client['loopDetector']; + vi.spyOn(loopDetector, 'addAndCheckHeuristicLoops').mockReturnValue( + true, + ); + vi.spyOn(loopDetector, 'getLastLoopType').mockReturnValue(null); + + mockTurnRunFn.mockReturnValue( + (async function* () { + yield { type: GeminiEventType.Content, value: 'Hello, world.' }; + })(), + ); + + const stream = client.sendMessageStream( + [{ text: 'trigger a heuristic loop' }], + new AbortController().signal, + 'prompt-message-display-heuristic-loop', + ); + for await (const _ of stream) { + // consume stream + } + + const finalCall = mockMessageBus.request.mock.calls.find( + ([request]) => + request.eventName === 'MessageDisplay' && request.input?.is_final, + ); + expect(finalCall).toBeDefined(); + expect(finalCall![0].input.displayed_text).toBe('Hello, world.'); + }); + + it('fires the final MessageDisplay flush when the turn stream yields an Error event', async () => { + const mockMessageBus = { + request: vi.fn().mockResolvedValue({}), + response: vi.fn(), + }; + vi.mocked(mockConfig.getDisableAllHooks).mockReturnValue(false); + vi.mocked(mockConfig.getMessageBus).mockReturnValue( + mockMessageBus as unknown as ReturnType, + ); + vi.mocked(mockConfig.hasHooksForEvent).mockImplementation( + (event: string) => event === 'MessageDisplay', + ); + + mockTurnRunFn.mockReturnValue( + (async function* () { + yield { type: GeminiEventType.Content, value: 'Hello, world.' }; + yield { + type: GeminiEventType.Error, + value: { error: { message: 'test error' } }, + }; + })(), + ); + + const stream = client.sendMessageStream( + [{ text: 'Hi' }], + new AbortController().signal, + 'prompt-message-display-error', + ); + for await (const _ of stream) { + // consume stream + } + + const finalCall = mockMessageBus.request.mock.calls.find( + ([request]) => + request.eventName === 'MessageDisplay' && request.input?.is_final, + ); + expect(finalCall).toBeDefined(); + expect(finalCall![0].input.displayed_text).toBe('Hello, world.'); + }); + + it('suppresses the final MessageDisplay flush when the signal is aborted before the stream ends', async () => { + const mockMessageBus = { + request: vi.fn().mockResolvedValue({}), + response: vi.fn(), + }; + vi.mocked(mockConfig.getDisableAllHooks).mockReturnValue(false); + vi.mocked(mockConfig.getMessageBus).mockReturnValue( + mockMessageBus as unknown as ReturnType, + ); + vi.mocked(mockConfig.hasHooksForEvent).mockImplementation( + (event: string) => event === 'MessageDisplay', + ); + + const controller = new AbortController(); + mockTurnRunFn.mockReturnValue( + (async function* () { + yield { type: GeminiEventType.Content, value: 'Hello, world.' }; + controller.abort(); + })(), + ); + + const stream = client.sendMessageStream( + [{ text: 'Hi' }], + controller.signal, + 'prompt-message-display-aborted', + ); + for await (const _ of stream) { + // consume stream + } + + const messageDisplayCalls = mockMessageBus.request.mock.calls.filter( + ([request]) => request.eventName === 'MessageDisplay', + ); + expect(messageDisplayCalls).toHaveLength(0); + }); + + it('suppresses the final MessageDisplay flush for a tool-call-only turn with no Content events', async () => { + const mockMessageBus = { + request: vi.fn().mockResolvedValue({}), + response: vi.fn(), + }; + vi.mocked(mockConfig.getDisableAllHooks).mockReturnValue(false); + vi.mocked(mockConfig.getMessageBus).mockReturnValue( + mockMessageBus as unknown as ReturnType, + ); + vi.mocked(mockConfig.hasHooksForEvent).mockImplementation( + (event: string) => event === 'MessageDisplay', + ); + + mockTurnRunFn.mockReturnValue( + (async function* () { + yield { + type: GeminiEventType.ToolCallRequest, + value: { + callId: '1', + name: 'read_file', + args: {}, + isClientInitiated: false, + prompt_id: 'prompt-message-display-tool-only', + }, + }; + })(), + ); + + const stream = client.sendMessageStream( + [{ text: 'Hi' }], + new AbortController().signal, + 'prompt-message-display-tool-only', + ); + for await (const _ of stream) { + // consume stream + } + + const messageDisplayCalls = mockMessageBus.request.mock.calls.filter( + ([request]) => request.eventName === 'MessageDisplay', + ); + expect(messageDisplayCalls).toHaveLength(0); + }); + it('ends the Stop hook loop when the blocking cap is reached', async () => { const mockMessageBus = { request: vi.fn().mockResolvedValue({ diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 36da7cf4763..81385a12ecc 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -2440,6 +2440,35 @@ export class GeminiClient { const messageDisplayId = messageDisplayEnabled ? randomUUID() : ''; let messageDisplayState: MessageDisplayState = createInitialMessageDisplayState(Date.now()); + // Final MessageDisplay flush: called from every exit out of the `for await` + // loop below (normal completion and each early `return turn`), so a hook + // script's `is_final: true` completion signal is never skipped just because + // the turn ended via loop detection or a stream error instead of falling off + // the bottom of the loop. Not gated on !turn.pendingToolCalls the way the + // Stop hook below is, since a message boundary and a Stop-worthy end-of-turn + // are different things here. It does re-send the same cumulative text as the + // last debounced flush when nothing changed since then — that's intentional, + // not a missed dedup: is_final is itself new information (it tells + // subscribers this message is done), so the event still needs to fire even + // when displayedText didn't change. Gated on there being any text at all so + // tool-call-only turns (no Content events ever arrived) don't fire a vacuous + // empty-text final event, and on !signal.aborted to match the Stop hook's + // cancellation guard below. + const flushFinalMessageDisplay = (): void => { + if ( + messageDisplayEnabled && + messageDisplayState.displayedText !== '' && + !signal.aborted + ) { + this.fireMessageDisplayHook( + messageBus, + messageDisplayId, + messageDisplayState.displayedText, + true, + signal, + ); + } + }; const resultStream = turn.run(model, requestToSend, signal); let didUpdateIdeContextState = false; @@ -2494,6 +2523,7 @@ export class GeminiClient { if (isTopLevelInteraction) endInteractionSpan('error', { errorMessage: 'loop detected' }); this.cancelPendingMemoryPrefetch(); + flushFinalMessageDisplay(); return turn; } @@ -2525,6 +2555,7 @@ export class GeminiClient { // finally cleanup catches this, but cancel explicitly to match // the cleanup pattern at other early-return sites. this.cancelPendingMemoryPrefetch(); + flushFinalMessageDisplay(); return turn; } // Update arena status on Finished events — stats are derived @@ -2589,36 +2620,12 @@ export class GeminiClient { // finally cleanup catches this, but cancel explicitly to match // the cleanup pattern at other early-return sites. this.cancelPendingMemoryPrefetch(); + flushFinalMessageDisplay(); return turn; } } - // Final MessageDisplay flush: this turn.run() stream is exhausted, so this - // message is done, regardless of whether pending tool calls will trigger a - // continuation (that continuation is its own message.run() call and gets its - // own message_id — see the const declarations above the loop). Not gated on - // !turn.pendingToolCalls the way the Stop hook below is, since a message - // boundary and a Stop-worthy end-of-turn are different things here. It does - // re-send the same cumulative text as the last debounced flush when nothing - // changed since then — that's intentional, not a missed dedup: is_final is - // itself new information (it tells subscribers this message is done), so the - // event still needs to fire even when displayedText didn't change. Gated on - // there being any text at all so tool-call-only turns (no Content events ever - // arrived) don't fire a vacuous empty-text final event, and on !signal.aborted - // to match the Stop hook's cancellation guard below. - if ( - messageDisplayEnabled && - messageDisplayState.displayedText !== '' && - !signal.aborted - ) { - this.fireMessageDisplayHook( - messageBus, - messageDisplayId, - messageDisplayState.displayedText, - true, - signal, - ); - } + flushFinalMessageDisplay(); // Track API completion time for thinking block idle cleanup this.lastApiCompletionTimestamp = Date.now(); From 63f84541fd08a40bfbfee034550d21eb14b2e977 Mon Sep 17 00:00:00 2001 From: Alex Yanchenko Date: Wed, 8 Jul 2026 16:28:30 +0200 Subject: [PATCH 05/10] fix(hooks): fire MessageDisplay on the ACP surface, coalesce delivery, drain is_final before turn end MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses the three findings from the local verification report on #6489: - ACP/qwen serve (Finding 1): the delivery logic now lives in a shared MessageDisplayDispatcher (packages/core), and Session.ts wires it into all four raw-stream loops (main prompt, Stop-hook continuation, cron tick, background notification) — these surfaces consume GeminiChat's stream directly and never enter GeminiClient.sendMessageStream, so they need their own fire sites. The daemon no longer advertises an event it never emits. - Slow-hook backlog (Finding 2): the per-message promise chain is replaced by coalescing delivery — at most one in-flight request plus one pending payload per message; newer flushes overwrite the pending slot, which is lossless because displayed_text is cumulative, and is_final is sticky. A slow hook now sees fewer, newer payloads instead of an ever-growing queue of stale ones. - Headless is_final drop (Finding 3): finish() resolves only once every enqueued payload has actually been delivered, and every exit out of the streaming loops awaits it (early returns, normal fall-through, and the enclosing finally for uncaught exceptions), so a short-lived -p process can no longer exit with the final payload still queued. As a consequence, is_final delivery now strictly precedes the Stop hook rather than racing it. Also: the failure log line carries the message_id, finish() is idempotent, the review-requested tests are added (mid-stream and final firings share one message_id; isFinal as the sole flush reason), and hooks.md gains a delivery-semantics contract covering coalescing, the drain guarantee, no is_final on cancellation, provisional displayed_text, and multiple messages per tool-using turn. Co-Authored-By: Claude Fable 5 --- docs/users/features/hooks.md | 12 +- .../acp-integration/session/Session.test.ts | 83 ++++++ .../src/acp-integration/session/Session.ts | 203 ++++++++++----- packages/core/src/core/client.test.ts | 77 +++++- packages/core/src/core/client.ts | 171 ++++--------- .../src/core/message-display-buffer.test.ts | 12 + .../core/message-display-dispatcher.test.ts | 238 ++++++++++++++++++ .../src/core/message-display-dispatcher.ts | 166 ++++++++++++ packages/core/src/index.ts | 1 + 9 files changed, 755 insertions(+), 208 deletions(-) create mode 100644 packages/core/src/core/message-display-dispatcher.test.ts create mode 100644 packages/core/src/core/message-display-dispatcher.ts diff --git a/docs/users/features/hooks.md b/docs/users/features/hooks.md index 98e21b67f6d..f39e833df24 100644 --- a/docs/users/features/hooks.md +++ b/docs/users/features/hooks.md @@ -581,9 +581,17 @@ The `permissionDecision` value controls whether the tool runs: } ``` -`displayed_text` is cumulative rather than a delta so hook scripts never need to reassemble chunks themselves — each firing carries the full text so far. Firing is debounced (at most every ~200ms) except for the final firing (`is_final: true`), which always fires immediately once the message ends, so the reply's tail is never dropped waiting on the debounce window. +`displayed_text` is cumulative rather than a delta so hook scripts never need to reassemble chunks themselves — each firing carries the full text so far. Firing is debounced (at most every ~200ms) except for the final firing (`is_final: true`), which always fires once the message ends, so the reply's tail is never dropped waiting on the debounce window. -**Note**: Fires in both the terminal UI and ACP (IDE/editor) sessions — they share the same underlying streaming event loop. +**Delivery semantics** — what a hook script can rely on: + +- **Slow hooks see fewer, newer payloads.** At most one hook execution per message is in flight at a time; while one runs, newer debounced payloads _replace_ the queued one rather than piling up behind it. A hook slower than the debounce window therefore skips intermediate snapshots — lossless, since each payload carries the full cumulative text. The `is_final` payload itself is never dropped. +- **`is_final` delivery completes before the turn does.** The turn's end (and the `Stop` hook, when it fires) waits for the final `MessageDisplay` delivery, so a headless run (`qwen -p ...`) exits only after your hook received it, and a consumer combining `MessageDisplay` with `Stop` always sees `is_final` first. +- **A cancelled turn fires no `is_final`.** Cancellation (Esc, an aborted request) stops firings without a terminal event — the message didn't end, it was abandoned. A consumer that buffers until `is_final` should treat cancellation-silence as its flush/discard signal (e.g. a timeout fallback). +- **`displayed_text` is provisional until `is_final`.** It reflects what has streamed so far; treat intermediate payloads as display state, not as authoritative final content. +- **A tool-using turn produces multiple messages.** Each model call gets its own `message_id` with its own `is_final: true` firing: the text before a tool call is one message, the continuation after the tool result is another. Model calls that produce no displayed text (tool-call-only) fire nothing. + +**Note**: Fires in the terminal UI, headless (`-p`), and ACP (IDE/editor/`qwen serve`) sessions, with the same payload contract on every surface. #### Stop diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index fcaa63246b0..7b6b6c69736 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -1682,6 +1682,89 @@ describe('Session', () => { ).toHaveBeenCalledWith(latestSnapshot); }); + it('fires MessageDisplay with cumulative non-thought text and is_final on the ACP prompt path', async () => { + // Regression: the ACP surface consumes GeminiChat's stream directly + // (never entering GeminiClient.sendMessageStream), so it must fire the + // MessageDisplay hook itself — without this, an IDE/daemon client sees + // the hook advertised but never receives an event. + const messageBus = { request: vi.fn().mockResolvedValue({}) }; + mockConfig.getMessageBus = vi.fn().mockReturnValue(messageBus); + mockConfig.getDisableAllHooks = vi.fn().mockReturnValue(false); + mockConfig.hasHooksForEvent = vi + .fn() + .mockImplementation((event: string) => event === 'MessageDisplay'); + mockChat.sendMessageStream = vi.fn().mockResolvedValue( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + candidates: [ + { + content: { + parts: [ + { text: 'Let me think...', thought: true }, + { text: 'Hello, ' }, + ], + }, + }, + ], + }, + }, + { + type: core.StreamEventType.CHUNK, + value: { + candidates: [{ content: { parts: [{ text: 'world.' }] } }], + }, + }, + ]), + ); + + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'hi' }], + }); + + const messageDisplayCalls = messageBus.request.mock.calls.filter( + ([request]) => request.eventName === 'MessageDisplay', + ); + expect(messageDisplayCalls.length).toBeGreaterThan(0); + const finalCall = messageDisplayCalls[messageDisplayCalls.length - 1][0]; + // Cumulative text of the displayed (non-thought) parts only. + expect(finalCall.input).toMatchObject({ + displayed_text: 'Hello, world.', + is_final: true, + }); + expect(finalCall.input.message_id).toEqual(expect.any(String)); + // Exactly one is_final firing for the message. + expect( + messageDisplayCalls.filter(([request]) => request.input.is_final), + ).toHaveLength(1); + }); + + it('does not fire MessageDisplay on the ACP prompt path when the hook is not registered', async () => { + const messageBus = { request: vi.fn().mockResolvedValue({}) }; + mockConfig.getMessageBus = vi.fn().mockReturnValue(messageBus); + mockConfig.getDisableAllHooks = vi.fn().mockReturnValue(false); + mockConfig.hasHooksForEvent = vi.fn().mockReturnValue(false); + mockChat.sendMessageStream = vi.fn().mockResolvedValue( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + candidates: [{ content: { parts: [{ text: 'Hello.' }] } }], + }, + }, + ]), + ); + + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'hi' }], + }); + + expect(messageBus.request).not.toHaveBeenCalled(); + }); + it('drains background task notifications through ACP after the prompt is idle', async () => { mockChat.sendMessageStream = vi .fn() diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index 2260023992c..f4a77bbec5b 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -72,6 +72,7 @@ import { createHookOutput, generateToolUseId, MessageBusType, + MessageDisplayDispatcher, getPlanModeSystemReminder, getArenaSystemReminder, getStartupContextLength, @@ -1826,6 +1827,9 @@ export class Session implements SessionContext { let usageMetadata: GenerateContentResponseUsageMetadata | null = null; const streamStartTime = Date.now(); + const messageDisplay = this.#createMessageDisplayDispatcher( + pendingSend.signal, + ); try { const sendResult = @@ -1870,6 +1874,9 @@ export class Session implements SessionContext { 'assistant', part.thought, ); + if (!part.thought) { + messageDisplay?.addChunk(part.text); + } } } @@ -1953,6 +1960,11 @@ export class Session implements SessionContext { } throw error; + } finally { + // Deliver is_final (skipped on abort) and drain before the + // turn proceeds, on every exit: normal end-of-stream, + // cancellation returns, and thrown stream errors alike. + await messageDisplay?.finish(); } if (usageMetadata) { @@ -2155,6 +2167,9 @@ export class Session implements SessionContext { const functionCalls: FunctionCall[] = []; let usageMetadata: GenerateContentResponseUsageMetadata | null = null; const streamStartTime = Date.now(); + const messageDisplay = this.#createMessageDisplayDispatcher( + pendingSend.signal, + ); try { const continueSendResult = @@ -2192,6 +2207,9 @@ export class Session implements SessionContext { 'assistant', part.thought, ); + if (!part.thought) { + messageDisplay?.addChunk(part.text); + } } } @@ -2245,6 +2263,10 @@ export class Session implements SessionContext { } throw error; + } finally { + // Same contract as the main prompt loop: is_final (skipped on + // abort) is delivered and drained on every exit path. + await messageDisplay?.finish(); } if (usageMetadata) { @@ -2312,6 +2334,33 @@ export class Session implements SessionContext { * stop reason when the provider send should be skipped because the request * was cancelled or the session token limit was exceeded. */ + /** + * Create the MessageDisplay hook dispatcher for one model call's streamed + * reply, or null when the hook isn't registered (the common case — keeps + * the streaming loops zero-cost). The ACP surface consumes GeminiChat's + * raw stream directly rather than going through + * GeminiClient.sendMessageStream, so it has to fire this hook itself — + * with the same contract as the terminal UI path in client.ts: debounced + * cumulative text, one message_id per model call, and an is_final firing + * on every non-aborted exit (delivered by awaiting `finish()` in a + * finally around each streaming loop). + */ + #createMessageDisplayDispatcher( + signal: AbortSignal, + ): MessageDisplayDispatcher | null { + const messageBus = this.config.getMessageBus?.(); + if ( + this.config.getDisableAllHooks?.() || + !messageBus || + !this.config.hasHooksForEvent?.('MessageDisplay') + ) { + return null; + } + return new MessageDisplayDispatcher(messageBus, signal, (message) => + debugLogger.warn(message), + ); + } + async #sendMessageStreamWithAutoCompression( promptId: string, message: Part[], @@ -3095,42 +3144,54 @@ export class Session implements SessionContext { this.loopTickResolver?.markDelivered(); } nextMessage = null; + const messageDisplay = this.#createMessageDisplayDispatcher( + ac.signal, + ); - for await (const resp of responseStream) { - if (ac.signal.aborted) return; + try { + for await (const resp of responseStream) { + if (ac.signal.aborted) return; - if ( - resp.type === StreamEventType.CHUNK && - resp.value.candidates && - resp.value.candidates.length > 0 - ) { - const candidate = resp.value.candidates[0]; - for (const part of candidate.content?.parts ?? []) { - if (!part.text) continue; - this.messageEmitter.emitMessage( - part.text, - 'assistant', - part.thought, - ); + if ( + resp.type === StreamEventType.CHUNK && + resp.value.candidates && + resp.value.candidates.length > 0 + ) { + const candidate = resp.value.candidates[0]; + for (const part of candidate.content?.parts ?? []) { + if (!part.text) continue; + this.messageEmitter.emitMessage( + part.text, + 'assistant', + part.thought, + ); + if (!part.thought) { + messageDisplay?.addChunk(part.text); + } + } } - } - if ( - resp.type === StreamEventType.CHUNK && - resp.value.usageMetadata - ) { - usageMetadata = resp.value.usageMetadata; - } + if ( + resp.type === StreamEventType.CHUNK && + resp.value.usageMetadata + ) { + usageMetadata = resp.value.usageMetadata; + } - if ( - resp.type === StreamEventType.CHUNK && - resp.value.functionCalls - ) { - functionCalls.push(...resp.value.functionCalls); - } - if (resp.type === StreamEventType.MODEL_FALLBACK) { - functionCalls.length = 0; + if ( + resp.type === StreamEventType.CHUNK && + resp.value.functionCalls + ) { + functionCalls.push(...resp.value.functionCalls); + } + if (resp.type === StreamEventType.MODEL_FALLBACK) { + functionCalls.length = 0; + } } + } finally { + // is_final (skipped on abort) delivered and drained on + // every exit path, same as the interactive prompt loops. + await messageDisplay?.finish(); } if (usageMetadata) { @@ -3400,49 +3461,59 @@ export class Session implements SessionContext { const responseStream = sendResult.responseStream; nextMessage = null; + const messageDisplay = this.#createMessageDisplayDispatcher( + ac.signal, + ); - for await (const resp of responseStream) { - if (ac.signal.aborted) { - await this.#emitBackgroundNotificationEndTurn('cancelled'); - return; - } + try { + for await (const resp of responseStream) { + if (ac.signal.aborted) { + await this.#emitBackgroundNotificationEndTurn('cancelled'); + return; + } - if ( - resp.type === StreamEventType.CHUNK && - resp.value.candidates && - resp.value.candidates.length > 0 - ) { - const candidate = resp.value.candidates[0]; - for (const part of candidate.content?.parts ?? []) { - if (!part.text) continue; - if (part.thought) { - await this.messageEmitter.emitMessage( - part.text, - 'assistant', - true, - ); - } else { - responseText += part.text; + if ( + resp.type === StreamEventType.CHUNK && + resp.value.candidates && + resp.value.candidates.length > 0 + ) { + const candidate = resp.value.candidates[0]; + for (const part of candidate.content?.parts ?? []) { + if (!part.text) continue; + if (part.thought) { + await this.messageEmitter.emitMessage( + part.text, + 'assistant', + true, + ); + } else { + responseText += part.text; + messageDisplay?.addChunk(part.text); + } } } - } - if ( - resp.type === StreamEventType.CHUNK && - resp.value.usageMetadata - ) { - usageMetadata = resp.value.usageMetadata; - } + if ( + resp.type === StreamEventType.CHUNK && + resp.value.usageMetadata + ) { + usageMetadata = resp.value.usageMetadata; + } - if ( - resp.type === StreamEventType.CHUNK && - resp.value.functionCalls - ) { - functionCalls.push(...resp.value.functionCalls); - } - if (resp.type === StreamEventType.MODEL_FALLBACK) { - functionCalls.length = 0; + if ( + resp.type === StreamEventType.CHUNK && + resp.value.functionCalls + ) { + functionCalls.push(...resp.value.functionCalls); + } + if (resp.type === StreamEventType.MODEL_FALLBACK) { + functionCalls.length = 0; + } } + } finally { + // is_final (skipped on abort) delivered and drained on every + // exit path, same as the interactive prompt loops. + await messageDisplay?.finish(); } if (responseText.length > 0) { diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 95642048c89..e383d341297 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -7661,13 +7661,6 @@ Other open files: await vi.advanceTimersByTimeAsync(MESSAGE_DISPLAY_DEBOUNCE_MS); releaseSecondChunk(); await consumed; - // fireMessageDisplayHook chains requests for the same message_id through - // a promise so a slow hook can't run concurrently with itself — that - // chain settles a few microtask ticks after the generator itself - // finishes, since it's deliberately not awaited by the caller. - for (let i = 0; i < 10; i++) { - await Promise.resolve(); - } expect(mockMessageBus.request).toHaveBeenCalledTimes(2); const [midStreamCall, finalCall] = mockMessageBus.request.mock.calls; @@ -7679,6 +7672,10 @@ Other open files: eventName: 'MessageDisplay', input: { displayed_text: 'Hello, world.', is_final: true }, }); + // Both firings belong to the same streamed message. + expect(finalCall[0].input.message_id).toBe( + midStreamCall[0].input.message_id, + ); }); it('logs and swallows a rejected MessageDisplay hook request', async () => { @@ -7715,18 +7712,72 @@ Other open files: for await (const _ of stream) { // consume stream } - // fireMessageDisplayHook is fire-and-forget; give its chained promise a - // tick to settle before asserting on the logger. - await Promise.resolve(); - await Promise.resolve(); + // The log line carries the message_id so a failure can be correlated + // to its turn when debug logging is enabled. expect(debugLogger.warn).toHaveBeenCalledWith( - expect.stringContaining( - 'MessageDisplay hook failed: Error: hook process failed', + expect.stringMatching( + /^MessageDisplay hook failed \[[0-9a-f-]{36}\]: Error: hook process failed$/, ), ); }); + it('does not end the turn until the final MessageDisplay payload has been delivered', async () => { + const mockMessageBus = { + request: vi.fn(), + response: vi.fn(), + }; + vi.mocked(mockConfig.getDisableAllHooks).mockReturnValue(false); + vi.mocked(mockConfig.getMessageBus).mockReturnValue( + mockMessageBus as unknown as ReturnType, + ); + vi.mocked(mockConfig.hasHooksForEvent).mockImplementation( + (event: string) => event === 'MessageDisplay', + ); + + // A slow hook: the final MessageDisplay request stays unresolved + // until the test releases it. + let releaseHook!: () => void; + mockMessageBus.request.mockImplementation( + () => + new Promise((resolve) => { + releaseHook = () => resolve({}); + }), + ); + mockTurnRunFn.mockReturnValue( + (async function* () { + yield { type: GeminiEventType.Content, value: 'Hello, world.' }; + })(), + ); + + const stream = client.sendMessageStream( + [{ text: 'Hi' }], + new AbortController().signal, + 'prompt-message-display-drain', + ); + let turnEnded = false; + const consumed = (async () => { + for await (const _ of stream) { + // consume stream + } + turnEnded = true; + })(); + + // Give the generator ample time to run to its end if it (wrongly) + // didn't wait for the hook delivery. + for (let i = 0; i < 20; i++) { + await Promise.resolve(); + } + expect(mockMessageBus.request).toHaveBeenCalledTimes(1); + // Regression: in a short-lived process (headless -p), returning here + // would drop the queued is_final payload on process exit. + expect(turnEnded).toBe(false); + + releaseHook(); + await consumed; + expect(turnEnded).toBe(true); + }); + it('fires the final MessageDisplay flush when the always-on loop-detection safety trips mid-stream', async () => { const mockMessageBus = { request: vi.fn().mockResolvedValue({}), diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 81385a12ecc..5fd9a81c3cc 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -135,13 +135,7 @@ import { import { partToString } from '../utils/partUtils.js'; import { createHookOutput, SessionStartSource } from '../hooks/types.js'; import fsPromises from 'node:fs/promises'; -import { randomUUID } from 'node:crypto'; -import { - createInitialMessageDisplayState, - stepMessageDisplay, - MESSAGE_DISPLAY_DEBOUNCE_MS, - type MessageDisplayState, -} from './message-display-buffer.js'; +import { MessageDisplayDispatcher } from './message-display-dispatcher.js'; // IDE integration import { ideContextStore } from '../ide/ideContext.js'; @@ -229,10 +223,6 @@ export class GeminiClient { private cachedGitStatus: string | null | undefined; private readonly surfacedRelevantAutoMemoryPaths = new Set(); private shutdownRequested = false; - // Per-message-id chain of in-flight MessageDisplay hook requests, so a slow - // hook command doesn't let concurrent processes for the same message pile - // up unboundedly — see fireMessageDisplayHook. - private readonly messageDisplayChains = new Map>(); private readonly loopDetector: LoopDetectionService; private lastPromptId: string | undefined = undefined; @@ -1232,58 +1222,6 @@ export class GeminiClient { } } - /** - * Fire one MessageDisplay batch through MessageBus, WITHOUT being awaited by the - * caller — the streaming loop calls this and immediately continues to the next - * event, so a slow or hung hook command never stalls the reply's display. Errors - * are caught and logged here since there is no caller left to observe them. - * - * Batches for the same messageId are chained (not run concurrently): each waits - * for the previous batch's hook process to finish before its own request goes - * out. This bounds concurrent hook processes per message to one, so a slow or - * hung `command` hook can't let them pile up, while still preserving arrival - * order (important since `displayed_text` is cumulative — an out-of-order - * delivery would let an older, shorter payload land after a newer one). - */ - private fireMessageDisplayHook( - messageBus: ReturnType, - messageId: string, - displayedText: string, - isFinal: boolean, - signal: AbortSignal, - ): void { - if (!messageBus) { - return; - } - const prior = this.messageDisplayChains.get(messageId) ?? Promise.resolve(); - const next = prior - .then(() => - messageBus.request( - { - type: MessageBusType.HOOK_EXECUTION_REQUEST, - eventName: 'MessageDisplay', - input: { - message_id: messageId, - displayed_text: displayedText, - is_final: isFinal, - }, - signal, - }, - MessageBusType.HOOK_EXECUTION_RESPONSE, - ), - ) - .then(() => undefined) - .catch((err) => { - this.config.getDebugLogger().warn(`MessageDisplay hook failed: ${err}`); - }) - .finally(() => { - if (this.messageDisplayChains.get(messageId) === next) { - this.messageDisplayChains.delete(messageId); - } - }); - this.messageDisplayChains.set(messageId, next); - } - async startChat( extraHistory?: Content[], sessionStartSource = extraHistory @@ -2042,6 +1980,10 @@ export class GeminiClient { // early-return) leaves this `false`, and the `finally` block aborts the // prefetch as a safety net. let normalCompletion = false; + // Declared outside the try so the finally block can close it out on + // uncaught-exception exits too; created (when the hook is registered) + // right before the turn's streaming loop below. + let messageDisplay: MessageDisplayDispatcher | null = null; try { if ( messageType === SendMessageType.UserQuery || @@ -2429,68 +2371,33 @@ export class GeminiClient { }; }; - // MessageDisplay hook: fires repeatedly as this turn's reply streams (before - // Stop, which fires once at the end). One id/buffer per turn.run() call — - // recursion into sendMessageStream (hook-forced continuations) naturally - // gets its own id since these locals are re-declared on each invocation. - const messageDisplayEnabled = + // MessageDisplay hook: fires repeatedly as this turn's reply streams + // (before Stop, which fires once at the end). One dispatcher — one + // message_id and one debounce accumulator — per turn.run() call; + // recursion into sendMessageStream (tool continuations, hook-forced + // continuations) naturally gets its own since this local is re-created + // on each invocation. `finish()` is awaited at every exit out of the + // `for await` loop below (normal completion and each early `return + // turn`) plus the outer finally, so a hook script's `is_final: true` + // completion signal is neither skipped when the turn ends via loop + // detection or a stream error, nor silently dropped by a process that + // exits (headless `-p`) before a slow hook's queue drained. Not gated + // on !turn.pendingToolCalls the way the Stop hook below is, since a + // message boundary and a Stop-worthy end-of-turn are different things. + messageDisplay = hooksEnabled && - !!messageBus && - this.config.hasHooksForEvent('MessageDisplay'); - const messageDisplayId = messageDisplayEnabled ? randomUUID() : ''; - let messageDisplayState: MessageDisplayState = - createInitialMessageDisplayState(Date.now()); - // Final MessageDisplay flush: called from every exit out of the `for await` - // loop below (normal completion and each early `return turn`), so a hook - // script's `is_final: true` completion signal is never skipped just because - // the turn ended via loop detection or a stream error instead of falling off - // the bottom of the loop. Not gated on !turn.pendingToolCalls the way the - // Stop hook below is, since a message boundary and a Stop-worthy end-of-turn - // are different things here. It does re-send the same cumulative text as the - // last debounced flush when nothing changed since then — that's intentional, - // not a missed dedup: is_final is itself new information (it tells - // subscribers this message is done), so the event still needs to fire even - // when displayedText didn't change. Gated on there being any text at all so - // tool-call-only turns (no Content events ever arrived) don't fire a vacuous - // empty-text final event, and on !signal.aborted to match the Stop hook's - // cancellation guard below. - const flushFinalMessageDisplay = (): void => { - if ( - messageDisplayEnabled && - messageDisplayState.displayedText !== '' && - !signal.aborted - ) { - this.fireMessageDisplayHook( - messageBus, - messageDisplayId, - messageDisplayState.displayedText, - true, - signal, - ); - } - }; + messageBus && + this.config.hasHooksForEvent('MessageDisplay') + ? new MessageDisplayDispatcher(messageBus, signal, (message) => + this.config.getDebugLogger().warn(message), + ) + : null; const resultStream = turn.run(model, requestToSend, signal); let didUpdateIdeContextState = false; for await (const event of resultStream) { - if (messageDisplayEnabled && event.type === GeminiEventType.Content) { - const step = stepMessageDisplay( - messageDisplayState, - event.value, - Date.now(), - MESSAGE_DISPLAY_DEBOUNCE_MS, - false, - ); - messageDisplayState = step.next; - if (step.flush) { - this.fireMessageDisplayHook( - messageBus, - messageDisplayId, - step.flush.displayedText, - step.flush.isFinal, - signal, - ); - } + if (messageDisplay && event.type === GeminiEventType.Content) { + messageDisplay.addChunk(event.value); } if (shouldUpdateIdeContextState && !didUpdateIdeContextState) { this.lastSentIdeContext = nextIdeContext; @@ -2523,7 +2430,7 @@ export class GeminiClient { if (isTopLevelInteraction) endInteractionSpan('error', { errorMessage: 'loop detected' }); this.cancelPendingMemoryPrefetch(); - flushFinalMessageDisplay(); + await messageDisplay?.finish(); return turn; } @@ -2555,7 +2462,7 @@ export class GeminiClient { // finally cleanup catches this, but cancel explicitly to match // the cleanup pattern at other early-return sites. this.cancelPendingMemoryPrefetch(); - flushFinalMessageDisplay(); + await messageDisplay?.finish(); return turn; } // Update arena status on Finished events — stats are derived @@ -2620,12 +2527,16 @@ export class GeminiClient { // finally cleanup catches this, but cancel explicitly to match // the cleanup pattern at other early-return sites. this.cancelPendingMemoryPrefetch(); - flushFinalMessageDisplay(); + await messageDisplay?.finish(); return turn; } } - flushFinalMessageDisplay(); + // Deliver `is_final: true` (and drain any still-queued mid-stream + // payload) BEFORE the Stop hook below fires, so consumers combining the + // two events can rely on receiving MessageDisplay's completion signal + // first. + await messageDisplay?.finish(); // Track API completion time for thinking block idle cleanup this.lastApiCompletionTimestamp = Date.now(); @@ -2904,9 +2815,15 @@ export class GeminiClient { return turn; } finally { restoreStrippedRetryEntries(); - // Belt-and-suspenders: abort the prefetch on any exit other than the - // bottom-of-try `return turn`. Catches uncaught exceptions and guards - // against future early-return sites that forget to call cancel. + // Belt-and-suspenders: close out the MessageDisplay dispatcher on any + // exit the explicit finish() sites above didn't cover (an uncaught + // exception thrown out of the streaming loop still ends the message, + // and buffering hook consumers need the is_final signal). finish() is + // idempotent, so on the normal paths this resolves immediately. + await messageDisplay?.finish(); + // Abort the prefetch on any exit other than the bottom-of-try + // `return turn`. Catches uncaught exceptions and guards against + // future early-return sites that forget to call cancel. if (!normalCompletion) { this.cancelPendingMemoryPrefetch(); } diff --git a/packages/core/src/core/message-display-buffer.test.ts b/packages/core/src/core/message-display-buffer.test.ts index 21c1e7fc1f2..2cd0adfc666 100644 --- a/packages/core/src/core/message-display-buffer.test.ts +++ b/packages/core/src/core/message-display-buffer.test.ts @@ -63,6 +63,18 @@ describe('messageDisplayBuffer', () => { expect(step.flush).toEqual({ displayedText: '', isFinal: true }); }); + it('flushes on isFinal even when the text is unchanged since the last flush and the window has not elapsed', () => { + let state = createInitialMessageDisplayState(0); + const first = stepMessageDisplay(state, 'Hello', 200, 200, false); + expect(first.flush).toEqual({ displayedText: 'Hello', isFinal: false }); + state = first.next; + + // Nothing new to say AND still inside the debounce window: isFinal must + // be the sole reason this flushes. + const final = stepMessageDisplay(state, '', 250, 200, true); + expect(final.flush).toEqual({ displayedText: 'Hello', isFinal: true }); + }); + it('isFinal flush carries the full cumulative text, including the final chunk', () => { let state = createInitialMessageDisplayState(0); state = stepMessageDisplay(state, 'Hello, ', 10, 200, false).next; diff --git a/packages/core/src/core/message-display-dispatcher.test.ts b/packages/core/src/core/message-display-dispatcher.test.ts new file mode 100644 index 00000000000..c261e90fdf5 --- /dev/null +++ b/packages/core/src/core/message-display-dispatcher.test.ts @@ -0,0 +1,238 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi } from 'vitest'; +import { MessageDisplayDispatcher } from './message-display-dispatcher.js'; +import { MESSAGE_DISPLAY_DEBOUNCE_MS } from './message-display-buffer.js'; +import type { MessageBus } from '../confirmation-bus/message-bus.js'; + +interface SentPayload { + message_id: string; + displayed_text: string; + is_final: boolean; +} + +/** + * A MessageBus stub whose `request` resolves only when the test releases it, + * so tests can hold a hook execution "in flight" while more flushes arrive. + */ +function createControlledBus() { + const sent: SentPayload[] = []; + const releases: Array<() => void> = []; + const request = vi.fn( + (message: { input: SentPayload }) => + new Promise((resolve) => { + sent.push(message.input); + releases.push(() => resolve({})); + }), + ); + return { + bus: { request } as unknown as MessageBus, + request, + sent, + /** Settle the oldest unresolved request. */ + release: async () => { + releases.shift()?.(); + // Let the dispatcher's .then/.finally continuations run. + await Promise.resolve(); + await Promise.resolve(); + await Promise.resolve(); + }, + }; +} + +function createDispatcher( + bus: MessageBus, + opts: { signal?: AbortSignal; warn?: (message: string) => void } = {}, +) { + return new MessageDisplayDispatcher( + bus, + opts.signal ?? new AbortController().signal, + opts.warn ?? (() => {}), + 0, + ); +} + +const PAST_DEBOUNCE = MESSAGE_DISPLAY_DEBOUNCE_MS + 1; + +describe('MessageDisplayDispatcher', () => { + it('delivers a due mid-stream flush and then the final flush, sharing one message_id', async () => { + const { bus, sent, release } = createControlledBus(); + const dispatcher = createDispatcher(bus); + + dispatcher.addChunk('Hello, ', PAST_DEBOUNCE); + await release(); + const finished = dispatcher.finish(); + await release(); + await finished; + + expect(sent).toHaveLength(2); + expect(sent[0]).toMatchObject({ + displayed_text: 'Hello, ', + is_final: false, + }); + expect(sent[1]).toMatchObject({ + displayed_text: 'Hello, ', + is_final: true, + }); + expect(sent[1].message_id).toBe(sent[0].message_id); + expect(sent[0].message_id).toBe(dispatcher.messageId); + }); + + it('coalesces flushes that arrive while a hook is in flight, keeping only the newest', async () => { + const { bus, sent, release } = createControlledBus(); + const dispatcher = createDispatcher(bus); + + // First due flush goes out and is held in flight. + dispatcher.addChunk('one ', PAST_DEBOUNCE); + expect(sent).toHaveLength(1); + + // Three more due flushes arrive while it's still running: each should + // overwrite the single pending slot, not queue up behind one another. + dispatcher.addChunk('two ', 2 * PAST_DEBOUNCE); + dispatcher.addChunk('three ', 3 * PAST_DEBOUNCE); + dispatcher.addChunk('four', 4 * PAST_DEBOUNCE); + expect(sent).toHaveLength(1); + + await release(); // first request settles -> pending (newest only) goes out + expect(sent).toHaveLength(2); + expect(sent[1]).toMatchObject({ + displayed_text: 'one two three four', + is_final: false, + }); + + await release(); + const finished = dispatcher.finish(); + await release(); + await finished; + + // Intermediate texts "one two " and "one two three " were superseded and + // never delivered — lossless, since displayed_text is cumulative. + expect(sent).toHaveLength(3); + expect(sent[2]).toMatchObject({ + displayed_text: 'one two three four', + is_final: true, + }); + }); + + it('lets the final flush overwrite a pending mid-stream payload, keeping is_final', async () => { + const { bus, sent, release } = createControlledBus(); + const dispatcher = createDispatcher(bus); + + dispatcher.addChunk('partial ', PAST_DEBOUNCE); // in flight + dispatcher.addChunk('more ', 2 * PAST_DEBOUNCE); // pending + const finished = dispatcher.finish(); // overwrites pending, is_final wins + await release(); + await release(); + await finished; + + expect(sent).toHaveLength(2); + expect(sent[1]).toMatchObject({ + displayed_text: 'partial more ', + is_final: true, + }); + }); + + it('finish() resolves only once every enqueued payload has been delivered', async () => { + const { bus, sent, release } = createControlledBus(); + const dispatcher = createDispatcher(bus); + + dispatcher.addChunk('text', PAST_DEBOUNCE); // in flight + let finishResolved = false; + const finished = dispatcher.finish().then(() => { + finishResolved = true; + }); + + await Promise.resolve(); + await Promise.resolve(); + expect(finishResolved).toBe(false); // mid-stream payload still in flight + + await release(); // mid-stream delivered -> final goes out + expect(finishResolved).toBe(false); // final still in flight + expect(sent).toHaveLength(2); + + await release(); + await finished; + expect(finishResolved).toBe(true); + }); + + it('finish() is idempotent — a second call neither re-fires is_final nor hangs', async () => { + const { bus, sent, release } = createControlledBus(); + const dispatcher = createDispatcher(bus); + + dispatcher.addChunk('text', 0); // within debounce window: no mid-stream flush + const first = dispatcher.finish(); + await release(); + await first; + await dispatcher.finish(); + + const finals = sent.filter((payload) => payload.is_final); + expect(finals).toHaveLength(1); + expect(sent).toHaveLength(1); + expect(sent[0]).toMatchObject({ displayed_text: 'text', is_final: true }); + }); + + it('fires nothing on finish() when no text ever streamed', async () => { + const { bus, request } = createControlledBus(); + const dispatcher = createDispatcher(bus); + + await dispatcher.finish(); + + expect(request).not.toHaveBeenCalled(); + }); + + it('suppresses the final flush and does not wait on in-flight delivery when aborted', async () => { + const { bus, sent } = createControlledBus(); + const controller = new AbortController(); + const dispatcher = createDispatcher(bus, { signal: controller.signal }); + + dispatcher.addChunk('text', PAST_DEBOUNCE); // in flight, never released + controller.abort(); + + // Resolves immediately despite the unsettled in-flight request. + await dispatcher.finish(); + + expect(sent).toHaveLength(1); + expect(sent.filter((payload) => payload.is_final)).toHaveLength(0); + }); + + it('logs a failed delivery with the message_id and still delivers the final flush', async () => { + const warn = vi.fn(); + const sent: SentPayload[] = []; + const request = vi.fn((message: { input: SentPayload }) => { + sent.push(message.input); + return message.input.is_final + ? Promise.resolve({}) + : Promise.reject(new Error('hook process failed')); + }); + const dispatcher = createDispatcher({ request } as unknown as MessageBus, { + warn, + }); + + dispatcher.addChunk('text', PAST_DEBOUNCE); // this delivery fails + await dispatcher.finish(); + + expect(warn).toHaveBeenCalledWith( + `MessageDisplay hook failed [${dispatcher.messageId}]: Error: hook process failed`, + ); + expect(sent).toHaveLength(2); + expect(sent[1]).toMatchObject({ displayed_text: 'text', is_final: true }); + }); + + it('ignores chunks that arrive after finish()', async () => { + const { bus, sent, release } = createControlledBus(); + const dispatcher = createDispatcher(bus); + + dispatcher.addChunk('text', 0); + const finished = dispatcher.finish(); + dispatcher.addChunk('late', 10 * PAST_DEBOUNCE); + await release(); + await finished; + + expect(sent).toHaveLength(1); + expect(sent[0]).toMatchObject({ displayed_text: 'text', is_final: true }); + }); +}); diff --git a/packages/core/src/core/message-display-dispatcher.ts b/packages/core/src/core/message-display-dispatcher.ts new file mode 100644 index 00000000000..bf6b9c3905e --- /dev/null +++ b/packages/core/src/core/message-display-dispatcher.ts @@ -0,0 +1,166 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import { randomUUID } from 'node:crypto'; +import type { MessageBus } from '../confirmation-bus/message-bus.js'; +import { + MessageBusType, + type HookExecutionRequest, + type HookExecutionResponse, +} from '../confirmation-bus/types.js'; +import { + createInitialMessageDisplayState, + stepMessageDisplay, + MESSAGE_DISPLAY_DEBOUNCE_MS, + type MessageDisplayState, +} from './message-display-buffer.js'; + +/** + * Owns the delivery side of the MessageDisplay hook for ONE streamed message: + * mints the `message_id`, folds streamed chunks through the pure + * {@link stepMessageDisplay} debounce, and dispatches due flushes through + * MessageBus without ever blocking the streaming loop that feeds it. + * + * Delivery is coalescing, not queueing: at most one hook request is in flight + * at a time, and at most one payload is held pending behind it. A newer flush + * simply overwrites the pending payload — lossless, because `displayed_text` + * is cumulative, so the newest payload strictly supersedes any older one. + * This bounds a slow hook's backlog to O(1) instead of letting undelivered + * batches accumulate for the length of the stream, and it means `is_final` + * is delivered at most one hook-execution behind the message actually ending. + * + * {@link finish} is idempotent and resolves only once every enqueued payload + * has actually been delivered (or failed), so callers can await it before + * ending the turn — without that, a short-lived process (headless `-p`) can + * exit while the final payload is still queued and silently drop it. + */ +export class MessageDisplayDispatcher { + readonly messageId: string = randomUUID(); + + private state: MessageDisplayState; + private pending: { displayedText: string; isFinal: boolean } | null = null; + private inFlight: Promise | null = null; + private finished = false; + private drainWaiters: Array<() => void> = []; + + constructor( + private readonly messageBus: MessageBus, + private readonly signal: AbortSignal, + private readonly warn: (message: string) => void, + nowMs: number = Date.now(), + ) { + this.state = createInitialMessageDisplayState(nowMs); + } + + /** + * Fold one streamed text chunk into the accumulator, firing a debounced + * mid-stream flush if one is due. Never blocks: dispatch happens in the + * background, and the caller's streaming loop continues immediately. + */ + addChunk(chunk: string, nowMs: number = Date.now()): void { + if (this.finished) { + return; + } + const step = stepMessageDisplay( + this.state, + chunk, + nowMs, + MESSAGE_DISPLAY_DEBOUNCE_MS, + false, + ); + this.state = step.next; + if (step.flush) { + this.enqueue(step.flush.displayedText, false); + } + } + + /** + * Close out this message: enqueue the `is_final: true` flush (skipped when + * no text ever streamed — a tool-call-only message — or when the turn was + * aborted, matching the Stop hook's cancellation guard), then wait for + * every enqueued payload to drain. Idempotent — extra calls just await the + * drain, so it is safe to call from both an explicit exit site and a + * `finally` block. The final flush intentionally re-sends the same + * cumulative text as the last debounced flush when nothing changed since + * then: `is_final` is itself new information (it tells subscribers this + * message is done), so the event still fires even when the text didn't. + */ + async finish(): Promise { + if (!this.finished) { + this.finished = true; + if (this.state.displayedText !== '' && !this.signal.aborted) { + this.enqueue(this.state.displayedText, true); + } + } + if (this.signal.aborted) { + // A cancelled turn never fires is_final (matching the Stop hook being + // skipped on abort) and shouldn't hold the turn's teardown hostage to a + // still-running hook process either — leave any in-flight mid-stream + // delivery to settle in the background. + return; + } + await this.drained(); + } + + /** + * Overwrite the single pending slot with the newest payload and kick the + * pump. `is_final` is sticky: once a final payload is pending it stays + * final even if (defensively) a non-final enqueue were to land after it. + */ + private enqueue(displayedText: string, isFinal: boolean): void { + this.pending = { + displayedText, + isFinal: isFinal || (this.pending?.isFinal ?? false), + }; + this.pump(); + } + + private pump(): void { + if (this.inFlight || !this.pending) { + return; + } + const payload = this.pending; + this.pending = null; + this.inFlight = this.messageBus + .request( + { + type: MessageBusType.HOOK_EXECUTION_REQUEST, + eventName: 'MessageDisplay', + input: { + message_id: this.messageId, + displayed_text: payload.displayedText, + is_final: payload.isFinal, + }, + signal: this.signal, + }, + MessageBusType.HOOK_EXECUTION_RESPONSE, + ) + .then(() => undefined) + .catch((err) => { + this.warn(`MessageDisplay hook failed [${this.messageId}]: ${err}`); + }) + .finally(() => { + this.inFlight = null; + if (this.pending) { + this.pump(); + } else { + for (const resolve of this.drainWaiters.splice(0)) { + resolve(); + } + } + }); + } + + /** Resolves once nothing is in flight and nothing is pending. */ + private drained(): Promise { + if (!this.inFlight && !this.pending) { + return Promise.resolve(); + } + return new Promise((resolve) => { + this.drainWaiters.push(resolve); + }); + } +} diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 2fe91f5a6b2..2c30e6bf7e3 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -67,6 +67,7 @@ export * from './core/geminiRequest.js'; export * from './core/inlineMediaLimit.js'; export * from './core/insightProtocol.js'; export * from './core/logger.js'; +export * from './core/message-display-dispatcher.js'; export * from './core/nonInteractiveToolExecutor.js'; export * from './core/prompts.js'; export * from './core/tokenLimits.js'; From a2b8fe5f9f128eff76d0ed7c4fa9dd3fcb3707ce Mon Sep 17 00:00:00 2001 From: Alex Yanchenko Date: Thu, 9 Jul 2026 07:45:42 +0200 Subject: [PATCH 06/10] fix(hooks): bound MessageDisplay drain wait, fix test gaps flagged in review finish() now gives up waiting on drain after 5s (MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS) instead of blocking turn teardown for up to the full 60s hook timeout, per the re-verification's S1 finding. Delivery keeps running in the background past the timeout; only the caller's wait is bounded. Also: add the config.ts bridge test for MessageDisplay field extraction (S5), and add the missing MessageDisplay/InstructionsLoaded entries to acpAgent.test.ts's HookEventName mock (S6). Co-Authored-By: Claude Sonnet 5 --- .../cli/src/acp-integration/acpAgent.test.ts | 2 + packages/core/src/config/config.test.ts | 77 +++++++++++++++++++ .../core/message-display-dispatcher.test.ts | 42 +++++++++- .../src/core/message-display-dispatcher.ts | 37 ++++++++- 4 files changed, 153 insertions(+), 5 deletions(-) diff --git a/packages/cli/src/acp-integration/acpAgent.test.ts b/packages/cli/src/acp-integration/acpAgent.test.ts index 611106a4a66..29767c91a50 100644 --- a/packages/cli/src/acp-integration/acpAgent.test.ts +++ b/packages/cli/src/acp-integration/acpAgent.test.ts @@ -295,6 +295,8 @@ vi.mock('@qwen-code/qwen-code-core', () => ({ StopFailure: 'StopFailure', TodoCreated: 'TodoCreated', TodoCompleted: 'TodoCompleted', + MessageDisplay: 'MessageDisplay', + InstructionsLoaded: 'InstructionsLoaded', }, buildInstallPlan: vi.fn((provider, inputs) => { const authType = inputs.protocol ?? provider.protocol; diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index a428aff4360..fcff878641c 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -55,6 +55,11 @@ import { ToolRegistry } from '../tools/tool-registry.js'; import { ToolNames } from '../tools/tool-names.js'; import { fireNotificationHook } from '../core/toolHookTriggers.js'; import type { MessageBus } from '../confirmation-bus/message-bus.js'; +import { + MessageBusType, + type HookExecutionRequest, + type HookExecutionResponse, +} from '../confirmation-bus/types.js'; import { loadServerHierarchicalMemory } from '../utils/memoryDiscovery.js'; import type { LoadServerHierarchicalMemoryOptions } from '../utils/memoryDiscovery.js'; import { readAutoMemoryIndex } from '../memory/store.js'; @@ -6742,4 +6747,76 @@ describe('Model Switching and Config Updates', () => { expect(buildContextUsage(0, 64000)).toBeUndefined(); }); }); + + describe('MessageDisplay dispatch through the hook execution bridge', () => { + it('extracts message_id/displayed_text/is_final from the request input and forwards them positionally', async () => { + const config = new Config({ ...baseParams }); + await config.initialize(); + + const fireMessageDisplayEvent = vi + .fn() + .mockResolvedValue({ finalOutput: undefined, allOutputs: [] }); + // @ts-expect-error - accessing private for testing + config['hookSystem'] = { fireMessageDisplayEvent }; + + const messageBus = config.getMessageBus(); + expect(messageBus).toBeDefined(); + + const response = await messageBus!.request< + HookExecutionRequest, + HookExecutionResponse + >( + { + type: MessageBusType.HOOK_EXECUTION_REQUEST, + eventName: 'MessageDisplay', + input: { + message_id: 'msg-123', + displayed_text: 'Hello, world', + is_final: true, + }, + }, + MessageBusType.HOOK_EXECUTION_RESPONSE, + ); + + expect(fireMessageDisplayEvent).toHaveBeenCalledWith( + 'msg-123', + 'Hello, world', + true, + undefined, + ); + expect(response.success).toBe(true); + }); + + it('defaults missing fields (empty message_id/text, is_final false) rather than throwing', async () => { + const config = new Config({ ...baseParams }); + await config.initialize(); + + const fireMessageDisplayEvent = vi + .fn() + .mockResolvedValue({ finalOutput: undefined, allOutputs: [] }); + // @ts-expect-error - accessing private for testing + config['hookSystem'] = { fireMessageDisplayEvent }; + + const messageBus = config.getMessageBus(); + const response = await messageBus!.request< + HookExecutionRequest, + HookExecutionResponse + >( + { + type: MessageBusType.HOOK_EXECUTION_REQUEST, + eventName: 'MessageDisplay', + input: {}, + }, + MessageBusType.HOOK_EXECUTION_RESPONSE, + ); + + expect(fireMessageDisplayEvent).toHaveBeenCalledWith( + '', + '', + false, + undefined, + ); + expect(response.success).toBe(true); + }); + }); }); diff --git a/packages/core/src/core/message-display-dispatcher.test.ts b/packages/core/src/core/message-display-dispatcher.test.ts index c261e90fdf5..a91ae5a08e2 100644 --- a/packages/core/src/core/message-display-dispatcher.test.ts +++ b/packages/core/src/core/message-display-dispatcher.test.ts @@ -5,7 +5,10 @@ */ import { describe, it, expect, vi } from 'vitest'; -import { MessageDisplayDispatcher } from './message-display-dispatcher.js'; +import { + MessageDisplayDispatcher, + MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS, +} from './message-display-dispatcher.js'; import { MESSAGE_DISPLAY_DEBOUNCE_MS } from './message-display-buffer.js'; import type { MessageBus } from '../confirmation-bus/message-bus.js'; @@ -222,6 +225,43 @@ describe('MessageDisplayDispatcher', () => { expect(sent[1]).toMatchObject({ displayed_text: 'text', is_final: true }); }); + it('gives up waiting on drain after the timeout and warns, while delivery keeps running in the background', async () => { + vi.useFakeTimers(); + try { + const warn = vi.fn(); + const { bus, sent, release } = createControlledBus(); + const dispatcher = createDispatcher(bus, { warn }); + + dispatcher.addChunk('text', PAST_DEBOUNCE); // in flight, never released + let finishResolved = false; + const finished = dispatcher.finish().then(() => { + finishResolved = true; + }); + + await vi.advanceTimersByTimeAsync(MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS - 1); + expect(finishResolved).toBe(false); + + await vi.advanceTimersByTimeAsync(1); + await finished; + + expect(finishResolved).toBe(true); + expect(warn).toHaveBeenCalledWith( + expect.stringContaining( + `MessageDisplay hook [${dispatcher.messageId}] still running after ${MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS}ms`, + ), + ); + + // finish() stopped waiting, but the in-flight delivery itself is still + // running and completes normally once released. + expect(sent).toHaveLength(1); + await release(); + expect(sent).toHaveLength(2); + expect(sent[1]).toMatchObject({ displayed_text: 'text', is_final: true }); + } finally { + vi.useRealTimers(); + } + }); + it('ignores chunks that arrive after finish()', async () => { const { bus, sent, release } = createControlledBus(); const dispatcher = createDispatcher(bus); diff --git a/packages/core/src/core/message-display-dispatcher.ts b/packages/core/src/core/message-display-dispatcher.ts index bf6b9c3905e..8d44eaf0cfa 100644 --- a/packages/core/src/core/message-display-dispatcher.ts +++ b/packages/core/src/core/message-display-dispatcher.ts @@ -18,6 +18,17 @@ import { type MessageDisplayState, } from './message-display-buffer.js'; +/** + * Ceiling on how long {@link MessageDisplayDispatcher.finish} waits for + * delivery to drain before letting the turn's teardown proceed anyway. Well + * short of `DEFAULT_HOOK_TIMEOUT` (60s, hookRunner.ts) because a turn can be + * blocked behind at most one in-flight hook execution — a slow or hung + * MessageDisplay hook shouldn't be able to freeze `qwen -p` or an ACP stream + * loop's `finally` for anywhere near that long. Delivery keeps running in the + * background past the timeout; this only bounds how long the caller waits. + */ +export const MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS = 5000; + /** * Owns the delivery side of the MessageDisplay hook for ONE streamed message: * mints the `message_id`, folds streamed chunks through the pure @@ -102,7 +113,7 @@ export class MessageDisplayDispatcher { // delivery to settle in the background. return; } - await this.drained(); + await this.drainWithTimeout(); } /** @@ -154,13 +165,31 @@ export class MessageDisplayDispatcher { }); } - /** Resolves once nothing is in flight and nothing is pending. */ - private drained(): Promise { + /** + * Resolves once nothing is in flight and nothing is pending, or after + * {@link MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS} elapses — whichever comes + * first — instead of waiting indefinitely behind a slow or hung hook + * process. Delivery is left running in the background; only the caller's + * wait is bounded. + */ + private drainWithTimeout(): Promise { if (!this.inFlight && !this.pending) { return Promise.resolve(); } return new Promise((resolve) => { - this.drainWaiters.push(resolve); + const timer = setTimeout(() => { + this.warn( + `MessageDisplay hook [${this.messageId}] still running after ` + + `${MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS}ms; continuing without ` + + 'waiting for it to finish (delivery continues in the background).', + ); + resolve(); + }, MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS); + timer.unref?.(); + this.drainWaiters.push(() => { + clearTimeout(timer); + resolve(); + }); }); } } From 2064017cbaa07d6cd1346beaac2ebf2c5fee4784 Mon Sep 17 00:00:00 2001 From: Alex Yanchenko Date: Thu, 9 Jul 2026 18:11:52 +0200 Subject: [PATCH 07/10] fix(hooks): dispatch MessageDisplay is_final alongside stale deliveries, share one drain budget MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-3 review findings on #6489: - finish() no longer queues the is_final payload behind an in-flight mid-stream delivery: the pending slot's supersession argument applies to the in-flight slot too, so the final payload is dispatched immediately, alongside the stale delivery if one is still running. is_final is handed to the hook the moment the message ends — before Stop — on every surface, and can no longer be dropped by a short-lived process exiting with it still queued (Finding 1). - The bounded drain wait is memoized: every finish() call (explicit, finally, or concurrent) shares one promise and one timer, so the teardown ceiling is MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS itself, not a multiple of it (Finding 2). - hooks.md delivery semantics rewritten to match the shipped behavior, including the headless orphaned-hook caveat and the unspecified completion order between an overlapped stale execution and the final one (Finding 3). - The dispatcher mirrors its warnings to console.warn itself (stderr on headless/ACP, ink patchConsole in the TUI) in addition to the injected debug-file sink, so hitting the drain timeout is visible by default (Finding 4). - A superseded mid-stream delivery that fails after the final was dispatched no longer warns; failures during streaming still do. - New tests: finish() twice while delivery is in flight (the exact client.ts sequence), concurrent finish() calls sharing one budget, is_final overtaking a held mid-stream delivery, and drain resolving on the final delivery alone. --- docs/users/features/hooks.md | 5 +- .../src/acp-integration/session/Session.ts | 2 + packages/core/src/core/client.test.ts | 10 + packages/core/src/core/client.ts | 12 +- .../core/message-display-dispatcher.test.ts | 247 +++++++++++++++--- .../src/core/message-display-dispatcher.ts | 198 +++++++++----- 6 files changed, 363 insertions(+), 111 deletions(-) diff --git a/docs/users/features/hooks.md b/docs/users/features/hooks.md index f39e833df24..bce984b7c36 100644 --- a/docs/users/features/hooks.md +++ b/docs/users/features/hooks.md @@ -585,8 +585,9 @@ The `permissionDecision` value controls whether the tool runs: **Delivery semantics** — what a hook script can rely on: -- **Slow hooks see fewer, newer payloads.** At most one hook execution per message is in flight at a time; while one runs, newer debounced payloads _replace_ the queued one rather than piling up behind it. A hook slower than the debounce window therefore skips intermediate snapshots — lossless, since each payload carries the full cumulative text. The `is_final` payload itself is never dropped. -- **`is_final` delivery completes before the turn does.** The turn's end (and the `Stop` hook, when it fires) waits for the final `MessageDisplay` delivery, so a headless run (`qwen -p ...`) exits only after your hook received it, and a consumer combining `MessageDisplay` with `Stop` always sees `is_final` first. +- **Slow hooks see fewer, newer payloads.** At most one mid-stream hook execution per message is in flight at a time; while one runs, newer debounced payloads _replace_ the queued one rather than piling up behind it. A hook slower than the debounce window therefore skips intermediate snapshots — lossless, since each payload carries the full cumulative text. +- **`is_final` is never queued behind a stale delivery.** The final payload is dispatched the moment the message ends — alongside a still-running mid-stream execution if there is one (the one exception to the one-at-a-time rule, justified the same way: the final cumulative text strictly supersedes whatever that execution is processing). Your hook always receives the `is_final` payload, and receives it before the `Stop` hook fires. One consequence for stateful hooks: when the final execution overlaps a superseded mid-stream one, their _completion_ order is unspecified — the stale execution may finish after the final one (even after `Stop`). Treat `is_final` as terminal per `message_id` and let the cumulative text win, rather than assuming the last execution to finish carries the newest state. +- **The turn waits for `is_final` delivery to complete — but not forever.** The turn's end (and the `Stop` hook, when it fires) waits up to 5 seconds for the final delivery to finish. A hook that completes within that budget keeps the strongest guarantee: a headless run (`qwen -p ...`) exits only after the hook finished, and the `is_final` execution completes before `Stop` starts. A slower hook still receives `is_final` first — only the wait for its completion is bounded: in the terminal UI or an ACP session the execution simply finishes in the background, while a headless run exits without waiting. The hook process is not killed on exit; it is left to finish on its own, so a script chaining `qwen -p … && next-step` can observe `next-step` starting while a slow hook is still running. Hitting this timeout prints a warning on stderr. - **A cancelled turn fires no `is_final`.** Cancellation (Esc, an aborted request) stops firings without a terminal event — the message didn't end, it was abandoned. A consumer that buffers until `is_final` should treat cancellation-silence as its flush/discard signal (e.g. a timeout fallback). - **`displayed_text` is provisional until `is_final`.** It reflects what has streamed so far; treat intermediate payloads as display state, not as authoritative final content. - **A tool-using turn produces multiple messages.** Each model call gets its own `message_id` with its own `is_final: true` firing: the text before a tool call is one message, the continuation after the tool result is another. Model calls that produce no displayed text (tool-call-only) fire nothing. diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index f4a77bbec5b..490a1ddf049 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -2356,6 +2356,8 @@ export class Session implements SessionContext { ) { return null; } + // The dispatcher mirrors warnings to console.warn itself; this sink + // only adds them to the debug-log file. return new MessageDisplayDispatcher(messageBus, signal, (message) => debugLogger.warn(message), ); diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index e383d341297..640175273ad 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -7686,6 +7686,9 @@ Other open files: warn: vi.fn(), error: vi.fn(), }; + const consoleWarnSpy = vi + .spyOn(console, 'warn') + .mockImplementation(() => {}); const mockMessageBus = { request: vi.fn().mockRejectedValue(new Error('hook process failed')), response: vi.fn(), @@ -7720,6 +7723,13 @@ Other open files: /^MessageDisplay hook failed \[[0-9a-f-]{36}\]: Error: hook process failed$/, ), ); + // Also surfaced on the console: the debug logger writes only to a + // gated log file, and a dropped/failed delivery is the moment a + // documented guarantee is at stake — it must be visible by default. + expect(consoleWarnSpy).toHaveBeenCalledWith( + expect.stringMatching(/^MessageDisplay hook failed/), + ); + consoleWarnSpy.mockRestore(); }); it('does not end the turn until the final MessageDisplay payload has been delivered', async () => { diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 5fd9a81c3cc..024f8382f7e 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -2384,6 +2384,8 @@ export class GeminiClient { // exits (headless `-p`) before a slow hook's queue drained. Not gated // on !turn.pendingToolCalls the way the Stop hook below is, since a // message boundary and a Stop-worthy end-of-turn are different things. + // The dispatcher mirrors warnings to console.warn itself; this sink + // only adds them to the debug-log file. messageDisplay = hooksEnabled && messageBus && @@ -2532,10 +2534,12 @@ export class GeminiClient { } } - // Deliver `is_final: true` (and drain any still-queued mid-stream - // payload) BEFORE the Stop hook below fires, so consumers combining the - // two events can rely on receiving MessageDisplay's completion signal - // first. + // Dispatch `is_final: true` — and wait, bounded by the shared drain + // budget, for its delivery to complete — BEFORE the Stop hook below + // fires, so a consumer combining the two events receives + // MessageDisplay's completion payload before Stop starts. The finish() + // in the outer finally is a no-op after this one (idempotent, and the + // drain budget is shared, not per-call). await messageDisplay?.finish(); // Track API completion time for thinking block idle cleanup diff --git a/packages/core/src/core/message-display-dispatcher.test.ts b/packages/core/src/core/message-display-dispatcher.test.ts index a91ae5a08e2..2812d2f97c5 100644 --- a/packages/core/src/core/message-display-dispatcher.test.ts +++ b/packages/core/src/core/message-display-dispatcher.test.ts @@ -24,7 +24,9 @@ interface SentPayload { */ function createControlledBus() { const sent: SentPayload[] = []; - const releases: Array<() => void> = []; + // Index-aligned with `sent`; entries are cleared (not removed) once + // settled so positions stay stable. + const releases: Array<(() => void) | undefined> = []; const request = vi.fn( (message: { input: SentPayload }) => new Promise((resolve) => { @@ -36,9 +38,17 @@ function createControlledBus() { bus: { request } as unknown as MessageBus, request, sent, - /** Settle the oldest unresolved request. */ - release: async () => { - releases.shift()?.(); + /** + * Settle one unresolved request: the oldest by default, or the request + * at `index` (position in `sent`) — so a test can settle the final + * delivery while an older mid-stream one is still held in flight. + */ + release: async (index?: number) => { + const i = index ?? releases.findIndex((r) => r !== undefined); + if (i >= 0) { + releases[i]?.(); + releases[i] = undefined; + } // Let the dispatcher's .then/.finally continuations run. await Promise.resolve(); await Promise.resolve(); @@ -121,13 +131,13 @@ describe('MessageDisplayDispatcher', () => { }); }); - it('lets the final flush overwrite a pending mid-stream payload, keeping is_final', async () => { + it('lets the final flush supersede a pending mid-stream payload, keeping is_final', async () => { const { bus, sent, release } = createControlledBus(); const dispatcher = createDispatcher(bus); dispatcher.addChunk('partial ', PAST_DEBOUNCE); // in flight dispatcher.addChunk('more ', 2 * PAST_DEBOUNCE); // pending - const finished = dispatcher.finish(); // overwrites pending, is_final wins + const finished = dispatcher.finish(); // drops pending, dispatches is_final await release(); await release(); await finished; @@ -139,7 +149,7 @@ describe('MessageDisplayDispatcher', () => { }); }); - it('finish() resolves only once every enqueued payload has been delivered', async () => { + it('finish() resolves only once the final payload has been delivered', async () => { const { bus, sent, release } = createControlledBus(); const dispatcher = createDispatcher(bus); @@ -148,16 +158,16 @@ describe('MessageDisplayDispatcher', () => { const finished = dispatcher.finish().then(() => { finishResolved = true; }); + expect(sent).toHaveLength(2); // final dispatched alongside the mid-stream one await Promise.resolve(); await Promise.resolve(); - expect(finishResolved).toBe(false); // mid-stream payload still in flight + expect(finishResolved).toBe(false); // final delivery still in flight - await release(); // mid-stream delivered -> final goes out - expect(finishResolved).toBe(false); // final still in flight - expect(sent).toHaveLength(2); + await release(); // mid-stream delivered; final still in flight + expect(finishResolved).toBe(false); - await release(); + await release(); // final delivered await finished; expect(finishResolved).toBe(true); }); @@ -203,30 +213,118 @@ describe('MessageDisplayDispatcher', () => { }); it('logs a failed delivery with the message_id and still delivers the final flush', async () => { - const warn = vi.fn(); - const sent: SentPayload[] = []; - const request = vi.fn((message: { input: SentPayload }) => { - sent.push(message.input); - return message.input.is_final - ? Promise.resolve({}) - : Promise.reject(new Error('hook process failed')); - }); - const dispatcher = createDispatcher({ request } as unknown as MessageBus, { - warn, - }); + const consoleWarnSpy = vi + .spyOn(console, 'warn') + .mockImplementation(() => {}); + try { + const warn = vi.fn(); + const sent: SentPayload[] = []; + const request = vi.fn((message: { input: SentPayload }) => { + sent.push(message.input); + return message.input.is_final + ? Promise.resolve({}) + : Promise.reject(new Error('hook process failed')); + }); + const dispatcher = createDispatcher( + { request } as unknown as MessageBus, + { warn }, + ); - dispatcher.addChunk('text', PAST_DEBOUNCE); // this delivery fails - await dispatcher.finish(); + dispatcher.addChunk('text', PAST_DEBOUNCE); // this delivery fails + // Let the failure settle while the message is still streaming — a + // failure noticed only after finish() dispatched the final payload is + // deliberately suppressed as superseded (see the next test). + await Promise.resolve(); + await Promise.resolve(); + await dispatcher.finish(); - expect(warn).toHaveBeenCalledWith( - `MessageDisplay hook failed [${dispatcher.messageId}]: Error: hook process failed`, - ); - expect(sent).toHaveLength(2); - expect(sent[1]).toMatchObject({ displayed_text: 'text', is_final: true }); + expect(warn).toHaveBeenCalledWith( + `MessageDisplay hook failed [${dispatcher.messageId}]: Error: hook process failed`, + ); + // The injected sink is typically the gated debug-file logger, so the + // dispatcher itself mirrors every warning to the console — a broken + // delivery must be visible by default, on every surface. + expect(consoleWarnSpy).toHaveBeenCalledWith( + `MessageDisplay hook failed [${dispatcher.messageId}]: Error: hook process failed`, + ); + expect(sent).toHaveLength(2); + expect(sent[1]).toMatchObject({ displayed_text: 'text', is_final: true }); + } finally { + consoleWarnSpy.mockRestore(); + } + }); + + it('does not warn when a superseded mid-stream delivery fails after the final was dispatched', async () => { + const consoleWarnSpy = vi + .spyOn(console, 'warn') + .mockImplementation(() => {}); + try { + const warn = vi.fn(); + const sent: SentPayload[] = []; + let rejectMidStream!: (err: Error) => void; + const request = vi.fn((message: { input: SentPayload }) => { + sent.push(message.input); + return message.input.is_final + ? Promise.resolve({}) + : new Promise((_resolve, reject) => { + rejectMidStream = reject; + }); + }); + const dispatcher = createDispatcher( + { request } as unknown as MessageBus, + { warn }, + ); + + dispatcher.addChunk('text', PAST_DEBOUNCE); // in flight, held + await dispatcher.finish(); // final dispatched alongside, settles fine + + // The stale delivery's outcome no longer matters — the final payload + // superseded it and was delivered. A late failure (e.g. the bus + // request's own timeout) must not alarm anyone about a turn that + // actually completed correctly. + rejectMidStream(new Error('request timed out')); + await Promise.resolve(); + await Promise.resolve(); + + expect(warn).not.toHaveBeenCalled(); + expect(consoleWarnSpy).not.toHaveBeenCalled(); + } finally { + consoleWarnSpy.mockRestore(); + } + }); + + it('shares one drain budget across concurrent finish() calls', async () => { + vi.useFakeTimers(); + const consoleWarnSpy = vi + .spyOn(console, 'warn') + .mockImplementation(() => {}); + try { + const warn = vi.fn(); + const { bus } = createControlledBus(); + const dispatcher = createDispatcher(bus, { warn }); + + dispatcher.addChunk('text', PAST_DEBOUNCE); // in flight, never released + const first = dispatcher.finish(); + const second = dispatcher.finish(); // concurrent, not sequential + + await vi.advanceTimersByTimeAsync(MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS); + await first; + await second; + + // Both calls shared ONE timer and produced ONE warning — a concurrent + // second call must not open a second timeout window of its own. + expect(warn).toHaveBeenCalledTimes(1); + } finally { + consoleWarnSpy.mockRestore(); + vi.useRealTimers(); + } }); it('gives up waiting on drain after the timeout and warns, while delivery keeps running in the background', async () => { vi.useFakeTimers(); + const consoleWarnSpy = vi + .spyOn(console, 'warn') + .mockImplementation(() => {}); try { const warn = vi.fn(); const { bus, sent, release } = createControlledBus(); @@ -238,6 +336,12 @@ describe('MessageDisplayDispatcher', () => { finishResolved = true; }); + // The final payload was dispatched immediately, alongside the stale + // mid-stream delivery — the timeout bounds waiting for the hook to + // finish executing, not whether it receives is_final. + expect(sent).toHaveLength(2); + expect(sent[1]).toMatchObject({ displayed_text: 'text', is_final: true }); + await vi.advanceTimersByTimeAsync(MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS - 1); expect(finishResolved).toBe(false); @@ -251,13 +355,88 @@ describe('MessageDisplayDispatcher', () => { ), ); - // finish() stopped waiting, but the in-flight delivery itself is still - // running and completes normally once released. - expect(sent).toHaveLength(1); + // finish() stopped waiting, but both deliveries are still running in + // the background and settle normally once released. + await release(); await release(); expect(sent).toHaveLength(2); - expect(sent[1]).toMatchObject({ displayed_text: 'text', is_final: true }); + + // The drain-timeout warning also reaches the console: the injected + // sink is typically the gated debug-file logger, and this is the + // moment a documented guarantee is being relaxed. + expect(consoleWarnSpy).toHaveBeenCalledWith( + expect.stringContaining( + `MessageDisplay hook [${dispatcher.messageId}] still running after ${MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS}ms`, + ), + ); + } finally { + consoleWarnSpy.mockRestore(); + vi.useRealTimers(); + } + }); + + it('dispatches is_final immediately alongside a stale in-flight mid-stream delivery instead of queueing behind it', async () => { + const { bus, sent } = createControlledBus(); + const dispatcher = createDispatcher(bus); + + dispatcher.addChunk('The quick', PAST_DEBOUNCE); // in flight, held + dispatcher.addChunk(' brown fox', 2 * PAST_DEBOUNCE); // pending, superseded + void dispatcher.finish(); + + // The final payload must not wait for the stale in-flight delivery to + // settle: it strictly supersedes it (cumulative text), and queueing + // behind it is what dropped is_final in short-lived processes. + expect(sent).toHaveLength(2); + expect(sent[1]).toMatchObject({ + displayed_text: 'The quick brown fox', + is_final: true, + }); + }); + + it('finish() resolves once the final delivery settles, even while a superseded mid-stream delivery is still running', async () => { + const warn = vi.fn(); + const { bus, sent, release } = createControlledBus(); + const dispatcher = createDispatcher(bus, { warn }); + + dispatcher.addChunk('stale', PAST_DEBOUNCE); // in flight, never released + const finished = dispatcher.finish(); + + expect(sent).toHaveLength(2); // final dispatched alongside the stale one + await release(1); // settle ONLY the final delivery + await finished; // must not wait on the stale delivery (or the timeout) + + expect(warn).not.toHaveBeenCalled(); + }); + + it('does not restart the drain budget when finish() is called again while delivery is still in flight', async () => { + vi.useFakeTimers(); + const consoleWarnSpy = vi + .spyOn(console, 'warn') + .mockImplementation(() => {}); + try { + const warn = vi.fn(); + const { bus } = createControlledBus(); + const dispatcher = createDispatcher(bus, { warn }); + + dispatcher.addChunk('text', PAST_DEBOUNCE); // in flight, never released + const first = dispatcher.finish(); + await vi.advanceTimersByTimeAsync(MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS); + await first; // budget spent, one warning + + // The client.ts sequence: an explicit finish() before the Stop hook, + // then a second from the outer finally. The second call must not buy + // the hung delivery another full timeout — the ceiling is the + // constant, not a multiple of it. + let secondResolved = false; + const second = dispatcher.finish().then(() => { + secondResolved = true; + }); + await vi.advanceTimersByTimeAsync(0); + expect(secondResolved).toBe(true); + await second; + expect(warn).toHaveBeenCalledTimes(1); } finally { + consoleWarnSpy.mockRestore(); vi.useRealTimers(); } }); diff --git a/packages/core/src/core/message-display-dispatcher.ts b/packages/core/src/core/message-display-dispatcher.ts index 8d44eaf0cfa..4daabde6320 100644 --- a/packages/core/src/core/message-display-dispatcher.ts +++ b/packages/core/src/core/message-display-dispatcher.ts @@ -19,13 +19,17 @@ import { } from './message-display-buffer.js'; /** - * Ceiling on how long {@link MessageDisplayDispatcher.finish} waits for - * delivery to drain before letting the turn's teardown proceed anyway. Well - * short of `DEFAULT_HOOK_TIMEOUT` (60s, hookRunner.ts) because a turn can be - * blocked behind at most one in-flight hook execution — a slow or hung - * MessageDisplay hook shouldn't be able to freeze `qwen -p` or an ACP stream - * loop's `finally` for anywhere near that long. Delivery keeps running in the - * background past the timeout; this only bounds how long the caller waits. + * Ceiling on how long {@link MessageDisplayDispatcher.finish} waits for the + * final payload's delivery to complete before letting the turn's teardown + * proceed anyway. Well short of `DEFAULT_HOOK_TIMEOUT` (60s, hookRunner.ts) + * because a slow or hung MessageDisplay hook shouldn't be able to freeze + * `qwen -p` or an ACP stream loop's `finally` for anywhere near that long. + * The budget is shared across finish() calls (client.ts calls it from an + * explicit exit site and again from a finally), so this constant is the + * ceiling itself, not a per-call increment. The hook has already received + * the `is_final` payload by the time this wait starts — the timeout only + * bounds how long the caller waits for the hook to finish executing, and + * delivery keeps running in the background past it. */ export const MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS = 5000; @@ -35,28 +39,54 @@ export const MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS = 5000; * {@link stepMessageDisplay} debounce, and dispatches due flushes through * MessageBus without ever blocking the streaming loop that feeds it. * - * Delivery is coalescing, not queueing: at most one hook request is in flight - * at a time, and at most one payload is held pending behind it. A newer flush - * simply overwrites the pending payload — lossless, because `displayed_text` - * is cumulative, so the newest payload strictly supersedes any older one. - * This bounds a slow hook's backlog to O(1) instead of letting undelivered - * batches accumulate for the length of the stream, and it means `is_final` - * is delivered at most one hook-execution behind the message actually ending. + * Mid-stream delivery is coalescing, not queueing: at most one hook request + * is in flight at a time, and at most one payload is held pending behind it. + * A newer flush simply overwrites the pending payload — lossless, because + * `displayed_text` is cumulative, so the newest payload strictly supersedes + * any older one. This bounds a slow hook's backlog to O(1) instead of letting + * undelivered batches accumulate for the length of the stream. * - * {@link finish} is idempotent and resolves only once every enqueued payload - * has actually been delivered (or failed), so callers can await it before - * ending the turn — without that, a short-lived process (headless `-p`) can - * exit while the final payload is still queued and silently drop it. + * The final payload is the one exception to the single-request rule: + * {@link finish} dispatches it immediately, alongside any still-running + * mid-stream delivery, rather than queueing behind it. The same supersession + * argument applies one slot further — an in-flight mid-stream payload carries + * strictly less information than the final one, so waiting for it to settle + * would only delay `is_final` by a full hook execution, and in a short-lived + * process (headless `-p`) could drop it entirely when the process exits with + * the final payload still queued. A hook may therefore see its last + * mid-stream execution overlap the final one. + * + * {@link finish} is idempotent and resolves once the final payload's + * delivery has completed, or after {@link MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS} + * — whichever comes first — so callers can await it before ending the turn + * without a hung hook holding the turn hostage. */ export class MessageDisplayDispatcher { readonly messageId: string = randomUUID(); private state: MessageDisplayState; - private pending: { displayedText: string; isFinal: boolean } | null = null; + /** Newest undelivered mid-stream text, coalesced behind {@link inFlight}. */ + private pending: string | null = null; + /** The one in-flight mid-stream delivery (single-request rule). */ private inFlight: Promise | null = null; + /** The final payload's delivery, dispatched by {@link finish}. */ + private finalDelivery: Promise | null = null; private finished = false; - private drainWaiters: Array<() => void> = []; + /** + * The one bounded wait on {@link finalDelivery}, shared by every finish() + * call — a single timeout budget, no matter how many times (or how + * concurrently) finish() is invoked. + */ + private drain: Promise | null = null; + /** + * @param warn Additional sink for delivery warnings — typically the + * surface's debug-file logger. The dispatcher always mirrors warnings to + * `console.warn` itself, so callers don't need (and shouldn't add) their + * own console wiring: a warning here means a documented delivery + * guarantee is at stake, which must be visible by default even when no + * debug-log session is active. + */ constructor( private readonly messageBus: MessageBus, private readonly signal: AbortSignal, @@ -84,26 +114,38 @@ export class MessageDisplayDispatcher { ); this.state = step.next; if (step.flush) { - this.enqueue(step.flush.displayedText, false); + this.pending = step.flush.displayedText; + this.pump(); } } /** - * Close out this message: enqueue the `is_final: true` flush (skipped when - * no text ever streamed — a tool-call-only message — or when the turn was - * aborted, matching the Stop hook's cancellation guard), then wait for - * every enqueued payload to drain. Idempotent — extra calls just await the + * Close out this message: dispatch the `is_final: true` payload (skipped + * when no text ever streamed — a tool-call-only message — or when the turn + * was aborted, matching the Stop hook's cancellation guard), then wait for + * its delivery to complete, bounded by the shared drain budget. Idempotent + * — extra calls just re-await the (already spent or already settled) * drain, so it is safe to call from both an explicit exit site and a - * `finally` block. The final flush intentionally re-sends the same - * cumulative text as the last debounced flush when nothing changed since - * then: `is_final` is itself new information (it tells subscribers this - * message is done), so the event still fires even when the text didn't. + * `finally` block without doubling the ceiling. The final flush + * intentionally re-sends the same cumulative text as the last debounced + * flush when nothing changed since then: `is_final` is itself new + * information (it tells subscribers this message is done), so the event + * still fires even when the text didn't. */ async finish(): Promise { if (!this.finished) { this.finished = true; if (this.state.displayedText !== '' && !this.signal.aborted) { - this.enqueue(this.state.displayedText, true); + // The final payload strictly supersedes any queued or in-flight + // mid-stream delivery (displayed_text is cumulative), so it never + // waits behind one: drop the pending payload and dispatch is_final + // NOW, alongside the stale delivery if one is still running. This + // is what keeps is_final's dispatch prompt — and ordered before the + // Stop hook — even when the hook is slower than the drain budget: + // the budget below bounds how long we wait for the hook to finish + // executing, not whether it receives the payload. + this.pending = null; + this.finalDelivery = this.dispatch(this.state.displayedText, true); } } if (this.signal.aborted) { @@ -116,34 +158,17 @@ export class MessageDisplayDispatcher { await this.drainWithTimeout(); } - /** - * Overwrite the single pending slot with the newest payload and kick the - * pump. `is_final` is sticky: once a final payload is pending it stays - * final even if (defensively) a non-final enqueue were to land after it. - */ - private enqueue(displayedText: string, isFinal: boolean): void { - this.pending = { - displayedText, - isFinal: isFinal || (this.pending?.isFinal ?? false), - }; - this.pump(); - } - - private pump(): void { - if (this.inFlight || !this.pending) { - return; - } - const payload = this.pending; - this.pending = null; - this.inFlight = this.messageBus + /** Send one payload through MessageBus; failures are logged, never thrown. */ + private dispatch(displayedText: string, isFinal: boolean): Promise { + return this.messageBus .request( { type: MessageBusType.HOOK_EXECUTION_REQUEST, eventName: 'MessageDisplay', input: { message_id: this.messageId, - displayed_text: payload.displayedText, - is_final: payload.isFinal, + displayed_text: displayedText, + is_final: isFinal, }, signal: this.signal, }, @@ -151,45 +176,76 @@ export class MessageDisplayDispatcher { ) .then(() => undefined) .catch((err) => { - this.warn(`MessageDisplay hook failed [${this.messageId}]: ${err}`); - }) - .finally(() => { - this.inFlight = null; - if (this.pending) { - this.pump(); - } else { - for (const resolve of this.drainWaiters.splice(0)) { - resolve(); - } + if (this.finished && !isFinal) { + // This delivery was superseded by the final payload before it + // settled; its outcome no longer matters, so a late failure (e.g. + // the bus request's own timeout) must not alarm anyone about a + // turn that completed correctly. + return; } + this.emitWarning( + `MessageDisplay hook failed [${this.messageId}]: ${err}`, + ); }); } /** - * Resolves once nothing is in flight and nothing is pending, or after + * Route a warning to the console AND the injected sink. The sink is + * typically a gated debug-file logger (a no-op without an active debug-log + * session), and these warnings fire exactly when a documented delivery + * guarantee is at stake — they must reach stderr by default on every + * surface (headless, ACP — stdout carries the protocol, stderr is free — + * and the TUI, where ink's patchConsole renders them above the app). + */ + private emitWarning(message: string): void { + // eslint-disable-next-line no-console + console.warn(message); + this.warn(message); + } + + private pump(): void { + if (this.inFlight || this.pending === null) { + return; + } + const displayedText = this.pending; + this.pending = null; + this.inFlight = this.dispatch(displayedText, false).finally(() => { + this.inFlight = null; + this.pump(); + }); + } + + /** + * Resolves once the final payload's delivery has settled, or after * {@link MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS} elapses — whichever comes - * first — instead of waiting indefinitely behind a slow or hung hook - * process. Delivery is left running in the background; only the caller's - * wait is bounded. + * first. A superseded mid-stream delivery still running in the background + * never holds the drain; only the final payload's own delivery does. The + * wait is memoized: every finish() call — sequential or concurrent — + * shares the same single promise and timer, so the ceiling is the constant + * itself, never a multiple of it, and a call after the delivery settled + * costs nothing. */ private drainWithTimeout(): Promise { - if (!this.inFlight && !this.pending) { + const delivery = this.finalDelivery; + if (!delivery) { return Promise.resolve(); } - return new Promise((resolve) => { + this.drain ??= new Promise((resolve) => { const timer = setTimeout(() => { - this.warn( + this.emitWarning( `MessageDisplay hook [${this.messageId}] still running after ` + `${MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS}ms; continuing without ` + - 'waiting for it to finish (delivery continues in the background).', + 'waiting for it to finish (the hook already received the final ' + + 'payload; its execution continues in the background).', ); resolve(); }, MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS); timer.unref?.(); - this.drainWaiters.push(() => { + void delivery.finally(() => { clearTimeout(timer); resolve(); }); }); + return this.drain; } } From 73add3d8e46a1091e1cc16e94de1353c67fb62dd Mon Sep 17 00:00:00 2001 From: Alex Yanchenko Date: Thu, 9 Jul 2026 23:40:43 +0200 Subject: [PATCH 08/10] refactor(core): consolidate MessageDisplay finish() calls, dedupe test spy setup client.ts: wrap the turn.run() streaming loop in try/finally so messageDisplay.finish() fires once instead of at each of the three early-return sites plus the post-loop path -- matching the pattern the four raw-stream loops in Session.ts already use for the same dispatcher. message-display-dispatcher.test.ts: centralize the console.warn spy setup/teardown in beforeEach/afterEach instead of five repeated per-test try/finally blocks. No behavior change: full client.test.ts (246/246) and the message-display-buffer/dispatcher suites (24/24) pass unchanged. --- packages/core/src/core/client.ts | 267 +++++++------- .../core/message-display-dispatcher.test.ts | 333 ++++++++---------- 2 files changed, 289 insertions(+), 311 deletions(-) diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 024f8382f7e..2238078ed2d 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -2397,151 +2397,152 @@ export class GeminiClient { const resultStream = turn.run(model, requestToSend, signal); let didUpdateIdeContextState = false; - for await (const event of resultStream) { - if (messageDisplay && event.type === GeminiEventType.Content) { - messageDisplay.addChunk(event.value); - } - if (shouldUpdateIdeContextState && !didUpdateIdeContextState) { - this.lastSentIdeContext = nextIdeContext; - this.forceFullIdeContext = false; - didUpdateIdeContextState = true; - } - - // Always-on safety checks (consecutive-identical tool-call guard, - // shell inspection stagnation, and per-turn tool-call cap). These fire - // before the skipLoopDetection gate so they cannot be bypassed by - // configuration. - const alwaysOnLoop = this.loopDetector.checkAlwaysOnSafeties(event); - if (alwaysOnLoop) { - // Drop every tool call collected before the guard fired so the run - // halts here instead of spawning a continuation that re-trips it. - // turn.pendingToolCalls is internal to this loop and is not read - // after the early return — stream consumers (the TUI scheduler and - // the non-interactive runner) build their own list from the yielded - // ToolCallRequest events and stop on LoopDetected. - turn.pendingToolCalls.length = 0; - const loopType = this.loopDetector.getLastLoopType(); - yield { - type: GeminiEventType.LoopDetected, - ...(loopType && { value: { loopType } }), - }; - if (arenaAgentClient) { - await arenaAgentClient.reportError('Loop detected'); + try { + for await (const event of resultStream) { + if (messageDisplay && event.type === GeminiEventType.Content) { + messageDisplay.addChunk(event.value); + } + if (shouldUpdateIdeContextState && !didUpdateIdeContextState) { + this.lastSentIdeContext = nextIdeContext; + this.forceFullIdeContext = false; + didUpdateIdeContextState = true; } - this.lastApiCompletionTimestamp = Date.now(); - if (isTopLevelInteraction) - endInteractionSpan('error', { errorMessage: 'loop detected' }); - this.cancelPendingMemoryPrefetch(); - await messageDisplay?.finish(); - return turn; - } - // Heuristic loop detection is opt-in: `model.skipLoopDetection` - // defaults to true (see settingsSchema) to avoid false-positive - // interruptions. Only the historically false-positive-prone heuristics - // (content/thought repetition, read-file and action stagnation, - // global-duplicate and alternating tool-call patterns) sit behind this - // flag. The precise consecutive-identical guard, shell inspection - // stagnation guard, and per-turn cap run unconditionally in - // checkAlwaysOnSafeties above, so the documented escape hatch only - // relaxes the heuristics (see nonInteractiveCli.ts). - const skipLoopDetection = this.config.getSkipLoopDetection(); - const heuristicLoop = - !skipLoopDetection && - this.loopDetector.addAndCheckHeuristicLoops(event); - if (heuristicLoop) { - const loopType = this.loopDetector.getLastLoopType(); - yield { - type: GeminiEventType.LoopDetected, - ...(loopType && { value: { loopType } }), - }; - if (arenaAgentClient) { - await arenaAgentClient.reportError('Loop detected'); + // Always-on safety checks (consecutive-identical tool-call guard, + // shell inspection stagnation, and per-turn tool-call cap). These fire + // before the skipLoopDetection gate so they cannot be bypassed by + // configuration. + const alwaysOnLoop = this.loopDetector.checkAlwaysOnSafeties(event); + if (alwaysOnLoop) { + // Drop every tool call collected before the guard fired so the run + // halts here instead of spawning a continuation that re-trips it. + // turn.pendingToolCalls is internal to this loop and is not read + // after the early return — stream consumers (the TUI scheduler and + // the non-interactive runner) build their own list from the yielded + // ToolCallRequest events and stop on LoopDetected. + turn.pendingToolCalls.length = 0; + const loopType = this.loopDetector.getLastLoopType(); + yield { + type: GeminiEventType.LoopDetected, + ...(loopType && { value: { loopType } }), + }; + if (arenaAgentClient) { + await arenaAgentClient.reportError('Loop detected'); + } + this.lastApiCompletionTimestamp = Date.now(); + if (isTopLevelInteraction) + endInteractionSpan('error', { errorMessage: 'loop detected' }); + this.cancelPendingMemoryPrefetch(); + return turn; } - this.lastApiCompletionTimestamp = Date.now(); - if (isTopLevelInteraction) - endInteractionSpan('error', { errorMessage: 'loop detected' }); - // finally cleanup catches this, but cancel explicitly to match - // the cleanup pattern at other early-return sites. - this.cancelPendingMemoryPrefetch(); - await messageDisplay?.finish(); - return turn; - } - // Update arena status on Finished events — stats are derived - // automatically from uiTelemetryService by the reporter. - if (arenaAgentClient && event.type === GeminiEventType.Finished) { - await arenaAgentClient.updateStatus(); - } - // Re-send a full IDE context blob on the next regular message — auto - // compaction inside chat.sendMessageStream may have summarized away - // the previous merged IDE context. - if (event.type === GeminiEventType.ChatCompressed) { - this.forceFullIdeContext = true; - // Auto-compaction summarized away the startup prelude. Rebuild it - // before the next turn so env/tool/MCP context isn't lost for the - // rest of the session (manual /compress gets this via startChat). - try { - await this.restoreStartupContextAfterCompaction(); - } catch (error) { - this.config - .getDebugLogger() - .warn( - `Failed to restore startup context after compaction: ${error}`, - ); + // Heuristic loop detection is opt-in: `model.skipLoopDetection` + // defaults to true (see settingsSchema) to avoid false-positive + // interruptions. Only the historically false-positive-prone heuristics + // (content/thought repetition, read-file and action stagnation, + // global-duplicate and alternating tool-call patterns) sit behind this + // flag. The precise consecutive-identical guard, shell inspection + // stagnation guard, and per-turn cap run unconditionally in + // checkAlwaysOnSafeties above, so the documented escape hatch only + // relaxes the heuristics (see nonInteractiveCli.ts). + const skipLoopDetection = this.config.getSkipLoopDetection(); + const heuristicLoop = + !skipLoopDetection && + this.loopDetector.addAndCheckHeuristicLoops(event); + if (heuristicLoop) { + const loopType = this.loopDetector.getLastLoopType(); + yield { + type: GeminiEventType.LoopDetected, + ...(loopType && { value: { loopType } }), + }; + if (arenaAgentClient) { + await arenaAgentClient.reportError('Loop detected'); + } + this.lastApiCompletionTimestamp = Date.now(); + if (isTopLevelInteraction) + endInteractionSpan('error', { errorMessage: 'loop detected' }); + // finally cleanup catches this, but cancel explicitly to match + // the cleanup pattern at other early-return sites. + this.cancelPendingMemoryPrefetch(); + return turn; } - void this.fireSessionStartHook(SessionStartSource.Compact) - .then((compactAdditionalContext) => { - if (!compactAdditionalContext || !this.chat) { - return; - } - this.lastSessionStartContext = compactAdditionalContext; - this.lastSessionStartSource = SessionStartSource.Compact; - this.chat.applySessionStartContext( - compactAdditionalContext, - SessionStartSource.Compact, - ); - }) - .catch((error) => { + // Update arena status on Finished events — stats are derived + // automatically from uiTelemetryService by the reporter. + if (arenaAgentClient && event.type === GeminiEventType.Finished) { + await arenaAgentClient.updateStatus(); + } + + // Re-send a full IDE context blob on the next regular message — auto + // compaction inside chat.sendMessageStream may have summarized away + // the previous merged IDE context. + if (event.type === GeminiEventType.ChatCompressed) { + this.forceFullIdeContext = true; + // Auto-compaction summarized away the startup prelude. Rebuild it + // before the next turn so env/tool/MCP context isn't lost for the + // rest of the session (manual /compress gets this via startChat). + try { + await this.restoreStartupContextAfterCompaction(); + } catch (error) { this.config .getDebugLogger() - .warn(`SessionStart hook failed: ${error}`); - }); - } - - yield event; - if (event.type === GeminiEventType.Error) { - this.forceFullIdeContext = true; - if (arenaAgentClient) { - const errorMsg = - event.value instanceof Error - ? event.value.message - : 'Unknown error'; - await arenaAgentClient.reportError(errorMsg); + .warn( + `Failed to restore startup context after compaction: ${error}`, + ); + } + void this.fireSessionStartHook(SessionStartSource.Compact) + .then((compactAdditionalContext) => { + if (!compactAdditionalContext || !this.chat) { + return; + } + this.lastSessionStartContext = compactAdditionalContext; + this.lastSessionStartSource = SessionStartSource.Compact; + this.chat.applySessionStartContext( + compactAdditionalContext, + SessionStartSource.Compact, + ); + }) + .catch((error) => { + this.config + .getDebugLogger() + .warn(`SessionStart hook failed: ${error}`); + }); } - this.lastApiCompletionTimestamp = Date.now(); - if (isTopLevelInteraction) { - // Sanitize: do not pass raw API error messages to span status - const errMsg = - event.value instanceof Error ? '[API error]' : 'unknown error'; - endInteractionSpan('error', { errorMessage: errMsg }); + + yield event; + if (event.type === GeminiEventType.Error) { + this.forceFullIdeContext = true; + if (arenaAgentClient) { + const errorMsg = + event.value instanceof Error + ? event.value.message + : 'Unknown error'; + await arenaAgentClient.reportError(errorMsg); + } + this.lastApiCompletionTimestamp = Date.now(); + if (isTopLevelInteraction) { + // Sanitize: do not pass raw API error messages to span status + const errMsg = + event.value instanceof Error ? '[API error]' : 'unknown error'; + endInteractionSpan('error', { errorMessage: errMsg }); + } + // finally cleanup catches this, but cancel explicitly to match + // the cleanup pattern at other early-return sites. + this.cancelPendingMemoryPrefetch(); + return turn; } - // finally cleanup catches this, but cancel explicitly to match - // the cleanup pattern at other early-return sites. - this.cancelPendingMemoryPrefetch(); - await messageDisplay?.finish(); - return turn; } + } finally { + // Fires on every exit from the loop above: normal completion, any of + // the three early returns, or an uncaught exception -- instead of one + // explicit call duplicated at each site. This is the pattern the four + // raw-stream loops in Session.ts already use for the same dispatcher. + // finish() is idempotent and dispatches is_final (bounded by the + // shared drain budget) BEFORE the Stop hook below fires; the + // belt-and-suspenders call in the outer finally further down is then + // a no-op. + await messageDisplay?.finish(); } - // Dispatch `is_final: true` — and wait, bounded by the shared drain - // budget, for its delivery to complete — BEFORE the Stop hook below - // fires, so a consumer combining the two events receives - // MessageDisplay's completion payload before Stop starts. The finish() - // in the outer finally is a no-op after this one (idempotent, and the - // drain budget is shared, not per-call). - await messageDisplay?.finish(); - // Track API completion time for thinking block idle cleanup this.lastApiCompletionTimestamp = Date.now(); diff --git a/packages/core/src/core/message-display-dispatcher.test.ts b/packages/core/src/core/message-display-dispatcher.test.ts index 2812d2f97c5..af3504a49d8 100644 --- a/packages/core/src/core/message-display-dispatcher.test.ts +++ b/packages/core/src/core/message-display-dispatcher.test.ts @@ -4,7 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { describe, it, expect, vi } from 'vitest'; +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; import { MessageDisplayDispatcher, MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS, @@ -72,6 +72,21 @@ function createDispatcher( const PAST_DEBOUNCE = MESSAGE_DISPLAY_DEBOUNCE_MS + 1; describe('MessageDisplayDispatcher', () => { + // Centralized here instead of per-test: every test that spies on + // console.warn (and the handful that also fake timers) needs the same + // restore, so a shared afterEach removes the repeated try/finally + // boilerplate. Restoring on tests that never touched these is a no-op. + let consoleWarnSpy: ReturnType; + + beforeEach(() => { + consoleWarnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {}); + }); + + afterEach(() => { + consoleWarnSpy.mockRestore(); + vi.useRealTimers(); + }); + it('delivers a due mid-stream flush and then the final flush, sharing one message_id', async () => { const { bus, sent, release } = createControlledBus(); const dispatcher = createDispatcher(bus); @@ -213,166 +228,136 @@ describe('MessageDisplayDispatcher', () => { }); it('logs a failed delivery with the message_id and still delivers the final flush', async () => { - const consoleWarnSpy = vi - .spyOn(console, 'warn') - .mockImplementation(() => {}); - try { - const warn = vi.fn(); - const sent: SentPayload[] = []; - const request = vi.fn((message: { input: SentPayload }) => { - sent.push(message.input); - return message.input.is_final - ? Promise.resolve({}) - : Promise.reject(new Error('hook process failed')); - }); - const dispatcher = createDispatcher( - { request } as unknown as MessageBus, - { warn }, - ); - - dispatcher.addChunk('text', PAST_DEBOUNCE); // this delivery fails - // Let the failure settle while the message is still streaming — a - // failure noticed only after finish() dispatched the final payload is - // deliberately suppressed as superseded (see the next test). - await Promise.resolve(); - await Promise.resolve(); - await dispatcher.finish(); - - expect(warn).toHaveBeenCalledWith( - `MessageDisplay hook failed [${dispatcher.messageId}]: Error: hook process failed`, - ); - // The injected sink is typically the gated debug-file logger, so the - // dispatcher itself mirrors every warning to the console — a broken - // delivery must be visible by default, on every surface. - expect(consoleWarnSpy).toHaveBeenCalledWith( - `MessageDisplay hook failed [${dispatcher.messageId}]: Error: hook process failed`, - ); - expect(sent).toHaveLength(2); - expect(sent[1]).toMatchObject({ displayed_text: 'text', is_final: true }); - } finally { - consoleWarnSpy.mockRestore(); - } + const warn = vi.fn(); + const sent: SentPayload[] = []; + const request = vi.fn((message: { input: SentPayload }) => { + sent.push(message.input); + return message.input.is_final + ? Promise.resolve({}) + : Promise.reject(new Error('hook process failed')); + }); + const dispatcher = createDispatcher( + { request } as unknown as MessageBus, + { warn }, + ); + + dispatcher.addChunk('text', PAST_DEBOUNCE); // this delivery fails + // Let the failure settle while the message is still streaming — a + // failure noticed only after finish() dispatched the final payload is + // deliberately suppressed as superseded (see the next test). + await Promise.resolve(); + await Promise.resolve(); + await dispatcher.finish(); + + expect(warn).toHaveBeenCalledWith( + `MessageDisplay hook failed [${dispatcher.messageId}]: Error: hook process failed`, + ); + // The injected sink is typically the gated debug-file logger, so the + // dispatcher itself mirrors every warning to the console — a broken + // delivery must be visible by default, on every surface. + expect(consoleWarnSpy).toHaveBeenCalledWith( + `MessageDisplay hook failed [${dispatcher.messageId}]: Error: hook process failed`, + ); + expect(sent).toHaveLength(2); + expect(sent[1]).toMatchObject({ displayed_text: 'text', is_final: true }); }); it('does not warn when a superseded mid-stream delivery fails after the final was dispatched', async () => { - const consoleWarnSpy = vi - .spyOn(console, 'warn') - .mockImplementation(() => {}); - try { - const warn = vi.fn(); - const sent: SentPayload[] = []; - let rejectMidStream!: (err: Error) => void; - const request = vi.fn((message: { input: SentPayload }) => { - sent.push(message.input); - return message.input.is_final - ? Promise.resolve({}) - : new Promise((_resolve, reject) => { - rejectMidStream = reject; - }); - }); - const dispatcher = createDispatcher( - { request } as unknown as MessageBus, - { warn }, - ); - - dispatcher.addChunk('text', PAST_DEBOUNCE); // in flight, held - await dispatcher.finish(); // final dispatched alongside, settles fine - - // The stale delivery's outcome no longer matters — the final payload - // superseded it and was delivered. A late failure (e.g. the bus - // request's own timeout) must not alarm anyone about a turn that - // actually completed correctly. - rejectMidStream(new Error('request timed out')); - await Promise.resolve(); - await Promise.resolve(); + const warn = vi.fn(); + const sent: SentPayload[] = []; + let rejectMidStream!: (err: Error) => void; + const request = vi.fn((message: { input: SentPayload }) => { + sent.push(message.input); + return message.input.is_final + ? Promise.resolve({}) + : new Promise((_resolve, reject) => { + rejectMidStream = reject; + }); + }); + const dispatcher = createDispatcher( + { request } as unknown as MessageBus, + { warn }, + ); + + dispatcher.addChunk('text', PAST_DEBOUNCE); // in flight, held + await dispatcher.finish(); // final dispatched alongside, settles fine + + // The stale delivery's outcome no longer matters — the final payload + // superseded it and was delivered. A late failure (e.g. the bus + // request's own timeout) must not alarm anyone about a turn that + // actually completed correctly. + rejectMidStream(new Error('request timed out')); + await Promise.resolve(); + await Promise.resolve(); - expect(warn).not.toHaveBeenCalled(); - expect(consoleWarnSpy).not.toHaveBeenCalled(); - } finally { - consoleWarnSpy.mockRestore(); - } + expect(warn).not.toHaveBeenCalled(); + expect(consoleWarnSpy).not.toHaveBeenCalled(); }); it('shares one drain budget across concurrent finish() calls', async () => { vi.useFakeTimers(); - const consoleWarnSpy = vi - .spyOn(console, 'warn') - .mockImplementation(() => {}); - try { - const warn = vi.fn(); - const { bus } = createControlledBus(); - const dispatcher = createDispatcher(bus, { warn }); - - dispatcher.addChunk('text', PAST_DEBOUNCE); // in flight, never released - const first = dispatcher.finish(); - const second = dispatcher.finish(); // concurrent, not sequential - - await vi.advanceTimersByTimeAsync(MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS); - await first; - await second; - - // Both calls shared ONE timer and produced ONE warning — a concurrent - // second call must not open a second timeout window of its own. - expect(warn).toHaveBeenCalledTimes(1); - } finally { - consoleWarnSpy.mockRestore(); - vi.useRealTimers(); - } + const warn = vi.fn(); + const { bus } = createControlledBus(); + const dispatcher = createDispatcher(bus, { warn }); + + dispatcher.addChunk('text', PAST_DEBOUNCE); // in flight, never released + const first = dispatcher.finish(); + const second = dispatcher.finish(); // concurrent, not sequential + + await vi.advanceTimersByTimeAsync(MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS); + await first; + await second; + + // Both calls shared ONE timer and produced ONE warning — a concurrent + // second call must not open a second timeout window of its own. + expect(warn).toHaveBeenCalledTimes(1); }); it('gives up waiting on drain after the timeout and warns, while delivery keeps running in the background', async () => { vi.useFakeTimers(); - const consoleWarnSpy = vi - .spyOn(console, 'warn') - .mockImplementation(() => {}); - try { - const warn = vi.fn(); - const { bus, sent, release } = createControlledBus(); - const dispatcher = createDispatcher(bus, { warn }); - - dispatcher.addChunk('text', PAST_DEBOUNCE); // in flight, never released - let finishResolved = false; - const finished = dispatcher.finish().then(() => { - finishResolved = true; - }); - - // The final payload was dispatched immediately, alongside the stale - // mid-stream delivery — the timeout bounds waiting for the hook to - // finish executing, not whether it receives is_final. - expect(sent).toHaveLength(2); - expect(sent[1]).toMatchObject({ displayed_text: 'text', is_final: true }); - - await vi.advanceTimersByTimeAsync(MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS - 1); - expect(finishResolved).toBe(false); - - await vi.advanceTimersByTimeAsync(1); - await finished; - - expect(finishResolved).toBe(true); - expect(warn).toHaveBeenCalledWith( - expect.stringContaining( - `MessageDisplay hook [${dispatcher.messageId}] still running after ${MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS}ms`, - ), - ); - - // finish() stopped waiting, but both deliveries are still running in - // the background and settle normally once released. - await release(); - await release(); - expect(sent).toHaveLength(2); - - // The drain-timeout warning also reaches the console: the injected - // sink is typically the gated debug-file logger, and this is the - // moment a documented guarantee is being relaxed. - expect(consoleWarnSpy).toHaveBeenCalledWith( - expect.stringContaining( - `MessageDisplay hook [${dispatcher.messageId}] still running after ${MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS}ms`, - ), - ); - } finally { - consoleWarnSpy.mockRestore(); - vi.useRealTimers(); - } + const warn = vi.fn(); + const { bus, sent, release } = createControlledBus(); + const dispatcher = createDispatcher(bus, { warn }); + + dispatcher.addChunk('text', PAST_DEBOUNCE); // in flight, never released + let finishResolved = false; + const finished = dispatcher.finish().then(() => { + finishResolved = true; + }); + + // The final payload was dispatched immediately, alongside the stale + // mid-stream delivery — the timeout bounds waiting for the hook to + // finish executing, not whether it receives is_final. + expect(sent).toHaveLength(2); + expect(sent[1]).toMatchObject({ displayed_text: 'text', is_final: true }); + + await vi.advanceTimersByTimeAsync(MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS - 1); + expect(finishResolved).toBe(false); + + await vi.advanceTimersByTimeAsync(1); + await finished; + + expect(finishResolved).toBe(true); + expect(warn).toHaveBeenCalledWith( + expect.stringContaining( + `MessageDisplay hook [${dispatcher.messageId}] still running after ${MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS}ms`, + ), + ); + + // finish() stopped waiting, but both deliveries are still running in + // the background and settle normally once released. + await release(); + await release(); + expect(sent).toHaveLength(2); + + // The drain-timeout warning also reaches the console: the injected + // sink is typically the gated debug-file logger, and this is the + // moment a documented guarantee is being relaxed. + expect(consoleWarnSpy).toHaveBeenCalledWith( + expect.stringContaining( + `MessageDisplay hook [${dispatcher.messageId}] still running after ${MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS}ms`, + ), + ); }); it('dispatches is_final immediately alongside a stale in-flight mid-stream delivery instead of queueing behind it', async () => { @@ -410,35 +395,27 @@ describe('MessageDisplayDispatcher', () => { it('does not restart the drain budget when finish() is called again while delivery is still in flight', async () => { vi.useFakeTimers(); - const consoleWarnSpy = vi - .spyOn(console, 'warn') - .mockImplementation(() => {}); - try { - const warn = vi.fn(); - const { bus } = createControlledBus(); - const dispatcher = createDispatcher(bus, { warn }); - - dispatcher.addChunk('text', PAST_DEBOUNCE); // in flight, never released - const first = dispatcher.finish(); - await vi.advanceTimersByTimeAsync(MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS); - await first; // budget spent, one warning - - // The client.ts sequence: an explicit finish() before the Stop hook, - // then a second from the outer finally. The second call must not buy - // the hung delivery another full timeout — the ceiling is the - // constant, not a multiple of it. - let secondResolved = false; - const second = dispatcher.finish().then(() => { - secondResolved = true; - }); - await vi.advanceTimersByTimeAsync(0); - expect(secondResolved).toBe(true); - await second; - expect(warn).toHaveBeenCalledTimes(1); - } finally { - consoleWarnSpy.mockRestore(); - vi.useRealTimers(); - } + const warn = vi.fn(); + const { bus } = createControlledBus(); + const dispatcher = createDispatcher(bus, { warn }); + + dispatcher.addChunk('text', PAST_DEBOUNCE); // in flight, never released + const first = dispatcher.finish(); + await vi.advanceTimersByTimeAsync(MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS); + await first; // budget spent, one warning + + // The client.ts sequence: an explicit finish() before the Stop hook, + // then a second from the outer finally. The second call must not buy + // the hung delivery another full timeout — the ceiling is the + // constant, not a multiple of it. + let secondResolved = false; + const second = dispatcher.finish().then(() => { + secondResolved = true; + }); + await vi.advanceTimersByTimeAsync(0); + expect(secondResolved).toBe(true); + await second; + expect(warn).toHaveBeenCalledTimes(1); }); it('ignores chunks that arrive after finish()', async () => { From a2ab71c2681eacb5f1b239251e310d47f3d229bc Mon Sep 17 00:00:00 2001 From: Alex Yanchenko Date: Fri, 10 Jul 2026 04:10:13 +0200 Subject: [PATCH 09/10] docs(hooks): clarify MessageDisplay cancellation timing (round-4 nit) --- docs/users/features/hooks.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/users/features/hooks.md b/docs/users/features/hooks.md index bce984b7c36..66d73b2e0b2 100644 --- a/docs/users/features/hooks.md +++ b/docs/users/features/hooks.md @@ -588,7 +588,7 @@ The `permissionDecision` value controls whether the tool runs: - **Slow hooks see fewer, newer payloads.** At most one mid-stream hook execution per message is in flight at a time; while one runs, newer debounced payloads _replace_ the queued one rather than piling up behind it. A hook slower than the debounce window therefore skips intermediate snapshots — lossless, since each payload carries the full cumulative text. - **`is_final` is never queued behind a stale delivery.** The final payload is dispatched the moment the message ends — alongside a still-running mid-stream execution if there is one (the one exception to the one-at-a-time rule, justified the same way: the final cumulative text strictly supersedes whatever that execution is processing). Your hook always receives the `is_final` payload, and receives it before the `Stop` hook fires. One consequence for stateful hooks: when the final execution overlaps a superseded mid-stream one, their _completion_ order is unspecified — the stale execution may finish after the final one (even after `Stop`). Treat `is_final` as terminal per `message_id` and let the cumulative text win, rather than assuming the last execution to finish carries the newest state. - **The turn waits for `is_final` delivery to complete — but not forever.** The turn's end (and the `Stop` hook, when it fires) waits up to 5 seconds for the final delivery to finish. A hook that completes within that budget keeps the strongest guarantee: a headless run (`qwen -p ...`) exits only after the hook finished, and the `is_final` execution completes before `Stop` starts. A slower hook still receives `is_final` first — only the wait for its completion is bounded: in the terminal UI or an ACP session the execution simply finishes in the background, while a headless run exits without waiting. The hook process is not killed on exit; it is left to finish on its own, so a script chaining `qwen -p … && next-step` can observe `next-step` starting while a slow hook is still running. Hitting this timeout prints a warning on stderr. -- **A cancelled turn fires no `is_final`.** Cancellation (Esc, an aborted request) stops firings without a terminal event — the message didn't end, it was abandoned. A consumer that buffers until `is_final` should treat cancellation-silence as its flush/discard signal (e.g. a timeout fallback). +- **Cancellation behaviour depends on timing.** A turn cancelled _while the message is still streaming_ fires no `is_final` — the message didn't end, it was abandoned, and a consumer that buffers until `is_final` should treat cancellation-silence as its flush/discard signal (e.g. a timeout fallback). Cancelling _after the message ended_ (during the drain wait) is different: `is_final` has already been dispatched, and the still-running hook execution may be terminated mid-flight (SIGTERM). - **`displayed_text` is provisional until `is_final`.** It reflects what has streamed so far; treat intermediate payloads as display state, not as authoritative final content. - **A tool-using turn produces multiple messages.** Each model call gets its own `message_id` with its own `is_final: true` firing: the text before a tool call is one message, the continuation after the tool result is another. Model calls that produce no displayed text (tool-call-only) fire nothing. From 1f005f0e9801e3244c627e1a2a7f78b45d48afde Mon Sep 17 00:00:00 2001 From: Alex Yanchenko Date: Fri, 10 Jul 2026 07:13:31 +0200 Subject: [PATCH 10/10] test(hooks): cover the 3 untested MessageDisplay dispatch sites, fix cancellation doc wording Adds MessageDisplay is_final coverage for the Stop-hook continuation loop, the in-session cron fire, and the background-notification loop, each with a normal-completion and an abort case. Adds three MessageDisplayDispatcher edge-case tests: a delivery settling just before the drain timeout, an abort arriving after a drain wait has already started, and addChunk called after abort but before finish(). Rewords the cancellation-timing doc bullet to state the actual criterion (abort signal state when finish() runs) rather than an approximation of it. --- docs/users/features/hooks.md | 2 +- .../acp-integration/session/Session.test.ts | 402 ++++++++++++++++++ .../core/message-display-dispatcher.test.ts | 90 +++- 3 files changed, 485 insertions(+), 9 deletions(-) diff --git a/docs/users/features/hooks.md b/docs/users/features/hooks.md index 66d73b2e0b2..af765b5c065 100644 --- a/docs/users/features/hooks.md +++ b/docs/users/features/hooks.md @@ -588,7 +588,7 @@ The `permissionDecision` value controls whether the tool runs: - **Slow hooks see fewer, newer payloads.** At most one mid-stream hook execution per message is in flight at a time; while one runs, newer debounced payloads _replace_ the queued one rather than piling up behind it. A hook slower than the debounce window therefore skips intermediate snapshots — lossless, since each payload carries the full cumulative text. - **`is_final` is never queued behind a stale delivery.** The final payload is dispatched the moment the message ends — alongside a still-running mid-stream execution if there is one (the one exception to the one-at-a-time rule, justified the same way: the final cumulative text strictly supersedes whatever that execution is processing). Your hook always receives the `is_final` payload, and receives it before the `Stop` hook fires. One consequence for stateful hooks: when the final execution overlaps a superseded mid-stream one, their _completion_ order is unspecified — the stale execution may finish after the final one (even after `Stop`). Treat `is_final` as terminal per `message_id` and let the cumulative text win, rather than assuming the last execution to finish carries the newest state. - **The turn waits for `is_final` delivery to complete — but not forever.** The turn's end (and the `Stop` hook, when it fires) waits up to 5 seconds for the final delivery to finish. A hook that completes within that budget keeps the strongest guarantee: a headless run (`qwen -p ...`) exits only after the hook finished, and the `is_final` execution completes before `Stop` starts. A slower hook still receives `is_final` first — only the wait for its completion is bounded: in the terminal UI or an ACP session the execution simply finishes in the background, while a headless run exits without waiting. The hook process is not killed on exit; it is left to finish on its own, so a script chaining `qwen -p … && next-step` can observe `next-step` starting while a slow hook is still running. Hitting this timeout prints a warning on stderr. -- **Cancellation behaviour depends on timing.** A turn cancelled _while the message is still streaming_ fires no `is_final` — the message didn't end, it was abandoned, and a consumer that buffers until `is_final` should treat cancellation-silence as its flush/discard signal (e.g. a timeout fallback). Cancelling _after the message ended_ (during the drain wait) is different: `is_final` has already been dispatched, and the still-running hook execution may be terminated mid-flight (SIGTERM). +- **Cancellation behaviour depends on timing.** A turn cancelled _before `is_final` dispatches_ fires no `is_final` — the message is treated as abandoned, and a consumer that buffers until `is_final` should treat cancellation-silence as its flush/discard signal (e.g. a timeout fallback). The criterion is the abort signal's state at the moment the turn ends, not whether every chunk had already streamed — an abort landing in the brief gap before that check can still suppress `is_final` for a message whose text had, in practice, finished arriving. Cancelling _after `is_final` has dispatched_ (during the drain wait) is different: the still-running hook execution may be terminated mid-flight (SIGTERM), but the payload itself has already been delivered. - **`displayed_text` is provisional until `is_final`.** It reflects what has streamed so far; treat intermediate payloads as display state, not as authoritative final content. - **A tool-using turn produces multiple messages.** Each model call gets its own `message_id` with its own `is_final: true` firing: the text before a tool call is one message, the continuation after the tool result is another. Model calls that produce no displayed text (tool-call-only) fire nothing. diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index 8a70356a56c..ce10d9250b2 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -2112,6 +2112,143 @@ describe('Session', () => { ); }); + it('fires MessageDisplay with cumulative text and a single is_final for a background notification response', async () => { + // The background-notification loop (Session.ts ~line 3638) creates its + // own MessageDisplayDispatcher, independent of the ACP prompt path's — + // a regression here would not be caught by the prompt-path test alone. + const messageBus = { request: vi.fn().mockResolvedValue({}) }; + mockConfig.getMessageBus = vi.fn().mockReturnValue(messageBus); + mockConfig.getDisableAllHooks = vi.fn().mockReturnValue(false); + mockConfig.hasHooksForEvent = vi + .fn() + .mockImplementation( + (eventName: string) => eventName === 'MessageDisplay', + ); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce(createEmptyStream()) + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + candidates: [ + { + content: { + parts: [{ text: 'I saw the background result.' }], + }, + }, + ], + }, + }, + ]), + ); + + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'start background work' }], + }); + + const callback = mockBackgroundTaskRegistry.setNotificationCallback.mock + .calls[0][0] as ( + displayText: string, + modelText: string, + meta: { agentId: string; status: string; toolUseId?: string }, + ) => void; + + callback( + 'Background agent "worker" completed.', + 'completed', + { + agentId: 'agent-1', + status: 'completed', + toolUseId: 'tool-1', + }, + ); + + await vi.waitFor(() => { + const finals = messageBus.request.mock.calls.filter( + ([request]) => + request.eventName === 'MessageDisplay' && request.input.is_final, + ); + expect(finals).toHaveLength(1); + }); + + const messageDisplayCalls = messageBus.request.mock.calls.filter( + ([request]) => request.eventName === 'MessageDisplay', + ); + const finalCall = messageDisplayCalls[messageDisplayCalls.length - 1][0]; + expect(finalCall.input).toMatchObject({ + displayed_text: 'I saw the background result.', + is_final: true, + }); + }); + + it('suppresses is_final for MessageDisplay when a background notification response is cancelled mid-stream', async () => { + let releaseNotification: () => void; + const notificationGate = new Promise((resolve) => { + releaseNotification = resolve; + }); + async function* notificationStream() { + yield { + type: core.StreamEventType.CHUNK, + value: { + candidates: [ + { content: { parts: [{ text: 'partial background reply' }] } }, + ], + }, + }; + await notificationGate; + } + + const messageBus = { request: vi.fn().mockResolvedValue({}) }; + mockConfig.getMessageBus = vi.fn().mockReturnValue(messageBus); + mockConfig.getDisableAllHooks = vi.fn().mockReturnValue(false); + mockConfig.hasHooksForEvent = vi + .fn() + .mockImplementation( + (eventName: string) => eventName === 'MessageDisplay', + ); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce(createEmptyStream()) + .mockResolvedValueOnce(notificationStream()); + + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'start background work' }], + }); + + const callback = mockBackgroundTaskRegistry.setNotificationCallback.mock + .calls[0][0] as ( + displayText: string, + modelText: string, + meta: { agentId: string; status: string; toolUseId?: string }, + ) => void; + + callback('done', '', { + agentId: 'agent-1', + status: 'completed', + }); + + // Wait until the notification's own streamed send has started (the + // dispatcher exists and has received the first chunk) rather than for + // a mid-stream MessageDisplay flush, which is debounced (~200ms) and + // may not be due yet by the time we cancel. + await vi.waitFor(() => + expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2), + ); + + await session.cancelPendingPrompt(); + releaseNotification!(); + + const finals = messageBus.request.mock.calls.filter( + ([request]) => + request.eventName === 'MessageDisplay' && request.input.is_final, + ); + expect(finals).toHaveLength(0); + }); + it('cancels an in-flight background notification prompt', async () => { const notificationCompression = { signal: undefined as AbortSignal | undefined, @@ -9073,6 +9210,126 @@ describe('Session', () => { ); expect(spawner).not.toHaveBeenCalled(); }); + + it('fires MessageDisplay with cumulative text and a single is_final for an in-session cron fire', async () => { + // The cron loop (Session.ts ~line 3321) creates its own + // MessageDisplayDispatcher, independent of the ACP prompt path's — + // a regression here would not be caught by the prompt-path test alone. + const messageBus = { request: vi.fn().mockResolvedValue({}) }; + const scheduler = schedulerFiring({ + prompt: 'nightly report', + runMode: 'shared', + }); + mockConfig.isCronEnabled = vi.fn().mockReturnValue(true); + mockConfig.getCronScheduler = vi.fn().mockReturnValue(scheduler); + mockConfig.getSubSessionSpawner = vi.fn().mockReturnValue(undefined); + mockConfig.getMessageBus = vi.fn().mockReturnValue(messageBus); + mockConfig.getDisableAllHooks = vi.fn().mockReturnValue(false); + mockConfig.hasHooksForEvent = vi + .fn() + .mockImplementation( + (eventName: string) => eventName === 'MessageDisplay', + ); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce(createEmptyStream()) + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + candidates: [ + { content: { parts: [{ text: 'cron result' }] } }, + ], + }, + }, + ]), + ); + + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'hello' }], + }); + + await vi.waitFor(() => { + const finals = messageBus.request.mock.calls.filter( + ([request]) => + request.eventName === 'MessageDisplay' && request.input.is_final, + ); + expect(finals).toHaveLength(1); + }); + + const messageDisplayCalls = messageBus.request.mock.calls.filter( + ([request]) => request.eventName === 'MessageDisplay', + ); + const finalCall = + messageDisplayCalls[messageDisplayCalls.length - 1][0]; + expect(finalCall.input).toMatchObject({ + displayed_text: 'cron result', + is_final: true, + }); + }); + + it('suppresses is_final for MessageDisplay when a cron fire is cancelled mid-stream', async () => { + let releaseCron: () => void; + const cronGate = new Promise((resolve) => { + releaseCron = resolve; + }); + async function* cronStream() { + yield { + type: core.StreamEventType.CHUNK, + value: { + candidates: [ + { content: { parts: [{ text: 'partial cron result' }] } }, + ], + }, + }; + await cronGate; + } + + const messageBus = { request: vi.fn().mockResolvedValue({}) }; + const scheduler = schedulerFiring({ + prompt: 'nightly report', + runMode: 'shared', + }); + mockConfig.isCronEnabled = vi.fn().mockReturnValue(true); + mockConfig.getCronScheduler = vi.fn().mockReturnValue(scheduler); + mockConfig.getSubSessionSpawner = vi.fn().mockReturnValue(undefined); + mockConfig.getMessageBus = vi.fn().mockReturnValue(messageBus); + mockConfig.getDisableAllHooks = vi.fn().mockReturnValue(false); + mockConfig.hasHooksForEvent = vi + .fn() + .mockImplementation( + (eventName: string) => eventName === 'MessageDisplay', + ); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce(createEmptyStream()) + .mockResolvedValueOnce(cronStream()); + + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'hello' }], + }); + + // Wait until the cron fire's own streamed send has started (the + // dispatcher exists and has received the first chunk) rather than + // for a mid-stream MessageDisplay flush, which is debounced + // (~200ms) and may not be due yet by the time we cancel. + await vi.waitFor(() => + expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2), + ); + + await session.cancelPendingPrompt(); + releaseCron!(); + + expect(scheduler.stop).toHaveBeenCalled(); + const finals = messageBus.request.mock.calls.filter( + ([request]) => + request.eventName === 'MessageDisplay' && request.input.is_final, + ); + expect(finals).toHaveLength(0); + }); }); describe('hooks', () => { @@ -9454,6 +9711,151 @@ describe('Session', () => { }, }); }); + + it('fires MessageDisplay with cumulative text and a single is_final during Stop hook continuation', async () => { + // The Stop-hook continuation loop (Session.ts ~line 2282) creates + // its own MessageDisplayDispatcher, independent of the main prompt + // loop's — a regression here would not be caught by the + // ACP-prompt-path test alone. + let stopHookCalls = 0; + const messageBus = { + request: vi.fn().mockImplementation(async (request) => { + if (request.eventName === 'Stop') { + stopHookCalls++; + return stopHookCalls === 1 + ? { + success: true, + output: { decision: 'block', reason: 'keep going' }, + } + : { success: true, output: {} }; + } + return { success: true, output: {} }; + }), + }; + mockConfig.getMessageBus = vi.fn().mockReturnValue(messageBus); + mockConfig.getDisableAllHooks = vi.fn().mockReturnValue(false); + mockConfig.hasHooksForEvent = vi + .fn() + .mockImplementation( + (eventName: string) => + eventName === 'Stop' || eventName === 'MessageDisplay', + ); + mockChat.getHistory = vi + .fn() + .mockReturnValue([ + { role: 'model', parts: [{ text: 'response text' }] }, + ]); + mockChat.getLastModelMessageText = vi + .fn() + .mockReturnValue('response text'); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce(createEmptyStream()) + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + candidates: [ + { content: { parts: [{ text: 'continued reply' }] } }, + ], + }, + }, + ]), + ); + + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'hello' }], + }); + + const messageDisplayCalls = messageBus.request.mock.calls.filter( + ([request]) => request.eventName === 'MessageDisplay', + ); + expect(messageDisplayCalls.length).toBeGreaterThan(0); + const finalCall = + messageDisplayCalls[messageDisplayCalls.length - 1][0]; + expect(finalCall.input).toMatchObject({ + displayed_text: 'continued reply', + is_final: true, + }); + expect( + messageDisplayCalls.filter(([request]) => request.input.is_final), + ).toHaveLength(1); + }); + + it('suppresses is_final for MessageDisplay when the turn is cancelled mid Stop-hook continuation', async () => { + let releaseContinuation: () => void; + const continuationGate = new Promise((resolve) => { + releaseContinuation = resolve; + }); + async function* continuationStream() { + yield { + type: core.StreamEventType.CHUNK, + value: { + candidates: [ + { content: { parts: [{ text: 'partial continuation' }] } }, + ], + }, + }; + await continuationGate; + } + + const messageBus = { + request: vi.fn().mockImplementation(async (request) => { + if (request.eventName === 'Stop') { + return { + success: true, + output: { decision: 'block', reason: 'keep going' }, + }; + } + return { success: true, output: {} }; + }), + }; + mockConfig.getMessageBus = vi.fn().mockReturnValue(messageBus); + mockConfig.getDisableAllHooks = vi.fn().mockReturnValue(false); + mockConfig.hasHooksForEvent = vi + .fn() + .mockImplementation( + (eventName: string) => + eventName === 'Stop' || eventName === 'MessageDisplay', + ); + mockChat.getHistory = vi + .fn() + .mockReturnValue([ + { role: 'model', parts: [{ text: 'response text' }] }, + ]); + mockChat.getLastModelMessageText = vi + .fn() + .mockReturnValue('response text'); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce(createEmptyStream()) + .mockResolvedValueOnce(continuationStream()); + + const promptPromise = session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'hello' }], + }); + + // Wait until the continuation's own streamed send has started + // (the dispatcher exists and has received the first chunk) rather + // than for a mid-stream MessageDisplay flush, which is debounced + // (~200ms) and may not be due yet by the time we cancel. + await vi.waitFor(() => + expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2), + ); + + await session.cancelPendingPrompt(); + releaseContinuation!(); + await promptPromise; + + const finals = messageBus.request.mock.calls.filter( + ([request]) => + request.eventName === 'MessageDisplay' && request.input.is_final, + ); + expect(finals).toHaveLength(0); + }); }); describe('PreToolUse hook', () => { diff --git a/packages/core/src/core/message-display-dispatcher.test.ts b/packages/core/src/core/message-display-dispatcher.test.ts index af3504a49d8..a3eb4336d34 100644 --- a/packages/core/src/core/message-display-dispatcher.test.ts +++ b/packages/core/src/core/message-display-dispatcher.test.ts @@ -236,10 +236,9 @@ describe('MessageDisplayDispatcher', () => { ? Promise.resolve({}) : Promise.reject(new Error('hook process failed')); }); - const dispatcher = createDispatcher( - { request } as unknown as MessageBus, - { warn }, - ); + const dispatcher = createDispatcher({ request } as unknown as MessageBus, { + warn, + }); dispatcher.addChunk('text', PAST_DEBOUNCE); // this delivery fails // Let the failure settle while the message is still streaming — a @@ -274,10 +273,9 @@ describe('MessageDisplayDispatcher', () => { rejectMidStream = reject; }); }); - const dispatcher = createDispatcher( - { request } as unknown as MessageBus, - { warn }, - ); + const dispatcher = createDispatcher({ request } as unknown as MessageBus, { + warn, + }); dispatcher.addChunk('text', PAST_DEBOUNCE); // in flight, held await dispatcher.finish(); // final dispatched alongside, settles fine @@ -294,6 +292,82 @@ describe('MessageDisplayDispatcher', () => { expect(consoleWarnSpy).not.toHaveBeenCalled(); }); + it('resolves the drain via the delivery settling just before the timeout, without warning', async () => { + vi.useFakeTimers(); + const warn = vi.fn(); + const { bus, release } = createControlledBus(); + const dispatcher = createDispatcher(bus, { warn }); + + dispatcher.addChunk('text', PAST_DEBOUNCE); // in flight, held + let finishResolved = false; + const finished = dispatcher.finish().then(() => { + finishResolved = true; + }); + + await vi.advanceTimersByTimeAsync(MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS - 1); + expect(finishResolved).toBe(false); + + // Settle the final delivery (index 1: dispatched alongside the stale + // mid-stream one) one tick before the drain timer would fire — the + // drain must resolve via delivery.finally clearing the timer, not via + // the timeout warning path. + await release(1); + await finished; + + expect(finishResolved).toBe(true); + expect(warn).not.toHaveBeenCalled(); + expect(consoleWarnSpy).not.toHaveBeenCalled(); + }); + + it('does not shorten an already-started drain wait when the signal aborts afterward', async () => { + // drainWithTimeout() has no abort check of its own — once finish() has + // captured `finalDelivery` and started the drain, an abort arriving + // later does not cut the wait short; only an abort present when + // finish() itself runs affects dispatch (see the "suppresses the final + // flush ... when aborted" test above). + vi.useFakeTimers(); + const controller = new AbortController(); + const { bus } = createControlledBus(); + const dispatcher = createDispatcher(bus, { signal: controller.signal }); + + dispatcher.addChunk('text', PAST_DEBOUNCE); // in flight, never released + let finishResolved = false; + const finished = dispatcher.finish().then(() => { + finishResolved = true; + }); + + controller.abort(); + await vi.advanceTimersByTimeAsync(MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS - 1); + expect(finishResolved).toBe(false); + + await vi.advanceTimersByTimeAsync(1); + await finished; + expect(finishResolved).toBe(true); + }); + + it('does not suppress a mid-stream flush from addChunk called after abort but before finish()', async () => { + // Neither addChunk nor dispatch() consult the abort signal — only + // finish()'s is_final dispatch does. Callers rely on their own + // streaming loop's abort check before calling addChunk (see Session.ts's + // `if (signal.aborted) return` guards); the dispatcher in isolation + // still fires a due mid-stream flush after the signal has aborted. + const controller = new AbortController(); + const { bus, sent, release } = createControlledBus(); + const dispatcher = createDispatcher(bus, { signal: controller.signal }); + + controller.abort(); + dispatcher.addChunk('text', PAST_DEBOUNCE); + + expect(sent).toHaveLength(1); + expect(sent[0]).toMatchObject({ displayed_text: 'text', is_final: false }); + + await release(); + await dispatcher.finish(); + + // finish() itself still suppresses is_final once aborted. + expect(sent.filter((payload) => payload.is_final)).toHaveLength(0); + }); + it('shares one drain budget across concurrent finish() calls', async () => { vi.useFakeTimers(); const warn = vi.fn();