K2 - how to associate a KtCodeFragment with KaModule

IntelliJ Platform
, ,
1

(the post below was drafted with a great deal of help from ClaudeAI - thanks in advance)

The JTE template project had an IntelliJ plugin enabling near-perfect Kotlin code editing - autocomplete, error highlighting, etc. It relied on some hacks with the opaque MultiHostInjector API, but they no longer work with changes to that API for the K2 compiler, which is now the default while K1 is deprecated ( Add support for new Kotlin K2 mode. · Issue #49 · casid/jte-intellij · GitHub ). I (no prior experience with any of the Injector APIs) am attempting to migrate to LanguageInjectionContributor/LanguageInjectionPerformer to restore the Kotlin developer experience we previously enjoyed, on the K2 compiler. The migration itself was smooth but we’ve hit a KaModule association problem that appears regardless of which injection API is used.
What we’ve tried and where each approach breaks:

  1. MODULE_ROOT_TYPE_KEY (K1 approach): inert in K2, no effect.
  2. resolveExtensionFileModule on the VirtualFile: getKaModule() returns the correct module, but IDE features (semantic highlighter, intentions) that analyze PSI elements from the injected file throw KaBaseIllegalPsiException: The element cannot be analyzed in the context of the current session — a session mismatch between where the PSI was created and where the analysis session was opened.
  3. ktFile.contextModule = kaModule (@KaExperimentalApi): throws IllegalArgumentException: 'contextModule' cannot be set for code fragments — because for a non-Kotlin host file, the Kotlin injection machinery always creates a KtCodeFragment rather than a KtFile, regardless of the "kt" extension hint passed to startInjecting.
  4. KtCodeFragment.context: constructor-only val, cannot be set post-construction, so we can’t provide a context element after the platform has already created the fragment.

The core question: For a non-Kotlin host file, the platform appears to always produce a KtCodeFragment for injected Kotlin. Is there a supported K2 mechanism for associating that fragment with the correct KaModule so that IDE features can analyze its PSI elements? Or is there a way to make the platform produce a KtFile instead of a KtCodeFragment for this case?
We have a YouTrack reference KTIJ-32613 which appears related. Happy to file a new ticket if this is a known gap.

3

Hi! injections in k2 are represented as KtCodeFragments. Could you please paste the exception from resolveExtensionFileModule on the VirtualFile: getKaModule() returns the correct module, but IDE features (semantic highlighter, intentions) that analyze PSI elements from here:

“the injected file throw KaBaseIllegalPsiException: The element cannot be analyzed in the context of the current session — a session mismatch between where the PSI was created and where the analysis session was opened.” and the code which produces it.

Generally, you need to run the analysis session created from that fragment and the fragment itself should already contain the context of the injected host.

Thanks

4

element_stacktrace.txt (6.9 KB)

Could you please paste the exception

Hi Anna. A full stacktrace is attached to this reply.

Our host language is a template language, with Kotlin-compatible imports declared at the top of the page along with some variable declarations (param). In isolation, the Kotlin code fragments throughout the page have little meaning, they require knowledge of the earlier imports to understand types, and in the case of variables, need the earlier param definitions (i.e. @param username : String) to reference the defined type. Our Injector collects all the Kotlin code scattered across the template file and produces a syntactically-valid Kotlin file in-memory, and that’s what we try to associate the individual fragments with.

The real question is - our code fragments, in isolation, are each meaningless, without imports and variable definitions. So how can we use the various Language Injection APIs or strategies to obtain the key editing features (autocomplete, error highlighting) in the IDE, with the K2 compiler? How do we utilize the full, valid Kotlin file we built from the template’s code fragments? From my original question, you can see various things we tried, I also tried implementing PsiReferenceContributor/Provider and KaResolveExtensionProvider. None of these worked, mostly due to something related to the injected file being isolated from (and invisible to) the currently-analyzed fragment.

Thank you for your help! It would be amazing to get this working under K2.

5

We indeed use KtCodeFragments as a backing implementation for Kotlin code injections. While it’s convenient for simpler injection cases, it has its own downsides. Apart from the contextModule problem, code fragments do not support import and package directives. So there is a real probability we will reconsider this decision in the future.

Until that happens, you can try using the refinedContextModule. It’s the same as contextModule, but it doesn’t perform any code fragment checks. The property is annotated with @KaImplementationDetail as it’s unsafe, and we don’t expect any external usage of it. At the same time, we don’t plan to remove the property or change it in any way in the foreseeable future.

It would be great if you could share the reproducible sample (e.g., a commit in jte-intellij with step-by-step instructions). There might be a better way to solve the issue, but it’s hard for me to help without first debugging the code.

6

Thanks Yan. I have attached our performInjection() implementation (of LanguageInjectionPerformer class), where we utilize refinedContextModule. The Injector.inject() method in there is where our custom injector does its startInjecting(), addPlace() and doneInjecting()

performInjectionDemo.txt (2.3 KB)

We’ve made significant progress, thanks mostly to two breakthroughs. First is refinedContextModule, allowing us to attach the right KaSourceModule to the injected file. Thanks for that! Secondly is creating a synthetic code fragment in the Annotator based on the original injectedFile, but where we add the imports manually, then analyze that instead of the original injected file:

  val syntheticFragment = KtBlockCodeFragment(
      project = injectedFile.project,
      name = injectedFile.name,
      text = injectedFile.text,
      imports = importFqNames.joinToString(KtCodeFragment.IMPORT_SEPARATOR),
      context = injectedFile
  )

  analyze(syntheticFragment) {
   ...
  }

Nice little trick, however autocomplete is handled directly by the native Kotlin plugin, so we don’t have the opportunity to intercept the file being analyzed. And autocomplete does work well in many places - it handles most everything in stdlib, and it also handles parameters or variables defined via FQNs (val myAdder : com.example.Adder). But because the injected file doesn’t include import statements, autocomplete fails on variables with just the base class name.

Digging deep into the Kotlin repo, we found KtPsiFactory.createBlockCodeFragment() whose code constructs a KtCodeBlockFragment with imports hard-coded to null. We suspect this may be what the MultiHostInjector calls, explaining why imports are always ignored. If this is, in fact, the right function call (or something similar), then I propose a simple update. You could keep it undocumented like refinedContextModule, or requiring an Experimental annotation, or something…

// requires a new key to be defined, something like
val PRESERVE_IMPORTS_KEY : Key<List<String>> = Key.create<List<String>>(".z.y.x.preservedImports")


fun createBlockCodeFragment(@NonNls text: String, context: PsiElement?): KtBlockCodeFragment {
    val imports = context?.getUserData(PRESERVE_IMPORTS_KEY)
        ?.takeIf { it.isNotEmpty() }
        ?.joinToString(KtCodeFragment.IMPORT_SEPARATOR)

    /*  or, to make the default null path more explicit
    val imports = if (context?.getUserData(PRESERVE_IMPORTS_KEY) != null) {
      context.getUserData(PRESERVE_IMPORTS_KEY)!!
         .joinToString(KtCodeFragment.IMPORT_SEPARATOR)
    } else {
      null
    }
    */

    return KtBlockCodeFragment(project, "fragment.kt", text, imports, context)
}

and then before startInjecting() (or before doneInjecting()), the developer would manually find the imports, then set them via putUserData:
host.putUserData(PRESERVE_IMPORTS_KEY, customImportStatementList)

This doesn’t require any changes to function signatures, it’s just a new Key<List<String>> that gets checked, it will not affect existing code at all.

I’m not questioning your decision to not utilize imports in code fragments, however at the same time I believe we’ve brought up a valid use case for including them - a popular, Kotlin-native template engine with an existing user base that relied on a working IDE plugin that broke with the K2 compiler. Therefore, I (and many other JTE users) would greatly appreciate you considering the change above. Or something similar, hopefully also with a small footprint and zero effect to any existing code using these APIs. Thanks!

7

Please follow https://youtrack.jetbrains.com/issue/KTIJ-39750