[goingforstudying-ctrl]

Ran into this while serving an agent off a shared Memory object in an async server: chunks of the recent conversation kept vanishing even though the token limit was nowhere near hit. Turns out the waterfall isn't safe against two puts landing on the same session at the same time.

_manage_queue snapshots the active queue, counts tokens, then flushes with archive_oldest_messages(n) — positional. When two aputs overlap, both waterfalls snapshot the queue before either archives. The second one's "oldest n" then lands on rows the first one deliberately kept, so active messages get archived early. With an unlucky interleaving it flushes the whole queue, including the message that was just added.

Repro (deterministic, using a store that staggers the two archive calls so the second select runs after the first commit):

class StaggeredStore(InMemoryChatStore):
    calls = 0
    async def archive_oldest_messages(self, key, n):
        type(self).calls += 1
        await asyncio.sleep(0.05 if type(self).calls == 1 else 0.3)
        return await super().archive_oldest_messages(key, n)

store = StaggeredStore()
await store.add_messages("s", [ChatMessage(role="user", content=f"seed{i} " + "word " * 12) for i in range(6)])
memory = Memory(token_limit=60, token_flush_size=40, chat_history_token_ratio=1.0,
                session_id="s", sql_store=store)
await asyncio.gather(
    memory.aput(ChatMessage(role="user", content="newA " + "word " * 12)),
    memory.aput(ChatMessage(role="user", content="newB " + "word " * 12)),
)

Back-to-back puts keep [seed4, seed5, newA, newB] active. The racy version archives everything — active queue comes back empty.

Fix: serialize the waterfall per event loop. aput/aput_messages/aset now run through a small asyncio lock, so each manage cycle sees the store as the previous one left it. The lock follows the running loop instead of being created once, because the sync wrappers go through asyncio_run and a cached asyncio.Lock would end up bound to a dead loop on the second put().

One thing I'm not fully sure about: this only serializes within a single Memory instance. Two instances pointed at the same session_id against a shared DB can still race at the store level — doing that properly would need something transactional in AsyncDBChatStore, which felt out of scope.

Tests: new regression test forces the interleaving with the gated store and asserts the concurrent result matches sequential puts (fails on main, passes here), plus a guard that sync put() works across loops. pytest tests/memory/ — 87 passed, ruff clean.

[0xamlab]

Reproduced the race on main with the gated store: two concurrent aputs left the active queue empty (all 8 messages archived). Checked out the branch and the same repro now matches the sequential run. Ran pytest llama-index-core/tests/memory/ — 87 passed — and pytest llama-index-core/tests/storage/chat_store/ — 25 passed — plus ruff check on the changed files; all clean. The loop-aware lock in memory.py:261 is the right fix because asyncio.Lock would otherwise bind to the sync wrapper's temporary loop. Good to merge.

[VANDRANKI]

Community review. I am not a maintainer, so this does not clear the merge gate.

I read the full diff and Memory._manage_queue / aput / aput_messages / aset in memory.py, and ran the code (llama-index-core at the PR head, Python 3.13, in-memory chat stores only; I did not run a SQL-backed store).

The race is real and the lock fixes the case it targets. The new test_concurrent_aput_does_not_over_archive passes on the PR head and fails with memory.py taken from origin/main. tests/memory and tests/agent on the PR head: 177 passed, 3 skipped.

Three things I found by running it, all with the PR's own StaggeredArchiveChatStore and the seed messages from the test:

1. A copy of the Memory object after any aput can no longer be deep-copied. _manage_lock_loop holds the event loop, and copy.deepcopy(memory) and memory.model_copy(deep=True) raise TypeError: cannot pickle '_contextvars.Context' object once a message has been put under a running loop. Both work on origin/main. A shallow model_copy() is fine. pickle.dumps(memory) and model_dump() still work (pickle logs "Removing unpickleable private attribute _manage_lock_loop"). I searched llama_index/core/agent, memory and workflow and found no deepcopy of a Memory, so this may only affect user code; I did not search the integrations. It is probably worth keeping the lock out of the instance state (for example a WeakKeyDictionary keyed by loop) or adding a __deepcopy__.

2. The lock is per Memory object and per event loop, so it does not cover two other ways to hit the same race. With the same staggered store and the same seed messages, sequential puts of newA then newB leave ['newA', 'newB'] active. On the PR head, both of these still end with an empty active queue (identical to origin/main):

   • two Memory objects with the same session_id and the same store, each doing one aput under asyncio.gather;
   • one Memory doing aput on the main loop while asyncio.to_thread(memory.put, ...) runs, since asyncio_run then uses a worker thread with its own event loop and therefore gets a different lock.
     The code comment says "per event loop", so this is by design, but the PR title and description read as if concurrent puts on one session are fixed in general. I would say in the description that it covers one Memory object on one loop.

3. Minor: aset takes the lock but areset does not. I did not test whether that matters.

I did not run a database-backed chat store, other concurrency levels than two puts, or CI.