import json import os import shutil import time import unicodedata from typing import Callable import requests _DEFAULT_DOWNLOADS_ROOT = "/slskd-downloads" _AUDIO_EXT = {"flac", "mp3", "m4a", "ogg", "opus", "wav", "aac"} _LOSSLESS_EXT = {"flac", "wav"} _SEARCH_POLLS = 40 # * 3s ≈ 2 min _MAX_XFER_POLLS = 1200 # * 3s ≈ 60 min absolute backstop for a slow-but-progressing peer _STALL_POLLS = 30 # * 3s ≈ 90s of no byte progress → abandon a queued/dead peer and fall through _POLL_SECONDS = 3 def _basename(path: str) -> str: return path.rsplit("\\", 1)[-1].rsplit("/", 1)[-1] def _dirname(path: str) -> str: if "\\" in path: return path.rsplit("\\", 1)[0] if "/" in path: return path.rsplit("/", 1)[0] return "" _ETA_CAP_SECONDS = 359_999 # ~100h; keeps the value well inside Job.downloadEtaSeconds (int4) def _eta_seconds(total_bytes: int, bytes_now: int, speed_bps: float) -> int | None: """Seconds until the transfer finishes at the current measured rate, or None when it can't be estimated (no throughput sample yet, unknown total, or already complete). Pure — unit-tested. A near-stalled peer drives speed_bps toward zero, so the estimate is capped: an uncapped value can exceed int4 and overflow the downloadEtaSeconds column write (poisoning the pipeline txn).""" if speed_bps <= 0 or total_bytes <= 0 or bytes_now >= total_bytes: return None return min(int((total_bytes - bytes_now) / speed_bps), _ETA_CAP_SECONDS) def _norm(s: str) -> str: """Lowercase, strip diacritics, and reduce to alphanumeric tokens, so 'Beyoncé' matches a peer's 'Beyonce' folder and punctuation ('I Am… Sasha Fierce') doesn't defeat the match.""" s = unicodedata.normalize("NFKD", s or "") s = "".join(ch for ch in s if not unicodedata.combining(ch)) s = "".join(ch if ch.isalnum() else " " for ch in s.lower()) return " ".join(s.split()) def _path_matches(needle: str, path: str) -> bool: """True if `needle` (an artist or album name) appears in a file path, tolerant of accents and punctuation. Every significant token (≥3 chars) must be present; for a very short needle (e.g. 'U2') the whole normalized string must appear.""" npath = _norm(path) tokens = [t for t in _norm(needle).split() if len(t) >= 3] if not tokens: whole = _norm(needle) return bool(whole) and whole in npath return all(t in npath for t in tokens) def _keep_matching(responses: list, needle: str) -> list: """Drop files — then now-empty responses — whose path doesn't match `needle`. Makes an album- or artist-only fallback search precise so we never grab a same-titled release by a different artist. Pure — unit-tested offline.""" out = [] for resp in responses: files = [f for f in (resp.get("files") or []) if _path_matches(needle, f.get("filename", ""))] if files: out.append({**resp, "files": files}) return out def _parse_search_responses(responses: list, artist: str = "") -> list[dict]: """Turn slskd search responses into album candidates (one per peer+directory), ordered so the peers that will FINISH first come first: a free upload slot, then the shortest estimated transfer time (album bytes ÷ the peer's advertised speed), then a shorter queue. Ranking on estimated time — not raw speed — matters because Soulseek FLAC rips of the same album vary ~2x in size (a 24-bit/bloated rip vs a standard CD rip), and Lyra scores them the SAME quality class (the search exposes no bit-depth), so a peer advertising high speed but serving huge files can lose a race — and blow past the download backstop — against a peer with a smaller standard rip. The pipeline tries candidates in this order, so this is what makes the Soulseek fall-through converge on a copy that actually completes. Pure — no I/O — unit-tested offline.""" scored: list[tuple] = [] for resp in responses: username = resp.get("username", "") has_slot = 0 if resp.get("hasFreeUploadSlot") else 1 # 0 sorts first (free slot = starts now) speed = int(resp.get("uploadSpeed") or 0) queue = int(resp.get("queueLength") or 0) by_dir: dict[str, list] = {} for f in resp.get("files", []) or []: name = f.get("filename", "") ext = name.rsplit(".", 1)[-1].lower() if "." in name else "" if ext not in _AUDIO_EXT: continue by_dir.setdefault(_dirname(name), []).append( {"filename": name, "size": f.get("size", 0), "ext": ext} ) for directory, dfiles in by_dir.items(): if not dfiles: continue # Estimated seconds to pull this album from this peer at its advertised rate. Unknown/ # zero speed → a low nominal (1 B/s) so total size still orders those peers last. total_bytes = sum(int(f["size"] or 0) for f in dfiles) est_seconds = total_bytes / (speed if speed > 0 else 1) lossless = all(f["ext"] in _LOSSLESS_EXT for f in dfiles) dirname = _basename(directory) # Trust the REQUESTED artist when it appears anywhere in the path — many rips put the # artist in a parent folder ("Rihanna\(2009) Rated R\…") or bury it ("2009 - Rihanna - # Rated R"), which the "Artist - Album" folder parse gets wrong; that tanked the # confidence score and lost real FLACs to clean-named MP3s. album stays the folder name # (fuzzy-matched downstream); artist is what the parse gets wrong. if artist and _path_matches(artist, dfiles[0]["filename"]): guess_artist, guess_album = artist, dirname elif " - " in dirname: guess_artist, guess_album = (p.strip() for p in dirname.split(" - ", 1)) else: guess_artist, guess_album = "", dirname candidate = { "source_ref": json.dumps({ "username": username, "files": [{"filename": f["filename"], "size": f["size"]} for f in dfiles], }), "title": guess_album, "artist": guess_artist, "track_count": len(dfiles), "format": "FLAC" if lossless else "MP3", "bitrate": None, } scored.append((has_slot, est_seconds, queue, candidate)) # ascending: free slot first, then shortest estimated transfer, then shortest queue scored.sort(key=lambda t: (t[0], t[1], t[2])) return [c for *_rest, c in scored] class SlskdClient: """Real Soulseek client talking to an slskd daemon's REST API. NOT unit-tested offline; see test_soulseek_live.py. slskd response/transfer JSON shapes may vary by version — validate with the live test. """ def __init__(self, config: dict): self._url = (config.get("slskd.url", "") or "").rstrip("/") self._key = config.get("slskd.api_key", "") # Where slskd (an external daemon) writes finished downloads. slskd saves into ITS # OWN directory, so this path must be a share/mount both slskd and this worker see. # Precedence: the slskd.downloadsDir Config row (Settings-overridable) > the # SLSKD_DOWNLOADS_ROOT env (the container mount point) > a sensible default. self._downloads_dir = ( config.get("slskd.downloadsDir") or os.environ.get("SLSKD_DOWNLOADS_ROOT") or _DEFAULT_DOWNLOADS_ROOT ) def is_configured(self) -> bool: return bool(self._url and self._key) def _headers(self) -> dict: return {"X-API-Key": self._key, "Content-Type": "application/json"} def _get(self, path: str) -> dict | list: r = requests.get(f"{self._url}{path}", headers=self._headers(), timeout=30) r.raise_for_status() return r.json() def _run_search(self, text: str) -> list: """POST a search, wait for it to complete, and return the raw responses (a list).""" r = requests.post( f"{self._url}/api/v0/searches", headers=self._headers(), json={"searchText": text}, timeout=30, ) r.raise_for_status() search_id = r.json()["id"] for _ in range(_SEARCH_POLLS): time.sleep(_POLL_SECONDS) state = self._get(f"/api/v0/searches/{search_id}") if state.get("isComplete") or "Completed" in str(state.get("state", "")): break responses = self._get(f"/api/v0/searches/{search_id}/responses") return responses if isinstance(responses, list) else [] def search_album(self, artist: str, album: str) -> list[dict]: responses = self._run_search(f"{artist} {album}") if not responses: # The Soulseek server silently drops searches containing certain blocklisted terms # (major-label artists / album titles like Adele, Rihanna, Beyoncé, "Lemonade") → # 0 responses. The blocked term poisons the whole query, but the OTHER field usually # isn't blocked, so retry with each field alone (normalized) and keep only results # whose path still matches the field we dropped. That recovers e.g. Beyoncé's # "I Am… Sasha Fierce" via the album title, without ever grabbing a same-titled album # by someone else. for query, needle in ((album, artist), (artist, album)): q = _norm(query) if not q: continue alt = _keep_matching(self._run_search(q), needle) if alt: responses = alt break return _parse_search_responses(responses, artist) def download(self, source_ref: str, dest: str, on_progress: Callable[[float], None]) -> dict: ref = json.loads(source_ref) username = ref["username"] files = ref["files"] wanted = {f["filename"] for f in files} os.makedirs(dest, exist_ok=True) xfer_ids: set[str] = set() # slskd transfer ids we started, for cancel-on-abandon try: return self._download(username, files, wanted, dest, on_progress, xfer_ids) except Exception: # Abandoning this peer (timeout / error / the pipeline moving on): cancel the transfers # we enqueued so they don't keep running and pile up in slskd as the job falls through # to other peers/sources for the same album. self._cancel_transfers(username, xfer_ids) raise def _download(self, username, files, wanted, dest, on_progress, xfer_ids: set) -> dict: r = requests.post( f"{self._url}/api/v0/transfers/downloads/{username}", headers=self._headers(), json=[{"filename": f["filename"], "size": f["size"]} for f in files], timeout=30, ) r.raise_for_status() on_progress(0.0) total = len(files) total_bytes = sum(int(f.get("size") or 0) for f in files) completed = False last_bytes = -1 stall = 0 speed = 0.0 # bytes/sec, EWMA-smoothed so the ETA doesn't jitter poll to poll prev_bytes = 0 prev_t = time.monotonic() for _ in range(_MAX_XFER_POLLS): time.sleep(_POLL_SECONDS) data = self._get(f"/api/v0/transfers/downloads/{username}") states: list[str] = [] bytes_now = 0 for directory in data.get("directories", []) if isinstance(data, dict) else []: for f in directory.get("files", []): if f.get("filename") in wanted: # only the files we enqueued states.append(str(f.get("state", ""))) bytes_now += int(f.get("bytesTransferred") or 0) if f.get("id"): # remember the transfer so we can cancel it if abandoned xfer_ids.add(str(f["id"])) if states: if any(bad in s for s in states for bad in ("Errored", "Rejected", "Cancelled")): raise RuntimeError("soulseek transfer failed") now = time.monotonic() dt = now - prev_t if dt > 0 and bytes_now >= prev_bytes: inst = (bytes_now - prev_bytes) / dt speed = inst if speed == 0.0 else 0.6 * speed + 0.4 * inst prev_bytes, prev_t = bytes_now, now done = sum(1 for s in states if "Completed" in s) on_progress(min(1.0, done / max(total, 1)), _eta_seconds(total_bytes, bytes_now, speed)) if done >= total: completed = True break # Stall detection: a peer making byte progress keeps its slot (up to the absolute # backstop above) so a slow-but-working transfer can finish; one that sends no bytes # for _STALL_POLLS (~90s) is queued/dead — abandon it fast and fall through. if bytes_now > last_bytes: last_bytes = bytes_now stall = 0 else: stall += 1 if stall >= _STALL_POLLS: raise TimeoutError(f"soulseek transfer stalled for {username}") if not completed: raise TimeoutError(f"soulseek download did not complete for {username}") # slskd saved the files into its OWN downloads dir — retrieve them into `dest` so the # pipeline (which verifies + imports against `dest`/staging) actually sees the audio. moved = self._retrieve(files, dest) if moved == 0: raise RuntimeError( f"soulseek transfer completed but no files were found under " f"{self._downloads_dir!r} to move into staging — check that the slskd " f"downloads dir is shared/mounted (SLSKD_DOWNLOADS_ROOT / slskd.downloadsDir)" ) return {"track_count": moved, "path": dest} def _cancel_transfers(self, username: str, xfer_ids: set) -> None: """Cancel the given in-flight slskd transfers for a peer. Best-effort: a cleanup failure must never mask the original download failure.""" for tid in xfer_ids: try: requests.delete( f"{self._url}/api/v0/transfers/downloads/{username}/{tid}", headers=self._headers(), timeout=15, ) except Exception: pass def _retrieve(self, files: list[dict], dest: str) -> int: """Move the just-finished slskd files out of its downloads dir into `dest`. slskd's on-disk layout can vary by version/config, so match tolerantly: for each wanted file, find on-disk candidates by basename anywhere under the downloads root, then prefer the one whose parent-folder name matches the remote directory and whose size matches. Returns the number of files actually moved. Offline-testable. """ root = self._downloads_dir if not root or not os.path.isdir(root): raise RuntimeError( f"slskd downloads dir not found: {root!r} — mount slskd's downloads " f"directory into the worker (SLSKD_DOWNLOADS_ROOT / slskd.downloadsDir)" ) os.makedirs(dest, exist_ok=True) wanted: dict[str, dict] = {} # basename(lower) -> {parent, size} for f in files: name = f.get("filename", "") wanted[_basename(name).lower()] = { "parent": _basename(_dirname(name)).lower(), "size": f.get("size"), } found: dict[str, list[str]] = {b: [] for b in wanted} for cur_root, _dirs, fnames in os.walk(root): for fn in fnames: key = fn.lower() if key in found: found[key].append(os.path.join(cur_root, fn)) moved = 0 for basename, paths in found.items(): if not paths: continue want = wanted[basename] def _score(p: str) -> tuple[int, int]: parent_ok = int(os.path.basename(os.path.dirname(p)).lower() == want["parent"]) size_ok = 0 try: size_ok = int(bool(want["size"]) and os.path.getsize(p) == want["size"]) except OSError: pass return (parent_ok, size_ok) best = max(paths, key=_score) shutil.move(best, os.path.join(dest, os.path.basename(best))) moved += 1 return moved