This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new 25282b57 [Feature][api][runtime] Support auto resolve memory reference
for passing data across actions. (#950)
25282b57 is described below
commit 25282b57bdf71b6c1c085059088367b0c131ce1d
Author: Jinkun Liu <[email protected]>
AuthorDate: Tue Sep 15 20:50:20 2026 -0700
[Feature][api][runtime] Support auto resolve memory reference for passing
data across actions. (#950)
---
.../java/org/apache/flink/agents/api/Event.java | 84 +++++-
.../org/apache/flink/agents/api/OutputEvent.java | 7 +-
.../flink/agents/api/context/MemoryObject.java | 8 +-
.../apache/flink/agents/api/context/MemoryRef.java | 80 +++++-
.../agents/api/CrossLanguageEventSnapshotTest.java | 26 ++
.../org/apache/flink/agents/api/EventTest.java | 49 ++++
.../agents/api/context/MemoryRefJsonTest.java | 57 ++++
.../java/agent_run_begin_event.json | 1 +
.../java/chat_request_event.json | 1 +
.../java/chat_response_event.json | 1 +
.../java/context_retrieval_request_event.json | 1 +
.../java/context_retrieval_response_event.json | 1 +
.../java/generic_event_with_attrs.json | 3 +-
.../java/input_event.json | 1 +
.../java/memory_ref_attachment_event.json | 14 +
.../java/output_event.json | 1 +
.../java/short_term_write_event.json | 1 +
.../java/tool_request_event.json | 1 +
.../java/tool_response_event.json | 1 +
.../python/agent_run_begin_event.json | 3 +-
.../python/chat_request_event.json | 3 +-
.../python/chat_response_event.json | 3 +-
.../python/context_retrieval_request_event.json | 3 +-
.../python/context_retrieval_response_event.json | 3 +-
.../python/generic_event_with_attrs.json | 3 +-
.../python/input_event.json | 3 +-
.../python/memory_ref_attachment_event.json | 14 +
.../python/output_event.json | 3 +-
.../python/python_only_subclass_event.json | 3 +-
.../python/short_term_write_event.json | 3 +-
.../python/tool_request_event.json | 3 +-
.../python/tool_response_event.json | 3 +-
python/flink_agents/api/events/event.py | 47 +++-
python/flink_agents/api/memory_object.py | 2 +-
python/flink_agents/api/memory_reference.py | 12 +-
.../tests/test_cross_language_event_snapshots.py | 27 ++
python/flink_agents/api/tests/test_event.py | 42 ++-
.../event_attachments_test.py | 103 +++++++
.../flink_agents/runtime/flink_runner_context.py | 2 +
.../runtime/memory/event_attachment_utils.py | 108 +++++++
python/flink_agents/runtime/python_java_utils.py | 7 +
.../runtime/tests/test_event_attachment_utils.py | 178 ++++++++++++
.../runtime/tests/test_flink_runner_context.py | 68 +++++
.../agents/runtime/context/RunnerContextImpl.java | 6 +
.../runtime/memory/EventAttachmentUtils.java | 146 ++++++++++
.../runtime/operator/ActionExecutionOperator.java | 10 +
.../agents/runtime/operator/JavaActionTask.java | 20 +-
.../runtime/python/utils/PythonActionExecutor.java | 6 +
.../runtime/actionstate/ActionStateSerdeTest.java | 8 +
.../runtime/memory/EventAttachmentUtilsTest.java | 310 +++++++++++++++++++++
.../operator/ActionExecutionOperatorTest.java | 80 ++++++
.../runtime/operator/JavaActionTaskTest.java | 188 +++++++++++++
.../python/utils/PythonActionExecutorTest.java | 36 +++
53 files changed, 1751 insertions(+), 43 deletions(-)
diff --git a/api/src/main/java/org/apache/flink/agents/api/Event.java
b/api/src/main/java/org/apache/flink/agents/api/Event.java
index 43286ac8..83538fe8 100644
--- a/api/src/main/java/org/apache/flink/agents/api/Event.java
+++ b/api/src/main/java/org/apache/flink/agents/api/Event.java
@@ -22,7 +22,13 @@ import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.core.JsonParser;
+import com.fasterxml.jackson.databind.DeserializationContext;
+import com.fasterxml.jackson.databind.JsonDeserializer;
+import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
+import org.apache.flink.agents.api.context.MemoryRef;
import javax.annotation.Nullable;
@@ -42,6 +48,11 @@ public class Event {
private final String type;
private final Map<String, Object> attributes;
+ // Keep the annotation on the field as well as the creator parameter so it
also applies when
+ // Jackson constructs Event subclasses whose creators do not declare
attachments.
+ @JsonDeserialize(contentUsing = AttachmentValueDeserializer.class)
+ private final Map<String, Object> attachments;
+
@Nullable private UUID upstreamEventId;
@Nullable private String upstreamActionName;
@@ -53,7 +64,7 @@ public class Event {
/** Unified event with user-defined type and attributes. */
public Event(String type, Map<String, Object> attributes) {
- this(UUID.randomUUID(), type, attributes);
+ this(UUID.randomUUID(), type, attributes, new HashMap<>());
}
/** Unified event with user-defined type and empty attributes. */
@@ -70,6 +81,7 @@ public class Event {
* @param id the existing Event ID
* @param type the Event type used for routing
* @param attributes the Event payload
+ * @param attachments key-value data passed between Actions through
sensory memory
* @param upstreamEventId the ID of the direct upstream Event, or {@code
null}
* @param upstreamActionName the name of the emitting Action, or {@code
null}
*/
@@ -78,6 +90,9 @@ public class Event {
@JsonProperty("id") UUID id,
@JsonProperty("type") String type,
@JsonProperty("attributes") Map<String, Object> attributes,
+ @JsonProperty("attachments")
+ @JsonDeserialize(contentUsing =
AttachmentValueDeserializer.class)
+ Map<String, Object> attachments,
@JsonProperty("upstreamEventId") @Nullable UUID upstreamEventId,
@JsonProperty("upstreamActionName") @Nullable String
upstreamActionName) {
if (type == null || type.isEmpty()) {
@@ -86,14 +101,31 @@ public class Event {
// Explicit null matches an omitted id: both mint a per-occurrence
UUID.
this.id = id != null ? id : UUID.randomUUID();
this.type = type;
- this.attributes = attributes != null ? attributes : new HashMap<>();
+ this.attributes = attributes != null ? new HashMap<>(attributes) : new
HashMap<>();
+ this.attachments = attachments != null ? new HashMap<>(attachments) :
new HashMap<>();
this.upstreamEventId = upstreamEventId;
this.upstreamActionName = upstreamActionName;
}
+ /** Reconstructs an Event with an existing identity, attachments, and no
upstream lineage. */
+ public Event(
+ UUID id, String type, Map<String, Object> attributes, Map<String,
Object> attachments) {
+ this(id, type, attributes, attachments, null, null);
+ }
+
+ /** Reconstructs an Event with an existing identity and optional
framework-managed lineage. */
+ public Event(
+ UUID id,
+ String type,
+ Map<String, Object> attributes,
+ @Nullable UUID upstreamEventId,
+ @Nullable String upstreamActionName) {
+ this(id, type, attributes, new HashMap<>(), upstreamEventId,
upstreamActionName);
+ }
+
/** Reconstructs an Event with an existing identity and no upstream
lineage. */
public Event(UUID id, String type, Map<String, Object> attributes) {
- this(id, type, attributes, null, null);
+ this(id, type, attributes, new HashMap<>(), null, null);
}
public UUID getId() {
@@ -110,6 +142,10 @@ public class Event {
return attributes;
}
+ public Map<String, Object> getAttachments() {
+ return attachments;
+ }
+
/** Returns the ID of the Event consumed by the Action that emitted this
Event. */
@Nullable
@JsonInclude(JsonInclude.Include.NON_NULL)
@@ -150,6 +186,20 @@ public class Event {
attributes.put(name, value);
}
+ public Object getAttachment(String name) {
+ return attachments.get(name);
+ }
+
+ /**
+ * Sets an attachment on this Event.
+ *
+ * <p>If {@code value} is a {@link MemoryRef}, it must reference sensory
memory and will not be
+ * wrapped again.
+ */
+ public void setAttachment(String name, Object value) {
+ attachments.put(name, value);
+ }
+
@JsonIgnore
public boolean hasSourceTimestamp() {
return sourceTimestamp != null;
@@ -166,9 +216,9 @@ public class Event {
}
/**
- * Creates a base Event from another Event, copying its identity, data,
and framework metadata.
- * Subclasses override this to reconstruct typed event objects with proper
field
- * deserialization.
+ * Creates a base Event from another Event, copying its identity, data,
attachments, and
+ * framework metadata. Subclasses override this to reconstruct typed event
objects with proper
+ * field deserialization.
*/
public static Event fromEvent(Event event) {
return reconstructFrom(
@@ -194,6 +244,8 @@ public class Event {
+ source.getId());
}
Event reconstructedEvent = reconstructed;
+ reconstructedEvent.attachments.clear();
+ reconstructedEvent.attachments.putAll(source.attachments);
reconstructedEvent.sourceTimestamp = source.sourceTimestamp;
reconstructedEvent.upstreamEventId = source.upstreamEventId;
reconstructedEvent.upstreamActionName = source.upstreamActionName;
@@ -211,6 +263,21 @@ public class Event {
return MAPPER.readValue(json, Event.class);
}
+ /** Deserializes one attachment value, preserving explicitly tagged memory
references. */
+ static final class AttachmentValueDeserializer extends
JsonDeserializer<Object> {
+
+ @Override
+ public Object deserialize(JsonParser parser, DeserializationContext
context)
+ throws IOException {
+ JsonNode node = parser.getCodec().readTree(parser);
+ if (node.isObject()
+ &&
MemoryRef.TYPE_VALUE.equals(node.path(MemoryRef.TYPE_FIELD).asText())) {
+ return parser.getCodec().treeToValue(node, MemoryRef.class);
+ }
+ return parser.getCodec().treeToValue(node, Object.class);
+ }
+ }
+
@Override
public boolean equals(Object o) {
if (this == o) return true;
@@ -218,11 +285,12 @@ public class Event {
Event other = (Event) o;
return Objects.equals(this.id, other.id)
&& Objects.equals(this.getType(), other.getType())
- && Objects.equals(this.attributes, other.attributes);
+ && Objects.equals(this.attributes, other.attributes)
+ && Objects.equals(this.attachments, other.attachments);
}
@Override
public int hashCode() {
- return Objects.hash(id, getType(), attributes);
+ return Objects.hash(id, getType(), attributes, attachments);
}
}
diff --git a/api/src/main/java/org/apache/flink/agents/api/OutputEvent.java
b/api/src/main/java/org/apache/flink/agents/api/OutputEvent.java
index 7c57d634..c139deee 100644
--- a/api/src/main/java/org/apache/flink/agents/api/OutputEvent.java
+++ b/api/src/main/java/org/apache/flink/agents/api/OutputEvent.java
@@ -27,7 +27,8 @@ import java.util.UUID;
/**
* Event representing a result from agent. By generating an OutputEvent,
actions can emit output
- * data.
+ * data. Attachments are only supported on events passed between actions and
cannot be carried by an
+ * OutputEvent.
*/
public class OutputEvent extends Event {
@@ -50,8 +51,12 @@ public class OutputEvent extends Event {
*
* @param event the base event containing the output data in attributes
* @return a typed OutputEvent
+ * @throws IllegalArgumentException if the source Event carries attachments
*/
public static OutputEvent fromEvent(Event event) {
+ if (!event.getAttachments().isEmpty()) {
+ throw new IllegalArgumentException("OutputEvent cannot carry
attachments.");
+ }
return reconstructFrom(event, OutputEvent::new);
}
diff --git
a/api/src/main/java/org/apache/flink/agents/api/context/MemoryObject.java
b/api/src/main/java/org/apache/flink/agents/api/context/MemoryObject.java
index 0f18471c..2b19062e 100644
--- a/api/src/main/java/org/apache/flink/agents/api/context/MemoryObject.java
+++ b/api/src/main/java/org/apache/flink/agents/api/context/MemoryObject.java
@@ -21,9 +21,9 @@ import java.util.List;
import java.util.Map;
/**
- * A representation of an object in the short-term memory. It is responsible
for accessing and
- * manipulating (direct or indirect) fields within the memory structure. A
direct field is a field
- * which stores primitive data directly, while an indirect filed is a field
which represents a
+ * A representation of an object in sensory or short-term memory. It is
responsible for accessing
+ * and manipulating (direct or indirect) fields within the memory structure. A
direct field is a
+ * field which stores primitive data directly, while an indirect filed is a
field which represents a
* nested object.Fields can be accessed using an absolute or relative path.
*/
public interface MemoryObject {
@@ -46,7 +46,7 @@ public interface MemoryObject {
/**
* Returns a MemoryObject that represents the path of the given reference.
*
- * @param ref a reference to a data item in the short-term memory
+ * @param ref a reference to a data item in the memory
* @return a MemoryObject instance pointing to the referenced field.
Returns null if the path of
* the reference does not exist.
* @throws Exception if an error occurs during state access
diff --git
a/api/src/main/java/org/apache/flink/agents/api/context/MemoryRef.java
b/api/src/main/java/org/apache/flink/agents/api/context/MemoryRef.java
index df6f6c18..062865f3 100644
--- a/api/src/main/java/org/apache/flink/agents/api/context/MemoryRef.java
+++ b/api/src/main/java/org/apache/flink/agents/api/context/MemoryRef.java
@@ -17,17 +17,38 @@
*/
package org.apache.flink.agents.api.context;
+import com.fasterxml.jackson.core.JsonGenerator;
+import com.fasterxml.jackson.core.JsonParser;
+import com.fasterxml.jackson.databind.DeserializationContext;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.SerializerProvider;
+import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
+import com.fasterxml.jackson.databind.annotation.JsonSerialize;
+import com.fasterxml.jackson.databind.deser.std.StdDeserializer;
+import com.fasterxml.jackson.databind.ser.std.StdSerializer;
+
+import java.io.IOException;
import java.io.Serializable;
+import java.util.LinkedHashMap;
+import java.util.Locale;
+import java.util.Map;
import java.util.Objects;
/**
- * A serializable, persistent reference to a specific data item in Short-Term
Memory. It acts as a
- * lightweight pointer, containing the path of the data, allowing for
efficient passing of large
- * objects between Actions.
+ * A serializable, persistent reference to a specific data item in a {@link
MemoryObject}. It acts
+ * as a lightweight pointer, containing the path of the data, allowing for
efficient passing of
+ * large objects between Actions.
*/
+@JsonSerialize(using = MemoryRef.Serializer.class)
+@JsonDeserialize(using = MemoryRef.Deserializer.class)
public final class MemoryRef implements Serializable {
private static final long serialVersionUID = 1L;
+ public static final String TYPE_FIELD = "@type";
+ public static final String TYPE_VALUE = "memory_ref";
+ public static final String MEMORY_TYPE_FIELD = "memory_type";
+ public static final String PATH_FIELD = "path";
+
private final MemoryObject.MemoryType type;
private final String path;
@@ -39,7 +60,7 @@ public final class MemoryRef implements Serializable {
/**
* Creates a new MemoryRef instance with the given path.
*
- * @param path The absolute path of the data in Short-Term Memory.
+ * @param path The absolute path of the data in memory.
* @return A new MemoryRef instance.
*/
public static MemoryRef create(MemoryObject.MemoryType type, String path) {
@@ -49,7 +70,7 @@ public final class MemoryRef implements Serializable {
/**
* Resolves the reference using the provided RunnerContext to get the
actual data.
*
- * @param ctx The current execution context, used to access Short-Term
Memory.
+ * @param ctx The current execution context, used to access memory.
* @return The deserialized, original data object.
* @throws Exception if the memory cannot be accessed or the data cannot
be resolved.
*/
@@ -67,6 +88,55 @@ public final class MemoryRef implements Serializable {
return path;
}
+ public MemoryObject.MemoryType getType() {
+ return type;
+ }
+
+ /** Serializes a {@link MemoryRef} to JSON. */
+ public static final class Serializer extends StdSerializer<MemoryRef> {
+
+ public Serializer() {
+ super(MemoryRef.class);
+ }
+
+ @Override
+ public void serialize(MemoryRef value, JsonGenerator generator,
SerializerProvider provider)
+ throws IOException {
+ Map<String, String> serialized = new LinkedHashMap<>();
+ serialized.put(TYPE_FIELD, TYPE_VALUE);
+ serialized.put(MEMORY_TYPE_FIELD,
value.getType().name().toLowerCase(Locale.ROOT));
+ serialized.put(PATH_FIELD, value.getPath());
+ generator.writeObject(serialized);
+ }
+ }
+
+ /** Deserializes a {@link MemoryRef} from JSON. */
+ public static final class Deserializer extends StdDeserializer<MemoryRef> {
+
+ public Deserializer() {
+ super(MemoryRef.class);
+ }
+
+ @Override
+ public MemoryRef deserialize(JsonParser parser, DeserializationContext
context)
+ throws IOException {
+ JsonNode node = parser.getCodec().readTree(parser);
+ JsonNode typeNode = node.get(MEMORY_TYPE_FIELD);
+ JsonNode pathNode = node.get(PATH_FIELD);
+ if (typeNode == null || typeNode.isNull() || pathNode == null ||
pathNode.isNull()) {
+ throw new IllegalArgumentException(
+ "MemoryRef JSON must contain non-null '"
+ + MEMORY_TYPE_FIELD
+ + "' and '"
+ + PATH_FIELD
+ + "' fields.");
+ }
+ MemoryObject.MemoryType memoryType =
+
MemoryObject.MemoryType.valueOf(typeNode.asText().toUpperCase(Locale.ROOT));
+ return create(memoryType, pathNode.asText());
+ }
+ }
+
@Override
public boolean equals(Object o) {
if (this == o) return true;
diff --git
a/api/src/test/java/org/apache/flink/agents/api/CrossLanguageEventSnapshotTest.java
b/api/src/test/java/org/apache/flink/agents/api/CrossLanguageEventSnapshotTest.java
index b09e6a13..1a6e42da 100644
---
a/api/src/test/java/org/apache/flink/agents/api/CrossLanguageEventSnapshotTest.java
+++
b/api/src/test/java/org/apache/flink/agents/api/CrossLanguageEventSnapshotTest.java
@@ -23,6 +23,8 @@ import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.flink.agents.api.agents.OutputSchema;
import org.apache.flink.agents.api.chat.messages.ChatMessage;
import org.apache.flink.agents.api.chat.messages.MessageRole;
+import org.apache.flink.agents.api.context.MemoryObject;
+import org.apache.flink.agents.api.context.MemoryRef;
import org.apache.flink.agents.api.event.AgentRunBeginEvent;
import org.apache.flink.agents.api.event.ChatRequestEvent;
import org.apache.flink.agents.api.event.ChatResponseEvent;
@@ -52,6 +54,7 @@ import java.util.UUID;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assumptions.assumeTrue;
@@ -68,6 +71,7 @@ class CrossLanguageEventSnapshotTest {
private static final String FIXED_TOOL_CALL_ID = "call_aaaa";
private static final String FIXED_TOOL_CALL_ID_NUMERIC = "call_bbbb";
private static final String FIXED_TOOL_CALL_ID_BOOL = "call_cccc";
+ private static final String MEMORY_REF_ATTACHMENT_EVENT_TYPE =
"_memory_ref_attachment_event";
private static final long FIXED_TIMESTAMP = 1_700_000_000_000L;
private static Path snapshotDir;
@@ -579,6 +583,28 @@ class CrossLanguageEventSnapshotTest {
assertEquals(7, attrs.get("count"));
}
+ @Test
+ void memoryRefAttachmentCrossLanguageSnapshotIsStable() throws Exception {
+ Event javaEvent =
+ new Event(
+ FIXED_EVENT_ID, MEMORY_REF_ATTACHMENT_EVENT_TYPE,
Map.of("value", "ping"));
+ javaEvent.setAttachment(
+ "payload", MemoryRef.create(MemoryObject.MemoryType.SENSORY,
"memory.path"));
+ if (regenerateRequested()) {
+ writeJavaSnapshot("memory_ref_attachment_event.json", javaEvent);
+ }
+ assertJavaSnapshotStable("memory_ref_attachment_event.json",
javaEvent);
+
+ Event pythonEvent =
readPythonSnapshot("memory_ref_attachment_event.json");
+ assertEquals(FIXED_EVENT_ID, pythonEvent.getId());
+ assertEquals(MEMORY_REF_ATTACHMENT_EVENT_TYPE, pythonEvent.getType());
+ assertEquals("ping", pythonEvent.getAttr("value"));
+ MemoryRef reference =
+ assertInstanceOf(MemoryRef.class,
pythonEvent.getAttachment("payload"));
+ assertEquals(MemoryObject.MemoryType.SENSORY, reference.getType());
+ assertEquals("memory.path", reference.getPath());
+ }
+
// ── Smoke ──────────────────────────────────────────────────────────────
@Test
diff --git a/api/src/test/java/org/apache/flink/agents/api/EventTest.java
b/api/src/test/java/org/apache/flink/agents/api/EventTest.java
index a87357f2..56439698 100644
--- a/api/src/test/java/org/apache/flink/agents/api/EventTest.java
+++ b/api/src/test/java/org/apache/flink/agents/api/EventTest.java
@@ -20,6 +20,8 @@ package org.apache.flink.agents.api;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.flink.agents.api.context.MemoryObject;
+import org.apache.flink.agents.api.context.MemoryRef;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
@@ -145,6 +147,38 @@ class EventTest {
assertNull(event.getUpstreamActionName());
}
+ @Test
+ void testUnifiedEventJsonRoundTripPreservesAttachments() throws Exception {
+ Map<String, Object> attachments = new HashMap<>();
+ attachments.put("payload", Map.of("memory_type", "sensory", "path",
"memory.path"));
+ Event original = new Event(UUID.randomUUID(), "MyEvent", new
HashMap<>(), attachments);
+
+ Event restored =
+
objectMapper.readValue(objectMapper.writeValueAsString(original), Event.class);
+
+ assertEquals(attachments, restored.getAttachments());
+ }
+
+ @Test
+ void testUnifiedEventJsonRoundTripPreservesMemoryRefAttachments() throws
Exception {
+ MemoryRef reference =
MemoryRef.create(MemoryObject.MemoryType.SENSORY, "memory.path");
+ Event original =
+ new Event(
+ UUID.randomUUID(),
+ "MyEvent",
+ new HashMap<>(),
+ Map.of("payload", reference));
+
+ String json = objectMapper.writeValueAsString(original);
+ JsonNode attachment =
objectMapper.readTree(json).get("attachments").get("payload");
+ Event restored = objectMapper.readValue(json, Event.class);
+
+ assertEquals("memory_ref", attachment.get("@type").asText());
+ assertEquals("sensory", attachment.get("memory_type").asText());
+ assertEquals("memory.path", attachment.get("path").asText());
+ assertEquals(reference, restored.getAttachment("payload"));
+ }
+
@Test
void testSubclassedEventJsonRoundTrip() throws Exception {
InputEvent original = new InputEvent("round trip");
@@ -317,6 +351,21 @@ class EventTest {
assertEquals("output_action", copy.getUpstreamActionName());
}
+ @Test
+ void testOutputEventFromEventRejectsAttachments() {
+ Event original =
+ new Event(
+ UUID.randomUUID(),
+ OutputEvent.EVENT_TYPE,
+ Map.of("output", "result"),
+ Map.of("payload", "attachment"));
+
+ IllegalArgumentException error =
+ assertThrows(IllegalArgumentException.class, () ->
OutputEvent.fromEvent(original));
+
+ assertEquals("OutputEvent cannot carry attachments.",
error.getMessage());
+ }
+
@Test
void
testCustomFactoryReconstructsSameOccurrenceWithoutManualMetadataCopy() {
UUID upstreamEventId = UUID.randomUUID();
diff --git
a/api/src/test/java/org/apache/flink/agents/api/context/MemoryRefJsonTest.java
b/api/src/test/java/org/apache/flink/agents/api/context/MemoryRefJsonTest.java
new file mode 100644
index 00000000..ee8359b1
--- /dev/null
+++
b/api/src/test/java/org/apache/flink/agents/api/context/MemoryRefJsonTest.java
@@ -0,0 +1,57 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.flink.agents.api.context;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+class MemoryRefJsonTest {
+
+ private final ObjectMapper objectMapper = new ObjectMapper();
+
+ @Test
+ void serializesAndDeserializesMemoryRef() throws Exception {
+ MemoryRef original = MemoryRef.create(MemoryObject.MemoryType.SENSORY,
"memory.path");
+
+ String json = objectMapper.writeValueAsString(original);
+ JsonNode node = objectMapper.readTree(json);
+ MemoryRef restored = objectMapper.readValue(json, MemoryRef.class);
+
+ assertEquals("memory_ref", node.get("@type").asText());
+ assertEquals("sensory", node.get("memory_type").asText());
+ assertEquals("memory.path", node.get("path").asText());
+ assertEquals(original, restored);
+ assertEquals(MemoryObject.MemoryType.SENSORY, restored.getType());
+ }
+
+ @Test
+ void rejectsInvalidMemoryRef() throws JsonProcessingException {
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ objectMapper.readValue(
+
"{\"memory_type\":\"unknown\",\"path\":\"memory.path\"}",
+ MemoryRef.class),
+ "No enum constant
org.apache.flink.agents.api.context.MemoryObject.MemoryType.UNKNOWN");
+ }
+}
diff --git
a/e2e-test/cross-language-event-snapshots/java/agent_run_begin_event.json
b/e2e-test/cross-language-event-snapshots/java/agent_run_begin_event.json
index 5c050f0a..096ab944 100644
--- a/e2e-test/cross-language-event-snapshots/java/agent_run_begin_event.json
+++ b/e2e-test/cross-language-event-snapshots/java/agent_run_begin_event.json
@@ -7,5 +7,6 @@
"user.address.city" : "SF"
}
},
+ "attachments" : { },
"type" : "_agent_run_begin_event"
}
diff --git
a/e2e-test/cross-language-event-snapshots/java/chat_request_event.json
b/e2e-test/cross-language-event-snapshots/java/chat_request_event.json
index 347c47e7..1a233418 100644
--- a/e2e-test/cross-language-event-snapshots/java/chat_request_event.json
+++ b/e2e-test/cross-language-event-snapshots/java/chat_request_event.json
@@ -9,5 +9,6 @@
"extra_args" : { }
} ]
},
+ "attachments" : { },
"type" : "_chat_request_event"
}
diff --git
a/e2e-test/cross-language-event-snapshots/java/chat_response_event.json
b/e2e-test/cross-language-event-snapshots/java/chat_response_event.json
index 3d5b4793..562b034b 100644
--- a/e2e-test/cross-language-event-snapshots/java/chat_response_event.json
+++ b/e2e-test/cross-language-event-snapshots/java/chat_response_event.json
@@ -11,5 +11,6 @@
"retry_count" : 0,
"total_retry_wait_sec" : 0
},
+ "attachments" : { },
"type" : "_chat_response_event"
}
diff --git
a/e2e-test/cross-language-event-snapshots/java/context_retrieval_request_event.json
b/e2e-test/cross-language-event-snapshots/java/context_retrieval_request_event.json
index ead03f8d..e9f868cd 100644
---
a/e2e-test/cross-language-event-snapshots/java/context_retrieval_request_event.json
+++
b/e2e-test/cross-language-event-snapshots/java/context_retrieval_request_event.json
@@ -5,5 +5,6 @@
"vector_store" : "test-store",
"max_results" : 5
},
+ "attachments" : { },
"type" : "_context_retrieval_request_event"
}
diff --git
a/e2e-test/cross-language-event-snapshots/java/context_retrieval_response_event.json
b/e2e-test/cross-language-event-snapshots/java/context_retrieval_response_event.json
index 90592d56..b7cd0caa 100644
---
a/e2e-test/cross-language-event-snapshots/java/context_retrieval_response_event.json
+++
b/e2e-test/cross-language-event-snapshots/java/context_retrieval_response_event.json
@@ -13,5 +13,6 @@
"score" : null
} ]
},
+ "attachments" : { },
"type" : "_context_retrieval_response_event"
}
diff --git
a/e2e-test/cross-language-event-snapshots/java/generic_event_with_attrs.json
b/e2e-test/cross-language-event-snapshots/java/generic_event_with_attrs.json
index 96b9fc0f..91f7216b 100644
--- a/e2e-test/cross-language-event-snapshots/java/generic_event_with_attrs.json
+++ b/e2e-test/cross-language-event-snapshots/java/generic_event_with_attrs.json
@@ -11,5 +11,6 @@
"k_dict" : {
"nested" : "value"
}
- }
+ },
+ "attachments" : { }
}
diff --git a/e2e-test/cross-language-event-snapshots/java/input_event.json
b/e2e-test/cross-language-event-snapshots/java/input_event.json
index 8150a1ce..72801563 100644
--- a/e2e-test/cross-language-event-snapshots/java/input_event.json
+++ b/e2e-test/cross-language-event-snapshots/java/input_event.json
@@ -3,5 +3,6 @@
"attributes" : {
"input" : "hello"
},
+ "attachments" : { },
"type" : "_input_event"
}
diff --git
a/e2e-test/cross-language-event-snapshots/java/memory_ref_attachment_event.json
b/e2e-test/cross-language-event-snapshots/java/memory_ref_attachment_event.json
new file mode 100644
index 00000000..43956e65
--- /dev/null
+++
b/e2e-test/cross-language-event-snapshots/java/memory_ref_attachment_event.json
@@ -0,0 +1,14 @@
+{
+ "id" : "00000000-0000-0000-0000-000000000001",
+ "type" : "_memory_ref_attachment_event",
+ "attributes" : {
+ "value" : "ping"
+ },
+ "attachments" : {
+ "payload" : {
+ "@type" : "memory_ref",
+ "memory_type" : "sensory",
+ "path" : "memory.path"
+ }
+ }
+}
diff --git a/e2e-test/cross-language-event-snapshots/java/output_event.json
b/e2e-test/cross-language-event-snapshots/java/output_event.json
index 3fb4269e..3d7941ed 100644
--- a/e2e-test/cross-language-event-snapshots/java/output_event.json
+++ b/e2e-test/cross-language-event-snapshots/java/output_event.json
@@ -3,5 +3,6 @@
"attributes" : {
"output" : "world"
},
+ "attachments" : { },
"type" : "_output_event"
}
diff --git
a/e2e-test/cross-language-event-snapshots/java/short_term_write_event.json
b/e2e-test/cross-language-event-snapshots/java/short_term_write_event.json
index a6f906c2..f62d61b1 100644
--- a/e2e-test/cross-language-event-snapshots/java/short_term_write_event.json
+++ b/e2e-test/cross-language-event-snapshots/java/short_term_write_event.json
@@ -7,5 +7,6 @@
"profile.m_a01" : null
}
},
+ "attachments" : { },
"type" : "_short_term_write_event"
}
diff --git
a/e2e-test/cross-language-event-snapshots/java/tool_request_event.json
b/e2e-test/cross-language-event-snapshots/java/tool_request_event.json
index 0f8ab2f3..0b88c932 100644
--- a/e2e-test/cross-language-event-snapshots/java/tool_request_event.json
+++ b/e2e-test/cross-language-event-snapshots/java/tool_request_event.json
@@ -10,5 +10,6 @@
}
} ]
},
+ "attachments" : { },
"type" : "_tool_request_event"
}
diff --git
a/e2e-test/cross-language-event-snapshots/java/tool_response_event.json
b/e2e-test/cross-language-event-snapshots/java/tool_response_event.json
index 04698abe..850dea38 100644
--- a/e2e-test/cross-language-event-snapshots/java/tool_response_event.json
+++ b/e2e-test/cross-language-event-snapshots/java/tool_response_event.json
@@ -18,5 +18,6 @@
"external_ids" : { },
"timestamp" : 1700000000000
},
+ "attachments" : { },
"type" : "_tool_response_event"
}
diff --git
a/e2e-test/cross-language-event-snapshots/python/agent_run_begin_event.json
b/e2e-test/cross-language-event-snapshots/python/agent_run_begin_event.json
index 79626d1d..6768af7e 100644
--- a/e2e-test/cross-language-event-snapshots/python/agent_run_begin_event.json
+++ b/e2e-test/cross-language-event-snapshots/python/agent_run_begin_event.json
@@ -7,5 +7,6 @@
"user.tier": "gold",
"user.address.city": "SF"
}
- }
+ },
+ "attachments": {}
}
diff --git
a/e2e-test/cross-language-event-snapshots/python/chat_request_event.json
b/e2e-test/cross-language-event-snapshots/python/chat_request_event.json
index ac880823..6d5ca620 100644
--- a/e2e-test/cross-language-event-snapshots/python/chat_request_event.json
+++ b/e2e-test/cross-language-event-snapshots/python/chat_request_event.json
@@ -13,5 +13,6 @@
],
"prompt_args": {},
"output_schema": null
- }
+ },
+ "attachments": {}
}
diff --git
a/e2e-test/cross-language-event-snapshots/python/chat_response_event.json
b/e2e-test/cross-language-event-snapshots/python/chat_response_event.json
index bafb2811..fae3956f 100644
--- a/e2e-test/cross-language-event-snapshots/python/chat_response_event.json
+++ b/e2e-test/cross-language-event-snapshots/python/chat_response_event.json
@@ -11,5 +11,6 @@
},
"retry_count": 0,
"total_retry_wait_sec": 0
- }
+ },
+ "attachments": {}
}
diff --git
a/e2e-test/cross-language-event-snapshots/python/context_retrieval_request_event.json
b/e2e-test/cross-language-event-snapshots/python/context_retrieval_request_event.json
index 357ce8bc..6ad4ef14 100644
---
a/e2e-test/cross-language-event-snapshots/python/context_retrieval_request_event.json
+++
b/e2e-test/cross-language-event-snapshots/python/context_retrieval_request_event.json
@@ -5,5 +5,6 @@
"query": "what is flink",
"vector_store": "test-store",
"max_results": 5
- }
+ },
+ "attachments": {}
}
diff --git
a/e2e-test/cross-language-event-snapshots/python/context_retrieval_response_event.json
b/e2e-test/cross-language-event-snapshots/python/context_retrieval_response_event.json
index 95e14f0a..3c0dc4e4 100644
---
a/e2e-test/cross-language-event-snapshots/python/context_retrieval_response_event.json
+++
b/e2e-test/cross-language-event-snapshots/python/context_retrieval_response_event.json
@@ -15,5 +15,6 @@
"score": null
}
]
- }
+ },
+ "attachments": {}
}
diff --git
a/e2e-test/cross-language-event-snapshots/python/generic_event_with_attrs.json
b/e2e-test/cross-language-event-snapshots/python/generic_event_with_attrs.json
index cfd461b3..1896e8ca 100644
---
a/e2e-test/cross-language-event-snapshots/python/generic_event_with_attrs.json
+++
b/e2e-test/cross-language-event-snapshots/python/generic_event_with_attrs.json
@@ -15,5 +15,6 @@
"k_dict": {
"nested": "value"
}
- }
+ },
+ "attachments": {}
}
diff --git a/e2e-test/cross-language-event-snapshots/python/input_event.json
b/e2e-test/cross-language-event-snapshots/python/input_event.json
index db24e4c5..35eab9af 100644
--- a/e2e-test/cross-language-event-snapshots/python/input_event.json
+++ b/e2e-test/cross-language-event-snapshots/python/input_event.json
@@ -3,5 +3,6 @@
"type": "_input_event",
"attributes": {
"input": "hello"
- }
+ },
+ "attachments": {}
}
diff --git
a/e2e-test/cross-language-event-snapshots/python/memory_ref_attachment_event.json
b/e2e-test/cross-language-event-snapshots/python/memory_ref_attachment_event.json
new file mode 100644
index 00000000..203d613c
--- /dev/null
+++
b/e2e-test/cross-language-event-snapshots/python/memory_ref_attachment_event.json
@@ -0,0 +1,14 @@
+{
+ "id": "00000000-0000-0000-0000-000000000001",
+ "type": "_memory_ref_attachment_event",
+ "attributes": {
+ "value": "ping"
+ },
+ "attachments": {
+ "payload": {
+ "@type": "memory_ref",
+ "memory_type": "sensory",
+ "path": "memory.path"
+ }
+ }
+}
diff --git a/e2e-test/cross-language-event-snapshots/python/output_event.json
b/e2e-test/cross-language-event-snapshots/python/output_event.json
index f4b48a74..7d0d2c4a 100644
--- a/e2e-test/cross-language-event-snapshots/python/output_event.json
+++ b/e2e-test/cross-language-event-snapshots/python/output_event.json
@@ -3,5 +3,6 @@
"type": "_output_event",
"attributes": {
"output": "world"
- }
+ },
+ "attachments": {}
}
diff --git
a/e2e-test/cross-language-event-snapshots/python/python_only_subclass_event.json
b/e2e-test/cross-language-event-snapshots/python/python_only_subclass_event.json
index a48448c1..087c32a4 100644
---
a/e2e-test/cross-language-event-snapshots/python/python_only_subclass_event.json
+++
b/e2e-test/cross-language-event-snapshots/python/python_only_subclass_event.json
@@ -4,5 +4,6 @@
"attributes": {
"value": "ping",
"count": 7
- }
+ },
+ "attachments": {}
}
diff --git
a/e2e-test/cross-language-event-snapshots/python/short_term_write_event.json
b/e2e-test/cross-language-event-snapshots/python/short_term_write_event.json
index 1a4ac90a..5588c317 100644
--- a/e2e-test/cross-language-event-snapshots/python/short_term_write_event.json
+++ b/e2e-test/cross-language-event-snapshots/python/short_term_write_event.json
@@ -7,5 +7,6 @@
"user.tier": "gold",
"profile.m_a01": null
}
- }
+ },
+ "attachments": {}
}
diff --git
a/e2e-test/cross-language-event-snapshots/python/tool_request_event.json
b/e2e-test/cross-language-event-snapshots/python/tool_request_event.json
index 2ac1fc51..d311b941 100644
--- a/e2e-test/cross-language-event-snapshots/python/tool_request_event.json
+++ b/e2e-test/cross-language-event-snapshots/python/tool_request_event.json
@@ -12,5 +12,6 @@
}
}
]
- }
+ },
+ "attachments": {}
}
diff --git
a/e2e-test/cross-language-event-snapshots/python/tool_response_event.json
b/e2e-test/cross-language-event-snapshots/python/tool_response_event.json
index 4ec6d0b0..65b2cbd6 100644
--- a/e2e-test/cross-language-event-snapshots/python/tool_response_event.json
+++ b/e2e-test/cross-language-event-snapshots/python/tool_response_event.json
@@ -19,5 +19,6 @@
"call_bbbb": null,
"call_cccc": null
}
- }
+ },
+ "attachments": {}
}
diff --git a/python/flink_agents/api/events/event.py
b/python/flink_agents/api/events/event.py
index 57fa10d1..ea269ca4 100644
--- a/python/flink_agents/api/events/event.py
+++ b/python/flink_agents/api/events/event.py
@@ -36,6 +36,8 @@ from pydantic import (
from pydantic_core import PydanticSerializationError
from pyflink.common import Row
+from flink_agents.api.memory_reference import MemoryRef
+
def _reconstruct_row_if_needed(data: Any) -> Any:
"""Recursively reconstruct pyflink Row objects from their JSON-serialized
dicts.
@@ -79,6 +81,8 @@ class Event(BaseModel, extra="allow"):
Event type string used for routing. Required for all events.
attributes : Dict[str, Any]
Key-value properties for the event data.
+ attachments : Dict[str, Any]
+ Key-value data passed between actions through sensory memory.
upstream_event_id : UUID | None
The ID of the direct upstream Event, or None.
upstream_action_name : str | None
@@ -88,6 +92,7 @@ class Event(BaseModel, extra="allow"):
id: UUID = Field(default_factory=uuid4, frozen=True)
type: str
attributes: Dict[str, Any] = Field(default_factory=dict)
+ attachments: Dict[str, Any] = Field(default_factory=dict)
upstream_event_id: UUID | None = Field(
default=None,
validation_alias=AliasChoices("upstream_event_id", "upstreamEventId"),
@@ -105,6 +110,20 @@ class Event(BaseModel, extra="allow"):
"""Treat explicit None like an omitted id and mint a per-occurrence
UUID."""
return uuid4() if value is None else value
+ @field_validator("attachments", mode="before")
+ @classmethod
+ def _deserialize_memory_ref_attachments(cls, attachments: Any) -> Any:
+ """Restore explicitly tagged memory-reference attachment values."""
+ if not isinstance(attachments, dict):
+ return attachments
+ return {
+ key: MemoryRef.model_validate(value)
+ if isinstance(value, dict)
+ and value.get(MemoryRef.TYPE_FIELD) == MemoryRef.TYPE_VALUE
+ else value
+ for key, value in attachments.items()
+ }
+
@staticmethod
def __serialize_unknown(field: Any) -> Dict[str, Any]:
"""Handle serialization of unknown types, specifically Row objects."""
@@ -143,20 +162,22 @@ class Event(BaseModel, extra="allow"):
@model_validator(mode="after")
def validate_serializable_fields(self) -> "Event":
- """Validate that all Event fields can be serialized."""
- self.model_dump_json()
+ """Validate JSON event fields without serializing raw attachments."""
+ self.model_dump_json(exclude={"attachments"})
return self
def __setattr__(self, name: str, value: Any) -> None:
super().__setattr__(name, value)
- # Ensure added property can be serialized.
- self.model_dump_json()
+ # Raw attachments are offloaded to sensory memory before sending.
Validate every
+ # other field here without serializing those payloads.
+ self.model_dump_json(exclude={"attachments"})
def reconstruct_from(self, source: "Event") -> Self:
"""Return a typed copy representing the same Event occurrence as
source."""
return self.model_copy(
update={
"id": source.id,
+ "attachments": dict(source.attachments),
"upstream_event_id": source.upstream_event_id,
"upstream_action_name": source.upstream_action_name,
}
@@ -174,6 +195,18 @@ class Event(BaseModel, extra="allow"):
"""Set an attribute value in the attributes map."""
self.attributes[name] = value
+ def get_attachment(self, name: str) -> Any:
+ """Get an attachment value from the attachments map."""
+ return self.attachments.get(name)
+
+ def set_attachment(self, name: str, value: Any) -> None:
+ """Set an attachment value in the attachments map.
+
+ If ``value`` is a :class:`MemoryRef`, it must reference sensory memory
and
+ will not be wrapped again.
+ """
+ self.attachments = {**self.attachments, name: value}
+
@classmethod
def from_event(cls, event: "Event") -> "Event":
"""Reconstruct a typed event from a base Event.
@@ -248,6 +281,9 @@ class OutputEvent(Event):
"""Event representing a result from agent. By generating an OutputEvent,
actions can emit output data.
+ Attachments are only supported on events passed between actions and cannot
+ be carried by an OutputEvent.
+
Attributes:
----------
output : Any
@@ -266,6 +302,9 @@ class OutputEvent(Event):
@classmethod
@override
def from_event(cls, event: Event) -> "OutputEvent":
+ if event.attachments:
+ msg = "OutputEvent cannot carry attachments."
+ raise ValueError(msg)
assert "output" in event.attributes
result = OutputEvent(output=event.attributes["output"])
return result.reconstruct_from(event)
diff --git a/python/flink_agents/api/memory_object.py
b/python/flink_agents/api/memory_object.py
index 169d2e16..96452c7b 100644
--- a/python/flink_agents/api/memory_object.py
+++ b/python/flink_agents/api/memory_object.py
@@ -91,7 +91,7 @@ def _validate(value: Any, where: str) -> None:
class MemoryType(Enum):
"""Memory types based on MemoryObject."""
- SENSORY = ("sensory",)
+ SENSORY = "sensory"
SHORT_TERM = "short_term"
diff --git a/python/flink_agents/api/memory_reference.py
b/python/flink_agents/api/memory_reference.py
index 5c37793d..d48e614b 100644
--- a/python/flink_agents/api/memory_reference.py
+++ b/python/flink_agents/api/memory_reference.py
@@ -17,9 +17,9 @@
#################################################################################
from __future__ import annotations
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, ClassVar
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, model_serializer
from flink_agents.api.memory_object import MemoryType
@@ -30,6 +30,9 @@ if TYPE_CHECKING:
class MemoryRef(BaseModel):
"""Reference to a specific data item in the Short-Term Memory."""
+ TYPE_FIELD: ClassVar[str] = "@type"
+ TYPE_VALUE: ClassVar[str] = "memory_ref"
+
memory_type: MemoryType = MemoryType.SHORT_TERM
path: str
@@ -53,6 +56,11 @@ class MemoryRef(BaseModel):
"""
return MemoryRef(memory_type=memory_type, path=path)
+ @model_serializer(mode="wrap")
+ def _serialize_with_discriminator(self, handler: Any) -> dict[str, Any]:
+ """Serialize this reference with its language-neutral type
discriminator."""
+ return {self.TYPE_FIELD: self.TYPE_VALUE, **handler(self)}
+
def resolve(self, ctx: RunnerContext) -> Any:
"""Resolve the reference to get the actual data.
diff --git
a/python/flink_agents/api/tests/test_cross_language_event_snapshots.py
b/python/flink_agents/api/tests/test_cross_language_event_snapshots.py
index 545cbfdf..dc8c83e5 100644
--- a/python/flink_agents/api/tests/test_cross_language_event_snapshots.py
+++ b/python/flink_agents/api/tests/test_cross_language_event_snapshots.py
@@ -35,6 +35,8 @@ from flink_agents.api.events.event import Event, InputEvent,
OutputEvent
from flink_agents.api.events.memory_event import MemoryEvent,
ShortTermWriteEvent
from flink_agents.api.events.run_event import AgentRunBeginEvent
from flink_agents.api.events.tool_event import ToolRequestEvent,
ToolResponseEvent
+from flink_agents.api.memory_object import MemoryType
+from flink_agents.api.memory_reference import MemoryRef
from flink_agents.api.vector_stores.vector_store import Document
_REPO_ROOT = Path(__file__).resolve().parents[4]
@@ -560,6 +562,31 @@ def test_python_only_subclass_event_snapshot_is_stable()
-> None:
)
+def test_memory_ref_attachment_cross_language_snapshot_is_stable() -> None:
+ python_event = _force_id(
+ Event(
+ type="_memory_ref_attachment_event",
+ attributes={"value": "ping"},
+ ),
+ _FIXED_EVENT_ID,
+ )
+ python_event.set_attachment(
+ "payload", MemoryRef.create(MemoryType.SENSORY, "memory.path")
+ )
+ if _regenerate_enabled():
+ _write_python_snapshot("memory_ref_attachment_event.json",
python_event)
+ _assert_python_snapshot_stable("memory_ref_attachment_event.json",
python_event)
+
+ java_event = _read_java_snapshot("memory_ref_attachment_event.json")
+ assert java_event.id == _FIXED_EVENT_ID
+ assert java_event.type == "_memory_ref_attachment_event"
+ assert java_event.get_attr("value") == "ping"
+ reference = java_event.get_attachment("payload")
+ assert isinstance(reference, MemoryRef)
+ assert reference.memory_type == MemoryType.SENSORY
+ assert reference.path == "memory.path"
+
+
# ── Smoke ───────────────────────────────────────────────────────────────
diff --git a/python/flink_agents/api/tests/test_event.py
b/python/flink_agents/api/tests/test_event.py
index 4cd3cb56..ce945461 100644
--- a/python/flink_agents/api/tests/test_event.py
+++ b/python/flink_agents/api/tests/test_event.py
@@ -25,6 +25,8 @@ from pydantic_core import PydanticSerializationError
from pyflink.common import Row
from flink_agents.api.events.event import Event, InputEvent, OutputEvent
+from flink_agents.api.memory_object import MemoryType
+from flink_agents.api.memory_reference import MemoryRef
class _CustomEvent(Event):
@@ -237,14 +239,31 @@ def test_output_event_from_event() -> None:
assert reconstructed.type == "_output_event"
+def test_output_event_from_event_rejects_attachments() -> None:
+ """Test OutputEvent.from_event rejects an invalid attached source Event."""
+ base = Event(
+ type="_output_event",
+ attributes={"output": 99},
+ attachments={"payload": "attachment"},
+ )
+
+ with pytest.raises(ValueError, match="OutputEvent cannot carry
attachments"):
+ OutputEvent.from_event(base)
+
+
# ── Unified Event tests ──────────────────────────────────────────────────
def test_unified_event_creation() -> None:
"""Test creating a unified event with type and attributes."""
- event = Event(type="MyEvent", attributes={"field1": "test", "field2": 42})
+ event = Event(
+ type="MyEvent",
+ attributes={"field1": "test", "field2": 42},
+ attachments={"field3": "0105", "field4": 2004},
+ )
assert event.type == "MyEvent"
assert event.attributes == {"field1": "test", "field2": 42}
+ assert event.attachments == {"field3": "0105", "field4": 2004}
assert event.get_type() == "MyEvent"
@@ -262,6 +281,14 @@ def test_unified_event_get_attr_set_attr() -> None:
assert event.get_attr("missing") is None
+def test_unified_event_get_attachment_set_attachment() -> None:
+ """Test get_attachment and set_attachment convenience methods."""
+ event = Event(type="TestEvent")
+ event.set_attachment("key", "value")
+ assert event.get_attachment("key") == "value"
+ assert event.get_attachment("missing") is None
+
+
def test_unified_event_from_json() -> None:
"""Test deserializing a unified event from JSON."""
data = {"type": "MyEvent", "attributes": {"x": 1}}
@@ -309,6 +336,19 @@ def test_unified_event_serialization_roundtrip() -> None:
assert restored.attributes == {"a": 1, "b": "two"}
+def test_unified_event_serialization_roundtrip_with_memory_ref_attachment() ->
None:
+ """Test that tagged memory-reference attachments retain their concrete
type."""
+ reference = MemoryRef.create(MemoryType.SENSORY, "memory.path")
+ original = Event(type="RoundTrip", attachments={"payload": reference})
+
+ parsed = json.loads(original.model_dump_json())
+ restored = Event.model_validate(parsed)
+
+ assert parsed["attachments"]["payload"][MemoryRef.TYPE_FIELD] ==
MemoryRef.TYPE_VALUE
+ assert restored.get_attachment("payload") == reference
+ assert isinstance(restored.get_attachment("payload"), MemoryRef)
+
+
def test_event_lineage_json_roundtrip_uses_java_field_names() -> None:
"""Test framework-managed lineage has a stable cross-language JSON
shape."""
upstream_event_id = UUID("00000000-0000-0000-0000-000000000001")
diff --git
a/python/flink_agents/e2e_tests/e2e_tests_integration/event_attachments_test.py
b/python/flink_agents/e2e_tests/e2e_tests_integration/event_attachments_test.py
new file mode 100644
index 00000000..a1dc7041
--- /dev/null
+++
b/python/flink_agents/e2e_tests/e2e_tests_integration/event_attachments_test.py
@@ -0,0 +1,103 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#################################################################################
+import os
+import sys
+import sysconfig
+from pathlib import Path
+from typing import Any
+
+from pyflink.common import Configuration
+from pyflink.datastream import KeySelector, StreamExecutionEnvironment
+
+from flink_agents.api.agents.agent import Agent
+from flink_agents.api.decorators import action
+from flink_agents.api.events.event import Event, InputEvent, OutputEvent
+from flink_agents.api.execution_environment import AgentsExecutionEnvironment
+from flink_agents.api.runner_context import RunnerContext
+
+current_dir = Path(__file__).parent
+os.environ["PYTHONPATH"] = (
+ f"{current_dir.parent.parent.parent}:{sysconfig.get_paths()['purelib']}"
+)
+
+
+class _KeySelector(KeySelector):
+ def get_key(self, value: dict[str, Any]) -> str:
+ return str(value["key"])
+
+
+class EventAttachmentsAgent(Agent):
+ @action(InputEvent.EVENT_TYPE)
+ @staticmethod
+ def send_attachment(event: Event, ctx: RunnerContext) -> None:
+ value = InputEvent.from_event(event).input
+ ctx.send_event(
+ Event(
+ type="AttachmentStep",
+ attributes={"kind": "inline"},
+ attachments={
+ "payload": {
+ "value": value,
+ "items": [1, 2, 3],
+ }
+ },
+ )
+ )
+
+ @action("AttachmentStep")
+ @staticmethod
+ def receive_attachment(event: Event, ctx: RunnerContext) -> None:
+ attachment = event.get_attachment("payload")
+ ctx.send_event(
+ OutputEvent(
+ output={
+ "kind": event.get_attr("kind"),
+ "payload": {
+ "value": attachment["value"],
+ "items": attachment["items"],
+ },
+ }
+ )
+ )
+
+
+def test_python_event_attachments_roundtrip_on_flink() -> None:
+ config = Configuration()
+ config.set_string("python.pythonpath", os.environ["PYTHONPATH"])
+ env = StreamExecutionEnvironment.get_execution_environment(config)
+ env.set_python_executable(sys.executable)
+ env.set_parallelism(1)
+ input_stream = env.from_collection(
+ [{"key": "k1", "value": {"message": "hello"}}]
+ )
+ agents_env = AgentsExecutionEnvironment.get_execution_environment(env=env)
+ output = (
+ agents_env.from_datastream(input_stream, _KeySelector())
+ .apply(EventAttachmentsAgent())
+ .to_datastream()
+ )
+
+ assert list(output.execute_and_collect()) == [
+ {
+ "kind": "inline",
+ "payload": {
+ "value": {"key": "k1", "value": {"message": "hello"}},
+ "items": [1, 2, 3],
+ },
+ }
+ ]
diff --git a/python/flink_agents/runtime/flink_runner_context.py
b/python/flink_agents/runtime/flink_runner_context.py
index 4e9fefa8..8fab1e7f 100644
--- a/python/flink_agents/runtime/flink_runner_context.py
+++ b/python/flink_agents/runtime/flink_runner_context.py
@@ -55,6 +55,7 @@ from flink_agents.runtime.durable_execution import (
)
from flink_agents.runtime.flink_memory_object import FlinkMemoryObject
from flink_agents.runtime.flink_metric_group import FlinkMetricGroup
+from flink_agents.runtime.memory.event_attachment_utils import
store_event_attachments
from flink_agents.runtime.memory.internal_base_long_term_memory import (
InternalBaseLongTermMemory,
)
@@ -514,6 +515,7 @@ class FlinkRunnerContext(RunnerContext, ExecutionReporter):
event : Event
The event to be processed by the agent system.
"""
+ store_event_attachments(event, self)
event_json = event.model_dump_json()
try:
self._j_runner_context.sendEventJson(event_json)
diff --git a/python/flink_agents/runtime/memory/event_attachment_utils.py
b/python/flink_agents/runtime/memory/event_attachment_utils.py
new file mode 100644
index 00000000..c7b555c2
--- /dev/null
+++ b/python/flink_agents/runtime/memory/event_attachment_utils.py
@@ -0,0 +1,108 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#################################################################################
+from __future__ import annotations
+
+import hashlib
+from typing import TYPE_CHECKING, Any
+
+from flink_agents.api.events.event import OutputEvent
+from flink_agents.api.memory_object import MemoryType, validate_memory_value
+from flink_agents.api.memory_reference import MemoryRef
+
+if TYPE_CHECKING:
+ from uuid import UUID
+
+ from flink_agents.api.events.event import Event
+ from flink_agents.api.runner_context import RunnerContext
+
+_ATTACHMENT_ROOT = "__event_attachments__"
+
+
+class EventAttachmentError(RuntimeError):
+ """Raised when an Event attachment cannot be stored or loaded."""
+
+
+def _hash_attachment_key(key: str) -> str:
+ return hashlib.sha256(key.encode("UTF-8")).hexdigest()
+
+
+def build_attachment_path(event_id: UUID, key: str) -> str:
+ """Build the canonical SensoryMemory path for one attachment."""
+ return f"{_ATTACHMENT_ROOT}.{event_id}.{_hash_attachment_key(key)}"
+
+
+def _attachment_context(event: Event, key: str, path: str | None = None) ->
str:
+ suffix = f", path={path}" if path is not None else ""
+ return f"event_id={event.id}, event_type={event.type}, key={key}{suffix}"
+
+
+def store_event_attachments(event: Event, ctx: RunnerContext) -> None:
+ """Store concrete attachment values in SensoryMemory and replace them with
refs."""
+ if not event.attachments:
+ return
+
+ if event.type == OutputEvent.EVENT_TYPE:
+ keys = ", ".join(sorted(event.attachments))
+ msg = f"Output events cannot carry attachments:
{_attachment_context(event, keys)}"
+ raise EventAttachmentError(msg)
+
+ pending: list[tuple[str, str, Any]] = []
+ for key, value in event.attachments.items():
+ if isinstance(value, MemoryRef):
+ if value.memory_type != MemoryType.SENSORY:
+ msg = (
+ "Event attachments must use sensory memory references: "
+ f"{_attachment_context(event, key, value.path)}"
+ )
+ raise EventAttachmentError(msg)
+ continue
+ path = build_attachment_path(event.id, key)
+ try:
+ validate_memory_value(path, value)
+ except Exception as exc:
+ msg = f"Invalid event attachment value:
{_attachment_context(event, key, path)}"
+ raise EventAttachmentError(msg) from exc
+ pending.append((key, path, value))
+
+ for key, path, value in pending:
+ event.attachments[key] = ctx.sensory_memory.set(path, value)
+
+
+def load_event_attachments(event: Event, ctx: RunnerContext) -> None:
+ """Load sensory refs in place immediately before a Python Action runs."""
+ for key, value in list(event.attachments.items()):
+ if not isinstance(value, MemoryRef):
+ continue
+
+ try:
+ sensory_memory = ctx.sensory_memory
+ attachment = sensory_memory.get(value)
+ attachment_missing = attachment is None and not
sensory_memory.is_exist(
+ value.path
+ )
+ except Exception as exc:
+ msg = f"Failed to load event attachment:
{_attachment_context(event, key, value.path)}"
+ raise EventAttachmentError(msg) from exc
+
+ if attachment_missing:
+ msg = (
+ "Event attachment does not exist in sensory memory: "
+ f"{_attachment_context(event, key, value.path)}"
+ )
+ raise EventAttachmentError(msg)
+ event.attachments[key] = attachment
diff --git a/python/flink_agents/runtime/python_java_utils.py
b/python/flink_agents/runtime/python_java_utils.py
index d13db0b7..22f21353 100644
--- a/python/flink_agents/runtime/python_java_utils.py
+++ b/python/flink_agents/runtime/python_java_utils.py
@@ -42,6 +42,7 @@ from flink_agents.runtime.java.java_resource_wrapper import (
JavaResourceContextWrapper,
JavaTool,
)
+from flink_agents.runtime.memory.event_attachment_utils import
load_event_attachments
def convert_to_python_object(bytesObject: bytes) -> Any:
@@ -67,6 +68,12 @@ def convert_json_to_python_event(event_json: str) -> Event:
return Event.from_json(event_json)
+def load_event_attachments_for_action(event: Event, ctx: Any) -> Event:
+ """Load attachment values before invoking a Python action."""
+ load_event_attachments(event, ctx)
+ return event
+
+
def wrap_to_input_event(bytesObject: bytes) -> str:
"""Wrap data to python input event and serialize as JSON.
diff --git a/python/flink_agents/runtime/tests/test_event_attachment_utils.py
b/python/flink_agents/runtime/tests/test_event_attachment_utils.py
new file mode 100644
index 00000000..6b871f6a
--- /dev/null
+++ b/python/flink_agents/runtime/tests/test_event_attachment_utils.py
@@ -0,0 +1,178 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#################################################################################
+import uuid
+
+import pytest
+
+from flink_agents.api.events.event import Event, OutputEvent
+from flink_agents.api.memory_object import MemoryType
+from flink_agents.api.memory_reference import MemoryRef
+from flink_agents.runtime.memory.event_attachment_utils import (
+ EventAttachmentError,
+ build_attachment_path,
+ load_event_attachments,
+ store_event_attachments,
+)
+from flink_agents.runtime.tests.local_memory_object import LocalMemoryObject
+
+
+class MockRunnerContext:
+ def __init__(self, memory: LocalMemoryObject) -> None:
+ self._memory = memory
+
+ @property
+ def sensory_memory(self) -> LocalMemoryObject:
+ return self._memory
+
+
+def test_store_event_attachments() -> None:
+ sensory_memory = LocalMemoryObject(MemoryType.SENSORY, {})
+ ctx = MockRunnerContext(sensory_memory)
+ event_id = uuid.uuid4()
+ payload = {"value": "original"}
+ event = Event.model_construct(
+ id=event_id,
+ type="AttachmentStep",
+ attributes={},
+ attachments={"payload": payload},
+ )
+
+ store_event_attachments(event, ctx)
+
+ attachment = event.get_attachment("payload")
+ assert isinstance(attachment, MemoryRef)
+ assert attachment.path == build_attachment_path(event_id, "payload")
+ assert sensory_memory.get(attachment) == payload
+
+
+def test_store_offloads_non_utf8_bytes_from_regular_event() -> None:
+ """Test that raw bytes are offloaded before JSON serialization is
required."""
+ sensory_memory = LocalMemoryObject(MemoryType.SENSORY, {})
+ ctx = MockRunnerContext(sensory_memory)
+ payload = b"\xff\x00"
+ event = Event(type="AttachmentStep", attachments={"payload": payload})
+
+ store_event_attachments(event, ctx)
+
+ attachment = event.get_attachment("payload")
+ assert isinstance(attachment, MemoryRef)
+ assert sensory_memory.get(attachment) == payload
+ event.model_dump_json()
+
+
+def test_store_rejects_non_sensory_memory_references() -> None:
+ sensory_memory = LocalMemoryObject(MemoryType.SENSORY, {})
+ ctx = MockRunnerContext(sensory_memory)
+ event = Event(
+ type="AttachmentStep",
+ attachments={
+ "payload": MemoryRef.create(MemoryType.SHORT_TERM,
"attachment.path")
+ },
+ )
+
+ with pytest.raises(EventAttachmentError) as exc_info:
+ store_event_attachments(event, ctx)
+
+ assert str(exc_info.value).startswith(
+ "Event attachments must use sensory memory references:"
+ )
+
+
+def test_store_rejects_output_event_attachments_before_storing_them() -> None:
+ sensory_memory = LocalMemoryObject(MemoryType.SENSORY, {})
+ ctx = MockRunnerContext(sensory_memory)
+ event_id = uuid.uuid4()
+ attachments = {"zeta": {"value": 2}, "alpha": {"value": 1}}
+ event = Event.model_construct(
+ id=event_id,
+ type=OutputEvent.EVENT_TYPE,
+ attributes={"output": "result"},
+ attachments=dict(attachments),
+ )
+
+ with pytest.raises(EventAttachmentError) as exc_info:
+ store_event_attachments(event, ctx)
+
+ message = str(exc_info.value)
+ assert message.startswith("Output events cannot carry attachments:")
+
+
+def test_load_event_attachments() -> None:
+ sensory_memory = LocalMemoryObject(MemoryType.SENSORY, {})
+ ctx = MockRunnerContext(sensory_memory)
+ event_id = uuid.uuid4()
+ payload = {"value": "original"}
+ reference = sensory_memory.set(build_attachment_path(event_id, "payload"),
payload)
+ event = Event.model_construct(
+ id=event_id,
+ type="AttachmentStep",
+ attributes={},
+ attachments={"payload": reference},
+ )
+
+ load_event_attachments(event, ctx)
+
+ assert event.get_attachment("payload") == payload
+
+
+def test_load_preserves_none_event_attachment() -> None:
+ sensory_memory = LocalMemoryObject(MemoryType.SENSORY, {})
+ ctx = MockRunnerContext(sensory_memory)
+ event_id = uuid.uuid4()
+ reference = sensory_memory.set(build_attachment_path(event_id, "payload"),
None)
+ event = Event.model_construct(
+ id=event_id,
+ type="AttachmentStep",
+ attributes={},
+ attachments={"payload": reference},
+ )
+
+ load_event_attachments(event, ctx)
+
+ assert event.get_attachment("payload") is None
+
+
+def test_load_rejects_missing_event_attachment() -> None:
+ sensory_memory = LocalMemoryObject(MemoryType.SENSORY, {})
+ ctx = MockRunnerContext(sensory_memory)
+ event = Event(
+ type="AttachmentStep",
+ attachments={
+ "payload": MemoryRef.create(MemoryType.SENSORY,
"missing.attachment")
+ },
+ )
+
+ with pytest.raises(EventAttachmentError) as exc_info:
+ load_event_attachments(event, ctx)
+
+ assert str(exc_info.value).startswith(
+ "Event attachment does not exist in sensory memory:"
+ )
+
+
+def test_build_attachment_path() -> None:
+ event_id = uuid.UUID("00000000-0000-0000-0000-000000000001")
+
+ path = build_attachment_path(event_id, "payload")
+
+ assert (
+ path
+ == "__event_attachments__."
+ + str(event_id)
+ + ".239f59ed55e737c77147cf55ad0c1b030b6d7ee748a7426952f9b852d5a935e5"
+ )
diff --git a/python/flink_agents/runtime/tests/test_flink_runner_context.py
b/python/flink_agents/runtime/tests/test_flink_runner_context.py
new file mode 100644
index 00000000..f283bab0
--- /dev/null
+++ b/python/flink_agents/runtime/tests/test_flink_runner_context.py
@@ -0,0 +1,68 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#################################################################################
+from typing import Any
+
+from flink_agents.api.events.event import Event
+from flink_agents.api.memory_reference import MemoryRef
+from flink_agents.runtime.flink_runner_context import FlinkRunnerContext
+
+
+class _FakeJavaMemoryRef:
+ def __init__(self, path: str) -> None:
+ self._path = path
+
+ def getPath(self) -> str:
+ return self._path
+
+
+class _FakeJavaSensoryMemory:
+ def __init__(self) -> None:
+ self.values: dict[str, Any] = {}
+
+ def set(self, path: str, value: Any) -> _FakeJavaMemoryRef:
+ self.values[path] = value
+ return _FakeJavaMemoryRef(path)
+
+
+class _FakeJavaRunnerContext:
+ def __init__(self) -> None:
+ self.sensory_memory = _FakeJavaSensoryMemory()
+ self.sent_event_json: str | None = None
+
+ def getSensoryMemory(self) -> _FakeJavaSensoryMemory:
+ return self.sensory_memory
+
+ def sendEventJson(self, event_json: str) -> None:
+ self.sent_event_json = event_json
+
+
+def test_send_event_offloads_attachments_before_forwarding() -> None:
+ java_context = _FakeJavaRunnerContext()
+ ctx = FlinkRunnerContext.__new__(FlinkRunnerContext)
+ ctx._j_runner_context = java_context
+ payload = {"value": "original"}
+ event = Event(type="AttachmentStep", attachments={"payload": payload})
+
+ ctx.send_event(event)
+
+ attachment = event.get_attachment("payload")
+ assert isinstance(attachment, MemoryRef)
+ assert java_context.sensory_memory.values[attachment.path] == payload
+ assert java_context.sent_event_json is not None
+ forwarded_event = Event.from_json(java_context.sent_event_json)
+ assert forwarded_event.get_attachment("payload") == attachment
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java
index 49d56e37..fe1c4ccf 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java
@@ -41,6 +41,7 @@ import
org.apache.flink.agents.runtime.actionstate.ActionState;
import org.apache.flink.agents.runtime.actionstate.CallResult;
import org.apache.flink.agents.runtime.lifecycle.ComponentExecutionListener;
import org.apache.flink.agents.runtime.memory.CachedMemoryStore;
+import org.apache.flink.agents.runtime.memory.EventAttachmentUtils;
import org.apache.flink.agents.runtime.memory.InteranlBaseLongTermMemory;
import org.apache.flink.agents.runtime.memory.MemoryEventBuilder;
import org.apache.flink.agents.runtime.memory.MemoryEventSettings;
@@ -215,6 +216,11 @@ public class RunnerContextImpl implements RunnerContext,
ExecutionReporter {
@Override
public void sendEvent(Event event) {
mailboxThreadChecker.run();
+ try {
+ EventAttachmentUtils.storeEventAttachments(event, this);
+ } catch (Exception e) {
+ throw new IllegalArgumentException("Failed to store event
attachments.", e);
+ }
try {
JsonUtils.checkSerializable(event);
} catch (JsonProcessingException e) {
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/memory/EventAttachmentUtils.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/memory/EventAttachmentUtils.java
new file mode 100644
index 00000000..6928d7b7
--- /dev/null
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/memory/EventAttachmentUtils.java
@@ -0,0 +1,146 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.flink.agents.runtime.memory;
+
+import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.api.OutputEvent;
+import org.apache.flink.agents.api.context.MemoryObject;
+import org.apache.flink.agents.api.context.MemoryRef;
+import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
+
+import java.nio.charset.StandardCharsets;
+import java.security.MessageDigest;
+import java.security.NoSuchAlgorithmException;
+import java.util.Map;
+import java.util.Objects;
+import java.util.UUID;
+import java.util.stream.Collectors;
+
+/** Stores event attachments in sensory memory while events cross action
boundaries. */
+public final class EventAttachmentUtils {
+
+ private static final String ATTACHMENT_ROOT = "__event_attachments__";
+
+ private EventAttachmentUtils() {}
+
+ /** Stores concrete attachment values and replaces them with
sensory-memory references. */
+ public static void storeEventAttachments(Event event, RunnerContext
context) throws Exception {
+ if (event.getAttachments().isEmpty()) {
+ return;
+ }
+
+ if (OutputEvent.EVENT_TYPE.equals(event.getType())) {
+ String keys =
+ event.getAttachments().keySet().stream()
+ .sorted()
+ .collect(Collectors.joining(", "));
+ throw new IllegalArgumentException(
+ "Output events cannot carry attachments: event_id="
+ + event.getId()
+ + ", event_type="
+ + event.getType()
+ + ", key="
+ + keys);
+ }
+
+ for (Map.Entry<String, Object> entry :
event.getAttachments().entrySet()) {
+ String key = entry.getKey();
+ Object value = entry.getValue();
+ if (value instanceof MemoryRef) {
+ MemoryRef reference = (MemoryRef) value;
+ if
(!MemoryObject.MemoryType.SENSORY.equals(reference.getType())) {
+ throw new IllegalArgumentException(
+ "Event attachments must use sensory memory
references: event_id="
+ + event.getId()
+ + ", event_type="
+ + event.getType()
+ + ", key="
+ + key
+ + ", memory_type="
+ + reference.getType());
+ }
+ continue;
+ }
+
+ MemoryRef reference =
+
context.getSensoryMemory().set(buildAttachmentPath(event.getId(), key), value);
+
+ event.getAttachments().put(key, reference);
+ }
+ }
+
+ /** Returns an Event with sensory-memory references resolved on a copy
when necessary. */
+ public static Event loadEventAttachments(
+ Event event, RunnerContext context, TypeSerializer<Event>
eventSerializer)
+ throws Exception {
+ boolean requiresResolution =
+
event.getAttachments().values().stream().anyMatch(MemoryRef.class::isInstance);
+ if (!requiresResolution) {
+ return event;
+ }
+
+ Event actionEvent =
+ Objects.requireNonNull(
+ eventSerializer,
+ "Event serializer is required to resolve
attachment references.")
+ .copy(event);
+ for (Map.Entry<String, Object> entry :
actionEvent.getAttachments().entrySet()) {
+ Object value = entry.getValue();
+ if (!(value instanceof MemoryRef)) {
+ continue;
+ }
+ MemoryRef reference = (MemoryRef) value;
+
+ MemoryObject attachment =
context.getSensoryMemory().get(reference);
+ if (attachment == null) {
+ throw new IllegalStateException(
+ "Event attachment does not exist in sensory memory: "
+ + reference.getPath());
+ }
+ actionEvent.getAttachments().put(entry.getKey(),
attachment.getValue());
+ }
+ return actionEvent;
+ }
+
+ /** Builds the sensory-memory path for one event attachment. */
+ public static String buildAttachmentPath(UUID eventId, String key) {
+ if (eventId == null) {
+ throw new IllegalArgumentException("Event attachment requires a
non-null event id.");
+ }
+ if (key == null) {
+ throw new IllegalArgumentException("Event attachment key must not
be null.");
+ }
+ return ATTACHMENT_ROOT + "." + eventId + "." + hashAttachmentKey(key);
+ }
+
+ private static String hashAttachmentKey(String key) {
+ try {
+ byte[] digest =
+ MessageDigest.getInstance("SHA-256")
+ .digest(key.getBytes(StandardCharsets.UTF_8));
+ StringBuilder sb = new StringBuilder(digest.length * 2);
+ for (byte value : digest) {
+ sb.append(String.format("%02x", value));
+ }
+ return sb.toString();
+ } catch (NoSuchAlgorithmException e) {
+ throw new IllegalStateException("SHA-256 is not available.", e);
+ }
+ }
+}
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
index 94c050a9..4d8d5d84 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
@@ -50,6 +50,7 @@ import
org.apache.flink.agents.runtime.trace.ExecutionEventLogger;
import org.apache.flink.agents.runtime.utils.EventUtil;
import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.api.common.operators.MailboxExecutor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.runtime.state.KeyGroupRange;
import org.apache.flink.runtime.state.StateInitializationContext;
@@ -137,6 +138,8 @@ public class ActionExecutionOperator<IN, OUT> extends
AbstractStreamOperator<OUT
private transient OperatorStateManager stateManager;
+ private transient TypeSerializer<Event> eventSerializer;
+
// Each job can only have one identifier and this identifier must be
consistent across restarts.
// We cannot use job id as the identifier here because user may change job
id by
// creating a savepoint, stop the job and then resume from savepoint.
@@ -189,6 +192,10 @@ public class ActionExecutionOperator<IN, OUT> extends
AbstractStreamOperator<OUT
public void open() throws Exception {
super.open();
+ eventSerializer =
+ TypeInformation.of(Event.class)
+
.createSerializer(getExecutionConfig().getSerializerConfig());
+
stateManager.initializeKeyedStates(getRuntimeContext(),
agentPlan.getConfig());
stateManager.initializeOperatorStates(getOperatorStateBackend());
@@ -510,6 +517,9 @@ public class ActionExecutionOperator<IN, OUT> extends
AbstractStreamOperator<OUT
ActionTask.ActionTaskResult actionTaskResult;
try {
+ if (actionTask instanceof JavaActionTask) {
+ ((JavaActionTask)
actionTask).setEventSerializer(eventSerializer);
+ }
actionTaskResult =
actionTask.invoke(
getRuntimeContext().getUserCodeClassLoader(),
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/JavaActionTask.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/JavaActionTask.java
index 702d82a8..f6b509ad 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/JavaActionTask.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/JavaActionTask.java
@@ -22,7 +22,9 @@ import
org.apache.flink.agents.api.trace.ExecutionTraceContext;
import org.apache.flink.agents.plan.JavaFunction;
import org.apache.flink.agents.plan.actions.Action;
import org.apache.flink.agents.runtime.context.JavaRunnerContextImpl;
+import org.apache.flink.agents.runtime.memory.EventAttachmentUtils;
import org.apache.flink.agents.runtime.python.utils.PythonActionExecutor;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
import java.util.Collections;
@@ -39,7 +41,10 @@ import static org.apache.flink.util.Preconditions.checkState;
*/
public class JavaActionTask extends ActionTask {
- private boolean executionStarted = false;
+ /** Event passed to the action, copied when attachment references need
resolution. */
+ private transient Event invocationEvent;
+
+ private transient TypeSerializer<Event> eventSerializer;
public JavaActionTask(Object key, Event event, Action action, long
sequenceNumber) {
super(key, event, action, sequenceNumber);
@@ -56,6 +61,10 @@ public class JavaActionTask extends ActionTask {
checkState(action.getExec() instanceof JavaFunction);
}
+ void setEventSerializer(TypeSerializer<Event> eventSerializer) {
+ this.eventSerializer = eventSerializer;
+ }
+
@Override
public ActionTaskResult invoke(ClassLoader userCodeClassLoader,
PythonActionExecutor executor)
throws Exception {
@@ -65,12 +74,15 @@ public class JavaActionTask extends ActionTask {
event,
key);
- if (!executionStarted) {
+ if (invocationEvent == null) {
runnerContext.checkNoPendingEvents();
- executionStarted = true;
+ invocationEvent =
+ EventAttachmentUtils.loadEventAttachments(
+ event, runnerContext, eventSerializer);
}
JavaRunnerContextImpl javaRunnerContext = (JavaRunnerContextImpl)
runnerContext;
+ Event actionEvent = invocationEvent;
ClassLoader cl = Thread.currentThread().getContextClassLoader();
boolean finished;
@@ -83,7 +95,7 @@ public class JavaActionTask extends ActionTask {
javaRunnerContext.getContinuationContext(),
() -> {
try {
- action.getExec().call(event,
runnerContext);
+ action.getExec().call(actionEvent,
runnerContext);
} catch (Exception e) {
throw new RuntimeException(e);
}
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java
index 3f97b961..0828f2db 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java
@@ -88,6 +88,8 @@ public class PythonActionExecutor implements AutoCloseable {
// =========== PYTHON AND JAVA OBJECT CONVERT ===========
private static final String CONVERT_JSON_TO_PYTHON_EVENT =
"python_java_utils.convert_json_to_python_event";
+ private static final String LOAD_EVENT_ATTACHMENTS_FOR_ACTION =
+ "python_java_utils.load_event_attachments_for_action";
private static final String CONVERT_TO_PYTHON_KEY_TEXT =
"python_java_utils.convert_to_python_key_text";
private static final String PICKLED_KEY_SERIALIZATION = "pickled";
@@ -244,6 +246,10 @@ public class PythonActionExecutor implements AutoCloseable
{
interpreter.invoke(
CONVERT_JSON_TO_PYTHON_EVENT,
eventJson));
+ interpreter.invoke(
+ LOAD_EVENT_ATTACHMENTS_FOR_ACTION,
+ pythonEventObject,
+ pythonRunnerContext);
PyObject calledResult =
scope.own(
(PyObject)
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java
index c1f23504..9afebd73 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java
@@ -25,6 +25,8 @@ import org.apache.flink.agents.api.InputEvent;
import org.apache.flink.agents.api.OutputEvent;
import org.apache.flink.agents.api.chat.messages.ChatMessage;
import org.apache.flink.agents.api.chat.messages.MessageRole;
+import org.apache.flink.agents.api.context.MemoryObject;
+import org.apache.flink.agents.api.context.MemoryRef;
import org.apache.flink.agents.api.context.MemoryUpdate;
import org.apache.flink.agents.api.event.AgentRunBeginEvent;
import org.apache.flink.agents.api.event.ChatRequestEvent;
@@ -58,6 +60,8 @@ public class ActionStateSerdeTest {
// Create test data
InputEvent inputEvent = new InputEvent("test input");
inputEvent.setAttr("testAttr", "testValue");
+ MemoryRef reference =
MemoryRef.create(MemoryObject.MemoryType.SENSORY, "attachment.path");
+ inputEvent.setAttachment("payload", reference);
OutputEvent outputEvent = new OutputEvent("test output");
outputEvent.setAttr("outputAttr", 123);
@@ -88,6 +92,10 @@ public class ActionStateSerdeTest {
InputEvent deserializedInputEvent = (InputEvent)
deserializedState.getTaskEvent();
assertEquals("test input", deserializedInputEvent.getInput());
assertEquals("testValue", deserializedInputEvent.getAttr("testAttr"));
+ Object restoredAttachment =
deserializedInputEvent.getAttachment("payload");
+ assertInstanceOf(MemoryRef.class, restoredAttachment);
+ assertEquals(reference, restoredAttachment);
+ assertEquals(MemoryObject.MemoryType.SENSORY, ((MemoryRef)
restoredAttachment).getType());
// Verify memoryUpdates
assertEquals(1, deserializedState.getSensoryMemoryUpdates().size());
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/memory/EventAttachmentUtilsTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/EventAttachmentUtilsTest.java
new file mode 100644
index 00000000..662b6061
--- /dev/null
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/EventAttachmentUtilsTest.java
@@ -0,0 +1,310 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.flink.agents.runtime.memory;
+
+import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.api.OutputEvent;
+import org.apache.flink.agents.api.configuration.ReadableConfiguration;
+import org.apache.flink.agents.api.context.DurableCallable;
+import org.apache.flink.agents.api.context.MemoryObject;
+import org.apache.flink.agents.api.context.MemoryRef;
+import org.apache.flink.agents.api.context.Outcome;
+import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.api.memory.BaseLongTermMemory;
+import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup;
+import org.apache.flink.agents.api.resource.Resource;
+import org.apache.flink.agents.api.resource.ResourceType;
+import org.apache.flink.api.common.ExecutionConfig;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotSame;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+class EventAttachmentUtilsTest {
+
+ private MemoryObject sensoryMemory;
+ private RunnerContext context;
+ private TypeSerializer<Event> eventSerializer;
+
+ @BeforeEach
+ void setUp() throws Exception {
+ sensoryMemory =
+ new MemoryObjectImpl(
+ MemoryObject.MemoryType.SENSORY,
+ new CachedMemoryStore(new ForTestMemoryMapState<>()),
+ MemoryObjectImpl.ROOT_KEY,
+ new LinkedList<>());
+ context = new MockRunnerContext(sensoryMemory);
+ eventSerializer =
+ TypeInformation.of(Event.class)
+ .createSerializer(new
ExecutionConfig().getSerializerConfig());
+ }
+
+ @Test
+ void storesEventAttachments() throws Exception {
+ UUID eventId = UUID.randomUUID();
+ Map<String, Object> payload = Map.of("value", "original");
+ Event event =
+ new Event(
+ eventId,
+ "AttachmentStep",
+ Map.of(),
+ new HashMap<>(Map.of("payload", payload)));
+
+ EventAttachmentUtils.storeEventAttachments(event, context);
+
+ Object attachment = event.getAttachment("payload");
+ assertTrue(attachment instanceof MemoryRef);
+ MemoryRef reference = (MemoryRef) attachment;
+ assertEquals(
+ EventAttachmentUtils.buildAttachmentPath(eventId, "payload"),
reference.getPath());
+ assertEquals(payload, sensoryMemory.get(reference).getValue());
+ }
+
+ @Test
+ void storesAttachmentsFromImmutableMap() throws Exception {
+ UUID eventId = UUID.randomUUID();
+ Map<String, Object> payload = Map.of("value", "original");
+ Map<String, Object> attachments = Map.of("payload", payload);
+ Event event = new Event(eventId, "AttachmentStep", Map.of(),
attachments);
+
+ EventAttachmentUtils.storeEventAttachments(event, context);
+
+ assertTrue(event.getAttachment("payload") instanceof MemoryRef);
+ assertEquals(payload, attachments.get("payload"));
+ }
+
+ @Test
+ void rejectsNonSensoryMemoryReferencesWhenStoringAttachments() {
+ Event event =
+ new Event(
+ UUID.randomUUID(),
+ "AttachmentStep",
+ Map.of(),
+ new HashMap<>(
+ Map.of(
+ "payload",
+ MemoryRef.create(
+
MemoryObject.MemoryType.SHORT_TERM,
+ "attachment.path"))));
+
+ IllegalArgumentException error =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
EventAttachmentUtils.storeEventAttachments(event, context));
+
+ assertTrue(
+ error.getMessage()
+ .startsWith("Event attachments must use sensory memory
references:"));
+ }
+
+ @Test
+ void rejectsOutputEventAttachmentsBeforeStoringThem() throws Exception {
+ UUID eventId = UUID.randomUUID();
+ Map<String, Object> attachments =
+ Map.of("zeta", Map.of("value", 2), "alpha", Map.of("value",
1));
+ Event event =
+ new Event(
+ eventId,
+ OutputEvent.EVENT_TYPE,
+ Map.of("output", "result"),
+ new HashMap<>(attachments));
+
+ IllegalArgumentException error =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
EventAttachmentUtils.storeEventAttachments(event, context));
+
+ assertTrue(error.getMessage().startsWith("Output events cannot carry
attachments:"));
+ }
+
+ @Test
+ void loadsEventAttachments() throws Exception {
+ UUID eventId = UUID.randomUUID();
+ Map<String, Object> payload = Map.of("value", "original");
+ MemoryRef reference =
+ sensoryMemory.set(
+ EventAttachmentUtils.buildAttachmentPath(eventId,
"payload"), payload);
+ Event event =
+ new Event(
+ eventId,
+ "AttachmentStep",
+ Map.of(),
+ new HashMap<>(Map.of("payload", reference)));
+
+ Event actionEvent =
+ EventAttachmentUtils.loadEventAttachments(event, context,
eventSerializer);
+
+ assertNotSame(event, actionEvent);
+ assertNotSame(event.getAttachments(), actionEvent.getAttachments());
+ assertEquals(payload, actionEvent.getAttachment("payload"));
+ assertSame(reference, event.getAttachment("payload"));
+ }
+
+ @Test
+ void rejectsMissingEventAttachment() {
+ Event event = new Event("AttachmentStep");
+ event.setAttachment(
+ "payload", MemoryRef.create(MemoryObject.MemoryType.SENSORY,
"missing.attachment"));
+
+ IllegalStateException error =
+ assertThrows(
+ IllegalStateException.class,
+ () ->
+ EventAttachmentUtils.loadEventAttachments(
+ event, context, eventSerializer));
+
+ assertTrue(
+ error.getMessage()
+ .startsWith("Event attachment does not exist in
sensory memory:"));
+ }
+
+ @Test
+ void loadsNullEventAttachment() throws Exception {
+ UUID eventId = UUID.randomUUID();
+ MemoryRef reference =
+ sensoryMemory.set(
+ EventAttachmentUtils.buildAttachmentPath(eventId,
"payload"), null);
+ Map<String, Object> attachments = new HashMap<>();
+ attachments.put("payload", reference);
+ Event event = new Event(eventId, "AttachmentStep", Map.of(),
attachments);
+
+ Event actionEvent =
+ EventAttachmentUtils.loadEventAttachments(event, context,
eventSerializer);
+
+ assertNull(actionEvent.getAttachment("payload"));
+ }
+
+ @Test
+ void returnsOriginalEventWhenNoReferencesNeedResolution() throws Exception
{
+ Event event = new Event("AttachmentStep", Map.of("value", "original"));
+ event.setAttachment("payload", Map.of("value", "inline"));
+
+ Event actionEvent =
+ EventAttachmentUtils.loadEventAttachments(event, context,
eventSerializer);
+
+ assertSame(event, actionEvent);
+ }
+
+ @Test
+ void buildsAttachmentPath() {
+ UUID eventId = UUID.fromString("00000000-0000-0000-0000-000000000001");
+
+ String path = EventAttachmentUtils.buildAttachmentPath(eventId,
"payload");
+
+ assertEquals(
+ "__event_attachments__."
+ + eventId
+ +
".239f59ed55e737c77147cf55ad0c1b030b6d7ee748a7426952f9b852d5a935e5",
+ path);
+ }
+
+ /** Mock RunnerContext for testing resolve(). */
+ static class MockRunnerContext implements RunnerContext {
+ private final MemoryObject memoryObject;
+
+ MockRunnerContext(MemoryObject memoryObject) {
+ this.memoryObject = memoryObject;
+ }
+
+ @Override
+ public MemoryObject getShortTermMemory() {
+ return null;
+ }
+
+ @Override
+ public BaseLongTermMemory getLongTermMemory() throws Exception {
+ return null;
+ }
+
+ @Override
+ public MemoryObject getSensoryMemory() {
+ return memoryObject;
+ }
+
+ @Override
+ public void sendEvent(org.apache.flink.agents.api.Event event) {}
+
+ @Override
+ public FlinkAgentsMetricGroup getAgentMetricGroup() {
+ return null;
+ }
+
+ @Override
+ public FlinkAgentsMetricGroup getActionMetricGroup() {
+ return null;
+ }
+
+ @Override
+ public Resource getResource(String name, ResourceType type) throws
Exception {
+ return null;
+ }
+
+ @Override
+ public ReadableConfiguration getConfig() {
+ return null;
+ }
+
+ @Override
+ public Map<String, Object> getActionConfig() {
+ return Map.of();
+ }
+
+ @Override
+ public Object getActionConfigValue(String key) {
+ return null;
+ }
+
+ @Override
+ public <T> T durableExecute(DurableCallable<T> callable) throws
Exception {
+ return callable.call();
+ }
+
+ @Override
+ public <T> T durableExecuteAsync(DurableCallable<T> callable) throws
Exception {
+ return callable.call();
+ }
+
+ @Override
+ public <T> List<Outcome<T>>
durableExecuteAllAsync(List<DurableCallable<T>> callables)
+ throws Exception {
+ List<Outcome<T>> outcomes = new ArrayList<>(callables.size());
+ for (DurableCallable<T> callable : callables) {
+ outcomes.add(Outcome.success(callable.call()));
+ }
+ return outcomes;
+ }
+
+ @Override
+ public void close() throws Exception {}
+ }
+}
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 1eaba97c..a135af77 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
@@ -26,6 +26,7 @@ import org.apache.flink.agents.api.OutputEvent;
import org.apache.flink.agents.api.configuration.AgentConfigOptions;
import org.apache.flink.agents.api.context.DurableCallable;
import org.apache.flink.agents.api.context.MemoryObject;
+import org.apache.flink.agents.api.context.MemoryRef;
import org.apache.flink.agents.api.context.RunnerContext;
import org.apache.flink.agents.api.event.ShortTermWriteEvent;
import org.apache.flink.agents.api.event.ToolRequestEvent;
@@ -192,6 +193,41 @@ public class ActionExecutionOperatorTest {
}
}
+ @Test
+ void testJavaEventAttachmentsAreOffloadedAndResolvedBetweenActions()
throws Exception {
+ AgentPlan agentPlan = TestAgent.getEventAttachmentAgentPlan();
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory<>(agentPlan, true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>)
testHarness.getOperator();
+
+ List<Event> eventsAtSendBoundary = new ArrayList<>();
+ operator.getEventRouter()
+ .addEventListener(
+ (context, event) -> {
+ if
(TestAgent.ATTACHMENT_EVENT_TYPE.equals(event.getType())) {
+ eventsAtSendBoundary.add(event);
+ }
+ });
+
+ testHarness.processElement(new StreamRecord<>(1L));
+ operator.waitInFlightEventsFinished();
+
+ assertThat(eventsAtSendBoundary).hasSize(1);
+ Event runtimeEvent = eventsAtSendBoundary.get(0);
+ Object reference =
runtimeEvent.getAttachment(TestAgent.ATTACHMENT_KEY);
+ assertThat(reference).isInstanceOf(MemoryRef.class);
+ List<StreamRecord<Object>> recordOutput =
+ (List<StreamRecord<Object>>) testHarness.getRecordOutput();
+ assertThat(recordOutput).hasSize(1);
+
assertThat(recordOutput.get(0).getValue()).isEqualTo(Map.of("value", 1L));
+ }
+ }
+
@Test
void testSameKeyDataAreProcessedInOrder() throws Exception {
try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
testHarness =
@@ -2963,6 +2999,9 @@ public class ActionExecutionOperatorTest {
public static class TestAgent {
+ private static final String ATTACHMENT_EVENT_TYPE = "AttachmentEvent";
+ private static final String ATTACHMENT_KEY = "payload";
+
/** Counter to track how many times the durable supplier is executed.
*/
public static final java.util.concurrent.atomic.AtomicInteger
DURABLE_CALL_COUNTER =
new java.util.concurrent.atomic.AtomicInteger(0);
@@ -3008,6 +3047,19 @@ public class ActionExecutionOperatorTest {
context.sendEvent(new MiddleEvent(inputData + 1));
}
+ public static void sendEventAttachment(Event event, RunnerContext
context) {
+ Event attachmentEvent = new Event(ATTACHMENT_EVENT_TYPE);
+ attachmentEvent.setAttachment(
+ ATTACHMENT_KEY, Map.of("value",
InputEvent.fromEvent(event).getInput()));
+ context.sendEvent(attachmentEvent);
+ }
+
+ public static void receiveEventAttachment(Event event, RunnerContext
context) {
+ Object attachment = event.getAttachment(ATTACHMENT_KEY);
+ event.setAttachment(ATTACHMENT_KEY, "mutated-by-action");
+ context.sendEvent(new OutputEvent(attachment));
+ }
+
public static void action2(MiddleEvent event, RunnerContext context) {
ACTION2_CALL_COUNTER.incrementAndGet();
try {
@@ -3355,6 +3407,34 @@ public class ActionExecutionOperatorTest {
return getAgentPlanWithConfig(new AgentConfiguration(),
testMemoryAccessOutOfMailbox);
}
+ public static AgentPlan getEventAttachmentAgentPlan() {
+ try {
+ Action sendAction =
+ new Action(
+ "sendEventAttachment",
+ new JavaFunction(
+ TestAgent.class,
+ "sendEventAttachment",
+ new Class<?>[] {Event.class,
RunnerContext.class}),
+
Collections.singletonList(InputEvent.EVENT_TYPE));
+ Action receiveAction =
+ new Action(
+ "receiveEventAttachment",
+ new JavaFunction(
+ TestAgent.class,
+ "receiveEventAttachment",
+ new Class<?>[] {Event.class,
RunnerContext.class}),
+
Collections.singletonList(ATTACHMENT_EVENT_TYPE));
+ return new AgentPlan(
+ Map.of(
+ sendAction.getName(), sendAction,
+ receiveAction.getName(), receiveAction));
+ } catch (Exception e) {
+ ExceptionUtils.rethrow(e);
+ }
+ return null;
+ }
+
public static AgentPlan getAgentPlanWithConfig(AgentConfiguration
config) {
return getAgentPlanWithConfig(config, false);
}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/JavaActionTaskTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/JavaActionTaskTest.java
new file mode 100644
index 00000000..466f90a4
--- /dev/null
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/JavaActionTaskTest.java
@@ -0,0 +1,188 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.flink.agents.runtime.operator;
+
+import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.api.InputEvent;
+import org.apache.flink.agents.api.context.DurableCallable;
+import org.apache.flink.agents.api.context.MemoryRef;
+import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.plan.AgentPlan;
+import org.apache.flink.agents.plan.JavaFunction;
+import org.apache.flink.agents.plan.actions.Action;
+import org.apache.flink.agents.runtime.actionstate.ActionState;
+import org.apache.flink.agents.runtime.actionstate.InMemoryActionStateStore;
+import org.apache.flink.agents.runtime.async.ContinuationActionExecutor;
+import org.apache.flink.agents.runtime.async.ContinuationContext;
+import org.apache.flink.agents.runtime.context.JavaRunnerContextImpl;
+import org.apache.flink.agents.runtime.context.RunnerContextImpl;
+import org.apache.flink.agents.runtime.memory.CachedMemoryStore;
+import org.apache.flink.agents.runtime.memory.ForTestMemoryMapState;
+import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl;
+import org.apache.flink.api.common.ExecutionConfig;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.apache.flink.runtime.metrics.groups.UnregisteredMetricGroups;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.Callable;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests for {@link JavaActionTask}. */
+public class JavaActionTaskTest {
+
+ private static final String ACTION_NAME = "capture-attachment";
+ private static final String KEY = "key";
+ private static final long SEQUENCE_NUMBER = 0L;
+
+ private static Event invokedEvent;
+ private static Object invokedAttachment;
+
+ private ContinuationActionExecutor continuationExecutor;
+ private DurableExecutionManager durableExecutionManager;
+ private JavaRunnerContextImpl runnerContext;
+ private Action action;
+ private TypeSerializer<Event> eventSerializer;
+
+ @BeforeEach
+ void setUp() throws Exception {
+ invokedEvent = null;
+ invokedAttachment = null;
+ continuationExecutor = new ContinuationActionExecutor(1);
+ durableExecutionManager = new DurableExecutionManager(new
InMemoryActionStateStore(false));
+ runnerContext = createContext();
+ action = createAction();
+ eventSerializer =
+ TypeInformation.of(Event.class)
+ .createSerializer(new
ExecutionConfig().getSerializerConfig());
+ }
+
+ @AfterEach
+ void tearDown() throws Exception {
+ try {
+ continuationExecutor.close();
+ } finally {
+ durableExecutionManager.close();
+ }
+ }
+
+ public static void captureAttachment(InputEvent event, RunnerContext
context) throws Exception {
+ invokedEvent = event;
+ invokedAttachment = event.getAttachment("payload");
+ event.setAttachment("payload", "mutated-by-action");
+ context.durableExecute(
+ new TestDurableCallable<>(
+ "persist-attachment-state", String.class, () ->
"persisted"));
+ }
+
+ @Test
+ void resolvesAttachmentsWithoutMutatingDurableEvent() throws Exception {
+ Map<String, Object> payload = Map.of("value", "original");
+ MemoryRef reference =
runnerContext.getSensoryMemory().set("attachment.payload", payload);
+ InputEvent runtimeEvent = new InputEvent(1L);
+ runtimeEvent.setAttachment("payload", reference);
+ JavaActionTask task = new JavaActionTask(KEY, runtimeEvent, action,
1L);
+ task.setRunnerContext(runnerContext);
+ task.setEventSerializer(eventSerializer);
+ durableExecutionManager.maybeInitActionState(KEY, SEQUENCE_NUMBER,
action, runtimeEvent);
+ ActionState actionState = getPersistedState(runtimeEvent);
+ durableExecutionManager.setupDurableExecutionContext(task,
actionState, SEQUENCE_NUMBER);
+
+ task.invoke(getClass().getClassLoader(), null);
+
+ ActionState persistedState = getPersistedState(runtimeEvent);
+
assertThat(invokedEvent).isInstanceOf(InputEvent.class).isNotSameAs(runtimeEvent);
+ assertThat(((InputEvent) invokedEvent).getInput()).isEqualTo(1L);
+ assertThat(invokedAttachment).isSameAs(payload);
+ assertThat(persistedState.getCallResultCount()).isOne();
+
assertThat(persistedState.getTaskEvent().getAttachment("payload")).isSameAs(reference);
+ }
+
+ private JavaRunnerContextImpl createContext() throws Exception {
+ JavaRunnerContextImpl context =
+ new JavaRunnerContextImpl(
+ new FlinkAgentsMetricGroupImpl(
+
UnregisteredMetricGroups.createUnregisteredOperatorMetricGroup()),
+ () -> {},
+ new AgentPlan(new HashMap<>(), new HashMap<>()),
+ null,
+ "test-job",
+ continuationExecutor);
+ context.setContinuationContext(new ContinuationContext());
+ context.switchActionContext(
+ ACTION_NAME,
+ new RunnerContextImpl.MemoryContext(
+ new CachedMemoryStore(new ForTestMemoryMapState<>()),
+ new CachedMemoryStore(new ForTestMemoryMapState<>())),
+ new ArrayList<>(),
+ KEY,
+ "observation",
+ false,
+ null);
+ return context;
+ }
+
+ private Action createAction() throws Exception {
+ return new Action(
+ ACTION_NAME,
+ new JavaFunction(
+ JavaActionTaskTest.class,
+ "captureAttachment",
+ new Class<?>[] {InputEvent.class,
RunnerContext.class}),
+ List.of(InputEvent.EVENT_TYPE));
+ }
+
+ private ActionState getPersistedState(Event event) throws Exception {
+ return durableExecutionManager.maybeGetActionState(KEY,
SEQUENCE_NUMBER, action, event);
+ }
+
+ private static final class TestDurableCallable<T> implements
DurableCallable<T> {
+
+ private final String id;
+ private final Class<T> resultClass;
+ private final Callable<T> callSupplier;
+
+ private TestDurableCallable(String id, Class<T> resultClass,
Callable<T> callSupplier) {
+ this.id = id;
+ this.resultClass = resultClass;
+ this.callSupplier = callSupplier;
+ }
+
+ @Override
+ public String getId() {
+ return id;
+ }
+
+ @Override
+ public Class<T> getResultClass() {
+ return resultClass;
+ }
+
+ @Override
+ public T call() throws Exception {
+ return callSupplier.call();
+ }
+ }
+}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutorTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutorTest.java
index 63575170..a17e9dc7 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutorTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutorTest.java
@@ -127,6 +127,42 @@ class PythonActionExecutorTest {
}
}
+ @Test
+ void loadsEventAttachmentsBeforeCallingPythonFunction() throws Exception {
+ PythonInterpreter interpreter = mock(PythonInterpreter.class);
+ PythonActionExecutor executor =
+ newExecutor(interpreter, mock(PythonRunnerContextImpl.class));
+ PythonFunction function = new PythonFunction("test_module",
"test_action");
+ PyObject runnerContext = mock(PyObject.class);
+ PyObject pythonEvent = mock(PyObject.class);
+ setField(executor, "pythonRunnerContext", runnerContext);
+ when(interpreter.invoke(same(CONVERT_JSON_TO_PYTHON_EVENT),
anyString()))
+ .thenReturn(pythonEvent);
+ when(interpreter.invoke(
+ CALL_PYTHON_FUNCTION,
+ "test_module",
+ "test_action",
+ new Object[] {pythonEvent, runnerContext}))
+ .thenReturn(null);
+
+ executor.executePythonFunction(function, new Event("test-event"));
+
+ InOrder invocationOrder = inOrder(interpreter);
+ invocationOrder
+ .verify(interpreter)
+ .invoke(
+ "python_java_utils.load_event_attachments_for_action",
+ pythonEvent,
+ runnerContext);
+ invocationOrder
+ .verify(interpreter)
+ .invoke(
+ CALL_PYTHON_FUNCTION,
+ "test_module",
+ "test_action",
+ new Object[] {pythonEvent, runnerContext});
+ }
+
@Test
void resolvesPickledPythonKeyTextFromPyFlinkKeyRow() throws Exception {
PythonInterpreter interpreter = mock(PythonInterpreter.class);
