Issues successfully adding folder to java main sourcesets as task action

Hi there,
Quite new to Gradle, still figuring some things out.

I’m developing a custom plugin for our company, consisting of two tasks that are exposed to the user.
Basically, the first task generates and copies some java source files to a subfolder of build/generated/sources.
The second task should add the folder to the java sourceset so that it is compiled along with the user’s sources.

I can get it working by adding it as a source to the compileJava task:

(AddToSourceSetsTask.kt)
open class AddToSourceSetsTask: DefaultTask() {
    val folder: File = project.file("build/generated/sources/...")

    @TaskAction
    fun addToSourceSets() {
        project.tasks.named("compileJava", JavaCompile::class.java) {
            it.source(folder)
        }
    }
}

But this seems to be the wrong way to go about this.
I have tried adding it as such:

(AddToSourceSetsTask.kt)
open class AddToSourceSetsTask: DefaultTask() {
    val folder: File = project.file("build/generated/sources/...")

    @TaskAction
    fun addToSourceSets() {
        project.convention.findPlugin(JavaPluginConvention::class.java)?.sourceSets?.getByName("main")?.allJava?.srcDir(folder)
    }
}

but when building, the generated source files are not actually compiled.

Any help would be appreciated.
Cheers!