Add support for external auth providers in code search

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/CodyAgentServer.kt

@@ -136,6 +136,8 @@ interface CodyAgentServer {
   fun testing_ignore_overridePolicy(params: ContextFilters?): CompletableFuture<Null?>
   @JsonRequest("extension/reset")
   fun extension_reset(params: Null?): CompletableFuture<Null?>
+  @JsonRequest("internal/getAuthHeaders")
+  fun internal_getAuthHeaders(params: String): CompletableFuture<Map<String, String>>
 
   // =============
   // Notifications

agent/scripts/reverse-proxy.py

@@ -22,11 +22,12 @@ async def proxy_handler(request):
             del headers['Transfer-Encoding']
 
         # Use value of 'Authorization: Bearer' to fill 'X-Forwarded-User' and remove 'Authorization' header
-        if 'Authorization' in headers:
-            match = re.match('Bearer (.*)', headers['Authorization'])
-            if match:
-                headers['X-Forwarded-User'] = match.group(1)
-            del headers['Authorization']
+
+        match = re.match('Bearer (.*)', headers['Authorization'])
+        if match:
+            headers['X-Forwarded-User'] = match.group(1)
+            if 'Authorization' in headers:
+                del headers['Authorization']
 
         # Forward the request to target
         async with session.request(

agent/src/agent.ts

@@ -14,6 +14,7 @@ import {
     currentAuthStatusAuthed,
     firstNonPendingAuthStatus,
     firstResultFromOperation,
+    getAuthHeaders,
     resolvedConfig,
     telemetryRecorder,
     waitUntilComplete,
@@ -1411,6 +1412,11 @@ export class Agent extends MessageHandler implements ExtensionClient {
             contextFiltersProvider.setTestingContextFilters(contextFilters)
             return null
         })
+
+        this.registerAuthenticatedRequest('internal/getAuthHeaders', async url => {
+            const config = await firstResultFromOperation(resolvedConfig)
+            return await getAuthHeaders(config.auth, new URL(url))
+        })
     }
 
     private pushPendingPromise(pendingPromise: Promise<unknown>): void {

jetbrains/build.gradle.kts

@@ -343,7 +343,9 @@ tasks {
       return destinationDir
     }
 
-    val sourcegraphDir = unzipCodeSearch()
+    val codeSearchDirOverride = System.getenv("CODE_SEARCH_DIR_OVERRIDE")
+    val sourcegraphDir: File =
+        if (codeSearchDirOverride != null) file(codeSearchDirOverride) else unzipCodeSearch()
     exec {
       workingDir(sourcegraphDir.toString())
       commandLine(*pnpmPath, "install", "--frozen-lockfile", "--fix-lockfile")

jetbrains/src/main/java/com/sourcegraph/find/FindPopupPanel.java

@@ -16,6 +16,7 @@
 import com.intellij.util.ui.JBUI;
 import com.intellij.util.ui.components.BorderLayoutPanel;
 import com.sourcegraph.Icons;
+import com.sourcegraph.cody.auth.CodyAuthService;
 import com.sourcegraph.common.NotificationGroups;
 import com.sourcegraph.common.ui.DumbAwareEDTAction;
 import com.sourcegraph.find.browser.BrowserAndLoadingPanel;
@@ -64,7 +65,10 @@ public FindPopupPanel(@NotNull Project project, @NotNull FindService findService
     browserAndLoadingPanel = new BrowserAndLoadingPanel(project);
     JSToJavaBridgeRequestHandler requestHandler =
         new JSToJavaBridgeRequestHandler(project, this, findService);
-    browser = JBCefApp.isSupported() ? new SourcegraphJBCefBrowser(requestHandler) : null;
+    String endpointUrl = CodyAuthService.getInstance(project).getEndpoint().getUrl();
+
+    browser =
+        JBCefApp.isSupported() ? new SourcegraphJBCefBrowser(requestHandler, endpointUrl) : null;
     if (browser == null) {
       showNoBrowserErrorNotification();
       Logger logger = Logger.getInstance(JSToJavaBridgeRequestHandler.class);

jetbrains/src/main/java/com/sourcegraph/find/FindService.java

@@ -11,6 +11,7 @@
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.wm.ex.WindowManagerEx;
 import com.intellij.util.ui.UIUtil;
+import com.sourcegraph.config.ConfigUtil;
 import com.sourcegraph.find.browser.BrowserAndLoadingPanel;
 import com.sourcegraph.find.browser.JavaToJSBridge;
 import java.awt.*;
@@ -31,6 +32,10 @@ public FindService(@NotNull Project project) {
     mainPanel = new FindPopupPanel(project, this);
   }
 
+  public static FindService getInstance(@NotNull Project project) {
+    return project.getService(FindService.class);
+  }
+
   public synchronized void showPopup() {
     createOrShowPopup();
   }
@@ -40,6 +45,15 @@ public void hidePopup() {
     hideMaterialUiOverlay();
   }
 
+  public void refreshConfiguration() {
+    JavaToJSBridge javaToJSBridge = mainPanel.getJavaToJSBridge();
+    if (javaToJSBridge != null) {
+      mainPanel
+          .getJavaToJSBridge()
+          .callJS("pluginSettingsChanged", ConfigUtil.getConfigAsJson(project));
+    }
+  }
+
   private void createOrShowPopup() {
     if (popup != null) {
       if (!popup.isVisible()) {

jetbrains/src/main/java/com/sourcegraph/find/browser/HttpSchemeHandler.java

@@ -1,9 +1,11 @@
 package com.sourcegraph.find.browser;
 
 import com.google.common.collect.ImmutableMap;
+import com.intellij.openapi.diagnostic.Logger;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.net.URL;
 import java.util.Map;
 import java.util.Optional;
 import org.cef.callback.CefCallback;
@@ -21,11 +23,19 @@ public class HttpSchemeHandler extends CefResourceHandlerAdapter {
   private int responseHeader = 400;
   private int offset = 0;
 
+  Logger logger = Logger.getInstance(HttpSchemeHandler.class);
+
   public boolean processRequest(@NotNull CefRequest request, @NotNull CefCallback callback) {
     String extension = getExtension(request.getURL());
     mimeType = getMimeType(extension);
     String url = request.getURL();
-    String path = url.replace("http://sourcegraph", "");
+    String path;
+    try {
+      path = new URL(url).getPath();
+    } catch (Exception ignored) {
+      logger.error("Failed to parse request url: " + url);
+      return false;
+    }
 
     if (mimeType != null) {
       data = loadResource(path);

jetbrains/src/main/java/com/sourcegraph/find/browser/JSToJavaBridgeRequestHandler.java

@@ -46,6 +46,8 @@ public JBCefJSQuery.Response handle(@NotNull JsonObject request) {
         case "getTheme":
           JsonObject currentThemeAsJson = ThemeUtil.getCurrentThemeAsJson();
           return createSuccessResponse(currentThemeAsJson);
+        case "getCustomRequestHeaders":
+          return createSuccessResponse(ConfigUtil.getAuthorizationHeadersAsJson(project));
         case "indicateSearchError":
           arguments = request.getAsJsonObject("arguments");
           // This must run on EDT (Event Dispatch Thread) because it changes the UI.

jetbrains/src/main/java/com/sourcegraph/find/browser/SourcegraphJBCefBrowser.java

@@ -2,7 +2,6 @@
 
 import com.intellij.openapi.util.Disposer;
 import com.intellij.ui.jcef.JBCefBrowser;
-import com.sourcegraph.cody.config.notification.CodySettingChangeListener;
 import com.sourcegraph.config.ThemeUtil;
 import javax.swing.*;
 import org.cef.CefApp;
@@ -11,23 +10,18 @@
 public class SourcegraphJBCefBrowser extends JBCefBrowser {
   private final JavaToJSBridge javaToJSBridge;
 
-  public SourcegraphJBCefBrowser(@NotNull JSToJavaBridgeRequestHandler requestHandler) {
-    super("http://sourcegraph/html/index.html");
-    // Create and set up JCEF browser
-    CefApp.getInstance()
-        .registerSchemeHandlerFactory("http", "sourcegraph", new HttpSchemeHandlerFactory());
+  public SourcegraphJBCefBrowser(
+      @NotNull JSToJavaBridgeRequestHandler requestHandler, String endpointUrl) {
+    super(endpointUrl.replaceAll("/+$", "").replace("https://", "http://") + "/html/index.html");
+
+    CefApp.getInstance().registerSchemeHandlerFactory("http", null, new HttpSchemeHandlerFactory());
 
     // Create bridges, set up handlers, then run init function
     String initJSCode = "window.initializeSourcegraph();";
     JSToJavaBridge jsToJavaBridge = new JSToJavaBridge(this, requestHandler, initJSCode);
     Disposer.register(this, jsToJavaBridge);
     javaToJSBridge = new JavaToJSBridge(this);
 
-    requestHandler
-        .getProject()
-        .getService(CodySettingChangeListener.class)
-        .setJavaToJSBridge(javaToJSBridge);
-
     UIManager.addPropertyChangeListener(
         propertyChangeEvent -> {
           if (propertyChangeEvent.getPropertyName().equals("lookAndFeel")) {

jetbrains/src/main/kotlin/com/sourcegraph/cody/auth/CodyAuthService.kt

@@ -4,6 +4,7 @@ import com.intellij.openapi.components.Service
 import com.intellij.openapi.components.service
 import com.intellij.openapi.project.Project
 import com.sourcegraph.config.ConfigUtil
+import com.sourcegraph.find.FindService
 
 @Service(Service.Level.PROJECT)
 class CodyAuthService(val project: Project) {
@@ -18,6 +19,7 @@ class CodyAuthService(val project: Project) {
 
   fun setActivated(isActivated: Boolean) {
     this.isActivated = isActivated
+    if (isActivated) FindService.getInstance(project).refreshConfiguration()
   }
 
   fun getEndpoint(): SourcegraphServerPath {

jetbrains/src/main/kotlin/com/sourcegraph/cody/config/notification/ChangeListener.kt

@@ -4,11 +4,9 @@ import com.intellij.openapi.Disposable
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.project.Project
 import com.intellij.util.messages.MessageBusConnection
-import com.sourcegraph.find.browser.JavaToJSBridge
 
 abstract class ChangeListener(protected val project: Project) : Disposable {
   protected val connection: MessageBusConnection = project.messageBus.connect()
-  var javaToJSBridge: JavaToJSBridge? = null
   protected val logger = Logger.getInstance(ChangeListener::class.java)
 
   override fun dispose() {

jetbrains/src/main/kotlin/com/sourcegraph/cody/config/notification/CodySettingChangeListener.kt

@@ -19,8 +19,6 @@ class CodySettingChangeListener(project: Project) : ChangeListener(project) {
         CodySettingChangeActionNotifier.TOPIC,
         object : CodySettingChangeActionNotifier {
           override fun afterAction(context: CodySettingChangeContext) {
-            // Notify JCEF about the config changes
-            javaToJSBridge?.callJS("pluginSettingsChanged", ConfigUtil.getConfigAsJson(project))
 
             if (context.oldCodyEnabled != context.newCodyEnabled) {
               if (context.newCodyEnabled) {

jetbrains/src/main/kotlin/com/sourcegraph/config/ConfigUtil.kt

@@ -11,16 +11,17 @@ import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.fileEditor.TextEditor
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.project.ProjectManager
+import com.sourcegraph.cody.agent.CodyAgentService
 import com.sourcegraph.cody.agent.protocol_generated.ExtensionConfiguration
 import com.sourcegraph.cody.auth.CodyAuthService
-import com.sourcegraph.cody.auth.CodySecureStore
 import com.sourcegraph.cody.auth.SourcegraphServerPath
 import com.sourcegraph.cody.config.CodyApplicationSettings
 import com.typesafe.config.ConfigFactory
 import com.typesafe.config.ConfigRenderOptions
 import com.typesafe.config.ConfigValueFactory
 import java.nio.file.Path
 import java.nio.file.Paths
+import java.util.concurrent.CompletableFuture
 import kotlin.io.path.readText
 import org.jetbrains.annotations.Contract
 import org.jetbrains.annotations.VisibleForTesting
@@ -110,12 +111,24 @@ object ConfigUtil {
   }
 
   @JvmStatic
-  fun getConfigAsJson(project: Project): JsonObject {
+  fun getAuthorizationHeadersAsJson(project: Project): JsonObject {
     val endpoint = CodyAuthService.getInstance(project).getEndpoint()
+    val authHeaders = CompletableFuture<Map<String, String>>()
+    CodyAgentService.withAgent(project) { agent ->
+      agent.server.internal_getAuthHeaders(endpoint.url).thenAccept { headers ->
+        authHeaders.complete(headers)
+      }
+    }
+
+    val jsonObject = JsonObject()
+    authHeaders.get().forEach { (key, value) -> jsonObject.addProperty(key, value) }
+    return jsonObject
+  }
+
+  @JvmStatic
+  fun getConfigAsJson(project: Project): JsonObject {
     return JsonObject().apply {
-      addProperty("instanceURL", endpoint.url)
-      addProperty("accessToken", CodySecureStore.getFromSecureStore(endpoint.url.toString()))
-      addProperty("customRequestHeadersAsString", "")
+      addProperty("instanceURL", CodyAuthService.getInstance(project).getEndpoint().url)
       addProperty("pluginVersion", getPluginVersion())
       addProperty("anonymousUserId", CodyApplicationSettings.instance.anonymousUserId)
     }

lib/shared/src/sourcegraph-api/utils.ts

@@ -36,17 +36,26 @@ export function toPartialUtf8String(buf: Buffer): { str: string; buf: Buffer } {
     }
 }
 
-export async function addAuthHeaders(auth: AuthCredentials, headers: Headers, url: URL): Promise<void> {
+export async function getAuthHeaders(auth: AuthCredentials, url: URL): Promise<Record<string, string>> {
     // We want to be sure we sent authorization headers only to the valid endpoint
     if (auth.credentials && url.host === new URL(auth.serverEndpoint).host) {
         if ('token' in auth.credentials) {
-            headers.set('Authorization', `token ${auth.credentials.token}`)
-        } else if (typeof auth.credentials.getHeaders === 'function') {
-            for (const [key, value] of Object.entries(await auth.credentials.getHeaders())) {
-                headers.set(key, value)
-            }
-        } else {
-            console.error('Cannot add headers: neither token nor headers found')
+            return { Authorization: `token ${auth.credentials.token}` }
+        }
+        if (typeof auth.credentials.getHeaders === 'function') {
+            return await auth.credentials.getHeaders()
         }
     }
+
+    console.error('Cannot add headers: neither token nor headers found')
+
+    return {}
+}
+
+export async function addAuthHeaders(auth: AuthCredentials, headers: Headers, url: URL): Promise<void> {
+    await getAuthHeaders(auth, url).then(authHeaders => {
+        for (const [key, value] of Object.entries(authHeaders)) {
+            headers.set(key, value)
+        }
+    })
 }

vscode/src/jsonrpc/agent-protocol.ts

@@ -277,6 +277,8 @@ export type ClientRequests = {
     // Called after the extension has been uninstalled by a user action.
     // Attempts to wipe out any state that the extension has stored.
     'extension/reset': [null, null]
+
+    'internal/getAuthHeaders': [string, Record<string, string>]
 }
 
 // ================