This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new 3159c36f7 [runtime][python][docs] Fix execution guards and
observability gaps in Operations guidance (#1305)
3159c36f7 is described below
commit 3159c36f72a1ec6cbb10cd7f5fa4d151802603d1
Author: Wenjin Xie <[email protected]>
AuthorDate: Thu Oct 22 16:45:33 2026 +0800
[runtime][python][docs] Fix execution guards and observability gaps in
Operations guidance (#1305)
Generated-by: Codex CLI 0.153.4 (gpt-6-astra)
---
docs/content/docs/operations/deployment.md | 25 +++
docs/content/docs/operations/monitoring.md | 10 +-
python/flink_agents/api/core_options.py | 9 +
.../flink_agents/plan/tests/test_configuration.py | 37 ++++
.../apache/flink/agents/runtime/CompileUtils.java | 7 +-
.../apache/flink/agents/runtime/ResourceCache.java | 9 +
.../runtime/metrics/BuiltInExecutionMetrics.java | 37 +++-
.../agents/runtime/metrics/BuiltInMetrics.java | 7 +
.../runtime/operator/ActionExecutionOperator.java | 25 ++-
.../agents/runtime/tracelog/Slf4jTraceLogger.java | 1 +
.../flink/agents/runtime/CompileUtilsTest.java | 108 ++++++++++
.../agents/runtime/metrics/BuiltInMetricsTest.java | 24 ++-
.../runtime/subagent/ReActAgentSubagentTest.java | 176 +++++++++++++++--
.../tracelog/Slf4jTraceLoggerFailureTest.java | 217 +++++++++++++++++++++
14 files changed, 655 insertions(+), 37 deletions(-)
diff --git a/docs/content/docs/operations/deployment.md
b/docs/content/docs/operations/deployment.md
index 46aad7c3e..777296c07 100644
--- a/docs/content/docs/operations/deployment.md
+++ b/docs/content/docs/operations/deployment.md
@@ -47,6 +47,31 @@ Follow [Setup]({{< ref "docs/get-started/setup" >}}) to
install Flink Agents for
the Java Flink or PyFlink job. When the Agent uses Actions or Resources
implemented
in the other language, also install their dependencies.
+### Execution mode
+
+Flink's default execution mode is `STREAMING`. Keep this mode for Flink Agents
+jobs, including jobs with bounded inputs; no additional configuration is needed
+unless the default has been overridden. To explicitly select it in Flink's
+`config.yaml`, set the following outside the `agent` section:
+
+```yaml
+execution:
+ runtime-mode: STREAMING
+```
+
+Flink Agents does not support Flink's batch keyed-state backend, which is
enabled
+by default in `BATCH` mode through `execution.batch-state-backend.enabled`.
This
+backend clears state when switching keys, while an Agent may still have
unfinished
+Actions for a previous key. Flink Agents rejects this backend before processing
+Agent inputs.
+
+`AUTOMATIC` selects `BATCH` when all sources are bounded, so it can encounter
the
+same restriction. This also applies if the execution mode is changed after the
+Agent pipeline is built. Ensure the final job configuration uses `STREAMING`.
+
+Setting `execution.batch-state-backend.enabled: false` avoids this particular
+backend restriction, but does not imply full support for every BATCH use case.
+
### Submit to Flink Cluster
Submit the Java Flink or PyFlink job to the cluster. If the Agent is defined
in YAML,
diff --git a/docs/content/docs/operations/monitoring.md
b/docs/content/docs/operations/monitoring.md
index 93b762bc7..75f20ba7d 100644
--- a/docs/content/docs/operations/monitoring.md
+++ b/docs/content/docs/operations/monitoring.md
@@ -85,7 +85,7 @@ Flink Agents provides built-in metrics for input runs,
Events, Actions, executio
| **Action** | action.\<action_name\>.numOfActiveActionExecutions | Action
executions that have started and not yet finished. [3] | Gauge |
| **Action** | action.\<action_name\>.routingDecisionLatencyMs | Time to
resolve a routing decision when a `ChatRequestEvent` names a `MODEL_ROUTER`,
including replayed decisions; see [Model Routing]({{< ref
"docs/development/calling-llms/model_routing#observability" >}}). | Histogram |
| **Agent** | traceLogTruncatedRecords | Trace Log records whose attributes
were truncated at `STANDARD` level, counted once per record. | Counter |
-| **Agent** | traceLogWriteFailures | Failed Trace Log writes. Writes are best
effort and do not fail the job. | Counter |
+| **Agent** | traceLogWriteFailures | Detected Trace Log write failures.
Writes are best effort and do not fail the job. | Counter |
1. Input-run outcomes and latency samples are local to a task attempt. Input
runs already active when a task is restored produce neither outcome counts nor
latency samples, and Action executions already active when the task is restored
produce no latency sample. An input restored from the pending queue records its
outcome and processing latency once it starts in the new attempt, but not its
queue or end-to-end latency.
2. `numOfPendingInputEvents` counts buffered inputs and `numOfActiveInputRuns`
counts runs; a run waiting for an asynchronous continuation stays active. Both
gauges are rebuilt from Flink state after a restore.
@@ -115,7 +115,7 @@ These metrics describe LLM, Tool, Skill, MCP Server and
sub-agent calls that the
| **Subagent** |
action.\<action_name\>.subagent.\<agent_name\>.subagentCallLatencyMs | Time
from handing the prompt to the sub-agent until its result is observed. [4] |
Histogram |
1. One sample is one framework invocation of `ChatModel`. A framework retry
produces another outcome and latency sample; retries inside a provider client
are not observed. A response that is later rejected for its finish reason or
structured-output parsing still counts as a successful call.
-2. A Tool call fails when preparing or invoking it raises an exception, when
it returns `ToolResponse.error(...)`, or when its return value cannot be
converted to text. Tool names that are not registered resources are counted
under `tool=unknown`.
+2. A Tool call fails when preparing or invoking it raises an exception, when
it returns `ToolResponse.error(...)`, or when its return value cannot be
converted to text. Calls to Tools registered in an Agent, including an internal
sub-agent, are reported under their Tool names. Unregistered Tool names are
counted under `tool=unknown`. Within one operator, calls with the same Action
name and Tool name share metrics.
3. A `load_skill` call or an MCP Tool call is also counted in its Tool scope.
`numOfSkillLoads` increments even when the load fails, for example for a
missing Skill. Skill names that do not resolve are counted under
`skill=unknown`.
4. Only delegations the model requests through the reserved `_subagent_` name
are reported, under the Subagent scope keyed by the registered sub-agent name
and not also under a Tool scope. A `_subagent_` name that does not resolve is
counted under `tool=unknown`. For delegations dispatched together in one batch,
the latency includes their overlap.
5. A Tool call has a latency sample only when its callable started in the
current task attempt.
@@ -123,7 +123,9 @@ These metrics describe LLM, Tool, Skill, MCP Server and
sub-agent calls that the
#### Token Usage Metrics
-Token usage metrics are recorded when chat models are invoked through
`ChatModelConnection`, which helps track LLM API usage and cost. In the
built-in chat flow, token usage metrics appear under
`action.chat_model_action.model.<model_name>`, not under the Action that sends
`ChatRequestEvent`.
+Token usage metrics are recorded in the built-in [chat flow]({{< ref
"docs/development/calling-llms/chat" >}}), including calls through
`ChatRequestEvent` and `ctx.chat`, when the response provides a model name and
both prompt and completion token counts. If any of these values is missing,
neither counter is updated.
+
+These metrics appear under `action.chat_model_action.model.<model_name>`,
where `<model_name>` is the model name reported in the response.
| Scope | Metrics | Description | Type |
|-------|---------|-------------|------|
@@ -235,7 +237,7 @@ You can read the logs of a running job in the Flink Web UI:
Trace Log records Events, Action executions, and component calls as structured
JSON. By default, it records only Events at `STANDARD` level. Configure
[recording targets](#recording-targets) to select other entities or change
their level.
-Trace Log writes are best effort. A failed write does not fail Event
processing or the job; the first failure is logged at `WARN`, later ones at
`DEBUG`, and each one increments `traceLogWriteFailures`.
+Trace Log writes are best effort. A detected write failure does not fail Event
processing or the job; the first detected failure is logged at `WARN`, later
ones at `DEBUG`, and each one increments `traceLogWriteFailures`. Failures
swallowed by the logging backend cannot be reflected in this metric.
### Recording Targets
diff --git a/python/flink_agents/api/core_options.py
b/python/flink_agents/api/core_options.py
index 6df03e061..aa1db965e 100644
--- a/python/flink_agents/api/core_options.py
+++ b/python/flink_agents/api/core_options.py
@@ -49,6 +49,15 @@ class LoggerType(Enum):
SLF4J = "slf4j"
FILE = "file"
+ @classmethod
+ def _missing_(cls, value: object) -> "LoggerType | None":
+ """Accept case-insensitive configuration values, matching Java."""
+ if isinstance(value, str):
+ for member in cls:
+ if member.value.casefold() == value.casefold():
+ return member
+ return None
+
class ConditionEvaluationFailureStrategy(Enum):
"""Behavior when evaluating an action trigger condition fails.
diff --git a/python/flink_agents/plan/tests/test_configuration.py
b/python/flink_agents/plan/tests/test_configuration.py
index db5b44d72..d2773f110 100644
--- a/python/flink_agents/plan/tests/test_configuration.py
+++ b/python/flink_agents/plan/tests/test_configuration.py
@@ -259,6 +259,43 @@ def test_get_with_null_and_default() -> None:
assert config.get(nullable_str) == "default"
[email protected](
+ ("value", "expected"),
+ [
+ ("FILE", LoggerType.FILE),
+ ("file", LoggerType.FILE),
+ ("File", LoggerType.FILE),
+ ("SLF4J", LoggerType.SLF4J),
+ ("slf4j", LoggerType.SLF4J),
+ ("Slf4j", LoggerType.SLF4J),
+ ],
+)
+def test_trace_log_output_type_from_yaml(
+ tmp_path: Path, value: str, expected: LoggerType
+) -> None:
+ config_file = tmp_path / "agent.yaml"
+ config_file.write_text(
+ yaml.safe_dump({"agent": {"trace-log": {"output-type": value}}}),
+ encoding="utf-8",
+ )
+ config = AgentConfiguration()
+ config.load_from_file(str(config_file))
+ assert config.get(AgentConfigOptions.TRACE_LOG_OUTPUT_TYPE) is expected
+
+ config.set(AgentConfigOptions.TRACE_LOG_OUTPUT_TYPE, expected)
+ serialized = config.model_dump_json()
+ assert json.loads(serialized)["conf_data"]["trace-log.output-type"] ==
value.lower()
+ restored = AgentConfiguration.model_validate_json(serialized)
+ assert restored.get(AgentConfigOptions.TRACE_LOG_OUTPUT_TYPE) is expected
+
+
[email protected]("value", ["unknown", "", 123])
+def test_trace_log_output_type_rejects_invalid_values(value: object) -> None:
+ config = AgentConfiguration({"trace-log.output-type": value})
+ with pytest.raises(ValueError, match="Invalid value for
trace-log.output-type"):
+ config.get(AgentConfigOptions.TRACE_LOG_OUTPUT_TYPE)
+
+
def test_trace_log_defaults_and_empty_targets_survive_json() -> None:
config = AgentConfiguration()
assert config.get(AgentConfigOptions.TRACE_LOG_TARGETS) == [
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/CompileUtils.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/CompileUtils.java
index d8e123ca1..2d64f2f09 100644
--- a/runtime/src/main/java/org/apache/flink/agents/runtime/CompileUtils.java
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/CompileUtils.java
@@ -130,10 +130,9 @@ public class CompileUtils {
* {@code FINISHED}; only the affected records go missing. See
* https://github.com/apache/flink-agents/issues/939.
*
- * <p>This only catches an explicitly configured {@code BATCH} runtime
mode. {@code
- * RuntimeExecutionMode.AUTOMATIC} (the default) that later resolves to
batch execution because
- * every source happens to be bounded is not detected here, since that
resolution has not
- * happened yet at graph-construction time.
+ * <p>This early check only catches an explicitly configured {@code BATCH}
runtime mode. The
+ * operator also validates the actual keyed-state backend during
initialization, covering {@code
+ * AUTOMATIC} with bounded sources and configuration changes after this
check.
*/
private static void checkBatchStateBackendCompatibility(KeyedStream<?, ?>
keyedInputStream) {
ReadableConfig configuration =
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/ResourceCache.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/ResourceCache.java
index fbe0271aa..41becd5e6 100644
--- a/runtime/src/main/java/org/apache/flink/agents/runtime/ResourceCache.java
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/ResourceCache.java
@@ -212,6 +212,15 @@ public class ResourceCache implements AutoCloseable {
return providers != null && providers.containsKey(name);
}
+ /**
+ * Checks registration along the same fallback chain as {@link
#getResource}, without creating
+ * resources.
+ */
+ public boolean hasResourceInHierarchy(String name, ResourceType type) {
+ return hasResource(name, type)
+ || (parent != null && parent.hasResourceInHierarchy(name,
type));
+ }
+
/**
* Resolves a resource by name and type, creating it from its provider if
not cached.
*
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/metrics/BuiltInExecutionMetrics.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/metrics/BuiltInExecutionMetrics.java
index af4e135ee..d3d996b97 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/metrics/BuiltInExecutionMetrics.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/metrics/BuiltInExecutionMetrics.java
@@ -20,6 +20,7 @@ package org.apache.flink.agents.runtime.metrics;
import org.apache.flink.agents.api.trace.TraceContext;
import org.apache.flink.agents.api.trace.TraceRecord;
+import org.apache.flink.agents.runtime.lifecycle.ComponentExecutionListener;
import java.time.Duration;
import java.time.Instant;
@@ -38,24 +39,40 @@ final class BuiltInExecutionMetrics {
BuiltInExecutionMetrics(
FlinkAgentsMetricGroupImpl agentMetricGroup, Predicate<String>
isRegisteredTool) {
this.agentMetricGroup = agentMetricGroup;
+ this.metricRecordersByEntityType = createRecorders(isRegisteredTool);
+ }
+
+ private static Map<String, ExecutionMetricRecorder> createRecorders(
+ Predicate<String> isRegisteredTool) {
ExecutionMetricRecorder llmMetricRecorder = new
LlmExecutionMetricRecorder();
ExecutionMetricRecorder toolMetricRecorder =
new ToolExecutionMetricRecorder(isRegisteredTool);
ExecutionMetricRecorder subagentMetricRecorder = new
SubagentExecutionMetricRecorder();
- this.metricRecordersByEntityType =
- Map.of(
- llmMetricRecorder.entityType(),
- llmMetricRecorder,
- toolMetricRecorder.entityType(),
- toolMetricRecorder,
- subagentMetricRecorder.entityType(),
- subagentMetricRecorder);
+ return Map.of(
+ llmMetricRecorder.entityType(),
+ llmMetricRecorder,
+ toolMetricRecorder.entityType(),
+ toolMetricRecorder,
+ subagentMetricRecorder.entityType(),
+ subagentMetricRecorder);
}
void executionRecordObserved(String actionName, TraceRecord record) {
+ executionRecordObserved(actionName, record,
metricRecordersByEntityType);
+ }
+
+ ComponentExecutionListener createListener(
+ String actionName, Predicate<String> isRegisteredTool) {
+ Map<String, ExecutionMetricRecorder> recorders =
createRecorders(isRegisteredTool);
+ // Recorders retain this Action's resource lookup; timing state
remains shared so the
+ // normal Action terminal record can clean up unfinished calls from
every scope.
+ return record -> executionRecordObserved(actionName, record,
recorders);
+ }
+
+ private void executionRecordObserved(
+ String actionName, TraceRecord record, Map<String,
ExecutionMetricRecorder> recorders) {
TraceContext traceContext = record.getContext();
- ExecutionMetricRecorder recorder =
- metricRecordersByEntityType.get(traceContext.getEntityType());
+ ExecutionMetricRecorder recorder =
recorders.get(traceContext.getEntityType());
if (isBlank(actionName) || recorder == null) {
return;
}
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 4c3ec5645..e85331633 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
@@ -24,6 +24,7 @@ import org.apache.flink.agents.api.trace.ExecutionReporter;
import org.apache.flink.agents.api.trace.TraceContext;
import org.apache.flink.agents.api.trace.TraceRecord;
import org.apache.flink.agents.plan.AgentPlan;
+import org.apache.flink.agents.runtime.lifecycle.ComponentExecutionListener;
import org.apache.flink.metrics.Counter;
import org.apache.flink.metrics.Meter;
@@ -153,6 +154,12 @@ public class BuiltInMetrics {
}
}
+ /** Creates a component listener bound to the resources of one Action
execution. */
+ public ComponentExecutionListener createComponentMetricListener(
+ String actionName, Predicate<String> isRegisteredTool) {
+ return executionMetrics.createListener(actionName, isRegisteredTool);
+ }
+
/** Returns the counter tracking Trace Log records with truncated
attributes. */
public Counter getTraceLogTruncatedRecordsCounter() {
return traceLogTruncatedRecords;
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 471da6ebc..5ea003abd 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
@@ -79,6 +79,7 @@ import
org.apache.flink.streaming.api.operators.ChainingStrategy;
import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
import org.apache.flink.streaming.api.operators.OperatorSnapshotFutures;
import org.apache.flink.streaming.api.operators.Output;
+import
org.apache.flink.streaming.api.operators.sorted.state.BatchExecutionKeyedStateBackend;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.runtime.tasks.ProcessingTimeService;
@@ -1174,6 +1175,15 @@ public class ActionExecutionOperator<IN, OUT> extends
AbstractStreamOperator<OUT
public void initializeState(StateInitializationContext context) throws
Exception {
super.initializeState(context);
+ // Check the actual backend after Flink resolves AUTOMATIC and any
configuration changes
+ // made after the Agent pipeline was built, before opening resources
or recovering Actions.
+ checkState(
+ !(getKeyedStateBackend() instanceof
BatchExecutionKeyedStateBackend),
+ "Flink Agents does not support Flink's batch keyed-state
backend"
+ + " (execution.batch-state-backend.enabled): switching
keys can discard"
+ + " unfinished Actions. Run this job with
RuntimeExecutionMode.STREAMING,"
+ + " or set execution.batch-state-backend.enabled to
false.");
+
int maxParallelism =
getRuntimeContext().getTaskInfo().getMaxNumberOfParallelSubtasks();
durableExecManager.maybeInitActionStateStore(
agentPlan.getConfig(), maxParallelism,
getActionStateKeySerializer());
@@ -1381,8 +1391,21 @@ public class ActionExecutionOperator<IN, OUT> extends
AbstractStreamOperator<OUT
* built-in writer and metrics first, followed by registered listeners.
*/
private ExecutionReportingContext
createExecutionReportingContext(ActionTask actionTask) {
+ // Bind to this execution's resources: the shared RunnerContext may
have switched to
+ // another Agent by the time an asynchronous call reports its outcome.
+ RunnerContextImpl.ResourceScope scope =
contextManager.getResourceScope(actionTask);
+ ResourceCache executionResources =
+ scope == null ? resourceCache : scope.getChildResourceCache();
+ ComponentExecutionListener metricListener =
+ builtInMetrics.createComponentMetricListener(
+ actionTask.getAction().getName(),
+ name ->
executionResources.hasResourceInHierarchy(name, ResourceType.TOOL));
List<ComponentExecutionListener> listeners = new ArrayList<>();
- listeners.add(record ->
observeExecutionRecord(actionTask.getAction().getName(), record));
+ listeners.add(
+ record -> {
+ traceLogWriter.appendAndFlush(record);
+ metricListener.onExecutionReported(record);
+ });
listeners.addAll(componentExecutionListeners);
return new ExecutionReportingContext(actionTask.getTraceContext(),
listeners);
}
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/tracelog/Slf4jTraceLogger.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/tracelog/Slf4jTraceLogger.java
index 2c8f68ace..c3716faaa 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/tracelog/Slf4jTraceLogger.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/tracelog/Slf4jTraceLogger.java
@@ -255,6 +255,7 @@ public class Slf4jTraceLogger implements TraceLogger {
.withFileName(traceLogFile)
.withAppend(true)
.setLayout(layout)
+ .setIgnoreExceptions(false)
.build();
appender.start();
configuration.addAppender(appender);
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/CompileUtilsTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/CompileUtilsTest.java
index c21cc6508..ba25351cf 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/CompileUtilsTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/CompileUtilsTest.java
@@ -21,11 +21,15 @@ import com.fasterxml.jackson.core.JsonProcessingException;
import org.apache.flink.agents.plan.AgentPlan;
import org.apache.flink.agents.runtime.operator.ActionExecutionOperatorTest;
import org.apache.flink.api.common.RuntimeExecutionMode;
+import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.common.typeinfo.Types;
+import org.apache.flink.api.connector.source.Boundedness;
+import org.apache.flink.api.connector.source.lib.NumberSequenceSource;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.configuration.ExecutionOptions;
+import org.apache.flink.configuration.RestartStrategyOptions;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.KeyedStream;
@@ -35,6 +39,10 @@ import org.apache.flink.types.Row;
import org.apache.flink.util.CloseableIterator;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+import org.junit.jupiter.params.provider.EnumSource;
import java.util.ArrayList;
import java.util.Collections;
@@ -137,6 +145,106 @@ public class CompileUtilsTest {
.doesNotThrowAnyException();
}
+ @ParameterizedTest
+ @CsvSource({"AUTOMATIC, false", "AUTOMATIC, true", "BATCH, true"})
+ @Timeout(60)
+ void rejectsBatchStateBackendAtRuntime(
+ RuntimeExecutionMode mode, boolean configureAfterConnecting)
throws Exception {
+ Configuration conf = new Configuration();
+ conf.set(RestartStrategyOptions.RESTART_STRATEGY, "none");
+ conf.set(
+ ExecutionOptions.RUNTIME_MODE,
+ configureAfterConnecting ? RuntimeExecutionMode.STREAMING :
mode);
+ try (StreamExecutionEnvironment env =
+ StreamExecutionEnvironment.getExecutionEnvironment(conf)) {
+ env.setParallelism(1);
+ DataStream<Object> output =
+ CompileUtils.connectToAgent(
+ env.fromData(0L, 1L, 2L, 3L, 4L).keyBy(value ->
value),
+ TEST_AGENT_PLAN);
+ if (configureAfterConnecting) {
+ env.setRuntimeMode(mode);
+ }
+
+ // Flink resolves AUTOMATIC only when generating the job graph.
The guard must also
+ // catch configuration changes after connectToAgent's early
validation.
+ assertThatThrownBy(
+ () -> {
+ try (CloseableIterator<Object> iterator =
+ output.executeAndCollect()) {
+ iterator.forEachRemaining(ignored -> {});
+ }
+ })
+ .hasStackTraceContaining(
+ "Flink Agents does not support Flink's batch
keyed-state backend")
+
.hasStackTraceContaining("execution.batch-state-backend.enabled")
+ .hasStackTraceContaining("RuntimeExecutionMode.STREAMING");
+ }
+ }
+
+ @ParameterizedTest
+ @EnumSource(RuntimeExecutionMode.class)
+ @Timeout(60)
+ void processesAllKeysWithCompatibleStateBackend(RuntimeExecutionMode mode)
throws Exception {
+ Configuration conf = new Configuration();
+ conf.set(ExecutionOptions.RUNTIME_MODE, mode);
+ conf.set(RestartStrategyOptions.RESTART_STRATEGY, "none");
+ // In STREAMING mode the default option must remain harmless. In batch
execution,
+ // disabling the batch-specific backend must preserve pending work
across key switches.
+ conf.set(ExecutionOptions.USE_BATCH_STATE_BACKEND, mode ==
RuntimeExecutionMode.STREAMING);
+ try (StreamExecutionEnvironment env =
+ StreamExecutionEnvironment.getExecutionEnvironment(conf)) {
+ env.setParallelism(1);
+ DataStream<Object> output =
+ CompileUtils.connectToAgent(
+ env.fromData(0L, 1L, 2L, 3L, 4L).keyBy(value ->
value),
+ TEST_AGENT_PLAN);
+ List<Object> results = new ArrayList<>();
+ try (CloseableIterator<Object> iterator =
output.executeAndCollect()) {
+ iterator.forEachRemaining(results::add);
+ }
+ assertThat(results).containsExactlyInAnyOrder(2L, 4L, 6L, 8L, 10L);
+ }
+ }
+
+ @Test
+ @Timeout(60)
+ void allowsAutomaticModeWithUnboundedSource() throws Exception {
+ Configuration conf = new Configuration();
+ conf.set(ExecutionOptions.RUNTIME_MODE,
RuntimeExecutionMode.AUTOMATIC);
+ conf.set(RestartStrategyOptions.RESTART_STRATEGY, "none");
+ try (StreamExecutionEnvironment env =
+ StreamExecutionEnvironment.getExecutionEnvironment(conf)) {
+ env.setParallelism(1);
+ DataStream<Object> output =
+ CompileUtils.connectToAgent(
+ env.fromSource(
+ new UnboundedSequenceSource(),
+ WatermarkStrategy.noWatermarks(),
+ "unbounded-input")
+ .keyBy(value -> value),
+ TEST_AGENT_PLAN);
+ List<Object> results = new ArrayList<>();
+ try (CloseableIterator<Object> iterator =
output.executeAndCollect()) {
+ iterator.forEachRemaining(results::add);
+ }
+ assertThat(results).containsExactlyInAnyOrder(2L, 4L, 6L, 8L, 10L);
+ }
+ }
+
+ private static class UnboundedSequenceSource extends NumberSequenceSource {
+ UnboundedSequenceSource() {
+ super(0, 4);
+ }
+
+ @Override
+ public Boundedness getBoundedness() {
+ // Report unbounded input to exercise AUTOMATIC's streaming
selection, but finish
+ // after five records so the test can verify every result without
cancelling the job.
+ return Boundedness.CONTINUOUS_UNBOUNDED;
+ }
+ }
+
@Test
void detectsPickledAndExplicitByteArrayPythonKeyTypes() {
assertThat(
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/metrics/BuiltInMetricsTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/metrics/BuiltInMetricsTest.java
index 590f7f0f2..676f901ba 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/metrics/BuiltInMetricsTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/metrics/BuiltInMetricsTest.java
@@ -23,9 +23,12 @@ import org.apache.flink.agents.api.trace.TraceContext;
import org.apache.flink.agents.api.trace.TraceRecord;
import org.apache.flink.agents.api.trace.TraceRecords;
import org.apache.flink.agents.plan.AgentPlan;
+import org.apache.flink.agents.runtime.lifecycle.ComponentExecutionListener;
import org.apache.flink.metrics.MetricGroup;
import org.apache.flink.runtime.metrics.groups.UnregisteredMetricGroups;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
import java.util.Map;
import java.util.UUID;
@@ -95,8 +98,9 @@ class BuiltInMetricsTest {
assertThat(actionMetricGroup.getCounter("numOfActionsExecuted").getCount()).isEqualTo(1L);
}
- @Test
- void actionTerminalDropsOnlyItsUnpairedChildLatencyState() {
+ @ParameterizedTest
+ @ValueSource(booleans = {false, true})
+ void actionTerminalDropsOnlyItsUnpairedChildLatencyState(boolean
scopedListeners) {
MetricGroup parentMetricGroup =
UnregisteredMetricGroups.createUnregisteredOperatorMetricGroup();
FlinkAgentsMetricGroupImpl metricGroup = new
FlinkAgentsMetricGroupImpl(parentMetricGroup);
@@ -121,11 +125,19 @@ class BuiltInMetricsTest {
Map.of());
metrics.restoreActionTask(completedAction, false);
- metrics.markExecutionRecord("restored_action",
TraceRecords.started(completedActionLlm));
- metrics.markExecutionRecord("restored_action",
TraceRecords.started(activeActionLlm));
+ ComponentExecutionListener completedListener =
+ scopedListeners
+ ?
metrics.createComponentMetricListener("restored_action", ignored -> true)
+ : record ->
metrics.markExecutionRecord("restored_action", record);
+ ComponentExecutionListener activeListener =
+ scopedListeners
+ ?
metrics.createComponentMetricListener("restored_action", ignored -> false)
+ : record ->
metrics.markExecutionRecord("restored_action", record);
+
completedListener.onExecutionReported(TraceRecords.started(completedActionLlm));
+
activeListener.onExecutionReported(TraceRecords.started(activeActionLlm));
metrics.markExecutionRecord("restored_action",
TraceRecords.succeeded(completedAction));
- metrics.markExecutionRecord("restored_action",
TraceRecords.succeeded(completedActionLlm));
- metrics.markExecutionRecord("restored_action",
TraceRecords.succeeded(activeActionLlm));
+
completedListener.onExecutionReported(TraceRecords.succeeded(completedActionLlm));
+
activeListener.onExecutionReported(TraceRecords.succeeded(activeActionLlm));
assertThat(
metricGroup
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/ReActAgentSubagentTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/ReActAgentSubagentTest.java
index f27e0680f..ea644893c 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/ReActAgentSubagentTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/ReActAgentSubagentTest.java
@@ -41,6 +41,7 @@ import org.apache.flink.agents.api.subagent.SubagentResult;
import org.apache.flink.agents.api.subagent.SubagentSetup;
import org.apache.flink.agents.api.tools.Tool;
import org.apache.flink.agents.plan.AgentPlan;
+import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl;
import org.apache.flink.agents.runtime.operator.ActionExecutionOperator;
import org.apache.flink.agents.runtime.operator.ActionExecutionOperatorFactory;
import org.apache.flink.api.common.typeinfo.TypeInformation;
@@ -52,9 +53,12 @@ import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.Timeout;
import org.junit.jupiter.api.condition.EnabledForJreRange;
import org.junit.jupiter.api.condition.JRE;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
import javax.annotation.Nullable;
+import java.lang.reflect.Field;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Deque;
@@ -62,6 +66,7 @@ import java.util.HashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;
+import java.util.function.Consumer;
import java.util.stream.Collectors;
import static org.assertj.core.api.Assertions.assertThat;
@@ -243,9 +248,10 @@ public class ReActAgentSubagentTest {
.anySatisfy(text -> assertThat(text).contains("reviewed: ok"));
}
- @Test
+ @ParameterizedTest
+ @ValueSource(booleans = {false, true})
@Timeout(60)
- void toolCallLoopRunsInsideTheChildScope() throws Exception {
+ void toolCallLoopRunsInsideTheChildScope(boolean inheritTool) throws
Exception {
ToolCallBlock lookup = new ToolCallBlock("call-1", "shout",
Map.of("text", "hello world"));
ChildScriptedConnection.REPLIES.add(assistantWithToolCall("let me
check", lookup));
ChildScriptedConnection.REPLIES.add(
@@ -253,13 +259,28 @@ public class ReActAgentSubagentTest {
Agent parent = new Agent();
registerConnections(parent);
- parent.addResource("child", ResourceType.AGENT,
toolCallingReactAgent());
+ parent.addResource("child", ResourceType.AGENT,
toolCallingReactAgent(!inheritTool));
+ if (inheritTool) {
+ parent.addResource(
+ "shout",
+ ResourceType.TOOL,
+
Tool.fromMethod(ReActAgentSubagentTest.class.getMethod("shout", String.class)));
+ }
parent.addAction(
new String[] {InputEvent.EVENT_TYPE},
ReActAgentSubagentTest.class.getMethod(
"callReActChild", Event.class, RunnerContext.class));
- List<Object> output = run(parent, 7L);
+ List<Object> output =
+ run(
+ parent,
+ 7L,
+ metrics -> {
+ assertThat(toolCounter(metrics, "shout",
"numOfToolCallsSucceeded"))
+ .isEqualTo(1);
+ assertThat(toolCounter(metrics, "unknown",
"numOfToolCallsSucceeded"))
+ .isZero();
+ });
// The tool schema was offered on the first round, the tool executed
inside the child
// scope, and the second round saw the tool result before answering.
@@ -275,6 +296,110 @@ public class ReActAgentSubagentTest {
});
}
+ @Test
+ @Timeout(60)
+ void parentAndChildToolsWithTheSameNameShareMetrics() throws Exception {
+ ReActAgent parent =
+ delegatingParent(null, null, toolCallingReactAgent(true),
List.of("shout"));
+ parent.addResource(
+ "shout",
+ ResourceType.TOOL,
+ Tool.fromMethod(
+ ReActAgentSubagentTest.class.getMethod("parentShout",
String.class)));
+ ParentScriptedConnection.REPLIES.add(
+ assistantWithToolCall(
+ "parent tool",
+ new ToolCallBlock("root-tool", "shout", Map.of("text",
"parent"))));
+ ParentScriptedConnection.REPLIES.add(assistantWithToolCall("delegate",
delegation()));
+ ParentScriptedConnection.REPLIES.add(ChatMessage.assistant("done"));
+ ChildScriptedConnection.REPLIES.add(
+ assistantWithToolCall(
+ "child tool",
+ new ToolCallBlock("child-tool", "shout",
Map.of("text", "child"))));
+ ChildScriptedConnection.REPLIES.add(ChatMessage.assistant("child
done"));
+
+ assertThat(
+ run(
+ parent,
+ 7L,
+ metrics -> {
+ assertThat(
+ toolCounter(
+ metrics,
+ "shout",
+
"numOfToolCallsSucceeded"))
+ .isEqualTo(2);
+ assertThat(
+ toolCounter(
+ metrics,
+ "unknown",
+
"numOfToolCallsSucceeded"))
+ .isZero();
+ assertThat(
+ metrics.getSubGroup(
+ "action",
"tool_call_action")
+
.getSubGroup("tool", "shout")
+
.getHistogram("toolCallLatencyMs")
+ .getCount())
+ .isEqualTo(2);
+ }))
+ .containsExactly("done");
+ assertThat(allTexts(ParentScriptedConnection.CONVERSATIONS.get(1)))
+ .anySatisfy(text ->
assertThat(text).contains("parent:parent"));
+ assertThat(allTexts(ChildScriptedConnection.CONVERSATIONS.get(1)))
+ .anySatisfy(text -> assertThat(text).contains("CHILD!"));
+ }
+
+ @Test
+ @Timeout(60)
+ void missingToolsRemainUnknownInTheCallingScope() throws Exception {
+ ReActAgent parent = delegatingParent(null, null,
toolCallingReactAgent(true), List.of());
+ // A Tool registered only in the child must not make the parent's
failed call look
+ // registered.
+ ParentScriptedConnection.REPLIES.add(
+ assistantWithToolCall(
+ "missing in parent",
+ new ToolCallBlock("root-tool", "shout", Map.of("text",
"parent"))));
+ ParentScriptedConnection.REPLIES.add(assistantWithToolCall("delegate",
delegation()));
+ ParentScriptedConnection.REPLIES.add(ChatMessage.assistant("done"));
+ ChildScriptedConnection.REPLIES.add(
+ assistantWithToolCall(
+ "missing everywhere",
+ new ToolCallBlock("child-tool", "missing", Map.of())));
+ ChildScriptedConnection.REPLIES.add(ChatMessage.assistant("child
done"));
+
+ assertThat(
+ run(
+ parent,
+ 7L,
+ metrics -> {
+ assertThat(
+ toolCounter(
+ metrics,
+ "unknown",
+
"numOfToolCallsFailed"))
+ .isEqualTo(2);
+ assertThat(
+ toolCounter(
+ metrics,
+ "shout",
+
"numOfToolCallsFailed"))
+ .isZero();
+ assertThat(
+ toolCounter(
+ metrics,
+ "missing",
+
"numOfToolCallsFailed"))
+ .isZero();
+ }))
+ .containsExactly("done");
+ }
+
+ @SuppressWarnings("unused")
+ public static String parentShout(@ToolParam(name = "text") String text) {
+ return "parent:" + text;
+ }
+
@Test
@Timeout(60)
void bareChildDoesNotInheritTheParentPrompt() throws Exception {
@@ -343,17 +468,25 @@ public class ReActAgentSubagentTest {
/** A ReActAgent parent configured as given, declaring a bare ReActAgent
child to its model. */
private static ReActAgent delegatingParent(
@Nullable Prompt prompt, @Nullable Class<?> outputSchema) {
+ return delegatingParent(prompt, outputSchema,
bareReactAgent("child_conn"), List.of());
+ }
+
+ private static ReActAgent delegatingParent(
+ @Nullable Prompt prompt,
+ @Nullable Class<?> outputSchema,
+ Agent child,
+ List<String> tools) {
ReActAgent parent =
new ReActAgent(
ResourceDescriptor.Builder.newBuilder(ScriptedModelSetup.class.getName())
.addInitialArgument("connection",
"parent_conn")
.addInitialArgument("subagents",
List.of("child"))
+ .addInitialArgument("tools", tools)
.build(),
prompt,
outputSchema);
registerConnections(parent);
- parent.addResource(
- "child", ResourceType.AGENT,
declaredChild(bareReactAgent("child_conn")));
+ parent.addResource("child", ResourceType.AGENT, declaredChild(child));
return parent;
}
@@ -371,7 +504,7 @@ public class ReActAgentSubagentTest {
.build();
}
- private static ReActAgent toolCallingReactAgent() throws Exception {
+ private static ReActAgent toolCallingReactAgent(boolean registerTool)
throws Exception {
ReActAgent child =
new ReActAgent(
ResourceDescriptor.Builder.newBuilder(ScriptedModelSetup.class.getName())
@@ -380,10 +513,12 @@ public class ReActAgentSubagentTest {
.build(),
null,
null);
- child.addResource(
- "shout",
- ResourceType.TOOL,
-
Tool.fromMethod(ReActAgentSubagentTest.class.getMethod("shout", String.class)));
+ if (registerTool) {
+ child.addResource(
+ "shout",
+ ResourceType.TOOL,
+
Tool.fromMethod(ReActAgentSubagentTest.class.getMethod("shout", String.class)));
+ }
return child;
}
@@ -440,8 +575,22 @@ public class ReActAgentSubagentTest {
.build());
}
- @SuppressWarnings("unchecked")
private static List<Object> run(Agent root, long input) throws Exception {
+ return run(root, input, metrics -> {});
+ }
+
+ private static long toolCounter(
+ FlinkAgentsMetricGroupImpl metrics, String tool, String counter) {
+ return metrics.getSubGroup("action", "tool_call_action")
+ .getSubGroup("tool", tool)
+ .getCounter(counter)
+ .getCount();
+ }
+
+ @SuppressWarnings("unchecked")
+ private static List<Object> run(
+ Agent root, long input, Consumer<FlinkAgentsMetricGroupImpl>
checkMetrics)
+ throws Exception {
try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
harness =
new KeyedOneInputStreamOperatorTestHarness<>(
new ActionExecutionOperatorFactory<>(new
AgentPlan(root), true),
@@ -451,6 +600,9 @@ public class ReActAgentSubagentTest {
harness.processElement(new StreamRecord<>(input));
((ActionExecutionOperator<Long, Object>) harness.getOperator())
.waitInFlightEventsFinished();
+ Field field =
ActionExecutionOperator.class.getDeclaredField("metricGroup");
+ field.setAccessible(true);
+ checkMetrics.accept((FlinkAgentsMetricGroupImpl)
field.get(harness.getOperator()));
return ((List<StreamRecord<Object>>) harness.getRecordOutput())
.stream().map(StreamRecord::getValue).collect(Collectors.toList());
}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/tracelog/Slf4jTraceLoggerFailureTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/tracelog/Slf4jTraceLoggerFailureTest.java
new file mode 100644
index 000000000..9be459ed7
--- /dev/null
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/tracelog/Slf4jTraceLoggerFailureTest.java
@@ -0,0 +1,217 @@
+/*
+ * 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.tracelog;
+
+import org.apache.flink.agents.api.InputEvent;
+import org.apache.flink.agents.api.logger.LoggerType;
+import org.apache.flink.agents.api.logger.TraceLoggerConfig;
+import org.apache.flink.agents.api.trace.TraceContext;
+import org.apache.flink.agents.api.trace.TraceRecord;
+import org.apache.flink.agents.runtime.metrics.BuiltInMetrics;
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.api.common.JobInfo;
+import org.apache.flink.api.common.TaskInfo;
+import org.apache.flink.metrics.SimpleCounter;
+import org.apache.flink.streaming.api.operators.StreamingRuntimeContext;
+import org.apache.logging.log4j.Level;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.core.Appender;
+import org.apache.logging.log4j.core.LogEvent;
+import org.apache.logging.log4j.core.LoggerContext;
+import org.apache.logging.log4j.core.appender.AbstractAppender;
+import org.apache.logging.log4j.core.appender.FileAppender;
+import org.apache.logging.log4j.core.appender.OutputStreamManager;
+import org.apache.logging.log4j.core.config.Configuration;
+import org.apache.logging.log4j.core.config.LoggerConfig;
+import org.apache.logging.log4j.core.layout.PatternLayout;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+import java.io.FilterOutputStream;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.lang.reflect.Method;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.List;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+/** Exercises write failures through the automatically configured Log4j2 file
appender. */
+class Slf4jTraceLoggerFailureTest {
+
+ @TempDir private Path tempDir;
+
+ @Test
+ void fileWriteFailuresAreCountedWithoutEscapingAndWarnOnlyOnce() throws
Exception {
+ LoggerContext context = (LoggerContext) LogManager.getContext(false);
+ Configuration configuration = context.getConfiguration();
+ String traceName = Slf4jTraceLogger.TRACE_LOGGER_NAME;
+ String writerName = TraceLogWriter.class.getName();
+ String appenderName = "FlinkAgentsTraceLogAppender";
+ LoggerConfig previousTrace = configuration.getLoggers().get(traceName);
+ LoggerConfig previousWriter =
configuration.getLoggers().get(writerName);
+ Appender previousAppender = configuration.getAppender(appenderName);
+ String previousLogFile = System.getProperty("log.file");
+ List<LogEvent> diagnostics = new ArrayList<>();
+ AbstractAppender diagnosticAppender =
+ new AbstractAppender(
+ "TraceWriteFailureDiagnostics",
+ null,
+ PatternLayout.createDefaultLayout(),
+ true,
+ null) {
+ @Override
+ public void append(LogEvent event) {
+ diagnostics.add(event.toImmutable());
+ }
+ };
+ FileAppender fileAppender = null;
+ try {
+ // Remove a prior test's logger so open() takes the production
auto-configuration path.
+ configuration.removeLogger(traceName);
+ configuration.removeLogger(writerName);
+ configuration.getAppenders().remove(appenderName);
+ diagnosticAppender.start();
+ LoggerConfig writerConfig = new LoggerConfig(writerName,
Level.DEBUG, false);
+ writerConfig.addAppender(diagnosticAppender, Level.DEBUG, null);
+ configuration.addLogger(writerName, writerConfig);
+ context.updateLoggers();
+ Path logFile = tempDir.resolve("taskmanager.log");
+ System.setProperty("log.file", logFile.toString());
+
+ StreamingRuntimeContext runtime =
mock(StreamingRuntimeContext.class);
+ JobInfo job = mock(JobInfo.class);
+ TaskInfo task = mock(TaskInfo.class);
+ when(runtime.getJobInfo()).thenReturn(job);
+ when(runtime.getTaskInfo()).thenReturn(task);
+ when(job.getJobId()).thenReturn(JobID.generate());
+ when(task.getTaskName()).thenReturn("agent");
+ BuiltInMetrics metrics = mock(BuiltInMetrics.class);
+ SimpleCounter failures = new SimpleCounter();
+
when(metrics.getTraceLogWriteFailuresCounter()).thenReturn(failures);
+ when(metrics.getTraceLogTruncatedRecordsCounter()).thenReturn(new
SimpleCounter());
+ try (TraceLogWriter writer =
+ TraceLogWriter.forTraceLogger(
+ new Slf4jTraceLogger(
+ TraceLoggerConfig.builder()
+ .loggerType(LoggerType.SLF4J)
+ .build()))) {
+ writer.open(runtime, metrics);
+ fileAppender = (FileAppender)
configuration.getAppender(appenderName);
+ assertThat(fileAppender).isNotNull();
+ writer.appendAndFlush(inputRecord("before-failure"));
+ assertThat(failures.getCount()).isZero();
+ Path traceFile = Path.of(logFile + ".trace-log.log");
+
assertThat(Files.readString(traceFile)).contains("before-failure");
+
+ // Replace only the destination stream; retain the real
logger/appender error path.
+ OutputStreamManager manager = fileAppender.getManager();
+ Method getStream =
OutputStreamManager.class.getDeclaredMethod("getOutputStream");
+ getStream.setAccessible(true);
+ FailingOutputStream stream =
+ new FailingOutputStream((OutputStream)
getStream.invoke(manager));
+ Method setStream =
+ OutputStreamManager.class.getDeclaredMethod(
+ "setOutputStream", OutputStream.class);
+ setStream.setAccessible(true);
+ setStream.invoke(manager, stream);
+ stream.failWrites = true;
+ for (int attempt = 1; attempt <= 2; attempt++) {
+ assertThatCode(() ->
writer.appendAndFlush(inputRecord("failed-write")))
+ .doesNotThrowAnyException();
+ assertThat(failures.getCount()).isEqualTo(attempt);
+ }
+ assertThat(diagnostics)
+ .filteredOn(
+ event ->
+ event.getMessage()
+ .getFormattedMessage()
+ .startsWith("Trace Log write
failed"))
+ .extracting(LogEvent::getLevel)
+ .containsExactly(Level.WARN, Level.DEBUG);
+ assertThat(fileAppender.ignoreExceptions()).isFalse();
+
+ stream.failWrites = false;
+ writer.appendAndFlush(inputRecord("after-recovery"));
+ assertThat(failures.getCount()).isEqualTo(2);
+
assertThat(Files.readString(traceFile)).contains("after-recovery");
+ }
+ } finally {
+ configuration.removeLogger(traceName);
+ configuration.removeLogger(writerName);
+ if (previousTrace != null) {
+ configuration.addLogger(traceName, previousTrace);
+ }
+ if (previousWriter != null) {
+ configuration.addLogger(writerName, previousWriter);
+ }
+ configuration.getAppenders().remove(appenderName);
+ if (previousAppender != null) {
+ configuration.addAppender(previousAppender);
+ }
+ context.updateLoggers();
+ if (fileAppender != null) {
+ fileAppender.stop();
+ }
+ diagnosticAppender.stop();
+ if (previousLogFile == null) {
+ System.clearProperty("log.file");
+ } else {
+ System.setProperty("log.file", previousLogFile);
+ }
+ }
+ }
+
+ private static TraceRecord inputRecord(String value) {
+ InputEvent event = new InputEvent(value);
+ return TraceRecord.create(
+ TraceContext.forEvent(null, event.getType(),
event.getId().toString(), null, null),
+ null,
+ null,
+ event.getAttributes());
+ }
+
+ private static class FailingOutputStream extends FilterOutputStream {
+ private boolean failWrites;
+
+ private FailingOutputStream(OutputStream delegate) {
+ super(delegate);
+ }
+
+ @Override
+ public void write(int value) throws IOException {
+ if (failWrites) {
+ throw new IOException("Simulated Trace Log disk write
failure");
+ }
+ out.write(value);
+ }
+
+ @Override
+ public void write(byte[] bytes, int offset, int length) throws
IOException {
+ if (failWrites) {
+ throw new IOException("Simulated Trace Log disk write
failure");
+ }
+ out.write(bytes, offset, length);
+ }
+ }
+}
