feat: support agent graceful shutdown

agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java

@@ -29,6 +29,7 @@
 import io.agentscope.core.hook.ReasoningChunkEvent;
 import io.agentscope.core.hook.SummaryChunkEvent;
 import io.agentscope.core.interruption.InterruptContext;
+import io.agentscope.core.interruption.InterruptSource;
 import io.agentscope.core.memory.InMemoryMemory;
 import io.agentscope.core.memory.LongTermMemory;
 import io.agentscope.core.memory.LongTermMemoryMode;
@@ -56,6 +57,9 @@
 import io.agentscope.core.rag.model.Document;
 import io.agentscope.core.rag.model.RetrieveConfig;
 import io.agentscope.core.session.Session;
+import io.agentscope.core.shutdown.AgentShuttingDownException;
+import io.agentscope.core.shutdown.GracefulShutdownManager;
+import io.agentscope.core.shutdown.PartialReasoningPolicy;
 import io.agentscope.core.skill.SkillBox;
 import io.agentscope.core.skill.SkillHook;
 import io.agentscope.core.state.AgentMetaState;
@@ -130,6 +134,8 @@
 public class ReActAgent extends StructuredOutputCapableAgent {
 
     private static final Logger log = LoggerFactory.getLogger(ReActAgent.class);
+    private static final GracefulShutdownManager shutdownManager =
+            GracefulShutdownManager.getInstance();
 
     // ==================== Core Dependencies ====================
 
@@ -224,8 +230,15 @@ public void saveTo(Session session, SessionKey sessionKey) {
      * @param session the session to load state from
      * @param sessionKey the session identifier
      */
+    @Override
+    public boolean loadIfExists(Session session, SessionKey sessionKey) {
+        shutdownManager.bindSession(this, session, sessionKey);
+        return super.loadIfExists(session, sessionKey);
+    }
+
     @Override
     public void loadFrom(Session session, SessionKey sessionKey) {
+        shutdownManager.bindSession(this, session, sessionKey);
         // Load memory if managed
         if (statePersistence.memoryManaged()) {
             memory.loadFrom(session, sessionKey);
@@ -440,10 +453,19 @@ private Mono<Msg> reasoning(int iter, boolean ignoreMaxIters) {
                 .onErrorResume(
                         InterruptedException.class,
                         error -> {
-                            // Save accumulated message before propagating interrupt
                             Msg msg = context.buildFinalMessage();
                             if (msg != null) {
-                                memory.addMessage(msg);
+                                boolean discard =
+                                        getInterruptSource() == InterruptSource.SYSTEM
+                                                && shutdownManager
+                                                                .getConfig()
+                                                                .partialReasoningPolicy()
+                                                        == PartialReasoningPolicy.DISCARD;
+                                // Manually interruption will save the msg, while system
+                                // interruption will discard on specific config
+                                if (!discard) {
+                                    memory.addMessage(msg);
+                                }
                             }
                             return Mono.error(error);
                         })
@@ -932,6 +954,11 @@ private Mono<Void> notifySummaryChunk(
 
     @Override
     protected Mono<Msg> handleInterrupt(InterruptContext context, Msg... originalArgs) {
+        if (context.getSource() == InterruptSource.SYSTEM) {
+            shutdownManager.saveOnInterruptObserved(this);
+            return Mono.error(new AgentShuttingDownException());
+        }
+
         String recoveryText = "I noticed that you have interrupted me. What can I do for you?";
 
         Msg recoveryMsg =

agentscope-core/src/main/java/io/agentscope/core/agent/AgentBase.java

@@ -23,6 +23,8 @@
 import io.agentscope.core.interruption.InterruptContext;
 import io.agentscope.core.interruption.InterruptSource;
 import io.agentscope.core.message.Msg;
+import io.agentscope.core.shutdown.GracefulShutdownHook;
+import io.agentscope.core.shutdown.GracefulShutdownManager;
 import io.agentscope.core.state.StateModule;
 import io.agentscope.core.tracing.TracerRegistry;
 import java.util.ArrayList;
@@ -93,14 +95,18 @@ public abstract class AgentBase implements StateModule, Agent {
     private final AtomicBoolean running = new AtomicBoolean(false);
     private final boolean checkRunning;
     private final List<Hook> hooks;
-    private static final List<Hook> systemHooks = new CopyOnWriteArrayList<>();
+    private static final List<Hook> systemHooks =
+            new CopyOnWriteArrayList<>(
+                    List.of(new GracefulShutdownHook(GracefulShutdownManager.getInstance())));
     private final Map<String, List<AgentBase>> hubSubscribers = new ConcurrentHashMap<>();
 
     // Interrupt state management (available to all agents)
     private final AtomicBoolean interruptFlag = new AtomicBoolean(false);
     private final AtomicReference<Msg> userInterruptMessage = new AtomicReference<>(null);
     // Hook non-null
     private static final Comparator<Hook> HOOK_COMPARATOR = Comparator.comparingInt(Hook::priority);
+    private final AtomicReference<InterruptSource> interruptSource =
+            new AtomicReference<>(InterruptSource.USER);
 
     /**
      * Constructor for AgentBase.
@@ -165,14 +171,7 @@ public final String getDescription() {
     @Override
     public final Mono<Msg> call(List<Msg> msgs) {
         return Mono.using(
-                () -> {
-                    if (checkRunning && !running.compareAndSet(false, true)) {
-                        throw new IllegalStateException(
-                                "Agent is still running, please wait for it to finish");
-                    }
-                    resetInterruptFlag();
-                    return this;
-                },
+                this::acquireExecution,
                 resource ->
                         TracerRegistry.get()
                                 .callAgent(
@@ -185,7 +184,7 @@ public final Mono<Msg> call(List<Msg> msgs) {
                                                         .onErrorResume(
                                                                 createErrorHandler(
                                                                         msgs.toArray(new Msg[0])))),
-                resource -> running.set(false),
+                this::releaseExecution,
                 true);
     }
 
@@ -201,14 +200,7 @@ public final Mono<Msg> call(List<Msg> msgs) {
     @Override
     public final Mono<Msg> call(List<Msg> msgs, Class<?> structuredOutputClass) {
         return Mono.using(
-                () -> {
-                    if (checkRunning && !running.compareAndSet(false, true)) {
-                        throw new IllegalStateException(
-                                "Agent is still running, please wait for it to finish");
-                    }
-                    resetInterruptFlag();
-                    return this;
-                },
+                this::acquireExecution,
                 resource ->
                         TracerRegistry.get()
                                 .callAgent(
@@ -225,7 +217,7 @@ public final Mono<Msg> call(List<Msg> msgs, Class<?> structuredOutputClass) {
                                                         .onErrorResume(
                                                                 createErrorHandler(
                                                                         msgs.toArray(new Msg[0])))),
-                resource -> running.set(false),
+                this::releaseExecution,
                 true);
     }
 
@@ -241,14 +233,7 @@ public final Mono<Msg> call(List<Msg> msgs, Class<?> structuredOutputClass) {
     @Override
     public final Mono<Msg> call(List<Msg> msgs, JsonNode schema) {
         return Mono.using(
-                () -> {
-                    if (checkRunning && !running.compareAndSet(false, true)) {
-                        throw new IllegalStateException(
-                                "Agent is still running, please wait for it to finish");
-                    }
-                    resetInterruptFlag();
-                    return this;
-                },
+                this::acquireExecution,
                 resource ->
                         TracerRegistry.get()
                                 .callAgent(
@@ -261,7 +246,7 @@ public final Mono<Msg> call(List<Msg> msgs, JsonNode schema) {
                                                         .onErrorResume(
                                                                 createErrorHandler(
                                                                         msgs.toArray(new Msg[0])))),
-                resource -> running.set(false),
+                this::releaseExecution,
                 true);
     }
 
@@ -318,6 +303,7 @@ public static void removeSystemHook(Hook hook) {
      */
     @Override
     public void interrupt() {
+        interruptSource.set(InterruptSource.USER);
         interruptFlag.set(true);
     }
 
@@ -329,12 +315,23 @@ public void interrupt() {
      */
     @Override
     public void interrupt(Msg msg) {
+        interruptSource.set(InterruptSource.USER);
         interruptFlag.set(true);
         if (msg != null) {
             userInterruptMessage.set(msg);
         }
     }
 
+    /**
+     * Interrupt execution with explicit source.
+     *
+     * @param source interruption source
+     */
+    public void interrupt(InterruptSource source) {
+        interruptSource.set(source != null ? source : InterruptSource.SYSTEM);
+        interruptFlag.set(true);
+    }
+
     /**
      * Check if the agent execution has been interrupted (reactive version).
      * Returns a Mono that completes normally if not interrupted, or errors with
@@ -376,6 +373,7 @@ protected Mono<Void> checkInterruptedAsync() {
     protected void resetInterruptFlag() {
         interruptFlag.set(false);
         userInterruptMessage.set(null);
+        interruptSource.set(InterruptSource.USER);
     }
 
     /**
@@ -386,11 +384,47 @@ protected void resetInterruptFlag() {
      */
     private InterruptContext createInterruptContext() {
         return InterruptContext.builder()
-                .source(InterruptSource.USER)
+                .source(interruptSource.get())
                 .userMessage(userInterruptMessage.get())
                 .build();
     }
 
+    /**
+     * Acquire execution resources for a {@code call()} invocation.
+     * Used as the {@code resourceSupplier} in {@link Mono#using} to guarantee that
+     * {@link #releaseExecution} is always called on completion, error, or cancellation.
+     *
+     * @return this agent instance
+     */
+    private AgentBase acquireExecution() {
+        if (checkRunning && !running.compareAndSet(false, true)) {
+            throw new IllegalStateException("Agent is still running, please wait for it to finish");
+        }
+        try {
+            resetInterruptFlag();
+            GracefulShutdownManager.getInstance().ensureAcceptingRequests();
+            GracefulShutdownManager.getInstance().registerRequest(this);
+        } catch (RuntimeException ex) {
+            if (checkRunning) {
+                running.set(false);
+            }
+            throw ex;
+        }
+        return this;
+    }
+
+    /**
+     * Release execution resources after a {@code call()} invocation.
+     * Used as the {@code resourceCleanup} in {@link Mono#using} — guaranteed to run
+     * regardless of how the reactive chain terminates (success, error, or cancel).
+     *
+     * @param resource the agent instance (ignored, uses {@code this})
+     */
+    private void releaseExecution(AgentBase resource) {
+        running.set(false);
+        GracefulShutdownManager.getInstance().unregisterRequest(this);
+    }
+
     /**
      * Create error handler for call() methods.
      * Handles InterruptedException specially and delegates to handleInterrupt,
@@ -420,6 +454,15 @@ protected AtomicBoolean getInterruptFlag() {
         return interruptFlag;
     }
 
+    /**
+     * Get current interruption source.
+     *
+     * @return interruption source
+     */
+    protected InterruptSource getInterruptSource() {
+        return interruptSource.get();
+    }
+
     /**
      * Observe a message without generating a reply.
      * This allows agents to receive messages from other agents or the environment

agentscope-core/src/main/java/io/agentscope/core/model/transport/HttpTransportFactory.java

@@ -47,8 +47,9 @@
  * HttpTransportFactory.register(myTransport);
  * }</pre>
  *
- * <p>All registered transports and the default transport will be automatically
- * closed when the JVM shuts down.
+ * <p>All registered transports and the default transport will be closed
+ * during JVM shutdown, orchestrated by {@code AgentScopeJvmShutdownHook}
+ * after all agents have finished processing.
  */
 public final class HttpTransportFactory {
 
@@ -60,9 +61,6 @@ public final class HttpTransportFactory {
     /** List of all managed transports for cleanup. */
     private static final List<HttpTransport> managedTransports = new CopyOnWriteArrayList<>();
 
-    /** Whether the shutdown hook has been registered. */
-    private static volatile boolean shutdownHookRegistered = false;
-
     /** Lock object for synchronization. */
     private static final Object lock = new Object();
 
@@ -85,7 +83,6 @@ public static HttpTransport getDefault() {
                 if (defaultTransport == null) {
                     defaultTransport = new JdkHttpTransport();
                     managedTransports.add(defaultTransport);
-                    registerShutdownHookIfNeeded();
                     log.debug("Created default HttpTransport: {}", defaultTransport);
                 }
             }
@@ -107,7 +104,6 @@ public static void setDefault(HttpTransport transport) {
             if (transport != null && !managedTransports.contains(transport)) {
                 managedTransports.add(transport);
             }
-            registerShutdownHookIfNeeded();
             log.debug("Set default HttpTransport: {}", transport);
         }
     }
@@ -124,7 +120,6 @@ public static void setDefault(HttpTransport transport) {
     public static void register(HttpTransport transport) {
         if (transport != null && !managedTransports.contains(transport)) {
             managedTransports.add(transport);
-            registerShutdownHookIfNeeded();
             log.debug("Registered HttpTransport for management: {}", transport);
         }
     }
@@ -146,30 +141,6 @@ public static boolean unregister(HttpTransport transport) {
         return removed;
     }
 
-    /**
-     * Register the JVM shutdown hook if not already registered.
-     */
-    private static void registerShutdownHookIfNeeded() {
-        if (!shutdownHookRegistered) {
-            synchronized (lock) {
-                if (!shutdownHookRegistered) {
-                    Runtime.getRuntime()
-                            .addShutdownHook(
-                                    new Thread(
-                                            () -> {
-                                                log.debug(
-                                                        "Shutdown hook triggered, closing all"
-                                                                + " HttpTransports");
-                                                shutdown();
-                                            },
-                                            "HttpTransportFactory-ShutdownHook"));
-                    shutdownHookRegistered = true;
-                    log.debug("Registered HttpTransport shutdown hook");
-                }
-            }
-        }
-    }
-
     /**
      * Shutdown and close all managed transports.
      *

agentscope-core/src/main/java/io/agentscope/core/session/InMemorySession.java

@@ -170,6 +170,15 @@ public void delete(SessionKey sessionKey) {
         sessions.remove(sessionKeyStr);
     }
 
+    @Override
+    public void delete(SessionKey sessionKey, String key) {
+        String sessionKeyStr = serializeSessionKey(sessionKey);
+        SessionData data = sessions.get(sessionKeyStr);
+        if (data != null) {
+            data.removeSingleState(key);
+        }
+    }
+
     /**
      * List all session keys.
      *
@@ -224,6 +233,10 @@ State getSingleState(String key) {
             return singleStates.get(key);
         }
 
+        void removeSingleState(String key) {
+            singleStates.remove(key);
+        }
+
         void setListState(String key, List<? extends State> values) {
             // Use List.copyOf for immutable copy to prevent external modification
             listStates.put(key, List.copyOf(values));