[cmd3] Scope scheduled commands to the running opmode, if one exists (#8492)

This prevents commands from outliving the opmodes in which they were
scheduled

commandsv3/src/main/java/org/wpilib/command3/BindingScope.java

@@ -25,6 +25,14 @@ interface BindingScope {
     return new ForCommand(scheduler, command);
   }
 
+  static BindingScope forOpmode(long opmodeId) {
+    if (opmodeId == 0) {
+      throw new IllegalArgumentException("Invalid OpMode ID provided");
+    }
+
+    return new ForOpmode(opmodeId);
+  }
+
   /** A global binding scope. Bindings in this scope are always active. */
   final class Global implements BindingScope {
     // No reason not to be a singleton.
@@ -49,4 +57,16 @@ interface BindingScope {
       return scheduler.isRunning(command);
     }
   }
+
+  /**
+   * A binding scoped to a running opmode.
+   *
+   * @param opmodeId The ID of the opmode that the binding is scoped to.
+   */
+  record ForOpmode(long opmodeId) implements BindingScope {
+    @Override
+    public boolean active() {
+      return OpModeFetcher.getFetcher().getOpModeId() == opmodeId;
+    }
+  }
 }

commandsv3/src/main/java/org/wpilib/command3/OpModeFetcher.java

@@ -0,0 +1,63 @@
+// Copyright (c) FIRST and other WPILib contributors.
+// Open Source Software; you can modify and/or share it under the terms of
+// the WPILib BSD license file in the root directory of this project.
+
+package org.wpilib.command3;
+
+import static org.wpilib.util.ErrorMessages.requireNonNullParam;
+
+import org.wpilib.driverstation.DriverStation;
+
+/**
+ * Helper class for fetching information about the current opmode. This is a package-private class
+ * so tests for this library don't need to hook into driverstation simulation and the HAL.
+ */
+abstract class OpModeFetcher {
+  private static volatile OpModeFetcher s_fetcher;
+
+  abstract long getOpModeId();
+
+  abstract String getOpModeName();
+
+  /**
+   * Gets the current fetcher implementation. If {@link #setFetcher(OpModeFetcher)} has not already
+   * been called to set an implementation, this will default to a {@link DriverStationOpModeFetcher}
+   * instance.
+   *
+   * @return The fetcher instance to use to get opmode information.
+   */
+  static OpModeFetcher getFetcher() {
+    // Default to pull from the DS unless otherwise specified
+    if (s_fetcher == null) {
+      synchronized (OpModeFetcher.class) {
+        if (s_fetcher == null) {
+          s_fetcher = new DriverStationOpModeFetcher();
+        }
+      }
+    }
+
+    return s_fetcher;
+  }
+
+  /**
+   * Sets the fetcher to use. In tests, this is reset before every test with an implementation that
+   * always returns an ID of 0 and an empty string for the opmode name.
+   *
+   * @param fetcher The fetcher implementation to set. Cannot be null.
+   */
+  static void setFetcher(OpModeFetcher fetcher) {
+    s_fetcher = requireNonNullParam(fetcher, "fetcher", "setFetcher");
+  }
+
+  static final class DriverStationOpModeFetcher extends OpModeFetcher {
+    @Override
+    public long getOpModeId() {
+      return DriverStation.getOpModeId();
+    }
+
+    @Override
+    public String getOpModeName() {
+      return DriverStation.getOpMode();
+    }
+  }
+}

commandsv3/src/main/java/org/wpilib/command3/Scheduler.java

