fix(statemachine): re-initialise StateMachineScript on restart

runelite-client/src/main/java/net/runelite/client/plugins/microbot/statemachine/StateMachineScript.java

@@ -73,7 +73,8 @@ public static Map<String, StateMachineScript<?>> getRegistry() {
     private long loopCount;
     private long transitionCount;
     private List<Transition<S>> transitions;
-    private boolean initialized;
+    private volatile boolean initialized;
+    private final Object lifecycleLock = new Object();
 
     /** Bounded ring buffer of recent transitions for the trace log. */
     private static final int MAX_TRACE_SIZE = 50;
@@ -136,8 +137,8 @@ protected S onError(S state, Exception e) {
      * @return false if the base guard paused execution, true otherwise
      */
     public final boolean step() {
-        if (!initialized) {
-            initialize();
+        if (!initialized && !initialize()) {
+            return false;
         }
 
         // Check the base Script guard (heartbeat, blocking events, pause, etc.)
@@ -213,13 +214,25 @@ public final boolean step() {
         return true;
     }
 
-    private void initialize() {
-        currentState = initialState();
-        stateEnteredAt = Instant.now();
-        transitions = Collections.unmodifiableList(defineTransitions());
-        initialized = true;
-        REGISTRY.put(getScriptName(), this);
+    private boolean initialize() {
+        synchronized (lifecycleLock) {
+            if (Thread.currentThread().isInterrupted()) {
+                return false;
+            }
+            currentState = initialState();
+            previousState = null;
+            lastTransitionReason = null;
+            lastTransitionAt = null;
+            loopCount = 0;
+            transitionCount = 0;
+            traceBuffer.clear();
+            stateEnteredAt = Instant.now();
+            transitions = Collections.unmodifiableList(defineTransitions());
+            initialized = true;
+            REGISTRY.put(getScriptName(), this);
+        }
         log.info("[{}] State machine initialized in state {}", getScriptName(), currentState);
+        return true;
     }
 
     /**
@@ -245,13 +258,18 @@ private void publishSnapshot(boolean allGuardsFalse) {
             }
         }
 
-        snapshot.set(new StateSnapshot<>(
+        StateSnapshot<S> next = new StateSnapshot<>(
                 currentState, previousState, lastTransitionReason,
                 lastTransitionAt, stateEnteredAt,
                 loopCount, transitionCount,
                 pending,
                 new ArrayList<>(traceBuffer)
-        ));
+        );
+        synchronized (lifecycleLock) {
+            if (initialized) {
+                snapshot.set(next);
+            }
+        }
     }
 
     /**
@@ -287,8 +305,12 @@ public void forceState(S newState, String reason) {
 
     @Override
     public void shutdown() {
-        REGISTRY.remove(getScriptName());
         super.shutdown();
+        synchronized (lifecycleLock) {
+            initialized = false;
+            REGISTRY.remove(getScriptName(), this);
+            snapshot.set(null);
+        }
     }
 
     private String getScriptName() {

runelite-client/src/test/java/net/runelite/client/plugins/microbot/statemachine/StateMachineScriptRestartTest.java

@@ -0,0 +1,128 @@
+package net.runelite.client.plugins.microbot.statemachine;
+
+import java.util.ArrayList;
+import java.util.List;
+import net.runelite.client.plugins.microbot.Microbot;
+import org.junit.Test;
+import org.mockito.MockedStatic;
+import static org.junit.Assert.*;
+import static org.mockito.Mockito.*;
+
+public class StateMachineScriptRestartTest {
+    private static final String NAME = Machine.class.getSimpleName();
+
+    @Test public void restartReRegistersAndBeginsFromInitialState() {
+        Machine machine = new Machine();
+        try (MockedStatic<Microbot> ignored = mockStatic(Microbot.class)) {
+            machine.advance = true;
+            assertTrue(machine.step());
+            assertEquals(Machine.State.WORKING, machine.getCurrentState());
+            assertSame(machine, StateMachineScript.getRegistry().get(NAME));
+
+            machine.shutdown();
+            assertFalse(StateMachineScript.getRegistry().containsKey(NAME));
+            assertNull(machine.getSnapshot());
+
+            machine.advance = false;
+            assertTrue(machine.step());
+            assertEquals(Machine.State.IDLE, machine.getCurrentState());
+            assertSame(machine, StateMachineScript.getRegistry().get(NAME));
+            StateSnapshot<Machine.State> snapshot = machine.getSnapshot();
+            assertEquals(Machine.State.IDLE, snapshot.currentState());
+            assertNull(snapshot.previousState());
+            assertEquals(1, snapshot.loopCount());
+            assertEquals(0, snapshot.transitionCount());
+            assertTrue(snapshot.recentTrace().isEmpty());
+            assertEquals(List.of(Machine.State.WORKING, Machine.State.IDLE), machine.visited);
+
+            machine.shutdown();
+            assertFalse(StateMachineScript.getRegistry().containsKey(NAME));
+        } finally {
+            machine.shutdown();
+            machine.dispose();
+        }
+    }
+
+    @Test public void staleInstanceStopDoesNotUnregisterItsReplacement() {
+        Machine first = new Machine();
+        Machine second = new Machine();
+        try (MockedStatic<Microbot> ignored = mockStatic(Microbot.class)) {
+            first.step();
+            second.step();
+            first.shutdown();
+            assertSame(second, StateMachineScript.getRegistry().get(NAME));
+            second.shutdown();
+            assertFalse(StateMachineScript.getRegistry().containsKey(NAME));
+        } finally {
+            first.shutdown();
+            second.shutdown();
+            first.dispose();
+            second.dispose();
+        }
+    }
+
+    @Test public void cancelledStepAfterStopDoesNotReRegister() {
+        Machine machine = new Machine();
+        try (MockedStatic<Microbot> ignored = mockStatic(Microbot.class)) {
+            machine.step();
+            machine.shutdown();
+            Thread.currentThread().interrupt();
+            try {
+                assertFalse(machine.step());
+            } finally {
+                Thread.interrupted();
+            }
+            assertFalse(StateMachineScript.getRegistry().containsKey(NAME));
+            assertNull(machine.getSnapshot());
+        } finally {
+            machine.shutdown();
+            machine.dispose();
+        }
+    }
+
+    @Test public void stopDuringStateActionDoesNotRepublishSnapshot() {
+        Machine machine = new Machine();
+        try (MockedStatic<Microbot> ignored = mockStatic(Microbot.class)) {
+            machine.stopInAction = true;
+            machine.step();
+            assertFalse(StateMachineScript.getRegistry().containsKey(NAME));
+            assertNull(machine.getSnapshot());
+        } finally {
+            machine.shutdown();
+            machine.dispose();
+        }
+    }
+
+    private static final class Machine extends StateMachineScript<Machine.State> {
+        enum State { IDLE, WORKING }
+
+        boolean advance;
+        boolean stopInAction;
+        final List<State> visited = new ArrayList<>();
+
+        @Override protected State onError(State state, Exception e) {
+            return State.WORKING;
+        }
+
+        @Override protected State initialState() {
+            return State.IDLE;
+        }
+
+        @Override protected List<Transition<State>> defineTransitions() {
+            return List.of(Transition.from(State.IDLE).when(() -> advance, "advance")
+                    .because("advance").goTo(State.WORKING));
+        }
+
+        @Override protected void onState(State state) {
+            visited.add(state);
+            if (stopInAction) {
+                shutdown();
+                throw new IllegalStateException("stopped");
+            }
+        }
+
+        void dispose() {
+            scheduledExecutorService.shutdownNow();
+        }
+    }
+}