diff --git a/argi-orchestration/argi-workflow-graph-core/src/main/java/io/github/agentic/ai/workflow/graph/scheduling/ScheduledAgentTask.java b/argi-orchestration/argi-workflow-graph-core/src/main/java/io/github/agentic/ai/workflow/graph/scheduling/ScheduledAgentTask.java index e2c6b43ce..728891dc6 100644 --- a/argi-orchestration/argi-workflow-graph-core/src/main/java/io/github/agentic/ai/workflow/graph/scheduling/ScheduledAgentTask.java +++ b/argi-orchestration/argi-workflow-graph-core/src/main/java/io/github/agentic/ai/workflow/graph/scheduling/ScheduledAgentTask.java @@ -41,6 +41,8 @@ public class ScheduledAgentTask { private static final Logger log = LoggerFactory.getLogger(CompiledGraph.class); + private final ScheduledAgentManager registrationManager; + private final TaskScheduler taskScheduler; private final CompiledGraph graph; @@ -60,10 +62,10 @@ public class ScheduledAgentTask { public ScheduledAgentTask(CompiledGraph graph, ScheduleConfig config) { this.graph = graph; this.config = config; - ScheduledAgentManager scheduledAgentManager = ScheduledAgentManagerFactory.getInstance().getManager(); - this.taskScheduler = scheduledAgentManager.getTaskScheduler(); - // Register with the active manager - this.taskId = scheduledAgentManager.registerTask(this); + this.registrationManager = ScheduledAgentManagerFactory.getInstance().getManager(); + this.taskScheduler = registrationManager.getTaskScheduler(); + // Retain the registration owner even if the factory's active manager changes. + this.taskId = registrationManager.registerTask(this); log.debug("Created ScheduledAgentTask with ID: {}", taskId); } @@ -113,7 +115,7 @@ public void stop() { } stopped = true; - // Unregister from active manager + // Unregister from the manager that owns this task. unregisterTask(); log.debug("Stopped and unregistered ScheduledAgentTask with ID: {}", taskId); @@ -214,7 +216,7 @@ private void notifyListeners(ScheduleLifecycleListener.ScheduleEvent event) { private void unregisterTask() { if (unregistered.compareAndSet(false, true)) { - ScheduledAgentManagerFactory.getInstance().getManager().unregisterTask(taskId); + registrationManager.unregisterTask(taskId); } } @@ -228,7 +230,7 @@ public boolean isStopped() { } /** - * Get the unique task ID assigned by the global manager + * Get the unique task ID assigned by the registration manager */ public String getTaskId() { return taskId; diff --git a/argi-orchestration/argi-workflow-graph-core/src/test/java/io/github/agentic/ai/workflow/graph/scheduling/ScheduledAgentTaskLifecycleTest.java b/argi-orchestration/argi-workflow-graph-core/src/test/java/io/github/agentic/ai/workflow/graph/scheduling/ScheduledAgentTaskLifecycleTest.java index b4b7684cd..c0db97376 100644 --- a/argi-orchestration/argi-workflow-graph-core/src/test/java/io/github/agentic/ai/workflow/graph/scheduling/ScheduledAgentTaskLifecycleTest.java +++ b/argi-orchestration/argi-workflow-graph-core/src/test/java/io/github/agentic/ai/workflow/graph/scheduling/ScheduledAgentTaskLifecycleTest.java @@ -246,7 +246,7 @@ private void registerRecordingManager() { ScheduledAgentManagerFactory.getInstance().registerProvider(() -> manager); } - private CompiledGraph graphFailingFirstAttempts(AtomicInteger attempts, int failures) throws Exception { + static CompiledGraph graphFailingFirstAttempts(AtomicInteger attempts, int failures) throws Exception { Map strategies = new HashMap<>(); strategies.put("result", (left, right) -> right); return new StateGraph("schedule-lifecycle-test", () -> strategies) @@ -276,9 +276,9 @@ private long count(CopyOnWriteArrayList return events.stream().filter(event::equals).count(); } - private static final class RecordingScheduledAgentManager implements ScheduledAgentManager { + static final class RecordingScheduledAgentManager implements ScheduledAgentManager { - private final RecordingTaskScheduler scheduler = new RecordingTaskScheduler(); + final RecordingTaskScheduler scheduler = new RecordingTaskScheduler(); private final Map tasks = new HashMap<>(); @@ -286,8 +286,15 @@ private static final class RecordingScheduledAgentManager implements ScheduledAg private final AtomicInteger sequence = new AtomicInteger(); + private int shutdownCount; + + private boolean rejectRegistrations; + @Override public String registerTask(ScheduledAgentTask task) { + if (rejectRegistrations) { + throw new IllegalStateException("Task registration rejected"); + } String taskId = "test-task-" + sequence.incrementAndGet(); tasks.put(taskId, task); return taskId; @@ -321,14 +328,23 @@ public TaskScheduler getTaskScheduler() { @Override public boolean isShutdown() { - return false; + return shutdownCount > 0; } @Override public void shutdown() { + shutdownCount++; tasks.clear(); } + void rejectRegistrations() { + rejectRegistrations = true; + } + + int shutdownCount() { + return shutdownCount; + } + int unregisterCount(String taskId) { AtomicInteger count = unregisterCounts.get(taskId); return count != null ? count.get() : 0; @@ -336,7 +352,7 @@ int unregisterCount(String taskId) { } - private static final class RecordingTaskScheduler implements TaskScheduler { + static final class RecordingTaskScheduler implements TaskScheduler { private final Queue scheduledTasks = new ArrayDeque<>(); diff --git a/argi-orchestration/argi-workflow-graph-core/src/test/java/io/github/agentic/ai/workflow/graph/scheduling/ScheduledAgentTaskManagerOwnershipTest.java b/argi-orchestration/argi-workflow-graph-core/src/test/java/io/github/agentic/ai/workflow/graph/scheduling/ScheduledAgentTaskManagerOwnershipTest.java new file mode 100644 index 000000000..e9b901d72 --- /dev/null +++ b/argi-orchestration/argi-workflow-graph-core/src/test/java/io/github/agentic/ai/workflow/graph/scheduling/ScheduledAgentTaskManagerOwnershipTest.java @@ -0,0 +1,158 @@ +/* + * 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.concurrent.atomic.AtomicInteger; + +import io.github.agentic.ai.workflow.graph.scheduling.ScheduledAgentTaskLifecycleTest.RecordingScheduledAgentManager; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +import static io.github.agentic.ai.workflow.graph.scheduling.ScheduleLifecycleListener.ScheduleEvent.EXECUTION_STARTED; +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.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** 复现 #214：切换全局管理器后，任务仍应向注册它的管理器注销。 */ +class ScheduledAgentTaskManagerOwnershipTest { + + private final ScheduledAgentManagerFactory factory = ScheduledAgentManagerFactory.getInstance(); + + private final ScheduledAgentManager originalManager = factory.getManager(); + + private final RecordingScheduledAgentManager owner = new RecordingScheduledAgentManager(); + + private final RecordingScheduledAgentManager replacement = new RecordingScheduledAgentManager(); + + /** 恢复测试前的管理器，避免全局 provider 影响其他测试。 */ + @AfterEach + void restoreOriginalManager() { + factory.registerProvider(() -> originalManager); + } + + /** 验证切换管理器后重复停止任务，只向原管理器注销一次。 */ + @Test + void repeatedStopAfterProviderSwitchUnregistersOnlyFromOwner() throws Exception { + factory.registerProvider(() -> owner); + final ScheduledAgentTask task = new ScheduledAgentTask(graphFailingFirstAttempts(new AtomicInteger(), 0), + ScheduleConfig.builder().fixedDelay(1000).initialDelay(0).build()).start(); + assertSame(task, owner.getTask(task.getTaskId()).orElseThrow()); + factory.registerProvider(() -> replacement); + + task.stop(); + task.stop(); + + assertTrue(task.isStopped()); + assertUnregisteredOnlyFromOwner(task); + } + + /** 验证一次性任务完成以及随后重复停止，都不会向新管理器注销。 */ + @Test + void oneTimeCompletionAfterProviderSwitchUnregistersOnlyFromOwner() throws Exception { + factory.registerProvider(() -> owner); + final AtomicInteger executions = new AtomicInteger(); + final ScheduledAgentTask task = new ScheduledAgentTask(graphFailingFirstAttempts(executions, 0), + ScheduleConfig.builder().initialDelay(0).build()).start(); + factory.registerProvider(() -> replacement); + + owner.scheduler.runNext(); + task.stop(); + task.stop(); + + assertEquals(1, executions.get()); + assertUnregisteredOnlyFromOwner(task); + } + + /** 验证任务在 start 返回前立即执行并切换 provider 时，仍保留原管理器归属。 */ + @Test + void immediateCompletionDuringProviderSwitchUnregistersOnlyFromOwner() throws Exception { + factory.registerProvider(() -> owner); + owner.scheduler.runOneTimeTasksImmediately(); + final AtomicInteger executions = new AtomicInteger(); + final ScheduledAgentTask task = new ScheduledAgentTask(graphFailingFirstAttempts(executions, 0), + ScheduleConfig.builder().initialDelay(0).addListener(new ScheduleLifecycleListener() { + @Override + public void onEvent(ScheduleEvent event, Object data) { + if (event == EXECUTION_STARTED) { + factory.registerProvider(() -> replacement); + } + } + }).build()).start(); + + task.stop(); + + assertEquals(1, executions.get()); + assertSame(replacement, factory.getManager()); + assertUnregisteredOnlyFromOwner(task); + } + + /** 验证两个管理器分配相同 taskId 时，停止旧任务不会误删新管理器中的任务。 */ + @Test + void stoppingOldTaskDoesNotRemoveReplacementTaskWithSameId() throws Exception { + factory.registerProvider(() -> owner); + final ScheduledAgentTask oldTask = new ScheduledAgentTask(graphFailingFirstAttempts(new AtomicInteger(), 0), + ScheduleConfig.builder().initialDelay(0).build()); + factory.registerProvider(() -> replacement); + final ScheduledAgentTask newTask = new ScheduledAgentTask(graphFailingFirstAttempts(new AtomicInteger(), 0), + ScheduleConfig.builder().initialDelay(0).build()); + assertEquals(oldTask.getTaskId(), newTask.getTaskId()); + assertSame(newTask, replacement.getTask(newTask.getTaskId()).orElseThrow()); + + oldTask.stop(); + oldTask.stop(); + + assertAll(() -> assertUnregisteredOnlyFromOwner(oldTask), + () -> assertSame(newTask, replacement.getTask(newTask.getTaskId()).orElse(null), + "Stopping the old task must preserve the replacement manager's task")); + } + + /** 验证注册失败不会遗留任务、关闭管理器，或影响新管理器后续注册和停止任务。 */ + @Test + void registrationFailureDoesNotAffectReplacementManager() throws Exception { + factory.registerProvider(() -> owner); + owner.rejectRegistrations(); + final var graph = graphFailingFirstAttempts(new AtomicInteger(), 0); + assertThrows(IllegalStateException.class, + () -> new ScheduledAgentTask(graph, ScheduleConfig.builder().initialDelay(0).build())); + assertEquals(0, owner.getActiveTaskCount()); + factory.registerProvider(() -> replacement); + final ScheduledAgentTask task = new ScheduledAgentTask(graph, + ScheduleConfig.builder().initialDelay(0).build()); + assertSame(task, replacement.getTask(task.getTaskId()).orElseThrow()); + + task.stop(); + task.stop(); + + assertAll(() -> assertEquals(0, owner.getActiveTaskCount()), + () -> assertEquals(0, owner.unregisterCount(task.getTaskId())), + () -> assertEquals(0, replacement.getActiveTaskCount()), + () -> assertEquals(1, replacement.unregisterCount(task.getTaskId())), + () -> assertEquals(0, owner.shutdownCount()), () -> assertEquals(0, replacement.shutdownCount())); + } + + /** 检查原管理器已清理任务、新管理器没有收到注销请求，且外部管理器未被关闭。 */ + private void assertUnregisteredOnlyFromOwner(ScheduledAgentTask task) { + assertAll(() -> assertTrue(owner.getTask(task.getTaskId()).isEmpty(), "Owner must release the task"), + () -> assertEquals(1, owner.unregisterCount(task.getTaskId()), "Owner must unregister exactly once"), + () -> assertEquals(0, replacement.unregisterCount(task.getTaskId()), + "Replacement manager must not receive the old task's unregister request"), + () -> assertEquals(0, owner.shutdownCount()), () -> assertEquals(0, replacement.shutdownCount())); + } + +}