diff --git a/spring-boot-cli/src/main/java/org/springframework/boot/cli/compiler/autoconfigure/SpringWebsocketCompilerAutoConfiguration.java b/spring-boot-cli/src/main/java/org/springframework/boot/cli/compiler/autoconfigure/SpringWebsocketCompilerAutoConfiguration.java
index bccbe3dbba0..1d6e68df705 100644
--- a/spring-boot-cli/src/main/java/org/springframework/boot/cli/compiler/autoconfigure/SpringWebsocketCompilerAutoConfiguration.java
+++ b/spring-boot-cli/src/main/java/org/springframework/boot/cli/compiler/autoconfigure/SpringWebsocketCompilerAutoConfiguration.java
@@ -47,9 +47,10 @@ public class SpringWebsocketCompilerAutoConfiguration extends CompilerAutoConfig
 	public void applyImports(ImportCustomizer imports) {
 		imports.addStarImports("org.springframework.messaging.handler.annotation",
 				"org.springframework.messaging.simp.config",
-				"org.springframework.web.socket",
 				"org.springframework.web.socket.handler",
-				"org.springframework.web.socket.config.annotation");
+				"org.springframework.web.socket.sockjs.transport.handler",
+				"org.springframework.web.socket.config.annotation").addImports(
+				"org.springframework.web.socket.WebSocketHandler");
 	}
 
 }