Skip to content

[Feature][api][runtime] Support auto resolve memory reference for passing data across actions. #950

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
JinkunLiu wants to merge 12 commits into
base: main
Choose a base branch
from
Open

[Feature][api][runtime] Support auto resolve memory reference for passing data across actions. #950

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
9bde274
[api][runtime] Add event attachments with automatic MemoryRef wrappin…
JinkunLiu Aug 2, 2026
1c9da03
[api][runtime] Test event attachments with automatic MemoryRef wrappi…
JinkunLiu Aug 2, 2026
600b901
[Fix] Preserve MemoryRef when restoring ActionState
JinkunLiu Aug 9, 2026
e4ffa8e
[Fix] Avoid JSON-serializing raw attachments before offload
JinkunLiu Aug 9, 2026
e2aa09e
[Fix] Make the Event own a mutable attachment map & attributes map
JinkunLiu Aug 9, 2026
e124231
[Fix] Resolve existing references according to their memory type
JinkunLiu Aug 16, 2026
ff464d9
Merge main into feature/476-auto-resolve-memory-ref
JinkunLiu Aug 16, 2026
3303fb9
[api][java][python] Reject attachments on OutputEvent
JinkunLiu Aug 22, 2026
ba79e7d
[runtime][java] Preserve attachment refs in durable action state
JinkunLiu Aug 22, 2026
28da154
[runtime][java] Use job serializer config for event copies
JinkunLiu Aug 23, 2026
6019ef9
[test] Test event attachment boundaries
JinkunLiu Aug 29, 2026
213ca02
Merge origin/main into feature/476-auto-resolve-memory-ref
JinkunLiu Aug 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
77 changes: 69 additions & 8 deletions api/src/main/java/org/apache/flink/agents/api/Event.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,13 @@
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;

Expand All @@ -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;

Expand All @@ -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. */
Expand All @@ -78,6 +89,9 @@ public 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()) {
Expand All @@ -86,14 +100,31 @@ public 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() {
Expand All @@ -110,6 +141,10 @@ public Map<String, Object> getAttributes() {
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)
Expand Down Expand Up @@ -150,6 +185,14 @@ public void setAttr(String name, Object value) {
attributes.put(name, value);
}

public Object getAttachment(String name) {
Comment thread
JinkunLiu marked this conversation as resolved.
return attachments.get(name);
}

public void setAttachment(String name, Object value) {
attachments.put(name, value);
}

@JsonIgnore
public boolean hasSourceTimestamp() {
return sourceTimestamp != null;
Expand All @@ -166,9 +209,9 @@ public void setSourceTimestamp(long timestamp) {
}

/**
* 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(
Expand All @@ -194,6 +237,8 @@ protected static <T extends Event> T reconstructFrom(
+ 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;
Expand All @@ -211,18 +256,34 @@ public static Event fromJson(String json) throws IOException {
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;
if (o == null || getClass() != o.getClass()) return false;
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);

@weiqingy weiqingy Aug 7, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Including attachments here, and in equals at :175, makes them part of event identity. The durable action-state key does not follow: ActionStateUtil.generateUUIDForEvent hashes event.getAttributes() only, so two events this line now distinguishes can still land on one ActionState.

A fan-out with durable execution on is where that shows: ctx.sendEvent(new Event("WorkItem", new HashMap<>(), Map.of("payload", item))) in a loop gives every sibling the same empty attributes, the same seqNum and the same action, so one state key covers all of them. Item 1 completes, item 2's lookup returns item 1's completed state, and ActionExecutionOperator.java:341 skips execution and replays item 1's output in its place.

Adding attachments to the key may just trade one problem for another, since a ref's path embeds the random event id the key deliberately avoids (buildAttachmentPath). I'm confident on the mechanism, less so on the odds, since it needs ACTION_STATE_STORE_BACKEND set plus siblings with equal attributes. Does that combination look reachable in practice?

@JinkunLiu JinkunLiu Aug 29, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for point that. I think the case you described can occur when durable execution is enabled.

However, I do not see a simple fix that fits within this PR. Maybe we need to redesign generateUUIDForEvent, perhaps by introducing a new ID to identify each event occurrence is better than deriving the Event identity from values carried.

Do you have any suggestions on this? I would appreciate your thoughts.

@weiqingy weiqingy Aug 30, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, an occurrence id is the right direction. But it cannot be Event.getId(), and it cannot be a fresh random one either.

EventRouter.wrapToInputEvent builds the root event with new InputEvent(input), which ends in UUID.randomUUID(). A record replayed from the source after a checkpoint gets a new id. So an id-based key would never find the old state, and every replayed record would re-run its whole chain. That is the case durable execution exists to avoid. ActionStateUtilTest.testGenerateKeyConsistency has pinned this since #138: two InputEvent("same-input") with different ids must produce the same key.

So the id has to do two things at once. It has to be unique per occurrence inside one (businessKey, seqNum, action), and it has to come out the same when the run is replayed from the source. A random id gives you the first and loses the second. Attribute content gives you the second and loses the first. Something derived from the event's position in the run gives you both.

There may already be a pattern to borrow. RunnerContextImpl.matchNextOrClearSubsequentCallResult identifies durable calls by their ordinal index in a persisted list, and uses functionId + argsDigest only to detect that replay diverged, clearing the later results when it does. Would that shape fit events too, with an ordinal in the key and the content hash demoted to a validation check?

If you go that way, an ordinal is the cheap option, but it is only replay-stable if dispatch order is deterministic across a restart, and I have not checked whether that holds once continuations and async calls interleave. A lineage path (parent occurrence, action name, index in that action's output list) sidesteps that question. upstreamEventId already carries the parent link, though it would need to point at the parent's derived identity rather than its random id.

One thing that may take some pressure off this PR: the collision does not need attachments. Two events with the same attributes sent to the same action already land on one ActionState on main today, since processEvent does not dedup and the key hashes attributes only.

Your PR makes it much easier to hit, because a fan-out that carries its payload in an attachment leaves the attributes identical by design. But the hole is already there. So this looks like a durable execution issue rather than an attachments one, and probably wants its own issue and its own tests rather than a fix squeezed in here.

Also my repro above does not compile, sorry. There is no Event(String, Map, Map) constructor. The simple version is Event e = new Event("WorkItem"); e.setAttachment("payload", item);

@wenjin272 wenjin272 Sep 2, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the detailed analysis. I created #1084 to track this separately. Since the collision already exists without attachments and likely requires a replay-stable occurrence identity, I agree that it should not be addressed in this PR.

}
}
7 changes: 6 additions & 1 deletion api/src/main/java/org/apache/flink/agents/api/OutputEvent.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,8 @@

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

Expand All @@ -50,8 +51,12 @@ public OutputEvent(
*
* @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);
}

Expand Down
70 changes: 70 additions & 0 deletions api/src/main/java/org/apache/flink/agents/api/context/MemoryRef.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
*/
@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;

Expand Down Expand Up @@ -67,6 +88,55 @@ public String getPath() {
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;
Expand Down
26 changes: 26 additions & 0 deletions api/src/test/java/org/apache/flink/agents/api/CrossLanguageEventSnapshotTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@
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;
Expand Down Expand Up @@ -52,6 +54,7 @@

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;
Expand All @@ -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;
Expand Down Expand Up @@ -579,6 +583,28 @@ void javaCanDeserializePythonOnlySubclassEventAsBaseEvent() throws Exception {
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
Expand Down
Loading
Loading