[weiqingy]

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.

Refusing a key Mem0 cannot isolate

Mem0 tests agent_id for truthiness instead of matching on it, so both an unbound key and an empty one would drop the filter entirely and widen an operation to every key sharing the job id and set name, which for a delete would remove another key's items. An add under such a key would be worse than mis-filtered: Mem0 omits agent_id from the stored metadata, leaving the item unattributed to any key.

The empty string is a legal Flink key, reachable by keying a stream on a string field that can be empty, so both states are now refused at the point of use. They carry separate messages because the remedies differ. Encoding the empty key instead was rejected: keeping already-stored memories readable requires the mapping to be the identity on non-empty keys, which leaves no unoccupied non-empty value for the empty key to take. Reserving a sentinel only moves which key fails, since a Flink key arrives as String.valueOf(key) and can be any string. A collision-free encoding would therefore have to transform every key and orphan every memory already written.

The unbound state is now None and null in both languages. Both previously defaulted the memory's own key to the empty string, so neither could distinguish an absent context from the empty key, and the Java wrapper additionally stamped that value onto a MemorySet whose own default was null, so an unbound set originating in Java arrived in Python looking bound and passed both guards.

Validation deliberately stays out of switch_context, so a job carrying an empty key still runs unless it actually uses long-term memory.

Scoping the management operations to the mailbox thread

Binding context onto the set makes the four set-scoped operations safe on a worker thread, but obtaining or deleting a set still reads the key currently in scope. Submitting either to durable_execute_async therefore preserved the same cross-key race one call earlier. getMemorySet and deleteMemorySet now run the runner context's mailbox-thread check, the same one getResource uses. The four set-scoped operations stay unchecked, since being usable from a worker thread is the point of binding the context to the set.

The check fires in Java as well as Python. Java's getMemorySet is a pure factory that stamps the Java-side fields and never crosses into Python, so a Python-only check would leave Java agents unguarded. Python receives the checker through the constructor rather than through a new method on the RunnerContext ABC, because Java's checkMailboxThread lives on the runtime class rather than on the API interface, and adding one to the Python ABC would put a user-facing API method there with no Java counterpart.

delete_memory_set keeps its name-based signature. It has no bound context and applies to the key currently in scope, which the mailbox-thread check now makes correct, so closing the gap no longer requires a public signature change. Its documentation records that it can target a different key than MemorySet.delete on a same-named set.

CI

The cross-language Mem0LongTermMemoryTest and SkillsIntegrationTest both skip themselves unless ACTION_API_KEY is set. The python-it job passes the secret; the java-it and cross-language jobs set only LOG_LEVEL, so both Java suites went green with every case skipped. Both jobs now receive it. Repository secrets are still withheld from pull_request runs on forks, so this restores the coverage for push builds and same-repository pull requests, matching what python-it already gets.

Tests

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
test_empty_partition_key_is_refused_by_set_operations	an empty key is refused on all four, and none reaches Mem0
test_empty_partition_key_is_refused_by_memory_set_management	obtaining and deleting a set are refused under an empty key
test_memory_set_management_before_any_context_switch_is_refused	the unbound state is refused rather than treated as a key
test_memory_set_management_is_refused_off_the_mailbox_thread	both management operations run the check
test_set_scoped_operations_run_without_the_mailbox_thread	the four set-scoped operations stay usable off the mailbox thread
testForwardedSetCarriesTheContextItWasObtainedIn	the Java bridge forwards the context bound at creation
testUnboundSetIsRefusedRatherThanWidened	the Java guard throws before forwarding
testMemorySetIsRefusedBeforeAnyContextSwitch	the unbound state is refused in Java rather than treated as a key
test_to_python_memory_set_carries_the_action_context	the bridge helper carries all four fields, not just the name
testEmptyKeyedSetIsRefusedRatherThanWidened	the Java guard refuses an empty key before marshalling the set
testMemorySetManagementIsRefusedForAnEmptyKey	obtaining and deleting are refused under an empty key in Java
testMemorySetManagementIsRefusedOffTheMailboxThread	both Java management operations run the check
testSetScopedOperationsRunWithoutTheMailboxThread	the four Java set-scoped operations stay usable off the mailbox thread

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. test_empty_context_key_is_used_consistently_for_mem0_operations asserted that an empty agent_id is forwarded to Mem0, recording the widening as intended. It is replaced by test_empty_partition_key_is_refused_by_set_operations, which asserts refusal and that Mem0 is never reached.

