Python: fix(redis): keep core message groups atomic when trimming history

python/packages/core/agent_framework/_compaction.py

@@ -159,7 +159,36 @@ def _is_reasoning_only_assistant(message: Message) -> bool:
     return all(content.type == "text_reasoning" for content in message.contents)
 
 
-def _unambiguous_tool_call_result_pairs(messages: Sequence[Message]) -> list[tuple[int, int]]:
+def _approval_request_identity(content: Content) -> str | None:
+    """Return the shared identity of an approval request/response wrapper."""
+    function_call = getattr(content, "function_call", None)
+    identity = getattr(function_call, "call_id", None) or getattr(function_call, "id", None)
+    if not isinstance(identity, str) or not identity:
+        identity = content.id
+    return identity if isinstance(identity, str) and identity else None
+
+
+def _unambiguous_approval_pairs(messages: Sequence[Message]) -> list[tuple[int, int]]:
+    """Pair approval responses with their unambiguous request occurrences."""
+    unmatched_request_indices: dict[str, list[int]] = {}
+    pairs: list[tuple[int, int]] = []
+    for message_index in range(len(messages)):
+        message = messages[message_index]
+        for content in message.contents:
+            if message.role == "assistant" and content.type == "function_approval_request":
+                identity = _approval_request_identity(content)
+                if identity:
+                    unmatched_request_indices.setdefault(identity, []).append(message_index)
+            elif content.type == "function_approval_response":
+                identity = _approval_request_identity(content)
+                candidates = unmatched_request_indices.get(identity) if identity else None
+                if candidates is not None and len(candidates) == 1:
+                    pairs.append((candidates.pop(), message_index))
+    return pairs
+
+
+def _unambiguous_tool_call_output_pairs(messages: Sequence[Message]) -> list[tuple[int, int]]:
+    """Pair tool results and approval requests with unambiguous declaration occurrences."""
     unmatched_declaration_indices: dict[tuple[str, str], list[int]] = {}
     pairs: list[tuple[int, int]] = []
 
@@ -174,6 +203,14 @@ def _unambiguous_tool_call_result_pairs(messages: Sequence[Message]) -> list[tup
             ):
                 unmatched_declaration_indices.setdefault((content.type, content.call_id), []).append(message_index)
                 continue
+            if message.role == "assistant" and content.type == "function_approval_request":
+                function_call = getattr(content, "function_call", None)
+                call_id = getattr(function_call, "call_id", None)
+                candidates = unmatched_declaration_indices.get(("function_call", call_id)) if call_id else None
+                if candidates is not None and len(candidates) == 1:
+                    # The same declaration also owns the later tool result.
+                    pairs.append((candidates[0], message_index))
+                continue
             if content.type not in ("function_result", "computer_tool_result") or not content.call_id:
                 continue
             call_type = "function_call" if content.type == "function_result" else "computer_tool_call"
@@ -240,7 +277,8 @@ def union(left: int, right: int) -> bool:
         return True
 
     linked = False
-    for declaration_message_index, result_message_index in _unambiguous_tool_call_result_pairs(messages):
+    pairs = [*_unambiguous_tool_call_output_pairs(messages), *_unambiguous_approval_pairs(messages)]
+    for declaration_message_index, result_message_index in pairs:
         declaration_span_index = span_by_message_index[declaration_message_index]
         result_span_index = span_by_message_index[result_message_index]
         if declaration_span_index < result_span_index and union(result_span_index, declaration_span_index):

python/packages/core/tests/core/test_compaction.py

@@ -274,6 +274,48 @@ def test_group_annotations_pair_nonadjacent_function_result_by_call_id() -> None
     assert _group_id(messages[1]) != call_group
 
 
