Reimplement agent lifecycle management #303

pkukielka · 2024-01-15T14:48:29Z

Test plan

This PR will require extensive testing of all existing features.
We need to go through whole https://github.com/sourcegraph/jetbrains/blob/main/TESTING.md
Also ideally use this changes locally for some time to try cach all kind of timing/lifecycle bugs.

Changes

Agent lifecycle is now managed fully by CodyAgentService.

Classes using agent DO NOT try to validate it state, restart, wait few seconds before doing some action, etc.
They also not do any assumptions about agent state anymore.
All they know/get is a CompletableFuture which will contain the agent when it's ready.

If agent will loose connection or crash in some other way it should be restarted automatically by CodyAgentService.

Other than that lifetime management and accesability of several objects components (like CodyAgentCodebase or server object) were changed.

Also, maybe not a huge win but:
30 files changed, 576 insertions(+), 659 deletions(-)
:)

pkukielka · 2024-01-15T14:49:16Z

src/main/java/com/sourcegraph/config/ApiAuthenticator.java

@@ -1,20 +0,0 @@
-package com.sourcegraph.config;


It was not used in any place.

mkondratek · 2024-01-15T16:01:02Z

src/main/java/com/sourcegraph/cody/CodyFileEditorListener.java

-              CodyAgent.getInitializedServer(source.getProject())
-                  // The timeout has been increased from 3 to 12.
-                  // This more like workaround than a fix to:
-                  // https://github.com/sourcegraph/jetbrains/issues/169


I assume we can say that this PR actually fixes: #169

There is a good chance for that.
Fas there any deterministic reproduction for the issue above?

I am afraid there was not 😕 I could have been related to the network

mkondratek · 2024-01-15T17:02:00Z

src/main/kotlin/com/sourcegraph/cody/config/notification/AccountSettingChangeListener.kt

