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

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

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
      */

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

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

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

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.
 

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(),
     )