diff --git a/compiler-plugin/compiler-plugin-backend/build.gradle.kts b/compiler-plugin/compiler-plugin-backend/build.gradle.kts
index 59b567bbb..d97bd2a8c 100644
--- a/compiler-plugin/compiler-plugin-backend/build.gradle.kts
+++ b/compiler-plugin/compiler-plugin-backend/build.gradle.kts
@@ -7,6 +7,7 @@ import org.jetbrains.kotlin.gradle.dsl.ExplicitApiMode
 plugins {
     alias(libs.plugins.conventions.jvm)
     alias(libs.plugins.compiler.specific.module)
+    alias(libs.plugins.conventions.no.psi.element)
 }
 
 kotlin {
diff --git a/compiler-plugin/compiler-plugin-backend/src/main/kotlin/kotlinx/rpc/codegen/extension/RpcStubGenerator.kt b/compiler-plugin/compiler-plugin-backend/src/main/kotlin/kotlinx/rpc/codegen/extension/RpcStubGenerator.kt
index da1c1876b..eaa734485 100644
--- a/compiler-plugin/compiler-plugin-backend/src/main/kotlin/kotlinx/rpc/codegen/extension/RpcStubGenerator.kt
+++ b/compiler-plugin/compiler-plugin-backend/src/main/kotlin/kotlinx/rpc/codegen/extension/RpcStubGenerator.kt
@@ -1245,7 +1245,7 @@ internal class RpcStubGenerator(
         val protobufMessage = owner.getAnnotation(ctx.withCodecAnnotation.owner.kotlinFqName)
 
         return if (protobufMessage != null) {
-            val classReference = protobufMessage.arguments.single() as? IrClassReference
+            val classReference = vsApi{ protobufMessage.argumentsVS }.single() as? IrClassReference
                 ?: error("Expected IrClassReference for ${ctx.withCodecAnnotation.owner.kotlinFqName} parameter")
 
             val codec = classReference.classType
diff --git a/compiler-plugin/compiler-plugin-cli/build.gradle.kts b/compiler-plugin/compiler-plugin-cli/build.gradle.kts
index 3e09f1c08..94d20e580 100644
--- a/compiler-plugin/compiler-plugin-cli/build.gradle.kts
+++ b/compiler-plugin/compiler-plugin-cli/build.gradle.kts
@@ -7,6 +7,7 @@ import org.jetbrains.kotlin.gradle.dsl.ExplicitApiMode
 plugins {
     alias(libs.plugins.conventions.jvm)
     alias(libs.plugins.compiler.specific.module)
+    alias(libs.plugins.conventions.no.psi.element)
 }
 
 kotlin {
diff --git a/compiler-plugin/compiler-plugin-common/build.gradle.kts b/compiler-plugin/compiler-plugin-common/build.gradle.kts
index a45dff435..add3e55f3 100644
--- a/compiler-plugin/compiler-plugin-common/build.gradle.kts
+++ b/compiler-plugin/compiler-plugin-common/build.gradle.kts
@@ -7,6 +7,7 @@ import org.jetbrains.kotlin.gradle.dsl.ExplicitApiMode
 plugins {
     alias(libs.plugins.conventions.jvm)
     alias(libs.plugins.compiler.specific.module)
+    alias(libs.plugins.conventions.no.psi.element)
 }
 
 kotlin {
diff --git a/compiler-plugin/compiler-plugin-k2/build.gradle.kts b/compiler-plugin/compiler-plugin-k2/build.gradle.kts
index 31d312196..f1184b1ae 100644
--- a/compiler-plugin/compiler-plugin-k2/build.gradle.kts
+++ b/compiler-plugin/compiler-plugin-k2/build.gradle.kts
@@ -10,6 +10,7 @@ import util.whenKotlinCompilerIsAtLeast
 plugins {
     alias(libs.plugins.conventions.jvm)
     alias(libs.plugins.compiler.specific.module)
+    alias(libs.plugins.conventions.no.psi.element)
 }
 
 tasks.jar {
diff --git a/compiler-plugin/compiler-plugin-k2/src/main/kotlin/kotlinx/rpc/codegen/checkers/FirWithCodecDeclarationChecker.kt b/compiler-plugin/compiler-plugin-k2/src/main/kotlin/kotlinx/rpc/codegen/checkers/FirWithCodecDeclarationChecker.kt
index d587fc161..30b0ce601 100644
--- a/compiler-plugin/compiler-plugin-k2/src/main/kotlin/kotlinx/rpc/codegen/checkers/FirWithCodecDeclarationChecker.kt
+++ b/compiler-plugin/compiler-plugin-k2/src/main/kotlin/kotlinx/rpc/codegen/checkers/FirWithCodecDeclarationChecker.kt
@@ -57,13 +57,14 @@ object FirWithCodecDeclarationChecker {
 
         val codecTargetClass = codecClassSymbol.findMessageCodecSuperType(context.session)
             .typeArguments.first().type
+            ?: error("Unexpected unresolved type argument for @WithCodec annotation")
 
-        if (codecTargetClass?.classId != declaration.symbol.classId) {
+        if (codecTargetClass.classId != declaration.symbol.classId) {
             reporter.reportOn(
                 source = withCodec.findArgumentByName(CODEC_ARGUMENT_NAME)?.source,
                 factory = FirGrpcDiagnostics.CODEC_TYPE_MISMATCH,
                 a = declaration.symbol.defaultType(),
-                b = kClassValue,
+                b = codecTargetClass,
                 context = context,
             )
         }
diff --git a/compiler-plugin/compiler-plugin-k2/src/main/templates/kotlinx/rpc/codegen/checkers/diagnostics/RpcKtDiagnosticsContainer.kt b/compiler-plugin/compiler-plugin-k2/src/main/templates/kotlinx/rpc/codegen/checkers/diagnostics/RpcKtDiagnosticsContainer.kt
index a60e1e32a..557cd594a 100644
--- a/compiler-plugin/compiler-plugin-k2/src/main/templates/kotlinx/rpc/codegen/checkers/diagnostics/RpcKtDiagnosticsContainer.kt
+++ b/compiler-plugin/compiler-plugin-k2/src/main/templates/kotlinx/rpc/codegen/checkers/diagnostics/RpcKtDiagnosticsContainer.kt
@@ -32,6 +32,7 @@ fun registerDiagnosticRendererFactories() {
 //##csm specific=[2.0.0...2.2.10]
     RootDiagnosticRendererFactory.registerFactory(FirRpcDiagnostics.getRendererFactoryVs())
     RootDiagnosticRendererFactory.registerFactory(FirRpcStrictModeDiagnostics.getRendererFactoryVs())
+    RootDiagnosticRendererFactory.registerFactory(FirGrpcDiagnostics.getRendererFactoryVs())
 //##csm /specific
 //##csm /registerDiagnosticRendererFactories
 }
diff --git a/gradle-conventions/src/main/kotlin/conventions-no-psi-element.gradle.kts b/gradle-conventions/src/main/kotlin/conventions-no-psi-element.gradle.kts
new file mode 100644
index 000000000..1d40f28be
--- /dev/null
+++ b/gradle-conventions/src/main/kotlin/conventions-no-psi-element.gradle.kts
@@ -0,0 +1,21 @@
+/*
+ * Copyright 2023-2025 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+import kotlin.io.path.readLines
+
+gradle.afterProject {
+    layout.projectDirectory.dir("src").asFileTree.visit {
+        if (isDirectory) return@visit
+        if (file.name.endsWith(".kt")) {
+            file.toPath().readLines().forEach { line ->
+                if (!line.trim().startsWith("//") && line.contains("PsiElement")) {
+                    throw GradleException(
+                        "Kotlin source file $file must not contain `PsiElement` reference " +
+                                "in the compiler plugin. Use KtElement instead."
+                    )
+                }
+            }
+        }
+    }
+}
diff --git a/versions-root/libs.versions.toml b/versions-root/libs.versions.toml
index 2766af467..654b82282 100644
--- a/versions-root/libs.versions.toml
+++ b/versions-root/libs.versions.toml
@@ -159,6 +159,7 @@ conventions-npm = { id = "conventions-npm" }
 conventions-root = { id = "conventions-root" }
 conventions-dokka-spec = { id = "conventions-dokka-spec" }
 conventions-dokka-public = { id = "conventions-dokka-public" }
+conventions-no-psi-element = { id = "conventions-no-psi-element" }
 compiler-specific-module = { id = "compiler-specific-module" }
 
 # gradle-plugin project