@@ -1,13 +1,13 @@
 package command
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
 	"os"
 	"os/exec"
 	"strings"
-	"time"
 
 	"alin.ovh/erl/output"
 )
@@ -15,7 +15,7 @@
 type Command interface {
 	Start() error
 	Wait() error
-	Stop() error
+	Stop(context.Context) error
 }
 
 type Cmd struct {
@@ -35,8 +35,6 @@ Stdout io.Writer
 	Stderr io.Writer
 }
 
-const timeout = 1 * time.Second
-
 func New(name string, args []string, options Options) *Cmd {
 	if options.Stdout == nil {
 		options.Stdout = os.Stdout
@@ -57,7 +55,7 @@ out: output.New(options.Output),
 	}
 }
 
-func (cmd *Cmd) Stop() error {
+func (cmd *Cmd) Stop(ctx context.Context) error {
 	if cmd.Cmd == nil || cmd.Process == nil || cmd.ProcessState != nil {
 		return nil
 	}
@@ -67,26 +65,32 @@ if err != nil {
 		return fmt.Errorf("error killing command: %v", err)
 	}
 
-	if cmd.ProcessState == nil {
-		cmd.out.Info("[command not stopped, waiting %s seconds before killing]\n", timeout)
-		t := timeout / 10
-		for range 10 {
-			time.Sleep(t)
-			if cmd.ProcessState != nil {
-				return nil
+	exited := make(chan struct{})
+	go func() {
+		cmd.Wait()
+		exited <- struct{}{}
+		close(exited)
+	}()
+
+	for {
+		select {
+		case <-ctx.Done():
+			err := ctx.Err()
+			if !errors.Is(err, context.DeadlineExceeded) {
+				return err
 			}
-		}
 
-		if cmd.ProcessState == nil {
-			cmd.out.Info("[command not stopped, killing]\n")
-			err := cmd.Process.Kill()
-			if err != nil {
+			cmd.out.Fail("[command not stopped, killing]\n")
+			err = cmd.Process.Kill()
+			if err != nil && !errors.Is(err, os.ErrProcessDone) {
 				return fmt.Errorf("error killing command: %v", err)
 			}
+
+			return nil
+		case <-exited:
+			return nil
 		}
 	}
-
-	return nil
 }
 
 func (cmd *Cmd) Exited() bool {
@@ -95,21 +99,18 @@ }
 
 func (cmd *Cmd) Wait() error {
 	err := cmd.Cmd.Wait()
-	if err != nil {
-		var exitErr *exec.ExitError
-		if errors.As(err, &exitErr) {
-			if cmd.Exited() {
-				cmd.out.Fail("[command exited with code %d]\n", cmd.ProcessState.ExitCode())
-			} else {
-				cmd.out.Fail("[command stopped]\n")
-			}
-		} else {
-			return fmt.Errorf("error waiting for command: %v", err)
-		}
+	var exitErr *exec.ExitError
+	if err != nil && !errors.As(err, &exitErr) {
+		return err
 	}
 
-	if cmd.Exited() && cmd.ProcessState.Success() {
+	switch {
+	case !cmd.Exited():
+		cmd.out.Info("[command stopped]\n")
+	case cmd.ProcessState.Success():
 		cmd.out.Success("[command finished]\n")
+	default:
+		cmd.out.Fail("[command exited with code %d]\n", cmd.ProcessState.ExitCode())
 	}
 
 	return nil