Tie keepalives to successful chat inference

1385bdd1ce3c8b3af2691f342db02a8735e56997

permissionBRICK <40219477+permissionBRICK@users.noreply.github.com>

Parents: 1154224f7b8872e08761c775d378cccf2c5de83b

public/script.js+32 -2 Viewed

         return Promise.resolve();
+    }
 +    let activeProviderRequest = null;
++
 +    async function beginProviderRequest() {
 +        activeProviderRequest = {
 +            id: uuidv4(),
 +            type,
 +            startedAt: Date.now(),
 +        };
 +        await eventSource.emit(event_types.GENERATION_REQUEST_STARTED, activeProviderRequest);
 +    }
++
 +    async function finishProviderRequest(success) {
 +        if (!activeProviderRequest) {
 +            return;
 +        }
 +        const request = activeProviderRequest;
 +        activeProviderRequest = null;
 +        await eventSource.emit(event_types.GENERATION_REQUEST_FINISHED, {
 +            ...request,
 +            success: !!success,
 +            finishedAt: Date.now(),
 +        });
 +    }
++
     /**
      * Saves itemized prompt bits and calls streaming or non-streaming generation API.
      * @returns {Promise<void|*|Awaited<*>|String|{fromStream}|string|undefined|Object>}

                 streamingProcessor.firstMessageText = '';
+            }
 +            await beginProviderRequest();
             streamingProcessor.generator = await sendStreamingRequest(type, generate_data, { jsonSchema });
             hideSwipeButtons();

             const isStreamFinished = streamingProcessor && !streamingProcessor.isStopped && streamingProcessor.isFinished;
             const isStreamWithToolCalls = streamingProcessor && Array.isArray(streamingProcessor.toolCalls) && streamingProcessor.toolCalls.length;
 +            await finishProviderRequest(isStreamFinished);
             if (canPerformToolCalls && isStreamFinished && isStreamWithToolCalls) {
                 const lastMessage = chat[chat.length - 1];
                 const hasToolCalls = ToolManager.hasToolCalls(streamingProcessor.toolCalls);

@@ -5389,7 +5415,10 @@ export async function Generate(type, { automatic_trigger, force_name2, quiet_pro
5389	5415	});
5390	5416	}
5391	5417	} else {
5392		- return await sendGenerationRequest(type, generate_data, { jsonSchema });
	5418	+ await beginProviderRequest();
	5419	+ const data = await sendGenerationRequest(type, generate_data, { jsonSchema });
	5420	+ await finishProviderRequest(!data?.error);
	5421	+ return data;
5393	5422	}
5394	5423	}
5395	5424

      * @param {Error|object} exception Error or response JSON
      * @throws {Error|object} Re-throws the exception
      */
     async function onError(exception) {
         // if the response JSON was thrown (novel|textgenerationwebui|kobold), show the error message
         if (typeof exception?.error?.message === 'string') {
             toastr.error(exception.error.message, t`Text generation error`, { timeOut: 10000, extendedTimeOut: 20000 });
+        }
 +        await finishProviderRequest(false);
         unblockGeneration(type);
         console.log(exception);
         streamingProcessor = null;

public/scripts/events.js+2 -0 Viewed

     GENERATION_STARTED: 'generation_started',
     GENERATION_STOPPED: 'generation_stopped',
     GENERATION_ENDED: 'generation_ended',
 +    GENERATION_REQUEST_STARTED: 'generation_request_started',
 +    GENERATION_REQUEST_FINISHED: 'generation_request_finished',
     SD_PROMPT_PROCESSING: 'sd_prompt_processing',
     EXTENSIONS_FIRST_LOAD: 'extensions_first_load',
     EXTENSION_SETTINGS_LOADED: 'extension_settings_loaded',

public/scripts/extensions/keepalive/index.js+89 -41 Viewed

 const HEARTBEAT_INTERVAL_MS = 20000;
 const RETRY_INTERVAL_MS = 5000;
 const SETTINGS_REFRESH_DELAY_MS = 300;
 +const MAIN_CHAT_GENERATION_TYPES = new Set(['normal', 'regenerate', 'swipe', 'continue']);
 // Extension-prompt key used to inject the keepalive message as a USER message at depth 0.
 const KEEPALIVE_INJECT_ID = 'keepalive_ping';

 let retryTimer = null;
 let preparingJob = false;
 let lifecycleVersion = 0;
 -let foregroundGenerationActive = false;
 let temporaryConnectionDepth = 0;
 let refreshAfterTemporaryConnection = false;
 +const activeMainRequests = new Map();
 +let latestSuccessfulRequestStartedAt = 0;
 // Whether keepalive is "armed" for the CURRENT chat. It stays dormant until the user sends a new
 // message or starts another foreground generation in the open chat. Merely loading a chat does not
 // prove that its prompt is cached server-side. Reset on page reload and whenever the chat changes.
 let armed = false;
 // Per-profile "time sinceof lastthe message",last keyedverified bysuccessful connectionmain-chat profileinference idrequest, (orkeyed NONE_KEY).by
 -// In-memory only: on reload every profile is treated as freshly active.
 +// connection profile id (or NONE_KEY). In-memory only; a reload starts disarmed.
 const lastActivity = {};
 // Last successful backend keepalive, keyed by connection profile. The backend returns this after a
 // sleeping tab wakes so we can distinguish a still-warm prompt from an expired one.

@@ -147,7 +149,7 @@ async function stopBackendJob() {
147	149	}
148	150	}
149	151
150	152	async function sendBackendHeartbeat(activity = false) {
151	153	if (temporaryConnectionDepth > 0 \|\| !armed \|\| !isKeepaliveEnabledForActive()) {
152	154	return;
153	155	}

     const version = lifecycleVersion;
     const profileKey = getActiveProfileKey();
     try {
         const response = await postKeepaliveBackend('heartbeat', { activity });
         const state = await response.json().catch(() => ({}));
         if (version !== lifecycleVersion) {
             return;

+    }
+}
 -function markActivity() {
 -    if (preparingJob || temporaryConnectionDepth > 0) {
 -        return;
 -    }
 -    lastActivity[getActiveProfileKey()] = Date.now();
 -    void sendBackendHeartbeat(true);
 -}
+-
 /**
  * A new user message armsmay keepalivelead beforeto itsa foregroundmain-chat generationrequest, beginsbut does not prove the provider accepted it.
 + * Arm the feature and refresh only the replay payload; the deadline remains unchanged until the
 + * request-success event arrives.
  */
 function onMessageSent() {
     armed = true;
     markActivityqueueJobRefresh();
+}
 /**

  */
 function onChatChanged() {
     armed = false;
 +    activeMainRequests.clear();
 +    latestSuccessfulRequestStartedAt = 0;
     void stopBackendJob();
+}

         return;
+    }
     const profileKey = getActiveProfileKey();
 +    if (!getLastWarmTime(profileKey)) {
 +        return;
 +    }
     if (isWarmPromptExpired(profileKey)) {
         disarmExpiredJob();
         return;
+    }
     const version = lifecycleVersion;
 +    const hadRegisteredJob = jobRegistered;
     try {
         if (isGenerating()) {
             throw new Error('Waiting for the active generation to finish');

         clearTimeout(retryTimer);
         jobRegistered = true;
     } catch (error) {
 -        jobRegistered = false;
 +        // A failed payload refresh does not mean the existing backend job disappeared. Keeping
 +        // this state is especially important while a foreground request is starting, because its
 +        // lifecycle handler still needs to pause that existing job before provider dispatch.
 +        jobRegistered = hadRegisteredJob;
         console.debug('[Keepalive] Could not register backend job:', error);
         clearTimeout(retryTimer);
         retryTimer = setTimeout(queueJobRefresh, RETRY_INTERVAL_MS);

     });
