github-featured/microsoft--agent-framework
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

Python: fix: prevent superlinear history growth by deduplicating messages in save_messages (#7242)

Browse Source 
* fix: prevent superlinear history growth by deduplicating messages in save_messages

* fix: address review feedback for history deduplication

* fix: Prevent superlinear history growth by deduplicating messages

* fix: add list[Message] type hints

* fix(sessions): resolve deduplication churn and collapsing of identical message

* fix(sessions): replace uuid/seen-set dedup with sequence aware filtering

* fix: use forward-scan sequence alignment in filter_new_messages

* fix(core): annotate new_msgs type to resolve pyright errors

---------

Co-authored-by: Eduard van Valkenburg <eavanvalkenburg@users.noreply.github.com>
Co-authored-by: Evan Mattson <35585003+moonbox3@users.noreply.github.com>
This commit is contained in:
pratik wayase
2026-08-19 04:40:16 +00:00
committed by GitHub
parent e74ac4613c
commit da11daebe5
6 changed files with 392 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/core/agent_framework/_sessions.py
+83 -7
View File
@@ -35,7 +35,7 @@ from contextvars import ContextVar, Token
from dataclasses import dataclass
from functools import partial
from pathlib import Path
from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias, TypeVar, cast
from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias, TypeVar, cast, TypeGuard
import msgspec
@@ -68,6 +68,7 @@ _MESSAGE_INJECTION_LOCK = threading.Lock()
JsonDumps: TypeAlias = Callable[[Any], str | bytes]
JsonLoads: TypeAlias = Callable[[str | bytes], Any]
ServiceSessionId: TypeAlias = Mapping[str, Any]
MessageIdentity: TypeAlias = tuple[str, ...]
StateT = TypeVar("StateT")
StateEncoder: TypeAlias = Callable[[Any], Mapping[str, Any]]
StateDecoder: TypeAlias = Callable[[Mapping[str, Any]], Any]
@@ -187,6 +188,62 @@ def _deduplicate_origin_session_ids(origin_session_ids: Iterable[str]) -> list[s
return unique_origin_session_ids
def get_message_identity(message: Message) -> MessageIdentity:
"""Return a stable identity for a message for deduplication.
Uses the message's ID if available, otherwise falls back to a hash of
its role and serialized contents to prevent duplicate persistence.
"""
msg_id = getattr(message, "message_id", None)
if msg_id is None:
msg_id = getattr(message, "id", None)
if msg_id is not None:
return ("id", str(msg_id))
try:
contents_data = [c.to_dict() for c in message.contents] if message.contents else []
serialized = json.dumps(contents_data, sort_keys=True, ensure_ascii=False)
return ("content", str(message.role), serialized)
except Exception:
return ("content", str(message.role), str(message.contents))
def _get_message_hash(message: Message) -> MessageIdentity:
"""Stable hash for sequence matching."""
return get_message_identity(message)
def filter_new_messages(existing: Sequence[Message], incoming: Sequence[Message]) -> list[Message]:
"""Filters incoming messages to only those that are truly new.
Handles both 'append-only' and 'full transcript replay' scenarios.
Prevents superlinear growth and preserves legitimate duplicate turns.
"""
if not existing:
return list(incoming)
existing_hashes = [_get_message_hash(m) for m in existing]
incoming_hashes = [_get_message_hash(m) for m in incoming]
if len(incoming) >= len(existing) and incoming_hashes[: len(existing_hashes)] == existing_hashes:
return list(incoming[len(existing) :])
try:
for i in range(len(incoming_hashes) - len(existing_hashes) + 1):
if incoming_hashes[i : i + len(existing_hashes)] == existing_hashes:
return list(incoming[i + len(existing_hashes) :])
except Exception:
logger.debug("sequence alignment check failed, falling back to set-based deduplication")
existing_set = set(existing_hashes)
new_msgs: list[Message] = []
for m, h in zip(incoming, incoming_hashes):
if h not in existing_set:
new_msgs.append(m)
existing_set.add(h)
return new_msgs
@dataclass(frozen=True, slots=True)
class _StateTypeRegistration:
cls: type[Any]
@@ -194,7 +251,6 @@ class _StateTypeRegistration:
encoder: StateEncoder
decoder: StateDecoder
_STATE_TYPE_REGISTRY: dict[str, _StateTypeRegistration] = {}
_STATE_CLASS_REGISTRY: dict[type[Any], _StateTypeRegistration] = {}
@@ -2095,10 +2151,13 @@ class InMemoryHistoryProvider(HistoryProvider):
) -> None:
"""Persist messages to session state."""
mark_feature_used(FeatureIndex.CORE_IN_MEMORY_HISTORY_PROVIDER)
if state is None:
if state is None or not messages:
return
existing = state.get("messages", [])
state["messages"] = [*existing, *messages]
new_messages = filter_new_messages(existing, messages)
if new_messages:
state["messages"] = [*existing, *new_messages]
@experimental(feature_id=ExperimentalFeature.FILE_HISTORY)
@@ -2257,9 +2316,26 @@ class FileHistoryProvider(HistoryProvider):
def _append_messages() -> None:
with file_lock:
if self.serialization_format == "json":
with file_path.open("a", encoding="utf-8") as file_handle:
for message in messages:
file_handle.write(f"{self._serialize_json_message(message)}\n")
existing_messages: list[Message] = []
if file_path.exists():
with file_path.open("r", encoding="utf-8") as f:
for line in f:
line = line.strip()
if not line:
continue
try:
payload = self.loads(line)
msg = Message.from_dict(dict(cast(Mapping[str, Any], payload)))
existing_messages.append(msg)
except Exception:
logger.debug("failed to parse history line for deduplication")
continue
new_messages = filter_new_messages(existing_messages, messages)
if new_messages:
with file_path.open("a", encoding="utf-8") as file_handle:
for message in new_messages:
file_handle.write(f"{self._serialize_json_message(message)}\n")
return
with file_path.open("ab") as file_handle:
for message in messages:
python/packages/core/tests/core/test_agent_hooks.py
+1 -1
View File
@@ -2072,7 +2072,7 @@ async def test_drained_and_discarded_attempt_flushes_on_allow(streaming: bool) -
expected_response = "update - hello there" if streaming else "test response - hello there"
assert final.text == expected_response
stored = cast("list[Message]", session.state[provider.source_id]["messages"])
assert [message.text for message in stored] == ["hello there", expected_response] * 2
assert [message.text for message in stored] == ["hello there", expected_response]
class _DrainThenTerminateWithoutResultMiddleware(AgentMiddleware):
python/packages/core/tests/core/test_middleware_with_agent.py
+4 -4
View File
@@ -1781,14 +1781,14 @@ class TestChatAgentSessionBehavior:
second_after = thread_states[3]
assert second_after["before_next"] is False
assert second_after["messages_count"] == 1 # Input messages unchanged
assert second_after["thread_count"] == 4 # Previous history + current input + current response
assert second_after["thread_count"] == 3 # Previous history (2) + current input (1)
assert second_after["messages_text"] == ["second message"]
# Thread should contain: first input + first response + second input + second response
# Thread should contain: first input + first response + second input
assert "first message" in second_after["thread_messages_text"]
assert "second message" in second_after["thread_messages_text"]
# Should have two "test response" entries (one for each run)
# "test response" should only appear once since the duplicate was correctly filtered
response_count = sum(1 for text in second_after["thread_messages_text"] if "test response" in text)
assert response_count == 2
assert response_count == 1
class TestChatAgentChatMiddleware:
python/packages/core/tests/core/test_sessions.py
+204 -1
View File
@@ -8,7 +8,7 @@ import time
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from typing import TYPE_CHECKING, Any, cast
from unittest.mock import patch
import msgspec
@@ -47,6 +47,9 @@ from agent_framework._sessions import (
from agent_framework._telemetry import FeatureIndex
from agent_framework.exceptions import MiddlewareException
if TYPE_CHECKING:
from agent_framework._agents import SupportsAgentRun
# ---------------------------------------------------------------------------
# SessionContext tests
# ---------------------------------------------------------------------------
@@ -1381,6 +1384,131 @@ class TestInMemoryHistoryProvider:
ctx.extend_messages("custom-source", [Message(role="user", contents=["test"])])
assert "custom-source" in ctx.context_messages
async def test_save_messages_deduplicates_identical_messages(self) -> None:
"""Test that save_messages does not re-append messages already in the store."""
provider = InMemoryHistoryProvider()
state: dict[str, Any] = {}
msg1 = Message(role="user", contents=["hello"])
msg2 = Message(role="assistant", contents=["hi there"])
await provider.save_messages("s1", [msg1, msg2], state=state)
assert len(state["messages"]) == 2
await provider.save_messages("s1", [msg1, msg2], state=state)
assert len(state["messages"]) == 2
async def test_save_messages_only_appends_new_messages(self) -> None:
"""Test that save_messages filters out old messages and only appends new ones"""
provider = InMemoryHistoryProvider()
state: dict[str, Any] = {}
msg1 = Message(role="user", contents=["hello"])
msg2 = Message(role="assistant", contents=["hi there"])
msg3 = Message(role="user", contents=["how are you?"])
await provider.save_messages("s1", [msg1, msg2], state=state)
assert len(state["messages"]) == 2
await provider.save_messages("s1", [msg1, msg2, msg3], state=state)
assert len(state["messages"]) == 3
assert state["messages"][2].text == "how are you?"
async def test_save_messages_different_roles_same_text_not_deduplicated(self) -> None:
"""Test that messages with the same text but different roles are kept separate."""
provider = InMemoryHistoryProvider()
state: dict[str, Any] = {}
msg1 = Message(role="user", contents=["ping"])
msg2 = Message(role="assistant", contents=["ping"])
await provider.save_messages("s1", [msg1, msg2], state=state)
assert len(state["messages"]) == 2
async def test_save_messages_deduplication_with_none_state(self) -> None:
"""Test that save_messages with None state does not raise."""
provider = InMemoryHistoryProvider()
msg = Message(role="user", contents=["hello"])
await provider.save_messages("s1", [msg], state=None)
async def test_full_loop_does_not_grow_superlinearly(self) -> None:
"""Regression test: a multi-round looped run must not re-persist the whole
conversation on every round."""
from agent_framework import AgentResponse
provider = InMemoryHistoryProvider()
session = AgentSession()
provider_state = session.state.setdefault(provider.source_id, {})
ctx1 = SessionContext(session_id="s1", input_messages=[Message(role="user", contents=["turn 1"])])
await provider.before_run(
agent=cast("SupportsAgentRun", None),
session=session,
context=ctx1,
state=provider_state,
)
ctx1._response = AgentResponse(messages=[Message(role="assistant", contents=["reply 1"])])
await provider.after_run(
agent=cast("SupportsAgentRun", None),
session=session,
context=ctx1,
state=provider_state,
)
ctx2 = SessionContext(session_id="s1", input_messages=[Message(role="user", contents=["turn 2"])])
await provider.before_run(
agent=cast("SupportsAgentRun", None),
session=session,
context=ctx2,
state=provider_state,
)
ctx2._response = AgentResponse(messages=[Message(role="assistant", contents=["reply 2"])])
await provider.after_run(
agent=cast("SupportsAgentRun", None),
session=session,
context=ctx2,
state=provider_state,
)
stored = session.state[provider.source_id]["messages"]
assert len(stored) == 4
texts = [m.text for m in stored]
assert texts == ["turn 1", "reply 1", "turn 2", "reply 2"]
async def test_save_messages_preserves_duplicate_content(self) -> None:
"""Two separate user 'yes' replies in the same batch must both be persisted."""
provider = InMemoryHistoryProvider()
state: dict[str, Any] = {}
yes_1 = Message(role="user", contents=["yes"])
yes_2 = Message(role="user", contents=["yes"])
await provider.save_messages("s1", [yes_1, yes_2], state=state)
assert len(state["messages"]) == 2
assert state["messages"][0].text == "yes"
assert state["messages"][1].text == "yes"
async def test_save_messages_handles_replayed_transcript_with_duplicates(self) -> None:
provider = InMemoryHistoryProvider()
state: dict[str, Any] = {}
msg_b = Message (role = "user", contents=["B"])
await provider.save_messages("s1", [msg_b], state = state)
assert len(state["messages"]) == 1
msg_a = Message(role="user", contents=["A"])
msg_c = Message(role="user", contents=["C"])
msg_b2 = Message(role="user", contents=["B"])
msg_d = Message(role="user", contents=["D"])
await provider.save_messages("s1", [msg_a, msg_b, msg_c, msg_b2, msg_d], state = state)
assert len(state["messages"]) == 4
texts = [m.text for m in state["messages"]]
assert texts == ["B", "C", "B", "D"]
class TestFileHistoryProvider:
def test_is_marked_experimental(self) -> None:
@@ -1665,6 +1793,81 @@ class TestFileHistoryProvider:
loaded = await provider.get_messages(session_id)
assert [message.text for message in loaded] == ["first", "second"]
async def test_save_messages_deduplicates_identical_messages(self, tmp_path: Path) -> None:
"""Test that FileHistoryProvider does not re-append already persisted messages."""
provider = FileHistoryProvider(tmp_path)
msg1 = Message(role="user", contents=["hello"])
msg2 = Message(role="assistant", contents=["hi there"])
await provider.save_messages("s1", [msg1, msg2])
loaded = await provider.get_messages("s1")
assert len(loaded) == 2
await provider.save_messages("s1", [msg1, msg2])
loaded = await provider.get_messages("s1")
assert len(loaded) == 2
async def test_save_messages_only_appends_new_messages(self, tmp_path: Path) -> None:
"""Test that FileHistoryProvider filters out old messages and only appends new ones"""
provider = FileHistoryProvider(tmp_path)
msg1 = Message(role="user", contents=["hello"])
msg2 = Message(role="assistant", contents=["hi there"])
msg3 = Message(role="user", contents=["how are you?"])
await provider.save_messages("s1", [msg1, msg2])
loaded = await provider.get_messages("s1")
assert len(loaded) == 2
await provider.save_messages("s1", [msg1, msg2, msg3])
loaded = await provider.get_messages("s1")
assert len(loaded) == 3
assert loaded[2].text == "how are you?"
async def test_save_messages_different_roles_same_text_not_deduplicated(self, tmp_path: Path) -> None:
"""Test that messages with the same text but different roles are kept separate."""
provider = FileHistoryProvider(tmp_path)
msg1 = Message(role="user", contents=["ping"])
msg2 = Message(role="assistant", contents=["ping"])
await provider.save_messages("s1", [msg1, msg2])
loaded = await provider.get_messages("s1")
assert len(loaded) == 2
async def test_deduplication_file_integrity(self, tmp_path: Path) -> None:
"""Test that deduplication writes the correct number of lines to the JSONL file."""
provider = FileHistoryProvider(tmp_path)
msg1 = Message(role="user", contents=["hello"])
msg2 = Message(role="assistant", contents=["hi there"])
msg3 = Message(role="user", contents=["follow-up"])
await provider.save_messages("s1", [msg1, msg2])
session_file = provider._session_file_path("s1")
raw_lines = (await asyncio.to_thread(session_file.read_text, encoding="utf-8")).splitlines()
assert len(raw_lines) == 2
await provider.save_messages("s1", [msg1, msg2, msg3])
raw_lines = (await asyncio.to_thread(session_file.read_text, encoding="utf-8")).splitlines()
assert len(raw_lines) == 3
async def test_save_messages_preserves_duplicate_content(self, tmp_path: Path) -> None:
"""Test that two identical user turns in the same batch are both persisted."""
provider = FileHistoryProvider(tmp_path)
yes_1 = Message(role="user", contents=["yes"])
yes_2 = Message(role="user", contents=["yes"])
await provider.save_messages("s1", [yes_1, yes_2])
loaded = await provider.get_messages("s1")
assert len(loaded) == 2
assert loaded[0].text == "yes"
assert loaded[1].text == "yes"
# ---------------------------------------------------------------------------
# Run-persistence gate tests
python/packages/redis/agent_framework_redis/_history_provider.py
+10 -6
View File
@@ -8,12 +8,13 @@ This module provides ``RedisHistoryProvider``, built on the new
from __future__ import annotations
import json
from collections.abc import Sequence
from typing import Any, ClassVar
import redis.asyncio as redis
from agent_framework import Message
from agent_framework._sessions import HistoryProvider
from agent_framework._sessions import HistoryProvider, filter_new_messages
from agent_framework._telemetry import mark_feature_used
from redis.credentials import CredentialProvider
@@ -170,7 +171,14 @@ class RedisHistoryProvider(HistoryProvider):
return
key = self._redis_key(session_id)
serialized_messages = [self._serialize_json(msg) for msg in messages]
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:
@@ -185,15 +193,11 @@ class RedisHistoryProvider(HistoryProvider):
@staticmethod
def _serialize_json(message: Message) -> str:
"""Serialize a Message to a JSON string for Redis storage."""
import json
return json.dumps(message.to_dict())
@staticmethod
def _deserialize_json(data: str) -> dict[str, Any]:
"""Deserialize a JSON string from Redis to a dict."""
import json
return json.loads(data)
async def clear(self, session_id: str | None) -> None:
python/packages/redis/tests/test_providers.py
+90
View File
@@ -608,3 +608,93 @@ class TestRedisHistoryProviderBeforeAfterRun:
) # type: ignore[arg-type]
mock_redis_client.pipeline.assert_not_called()
class TestRedisHistoryProviderDeduplication:
"""Tests for Redis save_messages deduplication and trimming behavior."""
async def test_deduplicates_identical_messages(self, mock_redis_client: MagicMock):
msg1 = Message(role="user", contents=["hello"])
msg2 = Message(role="assistant", contents=["hi there"])
mock_redis_client.lrange = AsyncMock(return_value=[json.dumps(msg1.to_dict()), json.dumps(msg2.to_dict())])
with patch("agent_framework_redis._history_provider.redis.from_url") as mock_from_url:
mock_from_url.return_value = mock_redis_client
provider = RedisHistoryProvider("mem", redis_url="redis://localhost:6379")
await provider.save_messages("s1", [msg1, msg2])
pipeline = mock_redis_client.pipeline.return_value.__aenter__.return_value
pipeline.rpush.assert_not_called()
pipeline.execute.assert_not_called()
async def test_only_appends_new_messages(self, mock_redis_client: MagicMock):
msg1 = Message(role="user", contents=["hello"])
msg2 = Message(role="assistant", contents=["hi there"])
msg3 = Message(role="user", contents=["how are you?"])
mock_redis_client.lrange = AsyncMock(return_value=[json.dumps(msg1.to_dict()), json.dumps(msg2.to_dict())])
with patch("agent_framework_redis._history_provider.redis.from_url") as mock_from_url:
mock_from_url.return_value = mock_redis_client
provider = RedisHistoryProvider("mem", redis_url="redis://localhost:6379")
await provider.save_messages("s1", [msg1, msg2, msg3])
pipeline = mock_redis_client.pipeline.return_value.__aenter__.return_value
assert pipeline.rpush.call_count == 1
call_args = pipeline.rpush.call_args[0]
pushed_msg_dict = json.loads(call_args[1])
assert pushed_msg_dict["contents"][0]["text"] == "how are you?"
async def test_different_roles_same_text_not_deduplicated(self, mock_redis_client: MagicMock):
msg1 = Message(role="user", contents=["ping"])
mock_redis_client.lrange = AsyncMock(return_value=[json.dumps(msg1.to_dict())])
with patch("agent_framework_redis._history_provider.redis.from_url") as mock_from_url:
mock_from_url.return_value = mock_redis_client
provider = RedisHistoryProvider("mem", redis_url="redis://localhost:6379")
msg2 = Message(role="assistant", contents=["ping"])
await provider.save_messages("s1", [msg1, msg2])
pipeline = mock_redis_client.pipeline.return_value.__aenter__.return_value
assert pipeline.rpush.call_count == 1
async def test_trimmed_messages_not_reappended(self, mock_redis_client: MagicMock):
"""Messages trimmed by max_messages should not be re-appended
when the caller resends the full transcript. Sequence matching
handles this without needing a :seen set."""
msg_old = Message(role="user", contents=["old"])
msg_new = Message(role="assistant", contents=["new"])
mock_redis_client.lrange = AsyncMock(return_value=[json.dumps(msg_new.to_dict())])
with patch("agent_framework_redis._history_provider.redis.from_url") as mock_from_url:
mock_from_url.return_value = mock_redis_client
provider = RedisHistoryProvider("mem", redis_url="redis://localhost:6379")
await provider.save_messages("s1", [msg_old, msg_new])
pipeline = mock_redis_client.pipeline.return_value.__aenter__.return_value
pipeline.rpush.assert_not_called()
async def test_preserves_duplicate_content(self, mock_redis_client: MagicMock):
"""Two separate user 'yes' replies must both be persisted."""
yes_1 = Message(role="user", contents=["yes"])
yes_2 = Message(role="user", contents=["yes"])
mock_redis_client.lrange = AsyncMock(return_value=[])
with patch("agent_framework_redis._history_provider.redis.from_url") as mock_from_url:
mock_from_url.return_value = mock_redis_client
provider = RedisHistoryProvider("mem", redis_url="redis://localhost:6379")
await provider.save_messages("s1", [yes_1, yes_2])
pipeline = mock_redis_client.pipeline.return_value.__aenter__.return_value
assert pipeline.rpush.call_count == 2