import { EventEmitter } from "events"; import * as fs from "fs"; import * as path from "path"; import * as crypto from "crypto"; import { app } from "electron"; import { Model, DownloadProgress, ModelManagerState, AVAILABLE_MODELS, } from "../../constants/models"; import { DownloadedModel } from "../../db/schema"; import { getDownloadedModelsRecord, createDownloadedModel, deleteDownloadedModel, validateDownloadedModels, validateModelFile, getValidDownloadedModels, } from "../../db/downloaded-models"; import { logger } from "../../main/logger"; interface ModelManagerEvents { "download-progress": (modelId: string, progress: DownloadProgress) => void; "download-complete": ( modelId: string, downloadedModel: DownloadedModel, ) => void; "download-error": (modelId: string, error: Error) => void; "download-cancelled": (modelId: string) => void; "model-deleted": (modelId: string) => void; } declare interface ModelManagerService { on( event: U, listener: ModelManagerEvents[U], ): this; emit( event: U, ...args: Parameters ): boolean; } class ModelManagerService extends EventEmitter { private state: ModelManagerState; private modelsDirectory: string; constructor() { super(); this.state = { activeDownloads: new Map(), }; // Create models directory in app data this.modelsDirectory = path.join(app.getPath("userData"), "models"); this.ensureModelsDirectory(); } // Initialize and validate models on startup async initialize(): Promise { try { const validation = await validateDownloadedModels(); if (validation.cleaned > 0) { logger.main.info("Cleaned up missing model records", { cleaned: validation.cleaned, valid: validation.valid.length, missing: validation.missing.map((m) => ({ id: m.id, path: m.localPath, })), }); } logger.main.info("Model manager initialized", { validModels: validation.valid.length, cleanedRecords: validation.cleaned, }); } catch (error) { logger.main.error("Error initializing model manager", { error }); } } private ensureModelsDirectory(): void { if (!fs.existsSync(this.modelsDirectory)) { fs.mkdirSync(this.modelsDirectory, { recursive: true }); logger.main.info("Created models directory", { path: this.modelsDirectory, }); } } // Get all available models from manifest getAvailableModels(): Model[] { return AVAILABLE_MODELS; } // Get downloaded models from database async getDownloadedModels(): Promise> { return await getDownloadedModelsRecord(); } // Get only valid downloaded models (files that exist on disk) async getValidDownloadedModels(): Promise> { const validModels = await getValidDownloadedModels(); const record: Record = {}; for (const model of validModels) { record[model.id] = model; } return record; } // Check if a model is downloaded and file exists async isModelDownloaded(modelId: string): Promise { return await validateModelFile(modelId); } // Get download progress for a model getDownloadProgress(modelId: string): DownloadProgress | null { return this.state.activeDownloads.get(modelId) || null; } // Get all active downloads getActiveDownloads(): DownloadProgress[] { return Array.from(this.state.activeDownloads.values()); } // Download a model async downloadModel(modelId: string): Promise { const model = AVAILABLE_MODELS.find((m) => m.id === modelId); if (!model) { throw new Error(`Model not found: ${modelId}`); } if (await this.isModelDownloaded(modelId)) { throw new Error(`Model already downloaded: ${modelId}`); } if (this.state.activeDownloads.has(modelId)) { throw new Error(`Download already in progress: ${modelId}`); } const abortController = new AbortController(); const downloadPath = path.join(this.modelsDirectory, model.filename); const progress: DownloadProgress = { modelId, progress: 0, status: "downloading", bytesDownloaded: 0, totalBytes: model.size, abortController, }; this.state.activeDownloads.set(modelId, progress); this.emit("download-progress", modelId, progress); try { logger.main.info("Starting model download", { modelId, size: model.sizeFormatted, url: model.downloadUrl, }); const response = await fetch(model.downloadUrl, { signal: abortController.signal, }); if (!response.ok) { throw new Error( `Failed to download: ${response.status} ${response.statusText}`, ); } const totalBytes = parseInt(response.headers.get("content-length") || "0") || model.size; progress.totalBytes = totalBytes; const fileStream = fs.createWriteStream(downloadPath); let bytesDownloaded = 0; let lastProgressEmit = 0; const reader = response.body?.getReader(); if (!reader) { throw new Error("Failed to get response reader"); } while (true) { const { done, value } = await reader.read(); if (done) break; if (abortController.signal.aborted) { fileStream.close(); fs.unlinkSync(downloadPath); throw new Error("Download cancelled"); } fileStream.write(value); bytesDownloaded += value.length; progress.bytesDownloaded = bytesDownloaded; progress.progress = Math.round((bytesDownloaded / totalBytes) * 100); // Emit progress every 1% or 1MB to avoid too many events const progressPercent = progress.progress; if ( progressPercent - lastProgressEmit >= 1 || bytesDownloaded - (lastProgressEmit * totalBytes) / 100 >= 1024 * 1024 ) { this.emit("download-progress", modelId, { ...progress }); lastProgressEmit = progressPercent; } } fileStream.end(); // Get actual file size (no validation against expected size) const stats = fs.statSync(downloadPath); logger.main.info("Download completed", { modelId, expectedSize: totalBytes, actualSize: stats.size, sizeDifference: Math.abs(stats.size - totalBytes), }); // Verify checksum if provided if (model.checksum) { const fileChecksum = await this.calculateFileChecksum(downloadPath); if (fileChecksum !== model.checksum) { fs.unlinkSync(downloadPath); throw new Error( `Checksum mismatch. Expected: ${model.checksum}, Got: ${fileChecksum}`, ); } } // Create downloaded model record in database const downloadedModel = await createDownloadedModel({ id: model.id, name: model.name, type: model.type, localPath: downloadPath, downloadedAt: new Date(), size: stats.size, checksum: model.checksum, }); // Clean up active download this.state.activeDownloads.delete(modelId); logger.main.info("Model download completed", { modelId, path: downloadPath, size: stats.size, }); this.emit("download-complete", modelId, downloadedModel); } catch (error) { // Clean up on error this.state.activeDownloads.delete(modelId); if (fs.existsSync(downloadPath)) { fs.unlinkSync(downloadPath); } const err = error instanceof Error ? error : new Error(String(error)); if (abortController.signal.aborted) { logger.main.info("Model download cancelled", { modelId }); this.emit("download-cancelled", modelId); } else { logger.main.error("Model download failed", { modelId, error: err.message, }); this.emit("download-error", modelId, err); } throw err; } } // Cancel a model download cancelDownload(modelId: string): void { const download = this.state.activeDownloads.get(modelId); if (!download) { throw new Error(`No active download found for model: ${modelId}`); } download.status = "cancelling"; download.abortController?.abort(); // Immediately remove from active downloads to prevent restart issues this.state.activeDownloads.delete(modelId); logger.main.info("Cancelled model download", { modelId }); this.emit("download-cancelled", modelId); } // Delete a downloaded model async deleteModel(modelId: string): Promise { const downloadedModels = await this.getDownloadedModels(); const downloadedModel = downloadedModels[modelId]; if (!downloadedModel) { throw new Error(`Model not found: ${modelId}`); } // Delete file if (fs.existsSync(downloadedModel.localPath)) { fs.unlinkSync(downloadedModel.localPath); logger.main.info("Deleted model file", { modelId, path: downloadedModel.localPath, }); } // Remove from database await deleteDownloadedModel(modelId); this.emit("model-deleted", modelId); } // Calculate file checksum (SHA-1) private async calculateFileChecksum(filePath: string): Promise { return new Promise((resolve, reject) => { const hash = crypto.createHash("sha1"); const stream = fs.createReadStream(filePath); stream.on("data", (data) => hash.update(data)); stream.on("end", () => resolve(hash.digest("hex"))); stream.on("error", reject); }); } // Get models directory path getModelsDirectory(): string { return this.modelsDirectory; } // Validate and clean up stale model records (can be called periodically) async validateAndCleanup(): Promise<{ cleaned: number; valid: number }> { try { const validation = await validateDownloadedModels(); if (validation.cleaned > 0) { logger.main.info("Periodic cleanup completed", { cleaned: validation.cleaned, valid: validation.valid.length, }); } return { cleaned: validation.cleaned, valid: validation.valid.length, }; } catch (error) { logger.main.error("Error during model validation cleanup", { error }); return { cleaned: 0, valid: 0 }; } } // Cleanup - cancel all active downloads cleanup(): void { logger.main.info("Cleaning up model downloads", { activeDownloads: this.state.activeDownloads.size, }); for (const [modelId] of this.state.activeDownloads) { try { this.cancelDownload(modelId); } catch (error) { logger.main.warn("Error cancelling download during cleanup", { modelId, error: error instanceof Error ? error.message : String(error), }); } } } } export { ModelManagerService };