Every assertion was checked against a deliberately reverted implementation. Reverting the suppression binding, the observation id binding, any of the refusal guards, the key binding itself, either mailbox check, or the order in which the thread check precedes the key check each leaves at least one test failing.

126 Python tests pass across api/memory, runtime/memory and runtime/tests. 80 Java tests pass across MemorySetTest, the runtime Mem0LongTermMemoryTest, MemoryRefTest, TestMemoryObservationFlush, ActionTaskContextManagerTest and ActionExecutionOperatorTest. ruff check and spotless:check are clean, and the touched Python files are ruff format 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 on push builds.

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 and delete_memory_set / deleteMemorySet keep their signatures and gain a documented per-action, mailbox-thread 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.

Mem0LongTermMemory gains a mailbox-thread checker parameter in both languages. It is a runtime class rather than public API, and the public RunnerContext interface is unchanged in both languages.

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. Code that obtains or deletes a memory set inside a DurableCallable or a durable_execute_async callable now raises on Python and on Java under JDK 21 and above; under JDK 11 the callable runs inline on the mailbox thread, so it does not raise there yet. A job keyed on a value that can be the empty string now raises when it uses long-term memory, rather than silently sharing one memory set across every key.

Documentation

• doc-needed
• doc-not-needed
• doc-included

Was this patch authored or co-authored using generative AI tooling?

• Yes
• No

Generated-by: Claude Code 2.1.251 (Claude Opus 5)

The first commit on the branch predates the model half of that convention and carries Generated-by: Claude Code 2.1.228.

[wenjin272]

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.

[wenjin272]

[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?

[weiqingy]

Good point, done in both languages.

Both management methods now run the same mailbox check getResource uses. The four set-scoped operations stay unchecked on purpose, since being safe on a worker thread is what binding the context to the set buys.

Java needed its own check rather than relying on the bridge, since getMemorySet there never enters Python. On the Python side I passed the checker through the constructor rather than adding one to the RunnerContext ABC, since Java's equivalent sits on the runtime class rather than the API interface. Curious whether you would have put it on the interface instead.

getLongTermMemory is left alone, since capturing the handle on the mailbox thread and using it inside the callable is normal.

This settles the delete_memory_set question without touching its signature. One caveat: on JDK 11 the durable callable runs inline, so the same code will not raise there yet.

[wenjin272]

[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.

[weiqingy]

PR description updated.

On refusing rather than encoding: to keep already-stored memories readable the mapping has to leave non-empty keys alone, which leaves no free non-empty value for "" to take. A reserved sentinel just moves which key breaks. Collision-free would mean rewriting every key's agent_id and orphaning what is already stored, so refusing looked like the better trade. Does that match how you would weigh it?

Both None and "" are refused now in both languages, with different messages. The unbound default was "" on both sides, so an unbound Java set already looked bound by the time it reached Python.

Two things worth knowing. On add, mem0 drops agent_id from the stored metadata, so items written under an empty key stay unrecoverable even after the key is fixed. And an existing test was asserting that agent_id="" gets forwarded, so this was pinned as intended behaviour.

[weiqingy]

@wenjin272 Thanks for the review! Good catch, and it turned out to be two jobs: the cross-language one you spotted, and java-it, which runs SkillsIntegrationTest behind the same gate. Both fixed here.

One limit worth knowing: secrets are withheld from pull_request runs on forks, so this restores coverage on push builds and same-repository PRs rather than on fork PRs.

[weiqingy]

The ut-java [ubuntu-latest] [java-11] [python-3.12] failure looks unrelated to this PR. Maven Central returned 502 Bad Gateway while resolving com.fasterxml.jackson.core:jackson-databind:2.18.2, so the build stopped during dependency resolution and no test ran. Evidence: job log.

That version is the existing jackson.version pin in the root pom.xml, this PR touches no pom.xml or dependency file, and the same job passed on the previous commit of this branch.

Re-triggering it once the current run finishes.

[wenjin272]

Thanks for addressing my comments. LGTM.

There are a few conflicts with main following the merge of #944. Could you please rebase and resolve them?

[weiqingy]

@wenjin272 Thanks for the review. Rebased onto main and pushed.