Automatically pump os.proc streams when SystemStreams are redirected

build.sc

@@ -153,7 +153,7 @@ object Deps {
   val junitInterface = ivy"com.github.sbt:junit-interface:0.13.3"
   val lambdaTest = ivy"de.tototec:de.tobiasroeser.lambdatest:0.8.0"
   val log4j2Core = ivy"org.apache.logging.log4j:log4j-core:2.23.1"
-  val osLib = ivy"com.lihaoyi::os-lib:0.10.2"
+  val osLib = ivy"com.lihaoyi::os-lib:0.10.2-11-6f1c69"
   val pprint = ivy"com.lihaoyi::pprint:0.9.0"
   val mainargs = ivy"com.lihaoyi::mainargs:0.7.0"
   val millModuledefsVersion = "0.10.9"
@@ -618,6 +618,7 @@ object main extends MillStableScalaModule with BuildInfo {
   )
 
   object api extends MillStableScalaModule with BuildInfo {
+    def moduleDeps = Seq(client)
     def buildInfoPackageName = "mill.api"
     def buildInfoMembers = Seq(
       BuildInfo.Value("millVersion", millVersion(), "Mill version."),

main/api/src/mill/api/SystemStreams.scala

@@ -1,6 +1,7 @@
 package mill.api
 
-import java.io.{InputStream, PrintStream}
+import java.io.{InputStream, OutputStream, PrintStream}
+import mill.main.client.InputPumper
 
 /**
  * Represents a set of streams that look similar to those provided by the
@@ -48,6 +49,18 @@ object SystemStreams {
 
   def originalErr: PrintStream = original.err
 
+  private class PumpedProcessInput extends os.ProcessInput {
+    def redirectFrom = ProcessBuilder.Redirect.PIPE
+    def processInput(processIn: => os.SubProcess.InputStream) = Some(
+      new InputPumper(() => System.in, () => processIn, true, () => true)
+    )
+  }
+
+  private class PumpedProcessOutput(dest: OutputStream) extends os.ProcessOutput {
+    def redirectTo = ProcessBuilder.Redirect.PIPE
+    def processOutput(processOut: => os.SubProcess.OutputStream) =
+      Some(new InputPumper(() => processOut, () => dest, false, () => true))
+  }
   def withStreams[T](systemStreams: SystemStreams)(t: => T): T = {
     val in = System.in
     val out = System.out
@@ -59,7 +72,13 @@ object SystemStreams {
       Console.withIn(systemStreams.in) {
         Console.withOut(systemStreams.out) {
           Console.withErr(systemStreams.err) {
-            t
+            os.Inherit.in.withValue(new PumpedProcessInput) {
+              os.Inherit.out.withValue(new PumpedProcessOutput(System.out)) {
+                os.Inherit.err.withValue(new PumpedProcessOutput(System.err)) {
+                  t
+                }
+              }
+            }
           }
         }
       }

main/client/src/mill/main/client/InputPumper.java

@@ -2,29 +2,34 @@
 
 import java.io.InputStream;
 import java.io.OutputStream;
+import java.util.function.Supplier;
 
 public class InputPumper implements Runnable{
-    private InputStream src;
-    private OutputStream dest;
+    private Supplier<InputStream> src0;
+    private Supplier<OutputStream> dest0;
+
     private Boolean checkAvailable;
     private java.util.function.BooleanSupplier runningCheck;
-    public InputPumper(InputStream src,
-                       OutputStream dest,
+    public InputPumper(Supplier<InputStream> src,
+                       Supplier<OutputStream> dest,
                        Boolean checkAvailable){
         this(src, dest, checkAvailable, () -> true);
     }
-    public InputPumper(InputStream src,
-                       OutputStream dest,
+    public InputPumper(Supplier<InputStream> src,
+                       Supplier<OutputStream> dest,
                        Boolean checkAvailable,
                        java.util.function.BooleanSupplier runningCheck){
-        this.src = src;
-        this.dest = dest;
+        this.src0 = src;
+        this.dest0 = dest;
         this.checkAvailable = checkAvailable;
         this.runningCheck = runningCheck;
     }
 
     boolean running = true;
     public void run() {
+        InputStream src = src0.get();
+        OutputStream dest = dest0.get();
+
         byte[] buffer = new byte[1024];
         try{
             while(running){

main/client/src/mill/main/client/MillClientMain.java

@@ -221,7 +221,7 @@ public static int run(
         InputStream outErr = ioSocket.getInputStream();
         OutputStream in = ioSocket.getOutputStream();
         ProxyStreamPumper outPump = new ProxyStreamPumper(outErr, stdout, stderr);
-        InputPumper inPump = new InputPumper(stdin, in, true);
+        InputPumper inPump = new InputPumper(() -> stdin, () -> in, true);
         Thread outThread = new Thread(outPump, "outPump");
         outThread.setDaemon(true);
         Thread inThread = new Thread(inPump, "inPump");

main/util/src/mill/util/Jvm.scala

@@ -223,46 +223,13 @@ object Jvm extends CoursierSupport {
       workingDir: os.Path,
       backgroundOutputs: Option[Tuple2[ProcessOutput, ProcessOutput]] = None
   ): SubProcess = {
-    // If System.in is fake, then we pump output manually rather than relying
-    // on `os.Inherit`. That is because `os.Inherit` does not follow changes
-    // to System.in/System.out/System.err, so the subprocess's streams get sent
-    // to the parent process's origin outputs even if we want to direct them
-    // elsewhere
-
-    if (!SystemStreams.isOriginal()) {
-      val process = os.proc(commandArgs).spawn(
-        cwd = workingDir,
-        env = envArgs,
-        stdin = if (backgroundOutputs.isEmpty) os.Pipe else "",
-        stdout = backgroundOutputs.map(_._1).getOrElse(os.Pipe),
-        stderr = backgroundOutputs.map(_._2).getOrElse(os.Pipe)
-      )
-
-      val sources = Seq(
-        (process.stdout, System.out, "spawnSubprocess.stdout", false, () => true),
-        (process.stderr, System.err, "spawnSubprocess.stderr", false, () => true),
-        (System.in, process.stdin, "spawnSubprocess.stdin", true, () => process.isAlive())
-      )
-
-      for ((std, dest, name, checkAvailable, runningCheck) <- sources) {
-        val t = new Thread(
-          new InputPumper(std, dest, checkAvailable, () => runningCheck()),
-          name
-        )
-        t.setDaemon(true)
-        t.start()
-      }
-
-      process
-    } else {
-      os.proc(commandArgs).spawn(
-        cwd = workingDir,
-        env = envArgs,
-        stdin = if (backgroundOutputs.isEmpty) os.Inherit else "",
-        stdout = backgroundOutputs.map(_._1).getOrElse(os.Inherit),
-        stderr = backgroundOutputs.map(_._2).getOrElse(os.Inherit)
-      )
-    }
+    os.proc(commandArgs).spawn(
+      cwd = workingDir,
+      env = envArgs,
+      stdin = if (backgroundOutputs.isEmpty) os.Inherit else "",
+      stdout = backgroundOutputs.map(_._1).getOrElse(os.Inherit),
+      stderr = backgroundOutputs.map(_._2).getOrElse(os.Inherit)
+    )
   }
 
   def runLocal(

runner/src/mill/runner/MillServerMain.scala

@@ -135,7 +135,7 @@ class Server[T](
     val pipedInput = new PipedInputStream()
     val pipedOutput = new PipedOutputStream()
     pipedOutput.connect(pipedInput)
-    val pumper = new InputPumper(in, pipedOutput, false)
+    val pumper = new InputPumper(() => in, () => pipedOutput, false)
     val pumperThread = new Thread(pumper, "proxyInputStreamThroughPumper")
     pumperThread.setDaemon(true)
     pumperThread.start()

scalajslib/worker/1/src/mill/scalajslib/worker/ScalaJSWorkerImpl.scala

@@ -305,7 +305,12 @@ class ScalaJSWorkerImpl extends ScalaJSWorkerApi {
 
           for ((std, dest, name, checkAvailable, runningCheck) <- sources) {
             val t = new Thread(
-              new mill.main.client.InputPumper(std, dest, checkAvailable, () => runningCheck()),
+              new mill.main.client.InputPumper(
+                () => std,
+                () => dest,
+                checkAvailable,
+                () => runningCheck()
+              ),
               name
             )
             t.setDaemon(true)