-            if (ConfigUtil.isCodyEnabled() && agentServer != null) {
-              agentServer.configurationDidChange(ConfigUtil.getAgentConfiguration(project))
+            CodyAgentService.applyAgentOnBackgroundThread(project) { agent ->
+              if (ConfigUtil.isCodyEnabled()) {


so previously depending on this boolean we were starting or stopping the agent. Is the behavior different now? so is the agent always started?

That is something I take liberty to remove but we have to test it. Agent is always started (or if not yet CodyAgentService.applyAgentOnBackgroundThread will wait for it to be).
I think starting/stoping should take place when cody settings change, as that is where you can enable/disable cody (and we have proper code for that in settings).

But if someone see a reason why that code was needed there I will be happy to hear it.

mkondratek · 2024-01-15T17:04:07Z

src/main/kotlin/com/sourcegraph/common/UpgradeToCodyProNotification.kt

+            .thenCompose { agent ->
+              agent.server.evaluateFeatureFlag(GetFeatureFlag("CodyProJetBrains"))
+            }
+            // TODO We should avoid using getNow if possible


can we have an issue for this one? or does it make sense to just add timeout(1, SECONDS)?

Timeout does not change anything there, because there is no quarantee if it will take 0, 2 or 10 sec.
I will create an issue.

mkondratek · 2024-01-15T17:09:12Z

src/main/kotlin/com/sourcegraph/cody/context/EmbeddingStatusView.kt

-    val tooltip = embeddingStatus.getTooltip(project)
-    if (tooltip.isNotEmpty()) {
-      embeddingStatusContent.setToolTipText(tooltip)
+    ApplicationManager.getApplication().invokeLater {


I see all the logic is wrapped in invokeLater. Shouldn't be a caller's responsibility to call it on a proper thread? We could then mark it as @RequiresEdt.

I have mixing feelings about that, let's discuss on the next standup.
I think this can go as it is for now.

mkondratek · 2024-01-15T17:18:49Z

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

+import javax.annotation.concurrent.GuardedBy
+
+@Service(Service.Level.APP)
+class CodyAgentService : Disposable {


beautiful 😍

mkondratek · 2024-01-15T17:25:34Z

src/main/kotlin/com/sourcegraph/cody/CodyToolWindowContent.kt

+        // Agent cannot gracefully recover when connection is lost, we need to restart it
+        logger.warn("Failed to load new chat, restarting agent", e)
+        CodyAgentService.getInstance(project).restartAgent(project)
+        Thread.sleep(5000)


is it needed to sleep here? startAgent has it's 15s timeout already ⌚

I wanted to avoid constant heavier utilisation of the CPU in case of the lack of the network.
In fact it would be better to have some listener which would say if network is back, I will check if we can do that. I will also move that code to agent service.

mkondratek · 2024-01-15T17:29:28Z

src/main/kotlin/com/sourcegraph/cody/CodyToolWindowContent.kt

-      if (server == null) {
-        setRecipesPanelError()
-      }
+    CodyAgentService.applyAgentOnBackgroundThread(project) { agent ->


I guess we can simplify the code here even more. Let's stick to the rule that it is the caller's responsibility to use a proper thread, hence:

remove CodyAgentService.applyAgentOnBackgroundThread(project) { ... } wrapping here

remove ApplicationManager.getApplication().executeOnPooledThread { ... } wrapping above (line 225, in refreshRecipes)

revamp loadCommands so it takes agent as a param

wrap loadCommands(agent) into CodyAgentService.applyAgentOnBackgroundThread(project) { ... } in line 225, refeshRecipes)

I agree that we can simplify this code a bit

I do not agree with Let's stick to the rule that it is the caller's responsibility to use a proper thread as a general rule, it can lead to too many errors/omissions. That is, unless we explicitly wrap methods bory in something like CompletableFuture, but with a distinct subtypes for EDT and Background (which I think could work nicely but it's a bit additional work and I'm not sure if it's currently worth it).

I'm happy to discuss this on the standup or a dedicated call.

ok, feel free to do it your way

exigow · 2024-01-16T09:15:34Z

src/main/kotlin/com/sourcegraph/cody/chat/Chat.kt

+        if (isFirstMessage.compareAndSet(false, true)) {
+          val contextMessages =
+              agentChatMessage.contextFiles
+                  ?.stream()


We can remove stream + collect

exigow · 2024-01-16T13:40:09Z

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

+private constructor(
+    val client: CodyAgentClient,
+    val server: CodyAgentServer,
+    val launcher: Launcher<CodyAgentServer>,


This lanucher is unused - can we remove that?

exigow · 2024-01-16T13:43:10Z

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

+
+  fun initializeRepoName() {
+    onRepositoryNameChange(RepoUtil.findRepositoryName(project, null))
+  }


This method should be placed below class fields.

exigow · 2024-01-16T13:43:42Z

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

@@ -9,19 +11,18 @@ import com.sourcegraph.common.ProjectFileUtils
 import com.sourcegraph.config.ConfigUtil
 import com.sourcegraph.vcs.RepoUtil

-class CodyAgentCodebase(private val underlying: CodyAgentServer, val project: Project) {
+@Service(Service.Level.PROJECT)
+class CodyAgentCodebase(val project: Project) {


Project should be private.

exigow

Looks good 👍.

I've added some minor comments and I'm looking forward for Cody commit bump to start testing. I understand the idea of these changes, especially after your presentation. I like this CodyAgentService which really simplifies usages (as confirmed by the net number of lines) and and forces us to think more about threading.

Sa1to

LGTM 🎉 FYI rebased #183 on this branch

pkukielka commented Jan 15, 2024

View reviewed changes

Reimplement agent lifecycle management

6a2ff15

pkukielka force-pushed the pkukielka/reimplement-agent-lifecycle-management branch from 379d969 to 6a2ff15 Compare January 15, 2024 14:52

exigow requested a review from a team January 15, 2024 15:00

pkukielka added 3 commits January 15, 2024 17:01

Minor fixes

2a2d86e

Refactor using spotless

30206af

Fix initial repo setting

c39fc72

mkondratek reviewed Jan 15, 2024

View reviewed changes

exigow assigned pkukielka Jan 16, 2024

exigow reviewed Jan 16, 2024

View reviewed changes

pkukielka mentioned this pull request Jan 16, 2024

Revisit and cleanup multithreaded code #307

Open

mkondratek approved these changes Jan 16, 2024

View reviewed changes

Few fixes for review

c494d27

exigow reviewed Jan 16, 2024

View reviewed changes

exigow approved these changes Jan 16, 2024

View reviewed changes

Fix for test

11fee97

Sa1to approved these changes Jan 16, 2024

View reviewed changes

Reformat

cab22bb

mkondratek approved these changes Jan 16, 2024

View reviewed changes

pkukielka merged commit 77e5c2f into main Jan 17, 2024
2 checks passed

pkukielka deleted the pkukielka/reimplement-agent-lifecycle-management branch January 17, 2024 09:23

mkondratek mentioned this pull request Jan 18, 2024

Investigate logging & telemetry disk usage #315

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Reimplement agent lifecycle management #303

Reimplement agent lifecycle management #303

pkukielka commented Jan 15, 2024 •

edited

Loading

pkukielka Jan 15, 2024

mkondratek Jan 15, 2024

pkukielka Jan 16, 2024

mkondratek Jan 16, 2024

mkondratek Jan 15, 2024

pkukielka Jan 16, 2024

mkondratek Jan 15, 2024

pkukielka Jan 16, 2024

mkondratek Jan 15, 2024

pkukielka Jan 16, 2024

mkondratek Jan 15, 2024

mkondratek Jan 15, 2024

pkukielka Jan 16, 2024 •

edited

Loading

mkondratek Jan 15, 2024

pkukielka Jan 16, 2024

mkondratek Jan 16, 2024

exigow Jan 16, 2024 •

edited

Loading

exigow Jan 16, 2024

exigow Jan 16, 2024

exigow Jan 16, 2024

exigow left a comment

Sa1to left a comment

Reimplement agent lifecycle management #303

Reimplement agent lifecycle management #303

Conversation

pkukielka commented Jan 15, 2024 • edited Loading

Test plan

Changes

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pkukielka Jan 16, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

exigow Jan 16, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

exigow left a comment

Choose a reason for hiding this comment

Sa1to left a comment

Choose a reason for hiding this comment

pkukielka commented Jan 15, 2024 •

edited

Loading

pkukielka Jan 16, 2024 •

edited

Loading

exigow Jan 16, 2024 •

edited

Loading