Improve task lifecycle handling and optimize TaskManagerImpl

Redesigned task lifecycle management to properly cascade parent-child removals and synchronize database persistence on task deletion. Refactored task management to address memory leaks, concurrency collisions, and pending-task cleanup issues. Enhanced lookup performance, logging accuracy, and test coverage with a comprehensive unit test suite validating all task scenarios.

.junie/plans/improve-task-manager.md

@@ -0,0 +1,213 @@
+---
+sessionId: session-260919-132150-1vmv
+---
+
+# Requirements
+
+### Overview & Goals
+The `TaskManager` subsystem in Net Neutrality manages the in-memory registry, execution dispatch, locking, persistence, and lifecycle transitions for Discord interaction tasks (buttons, modals, reactions, step-by-step flows, and scheduled background tasks).
+
+This plan addresses several bugs, synchronization oversights, memory leak potentials, and architectural discrepancies between `docs/tasks.md` and `TaskManagerImpl.java`:
+1. **Fix Parent-Child Lifecycle Propagation**: Correctly handle `REMOVE_SELF_AND_CHILDREN` when triggered by child tasks (`WithParent`).
+2. **Synchronize DB Persistence on Removal**: Ensure `PersistableTask` entries are purged from the database via `TaskService.unpersistTask()` when removed.
+3. **Prevent Memory Leaks in Pending Tasks**: Unify `pendingTasks` tracking and ensure `TaskCleanup` cleans up expired pending step tasks.
+4. **Fix Concurrency & Lock Cleanup Hazards**: Eliminate race conditions where lock eviction during task execution allows duplicate concurrent runs.
+5. **Optimize Lookup Performance**: Replace eager full-collection queries with short-circuiting streams for partial ID lookups.
+6. **Improve Test Coverage**: Introduce thorough unit tests covering all task lifecycle paths, locking behavior, and persistence hooks.
+
+### Scope
+- **In Scope**:
+  - `me.bobert.netneutrality.core.tasks.TaskManager` interface.
+  - `me.bobert.netneutrality.core.tasks.TaskManagerImpl` implementation.
+  - `me.bobert.netneutrality.scheduled.cleanup.TaskCleanup` expiration handling.
+  - Interaction deletion error-handling callbacks on `InteractionHook`.
+  - Comprehensive unit test suite for `TaskManagerImpl`.
+- **Out of Scope**:
+  - Refactoring concrete task business logic (e.g., individual music or reaction commands).
+  - Schema changes to the `OTask` entity or database migrations.
+
+### Functional Requirements
+- **FR-1**: When any child task implementing `WithParent` returns `Lifecycle.REMOVE_SELF_AND_CHILDREN`, `TaskManager` must remove the parent task and all associated child tasks from the registry.
+- **FR-2**: When any `PersistableTask` is removed from `TaskManager` (via lifecycle completion, manual removal, or cancellation), `TaskManager` must call `taskService.unpersistTask()` so the task is deleted from persistent storage.
+- **FR-3**: Pending step tasks must be tracked reliably, exposed for expiration checks, and cleaned up by `TaskCleanup` if abandoned.
+- **FR-4**: Concurrent task invocations sharing the same `lockId` must be serialized safely without lock collisions or premature lock disposal.
+- **FR-5**: `findTaskByPartialId` must short-circuit upon finding the first match without allocating full intermediate lists.
+- **FR-6**: Interaction deletion via `Deletable` must gracefully handle expired hooks and Discord API errors asynchronously.
+
+# Technical Design
+
+### Current Implementation
+`TaskManagerImpl` (`me.bobert.netneutrality.core.tasks.TaskManagerImpl`) implements `TaskManager` and `Lifecycle`. It maintains:
+- `tasks`: `ConcurrentHashMap<String, Task>` for active tasks.
+- `pendingTasks`: `CopyOnWriteArrayList<Task>` for unassigned step tasks.
+- `deletableTasksContext`: `ConcurrentHashMap<String, CommandContext>` for tracking interaction hooks to delete.
+- `tasksLock`: `ConcurrentHashMap<String, ReentrantLock>` for per-task/per-parent synchronization.
+- `taskService`: `TaskService` for loading and persisting `PersistableTask` entities.
+
+### Key Issues & Decisions
+
+#### 1. Parent/Child Task Lifecycle Resolution
+- **Issue**: `handleTaskLifecycle` only checks `if (task instanceof WithChildren)` for `REMOVE_SELF_AND_CHILDREN`. When a button task (which implements `WithParent<P>`, such as `LoadTrackFromPlaylistTask`) returns `REMOVE_SELF_AND_CHILDREN`, the check evaluates to `false`, leaving the parent and sibling buttons active in memory.
+- **Decision**: In `handleTaskLifecycle`, check if the task implements `WithParent<?>`. If so, obtain the parent task (`taskWithParent.getParentTask()`) and remove the parent along with all of its child tasks (`parent.getChildTasks()`). If the task directly implements `WithChildren`, remove all its child tasks and itself.
+
+#### 2. Automatic Persistence Synchronization
+- **Issue**: `docs/tasks.md` specifies that persistable tasks removed after execution are deleted from the database. Currently `TaskManagerImpl.removeTask` only removes from `tasks` map and never calls `taskService.unpersistTask(persistableTask)`.
+- **Decision**: In `removeTask(Task task)`, check `if (task instanceof PersistableTask persistableTask)` and invoke `taskService.unpersistTask(persistableTask)`.
+
+#### 3. Pending Task Registry & Cleanup
+- **Issue**: `pendingTasks` is stored in a `CopyOnWriteArrayList<Task>`. `getTasks()` ignores it, preventing `TaskCleanup` from pruning expired step tasks. `removeTask()` also ignores `pendingTasks`.
+- **Decision**: Store `pendingTasks` in a `ConcurrentHashMap<String, Task>` (or unify with active tasks / pending lookup). Update `getTasks()` or provide a method to include pending tasks in cleanup routines. Update `removeTask(Task task)` to remove from `pendingTasks` as well as `tasks`.
+
+#### 4. Concurrency & Lock Management
+- **Issue**: `removeTask` calls `tasksLock.remove(task.getLockId())` immediately. If another thread is waiting or about to acquire the lock for a sibling task, it can get a new lock instance from `computeIfAbsent`, breaking mutual exclusion.
+- **Decision**: Maintain lock instances safely; do not remove locks eagerly while in the middle of executing handlers. Persistable tasks with state should also participate in locking if needed, or maintain clear locking semantics.
+
+#### 5. Lookup API Optimization
+- **Issue**: `findTaskByPartialId` calls `findTasksByPartialId`, instantiating and populating a full `ArrayList` before taking the first element.
+- **Decision**: Implement `findTaskByPartialId` with `.filter(...).findFirst()` directly on the stream.
+
+### Architecture & Component Interaction Diagram
+
+```mermaid
+graph TD
+    subgraph Discord Events
+        BE[ButtonInteractionEvent] --> EH[EventsHandler]
+        ME[ModalInteractionEvent] --> EH
+        RE[MessageReactionAddEvent] --> EH
+        MSG[MessageReceivedEvent] --> EH
+    end
+
+    subgraph TaskManager Subsystem
+        EH -->|TaskContext| TM[TaskManagerImpl]
+        TM -->|Acquire Lock| TL[tasksLock Map]
+        TM -->|Lookup Task| TR[tasks & pendingTasks Maps]
+        TM -->|Execute| T[Task.execute]
+        T -->|Return Lifecycle| TM
+        TM -->|Cascade Cleanup| LC[handleTaskLifecycle]
+        LC -->|WithParent / WithChildren| RM[removeTask]
+        RM -->|Unpersist| TS[TaskService]
+        TS -->|Delete| DB[(Database / OTask)]
+        RM -->|Deletable| DH[InteractionHook.deleteOriginal]
+    end
+
+    subgraph Background Maintenance
+        TC[TaskCleanup] -->|Periodic Poll| TM
+        TM -->|Prune Expired| RM
+    end
+```
+
+### Proposed Code Changes & Contracts
+
+#### `TaskManagerImpl.java`
+```java
+private void handleTaskLifecycle(Task task, Task.Lifecycle status) {
+    if (status == null) {
+        throw new IllegalArgumentException("Task lifecycle status cannot be null");
+    }
+
+    switch (status) {
+        case REMOVE_SELF -> removeTask(task);
+        case REMOVE_SELF_AND_CHILDREN -> {
+            if (task instanceof WithParent<?> taskWithParent) {
+                Task parent = taskWithParent.getParentTask();
+                if (parent instanceof WithChildren parentWithChildren) {
+                    for (Task child : parentWithChildren.getChildTasks()) {
+                        removeTask(child);
+                    }
+                }
+                removeTask(parent);
+            } else if (task instanceof WithChildren taskWithChildren) {
+                for (Task child : taskWithChildren.getChildTasks()) {
+                    removeTask(child);
+                }
+                removeTask(task);
+            } else {
+                removeTask(task);
+            }
+        }
+        case KEEP, PARENT, BUILTIN -> {
+            // Maintained in active registry
+        }
+    }
+}
+
+@Override
+public void removeTask(Task task) {
+    if (task == null) return;
+    String id = task.getId();
+    boolean removed = (tasks.remove(id) != null);
+    removed |= (pendingTasks.remove(task) != null);
+
+    if (removed) {
+        logger.debug("Removing task: {}", id);
+        if (task instanceof PersistableTask persistableTask) {
+            taskService.unpersistTask(persistableTask);
+        }
+    }
+}
+```
+
+### File Structure Changes
+- **Modified**: `src/main/java/me/bobert/netneutrality/core/tasks/TaskManagerImpl.java`
+- **Modified**: `src/main/java/me/bobert/netneutrality/scheduled/cleanup/TaskCleanup.java`
+- **Added**: `src/test/java/me/bobert/netneutrality/core/tasks/TaskManagerImplTest.java`
+
+# Testing
+
+### Validation Approach
+Automated unit tests will be created in `src/test/java/me/bobert/netneutrality/core/tasks/TaskManagerImplTest.java` utilizing JUnit 5, AssertJ, and Mockito to validate all functional requirements without requiring live Discord connections or database instances.
+
+### Key Scenarios
+1. **Task Registration & Lookup**:
+   - Insert standard `Task`, verify `getTask(id)` and `hasTask(id)` return `true`.
+   - Insert duplicate task ID, verify rejected with user notification.
+2. **Parent / Child Lifecycle (`REMOVE_SELF_AND_CHILDREN`)**:
+   - Create parent task implementing `WithChildren` and child task implementing `WithParent`.
+   - Execute child task returning `REMOVE_SELF_AND_CHILDREN`.
+   - Verify parent task and all sibling child tasks are removed from the registry.
+3. **Persistable Task DB Synchronization**:
+   - Insert `PersistableTask`, verify `TaskService.persistTask` is called.
+   - Remove `PersistableTask` via lifecycle completion or `removeTask()`, verify `TaskService.unpersistTask` is invoked.
+4. **Step Task Flow**:
+   - Insert `StepTask` into pending registry.
+   - Reassign ID via `activateTask(task, oldId)`, verify task moves from pending to active registry under new ID.
+5. **Partial ID Search**:
+   - Query `findTaskByPartialId` and verify it matches correct instance without scanning beyond first match.
+6. **Task Expiration & Cleanup**:
+   - Mock expired active and pending tasks, trigger `TaskCleanup.run()`, verify expired tasks are pruned from both collections.
+7. **Concurrency & Locking**:
+   - Simulate concurrent executions for tasks sharing a `lockId`, confirming mutual exclusion.
+
+# Delivery Steps
+
+### ✓ Step 1: Fix Task Lifecycle Hierarchy and DB Persistence Synchronization
+`TaskManagerImpl` correctly propagates parent-child lifecycle actions and synchronizes task deletions with `TaskService`.
+
+- Update `handleTaskLifecycle` in `TaskManagerImpl.java` to support `REMOVE_SELF_AND_CHILDREN` for tasks implementing `WithParent`, resolving the parent task and removing both the parent and all sibling child tasks.
+- Modify `removeTask(Task task)` in `TaskManagerImpl.java` to invoke `taskService.unpersistTask(persistableTask)` when `task instanceof PersistableTask`, ensuring persisted records in the database are purged when tasks finish.
+- Refactor `deleteInteraction(Task task)` in `TaskManagerImpl.java` to cleanly handle `Deletable` logic for both standalone and parented tasks, adding failure logging to `InteractionHook.deleteOriginal().queue()`.
+
+### ✓ Step 2: Refactor Pending Task Management and Concurrency Locking
+Pending tasks and active task locks are managed safely without race conditions or memory leaks.
+
+- Refactor `pendingTasks` in `TaskManagerImpl.java` to use a `ConcurrentHashMap<String, Task>` keyed by task ID for O(1) lookups and consistent state management.
+- Update `removeTask(Task task)`, `getTask(String id)`, and `getTasks()` in `TaskManagerImpl.java` to include pending tasks so they can be retrieved, removed, and cleaned up when expired.
+- Update `TaskCleanup.java` to ensure expired pending tasks are identified and cleaned up.
+- Improve lock management in `TaskManagerImpl.handle(TaskContext context)` so that lock eviction does not interfere with concurrently waiting threads acquiring new lock instances.
+
+### ✓ Step 3: Optimize Task Query Performance and Diagnostic Logging
+Task lookup methods use short-circuiting stream pipelines, and diagnostic logging is corrected.
+
+- Optimize `findTaskByPartialId` in `TaskManagerImpl.java` to use `.findFirst()` directly rather than allocating a full list via `findTasksByPartialId`.
+- Fix typos in log messages and warnings across `TaskManagerImpl.java` (e.g. `"Aquiring lock"`, `"childs"`).
+- Ensure consistent ephemeral fallback messages and handling when unacknowledged interactions or missing tasks are encountered.
+
+### ✓ Step 4: Add Comprehensive Unit Tests for Task Management
+A complete unit test suite validates `TaskManagerImpl` functionality across all task types, lifecycles, and concurrency scenarios.
+
+- Create `src/test/java/me/bobert/netneutrality/core/tasks/TaskManagerImplTest.java` using JUnit 5 and Mockito.
+- Implement tests for standard task registration (`insert`), lookup (`getTask`, `findTaskByPartialId`), and removal (`removeTask`).
+- Implement tests for parent-child tasks (`WithChildren`, `WithParent`), verifying lock inheritance and `REMOVE_SELF_AND_CHILDREN` cascading removals.
+- Implement tests for `StepTask` lifecycle (pending insertion, ID reassignment via `activateTask`, and expiration).
+- Implement tests for `PersistableTask` lifecycle, verifying automatic database persistence on insert and unpersisting on removal.
+- Implement tests verifying thread synchronization and lock acquisition during task execution.

