Python: [Bug]: Automatic memory consolidation clears sessions still extracting memories #9218

import asyncio
import json
import tempfile
from datetime import timedelta

from agent_framework import Agent, AgentSession, ChatResponse, MemoryContextProvider, MemoryFileStore, Message, SessionContext


def reply(payload):
    return ChatResponse(messages=[Message(role="assistant", contents=[json.dumps(payload)])])


class Client:
    def __init__(self, interleaved):
        self.additional_properties = {}
        self.interleaved = interleaved
        self.consolidating = asyncio.Event()
        self.extracting_b = asyncio.Event()
        self.finish_consolidation = asyncio.Event()
        self.finish_extraction = asyncio.Event()
        self.consolidated_topics = []

    async def get_response(self, messages, **kwargs):
        system = messages[0].text if messages[0].role == "system" else ""
        if "consolidate one topic memory file" in system.lower():
            payload = json.loads(messages[-1].text)
            self.consolidated_topics.append(payload["topic"])
            self.consolidating.set()
            if self.interleaved:
                await self.finish_consolidation.wait()
            return reply({"summary": payload["summary"], "memories": payload["memories"]})
        if "extract durable memory candidates" in system.lower():
            if "B marker" not in messages[-1].text:
                return reply({"memories": []})
            self.extracting_b.set()
            if self.interleaved:
                await self.finish_extraction.wait()
            return reply({"memories": [{"topic": "New", "memory": "B new fact"}]})
        return reply("Reply")


async def reproduce(interleaved):
    with tempfile.TemporaryDirectory() as tmp:
        store = MemoryFileStore(tmp, owner_state_key="owner")
        client = Client(interleaved)
        provider = MemoryContextProvider(store=store, recent_turns=0, max_extractions=1,
            consolidation_min_sessions=2, consolidation_interval=timedelta(0), consolidation_client=client)
        agent = Agent(client=client, context_providers=[provider], default_options={"store": False})
        sessions = [AgentSession(session_id=sid) for sid in ["Seed", "A", "B", "C"]]
        for session in sessions:
            session.state["owner"] = "same-owner"
        seed, a, b, c = sessions
        context = SessionContext(session_id=seed.session_id, input_messages=[])
        await provider.before_run(agent=agent, session=seed, context=context, state={})
        write = next(tool for tool in context.tools if tool.name == "write_memory")
        await write.invoke(arguments={"topic": "Old", "memory": "Old fact"}, skip_parsing=True)
        await agent.run("ordinary turn", session=seed)
        if interleaved:
            first = asyncio.create_task(agent.run("ordinary turn", session=a))
            await asyncio.wait_for(client.consolidating.wait(), 5)
            second = asyncio.create_task(agent.run("B marker", session=b))
            await asyncio.wait_for(client.extracting_b.wait(), 5)
            client.finish_consolidation.set()
            await asyncio.wait_for(first, 5)
            client.finish_extraction.set()
            await asyncio.wait_for(second, 5)
        else:
            await agent.run("ordinary turn", session=a)
            await agent.run("B marker", session=b)
        state_after_b = store.read_state(b, source_id="memory")
        fact = store.get_topic(b, source_id="memory", topic="New").memories
        await agent.run("ordinary turn", session=c)
        return {"interleaved": interleaved, "after_b": state_after_b,
            "after_c": store.read_state(c, source_id="memory"),
            "new_fact": fact, "consolidated_topics": client.consolidated_topics}


async def main():
    print(json.dumps([await reproduce(False), await reproduce(True)], indent=2))


asyncio.run(main())