fix: exorcise goroutine-leaking util.After from the codebase (#211)

Author: johnstcn

cmd/attach/attach.go

@@ -14,7 +14,6 @@ import (
 
 	tea "github.com/charmbracelet/bubbletea"
 	"github.com/coder/agentapi/lib/httpapi"
-	"github.com/coder/agentapi/lib/util"
 	"github.com/coder/quartz"
 	"github.com/spf13/cobra"
 	sse "github.com/tmaxmax/go-sse"
@@ -239,9 +238,11 @@ func runAttach(remoteURL string) error {
 	}
 
 	p.Send(finishMsg{})
+	graceTimer := quartz.NewReal().NewTimer(1 * time.Second)
+	defer graceTimer.Stop()
 	select {
 	case <-pErrCh:
-	case <-util.After(quartz.NewReal(), 1*time.Second):
+	case <-graceTimer.C:
 	}
 
 	return err

lib/screentracker/pty_conversation.go

@@ -428,11 +428,14 @@ func (c *PTYConversation) writeStabilize(ctx context.Context, messageParts ...Me
 	}, func() (bool, error) {
 		screen := c.cfg.AgentIO.ReadScreen()
 		if screen != screenBeforeMessage {
+			stabilityTimer := c.cfg.Clock.NewTimer(1 * time.Second)
 			select {
 			case <-ctx.Done():
+				stabilityTimer.Stop()
 				return false, ctx.Err()
-			case <-util.After(c.cfg.Clock, 1*time.Second):
+			case <-stabilityTimer.C:
 			}
+			stabilityTimer.Stop()
 			newScreen := c.cfg.AgentIO.ReadScreen()
 			return newScreen == screen, nil
 		}
@@ -458,11 +461,14 @@ func (c *PTYConversation) writeStabilize(ctx context.Context, messageParts ...Me
 				return false, xerrors.Errorf("failed to write carriage return: %w", err)
 			}
 		}
+		crTimer := c.cfg.Clock.NewTimer(25 * time.Millisecond)
 		select {
 		case <-ctx.Done():
+			crTimer.Stop()
 			return false, ctx.Err()
-		case <-util.After(c.cfg.Clock, 25*time.Millisecond):
+		case <-crTimer.C:
 		}
+		crTimer.Stop()
 		screen := c.cfg.AgentIO.ReadScreen()
 
 		return screen != screenBeforeCarriageReturn, nil

lib/termexec/termexec.go

@@ -127,7 +127,9 @@ func (p *Process) ReadScreen() string {
 			return state
 		}
 		p.screenUpdateLock.RUnlock()
-		<-util.After(p.clock, 16*time.Millisecond)
+		t := p.clock.NewTimer(16 * time.Millisecond)
+		<-t.C
+		t.Stop()
 	}
 	return p.xp.State.String()
 }
@@ -152,9 +154,11 @@ func (p *Process) Close(logger *slog.Logger, timeout time.Duration) error {
 		close(exited)
 	}()
 
+	timeoutTimer := p.clock.NewTimer(timeout)
+	defer timeoutTimer.Stop()
 	var exitErr error
 	select {
-	case <-util.After(p.clock, timeout):
+	case <-timeoutTimer.C:
 		if err := p.execCmd.Process.Kill(); err != nil {
 			exitErr = xerrors.Errorf("failed to forcefully kill the process: %w", err)
 		}

lib/util/util.go

@@ -99,22 +99,3 @@ func OpenAPISchema[T ~string](r huma.Registry, enumName string, values []T) *hum
 	}
 	return &huma.Schema{Ref: fmt.Sprintf("#/components/schemas/%s", enumName)}
 }
-
-// After is a convenience function that returns a channel that will send the
-// time after the given duration has elapsed using the provided clock.
-// If clk is nil, a real clock will be used by default.
-// Note that this function spawns a goroutine that will remain alive until the
-// timer fires.
-func After(clk quartz.Clock, d time.Duration) <-chan time.Time {
-	if clk == nil {
-		clk = quartz.NewReal()
-	}
-	timer := clk.NewTimer(d)
-	ch := make(chan time.Time)
-	go func() {
-		defer timer.Stop()
-		defer close(ch)
-		ch <- <-timer.C
-	}()
-	return ch
-}