This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new 36532b1ad [runtime][api] Scope long-term memory operations to the
action that obtained the memory set (#1002)
36532b1ad is described below
commit 36532b1ad702d5644aa134daa367ac0c55e6092c
Author: Weiqing Yang <[email protected]>
AuthorDate: Wed Sep 2 02:45:21 2026 -0700
[runtime][api] Scope long-term memory operations to the action that
obtained the memory set (#1002)
Generated-by: Claude Code 2.1.251 (Claude Opus 5)
---
.github/workflows/ci.yml | 2 +
.../agents/api/memory/BaseLongTermMemory.java | 11 ++
.../apache/flink/agents/api/memory/MemorySet.java | 35 ++++
.../docs/development/memory/long_term_memory.md | 25 +++
.../integration/test/SkillsIntegrationTest.java | 16 +-
.../resource/test/Mem0LongTermMemoryTest.java | 5 +-
python/flink_agents/api/memory/long_term_memory.py | 26 +++
.../flink_agents/runtime/flink_runner_context.py | 1 +
.../runtime/memory/mem0/mem0_long_term_memory.py | 127 +++++++++++--
.../mem0/tests/test_mem0_long_term_memory.py | 20 +-
.../memory/mem0/tests/test_mem0_op_recording.py | 201 +++++++++++++++++++--
.../memory/mem0/tests/test_mem0_recording_hook.py | 6 +-
python/flink_agents/runtime/python_java_utils.py | 21 ++-
.../runtime/tests/test_python_java_utils.py | 10 +
.../agents/runtime/memory/Mem0LongTermMemory.java | 69 ++++++-
.../runtime/operator/PythonBridgeManager.java | 8 +-
.../runtime/memory/Mem0LongTermMemoryTest.java | 119 +++++++++++-
.../operator/ActionExecutionOperatorTest.java | 2 +-
18 files changed, 645 insertions(+), 59 deletions(-)
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index a71791ff9..cb6304076 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -242,6 +242,7 @@ jobs:
run: bash tools/start_ollama_server.sh
- name: Run Java IT
env:
+ ACTION_API_KEY: ${{ secrets.DASHSCOPE_API_KEY }}
LOG_LEVEL: INFO
run: tools/ut.sh -j -e -f ${{ matrix.flink-version }}
@@ -309,6 +310,7 @@ jobs:
mvn -B --no-transfer-progress -pl integrations/vector-stores/milvus
-am -Dspotless.skip=true -Drat.skip=true -Dtest=MilvusVectorStoreTest
-Dsurefire.failIfNoSpecifiedTests=false test
- name: Run e2e tests
env:
+ ACTION_API_KEY: ${{ secrets.DASHSCOPE_API_KEY }}
LOG_LEVEL: INFO
run: |
export ES_HOST="http://localhost:9200";
diff --git
a/api/src/main/java/org/apache/flink/agents/api/memory/BaseLongTermMemory.java
b/api/src/main/java/org/apache/flink/agents/api/memory/BaseLongTermMemory.java
index 4e35e4908..426424ac0 100644
---
a/api/src/main/java/org/apache/flink/agents/api/memory/BaseLongTermMemory.java
+++
b/api/src/main/java/org/apache/flink/agents/api/memory/BaseLongTermMemory.java
@@ -31,6 +31,12 @@ public interface BaseLongTermMemory extends AutoCloseable {
/**
* Gets the memory set by name. If it does not exist, the backend creates
it.
*
+ * <p>The returned set is bound to the calling action, so call this from
the action body itself;
+ * calling it from another thread throws {@link IllegalStateException}.
Obtain one per action
+ * rather than holding one across actions. Operating on a set whose
partition key is absent or
+ * empty throws rather than silently widening the operation to every
partition key, and this
+ * method itself throws unless a non-empty partition key is in scope.
+ *
* @param name the name of the memory set
* @return the memory set
*/
@@ -39,6 +45,11 @@ public interface BaseLongTermMemory extends AutoCloseable {
/**
* Deletes the memory set.
*
+ * <p>Unlike the set-scoped operations, this takes a name and applies to
the key currently in
+ * scope, so it must be called from the action body; calling it from
another thread throws
+ * {@link IllegalStateException}. It can target a different key than
{@link MemorySet#delete} on
+ * a same-named set would, and throws unless a non-empty partition key is
in scope.
+ *
* @param name the name of the memory set to delete
* @return true if the memory set was successfully deleted
*/
diff --git
a/api/src/main/java/org/apache/flink/agents/api/memory/MemorySet.java
b/api/src/main/java/org/apache/flink/agents/api/memory/MemorySet.java
index 2cfdb87da..0af97ae63 100644
--- a/api/src/main/java/org/apache/flink/agents/api/memory/MemorySet.java
+++ b/api/src/main/java/org/apache/flink/agents/api/memory/MemorySet.java
@@ -30,10 +30,18 @@ import java.util.Objects;
/**
* Represents a long term memory set, a named collection of memory items. Acts
as a thin proxy that
* delegates all operations to the bound {@link BaseLongTermMemory}.
+ *
+ * <p>A set also carries the action context it was obtained in. Operations run
on a worker thread
+ * when submitted through {@code durableExecuteAsync}, by which time the
owning long term memory may
+ * already have switched to another partition key, so they take the context
from the set rather than
+ * from it. A set must therefore be obtained per action and not reused across
actions.
*/
public class MemorySet {
private final String name;
private @JsonIgnore BaseLongTermMemory ltm;
+ private @JsonIgnore String partitionKey;
+ private @JsonIgnore String observationId = "";
+ private @JsonIgnore boolean observationSuppressed;
@JsonCreator
public MemorySet(@JsonProperty("name") String name) {
@@ -100,10 +108,37 @@ public class MemorySet {
this.ltm = ltm;
}
+ /**
+ * Binds this set to the action context it was obtained in. Called on the
mailbox thread when
+ * the set is created.
+ *
+ * @param partitionKey the partition key this set is scoped to
+ * @param observationId identifier for the owning action's observations
+ * @param observationSuppressed whether observation is suppressed for the
owning action
+ */
+ public void setActionContext(
+ String partitionKey, String observationId, boolean
observationSuppressed) {
+ this.partitionKey = partitionKey;
+ this.observationId = observationId;
+ this.observationSuppressed = observationSuppressed;
+ }
+
public String getName() {
return name;
}
+ public String getPartitionKey() {
+ return partitionKey;
+ }
+
+ public String getObservationId() {
+ return observationId;
+ }
+
+ public boolean isObservationSuppressed() {
+ return observationSuppressed;
+ }
+
@Override
public boolean equals(Object o) {
if (o == null || getClass() != o.getClass()) return false;
diff --git a/docs/content/docs/development/memory/long_term_memory.md
b/docs/content/docs/development/memory/long_term_memory.md
index 666c697b9..473390017 100644
--- a/docs/content/docs/development/memory/long_term_memory.md
+++ b/docs/content/docs/development/memory/long_term_memory.md
@@ -179,6 +179,29 @@ public static void processEvent(Event event, RunnerContext
ctx) throws Exception
{{< /tabs >}}
+{{< hint warning >}}
+A memory set is scoped to the key of the action that obtained it. Call
`get_memory_set` /
+`getMemorySet` inside each action that needs one, rather than caching a set
and reusing it
+in a later action. Reusing a set would apply another key's operations to the
key it was
+originally obtained for, and operating on a set that carries no scope raises
an error.
+
+Obtaining a memory set and deleting one both read the key currently in scope,
which only
+the action's own thread reads consistently. Both raise wherever they run on a
worker
+thread, so passing `getMemorySet` / `get_memory_set` or `deleteMemorySet` /
+`delete_memory_set` to `durableExecuteAsync` / `durable_execute_async` is
unsupported: in
+Python, and in Java on JDK 21 and above, the callable runs on a worker thread
and the call
+raises. On JDK 11 Java runs the callable inline instead, so the same code does
not raise
+today, but it is unsupported there too and will raise once the job moves to a
newer JDK.
+
+Operations on a set you already hold are unaffected, since they carry the key
the set was
+obtained under. That is why a memory set may be handed to a worker thread but
the long-term
+memory itself may not.
+
+Long-term memory also requires a non-empty partition key. A key of `""`
carries no
+isolation in the backing store, so obtaining or using a memory set under one
raises
+rather than writing memories that no key can read back.
+{{< /hint >}}
+
### Adding Items
{{< tabs "Adding Items" >}}
@@ -549,3 +572,5 @@ The isolation hierarchy works as follows:
This means you can reuse the same memory set name across different partitions,
and each partition will normally access only its own memories.
> **Note:** Partition-level isolation uses a textual identity derived from the
> logical key instead of `key.hashCode()`. Java keys use
> `String.valueOf(key)`. Default-serialized PyFlink keys are deserialized and
> use Python `str`; explicitly typed PyFlink keys use `String.valueOf`, except
> byte-array keys, which use Python's bytes representation. This avoids hash
> collisions, but distinct keys may still share a memory context if they
> produce the same text, for example when custom key types ha [...]
+
+An empty partition key is rejected rather than used. The backing store treats
an empty `agent_id` as no filter at all, so an empty key would widen every read
and delete to every key in the job, and items added under it would be stored
with no partition attribution to read back later. A job that can legitimately
produce an empty key must map it to a non-empty placeholder before keying the
stream; long-term memory cannot do that mapping itself without changing the
identity of keys whose me [...]
diff --git
a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/SkillsIntegrationTest.java
b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/SkillsIntegrationTest.java
index 65360910e..19c078a01 100644
---
a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/SkillsIntegrationTest.java
+++
b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/SkillsIntegrationTest.java
@@ -72,11 +72,20 @@ import static
org.apache.flink.agents.api.agents.AgentExecutionOptions.MAX_RETRI
*/
public class SkillsIntegrationTest {
+ /**
+ * Whether a usable key is present. GitHub sets the variable to the empty
string when the secret
+ * is unavailable, such as on a pull request from a fork, so an absent key
reaches the test as
+ * empty rather than as null.
+ */
+ private static boolean isApiKeySet() {
+ String apiKey = System.getenv("ACTION_API_KEY");
+ return apiKey != null && !apiKey.isEmpty();
+ }
+
@Test
public void testWorkflowWithSkills() throws Exception {
Assumptions.assumeTrue(
- System.getenv().get("ACTION_API_KEY") != null,
- "ACTION_API_KEY is required for the skills end-to-end test.");
+ isApiKeySet(), "ACTION_API_KEY is required for the skills
end-to-end test.");
StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment();
env.setParallelism(1);
@@ -120,8 +129,7 @@ public class SkillsIntegrationTest {
@Test
public void testReActAgentWithSkills() throws Exception {
Assumptions.assumeTrue(
- System.getenv().get("ACTION_API_KEY") != null,
- "ACTION_API_KEY is required for the skills end-to-end test.");
+ isApiKeySet(), "ACTION_API_KEY is required for the skills
end-to-end test.");
StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment();
env.setParallelism(1);
diff --git
a/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/Mem0LongTermMemoryTest.java
b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/Mem0LongTermMemoryTest.java
index 8f6ef1747..c6dce1550 100644
---
a/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/Mem0LongTermMemoryTest.java
+++
b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/Mem0LongTermMemoryTest.java
@@ -73,7 +73,10 @@ public class Mem0LongTermMemoryTest {
pythonReady = isPythonAvailable();
esConfigured = System.getenv("ES_HOST") != null;
milvusConfigured = System.getenv("MILVUS_URI") != null;
- apiKeySet = System.getenv("ACTION_API_KEY") != null;
+ // GitHub sets the variable to the empty string when the secret is
unavailable, such as
+ // on a pull request from a fork, so an absent key arrives as empty
rather than as null.
+ String apiKey = System.getenv("ACTION_API_KEY");
+ apiKeySet = apiKey != null && !apiKey.isEmpty();
}
@ParameterizedTest(name = "vectorStore={0}")
diff --git a/python/flink_agents/api/memory/long_term_memory.py
b/python/flink_agents/api/memory/long_term_memory.py
index 358d820e1..c3ecd8921 100644
--- a/python/flink_agents/api/memory/long_term_memory.py
+++ b/python/flink_agents/api/memory/long_term_memory.py
@@ -72,12 +72,25 @@ class MemorySetItem(BaseModel):
class MemorySet(BaseModel):
"""Represents a long term memory set contains memory items.
+ A set is bound to the action context it was obtained in. Operations run on
a
+ worker thread when submitted through ``durable_execute_async``, by which
time
+ the owning long term memory may already have switched to another partition
+ key, so they take the context from the set rather than from it. A set must
+ therefore be obtained per action and not reused across actions.
+
Attributes:
name: The name of this memory set.
+ partition_key: The partition key this set is scoped to.
+ observation_id: Identifier for the owning action's observations.
+ observation_suppressed: Whether observation is suppressed for the
owning
+ action.
"""
name: str
ltm: "BaseLongTermMemory" = Field(default=None, exclude=True)
+ partition_key: str | None = Field(default=None, exclude=True)
+ observation_id: str = Field(default="", exclude=True)
+ observation_suppressed: bool = Field(default=False, exclude=True)
def add(
self,
@@ -150,6 +163,13 @@ class BaseLongTermMemory(ABC, BaseModel):
def get_memory_set(self, name: str) -> MemorySet:
"""Get the memory set by name. If it does not exist, create it.
+ The returned set is bound to the calling action, so call this from the
action
+ body itself; calling it from another thread raises. Obtain one per
action rather
+ than holding one across actions. Operating on a set whose partition
key is
+ absent or empty raises rather than silently widening the operation to
every
+ partition key, and this method itself raises unless a non-empty
partition key is
+ in scope.
+
Args:
name: The name of the memory set.
@@ -161,6 +181,12 @@ class BaseLongTermMemory(ABC, BaseModel):
def delete_memory_set(self, name: str) -> bool:
"""Delete the memory set.
+ Unlike the set-scoped operations, this takes a name and applies to the
key
+ currently in scope, so it must be called from the action body; calling
it from
+ another thread raises. It can target a different key than
``MemorySet.delete``
+ on a same-named set would, and raises unless a non-empty partition key
is in
+ scope.
+
Args:
name: The name of the memory set.
diff --git a/python/flink_agents/runtime/flink_runner_context.py
b/python/flink_agents/runtime/flink_runner_context.py
index 284b3c6eb..cf44c10be 100644
--- a/python/flink_agents/runtime/flink_runner_context.py
+++ b/python/flink_agents/runtime/flink_runner_context.py
@@ -1271,6 +1271,7 @@ def _init_long_term_memory(
chat_model_name=chat_model_name,
embedding_model_name=embedding_model_name,
vector_store_name=vector_store_name,
+ mailbox_thread_checker=lambda:
ctx._j_runner_context.checkMailboxThread(),
)
diff --git a/python/flink_agents/runtime/memory/mem0/mem0_long_term_memory.py
b/python/flink_agents/runtime/memory/mem0/mem0_long_term_memory.py
index db87a3470..03a14f9e3 100644
--- a/python/flink_agents/runtime/memory/mem0/mem0_long_term_memory.py
+++ b/python/flink_agents/runtime/memory/mem0/mem0_long_term_memory.py
@@ -24,7 +24,7 @@ import queue
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
-from typing import Any, Dict, List
+from typing import Any, Callable, Dict, List
from pydantic import ConfigDict, Field, PrivateAttr, field_validator
from typing_extensions import override
@@ -123,6 +123,41 @@ def _create_flink_agents_config_classes() -> tuple:
return _FlinkAgentsLlmConfig, _FlinkAgentsEmbedderConfig
+def _reject_empty_partition_key(key: str) -> str:
+ """Return ``key`` unless it is empty, which cannot scope an operation.
+
+ Mem0 matches on ``agent_id`` only when it is truthy, so an empty key widens
+ every operation to all keys sharing the job id and set name, and stores
added
+ items with no ``agent_id`` at all.
+ """
+ if not key:
+ msg = (
+ "Long-term memory cannot be scoped to an empty partition key. Mem0
"
+ "ignores an empty agent_id, so the operation would reach every key
"
+ "sharing the job id and set name, and added items would be stored "
+ "unattributed. Key the stream by a non-empty value."
+ )
+ raise ValueError(msg)
+ return key
+
+
+def _bound_partition_key(memory_set: MemorySet) -> str:
+ """Return the partition key the set is scoped to.
+
+ An unbound set 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:
+ msg = (
+ f"Memory set {memory_set.name!r} is not bound to a partition key. "
+ "Obtain it with get_memory_set inside the action that uses it,
rather "
+ "than constructing it directly or reusing one across actions."
+ )
+ raise ValueError(msg)
+ return _reject_empty_partition_key(memory_set.partition_key)
+
+
class Mem0LongTermMemory(InternalBaseLongTermMemory):
"""Long-Term Memory backed by Mem0.
@@ -136,10 +171,19 @@ class Mem0LongTermMemory(InternalBaseLongTermMemory):
description="The runner context to retrieve resources.", exclude=True
)
+ mailbox_thread_checker: Callable[[], None] = Field(
+ description="Raises when called off the Flink mailbox thread.
Whole-memory-set "
+ "management reads the partition key currently in scope, which only the
mailbox "
+ "thread can read consistently.",
+ exclude=True,
+ )
+
job_id: str = Field(description="Unique identifier for the job.")
- key: str = Field(
- default="", description="Unique identifier for the keyed partition."
+ key: str | None = Field(
+ default=None,
+ description="Keyed partition currently in scope, or None before the
first "
+ "context switch.",
)
chat_model_name: str = Field(
@@ -184,6 +228,7 @@ class Mem0LongTermMemory(InternalBaseLongTermMemory):
chat_model_name: str,
embedding_model_name: str,
vector_store_name: str,
+ mailbox_thread_checker: Callable[[], None],
) -> None:
"""Initialize the Mem0-based Long-Term Memory.
@@ -194,6 +239,8 @@ class Mem0LongTermMemory(InternalBaseLongTermMemory):
embedding_model_name: Resource name of the embedding model.
vector_store_name: Resource name of a
``CollectionManageableVectorStore`` to back Mem0.
+ mailbox_thread_checker: Callable that raises when invoked off the
+ Flink mailbox thread.
"""
# Resolve metric group upfront on the main thread so that it is
# safe to use from any thread later.
@@ -205,6 +252,7 @@ class Mem0LongTermMemory(InternalBaseLongTermMemory):
)
super().__init__(
ctx=ctx,
+ mailbox_thread_checker=mailbox_thread_checker,
job_id=job_id,
chat_model_name=chat_model_name,
embedding_model_name=embedding_model_name,
@@ -425,29 +473,74 @@ class Mem0LongTermMemory(InternalBaseLongTermMemory):
return json.dumps([record.to_wire() for record in records],
ensure_ascii=False)
+ def _current_partition_key(self) -> str:
+ """Return the partition key currently in scope.
+
+ Whole-set management operations read the key from the memory rather
than
+ from a set, so they are only correct once an action has switched
context.
+ """
+ if self.key is None:
+ msg = (
+ "Long-term memory has no partition key in scope. Call this
from "
+ "an action body, which always runs under a partition key,
rather "
+ "than before the first action has run."
+ )
+ raise ValueError(msg)
+ return _reject_empty_partition_key(self.key)
+
@override
def get_memory_set(self, name: str) -> MemorySet:
"""Get the memory set by name.
+ 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. Only the mailbox thread reads that context
consistently,
+ so this method refuses to run on any other thread.
+
Args:
name: The name of the memory set.
Returns:
The memory set.
+
+ Raises:
+ ValueError: If the partition key in scope is absent or empty.
+ Exception: Called from a thread other than the mailbox thread. The
+ refusal originates on the Java side and reaches Python through
the
+ bridge, so the concrete type is whatever that marshals to.
"""
- return MemorySet(name=name, ltm=self)
+ self.mailbox_thread_checker()
+ return MemorySet(
+ name=name,
+ ltm=self,
+ partition_key=self._current_partition_key(),
+ observation_id=self._observation_id,
+ observation_suppressed=self._observation_suppressed,
+ )
@override
def delete_memory_set(self, name: str) -> bool:
"""Delete a memory set and all its items.
+ Takes a name rather than a ``MemorySet``, so it has no bound context
to read
+ and uses the key currently in scope. It therefore refuses to run
outside the
+ mailbox thread, and deleting a whole set can target a different key
than
+ ``MemorySet.delete`` on a set of the same name would.
+
Args:
name: The name of the memory set.
Returns:
True if the memory set was deleted.
+
+ Raises:
+ ValueError: If the partition key in scope is absent or empty.
+ Exception: Called from a thread other than the mailbox thread. The
+ refusal originates on the Java side and reaches Python through
the
+ bridge, so the concrete type is whatever that marshals to.
"""
- observation_key = self.key
+ self.mailbox_thread_checker()
+ observation_key = self._current_partition_key()
observation_id = self._observation_id
observation_enabled = (
self._update_observation_enabled and not
self._observation_suppressed
@@ -485,10 +578,10 @@ class Mem0LongTermMemory(InternalBaseLongTermMemory):
Returns:
List of IDs of the added memories.
"""
- observation_key = self.key
- observation_id = self._observation_id
+ observation_key = _bound_partition_key(memory_set)
+ observation_id = memory_set.observation_id
observation_enabled = (
- self._update_observation_enabled and not
self._observation_suppressed
+ self._update_observation_enabled and not
memory_set.observation_suppressed
)
if isinstance(memory_items, str):
memory_items = [memory_items]
@@ -550,10 +643,10 @@ class Mem0LongTermMemory(InternalBaseLongTermMemory):
Returns:
List of memory items.
"""
- observation_key = self.key
- observation_id = self._observation_id
+ observation_key = _bound_partition_key(memory_set)
+ observation_id = memory_set.observation_id
observation_enabled = (
- self._get_observation_enabled and not self._observation_suppressed
+ self._get_observation_enabled and not
memory_set.observation_suppressed
)
if ids is not None:
if isinstance(ids, str):
@@ -605,10 +698,10 @@ class Mem0LongTermMemory(InternalBaseLongTermMemory):
memory_set: The memory set to delete from.
ids: Optional ID or list of IDs. If None, deletes all items.
"""
- observation_key = self.key
- observation_id = self._observation_id
+ observation_key = _bound_partition_key(memory_set)
+ observation_id = memory_set.observation_id
observation_enabled = (
- self._update_observation_enabled and not
self._observation_suppressed
+ self._update_observation_enabled and not
memory_set.observation_suppressed
)
if ids is None:
self._mem0_instance.delete_all(
@@ -662,10 +755,10 @@ class Mem0LongTermMemory(InternalBaseLongTermMemory):
Returns:
List of matching memory items.
"""
- observation_key = self.key
- observation_id = self._observation_id
+ observation_key = _bound_partition_key(memory_set)
+ observation_id = memory_set.observation_id
observation_enabled = (
- self._search_observation_enabled and not
self._observation_suppressed
+ self._search_observation_enabled and not
memory_set.observation_suppressed
)
result = self._mem0_instance.search(
query=query,
diff --git
a/python/flink_agents/runtime/memory/mem0/tests/test_mem0_long_term_memory.py
b/python/flink_agents/runtime/memory/mem0/tests/test_mem0_long_term_memory.py
index 3f3faf4e5..bae18aca2 100644
---
a/python/flink_agents/runtime/memory/mem0/tests/test_mem0_long_term_memory.py
+++
b/python/flink_agents/runtime/memory/mem0/tests/test_mem0_long_term_memory.py
@@ -172,7 +172,11 @@ def ltm(mock_ctx):
chat_model_name="test_chat_model",
embedding_model_name="test_embedding_model",
vector_store_name="test_vector_store",
+ mailbox_thread_checker=lambda: None,
)
+ # Operations refuse an absent or empty key, so the fixture runs under the
key an
+ # action would have switched to.
+ mem0_ltm.switch_context("test_key", observation_id="test-action")
return mem0_ltm
@@ -330,14 +334,12 @@ def test_switch_context(ltm) -> None:
memory_set.add(items="Data for key_a")
ltm.switch_context("key_b", observation_id="action-b")
- # key_b should have no items in the same memory set name
- items = memory_set.get()
- # Items from key_a should not be visible under key_b
- # (They have different agent_id scoping)
- assert len(items) == 0
-
- # Reset context
- ltm.switch_context("", observation_id="action-empty")
+ # The set stays scoped to key_a, so it still reads key_a's item after the
+ # switch rather than following the current context.
+ assert len(memory_set.get()) == 1
+ # A set obtained under key_b is scoped to key_b, so key_a's items in the
+ # same-named set are not visible through it.
+ assert len(ltm.get_memory_set(name="context_set").get()) == 0
class MockChatModelWithTokenUsage:
@@ -459,6 +461,7 @@ def test_token_usage_reported_on_switch_context() -> None:
chat_model_name="test_chat_model",
embedding_model_name="test_embedding_model",
vector_store_name="test_vector_store",
+ mailbox_thread_checker=lambda: None,
)
# First switch_context triggers lazy init; no metrics yet.
@@ -519,6 +522,7 @@ def test_token_usage_flushed_on_close() -> None:
chat_model_name="test_chat_model",
embedding_model_name="test_embedding_model",
vector_store_name="test_vector_store",
+ mailbox_thread_checker=lambda: None,
)
ltm.switch_context("key_1", observation_id="action-1")
diff --git
a/python/flink_agents/runtime/memory/mem0/tests/test_mem0_op_recording.py
b/python/flink_agents/runtime/memory/mem0/tests/test_mem0_op_recording.py
index 4b83aec25..fa729ecc6 100644
--- a/python/flink_agents/runtime/memory/mem0/tests/test_mem0_op_recording.py
+++ b/python/flink_agents/runtime/memory/mem0/tests/test_mem0_op_recording.py
@@ -24,6 +24,9 @@ from threading import Event
from typing import Any
from unittest.mock import MagicMock
+import pytest
+
+from flink_agents.api.memory.long_term_memory import MemorySet
from flink_agents.runtime.memory.internal_base_long_term_memory import (
InternalBaseLongTermMemory,
)
@@ -34,7 +37,11 @@ def _make_ltm(mem0: Any) -> Mem0LongTermMemory:
ctx = MagicMock()
ctx.agent_metric_group = None
ltm = Mem0LongTermMemory.model_construct(
- ctx=ctx, job_id="job", key="partition", metric_group=None
+ ctx=ctx,
+ job_id="job",
+ key="partition",
+ metric_group=None,
+ mailbox_thread_checker=lambda: None,
)
ltm._mem0 = mem0
ltm._observation_id = "action"
@@ -172,16 +179,15 @@ def
test_context_switch_changes_observation_owner_and_current_suppression() -> N
"results": [{"event": "ADD", "id": "m1", "memory": "value"}]
}
ltm = _make_ltm(mem0)
- memory_set = ltm.get_memory_set("prefs")
ltm.switch_context(
"suppressed", observation_id="suppressed-action",
observation_suppressed=True
)
- ltm.add(memory_set, "ignored")
+ ltm.add(ltm.get_memory_set("prefs"), "ignored")
assert _drain(ltm, "suppressed", "suppressed-action") == []
ltm.switch_context("observed", observation_id="observed-action")
- ltm.add(memory_set, "recorded")
+ ltm.add(ltm.get_memory_set("prefs"), "recorded")
assert [record["id"] for record in _drain(ltm, "observed",
"observed-action")] == [
"m1"
]
@@ -190,25 +196,192 @@ def
test_context_switch_changes_observation_owner_and_current_suppression() -> N
assert ltm._search_observation_enabled is True
-def test_empty_context_key_is_used_consistently_for_mem0_operations() -> None:
+def test_empty_partition_key_is_refused_by_set_operations() -> None:
+ mem0 = MagicMock()
+ ltm = _make_ltm(mem0)
+ empty_keyed = MemorySet(name="prefs", ltm=ltm, partition_key="")
+
+ for op in (
+ lambda: ltm.add(empty_keyed, "input"),
+ lambda: ltm.get(empty_keyed),
+ lambda: ltm.delete(empty_keyed),
+ lambda: ltm.search(empty_keyed, "query", limit=5),
+ ):
+ with pytest.raises(ValueError, match="empty partition key"):
+ op()
+
+ mem0.add.assert_not_called()
+ mem0.get_all.assert_not_called()
+ mem0.delete_all.assert_not_called()
+ mem0.search.assert_not_called()
+
+
+def test_empty_partition_key_is_refused_by_memory_set_management() -> None:
+ mem0 = MagicMock()
+ ltm = _make_ltm(mem0)
+ ltm.switch_context("", observation_id="empty-action")
+
+ with pytest.raises(ValueError, match="empty partition key"):
+ ltm.get_memory_set("prefs")
+ with pytest.raises(ValueError, match="empty partition key"):
+ ltm.delete_memory_set("prefs")
+
+ mem0.delete_all.assert_not_called()
+
+
+def test_memory_set_management_before_any_context_switch_is_refused() -> None:
+ ltm = Mem0LongTermMemory.model_construct(
+ ctx=MagicMock(),
+ job_id="job",
+ metric_group=None,
+ mailbox_thread_checker=lambda: None,
+ )
+
+ with pytest.raises(ValueError, match="no partition key in scope"):
+ ltm.get_memory_set("prefs")
+
+
+def test_memory_set_stays_on_its_own_key_after_the_owner_switches() -> None:
mem0 = MagicMock()
mem0.add.return_value = {"results": []}
mem0.get_all.return_value = {"results": []}
mem0.search.return_value = {"results": []}
ltm = _make_ltm(mem0)
+
+ ltm.switch_context("owner", observation_id="owner-action")
memory_set = ltm.get_memory_set("prefs")
- ltm.switch_context("", observation_id="empty-action")
+ ltm.switch_context("other", observation_id="other-action")
+ ltm.add(memory_set, "input")
+ ltm.get(memory_set)
+ ltm.search(memory_set, "query", limit=5)
+ ltm.delete(memory_set)
+
+ assert mem0.add.call_args.kwargs["agent_id"] == "owner"
+ assert mem0.get_all.call_args.kwargs["agent_id"] == "owner"
+ assert mem0.search.call_args.kwargs["agent_id"] == "owner"
+ assert mem0.delete_all.call_args.kwargs["agent_id"] == "owner"
+
+
+def test_observations_stay_with_the_action_that_obtained_the_set() -> None:
+ mem0 = MagicMock()
+ mem0.add.return_value = {
+ "results": [{"event": "ADD", "id": "m1", "memory": "value"}]
+ }
+ mem0.get_all.return_value = {"results": [{"id": "m2", "memory": "stored"}]}
+ mem0.search.return_value = {"results": []}
+ ltm = _make_ltm(mem0)
+
+ ltm.switch_context("owner", observation_id="owner-action")
+ memory_set = ltm.get_memory_set("prefs")
+
+ ltm.switch_context("other", observation_id="other-action")
ltm.add(memory_set, "input")
ltm.get(memory_set)
ltm.search(memory_set, "query", limit=5)
ltm.delete(memory_set)
- ltm.delete_memory_set("prefs")
-
- assert mem0.add.call_args.kwargs["agent_id"] == ""
- assert mem0.get_all.call_args.kwargs["agent_id"] == ""
- assert mem0.search.call_args.kwargs["agent_id"] == ""
- assert [call.kwargs["agent_id"] for call in
mem0.delete_all.call_args_list] == [
- "",
- "",
+
+ assert _drain(ltm, "other", "other-action") == []
+ assert [record["op"] for record in _drain(ltm, "owner", "owner-action")]
== [
+ "ADD",
+ "GET",
+ "SEARCH",
+ "DELETE_SET",
]
+
+
+def test_unbound_memory_set_is_refused_rather_than_widened() -> None:
+ mem0 = MagicMock()
+ ltm = _make_ltm(mem0)
+ unbound = MemorySet(name="prefs", ltm=ltm)
+
+ for op in (
+ lambda: ltm.add(unbound, "input"),
+ lambda: ltm.get(unbound),
+ lambda: ltm.delete(unbound),
+ lambda: ltm.search(unbound, "query", limit=5),
+ ):
+ with pytest.raises(ValueError, match="not bound to a partition key"):
+ op()
+
+ mem0.add.assert_not_called()
+ mem0.get_all.assert_not_called()
+ mem0.delete_all.assert_not_called()
+ mem0.search.assert_not_called()
+
+
+def test_suppression_follows_the_set_not_the_current_context() -> None:
+ mem0 = MagicMock()
+ mem0.add.return_value = {
+ "results": [{"event": "ADD", "id": "m1", "memory": "value"}]
+ }
+ ltm = _make_ltm(mem0)
+
+ # Obtained while suppressed, used while the current context is not: the
set's
+ # own flag decides, so nothing is recorded.
+ ltm.switch_context(
+ "owner", observation_id="owner-action", observation_suppressed=True
+ )
+ suppressed_set = ltm.get_memory_set("prefs")
+ ltm.switch_context("owner", observation_id="live-action")
+ ltm.add(suppressed_set, "input")
+ assert _drain(ltm, "owner", "owner-action") == []
+
+ # And the reverse: obtained unsuppressed, used while the current context is
+ # suppressed, so the operation is still recorded.
+ unsuppressed_set = ltm.get_memory_set("prefs")
+ ltm.switch_context(
+ "owner", observation_id="quiet-action", observation_suppressed=True
+ )
+ ltm.add(unsuppressed_set, "input")
+ assert [record["id"] for record in _drain(ltm, "owner", "live-action")] ==
["m1"]
+
+
+def test_memory_set_management_is_refused_off_the_mailbox_thread() -> None:
+ def _refuse() -> None:
+ msg = "Expected to be running on the task mailbox thread, but was not."
+ raise RuntimeError(msg)
+
+ # No key in scope: only a checker that runs before the key is read can
produce the
+ # mailbox-thread message. That pins the intended order, because off the
mailbox
+ # thread the key read is itself unreliable.
+ mem0 = MagicMock()
+ ltm = Mem0LongTermMemory.model_construct(
+ ctx=MagicMock(),
+ job_id="job",
+ metric_group=None,
+ mailbox_thread_checker=_refuse,
+ )
+ ltm._mem0 = mem0
+
+ with pytest.raises(RuntimeError, match="task mailbox thread"):
+ ltm.get_memory_set("prefs")
+ with pytest.raises(RuntimeError, match="task mailbox thread"):
+ ltm.delete_memory_set("prefs")
+
+ mem0.delete_all.assert_not_called()
+
+
+def test_set_scoped_operations_run_without_the_mailbox_thread() -> None:
+ # A set carries the context it was obtained under, so operations on it are
safe
+ # to forward to a worker thread. Gating them would break durable async
execution.
+ calls = []
+ mem0 = MagicMock()
+ mem0.add.return_value = {"results": []}
+ mem0.get_all.return_value = {"results": []}
+ mem0.search.return_value = {"results": []}
+ ltm = _make_ltm(mem0)
+ ltm.mailbox_thread_checker = lambda: calls.append(None)
+ ltm.switch_context("owner", observation_id="owner-action")
+
+ memory_set = ltm.get_memory_set("prefs")
+ # Guards the assertion below from passing vacuously on a checker that
never runs.
+ assert len(calls) == 1
+ calls.clear()
+
+ ltm.add(memory_set, "input")
+ ltm.get(memory_set)
+ ltm.search(memory_set, "query", limit=5)
+ ltm.delete(memory_set)
+
+ assert calls == []
diff --git
a/python/flink_agents/runtime/memory/mem0/tests/test_mem0_recording_hook.py
b/python/flink_agents/runtime/memory/mem0/tests/test_mem0_recording_hook.py
index 765efe704..5c02b64a6 100644
--- a/python/flink_agents/runtime/memory/mem0/tests/test_mem0_recording_hook.py
+++ b/python/flink_agents/runtime/memory/mem0/tests/test_mem0_recording_hook.py
@@ -32,7 +32,11 @@ def _make_ltm() -> Mem0LongTermMemory:
ctx = MagicMock()
ctx.agent_metric_group = None
return Mem0LongTermMemory.model_construct(
- ctx=ctx, job_id="job", key="shared-partition", metric_group=None
+ ctx=ctx,
+ job_id="job",
+ key="shared-partition",
+ metric_group=None,
+ mailbox_thread_checker=lambda: None,
)
diff --git a/python/flink_agents/runtime/python_java_utils.py
b/python/flink_agents/runtime/python_java_utils.py
index fc1b06662..e1ecac32b 100644
--- a/python/flink_agents/runtime/python_java_utils.py
+++ b/python/flink_agents/runtime/python_java_utils.py
@@ -385,12 +385,23 @@ def get_long_term_memory(ctx: Any) -> Any:
return ctx.long_term_memory
-def to_python_memory_set(name: str) -> MemorySet:
- """Build a Python ``MemorySet`` from its name. Used by the Java
- ``Mem0LongTermMemory`` wrapper to forward calls into Python
``Mem0LongTermMemory``,
- which expects a ``MemorySet`` instance but only reads its ``name`` field.
+def to_python_memory_set(
+ name: str,
+ partition_key: str,
+ observation_id: str = "",
+ observation_suppressed: bool = False, # noqa: FBT001
+) -> MemorySet:
+ """Build a Python ``MemorySet`` from the fields the Java side holds. Used
by the
+ Java ``Mem0LongTermMemory`` wrapper to forward calls into Python
+ ``Mem0LongTermMemory``, which reads the action context off the set rather
than
+ off itself, so the context has to travel with each forwarded call.
"""
- return MemorySet(name=name)
+ return MemorySet(
+ name=name,
+ partition_key=partition_key,
+ observation_id=observation_id,
+ observation_suppressed=observation_suppressed,
+ )
def mem0_items_to_java(
diff --git a/python/flink_agents/runtime/tests/test_python_java_utils.py
b/python/flink_agents/runtime/tests/test_python_java_utils.py
index 20dce808b..d4264d46f 100644
--- a/python/flink_agents/runtime/tests/test_python_java_utils.py
+++ b/python/flink_agents/runtime/tests/test_python_java_utils.py
@@ -30,6 +30,7 @@ from flink_agents.runtime.python_java_utils import (
call_embedding_with_usage,
convert_to_python_key_text,
get_python_tool_metadata,
+ to_python_memory_set,
wrap_to_input_event,
)
@@ -99,3 +100,12 @@ def test_convert_to_python_key_text_uses_python_str() ->
None:
def test_convert_to_python_key_text_does_not_unpickle_explicit_bytes() -> None:
assert convert_to_python_key_text(b"N.", "explicit") == "b'N.'"
assert convert_to_python_key_text(b"\x80\x04N.", "explicit") ==
"b'\\x80\\x04N.'"
+
+
+def test_to_python_memory_set_carries_the_action_context() -> None:
+ memory_set = to_python_memory_set("prefs", "owner", "owner-action", True)
+
+ assert memory_set.name == "prefs"
+ assert memory_set.partition_key == "owner"
+ assert memory_set.observation_id == "owner-action"
+ assert memory_set.observation_suppressed is True
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemory.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemory.java
index f414422cc..bc4d16a19 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemory.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemory.java
@@ -45,23 +45,44 @@ public class Mem0LongTermMemory implements
InteranlBaseLongTermMemory {
private final PythonResourceAdapter adapter;
private PyObject pyMem0;
+ private final Runnable mailboxThreadChecker;
- public Mem0LongTermMemory(PythonResourceAdapter adapter, PyObject pyMem0) {
+ // Null until the first context switch, mirroring the Python side's own
default. A
+ // memory set obtained before then carries no key and is refused rather
than forwarded.
+ private String partitionKey;
+ private String observationId = "";
+ private boolean observationSuppressed;
+
+ public Mem0LongTermMemory(
+ PythonResourceAdapter adapter, PyObject pyMem0, Runnable
mailboxThreadChecker) {
this.adapter = adapter;
this.pyMem0 = pyMem0;
+ this.mailboxThreadChecker = mailboxThreadChecker;
}
@Override
public MemorySet getMemorySet(String name) {
// Mirrors Python's `Mem0LongTermMemory.get_memory_set`: a pure
factory that
- // returns a new MemorySet bound to this ltm; no Python call is needed.
+ // returns a new MemorySet bound to this ltm; no Python call is
needed. The
+ // current action context is copied onto the set so that operations
forwarded
+ // from a worker thread stay scoped to the action that obtained it,
which is
+ // only the right context to copy when the caller is the action itself.
+ mailboxThreadChecker.run();
MemorySet ms = new MemorySet(name);
ms.setLtm(this);
+ ms.setActionContext(currentPartitionKey(), observationId,
observationSuppressed);
return ms;
}
@Override
public boolean deleteMemorySet(String name) {
+ // Takes a name rather than a MemorySet, so it has no bound context
and the Python
+ // side uses the key currently in scope. It is therefore only correct
on the mailbox
+ // thread, and can target a different key than MemorySet.delete on a
same-named set.
+ // Both checks run here so a Java caller gets the failure in Java
rather than
+ // marshalled back from Python, which repeats them on its own side.
+ mailboxThreadChecker.run();
+ currentPartitionKey();
return (Boolean) adapter.callMethod(pyMem0, "delete_memory_set",
Map.of("name", name));
}
@@ -149,6 +170,9 @@ public class Mem0LongTermMemory implements
InteranlBaseLongTermMemory {
@Override
public void switchContext(
String partitionKey, String observationId, boolean
observationSuppressed) {
+ this.partitionKey = partitionKey;
+ this.observationId = observationId;
+ this.observationSuppressed = observationSuppressed;
adapter.callMethod(
pyMem0,
"switch_context",
@@ -180,7 +204,46 @@ public class Mem0LongTermMemory implements
InteranlBaseLongTermMemory {
}
private Object buildPyMemorySet(MemorySet memorySet) {
- return adapter.invoke(TO_PYTHON_MEMORY_SET, memorySet.getName());
+ // Mem0 ignores a falsy agent_id rather than matching on it, so
forwarding an
+ // unbound or empty-keyed set would widen the operation to every key
sharing the job
+ // id and set name, which for a delete means deleting another key's
items.
+ if (memorySet.getPartitionKey() == null) {
+ throw new IllegalStateException(
+ String.format(
+ "Memory set '%s' is not bound to a partition key.
Obtain it with"
+ + " getMemorySet inside the action that
uses it, rather than"
+ + " constructing it directly or reusing
one across actions.",
+ memorySet.getName()));
+ }
+ requireNonEmptyPartitionKey(memorySet.getPartitionKey());
+ return adapter.invoke(
+ TO_PYTHON_MEMORY_SET,
+ memorySet.getName(),
+ memorySet.getPartitionKey(),
+ memorySet.getObservationId(),
+ memorySet.isObservationSuppressed());
+ }
+
+ /** Returns the partition key in scope, refusing what Mem0 cannot scope an
operation to. */
+ private String currentPartitionKey() {
+ if (partitionKey == null) {
+ throw new IllegalStateException(
+ "Long-term memory has no partition key in scope. Call this
from an action"
+ + " body, which always runs under a partition key,
rather than before"
+ + " the first action has run.");
+ }
+ return requireNonEmptyPartitionKey(partitionKey);
+ }
+
+ private static String requireNonEmptyPartitionKey(String key) {
+ if (key.isEmpty()) {
+ throw new IllegalStateException(
+ "Long-term memory cannot be scoped to an empty partition
key. Mem0 ignores"
+ + " an empty agent_id, so the operation would
reach every key sharing"
+ + " the job id and set name, and added items would
be stored"
+ + " unattributed. Key the stream by a non-empty
value.");
+ }
+ return key;
}
@SuppressWarnings("unchecked")
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java
index 0383faee4..ba44a889a 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java
@@ -193,7 +193,7 @@ class PythonBridgeManager implements AutoCloseable {
initPythonActionExecutor(agentPlan, jobIdentifier);
}
if (mem0Configured) {
- wireLongTermMemory(agentPlan);
+ wireLongTermMemory(agentPlan, mailboxThreadChecker);
}
initialized = true;
}
@@ -240,7 +240,7 @@ class PythonBridgeManager implements AutoCloseable {
* {@code create_flink_runner_context} already initialised via {@code
_init_long_term_memory})
* and wrap it as a Java {@link Mem0LongTermMemory}.
*/
- private void wireLongTermMemory(AgentPlan agentPlan) {
+ private void wireLongTermMemory(AgentPlan agentPlan, Runnable
mailboxThreadChecker) {
PyObject pyCtx = pythonActionExecutor.getPythonRunnerContext();
Object pyLtm =
pythonInterpreter.invoke("python_java_utils.get_long_term_memory", pyCtx);
if (pyLtm == null) {
@@ -254,7 +254,9 @@ class PythonBridgeManager implements AutoCloseable {
LongTermMemoryOptions.Mem0.EMBEDDING_MODEL_SETUP.getKey(),
LongTermMemoryOptions.Mem0.VECTOR_STORE.getKey()));
}
- longTermMemory = new Mem0LongTermMemory(pythonResourceAdapter,
(PyObject) pyLtm);
+ longTermMemory =
+ new Mem0LongTermMemory(
+ pythonResourceAdapter, (PyObject) pyLtm,
mailboxThreadChecker);
MemoryEventSettings settings =
MemoryEventSettings.from(agentPlan.getConfigData());
longTermMemory.configureObservation(
settings.generate(MemoryEventSettings.MemoryOp.LONG_TERM_UPDATE),
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemoryTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemoryTest.java
index 72498d77c..1554be965 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemoryTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemoryTest.java
@@ -30,12 +30,14 @@ import pemja.core.object.PyObject;
import java.util.List;
import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
@@ -50,8 +52,12 @@ public class Mem0LongTermMemoryTest {
@BeforeEach
void setUp() {
mocks = MockitoAnnotations.openMocks(this);
- ltm = new Mem0LongTermMemory(mockAdapter, mockPyMem0);
- when(mockAdapter.invoke(eq("python_java_utils.to_python_memory_set"),
any()))
+ ltm = new Mem0LongTermMemory(mockAdapter, mockPyMem0, () -> {});
+ // Operations refuse an absent or empty key, so every test that does
not manage its
+ // own context runs under the key an action would have switched to.
+ ltm.switchContext("a-key", "an-action", false);
+ when(mockAdapter.invoke(
+ eq("python_java_utils.to_python_memory_set"), any(),
any(), any(), any()))
.thenReturn(mockPyMemorySet);
}
@@ -236,4 +242,113 @@ public class Mem0LongTermMemoryTest {
verify(mockAdapter).callMethod(mockPyMem0, "close", Map.of());
verify(mockPyMem0).close();
}
+
+ @Test
+ void testUnboundSetIsRefusedRatherThanWidened() {
+ MemorySet unbound = new MemorySet("notes");
+ unbound.setLtm(ltm);
+
+ assertThatThrownBy(() -> ltm.add(unbound, List.of("hello"), null))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("not bound to a partition key");
+ verify(mockAdapter, never())
+ .invoke(eq("python_java_utils.to_python_memory_set"), any(),
any(), any(), any());
+ }
+
+ @Test
+ void testEmptyKeyedSetIsRefusedRatherThanWidened() {
+ MemorySet emptyKeyed = new MemorySet("notes");
+ emptyKeyed.setLtm(ltm);
+ emptyKeyed.setActionContext("", "an-action", false);
+
+ assertThatThrownBy(() -> ltm.add(emptyKeyed, List.of("hello"), null))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("empty partition key");
+ verify(mockAdapter, never())
+ .invoke(eq("python_java_utils.to_python_memory_set"), any(),
any(), any(), any());
+ }
+
+ @Test
+ void testMemorySetManagementIsRefusedForAnEmptyKey() {
+ ltm.switchContext("", "an-action", false);
+
+ assertThatThrownBy(() -> ltm.getMemorySet("notes"))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("empty partition key");
+ assertThatThrownBy(() -> ltm.deleteMemorySet("notes"))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("empty partition key");
+ verify(mockAdapter, never()).callMethod(eq(mockPyMem0),
eq("delete_memory_set"), any());
+ }
+
+ @Test
+ void testMemorySetIsRefusedBeforeAnyContextSwitch() {
+ Mem0LongTermMemory fresh = new Mem0LongTermMemory(mockAdapter,
mockPyMem0, () -> {});
+
+ assertThatThrownBy(() -> fresh.getMemorySet("notes"))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("no partition key in scope");
+ }
+
+ @Test
+ void testForwardedSetCarriesTheContextItWasObtainedIn() throws Exception {
+ ltm.switchContext("owner", "owner-action", false);
+ MemorySet ms = ltm.getMemorySet("notes");
+
+ ltm.switchContext("other", "other-action", true);
+ ltm.add(ms, List.of("hello"), null);
+
+ verify(mockAdapter)
+ .invoke(
+ eq("python_java_utils.to_python_memory_set"),
+ eq("notes"),
+ eq("owner"),
+ eq("owner-action"),
+ eq(false));
+ }
+
+ @Test
+ void testMemorySetManagementIsRefusedOffTheMailboxThread() {
+ Mem0LongTermMemory guarded =
+ new Mem0LongTermMemory(
+ mockAdapter,
+ mockPyMem0,
+ () -> {
+ throw new IllegalStateException(
+ "Expected to be running on the task
mailbox thread, but was"
+ + " not.");
+ });
+ // No context switch: with no key in scope, only a checker that runs
before the key
+ // is read can produce the mailbox-thread message. That pins the
intended order,
+ // because off the mailbox thread the key read is itself unreliable.
+
+ assertThatThrownBy(() -> guarded.getMemorySet("notes"))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("task mailbox thread");
+ assertThatThrownBy(() -> guarded.deleteMemorySet("notes"))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("task mailbox thread");
+ verify(mockAdapter, never()).callMethod(eq(mockPyMem0),
eq("delete_memory_set"), any());
+ }
+
+ @Test
+ void testSetScopedOperationsRunWithoutTheMailboxThread() {
+ // A set carries the context it was obtained under, so operations on
it are safe to
+ // forward to a worker thread. Gating them would break durable async
execution.
+ AtomicInteger checkerCalls = new AtomicInteger();
+ Mem0LongTermMemory counting =
+ new Mem0LongTermMemory(mockAdapter, mockPyMem0,
checkerCalls::incrementAndGet);
+ counting.switchContext("a-key", "an-action", false);
+ MemorySet ms = counting.getMemorySet("notes");
+ // Guards the assertion below from passing vacuously on a checker that
never runs.
+ assertThat(checkerCalls.get()).isOne();
+ checkerCalls.set(0);
+
+ counting.add(ms, List.of("hello"), null);
+ counting.get(ms, null, null, null);
+ counting.delete(ms, null);
+ counting.search(ms, "query", 5, null, Map.of());
+
+ assertThat(checkerCalls.get()).isZero();
+ }
}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java
index b92f887e3..99bca76ff 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java
@@ -804,7 +804,7 @@ public class ActionExecutionOperatorTest {
private RuntimeException drainFailure;
private RecordingMem0LongTermMemory() {
- super(null, null);
+ super(null, null, () -> {});
}
@Override
