From ad45bc9072032cca5c1ac32d22ee482cede715ed Mon Sep 17 00:00:00 2001 From: Ashfaq <105435085+Ashfaqbs@users.noreply.github.com> Date: Sun, 30 Aug 2026 07:38:40 +0530 Subject: [PATCH 1/4] [plan][runtime] Do not retry or persist interrupted chat calls on cancellation ChatModelInvoker.chatWithRetries() treated a job-cancellation InterruptedException the same as any ordinary model failure, so it could be retried under ERROR_HANDLING_STRATEGY=RETRY and delay task shutdown. RunnerContextImpl.durableExecuteCompletionOnly()/executeAndFinalizeCurrentCall() also recorded the interruption as a completed durable result, so a stale interruption could be replayed as terminal after recovery instead of the call being re-executed. Both now special-case InterruptedException: restore the interrupt status and propagate immediately without retrying or finalizing the durable call, regardless of FAIL/RETRY/IGNORE. Ordinary failures are unaffected. Fixes #1070. Generated-by: Claude Code 2.1.226 (Claude Opus 4.6) --- .../agents/plan/actions/ChatModelInvoker.java | 6 + .../plan/actions/ChatModelInvokerTest.java | 109 ++++++++++++++++++ .../runtime/context/RunnerContextImpl.java | 12 ++ .../RunnerContextImplDurableExecuteTest.java | 47 ++++++++ 4 files changed, 174 insertions(+) create mode 100644 plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelInvokerTest.java diff --git a/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelInvoker.java b/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelInvoker.java index 8dfacee4e..332bf0d86 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelInvoker.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelInvoker.java @@ -182,6 +182,12 @@ public ChatMessage call() throws Exception { } return new ChatAttemptResult( model, chatModel, response, actualRetryCount, totalWaitTimeSec); + } catch (InterruptedException e) { + // A cancellation signal, not a model failure: restore the interrupt status and + // propagate immediately so task shutdown isn't delayed by retry backoff or an + // extra model call, regardless of the configured error-handling strategy. + Thread.currentThread().interrupt(); + throw e; } catch (Exception e) { if (strategy == Agent.ErrorHandlingStrategy.RETRY && attempt < numRetries) { actualRetryCount = attempt + 1; diff --git a/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelInvokerTest.java b/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelInvokerTest.java new file mode 100644 index 000000000..086320252 --- /dev/null +++ b/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelInvokerTest.java @@ -0,0 +1,109 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.flink.agents.plan.actions; + +import org.apache.flink.agents.api.agents.Agent; +import org.apache.flink.agents.api.agents.AgentExecutionOptions; +import org.apache.flink.agents.api.chat.model.BaseChatModelSetup; +import org.apache.flink.agents.api.configuration.ReadableConfiguration; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup; +import org.apache.flink.agents.api.resource.ResourceType; +import org.junit.jupiter.api.Test; + +import java.util.List; +import java.util.Map; +import java.util.UUID; + +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** Tests for {@link ChatModelInvoker}. */ +class ChatModelInvokerTest { + + @Test + void testChatWithRetriesDoesNotRetryOnInterruption() throws Exception { + RunnerContext ctx = mock(RunnerContext.class); + BaseChatModelSetup chatModel = mock(BaseChatModelSetup.class); + ReadableConfiguration config = mock(ReadableConfiguration.class); + when(ctx.getConfig()).thenReturn(config); + when(config.get(AgentExecutionOptions.CHAT_ASYNC)).thenReturn(false); + when(ctx.getResource("test-model", ResourceType.CHAT_MODEL)).thenReturn(chatModel); + when(ctx.getActionMetricGroup()).thenReturn(mock(FlinkAgentsMetricGroup.class)); + when(ctx.durableExecute(any())).thenThrow(new InterruptedException("cancelled")); + + // Clear any interrupt status left over from a previous test before asserting on it below. + Thread.interrupted(); + + assertThrows( + InterruptedException.class, + () -> + ChatModelInvoker.chatWithRetries( + UUID.randomUUID(), + "test-model", + "durable-call-id", + List.of(), + Map.of(), + null, + ctx, + Agent.ErrorHandlingStrategy.RETRY, + 3, + 0)); + + // Only the first attempt should have run: retry backoff must not consume more attempts + // after a cancellation interrupts the call. + verify(ctx, times(1)).durableExecute(any()); + assertTrue(Thread.interrupted(), "interrupt status should be restored on the thread"); + } + + @Test + void testChatWithRetriesRetriesOnOrdinaryFailure() throws Exception { + RunnerContext ctx = mock(RunnerContext.class); + BaseChatModelSetup chatModel = mock(BaseChatModelSetup.class); + ReadableConfiguration config = mock(ReadableConfiguration.class); + when(ctx.getConfig()).thenReturn(config); + when(config.get(AgentExecutionOptions.CHAT_ASYNC)).thenReturn(false); + when(ctx.getResource("test-model", ResourceType.CHAT_MODEL)).thenReturn(chatModel); + when(ctx.getActionMetricGroup()).thenReturn(mock(FlinkAgentsMetricGroup.class)); + when(ctx.durableExecute(any())).thenThrow(new RuntimeException("transient failure")); + + assertThrows( + ChatModelInvoker.ChatAttemptFailed.class, + () -> + ChatModelInvoker.chatWithRetries( + UUID.randomUUID(), + "test-model", + "durable-call-id", + List.of(), + Map.of(), + null, + ctx, + Agent.ErrorHandlingStrategy.RETRY, + 2, + 0)); + + // An ordinary failure must still consume the full retry budget (initial attempt + 2 + // retries), confirming the interruption fix doesn't disturb normal retry behavior. + verify(ctx, times(3)).durableExecute(any()); + } +} diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java b/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java index dcfc34ed0..a40e56b03 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java @@ -576,6 +576,12 @@ protected T durableExecuteCompletionOnly( Exception exception = null; try { result = executionCallable.call(); + } catch (InterruptedException e) { + // A cancellation signal, not a genuine call failure: leave the durable slot + // unfinished so recovery re-executes or reconciles the call instead of replaying a + // stale interruption as a completed success or failure. + Thread.currentThread().interrupt(); + throw e; } catch (Exception e) { exception = e; } @@ -939,6 +945,12 @@ protected T executeAndFinalizeCurrentCall( Exception exception = null; try { result = callSupplier.call(); + } catch (InterruptedException e) { + // A cancellation signal, not a genuine call failure: leave the pending call + // unfinalized so recovery re-executes or reconciles it instead of replaying a stale + // interruption as a completed success or failure. + Thread.currentThread().interrupt(); + throw e; } catch (Exception e) { exception = e; } diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextImplDurableExecuteTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextImplDurableExecuteTest.java index 7d4cfde21..0b767fb97 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextImplDurableExecuteTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextImplDurableExecuteTest.java @@ -72,6 +72,53 @@ void testDurableExecuteLegacyCall() throws Exception { assertSame(context.getDurableExecutionContext().getActionState(), lastPersistedState); } + @Test + void testDurableExecuteCompletionOnlyDoesNotPersistInterruption() { + RunnerContextImpl context = createContext(new ActionState(null)); + TestDurableCallable callable = + new TestDurableCallable<>( + "legacy-call", + String.class, + () -> { + throw new InterruptedException("cancelled"); + }); + + // Clear any interrupt status left over from a previous test before asserting on it below. + Thread.interrupted(); + + assertThrows(InterruptedException.class, () -> context.durableExecute(callable)); + + assertTrue(Thread.interrupted(), "interrupt status should be restored on the thread"); + // A cancellation must not be finalized as a durable success or failure: the slot stays + // unfinished so recovery re-executes the call instead of replaying a stale interruption. + assertEquals(0, persistCallCount.get()); + assertEquals(0, context.getDurableExecutionContext().getActionState().getCallResultCount()); + } + + @Test + void testDurableExecuteCompletionOnlyReExecutesPendingSlotDoesNotPersistInterruption() { + ActionState actionState = new ActionState(null); + actionState.addCallResult(CallResult.pending("tool-call", "")); + RunnerContextImpl context = createContext(actionState); + TestDurableCallable callable = + new TestDurableCallable<>( + "tool-call", + String.class, + () -> { + throw new InterruptedException("cancelled"); + }); + + Thread.interrupted(); + + assertThrows(InterruptedException.class, () -> context.durableExecute(callable)); + + assertTrue(Thread.interrupted(), "interrupt status should be restored on the thread"); + assertEquals(0, persistCallCount.get()); + CallResult pending = + context.getDurableExecutionContext().getActionState().getCallResults().get(0); + assertTrue(pending.isPending(), "interrupted pending slot should remain unfinalized"); + } + @Test void testDurableExecuteReconcilableSuccessCall() throws Exception { RunnerContextImpl context = createContext(new ActionState(null)); From c30ba89412850353340cc9048b6f01de519d6ae6 Mon Sep 17 00:00:00 2001 From: Ashfaq <105435085+Ashfaqbs@users.noreply.github.com> Date: Mon, 31 Aug 2026 12:08:34 +0530 Subject: [PATCH 2/4] [java][plan][runtime] Restore interrupt flag around retry backoff sleep; stop leaking it in tests Thread.sleep in the RETRY backoff wait sits inside a sibling catch(Exception) block, not the InterruptedException handler above it, so a cancel during that wait cleared the interrupt flag instead of restoring it before propagating. Wrap it the same way as the call above. Also add @AfterEach cleanup in ChatModelInvokerTest and RunnerContextImplDurableExecuteTest: the existing assertTrue(Thread.interrupted()) only clears the flag when the assertion before it passes, so a failing assertion left it set on the JUnit thread for later tests to trip over. Generated-by: Claude Code 2.1.226 (Claude Sonnet 5) --- .../flink/agents/plan/actions/ChatModelInvoker.java | 7 ++++++- .../flink/agents/plan/actions/ChatModelInvokerTest.java | 9 +++++++++ .../context/RunnerContextImplDurableExecuteTest.java | 8 ++++++++ 3 files changed, 23 insertions(+), 1 deletion(-) diff --git a/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelInvoker.java b/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelInvoker.java index 332bf0d86..a22212762 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelInvoker.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelInvoker.java @@ -200,7 +200,12 @@ public ChatMessage call() throws Exception { numRetries, currentWaitSec); if (currentWaitSec > 0) { - Thread.sleep(currentWaitSec * 1000L); + try { + Thread.sleep(currentWaitSec * 1000L); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + throw ie; + } totalWaitTimeSec += currentWaitSec; } continue; diff --git a/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelInvokerTest.java b/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelInvokerTest.java index 086320252..4c1dd2b90 100644 --- a/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelInvokerTest.java +++ b/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelInvokerTest.java @@ -24,6 +24,7 @@ import org.apache.flink.agents.api.context.RunnerContext; import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup; import org.apache.flink.agents.api.resource.ResourceType; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; import java.util.List; @@ -41,6 +42,14 @@ /** Tests for {@link ChatModelInvoker}. */ class ChatModelInvokerTest { + @AfterEach + void clearInterruptStatus() { + // Prevents a leftover interrupt flag (e.g. if the assertion below the interruption test + // ever fails before consuming it) from failing an unrelated later test's real Thread.sleep + // backoff with a spurious InterruptedException. + Thread.interrupted(); + } + @Test void testChatWithRetriesDoesNotRetryOnInterruption() throws Exception { RunnerContext ctx = mock(RunnerContext.class); diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextImplDurableExecuteTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextImplDurableExecuteTest.java index 0b767fb97..816b874d5 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextImplDurableExecuteTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextImplDurableExecuteTest.java @@ -26,6 +26,7 @@ import org.apache.flink.agents.runtime.actionstate.CallResult; import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl; import org.apache.flink.runtime.metrics.groups.UnregisteredMetricGroups; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -53,6 +54,13 @@ void setUp() { lastPersistedState = null; } + @AfterEach + void clearInterruptStatus() { + // Prevents a leftover interrupt flag (e.g. if an assertion below one of the interruption + // tests ever fails before consuming it) from leaking into an unrelated later test. + Thread.interrupted(); + } + @Test void testDurableExecuteLegacyCall() throws Exception { RunnerContextImpl context = createContext(new ActionState(null)); From 5803bb1d4ecd82d4c1061268a4e48d152559db19 Mon Sep 17 00:00:00 2001 From: Ashfaq <105435085+Ashfaqbs@users.noreply.github.com> Date: Wed, 2 Sep 2026 21:16:25 +0530 Subject: [PATCH 3/4] [plan] Add deterministic test for the retry backoff interrupt-flag fix The existing tests pass retryWaitIntervalSec = 0, so neither exercises the Thread.sleep block that c30ba894 wrapped in a try(InterruptedException)/restore/rethrow. Adds a test that sets the interrupt flag from within the mocked durableExecute call itself, so Thread.sleep throws immediately (no wall-clock cost) with RETRY, numRetries = 1, retryWaitIntervalSec = 1, and asserts the flag is restored rather than left cleared. Verified RED against the pre-c30ba894 sleep block (temporarily restored, spotless check skipped) and GREEN against the current fix. Generated-by: Claude Code 2.1.226 (Claude Sonnet 5) --- .../plan/actions/ChatModelInvokerTest.java | 43 +++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelInvokerTest.java b/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelInvokerTest.java index 4c1dd2b90..679a765c3 100644 --- a/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelInvokerTest.java +++ b/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelInvokerTest.java @@ -85,6 +85,49 @@ void testChatWithRetriesDoesNotRetryOnInterruption() throws Exception { assertTrue(Thread.interrupted(), "interrupt status should be restored on the thread"); } + @Test + void testChatWithRetriesRestoresInterruptFlagFromRetryBackoffSleep() throws Exception { + RunnerContext ctx = mock(RunnerContext.class); + BaseChatModelSetup chatModel = mock(BaseChatModelSetup.class); + ReadableConfiguration config = mock(ReadableConfiguration.class); + when(ctx.getConfig()).thenReturn(config); + when(config.get(AgentExecutionOptions.CHAT_ASYNC)).thenReturn(false); + when(ctx.getResource("test-model", ResourceType.CHAT_MODEL)).thenReturn(chatModel); + when(ctx.getActionMetricGroup()).thenReturn(mock(FlinkAgentsMetricGroup.class)); + // The interrupt fires from within the retry backoff's Thread.sleep, not from the call + // itself: set the flag first so Thread.sleep throws immediately, at no wall-clock cost. + when(ctx.durableExecute(any())) + .thenAnswer( + invocation -> { + Thread.currentThread().interrupt(); + throw new RuntimeException("transient failure"); + }); + + assertThrows( + InterruptedException.class, + () -> + ChatModelInvoker.chatWithRetries( + UUID.randomUUID(), + "test-model", + "durable-call-id", + List.of(), + Map.of(), + null, + ctx, + Agent.ErrorHandlingStrategy.RETRY, + 1, + 1)); + + // Only the first attempt should have run: the sleep before the retry throws before a + // second call is made. + verify(ctx, times(1)).durableExecute(any()); + // This is the only assertion that distinguishes the fix from the pre-fix code: the + // ChatAttemptFailed/times(1) shape above passes either way, since Thread.sleep still + // aborts the retry loop on both. Only the restored flag proves the backoff sleep's catch + // block restores it instead of leaving it cleared. + assertTrue(Thread.interrupted(), "interrupt status should be restored on the thread"); + } + @Test void testChatWithRetriesRetriesOnOrdinaryFailure() throws Exception { RunnerContext ctx = mock(RunnerContext.class); From 5d5cf629042f1b7ef6dd67406420a0e38802aec0 Mon Sep 17 00:00:00 2001 From: Ashfaq <105435085+Ashfaqbs@users.noreply.github.com> Date: Thu, 3 Sep 2026 21:04:17 +0530 Subject: [PATCH 4/4] fix(plan): correct stale comment referencing wrong test assertion weiqingy's nit on the LGTM review: the comment said ChatAttemptFailed, but this test asserts InterruptedException (line 106). The only ChatAttemptFailed assertion is in the sibling test testChatWithRetriesRetriesOnOrdinaryFailure, so the comment pointed a reader at the wrong test. --- .../apache/flink/agents/plan/actions/ChatModelInvokerTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelInvokerTest.java b/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelInvokerTest.java index 679a765c3..9b403e351 100644 --- a/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelInvokerTest.java +++ b/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelInvokerTest.java @@ -122,7 +122,7 @@ void testChatWithRetriesRestoresInterruptFlagFromRetryBackoffSleep() throws Exce // second call is made. verify(ctx, times(1)).durableExecute(any()); // This is the only assertion that distinguishes the fix from the pre-fix code: the - // ChatAttemptFailed/times(1) shape above passes either way, since Thread.sleep still + // InterruptedException/times(1) shape above passes either way, since Thread.sleep still // aborts the retry loop on both. Only the restored flag proves the backoff sleep's catch // block restores it instead of leaving it cleared. assertTrue(Thread.interrupted(), "interrupt status should be restored on the thread");