From 26e472f0d87eabb7b0e34489b926e056ef0a17e1 Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Wed, 5 Aug 2026 09:09:23 +0900 Subject: [PATCH] fix(memory): preserve repeated history provenance Retain callback history object identities before invocation so repeated or moved history items cannot be persisted as fresh input, while preserving post-callback replacement and bounded content matching behavior. Co-authored-by: Henry Su --- .../run_internal/session_persistence.py | 12 +- tests/memory/test_session.py | 37 +++++++ tests/test_agent_runner.py | 104 ++++++++++++++++++ 3 files changed, 152 insertions(+), 1 deletion(-) diff --git a/src/agents/run_internal/session_persistence.py b/src/agents/run_internal/session_persistence.py index b4c98d2747..3f44e7d2d7 100644 --- a/src/agents/run_internal/session_persistence.py +++ b/src/agents/run_internal/session_persistence.py @@ -217,6 +217,10 @@ async def prepare_input_with_session( ) history_for_callback = copy.deepcopy(converted_history) new_items_for_callback = copy.deepcopy(new_input_list) + # Keep the original history objects alive so their identities remain valid even if the + # callback removes them from the list it receives. + original_history_objects = list(history_for_callback) + original_history_object_ids = {id(item) for item in original_history_objects} combined = session_input_callback(history_for_callback, new_items_for_callback) if inspect.isawaitable(combined): combined = await combined @@ -246,12 +250,18 @@ async def prepare_input_with_session( new_key = _session_item_key(item) if _consume_reference(new_refs, new_key, item): new_counts[new_key] = max(new_counts.get(new_key, 0) - 1, 0) - appended.append(item) + if id(item) in original_history_object_ids: + prune_history_indexes.add(combined_index) + else: + appended.append(item) continue if _consume_reference(history_refs, history_key, item): history_counts[history_key] = max(history_counts.get(history_key, 0) - 1, 0) prune_history_indexes.add(combined_index) continue + if id(item) in original_history_object_ids: + prune_history_indexes.add(combined_index) + continue if history_counts.get(history_key, 0) > 0: history_counts[history_key] = history_counts.get(history_key, 0) - 1 prune_history_indexes.add(combined_index) diff --git a/tests/memory/test_session.py b/tests/memory/test_session.py index be761aea6e..3b180539b6 100644 --- a/tests/memory/test_session.py +++ b/tests/memory/test_session.py @@ -569,6 +569,43 @@ def filter_assistant_messages(history, new_input): session.close() +@pytest.mark.parametrize("runner_method", ["run", "run_sync", "run_streamed"]) +@pytest.mark.asyncio +async def test_session_callback_repeating_history_does_not_grow_session(runner_method): + with tempfile.TemporaryDirectory() as temp_dir: + db_path = Path(temp_dir) / "test_memory.db" + model = FakeModel() + agent = Agent(name="test", model=model) + session = SQLiteSession("session_repeat", db_path) + + def repeat_first(history, new_input): + if not history: + return new_input + return history + [history[0]] + new_input + + try: + for turn in range(3): + model.set_next_output([get_text_message(f"assistant {turn}")]) + await run_agent_async( + runner_method, + agent, + f"user {turn}", + session=session, + run_config=RunConfig(session_input_callback=repeat_first), + ) + + stored = await session.get_items() + user_messages = [item for item in stored if item.get("role") == "user"] + assert [item.get("content") for item in user_messages] == [ + "user 0", + "user 1", + "user 2", + ] + assert len(stored) == 6 + finally: + session.close() + + @pytest.mark.asyncio async def test_sqlite_session_unicode_content(): """Test that session correctly stores and retrieves unicode/non-ASCII content.""" diff --git a/tests/test_agent_runner.py b/tests/test_agent_runner.py index 8d343b6609..f651dfff9c 100644 --- a/tests/test_agent_runner.py +++ b/tests/test_agent_runner.py @@ -2676,6 +2676,110 @@ def callback( assert [cast(dict[str, Any], item).get("content") for item in session_items] == ["new"] +@pytest.mark.asyncio +async def test_prepare_input_with_session_repeated_history_keeps_equal_new_item() -> None: + history_item = cast(TResponseInputItem, {"role": "user", "content": "same"}) + session = SimpleListSession(history=[history_item]) + + def callback( + history: list[TResponseInputItem], new_input: list[TResponseInputItem] + ) -> list[TResponseInputItem]: + return [history[0], history[0], new_input[0]] + + prepared, session_items = await prepare_input_with_session("same", session, callback) + + assert [cast(dict[str, Any], item).get("content") for item in prepared] == [ + "same", + "same", + "same", + ] + assert [cast(dict[str, Any], item).get("content") for item in session_items] == ["same"] + + +@pytest.mark.asyncio +async def test_prepare_input_with_session_async_callback_moves_repeated_history_item() -> None: + history_item = cast(TResponseInputItem, {"role": "user", "content": "history"}) + session = SimpleListSession(history=[history_item]) + + async def callback( + history: list[TResponseInputItem], new_input: list[TResponseInputItem] + ) -> list[TResponseInputItem]: + await asyncio.sleep(0) + moved = history.pop(0) + return [moved, new_input[0], moved] + + prepared, session_items = await prepare_input_with_session("new", session, callback) + + assert [cast(dict[str, Any], item).get("content") for item in prepared] == [ + "history", + "new", + "history", + ] + assert [cast(dict[str, Any], item).get("content") for item in session_items] == ["new"] + + +@pytest.mark.asyncio +async def test_prepare_input_with_session_history_moved_to_new_input_stays_history() -> None: + history_item = cast(TResponseInputItem, {"role": "user", "content": "history"}) + session = SimpleListSession(history=[history_item]) + + def callback( + history: list[TResponseInputItem], new_input: list[TResponseInputItem] + ) -> list[TResponseInputItem]: + moved = history.pop(0) + new_input.insert(0, moved) + return new_input + [moved] + + prepared, session_items = await prepare_input_with_session("new", session, callback) + + assert [cast(dict[str, Any], item).get("content") for item in prepared] == [ + "history", + "new", + "history", + ] + assert [cast(dict[str, Any], item).get("content") for item in session_items] == ["new"] + + +@pytest.mark.asyncio +async def test_prepare_input_with_session_callback_replaces_history_item() -> None: + history_item = cast(TResponseInputItem, {"role": "user", "content": "history"}) + replacement = cast(TResponseInputItem, {"role": "user", "content": "summary"}) + session = SimpleListSession(history=[history_item]) + + def callback( + history: list[TResponseInputItem], new_input: list[TResponseInputItem] + ) -> list[TResponseInputItem]: + history[0] = replacement + return history + new_input + + prepared, session_items = await prepare_input_with_session("new", session, callback) + + assert [cast(dict[str, Any], item).get("content") for item in prepared] == [ + "summary", + "new", + ] + assert [cast(dict[str, Any], item).get("content") for item in session_items] == ["new"] + + +@pytest.mark.asyncio +async def test_prepare_input_with_session_extra_reconstructed_history_item_stays_new() -> None: + history_item = cast(TResponseInputItem, {"role": "user", "content": "history"}) + session = SimpleListSession(history=[history_item]) + + def callback( + history: list[TResponseInputItem], new_input: list[TResponseInputItem] + ) -> list[TResponseInputItem]: + rebuilt = cast(TResponseInputItem, dict(cast(dict[str, Any], history[0]))) + return [history[0], rebuilt, new_input[0]] + + _, session_items = await prepare_input_with_session("new", session, callback) + + assert [cast(dict[str, Any], item).get("content") for item in session_items] == [ + "history", + "new", + ] + + @pytest.mark.asyncio async def test_prepare_input_with_openai_conversation_strips_assistant_history_ids() -> None: class DummyOpenAIConversationsSession(OpenAIConversationsSession):