From 7bae16d4f273ebd7b51ee42960dc6d54d71320b7 Mon Sep 17 00:00:00 2001
From: Kalan <51868853+kalanchan@users.noreply.github.com>
Date: Tue, 12 Nov 2024 12:25:58 -0800
Subject: [PATCH] =?UTF-8?q?Fix=20extensionConfiguration/change=20type,=20r?=
 =?UTF-8?q?espect=20authStatus=20from=20cody=20=E2=80=A6=20(#2629)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

cherry picking (#2623)

## Changes

1. Bumped cody commit to match the one from `jb-v7.1.x` branch
2. Fixed `"extensionConfiguration/change"` endpoint type - issue was
introduced

[there](https://github.com/sourcegraph/jetbrains/commit/8ea63fc50cb91a78900df25ef2d9ceb8860082a2)
3. For the sake of correctness I also added check of returned auth state
and if auth fails I set the token as invalid.

## Test plan

1. Run IDE
4. Remove all account in the cody settings
5. Login panel shoul appear
6. Sign in to enterprise account
7. Make sure chat panel appears and in the account panel enterprise acc
is shown
8. Remove all account in the cody settings
9. Login panel should appear
10. Sign in to free dotcom account
11. Make sure chat panel appears and in the account panel free acc is
shown

## Test plan

<!-- All pull requests REQUIRE a test plan:
https://sourcegraph.com/docs/dev/background-information/testing_principles

Why does it matter?

These test plans are there to demonstrate that are following industry
standards which are important or critical for our customers.
They might be read by customers or an auditor. There are meant be simple
and easy to read. Simply explain what you did to ensure
your changes are correct!

Here are a non exhaustive list of test plan examples to help you:

- Making changes on a given feature or component:
- "Covered by existing tests" or "CI" for the shortest possible plan if
there is zero ambiguity
  - "Added new tests"
- "Manually tested" (if non trivial, share some output, logs, or
screenshot)
- Updating docs:
  - "previewed locally"
  - share a screenshot if you want to be thorough
- Updating deps, that would typically fail immediately in CI if
incorrect
  - "CI"
  - "locally tested"
-->

Co-authored-by: Piotr Kukiełka <pkukielka@users.noreply.github.com>
---
 .../sourcegraph/cody/agent/CodyAgentServer.kt |  2 +-
 .../cody/config/CodyAuthenticationManager.kt  | 36 +++++++++++++------
 2 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentServer.kt b/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentServer.kt
index ec424e743b..8dc29fe886 100644
--- a/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentServer.kt
+++ b/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentServer.kt
@@ -84,7 +84,7 @@ interface _SubsetGeneratedCodyAgentServer {
   @JsonRequest("extensionConfiguration/getSettingsSchema")
   fun extensionConfiguration_getSettingsSchema(params: Null?): CompletableFuture<String>
 
-  @JsonNotification("extensionConfiguration/change")
+  @JsonRequest("extensionConfiguration/change")
   fun extensionConfiguration_change(
       params: ExtensionConfiguration
   ): CompletableFuture<ProtocolAuthStatus?>
diff --git a/src/main/kotlin/com/sourcegraph/cody/config/CodyAuthenticationManager.kt b/src/main/kotlin/com/sourcegraph/cody/config/CodyAuthenticationManager.kt
index a73cbe42ad..5ff4c2b641 100644
--- a/src/main/kotlin/com/sourcegraph/cody/config/CodyAuthenticationManager.kt
+++ b/src/main/kotlin/com/sourcegraph/cody/config/CodyAuthenticationManager.kt
@@ -17,6 +17,7 @@ import com.intellij.openapi.project.ProjectManager
 import com.intellij.util.AuthData
 import com.intellij.util.concurrency.annotations.RequiresEdt
 import com.sourcegraph.cody.agent.CodyAgentService
+import com.sourcegraph.cody.agent.protocol_generated.ProtocolAuthenticatedAuthStatus
 import com.sourcegraph.cody.api.SourcegraphApiRequestExecutor
 import com.sourcegraph.cody.api.SourcegraphApiRequests
 import com.sourcegraph.cody.config.notification.AccountSettingChangeActionNotifier
@@ -212,8 +213,15 @@ class CodyAuthenticationManager :
       ProjectManager.getInstance().openProjects.forEach { project ->
         CodyAgentService.withAgentRestartIfNeeded(project) { agent ->
           if (!project.isDisposed) {
-            agent.server.extensionConfiguration_change(ConfigUtil.getAgentConfiguration(project))
-            publisher(project).afterAction(AccountSettingChangeContext(accessTokenChanged = true))
+            agent.server
+                .extensionConfiguration_change(ConfigUtil.getAgentConfiguration(project))
+                .thenApply { authStatus ->
+                  isTokenInvalid =
+                      CompletableFuture.completedFuture(
+                          authStatus !is ProtocolAuthenticatedAuthStatus)
+                  publisher(project)
+                      .afterAction(AccountSettingChangeContext(accessTokenChanged = true))
+                }
           }
         }
       }
@@ -236,15 +244,21 @@ class CodyAuthenticationManager :
     ProjectManager.getInstance().openProjects.forEach { project ->
       CodyAgentService.withAgentRestartIfNeeded(project) { agent ->
         if (!project.isDisposed) {
-          agent.server.extensionConfiguration_change(ConfigUtil.getAgentConfiguration(project))
-          if (serverUrlChanged || tierChanged || accountChanged) {
-            publisher(project)
-                .afterAction(
-                    AccountSettingChangeContext(
-                        serverUrlChanged = serverUrlChanged,
-                        accountTierChanged = tierChanged,
-                        accessTokenChanged = accountChanged))
-          }
+          agent.server
+              .extensionConfiguration_change(ConfigUtil.getAgentConfiguration(project))
+              .thenApply { authStatus ->
+                isTokenInvalid =
+                    CompletableFuture.completedFuture(
+                        authStatus !is ProtocolAuthenticatedAuthStatus)
+                if (serverUrlChanged || tierChanged || accountChanged) {
+                  publisher(project)
+                      .afterAction(
+                          AccountSettingChangeContext(
+                              serverUrlChanged = serverUrlChanged,
+                              accountTierChanged = tierChanged,
+                              accessTokenChanged = accountChanged))
+                }
+              }
         }
       }
     }