From db97f012b04f5e750f9a2e5770e8ed6118cad27c Mon Sep 17 00:00:00 2001 From: tomt610 Date: Sat, 10 Jan 2026 19:10:33 +0000 Subject: [PATCH 1/2] Refactor history injection to modify prompts instead of chat messages This prevents any risk of injected context being accidentally saved to the chat. Instead of modifying chat[].mes directly, we now: 1. Build a context map during GENERATION_STARTED 2. Inject into the prompt string (GENERATE_AFTER_COMBINE_PROMPTS) for text completion 3. Inject into the message array (CHAT_COMPLETION_PROMPT_READY) for chat completion The original chat messages are never modified. --- src/systems/generation/injector.js | 185 +++++++++++++++++++++-------- 1 file changed, 138 insertions(+), 47 deletions(-) diff --git a/src/systems/generation/injector.js b/src/systems/generation/injector.js index 0e87a70..2867f07 100644 --- a/src/systems/generation/injector.js +++ b/src/systems/generation/injector.js @@ -37,9 +37,8 @@ let currentSuppressionState = false; // Track last chat length we committed at to prevent duplicate commits from streaming let lastCommittedChatLength = -1; -// Store original message content for restoration after generation -// Map of message index -> original mes content -let originalMessageContent = new Map(); +// Store context map for prompt injection (used by event handlers) +let pendingContextMap = new Map(); /** * Builds a map of historical context data from ST chat messages with rpg_companion_swipes data. @@ -162,81 +161,171 @@ function buildHistoricalContextMap() { } /** - * Injects historical context into chat messages by modifying them in-place. - * Stores original content for restoration after generation. - * This approach works for ALL API types (text completion and chat completion). + * Prepares historical context for injection into prompts. + * This builds the context map and stores it for use by prompt event handlers. + * Does NOT modify the original chat messages. */ -function injectHistoricalContextIntoChat() { +function prepareHistoricalContextInjection() { const historyPersistence = extensionSettings.historyPersistence; if (!historyPersistence || !historyPersistence.enabled) { - // console.log('[RPG Companion] History persistence not enabled, skipping injection'); + pendingContextMap = new Map(); return; } if (currentSuppressionState || !extensionSettings.enabled) { - // console.log('[RPG Companion] Skipping history injection: suppressed or disabled'); + pendingContextMap = new Map(); return; } const context = getContext(); const chat = context.chat; if (!chat || chat.length < 2) { - // console.log('[RPG Companion] Chat too short, skipping history injection'); + pendingContextMap = new Map(); return; } - // Build the context map - const contextMap = buildHistoricalContextMap(); - if (contextMap.size === 0) { - // console.log('[RPG Companion] No historical context to inject'); - return; + // Build and store the context map for use by prompt handlers + pendingContextMap = buildHistoricalContextMap(); +} + +/** + * Injects historical context into a text completion prompt string. + * Searches for message content in the prompt and appends context after matches. + * + * @param {string} prompt - The text completion prompt + * @returns {string} - The modified prompt with injected context + */ +function injectContextIntoTextPrompt(prompt) { + if (pendingContextMap.size === 0) { + return prompt; } - // console.log(`[RPG Companion] Injecting historical context into ${contextMap.size} messages`); - - // Clear any previous stored content - originalMessageContent.clear(); - + const context = getContext(); + const chat = context.chat; + let modifiedPrompt = prompt; let injectedCount = 0; - for (const [msgIdx, ctxContent] of contextMap) { + + // Process each message that needs context injection + for (const [msgIdx, ctxContent] of pendingContextMap) { const message = chat[msgIdx]; if (!message || typeof message.mes !== 'string') { continue; } - // Store original content for restoration - originalMessageContent.set(msgIdx, message.mes); + // Find the message content in the prompt + // Use a portion of the message to find it (last 100 chars should be unique enough) + const searchContent = message.mes.length > 100 + ? message.mes.slice(-100) + : message.mes; + + const searchIndex = modifiedPrompt.lastIndexOf(searchContent); + if (searchIndex === -1) { + // Message not found in prompt (might be truncated) + continue; + } - // Modify the message in-place - message.mes = message.mes + ctxContent; + // Find the end of this message content in the prompt + const insertPosition = searchIndex + searchContent.length; + + // Insert the context after the message + modifiedPrompt = modifiedPrompt.slice(0, insertPosition) + ctxContent + modifiedPrompt.slice(insertPosition); injectedCount++; - // console.log(`[RPG Companion] Injected context into message ${msgIdx}`); } - // console.log(`[RPG Companion] Successfully injected historical context into ${injectedCount} messages`); + if (injectedCount > 0) { + console.log(`[RPG Companion] Injected historical context into ${injectedCount} positions in text prompt`); + } + + return modifiedPrompt; } /** - * Restores original message content after generation completes. - * This ensures the injected context doesn't persist in the actual chat data. + * Injects historical context into a chat completion message array. + * Modifies the content of messages in the array directly. + * + * @param {Array} chatMessages - The chat completion message array + * @returns {Array} - The modified message array with injected context */ -function restoreOriginalMessageContent() { - if (originalMessageContent.size === 0) { - return; +function injectContextIntoChatPrompt(chatMessages) { + if (pendingContextMap.size === 0 || !Array.isArray(chatMessages)) { + return chatMessages; } const context = getContext(); const chat = context.chat; + let injectedCount = 0; - // console.log(`[RPG Companion] Restoring ${originalMessageContent.size} messages to original content`); + // Process each message that needs context injection + for (const [msgIdx, ctxContent] of pendingContextMap) { + const originalMessage = chat[msgIdx]; + if (!originalMessage || typeof originalMessage.mes !== 'string') { + continue; + } - for (const [msgIdx, originalContent] of originalMessageContent) { - if (chat[msgIdx]) { - chat[msgIdx].mes = originalContent; + // Find this message in the chat completion array by matching content + // Use a portion of the message to find it + const searchContent = originalMessage.mes.length > 100 + ? originalMessage.mes.slice(-100) + : originalMessage.mes; + + for (const promptMsg of chatMessages) { + if (promptMsg.content && typeof promptMsg.content === 'string' && + promptMsg.content.includes(searchContent)) { + // Found the message - append context + promptMsg.content = promptMsg.content + ctxContent; + injectedCount++; + break; + } } } - originalMessageContent.clear(); + if (injectedCount > 0) { + console.log(`[RPG Companion] Injected historical context into ${injectedCount} messages in chat prompt`); + } + + return chatMessages; +} + +/** + * Event handler for GENERATE_AFTER_COMBINE_PROMPTS (text completion). + * Injects historical context into the prompt string. + * + * @param {Object} eventData - Event data with prompt property + */ +function onGenerateAfterCombinePrompts(eventData) { + if (!eventData || typeof eventData.prompt !== 'string') { + return; + } + + if (eventData.dryRun) { + return; + } + + eventData.prompt = injectContextIntoTextPrompt(eventData.prompt); + + // Clear the pending context after injection + pendingContextMap = new Map(); +} + +/** + * Event handler for CHAT_COMPLETION_PROMPT_READY. + * Injects historical context into the chat message array. + * + * @param {Object} eventData - Event data with chat property + */ +function onChatCompletionPromptReady(eventData) { + if (!eventData || !Array.isArray(eventData.chat)) { + return; + } + + if (eventData.dryRun) { + return; + } + + eventData.chat = injectContextIntoChatPrompt(eventData.chat); + + // Clear the pending context after injection + pendingContextMap = new Map(); } /** @@ -565,22 +654,24 @@ Ensure these details naturally reflect and influence the narrative. Character be // Set suppression state for the historical context injection currentSuppressionState = shouldSuppress; - // Inject historical context directly into chat messages - // This temporarily modifies messages and will be restored after generation - injectHistoricalContextIntoChat(); + // Prepare historical context for injection into prompts + // This builds the context map but does NOT modify original chat messages + prepareHistoricalContextInjection(); - // Register a one-time listener to restore messages after prompt is built - // Using .once() so it auto-removes after firing - eventSource.once(event_types.GENERATE_AFTER_COMBINE_PROMPTS, () => { - restoreOriginalMessageContent(); - }); + // 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 restore original message content. + * Called when generation ends to clean up any pending context. * This should be called from the GENERATION_ENDED event handler. */ export function onGenerationEndedCleanup() { - restoreOriginalMessageContent(); + // Clear any pending context that wasn't used (e.g., if generation was cancelled) + pendingContextMap = new Map(); } From b9a15722d63ed5117cc0d99c585e4c931639fe11 Mon Sep 17 00:00:00 2001 From: tomt610 Date: Sat, 10 Jan 2026 19:33:26 +0000 Subject: [PATCH 2/2] 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 | 12 +++++++- src/systems/generation/injector.js | 42 +++++++++++++++----------- src/systems/integration/sillytavern.js | 13 +++++--- 3 files changed, 44 insertions(+), 23 deletions(-) diff --git a/index.js b/index.js index 26170ce..85c6ced 100644 --- a/index.js +++ b/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({ diff --git a/src/systems/generation/injector.js b/src/systems/generation/injector.js index 2867f07..0986332 100644 --- a/src/systems/generation/injector.js +++ b/src/systems/generation/injector.js @@ -301,10 +301,14 @@ function onGenerateAfterCombinePrompts(eventData) { return; } + // Only inject if we have pending context + if (pendingContextMap.size === 0) { + return; + } + eventData.prompt = injectContextIntoTextPrompt(eventData.prompt); - - // Clear the pending context after injection - 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; } + // Only inject if we have pending context + if (pendingContextMap.size === 0) { + return; + } + eventData.chat = injectContextIntoChatPrompt(eventData.chat); - - // Clear the pending context after injection - 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. - * This should be called from the GENERATION_ENDED event handler. + * Initialize the history injection event listeners. + * These are persistent listeners that inject context into ALL generations + * while pendingContextMap has data. Should be called once at extension init. */ -export function onGenerationEndedCleanup() { - // Clear any pending context that wasn't used (e.g., if generation was cancelled) - pendingContextMap = new Map(); +export function initHistoryInjectionListeners() { + // Register persistent listeners for prompt injection + // 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'); } diff --git a/src/systems/integration/sillytavern.js b/src/systems/integration/sillytavern.js index 0064fff..880594f 100644 --- a/src/systems/integration/sillytavern.js +++ b/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(); +}