Blame for app/services/background_tasks.py

"""Run long operations in a background daemon thread.

Streamlit reruns the whole script on every interaction (including clicking a
different step in the sidebar). A long, synchronous operation — generating a
proposal, running the analysis — gets interrupted by that rerun, so the work is
lost. This helper runs the operation on an independent daemon thread whose life
is NOT tied to the script run, keeps the result in a module-level registry, and
lets the UI poll for completion. The user can switch steps freely; the work
keeps going and the result is waiting when they return.

IMPORTANT: the worker function must not touch ``st.session_state`` (Streamlit
calls from a thread with no script-run context fail). Pass everything it needs
as arguments and read the result back on the main thread.

The registry is shared by every visitor this process serves, so each run gets
its own id from :func:`new_task_id`, and the caller keeps that id in its own
``st.session_state``. Only that session can poll or collect the run, and a
session that forgets the id (a new job, a re-run) has abandoned it: clearing a
run that is still going discards whatever it produces.

API-key propagation: the API key now lives only in the per-session store
(``st.session_state``), which a worker thread cannot see. So before spawning the
worker we snapshot the launching session's config overrides on the MAIN thread
and re-install them into the worker's own thread-local (see
``app.config.export_session_overrides`` / ``install_thread_overrides``). This
makes ``get_settings()`` inside the worker resolve the same key/provider/model
the user configured — without leaking it to any other visitor's worker.

Usage log: for the same reason, the AI calls a worker makes cannot be logged to
the session's ``api_usage_log``. The worker logs them to a list of its own
(``app.services.llm_client.set_thread_usage_log``) that goes back with the
outcome, and :func:`pop` adds them to the log of the session collecting it.
"""

from __future__ import annotations

import threading
import time
import traceback
import uuid
from datetime import datetime, timezone
from typing import Any, Callable

from app import config
from app.services import diagnostics, llm_client
from app.utils.logging_utils import get_logger

_logger = get_logger("upwork_strategist.background_tasks")

_lock = threading.Lock()
_tasks: dict[str, dict[str, Any]] = {}

# A finished result waits here until its session collects it. One nobody ever
# collects (the tab was closed) is dropped after this long, so the shared
# registry cannot grow for the life of the process.
_RESULT_TTL_SECONDS = 60 * 60


def new_task_id(name: str) -> str:
    """A fresh id for one run of ``name`` (e.g. ``"proposal_generation"``).

    A fixed id would be shared by every visitor: one visitor's run would block
    another's, and whichever session polled first would collect it.
    """
    return f"{name}:{uuid.uuid4().hex}"


def _sweep_expired() -> None:
    """Drop finished results nobody collected in time. Call with ``_lock`` held."""
    cutoff = time.monotonic() - _RESULT_TTL_SECONDS
    expired = [
        task_id
        for task_id, t in _tasks.items()
        if t.get("finished_at", cutoff) < cutoff
    ]
    for task_id in expired:
        del _tasks[task_id]


def _report_crash(task_id: str, exc: Exception) -> None:
    """Log a worker crash and record it for diagnostics. Call in the ``except``.

    Without this a crash leaves no trace: the screen only learns that the run
    failed. The traceback goes to the server log, inside the message so the
    log's redaction filter scrubs it like any other line. The diagnostics
    event feeds the "Problem detected" panel and the report a user can share,
    so it names the exception's class only — never its message, which can
    quote dossier or job text.
    """
    task_name = task_id.split(":", 1)[0]
    error_type = type(exc).__name__
    _logger.error(
        "task=%s status=app_error error_type=%s\n%s",
        task_name,
        error_type,
        traceback.format_exc(),
    )
    # Same record shape as llm_client's, so a report renders one table.
    diagnostics.record_event(
        {
            "timestamp": datetime.now(timezone.utc).isoformat(timespec="seconds"),
            "task_name": task_name,
            "provider": None,
            "model": None,
            "used_api": False,
            "status": "app_error",
            "error_message": None,
            "duration_seconds": None,
            "estimated_input_tokens": None,
            "estimated_output_tokens": None,
            "request_id": None,
            "http_status": None,
            "error_type": error_type,
        }
    )


