Refactor historical context injection logic to support dynamic message indexing based on injection position

src/systems/generation/injector.js

@@ -44,8 +44,9 @@ let originalMessageContent = new Map();
 /**
  * Builds a map of historical context data from ST chat messages with rpg_companion_swipes data.
  * Returns a map keyed by message index with formatted context strings.
+ * The index stored depends on the injection position setting.
  *
- * @returns {Map<number, {context: string, isUserMessage: boolean}>} Map of message index to context data
+ * @returns {Map<number, string>} Map of target message index to formatted context string
  */
 function buildHistoricalContextMap() {
     const historyPersistence = extensionSettings.historyPersistence;
@@ -61,19 +62,41 @@ function buildHistoricalContextMap() {
 
     const trackerConfig = extensionSettings.trackerConfig;
     const userName = context.name1;
+    const position = historyPersistence.injectionPosition || 'assistant_message_end';
     const contextMap = new Map();
 
     // Determine how many messages to include (0 = all available)
     const messageCount = historyPersistence.messageCount || 0;
     const maxMessages = messageCount === 0 ? chat.length : Math.min(messageCount, chat.length);
 
-    // Start from the second-to-last message (skip the most recent one as it gets current context)
+    // Find the last assistant message - this is the one that gets current context via setExtensionPrompt
-    // and work backwards
+    // We should NOT add historical context to it
-    let processedCount = 0;
+    let lastAssistantIndex = -1;
+    for (let i = chat.length - 1; i >= 0; i--) {
+        if (!chat[i].is_user && !chat[i].is_system) {
+            lastAssistantIndex = i;
+            break;
+        }
+    }
 
-    for (let i = chat.length - 2; i >= 0 && (messageCount === 0 || processedCount < maxMessages); i--) {
+    // Iterate through messages to find those with tracker data
+    // Start from before the last assistant message
+    let processedCount = 0;
+    const startIndex = lastAssistantIndex > 0 ? lastAssistantIndex - 1 : chat.length - 2;
+
+    for (let i = startIndex; i >= 0 && (messageCount === 0 || processedCount < maxMessages); i--) {
         const message = chat[i];
 
+        // Skip system messages
+        if (message.is_system) {
+            continue;
+        }
+
+        // Only assistant messages have rpg_companion_swipes data
+        if (message.is_user) {
+            continue;
+        }
+
         // Get the rpg_companion_swipes data for current swipe
         const swipeData = message.extra?.rpg_companion_swipes;
         if (!swipeData) {
@@ -96,11 +119,32 @@ function buildHistoricalContextMap() {
         const preamble = historyPersistence.contextPreamble || '[Context at this point:]';
         const wrappedContext = `\n${preamble}\n${formattedContext}`;
 
-        // Store with message index and whether it's a user message
+        // Determine which message index to store based on injection position
-        contextMap.set(i, {
+        let targetIndex = i; // Default: the assistant message itself
-            context: wrappedContext,
+
-            isUserMessage: message.is_user
+        if (position === 'user_message_end') {
-        });
+            // Find the next user message after this assistant message
+            for (let j = i + 1; j < chat.length; j++) {
+                if (chat[j].is_user && !chat[j].is_system) {
+                    targetIndex = j;
+                    break;
+                }
+            }
+            // If no user message found after, skip this one
+            if (targetIndex === i) {
+                continue;
+            }
+        }
+        // For assistant_message_end, extra_user_message, extra_assistant_message:
+        // We inject into the assistant message itself (for now - extra messages handled differently)
+
+        // Store the context keyed by target index
+        // If multiple assistant messages map to the same user message, append
+        if (contextMap.has(targetIndex)) {
+            contextMap.set(targetIndex, contextMap.get(targetIndex) + wrappedContext);
+        } else {
+            contextMap.set(targetIndex, wrappedContext);
+        }
 
         processedCount++;
     }
@@ -141,41 +185,23 @@ function injectHistoricalContextIntoChat() {
 
     console.log(`[RPG Companion] Injecting historical context into ${contextMap.size} messages`);
 
-    const position = historyPersistence.injectionPosition || 'assistant_message_end';
-
     // Clear any previous stored content
     originalMessageContent.clear();
 
     let injectedCount = 0;
-    for (const [msgIdx, data] of contextMap) {
+    for (const [msgIdx, ctxContent] of contextMap) {
         const message = chat[msgIdx];
         if (!message || typeof message.mes !== 'string') {
             continue;
         }
 
-        const { context: ctxContent, isUserMessage } = data;
+        // Store original content for restoration
+        originalMessageContent.set(msgIdx, message.mes);
 
-        // Determine if we should inject based on position and message type
+        // Modify the message in-place
-        let shouldInject = false;
+        message.mes = message.mes + ctxContent;
-
+        injectedCount++;
-        if (position === 'user_message_end' && isUserMessage) {
+        console.log(`[RPG Companion] Injected context into message ${msgIdx}`);
-            shouldInject = true;
-        } else if (position === 'assistant_message_end' && !isUserMessage) {
-            shouldInject = true;
-        } else if (position === 'extra_user_message' || position === 'extra_assistant_message') {
-            // For these positions, inject regardless of message type
-            shouldInject = true;
-        }
-
-        if (shouldInject) {
-            // Store original content for restoration
-            originalMessageContent.set(msgIdx, message.mes);
-
-            // Modify the message in-place
-            message.mes = message.mes + ctxContent;
-            injectedCount++;
-            console.log(`[RPG Companion] Injected context into message ${msgIdx}`);
-        }
     }
 
     console.log(`[RPG Companion] Successfully injected historical context into ${injectedCount} messages`);
@@ -533,6 +559,14 @@ Ensure these details naturally reflect and influence the narrative. Character be
     // Inject historical context directly into chat messages
     // This temporarily modifies messages and will be restored after generation
     injectHistoricalContextIntoChat();
+
+    // Register a one-time listener to restore messages after prompt is built
+    // This ensures messages are restored even if generation is cancelled or for dry runs
+    const restoreAfterPrompt = () => {
+        restoreOriginalMessageContent();
+        eventSource.off(event_types.GENERATE_AFTER_COMBINE_PROMPTS, restoreAfterPrompt);
+    };
+    eventSource.once(event_types.GENERATE_AFTER_COMBINE_PROMPTS, restoreAfterPrompt);
 }
 
 /**

src/systems/generation/promptBuilder.js

@@ -1123,26 +1123,91 @@ export async function generateSeparateUpdatePrompt() {
     // Add chat history as separate user/assistant messages with per-message historical context
     const recentMessages = chat.slice(-depth);
     const startIndex = chat.length - depth;
+    const position = historyPersistence?.injectionPosition || 'assistant_message_end';
 
-    for (let i = 0; i < recentMessages.length; i++) {
+    // Build a map of which messages should get context based on position setting
-        const message = recentMessages[i];
+    // Key: message index in recentMessages, Value: context string
-        const chatIndex = startIndex + i;
+    const contextInjectionMap = new Map();
-        let content = message.mes;
+
+    if (historyPersistence?.enabled) {
+        // Find the last assistant message index (in recentMessages)
+        let lastAssistantIdx = -1;
+        for (let i = recentMessages.length - 1; i >= 0; i--) {
+            if (!recentMessages[i].is_user && !recentMessages[i].is_system) {
+                lastAssistantIdx = i;
+                break;
+            }
+        }
+
+        // Iterate through assistant messages to find tracker data
+        for (let i = 0; i < recentMessages.length; i++) {
+            const message = recentMessages[i];
+
+            // Skip user and system messages - only assistant messages have tracker data
+            if (message.is_user || message.is_system) {
+                continue;
+            }
+
+            // Skip the last assistant message - it gets current context elsewhere
+            if (i === lastAssistantIdx) {
+                continue;
+            }
 
-        // Append historical tracker context to this message if enabled and available
-        if (historyPersistence?.enabled && chatIndex < chat.length - 1) {
             const swipeData = message.extra?.rpg_companion_swipes;
-            if (swipeData) {
+            if (!swipeData) {
-                const currentSwipeId = message.swipe_id || 0;
+                continue;
-                const trackerData = swipeData[currentSwipeId];
+            }
-                if (trackerData) {
+
-                    const formattedContext = formatHistoricalTrackerData(trackerData, trackerConfig, userName);
+            const currentSwipeId = message.swipe_id || 0;
-                    if (formattedContext) {
+            const trackerData = swipeData[currentSwipeId];
-                        const preamble = historyPersistence.contextPreamble || '[Context at this point:]';
+            if (!trackerData) {
-                        content += `\n${preamble}\n${formattedContext}`;
+                continue;
+            }
+
+            const formattedContext = formatHistoricalTrackerData(trackerData, trackerConfig, userName);
+            if (!formattedContext) {
+                continue;
+            }
+
+            const preamble = historyPersistence.contextPreamble || '[Context at this point:]';
+            const wrappedContext = `\n${preamble}\n${formattedContext}`;
+
+            // Determine target message based on position
+            let targetIdx = i;
+
+            if (position === 'user_message_end') {
+                // Find next user message after this assistant message
+                for (let j = i + 1; j < recentMessages.length; j++) {
+                    if (recentMessages[j].is_user && !recentMessages[j].is_system) {
+                        targetIdx = j;
+                        break;
                     }
                 }
+                // If no user message found, skip
+                if (targetIdx === i) {
+                    continue;
+                }
             }
+            // For assistant_message_end, extra_user_message, extra_assistant_message:
+            // Inject into the assistant message itself
+
+            // Append to existing or create new entry
+            if (contextInjectionMap.has(targetIdx)) {
+                contextInjectionMap.set(targetIdx, contextInjectionMap.get(targetIdx) + wrappedContext);
+            } else {
+                contextInjectionMap.set(targetIdx, wrappedContext);
+            }
+        }
+    }
+
+    // Now build the messages array with injected context
+    for (let i = 0; i < recentMessages.length; i++) {
+        const message = recentMessages[i];
+        let content = message.mes;
+
+        // Add historical context if this message is a target
+        if (contextInjectionMap.has(i)) {
+            content += contextInjectionMap.get(i);
         }
 
         messages.push({