From 852676f8c049f2e399c5cde0cf3ea4fc67a7a3fb Mon Sep 17 00:00:00 2001 From: Yu Tang Date: Thu, 20 Aug 2026 13:57:13 +0800 Subject: [PATCH 1/3] [api][runtime][python] Add Agent Trace recording to Event Log (#924) Generated-by: Cursor (Grok 4.6) --- .../org/apache/flink/agents/api/Event.java | 3 +- .../apache/flink/agents/api/EventContext.java | 6 +- .../api/chat/model/BaseChatModelSetup.java | 2 +- .../api/configuration/AgentConfigOptions.java | 8 + .../agents/api/listener/EventListener.java | 5 +- .../flink/agents/api/logger/EventLogger.java | 22 +- .../tools/ToolExecutionMetadataProvider.java | 38 + .../api/trace/ExecutionLifecycleEvents.java | 108 +++ .../agents/api/trace/ExecutionReporter.java | 88 +++ .../agents/api/trace/ExecutionReporters.java | 117 +++ .../api/trace/ExecutionTraceContext.java | 213 ++++++ .../api/trace/LLMExecutionMetadataKeys.java | 26 + .../api/trace/ToolExecutionMetadataKeys.java | 32 + .../apache/flink/agents/api/EventTest.java | 17 + .../trace/ExecutionLifecycleEventsTest.java | 56 ++ .../api/trace/ExecutionReportersTest.java | 132 ++++ docs/content/docs/operations/configuration.md | 1 + docs/content/docs/operations/monitoring.md | 64 +- .../java/agent_plan_with_python_action.json | 1 + .../python/agent_plan_with_java_action.json | 1 + .../agents/integrations/mcp/MCPTool.java | 43 +- .../agents/integrations/mcp/MCPToolTest.java | 21 +- .../apache/flink/agents/plan/AgentPlan.java | 43 +- .../flink/agents/plan/JavaFunction.java | 14 +- .../agents/plan/actions/ChatModelAction.java | 82 ++- .../agents/plan/actions/ToolCallAction.java | 200 ++++- .../plan/resource/python/PythonMCPServer.java | 9 +- .../plan/resource/python/PythonMCPTool.java | 31 +- .../serializer/AgentPlanJsonDeserializer.java | 5 +- .../serializer/AgentPlanJsonSerializer.java | 2 + .../plan/AgentPlanDeclareMCPServerTest.java | 1 + .../agents/plan/FunctionToolPlanTest.java | 17 + .../flink/agents/plan/TestFunction.java | 14 + .../actions/ChatModelActionRetryTest.java | 213 ++++++ .../actions/ToolCallActionReportTest.java | 299 ++++++++ .../plan/actions/ToolCallActionTest.java | 16 + .../AgentPlanJsonSerializerTest.java | 21 + python/flink_agents/api/core_options.py | 6 + python/flink_agents/api/events/event.py | 10 +- .../api/tests/test_core_options.py | 1 + python/flink_agents/api/tests/test_event.py | 39 + .../api/tests/test_execution_reporter.py | 67 ++ python/flink_agents/api/tools/__init__.py | 9 +- .../tools/tool_execution_metadata_provider.py | 31 + python/flink_agents/api/trace/__init__.py | 42 ++ .../api/trace/execution_lifecycle_events.py | 46 ++ .../api/trace/execution_reporter.py | 145 ++++ .../api/trace/llm_execution_metadata_keys.py | 24 + .../api/trace/tool_execution_metadata_keys.py | 30 + .../flink_agents/cli/tests/test_trace_tree.py | 155 +++- python/flink_agents/cli/trace_tree.py | 136 +++- .../memory_event_logging_test.py | 13 +- .../python_event_logging_test.py | 25 +- python/flink_agents/e2e_tests/test_utils.py | 8 +- .../e2e_tests/test_utils_unit_test.py | 42 +- python/flink_agents/integrations/mcp/mcp.py | 17 +- .../integrations/mcp/tests/test_mcp.py | 11 +- .../plan/actions/chat_model_action.py | 73 +- .../plan/actions/tool_call_action.py | 126 +++- python/flink_agents/plan/agent_plan.py | 15 +- .../actions/test_chat_model_action_retry.py | 115 ++- .../tests/actions/test_tool_call_action.py | 147 +++- .../plan/tests/resources/agent_plan.json | 1 + .../plan/tests/test_agent_plan.py | 28 +- .../runtime/flink_runner_context.py | 71 +- .../runtime/remote_execution_environment.py | 8 +- .../flink_agents/runtime/skill/skill_tools.py | 31 +- .../runtime/skill/tests/test_load_skill.py | 23 + .../tests/test_flink_runner_context_trace.py | 83 +++ .../runtime/PythonMCPResourceDiscovery.java | 2 +- .../runtime/context/RunnerContextImpl.java | 126 +++- .../env/RemoteExecutionEnvironment.java | 8 +- .../runtime/eventlog/EventLogRecord.java | 44 +- .../EventLogRecordJsonDeserializer.java | 131 +++- .../EventLogRecordJsonSerializer.java | 139 ++-- .../runtime/eventlog/EventLogWriter.java | 179 +++++ .../runtime/eventlog/FileEventLogger.java | 51 +- .../runtime/eventlog/JsonTruncator.java | 63 +- .../runtime/eventlog/Slf4jEventLogger.java | 48 +- .../runtime/metrics/BuiltInMetrics.java | 8 + .../operator/ActionExecutionOperator.java | 176 +++-- .../agents/runtime/operator/ActionTask.java | 71 +- .../operator/ActionTaskContextManager.java | 40 +- .../agents/runtime/operator/EventRouter.java | 101 +-- .../runtime/operator/JavaActionTask.java | 7 + .../context/PythonRunnerContextImpl.java | 43 ++ .../python/operator/PythonActionTask.java | 20 +- .../operator/PythonGeneratorActionTask.java | 17 +- .../agents/runtime/skill/LoadSkillTool.java | 32 +- .../runtime/trace/ExecutionEventLogger.java | 43 ++ .../runtime/trace/ExecutionEventSink.java | 29 + .../runtime/trace/ReportedExecutionKey.java | 69 ++ ...unnerContextImplExecutionReporterTest.java | 185 +++++ .../eventlog/EventLogRecordJsonSerdeTest.java | 331 ++++++--- .../runtime/eventlog/EventLogWriterTest.java | 149 ++++ .../runtime/eventlog/FileEventLoggerTest.java | 104 +-- .../runtime/eventlog/JsonTruncatorTest.java | 37 +- .../eventlog/Slf4jEventLoggerTest.java | 65 +- .../operator/ActionExecutionOperatorTest.java | 683 +++++++++++++++++- .../ActionTaskContextManagerTest.java | 111 ++- .../EventLineageReconstructionTest.java | 17 +- .../runtime/operator/EventRouterTest.java | 38 +- .../operator/ExecutionTraceContextTest.java | 114 +++ .../runtime/skill/LoadSkillToolTest.java | 45 ++ 104 files changed, 6109 insertions(+), 766 deletions(-) create mode 100644 api/src/main/java/org/apache/flink/agents/api/tools/ToolExecutionMetadataProvider.java create mode 100644 api/src/main/java/org/apache/flink/agents/api/trace/ExecutionLifecycleEvents.java create mode 100644 api/src/main/java/org/apache/flink/agents/api/trace/ExecutionReporter.java create mode 100644 api/src/main/java/org/apache/flink/agents/api/trace/ExecutionReporters.java create mode 100644 api/src/main/java/org/apache/flink/agents/api/trace/ExecutionTraceContext.java create mode 100644 api/src/main/java/org/apache/flink/agents/api/trace/LLMExecutionMetadataKeys.java create mode 100644 api/src/main/java/org/apache/flink/agents/api/trace/ToolExecutionMetadataKeys.java create mode 100644 api/src/test/java/org/apache/flink/agents/api/trace/ExecutionLifecycleEventsTest.java create mode 100644 api/src/test/java/org/apache/flink/agents/api/trace/ExecutionReportersTest.java create mode 100644 plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionReportTest.java create mode 100644 python/flink_agents/api/tests/test_execution_reporter.py create mode 100644 python/flink_agents/api/tools/tool_execution_metadata_provider.py create mode 100644 python/flink_agents/api/trace/__init__.py create mode 100644 python/flink_agents/api/trace/execution_lifecycle_events.py create mode 100644 python/flink_agents/api/trace/execution_reporter.py create mode 100644 python/flink_agents/api/trace/llm_execution_metadata_keys.py create mode 100644 python/flink_agents/api/trace/tool_execution_metadata_keys.py create mode 100644 python/flink_agents/runtime/tests/test_flink_runner_context_trace.py create mode 100644 runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/EventLogWriter.java create mode 100644 runtime/src/main/java/org/apache/flink/agents/runtime/trace/ExecutionEventLogger.java create mode 100644 runtime/src/main/java/org/apache/flink/agents/runtime/trace/ExecutionEventSink.java create mode 100644 runtime/src/main/java/org/apache/flink/agents/runtime/trace/ReportedExecutionKey.java create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextImplExecutionReporterTest.java create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/EventLogWriterTest.java create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/operator/ExecutionTraceContextTest.java 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 2383428b8..43286ac88 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 @@ -83,7 +83,8 @@ public Event( if (type == null || type.isEmpty()) { throw new IllegalArgumentException("Event 'type' must not be null or empty."); } - this.id = id; + // Explicit null matches an omitted id: both mint a per-occurrence UUID. + this.id = id != null ? id : UUID.randomUUID(); this.type = type; this.attributes = attributes != null ? attributes : new HashMap<>(); this.upstreamEventId = upstreamEventId; 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/chat/model/BaseChatModelSetup.java b/api/src/main/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetup.java index 0c61bf6d4..3cb2e655b 100644 --- a/api/src/main/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetup.java +++ b/api/src/main/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetup.java @@ -195,7 +195,7 @@ public String getConnectionName() { return this.connectionName; } - @VisibleForTesting + /** Returns the configured model or deployment identifier used by this setup. */ public String getModel() { return model; } 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 18ecc170f..7fe4a93d8 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 @@ -134,6 +134,14 @@ public enum ConditionEvaluationFailureStrategy { 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/LLMExecutionMetadataKeys.java b/api/src/main/java/org/apache/flink/agents/api/trace/LLMExecutionMetadataKeys.java new file mode 100644 index 000000000..48433a619 --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/trace/LLMExecutionMetadataKeys.java @@ -0,0 +1,26 @@ +/* + * 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 LLM execution entity metadata. */ +public final class LLMExecutionMetadataKeys { + + public static final String MODEL = "model"; + + private LLMExecutionMetadataKeys() {} +} 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 8327884ac..a87357f29 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 @@ -197,10 +197,27 @@ 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")); } + @Test + void testFromJsonExplicitNullIdGeneratesUuid() throws IOException { + Event event = Event.fromJson("{\"id\":null,\"type\":\"x\"}"); + + assertNotNull(event.getId()); + assertEquals("x", event.getType()); + } + + @Test + void testConstructorExplicitNullIdGeneratesUuid() { + Event event = new Event(null, "x", Map.of()); + + assertNotNull(event.getId()); + assertEquals("x", event.getType()); + } + @Test void testFromJsonMissingType() { String json = "{\"attributes\":{\"msg\":\"hello\"}}"; 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/operations/configuration.md b/docs/content/docs/operations/configuration.md index c1918bb5a..e363829da 100644 --- a/docs/content/docs/operations/configuration.md +++ b/docs/content/docs/operations/configuration.md @@ -137,6 +137,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 7639d0cac..652c9bfe9 100644 --- a/docs/content/docs/operations/monitoring.md +++ b/docs/content/docs/operations/monitoring.md @@ -39,6 +39,7 @@ We offer data monitoring for built-in metrics, which includes events, actions, a | **Action** | action.\.numOfActionsExecuted | The total number of actions this operator has executed for a specific action name. | Count | | **Action** | action.\.numOfActionsExecutedPerSec | The number of actions this operator has executed per second for a specific action name. | Meter | | **Agent** | eventLogTruncatedEvents | Number of event log records whose payload was truncated at `STANDARD` level. Increments once per event, regardless of how many fields inside it were truncated. Use this to decide whether to raise truncation thresholds or move specific event types to `VERBOSE`. | Count | +| **Agent** | eventLogWriteFailures | Number of Event Log write attempts for which `append`, `flush`, or both failed. Event Log writes are best-effort and do not fail the job. | Count | #### Token Usage Metrics @@ -149,6 +150,8 @@ The system supports two types of event loggers: **SLF4J Event Log** (default) an By default, the SLF4J Event Log is used. If `baseLogDir` is configured, the system automatically switches to the File Event Log. +Event Log is an observability output. An `append` or `flush` failure does not fail Event processing or the Flink job. The first failure is logged at `WARN`, subsequent failures are logged at `DEBUG`, and every failed write attempt increments `eventLogWriteFailures`. + ### SLF4J Event Log (Default) The **SLF4J Event Log** outputs events through a dedicated SLF4J logger (`org.apache.flink.agents.EventLog`). On startup, the logger **automatically configures** log4j2 to write events to a separate file (`{log.file}.event-log.log`) in Flink's log directory, making them visible in Flink's Web UI **Logs** tab. **No manual log4j2 configuration is required.** @@ -178,29 +181,42 @@ 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. + +After fine-grained recovery, a cached durable LLM or Tool result is currently recorded as a new successful execution because cache reuse is not exposed to execution reporting. Distinguishing reused child executions is follow-up work. + +Example Trace record: ```json { "timestamp": "2024-01-15T10:30:00Z", "logLevel": "STANDARD", - "eventType": "MiddleEvent", - "event": { - "eventType": "MiddleEvent", - "id": "39361629-4f1d-4b62-b734-a48b181cb6e0", - "attributes": {}, - "type": "MiddleEvent", - "upstreamEventId": "dad5ed00-80e2-4746-8bdb-f126bad504b5", - "upstreamActionName": "action1" - } + "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": "_default_chat_model", + "entityMetadata": {"model": "qwen-max"}, + "eventId": "80f30736-2759-41d8-aa59-9c8ad481ab42", + "eventType": "_execution_finished_event", + "status": "success", + "eventAttributes": {} } ``` -`upstreamEventId` identifies the Event consumed by the Action that emitted the current Event, and `upstreamActionName` identifies that Action. The framework maintains both fields directly on the `event` object, outside business `attributes`. A root `InputEvent` omits both fields. +For an LLM execution, `entityName` identifies the configured ChatModel Resource, while +`entityMetadata.model` identifies the model or deployment configured on that Resource. This value +is the requested identifier and is not a provider-confirmed model identity. + +`upstreamEventId` identifies the Event consumed by the Action that emitted the current Event, and `upstreamActionName` identifies that Action. Both are top-level Event Log fields derived from framework-managed Event lineage and remain outside `eventAttributes`. A root `InputEvent` omits both fields. ### Trace Tree Reconstruction -The `flink-agents-trace-tree` command is installed with the Flink Agents Python wheel. It rebuilds InputEvent-rooted Trace Trees from a saved File Event Log. Pass either one log file for text output or a log directory for Trace Tree JSON: +The `flink-agents-trace-tree` command is installed with the Flink Agents Python wheel. It rebuilds InputEvent-rooted Trace Trees from business Events in a saved File Event Log and ignores execution lifecycle Events. The four lifecycle types `_execution_started_event`, `_execution_finished_event`, `_execution_failed_event`, and `_execution_reused_event` are reserved for the framework. A record is ignored only when its type, status, and execution identity match the corresponding framework lifecycle shape. A business Event that uses a reserved type without that shape is retained and reported with a reconstruction warning. The reader accepts both the current flat record shape and the previous nested `event` shape, including files that contain both formats. Pass either one log file for text output or a log directory for Trace Tree JSON: ```bash flink-agents-trace-tree /path/to/events-job-task-0.log @@ -320,6 +336,7 @@ consistently in both directions. | `MISSING_PARENT` | An Event references an upstream Event with no valid node in the Event Log. | | `CYCLE_DETECTED` | A lineage edge closes a cycle and is omitted from the reconstructed DAG. | | `MALFORMED_RECORD` | An Event Log record is invalid or partially written. | +| `RESERVED_EVENT_TYPE` | A business Event uses a framework-reserved lifecycle Event type. | | `UNREADABLE_FILE` | An Event Log file could not be read. | ### Event Log Levels @@ -334,7 +351,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 `eventType`, `type`, `id`, `upstreamEventId`, and `upstreamActionName`. The `attributes` envelope is also preserved, while large nested content inside it can still be truncated. +**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: @@ -352,10 +369,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": [ @@ -370,9 +386,9 @@ 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 | +| Event | `` value | |--------------------------|----------------------------------| | `InputEvent` | `_input_event` | | `OutputEvent` | `_output_event` | @@ -382,9 +398,15 @@ You can override the level for individual event types using the `event-log.type. | `ToolResponseEvent` | `_tool_response_event` | | `ContextRetrievalRequestEvent` | `_context_retrieval_request_event` | | `ContextRetrievalResponseEvent` | `_context_retrieval_response_event` | +| Execution lifecycle: started | `_execution_started_event` | +| Execution lifecycle: finished | `_execution_finished_event` | +| Execution lifecycle: failed | `_execution_failed_event` | +| Execution lifecycle: reused | `_execution_reused_event` | Each event type has its own independently overridable key, so a job-level override does not clobber other entries from `config.yaml`. +`event-log.trace.enabled` controls whether execution lifecycle Events are produced. When Trace recording is enabled, these Events still use the per-event-type level resolution above. Setting any `_execution_*` type to `OFF` suppresses that lifecycle Event and may make the recorded Trace incomplete. + Resolution is hierarchical — the resolver walks up dot-separated segments of the event type, mirroring Log4j's logger hierarchy. For a user-defined event type `com.example.myapp.OrderEvent`, the lookup order is: 1. `event-log.type.com.example.myapp.OrderEvent.level` (exact match) @@ -421,4 +443,8 @@ 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 record shape.** The nested `event` object is removed: `event.id` becomes the top-level `eventId`, and `event.attributes` becomes the top-level `eventAttributes`. `eventType` was already available at the top level and remains there. 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. +- **Python Event IDs now identify occurrences.** Python previously generated a deterministic UUID from Event content and regenerated it when the content changed. It now assigns a UUID4 to each Event occurrence, matching the identity semantics used by Agent Trace. Consumers must not rely on equal Event payloads producing equal IDs for deduplication. +- **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. +- **Event Log write failures are now best-effort.** Earlier versions propagated `append` or `flush` failures into Event processing. Write failures now leave the job running and are reported through the `eventLogWriteFailures` metric and a first-failure `WARN` log. 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 3b9a76e32..f49c807aa 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 97974fac7..6ae4d1f0a 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 98d967c34..df6801132 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 @@ -93,6 +93,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) { @@ -113,9 +116,18 @@ public AgentPlan( Map actions, Map> resourceProviders, AgentConfiguration config) { + this(actions, resourceProviders, config, null); + } + + public AgentPlan( + Map actions, + Map> resourceProviders, + AgentConfiguration config, + String agentName) { this.actions = Collections.unmodifiableMap(new LinkedHashMap<>(actions)); this.resourceProviders = resourceProviders; this.config = config; + this.agentName = agentName; } /** @@ -130,12 +142,17 @@ 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.actions = new LinkedHashMap<>(); this.resourceProviders = new HashMap<>(); this.config = config; extractActionsFromAgent(agent); extractResourceProvidersFromAgent(agent); this.actions = Collections.unmodifiableMap(new LinkedHashMap<>(actions)); + this.agentName = agentName != null ? agentName : defaultAgentName(agent); } public Map getActions() { @@ -154,6 +171,10 @@ public Map> getResourceProviders() { return resourceProviders; } + public String getAgentName() { + return agentName; + } + public AgentConfiguration getConfig() { return config; } @@ -172,9 +193,15 @@ private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundE AgentPlan agentPlan = new ObjectMapper().readValue(serializedStr, AgentPlan.class); this.actions = Collections.unmodifiableMap(new LinkedHashMap<>(agentPlan.getActions())); 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[] triggerConditions, @@ -353,7 +380,8 @@ private void extractJavaMCPServer(Method method) throws Exception { Iterable tools = (Iterable) listToolsMethod.invoke(mcpServer); - for (SerializableResource tool : tools) { + for (SerializableResource discoveredTool : tools) { + SerializableResource tool = attachMcpServerNameIfSupported(discoveredTool, name); Method getNameMethod = tool.getClass().getMethod("getName"); String toolName = (String) getNameMethod.invoke(tool); addResourceProvider( @@ -379,6 +407,19 @@ 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 associatedTool = method.invoke(tool, mcpServerName); + return associatedTool instanceof SerializableResource + ? (SerializableResource) associatedTool + : tool; + } catch (NoSuchMethodException ignored) { + 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 e8c663f34..13f89d876 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,9 @@ 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.api.trace.LLMExecutionMetadataKeys; import org.apache.flink.agents.plan.JavaFunction; import org.apache.flink.api.java.typeutils.RowTypeInfo; import org.apache.flink.types.Row; @@ -373,19 +376,36 @@ public ChatMessage call() throws Exception { return chatModel.chat(messages, promptArgs, Map.of()); } }; + Map llmMetadata = + chatModel.getModel() == null + ? Map.of() + : Map.of(LLMExecutionMetadataKeys.MODEL, chatModel.getModel()); for (int attempt = 0; attempt < numRetries + 1; attempt++) { try { - response = - chatAsync - ? ctx.durableExecuteAsync(callable) - : ctx.durableExecute(callable); + ExecutionReporters.started( + ctx, ExecutionReporter.EntityTypes.LLM, model, llmMetadata); + try { + response = + chatAsync + ? ctx.durableExecuteAsync(callable) + : ctx.durableExecute(callable); + Objects.requireNonNull(response, "ChatModel returned a null response."); + } catch (Throwable modelError) { + throw reportFailedAndPropagate( + ctx, + ExecutionReporter.EntityTypes.LLM, + model, + llmMetadata, + modelError, + ExecutionReporter.ProblemCategories.MODEL_CALL_FAILED); + } + ExecutionReporters.succeeded( + ctx, ExecutionReporter.EntityTypes.LLM, model, llmMetadata); recordChatTokenMetrics(chatModel, response, requestMetricGroup); - // 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( @@ -408,6 +428,7 @@ public ChatMessage call() throws Exception { Thread.sleep(currentWaitSec * 1000L); totalWaitTimeSec += currentWaitSec; } + continue; } else { LOG.debug( "Chat request {} failed, the input chat messages are {}.", @@ -416,6 +437,7 @@ public ChatMessage call() throws Exception { throw e; } } + break; } if (actualRetryCount > 0) { @@ -447,6 +469,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 (Throwable e) { + throw reportFailedAndPropagate( + ctx, + ExecutionReporter.EntityTypes.PARSER, + STRUCTURED_OUTPUT, + null, + e, + ExecutionReporter.ProblemCategories.MODEL_OUTPUT_PARSE_ERROR); + } + } + private static void processChatRequest(ChatRequestEvent event, RunnerContext ctx) throws Exception { chat( @@ -528,4 +569,31 @@ public static void processChatRequestOrToolResponse(Event event, RunnerContext c throw new RuntimeException(String.format("Unexpected type event %s", event)); } } + + /** + * Reports a nested execution failure, then always throws the original failure. The Exception + * return type exists so callers must {@code throw} the result and cannot fall through. + */ + private static Exception reportFailedAndPropagate( + RunnerContext ctx, + String entityType, + String entityName, + @Nullable Map entityMetadata, + Throwable error, + String problemCategory) + throws Exception { + if (entityMetadata == null) { + ExecutionReporters.failed(ctx, entityType, entityName, error, problemCategory); + } else { + ExecutionReporters.failed( + ctx, entityType, entityName, entityMetadata, error, problemCategory); + } + if (error instanceof Error) { + throw (Error) error; + } + if (error instanceof Exception) { + throw (Exception) error; + } + throw new RuntimeException(error); + } } 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..e0423d660 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,165 @@ 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); + } catch (Error e) { + ExecutionReporters.failed( + ctx, + ExecutionReporter.EntityTypes.TOOL, + name, + entityMetadata, + e, + ExecutionReporter.ProblemCategories.TOOL_CALL_FAILED); + throw e; + } + 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 caaee783a..ed785f9e6 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 @@ -50,6 +50,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 @@ -111,6 +114,6 @@ public AgentPlan deserialize(JsonParser parser, DeserializationContext ctx) } AgentConfiguration config = new AgentConfiguration(configData); - return new AgentPlan(actions, resourceProviders, config); + return new AgentPlan(actions, 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 4a1494ff1..25ee89932 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 f93a0836f..a1d7b0cf3 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,8 @@ 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.agents.api.trace.LLMExecutionMetadataKeys; import org.apache.flink.metrics.Counter; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -55,6 +57,9 @@ /** Tests for retry behavior in {@link ChatModelAction}. */ class ChatModelActionRetryTest { + private static final Map LLM_METADATA = + Map.of(LLMExecutionMetadataKeys.MODEL, "configured-model"); + @Mock private RunnerContext mockCtx; @Mock private BaseChatModelSetup mockChatModel; @@ -156,6 +161,132 @@ void chatRecordsTokenMetricsWithRequestScopedMetricGroup() throws Exception { verify(mockChatModel, never()).recordTokenMetrics(actionB, "provider-model", 100L, 50L); } + @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", LLM_METADATA); + verify(reporter) + .reportExecutionSucceeded( + ExecutionReporter.EntityTypes.LLM, "test-model", LLM_METADATA); + } + + @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", LLM_METADATA); + verify(reporter, times(2)) + .reportExecutionSucceeded( + ExecutionReporter.EntityTypes.LLM, "test-model", LLM_METADATA); + 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(LLM_METADATA), + 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", LLM_METADATA); + verify(reporter) + .reportExecutionFailed( + eq(ExecutionReporter.EntityTypes.LLM), + eq("test-model"), + eq(LLM_METADATA), + any(RuntimeException.class), + eq(ExecutionReporter.ProblemCategories.MODEL_CALL_FAILED)); + verify(reporter) + .reportExecutionSucceeded( + ExecutionReporter.EntityTypes.LLM, "test-model", LLM_METADATA); + } + @Test void chatRetriesWithExponentialBackoff() throws Exception { // 1 second base interval; fail once then succeed -> wait 1s (1 * 2^0) @@ -356,6 +487,88 @@ 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(chatModel.getModel()).thenReturn("configured-model"); + 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 ed316e483..a87e95e9c 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 @@ -213,6 +213,27 @@ 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 + public void testExplicitEmptyAgentNameIsPreserved() throws Exception { + AgentPlan plan = new AgentPlan(new TestAgent(), new AgentConfiguration(), ""); + + assertThat(plan.getAgentName()).isEqualTo(""); } @Test diff --git a/python/flink_agents/api/core_options.py b/python/flink_agents/api/core_options.py index d5cd70da6..64f6a727a 100644 --- a/python/flink_agents/api/core_options.py +++ b/python/flink_agents/api/core_options.py @@ -206,6 +206,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 2db294d7f..57fa10d1c 100644 --- a/python/flink_agents/api/events/event.py +++ b/python/flink_agents/api/events/event.py @@ -29,6 +29,7 @@ BaseModel, Field, SerializerFunctionWrapHandler, + field_validator, model_serializer, model_validator, ) @@ -72,7 +73,8 @@ class Event(BaseModel, extra="allow"): Attributes: ---------- id : UUID - Random version 4 UUID generated when the Event is created. + Random version 4 UUID generated when the Event is created. An omitted or + explicit ``None`` id is replaced with a new UUID, matching Java. type : str Event type string used for routing. Required for all events. attributes : Dict[str, Any] @@ -97,6 +99,12 @@ class Event(BaseModel, extra="allow"): serialization_alias="upstreamActionName", ) + @field_validator("id", mode="before") + @classmethod + def generate_id_when_explicitly_none(cls, value: Any) -> Any: + """Treat explicit None like an omitted id and mint a per-occurrence UUID.""" + return uuid4() if value is None else value + @staticmethod def __serialize_unknown(field: Any) -> Dict[str, Any]: """Handle serialization of unknown types, specifically Row objects.""" diff --git a/python/flink_agents/api/tests/test_core_options.py b/python/flink_agents/api/tests/test_core_options.py index 5a4e91d26..128e6f3a8 100644 --- a/python/flink_agents/api/tests/test_core_options.py +++ b/python/flink_agents/api/tests/test_core_options.py @@ -108,6 +108,7 @@ def test_agent_config_options_are_explicitly_declared() -> None: 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 condition_failure = options["CONDITION_EVALUATION_FAILURE_STRATEGY"] assert ( condition_failure.get_key() diff --git a/python/flink_agents/api/tests/test_event.py b/python/flink_agents/api/tests/test_event.py index d4b3df62b..4cd3cb56f 100644 --- a/python/flink_agents/api/tests/test_event.py +++ b/python/flink_agents/api/tests/test_event.py @@ -57,6 +57,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): @@ -110,6 +126,29 @@ def test_event_id_cannot_be_reassigned() -> None: event.id = uuid4() +def test_explicit_none_id_generates_uuid() -> None: + event = Event(id=None, type="x") + + assert event.id is not None + assert event.id.version == 4 + assert event.type == "x" + + +def test_from_json_explicit_null_id_generates_uuid() -> None: + event = Event.from_json('{"id":null,"type":"x"}') + + assert event.id is not None + assert event.id.version == 4 + assert event.type == "x" + + +def test_explicit_none_ids_are_distinct() -> None: + first = Event(id=None, type="x") + second = Event(id=None, type="x") + + assert first.id != second.id + + def test_event_json_serialization_with_row() -> None: event = InputEvent(input=Row({"test": "data"})) json_str = event.model_dump_json() 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..db7f8962d --- /dev/null +++ b/python/flink_agents/api/trace/__init__.py @@ -0,0 +1,42 @@ +################################################################################ +# 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_lifecycle_events import ( + ExecutionLifecycleEvents, +) +from flink_agents.api.trace.execution_reporter import ( + ExecutionEntityTypes, + ExecutionProblemCategories, + ExecutionReporter, + ExecutionReporters, +) +from flink_agents.api.trace.llm_execution_metadata_keys import ( + LLMExecutionMetadataKeys, +) +from flink_agents.api.trace.tool_execution_metadata_keys import ( + ToolExecutionMetadataKeys, +) + +__all__ = [ + "ExecutionEntityTypes", + "ExecutionLifecycleEvents", + "ExecutionProblemCategories", + "ExecutionReporter", + "ExecutionReporters", + "LLMExecutionMetadataKeys", + "ToolExecutionMetadataKeys", +] diff --git a/python/flink_agents/api/trace/execution_lifecycle_events.py b/python/flink_agents/api/trace/execution_lifecycle_events.py new file mode 100644 index 000000000..ae32791f0 --- /dev/null +++ b/python/flink_agents/api/trace/execution_lifecycle_events.py @@ -0,0 +1,46 @@ +################################################################################ +# 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 typing import ClassVar + + +class ExecutionLifecycleEvents: + """Framework-owned execution lifecycle Event types and statuses.""" + + EXECUTION_STARTED_EVENT_TYPE = "_execution_started_event" + EXECUTION_FINISHED_EVENT_TYPE = "_execution_finished_event" + EXECUTION_FAILED_EVENT_TYPE = "_execution_failed_event" + EXECUTION_REUSED_EVENT_TYPE = "_execution_reused_event" + + STATUS_STARTED = "started" + STATUS_SUCCESS = "success" + STATUS_FAILED = "failed" + STATUS_REUSED = "reused" + + _EXPECTED_STATUS_BY_EVENT_TYPE: ClassVar[dict[str, str]] = { + EXECUTION_STARTED_EVENT_TYPE: STATUS_STARTED, + EXECUTION_FINISHED_EVENT_TYPE: STATUS_SUCCESS, + EXECUTION_FAILED_EVENT_TYPE: STATUS_FAILED, + EXECUTION_REUSED_EVENT_TYPE: STATUS_REUSED, + } + EVENT_TYPES = frozenset(_EXPECTED_STATUS_BY_EVENT_TYPE) + + @classmethod + def expected_status(cls, event_type: str) -> str | None: + """Return the status required by a framework lifecycle Event type.""" + return cls._EXPECTED_STATUS_BY_EVENT_TYPE.get(event_type) 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/llm_execution_metadata_keys.py b/python/flink_agents/api/trace/llm_execution_metadata_keys.py new file mode 100644 index 000000000..1df601087 --- /dev/null +++ b/python/flink_agents/api/trace/llm_execution_metadata_keys.py @@ -0,0 +1,24 @@ +################################################################################ +# 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 LLMExecutionMetadataKeys: + """Shared field names for LLM execution entity metadata.""" + + MODEL = "model" 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/cli/tests/test_trace_tree.py b/python/flink_agents/cli/tests/test_trace_tree.py index 754896a57..5f3915d85 100644 --- a/python/flink_agents/cli/tests/test_trace_tree.py +++ b/python/flink_agents/cli/tests/test_trace_tree.py @@ -22,6 +22,7 @@ import pytest +from flink_agents.api.trace import ExecutionLifecycleEvents from flink_agents.cli.trace_tree import find_log_files @@ -31,6 +32,26 @@ def _record( upstream_event_id: str | None = None, upstream_action_name: str | None = None, attributes: dict | None = None, +) -> dict: + record = { + "timestamp": "2026-07-17T10:00:00Z", + "eventId": event_id, + "eventType": event_type, + "eventAttributes": attributes or {}, + } + if upstream_event_id is not None: + record["upstreamEventId"] = upstream_event_id + if upstream_action_name is not None: + record["upstreamActionName"] = upstream_action_name + return record + + +def _legacy_record( + event_id: str, + event_type: str, + upstream_event_id: str | None = None, + upstream_action_name: str | None = None, + attributes: dict | None = None, ) -> dict: event = { "id": event_id, @@ -48,6 +69,17 @@ def _record( } +def _execution_record(event_id: str, event_type: str, status: str) -> dict: + return { + **_record(event_id, event_type), + "inputRunId": "run-1", + "executionId": f"execution-{event_id}", + "entityType": "action", + "entityName": "test_action", + "status": status, + } + + def _write_log(path: Path, records: list[dict]) -> None: path.write_text("".join(json.dumps(record) + "\n" for record in records)) @@ -173,6 +205,110 @@ def test_reader_reconstructs_pretty_printed_records(tmp_path: Path) -> None: assert text_result.stderr == "" +def test_reader_merges_flat_and_legacy_records(tmp_path: Path) -> None: + log_path = tmp_path / "events.log" + _write_log( + log_path, + [ + _record("root", "_input_event"), + _legacy_record( + "child", + "ChildEvent", + "root", + "child_action", + {"value": 1}, + ), + _record( + "child", + "ChildEvent", + "root", + "child_action", + {"value": 1}, + ), + ], + ) + + result = _run_reader(log_path, "json") + trace_forest = json.loads(result.stdout) + + assert trace_forest["roots"] == ["root"] + assert trace_forest["nodes"]["root"]["actions"] == [ + {"name": "child_action", "children": ["child"]} + ] + assert trace_forest["nodes"]["child"]["observationCount"] == 2 + assert trace_forest["warnings"] == [] + assert result.stderr == "" + + +def test_reader_ignores_execution_lifecycle_records(tmp_path: Path) -> None: + log_path = tmp_path / "events.log" + _write_log( + log_path, + [ + _record("root", "_input_event"), + _execution_record( + "started", + ExecutionLifecycleEvents.EXECUTION_STARTED_EVENT_TYPE, + ExecutionLifecycleEvents.STATUS_STARTED, + ), + _execution_record( + "finished", + ExecutionLifecycleEvents.EXECUTION_FINISHED_EVENT_TYPE, + ExecutionLifecycleEvents.STATUS_SUCCESS, + ), + _execution_record( + "failed", + ExecutionLifecycleEvents.EXECUTION_FAILED_EVENT_TYPE, + ExecutionLifecycleEvents.STATUS_FAILED, + ), + _execution_record( + "reused", + ExecutionLifecycleEvents.EXECUTION_REUSED_EVENT_TYPE, + ExecutionLifecycleEvents.STATUS_REUSED, + ), + _record("child", "ChildEvent", "root", "child_action"), + ], + ) + + result = _run_reader(log_path, "json") + trace_forest = json.loads(result.stdout) + + assert set(trace_forest["nodes"]) == {"root", "child"} + assert trace_forest["warnings"] == [] + assert result.stderr == "" + + +def test_reader_retains_business_event_using_reserved_lifecycle_type( + tmp_path: Path, +) -> None: + log_path = tmp_path / "events.log" + reserved_type = ExecutionLifecycleEvents.EXECUTION_STARTED_EVENT_TYPE + _write_log( + log_path, + [ + _record("root", "_input_event"), + _record("reserved", reserved_type, "root", "reserved_action"), + _record("child", "ChildEvent", "reserved", "child_action"), + ], + ) + + result = _run_reader(log_path, "json") + trace_forest = json.loads(result.stdout) + + assert set(trace_forest["nodes"]) == {"root", "reserved", "child"} + assert trace_forest["nodes"]["root"]["actions"] == [ + {"name": "reserved_action", "children": ["reserved"]} + ] + assert trace_forest["nodes"]["reserved"]["actions"] == [ + {"name": "child_action", "children": ["child"]} + ] + assert [item["code"] for item in trace_forest["warnings"]] == [ + "RESERVED_EVENT_TYPE" + ] + assert trace_forest["warnings"][0]["eventId"] == "reserved" + assert "[RESERVED_EVENT_TYPE]" in result.stderr + + def test_reader_warns_on_truncated_tail_and_keeps_valid_records( tmp_path: Path, ) -> None: @@ -239,11 +375,26 @@ def test_reader_resynchronizes_after_malformed_pretty_record( ({}, None), ({"eventType": "BadEvent", "event": []}, None), ({"eventType": "BadEvent", "event": {"attributes": {}}}, None), - ({"eventType": 123, "event": {"id": "bad"}}, None), + ({"eventType": 123, "event": {"id": "bad"}}, "bad"), + ( + {"eventId": "bad-flat", "eventType": 123, "eventAttributes": {}}, + "bad-flat", + ), + ( + { + "eventType": "BadEvent", + "event": {"id": "bad-legacy", "attributes": []}, + }, + "bad-legacy", + ), ( { "eventType": "BadEvent", - "event": {"id": "bad", "upstreamEventId": []}, + "event": { + "id": "bad", + "attributes": {}, + "upstreamEventId": [], + }, }, "bad", ), diff --git a/python/flink_agents/cli/trace_tree.py b/python/flink_agents/cli/trace_tree.py index bae28f645..bd379f41a 100644 --- a/python/flink_agents/cli/trace_tree.py +++ b/python/flink_agents/cli/trace_tree.py @@ -22,9 +22,26 @@ from pathlib import Path from typing import Any, Iterator +from flink_agents.api.trace import ExecutionLifecycleEvents + INPUT_EVENT_TYPE = "_input_event" +def _is_execution_lifecycle_record(record: Any) -> bool: + """Return whether a record has the framework lifecycle Event shape.""" + if not isinstance(record, dict): + return False + expected_status = ExecutionLifecycleEvents.expected_status(record.get("eventType")) + return ( + expected_status is not None + and record.get("status") == expected_status + and all( + isinstance(record.get(field), str) and record[field] + for field in ("executionId", "entityType", "entityName") + ) + ) + + def _json_fingerprint(value: Any) -> str: """Return a deterministic, type-sensitive representation of JSON data.""" return json.dumps(value, ensure_ascii=False, separators=(",", ":"), sort_keys=True) @@ -82,6 +99,71 @@ def read_json_objects(path: Path, warnings: list[dict[str, Any]]) -> Iterator[An yield record +def _normalize_event_record( + record: Any, +) -> tuple[dict[str, Any] | None, str | None, str | None]: + """Normalize flat and legacy Event Log records for lineage reconstruction.""" + if not isinstance(record, dict): + return None, None, "record must be a JSON object" + + event_type = record.get("eventType") + flat_record = "eventId" in record or "eventAttributes" in record + if flat_record: + event_id_value = record.get("eventId") + event_attributes = record.get("eventAttributes") + upstream_event_id = record.get("upstreamEventId") + upstream_action_name = record.get("upstreamActionName") + event_id_field = "eventId" + attributes_field = "eventAttributes" + lineage_prefix = "" + else: + event = record.get("event") + if not isinstance(event, dict): + return None, None, "field 'event' must be a JSON object" + event_id_value = event.get("id") + event_attributes = event.get("attributes") + upstream_event_id = event.get("upstreamEventId") + upstream_action_name = event.get("upstreamActionName") + event_id_field = "event.id" + attributes_field = "event.attributes" + lineage_prefix = "event." + + if not isinstance(event_id_value, str) or not event_id_value: + return None, None, f"field '{event_id_field}' must be a non-empty string" + if not isinstance(event_type, str) or not event_type: + return None, event_id_value, "field 'eventType' must be a non-empty string" + if not isinstance(event_attributes, dict): + return ( + None, + event_id_value, + f"field '{attributes_field}' must be a JSON object", + ) + + for field_name, field_value in ( + ("upstreamEventId", upstream_event_id), + ("upstreamActionName", upstream_action_name), + ): + if field_value is not None and not isinstance(field_value, str): + return ( + None, + event_id_value, + f"field '{lineage_prefix}{field_name}' must be a string or null", + ) + + return ( + { + "eventId": event_id_value, + "eventType": event_type, + "timestamp": record.get("timestamp"), + "upstreamEventId": upstream_event_id, + "upstreamActionName": upstream_action_name, + "eventContent": event_attributes, + }, + event_id_value, + None, + ) + + def read_event_records( path: Path, warnings: list[dict[str, Any]] ) -> Iterator[dict[str, Any]]: @@ -93,28 +175,9 @@ def read_event_records( for log_file in log_files: for record in read_json_objects(log_file, warnings): - event_id: str | None = None - invalid_reason: str | None = None - if not isinstance(record, dict): - invalid_reason = "record must be a JSON object" - else: - event = record.get("event") - event_type = record.get("eventType") - if not isinstance(event, dict): - invalid_reason = "field 'event' must be a JSON object" - elif not isinstance(event.get("id"), str) or not event["id"]: - invalid_reason = "field 'event.id' must be a non-empty string" - elif not isinstance(event_type, str) or not event_type: - invalid_reason = "field 'eventType' must be a non-empty string" - else: - event_id = event["id"] - for field_name in ("upstreamEventId", "upstreamActionName"): - field_value = event.get(field_name) - if field_value is not None and not isinstance(field_value, str): - invalid_reason = ( - f"field 'event.{field_name}' must be a string or null" - ) - break + normalized_record, event_id, invalid_reason = _normalize_event_record( + record + ) if invalid_reason is not None: warnings.append( @@ -127,20 +190,21 @@ def read_event_records( ) continue - assert isinstance(record, dict) - event = record["event"] - assert isinstance(event, dict) - event_content = dict(event) - event_content.pop("upstreamEventId", None) - event_content.pop("upstreamActionName", None) - yield { - "eventId": event["id"], - "eventType": record["eventType"], - "timestamp": record.get("timestamp"), - "upstreamEventId": event.get("upstreamEventId"), - "upstreamActionName": event.get("upstreamActionName"), - "eventContent": event_content, - } + assert normalized_record is not None + if _is_execution_lifecycle_record(record): + continue + if normalized_record["eventType"] in ExecutionLifecycleEvents.EVENT_TYPES: + warnings.append( + warning( + "RESERVED_EVENT_TYPE", + event_id, + f"Event {event_id} uses reserved execution lifecycle type " + f"{normalized_record['eventType']} without the corresponding " + "execution lifecycle fields; it was retained as a business Event.", + file_path=log_file, + ) + ) + yield normalized_record def build_trace_forest( diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/memory_event_logging_test.py b/python/flink_agents/e2e_tests/e2e_tests_integration/memory_event_logging_test.py index 0e515f1da..5d9d4cd37 100644 --- a/python/flink_agents/e2e_tests/e2e_tests_integration/memory_event_logging_test.py +++ b/python/flink_agents/e2e_tests/e2e_tests_integration/memory_event_logging_test.py @@ -41,6 +41,7 @@ f"{_purelib}{os.pathsep}{_extra_pythonpath}" if _extra_pythonpath else _purelib ) + class MemoryWritingAgent(Agent): """Agent whose input action writes short-term memory.""" @@ -78,7 +79,7 @@ def test_memory_event_logging(tmp_path: Path) -> None: agents_env.get_config().set_str("baseLogDir", str(event_log_dir)) # Memory operation events stay at defaults; run-begin snapshots are opt-in. agents_env.get_config().set(AgentExecutionOptions.AGENT_RUN_BEGIN_EVENT, True) - # VERBOSE keeps the full event subtree. + # VERBOSE keeps the full event payload. agents_env.get_config().set_str("event-log.level", "VERBOSE") input_datastream = env.from_collection(inputs) @@ -98,15 +99,15 @@ def test_memory_event_logging(tmp_path: Path) -> None: assert EventType.ShortTermWriteEvent in types write = next(r for r in records if r["eventType"] == EventType.ShortTermWriteEvent) - assert write["event"]["attributes"]["value"]["user.tier"] == "gold" - assert write["event"]["attributes"]["key"] == "7" + assert write["eventAttributes"]["value"]["user.tier"] == "gold" + assert write["eventAttributes"]["key"] == "7" begins = [r for r in records if r["eventType"] == EventType.AgentRunBeginEvent] assert len(begins) == len(inputs) - assert {event["event"]["attributes"]["key"] for event in begins} == {"7"} + assert {event["eventAttributes"]["key"] for event in begins} == {"7"} # First run-begin: empty STM; a later one contains the previous run's write. - assert begins[0]["event"]["attributes"]["value"] == {} - assert begins[-1]["event"]["attributes"]["value"].get("user.tier") == "gold" + assert begins[0]["eventAttributes"]["value"] == {} + assert begins[-1]["eventAttributes"]["value"].get("user.tier") == "gold" # Adjacency: each run-begin line directly follows its InputEvent line (both are # emitted synchronously back-to-back) — the reconstruction anchor. 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 4b74a85bd..e5d870611 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 @@ -119,7 +119,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 @@ -130,18 +132,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( @@ -234,15 +236,17 @@ def test_event_lineage_reconstructs_trace_trees_from_runtime_logs( tmp_path: Path, ) -> None: """Test a real PyFlink job produces logs consumable by the Trace Tree reader.""" - event_log_dir = _run_event_logging_pipeline(tmp_path) + event_log_dir = _run_event_logging_pipeline( + tmp_path, config_overrides={"event-log.trace.enabled": "true"} + ) records = _read_log_records(event_log_dir) input_event_ids = { - record["event"]["id"] + record["eventId"] for record in records if record["eventType"] == InputEvent.EVENT_TYPE } output_event_ids = { - record["event"]["id"] + record["eventId"] for record in records if record["eventType"] == OutputEvent.EVENT_TYPE } @@ -253,6 +257,7 @@ def test_event_lineage_reconstructs_trace_trees_from_runtime_logs( assert input_event_ids assert len(output_event_ids) == len(input_event_ids) + assert any(record["eventType"] == "_execution_started_event" for record in records) assert set(trace_forest["roots"]) == input_event_ids assert set(trace_forest["nodes"]) == input_event_ids | output_event_ids assert trace_forest["warnings"] == [] 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 1065c785a..0b6e7f4ae 100644 --- a/python/flink_agents/plan/actions/chat_model_action.py +++ b/python/flink_agents/plan/actions/chat_model_action.py @@ -40,6 +40,12 @@ 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, + LLMExecutionMetadataKeys, +) 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 +272,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 +302,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, @@ -314,17 +350,36 @@ async def chat( response = None actual_retry_count = 0 total_wait_time_sec = 0 + llm_metadata = {LLMExecutionMetadataKeys.MODEL: chat_model.model} 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, llm_metadata + ) + 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, + llm_metadata, + model_error, + ExecutionProblemCategories.MODEL_CALL_FAILED, ) + raise + ExecutionReporters.succeeded( + ctx, ExecutionEntityTypes.LLM, model, llm_metadata + ) if ( request_metric_group is not None @@ -339,7 +394,9 @@ async def chat( request_metric_group, ) 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 27b903f6a..8c365b02f 100644 --- a/python/flink_agents/plan/agent_plan.py +++ b/python/flink_agents/plan/agent_plan.py @@ -72,6 +72,7 @@ class AgentPlan(BaseModel): actions: Dict[str, Action] resource_providers: Dict[ResourceType, Dict[str, ResourceProvider]] | None = None + agent_name: str | None = None config: AgentConfiguration | None = None @field_serializer("resource_providers") @@ -135,7 +136,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 = {} for action in _get_actions(agent) + BUILT_IN_ACTIONS: @@ -155,6 +158,7 @@ def from_agent(agent: Agent, config: AgentConfiguration) -> "AgentPlan": return AgentPlan( actions=actions, resource_providers=resource_providers, + agent_name=agent_name if agent_name is not None else agent.__class__.__name__, config=config, ) @@ -440,14 +444,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..f73648aa8 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,11 +36,19 @@ 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, + LLMExecutionMetadataKeys, +) from flink_agents.plan.actions.chat_model_action import ( chat, process_chat_request_or_tool_response, ) +_LLM_METADATA = {LLMExecutionMetadataKeys.MODEL: "configured-model"} + # ============================================================================ # Mock infrastructure # ============================================================================ @@ -100,6 +111,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, @@ -112,6 +127,7 @@ def _create_mock_runner_context( sent_events = [] metric_group = _MockMetricGroup() sensory_memory = _MockMemoryObject() + chat_model.model = "configured-model" config = MagicMock() option_values = { @@ -126,7 +142,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 +192,17 @@ 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, + _LLM_METADATA, + ) + ctx.report_execution_succeeded.assert_called_once_with( + ExecutionEntityTypes.LLM, + chat_model.connection, + _LLM_METADATA, + ) + 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 +249,17 @@ 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[2] == _LLM_METADATA + assert failed_args[-1] == ExecutionProblemCategories.MODEL_CALL_FAILED + ctx.report_execution_succeeded.assert_called_once_with( + ExecutionEntityTypes.LLM, + "test-model", + _LLM_METADATA, + ) def test_chat_exhausts_retries_and_raises(self) -> None: """All retries exhausted: exception raised, no event sent.""" @@ -246,6 +284,63 @@ 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[2] == _LLM_METADATA + 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", + _LLM_METADATA, + ) + ) + == 2 + ) class TestChatResponseEventRetryFields: @@ -325,9 +420,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 +472,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 +484,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 27781c8e4..e9cca2d2c 100644 --- a/python/flink_agents/plan/tests/resources/agent_plan.json +++ b/python/flink_agents/plan/tests/resources/agent_plan.json @@ -128,6 +128,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 1cb2cfcda..35226da00 100644 --- a/python/flink_agents/plan/tests/test_agent_plan.py +++ b/python/flink_agents/plan/tests/test_agent_plan.py @@ -70,6 +70,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" action = agent_plan.actions["increment"] assert action.name == "increment" func = action.exec @@ -96,6 +97,20 @@ def test_from_agent_preserves_raw_conditions() -> None: ] +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" + + +def test_from_agent_preserves_empty_agent_name() -> None: + agent_plan = AgentPlan.from_agent(AgentForTest(), AgentConfiguration(), "") + + assert agent_plan.agent_name == "" + + class InvalidAgent(Agent): @action(EventType.InputEvent) @staticmethod @@ -164,8 +179,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" ) @@ -545,10 +559,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", @@ -591,7 +609,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 2384dd6ea..955390843 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,22 @@ 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: + break + current = cause + return current + + @dataclass(frozen=True) class _PersistedCallResult: function_id: str @@ -241,7 +260,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 +420,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 ac5819eeb..a5d1515b2 100644 --- a/python/flink_agents/runtime/remote_execution_environment.py +++ b/python/flink_agents/runtime/remote_execution_environment.py @@ -96,7 +96,9 @@ def apply(self, agent: Agent | str) -> "AgentBuilder": if self.__agent_plan_json 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 " @@ -109,9 +111,9 @@ 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] - agent_plan_json = AgentPlan.from_agent(agent, self.__config).model_dump_json( - serialize_as_any=True - ) + agent_plan_json = AgentPlan.from_agent( + agent, self.__config, agent_name + ).model_dump_json(serialize_as_any=True) self.__validate_agent_plan_json(agent_plan_json) self.__agent_plan_json = agent_plan_json diff --git a/python/flink_agents/runtime/skill/skill_tools.py b/python/flink_agents/runtime/skill/skill_tools.py index 3d1a1b852..5178ab308 100644 --- a/python/flink_agents/runtime/skill/skill_tools.py +++ b/python/flink_agents/runtime/skill/skill_tools.py @@ -23,14 +23,27 @@ 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__) +_DEFAULT_SKILL_RESOURCE_PATH = "SKILL.md" + + +def _normalize_skill_resource_path(path: Any, *, missing: bool = False) -> str: + """Resolve a skill resource path the same way ``call()`` loads it.""" + if missing or path is None: + return _DEFAULT_SKILL_RESOURCE_PATH + return str(path) + class LoadSkillArgs(BaseModel): """Arguments for LoadSkillTool.""" @@ -47,7 +60,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 +86,20 @@ 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"]) + metadata[ToolExecutionMetadataKeys.SKILL_RESOURCE_PATH] = ( + _normalize_skill_resource_path( + parameters.get("path"), missing="path" not in parameters + ) + ) + return metadata + def call(self, *args: Any, **kwargs: Any) -> str: """Call the tool to load a skill.""" if args: @@ -81,7 +108,7 @@ def call(self, *args: Any, **kwargs: Any) -> str: parsed_args = LoadSkillArgs(**kwargs) skill_name = parsed_args.name - resource_path = parsed_args.path + resource_path = _normalize_skill_resource_path(parsed_args.path) logger.debug(f"Loading skill resource {resource_path} for {skill_name}.") manager = self._get_skill_manager() 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..be0246ce8 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,28 @@ 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" + + def test_execution_metadata_normalizes_omitted_path(self, tool: LoadSkillTool) -> None: + metadata = tool.get_tool_execution_metadata({"name": "github"}) + assert metadata[ToolExecutionMetadataKeys.SKILL_NAME] == "github" + assert metadata[ToolExecutionMetadataKeys.SKILL_RESOURCE_PATH] == "SKILL.md" + + def test_execution_metadata_normalizes_explicit_none_path( + self, tool: LoadSkillTool + ) -> None: + metadata = tool.get_tool_execution_metadata({"name": "github", "path": None}) + assert metadata[ToolExecutionMetadataKeys.SKILL_NAME] == "github" + assert metadata[ToolExecutionMetadataKeys.SKILL_RESOURCE_PATH] == "SKILL.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..083a3ef73 --- /dev/null +++ b/python/flink_agents/runtime/tests/test_flink_runner_context_trace.py @@ -0,0 +1,83 @@ +################################################################################ +# 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_explicit_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, + ) + + +def test_failed_execution_does_not_follow_implicit_context() -> None: + java_context = MagicMock() + ctx = FlinkRunnerContext.__new__(FlinkRunnerContext) + ctx._j_runner_context = java_context + context_error = ValueError("invalid model response") + reported_error = RuntimeError("structured output failed") + reported_error.__context__ = context_error + assert reported_error.__cause__ is None + + ctx.report_execution_failed( + ExecutionEntityTypes.PARSER, + "structured_output", + {}, + reported_error, + ExecutionProblemCategories.MODEL_OUTPUT_PARSE_ERROR, + ) + + java_context.reportExecutionFailedJson.assert_called_once_with( + ExecutionEntityTypes.PARSER, + "structured_output", + "{}", + "builtins.RuntimeError", + "structured output failed", + 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 2049e1e02..634ed2ae4 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; @@ -43,6 +46,8 @@ import org.apache.flink.agents.runtime.memory.MemoryObjectImpl; import org.apache.flink.agents.runtime.memory.MemoryValueObservation; 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; @@ -61,7 +66,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()); @@ -144,6 +149,10 @@ public CachedMemoryStore getSensoryMemStore() { /** Whether the fixed job-level configuration enables any LTM observation. */ private final boolean ltmObservationConfigured; + @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; @@ -175,17 +184,57 @@ public void switchActionContext( String contextKey, String observationId, boolean observationSuppressed) { + switchActionContext( + actionName, + memoryContext, + contextKey, + observationId, + observationSuppressed, + null, + null); + } + + public void switchActionContext( + String actionName, + MemoryContext memoryContext, + String contextKey, + @Nullable ExecutionTraceContext actionTraceContext, + @Nullable Map activeReportedExecutions) { + switchActionContext( + actionName, + memoryContext, + contextKey, + null, + false, + actionTraceContext, + activeReportedExecutions); + } + + public void switchActionContext( + String actionName, + MemoryContext memoryContext, + String contextKey, + @Nullable String observationId, + boolean observationSuppressed, + @Nullable ExecutionTraceContext actionTraceContext, + @Nullable Map activeReportedExecutions) { this.actionName = actionName; this.memoryContext = memoryContext; this.contextKey = contextKey; this.observationId = observationId; this.observationSuppressed = observationSuppressed; this.ltmObservationEnabled = !observationSuppressed && ltmObservationConfigured; + this.actionTraceContext = actionTraceContext; + this.activeReportedExecutions = activeReportedExecutions; if (ltm != null) { ltm.switchContext(contextKey, observationId, observationSuppressed); } } + public void setExecutionEventSink(@Nullable ExecutionEventSink executionEventSink) { + this.executionEventSink = executionEventSink; + } + public MemoryContext getMemoryContext() { return memoryContext; } @@ -317,6 +366,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..67d37c594 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,91 @@ 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); + } + } + UUID upstreamEventId = parseOptionalUuid(getText(rootNode, "upstreamEventId", false)); + String upstreamActionName = getText(rootNode, "upstreamActionName", false); + Event event = + new Event(eventId, eventType, attributes, upstreamEventId, upstreamActionName); + 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 +164,45 @@ 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 UUID parseOptionalUuid(String value) throws IOException { + return value == null ? null : parseUuid(value); + } + + 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..656595d11 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,68 @@ 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()); + } + gen.writeStringField("eventId", event.getId().toString()); + gen.writeStringField("eventType", event.getType()); + if (event.getUpstreamEventId() != null) { + gen.writeStringField("upstreamEventId", event.getUpstreamEventId().toString()); } - eventNode = reorderEventFields((ObjectNode) eventNode, record.getEvent(), mapper); - gen.writeTree(eventNode); + writeStringFieldIfPresent(gen, "upstreamActionName", event.getUpstreamActionName()); + 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..abffc76d1 --- /dev/null +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/EventLogWriter.java @@ -0,0 +1,179 @@ +/* + * 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.metrics.Counter; +import org.apache.flink.streaming.api.operators.StreamingRuntimeContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nullable; + +import java.util.concurrent.atomic.AtomicBoolean; + +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; + private final AtomicBoolean writeFailureWarned = new AtomicBoolean(); + + @Nullable private Counter writeFailuresCounter; + + 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)); + writeFailuresCounter = builtInMetrics.getEventLogWriteFailuresCounter(); + 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; + } + Exception writeError = null; + try { + eventLogger.append(eventContext, event, traceContext); + } catch (Exception appendError) { + writeError = appendError; + } + try { + eventLogger.flush(); + } catch (Exception flushError) { + if (writeError == null) { + writeError = flushError; + } else if (writeError != flushError) { + writeError.addSuppressed(flushError); + } + } + if (writeError != null) { + recordWriteFailure(writeError); + } + } + + private void recordWriteFailure(Exception writeError) { + if (writeFailuresCounter != null) { + writeFailuresCounter.inc(); + } + if (writeFailureWarned.compareAndSet(false, true)) { + LOG.warn( + "Event Log write failed and was ignored. Subsequent failures will be logged at DEBUG.", + writeError); + } else { + LOG.debug("Event Log write failed and was ignored.", writeError); + } + } + + @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..6e9988bfb 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,15 +28,18 @@ 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; import java.io.FileWriter; +import java.io.IOException; import java.io.PrintWriter; import java.nio.file.Files; import java.nio.file.Path; import java.nio.file.Paths; import java.util.Collections; +import java.util.Iterator; import java.util.Map; /** @@ -174,7 +177,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 +199,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 +208,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,13 +228,29 @@ 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) { throw new IllegalStateException("FileEventLogger not initialized. Call open() first."); } - // Flush the writer to ensure all data is written to the file - writer.flush(); + // checkError flushes first and exposes I/O failures otherwise swallowed by PrintWriter. + if (writer.checkError()) { + throw new IOException("Failed to flush the Event Log file."); + } } /** @@ -247,7 +266,7 @@ public void setTruncatedEventsCounter(Counter counter) { @Override public void close() throws Exception { if (writer != null) { - flush(); + // PrintWriter.close() flushes before releasing the underlying writer. writer.close(); } } diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/JsonTruncator.java b/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/JsonTruncator.java index fa89a95aa..e14ff35a5 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/JsonTruncator.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/JsonTruncator.java @@ -25,7 +25,6 @@ import java.util.ArrayList; import java.util.List; -import java.util.Set; /** * Truncates a Jackson {@link JsonNode} tree per configurable thresholds. @@ -44,22 +43,11 @@ *

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

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

The input node is treated entirely as Event payload. Structural Event Log fields are kept + * outside this node and are not affected by truncation. */ public class JsonTruncator { - private static final Set PROTECTED_FIELDS = - Set.of( - "eventType", - "type", - "id", - "upstreamEventId", - "upstreamActionName", - "attributes"); - private final int maxStringLength; private final int maxArrayElements; private final int maxDepth; @@ -78,53 +66,37 @@ public JsonTruncator(int maxStringLength, int maxArrayElements, int maxDepth) { } /** - * Truncates the given event node in place according to configured thresholds. - * - *

Protected fields ({@code eventType}, {@code type}, {@code id}, {@code upstreamEventId}, - * {@code upstreamActionName}, {@code attributes}) at the top level of the event node are never - * truncated as structural fields. The {@code attributes} envelope is additionally traversed so - * user payload stored inside it is truncated like any other field. + * Truncates the given Event payload in place according to configured thresholds. * - * @param eventNode the top-level event JSON object to truncate + * @param payloadNode the Event payload to truncate * @return {@code true} if any field was truncated, {@code false} if the node was unchanged */ - public boolean truncate(ObjectNode eventNode) { - if (eventNode == null) { + public boolean truncate(ObjectNode payloadNode) { + if (payloadNode == null) { return false; } - boolean truncated = truncateObject(eventNode, 1, true); - // Traverse the protected attributes envelope so its user payload still gets truncated. - JsonNode attributes = eventNode.get("attributes"); - if (attributes instanceof ObjectNode) { - truncated |= truncateObject((ObjectNode) attributes, 1, false); - } - return truncated; + return truncateObject(payloadNode, 1); } /** * Recursively truncates an object node. * * @param node the object node to process - * @param depth current depth (1 = top-level event node) - * @param isTopLevel whether this is the top-level event node (for protected field checks) + * @param depth current depth (1 = payload root) * @return true if any truncation occurred */ - private boolean truncateObject(ObjectNode node, int depth, boolean isTopLevel) { + private boolean truncateObject(ObjectNode node, int depth) { boolean truncated = false; // At max depth, collapse the entire object to retain only scalars if (maxDepth > 0 && depth >= maxDepth) { - return collapseAtMaxDepth(node, isTopLevel); + return collapseAtMaxDepth(node); } List fieldNames = new ArrayList<>(); node.fieldNames().forEachRemaining(fieldNames::add); for (String fieldName : fieldNames) { - if (isTopLevel && PROTECTED_FIELDS.contains(fieldName)) { - continue; - } - JsonNode child = node.get(fieldName); if (child == null) { continue; @@ -149,7 +121,7 @@ private boolean truncateObject(ObjectNode node, int depth, boolean isTopLevel) { truncated |= truncateArrayContents((ArrayNode) child, depth + 1); } } else if (child.isObject()) { - truncated |= truncateObject((ObjectNode) child, depth + 1, false); + truncated |= truncateObject((ObjectNode) child, depth + 1); } } @@ -206,7 +178,7 @@ private boolean truncateArrayContents(ArrayNode array, int depth) { truncated = true; } } else if (element.isObject()) { - truncated |= truncateObject((ObjectNode) element, depth, false); + truncated |= truncateObject((ObjectNode) element, depth); } else if (element.isArray()) { JsonNode replacement = truncateArray((ArrayNode) element); if (replacement != null) { @@ -228,7 +200,7 @@ private boolean truncateArrayContents(ArrayNode array, int depth) { * * @return true if any fields were dropped */ - private boolean collapseAtMaxDepth(ObjectNode node, boolean isTopLevel) { + private boolean collapseAtMaxDepth(ObjectNode node) { List fieldNames = new ArrayList<>(); node.fieldNames().forEachRemaining(fieldNames::add); @@ -237,12 +209,6 @@ private boolean collapseAtMaxDepth(ObjectNode node, boolean isTopLevel) { boolean hasNonScalar = false; for (String fieldName : fieldNames) { - if (isTopLevel && PROTECTED_FIELDS.contains(fieldName)) { - // Protected fields are kept as-is, even if non-scalar - scalarFields.set(fieldName, node.get(fieldName)); - continue; - } - JsonNode child = node.get(fieldName); if (child.isObject() || child.isArray()) { omittedCount++; @@ -267,9 +233,6 @@ private boolean collapseAtMaxDepth(ObjectNode node, boolean isTopLevel) { // Check if any scalar field was replaced boolean stringTruncated = false; for (String fieldName : fieldNames) { - if (isTopLevel && PROTECTED_FIELDS.contains(fieldName)) { - continue; - } JsonNode original = node.get(fieldName); if (original.isTextual()) { JsonNode replacement = truncateString(original.textValue()); 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/metrics/BuiltInMetrics.java b/runtime/src/main/java/org/apache/flink/agents/runtime/metrics/BuiltInMetrics.java index 6bbc03aa6..8e3a17cb8 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/metrics/BuiltInMetrics.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/metrics/BuiltInMetrics.java @@ -39,6 +39,8 @@ public class BuiltInMetrics { private final Counter eventLogTruncatedEvents; + private final Counter eventLogWriteFailures; + private final HashMap actionMetricGroups; public BuiltInMetrics(FlinkAgentsMetricGroupImpl parentMetricGroup, AgentPlan agentPlan) { @@ -51,6 +53,7 @@ public BuiltInMetrics(FlinkAgentsMetricGroupImpl parentMetricGroup, AgentPlan ag parentMetricGroup.getMeter("numOfActionsExecutedPerSec", numOfActionsExecuted); this.eventLogTruncatedEvents = parentMetricGroup.getCounter("eventLogTruncatedEvents"); + this.eventLogWriteFailures = parentMetricGroup.getCounter("eventLogWriteFailures"); this.actionMetricGroups = new HashMap<>(); for (String actionName : agentPlan.getActions().keySet()) { @@ -78,4 +81,9 @@ public void markActionExecuted(String actionName) { public Counter getEventLogTruncatedEventsCounter() { return eventLogTruncatedEvents; } + + /** Returns the counter tracking failed Event Log writes. */ + public Counter getEventLogWriteFailuresCounter() { + return eventLogWriteFailures; + } } 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 27e74620a..529de4179 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 @@ -22,6 +22,9 @@ import org.apache.flink.agents.api.agents.AgentExecutionOptions; import org.apache.flink.agents.api.context.MemoryUpdate; import org.apache.flink.agents.api.event.AgentRunBeginEvent; +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; @@ -29,6 +32,7 @@ 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.memory.MemoryEventBuilder; import org.apache.flink.agents.runtime.memory.MemoryObjectImpl; @@ -36,6 +40,7 @@ import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl; import org.apache.flink.agents.runtime.python.operator.PythonActionTask; import org.apache.flink.agents.runtime.python.utils.PythonActionExecutor; +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; @@ -116,6 +121,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; @@ -143,7 +152,9 @@ public ActionExecutionOperator( this.mailboxExecutor = mailboxExecutor; this.inputIsJava = inputIsJava; this.pythonKeyIsPickled = pythonKeyIsPickled; - 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); this.agentRunBeginEventEnabled = Boolean.TRUE.equals( @@ -211,8 +222,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()); @@ -265,7 +275,17 @@ private void processInputEvent(Object key, Event inputEvent) throws Exception { * `tryProcessActionTaskForKey` to continue processing. */ private void processEvent(Object key, String contextKey, Event event) throws Exception { - eventRouter.notifyEventProcessed(event); + processEvent( + key, + contextKey, + event, + ExecutionTraceContext.forInputRun(contextKey, agentPlan.getAgentName())); + } + + private void processEvent( + Object key, String contextKey, Event event, ExecutionTraceContext traceContext) + throws Exception { + eventRouter.notifyEventProcessed(event, traceContext); boolean isInputEvent = EventUtil.isInputEvent(event); if (EventUtil.isOutputEvent(event)) { @@ -287,14 +307,15 @@ private void processEvent(Object key, String contextKey, Event event) throws Exc // If the event is an InputEvent, we mark that the key is currently being processed. stateManager.addProcessingKey(key); stateManager.initOrIncSequenceNumber(); - tryEmitAgentRunBeginEvent(key, contextKey, event); + tryEmitAgentRunBeginEvent(key, contextKey, event, traceContext); } // We then obtain the triggered action and add ActionTasks to the waiting processing // queue. List triggerActions = eventRouter.getActionsTriggeredBy(event); if (triggerActions != null && !triggerActions.isEmpty()) { for (Action triggerAction : triggerActions) { - stateManager.addActionTask(createActionTask(key, triggerAction, event)); + stateManager.addActionTask( + createActionTask(key, triggerAction, event, traceContext)); } } } @@ -310,7 +331,8 @@ private void processEvent(Object key, String contextKey, Event event) throws Exc * Attempts to emit an {@link AgentRunBeginEvent} for the input before any action triggered by * that input executes. */ - private void tryEmitAgentRunBeginEvent(Object key, String contextKey, Event inputEvent) + private void tryEmitAgentRunBeginEvent( + Object key, String contextKey, Event inputEvent, ExecutionTraceContext traceContext) throws Exception { if (!agentRunBeginEventEnabled) { return; @@ -348,7 +370,7 @@ private void tryEmitAgentRunBeginEvent(Object key, String contextKey, Event inpu } beginEvent.setUpstreamEventId(inputEvent.getId()); beginEvent.setUpstreamActionName(AGENT_RUN_BEGIN_ACTION_NAME); - processEvent(key, contextKey, beginEvent); + processEvent(key, contextKey, beginEvent, traceContext); } private void tryProcessActionTaskForKey(Object key, String contextKey) { @@ -398,7 +420,8 @@ private void processActionTaskForKey(Object key, String contextKey) throws Excep stateManager.getSensoryMemState(), stateManager.getShortTermMemState(), pythonBridge.getPythonRunnerContext(), - ltm); + ltm, + executionEventLogger); long sequenceNumber = stateManager.getSequenceNumber(); boolean isFinished; @@ -430,6 +453,7 @@ private void processActionTaskForKey(Object key, String contextKey) throws Excep .getSensoryMemory() .set(memoryUpdate.getPath(), memoryUpdate.getValue()); } + notifyActionReused(actionTask); } else { // Initialize ActionState if not exists, or use existing one for recovery if (actionState == null) { @@ -440,49 +464,68 @@ private void processActionTaskForKey(Object key, String contextKey) throws Excep key, sequenceNumber, actionTask.action, actionTask.event); } - // Set up durable execution context for fine-grained recovery - durableExecManager.setupDurableExecutionContext( - actionTask, actionState, sequenceNumber); - - ActionTask.ActionTaskResult actionTaskResult; + notifyActionStarted(actionTask); try { - actionTaskResult = - actionTask.invoke( - getRuntimeContext().getUserCodeClassLoader(), - this.pythonBridge.getPythonActionExecutor()); - } catch (Throwable actionFailure) { + // Set up durable execution context for fine-grained recovery + durableExecManager.setupDurableExecutionContext( + actionTask, actionState, sequenceNumber); + + ActionTask.ActionTaskResult actionTaskResult; try { - actionTask.getRunnerContext().discardMemoryObservation(); - } catch (Throwable discardFailure) { - if (discardFailure != actionFailure) { - actionFailure.addSuppressed(discardFailure); + actionTaskResult = + actionTask.invoke( + getRuntimeContext().getUserCodeClassLoader(), + this.pythonBridge.getPythonActionExecutor()); + } catch (Throwable actionFailure) { + try { + actionTask.getRunnerContext().discardMemoryObservation(); + } catch (Throwable discardFailure) { + if (discardFailure != actionFailure) { + actionFailure.addSuppressed(discardFailure); + } } + ExceptionUtils.rethrowException(actionFailure); + throw new AssertionError("Unreachable after rethrowing action failure"); } - ExceptionUtils.rethrowException(actionFailure); - throw new AssertionError("Unreachable after rethrowing action failure"); - } - // 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); - outputEvents = actionTaskResult.getOutputEvents(); - durableExecManager.maybePersistTaskResult( - key, - sequenceNumber, - actionTask.action, - actionTask.event, - actionTask.getRunnerContext(), - actionTaskResult); - isFinished = actionTaskResult.isFinished(); - generatedActionTaskOpt = actionTaskResult.getGeneratedActionTask(); + // 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(); + if (isFinished) { + notifyActionFinished(actionTask); + } + } catch (Throwable t) { + try { + notifyActionFailed(actionTask, t); + } finally { + contextManager.completeActionExecution(actionTask); + } + ExceptionUtils.rethrowException(t); + // Unreachable; required for Java definite-assignment analysis. + return; + } } - for (Event actionOutputEvent : outputEvents) { - processEvent(key, contextKey, actionOutputEvent); + try { + for (Event actionOutputEvent : outputEvents) { + processEvent(key, contextKey, actionOutputEvent, actionTask.getTraceContext()); + } + } finally { + if (isFinished) { + contextManager.completeActionExecution(actionTask); + } } boolean currentInputEventFinished = false; @@ -563,8 +606,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(); @@ -629,11 +672,44 @@ 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, Throwable 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()); 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 e4972d5cb..4a1ee3716 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; @@ -41,17 +44,21 @@ * {@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; - /** Stable identifier for observations produced by this logical action execution. */ protected String observationId; + 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 @@ -60,14 +67,41 @@ public abstract class ActionTask { protected transient RunnerContextImpl runnerContext; public ActionTask(Object key, Event event, Action action) { - this(key, event, action, UUID.randomUUID().toString()); + this( + key, + event, + action, + UUID.randomUUID().toString(), + ExecutionTraceContext.forExecution( + null, null, null, ExecutionReporter.EntityTypes.ACTION, action.getName())); } protected ActionTask(Object key, Event event, Action action, String observationId) { + this( + key, + event, + action, + observationId, + ExecutionTraceContext.forExecution( + null, null, null, ExecutionReporter.EntityTypes.ACTION, action.getName())); + } + + protected ActionTask( + Object key, Event event, Action action, ExecutionTraceContext traceContext) { + this(key, event, action, UUID.randomUUID().toString(), traceContext); + } + + protected ActionTask( + Object key, + Event event, + Action action, + String observationId, + ExecutionTraceContext traceContext) { this.key = key; this.event = event; this.action = action; this.observationId = Objects.requireNonNull(observationId, "observationId"); + this.traceContext = Objects.requireNonNull(traceContext, "traceContext must not be null"); } public RunnerContextImpl getRunnerContext() { @@ -91,6 +125,32 @@ public String getObservationId() { return observationId; } + 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; @@ -99,12 +159,13 @@ public boolean equals(Object o) { return Objects.equals(this.key, other.key) && Objects.equals(this.event, other.event) && Objects.equals(this.action, other.action) - && Objects.equals(this.getObservationId(), other.getObservationId()); + && Objects.equals(this.getObservationId(), other.getObservationId()) + && Objects.equals(this.traceContext, other.traceContext); } @Override public int hashCode() { - return Objects.hash(key, event, action, getObservationId()); + return Objects.hash(key, event, action, getObservationId(), 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 9162a8448..0fd0caa67 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 @@ -18,6 +18,7 @@ package org.apache.flink.agents.runtime.operator; import org.apache.flink.agents.api.event.MemoryEvent; +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; @@ -31,6 +32,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; @@ -46,9 +49,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. *
* @@ -73,13 +78,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); } @@ -148,8 +155,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 @@ -179,7 +187,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 = @@ -207,6 +216,7 @@ void createAndSetRunnerContext( throw new IllegalStateException( "Unsupported action type: " + actionTask.action.getExec().getClass()); } + context.setExecutionEventSink(executionEventSink); RunnerContextImpl.MemoryContext memoryContext; if (actionTaskMemoryContexts.containsKey(actionTask)) { @@ -223,7 +233,9 @@ void createAndSetRunnerContext( memoryContext, contextKey, actionTask.getObservationId(), - MemoryEvent.isMemoryType(actionTask.event.getType())); + MemoryEvent.isMemoryType(actionTask.event.getType()), + actionTask.getTraceContext(), + getOrCreateActiveReportedExecutions(actionTask)); if (context instanceof JavaRunnerContextImpl) { ContinuationContext continuationContext; @@ -271,6 +283,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) { @@ -290,6 +303,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 1ed5025fa..c37e4a690 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,15 +24,11 @@ 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.condition.ActionMatcher; -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; @@ -50,21 +46,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 @@ -74,18 +67,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; @@ -97,14 +88,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<>(); this.actionMatcher = new ActionMatcher(agentPlan); } @@ -113,8 +108,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. */ @@ -124,20 +120,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. * @@ -229,23 +211,20 @@ List getActionsTriggeredBy(Event event) { /** * 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) { @@ -283,7 +262,7 @@ StreamRecord getReusedStreamRecord() { @VisibleForTesting @Nullable EventLogger getEventLogger() { - return eventLogger; + return eventLogWriter.getEventLogger(); } @VisibleForTesting @@ -291,34 +270,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 19bd677f9..4523d9be0 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 98a3c5af2..65399ab7c 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; @@ -44,6 +45,22 @@ protected PythonActionTask(Object key, Event event, Action action, String observ 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); + } + + protected PythonActionTask( + Object key, + Event event, + Action action, + String observationId, + ExecutionTraceContext traceContext) { + super(key, event, action, observationId, traceContext); + checkState(action.getExec() instanceof PythonFunction); + } + public ActionTaskResult invoke(ClassLoader userCodeClassLoader, PythonActionExecutor executor) throws Exception { LOG.debug( @@ -63,7 +80,8 @@ public ActionTaskResult invoke(ClassLoader userCodeClassLoader, PythonActionExec // submit an asynchronous task and return whether the action has been completed. ((PythonRunnerContextImpl) runnerContext).setPythonAwaitableRef(pythonAwaitableRef); ActionTask tempGeneratedActionTask = - new PythonGeneratorActionTask(key, event, action, getObservationId()); + new PythonGeneratorActionTask( + key, event, action, getObservationId(), 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 a6d78eae1..3b0f1a1e6 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,20 @@ public PythonGeneratorActionTask(Object key, Event event, Action action, String super(key, event, action, observationId); } + public PythonGeneratorActionTask( + Object key, Event event, Action action, ExecutionTraceContext traceContext) { + super(key, event, action, traceContext); + } + + public PythonGeneratorActionTask( + Object key, + Event event, + Action action, + String observationId, + ExecutionTraceContext traceContext) { + super(key, event, action, observationId, traceContext); + } + @Override public ActionTaskResult invoke(ClassLoader userCodeClassLoader, PythonActionExecutor executor) throws Exception { @@ -48,7 +63,7 @@ public ActionTaskResult invoke(ClassLoader userCodeClassLoader, PythonActionExec + "re-executing from beginning.", action.getName()); PythonActionTask freshTask = - new PythonActionTask(key, event, action, getObservationId()); + new PythonActionTask(key, event, action, getObservationId(), 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..abfeb306f 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,13 +65,23 @@ 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"))); + } + metadata.put( + ToolExecutionMetadataKeys.SKILL_RESOURCE_PATH, normalizeResourcePath(parameters)); + return metadata; + } + @Override public ToolResponse call(ToolParameters parameters) { String name = parameters.getParameter("name", String.class); - String path = - parameters.hasParameter("path") - ? parameters.getParameter("path", String.class) - : "SKILL.md"; + String path = normalizeResourcePath(parameters); SkillManager manager; try { @@ -143,4 +157,12 @@ private SkillManager resolveSkillManager() throws Exception { } return null; } + + private static String normalizeResourcePath(ToolParameters parameters) { + if (parameters == null || !parameters.hasParameter("path")) { + return "SKILL.md"; + } + String path = parameters.getParameter("path", String.class); + return path == null ? "SKILL.md" : path; + } } 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..be3090ae5 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextImplExecutionReporterTest.java @@ -0,0 +1,185 @@ +/* + * 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.plan.AgentPlan; +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, () -> {}, emptyAgentPlan(), 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, () -> {}, emptyAgentPlan(), 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, () -> {}, emptyAgentPlan(), 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 AgentPlan emptyAgentPlan() { + return new AgentPlan(new HashMap<>(), new HashMap<>()); + } + + 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 0af92189e..1e6549125 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,13 +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}. @@ -47,130 +56,224 @@ 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); UUID upstreamEventId = UUID.randomUUID(); originalEvent.setUpstreamEventId(upstreamEventId); originalEvent.setUpstreamActionName("custom_action"); - 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); @@ -184,20 +287,14 @@ void testDeserializeCustomEvent() throws Exception { @Test void testRoundTripSerialization() throws Exception { - // Given InputEvent originalEvent = new InputEvent("round trip test"); - EventContext originalContext = new EventContext(originalEvent); - EventLogRecord originalRecord = new EventLogRecord(originalContext, originalEvent); + EventLogRecord originalRecord = record(originalEvent, null); - // 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(originalEvent.getType(), deserializedRecord.getEventContext().getEventType()); assertEquals(InputEvent.EVENT_TYPE, deserializedRecord.getEvent().getType()); - InputEvent deserializedEvent = InputEvent.fromEvent(deserializedRecord.getEvent()); assertEquals(originalEvent.getInput(), deserializedEvent.getInput()); } @@ -208,83 +305,101 @@ void testRoundTripLineageFields() throws Exception { Event originalEvent = new Event("ChildEvent"); originalEvent.setUpstreamEventId(upstreamEventId); originalEvent.setUpstreamActionName("child_action"); - EventLogRecord originalRecord = - new EventLogRecord(new EventContext(originalEvent), originalEvent); + EventLogRecord originalRecord = record(originalEvent, null); String json = objectMapper.writeValueAsString(originalRecord); EventLogRecord deserializedRecord = objectMapper.readValue(json, EventLogRecord.class); - JsonNode eventNode = objectMapper.readTree(json).get("event"); - assertEquals(upstreamEventId.toString(), eventNode.get("upstreamEventId").asText()); - assertEquals("child_action", eventNode.get("upstreamActionName").asText()); + JsonNode jsonNode = objectMapper.readTree(json); + assertEquals(upstreamEventId.toString(), jsonNode.get("upstreamEventId").asText()); + assertEquals("child_action", jsonNode.get("upstreamActionName").asText()); assertEquals(upstreamEventId, deserializedRecord.getEvent().getUpstreamEventId()); assertEquals("child_action", deserializedRecord.getEvent().getUpstreamActionName()); } @Test void testSerializeUnifiedEvent() throws Exception { - // Given - a unified event with user-defined type - java.util.Map attrs = new java.util.HashMap<>(); + Map attrs = new HashMap<>(); attrs.put("msg", "hello"); Event unifiedEvent = new Event("MyCustomEvent", attrs); - EventContext context = new EventContext(unifiedEvent); - EventLogRecord record = new EventLogRecord(context, unifiedEvent); + EventLogRecord record = record(unifiedEvent, null); - // 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()); + assertEquals("MyCustomEvent", jsonNode.get("eventType").asText()); + assertEquals("hello", jsonNode.get("eventAttributes").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<>(); + Map attrs = new 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); + EventLogRecord originalRecord = record(originalEvent, null); 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) + EventContext deserializedContext = deserializedRecord.getEventContext(); 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..e682e35d7 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/EventLogWriterTest.java @@ -0,0 +1,149 @@ +/* + * 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.apache.flink.agents.runtime.metrics.BuiltInMetrics; +import org.apache.flink.metrics.SimpleCounter; +import org.apache.flink.streaming.api.operators.StreamingRuntimeContext; +import org.junit.jupiter.api.Test; + +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.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; + +/** Contract tests for {@link EventLogWriter}. */ +class EventLogWriterTest { + + @Test + void appendAndFlushAttemptsBothOperationsAndCountsOneFailure() throws Exception { + EventLogger mockLogger = mock(EventLogger.class); + EventLogWriter writer = EventLogWriter.forEventLogger(mockLogger); + SimpleCounter writeFailures = openWithWriteFailureCounter(writer); + InputEvent inputEvent = new InputEvent(1L); + EventContext eventContext = new EventContext(inputEvent); + doThrow(new RuntimeException("append failed")) + .when(mockLogger) + .append(any(EventContext.class), eq(inputEvent), isNull()); + doThrow(new RuntimeException("flush failed")).when(mockLogger).flush(); + + assertThatCode(() -> writer.appendBusinessEventAndFlush(eventContext, inputEvent, null)) + .doesNotThrowAnyException(); + verify(mockLogger).flush(); + assertThat(writeFailures.getCount()).isEqualTo(1); + } + + @Test + void appendAndFlushIgnoresFlushFailure() throws Exception { + EventLogger mockLogger = mock(EventLogger.class); + EventLogWriter writer = EventLogWriter.forEventLogger(mockLogger); + SimpleCounter writeFailures = openWithWriteFailureCounter(writer); + 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(); + assertThat(writeFailures.getCount()).isEqualTo(1); + } + + @Test + void appendAndFlushCountsEveryFailedWriteAttempt() throws Exception { + EventLogger mockLogger = mock(EventLogger.class); + EventLogWriter writer = EventLogWriter.forEventLogger(mockLogger); + SimpleCounter writeFailures = openWithWriteFailureCounter(writer); + InputEvent inputEvent = new InputEvent(1L); + EventContext eventContext = new EventContext(inputEvent); + doThrow(new RuntimeException("append failed")) + .when(mockLogger) + .append(any(EventContext.class), eq(inputEvent), isNull()); + + writer.appendBusinessEventAndFlush(eventContext, inputEvent, null); + writer.appendBusinessEventAndFlush(eventContext, inputEvent, null); + + verify(mockLogger, times(2)).flush(); + assertThat(writeFailures.getCount()).isEqualTo(2); + } + + @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(); + } + + private static SimpleCounter openWithWriteFailureCounter(EventLogWriter writer) + throws Exception { + SimpleCounter writeFailures = new SimpleCounter(); + BuiltInMetrics builtInMetrics = mock(BuiltInMetrics.class); + when(builtInMetrics.getEventLogWriteFailuresCounter()).thenReturn(writeFailures); + writer.open(mock(StreamingRuntimeContext.class), builtInMetrics); + return writeFailures; + } +} 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 05bbdbb44..dd2d24a6c 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(); @@ -177,10 +181,10 @@ void testAppendWithCustomEvent() throws Exception { UUID upstreamEventId = UUID.randomUUID(); customEvent.setUpstreamEventId(upstreamEventId); customEvent.setUpstreamActionName("custom_action"); - EventContext context = new EventContext(customEvent); + ExecutionTraceContext context = null; // When - logger.append(context, customEvent); + append(logger, customEvent, context); logger.flush(); // Then @@ -190,9 +194,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()); @@ -216,16 +220,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(); @@ -236,13 +240,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); @@ -259,8 +260,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 @@ -269,8 +270,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(); @@ -298,11 +299,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 = @@ -326,7 +325,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) @@ -361,9 +360,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(); @@ -372,9 +371,13 @@ void testStandardLevelTruncation() throws Exception { JsonNode jsonNode = objectMapper.readTree(lines.get(0)); assertEquals("STANDARD", jsonNode.get("logLevel").asText()); + assertEquals( + event.getId().toString(), + jsonNode.get("eventId").textValue(), + "Top-level Event identity should not be truncated"); // 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"), @@ -395,9 +398,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(); @@ -408,7 +411,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"); @@ -427,9 +430,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(); @@ -451,12 +454,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(); @@ -466,13 +469,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 @@ -480,9 +482,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(); @@ -491,7 +493,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()); } @@ -533,11 +537,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(); @@ -555,6 +559,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/JsonTruncatorTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/JsonTruncatorTest.java index 504b2d342..1c51c0cdf 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/JsonTruncatorTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/JsonTruncatorTest.java @@ -133,40 +133,19 @@ void testDisabledThresholds() { } @Test - void testProtectedFields() { - JsonTruncator truncator = new JsonTruncator(5, 2, 2); + void testFormerEnvelopeFieldNamesAreTruncatedAsPayload() { + JsonTruncator truncator = new JsonTruncator(5, 0, 0); ObjectNode node = MAPPER.createObjectNode(); - // Protected fields should never be truncated - node.put("eventType", "org.apache.flink.agents.api.event.ChatRequestEvent"); - node.put("type", "org.apache.flink.agents.api.event.ChatRequestEvent"); - node.put("id", "a-very-long-identifier-that-exceeds-the-limit"); - node.put("upstreamEventId", "a-very-long-upstream-event-identifier"); - node.put("upstreamActionName", "a-very-long-action-name"); - ObjectNode attributes = MAPPER.createObjectNode(); - attributes.put("key", "business payload should be truncated"); - attributes.set("nested", MAPPER.createObjectNode().put("deep", "data")); - node.set("attributes", attributes); - // Non-protected field should be truncated - node.put("content", "This should be truncated"); + node.put("eventType", "a-long-event-type"); + node.put("id", "a-long-identifier"); + node.put("attributes", "a-long-attribute-value"); boolean result = truncator.truncate(node); assertThat(result).isTrue(); - // Protected fields remain untouched - assertThat(node.get("eventType").asText()) - .isEqualTo("org.apache.flink.agents.api.event.ChatRequestEvent"); - assertThat(node.get("type").asText()) - .isEqualTo("org.apache.flink.agents.api.event.ChatRequestEvent"); - assertThat(node.get("id").asText()) - .isEqualTo("a-very-long-identifier-that-exceeds-the-limit"); - assertThat(node.get("upstreamEventId").asText()) - .isEqualTo("a-very-long-upstream-event-identifier"); - assertThat(node.get("upstreamActionName").asText()).isEqualTo("a-very-long-action-name"); - assertThat(node.get("attributes").isObject()).isTrue(); - assertThat(node.get("attributes").get("key").get("truncatedString").asText()) - .isEqualTo("busin..."); - // Non-protected field is truncated - assertThat(node.get("content").get("truncatedString")).isNotNull(); + assertThat(node.get("eventType").get("truncatedString").asText()).isEqualTo("a-lon..."); + assertThat(node.get("id").get("truncatedString").asText()).isEqualTo("a-lon..."); + assertThat(node.get("attributes").get("truncatedString").asText()).isEqualTo("a-lon..."); } @Test 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..4c732d0b4 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()); @@ -228,6 +230,35 @@ void testLogLevelAppearsInJson() throws Exception { assertNotNull(jsonNode.get("logLevel"), "logLevel field should be present in JSON output"); } + @Test + void testStandardLevelTruncatesOnlyEventAttributes() throws Exception { + Map agentConfig = new HashMap<>(); + agentConfig.put(AgentConfigOptions.EVENT_LOG_LEVEL.getKey(), "STANDARD"); + agentConfig.put(AgentConfigOptions.EVENT_LOG_MAX_STRING_LENGTH.getKey(), 10); + EventLoggerConfig config = + EventLoggerConfig.builder() + .loggerType(LoggerType.SLF4J) + .property(EventLoggerConfig.AGENT_CONFIG_PROPERTY_KEY, agentConfig) + .build(); + logger = new Slf4jEventLogger(config); + logger.open(openParams); + + InputEvent inputEvent = + new InputEvent("this is a very long string that exceeds ten characters"); + append(inputEvent, null); + + List messages = testAppender.getMessages(); + assertEquals(1, messages.size()); + JsonNode jsonNode = objectMapper.readTree(messages.get(0)); + assertEquals( + inputEvent.getId().toString(), + jsonNode.get("eventId").asText(), + "Top-level Event identity should not be truncated"); + assertTrue( + jsonNode.get("eventAttributes").get("input").has("truncatedString"), + "Long Event payload strings should be truncated"); + } + @Test void testDefaultIsNotPrettyPrinted() throws Exception { // Default config should produce single-line JSON, matching AgentConfigOptions.PRETTY_PRINT @@ -237,7 +268,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 +291,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 +312,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 2826a9bb4..d92fa43d9 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 @@ -28,19 +28,32 @@ import org.apache.flink.agents.api.context.MemoryObject; import org.apache.flink.agents.api.context.RunnerContext; import org.apache.flink.agents.api.event.ShortTermWriteEvent; +import org.apache.flink.agents.api.event.ToolRequestEvent; 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.memory.MemorySet; +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.AgentConfiguration; import org.apache.flink.agents.plan.AgentPlan; import org.apache.flink.agents.plan.JavaFunction; import org.apache.flink.agents.plan.actions.Action; +import org.apache.flink.agents.plan.actions.ToolCallAction; +import org.apache.flink.agents.plan.resourceprovider.JavaSerializableResourceProvider; +import org.apache.flink.agents.plan.resourceprovider.ResourceProvider; +import org.apache.flink.agents.plan.tools.FunctionTool; import org.apache.flink.agents.runtime.actionstate.ActionState; import org.apache.flink.agents.runtime.actionstate.ActionStateSerde; import org.apache.flink.agents.runtime.actionstate.CallResult; import org.apache.flink.agents.runtime.actionstate.InMemoryActionStateStore; import org.apache.flink.agents.runtime.eventlog.FileEventLogger; +import org.apache.flink.agents.runtime.eventlog.Slf4jEventLogger; import org.apache.flink.agents.runtime.memory.Mem0LongTermMemory; import org.apache.flink.api.common.typeinfo.TypeInformation; import org.apache.flink.api.java.functions.KeySelector; @@ -52,6 +65,7 @@ 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; @@ -85,6 +99,13 @@ void resetReconcilableFixtures() { TestAgent.resetReconcilableRecoveryFixture(); TestAgent.resetMixedRecoveryFixture(); TestAgent.FOLLOWING_ACTION_EXECUTED.set(false); + RecordingEventLogger.reset(); + EventLoggerFactory.registerFactory(LoggerType.SLF4J, config -> new RecordingEventLogger()); + } + + @AfterEach + void restoreEventLoggerFactory() { + EventLoggerFactory.registerFactory(LoggerType.SLF4J, Slf4jEventLogger::new); } @Test @@ -389,11 +410,18 @@ void testMemoryAccessProhibitedOutsideMailboxThread() throws Exception { } @Test - void testMailboxSubmittedActionTaskPropagatesError() throws Exception { + void testMailboxSubmittedActionTaskPropagatesErrorAndClosesActionLifecycle() throws Exception { + AgentPlan basePlan = TestAgent.getLinkageErrorAgentPlan(); + AgentPlan agentPlan = + new AgentPlan( + basePlan.getActions(), + basePlan.getResourceProviders(), + traceEnabledConfig(), + basePlan.getAgentName()); + try (KeyedOneInputStreamOperatorTestHarness testHarness = new KeyedOneInputStreamOperatorTestHarness<>( - new ActionExecutionOperatorFactory( - TestAgent.getLinkageErrorAgentPlan(), true), + new ActionExecutionOperatorFactory(agentPlan, true), (KeySelector) value -> value, TypeInformation.of(Long.class))) { testHarness.open(); @@ -407,6 +435,66 @@ void testMailboxSubmittedActionTaskPropagatesError() throws Exception { .isInstanceOf(NoClassDefFoundError.class) .hasMessageContaining("synthetic missing runtime dependency"); } + + RecordedEvent started = + findRecordedLifecycleEvent( + ExecutionLifecycleEvents.EXECUTION_STARTED_EVENT_TYPE, + "linkageErrorAction", + ExecutionLifecycleEvents.STATUS_STARTED); + RecordedEvent failed = + findRecordedLifecycleEvent( + ExecutionLifecycleEvents.EXECUTION_FAILED_EVENT_TYPE, + "linkageErrorAction", + ExecutionLifecycleEvents.STATUS_FAILED); + assertThat(failed.traceContext().getExecutionId()) + .isEqualTo(started.traceContext().getExecutionId()); + assertThat(failed.event.getAttr("errorType")) + .isEqualTo(NoClassDefFoundError.class.getName()); + } + + @Test + void testToolLinkageErrorEmitsFailedLifecycleBeforeActionFailure() throws Exception { + AgentPlan basePlan = TestAgent.getLinkageErrorToolAgentPlan(); + AgentPlan agentPlan = + new AgentPlan( + basePlan.getActions(), + basePlan.getResourceProviders(), + traceEnabledConfig(), + basePlan.getAgentName()); + + 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<>(0L)); + assertThatThrownBy(() -> operator.waitInFlightEventsFinished()) + .hasCauseInstanceOf(ActionExecutionOperator.ActionTaskExecutionException.class) + .rootCause() + .isInstanceOf(NoClassDefFoundError.class) + .hasMessageContaining("synthetic missing runtime dependency"); + } + + RecordedEvent started = + findRecordedLifecycleEvent( + ExecutionLifecycleEvents.EXECUTION_STARTED_EVENT_TYPE, + "linkageErrorTool", + ExecutionLifecycleEvents.STATUS_STARTED); + RecordedEvent failed = + findRecordedLifecycleEvent( + ExecutionLifecycleEvents.EXECUTION_FAILED_EVENT_TYPE, + "linkageErrorTool", + ExecutionLifecycleEvents.STATUS_FAILED); + assertThat(started.traceContext().getEntityType()) + .isEqualTo(ExecutionReporter.EntityTypes.TOOL); + assertThat(failed.traceContext().getExecutionId()) + .isEqualTo(started.traceContext().getExecutionId()); + assertThat(failed.event.getAttr("errorType")) + .isEqualTo(NoClassDefFoundError.class.getName()); } @Test @@ -678,10 +766,12 @@ 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()); } } @@ -707,9 +797,144 @@ private static Map recordEventSnapshotsByType( return snapshotsByType; } + /** Fails while processing an Action's output Event. */ + public static class FailingMiddleEventListener implements EventListener { + @Override + public void onEventProcessed(EventContext context, Event event) { + if (TestAgent.MiddleEvent.EVENT_TYPE.equals(event.getType())) { + throw new IllegalStateException("Failed to process Action output Event"); + } + } + } + + /** 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); @@ -739,13 +964,213 @@ 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.getResourceProviders(), + config, + basePlan.getAgentName()); + + 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 testActionFinishesBeforeItsOutputEventIsProcessed() throws Exception { + AgentConfiguration config = traceEnabledConfig(); + config.set( + AgentConfigOptions.EVENT_LISTENERS, + List.of(FailingMiddleEventListener.class.getName())); + 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)); + assertThatThrownBy(operator::waitInFlightEventsFinished) + .hasCauseInstanceOf(ActionExecutionOperator.ActionTaskExecutionException.class) + .rootCause() + .hasMessage("Failed to process Action output Event"); + } + + RecordedEvent started = + findRecordedLifecycleEvent( + ExecutionLifecycleEvents.EXECUTION_STARTED_EVENT_TYPE, + "action1", + ExecutionLifecycleEvents.STATUS_STARTED); + RecordedEvent finished = + findRecordedLifecycleEvent( + ExecutionLifecycleEvents.EXECUTION_FINISHED_EVENT_TYPE, + "action1", + ExecutionLifecycleEvents.STATUS_SUCCESS); + assertThat(finished.traceContext().getExecutionId()) + .isEqualTo(started.traceContext().getExecutionId()); + assertThat(RecordingEventLogger.events()) + .noneMatch( + record -> + ExecutionLifecycleEvents.EXECUTION_FAILED_EVENT_TYPE.equals( + record.event.getType()) + && "action1".equals(record.traceContext().getEntityName())); + } + + @Test + void testActionContinuationEmitsStartedLifecycleEventOnce() throws Exception { + final AgentConfiguration config = traceEnabledConfig(); + AgentPlan basePlan = TestAgent.getAsyncAgentPlan(false); + AgentPlan agentPlan = + new AgentPlan( + basePlan.getActions(), + basePlan.getResourceProviders(), + config, + basePlan.getAgentName()); + + 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 void testDoesNotPruneBeforeCheckpointComplete() throws Exception { AgentPlan agentPlanWithStateStore = TestAgent.getAgentPlan(false); @@ -802,6 +1227,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"; @@ -1134,6 +1622,142 @@ void testReplaySkipsCompletedActions() throws Exception { } } + @Test + void testActionStateStoreReplayRecordsReusedExecutions() throws Exception { + AgentConfiguration config = traceEnabledConfig(); + AgentPlan basePlan = TestAgent.getAgentPlan(false); + AgentPlan agentPlanWithStateStore = + new AgentPlan( + basePlan.getActions(), + basePlan.getResourceProviders(), + config, + basePlan.getAgentName()); + InMemoryActionStateStore actionStateStore; + String originalInputRunId; + UUID originalMiddleEventId; + UUID originalOutputEventId; + + try (KeyedOneInputStreamOperatorTestHarness testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory<>( + agentPlanWithStateStore, true, new InMemoryActionStateStore(false)), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) testHarness.getOperator(); + actionStateStore = + (InMemoryActionStateStore) + operator.getDurableExecutionManager().getActionStateStore(); + + long inputValue = 7L; + 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<>( + agentPlanWithStateStore, true, actionStateStore), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) testHarness.getOperator(); + + long inputValue = 7L; + testHarness.processElement(new StreamRecord<>(inputValue)); + operator.waitInFlightEventsFinished(); + + List> outputRecords = + (List>) testHarness.getRecordOutput(); + assertThat(outputRecords).hasSize(1); + assertThat(outputRecords.get(0).getValue()).isEqualTo((inputValue + 1) * 2); + assertThat(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()); + } + } + @Test void testReplayRebindsOutputLineage() throws Exception { AgentPlan agentPlan = TestAgent.getAgentPlan(false); @@ -2436,6 +3060,20 @@ public static AgentPlan getDurableExceptionAgentPlan() { return null; } + public static void requestLinkageErrorTool(Event event, RunnerContext context) { + Map function = new HashMap<>(); + function.put("name", "linkageErrorTool"); + function.put("arguments", new HashMap()); + Map toolCall = new HashMap<>(); + toolCall.put("id", "call-1"); + toolCall.put("function", function); + context.sendEvent(new ToolRequestEvent("unused-model", List.of(toolCall))); + } + + public static String linkageErrorTool() { + throw new NoClassDefFoundError("synthetic missing runtime dependency"); + } + public static AgentPlan getLinkageErrorAgentPlan() { try { Map actions = new HashMap<>(); @@ -2457,6 +3095,41 @@ public static AgentPlan getLinkageErrorAgentPlan() { return null; } + public static AgentPlan getLinkageErrorToolAgentPlan() { + try { + Action requestAction = + new Action( + "requestLinkageErrorTool", + new JavaFunction( + TestAgent.class, + "requestLinkageErrorTool", + new Class[] {Event.class, RunnerContext.class}), + Collections.singletonList(InputEvent.EVENT_TYPE)); + Action toolCallAction = ToolCallAction.getToolCallAction(); + Map actions = new LinkedHashMap<>(); + actions.put(requestAction.getName(), requestAction); + actions.put(toolCallAction.getName(), toolCallAction); + + FunctionTool tool = + FunctionTool.fromStaticMethod( + "Throws a linkage error.", + TestAgent.class.getMethod("linkageErrorTool")); + Map tools = new HashMap<>(); + tools.put( + "linkageErrorTool", + JavaSerializableResourceProvider.createResourceProvider( + "linkageErrorTool", ResourceType.TOOL, tool)); + Map> resourceProviders = + new HashMap<>(); + resourceProviders.put(ResourceType.TOOL, tools); + + return new AgentPlan(actions, resourceProviders); + } catch (Exception e) { + ExceptionUtils.rethrow(e); + } + return null; + } + public static AgentPlan getBestEffortMemoryObservationPlan() { try { Action action = 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 c5894df0f..468a3ed7f 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; @@ -29,10 +31,19 @@ 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.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; @@ -163,6 +174,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 @@ -185,6 +197,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(); } } @@ -244,17 +334,29 @@ 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); + invokeCreateAndSetRunnerContext(mgr, task, null, null); } - @SuppressWarnings("unchecked") private static void invokeCreateAndSetRunnerContext( ActionTaskContextManager mgr, ActionTask task, InteranlBaseLongTermMemory longTermMemory) { + invokeCreateAndSetRunnerContext(mgr, task, longTermMemory, null); + } + + private static void invokeCreateAndSetRunnerContext( + ActionTaskContextManager mgr, ActionTask task, ExecutionEventSink executionEventSink) { + invokeCreateAndSetRunnerContext(mgr, task, null, executionEventSink); + } + + @SuppressWarnings("unchecked") + private static void invokeCreateAndSetRunnerContext( + ActionTaskContextManager mgr, + ActionTask task, + InteranlBaseLongTermMemory longTermMemory, + ExecutionEventSink executionEventSink) { AgentPlan plan = newEmptyAgentPlan(); ResourceCache cache = mock(ResourceCache.class); FlinkAgentsMetricGroupImpl metricGroup = @@ -272,7 +374,8 @@ private static void invokeCreateAndSetRunnerContext( sensoryMem, shortTermMem, /* pythonRunnerContext */ null, - longTermMemory); + longTermMemory, + executionEventSink); } private static AgentPlan newEmptyAgentPlan() { diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/EventLineageReconstructionTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/EventLineageReconstructionTest.java index f4af79331..f0210057c 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/EventLineageReconstructionTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/EventLineageReconstructionTest.java @@ -72,12 +72,10 @@ void recordsMinimalEventLineageInTheEventLog(@TempDir Path logDir) throws Except MAPPER.treeToValue(recordsByType.get(InputEvent.EVENT_TYPE), EventLogRecord.class); EventLogRecord outputRecord = MAPPER.treeToValue(recordsByType.get(OutputEvent.EVENT_TYPE), EventLogRecord.class); - JsonNode input = recordsByType.get(InputEvent.EVENT_TYPE).get("event"); + JsonNode input = recordsByType.get(InputEvent.EVENT_TYPE); JsonNode middle = - recordsByType - .get(ActionExecutionOperatorTest.TestAgent.MiddleEvent.EVENT_TYPE) - .get("event"); - JsonNode output = recordsByType.get(OutputEvent.EVENT_TYPE).get("event"); + recordsByType.get(ActionExecutionOperatorTest.TestAgent.MiddleEvent.EVENT_TYPE); + JsonNode output = recordsByType.get(OutputEvent.EVENT_TYPE); assertThat(recordsByType).hasSize(3); assertThat(inputRecord.getEvent().getType()).isEqualTo(InputEvent.EVENT_TYPE); @@ -87,16 +85,17 @@ void recordsMinimalEventLineageInTheEventLog(@TempDir Path logDir) throws Except assertThat(middle.get("upstreamEventId").isTextual()).isTrue(); assertThat(middle.get("upstreamActionName").isTextual()).isTrue(); - assertThat(middle.get("upstreamEventId").asText()).isEqualTo(input.get("id").asText()); + assertThat(middle.get("upstreamEventId").asText()).isEqualTo(input.get("eventId").asText()); assertThat(middle.get("upstreamActionName").asText()).isEqualTo("action1"); assertThat(output.get("upstreamEventId").isTextual()).isTrue(); assertThat(output.get("upstreamActionName").isTextual()).isTrue(); - assertThat(output.get("upstreamEventId").asText()).isEqualTo(middle.get("id").asText()); + assertThat(output.get("upstreamEventId").asText()) + .isEqualTo(middle.get("eventId").asText()); assertThat(output.get("upstreamActionName").asText()).isEqualTo("action2"); - assertThat(middle.get("attributes").has("upstreamEventId")).isFalse(); - assertThat(middle.get("attributes").has("upstreamActionName")).isFalse(); + assertThat(middle.get("eventAttributes").has("upstreamEventId")).isFalse(); + assertThat(middle.get("eventAttributes").has("upstreamActionName")).isFalse(); } private static AgentConfiguration fileLoggerConfig(Path logDir) { 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 dc224ba8c..1cba93622 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..f131cd07e 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,50 @@ 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 executionMetadataNormalizesOmittedPathToSkillMd() { + Map metadata = + tool(contextWithSkills()).getToolExecutionMetadata(args("github", null)); + + assertEquals("github", metadata.get(ToolExecutionMetadataKeys.SKILL_NAME)); + assertEquals("SKILL.md", metadata.get(ToolExecutionMetadataKeys.SKILL_RESOURCE_PATH)); + } + + @Test + void executionMetadataNormalizesExplicitNullPathToSkillMd() { + Map parameters = new HashMap<>(); + parameters.put("name", "github"); + parameters.put("path", null); + + Map metadata = + tool(contextWithSkills()).getToolExecutionMetadata(new ToolParameters(parameters)); + + assertEquals("github", metadata.get(ToolExecutionMetadataKeys.SKILL_NAME)); + assertEquals("SKILL.md", metadata.get(ToolExecutionMetadataKeys.SKILL_RESOURCE_PATH)); + } + + @Test + void explicitNullPathLoadsSkillContentEnvelope() { + Map parameters = new HashMap<>(); + parameters.put("name", "github"); + parameters.put("path", null); + LoadSkillTool t = tool(contextWithSkills()); + + ToolResponse resp = t.call(new ToolParameters(parameters)); + + String out = (String) resp.getResult(); + assertTrue(out.startsWith("")); + } + @Test void missingResourceReportsAvailable() { LoadSkillTool t = tool(contextWithSkills()); From cfcd88f5c9eccd5145d158c7b28fd43ff1def40f Mon Sep 17 00:00:00 2001 From: yanjin <89272471+yanjin24@users.noreply.github.com> Date: Thu, 20 Aug 2026 15:21:57 +0800 Subject: [PATCH 2/3] [python] Fix jar URI construction in execution_environment.py (#1025) --- python/flink_agents/api/execution_environment.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/flink_agents/api/execution_environment.py b/python/flink_agents/api/execution_environment.py index e9ffac9a9..8750118fb 100644 --- a/python/flink_agents/api/execution_environment.py +++ b/python/flink_agents/api/execution_environment.py @@ -137,7 +137,7 @@ def get_execution_environment( if common_lib.is_dir(): for jar_file in common_lib.iterdir(): if jar_file.is_file() and str(jar_file).endswith(".jar"): - env.add_jars(f"file://{jar_file}") + env.add_jars(jar_file.resolve().as_uri()) else: err_msg = "Flink Agents common JAR not found." raise FileNotFoundError(err_msg) @@ -150,7 +150,7 @@ def get_execution_environment( if version_lib.is_dir(): for jar_file in version_lib.iterdir(): if jar_file.is_file() and str(jar_file).endswith(".jar"): - env.add_jars(f"file://{jar_file}") + env.add_jars(jar_file.resolve().as_uri()) else: err_msg = f"Flink Agents dist JAR for Flink {major_version} not found." raise FileNotFoundError(err_msg) From a6213bdb7f991544877ab77e00ddf865e9c9d633 Mon Sep 17 00:00:00 2001 From: Edson <113824357+Zhuoxi2000@users.noreply.github.com> Date: Thu, 20 Aug 2026 03:57:39 -0400 Subject: [PATCH 3/3] [integration][ollama] Tolerate tool schemas without a 'required' key (#1026) --- .../ollama/OllamaChatModelConnection.java | 8 +- .../ollama/OllamaChatModelConnectionTest.java | 104 ++++++++++++++++++ 2 files changed, 110 insertions(+), 2 deletions(-) create mode 100644 integrations/chat-models/ollama/src/test/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnectionTest.java diff --git a/integrations/chat-models/ollama/src/main/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnection.java b/integrations/chat-models/ollama/src/main/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnection.java index de6f946f9..e86f12779 100644 --- a/integrations/chat-models/ollama/src/main/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnection.java +++ b/integrations/chat-models/ollama/src/main/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnection.java @@ -106,8 +106,9 @@ public OllamaChatModelConnection(String endpoint, ResourceContext resourceContex * @return List of Ollama compatible tool specifications * @throws RuntimeException if schema parsing or conversion fails */ + // Package-visible for unit testing of the schema conversion. @SuppressWarnings("unchecked") - private List convertToOllamaTools(List tools) { + List convertToOllamaTools(List tools) { final ObjectMapper mapper = new ObjectMapper(); final List ollamaTools = new ArrayList<>(); try { @@ -118,7 +119,10 @@ private List convertToOllamaTools(List tools) { final Map> properties = (Map>) schema.get("properties"); - final List required = (List) schema.get("required"); + // "required" is optional in JSON Schema, and SchemaUtils only emits it when at + // least one parameter is required — treat a missing list as empty (#1014). + final List required = + (List) schema.getOrDefault("required", Collections.emptyList()); Map propertiesMap = new HashMap<>(); diff --git a/integrations/chat-models/ollama/src/test/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnectionTest.java b/integrations/chat-models/ollama/src/test/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnectionTest.java new file mode 100644 index 000000000..06570c29c --- /dev/null +++ b/integrations/chat-models/ollama/src/test/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnectionTest.java @@ -0,0 +1,104 @@ +/* + * 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.integrations.chatmodels.ollama; + +import io.github.ollama4j.tools.Tools; +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.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.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Unit tests for {@link OllamaChatModelConnection}'s tool-schema conversion — no network access. + */ +class OllamaChatModelConnectionTest { + + private static final ResourceContext NOOP = ResourceContext.fromGetResource((a, b) -> null); + + private static OllamaChatModelConnection connection() { + ResourceDescriptor desc = + ResourceDescriptor.Builder.newBuilder(OllamaChatModelConnection.class.getName()) + .addInitialArgument("endpoint", "http://localhost:11434") + .build(); + return new OllamaChatModelConnection(desc, NOOP); + } + + /** Minimal tool carrying only metadata; never invoked in these tests. */ + private static final class SchemaOnlyTool extends Tool { + SchemaOnlyTool(String inputSchema) { + super(new ToolMetadata("add", "Add two numbers.", inputSchema)); + } + + @Override + public ToolType getToolType() { + return ToolType.FUNCTION; + } + + @Override + public ToolResponse call(ToolParameters parameters) { + throw new UnsupportedOperationException("not invoked in this test"); + } + } + + @Test + @DisplayName("A schema without a 'required' key converts with every property optional") + void testSchemaWithoutRequiredKey() { + // SchemaUtils only emits "required" when at least one parameter is required, so an + // all-optional @Tool produces exactly this shape (#1014). + String schema = + "{\"type\":\"object\",\"properties\":{" + + "\"a\":{\"type\":\"integer\"},\"b\":{\"type\":\"integer\"}}}"; + + List converted = + connection().convertToOllamaTools(List.of(new SchemaOnlyTool(schema))); + + assertThat(converted).hasSize(1); + Tools.Tool tool = converted.get(0); + assertThat(tool.getToolSpec().getParameters().getProperties()) + .containsOnlyKeys("a", "b") + .allSatisfy((name, property) -> assertThat(property.isRequired()).isFalse()); + } + + @Test + @DisplayName("A schema with a 'required' key still marks the listed parameters required") + void testSchemaWithRequiredKey() { + String schema = + "{\"type\":\"object\",\"properties\":{" + + "\"a\":{\"type\":\"integer\"},\"b\":{\"type\":\"integer\"}}," + + "\"required\":[\"a\"]}"; + + List converted = + connection().convertToOllamaTools(List.of(new SchemaOnlyTool(schema))); + + assertThat(converted).hasSize(1); + Tools.Tool tool = converted.get(0); + assertThat(tool.getToolSpec().getParameters().getProperties().get("a").isRequired()) + .isTrue(); + assertThat(tool.getToolSpec().getParameters().getProperties().get("b").isRequired()) + .isFalse(); + } +}