Fix keyOf to handle unresolved type variables

misk-inject/api/misk-inject.api

@@ -108,6 +108,7 @@ public final class misk/inject/FakeSwitchModule : misk/inject/KAbstractModule {
 
 public final class misk/inject/GuiceKt {
 	public static final fun asSingleton (Lcom/google/inject/binder/ScopedBindingBuilder;)V
+	public static final fun containsTypeVariable (Ljava/lang/reflect/Type;)Z
 	public static final fun getSetOf (Lcom/google/inject/Injector;Lkotlin/reflect/KClass;Lkotlin/reflect/KClass;)Ljava/util/Set;
 	public static synthetic fun getSetOf$default (Lcom/google/inject/Injector;Lkotlin/reflect/KClass;Lkotlin/reflect/KClass;ILjava/lang/Object;)Ljava/util/Set;
 	public static final fun listOfType (Lcom/google/inject/TypeLiteral;)Lcom/google/inject/TypeLiteral;

misk-inject/src/main/kotlin/misk/inject/Guice.kt

@@ -8,8 +8,10 @@ import com.google.inject.binder.ScopedBindingBuilder
 import com.google.inject.util.Types
 import jakarta.inject.Inject
 import jakarta.inject.Singleton
+import java.lang.reflect.GenericArrayType
 import java.lang.reflect.ParameterizedType
 import java.lang.reflect.Type
+import java.lang.reflect.TypeVariable
 import java.lang.reflect.WildcardType
 import kotlin.reflect.KClass
 import kotlin.reflect.KType
@@ -105,14 +107,41 @@ inline fun <reified T : Any> keyOf(a: KClass<out Annotation>?): Key<T> = keyOf(a
  * This is the primary function for creating qualified keys. It handles both type-based and instance-based qualifiers,
  * as well as unqualified bindings.
  *
+ * Uses [TypeLiteral] to preserve generic type information when the type is fully specified. Falls back to the erased
+ * class when the type contains unresolved type variables (e.g., when called with a non-reified type parameter from an
+ * enclosing generic class).
+ *
  * @param qualifier The [BindingQualifier] to use, or null for an unqualified binding
  * @return A [Key] for the specified type and qualifier
  */
-inline fun <reified T : Any> keyOf(qualifier: BindingQualifier? = null): Key<T> =
-  when (qualifier) {
-    is BindingQualifier.InstanceQualifier -> Key.get(object : TypeLiteral<T>() {}, qualifier.annotation)
-    is BindingQualifier.TypeClassifier -> Key.get(object : TypeLiteral<T>() {}, qualifier.type.java)
-    null -> Key.get(object : TypeLiteral<T>() {})
+inline fun <reified T : Any> keyOf(qualifier: BindingQualifier? = null): Key<T> {
+  val typeLiteral = object : TypeLiteral<T>() {}
+  return if (typeLiteral.type.containsTypeVariable()) {
+    when (qualifier) {
+      is BindingQualifier.InstanceQualifier -> Key.get(T::class.java, qualifier.annotation)
+      is BindingQualifier.TypeClassifier -> Key.get(T::class.java, qualifier.type.java)
+      null -> Key.get(T::class.java)
+    }
+  } else {
+    when (qualifier) {
+      is BindingQualifier.InstanceQualifier -> Key.get(typeLiteral, qualifier.annotation)
+      is BindingQualifier.TypeClassifier -> Key.get(typeLiteral, qualifier.type.java)
+      null -> Key.get(typeLiteral)
+    }
+  }
+}
+
+/**
+ * Checks whether a [Type] contains any unresolved [TypeVariable]s, recursively inspecting parameterized types, wildcard
+ * types, and generic array types.
+ */
+fun Type.containsTypeVariable(): Boolean =
+  when (this) {
+    is TypeVariable<*> -> true
+    is ParameterizedType -> actualTypeArguments.any { it.containsTypeVariable() }
+    is WildcardType -> upperBounds.any { it.containsTypeVariable() } || lowerBounds.any { it.containsTypeVariable() }
+    is GenericArrayType -> genericComponentType.containsTypeVariable()
+    else -> false
   }
 
 /**

misk-inject/src/test/kotlin/misk/inject/KeyOfTest.kt

@@ -0,0 +1,197 @@
+package misk.inject
+
+import com.google.inject.Guice
+import com.google.inject.Key
+import com.google.inject.TypeLiteral
+import com.google.inject.name.Names
+import java.lang.reflect.ParameterizedType
+import org.assertj.core.api.Assertions.assertThat
+import org.junit.jupiter.api.Test
+
+class KeyOfTest {
+  @Test
+  fun `keyOf with simple type creates key`() {
+    val key = keyOf<String>()
+    assertThat(key).isEqualTo(Key.get(String::class.java))
+  }
+
+  @Test
+  fun `keyOf with simple type preserves type via TypeLiteral`() {
+    val key = keyOf<String>()
+    // TypeLiteral-based key and class-based key are equal for non-generic types.
+    assertThat(key).isEqualTo(Key.get(object : TypeLiteral<String>() {}))
+  }
+
+  @Test
+  fun `keyOf with fully specified generic type preserves type parameters`() {
+    val key = keyOf<List<String>>()
+    // The key should use TypeLiteral and preserve the generic type parameter.
+    val typeLiteral = key.typeLiteral
+    assertThat(typeLiteral.type).isInstanceOf(ParameterizedType::class.java)
+    val paramType = typeLiteral.type as ParameterizedType
+    assertThat(paramType.rawType).isEqualTo(List::class.java)
+    // Kotlin's List<String> maps to java.util.List<? extends String> due to declaration-site variance.
+    assertThat(paramType.actualTypeArguments).hasSize(1)
+  }
+
+  @Test
+  fun `keyOf with fully specified generic type is distinct from erased type`() {
+    val genericKey = keyOf<List<String>>()
+    val erasedKey = Key.get(List::class.java)
+    // A TypeLiteral-based key with generics should differ from the erased class key.
+    assertThat(genericKey).isNotEqualTo(erasedKey)
+  }
+
+  @Test
+  fun `keyOf with annotation type qualifier`() {
+    val key = keyOf<String>(TestAnnotation::class)
+    assertThat(key.typeLiteral.rawType).isEqualTo(String::class.java)
+    assertThat(key.annotation).isNull()
+    assertThat(key.annotationType).isEqualTo(TestAnnotation::class.java)
+  }
+
+  @Test
+  fun `keyOf with annotation instance qualifier`() {
+    val named = Names.named("test")
+    val key = keyOf<String>(named)
+    assertThat(key.typeLiteral.rawType).isEqualTo(String::class.java)
+    assertThat(key.annotation).isEqualTo(named)
+  }
+
+  @Test
+  fun `keyOf with BindingQualifier TypeClassifier`() {
+    val qualifier = TestAnnotation::class.qualifier
+    val key = keyOf<String>(qualifier)
+    assertThat(key.annotationType).isEqualTo(TestAnnotation::class.java)
+  }
+
+  @Test
+  fun `keyOf with BindingQualifier InstanceQualifier`() {
+    val named = Names.named("test")
+    val qualifier = named.qualifier
+    val key = keyOf<String>(qualifier)
+    assertThat(key.annotation).isEqualTo(named)
+  }
+
+  @Test
+  fun `keyOf with null qualifier creates unqualified key`() {
+    val key = keyOf<String>(null as BindingQualifier?)
+    assertThat(key).isEqualTo(Key.get(String::class.java))
+    assertThat(key.annotation).isNull()
+    assertThat(key.annotationType).isNull()
+  }
+
+  @Test
+  fun `keyOf with generic type preserves type and can be used in Guice bindings`() {
+    val key = keyOf<List<String>>()
+    val injector =
+      Guice.createInjector(
+        object : KAbstractModule() {
+          override fun configure() {
+            bind(key).toInstance(listOf("hello", "world"))
+          }
+        }
+      )
+    val result = injector.getInstance(key)
+    assertThat(result).containsExactly("hello", "world")
+  }
+
+  /**
+   * This is the regression test for the bug fixed in this PR. When `keyOf` is called from a generic class with an
+   * unresolved type variable (e.g., `keyOf<SomeType<T>>()` where T is a class type parameter, not reified), the
+   * TypeLiteral captures an unresolved TypeVariable. Without the fix, Guice would throw `Key is not fully specified`
+   * because it cannot handle TypeVariables.
+   *
+   * The fix detects unresolved type variables and falls back to the erased class, which Guice can handle.
+   */
+  @Test
+  fun `keyOf with unresolved type variable falls back to erased class`() {
+    // Simulate the pattern from BatchReceiverModule: a generic class calling keyOf<Wrapper<T>>()
+    // where T is a class type parameter (not reified).
+    val helper = GenericKeyCreator<String>()
+    // This should not throw "Key is not fully specified".
+    val key = helper.createKey()
+    // The key should use the erased class (Wrapper) since the type variable can't be resolved.
+    assertThat(key.typeLiteral.rawType).isEqualTo(Wrapper::class.java)
+  }
+
+  @Test
+  fun `keyOf with unresolved type variable and annotation type qualifier`() {
+    val helper = GenericKeyCreator<String>()
+    val key = helper.createKeyWithAnnotationType()
+    assertThat(key.typeLiteral.rawType).isEqualTo(Wrapper::class.java)
+    assertThat(key.annotationType).isEqualTo(TestAnnotation::class.java)
+  }
+
+  @Test
+  fun `keyOf with unresolved type variable and annotation instance qualifier`() {
+    val helper = GenericKeyCreator<String>()
+    val named = Names.named("test")
+    val key = helper.createKeyWithAnnotationInstance(named)
+    assertThat(key.typeLiteral.rawType).isEqualTo(Wrapper::class.java)
+    assertThat(key.annotation).isEqualTo(named)
+  }
+
+  @Test
+  fun `keyOf with unresolved type variable can be used in Guice bindings`() {
+    val helper = GenericKeyCreator<String>()
+    val key = helper.createKey()
+    // The erased key can be bound and retrieved from an injector.
+    val injector =
+      Guice.createInjector(
+        object : KAbstractModule() {
+          override fun configure() {
+            bind(key).toInstance(Wrapper("hello"))
+          }
+        }
+      )
+    @Suppress("UNCHECKED_CAST") val result = injector.getInstance(key) as Wrapper<String>
+    assertThat(result.value).isEqualTo("hello")
+  }
+
+  @Test
+  fun `containsTypeVariable returns false for simple class`() {
+    assertThat(String::class.java.containsTypeVariable()).isFalse()
+  }
+
+  @Test
+  fun `containsTypeVariable returns false for fully specified parameterized type`() {
+    val type = object : TypeLiteral<List<String>>() {}.type
+    assertThat(type.containsTypeVariable()).isFalse()
+  }
+
+  @Test
+  fun `containsTypeVariable returns true for type variable`() {
+    // Get a TypeVariable from a generic class's type parameters.
+    val typeVar = Wrapper::class.java.typeParameters[0]
+    assertThat(typeVar.containsTypeVariable()).isTrue()
+  }
+
+  @Test
+  fun `containsTypeVariable returns true for parameterized type with type variable`() {
+    // Get a field type like Wrapper<T> where T is unresolved.
+    val wrapperField = GenericFieldHolder::class.java.getDeclaredField("field")
+    val fieldType = wrapperField.genericType
+    assertThat(fieldType.containsTypeVariable()).isTrue()
+  }
+
+  /** A generic wrapper class used to test keyOf with unresolved type variables. */
+  class Wrapper<T>(val value: T)
+
+  /**
+   * Simulates the pattern from BatchReceiverModule where a generic class calls keyOf<SomeType<T>>() with an unresolved
+   * class type parameter T.
+   */
+  class GenericKeyCreator<T : Any> {
+    fun createKey(): Key<Wrapper<T>> = keyOf<Wrapper<T>>()
+
+    fun createKeyWithAnnotationType(): Key<Wrapper<T>> = keyOf<Wrapper<T>>(TestAnnotation::class)
+
+    fun createKeyWithAnnotationInstance(annotation: Annotation): Key<Wrapper<T>> = keyOf<Wrapper<T>>(annotation)
+  }
+
+  /** Helper class with a generic field for testing containsTypeVariable on field types. */
+  class GenericFieldHolder<T> {
+    @JvmField var field: Wrapper<T>? = null
+  }
+}