Move keepalive scheduling to backend
| @@ -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`). | |
| 12 | 12 | ||
| 13 | Related: the RunPod worker image (ComfyUI + custom nodes + boot-time model | 13 | Related: the RunPod worker image (ComfyUI + custom nodes + boot-time model |
| @@ -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 entry | 15 | POST /lazy/warmup -> start pod for the active catalog entry |
| 16 | POST /lazy/shutdown -> terminate pod now | 16 | 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 pod | 19 | never starts or provisions a pod |
| 20 | 20 | ||
| 21 | Start semantics: only POST /lazy/warmup may provision or change models. ComfyUI | 21 | Start semantics: only POST /lazy/warmup may provision or change models. ComfyUI |
| 22 | requests return 503 unless a pod is already ready, allowing SillyTavern's image | 22 | requests return 503 unless a pod is already ready, allowing SillyTavern's image |
| 23 | target fallback chain to continue without waking this target. The idle timer | 23 | target fallback chain to continue without waking this target. Browser heartbeats |
| 24 | lives here (NAS side): IDLE_SECONDS after the last activity signal the pod is | 24 | grant a short lease during which this proxy sends keepalive probes on a stable |
| 25 | terminated, frontend or no frontend. | 25 | backend cadence. The idle timer lives here (NAS side): IDLE_SECONDS after the |
| 26 | last activity signal the pod is terminated, frontend or no frontend. | ||
| 26 | """ | 27 | """ |
| 27 | import json | 28 | import json |
| 28 | import os | 29 | import os |
| @@ -35,6 +36,8 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer | |||
| 35 | RUNPOD_KEY = os.environ['RUNPOD_KEY'] | 36 | RUNPOD_KEY = os.environ['RUNPOD_KEY'] |
| 36 | LISTEN_PORT = int(os.environ.get('LISTEN_PORT', '8189')) | 37 | LISTEN_PORT = int(os.environ.get('LISTEN_PORT', '8189')) |
| 37 | IDLE_SECONDS = int(os.environ.get('IDLE_SECONDS', '900')) | 38 | IDLE_SECONDS = int(os.environ.get('IDLE_SECONDS', '900')) |
| 39 | KEEPALIVE_SECONDS = max(1, int(os.environ.get('KEEPALIVE_SECONDS', '60'))) | ||
| 40 | FRONTEND_LEASE_SECONDS = max(1, int(os.environ.get('FRONTEND_LEASE_SECONDS', '300'))) | ||
| 38 | DATACENTERS = [d for d in os.environ.get('DATACENTERS', '').split(',') if d] | 41 | DATACENTERS = [d for d in os.environ.get('DATACENTERS', '').split(',') if d] |
| 39 | CLOUD_TYPE = os.environ.get('CLOUD_TYPE', 'SECURE') | 42 | CLOUD_TYPE = os.environ.get('CLOUD_TYPE', 'SECURE') |
| 40 | IMAGE = os.environ.get('IMAGE', 'ghcr.io/permissionbrick/comfyui-runpod-worker:latest') | 43 | IMAGE = 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 warmups | 64 | '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()}) |
| 373 | 383 | ||
| 374 | 384 | ||
| 385 | def 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 | |||
| 394 | def 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 | |||
| 413 | def 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 | |||
| 375 | def status_body(): | 428 | def 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() |
| 410 | 465 | ||
| 411 | 466 | ||
| @@ -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): | |||
| 516 | 571 | ||
| 517 | def main(): | 572 | def 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 or | 575 | # 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() |
| @@ -1,5 +1,5 @@ | |||
| 1 | #!/bin/sh | 1 | #!/bin/sh |
| 2 | # CLI for the runpod-lazy proxy: status | warmup | shutdown | ping | 2 | # 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 pod | 3 | # 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. |
| 5 | CMD="${1:-status}" | 5 | CMD="${1:-status}" |
| @@ -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') |
| 75 | 77 | ||
| 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) | ||
| 79 | 83 | ||
| 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) |
| 102 | 106 | ||
| 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') |
| @@ -1,18 +1,20 @@ | |||
| 1 | import { extension_settings, renderExtensionTemplateAsync } from '../../extensions.js'; | 1 | import { extension_settings, renderExtensionTemplateAsync } from '../../extensions.js'; |
| 2 | import { | 2 | import { |
| 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'; |
| 15 | import { createGenerationParameters, getChatCompletionModel, oai_settings } from '../../openai.js'; | ||
| 16 | import { horde_settings } from '../../horde.js'; | ||
| 17 | import { uuidv4 } from '../../utils.js'; | ||
| 16 | 18 | ||
| 17 | export { init }; | 19 | export { init }; |
| 18 | 20 | ||
| @@ -22,6 +24,9 @@ const MODULE = 'keepalive'; | |||
| 22 | // still processed by the backend (which is what refreshes the server-side | 24 | // 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. |
| 24 | const KEEPALIVE_RESPONSE_LENGTH = 1; | 26 | const KEEPALIVE_RESPONSE_LENGTH = 1; |
| 27 | const HEARTBEAT_INTERVAL_MS = 20000; | ||
| 28 | const RETRY_INTERVAL_MS = 5000; | ||
| 29 | const SETTINGS_REFRESH_DELAY_MS = 300; | ||
| 25 | 30 | ||
| 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. |
| 27 | const KEEPALIVE_INJECT_ID = 'keepalive_ping'; | 32 | const 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. |
| 38 | const EXPIRED_GRACE_MS = 3000; | 43 | const EXPIRED_GRACE_MS = 3000; |
| 39 | 44 | ||
| 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. | ||
| 43 | const MAX_TIMER_DRIFT_MS = 5000; | ||
| 44 | |||
| 45 | const defaultSettings = { | 45 | const 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 | }; |
| 52 | 52 | ||
| 53 | let idleTimer = null; | 53 | let heartbeatTimer = null; |
| 54 | let keepaliveActive = false; | 54 | let jobRegistered = false; |
| 55 | let refreshPromise = null; | ||
| 56 | let refreshRequested = false; | ||
| 57 | let refreshTimer = null; | ||
| 58 | let preparingJob = false; | ||
| 59 | let lifecycleVersion = 0; | ||
| 55 | // Whether keepalive is "armed" for the CURRENT chat. It stays dormant until the user sends a new | 60 | // 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 already | 61 | // 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 would | 62 | // 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. |
| 59 | let armed = false; | 64 | let 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). | ||
| 62 | let 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. |
| 65 | const lastActivity = {}; | 67 | const lastActivity = {}; |
| 66 | 68 | ||
| 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. | ||
| 71 | const 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 | } |
| 92 | 98 | ||
| 93 | /** | 99 | function getTimeoutMs() { |
| 94 | * Clears the pending idle timer, if any. | 100 | return Math.max(60, Number(extension_settings[MODULE].timeoutSeconds) || 295) * 1000; |
| 95 | */ | ||
| 96 | function clearIdleTimer() { | ||
| 97 | if (idleTimer) { | ||
| 98 | clearTimeout(idleTimer); | ||
| 99 | idleTimer = null; | ||
| 100 | } | 101 | } |
| 102 | |||
| 103 | async 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 | } |
| 102 | 110 | ||
| 103 | /** | 111 | async function stopBackendJob() { |
| 104 | * (Re)schedules the keepalive for the ACTIVE profile based on how long it has been since | 112 | lifecycleVersion++; |
| 105 | * that profile's last message. Does nothing (and leaves the timer cleared) until keepalive | 113 | 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 own | 115 | await postKeepaliveBackend('stop', {}); |
| 108 | * independent idle clock. | 116 | } catch (error) { |
| 109 | */ | 117 | console.debug('[Keepalive] Could not stop backend job:', error); |
| 110 | function scheduleIdleTimer() { | 118 | } |
| 111 | clearIdleTimer(); | 119 | } |
| 112 | if (!armed || !isKeepaliveEnabledForActive()) { | 120 | |
| 113 | scheduledFireTime = null; | 121 | async 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 | } |
| 123 | 135 | ||
| 124 | /** | ||
| 125 | * Marks the active profile as just-used (resets its idle clock) and reschedules its keepalive. | ||
| 126 | */ | ||
| 127 | function markActivity() { | 136 | function 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 | } |
| 131 | 143 | ||
| 132 | /** | 144 | /** |
| @@ -146,76 +158,162 @@ function onMessageSent() { | |||
| 146 | */ | 158 | */ |
| 147 | function onChatChanged() { | 159 | function onChatChanged() { |
| 148 | armed = false; | 160 | armed = false; |
| 149 | scheduleIdleTimer(); | 161 | void stopBackendJob(); |
| 150 | } | 162 | } |
| 151 | 163 | ||
| 152 | /** | 164 | /** |
| 153 | * Fires a background ("quiet") generation against the active model to keep the | 165 | * 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 the | 167 | * @param {object} payload Prepared generation payload. |
| 156 | * completion is capped to a single token (KEEPALIVE_RESPONSE_LENGTH) so we never | 168 | * @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 | */ |
| 171 | async function fireKeepalive() { | 170 | function 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 | */ | ||
| 193 | async 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 | |||
| 251 | async function registerBackendJob() { | ||
| 172 | if (!armed || !isKeepaliveEnabledForActive()) { | 252 | if (!armed || !isKeepaliveEnabledForActive()) { |
| 253 | await stopBackendJob(); | ||
| 173 | return; | 254 | return; |
| 174 | } | 255 | } |
| 175 | 256 | 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 has | 258 | if (isGenerating()) { |
| 178 | // almost certainly expired. Don't ping on a stale assumption — go dormant until the next | 259 | 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 | } | ||
| 187 | 288 | ||
| 188 | // Never start while any generation is in progress, or while a previous | 289 | function 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 | } | ||
| 194 | 303 | ||
| 195 | keepaliveActive = true; | 304 | function scheduleJobRefresh() { |
| 196 | setSendButtonState(true); | 305 | clearTimeout(refreshTimer); |
| 197 | deactivateSendButtons(); | 306 | refreshTimer = setTimeout(queueJobRefresh, SETTINGS_REFRESH_DELAY_MS); |
| 307 | } | ||
| 198 | 308 | ||
| 199 | try { | 309 | function setupHeartbeatLoop() { |
| 200 | // Deliver the keepalive as a USER message at depth 0 (the quiet-prompt path forces a | 310 | clearInterval(heartbeatTimer); |
| 201 | // system role — openai.js hardcodes role:'system' for it), using the same depth+role | 311 | heartbeatTimer = setInterval(() => void sendBackendHeartbeat(), HEARTBEAT_INTERVAL_MS); |
| 202 | // injection Author's Note uses. The quiet prompt itself stays empty so the full chat | 312 | 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 | } |
| 220 | 318 | ||
| 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 | }); |
| 298 | 400 | ||
| 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 | }); |
| 304 | 408 | ||
| 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 | }); |
| 309 | 416 | ||
| 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 | } |
| 324 | 435 | ||
| 436 | function onGenerationStarted(_type, _options, dryRun) { | ||
| 437 | if (!dryRun) { | ||
| 438 | markActivity(); | ||
| 439 | } | ||
| 440 | } | ||
| 441 | |||
| 442 | function onGenerationEnded() { | ||
| 443 | markActivity(); | ||
| 444 | if (armed && isKeepaliveEnabledForActive()) { | ||
| 445 | queueJobRefresh(); | ||
| 446 | } | ||
| 447 | } | ||
| 448 | |||
| 449 | function onChatContentChanged() { | ||
| 450 | markActivity(); | ||
| 451 | if (armed && isKeepaliveEnabledForActive()) { | ||
| 452 | queueJobRefresh(); | ||
| 453 | } | ||
| 454 | } | ||
| 455 | |||
| 456 | async 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 | |||
| 325 | async function init() { | 464 | async 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); |
| 328 | 467 | ||
| 329 | loadSettings(); | 468 | loadSettings(); |
| 330 | setupListeners(); | 469 | setupListeners(); |
| 470 | setupHeartbeatLoop(); | ||
| 331 | 471 | ||
| 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); |
| 334 | 474 | ||
| 335 | // These reset the active profile's idle clock (so generation time isn't counted as idle) but | 475 | // 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 while | 476 | // 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 | } | ||
| 346 | 485 | ||
| 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); |
| 349 | 488 | ||
| 350 | // When the active profile changes, re-target the idle timer at the new profile using | 489 | // 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); | ||
| 353 | 495 | ||
| 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 | } |
| @@ -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" |
| @@ -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" /> |
| 18 | 18 | ||
| 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> |
| 20 | 20 | ||
| 21 | <hr> | 21 | <hr> |
| 22 | 22 | ||
| @@ -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). */ |
| 1542 | const RUNPOD_POLL_IDLE_MS = 30000; | 1542 | const RUNPOD_POLL_IDLE_MS = 30000; |
| 1543 | const RUNPOD_POLL_BUSY_MS = 5000; | 1543 | const 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. */ |
| 1545 | const RUNPOD_PING_MS = 60000; | 1545 | const RUNPOD_HEARTBEAT_MS = 20000; |
| 1546 | 1546 | ||
| 1547 | let runpodStatusTimer = null; | 1547 | let runpodStatusTimer = null; |
| 1548 | let runpodPingTimer = null; | 1548 | let runpodHeartbeatTimer = null; |
| 1549 | let runpodLastPhase = 'red'; | 1549 | let runpodLastPhase = 'red'; |
| 1550 | 1550 | ||
| 1551 | function getRunpodLazyUrl() { | 1551 | function 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 | } |
| 1692 | 1692 | ||
| 1693 | function runpodKeepalive() { | 1693 | function 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 unnecessary | 1698 | // 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 | } |
| 1702 | 1703 | ||
| 1703 | /** (Re)starts the polling/keepalive loops based on the configured URL. */ | 1704 | /** (Re)starts the status polling and frontend-heartbeat loops. */ |
| 1704 | function setupRunpodLoops() { | 1705 | function 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 | } |
| 1719 | 1721 | ||
| 1720 | function onRunpodLazyUrlInput() { | 1722 | function onRunpodLazyUrlInput() { |
| @@ -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 "%model%" resolves to; downloads = one "subpath url" 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 "%model%" resolves to; downloads = one "subpath url" 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> |
| @@ -0,0 +1,264 @@ | |||
| 1 | import express from 'express'; | ||
| 2 | import https from 'node:https'; | ||
| 3 | import fetch from 'node-fetch'; | ||
| 4 | |||
| 5 | export const router = express.Router(); | ||
| 6 | |||
| 7 | export const FRONTEND_LEASE_MS = 5 * 60 * 1000; | ||
| 8 | const MIN_INTERVAL_MS = 60 * 1000; | ||
| 9 | const MAX_INTERVAL_MS = 24 * 60 * 60 * 1000; | ||
| 10 | const MAX_JOB_ID_LENGTH = 128; | ||
| 11 | const 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 | ]); | ||
| 18 | const localHttpsAgent = new https.Agent({ rejectUnauthorized: false }); | ||
| 19 | |||
| 20 | function 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 | */ | ||
| 31 | export 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 | |||
| 150 | async 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 | |||
| 171 | const scheduler = new KeepaliveScheduler(); | ||
| 172 | |||
| 173 | function getJobKey(request, id) { | ||
| 174 | return `${request.user.profile.handle}:${id}`; | ||
| 175 | } | ||
| 176 | |||
| 177 | function 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 | |||
| 185 | function 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 | |||
| 219 | router.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 | |||
| 243 | router.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 | |||
| 256 | router.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 | }); | ||
| @@ -51,6 +51,7 @@ import { router as dataMaidRouter } from './endpoints/data-maid.js'; | |||
| 51 | import { router as backupsRouter } from './endpoints/backups.js'; | 51 | import { router as backupsRouter } from './endpoints/backups.js'; |
| 52 | import { router as imageMetadataRouter } from './endpoints/image-metadata.js'; | 52 | import { router as imageMetadataRouter } from './endpoints/image-metadata.js'; |
| 53 | import { router as volcengineRouter } from './endpoints/volcengine.js'; | 53 | import { router as volcengineRouter } from './endpoints/volcengine.js'; |
| 54 | import { router as keepaliveRouter } from './endpoints/keepalive.js'; | ||
| 54 | 55 | ||
| 55 | /** | 56 | /** |
| 56 | * @typedef {object} ServerStartupResult | 57 | * @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 | } |
| 189 | 191 | ||
| 190 | /** | 192 | /** |
| @@ -0,0 +1,80 @@ | |||
| 1 | import { describe, expect, jest, test } from '@jest/globals'; | ||
| 2 | import { FRONTEND_LEASE_MS, KeepaliveScheduler } from '../src/endpoints/keepalive.js'; | ||
| 3 | |||
| 4 | function 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 | |||
| 16 | describe('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 | }); | ||