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	} else {	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
      */
-    function onError(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+90 -42 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 since last message", keyed by connection profile id (or NONE_KEY).
+// Per-profile time of the last verified successful main-chat inference request, keyed 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	async function sendBackendHeartbeat(activity = false) {	152	async function sendBackendHeartbeat() {
151	if (temporaryConnectionDepth > 0 \|\| !armed \|\| !isKeepaliveEnabledForActive()) {	153	if (temporaryConnectionDepth > 0 \|\| !armed \|\| !isKeepaliveEnabledForActive()) {
152	return;	154	return;
153	}	155	}

     const version = lifecycleVersion;
     const profileKey = getActiveProfileKey();
     try {
-        const response = await postKeepaliveBackend('heartbeat', { activity });
+        const response = await postKeepaliveBackend('heartbeat', {});
         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 arms keepalive before its foreground generation begins.
+ * A new user message may lead to a main-chat request, but 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;
-    markActivity();
+    queueJobRefresh();
+}
 /**

  */
 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);

     });
+}
-function onGenerationStarted(type, _options, dryRun) {
+async function pauseBackendJob() {
-    // 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.
-    if (dryRun || type === 'quiet') {
         return;
+    }
-    // Regenerates, swipes, and continues send the current prompt to the provider without emitting
+    try {
-    // MESSAGE_SENT. They are just as valid proof of a warm prompt as a newly typed message.
+        const response = await postKeepaliveBackend('pause', {});
-    foregroundGenerationActive = true;
+        if (response.status === 404) {
-    armed = true;
+            jobRegistered = false;
-    markActivity();
+        }
+    } catch (error) {
+        console.debug('[Keepalive] Could not pause backend job:', error);
+    }
+}
-function onGenerationEnded() {
+async function resumeBackendJob(activityStartedAt = null) {
-    if (!foregroundGenerationActive) {
+    if (!jobRegistered) {
         return;
+    }
-    foregroundGenerationActive = false;
+    const elapsedMs = Number.isFinite(activityStartedAt)
-    markActivity();
+        ? Math.max(0, Date.now() - activityStartedAt)
-    scheduleGenerationJobRefresh();
+        : 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;
+    }
+    armed = true;
+    activeMainRequests.set(request.id, {
+        profileKey: getActiveProfileKey(),
+        startedAt: Number(request.startedAt) || Date.now(),
+    });
+    if (activeMainRequests.size === 1) {
+        await pauseBackendJob();
+    }
+}
-function onMessageReceived(_messageId, type) {
+async function onGenerationRequestFinished(request) {
-    // Image-generation extensions add their results as chat messages. Those messages neither used
+    const activeRequest = activeMainRequests.get(request?.id);
-    // nor refreshed the foreground text cache, so they must not move its keepalive deadline.
+    if (!activeRequest) {
-    if (type === 'extension') {
+        return;
+    }
+    activeMainRequests.delete(request.id);
+    if (request.success) {
+        const startedAt = activeRequest.startedAt;
+        lastActivity[activeRequest.profileKey] = Math.max(lastActivity[activeRequest.profileKey] || 0, startedAt);
+        latestSuccessfulRequestStartedAt = Math.max(latestSuccessfulRequestStartedAt, startedAt);
+        armed = true;
+    }
+    if (activeMainRequests.size > 0) {
         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 before generation starts. Foreground generations that do
+    // A user message arms keepalive, but only the exact provider-request lifecycle below is allowed
-    // not send a new message (regenerate/swipe/continue) are armed by onGenerationStarted instead.
+    // to move its deadline. Regenerate/swipe/continue requests arm it when their 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_RECEIVED, onMessageReceived);
+    eventSource.on(event_types.GENERATION_REQUEST_STARTED, onGenerationRequestStarted);
-    eventSource.on(event_types.GENERATION_STARTED, onGenerationStarted);
+    eventSource.on(event_types.GENERATION_REQUEST_FINISHED, onGenerationRequestFinished);
-    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.5",
+    "version": "1.1.6",
     "homePage": "https://github.com/SillyTavern/SillyTavern",
     "hooks": {
         "activate": "init"

src/endpoints/keepalive.js+91 -17 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;
-        this.#jobs.set(key, {
+        const requestedNextRun = now + delayMs;
-            ...job,
+        if (existing) {
-            intervalMs,
+            // Re-registering updates the replay payload, but cannot claim that the provider cache
-            lastSeen: now,
+            // was refreshed. Mutate the existing object so an in-flight dispatch can still clear
-            nextRun: now + delayMs,
+            // its state, while preserving an earlier deadline and any foreground-request pause.
-            inFlight: false,
+            Object.assign(existing, job, {
-            lastRun: null,
+                intervalMs,
-        });
+                lastSeen: now,
+                nextRun: Math.min(existing.nextRun, requestedNextRun),
+            });
+        } else {
+            this.#jobs.set(key, {
+                ...job,
+                intervalMs,
+                lastSeen: now,
+                nextRun: requestedNextRun,
+                inFlight: false,
+                paused: false,
+                lastRun: null,
+            });
+        }
         this.#schedule();
         return this.#jobs.get(key);
+    }
-    heartbeat(key, activity = false, request = null) {
+    heartbeat(key, request = null) {
         const job = this.#jobs.get(key);
         if (!job) {
             return null;

@@ -84,9 +99,6 @@ export class KeepaliveScheduler {
84	return null;	99	return null;
85	}	100	}
86	job.lastSeen = now;	101	job.lastSeen = now;
87	if (activity) {
88	job.nextRun = now + job.intervalMs;
89	}
90	if (request) {	102	if (request) {
91	job.request = request;	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));
+        const job = scheduler.heartbeat(key, 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('activity resets the next keepalive deadline', () => {
+    test('re-registering refreshes payload without postponing 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.heartbeat('user:tab', true);
+        scheduler.pause('user:tab');
+        setNow(70000);
+        scheduler.resume('user:tab', 30000);
+        expect(scheduler.get('user:tab').nextRun).toBe(100000);
+    });
-        expect(scheduler.get('user:tab').nextRun).toBe(110000);
+    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(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 () => {