From e652c694f82c8474bfe39ca266fdb788f2e7bd0d Mon Sep 17 00:00:00 2001 From: Gerard-Devlin Date: Thu, 11 Jun 2026 01:10:04 +0800 Subject: [PATCH] feat: enhance message handling to sanitize stored user prompts and improve context display --- .../stores/claude-chat-send-prompt.test.ts | 44 ++++++++++++ .../components/claude-chat/chat-messages.tsx | 13 ++-- apps/desktop/src/stores/claude-chat-store.ts | 68 +++++++++++++++++-- 3 files changed, 117 insertions(+), 8 deletions(-) diff --git a/apps/desktop/src/__tests__/stores/claude-chat-send-prompt.test.ts b/apps/desktop/src/__tests__/stores/claude-chat-send-prompt.test.ts index 3a7c68a..25387df 100644 --- a/apps/desktop/src/__tests__/stores/claude-chat-send-prompt.test.ts +++ b/apps/desktop/src/__tests__/stores/claude-chat-send-prompt.test.ts @@ -321,4 +321,48 @@ describe("useClaudeChatStore.resumeSession", () => { expect(state.totalInputTokens).toBe(24); expect(state.totalOutputTokens).toBe(12); }); + + it("hides internal file and pasted-image context when restoring history", async () => { + const tempImagePath = [ + "C:\\Users\\Devlin\\AppData\\Local\\Temp", + "ClaudePrism", + "chat-pastes", + "1781110224092-1-paste-1781110223586-1.png", + ].join("\\"); + const restoredPrompt = [ + "[Currently open file: main.tex]", + "[Selection: Pasted image]", + "[Selected text:", + `[Temporary pasted image: ${tempImagePath}]`, + "Use this image file as visual context for the user's message.", + "]", + "", + "Please inspect this image", + ].join("\n"); + + vi.mocked(invoke).mockResolvedValueOnce([ + { + type: "user", + message: { + content: restoredPrompt, + }, + }, + { + type: "assistant", + message: { + content: [{ type: "text", text: "OK" }], + }, + }, + ]); + + await useClaudeChatStore.getState().resumeSession("session-with-image"); + + const state = useClaudeChatStore.getState(); + const userContent = state.messages[0].message?.content as any; + + expect(userContent).toBe("Pasted image\nPlease inspect this image"); + expect(userContent).not.toContain("[Currently open file:"); + expect(userContent).not.toContain("[Temporary pasted image:"); + expect(state.title).toBe("Please inspect this image"); + }); }); diff --git a/apps/desktop/src/components/claude-chat/chat-messages.tsx b/apps/desktop/src/components/claude-chat/chat-messages.tsx index 22fc238..6d64356 100644 --- a/apps/desktop/src/components/claude-chat/chat-messages.tsx +++ b/apps/desktop/src/components/claude-chat/chat-messages.tsx @@ -215,10 +215,15 @@ const UserMessage: FC<{ message: ClaudeStreamMessage }> = ({ message }) => { if (!textContent) return null; - // Parse leading @file:line:col or ~@file:line context reference - const contextMatch = textContent.match(/^(~?@[^\n]+)\n([\s\S]*)$/); - const contextLabel = contextMatch?.[1] ?? null; - const bodyText = contextMatch ? contextMatch[2] : textContent; + const firstLineMatch = textContent.match(/^([^\n]+)\n([\s\S]*)$/); + const firstLine = firstLineMatch?.[1]?.trim() ?? ""; + const hasContextLabel = + firstLine.startsWith("@") || + firstLine.startsWith("~@") || + /^Pasted image(?: \d+)?(?:, Pasted image(?: \d+)?)*$/.test(firstLine); + const contextLabel = hasContextLabel ? firstLine : null; + const bodyText = + hasContextLabel && firstLineMatch ? firstLineMatch[2] : textContent; // Parse error block patterns for styled rendering: // Lint single: "[Lint error in FILE:LINE]\n[Error: MSG]\n\nPrompt" diff --git a/apps/desktop/src/stores/claude-chat-store.ts b/apps/desktop/src/stores/claude-chat-store.ts index 80df88f..3fc5317 100644 --- a/apps/desktop/src/stores/claude-chat-store.ts +++ b/apps/desktop/src/stores/claude-chat-store.ts @@ -240,7 +240,10 @@ function stringifyBlockContent(value: unknown): string { } function messageContentText(message: ClaudeStreamMessage): string { - const blocks = message.message?.content ?? []; + const rawContent = (message.message as any)?.content; + if (typeof rawContent === "string") return rawContent.trim(); + + const blocks = rawContent ?? []; const parts: string[] = []; for (const block of blocks) { if (block.type === "text" && block.text?.trim()) { @@ -258,6 +261,62 @@ function messageContentText(message: ClaudeStreamMessage): string { return parts.join("\n").trim(); } +function displayTextForStoredUserPrompt(text: string): string { + const normalized = text.replace(/\r\n/g, "\n"); + if (!/^\[(?:Currently open file|File): [^\]\n]*\]/.test(normalized)) { + return text; + } + + const contextEnd = normalized.lastIndexOf("]\n\n"); + if (contextEnd < 0) return text; + + const contextText = normalized.slice(0, contextEnd + 1); + const body = normalized.slice(contextEnd + 3); + const selectionMatch = contextText.match( + /(?:^|\n)\[Selection: ([^\]\n]+)\]/, + ); + const contextLabel = selectionMatch?.[1]?.trim(); + + if (!contextLabel) return body; + return body.trim() ? `${contextLabel}\n${body}` : contextLabel; +} + +function sanitizeStoredUserMessageForDisplay( + message: ClaudeStreamMessage, +): ClaudeStreamMessage { + if (message.type !== "user") return message; + + const rawContent = (message.message as any)?.content; + if (typeof rawContent === "string") { + const displayText = displayTextForStoredUserPrompt(rawContent); + return displayText === rawContent + ? message + : { + ...message, + message: { ...message.message, content: displayText as any }, + }; + } + + if (!Array.isArray(rawContent)) return message; + + let changed = false; + const content = rawContent.map((block) => { + if (block.type !== "text" || typeof block.text !== "string") { + return block; + } + + const displayText = displayTextForStoredUserPrompt(block.text); + if (displayText === block.text) return block; + + changed = true; + return { ...block, text: displayText }; + }); + + return changed + ? { ...message, message: { ...message.message, content } } + : message; +} + function buildProviderSwitchContext( messages: ClaudeStreamMessage[], maxChars = 18000, @@ -1150,14 +1209,15 @@ export const useClaudeChatStore = create()((set, get) => ({ }); // Filter to displayable message types and map to ClaudeStreamMessage - const messages: ClaudeStreamMessage[] = []; + const rawMessages: ClaudeStreamMessage[] = []; for (const entry of history) { const type = entry.type; if (type === "user" || type === "assistant" || type === "result") { - messages.push(entry as ClaudeStreamMessage); + rawMessages.push(entry as ClaudeStreamMessage); } } + const messages = rawMessages.map(sanitizeStoredUserMessageForDisplay); const totals = usageTotalsForMessages(messages); const providerKey = inferProviderKeyFromHistory(history); const selectedProviderCredentialId = @@ -1169,7 +1229,7 @@ export const useClaudeChatStore = create()((set, get) => ({ ...applyTabUpdate(s, activeTabId, { messages, providerKey, - title: sessionTitle ?? titleForMessages(messages) ?? "New Chat", + title: sessionTitle ?? titleForMessages(rawMessages) ?? "New Chat", totalInputTokens: totals.inputTokens, totalOutputTokens: totals.outputTokens, }),