Fix test flake

private/buf/buflsp/buflsp.go

@@ -61,6 +61,17 @@
 	logger.Info("starting LSP server")
 
 	conn := jsonrpc2.NewConn(stream)
+	// connCtx is a context scoped to the connection's lifetime. It is cancelled
+	// when the connection is done (or when ctx is cancelled), so that background
+	// goroutines (e.g. RunChecks) do not outlive the connection.
+	connCtx, connCancel := context.WithCancel(context.Background())
+	go func() {
+		select {
+		case <-ctx.Done():
+		case <-conn.Done():
+		}
+		connCancel()
+	}()
 	lsp := &lsp{
 		conn: conn,
 		client: protocol.ClientDispatcher(
@@ -76,6 +87,7 @@
 		queryExecutor: queryExecutor,
 		opener:        source.NewMap(nil),
 		irSession:     new(ir.Session),
+		connCtx:       connCtx,
 	}
 	lsp.fileManager = newFileManager(lsp)
 	lsp.workspaceManager = newWorkspaceManager(lsp)
@@ -103,6 +115,7 @@
 	conn      jsonrpc2.Conn
 	client    protocol.Client
 	container appext.Container
+	connCtx   context.Context // cancelled when the connection is done
 
 	logger           *slog.Logger
 	bufVersion       string // buf version, set at server creation

private/buf/buflsp/file.go

@@ -1062,7 +1062,7 @@ func (f *file) RunChecks(ctx context.Context) {
 	}
 
 	const checkTimeout = 30 * time.Second
-	ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), checkTimeout)
+	ctx, cancel := context.WithTimeout(f.lsp.connCtx, checkTimeout)
 	f.cancelChecks = cancel
 
 	go func() {
@@ -1079,7 +1079,7 @@ func (f *file) RunChecks(ctx context.Context) {
 			); err != nil {
 				var fileAnnotationSet bufanalysis.FileAnnotationSet
 				if !errors.As(err, &fileAnnotationSet) {
-					if errors.Is(err, context.Canceled) {
+					if errors.Is(err, context.Canceled) || ctx.Err() != nil {
 						f.lsp.logger.DebugContext(ctx, "checks cancelled", slog.String("uri", f.uri.Filename()), xslog.ErrorAttr(err))
 					} else if errors.Is(err, context.DeadlineExceeded) {
 						f.lsp.logger.WarnContext(ctx, "checks deadline exceeded", slog.String("uri", f.uri.Filename()), xslog.ErrorAttr(err))

private/buf/buflsp/image.go

@@ -85,7 +85,7 @@ func buildImage(
 	var diagnostics []protocol.Diagnostic
 	compiled, err := compiler.Compile(ctx, path)
 	if err != nil {
-		logger.Warn("error building image", slog.String("path", path), xslog.ErrorAttr(err))
+		logger.WarnContext(ctx, "error building image", slog.String("path", path), xslog.ErrorAttr(err))
 		var errorWithPos reporter.ErrorWithPos
 		if errors.As(err, &errorWithPos) {
 			diagnostics = []protocol.Diagnostic{newDiagnostic(errorWithPos, false, opener, logger)}