Spaces:
Running
Running
File size: 4,513 Bytes
6303ae6 b5d28f1 6303ae6 b5d28f1 6303ae6 b5d28f1 6303ae6 b5d28f1 6303ae6 b5d28f1 6303ae6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | """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.
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. The LLM
client's usage logging is already guarded against a missing session.
"""
from __future__ import annotations
import threading
from typing import Any, Callable
from app import config
_lock = threading.Lock()
_tasks: dict[str, dict[str, Any]] = {}
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).
"""
with _lock:
cur = _tasks.get(task_id)
if cur and cur.get("status") == "running":
return False
_tasks[task_id] = {"status": "running", "result": None, "error": None}
# 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()
def _run() -> None:
if session_overrides:
config.install_thread_overrides(session_overrides)
try:
res = fn(*args, **kwargs)
with _lock:
_tasks[task_id] = {"status": "done", "result": res, "error": None}
except Exception as exc: # noqa: BLE001 - surface as a clean error status
with _lock:
_tasks[task_id] = {"status": "error", "result": None, "error": str(exc)}
finally:
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)."""
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."""
with _lock:
t = _tasks.get(task_id)
if t and t.get("status") in ("done", "error"):
_tasks.pop(task_id, None)
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)
|