diff --git a/agent/context_references.py b/agent/context_references.py index dea5fd4c3a5e..8981aa472f50 100644 --- a/agent/context_references.py +++ b/agent/context_references.py @@ -19,6 +19,7 @@ REFERENCE_PATTERN = re.compile( rf"(?diff|staged)\b|(?Pfile|folder|git|url):(?P{_QUOTED_REFERENCE_VALUE}(?::\d+(?:-\d+)?)?|\S+))" ) TRAILING_PUNCTUATION = ",.;!?" +_NEEDS_QUOTING = re.compile(r"""[\s()\[\]{}<>"'`]""") _SENSITIVE_HOME_DIRS = (".ssh", ".aws", ".gnupg", ".kube", ".docker", ".azure", ".config/gh") _SENSITIVE_HERMES_DIRS = (Path("skills") / ".hub",) _SENSITIVE_HOME_FILES = ( @@ -60,6 +61,21 @@ class ContextReferenceResult: blocked: bool = False +def format_reference_value(value: str) -> str: + """Quote a reference value so ``REFERENCE_PATTERN`` reads it back whole. + + The unquoted alternative in the pattern is ``\\S+``, so a path containing a + space parses as a truncated ref with the tail left behind as loose text. + Mirrors ``formatRefValue`` in the desktop's directive-text.tsx. + """ + if not _NEEDS_QUOTING.search(value): + return value + for quote in ("`", '"', "'"): + if quote not in value: + return f"{quote}{value}{quote}" + return value + + def parse_context_references(message: str) -> list[ContextReference]: refs: list[ContextReference] = [] if not message: diff --git a/apps/desktop/e2e/image-attachment-resume.spec.ts b/apps/desktop/e2e/image-attachment-resume.spec.ts new file mode 100644 index 000000000000..cfa9829a4e2e --- /dev/null +++ b/apps/desktop/e2e/image-attachment-resume.spec.ts @@ -0,0 +1,172 @@ +/** + * Regression coverage for an attached image in a durable session. The gateway + * persists the turn, the builder exits, and desktop renders it from SessionDB + * for the first time — the "quit and relaunch" case, where the transcript used + * to come back as vision-enrichment prose instead of a thumbnail. + * + * The fixture pins `image_input_mode: native` because that is the majority + * routing path (any vision-capable model) and the one where a text-only + * persist override is silently dropped. The image also sits behind directory + * and file names containing spaces, mirroring the macOS composer's + * `~/Library/Application Support/...` staging path. + */ + +import * as fs from 'node:fs' +import * as path from 'node:path' + +import { + buildAppEnv, + createSandbox, + launchDesktop, + type Sandbox, + waitForAppReady, + writeEnvFile, + writeMockProviderConfig, +} from './fixtures' +import { type MockServer, startMockServer } from './mock-server' +import { RealSessionBuilder } from './real-session-builder' +import { type ElectronApplication, expect, type Page, test } from './test' + +// A seeded session has no generated title, so every label falls back to the +// session preview — the first 60 characters of the first user message. +const SESSION_TITLE = 'E2E attached image session' +const CAPTION = 'E2E attached image must survive a relaunch' +const IMAGE_DIR = 'Application Support/e2e shots' +const IMAGE_NAME = 'e2e capture.png' +const NATIVE_IMAGE_CONFIG = 'agent:\n image_input_mode: native' + +/** A 160x100 framed magenta block — small, but visible in the screenshots. */ +const PNG_BASE64 = + 'iVBORw0KGgoAAAANSUhEUgAAAKAAAABkCAIAAACO1KzYAAAA30lEQVR42u3dwQ2AIBAAQTAWAx1iBXYI7diCuWhEMvP2dZsj+CL30hLr2oxAYARGYARGYARGYIERGIH53n7nozpOk5rQqIcNdkQjMAIjMNPeomP3N54V+5exwY5oBEZgBEZgBEZggREYgREYgREYgQVGYARGYARGYARGYIERGIERGIERGIEFRmAERmAERmAEFhiBERiBERiBERiBBUZgBEZgBEZgBBYYgREYgREYgRFYYCMQGIERmPSj94Njb9ligxEYgRFYYJaQe2mmYIMRGIERGIERGIEFRmAERmDedAFtjAtAGWDnoAAAAABJRU5ErkJggg==' + +interface SeededFixture { + app: ElectronApplication + mock: MockServer + page: Page + sandbox: Sandbox + cleanup: () => Promise +} + +function writeImage(sandbox: Sandbox): string { + const dir = path.join(sandbox.root, IMAGE_DIR) + fs.mkdirSync(dir, { recursive: true }) + + const imagePath = path.join(dir, IMAGE_NAME) + fs.writeFileSync(imagePath, Buffer.from(PNG_BASE64, 'base64')) + + return imagePath +} + +async function setupSeededDesktop(): Promise { + const mock = await startMockServer() + const sandbox = createSandbox('image-attachment') + writeMockProviderConfig(sandbox.hermesHome, mock.url, undefined, NATIVE_IMAGE_CONFIG) + writeEnvFile(sandbox.hermesHome) + + const builder = await RealSessionBuilder.start(sandbox.hermesHome) + + try { + await builder.createSession({ + title: SESSION_TITLE, + turns: [{ images: [writeImage(sandbox)], text: CAPTION }], + }) + } finally { + await builder.close() + } + + const { app, page } = await launchDesktop(buildAppEnv(sandbox)) + + return { + app, + mock, + page, + sandbox, + cleanup: async () => { + await app.close().catch(() => undefined) + await mock.close() + sandbox.cleanup() + }, + } +} + +function sessionRow(page: Page) { + return page.locator('[data-slot="sidebar"] button').filter({ hasText: CAPTION }).first() +} + +async function openSeededSession(page: Page): Promise { + const row = sessionRow(page) + await row.waitFor({ state: 'visible', timeout: 60_000 }) + await row.click() + await page.waitForFunction( + expected => (document.querySelector('[data-slot="aui_thread-viewport"]')?.textContent ?? '').includes(expected), + CAPTION, + { timeout: 30_000 }, + ) +} + +async function openNewSession(page: Page): Promise { + await page.locator('[data-slot="sidebar"] button[aria-label="New session"]').first().click() + await page.waitForFunction( + expected => !(document.querySelector('[data-slot="aui_thread-viewport"]')?.textContent ?? '').includes(expected), + CAPTION, + { timeout: 15_000 }, + ) +} + +async function transcriptText(page: Page): Promise { + return page.evaluate(() => document.querySelector('[data-slot="aui_thread-viewport"]')?.textContent ?? '') +} + +async function assertRendersThumbnail(page: Page, label: string): Promise { + const thumbnail = page.locator('[data-slot="aui_directive-image"] img') + await expect(thumbnail, `${label}: the attachment should render as an image`).toHaveCount(1) + await expect(thumbnail, `${label}: the thumbnail should resolve off disk`).toHaveAttribute('src', /^data:image\//) + + const text = await transcriptText(page) + expect(text, `${label}: the caption should survive alongside the image`).toContain(CAPTION) + // A broken ref falls back to a chip whose label leaks the path, and a + // flattened multimodal turn leaves the agent's placeholder behind. + expect(text, `${label}: the raw image path should not leak into the transcript`).not.toContain(IMAGE_NAME) + expect(text, `${label}: the image directive should not render literally`).not.toContain('@image:') + expect(text, `${label}: the flattening placeholder should not render`).not.toContain('[screenshot]') +} + +test.describe('attached image resume', () => { + let fixture: SeededFixture | null = null + + test.afterEach(async () => { + await fixture?.cleanup() + fixture = null + }) + + test('renders a persisted attachment as a thumbnail on first open and after a cold reload', async ({}, testInfo) => { + // Seeding through the real gateway plus two full app boots does not fit the + // default per-test budget on a cold runner. + test.slow() + + fixture = await setupSeededDesktop() + await waitForAppReady(fixture, 120_000) + + // The sidebar labels a session by its preview, so the caption has to lead + // the persisted turn — a leading directive reads as a truncated file path. + const row = sessionRow(fixture.page) + await row.waitFor({ state: 'visible', timeout: 60_000 }) + + const label = (await row.textContent())?.trim() ?? '' + expect(label.startsWith(CAPTION), `sidebar label should open with the caption: ${label}`).toBe(true) + + await openSeededSession(fixture.page) + await assertRendersThumbnail(fixture.page, 'first open') + await fixture.page.screenshot({ path: testInfo.outputPath('attachment-first-open.png') }) + + // A reload drops every cached attachment ref, so the transcript has to come + // back from the persisted turn alone. + await fixture.page.reload() + await waitForAppReady(fixture, 120_000) + await openNewSession(fixture.page) + + await openSeededSession(fixture.page) + await assertRendersThumbnail(fixture.page, 'cold reload') + await fixture.page.screenshot({ path: testInfo.outputPath('attachment-cold-reload.png') }) + }) +}) diff --git a/apps/desktop/e2e/real-session-builder.ts b/apps/desktop/e2e/real-session-builder.ts index 488291695b56..0f670e9698bf 100644 --- a/apps/desktop/e2e/real-session-builder.ts +++ b/apps/desktop/e2e/real-session-builder.ts @@ -28,11 +28,18 @@ interface CreatedSession { stored_session_id: string } +export interface RealSessionTurn { + /** Local image paths attached before the prompt, as the composer would. */ + images?: readonly string[] + text: string +} + export interface RealSessionSpec { - /** Human-visible sidebar title, persisted by the first completed turn. */ + /** Session label. The durable row stores no title, so clients fall back to + * the preview (the first 60 characters of the first user message). */ title: string /** Each item becomes one real user prompt followed by the mock provider's reply. */ - turns: readonly string[] + turns: readonly (RealSessionTurn | string)[] } export interface RealSession { @@ -107,7 +114,13 @@ export class RealSessionBuilder { const runtimeId = requireString(created, 'session_id') const sessionId = requireString(created, 'stored_session_id') - for (const text of spec.turns) { + for (const turn of spec.turns) { + const { images = [], text } = typeof turn === 'string' ? { text: turn } : turn + + for (const image of images) { + await this.request('image.attach', { session_id: runtimeId, path: image }) + } + const completion = this.waitForEvent( frame => frame.params?.type === 'message.complete' && frame.params.session_id === runtimeId, ) diff --git a/apps/desktop/src/app/session/hooks/use-session-actions.test.tsx b/apps/desktop/src/app/session/hooks/use-session-actions.test.tsx index 2472dccb1132..603b45a1a2bd 100644 --- a/apps/desktop/src/app/session/hooks/use-session-actions.test.tsx +++ b/apps/desktop/src/app/session/hooks/use-session-actions.test.tsx @@ -1307,6 +1307,76 @@ describe('resumeSession warm-cache mapping integrity', () => { expect(runtimeIdByStoredSessionIdRef.current.get('stored-A')).toBe('rt-A') }) + it('preserves cached image attachments through an idle persisted transcript refresh', async () => { + const runtimeIdByStoredSessionIdRef: MutableRefObject> = { + current: new Map([['stored-A', 'rt-A']]) + } + + const state = clientState('stored-A') + state.messages = [ + { + id: 'cached-user', + role: 'user', + parts: [{ type: 'text', text: 'describe this image' }], + attachmentRefs: ['@image:/tmp/photo.png'] + }, + { + id: 'cached-assistant', + role: 'assistant', + parts: [{ type: 'text', text: 'It is a photo.' }] + } + ] + + const sessionStateByRuntimeIdRef: MutableRefObject> = { + current: new Map([['rt-A', state]]) + } + + const persistedMessages = [ + { content: 'describe this image', role: 'user', timestamp: 1 }, + { content: 'It is a photo.', role: 'assistant', timestamp: 2 } + ] + + vi.mocked(getSessionMessages).mockResolvedValue({ + messages: persistedMessages, + session_id: 'stored-A' + } as never) + + const requestGateway = vi.fn(async (method: string) => { + if (method === 'session.activate') { + return { + session_id: 'rt-A', + session_key: 'stored-A', + resumed: 'stored-A', + message_count: persistedMessages.length, + messages: persistedMessages, + running: false, + info: {} + } as never + } + + return {} as never + }) + + let resumedState: ClientSessionState | undefined + let resume: ((storedSessionId: string, replaceRoute?: boolean) => Promise) | null = null + + render( + (resume = ready)} + onStateUpdate={(_sessionId, next) => (resumedState = next)} + requestGateway={requestGateway} + runtimeIdByStoredSessionIdRef={runtimeIdByStoredSessionIdRef} + sessionStateByRuntimeIdRef={sessionStateByRuntimeIdRef} + /> + ) + await waitFor(() => expect(resume).not.toBeNull()) + await resume!('stored-A', true) + + expect(requestGateway.mock.calls.map(([method]) => method)).toContain('session.activate') + expect(getSessionMessages).toHaveBeenCalledWith('stored-A', undefined) + expect(resumedState?.messages[0]?.attachmentRefs).toEqual(['@image:/tmp/photo.png']) + }) + it('repairs an idle warm cache from a divergent equal-length persisted transcript', async () => { const runtimeIdByStoredSessionIdRef: MutableRefObject> = { current: new Map([['stored-A', 'rt-A']]) diff --git a/apps/desktop/src/app/session/hooks/use-session-actions/utils.test.ts b/apps/desktop/src/app/session/hooks/use-session-actions/utils.test.ts index 277a9c96b642..6da126b99334 100644 --- a/apps/desktop/src/app/session/hooks/use-session-actions/utils.test.ts +++ b/apps/desktop/src/app/session/hooks/use-session-actions/utils.test.ts @@ -317,6 +317,52 @@ describe('reconcileResumeMessages', () => { const [out] = reconcileResumeMessages(next, previous) expect(out.parts.some(p => p.type === 'reasoning')).toBe(true) }) + + it('preserves attachment refs for a matching user turn', () => { + const next = [msg('stored-user', 'user', 'describe this image')] + + const previous = [ + msg('live-user', 'user', 'describe this image', { + attachmentRefs: ['@image:/tmp/photo.png'] + }) + ] + + const [out] = reconcileResumeMessages(next, previous) + + expect(out.attachmentRefs).toEqual(['@image:/tmp/photo.png']) + }) + + it('does not overwrite attachment refs already present on the resumed message', () => { + const next = [ + msg('stored-user', 'user', 'describe this image', { + attachmentRefs: ['@image:/tmp/authoritative.png'] + }) + ] + + const previous = [ + msg('live-user', 'user', 'describe this image', { + attachmentRefs: ['@image:/tmp/cached.png'] + }) + ] + + const [out] = reconcileResumeMessages(next, previous) + + expect(out.attachmentRefs).toEqual(['@image:/tmp/authoritative.png']) + }) + + it('does not preserve attachment refs when the user text differs', () => { + const next = [msg('stored-user', 'user', 'a different prompt')] + + const previous = [ + msg('live-user', 'user', 'describe this image', { + attachmentRefs: ['@image:/tmp/photo.png'] + }) + ] + + const [out] = reconcileResumeMessages(next, previous) + + expect(out.attachmentRefs).toBeUndefined() + }) }) describe('preserveLocalPendingTurnMessages', () => { @@ -440,9 +486,91 @@ describe('preserveLocalPendingTurnMessages', () => { 'user-optimistic' ]) }) + + // #70720: the gateway persists an attached image as a leading `@image:` + // directive line, while the local optimistic composer keeps it as separate + // `attachmentRefs`. A naive text compare (chatMessageText a === b) therefore + // always mismatched whenever an image was attached and re-appended the + // optimistic row as a distinct, duplicate user bubble. Both sides must now + // reduce to the same visible text via textWithoutImageRefs. + it('does not duplicate the optimistic image turn when the persisted turn carries @image refs', () => { + const previous = [ + msg('1-user', 'user', 'first'), + msg('2-assistant', 'assistant', 'first answer'), + msg('user-optimistic', 'user', 'what is in this photo?', { + attachmentRefs: ['@image:/tmp/cat.png'] + }) + ] + + const next = [ + msg('1-user-stored', 'user', 'first'), + msg('2-assistant-stored', 'assistant', 'first answer'), + msg('3-user-stored', 'user', '@image:/tmp/cat.png\nwhat is in this photo?') + ] + + expect(preserveLocalPendingTurnMessages(next, previous)).toBe(next) + }) + + it('still keeps a genuinely uncommitted optimistic image turn when the persisted text differs', () => { + const previous = [ + msg('1-user', 'user', 'first'), + msg('2-assistant', 'assistant', 'first answer'), + msg('user-optimistic', 'user', 'a different caption', { + attachmentRefs: ['@image:/tmp/cat.png'] + }) + ] + + // Persisted turn has a different caption — the optimistic row is still + // uncommitted and must survive (image-aware compare must not over-correct). + const next = [ + msg('1-user-stored', 'user', 'first'), + msg('2-assistant-stored', 'assistant', 'first answer'), + msg('3-user-stored', 'user', '@image:/tmp/cat.png\nwhat is in this photo?') + ] + + expect(preserveLocalPendingTurnMessages(next, previous).map(message => message.id)).toEqual([ + '1-user-stored', + '2-assistant-stored', + '3-user-stored', + 'user-optimistic' + ]) + }) }) describe('appendLiveSessionProjection', () => { + it('does not duplicate the inflight user when the persisted turn carries @image refs', () => { + // By the time a stored transcript reaches appendLiveSessionProjection it + // has already been run through toChatMessages, so the @image directive has + // been lifted into attachmentRefs and the visible text is the bare caption. + const stored = [ + msg('stored-user', 'user', 'current running prompt', { + attachmentRefs: ['@image:/tmp/cat.png'] + }), + msg('stored-assistant', 'assistant', 'earlier answer') + ] + + const restored = appendLiveSessionProjection(stored, { + session_id: 'runtime-1', + inflight: { + user: 'current running prompt', + assistant: 'partial answer', + streaming: true + } + }) + + // The persisted user already carries the same visible text (the attachment + // lives in attachmentRefs on both sides), so the inflight *user* projection + // must be suppressed — exactly one user row, no duplicated bubble stacked + // on top of the persisted one. The live assistant tail is still projected. + const userRows = restored.filter(message => message.role === 'user') + expect(userRows).toHaveLength(1) + expect(userRows[0].id).toBe('stored-user') + const userText = userRows[0].parts + .map(part => ('text' in part ? part.text : '')) + .join('') + expect(userText).toBe('current running prompt') + }) + it('restores the running turn and accepted queued prompt after a renderer restart', () => { const stored = [msg('stored-user', 'user', 'earlier'), msg('stored-assistant', 'assistant', 'earlier answer')] diff --git a/apps/desktop/src/app/session/hooks/use-session-actions/utils.ts b/apps/desktop/src/app/session/hooks/use-session-actions/utils.ts index 553f473b9948..bc5cb6cec3b5 100644 --- a/apps/desktop/src/app/session/hooks/use-session-actions/utils.ts +++ b/apps/desktop/src/app/session/hooks/use-session-actions/utils.ts @@ -1,7 +1,7 @@ import { getSession } from '@/hermes' import { assistantTextPart, type ChatMessage, chatMessageText, textPart } from '@/lib/chat-messages' import { normalizePersonalityValue } from '@/lib/chat-runtime' -import { embeddedImageUrls, textWithoutEmbeddedImages } from '@/lib/embedded-images' +import { embeddedImageUrls, textWithoutEmbeddedImages, textWithoutImageRefs } from '@/lib/embedded-images' import { reconcileApprovalModeForProfile } from '@/store/approval-mode' import { requestDesktopOnboardingForCredentialWarning } from '@/store/onboarding' import { $activeGatewayProfile, $profiles, normalizeProfileKey } from '@/store/profile' @@ -211,6 +211,10 @@ export function reconcileResumeMessages(nextMessages: ChatMessage[], previousMes if (nextText === previousVisibleText || nextText === previousText.trim()) { preserved = preserveReasoningParts(preserved, previous) + + if (message.role === 'user' && preserved.attachmentRefs === undefined && previous.attachmentRefs?.length) { + preserved = { ...preserved, attachmentRefs: [...previous.attachmentRefs] } + } } const previousImages = embeddedImageUrls(previousText) @@ -306,7 +310,7 @@ export function preserveLocalPendingTurnMessages( if ( isOptimisticUser && latestAuthoritativeUser && - chatMessageText(latestAuthoritativeUser).trim() === chatMessageText(message).trim() + textWithoutImageRefs(chatMessageText(latestAuthoritativeUser)) === textWithoutImageRefs(chatMessageText(message)) ) { continue } @@ -318,7 +322,7 @@ export function preserveLocalPendingTurnMessages( continue } - if (chatMessageText(authoritative).trim() === chatMessageText(message).trim()) { + if (textWithoutImageRefs(chatMessageText(authoritative)) === textWithoutImageRefs(chatMessageText(message))) { continue } } @@ -359,7 +363,8 @@ export function appendLiveSessionProjection( // Only suppress the projection when the latest authoritative user row is the // same turn — older identical prompts must not hide a newly accepted repeat. const latestUser = [...messages].reverse().find(message => message.role === 'user') - const inflightUserAlreadyPersisted = latestUser && chatMessageText(latestUser).trim() === inflightUser + const inflightUserAlreadyPersisted = + latestUser && textWithoutImageRefs(chatMessageText(latestUser)) === textWithoutImageRefs(inflightUser) if (inflightUser && !inflightUserAlreadyPersisted) { projected.push({ diff --git a/apps/desktop/src/components/assistant-ui/directive-text.tsx b/apps/desktop/src/components/assistant-ui/directive-text.tsx index fe3c8e7f1657..ffa63a7b8877 100644 --- a/apps/desktop/src/components/assistant-ui/directive-text.tsx +++ b/apps/desktop/src/components/assistant-ui/directive-text.tsx @@ -351,19 +351,32 @@ export function DirectiveContent({ text }: { text: string }) { const { cleanedText, images } = useMemo(() => safeEmbeddedImages(text ?? ''), [text]) const segments = useMemo(() => safeDirectiveSegments(cleanedText), [cleanedText]) + // `@image:` directives render as a block-level thumbnail row (like + // embedded base64 images below), not inline mid-text — otherwise a large + // thumbnail gets wedged between words and breaks the text's line flow. + const imageSegments = useMemo( + () => + segments.filter( + (segment): segment is Extract => + segment.kind === 'mention' && segment.type === 'image' + ), + [segments] + ) + return ( {segments.map((segment, index) => segment.kind === 'text' ? ( {segment.text} - ) : segment.type === 'image' ? ( - - ) : ( + ) : segment.type === 'image' ? null : ( ) )} - {images.length > 0 && ( + {(imageSegments.length > 0 || images.length > 0) && ( + {imageSegments.map((segment, index) => ( + + ))} {images.map((src, index) => ( = ({ id, label }) => { return ( { expect(chatMessageText(message)).toBe('Here you go.') }) + it('lifts @image directive lines into attachmentRefs instead of inline text', () => { + const [message] = toChatMessages([ + { + role: 'user', + content: '@image:/tmp/cat.png\nwhat is in this photo?', + timestamp: 1 + } + ]) + + expect(chatMessageText(message)).toBe('what is in this photo?') + expect((message as { attachmentRefs?: string[] }).attachmentRefs).toEqual(['@image:/tmp/cat.png']) + }) + + it('keeps a user turn that carried only an attached image (no caption)', () => { + const [message] = toChatMessages([ + { + role: 'user', + content: '@image:/tmp/cat.png', + timestamp: 1 + } + ]) + + // The bubble has no visible text, but must survive the empty-turn filter + // because it carries attachment refs — otherwise a stand-alone attachment + // vanishes from the transcript after a session switch / restart. + expect(chatMessageText(message)).toBe('') + expect((message as { attachmentRefs?: string[] }).attachmentRefs).toEqual(['@image:/tmp/cat.png']) + }) + + it('renders a native-vision turn as caption plus thumbnail, not raw placeholder text', () => { + // How a turn sent to a natively-vision-capable model comes back out of the + // session store: a backtick-quoted ref (the path has spaces) and the + // `[screenshot]` stand-in left by flattening the parts list. + const ref = '@image:`/Users/me/Library/Application Support/Hermes/composer-images/a.png`' + + const [message] = toChatMessages([ + { + role: 'user', + content: `${ref}\nwhat is in this photo?\n[screenshot]`, + timestamp: 1 + } + ]) + + expect(chatMessageText(message)).toBe('what is in this photo?') + expect((message as { attachmentRefs?: string[] }).attachmentRefs).toEqual([ref]) + }) + + it('leaves a plain user prompt without attachment refs untouched', () => { + const [message] = toChatMessages([ + { + role: 'user', + content: 'just a question', + timestamp: 1 + } + ]) + + expect((message as { attachmentRefs?: string[] }).attachmentRefs).toBeUndefined() + }) + it('coerces non-string message content without throwing', () => { const [message] = toChatMessages([ { diff --git a/apps/desktop/src/lib/chat-messages.ts b/apps/desktop/src/lib/chat-messages.ts index 0019c3d4b5ef..8c01d3214b36 100644 --- a/apps/desktop/src/lib/chat-messages.ts +++ b/apps/desktop/src/lib/chat-messages.ts @@ -1,6 +1,7 @@ import type { ThreadMessageLike } from '@assistant-ui/react' import type { BillingBlock } from '@hermes/shared' +import { extractImageRefs } from '@/lib/embedded-images' import { dedupeGeneratedImageEchoesInParts } from '@/lib/generated-images' import { mediaDisplayLabel, mediaMarkdownHref } from '@/lib/media' import { normalize } from '@/lib/text' @@ -931,7 +932,7 @@ export function toChatMessages(messages: SessionMessage[]): ChatMessage[] { const content = message.content || message.text || message.context || message.name - const displayContent = transcriptContent( + const rawDisplayContent = transcriptContent( message.display_kind, timelineDisplayContent(message, displayContentForMessage(message.role, content)) ) @@ -941,6 +942,17 @@ export function toChatMessages(messages: SessionMessage[]): ChatMessage[] { ? 'system' : message.role + // Persisted user turns carry `@image:` directive lines inline in + // the text (see tui_gateway/server.py's persist-time rewrite). The + // read-only bubble clamps its body to ~2 lines, and a large inline image + // thumbnail pushes any caption text below the clamp's visible area — so + // pull image refs out into `attachmentRefs` (same shape the local + // optimistic composer already uses) and render them via the dedicated + // attachments row below the bubble instead. + const imageRefExtraction = displayRole === 'user' && rawDisplayContent ? extractImageRefs(rawDisplayContent) : null + const displayContent = imageRefExtraction ? imageRefExtraction.cleanedText : rawDisplayContent + const extractedAttachmentRefs = imageRefExtraction?.refs.length ? imageRefExtraction.refs : undefined + const parts: ChatMessagePart[] = [] const reasoning = @@ -960,7 +972,7 @@ export function toChatMessages(messages: SessionMessage[]): ChatMessage[] { parts.push(...message.tool_calls.map((call, callIndex) => toolPartFromStoredCall(call, callIndex))) } - if (!parts.length) { + if (!parts.length && !extractedAttachmentRefs?.length) { if (message.role !== 'assistant') { flushPendingTools(index) activeAssistantIndex = null @@ -1010,7 +1022,8 @@ export function toChatMessages(messages: SessionMessage[]): ChatMessage[] { id: `${message.timestamp || Date.now()}-${index}-${displayRole}`, role: displayRole, parts, - timestamp: message.timestamp + timestamp: message.timestamp, + ...(extractedAttachmentRefs ? { attachmentRefs: extractedAttachmentRefs } : {}) }) activeAssistantIndex = message.role === 'assistant' ? result.length - 1 : null @@ -1022,7 +1035,9 @@ export function toChatMessages(messages: SessionMessage[]): ChatMessage[] { ) return withUniqueToolCallIds( - withoutGeneratedImageEchoes.filter(m => chatMessageText(m).trim() || m.parts.some(part => part.type !== 'text')) + withoutGeneratedImageEchoes.filter( + m => chatMessageText(m).trim() || m.parts.some(part => part.type !== 'text') || m.attachmentRefs?.length + ) ) } diff --git a/apps/desktop/src/lib/embedded-images.test.ts b/apps/desktop/src/lib/embedded-images.test.ts index c51742783b04..7d6a0b75e824 100644 --- a/apps/desktop/src/lib/embedded-images.test.ts +++ b/apps/desktop/src/lib/embedded-images.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it } from 'vitest' -import { extractEmbeddedImages } from './embedded-images' +import { extractEmbeddedImages, extractImageRefs, textWithoutImageRefs } from './embedded-images' const SAMPLE_PNG_DATA_URL = 'data:image/png;base64,' + 'A'.repeat(120) @@ -42,3 +42,72 @@ describe('extractEmbeddedImages', () => { expect(result.images[0]).toHaveLength(hugeDataUrl.length) }) }) + +describe('textWithoutImageRefs', () => { + it('leaves plain text untouched', () => { + expect(textWithoutImageRefs('just a question')).toBe('just a question') + }) + + it('strips a single leading @image directive line', () => { + expect(textWithoutImageRefs('@image:/tmp/cat.png\nwhat is this?')).toBe('what is this?') + }) + + it('strips multiple @image directive lines and trims', () => { + const input = '@image:/tmp/a.png\n@image:/tmp/b.png\n describe both ' + + expect(textWithoutImageRefs(input)).toBe('describe both') + }) + + it('does not treat an inline @image mention as a directive line', () => { + // Only full-line leading directives are stripped, matching the gateway's + // persist-time rewrite. A bare mention mid-prose is preserved. + expect(textWithoutImageRefs('see @image:/tmp/cat.png here')).toBe('see @image:/tmp/cat.png here') + }) +}) + +describe('extractImageRefs', () => { + it('returns the text untouched and no refs when there are no directives', () => { + expect(extractImageRefs('a normal prompt')).toEqual({ cleanedText: 'a normal prompt', refs: [] }) + }) + + it('lifts leading @image directive lines into refs and clears the text', () => { + const result = extractImageRefs('@image:/tmp/cat.png\nwhat do you see?') + + expect(result).toEqual({ cleanedText: 'what do you see?', refs: ['@image:/tmp/cat.png'] }) + }) + + it('collects multiple refs in order', () => { + const result = extractImageRefs('@image:/tmp/a.png\n@image:/tmp/b.png\ncompare them') + + expect(result.cleanedText).toBe('compare them') + expect(result.refs).toEqual(['@image:/tmp/a.png', '@image:/tmp/b.png']) + }) + + it('keeps only the directive lines when there is no trailing text', () => { + const result = extractImageRefs('@image:/tmp/only.png') + + expect(result).toEqual({ cleanedText: '', refs: ['@image:/tmp/only.png'] }) + }) + + it('lifts a backtick-quoted ref so a path with spaces survives intact', () => { + const ref = '@image:`/Users/me/Library/Application Support/Hermes/composer-images/a.png`' + const result = extractImageRefs(`${ref}\nwhat is this?`) + + expect(result).toEqual({ cleanedText: 'what is this?', refs: [ref] }) + }) + + it('drops the [screenshot] placeholder a native-vision turn leaves behind', () => { + // Flattening a parts list replaces each image part with `[screenshot]`; the + // lifted ref already renders that same attachment. + const result = extractImageRefs('@image:/tmp/cat.png\nwhat is in this photo?\n[screenshot]') + + expect(result).toEqual({ cleanedText: 'what is in this photo?', refs: ['@image:/tmp/cat.png'] }) + }) + + it('keeps [screenshot] when the message carries no image refs', () => { + expect(extractImageRefs('[screenshot]\nlook at the attached capture')).toEqual({ + cleanedText: '[screenshot]\nlook at the attached capture', + refs: [] + }) + }) +}) diff --git a/apps/desktop/src/lib/embedded-images.ts b/apps/desktop/src/lib/embedded-images.ts index 9b75eeae1403..144a602091e9 100644 --- a/apps/desktop/src/lib/embedded-images.ts +++ b/apps/desktop/src/lib/embedded-images.ts @@ -160,3 +160,50 @@ export function embeddedImageUrls(text: string): string[] { export function textWithoutEmbeddedImages(text: string): string { return extractEmbeddedImages(text).cleanedText } + +// The gateway persists attached images as `@image:` directive lines +// (see tui_gateway/server.py's persist-time rewrite), prepended before the +// user's own text. The composer's own optimistic/local turn never carries +// this prefix — it keeps the attachment as separate `attachmentRefs` +// metadata, not inline text. Comparing raw chatMessageText between the +// optimistic turn and the authoritative (persisted) turn therefore always +// mismatches whenever an image was attached, which defeats the "is this the +// same turn" checks in preserveLocalPendingTurnMessages / appendLiveSessionProjection +// and re-appends the optimistic row as if it were a distinct, unconfirmed +// turn — a duplicated user bubble. Strip the directive line(s) before any +// such equality comparison so both sides reduce to the same visible text. +const IMAGE_REF_LINE_RE = /^@image:[^\n]*\n?/gm + +export function textWithoutImageRefs(text: string): string { + return text.replace(IMAGE_REF_LINE_RE, '').trim() +} + +// Same directive lines as textWithoutImageRefs, but keeps them instead of +// discarding — used when converting persisted server messages into +// ChatMessage/ThreadMessageLike shape, where `@image:` refs need to +// move from inline text into the `attachmentRefs` metadata field (mirroring +// how the local optimistic composer represents attachments) rather than stay +// embedded in the bubble's clamped text body, where a large inline thumbnail +// pushes the caption text out of the clamp's visible area. +// Native-vision turns are stored as a parts list, which the session store +// flattens by replacing each image part with a literal `[screenshot]` line. The +// `@image:` ref describes that same attachment, so keeping both renders the +// placeholder as stray text under the thumbnail. Drop it only when a ref was +// actually lifted, so a `[screenshot]` in a message without attachments stays. +const SCREENSHOT_PLACEHOLDER_LINE_RE = /^\[screenshot\]\n?/gm + +export function extractImageRefs(text: string): { cleanedText: string; refs: string[] } { + const refs: string[] = [] + + let cleanedText = text.replace(IMAGE_REF_LINE_RE, match => { + refs.push(match.trim()) + + return '' + }) + + if (refs.length) { + cleanedText = cleanedText.replace(SCREENSHOT_PLACEHOLDER_LINE_RE, '') + } + + return { cleanedText: cleanedText.trim(), refs } +} diff --git a/tests/agent/test_context_references.py b/tests/agent/test_context_references.py index 1fab3ef1b74e..51f66fb937af 100644 --- a/tests/agent/test_context_references.py +++ b/tests/agent/test_context_references.py @@ -445,3 +445,25 @@ async def test_canonical_guard_fails_closed_when_lookup_raises(tmp_path: Path, m "credential deny-list" in warning or "sensitive credential" in warning for warning in result.warnings ) + + +@pytest.mark.parametrize( + "value", + [ + "/tmp/plain.png", + "/Users/me/Library/Application Support/Hermes/composer-images/a.png", + r"C:\Users\John Doe\Pictures\cat.png", + "/tmp/report (final).pdf", + "/tmp/it's here.png", + '/tmp/say "hi".png', + ], +) +def test_format_reference_value_round_trips_through_the_parser(value): + """Whatever the path contains, the formatted ref must parse back whole — + an unquoted value stops at the first space and strands the tail as text.""" + from agent.context_references import REFERENCE_PATTERN, format_reference_value + + match = REFERENCE_PATTERN.search(f"@file:{format_reference_value(value)}") + + assert match is not None + assert match.group("value").strip("`\"'") == value diff --git a/tests/test_lazy_session_regressions.py b/tests/test_lazy_session_regressions.py index 1e5a75206e1d..ec684254c62d 100644 --- a/tests/test_lazy_session_regressions.py +++ b/tests/test_lazy_session_regressions.py @@ -140,7 +140,7 @@ class TestSyncSessionKeyAfterAutoCompress: self.session_id = "pre-compress-key" self._cached_system_prompt = "" - def run_conversation(self, prompt, conversation_history=None, stream_callback=None): + def run_conversation(self, prompt, conversation_history=None, stream_callback=None, **_kwargs): # Simulate what _compress_context does: rotate session_id self.session_id = "post-compress-key" return { diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index 5aa76b437315..f70291c3e6c1 100644 --- a/tests/test_tui_gateway_server.py +++ b/tests/test_tui_gateway_server.py @@ -3313,7 +3313,7 @@ def test_prompt_submit_empty_truncation_allowed_with_confirm(monkeypatch): class _Agent: def run_conversation( - self, prompt, conversation_history=None, stream_callback=None + self, prompt, conversation_history=None, stream_callback=None, **_kwargs ): seen["prompt"] = prompt seen["history"] = conversation_history @@ -3709,9 +3709,7 @@ class _RecordingAgent: def clear_interrupt(self): return None - def run_conversation( - self, prompt, conversation_history=None, stream_callback=None - ): + def run_conversation(self, prompt, conversation_history=None, stream_callback=None, **_kwargs): self._turns.append(prompt) return {"final_response": "", "messages": []} @@ -3820,9 +3818,7 @@ def test_run_prompt_submit_requeues_all_unstarted_notifications_with_real_thread return thread class _BlockingNotificationAgent(_RecordingAgent): - def run_conversation( - self, prompt, conversation_history=None, stream_callback=None - ): + def run_conversation(self, prompt, conversation_history=None, stream_callback=None, **_kwargs): turns.append(prompt) if "proc_batch_1" in prompt: nested_started.set() @@ -6608,9 +6604,7 @@ def test_prompt_submit_sets_approval_session_key(monkeypatch): captured = {} class _Agent: - def run_conversation( - self, prompt, conversation_history=None, stream_callback=None - ): + def run_conversation(self, prompt, conversation_history=None, stream_callback=None, **_kwargs): captured["session_key"] = get_current_session_key(default="") return { "final_response": "ok", @@ -6650,9 +6644,7 @@ def test_prompt_submit_expands_context_refs(monkeypatch): base_url = "" api_key = "" - def run_conversation( - self, prompt, conversation_history=None, stream_callback=None - ): + def run_conversation(self, prompt, conversation_history=None, stream_callback=None, **_kwargs): captured["prompt"] = prompt return { "final_response": "ok", @@ -7521,9 +7513,7 @@ def test_prompt_submit_history_version_mismatch_surfaces_warning(monkeypatch): session_ref = {"s": None} class _RacyAgent: - def run_conversation( - self, prompt, conversation_history=None, stream_callback=None - ): + def run_conversation(self, prompt, conversation_history=None, stream_callback=None, **_kwargs): # Simulate: something external bumped history_version # while we were running. with session_ref["s"]["history_lock"]: @@ -7584,9 +7574,7 @@ def test_prompt_submit_sanitizes_bracketed_paste_before_agent(monkeypatch): captured: dict[str, str] = {} class _Agent: - def run_conversation( - self, prompt, conversation_history=None, stream_callback=None - ): + def run_conversation(self, prompt, conversation_history=None, stream_callback=None, **_kwargs): captured["prompt"] = prompt return { "final_response": "ok", @@ -7628,9 +7616,7 @@ def test_prompt_submit_history_version_match_persists_normally(monkeypatch): """Regression guard: the backstop does not affect the happy path.""" class _Agent: - def run_conversation( - self, prompt, conversation_history=None, stream_callback=None - ): + def run_conversation(self, prompt, conversation_history=None, stream_callback=None, **_kwargs): return { "final_response": "reply", "messages": [{"role": "assistant", "content": "reply"}], @@ -7681,9 +7667,7 @@ def test_prompt_submit_can_truncate_before_user_ordinal(monkeypatch): seen = {} class _Agent: - def run_conversation( - self, prompt, conversation_history=None, stream_callback=None - ): + def run_conversation(self, prompt, conversation_history=None, stream_callback=None, **_kwargs): seen["prompt"] = prompt seen["history"] = conversation_history return { @@ -9671,9 +9655,7 @@ def test_prompt_submit_auto_titles_session_on_complete(monkeypatch): api_key = object() api_mode = "codex_responses" - def run_conversation( - self, prompt, conversation_history=None, stream_callback=None - ): + def run_conversation(self, prompt, conversation_history=None, stream_callback=None, **_kwargs): return { "final_response": "Rome was founded in 753 BC.", "messages": [ @@ -9716,9 +9698,7 @@ def test_prompt_submit_skips_auto_title_when_interrupted(monkeypatch): """maybe_auto_title must NOT be called when the agent was interrupted.""" class _Agent: - def run_conversation( - self, prompt, conversation_history=None, stream_callback=None - ): + def run_conversation(self, prompt, conversation_history=None, stream_callback=None, **_kwargs): return { "final_response": "partial answer", "interrupted": True, @@ -9748,9 +9728,7 @@ def test_prompt_submit_skips_auto_title_when_response_empty(monkeypatch): """maybe_auto_title must NOT be called when the agent returns an empty reply.""" class _Agent: - def run_conversation( - self, prompt, conversation_history=None, stream_callback=None - ): + def run_conversation(self, prompt, conversation_history=None, stream_callback=None, **_kwargs): return { "final_response": "", "messages": [], @@ -9781,9 +9759,7 @@ def test_prompt_submit_surfaces_backend_error_as_visible_text(monkeypatch): instead of emitting a blank message.complete turn.""" class _Agent: - def run_conversation( - self, prompt, conversation_history=None, stream_callback=None - ): + def run_conversation(self, prompt, conversation_history=None, stream_callback=None, **_kwargs): return { "final_response": None, "messages": [], @@ -9828,9 +9804,7 @@ def test_prompt_submit_preserves_empty_response_without_error(monkeypatch): semantics owned by downstream handlers.""" class _Agent: - def run_conversation( - self, prompt, conversation_history=None, stream_callback=None - ): + def run_conversation(self, prompt, conversation_history=None, stream_callback=None, **_kwargs): return { "final_response": None, "messages": [], @@ -9993,7 +9967,7 @@ def test_session_activate_returns_inflight_stream_before_completion(monkeypatch) class _Agent: model = "model-live" - def run_conversation(self, prompt, conversation_history=None, stream_callback=None): + def run_conversation(self, prompt, conversation_history=None, stream_callback=None, **_kwargs): assert prompt == "write a long answer" assert conversation_history == [] stream_callback("partial ") @@ -11317,7 +11291,7 @@ def test_notification_poller_delivers_completion(monkeypatch): emitted = [] class _Agent: - def run_conversation(self, prompt, conversation_history=None, stream_callback=None): + def run_conversation(self, prompt, conversation_history=None, stream_callback=None, **_kwargs): turns.append(prompt) return { "final_response": "ok", @@ -11388,7 +11362,7 @@ def test_notification_poller_skips_consumed(monkeypatch): turns = [] class _Agent: - def run_conversation(self, prompt, conversation_history=None, stream_callback=None): + def run_conversation(self, prompt, conversation_history=None, stream_callback=None, **_kwargs): turns.append(prompt) return {"final_response": "ok", "messages": []} @@ -13229,3 +13203,185 @@ def test_clarify_timeout_seconds_maps_non_positive_to_unlimited(monkeypatch, con monkeypatch.setattr("tools.clarify_gateway.get_clarify_timeout", lambda: configured) assert server._clarify_timeout_seconds() == expected + + +def test_build_persist_message_with_image_refs_without_images_returns_text(monkeypatch): + """#70720: when no images are attached the persisted message is the raw + prompt — no @image directive prefix is introduced.""" + assert server._build_persist_message_with_image_refs("what is this?", []) == "what is this?" + assert server._build_persist_message_with_image_refs("", []) == "" + + +def test_build_persist_message_with_image_refs_appends_existing_paths(monkeypatch, tmp_path): + """Attached images that still exist on disk are persisted as trailing + ``@image:`` directive lines so the desktop renders them after a + restart (instead of the vision-only enrichment that silently breaks).""" + img = tmp_path / "cat.png" + img.write_bytes(b"\x89PNG") + + result = server._build_persist_message_with_image_refs("what is in this photo?", [str(img)]) + + assert result == f"what is in this photo?\n@image:{img}" + + +def test_build_persist_message_keeps_the_caption_on_the_first_line(tmp_path): + """Session previews are the first 60 characters of the first user message, + so a leading directive would title the session with a truncated file path + in the sidebar, switcher, and command palette.""" + img = tmp_path / "cat.png" + img.write_bytes(b"png") + + result = server._build_persist_message_with_image_refs("what is in this photo?", [str(img)]) + + assert result.split("\n", 1)[0] == "what is in this photo?" + + +def test_build_persist_message_with_image_refs_skips_missing_paths(monkeypatch, tmp_path): + """Only paths that still exist are persisted; a missing file must not + inject a dangling @image ref into the transcript.""" + existing = tmp_path / "a.png" + existing.write_bytes(b"png") + missing = str(tmp_path / "gone.png") + + result = server._build_persist_message_with_image_refs("compare them", [str(existing), missing]) + + assert result == f"compare them\n@image:{existing}" + + +def test_build_persist_message_with_image_refs_without_text_is_refs_only(monkeypatch, tmp_path): + """A stand-alone attachment (no caption) persists as just the directive + line, so a bare image survives in history and is not dropped as empty.""" + img = tmp_path / "only.png" + img.write_bytes(b"png") + + assert server._build_persist_message_with_image_refs("", [str(img)]) == f"@image:{img}" + + +def test_build_persist_message_quotes_paths_containing_spaces(tmp_path): + """The unquoted alternative in the directive pattern is ``\\S+``, so a path + with a space parses as a truncated ref with the tail left as loose text. + Desktop composer images live in the app's userData dir, which on macOS is + ``~/Library/Application Support/...`` — a space every time.""" + img_dir = tmp_path / "Application Support" / "Hermes" / "composer-images" + img_dir.mkdir(parents=True) + img = img_dir / "cat.png" + img.write_bytes(b"png") + + result = server._build_persist_message_with_image_refs("what is this?", [str(img)]) + + assert result == f"what is this?\n@image:`{img}`" + + +def test_persist_user_message_mirrors_the_shape_sent_to_the_model(tmp_path): + """A native-vision turn sends ``content`` as a parts list, and the session + store ignores a plain-string override for a list payload. The override must + mirror the list shape (ref text + the original image parts) or it is + silently dropped and the attachment never reaches history.""" + img = tmp_path / "cat.png" + img.write_bytes(b"png") + image_part = {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}} + native_parts = [{"type": "text", "text": "api-only text"}, image_part] + + override = server._build_persist_user_message("what is this?", [str(img)], native_parts) + + assert override == [{"type": "text", "text": f"what is this?\n@image:{img}"}, image_part] + + +def test_persist_user_message_stays_a_string_for_text_mode(tmp_path): + """Text-mode (vision-preprocessed) turns send a string, so the override + stays a string — the shape the session store rewrites directly.""" + img = tmp_path / "cat.png" + img.write_bytes(b"png") + + override = server._build_persist_user_message("what is this?", [str(img)], "enriched api-only text") + + assert override == f"what is this?\n@image:{img}" + + +def test_native_vision_turn_persists_a_renderable_image_ref(tmp_path): + """End to end through the real session-store flush: whichever image input + mode the turn used, the durable row carries an ``@image:`` ref the desktop + can render after a restart.""" + from unittest.mock import MagicMock + + from agent.image_routing import build_native_content_parts + from run_agent import AIAgent + + img_dir = tmp_path / "Application Support" / "composer-images" + img_dir.mkdir(parents=True) + img = img_dir / "cat.png" + img.write_bytes( + bytes.fromhex( + "89504e470d0a1a0a0000000d494844520000000100000001080600000" + "01f15c4890000000a49444154789c6360000002000100ffff0300000600" + "0557bfabd40000000049454e44ae426082" + ) + ) + native_parts, skipped = build_native_content_parts("what is in this photo?", [str(img)]) + assert not skipped + + agent = AIAgent.__new__(AIAgent) + agent._session_db = MagicMock() + agent._session_db_created = True + agent.session_id = "s-1" + agent._last_flushed_db_idx = 0 + agent._persist_disabled = False + agent._flushed_db_message_ids = set() + agent._flushed_db_message_session_id = None + agent._pending_cli_user_message = None + agent._persist_user_message_timestamp = None + agent._persist_user_message_idx = 0 + agent._persist_user_message_override = server._build_persist_user_message( + "what is in this photo?", [str(img)], native_parts + ) + + agent._flush_messages_to_session_db([{"role": "user", "content": native_parts}], []) + + written = agent._session_db.append_message.call_args.kwargs["content"] + assert f"@image:`{img}`" in written + assert "what is in this photo?" in written + # The model keeps the pixels for the rest of the session. + assert any(part.get("type") == "image_url" for part in agent._persist_user_message_override) + + +def test_prompt_submit_passes_persist_user_message_to_agent(monkeypatch): + """#70720: _run_prompt_submit must forward the (image-ref-aware) persisted + user message to run_conversation via persist_user_message, so the gateway + stores the UI-recognizable form instead of the vision enrichment.""" + captured = {} + + class _Agent: + def run_conversation(self, prompt, conversation_history=None, stream_callback=None, **_kwargs): + captured["persist_user_message"] = _kwargs.get("persist_user_message") + return { + "final_response": "reply", + "messages": [{"role": "assistant", "content": "reply"}], + } + + class _ImmediateThread: + def __init__(self, target=None, daemon=None): + self._target = target + + def start(self): + self._target() + + server._sessions["sid"] = _session(agent=_Agent()) + try: + monkeypatch.setattr(server.threading, "Thread", _ImmediateThread) + monkeypatch.setattr(server, "_get_usage", lambda _a: {}) + monkeypatch.setattr(server, "render_message", lambda _t, _c: "") + monkeypatch.setattr(server, "_emit", lambda *a: None) + + resp = server.handle_request( + { + "id": "1", + "method": "prompt.submit", + "params": {"session_id": "sid", "text": "hi"}, + } + ) + assert resp.get("result") + + # Without attachments the persist form equals the raw prompt. + assert captured.get("persist_user_message") == "hi" + finally: + server._sessions.pop("sid", None) diff --git a/tui_gateway/server.py b/tui_gateway/server.py index 0c48059bf241..5de8d4fc6abf 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -5739,6 +5739,49 @@ def _enrich_with_attached_images(user_text: str, image_paths: list[str]) -> str: return text or "What do you see in this image?" +def _build_persist_message_with_image_refs(user_text: str, image_paths: list[str]) -> str: + """Build the clean, UI-recognizable version of the user's message for + persisting to session history. Uses ``@image:`` directives — the + format the desktop client (directive-text.tsx / HERMES_DIRECTIVE_RE) + actually parses and renders as an image — unlike + ``_enrich_with_attached_images``, which embeds a vision description and + an ``image_url:`` hint meant only for the model and must never be + persisted as-is (it silently breaks image rendering after a full + restart, and reorders image/text on live session-switch reconciliation). + + The caption leads and the directives trail: session previews are the first + 60 characters of the first user message (``list_sessions_rich``), so a + leading directive would label the session with a truncated file path in the + sidebar, switcher, and command palette. Clients lift the refs out of the + body by line, so their position does not affect how the turn renders. + """ + from agent.context_references import format_reference_value + + text = user_text or "" + refs = "\n".join(f"@image:{format_reference_value(p)}" for p in image_paths if Path(p).exists()) + if not refs: + return text + return f"{text}\n{refs}" if text else refs + + +def _build_persist_user_message(user_text: str, image_paths: list[str], run_message: Any) -> Any: + """Shape the persisted user turn to match what was sent to the model. + + Native-vision turns send ``content`` as a parts list, and + ``_flush_messages_to_session_db`` deliberately ignores a plain-string + override for a list payload (a text override must not erase a turn's + image/audio summary). So mirror the shape: replace only the text part with + the ``@image:`` ref form and keep the image parts, so the model still has + the pixels for the rest of the session. Any API-only text part (the + barge-in note) is dropped along the way, which is the point of the override. + """ + persist_text = _build_persist_message_with_image_refs(user_text, image_paths) + if not isinstance(run_message, list): + return persist_text + image_parts = [p for p in run_message if not (isinstance(p, dict) and p.get("type") == "text")] + return [{"type": "text", "text": persist_text}, *image_parts] + + def _content_display_text(content: Any) -> str: if content is None: return "" @@ -10997,6 +11040,9 @@ def _run_prompt_submit( run_kwargs = { "conversation_history": list(history), "stream_callback": _stream, + "persist_user_message": ( + _build_persist_user_message(prompt, images, run_message) if images else prompt + ), } try: if "task_id" in inspect.signature(agent.run_conversation).parameters: