[Feature][api][runtime] Support auto resolve memory reference for passing data across actions. - #950
[Feature][api][runtime] Support auto resolve memory reference for passing data across actions.#950JinkunLiu wants to merge 12 commits into
Conversation
There was a problem hiding this comment.
Pull request overview
This PR introduces Event attachments as a first-class, cross-language concept (Java + Python) and wires the runtime so that attachment payloads can be automatically stored as MemoryRefs when events are sent, then resolved back to concrete values before actions run, enabling efficient data passing across action boundaries.
Changes:
- Add an
attachmentsmap to the unifiedEventmodel in both Java and Python, plus (de)serialization support forMemoryRefin attachments. - Add runtime utilities to store attachments into sensory memory (
store_event_attachments) and resolve them back (load_event_attachments), integrated into both Java and Python runner contexts and the Python action executor bridge. - Extend unit tests, cross-language snapshot tests, and add an end-to-end integration test covering attachment round-trips.
Reviewed changes
Copilot reviewed 49 out of 49 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| runtime/src/test/java/org/apache/flink/agents/runtime/memory/EventAttachmentUtilsTest.java | Adds Java unit tests for storing/loading attachments via sensory memory. |
| runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java | Resolves attachments on the Python side before invoking Python actions. |
| runtime/src/main/java/org/apache/flink/agents/runtime/operator/JavaActionTask.java | Resolves attachments before invoking Java actions. |
| runtime/src/main/java/org/apache/flink/agents/runtime/memory/EventAttachmentUtils.java | Implements Java-side attachment store/load utilities backed by sensory memory. |
| runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java | Automatically stores attachments as MemoryRefs when sending events. |
| python/flink_agents/runtime/tests/test_event_attachment_utils.py | Adds Python unit tests for storing/loading attachments via sensory memory. |
| python/flink_agents/runtime/python_java_utils.py | Adds a helper invoked from Java to load attachments before Python actions. |
| python/flink_agents/runtime/memory/event_attachment_utils.py | Implements Python-side attachment store/load utilities backed by sensory memory. |
| python/flink_agents/runtime/flink_runner_context.py | Automatically stores attachments when sending events from Python. |
| python/flink_agents/e2e_tests/e2e_tests_integration/event_attachments_test.py | Adds an e2e test validating attachment round-trip through Flink execution. |
| python/flink_agents/api/tests/test_event.py | Extends Event API tests to cover attachments getters/setters and JSON behavior. |
| python/flink_agents/api/tests/test_cross_language_event_snapshots.py | Validates cross-language snapshots include a MemoryRef attachment. |
| python/flink_agents/api/memory_object.py | Adjusts MemoryType enum representation (notably SENSORY). |
| python/flink_agents/api/events/tool_event.py | Preserves attachments when reconstructing typed tool events from base events. |
| python/flink_agents/api/events/event.py | Adds attachments field + accessors; parses MemoryRef objects from JSON attachments. |
| python/flink_agents/api/events/context_retrieval_event.py | Preserves attachments when reconstructing typed context retrieval events. |
| python/flink_agents/api/events/chat_event.py | Preserves attachments when reconstructing typed chat events. |
| e2e-test/cross-language-event-snapshots/python/tool_response_event.json | Updates Python snapshot with attachments section. |
| e2e-test/cross-language-event-snapshots/python/tool_request_event.json | Updates Python snapshot with attachments section. |
| e2e-test/cross-language-event-snapshots/python/python_only_subclass_event.json | Updates Python snapshot with attachments section. |
| e2e-test/cross-language-event-snapshots/python/output_event.json | Updates Python snapshot with attachments section. |
| e2e-test/cross-language-event-snapshots/python/input_event.json | Updates Python snapshot with attachments section. |
| e2e-test/cross-language-event-snapshots/python/generic_event_with_attrs.json | Updates Python snapshot with attachments section. |
| e2e-test/cross-language-event-snapshots/python/context_retrieval_response_event.json | Updates Python snapshot with attachments section. |
| e2e-test/cross-language-event-snapshots/python/context_retrieval_request_event.json | Updates Python snapshot with attachments section. |
| e2e-test/cross-language-event-snapshots/python/chat_response_event.json | Updates Python snapshot with attachments section. |
| e2e-test/cross-language-event-snapshots/python/chat_request_event.json | Updates Python snapshot with attachments section. |
| e2e-test/cross-language-event-snapshots/java/tool_response_event.json | Updates Java snapshot with attachments section. |
| e2e-test/cross-language-event-snapshots/java/tool_request_event.json | Updates Java snapshot with attachments section. |
| e2e-test/cross-language-event-snapshots/java/output_event.json | Updates Java snapshot with attachments section. |
| e2e-test/cross-language-event-snapshots/java/input_event.json | Updates Java snapshot with attachments section. |
| e2e-test/cross-language-event-snapshots/java/generic_event_with_attrs.json | Updates Java snapshot with attachments section. |
| e2e-test/cross-language-event-snapshots/java/context_retrieval_response_event.json | Updates Java snapshot with attachments section. |
| e2e-test/cross-language-event-snapshots/java/context_retrieval_request_event.json | Updates Java snapshot with attachments section. |
| e2e-test/cross-language-event-snapshots/java/chat_response_event.json | Updates Java snapshot with attachments section. |
| e2e-test/cross-language-event-snapshots/java/chat_request_event.json | Updates Java snapshot with attachments section. |
| api/src/test/java/org/apache/flink/agents/api/EventTest.java | Adds Java tests ensuring attachments (including MemoryRef) survive JSON round-trips. |
| api/src/test/java/org/apache/flink/agents/api/CrossLanguageEventSnapshotTest.java | Extends Java cross-language snapshot tests to include MemoryRef attachments. |
| api/src/test/java/org/apache/flink/agents/api/context/MemoryRefJsonTest.java | Adds Java tests for MemoryRef JSON (de)serialization. |
| api/src/main/java/org/apache/flink/agents/api/OutputEvent.java | Copies attachments when converting from base Event. |
| api/src/main/java/org/apache/flink/agents/api/InputEvent.java | Copies attachments when converting from base Event. |
| api/src/main/java/org/apache/flink/agents/api/event/ToolResponseEvent.java | Copies attachments when converting from base Event. |
| api/src/main/java/org/apache/flink/agents/api/event/ToolRequestEvent.java | Copies attachments when converting from base Event. |
| api/src/main/java/org/apache/flink/agents/api/event/ContextRetrievalResponseEvent.java | Copies attachments when converting from base Event. |
| api/src/main/java/org/apache/flink/agents/api/event/ContextRetrievalRequestEvent.java | Copies attachments when converting from base Event. |
| api/src/main/java/org/apache/flink/agents/api/event/ChatResponseEvent.java | Copies attachments when converting from base Event. |
| api/src/main/java/org/apache/flink/agents/api/event/ChatRequestEvent.java | Copies attachments when converting from base Event. |
| api/src/main/java/org/apache/flink/agents/api/Event.java | Adds attachments field + accessors; converts attachment maps into MemoryRef on JSON load. |
| api/src/main/java/org/apache/flink/agents/api/context/MemoryRef.java | Adds JSON (de)serialization for MemoryRef including memory_type and path. |
Suppressed comments (1)
api/src/main/java/org/apache/flink/agents/api/context/MemoryRef.java:143
MemoryRef.equals/hashCodeignoretype, so a SHORT_TERM ref and a SENSORY ref with the same path compare equal. Withmemory_typenow part of the JSON representation and cross-language contract, equality should include bothtypeandpathto avoid collisions and incorrect map/set behavior.
@Override
public boolean equals(Object o) {
if (this == o) return true;
if (o == null || getClass() != o.getClass()) return false;
MemoryRef memoryRef = (MemoryRef) o;
return path.equals(memoryRef.path);
}
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
…ng and resolution
1cb02a3 to
1c9da03
Compare
|
Hello @wenjin272, Could you take a look at this PR when you have time? Thank you very much |
|
Thanks for taking this on @JinkunLiu. I'll review it ASAP. |
wenjin272
left a comment
There was a problem hiding this comment.
Thank you, @JinkunLiu, for contributing this useful cross-action attachment feature and for covering both the Java and Python paths with unit, cross-language, and end-to-end tests. The overall direction looks promising. I left four comments around durable-recovery type fidelity, Python pre-offload serialization, attachment-map ownership, and type-aware MemoryRef resolution. Please take a look when you have time.
| public static Event fromJson(String json) throws IOException { | ||
| return MAPPER.readValue(json, Event.class); | ||
| Event event = MAPPER.readValue(json, Event.class); | ||
| for (Map.Entry<String, Object> entry : event.getAttachments().entrySet()) { |
There was a problem hiding this comment.
[P1] Preserve MemoryRef when restoring ActionState
Thanks for adding JSON support for MemoryRef. This conversion only runs when Event.fromJson() is called explicitly. Durable recovery instead deserializes the enclosing ActionState directly through ActionStateSerde, so attachment values declared as Object are restored as LinkedHashMap rather than MemoryRef. loadEventAttachments() then skips them, and the recovered action receives the reference-shaped map instead of the original payload.
A possible implementation sketch would be to bind a content deserializer directly to the attachments values:
@JsonCreator
public Event(
// ...
@JsonProperty("attachments")
@JsonDeserialize(contentUsing = AttachmentValueDeserializer.class)
Map<String, Object> attachments) {
// ...
}
static final class AttachmentValueDeserializer extends JsonDeserializer<Object> {
@Override
public Object deserialize(JsonParser p, DeserializationContext ctxt)
throws IOException {
JsonNode node = p.getCodec().readTree(p);
if (node.isObject()
&& "memory_ref".equals(node.path("@type").asText())) {
return p.getCodec().treeToValue(node, MemoryRef.class);
}
return p.getCodec().treeToValue(node, Object.class);
}
}The MemoryRef serializer would emit a language-neutral discriminator, for example {"@type": "memory_ref", "memory_type": "sensory", "path": "..."}. Merely adding the discriminator to the JSON is not sufficient: because the declared map value type is still Object, Jackson also needs this property-level deserialization hook (or an equivalent central hook) to select MemoryRef.Deserializer. This should use the same discriminator on the Python side, and an ActionStateSerde round-trip test should verify that the restored attachment is still a MemoryRef.
| id: UUID = Field(default=None) | ||
| type: str | ||
| attributes: Dict[str, Any] = Field(default_factory=dict) | ||
| attachments: Dict[str, Any] = Field(default_factory=dict) |
There was a problem hiding this comment.
[P2] Avoid JSON-serializing raw attachments before offload
Thanks for adding the Python attachment API. Because attachments is a regular Pydantic field, Event construction immediately includes the raw values in _generate_content_based_id(), and validate_and_set_id() serializes the entire Event again before store_event_attachments() can offload anything. This preserves the full JSON SerDe cost and also rejects valid memory payloads that are not JSON-serializable—for example, attachments={"payload": b"\xff\x00"} fails while constructing the Event. Could we keep raw attachments out of this generic JSON validation path and use an ID/offload strategy that does not require serializing the payload first? A regression test with non-UTF-8 bytes would help cover this.
There was a problem hiding this comment.
Yes, this is a clear bug, and I have fixed it on this commit.
While working on the fix, I also noticed that Python generated event IDs from an MD5 hash of serialized event content. As a result, two events with identical attributes could receive the same ID, which could cause attachment values to overwrite each other because their MemoryRef paths are derived from the event ID.
I have also updated Python to generate event IDs with uuid.uuid4(), matching Java’s UUID.randomUUID() behavior. If Java and Python intentionally used different ID-generation strategies for historical reasons, please let me know.
| this.id = id; | ||
| this.type = type; | ||
| this.attributes = attributes != null ? attributes : new HashMap<>(); | ||
| this.attachments = attachments != null ? attachments : new HashMap<>(); |
There was a problem hiding this comment.
[P2] Make the Event own a mutable attachment map
Thanks for the update. It looks like the immutable-map issue can still occur in the current version: the constructor stores the caller's map directly, while storeEventAttachments() and loadEventAttachments() later mutate it with put(). A common call such as new Event(..., Map.of("payload", value)) therefore writes the value to sensory memory and then fails with UnsupportedOperationException; a mutable shared map is modified behind the caller's back instead. Could we defensively copy it with new HashMap<>(attachments) and add a Map.of(...) regression test?
There was a problem hiding this comment.
Fix it in both attachments and attributes. e2aa09e0
| } | ||
| MemoryRef reference = (MemoryRef) value; | ||
|
|
||
| MemoryObject attachment = context.getSensoryMemory().get(reference); |
There was a problem hiding this comment.
[P2] Resolve existing references according to their memory type
Thanks for centralizing attachment resolution here. storeEventAttachments() accepts and skips every existing MemoryRef, but this load path always queries sensory memory. A SHORT_TERM reference is therefore accepted on send and then looked up in the wrong store; the Python equivalent also replaces a missing lookup with None without reporting it. Could we resolve through MemoryRef.resolve(context) (or dispatch on memory_type), or alternatively reject non-sensory references explicitly when storing? The Python path should also treat a missing resolved value as an error.
There was a problem hiding this comment.
Thanks for you point that,Fix it in e1242310
weiqingy
left a comment
There was a problem hiding this comment.
Thanks for taking this on. A few questions inline.
| @Override | ||
| public int hashCode() { | ||
| return Objects.hash(id, getType(), attributes); | ||
| return Objects.hash(id, getType(), attributes, attachments); |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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);
There was a problem hiding this comment.
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.
| */ | ||
| public static OutputEvent fromEvent(Event event) { | ||
| OutputEvent result = new OutputEvent(event.getId(), new HashMap<>(event.getAttributes())); | ||
| result.getAttachments().putAll(event.getAttachments()); |
There was a problem hiding this comment.
Both runtimes reject an OutputEvent carrying attachments before storing them (EventAttachmentUtils.java:46-58, event_attachment_utils.py:59-62), so there is no Java/Python gap here to close. What is left is internal: this copy, and the one at event.py:254, only ever build an object sendEvent refuses, and the output_event.json snapshots now pin that shape as a fixture.
What is the intended contract for attachments on OutputEvent? That answer decides whether the rejection moves or the copy does.
There was a problem hiding this comment.
Resolved, so I am closing this one out. fromEvent now rejects attachments, Python raises the same message, both output_event.json snapshots are clean, and there are tests on both sides.
| .to_datastream() | ||
| ) | ||
|
|
||
| assert list(output.execute_and_collect()) == [ |
There was a problem hiding this comment.
I went looking for what this assertion would catch. It checks the final OutputEvent payload, which is identical whether the attachment travelled offloaded or inline, so a regression that skipped the offload entirely would still pass. Nothing here observes a MemoryRef in flight.
None of the four new wiring points is covered by a test that carries an attachment either: RunnerContextImpl.java:154, JavaActionTask.java:60, PythonActionExecutor.java:139-141, flink_runner_context.py:301. Two even look removable without failing anything. Without flink_runner_context.py:301, Java's sendEvent offloads the dict instead and the output is unchanged. Without JavaActionTask.java:60, nothing fails either, since no test runs a Java action against an event carrying an attachment.
What would you want a test to pin down here? A MemoryRef at the send boundary and the resolved value at the receive boundary is the shape I'd reach for, and ActionExecutionOperatorTest looks like it could host it.
There was a problem hiding this comment.
Thanks for pointing this out. Explicitly testing the behavior at the action boundaries is necessary here.
Added new test at 28da154e
There was a problem hiding this comment.
Thanks, this covers it. I deleted the offload call and then moved it after serialization, and the Python test failed both times, so it pins the boundary and not just the payload.
PythonActionExecutor.java:139-140 is the one wiring point still reached only by the e2e test, which ut.sh skips by default. Is e2e the right level for the bridge, or would a small unit test there be worth having?
Generated-by: Codex (GPT-5)
ef57662 to
ad7180f
Compare
ad7180f to
213ca02
Compare
weiqingy
left a comment
There was a problem hiding this comment.
Thanks for working through these. The new boundary tests are the shape I was hoping for. I replied on the event id question, and left two small things inline.
| runtimeEvent); | ||
| assertThat(actionState).isNotNull(); | ||
| assertThat(actionState.getTaskEvent().getAttachment(TestAgent.ATTACHMENT_KEY)) | ||
| .isSameAs(reference); |
There was a problem hiding this comment.
I think this one can never fail. actionState.getTaskEvent() is the same object as runtimeEvent.
processEvent passes one event instance both to eventRouter.notifyEventProcessed (ActionExecutionOperator.java:298) and to createActionTask (:328). ActionTask.event is final (ActionTask.java:54), and the resolved copy goes to a different field, JavaActionTask.invocationEvent (JavaActionTask.java:45, assigned at :75), so the task event is never replaced. new ActionState(event) keeps that same reference (DurableExecutionManager.java:231), and InMemoryActionStateStore is a plain HashMap with no serialization. So this compares reference with itself.
Line 176 is the check that does the work. It already fails if resolution leaks back into the runtime event.
There is also a shape question. The identity only holds because this store does not serialize. A Kafka or Fluss store would round-trip the event and break isSameAs even though the contract still holds. ActionStateSerdeTest:96-97 uses assertInstanceOf plus assertEquals for the same idea. Would that fit better here?
|
|
||
| MemoryObject attachment = context.getSensoryMemory().get(reference); | ||
| if (attachment == null) { | ||
| throw new IllegalStateException( |
There was a problem hiding this comment.
nit: Python covers this branch with test_load_rejects_missing_event_attachment (test_event_attachment_utils.py:133), but I could not find a Java twin. EventAttachmentUtilsTest has seven tests and none reach the null path. Worth a short one next to loadsEventAttachments?
wenjin272
left a comment
There was a problem hiding this comment.
Thanks for the update and for adding the boundary tests. The attachment-related review comments look well addressed overall. I left one follow-up about handling a stored None value.
| msg = f"Failed to load event attachment: {_attachment_context(event, key, value.path)}" | ||
| raise EventAttachmentError(msg) from exc | ||
|
|
||
| if attachment is None: |
There was a problem hiding this comment.
None is a valid memory value, but get() returns None both for a stored None and for a missing path. As a result, a valid None attachment is incorrectly reported as missing here. Could we use is_exist(value.path) to distinguish these cases and add a regression test?
Linked issue: #476
Purpose of change
Add event attachments with automatic MemoryRef wrapping and resolution
Tests
API
Documentation
doc-neededAdd docs after review~
doc-not-neededdoc-included