From b9a15722d63ed5117cc0d99c585e4c931639fe11 Mon Sep 17 00:00:00 2001 From: tomt610 Date: Sat, 10 Jan 2026 19:33:26 +0000 Subject: [PATCH] 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(); +}