81f65203 · permissionBRICK/SillyTavern

     AND_ALL: 3,
 };
+/**
+ * @enum {number} Possible states of the WI evaluation
+ */
+const scan_state = {
+    /**
+     * The scan will be stopped.
+     */
+    NONE: 0,
+    /**
+     * Initial state.
+     */
+    INITIAL: 1,
+    /**
+     * The scan is triggered by a recursion step.
+     */
+    RECURSION: 2,
+    /**
+     * The scan is triggered by a min activations depth skew.
+     */
+    MIN_ACTIVATIONS: 2,
+};
 const WI_ENTRY_EDIT_TEMPLATE = $('#entry_edit_template .world_entry');
 let world_info = {};

     #recurseBuffer = [];
     /**
+     * @type {string[]} Array of strings added by prompt injections that are valid for the current scan
+     */
+    #injectBuffer = [];
+    /**
      * @type {number} The skew of the global scan depth. Used in "min activations"
      */
     #skew = 0;

     /**
      * Gets all messages up to the given depth + recursion buffer.
      * @param {WIScanEntry} entry The entry that triggered the scan
+     * @param {number} scanState The state of the scan
      * @returns {string} A slice of buffer until the given depth (inclusive)
      */
-    get(entry) {
+    get(entry, scanState) {
         let depth = entry.scanDepth ?? this.getDepth();
         if (depth <= this.#startDepth) {
             return '';

         let result = this.#depthBuffer.slice(this.#startDepth, depth).join('\n');
-        if (this.#recurseBuffer.length > 0) {
+        if (this.#injectBuffer.length > 0) {
+            result += '\n' + this.#injectBuffer.join('\n');
+        }
+        // Min activations should not include the recursion buffer
+        if (this.#recurseBuffer.length > 0 && scanState !== scan_state.MIN_ACTIVATIONS) {
             result += '\n' + this.#recurseBuffer.join('\n');
+        }

+    }
     /**
+     * Adds an injection to the buffer.
+     * @param {string} message The injection to add
+     */
+    addInject(message) {
+        this.#injectBuffer.push(message);
+    }
+    /**
      * Increments skew and sets startDepth to previous depth.
      */
     advanceScanPosition() {

     /**
      * Gets the match score for the given entry.
      * @param {WIScanEntry} entry Entry to check
+     * @param {number} scanState The state of the scan
      * @returns {number} The number of key activations for the given entry
      */
-    getScore(entry) {
+    getScore(entry, scanState) {
-        const bufferState = this.get(entry);
+        const bufferState = this.get(entry, scanState);
         let numberOfPrimaryKeys = 0;
         let numberOfSecondaryKeys = 0;
         let primaryScore = 0;

@@ -3503,12 +3545,12 @@ async function checkWorldInfo(chat, maxContext, isDryRun) {
3503	if (context.extensionPrompts[key]?.scan) {	3545	if (context.extensionPrompts[key]?.scan) {
3504	const prompt = getExtensionPromptByName(key);	3546	const prompt = getExtensionPromptByName(key);
3505	if (prompt) {	3547	if (prompt) {
3506	buffer.addRecurse(prompt);	3548	buffer.addInject(prompt);
3507	}	3549	}
3508	}	3550	}
3509	}	3551	}
3510		3552
3511	let needsToScan = true;	3553	let scanState = scan_state.INITIAL;
3512	let token_budget_overflowed = false;	3554	let token_budget_overflowed = false;
3513	let count = 0;	3555	let count = 0;
3514	let allActivatedEntries = new Set();	3556	let allActivatedEntries = new Set();

         return { worldInfoBefore: '', worldInfoAfter: '', WIDepthEntries: [], EMEntries: [], allActivatedEntries: new Set() };
+    }
-    while (needsToScan) {
+    while (scanState) {
-        // Track how many times the loop has run
+        // Track how many times the loop has run. May be useful for debugging.
+        // eslint-disable-next-line no-unused-vars
         count++;
         let activatedNow = new Set();

                 continue;
+            }
-            if (allActivatedEntries.has(entry) || entry.disable == true || (count > 1 && world_info_recursive && entry.excludeRecursion) || (count == 1 && entry.delayUntilRecursion)) {
+            if (allActivatedEntries.has(entry) || entry.disable == true) {
+                continue;
+            }
+            // Only use checks for recursion flags if the scan step was activated by recursion
+            if (scanState !== scan_state.RECURSION && entry.delayUntilRecursion) {
+                console.debug(`WI entry ${entry.uid} suppressed by delay until recursion`, entry);
+                continue;
+            }
+            if (scanState === scan_state.RECURSION && world_info_recursive && entry.excludeRecursion) {
+                console.debug(`WI entry ${entry.uid} suppressed by exclude recursion`, entry);
                 continue;
+            }

                 primary: for (let key of entry.key) {
                     const substituted = substituteParams(key);
-                    const textToScan = buffer.get(entry);
+                    const textToScan = buffer.get(entry, scanState);
                     if (substituted && buffer.matchKeys(textToScan, substituted.trim(), entry)) {
                         console.debug(`WI UID ${entry.uid} found by primary match: ${substituted}.`);

+            }
+        }
-        needsToScan = world_info_recursive && activatedNow.size > 0;
+        scanState = world_info_recursive && activatedNow.size > 0 ? scan_state.RECURSION : scan_state.NONE;
         const newEntries = [...activatedNow]
             .sort((a, b) => sortedEntries.indexOf(a) - sortedEntries.indexOf(b));
         let newContent = '';
         const textToScanTokens = await getTokenCountAsync(allActivatedText);
         const probabilityChecksBefore = failedProbabilityChecks.size;
-        filterByInclusionGroups(newEntries, allActivatedEntries, buffer);
+        filterByInclusionGroups(newEntries, allActivatedEntries, buffer, scanState);
         console.debug('-- PROBABILITY CHECKS BEGIN --');
         for (const entry of newEntries) {

@@ -3697,7 +3751,7 @@ async function checkWorldInfo(chat, maxContext, isDryRun) {
3697	console.log('Alerting');	3751	console.log('Alerting');
3698	toastr.warning(`World info budget reached after ${allActivatedEntries.size} entries.`, 'World Info');	3752	toastr.warning(`World info budget reached after ${allActivatedEntries.size} entries.`, 'World Info');
3699	}	3753	}
3700	needsToScan = false;	3754	scanState = scan_state.NONE;
3701	token_budget_overflowed = true;	3755	token_budget_overflowed = true;
3702	break;	3756	break;
3703	}	3757	}

         if ((probabilityChecksAfter - probabilityChecksBefore) === activatedNow.size) {
             console.debug('WI probability checks failed for all activated entries, stopping');
-            needsToScan = false;
+            scanState = scan_state.NONE;
+        }
         if (newEntries.length === 0) {
             console.debug('No new entries activated, stopping');
-            needsToScan = false;
+            scanState = scan_state.NONE;
+        }
-        if (needsToScan) {
+        if (scanState) {
             const text = newEntries
                 .filter(x => !failedProbabilityChecks.has(x))
                 .filter(x => !x.preventRecursion)

+        }
         // world_info_min_activations
-        if (!needsToScan && !token_budget_overflowed) {
+        if (!scanState && !token_budget_overflowed) {
             if (world_info_min_activations > 0 && (allActivatedEntries.size < world_info_min_activations)) {
                 let over_max = (
                     world_info_min_activations_depth_max > 0 &&

@@ -3736,7 +3790,7 @@ async function checkWorldInfo(chat, maxContext, isDryRun) {
3736	) \|\| (buffer.getDepth() > chat.length);	3790	) \|\| (buffer.getDepth() > chat.length);
3737		3791
3738	if (!over_max) {	3792	if (!over_max) {
3739	needsToScan = true; // loop	3793	scanState = scan_state.MIN_ACTIVATIONS; // loop
3740	buffer.advanceScanPosition();	3794	buffer.advanceScanPosition();
3741	}	3795	}
3742	}	3796	}

  * @param {Record<string, WIScanEntry[]>} groups The groups to filter
  * @param {WorldInfoBuffer} buffer The buffer to use for scoring
  * @param {(entry: WIScanEntry) => void} removeEntry The function to remove an entry
+ * @param {number} scanState The current scan state
  */
-function filterGroupsByScoring(groups, buffer, removeEntry) {
+function filterGroupsByScoring(groups, buffer, removeEntry, scanState) {
     for (const [key, group] of Object.entries(groups)) {
         // Group scoring is disabled both globally and for the group entries
         if (!world_info_use_group_scoring && !group.some(x => x.useGroupScoring)) {

             continue;
+        }
-        const scores = group.map(entry => buffer.getScore(entry));
+        const scores = group.map(entry => buffer.getScore(entry, scanState));
         const maxScore = Math.max(...scores);
         console.debug(`Group '${key}' max score: ${maxScore}`);
         //console.table(group.map((entry, i) => ({ uid: entry.uid, key: JSON.stringify(entry.key), score: scores[i] })));

  * @param {object[]} newEntries Entries activated on current recursion level
  * @param {Set<object>} allActivatedEntries Set of all activated entries
  * @param {WorldInfoBuffer} buffer The buffer to use for scanning
+ * @param {number} scanState The current scan state
  */
-function filterByInclusionGroups(newEntries, allActivatedEntries, buffer) {
+function filterByInclusionGroups(newEntries, allActivatedEntries, buffer, scanState) {
     console.debug('-- INCLUSION GROUP CHECKS BEGIN --');
     const grouped = newEntries.filter(x => x.group).reduce((acc, item) => {
         item.group.split(/,\s*/).filter(x => x).forEach(group => {

@@ -3891,7 +3947,7 @@ function filterByInclusionGroups(newEntries, allActivatedEntries, buffer) {
3891	}	3947	}
3892	}	3948	}
3893		3949
3894	filterGroupsByScoring(grouped, buffer, removeEntry);	3950	filterGroupsByScoring(grouped, buffer, removeEntry, scanState);
3895		3951
3896	for (const [key, group] of Object.entries(grouped)) {	3952	for (const [key, group] of Object.entries(grouped)) {
3897	console.debug(`Checking inclusion group '${key}' with ${group.length} entries`, group);	3953	console.debug(`Checking inclusion group '${key}' with ${group.length} entries`, group);

Merge pull request #2468 from SillyTavern/wi-scan-state Fix min activations for non-recursable entries