diff --git a/packages/python/README.md b/packages/python/README.md index 9fdd7f19..ecb19b04 100644 --- a/packages/python/README.md +++ b/packages/python/README.md @@ -89,19 +89,17 @@ print(price.total_price) ### `UpdatePrices` -`UpdatePrices` can be used to periodically update the price data by downloading it from GitHub +`UpdatePrices` periodically updates the price data by downloading it from GitHub. Please note: - this functionality is explicitly opt-in - we download data directly from GitHub (`https://raw.githubusercontent.com/pydantic/genai-prices/refs/heads/main/prices/data.json`) so we don't and can't monitor requests or gather telemetry -At the time of writing, the `data.json` file -downloaded by `UpdatePrices` is around 26KB when compressed, so is generally very quick to download. +At the time of writing, the `data.json` file downloaded is around 26KB when compressed, so is generally very +quick to download. By default the first fetch happens immediately in the background, then every hour after that. -By default `UpdatePrices` downloads price data immediately after it's started in the background, then every hour after that. - -Usage with `UpdatePrices` as as context manager: +Usage as a context manager: ```py from genai_prices import UpdatePrices, Usage, calc_price @@ -112,21 +110,57 @@ with UpdatePrices() as update_prices: print(p) ``` -Usage with `UpdatePrices` as a simple class: +Or by calling `start()` / `stop()` yourself: ```py from genai_prices import UpdatePrices, Usage, calc_price update_prices = UpdatePrices() -update_prices.start(wait=True) # start updating prices, optionally wait for prices to have updated +update_prices.start(wait=True) # optionally wait for the first update p = calc_price(Usage(input_tokens=123, output_tokens=456), 'gpt-5') print(p) -update_prices.stop() # stop updating prices +update_prices.stop() ``` -Only one `UpdatePrices` instance can be running at a time. +A single shared, process-wide updater backs every `UpdatePrices` instance, so it is safe to create and start them +from anywhere — including from several libraries in the same process. Starting an instance is a **claim** on that +one updater, not a private thread: the first `start()` launches it, later ones join it, and the background thread +stops (reverting prices to the data bundled with the installed package) only when the **last** started instance is +stopped. + +**For libraries and integrations** (e.g. Logfire, Pydantic AI): create one `UpdatePrices()` with no arguments, +`start()` it at startup, and `stop()` it on shutdown. If several libraries do this they share the one updater +rather than spawning duplicate threads, and each one's claim independently keeps it alive until stopped. Leave +configuration to the application. + +**For application authors** who need a custom URL or refresh interval: pass them, and start early in startup, +before the libraries initialize: + +```py +UpdatePrices(url='https://my-mirror.example/prices.json', update_interval=1800).start() +``` + +Configuration is **first-wins**: the first instance to start fixes `url`, `update_interval` and `request_timeout` +for the updater's lifetime. Starting another instance with different settings logs a warning on the `genai-prices` +logger and joins the running updater; its settings are ignored. Starting early ensures your configuration is the +one that takes effect. + +`stop()` is idempotent and never raises — a background fetch error is logged instead (use `wait()` if you want it +raised). Stopping the last claim reverts prices to the bundled data immediately. If a fetch is in flight, the +daemon thread is given a short grace period to exit and is otherwise abandoned with a warning log: it exits once +the fetch completes, and its result is discarded — a stopped updater can never install prices afterwards. Other +updater API calls are never blocked, and `calc_price` is always unaffected. + +`start()` does not wait for the download (unless you pass `wait`). Until the first fetch completes, `calc_price` +keeps using the data bundled with the installed package, so prices for models released after that snapshot may be +missing for the first moments of the process. Once the fetch lands, every subsequent calculation uses the fresh +data — prices computed before then are not recalculated. If you need fresh prices before calculating (e.g. in a +short-lived script), pass `wait` to `start()`, or call `wait_prices_updated_sync()` / +`wait_prices_updated_async()` — these never raise and return `False` if the update failed (the error is logged on +the `genai-prices` logger), so your calculations simply fall back to the bundled data. -If you'd like to wait for prices to be updated without access to the `UpdatePrices` instance, you can use the `wait_prices_updated_sync` function: +You can wait for prices to be updated from anywhere — without access to the `UpdatePrices` instance — with +`wait_prices_updated_sync`: ```py from genai_prices import wait_prices_updated_sync @@ -135,7 +169,7 @@ wait_prices_updated_sync() ... ``` -Or it's async variant, `wait_prices_updated_async`. +Or its async variant, `wait_prices_updated_async`. ### CLI Usage diff --git a/packages/python/genai_prices/__init__.py b/packages/python/genai_prices/__init__.py index c9a7df1c..ac7f7812 100644 --- a/packages/python/genai_prices/__init__.py +++ b/packages/python/genai_prices/__init__.py @@ -6,10 +6,21 @@ from . import data_snapshot, types from .types import Usage -from .update_prices import UpdatePrices, wait_prices_updated_async, wait_prices_updated_sync +from .update_prices import ( + UpdatePrices, + wait_prices_updated_async, + wait_prices_updated_sync, +) __version__ = _metadata_version('genai_prices') -__all__ = 'Usage', 'calc_price', 'UpdatePrices', 'wait_prices_updated_sync', 'wait_prices_updated_async', '__version__' +__all__ = ( + 'Usage', + 'calc_price', + 'UpdatePrices', + 'wait_prices_updated_sync', + 'wait_prices_updated_async', + '__version__', +) @overload diff --git a/packages/python/genai_prices/update_prices.py b/packages/python/genai_prices/update_prices.py index 8a4e6b69..87fca8ca 100644 --- a/packages/python/genai_prices/update_prices.py +++ b/packages/python/genai_prices/update_prices.py @@ -2,6 +2,7 @@ import asyncio import logging +import os import threading from dataclasses import dataclass, field from time import time @@ -19,107 +20,275 @@ logger = logging.getLogger('genai-prices') DEFAULT_UPDATE_URL = 'https://raw.githubusercontent.com/pydantic/genai-prices/refs/heads/main/prices/data.json' -_global_update_prices: UpdatePrices | None = None +DEFAULT_UPDATE_INTERVAL = 3600.0 + + +def _default_request_timeout() -> httpx2.Timeout: + return httpx2.Timeout(timeout=10, connect=5) + + +# There is exactly one price snapshot per process (data_snapshot.set_custom_snapshot is a global), +# so "updating prices" is inherently a singleton activity. Every `UpdatePrices` is therefore not an +# updater of its own but a *claim* on one shared, process-wide updater guarded by `_lock`: starting +# any instance acquires a claim (ref-count++), and the background thread stops only when the last +# claim is released. +# +# This module-level state would not survive os.fork() on its own (threads die with the parent, and +# locks can be inherited in a held state, e.g. under gunicorn with preload_app=True). Fork hooks +# registered lazily in _register_fork_hooks() keep it consistent: the lock is held across the fork and +# reinitialized in the child, and a running updater is restarted in place - see _fork_after_in_child. +_updater: _BackgroundUpdater | None = None +_ref_count = 0 +_lock = threading.RLock() + +_fork_hooks_registered = False + + +def _register_fork_hooks() -> None: + """Keep the updater working across os.fork() (e.g. gunicorn with preload_app=True). + + Called under `_lock` from `_BackgroundUpdater._start_thread`, so registration happens at most + once, and only in processes that actually start an updater. + """ + global _fork_hooks_registered + if _fork_hooks_registered or not hasattr(os, 'register_at_fork'): + return + _fork_hooks_registered = True + os.register_at_fork(before=_fork_before, after_in_parent=_fork_after_in_parent, after_in_child=_fork_after_in_child) + + +def _fork_before() -> None: + # Hold the lock across the fork so the child inherits consistent bookkeeping rather than a + # torn, mid-mutation state. + _lock.acquire() + + +def _fork_after_in_parent() -> None: + _lock.release() + + +def _fork_after_in_child() -> None: + global _lock, _updater, _ref_count + + # The child inherits the lock acquired in _fork_before; replace it with a fresh one rather + # than releasing the inherited object. + _lock = threading.RLock() + updater = _updater + if updater is None: + return + try: + # The parent's updater thread does not survive the fork; restart it on the same instance, + # preserving identity so existing claims remain valid. + updater._revive_after_fork() # pyright: ignore[reportPrivateUsage] + except Exception: + _updater = None + _ref_count = 0 + data_snapshot.set_custom_snapshot(None) + logger.warning('Failed to restart the genai-prices background updater after fork', exc_info=True) + + +_STOPPED_THREAD_JOIN_TIMEOUT = 5.0 + + +def _join_stopped_updater_thread(thread: threading.Thread | None) -> None: + """Give a stopped updater thread a short grace period to exit, then abandon it. + + Called outside `_lock` so a thread blocked on an in-flight fetch never stalls other updater + API calls. Abandonment is safe: the thread is a daemon, it exits as soon as the fetch + completes, and a stopped updater can never install its result (see `_install_snapshot`). + """ + if thread is None: + return + thread.join(timeout=_STOPPED_THREAD_JOIN_TIMEOUT) + if thread.is_alive(): + logger.warning( + 'genai-prices background updater thread did not exit within %.0f seconds (a fetch is ' + 'likely in flight); abandoning the daemon thread. It will exit once the fetch ' + 'completes, without updating prices.', + _STOPPED_THREAD_JOIN_TIMEOUT, + ) def wait_prices_updated_sync(timeout: float | None = None) -> bool: - """Synchronously wait for prices to be updated. + """Synchronously wait for prices to be updated by the shared background updater. + + Never raises: if the update attempt failed, this returns False - the error is logged on the + `genai-prices` logger. Args: - timeout: The maximum time to wait for prices to be updated. Defaults to None which waits indefinitely. + timeout: The maximum time to wait for prices to be updated. Defaults to None which waits + indefinitely. Returns: - True if prices were updated, False otherwise. + True if prices were updated, False otherwise (including when the update failed, or when no + updater is running). """ - if _global_update_prices: - return _global_update_prices.wait(timeout) + with _lock: + updater = _updater + + if updater is not None: + return updater._wait_updated(timeout) # pyright: ignore[reportPrivateUsage] return False async def wait_prices_updated_async(timeout: float | None = None) -> bool: - """Asynchronously wait for prices to be updated. + """Asynchronously wait for prices to be updated by the shared background updater. + + Never raises: if the update attempt failed, this returns False - the error is logged on the + `genai-prices` logger. Args: timeout: The maximum time to wait for prices to be updated. Defaults to None which waits indefinitely. Returns: - True if prices were updated, False otherwise. + True if prices were updated, False otherwise (including when the update failed, or when no + updater is running). """ return await asyncio.to_thread(wait_prices_updated_sync, timeout) @dataclass class UpdatePrices: - """Update prices in the background using a daemon thread. + """Periodically update price data by downloading it in a background daemon thread. - Can be used either as a context manager or as a simple class, where you'll need to call start() and stop() manually. + A single shared, process-wide updater backs every `UpdatePrices` instance: starting an instance + is a *claim* on that one updater, not a private thread. It is therefore safe to create and start + `UpdatePrices` from anywhere - including from several libraries in the same process. The first + `start()` launches the updater; later ones join it and bump a reference count; the background + thread stops (and prices revert to the bundled data) only when the **last** started instance is + stopped. + + Configuration is first-wins: the first instance to start fixes `url`, `update_interval` and + `request_timeout` for the updater's lifetime. Starting another instance with different settings + logs a warning and joins the running updater; its settings are ignored. Application authors who + need a custom URL should start their instance early, before any library does. + + Can be used as a context manager (`with UpdatePrices(): ...`) or by calling `start()`/`stop()`. """ - update_interval: float = 3600 + update_interval: float = DEFAULT_UPDATE_INTERVAL """How often to update prices in seconds.""" url: str = DEFAULT_UPDATE_URL """The URL to fetch prices from.""" - request_timeout: httpx2.Timeout = field(default_factory=lambda: httpx2.Timeout(timeout=10, connect=5)) + request_timeout: httpx2.Timeout = field(default_factory=_default_request_timeout) """The timeout for HTTP requests.""" - _stop_event: threading.Event = field(default_factory=threading.Event) - _prices_updated: threading.Event = field(default_factory=threading.Event) - _thread: threading.Thread | None = field(default=None, init=False) - _background_exc: Exception | None = field(default=None, init=False) + _claimed: bool = field(default=False, init=False) + _updater: _BackgroundUpdater | None = field(default=None, init=False) def start(self, *, wait: bool | float = False): - """Start the background task. + """Acquire this instance's claim on the shared background updater. + + The first claim starts the updater; later claims join it. Starting the same instance twice + raises `RuntimeError`. This does not wait for the download unless `wait` is passed: until + the first fetch completes, price calculations keep using the bundled data. Args: - wait: Whether to wait for the prices to be updated before returning, if an int is passed - wait for that many seconds, if `True` wait for 30 seconds. + wait: Whether to wait for prices to be updated before returning; if an int/float is + passed wait that many seconds, if `True` wait for 30 seconds. """ - global _global_update_prices - - if self._thread is not None: - raise RuntimeError('UpdatePrices background task already started') - - if _global_update_prices is not None: - raise RuntimeError( - 'UpdatePrices global task already started, only one UpdatePrices can be active at a time' - ) + global _updater, _ref_count + + with _lock: + if self._claimed: + raise RuntimeError('UpdatePrices background task already started') + + if _updater is None: + updater = _BackgroundUpdater( + url=self.url, update_interval=self.update_interval, request_timeout=self.request_timeout + ) + _updater = updater + _ref_count = 1 + try: + updater._start_thread() # pyright: ignore[reportPrivateUsage] + except Exception: + _updater = None + _ref_count = 0 + raise + else: + if ( + _updater.url != self.url + or _updater.update_interval != self.update_interval + or _updater.request_timeout != self.request_timeout + ): + logger.warning( + 'A genai-prices background updater is already running (url=%r, update_interval=%r, ' + 'request_timeout=%r); ignoring the different configuration of this UpdatePrices ' + '(url=%r, update_interval=%r, request_timeout=%r) and joining the existing updater. ' + 'Start the updater before any other caller if you need custom settings.', + _updater.url, + _updater.update_interval, + _updater.request_timeout, + self.url, + self.update_interval, + self.request_timeout, + ) + _ref_count += 1 + + self._claimed = True + self._updater = _updater - _global_update_prices = self - self._prices_updated.clear() - self._stop_event.clear() - self._background_exc = None - self._thread = threading.Thread(target=self._background_task, daemon=True, name='genai_prices:update') - self._thread.start() if wait: self.wait(timeout=30 if wait is True else wait) + def stop(self): + """Release this instance's claim on the shared background updater. + + Stops the updater (reverting prices to the bundled data) only if this was the last open + claim; a claim held by another caller keeps it running. A no-op if this instance was never + started or has already been stopped. + + Never raises: a stored background fetch error is logged instead - use `wait()` if you want + the error raised. Returns promptly; if a fetch is in flight the daemon thread is given a + short grace period to exit before being abandoned with a warning log, and its result is + discarded - a stopped updater can never install prices afterwards. + """ + global _updater, _ref_count + + with _lock: + if not self._claimed: + return + self._claimed = False + + updater = self._updater + self._updater = None + if updater is None or _updater is not updater: + return + + _ref_count -= 1 + if _ref_count > 0: + return + + _updater = None + _ref_count = 0 + thread = updater._stop_and_detach() # pyright: ignore[reportPrivateUsage] + + # Join outside the lock so a thread blocked on an in-flight fetch never stalls other + # updater API calls; re-publication is already fenced off by _stop_and_detach(). + _join_stopped_updater_thread(thread) + exc = updater._background_exc # pyright: ignore[reportPrivateUsage] + if exc: + updater._background_exc = None # pyright: ignore[reportPrivateUsage] + logger.error('Error from genai-prices background updater while stopping', exc_info=exc) + def wait(self, timeout: float | None = None) -> bool: - """Wait for the prices to be updated in the background task. + """Wait for prices to be updated by the shared background updater. + + Raises the stored exception if the update attempt failed. Returns False if this instance is + not currently started. Args: timeout: The maximum time to wait for the prices to be updated in seconds. """ - prices_updated = self._prices_updated.wait(timeout=timeout) - exc = self._background_exc - if exc: - self._background_exc = None - raise exc - return prices_updated + updater = self._updater + if updater is None: + return False + return updater._wait_raising(timeout) # pyright: ignore[reportPrivateUsage] - def stop(self): - """Stop the background task.""" - global _global_update_prices - - _global_update_prices = None - if self._thread is not None: - self._stop_event.set() - self._thread.join() - self._thread = None - # Clear after the thread exits so an in-flight fetch cannot reinstall a snapshot after stop(). - data_snapshot.set_custom_snapshot(None) - if self._background_exc: - exc = self._background_exc - self._background_exc = None - raise exc + def fetch(self) -> data_snapshot.DataSnapshot | None: + """Fetch the latest provider data from the configured URL (does not start a background task).""" + return _BackgroundUpdater( + url=self.url, update_interval=self.update_interval, request_timeout=self.request_timeout + ).fetch() def __enter__(self): self.start() @@ -128,25 +297,94 @@ def __enter__(self): def __exit__(self, *_args: object): self.stop() + +@dataclass +class _BackgroundUpdater: + """The single process-wide updater. Internal: consumers interact via `UpdatePrices`.""" + + url: str + update_interval: float + request_timeout: httpx2.Timeout + _stop_event: threading.Event = field(default_factory=threading.Event) + _prices_updated: threading.Event = field(default_factory=threading.Event) + _thread: threading.Thread | None = field(default=None) + _background_exc: Exception | None = field(default=None) + _update_succeeded: bool = field(default=False) + + def _start_thread(self) -> None: + if self._thread is not None: + raise RuntimeError('genai-prices background task already started') + + _register_fork_hooks() + self._prices_updated.clear() + self._stop_event.clear() + self._background_exc = None + self._update_succeeded = False + self._thread = threading.Thread(target=self._background_task, daemon=True, name='genai_prices:update') + self._thread.start() + + def _wait_updated(self, timeout: float | None) -> bool: + # Never raises and does not consume the stored background exception: failures are already + # logged by the background task, and the exception stays stored for stop() to surface. + return self._prices_updated.wait(timeout=timeout) and self._update_succeeded + + def _wait_raising(self, timeout: float | None) -> bool: + # Like _wait_updated, but raises (and consumes) the stored background exception - the + # behaviour of UpdatePrices.wait(). + prices_updated = self._prices_updated.wait(timeout=timeout) + exc = self._background_exc + if exc: + self._background_exc = None + raise exc + return prices_updated and self._update_succeeded + + def _revive_after_fork(self) -> None: + # Threads do not survive fork, and the inherited events may have been captured in an + # arbitrary state (e.g. mid-install); recreate them and restart the background thread on + # this same instance. The snapshot is fenced by the module `_lock`, which is itself + # replaced with a fresh one in `_fork_after_in_child`. + self._thread = None + self._stop_event = threading.Event() + self._prices_updated = threading.Event() + self._start_thread() + + def _stop_and_detach(self) -> threading.Thread | None: + """Signal the background thread to stop and revert prices to the bundled data. + + Called with the module `_lock` held. Does not join; returns the (possibly still draining) + thread so the caller can wait on it after releasing `_lock`. Setting the stop event before + clearing the snapshot, both ordered against `_install_snapshot` via `_lock`, guarantees an + in-flight fetch can finish but can never install its result afterwards. + """ + thread = self._thread + self._thread = None + self._stop_event.set() + data_snapshot.set_custom_snapshot(None) + return thread + def _background_task(self) -> None: logger.info('Starting genai-prices background task') try: while True: try: - self._update_prices() - self._prices_updated.set() + installed = self._update_prices() self._background_exc = None + # Reflect whether the snapshot was actually installed: a fetch discarded by the + # stop fence (see _install_snapshot) must not report success to waiters. Set + # before signaling the event so waiters observing the event see the flag. + self._update_succeeded = installed + self._prices_updated.set() except Exception as e: self._background_exc = e self._prices_updated.set() - logger.error('Error updating genai-prices in the background (%s): %s', type(e).__name__, e) + logger.exception('Error updating genai-prices in the background') if self._stop_event.wait(self.update_interval): break finally: logger.info('genai-prices background task stopped') - def _update_prices(self): + def _update_prices(self) -> bool: start = time() snapshot = self.fetch() interval = time() - start @@ -155,7 +393,20 @@ def _update_prices(self): else: logger.info('Successfully fetched null snapshot in %.2f seconds', interval) - data_snapshot.set_custom_snapshot(snapshot) + return self._install_snapshot(snapshot) + + def _install_snapshot(self, snapshot: data_snapshot.DataSnapshot | None) -> bool: + # Fencing check: never publish after stop. The stop path sets _stop_event and then clears + # the snapshot under `_lock` (see _stop_and_detach), so taking the same lock here and + # re-checking the event makes publish-after-stop impossible rather than merely unlikely. + # Returns whether the snapshot was installed, so a discarded fetch is not reported as a + # successful update to waiters. + with _lock: + if self._stop_event.is_set(): + logger.info('genai-prices updater was stopped during the fetch; discarding the fetched snapshot') + return False + data_snapshot.set_custom_snapshot(snapshot) + return True def fetch(self) -> data_snapshot.DataSnapshot | None: """Fetches the latest provider data from the configured URL.""" diff --git a/tests/test_custom_prices.py b/tests/test_custom_prices.py index 8b6da9e8..717247a4 100644 --- a/tests/test_custom_prices.py +++ b/tests/test_custom_prices.py @@ -9,7 +9,27 @@ from genai_prices import Usage, calc_price, data, types from genai_prices.data_snapshot import DataSnapshot, set_custom_snapshot -from genai_prices.update_prices import UpdatePrices + + +class _CustomSnapshotSource: + """Test helper that installs a custom snapshot for the duration of a `with` block. + + Replaces subclassing the (deprecated) UpdatePrices: these tests only need a custom snapshot + active, not a background updater thread. Subclasses implement `fetch()`. + """ + + def fetch(self) -> DataSnapshot | None: + raise NotImplementedError + + def __enter__(self) -> _CustomSnapshotSource: + set_custom_snapshot(self.fetch()) + return self + + def wait(self, _timeout: float | None = None) -> bool: + return True + + def __exit__(self, *_args: object) -> None: + set_custom_snapshot(None) @dataclass @@ -23,7 +43,7 @@ def calc_price(self, usage: types.AbstractUsage) -> types.CalcPrice: return price -class AltUpdatePrices(UpdatePrices): +class AltUpdatePrices(_CustomSnapshotSource): def fetch(self) -> DataSnapshot | None: custom_providers = [ types.Provider( @@ -53,7 +73,7 @@ def fetch(self) -> DataSnapshot | None: return DataSnapshot(providers=custom_providers, from_auto_update=False) -class ExtraUpdatePrices(UpdatePrices): +class ExtraUpdatePrices(_CustomSnapshotSource): def fetch(self) -> DataSnapshot | None: providers = deepcopy(data.providers) openai = next(provider for provider in providers if provider.id == 'openai') diff --git a/tests/test_tiered_pricing.py b/tests/test_tiered_pricing.py index ae7329ef..e8fca34b 100644 --- a/tests/test_tiered_pricing.py +++ b/tests/test_tiered_pricing.py @@ -5,17 +5,30 @@ import pytest from inline_snapshot import snapshot -from genai_prices import Usage, calc_price, types +from genai_prices import Usage, calc_price, data_snapshot, types from genai_prices.data_snapshot import DataSnapshot -from genai_prices.update_prices import UpdatePrices pytestmark = pytest.mark.anyio -class MultiTierUpdatePrices(UpdatePrices): - """Custom UpdatePrices that injects a multi-tier pricing model for testing.""" +class MultiTierUpdatePrices: + """Test helper that installs a multi-tier pricing snapshot for the duration of a `with` block. - def fetch(self) -> DataSnapshot | None: + These tests only need a custom snapshot to be active, not a real background updater thread, so + this sets the snapshot directly rather than going through the (deprecated) UpdatePrices class. + """ + + def __enter__(self) -> MultiTierUpdatePrices: + data_snapshot.set_custom_snapshot(self._build_snapshot()) + return self + + def wait(self, _timeout: float | None = None) -> bool: + return True + + def __exit__(self, *_args: object) -> None: + data_snapshot.set_custom_snapshot(None) + + def _build_snapshot(self) -> DataSnapshot: """Create a mock provider with a multi-tier pricing model. Pricing structure (threshold-based with different tiers for input vs output): diff --git a/tests/test_update_prices.py b/tests/test_update_prices.py index 4d92f526..265ecf30 100644 --- a/tests/test_update_prices.py +++ b/tests/test_update_prices.py @@ -1,5 +1,7 @@ +import os import threading from decimal import Decimal +from time import monotonic import httpx2 import pytest @@ -10,9 +12,11 @@ Usage, calc_price, data_snapshot, + update_prices as update_prices_module, wait_prices_updated_async, wait_prices_updated_sync, ) +from genai_prices.update_prices import DEFAULT_UPDATE_URL pytestmark = pytest.mark.anyio @@ -34,7 +38,7 @@ def raise_for_status(self) -> None: def fake_get(url: str, timeout: httpx2.Timeout) -> Response: assert url in { 'https://example.test/prices.json', - 'https://raw.githubusercontent.com/pydantic/genai-prices/refs/heads/main/prices/data.json', + DEFAULT_UPDATE_URL, } assert timeout is not None return Response(content) @@ -72,40 +76,132 @@ def test_update_prices_wait_on_start(monkeypatch: pytest.MonkeyPatch): assert price.total_price == snapshot(Decimal('0.0035')) assert price.provider.id == snapshot('openai') assert price.auto_update_timestamp is not None + assert data_snapshot._custom_snapshot is None def test_wait_prices_updated_sync(monkeypatch: pytest.MonkeyPatch): _mock_update_prices_get(monkeypatch) assert data_snapshot._custom_snapshot is None with UpdatePrices(): - wait_prices_updated_sync() + assert wait_prices_updated_sync(timeout=5) assert data_snapshot._custom_snapshot is not None - price = calc_price(Usage(input_tokens=1000, output_tokens=100), model_ref='gpt-4o', provider_id='openai') - assert price.input_price == snapshot(Decimal('0.0025')) - assert price.output_price == snapshot(Decimal('0.001')) - assert price.total_price == snapshot(Decimal('0.0035')) - assert price.provider.id == snapshot('openai') - assert price.auto_update_timestamp is not None + assert data_snapshot._custom_snapshot is None async def test_wait_prices_updated_async(monkeypatch: pytest.MonkeyPatch): _mock_update_prices_get(monkeypatch) assert data_snapshot._custom_snapshot is None with UpdatePrices(): - await wait_prices_updated_async() + assert await wait_prices_updated_async(timeout=5) assert data_snapshot._custom_snapshot is not None - price = calc_price(Usage(input_tokens=1000, output_tokens=100), model_ref='gpt-4o', provider_id='openai') - assert price.input_price == snapshot(Decimal('0.0025')) - assert price.output_price == snapshot(Decimal('0.001')) - assert price.total_price == snapshot(Decimal('0.0035')) - assert price.provider.id == snapshot('openai') - assert price.auto_update_timestamp is not None + assert data_snapshot._custom_snapshot is None + + +def test_ref_counted_across_instances(monkeypatch: pytest.MonkeyPatch): + _mock_update_prices_get(monkeypatch) + up1 = UpdatePrices() + up2 = UpdatePrices() + up1.start() + up2.start() + try: + assert wait_prices_updated_sync(timeout=5) + assert data_snapshot._custom_snapshot is not None + + # Releasing one claim keeps the updater alive for the other. + up1.stop() + assert wait_prices_updated_sync(timeout=0) + assert data_snapshot._custom_snapshot is not None + + up1.stop() # idempotent no-op + assert data_snapshot._custom_snapshot is not None + + up2.stop() + assert data_snapshot._custom_snapshot is None + assert not wait_prices_updated_sync(timeout=0) + finally: + up1.stop() + up2.stop() + data_snapshot.set_custom_snapshot(None) + + +def test_first_wins_config(monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture): + _mock_update_prices_get(monkeypatch) + up1 = UpdatePrices(url='https://example.test/prices.json') + up1.start() + try: + with caplog.at_level('WARNING', logger='genai-prices'): + up2 = UpdatePrices(url=DEFAULT_UPDATE_URL, update_interval=1) + up2.start() + try: + # The first instance's configuration wins; the second joins the running updater. + assert update_prices_module._updater is not None + assert update_prices_module._updater.url == 'https://example.test/prices.json' + assert update_prices_module._updater.update_interval == 3600 + assert any('already running' in record.message for record in caplog.records) + assert wait_prices_updated_sync(timeout=5) + finally: + up2.stop() + finally: + up1.stop() + data_snapshot.set_custom_snapshot(None) + assert data_snapshot._custom_snapshot is None + + +def test_double_start_same_instance_raises(monkeypatch: pytest.MonkeyPatch): + _mock_update_prices_get(monkeypatch) + update_prices = UpdatePrices() + update_prices.start() + try: + with pytest.raises(RuntimeError, match='already started'): + update_prices.start() + finally: + update_prices.stop() + data_snapshot.set_custom_snapshot(None) + + +def test_stop_on_unstarted_instance_is_noop(monkeypatch: pytest.MonkeyPatch): + _mock_update_prices_get(monkeypatch) + with UpdatePrices() as update_prices: + assert update_prices.wait(timeout=5) + # stop() on a never-started instance does not touch the live updater. + UpdatePrices().stop() + assert update_prices_module._updater is not None + assert data_snapshot._custom_snapshot is not None + assert data_snapshot._custom_snapshot is None -def test_update_prices_stop_clears_snapshot_after_in_flight_fetch(monkeypatch: pytest.MonkeyPatch) -> None: +def test_restopping_does_not_affect_a_new_updater(monkeypatch: pytest.MonkeyPatch): + _mock_update_prices_get(monkeypatch) + up1 = UpdatePrices() + up1.start() + assert wait_prices_updated_sync(timeout=5) + stale_updater = update_prices_module._updater + up1.stop() + assert data_snapshot._custom_snapshot is None + + up2 = UpdatePrices() + up2.start() + try: + assert wait_prices_updated_sync(timeout=5) + new_updater = update_prices_module._updater + assert new_updater is not None and new_updater is not stale_updater + + # up1 was already released; stopping it again must not release up2's claim. + up1.stop() + assert update_prices_module._updater is new_updater + assert update_prices_module._ref_count == 1 + assert data_snapshot._custom_snapshot is not None + finally: + up2.stop() + data_snapshot.set_custom_snapshot(None) + assert data_snapshot._custom_snapshot is None + + +def test_stop_does_not_block_on_in_flight_fetch_and_discards_result( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +): fetch_started = threading.Event() allow_fetch_return = threading.Event() - stop_errors: list[BaseException] = [] class Response: content = PROVIDER_ARRAY_PAYLOAD @@ -114,33 +210,75 @@ def raise_for_status(self) -> None: pass def fake_get(url: str, timeout: httpx2.Timeout) -> Response: - assert url == 'https://example.test/prices.json' + assert url == DEFAULT_UPDATE_URL assert timeout is not None fetch_started.set() assert allow_fetch_return.wait(timeout=5) return Response() - def stop_update_prices(update_prices: UpdatePrices) -> None: - try: + monkeypatch.setattr(httpx2, 'get', fake_get) + monkeypatch.setattr(update_prices_module, '_STOPPED_THREAD_JOIN_TIMEOUT', 0.05) + + update_prices = UpdatePrices() + update_prices.start() + try: + assert fetch_started.wait(timeout=5) + (thread,) = [t for t in threading.enumerate() if t.name == 'genai_prices:update'] + + start = monotonic() + with caplog.at_level('WARNING', logger='genai-prices'): update_prices.stop() - except BaseException as exc: - stop_errors.append(exc) + assert monotonic() - start < 2 # did not wait for the in-flight fetch + assert any('abandoning the daemon thread' in record.message for record in caplog.records) + assert data_snapshot._custom_snapshot is None + + # The drained fetch completes but its snapshot is discarded by the fencing check. + allow_fetch_return.set() + thread.join(timeout=5) + assert not thread.is_alive() + assert data_snapshot._custom_snapshot is None + finally: + allow_fetch_return.set() + update_prices.stop() + data_snapshot.set_custom_snapshot(None) + + +def test_waiter_gets_false_when_stop_discards_in_flight_fetch(monkeypatch: pytest.MonkeyPatch): + # A fetch discarded by the stop fence must not be reported to waiters as a successful update. + fetch_started = threading.Event() + allow_fetch_return = threading.Event() + + class Response: + content = PROVIDER_ARRAY_PAYLOAD + + def raise_for_status(self) -> None: + pass + + def fake_get(url: str, timeout: httpx2.Timeout) -> Response: + assert url == DEFAULT_UPDATE_URL + assert timeout is not None + fetch_started.set() + assert allow_fetch_return.wait(timeout=5) + return Response() monkeypatch.setattr(httpx2, 'get', fake_get) - update_prices = UpdatePrices(url='https://example.test/prices.json', update_interval=3600) + monkeypatch.setattr(update_prices_module, '_STOPPED_THREAD_JOIN_TIMEOUT', 0.05) + + update_prices = UpdatePrices() update_prices.start() try: assert fetch_started.wait(timeout=5) + # A waiter that captured the updater before stop() observes this instance directly. + updater = update_prices_module._updater + assert updater is not None - stop_thread = threading.Thread(target=stop_update_prices, args=(update_prices,)) - stop_thread.start() - assert update_prices._stop_event.wait(timeout=5) - allow_fetch_return.set() - stop_thread.join(timeout=5) + update_prices.stop() # sets the stop fence and reverts before the fetch returns + allow_fetch_return.set() # the in-flight fetch now completes but must be discarded - assert not stop_thread.is_alive() - if stop_errors: - raise stop_errors[0] + # The background loop signals the event but must report the discarded fetch as not-updated. + assert updater._prices_updated.wait(timeout=5) + assert updater._update_succeeded is False + assert updater._wait_updated(0) is False assert data_snapshot._custom_snapshot is None finally: allow_fetch_return.set() @@ -150,31 +288,123 @@ def stop_update_prices(update_prices: UpdatePrices) -> None: @pytest.mark.default_cassette('fail.yaml') @pytest.mark.vcr() -def test_update_prices_failed(): +def test_wait_raises_on_failed_fetch(): assert data_snapshot._custom_snapshot is None with UpdatePrices(url='https://demo-endpoints.pydantic.workers.dev/bin?status=404') as update_prices: with pytest.raises(httpx2.HTTPStatusError): - update_prices.wait() + update_prices.wait(timeout=5) assert data_snapshot._custom_snapshot is None -@pytest.mark.default_cassette('fail.yaml') -@pytest.mark.vcr() -def test_update_prices_failed_stop(): - assert data_snapshot._custom_snapshot is None - update_prices = UpdatePrices(url='https://demo-endpoints.pydantic.workers.dev/bin?status=404') +def test_stop_logs_and_does_not_raise_after_failed_fetch( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +): + def fake_get(*_args: object, **_kwargs: object) -> object: + raise httpx2.ConnectError('network down') + + monkeypatch.setattr(httpx2, 'get', fake_get) + update_prices = UpdatePrices() + update_prices.start() + try: + # Wait on the internal event rather than wait(), which would consume and raise the stored + # exception that stop() is expected to log instead. + updater = update_prices_module._updater + assert updater is not None + assert updater._prices_updated.wait(timeout=5) + finally: + with caplog.at_level('ERROR', logger='genai-prices'): + update_prices.stop() # never raises + data_snapshot.set_custom_snapshot(None) + assert any('while stopping' in record.message for record in caplog.records) + + +def test_wait_prices_updated_sync_returns_false_on_failed_fetch( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +): + def fake_get(*_args: object, **_kwargs: object) -> object: + raise httpx2.ConnectError('network down') + + monkeypatch.setattr(httpx2, 'get', fake_get) + update_prices = UpdatePrices() update_prices.start() - with pytest.raises(httpx2.HTTPStatusError): + try: + # Never raises and returns False on failure, for every waiter — not just the first. + assert wait_prices_updated_sync(timeout=5) is False + assert wait_prices_updated_sync(timeout=0) is False + assert data_snapshot._custom_snapshot is None + finally: + with caplog.at_level('ERROR', logger='genai-prices'): + update_prices.stop() + data_snapshot.set_custom_snapshot(None) + # The stored exception is not consumed by the waiters above, so stop() still logs it. + assert any('while stopping' in record.message for record in caplog.records) + + +@pytest.mark.skipif(not hasattr(os, 'fork'), reason='requires os.fork') +def test_forked_child_restarts_shared_updater(monkeypatch: pytest.MonkeyPatch): + _mock_update_prices_get(monkeypatch) + update_prices = UpdatePrices() + update_prices.start() + try: + assert wait_prices_updated_sync(timeout=5) + + pid = os.fork() + if pid == 0: + # Child: never return into pytest - report via the exit code. + try: + ok = ( + wait_prices_updated_sync(timeout=5) + and any(t.name == 'genai_prices:update' and t.is_alive() for t in threading.enumerate()) + and data_snapshot._custom_snapshot is not None + and update_prices_module._updater is not None + ) + os._exit(0 if ok else 1) + except BaseException: + os._exit(2) + + _, status = os.waitpid(pid, 0) + assert os.waitstatus_to_exitcode(status) == 0 + finally: update_prices.stop() - assert data_snapshot._custom_snapshot is None + data_snapshot.set_custom_snapshot(None) -def test_update_prices_multiple(monkeypatch: pytest.MonkeyPatch): +def test_concurrent_start_stop(monkeypatch: pytest.MonkeyPatch): _mock_update_prices_get(monkeypatch) - with UpdatePrices(): - with pytest.raises( - RuntimeError, - match='UpdatePrices global task already started, only one UpdatePrices can be active at a time', - ): - with UpdatePrices(): - pass + barrier = threading.Barrier(16) + errors: list[BaseException] = [] + + def worker() -> None: + try: + barrier.wait(timeout=5) + for _ in range(10): + update_prices = UpdatePrices() + update_prices.start() + update_prices.stop() + except BaseException as exc: # pragma: no cover + errors.append(exc) + + threads = [threading.Thread(target=worker) for _ in range(16)] + try: + for thread in threads: + thread.start() + finally: + for thread in threads: + thread.join(timeout=30) + + assert not [t for t in threads if t.is_alive()] + assert not errors + assert update_prices_module._updater is None + assert update_prices_module._ref_count == 0 + assert not [t for t in threading.enumerate() if t.name == 'genai_prices:update'] + assert data_snapshot._custom_snapshot is None + + # The shared updater can still be acquired cleanly after the churn. + update_prices = UpdatePrices() + update_prices.start() + try: + assert wait_prices_updated_sync(timeout=5) + assert data_snapshot._custom_snapshot is not None + finally: + update_prices.stop() + data_snapshot.set_custom_snapshot(None)