Isolate text keepalive from image generation

ad7958ed1cbe509a6ebeac4e41fb206b52a46942

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

Parents: 0aa62e9004c1517fe3801ca9fef82ecf60d6bc29
4 files changed, +52 -8Showing whitespace changes
public/scripts/events.js+2 -0
@@ -79,6 +79,8 @@ export const event_types = {
7979 ONLINE_STATUS_CHANGED: 'online_status_changed',
8080 IMAGE_SWIPED: 'image_swiped',
8181 CONNECTION_PROFILE_LOADED: 'connection_profile_loaded',
82+ CONNECTION_PROFILE_TEMPORARY_STARTED: 'connection_profile_temporary_started',
83+ CONNECTION_PROFILE_TEMPORARY_ENDED: 'connection_profile_temporary_ended',
8284 CONNECTION_PROFILE_CREATED: 'connection_profile_created',
8385 CONNECTION_PROFILE_DELETED: 'connection_profile_deleted',
8486 CONNECTION_PROFILE_UPDATED: 'connection_profile_updated',
public/scripts/extensions/keepalive/index.js+43 -6
@@ -54,6 +54,9 @@ let generationRefreshTimer = null;
5454let retryTimer = null;
5555let preparingJob = false;
5656let lifecycleVersion = 0;
57+let foregroundGenerationActive = false;
58+let temporaryConnectionDepth = 0;
59+let refreshAfterTemporaryConnection = false;
5760// Whether keepalive is "armed" for the CURRENT chat. It stays dormant until the user sends a new
5861// message or starts another foreground generation in the open chat. Merely loading a chat does not
5962// prove that its prompt is cached server-side. Reset on page reload and whenever the chat changes.
@@ -145,7 +148,7 @@ async function stopBackendJob() {
145148}
146149
147150async function sendBackendHeartbeat(activity = false) {
148151 if (temporaryConnectionDepth > 0 || !armed || !isKeepaliveEnabledForActive()) {
149152 return;
150153 }
151154 // Registration is normally created as soon as a response finishes. If that one-shot event
@@ -182,7 +185,7 @@ async function sendBackendHeartbeat(activity = false) {
182185}
183186
184187function markActivity() {
185188 if (preparingJob || temporaryConnectionDepth > 0) {
186189 return;
187190 }
188191 lastActivity[getActiveProfileKey()] = Date.now();
@@ -296,6 +299,10 @@ async function prepareKeepaliveRequest() {
296299}
297300
298301async function registerBackendJob() {
302+ if (temporaryConnectionDepth > 0) {
303+ refreshAfterTemporaryConnection = true;
304+ return;
305+ }
299306 if (!armed || !isKeepaliveEnabledForActive()) {
300307 await stopBackendJob();
301308 return;
@@ -500,23 +507,34 @@ function setupListeners() {
500507}
501508
502509function onGenerationStarted(type, _options, dryRun) {
503- if (dryRun) {
510+ // Quiet/background generations (image prompts, summaries, helpers, etc.) use a different model
511+ // or context and do not refresh the foreground text-chat cache.
512+ if (dryRun || type === 'quiet') {
504513 return;
505514 }
506515 // Regenerates, swipes, and continues send the current prompt to the provider without emitting
507516 // MESSAGE_SENT. They are just as valid proof of a warm prompt as a newly typed message.
508- if (type !== 'quiet') {
517+ foregroundGenerationActive = true;
509518 armed = true;
510- }
511519 markActivity();
512520}
513521
514522function onGenerationEnded() {
523+ if (!foregroundGenerationActive) {
524+ return;
525+ }
526+ foregroundGenerationActive = false;
515527 markActivity();
516528 scheduleGenerationJobRefresh();
517529}
518530
519531function onMessageReceived(_messageId, type) {
532+ // Image-generation extensions add their results as chat messages. Those messages neither used
533+ // nor refreshed the foreground text cache, so they must not move its keepalive deadline.
534+ if (type === 'extension') {
535+ return;
536+ }
537+ foregroundGenerationActive = false;
520538 markActivity();
521539 scheduleGenerationJobRefresh();
522540}
@@ -529,6 +547,9 @@ function onChatContentChanged() {
529547}
530548
531549async function onConnectionChanged() {
550+ if (temporaryConnectionDepth > 0) {
551+ return;
552+ }
532553 // Remove the previous connection's request before assembling one with the new settings.
533554 await stopBackendJob();
534555 if (armed && isKeepaliveEnabledForActive()) {
@@ -536,6 +557,18 @@ async function onConnectionChanged() {
536557 }
537558}
538559
560+function onTemporaryConnectionStarted() {
561+ temporaryConnectionDepth++;
562+}
563+
564+function onTemporaryConnectionEnded() {
565+ temporaryConnectionDepth = Math.max(0, temporaryConnectionDepth - 1);
566+ if (temporaryConnectionDepth === 0 && refreshAfterTemporaryConnection) {
567+ refreshAfterTemporaryConnection = false;
568+ queueJobRefresh();
569+ }
570+}
571+
539572async function init() {
540573 const settingsHtml = await renderExtensionTemplateAsync(MODULE, 'settings');
541574 $('#extensions_settings2').append(settingsHtml);
@@ -554,6 +587,7 @@ async function init() {
554587 eventSource.on(event_types.MESSAGE_RECEIVED, onMessageReceived);
555588 eventSource.on(event_types.GENERATION_STARTED, onGenerationStarted);
556589 eventSource.on(event_types.GENERATION_ENDED, onGenerationEnded);
590+ eventSource.on(event_types.GENERATION_STOPPED, () => { foregroundGenerationActive = false; });
557591 eventSource.on(event_types.MESSAGE_EDITED, onChatContentChanged);
558592 eventSource.on(event_types.MESSAGE_DELETED, onChatContentChanged);
559593 eventSource.on(event_types.MESSAGE_UPDATED, onChatContentChanged);
@@ -563,6 +597,9 @@ async function init() {
563597 eventSource.on(event_types.CHAT_CHANGED, onChatChanged);
564598
565599 // When the active profile changes, replace the backend job without marking it as activity.
600+ // Temporary profile switches used by background image-prompt generation are explicitly ignored.
601+ eventSource.on(event_types.CONNECTION_PROFILE_TEMPORARY_STARTED, onTemporaryConnectionStarted);
602+ eventSource.on(event_types.CONNECTION_PROFILE_TEMPORARY_ENDED, onTemporaryConnectionEnded);
566603 eventSource.on(event_types.CONNECTION_PROFILE_LOADED, onConnectionChanged);
567604 eventSource.on(event_types.MAIN_API_CHANGED, onConnectionChanged);
568605 eventSource.on(event_types.CHATCOMPLETION_SOURCE_CHANGED, onConnectionChanged);
public/scripts/extensions/keepalive/manifest.json+1 -1
@@ -6,7 +6,7 @@
66 "js": "index.js",
77 "css": "",
88 "author": "SillyTavern",
99 "version": "1.1.34",
1010 "homePage": "https://github.com/SillyTavern/SillyTavern",
1111 "hooks": {
1212 "activate": "init"
public/scripts/extensions/stable-diffusion/index.js+6 -1
@@ -4714,14 +4714,19 @@ async function withConnectionProfile(targetProfileId, callback) {
47144714 await waitUntilCondition(() => online_status !== 'no_connection', 10000, 100, { rejectOnTimeout: false });
47154715 };
47164716
4717- await switchToProfile(targetProfileId);
4717+ // Background image-prompt work must not disturb foreground features that follow the user's
4718+ // active connection profile (notably text prompt-cache keepalives).
4719+ await eventSource.emit(event_types.CONNECTION_PROFILE_TEMPORARY_STARTED);
47184720 try {
4721+ await switchToProfile(targetProfileId);
47194722 return await callback();
47204723 } finally {
47214724 try {
47224725 await switchToProfile(currentProfileId);
47234726 } catch (err) {
47244727 console.error('SD: failed to restore the previous connection profile after image-prompt generation', err);
4728+ } finally {
4729+ await eventSource.emit(event_types.CONNECTION_PROFILE_TEMPORARY_ENDED);
47254730 }
47264731 }
47274732}