Register (de)serializers

Some generated protocol messages require special handling.

src/main/kotlin/com/sourcegraph/cody/agent/CodyAgent.kt

@@ -112,20 +112,23 @@ private constructor(
         try {
           return server
               .initialize(
-                  ClientInfoFactory.build(
+                  ClientInfo(
+                      name = "JetBrains",
                       version = ConfigUtil.getPluginVersion(),
                       ideVersion = ApplicationInfo.getInstance().build.toString(),
                       workspaceRootUri =
                           ConfigUtil.getWorkspaceRootPath(project).toUri().toString(),
                       extensionConfiguration = ConfigUtil.getAgentConfiguration(project),
                       capabilities =
-                          ClientCapabilitiesFactory.build(
-                              edit = "enabled",
-                              editWorkspace = "enabled",
-                              codeLenses = "enabled",
-                              showDocument = "enabled",
-                              ignore = "enabled",
-                              untitledDocuments = "enabled")))
+                          ClientCapabilities(
+                              edit = ClientCapabilities.EditEnum.Enabled,
+                              editWorkspace = ClientCapabilities.EditWorkspaceEnum.Enabled,
+                              codeLenses = ClientCapabilities.CodeLensesEnum.Enabled,
+                              showDocument = ClientCapabilities.ShowDocumentEnum.Enabled,
+                              ignore = ClientCapabilities.IgnoreEnum.Enabled,
+                              untitledDocuments = ClientCapabilities.UntitledDocumentsEnum.Enabled,
+                              codeActions = ClientCapabilities.CodeActionsEnum.Enabled),
+                  ))
               .thenApply { info ->
                 logger.warn("Connected to Cody agent " + info.name)
                 server.initialized()
@@ -271,16 +274,24 @@ private constructor(
     ): Launcher<CodyAgentServer> {
       return Launcher.Builder<CodyAgentServer>()
           .configureGson { gsonBuilder ->
-            gsonBuilder
-                // emit `null` instead of leaving fields undefined because Cody
-                // VSC has many `=== null` checks that return false for undefined fields.
-                .serializeNulls()
-                .registerTypeAdapter(CompletionItemID::class.java, CompletionItemIDSerializer)
-                .registerTypeAdapter(ContextItem::class.java, ContextItem.deserializer)
-                .registerTypeAdapter(Speaker::class.java, speakerDeserializer)
-                .registerTypeAdapter(Speaker::class.java, speakerSerializer)
-                .registerTypeAdapter(URI::class.java, uriDeserializer)
-                .registerTypeAdapter(URI::class.java, uriSerializer)
+            run {
+              gsonBuilder
+                  // emit `null` instead of leaving fields undefined because Cody
+                  // VSC has many `=== null` checks that return false for undefined fields.
+                  .serializeNulls()
+                  .registerTypeAdapter(ContextItem::class.java, ContextItem.deserializer)
+                  .registerTypeAdapter(CompletionItemID::class.java, CompletionItemIDSerializer)
+                  // TODO: Remove legacy enum conversions
+                  .registerTypeAdapter(Speaker::class.java, speakerDeserializer)
+                  .registerTypeAdapter(Speaker::class.java, speakerSerializer)
+                  .registerTypeAdapter(URI::class.java, uriDeserializer)
+                  .registerTypeAdapter(URI::class.java, uriSerializer)
+
+              ProtocolTypeAdapters.register(gsonBuilder)
+              // This ensures that by default all enums are always serialized to their string
+              // equivalents
+              gsonBuilder.registerTypeAdapterFactory(EnumTypeAdapterFactory())
+            }
           }
           .setRemoteInterface(CodyAgentServer::class.java)
           .traceMessages(traceWriter())

src/main/kotlin/com/sourcegraph/cody/agent/EnumTypeAdapterFactory.kt

@@ -0,0 +1,57 @@
+package com.sourcegraph.cody.agent
+
+import com.google.gson.*
+import com.google.gson.annotations.SerializedName
+import com.google.gson.reflect.TypeToken
+import com.google.gson.stream.JsonReader
+import com.google.gson.stream.JsonWriter
+import java.io.IOException
+import java.lang.reflect.Field
+
+class EnumTypeAdapterFactory : TypeAdapterFactory {
+  override fun <T> create(gson: Gson, type: TypeToken<T>): TypeAdapter<T>? {
+    val rawType = type.rawType as? Class<*> ?: return null
+    if (!rawType.isEnum) {
+      return null
+    }
+    @Suppress("UNCHECKED_CAST")
+    return EnumTypeAdapter(rawType as Class<out Enum<*>>) as TypeAdapter<T>
+  }
+}
+
+class EnumTypeAdapter<T : Enum<T>>(private val classOfT: Class<T>) : TypeAdapter<T>() {
+  private val nameToConstant: Map<String, T> = HashMap()
+  private val constantToName: Map<T, String> = HashMap()
+
+  init {
+    for (constant in classOfT.enumConstants) {
+      val name = getSerializedName(constant) ?: constant.name
+      (nameToConstant as HashMap)[name.lowercase()] = constant
+      (constantToName as HashMap)[constant] = name
+    }
+  }
+
+  private fun getSerializedName(enumConstant: T): String? {
+    return try {
+      val field: Field = classOfT.getField(enumConstant.name)
+      field.getAnnotation(SerializedName::class.java)?.value
+    } catch (e: NoSuchFieldException) {
+      null
+    }
+  }
+
+  override fun write(out: JsonWriter, value: T?) {
+    if (value == null) {
+      out.nullValue()
+    } else {
+      out.value(constantToName[value])
+    }
+  }
+
+  @Throws(IOException::class)
+  override fun read(`in`: JsonReader): T? {
+    val value = `in`.nextString()
+    return nameToConstant[value.lowercase()]
+        ?: throw JsonParseException("Unknown enum value: $value")
+  }
+}