From 13e22c1e05b8aa61eecdb15dfd4c4bee75bc02b0 Mon Sep 17 00:00:00 2001
From: Kurvaz <fillthecatoffical@gmail.com>
Date: Wed, 7 Jan 2026 15:49:37 -0700
Subject: [PATCH] Added persistence to retry

---
 .../migrations/006_story_retry_state.sql      |   2 +
 src-tauri/src/lib.rs                          |   6 +
 src/lib/components/story/ActionInput.svelte   | 160 +++++++++---
 src/lib/services/database.ts                  |  28 +++
 src/lib/services/export.ts                    |   1 +
 src/lib/stores/story.svelte.ts                | 107 +++++++-
 src/lib/stores/ui.svelte.ts                   | 234 ++++++++++++++++--
 src/lib/types/index.ts                        |  22 ++
 8 files changed, 501 insertions(+), 59 deletions(-)
 create mode 100644 src-tauri/migrations/006_story_retry_state.sql

diff --git a/src-tauri/migrations/006_story_retry_state.sql b/src-tauri/migrations/006_story_retry_state.sql
new file mode 100644
index 0000000..7cba5aa
--- /dev/null
+++ b/src-tauri/migrations/006_story_retry_state.sql
@@ -0,0 +1,2 @@
+-- Add retry_state column to stories for persistent retry functionality
+ALTER TABLE stories ADD COLUMN retry_state TEXT;
diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs
index 9f9f70a..454bd5a 100644
--- a/src-tauri/src/lib.rs
+++ b/src-tauri/src/lib.rs
@@ -40,6 +40,12 @@ pub fn run() {
             sql: include_str!("../migrations/005_story_beats_resolved_at.sql"),
             kind: MigrationKind::Up,
         },
+        Migration {
+            version: 6,
+            description: "add_story_retry_state",
+            sql: include_str!("../migrations/006_story_retry_state.sql"),
+            kind: MigrationKind::Up,
+        },
     ];
 
     tauri::Builder::default()
diff --git a/src/lib/components/story/ActionInput.svelte b/src/lib/components/story/ActionInput.svelte
index 0705f36..8d8ec2b 100644
--- a/src/lib/components/story/ActionInput.svelte
+++ b/src/lib/components/story/ActionInput.svelte
@@ -27,6 +27,7 @@
   let isRawActionChoice = $state(false); // True when submitting an AI-generated choice (no prefix/suffix)
   let stopRequested = false;
   let activeAbortController: AbortController | null = null;
+  let isRetryingLastMessage = $state(false); // Hide stop button during completed-message retries
 
   // In creative writing mode, show different input style
   const isCreativeMode = $derived(story.storyMode === 'creative-writing');
@@ -880,6 +881,10 @@
       log('Stop ignored (not generating)');
       return;
     }
+    if (isRetryingLastMessage) {
+      log('Stop ignored (retrying completed message)');
+      return;
+    }
 
     stopRequested = true;
     activeAbortController?.abort();
@@ -903,18 +908,42 @@
     ui.clearGenerationError();
     ui.clearSuggestions(story.currentStory?.id);
     ui.clearActionChoices(story.currentStory?.id);
-    ui.restoreActivationData(backup.activationData, backup.storyPosition);
+    if (backup.hasFullState) {
+      ui.restoreActivationData(backup.activationData, backup.storyPosition);
+    }
     ui.setLastLorebookRetrieval(null);
 
     try {
-      await story.restoreFromRetryBackup({
-        entries: backup.entries,
-        characters: backup.characters,
-        locations: backup.locations,
-        items: backup.items,
-        storyBeats: backup.storyBeats,
-        lorebookEntries: backup.lorebookEntries,
-      });
+      if (backup.hasFullState) {
+        await story.restoreFromRetryBackup({
+          entries: backup.entries,
+          characters: backup.characters,
+          locations: backup.locations,
+          items: backup.items,
+          storyBeats: backup.storyBeats,
+          lorebookEntries: backup.lorebookEntries,
+        });
+      } else {
+        // Persistent restore - delete entries and entities created after backup
+        // Clear activation data but don't save yet - let the next action rebuild it
+        ui.clearActivationData();
+
+        log('Persistent stop restore: deleting entries from position', backup.entryCountBeforeAction);
+        await story.deleteEntriesFromPosition(backup.entryCountBeforeAction);
+
+        if (backup.hasEntityIds) {
+          log('Persistent stop restore: deleting entities created after backup');
+          await story.deleteEntitiesCreatedAfterBackup({
+            characterIds: backup.characterIds,
+            locationIds: backup.locationIds,
+            itemIds: backup.itemIds,
+            storyBeatIds: backup.storyBeatIds,
+            lorebookEntryIds: backup.lorebookEntryIds,
+          });
+        } else {
+          log('Persistent stop restore: skipping entity cleanup (no ID snapshot)');
+        }
+      }
 
       await tick();
       actionType = backup.actionType;
@@ -924,7 +953,7 @@
       log('Stop restore failed', error);
       console.error('Stop restore failed:', error);
     } finally {
-      ui.clearRetryBackup();
+      ui.clearRetryBackup(true); // Clear from DB since user explicitly stopped
     }
   }
 
@@ -970,6 +999,7 @@
   /**
    * Retry the last user message by restoring to the backup state
    * and regenerating with the same user action.
+   * Supports both full state restore (in-memory backup) and entry-only restore (persistent backup).
    */
   async function handleRetryLastMessage() {
     log('handleRetryLastMessage called', {
@@ -987,12 +1017,14 @@
     // Verify backup is for current story
     if (backup.storyId !== story.currentStory.id) {
       log('Backup is for different story, clearing');
-      ui.clearRetryBackup();
+      ui.clearRetryBackup(false); // Just clear in-memory, don't touch DB
       return;
     }
 
     log('Restoring from backup and regenerating', {
+      hasFullState: backup.hasFullState,
       backupEntriesCount: backup.entries.length,
+      entryCountBeforeAction: backup.entryCountBeforeAction,
       currentEntriesCount: story.entries.length,
       userAction: backup.userActionContent.substring(0, 50),
     });
@@ -1004,23 +1036,46 @@
     ui.clearSuggestions(story.currentStory?.id);
     ui.clearActionChoices(story.currentStory?.id);
 
-    // Restore activation data from backup to preserve lorebook stickiness state
-    // This ensures entries that were "sticky" before the user action remain sticky
-    ui.restoreActivationData(backup.activationData, backup.storyPosition);
-
     // Clear lorebook retrieval debug state since it's now stale
     ui.setLastLorebookRetrieval(null);
 
     try {
-      // Restore story state from backup
-      await story.restoreFromRetryBackup({
-        entries: backup.entries,
-        characters: backup.characters,
-        locations: backup.locations,
-        items: backup.items,
-        storyBeats: backup.storyBeats,
-        lorebookEntries: backup.lorebookEntries,
-      });
+      if (backup.hasFullState) {
+        // Full state restore (in-memory backup with snapshots)
+        // Restore activation data from backup to preserve lorebook stickiness state
+        ui.restoreActivationData(backup.activationData, backup.storyPosition);
+
+        // Restore story state from backup
+        await story.restoreFromRetryBackup({
+          entries: backup.entries,
+          characters: backup.characters,
+          locations: backup.locations,
+          items: backup.items,
+          storyBeats: backup.storyBeats,
+          lorebookEntries: backup.lorebookEntries,
+        });
+      } else {
+        // Persistent restore (backup without full snapshots, but with entity IDs)
+        // Clear activation data but don't save yet - generation will rebuild it
+        ui.clearActivationData();
+
+        log('Persistent restore: deleting entries from position', backup.entryCountBeforeAction);
+        await story.deleteEntriesFromPosition(backup.entryCountBeforeAction);
+
+        if (backup.hasEntityIds) {
+          // Delete entities that were created after the backup (AI extractions)
+          log('Persistent restore: deleting entities created after backup');
+          await story.deleteEntitiesCreatedAfterBackup({
+            characterIds: backup.characterIds,
+            locationIds: backup.locationIds,
+            itemIds: backup.itemIds,
+            storyBeatIds: backup.storyBeatIds,
+            lorebookEntryIds: backup.lorebookEntryIds,
+          });
+        } else {
+          log('Persistent restore: skipping entity cleanup (no ID snapshot)');
+        }
+      }
 
       // Wait for state to sync
       await tick();
@@ -1037,7 +1092,12 @@
 
       // Regenerate
       if (!settings.needsApiKey) {
-        await generateResponse(userActionEntry.id, backup.userActionContent);
+        isRetryingLastMessage = true;
+        try {
+          await generateResponse(userActionEntry.id, backup.userActionContent);
+        } finally {
+          isRetryingLastMessage = false;
+        }
       }
     } catch (error) {
       log('Retry last message failed', error);
@@ -1049,7 +1109,7 @@
    * Dismiss the retry backup (user doesn't want to retry)
    */
   function dismissRetryBackup() {
-    ui.clearRetryBackup();
+    ui.clearRetryBackup(true); // Clear from DB since user explicitly dismissed
   }
 
   function handleKeydown(event: KeyboardEvent) {
@@ -1112,13 +1172,23 @@
         ></textarea>
       </div>
       {#if ui.isGenerating}
-        <button
-          onclick={handleStopGeneration}
-          class="btn self-stretch px-3 sm:px-4 py-3 min-h-[44px] min-w-[44px] bg-red-500/20 text-red-400 hover:bg-red-500/30"
-          title="Stop generation"
-        >
-          <Square class="h-5 w-5" />
-        </button>
+        {#if !isRetryingLastMessage}
+          <button
+            onclick={handleStopGeneration}
+            class="btn self-stretch px-3 sm:px-4 py-3 min-h-[44px] min-w-[44px] bg-red-500/20 text-red-400 hover:bg-red-500/30"
+            title="Stop generation"
+          >
+            <Square class="h-5 w-5" />
+          </button>
+        {:else}
+          <button
+            disabled
+            class="btn self-stretch px-3 sm:px-4 py-3 min-h-[44px] min-w-[44px] bg-red-500/20 text-red-400 opacity-50 cursor-not-allowed"
+            title="Stop disabled during retry"
+          >
+            <Square class="h-5 w-5" />
+          </button>
+        {/if}
       {:else}
         <button
           onclick={handleSubmit}
@@ -1196,13 +1266,23 @@
         ></textarea>
       </div>
       {#if ui.isGenerating}
-        <button
-          onclick={handleStopGeneration}
-          class="btn self-stretch px-3 sm:px-4 py-3 min-h-[44px] min-w-[44px] bg-red-500/20 text-red-400 hover:bg-red-500/30"
-          title="Stop generation"
-        >
-          <Square class="h-5 w-5" />
-        </button>
+        {#if !isRetryingLastMessage}
+          <button
+            onclick={handleStopGeneration}
+            class="btn self-stretch px-3 sm:px-4 py-3 min-h-[44px] min-w-[44px] bg-red-500/20 text-red-400 hover:bg-red-500/30"
+            title="Stop generation"
+          >
+            <Square class="h-5 w-5" />
+          </button>
+        {:else}
+          <button
+            disabled
+            class="btn self-stretch px-3 sm:px-4 py-3 min-h-[44px] min-w-[44px] bg-red-500/20 text-red-400 opacity-50 cursor-not-allowed"
+            title="Stop disabled during retry"
+          >
+            <Square class="h-5 w-5" />
+          </button>
+        {/if}
       {:else}
         <button
           onclick={handleSubmit}
diff --git a/src/lib/services/database.ts b/src/lib/services/database.ts
index f74276d..06bc4ad 100644
--- a/src/lib/services/database.ts
+++ b/src/lib/services/database.ts
@@ -14,6 +14,7 @@ import type {
   EntryType,
   EntryState,
   EntryPreview,
+  PersistentRetryState,
 } from '$lib/types';
 
 class DatabaseService {
@@ -119,6 +120,10 @@ class DatabaseService {
       setClauses.push('memory_config = ?');
       values.push(updates.memoryConfig ? JSON.stringify(updates.memoryConfig) : null);
     }
+    if (updates.retryState !== undefined) {
+      setClauses.push('retry_state = ?');
+      values.push(updates.retryState ? JSON.stringify(updates.retryState) : null);
+    }
 
     values.push(id);
     await db.execute(
@@ -127,6 +132,28 @@ class DatabaseService {
     );
   }
 
+  /**
+   * Save retry state for a story.
+   */
+  async saveRetryState(storyId: string, retryState: PersistentRetryState): Promise<void> {
+    const db = await this.getDb();
+    await db.execute(
+      'UPDATE stories SET retry_state = ? WHERE id = ?',
+      [JSON.stringify(retryState), storyId]
+    );
+  }
+
+  /**
+   * Clear retry state for a story.
+   */
+  async clearRetryState(storyId: string): Promise<void> {
+    const db = await this.getDb();
+    await db.execute(
+      'UPDATE stories SET retry_state = NULL WHERE id = ?',
+      [storyId]
+    );
+  }
+
   async deleteStory(id: string): Promise<void> {
     const db = await this.getDb();
     await db.execute('DELETE FROM stories WHERE id = ?', [id]);
@@ -844,6 +871,7 @@ class DatabaseService {
       updatedAt: row.updated_at,
       settings: row.settings ? JSON.parse(row.settings) : null,
       memoryConfig: row.memory_config ? JSON.parse(row.memory_config) : null,
+      retryState: row.retry_state ? JSON.parse(row.retry_state) : null,
     };
   }
 
diff --git a/src/lib/services/export.ts b/src/lib/services/export.ts
index 3931e23..428ba78 100644
--- a/src/lib/services/export.ts
+++ b/src/lib/services/export.ts
@@ -225,6 +225,7 @@ class ExportService {
         mode: data.story.mode || 'adventure',
         settings: data.story.settings,
         memoryConfig: data.story.memoryConfig || null,
+        retryState: null, // Clear retry state on import
       };
 
       await database.createStory(importedStory);
diff --git a/src/lib/stores/story.svelte.ts b/src/lib/stores/story.svelte.ts
index b999ca3..59ec597 100644
--- a/src/lib/stores/story.svelte.ts
+++ b/src/lib/stores/story.svelte.ts
@@ -301,8 +301,13 @@ class StoryStore {
     // Load persisted activation data for this story (stickiness tracking)
     await ui.loadActivationData(storyId);
 
-    // Clear retry backup from previous story
-    ui.clearRetryBackup();
+    // Set current story ID for retry backup tracking
+    ui.setCurrentRetryStoryId(storyId);
+
+    // Load retry state from DB if we don't have an in-memory backup for this story
+    if (story.retryState) {
+      ui.loadRetryBackupFromPersistent(storyId, story.retryState);
+    }
 
     // Validate and repair chapter integrity (handles orphaned references)
     await this.validateChapterIntegrity();
@@ -340,6 +345,7 @@ class StoryStore {
       mode,
       settings: null,
       memoryConfig: DEFAULT_MEMORY_CONFIG,
+      retryState: null,
     });
 
     this.allStories = [storyData, ...this.allStories];
@@ -370,6 +376,7 @@ class StoryStore {
       mode,
       settings: settings ?? null,
       memoryConfig: DEFAULT_MEMORY_CONFIG,
+      retryState: null,
     });
 
     this.allStories = [storyData, ...this.allStories];
@@ -503,6 +510,97 @@ class StoryStore {
     await database.updateStory(this.currentStory.id, {});
   }
 
+  /**
+   * Delete all entries from a given position onward.
+   * Used for entry-only retry restore (persistent retry).
+   */
+  async deleteEntriesFromPosition(position: number): Promise<void> {
+    if (!this.currentStory) throw new Error('No story loaded');
+
+    // Find entries to delete (position >= the given position)
+    const entriesToDelete = this.entries.filter(e => e.position >= position);
+
+    log('Deleting entries from position', {
+      position,
+      entriesToDelete: entriesToDelete.length,
+      totalEntries: this.entries.length,
+    });
+
+    // Delete from database
+    for (const entry of entriesToDelete) {
+      await database.deleteStoryEntry(entry.id);
+    }
+
+    // Update in-memory state
+    this.entries = this.entries.filter(e => e.position < position);
+
+    // Update story's updatedAt
+    await database.updateStory(this.currentStory.id, {});
+  }
+
+  /**
+   * Delete entities that were created after the backup.
+   * Used for persistent retry restore to remove AI-extracted entities.
+   * Compares current entity IDs against the saved ID lists and deletes any not in the lists.
+   */
+  async deleteEntitiesCreatedAfterBackup(savedIds: {
+    characterIds: string[];
+    locationIds: string[];
+    itemIds: string[];
+    storyBeatIds: string[];
+    lorebookEntryIds: string[];
+  }): Promise<void> {
+    if (!this.currentStory) throw new Error('No story loaded');
+
+    const characterIdsSet = new Set(savedIds.characterIds);
+    const locationIdsSet = new Set(savedIds.locationIds);
+    const itemIdsSet = new Set(savedIds.itemIds);
+    const storyBeatIdsSet = new Set(savedIds.storyBeatIds);
+    const lorebookEntryIdsSet = new Set(savedIds.lorebookEntryIds);
+
+    // Find entities to delete (not in saved lists)
+    const charactersToDelete = this.characters.filter(c => !characterIdsSet.has(c.id));
+    const locationsToDelete = this.locations.filter(l => !locationIdsSet.has(l.id));
+    const itemsToDelete = this.items.filter(i => !itemIdsSet.has(i.id));
+    const storyBeatsToDelete = this.storyBeats.filter(sb => !storyBeatIdsSet.has(sb.id));
+    const lorebookEntriesToDelete = this.lorebookEntries.filter(le => !lorebookEntryIdsSet.has(le.id));
+
+    log('Deleting entities created after backup', {
+      characters: charactersToDelete.length,
+      locations: locationsToDelete.length,
+      items: itemsToDelete.length,
+      storyBeats: storyBeatsToDelete.length,
+      lorebookEntries: lorebookEntriesToDelete.length,
+    });
+
+    // Delete from database
+    for (const character of charactersToDelete) {
+      await database.deleteCharacter(character.id);
+    }
+    for (const location of locationsToDelete) {
+      await database.deleteLocation(location.id);
+    }
+    for (const item of itemsToDelete) {
+      await database.deleteItem(item.id);
+    }
+    for (const storyBeat of storyBeatsToDelete) {
+      await database.deleteStoryBeat(storyBeat.id);
+    }
+    for (const lorebookEntry of lorebookEntriesToDelete) {
+      await database.deleteEntry(lorebookEntry.id);
+    }
+
+    // Update in-memory state
+    this.characters = this.characters.filter(c => characterIdsSet.has(c.id));
+    this.locations = this.locations.filter(l => locationIdsSet.has(l.id));
+    this.items = this.items.filter(i => itemIdsSet.has(i.id));
+    this.storyBeats = this.storyBeats.filter(sb => storyBeatIdsSet.has(sb.id));
+    this.lorebookEntries = this.lorebookEntries.filter(le => lorebookEntryIdsSet.has(le.id));
+
+    // Update story's updatedAt
+    await database.updateStory(this.currentStory.id, {});
+  }
+
   // Add a character
   async addCharacter(name: string, description?: string, relationship?: string): Promise<Character> {
     if (!this.currentStory) throw new Error('No story loaded');
@@ -1143,8 +1241,8 @@ class StoryStore {
     this.chapters = [];
     this.checkpoints = [];
 
-    // Clear retry backup since it belongs to the old story
-    ui.clearRetryBackup();
+    // Clear current retry story ID (backups are kept per-story)
+    ui.setCurrentRetryStoryId(null);
   }
 
   // Update story mode
@@ -1449,6 +1547,7 @@ class StoryStore {
         systemPromptOverride: data.systemPrompt,
       },
       memoryConfig: DEFAULT_MEMORY_CONFIG,
+      retryState: null,
     });
 
     this.allStories = [storyData, ...this.allStories];
diff --git a/src/lib/stores/ui.svelte.ts b/src/lib/stores/ui.svelte.ts
index 88b9ed5..ba68b0e 100644
--- a/src/lib/stores/ui.svelte.ts
+++ b/src/lib/stores/ui.svelte.ts
@@ -1,4 +1,4 @@
-import type { ActivePanel, SidebarTab, UIState, EntryType, StoryEntry, Character, Location, Item, StoryBeat, Entry } from '$lib/types';
+import type { ActivePanel, SidebarTab, UIState, EntryType, StoryEntry, Character, Location, Item, StoryBeat, Entry, ActionInputType } from '$lib/types';
 import type { ActionChoice } from '$lib/services/ai/actionChoices';
 import type { StorySuggestion } from '$lib/services/ai/suggestions';
 import type { StyleReviewResult } from '$lib/services/ai/styleReviewer';
@@ -6,8 +6,7 @@ import type { EntryRetrievalResult, ActivationTracker } from '$lib/services/ai/e
 import type { SyncMode } from '$lib/types/sync';
 import { SimpleActivationTracker } from '$lib/services/ai/entryRetrieval';
 import { database } from '$lib/services/database';
-
-type ActionInputType = 'do' | 'say' | 'think' | 'story' | 'free';
+import { SvelteMap } from 'svelte/reactivity';
 
 // Debug log entry for request/response logging
 export interface DebugLogEntry {
@@ -25,6 +24,7 @@ export interface RetryBackup {
   storyId: string;
   timestamp: number;
   // State snapshots (captured BEFORE user action is added)
+  // These may be empty if loaded from persistent storage (entry-only restore)
   entries: StoryEntry[];
   characters: Character[];
   locations: Location[];
@@ -39,6 +39,18 @@ export interface RetryBackup {
   // Lorebook activation tracking data (for stickiness preservation)
   activationData: Record<string, number>;
   storyPosition: number;
+  // Next entry position at time of backup - used for entry-only restore
+  entryCountBeforeAction: number;
+  // Flag indicating if this has full state snapshots (in-memory) or just entry data (from DB)
+  hasFullState: boolean;
+  // Flag indicating if entity ID snapshots are present for safe cleanup
+  hasEntityIds: boolean;
+  // Entity IDs for persistent restore - delete any entities not in these lists
+  characterIds: string[];
+  locationIds: string[];
+  itemIds: string[];
+  storyBeatIds: string[];
+  lorebookEntryIds: string[];
 }
 
 // Error state for retry functionality
@@ -86,8 +98,25 @@ class UIStore {
   // Error state for retry
   lastGenerationError = $state<GenerationError | null>(null);
 
-  // Retry backup - captures state before each user message for "retry last message" feature
-  retryBackup = $state<RetryBackup | null>(null);
+  // Retry backups - per-story backups for "retry last message" feature
+  // Stored by storyId so they persist across story switches within a session
+  private retryBackups = new SvelteMap<string, RetryBackup>();
+  private currentRetryStoryId = $state<string | null>(null);
+  retryStateWrite = Promise.resolve();
+
+  // Computed getter for current story's retry backup
+  get retryBackup(): RetryBackup | null {
+    if (!this.currentRetryStoryId) return null;
+    return this.retryBackups.get(this.currentRetryStoryId) ?? null;
+  }
+
+  /**
+   * Set the current story ID for retry backup tracking.
+   * Called when switching stories to ensure the correct backup is returned.
+   */
+  setCurrentRetryStoryId(storyId: string | null) {
+    this.currentRetryStoryId = storyId;
+  }
 
   // RPG action choices (displayed after narration)
   actionChoices = $state<ActionChoice[]>([]);
@@ -231,6 +260,7 @@ class UIStore {
    * Create a backup of the current story state before a user message.
    * This captures the state BEFORE the user action is added, so we can restore to this point.
    * Also captures lorebook activation data for stickiness preservation.
+   * Persists a lightweight version to the database for cross-session retry.
    */
   createRetryBackup(
     storyId: string,
@@ -245,10 +275,20 @@ class UIStore {
     actionType: ActionInputType,
     wasRawActionChoice: boolean
   ) {
-    // Clear old backup and create new one
-    this.retryBackup = {
+    const timestamp = Date.now();
+    const nextEntryPosition = entries.reduce((max, entry) => Math.max(max, entry.position ?? -1), -1) + 1;
+
+    // Extract entity IDs for persistent restore
+    const characterIds = characters.map(c => c.id);
+    const locationIds = locations.map(l => l.id);
+    const itemIds = items.map(i => i.id);
+    const storyBeatIds = storyBeats.map(sb => sb.id);
+    const lorebookEntryIds = lorebookEntries.map(le => le.id);
+
+    // Create new backup and store by story ID
+    const backup: RetryBackup = {
       storyId,
-      timestamp: Date.now(),
+      timestamp,
       // Deep copy arrays to avoid mutation issues
       entries: [...entries],
       characters: [...characters],
@@ -263,7 +303,38 @@ class UIStore {
       // Capture activation data for lorebook stickiness preservation
       activationData: { ...this.activationData },
       storyPosition: this.currentStoryPosition,
+      // New fields for persistent retry
+      entryCountBeforeAction: nextEntryPosition,
+      hasFullState: true,
+      hasEntityIds: true,
+      // Entity IDs for persistent restore
+      characterIds,
+      locationIds,
+      itemIds,
+      storyBeatIds,
+      lorebookEntryIds,
     };
+    this.retryBackups.set(storyId, backup);
+    this.currentRetryStoryId = storyId;
+
+    // Persist lightweight version to database (includes entity IDs for full restore)
+    this.queueRetryStateWrite(
+      () => database.saveRetryState(storyId, {
+        timestamp,
+        entryCountBeforeAction: nextEntryPosition,
+        userActionContent,
+        rawInput,
+        actionType,
+        wasRawActionChoice,
+        characterIds,
+        locationIds,
+        itemIds,
+        storyBeatIds,
+        lorebookEntryIds,
+      }),
+      'persist'
+    );
+
     console.log('[UI] Retry backup created', {
       storyId,
       entriesCount: entries.length,
@@ -273,11 +344,122 @@ class UIStore {
   }
 
   /**
-   * Clear the retry backup (called when switching stories or if user doesn't want to retry).
+   * Clear the retry backup for a story.
+   * @param clearFromDb - If true, also clears from database (use for explicit dismissal/use).
+   * @param storyId - Optional story ID. If not provided, clears the current story's backup.
    */
-  clearRetryBackup() {
-    this.retryBackup = null;
-    console.log('[UI] Retry backup cleared');
+  clearRetryBackup(clearFromDb: boolean = false, storyId?: string) {
+    const targetStoryId = storyId ?? this.currentRetryStoryId;
+
+    if (targetStoryId) {
+      this.retryBackups.delete(targetStoryId);
+
+      // Only clear from database if explicitly requested (user dismissed or used retry)
+      if (clearFromDb) {
+        this.queueRetryStateWrite(
+          () => database.clearRetryState(targetStoryId),
+          'clear'
+        );
+      }
+    }
+
+    console.log('[UI] Retry backup cleared', { clearFromDb, storyId: targetStoryId });
+  }
+
+  private queueRetryStateWrite(task: () => Promise<void>, label: string) {
+    this.retryStateWrite = this.retryStateWrite
+      .catch(() => {})
+      .then(task)
+      .catch(err => {
+        console.warn(`[UI] Failed to ${label} retry state:`, err);
+      });
+  }
+
+  /**
+   * Load retry backup from persistent state (called when a story is loaded).
+   * Creates a partial RetryBackup with hasFullState=false for entity-aware restore.
+   * Only loads if there isn't already an in-memory backup for this story.
+   */
+  loadRetryBackupFromPersistent(storyId: string, retryState: {
+    timestamp: number;
+    entryCountBeforeAction: number;
+    userActionContent: string;
+    rawInput: string;
+    actionType: ActionInputType;
+    wasRawActionChoice: boolean;
+    characterIds?: string[];
+    locationIds?: string[];
+    itemIds?: string[];
+    storyBeatIds?: string[];
+    lorebookEntryIds?: string[];
+  }) {
+    // Skip if we already have an in-memory backup for this story (it's more complete)
+    if (this.retryBackups.has(storyId)) {
+      console.log('[UI] Skipping persistent retry state load - in-memory backup exists', { storyId });
+      return;
+    }
+
+    // Validate required fields exist
+    if (
+      typeof retryState.timestamp !== 'number' ||
+      typeof retryState.entryCountBeforeAction !== 'number' ||
+      typeof retryState.userActionContent !== 'string' ||
+      typeof retryState.rawInput !== 'string' ||
+      typeof retryState.actionType !== 'string' ||
+      typeof retryState.wasRawActionChoice !== 'boolean'
+    ) {
+      console.warn('[UI] Invalid persistent retry state, skipping load', { storyId, retryState });
+      return;
+    }
+
+    const hasEntityIds = Array.isArray(retryState.characterIds)
+      && Array.isArray(retryState.locationIds)
+      && Array.isArray(retryState.itemIds)
+      && Array.isArray(retryState.storyBeatIds)
+      && Array.isArray(retryState.lorebookEntryIds);
+
+    const backup: RetryBackup = {
+      storyId,
+      timestamp: retryState.timestamp,
+      // Empty state arrays - will use ID-based restore
+      entries: [],
+      characters: [],
+      locations: [],
+      items: [],
+      storyBeats: [],
+      lorebookEntries: [],
+      // User input data
+      userActionContent: retryState.userActionContent,
+      rawInput: retryState.rawInput,
+      actionType: retryState.actionType,
+      wasRawActionChoice: retryState.wasRawActionChoice,
+      // Empty activation data
+      activationData: {},
+      storyPosition: 0,
+      // Persistent retry fields
+      entryCountBeforeAction: retryState.entryCountBeforeAction,
+      hasFullState: false, // Indicates ID-based restore
+      hasEntityIds,
+      // Entity IDs for restore - delete any entities not in these lists
+      characterIds: retryState.characterIds ?? [],
+      locationIds: retryState.locationIds ?? [],
+      itemIds: retryState.itemIds ?? [],
+      storyBeatIds: retryState.storyBeatIds ?? [],
+      lorebookEntryIds: retryState.lorebookEntryIds ?? [],
+    };
+    this.retryBackups.set(storyId, backup);
+    console.log('[UI] Retry backup loaded from persistent state', {
+      storyId,
+      entryCountBeforeAction: retryState.entryCountBeforeAction,
+      userAction: retryState.userActionContent.substring(0, 50),
+      entityIds: {
+        characters: backup.characterIds.length,
+        locations: backup.locationIds.length,
+        items: backup.itemIds.length,
+        storyBeats: backup.storyBeatIds.length,
+        lorebookEntries: backup.lorebookEntryIds.length,
+      },
+    });
   }
 
   /**
@@ -305,12 +487,34 @@ class UIStore {
    * Used when editing the last user message to retry with new content.
    */
   updateRetryBackupContent(newContent: string) {
-    if (this.retryBackup) {
-      this.retryBackup = {
-        ...this.retryBackup,
+    const backup = this.retryBackup;
+    if (backup && this.currentRetryStoryId) {
+      const updatedBackup: RetryBackup = {
+        ...backup,
         userActionContent: newContent,
         rawInput: newContent,
       };
+      this.retryBackups.set(this.currentRetryStoryId, updatedBackup);
+
+      // Also persist the updated content to the database
+      const storyId = this.currentRetryStoryId;
+      this.queueRetryStateWrite(
+        () => database.saveRetryState(storyId, {
+          timestamp: backup.timestamp,
+          entryCountBeforeAction: backup.entryCountBeforeAction,
+          userActionContent: newContent,
+          rawInput: newContent,
+          actionType: backup.actionType,
+          wasRawActionChoice: backup.wasRawActionChoice,
+          characterIds: backup.characterIds,
+          locationIds: backup.locationIds,
+          itemIds: backup.itemIds,
+          storyBeatIds: backup.storyBeatIds,
+          lorebookEntryIds: backup.lorebookEntryIds,
+        }),
+        'update'
+      );
+
       console.log('[UI] Retry backup content updated', {
         newContent: newContent.substring(0, 50),
       });
diff --git a/src/lib/types/index.ts b/src/lib/types/index.ts
index 5477173..b25a7b1 100644
--- a/src/lib/types/index.ts
+++ b/src/lib/types/index.ts
@@ -15,6 +15,28 @@ export interface Story {
   updatedAt: number;
   settings: StorySettings | null;
   memoryConfig: MemoryConfig | null;
+  retryState: PersistentRetryState | null;
+}
+
+// Persistent retry state - lightweight version saved to database
+export type ActionInputType = 'do' | 'say' | 'think' | 'story' | 'free';
+
+export interface PersistentRetryState {
+  timestamp: number;
+  // The next entry position before user action was added (max position + 1)
+  // On retry, delete entries from this position onward
+  entryCountBeforeAction: number;
+  // The user's input data
+  userActionContent: string;
+  rawInput: string;
+  actionType: ActionInputType;
+  wasRawActionChoice: boolean;
+  // Entity IDs that existed before the action - on restore, delete any not in these lists
+  characterIds: string[];
+  locationIds: string[];
+  itemIds: string[];
+  storyBeatIds: string[];
+  lorebookEntryIds: string[];
 }
 
 export interface MemoryConfig {