"""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)