fix(slack): rehydrate thread context after gateway restart

Persistent sessions survive gateway restarts, but thread replies posted
while the gateway was DOWN never reached the session — and the adapter
had no way to notice, so the conversation silently resumed with a hole
in it.

On the first ordinary reply per thread after a restart (tracked by a
fresh-process _thread_rehydration_checked set), fetch the thread delta
past the persisted per-session watermark and inject any missed messages
as part of the new turn via channel_context. Exactly-once per thread
per process; when the watermark is empty (pre-feature sessions) the
check is a no-op. Steady-state replies keep advancing the watermark so
rehydration never re-injects messages the session already carries as
ordinary turns. Prior history is never rewritten (prompt caching safe).

Builds on the persisted watermark introduced for #23918.

Salvaged from #33215 by @vexclawx31, reworked from a repeated
full-thread injection guard into a watermark-delta injection so
rehydration adds only what the session actually missed.
This commit is contained in:
vexclawx31 2026-07-22 04:46:06 -07:00 committed by Teknium
parent ad4034711d
commit fc0009b9ba
2 changed files with 183 additions and 1 deletions

View file

@ -705,6 +705,14 @@ class SlackAdapter(BasePlatformAdapter):
# Cache for _fetch_thread_context results: cache_key → _ThreadContextCache
self._thread_context_cache: Dict[str, _ThreadContextCache] = {}
self._THREAD_CACHE_TTL = 60.0
# Persistent sessions survive gateway restarts, but messages that
# arrived while the gateway was DOWN never reached the session.
# Track which threads have been rehydration-checked this process so
# the first ordinary reply after a restart injects the missed delta
# exactly once (#63530 restart gap / rehydration). Keys follow the
# thread session-key scoping.
self._thread_rehydration_checked: set = set()
self._THREAD_REHYDRATION_CHECKED_MAX = 5000
# Track message IDs that should get reaction lifecycle (DMs / @mentions).
self._reacting_message_ids: set = set()
# Track active Assistant statuses by (team_id, channel_id, thread_ts)
@ -3867,6 +3875,9 @@ class SlackAdapter(BasePlatformAdapter):
watermark_ts=ts,
team_id=team_id,
)
self._mark_thread_rehydration_checked(
channel_id, event_thread_ts, user_id, team_id
)
elif is_thread_reply and has_active_thread_session and is_mentioned:
# Explicit @mention on an active thread is a fresh intent signal:
# the user expects the bot to read the CURRENT thread state, which
@ -3895,6 +3906,60 @@ class SlackAdapter(BasePlatformAdapter):
watermark_ts=ts,
team_id=team_id,
)
self._mark_thread_rehydration_checked(
channel_id, event_thread_ts, user_id, team_id
)
elif is_thread_reply and has_active_thread_session:
# Restart rehydration (#63530 restart gap / #33215): persistent
# sessions survive gateway restarts, but thread replies posted
# while the gateway was down never reached the session. On the
# FIRST ordinary reply per thread in this process, fetch the
# delta past the persisted watermark and inject anything missed
# as part of this new turn. Checked at most once per thread per
# process; a non-empty watermark plus an empty delta costs one
# cached conversations.replies call.
rehydration_key = self._thread_rehydration_key(
channel_id, event_thread_ts, user_id, team_id
)
if rehydration_key not in self._thread_rehydration_checked:
watermark_ts = self._get_thread_watermark(
channel_id=channel_id,
thread_ts=event_thread_ts,
user_id=user_id,
team_id=team_id,
)
if watermark_ts:
thread_context = await self._fetch_thread_context(
channel_id=channel_id,
thread_ts=event_thread_ts,
current_ts=ts,
team_id=team_id,
after_ts=watermark_ts,
force_refresh=True,
)
if thread_context:
channel_context = thread_context
self._set_thread_watermark(
channel_id=channel_id,
thread_ts=event_thread_ts,
user_id=user_id,
watermark_ts=ts,
team_id=team_id,
)
self._mark_thread_rehydration_checked(
channel_id, event_thread_ts, user_id, team_id
)
else:
# Steady state: keep the watermark advancing so a future
# refresh/rehydration never re-injects messages the session
# already carries as ordinary turns.
self._set_thread_watermark(
channel_id=channel_id,
thread_ts=event_thread_ts,
user_id=user_id,
watermark_ts=ts,
team_id=team_id,
)
# Determine message type
msg_type = MessageType.TEXT
@ -5502,6 +5567,46 @@ class SlackAdapter(BasePlatformAdapter):
def _thread_watermark_key(self, channel_id: str, thread_ts: str) -> str:
return f"slack_thread_watermark:{channel_id}:{thread_ts}"
def _thread_rehydration_key(
self,
channel_id: str,
thread_ts: str,
user_id: str,
team_id: str = "",
) -> str:
"""Per-process key for the once-per-thread restart-rehydration check.
Scoped like the session key: when ``thread_sessions_per_user`` is on,
each user's thread session rehydrates independently.
"""
key = f"{team_id}:{channel_id}:{thread_ts}"
store_cfg = getattr(getattr(self, "_session_store", None), "config", None)
if getattr(store_cfg, "thread_sessions_per_user", False):
key = f"{key}:{user_id}"
return key
def _mark_thread_rehydration_checked(
self,
channel_id: str,
thread_ts: str,
user_id: str,
team_id: str = "",
) -> None:
"""Record that this thread's restart-rehydration check has run."""
self._thread_rehydration_checked.add(
self._thread_rehydration_key(channel_id, thread_ts, user_id, team_id)
)
if (
len(self._thread_rehydration_checked)
> self._THREAD_REHYDRATION_CHECKED_MAX
):
excess = (
len(self._thread_rehydration_checked)
- self._THREAD_REHYDRATION_CHECKED_MAX // 2
)
for old_key in list(self._thread_rehydration_checked)[:excess]:
self._thread_rehydration_checked.discard(old_key)
def _get_thread_watermark(
self,
channel_id: str,

View file

@ -3605,11 +3605,14 @@ class TestThreadReplyHandling:
self, adapter_with_session_store, mock_session_store
):
"""Unmentioned replies in active threads keep the existing behavior:
no thread re-fetch, no context injection."""
no thread re-fetch, no context injection (once the one-shot restart
rehydration check has found no watermark)."""
mock_session_store._entries = {"any": MagicMock()}
adapter_with_session_store._has_active_session_for_thread = MagicMock(
return_value=True
)
# No persisted watermark → rehydration check is a no-op.
mock_session_store.get_session_metadata = MagicMock(return_value="")
adapter_with_session_store._app.client.conversations_replies = AsyncMock()
adapter_with_session_store._fetch_thread_parent_text = AsyncMock(
return_value=""
@ -3630,6 +3633,80 @@ class TestThreadReplyHandling:
msg_event = adapter_with_session_store.handle_message.call_args[0][0]
assert msg_event.channel_context is None
@pytest.mark.asyncio
async def test_restart_rehydrates_thread_delta_once(
self, adapter_with_session_store, mock_session_store
):
"""After a gateway restart (fresh adapter instance, persisted session
+ watermark), the FIRST ordinary thread reply injects messages the
session missed while the gateway was down exactly once. Subsequent
replies do not re-fetch."""
mock_session_store._entries = {"any": MagicMock()}
adapter_with_session_store._has_active_session_for_thread = MagicMock(
return_value=True
)
# Persisted watermark survives the restart via the session store.
metadata = {"slack_thread_watermark:C123:123.000": "123.100"}
mock_session_store.get_session_metadata = MagicMock(
side_effect=lambda sk, k, d=None: metadata.get(k, d)
)
mock_session_store.set_session_metadata = MagicMock(
side_effect=lambda sk, k, v: metadata.__setitem__(k, v) or True
)
adapter_with_session_store._app.client.conversations_replies = AsyncMock(
return_value={
"messages": [
{"ts": "123.000", "user": "U_PARENT", "text": "Original question"},
{"ts": "123.100", "user": "U_USER", "text": "Old context"},
{"ts": "123.200", "user": "U_OTHER", "text": "Missed while down"},
{"ts": "123.456", "user": "U_USER", "text": "please continue"},
]
}
)
adapter_with_session_store._user_name_cache = {
("T_TEAM", "U_PARENT"): "Parent",
("T_TEAM", "U_USER"): "User",
("T_TEAM", "U_OTHER"): "Other",
}
# Fresh adapter instance == empty _thread_rehydration_checked, which
# is exactly the post-restart state.
assert adapter_with_session_store._thread_rehydration_checked == set()
await adapter_with_session_store._handle_slack_message({
"text": "please continue",
"user": "U_USER",
"channel": "C123",
"ts": "123.456",
"thread_ts": "123.000",
"channel_type": "channel",
"team": "T_TEAM",
})
first_event = adapter_with_session_store.handle_message.call_args[0][0]
assert first_event.text == "please continue"
assert "Missed while down" in first_event.channel_context
assert "Old context" not in first_event.channel_context
assert metadata["slack_thread_watermark:C123:123.000"] == "123.456"
# Second ordinary reply: no re-fetch, no injection.
adapter_with_session_store.handle_message.reset_mock()
adapter_with_session_store._app.client.conversations_replies.reset_mock()
await adapter_with_session_store._handle_slack_message({
"text": "and another thing",
"user": "U_USER",
"channel": "C123",
"ts": "123.500",
"thread_ts": "123.000",
"channel_type": "channel",
"team": "T_TEAM",
})
adapter_with_session_store._app.client.conversations_replies.assert_not_called()
second_event = adapter_with_session_store.handle_message.call_args[0][0]
assert second_event.channel_context is None
# Watermark keeps advancing in steady state.
assert metadata["slack_thread_watermark:C123:123.000"] == "123.500"
@pytest.mark.asyncio
async def test_top_level_message_requires_mention_even_with_session(
self, adapter_with_session_store, mock_session_store