From 8c0cc39ad1775d05ea0e261d3877d5b879d87da0 Mon Sep 17 00:00:00 2001 From: NanPan <111261006+poijygfdyy@users.noreply.github.com> Date: Wed, 7 Oct 2026 08:18:55 +0800 Subject: [PATCH] fix(wechat): drain poll-state persistence across cancellation (#6389) * fix(wechat): retain cursor persistence across cancellation * fix(wechat): drain expired auth-state persistence * fix(wechat): drain the complete token expiry transition --------- Co-authored-by: Willem Jiang --- README.md | 1 + backend/app/channels/AGENTS.md | 6 +- backend/app/channels/wechat.py | 21 ++- ..._wechat_cursor_persistence_cancellation.py | 143 ++++++++++++++++++ 4 files changed, 160 insertions(+), 11 deletions(-) create mode 100644 backend/tests/test_wechat_cursor_persistence_cancellation.py diff --git a/README.md b/README.md index f4a5191633..4202ddf8d5 100644 --- a/README.md +++ b/README.md @@ -970,6 +970,7 @@ DINGTALK_CLIENT_SECRET=your_client_secret 5. For Docker Compose deployments, keep `state_dir` on a persistent volume so the `get_updates_buf` cursor and saved auth state survive restarts. 6. Outbound images/files enforce `max_outbound_image_bytes` / `max_outbound_file_bytes` (20 MiB / 50 MiB defaults) while reading, including files that grow after resolution. Oversize reads are rejected before encryption/upload instead of sending a truncated prefix. Non-positive limits disable the corresponding cap. 7. `allowed_users` takes iLink user IDs. Unset, `null`, `[]`, or a blank string allows everyone. A single ID is one entry, not a sequence of characters, and an unquoted integer-valued number is stored as that integer's text. A scalar string containing commas or interior whitespace logs a warning but remains one literal ID; use a YAML list for multiple IDs. Invalid entries are ignored with a warning; any other configured value that yields no valid ID denies every user and logs an error. `/connect` is still accepted before that check, and a denied sender is dropped before inbound media is downloaded. +8. Shutdown waits for in-flight cursor writes. On token expiry, DeerFlow persists the cursor reset and removes the saved token before completing poller cancellation. **WeCom Setup** diff --git a/backend/app/channels/AGENTS.md b/backend/app/channels/AGENTS.md index 03b53dbacd..716ea8aeb3 100644 --- a/backend/app/channels/AGENTS.md +++ b/backend/app/channels/AGENTS.md @@ -2,8 +2,8 @@ Bridges external messaging platforms (Feishu, Slack, Telegram, Discord, DingTalk, GitHub) to the DeerFlow agent via Gateway's LangGraph-compatible API. -WeChat's off-loop `_read_outbound_bytes` owns cap + 1 reads and rejects overflow -with `None` before encryption/upload. Non-positive caps stay unlimited. +WeChat `_read_outbound_bytes` reads cap+1 off-loop; overflow returns `None` before encryption/upload; caps <=0 are unlimited. +Cursor writes drain before cancellation; token expiry drains cursor/auth clearing under `_auth_lock`. **Architecture**: Channels communicate with Gateway through the `langgraph-sdk` HTTP client (same as the frontend), ensuring threads are created and managed server-side. The internal SDK client injects process-local internal auth plus a matching CSRF cookie/header pair so Gateway accepts state-changing thread/run requests from channel workers without relying on browser session cookies. @@ -62,7 +62,7 @@ The cached value is reused for both the blocking (`runs.wait`) and streaming (`_ - No public IP, OAuth callback URL, or provider webhook route is required by the current implementation. - Telegram binds via deep-link `/start ` over long polling; Slack, Discord, Feishu/Lark, DingTalk, WeChat, and WeCom use `/connect ` through their outbound workers. - WeChat `polling_timeout`, `polling_retry_delay`, `qrcode_poll_interval`, and `qrcode_poll_timeout` require positive finite seconds; invalid values use defaults to prevent hot loops or endless sleeps. -- `allowed_users.py` shares WeChat/Telegram fail-closed parsing with channel-specific ID coercion: unset/null/empty collections/blank string allow all; invalid entries warn, all-invalid values deny all and log an error. WeChat stores integer-valued numbers as integer text, warns on scalar commas/interior whitespace without splitting, and rejects denied senders before media downloads. +- `allowed_users.py` uses WeChat/Telegram ID coercers: unset/null/empty collections/blank strings allow all; invalid entries warn; all-invalid input denies all with an error. WeChat stores integer-valued numbers as integer text, warns without splitting scalar separators, and gates users before media downloads. - WeCom serializes `start()` and `stop()` for each channel instance. The SDK `connect()` task covers connection setup only; after the handshake, the SDK owns a separate receive task. Shutdown cancels an in-progress connection attempt and awaits the SDK's actual asynchronous receive-task/socket cleanup before releasing lifecycle state or allowing a restart. Cancellation of `stop()` still propagates, but only after owned cleanup finishes and lifecycle references are cleared; real connection failures remain reported by `_on_ws_task_done`. - WeCom outbound content is capped at the protocol's 20480 UTF-8 bytes: stream replies clip on a character boundary with a truncation marker (one stream carries the whole reply and cannot split mid-way), while proactive pushes split into at most 10 sequential markdown messages per push, with the remaining tail clipped and marker-terminated. Both paths measure bytes, not characters. A per-chat send lock serializes each split batch end to end, because manager workers run concurrently and two long pushes to the same chat would otherwise interleave chunks; locks are reference-counted and reclaimed once no sender holds or waits on them, so the registry does not grow for the life of the Gateway. - Frontend APIs: `GET /api/channels/providers`, `GET /api/channels/connections`, `POST /api/channels/{provider}/connect`, and `DELETE /api/channels/connections/{connection_id}`. diff --git a/backend/app/channels/wechat.py b/backend/app/channels/wechat.py index f6dd43118a..6a1e038cab 100644 --- a/backend/app/channels/wechat.py +++ b/backend/app/channels/wechat.py @@ -652,12 +652,7 @@ class WechatChannel(Channel): if ret not in (0, None): errcode = data.get("errcode") if errcode == -14: - self._bot_token = "" - self._get_updates_buf = "" - await asyncio.to_thread(self._save_state) - await asyncio.to_thread(self._save_auth_state, status="expired", bot_token="") - logger.error("[WeChat] bot token expired; scan again or update bot_token and restart the channel") - self._running = False + await await_drained(self._persist_expired_auth()) break logger.warning( "[WeChat] getupdates returned ret=%s errcode=%s errmsg=%s", @@ -693,13 +688,23 @@ class WechatChannel(Channel): next_buf = data.get("get_updates_buf") if isinstance(next_buf, str) and next_buf != self._get_updates_buf: self._get_updates_buf = next_buf - await asyncio.to_thread(self._save_state) + await await_drained(asyncio.to_thread(self._save_state)) except asyncio.CancelledError: raise except Exception: logger.exception("[WeChat] polling loop failed") await asyncio.sleep(self._retry_delay) + async def _persist_expired_auth(self) -> None: + """Persist both expiry writes under the auth lock; the poller drains this whole operation.""" + async with self._auth_lock: + self._bot_token = "" + self._get_updates_buf = "" + await asyncio.to_thread(self._save_state) + await self._save_auth_state_drained(status="expired", bot_token="") + logger.error("[WeChat] bot token expired; scan again or update bot_token and restart the channel") + self._running = False + async def _handle_update(self, raw_message: Any) -> None: if not isinstance(raw_message, dict): return @@ -843,7 +848,7 @@ class WechatChannel(Channel): qrcode: str | None = None, qrcode_img_content: str | None = None, ) -> dict[str, Any]: - """Persist QR auth state before propagating caller cancellation. + """Persist auth state before propagating caller cancellation. The write is a small local state-file update and is intentionally drained while _auth_lock is held so cancellation cannot expose an in-memory diff --git a/backend/tests/test_wechat_cursor_persistence_cancellation.py b/backend/tests/test_wechat_cursor_persistence_cancellation.py new file mode 100644 index 0000000000..1b335966b2 --- /dev/null +++ b/backend/tests/test_wechat_cursor_persistence_cancellation.py @@ -0,0 +1,143 @@ +from __future__ import annotations + +import asyncio +import json +import threading +from pathlib import Path + +import pytest + +from app.channels.message_bus import MessageBus +from app.channels.wechat import WechatChannel + + +@pytest.mark.asyncio +async def test_poll_cursor_persistence_drains_across_cancellation(monkeypatch, tmp_path: Path) -> None: + channel = WechatChannel( + MessageBus(), + config={"bot_token": "test-token", "state_dir": str(tmp_path)}, + ) + channel._running = True + + started = threading.Event() + release = threading.Event() + original_save_state = channel._save_state + requests = 0 + + async def ensure_authenticated() -> bool: + return True + + async def request_json(*_args, **_kwargs): + nonlocal requests + requests += 1 + if requests == 1: + return {"ret": 0, "msgs": [], "get_updates_buf": "cursor-next"} + await asyncio.Future() + + def blocking_save_state() -> None: + started.set() + assert release.wait(timeout=5) + original_save_state() + + monkeypatch.setattr(channel, "_ensure_authenticated", ensure_authenticated) + monkeypatch.setattr(channel, "_request_json", request_json) + monkeypatch.setattr(channel, "_save_state", blocking_save_state) + + task = asyncio.create_task(channel._poll_loop()) + try: + assert await asyncio.to_thread(started.wait, 5) + + task.cancel() + await asyncio.sleep(0) + task.cancel() + await asyncio.sleep(0) + assert not task.done() + + release.set() + with pytest.raises(asyncio.CancelledError): + await task + finally: + release.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + persisted = json.loads((tmp_path / "wechat-getupdates.json").read_text(encoding="utf-8")) + assert persisted["get_updates_buf"] == "cursor-next" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("blocked_write", [None, "cursor", "auth"], ids=["no-cancellation", "cancel-during-cursor", "cancel-during-auth"]) +async def test_poll_token_expiry_finishes_both_writes_before_cancellation(monkeypatch, tmp_path: Path, caplog, blocked_write: str | None) -> None: + channel = WechatChannel(MessageBus(), config={"bot_token": "expired-token", "state_dir": str(tmp_path)}) + channel._get_updates_buf = "cursor-before-expiry" + await asyncio.to_thread(channel._save_state) + await asyncio.to_thread(channel._save_auth_state, status="confirmed", bot_token="expired-token") + channel._running = True + + started = threading.Event() + release = threading.Event() + original_save_state = channel._save_state + original_save_auth_state = channel._save_auth_state + auth_lock_observations = [] + + async def ensure_authenticated() -> bool: + return True + + async def request_json(*_args, **_kwargs): + return {"ret": 1, "errcode": -14} + + def save_state() -> None: + if blocked_write == "cursor": + started.set() + assert release.wait(timeout=5) + original_save_state() + + def save_auth_state(**kwargs): + auth_lock_observations.append(channel._auth_lock.locked()) + if blocked_write == "auth": + started.set() + assert release.wait(timeout=5) + return original_save_auth_state(**kwargs) + + monkeypatch.setattr(channel, "_ensure_authenticated", ensure_authenticated) + monkeypatch.setattr(channel, "_request_json", request_json) + monkeypatch.setattr(channel, "_save_state", save_state) + monkeypatch.setattr(channel, "_save_auth_state", save_auth_state) + + task = asyncio.create_task(channel._poll_loop()) + try: + if blocked_write is None: + await asyncio.wait_for(task, timeout=5) + else: + assert await asyncio.to_thread(started.wait, 5) + task.cancel() + await asyncio.sleep(0) + task.cancel() + await asyncio.sleep(0) + assert not task.done() + + release.set() + with pytest.raises(asyncio.CancelledError): + await task + finally: + release.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + cursor = json.loads((tmp_path / "wechat-getupdates.json").read_text(encoding="utf-8")) + auth = json.loads((tmp_path / "wechat-auth.json").read_text(encoding="utf-8")) + assert cursor["get_updates_buf"] == "" + assert auth["status"] == "expired" + assert "bot_token" not in auth + assert channel._bot_token == "" + assert not channel._running + assert "bot token expired" in caplog.text + assert auth_lock_observations == [True] + assert not channel._auth_lock.locked() + + restored = WechatChannel(MessageBus(), config={"state_dir": str(tmp_path)}) + await asyncio.to_thread(restored._load_state) + assert restored._bot_token == "" + assert restored._get_updates_buf == ""