Keep expired keepalives dormant on wake

0aa62e9004c1517fe3801ca9fef82ecf60d6bc29

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

Parents: 8346f51967f3e6fdf658f3f24949220d553991a4

public/scripts/extensions/keepalive/index.js+55 -8 Viewed

 // Idle-tracking key used when no connection profile is active (manual API panel).
 const NONE_KEY = '__none__';
-// When the active profile is already past its timeout (e.g. you just switched back
-// to a profile that went cold), wait this short grace before pinging so an imminent
-// user message can pre-empt the keepalive.
-const EXPIRED_GRACE_MS = 3000;
 const defaultSettings = {
     enabled: false,
     timeoutSeconds: 295,

 // Per-profile "time since last message", keyed by connection profile id (or NONE_KEY).
 // In-memory only: on reload every profile is treated as freshly active.
 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.
+const lastKeepalive = {};
 // Per-document ids keep duplicated tabs independent. Jobs orphaned by a closed or reloaded
 // document are deliberately cleaned up by the backend lease.

     return Math.max(60, Number(extension_settings[MODULE].timeoutSeconds) || 295) * 1000;
+}
+function updateLastKeepalive(profileKey, value) {
+    const timestamp = Number(value);
+    if (Number.isFinite(timestamp) && timestamp > 0) {
+        lastKeepalive[profileKey] = Math.max(lastKeepalive[profileKey] || 0, timestamp);
+    }
+}
+function getLastWarmTime(profileKey = getActiveProfileKey()) {
+    return Math.max(lastActivity[profileKey] || 0, lastKeepalive[profileKey] || 0);
+}
+function isWarmPromptExpired(profileKey = getActiveProfileKey()) {
+    const lastWarmTime = getLastWarmTime(profileKey);
+    return lastWarmTime > 0 && (Date.now() - lastWarmTime) >= getTimeoutMs();
+}
+function disarmExpiredJob() {
+    armed = false;
+    jobRegistered = false;
+    refreshRequested = false;
+    lifecycleVersion++;
+    clearTimeout(retryTimer);
+    clearTimeout(refreshTimer);
+    clearTimeout(generationRefreshTimer);
+}
 async function postKeepaliveBackend(path, body) {
     return fetch(`/api/keepalive/${path}`, {
         method: 'POST',

     // was missed (or registration failed), the regular heartbeat must self-heal instead of doing
     // nothing forever.
     if (!jobRegistered) {
+        if (isWarmPromptExpired()) {
+            disarmExpiredJob();
+            return;
+        }
         queueJobRefresh();
         return;
+    }
+    const version = lifecycleVersion;
+    const profileKey = getActiveProfileKey();
     try {
         const response = await postKeepaliveBackend('heartbeat', { activity });
+        const state = await response.json().catch(() => ({}));
+        if (version !== lifecycleVersion) {
+            return;
+        }
+        updateLastKeepalive(profileKey, state.lastRun);
         if (response.status === 404) {
             jobRegistered = false;
+            if (isWarmPromptExpired(profileKey)) {
+                disarmExpiredJob();
+                return;
+            }
             queueJobRefresh();
+        }
     } catch (error) {

         await stopBackendJob();
         return;
+    }
+    const profileKey = getActiveProfileKey();
+    if (isWarmPromptExpired(profileKey)) {
+        disarmExpiredJob();
+        return;
+    }
     const version = lifecycleVersion;
     try {
         if (isGenerating()) {

             return;
+        }
         const timeoutMs = getTimeoutMs();
-        const last = lastActivity[getActiveProfileKey()] ?? Date.now();
+        const last = getLastWarmTime(profileKey) || Date.now();
         const remaining = timeoutMs - (Date.now() - last);
-        const delayMs = remaining > 0 ? remaining : EXPIRED_GRACE_MS;
+        if (remaining <= 0) {
+            disarmExpiredJob();
+            return;
+        }
         const response = await postKeepaliveBackend('register', {
             ...request,
             intervalMs: timeoutMs,
-            delayMs,
+            delayMs: remaining,
         });
         if (!response.ok) {
             throw new Error(`HTTP ${response.status}`);

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

     "js": "index.js",
     "css": "",
     "author": "SillyTavern",
-    "version": "1.1.2",
+    "version": "1.1.3",
     "homePage": "https://github.com/SillyTavern/SillyTavern",
     "hooks": {
         "activate": "init"

public/scripts/extensions/keepalive/settings.html+1 -1 Viewed

             <label for="keepalive_message" data-i18n="ext_keepalive_message">Keepalive message</label>
             <input id="keepalive_message" class="text_pole" type="text" />
-            <small data-i18n="ext_keepalive_help">After this many seconds without activity on the active connection profile, the current chat plus this message (sent as a hidden user message) is submitted to that profile to keep its server-side prompt cache warm. The completion is capped to a single token and is never saved to the chat. The server owns the keepalive timer, while this tab renews its lease every 20 seconds; if the tab sleeps or closes, the server continues on schedule for up to five minutes after the last heartbeat. Only the profile you're currently using is pinged. Keepalive stays off until you send a message or start a foreground generation (including regenerate, swipe, or continue); merely opening, switching to, or reloading a chat is not proof its prompt is cached.</small>
+            <small data-i18n="ext_keepalive_help">After this many seconds without activity on the active connection profile, the current chat plus this message (sent as a hidden user message) is submitted to that profile to keep its server-side prompt cache warm. The completion is capped to a single token and is never saved to the chat. The server owns the keepalive timer, while this tab renews its lease every 20 seconds; if the tab sleeps or closes, the server continues on schedule for up to five minutes after the last heartbeat. Waking a tab after both its lease and cache window have expired does not restart keepalives until a new message, regenerate, swipe, or continue. Only the profile you're currently using is pinged. Merely opening, switching to, or reloading a chat is not proof its prompt is cached.</small>
             <hr>

src/endpoints/keepalive.js+51 -6 Viewed

 export const FRONTEND_LEASE_MS = 5 * 60 * 1000;
 const MIN_INTERVAL_MS = 60 * 1000;
 const MAX_INTERVAL_MS = 24 * 60 * 60 * 1000;
+const EXPIRED_JOB_RETENTION_MS = MAX_INTERVAL_MS;
+const MAX_EXPIRED_JOBS = 1000;
 const MAX_JOB_ID_LENGTH = 128;
 const ALLOWED_ENDPOINTS = new Set([
     '/api/backends/chat-completions/generate',

  */
 export class KeepaliveScheduler {
     #jobs = new Map();
+    #expiredJobs = new Map();
     #timer = null;
     #now;
     #dispatch;

     register(key, job) {
         const now = this.#now();
+        this.#pruneExpiredJobs(now);
+        this.#expiredJobs.delete(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)

             lastSeen: now,
             nextRun: now + delayMs,
             inFlight: false,
+            lastRun: null,
         });
         this.#schedule();
         return this.#jobs.get(key);

@@ -73,7 +79,7 @@ export class KeepaliveScheduler {
73	}	79	}
74	const now = this.#now();	80	const now = this.#now();
75	if ((now - job.lastSeen) >= FRONTEND_LEASE_MS) {	81	if ((now - job.lastSeen) >= FRONTEND_LEASE_MS) {
76	this.#jobs.delete(key);	82	this.#expire(key, job, now);
77	this.#schedule();	83	this.#schedule();
78	return null;	84	return null;
79	}	85	}

     stop(key) {
         const removed = this.#jobs.delete(key);
+        this.#expiredJobs.delete(key);
         this.#schedule();
         return removed;
+    }

         return this.#jobs.get(key) ?? null;
+    }
+    getExpired(key) {
+        this.#pruneExpiredJobs(this.#now());
+        return this.#expiredJobs.get(key) ?? null;
+    }
     async tick(now = this.#now()) {
         const pending = [];
         for (const [key, job] of this.#jobs) {
             if ((now - job.lastSeen) >= FRONTEND_LEASE_MS) {
-                this.#jobs.delete(key);
+                this.#expire(key, job, now);
                 continue;
+            }
             if (job.nextRun > now) {

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

         await Promise.all(pending);
+    }
+    #expire(key, job, now) {
+        this.#jobs.delete(key);
+        this.#expiredJobs.set(key, {
+            expiredAt: now,
+            intervalMs: job.intervalMs,
+            lastRun: job.lastRun,
+        });
+        this.#pruneExpiredJobs(now);
+    }
+    #pruneExpiredJobs(now) {
+        for (const [key, job] of this.#expiredJobs) {
+            if ((now - job.expiredAt) >= EXPIRED_JOB_RETENTION_MS) {
+                this.#expiredJobs.delete(key);
+            }
+        }
+        while (this.#expiredJobs.size > MAX_EXPIRED_JOBS) {
+            this.#expiredJobs.delete(this.#expiredJobs.keys().next().value);
+        }
+    }
     #schedule() {
         if (this.#timer) {
             this.#clearTimer(this.#timer);

 router.post('/heartbeat', (request, response) => {
     try {
         const id = getJobId(request);
-        const job = scheduler.heartbeat(getJobKey(request, id), !!request.body?.activity, getReplayRequest(request));
+        const key = getJobKey(request, id);
+        const job = scheduler.heartbeat(key, !!request.body?.activity, getReplayRequest(request));
         if (!job) {
-            return response.sendStatus(404);
+            const expiredJob = scheduler.getExpired(key);
+            return response.status(404).send({
+                expired: !!expiredJob,
+                lastRun: expiredJob?.lastRun ?? null,
+            });
+        }
-        return response.send({ ok: true, nextRun: job.nextRun, leaseMs: FRONTEND_LEASE_MS });
+        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) });
+    }

tests/keepalive-scheduler.test.js+22 -0 Viewed

         expect(scheduler.heartbeat('user:tab')).toBeNull();
         expect(scheduler.get('user:tab')).toBeNull();
+        expect(scheduler.getExpired('user:tab')).toEqual({
+            expiredAt: FRONTEND_LEASE_MS,
+            intervalMs: 60000,
+            lastRun: null,
+        });
     });
     test('accepts an immediate first deadline', () => {

         await scheduler.tick();
         expect(dispatch).toHaveBeenCalledTimes(2);
         expect(scheduler.get('user:tab').nextRun).toBe(240000);
+        expect(scheduler.get('user:tab').lastRun).toBe(190000);
     });
     test('activity resets the next keepalive deadline', () => {

         expect(scheduler.get('user:tab')).toBeNull();
         expect(dispatch).not.toHaveBeenCalled();
     });
+    test('retains the last successful keepalive after the frontend lease expires', async () => {
+        const { scheduler, setNow } = createScheduler();
+        scheduler.register('user:tab', { id: 'tab', intervalMs: 60000, delayMs: 60000 });
+        setNow(240000);
+        await scheduler.tick();
+        setNow(FRONTEND_LEASE_MS);
+        await scheduler.tick();
+        expect(scheduler.getExpired('user:tab')).toEqual({
+            expiredAt: FRONTEND_LEASE_MS,
+            intervalMs: 60000,
+            lastRun: 240000,
+        });
+    });
 });