From 9cc1989bec4a6a3631897c076cbf5ab2bbb23360 Mon Sep 17 00:00:00 2001 From: Darkatse <39111369+Darkatse@users.noreply.github.com> Date: Sat, 4 Apr 2026 22:40:36 -0400 Subject: [PATCH] fix (generate): Strictly enforce `Sillytavernget.Context().generate` as a single-flight operation. --- src/script.js | 37 ++++- src/scripts/st-context.js | 9 +- src/scripts/util/generation-idle-gate.js | 52 +++++++ src/scripts/util/safe-generate.js | 31 +++++ tests/generation-lifecycle-contract.test.mjs | 19 +++ tests/generation-single-flight.test.mjs | 136 +++++++++++++++++++ tests/st-context-contract.test.mjs | 37 +++++ 7 files changed, 319 insertions(+), 2 deletions(-) create mode 100644 src/scripts/util/generation-idle-gate.js create mode 100644 src/scripts/util/safe-generate.js create mode 100644 tests/generation-lifecycle-contract.test.mjs create mode 100644 tests/generation-single-flight.test.mjs create mode 100644 tests/st-context-contract.test.mjs diff --git a/src/script.js b/src/script.js index de9b54df..09e94292 100644 --- a/src/script.js +++ b/src/script.js @@ -326,6 +326,7 @@ import { initAccessibility } from './scripts/a11y.js'; import { applyStreamFadeIn } from './scripts/util/stream-fadein.js'; import { initDomHandlers } from './scripts/dom-handlers.js'; import { SimpleMutex } from './scripts/util/SimpleMutex.js'; +import { createGenerationIdleGate } from './scripts/util/generation-idle-gate.js'; import { AudioPlayer } from './scripts/audio-player.js'; import { MacroEnvBuilder } from './scripts/macros/engine/MacroEnvBuilder.js'; import { MacroEngine } from './scripts/macros/engine/MacroEngine.js'; @@ -4514,6 +4515,31 @@ function removeLastMessage() { * @property {JsonSchema} [jsonSchema] JSON schema to use for the structured generation. Usually requires a special instruction. */ +const generationIdleGate = createGenerationIdleGate(); +let generationInFlightCount = 0; + +export function waitForGenerationIdle() { + return generationIdleGate.wait(); +} + +function enterGeneration() { + if (generationInFlightCount === 0) { + generationIdleGate.markBusy(); + } + generationInFlightCount += 1; +} + +function exitGeneration() { + if (generationInFlightCount <= 0) { + throw new Error('Generation in-flight counter underflow'); + } + + generationInFlightCount -= 1; + if (generationInFlightCount === 0) { + generationIdleGate.markIdle(); + } +} + /** * MARK:Generate() * Runs a generation using the current chat context. @@ -4522,7 +4548,16 @@ function removeLastMessage() { * @param {boolean} dryRun Whether to actually generate a message or just assemble the prompt * @returns {Promise} Returns a promise that resolves when the text is done generating. */ -export async function Generate(type, { automatic_trigger, force_name2, quiet_prompt, quietToLoud, skipWIAN, force_chid, signal, quietImage, quietName, jsonSchema = null, depth = 0 } = {}, dryRun = false) { +export async function Generate(type, options = {}, dryRun = false) { + enterGeneration(); + try { + return await GenerateInternal(type, options, dryRun); + } finally { + exitGeneration(); + } +} + +async function GenerateInternal(type, { automatic_trigger, force_name2, quiet_prompt, quietToLoud, skipWIAN, force_chid, signal, quietImage, quietName, jsonSchema = null, depth = 0 } = {}, dryRun = false) { console.log('Generate entered'); setGenerationProgress(0); generation_started = new Date(); diff --git a/src/scripts/st-context.js b/src/scripts/st-context.js index 54947c4c..51fa84c2 100644 --- a/src/scripts/st-context.js +++ b/src/scripts/st-context.js @@ -67,6 +67,7 @@ import { scrollOnMediaLoad, getOneCharacter, getCharacterSource, + waitForGenerationIdle, } from '../script.js'; import { extension_settings, @@ -104,6 +105,12 @@ import { ConnectionManagerRequestService } from './extensions/shared.js'; import { updateReasoningUI, parseReasoningFromString, getReasoningTemplateByName } from './reasoning.js'; import { IGNORE_SYMBOL } from './constants.js'; import { macros } from './macros/macro-system.js'; +import { createSafeGenerate } from './util/safe-generate.js'; + +const generateSafely = createSafeGenerate({ + waitForIdle: waitForGenerationIdle, + generate: Generate, +}); export function getContext() { return { @@ -133,7 +140,7 @@ export function getContext() { addOneMessage, deleteLastMessage, deleteMessage, - generate: Generate, + generate: generateSafely, sendStreamingRequest, sendGenerationRequest, stopGeneration, diff --git a/src/scripts/util/generation-idle-gate.js b/src/scripts/util/generation-idle-gate.js new file mode 100644 index 00000000..1d27104f --- /dev/null +++ b/src/scripts/util/generation-idle-gate.js @@ -0,0 +1,52 @@ +/** + * A tiny "idle gate" primitive for generation lifecycles. + * + * Semantics: + * - `wait()` resolves only when the gate is idle. + * - `markBusy()` transitions idle -> busy and makes `wait()` pending. + * - `markIdle()` transitions busy -> idle and resolves pending `wait()`. + * + * The gate is intentionally minimal and side-effect free, so it can be reused + * across vendor and host layers and unit-tested in isolation. + * + * @returns {{ + * wait: () => Promise; + * markBusy: () => void; + * markIdle: () => void; + * }} + */ +export function createGenerationIdleGate() { + /** @type {Promise} */ + let idlePromise = Promise.resolve(); + /** @type {(() => void) | null} */ + let resolveIdlePromise = null; + + const wait = () => idlePromise; + + const markBusy = () => { + if (resolveIdlePromise) { + return; + } + + idlePromise = new Promise((resolve) => { + resolveIdlePromise = resolve; + }); + }; + + const markIdle = () => { + const resolve = resolveIdlePromise; + if (!resolve) { + return; + } + + resolveIdlePromise = null; + resolve(); + }; + + return { + wait, + markBusy, + markIdle, + }; +} + diff --git a/src/scripts/util/safe-generate.js b/src/scripts/util/safe-generate.js new file mode 100644 index 00000000..a14f636d --- /dev/null +++ b/src/scripts/util/safe-generate.js @@ -0,0 +1,31 @@ +/** + * Creates a "single-flight" wrapper for generation calls. + * + * Goals: + * - Prevent re-entrancy: wait for the previous generation to fully settle. + * - Serialize concurrent callers: multiple scripts calling `generate()` at once + * should run in order, not in parallel. + * + * @template {(...args: any[]) => any} TGenerate + * @param {{ + * waitForIdle: () => Promise; + * generate: TGenerate; + * }} deps + * @returns {(...args: Parameters) => Promise>>} + */ +export function createSafeGenerate({ waitForIdle, generate }) { + let queue = Promise.resolve(); + + return function safeGenerate(...args) { + const run = queue + .catch(() => {}) + .then(async () => { + await waitForIdle(); + return generate(...args); + }); + + queue = run.catch(() => {}); + return run; + }; +} + diff --git a/tests/generation-lifecycle-contract.test.mjs b/tests/generation-lifecycle-contract.test.mjs new file mode 100644 index 00000000..5666c544 --- /dev/null +++ b/tests/generation-lifecycle-contract.test.mjs @@ -0,0 +1,19 @@ +import test from 'node:test'; +import assert from 'node:assert/strict'; +import { readFile } from 'node:fs/promises'; +import path from 'node:path'; +import { fileURLToPath } from 'node:url'; + +const REPO_ROOT = path.resolve(path.dirname(fileURLToPath(import.meta.url)), '..'); + +test('Generate wrapper uses an in-flight lifecycle gate around GenerateInternal()', async () => { + const source = await readFile(path.join(REPO_ROOT, 'src/script.js'), 'utf8'); + + assert.match(source, /const generationIdleGate = createGenerationIdleGate\(\);/); + assert.match(source, /let generationInFlightCount = 0;/); + assert.match(source, /export function waitForGenerationIdle\(\)\s*{\s*return generationIdleGate\.wait\(\);\s*}/s); + assert.match(source, /function enterGeneration\(\)\s*{\s*if \(generationInFlightCount === 0\) {\s*generationIdleGate\.markBusy\(\);\s*}\s*generationInFlightCount \+= 1;\s*}/s); + assert.match(source, /function exitGeneration\(\)\s*{\s*if \(generationInFlightCount <= 0\) {\s*throw new Error\('Generation in-flight counter underflow'\);\s*}\s*generationInFlightCount -= 1;\s*if \(generationInFlightCount === 0\) {\s*generationIdleGate\.markIdle\(\);\s*}\s*}/s); + assert.match(source, /export async function Generate\(type, options = \{\}, dryRun = false\)\s*{\s*enterGeneration\(\);\s*try {\s*return await GenerateInternal\(type, options, dryRun\);\s*} finally {\s*exitGeneration\(\);\s*}\s*}/s); + assert.match(source, /async function GenerateInternal\(/); +}); diff --git a/tests/generation-single-flight.test.mjs b/tests/generation-single-flight.test.mjs new file mode 100644 index 00000000..51b6a1fc --- /dev/null +++ b/tests/generation-single-flight.test.mjs @@ -0,0 +1,136 @@ +import test from 'node:test'; +import assert from 'node:assert/strict'; +import path from 'node:path'; +import { fileURLToPath, pathToFileURL } from 'node:url'; + +const REPO_ROOT = path.resolve(path.dirname(fileURLToPath(import.meta.url)), '..'); + +async function importFresh(modulePath) { + const url = `${pathToFileURL(modulePath).href}?t=${Date.now()}-${Math.random()}`; + return import(url); +} + +function nextMicrotask() { + return Promise.resolve(); +} + +function nextTick() { + return new Promise((resolve) => { + setImmediate(resolve); + }); +} + +test('generation idle gate blocks until markIdle()', async () => { + const { createGenerationIdleGate } = await importFresh( + path.join(REPO_ROOT, 'src/scripts/util/generation-idle-gate.js'), + ); + + const gate = createGenerationIdleGate(); + await gate.wait(); + + gate.markBusy(); + const pending = gate.wait(); + + let resolved = false; + pending.then(() => { + resolved = true; + }); + + await nextMicrotask(); + assert.equal(resolved, false); + + gate.markBusy(); + assert.equal(gate.wait(), pending); + + gate.markIdle(); + await pending; + assert.equal(resolved, true); + + gate.markIdle(); + await gate.wait(); +}); + +test('safe generate waits for idle and serializes callers', async () => { + const { createGenerationIdleGate } = await importFresh( + path.join(REPO_ROOT, 'src/scripts/util/generation-idle-gate.js'), + ); + const { createSafeGenerate } = await importFresh( + path.join(REPO_ROOT, 'src/scripts/util/safe-generate.js'), + ); + + const gate = createGenerationIdleGate(); + gate.markBusy(); + + const calls = []; + const generate = async (...args) => { + calls.push(args); + return args.join(':'); + }; + + const safeGenerate = createSafeGenerate({ + waitForIdle: gate.wait, + generate, + }); + + const run = safeGenerate('regenerate'); + await nextMicrotask(); + assert.equal(calls.length, 0); + + gate.markIdle(); + const result = await run; + assert.deepEqual(calls, [['regenerate']]); + assert.equal(result, 'regenerate'); + + const started = []; + const resolvers = []; + const slowGenerate = (label) => { + started.push(label); + return new Promise((resolve) => { + resolvers.push(resolve); + }); + }; + + const safeSlowGenerate = createSafeGenerate({ + waitForIdle: gate.wait, + generate: slowGenerate, + }); + + const first = safeSlowGenerate('first'); + await nextTick(); + assert.deepEqual(started, ['first']); + + const second = safeSlowGenerate('second'); + await nextTick(); + assert.deepEqual(started, ['first']); + + resolvers[0]('ok-1'); + await nextTick(); + assert.deepEqual(started, ['first', 'second']); + + resolvers[1]('ok-2'); + assert.equal(await first, 'ok-1'); + assert.equal(await second, 'ok-2'); +}); + +test('safe generate does not stall after a rejected run', async () => { + const { createSafeGenerate } = await importFresh( + path.join(REPO_ROOT, 'src/scripts/util/safe-generate.js'), + ); + + let callCount = 0; + const generate = async () => { + callCount += 1; + if (callCount === 1) { + throw new Error('boom'); + } + return 'ok'; + }; + + const safeGenerate = createSafeGenerate({ + waitForIdle: () => Promise.resolve(), + generate, + }); + + await assert.rejects(() => safeGenerate('first'), /boom/); + assert.equal(await safeGenerate('second'), 'ok'); +}); diff --git a/tests/st-context-contract.test.mjs b/tests/st-context-contract.test.mjs new file mode 100644 index 00000000..4a242ea4 --- /dev/null +++ b/tests/st-context-contract.test.mjs @@ -0,0 +1,37 @@ +import test from 'node:test'; +import assert from 'node:assert/strict'; +import { readFile } from 'node:fs/promises'; +import path from 'node:path'; +import { fileURLToPath } from 'node:url'; + +const REPO_ROOT = path.resolve(path.dirname(fileURLToPath(import.meta.url)), '..'); + +test('SillyTavern global wiring still exposes getContext()', async () => { + const source = await readFile(path.join(REPO_ROOT, 'src/script.js'), 'utf8'); + + assert.match(source, /globalThis\.SillyTavern\s*=\s*\{/); + assert.match(source, /\bgetContext\b/); +}); + +test('SillyTavern context contract still exposes generate + stopGeneration', async () => { + const source = await readFile(path.join(REPO_ROOT, 'src/scripts/st-context.js'), 'utf8'); + + const contextStart = source.indexOf('export function getContext()'); + assert.ok(contextStart >= 0); + + const slice = source.slice(contextStart, contextStart + 4000); + + assert.match(slice, /\bgenerate\s*:\s*[^,\n]+,/); + assert.match(slice, /\bstopGeneration\b/); + assert.match(slice, /\beventSource\b/); + assert.match(slice, /\beventTypes\b/); +}); + +test('SillyTavern context generate is wired through safeGenerate', async () => { + const source = await readFile(path.join(REPO_ROOT, 'src/scripts/st-context.js'), 'utf8'); + + assert.match(source, /import\s+\{\s*createSafeGenerate\s*\}\s+from\s+'\.\/util\/safe-generate\.js';/); + assert.match(source, /import[\s\S]*\bwaitForGenerationIdle\b[\s\S]*from\s+'\.\.\/script\.js';/); + assert.match(source, /const generateSafely = createSafeGenerate\(\{\s*waitForIdle:\s*waitForGenerationIdle,\s*generate:\s*Generate,\s*\}\);/s); + assert.match(source, /\bgenerate\s*:\s*generateSafely,/); +});