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

[JinkunLiu]

Linked issue: #476

Purpose of change

Add event attachments with automatic MemoryRef wrapping and resolution

Tests

• Java and Python unit tests for resolving and restoring event attachments.
• Event serialization and deserialization tests.
• Cross-language snapshot tests covering Java-to-Python and Python-to-Java event conversion.
• End-to-end tests covering actions that consume and produce MemoryRef attachments.

API

• Add public APIs for event attachments in both Java and Python
• Add JSON serialization/deserialization support to MemoryRef.

Documentation

• doc-needed
  Add docs after review～
• doc-not-needed
• doc-included

[Copilot]

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 attachments map to the unified Event model in both Java and Python, plus (de)serialization support for MemoryRef in 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/hashCode ignore type, so a SHORT_TERM ref and a SENSORY ref with the same path compare equal. With memory_type now part of the JSON representation and cross-language contract, equality should include both type and path to 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.

[JinkunLiu]

Hello @wenjin272, Could you take a look at this PR when you have time? Thank you very much

[wenjin272]

Thanks for taking this on @JinkunLiu. I'll review it ASAP.

[wenjin272]

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.

[wenjin272]

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

[JinkunLiu]

Done，with this commit

[wenjin272]

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

[JinkunLiu]

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.

[wenjin272]

[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?

[JinkunLiu]

Fix it in both attachments and attributes. e2aa09e0

[wenjin272]

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

[weiqingy]

Thanks for taking this on. A few questions inline.

[weiqingy]

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?

[weiqingy]

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.

[weiqingy]

This writes the resolved value back into event.getAttachments(), and that event is ActionTask.event, the instance the runtime owns (JavaActionTask.java:60). Python resolves against a per-invocation copy instead (PythonActionExecutor.java:138-141), so its event keeps the refs.

With durable execution on, that undoes the offload: maybeInitActionState stores the live event before invoke() runs (DurableExecutionManager.java:219) and ActionState holds it by reference, so every later persist writes the payload inline instead of the ref. There may be a second effect on a heap backend, where ActionExecutionOperator.java:271-273 shares one event across sibling tasks, though that depends on ListState value semantics I did not run.

Was the in-place write deliberate, or would resolving into a copy the action owns work here?

[weiqingy]

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.