Spaces:
Sleeping
Sleeping
| """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) | |