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..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 @@ -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; @@ -194,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 new file mode 100644 index 000000000..9b403e351 --- /dev/null +++ b/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelInvokerTest.java @@ -0,0 +1,161 @@ +/* + * 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.AfterEach; +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 { + + @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); + 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 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 + // 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"); + } + + @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..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)); @@ -72,6 +80,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));