IntelliJ reacts to attribution enable/disable from site config

build.gradle.kts

@@ -50,6 +50,7 @@ dependencies {
   implementation("org.commonmark:commonmark-ext-gfm-tables:0.21.0")
   implementation("org.eclipse.lsp4j:org.eclipse.lsp4j.jsonrpc:0.21.0")
   implementation("com.googlecode.java-diff-utils:diffutils:1.3.0")
+  testImplementation("org.awaitility:awaitility-kotlin:4.2.0")
 }
 
 spotless {
@@ -386,4 +387,6 @@ tasks {
         listOf(
             properties("pluginVersion").split('-').getOrElse(1) { "default" }.split('.').first()))
   }
+
+  test { dependsOn(project.tasks.getByPath("buildCody")) }
 }

src/main/java/com/sourcegraph/cody/agent/CodyAgentClient.java

@@ -18,6 +18,8 @@ public class CodyAgentClient {
   private static final Logger logger = Logger.getInstance(CodyAgentClient.class);
   // Callback that is invoked when the agent sends a "chat/updateMessageInProgress" notification.
   @Nullable public Consumer<WebviewPostMessageParams> onNewMessage;
+  // Callback that is invoked when the agent sends a "setConfigFeatures" message.
+  @Nullable public ConfigFeaturesObserver onSetConfigFeatures;
   @Nullable public Editor editor;
 
   /**
@@ -66,5 +68,11 @@ public void webviewPostMessage(WebviewPostMessageParams params) {
       logger.debug("onNewMessage is null or message type is not transcript");
       logger.debug(String.format("webview/postMessage %s: %s", params.getId(), extensionMessage));
     }
+
+    if (onSetConfigFeatures != null
+        && extensionMessage.getType().equals(ExtensionMessage.Type.SET_CONFIG_FEATURES)) {
+      ApplicationManager.getApplication()
+          .invokeLater(() -> onSetConfigFeatures.update(extensionMessage.getConfigFeatures()));
+    }
   }
 }

src/main/java/com/sourcegraph/cody/agent/ConfigFeaturesObserver.java

@@ -0,0 +1,13 @@
+package com.sourcegraph.cody.agent;
+
+/**
+ * {@link ConfigFeaturesObserver} can be notified of changes in {@link ConfigFeatures} from the
+ * agent.
+ *
+ * <p>This can be attached to {@link CurrentConfigFeatures}, which multiplexes notifications from
+ * {@link CodyAgentClient#onConfigFeatures}.
+ */
+@FunctionalInterface
+public interface ConfigFeaturesObserver {
+  void update(ConfigFeatures newConfigFeatures);
+}

src/main/java/com/sourcegraph/cody/agent/CurrentConfigFeatures.java

@@ -0,0 +1,103 @@
+package com.sourcegraph.cody.agent;
+
+import com.intellij.openapi.components.Service;
+import com.sourcegraph.cody.vscode.CancellationToken;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
+
+/**
+ * {@link CurrentConfigFeatures} distributes the information about Cody feature configuration, like
+ * whether attribution is turned on or off, so that UI can adapt accordingly.
+ *
+ * <p>These features are turned on/off on the Sourcegraph instance and are fetched periodically by
+ * the agent.
+ *
+ * <p>Observers implementing {@link Consumer<ConfigFeatures>} can {@link #attach} and then will be
+ * notified of each config feature. Note: the observers will be notified irrespective of whether
+ * value of config feature is the same or different from the previous value, and need to
+ * de-duplicate individually if needed.
+ */
+@Service(Service.Level.PROJECT)
+public final class CurrentConfigFeatures implements ConfigFeaturesObserver {
+
+  /** Most recent {@link ConfigFeatures}. */
+  private final AtomicReference<ConfigFeatures> features =
+      new AtomicReference<>(new ConfigFeatures(false));
+
+  /**
+   * Observers that are attached (see {@link #attach}) and receive updates
+   * ({@link ConfigFeaturesObserver#update).
+   *
+   * <p>{@link IdentityObserver} is used here in order to provide precise
+   * dispose semantics (removal from this set) irrespecitve of {@link #equals}
+   * behavior on the delegate {@link ConfigFeaturesObserver}.
+   */
+  private final Set<IdentityObserver> observers = ConcurrentHashMap.newKeySet();
+
+  /** Retrieve the most recent {@link ConfigFeatures} value. */
+  public ConfigFeatures get() {
+    return features.get();
+  }
+
+  /**
+   * New {@link ConfigFeatures} arrive from the agent. This method updates state and notifies all
+   * observers.
+   */
+  @Override
+  public void update(ConfigFeatures configFeatures) {
+    this.features.set(configFeatures);
+    observers.forEach((observer) -> observer.update(configFeatures));
+  }
+
+  /**
+   * Given listener will be given new {@link ConfigFeatures} whenever they arrive. Observation
+   * relationship is ended once the returned cleanup {@link CancellationToken} is disposed.
+   */
+  public CancellationToken attach(ConfigFeaturesObserver observer) {
+    IdentityObserver id = new IdentityObserver(observer);
+    observers.add(id);
+    CancellationToken cancellation = new CancellationToken();
+    cancellation.onFinished((disposedOrAborted) -> observers.remove(id));
+    return cancellation;
+  }
+
+  /**
+   * {@link IdentityObserver} wraps {@link ConfigFeaturesObserver} reimplementing {@link #equals}
+   * with identity for precise cleanup. This way cleanup {@link Runnable} returned from {@link
+   * #attach} can drop precisely that given {@link ConfigFeaturesObserver} irrespective of the
+   * {@link #equals} semantics implemented.
+   */
+  private static class IdentityObserver implements ConfigFeaturesObserver {
+    final ConfigFeaturesObserver delegate;
+
+    IdentityObserver(ConfigFeaturesObserver delegate) {
+      this.delegate = delegate;
+    }
+
+    @Override
+    public void update(ConfigFeatures newConfigFeatures) {
+      delegate.update(newConfigFeatures);
+    }
+
+    @Override
+    public boolean equals(Object other) {
+      if (!(other instanceof IdentityObserver)) {
+        return false;
+      }
+      IdentityObserver that = (IdentityObserver) other;
+      return this.delegate == that.delegate;
+    }
+
+    /**
+     * {@link #delegate#hashCode} meets the {@link #equals} / {@link #hashCode} contract since
+     * {@link #equals} uses identity semantics on {@link IdentityObserver} which is stronger than
+     * identity semantics on {@link #delegate}.
+     */
+    @Override
+    public int hashCode() {
+      return delegate.hashCode();
+    }
+  }
+}

src/main/java/com/sourcegraph/cody/agent/WebviewPostMessageParams.kt

@@ -31,13 +31,19 @@ data class ExtensionMessage(
     val isTranscriptError: Boolean? = null,
     val customPrompts: List<List<Any>>? = null,
     val context: Any? = null,
-    val errors: String?
+    val errors: String?,
+    val configFeatures: ConfigFeatures? = null,
 ) {
 
   object Type {
     const val TRANSCRIPT = "transcript"
     const val ERRORS = "errors"
+    const val SET_CONFIG_FEATURES = "setConfigFeatures"
   }
 }
 
 data class WebviewPostMessageParams(val id: String, val message: ExtensionMessage)
+
+data class ConfigFeatures(
+    val attribution: Boolean,
+)

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

@@ -1,6 +1,7 @@
 package com.sourcegraph.cody.agent
 
 import com.intellij.ide.plugins.PluginManagerCore
+import com.intellij.openapi.components.service
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.extensions.PluginId
 import com.intellij.openapi.project.Project
@@ -94,6 +95,7 @@ private constructor(
       try {
         val conn = startAgentProcess()
         val client = CodyAgentClient()
+        client.onSetConfigFeatures = project.service<CurrentConfigFeatures>()
         val launcher = startAgentLauncher(conn, client)
         val server = launcher.remoteProxy
         val listeningToJsonRpc = launcher.startListening()

src/test/kotlin/com/sourcegraph/cody/cody/agent/CodyAgentClientTest.kt

@@ -0,0 +1,61 @@
+package com.sourcegraph.cody.cody.agent
+
+import com.intellij.testFramework.PlatformTestUtil
+import com.intellij.testFramework.fixtures.BasePlatformTestCase
+import com.sourcegraph.cody.agent.*
+import java.util.concurrent.locks.ReentrantLock
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@RunWith(JUnit4::class)
+class CodyAgentClientTest : BasePlatformTestCase() {
+  companion object {
+    const val WEBVIEW_ID: String = "unused-webview-id"
+  }
+
+  @Volatile var lastMessage: ConfigFeatures? = null
+  // Use lock/condition to synchronize between observer being invoked
+  // and the test being able to assert.
+  val lock = ReentrantLock()
+  val condition = lock.newCondition()
+
+  fun client(): CodyAgentClient {
+    val client = CodyAgentClient()
+    client.onSetConfigFeatures = ConfigFeaturesObserver {
+      lock.lock()
+      try {
+        lastMessage = it
+        condition.signal()
+      } finally {
+        lock.unlock()
+      }
+    }
+    return client
+  }
+
+  @Test
+  fun `notifies observer`() {
+    val expected = ConfigFeatures(attribution = true)
+    client()
+        .webviewPostMessage(
+            WebviewPostMessageParams(
+                id = WEBVIEW_ID,
+                message =
+                    ExtensionMessage(
+                        type = ExtensionMessage.Type.SET_CONFIG_FEATURES,
+                        errors = null,
+                        configFeatures = expected,
+                    )))
+    PlatformTestUtil.dispatchAllEventsInIdeEventQueue()
+    lock.lock()
+    try {
+      if (expected != lastMessage) {
+        condition.await()
+      }
+      assertEquals(expected, lastMessage)
+    } finally {
+      lock.unlock()
+    }
+  }
+}

src/test/kotlin/com/sourcegraph/cody/cody/agent/CurrentConfigFeaturesTest.kt

@@ -0,0 +1,86 @@
+package com.sourcegraph.cody.cody.agent
+
+import com.sourcegraph.cody.agent.ConfigFeatures
+import com.sourcegraph.cody.agent.ConfigFeaturesObserver
+import com.sourcegraph.cody.agent.CurrentConfigFeatures
+import java.util.concurrent.CopyOnWriteArrayList
+import org.awaitility.kotlin.await
+import org.awaitility.kotlin.until
+import org.hamcrest.CoreMatchers.hasItems
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@RunWith(JUnit4::class)
+class CurrentConfigFeaturesTest {
+
+  @Test
+  fun `observer is notified`() {
+    val current = CurrentConfigFeatures()
+    val observer = FakeObserver()
+    val expected = ConfigFeatures(attribution = true)
+    current.attach(observer)
+    current.update(expected)
+    assertThat(observer.features, hasItems(expected))
+  }
+
+  @Test
+  fun `observer is eventually not notified after cancelling`() {
+    val current = CurrentConfigFeatures()
+    val cancelledObserver = FakeObserver()
+    current.attach(cancelledObserver).dispose()
+    // dispose is async so got to await update
+    // (which is synchronous) not to take effect:
+    await until
+        {
+          val beforeCount = cancelledObserver.features.size
+          current.update(ConfigFeatures(attribution = true))
+          val afterCount = cancelledObserver.features.size
+          afterCount == beforeCount
+        }
+  }
+
+  @Test
+  fun `other observers keep being notified after one is cancelled`() {
+    val current = CurrentConfigFeatures()
+    val cancelledObserver = FakeObserver()
+    current.attach(cancelledObserver).dispose()
+    val activeObserver = FakeObserver()
+    current.attach(activeObserver) // No call to dispose.
+    // Wait until the cancelledObserver is cancelled (like previously).
+    await until
+        {
+          val beforeCount = cancelledObserver.features.size
+          current.update(ConfigFeatures(attribution = true))
+          val afterCount = cancelledObserver.features.size
+          afterCount == beforeCount
+        }
+    val beforeCount = activeObserver.features.size
+    current.update(ConfigFeatures(attribution = true))
+    val afterCount = activeObserver.features.size
+    assertEquals(beforeCount + 1, afterCount)
+  }
+
+  class FakeObserver : ConfigFeaturesObserver {
+    val features = CopyOnWriteArrayList<ConfigFeatures?>()
+
+    override fun update(newConfigFeatures: ConfigFeatures?) {
+      features.add(newConfigFeatures)
+    }
+
+    /**
+     * Using a total equality allows us to test for observing cancellation even when observer
+     * objects are equal.
+     */
+    override fun equals(other: Any?): Boolean {
+      return other is FakeObserver
+    }
+
+    /** Hash code meeting the contract of whole class being a single equality group. */
+    override fun hashCode(): Int {
+      return FakeObserver::class.hashCode()
+    }
+  }
+}