import { defineStore } from 'pinia' import { computed, ref } from 'vue' import type { AssetDownloadWsMessage } from '@/schemas/apiSchema' import { api } from '@/scripts/api' export interface AssetDownload { taskId: string assetId: string assetName: string bytesTotal: number bytesDownloaded: number progress: number status: 'created' | 'running' | 'completed' | 'failed' error?: string } interface CompletedDownload { taskId: string modelType: string timestamp: number } const PROCESSED_TASK_CLEANUP_MS = 60000 const MAX_COMPLETED_DOWNLOADS = 10 export const useAssetDownloadStore = defineStore('assetDownload', () => { /** Map of task IDs to their download progress data */ const downloads = ref>(new Map()) /** Map of task IDs to model types, used to track which model type to refresh after download completes */ const pendingModelTypes = new Map() /** Set of task IDs that have reached a terminal state (completed/failed), prevents duplicate processing */ const processedTaskIds = new Set() /** Reactive signal for completed downloads */ const completedDownloads = ref([]) const downloadList = computed(() => Array.from(downloads.value.values())) const activeDownloads = computed(() => downloadList.value.filter( (d) => d.status === 'created' || d.status === 'running' ) ) const finishedDownloads = computed(() => downloadList.value.filter( (d) => d.status === 'completed' || d.status === 'failed' ) ) const hasActiveDownloads = computed(() => activeDownloads.value.length > 0) const hasDownloads = computed(() => downloads.value.size > 0) /** * Associates a download task with its model type for later use when the download completes. * Intended for external callers (e.g., useUploadModelWizard) to register async downloads. */ function trackDownload(taskId: string, modelType: string) { pendingModelTypes.set(taskId, modelType) } /** * Handles asset download WebSocket events. Updates download progress, manages toast notifications, * and tracks completed downloads. Prevents duplicate processing of terminal states (completed/failed). */ function handleAssetDownload(e: CustomEvent) { const data = e.detail if (data.status === 'completed' || data.status === 'failed') { if (processedTaskIds.has(data.task_id)) return processedTaskIds.add(data.task_id) } const download: AssetDownload = { taskId: data.task_id, assetId: data.asset_id, assetName: data.asset_name, bytesTotal: data.bytes_total, bytesDownloaded: data.bytes_downloaded, progress: data.progress, status: data.status, error: data.error } downloads.value.set(data.task_id, download) if (data.status === 'completed') { const modelType = pendingModelTypes.get(data.task_id) if (modelType) { const newDownload: CompletedDownload = { taskId: data.task_id, modelType, timestamp: Date.now() } const updated = [...completedDownloads.value, newDownload] if (updated.length > MAX_COMPLETED_DOWNLOADS) { updated.shift() } completedDownloads.value = updated pendingModelTypes.delete(data.task_id) } setTimeout( () => processedTaskIds.delete(data.task_id), PROCESSED_TASK_CLEANUP_MS ) } else if (data.status === 'failed') { pendingModelTypes.delete(data.task_id) setTimeout( () => processedTaskIds.delete(data.task_id), PROCESSED_TASK_CLEANUP_MS ) } } api.addEventListener('asset_download', handleAssetDownload) function clearFinishedDownloads() { for (const download of finishedDownloads.value) { downloads.value.delete(download.taskId) } } return { activeDownloads, finishedDownloads, hasActiveDownloads, hasDownloads, downloadList, completedDownloads, trackDownload, clearFinishedDownloads } })