fix(scheduling): make task registration mutually exclusive with shutdown (issue #216)

argi-orchestration/argi-workflow-graph-core/src/main/java/io/github/agentic/ai/workflow/graph/scheduling/DefaultScheduledAgentManager.java

@@ -50,10 +50,18 @@ public class DefaultScheduledAgentManager implements ScheduledAgentManager {
 	private volatile boolean shutdown = false;
 
 	private DefaultScheduledAgentManager() {
-		this.taskScheduler = createTaskScheduler();
+		this(createTaskScheduler());
 		log.info("Default Scheduled Agent Manager initialized with shared task scheduler");
 	}
 
+	/**
+	 * Test seam: the eager singleton cannot be resurrected and its real scheduler cannot
+	 * drive deterministic interleaving tests, so tests inject their own scheduler.
+	 */
+	DefaultScheduledAgentManager(TaskScheduler taskScheduler) {
+		this.taskScheduler = taskScheduler;
+	}
+
 	/**
 	 * Get the singleton instance of the default manager
 	 */
@@ -82,6 +90,12 @@ public String registerTask(ScheduledAgentTask task) {
 		String taskId = String.format("agent-task-%s-%d", task.getName(), taskIdGenerator.getAndIncrement());
 		lock.writeLock().lock();
 		try {
+			// shutdown() sets the flag under the same write lock, so re-checking here
+			// makes the check-and-register atomic with shutdown: a registration that
+			// passed the fast-fail check cannot be interleaved into a shut-down manager.
+			if (shutdown) {
+				throw new IllegalStateException("Default Scheduled Agent Manager is shut down");
+			}
 			activeTasks.put(taskId, task);
 			log.debug("Registered scheduled agent task: {}", taskId);
 			return taskId;
@@ -204,7 +218,7 @@ public boolean isShutdown() {
 		return shutdown;
 	}
 
-	private TaskScheduler createTaskScheduler() {
+	private static TaskScheduler createTaskScheduler() {
 		ThreadPoolTaskScheduler scheduler = new ThreadPoolTaskScheduler();
 		scheduler.setPoolSize(3);
 		scheduler.setThreadNamePrefix("agent-scheduler-");

argi-orchestration/argi-workflow-graph-core/src/main/java/io/github/agentic/ai/workflow/graph/scheduling/ScheduledAgentManager.java

@@ -34,6 +34,9 @@ public interface ScheduledAgentManager {
 	 * Register a scheduled agent task
 	 * @param task the task to register
 	 * @return the unique task ID assigned to this task
+	 * @throws IllegalStateException if the manager is shut down; the rejection is decided
+	 * atomically against shutdown, so a registration racing with an in-progress shutdown
+	 * must also be rejected instead of leaving a stranded task
 	 */
 	String registerTask(ScheduledAgentTask task);
 

argi-orchestration/argi-workflow-graph-core/src/test/java/io/github/agentic/ai/workflow/graph/scheduling/ScheduledAgentManagerShutdownMutexTest.java

@@ -0,0 +1,197 @@
+/*
+ * Copyright 2025-2026 the original author or authors.
+ *
+ * Licensed 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
+ *
+ *     https://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 io.github.agentic.ai.workflow.graph.scheduling;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+
+import io.github.agentic.ai.workflow.graph.CompiledGraph;
+import io.github.agentic.ai.workflow.graph.KeyStrategy;
+import io.github.agentic.ai.workflow.graph.StateGraph;
+import io.github.agentic.ai.workflow.graph.action.AsyncNodeAction;
+import io.github.agentic.ai.workflow.graph.scheduling.ScheduledAgentTaskLifecycleTest.RecordingTaskScheduler;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+
+import static io.github.agentic.ai.workflow.graph.StateGraph.END;
+import static io.github.agentic.ai.workflow.graph.StateGraph.START;
+import static io.github.agentic.ai.workflow.graph.scheduling.ScheduleLifecycleListener.ScheduleEvent.STOPPED;
+import static io.github.agentic.ai.workflow.graph.scheduling.ScheduledAgentTaskLifecycleTest.graphFailingFirstAttempts;
+import static org.junit.jupiter.api.Assertions.assertAll;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+/** 复现 #216：注册与 shutdown 必须原子互斥，通过锁外检查的注册不得滞留进已关闭的管理器。 */
+class ScheduledAgentManagerShutdownMutexTest {
+
+	private final ScheduledAgentManagerFactory factory = ScheduledAgentManagerFactory.getInstance();
+
+	private final ScheduledAgentManager originalManager = factory.getManager();
+
+	/** 恢复测试前的管理器，避免全局 provider 影响其他测试。 */
+	@AfterEach
+	void restoreOriginalManager() {
+		factory.registerProvider(() -> originalManager);
+	}
+
+	/** 验证 shutdown 期间通过锁外检查的注册，在关闭完成后被锁内复查拒绝且不产生滞留记录。 */
+	@Test
+	void aRegistrationPassingTheCheckDuringShutdownCannotEnterActiveTasks() throws Exception {
+		DefaultScheduledAgentManager manager = new DefaultScheduledAgentManager(new RecordingTaskScheduler());
+		factory.registerProvider(() -> manager);
+		CountDownLatch stopCallbackEntered = new CountDownLatch(1);
+		CountDownLatch releaseStopCallback = new CountDownLatch(1);
+		new ScheduledAgentTask(graphFailingFirstAttempts(new AtomicInteger(), 0),
+				ScheduleConfig.builder().initialDelay(0).addListener(new ScheduleLifecycleListener() {
+					@Override
+					public void onEvent(ScheduleEvent event, Object data) {
+						if (event == STOPPED) {
+							stopCallbackEntered.countDown();
+							try {
+								releaseStopCallback.await(10, TimeUnit.SECONDS);
+							}
+							catch (InterruptedException e) {
+								Thread.currentThread().interrupt();
+							}
+						}
+					}
+				}).build());
+		assertEquals(1, manager.getActiveTaskCount());
+
+		CountDownLatch nameRead = new CountDownLatch(1);
+		CountDownLatch releaseName = new CountDownLatch(1);
+		AtomicReference<Throwable> registrationFailure = new AtomicReference<>();
+		Thread shutdownThread = new Thread(manager::shutdown, "shutdown-thread");
+		try {
+			shutdownThread.start();
+			// shutdown 持有写锁停在 stop 回调里，此刻 shutdown 标志仍为 false。
+			assertTrue(stopCallbackEntered.await(5, TimeUnit.SECONDS),
+					"shutdown must hold the write lock inside the stop callback");
+			// registerTask 在锁外检查与加锁之间调用 task.getName()，用它做可控缝隙。
+			Thread registrationThread = new Thread(() -> {
+				try {
+					new ScheduledAgentTask(graphWithNameGatedBy(nameRead, releaseName),
+							ScheduleConfig.builder().initialDelay(0).build());
+				}
+				catch (Throwable throwable) {
+					registrationFailure.set(throwable);
+				}
+			}, "registration-thread");
+			registrationThread.start();
+			assertTrue(nameRead.await(5, TimeUnit.SECONDS),
+					"registration must pass the pre-lock check while shutdown is in progress");
+
+			releaseStopCallback.countDown();
+			shutdownThread.join(TimeUnit.SECONDS.toMillis(5));
+			releaseName.countDown();
+			registrationThread.join(TimeUnit.SECONDS.toMillis(5));
+
+			assertAll(() -> assertFalse(shutdownThread.isAlive(), "shutdown must finish"),
+					() -> assertFalse(registrationThread.isAlive(), "registration must finish"),
+					() -> assertTrue(registrationFailure.get() instanceof IllegalStateException,
+							"A registration racing with shutdown must be rejected, got: " + registrationFailure.get()),
+					() -> assertEquals(0, manager.getActiveTaskCount(),
+							"No task may be stranded in a shut-down manager's activeTasks"),
+					() -> assertTrue(manager.isShutdown()));
+		}
+		finally {
+			releaseStopCallback.countDown();
+			releaseName.countDown();
+			shutdownThread.join(TimeUnit.SECONDS.toMillis(5));
+		}
+	}
+
+	/** 验证关闭完成后的注册被明确拒绝，且不留任何任务记录。 */
+	@Test
+	void aRegistrationAfterShutdownIsRejectedAndLeavesNoResidue() throws Exception {
+		DefaultScheduledAgentManager manager = new DefaultScheduledAgentManager(new RecordingTaskScheduler());
+		factory.registerProvider(() -> manager);
+		new ScheduledAgentTask(graphFailingFirstAttempts(new AtomicInteger(), 0),
+				ScheduleConfig.builder().initialDelay(0).build());
+		assertEquals(1, manager.getActiveTaskCount());
+
+		manager.shutdown();
+		assertThrows(IllegalStateException.class, () -> new ScheduledAgentTask(
+				graphFailingFirstAttempts(new AtomicInteger(), 0), ScheduleConfig.builder().initialDelay(0).build()));
+
+		assertAll(() -> assertTrue(manager.isShutdown()), () -> assertEquals(0, manager.getActiveTaskCount()));
+	}
+
+	/** 验证并发 stop 与 shutdown 保持幂等：结束后无残留记录，所有任务都已停止。 */
+	@Test
+	void concurrentStopsAndShutdownsRemainIdempotentWithoutResidue() throws Exception {
+		DefaultScheduledAgentManager manager = new DefaultScheduledAgentManager(new RecordingTaskScheduler());
+		factory.registerProvider(() -> manager);
+		ScheduledAgentTask first = new ScheduledAgentTask(graphFailingFirstAttempts(new AtomicInteger(), 0),
+				ScheduleConfig.builder().initialDelay(0).build());
+		ScheduledAgentTask second = new ScheduledAgentTask(graphFailingFirstAttempts(new AtomicInteger(), 0),
+				ScheduleConfig.builder().initialDelay(0).build());
+		assertEquals(2, manager.getActiveTaskCount());
+
+		AtomicReference<Throwable> failure = new AtomicReference<>();
+		Thread[] workers = new Thread[] { new Thread(manager::shutdown, "shutdown-1"),
+				new Thread(manager::shutdown, "shutdown-2"), new Thread(first::stop, "stop-1"),
+				new Thread(first::stop, "stop-1-again"), new Thread(second::stop, "stop-2"),
+				new Thread(second::stop, "stop-2-again") };
+		for (Thread worker : workers) {
+			worker.setUncaughtExceptionHandler((thread, throwable) -> failure.compareAndSet(null, throwable));
+			worker.start();
+		}
+		for (Thread worker : workers) {
+			worker.join(TimeUnit.SECONDS.toMillis(5));
+		}
+
+		assertAll(() -> assertEquals(null, failure.get(), "concurrent stop and shutdown must not throw"),
+				() -> assertTrue(manager.isShutdown()),
+				() -> assertEquals(0, manager.getActiveTaskCount(), "no task may survive shutdown"),
+				() -> assertTrue(first.isStopped()), () -> assertTrue(second.isStopped()));
+	}
+
+	/**
+	 * Builds a graph whose {@code getName()} is gated on the given latches, hooking the
+	 * gap between the shutdown check and the write lock inside {@code registerTask}.
+	 */
+	private static CompiledGraph graphWithNameGatedBy(CountDownLatch entered, CountDownLatch release) throws Exception {
+		Map<String, KeyStrategy> strategies = new HashMap<>();
+		strategies.put("result", (left, right) -> right);
+		StateGraph graph = new StateGraph("name-gated-registration", () -> strategies) {
+
+			@Override
+			public String getName() {
+				entered.countDown();
+				try {
+					release.await(10, TimeUnit.SECONDS);
+				}
+				catch (InterruptedException e) {
+					Thread.currentThread().interrupt();
+				}
+				return "name-gated-registration";
+			}
+
+		};
+		graph.addNode("attempt", AsyncNodeAction.node_async(state -> Map.of("result", "ok")))
+			.addEdge(START, "attempt")
+			.addEdge("attempt", END);
+		return graph.compile();
+	}
+
+}