+def test_group_annotations_keep_approval_request_response_and_result_atomic() -> None:
+    function_call = Content.from_function_call(id="c1", call_id="c1", name="f", arguments="{}")
+    request = Content.from_function_approval_request(id="c1", function_call=function_call)
+    response = Content.from_function_approval_response(approved=True, id="c1", function_call=function_call)
+    messages = [
+        Message(role="assistant", contents=[function_call, request]),
+        Message(role="user", contents=[response]),
+        _tool_result("c1", "ok"),
+        Message(role="assistant", contents=["done"]),
+    ]
+
+    annotate_message_groups(messages)
+
+    approval_group = _group_id(messages[0])
+    assert approval_group is not None
+    assert _group_id(messages[1]) == approval_group
+    assert _group_id(messages[2]) == approval_group
+    assert _group_id(messages[3]) != approval_group
+
+
+def test_group_annotations_link_separate_approval_request_to_original_call() -> None:
+    function_call = Content.from_function_call(id="c1", call_id="c1", name="f", arguments="{}")
+    request = Content.from_function_approval_request(id="c1", function_call=function_call)
+    response = Content.from_function_approval_response(approved=True, id="c1", function_call=function_call)
+    messages = [
+        Message(role="assistant", contents=[function_call]),
+        Message(role="assistant", contents=[request]),
+        Message(role="user", contents=[response]),
+        _tool_result("c1", "ok"),
+        Message(role="assistant", contents=["done"]),
+    ]
+
+    annotate_message_groups(messages)
+
+    approval_group = _group_id(messages[0])
+    assert approval_group is not None
+    assert _group_id(messages[1]) == approval_group
+    assert _group_id(messages[2]) == approval_group
+    assert _group_id(messages[3]) == approval_group
+    assert _group_id(messages[4]) != approval_group
+
+
 @pytest.mark.parametrize("screenshot", [Content.from_data(b"image", "image/png"), None])
 def test_group_annotations_keep_computer_call_result_and_reasoning_atomic(screenshot: Content | None) -> None:
     messages = [

python/packages/redis/agent_framework_redis/_history_provider.py

@@ -11,20 +11,101 @@
 import json
 import warnings
 from collections.abc import Awaitable, Sequence
+from copy import copy
 from inspect import isawaitable
 from typing import Any, ClassVar, Literal, TypeVar, cast
 
 import redis.asyncio as redis
 from agent_framework import Message
+from agent_framework._compaction import (
+    _unambiguous_approval_pairs,  # pyright: ignore[reportPrivateUsage]
+    group_messages,
+)
 from agent_framework._filesystem import _storage_key_segment  # pyright: ignore[reportPrivateUsage]
 from agent_framework._sessions import HistoryProvider, filter_new_messages
 from agent_framework._telemetry import mark_feature_used
 from redis.credentials import CredentialProvider
+from redis.exceptions import WatchError
 
 from ._feature_usage import FeatureIndex
 
 _T = TypeVar("_T")
 
+# Optimistic-concurrency retries for a single ``save_messages`` call. Each retry recomputes the
+# trim boundary against the key state after a competing writer modified it (WATCH conflict).
+_SAVE_WATCH_RETRIES = 8
+
+
+def _has_unfinished_approval(messages: Sequence[Message]) -> bool:
+    """Return whether an approval group still lacks at least one response."""
+    request_count = sum(
+        content.type == "function_approval_request" for message in messages for content in message.contents
+    )
+    paired_request_count = len(_unambiguous_approval_pairs(messages))
+    return request_count > paired_request_count
+
+
+def _trim_start_index(messages: Sequence[Message], max_messages: int) -> int:
+    """Return the oldest index whose retained suffix contains only complete core groups.
+
+    The raw count boundary advances past every group that crosses it, including linked
+    non-contiguous spans. Advancing can cross another linked group, so repeat until the
+    boundary is stable. A group larger than the limit is dropped rather than split.
+    Pending approval groups are retained whole across saves until their responses arrive,
+    even when that temporarily exceeds the limit.
+    """
+    count = len(messages)
+    if count <= max_messages:
+        return 0
+
+    # group_messages assigns missing message ids to form stable group ids. Shallow
+    # copies isolate that top-level mutation without copying the message contents.
+    spans = group_messages([copy(message) for message in messages])
+    group_bounds: dict[str, list[int]] = {}
+    group_members: dict[str, list[Message]] = {}
+    for span in spans:
+        group_id = str(span["group_id"])
+        start_index = int(span["start_index"])
+        end_index = int(span["end_index"])
+        bounds = group_bounds.setdefault(group_id, [start_index, end_index])
+        bounds[0] = min(bounds[0], start_index)
+        bounds[1] = max(bounds[1], end_index)
+        group_members.setdefault(group_id, []).extend(messages[start_index : end_index + 1])
+
+    start = count - max_messages
+    pending_starts = [
+        group_bounds[group_id][0] for group_id, members in group_members.items() if _has_unfinished_approval(members)
+    ]
+    boundary_limit = count
+    if pending_starts:
+        # A pending approval may need its original declaration on a later save.
+        # Retain the whole suffix containing it, including any overlapping groups.
+        protected_start = min(pending_starts)
+        while True:
+            next_start = min(
+                (
+                    group_start
+                    for group_start, group_end in group_bounds.values()
+                    if group_start < protected_start <= group_end
+                ),
+                default=protected_start,
+            )
+            if next_start == protected_start:
+                break
+            protected_start = next_start
+        start = min(start, protected_start)
+        boundary_limit = protected_start
+
+    while start < boundary_limit:
+        next_start = start
+        for group_start, group_end in group_bounds.values():
+            if group_start < start <= group_end:
+                next_start = max(next_start, group_end + 1)
+        if next_start == start:
+            break
+        start = next_start
+    return start
+
 
 async def _redis_result(value: Awaitable[_T] | _T) -> _T:
     """Await a redis-py command result that is annotated as the sync/async union.
@@ -316,23 +397,53 @@ async def save_messages(
             # history is what clear() is for.
             return
 
-        existing_messages = await self.get_messages(session_id, state=state, **kwargs)
-        new_messages = filter_new_messages(existing_messages, messages)
-
-        if not new_messages:
-            return
-
-        serialized_messages = [self._serialize_json(msg) for msg in new_messages]
-
-        async with self._redis_client.pipeline(transaction=True) as pipe:
-            for serialized in serialized_messages:
-                await _redis_result(pipe.rpush(key, serialized))
-            await pipe.execute()
-
-        if self.max_messages is not None:
-            current_count: int = await _redis_result(self._redis_client.llen(key))
-            if current_count > self.max_messages:
-                await _redis_result(self._redis_client.ltrim(key, -self.max_messages, -1))
+        # Retention trims the oldest messages, and the cut must not split a tool call/result pair.
+        # Read, append, and trim are guarded by WATCH so a concurrent save to the same key cannot
+        # make this one trim from a stale absolute index and delete the whole session history: if
+        # the key changed under us, the transaction aborts and we recompute against the new state.
+        last_conflict: WatchError | None = None
+        for _ in range(_SAVE_WATCH_RETRIES):
+            async with self._redis_client.pipeline(transaction=True) as pipe:
+                await _redis_result(pipe.watch(key))
+                # Read through the pipeline's own reserved connection. `WATCH` has
+                # already pinned it, and going back out through the parent client would
+                # need a second one: a borrowed client with `max_connections=1` would
+                # then fail every non-empty save.
+                watched = cast("list[str]", await _redis_result(cast("Any", pipe).lrange(key, 0, -1)))
+                existing_messages = [Message.from_dict(self._deserialize_json(row)) for row in watched]
+                new_messages = filter_new_messages(existing_messages, messages)
+                if not new_messages:
+                    await _redis_result(pipe.unwatch())
+                    return
+
+                serialized_messages = [self._serialize_json(msg) for msg in new_messages]
+                # ``existing_messages`` plus ``new_messages`` is exactly the list being written
+                # (existing entries followed by the appended ones), so its indices line up with the
+                # Redis list we are about to trim.
+                full_messages: list[Message] = list(existing_messages) + list(new_messages)
+                trim_start = -1
+                if self.max_messages is not None and len(full_messages) > self.max_messages:
+                    trim_start = _trim_start_index(full_messages, self.max_messages)
+
+                pipe.multi()
+                for serialized in serialized_messages:
+                    await _redis_result(pipe.rpush(key, serialized))
+                if trim_start >= 0:
+                    await _redis_result(pipe.ltrim(key, trim_start, -1))
+                try:
+                    await pipe.execute()
+                except WatchError as exc:
+                    # The key changed between WATCH and EXEC. Recompute against the new state.
+                    last_conflict = exc
+                    continue
+                return
+        # Every attempt aborted. Falling off the end would report success while
+        # nothing was persisted, so surface the conflict that defeated the last one.
+        raise (
+            last_conflict
+            if last_conflict is not None
+            else RuntimeError(f"Redis history save for {key!r} did not persist after {_SAVE_WATCH_RETRIES} attempts")
+        )
 
     @staticmethod
     def _serialize_json(message: Message) -> str: