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. * + *
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 2cfdb87daf..0af97ae63e 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 666c697b9e..4733900173 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 have non-unique `toString()` or `__str__()` implementations. Key representations should therefore be stable and unique, and this isolation should not be treated as a security boundary. + +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 memories are already stored. 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 65360910ef..19c078a01f 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 @@ */ 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 void testWorkflowWithSkills() throws Exception { @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 8f6ef17474..c6dce1550b 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 Mem0LongTermMemoryTest() throws IOException { 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 358d820e16..c3ecd8921f 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 @@ 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; 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 284b3c6eb7..cf44c10bee 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 db87a3470d..03a14f9e35 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 @@ 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 validate_config(cls, v: Any, values: Any) -> Any: 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 @@ def __init__( 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 @@ def __init__( 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 @@ def __init__( ) 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 @@ def drain_ltm_observation_records(self, key: str, observation_id: str) -> str: 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 @@ 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 +643,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 +698,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 +755,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 3f3faf4e59..bae18aca26 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 get_resource(name: str, resource_type: Any, **kwargs: Any) -> Any: 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 get_resource(name: str, resource_type: Any, **kwargs: Any) -> Any: 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 4b83aec254..fa729ecc67 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, ) @@ -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 765efe7047..5c02b64a60 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 fc1b066622..e1ecac32b5 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 20dce808b4..d4264d46f8 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 f414422cce..bc4d16a195 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 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", @@ -180,7 +204,46 @@ public void close() throws Exception { } 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 0383faee4f..ba44a889a2 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 @@ void open( initPythonActionExecutor(agentPlan, jobIdentifier); } if (mem0Configured) { - wireLongTermMemory(agentPlan); + wireLongTermMemory(agentPlan, mailboxThreadChecker); } initialized = true; } @@ -240,7 +240,7 @@ private boolean isMem0Configured(AgentPlan agentPlan) { * {@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 @@ private void wireLongTermMemory(AgentPlan agentPlan) { 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 72498d77ce..1554be965a 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 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 @@ void testCloseReleasesPythonObjectWhenLogicalCleanupFails() throws Exception { 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 b92f887e3f..99bca76ffa 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 @@ private static final class RecordingMem0LongTermMemory extends Mem0LongTermMemor private RuntimeException drainFailure; private RecordingMem0LongTermMemory() { - super(null, null); + super(null, null, () -> {}); } @Override