[runtime][api] Scope long-term memory operations to the action that obtained the memory set - #1002
[runtime][api] Scope long-term memory operations to the action that obtained the memory set#1002weiqingy wants to merge 1 commit into
Conversation
…btained the set Mem0LongTermMemory holds the partition key and observation context in mutable fields that the mailbox thread overwrites on every context switch. Operations submitted through durable_execute_async run on a worker thread and read those fields themselves, so an operation can be attributed to whichever key the mailbox thread most recently switched to. The key reaches Mem0 as agent_id, which is the only isolation boundary between two keys that share a job id and a memory set name. Bind the partition key, observation id and suppression flag onto the MemorySet when it is created on the mailbox thread, and have add, get, search and delete take them from the set. The Java wrapper forwards into the same Python object and rebuilt the Python set from its name alone, so it now records the context on switch and carries it across the bridge. A set is therefore scoped to one action and must not be reused across actions. Two existing tests reused one across a context switch and expected the new key to apply; they now assert that the set keeps its own key instead. Operating on a set that carries no binding raises rather than proceeding. Mem0 ignores a falsy agent_id instead of matching on it, so an unbound set would widen an operation to every key sharing the job id and set name, which for a delete would remove another key's items. delete_memory_set takes a name rather than a MemorySet and still reads the shared field, which its docstring now records. Generated-by: Claude Code 2.1.228
github-actions
Bot
added
doc-included
There was a problem hiding this comment.
Thanks for taking this on @weiqingy.
BTW, while reviewing the change, I noticed that the cross-language job does not provide ACTION_API_KEY like python-it, which caused the Java LTM E2E test to be skipped. I think we can fix this in this PR or in a separate one.
|
|
||
| The current partition key and observation context are copied onto the set | ||
| so that operations submitted to a worker thread stay scoped to the action | ||
| that obtained it. Must be called on the mailbox thread. |
There was a problem hiding this comment.
[P1] Enforce mailbox-thread access for memory-set management
I think get_memory_set and delete_memory_set, as memory-set-level management operations, should only run on the mailbox thread. Otherwise, passing either to durable_execute_async can preserve the same cross-key race. Could we enforce this in both Java and Python, similar to RunnerContextImpl#getResource?
| would widen every operation to all keys sharing the job id and set name, which | ||
| for a delete means deleting another key's items. Refuse the operation instead. | ||
| """ | ||
| if memory_set.partition_key is None: |
There was a problem hiding this comment.
[P1] Handle empty partition keys safely
"" is a valid Flink key, but Mem0 ignores a falsy agent_id, dropping partition isolation. Please reject empty keys or encode them into a non-empty, collision-free value. Also use None/null, rather than "", for the unbound state.
2 participants
Linked issue: #998
Purpose of change
Mem0LongTermMemoryis shared across partition keys and holds the current key, observation id and suppression flag in mutable fields.switch_contextwrites them on the mailbox thread before each action, whileadd,get,searchanddeletecan run on a worker thread when an action passes them todurable_execute_async, and those operations read the fields themselves. The key they see is whatever the mailbox thread wrote most recently, not the key of the action that submitted the work.The key is the isolation boundary rather than observability metadata: it reaches Mem0 as
agent_id, and two keys sharing a job id and a memory set name are separated byagent_idalone. Reading the wrong one means one key's items land in, or are read from, another key's set. The observation id and suppression flag come from the same fields, so observations can be misattributed independently.This binds the partition key, observation id and suppression flag onto the
MemorySetwhen it is created, and has the four operations take them from the set. The four already receivememory_setas their first parameter, so nothing new is threaded through. A set is consequently scoped to one action and must not be held across actions.Java is covered in the same change rather than separately. Its operations all delegate into the same Python object, and the bridge rebuilt the Python set from its name alone, so once the Python side sources the key from the set, a Java-originated call would arrive without one. The Java wrapper now records the context when it switches and carries it across the bridge, and
to_python_memory_setrequires the key rather than defaulting it.Operating on a set that carries no binding raises. Mem0 tests
agent_idfor truthiness instead of matching on it, so an unbound set would widen an operation to every key sharing the job id and set name, which for a delete removes another key's items. An empty string is a legal partition key and is unaffected: the check is onNoneandnullonly.delete_memory_setis unchanged. It takes a name rather than aMemorySet, so it has no bound context and applies to the key currently in scope, which its documentation now records along with the fact that it can target a different key thanMemorySet.deleteon a same-named set. Giving it the same isolation means changing its signature, which is a public API question left open on #998.Tests
Three new tests and two adapted, plus one new Java test.
test_memory_set_stays_on_its_own_key_after_the_owner_switchestest_observations_stay_with_the_action_that_obtained_the_settest_suppression_follows_the_set_not_the_current_contexttest_unbound_memory_set_is_refused_rather_than_widenedtestForwardedSetCarriesTheContextItWasObtainedIntestUnboundSetIsRefusedRatherThanWidenedtest_switch_contextandtest_context_switch_changes_observation_owner_and_current_suppressionreused a single set across a context switch and expected the current key to apply. They now assert that a set keeps its own key, and obtain one per context.Each assertion was checked against a deliberately reverted implementation: reverting the suppression binding, the observation id binding, either guard body, or the key binding itself each leaves at least one test failing, and reverting the key binding fails five including both parameterisations of
test_switch_context.122 Python tests pass across
api/memory,runtime/memoryandruntime/tests. 75 Java tests pass acrossMemorySetTest,Mem0LongTermMemoryTest,MemoryRefTest,TestMemoryObservationFlush,ActionTaskContextManagerTestandActionExecutionOperatorTest.ruff check,ruff format --checkandspotless:checkare clean.The long-term memory e2e test is gated on
ACTION_API_KEYand needs a Flink cluster plus a local Ollama embedding model, so end-to-end coverage comes from CI.API
MemorySetgains three fields in both languages, excluded from serialization like the existing long term memory back-reference. Java addssetActionContextand three getters.BaseLongTermMemory.get_memory_set/getMemorySetkeeps its signature and gains the documented per-action scope. The runtime bridge helperto_python_memory_setnow requires the partition key; its only caller is the Java wrapper, and a version mismatch fails loudly with a missing-argument error rather than silently.Existing code that obtains a memory set inside the action using it is unaffected. Code that cached one across actions was already reading whichever key happened to be current, and now raises or stays on its original key instead.
Documentation
doc-neededdoc-not-neededdoc-included