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 <willem.jiang@gmail.com>

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**
 

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 <code>` over long polling; Slack, Discord, Feishu/Lark, DingTalk, WeChat, and WeCom use `/connect <code>` 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}`.

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

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 == ""