From ffbcfa1597c0feefec3c12ab6c72f78cd4f1d29e Mon Sep 17 00:00:00 2001 From: Brooklyn Nicholson Date: Fri, 24 Jul 2026 20:15:39 -0500 Subject: [PATCH] fix(desktop): persist the image ref for natively-vision-capable models too MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A turn routed to a model that takes pixels directly sends `content` as a parts list, and the session store deliberately ignores a plain-string persist override for a list payload — a text override must not erase a turn's image summary. So the override was dropped for every user on a vision-capable main model, and the durable row kept only the caption plus a literal `[Image attached at: ...]` / `[screenshot]`, which the renderer cannot turn back into an image. Only vision-preprocessed (text-mode) turns were actually fixed. Mirror the shape instead: swap the text part for the `@image:` ref form and keep the image parts, so the model still has the pixels for the rest of the session, and drop the `[screenshot]` stand-in on the way into the bubble when a ref was lifted from the same message. --- apps/desktop/src/lib/chat-messages.test.ts | 18 +++++ apps/desktop/src/lib/embedded-images.test.ts | 22 ++++++ apps/desktop/src/lib/embedded-images.ts | 24 +++++-- tests/test_tui_gateway_server.py | 72 ++++++++++++++++++++ tui_gateway/server.py | 20 +++++- 5 files changed, 148 insertions(+), 8 deletions(-) diff --git a/apps/desktop/src/lib/chat-messages.test.ts b/apps/desktop/src/lib/chat-messages.test.ts index 9e967307561c..14bad6075bc7 100644 --- a/apps/desktop/src/lib/chat-messages.test.ts +++ b/apps/desktop/src/lib/chat-messages.test.ts @@ -162,6 +162,24 @@ describe('toChatMessages', () => { 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([ { diff --git a/apps/desktop/src/lib/embedded-images.test.ts b/apps/desktop/src/lib/embedded-images.test.ts index 83c3e21bbf1d..7d6a0b75e824 100644 --- a/apps/desktop/src/lib/embedded-images.test.ts +++ b/apps/desktop/src/lib/embedded-images.test.ts @@ -88,4 +88,26 @@ describe('extractImageRefs', () => { 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 8c5e4d8c28ea..144a602091e9 100644 --- a/apps/desktop/src/lib/embedded-images.ts +++ b/apps/desktop/src/lib/embedded-images.ts @@ -185,15 +185,25 @@ export function textWithoutImageRefs(text: string): string { // 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[] = [] - const cleanedText = text - .replace(IMAGE_REF_LINE_RE, match => { - refs.push(match.trim()) - return '' - }) - .trim() + let cleanedText = text.replace(IMAGE_REF_LINE_RE, match => { + refs.push(match.trim()) - return { cleanedText, refs } + return '' + }) + + if (refs.length) { + cleanedText = cleanedText.replace(SCREENSHOT_PLACEHOLDER_LINE_RE, '') + } + + return { cleanedText: cleanedText.trim(), refs } } diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index 99c811c484d6..3644e972de68 100644 --- a/tests/test_tui_gateway_server.py +++ b/tests/test_tui_gateway_server.py @@ -13260,6 +13260,78 @@ def test_build_persist_message_quotes_paths_containing_spaces(tmp_path): assert result == f"@image:`{img}`\nwhat is this?" +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"@image:{img}\nwhat is this?"}, 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"@image:{img}\nwhat is this?" + + +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 diff --git a/tui_gateway/server.py b/tui_gateway/server.py index 22b3764f7fcb..2345eb7a983e 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -5758,6 +5758,24 @@ def _build_persist_message_with_image_refs(user_text: str, image_paths: list[str return f"{refs}\n{text}" 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 "" @@ -11017,7 +11035,7 @@ def _run_prompt_submit( "conversation_history": list(history), "stream_callback": _stream, "persist_user_message": ( - _build_persist_message_with_image_refs(prompt, images) if images else prompt + _build_persist_user_message(prompt, images, run_message) if images else prompt ), } try: