From d0f5ef70416da4fad0707e77a5defe22daefdd9f Mon Sep 17 00:00:00 2001 From: alelpoan Date: Fri, 24 Jul 2026 20:03:30 -0500 Subject: [PATCH] fix(desktop): persist @image: refs instead of the vision-enrichment text The desktop gateway passed the vision-enriched, model-only message text (carrying an `image_url:` hint) straight into run_conversation as the persisted user turn. The renderer only parses `@image:`, so it could not rebuild the attachment from history: after a restart the image was gone and only the caption survived, and on a live session switch the warm cache disagreed with the authoritative text and the frontend "rescued" the image by appending it after the caption. run_conversation already supports persist_user_message for exactly this "what the model sees" vs "what gets stored" split; it was simply never wired up for the attachment path. --- .../hooks/use-session-actions/utils.test.ts | 82 ++++++++++ .../hooks/use-session-actions/utils.ts | 9 +- .../assistant-ui/directive-text.tsx | 19 ++- apps/desktop/src/lib/chat-messages.test.ts | 41 +++++ apps/desktop/src/lib/chat-messages.ts | 23 ++- apps/desktop/src/lib/embedded-images.test.ts | 49 +++++- apps/desktop/src/lib/embedded-images.ts | 37 +++++ tests/test_lazy_session_regressions.py | 2 +- tests/test_tui_gateway_server.py | 143 ++++++++++++------ tui_gateway/server.py | 22 +++ 10 files changed, 369 insertions(+), 58 deletions(-) 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..5c3dc83a6de8 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 @@ -440,9 +440,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..f7287497bc59 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' @@ -306,7 +306,7 @@ export function preserveLocalPendingTurnMessages( if ( isOptimisticUser && latestAuthoritativeUser && - chatMessageText(latestAuthoritativeUser).trim() === chatMessageText(message).trim() + textWithoutImageRefs(chatMessageText(latestAuthoritativeUser)) === textWithoutImageRefs(chatMessageText(message)) ) { continue } @@ -318,7 +318,7 @@ export function preserveLocalPendingTurnMessages( continue } - if (chatMessageText(authoritative).trim() === chatMessageText(message).trim()) { + if (textWithoutImageRefs(chatMessageText(authoritative)) === textWithoutImageRefs(chatMessageText(message))) { continue } } @@ -359,7 +359,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..4f9f489d872c 100644 --- a/apps/desktop/src/components/assistant-ui/directive-text.tsx +++ b/apps/desktop/src/components/assistant-ui/directive-text.tsx @@ -351,19 +351,28 @@ 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 = segments.filter( + (segment): segment is Extract => + segment.kind === 'mention' && segment.type === 'image' +) + 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('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..83c3e21bbf1d 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,50 @@ 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'] }) + }) +}) diff --git a/apps/desktop/src/lib/embedded-images.ts b/apps/desktop/src/lib/embedded-images.ts index 9b75eeae1403..8c5e4d8c28ea 100644 --- a/apps/desktop/src/lib/embedded-images.ts +++ b/apps/desktop/src/lib/embedded-images.ts @@ -160,3 +160,40 @@ 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. +export function extractImageRefs(text: string): { cleanedText: string; refs: string[] } { + const refs: string[] = [] + const cleanedText = text + .replace(IMAGE_REF_LINE_RE, match => { + refs.push(match.trim()) + + return '' + }) + .trim() + + return { cleanedText, refs } +} 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..fcf22be117c9 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,86 @@ 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_prepends_existing_paths(monkeypatch, tmp_path): + """Attached images that still exist on disk are persisted as leading + ``@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"@image:{img}\nwhat 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"@image:{existing}\ncompare them" + + +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_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..9f55c63c58f1 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -5739,6 +5739,23 @@ 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). + """ + text = user_text or "" + refs = "\n".join(f"@image:{p}" for p in image_paths if Path(p).exists()) + if not refs: + return text + return f"{refs}\n{text}" if text else refs + + def _content_display_text(content: Any) -> str: if content is None: return "" @@ -10997,6 +11014,11 @@ def _run_prompt_submit( run_kwargs = { "conversation_history": list(history), "stream_callback": _stream, + "persist_user_message": ( + _build_persist_message_with_image_refs(prompt, images) + if images + else prompt + ), } try: if "task_id" in inspect.signature(agent.run_conversation).parameters: