Move keepalive scheduling to backend

18bf41c8bdd4f16a5745803d1dc7f4779c15387b

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

12 files changed, +719 -139Showing whitespace changes
homelab/README.md+1 -1
@@ -7,7 +7,7 @@ changing one, update both (scp to the NAS path + commit here).
7| File | Deployed at | Purpose |7| File | Deployed at | Purpose |
8|---|---|---|8|---|---|---|
9| `whisper-lazy-proxy.py` | `/mnt/user/appdata/whisper-lazy/proxy.py` (container `whisper-lazy`, :10300) | Wyoming lazy proxy: answers HA availability checks from cache, starts `faster-whisper` on demand, stops it after 15 min idle (frees GPU VRAM). |9| `whisper-lazy-proxy.py` | `/mnt/user/appdata/whisper-lazy/proxy.py` (container `whisper-lazy`, :10300) | Wyoming lazy proxy: answers HA availability checks from cache, starts `faster-whisper` on demand, stops it after 15 min idle (frees GPU VRAM). |
10| `runpod-lazy-proxy.py` | `/mnt/user/appdata/runpod-lazy/proxy.py` (container `runpod-lazy`, :8189) | ComfyUI-compatible on-demand proxy for a volumeless RunPod GPU pod: only `/lazy/warmup` starts or provisions, while cold ComfyUI requests return 503 for provider failover. `/lazy/status|warmup|shutdown|ping|catalog` provide SillyTavern controls and a NAS-side 15-min idle terminate. Model-set changes requested by Warm up download in place via the worker image's model-manager (port 8189, LRU disk eviction); pods are only recreated when none exists or the image predates the manager. |10| `runpod-lazy-proxy.py` | `/mnt/user/appdata/runpod-lazy/proxy.py` (container `runpod-lazy`, :8189) | ComfyUI-compatible on-demand proxy for a volumeless RunPod GPU pod: only `/lazy/warmup` starts or provisions, while cold ComfyUI requests return 503 for provider failover. `/lazy/ping` records 20-second frontend heartbeats; the proxy owns the 60-second upstream keepalive cadence and continues it through browser timer pauses up to a 5-minute lease. A NAS-side idle timer terminates the pod 15 minutes after keepalives and other activity stop. Model-set changes requested by Warm up download in place via the worker image's model-manager (port 8189, LRU disk eviction); pods are only recreated when none exists or the image predates the manager. |
11| `runpod-pod.sh` | `/boot/config/runpod-pod.sh` | CLI for the runpod-lazy control API (`status|warmup|shutdown|ping|watch`). |11| `runpod-pod.sh` | `/boot/config/runpod-pod.sh` | CLI for the runpod-lazy control API (`status|warmup|shutdown|ping|watch`). |
1212
13Related: the RunPod worker image (ComfyUI + custom nodes + boot-time model13Related: the RunPod worker image (ComfyUI + custom nodes + boot-time model
homelab/runpod-lazy-proxy.py+62 -6
@@ -14,15 +14,16 @@ Control API (CORS-enabled; ST warmup UI + runpod-pod.sh CLI):
14 GET /lazy/status -> {"state": red|orange|green, "model": ..., ...}14 GET /lazy/status -> {"state": red|orange|green, "model": ..., ...}
15 POST /lazy/warmup -> start pod for the active catalog entry15 POST /lazy/warmup -> start pod for the active catalog entry
16 POST /lazy/shutdown -> terminate pod now16 POST /lazy/shutdown -> terminate pod now
17 POST /lazy/ping -> keepalive for a ready pod (never starts a pod)17 POST /lazy/ping -> record a frontend heartbeat (never starts a pod)
18 POST /lazy/catalog -> {"models": [...], "active": "<value>"} store catalog;18 POST /lazy/catalog -> {"models": [...], "active": "<value>"} store catalog;
19 never starts or provisions a pod19 never starts or provisions a pod
2020
21Start semantics: only POST /lazy/warmup may provision or change models. ComfyUI21Start semantics: only POST /lazy/warmup may provision or change models. ComfyUI
22requests return 503 unless a pod is already ready, allowing SillyTavern's image22requests return 503 unless a pod is already ready, allowing SillyTavern's image
23target fallback chain to continue without waking this target. The idle timer23target fallback chain to continue without waking this target. Browser heartbeats
24lives here (NAS side): IDLE_SECONDS after the last activity signal the pod is24grant a short lease during which this proxy sends keepalive probes on a stable
25terminated, frontend or no frontend.25backend cadence. The idle timer lives here (NAS side): IDLE_SECONDS after the
26last activity signal the pod is terminated, frontend or no frontend.
26"""27"""
27import json28import json
28import os29import os
@@ -35,6 +36,8 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
35RUNPOD_KEY = os.environ['RUNPOD_KEY']36RUNPOD_KEY = os.environ['RUNPOD_KEY']
36LISTEN_PORT = int(os.environ.get('LISTEN_PORT', '8189'))37LISTEN_PORT = int(os.environ.get('LISTEN_PORT', '8189'))
37IDLE_SECONDS = int(os.environ.get('IDLE_SECONDS', '900'))38IDLE_SECONDS = int(os.environ.get('IDLE_SECONDS', '900'))
39KEEPALIVE_SECONDS = max(1, int(os.environ.get('KEEPALIVE_SECONDS', '60')))
40FRONTEND_LEASE_SECONDS = max(1, int(os.environ.get('FRONTEND_LEASE_SECONDS', '300')))
38DATACENTERS = [d for d in os.environ.get('DATACENTERS', '').split(',') if d]41DATACENTERS = [d for d in os.environ.get('DATACENTERS', '').split(',') if d]
39CLOUD_TYPE = os.environ.get('CLOUD_TYPE', 'SECURE')42CLOUD_TYPE = os.environ.get('CLOUD_TYPE', 'SECURE')
40IMAGE = os.environ.get('IMAGE', 'ghcr.io/permissionbrick/comfyui-runpod-worker:latest')43IMAGE = os.environ.get('IMAGE', 'ghcr.io/permissionbrick/comfyui-runpod-worker:latest')
@@ -55,6 +58,8 @@ state = {
55 'gpu': None,58 'gpu': None,
56 'since': time.time(),59 'since': time.time(),
57 'last': time.monotonic(),60 'last': time.monotonic(),
61 'frontend_last_seen': None,
62 'keepalive_last': None,
58 'ensuring': 0, # active ensure_pod waiters (status skips probing then)63 'ensuring': 0, # active ensure_pod waiters (status skips probing then)
59 'control_epoch': 0, # incremented by shutdown to cancel in-flight warmups64 'control_epoch': 0, # incremented by shutdown to cancel in-flight warmups
60 'lock': threading.Lock(),65 'lock': threading.Lock(),
@@ -364,6 +369,11 @@ def idle_reaper():
364 time.sleep(30)369 time.sleep(30)
365 idle = time.monotonic() - state['last']370 idle = time.monotonic() - state['last']
366 if idle > IDLE_SECONDS:371 if idle > IDLE_SECONDS:
372 # A heartbeat may arrive just before the next scheduled backend
373 # probe. Its live lease protects the pod until that probe catches
374 # up, without relying on the browser to refresh the idle clock.
375 if frontend_lease_seconds_left() > 0:
376 continue
367 pod_id = state['pod_id'] or find_pod()[0]377 pod_id = state['pod_id'] or find_pod()[0]
368 if pod_id:378 if pod_id:
369 log(f'idle {int(idle)}s - terminating pod')379 log(f'idle {int(idle)}s - terminating pod')
@@ -372,6 +382,49 @@ def idle_reaper():
372 state.update({'pod_id': None, 'model': None, 'gpu': None, 'phase': 'red', 'since': time.time()})382 state.update({'pod_id': None, 'model': None, 'gpu': None, 'phase': 'red', 'since': time.time()})
373383
374384
385def frontend_lease_seconds_left(now=None):
386 """Returns the remaining frontend heartbeat lease in whole seconds."""
387 now = time.monotonic() if now is None else now
388 last_seen = state['frontend_last_seen']
389 if last_seen is None:
390 return 0
391 return max(0, FRONTEND_LEASE_SECONDS - int(now - last_seen))
392
393
394def maintain_frontend_keepalive(now=None):
395 """Probes and refreshes a ready pod while a frontend heartbeat lease is live."""
396 now = time.monotonic() if now is None else now
397 with state['lock']:
398 pod_id = state['pod_id']
399 lease_active = frontend_lease_seconds_left(now) > 0
400 ready = state['phase'] == 'green' and state['ensuring'] == 0
401 if not pod_id or not lease_active or not ready or not upstream_ready(pod_id):
402 return False
403 with state['lock']:
404 # The network probe runs outside the lock. Re-check that its result still
405 # belongs to the current pod and a still-live frontend lease.
406 if state['pod_id'] != pod_id or frontend_lease_seconds_left(now) <= 0:
407 return False
408 state['last'] = now
409 state['keepalive_last'] = now
410 return True
411
412
413def keepalive_scheduler():
414 """Runs keepalive probes on a monotonic backend cadence without timer drift."""
415 next_tick = time.monotonic() + KEEPALIVE_SECONDS
416 while True:
417 time.sleep(max(0, next_tick - time.monotonic()))
418 maintain_frontend_keepalive()
419 next_tick += KEEPALIVE_SECONDS
420 now = time.monotonic()
421 if next_tick <= now:
422 # Skip missed ticks after a long process pause instead of issuing a
423 # burst of stale keepalives, while keeping future ticks aligned.
424 missed = int((now - next_tick) // KEEPALIVE_SECONDS) + 1
425 next_tick += missed * KEEPALIVE_SECONDS
426
427
375def status_body():428def status_body():
376 pod_id = state['pod_id']429 pod_id = state['pod_id']
377 previous_phase = state['phase']430 previous_phase = state['phase']
@@ -406,6 +459,8 @@ def status_body():
406 'since': state['since'],459 'since': state['since'],
407 'url': pod_url(state['pod_id']) if state['pod_id'] else None,460 'url': pod_url(state['pod_id']) if state['pod_id'] else None,
408 'idle_seconds_left': max(0, IDLE_SECONDS - int(time.monotonic() - state['last'])) if state['pod_id'] else 0,461 'idle_seconds_left': max(0, IDLE_SECONDS - int(time.monotonic() - state['last'])) if state['pod_id'] else 0,
462 'frontend_lease_seconds_left': frontend_lease_seconds_left(),
463 'keepalive_active': bool(state['pod_id'] and state['phase'] == 'green' and frontend_lease_seconds_left() > 0),
409 }).encode()464 }).encode()
410465
411466
@@ -441,8 +496,8 @@ class Proxy(BaseHTTPRequestHandler):
441 if path in ('/lazy/ping', '/lazy/warmup', '/lazy/shutdown', '/lazy/catalog') and self.command != 'POST':496 if path in ('/lazy/ping', '/lazy/warmup', '/lazy/shutdown', '/lazy/catalog') and self.command != 'POST':
442 return self._reply(405, b'{"error": "method not allowed"}')497 return self._reply(405, b'{"error": "method not allowed"}')
443 if path == '/lazy/ping':498 if path == '/lazy/ping':
444 if state['pod_id'] and state['phase'] == 'green':499 with state['lock']:
445 state['last'] = time.monotonic()500 state['frontend_last_seen'] = time.monotonic()
446 return self._reply(200, b'{"ok": true}')501 return self._reply(200, b'{"ok": true}')
447 if path == '/lazy/warmup':502 if path == '/lazy/warmup':
448 with state['lock']:503 with state['lock']:
@@ -516,6 +571,7 @@ class Proxy(BaseHTTPRequestHandler):
516571
517def main():572def main():
518 threading.Thread(target=idle_reaper, daemon=True).start()573 threading.Thread(target=idle_reaper, daemon=True).start()
574 threading.Thread(target=keepalive_scheduler, daemon=True).start()
519 # Adopt a pre-existing pod after a proxy restart without provisioning or575 # Adopt a pre-existing pod after a proxy restart without provisioning or
520 # changing it. Read-only status polls will update its readiness phase.576 # changing it. Read-only status polls will update its readiness phase.
521 pod_id, have, gpu = find_pod()577 pod_id, have, gpu = find_pod()
homelab/runpod-pod.sh+1 -1
@@ -1,5 +1,5 @@
1#!/bin/sh1#!/bin/sh
2# CLI for the runpod-lazy proxy: status | warmup | shutdown | ping2# CLI for the runpod-lazy proxy: status | warmup | shutdown | ping (renew frontend lease)
3# Usable by anyone (or any AI) without SillyTavern. Proxy owns the pod3# Usable by anyone (or any AI) without SillyTavern. Proxy owns the pod
4# lifecycle; this only talks to its control API.4# lifecycle; this only talks to its control API.
5CMD="${1:-status}"5CMD="${1:-status}"
homelab/tests/test_runpod_lazy_proxy.py+36 -0
@@ -27,6 +27,8 @@ class RunpodLazyProxyTest(unittest.TestCase):
27 'gpu': None,27 'gpu': None,
28 'since': 0,28 'since': 0,
29 'last': 100,29 'last': 100,
30 'frontend_last_seen': None,
31 'keepalive_last': None,
30 'ensuring': 0,32 'ensuring': 0,
31 'control_epoch': 0,33 'control_epoch': 0,
32 })34 })
@@ -73,9 +75,11 @@ class RunpodLazyProxyTest(unittest.TestCase):
73 self.assertEqual(status, 200)75 self.assertEqual(status, 200)
74 self.assertEqual(body['state'], 'red')76 self.assertEqual(body['state'], 'red')
7577
78 with mock.patch.object(proxy.time, 'monotonic', return_value=120):
76 status, _ = self.request('POST', '/lazy/ping')79 status, _ = self.request('POST', '/lazy/ping')
77 self.assertEqual(status, 200)80 self.assertEqual(status, 200)
78 self.assertEqual(proxy.state['last'], 100)81 self.assertEqual(proxy.state['last'], 100)
82 self.assertEqual(proxy.state['frontend_last_seen'], 120)
7983
80 catalog = {84 catalog = {
81 'models': [{'name': 'Test', 'value': 'test.safetensors', 'files': []}],85 'models': [{'name': 'Test', 'value': 'test.safetensors', 'files': []}],
@@ -100,6 +104,38 @@ class RunpodLazyProxyTest(unittest.TestCase):
100 self.assertEqual(status, 200)104 self.assertEqual(status, 200)
101 self.assertEqual(proxy.state['last'], 100)105 self.assertEqual(proxy.state['last'], 100)
102106
107 def test_backend_maintains_keepalive_during_frontend_lease(self):
108 proxy.state.update({
109 'pod_id': 'ready-pod',
110 'phase': 'green',
111 'last': 100,
112 'frontend_last_seen': 100,
113 })
114
115 with mock.patch.object(proxy, 'upstream_ready', return_value=True) as upstream_ready:
116 maintained = proxy.maintain_frontend_keepalive(now=399)
117
118 self.assertTrue(maintained)
119 self.assertEqual(proxy.state['last'], 399)
120 self.assertEqual(proxy.state['keepalive_last'], 399)
121 upstream_ready.assert_called_once_with('ready-pod')
122
123 def test_backend_stops_keepalive_after_frontend_lease_expires(self):
124 proxy.state.update({
125 'pod_id': 'ready-pod',
126 'phase': 'green',
127 'last': 100,
128 'frontend_last_seen': 100,
129 })
130
131 with mock.patch.object(proxy, 'upstream_ready') as upstream_ready:
132 maintained = proxy.maintain_frontend_keepalive(now=401)
133
134 self.assertFalse(maintained)
135 self.assertEqual(proxy.state['last'], 100)
136 self.assertIsNone(proxy.state['keepalive_last'])
137 upstream_ready.assert_not_called()
138
103 def test_get_warmup_is_rejected_without_provisioning(self):139 def test_get_warmup_is_rejected_without_provisioning(self):
104 with mock.patch.object(proxy, 'ensure_pod') as ensure_pod:140 with mock.patch.object(proxy, 'ensure_pod') as ensure_pod:
105 status, body = self.request('GET', '/lazy/warmup')141 status, body = self.request('GET', '/lazy/warmup')
public/scripts/extensions/keepalive/index.js+258 -118
@@ -1,18 +1,20 @@
1import { extension_settings, renderExtensionTemplateAsync } from '../../extensions.js';1import { extension_settings, renderExtensionTemplateAsync } from '../../extensions.js';
2import {2import {
3 activateSendButtons,3 Generate,
4 deactivateSendButtons,
5 eventSource,4 eventSource,
6 event_types,5 event_types,
7 extension_prompt_roles,6 extension_prompt_roles,
8 extension_prompt_types,7 extension_prompt_types,
9 generateQuietPrompt,8 getGenerateUrl,
9 getRequestHeaders,
10 isGenerating,10 isGenerating,
11 main_api,
11 saveSettingsDebounced,12 saveSettingsDebounced,
12 setExtensionPrompt,13 setExtensionPrompt,
13 setSendButtonState,
14 streamingProcessor,
15} from '../../../script.js';14} from '../../../script.js';
15import { createGenerationParameters, getChatCompletionModel, oai_settings } from '../../openai.js';
16import { horde_settings } from '../../horde.js';
17import { uuidv4 } from '../../utils.js';
1618
17export { init };19export { init };
1820
@@ -22,6 +24,9 @@ const MODULE = 'keepalive';
22// still processed by the backend (which is what refreshes the server-side24// still processed by the backend (which is what refreshes the server-side
23// cache), but we never pay for a full completion.25// cache), but we never pay for a full completion.
24const KEEPALIVE_RESPONSE_LENGTH = 1;26const KEEPALIVE_RESPONSE_LENGTH = 1;
27const HEARTBEAT_INTERVAL_MS = 20000;
28const RETRY_INTERVAL_MS = 5000;
29const SETTINGS_REFRESH_DELAY_MS = 300;
2530
26// Extension-prompt key used to inject the keepalive message as a USER message at depth 0.31// Extension-prompt key used to inject the keepalive message as a USER message at depth 0.
27const KEEPALIVE_INJECT_ID = 'keepalive_ping';32const KEEPALIVE_INJECT_ID = 'keepalive_ping';
@@ -37,11 +42,6 @@ const NONE_KEY = '__none__';
37// user message can pre-empt the keepalive.42// user message can pre-empt the keepalive.
38const EXPIRED_GRACE_MS = 3000;43const EXPIRED_GRACE_MS = 3000;
3944
40// If the idle timer fires more than this many ms later than it was scheduled to, assume the
41// machine slept / the tab was throttled — far more wall-clock time passed than the interval,
42// so the cache almost certainly lapsed. Skip the ping and go dormant rather than guess.
43const MAX_TIMER_DRIFT_MS = 5000;
44
45const defaultSettings = {45const defaultSettings = {
46 enabled: false,46 enabled: false,
47 timeoutSeconds: 295,47 timeoutSeconds: 295,
@@ -50,20 +50,26 @@ const defaultSettings = {
50 profiles: {},50 profiles: {},
51};51};
5252
53let idleTimer = null;53let heartbeatTimer = null;
54let keepaliveActive = false;54let jobRegistered = false;
55let refreshPromise = null;
56let refreshRequested = false;
57let refreshTimer = null;
58let preparingJob = false;
59let lifecycleVersion = 0;
55// Whether keepalive is "armed" for the CURRENT chat. It stays dormant until the user sends a new60// Whether keepalive is "armed" for the CURRENT chat. It stays dormant until the user sends a new
56// message in the open chat — on a freshly loaded or switched-to chat (even with history already61// message in the open chat — on a freshly loaded or switched-to chat (even with history already
57// present) we have no way to know whether its prompt is still cached server-side, so pinging would62// present) we have no way to know whether its prompt is still cached server-side, so pinging would
58// be pointless. Reset on page reload and whenever the chat changes.63// be pointless. Reset on page reload and whenever the chat changes.
59let armed = false;64let armed = false;
60// Wall-clock time (Date.now() + delay) the pending idle timer is expected to fire, or null when
61// none is scheduled. Used to detect a timer that fired far too late (sleep/throttle).
62let scheduledFireTime = null;
63// Per-profile "time since last message", keyed by connection profile id (or NONE_KEY).65// Per-profile "time since last message", keyed by connection profile id (or NONE_KEY).
64// In-memory only: on reload every profile is treated as freshly active.66// In-memory only: on reload every profile is treated as freshly active.
65const lastActivity = {};67const lastActivity = {};
6668
69// Per-document ids keep duplicated tabs independent. Jobs orphaned by a closed or reloaded
70// document are deliberately cleaned up by the backend lease.
71const keepaliveJobId = uuidv4().replaceAll('-', '');
72
67/**73/**
68 * @returns {string} The active connection profile id, or NONE_KEY for a manual connection.74 * @returns {string} The active connection profile id, or NONE_KEY for a manual connection.
69 */75 */
@@ -90,43 +96,49 @@ function isKeepaliveEnabledForActive() {
90 return perProfile === undefined ? true : !!perProfile;96 return perProfile === undefined ? true : !!perProfile;
91}97}
9298
93/**99function getTimeoutMs() {
94 * Clears the pending idle timer, if any.100 return Math.max(60, Number(extension_settings[MODULE].timeoutSeconds) || 295) * 1000;
95 */
96function clearIdleTimer() {
97 if (idleTimer) {
98 clearTimeout(idleTimer);
99 idleTimer = null;
100}101}
102
103async function postKeepaliveBackend(path, body) {
104 return fetch(`/api/keepalive/${path}`, {
105 method: 'POST',
106 headers: getRequestHeaders(),
107 body: JSON.stringify({ id: keepaliveJobId, ...body }),
108 });
101}109}
102110
103/**111async function stopBackendJob() {
104 * (Re)schedules the keepalive for the ACTIVE profile based on how long it has been since112 lifecycleVersion++;
105 * that profile's last message. Does nothing (and leaves the timer cleared) until keepalive113 jobRegistered = false;
106 * has been armed for the current chat, or when it is disabled for the active connection.114 try {
107 * Called on activity AND whenever the active profile changes, so each profile keeps its own115 await postKeepaliveBackend('stop', {});
108 * independent idle clock.116 } catch (error) {
109 */117 console.debug('[Keepalive] Could not stop backend job:', error);
110function scheduleIdleTimer() {118 }
111 clearIdleTimer();119}
112 if (!armed || !isKeepaliveEnabledForActive()) {120
113 scheduledFireTime = null;121async function sendBackendHeartbeat(activity = false) {
122 if (!jobRegistered || !armed || !isKeepaliveEnabledForActive()) {
114 return;123 return;
115 }124 }
116 const timeoutMs = Math.max(60, Number(extension_settings[MODULE].timeoutSeconds) || 295) * 1000;125 try {
117 const last = lastActivity[getActiveProfileKey()] ?? Date.now();126 const response = await postKeepaliveBackend('heartbeat', { activity });
118 const remaining = timeoutMs - (Date.now() - last);127 if (response.status === 404) {
119 const delay = remaining > 0 ? remaining : EXPIRED_GRACE_MS;128 jobRegistered = false;
120 scheduledFireTime = Date.now() + delay;129 queueJobRefresh();
121 idleTimer = setTimeout(fireKeepalive, delay);130 }
131 } catch (error) {
132 console.debug('[Keepalive] Backend heartbeat failed:', error);
133 }
122}134}
123135
124/**
125 * Marks the active profile as just-used (resets its idle clock) and reschedules its keepalive.
126 */
127function markActivity() {136function markActivity() {
137 if (preparingJob) {
138 return;
139 }
128 lastActivity[getActiveProfileKey()] = Date.now();140 lastActivity[getActiveProfileKey()] = Date.now();
129 scheduleIdleTimer();141 void sendBackendHeartbeat(true);
130}142}
131143
132/**144/**
@@ -146,76 +158,162 @@ function onMessageSent() {
146 */158 */
147function onChatChanged() {159function onChatChanged() {
148 armed = false;160 armed = false;
149 scheduleIdleTimer();161 void stopBackendJob();
150}162}
151163
152/**164/**
153 * Fires a background ("quiet") generation against the active model to keep the165 * Caps a prepared request to a single, non-streaming completion token.
154 * server-side prompt cache warm. The full prompt prefix (current chat context +166 * @param {string} endpoint Backend generation endpoint.
155 * the keepalive message) is sent — which is what refreshes the cache — while the167 * @param {object} payload Prepared generation payload.
156 * completion is capped to a single token (KEEPALIVE_RESPONSE_LENGTH) so we never168 * @returns {object} Cloned and capped payload.
157 * pay for a real response.
158 *
159 * Only the CURRENTLY ACTIVE profile is ever pinged — keepalive never switches the
160 * connection in the background. Each profile's idle clock is tracked separately, so
161 * switching profiles re-targets the timer at whichever profile is now active.
162 *
163 * Quiet generations are forced non-streaming by the core (Generate excludes
164 * type === 'quiet' from the StreamingProcessor) and are never saved to chat, so the
165 * output is simply discarded. generateQuietPrompt's responseLength applies a *global*
166 * temporary response-length override for the duration of the request, so we hold the
167 * same generation lock the blocking summary path uses — setSendButtonState(true) flips
168 * `is_send_press` (which every user send entry point checks and bails on) and
169 * deactivateSendButtons() reflects the busy state — and release both in `finally`.
170 */169 */
171async function fireKeepalive() {170function capKeepalivePayload(endpoint, payload) {
171 const capped = structuredClone(payload);
172 capped.stream = false;
173 for (const key of ['max_tokens', 'max_completion_tokens', 'max_length', 'max_new_tokens', 'n_predict']) {
174 if (Object.hasOwn(capped, key)) {
175 capped[key] = KEEPALIVE_RESPONSE_LENGTH;
176 }
177 }
178 if (endpoint === '/api/backends/chat-completions/generate'
179 && !Object.hasOwn(capped, 'max_tokens')
180 && !Object.hasOwn(capped, 'max_completion_tokens')) {
181 capped.max_tokens = KEEPALIVE_RESPONSE_LENGTH;
182 }
183 if (Object.hasOwn(capped, 'n')) {
184 capped.n = 1;
185 }
186 return capped;
187}
188
189/**
190 * Dry-runs prompt assembly and returns a request the backend can replay later.
191 * @returns {Promise<{endpoint: string, payload: object}>}
192 */
193async function prepareKeepaliveRequest() {
194 const api = main_api;
195 let capturedData = null;
196 const captureData = (data, dryRun) => {
197 if (dryRun && preparingJob) {
198 capturedData = structuredClone(data);
199 }
200 };
201
202 preparingJob = true;
203 eventSource.on(event_types.GENERATE_AFTER_DATA, captureData);
204 setExtensionPrompt(KEEPALIVE_INJECT_ID, extension_settings[MODULE].message, extension_prompt_types.IN_CHAT, 0, false, extension_prompt_roles.USER);
205 try {
206 await Generate('quiet', { quiet_prompt: '', force_name2: true }, true);
207 } finally {
208 setExtensionPrompt(KEEPALIVE_INJECT_ID, '', extension_prompt_types.IN_CHAT, 0, false, extension_prompt_roles.USER);
209 eventSource.removeListener(event_types.GENERATE_AFTER_DATA, captureData);
210 preparingJob = false;
211 }
212
213 if (!capturedData || api !== main_api) {
214 throw new Error('Active connection changed while preparing the keepalive request');
215 }
216
217 let endpoint;
218 let payload;
219 if (api === 'openai') {
220 endpoint = '/api/backends/chat-completions/generate';
221 const model = getChatCompletionModel(oai_settings);
222 ({ generate_data: payload } = await createGenerationParameters(oai_settings, model, 'quiet', capturedData.prompt));
223 await eventSource.emit(event_types.CHAT_COMPLETION_SETTINGS_READY, payload);
224 } else if (api === 'koboldhorde') {
225 endpoint = '/api/horde/generate-text';
226 const params = structuredClone(capturedData);
227 const prompt = params.prompt;
228 delete params.prompt;
229 Object.assign(params, {
230 n: 1,
231 max_length: KEEPALIVE_RESPONSE_LENGTH,
232 frmtadsnsp: false,
233 frmtrmblln: false,
234 frmtrmspch: false,
235 frmttriminc: false,
236 });
237 payload = {
238 prompt,
239 params,
240 trusted_workers: horde_settings.trusted_workers_only,
241 models: horde_settings.models,
242 };
243 } else {
244 endpoint = getGenerateUrl(api);
245 payload = capturedData;
246 }
247
248 return { endpoint, payload: capKeepalivePayload(endpoint, payload) };
249}
250
251async function registerBackendJob() {
172 if (!armed || !isKeepaliveEnabledForActive()) {252 if (!armed || !isKeepaliveEnabledForActive()) {
253 await stopBackendJob();
173 return;254 return;
174 }255 }
175256 const version = lifecycleVersion;
176 // If the timer fired much later than scheduled (machine sleep, background-tab throttling,257 try {
177 // etc.), far more wall-clock time has elapsed than the idle interval, so the cache has258 if (isGenerating()) {
178 // almost certainly expired. Don't ping on a stale assumption — go dormant until the next259 throw new Error('Waiting for the active generation to finish');
179 // sent message re-arms keepalive, exactly like a freshly opened chat.260 }
180 if (scheduledFireTime !== null && (Date.now() - scheduledFireTime) > MAX_TIMER_DRIFT_MS) {261 const request = await prepareKeepaliveRequest();
181 console.debug('[Keepalive] Idle timer fired', Date.now() - scheduledFireTime, 'ms late; skipping and disarming until next activity.');262 if (version !== lifecycleVersion || !armed || !isKeepaliveEnabledForActive()) {
182 armed = false;263 return;
183 clearIdleTimer();264 }
184 scheduledFireTime = null;265 const timeoutMs = getTimeoutMs();
266 const last = lastActivity[getActiveProfileKey()] ?? Date.now();
267 const remaining = timeoutMs - (Date.now() - last);
268 const delayMs = remaining > 0 ? remaining : EXPIRED_GRACE_MS;
269 const response = await postKeepaliveBackend('register', {
270 ...request,
271 intervalMs: timeoutMs,
272 delayMs,
273 });
274 if (!response.ok) {
275 throw new Error(`HTTP ${response.status}`);
276 }
277 if (version !== lifecycleVersion || !armed || !isKeepaliveEnabledForActive()) {
278 await stopBackendJob();
185 return;279 return;
186 }280 }
281 jobRegistered = true;
282 } catch (error) {
283 jobRegistered = false;
284 console.debug('[Keepalive] Could not register backend job:', error);
285 setTimeout(queueJobRefresh, RETRY_INTERVAL_MS);
286 }
287}
187288
188 // Never start while any generation is in progress, or while a previous289function queueJobRefresh() {
189 // keepalive is still running. Reschedule and bail.290 refreshRequested = true;
190 if (keepaliveActive || isGenerating() || streamingProcessor) {291 if (refreshPromise) {
191 scheduleIdleTimer();
192 return;292 return;
193 }293 }
294 refreshPromise = (async () => {
295 while (refreshRequested) {
296 refreshRequested = false;
297 await registerBackendJob();
298 }
299 })().finally(() => {
300 refreshPromise = null;
301 });
302}
194303
195 keepaliveActive = true;304function scheduleJobRefresh() {
196 setSendButtonState(true);305 clearTimeout(refreshTimer);
197 deactivateSendButtons();306 refreshTimer = setTimeout(queueJobRefresh, SETTINGS_REFRESH_DELAY_MS);
307}
198308
199 try {309function setupHeartbeatLoop() {
200 // Deliver the keepalive as a USER message at depth 0 (the quiet-prompt path forces a310 clearInterval(heartbeatTimer);
201 // system role — openai.js hardcodes role:'system' for it), using the same depth+role311 heartbeatTimer = setInterval(() => void sendBackendHeartbeat(), HEARTBEAT_INTERVAL_MS);
202 // injection Author's Note uses. The quiet prompt itself stays empty so the full chat312 document.addEventListener('visibilitychange', () => {
203 // context is still assembled (which is what keeps the cached prefix warm).313 if (document.visibilityState === 'visible') {
204 setExtensionPrompt(KEEPALIVE_INJECT_ID, extension_settings[MODULE].message, extension_prompt_types.IN_CHAT, 0, false, extension_prompt_roles.USER);314 void sendBackendHeartbeat();
205 await generateQuietPrompt({ quietPrompt: '', responseLength: KEEPALIVE_RESPONSE_LENGTH });
206 // The ping kept this profile warm; treat it as fresh activity so the next
207 // keepalive for it is a full timeout away.
208 lastActivity[getActiveProfileKey()] = Date.now();
209 } catch (error) {
210 // Best-effort: any failure here is harmless and intentionally swallowed.
211 console.debug('[Keepalive] Background generation ended:', error);
212 } finally {
213 setExtensionPrompt(KEEPALIVE_INJECT_ID, '', extension_prompt_types.IN_CHAT, 0, false, extension_prompt_roles.USER);
214 setSendButtonState(false);
215 activateSendButtons();
216 keepaliveActive = false;
217 scheduleIdleTimer();
218 }315 }
316 });
219}317}
220318
221/**319/**
@@ -293,18 +391,27 @@ function setupListeners() {
293 $('#keepalive_enabled').on('change', function () {391 $('#keepalive_enabled').on('change', function () {
294 extension_settings[MODULE].enabled = !!$(this).prop('checked');392 extension_settings[MODULE].enabled = !!$(this).prop('checked');
295 saveSettingsDebounced();393 saveSettingsDebounced();
296 scheduleIdleTimer();394 if (extension_settings[MODULE].enabled && armed) {
395 queueJobRefresh();
396 } else {
397 void stopBackendJob();
398 }
297 });399 });
298400
299 $('#keepalive_timeout').on('input', function () {401 $('#keepalive_timeout').on('input', function () {
300 extension_settings[MODULE].timeoutSeconds = Number($(this).val()) || 295;402 extension_settings[MODULE].timeoutSeconds = Number($(this).val()) || 295;
301 saveSettingsDebounced();403 saveSettingsDebounced();
302 scheduleIdleTimer();404 if (armed && isKeepaliveEnabledForActive()) {
405 scheduleJobRefresh();
406 }
303 });407 });
304408
305 $('#keepalive_message').on('input', function () {409 $('#keepalive_message').on('input', function () {
306 extension_settings[MODULE].message = String($(this).val());410 extension_settings[MODULE].message = String($(this).val());
307 saveSettingsDebounced();411 saveSettingsDebounced();
412 if (armed && isKeepaliveEnabledForActive()) {
413 scheduleJobRefresh();
414 }
308 });415 });
309416
310 // Per-profile toggles are rendered dynamically; use a delegated handler.417 // Per-profile toggles are rendered dynamically; use a delegated handler.
@@ -315,46 +422,79 @@ function setupListeners() {
315 }422 }
316 extension_settings[MODULE].profiles[profileId] = !!$(this).prop('checked');423 extension_settings[MODULE].profiles[profileId] = !!$(this).prop('checked');
317 saveSettingsDebounced();424 saveSettingsDebounced();
318 // If the toggled profile is the one currently active, reschedule immediately.425 // If the toggled profile is the one currently active, update its backend job immediately.
319 if (profileId === extension_settings.connectionManager?.selectedProfile) {426 if (profileId === extension_settings.connectionManager?.selectedProfile) {
320 scheduleIdleTimer();427 if (isKeepaliveEnabledForActive() && armed) {
428 queueJobRefresh();
429 } else {
430 void stopBackendJob();
431 }
321 }432 }
322 });433 });
323}434}
324435
436function onGenerationStarted(_type, _options, dryRun) {
437 if (!dryRun) {
438 markActivity();
439 }
440}
441
442function onGenerationEnded() {
443 markActivity();
444 if (armed && isKeepaliveEnabledForActive()) {
445 queueJobRefresh();
446 }
447}
448
449function onChatContentChanged() {
450 markActivity();
451 if (armed && isKeepaliveEnabledForActive()) {
452 queueJobRefresh();
453 }
454}
455
456async function onConnectionChanged() {
457 // Remove the previous connection's request before assembling one with the new settings.
458 await stopBackendJob();
459 if (armed && isKeepaliveEnabledForActive()) {
460 queueJobRefresh();
461 }
462}
463
325async function init() {464async function init() {
326 const settingsHtml = await renderExtensionTemplateAsync(MODULE, 'settings');465 const settingsHtml = await renderExtensionTemplateAsync(MODULE, 'settings');
327 $('#extensions_settings2').append(settingsHtml);466 $('#extensions_settings2').append(settingsHtml);
328467
329 loadSettings();468 loadSettings();
330 setupListeners();469 setupListeners();
470 setupHeartbeatLoop();
331471
332 // Only the user sending a new message arms keepalive (for the current chat).472 // Only the user sending a new message arms keepalive (for the current chat).
333 eventSource.on(event_types.MESSAGE_SENT, onMessageSent);473 eventSource.on(event_types.MESSAGE_SENT, onMessageSent);
334474
335 // These reset the active profile's idle clock (so generation time isn't counted as idle) but475 // Activity resets the backend deadline but never arms keepalive on its own. A completed
336 // never arm keepalive on their own — they also fire for background/quiet generations and while476 // generation or a chat edit also refreshes the stored request so the server always replays
337 // a chat is loading. They only re-arm the timer if a sent message already armed it.477 // the latest prompt assembled by the UI.
338 const activityEvents = [478 eventSource.on(event_types.MESSAGE_RECEIVED, markActivity);
339 event_types.MESSAGE_RECEIVED,479 eventSource.on(event_types.GENERATION_STARTED, onGenerationStarted);
340 event_types.GENERATION_STARTED,480 eventSource.on(event_types.GENERATION_ENDED, onGenerationEnded);
341 event_types.GENERATION_ENDED,481 eventSource.on(event_types.MESSAGE_EDITED, onChatContentChanged);
342 ];482 eventSource.on(event_types.MESSAGE_DELETED, onChatContentChanged);
343 for (const event of activityEvents) {483 eventSource.on(event_types.MESSAGE_UPDATED, onChatContentChanged);
344 eventSource.on(event, markActivity);484 eventSource.on(event_types.MESSAGE_SWIPED, onChatContentChanged);
345 }
346485
347 // Opening/switching a chat disarms keepalive — it must be re-armed by a new sent message.486 // Opening/switching a chat disarms keepalive — it must be re-armed by a new sent message.
348 eventSource.on(event_types.CHAT_CHANGED, onChatChanged);487 eventSource.on(event_types.CHAT_CHANGED, onChatChanged);
349488
350 // When the active profile changes, re-target the idle timer at the new profile using489 // When the active profile changes, replace the backend job without marking it as activity.
351 // ITS own last-activity time (do NOT mark activity — switching is not using it).490 eventSource.on(event_types.CONNECTION_PROFILE_LOADED, onConnectionChanged);
352 eventSource.on(event_types.CONNECTION_PROFILE_LOADED, scheduleIdleTimer);491 eventSource.on(event_types.MAIN_API_CHANGED, onConnectionChanged);
492 eventSource.on(event_types.CHATCOMPLETION_SOURCE_CHANGED, onConnectionChanged);
493 eventSource.on(event_types.CHATCOMPLETION_MODEL_CHANGED, onConnectionChanged);
494 eventSource.on(event_types.PRESET_CHANGED, onConnectionChanged);
353495
354 // Keep the per-profile toggle list in sync with Connection Manager.496 // Keep the per-profile toggle list in sync with Connection Manager.
355 eventSource.on(event_types.CONNECTION_PROFILE_CREATED, renderProfileToggles);497 eventSource.on(event_types.CONNECTION_PROFILE_CREATED, renderProfileToggles);
356 eventSource.on(event_types.CONNECTION_PROFILE_UPDATED, renderProfileToggles);498 eventSource.on(event_types.CONNECTION_PROFILE_UPDATED, renderProfileToggles);
357 eventSource.on(event_types.CONNECTION_PROFILE_DELETED, renderProfileToggles);499 eventSource.on(event_types.CONNECTION_PROFILE_DELETED, renderProfileToggles);
358
359 scheduleIdleTimer();
360}500}
public/scripts/extensions/keepalive/manifest.json+1 -1
@@ -6,7 +6,7 @@
6 "js": "index.js",6 "js": "index.js",
7 "css": "",7 "css": "",
8 "author": "SillyTavern",8 "author": "SillyTavern",
9 "version": "1.0.0",9 "version": "1.1.0",
10 "homePage": "https://github.com/SillyTavern/SillyTavern",10 "homePage": "https://github.com/SillyTavern/SillyTavern",
11 "hooks": {11 "hooks": {
12 "activate": "init"12 "activate": "init"
public/scripts/extensions/keepalive/settings.html+1 -1
@@ -16,7 +16,7 @@
16 <label for="keepalive_message" data-i18n="ext_keepalive_message">Keepalive message</label>16 <label for="keepalive_message" data-i18n="ext_keepalive_message">Keepalive message</label>
17 <input id="keepalive_message" class="text_pole" type="text" />17 <input id="keepalive_message" class="text_pole" type="text" />
1818
19 <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. Sending is briefly locked while it runs. Only the profile you're currently using is pinged; each profile has its own idle timer. Keepalive stays off until you send a new message in the current chat (opening, switching to, or reloading a chat isn't proof its prompt is still cached), and it also skips a ping if its timer fired far later than scheduled (e.g. after the machine slept).</small>19 <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 new message in the current chat (opening, switching to, or reloading a chat isn't proof its prompt is still cached).</small>
2020
21 <hr>21 <hr>
2222
public/scripts/extensions/stable-diffusion/index.js+12 -10
@@ -1541,11 +1541,11 @@ async function fetchReferenceImageBase64(refImage) {
1541/** Poll cadence for the pod status indicator (faster while it is starting). */1541/** Poll cadence for the pod status indicator (faster while it is starting). */
1542const RUNPOD_POLL_IDLE_MS = 30000;1542const RUNPOD_POLL_IDLE_MS = 30000;
1543const RUNPOD_POLL_BUSY_MS = 5000;1543const RUNPOD_POLL_BUSY_MS = 5000;
1544/** Keepalive cadence for a pod that the status API has confirmed is ready. */1544/** Frontend heartbeat cadence. The proxy tolerates browser timer pauses for five minutes. */
1545const RUNPOD_PING_MS = 60000;1545const RUNPOD_HEARTBEAT_MS = 20000;
15461546
1547let runpodStatusTimer = null;1547let runpodStatusTimer = null;
1548let runpodPingTimer = null;1548let runpodHeartbeatTimer = null;
1549let runpodLastPhase = 'red';1549let runpodLastPhase = 'red';
15501550
1551function getRunpodLazyUrl() {1551function getRunpodLazyUrl() {
@@ -1690,20 +1690,21 @@ async function runpodControl(action) {
1690 runpodStatusTimer = setTimeout(pollRunpodStatus, RUNPOD_POLL_BUSY_MS);1690 runpodStatusTimer = setTimeout(pollRunpodStatus, RUNPOD_POLL_BUSY_MS);
1691}1691}
16921692
1693function runpodKeepalive() {1693function runpodHeartbeat() {
1694 const url = getRunpodLazyUrl();1694 const url = getRunpodLazyUrl();
1695 if (!url || runpodLastPhase !== 'green') {1695 if (!url) {
1696 return;1696 return;
1697 }1697 }
1698 // The proxy also rejects wake-on-ping; this phase guard avoids unnecessary1698 // This only renews the proxy's frontend lease and never starts a pod. The
1699 // requests whenever the last read-only status poll says the pod is down.1699 // proxy owns the stable upstream keepalive cadence so browser timer
1700 // throttling cannot directly skew it.
1700 fetch(`${url}/lazy/ping`, { method: 'POST', signal: AbortSignal.timeout(5000) }).catch(() => { });1701 fetch(`${url}/lazy/ping`, { method: 'POST', signal: AbortSignal.timeout(5000) }).catch(() => { });
1701}1702}
17021703
1703/** (Re)starts the polling/keepalive loops based on the configured URL. */1704/** (Re)starts the status polling and frontend-heartbeat loops. */
1704function setupRunpodLoops() {1705function setupRunpodLoops() {
1705 clearTimeout(runpodStatusTimer);1706 clearTimeout(runpodStatusTimer);
1706 clearInterval(runpodPingTimer);1707 clearInterval(runpodHeartbeatTimer);
1707 ensureRunpodBarDot();1708 ensureRunpodBarDot();
1708 const barDot = document.getElementById('sd_runpod_bar_dot');1709 const barDot = document.getElementById('sd_runpod_bar_dot');
1709 if (barDot) {1710 if (barDot) {
@@ -1714,7 +1715,8 @@ function setupRunpodLoops() {
1714 return;1715 return;
1715 }1716 }
1716 pollRunpodStatus();1717 pollRunpodStatus();
1717 runpodPingTimer = setInterval(runpodKeepalive, RUNPOD_PING_MS);1718 runpodHeartbeat();
1719 runpodHeartbeatTimer = setInterval(runpodHeartbeat, RUNPOD_HEARTBEAT_MS);
1718}1720}
17191721
1720function onRunpodLazyUrlInput() {1722function onRunpodLazyUrlInput() {
public/scripts/extensions/stable-diffusion/settings.html+1 -1
@@ -100,7 +100,7 @@
100 <span data-i18n="Copy ComfyUI URL">Copy ComfyUI URL</span>100 <span data-i18n="Copy ComfyUI URL">Copy ComfyUI URL</span>
101 </div>101 </div>
102 </div>102 </div>
103 <small data-i18n="sd_runpod_small">The pod starts only when you press Warm up. A SillyTavern tab keeps an already-ready pod alive; it shuts down 15 minutes after the last activity. Image requests fail over instead of starting a stopped pod. Red = off, orange = starting/downloading models, green = ready.</small>103 <small data-i18n="sd_runpod_small">The pod starts only when you press Warm up. A SillyTavern tab sends a heartbeat every 20 seconds; the proxy keeps an already-ready pod alive on its own stable timer and tolerates browser sleep for up to 5 minutes. The pod shuts down 15 minutes after keepalives and other activity stop. Image requests fail over instead of starting a stopped pod. Red = off, orange = starting/downloading models, green = ready.</small>
104 <h5 data-i18n="Model catalog">Model catalog</h5>104 <h5 data-i18n="Model catalog">Model catalog</h5>
105 <small data-i18n="sd_runpod_models_small">When the ComfyUI URL points at the proxy, the Model dropdown lists these entries instead of asking the backend. Press Warm up to apply the selected model and download its files; changing the selection or generating an image never starts or provisions a stopped pod. Filename = what &quot;%model%&quot; resolves to; downloads = one &quot;subpath url&quot; per line, including companion files (text encoder, VAE, LoRA).</small>105 <small data-i18n="sd_runpod_models_small">When the ComfyUI URL points at the proxy, the Model dropdown lists these entries instead of asking the backend. Press Warm up to apply the selected model and download its files; changing the selection or generating an image never starts or provisions a stopped pod. Filename = what &quot;%model%&quot; resolves to; downloads = one &quot;subpath url&quot; per line, including companion files (text encoder, VAE, LoRA).</small>
106 <div id="sd_runpod_models_list" class="flex-container flexFlowColumn marginTopBot5"></div>106 <div id="sd_runpod_models_list" class="flex-container flexFlowColumn marginTopBot5"></div>
src/endpoints/keepalive.js+264 -0
@@ -0,0 +1,264 @@
1import express from 'express';
2import https from 'node:https';
3import fetch from 'node-fetch';
4
5export const router = express.Router();
6
7export const FRONTEND_LEASE_MS = 5 * 60 * 1000;
8const MIN_INTERVAL_MS = 60 * 1000;
9const MAX_INTERVAL_MS = 24 * 60 * 60 * 1000;
10const MAX_JOB_ID_LENGTH = 128;
11const ALLOWED_ENDPOINTS = new Set([
12 '/api/backends/chat-completions/generate',
13 '/api/backends/text-completions/generate',
14 '/api/backends/kobold/generate',
15 '/api/novelai/generate',
16 '/api/horde/generate-text',
17]);
18const localHttpsAgent = new https.Agent({ rejectUnauthorized: false });
19
20function advanceDeadline(deadline, interval, now) {
21 if (deadline > now) {
22 return deadline;
23 }
24 const missedIntervals = Math.floor((now - deadline) / interval) + 1;
25 return deadline + missedIntervals * interval;
26}
27
28/**
29 * Backend-owned scheduler for prompt-cache keepalive requests.
30 */
31export class KeepaliveScheduler {
32 #jobs = new Map();
33 #timer = null;
34 #now;
35 #dispatch;
36 #setTimer;
37 #clearTimer;
38
39 constructor({
40 now = Date.now,
41 dispatch = dispatchKeepaliveRequest,
42 setTimer = setTimeout,
43 clearTimer = clearTimeout,
44 } = {}) {
45 this.#now = now;
46 this.#dispatch = dispatch;
47 this.#setTimer = setTimer;
48 this.#clearTimer = clearTimer;
49 }
50
51 register(key, job) {
52 const now = this.#now();
53 const intervalMs = Math.min(MAX_INTERVAL_MS, Math.max(MIN_INTERVAL_MS, Number(job.intervalMs) || MIN_INTERVAL_MS));
54 const requestedDelay = Number(job.delayMs);
55 const delayMs = Number.isFinite(requestedDelay)
56 ? Math.min(intervalMs, Math.max(0, requestedDelay))
57 : intervalMs;
58 this.#jobs.set(key, {
59 ...job,
60 intervalMs,
61 lastSeen: now,
62 nextRun: now + delayMs,
63 inFlight: false,
64 });
65 this.#schedule();
66 return this.#jobs.get(key);
67 }
68
69 heartbeat(key, activity = false, request = null) {
70 const job = this.#jobs.get(key);
71 if (!job) {
72 return null;
73 }
74 const now = this.#now();
75 if ((now - job.lastSeen) >= FRONTEND_LEASE_MS) {
76 this.#jobs.delete(key);
77 this.#schedule();
78 return null;
79 }
80 job.lastSeen = now;
81 if (activity) {
82 job.nextRun = now + job.intervalMs;
83 }
84 if (request) {
85 job.request = request;
86 }
87 this.#schedule();
88 return job;
89 }
90
91 stop(key) {
92 const removed = this.#jobs.delete(key);
93 this.#schedule();
94 return removed;
95 }
96
97 get(key) {
98 return this.#jobs.get(key) ?? null;
99 }
100
101 async tick(now = this.#now()) {
102 const pending = [];
103 for (const [key, job] of this.#jobs) {
104 if ((now - job.lastSeen) >= FRONTEND_LEASE_MS) {
105 this.#jobs.delete(key);
106 continue;
107 }
108 if (job.nextRun > now) {
109 continue;
110 }
111
112 job.nextRun = advanceDeadline(job.nextRun, job.intervalMs, now);
113 if (job.inFlight) {
114 continue;
115 }
116 job.inFlight = true;
117 pending.push(Promise.resolve(this.#dispatch(job)).catch(error => {
118 console.debug(`[Keepalive] Backend request failed for ${job.id}:`, error);
119 }).finally(() => {
120 job.inFlight = false;
121 this.#schedule();
122 }));
123 }
124 this.#schedule();
125 await Promise.all(pending);
126 }
127
128 #schedule() {
129 if (this.#timer) {
130 this.#clearTimer(this.#timer);
131 this.#timer = null;
132 }
133 if (!this.#jobs.size) {
134 return;
135 }
136
137 let nextWake = Infinity;
138 for (const job of this.#jobs.values()) {
139 nextWake = Math.min(nextWake, job.nextRun, job.lastSeen + FRONTEND_LEASE_MS);
140 }
141 const delay = Math.max(0, nextWake - this.#now());
142 this.#timer = this.#setTimer(() => {
143 this.#timer = null;
144 void this.tick();
145 }, delay);
146 this.#timer?.unref?.();
147 }
148}
149
150async function dispatchKeepaliveRequest(job) {
151 const controller = new AbortController();
152 const timeout = setTimeout(() => controller.abort(new Error('Keepalive request timed out')), 120000);
153 timeout.unref?.();
154 try {
155 const response = await fetch(`${job.request.origin}${job.endpoint}`, {
156 method: 'POST',
157 headers: job.request.headers,
158 body: JSON.stringify(job.payload),
159 signal: controller.signal,
160 agent: job.request.secure ? localHttpsAgent : undefined,
161 });
162 await response.arrayBuffer();
163 if (!response.ok) {
164 throw new Error(`HTTP ${response.status}`);
165 }
166 } finally {
167 clearTimeout(timeout);
168 }
169}
170
171const scheduler = new KeepaliveScheduler();
172
173function getJobKey(request, id) {
174 return `${request.user.profile.handle}:${id}`;
175}
176
177function getJobId(request) {
178 const id = String(request.body?.id ?? '');
179 if (!id || id.length > MAX_JOB_ID_LENGTH || !/^[a-zA-Z0-9_-]+$/.test(id)) {
180 throw new Error('Invalid keepalive job id');
181 }
182 return id;
183}
184
185function getReplayRequest(request) {
186 const localPort = Number(request.socket.localPort);
187 if (!Number.isInteger(localPort) || localPort < 1 || localPort > 65535) {
188 throw new Error('Unable to determine local server port');
189 }
190 let localAddress = request.socket.localAddress || '127.0.0.1';
191 if (localAddress === '::' || localAddress === '0:0:0:0:0:0:0:0') {
192 localAddress = '::1';
193 } else if (localAddress === '0.0.0.0') {
194 localAddress = '127.0.0.1';
195 } else if (localAddress.startsWith('::ffff:')) {
196 localAddress = localAddress.slice('::ffff:'.length);
197 }
198 const loopbackHost = localAddress.includes(':') ? `[${localAddress}]` : localAddress;
199 const headers = {
200 'Content-Type': 'application/json',
201 'Cookie': request.get('cookie') || '',
202 'X-CSRF-Token': request.get('x-csrf-token') || '',
203 'Host': request.get('host') || '',
204 };
205 for (const name of ['authorization', 'origin', 'referer', 'user-agent']) {
206 const value = request.get(name);
207 if (value) {
208 headers[name] = value;
209 }
210 }
211 const secure = !!request.socket.encrypted;
212 return {
213 origin: `${secure ? 'https' : 'http'}://${loopbackHost}:${localPort}`,
214 headers,
215 secure,
216 };
217}
218
219router.post('/register', (request, response) => {
220 try {
221 const id = getJobId(request);
222 const endpoint = String(request.body?.endpoint ?? '');
223 if (!ALLOWED_ENDPOINTS.has(endpoint)) {
224 return response.status(400).send({ error: 'Unsupported keepalive endpoint' });
225 }
226 if (!request.body?.payload || typeof request.body.payload !== 'object' || Array.isArray(request.body.payload)) {
227 return response.status(400).send({ error: 'Missing keepalive payload' });
228 }
229 const job = scheduler.register(getJobKey(request, id), {
230 id,
231 endpoint,
232 payload: request.body.payload,
233 intervalMs: request.body.intervalMs,
234 delayMs: request.body.delayMs,
235 request: getReplayRequest(request),
236 });
237 return response.send({ ok: true, nextRun: job.nextRun, leaseMs: FRONTEND_LEASE_MS });
238 } catch (error) {
239 return response.status(400).send({ error: String(error.message || error) });
240 }
241});
242
243router.post('/heartbeat', (request, response) => {
244 try {
245 const id = getJobId(request);
246 const job = scheduler.heartbeat(getJobKey(request, id), !!request.body?.activity, getReplayRequest(request));
247 if (!job) {
248 return response.sendStatus(404);
249 }
250 return response.send({ ok: true, nextRun: job.nextRun, leaseMs: FRONTEND_LEASE_MS });
251 } catch (error) {
252 return response.status(400).send({ error: String(error.message || error) });
253 }
254});
255
256router.post('/stop', (request, response) => {
257 try {
258 const id = getJobId(request);
259 scheduler.stop(getJobKey(request, id));
260 return response.send({ ok: true });
261 } catch (error) {
262 return response.status(400).send({ error: String(error.message || error) });
263 }
264});
src/server-startup.js+2 -0
@@ -51,6 +51,7 @@ import { router as dataMaidRouter } from './endpoints/data-maid.js';
51import { router as backupsRouter } from './endpoints/backups.js';51import { router as backupsRouter } from './endpoints/backups.js';
52import { router as imageMetadataRouter } from './endpoints/image-metadata.js';52import { router as imageMetadataRouter } from './endpoints/image-metadata.js';
53import { router as volcengineRouter } from './endpoints/volcengine.js';53import { router as volcengineRouter } from './endpoints/volcengine.js';
54import { router as keepaliveRouter } from './endpoints/keepalive.js';
5455
55/**56/**
56 * @typedef {object} ServerStartupResult57 * @typedef {object} ServerStartupResult
@@ -185,6 +186,7 @@ export function setupPrivateEndpoints(app) {
185 app.use('/api/data-maid', dataMaidRouter);186 app.use('/api/data-maid', dataMaidRouter);
186 app.use('/api/backups', backupsRouter);187 app.use('/api/backups', backupsRouter);
187 app.use('/api/image-metadata', imageMetadataRouter);188 app.use('/api/image-metadata', imageMetadataRouter);
189 app.use('/api/keepalive', keepaliveRouter);
188}190}
189191
190/**192/**
tests/keepalive-scheduler.test.js+80 -0
@@ -0,0 +1,80 @@
1import { describe, expect, jest, test } from '@jest/globals';
2import { FRONTEND_LEASE_MS, KeepaliveScheduler } from '../src/endpoints/keepalive.js';
3
4function createScheduler() {
5 let now = 0;
6 const dispatch = jest.fn().mockResolvedValue(undefined);
7 const scheduler = new KeepaliveScheduler({
8 now: () => now,
9 dispatch,
10 setTimer: () => ({ unref() { } }),
11 clearTimer: () => { },
12 });
13 return { scheduler, dispatch, setNow: value => { now = value; } };
14}
15
16describe('KeepaliveScheduler', () => {
17 test('heartbeats renew the lease without moving the backend deadline', () => {
18 const { scheduler, setNow } = createScheduler();
19 scheduler.register('user:tab', { id: 'tab', intervalMs: 60000, delayMs: 60000 });
20
21 setNow(20000);
22 scheduler.heartbeat('user:tab');
23
24 expect(scheduler.get('user:tab').lastSeen).toBe(20000);
25 expect(scheduler.get('user:tab').nextRun).toBe(60000);
26 });
27
28 test('a heartbeat cannot revive an already expired lease', () => {
29 const { scheduler, setNow } = createScheduler();
30 scheduler.register('user:tab', { id: 'tab', intervalMs: 60000, delayMs: 60000 });
31
32 setNow(FRONTEND_LEASE_MS);
33
34 expect(scheduler.heartbeat('user:tab')).toBeNull();
35 expect(scheduler.get('user:tab')).toBeNull();
36 });
37
38 test('accepts an immediate first deadline', () => {
39 const { scheduler } = createScheduler();
40 scheduler.register('user:tab', { id: 'tab', intervalMs: 60000, delayMs: 0 });
41
42 expect(scheduler.get('user:tab').nextRun).toBe(0);
43 });
44
45 test('dispatches on the backend deadline and stays aligned after a late tick', async () => {
46 const { scheduler, dispatch, setNow } = createScheduler();
47 scheduler.register('user:tab', { id: 'tab', intervalMs: 60000, delayMs: 60000 });
48
49 setNow(60000);
50 await scheduler.tick();
51 expect(dispatch).toHaveBeenCalledTimes(1);
52 expect(scheduler.get('user:tab').nextRun).toBe(120000);
53
54 setNow(190000);
55 await scheduler.tick();
56 expect(dispatch).toHaveBeenCalledTimes(2);
57 expect(scheduler.get('user:tab').nextRun).toBe(240000);
58 });
59
60 test('activity resets the next keepalive deadline', () => {
61 const { scheduler, setNow } = createScheduler();
62 scheduler.register('user:tab', { id: 'tab', intervalMs: 60000, delayMs: 60000 });
63
64 setNow(50000);
65 scheduler.heartbeat('user:tab', true);
66
67 expect(scheduler.get('user:tab').nextRun).toBe(110000);
68 });
69
70 test('removes a job after five minutes without a frontend heartbeat', async () => {
71 const { scheduler, dispatch, setNow } = createScheduler();
72 scheduler.register('user:tab', { id: 'tab', intervalMs: 60000, delayMs: 60000 });
73
74 setNow(FRONTEND_LEASE_MS);
75 await scheduler.tick();
76
77 expect(scheduler.get('user:tab')).toBeNull();
78 expect(dispatch).not.toHaveBeenCalled();
79 });
80});