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 84f3ab9f2 [runtime][java] Reuse saved Kafka action results after 
recovery (#1161)
84f3ab9f2 is described below

commit 84f3ab9f289438527c4572cf8128208b06a18502
Author: Stefan Wang <[email protected]>
AuthorDate: Mon Oct 12 08:38:13 2026 +0100

    [runtime][java] Reuse saved Kafka action results after recovery (#1161)
    
    Generated-by: GitHub Copilot CLI (gpt-6-astra)
    Signed-off-by: 1fanwang <[email protected]>
---
 .../runtime/actionstate/ActionStateStore.java      |  16 +
 .../runtime/actionstate/KafkaActionStateStore.java |  92 +++++-
 .../runtime/operator/ActionExecutionOperator.java  |   7 +-
 .../operator/ActionExecutionOperatorFactory.java   |   2 +-
 .../runtime/operator/DurableExecutionManager.java  |  11 +-
 .../KafkaActionStateStoreRecoveryTest.java         | 336 +++++++++++++++++++++
 .../actionstate/KafkaActionStateStoreTest.java     | 282 +++++++++++++++++
 .../operator/DurableExecutionManagerTest.java      |  27 +-
 8 files changed, 741 insertions(+), 32 deletions(-)

diff --git 
a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateStore.java
 
b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateStore.java
index 844a37e6a..583c87b65 100644
--- 
a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateStore.java
+++ 
b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateStore.java
@@ -22,6 +22,7 @@ import org.apache.flink.agents.plan.actions.Action;
 
 import java.io.IOException;
 import java.util.List;
+import java.util.Map;
 import java.util.function.IntPredicate;
 
 /** Interface for storing and retrieving the state of actions performed by 
agents. */
@@ -124,4 +125,19 @@ public interface ActionStateStore extends AutoCloseable {
     default Object getRecoveryMarker() {
         return null;
     }
+
+    /**
+     * Get a marker object using the per-key sequence numbers reflected in the 
current Flink
+     * checkpoint.
+     *
+     * <p>Implementations that do not need the sequence numbers delegate to 
{@link
+     * #getRecoveryMarker()} by default.
+     *
+     * @param lastCompletedSequenceNumbers per-key sequence numbers reflected 
in the current Flink
+     *     checkpoint
+     * @return a marker object, or null if not supported
+     */
+    default Object getRecoveryMarker(Map<?, Long> 
lastCompletedSequenceNumbers) {
+        return getRecoveryMarker();
+    }
 }
diff --git 
a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStore.java
 
b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStore.java
index 60a72a80f..773e37ca0 100644
--- 
a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStore.java
+++ 
b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStore.java
@@ -38,6 +38,7 @@ import org.apache.kafka.clients.producer.KafkaProducer;
 import org.apache.kafka.clients.producer.Producer;
 import org.apache.kafka.clients.producer.ProducerConfig;
 import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.serialization.StringDeserializer;
@@ -55,6 +56,7 @@ import java.util.Map;
 import java.util.Properties;
 import java.util.Set;
 import java.util.UUID;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.function.IntPredicate;
 
@@ -93,8 +95,7 @@ public class KafkaActionStateStore implements 
ActionStateStore {
     // In memory action state for quick state retrieval
     private final Map<String, ActionState> actionStates;
 
-    // Record the lastest sequence number for each key that should be 
considered as valid
-    private final Map<String, Long> latestKeySeqNum;
+    private final Map<String, StateRecordOffset> latestStateOffsets;
 
     // Kafka producer
     private final Producer<String, ActionState> producer;
@@ -146,7 +147,7 @@ public class KafkaActionStateStore implements 
ActionStateStore {
         this.topic = topic;
         this.topicMetadataLoader = () -> testTopicMetadata(consumer, topic);
         this.topicMetadata = testTopicMetadata(consumer, topic);
-        this.latestKeySeqNum = new HashMap<>();
+        this.latestStateOffsets = new HashMap<>();
         this.agentConfiguration = agentConfiguration;
         this.tombstoneEnabled = 
agentConfiguration.get(KAFKA_ACTION_STATE_TOMBSTONE_ENABLED);
         this.keyEncoder = Preconditions.checkNotNull(keyEncoder, "keyEncoder 
cannot be null");
@@ -166,7 +167,7 @@ public class KafkaActionStateStore implements 
ActionStateStore {
             AgentConfiguration agentConfiguration, ActionStateKeyEncoder 
keyEncoder) {
         this.keyEncoder = Preconditions.checkNotNull(keyEncoder, "keyEncoder 
cannot be null");
         this.actionStates = new HashMap<>();
-        this.latestKeySeqNum = new HashMap<>();
+        this.latestStateOffsets = new HashMap<>();
         this.agentConfiguration = agentConfiguration;
         this.tombstoneEnabled = 
agentConfiguration.get(KAFKA_ACTION_STATE_TOMBSTONE_ENABLED);
         String cleanupControlTopic =
@@ -224,13 +225,21 @@ public class KafkaActionStateStore implements 
ActionStateStore {
         try {
             ProducerRecord<String, ActionState> kafkaRecord =
                     new ProducerRecord<>(topic, stateKey, state);
-            producer.send(kafkaRecord);
-            actionStates.put(stateKey, state);
+            Future<RecordMetadata> sendResult = producer.send(kafkaRecord);
             producer.flush();
+            RecordMetadata metadata = sendResult.get();
+            actionStates.put(stateKey, state);
+            latestStateOffsets.put(
+                    stateKey, new StateRecordOffset(metadata.partition(), 
metadata.offset()));
             LOG.debug(
                     "Stored action state to Kafka: key={}, isCompleted={}",
                     stateKey,
                     state.isCompleted());
+        } catch (InterruptedException e) {
+            // Future.get() clears the interrupt flag before it throws, so 
restore it here.
+            // Wrapping it would leave the caller with an error but no visible 
cancellation.
+            Thread.currentThread().interrupt();
+            throw e;
         } catch (Exception e) {
             throw new RuntimeException("Failed to send action state to Kafka", 
e);
         }
@@ -251,15 +260,15 @@ public class KafkaActionStateStore implements 
ActionStateStore {
         boolean hasDivergence = checkDivergence(businessKeyIdentity, seqNum);
 
         if (!actionStates.containsKey(stateKey) || hasDivergence) {
-            // Clean up this key's states with sequence number greater than 
the requested seqNum.
-            actionStates
-                    .keySet()
-                    .removeIf(
-                            cachedKey ->
-                                    
ActionStateUtil.matchesBusinessKeyIdentityWithSeqNum(
-                                            cachedKey,
-                                            businessKeyIdentity,
-                                            stateSeqNum -> stateSeqNum > 
seqNum));
+            List<String> keysToRemove = new ArrayList<>();
+            for (String cachedKey : actionStates.keySet()) {
+                if (ActionStateUtil.matchesBusinessKeyIdentityWithSeqNum(
+                        cachedKey, businessKeyIdentity, stateSeqNum -> 
stateSeqNum > seqNum)) {
+                    keysToRemove.add(cachedKey);
+                }
+            }
+            actionStates.keySet().removeAll(keysToRemove);
+            latestStateOffsets.keySet().removeAll(keysToRemove);
         }
 
         ActionState result = actionStates.get(stateKey);
@@ -348,8 +357,12 @@ public class KafkaActionStateStore implements 
ActionStateStore {
                     if (record.value() == null) {
                         // Tombstone record - remove the key from cache
                         actionStates.remove(record.key());
+                        latestStateOffsets.remove(record.key());
                     } else {
                         actionStates.put(record.key(), record.value());
+                        latestStateOffsets.put(
+                                record.key(),
+                                new StateRecordOffset(record.partition(), 
record.offset()));
                     }
                 }
 
@@ -568,14 +581,32 @@ public class KafkaActionStateStore implements 
ActionStateStore {
 
         // Remove from in-memory cache (always, regardless of tombstone 
success)
         actionStates.keySet().removeAll(keysToPrune);
+        latestStateOffsets.keySet().removeAll(keysToPrune);
 
         LOG.debug("Pruned state for key: {} up to sequence number: {}", key, 
seqNum);
     }
 
-    /** Returns a versioned marker containing the Kafka topic identity and 
current end offsets. */
+    /**
+     * Returns a versioned marker containing per-partition replay start 
offsets. For partitions
+     * whose latest durable state is already reflected in the Flink 
checkpoint, the end offset is
+     * used; for partitions with newer durable state, the marker rewinds to 
that state's earliest
+     * record offset.
+     */
     @Override
     public Object getRecoveryMarker() {
+        return getRecoveryMarker(Map.of());
+    }
+
+    @Override
+    public Object getRecoveryMarker(Map<?, Long> lastCompletedSequenceNumbers) 
{
         try {
+            Map<String, Long> lastCompletedSequenceNumbersByIdentity = new 
HashMap<>();
+            lastCompletedSequenceNumbers.forEach(
+                    (key, sequenceNumber) ->
+                            lastCompletedSequenceNumbersByIdentity.merge(
+                                    
keyEncoder.generateBusinessKeyIdentity(key),
+                                    sequenceNumber,
+                                    Math::max));
             KafkaTopicMetadata currentTopicMetadata = 
loadVerifiedTopicMetadata();
             List<TopicPartition> partitions = new ArrayList<>();
             for (Integer partition : currentTopicMetadata.getPartitions()) {
@@ -586,6 +617,12 @@ public class KafkaActionStateStore implements 
ActionStateStore {
             for (Map.Entry<TopicPartition, Long> entry : 
endOffsets.entrySet()) {
                 recoveryOffsets.put(entry.getKey().partition(), 
entry.getValue());
             }
+            latestStateOffsets.forEach(
+                    (stateKey, offset) -> {
+                        if (needsReplay(stateKey, 
lastCompletedSequenceNumbersByIdentity)) {
+                            recoveryOffsets.merge(offset.partition, 
offset.offset, Math::min);
+                        }
+                    });
             if 
(!recoveryOffsets.keySet().equals(currentTopicMetadata.getPartitions())) {
                 throw new IllegalStateException(
                         String.format(
@@ -780,4 +817,27 @@ public class KafkaActionStateStore implements 
ActionStateStore {
     private interface TopicMetadataLoader {
         KafkaTopicMetadata load() throws Exception;
     }
+
+    private boolean needsReplay(
+            String stateKey, Map<String, Long> 
lastCompletedSequenceNumbersByIdentity) {
+        String businessKeyIdentity = 
ActionStateUtil.businessKeyIdentityOf(stateKey);
+        if (businessKeyIdentity == null) {
+            return true;
+        }
+
+        Long checkpointedSeqNum = 
lastCompletedSequenceNumbersByIdentity.get(businessKeyIdentity);
+        return checkpointedSeqNum == null
+                || !ActionStateUtil.matchesBusinessKeyIdentityWithSeqNum(
+                        stateKey, businessKeyIdentity, seqNum -> seqNum <= 
checkpointedSeqNum);
+    }
+
+    private static final class StateRecordOffset {
+        private final int partition;
+        private final long offset;
+
+        private StateRecordOffset(int partition, long offset) {
+            this.partition = partition;
+            this.offset = offset;
+        }
+    }
 }
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 7df1cc5d2..4368d87ca 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
@@ -1214,9 +1214,10 @@ public class ActionExecutionOperator<IN, OUT> extends 
AbstractStreamOperator<OUT
     }
 
     private void snapshotStateInternal(StateSnapshotContext context) throws 
Exception {
-        durableExecManager.snapshotRecoveryMarker();
-        durableExecManager.snapshotLastCompletedSequenceNumbers(
-                getKeyedStateBackend(), context.getCheckpointId());
+        Map<Object, Long> lastCompletedSequenceNumbers =
+                durableExecManager.snapshotLastCompletedSequenceNumbers(
+                        getKeyedStateBackend(), context.getCheckpointId());
+        
durableExecManager.snapshotRecoveryMarker(lastCompletedSequenceNumbers);
 
         super.snapshotState(context);
     }
diff --git 
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorFactory.java
 
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorFactory.java
index 3eeec4bf1..dc8faf66a 100644
--- 
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorFactory.java
+++ 
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorFactory.java
@@ -53,7 +53,7 @@ public class ActionExecutionOperatorFactory<IN, OUT> extends 
AbstractStreamOpera
     }
 
     @VisibleForTesting
-    protected ActionExecutionOperatorFactory(
+    public ActionExecutionOperatorFactory(
             AgentPlan agentPlan, Boolean inputIsJava, ActionStateStore 
actionStateStore) {
         this(agentPlan, inputIsJava, false, actionStateStore);
     }
diff --git 
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/DurableExecutionManager.java
 
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/DurableExecutionManager.java
index 864ab00ee..81dc72b04 100644
--- 
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/DurableExecutionManager.java
+++ 
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/DurableExecutionManager.java
@@ -408,9 +408,10 @@ class DurableExecutionManager implements 
ActionStatePersister, AutoCloseable {
         }
     }
 
-    void snapshotRecoveryMarker() throws Exception {
+    void snapshotRecoveryMarker(Map<Object, Long> 
lastCompletedSequenceNumbers) throws Exception {
         if (actionStateStore != null) {
-            Object recoveryMarker = actionStateStore.getRecoveryMarker();
+            Object recoveryMarker =
+                    
actionStateStore.getRecoveryMarker(lastCompletedSequenceNumbers);
             if (recoveryMarker != null) {
                 recoveryMarkerOpState.update(List.of(recoveryMarker));
             }
@@ -437,12 +438,13 @@ class DurableExecutionManager implements 
ActionStatePersister, AutoCloseable {
      *
      * @param keyedStateBackend the keyed state backend to scan.
      * @param checkpointId the id of the checkpoint being snapshotted.
+     * @return the completed sequence numbers captured for the checkpoint.
      */
     @SuppressWarnings("unchecked")
-    void snapshotLastCompletedSequenceNumbers(
+    Map<Object, Long> snapshotLastCompletedSequenceNumbers(
             KeyedStateBackend<?> keyedStateBackend, long checkpointId) throws 
Exception {
         if (actionStateStore == null) {
-            return;
+            return Map.of();
         }
         HashMap<Object, Long> keyToSeqNum = new HashMap<>();
         ((KeyedStateBackend<Object>) keyedStateBackend)
@@ -458,6 +460,7 @@ class DurableExecutionManager implements 
ActionStatePersister, AutoCloseable {
                             }
                         });
         checkpointIdToSeqNums.put(checkpointId, keyToSeqNum);
+        return keyToSeqNum;
     }
 
     // --- Durable execution context map accessors ---
diff --git 
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStoreRecoveryTest.java
 
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStoreRecoveryTest.java
new file mode 100644
index 000000000..dd85cc3af
--- /dev/null
+++ 
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStoreRecoveryTest.java
@@ -0,0 +1,336 @@
+/*
+ * 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.actionstate;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.api.InputEvent;
+import org.apache.flink.agents.api.OutputEvent;
+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.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.runtime.operator.ActionExecutionOperator;
+import org.apache.flink.agents.runtime.operator.ActionExecutionOperatorFactory;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.MockConsumer;
+import org.apache.kafka.clients.producer.MockProducer;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.PartitionInfo;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.StringSerializer;
+import org.junit.jupiter.api.Test;
+import org.testcontainers.DockerClientFactory;
+import org.testcontainers.kafka.KafkaContainer;
+import org.testcontainers.utility.DockerImageName;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static 
org.apache.flink.agents.api.configuration.AgentConfigOptions.KAFKA_ACTION_STATE_TOPIC;
+import static 
org.apache.flink.agents.api.configuration.AgentConfigOptions.KAFKA_ACTION_STATE_TOPIC_NUM_PARTITIONS;
+import static 
org.apache.flink.agents.api.configuration.AgentConfigOptions.KAFKA_BOOTSTRAP_SERVERS;
+import static 
org.apache.flink.agents.runtime.actionstate.ActionStateTestUtils.createKeyEncoder;
+import static 
org.apache.kafka.clients.consumer.internals.AutoOffsetResetStrategy.EARLIEST;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.junit.jupiter.api.Assumptions.assumeTrue;
+
+/** End-to-end recovery test for a queued action with a persisted durable call 
result. */
+public class KafkaActionStateStoreRecoveryTest {
+
+    private static final String TOPIC = "test-action-state-recovery";
+    private static final int MAX_PARALLELISM = 128;
+    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+
+    private static final AtomicInteger EXTERNAL_CALLS = new AtomicInteger();
+    private static final AtomicBoolean FAIL_AFTER_DURABLE_CALL = new 
AtomicBoolean();
+
+    public static void durableCall(Event event, RunnerContext context) throws 
Exception {
+        Long input = (Long) InputEvent.fromEvent(event).getInput();
+        Long result =
+                context.durableExecute(
+                        new DurableCallable<Long>() {
+                            @Override
+                            public String getId() {
+                                return "multiply";
+                            }
+
+                            @Override
+                            public Class<Long> getResultClass() {
+                                return Long.class;
+                            }
+
+                            @Override
+                            public Long call() {
+                                return callExternal(input);
+                            }
+                        });
+        if (FAIL_AFTER_DURABLE_CALL.get()) {
+            throw new IllegalStateException("action failed after persisting 
its durable result");
+        }
+        context.sendEvent(new OutputEvent(result));
+    }
+
+    private static Long callExternal(Long input) {
+        EXTERNAL_CALLS.incrementAndGet();
+        return input * 10;
+    }
+
+    private static AgentPlan plan() throws Exception {
+        Action action =
+                new Action(
+                        "durableCall",
+                        new JavaFunction(
+                                KafkaActionStateStoreRecoveryTest.class,
+                                "durableCall",
+                                new Class<?>[] {Event.class, 
RunnerContext.class}),
+                        Collections.singletonList(InputEvent.EVENT_TYPE));
+        Map<String, Action> actions = new HashMap<>();
+        actions.put(action.getName(), action);
+        AgentConfiguration config = new AgentConfiguration();
+        config.set(AgentExecutionOptions.PARALLEL_EXECUTION_ENABLED, false);
+        return new AgentPlan(actions, new HashMap<>(), config);
+    }
+
+    @Test
+    void queuedActionUsesRecoveredDurableResultAfterRestore() throws Exception 
{
+        MockProducer<String, ActionState> producer =
+                new MockProducer<>(
+                        true,
+                        new ActionStateKeyPartitioner(),
+                        new StringSerializer(),
+                        new ActionStateKafkaSeder());
+        KafkaActionStateStore store =
+                new KafkaActionStateStore(
+                        new HashMap<>(),
+                        new AgentConfiguration(),
+                        producer,
+                        consumerWithLiveEndOffsets(producer),
+                        TOPIC,
+                        createKeyEncoder(MAX_PARALLELISM));
+
+        EXTERNAL_CALLS.set(0);
+        AgentPlan agentPlan = plan();
+        Action action = agentPlan.getActions().values().iterator().next();
+        OperatorSubtaskState checkpoint;
+
+        try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> 
harness =
+                new KeyedOneInputStreamOperatorTestHarness<>(
+                        new ActionExecutionOperatorFactory<>(agentPlan, true, 
store),
+                        (KeySelector<Long, Long>) value -> 1L,
+                        TypeInformation.of(Long.class),
+                        MAX_PARALLELISM,
+                        1,
+                        0)) {
+            harness.open();
+            ActionExecutionOperator<Long, Object> operator =
+                    (ActionExecutionOperator<Long, Object>) 
harness.getOperator();
+
+            harness.processElement(new StreamRecord<>(3L));
+            operator.waitInFlightEventsFinished();
+            assertThat(harness.getOutput())
+                    .singleElement()
+                    .isInstanceOfSatisfying(
+                            StreamRecord.class,
+                            record -> 
assertThat(record.getValue()).isEqualTo(30L));
+            harness.getOutput().clear();
+
+            InputEvent inputEvent = new InputEvent(7L);
+            ActionState pendingState = new ActionState(inputEvent);
+            pendingState.addCallResult(
+                    new CallResult("multiply", 
OBJECT_MAPPER.writeValueAsBytes(callExternal(7L))));
+            store.put(1L, 1L, action, inputEvent, pendingState);
+
+            harness.processElement(new StreamRecord<>(7L));
+            assertThat(harness.getOutput()).isEmpty();
+            checkpoint = harness.snapshot(1L, 1L);
+            assertThat(EXTERNAL_CALLS.get()).isEqualTo(2);
+        }
+
+        MockConsumer<String, ActionState> recoveryConsumer = 
recoveryConsumer(producer.history());
+        try (KafkaActionStateStore recoveredStore =
+                        new KafkaActionStateStore(
+                                new HashMap<>(),
+                                new AgentConfiguration(),
+                                new MockProducer<>(
+                                        true,
+                                        new ActionStateKeyPartitioner(),
+                                        new StringSerializer(),
+                                        new ActionStateKafkaSeder()),
+                                recoveryConsumer,
+                                TOPIC,
+                                createKeyEncoder(MAX_PARALLELISM));
+                KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> 
restored =
+                        new KeyedOneInputStreamOperatorTestHarness<>(
+                                new ActionExecutionOperatorFactory<>(
+                                        agentPlan, true, recoveredStore),
+                                (KeySelector<Long, Long>) value -> 1L,
+                                TypeInformation.of(Long.class),
+                                MAX_PARALLELISM,
+                                1,
+                                0)) {
+            restored.initializeState(checkpoint);
+            restored.open();
+            ActionExecutionOperator<Long, Object> operator =
+                    (ActionExecutionOperator<Long, Object>) 
restored.getOperator();
+            operator.waitInFlightEventsFinished();
+
+            assertThat(restored.getOutput())
+                    .singleElement()
+                    .isInstanceOfSatisfying(
+                            StreamRecord.class,
+                            record -> 
assertThat(record.getValue()).isEqualTo(70L));
+            assertThat(EXTERNAL_CALLS.get()).isEqualTo(2);
+        }
+    }
+
+    @Test
+    void queuedActionReusesResultFromPreviousAttemptWithKafka() throws 
Exception {
+        assumeTrue(
+                DockerClientFactory.instance().isDockerAvailable(),
+                "Docker is required for the Kafka recovery integration test");
+
+        EXTERNAL_CALLS.set(0);
+        try (KafkaContainer kafka =
+                new 
KafkaContainer(DockerImageName.parse("apache/kafka-native:3.8.0"))) {
+            kafka.start();
+            AgentConfiguration configuration = new AgentConfiguration();
+            configuration.set(KAFKA_BOOTSTRAP_SERVERS, 
kafka.getBootstrapServers());
+            configuration.set(KAFKA_ACTION_STATE_TOPIC, TOPIC);
+            configuration.set(KAFKA_ACTION_STATE_TOPIC_NUM_PARTITIONS, 1);
+            AgentPlan agentPlan = plan();
+            OperatorSubtaskState queuedCheckpoint;
+            OperatorSubtaskState pendingResultCheckpoint;
+
+            // The first attempt saves its result after the checkpoint, then 
fails mid-action.
+            try (KafkaActionStateStore store =
+                            new KafkaActionStateStore(
+                                    configuration, 
createKeyEncoder(MAX_PARALLELISM));
+                    KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> 
harness =
+                            kafkaHarness(agentPlan, store)) {
+                harness.open();
+                harness.processElement(new StreamRecord<>(7L));
+                queuedCheckpoint = harness.snapshot(1L, 1L);
+                ActionExecutionOperator<Long, Object> operator =
+                        (ActionExecutionOperator<Long, Object>) 
harness.getOperator();
+                FAIL_AFTER_DURABLE_CALL.set(true);
+                assertThatThrownBy(operator::waitInFlightEventsFinished)
+                        .hasRootCauseMessage("action failed after persisting 
its durable result");
+                assertThat(EXTERNAL_CALLS.get()).isEqualTo(1);
+                assertThat(harness.getOutput()).isEmpty();
+            } finally {
+                FAIL_AFTER_DURABLE_CALL.set(false);
+            }
+
+            // Restore the queued action and take another checkpoint before 
dispatching it.
+            // Its durable result now precedes the Kafka end offset captured 
by this checkpoint.
+            try (KafkaActionStateStore store =
+                            new KafkaActionStateStore(
+                                    configuration, 
createKeyEncoder(MAX_PARALLELISM));
+                    KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> 
harness =
+                            kafkaHarness(agentPlan, store)) {
+                harness.initializeState(queuedCheckpoint);
+                harness.open();
+                pendingResultCheckpoint = harness.snapshot(2L, 2L);
+                assertThat(EXTERNAL_CALLS.get()).isEqualTo(1);
+                assertThat(harness.getOutput()).isEmpty();
+            }
+
+            try (KafkaActionStateStore store =
+                            new KafkaActionStateStore(
+                                    configuration, 
createKeyEncoder(MAX_PARALLELISM));
+                    KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> 
harness =
+                            kafkaHarness(agentPlan, store)) {
+                harness.initializeState(pendingResultCheckpoint);
+                harness.open();
+                ActionExecutionOperator<Long, Object> operator =
+                        (ActionExecutionOperator<Long, Object>) 
harness.getOperator();
+                operator.waitInFlightEventsFinished();
+
+                assertThat(harness.getOutput())
+                        .singleElement()
+                        .isInstanceOfSatisfying(
+                                StreamRecord.class,
+                                record -> 
assertThat(record.getValue()).isEqualTo(70L));
+                assertThat(EXTERNAL_CALLS.get()).isEqualTo(1);
+            }
+        }
+    }
+
+    private static KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> 
kafkaHarness(
+            AgentPlan agentPlan, KafkaActionStateStore store) throws Exception 
{
+        return new KeyedOneInputStreamOperatorTestHarness<>(
+                new ActionExecutionOperatorFactory<>(agentPlan, true, store),
+                (KeySelector<Long, Long>) value -> 1L,
+                TypeInformation.of(Long.class),
+                MAX_PARALLELISM,
+                1,
+                0);
+    }
+
+    /**
+     * Emulates a real Kafka broker: the store's consumer reads end offsets 
from the same topic the
+     * producer writes, so a marker taken mid-run reflects exactly the records 
already flushed.
+     */
+    private static MockConsumer<String, ActionState> 
consumerWithLiveEndOffsets(
+            MockProducer<String, ActionState> producer) {
+        return new MockConsumer<>(EARLIEST.name()) {
+            {
+                updatePartitions(TOPIC, List.of(new PartitionInfo(TOPIC, 0, 
null, null, null)));
+            }
+
+            @Override
+            public synchronized Map<TopicPartition, Long> endOffsets(
+                    Collection<TopicPartition> partitions) {
+                Map<TopicPartition, Long> result = new HashMap<>();
+                for (TopicPartition topicPartition : partitions) {
+                    result.put(topicPartition, (long) 
producer.history().size());
+                }
+                return result;
+            }
+        };
+    }
+
+    private static MockConsumer<String, ActionState> recoveryConsumer(
+            List<ProducerRecord<String, ActionState>> history) {
+        MockConsumer<String, ActionState> consumer = new 
MockConsumer<>(EARLIEST.name());
+        consumer.updatePartitions(TOPIC, List.of(new PartitionInfo(TOPIC, 0, 
null, null, null)));
+        consumer.assign(List.of(new TopicPartition(TOPIC, 0)));
+        consumer.updateBeginningOffsets(Map.of(new TopicPartition(TOPIC, 0), 
0L));
+        consumer.updateEndOffsets(Map.of(new TopicPartition(TOPIC, 0), (long) 
history.size()));
+        for (int i = 0; i < history.size(); i++) {
+            ProducerRecord<String, ActionState> record = history.get(i);
+            consumer.addRecord(new ConsumerRecord<>(TOPIC, 0, i, record.key(), 
record.value()));
+        }
+        return consumer;
+    }
+}
diff --git 
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStoreTest.java
 
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStoreTest.java
index 28a6aebbf..bc01e2e4a 100644
--- 
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStoreTest.java
+++ 
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStoreTest.java
@@ -46,6 +46,7 @@ import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
+import org.mockito.InOrder;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -54,7 +55,9 @@ import java.util.List;
 import java.util.Map;
 import java.util.Properties;
 import java.util.Set;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -67,9 +70,12 @@ import static 
org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.assertj.core.api.Assertions.catchThrowable;
 import static org.assertj.core.api.Assertions.entry;
 import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 /** Unit tests for {@link KafkaActionStateStore}. */
 public class KafkaActionStateStoreTest {
@@ -203,6 +209,74 @@ public class KafkaActionStateStoreTest {
         assertThat(record.value()).isEqualTo(testActionState);
     }
 
+    @Test
+    void testPutFlushesBeforeAwaitingMetadata() throws Exception {
+        Producer<String, ActionState> producer = mock(Producer.class);
+        Future<RecordMetadata> future = mock(Future.class);
+        when(producer.send(any())).thenReturn(future);
+        when(future.get())
+                .thenReturn(new RecordMetadata(new TopicPartition(TEST_TOPIC, 
0), 0L, 0, 0L, 0, 0));
+        mockConsumer.updateEndOffsets(
+                Map.of(
+                        new TopicPartition(TEST_TOPIC, 0), 1L,
+                        new TopicPartition(TEST_TOPIC, 1), 1L));
+        try (KafkaActionStateStore store =
+                new KafkaActionStateStore(
+                        actionStates,
+                        new AgentConfiguration(),
+                        producer,
+                        mockConsumer,
+                        TEST_TOPIC,
+                        createKeyEncoder(MAX_PARALLELISM))) {
+            store.put(TEST_KEY, 1L, testAction, testEvent, testActionState);
+
+            InOrder order = inOrder(producer, future);
+            order.verify(producer).send(any());
+            order.verify(producer).flush();
+            order.verify(future).get();
+            assertThat(store.get(TEST_KEY, 1L, testAction, 
testEvent)).isEqualTo(testActionState);
+            assertThat(((KafkaActionStateRecoveryMarker) 
store.getRecoveryMarker()).getOffsets())
+                    .containsExactlyInAnyOrderEntriesOf(Map.of(0, 0L, 1, 1L));
+        }
+    }
+
+    @ParameterizedTest
+    @ValueSource(booleans = {true, false})
+    void testPutDoesNotCacheFailedWrite(boolean failOnFlush) throws Exception {
+        Producer<String, ActionState> producer = mock(Producer.class);
+        Future<RecordMetadata> future = mock(Future.class);
+        when(producer.send(any())).thenReturn(future);
+        RuntimeException failure = new RuntimeException("simulated broker 
failure");
+        if (failOnFlush) {
+            doThrow(failure).when(producer).flush();
+            when(future.get())
+                    .thenReturn(
+                            new RecordMetadata(new TopicPartition(TEST_TOPIC, 
0), 0L, 0, 0L, 0, 0));
+        } else {
+            when(future.get()).thenThrow(new ExecutionException(failure));
+        }
+        mockConsumer.updateEndOffsets(
+                Map.of(
+                        new TopicPartition(TEST_TOPIC, 0), 1L,
+                        new TopicPartition(TEST_TOPIC, 1), 1L));
+        try (KafkaActionStateStore store =
+                new KafkaActionStateStore(
+                        actionStates,
+                        new AgentConfiguration(),
+                        producer,
+                        mockConsumer,
+                        TEST_TOPIC,
+                        createKeyEncoder(MAX_PARALLELISM))) {
+            assertThatThrownBy(
+                            () -> store.put(TEST_KEY, 1L, testAction, 
testEvent, testActionState))
+                    .isInstanceOf(RuntimeException.class)
+                    .hasRootCause(failure);
+            assertThat(actionStates).isEmpty();
+            assertThat(((KafkaActionStateRecoveryMarker) 
store.getRecoveryMarker()).getOffsets())
+                    .containsExactlyInAnyOrderEntriesOf(Map.of(0, 1L, 1, 1L));
+        }
+    }
+
     @Test
     void testGetNonExistentActionState() throws Exception {
         actionStates.put(
@@ -245,6 +319,70 @@ public class KafkaActionStateStoreTest {
         assertNull(actionStateStore.get(TEST_KEY, 4L, testAction, testEvent));
     }
 
+    @Test
+    void testPutRestoresInterruptFlagWhenSendIsInterrupted() throws Exception {
+        // Send finishes like a cancelled FutureTask: get() clears the 
thread's interrupt flag
+        // before throwing InterruptedException. put() must restore the flag.
+        MockProducer<String, ActionState> interruptingProducer =
+                new MockProducer<>(
+                        false,
+                        new ActionStateKeyPartitioner(),
+                        new StringSerializer(),
+                        new ActionStateKafkaSeder()) {
+                    @Override
+                    public Future<RecordMetadata> send(ProducerRecord<String, 
ActionState> record) {
+                        return new Future<>() {
+                            @Override
+                            public RecordMetadata get() throws 
InterruptedException {
+                                Thread.interrupted(); // clear the flag like 
FutureTask.get()
+                                throw new InterruptedException("simulated 
interruption");
+                            }
+
+                            @Override
+                            public RecordMetadata get(long timeout, TimeUnit 
unit)
+                                    throws InterruptedException {
+                                return get();
+                            }
+
+                            @Override
+                            public boolean cancel(boolean 
mayInterruptIfRunning) {
+                                return false;
+                            }
+
+                            @Override
+                            public boolean isCancelled() {
+                                return false;
+                            }
+
+                            @Override
+                            public boolean isDone() {
+                                return false;
+                            }
+                        };
+                    }
+                };
+        KafkaActionStateStore store =
+                new KafkaActionStateStore(
+                        new HashMap<>(),
+                        new AgentConfiguration(),
+                        interruptingProducer,
+                        mockConsumer,
+                        TEST_TOPIC,
+                        createKeyEncoder(MAX_PARALLELISM));
+        try {
+            Thread.currentThread().interrupt();
+            Throwable thrown =
+                    catchThrowable(
+                            () -> store.put(TEST_KEY, 1L, testAction, 
testEvent, testActionState));
+            assertThat(thrown).isInstanceOf(InterruptedException.class);
+            assertThat(Thread.currentThread().isInterrupted())
+                    .as("The interrupt flag cleared by Future.get() must be 
restored")
+                    .isTrue();
+        } finally {
+            Thread.interrupted(); // clear the flag so the remaining tests run 
clean
+        }
+    }
+
     @Test
     void testGetCleansFutureStateForKeyContainingUnderscore() throws Exception 
{
         String flinkKey = "user_123";
@@ -290,6 +428,150 @@ public class KafkaActionStateStoreTest {
                 .containsEntry(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, 
false);
     }
 
+    @Test
+    void testRebuildStateKeepsPendingResultBeforeRecoveryMarker() throws 
Exception {
+        ActionState pendingState = new ActionState(testEvent);
+        actionStateStore.put(TEST_KEY, 1L, testAction, testEvent, 
pendingState);
+
+        ProducerRecord<String, ActionState> persisted = 
mockProducer.history().get(0);
+        mockConsumer.updatePartitions(
+                TEST_TOPIC,
+                List.of(
+                        new PartitionInfo(TEST_TOPIC, 0, null, null, null),
+                        new PartitionInfo(TEST_TOPIC, 1, null, null, null)));
+        mockConsumer.updateEndOffsets(
+                Map.of(
+                        new TopicPartition(TEST_TOPIC, 0),
+                        1L,
+                        new TopicPartition(TEST_TOPIC, 1),
+                        1L));
+
+        Object recoveryMarkerObj = actionStateStore.getRecoveryMarker();
+        KafkaActionStateRecoveryMarker recoveryMarker =
+                (KafkaActionStateRecoveryMarker) recoveryMarkerObj;
+        int recoveryPartition =
+                recoveryMarker.getOffsets().entrySet().stream()
+                        .filter(entry -> entry.getValue() == 0L)
+                        .map(Map.Entry::getKey)
+                        .findFirst()
+                        .orElse(0);
+        TopicPartition topicPartition = new TopicPartition(TEST_TOPIC, 
recoveryPartition);
+
+        MockConsumer<String, ActionState> recoveryConsumer = new 
MockConsumer<>(EARLIEST.name());
+        recoveryConsumer.updatePartitions(
+                TEST_TOPIC,
+                List.of(
+                        new PartitionInfo(TEST_TOPIC, 0, null, null, null),
+                        new PartitionInfo(TEST_TOPIC, 1, null, null, null)));
+        recoveryConsumer.assign(List.of(topicPartition));
+        recoveryConsumer.updateBeginningOffsets(
+                Map.of(
+                        new TopicPartition(TEST_TOPIC, 0), 0L,
+                        new TopicPartition(TEST_TOPIC, 1), 0L));
+        recoveryConsumer.updateEndOffsets(
+                Map.of(
+                        new TopicPartition(TEST_TOPIC, 0), 1L,
+                        new TopicPartition(TEST_TOPIC, 1), 1L));
+        recoveryConsumer.addRecord(
+                new ConsumerRecord<>(
+                        TEST_TOPIC, recoveryPartition, 0L, persisted.key(), 
persisted.value()));
+
+        Map<String, ActionState> recoveredStates = new HashMap<>();
+        try (KafkaActionStateStore recoveredStore =
+                new KafkaActionStateStore(
+                        recoveredStates,
+                        new AgentConfiguration(),
+                        null,
+                        recoveryConsumer,
+                        TEST_TOPIC,
+                        createKeyEncoder(MAX_PARALLELISM))) {
+            recoveredStore.rebuildState(List.of(recoveryMarker));
+            assertThat(recoveredStore.get(TEST_KEY, 1L, testAction, testEvent))
+                    .as("Recovery must reuse the saved pending result instead 
of skipping it")
+                    .isEqualTo(pendingState);
+        }
+    }
+
+    @Test
+    void 
testRecoveryMarkerUsesSnapshotCompletedSequencesWithoutRetainingThem() throws 
Exception {
+        ActionState completedState = new ActionState(testEvent);
+        completedState.markCompleted();
+        ActionState pendingState = new ActionState(new InputEvent("pending"));
+
+        actionStateStore.put(TEST_KEY, 1L, testAction, testEvent, 
completedState);
+        actionStateStore.put(TEST_KEY, 2L, testAction, testEvent, 
pendingState);
+
+        mockConsumer.updatePartitions(
+                TEST_TOPIC,
+                List.of(
+                        new PartitionInfo(TEST_TOPIC, 0, null, null, null),
+                        new PartitionInfo(TEST_TOPIC, 1, null, null, null)));
+        mockConsumer.updateEndOffsets(
+                Map.of(
+                        new TopicPartition(TEST_TOPIC, 0),
+                        2L,
+                        new TopicPartition(TEST_TOPIC, 1),
+                        2L));
+
+        Object recoveryMarkerObj =
+                
actionStateStore.getRecoveryMarker(Collections.singletonMap(TEST_KEY, 1L));
+        KafkaActionStateRecoveryMarker recoveryMarker =
+                (KafkaActionStateRecoveryMarker) recoveryMarkerObj;
+
+        
assertThat(recoveryMarker.getOffsets().values()).contains(1L).doesNotContain(0L);
+
+        KafkaActionStateRecoveryMarker markerWithoutCompletedSequences =
+                (KafkaActionStateRecoveryMarker)
+                        
actionStateStore.getRecoveryMarker(Collections.emptyMap());
+        
assertThat(markerWithoutCompletedSequences.getOffsets().values()).contains(0L);
+    }
+
+    @Test
+    void testRecoveryMarkerUsesRestoredCompletedSequencesBeforePruning() 
throws Exception {
+        TopicPartition partition = new TopicPartition(TEST_TOPIC, 0);
+        mockConsumer.updatePartitions(
+                TEST_TOPIC, List.of(new PartitionInfo(TEST_TOPIC, 0, null, 
null, null)));
+        actionStateStore = createStore(MAX_PARALLELISM);
+        actionStateStore.put("pending-key", 0L, testAction, testEvent, 
testActionState);
+        actionStateStore.put("completed-key", 0L, testAction, testEvent, 
testActionState);
+        mockConsumer.updateEndOffsets(Map.of(partition, 2L));
+        Object marker = 
actionStateStore.getRecoveryMarker(Map.of("completed-key", 0L));
+        assertThat(((KafkaActionStateRecoveryMarker) marker).getOffsets())
+                .containsExactly(entry(0, 0L));
+
+        MockConsumer<String, ActionState> recoveryConsumer = new 
MockConsumer<>(EARLIEST.name());
+        recoveryConsumer.updatePartitions(
+                TEST_TOPIC, List.of(new PartitionInfo(TEST_TOPIC, 0, null, 
null, null)));
+        recoveryConsumer.assign(List.of(partition));
+        recoveryConsumer.updateBeginningOffsets(Map.of(partition, 0L));
+        recoveryConsumer.updateEndOffsets(Map.of(partition, 2L));
+        List<ProducerRecord<String, ActionState>> history = 
mockProducer.history();
+        for (int offset = 0; offset < history.size(); offset++) {
+            ProducerRecord<String, ActionState> record = history.get(offset);
+            recoveryConsumer.addRecord(
+                    new ConsumerRecord<>(TEST_TOPIC, 0, offset, record.key(), 
record.value()));
+        }
+        Map<String, ActionState> recoveredStates = new HashMap<>();
+        try (KafkaActionStateStore recoveredStore =
+                new KafkaActionStateStore(
+                        recoveredStates,
+                        new AgentConfiguration(),
+                        null,
+                        recoveryConsumer,
+                        TEST_TOPIC,
+                        createKeyEncoder(MAX_PARALLELISM))) {
+            recoveredStore.rebuildState(List.of(marker));
+            assertThat(recoveredStates).hasSize(2);
+
+            KafkaActionStateRecoveryMarker nextMarker =
+                    (KafkaActionStateRecoveryMarker)
+                            recoveredStore.getRecoveryMarker(
+                                    Map.of("pending-key", 0L, "completed-key", 
0L));
+            assertThat(nextMarker.getOffsets()).containsExactly(entry(0, 2L));
+            assertThat(recoveredStates).hasSize(2);
+        }
+    }
+
     @Test
     void testPruneState() throws Exception {
         // Arrange
diff --git 
a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/DurableExecutionManagerTest.java
 
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/DurableExecutionManagerTest.java
index 6c5b6e51d..c2bc56446 100644
--- 
a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/DurableExecutionManagerTest.java
+++ 
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/DurableExecutionManagerTest.java
@@ -25,6 +25,7 @@ import org.apache.flink.agents.plan.AgentPlan;
 import org.apache.flink.agents.plan.actions.Action;
 import org.apache.flink.agents.runtime.actionstate.ActionState;
 import org.apache.flink.agents.runtime.actionstate.ActionStateKeyEncoder;
+import org.apache.flink.agents.runtime.actionstate.ActionStateStore;
 import org.apache.flink.agents.runtime.actionstate.InMemoryActionStateStore;
 import org.apache.flink.agents.runtime.actionstate.KafkaActionStateStore;
 import org.apache.flink.agents.runtime.context.RunnerContextImpl;
@@ -106,7 +107,19 @@ class DurableExecutionManagerTest {
         assertThat(dem.maybeGetActionState("k", 0L, action, event)).isNull();
         dem.maybeInitActionState("k", 0L, action, event);
         dem.notifyCheckpointComplete(1L);
-        dem.snapshotRecoveryMarker();
+        dem.snapshotRecoveryMarker(Map.of());
+        dem.close();
+    }
+
+    @Test
+    void snapshotRecoveryMarkerPassesCompletedSequencesToStore() throws 
Exception {
+        ActionStateStore store = mock(ActionStateStore.class);
+        DurableExecutionManager dem = new DurableExecutionManager(store);
+        Map<Object, Long> lastCompletedSequenceNumbers = Map.of("key", 3L);
+
+        dem.snapshotRecoveryMarker(lastCompletedSequenceNumbers);
+
+        verify(store).getRecoveryMarker(lastCompletedSequenceNumbers);
         dem.close();
     }
 
@@ -221,17 +234,15 @@ class DurableExecutionManagerTest {
                 .when(backend)
                 .applyToAllKeys(any(), any(), any(), any());
 
-        dem.snapshotLastCompletedSequenceNumbers(backend, 1000L);
+        Map<Object, Long> lastCompletedSequenceNumbers =
+                dem.snapshotLastCompletedSequenceNumbers(backend, 1000L);
 
-        assertThat(dem.getCheckpointIdToSeqNums())
-                .containsOnlyKeys(1000L)
-                .extractingByKey(1000L)
-                .asInstanceOf(
-                        org.assertj.core.api.InstanceOfAssertFactories.map(
-                                Object.class, Long.class))
+        assertThat(lastCompletedSequenceNumbers)
                 .containsEntry("k1", 10L)
                 .containsEntry("k2", 20L)
                 .hasSize(2);
+        assertThat(dem.getCheckpointIdToSeqNums())
+                .containsEntry(1000L, lastCompletedSequenceNumbers);
 
         dem.close();
     }

Reply via email to