diff --git a/package-lock.json b/package-lock.json
index 72f3c79..509eb9d 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -1,12 +1,12 @@
 {
   "name": "codepilot",
-  "version": "0.17.2",
+  "version": "0.17.3",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "codepilot",
-      "version": "0.17.2",
+      "version": "0.17.3",
       "dependencies": {
         "@ai-sdk/anthropic": "^3.0.47",
         "@ai-sdk/google": "^3.0.31",
diff --git a/package.json b/package.json
index 8695091..9d49cce 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "codepilot",
-  "version": "0.17.2",
+  "version": "0.17.3",
   "private": true,
   "author": {
     "name": "op7418",
diff --git a/src/app/api/chat/messages/route.ts b/src/app/api/chat/messages/route.ts
index 7057c77..acca080 100644
--- a/src/app/api/chat/messages/route.ts
+++ b/src/app/api/chat/messages/route.ts
@@ -1,5 +1,5 @@
 import { NextRequest, NextResponse } from 'next/server';
-import { addMessage, updateMessageContent, getSession } from '@/lib/db';
+import { addMessage, updateMessageContent, updateMessageBySessionAndHint, getSession } from '@/lib/db';
 
 export const runtime = 'nodejs';
 
@@ -42,24 +42,54 @@ export async function POST(request: NextRequest) {
  * PUT /api/chat/messages
  * Update the content of an existing message.
  * Used to replace image-gen-request with image-gen-result after generation.
+ *
+ * Tries message_id first. If no rows updated (temp ID), falls back to
+ * searching by session_id + prompt_hint for the real DB message.
  */
 export async function PUT(request: NextRequest) {
   try {
     const body = await request.json();
-    const { message_id, content } = body as {
+    const { message_id, content, session_id, prompt_hint } = body as {
       message_id: string;
       content: string;
+      session_id?: string;
+      prompt_hint?: string;
     };
 
-    if (!message_id || !content) {
+    if (!content) {
       return NextResponse.json(
-        { error: 'message_id and content are required' },
+        { error: 'content is required' },
         { status: 400 },
       );
     }
 
-    updateMessageContent(message_id, content);
-    return NextResponse.json({ ok: true });
+    // Try direct update by message_id
+    let changes = 0;
+    let updatedMessageId = message_id;
+    let fallbackUsed = false;
+
+    if (message_id) {
+      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);
+      changes = result.changes;
+      if (result.messageId) {
+        updatedMessageId = result.messageId;
+        fallbackUsed = true;
+      }
+    }
+
+    if (changes === 0) {
+      return NextResponse.json(
+        { error: 'Message not found' },
+        { status: 404 },
+      );
+    }
+
+    return NextResponse.json({ ok: true, updated_message_id: updatedMessageId, fallback_used: fallbackUsed });
   } catch (error) {
     const msg = error instanceof Error ? error.message : 'Failed to update message';
     return NextResponse.json({ error: msg }, { status: 500 });
diff --git a/src/components/chat/ImageGenConfirmation.tsx b/src/components/chat/ImageGenConfirmation.tsx
index ebeb9e9..be59efe 100644
--- a/src/components/chat/ImageGenConfirmation.tsx
+++ b/src/components/chat/ImageGenConfirmation.tsx
@@ -143,7 +143,7 @@ export function ImageGenConfirmation({
         }
 
         // Persist result to DB by replacing image-gen-request with image-gen-result
-        if (messageId) {
+        {
           const resultBlock = JSON.stringify({
             status: 'completed',
             prompt,
@@ -158,8 +158,10 @@ export function ImageGenConfirmation({
             method: 'PUT',
             headers: { 'Content-Type': 'application/json' },
             body: JSON.stringify({
-              message_id: messageId,
+              message_id: messageId || '',
               content: '```image-gen-result\n' + resultBlock + '\n```',
+              session_id: sessionId,
+              prompt_hint: initialPrompt,
             }),
           }).catch(() => {});
         }
diff --git a/src/lib/db.ts b/src/lib/db.ts
index a5e3f13..dd0b1d6 100644
--- a/src/lib/db.ts
+++ b/src/lib/db.ts
@@ -565,9 +565,32 @@ export function addMessage(
   return db.prepare('SELECT * FROM messages WHERE id = ?').get(id) as Message;
 }
 
-export function updateMessageContent(messageId: string, content: string): void {
+export function updateMessageContent(messageId: string, content: string): number {
   const db = getDb();
-  db.prepare('UPDATE messages SET content = ? WHERE id = ?').run(content, messageId);
+  const result = db.prepare('UPDATE messages SET content = ? WHERE id = ?').run(content, messageId);
+  return result.changes;
+}
+
+/**
+ * Find the most recent assistant message in a session that contains a given text snippet,
+ * update its content, and return the real message ID. Used as fallback when the frontend
+ * only has a temporary message ID.
+ */
+export function updateMessageBySessionAndHint(
+  sessionId: string,
+  promptHint: string,
+  content: 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 };
+
+  const result = db.prepare('UPDATE messages SET content = ? WHERE id = ?').run(content, row.id);
+  return { changes: result.changes, messageId: row.id };
 }
 
 export function clearSessionMessages(sessionId: string): void {