weiqingy commented on code in PR #1071: URL: https://github.com/apache/flink-agents/pull/1071#discussion_r3892137327
########## 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"); Review Comment: nit: `Thread.interrupted()` on this line clears the flag, but it only runs if the `verify` above it passes. If that `verify` ever fails, the flag stays set on the JUnit thread. `ChatModelActionRetryTest.chatRetriesWithExponentialBackoff` is in this same package and drives a real one second backoff, so it can then fail with an unrelated `InterruptedException` and send someone chasing the wrong test. `RunnerContextImplDurableExecuteTest:91` and `:115` have the same shape. Would an `@AfterEach` calling `Thread.interrupted()` be worth adding, so the cleanup runs either way? ########## 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) { Review Comment: `Thread.sleep` on line 203 sits inside this `catch (Exception e)` block, so a cancel during the backoff wait throws from in here rather than from the call above. A catch block is not covered by its own sibling catch, and `Thread.sleep` clears the interrupt status when it throws. The call still stops, so this is only about the flag, but on this one path it ends up cleared rather than restored. `RETRY_WAIT_INTERVAL` defaults to 1, so under RETRY there is a one second window on every retry. Something like this, if useful: ```java try { Thread.sleep(currentWaitSec * 1000L); } catch (InterruptedException ie) { Thread.currentThread().interrupt(); throw ie; } ``` Worth restoring the flag there too? ########## runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java: ########## @@ -939,6 +945,12 @@ protected <T> 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; Review Comment: Tool calls run through the two methods you patched, so this rethrow reaches them too, but the tool path still finishes normally after a cancel. `ToolCallAction.java:256` catches the `InterruptedException`, line 257 records it as a tool error, and the loop moves on to the next tool. Line 85 then sends the `ToolResponseEvent` anyway, which drives another chat call, and the action is persisted as finished. `executeParallel` has the same shape at line 209. That catch predates this PR. Is the tool path meant to be in scope here, or is it worth a separate issue? ########## runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java: ########## @@ -576,6 +576,12 @@ protected <T> 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; Review Comment: `ModelRoutingResolver.java:97` runs the routing strategy through `ctx.durableExecute`, so an interrupt there reaches this rethrow. `ChatModelAction.java:609` then catches it, and under IGNORE line 621 returns normally. I could not trigger this with in-tree code. It needs a `MODEL_ROUTER` resource, and the only routing strategy in the repo today does no I/O. But `RoutingStrategy` is a user extension point, and the comment on line 612 already expects strategies that do I/O. The issue's first bullet asks for cancellation to propagate under IGNORE too. Should that catch let `InterruptedException` through before the IGNORE check? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
