diff --git a/.qwen/skills/triage/references/pr-workflow.md b/.qwen/skills/triage/references/pr-workflow.md index de61460cf83..82a9a01f9bb 100644 --- a/.qwen/skills/triage/references/pr-workflow.md +++ b/.qwen/skills/triage/references/pr-workflow.md @@ -363,6 +363,7 @@ gh pr review "$PR_NUMBER" --repo "$REPO" --request-changes --body "Needs some re ``` Genuinely unsure, or `GUARD` blocked approval — **don't approve or reject**, but **never defer silently**. Post an explicit defer comment that: + 1. States you are escalating to the maintainer. 2. Names the specific reason(s) for uncertainty — what you cannot resolve from the diff, tests, and PR description. 3. @mentions the maintainer (use `$QWEN_MAINTAINER_HANDLE` if set, or the most recent human reviewer). diff --git a/docs/users/features/hooks.md b/docs/users/features/hooks.md index 1f49a0c12b2..af765b5c065 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 | @@ -268,6 +269,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:** @@ -565,6 +567,33 @@ 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 once the message ends, so the reply's tail is never dropped waiting on the debounce window. + +**Delivery semantics** — what a hook script can rely on: + +- **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 _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. + +**Note**: Fires in the terminal UI, headless (`-p`), and ACP (IDE/editor/`qwen serve`) sessions, with the same payload contract on every surface. + #### 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 e923c627c5a..cb94f844b48 100644 --- a/packages/acp-bridge/src/status.ts +++ b/packages/acp-bridge/src/status.ts @@ -948,6 +948,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/acp-integration/acpAgent.test.ts b/packages/cli/src/acp-integration/acpAgent.test.ts index cd82042534a..8dc19a4dba9 100644 --- a/packages/cli/src/acp-integration/acpAgent.test.ts +++ b/packages/cli/src/acp-integration/acpAgent.test.ts @@ -333,6 +333,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/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index 018b0ac0926..9c4f03c6fc8 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -1915,6 +1915,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() @@ -2030,6 +2113,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, @@ -8822,6 +9042,130 @@ describe('Session', () => { ); }); + describe('in-session cron MessageDisplay', () => { + /** Mock scheduler that delivers exactly one in-session job through `start`. */ + function schedulerFiring(job: { prompt: string }) { + return { + size: 1, + hasPendingWork: true, + start: vi.fn((callback: (j: typeof job) => void) => callback(job)), + stop: vi.fn(), + getExitSummary: vi.fn().mockReturnValue(undefined), + }; + } + + it('fires MessageDisplay with cumulative text and a single is_final for an in-session cron fire', async () => { + // The cron loop (Session.ts #executeCronPromptInner) 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' }); + mockConfig.isCronEnabled = vi.fn().mockReturnValue(true); + mockConfig.getCronScheduler = vi.fn().mockReturnValue(scheduler); + 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' }); + mockConfig.isCronEnabled = vi.fn().mockReturnValue(true); + mockConfig.getCronScheduler = vi.fn().mockReturnValue(scheduler); + 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', () => { describe('PermissionDenied hook', () => { it('fires PermissionDenied hooks for AUTO classifier blocks', async () => { @@ -9201,6 +9545,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/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index c06c622fed1..172f00be5a0 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -74,6 +74,7 @@ import { createHookOutput, generateToolUseId, MessageBusType, + MessageDisplayDispatcher, getPlanModeSystemReminder, getArenaSystemReminder, getStartupContextLength, @@ -1953,6 +1954,9 @@ export class Session implements SessionContext { let usageMetadata: GenerateContentResponseUsageMetadata | null = null; const streamStartTime = Date.now(); + const messageDisplay = this.#createMessageDisplayDispatcher( + pendingSend.signal, + ); try { const sendResult = @@ -1997,6 +2001,9 @@ export class Session implements SessionContext { 'assistant', part.thought, ); + if (!part.thought) { + messageDisplay?.addChunk(part.text); + } } } @@ -2080,6 +2087,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) { @@ -2282,6 +2294,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 = @@ -2319,6 +2334,9 @@ export class Session implements SessionContext { 'assistant', part.thought, ); + if (!part.thought) { + messageDisplay?.addChunk(part.text); + } } } @@ -2372,6 +2390,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) { @@ -2439,6 +2461,35 @@ 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; + } + // 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), + ); + } + async #sendMessageStreamWithAutoCompression( promptId: string, message: Part[], @@ -3220,42 +3271,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) { @@ -3525,49 +3588,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/cli/src/config/settingsSchema.ts b/packages/cli/src/config/settingsSchema.ts index faa0f2db20d..2a73f4e9951 100644 --- a/packages/cli/src/config/settingsSchema.ts +++ b/packages/cli/src/config/settingsSchema.ts @@ -2931,6 +2931,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.test.ts b/packages/core/src/config/config.test.ts index 38eb7e7a502..e5366ac6409 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'; @@ -6878,4 +6883,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/config/config.ts b/packages/core/src/config/config.ts index 6090cdd35df..1e073904209 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -2252,6 +2252,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 e1be9872490..193c47e1dad 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, @@ -7542,6 +7543,453 @@ 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('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 }, + }); + // 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 () => { + const debugLogger = { + isEnabled: vi.fn().mockReturnValue(true), + debug: vi.fn(), + info: vi.fn(), + 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(), + }; + 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 + } + + // 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.stringMatching( + /^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 () => { + 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({}), + 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 d050d988156..46290172731 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -128,6 +128,7 @@ import { import { partToString } from '../utils/partUtils.js'; import { createHookOutput, SessionStartSource } from '../hooks/types.js'; import fsPromises from 'node:fs/promises'; +import { MessageDisplayDispatcher } from './message-display-dispatcher.js'; // IDE integration import { ideContextStore } from '../ide/ideContext.js'; @@ -1971,6 +1972,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 || @@ -2358,137 +2363,176 @@ export class GeminiClient { }; }; + // 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. + // The dispatcher mirrors warnings to console.warn itself; this sink + // only adds them to the debug-log file. + messageDisplay = + hooksEnabled && + 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 (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(); - 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(); - 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(); - 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(); } // Track API completion time for thinking block idle cleanup @@ -2768,9 +2812,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 new file mode 100644 index 00000000000..2cd0adfc666 --- /dev/null +++ b/packages/core/src/core/message-display-buffer.test.ts @@ -0,0 +1,122 @@ +/** + * @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('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; + 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..2711fbe0c73 --- /dev/null +++ b/packages/core/src/core/message-display-buffer.ts @@ -0,0 +1,77 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * 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; + +/** + * 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/core/message-display-dispatcher.test.ts b/packages/core/src/core/message-display-dispatcher.test.ts new file mode 100644 index 00000000000..a3eb4336d34 --- /dev/null +++ b/packages/core/src/core/message-display-dispatcher.test.ts @@ -0,0 +1,508 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; +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'; + +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[] = []; + // 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) => { + sent.push(message.input); + releases.push(() => resolve({})); + }), + ); + return { + bus: { request } as unknown as MessageBus, + request, + sent, + /** + * 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(); + 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', () => { + // 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); + + 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 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(); // drops pending, dispatches is_final + 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 the final 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; + }); + expect(sent).toHaveLength(2); // final dispatched alongside the mid-stream one + + await Promise.resolve(); + await Promise.resolve(); + expect(finishResolved).toBe(false); // final delivery still in flight + + await release(); // mid-stream delivered; final still in flight + expect(finishResolved).toBe(false); + + await release(); // final delivered + 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 + // 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 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(); + }); + + 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(); + 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 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 () => { + 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 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 () => { + 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..4daabde6320 --- /dev/null +++ b/packages/core/src/core/message-display-dispatcher.ts @@ -0,0 +1,251 @@ +/** + * @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'; + +/** + * 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; + +/** + * 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. + * + * 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. + * + * 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; + /** 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; + /** + * 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, + 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.pending = step.flush.displayedText; + this.pump(); + } + } + + /** + * 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 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) { + // 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) { + // 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.drainWithTimeout(); + } + + /** 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: displayedText, + is_final: isFinal, + }, + signal: this.signal, + }, + MessageBusType.HOOK_EXECUTION_RESPONSE, + ) + .then(() => undefined) + .catch((err) => { + 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}`, + ); + }); + } + + /** + * 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. 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 { + const delivery = this.finalDelivery; + if (!delivery) { + return Promise.resolve(); + } + this.drain ??= new Promise((resolve) => { + const timer = setTimeout(() => { + this.emitWarning( + `MessageDisplay hook [${this.messageId}] still running after ` + + `${MESSAGE_DISPLAY_DRAIN_TIMEOUT_MS}ms; continuing without ` + + '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?.(); + void delivery.finally(() => { + clearTimeout(timer); + resolve(); + }); + }); + return this.drain; + } +} 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 9fd22f70e36..097be1546f4 100644 --- a/packages/core/src/hooks/hookEventHandler.test.ts +++ b/packages/core/src/hooks/hookEventHandler.test.ts @@ -531,6 +531,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 ca8b5b99b95..2c0868c42dd 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, @@ -251,6 +252,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 5a9bf02c41d..529161c9f3f 100644 --- a/packages/core/src/hooks/hookPlanner.test.ts +++ b/packages/core/src/hooks/hookPlanner.test.ts @@ -111,6 +111,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 94df2ef3d19..fd325692b35 100644 --- a/packages/core/src/hooks/hookPlanner.ts +++ b/packages/core/src/hooks/hookPlanner.ts @@ -74,6 +74,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 ee82115dee9..ed9dfb50435 100644 --- a/packages/core/src/hooks/hookSystem.test.ts +++ b/packages/core/src/hooks/hookSystem.test.ts @@ -91,6 +91,7 @@ describe('HookSystem', () => { fireInstructionsLoadedEvent: vi.fn(), fireUserPromptExpansionEvent: vi.fn(), fireStopEvent: vi.fn(), + fireMessageDisplayEvent: vi.fn(), fireSessionStartEvent: vi.fn(), fireSessionEndEvent: vi.fn(), firePreToolUseEvent: vi.fn(), @@ -382,6 +383,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 8145288b72b..c3ab1fa706e 100644 --- a/packages/core/src/hooks/hookSystem.ts +++ b/packages/core/src/hooks/hookSystem.ts @@ -220,6 +220,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 6448fac6327..30243b869f2 100644 --- a/packages/core/src/hooks/types.ts +++ b/packages/core/src/hooks/types.ts @@ -40,6 +40,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 @@ -968,6 +970,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/core/src/index.ts b/packages/core/src/index.ts index 3c6e776ec43..892b708b991 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'; diff --git a/packages/vscode-ide-companion/schemas/settings.schema.json b/packages/vscode-ide-companion/schemas/settings.schema.json index bef1cab2e46..25148c6d558 100644 --- a/packages/vscode-ide-companion/schemas/settings.schema.json +++ b/packages/vscode-ide-companion/schemas/settings.schema.json @@ -1661,6 +1661,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",