"""In-memory async job store. Personal-scale: jobs are lost on restart. Generic — knows nothing about musicfetch; callers pass a no-arg `fn`. A job can also park at `awaiting_choice`: the worker produced a list of `choices` and waits for the client to pick one (`take_choice`). A parked job expires after CHOICE_TTL seconds so the store never fills with dead prompts.""" import os import time import uuid from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from typing import Any, Callable, Optional _EXECUTOR = ThreadPoolExecutor(max_workers=2) JOBS: "dict[str, Job]" = {} _MAX_JOBS = 200 # cap to bound memory CHOICE_TTL = float(os.environ.get("MUSICFETCH_CHOICE_TTL", "900")) # 15 minutes class ChoiceError(Exception): """The client's pick cannot be used. The text is safe to show a user.""" @dataclass class Job: id: str # queued | running | awaiting_choice | done | failed | expired status: str hit: Any message: str result: Optional[dict] = None error: Optional[str] = None choices: Optional[list] = None # [{"n": 1, "label": "..."}] shown to the client pending: Optional[list] = None # server-side data behind each choice expires_at: Optional[float] = None created_at: float = field(default_factory=time.time) updated_at: float = field(default_factory=time.time) def _touch(job: "Job", **changes): for k, v in changes.items(): setattr(job, k, v) job.updated_at = time.time() def _evict_if_needed(): # Post-condition: len(JOBS) <= _MAX_JOBS (evicts oldest overflow entries). if len(JOBS) <= _MAX_JOBS: return for jid in sorted(JOBS, key=lambda j: JOBS[j].created_at)[: len(JOBS) - _MAX_JOBS]: JOBS.pop(jid, None) def create_job(hit: Any, message: str) -> "Job": job = Job(id=uuid.uuid4().hex[:8], status="queued", hit=hit, message=message) JOBS[job.id] = job _evict_if_needed() return job def get_job(job_id: str) -> Optional["Job"]: job = JOBS.get(job_id) if job is not None and _is_stale(job): _touch(job, status="expired", choices=None, pending=None, message="That choice list has expired. Please search again.") return job def _is_stale(job: "Job") -> bool: return (job.status == "awaiting_choice" and job.expires_at is not None and time.time() > job.expires_at) def _fail_text(fail_message, exc: Exception) -> str: return fail_message(exc) if callable(fail_message) else fail_message def run_job(job_id: str, fn: Callable[[], dict], done_message: "str | Callable[[dict], str]", fail_message: "str | Callable[[Exception], str]" = "Something went wrong while fetching.") -> None: def _task(): job = JOBS.get(job_id) if job is None: return _touch(job, status="running") try: result = fn() msg = done_message(result) if callable(done_message) else done_message _touch(job, status="done", result=result, message=msg) except Exception as e: # noqa: BLE001 — record any failure on the job _touch(job, status="failed", error=f"{type(e).__name__}: {e}", message=_fail_text(fail_message, e)) _EXECUTOR.submit(_task) def await_choice(job_id: str, choices: list, pending: list, message: str) -> None: """Park a running job until the client picks one of `choices`.""" job = JOBS.get(job_id) if job is None: return _touch(job, status="awaiting_choice", choices=choices, pending=pending, message=message, expires_at=time.time() + CHOICE_TTL) def requeue(job: "Job", message: str) -> None: """Re-arm a parked job for a second piece of work (the client's pick).""" _touch(job, status="queued", message=message, error=None) def take_choice(job_id: str, n: int) -> tuple["Job", Any]: """Return (job, the picked pending item) and clear the prompt. Raises ChoiceError with a user-safe reason when the pick cannot be used.""" job = get_job(job_id) if job is None: raise KeyError(job_id) if job.status == "expired": raise ChoiceError("That choice list has expired. Please search again.") if job.status != "awaiting_choice": raise ChoiceError(f"That search is not waiting for a choice (it is {job.status}).") pending = job.pending or [] if not isinstance(n, int) or n < 1 or n > len(pending): raise ChoiceError(f"Pick a number between 1 and {len(pending)}.") item = pending[n - 1] # Leave awaiting_choice in the same touch that clears the list, so a second # request cannot take the same choice. _touch(job, status="queued", choices=None, pending=None, expires_at=None) return job, item def run_choice_job(job_id: str, fn: Callable[[], dict], choice_message: Callable[[dict], str], done_message: "str | Callable[[dict], str]", fail_message: "str | Callable[[Exception], str]" = "Something went wrong while searching.") -> None: """Like run_job, but a result carrying a non-empty "choices" list parks the job at awaiting_choice instead of finishing it.""" def _task(): job = JOBS.get(job_id) if job is None: return _touch(job, status="running") try: result = fn() # `pending` holds indexer guids and `choices` is already top-level, # so neither belongs in the public `result` body. public = {k: v for k, v in result.items() if k not in ("pending", "choices")} choices = result.get("choices") or [] if choices: await_choice(job_id, choices, result.get("pending") or [], choice_message(result)) _touch(job, result=public) return msg = done_message(result) if callable(done_message) else done_message _touch(job, status="done", result=public, message=msg) except Exception as e: # noqa: BLE001 — record any failure on the job _touch(job, status="failed", error=f"{type(e).__name__}: {e}", message=_fail_text(fail_message, e)) _EXECUTOR.submit(_task)