From 4b2bf6303a3b2d94eff9339ab781095ac1ff013c Mon Sep 17 00:00:00 2001 From: weiqingy Date: Tue, 11 Aug 2026 21:26:21 -0700 Subject: [PATCH] [runtime][api] Scope long-term memory operations to the action that obtained 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 --- .../agents/api/memory/BaseLongTermMemory.java | 9 ++ .../flink/agents/api/memory/MemorySet.java | 35 ++++++ .../development/memory/long_term_memory.md | 7 ++ .../api/memory/long_term_memory.py | 24 ++++ .../memory/mem0/mem0_long_term_memory.py | 58 +++++++--- .../mem0/tests/test_mem0_long_term_memory.py | 11 +- .../mem0/tests/test_mem0_op_recording.py | 106 +++++++++++++++++- .../flink_agents/runtime/python_java_utils.py | 21 +++- .../runtime/tests/test_python_java_utils.py | 10 ++ .../runtime/memory/Mem0LongTermMemory.java | 35 +++++- .../memory/Mem0LongTermMemoryTest.java | 33 +++++- 11 files changed, 319 insertions(+), 30 deletions(-) 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..690513988 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,11 @@ public interface BaseLongTermMemory extends AutoCloseable { /** * Gets the memory set by name. If it does not exist, the backend creates it. * + *

The returned set is bound to the calling action, so call this from the action body itself, + * not from a callback running on another thread, and obtain one per action rather than holding + * one across actions. Operating on a set that carries no binding throws rather than silently + * widening the operation to every partition key. + * * @param name the name of the memory set * @return the memory set */ @@ -39,6 +44,10 @@ public interface BaseLongTermMemory extends AutoCloseable { /** * Deletes 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 rather than from another thread, and can + * target a different key than {@link MemorySet#delete} on a same-named set would. + * * @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 @@ /** * 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}. + * + *

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 void setLtm(BaseLongTermMemory ltm) { 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..27dcab87b 100644 --- a/docs/content/docs/development/memory/long_term_memory.md +++ b/docs/content/docs/development/memory/long_term_memory.md @@ -179,6 +179,13 @@ 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. +{{< /hint >}} + ### Adding Items {{< tabs "Adding Items" >}} diff --git a/python/flink_agents/api/memory/long_term_memory.py b/python/flink_agents/api/memory/long_term_memory.py index 358d820e1..e06817497 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,12 @@ 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, not from a callback running on another thread, and obtain one per + action rather than holding one across actions. Operating on a set that carries + no binding raises rather than silently widening the operation to every + partition key. + Args: name: The name of the memory set. @@ -161,6 +180,11 @@ def get_memory_set(self, name: str) -> MemorySet: 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 rather than from + another thread, and can target a different key than ``MemorySet.delete`` on a + same-named set would. + Args: name: The name of the memory set. 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..9453d63da 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 @@ -123,6 +123,23 @@ def validate_config(cls, v: Any, values: Any) -> Any: return _FlinkAgentsLlmConfig, _FlinkAgentsEmbedderConfig +def _bound_partition_key(memory_set: MemorySet) -> str: + """Return the partition key the set is scoped to. + + Mem0 ignores a falsy ``agent_id`` rather than matching on it, so 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 memory_set.partition_key + + class Mem0LongTermMemory(InternalBaseLongTermMemory): """Long-Term Memory backed by Mem0. @@ -429,18 +446,33 @@ def drain_ltm_observation_records(self, key: str, observation_id: str) -> str: 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. Must be called on the mailbox thread. + Args: name: The name of the memory set. Returns: The memory set. """ - return MemorySet(name=name, ltm=self) + return MemorySet( + name=name, + ltm=self, + partition_key=self.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 is therefore only correct on 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. @@ -485,10 +517,10 @@ def add( 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 +582,10 @@ def get( 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 +637,10 @@ def delete(self, memory_set: MemorySet, ids: str | List[str] | None = None) -> N 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 +694,10 @@ def search( 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..58a104039 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 @@ -330,11 +330,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 + # 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 # Reset context ltm.switch_context("", observation_id="action-empty") 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..d39f2083f 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 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, ) @@ -172,16 +175,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" ] @@ -196,9 +198,9 @@ def test_empty_context_key_is_used_consistently_for_mem0_operations() -> None: mem0.get_all.return_value = {"results": []} mem0.search.return_value = {"results": []} ltm = _make_ltm(mem0) - memory_set = ltm.get_memory_set("prefs") ltm.switch_context("", observation_id="empty-action") + memory_set = ltm.get_memory_set("prefs") ltm.add(memory_set, "input") ltm.get(memory_set) ltm.search(memory_set, "query", limit=5) @@ -212,3 +214,99 @@ def test_empty_context_key_is_used_consistently_for_mem0_operations() -> None: "", "", ] + + +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("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) + + 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"] 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 @@ 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 11d95ac8d..d82f58ef0 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 @@ -46,6 +46,12 @@ public class Mem0LongTermMemory implements InteranlBaseLongTermMemory { private final PythonResourceAdapter adapter; private final PyObject pyMem0; + // Defaults mirror the Python side's own defaults, so a set obtained before any + // context switch forwards the same values Python would have used itself. + private String partitionKey = ""; + private String observationId = ""; + private boolean observationSuppressed; + public Mem0LongTermMemory(PythonResourceAdapter adapter, PyObject pyMem0) { this.adapter = adapter; this.pyMem0 = pyMem0; @@ -54,14 +60,20 @@ public Mem0LongTermMemory(PythonResourceAdapter adapter, PyObject pyMem0) { @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. MemorySet ms = new MemorySet(name); ms.setLtm(this); + ms.setActionContext(partitionKey, 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. return (Boolean) adapter.callMethod(pyMem0, "delete_memory_set", Map.of("name", name)); } @@ -149,6 +161,9 @@ public void configureObservation( @Override public void switchContext( String partitionKey, String observationId, boolean observationSuppressed) { + this.partitionKey = partitionKey; + this.observationId = observationId; + this.observationSuppressed = observationSuppressed; adapter.callMethod( pyMem0, "switch_context", @@ -173,7 +188,23 @@ public void close() { } 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 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())); + } + return adapter.invoke( + TO_PYTHON_MEMORY_SET, + memorySet.getName(), + memorySet.getPartitionKey(), + memorySet.getObservationId(), + memorySet.isObservationSuppressed()); } @SuppressWarnings("unchecked") 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 f2575be45..46b9ffe02 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 @@ -32,8 +32,10 @@ import java.util.Map; 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.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -49,7 +51,8 @@ public class Mem0LongTermMemoryTest { void setUp() { mocks = MockitoAnnotations.openMocks(this); ltm = new Mem0LongTermMemory(mockAdapter, mockPyMem0); - when(mockAdapter.invoke(eq("python_java_utils.to_python_memory_set"), any())) + when(mockAdapter.invoke( + eq("python_java_utils.to_python_memory_set"), any(), any(), any(), any())) .thenReturn(mockPyMemorySet); } @@ -220,4 +223,32 @@ void testSwitchContextAndCloseForward() { eq(Map.of("key", "k1", "observation_id", "observation-1"))); verify(mockAdapter).callMethod(eq(mockPyMem0), eq("close"), eq(Map.of())); } + + @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()).callMethod(eq(mockPyMem0), eq("add"), any()); + } + + @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)); + } }