| 1 | import { existsSync, mkdirSync } from "node:fs"; |
| 2 | import { basename, extname, join } from "node:path"; |
| 3 | import type { BrowserDownloadState, BrowserDownloadView } from "../../shared/ipc.js"; |
| 4 | import type { Logger } from "../log.js"; |
| 5 | |
| 6 | export interface DownloadItemLike { |
| 7 | getURL(): string; |
| 8 | getFilename(): string; |
| 9 | getSavePath(): string; |
| 10 | setSavePath(path: string): void; |
| 11 | getState(): BrowserDownloadState; |
| 12 | getReceivedBytes(): number; |
| 13 | getTotalBytes(): number; |
| 14 | on(event: "updated" | "done", listener: (event: unknown, state: string) => void): unknown; |
| 15 | } |
| 16 | |
| 17 | export interface DownloadTrackerDeps { |
| 18 | tabForWebContents(webContentsId: number): { id: string; taskId: string } | undefined; |
| 19 | defaultDirectory(taskId: string): string; |
| 20 | onUpdate(download: BrowserDownloadView): void; |
| 21 | log: Logger; |
| 22 | exists?(path: string): boolean; |
| 23 | mkdir?(path: string): void; |
| 24 | setTimeout?: typeof setTimeout; |
| 25 | clearTimeout?: typeof clearTimeout; |
| 26 | } |
| 27 | |
| 28 | export interface HostDownload { |
| 29 | id: string; |
| 30 | url: string; |
| 31 | path: string; |
| 32 | state: BrowserDownloadState; |
| 33 | bytes: number; |
| 34 | } |
| 35 | |
| 36 | interface Waiter { |
| 37 | tabId: string; |
| 38 | resolve(): void; |
| 39 | } |
| 40 | |
| 41 | export class DownloadTracker { |
| 42 | private readonly downloads = new Map<string, BrowserDownloadView>(); |
| 43 | private readonly taskDirectories = new Map<string, string>(); |
| 44 | private readonly waiters = new Set<Waiter>(); |
| 45 | private readonly reservedPaths = new Set<string>(); |
| 46 | private counter = 0; |
| 47 | |
| 48 | constructor(private readonly deps: DownloadTrackerDeps) {} |
| 49 | |
| 50 | // Screenshots and acts name the task's scratch directory; downloads of |
| 51 | // that task's tabs land there so Go can hand the file over as it does |
| 52 | // captures. Without one the shell's own per-task folder is used. |
| 53 | setTaskDirectory(taskId: string, directory: string): void { |
| 54 | this.taskDirectories.set(taskId, directory); |
| 55 | } |
| 56 | |
| 57 | handleWillDownload(item: DownloadItemLike, webContentsId: number): void { |
| 58 | const tab = this.deps.tabForWebContents(webContentsId); |
| 59 | if (!tab) { |
| 60 | this.deps.log.warn(`download from unknown webContents ${webContentsId} dropped`); |
| 61 | return; |
| 62 | } |
| 63 | const directory = this.taskDirectories.get(tab.taskId) ?? this.deps.defaultDirectory(tab.taskId); |
| 64 | try { |
| 65 | (this.deps.mkdir ?? ((path: string) => mkdirSync(path, { recursive: true })))(directory); |
| 66 | } catch (error) { |
| 67 | this.deps.log.warn(`download directory ${directory} unavailable: ${String(error)}`); |
| 68 | } |
| 69 | const path = this.uniquePath(directory, item.getFilename()); |
| 70 | this.reservedPaths.add(path); |
| 71 | item.setSavePath(path); |
| 72 | this.counter += 1; |
| 73 | const id = `dl-${this.counter}`; |
| 74 | const record: BrowserDownloadView = { |
| 75 | id, |
| 76 | tabId: tab.id, |
| 77 | url: item.getURL(), |
| 78 | filename: basename(path), |
| 79 | path, |
| 80 | state: "progressing", |
| 81 | received: 0, |
| 82 | total: item.getTotalBytes(), |
| 83 | }; |
| 84 | this.downloads.set(id, record); |
| 85 | this.deps.onUpdate({ ...record }); |
| 86 | item.on("updated", (_event, state) => this.update(record, item, state)); |
| 87 | item.on("done", (_event, state) => { |
| 88 | this.update(record, item, state); |
| 89 | this.settle(); |
| 90 | }); |
| 91 | } |
| 92 | |
| 93 | list(tabId: string): BrowserDownloadView[] { |
| 94 | return [...this.downloads.values()].filter((download) => download.tabId === tabId).map((download) => ({ ...download })); |
| 95 | } |
| 96 | |
| 97 | hostList(tabId: string): HostDownload[] { |
| 98 | return this.list(tabId).map((download) => ({ id: download.id, url: download.url, path: download.path, state: download.state, bytes: download.received })); |
| 99 | } |
| 100 | |
| 101 | // Resolves once no download of the tab is in progress or the wait expires. |
| 102 | wait(tabId: string, waitForMs: number): Promise<HostDownload[]> { |
| 103 | if (waitForMs <= 0 || !this.inProgress(tabId)) return Promise.resolve(this.hostList(tabId)); |
| 104 | const schedule = this.deps.setTimeout ?? setTimeout; |
| 105 | const cancel = this.deps.clearTimeout ?? clearTimeout; |
| 106 | return new Promise((resolve) => { |
| 107 | const waiter: Waiter = { |
| 108 | tabId, |
| 109 | resolve: () => { |
| 110 | cancel(timer); |
| 111 | this.waiters.delete(waiter); |
| 112 | resolve(this.hostList(tabId)); |
| 113 | }, |
| 114 | }; |
| 115 | const timer = schedule(() => waiter.resolve(), waitForMs); |
| 116 | this.waiters.add(waiter); |
| 117 | }); |
| 118 | } |
| 119 | |
| 120 | forgetTab(tabId: string): void { |
| 121 | for (const [id, download] of this.downloads) { |
| 122 | if (download.tabId === tabId && download.state !== "progressing") this.downloads.delete(id); |
| 123 | } |
| 124 | } |
| 125 | |
| 126 | private inProgress(tabId: string): boolean { |
| 127 | return this.list(tabId).some((download) => download.state === "progressing"); |
| 128 | } |
| 129 | |
| 130 | private update(record: BrowserDownloadView, item: DownloadItemLike, state: string): void { |
| 131 | record.state = state === "completed" || state === "cancelled" || state === "interrupted" || state === "progressing" ? state : item.getState(); |
| 132 | record.received = item.getReceivedBytes(); |
| 133 | record.total = item.getTotalBytes(); |
| 134 | const savePath = item.getSavePath(); |
| 135 | if (savePath !== "") { |
| 136 | record.path = savePath; |
| 137 | record.filename = basename(savePath); |
| 138 | } |
| 139 | this.deps.onUpdate({ ...record }); |
| 140 | if (record.state === "cancelled" || record.state === "interrupted") this.reservedPaths.delete(record.path); |
| 141 | } |
| 142 | |
| 143 | private settle(): void { |
| 144 | for (const waiter of [...this.waiters]) { |
| 145 | if (!this.inProgress(waiter.tabId)) waiter.resolve(); |
| 146 | } |
| 147 | } |
| 148 | |
| 149 | private uniquePath(directory: string, filename: string): string { |
| 150 | const exists = this.deps.exists ?? existsSync; |
| 151 | const safe = basename(filename) || "download"; |
| 152 | const ext = extname(safe); |
| 153 | const stem = ext ? safe.slice(0, -ext.length) : safe; |
| 154 | let candidate = join(directory, safe); |
| 155 | for (let n = 1; exists(candidate) || this.reservedPaths.has(candidate); n += 1) candidate = join(directory, `${stem}-${n}${ext}`); |
| 156 | return candidate; |
| 157 | } |
| 158 | } |
| 159 |