def _has_streamlit_runtime() -> bool:
    """True only inside a real Streamlit script run (not tests / bare mode)."""
    try:
        from streamlit.runtime.scriptrunner import get_script_run_ctx

        return get_script_run_ctx() is not None
    except Exception:  # noqa: BLE001
        return False


def start(task_id: str, fn: Callable[..., Any], *args: Any, **kwargs: Any) -> bool:
    """Start ``fn`` under ``task_id``.

    In the live app this runs on a daemon thread (so a Streamlit rerun can't
    kill it). Outside a real Streamlit runtime — unit tests, bare mode — it
    runs SYNCHRONOUSLY so callers that expect an immediate result still work.

    Returns True if a new run was started, False if one was already running
    (so a double-click or a rerun won't launch duplicate work).
    """
    record = {"status": "running", "result": None, "error": None}
    with _lock:
        _sweep_expired()
        cur = _tasks.get(task_id)
        if cur and cur.get("status") == "running":
            return False
        _tasks[task_id] = record

    # Snapshot the launching session's config (incl. the session-only API key)
    # HERE, on the main thread, while st.session_state is reachable. The worker
    # re-installs it into its own thread-local so get_settings() resolves the
    # right key. Empty when no session (tests / bare mode).
    session_overrides = config.export_session_overrides()
    # The usage-log entries of the AI calls this run makes. They go back with
    # the outcome, and pop() adds them to the collecting session's log.
    usage_log: list[dict] = []

    def _finish(outcome: dict[str, Any]) -> None:
        # Stored only while this run is still the one registered: a run its
        # session cleared (abandoned) must never come back as a result.
        outcome["finished_at"] = time.monotonic()
        outcome["usage_log"] = usage_log
        with _lock:
            if _tasks.get(task_id) is record:
                _tasks[task_id] = outcome

    def _run() -> None:
        if session_overrides:
            config.install_thread_overrides(session_overrides)
        previous_log = llm_client.set_thread_usage_log(usage_log)
        try:
            res = fn(*args, **kwargs)
            _finish({"status": "done", "result": res, "error": None})
        except Exception as exc:  # noqa: BLE001 - surface as a clean error status
            _report_crash(task_id, exc)
            # The exception itself goes back too, so the screen can tell a
            # provider failure from a fault in the app.
            _finish({
                "status": "error",
                "result": None,
                "error": str(exc),
                "error_type": type(exc).__name__,
                "exc": exc,
            })
        finally:
            llm_client.set_thread_usage_log(previous_log)
            if session_overrides:
                config.clear_thread_overrides()

    if _has_streamlit_runtime():
        threading.Thread(target=_run, daemon=True).start()
    else:
        _run()  # synchronous fallback for tests / bare mode
    return True


def status(task_id: str) -> dict[str, Any]:
    """Return ``{"status", "result", "error"}`` for ``task_id`` (status idle if unknown).

    A failed run also carries ``error_type`` and the exception as ``exc``.
    """
    with _lock:
        t = _tasks.get(task_id)
        return dict(t) if t else {"status": "idle", "result": None, "error": None}


def is_running(task_id: str) -> bool:
    return status(task_id).get("status") == "running"


def pop(task_id: str) -> dict[str, Any]:
    """Read and REMOVE a finished task's record (done/error). Running tasks stay.

    Call it on the main thread of the session that started the run: the AI
    calls the run made are added to that session's usage log here.
    """
    with _lock:
        t = _tasks.get(task_id)
        finished = bool(t and t.get("status") in ("done", "error"))
        if finished:
            _tasks.pop(task_id, None)
    if finished:
        llm_client.add_to_usage_log(t.get("usage_log") or ())
    return dict(t) if t else {"status": "idle", "result": None, "error": None}


def clear(task_id: str) -> None:
    with _lock:
        _tasks.pop(task_id, None)