diff --git a/src/loop.c b/src/loop.c
index f32391e0..39ade3ac 100644
--- a/src/loop.c
+++ b/src/loop.c
@@ -130,7 +130,11 @@ int mosquitto_main_loop(struct mosquitto_db *db, int *listensock, int listensock
 			}
 			context->pollfd_index = -1;
 
+#ifdef WITH_WEBSOCKETS
+			if(context->sock != INVALID_SOCKET || context->wsi){
+#else
 			if(context->sock != INVALID_SOCKET){
+#endif
 #ifdef WITH_BRIDGE
 				if(context->bridge){
 					_mosquitto_check_keepalive(db, context);
@@ -348,6 +352,9 @@ static void do_disconnect(struct mosquitto_db *db, struct mosquitto *context)
 		}
 	}
 	mqtt3_context_disconnect(db, context);
+	if(context->clean_session){
+		mqtt3_context_cleanup(db, context, true);
+	}
 }
 
 /* Error ocurred, probably an fd has been closed. 
diff --git a/src/websockets.c b/src/websockets.c
index 159d623e..d31ed5ed 100644
--- a/src/websockets.c
+++ b/src/websockets.c
@@ -129,7 +129,11 @@ static int callback_mqtt(struct libwebsocket_context *context,
 
 		case LWS_CALLBACK_CLOSED:
 			mosq = u->mosq;
-			mqtt3_context_cleanup(db, mosq, true);
+			mqtt3_context_disconnect(db, mosq);
+			if(context->clean_session){
+				mqtt3_context_cleanup(db, mosq, true);
+			}
+			mosq->wsi = NULL;
 			break;
 
 		case LWS_CALLBACK_SERVER_WRITEABLE: