1
import { extension_settings, renderExtensionTemplateAsync } from '../../extensions.js';
2
import {
3
    activateSendButtons,
4
    deactivateSendButtons,
5
    eventSource,
6
    event_types,
7
    extension_prompt_roles,
8
    extension_prompt_types,
9
    generateQuietPrompt,
10
    isGenerating,
11
    saveSettingsDebounced,
12
    setExtensionPrompt,
13
    setSendButtonState,
14
    streamingProcessor,
15
} from '../../../script.js';
16

17
export { init };
18

19
const MODULE = 'keepalive';
20

21
// Cap the keepalive completion to a single token. The full prompt prefix is
22
// still processed by the backend (which is what refreshes the server-side
23
// cache), but we never pay for a full completion.
24
const KEEPALIVE_RESPONSE_LENGTH = 1;
25

26
// Extension-prompt key used to inject the keepalive message as a USER message at depth 0.
27
const KEEPALIVE_INJECT_ID = 'keepalive_ping';
28

29
// Original default message; used to migrate untouched installs to the current default wording.
30
const LEGACY_DEFAULT_MESSAGE = 'Keep the cache warm.';
31

32
// Idle-tracking key used when no connection profile is active (manual API panel).
33
const NONE_KEY = '__none__';
34

35
// When the active profile is already past its timeout (e.g. you just switched back
36
// to a profile that went cold), wait this short grace before pinging so an imminent
37
// user message can pre-empt the keepalive.
38
const EXPIRED_GRACE_MS = 3000;
39

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 = {
46
    enabled: false,
47
    timeoutSeconds: 295,
48
    message: '[Note: just reply with an empty response to keep the session alive]',
49
    // Per-profile enable map: { [profileId]: boolean }. Unlisted profiles default to enabled.
50
    profiles: {},
51
};
52

53
let idleTimer = null;
54
let keepaliveActive = false;
55
// 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
57
// 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.
59
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).
64
// In-memory only: on reload every profile is treated as freshly active.
65
const lastActivity = {};
66

67
/**
68
 * @returns {string} The active connection profile id, or NONE_KEY for a manual connection.
69
 */
70
function getActiveProfileKey() {
71
    return extension_settings.connectionManager?.selectedProfile || NONE_KEY;
72
}
73

74
/**
75
 * Whether keepalive should run for the CURRENTLY ACTIVE connection. Requires the global
76
 * master switch; for a real profile it also requires that profile's per-profile toggle
77
 * (default on). A manual connection (no profile selected) is governed by the master switch alone.
78
 * @returns {boolean}
79
 */
80
function isKeepaliveEnabledForActive() {
81
    const s = extension_settings[MODULE];
82
    if (!s.enabled) {
83
        return false;
84
    }
85
    const profileId = extension_settings.connectionManager?.selectedProfile;
86
    if (!profileId) {
87
        return true;
88
    }
89
    const perProfile = s.profiles?.[profileId];
90
    return perProfile === undefined ? true : !!perProfile;
91
}
92

93
/**
94
 * Clears the pending idle timer, if any.
95
 */
96
function clearIdleTimer() {
97
    if (idleTimer) {
98
        clearTimeout(idleTimer);
99
        idleTimer = null;
100
    }
101
}
102

103
/**
104
 * (Re)schedules the keepalive for the ACTIVE profile based on how long it has been since
105
 * that profile's last message. Does nothing (and leaves the timer cleared) until keepalive
106
 * has been armed for the current chat, or when it is disabled for the active connection.
107
 * Called on activity AND whenever the active profile changes, so each profile keeps its own
108
 * independent idle clock.
109
 */
110
function scheduleIdleTimer() {
111
    clearIdleTimer();
112
    if (!armed || !isKeepaliveEnabledForActive()) {
113
        scheduledFireTime = null;
114
        return;
115
    }
116
    const timeoutMs = Math.max(60, Number(extension_settings[MODULE].timeoutSeconds) || 295) * 1000;
117
    const last = lastActivity[getActiveProfileKey()] ?? Date.now();
118
    const remaining = timeoutMs - (Date.now() - last);
119
    const delay = remaining > 0 ? remaining : EXPIRED_GRACE_MS;
120
    scheduledFireTime = Date.now() + delay;
121
    idleTimer = setTimeout(fireKeepalive, delay);
122
}
123

124
/**
125
 * Marks the active profile as just-used (resets its idle clock) and reschedules its keepalive.
126
 */
