weiqingy opened a new pull request, #1002: URL: https://github.com/apache/flink-agents/pull/1002
Linked issue: #998 ### Purpose of change `Mem0LongTermMemory` is shared across partition keys and holds the current key, observation id and suppression flag in mutable fields. `switch_context` writes them on the mailbox thread before each action, while `add`, `get`, `search` and `delete` can run on a worker thread when an action passes them to `durable_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 by `agent_id` alone. 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 `MemorySet` when it is created, and has the four operations take them from the set. The four already receive `memory_set` as 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_set` requires the key rather than defaulting it. Operating on a set that carries no binding raises. Mem0 tests `agent_id` for 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 on `None` and `null` only. `delete_memory_set` is unchanged. It takes a name rather than a `MemorySet`, 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 than `MemorySet.delete` on 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 | What it pins | |---|---| | `test_memory_set_stays_on_its_own_key_after_the_owner_switches` | all four operations keep the set's key after the owner switches | | `test_observations_stay_with_the_action_that_obtained_the_set` | observation ownership follows the set for add, get, search and delete | | `test_suppression_follows_the_set_not_the_current_context` | the set's own suppression flag decides, in both directions | | `test_unbound_memory_set_is_refused_rather_than_widened` | all four operations raise, and none reaches Mem0 | | `testForwardedSetCarriesTheContextItWasObtainedIn` | the Java bridge forwards the context bound at creation | | `testUnboundSetIsRefusedRatherThanWidened` | the Java guard throws before forwarding | `test_switch_context` and `test_context_switch_changes_observation_owner_and_current_suppression` reused 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/memory` and `runtime/tests`. 75 Java tests pass across `MemorySetTest`, `Mem0LongTermMemoryTest`, `MemoryRefTest`, `TestMemoryObservationFlush`, `ActionTaskContextManagerTest` and `ActionExecutionOperatorTest`. `ruff check`, `ruff format --check` and `spotless:check` are clean. The long-term memory e2e test is gated on `ACTION_API_KEY` and needs a Flink cluster plus a local Ollama embedding model, so end-to-end coverage comes from CI. ### API `MemorySet` gains three fields in both languages, excluded from serialization like the existing long term memory back-reference. Java adds `setActionContext` and three getters. `BaseLongTermMemory.get_memory_set` / `getMemorySet` keeps its signature and gains the documented per-action scope. The runtime bridge helper `to_python_memory_set` now 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-needed` - [ ] `doc-not-needed` - [x] `doc-included` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
