[runtime][api] Scope long-term memory operations to the action that obtained the memory set

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.
      *
+     * <p>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.
      *
+     * <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 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
      */

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}.
+ *
+ * <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 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;

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" >}}

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.
 

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,

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")