src/main/java/me/bobert/netneutrality/core/tasks/TaskManagerImpl.java

@@ -23,13 +23,13 @@ import net.dv8tion.jda.api.interactions.InteractionHook;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
@@ -41,7 +41,7 @@ public class TaskManagerImpl implements TaskManager {
     private static final Logger logger = LogManager.getLogger(TaskManagerImpl.class);
     private final TaskService taskService;
     private final Map<String, Task> tasks;
-    private final List<Task> pendingTasks;
+    private final Map<String, Task> pendingTasks;
     private final Map<String, CommandContext> deletableTasksContext;
     private final Map<String, ReentrantLock> tasksLock;
     private final ApplicationContext context;
@@ -55,7 +55,7 @@ public class TaskManagerImpl implements TaskManager {
     public TaskManagerImpl(ApplicationContext context) {
         Objects.requireNonNull(context, "ApplicationContext must not be null");
         this.tasks = new ConcurrentHashMap<>();
-        this.pendingTasks = new CopyOnWriteArrayList<>();
+        this.pendingTasks = new ConcurrentHashMap<>();
         this.deletableTasksContext = new ConcurrentHashMap<>();
         this.tasksLock = new ConcurrentHashMap<>();
         this.taskService = new TaskService(this);
@@ -86,7 +86,14 @@ public class TaskManagerImpl implements TaskManager {
 
     @Override
     public Task getTask(String id) {
-        return tasks.get(id);
+        if (id == null) {
+            return null;
+        }
+        Task task = tasks.get(id);
+        if (task != null) {
+            return task;
+        }
+        return pendingTasks.get(id);
     }
 
     /**
@@ -96,12 +103,17 @@ public class TaskManagerImpl implements TaskManager {
      */
     @Override
     public boolean hasTask(String id) {
-        return tasks.containsKey(id) || pendingTasks.stream().anyMatch(task -> task.getId().equals(id));
+        if (id == null) {
+            return false;
+        }
+        return tasks.containsKey(id) || pendingTasks.containsKey(id);
     }
 
     @Override
     public Collection<Task> getTasks() {
-        return tasks.values();
+        List<Task> allTasks = new ArrayList<>(tasks.values());
+        allTasks.addAll(pendingTasks.values());
+        return allTasks;
     }
 
     @Override
@@ -123,7 +135,7 @@ public class TaskManagerImpl implements TaskManager {
             }
 
             if (task instanceof StepTask) {
-                pendingTasks.add(task);
+                pendingTasks.put(id, task);
             } else {
                 tasks.put(id, task);
             }
@@ -136,7 +148,7 @@ public class TaskManagerImpl implements TaskManager {
         if (task instanceof WithChildren taskWithChildren) {
             List<Task> childrens = taskWithChildren.getChildTasks();
             if (childrens == null || childrens.isEmpty()) {
-                logger.warn("Parent task '{}' doesn't create any childs", task.getName());
+                logger.warn("Parent task '{}' doesn't create any children", task.getName());
                 return;
             }
 
@@ -148,23 +160,37 @@ public class TaskManagerImpl implements TaskManager {
 
     @Override
     public void activateTask(Task task, String oldId) {
-        if (pendingTasks.remove(task)) {
+        if (task == null) {
+            return;
+        }
+        boolean removed = oldId != null && pendingTasks.remove(oldId) != null;
+        if (!removed) {
+            removed = pendingTasks.remove(task.getId()) != null;
+        }
+        if (!removed) {
+            pendingTasks.values().remove(task);
+        }
         tasks.put(task.getId(), task);
         logger.info("Updating task {} with new id: {} (was: {})",
                 task.getName(), task.getId(), oldId);
     }
-    }
 
     @Override
     public void removeTask(Task task) {
+        if (task == null) {
+            return;
+        }
         String id = task.getId();
 
-        if (tasks.remove(id) != null) {
+        boolean removed = tasks.remove(id) != null;
+        removed |= pendingTasks.remove(id) != null;
+        removed |= pendingTasks.values().remove(task);
+
+        if (removed) {
             logger.debug("Removing task: {}", id);
 
-            ReentrantLock lock = tasksLock.remove(task.getLockId());
+            if (task instanceof PersistableTask persistableTask) {
-            if (lock != null) {
+                taskService.unpersistTask(persistableTask);
-                logger.debug("Cleaning up lock for removed task: {}", task.getId());
             }
         }
     }
@@ -186,11 +212,11 @@ public class TaskManagerImpl implements TaskManager {
 
             try {
                 if (lock != null) {
-                    logger.debug("Aquiring lock for task: {}", taskId);
+                    logger.debug("Acquiring lock for task: {}", taskId);
                     lock.lock();
 
                     if (!tasks.containsKey(taskId)) {
-                        logger.debug("Task {} is no longer available after aquiring lock", taskId);
+                        logger.debug("Task {} is no longer available after acquiring lock", taskId);
                         return;
                     }
                 }
@@ -224,16 +250,25 @@ public class TaskManagerImpl implements TaskManager {
 
     @Override
     public <E extends Task> Optional<E> findTaskByPartialId(Class<E> clazz, String partialId) {
-        List<E> tasks = findTasksByPartialId(clazz, partialId);
+        if (clazz == null || partialId == null) {
-        return Optional.ofNullable(tasks.isEmpty() ? null : tasks.getFirst());
+            return Optional.empty();
+        }
+        return tasks.values().stream()
+                .filter(clazz::isInstance)
+                .map(clazz::cast)
+                .filter(task -> task.getId() != null && task.getId().contains(partialId))
+                .findFirst();
     }
 
     @Override
     public <E extends Task> List<E> findTasksByPartialId(Class<E> clazz, String partialId) {
+        if (clazz == null || partialId == null) {
+            return List.of();
+        }
         return tasks.values().stream()
                 .filter(clazz::isInstance)
                 .map(clazz::cast)
-            .filter(task -> task.getId().contains(partialId))
+                .filter(task -> task.getId() != null && task.getId().contains(partialId))
                 .collect(Collectors.toList());
     }
 
@@ -272,6 +307,9 @@ public class TaskManagerImpl implements TaskManager {
 
 
     private void deleteInteraction(Task task) {
+        if (task == null) {
+            return;
+        }
         if (task instanceof WithParent<?> taskWithParent) {
             Task parent = taskWithParent.getParentTask();
 
@@ -289,18 +327,16 @@ public class TaskManagerImpl implements TaskManager {
     private void deleteInteraction(String taskId, Deletable deletable) {
         logger.debug("Deleting interaction: {}", taskId);
 
-        CommandContext context = deletableTasksContext.get(taskId);
+        CommandContext commandContext = deletableTasksContext.remove(taskId);
-        if (context != null) {
+        if (commandContext != null) {
-            deletableTasksContext.remove(taskId);
+            InteractionHook hook = commandContext.getInteractionHook();
-
-            InteractionHook hook = context.getInteractionHook();
             if (hook != null) {
                 if (hook.isExpired()) {
-                    context.replyEphemeral("This interaction has expired, can't remove it");
+                    commandContext.replyEphemeral("This interaction has expired, can't remove it");
-                }
+                } else {
-
+                    hook.deleteOriginal().queue(null, error ->
-                if (!hook.isExpired()) {
+                            logger.warn("Failed to delete interaction for task {}: {}", taskId, error.getMessage())
-                    hook.deleteOriginal().queue();
+                    );
                 }
             }
         }
@@ -319,12 +355,32 @@ public class TaskManagerImpl implements TaskManager {
                 break;
 
             case REMOVE_SELF_AND_CHILDREN:
-                if (task instanceof WithChildren taskWithChildren) {
+                if (task instanceof WithParent<?> taskWithParent) {
-                    for (Task child : taskWithChildren.getChildTasks()) {
+                    Task parent = taskWithParent.getParentTask();
+                    if (parent != null) {
+                        if (parent instanceof WithChildren parentWithChildren) {
+                            List<Task> childTasks = parentWithChildren.getChildTasks();
+                            if (childTasks != null) {
+                                for (Task child : childTasks) {
+                                    removeTask(child);
+                                }
+                            }
+                        }
+                        removeTask(parent);
+                    } else {
+                        removeTask(task);
+                    }
+                } else if (task instanceof WithChildren taskWithChildren) {
+                    List<Task> childTasks = taskWithChildren.getChildTasks();
+                    if (childTasks != null) {
+                        for (Task child : childTasks) {
                             removeTask(child);
                         }
                     }
                     removeTask(task);
+                } else {
+                    removeTask(task);
+                }
                 break;
 
             case KEEP:

src/test/java/me/bobert/netneutrality/core/tasks/TaskManagerImplTest.java

@@ -0,0 +1,498 @@
+package me.bobert.netneutrality.core.tasks;
+
+/*
+ * Copyright (c) anthorob
+ *
+ * https://discord.gg/qUa6zNFZsd
+ */
+
+import me.bobert.netneutrality.bootstrap.ApplicationContext;
+import me.bobert.netneutrality.context.tasks.ButtonTaskContext;
+import me.bobert.netneutrality.core.context.CommandContext;
+import me.bobert.netneutrality.core.context.InteractionContext;
+import me.bobert.netneutrality.core.context.TaskContext;
+import me.bobert.netneutrality.core.tasks.impl.AbstractTask;
+import me.bobert.netneutrality.core.tasks.trait.Deletable;
+import me.bobert.netneutrality.core.tasks.trait.WithButton;
+import me.bobert.netneutrality.core.tasks.trait.WithChildren;
+import me.bobert.netneutrality.core.tasks.trait.WithParent;
+import me.bobert.netneutrality.db.service.TaskService;
+import me.bobert.netneutrality.scheduled.cleanup.TaskCleanup;
+import net.dv8tion.jda.api.components.buttons.Button;
+import net.dv8tion.jda.api.interactions.InteractionHook;
+import net.dv8tion.jda.api.requests.RestAction;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+
+import java.lang.reflect.Field;
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Consumer;
+
+import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.*;
+
+@ExtendWith(MockitoExtension.class)
+class TaskManagerImplTest {
+
+    @Mock
+    private ApplicationContext applicationContext;
+
+    @Mock
+    private CommandContext commandContext;
+
+    @Mock
+    private ButtonTaskContext buttonTaskContext;
+
+    @Mock
+    private TaskService taskService;
+
+    private TaskManagerImpl taskManager;
+
+    @BeforeEach
+    void setUp() throws Exception {
+        taskManager = new TaskManagerImpl(applicationContext);
+
+        Field taskServiceField = TaskManagerImpl.class.getDeclaredField("taskService");
+        taskServiceField.setAccessible(true);
+        taskServiceField.set(taskManager, taskService);
+    }
+
+    private static class TestTask extends AbstractTask {
+        private final String id;
+        private final String lockId;
+        private boolean expired;
+        private Lifecycle lifecycleToReturn = Lifecycle.REMOVE_SELF;
+        private Consumer<TaskContext> onExecute;
+
+        TestTask(TaskManager taskManager, String id) {
+            this(taskManager, id, id);
+        }
+
+        TestTask(TaskManager taskManager, String id, String lockId) {
+            super(taskManager);
+            this.id = id;
+            this.lockId = lockId;
+        }
+
+        @Override
+        public String getId() {
+            return id;
+        }
+
+        @Override
+        public String getName() {
+            return "TestTask-" + id;
+        }
+
+        @Override
+        public String getLockId() {
+            return lockId;
+        }
+
+        @Override
+        public boolean isInsertable() {
+            return true;
+        }
+
+        @Override
+        public boolean isExpired() {
+            return expired;
+        }
+
+        public void setExpired(boolean expired) {
+            this.expired = expired;
+        }
+
+        public void setLifecycleToReturn(Lifecycle lifecycleToReturn) {
+            this.lifecycleToReturn = lifecycleToReturn;
+        }
+
+        public void setOnExecute(Consumer<TaskContext> onExecute) {
+            this.onExecute = onExecute;
+        }
+
+        @Override
+        public Lifecycle execute(TaskContext context) {
+            if (onExecute != null) {
+                onExecute.accept(context);
+            }
+            return lifecycleToReturn;
+        }
+
+        @Override
+        public void replyToInteraction(CommandContext context) {
+        }
+    }
+
+    private static class TestParentTask extends TestTask implements WithChildren {
+        private final List<Task> childTasks;
+
+        TestParentTask(TaskManager taskManager, String id, List<Task> childTasks) {
+            super(taskManager, id);
+            this.childTasks = childTasks;
+        }
+
+        @Override
+        public List<Task> getChildTasks() {
+            return childTasks;
+        }
+    }
+
+    private static class TestChildTask extends TestTask implements WithParent<TestParentTask> {
+        private final TestParentTask parent;
+
+        TestChildTask(TaskManager taskManager, String id, TestParentTask parent) {
+            super(taskManager, id, parent != null ? parent.getLockId() : id);
+            this.parent = parent;
+        }
+
+        @Override
+        public TestParentTask getParentTask() {
+            return parent;
+        }
+    }
+
+    private static class TestButtonTask extends TestTask implements WithButton {
+        TestButtonTask(TaskManager taskManager, String id) {
+            super(taskManager, id);
+        }
+
+        @Override
+        public Button createButton() {
+            return Button.primary(getId(), getName());
+        }
+    }
+
+    private static class TestStepTask extends TestTask implements StepTask {
+        TestStepTask(TaskManager taskManager, String id) {
+            super(taskManager, id);
+        }
+
+        @Override
+        public StepType getStepType() {
+            return StepType.STRING;
+        }
+
+        @Override
+        public String getStepQuestion() {
+            return "Question?";
+        }
+
+        @Override
+        public String getStepDescription() {
+            return "Description";
+        }
+
+        @Override
+        public StepResult executeStep(Object answer) {
+            return StepResult.success("Done");
+        }
+    }
+
+    private static class TestPersistableTask extends TestTask implements PersistableTask {
+        TestPersistableTask(TaskManager taskManager, String id) {
+            super(taskManager, id);
+        }
+
+        @Override
+        public void save() {
+        }
+
+        @Override
+        public void delete() {
+        }
+
+        @Override
+        public String toPayload() {
+            return "{}";
+        }
+    }
+
+    private static class TestDeletableTask extends TestTask implements Deletable {
+        private final AtomicBoolean deleted = new AtomicBoolean(false);
+
+        TestDeletableTask(TaskManager taskManager, String id) {
+            super(taskManager, id);
+        }
+
+        @Override
+        public void onInteractionDelete() {
+            deleted.set(true);
+        }
+
+        public boolean isDeleted() {
+            return deleted.get();
+        }
+    }
+
+    @Test
+    void testInsertAndLookup() {
+        TestTask task = new TestTask(taskManager, "task-1");
+
+        taskManager.insert(commandContext, task);
+
+        assertTrue(taskManager.hasTask("task-1"));
+        assertEquals(task, taskManager.getTask("task-1"));
+        assertTrue(taskManager.getTasks().contains(task));
+    }
+
+    @Test
+    void testInsertDuplicateRejected() {
+        TestTask task1 = new TestTask(taskManager, "duplicate-id");
+        TestTask task2 = new TestTask(taskManager, "duplicate-id");
+
+        taskManager.insert(commandContext, task1);
+        taskManager.insert(commandContext, task2);
+
+        verify(commandContext).reply("A task with the same ID already exists!");
+        assertEquals(1, taskManager.getTasks().size());
+    }
+
+    @Test
+    void testRemoveTask() {
+        TestTask task = new TestTask(taskManager, "task-to-remove");
+
+        taskManager.insert(commandContext, task);
+        assertTrue(taskManager.hasTask("task-to-remove"));
+
+        taskManager.removeTask(task);
+        assertFalse(taskManager.hasTask("task-to-remove"));
+        assertNull(taskManager.getTask("task-to-remove"));
+    }
+
+    @Test
+    void testParentAndChildInsertion() {
+        TestChildTask child1 = new TestChildTask(taskManager, "child-1", null);
+        TestChildTask child2 = new TestChildTask(taskManager, "child-2", null);
+        TestParentTask parent = new TestParentTask(taskManager, "parent-1", List.of(child1, child2));
+
+        taskManager.insert(commandContext, parent);
+
+        assertTrue(taskManager.hasTask("parent-1"));
+        assertTrue(taskManager.hasTask("child-1"));
+        assertTrue(taskManager.hasTask("child-2"));
+    }
+
+    @Test
+    void testChildRemoveSelfAndChildrenCascades() {
+        java.util.List<Task> children = new java.util.ArrayList<>();
+        TestParentTask parent = new TestParentTask(taskManager, "parent-task", children);
+        TestChildTask child1 = new TestChildTask(taskManager, "child-task-1", parent);
+        TestChildTask child2 = new TestChildTask(taskManager, "child-task-2", parent);
+        children.add(child1);
+        children.add(child2);
+
+        taskManager.insert(commandContext, parent);
+        assertTrue(taskManager.hasTask("parent-task"));
+        assertTrue(taskManager.hasTask("child-task-1"));
+        assertTrue(taskManager.hasTask("child-task-2"));
+
+        child1.setLifecycleToReturn(Task.Lifecycle.REMOVE_SELF_AND_CHILDREN);
+        when(buttonTaskContext.getId()).thenReturn("child-task-1");
+        when(buttonTaskContext.isAcknowledged()).thenReturn(true);
+
+        taskManager.handle(buttonTaskContext);
+
+        assertFalse(taskManager.hasTask("parent-task"));
+        assertFalse(taskManager.hasTask("child-task-1"));
+        assertFalse(taskManager.hasTask("child-task-2"));
+    }
+
+    @Test
+    void testParentRemoveSelfAndChildrenCascades() {
+        java.util.List<Task> children = new java.util.ArrayList<>();
+        TestParentTask parent = new TestParentTask(taskManager, "parent-task", children);
+        TestChildTask child1 = new TestChildTask(taskManager, "child-task-1", parent);
+        TestChildTask child2 = new TestChildTask(taskManager, "child-task-2", parent);
+        children.add(child1);
+        children.add(child2);
+
+        taskManager.insert(commandContext, parent);
+        parent.setLifecycleToReturn(Task.Lifecycle.REMOVE_SELF_AND_CHILDREN);
+        when(buttonTaskContext.getId()).thenReturn("parent-task");
+        when(buttonTaskContext.isAcknowledged()).thenReturn(true);
+
+        taskManager.handle(buttonTaskContext);
+
+        assertFalse(taskManager.hasTask("parent-task"));
+        assertFalse(taskManager.hasTask("child-task-1"));
+        assertFalse(taskManager.hasTask("child-task-2"));
+    }
+
+    @Test
+    void testChildRemoveSelfRemovesOnlySelf() {
+        java.util.List<Task> children = new java.util.ArrayList<>();
+        TestParentTask parent = new TestParentTask(taskManager, "parent-task", children);
+        TestChildTask child1 = new TestChildTask(taskManager, "child-task-1", parent);
+        TestChildTask child2 = new TestChildTask(taskManager, "child-task-2", parent);
+        children.add(child1);
+        children.add(child2);
+
+        taskManager.insert(commandContext, parent);
+        child1.setLifecycleToReturn(Task.Lifecycle.REMOVE_SELF);
+        when(buttonTaskContext.getId()).thenReturn("child-task-1");
+        when(buttonTaskContext.isAcknowledged()).thenReturn(true);
+
+        taskManager.handle(buttonTaskContext);
+
+        assertFalse(taskManager.hasTask("child-task-1"));
+        assertTrue(taskManager.hasTask("parent-task"));
+        assertTrue(taskManager.hasTask("child-task-2"));
+    }
+
+    @Test
+    void testStepTaskPendingAndActivation() {
+        TestStepTask stepTask = new TestStepTask(taskManager, "initial-uuid");
+
+        taskManager.insert(commandContext, stepTask);
+
+        assertTrue(taskManager.hasTask("initial-uuid"));
+        assertEquals(stepTask, taskManager.getTask("initial-uuid"));
+        assertTrue(taskManager.getTasks().contains(stepTask));
+
+        TestStepTask activatedTask = new TestStepTask(taskManager, "message-id-123");
+        taskManager.activateTask(activatedTask, "initial-uuid");
+
+        assertTrue(taskManager.hasTask("message-id-123"));
+        assertEquals(activatedTask, taskManager.getTask("message-id-123"));
+        assertFalse(taskManager.hasTask("initial-uuid"));
+    }
+
+    @Test
+    void testTaskCleanupPrunesExpiredStepTasks() {
+        TestStepTask expiredStepTask = new TestStepTask(taskManager, "expired-step");
+        expiredStepTask.setExpired(true);
+
+        TestStepTask activeStepTask = new TestStepTask(taskManager, "active-step");
+        activeStepTask.setExpired(false);
+
+        taskManager.insert(commandContext, expiredStepTask);
+        taskManager.insert(commandContext, activeStepTask);
+
+        TaskCleanup cleanup = new TaskCleanup(taskManager);
+        cleanup.run();
+
+        assertFalse(taskManager.hasTask("expired-step"));
+        assertTrue(taskManager.hasTask("active-step"));
+    }
+
+    @Test
+    void testPersistableTaskUnpersistedOnRemoval() {
+        TestPersistableTask persistableTask = new TestPersistableTask(taskManager, "persist-1");
+
+        taskManager.insert(commandContext, persistableTask);
+        verify(taskService).persistTask(persistableTask);
+
+        taskManager.removeTask(persistableTask);
+        verify(taskService).unpersistTask(persistableTask);
+    }
+
+    @Test
+    void testFindTaskByPartialId() {
+        TestButtonTask button1 = new TestButtonTask(taskManager, "btn_action_play");
+        TestButtonTask button2 = new TestButtonTask(taskManager, "btn_action_pause");
+        TestTask regularTask = new TestTask(taskManager, "regular_play");
+
+        taskManager.insert(commandContext, button1);
+        taskManager.insert(commandContext, button2);
+        taskManager.insert(commandContext, regularTask);
+
+        Optional<TestButtonTask> found = taskManager.findTaskByPartialId(TestButtonTask.class, "action_play");
+        assertTrue(found.isPresent());
+        assertEquals(button1, found.get());
+
+        List<TestButtonTask> allButtons = taskManager.findTasksByPartialId(TestButtonTask.class, "btn_action");
+        assertEquals(2, allButtons.size());
+        assertTrue(allButtons.contains(button1));
+        assertTrue(allButtons.contains(button2));
+
+        WithButton foundButton = taskManager.findButtonTaskByClass(TestButtonTask.class);
+        assertNotNull(foundButton);
+    }
+
+    @Test
+    void testConcurrentExecutionSharingLock() throws Exception {
+        int threadCount = 4;
+        ExecutorService executor = Executors.newFixedThreadPool(threadCount);
+        CountDownLatch startLatch = new CountDownLatch(1);
+        CountDownLatch finishLatch = new CountDownLatch(threadCount);
+
+        AtomicInteger concurrentExecutions = new AtomicInteger(0);
+        AtomicInteger maxConcurrentSeen = new AtomicInteger(0);
+
+        String sharedLockId = "shared-lock-123";
+
+        for (int i = 0; i < threadCount; i++) {
+            String taskId = "task-lock-" + i;
+            TestTask task = new TestTask(taskManager, taskId, sharedLockId);
+            task.setLifecycleToReturn(Task.Lifecycle.KEEP);
+            task.setOnExecute(ctx -> {
+                int active = concurrentExecutions.incrementAndGet();
+                maxConcurrentSeen.updateAndGet(curr -> Math.max(curr, active));
+                try {
+                    Thread.sleep(50);
+                } catch (InterruptedException ignored) {
+                }
+                concurrentExecutions.decrementAndGet();
+            });
+
+            taskManager.insert(commandContext, task);
+
+            executor.submit(() -> {
+                try {
+                    startLatch.await();
+                    TaskContext mockContext = mock(TaskContext.class);
+                    when(mockContext.getId()).thenReturn(taskId);
+                    when(mockContext.isAcknowledged()).thenReturn(true);
+                    taskManager.handle(mockContext);
+                } catch (Exception e) {
+                    e.printStackTrace();
+                } finally {
+                    finishLatch.countDown();
+                }
+            });
+        }
+
+        startLatch.countDown();
+        boolean completed = finishLatch.await(5, TimeUnit.SECONDS);
+        executor.shutdown();
+
+        assertTrue(completed);
+        assertEquals(1, maxConcurrentSeen.get(), "Tasks sharing a lockId must be executed mutually exclusively");
+    }
+
+    @Test
+    void testDeletableInteractionExecution() {
+        TestDeletableTask deletableTask = new TestDeletableTask(taskManager, "deletable-1");
+        deletableTask.setLifecycleToReturn(Task.Lifecycle.REMOVE_SELF);
+
+        InteractionHook mockHook = mock(InteractionHook.class);
+        when(mockHook.isExpired()).thenReturn(false);
+        @SuppressWarnings("unchecked")
+        RestAction<Void> restAction = mock(RestAction.class);
+        when(mockHook.deleteOriginal()).thenReturn(restAction);
+
+        when(commandContext.getInteractionHook()).thenReturn(mockHook);
+        taskManager.insert(commandContext, deletableTask);
+
+        when(buttonTaskContext.getId()).thenReturn("deletable-1");
+        when(buttonTaskContext.isAcknowledged()).thenReturn(true);
+
+        taskManager.handle(buttonTaskContext);
+
+        assertTrue(deletableTask.isDeleted());
+        verify(mockHook).deleteOriginal();
+        verify(restAction).queue(any(), any());
+    }
+}