diff --git a/agentscope-core/src/main/java/io/agentscope/core/shutdown/AgentScopeJvmShutdownHook.java b/agentscope-core/src/main/java/io/agentscope/core/shutdown/AgentScopeJvmShutdownHook.java index 0354179c0..ed3167b60 100644 --- a/agentscope-core/src/main/java/io/agentscope/core/shutdown/AgentScopeJvmShutdownHook.java +++ b/agentscope-core/src/main/java/io/agentscope/core/shutdown/AgentScopeJvmShutdownHook.java @@ -40,7 +40,20 @@ public final class AgentScopeJvmShutdownHook { private AgentScopeJvmShutdownHook() {} + /** + * Resets the REGISTERED flag for testing purposes. + * + *

This method is intended for testing only. It allows tests to verify + * the JVM hook registration behavior in isolation. + */ + static void resetForTesting() { + REGISTERED.set(false); + } + public static void register(GracefulShutdownManager manager) { + if (!manager.getConfig().isRegister()) { + return; + } if (!REGISTERED.compareAndSet(false, true)) { return; } diff --git a/agentscope-core/src/main/java/io/agentscope/core/shutdown/GracefulShutdownConfig.java b/agentscope-core/src/main/java/io/agentscope/core/shutdown/GracefulShutdownConfig.java index 960463cb9..8c0e0ea43 100644 --- a/agentscope-core/src/main/java/io/agentscope/core/shutdown/GracefulShutdownConfig.java +++ b/agentscope-core/src/main/java/io/agentscope/core/shutdown/GracefulShutdownConfig.java @@ -25,6 +25,7 @@ *

* * @param shutdownTimeout maximum duration to wait for shutdown completion; @@ -32,9 +33,12 @@ * to complete; if specified, must be a positive duration * @param partialReasoningPolicy policy for handling incomplete reasoning results during shutdown; * cannot be null + * @param isRegister Policy for control register jvm shutdown hook */ public record GracefulShutdownConfig( - Duration shutdownTimeout, PartialReasoningPolicy partialReasoningPolicy) { + Duration shutdownTimeout, + PartialReasoningPolicy partialReasoningPolicy, + boolean isRegister) { /** * Default configuration instance. @@ -43,10 +47,25 @@ public record GracefulShutdownConfig( * */ public static final GracefulShutdownConfig DEFAULT = - new GracefulShutdownConfig(null, PartialReasoningPolicy.SAVE); + new GracefulShutdownConfig(null, PartialReasoningPolicy.SAVE, false); + + /** + * + * + *

Uses the following settings: + *

+ */ + public GracefulShutdownConfig( + Duration shutdownTimeout, PartialReasoningPolicy partialReasoningPolicy) { + this(shutdownTimeout, partialReasoningPolicy, false); + } /** * Compact constructor for validation. diff --git a/agentscope-core/src/test/java/io/agentscope/core/shutdown/GracefulShutdownTest.java b/agentscope-core/src/test/java/io/agentscope/core/shutdown/GracefulShutdownTest.java index 2a1c5f44b..bde9fa7ce 100644 --- a/agentscope-core/src/test/java/io/agentscope/core/shutdown/GracefulShutdownTest.java +++ b/agentscope-core/src/test/java/io/agentscope/core/shutdown/GracefulShutdownTest.java @@ -113,6 +113,35 @@ void configRejectsInvalidTimeout() { void configRejectsNullPolicy() { assertThrows(NullPointerException.class, () -> new GracefulShutdownConfig(null, null)); } + + @Test + @DisplayName("DEFAULT config has isRegister set to false") + void defaultConfigJvmHookDisabled() { + GracefulShutdownConfig cfg = GracefulShutdownConfig.DEFAULT; + assertFalse(cfg.isRegister()); + } + + @Test + @DisplayName("2-param constructor sets isRegister to false") + void twoParamConstructorJvmHookEnabled() { + GracefulShutdownConfig cfg = + new GracefulShutdownConfig(Duration.ofSeconds(10), PartialReasoningPolicy.SAVE); + assertFalse(cfg.isRegister()); + } + + @Test + @DisplayName("3-param constructor allows explicit isRegister value") + void threeParamConstructorExplicitValue() { + GracefulShutdownConfig cfgWithHook = + new GracefulShutdownConfig( + Duration.ofSeconds(10), PartialReasoningPolicy.SAVE, true); + assertTrue(cfgWithHook.isRegister()); + + GracefulShutdownConfig cfgWithoutHook = + new GracefulShutdownConfig( + Duration.ofSeconds(10), PartialReasoningPolicy.SAVE, false); + assertFalse(cfgWithoutHook.isRegister()); + } } // ==================== GracefulShutdownManager ==================== @@ -576,6 +605,40 @@ void unrelatedEventsPassThrough() { } } + // ==================== AgentScopeJvmShutdownHook ==================== + + @Nested + @DisplayName("AgentScopeJvmShutdownHook") + class JvmShutdownHookTests { + + @Test + @DisplayName("register does nothing when isRegister is false") + void registerSkippedWhenDisabled() { + AgentScopeJvmShutdownHook.resetForTesting(); + manager.setConfig( + new GracefulShutdownConfig( + Duration.ofSeconds(10), PartialReasoningPolicy.SAVE, false)); + + // Should not throw and should not register the hook + assertDoesNotThrow(() -> AgentScopeJvmShutdownHook.register(manager)); + } + + @Test + @DisplayName("register is idempotent when isRegister is true") + void registerIdempotentWhenEnabled() { + AgentScopeJvmShutdownHook.resetForTesting(); + manager.setConfig( + new GracefulShutdownConfig( + Duration.ofSeconds(10), PartialReasoningPolicy.SAVE, true)); + + // First call should succeed + assertDoesNotThrow(() -> AgentScopeJvmShutdownHook.register(manager)); + + // Second call should also succeed (idempotent) + assertDoesNotThrow(() -> AgentScopeJvmShutdownHook.register(manager)); + } + } + // ==================== AgentShuttingDownException ==================== @Nested diff --git a/agentscope-examples/graceful-shutdown/src/main/java/io/agentscope/examples/shutdown/e2e/AgentService.java b/agentscope-examples/graceful-shutdown/src/main/java/io/agentscope/examples/shutdown/e2e/AgentService.java index 681180d1e..b93f0cf74 100644 --- a/agentscope-examples/graceful-shutdown/src/main/java/io/agentscope/examples/shutdown/e2e/AgentService.java +++ b/agentscope-examples/graceful-shutdown/src/main/java/io/agentscope/examples/shutdown/e2e/AgentService.java @@ -89,6 +89,9 @@ public class AgentService { @Value("${agent.shutdown-partial-reasoning-policy:save}") private String partialReasoningPolicy; + @Value("${agent.shutdown-hook-is-register:true}") + private boolean isRegister; + @PostConstruct public void initShutdownConfig() { Duration gracefulShutdownTime = @@ -100,7 +103,7 @@ public void initShutdownConfig() { GracefulShutdownManager.getInstance() .setConfig( new GracefulShutdownConfig( - gracefulShutdownTime, shutdownPartialReasoningPolicy)); + gracefulShutdownTime, shutdownPartialReasoningPolicy, isRegister)); } /** diff --git a/agentscope-examples/graceful-shutdown/src/main/java/io/agentscope/examples/shutdown/smoke/GracefulShutdownExample.java b/agentscope-examples/graceful-shutdown/src/main/java/io/agentscope/examples/shutdown/smoke/GracefulShutdownExample.java index 52ae0ef5b..e14500b29 100644 --- a/agentscope-examples/graceful-shutdown/src/main/java/io/agentscope/examples/shutdown/smoke/GracefulShutdownExample.java +++ b/agentscope-examples/graceful-shutdown/src/main/java/io/agentscope/examples/shutdown/smoke/GracefulShutdownExample.java @@ -129,7 +129,8 @@ private static void scenario1_toolExecutionTimeout(String apiKey, Session sessio shutdownManager.resetForTesting(); shutdownManager.setConfig( - new GracefulShutdownConfig(Duration.ofSeconds(5), PartialReasoningPolicy.SAVE)); + new GracefulShutdownConfig( + Duration.ofSeconds(5), PartialReasoningPolicy.SAVE, true)); Toolkit toolkit = new Toolkit(); toolkit.registerTool(new DataAnalysisTool()); diff --git a/agentscope-examples/graceful-shutdown/src/main/resources/application.yml b/agentscope-examples/graceful-shutdown/src/main/resources/application.yml index 9b965dcd3..25a6223ff 100644 --- a/agentscope-examples/graceful-shutdown/src/main/resources/application.yml +++ b/agentscope-examples/graceful-shutdown/src/main/resources/application.yml @@ -26,4 +26,5 @@ dashscope: agent: model-name: qwen3.5-plus shutdown-timeout-seconds: -1 - shutdown-partial-reasoning-policy: save \ No newline at end of file + shutdown-partial-reasoning-policy: save + shutdown-hook-is-register: true \ No newline at end of file