127
function markActivity() {
128
    lastActivity[getActiveProfileKey()] = Date.now();
129
    scheduleIdleTimer();
130
}
131

132
/**
133
 * The user sent a new message in the current chat — the ONLY thing that arms keepalive.
134
 * Background/quiet generations, swipes, regenerations, and loading chat history deliberately do
135
 * NOT arm it: none of those prove the user is actively continuing THIS chat from a known-cached
136
 * prompt (GENERATION_STARTED in particular fires for every quiet/background generation too).
137
 */
138
function onMessageSent() {
139
    armed = true;
140
    markActivity();
141
}
142

143
/**
144
 * Opening or switching chats disarms keepalive: a freshly loaded chat is not proof its prompt is
145
 * still cached server-side, so we wait for a new sent message in that chat before running again.
146
 */
147
function onChatChanged() {
148
    armed = false;
149
    scheduleIdleTimer();
150
}
151

152
/**
153
 * Fires a background ("quiet") generation against the active model to keep the
154
 * server-side prompt cache warm. The full prompt prefix (current chat context +
155
 * the keepalive message) is sent — which is what refreshes the cache — while the
156
 * completion is capped to a single token (KEEPALIVE_RESPONSE_LENGTH) so we never
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
 */
171
async function fireKeepalive() {
172
    if (!armed || !isKeepaliveEnabledForActive()) {
173
        return;
174
    }
175

176
    // If the timer fired much later than scheduled (machine sleep, background-tab throttling,
177
    // etc.), far more wall-clock time has elapsed than the idle interval, so the cache has
178
    // almost certainly expired. Don't ping on a stale assumption — go dormant until the next
179
    // sent message re-arms keepalive, exactly like a freshly opened chat.
180
    if (scheduledFireTime !== null && (Date.now() - scheduledFireTime) > MAX_TIMER_DRIFT_MS) {
181
        console.debug('[Keepalive] Idle timer fired', Date.now() - scheduledFireTime, 'ms late; skipping and disarming until next activity.');
182
        armed = false;
183
        clearIdleTimer();
184
        scheduledFireTime = null;
185
        return;
186
    }
187

188
    // Never start while any generation is in progress, or while a previous
189
    // keepalive is still running. Reschedule and bail.
190
    if (keepaliveActive || isGenerating() || streamingProcessor) {
191
        scheduleIdleTimer();
192
        return;
193
    }
194

195
    keepaliveActive = true;
196
    setSendButtonState(true);
197
    deactivateSendButtons();
198

199
    try {
200
        // Deliver the keepalive as a USER message at depth 0 (the quiet-prompt path forces a
201
        // system role — openai.js hardcodes role:'system' for it), using the same depth+role
202
        // injection Author's Note uses. The quiet prompt itself stays empty so the full chat
203
        // context is still assembled (which is what keeps the cached prefix warm).
204
        setExtensionPrompt(KEEPALIVE_INJECT_ID, extension_settings[MODULE].message, extension_prompt_types.IN_CHAT, 0, false, extension_prompt_roles.USER);
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
    }
219
}
220

221
/**
222
 * Renders the per-profile enable checkboxes from the current Connection Manager profiles.
223
 */
224
function renderProfileToggles() {
225
    const container = $('#keepalive_profiles_list');
226
    if (!container.length) {
227
        return;
228
    }
229
    container.empty();
230

231
    const profiles = extension_settings.connectionManager?.profiles;
232
    if (!Array.isArray(profiles) || profiles.length === 0) {
233
        container.append(
234
            $('<small>')
235
                .attr('data-i18n', 'ext_keepalive_no_profiles')
236
                .text('No connection profiles found. Keepalive uses the active model, governed by the master toggle above.'),
237
        );
238
        return;
239
    }
240

241
    const perProfile = extension_settings[MODULE].profiles || {};
242
    const sorted = profiles.slice().sort((a, b) => String(a.name).localeCompare(String(b.name)));
243
    for (const profile of sorted) {
244
        const enabled = perProfile[profile.id] === undefined ? true : !!perProfile[profile.id];
245
        const label = $('<label>').addClass('checkbox_label');
246
        const input = $('<input>')
247
            .attr('type', 'checkbox')
248
            .addClass('keepalive_profile_toggle')
249
            .attr('data-profile-id', profile.id)
250
            .prop('checked', enabled);
251
        const span = $('<span>').text(profile.name);
252
        label.append(input, span);
253
        container.append(label);
254
    }
255
}
256

257
/**
258
 * Loads settings into the global store and hydrates the UI controls.
259
 */
