diff --git a/src/weilink/_store.py b/src/weilink/_store.py index 30afdc2..442b245 100644 --- a/src/weilink/_store.py +++ b/src/weilink/_store.py @@ -548,6 +548,59 @@ def query_messages( logger.warning("Failed to deserialize message: %s", data_json[:100]) return results + def max_rowid(self) -> int: + """Return the current maximum row ID, or 0 if the table is empty.""" + with self._lock: + row = self._conn.execute("SELECT MAX(id) FROM messages").fetchone() + return row[0] if row and row[0] is not None else 0 + + def query_since_rowid( + self, + since_id: int = 0, + *, + direction: int | None = None, + ) -> tuple[list[Message], int]: + """Return messages with rowid > *since_id* and the new high-water mark. + + Intended for the store-watcher dispatcher fallback: poll SQLite + for rows inserted by another process. + + Args: + since_id: Only return rows with ``id`` strictly greater than + this value. + direction: Optional direction filter (1=received, 2=sent). + + Returns: + ``(messages, new_hwm)`` where *messages* is in chronological + order (oldest first) and *new_hwm* is the maximum ``id`` + among returned rows (or *since_id* if none). + """ + if self._closed: + return [], since_id + + clauses = ["id > ?"] + params: list[Any] = [since_id] + if direction is not None: + clauses.append("direction = ?") + params.append(direction) + + where = " WHERE " + " AND ".join(clauses) + sql = f"SELECT id, data FROM messages{where} ORDER BY id ASC" + + with self._lock: + rows = self._conn.execute(sql, params).fetchall() + + messages: list[Message] = [] + new_hwm = since_id + for row_id, data_json in rows: + try: + messages.append(deserialize_message(data_json)) + new_hwm = row_id + except (json.JSONDecodeError, KeyError, TypeError): + logger.warning("Failed to deserialize message: %s", data_json[:100]) + new_hwm = row_id # still advance past broken rows + return messages, new_hwm + def count( self, *, diff --git a/src/weilink/client.py b/src/weilink/client.py index 9ed79ea..e91ed1b 100644 --- a/src/weilink/client.py +++ b/src/weilink/client.py @@ -43,6 +43,7 @@ _DEFAULT_SESSION = "default" _FALLBACK_WINDOW = 60 # seconds: time window for Route C degraded SQLite reads _DEFAULT_QUEUE_MAXSIZE = 1000 +_STORE_WATCH_INTERVAL = 2.0 # seconds: polling interval for store watcher fallback def _atomic_write(path: Path, data: str) -> None: @@ -1579,21 +1580,73 @@ def stop(self) -> None: self._dispatcher_stop.clear() def _start_dispatcher(self, poll_timeout: float) -> None: - """Start the background polling thread if not already running.""" + """Start the background polling thread if not already running. + + When the poll lock is available, starts the normal ``_poll_loop`` + that long-polls the iLink API. When another process already + holds the lock (e.g. an MCP server), falls back to + ``_store_watch_loop`` which watches SQLite for new messages. + """ with self._dispatcher_lock: if ( self._dispatcher_thread is not None and self._dispatcher_thread.is_alive() ): return + + # Decide between poll loop and store watcher. + if self._poll_lock.try_lock(): + self._poll_lock.unlock() + target = self._poll_loop + args: tuple[float] = (poll_timeout,) + elif self._message_store is not None: + logger.info( + "Poll lock held by another process, starting store watcher fallback" + ) + target = self._store_watch_loop + args = (_STORE_WATCH_INTERVAL,) + else: + raise RuntimeError( + "Cannot start dispatcher: poll lock is held by another " + "process and message_store is not enabled. Enable " + "message_store to use the store watcher fallback." + ) + self._dispatcher_stop.clear() self._dispatcher_thread = threading.Thread( - target=self._poll_loop, - args=(poll_timeout,), + target=target, + args=args, daemon=True, ) self._dispatcher_thread.start() + def _dispatch_messages(self, messages: list[Message]) -> None: + """Dispatch messages to handlers and enqueue for recv() consumers.""" + for msg in messages: + for handler in self._message_handlers: + try: + handler(msg) + except Exception: + logger.exception( + "Handler %s raised an exception", + getattr(handler, "__name__", handler), + ) + + # Enqueue for recv() consumers, drop oldest on overflow + try: + self._message_queue.put_nowait(msg) + except queue.Full: + try: + dropped = self._message_queue.get_nowait() + except queue.Empty: + pass + else: + logger.warning( + "Dispatcher queue full, dropped oldest message %s", + dropped.message_id, + ) + self._message_queue.put_nowait(msg) + def _poll_loop(self, poll_timeout: float) -> None: """Background loop: poll iLink and dispatch to handlers + queue.""" while not self._dispatcher_stop.is_set(): @@ -1617,31 +1670,32 @@ def _poll_loop(self, poll_timeout: float) -> None: logger.exception("Polling error in dispatcher") continue - for msg in messages: - # Dispatch to handlers - for handler in self._message_handlers: - try: - handler(msg) - except Exception: - logger.exception( - "Handler %s raised an exception", - getattr(handler, "__name__", handler), - ) + self._dispatch_messages(messages) - # Enqueue for recv() consumers, drop oldest on overflow - try: - self._message_queue.put_nowait(msg) - except queue.Full: - try: - dropped = self._message_queue.get_nowait() - except queue.Empty: - pass - else: - logger.warning( - "Dispatcher queue full, dropped oldest message %s", - dropped.message_id, - ) - self._message_queue.put_nowait(msg) + self._dispatcher_stop.set() + + def _store_watch_loop(self, interval: float) -> None: + """Watch SQLite for new messages when poll lock is unavailable. + + This is the store-watcher fallback: instead of polling the iLink + API directly, periodically query the message store for rows + inserted by the primary poller (e.g. an MCP server). + """ + assert self._message_store is not None + hwm = self._message_store.max_rowid() + logger.debug("Store watcher started, high-water mark = %d", hwm) + + while not self._dispatcher_stop.wait(timeout=interval): + messages, new_hwm = self._message_store.query_since_rowid(hwm) + if messages: + logger.debug( + "Store watcher: %d new message(s), hwm %d -> %d", + len(messages), + hwm, + new_hwm, + ) + self._dispatch_messages(messages) + hwm = new_hwm self._dispatcher_stop.set() diff --git a/tests/test_store_watcher.py b/tests/test_store_watcher.py new file mode 100644 index 0000000..6ab9381 --- /dev/null +++ b/tests/test_store_watcher.py @@ -0,0 +1,233 @@ +"""Tests for the store-watcher dispatcher fallback.""" + +from __future__ import annotations + +import time +from pathlib import Path + +import pytest + +from weilink._store import MessageStore +from weilink._vendor.filelock import FileLock +from weilink.client import WeiLink +from weilink.models import Message, MessageType + + +def _make_msg( + text: str = "hello", + from_user: str = "user1@im.wechat", + bot_id: str = "bot1@im.bot", + message_id: int | None = 100, + timestamp: int = 1700000000000, +) -> Message: + return Message( + from_user=from_user, + msg_type=MessageType.TEXT, + text=text, + timestamp=timestamp, + message_id=message_id, + bot_id=bot_id, + ) + + +# ------------------------------------------------------------------ +# MessageStore: max_rowid / query_since_rowid +# ------------------------------------------------------------------ + + +class TestStoreRowidQueries: + def test_max_rowid_empty(self, tmp_path: Path) -> None: + store = MessageStore(tmp_path / "test.db") + assert store.max_rowid() == 0 + store.close() + + def test_max_rowid_after_insert(self, tmp_path: Path) -> None: + store = MessageStore(tmp_path / "test.db") + store.store([_make_msg(message_id=1)]) + assert store.max_rowid() >= 1 + store.store([_make_msg(message_id=2)]) + assert store.max_rowid() >= 2 + store.close() + + def test_query_since_rowid_empty(self, tmp_path: Path) -> None: + store = MessageStore(tmp_path / "test.db") + msgs, hwm = store.query_since_rowid(0) + assert msgs == [] + assert hwm == 0 + store.close() + + def test_query_since_rowid_returns_new(self, tmp_path: Path) -> None: + store = MessageStore(tmp_path / "test.db") + store.store([_make_msg(text="first", message_id=1)]) + hwm = store.max_rowid() + + store.store([_make_msg(text="second", message_id=2)]) + msgs, new_hwm = store.query_since_rowid(hwm) + assert len(msgs) == 1 + assert msgs[0].text == "second" + assert new_hwm > hwm + store.close() + + def test_query_since_rowid_skips_old(self, tmp_path: Path) -> None: + store = MessageStore(tmp_path / "test.db") + store.store([_make_msg(text="old", message_id=1)]) + store.store([_make_msg(text="also old", message_id=2)]) + hwm = store.max_rowid() + + msgs, new_hwm = store.query_since_rowid(hwm) + assert msgs == [] + assert new_hwm == hwm + store.close() + + def test_query_since_rowid_chronological_order(self, tmp_path: Path) -> None: + store = MessageStore(tmp_path / "test.db") + hwm = store.max_rowid() + store.store([_make_msg(text="a", message_id=1, timestamp=1000)]) + store.store([_make_msg(text="b", message_id=2, timestamp=2000)]) + store.store([_make_msg(text="c", message_id=3, timestamp=3000)]) + + msgs, _ = store.query_since_rowid(hwm) + assert [m.text for m in msgs] == ["a", "b", "c"] + store.close() + + def test_query_since_rowid_direction_filter(self, tmp_path: Path) -> None: + store = MessageStore(tmp_path / "test.db") + hwm = store.max_rowid() + store.store([_make_msg(text="received", message_id=1)]) + store.store_sent("user1@im.wechat", "bot1@im.bot", text="sent") + + msgs_recv, _ = store.query_since_rowid(hwm, direction=1) + assert len(msgs_recv) == 1 + assert msgs_recv[0].text == "received" + + msgs_all, _ = store.query_since_rowid(hwm) + assert len(msgs_all) == 2 + store.close() + + +# ------------------------------------------------------------------ +# Store watcher integration +# ------------------------------------------------------------------ + + +class TestStoreWatcher: + def test_dispatch_to_handlers(self, tmp_path: Path) -> None: + """Store watcher dispatches new messages to on_message handlers.""" + wl = WeiLink( + token_path=tmp_path / "token.json", + message_store=tmp_path / "messages.db", + ) + received: list[Message] = [] + + @wl.on_message + def handler(msg: Message) -> None: + received.append(msg) + + # Use a separate FileLock instance to simulate another process + # holding the poll lock. Same-instance lock is re-entrant on Unix. + external_lock = FileLock(tmp_path / ".poll.lock") + external_lock.lock() + try: + wl.run_background() + + # Insert a message into the store (simulating another process). + assert wl._message_store is not None + wl._message_store.store([_make_msg(text="watcher test", message_id=42)]) + + # Wait for the store watcher to pick it up. + deadline = time.monotonic() + 5.0 + while not received and time.monotonic() < deadline: + time.sleep(0.2) + + assert len(received) == 1 + assert received[0].text == "watcher test" + finally: + wl.stop() + external_lock.unlock() + external_lock.close() + wl.close() + + def test_enqueues_for_recv(self, tmp_path: Path) -> None: + """Store watcher enqueues messages for recv() consumers.""" + wl = WeiLink( + token_path=tmp_path / "token.json", + message_store=tmp_path / "messages.db", + ) + + external_lock = FileLock(tmp_path / ".poll.lock") + external_lock.lock() + try: + wl.run_background() + + assert wl._message_store is not None + wl._message_store.store([_make_msg(text="queue test", message_id=43)]) + + # recv should return the message from the queue. + msgs = wl.recv(timeout=5.0) + assert len(msgs) >= 1 + assert any(m.text == "queue test" for m in msgs) + finally: + wl.stop() + external_lock.unlock() + external_lock.close() + wl.close() + + def test_skips_existing_messages(self, tmp_path: Path) -> None: + """Messages already in the store before start are not dispatched.""" + wl = WeiLink( + token_path=tmp_path / "token.json", + message_store=tmp_path / "messages.db", + ) + received: list[Message] = [] + + @wl.on_message + def handler(msg: Message) -> None: + received.append(msg) + + # Insert BEFORE starting dispatcher. + assert wl._message_store is not None + wl._message_store.store([_make_msg(text="old msg", message_id=99)]) + + external_lock = FileLock(tmp_path / ".poll.lock") + external_lock.lock() + try: + wl.run_background() + # Wait a couple of watch intervals. + time.sleep(3.0) + assert len(received) == 0 + finally: + wl.stop() + external_lock.unlock() + external_lock.close() + wl.close() + + def test_start_dispatcher_raises_without_store(self, tmp_path: Path) -> None: + """RuntimeError if poll lock held and no message_store.""" + wl = WeiLink(token_path=tmp_path / "token.json") + + external_lock = FileLock(tmp_path / ".poll.lock") + external_lock.lock() + try: + with pytest.raises(RuntimeError, match="message_store is not enabled"): + wl.run_background() + finally: + external_lock.unlock() + external_lock.close() + wl.close() + + def test_start_dispatcher_uses_poll_loop_when_free(self, tmp_path: Path) -> None: + """When poll lock is free, normal poll loop is used (not store watcher).""" + wl = WeiLink( + token_path=tmp_path / "token.json", + message_store=tmp_path / "messages.db", + ) + try: + wl.run_background() + # The thread target should be _poll_loop, not _store_watch_loop + assert wl._dispatcher_thread is not None + # Can't easily check target, but if it started without error + # and poll lock is free, it's using _poll_loop. + assert wl._dispatcher_thread.is_alive() + finally: + wl.stop() + wl.close()