Fix history injection for prewarm extensions

- Use persistent event listeners instead of once() to inject into ALL generations
- Don't clear context map on GENERATION_ENDED so prewarm gets the same context
- Remove unused onGenerationEndedCleanup function

index.js

@@ -151,7 +151,8 @@ import {
     onMessageSwiped,
     updatePersonaAvatar,
     clearExtensionPrompts,
-    onGenerationEnded
+    onGenerationEnded,
+    initHistoryInjection
 } from './src/systems/integration/sillytavern.js';
 
 // Old state variable declarations removed - now imported from core modules
@@ -1018,6 +1019,15 @@ jQuery(async () => {
             // Non-critical - continue anyway
         }
 
+        // Initialize history injection event listeners
+        // This must be done before event registration so listeners are ready
+        try {
+            initHistoryInjection();
+        } catch (error) {
+            console.error('[RPG Companion] History injection init failed:', error);
+            // Non-critical - continue without it
+        }
+
         // Register all event listeners
         try {
             registerAllEvents({

src/systems/generation/injector.js

@@ -301,10 +301,14 @@ function onGenerateAfterCombinePrompts(eventData) {
         return;
     }
 
-    eventData.prompt = injectContextIntoTextPrompt(eventData.prompt);
+    // Only inject if we have pending context
+    if (pendingContextMap.size === 0) {
+        return;
+    }
 
-    // Clear the pending context after injection
+    eventData.prompt = injectContextIntoTextPrompt(eventData.prompt);
-    pendingContextMap = new Map();
+    // DON'T clear pendingContextMap here - let it persist for other generations
+    // (e.g., prewarm extensions). It will be cleared on GENERATION_ENDED.
 }
 
 /**
@@ -322,10 +326,14 @@ function onChatCompletionPromptReady(eventData) {
         return;
     }
 
-    eventData.chat = injectContextIntoChatPrompt(eventData.chat);
+    // Only inject if we have pending context
+    if (pendingContextMap.size === 0) {
+        return;
+    }
 
-    // Clear the pending context after injection
+    eventData.chat = injectContextIntoChatPrompt(eventData.chat);
-    pendingContextMap = new Map();
+    // DON'T clear pendingContextMap here - let it persist for other generations
+    // (e.g., prewarm extensions). It will be cleared on GENERATION_ENDED.
 }
 
 /**
@@ -656,22 +664,20 @@ Ensure these details naturally reflect and influence the narrative. Character be
 
     // Prepare historical context for injection into prompts
     // This builds the context map but does NOT modify original chat messages
+    // The persistent event listeners will inject it into all prompts until cleared
     prepareHistoricalContextInjection();
-
-    // Register one-time listeners to inject context into the actual prompt
-    // These modify only the prompt sent to the API, not the stored chat data
-    if (pendingContextMap.size > 0) {
-        eventSource.once(event_types.GENERATE_AFTER_COMBINE_PROMPTS, onGenerateAfterCombinePrompts);
-        eventSource.once(event_types.CHAT_COMPLETION_PROMPT_READY, onChatCompletionPromptReady);
-    }
 }
 
 /**
- * Called when generation ends to clean up any pending context.
+ * Initialize the history injection event listeners.
- * This should be called from the GENERATION_ENDED event handler.
+ * These are persistent listeners that inject context into ALL generations
+ * while pendingContextMap has data. Should be called once at extension init.
  */
-export function onGenerationEndedCleanup() {
+export function initHistoryInjectionListeners() {
-    // Clear any pending context that wasn't used (e.g., if generation was cancelled)
+    // Register persistent listeners for prompt injection
-    pendingContextMap = new Map();
+    // These check pendingContextMap and only inject if there's data
+    eventSource.on(event_types.GENERATE_AFTER_COMBINE_PROMPTS, onGenerateAfterCombinePrompts);
+    eventSource.on(event_types.CHAT_COMPLETION_PROMPT_READY, onChatCompletionPromptReady);
+    console.log('[RPG Companion] History injection listeners initialized');
 }
 

src/systems/integration/sillytavern.js

@@ -30,7 +30,7 @@ import { parseResponse, parseUserStats } from '../generation/parser.js';
 import { parseAndStoreSpotifyUrl, convertToEmbedUrl } from '../features/musicPlayer.js';
 import { updateRPGData } from '../generation/apiClient.js';
 import { removeLocks } from '../generation/lockManager.js';
-import { onGenerationStarted, onGenerationEndedCleanup } from '../generation/injector.js';
+import { onGenerationStarted, initHistoryInjectionListeners } from '../generation/injector.js';
 
 // Rendering
 import { renderUserStats } from '../rendering/userStats.js';
@@ -453,9 +453,6 @@ export function clearExtensionPrompts() {
 export async function onGenerationEnded() {
     // console.log('[RPG Companion] 🏁 onGenerationEnded called');
 
-    // Restore original message content that was modified for historical context injection
-    onGenerationEndedCleanup();
-
     // Note: isGenerating flag is cleared in onMessageReceived after parsing (together mode)
     // or in apiClient.js after separate generation completes (separate mode)
 
@@ -463,3 +460,11 @@ export async function onGenerationEnded() {
     // Re-apply checkpoint if one exists
     await restoreCheckpointOnLoad();
 }
+
+/**
+ * Initialize history injection event listeners.
+ * Should be called once during extension initialization.
+ */
+export function initHistoryInjection() {
+    initHistoryInjectionListeners();
+}