+}
 async function onGenerationStartedpauseBackendJob(type, _options, dryRun) {
 -    // Quiet/background generations (image prompts, summaries, helpers, etc.) use a different model
 +    if (!jobRegistered) {
 -    // or context and do not refresh the foreground text-chat cache.
 +        return;
 -    if (dryRun || type === 'quiet') {
 +    }
 +    try {
 +        const response = await postKeepaliveBackend('pause', {});
 +        if (response.status === 404) {
 +            jobRegistered = false;
 +        }
 +    } catch (error) {
 +        console.debug('[Keepalive] Could not pause backend job:', error);
 +    }
 +}
++
 +async function resumeBackendJob(activityStartedAt = null) {
 +    if (!jobRegistered) {
 +        return;
 +    }
 +    const elapsedMs = Number.isFinite(activityStartedAt)
 +        ? Math.max(0, Date.now() - activityStartedAt)
 +        : null;
 +    const profileKey = getActiveProfileKey();
 +    try {
 +        const response = await postKeepaliveBackend('resume', {
 +            activityElapsedMs: elapsedMs,
 +        });
 +        const state = await response.json().catch(() => ({}));
 +        updateLastKeepalive(profileKey, state.lastRun);
 +        if (response.status === 404) {
 +            jobRegistered = false;
 +        }
 +    } catch (error) {
 +        console.debug('[Keepalive] Could not resume backend job:', error);
 +    }
 +}
++
 +async function onGenerationRequestStarted(request) {
 +    if (temporaryConnectionDepth > 0 || !request?.id || !MAIN_CHAT_GENERATION_TYPES.has(request.type)) {
         return;
+    }
 -    // Regenerates, swipes, and continues send the current prompt to the provider without emitting
 -    // MESSAGE_SENT. They are just as valid proof of a warm prompt as a newly typed message.
 -    foregroundGenerationActive = true;
     armed = true;
 -    markActivity();
 +    activeMainRequests.set(request.id, {
 +        profileKey: getActiveProfileKey(),
 +        startedAt: Number(request.startedAt) || Date.now(),
 +    });
 +    if (activeMainRequests.size === 1) {
 +        await pauseBackendJob();
 +    }
+}
 async function onGenerationEndedonGenerationRequestFinished(request) {
 -    if (!foregroundGenerationActive) {
 +    const activeRequest = activeMainRequests.get(request?.id);
 +    if (!activeRequest) {
         return;
+    }
 -    foregroundGenerationActive = false;
 +    activeMainRequests.delete(request.id);
 -    markActivity();
++
 -    scheduleGenerationJobRefresh();
 +    if (request.success) {
 +        const startedAt = activeRequest.startedAt;
 +        lastActivity[activeRequest.profileKey] = Math.max(lastActivity[activeRequest.profileKey] || 0, startedAt);
 +        latestSuccessfulRequestStartedAt = Math.max(latestSuccessfulRequestStartedAt, startedAt);
 +        armed = true;
+    }
 -function onMessageReceived(_messageId, type) {
 +    if (activeMainRequests.size > 0) {
 -    // Image-generation extensions add their results as chat messages. Those messages neither used
 -    // nor refreshed the foreground text cache, so they must not move its keepalive deadline.
 -    if (type === 'extension') {
         return;
+    }
 -    foregroundGenerationActive = false;
++
 -    markActivity();
 +    const successfulStartedAt = latestSuccessfulRequestStartedAt || null;
 +    latestSuccessfulRequestStartedAt = 0;
 +    await resumeBackendJob(successfulStartedAt);
++
 +    // A failed request may still have changed local chat content, so refresh the replay payload,
 +    // but registerBackendJob() will retain the old deadline. A successful request supplies the
 +    // verified activity timestamp used for the new deadline.
     scheduleGenerationJobRefresh();
+}

     setupListeners();
     setupHeartbeatLoop();
     // A new user message arms keepalive, beforebut generationonly starts.the Foregroundexact generationsprovider-request thatlifecycle dobelow is allowed
     // not sendto amove newits messagedeadline. (regenerateRegenerate/swipe/continue) arerequests armedarm byit onGenerationStartedwhen insteadtheir request starts.
     eventSource.on(event_types.MESSAGE_SENT, onMessageSent);
     // Only actual text-provider generations reset the cache deadline. Local chat edits still
     // refresh the stored request so the server replays the latest prompt without pretending that
     // those edits refreshed the provider cache.
     eventSource.on(event_types.MESSAGE_RECEIVEDGENERATION_REQUEST_STARTED, onMessageReceivedonGenerationRequestStarted);
     eventSource.on(event_types.GENERATION_STARTEDGENERATION_REQUEST_FINISHED, onGenerationStartedonGenerationRequestFinished);
 -    eventSource.on(event_types.GENERATION_ENDED, onGenerationEnded);
 -    eventSource.on(event_types.GENERATION_STOPPED, () => { foregroundGenerationActive = false; });
     eventSource.on(event_types.MESSAGE_EDITED, onChatContentChanged);
     eventSource.on(event_types.MESSAGE_DELETED, onChatContentChanged);
     eventSource.on(event_types.MESSAGE_UPDATED, onChatContentChanged);

public/scripts/extensions/keepalive/manifest.json+1 -1 Viewed

     "js": "index.js",
     "css": "",
     "author": "SillyTavern",
     "version": "1.1.56",
     "homePage": "https://github.com/SillyTavern/SillyTavern",
     "hooks": {
         "activate": "init"

src/endpoints/keepalive.js+84 -10 Viewed

 export const FRONTEND_LEASE_MS = 5 * 60 * 1000;
 const MIN_INTERVAL_MS = 60 * 1000;
 const MAX_INTERVAL_MS = 24 * 60 * 60 * 1000;
 +const FAILED_REQUEST_RETRY_MS = 5 * 1000;
 const EXPIRED_JOB_RETENTION_MS = MAX_INTERVAL_MS;
 const MAX_EXPIRED_JOBS = 1000;
 const MAX_JOB_ID_LENGTH = 128;

         const now = this.#now();
         this.#pruneExpiredJobs(now);
         this.#expiredJobs.delete(key);
 +        const existing = this.#jobs.get(key);
         const intervalMs = Math.min(MAX_INTERVAL_MS, Math.max(MIN_INTERVAL_MS, Number(job.intervalMs) || MIN_INTERVAL_MS));
         const requestedDelay = Number(job.delayMs);
         const delayMs = Number.isFinite(requestedDelay)
             ? Math.min(intervalMs, Math.max(0, requestedDelay))
             : intervalMs;
 +        const requestedNextRun = now + delayMs;
 +        if (existing) {
 +            // Re-registering updates the replay payload, but cannot claim that the provider cache
 +            // was refreshed. Mutate the existing object so an in-flight dispatch can still clear
 +            // its state, while preserving an earlier deadline and any foreground-request pause.
 +            Object.assign(existing, job, {
 +                intervalMs,
 +                lastSeen: now,
 +                nextRun: Math.min(existing.nextRun, requestedNextRun),
 +            });
 +        } else {
             this.#jobs.set(key, {
                 ...job,
                 intervalMs,
                 lastSeen: now,
                             nextRun: now + delayMsrequestedNextRun,
                 inFlight: false,
 +                paused: false,
                 lastRun: null,
             });
 +        }
         this.#schedule();
         return this.#jobs.get(key);
+    }
     heartbeat(key, activity = false, request = null) {
         const job = this.#jobs.get(key);
         if (!job) {
             return null;

@@ -84,9 +99,6 @@ export class KeepaliveScheduler {
84	99	return null;
85	100	}
86	101	job.lastSeen = now;
87		- if (activity) {
88		- job.nextRun = now + job.intervalMs;
89		- }
90	102	if (request) {
91	103	job.request = request;
92	104	}

         return job;
+    }
 +    pause(key) {
 +        const job = this.#jobs.get(key);
 +        if (!job) {
 +            return null;
 +        }
 +        job.lastSeen = this.#now();
 +        job.paused = true;
 +        this.#schedule();
 +        return job;
 +    }
++
 +    resume(key, activityElapsedMs = null) {
 +        const job = this.#jobs.get(key);
 +        if (!job) {
 +            return null;
 +        }
 +        const now = this.#now();
 +        job.lastSeen = now;
 +        job.paused = false;
 +        const elapsedMs = Number(activityElapsedMs);
 +        if (activityElapsedMs !== null && activityElapsedMs !== undefined && Number.isFinite(elapsedMs) && elapsedMs >= 0) {
 +            job.nextRun = now + Math.max(0, job.intervalMs - elapsedMs);
 +        }
 +        this.#schedule();
 +        return job;
 +    }
++
     stop(key) {
         const removed = this.#jobs.delete(key);
         this.#expiredJobs.delete(key);

             if (job.nextRun > now) {
                 continue;
+            }
+-
 +            if (job.inFlight || job.paused) {
 -            job.nextRun = advanceDeadline(job.nextRun, job.intervalMs, now);
 -            if (job.inFlight) {
                 continue;
+            }
 +            const scheduledRun = job.nextRun;
             job.inFlight = true;
             pending.push(Promise.resolve(this.#dispatch(job)).then(() => {
                 const completedAt = this.#now();
                 job.lastRun = completedAt;
 +                job.nextRun = advanceDeadline(scheduledRun, job.intervalMs, completedAt);
                 const expiredJob = this.#expiredJobs.get(key);
                 if (expiredJob) {
                     expiredJob.lastRun = completedAt;
+                }
             }).catch(error => {
                 console.debug(`[Keepalive] Backend request failed for ${job.id}:`, error);
 +                job.nextRun = this.#now() + FAILED_REQUEST_RETRY_MS;
             }).finally(() => {
                 job.inFlight = false;
                 this.#schedule();

         let nextWake = Infinity;
         for (const job of this.#jobs.values()) {
 -            nextWake = Math.min(nextWake, job.nextRun, job.lastSeen + FRONTEND_LEASE_MS);
 +            const nextRun = (job.inFlight || job.paused) ? Infinity : job.nextRun;
 +            nextWake = Math.min(nextWake, nextRun, job.lastSeen + FRONTEND_LEASE_MS);
+        }
         const delay = Math.max(0, nextWake - this.#now());
         this.#timer = this.#setTimer(() => {

     try {
         const id = getJobId(request);
         const key = getJobKey(request, id);
         const job = scheduler.heartbeat(key, !!request.body?.activity, getReplayRequest(request));
         if (!job) {
             const expiredJob = scheduler.getExpired(key);
             return response.status(404).send({

+    }
 });
 +router.post('/pause', (request, response) => {
 +    try {
 +        const id = getJobId(request);
 +        const job = scheduler.pause(getJobKey(request, id));
 +        if (!job) {
 +            return response.status(404).send({ expired: false });
 +        }
 +        console.info(`[Keepalive] Paused job ${id} while foreground inference is running`);
 +        return response.send({ ok: true, nextRun: job.nextRun, lastRun: job.lastRun, leaseMs: FRONTEND_LEASE_MS });
 +    } catch (error) {
 +        return response.status(400).send({ error: String(error.message || error) });
 +    }
 +});
++
 +router.post('/resume', (request, response) => {
 +    try {
 +        const id = getJobId(request);
 +        const job = scheduler.resume(getJobKey(request, id), request.body?.activityElapsedMs);
 +        if (!job) {
 +            const expiredJob = scheduler.getExpired(getJobKey(request, id));
 +            return response.status(404).send({
 +                expired: !!expiredJob,
 +                lastRun: expiredJob?.lastRun ?? null,
 +            });
 +        }
 +        const reset = request.body?.activityElapsedMs !== null && request.body?.activityElapsedMs !== undefined;
 +        console.info(`[Keepalive] Resumed job ${id}; deadline ${reset ? 'reset by successful foreground inference' : 'preserved after unsuccessful foreground inference'}`);
 +        return response.send({ ok: true, nextRun: job.nextRun, lastRun: job.lastRun, leaseMs: FRONTEND_LEASE_MS });
 +    } catch (error) {
 +        return response.status(400).send({ error: String(error.message || error) });
 +    }
 +});
++
 router.post('/stop', (request, response) => {
     try {
         const id = getJobId(request);

tests/keepalive-scheduler.test.js+51 -3 Viewed

         expect(scheduler.get('user:tab').lastRun).toBe(190000);
     });
     test('activityre-registering resetsrefreshes thepayload nextwithout keepalivepostponing the deadline', () => {
 +        const { scheduler, setNow } = createScheduler();
 +        scheduler.register('user:tab', { id: 'tab', payload: { version: 1 }, intervalMs: 60000, delayMs: 60000 });
++
 +        setNow(50000);
 +        scheduler.register('user:tab', { id: 'tab', payload: { version: 2 }, intervalMs: 60000, delayMs: 60000 });
++
 +        expect(scheduler.get('user:tab').nextRun).toBe(60000);
 +        expect(scheduler.get('user:tab').payload).toEqual({ version: 2 });
 +    });
++
 +    test('pauses a due keepalive while foreground inference is running', async () => {
 +        const { scheduler, dispatch, setNow } = createScheduler();
 +        scheduler.register('user:tab', { id: 'tab', intervalMs: 60000, delayMs: 60000 });
++
 +        setNow(50000);
 +        scheduler.pause('user:tab');
 +        setNow(70000);
 +        await scheduler.tick();
 +        expect(dispatch).not.toHaveBeenCalled();
++
 +        scheduler.resume('user:tab');
 +        await scheduler.tick();
 +        expect(dispatch).toHaveBeenCalledTimes(1);
 +    });
++
 +    test('only a successful foreground request advances the deadline', () => {
         const { scheduler, setNow } = createScheduler();
         scheduler.register('user:tab', { id: 'tab', intervalMs: 60000, delayMs: 60000 });
         setNow(50000);
         scheduler.heartbeatpause('user:tab', true);
 +        setNow(70000);
 +        scheduler.resume('user:tab', 30000);
++
 +        expect(scheduler.get('user:tab').nextRun).toBe(100000);
 +    });
++
 +    test('a failed keepalive retries without advancing a full interval', async () => {
 +        let now = 0;
 +        const dispatch = jest.fn().mockRejectedValue(new Error('HTTP 502'));
 +        const scheduler = new KeepaliveScheduler({
 +            now: () => now,
 +            dispatch,
 +            setTimer: () => ({ unref() { } }),
 +            clearTimer: () => { },
 +        });
 +        const debug = jest.spyOn(console, 'debug').mockImplementation(() => { });
 +        scheduler.register('user:tab', { id: 'tab', intervalMs: 60000, delayMs: 60000 });
++
 +        now = 60000;
 +        await scheduler.tick();
 -        expect(scheduler.get('user:tab').nextRun).toBe(110000);
 +        expect(dispatch).toHaveBeenCalledTimes(1);
 +        expect(scheduler.get('user:tab').nextRun).toBe(65000);
 +        debug.mockRestore();
     });
     test('removes a job after five minutes without a frontend heartbeat', async () => {