From 8cce9889eb8f3b9f395f5d1f23d048ce317380aa Mon Sep 17 00:00:00 2001 From: 7418 <7418@openclaw.ai> Date: Thu, 12 Mar 2026 17:26:34 +0800 Subject: [PATCH] fix: Design Agent session isolation, image write-back robustness, and new-session flow MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Image reference store session-scoped: - Rewrite image-ref-store.ts: all keys/functions now scoped by sessionId - Thread sessionId through MessageList → StreamingMessage/MessageItem → ImageGenConfirmation - ChatView loads last-generated images for current session on mount Image write-back robustness: - Remove localStorage caching from ImageGenConfirmation (root cause of stale results) - Add rawRequestBlock for exact content matching in DB fallback query - updateMessageBySessionAndHint prefers exact block match, falls back to prompt hint - PUT route passes raw_request_block to DB function - Add retry-after-3s when first PUT fails and no messageId available - Exclude already-completed messages (NOT LIKE image-gen-result) from update query New session page (/chat) fix: - Add createdSessionId state, set after session creation, pass to MessageList - Prevents first-message image-gen-request from inheriting stale panel sessionId Co-Authored-By: Claude Opus 4.6 --- src/app/api/chat/messages/route.ts | 9 +- src/app/chat/page.tsx | 19 +++-- src/components/chat/ChatView.tsx | 23 ++++-- src/components/chat/ImageGenConfirmation.tsx | 85 +++++++------------ src/components/chat/MessageItem.tsx | 15 ++-- src/components/chat/MessageList.tsx | 3 +- src/components/chat/StreamingMessage.tsx | 8 +- src/lib/db.ts | 87 +++++++++++++++++--- src/lib/image-ref-store.ts | 46 +++++++---- 9 files changed, 187 insertions(+), 108 deletions(-) diff --git a/src/app/api/chat/messages/route.ts b/src/app/api/chat/messages/route.ts index acca080..5ce7416 100644 --- a/src/app/api/chat/messages/route.ts +++ b/src/app/api/chat/messages/route.ts @@ -49,11 +49,12 @@ export async function POST(request: NextRequest) { export async function PUT(request: NextRequest) { try { const body = await request.json(); - const { message_id, content, session_id, prompt_hint } = body as { + const { message_id, content, session_id, prompt_hint, raw_request_block } = body as { message_id: string; content: string; session_id?: string; prompt_hint?: string; + raw_request_block?: string; }; if (!content) { @@ -72,9 +73,9 @@ export async function PUT(request: NextRequest) { changes = updateMessageContent(message_id, content); } - // Fallback: search by session + prompt hint - if (changes === 0 && session_id && prompt_hint) { - const result = updateMessageBySessionAndHint(session_id, prompt_hint, content); + // Fallback: search by session + exact block or prompt hint + if (changes === 0 && session_id && (raw_request_block || prompt_hint)) { + const result = updateMessageBySessionAndHint(session_id, content, raw_request_block, prompt_hint); changes = result.changes; if (result.messageId) { updatedMessageId = result.messageId; diff --git a/src/app/chat/page.tsx b/src/app/chat/page.tsx index eae153e..f413131 100644 --- a/src/app/chat/page.tsx +++ b/src/app/chat/page.tsx @@ -39,23 +39,25 @@ export default function NewChatPage() { const [permissionResolved, setPermissionResolved] = useState<'allow' | 'deny' | null>(null); const [streamingToolOutput, setStreamingToolOutput] = useState(''); const [permissionProfile, setPermissionProfile] = useState<'default' | 'full_access'>('default'); + const [createdSessionId, setCreatedSessionId] = useState(); const abortControllerRef = useRef(null); // Effort level — lifted here so the first message includes it const [selectedEffort, setSelectedEffort] = useState(undefined); - // Thinking mode from app settings + // Provider options (thinking mode + 1M context) const [thinkingMode, setThinkingMode] = useState('adaptive'); + const [context1m, setContext1m] = useState(false); - // Fetch thinking mode from app settings + // Fetch provider-specific options useEffect(() => { - fetch('/api/settings/app') + const pid = currentProviderId || 'env'; + fetch(`/api/providers/options?providerId=${encodeURIComponent(pid)}`) .then(r => r.ok ? r.json() : null) .then(data => { - if (data?.settings?.thinking_mode) { - setThinkingMode(data.settings.thinking_mode); - } + setThinkingMode(data?.options?.thinking_mode || 'adaptive'); + setContext1m(!!data?.options?.context_1m); }) .catch(() => {}); - }, []); + }, [currentProviderId]); const stopStreaming = useCallback(() => { abortControllerRef.current?.abort(); @@ -150,6 +152,7 @@ export default function NewChatPage() { const { session }: SessionResponse = await createRes.json(); sessionId = session.id; + setCreatedSessionId(sessionId); // Notify ChatListPanel to refresh immediately window.dispatchEvent(new CustomEvent('session-created')); @@ -183,6 +186,7 @@ export default function NewChatPage() { ...(systemPromptAppend ? { systemPromptAppend } : {}), ...(selectedEffort ? { effort: selectedEffort } : {}), ...(thinkingConfig ? { thinking: thinkingConfig } : {}), + ...(context1m ? { context_1m: true } : {}), ...(displayOverride ? { displayOverride } : {}), }), signal: controller.signal, @@ -394,6 +398,7 @@ export default function NewChatPage() { messages={messages} streamingContent={streamingContent} isStreaming={isStreaming} + sessionId={createdSessionId} toolUses={toolUses} toolResults={toolResults} streamingToolOutput={streamingToolOutput} diff --git a/src/components/chat/ChatView.tsx b/src/components/chat/ChatView.tsx index db0229c..32b438a 100644 --- a/src/components/chat/ChatView.tsx +++ b/src/components/chat/ChatView.tsx @@ -14,7 +14,7 @@ import { usePanel } from '@/hooks/usePanel'; import { useTranslation } from '@/hooks/useTranslation'; import { PermissionPrompt } from './PermissionPrompt'; import { BatchExecutionDashboard, BatchContextSync } from './batch-image-gen'; -import { setLastGeneratedImages } from '@/lib/image-ref-store'; +import { setLastGeneratedImages, loadLastGenerated } from '@/lib/image-ref-store'; import { useChatCommands } from '@/hooks/useChatCommands'; import { useAssistantTrigger } from '@/hooks/useAssistantTrigger'; import { useStreamSubscription } from '@/hooks/useStreamSubscription'; @@ -53,20 +53,28 @@ export function ChatView({ sessionId, initialMessages = [], initialHasMore = fal const [currentProviderId, setCurrentProviderId] = useState(() => providerId || (typeof window !== 'undefined' ? localStorage.getItem('codepilot:last-provider-id') : null) || ''); const [selectedEffort, setSelectedEffort] = useState(undefined); const [thinkingMode, setThinkingMode] = useState('adaptive'); + const [context1m, setContext1m] = useState(false); // Sync model/provider when session data loads useEffect(() => { if (modelName) setCurrentModel(modelName); }, [modelName]); useEffect(() => { if (providerId) setCurrentProviderId(providerId); }, [providerId]); - // Fetch thinking mode from app settings + // Fetch provider-specific options (thinking mode + 1M context) useEffect(() => { - fetch('/api/settings/app') + const pid = currentProviderId || 'env'; + fetch(`/api/providers/options?providerId=${encodeURIComponent(pid)}`) .then(r => r.ok ? r.json() : null) - .then(data => { if (data?.settings?.thinking_mode) setThinkingMode(data.settings.thinking_mode); }) + .then(data => { + setThinkingMode(data?.options?.thinking_mode || 'adaptive'); + setContext1m(!!data?.options?.context_1m); + }) .catch(() => {}); - }, []); + }, [currentProviderId]); useEffect(() => { if (initialPermissionProfile) setPermissionProfile(initialPermissionProfile); }, [initialPermissionProfile]); + // Restore session-scoped last-generated images from sessionStorage + useEffect(() => { loadLastGenerated(sessionId); }, [sessionId]); + // Stream snapshot from the manager — drives all streaming UI const [streamSnapshot, setStreamSnapshot] = useState( () => getSnapshot(sessionId) @@ -288,6 +296,7 @@ export function ChatView({ sessionId, initialMessages = [], initialHasMore = fal pendingImageNotices: notices, effort: selectedEffort, thinking: buildThinkingConfig(), + context1m, displayOverride, onModeChanged: (sdkMode) => { const uiMode = sdkMode === 'plan' ? 'plan' : 'code'; @@ -302,7 +311,7 @@ export function ChatView({ sessionId, initialMessages = [], initialHasMore = fal }, }); }, - [sessionId, isStreaming, mode, currentModel, currentProviderId, selectedEffort, buildThinkingConfig, handleModeChange] + [sessionId, isStreaming, mode, currentModel, currentProviderId, selectedEffort, context1m, buildThinkingConfig, handleModeChange] ); sendMessageRef.current = sendMessage; @@ -321,7 +330,7 @@ export function ChatView({ sessionId, initialMessages = [], initialHasMore = fal const notice = `[Image generation completed]\n- Prompt: "${detail.prompt}"\n- Aspect ratio: ${detail.aspectRatio}\n- Resolution: ${detail.resolution}${pathInfo}`; if (paths.length > 0) { - setLastGeneratedImages(paths); + setLastGeneratedImages(sessionId, paths); } pendingImageNoticesRef.current.push(notice); diff --git a/src/components/chat/ImageGenConfirmation.tsx b/src/components/chat/ImageGenConfirmation.tsx index bbb60a1..c692f1c 100644 --- a/src/components/chat/ImageGenConfirmation.tsx +++ b/src/components/chat/ImageGenConfirmation.tsx @@ -19,29 +19,29 @@ const RESOLUTIONS = ['1K', '2K', '4K'] as const; interface ImageGenConfirmationProps { messageId?: string; + sessionId?: string; initialPrompt: string; initialAspectRatio: string; initialResolution: string; + /** The original raw ```image-gen-request...``` block — used for exact DB matching */ + rawRequestBlock?: string; referenceImages?: ReferenceImage[]; } type Status = 'idle' | 'generating' | 'completed' | 'error'; -/** Stable key for persisting generation results in localStorage */ -function storageKey(prompt: string, sessionId?: string): string { - const prefix = sessionId ? `${sessionId}:` : ''; - return `imggen:${prefix}${prompt.slice(0, 80)}`; -} - export function ImageGenConfirmation({ messageId, + sessionId: sessionIdProp, initialPrompt, initialAspectRatio, initialResolution, + rawRequestBlock, referenceImages, }: ImageGenConfirmationProps) { const { t } = useTranslation(); - const { sessionId } = usePanel(); + const { sessionId: panelSessionId } = usePanel(); + const sessionId = sessionIdProp || panelSessionId; const [prompt, setPrompt] = useState(initialPrompt); const [aspectRatio, setAspectRatio] = useState( ASPECT_RATIOS.includes(initialAspectRatio as typeof ASPECT_RATIOS[number]) @@ -58,21 +58,6 @@ export function ImageGenConfirmation({ const [error, setError] = useState(null); const abortRef = useRef(null); - // Restore completed result from localStorage on mount - useEffect(() => { - try { - const saved = localStorage.getItem(storageKey(initialPrompt, sessionId)); - if (saved) { - const parsed: ImageGenResult = JSON.parse(saved); - if (parsed.images && parsed.images.length > 0) { - setResult(parsed); - setStatus('completed'); - } - } - } catch { - // ignore - } - }, [initialPrompt, sessionId]); const handleStop = useCallback(() => { if (abortRef.current) { @@ -127,23 +112,9 @@ export function ImageGenConfirmation({ setResult(genResult); setStatus('completed'); - // Store lightweight version in localStorage - try { - const storable = { - id: genResult.id, - text: genResult.text, - images: genResult.images.map(img => ({ - mimeType: img.mimeType, - localPath: img.localPath, - data: '', - })), - }; - localStorage.setItem(storageKey(initialPrompt, sessionId), JSON.stringify(storable)); - } catch { - // storage full - } - - // Persist result to DB by replacing image-gen-request with image-gen-result + // Persist result to DB by replacing image-gen-request with image-gen-result. + // During streaming the assistant message may not yet be in DB (no messageId), + // so retry once after a short delay to give the stream time to complete. { const resultBlock = JSON.stringify({ status: 'completed', @@ -155,16 +126,29 @@ export function ImageGenConfirmation({ localPath: img.localPath, })), }); - fetch('/api/chat/messages', { + const persistBody = { + message_id: messageId || '', + content: '```image-gen-result\n' + resultBlock + '\n```', + session_id: sessionId, + prompt_hint: initialPrompt, + // Pass the raw block for exact content matching when messageId is unavailable + raw_request_block: rawRequestBlock, + }; + const doPut = () => fetch('/api/chat/messages', { method: 'PUT', headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ - message_id: messageId || '', - content: '```image-gen-result\n' + resultBlock + '\n```', - session_id: sessionId, - prompt_hint: initialPrompt, - }), - }).catch(() => {}); + body: JSON.stringify(persistBody), + }); + doPut().then(r => { + if (!r.ok && !messageId) { + // Retry after 3s — message should be persisted by then + setTimeout(() => doPut().catch(() => {}), 3000); + } + }).catch(() => { + if (!messageId) { + setTimeout(() => doPut().catch(() => {}), 3000); + } + }); } // Defer event dispatch so React commits setResult/setStatus before @@ -199,12 +183,7 @@ export function ImageGenConfirmation({ const handleRegenerate = useCallback(() => { setResult(null); setStatus('idle'); - try { - localStorage.removeItem(storageKey(initialPrompt, sessionId)); - } catch { - // ignore - } - }, [initialPrompt, sessionId]); + }, []); // ── Completed: show result only ── if (status === 'completed' && result && result.images.length > 0) { diff --git a/src/components/chat/MessageItem.tsx b/src/components/chat/MessageItem.tsx index 2f18d84..219ea6d 100644 --- a/src/components/chat/MessageItem.tsx +++ b/src/components/chat/MessageItem.tsx @@ -26,7 +26,7 @@ interface ImageGenRequest { useLastGenerated?: boolean; } -function parseImageGenRequest(text: string): { beforeText: string; request: ImageGenRequest; afterText: string } | null { +function parseImageGenRequest(text: string): { beforeText: string; request: ImageGenRequest; afterText: string; rawBlock: string } | null { const regex = /```image-gen-request\s*\n?([\s\S]*?)\n?\s*```/; const match = text.match(regex); if (!match) return null; @@ -55,6 +55,7 @@ function parseImageGenRequest(text: string): { beforeText: string; request: Imag useLastGenerated: json.useLastGenerated === true, }, afterText, + rawBlock: match[0], }; } catch { return null; @@ -126,6 +127,7 @@ function parseBatchPlan(text: string): { beforeText: string; plan: PlannerOutput interface MessageItemProps { message: Message; + sessionId?: string; } interface ToolBlock { @@ -317,7 +319,7 @@ function TokenUsageDisplay({ usage }: { usage: TokenUsage }) { const COLLAPSE_HEIGHT = 300; -export const MessageItem = memo(function MessageItem({ message }: MessageItemProps) { +export const MessageItem = memo(function MessageItem({ message, sessionId }: MessageItemProps) { const isUser = message.role === 'user'; // Collapse/expand state for long user messages (hooks must be called unconditionally) @@ -427,7 +429,7 @@ export const MessageItem = memo(function MessageItem({ message }: MessageItemPro )} - ) : + ) : )} @@ -445,7 +447,7 @@ export const MessageItem = memo(function MessageItem({ message }: MessageItemPro * Memoized assistant message content — avoids re-running parseBatchPlan / parseImageGenResult / * parseImageGenRequest on every render when only unrelated props change. */ -const AssistantContent = memo(function AssistantContent({ displayText, messageId }: { displayText: string; messageId: string }) { +const AssistantContent = memo(function AssistantContent({ displayText, messageId, sessionId }: { displayText: string; messageId: string; sessionId?: string }) { return useMemo(() => { // Try batch-plan first (Image Agent batch mode) const batchPlanResult = parseBatchPlan(displayText); @@ -512,6 +514,7 @@ const AssistantContent = memo(function AssistantContent({ displayText, messageId if (parsed) { const refs = buildReferenceImages( messageId, + sessionId || '', parsed.request.useLastGenerated || false, parsed.request.referenceImages, ); @@ -520,9 +523,11 @@ const AssistantContent = memo(function AssistantContent({ displayText, messageId {parsed.beforeText && {parsed.beforeText}} 0 ? refs : undefined} /> {parsed.afterText && {parsed.afterText}} @@ -535,5 +540,5 @@ const AssistantContent = memo(function AssistantContent({ displayText, messageId .replace(/```batch-plan[\s\S]*?```/g, '') .trim(); return stripped ? {stripped} : null; - }, [displayText, messageId]); + }, [displayText, messageId, sessionId]); }); diff --git a/src/components/chat/MessageList.tsx b/src/components/chat/MessageList.tsx index f0866b2..1cd1fe6 100644 --- a/src/components/chat/MessageList.tsx +++ b/src/components/chat/MessageList.tsx @@ -256,7 +256,7 @@ export function MessageList({ return (
- + {rewindSdkUuid && sessionId && !isStreaming && ( )} @@ -268,6 +268,7 @@ export function MessageList({ {parsed.beforeText && {parsed.beforeText}} 0 ? refs : undefined} /> {parsed.afterText && {parsed.afterText}} diff --git a/src/lib/db.ts b/src/lib/db.ts index 3569852..2fbfb53 100644 --- a/src/lib/db.ts +++ b/src/lib/db.ts @@ -413,6 +413,9 @@ function migrateDb(db: Database.Database): void { if (!provColNames.includes('role_models_json')) { db.exec("ALTER TABLE api_providers ADD COLUMN role_models_json TEXT NOT NULL DEFAULT '{}'"); } + if (!provColNames.includes('options_json')) { + db.exec("ALTER TABLE api_providers ADD COLUMN options_json TEXT NOT NULL DEFAULT '{}'"); + } } // Create provider_models table @@ -872,25 +875,46 @@ export function updateMessageContent(messageId: string, content: string): number } /** - * Find the most recent assistant message in a session that contains a given text snippet, + * Find the most recent assistant message in a session that contains an image-gen-request, * update its content, and return the real message ID. Used as fallback when the frontend * only has a temporary message ID. + * + * Prefers exact match on rawRequestBlock (the full ```image-gen-request...``` fence). + * Falls back to prompt hint prefix match if rawRequestBlock is unavailable or doesn't match. */ export function updateMessageBySessionAndHint( sessionId: string, - promptHint: string, content: string, + rawRequestBlock?: string, + promptHint?: string, ): { changes: number; messageId?: string } { const db = getDb(); - // Find the latest assistant message containing the prompt hint within an image-gen-request block - const row = db.prepare( - "SELECT id FROM messages WHERE session_id = ? AND role = 'assistant' AND content LIKE '%image-gen-request%' AND content LIKE ? ORDER BY created_at DESC LIMIT 1" - ).get(sessionId, `%${promptHint.slice(0, 60)}%`) as { id: string } | undefined; - if (!row) return { changes: 0 }; + // Strategy 1: Exact match on the raw ```image-gen-request...``` block content. + // This is unambiguous even when multiple requests share the same prompt. + if (rawRequestBlock) { + const escaped = rawRequestBlock.replace(/%/g, '\\%').replace(/_/g, '\\_'); + const row = db.prepare( + "SELECT id FROM messages WHERE session_id = ? AND role = 'assistant' AND content LIKE ? ESCAPE '\\' AND content NOT LIKE '%image-gen-result%' ORDER BY created_at DESC LIMIT 1" + ).get(sessionId, `%${escaped}%`) as { id: string } | undefined; + if (row) { + const result = db.prepare('UPDATE messages SET content = ? WHERE id = ?').run(content, row.id); + return { changes: result.changes, messageId: row.id }; + } + } - const result = db.prepare('UPDATE messages SET content = ? WHERE id = ?').run(content, row.id); - return { changes: result.changes, messageId: row.id }; + // Strategy 2: Fallback to prompt hint prefix match (legacy path). + if (promptHint) { + const row = db.prepare( + "SELECT id FROM messages WHERE session_id = ? AND role = 'assistant' AND content LIKE '%image-gen-request%' AND content NOT LIKE '%image-gen-result%' AND content LIKE ? ORDER BY created_at DESC LIMIT 1" + ).get(sessionId, `%${promptHint.slice(0, 60)}%`) as { id: string } | undefined; + if (row) { + const result = db.prepare('UPDATE messages SET content = ? WHERE id = ?').run(content, row.id); + return { changes: result.changes, messageId: row.id }; + } + } + + return { changes: 0 }; } export function clearSessionMessages(sessionId: string): void { @@ -1055,8 +1079,8 @@ export function createProvider(data: CreateProviderRequest): ApiProvider { const sortOrder = (maxRow.max_order ?? -1) + 1; db.prepare( - `INSERT INTO api_providers (id, name, provider_type, protocol, base_url, api_key, is_active, sort_order, extra_env, headers_json, env_overrides_json, role_models_json, notes, created_at, updated_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)` + `INSERT INTO api_providers (id, name, provider_type, protocol, base_url, api_key, is_active, sort_order, extra_env, headers_json, env_overrides_json, role_models_json, options_json, notes, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)` ).run( id, data.name, @@ -1070,6 +1094,7 @@ export function createProvider(data: CreateProviderRequest): ApiProvider { data.headers_json || '{}', data.env_overrides_json || '', data.role_models_json || '{}', + data.options_json || '{}', data.notes || '', now, now, @@ -1093,14 +1118,15 @@ export function updateProvider(id: string, data: UpdateProviderRequest): ApiProv const headersJson = data.headers_json ?? existing.headers_json; const envOverridesJson = data.env_overrides_json ?? existing.env_overrides_json; const roleModelsJson = data.role_models_json ?? existing.role_models_json; + const optionsJson = data.options_json ?? existing.options_json; const notes = data.notes ?? existing.notes; const sortOrder = data.sort_order ?? existing.sort_order; db.prepare( `UPDATE api_providers SET name = ?, provider_type = ?, protocol = ?, base_url = ?, api_key = ?, - extra_env = ?, headers_json = ?, env_overrides_json = ?, role_models_json = ?, + extra_env = ?, headers_json = ?, env_overrides_json = ?, role_models_json = ?, options_json = ?, notes = ?, sort_order = ?, updated_at = ? WHERE id = ?` - ).run(name, providerType, protocol, baseUrl, apiKey, extraEnv, headersJson, envOverridesJson, roleModelsJson, notes, sortOrder, now, id); + ).run(name, providerType, protocol, baseUrl, apiKey, extraEnv, headersJson, envOverridesJson, roleModelsJson, optionsJson, notes, sortOrder, now, id); return getProvider(id); } @@ -1111,6 +1137,41 @@ export function deleteProvider(id: string): boolean { return result.changes > 0; } +// ── Provider Options ──────────────────────────────────────────── + +/** + * Get options for a provider. For 'env' provider, reads from settings table. + * For DB providers, reads from options_json column. + */ +export function getProviderOptions(providerId: string): import('@/types').ProviderOptions { + if (providerId === 'env') { + const thinkingMode = getSetting('thinking_mode') || 'adaptive'; + const context1m = getSetting('context_1m') === 'true'; + return { thinking_mode: thinkingMode as 'adaptive' | 'enabled' | 'disabled', context_1m: context1m }; + } + const provider = getProvider(providerId); + if (!provider) return {}; + try { + return JSON.parse(provider.options_json || '{}'); + } catch { return {}; } +} + +/** + * Set options for a provider. For 'env' provider, writes to settings table. + * For DB providers, writes to options_json column. + */ +export function setProviderOptions(providerId: string, options: import('@/types').ProviderOptions): void { + if (providerId === 'env') { + if (options.thinking_mode !== undefined) setSetting('thinking_mode', options.thinking_mode); + if (options.context_1m !== undefined) setSetting('context_1m', options.context_1m ? 'true' : ''); + return; + } + const db = getDb(); + const now = new Date().toISOString().replace('T', ' ').split('.')[0]; + db.prepare('UPDATE api_providers SET options_json = ?, updated_at = ? WHERE id = ?') + .run(JSON.stringify(options), now, providerId); +} + // ── Provider Models ───────────────────────────────────────────── export function getModelsForProvider(providerId: string): import('@/types').ProviderModel[] { diff --git a/src/lib/image-ref-store.ts b/src/lib/image-ref-store.ts index 5211e8a..0af5b6d 100644 --- a/src/lib/image-ref-store.ts +++ b/src/lib/image-ref-store.ts @@ -3,28 +3,41 @@ * * Keys: * PENDING_KEY – images uploaded via the input bar, not yet bound to a message - * LAST_GENERATED_KEY – file paths of the most recently generated images (sessionStorage-backed) + * lastGenKey(sid) – file paths of the most recently generated images, scoped per session * – images bound to a specific assistant message */ import type { ReferenceImage } from '@/types'; export const PENDING_KEY = '__pending__'; -export const LAST_GENERATED_KEY = '__last_generated__'; const store = new Map(); // ── sessionStorage persistence for last-generated paths ── -const SS_KEY = 'imgref:last_generated'; +/** Build a session-scoped key for the in-memory store. */ +function lastGenKey(sessionId: string): string { + return `__last_generated__:${sessionId}`; +} -function loadLastGenerated(): void { - if (typeof window === 'undefined') return; +/** Build a session-scoped sessionStorage key. */ +function ssKey(sessionId: string): string { + return `imgref:last_generated:${sessionId}`; +} + +/** + * Restore last-generated images for a specific session from sessionStorage. + * Called when a ChatView mounts with a known sessionId. + */ +export function loadLastGenerated(sessionId: string): void { + if (typeof window === 'undefined' || !sessionId) return; + const key = lastGenKey(sessionId); + if (store.has(key)) return; // already loaded try { - const raw = sessionStorage.getItem(SS_KEY); + const raw = sessionStorage.getItem(ssKey(sessionId)); if (raw) { const arr: ReferenceImage[] = JSON.parse(raw); if (Array.isArray(arr) && arr.length > 0) { - store.set(LAST_GENERATED_KEY, arr); + store.set(key, arr); } } } catch { @@ -32,9 +45,6 @@ function loadLastGenerated(): void { } } -// Auto-restore on module load (client only) -loadLastGenerated(); - // ── Public helpers ── /** Set reference images for a given key. */ @@ -69,15 +79,15 @@ export function transferPendingToMessage(messageId: string): void { } /** - * Store generated image paths and persist to sessionStorage. + * Store generated image paths, scoped to a session, and persist to sessionStorage. * Called when image generation completes. */ -export function setLastGeneratedImages(paths: string[]): void { +export function setLastGeneratedImages(sessionId: string, paths: string[]): void { const images: ReferenceImage[] = paths.map(p => ({ mimeType: 'image/png', localPath: p })); - store.set(LAST_GENERATED_KEY, images); + store.set(lastGenKey(sessionId), images); if (typeof window !== 'undefined') { try { - sessionStorage.setItem(SS_KEY, JSON.stringify(images)); + sessionStorage.setItem(ssKey(sessionId), JSON.stringify(images)); } catch { // storage full } @@ -88,11 +98,13 @@ export function setLastGeneratedImages(paths: string[]): void { * Single merge entry-point: build a unified ReferenceImage[] for a given context. * * @param key Store key (PENDING_KEY or message.id) + * @param sessionId Session ID to scope last-generated lookup * @param useLastGenerated Whether the LLM requested editing last-generated images * @param extraPaths Additional file paths from the parsed request (referenceImages field) */ export function buildReferenceImages( key: string, + sessionId: string, useLastGenerated: boolean, extraPaths?: string[], ): ReferenceImage[] { @@ -104,9 +116,9 @@ export function buildReferenceImages( result.push(...uploaded); } - // 2. Last-generated images (if LLM requested) - if (useLastGenerated) { - const lastGen = store.get(LAST_GENERATED_KEY); + // 2. Last-generated images (if LLM requested), scoped to this session + if (useLastGenerated && sessionId) { + const lastGen = store.get(lastGenKey(sessionId)); if (lastGen) { result.push(...lastGen); }