@@ -271,11 +271,26 @@ public final class Scheduler implements ProtobufSerializable {
   // innermost commands that actually _do_ something to start running hundreds of milliseconds after
   // their root ancestor was scheduled.
   public ScheduleResult schedule(Command command) {
+    // Get the narrowest binding scope.
+    // This prevents commands from outliving the opmodes that scheduled them, or from outliving
+    // their parents (eg if someone writes a command that manually calls schedule(Command) instead
+    // of using triggers to do so).
+    Command currentCommand = currentCommand();
+    long currentOpmode = OpModeFetcher.getFetcher().getOpModeId();
+
+    BindingScope scope;
+    if (currentCommand != null) {
+      scope = BindingScope.forCommand(this, currentCommand);
+    } else if (currentOpmode != 0) {
+      scope = BindingScope.forOpmode(currentOpmode);
+    } else {
+      scope = BindingScope.global();
+    }
+
     // Note: we use a throwable here instead of Thread.currentThread().getStackTrace() for easier
     //       stack frame filtering and modification.
     var binding =
-        new Binding(
-            BindingScope.global(), BindingType.IMMEDIATE, command, new Throwable().getStackTrace());
+        new Binding(scope, BindingType.IMMEDIATE, command, new Throwable().getStackTrace());
 
     return schedule(binding);
   }

commandsv3/src/main/java/org/wpilib/command3/Trigger.java

@@ -351,17 +351,22 @@ public class Trigger implements BooleanSupplier {
   }
 
   private void addBinding(BindingType bindingType, Command command) {
-    BindingScope scope =
-        switch (m_scheduler.currentCommand()) {
-          case Command c -> {
-            // A command is creating a binding - make it scoped to that specific command
-            yield BindingScope.forCommand(m_scheduler, c);
-          }
-          case null -> {
-            // Creating a binding outside a command - it's global in scope
-            yield BindingScope.global();
-          }
-        };
+    Command currentCommand = m_scheduler.currentCommand();
+    long currentOpmode = OpModeFetcher.getFetcher().getOpModeId();
+
+    BindingScope scope;
+    if (currentCommand != null) {
+      // A command is creating a binding - make it scoped to that specific command.
+      // The binding will be removed when the command exits.
+      scope = BindingScope.forCommand(m_scheduler, currentCommand);
+    } else if (currentOpmode != 0) {
+      // An opmode is currently running; scope the binding to just that mode.
+      // The binding will be removed when the opmode exits.
+      scope = BindingScope.forOpmode(currentOpmode);
+    } else {
+      // No opmode selected and no command is running; the binding is global in scope.
+      scope = BindingScope.global();
+    }
 
     addBinding(scope, bindingType, command);
   }

commandsv3/src/test/java/org/wpilib/command3/MockHardwareExtension.java

@@ -0,0 +1,27 @@
+// Copyright (c) FIRST and other WPILib contributors.
+// Open Source Software; you can modify and/or share it under the terms of
+// the WPILib BSD license file in the root directory of this project.
+
+package org.wpilib.command3;
+
+import org.junit.jupiter.api.extension.BeforeEachCallback;
+import org.junit.jupiter.api.extension.ExtensionContext;
+
+public class MockHardwareExtension implements BeforeEachCallback {
+  @Override
+  public void beforeEach(ExtensionContext context) {
+    OpModeFetcher.setFetcher(new MockOpModeFetcher());
+  }
+
+  private static final class MockOpModeFetcher extends OpModeFetcher {
+    @Override
+    long getOpModeId() {
+      return 0;
+    }
+
+    @Override
+    String getOpModeName() {
+      return "";
+    }
+  }
+}

commandsv3/src/test/java/org/wpilib/command3/TriggerTest.java

@@ -184,6 +184,43 @@ class TriggerTest extends CommandTestBase {
         "Command should have been canceled when scope became inactive");
   }
 
+  @Test
+  void bindingScopesToOpmodeIfAvailable() {
+    var fetcher =
+        new OpModeFetcher() {
+          long m_id = 12345;
+
+          void clear() {
+            m_id = 0;
+          }
+
+          @Override
+          long getOpModeId() {
+            return m_id;
+          }
+
+          @Override
+          String getOpModeName() {
+            return "This is an opmode!";
+          }
+        };
+    OpModeFetcher.setFetcher(fetcher);
+
+    var triggerSignal = new AtomicBoolean(false);
+    var trigger = new Trigger(m_scheduler, triggerSignal::get);
+
+    var command = Command.noRequirements().executing(Coroutine::park).named("Command");
+    trigger.whileTrue(command);
+
+    triggerSignal.set(true);
+    m_scheduler.run();
+    assertTrue(m_scheduler.isRunning(command), "Command should have started when triggered");
+
+    fetcher.clear();
+    m_scheduler.run();
+    assertFalse(m_scheduler.isRunning(command), "Command should have stopped when opmode exited");
+  }
+
   // The scheduler lifecycle polls triggers at the start of `run()`
   // Even though the trigger condition is set, the command exits and the trigger's scope goes
   // inactive before the next `run()` call can poll the trigger

commandsv3/src/test/resources/META-INF/services/org.junit.jupiter.api.extension.Extension

@@ -0,0 +1 @@
+org.wpilib.command3.MockHardwareExtension