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 + } + ] +}