Faits :
A.
Par jugement du 25 septembre 2019, le Tribunal correctionnel de l'arrondissement de l'Est vaudois a reconnu A._ coupable de contravention, d'infraction simple et d'infraction grave à la LStup ainsi que d'infractions à la LCR. Il a révoqué le sursis octroyé le 6 mars 2013 par le Regionalgericht Bern-Mittelland, et l'a condamné à une peine privative de liberté d'ensemble de 40 mois, sous déduction de la détention subie avant jugement et de la détention subie dans des conditions illicites, ainsi qu'à une amende de 1'000 fr., la peine privative de liberté de substitution étant fixée à 20 jours.
Par jugement du 27 janvier 2020, la Cour d'appel pénale du Tribunal cantonal vaudois a rejeté l'appel formé par A._ contre le jugement de première instance.
B.
Par arrêt du 15 mai 2020 (6B_291/2020), le Tribunal fédéral a partiellement admis le recours formé par A._ contre le jugement du 27 janvier 2020, l'a annulé s'agissant de la peine et a renvoyé la cause à la cour cantonale pour nouvelle décision sur ce point. Pour le surplus, le recours a été rejeté dans la mesure de sa recevabilité.
C.
Par jugement du 22 juin 2020, rendu à la suite de l'arrêt de renvoi du 15 mai 2020, la Cour d'appel pénale du Tribunal cantonal vaudois a rejeté l'appel et a maintenu son dispositif.
D.
A._ forme un recours en matière pénale auprès du Tribunal fédéral contre le jugement du 22 juin 2020 et conclut, avec suite de frais et dépens, au prononcé d'une peine privative de liberté qui ne saurait excéder la durée de la détention déjà subie et au prononcé de sa liberté immédiate. Subsidiairement, il conclut à l'annulation du jugement et au renvoi de la cause à la cour cantonale pour nouvelle décision. Il sollicite en outre le bénéfice de l'assistance judiciaire.
E.
Invités à se déterminer sur le recours de A._, la cour cantonale et le Ministère public vaudois y ont renoncé.

