From 0b05bc6e800a135a947bfd5b4d74322fe37361d5 Mon Sep 17 00:00:00 2001 From: Yannic Charlon <52761674+JustYannicc@users.noreply.github.com> Date: Fri, 3 Jul 2026 23:35:16 +0200 Subject: [PATCH] perf(settings): clean up Ollama pull listeners --- scripts/test-ollama-pull-listener-cleanup.mjs | 278 ++++++++++++++++++ src/main/preload.ts | 12 +- src/renderer/src/settings/AITab.tsx | 34 ++- .../src/settings/ollamaPullProgress.ts | 110 +++++++ src/renderer/types/electron.d.ts | 6 +- 5 files changed, 418 insertions(+), 22 deletions(-) create mode 100644 scripts/test-ollama-pull-listener-cleanup.mjs create mode 100644 src/renderer/src/settings/ollamaPullProgress.ts diff --git a/scripts/test-ollama-pull-listener-cleanup.mjs b/scripts/test-ollama-pull-listener-cleanup.mjs new file mode 100644 index 00000000..bba9a4a6 --- /dev/null +++ b/scripts/test-ollama-pull-listener-cleanup.mjs @@ -0,0 +1,278 @@ +#!/usr/bin/env node + +import test from 'node:test'; +import assert from 'node:assert/strict'; +import { EventEmitter } from 'node:events'; +import path from 'node:path'; +import { importTs } from './lib/ts-import.mjs'; + +const REMOUNT_CYCLES = 5; +const CURRENT_REQUEST_ID = 'ollama-pull-current'; +const OTHER_REQUEST_ID = 'ollama-pull-stale'; +const PREFERRED_MODEL = 'llama3.2'; + +const { + registerOllamaPullListeners, + toOllamaPullProgressState, +} = await importTs(path.resolve('src/renderer/src/settings/ollamaPullProgress.ts')); + +function createBridge({ returnCleanup }) { + const emitter = new EventEmitter(); + + function subscribe(channel, callback) { + const listener = (data) => callback(data); + emitter.on(channel, listener); + if (!returnCleanup) return undefined; + return () => emitter.removeListener(channel, listener); + } + + return { + onOllamaPullProgress: (callback) => subscribe('ollama-pull-progress', callback), + onOllamaPullDone: (callback) => subscribe('ollama-pull-done', callback), + onOllamaPullError: (callback) => subscribe('ollama-pull-error', callback), + emitProgress: (data) => emitter.emit('ollama-pull-progress', data), + emitDone: (data) => emitter.emit('ollama-pull-done', data), + emitError: (data) => emitter.emit('ollama-pull-error', data), + listenerCounts: () => ({ + progress: emitter.listenerCount('ollama-pull-progress'), + done: emitter.listenerCount('ollama-pull-done'), + error: emitter.listenerCount('ollama-pull-error'), + }), + }; +} + +function createStateRecorder() { + let activeRequestId = CURRENT_REQUEST_ID; + let preferredModel = PREFERRED_MODEL; + const metrics = { + progressWrites: 0, + progressResets: 0, + doneRefreshes: 0, + errorMessages: 0, + errorClearsScheduled: 0, + clearActivePulls: 0, + pullingModelClears: 0, + refreshedModels: [], + errors: [], + }; + + return { + get activeRequestId() { + return activeRequestId; + }, + get metrics() { + return metrics; + }, + resetActivePull() { + activeRequestId = CURRENT_REQUEST_ID; + preferredModel = PREFERRED_MODEL; + }, + optionsForFixed(bridge) { + return { + bridge, + getActiveRequestId: () => activeRequestId, + getPreferredModel: () => preferredModel, + clearActivePull: () => { + activeRequestId = null; + preferredModel = undefined; + metrics.clearActivePulls += 1; + }, + setPullingModel: (modelName) => { + if (modelName === null) metrics.pullingModelClears += 1; + }, + setPullProgress: (progress) => { + if (progress.status === '' && progress.percent === 0) { + metrics.progressResets += 1; + } else { + metrics.progressWrites += 1; + } + }, + setOllamaError: (error) => { + if (error) { + metrics.errorMessages += 1; + metrics.errors.push(error); + } + }, + scheduleErrorClear: () => { + metrics.errorClearsScheduled += 1; + }, + refreshOllamaStatus: (modelName) => { + metrics.doneRefreshes += 1; + metrics.refreshedModels.push(modelName); + }, + }; + }, + optionsForLegacy(bridge) { + return { + bridge, + getPreferredModel: () => preferredModel, + clearActivePull: () => { + activeRequestId = null; + preferredModel = undefined; + metrics.clearActivePulls += 1; + }, + setPullingModel: (modelName) => { + if (modelName === null) metrics.pullingModelClears += 1; + }, + setPullProgress: (progress) => { + if (progress.status === '' && progress.percent === 0) { + metrics.progressResets += 1; + } else { + metrics.progressWrites += 1; + } + }, + setOllamaError: (error) => { + if (error) { + metrics.errorMessages += 1; + metrics.errors.push(error); + } + }, + scheduleErrorClear: () => { + metrics.errorClearsScheduled += 1; + }, + refreshOllamaStatus: (modelName) => { + metrics.doneRefreshes += 1; + metrics.refreshedModels.push(modelName); + }, + }; + }, + }; +} + +function registerLegacyAITabPullListeners({ + bridge, + clearActivePull, + getPreferredModel, + refreshOllamaStatus, + scheduleErrorClear, + setOllamaError, + setPullingModel, + setPullProgress, +}) { + bridge.onOllamaPullProgress((data) => { + setPullProgress(toOllamaPullProgressState(data)); + }); + bridge.onOllamaPullDone(() => { + const preferredModel = getPreferredModel(); + clearActivePull(); + setPullingModel(null); + setPullProgress({ status: '', percent: 0 }); + refreshOllamaStatus(preferredModel); + }); + bridge.onOllamaPullError((data) => { + clearActivePull(); + setPullingModel(null); + setPullProgress({ status: '', percent: 0 }); + setOllamaError(data.error); + scheduleErrorClear(); + }); + return () => {}; +} + +function mountAndUnmountRepeatedly(register, bridge, state, cycles) { + for (let index = 0; index < cycles; index += 1) { + const cleanup = register(bridge, state); + cleanup(); + } + return register(bridge, state); +} + +function progressPayload(overrides = {}) { + return { + requestId: CURRENT_REQUEST_ID, + status: 'pulling manifest', + digest: 'sha256:test', + total: 100, + completed: 25, + ...overrides, + }; +} + +function runLifecycleScenario({ label, bridge, register }) { + const state = createStateRecorder(); + const cleanupActiveMount = mountAndUnmountRepeatedly(register, bridge, state, REMOUNT_CYCLES); + + bridge.emitProgress(progressPayload()); + state.resetActivePull(); + bridge.emitDone({ requestId: CURRENT_REQUEST_ID }); + state.resetActivePull(); + bridge.emitError({ requestId: CURRENT_REQUEST_ID, error: 'pull failed' }); + cleanupActiveMount(); + + const result = { + label, + remountCycles: REMOUNT_CYCLES, + metrics: state.metrics, + listenerCountsAfterCleanup: bridge.listenerCounts(), + }; + console.log(`[ollama-pull ${label}] ${JSON.stringify(result)}`); + return result; +} + +test('AITab-like remount churn does not multiply Ollama pull listeners after cleanup', () => { + const before = runLifecycleScenario({ + label: 'before-inline-listeners', + bridge: createBridge({ returnCleanup: false }), + register: (bridge, state) => registerLegacyAITabPullListeners(state.optionsForLegacy(bridge)), + }); + const after = runLifecycleScenario({ + label: 'after-cleanup-listeners', + bridge: createBridge({ returnCleanup: true }), + register: (bridge, state) => registerOllamaPullListeners(state.optionsForFixed(bridge)), + }); + + assert.deepEqual({ + progressWrites: before.metrics.progressWrites, + doneRefreshes: before.metrics.doneRefreshes, + errorMessages: before.metrics.errorMessages, + }, { + progressWrites: REMOUNT_CYCLES + 1, + doneRefreshes: REMOUNT_CYCLES + 1, + errorMessages: REMOUNT_CYCLES + 1, + }); + assert.deepEqual(before.listenerCountsAfterCleanup, { + progress: REMOUNT_CYCLES + 1, + done: REMOUNT_CYCLES + 1, + error: REMOUNT_CYCLES + 1, + }); + + assert.deepEqual({ + progressWrites: after.metrics.progressWrites, + doneRefreshes: after.metrics.doneRefreshes, + errorMessages: after.metrics.errorMessages, + }, { + progressWrites: 1, + doneRefreshes: 1, + errorMessages: 1, + }); + assert.deepEqual(after.listenerCountsAfterCleanup, { progress: 0, done: 0, error: 0 }); + assert.deepEqual(after.metrics.refreshedModels, [PREFERRED_MODEL]); +}); + +test('Ollama pull listeners ignore stale request ids and dedupe identical progress', () => { + const bridge = createBridge({ returnCleanup: true }); + const state = createStateRecorder(); + const cleanup = registerOllamaPullListeners(state.optionsForFixed(bridge)); + + bridge.emitProgress(progressPayload({ requestId: OTHER_REQUEST_ID, completed: 75 })); + bridge.emitDone({ requestId: OTHER_REQUEST_ID }); + bridge.emitError({ requestId: OTHER_REQUEST_ID, error: 'stale failure' }); + assert.deepEqual({ + progressWrites: state.metrics.progressWrites, + doneRefreshes: state.metrics.doneRefreshes, + errorMessages: state.metrics.errorMessages, + }, { + progressWrites: 0, + doneRefreshes: 0, + errorMessages: 0, + }); + + bridge.emitProgress(progressPayload({ completed: 25 })); + bridge.emitProgress(progressPayload({ completed: 25 })); + bridge.emitProgress(progressPayload({ completed: 25.1 })); + bridge.emitProgress(progressPayload({ status: 'pulling layers', completed: 25.1 })); + + assert.equal(state.metrics.progressWrites, 2); + cleanup(); + assert.deepEqual(bridge.listenerCounts(), { progress: 0, done: 0, error: 0 }); +}); diff --git a/src/main/preload.ts b/src/main/preload.ts index eff106ef..589b16ef 100644 --- a/src/main/preload.ts +++ b/src/main/preload.ts @@ -1108,13 +1108,19 @@ const electronAPI = { ollamaOpenDownload: (): Promise => ipcRenderer.invoke('ollama-open-download'), onOllamaPullProgress: (callback: (data: { requestId: string; status: string; digest: string; total: number; completed: number }) => void) => { - ipcRenderer.on('ollama-pull-progress', (_event: any, data: any) => callback(data)); + const listener = (_event: any, data: any) => callback(data); + ipcRenderer.on('ollama-pull-progress', listener); + return () => { ipcRenderer.removeListener('ollama-pull-progress', listener); }; }, onOllamaPullDone: (callback: (data: { requestId: string }) => void) => { - ipcRenderer.on('ollama-pull-done', (_event: any, data: any) => callback(data)); + const listener = (_event: any, data: any) => callback(data); + ipcRenderer.on('ollama-pull-done', listener); + return () => { ipcRenderer.removeListener('ollama-pull-done', listener); }; }, onOllamaPullError: (callback: (data: { requestId: string; error: string }) => void) => { - ipcRenderer.on('ollama-pull-error', (_event: any, data: any) => callback(data)); + const listener = (_event: any, data: any) => callback(data); + ipcRenderer.on('ollama-pull-error', listener); + return () => { ipcRenderer.removeListener('ollama-pull-error', listener); }; }, // ─── Hyper Key ────────────────────────────────────────────────── diff --git a/src/renderer/src/settings/AITab.tsx b/src/renderer/src/settings/AITab.tsx index c27befdd..4dd1028f 100644 --- a/src/renderer/src/settings/AITab.tsx +++ b/src/renderer/src/settings/AITab.tsx @@ -38,6 +38,7 @@ import { getCachedElevenLabsVoices, setCachedElevenLabsVoices, } from '../utils/voice-cache'; +import { registerOllamaPullListeners } from './ollamaPullProgress'; const getProviderOptions = (t: (key: string) => string) => [ { id: 'openai' as const, label: t('settings.ai.llm.provider.openai'), description: t('settings.ai.llm.providerDescriptions.openai') }, @@ -282,6 +283,7 @@ const AITab: React.FC = () => { const settingsRef = useRef(null); const pullingModelRef = useRef(null); + const pullRequestIdRef = useRef(null); const selectingOllamaDefaultRef = useRef(false); useEffect(() => { @@ -584,27 +586,27 @@ const AITab: React.FC = () => { }, [settings?.ai?.ollamaBaseUrl, settings?.ai?.provider, refreshOllamaStatus]); useEffect(() => { - window.electron.onOllamaPullProgress((data) => { - const percent = data.total > 0 ? Math.round((data.completed / data.total) * 100) : 0; - setPullProgress({ status: data.status, percent }); - }); - window.electron.onOllamaPullDone(() => { - const preferredModel = pullingModelRef.current || undefined; - pullingModelRef.current = null; - setPullingModel(null); - setPullProgress({ status: '', percent: 0 }); - refreshOllamaStatus(preferredModel); - }); - window.electron.onOllamaPullError((data) => { - setPullingModel(null); - setPullProgress({ status: '', percent: 0 }); - setOllamaError(data.error); - setTimeout(() => setOllamaError(null), 5000); + return registerOllamaPullListeners({ + bridge: window.electron, + getActiveRequestId: () => pullRequestIdRef.current, + getPreferredModel: () => pullingModelRef.current || undefined, + clearActivePull: () => { + pullRequestIdRef.current = null; + pullingModelRef.current = null; + }, + setPullingModel, + setPullProgress, + setOllamaError, + scheduleErrorClear: () => { + setTimeout(() => setOllamaError(null), 5000); + }, + refreshOllamaStatus, }); }, [refreshOllamaStatus]); const handlePull = (modelName: string) => { const requestId = `ollama-pull-${Date.now()}`; + pullRequestIdRef.current = requestId; pullingModelRef.current = modelName; setPullingModel(modelName); setPullProgress({ status: t('settings.ai.llm.ollama.startingDownload'), percent: 0 }); diff --git a/src/renderer/src/settings/ollamaPullProgress.ts b/src/renderer/src/settings/ollamaPullProgress.ts new file mode 100644 index 00000000..9e17015a --- /dev/null +++ b/src/renderer/src/settings/ollamaPullProgress.ts @@ -0,0 +1,110 @@ +export type OllamaPullProgressEvent = { + requestId: string; + status: string; + digest: string; + total: number; + completed: number; +}; + +export type OllamaPullDoneEvent = { + requestId: string; +}; + +export type OllamaPullErrorEvent = { + requestId: string; + error: string; +}; + +export type OllamaPullProgressState = { + status: string; + percent: number; +}; + +type MaybeCleanup = (() => void) | void; + +type OllamaPullEventBridge = { + onOllamaPullProgress: (callback: (data: OllamaPullProgressEvent) => void) => MaybeCleanup; + onOllamaPullDone: (callback: (data: OllamaPullDoneEvent) => void) => MaybeCleanup; + onOllamaPullError: (callback: (data: OllamaPullErrorEvent) => void) => MaybeCleanup; +}; + +type RegisterOllamaPullListenersOptions = { + bridge: OllamaPullEventBridge; + getActiveRequestId: () => string | null; + getPreferredModel: () => string | undefined; + clearActivePull: () => void; + setPullingModel: (modelName: string | null) => void; + setPullProgress: (progress: OllamaPullProgressState) => void; + setOllamaError: (error: string | null) => void; + scheduleErrorClear: () => void; + refreshOllamaStatus: (preferredModelName?: string) => void; +}; + +function isActivePullEvent(eventRequestId: string, activeRequestId: string | null): boolean { + return Boolean(activeRequestId && eventRequestId === activeRequestId); +} + +export function toOllamaPullProgressState(data: OllamaPullProgressEvent): OllamaPullProgressState { + return { + status: data.status, + percent: data.total > 0 ? Math.round((data.completed / data.total) * 100) : 0, + }; +} + +export function registerOllamaPullListeners({ + bridge, + clearActivePull, + getActiveRequestId, + getPreferredModel, + refreshOllamaStatus, + scheduleErrorClear, + setOllamaError, + setPullingModel, + setPullProgress, +}: RegisterOllamaPullListenersOptions): () => void { + let lastProgress: (OllamaPullProgressState & { requestId: string }) | null = null; + + const cleanupProgress = bridge.onOllamaPullProgress((data) => { + if (!isActivePullEvent(data.requestId, getActiveRequestId())) return; + + const nextProgress = toOllamaPullProgressState(data); + if ( + lastProgress?.requestId === data.requestId && + lastProgress.status === nextProgress.status && + lastProgress.percent === nextProgress.percent + ) { + return; + } + + lastProgress = { ...nextProgress, requestId: data.requestId }; + setPullProgress(nextProgress); + }); + + const cleanupDone = bridge.onOllamaPullDone((data) => { + if (!isActivePullEvent(data.requestId, getActiveRequestId())) return; + + const preferredModel = getPreferredModel(); + lastProgress = null; + clearActivePull(); + setPullingModel(null); + setPullProgress({ status: '', percent: 0 }); + refreshOllamaStatus(preferredModel); + }); + + const cleanupError = bridge.onOllamaPullError((data) => { + if (!isActivePullEvent(data.requestId, getActiveRequestId())) return; + + lastProgress = null; + clearActivePull(); + setPullingModel(null); + setPullProgress({ status: '', percent: 0 }); + setOllamaError(data.error); + scheduleErrorClear(); + }); + + return () => { + cleanupProgress?.(); + cleanupDone?.(); + cleanupError?.(); + }; +} diff --git a/src/renderer/types/electron.d.ts b/src/renderer/types/electron.d.ts index 203a7699..e1a33178 100644 --- a/src/renderer/types/electron.d.ts +++ b/src/renderer/types/electron.d.ts @@ -1344,9 +1344,9 @@ export interface ElectronAPI { ollamaPull: (requestId: string, modelName: string) => Promise; ollamaDelete: (modelName: string) => Promise<{ success: boolean; error: string | null }>; ollamaOpenDownload: () => Promise; - onOllamaPullProgress: (callback: (data: { requestId: string; status: string; digest: string; total: number; completed: number }) => void) => void; - onOllamaPullDone: (callback: (data: { requestId: string }) => void) => void; - onOllamaPullError: (callback: (data: { requestId: string; error: string }) => void) => void; + onOllamaPullProgress: (callback: (data: { requestId: string; status: string; digest: string; total: number; completed: number }) => void) => (() => void); + onOllamaPullDone: (callback: (data: { requestId: string }) => void) => (() => void); + onOllamaPullError: (callback: (data: { requestId: string; error: string }) => void) => (() => void); // Hyper Key onHyperKeyCombo: (callback: (key: string) => void) => (() => void);