From 4cdc58ad4ceb800af11c308f1618a2e3d5e0f878 Mon Sep 17 00:00:00 2001 From: teknium1 <127238744+teknium1@users.noreply.github.com> Date: Thu, 24 Sep 2026 06:53:11 -0700 Subject: [PATCH] fix(desktop): Stop shows the same partial reply the session saves Stop seals the live bubble at the click and mutateStream drops every later message.delta (queued-but-unflushed ones too), but the agent keeps streaming until it honours the interrupt and persists every delta it delivered. The interrupted message.complete carrying that persisted partial was ignored, so state.db (and the next turn's context) held words the user never saw. The interrupted completion now extends the stopped bubble to the persisted partial, or adds a bubble when nothing had been painted. Extend-only: a shorter or different text never replaces what was shown. Refs #121594 --- .../e2e/core/interrupted-reply.spec.ts | 199 ++++++++++++++++++ .../gateway-event/message-stream.ts | 3 +- .../use-message-stream/gateway-event/types.ts | 3 +- .../session/hooks/use-message-stream/index.ts | 9 +- .../interrupted-reply.test.tsx | 135 ++++++++++++ .../use-message-stream/interrupted-reply.ts | 66 ++++++ 6 files changed, 411 insertions(+), 4 deletions(-) create mode 100644 apps/desktop/e2e/core/interrupted-reply.spec.ts create mode 100644 apps/desktop/src/app/session/hooks/use-message-stream/interrupted-reply.test.tsx create mode 100644 apps/desktop/src/app/session/hooks/use-message-stream/interrupted-reply.ts diff --git a/apps/desktop/e2e/core/interrupted-reply.spec.ts b/apps/desktop/e2e/core/interrupted-reply.spec.ts new file mode 100644 index 0000000000..ee718a00f2 --- /dev/null +++ b/apps/desktop/e2e/core/interrupted-reply.spec.ts @@ -0,0 +1,199 @@ +/** + * Stop mid-stream: the reply on screen is the reply the session saved (#121594). + * + * Stop seals the live bubble at the click and the renderer drops every later + * message.delta, but the agent keeps streaming until it honours the interrupt + * and persists everything it delivered — state.db and the model's next-turn + * context then hold words the user never saw. The interrupted message.complete + * carries that persisted partial; the bubble must extend to it. + * + * Each run streams numbered words (`wNNN`) from the scripted provider, presses + * the real Stop button mid-stream, waits for the interrupted completion on the + * wire and the assistant row in state.db, then compares the last word rendered + * with the last word persisted. Red on base: the bubble was 3-5 words short at + * the completion in every run; the first chat of a launch stayed short for + * good (later chats were repaired only by an unrelated transcript re-read + * ~0.4-0.7 s later). HERMES_E2E_INTERRUPT_RUNS / _OUT drive the N-run A/B. + */ + +import fs from 'node:fs' +import path from 'node:path' +import { DatabaseSync } from 'node:sqlite' + +import { expect, type Page, test } from '@playwright/test' + +import { + coreAppEnv, + createCoreSandbox, + currentSessionId, + launchCoreApp, + recordWebSockets, + send, + waitForInteractive, + writeProviderHome +} from './harness' +import { startScriptedProvider } from './provider' + +const nonce = Math.random() + .toString(36) + .slice(2, 8) + .replace(/[^a-z0-9]/g, 'x') + .padEnd(4, 'q') + +const U = (n: number) => `U${n}-${nonce}` +const A = (n: number) => `A${n}-${nonce}` +const RUNS = Number(process.env.HERMES_E2E_INTERRUPT_RUNS || 2) +const WORDS = 400 +const STOP_AFTER = 15 +const word = (i: number) => `w${String(i).padStart(3, '0')}` + +const lastWord = (text: string) => { + const all = text.match(/\bw(\d{3})\b/g) ?? [] + + return all.length ? Number(all.at(-1)!.slice(1)) : -1 +} + +function persistedReply(dbPath: string, marker: string): null | string { + if (!fs.existsSync(dbPath)) { + return null + } + + const db = new DatabaseSync(dbPath, { readOnly: true }) + + try { + const row = db + .prepare("SELECT content FROM messages WHERE role = 'assistant' AND content LIKE ? ORDER BY id DESC LIMIT 1") + .get(`%${marker}%`) as undefined | { content: string } + + return row?.content ?? null + } catch { + return null + } finally { + db.close() + } +} + +async function renderedReply(page: Page, marker: string): Promise { + return page.evaluate(marker => { + const viewport = document.querySelector('[data-slot="aui_thread-viewport"]') + const bubbles = [...(viewport?.querySelectorAll('[data-slot="aui_assistant-message-root"]') ?? [])] + + return bubbles + .map(el => (el as HTMLElement).innerText) + .filter(text => text.includes(marker)) + .join('\n') + }, marker) +} + +test('Stop mid-stream renders exactly the partial reply the session persisted', async () => { + test.setTimeout(120_000 + RUNS * 45_000) + const provider = await startScriptedProvider() + const sandbox = createCoreSandbox('interrupt') + writeProviderHome(sandbox.hermesHome, provider.url) + const dbPath = path.join(sandbox.hermesHome, 'state.db') + const { app, page } = await launchCoreApp(coreAppEnv(sandbox)) + const ws = recordWebSockets(page) + + const results: { + run: number + atStop: number + atComplete: number + healMs: number + wire: number + rendered: number + persisted: number + status: string + }[] = [] + + try { + await waitForInteractive(app, page) + + for (let run = 1; run <= RUNS; run++) { + if (run > 1) { + await page.evaluate(() => { + window.location.hash = '#/' + }) + await expect.poll(() => currentSessionId(page)).toBe('') + } + + provider.script(U(run), [{ text: [`${A(run)} `, ...Array.from({ length: WORDS }, (_, i) => `${word(i + 1)} `)] }]) + await send(page, `${U(run)} count`, 'Enter', ws) + await provider.streamStarted(U(run)) + await page.waitForFunction( + ([marker, needle]) => + [...document.querySelectorAll('[data-slot="aui_assistant-message-root"]')].some( + el => (el as HTMLElement).innerText.includes(marker) && (el as HTMLElement).innerText.includes(needle) + ), + [A(run), word(STOP_AFTER)] as const, + { polling: 'raf', timeout: 60_000 } + ) + await page + .locator('[data-slot="composer-root"] button[aria-label="Stop"]') + .filter({ visible: true }) + .first() + .click() + const atStop = lastWord(await renderedReply(page, A(run))) + + const findComplete = () => + ws.events.find(e => e.type === 'message.complete' && String(e.payload?.text ?? '').includes(A(run))) + + await expect + .poll(() => Boolean(findComplete()), { intervals: [20], message: `completion for ${U(run)}` }) + .toBe(true) + const completeAt = Date.now() + const complete = findComplete()! + const atComplete = lastWord(await renderedReply(page, A(run))) + + await expect.poll(() => persistedReply(dbPath, A(run)), { message: `assistant row for ${U(run)}` }).not.toBeNull() + const persisted = lastWord(persistedReply(dbPath, A(run))!) + // Bounded observation window: how long after the completion the screen + // first matches state.db (-1: never). A dropped tail never heals from + // the completion alone, so polling cannot mask it. + let rendered = -1 + let healMs = -1 + + await expect + .poll( + async () => { + rendered = lastWord(await renderedReply(page, A(run))) + healMs = rendered === persisted ? Date.now() - completeAt : -1 + + return rendered + }, + { timeout: 5_000, intervals: [20] } + ) + .toBe(persisted) + .catch(() => undefined) + + const streamed = ws.events + .filter(e => e.type === 'message.delta' && e.sessionId === complete.sessionId) + .map(e => String(e.payload?.text ?? '')) + .join('') + + results.push({ + run, + atStop, + atComplete, + healMs, + wire: lastWord(streamed.slice(streamed.lastIndexOf(A(run)))), + rendered, + persisted, + status: String(complete.payload?.status ?? '') + }) + } + } finally { + const out = process.env.HERMES_E2E_INTERRUPT_OUT + + if (out) { + fs.writeFileSync(out, JSON.stringify(results, null, 2)) + } + + await app.close().catch(() => undefined) + await provider.close() + sandbox.cleanup() + } + + // Stop must land mid-stream for the run to mean anything. + expect(results.every(r => r.status === 'interrupted' && r.persisted < WORDS)).toBe(true) + expect(results.filter(r => r.rendered !== r.persisted)).toEqual([]) +}) diff --git a/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event/message-stream.ts b/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event/message-stream.ts index 5628bb088d..202d76f024 100644 --- a/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event/message-stream.ts +++ b/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event/message-stream.ts @@ -372,7 +372,8 @@ export function handleMessageStreamEvent(ctx: GatewayEventContext): boolean { failure, occurredAt, payload?.persisted_turn, - Boolean(payload?.response_transformed) + Boolean(payload?.response_transformed), + typeof payload?.status === 'string' ? payload.status : undefined ) // Onboarding's first build: between turns is the only moment Setup may diff --git a/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event/types.ts b/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event/types.ts index 6b9803bb48..0bde5ab788 100644 --- a/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event/types.ts +++ b/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event/types.ts @@ -22,7 +22,8 @@ export interface GatewayEventDeps { failure?: { error: string; partial: boolean }, occurredAt?: number, persistedTurn?: PersistedTurn | null, - responseTransformed?: boolean + responseTransformed?: boolean, + status?: string ) => void failAssistantMessage: ( sessionId: string, diff --git a/apps/desktop/src/app/session/hooks/use-message-stream/index.ts b/apps/desktop/src/app/session/hooks/use-message-stream/index.ts index 7944ecdee5..0c0e71eae3 100644 --- a/apps/desktop/src/app/session/hooks/use-message-stream/index.ts +++ b/apps/desktop/src/app/session/hooks/use-message-stream/index.ts @@ -39,6 +39,7 @@ import type { ClientSessionState } from '../../../types' import { collapseDuplicateFinalAfterToolInterim, type DuplicateFinalCollapse } from './collapse-duplicate-final' import { useGatewayEventHandler } from './gateway-event' import { handleServerRequest as dispatchServerRequest } from './gateway-event/server-requests' +import { extendInterruptedReply } from './interrupted-reply' import { currentResponseParts, mergeCurrentResponseText } from './response-parts' import { completionErrorText, delegateTaskPayloads, MAX_STREAM_FLUSH_GAP_MS, STREAM_DELTA_FLUSH_MS } from './utils' @@ -625,7 +626,8 @@ export function useMessageStream({ failure?: { error: string; partial: boolean; surface?: ErrorSurface | null }, occurredAt = Date.now() / 1000, persistedTurn?: PersistedTurn | null, - responseTransformed?: boolean + responseTransformed?: boolean, + status?: string ) => { let shouldHydrate = false @@ -633,10 +635,13 @@ export function useMessageStream({ // Late completion from an already-cancelled turn: cancelRun has // already finalized the bubble (kept the partial text, dropped it if // empty). Re-running the dedupe below would replace the partial with - // the just-cancelled full text, so we settle and bail instead. + // the just-cancelled full text, so we settle and bail instead — only + // extending the bubble to the partial the agent persisted (#121594). if (state.interrupted) { return { ...state, + messages: + status === 'interrupted' ? extendInterruptedReply(state.messages, text, occurredAt) : state.messages, awaitingResponse: false, busy: false, needsInput: false, diff --git a/apps/desktop/src/app/session/hooks/use-message-stream/interrupted-reply.test.tsx b/apps/desktop/src/app/session/hooks/use-message-stream/interrupted-reply.test.tsx new file mode 100644 index 0000000000..349306fc82 --- /dev/null +++ b/apps/desktop/src/app/session/hooks/use-message-stream/interrupted-reply.test.tsx @@ -0,0 +1,135 @@ +import type { GatewayEvent } from '@hermes/shared' +// #121594: Stop seals the live bubble at the click and drops every later +// message.delta, but the agent keeps streaming until it honours the interrupt +// and persists everything it delivered (state.db and the next turn's context). +// Its interrupted message.complete carries that persisted partial. Contract: +// the screen shows exactly what the session saved — extend-only, never +// shortened or rewritten. +import { act, cleanup } from '@testing-library/react' +import { afterEach, describe, expect, it, vi } from 'vitest' + +import { finalizeUserInterruptedMessages } from '@/app/session/hooks/use-prompt-actions/rewind' +import { chatMessageText } from '@/lib/chat-messages' + +import { type MessageStreamHarness, renderMessageStream } from './test-harness' +import { STREAM_DELTA_FLUSH_MS } from './utils' + +const SID = 'interrupted-reply-session' + +let stream: MessageStreamHarness + +async function mountHarness() { + vi.useFakeTimers() + stream = renderMessageStream(SID) + await act(async () => { + await Promise.resolve() + }) +} + +const flushDeltas = async () => { + await act(async () => { + await vi.advanceTimersByTimeAsync(STREAM_DELTA_FLUSH_MS) + }) +} + +const emit = (event: GatewayEvent) => act(() => stream.handleEvent(event)) + +// The state transform both Stop paths (use-prompt-actions and the session +// tile) apply at the click. +const pressStop = () => + act(() => { + const state = stream.state() + stream.states.set(SID, { + ...state, + messages: finalizeUserInterruptedMessages(state.messages, state.streamId), + busy: false, + awaitingResponse: false, + streamId: null, + pendingBranchGroup: null, + needsInput: false, + interrupted: true, + turnStartedAt: null, + turnLive: false + }) + }) + +const assistantTexts = () => + stream + .state() + .messages.filter(message => message.role === 'assistant' && !message.hidden) + .map(message => chatMessageText(message)) + +describe('Stop: the screen shows the partial the session saves (#121594)', () => { + afterEach(() => { + cleanup() + vi.useRealTimers() + vi.restoreAllMocks() + }) + + it('late deltas after Stop reach the bubble through the interrupted completion', async () => { + await mountHarness() + emit({ payload: {}, session_id: SID, type: 'message.start' }) + emit({ payload: { text: 'The quick brown ' }, session_id: SID, type: 'message.delta' }) + await flushDeltas() + + pressStop() + // In flight when Stop landed: dropped by the renderer, persisted by the agent. + emit({ payload: { text: 'fox jumps ' }, session_id: SID, type: 'message.delta' }) + await flushDeltas() + expect(assistantTexts()).toEqual(['The quick brown ']) + + emit({ + payload: { status: 'interrupted', text: 'The quick brown fox jumps' }, + session_id: SID, + type: 'message.complete' + }) + + expect(assistantTexts()).toEqual(['The quick brown fox jumps']) + expect(stream.state().busy).toBe(false) + }) + + it('adds the bubble when Stop landed before anything was painted', async () => { + await mountHarness() + emit({ payload: {}, session_id: SID, type: 'message.start' }) + // Queued for the next flush, not yet painted. + emit({ payload: { text: 'Hello' }, session_id: SID, type: 'message.delta' }) + pressStop() + await flushDeltas() + expect(assistantTexts()).toEqual([]) + + emit({ payload: { status: 'interrupted', text: 'Hello there' }, session_id: SID, type: 'message.complete' }) + + expect(assistantTexts()).toEqual(['Hello there']) + }) + + it('never shortens or rewrites the shown partial', async () => { + await mountHarness() + emit({ payload: {}, session_id: SID, type: 'message.start' }) + emit({ payload: { text: 'Alpha beta gamma' }, session_id: SID, type: 'message.delta' }) + await flushDeltas() + pressStop() + + emit({ payload: { status: 'interrupted', text: 'Alpha beta' }, session_id: SID, type: 'message.complete' }) + expect(assistantTexts()).toEqual(['Alpha beta gamma']) + + emit({ payload: { status: 'interrupted', text: 'Something else' }, session_id: SID, type: 'message.complete' }) + expect(assistantTexts()).toEqual(['Alpha beta gamma']) + }) + + it('does not repeat the pre-tool text when Stop lands during a tool', async () => { + await mountHarness() + emit({ payload: {}, session_id: SID, type: 'message.start' }) + emit({ payload: { text: 'Let me check.' }, session_id: SID, type: 'message.delta' }) + await flushDeltas() + emit({ + payload: { args: { command: 'ls' }, name: 'terminal', tool_id: 'call-1' }, + session_id: SID, + type: 'tool.start' + }) + pressStop() + + emit({ payload: { status: 'interrupted', text: 'Let me check.' }, session_id: SID, type: 'message.complete' }) + + expect(assistantTexts()).toEqual(['Let me check.']) + }) +}) diff --git a/apps/desktop/src/app/session/hooks/use-message-stream/interrupted-reply.ts b/apps/desktop/src/app/session/hooks/use-message-stream/interrupted-reply.ts new file mode 100644 index 0000000000..e2af0431a7 --- /dev/null +++ b/apps/desktop/src/app/session/hooks/use-message-stream/interrupted-reply.ts @@ -0,0 +1,66 @@ +import { + assistantTextPart, + type ChatMessage, + chatMessageText, + mergeFinalAssistantText, + renderMediaTags +} from '@/lib/chat-messages' +import { generatedImageEchoSources, stripGeneratedImageEchoes } from '@/lib/generated-images' + +const flat = (text: string) => text.replace(/\s+/g, ' ').trim() + +/** + * Stop seals the live bubble at the click and drops every later delta, but the + * agent keeps streaming until it honours the interrupt and persists all of it + * (state.db, the next turn's context). Its interrupted `message.complete` + * carries that persisted partial (#121594). When it extends what the turn's + * bubble shows, the bubble takes it; when Stop landed before anything was + * painted, a bubble is added. Extend-only: a shorter or different text never + * replaces what the user saw. + */ +export function extendInterruptedReply(messages: ChatMessage[], rawText: string, occurredAt: number): ChatMessage[] { + const text = renderMediaTags(rawText).trim() + + if (!text) { + return messages + } + + const lastUserIndex = messages.findLastIndex(message => message.role === 'user') + + const turn = messages.filter( + (message, index) => index > lastUserIndex && message.role === 'assistant' && !message.hidden + ) + + // Stop keeps a painted live bubble as a settled, non-interim row; an + // interim tail means the live segment had nothing painted and was dropped. + const target = turn.at(-1) + + if (target && !target.interim) { + const visible = stripGeneratedImageEchoes(text, generatedImageEchoSources(target.parts)).trim() + const parts = mergeFinalAssistantText(target.parts, visible, occurredAt) + const shown = flat(chatMessageText(target)) + const persisted = flat(chatMessageText({ ...target, parts })) + + if (persisted.length <= shown.length || !persisted.startsWith(shown)) { + return messages + } + + return messages.map(message => (message === target ? { ...message, parts } : message)) + } + + if (turn.some(message => flat(chatMessageText(message)).includes(flat(text)))) { + return messages + } + + return [ + ...messages, + { + id: `assistant-interrupted-${Date.now()}`, + role: 'assistant', + parts: [{ ...assistantTextPart(text, occurredAt), completedAt: occurredAt }], + timestamp: occurredAt, + completedAt: occurredAt, + pending: false + } + ] +}