PsiAugmentProvider for Kotlin class?

How does one use a PsiAugmentProvider to augment a Kotlin class?

I tried the following but it does not find the method when writing eg: this.myMethod() causes [UNRESOLVED_REFERENCE]

override fun <Psi : PsiElement?> getAugments(element: PsiElement, type: Class<Psi>): MutableList<Psi> {
    if (type == PsiMethod::class.java)
        if (element is KtLightClass && element.hasAnnotation("com.example.MyAnnotation")) {
            val method = LightMethodBuilder(element.manager, JavaLanguage.INSTANCE, "myMethod")
                .apply {
                    containingClass = element
                    setMethodReturnType("Unit")
                    setModifiers("public")
                }

            return mutableListOf(KtLightMethodImpl.create(method, null, element) as Psi)
        }

    return super.getAugments(element, type) as MutableList<Psi>
}