Add leak detector (#811)

* feat: add goleak detection in test suites to prevent goroutine leaks

* docs: add instructions for integrating goleak in Go package tests

* feat: integrate goleak detection in end-to-end tests to prevent goroutine leaks

* feat: refactor end-to-end tests to check for goroutine leaks using a dedicated function

Author: ilopezluna

AGENTS.md

@@ -23,6 +23,34 @@ If any step fails, fix the issue and re-run before committing.
 - **golangci-lint v2.7.2+** — [Install instructions](https://golangci-lint.run/welcome/install/)
 - **ShellCheck** — `brew install shellcheck` (macOS) or `apt-get install shellcheck` (Linux)
 
+## Testing
+
+When adding a new Go package that contains tests, include a
+`testmain_test.go` file in the package directory to enable automatic
+goroutine leak detection via `go.uber.org/goleak`. Match the `package`
+declaration to the existing test files in that directory (use the external
+`package foo_test` form if that is what the other test files use).
+
+```go
+package <pkg> // or <pkg>_test — match the existing test files
+
+import (
+	"testing"
+
+	"go.uber.org/goleak"
+)
+
+// TestMain runs goleak after the test suite to detect goroutine leaks.
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m)
+}
+```
+
+If the package already has a `TestMain` that does non-trivial setup (e.g.
+starting a server), integrate goleak manually using the
+setup → `m.Run()` → teardown → `goleak.Find` → `os.Exit` pattern instead
+of calling `goleak.VerifyTestMain`.
+
 ## Project documentation
 
 - [README.md](./README.md) — project overview, building from source, API examples, and Makefile usage

cmd/cli/commands/testmain_test.go

@@ -0,0 +1,12 @@
+package commands
+
+import (
+	"testing"
+
+	"go.uber.org/goleak"
+)
+
+// TestMain runs goleak after the test suite to detect goroutine leaks.
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m)
+}

cmd/cli/desktop/testmain_test.go

@@ -0,0 +1,12 @@
+package desktop
+
+import (
+	"testing"
+
+	"go.uber.org/goleak"
+)
+
+// TestMain runs goleak after the test suite to detect goroutine leaks.
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m)
+}

cmd/cli/search/testmain_test.go

@@ -0,0 +1,12 @@
+package search
+
+import (
+	"testing"
+
+	"go.uber.org/goleak"
+)
+
+// TestMain runs goleak after the test suite to detect goroutine leaks.
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m)
+}

e2e/e2e_test.go

@@ -26,6 +26,8 @@ import (
 	"testing"
 	"time"
 
+	"go.uber.org/goleak"
+
 	"github.com/docker/model-runner/cmd/cli/desktop"
 	"github.com/docker/model-runner/pkg/ollama"
 )
@@ -93,7 +95,6 @@ func runNative(m *testing.M, root string) int {
 	fmt.Fprintf(os.Stderr, "e2e: starting model-runner on port %d\n", port)
 
 	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
 
 	server := exec.CommandContext(ctx, serverBin)
 	// TODO: os.Interrupt is not supported on Windows. When Windows e2e
@@ -111,14 +112,17 @@ func runNative(m *testing.M, root string) int {
 
 	if err := server.Start(); err != nil {
 		fmt.Fprintf(os.Stderr, "e2e: failed to start server: %v\n", err)
+		cancel()
 		return 1
 	}
-	defer func() {
-		cancel()
-		_ = server.Wait()
-	}()
 
-	return waitAndRunTests(m)
+	code := waitAndRunTests(m)
+
+	// Shut down the server and drain its goroutines before the leak check.
+	cancel()
+	_ = server.Wait()
+
+	return checkLeaks(code)
 }
 
 // runDocker builds the Docker image and CLI from source, then lets the CLI
@@ -153,7 +157,23 @@ func runDocker(m *testing.M, root string) int {
 	}
 
 	serverURL = "http://localhost:12434"
-	return waitAndRunTests(m)
+	return checkLeaks(waitAndRunTests(m))
+}
+
+// checkLeaks closes idle HTTP connections and checks for goroutine leaks in
+// the test harness. It returns code unchanged when code != 0 — a failing run
+// already reports errors and extra leak noise adds no value.
+func checkLeaks(code int) int {
+	// Close idle keep-alive connections so the default HTTP transport does
+	// not leave goroutines that would be false-positive leaks.
+	http.DefaultClient.CloseIdleConnections()
+	if code == 0 {
+		if err := goleak.Find(); err != nil {
+			fmt.Fprintf(os.Stderr, "e2e: goroutine leak detected: %v\n", err)
+			return 1
+		}
+	}
+	return code
 }
 
 func waitAndRunTests(m *testing.M) int {

go.mod

@@ -36,6 +36,7 @@ require (
 	github.com/testcontainers/testcontainers-go v0.41.0
 	github.com/testcontainers/testcontainers-go/modules/registry v0.41.0
 	go.opentelemetry.io/otel v1.42.0
+	go.uber.org/goleak v1.3.0
 	go.uber.org/mock v0.6.0
 	golang.org/x/sync v0.20.0
 	golang.org/x/sys v0.42.0

pkg/anthropic/testmain_test.go

@@ -0,0 +1,12 @@
+package anthropic
+
+import (
+	"testing"
+
+	"go.uber.org/goleak"
+)
+
+// TestMain runs goleak after the test suite to detect goroutine leaks.
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m)
+}

pkg/distribution/builder/testmain_test.go

@@ -0,0 +1,12 @@
+package builder_test
+
+import (
+	"testing"
+
+	"go.uber.org/goleak"
+)
+
+// TestMain runs goleak after the test suite to detect goroutine leaks.
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m)
+}

pkg/distribution/distribution/testmain_test.go

@@ -0,0 +1,12 @@
+package distribution
+
+import (
+	"testing"
+
+	"go.uber.org/goleak"
+)
+
+// TestMain runs goleak after the test suite to detect goroutine leaks.
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m)
+}

pkg/distribution/files/testmain_test.go

@@ -0,0 +1,12 @@
+package files
+
+import (
+	"testing"
+
+	"go.uber.org/goleak"
+)
+
+// TestMain runs goleak after the test suite to detect goroutine leaks.
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m)
+}