Considérant en droit :
1.
Le recourant reproche à la cour cantonale de ne pas avoir réduit sa peine au regard de la circonstance atténuante de l'<ref-law>, alors qu'elle ne disposait d'aucune marge de manoeuvre sur ce point eu égard à l'arrêt de renvoi sur ce point.
1.1.
1.1.1. Aux termes de l'art. 107 al. 2 1ère phrase LTF, si le Tribunal fédéral admet le recours, il statue lui-même sur le fond ou renvoie l'affaire à l'autorité précédente pour qu'elle prenne une nouvelle décision. Le principe de l'autorité de l'arrêt de renvoi découle du droit fédéral non écrit (<ref-ruling> consid. 5.3.3 p. 222; <ref-ruling> consid. 2.1 p. 335). Conformément à ce principe, l'autorité cantonale à laquelle la cause est renvoyée par le Tribunal fédéral est tenue de fonder sa nouvelle décision sur les considérants de droit de l'arrêt du Tribunal fédéral. Elle est ainsi liée par ce qui a déjà été définitivement tranché par le Tribunal fédéral et par les constatations de fait qui n'ont pas été attaquées devant lui ou l'ont été sans succès (<ref-ruling> consid. 5.2 p. 94; <ref-ruling> consid. 3b p. 277; cf. aussi arrêt 6B_1114/2019 du 20 novembre 2019 consid. 1.1). La motivation de l'arrêt de renvoi détermine dans quelle mesure la cour cantonale est liée à la première décision, décision de renvoi qui fixe aussi bien le cadre du nouvel état de fait que celui de la nouvelle motivation juridique (ATF <ref-ruling> consid. 2 p. 335). Les faits nouveaux ne peuvent être pris en considération que sur les points ayant fait l'objet du renvoi, lesquels ne peuvent être ni étendus, ni fondés sur une base juridique nouvelle (ATF <ref-ruling> consid. 2 p. 335; arrêts 6B_527/2020 du 29 septembre 2020 consid. 1.1; 6B_1114/2019 du 20 novembre 2019 consid. 1.1).
1.1.2. Aux termes de l'<ref-law>, le juge peut atténuer librement la peine dans le cas d'une infraction visée à l'al. 2, si l'auteur est dépendant et que cette infraction aurait dû servir au financement de sa propre consommation de stupéfiants.
1.2.
1.2.1. Dans le jugement sur appel du 27 janvier 2020, la cour cantonale a fixé à 30 mois la peine de base pour l'infraction grave à la LStup dont le recourant a été reconnu coupable. Par l'effet du concours, elle a augmenté cette peine de 2 mois pour l'infraction simple à la LStup et de 2 mois supplémentaires en raison du concours avec les infractions à la LCR, qui justifiaient des peines privatives de liberté en raison des récidives innombrables dans ce domaine. A cela s'ajoutait encore une aggravation de la peine de l'ordre de 6 mois pour tenir compte de la révocation du sursis à la peine de 8 mois prononcée par le Regionalgericht Bern-Mittelland. Elle a ainsi fixé la peine d'ensemble infligée au recourant à 40 mois.
1.2.2. Par arrêt de renvoi du 15 mai 2020, le Tribunal fédéral a retenu que, alors même que la cour cantonale avait mis le recourant au bénéfice d'une circonstance atténuante (<ref-law>), elle n'avait pas réduit la peine privative de liberté de base de 30 mois qu'elle avait fixée.
1.3. Dans son jugement sur renvoi, la cour cantonale expose qu'en fixant la peine de base pour l'infraction grave à la LStup, elle avait " à l'esprit qu'il fallait prendre en considération (...) à décharge le cas d'atténuation de l'<ref-law> ". Elle relève que, pour arrêter la peine privative de liberté de 30 mois, les circonstances à charge (trafic de drogue dure et antécédents) et à décharge (bonne collaboration et toxicomanie) se compensent. Dans un nouveau raisonnement, la cour cantonale expose que, si elle n'avait pas eu à tenir compte de la toxicomanie et de l'atténuation résultant de l'<ref-law>, elle aurait fixé la peine privative de liberté de base à 36 mois.
1.4. Cette motivation est contradictoire et ne résiste pas à l'examen.
Il est établi que les conditions de l'<ref-law> sont réalisées et que les juges précédents n'ont pas renoncé à atténuer la peine, malgré le caractère facultatif de la démarche.
Or, alors qu'il ressort expressément du jugement cantonal du 27 janvier 2020 que la peine arrêtée à 30 mois pour infraction grave à la LStup (<ref-law>), a été aggravée à plusieurs égards pour différents motifs, la cour cantonale n'a à aucun moment exposé que la peine était atténuée en raison de la circonstance atténuante prévue à l'<ref-law>. Dans la mesure où le Tribunal fédéral a annulé le jugement concernant la peine et renvoyé la cause à la cour cantonale au motif qu'elle n'avait pas réduit la peine privative de liberté de base, la cour cantonale ne pouvait faire l'économie d'une atténuation de peine et se contenter de prétendre " avoir à l'esprit " une circonstance atténuante comme facteur à décharge.
L'arrêt de renvoi fixait le cadre de la nouvelle motivation juridique, laquelle devait porter sur une atténuation de peine et ne visait pas une justification a posteriori des circonstances à charge et à décharge qui auraient été retenues.
Le recours doit être admis, le jugement cantonal annulé s'agissant de la peine et la cause renvoyée à la cour cantonale pour qu'elle atténue la peine prononcée pour infraction grave à la LStup, en tenant compte de la circonstance atténuante déduite de l'<ref-law>, disposition qui figure expressément dans le dispositif du jugement.
2.
Le recours doit être admis. Le recourant, qui obtient gain de cause, ne supporte pas de frais judiciaires (<ref-law>). Il peut prétendre à une indemnité de dépens, à la charge du canton de Vaud. La demande d'assistance judiciaire devient ainsi sans objet.

Par ces motifs, le Tribunal fédéral prononce :
1.
Le recours est admis, le jugement attaqué est annulé s'agissant de la peine et la cause est renvoyée à l'autorité cantonale pour nouvelle décision.
2.
Il n'est pas perçu de frais judiciaires.
3.
Le canton de Vaud versera au conseil du recourant une indemnité de 3'000 fr. à titre de dépens pour la procédure devant le Tribunal fédéral.
4.
Le présent arrêt est communiqué aux parties et à la Cour d'appel pénale du Tribunal cantonal du canton de Vaud.
Lausanne, le 16 novembre 2020
Au nom de la Cour de droit pénal
du Tribunal fédéral suisse
Le Président : Denys
La Greffière : Klinke