From b22c1511e18388ee584c3e2bc9a13e1ea02f9eb5 Mon Sep 17 00:00:00 2001 From: aloktripathi1 Date: Tue, 23 Jun 2026 22:13:13 +0530 Subject: [PATCH 1/2] fix(prompt): start PromptCache workers lazily instead of on import PromptCache spawned 2 background _RefreshWorker daemon threads eagerly at construction, including for the module-level singleton created at import. But refresh_async() -- the only thing that feeds those threads -- is never called in the SDK, so every user importing fi.prompt got 2 permanently idle threads. _TaskManager now starts workers lazily on first submit() via double-checked locking. No public API change; refresh_async() works as before, threads just start on demand. Adds 4 regression tests covering lazy startup, on-demand worker creation, idempotent start, and safe shutdown before any submit. --- python/fi/prompt/cache.py | 21 +++++++++++++---- python/tests/test_prompt_cache.py | 38 +++++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 5 deletions(-) create mode 100644 python/tests/test_prompt_cache.py diff --git a/python/fi/prompt/cache.py b/python/fi/prompt/cache.py index 295c4b0..9c729a8 100644 --- a/python/fi/prompt/cache.py +++ b/python/fi/prompt/cache.py @@ -77,16 +77,27 @@ class _TaskManager: """Manages background refresh workers and graceful shutdown.""" def __init__(self, num_workers: int): + self._num_workers = num_workers self._queue: "Queue[Callable[[], None]]" = Queue() - self._workers = [_RefreshWorker(self._queue, i) for i in range(num_workers)] - for w in self._workers: - w.start() - - atexit.register(self._shutdown) + self._workers: list = [] + self._start_lock = threading.Lock() + self._started = False # Public API ----------------------------------------------------------- + def _ensure_started(self): + if not self._started: + with self._start_lock: + if not self._started: + for i in range(self._num_workers): + w = _RefreshWorker(self._queue, i) + w.start() + self._workers.append(w) + atexit.register(self._shutdown) + self._started = True + def submit(self, task: Callable[[], None]): + self._ensure_started() self._queue.put(task) # Private -------------------------------------------------------------- diff --git a/python/tests/test_prompt_cache.py b/python/tests/test_prompt_cache.py new file mode 100644 index 0000000..398052f --- /dev/null +++ b/python/tests/test_prompt_cache.py @@ -0,0 +1,38 @@ +"""Tests for lazy worker initialization in _TaskManager.""" + +import threading +import time + +import pytest + +from fi.prompt.cache import _TaskManager + + +class TestLazyWorkerInit: + def test_no_workers_at_init(self): + tm = _TaskManager(num_workers=2) + assert tm._workers == [] + assert tm._started is False + + def test_workers_start_on_first_submit(self): + tm = _TaskManager(num_workers=2) + done = threading.Event() + tm.submit(done.set) + assert done.wait(timeout=3), "task never executed" + assert len(tm._workers) == 2 + assert tm._started is True + tm._shutdown() + + def test_ensure_started_is_idempotent(self): + tm = _TaskManager(num_workers=2) + tm._ensure_started() + tm._ensure_started() + tm._ensure_started() + assert len(tm._workers) == 2 + tm._shutdown() + + def test_shutdown_safe_before_any_submit(self): + tm = _TaskManager(num_workers=2) + # _shutdown must not raise even when workers were never started + tm._shutdown() + assert tm._workers == [] From e9d65c34da7955d2a395845f8c65841dfcb55b01 Mon Sep 17 00:00:00 2001 From: aloktripathi1 Date: Mon, 10 Aug 2026 17:09:33 +0530 Subject: [PATCH 2/2] refactor: remove unused refresh_async instead of making it lazy Per review: refresh_async had zero callers anywhere in the codebase and wasn't part of the public API (not exported from fi/prompt/__init__.py, not referenced in README or examples). Removing it fixes the same bug with a smaller diff than making it lazy, and matches the existing intent already documented in client.py: "Async refresh logic is not applied here to simplify the fallback flow." Deleted _RefreshWorker, _TaskManager, refresh_async(), _refreshing_keys, and the max_workers param. get(), get_stale(), set(), invalidate(), make_key(), and the prompt_cache singleton are untouched -- none of them referenced the deleted code. Replaced the 4 implementation-detail tests (which asserted on internal attributes like _workers/_started) with a single subprocess-based black-box test that checks actual thread names after importing fi.prompt. This is robust to refactors and catches the real regression class -- any new eager thread anywhere on import, not just this specific mechanism. --- python/fi/prompt/cache.py | 106 +----------------------------- python/tests/test_prompt_cache.py | 45 +++---------- 2 files changed, 13 insertions(+), 138 deletions(-) diff --git a/python/fi/prompt/cache.py b/python/fi/prompt/cache.py index 9c729a8..58da235 100644 --- a/python/fi/prompt/cache.py +++ b/python/fi/prompt/cache.py @@ -1,11 +1,9 @@ from __future__ import annotations -import atexit import logging import threading import time -from queue import Empty, Queue -from typing import Callable, Dict, Optional, Tuple +from typing import Dict, Optional # We deliberately import via string to avoid circular import at runtime from typing import TYPE_CHECKING @@ -19,7 +17,6 @@ # --------------------------------------------------------------------------- DEFAULT_TTL_SEC = 60 * 5 # 5 minutes -DEFAULT_REFRESH_WORKERS = 2 logger = logging.getLogger("fi.prompt.cache") @@ -42,82 +39,6 @@ def is_stale(self) -> bool: return time.time() >= self.expiry -# --------------------------------------------------------------------------- -# Background refresh infra (Queue + workers) -# --------------------------------------------------------------------------- - - -class _RefreshWorker(threading.Thread): - """Continuously processes refresh callables from the shared queue.""" - - def __init__(self, q: "Queue[Callable[[], None]]", identifier: int): - super().__init__(daemon=True, name=f"PromptCacheWorker-{identifier}") - self._queue = q - self._running = True - - def run(self) -> None: # noqa: D401 – imperative mood fine - while self._running: - try: - task = self._queue.get(timeout=1) - except Empty: - continue # check _running flag again - - try: - task() - except Exception as exc: # pragma: no cover – log + continue - logger.warning("Prompt cache refresh task failed: %s", exc, exc_info=True) - finally: - self._queue.task_done() - - def stop(self) -> None: - self._running = False - - -class _TaskManager: - """Manages background refresh workers and graceful shutdown.""" - - def __init__(self, num_workers: int): - self._num_workers = num_workers - self._queue: "Queue[Callable[[], None]]" = Queue() - self._workers: list = [] - self._start_lock = threading.Lock() - self._started = False - - # Public API ----------------------------------------------------------- - - def _ensure_started(self): - if not self._started: - with self._start_lock: - if not self._started: - for i in range(self._num_workers): - w = _RefreshWorker(self._queue, i) - w.start() - self._workers.append(w) - atexit.register(self._shutdown) - self._started = True - - def submit(self, task: Callable[[], None]): - self._ensure_started() - self._queue.put(task) - - # Private -------------------------------------------------------------- - - def _shutdown(self): - logger.debug("Shutting down PromptCache workers …") - for w in self._workers: - w.stop() - # Drain queue quickly - while not self._queue.empty(): - try: - self._queue.get_nowait() - self._queue.task_done() - except Empty: - break - for w in self._workers: - w.join(timeout=1) - logger.debug("PromptCache workers shut down.") - - # --------------------------------------------------------------------------- # Public cache API # --------------------------------------------------------------------------- @@ -126,12 +47,10 @@ def _shutdown(self): class PromptCache: """Thread-safe, stale-while-revalidate cache for `PromptTemplate` objects.""" - def __init__(self, ttl_sec: int = DEFAULT_TTL_SEC, max_workers: int = DEFAULT_REFRESH_WORKERS): + def __init__(self, ttl_sec: int = DEFAULT_TTL_SEC): self._ttl_sec = ttl_sec self._store: Dict[str, _CacheItem] = {} self._lock = threading.Lock() # protects _store mutations - self._refreshing_keys: set[str] = set() - self._tm = _TaskManager(max_workers) # ------------------------------- helpers ---------------------------- @@ -172,26 +91,7 @@ def invalidate(self, key_prefix: str): for k in to_delete: del self._store[k] - # -------------------------- refresh management --------------------- - - def refresh_async(self, key: str, fetch_fn: Callable[[], "PromptTemplate"]): - """Schedule a refresh if one is not already in-flight.""" - with self._lock: - if key in self._refreshing_keys: - return - self._refreshing_keys.add(key) - - def _task(): - try: - tpl = fetch_fn() - self.set(key, tpl) - finally: - with self._lock: - self._refreshing_keys.discard(key) - - self._tm.submit(_task) - # Global singleton -------------------------------------------------------- -prompt_cache = PromptCache() \ No newline at end of file +prompt_cache = PromptCache() diff --git a/python/tests/test_prompt_cache.py b/python/tests/test_prompt_cache.py index 398052f..efbb322 100644 --- a/python/tests/test_prompt_cache.py +++ b/python/tests/test_prompt_cache.py @@ -1,38 +1,13 @@ -"""Tests for lazy worker initialization in _TaskManager.""" +"""Tests for fi.prompt.cache.""" -import threading -import time +import subprocess +import sys -import pytest -from fi.prompt.cache import _TaskManager - - -class TestLazyWorkerInit: - def test_no_workers_at_init(self): - tm = _TaskManager(num_workers=2) - assert tm._workers == [] - assert tm._started is False - - def test_workers_start_on_first_submit(self): - tm = _TaskManager(num_workers=2) - done = threading.Event() - tm.submit(done.set) - assert done.wait(timeout=3), "task never executed" - assert len(tm._workers) == 2 - assert tm._started is True - tm._shutdown() - - def test_ensure_started_is_idempotent(self): - tm = _TaskManager(num_workers=2) - tm._ensure_started() - tm._ensure_started() - tm._ensure_started() - assert len(tm._workers) == 2 - tm._shutdown() - - def test_shutdown_safe_before_any_submit(self): - tm = _TaskManager(num_workers=2) - # _shutdown must not raise even when workers were never started - tm._shutdown() - assert tm._workers == [] +def test_importing_fi_prompt_spawns_no_threads(): + code = "import threading, fi.prompt; print([t.name for t in threading.enumerate()])" + result = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, text=True, check=True, + ) + assert "PromptCacheWorker" not in result.stdout