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..540b6e959 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 @@ -62,7 +62,7 @@ public Event( if (type == null || type.isEmpty()) { throw new IllegalArgumentException("Event 'type' must not be null or empty."); } - this.id = id; + this.id = id != null ? id : UUID.randomUUID(); this.type = type; this.attributes = attributes != null ? attributes : new HashMap<>(); } diff --git a/api/src/main/java/org/apache/flink/agents/api/EventContext.java b/api/src/main/java/org/apache/flink/agents/api/EventContext.java index f50f1700f..1b662ed3c 100644 --- a/api/src/main/java/org/apache/flink/agents/api/EventContext.java +++ b/api/src/main/java/org/apache/flink/agents/api/EventContext.java @@ -23,12 +23,12 @@ import java.time.Instant; -/** Contextual information about an event, such as its type and timestamp. */ +/** Runtime context for one observed event occurrence. */ public class EventContext { - /** The routing key for the event, matching the {@code EVENT_TYPE} constant or type string. */ + /** The event type observed by the runtime, matching {@link Event#getType()}. */ private final String eventType; - // Timestamp of when the event occurred + /** Timestamp of when the event occurrence was observed by the runtime. */ private final String timestamp; public EventContext(Event event) { diff --git a/api/src/main/java/org/apache/flink/agents/api/configuration/AgentConfigOptions.java b/api/src/main/java/org/apache/flink/agents/api/configuration/AgentConfigOptions.java index c39997da1..3fc8b87ba 100644 --- a/api/src/main/java/org/apache/flink/agents/api/configuration/AgentConfigOptions.java +++ b/api/src/main/java/org/apache/flink/agents/api/configuration/AgentConfigOptions.java @@ -117,6 +117,14 @@ public class AgentConfigOptions { public static final ConfigOption EVENT_LOG_LEVEL = new ConfigOption<>("event-log.level", EventLogLevel.class, EventLogLevel.STANDARD); + /** + * Whether to persist Agent Trace context and execution lifecycle Events in the Event Log. + * Business Events continue to be logged without trace context when this option is disabled. + * In-process execution reporting, including built-in metrics, is unaffected. + */ + public static final ConfigOption EVENT_LOG_TRACE_ENABLED = + new ConfigOption<>("event-log.trace.enabled", Boolean.class, false); + /** * The maximum string length for event payloads when logging at STANDARD level. Strings * exceeding this length will be truncated. Defaults to 2000. diff --git a/api/src/main/java/org/apache/flink/agents/api/listener/EventListener.java b/api/src/main/java/org/apache/flink/agents/api/listener/EventListener.java index 6c5b9fb27..6c19595dd 100644 --- a/api/src/main/java/org/apache/flink/agents/api/listener/EventListener.java +++ b/api/src/main/java/org/apache/flink/agents/api/listener/EventListener.java @@ -22,11 +22,12 @@ import org.apache.flink.agents.api.EventContext; /** - * Interface for event listeners that are notified when events are received for processing. + * Interface for event listeners that are notified when business events are received for processing. * *

EventListener provides a callback mechanism triggered at the beginning of event processing. * This is useful for monitoring, metrics collection, debugging, or triggering side effects based on - * event reception. + * event reception. Runtime observability records such as execution trace lifecycle events are + * written to Event Log but are not delivered to this listener. * *

Event listeners are executed synchronously when an event is received, before any actions are * triggered. Implementations should be lightweight and avoid blocking operations to prevent diff --git a/api/src/main/java/org/apache/flink/agents/api/logger/EventLogger.java b/api/src/main/java/org/apache/flink/agents/api/logger/EventLogger.java index 8efb111e3..bfa324e62 100644 --- a/api/src/main/java/org/apache/flink/agents/api/logger/EventLogger.java +++ b/api/src/main/java/org/apache/flink/agents/api/logger/EventLogger.java @@ -20,9 +20,12 @@ import org.apache.flink.agents.api.Event; import org.apache.flink.agents.api.EventContext; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; + +import javax.annotation.Nullable; /** - * Interface for logging action events in the Flink Agents framework. + * Interface for logging events in the Flink Agents framework. * *

EventLogger provides a unified interface for capturing, filtering, and persisting events as * they flow through the agent execution pipeline. Implementations can target different storage @@ -49,11 +52,24 @@ public interface EventLogger extends AutoCloseable { * should perform the append operation efficiently, as it may be called frequently during event * processing. * - * @param context metadata and other contextual information associated with the event + * @param eventContext runtime context associated with the event occurrence * @param event the event to be logged * @throws Exception if the append operation fails */ - void append(EventContext context, Event event) throws Exception; + void append(EventContext eventContext, Event event) throws Exception; + + /** + * Appends an event with runtime event context and optional trace context. + * + *

Implementations that persist trace context should override this method. Existing + * implementations can keep implementing {@link #append(EventContext, Event)}; the default + * behavior preserves the previous contract and ignores the optional trace context. + */ + default void append( + EventContext eventContext, Event event, @Nullable ExecutionTraceContext traceContext) + throws Exception { + append(eventContext, event); + } /** * Flush any buffered events to the underlying storage. diff --git a/api/src/main/java/org/apache/flink/agents/api/tools/ToolExecutionMetadataProvider.java b/api/src/main/java/org/apache/flink/agents/api/tools/ToolExecutionMetadataProvider.java new file mode 100644 index 000000000..3b808d9fd --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/tools/ToolExecutionMetadataProvider.java @@ -0,0 +1,38 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.flink.agents.api.tools; + +import java.util.Map; + +/** + * Optional Tool capability for adding structured metadata to Tool executions. + * + *

The returned metadata is supplemental: it must not define execution identity, parent/child + * relationships, or event payload. Implementations should keep values small, stable, and JSON + * serializable. + */ +public interface ToolExecutionMetadataProvider { + + /** + * Returns additional metadata for the current tool call. + * + * @param parameters tool call parameters associated with the execution + * @return small structured metadata to merge into execution {@code entityMetadata} + */ + Map getToolExecutionMetadata(ToolParameters parameters); +} diff --git a/api/src/main/java/org/apache/flink/agents/api/trace/ExecutionLifecycleEvents.java b/api/src/main/java/org/apache/flink/agents/api/trace/ExecutionLifecycleEvents.java new file mode 100644 index 000000000..dd00a5990 --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/trace/ExecutionLifecycleEvents.java @@ -0,0 +1,108 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.agents.api.trace; + +import org.apache.flink.agents.api.Event; + +import javax.annotation.Nullable; + +import java.util.Collections; +import java.util.HashMap; +import java.util.IdentityHashMap; +import java.util.Map; +import java.util.Set; + +/** Event factory for execution lifecycle reports in the trace model. */ +public final class ExecutionLifecycleEvents { + + public static final String EXECUTION_STARTED_EVENT_TYPE = "_execution_started_event"; + public static final String EXECUTION_FINISHED_EVENT_TYPE = "_execution_finished_event"; + public static final String EXECUTION_FAILED_EVENT_TYPE = "_execution_failed_event"; + public static final String EXECUTION_REUSED_EVENT_TYPE = "_execution_reused_event"; + + public static final String STATUS_STARTED = "started"; + public static final String STATUS_SUCCESS = "success"; + public static final String STATUS_FAILED = "failed"; + public static final String STATUS_REUSED = "reused"; + public static final String STATUS_ATTRIBUTE = "status"; + public static final String PROBLEM_CATEGORY_ATTRIBUTE = "problemCategory"; + + private ExecutionLifecycleEvents() {} + + public static Event executionStarted() { + return eventWithStatus(EXECUTION_STARTED_EVENT_TYPE, STATUS_STARTED); + } + + public static Event executionFinished() { + return eventWithStatus(EXECUTION_FINISHED_EVENT_TYPE, STATUS_SUCCESS); + } + + public static Event executionReused() { + return eventWithStatus(EXECUTION_REUSED_EVENT_TYPE, STATUS_REUSED); + } + + /** Returns whether the given type identifies an execution lifecycle event. */ + public static boolean isExecutionLifecycleEvent(String eventType) { + return EXECUTION_STARTED_EVENT_TYPE.equals(eventType) + || EXECUTION_FINISHED_EVENT_TYPE.equals(eventType) + || EXECUTION_FAILED_EVENT_TYPE.equals(eventType) + || EXECUTION_REUSED_EVENT_TYPE.equals(eventType); + } + + public static Event executionFailed(Throwable error) { + return executionFailed(error, null); + } + + public static Event executionFailed(Throwable error, @Nullable String problemCategory) { + Throwable rootCause = rootCause(error); + return executionFailed( + rootCause.getClass().getName(), rootCause.getMessage(), problemCategory); + } + + public static Event executionFailed( + String errorType, @Nullable String errorMessage, @Nullable String problemCategory) { + Map attributes = new HashMap<>(); + attributes.put(STATUS_ATTRIBUTE, STATUS_FAILED); + if (problemCategory != null) { + attributes.put(PROBLEM_CATEGORY_ATTRIBUTE, problemCategory); + } + if (errorType != null && !errorType.isEmpty()) { + attributes.put("errorType", errorType); + } + if (errorMessage != null && !errorMessage.isEmpty()) { + attributes.put("errorMessage", errorMessage); + } + return new Event(EXECUTION_FAILED_EVENT_TYPE, attributes); + } + + private static Event eventWithStatus(String eventType, String status) { + Map attributes = new HashMap<>(); + attributes.put(STATUS_ATTRIBUTE, status); + return new Event(eventType, attributes); + } + + private static Throwable rootCause(Throwable error) { + Throwable current = error; + Set visited = Collections.newSetFromMap(new IdentityHashMap<>()); + while (visited.add(current) && current.getCause() != null) { + current = current.getCause(); + } + return current; + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/trace/ExecutionReporter.java b/api/src/main/java/org/apache/flink/agents/api/trace/ExecutionReporter.java new file mode 100644 index 000000000..a3fb6f11b --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/trace/ExecutionReporter.java @@ -0,0 +1,88 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.flink.agents.api.trace; + +import javax.annotation.Nullable; + +import java.util.Map; + +/** + * Optional capability for reporting logical executions nested inside the current action. + * + *

Implementations decide how reports are consumed or ignored. Callers should provide stable + * entity type/name pairs and keep metadata small, structured, serializable, and stable for equality + * matching between the start and terminal reports of the same logical execution. + */ +public interface ExecutionReporter { + + /** Shared entity type names for execution reports. */ + final class EntityTypes { + public static final String ACTION = "action"; + public static final String LLM = "llm"; + public static final String PARSER = "parser"; + public static final String TOOL = "tool"; + + private EntityTypes() {} + } + + /** Shared low-cardinality problem categories for failed execution reports. */ + final class ProblemCategories { + public static final String ACTION_EXECUTION_FAILED = "action_execution_failed"; + public static final String MODEL_CALL_FAILED = "model_call_failed"; + public static final String MODEL_OUTPUT_PARSE_ERROR = "model_output_parse_error"; + public static final String TOOL_CALL_FAILED = "tool_call_failed"; + + private ProblemCategories() {} + } + + /** + * Reports that a logical execution started within the current action. + * + * @param entityType stable category of the reported execution, such as LLM, parser, or tool + * @param entityName stable name of the reported execution, such as model or tool name + * @param entityMetadata small structured metadata used to distinguish repeated executions with + * the same type/name + */ + void reportExecutionStarted( + String entityType, String entityName, Map entityMetadata) + throws Exception; + + /** + * Reports that a previously started logical execution completed successfully. + * + *

The entity type/name/metadata should match the corresponding start report when one was + * reported. + */ + void reportExecutionSucceeded( + String entityType, String entityName, Map entityMetadata) + throws Exception; + + /** + * Reports that a logical execution failed. + * + *

The entity type/name/metadata should match the corresponding start report when one was + * reported. The problem category should be a stable, low-cardinality classification. + */ + void reportExecutionFailed( + String entityType, + String entityName, + Map entityMetadata, + Throwable error, + @Nullable String problemCategory) + throws Exception; +} diff --git a/api/src/main/java/org/apache/flink/agents/api/trace/ExecutionReporters.java b/api/src/main/java/org/apache/flink/agents/api/trace/ExecutionReporters.java new file mode 100644 index 000000000..87ca12445 --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/trace/ExecutionReporters.java @@ -0,0 +1,117 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.flink.agents.api.trace; + +import org.apache.flink.agents.api.context.RunnerContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nullable; + +import java.util.Map; + +/** + * Utilities for reporting optional nested executions through a {@link RunnerContext}. + * + *

These helpers are best-effort: reporter failures are ignored so execution reporting never + * changes action success or failure. When reporting a failed execution, reporter failures are also + * attached as suppressed exceptions to the business error so callers that later throw that error + * can still inspect the reporting failure. + */ +public final class ExecutionReporters { + + private static final Logger LOG = LoggerFactory.getLogger(ExecutionReporters.class); + + private static final Map EMPTY_METADATA = Map.of(); + + private ExecutionReporters() {} + + public static void started(RunnerContext ctx, String entityType, String entityName) { + started(ctx, entityType, entityName, EMPTY_METADATA); + } + + public static void started( + RunnerContext ctx, + String entityType, + String entityName, + Map entityMetadata) { + report( + ctx, + reporter -> reporter.reportExecutionStarted(entityType, entityName, entityMetadata), + null); + } + + public static void succeeded(RunnerContext ctx, String entityType, String entityName) { + succeeded(ctx, entityType, entityName, EMPTY_METADATA); + } + + public static void succeeded( + RunnerContext ctx, + String entityType, + String entityName, + Map entityMetadata) { + report( + ctx, + reporter -> + reporter.reportExecutionSucceeded(entityType, entityName, entityMetadata), + null); + } + + public static void failed( + RunnerContext ctx, + String entityType, + String entityName, + Throwable error, + @Nullable String problemCategory) { + failed(ctx, entityType, entityName, EMPTY_METADATA, error, problemCategory); + } + + public static void failed( + RunnerContext ctx, + String entityType, + String entityName, + Map entityMetadata, + Throwable error, + @Nullable String problemCategory) { + report( + ctx, + reporter -> + reporter.reportExecutionFailed( + entityType, entityName, entityMetadata, error, problemCategory), + error); + } + + private static void report( + RunnerContext ctx, ReporterCall reporterCall, @Nullable Throwable businessError) { + if (ctx instanceof ExecutionReporter) { + try { + reporterCall.report((ExecutionReporter) ctx); + } catch (Exception reportError) { + if (businessError != null && businessError != reportError) { + businessError.addSuppressed(reportError); + } + LOG.debug("Execution reporting failed and was ignored.", reportError); + } + } + } + + @FunctionalInterface + private interface ReporterCall { + void report(ExecutionReporter reporter) throws Exception; + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/trace/ExecutionTraceContext.java b/api/src/main/java/org/apache/flink/agents/api/trace/ExecutionTraceContext.java new file mode 100644 index 000000000..a8fb8f8ae --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/trace/ExecutionTraceContext.java @@ -0,0 +1,213 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.agents.api.trace; + +import javax.annotation.Nullable; + +import java.io.Serializable; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Objects; +import java.util.UUID; + +/** Runtime context used to propagate input-run and execution-lineage information. */ +public final class ExecutionTraceContext implements Serializable { + private static final long serialVersionUID = 1L; + + private final String inputRunId; + private final String businessKey; + private final String agentName; + private final String executionId; + private final String parentExecutionId; + private final String entityType; + private final String entityName; + private final Map entityMetadata; + + private ExecutionTraceContext( + String inputRunId, + String businessKey, + String agentName, + String executionId, + String parentExecutionId, + String entityType, + String entityName, + Map entityMetadata) { + this.inputRunId = inputRunId; + this.businessKey = businessKey; + this.agentName = agentName; + this.executionId = executionId; + this.parentExecutionId = parentExecutionId; + this.entityType = entityType; + this.entityName = entityName; + this.entityMetadata = copyMetadata(entityMetadata); + } + + public static ExecutionTraceContext forInputRun(String businessKey) { + return forInputRun(businessKey, null); + } + + public static ExecutionTraceContext forInputRun(String businessKey, String agentName) { + return new ExecutionTraceContext( + UUID.randomUUID().toString(), businessKey, agentName, null, null, null, null, null); + } + + /** Creates an execution context without Agent identity. */ + public static ExecutionTraceContext forExecution( + String inputRunId, + String businessKey, + String parentExecutionId, + String entityType, + String entityName) { + return new ExecutionTraceContext( + inputRunId, + businessKey, + null, + UUID.randomUUID().toString(), + parentExecutionId, + entityType, + entityName, + null); + } + + /** Creates a sibling Action execution from the input scope carried by another trace context. */ + public static ExecutionTraceContext forAction( + ExecutionTraceContext sourceContext, String actionName) { + Objects.requireNonNull(sourceContext, "sourceContext"); + return new ExecutionTraceContext( + sourceContext.inputRunId, + sourceContext.businessKey, + sourceContext.agentName, + UUID.randomUUID().toString(), + null, + ExecutionReporter.EntityTypes.ACTION, + actionName, + null); + } + + public static ExecutionTraceContext fromExistingIds( + @Nullable String inputRunId, + @Nullable String businessKey, + @Nullable String agentName, + @Nullable String executionId, + @Nullable String parentExecutionId, + @Nullable String entityType, + @Nullable String entityName, + @Nullable Map entityMetadata) { + return new ExecutionTraceContext( + inputRunId, + businessKey, + agentName, + executionId, + parentExecutionId, + entityType, + entityName, + entityMetadata); + } + + /** Creates a child execution linked to this execution through {@code parentExecutionId}. */ + public ExecutionTraceContext childExecution(String entityType, String entityName) { + return childExecution(entityType, entityName, null); + } + + /** Creates a child execution with entity metadata and the same containment semantics. */ + public ExecutionTraceContext childExecution( + String entityType, String entityName, Map entityMetadata) { + return new ExecutionTraceContext( + inputRunId, + businessKey, + agentName, + UUID.randomUUID().toString(), + executionId, + entityType, + entityName, + entityMetadata); + } + + public String getInputRunId() { + return inputRunId; + } + + public String getBusinessKey() { + return businessKey; + } + + public String getAgentName() { + return agentName; + } + + public String getExecutionId() { + return executionId; + } + + public String getParentExecutionId() { + return parentExecutionId; + } + + public String getEntityType() { + return entityType; + } + + public String getEntityName() { + return entityName; + } + + public Map getEntityMetadata() { + return Collections.unmodifiableMap(entityMetadata); + } + + private static Map copyMetadata(Map entityMetadata) { + if (entityMetadata == null) { + return new LinkedHashMap<>(); + } + return new LinkedHashMap<>(entityMetadata); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof ExecutionTraceContext)) { + return false; + } + ExecutionTraceContext that = (ExecutionTraceContext) o; + return Objects.equals(inputRunId, that.inputRunId) + && Objects.equals(businessKey, that.businessKey) + && Objects.equals(agentName, that.agentName) + && Objects.equals(executionId, that.executionId) + && Objects.equals(parentExecutionId, that.parentExecutionId) + && Objects.equals(entityType, that.entityType) + && Objects.equals(entityName, that.entityName) + && Objects.equals(entityMetadata, that.entityMetadata); + } + + @Override + public int hashCode() { + return Objects.hash( + inputRunId, + businessKey, + agentName, + executionId, + parentExecutionId, + entityType, + entityName, + entityMetadata); + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/trace/ToolExecutionMetadataKeys.java b/api/src/main/java/org/apache/flink/agents/api/trace/ToolExecutionMetadataKeys.java new file mode 100644 index 000000000..85f910c9f --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/trace/ToolExecutionMetadataKeys.java @@ -0,0 +1,32 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.flink.agents.api.trace; + +/** Shared field names for Tool execution entity metadata. */ +public final class ToolExecutionMetadataKeys { + + public static final String TOOL_REQUEST_EVENT_ID = "toolRequestEventId"; + public static final String TOOL_CALL_ID = "toolCallId"; + public static final String EXTERNAL_ID = "externalId"; + public static final String TOOL_TYPE = "toolType"; + public static final String MCP_SERVER = "mcpServer"; + public static final String SKILL_NAME = "skillName"; + public static final String SKILL_RESOURCE_PATH = "skillResourcePath"; + + private ToolExecutionMetadataKeys() {} +} 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..ce9e765e8 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 @@ -150,6 +150,7 @@ void testFromJsonWithValidUnifiedEvent() throws IOException { String json = "{\"type\":\"MyCustomEvent\",\"attributes\":{\"msg\":\"hello\"}}"; Event event = Event.fromJson(json); + assertNotNull(event.getId()); assertEquals("MyCustomEvent", event.getType()); assertEquals("hello", event.getAttr("msg")); } diff --git a/api/src/test/java/org/apache/flink/agents/api/trace/ExecutionLifecycleEventsTest.java b/api/src/test/java/org/apache/flink/agents/api/trace/ExecutionLifecycleEventsTest.java new file mode 100644 index 000000000..bec417e39 --- /dev/null +++ b/api/src/test/java/org/apache/flink/agents/api/trace/ExecutionLifecycleEventsTest.java @@ -0,0 +1,56 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.flink.agents.api.trace; + +import org.apache.flink.agents.api.Event; +import org.junit.jupiter.api.Test; + +import java.time.Duration; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively; + +/** Tests for {@link ExecutionLifecycleEvents}. */ +class ExecutionLifecycleEventsTest { + + @Test + void executionFailedUsesDeepestCause() { + IllegalArgumentException root = new IllegalArgumentException("root"); + RuntimeException error = new RuntimeException("outer", root); + + Event event = ExecutionLifecycleEvents.executionFailed(error); + + assertThat(event.getAttr("errorType")).isEqualTo(root.getClass().getName()); + assertThat(event.getAttr("errorMessage")).isEqualTo("root"); + } + + @Test + void executionFailedTerminatesForCyclicCauseChain() { + RuntimeException first = new RuntimeException("first"); + RuntimeException second = new RuntimeException("second", first); + first.initCause(second); + + Event event = + assertTimeoutPreemptively( + Duration.ofSeconds(1), + () -> ExecutionLifecycleEvents.executionFailed(first)); + + assertThat(event.getAttr("errorType")).isEqualTo(first.getClass().getName()); + assertThat(event.getAttr("errorMessage")).isEqualTo("first"); + } +} diff --git a/api/src/test/java/org/apache/flink/agents/api/trace/ExecutionReportersTest.java b/api/src/test/java/org/apache/flink/agents/api/trace/ExecutionReportersTest.java new file mode 100644 index 000000000..469903634 --- /dev/null +++ b/api/src/test/java/org/apache/flink/agents/api/trace/ExecutionReportersTest.java @@ -0,0 +1,132 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.flink.agents.api.trace; + +import org.apache.flink.agents.api.context.RunnerContext; +import org.junit.jupiter.api.Test; + +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatCode; +import static org.mockito.ArgumentMatchers.anyMap; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.same; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.withSettings; + +/** Tests for {@link ExecutionReporters}. */ +class ExecutionReportersTest { + + @Test + void startedAndSucceededIgnoreReporterFailures() throws Exception { + RunnerContext ctx = mockReportingContext(); + ExecutionReporter reporter = (ExecutionReporter) ctx; + Exception startedError = new Exception("started failed"); + Exception succeededError = new Exception("succeeded failed"); + doThrow(startedError) + .when(reporter) + .reportExecutionStarted( + eq(ExecutionReporter.EntityTypes.LLM), eq("model-a"), anyMap()); + doThrow(succeededError) + .when(reporter) + .reportExecutionSucceeded( + eq(ExecutionReporter.EntityTypes.LLM), eq("model-a"), anyMap()); + + assertThatCode( + () -> + ExecutionReporters.started( + ctx, ExecutionReporter.EntityTypes.LLM, "model-a")) + .doesNotThrowAnyException(); + assertThatCode( + () -> + ExecutionReporters.succeeded( + ctx, ExecutionReporter.EntityTypes.LLM, "model-a")) + .doesNotThrowAnyException(); + + verify(reporter) + .reportExecutionStarted( + eq(ExecutionReporter.EntityTypes.LLM), eq("model-a"), anyMap()); + verify(reporter) + .reportExecutionSucceeded( + eq(ExecutionReporter.EntityTypes.LLM), eq("model-a"), anyMap()); + } + + @Test + void failedIgnoresReporterFailureAndSuppressesItOnBusinessError() throws Exception { + RunnerContext ctx = mockReportingContext(); + ExecutionReporter reporter = (ExecutionReporter) ctx; + RuntimeException businessError = new RuntimeException("business failed"); + Exception reportError = new Exception("report failed"); + doThrow(reportError) + .when(reporter) + .reportExecutionFailed( + eq(ExecutionReporter.EntityTypes.TOOL), + eq("search"), + anyMap(), + same(businessError), + eq(ExecutionReporter.ProblemCategories.TOOL_CALL_FAILED)); + + assertThatCode( + () -> + ExecutionReporters.failed( + ctx, + ExecutionReporter.EntityTypes.TOOL, + "search", + Map.of("toolCallId", "call-1"), + businessError, + ExecutionReporter.ProblemCategories.TOOL_CALL_FAILED)) + .doesNotThrowAnyException(); + + assertThat(businessError.getSuppressed()).containsExactly(reportError); + } + + @Test + void helpersIgnoreContextsWithoutExecutionReporter() { + RunnerContext ctx = mock(RunnerContext.class); + RuntimeException businessError = new RuntimeException("business failed"); + + assertThatCode( + () -> + ExecutionReporters.started( + ctx, ExecutionReporter.EntityTypes.LLM, "model-a")) + .doesNotThrowAnyException(); + assertThatCode( + () -> + ExecutionReporters.succeeded( + ctx, ExecutionReporter.EntityTypes.LLM, "model-a")) + .doesNotThrowAnyException(); + assertThatCode( + () -> + ExecutionReporters.failed( + ctx, + ExecutionReporter.EntityTypes.LLM, + "model-a", + businessError, + ExecutionReporter.ProblemCategories.MODEL_CALL_FAILED)) + .doesNotThrowAnyException(); + + assertThat(businessError).hasNoSuppressedExceptions(); + } + + private static RunnerContext mockReportingContext() { + return mock(RunnerContext.class, withSettings().extraInterfaces(ExecutionReporter.class)); + } +} diff --git a/docs/content/docs/development/workflow_agent.md b/docs/content/docs/development/workflow_agent.md index 74589a1fe..acc2f515c 100644 --- a/docs/content/docs/development/workflow_agent.md +++ b/docs/content/docs/development/workflow_agent.md @@ -659,8 +659,8 @@ class MyEvent(Event): 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. + # Preserve the base event id so event logs, listeners, correlation, + # and deduplication stay consistent. result.id = event.id return result @@ -712,8 +712,7 @@ listeners, correlation, deduplication, and downstream timestamp propagation stay 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. + in both languages. - **`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. @@ -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/configuration.md b/docs/content/docs/operations/configuration.md index 9ba9e200e..1c73a6ca6 100644 --- a/docs/content/docs/operations/configuration.md +++ b/docs/content/docs/operations/configuration.md @@ -136,6 +136,7 @@ Here is the list of all built-in core configuration options. | `num-async-threads` | os cpu count * 2 | int | The thread pool size for async executor. | | `job-identifier` | none | String | The unique identifier of job, remaining consistent after restoring from a savepoint. If not set, uses flink job id. | | `event-log.level` | STANDARD | EventLogLevel | Global default verbosity for the [Event Log]({{< ref "docs/operations/monitoring#event-log" >}}). Valid values: `OFF` (skip event), `STANDARD` (payload may be truncated/summarized to keep logs concise), `VERBOSE` (full payload). Can be overridden per event type — see [Per-event-type log levels]({{< ref "docs/operations/monitoring#per-event-type-log-levels" >}}). | +| `event-log.trace.enabled` | false | boolean | Whether to persist Agent Trace information in the Event Log. When enabled, business Events include trace context and Action/LLM/Parser/Tool lifecycle Events are logged. | | `event-log.type..level` | (inherits) | EventLogLevel | Override the log level for a specific event type. `` is the event's routing type string (the same value that appears as `eventType` in the JSON log, e.g., `_chat_request_event` for built-ins, or `com.example.myapp.OrderEvent` for user-defined types). For dotted types, resolution walks up dot segments before falling back to `event-log.level`. See [Per-event-type log levels]({{< ref "docs/operations/monitoring#per-event-type-log-levels" >}}) for examples. | | `event-log.standard.max-string-length` | 2000 | int | At `STANDARD` level, strings in the event payload longer than this are truncated. Has no effect at `VERBOSE`. | | `event-log.standard.max-array-elements` | 20 | int | At `STANDARD` level, arrays in the event payload with more than this many elements are truncated. Has no effect at `VERBOSE`. | diff --git a/docs/content/docs/operations/monitoring.md b/docs/content/docs/operations/monitoring.md index 021452125..dfcc82ce7 100644 --- a/docs/content/docs/operations/monitoring.md +++ b/docs/content/docs/operations/monitoring.md @@ -178,17 +178,27 @@ By default, all File-based Event Logs are stored in the `flink-agents` subdirect The JSON record format described here applies to both the SLF4J and File event loggers. The SLF4J logger adds `jobId`, `taskName`, and `subtaskId` fields on top (see [SLF4J Event Log](#slf4j-event-log-default)); the File logger encodes those values in the file path instead. -Each record contains a top-level `timestamp`, the resolved `logLevel`, and a top-level `eventType` routing key (mirrors `event.eventType`), followed by the full event object. The top-level `eventType` makes it easy for downstream tools (e.g. `grep`, `jq`, log shippers) to filter by event type without parsing nested JSON: +Each record is a flat JSON object. Framework-owned field names use camelCase, consistent with the existing Event Log format. A record contains the event occurrence time, the resolved `logLevel`, optional execution trace fields, the event identity, and the event payload in `eventAttributes`. The flat `eventType` field makes it easy for downstream tools (e.g. `grep`, `jq`, log shippers) to filter by event type without parsing nested JSON. + +Agent Trace persistence is disabled by default. Set `event-log.trace.enabled: true` to add trace context to business Event records and persist Action, LLM, Parser, and Tool lifecycle Events. When Trace persistence is disabled, business Events continue to be logged without the trace fields shown below. + +Example Trace record: ```json { "timestamp": "2024-01-15T10:30:00Z", "logLevel": "STANDARD", - "eventType": "_input_event", - "event": { - "eventType": "_input_event", - "...": "..." - } + "inputRunId": "7f1b5d20-86c6-4a5d-b65a-9c41757f2e11", + "businessKey": "order-1001", + "agentName": "ReActAgent", + "executionId": "9cb3c3df-1d3b-4c45-ae7d-1b28a1b86522", + "parentExecutionId": "55cc59a8-f4e8-4f15-badf-4833f8a8a97a", + "entityType": "llm", + "entityName": "qwen-max", + "eventId": "80f30736-2759-41d8-aa59-9c8ad481ab42", + "eventType": "_execution_finished_event", + "status": "success", + "eventAttributes": {} } ``` @@ -204,7 +214,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`, trace fields such as `inputRunId` and `executionId`, `eventId`, `eventType`, and lifecycle fields such as `status` and `problemCategory`. Truncation only applies to large nested content under `eventAttributes` (long strings, big arrays, deeply nested objects). **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: @@ -222,10 +232,9 @@ Example record at `STANDARD` with a long string and a large array truncated: { "timestamp": "2024-01-15T10:30:00Z", "logLevel": "STANDARD", + "eventId": "...", "eventType": "_chat_request_event", - "event": { - "eventType": "_chat_request_event", - "id": "...", + "eventAttributes": { "model": "gpt-4", "messages": { "truncatedList": [ @@ -240,7 +249,7 @@ Example record at `STANDARD` with a long string and a large array truncated: ### Per-event-type log levels -You can override the level for individual event types using the `event-log.type..level` config key, where `` is the event's routing type string (the same string that appears as `eventType` in the JSON log). Built-in events use short snake-cased names such as: +You can override the level for individual event types using the `event-log.type..level` config key, where `` is the event's routing type string (the same string that appears as `eventType` in the JSON log). Although the field name uses camelCase, built-in Event type values remain snake-cased: | Event class | `` value | |--------------------------|----------------------------------| @@ -291,4 +300,6 @@ Other per-type levels from `config.yaml` are preserved — the `-D` flag only ov ### Compatibility Notes - **Default behavior changed.** Before this feature, every event was logged in full. The new default is `STANDARD`, which truncates large payloads. To restore the previous behavior either globally or per type, set the level to `VERBOSE`. -- **Old log records still parse.** Records written before this feature have no `logLevel` or top-level `eventType`. They deserialize correctly and are treated as `VERBOSE` (their payloads were never truncated). +- **Old log records still parse.** Records written in the previous nested format (`eventType` plus `event`) continue to deserialize. They do not contain run or execution context and therefore cannot reconstruct a complete Agent Trace. +- **External consumers must migrate to the flat field names.** New records expose fields such as `eventType` and `eventAttributes` at the top level. Compatibility in the framework deserializer does not automatically update existing `jq` expressions, log-shipper mappings, or downstream queries that read the previous nested JSON shape directly. +- **Existing pending ActionTask state is not compatible with the new schema.** Agent Trace adds execution identity and Action lifecycle state to pending `ActionTask` state. Restoring savepoints containing that state from before this change would require a versioned `ActionTask` state serializer, which is not included in this feature. diff --git a/e2e-test/cross-language-agent-plan-snapshots/java/agent_plan_with_python_action.json b/e2e-test/cross-language-agent-plan-snapshots/java/agent_plan_with_python_action.json index 8b6e28867..2dacb24df 100644 --- a/e2e-test/cross-language-agent-plan-snapshots/java/agent_plan_with_python_action.json +++ b/e2e-test/cross-language-agent-plan-snapshots/java/agent_plan_with_python_action.json @@ -1,4 +1,5 @@ { + "agent_name" : "Agent", "actions" : { "chat_model_action" : { "name" : "chat_model_action", diff --git a/e2e-test/cross-language-agent-plan-snapshots/python/agent_plan_with_java_action.json b/e2e-test/cross-language-agent-plan-snapshots/python/agent_plan_with_java_action.json index 19dee444a..574599350 100644 --- a/e2e-test/cross-language-agent-plan-snapshots/python/agent_plan_with_java_action.json +++ b/e2e-test/cross-language-agent-plan-snapshots/python/agent_plan_with_java_action.json @@ -1,4 +1,5 @@ { + "agent_name": "Agent", "actions": { "handle": { "name": "handle", diff --git a/integrations/mcp/src/main/java/org/apache/flink/agents/integrations/mcp/MCPTool.java b/integrations/mcp/src/main/java/org/apache/flink/agents/integrations/mcp/MCPTool.java index 286fb2d44..05d732fbb 100644 --- a/integrations/mcp/src/main/java/org/apache/flink/agents/integrations/mcp/MCPTool.java +++ b/integrations/mcp/src/main/java/org/apache/flink/agents/integrations/mcp/MCPTool.java @@ -22,12 +22,15 @@ import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.annotation.JsonProperty; import org.apache.flink.agents.api.tools.Tool; +import org.apache.flink.agents.api.tools.ToolExecutionMetadataProvider; import org.apache.flink.agents.api.tools.ToolMetadata; import org.apache.flink.agents.api.tools.ToolParameters; import org.apache.flink.agents.api.tools.ToolResponse; import org.apache.flink.agents.api.tools.ToolType; +import org.apache.flink.agents.api.trace.ToolExecutionMetadataKeys; import java.util.HashMap; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Objects; @@ -38,25 +41,42 @@ *

This represents a single tool from an MCP server. It extends the base Tool class and delegates * actual execution to the MCP server. */ -public class MCPTool extends Tool { +public class MCPTool extends Tool implements ToolExecutionMetadataProvider { private static final String FIELD_MCP_SERVER = "mcpServer"; + private static final String FIELD_MCP_SERVER_NAME = "mcpServerName"; @JsonProperty(FIELD_MCP_SERVER) private final MCPServer mcpServer; + @JsonProperty(FIELD_MCP_SERVER_NAME) + private final String mcpServerName; + + public MCPTool(ToolMetadata metadata, MCPServer mcpServer) { + this(metadata, mcpServer, null); + } + /** * Create a new MCPTool. * * @param metadata The tool metadata * @param mcpServer The MCP server reference + * @param mcpServerName The AgentPlan resource name of the MCP server, if known */ @JsonCreator public MCPTool( @JsonProperty("metadata") ToolMetadata metadata, - @JsonProperty(FIELD_MCP_SERVER) MCPServer mcpServer) { + @JsonProperty(FIELD_MCP_SERVER) MCPServer mcpServer, + @JsonProperty(FIELD_MCP_SERVER_NAME) String mcpServerName) { super(metadata); this.mcpServer = Objects.requireNonNull(mcpServer, "mcpServer cannot be null"); + this.mcpServerName = mcpServerName; + } + + /** Returns a copy of this tool associated with the given MCP server resource name. */ + @JsonIgnore + public MCPTool withMcpServerName(String mcpServerName) { + return new MCPTool(metadata, mcpServer, mcpServerName); } @Override @@ -106,18 +126,33 @@ public MCPServer getMcpServer() { return mcpServer; } + /** Returns the AgentPlan resource name of the MCP server, if available. */ + public String getMcpServerName() { + return mcpServerName; + } + + @Override + public Map getToolExecutionMetadata(ToolParameters parameters) { + Map metadata = new LinkedHashMap<>(); + if (mcpServerName != null) { + metadata.put(ToolExecutionMetadataKeys.MCP_SERVER, mcpServerName); + } + return metadata; + } + @Override public boolean equals(Object o) { if (this == o) return true; if (o == null || getClass() != o.getClass()) return false; MCPTool mcpTool = (MCPTool) o; return Objects.equals(metadata, mcpTool.metadata) - && Objects.equals(mcpServer, mcpTool.mcpServer); + && Objects.equals(mcpServer, mcpTool.mcpServer) + && Objects.equals(mcpServerName, mcpTool.mcpServerName); } @Override public int hashCode() { - return Objects.hash(metadata, mcpServer); + return Objects.hash(metadata, mcpServer, mcpServerName); } @Override diff --git a/integrations/mcp/src/test/java/org/apache/flink/agents/integrations/mcp/MCPToolTest.java b/integrations/mcp/src/test/java/org/apache/flink/agents/integrations/mcp/MCPToolTest.java index 2d1b233f2..73bd2a0d3 100644 --- a/integrations/mcp/src/test/java/org/apache/flink/agents/integrations/mcp/MCPToolTest.java +++ b/integrations/mcp/src/test/java/org/apache/flink/agents/integrations/mcp/MCPToolTest.java @@ -20,12 +20,16 @@ import com.fasterxml.jackson.databind.ObjectMapper; import org.apache.flink.agents.api.tools.ToolMetadata; +import org.apache.flink.agents.api.tools.ToolParameters; import org.apache.flink.agents.api.tools.ToolType; +import org.apache.flink.agents.api.trace.ToolExecutionMetadataKeys; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.condition.DisabledOnJre; import org.junit.jupiter.api.condition.JRE; +import java.util.Map; + import static com.fasterxml.jackson.databind.DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES; import static org.assertj.core.api.Assertions.assertThat; @@ -61,6 +65,7 @@ void testJsonDeserializationWithDefaultMapper() throws Exception { assertThat(deserialized.getName()).isEqualTo("add"); assertThat(deserialized.getMetadata()).isEqualTo(metadata); assertThat(deserialized.getMcpServer()).isEqualTo(server); + assertThat(deserialized.getMcpServerName()).isNull(); } @Test @@ -148,7 +153,7 @@ void testJsonSerialization() throws Exception { ToolMetadata metadata = new ToolMetadata("multiply", "Multiply two numbers", "{\"type\":\"object\"}"); MCPServer server = new MCPServer(DEFAULT_ENDPOINT); - MCPTool original = new MCPTool(metadata, server); + MCPTool original = new MCPTool(metadata, server, "testMcpServer"); ObjectMapper mapper = new ObjectMapper(); // Configure to ignore unknown properties during deserialization @@ -161,6 +166,7 @@ void testJsonSerialization() throws Exception { assertThat(deserialized.getName()).isEqualTo(original.getName()); assertThat(deserialized.getMetadata()).isEqualTo(original.getMetadata()); assertThat(deserialized.getMcpServer()).isEqualTo(original.getMcpServer()); + assertThat(deserialized.getMcpServerName()).isEqualTo("testMcpServer"); } @Test @@ -185,4 +191,17 @@ void testToolWithComplexSchema() { assertThat(tool.getMetadata().getInputSchema()).contains("param3"); assertThat(tool.getMetadata().getInputSchema()).contains("required"); } + + @Test + @DisabledOnJre(JRE.JAVA_11) + @DisplayName("Execution metadata contains MCP server context") + void executionMetadataContainsMcpServerContext() { + ToolMetadata metadata = new ToolMetadata("add", "Add two numbers", "{\"type\":\"object\"}"); + MCPServer server = new MCPServer(DEFAULT_ENDPOINT); + + MCPTool tool = new MCPTool(metadata, server, "testMcpServer"); + + assertThat(tool.getToolExecutionMetadata(new ToolParameters(Map.of()))) + .containsEntry(ToolExecutionMetadataKeys.MCP_SERVER, "testMcpServer"); + } } diff --git a/plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java b/plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java index 3b08bf2cf..b4a22f3dc 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java @@ -97,6 +97,9 @@ public class AgentPlan implements Serializable { /** Two-level mapping of resource type to resource name to resource provider. */ private Map> resourceProviders; + /** User-visible agent identity used for observability. */ + private String agentName; + private AgentConfiguration config; public AgentPlan(Map actions, Map> actionsByEvent) { @@ -121,10 +124,20 @@ public AgentPlan( Map> actionsByEvent, Map> resourceProviders, AgentConfiguration config) { + this(actions, actionsByEvent, resourceProviders, config, null); + } + + public AgentPlan( + Map actions, + Map> actionsByEvent, + Map> resourceProviders, + AgentConfiguration config, + String agentName) { this.actions = actions; this.actionsByEvent = actionsByEvent; this.resourceProviders = resourceProviders; this.config = config; + this.agentName = agentName; } /** @@ -139,10 +152,15 @@ public AgentPlan(Agent agent) throws Exception { } public AgentPlan(Agent agent, AgentConfiguration config) throws Exception { + this(agent, config, defaultAgentName(agent)); + } + + public AgentPlan(Agent agent, AgentConfiguration config, String agentName) throws Exception { this(new HashMap<>(), new HashMap<>()); extractActionsFromAgent(agent); extractResourceProvidersFromAgent(agent); this.config = config; + this.agentName = agentName != null ? agentName : defaultAgentName(agent); } public Map getActions() { @@ -165,6 +183,10 @@ public Map> getResourceProviders() { return resourceProviders; } + public String getAgentName() { + return agentName; + } + public List getActionsTriggeredBy(String eventType) { return actionsByEvent.get(eventType); } @@ -188,9 +210,15 @@ private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundE this.actions = agentPlan.getActions(); this.actionsByEvent = agentPlan.getActionsByEvent(); this.resourceProviders = agentPlan.getResourceProviders(); + this.agentName = agentPlan.getAgentName(); this.config = agentPlan.getConfig(); } + private static String defaultAgentName(Agent agent) { + String simpleName = agent.getClass().getSimpleName(); + return simpleName == null || simpleName.isEmpty() ? agent.getClass().getName() : simpleName; + } + private void extractActions( String actionName, String[] triggerEntries, @@ -391,6 +419,7 @@ private void extractJavaMCPServer(Method method) throws Exception { (Iterable) listToolsMethod.invoke(mcpServer); for (SerializableResource tool : tools) { + tool = attachMcpServerNameIfSupported(tool, name); Method getNameMethod = tool.getClass().getMethod("getName"); String toolName = (String) getNameMethod.invoke(tool); addResourceProvider( @@ -416,6 +445,20 @@ private void extractJavaMCPServer(Method method) throws Exception { closeMethod.invoke(mcpServer); } + private static SerializableResource attachMcpServerNameIfSupported( + SerializableResource tool, String mcpServerName) throws Exception { + try { + Method method = tool.getClass().getMethod("withMcpServerName", String.class); + Object result = method.invoke(tool, mcpServerName); + if (result instanceof SerializableResource) { + return (SerializableResource) result; + } + return tool; + } catch (NoSuchMethodException e) { + return tool; + } + } + private void extractResourceProvidersFromAgent(Agent agent) throws Exception { Class agentClass = agent.getClass(); diff --git a/plan/src/main/java/org/apache/flink/agents/plan/JavaFunction.java b/plan/src/main/java/org/apache/flink/agents/plan/JavaFunction.java index b99e4328b..53f9ceeb1 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/JavaFunction.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/JavaFunction.java @@ -20,6 +20,7 @@ import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.annotation.JsonProperty; +import java.lang.reflect.InvocationTargetException; import java.lang.reflect.Method; import java.util.Arrays; import java.util.Objects; @@ -103,7 +104,18 @@ public int hashCode() { @Override public Object call(Object... args) throws Exception { - return getMethod().invoke(null, args); + try { + return getMethod().invoke(null, args); + } catch (InvocationTargetException e) { + Throwable target = e.getTargetException(); + if (target instanceof Exception) { + throw (Exception) target; + } + if (target instanceof Error) { + throw (Error) target; + } + throw new RuntimeException(target); + } } @Override diff --git a/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java b/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java index df28c1d41..f06db909f 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java @@ -38,6 +38,8 @@ import org.apache.flink.agents.api.resource.ResourceType; import org.apache.flink.agents.api.skills.Skills; import org.apache.flink.agents.api.tools.ToolResponse; +import org.apache.flink.agents.api.trace.ExecutionReporter; +import org.apache.flink.agents.api.trace.ExecutionReporters; import org.apache.flink.agents.plan.JavaFunction; import org.apache.flink.api.java.typeutils.RowTypeInfo; import org.apache.flink.types.Row; @@ -368,16 +370,27 @@ public ChatMessage call() throws Exception { for (int attempt = 0; attempt < numRetries + 1; attempt++) { try { - response = - chatAsync - ? ctx.durableExecuteAsync(callable) - : ctx.durableExecute(callable); + ExecutionReporters.started(ctx, ExecutionReporter.EntityTypes.LLM, model); + try { + response = + chatAsync + ? ctx.durableExecuteAsync(callable) + : ctx.durableExecute(callable); + Objects.requireNonNull(response, "ChatModel returned a null response."); + } catch (Exception modelError) { + ExecutionReporters.failed( + ctx, + ExecutionReporter.EntityTypes.LLM, + model, + modelError, + ExecutionReporter.ProblemCategories.MODEL_CALL_FAILED); + throw modelError; + } + ExecutionReporters.succeeded(ctx, ExecutionReporter.EntityTypes.LLM, model); recordChatTokenMetrics(chatModel, response); - // only generate structured output for final response. if (outputSchema != null && response.getToolCalls().isEmpty()) { - response = generateStructuredOutput(response, outputSchema); + response = generateStructuredOutputWithReport(ctx, response, outputSchema); } - break; } catch (Exception e) { if (strategy == Agent.ErrorHandlingStrategy.IGNORE) { LOG.warn( @@ -400,6 +413,7 @@ public ChatMessage call() throws Exception { Thread.sleep(currentWaitSec * 1000L); totalWaitTimeSec += currentWaitSec; } + continue; } else { LOG.debug( "Chat request {} failed, the input chat messages are {}.", @@ -408,6 +422,7 @@ public ChatMessage call() throws Exception { throw e; } } + break; } if (actualRetryCount > 0) { @@ -439,6 +454,25 @@ public ChatMessage call() throws Exception { } } + private static ChatMessage generateStructuredOutputWithReport( + RunnerContext ctx, ChatMessage response, Object outputSchema) throws Exception { + ExecutionReporters.started(ctx, ExecutionReporter.EntityTypes.PARSER, STRUCTURED_OUTPUT); + try { + ChatMessage structuredResponse = generateStructuredOutput(response, outputSchema); + ExecutionReporters.succeeded( + ctx, ExecutionReporter.EntityTypes.PARSER, STRUCTURED_OUTPUT); + return structuredResponse; + } catch (Exception e) { + ExecutionReporters.failed( + ctx, + ExecutionReporter.EntityTypes.PARSER, + STRUCTURED_OUTPUT, + e, + ExecutionReporter.ProblemCategories.MODEL_OUTPUT_PARSE_ERROR); + throw e; + } + } + private static void processChatRequest(ChatRequestEvent event, RunnerContext ctx) throws Exception { chat( diff --git a/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java b/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java index bc82e1b61..d0d6f1f68 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java @@ -27,19 +27,30 @@ import org.apache.flink.agents.api.event.ToolResponseEvent; import org.apache.flink.agents.api.resource.ResourceType; import org.apache.flink.agents.api.tools.Tool; +import org.apache.flink.agents.api.tools.ToolExecutionMetadataProvider; import org.apache.flink.agents.api.tools.ToolParameterInjection; import org.apache.flink.agents.api.tools.ToolParameterSource; import org.apache.flink.agents.api.tools.ToolParameters; import org.apache.flink.agents.api.tools.ToolResponse; +import org.apache.flink.agents.api.tools.ToolType; +import org.apache.flink.agents.api.trace.ExecutionReporter; +import org.apache.flink.agents.api.trace.ExecutionReporters; +import org.apache.flink.agents.api.trace.ToolExecutionMetadataKeys; import org.apache.flink.agents.plan.JavaFunction; import org.apache.flink.agents.plan.tools.FunctionTool; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import java.util.HashMap; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.UUID; /** Built-in action for processing tool call. */ public class ToolCallAction { + private static final Logger LOG = LoggerFactory.getLogger(ToolCallAction.class); + public static Action getToolCallAction() throws Exception { return new Action( "tool_call_action", @@ -72,11 +83,11 @@ public static void processToolRequest(Event event, RunnerContext ctx) { } Tool tool = null; - String diagnosticError = null; + Exception preparationError = null; try { tool = (Tool) ctx.getResource(name, ResourceType.TOOL); } catch (Exception e) { - diagnosticError = e.getMessage(); + preparationError = e; } if (tool != null) { @@ -84,52 +95,156 @@ public static void processToolRequest(Event event, RunnerContext ctx) { // Framework-owned injected args must win over model-provided values so hidden // context such as tenant ids cannot be spoofed by a tool call payload. mergedArguments.putAll(resolveInjectedArguments(tool, ctx)); - ToolResponse response; - final Tool toolRef = tool; - final Map callArguments = mergedArguments; - DurableCallable callable = - new DurableCallable<>() { - @Override - public String getId() { - return "tool-call"; - } - - @Override - public Class getResultClass() { - return ToolResponse.class; - } - - @Override - public ToolResponse call() throws Exception { - return toolRef.call(new ToolParameters(callArguments)); - } - }; - response = - toolCallAsync - ? ctx.durableExecuteAsync(callable) - : ctx.durableExecute(callable); - success.put(id, response.isSuccess()); - responses.put(id, response); - if (!response.isSuccess() && response.getError() != null) { - error.put(id, response.getError()); - } } catch (Exception e) { - success.put(id, false); - responses.put( - id, ToolResponse.error(String.format("Tool %s execute failed.", name))); - error.put(id, e.getMessage()); + preparationError = e; + } + } + + ToolParameters metadataParameters = new ToolParameters(mergedArguments); + Map entityMetadata = + toolEntityMetadata( + toolRequest.getId(), + id, + externalIds.get(id), + name, + tool, + metadataParameters); + ExecutionReporters.started( + ctx, ExecutionReporter.EntityTypes.TOOL, name, entityMetadata); + + if (tool == null || preparationError != null) { + Exception failure = + preparationError != null + ? preparationError + : new IllegalArgumentException("Tool does not exist."); + success.put(id, false); + responses.put( + id, + ToolResponse.error( + String.format( + tool == null + ? "Tool %s does not exist." + : "Tool %s execute failed.", + name))); + String failureMessage = failure.getMessage(); + if (failureMessage == null && tool == null) { + failureMessage = "Tool does not exist."; } - } else { + error.put(id, failureMessage); + ExecutionReporters.failed( + ctx, + ExecutionReporter.EntityTypes.TOOL, + name, + entityMetadata, + failure, + ExecutionReporter.ProblemCategories.TOOL_CALL_FAILED); + continue; + } + + ToolResponse response = null; + try { + final Tool toolRef = tool; + final Map callArguments = mergedArguments; + DurableCallable callable = + new DurableCallable<>() { + @Override + public String getId() { + return "tool-call"; + } + + @Override + public Class getResultClass() { + return ToolResponse.class; + } + + @Override + public ToolResponse call() throws Exception { + return toolRef.call(new ToolParameters(callArguments)); + } + }; + response = + toolCallAsync + ? ctx.durableExecuteAsync(callable) + : ctx.durableExecute(callable); + success.put(id, response.isSuccess()); + responses.put(id, response); + if (!response.isSuccess() && response.getError() != null) { + error.put(id, response.getError()); + } + } catch (Exception e) { success.put(id, false); responses.put( - id, ToolResponse.error(String.format("Tool %s does not exist.", name))); - error.put(id, diagnosticError != null ? diagnosticError : "Tool does not exist."); + id, ToolResponse.error(String.format("Tool %s execute failed.", name))); + error.put(id, e.getMessage()); + ExecutionReporters.failed( + ctx, + ExecutionReporter.EntityTypes.TOOL, + name, + entityMetadata, + e, + ExecutionReporter.ProblemCategories.TOOL_CALL_FAILED); + } + if (response != null) { + if (response.isSuccess()) { + ExecutionReporters.succeeded( + ctx, ExecutionReporter.EntityTypes.TOOL, name, entityMetadata); + } else { + ExecutionReporters.failed( + ctx, + ExecutionReporter.EntityTypes.TOOL, + name, + entityMetadata, + new RuntimeException(response.getError()), + ExecutionReporter.ProblemCategories.TOOL_CALL_FAILED); + } } } ctx.sendEvent( new ToolResponseEvent(toolRequest.getId(), responses, success, error, externalIds)); } + private static Map toolEntityMetadata( + UUID toolRequestEventId, + String toolCallId, + String externalId, + String toolName, + Tool tool, + ToolParameters parameters) { + Map metadata = new LinkedHashMap<>(); + metadata.put( + ToolExecutionMetadataKeys.TOOL_REQUEST_EVENT_ID, toolRequestEventId.toString()); + metadata.put(ToolExecutionMetadataKeys.TOOL_CALL_ID, toolCallId); + if (externalId != null) { + metadata.put(ToolExecutionMetadataKeys.EXTERNAL_ID, externalId); + } + ToolType toolType = tool == null ? null : tool.getToolType(); + if (toolType != null) { + metadata.put(ToolExecutionMetadataKeys.TOOL_TYPE, toolType.getValue()); + } + if (tool instanceof ToolExecutionMetadataProvider) { + Map extra; + try { + extra = ((ToolExecutionMetadataProvider) tool).getToolExecutionMetadata(parameters); + } catch (RuntimeException e) { + LOG.debug("Failed to collect execution metadata for tool {}.", toolName, e); + extra = Map.of(); + } + if (extra != null && !extra.isEmpty()) { + mergeSupplementalMetadata(metadata, extra); + } + } + return metadata; + } + + private static void mergeSupplementalMetadata( + Map target, Map supplemental) { + for (Map.Entry entry : supplemental.entrySet()) { + if (entry.getKey() != null && entry.getValue() != null) { + target.putIfAbsent(entry.getKey(), entry.getValue()); + } + } + } + private static Map resolveInjectedArguments(Tool tool, RunnerContext ctx) throws Exception { Map result = new HashMap<>(); diff --git a/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPServer.java b/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPServer.java index 47009c05a..6bfd95d2d 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPServer.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPServer.java @@ -26,6 +26,8 @@ import org.apache.flink.agents.api.resource.python.PythonResourceWrapper; import pemja.core.object.PyObject; +import javax.annotation.Nullable; + import java.util.ArrayList; import java.util.Collections; import java.util.List; @@ -55,12 +57,17 @@ public PythonMCPServer( @SuppressWarnings("unchecked") public List listTools() { + return listTools(null); + } + + @SuppressWarnings("unchecked") + public List listTools(@Nullable String mcpServerName) { Object result = adapter.callMethod(server, "list_tools", Collections.emptyMap()); if (result instanceof List) { List pythonTools = (List) result; List tools = new ArrayList<>(pythonTools.size()); for (Object pyTool : pythonTools) { - tools.add(new PythonMCPTool(adapter, (PyObject) pyTool)); + tools.add(new PythonMCPTool(adapter, (PyObject) pyTool, mcpServerName)); } return tools; } diff --git a/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPTool.java b/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPTool.java index 9cbd0c585..8a831db5a 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPTool.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPTool.java @@ -21,20 +21,27 @@ import org.apache.flink.agents.api.resource.python.PythonResourceAdapter; import org.apache.flink.agents.api.resource.python.PythonResourceWrapper; import org.apache.flink.agents.api.tools.Tool; +import org.apache.flink.agents.api.tools.ToolExecutionMetadataProvider; import org.apache.flink.agents.api.tools.ToolMetadata; import org.apache.flink.agents.api.tools.ToolParameters; import org.apache.flink.agents.api.tools.ToolResponse; import org.apache.flink.agents.api.tools.ToolType; +import org.apache.flink.agents.api.trace.ToolExecutionMetadataKeys; import pemja.core.object.PyObject; +import javax.annotation.Nullable; + import java.util.HashMap; +import java.util.LinkedHashMap; import java.util.Map; -public class PythonMCPTool extends Tool implements PythonResourceWrapper { +public class PythonMCPTool extends Tool + implements PythonResourceWrapper, ToolExecutionMetadataProvider { private static final String GET_JAVA_TOOL_META = "python_java_utils.get_java_tool_metadata_from_tool"; private final PyObject tool; private final PythonResourceAdapter adapter; + @Nullable private final String mcpServerName; /** * Creates a new PythonMCPServer. @@ -44,9 +51,22 @@ public class PythonMCPTool extends Tool implements PythonResourceWrapper { * @param tool The Python MCP tool object */ public PythonMCPTool(PythonResourceAdapter adapter, PyObject tool) { + this(adapter, tool, null); + } + + /** + * Creates a new Python MCP tool associated with a server resource name. + * + * @param adapter The Python resource adapter + * @param tool The Python MCP tool object + * @param mcpServerName The AgentPlan resource name of the MCP server that exposed this tool + */ + public PythonMCPTool( + PythonResourceAdapter adapter, PyObject tool, @Nullable String mcpServerName) { super(getToolMetadata(adapter, tool)); this.tool = tool; this.adapter = adapter; + this.mcpServerName = mcpServerName; } @SuppressWarnings("unchecked") @@ -91,4 +111,13 @@ public void setMetricGroup(FlinkAgentsMetricGroup metricGroup) { public ToolType getToolType() { return ToolType.MCP; } + + @Override + public Map getToolExecutionMetadata(ToolParameters parameters) { + Map metadata = new LinkedHashMap<>(); + if (mcpServerName != null) { + metadata.put(ToolExecutionMetadataKeys.MCP_SERVER, mcpServerName); + } + return metadata; + } } diff --git a/plan/src/main/java/org/apache/flink/agents/plan/serializer/AgentPlanJsonDeserializer.java b/plan/src/main/java/org/apache/flink/agents/plan/serializer/AgentPlanJsonDeserializer.java index a39503ec7..6fd139cb0 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/serializer/AgentPlanJsonDeserializer.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/serializer/AgentPlanJsonDeserializer.java @@ -51,6 +51,9 @@ public AgentPlan deserialize(JsonParser parser, DeserializationContext ctx) throws IOException, JacksonException { ObjectCodec codec = parser.getCodec(); JsonNode node = codec.readTree(parser); + JsonNode agentNameNode = node.get("agent_name"); + String agentName = + agentNameNode != null && !agentNameNode.isNull() ? agentNameNode.asText() : null; JsonNode actionsNode = node.get("actions"); // Deserialize actions @@ -133,6 +136,6 @@ public AgentPlan deserialize(JsonParser parser, DeserializationContext ctx) } AgentConfiguration config = new AgentConfiguration(configData); - return new AgentPlan(actions, actionsByEvent, resourceProviders, config); + return new AgentPlan(actions, actionsByEvent, resourceProviders, config, agentName); } } diff --git a/plan/src/main/java/org/apache/flink/agents/plan/serializer/AgentPlanJsonSerializer.java b/plan/src/main/java/org/apache/flink/agents/plan/serializer/AgentPlanJsonSerializer.java index 6b887f0c6..68ffe6fa3 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/serializer/AgentPlanJsonSerializer.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/serializer/AgentPlanJsonSerializer.java @@ -41,6 +41,8 @@ public void serialize( throws IOException { jsonGenerator.writeStartObject(); + jsonGenerator.writeStringField("agent_name", agentPlan.getAgentName()); + // Serialize actions jsonGenerator.writeFieldName("actions"); jsonGenerator.writeStartObject(); diff --git a/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanDeclareMCPServerTest.java b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanDeclareMCPServerTest.java index 70c8f7aa2..501bfd4d6 100644 --- a/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanDeclareMCPServerTest.java +++ b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanDeclareMCPServerTest.java @@ -259,6 +259,7 @@ void retrieveMCPToolAdd() throws Exception { MCPTool mcpTool = (MCPTool) tool; assertEquals("add", mcpTool.getName()); + assertEquals("testMcpServer", mcpTool.getMcpServerName()); // Verify description starts with expected text assertTrue( mcpTool.getMetadata() diff --git a/plan/src/test/java/org/apache/flink/agents/plan/FunctionToolPlanTest.java b/plan/src/test/java/org/apache/flink/agents/plan/FunctionToolPlanTest.java index 00f83da4f..8ff48fbd4 100644 --- a/plan/src/test/java/org/apache/flink/agents/plan/FunctionToolPlanTest.java +++ b/plan/src/test/java/org/apache/flink/agents/plan/FunctionToolPlanTest.java @@ -151,6 +151,23 @@ void javaFunctionMapping() throws Exception { assertEquals(4.0, (Double) ok.getResult(), 1e-9); } + @Test + @DisplayName("FunctionTool reports user exception message from Java function") + void javaFunctionToolReportsUserExceptionMessage() throws Exception { + Method m = + FunctionToolPlanTest.class.getMethod( + "calc", double.class, double.class, String.class); + FunctionTool tool = FunctionTool.fromStaticMethod("desc", m); + + ToolResponse error = + tool.call( + new ToolParameters( + new HashMap<>(Map.of("a", 10, "b", 0, "operation", "div")))); + + assertFalse(error.isSuccess()); + assertEquals("Division by zero", error.getError()); + } + @Test @DisplayName("Injected @ToolParam is hidden from metadata schema") void injectedToolParamHiddenFromMetadata() throws Exception { diff --git a/plan/src/test/java/org/apache/flink/agents/plan/TestFunction.java b/plan/src/test/java/org/apache/flink/agents/plan/TestFunction.java index 578d5b610..b72264b72 100644 --- a/plan/src/test/java/org/apache/flink/agents/plan/TestFunction.java +++ b/plan/src/test/java/org/apache/flink/agents/plan/TestFunction.java @@ -30,6 +30,10 @@ public static int add(int a, int b) { return a + b; } + public static void failWithUserException() { + throw new IllegalStateException("user failure from reflected function"); + } + @Test public void testJavaFunctionExecution() throws Exception { Function func = @@ -49,6 +53,16 @@ public void testJavaFunctionInitByClass() throws Exception { Assertions.assertEquals(3, result); } + @Test + public void testJavaFunctionUnwrapsInvocationTargetException() throws Exception { + Function func = + new JavaFunction(TestFunction.class, "failWithUserException", new Class[] {}); + + IllegalStateException exception = + Assertions.assertThrows(IllegalStateException.class, func::call); + Assertions.assertEquals("user failure from reflected function", exception.getMessage()); + } + public static void check_class(InputEvent a, OutputEvent b) {} @Test diff --git a/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRetryTest.java b/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRetryTest.java index 8c1395059..a978f0d59 100644 --- a/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRetryTest.java +++ b/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRetryTest.java @@ -31,6 +31,7 @@ import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup; import org.apache.flink.agents.api.resource.ResourceType; import org.apache.flink.agents.api.tools.ToolResponse; +import org.apache.flink.agents.api.trace.ExecutionReporter; import org.apache.flink.metrics.Counter; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -127,6 +128,129 @@ void chatSucceedsWithoutRetry_retryCountIsZero() throws Exception { verify(mockActionMetricGroup, never()).getSubGroup(anyString(), anyString()); } + @Test + void chatReportsLlmExecution() throws Exception { + RunnerContext reportingCtx = reportingRunnerContext(); + BaseChatModelSetup chatModel = configureReportingChatContext(reportingCtx); + when(chatModel.chat(any(), any(), any())) + .thenReturn(new ChatMessage(MessageRole.ASSISTANT, "hello")); + + ChatModelAction.chat( + UUID.randomUUID(), + "test-model", + List.of(new ChatMessage(MessageRole.USER, "hi")), + Map.of(), + null, + reportingCtx); + + ExecutionReporter reporter = (ExecutionReporter) reportingCtx; + verify(reporter) + .reportExecutionStarted(ExecutionReporter.EntityTypes.LLM, "test-model", Map.of()); + verify(reporter) + .reportExecutionSucceeded( + ExecutionReporter.EntityTypes.LLM, "test-model", Map.of()); + } + + @Test + void chatReportsStructuredOutputParserExecution() throws Exception { + RunnerContext reportingCtx = reportingRunnerContext(); + BaseChatModelSetup chatModel = configureReportingChatContext(reportingCtx); + when(chatModel.chat(any(), any(), any())) + .thenReturn(new ChatMessage(MessageRole.ASSISTANT, "{\"answer\":\"42\"}")); + + ChatModelAction.chat( + UUID.randomUUID(), + "test-model", + List.of(new ChatMessage(MessageRole.USER, "hi")), + Map.of(), + Map.class, + reportingCtx); + + ExecutionReporter reporter = (ExecutionReporter) reportingCtx; + verify(reporter) + .reportExecutionStarted( + ExecutionReporter.EntityTypes.PARSER, Agent.STRUCTURED_OUTPUT, Map.of()); + verify(reporter) + .reportExecutionSucceeded( + ExecutionReporter.EntityTypes.PARSER, Agent.STRUCTURED_OUTPUT, Map.of()); + } + + @Test + void chatRetriesStructuredOutputParseErrorWithoutFailingLlm() throws Exception { + RunnerContext reportingCtx = reportingRunnerContext(); + BaseChatModelSetup chatModel = configureReportingChatContext(reportingCtx); + when(reportingCtx.getConfig()) + .thenReturn(readableConfig(Agent.ErrorHandlingStrategy.RETRY, 1, 0)); + when(chatModel.chat(any(), any(), any())) + .thenReturn( + new ChatMessage(MessageRole.ASSISTANT, "not-json"), + new ChatMessage(MessageRole.ASSISTANT, "{\"answer\":\"42\"}")); + + ChatModelAction.chat( + UUID.randomUUID(), + "test-model", + List.of(new ChatMessage(MessageRole.USER, "hi")), + Map.of(), + Map.class, + reportingCtx); + + verify(chatModel, times(2)).chat(any(), any(), any()); + + ExecutionReporter reporter = (ExecutionReporter) reportingCtx; + verify(reporter, times(2)) + .reportExecutionStarted(ExecutionReporter.EntityTypes.LLM, "test-model", Map.of()); + verify(reporter, times(2)) + .reportExecutionSucceeded( + ExecutionReporter.EntityTypes.LLM, "test-model", Map.of()); + verify(reporter) + .reportExecutionFailed( + eq(ExecutionReporter.EntityTypes.PARSER), + eq(Agent.STRUCTURED_OUTPUT), + eq(Map.of()), + any(Exception.class), + eq(ExecutionReporter.ProblemCategories.MODEL_OUTPUT_PARSE_ERROR)); + verify(reporter, never()) + .reportExecutionFailed( + eq(ExecutionReporter.EntityTypes.LLM), + eq("test-model"), + eq(Map.of()), + any(Throwable.class), + any()); + } + + @Test + void chatReportsEachRetriedModelInvocation() throws Exception { + RunnerContext reportingCtx = reportingRunnerContext(); + BaseChatModelSetup chatModel = configureReportingChatContext(reportingCtx); + when(reportingCtx.getConfig()) + .thenReturn(readableConfig(Agent.ErrorHandlingStrategy.RETRY, 1, 0)); + when(chatModel.chat(any(), any(), any())) + .thenThrow(new RuntimeException("transient error")) + .thenReturn(new ChatMessage(MessageRole.ASSISTANT, "success")); + + ChatModelAction.chat( + UUID.randomUUID(), + "test-model", + List.of(new ChatMessage(MessageRole.USER, "hi")), + Map.of(), + null, + reportingCtx); + + ExecutionReporter reporter = (ExecutionReporter) reportingCtx; + verify(reporter, times(2)) + .reportExecutionStarted(ExecutionReporter.EntityTypes.LLM, "test-model", Map.of()); + verify(reporter) + .reportExecutionFailed( + eq(ExecutionReporter.EntityTypes.LLM), + eq("test-model"), + eq(Map.of()), + any(RuntimeException.class), + eq(ExecutionReporter.ProblemCategories.MODEL_CALL_FAILED)); + verify(reporter) + .reportExecutionSucceeded( + ExecutionReporter.EntityTypes.LLM, "test-model", Map.of()); + } + @Test void chatRetriesWithExponentialBackoff() throws Exception { // 1 second base interval; fail once then succeed -> wait 1s (1 * 2^0) @@ -327,6 +451,87 @@ public String getStr(String key, String defaultValue) { }); } + private RunnerContext reportingRunnerContext() { + return mock(RunnerContext.class, withSettings().extraInterfaces(ExecutionReporter.class)); + } + + private BaseChatModelSetup configureReportingChatContext(RunnerContext reportingCtx) + throws Exception { + BaseChatModelSetup chatModel = mock(BaseChatModelSetup.class); + MemoryObject memory = createStatefulMemoryObject(); + + when(chatModel.getConnectionName()).thenReturn("test-connection"); + when(reportingCtx.getResource(anyString(), eq(ResourceType.CHAT_MODEL))) + .thenReturn(chatModel); + when(reportingCtx.getSensoryMemory()).thenReturn(memory); + when(reportingCtx.getActionMetricGroup()).thenReturn(mockActionMetricGroup); + when(reportingCtx.durableExecute(any())) + .thenAnswer(inv -> inv.>getArgument(0).call()); + doAnswer(inv -> sentEvents.add(inv.getArgument(0))).when(reportingCtx).sendEvent(any()); + when(reportingCtx.getConfig()).thenReturn(readableConfig(Agent.ErrorHandlingStrategy.FAIL)); + return chatModel; + } + + private org.apache.flink.agents.api.configuration.ReadableConfiguration readableConfig( + Agent.ErrorHandlingStrategy errorHandlingStrategy) { + return readableConfig(errorHandlingStrategy, 0, 0); + } + + private org.apache.flink.agents.api.configuration.ReadableConfiguration readableConfig( + Agent.ErrorHandlingStrategy errorHandlingStrategy, + int maxRetries, + int retryWaitIntervalSec) { + return new org.apache.flink.agents.api.configuration.ReadableConfiguration() { + @Override + @SuppressWarnings("unchecked") + public T get(org.apache.flink.agents.api.configuration.ConfigOption option) { + if (option == AgentExecutionOptions.ERROR_HANDLING_STRATEGY) { + return (T) errorHandlingStrategy; + } + if (option == AgentExecutionOptions.MAX_RETRIES) { + return (T) Integer.valueOf(maxRetries); + } + if (option == AgentExecutionOptions.RETRY_WAIT_INTERVAL) { + return (T) Integer.valueOf(retryWaitIntervalSec); + } + if (option == AgentExecutionOptions.CHAT_ASYNC) { + return (T) Boolean.FALSE; + } + return option.getDefaultValue(); + } + + @Override + public Integer getInt(String key, Integer defaultValue) { + return defaultValue; + } + + @Override + public Long getLong(String key, Long defaultValue) { + return defaultValue; + } + + @Override + public Float getFloat(String key, Float defaultValue) { + return defaultValue; + } + + @Override + public Double getDouble(String key, Double defaultValue) { + return defaultValue; + } + + @Override + public Boolean getBool(String key, Boolean defaultValue) { + return defaultValue; + } + + @Override + public String getStr(String key, String defaultValue) { + return defaultValue; + } + }; + } + /** * Creates a stateful MemoryObject backed by a HashMap, supporting isExist/get/set operations * needed by the retry stats accumulation logic. diff --git a/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionReportTest.java b/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionReportTest.java new file mode 100644 index 000000000..08348327b --- /dev/null +++ b/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionReportTest.java @@ -0,0 +1,299 @@ +/* + * 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.plan.actions; + +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.agents.AgentExecutionOptions; +import org.apache.flink.agents.api.context.DurableCallable; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.event.ToolRequestEvent; +import org.apache.flink.agents.api.event.ToolResponseEvent; +import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.api.tools.Tool; +import org.apache.flink.agents.api.tools.ToolExecutionMetadataProvider; +import org.apache.flink.agents.api.tools.ToolMetadata; +import org.apache.flink.agents.api.tools.ToolParameters; +import org.apache.flink.agents.api.tools.ToolResponse; +import org.apache.flink.agents.api.tools.ToolType; +import org.apache.flink.agents.api.trace.ExecutionReporter; +import org.apache.flink.agents.api.trace.ToolExecutionMetadataKeys; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.mockito.Mockito.withSettings; + +/** Tests for tool-call execution reports. */ +class ToolCallActionReportTest { + + @Test + void processToolRequestReportsEachToolCall() throws Exception { + RunnerContext ctx = + mock(RunnerContext.class, withSettings().extraInterfaces(ExecutionReporter.class)); + List sentEvents = new ArrayList<>(); + Tool tool = + new ReportingTool( + ToolType.MCP, + Map.of(ToolExecutionMetadataKeys.MCP_SERVER, "search-server"), + ToolResponse.success("ok")); + when(ctx.getResource("search", ResourceType.TOOL)).thenReturn(tool); + when(ctx.getConfig()).thenReturn(toolCallConfig()); + when(ctx.durableExecute(any())) + .thenAnswer(inv -> inv.>getArgument(0).call()); + doAnswer(inv -> sentEvents.add(inv.getArgument(0))).when(ctx).sendEvent(any()); + + Map function = new LinkedHashMap<>(); + function.put("name", "search"); + function.put("arguments", Map.of("query", "flink")); + Map toolCall = new LinkedHashMap<>(); + toolCall.put("id", "call-1"); + toolCall.put("original_id", "external-call-1"); + toolCall.put("function", function); + ToolRequestEvent request = new ToolRequestEvent("test-model", List.of(toolCall)); + + ToolCallAction.processToolRequest(request, ctx); + + Map metadata = new LinkedHashMap<>(); + metadata.put(ToolExecutionMetadataKeys.TOOL_REQUEST_EVENT_ID, request.getId().toString()); + metadata.put(ToolExecutionMetadataKeys.TOOL_CALL_ID, "call-1"); + metadata.put(ToolExecutionMetadataKeys.EXTERNAL_ID, "external-call-1"); + metadata.put(ToolExecutionMetadataKeys.TOOL_TYPE, ToolType.MCP.getValue()); + metadata.put(ToolExecutionMetadataKeys.MCP_SERVER, "search-server"); + ExecutionReporter reporter = (ExecutionReporter) ctx; + verify(reporter) + .reportExecutionStarted(ExecutionReporter.EntityTypes.TOOL, "search", metadata); + verify(reporter) + .reportExecutionSucceeded(ExecutionReporter.EntityTypes.TOOL, "search", metadata); + + assertThat(sentEvents).hasSize(1); + assertThat(sentEvents.get(0)).isInstanceOf(ToolResponseEvent.class); + } + + @Test + void processToolRequestMarksErrorResponseAsFailed() throws Exception { + RunnerContext ctx = + mock(RunnerContext.class, withSettings().extraInterfaces(ExecutionReporter.class)); + List sentEvents = new ArrayList<>(); + Tool tool = mock(Tool.class); + when(tool.call(any())).thenReturn(ToolResponse.error("tool rejected request")); + when(ctx.getResource("search", ResourceType.TOOL)).thenReturn(tool); + when(ctx.getConfig()).thenReturn(toolCallConfig()); + when(ctx.durableExecute(any())) + .thenAnswer(inv -> inv.>getArgument(0).call()); + doAnswer(inv -> sentEvents.add(inv.getArgument(0))).when(ctx).sendEvent(any()); + + Map function = new LinkedHashMap<>(); + function.put("name", "search"); + function.put("arguments", Map.of("query", "flink")); + Map toolCall = new LinkedHashMap<>(); + toolCall.put("id", "call-1"); + toolCall.put("function", function); + ToolRequestEvent request = new ToolRequestEvent("test-model", List.of(toolCall)); + + ToolCallAction.processToolRequest(request, ctx); + + Map metadata = new LinkedHashMap<>(); + metadata.put(ToolExecutionMetadataKeys.TOOL_REQUEST_EVENT_ID, request.getId().toString()); + metadata.put(ToolExecutionMetadataKeys.TOOL_CALL_ID, "call-1"); + ExecutionReporter reporter = (ExecutionReporter) ctx; + verify(reporter) + .reportExecutionFailed( + eq(ExecutionReporter.EntityTypes.TOOL), + eq("search"), + eq(metadata), + any(Throwable.class), + eq(ExecutionReporter.ProblemCategories.TOOL_CALL_FAILED)); + verify(reporter, never()) + .reportExecutionSucceeded(ExecutionReporter.EntityTypes.TOOL, "search", metadata); + + ToolResponseEvent responseEvent = (ToolResponseEvent) sentEvents.get(0); + assertThat(responseEvent.getSuccess()).containsEntry("call-1", false); + assertThat(responseEvent.getError()).containsEntry("call-1", "tool rejected request"); + } + + @Test + void processLoadSkillToolRequestAddsSkillMetadata() throws Exception { + RunnerContext ctx = + mock(RunnerContext.class, withSettings().extraInterfaces(ExecutionReporter.class)); + Tool tool = + new ReportingTool( + ToolType.FUNCTION, + Map.of( + ToolExecutionMetadataKeys.SKILL_NAME, + "math-calculator", + ToolExecutionMetadataKeys.SKILL_RESOURCE_PATH, + "README.md"), + ToolResponse.success("skill content")); + when(ctx.getResource("load_skill", ResourceType.TOOL)).thenReturn(tool); + when(ctx.getConfig()).thenReturn(toolCallConfig()); + when(ctx.durableExecute(any())) + .thenAnswer(inv -> inv.>getArgument(0).call()); + + Map function = new LinkedHashMap<>(); + function.put("name", "load_skill"); + function.put("arguments", Map.of("name", "math-calculator", "path", "README.md")); + Map toolCall = new LinkedHashMap<>(); + toolCall.put("id", "call-1"); + toolCall.put("function", function); + ToolRequestEvent request = new ToolRequestEvent("test-model", List.of(toolCall)); + + ToolCallAction.processToolRequest(request, ctx); + + Map metadata = new LinkedHashMap<>(); + metadata.put(ToolExecutionMetadataKeys.TOOL_REQUEST_EVENT_ID, request.getId().toString()); + metadata.put(ToolExecutionMetadataKeys.TOOL_CALL_ID, "call-1"); + metadata.put(ToolExecutionMetadataKeys.TOOL_TYPE, ToolType.FUNCTION.getValue()); + metadata.put(ToolExecutionMetadataKeys.SKILL_NAME, "math-calculator"); + metadata.put(ToolExecutionMetadataKeys.SKILL_RESOURCE_PATH, "README.md"); + verify((ExecutionReporter) ctx) + .reportExecutionStarted(ExecutionReporter.EntityTypes.TOOL, "load_skill", metadata); + } + + @Test + void executionMetadataCannotMutateToolCallParameters() throws Exception { + RunnerContext ctx = + mock(RunnerContext.class, withSettings().extraInterfaces(ExecutionReporter.class)); + List sentEvents = new ArrayList<>(); + when(ctx.getResource("search", ResourceType.TOOL)).thenReturn(new MutatingMetadataTool()); + when(ctx.getConfig()).thenReturn(toolCallConfig()); + when(ctx.durableExecute(any())) + .thenAnswer(inv -> inv.>getArgument(0).call()); + doAnswer(inv -> sentEvents.add(inv.getArgument(0))).when(ctx).sendEvent(any()); + + Map function = new LinkedHashMap<>(); + function.put("name", "search"); + function.put("arguments", Map.of("query", "flink")); + Map toolCall = new LinkedHashMap<>(); + toolCall.put("id", "call-1"); + toolCall.put("function", function); + + ToolCallAction.processToolRequest( + new ToolRequestEvent("test-model", List.of(toolCall)), ctx); + + ToolResponseEvent responseEvent = (ToolResponseEvent) sentEvents.get(0); + assertThat(responseEvent.getResponses().get("call-1").getResult()).isEqualTo("flink"); + } + + private static org.apache.flink.agents.api.configuration.ReadableConfiguration + toolCallConfig() { + return new org.apache.flink.agents.api.configuration.ReadableConfiguration() { + @Override + @SuppressWarnings("unchecked") + public T get(org.apache.flink.agents.api.configuration.ConfigOption option) { + if (option == AgentExecutionOptions.TOOL_CALL_ASYNC) { + return (T) Boolean.FALSE; + } + return option.getDefaultValue(); + } + + @Override + public Integer getInt(String key, Integer defaultValue) { + return defaultValue; + } + + @Override + public Long getLong(String key, Long defaultValue) { + return defaultValue; + } + + @Override + public Float getFloat(String key, Float defaultValue) { + return defaultValue; + } + + @Override + public Double getDouble(String key, Double defaultValue) { + return defaultValue; + } + + @Override + public Boolean getBool(String key, Boolean defaultValue) { + return defaultValue; + } + + @Override + public String getStr(String key, String defaultValue) { + return defaultValue; + } + }; + } + + private static final class ReportingTool extends Tool implements ToolExecutionMetadataProvider { + private final ToolType toolType; + private final Map entityMetadata; + private final ToolResponse response; + + private ReportingTool( + ToolType toolType, Map entityMetadata, ToolResponse response) { + super(new ToolMetadata("search", "Search", "{\"type\":\"object\"}")); + this.toolType = toolType; + this.entityMetadata = entityMetadata; + this.response = response; + } + + @Override + public ToolType getToolType() { + return toolType; + } + + @Override + public ToolResponse call(ToolParameters parameters) { + return response; + } + + @Override + public Map getToolExecutionMetadata(ToolParameters parameters) { + return entityMetadata; + } + } + + private static final class MutatingMetadataTool extends Tool + implements ToolExecutionMetadataProvider { + + private MutatingMetadataTool() { + super(new ToolMetadata("search", "Search", "{\"type\":\"object\"}")); + } + + @Override + public ToolType getToolType() { + return ToolType.FUNCTION; + } + + @Override + public ToolResponse call(ToolParameters parameters) { + return ToolResponse.success(parameters.getParameter("query")); + } + + @Override + public Map getToolExecutionMetadata(ToolParameters parameters) { + parameters.addParameter("query", "mutated"); + return Map.of(); + } + } +} diff --git a/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionTest.java b/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionTest.java index 750f12adc..e1c62556f 100644 --- a/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionTest.java +++ b/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionTest.java @@ -378,6 +378,22 @@ public Resource getResource(String name, ResourceType type) { assertThat(response.getError()).containsEntry("call-1", "missing resource"); } + @Test + void processToolRequestUsesFallbackWhenMissingToolErrorHasNoMessage() { + FakeRunnerContext ctx = + new FakeRunnerContext() { + @Override + public Resource getResource(String name, ResourceType type) { + throw new RuntimeException(); + } + }; + + ToolCallAction.processToolRequest(toolRequest("queryOrder"), ctx); + + ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(response.getError()).containsEntry("call-1", "Tool does not exist."); + } + private static ToolRequestEvent toolRequest(String toolName) { return new ToolRequestEvent( "model", diff --git a/plan/src/test/java/org/apache/flink/agents/plan/serializer/AgentPlanJsonSerializerTest.java b/plan/src/test/java/org/apache/flink/agents/plan/serializer/AgentPlanJsonSerializerTest.java index ef18cdadc..139ffa96f 100644 --- a/plan/src/test/java/org/apache/flink/agents/plan/serializer/AgentPlanJsonSerializerTest.java +++ b/plan/src/test/java/org/apache/flink/agents/plan/serializer/AgentPlanJsonSerializerTest.java @@ -249,6 +249,20 @@ public void testSerializeAgentPlanCreatedFromAgent() throws Exception { // Verify that config data from AgentConfiguration is present assertThat(json).contains("\"conf_data\":{\"config.key\":\"config.value\"}"); + assertThat(json).contains("\"agent_name\":\"TestAgent\""); + } + + @Test + public void testSerializeDeserializeAgentName() throws Exception { + AgentPlan original = + new AgentPlan(new TestAgent(), new AgentConfiguration(), "named-agent"); + + ObjectMapper mapper = new ObjectMapper(); + String json = mapper.writeValueAsString(original); + AgentPlan deserialized = mapper.readValue(json, AgentPlan.class); + + assertThat(json).contains("\"agent_name\":\"named-agent\""); + assertThat(deserialized.getAgentName()).isEqualTo("named-agent"); } @Test diff --git a/python/flink_agents/api/core_options.py b/python/flink_agents/api/core_options.py index d196777c4..df888fc8a 100644 --- a/python/flink_agents/api/core_options.py +++ b/python/flink_agents/api/core_options.py @@ -190,6 +190,12 @@ class AgentConfigOptions: default=EventLogLevel.STANDARD, ) + EVENT_LOG_TRACE_ENABLED = ConfigOption( + key="event-log.trace.enabled", + config_type=bool, + default=False, + ) + EVENT_LOG_MAX_STRING_LENGTH = ConfigOption( key="event-log.standard.max-string-length", config_type=int, diff --git a/python/flink_agents/api/events/event.py b/python/flink_agents/api/events/event.py index 77e7e5faf..013f98f01 100644 --- a/python/flink_agents/api/events/event.py +++ b/python/flink_agents/api/events/event.py @@ -15,7 +15,6 @@ # See the License for the specific language governing permissions and # limitations under the License. ################################################################################# -import hashlib import json from typing import Any, ClassVar, Dict @@ -23,7 +22,7 @@ from typing import override except ImportError: from typing_extensions import override -from uuid import UUID +from uuid import UUID, uuid4 from pydantic import BaseModel, Field, model_validator from pydantic_core import PydanticSerializationError @@ -66,15 +65,14 @@ class Event(BaseModel, extra="allow"): Attributes: ---------- id : UUID - Unique identifier for the event, generated deterministically based on - event content. + Unique identifier for the event. type : str Event type string used for routing. Required for all events. attributes : Dict[str, Any] Key-value properties for the event data. """ - id: UUID = Field(default=None) + id: UUID = Field(default_factory=uuid4) type: str attributes: Dict[str, Any] = Field(default_factory=dict) @@ -98,23 +96,11 @@ 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_validator(mode="after") def validate_and_set_id(self) -> "Event": - """Validate that fields are serializable and generate content-based ID.""" + """Validate that fields are serializable and ensure the event has an ID.""" if self.id is None: - object.__setattr__(self, "id", self._generate_content_based_id()) + object.__setattr__(self, "id", uuid4()) self.model_dump_json() return self @@ -122,9 +108,6 @@ 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 get_type(self) -> str: """Return the event type string used for routing.""" diff --git a/python/flink_agents/api/tests/test_core_options.py b/python/flink_agents/api/tests/test_core_options.py index f1e773a83..7ee111dec 100644 --- a/python/flink_agents/api/tests/test_core_options.py +++ b/python/flink_agents/api/tests/test_core_options.py @@ -101,10 +101,11 @@ def test_agent_config_options_are_explicitly_declared() -> None: from flink_agents.api.core_options import AgentConfigOptions, EventLogLevel options = _collect_config_options(AgentConfigOptions) - assert len(options) == 23 + assert len(options) == 24 assert options["BASE_LOG_DIR"].get_key() == "baseLogDir" assert options["KAFKA_BOOTSTRAP_SERVERS"].get_default_value() == "localhost:9092" assert options["EVENT_LOG_LEVEL"].get_default_value() is EventLogLevel.STANDARD + assert options["EVENT_LOG_TRACE_ENABLED"].get_default_value() is False def test_unknown_agent_config_option_raises_attribute_error() -> None: diff --git a/python/flink_agents/api/tests/test_event.py b/python/flink_agents/api/tests/test_event.py index 110106ea4..1884a482b 100644 --- a/python/flink_agents/api/tests/test_event.py +++ b/python/flink_agents/api/tests/test_event.py @@ -40,6 +40,22 @@ def test_event_setattr_serializable() -> None: event.c = Event(type="nested") +def test_same_content_creates_distinct_event_occurrences() -> None: + first = Event(type="test", a=1) + second = Event(type="test", a=1) + + assert first.id != second.id + + +def test_mutating_event_payload_preserves_occurrence_id() -> None: + event = Event(type="test", a=1) + event_id = event.id + + event.a = 2 + + assert event.id == event_id + + def test_event_setattr_non_serializable() -> None: event = Event(type="test", a=1) with pytest.raises(PydanticSerializationError): diff --git a/python/flink_agents/api/tests/test_execution_reporter.py b/python/flink_agents/api/tests/test_execution_reporter.py new file mode 100644 index 000000000..a31810484 --- /dev/null +++ b/python/flink_agents/api/tests/test_execution_reporter.py @@ -0,0 +1,67 @@ +################################################################################ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +################################################################################# +from unittest.mock import MagicMock + +from flink_agents.api.trace import ( + ExecutionEntityTypes, + ExecutionProblemCategories, + ExecutionReporter, + ExecutionReporters, +) + + +def test_failed_reporter_uses_metadata_before_error() -> None: + ctx = MagicMock(spec=ExecutionReporter) + metadata = {"toolCallId": "call-1"} + error = RuntimeError("boom") + + ExecutionReporters.failed( + ctx, + ExecutionEntityTypes.TOOL, + "search", + metadata, + error, + ExecutionProblemCategories.TOOL_CALL_FAILED, + ) + + ctx.report_execution_failed.assert_called_once_with( + ExecutionEntityTypes.TOOL, + "search", + metadata, + error, + ExecutionProblemCategories.TOOL_CALL_FAILED, + ) + + +def test_reporters_ignore_context_without_execution_reporter() -> None: + ctx = MagicMock() + + ExecutionReporters.started(ctx, ExecutionEntityTypes.LLM, "model") + ExecutionReporters.succeeded(ctx, ExecutionEntityTypes.LLM, "model") + ExecutionReporters.failed( + ctx, + ExecutionEntityTypes.LLM, + "model", + None, + RuntimeError("boom"), + ExecutionProblemCategories.MODEL_CALL_FAILED, + ) + + ctx.report_execution_started.assert_not_called() + ctx.report_execution_succeeded.assert_not_called() + ctx.report_execution_failed.assert_not_called() diff --git a/python/flink_agents/api/tools/__init__.py b/python/flink_agents/api/tools/__init__.py index 1c6f61a92..78d07b8ba 100644 --- a/python/flink_agents/api/tools/__init__.py +++ b/python/flink_agents/api/tools/__init__.py @@ -15,9 +15,16 @@ # See the License for the specific language governing permissions and # limitations under the License. ################################################################################# +from flink_agents.api.tools.tool_execution_metadata_provider import ( + ToolExecutionMetadataProvider, +) from flink_agents.api.tools.tool_parameter_injection import ( InjectedArg, ToolParameterSource, ) -__all__ = ["InjectedArg", "ToolParameterSource"] +__all__ = [ + "InjectedArg", + "ToolExecutionMetadataProvider", + "ToolParameterSource", +] diff --git a/python/flink_agents/api/tools/tool_execution_metadata_provider.py b/python/flink_agents/api/tools/tool_execution_metadata_provider.py new file mode 100644 index 000000000..a326a9435 --- /dev/null +++ b/python/flink_agents/api/tools/tool_execution_metadata_provider.py @@ -0,0 +1,31 @@ +################################################################################ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +################################################################################# +from abc import ABC, abstractmethod +from collections.abc import Mapping +from typing import Any + + +class ToolExecutionMetadataProvider(ABC): + """Optional capability for supplying stable Tool execution metadata.""" + + @abstractmethod + def get_tool_execution_metadata( + self, parameters: Mapping[str, Any] + ) -> Mapping[str, Any]: + """Return metadata associated with one Tool execution.""" diff --git a/python/flink_agents/api/trace/__init__.py b/python/flink_agents/api/trace/__init__.py new file mode 100644 index 000000000..b40d8f177 --- /dev/null +++ b/python/flink_agents/api/trace/__init__.py @@ -0,0 +1,34 @@ +################################################################################ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +################################################################################# +from flink_agents.api.trace.execution_reporter import ( + ExecutionEntityTypes, + ExecutionProblemCategories, + ExecutionReporter, + ExecutionReporters, +) +from flink_agents.api.trace.tool_execution_metadata_keys import ( + ToolExecutionMetadataKeys, +) + +__all__ = [ + "ExecutionEntityTypes", + "ExecutionProblemCategories", + "ExecutionReporter", + "ExecutionReporters", + "ToolExecutionMetadataKeys", +] diff --git a/python/flink_agents/api/trace/execution_reporter.py b/python/flink_agents/api/trace/execution_reporter.py new file mode 100644 index 000000000..926f1ae39 --- /dev/null +++ b/python/flink_agents/api/trace/execution_reporter.py @@ -0,0 +1,145 @@ +################################################################################ +# 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 logging +from abc import ABC, abstractmethod +from collections.abc import Callable, Mapping +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from flink_agents.api.runner_context import RunnerContext + +logger = logging.getLogger(__name__) + +_EMPTY_METADATA: Mapping[str, Any] = {} + + +class ExecutionEntityTypes: + """Shared entity type names for execution reports.""" + + ACTION = "action" + LLM = "llm" + PARSER = "parser" + TOOL = "tool" + + +class ExecutionProblemCategories: + """Shared low-cardinality problem categories for failed execution reports.""" + + ACTION_EXECUTION_FAILED = "action_execution_failed" + MODEL_CALL_FAILED = "model_call_failed" + MODEL_OUTPUT_PARSE_ERROR = "model_output_parse_error" + TOOL_CALL_FAILED = "tool_call_failed" + + +class ExecutionReporter(ABC): + """Optional capability for reporting executions nested inside an action.""" + + @abstractmethod + def report_execution_started( + self, + entity_type: str, + entity_name: str, + entity_metadata: Mapping[str, Any] | None = None, + ) -> None: + """Report that a logical execution started.""" + + @abstractmethod + def report_execution_succeeded( + self, + entity_type: str, + entity_name: str, + entity_metadata: Mapping[str, Any] | None = None, + ) -> None: + """Report that a logical execution completed successfully.""" + + @abstractmethod + def report_execution_failed( + self, + entity_type: str, + entity_name: str, + entity_metadata: Mapping[str, Any] | None, + error: BaseException, + problem_category: str | None = None, + ) -> None: + """Report that a logical execution failed.""" + + +class ExecutionReporters: + """Best-effort helpers for contexts that implement ExecutionReporter.""" + + @staticmethod + def started( + ctx: "RunnerContext", + entity_type: str, + entity_name: str, + entity_metadata: Mapping[str, Any] | None = None, + ) -> None: + """Report the start of a nested execution if the context supports it.""" + ExecutionReporters._report( + ctx, + lambda reporter: reporter.report_execution_started( + entity_type, entity_name, entity_metadata or _EMPTY_METADATA + ), + ) + + @staticmethod + def succeeded( + ctx: "RunnerContext", + entity_type: str, + entity_name: str, + entity_metadata: Mapping[str, Any] | None = None, + ) -> None: + """Report successful completion if the context supports it.""" + ExecutionReporters._report( + ctx, + lambda reporter: reporter.report_execution_succeeded( + entity_type, entity_name, entity_metadata or _EMPTY_METADATA + ), + ) + + @staticmethod + def failed( + ctx: "RunnerContext", + entity_type: str, + entity_name: str, + entity_metadata: Mapping[str, Any] | None, + error: BaseException, + problem_category: str | None = None, + ) -> None: + """Report failed completion if the context supports it.""" + ExecutionReporters._report( + ctx, + lambda reporter: reporter.report_execution_failed( + entity_type, + entity_name, + entity_metadata or _EMPTY_METADATA, + error, + problem_category, + ), + ) + + @staticmethod + def _report( + ctx: "RunnerContext", report: Callable[[ExecutionReporter], None] + ) -> None: + if not isinstance(ctx, ExecutionReporter): + return + try: + report(ctx) + except Exception: + logger.debug("Execution reporting failed and was ignored.", exc_info=True) diff --git a/python/flink_agents/api/trace/tool_execution_metadata_keys.py b/python/flink_agents/api/trace/tool_execution_metadata_keys.py new file mode 100644 index 000000000..1308cb7be --- /dev/null +++ b/python/flink_agents/api/trace/tool_execution_metadata_keys.py @@ -0,0 +1,30 @@ +################################################################################ +# 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. +################################################################################# + + +class ToolExecutionMetadataKeys: + """Shared field names for Tool execution entity metadata.""" + + TOOL_REQUEST_EVENT_ID = "toolRequestEventId" + TOOL_CALL_ID = "toolCallId" + EXTERNAL_ID = "externalId" + TOOL_TYPE = "toolType" + MCP_SERVER = "mcpServer" + SKILL_NAME = "skillName" + SKILL_RESOURCE_PATH = "skillResourcePath" 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..f5e81e944 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 @@ -117,7 +117,9 @@ def test_python_event_logging(tmp_path: Path) -> None: if line.strip(): record = json.loads(line) record_line = line - event_payload = record.get("event", {}) + event_payload = record.get("eventAttributes") + if event_payload is None: + event_payload = record.get("event", {}).get("attributes", {}) if "processed_review" in json.dumps(event_payload): has_processed_review = True break @@ -128,18 +130,18 @@ def test_python_event_logging(tmp_path: Path) -> None: assert record_line is not None, "Event log file is empty." assert "timestamp" in record assert "logLevel" in record + assert "eventId" in record assert "eventType" in record - assert "event" in record - assert "eventType" in record["event"] + assert "eventAttributes" in record assert has_processed_review, "Log should contain processed review content" + event_id_idx = record_line.find('"eventId"') event_type_idx = record_line.find('"eventType"') - id_idx = record_line.find('"id"') - attributes_idx = record_line.find('"attributes"') + attributes_idx = record_line.find('"eventAttributes"') + assert event_id_idx != -1 assert event_type_idx != -1 - assert id_idx != -1 assert attributes_idx != -1 - assert event_type_idx < id_idx < attributes_idx + assert event_id_idx < event_type_idx < attributes_idx def _run_event_logging_pipeline( diff --git a/python/flink_agents/e2e_tests/test_utils.py b/python/flink_agents/e2e_tests/test_utils.py index 68ef8009c..3f0fba07b 100644 --- a/python/flink_agents/e2e_tests/test_utils.py +++ b/python/flink_agents/e2e_tests/test_utils.py @@ -64,9 +64,13 @@ def collect_tool_invocations(log_dir: str | Path) -> list[dict]: if not line.strip(): continue record = json.loads(line) - if record.get("eventType") != "_tool_request_event": + event_type = record.get("eventType") + if event_type != "_tool_request_event": continue - tool_calls = record["event"]["attributes"].get("tool_calls", []) + attributes = record.get("eventAttributes") + if attributes is None: + attributes = record["event"].get("attributes", {}) + tool_calls = attributes.get("tool_calls", []) for tool_call in tool_calls: function = tool_call["function"] invocations.append( diff --git a/python/flink_agents/e2e_tests/test_utils_unit_test.py b/python/flink_agents/e2e_tests/test_utils_unit_test.py index f15089a4f..ec446f311 100644 --- a/python/flink_agents/e2e_tests/test_utils_unit_test.py +++ b/python/flink_agents/e2e_tests/test_utils_unit_test.py @@ -31,12 +31,9 @@ def _tool_request_record(tool_calls: list) -> dict: return { "timestamp": "2026-05-31T00:00:00Z", "logLevel": "STANDARD", + "eventId": "00000000-0000-0000-0000-000000000001", "eventType": "_tool_request_event", - "event": { - "eventType": "_tool_request_event", - "id": "00000000-0000-0000-0000-000000000001", - "attributes": {"model": "qwen3:1.7b", "tool_calls": tool_calls}, - }, + "eventAttributes": {"model": "qwen3:1.7b", "tool_calls": tool_calls}, } @@ -64,8 +61,9 @@ def test_collect_tool_invocations(tmp_path: Path) -> None: other_event = { "timestamp": "2026-05-31T00:00:00Z", "logLevel": "STANDARD", + "eventId": "00000000-0000-0000-0000-000000000002", "eventType": "_input_event", - "event": {"eventType": "_input_event", "id": "x", "attributes": {}}, + "eventAttributes": {}, } _write_log( log_dir, @@ -89,17 +87,41 @@ def test_collect_tool_invocations_no_tool(tmp_path: Path) -> None: { "timestamp": "2026-05-31T00:00:00Z", "logLevel": "STANDARD", + "eventId": "00000000-0000-0000-0000-000000000003", "eventType": "_output_event", + "eventAttributes": {}, + } + ], + ) + + assert collect_tool_invocations(log_dir) == [] + + +def test_collect_tool_invocations_legacy_format(tmp_path: Path) -> None: + """Parser remains compatible with the old nested event log format.""" + log_dir = tmp_path / "event_logs" + _write_log( + log_dir, + [ + { + "timestamp": "2026-05-31T00:00:00Z", + "logLevel": "STANDARD", + "eventType": "_tool_request_event", "event": { - "eventType": "_output_event", - "id": "y", - "attributes": {}, + "eventType": "_tool_request_event", + "id": "00000000-0000-0000-0000-000000000001", + "attributes": { + "model": "qwen3:1.7b", + "tool_calls": [_function_tool_call("add", {"a": 1, "b": 2})], + }, }, } ], ) - assert collect_tool_invocations(log_dir) == [] + assert collect_tool_invocations(log_dir) == [ + {"name": "add", "arguments": {"a": 1, "b": 2}} + ] def test_assert_tool_invoked_dict_args() -> None: diff --git a/python/flink_agents/integrations/mcp/mcp.py b/python/flink_agents/integrations/mcp/mcp.py index d45aab739..aa80073f2 100644 --- a/python/flink_agents/integrations/mcp/mcp.py +++ b/python/flink_agents/integrations/mcp/mcp.py @@ -20,7 +20,7 @@ from abc import ABC from contextlib import AsyncExitStack, asynccontextmanager from datetime import timedelta -from typing import Any, AsyncIterator, Dict, List +from typing import Any, AsyncIterator, Dict, List, Mapping from urllib.parse import urlparse import cloudpickle @@ -39,17 +39,22 @@ from flink_agents.api.chat_message import ChatMessage, MessageRole from flink_agents.api.prompts.prompt import Prompt from flink_agents.api.resource import Resource, ResourceType +from flink_agents.api.tools import ToolExecutionMetadataProvider from flink_agents.api.tools.tool import Tool, ToolMetadata, ToolType +from flink_agents.api.trace import ( + ToolExecutionMetadataKeys, +) from flink_agents.integrations.mcp.utils import extract_mcp_content_item -class MCPTool(Tool): +class MCPTool(Tool, ToolExecutionMetadataProvider): """MCP tool definition that can be called directly. This represents a single tool from an MCP server. """ mcp_server: "MCPServer" = Field(default=None) + mcp_server_name: str | None = None @classmethod @override @@ -67,6 +72,14 @@ def call(self, *args: Any, **kwargs: Any) -> Any: self.mcp_server.call_tool_async(self.metadata.name, *args, **kwargs) ) + @override + def get_tool_execution_metadata( + self, parameters: Mapping[str, Any] + ) -> Mapping[str, Any]: + if self.mcp_server_name is None: + return {} + return {ToolExecutionMetadataKeys.MCP_SERVER: self.mcp_server_name} + @override def close(self) -> None: if self.mcp_server is not None: diff --git a/python/flink_agents/integrations/mcp/tests/test_mcp.py b/python/flink_agents/integrations/mcp/tests/test_mcp.py index 1bb81ce8a..1b013ebb3 100644 --- a/python/flink_agents/integrations/mcp/tests/test_mcp.py +++ b/python/flink_agents/integrations/mcp/tests/test_mcp.py @@ -27,6 +27,7 @@ from flink_agents.api.chat_message import ChatMessage, MessageRole from flink_agents.api.tools.tool import ToolMetadata +from flink_agents.api.trace import ToolExecutionMetadataKeys from flink_agents.integrations.mcp.mcp import MCPServer, MCPTool @@ -140,7 +141,11 @@ def test_mcp_tool_roundtrip_preserves_metadata() -> None: "required": ["a", "b"], }, ) - tool = MCPTool(metadata=metadata, mcp_server=MCPServer(endpoint="http://x")) + tool = MCPTool( + metadata=metadata, + mcp_server=MCPServer(endpoint="http://x"), + mcp_server_name="calculator_server", + ) dumped = tool.model_dump() assert "metadata" in dumped, "serialized form must expose `metadata` key" @@ -149,3 +154,7 @@ def test_mcp_tool_roundtrip_preserves_metadata() -> None: restored = MCPTool.model_validate(dumped) assert restored.metadata == metadata assert restored.name == "add" + assert restored.mcp_server_name == "calculator_server" + assert restored.get_tool_execution_metadata({}) == { + ToolExecutionMetadataKeys.MCP_SERVER: "calculator_server" + } diff --git a/python/flink_agents/plan/actions/chat_model_action.py b/python/flink_agents/plan/actions/chat_model_action.py index de67881b3..548e5491e 100644 --- a/python/flink_agents/plan/actions/chat_model_action.py +++ b/python/flink_agents/plan/actions/chat_model_action.py @@ -40,6 +40,11 @@ from flink_agents.api.memory_object import MemoryObject from flink_agents.api.resource import ResourceType from flink_agents.api.runner_context import RunnerContext +from flink_agents.api.trace import ( + ExecutionEntityTypes, + ExecutionProblemCategories, + ExecutionReporters, +) from flink_agents.plan.actions.action import Action from flink_agents.plan.actions.utils import support_async from flink_agents.plan.function import PythonFunction @@ -266,6 +271,29 @@ def _generate_structured_output( return response +def _generate_structured_output_with_report( + ctx: RunnerContext, response: ChatMessage, output_schema: OutputSchema +) -> ChatMessage: + ExecutionReporters.started(ctx, ExecutionEntityTypes.PARSER, STRUCTURED_OUTPUT) + try: + structured_response = _generate_structured_output(response, output_schema) + except Exception as e: + ExecutionReporters.failed( + ctx, + ExecutionEntityTypes.PARSER, + STRUCTURED_OUTPUT, + {}, + e, + ExecutionProblemCategories.MODEL_OUTPUT_PARSE_ERROR, + ) + raise + else: + ExecutionReporters.succeeded( + ctx, ExecutionEntityTypes.PARSER, STRUCTURED_OUTPUT + ) + return structured_response + + def _clean_llm_response(raw_response: str) -> str: trimmed = raw_response.strip() if trimmed.startswith("```"): @@ -273,6 +301,13 @@ def _clean_llm_response(raw_response: str) -> str: return trimmed +def _require_model_response(response: ChatMessage | None) -> ChatMessage: + if response is None: + error_message = "ChatModel returned a null response." + raise ValueError(error_message) + return response + + async def chat( initial_request_id: UUID, model: str, @@ -316,14 +351,28 @@ async def chat( for attempt in range(num_retries + 1): try: - if chat_async: - response = await ctx.durable_execute_async( - chat_model.chat, messages, prompt_args=prompt_args - ) - else: - response = ctx.durable_execute( - chat_model.chat, messages, prompt_args=prompt_args + ExecutionReporters.started(ctx, ExecutionEntityTypes.LLM, model) + try: + if chat_async: + response = await ctx.durable_execute_async( + chat_model.chat, messages, prompt_args=prompt_args + ) + else: + response = ctx.durable_execute( + chat_model.chat, messages, prompt_args=prompt_args + ) + response = _require_model_response(response) + except Exception as model_error: + ExecutionReporters.failed( + ctx, + ExecutionEntityTypes.LLM, + model, + {}, + model_error, + ExecutionProblemCategories.MODEL_CALL_FAILED, ) + raise + ExecutionReporters.succeeded(ctx, ExecutionEntityTypes.LLM, model) if ( response.extra_args.get("model_name") @@ -336,7 +385,9 @@ async def chat( response.extra_args["completionTokens"], ) if output_schema is not None and len(response.tool_calls) == 0: - response = _generate_structured_output(response, output_schema) + response = _generate_structured_output_with_report( + ctx, response, output_schema + ) break except Exception as e: if error_handling_strategy == ErrorHandlingStrategy.IGNORE: diff --git a/python/flink_agents/plan/actions/tool_call_action.py b/python/flink_agents/plan/actions/tool_call_action.py index f605c4807..3b06c7e41 100644 --- a/python/flink_agents/plan/actions/tool_call_action.py +++ b/python/flink_agents/plan/actions/tool_call_action.py @@ -16,6 +16,7 @@ # limitations under the License. ################################################################################# import logging +from typing import Any from flink_agents.api.core_options import AgentExecutionOptions from flink_agents.api.events.event import Event @@ -23,10 +24,17 @@ from flink_agents.api.memory_object import MemoryObject from flink_agents.api.resource import ResourceType from flink_agents.api.runner_context import RunnerContext +from flink_agents.api.tools import ToolExecutionMetadataProvider from flink_agents.api.tools.tool_parameter_injection import ( InjectedArg, ToolParameterSource, ) +from flink_agents.api.trace import ( + ExecutionEntityTypes, + ExecutionProblemCategories, + ExecutionReporters, + ToolExecutionMetadataKeys, +) from flink_agents.plan.actions.action import Action from flink_agents.plan.function import PythonFunction from flink_agents.plan.tools.function_tool import FunctionTool @@ -34,6 +42,41 @@ _logger = logging.getLogger(__name__) +def _tool_entity_metadata( + tool_request_event_id: object, + tool_call_id: object, + external_id: object, + tool_name: str, + tool: object | None, + kwargs: dict[str, Any], +) -> dict[str, Any]: + metadata: dict[str, Any] = { + ToolExecutionMetadataKeys.TOOL_REQUEST_EVENT_ID: str(tool_request_event_id), + ToolExecutionMetadataKeys.TOOL_CALL_ID: str(tool_call_id), + } + if external_id is not None: + metadata[ToolExecutionMetadataKeys.EXTERNAL_ID] = str(external_id) + tool_type = tool.tool_type() if tool is not None else None + if tool_type is not None: + metadata[ToolExecutionMetadataKeys.TOOL_TYPE] = getattr( + tool_type, "value", str(tool_type) + ) + if isinstance(tool, ToolExecutionMetadataProvider): + try: + supplemental = tool.get_tool_execution_metadata(dict(kwargs)) or {} + except Exception: + _logger.debug( + "Failed to collect execution metadata for tool %s.", + tool_name, + exc_info=True, + ) + supplemental = {} + for key, value in supplemental.items(): + if key is not None and value is not None: + metadata.setdefault(key, value) + return metadata + + async def process_tool_request(event: Event, ctx: RunnerContext) -> None: """Built-in action for processing tool call requests.""" event = ToolRequestEvent.from_event(event) @@ -52,38 +95,73 @@ async def process_tool_request(event: Event, ctx: RunnerContext) -> None: name = tool_call["function"]["name"] kwargs = tool_call["function"]["arguments"] external_id = tool_call.get("original_id") + external_ids[call_id] = external_id + call_kwargs = dict(kwargs or {}) + tool = None + preparation_error = None try: tool = ctx.get_resource(name, ResourceType.TOOL) except Exception as e: - tool = None - error[call_id] = str(e) - if not tool: - responses[call_id] = f"Tool `{name}` does not exist." - success[call_id] = False - error.setdefault(call_id, f"Tool `{name}` does not exist.") - external_ids[call_id] = external_id - continue - else: + preparation_error = e + if tool is not None: try: - call_kwargs = dict(kwargs or {}) # Framework-owned injected args must win over model-provided values so # hidden context such as tenant ids cannot be spoofed by tool calls. call_kwargs.update(_resolve_injected_arguments(tool, ctx)) - if tool_call_async: - response = await ctx.durable_execute_async( - tool.call, **call_kwargs - ) - else: - response = ctx.durable_execute(tool.call, **call_kwargs) - responses[call_id] = response - success[call_id] = True except Exception as e: - responses[call_id] = f"Tool `{name}` execute failed." - success[call_id] = False - error[call_id] = str(e) + preparation_error = e - external_ids[call_id] = external_id + entity_metadata = _tool_entity_metadata( + event.id, call_id, external_id, name, tool, call_kwargs + ) + ExecutionReporters.started( + ctx, ExecutionEntityTypes.TOOL, name, entity_metadata + ) + + if not tool or preparation_error is not None: + failure = preparation_error or RuntimeError( + f"Tool `{name}` does not exist." + ) + responses[call_id] = ( + f"Tool `{name}` does not exist." + if not tool + else f"Tool `{name}` execute failed." + ) + success[call_id] = False + error[call_id] = str(failure) + ExecutionReporters.failed( + ctx, + ExecutionEntityTypes.TOOL, + name, + entity_metadata, + failure, + ExecutionProblemCategories.TOOL_CALL_FAILED, + ) + continue + + try: + if tool_call_async: + response = await ctx.durable_execute_async(tool.call, **call_kwargs) + else: + response = ctx.durable_execute(tool.call, **call_kwargs) + responses[call_id] = response + success[call_id] = True + ExecutionReporters.succeeded( + ctx, ExecutionEntityTypes.TOOL, name, entity_metadata + ) + except Exception as e: + responses[call_id] = f"Tool `{name}` execute failed." + success[call_id] = False + error[call_id] = str(e) + ExecutionReporters.failed( + ctx, + ExecutionEntityTypes.TOOL, + name, + entity_metadata, + e, + ExecutionProblemCategories.TOOL_CALL_FAILED, + ) ctx.send_event( ToolResponseEvent( request_id=event.id, @@ -125,9 +203,7 @@ def _resolve_injected_argument(injection: InjectedArg, ctx: RunnerContext) -> ob def _get_memory_value(memory: MemoryObject, source: str, path: str) -> object: if memory is None: - msg = ( - f"Cannot inject tool parameter from {source} because memory is not initialized." - ) + msg = f"Cannot inject tool parameter from {source} because memory is not initialized." raise ValueError(msg) if not memory.is_exist(path): msg = f"Missing memory path for injected tool parameter: {path}" diff --git a/python/flink_agents/plan/agent_plan.py b/python/flink_agents/plan/agent_plan.py index 747aa8577..493714a0d 100644 --- a/python/flink_agents/plan/agent_plan.py +++ b/python/flink_agents/plan/agent_plan.py @@ -75,6 +75,7 @@ class AgentPlan(BaseModel): actions: Dict[str, Action] actions_by_event: Dict[str, List[str]] resource_providers: Dict[ResourceType, Dict[str, ResourceProvider]] | None = None + agent_name: str | None = None config: AgentConfiguration | None = None @field_serializer("resource_providers") @@ -138,7 +139,9 @@ def __custom_deserialize(self) -> "AgentPlan": return self @staticmethod - def from_agent(agent: Agent, config: AgentConfiguration) -> "AgentPlan": + def from_agent( + agent: Agent, config: AgentConfiguration, agent_name: str | None = None + ) -> "AgentPlan": """Build a AgentPlan from user defined agent.""" actions = {} actions_by_event = {} @@ -164,6 +167,7 @@ def from_agent(agent: Agent, config: AgentConfiguration) -> "AgentPlan": actions=actions, actions_by_event=actions_by_event, resource_providers=resource_providers, + agent_name=agent_name or agent.__class__.__name__, config=config, ) @@ -299,10 +303,7 @@ def _get_actions(agent: Agent) -> List[Action]: Action( name=name, exec=_to_plan_function(action_tuple[1]), - trigger_conditions=[ - _resolve_event_type(et) - for et in action_tuple[0] - ], + trigger_conditions=[_resolve_event_type(et) for et in action_tuple[0]], config=action_tuple[2], ) ) @@ -480,14 +481,13 @@ def get_skill_dirs(self, *skill_names: str) -> List[str]: ] ) - resource_providers.extend( - [ + for tool in mcp_server.list_tools(): + tool.mcp_server_name = name + resource_providers.append( PythonSerializableResourceProvider.from_resource( name=tool.name, resource=tool ) - for tool in mcp_server.list_tools() - ] - ) + ) mcp_server.close() diff --git a/python/flink_agents/plan/tests/actions/test_chat_model_action_retry.py b/python/flink_agents/plan/tests/actions/test_chat_model_action_retry.py index 204cecb7b..85474d8b9 100644 --- a/python/flink_agents/plan/tests/actions/test_chat_model_action_retry.py +++ b/python/flink_agents/plan/tests/actions/test_chat_model_action_retry.py @@ -20,11 +20,14 @@ import asyncio import time from typing import Any, Sequence -from unittest.mock import MagicMock +from unittest.mock import MagicMock, call from uuid import uuid4 import pytest +from pydantic import BaseModel +from flink_agents.api.agents.agent import STRUCTURED_OUTPUT +from flink_agents.api.agents.react_agent import OutputSchema from flink_agents.api.chat_message import ChatMessage, MessageRole from flink_agents.api.core_options import ( AgentExecutionOptions, @@ -33,6 +36,11 @@ from flink_agents.api.events.chat_event import ChatResponseEvent from flink_agents.api.events.tool_event import ToolResponseEvent from flink_agents.api.metric_group import Counter, MetricGroup +from flink_agents.api.trace import ( + ExecutionEntityTypes, + ExecutionProblemCategories, + ExecutionReporter, +) from flink_agents.plan.actions.chat_model_action import ( chat, process_chat_request_or_tool_response, @@ -100,6 +108,10 @@ def set(self, path: str, value: Any) -> None: self._store[path] = value +class _StructuredResult(BaseModel): + result: int + + def _create_mock_runner_context( chat_model: Any, max_retries: int = 3, @@ -126,7 +138,7 @@ def _create_mock_runner_context( ) ) - ctx = MagicMock() + ctx = MagicMock(spec=ExecutionReporter) ctx.config = config ctx.sensory_memory = sensory_memory ctx.action_metric_group = metric_group @@ -176,6 +188,13 @@ def test_chat_succeeds_without_retry(self) -> None: # No retry metrics should be recorded assert len(metric_group._sub_groups) == 0 + ctx.report_execution_started.assert_called_once_with( + ExecutionEntityTypes.LLM, chat_model.connection, {} + ) + ctx.report_execution_succeeded.assert_called_once_with( + ExecutionEntityTypes.LLM, chat_model.connection, {} + ) + ctx.report_execution_failed.assert_not_called() def test_chat_retries_with_exponential_backoff(self) -> None: """Fail once then succeed: 1s interval, 1 retry -> wait 1s (1 * 2^0).""" @@ -222,6 +241,14 @@ def mock_chat(messages: Sequence[ChatMessage], **kwargs: Any) -> ChatMessage: model_group = metric_group.get_sub_group("model", chat_model.connection) assert model_group.get_counter("retryCount").get_count() == 1 assert model_group.get_counter("retryWaitSec").get_count() == 1 + assert ctx.report_execution_started.call_count == 2 + ctx.report_execution_failed.assert_called_once() + failed_args = ctx.report_execution_failed.call_args.args + assert failed_args[0] == ExecutionEntityTypes.LLM + assert failed_args[-1] == ExecutionProblemCategories.MODEL_CALL_FAILED + ctx.report_execution_succeeded.assert_called_once_with( + ExecutionEntityTypes.LLM, "test-model", {} + ) def test_chat_exhausts_retries_and_raises(self) -> None: """All retries exhausted: exception raised, no event sent.""" @@ -246,6 +273,58 @@ def test_chat_exhausts_retries_and_raises(self) -> None: ) assert len(sent_events) == 0 + assert ctx.report_execution_started.call_count == 3 + assert ctx.report_execution_failed.call_count == 3 + for failed_call in ctx.report_execution_failed.call_args_list: + assert failed_call.args[0] == ExecutionEntityTypes.LLM + assert failed_call.args[-1] == ExecutionProblemCategories.MODEL_CALL_FAILED + ctx.report_execution_succeeded.assert_not_called() + + def test_structured_output_parse_error_retries_without_failing_llm( + self, + ) -> None: + chat_model = MagicMock() + chat_model.chat = MagicMock( + side_effect=[ + ChatMessage(role=MessageRole.ASSISTANT, content="not-json"), + ChatMessage(role=MessageRole.ASSISTANT, content='{"result": 42}'), + ] + ) + + ctx, sent_events, _, _ = _create_mock_runner_context( + chat_model, max_retries=1, retry_wait_interval_sec=0 + ) + + asyncio.run( + chat( + uuid4(), + "test-model", + [ChatMessage(role=MessageRole.USER, content="hi")], + {}, + OutputSchema(output_schema=_StructuredResult), + ctx, + ) + ) + + assert chat_model.chat.call_count == 2 + assert len(sent_events) == 1 + response = sent_events[0].response + assert response.extra_args[STRUCTURED_OUTPUT].result == 42 + + ctx.report_execution_failed.assert_called_once() + failed_args = ctx.report_execution_failed.call_args.args + assert failed_args[0] == ExecutionEntityTypes.PARSER + assert failed_args[1] == STRUCTURED_OUTPUT + assert failed_args[-1] == ExecutionProblemCategories.MODEL_OUTPUT_PARSE_ERROR + + assert ctx.report_execution_started.call_count == 4 + assert ctx.report_execution_succeeded.call_count == 3 + assert ( + ctx.report_execution_succeeded.call_args_list.count( + call(ExecutionEntityTypes.LLM, "test-model", {}) + ) + == 2 + ) class TestChatResponseEventRetryFields: @@ -325,9 +404,9 @@ def mock_chat(messages: Sequence[ChatMessage], **kwargs: Any) -> ChatMessage: "_TOOL_CALL_CONTEXT", { str(initial_request_id): [ - ChatMessage( - role=MessageRole.USER, content="hi" - ).model_dump(mode="json") + ChatMessage(role=MessageRole.USER, content="hi").model_dump( + mode="json" + ) ] }, ) @@ -377,9 +456,9 @@ def mock_chat(messages: Sequence[ChatMessage], **kwargs: Any) -> ChatMessage: "_TOOL_CALL_CONTEXT", { str(initial_request_id): [ - ChatMessage( - role=MessageRole.USER, content="hi" - ).model_dump(mode="json") + ChatMessage(role=MessageRole.USER, content="hi").model_dump( + mode="json" + ) ] }, ) @@ -389,7 +468,9 @@ def mock_chat(messages: Sequence[ChatMessage], **kwargs: Any) -> ChatMessage: responses={tool_call_id: "Tool `query_order` execute failed."}, external_ids={}, success={tool_call_id: False}, - error={tool_call_id: "Missing config for injected tool parameter: tenant_id"}, + error={ + tool_call_id: "Missing config for injected tool parameter: tenant_id" + }, ) asyncio.run(process_chat_request_or_tool_response(tool_response_event, ctx)) diff --git a/python/flink_agents/plan/tests/actions/test_tool_call_action.py b/python/flink_agents/plan/tests/actions/test_tool_call_action.py index d9690981a..3cea0b0ed 100644 --- a/python/flink_agents/plan/tests/actions/test_tool_call_action.py +++ b/python/flink_agents/plan/tests/actions/test_tool_call_action.py @@ -17,12 +17,20 @@ ################################################################################# import asyncio from typing import Any +from unittest.mock import MagicMock from flink_agents.api.core_options import AgentExecutionOptions from flink_agents.api.events.tool_event import ToolRequestEvent, ToolResponseEvent from flink_agents.api.memory_object import MemoryObject from flink_agents.api.resource import ResourceType -from flink_agents.api.tools import InjectedArg +from flink_agents.api.tools import InjectedArg, ToolExecutionMetadataProvider +from flink_agents.api.tools.tool import ToolType +from flink_agents.api.trace import ( + ExecutionEntityTypes, + ExecutionProblemCategories, + ExecutionReporter, + ToolExecutionMetadataKeys, +) from flink_agents.plan.actions.tool_call_action import process_tool_request from flink_agents.plan.configuration import AgentConfiguration from flink_agents.plan.function import PythonFunction @@ -87,7 +95,9 @@ def get(self, path_or_ref: Any) -> Any: def set(self, path: str, value: Any) -> Any: raise NotImplementedError - def new_object(self, path: str, *, overwrite: bool = False) -> "_NestedMemoryObject": + def new_object( + self, path: str, *, overwrite: bool = False + ) -> "_NestedMemoryObject": raise NotImplementedError def is_exist(self, path: str) -> bool: @@ -188,7 +198,10 @@ def test_tool_call_action_reports_missing_config_injected_arg() -> None: { "id": "call-1", "type": "function", - "function": {"name": "query_order", "arguments": {"order_id": "order-1"}}, + "function": { + "name": "query_order", + "arguments": {"order_id": "order-1"}, + }, } ], ) @@ -268,7 +281,9 @@ def test_tool_call_action_exposes_wrong_config_type() -> None: response = ToolResponseEvent.from_event(ctx.sent_events[0]) assert response.responses["call-1"] == "Tool `query_order` execute failed." assert response.success["call-1"] is False - assert response.error["call-1"] == "'_WrongConfig' object has no attribute 'conf_data'" + assert ( + response.error["call-1"] == "'_WrongConfig' object has no attribute 'conf_data'" + ) def test_tool_call_action_uses_sync_execution_in_test_context() -> None: @@ -277,6 +292,125 @@ def test_tool_call_action_uses_sync_execution_in_test_context() -> None: assert ctx.config.get(AgentExecutionOptions.TOOL_CALL_ASYNC) is False +def test_tool_call_reports_started_and_succeeded() -> None: + tool = MagicMock() + tool.tool_type.return_value = ToolType.FUNCTION + tool.call = MagicMock(return_value="result") + ctx, sent_events = trace_context(tool) + request = ToolRequestEvent(model="model-a", tool_calls=[trace_tool_call()]) + + asyncio.run(process_tool_request(request, ctx)) + + assert len(sent_events) == 1 + metadata = { + ToolExecutionMetadataKeys.TOOL_REQUEST_EVENT_ID: str(request.id), + ToolExecutionMetadataKeys.TOOL_CALL_ID: "call-1", + ToolExecutionMetadataKeys.EXTERNAL_ID: "external-call-1", + ToolExecutionMetadataKeys.TOOL_TYPE: "function", + } + ctx.report_execution_started.assert_called_once_with( + ExecutionEntityTypes.TOOL, "search", metadata + ) + ctx.report_execution_succeeded.assert_called_once_with( + ExecutionEntityTypes.TOOL, "search", metadata + ) + ctx.report_execution_failed.assert_not_called() + + +def test_tool_call_reports_failed() -> None: + tool = MagicMock() + tool.tool_type.return_value = ToolType.FUNCTION + tool.call = MagicMock(side_effect=RuntimeError("boom")) + ctx, _ = trace_context(tool) + request = ToolRequestEvent(model="model-a", tool_calls=[trace_tool_call()]) + + asyncio.run(process_tool_request(request, ctx)) + + ctx.report_execution_failed.assert_called_once() + args = ctx.report_execution_failed.call_args.args + assert args[0] == ExecutionEntityTypes.TOOL + assert args[1] == "search" + assert args[2][ToolExecutionMetadataKeys.TOOL_CALL_ID] == "call-1" + assert isinstance(args[3], RuntimeError) + assert args[4] == ExecutionProblemCategories.TOOL_CALL_FAILED + + +def test_tool_call_includes_provider_metadata() -> None: + class MetadataTool(ToolExecutionMetadataProvider): + @staticmethod + def tool_type() -> ToolType: + return ToolType.MCP + + @staticmethod + def call(**kwargs: object) -> str: + return "result" + + def get_tool_execution_metadata( + self, parameters: dict[str, object] + ) -> dict[str, object]: + return {ToolExecutionMetadataKeys.MCP_SERVER: "search_server"} + + ctx, _ = trace_context(MetadataTool()) + request = ToolRequestEvent(model="model-a", tool_calls=[trace_tool_call()]) + + asyncio.run(process_tool_request(request, ctx)) + + metadata = ctx.report_execution_started.call_args.args[2] + assert metadata[ToolExecutionMetadataKeys.MCP_SERVER] == "search_server" + + +def test_tool_execution_metadata_cannot_mutate_call_arguments() -> None: + class MutatingMetadataTool(ToolExecutionMetadataProvider): + @staticmethod + def tool_type() -> ToolType: + return ToolType.FUNCTION + + @staticmethod + def call(**kwargs: object) -> object: + return kwargs["query"] + + def get_tool_execution_metadata( + self, parameters: dict[str, object] + ) -> dict[str, object]: + parameters["query"] = "mutated" + return {} + + ctx, sent_events = trace_context(MutatingMetadataTool()) + + asyncio.run( + process_tool_request( + ToolRequestEvent(model="model-a", tool_calls=[trace_tool_call()]), ctx + ) + ) + + response = ToolResponseEvent.from_event(sent_events[0]) + assert response.responses["call-1"] == "flink" + + +def trace_context(tool: object) -> tuple[MagicMock, list[ToolResponseEvent]]: + sent_events = [] + config = MagicMock() + config.get = MagicMock( + side_effect=lambda option: False + if option is AgentExecutionOptions.TOOL_CALL_ASYNC + else option.get_default_value() + ) + ctx = MagicMock(spec=ExecutionReporter) + ctx.config = config + ctx.get_resource = MagicMock(return_value=tool) + ctx.durable_execute = MagicMock(side_effect=lambda fn, **kwargs: fn(**kwargs)) + ctx.send_event = MagicMock(side_effect=lambda event: sent_events.append(event)) + return ctx, sent_events + + +def trace_tool_call() -> dict: + return { + "id": "call-1", + "original_id": "external-call-1", + "function": {"name": "search", "arguments": {"query": "flink"}}, + } + + def tool_request() -> ToolRequestEvent: return ToolRequestEvent( model="model", @@ -284,7 +418,10 @@ def tool_request() -> ToolRequestEvent: { "id": "call-1", "type": "function", - "function": {"name": "query_order", "arguments": {"order_id": "order-1"}}, + "function": { + "name": "query_order", + "arguments": {"order_id": "order-1"}, + }, } ], ) diff --git a/python/flink_agents/plan/tests/resources/agent_plan.json b/python/flink_agents/plan/tests/resources/agent_plan.json index 6e545ab1d..f0fa12320 100644 --- a/python/flink_agents/plan/tests/resources/agent_plan.json +++ b/python/flink_agents/plan/tests/resources/agent_plan.json @@ -149,6 +149,7 @@ } } }, + "agent_name": "MyAgent", "config": { "conf_data": { "mock.key": "mock.value" diff --git a/python/flink_agents/plan/tests/test_agent_plan.py b/python/flink_agents/plan/tests/test_agent_plan.py index 04fdb6c79..5a1da7429 100644 --- a/python/flink_agents/plan/tests/test_agent_plan.py +++ b/python/flink_agents/plan/tests/test_agent_plan.py @@ -68,6 +68,7 @@ def increment(event: Event, ctx: RunnerContext) -> None: def test_from_agent(): agent = AgentForTest() agent_plan = AgentPlan.from_agent(agent, AgentConfiguration()) + assert agent_plan.agent_name == "AgentForTest" actions = agent_plan.get_actions(InputEvent.EVENT_TYPE) assert len(actions) == 1 action = actions[0] @@ -79,6 +80,14 @@ def test_from_agent(): assert action.trigger_conditions == [InputEvent.EVENT_TYPE] +def test_from_agent_uses_explicit_agent_name() -> None: + agent_plan = AgentPlan.from_agent( + AgentForTest(), AgentConfiguration(), "registered_agent" + ) + + assert agent_plan.agent_name == "registered_agent" + + class InvalidAgent(Agent): @action(EventType.InputEvent) @staticmethod @@ -148,8 +157,7 @@ def test_action_inherited_from_parent_agent_class_is_rejected() -> None: _JAVA_HANDLER_QUALNAME = ( - "org.apache.flink.agents.runtime.operator." - "CrossLanguageActionRuntimeTest$Handlers" + "org.apache.flink.agents.runtime.operator.CrossLanguageActionRuntimeTest$Handlers" ) @@ -509,10 +517,14 @@ def test_get_resource() -> None: def test_add_action_and_resource_to_agent() -> None: my_agent = Agent() my_agent.add_action( - name="first_action", trigger_conditions=["_input_event"], func=MyAgent.first_action + name="first_action", + trigger_conditions=["_input_event"], + func=MyAgent.first_action, ) my_agent.add_action( - name="second_action", trigger_conditions=["_input_event", "_my_event"], func=MyAgent.second_action + name="second_action", + trigger_conditions=["_input_event", "_my_event"], + func=MyAgent.second_action, ) my_agent.add_resource( name="mock", @@ -555,7 +567,7 @@ def test_add_action_and_resource_to_agent() -> None: ), ) agent_plan = AgentPlan.from_agent( - my_agent, AgentConfiguration({"mock.key": "mock.value"}) + my_agent, AgentConfiguration({"mock.key": "mock.value"}), "MyAgent" ) json_value = agent_plan.model_dump_json(serialize_as_any=True, indent=4) with Path.open(Path(f"{current_dir}/resources/agent_plan.json")) as f: diff --git a/python/flink_agents/runtime/flink_runner_context.py b/python/flink_agents/runtime/flink_runner_context.py index ada5b8e30..e4c54b712 100644 --- a/python/flink_agents/runtime/flink_runner_context.py +++ b/python/flink_agents/runtime/flink_runner_context.py @@ -15,8 +15,10 @@ # See the License for the specific language governing permissions and # limitations under the License. ################################################################################# +import json import logging import os +from collections.abc import Mapping from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from functools import partial @@ -38,6 +40,7 @@ AsyncExecutionResult, RunnerContext, ) +from flink_agents.api.trace import ExecutionReporter from flink_agents.runtime.durable_execution import ( _compute_args_digest, _compute_function_id, @@ -56,6 +59,24 @@ logger = logging.getLogger(__name__) +def _error_type(error: BaseException) -> str: + return f"{error.__class__.__module__}.{error.__class__.__qualname__}" + + +def _root_cause(error: BaseException) -> BaseException: + current = error + visited: set[int] = set() + while id(current) not in visited: + visited.add(id(current)) + cause = current.__cause__ + if cause is None and not current.__suppress_context__: + cause = current.__context__ + if cause is None: + break + current = cause + return current + + @dataclass(frozen=True) class _PersistedCallResult: function_id: str @@ -241,7 +262,7 @@ def __await__(self) -> Any: return result -class FlinkRunnerContext(RunnerContext): +class FlinkRunnerContext(RunnerContext, ExecutionReporter): """Providing context for agent execution in Flink Environment. This context allows access to event handling and provides fine-grained @@ -401,6 +422,56 @@ def action_metric_group(self) -> FlinkMetricGroup: """ return FlinkMetricGroup(self._j_runner_context.getActionMetricGroup()) + @override + def report_execution_started( + self, + entity_type: str, + entity_name: str, + entity_metadata: Mapping[str, Any] | None = None, + ) -> None: + self._j_runner_context.reportExecutionStartedJson( + entity_type, + entity_name, + self._entity_metadata_json(entity_metadata), + ) + + @override + def report_execution_succeeded( + self, + entity_type: str, + entity_name: str, + entity_metadata: Mapping[str, Any] | None = None, + ) -> None: + self._j_runner_context.reportExecutionSucceededJson( + entity_type, + entity_name, + self._entity_metadata_json(entity_metadata), + ) + + @override + def report_execution_failed( + self, + entity_type: str, + entity_name: str, + entity_metadata: Mapping[str, Any] | None, + error: BaseException, + problem_category: str | None = None, + ) -> None: + root_error = _root_cause(error) + error_message = str(root_error) + self._j_runner_context.reportExecutionFailedJson( + entity_type, + entity_name, + self._entity_metadata_json(entity_metadata), + _error_type(root_error), + error_message or None, + problem_category, + ) + + @staticmethod + def _entity_metadata_json(entity_metadata: Mapping[str, Any] | None) -> str: + return json.dumps(dict(entity_metadata or {})) + def _try_get_cached_result( self, func: Callable, args: tuple, kwargs: dict ) -> tuple[bool, Any]: diff --git a/python/flink_agents/runtime/remote_execution_environment.py b/python/flink_agents/runtime/remote_execution_environment.py index b14ed0eb3..8c24af9b7 100644 --- a/python/flink_agents/runtime/remote_execution_environment.py +++ b/python/flink_agents/runtime/remote_execution_environment.py @@ -94,7 +94,9 @@ def apply(self, agent: Agent | str) -> "AgentBuilder": if self.__agent_plan is not None: err_msg = "RemoteAgentBuilder doesn't support apply multiple agents yet." raise RuntimeError(err_msg) + agent_name = None if isinstance(agent, str): + agent_name = agent if agent not in self.__agents: msg = ( f"No agent named {agent!r} is registered on this " @@ -107,7 +109,7 @@ def apply(self, agent: Agent | str) -> "AgentBuilder": for type, name_to_resource in self.__resources.items(): agent.resources[type] = name_to_resource | agent.resources[type] - self.__agent_plan = AgentPlan.from_agent(agent, self.__config) + self.__agent_plan = AgentPlan.from_agent(agent, self.__config, agent_name) return self diff --git a/python/flink_agents/runtime/skill/skill_tools.py b/python/flink_agents/runtime/skill/skill_tools.py index 3d1a1b852..32580fe44 100644 --- a/python/flink_agents/runtime/skill/skill_tools.py +++ b/python/flink_agents/runtime/skill/skill_tools.py @@ -23,11 +23,15 @@ from typing import TYPE_CHECKING, Any if TYPE_CHECKING: + from collections.abc import Mapping + from flink_agents.runtime.skill.skill_manager import SkillManager from pydantic import BaseModel, Field +from flink_agents.api.tools import ToolExecutionMetadataProvider from flink_agents.api.tools.tool import Tool, ToolMetadata, ToolType +from flink_agents.api.trace import ToolExecutionMetadataKeys logger = logging.getLogger(__name__) @@ -47,7 +51,7 @@ class LoadSkillArgs(BaseModel): ) -class LoadSkillTool(Tool): +class LoadSkillTool(Tool, ToolExecutionMetadataProvider): """Tool for loading skill content and resources. Accesses the SkillManager through the runtime ResourceContext @@ -73,6 +77,19 @@ def tool_type(cls) -> ToolType: """Return tool type of class.""" return ToolType.FUNCTION + def get_tool_execution_metadata( + self, parameters: Mapping[str, Any] + ) -> Mapping[str, Any]: + """Describe the requested skill resource for execution tracing.""" + metadata = {} + if "name" in parameters: + metadata[ToolExecutionMetadataKeys.SKILL_NAME] = str(parameters["name"]) + if "path" in parameters: + metadata[ToolExecutionMetadataKeys.SKILL_RESOURCE_PATH] = str( + parameters["path"] + ) + return metadata + def call(self, *args: Any, **kwargs: Any) -> str: """Call the tool to load a skill.""" if args: diff --git a/python/flink_agents/runtime/skill/tests/test_load_skill.py b/python/flink_agents/runtime/skill/tests/test_load_skill.py index 76db7583d..f7b8bd5da 100644 --- a/python/flink_agents/runtime/skill/tests/test_load_skill.py +++ b/python/flink_agents/runtime/skill/tests/test_load_skill.py @@ -24,6 +24,7 @@ from flink_agents.api.resource_context import ResourceContext from flink_agents.api.skills import Skills +from flink_agents.api.trace import ToolExecutionMetadataKeys from flink_agents.runtime.skill.skill_manager import SkillManager from flink_agents.runtime.skill.skill_tools import LoadSkillTool @@ -86,6 +87,16 @@ def test_load_resource_not_found(self, tool: LoadSkillTool) -> None: assert "not found" in result.lower() assert "scripts/generate_image.py" in result + def test_execution_metadata_describes_requested_resource( + self, tool: LoadSkillTool + ) -> None: + """Execution metadata identifies the explicitly requested skill resource.""" + metadata = tool.get_tool_execution_metadata( + {"name": "github", "path": "README.md"} + ) + assert metadata[ToolExecutionMetadataKeys.SKILL_NAME] == "github" + assert metadata[ToolExecutionMetadataKeys.SKILL_RESOURCE_PATH] == "README.md" + # -- skill not found ----------------------------------------------------- def test_skill_not_found(self, tool: LoadSkillTool) -> None: diff --git a/python/flink_agents/runtime/tests/test_flink_runner_context_trace.py b/python/flink_agents/runtime/tests/test_flink_runner_context_trace.py new file mode 100644 index 000000000..03256e447 --- /dev/null +++ b/python/flink_agents/runtime/tests/test_flink_runner_context_trace.py @@ -0,0 +1,56 @@ +################################################################################ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +################################################################################# +from unittest.mock import MagicMock + +from flink_agents.api.trace import ( + ExecutionEntityTypes, + ExecutionProblemCategories, + ExecutionReporter, +) +from flink_agents.runtime.flink_runner_context import FlinkRunnerContext + + +def test_flink_runner_context_is_execution_reporter() -> None: + assert issubclass(FlinkRunnerContext, ExecutionReporter) + + +def test_failed_execution_reports_deepest_cause() -> None: + java_context = MagicMock() + ctx = FlinkRunnerContext.__new__(FlinkRunnerContext) + ctx._j_runner_context = java_context + root_error = ValueError("invalid model response") + wrapped_error = RuntimeError("structured output failed") + wrapped_error.__cause__ = root_error + + ctx.report_execution_failed( + ExecutionEntityTypes.PARSER, + "structured_output", + {}, + wrapped_error, + ExecutionProblemCategories.MODEL_OUTPUT_PARSE_ERROR, + ) + + java_context.reportExecutionFailedJson.assert_called_once_with( + ExecutionEntityTypes.PARSER, + "structured_output", + "{}", + "builtins.ValueError", + "invalid model response", + ExecutionProblemCategories.MODEL_OUTPUT_PARSE_ERROR, + ) diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/PythonMCPResourceDiscovery.java b/runtime/src/main/java/org/apache/flink/agents/runtime/PythonMCPResourceDiscovery.java index 356ec67be..b9e0b48dc 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/PythonMCPResourceDiscovery.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/PythonMCPResourceDiscovery.java @@ -72,7 +72,7 @@ public static void discoverPythonMCPResources( PythonMCPServer server = (PythonMCPServer) provider.provide(cache.getResourceContext()); - for (PythonMCPTool tool : server.listTools()) { + for (PythonMCPTool tool : server.listTools(provider.getName())) { cache.put(tool.getName(), TOOL, tool); } for (PythonMCPPrompt prompt : server.listPrompts()) { diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java b/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java index 1395bdc56..a26a98a13 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java @@ -30,6 +30,9 @@ import org.apache.flink.agents.api.memory.BaseLongTermMemory; import org.apache.flink.agents.api.resource.Resource; import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.api.trace.ExecutionLifecycleEvents; +import org.apache.flink.agents.api.trace.ExecutionReporter; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; import org.apache.flink.agents.plan.AgentPlan; import org.apache.flink.agents.plan.actions.Action; import org.apache.flink.agents.plan.utils.JsonUtils; @@ -40,6 +43,8 @@ import org.apache.flink.agents.runtime.memory.InteranlBaseLongTermMemory; import org.apache.flink.agents.runtime.memory.MemoryObjectImpl; import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl; +import org.apache.flink.agents.runtime.trace.ExecutionEventSink; +import org.apache.flink.agents.runtime.trace.ReportedExecutionKey; import org.apache.flink.util.Preconditions; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -57,7 +62,7 @@ * The implementation class of {@link RunnerContext}, which serves as the execution context for * actions. */ -public class RunnerContextImpl implements RunnerContext { +public class RunnerContextImpl implements RunnerContext, ExecutionReporter { private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper().registerModule(new JavaTimeModule()); @@ -105,6 +110,10 @@ public CachedMemoryStore getSensoryMemStore() { protected String actionName; protected InteranlBaseLongTermMemory ltm; + @Nullable protected ExecutionTraceContext actionTraceContext; + @Nullable protected ExecutionEventSink executionEventSink; + @Nullable private Map activeReportedExecutions; + /** Context for fine-grained durable execution, may be null if not enabled. */ @Nullable protected DurableExecutionContext durableExecutionContext; @@ -125,13 +134,28 @@ public void setLongTermMemory(InteranlBaseLongTermMemory ltm) { } public void switchActionContext(String actionName, MemoryContext memoryContext, String key) { + switchActionContext(actionName, memoryContext, key, null, null); + } + + public void switchActionContext( + String actionName, + MemoryContext memoryContext, + String key, + @Nullable ExecutionTraceContext actionTraceContext, + @Nullable Map activeReportedExecutions) { this.actionName = actionName; this.memoryContext = memoryContext; + this.actionTraceContext = actionTraceContext; + this.activeReportedExecutions = activeReportedExecutions; if (ltm != null) { ltm.switchContext(key); } } + public void setExecutionEventSink(@Nullable ExecutionEventSink executionEventSink) { + this.executionEventSink = executionEventSink; + } + public MemoryContext getMemoryContext() { return memoryContext; } @@ -191,6 +215,81 @@ public List getShortTermMemoryUpdates() { return List.copyOf(memoryContext.getShortTermMemoryUpdates()); } + @Override + public void reportExecutionStarted( + String entityType, String entityName, Map entityMetadata) + throws Exception { + reportChildExecution( + entityType, + entityName, + entityMetadata, + ExecutionLifecycleEvents.executionStarted()); + } + + @Override + public void reportExecutionSucceeded( + String entityType, String entityName, Map entityMetadata) + throws Exception { + reportChildExecution( + entityType, + entityName, + entityMetadata, + ExecutionLifecycleEvents.executionFinished()); + } + + @Override + public void reportExecutionFailed( + String entityType, + String entityName, + Map entityMetadata, + Throwable error, + @Nullable String problemCategory) + throws Exception { + reportChildExecution( + entityType, + entityName, + entityMetadata, + ExecutionLifecycleEvents.executionFailed(error, problemCategory)); + } + + protected void reportChildExecution( + String entityType, String entityName, Map entityMetadata, Event event) { + mailboxThreadChecker.run(); + if (actionTraceContext == null + || executionEventSink == null + || activeReportedExecutions == null) { + return; + } + + ReportedExecutionKey key = new ReportedExecutionKey(entityType, entityName, entityMetadata); + ExecutionTraceContext reportTraceContext; + if (ExecutionLifecycleEvents.EXECUTION_STARTED_EVENT_TYPE.equals(event.getType())) { + reportTraceContext = + actionTraceContext.childExecution( + entityType, entityName, key.getEntityMetadata()); + ExecutionTraceContext previous = activeReportedExecutions.put(key, reportTraceContext); + if (previous != null) { + LOG.debug( + "Execution start report for {}:{} replaced an active report with the same metadata.", + entityType, + entityName); + } + } else { + reportTraceContext = activeReportedExecutions.remove(key); + if (reportTraceContext == null) { + LOG.debug( + "Execution terminal report for {}:{} has no matching start report; emitting it with a new execution id.", + entityType, + entityName); + reportTraceContext = + actionTraceContext.childExecution( + entityType, entityName, key.getEntityMetadata()); + } + } + + executionEventSink.emit(event, reportTraceContext); + } + @Override public MemoryObject getSensoryMemory() throws Exception { mailboxThreadChecker.run(); diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/env/RemoteExecutionEnvironment.java b/runtime/src/main/java/org/apache/flink/agents/runtime/env/RemoteExecutionEnvironment.java index 80d15f6c8..c89e4c556 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/env/RemoteExecutionEnvironment.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/env/RemoteExecutionEnvironment.java @@ -181,10 +181,14 @@ private StreamTableEnvironment getTableEnvironment() { @Override public AgentBuilder apply(Agent agent) { + return apply(agent, null); + } + + private AgentBuilder apply(Agent agent, @Nullable String agentName) { try { // Inspect resources registered in environment to agent. agent.addResourcesIfAbsent(resources); - this.agentPlan = new AgentPlan(agent, config); + this.agentPlan = new AgentPlan(agent, config, agentName); return this; } catch (Exception e) { throw new RuntimeException("Failed to create agent plan from agent", e); @@ -201,7 +205,7 @@ public AgentBuilder apply(String agentName) { + "'; no agent with that name is registered on the environment. " + "Did you forget to call env.loadYaml(...) or env.getAgents().put(...) first?"); } - return apply(agent); + return apply(agent, agentName); } @Override diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/EventLogRecord.java b/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/EventLogRecord.java index 402c3f52a..ec8982a08 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/EventLogRecord.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/EventLogRecord.java @@ -22,33 +22,43 @@ import com.fasterxml.jackson.databind.annotation.JsonSerialize; import org.apache.flink.agents.api.Event; import org.apache.flink.agents.api.EventContext; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; -/** - * Represents a record in the event log, containing the event context and the event itself. - * - *

This class is used to encapsulate the details of an event as it is logged, allowing for - * structured logging and retrieval of event information. - * - *

The class uses custom JSON serialization/deserialization to handle polymorphic Event types by - * leveraging the eventType information stored in the EventContext. - */ +import javax.annotation.Nullable; + +import java.util.Objects; + +/** A normalized Event Log record ready to be persisted. */ @JsonSerialize(using = EventLogRecordJsonSerializer.class) @JsonDeserialize(using = EventLogRecordJsonDeserializer.class) public class EventLogRecord { - private final EventContext context; + private final EventContext eventContext; + @Nullable private final ExecutionTraceContext traceContext; private final Event event; - public EventLogRecord(EventContext context, Event event) { - this.context = context; - this.event = event; + public EventLogRecord( + EventContext eventContext, @Nullable ExecutionTraceContext traceContext, Event event) { + this.eventContext = Objects.requireNonNull(eventContext, "eventContext"); + this.event = Objects.requireNonNull(event, "event"); + if (!event.getType().equals(eventContext.getEventType())) { + throw new IllegalArgumentException( + "EventContext event type must match Event type. context=" + + eventContext.getEventType() + + ", event=" + + event.getType()); + } + this.traceContext = traceContext; + } + + public EventContext getEventContext() { + return eventContext; } - /** Returns the event context. */ - public EventContext getContext() { - return context; + @Nullable + public ExecutionTraceContext getExecutionTraceContext() { + return traceContext; } - /** Returns the event. */ public Event getEvent() { return event; } diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/EventLogRecordJsonDeserializer.java b/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/EventLogRecordJsonDeserializer.java index ac1e69439..f99b67bc3 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/EventLogRecordJsonDeserializer.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/EventLogRecordJsonDeserializer.java @@ -19,6 +19,7 @@ package org.apache.flink.agents.runtime.eventlog; import com.fasterxml.jackson.core.JsonParser; +import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.DeserializationContext; import com.fasterxml.jackson.databind.JsonDeserializer; import com.fasterxml.jackson.databind.JsonNode; @@ -26,8 +27,13 @@ import com.fasterxml.jackson.databind.node.ObjectNode; import org.apache.flink.agents.api.Event; import org.apache.flink.agents.api.EventContext; +import org.apache.flink.agents.api.trace.ExecutionLifecycleEvents; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; import java.io.IOException; +import java.util.HashMap; +import java.util.Map; +import java.util.UUID; /** * Custom JSON deserializer for {@link EventLogRecord}. @@ -35,7 +41,7 @@ *

This deserializer reconstructs EventLogRecord instances from JSON format by: * *

    - *
  • Deserializing the EventContext (contains eventType and timestamp) + *
  • Deserializing normalized Event Log fields *
  • Deserializing the event as a base {@link Event} with type and attributes *
* @@ -57,21 +63,88 @@ public EventLogRecord deserialize(JsonParser parser, DeserializationContext cont throw new IOException("Missing 'timestamp' field in EventLogRecord JSON"); } + if (rootNode.has("eventAttributes")) { + return deserializeNormalizedRecord(mapper, rootNode, timestampNode.asText()); + } + + return deserializeLegacyRecord(mapper, rootNode, timestampNode.asText()); + } + + private static EventLogRecord deserializeNormalizedRecord( + ObjectMapper mapper, JsonNode rootNode, String timestamp) throws IOException { + String eventType = getText(rootNode, "eventType", true); + UUID eventId = parseUuid(getText(rootNode, "eventId", true)); + JsonNode attributesNode = rootNode.get("eventAttributes"); + if (attributesNode == null || !attributesNode.isObject()) { + throw new IOException( + "Field 'eventAttributes' must be an object in EventLogRecord JSON"); + } + Map attributes = + mapper.convertValue(attributesNode, new TypeReference>() {}); + boolean executionLifecycleEvent = + ExecutionLifecycleEvents.isExecutionLifecycleEvent(eventType); + String status = executionLifecycleEvent ? getText(rootNode, "status", false) : null; + String problemCategory = + executionLifecycleEvent ? getText(rootNode, "problemCategory", false) : null; + if (executionLifecycleEvent) { + if (status != null) { + attributes.putIfAbsent(ExecutionLifecycleEvents.STATUS_ATTRIBUTE, status); + } + if (problemCategory != null) { + attributes.putIfAbsent( + ExecutionLifecycleEvents.PROBLEM_CATEGORY_ATTRIBUTE, problemCategory); + } + } + Event event = new Event(eventId, eventType, attributes); + return new EventLogRecord( + new EventContext(eventType, timestamp), traceContext(mapper, rootNode), event); + } + + private static EventLogRecord deserializeLegacyRecord( + ObjectMapper mapper, JsonNode rootNode, String timestamp) throws IOException { // Deserialize event as base Event. Any top-level "logLevel" field present in older log - // files is silently ignored — it is no longer part of the record. + // files is silently ignored — it is not part of EventLogRecord. JsonNode eventNode = rootNode.get("event"); if (eventNode == null) { throw new IOException("Missing 'event' field in EventLogRecord JSON"); } - String eventType = getEventType(eventNode); - + String eventType = getLegacyEventType(eventNode); Event event = mapper.treeToValue(stripMetaFields(eventNode), Event.class); - EventContext eventContext = new EventContext(eventType, timestampNode.asText()); + return new EventLogRecord(new EventContext(eventType, timestamp), null, event); + } - return new EventLogRecord(eventContext, event); + private static ExecutionTraceContext traceContext(ObjectMapper mapper, JsonNode rootNode) + throws IOException { + String inputRunId = getText(rootNode, "inputRunId", false); + String businessKey = getText(rootNode, "businessKey", false); + String agentName = getText(rootNode, "agentName", false); + String executionId = getText(rootNode, "executionId", false); + String parentExecutionId = getText(rootNode, "parentExecutionId", false); + String entityType = getText(rootNode, "entityType", false); + String entityName = getText(rootNode, "entityName", false); + Map entityMetadata = getObjectMap(mapper, rootNode.get("entityMetadata")); + if (inputRunId == null + && businessKey == null + && agentName == null + && executionId == null + && parentExecutionId == null + && entityType == null + && entityName == null + && (entityMetadata == null || entityMetadata.isEmpty())) { + return null; + } + return ExecutionTraceContext.fromExistingIds( + inputRunId, + businessKey, + agentName, + executionId, + parentExecutionId, + entityType, + entityName, + entityMetadata); } - private static String getEventType(JsonNode eventNode) throws IOException { + private static String getLegacyEventType(JsonNode eventNode) throws IOException { JsonNode eventTypeNode = eventNode.get("eventType"); if (eventTypeNode == null || !eventTypeNode.isTextual()) { throw new IOException("Missing 'eventType' field in event JSON"); @@ -88,4 +161,41 @@ private static JsonNode stripMetaFields(JsonNode eventNode) { } return eventNode; } + + private static String getText(JsonNode node, String field, boolean required) + throws IOException { + JsonNode value = node.get(field); + if (value == null || value.isNull()) { + if (required) { + throw new IOException("Missing '" + field + "' field in EventLogRecord JSON"); + } + return null; + } + if (!value.isTextual()) { + throw new IOException("Field '" + field + "' must be textual in EventLogRecord JSON"); + } + return value.asText(); + } + + private static UUID parseUuid(String value) throws IOException { + try { + return UUID.fromString(value); + } catch (IllegalArgumentException e) { + throw new IOException("Invalid 'eventId' field in EventLogRecord JSON", e); + } + } + + private static Map getObjectMap(ObjectMapper mapper, JsonNode node) + throws IOException { + if (node == null || node.isNull()) { + return null; + } + if (!node.isObject()) { + throw new IOException( + "Field 'entityMetadata' must be an object in EventLogRecord JSON"); + } + Map result = + mapper.convertValue(node, new TypeReference>() {}); + return result == null ? new HashMap<>() : result; + } } diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/EventLogRecordJsonSerializer.java b/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/EventLogRecordJsonSerializer.java index 0a035603c..49de9d10b 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/EventLogRecordJsonSerializer.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/EventLogRecordJsonSerializer.java @@ -19,39 +19,35 @@ package org.apache.flink.agents.runtime.eventlog; import com.fasterxml.jackson.core.JsonGenerator; -import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.JsonSerializer; -import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.SerializerProvider; -import com.fasterxml.jackson.databind.node.ObjectNode; import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.trace.ExecutionLifecycleEvents; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; import java.io.IOException; -import java.util.Iterator; +import java.util.LinkedHashMap; import java.util.Map; /** * Custom JSON serializer for {@link EventLogRecord}. * - *

This serializer handles the serialization of EventLogRecord instances to JSON format suitable - * for structured logging. The serialization includes: - * - *

    - *
  • Top-level timestamp - *
  • Top-level eventType (routing key, mirrors {@code event.eventType}) - *
  • Event data serialized as a standard JSON object - *
- * - *

The resulting JSON structure is: + *

This serializer emits the normalized Event Log record shape used by downstream trace and + * aggregation consumers. Event identity, event type, attributes, input-run context, and execution + * hierarchy context are flattened at the top level. * * 

{@code
  * {
  *   "timestamp": "2024-01-15T10:30:00Z",
- *   "eventType": "_input_event",
- *   "event": {
- *     "eventType": "_input_event",
- *     // Event fields serialized normally
- *   }
+ *   "inputRunId": "...",
+ *   "businessKey": "...",
+ *   "executionId": "...",
+ *   "entityType": "action",
+ *   "entityName": "process",
+ *   "eventId": "...",
+ *   "eventType": "_execution_started_event",
+ *   "status": "started",
+ *   "eventAttributes": {}
  * }
  * }
*/ @@ -61,73 +57,64 @@ public class EventLogRecordJsonSerializer extends JsonSerializer public void serialize(EventLogRecord record, JsonGenerator gen, SerializerProvider serializers) throws IOException { - ObjectMapper mapper = (ObjectMapper) gen.getCodec(); - if (mapper == null) { - mapper = new ObjectMapper(); - } - + Event event = record.getEvent(); + ExecutionTraceContext traceContext = record.getExecutionTraceContext(); gen.writeStartObject(); - gen.writeStringField("timestamp", record.getContext().getTimestamp()); - gen.writeStringField("eventType", record.getEvent().getType()); - - gen.writeFieldName("event"); - JsonNode eventNode = buildEventNode(record.getEvent(), mapper); - if (!eventNode.isObject()) { - throw new IllegalStateException( - "Event log payload must be a JSON object, but was: " + eventNode.getNodeType()); + gen.writeStringField("timestamp", record.getEventContext().getTimestamp()); + if (traceContext != null) { + writeStringFieldIfPresent(gen, "inputRunId", traceContext.getInputRunId()); + writeStringFieldIfPresent(gen, "businessKey", traceContext.getBusinessKey()); + writeStringFieldIfPresent(gen, "agentName", traceContext.getAgentName()); + writeStringFieldIfPresent(gen, "executionId", traceContext.getExecutionId()); + writeStringFieldIfPresent( + gen, "parentExecutionId", traceContext.getParentExecutionId()); + writeStringFieldIfPresent(gen, "entityType", traceContext.getEntityType()); + writeStringFieldIfPresent(gen, "entityName", traceContext.getEntityName()); + writeMapFieldIfPresent(gen, "entityMetadata", traceContext.getEntityMetadata()); } - eventNode = reorderEventFields((ObjectNode) eventNode, record.getEvent(), mapper); - gen.writeTree(eventNode); + gen.writeStringField("eventId", event.getId().toString()); + gen.writeStringField("eventType", event.getType()); + writeStringFieldIfPresent( + gen, + "status", + executionLifecycleAttribute(event, ExecutionLifecycleEvents.STATUS_ATTRIBUTE)); + writeStringFieldIfPresent( + gen, + "problemCategory", + executionLifecycleAttribute( + event, ExecutionLifecycleEvents.PROBLEM_CATEGORY_ATTRIBUTE)); + gen.writeObjectField("eventAttributes", eventAttributes(event)); gen.writeEndObject(); } - private JsonNode buildEventNode(Event event, ObjectMapper mapper) { - JsonNode eventNode = mapper.valueToTree(event); - if (eventNode.isObject()) { - ObjectNode objectNode = (ObjectNode) eventNode; - objectNode.put("eventType", event.getType()); - objectNode.remove("sourceTimestamp"); + private static Map eventAttributes(Event event) { + Map attributes = new LinkedHashMap<>(event.getAttributes()); + if (ExecutionLifecycleEvents.isExecutionLifecycleEvent(event.getType())) { + attributes.remove(ExecutionLifecycleEvents.STATUS_ATTRIBUTE); + attributes.remove(ExecutionLifecycleEvents.PROBLEM_CATEGORY_ATTRIBUTE); } - return eventNode; + return attributes; } - private ObjectNode reorderEventFields(ObjectNode original, Event event, ObjectMapper mapper) { - ObjectNode ordered = mapper.createObjectNode(); - - // eventType — routing key (user-defined string) - JsonNode eventTypeNode = original.get("eventType"); - if (eventTypeNode != null) { - ordered.set("eventType", eventTypeNode); - } else { - ordered.put("eventType", event.getType()); - } - - JsonNode idNode = original.get("id"); - if (idNode != null) { - ordered.set("id", idNode); - } else if (event.getId() != null) { - ordered.put("id", event.getId().toString()); + private static String executionLifecycleAttribute(Event event, String name) { + if (!ExecutionLifecycleEvents.isExecutionLifecycleEvent(event.getType())) { + return null; } + Object value = event.getAttr(name); + return value == null ? null : String.valueOf(value); + } - JsonNode attributesNode = original.get("attributes"); - if (attributesNode != null) { - ordered.set("attributes", attributesNode); - } else { - ordered.putObject("attributes"); + private static void writeStringFieldIfPresent(JsonGenerator gen, String field, String value) + throws IOException { + if (value != null) { + gen.writeStringField(field, value); } + } - Iterator> fields = original.fields(); - while (fields.hasNext()) { - Map.Entry entry = fields.next(); - String fieldName = entry.getKey(); - if ("sourceTimestamp".equals(fieldName)) { - continue; - } - if (!ordered.has(fieldName)) { - ordered.set(fieldName, entry.getValue()); - } + private static void writeMapFieldIfPresent( + JsonGenerator gen, String field, Map value) throws IOException { + if (value != null && !value.isEmpty()) { + gen.writeObjectField(field, value); } - - return ordered; } } diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/EventLogWriter.java b/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/EventLogWriter.java new file mode 100644 index 000000000..d652c63f5 --- /dev/null +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/EventLogWriter.java @@ -0,0 +1,147 @@ +/* + * 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.eventlog; + +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.EventContext; +import org.apache.flink.agents.api.logger.EventLogger; +import org.apache.flink.agents.api.logger.EventLoggerConfig; +import org.apache.flink.agents.api.logger.EventLoggerFactory; +import org.apache.flink.agents.api.logger.EventLoggerOpenParams; +import org.apache.flink.agents.api.logger.LoggerType; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; +import org.apache.flink.agents.plan.AgentPlan; +import org.apache.flink.agents.runtime.metrics.BuiltInMetrics; +import org.apache.flink.annotation.Internal; +import org.apache.flink.annotation.VisibleForTesting; +import org.apache.flink.streaming.api.operators.StreamingRuntimeContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nullable; + +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.BASE_LOG_DIR; +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.EVENT_LOGGER_TYPE; +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.EVENT_LOG_TRACE_ENABLED; + +/** Operator-scoped writer that owns the physical Event Log logger lifecycle. */ +@Internal +public final class EventLogWriter implements AutoCloseable { + + private static final Logger LOG = LoggerFactory.getLogger(EventLogWriter.class); + + @Nullable private final EventLogger eventLogger; + private final boolean traceEnabled; + + public static EventLogWriter create(AgentPlan agentPlan) { + return new EventLogWriter( + createEventLogger(agentPlan), agentPlan.getConfig().get(EVENT_LOG_TRACE_ENABLED)); + } + + @VisibleForTesting + public static EventLogWriter forEventLogger(@Nullable EventLogger eventLogger) { + return forEventLogger(eventLogger, true); + } + + @VisibleForTesting + public static EventLogWriter forEventLogger( + @Nullable EventLogger eventLogger, boolean traceEnabled) { + return new EventLogWriter(eventLogger, traceEnabled); + } + + private EventLogWriter(@Nullable EventLogger eventLogger, boolean traceEnabled) { + this.eventLogger = eventLogger; + this.traceEnabled = traceEnabled; + } + + public void open(StreamingRuntimeContext runtimeContext, BuiltInMetrics builtInMetrics) + throws Exception { + if (eventLogger == null) { + return; + } + eventLogger.open(new EventLoggerOpenParams(runtimeContext)); + if (eventLogger instanceof FileEventLogger) { + ((FileEventLogger) eventLogger) + .setTruncatedEventsCounter(builtInMetrics.getEventLogTruncatedEventsCounter()); + } else if (eventLogger instanceof Slf4jEventLogger) { + ((Slf4jEventLogger) eventLogger) + .setTruncatedEventsCounter(builtInMetrics.getEventLogTruncatedEventsCounter()); + } + } + + /** Appends and flushes a business Event best-effort. */ + public void appendBusinessEventAndFlush( + EventContext eventContext, Event event, @Nullable ExecutionTraceContext traceContext) { + appendAndFlush(eventContext, event, traceEnabled ? traceContext : null); + } + + /** Appends and flushes an execution lifecycle Event when Trace recording is enabled. */ + public void appendExecutionEventAndFlush(Event event, ExecutionTraceContext traceContext) { + if (!traceEnabled) { + return; + } + appendAndFlush(new EventContext(event), event, traceContext); + } + + private void appendAndFlush( + EventContext eventContext, Event event, @Nullable ExecutionTraceContext traceContext) { + if (eventLogger == null) { + return; + } + try { + eventLogger.append(eventContext, event, traceContext); + eventLogger.flush(); + } catch (Exception logError) { + LOG.debug("Event Log write failed and was ignored.", logError); + } + } + + @VisibleForTesting + @Nullable + public EventLogger getEventLogger() { + return eventLogger; + } + + private static EventLogger createEventLogger(AgentPlan agentPlan) { + // Honor the EVENT_LOGGER_TYPE config, defaulting to SLF4J so events surface in the Flink + // Web UI by default. An explicit baseLogDir forces the file logger for backward + // compatibility with the existing file-based logging path. + LoggerType loggerType = agentPlan.getConfig().get(EVENT_LOGGER_TYPE); + String baseLogDir = agentPlan.getConfig().get(BASE_LOG_DIR); + if (baseLogDir != null && !baseLogDir.trim().isEmpty()) { + loggerType = LoggerType.FILE; + } + // The full agent config is the single source of truth for logger settings (baseLogDir, + // prettyPrint, event-log levels, truncation limits). Each logger pulls what it needs. + EventLoggerConfig config = + EventLoggerConfig.builder() + .loggerType(loggerType) + .property( + EventLoggerConfig.AGENT_CONFIG_PROPERTY_KEY, + agentPlan.getConfig().getConfData()) + .build(); + return EventLoggerFactory.createLogger(config); + } + + @Override + public void close() throws Exception { + if (eventLogger != null) { + eventLogger.close(); + } + } +} diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/FileEventLogger.java b/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/FileEventLogger.java index 9cedb6326..37180201f 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/FileEventLogger.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/FileEventLogger.java @@ -28,6 +28,7 @@ import org.apache.flink.agents.api.logger.EventLogger; import org.apache.flink.agents.api.logger.EventLoggerConfig; import org.apache.flink.agents.api.logger.EventLoggerOpenParams; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; import org.apache.flink.metrics.Counter; import java.io.BufferedWriter; @@ -37,6 +38,7 @@ import java.nio.file.Path; import java.nio.file.Paths; import java.util.Collections; +import java.util.Iterator; import java.util.Map; /** @@ -174,7 +176,13 @@ private String generateSubTaskLogFilePath(EventLoggerOpenParams params, String b } @Override - public void append(EventContext context, Event event) throws Exception { + public void append(EventContext eventContext, Event event) throws Exception { + append(eventContext, event, null); + } + + @Override + public void append(EventContext eventContext, Event event, ExecutionTraceContext traceContext) + throws Exception { if (writer == null) { throw new IllegalStateException("FileEventLogger not initialized. Call open() first."); } @@ -190,7 +198,7 @@ public void append(EventContext context, Event event) throws Exception { // All events should be JSON serializable; we already check this when sending events // to context (RunnerContextImpl.sendEvent). - EventLogRecord record = new EventLogRecord(context, event); + EventLogRecord record = new EventLogRecord(eventContext, traceContext, event); JsonNode tree = MAPPER.valueToTree(record); if (!(tree instanceof ObjectNode)) { throw new IllegalStateException( @@ -199,24 +207,18 @@ public void append(EventContext context, Event event) throws Exception { } ObjectNode rootNode = (ObjectNode) tree; - // Truncate the event subtree at STANDARD level. + // Truncate event attributes at STANDARD level. if (level == EventLogLevel.STANDARD && truncator != null) { - JsonNode eventNode = rootNode.get("event"); - if (eventNode instanceof ObjectNode) { - boolean truncated = truncator.truncate((ObjectNode) eventNode); + JsonNode attributesNode = rootNode.get("eventAttributes"); + if (attributesNode instanceof ObjectNode) { + boolean truncated = truncator.truncate((ObjectNode) attributesNode); if (truncated && truncatedEventsCounter != null) { truncatedEventsCounter.inc(); } } } - // Rebuild the top-level object so logLevel sits between timestamp and eventType, - // matching the documented JSON layout. - ObjectNode ordered = MAPPER.createObjectNode(); - ordered.set("timestamp", rootNode.get("timestamp")); - ordered.put("logLevel", level.name()); - ordered.set("eventType", rootNode.get("eventType")); - ordered.set("event", rootNode.get("event")); + ObjectNode ordered = withLogLevel(rootNode, level); String json = prettyPrint @@ -225,6 +227,20 @@ public void append(EventContext context, Event event) throws Exception { writer.println(json); } + private static ObjectNode withLogLevel(ObjectNode rootNode, EventLogLevel level) { + ObjectNode ordered = MAPPER.createObjectNode(); + ordered.set("timestamp", rootNode.get("timestamp")); + ordered.put("logLevel", level.name()); + Iterator> fields = rootNode.fields(); + while (fields.hasNext()) { + Map.Entry field = fields.next(); + if (!"timestamp".equals(field.getKey())) { + ordered.set(field.getKey(), field.getValue()); + } + } + return ordered; + } + @Override public void flush() throws Exception { if (writer == null) { diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/Slf4jEventLogger.java b/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/Slf4jEventLogger.java index d5f00f872..769c9fd28 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/Slf4jEventLogger.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/Slf4jEventLogger.java @@ -28,6 +28,7 @@ import org.apache.flink.agents.api.logger.EventLogger; import org.apache.flink.agents.api.logger.EventLoggerConfig; import org.apache.flink.agents.api.logger.EventLoggerOpenParams; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; import org.apache.flink.metrics.Counter; import org.apache.logging.log4j.Level; import org.apache.logging.log4j.LogManager; @@ -40,6 +41,7 @@ import org.slf4j.LoggerFactory; import java.util.Collections; +import java.util.Iterator; import java.util.Map; /** @@ -147,7 +149,13 @@ private static int getIntFromConfig(Map config, String key, int } @Override - public void append(EventContext context, Event event) throws Exception { + public void append(EventContext eventContext, Event event) throws Exception { + append(eventContext, event, null); + } + + @Override + public void append(EventContext eventContext, Event event, ExecutionTraceContext traceContext) + throws Exception { // Resolve log level and skip OFF events. EventLogLevel level = levelResolver != null @@ -157,7 +165,7 @@ public void append(EventContext context, Event event) throws Exception { return; } - EventLogRecord record = new EventLogRecord(context, event); + EventLogRecord record = new EventLogRecord(eventContext, traceContext, event); JsonNode tree = MAPPER.valueToTree(record); if (!(tree instanceof ObjectNode)) { throw new IllegalStateException( @@ -166,27 +174,18 @@ public void append(EventContext context, Event event) throws Exception { } ObjectNode rootNode = (ObjectNode) tree; - // Truncate the event subtree at STANDARD level. + // Truncate event attributes at STANDARD level. if (level == EventLogLevel.STANDARD && truncator != null) { - JsonNode eventNode = rootNode.get("event"); - if (eventNode instanceof ObjectNode) { - boolean truncated = truncator.truncate((ObjectNode) eventNode); + JsonNode attributesNode = rootNode.get("eventAttributes"); + if (attributesNode instanceof ObjectNode) { + boolean truncated = truncator.truncate((ObjectNode) attributesNode); if (truncated && truncatedEventsCounter != null) { truncatedEventsCounter.inc(); } } } - // Rebuild the top-level object so logLevel/subtask context sit alongside the documented - // JSON layout: timestamp, logLevel, eventType, subtask context, event. - ObjectNode ordered = MAPPER.createObjectNode(); - ordered.set("timestamp", rootNode.get("timestamp")); - ordered.put("logLevel", level.name()); - ordered.set("eventType", rootNode.get("eventType")); - ordered.put("jobId", jobId); - ordered.put("taskName", taskName); - ordered.put("subtaskId", subtaskId); - ordered.set("event", rootNode.get("event")); + ObjectNode ordered = withLogLevelAndSubtaskContext(rootNode, level); String json = prettyPrint @@ -195,6 +194,23 @@ public void append(EventContext context, Event event) throws Exception { EVENT_LOG.info(json); } + private ObjectNode withLogLevelAndSubtaskContext(ObjectNode rootNode, EventLogLevel level) { + ObjectNode ordered = MAPPER.createObjectNode(); + ordered.set("timestamp", rootNode.get("timestamp")); + ordered.put("logLevel", level.name()); + ordered.put("jobId", jobId); + ordered.put("taskName", taskName); + ordered.put("subtaskId", subtaskId); + Iterator> fields = rootNode.fields(); + while (fields.hasNext()) { + Map.Entry field = fields.next(); + if (!"timestamp".equals(field.getKey())) { + ordered.set(field.getKey(), field.getValue()); + } + } + return ordered; + } + @Override public void flush() throws Exception { // No-op: SLF4J/log4j2 handles flushing 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..6cba0de64 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 @@ -21,6 +21,9 @@ import org.apache.flink.agents.api.OutputEvent; import org.apache.flink.agents.api.agents.AgentExecutionOptions; import org.apache.flink.agents.api.context.MemoryUpdate; +import org.apache.flink.agents.api.trace.ExecutionLifecycleEvents; +import org.apache.flink.agents.api.trace.ExecutionReporter; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; import org.apache.flink.agents.plan.AgentPlan; import org.apache.flink.agents.plan.JavaFunction; import org.apache.flink.agents.plan.PythonFunction; @@ -28,10 +31,12 @@ import org.apache.flink.agents.runtime.ResourceCache; import org.apache.flink.agents.runtime.actionstate.ActionState; import org.apache.flink.agents.runtime.actionstate.ActionStateStore; +import org.apache.flink.agents.runtime.eventlog.EventLogWriter; import org.apache.flink.agents.runtime.memory.Mem0LongTermMemory; import org.apache.flink.agents.runtime.metrics.BuiltInMetrics; import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl; import org.apache.flink.agents.runtime.python.operator.PythonActionTask; +import org.apache.flink.agents.runtime.trace.ExecutionEventLogger; import org.apache.flink.agents.runtime.utils.EventUtil; import org.apache.flink.annotation.VisibleForTesting; import org.apache.flink.api.common.operators.MailboxExecutor; @@ -107,6 +112,10 @@ public class ActionExecutionOperator extends AbstractStreamOperator eventRouter; + private final transient ExecutionEventLogger executionEventLogger; + + private final transient EventLogWriter eventLogWriter; + private final transient DurableExecutionManager durableExecManager; private transient OperatorStateManager stateManager; @@ -127,7 +136,9 @@ public ActionExecutionOperator( this.agentPlan = agentPlan; this.processingTimeService = processingTimeService; this.mailboxExecutor = mailboxExecutor; - this.eventRouter = new EventRouter<>(agentPlan, inputIsJava); + this.eventLogWriter = EventLogWriter.create(agentPlan); + this.eventRouter = new EventRouter<>(agentPlan, inputIsJava, eventLogWriter); + this.executionEventLogger = ExecutionEventLogger.forEventLogWriter(eventLogWriter); this.durableExecManager = new DurableExecutionManager(actionStateStore); OperatorUtils.setChainStrategy(this, ChainingStrategy.ALWAYS); } @@ -192,8 +203,7 @@ public void open() throws Exception { mailboxProcessor = getMailboxProcessor(); - // Initialize the event logger if it is set. - eventRouter.initEventLogger(getRuntimeContext()); + eventLogWriter.open(getRuntimeContext(), builtInMetrics); // Initialize user event listeners from configuration eventRouter.initEventListeners(getRuntimeContext()); @@ -241,7 +251,16 @@ public void processElement(StreamRecord record) throws Exception { * `tryProcessActionTaskForKey` to continue processing. */ private void processEvent(Object key, Event event) throws Exception { - eventRouter.notifyEventProcessed(event); + processEvent( + key, + event, + ExecutionTraceContext.forInputRun( + businessKeyFromKey(key), agentPlan.getAgentName())); + } + + private void processEvent(Object key, Event event, ExecutionTraceContext traceContext) + throws Exception { + eventRouter.notifyEventProcessed(event, traceContext); boolean isInputEvent = EventUtil.isInputEvent(event); if (EventUtil.isOutputEvent(event)) { @@ -269,7 +288,8 @@ private void processEvent(Object key, Event event) throws Exception { List triggerActions = eventRouter.getActionsTriggeredBy(event, agentPlan); if (triggerActions != null && !triggerActions.isEmpty()) { for (Action triggerAction : triggerActions) { - stateManager.addActionTask(createActionTask(key, triggerAction, event)); + stateManager.addActionTask( + createActionTask(key, triggerAction, event, traceContext)); } } } @@ -327,10 +347,12 @@ private void processActionTaskForKey(Object key) throws Exception { stateManager.getSensoryMemState(), stateManager.getShortTermMemState(), pythonBridge.getPythonRunnerContext(), - ltm); + ltm, + executionEventLogger); long sequenceNumber = stateManager.getSequenceNumber(); boolean isFinished; + boolean notifyFinished = false; List outputEvents; Optional generatedActionTaskOpt = Optional.empty(); ActionState actionState = @@ -359,6 +381,7 @@ private void processActionTaskForKey(Object key) throws Exception { .getSensoryMemory() .set(memoryUpdate.getPath(), memoryUpdate.getValue()); } + notifyActionReused(actionTask); } else { // Initialize ActionState if not exists, or use existing one for recovery if (actionState == null) { @@ -369,36 +392,54 @@ private void processActionTaskForKey(Object key) throws Exception { key, sequenceNumber, actionTask.action, actionTask.event); } - // Set up durable execution context for fine-grained recovery - durableExecManager.setupDurableExecutionContext( - actionTask, actionState, sequenceNumber); - - ActionTask.ActionTaskResult actionTaskResult = - actionTask.invoke( - getRuntimeContext().getUserCodeClassLoader(), - this.pythonBridge.getPythonActionExecutor()); - - // We remove the contexts from the map after the task is processed. They will be added - // back later if the action task has a generated action task, meaning it is not - // finished. - contextManager.removeMemoryContext(actionTask); - durableExecManager.removeDurableContext(actionTask); - contextManager.removeContinuationContext(actionTask); - contextManager.removePythonAwaitableRef(actionTask); - durableExecManager.maybePersistTaskResult( - key, - sequenceNumber, - actionTask.action, - actionTask.event, - actionTask.getRunnerContext(), - actionTaskResult); - isFinished = actionTaskResult.isFinished(); - outputEvents = actionTaskResult.getOutputEvents(); - generatedActionTaskOpt = actionTaskResult.getGeneratedActionTask(); + notifyActionStarted(actionTask); + try { + // Set up durable execution context for fine-grained recovery + durableExecManager.setupDurableExecutionContext( + actionTask, actionState, sequenceNumber); + + ActionTask.ActionTaskResult actionTaskResult = + actionTask.invoke( + getRuntimeContext().getUserCodeClassLoader(), + this.pythonBridge.getPythonActionExecutor()); + + // Drop task-local contexts after each step; continuations transfer them back. + contextManager.removeMemoryContext(actionTask); + durableExecManager.removeDurableContext(actionTask); + contextManager.removeContinuationContext(actionTask); + contextManager.removePythonAwaitableRef(actionTask); + durableExecManager.maybePersistTaskResult( + key, + sequenceNumber, + actionTask.action, + actionTask.event, + actionTask.getRunnerContext(), + actionTaskResult); + isFinished = actionTaskResult.isFinished(); + outputEvents = actionTaskResult.getOutputEvents(); + generatedActionTaskOpt = actionTaskResult.getGeneratedActionTask(); + notifyFinished = isFinished; + } catch (Exception e) { + try { + notifyActionFailed(actionTask, e); + } finally { + contextManager.completeActionExecution(actionTask); + } + throw e; + } } - for (Event actionOutputEvent : outputEvents) { - processEvent(key, actionOutputEvent); + try { + for (Event actionOutputEvent : outputEvents) { + processEvent(key, actionOutputEvent, actionTask.getTraceContext()); + } + if (notifyFinished) { + notifyActionFinished(actionTask); + } + } finally { + if (isFinished) { + contextManager.completeActionExecution(actionTask); + } } boolean currentInputEventFinished = false; @@ -478,8 +519,8 @@ public void close() throws Exception { if (pythonBridge != null) { pythonBridge.close(); } - if (eventRouter != null) { - eventRouter.close(); + if (eventLogWriter != null) { + eventLogWriter.close(); } if (durableExecManager != null) { durableExecManager.close(); @@ -544,17 +585,54 @@ private void checkMailboxThread() { "Expected to be running on the task mailbox thread, but was not."); } - private ActionTask createActionTask(Object key, Action action, Event event) { + private void notifyActionStarted(ActionTask actionTask) { + if (actionTask.hasExecutionStartedEventEmitted()) { + return; + } + notifyExecutionLifecycleEvent( + actionTask.getTraceContext(), ExecutionLifecycleEvents.executionStarted()); + actionTask.markExecutionStartedEventEmitted(); + } + + private void notifyActionFinished(ActionTask actionTask) { + notifyExecutionLifecycleEvent( + actionTask.getTraceContext(), ExecutionLifecycleEvents.executionFinished()); + } + + private void notifyActionReused(ActionTask actionTask) { + notifyExecutionLifecycleEvent( + actionTask.getTraceContext(), ExecutionLifecycleEvents.executionReused()); + } + + private void notifyActionFailed(ActionTask actionTask, Exception error) { + notifyExecutionLifecycleEvent( + actionTask.getTraceContext(), + ExecutionLifecycleEvents.executionFailed( + error, ExecutionReporter.ProblemCategories.ACTION_EXECUTION_FAILED)); + } + + private void notifyExecutionLifecycleEvent(ExecutionTraceContext traceContext, Event event) { + executionEventLogger.emit(event, traceContext); + } + + private ActionTask createActionTask( + Object key, Action action, Event event, ExecutionTraceContext sourceTraceContext) { + ExecutionTraceContext actionTraceContext = + ExecutionTraceContext.forAction(sourceTraceContext, action.getName()); if (action.getExec() instanceof JavaFunction) { - return new JavaActionTask(key, event, action); + return new JavaActionTask(key, event, action, actionTraceContext); } else if (action.getExec() instanceof PythonFunction) { - return new PythonActionTask(key, event, action); + return new PythonActionTask(key, event, action, actionTraceContext); } else { throw new IllegalStateException( "Unsupported action type: " + action.getExec().getClass()); } } + private static String businessKeyFromKey(Object key) { + return key == null ? null : String.valueOf(key); + } + private void tryResumeProcessActionTasks() throws Exception { Iterable keys = stateManager.getProcessingKeys(); if (keys != null) { 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..627744adc 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 @@ -18,6 +18,8 @@ package org.apache.flink.agents.runtime.operator; import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.trace.ExecutionReporter; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; import org.apache.flink.agents.plan.actions.Action; import org.apache.flink.agents.runtime.context.RunnerContextImpl; import org.apache.flink.agents.runtime.python.utils.PythonActionExecutor; @@ -26,6 +28,7 @@ import javax.annotation.Nullable; +import java.io.Serializable; import java.util.List; import java.util.Objects; import java.util.Optional; @@ -40,13 +43,18 @@ * {@code ActionTask} via {@link ActionTaskResult#getGeneratedActionTask()} and continue executing * it. */ -public abstract class ActionTask { +public abstract class ActionTask implements Serializable { + + private static final long serialVersionUID = 1L; protected static final Logger LOG = LoggerFactory.getLogger(ActionTask.class); protected final Object key; protected final Event event; protected final Action action; + protected final ExecutionTraceContext traceContext; + + private boolean executionStartedEventEmitted; /** * Since RunnerContextImpl contains references to the Operator and state, it should not be * serialized and included in the state with ActionTask. Instead, we should check if a valid @@ -55,9 +63,20 @@ public abstract class ActionTask { protected transient RunnerContextImpl runnerContext; public ActionTask(Object key, Event event, Action action) { + this( + key, + event, + action, + ExecutionTraceContext.forExecution( + null, null, null, ExecutionReporter.EntityTypes.ACTION, action.getName())); + } + + protected ActionTask( + Object key, Event event, Action action, ExecutionTraceContext traceContext) { this.key = key; this.event = event; this.action = action; + this.traceContext = Objects.requireNonNull(traceContext, "traceContext must not be null"); } public RunnerContextImpl getRunnerContext() { @@ -72,6 +91,32 @@ public Object getKey() { return key; } + ExecutionTraceContext getTraceContext() { + return traceContext; + } + + void inheritLifecycleState(ActionTask source) { + if (source == this) { + return; + } + this.executionStartedEventEmitted = source.executionStartedEventEmitted; + } + + /** + * Returns whether the started lifecycle event has already been emitted for this execution. + * + *

This state is part of the pending continuation task so a resumed continuation does not + * emit duplicate started events for the same execution. + */ + boolean hasExecutionStartedEventEmitted() { + return executionStartedEventEmitted; + } + + /** Marks the started lifecycle event as emitted for this execution. */ + void markExecutionStartedEventEmitted() { + executionStartedEventEmitted = true; + } + @Override public boolean equals(Object o) { if (this == o) return true; @@ -79,12 +124,13 @@ public boolean equals(Object o) { ActionTask other = (ActionTask) o; return Objects.equals(this.key, other.key) && Objects.equals(this.event, other.event) - && Objects.equals(this.action, other.action); + && Objects.equals(this.action, other.action) + && Objects.equals(this.traceContext, other.traceContext); } @Override public int hashCode() { - return Objects.hash(key, event, action); + return Objects.hash(key, event, action, traceContext); } /** Invokes the action task. */ diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java index eba25d5d4..538250053 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java @@ -17,6 +17,7 @@ */ package org.apache.flink.agents.runtime.operator; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; import org.apache.flink.agents.plan.AgentPlan; import org.apache.flink.agents.plan.JavaFunction; import org.apache.flink.agents.plan.PythonFunction; @@ -30,6 +31,8 @@ import org.apache.flink.agents.runtime.memory.MemoryObjectImpl; import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl; import org.apache.flink.agents.runtime.python.context.PythonRunnerContextImpl; +import org.apache.flink.agents.runtime.trace.ExecutionEventSink; +import org.apache.flink.agents.runtime.trace.ReportedExecutionKey; import org.apache.flink.api.common.state.MapState; import javax.annotation.Nullable; @@ -45,9 +48,11 @@ *

    *
  • The shared (Java) {@link RunnerContextImpl} that is reused across action tasks via {@link * RunnerContextImpl#switchActionContext}. - *
  • Three per-{@link ActionTask} maps that survive across the boundary between a finishing - * action and the action it generates: memory contexts, continuation contexts (for async Java + *
  • Three per-{@link ActionTask} maps that survive across the boundary between one task and its + * generated continuation task: memory contexts, continuation contexts (for async Java * actions), and Python awaitable references. + *
  • Active child-execution reports, keyed by Action execution id, that pair start and terminal + * reports across continuation tasks without entering Flink state. *
  • The {@link ContinuationActionExecutor} thread pool used to run async Java continuations. *
* @@ -72,13 +77,15 @@ class ActionTaskContextManager implements AutoCloseable { private final Map actionTaskMemoryContexts; private final Map continuationContexts; private final Map pythonAwaitableRefs; - + private final Map> + activeReportedExecutionsByActionExecutionId; private ContinuationActionExecutor continuationActionExecutor; ActionTaskContextManager(int numAsyncThreads) { this.actionTaskMemoryContexts = new HashMap<>(); this.continuationContexts = new HashMap<>(); this.pythonAwaitableRefs = new HashMap<>(); + this.activeReportedExecutionsByActionExecutionId = new HashMap<>(); this.continuationActionExecutor = new ContinuationActionExecutor(numAsyncThreads); } @@ -147,8 +154,9 @@ RunnerContextImpl createOrGetRunnerContext( *
  • Selects a Java or Python runner context based on the action's {@code Exec} type. *
  • Reuses any existing {@link RunnerContextImpl.MemoryContext} for this task; otherwise * builds a fresh one backed by the supplied sensory/short-term memory states. + *
  • Wires the runtime-level execution event sink onto the runner context. *
  • Calls {@link RunnerContextImpl#switchActionContext} so the shared context now points at - * this action's name, memory, and key namespace. + * this action's name, memory, key namespace, trace context, and reported-execution state. *
  • For Java contexts, attaches a continuation context (re-used if the task is resuming * from an async suspend, fresh otherwise). *
  • For Python contexts, attaches the per-task awaitable reference (or {@code null} if the @@ -178,7 +186,8 @@ void createAndSetRunnerContext( MapState sensoryMemState, MapState shortTermMemState, PythonRunnerContextImpl pythonRunnerContext, - @Nullable InteranlBaseLongTermMemory longTermMemory) { + @Nullable InteranlBaseLongTermMemory longTermMemory, + @Nullable ExecutionEventSink executionEventSink) { RunnerContextImpl context; if (actionTask.action.getExec() instanceof JavaFunction) { context = @@ -206,6 +215,7 @@ void createAndSetRunnerContext( throw new IllegalStateException( "Unsupported action type: " + actionTask.action.getExec().getClass()); } + context.setExecutionEventSink(executionEventSink); RunnerContextImpl.MemoryContext memoryContext; if (actionTaskMemoryContexts.containsKey(actionTask)) { @@ -218,7 +228,11 @@ void createAndSetRunnerContext( } context.switchActionContext( - actionTask.action.getName(), memoryContext, String.valueOf(key.hashCode())); + actionTask.action.getName(), + memoryContext, + String.valueOf(key.hashCode()), + actionTask.getTraceContext(), + getOrCreateActiveReportedExecutions(actionTask)); if (context instanceof JavaRunnerContextImpl) { ContinuationContext continuationContext; @@ -266,6 +280,7 @@ RunnerContextImpl.MemoryContext removeMemoryContext(ActionTask actionTask) { void transferContexts( ActionTask fromTask, ActionTask toTask, DurableExecutionManager durableExecManager) { putMemoryContext(toTask, fromTask.getRunnerContext().getMemoryContext()); + toTask.inheritLifecycleState(fromTask); RunnerContextImpl.DurableExecutionContext durableContext = fromTask.getRunnerContext().getDurableExecutionContext(); if (durableContext != null) { @@ -285,6 +300,21 @@ void transferContexts( } } + void completeActionExecution(ActionTask actionTask) { + activeReportedExecutionsByActionExecutionId.remove( + actionTask.getTraceContext().getExecutionId()); + } + + private Map getOrCreateActiveReportedExecutions( + ActionTask actionTask) { + String executionId = actionTask.getTraceContext().getExecutionId(); + if (executionId == null) { + throw new IllegalStateException("Action execution id must not be null."); + } + return activeReportedExecutionsByActionExecutionId.computeIfAbsent( + executionId, ignored -> new HashMap<>()); + } + @Nullable ContinuationContext getContinuationContext(ActionTask actionTask) { return continuationContexts.get(actionTask); diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/EventRouter.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/EventRouter.java index 7eecc5d26..f575e84ef 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/EventRouter.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/EventRouter.java @@ -24,14 +24,10 @@ import org.apache.flink.agents.api.OutputEvent; import org.apache.flink.agents.api.listener.EventListener; import org.apache.flink.agents.api.logger.EventLogger; -import org.apache.flink.agents.api.logger.EventLoggerConfig; -import org.apache.flink.agents.api.logger.EventLoggerFactory; -import org.apache.flink.agents.api.logger.EventLoggerOpenParams; -import org.apache.flink.agents.api.logger.LoggerType; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; import org.apache.flink.agents.plan.AgentPlan; import org.apache.flink.agents.plan.actions.Action; -import org.apache.flink.agents.runtime.eventlog.FileEventLogger; -import org.apache.flink.agents.runtime.eventlog.Slf4jEventLogger; +import org.apache.flink.agents.runtime.eventlog.EventLogWriter; import org.apache.flink.agents.runtime.metrics.BuiltInMetrics; import org.apache.flink.agents.runtime.operator.queue.SegmentedQueue; import org.apache.flink.agents.runtime.python.utils.PythonActionExecutor; @@ -49,21 +45,18 @@ import java.util.ArrayList; import java.util.List; -import static org.apache.flink.agents.api.configuration.AgentConfigOptions.BASE_LOG_DIR; import static org.apache.flink.agents.api.configuration.AgentConfigOptions.EVENT_LISTENERS; -import static org.apache.flink.agents.api.configuration.AgentConfigOptions.EVENT_LOGGER_TYPE; import static org.apache.flink.util.Preconditions.checkState; /** * Handles event-side concerns for {@link ActionExecutionOperator}: input/output transformation - * between Java/Python representations, action lookup against the {@link AgentPlan}, event-logger - * and event-listener notification, and watermark draining via the per-key segment queue. + * between Java/Python representations, action lookup against the {@link AgentPlan}, Event Log + * writing, event-listener notification, and watermark draining via the per-key segment queue. * *

    Owned state: * *

      - *
    • The {@link EventLogger} created from the agent plan's logging configuration (may be {@code - * null} when logging is disabled). + *
    • The shared {@link EventLogWriter} owned by the operator. *
    • The list of registered {@link EventListener}s. *
    • A reused {@link StreamRecord} used to emit outputs without per-record allocation. *
    • The {@link SegmentedQueue} that orders watermarks behind in-flight keys so a watermark is @@ -73,18 +66,16 @@ * *

      Lifecycle: instantiated in the operator constructor (which decides {@link #inputIsJava}). * {@link #open(BuiltInMetrics)} runs from the operator's {@code open()} once metrics are available. - * {@link #initEventLogger} also runs from the operator's {@code open()} once the runtime context is - * available (after metrics have been built). {@link #close()} closes the event logger. * *

      Design constraint: package-private; no manager-to-manager held references. * * @param input record type * @param output record type */ -class EventRouter implements AutoCloseable { +class EventRouter { private final boolean inputIsJava; - private final EventLogger eventLogger; + private final EventLogWriter eventLogWriter; private final List eventListeners; private final AgentPlan agentPlan; private StreamRecord reusedStreamRecord; @@ -92,14 +83,18 @@ class EventRouter implements AutoCloseable { private BuiltInMetrics builtInMetrics; EventRouter(AgentPlan agentPlan, boolean inputIsJava) { - this(agentPlan, inputIsJava, createEventLogger(agentPlan)); + this(agentPlan, inputIsJava, EventLogWriter.create(agentPlan)); } @VisibleForTesting EventRouter(AgentPlan agentPlan, boolean inputIsJava, EventLogger eventLogger) { + this(agentPlan, inputIsJava, EventLogWriter.forEventLogger(eventLogger)); + } + + EventRouter(AgentPlan agentPlan, boolean inputIsJava, EventLogWriter eventLogWriter) { this.agentPlan = agentPlan; this.inputIsJava = inputIsJava; - this.eventLogger = eventLogger; + this.eventLogWriter = eventLogWriter; this.eventListeners = new ArrayList<>(); } @@ -107,8 +102,9 @@ class EventRouter implements AutoCloseable { * Initializes mutable runtime state that depends on metrics being available. * *

      Allocates the reused stream record and the segmented watermark queue, and stores the - * supplied {@link BuiltInMetrics} for use in {@link #notifyEventProcessed(Event)}. Called from - * the operator's {@code open()} once metric groups are constructed. + * supplied {@link BuiltInMetrics} for use in {@link #notifyEventProcessed(Event, + * ExecutionTraceContext)}. Called from the operator's {@code open()} once metric groups are + * constructed. * * @param builtInMetrics the operator's built-in metrics handle. */ @@ -118,20 +114,6 @@ void open(BuiltInMetrics builtInMetrics) { this.builtInMetrics = builtInMetrics; } - void initEventLogger(StreamingRuntimeContext runtimeContext) throws Exception { - if (eventLogger == null) { - return; - } - eventLogger.open(new EventLoggerOpenParams(runtimeContext)); - if (eventLogger instanceof FileEventLogger) { - ((FileEventLogger) eventLogger) - .setTruncatedEventsCounter(builtInMetrics.getEventLogTruncatedEventsCounter()); - } else if (eventLogger instanceof Slf4jEventLogger) { - ((Slf4jEventLogger) eventLogger) - .setTruncatedEventsCounter(builtInMetrics.getEventLogTruncatedEventsCounter()); - } - } - /** * Initializes the {@link EventListener}s configured for this agent. * @@ -223,23 +205,20 @@ List getActionsTriggeredBy(Event event, AgentPlan agentPlan) { /** * Notifies the configured event sinks (logger, listeners, metrics) that an event was processed. * - *

      If event logging is enabled, appends and immediately flushes the event. Then notifies - * every registered {@link EventListener}. Finally increments the {@code eventProcessed} - * built-in metric. The event logger is flushed per call as a temporary measure pending a - * batched flush mechanism. + *

      If event logging is enabled, appends and immediately flushes the event best-effort. Then + * notifies every registered {@link EventListener}. Finally increments the {@code + * eventProcessed} built-in metric. The event logger is flushed per call as a temporary measure + * pending a batched flush mechanism. * * @param event the event that was just processed. */ void notifyEventProcessed(Event event) throws Exception { + notifyEventProcessed(event, null); + } + + void notifyEventProcessed(Event event, ExecutionTraceContext traceContext) throws Exception { EventContext eventContext = new EventContext(event); - if (eventLogger != null) { - // If event logging is enabled, we log the event along with its context. - eventLogger.append(eventContext, event); - // For now, we flush the event logger after each event to ensure immediate logging. - // This is a temporary solution to ensure that events are logged immediately. - // TODO: In the future, we may want to implement a more efficient batching mechanism. - eventLogger.flush(); - } + eventLogWriter.appendBusinessEventAndFlush(eventContext, event, traceContext); if (eventListeners != null) { // Notify all registered event listeners about the event. for (EventListener listener : eventListeners) { @@ -277,7 +256,7 @@ StreamRecord getReusedStreamRecord() { @VisibleForTesting @Nullable EventLogger getEventLogger() { - return eventLogger; + return eventLogWriter.getEventLogger(); } @VisibleForTesting @@ -285,34 +264,6 @@ void addEventListener(EventListener listener) { eventListeners.add(listener); } - private static EventLogger createEventLogger(AgentPlan agentPlan) { - // Honor the EVENT_LOGGER_TYPE config, defaulting to SLF4J so events surface in the Flink - // Web UI by default. An explicit baseLogDir forces the file logger for backward - // compatibility with the existing file-based logging path. - LoggerType loggerType = agentPlan.getConfig().get(EVENT_LOGGER_TYPE); - String baseLogDir = agentPlan.getConfig().get(BASE_LOG_DIR); - if (baseLogDir != null && !baseLogDir.trim().isEmpty()) { - loggerType = LoggerType.FILE; - } - // The full agent config is the single source of truth for logger settings (baseLogDir, - // prettyPrint, event-log levels, truncation limits). Each logger pulls what it needs. - EventLoggerConfig config = - EventLoggerConfig.builder() - .loggerType(loggerType) - .property( - EventLoggerConfig.AGENT_CONFIG_PROPERTY_KEY, - agentPlan.getConfig().getConfData()) - .build(); - return EventLoggerFactory.createLogger(config); - } - - @Override - public void close() throws Exception { - if (eventLogger != null) { - eventLogger.close(); - } - } - @FunctionalInterface interface WatermarkEmitter { void emit(Watermark mark) throws Exception; diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/JavaActionTask.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/JavaActionTask.java index 11724ce68..4fdba5edf 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/JavaActionTask.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/JavaActionTask.java @@ -18,6 +18,7 @@ package org.apache.flink.agents.runtime.operator; import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; import org.apache.flink.agents.plan.JavaFunction; import org.apache.flink.agents.plan.actions.Action; import org.apache.flink.agents.runtime.context.JavaRunnerContextImpl; @@ -45,6 +46,12 @@ public JavaActionTask(Object key, Event event, Action action) { checkState(action.getExec() instanceof JavaFunction); } + public JavaActionTask( + Object key, Event event, Action action, ExecutionTraceContext traceContext) { + super(key, event, action, traceContext); + checkState(action.getExec() instanceof JavaFunction); + } + @Override public ActionTaskResult invoke(ClassLoader userCodeClassLoader, PythonActionExecutor executor) throws Exception { diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/python/context/PythonRunnerContextImpl.java b/runtime/src/main/java/org/apache/flink/agents/runtime/python/context/PythonRunnerContextImpl.java index f14d4bcff..633538725 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/python/context/PythonRunnerContextImpl.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/python/context/PythonRunnerContextImpl.java @@ -18,7 +18,10 @@ package org.apache.flink.agents.runtime.python.context; +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.ObjectMapper; import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.trace.ExecutionLifecycleEvents; import org.apache.flink.agents.plan.AgentPlan; import org.apache.flink.agents.runtime.ResourceCache; import org.apache.flink.agents.runtime.context.RunnerContextImpl; @@ -27,11 +30,18 @@ import javax.annotation.concurrent.NotThreadSafe; import java.io.IOException; +import java.util.LinkedHashMap; +import java.util.Map; /** A specialized {@link RunnerContext} that is specifically used when executing Python actions. */ @NotThreadSafe public class PythonRunnerContextImpl extends RunnerContextImpl { + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + + private static final TypeReference> METADATA_TYPE = + new TypeReference<>() {}; + /** * Reference to the Python awaitable object in the interpreter. This is set when a Python action * yields an awaitable and is used by PythonGeneratorActionTask to resume execution. @@ -61,6 +71,31 @@ public void sendEventJson(String eventJson) throws IOException { sendEvent(event); } + public void reportExecutionStartedJson( + String entityType, String entityName, String entityMetadataJson) throws Exception { + reportExecutionStarted(entityType, entityName, parseEntityMetadata(entityMetadataJson)); + } + + public void reportExecutionSucceededJson( + String entityType, String entityName, String entityMetadataJson) throws Exception { + reportExecutionSucceeded(entityType, entityName, parseEntityMetadata(entityMetadataJson)); + } + + public void reportExecutionFailedJson( + String entityType, + String entityName, + String entityMetadataJson, + String errorType, + String errorMessage, + String problemCategory) + throws Exception { + reportChildExecution( + entityType, + entityName, + parseEntityMetadata(entityMetadataJson), + ExecutionLifecycleEvents.executionFailed(errorType, errorMessage, problemCategory)); + } + public void checkMailboxThread() { // this method will be invoked by PythonActionExecutor's python interpreter. this.mailboxThreadChecker.run(); @@ -73,4 +108,12 @@ public String getPythonAwaitableRef() { public void setPythonAwaitableRef(String pythonAwaitableRef) { this.pythonAwaitableRef = pythonAwaitableRef; } + + private static Map parseEntityMetadata(String entityMetadataJson) + throws IOException { + if (entityMetadataJson == null || entityMetadataJson.isEmpty()) { + return Map.of(); + } + return OBJECT_MAPPER.readValue(entityMetadataJson, METADATA_TYPE); + } } diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonActionTask.java b/runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonActionTask.java index ce43195a5..5ec4a81b0 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonActionTask.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonActionTask.java @@ -18,6 +18,7 @@ package org.apache.flink.agents.runtime.python.operator; import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; import org.apache.flink.agents.plan.PythonFunction; import org.apache.flink.agents.plan.actions.Action; import org.apache.flink.agents.runtime.operator.ActionTask; @@ -39,6 +40,12 @@ public PythonActionTask(Object key, Event event, Action action) { checkState(action.getExec() instanceof PythonFunction); } + public PythonActionTask( + Object key, Event event, Action action, ExecutionTraceContext traceContext) { + super(key, event, action, traceContext); + checkState(action.getExec() instanceof PythonFunction); + } + public ActionTaskResult invoke(ClassLoader userCodeClassLoader, PythonActionExecutor executor) throws Exception { LOG.debug( @@ -58,7 +65,8 @@ public ActionTaskResult invoke(ClassLoader userCodeClassLoader, PythonActionExec // The Python action generates an awaitable. We need to execute it once, which will // submit an asynchronous task and return whether the action has been completed. ((PythonRunnerContextImpl) runnerContext).setPythonAwaitableRef(pythonAwaitableRef); - ActionTask tempGeneratedActionTask = new PythonGeneratorActionTask(key, event, action); + ActionTask tempGeneratedActionTask = + new PythonGeneratorActionTask(key, event, action, traceContext); tempGeneratedActionTask.setRunnerContext(runnerContext); return tempGeneratedActionTask.invoke(userCodeClassLoader, executor); } diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonGeneratorActionTask.java b/runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonGeneratorActionTask.java index 35296c1cc..a80db79dc 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonGeneratorActionTask.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonGeneratorActionTask.java @@ -18,6 +18,7 @@ package org.apache.flink.agents.runtime.python.operator; import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; import org.apache.flink.agents.plan.actions.Action; import org.apache.flink.agents.runtime.operator.ActionTask; import org.apache.flink.agents.runtime.python.context.PythonRunnerContextImpl; @@ -30,6 +31,11 @@ public PythonGeneratorActionTask(Object key, Event event, Action action) { super(key, event, action); } + public PythonGeneratorActionTask( + Object key, Event event, Action action, ExecutionTraceContext traceContext) { + super(key, event, action, traceContext); + } + @Override public ActionTaskResult invoke(ClassLoader userCodeClassLoader, PythonActionExecutor executor) throws Exception { @@ -47,7 +53,7 @@ public ActionTaskResult invoke(ClassLoader userCodeClassLoader, PythonActionExec "Python awaitable ref is null for action {} (likely restored from checkpoint), " + "re-executing from beginning.", action.getName()); - PythonActionTask freshTask = new PythonActionTask(key, event, action); + PythonActionTask freshTask = new PythonActionTask(key, event, action, traceContext); freshTask.setRunnerContext(runnerContext); return freshTask.invoke(userCodeClassLoader, executor); } diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/skill/LoadSkillTool.java b/runtime/src/main/java/org/apache/flink/agents/runtime/skill/LoadSkillTool.java index 6e08bd769..ef4ab42c3 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/skill/LoadSkillTool.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/skill/LoadSkillTool.java @@ -21,14 +21,18 @@ import org.apache.flink.agents.api.resource.ResourceContext; import org.apache.flink.agents.api.resource.ResourceDescriptor; import org.apache.flink.agents.api.tools.Tool; +import org.apache.flink.agents.api.tools.ToolExecutionMetadataProvider; import org.apache.flink.agents.api.tools.ToolMetadata; import org.apache.flink.agents.api.tools.ToolParameters; import org.apache.flink.agents.api.tools.ToolResponse; import org.apache.flink.agents.api.tools.ToolType; +import org.apache.flink.agents.api.trace.ToolExecutionMetadataKeys; import org.apache.flink.agents.runtime.resource.ResourceContextImpl; import java.nio.file.Path; +import java.util.LinkedHashMap; import java.util.List; +import java.util.Map; /** * Built-in tool that returns the body of a SKILL.md (default) or a specific bundled resource. @@ -36,7 +40,7 @@ *

      Mirrors the Python {@code flink_agents.runtime.skill.skill_tools.LoadSkillTool}. Auto-loaded * by {@code AgentPlan.addSkills} when an agent declares any {@code @Skills} method. */ -public class LoadSkillTool extends Tool { +public class LoadSkillTool extends Tool implements ToolExecutionMetadataProvider { private static final String DESCRIPTION = "Load a skill's content or a specific resource. Use this to access skill instructions and resources."; @@ -61,6 +65,22 @@ public ToolType getToolType() { return ToolType.FUNCTION; } + @Override + public Map getToolExecutionMetadata(ToolParameters parameters) { + Map metadata = new LinkedHashMap<>(); + if (parameters.hasParameter("name")) { + metadata.put( + ToolExecutionMetadataKeys.SKILL_NAME, + String.valueOf(parameters.getParameter("name"))); + } + if (parameters.hasParameter("path")) { + metadata.put( + ToolExecutionMetadataKeys.SKILL_RESOURCE_PATH, + String.valueOf(parameters.getParameter("path"))); + } + return metadata; + } + @Override public ToolResponse call(ToolParameters parameters) { String name = parameters.getParameter("name", String.class); diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/trace/ExecutionEventLogger.java b/runtime/src/main/java/org/apache/flink/agents/runtime/trace/ExecutionEventLogger.java new file mode 100644 index 000000000..21e092114 --- /dev/null +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/trace/ExecutionEventLogger.java @@ -0,0 +1,43 @@ +/* + * 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.trace; + +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; +import org.apache.flink.agents.runtime.eventlog.EventLogWriter; +import org.apache.flink.annotation.Internal; + +/** Logs execution events independently from the business-event router. */ +@Internal +public final class ExecutionEventLogger implements ExecutionEventSink { + + private final EventLogWriter eventLogWriter; + + public static ExecutionEventLogger forEventLogWriter(EventLogWriter eventLogWriter) { + return new ExecutionEventLogger(eventLogWriter); + } + + private ExecutionEventLogger(EventLogWriter eventLogWriter) { + this.eventLogWriter = eventLogWriter; + } + + @Override + public void emit(Event event, ExecutionTraceContext traceContext) { + eventLogWriter.appendExecutionEventAndFlush(event, traceContext); + } +} diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/trace/ExecutionEventSink.java b/runtime/src/main/java/org/apache/flink/agents/runtime/trace/ExecutionEventSink.java new file mode 100644 index 000000000..3eee1b24f --- /dev/null +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/trace/ExecutionEventSink.java @@ -0,0 +1,29 @@ +/* + * 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.trace; + +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; +import org.apache.flink.annotation.Internal; + +/** Runtime bridge for emitting execution lifecycle events to the event log pipeline. */ +@Internal +@FunctionalInterface +public interface ExecutionEventSink { + void emit(Event event, ExecutionTraceContext traceContext); +} diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/trace/ReportedExecutionKey.java b/runtime/src/main/java/org/apache/flink/agents/runtime/trace/ReportedExecutionKey.java new file mode 100644 index 000000000..105d2784e --- /dev/null +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/trace/ReportedExecutionKey.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.trace; + +import org.apache.flink.annotation.Internal; + +import java.io.Serializable; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Objects; + +/** Runtime key used to match lifecycle reports for the same reported execution. */ +@Internal +public final class ReportedExecutionKey implements Serializable { + private static final long serialVersionUID = 1L; + + private final String entityType; + private final String entityName; + private final Map entityMetadata; + + public ReportedExecutionKey( + String entityType, String entityName, Map entityMetadata) { + this.entityType = entityType; + this.entityName = entityName; + this.entityMetadata = + entityMetadata == null || entityMetadata.isEmpty() + ? Map.of() + : Collections.unmodifiableMap(new LinkedHashMap<>(entityMetadata)); + } + + public Map getEntityMetadata() { + return entityMetadata; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof ReportedExecutionKey)) { + return false; + } + ReportedExecutionKey that = (ReportedExecutionKey) o; + return Objects.equals(entityType, that.entityType) + && Objects.equals(entityName, that.entityName) + && Objects.equals(entityMetadata, that.entityMetadata); + } + + @Override + public int hashCode() { + return Objects.hash(entityType, entityName, entityMetadata); + } +} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextImplExecutionReporterTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextImplExecutionReporterTest.java new file mode 100644 index 000000000..e7e67e5a4 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextImplExecutionReporterTest.java @@ -0,0 +1,178 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.flink.agents.runtime.context; + +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.trace.ExecutionLifecycleEvents; +import org.apache.flink.agents.api.trace.ExecutionReporter; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; +import org.apache.flink.agents.runtime.python.context.PythonRunnerContextImpl; +import org.apache.flink.agents.runtime.trace.ReportedExecutionKey; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for execution reports emitted from {@link RunnerContextImpl}. */ +class RunnerContextImplExecutionReporterTest { + + @Test + void reportedExecutionReusesChildTraceContextBetweenStartAndFinish() throws Exception { + List reports = new ArrayList<>(); + RunnerContextImpl runnerContext = new RunnerContextImpl(null, () -> {}, null, null, "job"); + ExecutionTraceContext actionTraceContext = + ExecutionTraceContext.forInputRun("business-key", "agent") + .childExecution("action", "chat_model_action"); + runnerContext.setExecutionEventSink( + (event, context) -> reports.add(new RecordedReport(event, context))); + runnerContext.switchActionContext( + "chat_model_action", null, "business-key", actionTraceContext, new HashMap<>()); + + runnerContext.reportExecutionStarted( + ExecutionReporter.EntityTypes.LLM, "model-a", Map.of()); + runnerContext.reportExecutionSucceeded( + ExecutionReporter.EntityTypes.LLM, "model-a", Map.of()); + + assertThat(reports).hasSize(2); + RecordedReport started = reports.get(0); + RecordedReport finished = reports.get(1); + + assertThat(started.event.getType()) + .isEqualTo(ExecutionLifecycleEvents.EXECUTION_STARTED_EVENT_TYPE); + assertThat(finished.event.getType()) + .isEqualTo(ExecutionLifecycleEvents.EXECUTION_FINISHED_EVENT_TYPE); + assertThat(started.status()).isEqualTo(ExecutionLifecycleEvents.STATUS_STARTED); + assertThat(finished.status()).isEqualTo(ExecutionLifecycleEvents.STATUS_SUCCESS); + + assertThat(started.traceContext().getExecutionId()).isNotBlank(); + assertThat(finished.traceContext().getExecutionId()) + .isEqualTo(started.traceContext().getExecutionId()); + assertThat(started.traceContext().getParentExecutionId()) + .isEqualTo(actionTraceContext.getExecutionId()); + assertThat(started.traceContext().getEntityType()) + .isEqualTo(ExecutionReporter.EntityTypes.LLM); + assertThat(started.traceContext().getEntityName()).isEqualTo("model-a"); + } + + @Test + void reportedExecutionStateFollowsActionContextAcrossSwitches() throws Exception { + List reports = new ArrayList<>(); + RunnerContextImpl runnerContext = new RunnerContextImpl(null, () -> {}, null, null, "job"); + runnerContext.setExecutionEventSink( + (event, context) -> reports.add(new RecordedReport(event, context))); + + ExecutionTraceContext actionA = + ExecutionTraceContext.forInputRun("business-key", "agent") + .childExecution("action", "chat_model_action"); + ExecutionTraceContext actionB = + ExecutionTraceContext.forInputRun("business-key", "agent") + .childExecution("action", "tool_call_action"); + Map activeReportsA = new HashMap<>(); + Map activeReportsB = new HashMap<>(); + + runnerContext.switchActionContext( + "chat_model_action", null, "business-key", actionA, activeReportsA); + runnerContext.reportExecutionStarted( + ExecutionReporter.EntityTypes.LLM, "model-a", Map.of()); + + runnerContext.switchActionContext( + "tool_call_action", null, "business-key", actionB, activeReportsB); + runnerContext.reportExecutionStarted( + ExecutionReporter.EntityTypes.TOOL, "search", Map.of("toolCallId", "call-1")); + + runnerContext.switchActionContext( + "chat_model_action", null, "business-key", actionA, activeReportsA); + runnerContext.reportExecutionSucceeded( + ExecutionReporter.EntityTypes.LLM, "model-a", Map.of()); + + assertThat(reports).hasSize(3); + RecordedReport actionAStarted = reports.get(0); + RecordedReport actionBStarted = reports.get(1); + RecordedReport actionAFinished = reports.get(2); + + assertThat(actionAFinished.traceContext().getExecutionId()) + .isEqualTo(actionAStarted.traceContext().getExecutionId()); + assertThat(actionAFinished.traceContext().getParentExecutionId()) + .isEqualTo(actionA.getExecutionId()); + assertThat(actionBStarted.traceContext().getExecutionId()) + .isNotEqualTo(actionAStarted.traceContext().getExecutionId()); + } + + @Test + void pythonReporterBridgePreservesMetadataAndPythonErrorFields() throws Exception { + List reports = new ArrayList<>(); + PythonRunnerContextImpl runnerContext = + new PythonRunnerContextImpl(null, () -> {}, null, null, "job"); + ExecutionTraceContext actionTraceContext = + ExecutionTraceContext.forInputRun("business-key", "agent") + .childExecution("action", "tool_call_action"); + runnerContext.setExecutionEventSink( + (event, context) -> reports.add(new RecordedReport(event, context))); + runnerContext.switchActionContext( + "tool_call_action", null, "business-key", actionTraceContext, new HashMap<>()); + + String metadata = "{\"toolCallId\":\"call-1\",\"toolType\":\"function\"}"; + runnerContext.reportExecutionStartedJson( + ExecutionReporter.EntityTypes.TOOL, "search", metadata); + runnerContext.reportExecutionFailedJson( + ExecutionReporter.EntityTypes.TOOL, + "search", + metadata, + "builtins.ValueError", + "bad response", + ExecutionReporter.ProblemCategories.TOOL_CALL_FAILED); + + assertThat(reports).hasSize(2); + RecordedReport started = reports.get(0); + RecordedReport failed = reports.get(1); + + assertThat(failed.traceContext().getExecutionId()) + .isEqualTo(started.traceContext().getExecutionId()); + assertThat(failed.traceContext().getEntityMetadata()) + .containsEntry("toolCallId", "call-1") + .containsEntry("toolType", "function"); + assertThat(failed.event.getType()) + .isEqualTo(ExecutionLifecycleEvents.EXECUTION_FAILED_EVENT_TYPE); + assertThat(failed.event.getAttr("errorType")).isEqualTo("builtins.ValueError"); + assertThat(failed.event.getAttr("errorMessage")).isEqualTo("bad response"); + assertThat(failed.event.getAttr(ExecutionLifecycleEvents.PROBLEM_CATEGORY_ATTRIBUTE)) + .isEqualTo(ExecutionReporter.ProblemCategories.TOOL_CALL_FAILED); + } + + private static class RecordedReport { + private final Event event; + private final ExecutionTraceContext traceContext; + + private RecordedReport(Event event, ExecutionTraceContext traceContext) { + this.event = event; + this.traceContext = traceContext; + } + + private ExecutionTraceContext traceContext() { + return traceContext; + } + + private String status() { + return (String) event.getAttr(ExecutionLifecycleEvents.STATUS_ATTRIBUTE); + } + } +} 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..612bd4cd6 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 @@ -25,14 +25,22 @@ import org.apache.flink.agents.api.EventContext; import org.apache.flink.agents.api.InputEvent; import org.apache.flink.agents.api.OutputEvent; +import org.apache.flink.agents.api.trace.ExecutionLifecycleEvents; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import java.util.HashMap; +import java.util.LinkedHashMap; import java.util.Map; import java.util.UUID; -import static org.junit.jupiter.api.Assertions.*; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; /** * Unit tests for {@link EventLogRecordJsonSerializer} and {@link EventLogRecordJsonDeserializer}. @@ -48,127 +56,221 @@ void setUp() { @Test void testSerializeInputEvent() throws Exception { - // Given InputEvent inputEvent = new InputEvent("test input data"); - EventContext context = new EventContext(inputEvent); - EventLogRecord record = new EventLogRecord(context, inputEvent); + EventLogRecord record = record(inputEvent, null); - // When String json = objectMapper.writeValueAsString(record); - - // Then - assertNotNull(json); JsonNode jsonNode = objectMapper.readTree(json); - // Verify structure assertTrue(jsonNode.has("timestamp")); - assertTrue(jsonNode.has("event")); - - // Verify event - JsonNode eventNode = jsonNode.get("event"); - assertTrue(eventNode.has("eventType")); - assertEquals(InputEvent.EVENT_TYPE, eventNode.get("eventType").asText()); - assertFalse(eventNode.has("sourceTimestamp")); - // Data is stored in attributes - assertTrue(eventNode.has("attributes")); - assertEquals("test input data", eventNode.get("attributes").get("input").asText()); + assertTrue(jsonNode.has("eventId")); + assertEquals(InputEvent.EVENT_TYPE, jsonNode.get("eventType").asText()); + assertEquals("test input data", jsonNode.get("eventAttributes").get("input").asText()); + assertFalse(jsonNode.has("event")); + assertFalse(jsonNode.has("inputRunId")); + assertFalse(jsonNode.has("executionId")); + } + + @Test + void testRecordUsesEventContextTimestamp() throws Exception { + InputEvent inputEvent = new InputEvent("test input data"); + EventContext eventContext = new EventContext(InputEvent.EVENT_TYPE, "2026-01-01T00:00:00Z"); + + EventLogRecord record = record(eventContext, inputEvent, null); + String json = objectMapper.writeValueAsString(record); + JsonNode jsonNode = objectMapper.readTree(json); + + assertEquals(InputEvent.EVENT_TYPE, record.getEventContext().getEventType()); + assertEquals("2026-01-01T00:00:00Z", record.getEventContext().getTimestamp()); + assertEquals("2026-01-01T00:00:00Z", jsonNode.get("timestamp").asText()); + } + + @Test + void testRecordRejectsMismatchedEventContextType() { + InputEvent inputEvent = new InputEvent("test input data"); + EventContext eventContext = new EventContext("OtherEvent", "2026-01-01T00:00:00Z"); + + assertThrows(IllegalArgumentException.class, () -> record(eventContext, inputEvent, null)); } @Test void testSerializeOutputEvent() throws Exception { - // Given OutputEvent outputEvent = new OutputEvent("test output data"); - EventContext context = new EventContext(outputEvent); - EventLogRecord record = new EventLogRecord(context, outputEvent); + EventLogRecord record = record(outputEvent, null); - // When String json = objectMapper.writeValueAsString(record); - - // Then JsonNode jsonNode = objectMapper.readTree(json); - assertEquals(OutputEvent.EVENT_TYPE, jsonNode.get("event").get("eventType").asText()); - assertEquals( - "test output data", jsonNode.get("event").get("attributes").get("output").asText()); + + assertEquals(OutputEvent.EVENT_TYPE, jsonNode.get("eventType").asText()); + assertEquals("test output data", jsonNode.get("eventAttributes").get("output").asText()); } @Test void testSerializeCustomEvent() throws Exception { - // Given CustomTestEvent customEvent = new CustomTestEvent("custom data", 42, true); - EventContext context = new EventContext(customEvent); - EventLogRecord record = new EventLogRecord(context, customEvent); + EventLogRecord record = record(customEvent, null); - // When String json = objectMapper.writeValueAsString(record); - - // Then JsonNode jsonNode = objectMapper.readTree(json); - assertEquals(CustomTestEvent.EVENT_TYPE, jsonNode.get("event").get("eventType").asText()); - JsonNode attrsNode = jsonNode.get("event").get("attributes"); + assertEquals(CustomTestEvent.EVENT_TYPE, jsonNode.get("eventType").asText()); + JsonNode attrsNode = jsonNode.get("eventAttributes"); assertEquals("custom data", attrsNode.get("customData").asText()); assertEquals(42, attrsNode.get("customNumber").asInt()); - assertEquals(true, attrsNode.get("customFlag").asBoolean()); + assertTrue(attrsNode.get("customFlag").asBoolean()); + } + + @Test + void testBusinessEventKeepsStatusAttributesAsPayload() throws Exception { + Map attributes = new LinkedHashMap<>(); + attributes.put(ExecutionLifecycleEvents.STATUS_ATTRIBUTE, "business-status"); + attributes.put(ExecutionLifecycleEvents.PROBLEM_CATEGORY_ATTRIBUTE, "business-category"); + Event businessEvent = new Event("BusinessStatusEvent", attributes); + ExecutionTraceContext traceContext = + ExecutionTraceContext.fromExistingIds( + "input-run-1", + "business-key-1", + "agent-a", + "execution-1", + "parent-execution-1", + "action", + "process", + Map.of()); + EventLogRecord record = record(businessEvent, traceContext); + + String json = objectMapper.writeValueAsString(record); + JsonNode jsonNode = objectMapper.readTree(json); + + assertFalse(jsonNode.has("status")); + assertFalse(jsonNode.has("problemCategory")); + assertEquals( + "business-status", + jsonNode.get("eventAttributes") + .get(ExecutionLifecycleEvents.STATUS_ATTRIBUTE) + .asText()); + assertEquals( + "business-category", + jsonNode.get("eventAttributes") + .get(ExecutionLifecycleEvents.PROBLEM_CATEGORY_ATTRIBUTE) + .asText()); + + EventLogRecord deserializedRecord = objectMapper.readValue(json, EventLogRecord.class); + assertEquals( + "business-status", + deserializedRecord.getEvent().getAttr(ExecutionLifecycleEvents.STATUS_ATTRIBUTE)); + assertEquals( + "business-category", + deserializedRecord + .getEvent() + .getAttr(ExecutionLifecycleEvents.PROBLEM_CATEGORY_ATTRIBUTE)); + } + + @Test + void testSerializeAndDeserializeExecutionLifecycleFields() throws Exception { + Map entityMetadata = new LinkedHashMap<>(); + entityMetadata.put("mcpServer", "demo-server"); + + Event lifecycleEvent = + ExecutionLifecycleEvents.executionFailed( + new IllegalStateException("model returned malformed JSON"), + "model_output_parse_error"); + ExecutionTraceContext traceContext = + ExecutionTraceContext.fromExistingIds( + "input-run-1", + "business-key-1", + "agent-a", + "execution-1", + "parent-execution-1", + "action", + "process", + entityMetadata); + EventLogRecord record = record(lifecycleEvent, traceContext); + assertEquals(traceContext, record.getExecutionTraceContext()); + + String json = objectMapper.writeValueAsString(record); + JsonNode jsonNode = objectMapper.readTree(json); + + assertEquals("input-run-1", jsonNode.get("inputRunId").asText()); + assertEquals("business-key-1", jsonNode.get("businessKey").asText()); + assertEquals("agent-a", jsonNode.get("agentName").asText()); + assertEquals("execution-1", jsonNode.get("executionId").asText()); + assertEquals("parent-execution-1", jsonNode.get("parentExecutionId").asText()); + assertEquals("action", jsonNode.get("entityType").asText()); + assertEquals("process", jsonNode.get("entityName").asText()); + assertEquals("demo-server", jsonNode.get("entityMetadata").get("mcpServer").asText()); + assertEquals( + ExecutionLifecycleEvents.EXECUTION_FAILED_EVENT_TYPE, + jsonNode.get("eventType").asText()); + assertEquals("failed", jsonNode.get("status").asText()); + assertEquals("model_output_parse_error", jsonNode.get("problemCategory").asText()); + assertFalse(jsonNode.get("eventAttributes").has(ExecutionLifecycleEvents.STATUS_ATTRIBUTE)); + assertFalse( + jsonNode.get("eventAttributes") + .has(ExecutionLifecycleEvents.PROBLEM_CATEGORY_ATTRIBUTE)); + assertEquals( + IllegalStateException.class.getName(), + jsonNode.get("eventAttributes").get("errorType").asText()); + + EventLogRecord deserializedRecord = objectMapper.readValue(json, EventLogRecord.class); + ExecutionTraceContext deserializedTraceContext = + deserializedRecord.getExecutionTraceContext(); + assertEquals(traceContext, deserializedTraceContext); + assertNotNull(deserializedTraceContext); + assertEquals("input-run-1", deserializedTraceContext.getInputRunId()); + assertEquals("business-key-1", deserializedTraceContext.getBusinessKey()); + assertEquals("agent-a", deserializedTraceContext.getAgentName()); + assertEquals("execution-1", deserializedTraceContext.getExecutionId()); + assertEquals("parent-execution-1", deserializedTraceContext.getParentExecutionId()); + assertEquals("action", deserializedTraceContext.getEntityType()); + assertEquals("process", deserializedTraceContext.getEntityName()); + assertEquals("demo-server", deserializedTraceContext.getEntityMetadata().get("mcpServer")); + assertEquals( + "failed", + deserializedRecord.getEvent().getAttr(ExecutionLifecycleEvents.STATUS_ATTRIBUTE)); + assertEquals( + "model_output_parse_error", + deserializedRecord + .getEvent() + .getAttr(ExecutionLifecycleEvents.PROBLEM_CATEGORY_ATTRIBUTE)); } @Test void testDeserializeInputEvent() throws Exception { - // Given InputEvent originalEvent = new InputEvent("test input data"); - EventContext originalContext = new EventContext(originalEvent); - EventLogRecord originalRecord = new EventLogRecord(originalContext, originalEvent); + EventLogRecord originalRecord = record(originalEvent, null); String json = objectMapper.writeValueAsString(originalRecord); - // When EventLogRecord deserializedRecord = objectMapper.readValue(json, EventLogRecord.class); - // Then - assertNotNull(deserializedRecord); - assertNotNull(deserializedRecord.getContext()); - assertNotNull(deserializedRecord.getEvent()); - - // Verify context - EventContext deserializedContext = deserializedRecord.getContext(); - assertEquals(InputEvent.EVENT_TYPE, deserializedContext.getEventType()); - assertNotNull(deserializedContext.getTimestamp()); - - // Verify event via fromEvent + assertNotNull(deserializedRecord.getEventContext().getTimestamp()); + assertEquals(InputEvent.EVENT_TYPE, deserializedRecord.getEventContext().getEventType()); + assertNull(deserializedRecord.getExecutionTraceContext()); Event deserializedEvent = deserializedRecord.getEvent(); assertEquals(InputEvent.EVENT_TYPE, deserializedEvent.getType()); - InputEvent inputEvent = InputEvent.fromEvent(deserializedEvent); - assertEquals("test input data", inputEvent.getInput()); + assertEquals("test input data", InputEvent.fromEvent(deserializedEvent).getInput()); } @Test void testDeserializeOutputEvent() throws Exception { - // Given OutputEvent originalEvent = new OutputEvent("test output data"); - EventContext originalContext = new EventContext(originalEvent); - EventLogRecord originalRecord = new EventLogRecord(originalContext, originalEvent); + EventLogRecord originalRecord = record(originalEvent, null); String json = objectMapper.writeValueAsString(originalRecord); - // When EventLogRecord deserializedRecord = objectMapper.readValue(json, EventLogRecord.class); - // Then Event deserializedEvent = deserializedRecord.getEvent(); assertEquals(OutputEvent.EVENT_TYPE, deserializedEvent.getType()); - OutputEvent outputEvent = OutputEvent.fromEvent(deserializedEvent); - assertEquals("test output data", outputEvent.getOutput()); + assertEquals("test output data", OutputEvent.fromEvent(deserializedEvent).getOutput()); } @Test void testDeserializeCustomEvent() throws Exception { - // Given CustomTestEvent originalEvent = new CustomTestEvent("custom data", 42, true); - EventContext originalContext = new EventContext(originalEvent); - EventLogRecord originalRecord = new EventLogRecord(originalContext, originalEvent); + EventLogRecord originalRecord = record(originalEvent, null); String json = objectMapper.writeValueAsString(originalRecord); - // When EventLogRecord deserializedRecord = objectMapper.readValue(json, EventLogRecord.class); - // Then — deserialized as base Event; use fromEvent for typed access Event deserializedEvent = deserializedRecord.getEvent(); assertEquals(CustomTestEvent.EVENT_TYPE, deserializedEvent.getType()); CustomTestEvent customEvent = CustomTestEvent.fromEvent(deserializedEvent); @@ -177,90 +279,58 @@ void testDeserializeCustomEvent() throws Exception { assertTrue(customEvent.isCustomFlag()); } - @Test - void testRoundTripSerialization() throws Exception { - // Given - InputEvent originalEvent = new InputEvent("round trip test"); - EventContext originalContext = new EventContext(originalEvent); - EventLogRecord originalRecord = new EventLogRecord(originalContext, originalEvent); - - // When - serialize and deserialize - String json = objectMapper.writeValueAsString(originalRecord); - EventLogRecord deserializedRecord = objectMapper.readValue(json, EventLogRecord.class); - - // Then - verify all data is preserved - assertEquals( - originalContext.getEventType(), deserializedRecord.getContext().getEventType()); - assertEquals(InputEvent.EVENT_TYPE, deserializedRecord.getEvent().getType()); - - InputEvent deserializedEvent = InputEvent.fromEvent(deserializedRecord.getEvent()); - assertEquals(originalEvent.getInput(), deserializedEvent.getInput()); - } - - @Test - void testSerializeUnifiedEvent() throws Exception { - // Given - a unified event with user-defined type - java.util.Map attrs = new java.util.HashMap<>(); - attrs.put("msg", "hello"); - Event unifiedEvent = new Event("MyCustomEvent", attrs); - EventContext context = new EventContext(unifiedEvent); - EventLogRecord record = new EventLogRecord(context, unifiedEvent); - - // When - String json = objectMapper.writeValueAsString(record); - - // Then - JsonNode jsonNode = objectMapper.readTree(json); - JsonNode eventNode = jsonNode.get("event"); - - // eventType should be the user-defined type string - assertEquals("MyCustomEvent", eventNode.get("eventType").asText()); - // attributes should be present - assertEquals("hello", eventNode.get("attributes").get("msg").asText()); - } - - @Test - void testDeserializeUnifiedEvent() throws Exception { - // Given - a unified event serialized via the EventLogRecord serializer - java.util.Map attrs = new java.util.HashMap<>(); - attrs.put("key", "value"); - attrs.put("count", 42); - Event originalEvent = new Event("CustomType", attrs); - EventContext originalContext = new EventContext(originalEvent); - EventLogRecord originalRecord = new EventLogRecord(originalContext, originalEvent); - String json = objectMapper.writeValueAsString(originalRecord); - - // When - EventLogRecord deserializedRecord = objectMapper.readValue(json, EventLogRecord.class); - - // Then - EventContext deserializedContext = deserializedRecord.getContext(); - // eventType is the routing key (user-defined string) - assertEquals("CustomType", deserializedContext.getEventType()); - - // The event should be deserialized as a base Event with the type field set - Event deserializedEvent = deserializedRecord.getEvent(); - assertEquals("CustomType", deserializedEvent.getType()); - } - @Test void testRoundTripUnifiedEvent() throws Exception { - // Given - java.util.Map attrs = new java.util.HashMap<>(); + Map attrs = new HashMap<>(); attrs.put("x", 1); attrs.put("y", "two"); Event originalEvent = new Event("RoundTripEvent", attrs); - EventContext context = new EventContext(originalEvent); - EventLogRecord record = new EventLogRecord(context, originalEvent); + EventLogRecord record = record(originalEvent, null); - // When - serialize and deserialize String json = objectMapper.writeValueAsString(record); EventLogRecord deserialized = objectMapper.readValue(json, EventLogRecord.class); - // Then - assertEquals("RoundTripEvent", deserialized.getContext().getEventType()); + assertNotNull(deserialized.getEventContext().getTimestamp()); Event event = deserialized.getEvent(); assertEquals("RoundTripEvent", event.getType()); + assertEquals(1, ((Number) event.getAttr("x")).intValue()); + assertEquals("two", event.getAttr("y")); + } + + @Test + void testDeserializeLegacyRecord() throws Exception { + UUID eventId = UUID.randomUUID(); + String json = + "{" + + "\"timestamp\":\"2026-01-01T00:00:00Z\"," + + "\"event\":{" + + "\"id\":\"" + + eventId + + "\"," + + "\"type\":\"LegacyType\"," + + "\"eventType\":\"LegacyType\"," + + "\"eventClass\":\"LegacyEvent\"," + + "\"attributes\":{\"key\":\"value\"}" + + "}" + + "}"; + + EventLogRecord record = objectMapper.readValue(json, EventLogRecord.class); + + assertEquals("2026-01-01T00:00:00Z", record.getEventContext().getTimestamp()); + assertEquals("LegacyType", record.getEventContext().getEventType()); + assertEquals("LegacyType", record.getEvent().getType()); + assertEquals(eventId, record.getEvent().getId()); + assertEquals("value", record.getEvent().getAttr("key")); + assertNull(record.getExecutionTraceContext()); + } + + private static EventLogRecord record(Event event, ExecutionTraceContext executionTraceContext) { + return record(new EventContext(event), event, executionTraceContext); + } + + private static EventLogRecord record( + EventContext eventContext, Event event, ExecutionTraceContext executionTraceContext) { + return new EventLogRecord(eventContext, executionTraceContext, event); } /** Custom test event class using the attributes-based pattern. */ diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/EventLogWriterTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/EventLogWriterTest.java new file mode 100644 index 000000000..d0addfe29 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/EventLogWriterTest.java @@ -0,0 +1,112 @@ +/* + * 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.eventlog; + +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.EventContext; +import org.apache.flink.agents.api.InputEvent; +import org.apache.flink.agents.api.logger.EventLogger; +import org.apache.flink.agents.api.trace.ExecutionLifecycleEvents; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThatCode; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.isNull; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; + +/** Contract tests for {@link EventLogWriter}. */ +class EventLogWriterTest { + + @Test + void appendAndFlushIgnoresAppendFailure() throws Exception { + EventLogger mockLogger = mock(EventLogger.class); + EventLogWriter writer = EventLogWriter.forEventLogger(mockLogger); + InputEvent inputEvent = new InputEvent(1L); + EventContext eventContext = new EventContext(inputEvent); + doThrow(new RuntimeException("append failed")) + .when(mockLogger) + .append(any(EventContext.class), eq(inputEvent), isNull()); + + assertThatCode(() -> writer.appendBusinessEventAndFlush(eventContext, inputEvent, null)) + .doesNotThrowAnyException(); + verify(mockLogger, never()).flush(); + } + + @Test + void appendAndFlushIgnoresFlushFailure() throws Exception { + EventLogger mockLogger = mock(EventLogger.class); + EventLogWriter writer = EventLogWriter.forEventLogger(mockLogger); + InputEvent inputEvent = new InputEvent(1L); + EventContext eventContext = new EventContext(inputEvent); + doThrow(new RuntimeException("flush failed")).when(mockLogger).flush(); + + assertThatCode(() -> writer.appendBusinessEventAndFlush(eventContext, inputEvent, null)) + .doesNotThrowAnyException(); + verify(mockLogger).append(any(EventContext.class), eq(inputEvent), isNull()); + verify(mockLogger).flush(); + } + + @Test + void traceDisabledKeepsBusinessEventWithoutTraceContext() throws Exception { + EventLogger mockLogger = mock(EventLogger.class); + EventLogWriter writer = EventLogWriter.forEventLogger(mockLogger, false); + InputEvent inputEvent = new InputEvent(1L); + EventContext eventContext = new EventContext(inputEvent); + ExecutionTraceContext traceContext = + ExecutionTraceContext.forInputRun("business-key", "agent"); + + writer.appendBusinessEventAndFlush(eventContext, inputEvent, traceContext); + + verify(mockLogger).append(eq(eventContext), eq(inputEvent), isNull()); + verify(mockLogger).flush(); + } + + @Test + void traceDisabledSkipsExecutionEvent() { + EventLogger mockLogger = mock(EventLogger.class); + EventLogWriter writer = EventLogWriter.forEventLogger(mockLogger, false); + + writer.appendExecutionEventAndFlush( + ExecutionLifecycleEvents.executionStarted(), + ExecutionTraceContext.forInputRun("business-key", "agent") + .childExecution("action", "action1")); + + verifyNoInteractions(mockLogger); + } + + @Test + void traceEnabledWritesExecutionEventWithTraceContext() throws Exception { + EventLogger mockLogger = mock(EventLogger.class); + EventLogWriter writer = EventLogWriter.forEventLogger(mockLogger, true); + ExecutionTraceContext traceContext = + ExecutionTraceContext.forInputRun("business-key", "agent") + .childExecution("action", "action1"); + Event executionEvent = ExecutionLifecycleEvents.executionStarted(); + + writer.appendExecutionEventAndFlush(executionEvent, traceContext); + + verify(mockLogger).append(any(EventContext.class), eq(executionEvent), eq(traceContext)); + verify(mockLogger).flush(); + } +} 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..360b373fd 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 @@ -26,9 +26,11 @@ 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.EventLogger; import org.apache.flink.agents.api.logger.EventLoggerConfig; import org.apache.flink.agents.api.logger.EventLoggerOpenParams; import org.apache.flink.agents.api.logger.LoggerType; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; import org.apache.flink.api.common.JobID; import org.apache.flink.api.common.JobInfo; import org.apache.flink.api.common.TaskInfo; @@ -121,9 +123,9 @@ void testOpenCreatesLogFile() throws Exception { void testAppendWritesJsonEvents() throws Exception { logger.open(openParams); InputEvent inputEvent = new InputEvent("test input"); - EventContext context = new EventContext(inputEvent); + ExecutionTraceContext context = null; - logger.append(context, inputEvent); + append(logger, inputEvent, context); logger.flush(); Path logFile = getExpectedLogFilePath(); @@ -133,10 +135,12 @@ void testAppendWritesJsonEvents() throws Exception { EventLogRecord deserializedRecord = objectMapper.readValue(lines.get(0), EventLogRecord.class); assertNotNull(deserializedRecord, "Deserialized record should not be null"); - assertNotNull(deserializedRecord.getContext(), "Deserialized context should not be null"); + assertNotNull( + deserializedRecord.getEventContext().getTimestamp(), + "Deserialized timestamp should not be null"); assertNotNull(deserializedRecord.getEvent(), "Deserialized event should not be null"); - assertEquals(InputEvent.EVENT_TYPE, deserializedRecord.getContext().getEventType()); + assertEquals(InputEvent.EVENT_TYPE, deserializedRecord.getEvent().getType()); assertEquals(InputEvent.EVENT_TYPE, deserializedRecord.getEvent().getType()); InputEvent deserializedInput = InputEvent.fromEvent(deserializedRecord.getEvent()); assertEquals("test input", deserializedInput.getInput()); @@ -147,11 +151,11 @@ void testAppendMultipleEvents() throws Exception { logger.open(openParams); InputEvent inputEvent = new InputEvent("input data"); OutputEvent outputEvent = new OutputEvent("output data"); - EventContext inputContext = new EventContext(inputEvent); - EventContext outputContext = new EventContext(outputEvent); + ExecutionTraceContext inputContext = null; + ExecutionTraceContext outputContext = null; - logger.append(inputContext, inputEvent); - logger.append(outputContext, outputEvent); + append(logger, inputEvent, inputContext); + append(logger, outputEvent, outputContext); logger.flush(); Path logFile = getExpectedLogFilePath(); @@ -174,10 +178,10 @@ void testAppendWithCustomEvent() throws Exception { // Given logger.open(openParams); TestCustomEvent customEvent = new TestCustomEvent("custom data", 42); - EventContext context = new EventContext(customEvent); + ExecutionTraceContext context = null; // When - logger.append(context, customEvent); + append(logger, customEvent, context); logger.flush(); // Then @@ -187,9 +191,9 @@ void testAppendWithCustomEvent() throws Exception { // Verify JSON structure JsonNode jsonNode = objectMapper.readTree(lines.get(0)); - assertEquals(TestCustomEvent.EVENT_TYPE, jsonNode.get("event").get("eventType").asText()); + assertEquals(TestCustomEvent.EVENT_TYPE, jsonNode.get("eventType").asText()); - JsonNode attrsNode = jsonNode.get("event").get("attributes"); + JsonNode attrsNode = jsonNode.get("eventAttributes"); assertEquals("custom data", attrsNode.get("customData").asText()); assertEquals(42, attrsNode.get("customNumber").asInt()); @@ -210,16 +214,16 @@ void testAppendInAppendMode() throws Exception { // Given - first session logger.open(openParams); InputEvent event1 = new InputEvent("first event"); - EventContext context1 = new EventContext(event1); - logger.append(context1, event1); + ExecutionTraceContext context1 = null; + append(logger, event1, context1); logger.close(); // When - second session (append mode) FileEventLogger secondLogger = new FileEventLogger(config); secondLogger.open(openParams); InputEvent event2 = new InputEvent("second event"); - EventContext context2 = new EventContext(event2); - secondLogger.append(context2, event2); + ExecutionTraceContext context2 = null; + append(secondLogger, event2, context2); secondLogger.flush(); secondLogger.close(); @@ -230,13 +234,10 @@ void testAppendInAppendMode() throws Exception { // Verify JSON structure JsonNode firstEventJson = objectMapper.readTree(lines.get(0)); - assertEquals( - "first event", firstEventJson.get("event").get("attributes").get("input").asText()); + assertEquals("first event", firstEventJson.get("eventAttributes").get("input").asText()); JsonNode secondEventJson = objectMapper.readTree(lines.get(1)); - assertEquals( - "second event", - secondEventJson.get("event").get("attributes").get("input").asText()); + assertEquals("second event", secondEventJson.get("eventAttributes").get("input").asText()); // Verify deserialization via fromEvent EventLogRecord firstRecord = objectMapper.readValue(lines.get(0), EventLogRecord.class); @@ -253,8 +254,8 @@ void testMultipleSubTasks() throws Exception { // Given - subtask 0 logger.open(openParams); InputEvent event1 = new InputEvent("subtask 0 event"); - EventContext context1 = new EventContext(event1); - logger.append(context1, event1); + ExecutionTraceContext context1 = null; + append(logger, event1, context1); logger.flush(); // Given - subtask 1 @@ -263,8 +264,8 @@ void testMultipleSubTasks() throws Exception { EventLoggerOpenParams openParams2 = new EventLoggerOpenParams(runtimeContext); logger2.open(openParams2); InputEvent event2 = new InputEvent("subtask 1 event"); - EventContext context2 = new EventContext(event2); - logger2.append(context2, event2); + ExecutionTraceContext context2 = null; + append(logger2, event2, context2); logger2.flush(); logger2.close(); @@ -292,11 +293,9 @@ void testMultipleSubTasks() throws Exception { JsonNode subtask1EventJson = objectMapper.readTree(subtask1Lines.get(0)); assertEquals( - "subtask 0 event", - subtask0EventJson.get("event").get("attributes").get("input").asText()); + "subtask 0 event", subtask0EventJson.get("eventAttributes").get("input").asText()); assertEquals( - "subtask 1 event", - subtask1EventJson.get("event").get("attributes").get("input").asText()); + "subtask 1 event", subtask1EventJson.get("eventAttributes").get("input").asText()); // Verify deserialization via fromEvent EventLogRecord subtask0Record = @@ -320,7 +319,7 @@ void testPrettyPrintOutputsFormattedJson() throws Exception { logger.open(openParams); InputEvent inputEvent = new InputEvent("test input"); - logger.append(new EventContext(inputEvent), inputEvent); + append(logger, inputEvent, null); logger.flush(); // Then - output should be valid JSON spanning multiple lines (pretty-printed) @@ -355,9 +354,9 @@ void testStandardLevelTruncation() throws Exception { // Use a custom event with a very long string field TestCustomEvent event = new TestCustomEvent("this is a very long string that exceeds 10", 1); - EventContext context = new EventContext(event); + ExecutionTraceContext context = null; - logger.append(context, event); + append(logger, event, context); logger.flush(); Path logFile = getExpectedLogFilePath(); @@ -368,7 +367,7 @@ void testStandardLevelTruncation() throws Exception { assertEquals("STANDARD", jsonNode.get("logLevel").asText()); // The customData field (inside attributes) should be truncated - JsonNode attrsNode = jsonNode.get("event").get("attributes"); + JsonNode attrsNode = jsonNode.get("eventAttributes"); JsonNode customDataNode = attrsNode.get("customData"); assertTrue( customDataNode.has("truncatedString"), @@ -389,9 +388,9 @@ void testVerboseLevelNoTruncation() throws Exception { TestCustomEvent event = new TestCustomEvent("this is a very long string that exceeds 10", 1); - EventContext context = new EventContext(event); + ExecutionTraceContext context = null; - logger.append(context, event); + append(logger, event, context); logger.flush(); Path logFile = getExpectedLogFilePath(); @@ -402,7 +401,7 @@ void testVerboseLevelNoTruncation() throws Exception { assertEquals("VERBOSE", jsonNode.get("logLevel").asText()); // The customData field (inside attributes) should NOT be truncated - JsonNode attrsNode = jsonNode.get("event").get("attributes"); + JsonNode attrsNode = jsonNode.get("eventAttributes"); assertTrue( attrsNode.get("customData").isTextual(), "String should be preserved at VERBOSE level"); @@ -421,9 +420,9 @@ void testOffLevelSkipsEvent() throws Exception { logger.open(openParams); InputEvent event = new InputEvent("should not be logged"); - EventContext context = new EventContext(event); + ExecutionTraceContext context = null; - logger.append(context, event); + append(logger, event, context); logger.flush(); Path logFile = getExpectedLogFilePath(); @@ -445,12 +444,12 @@ void testPerTypeLevelOverride() throws Exception { // InputEvent should be VERBOSE (no truncation) InputEvent inputEvent = new InputEvent("this is a very long string that exceeds 10"); - logger.append(new EventContext(inputEvent), inputEvent); + append(logger, inputEvent, null); // TestCustomEvent should be STANDARD (truncated) TestCustomEvent customEvent = new TestCustomEvent("this is a very long string that exceeds 10", 1); - logger.append(new EventContext(customEvent), customEvent); + append(logger, customEvent, null); logger.flush(); Path logFile = getExpectedLogFilePath(); @@ -460,13 +459,12 @@ void testPerTypeLevelOverride() throws Exception { // InputEvent at VERBOSE - no truncation (data lives in attributes) JsonNode inputJson = objectMapper.readTree(lines.get(0)); assertEquals("VERBOSE", inputJson.get("logLevel").asText()); - assertTrue(inputJson.get("event").get("attributes").get("input").isTextual()); + assertTrue(inputJson.get("eventAttributes").get("input").isTextual()); // TestCustomEvent at STANDARD - truncated (data lives in attributes) JsonNode customJson = objectMapper.readTree(lines.get(1)); assertEquals("STANDARD", customJson.get("logLevel").asText()); - assertTrue( - customJson.get("event").get("attributes").get("customData").has("truncatedString")); + assertTrue(customJson.get("eventAttributes").get("customData").has("truncatedString")); } @Test @@ -474,9 +472,9 @@ void testJsonOutputHasNewFields() throws Exception { // Given - default config logger.open(openParams); InputEvent event = new InputEvent("test"); - EventContext context = new EventContext(event); + ExecutionTraceContext context = null; - logger.append(context, event); + append(logger, event, context); logger.flush(); Path logFile = getExpectedLogFilePath(); @@ -485,7 +483,9 @@ void testJsonOutputHasNewFields() throws Exception { // Verify new top-level fields exist assertTrue(jsonNode.has("logLevel"), "JSON should have logLevel field"); + assertTrue(jsonNode.has("eventId"), "JSON should have eventId field"); assertTrue(jsonNode.has("eventType"), "JSON should have eventType field"); + assertTrue(jsonNode.has("eventAttributes"), "JSON should have eventAttributes field"); assertEquals(InputEvent.EVENT_TYPE, jsonNode.get("eventType").asText()); assertNotNull(jsonNode.get("logLevel").asText()); } @@ -527,11 +527,11 @@ void testHierarchicalInheritance() throws Exception { // EventA has explicit VERBOSE override — should be logged TestNamespacedEventA eventA = new TestNamespacedEventA("should be logged"); - logger.append(new EventContext(eventA), eventA); + append(logger, eventA, null); // EventB inherits OFF from namespace level — should NOT be logged TestNamespacedEventB eventB = new TestNamespacedEventB("should not be logged"); - logger.append(new EventContext(eventB), eventB); + append(logger, eventB, null); logger.flush(); Path logFile = getExpectedLogFilePath(); @@ -549,6 +549,12 @@ private Path getExpectedLogFilePath() { "events-%s-%s-%d.log", testJobId.toString(), testTaskName, testSubTaskId)); } + private static void append( + EventLogger logger, Event event, ExecutionTraceContext executionTraceContext) + throws Exception { + logger.append(new EventContext(event), event, executionTraceContext); + } + /** Custom test event class using the attributes-based pattern. */ public static class TestCustomEvent extends Event { public static final String EVENT_TYPE = "TestCustomEvent"; diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/Slf4jEventLoggerTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/Slf4jEventLoggerTest.java index d3dc51f5a..2845e700a 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/Slf4jEventLoggerTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/Slf4jEventLoggerTest.java @@ -20,6 +20,7 @@ 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; import org.apache.flink.agents.api.InputEvent; import org.apache.flink.agents.api.OutputEvent; @@ -27,6 +28,7 @@ import org.apache.flink.agents.api.logger.EventLoggerConfig; import org.apache.flink.agents.api.logger.EventLoggerOpenParams; import org.apache.flink.agents.api.logger.LoggerType; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; import org.apache.flink.api.common.JobID; import org.apache.flink.api.common.JobInfo; import org.apache.flink.api.common.TaskInfo; @@ -127,9 +129,9 @@ void testAppendWritesJsonWithSubtaskContext() throws Exception { logger.open(openParams); InputEvent inputEvent = new InputEvent("test input"); - EventContext context = new EventContext(inputEvent); + ExecutionTraceContext context = null; - logger.append(context, inputEvent); + append(inputEvent, context); List messages = testAppender.getMessages(); assertEquals(1, messages.size(), "Should have logged one message"); @@ -141,7 +143,8 @@ void testAppendWritesJsonWithSubtaskContext() throws Exception { assertEquals(testSubTaskId, jsonNode.get("subtaskId").asInt()); // Verify event content assertNotNull(jsonNode.get("timestamp")); - assertNotNull(jsonNode.get("event")); + assertNotNull(jsonNode.get("eventId")); + assertNotNull(jsonNode.get("eventAttributes")); assertEquals(InputEvent.EVENT_TYPE, jsonNode.get("eventType").asText()); } @@ -154,18 +157,17 @@ void testAppendMultipleEvents() throws Exception { InputEvent inputEvent = new InputEvent("input data"); OutputEvent outputEvent = new OutputEvent("output data"); - logger.append(new EventContext(inputEvent), inputEvent); - logger.append(new EventContext(outputEvent), outputEvent); + append(inputEvent, null); + append(outputEvent, null); List messages = testAppender.getMessages(); assertEquals(2, messages.size(), "Should have logged two messages"); JsonNode inputJson = objectMapper.readTree(messages.get(0)); - assertEquals("input data", inputJson.get("event").get("attributes").get("input").asText()); + assertEquals("input data", inputJson.get("eventAttributes").get("input").asText()); JsonNode outputJson = objectMapper.readTree(messages.get(1)); - assertEquals( - "output data", outputJson.get("event").get("attributes").get("output").asText()); + assertEquals("output data", outputJson.get("eventAttributes").get("output").asText()); } @Test @@ -181,7 +183,7 @@ void testRootLevelOffSkipsAllEvents() throws Exception { logger.open(openParams); InputEvent inputEvent = new InputEvent("input data"); - logger.append(new EventContext(inputEvent), inputEvent); + append(inputEvent, null); List messages = testAppender.getMessages(); assertTrue(messages.isEmpty(), "No events should be logged when root level is OFF"); @@ -202,15 +204,15 @@ void testPerEventTypeOffSkipsMatchingEvents() throws Exception { InputEvent inputEvent = new InputEvent("input data"); OutputEvent outputEvent = new OutputEvent("output data"); - logger.append(new EventContext(inputEvent), inputEvent); - logger.append(new EventContext(outputEvent), outputEvent); + append(inputEvent, null); + append(outputEvent, null); List messages = testAppender.getMessages(); assertEquals( 1, messages.size(), "Only OutputEvent should be logged when InputEvent is OFF"); JsonNode jsonNode = objectMapper.readTree(messages.get(0)); - assertEquals("output data", jsonNode.get("event").get("attributes").get("output").asText()); + assertEquals("output data", jsonNode.get("eventAttributes").get("output").asText()); } @Test @@ -220,7 +222,7 @@ void testLogLevelAppearsInJson() throws Exception { logger.open(openParams); InputEvent inputEvent = new InputEvent("test input"); - logger.append(new EventContext(inputEvent), inputEvent); + append(inputEvent, null); List messages = testAppender.getMessages(); assertEquals(1, messages.size()); @@ -237,7 +239,7 @@ void testDefaultIsNotPrettyPrinted() throws Exception { logger.open(openParams); InputEvent inputEvent = new InputEvent("test input"); - logger.append(new EventContext(inputEvent), inputEvent); + append(inputEvent, null); List messages = testAppender.getMessages(); assertEquals(1, messages.size()); @@ -260,7 +262,7 @@ void testEnablePrettyPrint() throws Exception { logger.open(openParams); InputEvent inputEvent = new InputEvent("test input"); - logger.append(new EventContext(inputEvent), inputEvent); + append(inputEvent, null); List messages = testAppender.getMessages(); assertEquals(1, messages.size()); @@ -281,6 +283,10 @@ void testFlushAndCloseAreNoOps() throws Exception { assertDoesNotThrow(() -> logger.close(), "close() should not throw"); } + private void append(Event event, ExecutionTraceContext executionTraceContext) throws Exception { + logger.append(new EventContext(event), event, executionTraceContext); + } + /** A log4j2 appender that captures log messages for testing. */ private static class TestAppender extends AbstractAppender { 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..c2b916bb9 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 @@ -27,8 +27,14 @@ import org.apache.flink.agents.api.context.MemoryObject; import org.apache.flink.agents.api.context.RunnerContext; import org.apache.flink.agents.api.listener.EventListener; +import org.apache.flink.agents.api.logger.EventLogger; import org.apache.flink.agents.api.logger.EventLoggerConfig; +import org.apache.flink.agents.api.logger.EventLoggerFactory; +import org.apache.flink.agents.api.logger.EventLoggerOpenParams; import org.apache.flink.agents.api.logger.LoggerType; +import org.apache.flink.agents.api.trace.ExecutionLifecycleEvents; +import org.apache.flink.agents.api.trace.ExecutionReporter; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; import org.apache.flink.agents.plan.AgentConfiguration; import org.apache.flink.agents.plan.AgentPlan; import org.apache.flink.agents.plan.JavaFunction; @@ -37,6 +43,7 @@ import org.apache.flink.agents.runtime.actionstate.CallResult; import org.apache.flink.agents.runtime.actionstate.InMemoryActionStateStore; import org.apache.flink.agents.runtime.eventlog.FileEventLogger; +import org.apache.flink.agents.runtime.eventlog.Slf4jEventLogger; import org.apache.flink.api.common.typeinfo.TypeInformation; import org.apache.flink.api.java.functions.KeySelector; import org.apache.flink.runtime.checkpoint.OperatorSubtaskState; @@ -47,14 +54,17 @@ import org.apache.flink.streaming.util.AbstractStreamOperatorTestHarness; import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness; import org.apache.flink.util.ExceptionUtils; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import java.lang.reflect.Field; +import java.util.ArrayList; 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; @@ -73,6 +83,13 @@ public class ActionExecutionOperatorTest { void resetReconcilableFixtures() { TestAgent.resetReconcilableRecoveryFixture(); TestAgent.resetMixedRecoveryFixture(); + RecordingEventLogger.reset(); + EventLoggerFactory.registerFactory(LoggerType.SLF4J, config -> new RecordingEventLogger()); + } + + @AfterEach + void restoreEventLoggerFactory() { + EventLoggerFactory.registerFactory(LoggerType.SLF4J, Slf4jEventLogger::new); } @Test @@ -369,16 +386,143 @@ agentPlanWithStateStore, true, new InMemoryActionStateStore(false)), /** A EventListener for unit test */ public static class TestEventListener implements EventListener { public boolean called = false; + public final List eventTypes = new ArrayList<>(); @Override public void onEventProcessed(EventContext context, Event event) { this.called = true; + this.eventTypes.add(event.getType()); + } + } + + /** Records events appended to Event Log for assertions. */ + public static class RecordingEventLogger implements EventLogger { + private static final List EVENTS = new ArrayList<>(); + private static int createdCount; + private static int openCount; + private static int flushCount; + private static int closeCount; + + public RecordingEventLogger() { + createdCount++; + } + + static void reset() { + EVENTS.clear(); + createdCount = 0; + openCount = 0; + flushCount = 0; + closeCount = 0; + } + + static List events() { + return List.copyOf(EVENTS); + } + + static int createdCount() { + return createdCount; + } + + static int openCount() { + return openCount; + } + + static int flushCount() { + return flushCount; + } + + static int closeCount() { + return closeCount; + } + + @Override + public void open(EventLoggerOpenParams params) { + openCount++; + } + + @Override + public void append(EventContext eventContext, Event event) { + append(eventContext, event, null); + } + + @Override + public void append( + EventContext eventContext, Event event, ExecutionTraceContext traceContext) { + EVENTS.add(new RecordedEvent(event, traceContext)); + } + + @Override + public void flush() { + flushCount++; + } + + @Override + public void close() { + closeCount++; + } + } + + private static class RecordedEvent { + private final Event event; + private final ExecutionTraceContext traceContext; + + private RecordedEvent(Event event, ExecutionTraceContext traceContext) { + this.event = event; + this.traceContext = traceContext; } + + private ExecutionTraceContext traceContext() { + if (traceContext == null) { + throw new AssertionError("Missing ExecutionTraceContext"); + } + return traceContext; + } + + private String status() { + return (String) event.getAttr(ExecutionLifecycleEvents.STATUS_ATTRIBUTE); + } + + private String problemCategory() { + return (String) event.getAttr(ExecutionLifecycleEvents.PROBLEM_CATEGORY_ATTRIBUTE); + } + } + + private static RecordedEvent findRecordedLifecycleEvent( + String eventType, String entityName, String status) { + return RecordingEventLogger.events().stream() + .filter(record -> eventType.equals(record.event.getType())) + .filter(record -> entityName.equals(record.traceContext().getEntityName())) + .filter(record -> status.equals(record.status())) + .findFirst() + .orElseThrow( + () -> + new AssertionError( + String.format( + "Missing lifecycle event type=%s entity=%s status=%s in %s", + eventType, + entityName, + status, + RecordingEventLogger.events().stream() + .map( + record -> + record.event.getType() + + "/" + + record.traceContext() + .getEntityName() + + "/" + + record.status()) + .collect(Collectors.toList())))); + } + + private static AgentConfiguration traceEnabledConfig() { + AgentConfiguration config = new AgentConfiguration(); + config.set(AgentConfigOptions.EVENT_LOG_TRACE_ENABLED, true); + return config; } @Test void testEventListenersFromAgentConfig() throws Exception { - final AgentConfiguration config = new AgentConfiguration(); + final AgentConfiguration config = traceEnabledConfig(); config.set(AgentConfigOptions.EVENT_LISTENERS, List.of(TestEventListener.class.getName())); final AgentPlan agentPlan = TestAgent.getAgentPlanWithConfig(config); @@ -408,11 +552,167 @@ void testEventListenersFromAgentConfig() throws Exception { // process a some element to trigger the operator logic testHarness.processElement(new StreamRecord<>(1L)); + operator.waitInFlightEventsFinished(); // listener should have been invoked after element processing - called = ((TestEventListener) listener).called; + TestEventListener testEventListener = (TestEventListener) listener; + called = testEventListener.called; assertThat(called).isTrue(); + assertThat(testEventListener.eventTypes) + .noneMatch(ExecutionLifecycleEvents::isExecutionLifecycleEvent); + assertThat(RecordingEventLogger.events()) + .anyMatch( + record -> + ExecutionLifecycleEvents.isExecutionLifecycleEvent( + record.event.getType())); + } + } + + @Test + void testActionLifecycleEventsCarryExecutionContext() throws Exception { + final AgentConfiguration config = traceEnabledConfig(); + final AgentPlan agentPlan = TestAgent.getAgentPlanWithConfig(config); + + try (KeyedOneInputStreamOperatorTestHarness testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory(agentPlan, true), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) testHarness.getOperator(); + + testHarness.processElement(new StreamRecord<>(1L)); + operator.waitInFlightEventsFinished(); + } + + RecordedEvent action1Started = + findRecordedLifecycleEvent( + ExecutionLifecycleEvents.EXECUTION_STARTED_EVENT_TYPE, + "action1", + ExecutionLifecycleEvents.STATUS_STARTED); + RecordedEvent action1Finished = + findRecordedLifecycleEvent( + ExecutionLifecycleEvents.EXECUTION_FINISHED_EVENT_TYPE, + "action1", + ExecutionLifecycleEvents.STATUS_SUCCESS); + RecordedEvent action2Started = + findRecordedLifecycleEvent( + ExecutionLifecycleEvents.EXECUTION_STARTED_EVENT_TYPE, + "action2", + ExecutionLifecycleEvents.STATUS_STARTED); + RecordedEvent action2Finished = + findRecordedLifecycleEvent( + ExecutionLifecycleEvents.EXECUTION_FINISHED_EVENT_TYPE, + "action2", + ExecutionLifecycleEvents.STATUS_SUCCESS); + + assertThat(action1Started.traceContext().getInputRunId()).isNotBlank(); + assertThat(action1Started.traceContext().getBusinessKey()).isEqualTo("1"); + assertThat(action1Started.traceContext().getEntityType()).isEqualTo("action"); + assertThat(action1Started.status()).isEqualTo(ExecutionLifecycleEvents.STATUS_STARTED); + assertThat(action1Finished.traceContext().getExecutionId()) + .isEqualTo(action1Started.traceContext().getExecutionId()); + assertThat(action1Finished.status()).isEqualTo(ExecutionLifecycleEvents.STATUS_SUCCESS); + + assertThat(action2Started.traceContext().getInputRunId()) + .isEqualTo(action1Started.traceContext().getInputRunId()); + assertThat(action2Started.traceContext().getParentExecutionId()).isNull(); + assertThat(action2Finished.traceContext().getExecutionId()) + .isEqualTo(action2Started.traceContext().getExecutionId()); + + RecordedEvent middleEvent = + RecordingEventLogger.events().stream() + .filter( + record -> + TestAgent.MiddleEvent.EVENT_TYPE.equals( + record.event.getType())) + .findFirst() + .orElseThrow(); + assertThat(middleEvent.traceContext().getExecutionId()) + .isEqualTo(action1Started.traceContext().getExecutionId()); + assertThat(middleEvent.traceContext().getEntityName()).isEqualTo("action1"); + + RecordedEvent outputEvent = + RecordingEventLogger.events().stream() + .filter(record -> OutputEvent.EVENT_TYPE.equals(record.event.getType())) + .findFirst() + .orElseThrow(); + assertThat(outputEvent.traceContext().getExecutionId()) + .isEqualTo(action2Started.traceContext().getExecutionId()); + assertThat(outputEvent.traceContext().getEntityName()).isEqualTo("action2"); + } + + @Test + void testActionFailureLifecycleEventCarriesProblemCategory() throws Exception { + final AgentConfiguration config = traceEnabledConfig(); + AgentPlan basePlan = TestAgent.getDurableExceptionUncaughtAgentPlan(); + AgentPlan agentPlan = + new AgentPlan( + basePlan.getActions(), + basePlan.getActionsByEvent(), + basePlan.getResourceProviders(), + config); + + try (KeyedOneInputStreamOperatorTestHarness testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory(agentPlan, true), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) testHarness.getOperator(); + + testHarness.processElement(new StreamRecord<>(1L)); + assertThatThrownBy(operator::waitInFlightEventsFinished) + .hasCauseInstanceOf(ActionExecutionOperator.ActionTaskExecutionException.class); + } + + RecordedEvent failed = + findRecordedLifecycleEvent( + ExecutionLifecycleEvents.EXECUTION_FAILED_EVENT_TYPE, + "durableExceptionUncaughtAction", + ExecutionLifecycleEvents.STATUS_FAILED); + assertThat(failed.problemCategory()) + .isEqualTo(ExecutionReporter.ProblemCategories.ACTION_EXECUTION_FAILED); + assertThat(failed.event.getAttr("errorType")) + .isEqualTo(IllegalStateException.class.getName()); + assertThat(String.valueOf(failed.event.getAttr("errorMessage"))) + .contains("Simulated LLM failure"); + } + + @Test + void testActionContinuationEmitsStartedLifecycleEventOnce() throws Exception { + final AgentConfiguration config = traceEnabledConfig(); + AgentPlan basePlan = TestAgent.getAsyncAgentPlan(false); + AgentPlan agentPlan = + new AgentPlan( + basePlan.getActions(), + basePlan.getActionsByEvent(), + basePlan.getResourceProviders(), + config); + + try (KeyedOneInputStreamOperatorTestHarness testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory(agentPlan, true), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) testHarness.getOperator(); + + testHarness.processElement(new StreamRecord<>(5L)); + operator.waitInFlightEventsFinished(); } + + assertThat(RecordingEventLogger.events()) + .filteredOn( + record -> + ExecutionLifecycleEvents.EXECUTION_STARTED_EVENT_TYPE.equals( + record.event.getType()) + && "asyncAction1" + .equals(record.traceContext().getEntityName())) + .hasSize(1); } @Test @@ -471,6 +771,69 @@ void testDoesNotPruneSeqsInFlight() throws Exception { } } + @Test + void testBusinessAndExecutionEventsShareSingleEventLogger() throws Exception { + AgentPlan agentPlan = TestAgent.getAgentPlanWithConfig(traceEnabledConfig()); + try (KeyedOneInputStreamOperatorTestHarness testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory(agentPlan, true), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) testHarness.getOperator(); + + testHarness.processElement(new StreamRecord<>(1L)); + operator.waitInFlightEventsFinished(); + + assertThat(RecordingEventLogger.createdCount()).isEqualTo(1); + assertThat(RecordingEventLogger.openCount()).isEqualTo(1); + assertThat(RecordingEventLogger.flushCount()) + .isEqualTo(RecordingEventLogger.events().size()); + assertThat(RecordingEventLogger.events()) + .anySatisfy( + record -> + assertThat(record.event.getType()) + .isEqualTo(InputEvent.EVENT_TYPE)); + assertThat(RecordingEventLogger.events()) + .anySatisfy( + record -> + assertThat(record.event.getType()) + .isEqualTo( + ExecutionLifecycleEvents + .EXECUTION_STARTED_EVENT_TYPE)); + } + + assertThat(RecordingEventLogger.closeCount()).isEqualTo(1); + } + + @Test + void testTraceRecordingIsDisabledByDefault() throws Exception { + try (KeyedOneInputStreamOperatorTestHarness testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory(TestAgent.getAgentPlan(false), true), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) testHarness.getOperator(); + + testHarness.processElement(new StreamRecord<>(1L)); + operator.waitInFlightEventsFinished(); + + assertThat(RecordingEventLogger.events()).isNotEmpty(); + assertThat(RecordingEventLogger.events()) + .allSatisfy( + record -> { + assertThat(record.traceContext).isNull(); + assertThat( + ExecutionLifecycleEvents.isExecutionLifecycleEvent( + record.event.getType())) + .isFalse(); + }); + } + } + @Test void testEventLogBaseDirFromAgentConfig() throws Exception { String baseLogDir = "/tmp/flink-agents-test"; @@ -726,8 +1089,18 @@ void testEarlierCheckpointReplayKeepsDurableState() throws Exception { @Test void testActionStateStoreReplayIncurNoFunctionCall() throws Exception { - AgentPlan agentPlanWithStateStore = TestAgent.getAgentPlan(false); + AgentConfiguration config = traceEnabledConfig(); + AgentPlan basePlan = TestAgent.getAgentPlan(false); + AgentPlan agentPlanWithStateStore = + new AgentPlan( + basePlan.getActions(), + basePlan.getActionsByEvent(), + basePlan.getResourceProviders(), + config); InMemoryActionStateStore actionStateStore; + String originalInputRunId; + UUID originalMiddleEventId; + UUID originalOutputEventId; try (KeyedOneInputStreamOperatorTestHarness testHarness = new KeyedOneInputStreamOperatorTestHarness<>( new ActionExecutionOperatorFactory<>( @@ -744,10 +1117,37 @@ agentPlanWithStateStore, true, new InMemoryActionStateStore(false)), Long inputValue = 7L; - // First processing - this will execute the actual functions and store state + // First processing executes the Action functions and stores their results. testHarness.processElement(new StreamRecord<>(inputValue)); operator.waitInFlightEventsFinished(); - } + + RecordedEvent action1Started = + findRecordedLifecycleEvent( + ExecutionLifecycleEvents.EXECUTION_STARTED_EVENT_TYPE, + "action1", + ExecutionLifecycleEvents.STATUS_STARTED); + originalInputRunId = action1Started.traceContext().getInputRunId(); + originalMiddleEventId = + RecordingEventLogger.events().stream() + .filter( + record -> + TestAgent.MiddleEvent.EVENT_TYPE.equals( + record.event.getType())) + .findFirst() + .orElseThrow() + .event + .getId(); + originalOutputEventId = + RecordingEventLogger.events().stream() + .filter(record -> OutputEvent.EVENT_TYPE.equals(record.event.getType())) + .findFirst() + .orElseThrow() + .event + .getId(); + } + + RecordingEventLogger.reset(); + try (KeyedOneInputStreamOperatorTestHarness testHarness = new KeyedOneInputStreamOperatorTestHarness<>( new ActionExecutionOperatorFactory<>( @@ -760,7 +1160,7 @@ agentPlanWithStateStore, true, new InMemoryActionStateStore(false)), Long inputValue = 7L; - // First processing - this will execute the actual functions and store state + // Replay processing reuses the completed Action results from durable state. testHarness.processElement(new StreamRecord<>(inputValue)); operator.waitInFlightEventsFinished(); // Verify first output is correct @@ -772,6 +1172,60 @@ agentPlanWithStateStore, true, new InMemoryActionStateStore(false)), // The action state store should only have one entry assertThat(actionStateStore.getKeyedActionStates().get(String.valueOf(inputValue))) .hasSize(2); + + List replayEvents = RecordingEventLogger.events(); + assertThat(replayEvents) + .filteredOn( + record -> + ExecutionLifecycleEvents.EXECUTION_REUSED_EVENT_TYPE.equals( + record.event.getType())) + .extracting(record -> record.traceContext().getEntityName()) + .containsExactlyInAnyOrder("action1", "action2"); + assertThat(replayEvents) + .noneMatch( + record -> + ExecutionLifecycleEvents.EXECUTION_FINISHED_EVENT_TYPE.equals( + record.event.getType()) + && ExecutionLifecycleEvents.STATUS_SUCCESS.equals( + record.status()) + && List.of("action1", "action2") + .contains( + record.traceContext().getEntityName())); + + RecordedEvent action1Reused = + findRecordedLifecycleEvent( + ExecutionLifecycleEvents.EXECUTION_REUSED_EVENT_TYPE, + "action1", + ExecutionLifecycleEvents.STATUS_REUSED); + RecordedEvent action2Reused = + findRecordedLifecycleEvent( + ExecutionLifecycleEvents.EXECUTION_REUSED_EVENT_TYPE, + "action2", + ExecutionLifecycleEvents.STATUS_REUSED); + RecordedEvent replayedMiddleEvent = + replayEvents.stream() + .filter( + record -> + TestAgent.MiddleEvent.EVENT_TYPE.equals( + record.event.getType())) + .findFirst() + .orElseThrow(); + RecordedEvent replayedOutputEvent = + replayEvents.stream() + .filter(record -> OutputEvent.EVENT_TYPE.equals(record.event.getType())) + .findFirst() + .orElseThrow(); + + assertThat(action1Reused.traceContext().getInputRunId()) + .isNotEqualTo(originalInputRunId); + assertThat(action2Reused.traceContext().getInputRunId()) + .isEqualTo(action1Reused.traceContext().getInputRunId()); + assertThat(replayedMiddleEvent.event.getId()).isEqualTo(originalMiddleEventId); + assertThat(replayedMiddleEvent.traceContext().getExecutionId()) + .isEqualTo(action1Reused.traceContext().getExecutionId()); + assertThat(replayedOutputEvent.event.getId()).isEqualTo(originalOutputEventId); + assertThat(replayedOutputEvent.traceContext().getExecutionId()) + .isEqualTo(action2Reused.traceContext().getExecutionId()); } } diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java index 8fec7cb63..18997f423 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java @@ -18,6 +18,8 @@ package org.apache.flink.agents.runtime.operator; import org.apache.flink.agents.api.InputEvent; +import org.apache.flink.agents.api.trace.ExecutionReporter; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; import org.apache.flink.agents.plan.AgentPlan; import org.apache.flink.agents.plan.actions.Action; import org.apache.flink.agents.runtime.ResourceCache; @@ -28,10 +30,19 @@ import org.apache.flink.agents.runtime.context.RunnerContextImpl; import org.apache.flink.agents.runtime.memory.MemoryObjectImpl; import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl; +import org.apache.flink.agents.runtime.trace.ExecutionEventSink; +import org.apache.flink.api.common.serialization.SerializerConfigImpl; import org.apache.flink.api.common.state.MapState; +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.common.typeutils.TypeSerializer; +import org.apache.flink.core.memory.DataInputDeserializer; +import org.apache.flink.core.memory.DataOutputSerializer; import org.junit.jupiter.api.Test; +import java.util.ArrayList; import java.util.HashMap; +import java.util.List; +import java.util.Map; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -145,6 +156,7 @@ void transferContextsCopiesMemoryAndContinuationToNewTask() throws Exception { invokeCreateAndSetRunnerContext(mgr, from); RunnerContextImpl.MemoryContext fromMemCtx = from.getRunnerContext().getMemoryContext(); assertThat(fromMemCtx).isNotNull(); + from.markExecutionStartedEventEmitted(); // transferContexts (ActionTaskContextManager.java:266-286) copies but does NOT // remove from source. The from-side continuation map is never populated (the @@ -167,6 +179,84 @@ void transferContextsCopiesMemoryAndContinuationToNewTask() throws Exception { // — the source carries its continuation on its runner context, not on the // manager's map. assertThat(mgr.hasContinuationContext(from)).isFalse(); + + // (d) Persisted Action lifecycle state follows the continuation task. + assertThat(to.hasExecutionStartedEventEmitted()).isTrue(); + } + } + + @Test + void reportedExecutionStateFollowsActionExecutionAcrossContinuationTasks() throws Exception { + try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) { + Action action = TestActions.noopAction(); + ActionTask from = new JavaActionTask("k", new InputEvent(1L), action); + ActionTask to = + new JavaActionTask("k", new InputEvent(1L), action, from.getTraceContext()); + List reports = new ArrayList<>(); + ExecutionEventSink sink = (event, context) -> reports.add(context); + + invokeCreateAndSetRunnerContext(mgr, from, sink); + from.getRunnerContext() + .reportExecutionStarted( + ExecutionReporter.EntityTypes.TOOL, "slow-tool", Map.of()); + + mgr.transferContexts(from, to, new DurableExecutionManager(null)); + invokeCreateAndSetRunnerContext(mgr, to, sink); + to.getRunnerContext() + .reportExecutionSucceeded( + ExecutionReporter.EntityTypes.TOOL, "slow-tool", Map.of()); + + assertThat(reports).hasSize(2); + assertThat(reports.get(1).getExecutionId()).isEqualTo(reports.get(0).getExecutionId()); + } + } + + @Test + void completingActionExecutionDropsReportedExecutionState() throws Exception { + try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) { + ActionTask task = new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction()); + List reports = new ArrayList<>(); + ExecutionEventSink sink = (event, context) -> reports.add(context); + + invokeCreateAndSetRunnerContext(mgr, task, sink); + task.getRunnerContext() + .reportExecutionStarted(ExecutionReporter.EntityTypes.LLM, "model-a", Map.of()); + + mgr.completeActionExecution(task); + invokeCreateAndSetRunnerContext(mgr, task, sink); + task.getRunnerContext() + .reportExecutionSucceeded( + ExecutionReporter.EntityTypes.LLM, "model-a", Map.of()); + + assertThat(reports).hasSize(2); + assertThat(reports.get(1).getExecutionId()) + .isNotEqualTo(reports.get(0).getExecutionId()); + } + } + + @Test + void activeExecutionReportsDoNotEnterActionTaskState() throws Exception { + try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) { + ActionTask task = new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction()); + invokeCreateAndSetRunnerContext(mgr, task, (event, context) -> {}); + task.getRunnerContext() + .reportExecutionStarted( + ExecutionReporter.EntityTypes.TOOL, + "search", + Map.of("toolCallId", "call-1")); + task.markExecutionStartedEventEmitted(); + + TypeSerializer serializer = + TypeInformation.of(ActionTask.class) + .createSerializer(new SerializerConfigImpl()); + DataOutputSerializer output = new DataOutputSerializer(512); + serializer.serialize(task, output); + ActionTask restored = + serializer.deserialize(new DataInputDeserializer(output.getCopyOfBuffer())); + + assertThat(restored.getTraceContext()).isEqualTo(task.getTraceContext()); + assertThat(restored.hasExecutionStartedEventEmitted()).isTrue(); + assertThat(restored.getRunnerContext()).isNull(); } } @@ -226,9 +316,14 @@ void closeIsIdempotent() throws Exception { * Shared helper: install a runner context on {@code task} using mocked collaborators. Used by * tests that need a fully wired runner context but do not care about the collaborator details. */ - @SuppressWarnings("unchecked") private static void invokeCreateAndSetRunnerContext( ActionTaskContextManager mgr, ActionTask task) { + invokeCreateAndSetRunnerContext(mgr, task, null); + } + + @SuppressWarnings("unchecked") + private static void invokeCreateAndSetRunnerContext( + ActionTaskContextManager mgr, ActionTask task, ExecutionEventSink executionEventSink) { AgentPlan plan = newEmptyAgentPlan(); ResourceCache cache = mock(ResourceCache.class); FlinkAgentsMetricGroupImpl metricGroup = @@ -246,7 +341,8 @@ private static void invokeCreateAndSetRunnerContext( sensoryMem, shortTermMem, /* pythonRunnerContext */ null, - /* longTermMemory */ null); + /* longTermMemory */ null, + executionEventSink); } private static AgentPlan newEmptyAgentPlan() { diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/EventRouterTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/EventRouterTest.java index 97f994b42..ccb924421 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/EventRouterTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/EventRouterTest.java @@ -31,6 +31,7 @@ import org.apache.flink.agents.runtime.python.utils.PythonActionExecutor; import org.apache.flink.streaming.api.watermark.Watermark; import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; import org.mockito.InOrder; import java.util.ArrayList; @@ -39,9 +40,12 @@ import java.util.Map; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatCode; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.isNull; import static org.mockito.Mockito.RETURNS_DEEP_STUBS; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.inOrder; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; @@ -98,8 +102,15 @@ void notifyEventProcessedNotifiesAllListeners() throws Exception { InputEvent inputEvent = new InputEvent(7L); router.notifyEventProcessed(inputEvent); - verify(listener1).onEventProcessed(any(EventContext.class), eq(inputEvent)); - verify(listener2).onEventProcessed(any(EventContext.class), eq(inputEvent)); + ArgumentCaptor logContext = ArgumentCaptor.forClass(EventContext.class); + ArgumentCaptor listener1Context = ArgumentCaptor.forClass(EventContext.class); + ArgumentCaptor listener2Context = ArgumentCaptor.forClass(EventContext.class); + verify(mockLogger).append(logContext.capture(), eq(inputEvent), isNull()); + verify(listener1).onEventProcessed(listener1Context.capture(), eq(inputEvent)); + verify(listener2).onEventProcessed(listener2Context.capture(), eq(inputEvent)); + assertThat(logContext.getValue().getEventType()).isEqualTo(InputEvent.EVENT_TYPE); + assertThat(listener1Context.getValue()).isSameAs(logContext.getValue()); + assertThat(listener2Context.getValue()).isSameAs(logContext.getValue()); } /** @@ -138,10 +149,31 @@ void notifyEventProcessedAppendsAndFlushesLogger() throws Exception { router.notifyEventProcessed(inputEvent); InOrder ordered = inOrder(mockLogger); - ordered.verify(mockLogger).append(any(EventContext.class), eq(inputEvent)); + ordered.verify(mockLogger).append(any(EventContext.class), eq(inputEvent), isNull()); ordered.verify(mockLogger).flush(); } + @Test + void notifyEventProcessedIgnoresEventLogWriteFailure() throws Exception { + AgentPlan plan = new AgentPlan(new HashMap<>(), new HashMap<>()); + EventLogger mockLogger = mock(EventLogger.class); + EventRouter router = + new EventRouter<>(plan, /* inputIsJava */ true, mockLogger); + EventListener listener = mock(EventListener.class); + router.addEventListener(listener); + BuiltInMetrics spyMetrics = spy(makeMetrics()); + router.open(spyMetrics); + + InputEvent inputEvent = new InputEvent(3L); + doThrow(new RuntimeException("log failed")) + .when(mockLogger) + .append(any(EventContext.class), eq(inputEvent), isNull()); + + assertThatCode(() -> router.notifyEventProcessed(inputEvent)).doesNotThrowAnyException(); + verify(listener).onEventProcessed(any(EventContext.class), eq(inputEvent)); + verify(spyMetrics).markEventProcessed(); + } + /** * Verifies watermarks are drained in arrival order — even when the keys ahead of them finish * out-of-order. The {@link SegmentedQueue} closes a segment when a watermark is added, so the diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ExecutionTraceContextTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ExecutionTraceContextTest.java new file mode 100644 index 000000000..d5f50a2f9 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ExecutionTraceContextTest.java @@ -0,0 +1,114 @@ +/* + * 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.trace.ExecutionLifecycleEvents; +import org.apache.flink.agents.api.trace.ExecutionReporter; +import org.apache.flink.agents.api.trace.ExecutionTraceContext; +import org.apache.flink.api.common.serialization.SerializerConfigImpl; +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.common.typeutils.TypeSerializer; +import org.apache.flink.core.memory.DataInputDeserializer; +import org.apache.flink.core.memory.DataOutputSerializer; +import org.junit.jupiter.api.Test; + +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Unit tests for {@link ExecutionTraceContext}. */ +class ExecutionTraceContextTest { + + @Test + void inputRunContextCarriesInputScopeWithoutExecutionEntity() { + ExecutionTraceContext traceContext = + ExecutionTraceContext.forInputRun("business-key-1", "review-agent"); + + assertThat(traceContext.getInputRunId()).isNotBlank(); + assertThat(traceContext.getBusinessKey()).isEqualTo("business-key-1"); + assertThat(traceContext.getAgentName()).isEqualTo("review-agent"); + assertThat(traceContext.getExecutionId()).isNull(); + assertThat(traceContext.getParentExecutionId()).isNull(); + assertThat(traceContext.getEntityType()).isNull(); + assertThat(traceContext.getEntityName()).isNull(); + } + + @Test + void childExecutionKeepsInputScopeAndUsesCurrentExecutionAsParent() { + ExecutionTraceContext inputRunContext = ExecutionTraceContext.forInputRun("key"); + ExecutionTraceContext actionContext = inputRunContext.childExecution("action", "classify"); + ExecutionTraceContext toolContext = actionContext.childExecution("tool", "search"); + + assertThat(actionContext.getInputRunId()).isEqualTo(inputRunContext.getInputRunId()); + assertThat(actionContext.getBusinessKey()).isEqualTo("key"); + assertThat(actionContext.getExecutionId()).isNotBlank(); + assertThat(actionContext.getParentExecutionId()).isNull(); + assertThat(actionContext.getEntityType()).isEqualTo("action"); + assertThat(actionContext.getEntityName()).isEqualTo("classify"); + + assertThat(toolContext.getInputRunId()).isEqualTo(inputRunContext.getInputRunId()); + assertThat(toolContext.getExecutionId()).isNotBlank(); + assertThat(toolContext.getExecutionId()).isNotEqualTo(actionContext.getExecutionId()); + assertThat(toolContext.getParentExecutionId()).isEqualTo(actionContext.getExecutionId()); + assertThat(toolContext.getEntityType()).isEqualTo("tool"); + assertThat(toolContext.getEntityName()).isEqualTo("search"); + } + + @Test + void lifecycleEventCanCarryExecutionStatusAndProblemCategory() { + Event event = + ExecutionLifecycleEvents.executionFailed( + new RuntimeException("boom"), + ExecutionReporter.ProblemCategories.ACTION_EXECUTION_FAILED); + + assertThat(event.getAttr(ExecutionLifecycleEvents.STATUS_ATTRIBUTE)) + .isEqualTo(ExecutionLifecycleEvents.STATUS_FAILED); + assertThat(event.getAttr(ExecutionLifecycleEvents.PROBLEM_CATEGORY_ATTRIBUTE)) + .isEqualTo(ExecutionReporter.ProblemCategories.ACTION_EXECUTION_FAILED); + } + + @Test + void entityMetadataSurvivesFlinkStateSerialization() throws Exception { + ExecutionTraceContext traceContext = + ExecutionTraceContext.fromExistingIds( + "run", + "business-key", + "agent", + "execution", + "parent", + ExecutionReporter.EntityTypes.TOOL, + "search", + Map.of("toolCallId", "call-1")); + TypeSerializer serializer = + TypeInformation.of(ExecutionTraceContext.class) + .createSerializer(new SerializerConfigImpl()); + DataOutputSerializer output = new DataOutputSerializer(256); + + serializer.serialize(traceContext, output); + ExecutionTraceContext restored = + serializer.deserialize(new DataInputDeserializer(output.getCopyOfBuffer())); + + assertThat(restored).isEqualTo(traceContext); + assertThat(restored.getEntityMetadata()).containsEntry("toolCallId", "call-1"); + assertThatThrownBy(() -> restored.getEntityMetadata().put("new-key", "new-value")) + .isInstanceOf(UnsupportedOperationException.class); + } +} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/skill/LoadSkillToolTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/skill/LoadSkillToolTest.java index f048920df..cbe4ce3b4 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/skill/LoadSkillToolTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/skill/LoadSkillToolTest.java @@ -24,6 +24,7 @@ import org.apache.flink.agents.api.skills.Skills; import org.apache.flink.agents.api.tools.ToolParameters; import org.apache.flink.agents.api.tools.ToolResponse; +import org.apache.flink.agents.api.trace.ToolExecutionMetadataKeys; import org.apache.flink.agents.runtime.resource.ResourceContextImpl; import org.junit.jupiter.api.Test; @@ -96,6 +97,15 @@ void resourcePathReturnsRawContent() { assertTrue(out.length() > 0); } + @Test + void executionMetadataDescribesRequestedSkillResource() { + Map metadata = + tool(contextWithSkills()).getToolExecutionMetadata(args("github", "README.md")); + + assertEquals("github", metadata.get(ToolExecutionMetadataKeys.SKILL_NAME)); + assertEquals("README.md", metadata.get(ToolExecutionMetadataKeys.SKILL_RESOURCE_PATH)); + } + @Test void missingResourceReportsAvailable() { LoadSkillTool t = tool(contextWithSkills());