IntelliJ reacts to attribution enable/disable from site config

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,12 @@
+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}.
+ */
+public interface ConfigFeaturesObserver {
+  void update(ConfigFeatures newConfigFeatures);
+}

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

@@ -0,0 +1,100 @@
+package com.sourcegraph.cody.agent;
+
+import com.intellij.openapi.components.Service;
+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 Runnable} is {@link Runnable#run()}.
+   */
+  public Runnable attach(ConfigFeaturesObserver observer) {
+    IdentityObserver id = new IdentityObserver(observer);
+    observers.add(id);
+    return () -> observers.remove(id);
+  }
+
+  /**
+   * {@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

@@ -10,8 +10,8 @@ import com.sourcegraph.cody.agent.protocol.ContextFile
  */
 data class WebviewMessage(
     val command: String,
-    val text: String? = null,
-    val submitType: String? = null, // One of: "user", "suggestion", "example"
+    val text: String,
+    val submitType: String, // One of: "user", "suggestion", "example"
     val addEnhancedContext: Boolean? = null,
     val contextFiles: List<ContextFile>? = null,
     val error: ChatError? = null,
@@ -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?,
 ) {
 
   object Type {
     const val TRANSCRIPT = "transcript"
     const val ERRORS = "errors"
+    const val SET_CONFIG_FEATURES = "setConfigFeatures"
   }
 }
 
+data class ConfigFeatures(
+    val attribution: Boolean,
+)
+
 data class WebviewPostMessageParams(val id: String, val message: ExtensionMessage)

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()