diff --git a/tests/test_watch_dispatch.py b/tests/test_watch_dispatch.py index 901d2f0c..94759029 100644 --- a/tests/test_watch_dispatch.py +++ b/tests/test_watch_dispatch.py @@ -68,11 +68,6 @@ def _register_runner(session: ChatSession) -> tuple[Any, Any]: return runner, captured["fn"] -def _watch_row(active: bool = True, watch_id: str = "watch-1") -> dict[str, Any]: - """Minimal storage row shape ``valid_until`` predicates re-check.""" - return {"watch_id": watch_id, "active": active} - - # --------------------------------------------------------------------------- # Enqueue shape # --------------------------------------------------------------------------- @@ -239,13 +234,13 @@ class TestValidUntil: from turnstone.core import session as session_mod - # Storage stub returns active=False at drain time. - get_calls: list[str] = [] + # Storage stub returns False at drain time. + is_active_calls: list[str] = [] class _StubStorage: - def get_watch(self, watch_id: str) -> dict[str, Any] | None: - get_calls.append(watch_id) - return _watch_row(active=False, watch_id=watch_id) + def is_watch_active(self, watch_id: str) -> bool: + is_active_calls.append(watch_id) + return False monkeypatch.setattr(session_mod, "get_storage", lambda: _StubStorage()) @@ -254,7 +249,7 @@ class TestValidUntil: out = session._nudge_queue.drain({"any"}) assert out == [] # Predicate ran once. - assert get_calls == ["watch-1"] + assert is_active_calls == ["watch-1"] def test_valid_until_drops_when_watch_missing(self, tmp_db, monkeypatch): session = _make_session_for_dispatch() @@ -263,8 +258,8 @@ class TestValidUntil: from turnstone.core import session as session_mod class _StubStorage: - def get_watch(self, watch_id: str) -> dict[str, Any] | None: - return None # row gone (deleted) + def is_watch_active(self, watch_id: str) -> bool: + return False # row gone (deleted) monkeypatch.setattr(session_mod, "get_storage", lambda: _StubStorage()) @@ -284,7 +279,7 @@ class TestValidUntil: from turnstone.core import session as session_mod class _BoomStorage: - def get_watch(self, watch_id: str) -> dict[str, Any] | None: + def is_watch_active(self, watch_id: str) -> bool: raise RuntimeError("storage down") monkeypatch.setattr(session_mod, "get_storage", lambda: _BoomStorage()) @@ -303,8 +298,8 @@ class TestValidUntil: from turnstone.core import session as session_mod class _StubStorage: - def get_watch(self, watch_id: str) -> dict[str, Any] | None: - return _watch_row(active=True, watch_id=watch_id) + def is_watch_active(self, watch_id: str) -> bool: + return True monkeypatch.setattr(session_mod, "get_storage", lambda: _StubStorage()) @@ -339,8 +334,8 @@ class TestConcurrency: from turnstone.core import session as session_mod class _ActiveStorage: - def get_watch(self, watch_id: str) -> dict[str, Any]: - return _watch_row(active=True, watch_id=watch_id) + def is_watch_active(self, watch_id: str) -> bool: + return True monkeypatch.setattr(session_mod, "get_storage", lambda: _ActiveStorage()) diff --git a/tests/test_watch_integration.py b/tests/test_watch_integration.py index 36d5d214..183d87f6 100644 --- a/tests/test_watch_integration.py +++ b/tests/test_watch_integration.py @@ -56,12 +56,12 @@ def _make_session() -> ChatSession: def _stub_storage(active: bool = True) -> Any: """Minimal storage stub providing the surface the dispatch closure - touches — ``get_watch`` for the ``valid_until`` predicate. + touches — ``is_watch_active`` for the ``valid_until`` predicate. """ class _S: - def get_watch(self, watch_id: str) -> dict[str, Any]: - return {"watch_id": watch_id, "active": active} + def is_watch_active(self, watch_id: str) -> bool: + return active return _S() diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 0f414f56..3ab0bd36 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -1392,7 +1392,7 @@ class ChatSession: :data:`_WATCH_QUEUE_SOFT_CAP` + drop-oldest-on-saturation, standing in for the deleted ``_watch_pending`` maxsize bound. - a ``valid_until`` predicate that re-checks - ``storage.get_watch(watch_id)["active"]`` at drain time so a + ``storage.is_watch_active(watch_id)`` at drain time so a cancelled watch's last splat doesn't ride out a future wake. - producer-side :func:`sanitize_payload` over the whole formatted message so steering-vector / control-char payloads @@ -1427,15 +1427,13 @@ class ChatSession: # Re-checked at drain time outside the queue lock — if # the watch was cancelled between fire and drain, the # entry gets dropped silently rather than splicing a - # stale result onto the user's next turn. + # stale result onto the user's next turn. Single-column + # ``is_watch_active`` avoids the full-row marshal of + # ``get_watch`` on this hot path. try: - storage = get_storage() - row = storage.get_watch(bound_watch_id) + return get_storage().is_watch_active(bound_watch_id) except Exception: return False - if row is None: - return False - return bool(row.get("active", False)) nudge_queue.enqueue("watch_triggered", sanitized, "any", valid_until=_still_active) diff --git a/turnstone/core/storage/_postgresql.py b/turnstone/core/storage/_postgresql.py index add10deb..b040a099 100644 --- a/turnstone/core/storage/_postgresql.py +++ b/turnstone/core/storage/_postgresql.py @@ -1695,6 +1695,16 @@ class PostgreSQLBackend: return None return dict(row._mapping) + def is_watch_active(self, watch_id: str) -> bool: + + with self._conn() as conn: + row = conn.execute( + sa.select(watches.c.active).where(watches.c.watch_id == watch_id) + ).fetchone() + if row is None: + return False + return bool(row[0]) + def list_watches_for_ws(self, ws_id: str) -> list[dict[str, Any]]: with self._conn() as conn: diff --git a/turnstone/core/storage/_protocol.py b/turnstone/core/storage/_protocol.py index a48eea29..fbda5063 100644 --- a/turnstone/core/storage/_protocol.py +++ b/turnstone/core/storage/_protocol.py @@ -928,6 +928,16 @@ class StorageBackend(Protocol): """Return watch dict or None.""" ... + def is_watch_active(self, watch_id: str) -> bool: + """Return True iff the watch exists and its ``active`` flag is set. + + Single-column read for hot paths that only need the active flag + (e.g. the watch-dispatch ``valid_until`` predicate) without + paying for the full row marshal that ``get_watch`` does. + Returns ``False`` if the watch row is missing. + """ + ... + def list_watches_for_ws(self, ws_id: str) -> list[dict[str, Any]]: """Return active watches for a workstream, ordered by created DESC.""" ... diff --git a/turnstone/core/storage/_sqlite.py b/turnstone/core/storage/_sqlite.py index 1adc8fee..fbc44084 100644 --- a/turnstone/core/storage/_sqlite.py +++ b/turnstone/core/storage/_sqlite.py @@ -1842,6 +1842,16 @@ class SQLiteBackend: return None return dict(row._mapping) + def is_watch_active(self, watch_id: str) -> bool: + + with self._conn() as conn: + row = conn.execute( + sa.select(watches.c.active).where(watches.c.watch_id == watch_id) + ).fetchone() + if row is None: + return False + return bool(row[0]) + def list_watches_for_ws(self, ws_id: str) -> list[dict[str, Any]]: with self._conn() as conn: diff --git a/turnstone/core/watch.py b/turnstone/core/watch.py index 981c713e..bf23371d 100644 --- a/turnstone/core/watch.py +++ b/turnstone/core/watch.py @@ -266,7 +266,7 @@ class WatchRunner: The fn signature is ``(message, watch_id)`` — the runner passes the originating ``watch_id`` so dispatch closures can capture per-watch metadata (e.g. a ``valid_until`` predicate that - re-checks ``storage.get_watch(watch_id)["active"]`` before + re-checks ``storage.is_watch_active(watch_id)`` before delivering a stale entry). """ with self._dispatch_lock: