From 7f18aa5aad1ed4171aa22d6f64f22baffbe83882 Mon Sep 17 00:00:00 2001
From: L4XB <lukas.buck@e-mail.de>
Date: Wed, 9 Sep 2026 22:31:28 +0200
Subject: [PATCH] fix: keep the newer SessionWaiter registration when an older
 one finishes

register_wait() lets a newer waiter replace the registry entry of the
same session_id, but the older waiter's _cleanup() popped the entry
unconditionally. When the older waiter finished, timed out or was
cancelled, the newer waiter disappeared from USER_SESSIONS while its
future kept waiting, so trigger() could no longer reach it and it could
only time out.

Only remove the registry entry if it still points at the waiter being
cleaned up.

Fixes #9996
---
 astrbot/core/utils/session_waiter.py |  7 ++-
 tests/test_session_waiter.py         | 72 ++++++++++++++++++++++++++++
 2 files changed, 78 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_session_waiter.py

diff --git a/astrbot/core/utils/session_waiter.py b/astrbot/core/utils/session_waiter.py
index b327a61843..9b2292d575 100644
--- a/astrbot/core/utils/session_waiter.py
+++ b/astrbot/core/utils/session_waiter.py
@@ -143,7 +143,12 @@ async def register_wait(
 
     def _cleanup(self, error: Exception | None = None) -> None:
         """清理会话"""
-        USER_SESSIONS.pop(self.session_id, None)
+        # register_wait() lets a newer waiter replace this one under the same
+        # session_id. Only drop the registry entry if it still points at this
+        # waiter, otherwise the newer waiter becomes unreachable for trigger()
+        # and its future can only time out.
+        if USER_SESSIONS.get(self.session_id) is self:
+            USER_SESSIONS.pop(self.session_id, None)
         try:
             FILTERS.remove(self.session_filter)
         except ValueError:
diff --git a/tests/test_session_waiter.py b/tests/test_session_waiter.py
new file mode 100644
index 0000000000..4b3ba05ca8
--- /dev/null
+++ b/tests/test_session_waiter.py
@@ -0,0 +1,72 @@
+import asyncio
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from astrbot.core.utils import session_waiter as sw
+
+
+@pytest.fixture(autouse=True)
+def _clean_registry():
+    sw.USER_SESSIONS.clear()
+    sw.FILTERS.clear()
+    yield
+    sw.USER_SESSIONS.clear()
+    sw.FILTERS.clear()
+
+
+def _waiter(session_id: str) -> sw.SessionWaiter:
+    waiter = sw.SessionWaiter(sw.DefaultSessionFilter(), session_id, False)
+    sw.FILTERS.append(waiter.session_filter)
+    return waiter
+
+
+async def _finish(waiter: sw.SessionWaiter) -> None:
+    """Stop the waiter and let its keep-alive task exit."""
+    waiter.session_controller.stop()
+    event = waiter.session_controller.current_event
+    if event is not None:
+        event.set()
+    await asyncio.sleep(0)
+
+
+@pytest.mark.asyncio
+async def test_finished_waiter_keeps_newer_registration_for_same_session():
+    old, new = _waiter("same"), _waiter("same")
+    handler = AsyncMock()
+
+    task_old = asyncio.create_task(old.register_wait(handler, timeout=10))
+    await asyncio.sleep(0)
+    task_new = asyncio.create_task(new.register_wait(handler, timeout=10))
+    await asyncio.sleep(0)
+    assert sw.USER_SESSIONS["same"] is new
+
+    await _finish(old)
+    await task_old
+
+    # The newer waiter must still be registered and reachable.
+    assert sw.USER_SESSIONS["same"] is new
+    assert not new.session_controller.future.done()
+
+    event = MagicMock()
+    event.get_messages.return_value = []
+    await sw.SessionWaiter.trigger("same", event)
+    handler.assert_awaited_once()
+    assert handler.await_args.args[0] is new.session_controller
+
+    await _finish(new)
+    await task_new
+    assert "same" not in sw.USER_SESSIONS
+
+
+@pytest.mark.asyncio
+async def test_finished_waiter_removes_its_own_registration():
+    waiter = _waiter("solo")
+    task = asyncio.create_task(waiter.register_wait(AsyncMock(), timeout=10))
+    await asyncio.sleep(0)
+    assert sw.USER_SESSIONS["solo"] is waiter
+
+    await _finish(waiter)
+    await task
+    assert "solo" not in sw.USER_SESSIONS
+    assert waiter.session_filter not in sw.FILTERS