260
function loadSettings() {
261
    if (extension_settings[MODULE] === undefined) {
262
        extension_settings[MODULE] = {};
263
    }
264

265
    for (const key of Object.keys(defaultSettings)) {
266
        if (extension_settings[MODULE][key] === undefined) {
267
            extension_settings[MODULE][key] = structuredClone(defaultSettings[key]);
268
        }
269
    }
270

271
    // Guard against a corrupted/legacy value for the per-profile map.
272
    if (typeof extension_settings[MODULE].profiles !== 'object' || extension_settings[MODULE].profiles === null) {
273
        extension_settings[MODULE].profiles = {};
274
    }
275

276
    // Migrate the original default message to the current default so installs that never
277
    // customized it pick up the improved wording (custom messages are left untouched).
278
    if (extension_settings[MODULE].message === LEGACY_DEFAULT_MESSAGE) {
279
        extension_settings[MODULE].message = defaultSettings.message;
280
        saveSettingsDebounced();
281
    }
282

283
    $('#keepalive_enabled').prop('checked', extension_settings[MODULE].enabled);
284
    $('#keepalive_timeout').val(extension_settings[MODULE].timeoutSeconds);
285
    $('#keepalive_message').val(extension_settings[MODULE].message);
286
    renderProfileToggles();
287
}
288

289
/**
290
 * Wires up the settings UI control handlers.
291
 */
292
function setupListeners() {
293
    $('#keepalive_enabled').on('change', function () {
294
        extension_settings[MODULE].enabled = !!$(this).prop('checked');
295
        saveSettingsDebounced();
296
        scheduleIdleTimer();
297
    });
298

299
    $('#keepalive_timeout').on('input', function () {
300
        extension_settings[MODULE].timeoutSeconds = Number($(this).val()) || 295;
301
        saveSettingsDebounced();
302
        scheduleIdleTimer();
303
    });
304

305
    $('#keepalive_message').on('input', function () {
306
        extension_settings[MODULE].message = String($(this).val());
307
        saveSettingsDebounced();
308
    });
309

310
    // Per-profile toggles are rendered dynamically; use a delegated handler.
311
    $('#keepalive_profiles_list').on('change', '.keepalive_profile_toggle', function () {
312
        const profileId = $(this).attr('data-profile-id');
313
        if (!profileId) {
314
            return;
315
        }
316
        extension_settings[MODULE].profiles[profileId] = !!$(this).prop('checked');
317
        saveSettingsDebounced();
318
        // If the toggled profile is the one currently active, reschedule immediately.
319
        if (profileId === extension_settings.connectionManager?.selectedProfile) {
320
            scheduleIdleTimer();
321
        }
322
    });
323
}
324

325
async function init() {
326
    const settingsHtml = await renderExtensionTemplateAsync(MODULE, 'settings');
327
    $('#extensions_settings2').append(settingsHtml);
328

329
    loadSettings();
330
    setupListeners();
331

332
    // Only the user sending a new message arms keepalive (for the current chat).
333
    eventSource.on(event_types.MESSAGE_SENT, onMessageSent);
334

335
    // These reset the active profile's idle clock (so generation time isn't counted as idle) but
336
    // never arm keepalive on their own — they also fire for background/quiet generations and while
337
    // a chat is loading. They only re-arm the timer if a sent message already armed it.
338
    const activityEvents = [
339
        event_types.MESSAGE_RECEIVED,
340
        event_types.GENERATION_STARTED,
341
        event_types.GENERATION_ENDED,
342
    ];
343
    for (const event of activityEvents) {
344
        eventSource.on(event, markActivity);
345
    }
346

347
    // Opening/switching a chat disarms keepalive — it must be re-armed by a new sent message.
348
    eventSource.on(event_types.CHAT_CHANGED, onChatChanged);
349

350
    // When the active profile changes, re-target the idle timer at the new profile using
351
    // ITS own last-activity time (do NOT mark activity — switching is not using it).
352
    eventSource.on(event_types.CONNECTION_PROFILE_LOADED, scheduleIdleTimer);
353

354
    // Keep the per-profile toggle list in sync with Connection Manager.
355
    eventSource.on(event_types.CONNECTION_PROFILE_CREATED, renderProfileToggles);
356
    eventSource.on(event_types.CONNECTION_PROFILE_UPDATED, renderProfileToggles);
357
    eventSource.on(event_types.CONNECTION_PROFILE_DELETED, renderProfileToggles);
358

359
    scheduleIdleTimer();
360
}