@@ -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

@@ -206,7 +206,7 @@
 		// Give the command a moment to actually start
 		time.Sleep(100 * time.Millisecond)
 
-		err = cmd.Stop()
+		err = cmd.Stop(t.Context())
 		if err != nil {
 			t.Errorf("unexpected error stopping command: %v", err)
 		}
@@ -239,7 +239,7 @@ Stderr: &stderr,
 		})
 
 		// Try to stop before starting
-		err := cmd.Stop()
+		err := cmd.Stop(t.Context())
 		if err != nil {
 			t.Errorf("stopping non-running command should not error: %v", err)
 		}
@@ -265,7 +265,7 @@ if err != nil {
 			t.Fatalf("failed to wait for command: %v", err)
 		}
 
-		err = cmd.Stop()
+		err = cmd.Stop(t.Context())
 		if err != nil {
 			t.Errorf("stopping finished command should not error: %v", err)
 		}
@@ -310,7 +310,7 @@ t.Error("running command should not be marked as exited")
 		}
 
 		// Clean up
-		cmd.Stop()
+		cmd.Stop(t.Context())
 	})
 
 	t.Run("finished command", func(t *testing.T) {
@@ -437,7 +437,7 @@ if err != nil {
 			t.Errorf("Wait() failed: %v", err)
 		}
 
-		err = cmd.Stop()
+		err = cmd.Stop(t.Context())
 		if err != nil {
 			t.Errorf("Stop() failed: %v", err)
 		}

@@ -169,6 +169,9 @@ logger = log.New(io.Discard, "", 0)
 	}
 
 	cmd := command.New(program, args, copts)
-	sm := state.New(cmd, log.New(os.Stderr, "state: ", log.Lmsgprefix))
+	sm := state.New(cmd, state.Options{
+		Logger:  log.New(os.Stderr, "state: ", log.Lmsgprefix),
+		Timeout: time.Second * 5,
+	})
 	Start(ctx, logger, watcher, sm)
 }

@@ -1,7 +1,9 @@
 package state
 
 import (
+	"context"
 	"log"
+	"time"
 
 	"alin.ovh/erl/command"
 )
@@ -57,15 +59,31 @@
 type Action func()
 
 type StateMachine struct {
-	log          *log.Logger
+	log     *log.Logger
+	timeout time.Duration
+
 	currentState State
 	transitions  map[State]map[Event]State
 	actions      map[State]map[Event]Action
 }
 
-func New(cmd command.Command, log *log.Logger) *StateMachine {
+const DefaultTimeout = 1 * time.Second
+
+type Options struct {
+	Timeout time.Duration
+	Logger  *log.Logger
+}
+
+func New(cmd command.Command, opts Options) *StateMachine {
+	if opts.Timeout == 0 {
+		opts.Timeout = DefaultTimeout
+	}
+	log := opts.Logger
+
 	sm := &StateMachine{
-		log:          log,
+		log:     opts.Logger,
+		timeout: opts.Timeout,
+
 		currentState: NotStarted,
 		transitions:  make(map[State]map[Event]State),
 		actions:      make(map[State]map[Event]Action),
@@ -108,7 +126,10 @@
 		go wait(cmd)
 	}
 	stop := func() {
-		err := cmd.Stop()
+		ctx, cancel := context.WithTimeout(context.Background(), sm.timeout)
+		defer cancel()
+
+		err := cmd.Stop(ctx)
 		if err != nil {
 			log.Printf("Error stopping command: %v\n", err)
 		}

@@ -2,6 +2,7 @@ package state
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"log"
 	"sync"
@@ -48,7 +49,7 @@
 	return m.waitErr
 }
 
-func (m *MockCommand) Stop() error {
+func (m *MockCommand) Stop(_ context.Context) error {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 	m.stopCalled = true
@@ -148,7 +149,9 @@ mockCmd := NewMockCommand()
 	var buf bytes.Buffer
 	logger := log.New(&buf, "", 0)
 
-	sm := New(mockCmd, logger)
+	sm := New(mockCmd, Options{
+		Logger: logger,
+	})
 
 	return sm, mockCmd, &buf
 }