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 e7fbde464..2383428b8 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 @@ -20,14 +20,18 @@ 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.databind.ObjectMapper; +import javax.annotation.Nullable; + import java.io.IOException; import java.util.HashMap; import java.util.Map; import java.util.Objects; import java.util.UUID; +import java.util.function.BiFunction; /** Base class for all event types in the system. */ public class Event { @@ -38,6 +42,9 @@ public class Event { private final String type; private final Map attributes; + @Nullable private UUID upstreamEventId; + @Nullable private String upstreamActionName; + /** * Runtime-internal timestamp from the source record. Not part of the cross-language event * contract; used by the Flink runtime for timestamp propagation. @@ -54,17 +61,38 @@ public Event(String type) { this(type, new HashMap<>()); } + /** + * Reconstructs an Event with an existing identity and optional framework-managed lineage. + * + *

The lineage values support deserialization and reconstruction. When an Action emits this + * Event, the runtime overwrites them with the current trigger Event ID and Action name. + * + * @param id the existing Event ID + * @param type the Event type used for routing + * @param attributes the Event payload + * @param upstreamEventId the ID of the direct upstream Event, or {@code null} + * @param upstreamActionName the name of the emitting Action, or {@code null} + */ @JsonCreator public Event( @JsonProperty("id") UUID id, @JsonProperty("type") String type, - @JsonProperty("attributes") Map attributes) { + @JsonProperty("attributes") Map attributes, + @JsonProperty("upstreamEventId") @Nullable UUID upstreamEventId, + @JsonProperty("upstreamActionName") @Nullable String upstreamActionName) { if (type == null || type.isEmpty()) { throw new IllegalArgumentException("Event 'type' must not be null or empty."); } this.id = id; this.type = type; this.attributes = attributes != null ? attributes : new HashMap<>(); + this.upstreamEventId = upstreamEventId; + this.upstreamActionName = upstreamActionName; + } + + /** Reconstructs an Event with an existing identity and no upstream lineage. */ + public Event(UUID id, String type, Map attributes) { + this(id, type, attributes, null, null); } public UUID getId() { @@ -81,6 +109,38 @@ public Map getAttributes() { return attributes; } + /** Returns the ID of the Event consumed by the Action that emitted this Event. */ + @Nullable + @JsonInclude(JsonInclude.Include.NON_NULL) + public UUID getUpstreamEventId() { + return upstreamEventId; + } + + /** + * Sets the framework-managed ID of the Event consumed by the emitting Action. + * + *

The runtime overwrites this value when an Action emits the Event. + */ + public void setUpstreamEventId(@Nullable UUID upstreamEventId) { + this.upstreamEventId = upstreamEventId; + } + + /** Returns the name of the Action that emitted this Event. */ + @Nullable + @JsonInclude(JsonInclude.Include.NON_NULL) + public String getUpstreamActionName() { + return upstreamActionName; + } + + /** + * Sets the framework-managed name of the Action that emitted this Event. + * + *

The runtime overwrites this value when an Action emits the Event. + */ + public void setUpstreamActionName(@Nullable String upstreamActionName) { + this.upstreamActionName = upstreamActionName; + } + public Object getAttr(String name) { return attributes.get(name); } @@ -105,16 +165,38 @@ public void setSourceTimestamp(long timestamp) { } /** - * Creates a base Event from another Event, copying id, type, and attributes. Subclasses - * override this to reconstruct typed event objects with proper field deserialization. + * 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. */ public static Event fromEvent(Event event) { - Event copy = - new Event(event.getId(), event.getType(), new HashMap<>(event.getAttributes())); - if (event.hasSourceTimestamp()) { - copy.setSourceTimestamp(event.getSourceTimestamp()); + return reconstructFrom( + event, (id, attributes) -> new Event(id, event.getType(), attributes)); + } + + /** + * Reconstructs a typed Event while preserving the source identity and framework metadata. The + * factory receives the source ID and a copy of its attributes. + */ + protected static T reconstructFrom( + Event source, BiFunction, T> factory) { + Objects.requireNonNull(source, "source Event must not be null"); + Objects.requireNonNull(factory, "Event reconstruction factory must not be null"); + + T reconstructed = + Objects.requireNonNull( + factory.apply(source.getId(), new HashMap<>(source.getAttributes())), + "Event reconstruction factory must not return null"); + if (!Objects.equals(source.getId(), reconstructed.getId())) { + throw new IllegalStateException( + "Reconstructing the same Event occurrence must preserve Event ID " + + source.getId()); } - return copy; + Event reconstructedEvent = reconstructed; + reconstructedEvent.sourceTimestamp = source.sourceTimestamp; + reconstructedEvent.upstreamEventId = source.upstreamEventId; + reconstructedEvent.upstreamActionName = source.upstreamActionName; + return reconstructed; } /** diff --git a/api/src/main/java/org/apache/flink/agents/api/InputEvent.java b/api/src/main/java/org/apache/flink/agents/api/InputEvent.java index d07370215..3bf55a1f6 100644 --- a/api/src/main/java/org/apache/flink/agents/api/InputEvent.java +++ b/api/src/main/java/org/apache/flink/agents/api/InputEvent.java @@ -22,7 +22,6 @@ import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.annotation.JsonProperty; -import java.util.HashMap; import java.util.Map; import java.util.UUID; @@ -50,11 +49,7 @@ public InputEvent( * @return a typed InputEvent */ public static InputEvent fromEvent(Event event) { - InputEvent result = new InputEvent(event.getId(), new HashMap<>(event.getAttributes())); - if (event.hasSourceTimestamp()) { - result.setSourceTimestamp(event.getSourceTimestamp()); - } - return result; + return reconstructFrom(event, InputEvent::new); } @JsonIgnore 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 d7a7e0f4a..7c57d6341 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 @@ -22,7 +22,6 @@ import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.annotation.JsonProperty; -import java.util.HashMap; import java.util.Map; import java.util.UUID; @@ -53,11 +52,7 @@ public OutputEvent( * @return a typed OutputEvent */ public static OutputEvent fromEvent(Event event) { - OutputEvent result = new OutputEvent(event.getId(), new HashMap<>(event.getAttributes())); - if (event.hasSourceTimestamp()) { - result.setSourceTimestamp(event.getSourceTimestamp()); - } - return result; + return reconstructFrom(event, OutputEvent::new); } @JsonIgnore diff --git a/api/src/main/java/org/apache/flink/agents/api/event/ChatRequestEvent.java b/api/src/main/java/org/apache/flink/agents/api/event/ChatRequestEvent.java index 4e05d1c88..54ee175bf 100644 --- a/api/src/main/java/org/apache/flink/agents/api/event/ChatRequestEvent.java +++ b/api/src/main/java/org/apache/flink/agents/api/event/ChatRequestEvent.java @@ -29,7 +29,6 @@ import java.util.ArrayList; import java.util.Collections; -import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.UUID; @@ -96,12 +95,7 @@ private static Map normalizeAttributes(Map attri * @return a typed ChatRequestEvent */ public static ChatRequestEvent fromEvent(Event event) { - ChatRequestEvent result = - new ChatRequestEvent(event.getId(), new HashMap<>(event.getAttributes())); - if (event.hasSourceTimestamp()) { - result.setSourceTimestamp(event.getSourceTimestamp()); - } - return result; + return reconstructFrom(event, ChatRequestEvent::new); } @JsonIgnore diff --git a/api/src/main/java/org/apache/flink/agents/api/event/ChatResponseEvent.java b/api/src/main/java/org/apache/flink/agents/api/event/ChatResponseEvent.java index 8dab3b1d8..6a380e7ca 100644 --- a/api/src/main/java/org/apache/flink/agents/api/event/ChatResponseEvent.java +++ b/api/src/main/java/org/apache/flink/agents/api/event/ChatResponseEvent.java @@ -25,7 +25,6 @@ import org.apache.flink.agents.api.Event; import org.apache.flink.agents.api.chat.messages.ChatMessage; -import java.util.HashMap; import java.util.Map; import java.util.UUID; @@ -75,12 +74,7 @@ private static Map normalizeAttributes(Map attri * @return a typed ChatResponseEvent */ public static ChatResponseEvent fromEvent(Event event) { - ChatResponseEvent result = - new ChatResponseEvent(event.getId(), new HashMap<>(event.getAttributes())); - if (event.hasSourceTimestamp()) { - result.setSourceTimestamp(event.getSourceTimestamp()); - } - return result; + return reconstructFrom(event, ChatResponseEvent::new); } @JsonIgnore diff --git a/api/src/main/java/org/apache/flink/agents/api/event/ContextRetrievalRequestEvent.java b/api/src/main/java/org/apache/flink/agents/api/event/ContextRetrievalRequestEvent.java index 387319c95..255803081 100644 --- a/api/src/main/java/org/apache/flink/agents/api/event/ContextRetrievalRequestEvent.java +++ b/api/src/main/java/org/apache/flink/agents/api/event/ContextRetrievalRequestEvent.java @@ -23,7 +23,6 @@ import com.fasterxml.jackson.annotation.JsonProperty; import org.apache.flink.agents.api.Event; -import java.util.HashMap; import java.util.Map; import java.util.UUID; @@ -59,13 +58,7 @@ public ContextRetrievalRequestEvent( * @return a typed ContextRetrievalRequestEvent */ public static ContextRetrievalRequestEvent fromEvent(Event event) { - ContextRetrievalRequestEvent result = - new ContextRetrievalRequestEvent( - event.getId(), new HashMap<>(event.getAttributes())); - if (event.hasSourceTimestamp()) { - result.setSourceTimestamp(event.getSourceTimestamp()); - } - return result; + return reconstructFrom(event, ContextRetrievalRequestEvent::new); } @JsonIgnore diff --git a/api/src/main/java/org/apache/flink/agents/api/event/ContextRetrievalResponseEvent.java b/api/src/main/java/org/apache/flink/agents/api/event/ContextRetrievalResponseEvent.java index a71fd7031..6571cd2ef 100644 --- a/api/src/main/java/org/apache/flink/agents/api/event/ContextRetrievalResponseEvent.java +++ b/api/src/main/java/org/apache/flink/agents/api/event/ContextRetrievalResponseEvent.java @@ -26,7 +26,6 @@ import org.apache.flink.agents.api.vectorstores.Document; import java.util.ArrayList; -import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.UUID; @@ -82,13 +81,7 @@ private static Map normalizeAttributes(Map attri * @return a typed ContextRetrievalResponseEvent */ public static ContextRetrievalResponseEvent fromEvent(Event event) { - ContextRetrievalResponseEvent result = - new ContextRetrievalResponseEvent( - event.getId(), new HashMap<>(event.getAttributes())); - if (event.hasSourceTimestamp()) { - result.setSourceTimestamp(event.getSourceTimestamp()); - } - return result; + return reconstructFrom(event, ContextRetrievalResponseEvent::new); } @JsonIgnore diff --git a/api/src/main/java/org/apache/flink/agents/api/event/ToolRequestEvent.java b/api/src/main/java/org/apache/flink/agents/api/event/ToolRequestEvent.java index 9a24ef7ea..6ab53ba9b 100644 --- a/api/src/main/java/org/apache/flink/agents/api/event/ToolRequestEvent.java +++ b/api/src/main/java/org/apache/flink/agents/api/event/ToolRequestEvent.java @@ -23,7 +23,6 @@ import com.fasterxml.jackson.annotation.JsonProperty; import org.apache.flink.agents.api.Event; -import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.UUID; @@ -54,12 +53,7 @@ public ToolRequestEvent( */ @SuppressWarnings("unchecked") public static ToolRequestEvent fromEvent(Event event) { - ToolRequestEvent result = - new ToolRequestEvent(event.getId(), new HashMap<>(event.getAttributes())); - if (event.hasSourceTimestamp()) { - result.setSourceTimestamp(event.getSourceTimestamp()); - } - return result; + return reconstructFrom(event, ToolRequestEvent::new); } @JsonIgnore diff --git a/api/src/main/java/org/apache/flink/agents/api/event/ToolResponseEvent.java b/api/src/main/java/org/apache/flink/agents/api/event/ToolResponseEvent.java index 896f3bcb7..9c2884913 100644 --- a/api/src/main/java/org/apache/flink/agents/api/event/ToolResponseEvent.java +++ b/api/src/main/java/org/apache/flink/agents/api/event/ToolResponseEvent.java @@ -98,12 +98,7 @@ private static Map normalizeAttributes(Map attri * @return a typed ToolResponseEvent */ public static ToolResponseEvent fromEvent(Event event) { - ToolResponseEvent result = - new ToolResponseEvent(event.getId(), new HashMap<>(event.getAttributes())); - if (event.hasSourceTimestamp()) { - result.setSourceTimestamp(event.getSourceTimestamp()); - } - return result; + return reconstructFrom(event, ToolResponseEvent::new); } @JsonIgnore 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 c4a4ca104..8327884ac 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 @@ -35,6 +35,18 @@ class EventTest { private ObjectMapper objectMapper; + private static class CustomPayloadEvent extends Event { + private static final String EVENT_TYPE = "CustomPayloadEvent"; + + private CustomPayloadEvent(UUID id, Map attributes) { + super(id, EVENT_TYPE, attributes); + } + + public static CustomPayloadEvent fromEvent(Event event) { + return reconstructFrom(event, CustomPayloadEvent::new); + } + } + @BeforeEach void setUp() { objectMapper = new ObjectMapper(); @@ -129,6 +141,8 @@ void testUnifiedEventJsonDeserialization() throws Exception { assertEquals(id, event.getId()); assertEquals("MyEvent", event.getType()); assertEquals(1, event.getAttr("x")); + assertNull(event.getUpstreamEventId()); + assertNull(event.getUpstreamActionName()); } @Test @@ -143,6 +157,39 @@ void testSubclassedEventJsonRoundTrip() throws Exception { assertEquals(InputEvent.EVENT_TYPE, deserialized.getType()); } + @Test + void testLineageJsonRoundTrip() throws Exception { + UUID eventId = UUID.randomUUID(); + UUID upstreamEventId = UUID.randomUUID(); + Event original = + new Event( + eventId, + "ChildEvent", + Map.of("key", "value"), + upstreamEventId, + "child_action"); + + String json = objectMapper.writeValueAsString(original); + JsonNode node = objectMapper.readTree(json); + Event deserialized = objectMapper.readValue(json, Event.class); + + assertEquals(eventId.toString(), node.get("id").asText()); + assertEquals(upstreamEventId.toString(), node.get("upstreamEventId").asText()); + assertEquals("child_action", node.get("upstreamActionName").asText()); + assertEquals(eventId, deserialized.getId()); + assertEquals("value", deserialized.getAttr("key")); + assertEquals(upstreamEventId, deserialized.getUpstreamEventId()); + assertEquals("child_action", deserialized.getUpstreamActionName()); + } + + @Test + void testRootEventOmitsLineageFieldsFromJson() throws Exception { + JsonNode node = objectMapper.readTree(objectMapper.writeValueAsString(new InputEvent(1L))); + + assertFalse(node.has("upstreamEventId")); + assertFalse(node.has("upstreamActionName")); + } + // ── fromJson ─────────────────────────────────────────────────────────── @Test @@ -218,4 +265,70 @@ void testSourceTimestamp() { assertTrue(event.hasSourceTimestamp()); assertEquals(123456789L, event.getSourceTimestamp()); } + + @Test + void testFromEventReconstructsSameOccurrence() { + UUID upstreamEventId = UUID.randomUUID(); + Event original = new Event("Test"); + original.setSourceTimestamp(123456789L); + original.setUpstreamEventId(upstreamEventId); + original.setUpstreamActionName("test_action"); + + Event copy = Event.fromEvent(original); + + assertEquals(original.getId(), copy.getId()); + assertEquals(123456789L, copy.getSourceTimestamp()); + assertEquals(upstreamEventId, copy.getUpstreamEventId()); + assertEquals("test_action", copy.getUpstreamActionName()); + } + + @Test + void testTypedFromEventCopiesLineage() { + UUID upstreamEventId = UUID.randomUUID(); + Event original = + new Event( + UUID.randomUUID(), + OutputEvent.EVENT_TYPE, + new HashMap<>(Map.of("output", "result"))); + original.setUpstreamEventId(upstreamEventId); + original.setUpstreamActionName("output_action"); + + OutputEvent copy = OutputEvent.fromEvent(original); + + assertEquals(original.getId(), copy.getId()); + assertEquals(upstreamEventId, copy.getUpstreamEventId()); + assertEquals("output_action", copy.getUpstreamActionName()); + } + + @Test + void testCustomFactoryReconstructsSameOccurrenceWithoutManualMetadataCopy() { + UUID upstreamEventId = UUID.randomUUID(); + Event original = + new Event( + UUID.randomUUID(), + CustomPayloadEvent.EVENT_TYPE, + new HashMap<>(Map.of("value", "result"))); + original.setSourceTimestamp(123456789L); + original.setUpstreamEventId(upstreamEventId); + original.setUpstreamActionName("custom_action"); + + CustomPayloadEvent copy = CustomPayloadEvent.fromEvent(original); + + assertEquals(original.getId(), copy.getId()); + assertEquals("result", copy.getAttr("value")); + assertEquals(123456789L, copy.getSourceTimestamp()); + assertEquals(upstreamEventId, copy.getUpstreamEventId()); + assertEquals("custom_action", copy.getUpstreamActionName()); + } + + @Test + void testSameOccurrenceFactoryRejectsChangedEventId() { + Event original = new Event("Test"); + + assertThrows( + IllegalStateException.class, + () -> + Event.reconstructFrom( + original, (id, attributes) -> new Event("Test", attributes))); + } } diff --git a/docs/content/docs/development/workflow_agent.md b/docs/content/docs/development/workflow_agent.md index 74589a1fe..3eb917faa 100644 --- a/docs/content/docs/development/workflow_agent.md +++ b/docs/content/docs/development/workflow_agent.md @@ -631,6 +631,13 @@ public static void handleMyEvent(Event event, RunnerContext ctx) { {{< /tabs >}} +{{< hint info >}} +`upstreamEventId` and `upstreamActionName` (`upstream_event_id` and +`upstream_action_name` in Python) are framework-managed lineage metadata. User code should keep +user data in `attributes`. Values accepted during deserialization or reconstruction are +overwritten when an Action emits the Event. +{{< /hint >}} + ### JSON Serialization Events are serialized as JSON when passed between Python actions or across the Java-Python boundary. This means attribute values of non-trivial types (such as Pydantic models) lose their type information and arrive as plain `dict` objects. Users must manually reconstruct the typed object: @@ -658,11 +665,8 @@ class MyEvent(Event): @override def from_event(cls, event: Event) -> "MyEvent": assert "value" in event.attributes - result = MyEvent(value=event.attributes["value"]) - # Preserve the base event id. Assign it last: the content-based id is - # regenerated whenever another field changes. - result.id = event.id - return result + result = cls(value=event.attributes["value"]) + return result.reconstruct_from(event) @property def value(self) -> str: @@ -688,13 +692,7 @@ public class MyEvent extends Event { } public static MyEvent fromEvent(Event event) { - // Preserve the base event id (and sourceTimestamp) so event logs, listeners, - // correlation, deduplication, and timestamp propagation stay consistent. - MyEvent result = new MyEvent(event.getId(), new HashMap<>(event.getAttributes())); - if (event.hasSourceTimestamp()) { - result.setSourceTimestamp(event.getSourceTimestamp()); - } - return result; + return reconstructFrom(event, MyEvent::new); } public String getValue() { @@ -707,16 +705,17 @@ public class MyEvent extends Event { {{< /tabs >}} {{< hint info >}} -When reconstructing a typed event, preserve the base `Event` metadata so that event logs, -listeners, correlation, deduplication, and downstream timestamp propagation stay consistent with -built-in events: - -- **`id`**: copy the source event's `id` onto the reconstructed event, as all built-in events do - in both languages. In Python, assign `result.id = event.id` **last**, because the content-based - `id` is regenerated whenever any other field changes. -- **`sourceTimestamp`** (Java only): carry it over with `setSourceTimestamp(...)` when - `hasSourceTimestamp()` is true, matching built-in Java events. This field is runtime-internal and - used for timestamp propagation; the Python `Event` has no equivalent. +Typed reconstruction represents the same Event occurrence, so it must preserve the base Event's +identity and framework-managed metadata: + +- **Python**: return `reconstruct_from(event)` after constructing the typed object. + It returns a new typed object with the Event's UUIDv4 `id`, `upstream_event_id`, and + `upstream_action_name`; the returned Event's `id` remains immutable. +- **Java**: implement the typed constructor that accepts `(UUID id, Map + attributes)`, then have `fromEvent` return + `reconstructFrom(event, MyEvent::new)`. The framework supplies the original ID and + attributes, validates that the ID is preserved, and carries over `sourceTimestamp`, + `upstreamEventId`, and `upstreamActionName`. {{< /hint >}} {{< hint info >}} @@ -756,4 +755,4 @@ private static Map normalizeAttributes(Map attri There are several built-in `Event` and `Action` in Flink-Agents: * See [Chat Models]({{< ref "docs/development/chat_models#built-in-events-and-actions" >}}) for how to chat with a LLM leveraging built-in action and events. * See [Tool Use]({{< ref "docs/development/tool_use#built-in-events-and-actions" >}}) for how to programmatically use a tool leveraging built-in action and events. -* See [Vector Stores]({{< ref "docs/development/vector_stores#built-in-events-and-actions" >}}) for how to retrieve context from vector stores leveraging built-in action and events. \ No newline at end of file +* See [Vector Stores]({{< ref "docs/development/vector_stores#built-in-events-and-actions" >}}) for how to retrieve context from vector stores leveraging built-in action and events. diff --git a/docs/content/docs/operations/deployment.md b/docs/content/docs/operations/deployment.md index 4178b94de..384e2ab0e 100644 --- a/docs/content/docs/operations/deployment.md +++ b/docs/content/docs/operations/deployment.md @@ -94,6 +94,8 @@ After recovery from a checkpoint, Flink Agents reprocess events that arrived aft To ensure exactly-once action consistency, you must configure an external action state store. Flink Agents record action state in this store on a per-action basis. After recovering from a checkpoint, Flink Agents consult the external store and will not re-execute actions that were already completed. This guarantees each action is executed exactly once after recovering from a checkpoint. +When a completed Action is reused during recovery, its stored output Events keep their original Event IDs, while their lineage is rebound to the Event that triggers the reused Action in the recovered execution. + The same persisted action state is also used by fine-grained durable execution. {{< hint info >}} diff --git a/docs/content/docs/operations/monitoring.md b/docs/content/docs/operations/monitoring.md index 021452125..410317546 100644 --- a/docs/content/docs/operations/monitoring.md +++ b/docs/content/docs/operations/monitoring.md @@ -184,14 +184,125 @@ Each record contains a top-level `timestamp`, the resolved `logLevel`, and a top { "timestamp": "2024-01-15T10:30:00Z", "logLevel": "STANDARD", - "eventType": "_input_event", + "eventType": "MiddleEvent", "event": { - "eventType": "_input_event", - "...": "..." + "eventType": "MiddleEvent", + "id": "39361629-4f1d-4b62-b734-a48b181cb6e0", + "attributes": {}, + "type": "MiddleEvent", + "upstreamEventId": "dad5ed00-80e2-4746-8bdb-f126bad504b5", + "upstreamActionName": "action1" } } ``` +`upstreamEventId` identifies the Event consumed by the Action that emitted the current Event, and `upstreamActionName` identifies that Action. The framework maintains both fields directly on the `event` object, outside business `attributes`. A root `InputEvent` omits both fields. + +### Trace Tree Reconstruction + +The `flink-agents-trace-tree` command is installed with the Flink Agents Python wheel. It rebuilds InputEvent-rooted Trace Trees from a saved File Event Log. Pass either one log file for text output or a log directory for Trace Tree JSON: + +```bash +flink-agents-trace-tree /path/to/events-job-task-0.log +flink-agents-trace-tree /path/to/event-log-directory --format json +``` + +For example, the text output for an `InputEvent -> MiddleEvent -> OutputEvent` lineage is: + +```text +Trace Tree 1 + _input_event (dad5ed00-80e2-4746-8bdb-f126bad504b5) + [Action: action1] + MiddleEvent (39361629-4f1d-4b62-b734-a48b181cb6e0) + [Action: action2] + _output_event (f452ce08-c672-4c9c-841f-d81d15a900c5) +``` + +In JSON output: + +- `roots` lists the Event IDs of the root `InputEvent`s, one per reconstructed Trace Tree. +- `nodes` maps each logical Event ID to one Event node with the following fields: + - `eventId`: the logical Event ID of the node. + - `eventType`: the top-level `eventType` routing key from the Event Log record. + - `timestamp`: the `timestamp` of the first Event Log record observed for this Event ID. + - `observationCount`: Number of times this Event was recorded in the Event Log. Values greater + than 1 usually mean the Event was reused during replay. + - `upstreamEdges`: the distinct causal edges observed for this Event, each with an + `upstreamEventId` and an `upstreamActionName`. Repeated observations of the same + `(Event ID, upstream Event ID, upstream Action name)` edge are deduplicated. + - `actions`: the virtual Action nodes triggered by this Event, each with the Action `name` and + the `children` Event IDs it emitted. +- `warnings` contains the reconstruction warnings described below. + +Distinct edges are retained, so reused Events can be shared by multiple InputEvent-rooted +branches and the combined result is a DAG rather than a strict tree. For example, two roots that +share one reused output Event through the same Action produce: + +```json +{ + "roots": ["root-1", "root-2"], + "nodes": { + "root-1": { + "eventId": "root-1", + "eventType": "_input_event", + "timestamp": "2024-01-15T10:30:00Z", + "observationCount": 1, + "upstreamEdges": [], + "actions": [ + {"name": "action1", "children": ["shared-1"]} + ] + }, + "root-2": { + "eventId": "root-2", + "eventType": "_input_event", + "timestamp": "2024-01-15T10:31:00Z", + "observationCount": 1, + "upstreamEdges": [], + "actions": [ + {"name": "action1", "children": ["shared-1"]} + ] + }, + "shared-1": { + "eventId": "shared-1", + "eventType": "_output_event", + "timestamp": "2024-01-15T10:30:01Z", + "observationCount": 2, + "upstreamEdges": [ + {"upstreamEventId": "root-1", "upstreamActionName": "action1"}, + {"upstreamEventId": "root-2", "upstreamActionName": "action1"} + ], + "actions": [] + } + }, + "warnings": [] +} +``` + +{{< hint info >}} +During ActionStateStore recovery, a reused output Event keeps its Event ID when a completed +Action result is replayed, while its lineage is rebound to the recovered execution's current +triggering Event. The same logical Event ID may therefore appear in multiple physical Event Log +records with distinct upstream edges, as in the example above. Reusing an Event ID with the same +Event type and content is valid. +{{< /hint >}} + +The reader derives virtual Action nodes from `upstreamActionName`; they are not separate Event Log +records. Log order controls display order only and does not add execution-order semantics. + +Reconstruction warnings are written to standard error and included in JSON output while valid +InputEvent-rooted branches are retained. The reader keeps valid records from the same input and +continues with other Event Log files. + +| Warning Code | Trigger Condition | +|------------------------|----------------------------------------------------------------------------| +| `EVENT_ID_CONFLICT` | Records carrying the same Event ID disagree on the Event type or content. | +| `INVALID_ROOT_LINEAGE` | An InputEvent carries upstream lineage, which a root Event must not have. | +| `UNLINKED_EVENT` | A non-InputEvent has no upstream Event. | +| `MISSING_ACTION_NAME` | An Event has an upstream Event but no upstream Action name. | +| `MISSING_PARENT` | An Event references an upstream Event with no valid node in the Event Log. | +| `MALFORMED_RECORD` | An Event Log record is invalid or partially written. | +| `UNREADABLE_FILE` | An Event Log file could not be read. | + ### Event Log Levels Each event type is logged at a configurable verbosity. Three levels are supported: @@ -204,7 +315,7 @@ Each event type is logged at a configurable verbosity. Three levels are supporte The global default is set by [`event-log.level`]({{< ref "docs/operations/configuration#core-options" >}}). At `STANDARD` level, the payload is shrunk along three independent axes — long strings, large arrays, and deep nesting — controlled by `event-log.standard.max-string-length`, `event-log.standard.max-array-elements`, and `event-log.standard.max-depth` respectively. Setting any threshold to `0` disables that specific truncation; setting all three to `0` makes `STANDARD` behave identically to `VERBOSE` (apart from the `logLevel` label). The exact truncation strategy may evolve over time; the contract is only that `STANDARD` keeps logs concise while `VERBOSE` preserves the full payload. -**Fields that are never truncated.** Structural and identifying fields are always preserved in full so log consumers can still group, route, and correlate records: `timestamp`, `logLevel`, top-level `eventType`, and the event's own `id`, `type`, and short scalar fields. Truncation only applies to large nested content (long strings, big arrays, deeply nested objects). +**Fields that are never truncated.** Structural and identifying fields are always preserved in full so log consumers can still group, route, and correlate records: `timestamp`, `logLevel`, top-level `eventType`, and the event's own `eventType`, `type`, `id`, `upstreamEventId`, and `upstreamActionName`. The `attributes` envelope is also preserved, while large nested content inside it can still be truncated. **Truncation wrapper format.** When a field is truncated at `STANDARD` level it is replaced by a JSON object that records what was retained and what was dropped. This keeps the record valid JSON and lets downstream tooling detect truncation programmatically: diff --git a/python/flink_agents/api/events/chat_event.py b/python/flink_agents/api/events/chat_event.py index a3a3deec1..8d0d58ea9 100644 --- a/python/flink_agents/api/events/chat_event.py +++ b/python/flink_agents/api/events/chat_event.py @@ -83,8 +83,7 @@ def from_event(cls, event: Event) -> "ChatRequestEvent": prompt_args=event.attributes.get("prompt_args"), output_schema=output_schema_raw, ) - result.id = event.id - return result + return result.reconstruct_from(event) @property def model(self) -> str: @@ -160,8 +159,7 @@ def from_event(cls, event: Event) -> "ChatResponseEvent": retry_count=event.attributes.get("retry_count", 0), total_retry_wait_sec=event.attributes.get("total_retry_wait_sec", 0), ) - result.id = event.id - return result + return result.reconstruct_from(event) @property def request_id(self) -> UUID: diff --git a/python/flink_agents/api/events/context_retrieval_event.py b/python/flink_agents/api/events/context_retrieval_event.py index a8245a4ec..d014e1854 100644 --- a/python/flink_agents/api/events/context_retrieval_event.py +++ b/python/flink_agents/api/events/context_retrieval_event.py @@ -63,8 +63,7 @@ def from_event(cls, event: Event) -> "ContextRetrievalRequestEvent": vector_store=event.attributes["vector_store"], max_results=event.attributes.get("max_results", 3), ) - result.id = event.id - return result + return result.reconstruct_from(event) @property def query(self) -> str: @@ -124,8 +123,7 @@ def from_event(cls, event: Event) -> "ContextRetrievalResponseEvent": query=event.attributes["query"], documents=documents, ) - result.id = event.id - return result + return result.reconstruct_from(event) @property def request_id(self) -> UUID: diff --git a/python/flink_agents/api/events/event.py b/python/flink_agents/api/events/event.py index 77e7e5faf..2db294d7f 100644 --- a/python/flink_agents/api/events/event.py +++ b/python/flink_agents/api/events/event.py @@ -15,17 +15,23 @@ # See the License for the specific language governing permissions and # limitations under the License. ################################################################################# -import hashlib import json from typing import Any, ClassVar, Dict try: - from typing import override + from typing import Self, override except ImportError: - from typing_extensions import override -from uuid import UUID - -from pydantic import BaseModel, Field, model_validator + from typing_extensions import Self, override +from uuid import UUID, uuid4 + +from pydantic import ( + AliasChoices, + BaseModel, + Field, + SerializerFunctionWrapHandler, + model_serializer, + model_validator, +) from pydantic_core import PydanticSerializationError from pyflink.common import Row @@ -66,17 +72,30 @@ class Event(BaseModel, extra="allow"): Attributes: ---------- id : UUID - Unique identifier for the event, generated deterministically based on - event content. + Random version 4 UUID generated when the Event is created. type : str Event type string used for routing. Required for all events. attributes : Dict[str, Any] Key-value properties for the event data. + upstream_event_id : UUID | None + The ID of the direct upstream Event, or None. + upstream_action_name : str | None + The name of the emitting Action, or None. """ - id: UUID = Field(default=None) + id: UUID = Field(default_factory=uuid4, frozen=True) type: str attributes: Dict[str, Any] = Field(default_factory=dict) + upstream_event_id: UUID | None = Field( + default=None, + validation_alias=AliasChoices("upstream_event_id", "upstreamEventId"), + serialization_alias="upstreamEventId", + ) + upstream_action_name: str | None = Field( + default=None, + validation_alias=AliasChoices("upstream_action_name", "upstreamActionName"), + serialization_alias="upstreamActionName", + ) @staticmethod def __serialize_unknown(field: Any) -> Dict[str, Any]: @@ -98,23 +117,25 @@ def model_dump_json(self, **kwargs: Any) -> str: kwargs["fallback"] = self.__serialize_unknown return super().model_dump_json(**kwargs) - def _generate_content_based_id(self) -> UUID: - """Generate a deterministic UUID based on event content using MD5 hash. - - Similar to Java's UUID.nameUUIDFromBytes(), uses MD5 for version 3 UUID. - """ - # Serialize content excluding 'id' to avoid circular dependency - content_json = super().model_dump_json( - exclude={"id"}, fallback=self.__serialize_unknown - ) - md5_hash = hashlib.md5(content_json.encode()).digest() - return UUID(bytes=md5_hash, version=3) + @model_serializer(mode="wrap") + def _serialize_event( + self, handler: SerializerFunctionWrapHandler + ) -> Dict[str, Any]: + """Use cross-language names only for lineage and omit empty lineage.""" + serialized: Dict[str, Any] = handler(self) + missing = object() + for field_name, alias in ( + ("upstream_event_id", "upstreamEventId"), + ("upstream_action_name", "upstreamActionName"), + ): + value = serialized.pop(field_name, serialized.pop(alias, missing)) + if value is not missing and value is not None: + serialized[alias] = value + return serialized @model_validator(mode="after") - def validate_and_set_id(self) -> "Event": - """Validate that fields are serializable and generate content-based ID.""" - if self.id is None: - object.__setattr__(self, "id", self._generate_content_based_id()) + def validate_serializable_fields(self) -> "Event": + """Validate that all Event fields can be serialized.""" self.model_dump_json() return self @@ -122,9 +143,16 @@ def __setattr__(self, name: str, value: Any) -> None: super().__setattr__(name, value) # Ensure added property can be serialized. self.model_dump_json() - # Regenerate ID if content changed (but not if setting 'id' itself) - if name != "id": - object.__setattr__(self, "id", self._generate_content_based_id()) + + 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, + "upstream_event_id": source.upstream_event_id, + "upstream_action_name": source.upstream_action_name, + } + ) def get_type(self) -> str: """Return the event type string used for routing.""" @@ -200,8 +228,7 @@ def __init__(self, input: Any) -> None: def from_event(cls, event: Event) -> "InputEvent": assert "input" in event.attributes result = InputEvent(input=event.attributes["input"]) - result.id = event.id - return result + return result.reconstruct_from(event) @property def input(self) -> Any: @@ -233,8 +260,7 @@ def __init__(self, output: Any) -> None: def from_event(cls, event: Event) -> "OutputEvent": assert "output" in event.attributes result = OutputEvent(output=event.attributes["output"]) - result.id = event.id - return result + return result.reconstruct_from(event) @property def output(self) -> Any: diff --git a/python/flink_agents/api/events/tool_event.py b/python/flink_agents/api/events/tool_event.py index b7396b6f7..bd97a1c54 100644 --- a/python/flink_agents/api/events/tool_event.py +++ b/python/flink_agents/api/events/tool_event.py @@ -58,8 +58,7 @@ def from_event(cls, event: Event) -> "ToolRequestEvent": model=event.attributes["model"], tool_calls=event.attributes["tool_calls"], ) - result.id = event.id - return result + return result.reconstruct_from(event) @property def model(self) -> str: @@ -120,13 +119,10 @@ def from_event(cls, event: Event) -> "ToolResponseEvent": request_id=event.attributes["request_id"], responses=responses, external_ids=event.attributes.get("external_ids", {}), - success=event.attributes.get( - "success", dict.fromkeys(responses, True) - ), + success=event.attributes.get("success", dict.fromkeys(responses, True)), error=event.attributes.get("error", {}), ) - result.id = event.id - return result + return result.reconstruct_from(event) @property def request_id(self) -> UUID: 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 315301aeb..83648967a 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 @@ -50,8 +50,7 @@ def _regenerate_enabled() -> bool: def _force_id(event: Event, fixed_id: UUID) -> Event: - object.__setattr__(event, "id", fixed_id) - return event + return event.model_copy(update={"id": fixed_id}) def _write_python_snapshot(name: str, event: Event) -> None: diff --git a/python/flink_agents/api/tests/test_event.py b/python/flink_agents/api/tests/test_event.py index 110106ea4..d4b3df62b 100644 --- a/python/flink_agents/api/tests/test_event.py +++ b/python/flink_agents/api/tests/test_event.py @@ -16,16 +16,33 @@ # limitations under the License. ################################################################################# import json -from typing import Any, Type +from typing import Any, ClassVar, Type +from uuid import UUID, uuid4 import pytest -from pydantic import ValidationError +from pydantic import Field, ValidationError from pydantic_core import PydanticSerializationError from pyflink.common import Row from flink_agents.api.events.event import Event, InputEvent, OutputEvent +class _CustomEvent(Event): + EVENT_TYPE: ClassVar[str] = "custom_event" + + def __init__(self, value: str) -> None: + super().__init__(type=self.EVENT_TYPE, attributes={"value": value}) + + @classmethod + def from_event(cls, event: Event) -> "_CustomEvent": + result = cls(value=event.attributes["value"]) + return result.reconstruct_from(event) + + +class _CustomAliasedEvent(Event): + custom_value: str = Field(serialization_alias="customValue") + + def test_event_init_serializable() -> None: Event(type="test", a=1, b=InputEvent(input=1), c=OutputEvent(output="111")) @@ -52,7 +69,11 @@ def test_input_event_ignore_row_unserializable() -> None: def test_event_row_with_non_serializable_fails() -> None: with pytest.raises(ValidationError): - Event(type="test", row_field=Row({"a": 1}), non_serializable_field=Type[InputEvent]) + Event( + type="test", + row_field=Row({"a": 1}), + non_serializable_field=Type[InputEvent], + ) def test_event_multiple_rows_serializable() -> None: @@ -64,6 +85,31 @@ def test_event_setattr_row_serializable() -> None: event.row_field = Row({"key": "value"}) +def test_events_with_same_content_have_distinct_random_ids() -> None: + first = OutputEvent(output="same") + second = OutputEvent(output="same") + + assert first.id != second.id + assert first.id.version == 4 + assert second.id.version == 4 + + +def test_event_id_does_not_change_with_event_content() -> None: + event = Event(type="test", a=1) + event_id = event.id + + event.a = 2 + + assert event.id == event_id + + +def test_event_id_cannot_be_reassigned() -> None: + event = Event(type="test") + + with pytest.raises(ValidationError, match="Field is frozen"): + event.id = uuid4() + + def test_event_json_serialization_with_row() -> None: event = InputEvent(input=Row({"test": "data"})) json_str = event.model_dump_json() @@ -191,6 +237,27 @@ def test_unified_event_from_json_missing_type() -> None: Event.from_json(json.dumps({"attributes": {}})) +def test_lineage_serialization_respects_field_filters() -> None: + """Test lineage aliases do not bypass Pydantic include and exclude filters.""" + event = Event( + type="ChildEvent", + upstreamEventId=UUID("00000000-0000-0000-0000-000000000001"), + upstreamActionName="child_action", + ) + + assert event.model_dump(include={"type"}) == {"type": "ChildEvent"} + assert json.loads(event.model_dump_json(include={"type"})) == {"type": "ChildEvent"} + + excluded_fields = {"upstream_event_id", "upstream_action_name"} + dumped = event.model_dump(exclude=excluded_fields) + json_dumped = json.loads(event.model_dump_json(exclude=excluded_fields)) + + assert "upstreamEventId" not in dumped + assert "upstreamActionName" not in dumped + assert "upstreamEventId" not in json_dumped + assert "upstreamActionName" not in json_dumped + + def test_unified_event_serialization_roundtrip() -> None: """Test that unified events survive JSON serialization/deserialization.""" original = Event(type="RoundTrip", attributes={"a": 1, "b": "two"}) @@ -203,6 +270,106 @@ def test_unified_event_serialization_roundtrip() -> None: assert restored.attributes == {"a": 1, "b": "two"} +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") + event = Event(type="ChildEvent") + event_id = event.id + + event.upstream_event_id = upstream_event_id + event.upstream_action_name = "child_action" + + parsed = json.loads(event.model_dump_json()) + restored = Event.from_json(json.dumps(parsed)) + + assert event.id == event_id + assert parsed["upstreamEventId"] == str(upstream_event_id) + assert parsed["upstreamActionName"] == "child_action" + assert "upstream_event_id" not in parsed + assert "upstream_action_name" not in parsed + assert restored.upstream_event_id == upstream_event_id + assert restored.upstream_action_name == "child_action" + + +def test_event_lineage_aliases_do_not_enable_custom_aliases_by_default() -> None: + """Test only lineage uses cross-language aliases unless explicitly requested.""" + upstream_event_id = UUID("00000000-0000-0000-0000-000000000001") + event = _CustomAliasedEvent( + type="CustomAliasedEvent", + custom_value="value", + upstreamEventId=upstream_event_id, + upstreamActionName="custom_action", + ) + + default_json = json.loads(event.model_dump_json()) + aliased_json = json.loads(event.model_dump_json(by_alias=True)) + + assert default_json["custom_value"] == "value" + assert "customValue" not in default_json + assert default_json["upstreamEventId"] == str(upstream_event_id) + assert default_json["upstreamActionName"] == "custom_action" + assert aliased_json["customValue"] == "value" + assert "custom_value" not in aliased_json + + +def test_root_event_omits_lineage_fields_from_json() -> None: + """Test a root Event does not serialize empty lineage fields.""" + parsed = json.loads(InputEvent(input="root").model_dump_json()) + + assert "upstreamEventId" not in parsed + assert "upstreamActionName" not in parsed + + +def test_typed_from_event_preserves_lineage() -> None: + """Test typed reconstruction keeps framework-managed lineage metadata.""" + upstream_event_id = UUID("00000000-0000-0000-0000-000000000001") + event = Event( + type="_output_event", + attributes={"output": "result"}, + upstreamEventId=upstream_event_id, + upstreamActionName="output_action", + ) + + reconstructed = OutputEvent.from_event(event) + + assert reconstructed.upstream_event_id == upstream_event_id + assert reconstructed.upstream_action_name == "output_action" + + +def test_custom_typed_from_event_preserves_identity_and_lineage() -> None: + """Test custom typed reconstruction follows the framework metadata contract.""" + upstream_event_id = UUID("00000000-0000-0000-0000-000000000001") + event = Event( + type=_CustomEvent.EVENT_TYPE, + attributes={"value": "result"}, + upstreamEventId=upstream_event_id, + upstreamActionName="custom_action", + ) + + reconstructed = _CustomEvent.from_event(event) + + assert reconstructed.id == event.id + assert reconstructed.upstream_event_id == upstream_event_id + assert reconstructed.upstream_action_name == "custom_action" + + +def test_same_occurrence_reconstruction_returns_a_copy() -> None: + """Test occurrence reconstruction does not mutate the typed draft.""" + source = Event(type="_output_event", attributes={"output": "result"}) + source.upstream_action_name = "output_action" + draft = OutputEvent(output="result") + draft_id = draft.id + + reconstructed = draft.reconstruct_from(source) + + assert reconstructed is not draft + assert reconstructed.id == source.id + assert reconstructed.upstream_action_name == "output_action" + assert draft.id == draft_id + assert draft.id != source.id + assert draft.upstream_action_name is None + + def test_unified_event_serialization_roundtrip_with_row() -> None: """Test that unified events with Row fields survive JSON roundtrip.""" original = Event( diff --git a/python/flink_agents/cli/__init__.py b/python/flink_agents/cli/__init__.py new file mode 100644 index 000000000..6df135cc4 --- /dev/null +++ b/python/flink_agents/cli/__init__.py @@ -0,0 +1,18 @@ +################################################################################ +# 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. +################################################################################ +"""Command-line interfaces shipped with Flink Agents.""" diff --git a/python/flink_agents/cli/tests/__init__.py b/python/flink_agents/cli/tests/__init__.py new file mode 100644 index 000000000..65b48d4d7 --- /dev/null +++ b/python/flink_agents/cli/tests/__init__.py @@ -0,0 +1,17 @@ +################################################################################ +# 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. +################################################################################ diff --git a/python/flink_agents/cli/tests/test_trace_tree.py b/python/flink_agents/cli/tests/test_trace_tree.py new file mode 100644 index 000000000..e63df1f8f --- /dev/null +++ b/python/flink_agents/cli/tests/test_trace_tree.py @@ -0,0 +1,466 @@ +################################################################################ +# 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 json +import subprocess +import sys +from pathlib import Path + +import pytest + +from flink_agents.cli.trace_tree import find_log_files + + +def _record( + event_id: str, + event_type: str, + upstream_event_id: str | None = None, + upstream_action_name: str | None = None, + attributes: dict | None = None, +) -> dict: + event = { + "id": event_id, + "eventType": event_type, + "attributes": attributes or {}, + } + if upstream_event_id is not None: + event["upstreamEventId"] = upstream_event_id + if upstream_action_name is not None: + event["upstreamActionName"] = upstream_action_name + return { + "timestamp": "2026-07-17T10:00:00Z", + "eventType": event_type, + "event": event, + } + + +def _write_log(path: Path, records: list[dict]) -> None: + path.write_text("".join(json.dumps(record) + "\n" for record in records)) + + +def _write_pretty_log(path: Path, records: list[dict]) -> None: + path.write_text("".join(json.dumps(record, indent=2) + "\n" for record in records)) + + +def _run_reader(log_path: Path, output_format: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [ + sys.executable, + "-m", + "flink_agents.cli.trace_tree", + str(log_path), + "--format", + output_format, + ], + check=True, + capture_output=True, + text=True, + ) + + +def test_reader_reconstructs_text_and_json_in_log_order(tmp_path: Path) -> None: + log_path = tmp_path / "events.log" + _write_log( + log_path, + [ + _record("root", "_input_event"), + _record("child-1", "ChildEvent", "root", "branch_action"), + _record("child-2", "ChildEvent", "root", "branch_action"), + _record("output", "_output_event", "child-1", "output_action"), + _record("child-3", "ChildEvent", "root", "second_action"), + ], + ) + + json_result = _run_reader(log_path, "json") + trace_forest = json.loads(json_result.stdout) + text_result = _run_reader(log_path, "text") + + assert trace_forest["roots"] == ["root"] + assert trace_forest["nodes"]["root"]["actions"] == [ + {"name": "branch_action", "children": ["child-1", "child-2"]}, + {"name": "second_action", "children": ["child-3"]}, + ] + assert trace_forest["nodes"]["child-1"]["actions"] == [ + {"name": "output_action", "children": ["output"]} + ] + assert trace_forest["warnings"] == [] + assert json_result.stderr == "" + expected_text_order = [ + "_input_event (root)", + "[Action: branch_action]", + "ChildEvent (child-1)", + "[Action: output_action]", + "_output_event (output)", + "ChildEvent (child-2)", + "[Action: second_action]", + "ChildEvent (child-3)", + ] + positions = [text_result.stdout.index(item) for item in expected_text_order] + assert positions == sorted(positions) + assert text_result.stderr == "" + + +def test_reader_reconstructs_pretty_printed_records(tmp_path: Path) -> None: + log_path = tmp_path / "events.log" + _write_pretty_log( + log_path, + [ + _record("root", "_input_event"), + _record("child-1", "ChildEvent", "root", "branch_action"), + _record("child-2", "ChildEvent", "root", "branch_action"), + ], + ) + + json_result = _run_reader(log_path, "json") + trace_forest = json.loads(json_result.stdout) + text_result = _run_reader(log_path, "text") + + assert trace_forest["roots"] == ["root"] + assert trace_forest["nodes"]["root"]["actions"] == [ + {"name": "branch_action", "children": ["child-1", "child-2"]} + ] + assert trace_forest["warnings"] == [] + assert "_input_event (root)" in text_result.stdout + assert "[Action: branch_action]" in text_result.stdout + assert text_result.stdout.index("ChildEvent (child-1)") < text_result.stdout.index( + "ChildEvent (child-2)" + ) + assert json_result.stderr == "" + assert text_result.stderr == "" + + +def test_reader_warns_on_truncated_tail_and_keeps_valid_records( + tmp_path: Path, +) -> None: + log_path = tmp_path / "events.log" + root = _record("root", "_input_event") + child = _record("child", "ChildEvent", "root", "child_action") + log_path.write_text( + json.dumps(root) + "\n" + json.dumps(child) + '\n{"event":', + encoding="utf-8", + ) + + result = _run_reader(log_path, "json") + trace_forest = json.loads(result.stdout) + + assert trace_forest["roots"] == ["root"] + assert trace_forest["nodes"]["root"]["actions"] == [ + {"name": "child_action", "children": ["child"]} + ] + assert trace_forest["warnings"] == [ + { + "code": "MALFORMED_RECORD", + "message": ( + f"Could not decode an Event Log record in {log_path} " + "at line 3, column 10: Expecting value." + ), + "filePath": str(log_path), + "lineNumber": 3, + "columnNumber": 10, + } + ] + assert "[MALFORMED_RECORD]" in result.stderr + + +def test_reader_resynchronizes_after_malformed_pretty_record( + tmp_path: Path, +) -> None: + log_path = tmp_path / "events.log" + root = _record("root", "_input_event") + child = _record("child", "ChildEvent", "root", "child_action") + log_path.write_text( + json.dumps(root, indent=2) + + "\n" + + '{"event": invalid}\n' + + json.dumps(child, indent=2) + + "\n", + encoding="utf-8", + ) + + result = _run_reader(log_path, "json") + trace_forest = json.loads(result.stdout) + + assert trace_forest["roots"] == ["root"] + assert trace_forest["nodes"]["root"]["actions"] == [ + {"name": "child_action", "children": ["child"]} + ] + assert [item["code"] for item in trace_forest["warnings"]] == ["MALFORMED_RECORD"] + assert "[MALFORMED_RECORD]" in result.stderr + + +@pytest.mark.parametrize( + ("malformed_record", "expected_event_id"), + [ + ([], None), + ({}, None), + ({"eventType": "BadEvent", "event": []}, None), + ({"eventType": "BadEvent", "event": {"attributes": {}}}, None), + ({"eventType": 123, "event": {"id": "bad"}}, None), + ( + { + "eventType": "BadEvent", + "event": {"id": "bad", "upstreamEventId": []}, + }, + "bad", + ), + ], +) +def test_reader_warns_on_invalid_record_shape_and_continues( + tmp_path: Path, malformed_record: object, expected_event_id: str | None +) -> None: + log_path = tmp_path / "events.log" + _write_log( + log_path, + [ + _record("root", "_input_event"), + malformed_record, + _record("child", "ChildEvent", "root", "child_action"), + ], + ) + + result = _run_reader(log_path, "json") + trace_forest = json.loads(result.stdout) + + assert trace_forest["roots"] == ["root"] + assert trace_forest["nodes"]["root"]["actions"] == [ + {"name": "child_action", "children": ["child"]} + ] + assert [item["code"] for item in trace_forest["warnings"]] == ["MALFORMED_RECORD"] + malformed_warning = trace_forest["warnings"][0] + assert malformed_warning["filePath"] == str(log_path) + assert malformed_warning["message"].startswith( + f"Invalid Event Log record in {log_path}: " + ) + assert "lineNumber" not in malformed_warning + assert "columnNumber" not in malformed_warning + if expected_event_id is None: + assert "eventId" not in malformed_warning + else: + assert malformed_warning["eventId"] == expected_event_id + + +def test_directory_discovery_ignores_unrelated_log_files(tmp_path: Path) -> None: + unrelated_log = tmp_path / "taskmanager.log" + unrelated_log.write_text("not an Event Log", encoding="utf-8") + + assert find_log_files(tmp_path) == [] + + +def test_reader_warns_on_unreadable_file_and_keeps_other_files( + tmp_path: Path, +) -> None: + unreadable_log = tmp_path / "events-1-unreadable.log" + valid_log = tmp_path / "events-2-valid.log" + unreadable_log.write_bytes(b"\xff") + _write_log(valid_log, [_record("root", "_input_event")]) + + result = _run_reader(tmp_path, "json") + trace_forest = json.loads(result.stdout) + + assert trace_forest["roots"] == ["root"] + assert trace_forest["warnings"] == [ + { + "code": "UNREADABLE_FILE", + "message": f"Could not read Event Log file {unreadable_log}.", + "filePath": str(unreadable_log), + } + ] + assert "[UNREADABLE_FILE]" in result.stderr + + +def test_reader_models_reused_event_as_dag_and_deduplicates_edges( + tmp_path: Path, +) -> None: + log_path = tmp_path / "events.log" + reused_record = _record( + "shared-output", + "SharedEvent", + "root-1", + "shared_action", + {"value": 1}, + ) + _write_log( + log_path, + [ + _record("root-1", "_input_event"), + reused_record, + reused_record, + _record("root-2", "_input_event"), + _record( + "shared-output", + "SharedEvent", + "root-2", + "shared_action", + {"value": 1}, + ), + ], + ) + + json_result = _run_reader(log_path, "json") + trace_forest = json.loads(json_result.stdout) + text_result = _run_reader(log_path, "text") + + assert trace_forest["roots"] == ["root-1", "root-2"] + assert trace_forest["nodes"]["shared-output"]["observationCount"] == 3 + assert trace_forest["nodes"]["shared-output"]["upstreamEdges"] == [ + { + "upstreamEventId": "root-1", + "upstreamActionName": "shared_action", + }, + { + "upstreamEventId": "root-2", + "upstreamActionName": "shared_action", + }, + ] + assert trace_forest["nodes"]["root-1"]["actions"] == [ + {"name": "shared_action", "children": ["shared-output"]} + ] + assert trace_forest["nodes"]["root-2"]["actions"] == [ + {"name": "shared_action", "children": ["shared-output"]} + ] + assert trace_forest["warnings"] == [] + assert text_result.stdout.count("SharedEvent (shared-output)") == 2 + assert json_result.stderr == "" + assert text_result.stderr == "" + + +def test_text_reader_handles_trace_deeper_than_recursion_limit(tmp_path: Path) -> None: + log_path = tmp_path / "events.log" + depth = sys.getrecursionlimit() + 10 + records = [_record("event-0", "_input_event")] + records.extend( + _record( + f"event-{index}", + "MiddleEvent", + f"event-{index - 1}", + f"action-{index}", + ) + for index in range(1, depth + 1) + ) + _write_log(log_path, records) + + result = _run_reader(log_path, "text") + + assert f"MiddleEvent (event-{depth})" in result.stdout + assert result.stderr == "" + + +def test_reader_compares_event_content_by_json_type(tmp_path: Path) -> None: + log_path = tmp_path / "events.log" + _write_log( + log_path, + [ + _record("root", "_input_event"), + _record( + "type-sensitive", + "ChildEvent", + "root", + "child_action", + {"value": True}, + ), + _record( + "type-sensitive", + "ChildEvent", + "root", + "child_action", + {"value": 1}, + ), + _record( + "equivalent", + "ChildEvent", + "root", + "child_action", + {"nested": {"first": 1, "second": 2}}, + ), + _record( + "equivalent", + "ChildEvent", + "root", + "child_action", + {"nested": {"second": 2, "first": 1}}, + ), + ], + ) + + result = _run_reader(log_path, "json") + trace_forest = json.loads(result.stdout) + + assert "type-sensitive" not in trace_forest["nodes"] + assert trace_forest["nodes"]["equivalent"]["observationCount"] == 2 + assert [item["code"] for item in trace_forest["warnings"]] == ["EVENT_ID_CONFLICT"] + + +def test_reader_warns_and_keeps_valid_input_tree(tmp_path: Path) -> None: + log_path = tmp_path / "events.log" + _write_log( + log_path, + [ + _record("root", "_input_event"), + _record("valid-child", "ChildEvent", "root", "valid_action"), + _record("unlinked", "UnlinkedEvent"), + _record("missing", "MissingEvent", "absent", "missing_action"), + _record( + "type-conflict", + "FirstType", + "root", + "conflicting_action", + {"value": 1}, + ), + _record( + "type-conflict", + "SecondType", + "root", + "conflicting_action", + {"value": 1}, + ), + _record( + "content-conflict", + "ConflictingEvent", + "root", + "conflicting_action", + {"value": 1}, + ), + _record( + "content-conflict", + "ConflictingEvent", + "root", + "conflicting_action", + {"value": 2}, + ), + ], + ) + + result = _run_reader(log_path, "json") + trace_forest = json.loads(result.stdout) + warning_codes = {warning["code"] for warning in trace_forest["warnings"]} + + assert trace_forest["roots"] == ["root"] + assert trace_forest["nodes"]["root"]["actions"] == [ + {"name": "valid_action", "children": ["valid-child"]} + ] + assert "unlinked" in trace_forest["nodes"] + assert "missing" in trace_forest["nodes"] + assert "type-conflict" not in trace_forest["nodes"] + assert "content-conflict" not in trace_forest["nodes"] + assert warning_codes == { + "EVENT_ID_CONFLICT", + "MISSING_PARENT", + "UNLINKED_EVENT", + } + assert result.stderr.count("EVENT_ID_CONFLICT") == 2 + assert "MISSING_PARENT" in result.stderr + assert "UNLINKED_EVENT" in result.stderr diff --git a/python/flink_agents/cli/trace_tree.py b/python/flink_agents/cli/trace_tree.py new file mode 100644 index 000000000..b8227266f --- /dev/null +++ b/python/flink_agents/cli/trace_tree.py @@ -0,0 +1,333 @@ +#!/usr/bin/env python3 +# 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. +"""Reconstruct InputEvent-rooted Trace Trees from an Event Log.""" + +import argparse +import json +import sys +from collections import defaultdict +from pathlib import Path +from typing import Any, Iterator + +INPUT_EVENT_TYPE = "_input_event" + + +def _json_fingerprint(value: Any) -> str: + """Return a deterministic, type-sensitive representation of JSON data.""" + return json.dumps(value, ensure_ascii=False, separators=(",", ":"), sort_keys=True) + + +def find_log_files(path: Path) -> list[Path]: + """Return Event Log files in deterministic file-name order.""" + if path.is_file(): + return [path] + return sorted(path.glob("events-*.log")) + + +def read_json_objects(path: Path, warnings: list[dict[str, Any]]) -> Iterator[Any]: + """Read consecutive JSON objects while retaining recoverable file warnings.""" + try: + content = path.read_text(encoding="utf-8") + except (OSError, UnicodeError): + warnings.append( + warning( + "UNREADABLE_FILE", + None, + f"Could not read Event Log file {path}.", + file_path=path, + ) + ) + return + + decoder = json.JSONDecoder() + position = 0 + while position < len(content): + while position < len(content) and content[position].isspace(): + position += 1 + if position == len(content): + return + try: + record, position = decoder.raw_decode(content, position) + except json.JSONDecodeError as error: + warnings.append( + warning( + "MALFORMED_RECORD", + None, + f"Could not decode an Event Log record in {path} " + f"at line {error.lineno}, column {error.colno}: {error.msg}.", + file_path=path, + line_number=error.lineno, + column_number=error.colno, + ) + ) + next_record = content.find("\n{", max(error.pos, position + 1)) + if next_record < 0: + return + position = next_record + 1 + continue + + yield record + + +def read_event_records( + path: Path, warnings: list[dict[str, Any]] +) -> Iterator[dict[str, Any]]: + """Read the fields needed to reconstruct Trace Trees.""" + log_files = find_log_files(path) + if not log_files: + message = f"No Event Log files found at {path}" + raise FileNotFoundError(message) + + for log_file in log_files: + for record in read_json_objects(log_file, warnings): + event_id: str | None = None + invalid_reason: str | None = None + if not isinstance(record, dict): + invalid_reason = "record must be a JSON object" + else: + event = record.get("event") + event_type = record.get("eventType") + if not isinstance(event, dict): + invalid_reason = "field 'event' must be a JSON object" + elif not isinstance(event.get("id"), str) or not event["id"]: + invalid_reason = "field 'event.id' must be a non-empty string" + elif not isinstance(event_type, str) or not event_type: + invalid_reason = "field 'eventType' must be a non-empty string" + else: + event_id = event["id"] + for field_name in ("upstreamEventId", "upstreamActionName"): + field_value = event.get(field_name) + if field_value is not None and not isinstance(field_value, str): + invalid_reason = ( + f"field 'event.{field_name}' must be a string or null" + ) + break + + if invalid_reason is not None: + warnings.append( + warning( + "MALFORMED_RECORD", + event_id, + f"Invalid Event Log record in {log_file}: {invalid_reason}.", + file_path=log_file, + ) + ) + continue + + assert isinstance(record, dict) + event = record["event"] + assert isinstance(event, dict) + event_content = dict(event) + event_content.pop("upstreamEventId", None) + event_content.pop("upstreamActionName", None) + yield { + "eventId": event["id"], + "eventType": record["eventType"], + "timestamp": record.get("timestamp"), + "upstreamEventId": event.get("upstreamEventId"), + "upstreamActionName": event.get("upstreamActionName"), + "eventContent": event_content, + } + + +def build_trace_forest( + records: Iterator[dict[str, Any]], + warnings: list[dict[str, Any]] | None = None, +) -> dict[str, Any]: + """Build valid InputEvent trees while retaining auditable invalid nodes.""" + records_by_id: dict[str, list[dict[str, Any]]] = defaultdict(list) + for record in records: + records_by_id[record["eventId"]].append(record) + + if warnings is None: + warnings = [] + nodes: dict[str, dict[str, Any]] = {} + lineage_edges_by_id: dict[str, list[tuple[str | None, str | None]]] = {} + for event_id, matching_records in records_by_id.items(): + first_record = matching_records[0] + first_content_fingerprint = _json_fingerprint(first_record["eventContent"]) + if any( + record["eventType"] != first_record["eventType"] + or _json_fingerprint(record["eventContent"]) != first_content_fingerprint + for record in matching_records[1:] + ): + warnings.append( + warning( + "EVENT_ID_CONFLICT", + event_id, + f"Event ID {event_id} has inconsistent Event type or content " + f"across {len(matching_records)} records.", + ) + ) + continue + + lineage_edges: list[tuple[str | None, str | None]] = [] + seen_edges: set[tuple[str | None, str | None]] = set() + for record in matching_records: + edge = (record["upstreamEventId"], record["upstreamActionName"]) + if edge not in seen_edges: + seen_edges.add(edge) + lineage_edges.append(edge) + lineage_edges_by_id[event_id] = lineage_edges + + nodes[event_id] = { + "eventId": event_id, + "eventType": first_record["eventType"], + "timestamp": first_record["timestamp"], + "observationCount": len(matching_records), + "upstreamEdges": [ + { + "upstreamEventId": upstream_event_id, + "upstreamActionName": upstream_action_name, + } + for upstream_event_id, upstream_action_name in lineage_edges + if upstream_event_id is not None or upstream_action_name is not None + ], + "actions": [], + } + + roots: list[str] = [] + action_indexes: dict[str, dict[str, dict[str, Any]]] = defaultdict(dict) + for event_id, node in nodes.items(): + lineage_edges = lineage_edges_by_id[event_id] + + if node["eventType"] == INPUT_EVENT_TYPE: + if (None, None) in lineage_edges: + roots.append(event_id) + for upstream_event_id, upstream_action_name in lineage_edges: + if upstream_event_id is None and upstream_action_name is None: + continue + warnings.append( + warning( + "INVALID_ROOT_LINEAGE", + event_id, + f"InputEvent {event_id} must not have upstream lineage.", + ) + ) + continue + + for upstream_event_id, upstream_action_name in lineage_edges: + if upstream_event_id is None: + warnings.append( + warning( + "UNLINKED_EVENT", + event_id, + f"Non-InputEvent {event_id} has no upstream Event.", + ) + ) + continue + if upstream_action_name is None: + warnings.append( + warning( + "MISSING_ACTION_NAME", + event_id, + f"Event {event_id} has no upstream Action name.", + ) + ) + continue + if upstream_event_id not in nodes: + warnings.append( + warning( + "MISSING_PARENT", + event_id, + f"Event {event_id} references missing parent {upstream_event_id}.", + ) + ) + continue + + action = action_indexes[upstream_event_id].get(upstream_action_name) + if action is None: + action = {"name": upstream_action_name, "children": []} + action_indexes[upstream_event_id][upstream_action_name] = action + nodes[upstream_event_id]["actions"].append(action) + action["children"].append(event_id) + + return {"roots": roots, "nodes": nodes, "warnings": warnings} + + +def warning( + code: str, + event_id: str | None, + message: str, + *, + file_path: Path | None = None, + line_number: int | None = None, + column_number: int | None = None, +) -> dict[str, Any]: + """Create one machine-readable reconstruction warning.""" + item: dict[str, Any] = {"code": code} + if event_id is not None: + item["eventId"] = event_id + item["message"] = message + if file_path is not None: + item["filePath"] = str(file_path) + if line_number is not None: + item["lineNumber"] = line_number + if column_number is not None: + item["columnNumber"] = column_number + return item + + +def render_text(trace_forest: dict[str, Any]) -> str: + """Render valid Trace Trees without assigning meaning to sibling order.""" + lines: list[str] = [] + + def render_event(event_id: str, indent: str) -> None: + stack: list[tuple[str, Any, str]] = [("event", event_id, indent)] + while stack: + item_type, item, item_indent = stack.pop() + if item_type == "event": + node = trace_forest["nodes"][item] + lines.append(f"{item_indent}{node['eventType']} ({item})") + stack.extend( + ("action", action, item_indent) + for action in reversed(node["actions"]) + ) + else: + lines.append(f"{item_indent} [Action: {item['name']}]") + stack.extend( + ("event", child_id, item_indent + " ") + for child_id in reversed(item["children"]) + ) + + for root_number, root_id in enumerate(trace_forest["roots"], start=1): + if lines: + lines.append("") + lines.append(f"Trace Tree {root_number}") + render_event(root_id, " ") + return "\n".join(lines) + + +def main() -> None: + """Run the command-line reader.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("path", type=Path, help="Event Log file or directory") + parser.add_argument("--format", choices=("text", "json"), default="text") + args = parser.parse_args() + + warnings: list[dict[str, Any]] = [] + trace_forest = build_trace_forest(read_event_records(args.path, warnings), warnings) + for item in trace_forest["warnings"]: + print(f"[{item['code']}] {item['message']}", file=sys.stderr) + + if args.format == "json": + print(json.dumps(trace_forest, indent=2)) + else: + print(render_text(trace_forest)) + + +if __name__ == "__main__": + main() diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/flink_integration_agent.py b/python/flink_agents/e2e_tests/e2e_tests_integration/flink_integration_agent.py index b464754f5..a90145fdb 100644 --- a/python/flink_agents/e2e_tests/e2e_tests_integration/flink_integration_agent.py +++ b/python/flink_agents/e2e_tests/e2e_tests_integration/flink_integration_agent.py @@ -65,7 +65,8 @@ def __init__(self, value: Any) -> None: def from_event(cls, event: "Event") -> "MyEvent": """Reconstruct a MyEvent from a generic Event.""" assert "value" in event.attributes, "Missing 'value' in event attributes" - return cls(value=event.attributes["value"]) + result = cls(value=event.attributes["value"]) + return result.reconstruct_from(event) @property def value(self) -> Any: diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/long_term_memory_test.py b/python/flink_agents/e2e_tests/e2e_tests_integration/long_term_memory_test.py index 3db2ae010..b2d7bfe4b 100644 --- a/python/flink_agents/e2e_tests/e2e_tests_integration/long_term_memory_test.py +++ b/python/flink_agents/e2e_tests/e2e_tests_integration/long_term_memory_test.py @@ -120,7 +120,8 @@ def __init__(self, value: Any) -> None: def from_event(cls, event: Event) -> "MyEvent": """Reconstruct a MyEvent from a generic Event.""" assert "value" in event.attributes - return MyEvent(value=Record.model_validate(event.attributes["value"])) + result = cls(value=Record.model_validate(event.attributes["value"])) + return result.reconstruct_from(event) @property def value(self) -> Any: diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/python_event_logging_test.py b/python/flink_agents/e2e_tests/e2e_tests_integration/python_event_logging_test.py index 2aade7d80..4b74a85bd 100644 --- a/python/flink_agents/e2e_tests/e2e_tests_integration/python_event_logging_test.py +++ b/python/flink_agents/e2e_tests/e2e_tests_integration/python_event_logging_test.py @@ -18,6 +18,8 @@ import json import os import shutil +import subprocess +import sys import sysconfig import tempfile from pathlib import Path @@ -206,6 +208,87 @@ def _read_log_records(event_log_dir: Path) -> list[dict]: return records +def _run_trace_tree_reader( + event_log_dir: Path, output_format: str +) -> subprocess.CompletedProcess[str]: + """Run the public Trace Tree reader CLI against Runtime Event Logs.""" + result = subprocess.run( + [ + sys.executable, + "-m", + "flink_agents.cli.trace_tree", + str(event_log_dir), + "--format", + output_format, + ], + check=False, + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 0, result.stderr + return result + + +def test_event_lineage_reconstructs_trace_trees_from_runtime_logs( + tmp_path: Path, +) -> None: + """Test a real PyFlink job produces logs consumable by the Trace Tree reader.""" + event_log_dir = _run_event_logging_pipeline(tmp_path) + records = _read_log_records(event_log_dir) + input_event_ids = { + record["event"]["id"] + for record in records + if record["eventType"] == InputEvent.EVENT_TYPE + } + output_event_ids = { + record["event"]["id"] + for record in records + if record["eventType"] == OutputEvent.EVENT_TYPE + } + + json_result = _run_trace_tree_reader(event_log_dir, "json") + trace_forest = json.loads(json_result.stdout) + text_result = _run_trace_tree_reader(event_log_dir, "text") + + assert input_event_ids + assert len(output_event_ids) == len(input_event_ids) + assert set(trace_forest["roots"]) == input_event_ids + assert set(trace_forest["nodes"]) == input_event_ids | output_event_ids + assert trace_forest["warnings"] == [] + assert json_result.stderr == "" + + for root_id in trace_forest["roots"]: + root = trace_forest["nodes"][root_id] + assert root["eventType"] == InputEvent.EVENT_TYPE + assert root["upstreamEdges"] == [] + assert root["observationCount"] == 1 + assert len(root["actions"]) == 1 + + action_node = root["actions"][0] + assert action_node["name"] == "process_input" + assert len(action_node["children"]) == 1 + + child_id = action_node["children"][0] + child = trace_forest["nodes"][child_id] + assert child_id in output_event_ids + assert child["eventType"] == OutputEvent.EVENT_TYPE + assert child["upstreamEdges"] == [ + { + "upstreamEventId": root_id, + "upstreamActionName": "process_input", + } + ] + assert child["observationCount"] == 1 + assert child["actions"] == [] + + assert text_result.stdout.count("Trace Tree ") == len(input_event_ids) + assert text_result.stdout.count("[Action: process_input]") == len(input_event_ids) + for event_id in input_event_ids | output_event_ids: + assert f"({event_id})" in text_result.stdout + assert text_result.stderr == "" + + def test_event_log_verbose_level(tmp_path: Path) -> None: """Test that VERBOSE log level writes events without truncation.""" event_log_dir = _run_event_logging_pipeline( diff --git a/python/flink_agents/runtime/tests/test_python_java_utils.py b/python/flink_agents/runtime/tests/test_python_java_utils.py index fa708353b..3d6b0decb 100644 --- a/python/flink_agents/runtime/tests/test_python_java_utils.py +++ b/python/flink_agents/runtime/tests/test_python_java_utils.py @@ -17,15 +17,19 @@ ################################################################################# import json +import cloudpickle + from flink_agents.api.decorators import tool from flink_agents.api.embedding_models.embedding_model import ( EmbeddingResult, EmbeddingTokenUsage, ) +from flink_agents.api.events.event import Event from flink_agents.api.tools import InjectedArg from flink_agents.runtime.python_java_utils import ( call_embedding_with_usage, get_python_tool_metadata, + wrap_to_input_event, ) @@ -67,3 +71,14 @@ def test_call_embedding_with_usage_returns_pemja_safe_primitives() -> None: "embeddings": [0.1, 0.2], "token_usage": {"prompt_tokens": 7, "total_tokens": 9}, } + + +def test_wrap_to_input_event_assigns_unique_random_id_per_record() -> None: + payload = cloudpickle.dumps("same input") + + first = Event.from_json(wrap_to_input_event(payload)) + second = Event.from_json(wrap_to_input_event(payload)) + + assert first.id != second.id + assert first.id.version == 4 + assert second.id.version == 4 diff --git a/python/pyproject.toml b/python/pyproject.toml index 2d2d8d41c..0f7005d48 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -61,6 +61,9 @@ dependencies = [ "onnxruntime<1.24.1;python_version<'3.11'", ] +[project.scripts] +flink-agents-trace-tree = "flink_agents.cli.trace_tree:main" + [tool.setuptools] include-package-data = true diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/JsonTruncator.java b/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/JsonTruncator.java index 89e49ceeb..fa89a95aa 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/JsonTruncator.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/JsonTruncator.java @@ -24,8 +24,6 @@ import com.fasterxml.jackson.databind.node.ObjectNode; import java.util.ArrayList; -import java.util.Arrays; -import java.util.HashSet; import java.util.List; import java.util.Set; @@ -46,14 +44,21 @@ *

Setting any threshold to {@code 0} disables that specific truncation strategy. If all * thresholds are {@code 0}, no truncation occurs. * - *

Protected fields at the top level of the event node ({@code eventType}, {@code id}, {@code - * attributes}) are never truncated as structural fields. The {@code attributes} envelope is - * additionally traversed so user payload stored inside it is truncated like any other field. + *

Protected fields at the top level of the event node ({@code eventType}, {@code type}, {@code + * id}, {@code upstreamEventId}, {@code upstreamActionName}, {@code attributes}) are never truncated + * as structural fields. The {@code attributes} envelope is additionally traversed so user payload + * stored inside it is truncated like any other field. */ public class JsonTruncator { private static final Set PROTECTED_FIELDS = - new HashSet<>(Arrays.asList("eventType", "id", "attributes")); + Set.of( + "eventType", + "type", + "id", + "upstreamEventId", + "upstreamActionName", + "attributes"); private final int maxStringLength; private final int maxArrayElements; @@ -75,9 +80,10 @@ public JsonTruncator(int maxStringLength, int maxArrayElements, int maxDepth) { /** * Truncates the given event node in place according to configured thresholds. * - *

Protected fields ({@code eventType}, {@code id}, {@code attributes}) at the top level of - * the event node are never truncated as structural fields. The {@code attributes} envelope is - * additionally traversed so user payload stored inside it is truncated like any other field. + *

Protected fields ({@code eventType}, {@code type}, {@code id}, {@code upstreamEventId}, + * {@code upstreamActionName}, {@code attributes}) at the top level of the event node are never + * truncated as structural fields. The {@code attributes} envelope is additionally traversed so + * user payload stored inside it is truncated like any other field. * * @param eventNode the top-level event JSON object to truncate * @return {@code true} if any field was truncated, {@code false} if the node was unchanged 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 46d09e0a2..c6eb65167 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 @@ -345,7 +345,7 @@ private void processActionTaskForKey(Object key) throws Exception { actionTask.action.getName(), key); isFinished = true; - outputEvents = actionState.getOutputEvents(); + outputEvents = actionTask.finalizeOutputEvents(actionState.getOutputEvents()); for (MemoryUpdate memoryUpdate : actionState.getShortTermMemoryUpdates()) { actionTask .getRunnerContext() @@ -385,6 +385,7 @@ private void processActionTaskForKey(Object key) throws Exception { durableExecManager.removeDurableContext(actionTask); contextManager.removeContinuationContext(actionTask); contextManager.removePythonAwaitableRef(actionTask); + outputEvents = actionTaskResult.getOutputEvents(); durableExecManager.maybePersistTaskResult( key, sequenceNumber, @@ -393,7 +394,6 @@ private void processActionTaskForKey(Object key) throws Exception { actionTask.getRunnerContext(), actionTaskResult); isFinished = actionTaskResult.isFinished(); - outputEvents = actionTaskResult.getOutputEvents(); generatedActionTaskOpt = actionTaskResult.getGeneratedActionTask(); } diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTask.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTask.java index 053b9d955..36d92ea43 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTask.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTask.java @@ -91,6 +91,30 @@ public int hashCode() { public abstract ActionTaskResult invoke( ClassLoader userCodeClassLoader, PythonActionExecutor executor) throws Exception; + /** + * Validates and binds output Events to this task's Action and trigger Event. + * + *

All outputs are validated before mutation to avoid partial updates. Existing lineage is + * overwritten, including during replay. + */ + List finalizeOutputEvents(List outputEvents) { + for (Event outputEvent : outputEvents) { + if (Objects.equals(outputEvent.getId(), event.getId())) { + throw new IllegalArgumentException( + "Action '" + + action.getName() + + "' cannot emit its triggering Event " + + event.getId() + + "; output Event IDs must differ from the triggering Event ID."); + } + } + for (Event outputEvent : outputEvents) { + outputEvent.setUpstreamEventId(event.getId()); + outputEvent.setUpstreamActionName(action.getName()); + } + return outputEvents; + } + public class ActionTaskResult { private final boolean finished; private final List outputEvents; @@ -101,7 +125,7 @@ public ActionTaskResult( List outputEvents, @Nullable ActionTask generatedActionTask) { this.finished = finished; - this.outputEvents = outputEvents; + this.outputEvents = finalizeOutputEvents(outputEvents); this.generatedActionTaskOpt = Optional.ofNullable(generatedActionTask); } diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/EventLogRecordJsonSerdeTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/EventLogRecordJsonSerdeTest.java index 1f7a1192f..0af92189e 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/EventLogRecordJsonSerdeTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/EventLogRecordJsonSerdeTest.java @@ -28,7 +28,6 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; -import java.util.HashMap; import java.util.Map; import java.util.UUID; @@ -161,6 +160,9 @@ void testDeserializeOutputEvent() throws Exception { void testDeserializeCustomEvent() throws Exception { // Given CustomTestEvent originalEvent = new CustomTestEvent("custom data", 42, true); + UUID upstreamEventId = UUID.randomUUID(); + originalEvent.setUpstreamEventId(upstreamEventId); + originalEvent.setUpstreamActionName("custom_action"); EventContext originalContext = new EventContext(originalEvent); EventLogRecord originalRecord = new EventLogRecord(originalContext, originalEvent); String json = objectMapper.writeValueAsString(originalRecord); @@ -175,6 +177,9 @@ void testDeserializeCustomEvent() throws Exception { assertEquals("custom data", customEvent.getCustomData()); assertEquals(42, customEvent.getCustomNumber()); assertTrue(customEvent.isCustomFlag()); + assertEquals(originalEvent.getId(), customEvent.getId()); + assertEquals(upstreamEventId, customEvent.getUpstreamEventId()); + assertEquals("custom_action", customEvent.getUpstreamActionName()); } @Test @@ -197,6 +202,25 @@ void testRoundTripSerialization() throws Exception { assertEquals(originalEvent.getInput(), deserializedEvent.getInput()); } + @Test + void testRoundTripLineageFields() throws Exception { + UUID upstreamEventId = UUID.randomUUID(); + Event originalEvent = new Event("ChildEvent"); + originalEvent.setUpstreamEventId(upstreamEventId); + originalEvent.setUpstreamActionName("child_action"); + EventLogRecord originalRecord = + new EventLogRecord(new EventContext(originalEvent), originalEvent); + + String json = objectMapper.writeValueAsString(originalRecord); + EventLogRecord deserializedRecord = objectMapper.readValue(json, EventLogRecord.class); + + JsonNode eventNode = objectMapper.readTree(json).get("event"); + assertEquals(upstreamEventId.toString(), eventNode.get("upstreamEventId").asText()); + assertEquals("child_action", eventNode.get("upstreamActionName").asText()); + assertEquals(upstreamEventId, deserializedRecord.getEvent().getUpstreamEventId()); + assertEquals("child_action", deserializedRecord.getEvent().getUpstreamActionName()); + } + @Test void testSerializeUnifiedEvent() throws Exception { // Given - a unified event with user-defined type @@ -279,12 +303,7 @@ private CustomTestEvent(UUID id, Map attributes) { } public static CustomTestEvent fromEvent(Event event) { - CustomTestEvent result = - new CustomTestEvent(event.getId(), new HashMap<>(event.getAttributes())); - if (event.hasSourceTimestamp()) { - result.setSourceTimestamp(event.getSourceTimestamp()); - } - return result; + return reconstructFrom(event, CustomTestEvent::new); } @JsonIgnore diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/FileEventLoggerTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/FileEventLoggerTest.java index 9156aee72..05bbdbb44 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/FileEventLoggerTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/FileEventLoggerTest.java @@ -174,6 +174,9 @@ void testAppendWithCustomEvent() throws Exception { // Given logger.open(openParams); TestCustomEvent customEvent = new TestCustomEvent("custom data", 42); + UUID upstreamEventId = UUID.randomUUID(); + customEvent.setUpstreamEventId(upstreamEventId); + customEvent.setUpstreamActionName("custom_action"); EventContext context = new EventContext(customEvent); // When @@ -203,6 +206,9 @@ void testAppendWithCustomEvent() throws Exception { TestCustomEvent.fromEvent(deserializedRecord.getEvent()); assertEquals("custom data", deserializedEvent.getCustomData()); assertEquals(42, deserializedEvent.getCustomNumber()); + assertEquals(customEvent.getId(), deserializedEvent.getId()); + assertEquals(upstreamEventId, deserializedEvent.getUpstreamEventId()); + assertEquals("custom_action", deserializedEvent.getUpstreamActionName()); } @Test @@ -564,12 +570,7 @@ private TestCustomEvent(UUID id, Map attributes) { } public static TestCustomEvent fromEvent(Event event) { - TestCustomEvent result = - new TestCustomEvent(event.getId(), new HashMap<>(event.getAttributes())); - if (event.hasSourceTimestamp()) { - result.setSourceTimestamp(event.getSourceTimestamp()); - } - return result; + return reconstructFrom(event, TestCustomEvent::new); } @JsonIgnore diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/JsonTruncatorTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/JsonTruncatorTest.java index bd05b9a10..504b2d342 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/JsonTruncatorTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/JsonTruncatorTest.java @@ -138,9 +138,12 @@ void testProtectedFields() { ObjectNode node = MAPPER.createObjectNode(); // Protected fields should never be truncated node.put("eventType", "org.apache.flink.agents.api.event.ChatRequestEvent"); + node.put("type", "org.apache.flink.agents.api.event.ChatRequestEvent"); node.put("id", "a-very-long-identifier-that-exceeds-the-limit"); + node.put("upstreamEventId", "a-very-long-upstream-event-identifier"); + node.put("upstreamActionName", "a-very-long-action-name"); ObjectNode attributes = MAPPER.createObjectNode(); - attributes.put("key", "value"); + attributes.put("key", "business payload should be truncated"); attributes.set("nested", MAPPER.createObjectNode().put("deep", "data")); node.set("attributes", attributes); // Non-protected field should be truncated @@ -152,10 +155,16 @@ void testProtectedFields() { // Protected fields remain untouched assertThat(node.get("eventType").asText()) .isEqualTo("org.apache.flink.agents.api.event.ChatRequestEvent"); + assertThat(node.get("type").asText()) + .isEqualTo("org.apache.flink.agents.api.event.ChatRequestEvent"); assertThat(node.get("id").asText()) .isEqualTo("a-very-long-identifier-that-exceeds-the-limit"); + assertThat(node.get("upstreamEventId").asText()) + .isEqualTo("a-very-long-upstream-event-identifier"); + assertThat(node.get("upstreamActionName").asText()).isEqualTo("a-very-long-action-name"); assertThat(node.get("attributes").isObject()).isTrue(); - assertThat(node.get("attributes").get("key").asText()).isEqualTo("value"); + assertThat(node.get("attributes").get("key").get("truncatedString").asText()) + .isEqualTo("busin..."); // Non-protected field is truncated assertThat(node.get("content").get("truncatedString")).isNotNull(); } 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 22c0072f8..79a18931e 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 @@ -17,6 +17,7 @@ */ package org.apache.flink.agents.runtime.operator; +import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; import org.apache.flink.agents.api.Event; import org.apache.flink.agents.api.EventContext; @@ -34,6 +35,7 @@ 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.ActionStateSerde; import org.apache.flink.agents.runtime.actionstate.CallResult; import org.apache.flink.agents.runtime.actionstate.InMemoryActionStateStore; import org.apache.flink.agents.runtime.eventlog.FileEventLogger; @@ -50,11 +52,13 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import java.io.IOException; import java.lang.reflect.Field; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.UUID; import java.util.concurrent.Callable; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -376,6 +380,28 @@ public void onEventProcessed(EventContext context, Event event) { } } + private static final class EventSnapshot { + private final UUID id; + private final UUID upstreamEventId; + private final String upstreamActionName; + + private EventSnapshot(Event event) { + this.id = event.getId(); + this.upstreamEventId = event.getUpstreamEventId(); + this.upstreamActionName = event.getUpstreamActionName(); + } + } + + private static Map recordEventSnapshotsByType( + ActionExecutionOperator operator) { + Map snapshotsByType = new HashMap<>(); + operator.getEventRouter() + .addEventListener( + (context, event) -> + snapshotsByType.put(event.getType(), new EventSnapshot(event))); + return snapshotsByType; + } + @Test void testEventListenersFromAgentConfig() throws Exception { final AgentConfiguration config = new AgentConfiguration(); @@ -581,6 +607,36 @@ agentPlanWithStateStore, true, new InMemoryActionStateStore(false)), } } + @Test + void testCompletedActionStatePersistsOutputEventLineage() throws Exception { + AgentPlan agentPlan = TestAgent.getAgentPlan(false); + SerializingActionStateStore actionStateStore = new SerializingActionStateStore(); + + try (KeyedOneInputStreamOperatorTestHarness testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory<>(agentPlan, true, actionStateStore), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) testHarness.getOperator(); + + testHarness.processElement(new StreamRecord<>(3L)); + operator.waitInFlightEventsFinished(); + } + + assertThat(actionStateStore.getCompletedStateBytes()).hasSize(2); + for (Map.Entry entry : + actionStateStore.getCompletedStateBytes().entrySet()) { + JsonNode state = OBJECT_MAPPER.readTree(entry.getValue()); + JsonNode outputEvent = state.path("outputEvents").get(0); + + assertThat(outputEvent.path("upstreamEventId").asText()) + .isEqualTo(state.path("taskEvent").path("id").asText()); + assertThat(outputEvent.path("upstreamActionName").asText()).isEqualTo(entry.getKey()); + } + } + @Test void testActionStateStoreStateManagement() throws Exception { AgentPlan agentPlanWithStateStore = TestAgent.getAgentPlan(false); @@ -725,54 +781,113 @@ void testEarlierCheckpointReplayKeepsDurableState() throws Exception { } @Test - void testActionStateStoreReplayIncurNoFunctionCall() throws Exception { - AgentPlan agentPlanWithStateStore = TestAgent.getAgentPlan(false); - InMemoryActionStateStore actionStateStore; + void testReplaySkipsCompletedActions() throws Exception { + AgentPlan agentPlan = TestAgent.getAgentPlan(false); + long inputValue = 7L; + InMemoryActionStateStore actionStateStore = new InMemoryActionStateStore(false); + TestAgent.ACTION1_CALL_COUNTER.set(0); + TestAgent.ACTION2_CALL_COUNTER.set(0); + try (KeyedOneInputStreamOperatorTestHarness testHarness = new KeyedOneInputStreamOperatorTestHarness<>( - new ActionExecutionOperatorFactory<>( - agentPlanWithStateStore, true, new InMemoryActionStateStore(false)), + new ActionExecutionOperatorFactory<>(agentPlan, true, actionStateStore), (KeySelector) value -> value, TypeInformation.of(Long.class))) { testHarness.open(); ActionExecutionOperator operator = (ActionExecutionOperator) testHarness.getOperator(); - actionStateStore = - (InMemoryActionStateStore) - operator.getDurableExecutionManager().getActionStateStore(); + testHarness.processElement(new StreamRecord<>(inputValue)); + operator.waitInFlightEventsFinished(); - Long inputValue = 7L; + assertThat(TestAgent.ACTION1_CALL_COUNTER.get()).isEqualTo(1); + assertThat(TestAgent.ACTION2_CALL_COUNTER.get()).isEqualTo(1); + } + + try (KeyedOneInputStreamOperatorTestHarness testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory<>(agentPlan, true, actionStateStore), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) testHarness.getOperator(); - // First processing - this will execute the actual functions and store state testHarness.processElement(new StreamRecord<>(inputValue)); operator.waitInFlightEventsFinished(); + + List> outputRecords = + (List>) testHarness.getRecordOutput(); + assertThat(outputRecords).hasSize(1); + assertThat(outputRecords.get(0).getValue()).isEqualTo((inputValue + 1) * 2); + assertThat(TestAgent.ACTION1_CALL_COUNTER.get()) + .as("Completed action1 must not be re-executed during replay") + .isEqualTo(1); + assertThat(TestAgent.ACTION2_CALL_COUNTER.get()) + .as("Completed action2 must not be re-executed during replay") + .isEqualTo(1); } + } + + @Test + void testReplayRebindsOutputLineage() throws Exception { + AgentPlan agentPlan = TestAgent.getAgentPlan(false); + long inputValue = 7L; + InMemoryActionStateStore actionStateStore = new InMemoryActionStateStore(false); + Map firstSnapshots; try (KeyedOneInputStreamOperatorTestHarness testHarness = new KeyedOneInputStreamOperatorTestHarness<>( - new ActionExecutionOperatorFactory<>( - agentPlanWithStateStore, true, actionStateStore), + new ActionExecutionOperatorFactory<>(agentPlan, true, actionStateStore), (KeySelector) value -> value, TypeInformation.of(Long.class))) { testHarness.open(); ActionExecutionOperator operator = (ActionExecutionOperator) testHarness.getOperator(); - Long inputValue = 7L; + firstSnapshots = recordEventSnapshotsByType(operator); - // First processing - this will execute the actual functions and store state + // Execute the actions and persist their completed states. testHarness.processElement(new StreamRecord<>(inputValue)); operator.waitInFlightEventsFinished(); - // Verify first output is correct - List> recordOutput = - (List>) testHarness.getRecordOutput(); - assertThat(recordOutput.size()).isEqualTo(1); - assertThat(recordOutput.get(0).getValue()).isEqualTo((inputValue + 1) * 2); + } - // The action state store should only have one entry - assertThat(actionStateStore.getKeyedActionStates().get(String.valueOf(inputValue))) - .hasSize(2); + Map replaySnapshots; + try (KeyedOneInputStreamOperatorTestHarness testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory<>(agentPlan, true, actionStateStore), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) testHarness.getOperator(); + replaySnapshots = recordEventSnapshotsByType(operator); + + // Replay the same input and reuse the completed action states. + testHarness.processElement(new StreamRecord<>(inputValue)); + operator.waitInFlightEventsFinished(); } + + EventSnapshot firstInput = firstSnapshots.get(InputEvent.EVENT_TYPE); + EventSnapshot replayInput = replaySnapshots.get(InputEvent.EVENT_TYPE); + EventSnapshot firstMiddle = firstSnapshots.get(TestAgent.MiddleEvent.EVENT_TYPE); + EventSnapshot replayMiddle = replaySnapshots.get(TestAgent.MiddleEvent.EVENT_TYPE); + EventSnapshot firstOutput = firstSnapshots.get(OutputEvent.EVENT_TYPE); + EventSnapshot replayOutput = replaySnapshots.get(OutputEvent.EVENT_TYPE); + + assertThat(firstInput.id).isNotEqualTo(replayInput.id); + + assertThat(replayMiddle.id).isEqualTo(firstMiddle.id); + assertThat(firstMiddle.upstreamEventId).isEqualTo(firstInput.id); + assertThat(replayMiddle.upstreamEventId).isEqualTo(replayInput.id); + assertThat(firstMiddle.upstreamEventId).isNotEqualTo(replayMiddle.upstreamEventId); + assertThat(firstMiddle.upstreamActionName).isEqualTo("action1"); + assertThat(replayMiddle.upstreamActionName).isEqualTo("action1"); + + assertThat(replayOutput.id).isEqualTo(firstOutput.id); + assertThat(firstOutput.upstreamEventId).isEqualTo(firstMiddle.id); + assertThat(replayOutput.upstreamEventId).isEqualTo(replayMiddle.id); + assertThat(firstOutput.upstreamActionName).isEqualTo("action2"); + assertThat(replayOutput.upstreamActionName).isEqualTo("action2"); } @Test @@ -1439,6 +1554,13 @@ public static class TestAgent { public static final java.util.concurrent.atomic.AtomicInteger DURABLE_CALL_COUNTER = new java.util.concurrent.atomic.AtomicInteger(0); + /** Counters used to verify that completed Actions are not re-executed during replay. */ + public static final java.util.concurrent.atomic.AtomicInteger ACTION1_CALL_COUNTER = + new java.util.concurrent.atomic.AtomicInteger(0); + + public static final java.util.concurrent.atomic.AtomicInteger ACTION2_CALL_COUNTER = + new java.util.concurrent.atomic.AtomicInteger(0); + public static class MiddleEvent extends Event { public static final String EVENT_TYPE = "MiddleEvent"; @@ -1455,6 +1577,7 @@ public Long getNum() { } public static void action1(Event event, RunnerContext context) { + ACTION1_CALL_COUNTER.incrementAndGet(); Long inputData = (Long) InputEvent.fromEvent(event).getInput(); try { MemoryObject mem = context.getShortTermMemory(); @@ -1466,6 +1589,7 @@ public static void action1(Event event, RunnerContext context) { } public static void action2(MiddleEvent event, RunnerContext context) { + ACTION2_CALL_COUNTER.incrementAndGet(); try { MemoryObject mem = context.getShortTermMemory(); Long tmp = (Long) mem.get("tmp").getValue(); @@ -2154,6 +2278,27 @@ private List getPrunedSeqNums() { } } + private static class SerializingActionStateStore extends InMemoryActionStateStore { + private final Map completedStateBytes = new HashMap<>(); + + private SerializingActionStateStore() { + super(false); + } + + @Override + public void put(Object key, long seqNum, Action action, Event event, ActionState state) + throws IOException { + if (state.isCompleted()) { + completedStateBytes.put(action.getName(), ActionStateSerde.serialize(state)); + } + super.put(key, seqNum, action, event, state); + } + + private Map getCompletedStateBytes() { + return completedStateBytes; + } + } + private static void assertMailboxSizeAndRun(TaskMailbox mailbox, int expectedSize) throws Exception { assertThat(mailbox.size()).isEqualTo(expectedSize); diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskTest.java new file mode 100644 index 000000000..998998ba6 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskTest.java @@ -0,0 +1,69 @@ +/* + * 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.plan.actions.Action; +import org.junit.jupiter.api.Test; + +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Tests for output finalization owned by {@link ActionTask}. */ +class ActionTaskTest { + + @Test + void resultFinalizesOutputLineage() { + Event triggeringEvent = new InputEvent(1L); + Action action = TestActions.noopAction(); + ActionTask task = new JavaActionTask("key", triggeringEvent, action); + Event outputEvent = new Event("result"); + + ActionTask.ActionTaskResult result = + task.new ActionTaskResult(true, List.of(outputEvent), null); + + assertThat(result.getOutputEvents()).containsExactly(outputEvent); + assertThat(outputEvent.getUpstreamEventId()).isEqualTo(triggeringEvent.getId()); + assertThat(outputEvent.getUpstreamActionName()).isEqualTo(action.getName()); + } + + @Test + void resultRejectsSelfLoopBeforeMutatingAnyOutput() { + Event triggeringEvent = new InputEvent(1L); + Action action = TestActions.noopAction(); + ActionTask task = new JavaActionTask("key", triggeringEvent, action); + Event validOutput = new Event("result"); + + assertThatThrownBy( + () -> + task + .new ActionTaskResult( + true, List.of(validOutput, triggeringEvent), null)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining(action.getName()) + .hasMessageContaining(triggeringEvent.getId().toString()); + + assertThat(validOutput.getUpstreamEventId()).isNull(); + assertThat(validOutput.getUpstreamActionName()).isNull(); + assertThat(triggeringEvent.getUpstreamEventId()).isNull(); + assertThat(triggeringEvent.getUpstreamActionName()).isNull(); + } +} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/EventLineageReconstructionTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/EventLineageReconstructionTest.java new file mode 100644 index 000000000..72e26fa44 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/EventLineageReconstructionTest.java @@ -0,0 +1,127 @@ +/* + * 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 com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.flink.agents.api.InputEvent; +import org.apache.flink.agents.api.OutputEvent; +import org.apache.flink.agents.api.configuration.AgentConfigOptions; +import org.apache.flink.agents.api.logger.EventLogLevel; +import org.apache.flink.agents.api.logger.LoggerType; +import org.apache.flink.agents.plan.AgentConfiguration; +import org.apache.flink.agents.plan.AgentPlan; +import org.apache.flink.agents.runtime.eventlog.EventLogRecord; +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import java.util.stream.Stream; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Verifies that the event log alone contains the minimal causal chain between Events. */ +class EventLineageReconstructionTest { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + @Test + void recordsMinimalEventLineageInTheEventLog(@TempDir Path logDir) throws Exception { + AgentPlan agentPlan = + ActionExecutionOperatorTest.TestAgent.getAgentPlanWithConfig( + fileLoggerConfig(logDir)); + + try (KeyedOneInputStreamOperatorTestHarness harness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory<>(agentPlan, true), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + harness.open(); + harness.processElement(new StreamRecord<>(1L)); + ((ActionExecutionOperator) harness.getOperator()) + .waitInFlightEventsFinished(); + } + + Map recordsByType = readRecordsByType(logDir); + EventLogRecord inputRecord = + MAPPER.treeToValue(recordsByType.get(InputEvent.EVENT_TYPE), EventLogRecord.class); + EventLogRecord outputRecord = + MAPPER.treeToValue(recordsByType.get(OutputEvent.EVENT_TYPE), EventLogRecord.class); + JsonNode input = recordsByType.get(InputEvent.EVENT_TYPE).get("event"); + JsonNode middle = + recordsByType + .get(ActionExecutionOperatorTest.TestAgent.MiddleEvent.EVENT_TYPE) + .get("event"); + JsonNode output = recordsByType.get(OutputEvent.EVENT_TYPE).get("event"); + + assertThat(recordsByType).hasSize(3); + assertThat(inputRecord.getEvent().getType()).isEqualTo(InputEvent.EVENT_TYPE); + assertThat(outputRecord.getEvent().getType()).isEqualTo(OutputEvent.EVENT_TYPE); + assertThat(input.has("upstreamEventId")).isFalse(); + assertThat(input.has("upstreamActionName")).isFalse(); + + assertThat(middle.get("upstreamEventId").isTextual()).isTrue(); + assertThat(middle.get("upstreamActionName").isTextual()).isTrue(); + assertThat(middle.get("upstreamEventId").asText()).isEqualTo(input.get("id").asText()); + assertThat(middle.get("upstreamActionName").asText()).isEqualTo("action1"); + + assertThat(output.get("upstreamEventId").isTextual()).isTrue(); + assertThat(output.get("upstreamActionName").isTextual()).isTrue(); + assertThat(output.get("upstreamEventId").asText()).isEqualTo(middle.get("id").asText()); + assertThat(output.get("upstreamActionName").asText()).isEqualTo("action2"); + + assertThat(middle.get("attributes").has("upstreamEventId")).isFalse(); + assertThat(middle.get("attributes").has("upstreamActionName")).isFalse(); + } + + private static AgentConfiguration fileLoggerConfig(Path logDir) { + AgentConfiguration config = new AgentConfiguration(); + config.set(AgentConfigOptions.EVENT_LOGGER_TYPE, LoggerType.FILE); + config.set(AgentConfigOptions.BASE_LOG_DIR, logDir.toString()); + config.set(AgentConfigOptions.EVENT_LOG_LEVEL, EventLogLevel.STANDARD); + config.set(AgentConfigOptions.EVENT_LOG_MAX_STRING_LENGTH, 3); + return config; + } + + private static Map readRecordsByType(Path logDir) throws Exception { + List lines; + try (Stream files = Files.list(logDir)) { + List logFiles = + files.filter(path -> path.getFileName().toString().endsWith(".log")) + .collect(Collectors.toList()); + assertThat(logFiles).hasSize(1); + lines = Files.readAllLines(logFiles.get(0)); + } + + Map recordsByType = new LinkedHashMap<>(); + for (String line : lines) { + JsonNode record = MAPPER.readTree(line); + recordsByType.put(record.get("eventType").asText(), record); + } + return recordsByType; + } +}