diff --git a/build.gradle.kts b/build.gradle.kts index e9d199f93d..2e0c28ea1a 100644 --- a/build.gradle.kts +++ b/build.gradle.kts @@ -152,7 +152,7 @@ spotless { ktfmt() trimTrailingWhitespace() target("src/**/*.kt") - targetExclude("src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/**/*.kt") + targetExclude("src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/**") toggleOffOn() } } @@ -499,6 +499,7 @@ tasks { buildPlugin { dependsOn(project.tasks.getByPath("buildCody")) + composedJar.get().exclude("com/intellij/codeInsight/inline/completion/**") from( fileTree(buildCodyDir) { include("*") @@ -580,11 +581,6 @@ tasks { test { dependsOn(project.tasks.getByPath("buildCody")) } - configurations { - create("integrationTestImplementation") { extendsFrom(configurations.testImplementation.get()) } - create("integrationTestRuntimeClasspath") { extendsFrom(configurations.testRuntimeOnly.get()) } - } - sourceSets { create("integrationTest") { kotlin.srcDir("src/integrationTest/kotlin") diff --git a/src/main/kotlin/com/intellij/codeInsight/inline/completion/InlineCompletionEvent.kt b/src/main/kotlin/com/intellij/codeInsight/inline/completion/InlineCompletionEvent.kt new file mode 100644 index 0000000000..8965e6ba30 --- /dev/null +++ b/src/main/kotlin/com/intellij/codeInsight/inline/completion/InlineCompletionEvent.kt @@ -0,0 +1,92 @@ +// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the +// Apache 2.0 license. +package com.intellij.codeInsight.inline.completion + +import com.intellij.codeInsight.lookup.LookupElement +import com.intellij.codeInsight.lookup.LookupEvent +import com.intellij.openapi.actionSystem.DataContext +import com.intellij.openapi.application.runReadAction +import com.intellij.openapi.editor.Caret +import com.intellij.openapi.editor.Document +import com.intellij.openapi.editor.Editor +import com.intellij.openapi.project.Project +import com.intellij.openapi.util.UserDataHolderBase +import com.intellij.psi.PsiDocumentManager +import com.intellij.psi.PsiFile +import com.intellij.psi.impl.source.PsiFileImpl +import com.intellij.psi.util.PsiUtilBase +import com.intellij.util.concurrency.annotations.RequiresBlockingContext + +class InlineCompletionRequest( + val event: InlineCompletionEvent, + val file: PsiFile, + val editor: Editor, + val document: Document, + val startOffset: Int, + val endOffset: Int, + val lookupElement: LookupElement? = null, +) : UserDataHolderBase() + +/** + * Be aware that creating your own event is unsafe for a while and might face compatibility issues + */ +interface InlineCompletionEvent { + + @RequiresBlockingContext fun toRequest(): InlineCompletionRequest? + + /** A class representing a direct call in the code editor by [InsertInlineCompletionAction]. */ + class DirectCall( + val editor: Editor, + val caret: Caret, + val context: DataContext? = null, + ) : InlineCompletionEvent { + override fun toRequest(): InlineCompletionRequest? { + val offset = runReadAction { caret.offset } + val project = editor.project ?: return null + val file = getPsiFile(caret, project) ?: return null + return InlineCompletionRequest(this, file, editor, editor.document, offset, offset) + } + } + + sealed interface InlineLookupEvent : InlineCompletionEvent { + val event: LookupEvent + + override fun toRequest(): InlineCompletionRequest? { + val editor = runReadAction { event.lookup?.editor } ?: return null + val caretModel = editor.caretModel + if (caretModel.caretCount != 1) return null + + val project = editor.project ?: return null + + val (file, offset) = + runReadAction { getPsiFile(caretModel.currentCaret, project) to caretModel.offset } + if (file == null) return null + + return InlineCompletionRequest( + this, file, editor, editor.document, offset, offset, event.item) + } + } +} + +@RequiresBlockingContext +private fun getPsiFile(caret: Caret, project: Project): PsiFile? { + return runReadAction { + val file = + PsiDocumentManager.getInstance(project).getPsiFile(caret.editor.document) + ?: return@runReadAction null + // * [PsiUtilBase] takes into account injected [PsiFile] (like in Jupyter Notebooks) + // * However, it loads a file into the memory, which is expensive + // * Some tests forbid loading a file when tearing down + // * On tearing down, Lookup Cancellation happens, which causes the event + // * Existence of [treeElement] guarantees that it's in the memory + if (file.isLoadedInMemory()) { + PsiUtilBase.getPsiFileInEditor(caret, project) + } else { + file + } + } +} + +private fun PsiFile.isLoadedInMemory(): Boolean { + return (this as? PsiFileImpl)?.treeElement != null +} diff --git a/src/main/kotlin/com/intellij/codeInsight/inline/completion/InlineCompletionSuggestion.kt b/src/main/kotlin/com/intellij/codeInsight/inline/completion/InlineCompletionSuggestion.kt new file mode 100644 index 0000000000..23c99b3e15 --- /dev/null +++ b/src/main/kotlin/com/intellij/codeInsight/inline/completion/InlineCompletionSuggestion.kt @@ -0,0 +1,34 @@ +// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the +// Apache 2.0 license. +package com.intellij.codeInsight.inline.completion + +import com.intellij.codeInsight.inline.completion.elements.InlineCompletionElement +import com.intellij.openapi.util.UserDataHolderBase +import kotlinx.coroutines.flow.Flow +import kotlinx.coroutines.flow.FlowCollector +import kotlinx.coroutines.flow.emptyFlow +import kotlinx.coroutines.flow.flow + +/** + * Abstract class representing an inline completion suggestion. + * + * Provides the suggestion flow for generating only one suggestion. + * + * @see InlineCompletionElement + */ +abstract class InlineCompletionSuggestion : UserDataHolderBase() { + abstract val suggestionFlow: Flow + + class Default(override val suggestionFlow: Flow) : + InlineCompletionSuggestion() + + companion object { + fun empty(): InlineCompletionSuggestion = Default(emptyFlow()) + + fun withFlow( + buildSuggestion: suspend FlowCollector.() -> Unit + ): InlineCompletionSuggestion { + return Default(flow(buildSuggestion)) + } + } +} diff --git a/src/main/kotlin/com/intellij/codeInsight/inline/completion/elements/InlineCompletionElement.kt b/src/main/kotlin/com/intellij/codeInsight/inline/completion/elements/InlineCompletionElement.kt new file mode 100644 index 0000000000..b9a181bb83 --- /dev/null +++ b/src/main/kotlin/com/intellij/codeInsight/inline/completion/elements/InlineCompletionElement.kt @@ -0,0 +1,5 @@ +// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the +// Apache 2.0 license. +package com.intellij.codeInsight.inline.completion.elements + +interface InlineCompletionElement diff --git a/src/main/kotlin/com/intellij/codeInsight/inline/completion/elements/InlineCompletionGrayTextElement.kt b/src/main/kotlin/com/intellij/codeInsight/inline/completion/elements/InlineCompletionGrayTextElement.kt new file mode 100644 index 0000000000..2a9097d86b --- /dev/null +++ b/src/main/kotlin/com/intellij/codeInsight/inline/completion/elements/InlineCompletionGrayTextElement.kt @@ -0,0 +1,5 @@ +// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the +// Apache 2.0 license. +package com.intellij.codeInsight.inline.completion.elements + +class InlineCompletionGrayTextElement(val text: String) : InlineCompletionElement diff --git a/src/main/kotlin/com/sourcegraph/cody/autocomplete/CodyAutocompleteManager.kt b/src/main/kotlin/com/sourcegraph/cody/autocomplete/CodyAutocompleteManager.kt index 0e40bb1d6f..b369e981c1 100644 --- a/src/main/kotlin/com/sourcegraph/cody/autocomplete/CodyAutocompleteManager.kt +++ b/src/main/kotlin/com/sourcegraph/cody/autocomplete/CodyAutocompleteManager.kt @@ -5,6 +5,7 @@ import com.github.difflib.patch.DeltaType import com.github.difflib.patch.Patch import com.intellij.codeInsight.hint.HintManager import com.intellij.openapi.application.ApplicationManager +import com.intellij.openapi.client.ClientSessionsManager import com.intellij.openapi.command.CommandProcessor import com.intellij.openapi.command.WriteCommandAction import com.intellij.openapi.components.Service @@ -14,7 +15,6 @@ import com.intellij.openapi.editor.Editor import com.intellij.openapi.editor.Inlay import com.intellij.openapi.editor.InlayModel import com.intellij.openapi.editor.colors.EditorColorsManager -import com.intellij.openapi.fileEditor.FileDocumentManager import com.intellij.openapi.keymap.KeymapUtil import com.intellij.openapi.project.Project import com.intellij.openapi.ui.popup.Balloon @@ -27,29 +27,15 @@ import com.intellij.util.concurrency.annotations.RequiresEdt import com.sourcegraph.cody.Icons import com.sourcegraph.cody.agent.CodyAgentService import com.sourcegraph.cody.agent.protocol.AutocompleteItem -import com.sourcegraph.cody.agent.protocol.AutocompleteParams import com.sourcegraph.cody.agent.protocol.AutocompleteResult -import com.sourcegraph.cody.agent.protocol.AutocompleteTriggerKind import com.sourcegraph.cody.agent.protocol.CompletionItemParams -import com.sourcegraph.cody.agent.protocol.ErrorCode -import com.sourcegraph.cody.agent.protocol.ErrorCodeUtils.toErrorCode -import com.sourcegraph.cody.agent.protocol.ProtocolTextDocument.Companion.uriFor -import com.sourcegraph.cody.agent.protocol.RateLimitError.Companion.toRateLimitError -import com.sourcegraph.cody.agent.protocol.SelectedCompletionInfo -import com.sourcegraph.cody.agent.protocol_extensions.Position -import com.sourcegraph.cody.agent.protocol_generated.Position -import com.sourcegraph.cody.agent.protocol_generated.Range +import com.sourcegraph.cody.autocomplete.Utils.triggerAutocompleteAsync import com.sourcegraph.cody.autocomplete.render.AutocompleteRendererType import com.sourcegraph.cody.autocomplete.render.CodyAutocompleteBlockElementRenderer import com.sourcegraph.cody.autocomplete.render.CodyAutocompleteElementRenderer import com.sourcegraph.cody.autocomplete.render.CodyAutocompleteSingleLineRenderer import com.sourcegraph.cody.autocomplete.render.InlayModelUtil.getAllInlaysForEditor import com.sourcegraph.cody.config.CodyAuthenticationManager -import com.sourcegraph.cody.ignore.ActionInIgnoredFileNotification -import com.sourcegraph.cody.ignore.IgnoreOracle -import com.sourcegraph.cody.ignore.IgnorePolicy -import com.sourcegraph.cody.statusbar.CodyStatus -import com.sourcegraph.cody.statusbar.CodyStatusService.Companion.notifyApplication import com.sourcegraph.cody.statusbar.CodyStatusService.Companion.resetApplication import com.sourcegraph.cody.vscode.CancellationToken import com.sourcegraph.cody.vscode.InlineCompletionTriggerKind @@ -57,7 +43,6 @@ import com.sourcegraph.cody.vscode.IntelliJTextDocument import com.sourcegraph.cody.vscode.TextDocument import com.sourcegraph.common.CodyBundle import com.sourcegraph.common.CodyBundle.fmt -import com.sourcegraph.common.UpgradeToCodyProNotification import com.sourcegraph.config.ConfigUtil.isCodyEnabled import com.sourcegraph.config.UserLevelConfig import com.sourcegraph.utils.CodyEditorUtil.getAllOpenEditors @@ -67,13 +52,8 @@ import com.sourcegraph.utils.CodyEditorUtil.isCommandExcluded import com.sourcegraph.utils.CodyEditorUtil.isEditorValidForAutocomplete import com.sourcegraph.utils.CodyEditorUtil.isImplicitAutocompleteEnabledForEditor import com.sourcegraph.utils.CodyFormatter -import java.util.concurrent.CancellationException -import java.util.concurrent.CompletableFuture -import java.util.concurrent.CompletionException -import java.util.concurrent.TimeUnit import java.util.concurrent.atomic.AtomicReference import java.util.stream.Collectors -import org.eclipse.lsp4j.jsonrpc.ResponseErrorException /** Responsible for triggering and clearing inline code completions (the autocomplete feature). */ @Service @@ -160,6 +140,10 @@ class CodyAutocompleteManager { } return } + val isRemoteDev = ClientSessionsManager.getAppSession()?.isRemote ?: false + if (isRemoteDev) { + return + } if (isTriggeredImplicitly && !isImplicitAutocompleteEnabledForEditor(editor)) { return } @@ -197,110 +181,11 @@ class CodyAutocompleteManager { triggerKind, cancellationToken, lookupString, - originalText) - } - - /** Asynchronously triggers auto-complete for the given editor and offset. */ - private fun triggerAutocompleteAsync( - project: Project, - editor: Editor, - offset: Int, - textDocument: TextDocument, - triggerKind: InlineCompletionTriggerKind, - cancellationToken: CancellationToken, - lookupString: String?, - originalText: String - ): CompletableFuture { - val position = textDocument.positionAt(offset) - val lineNumber = editor.document.getLineNumber(offset) - var startPosition = 0 - if (!lookupString.isNullOrEmpty()) { - startPosition = findLastCommonSuffixElementPosition(originalText, lookupString) - } - - val virtualFile = - FileDocumentManager.getInstance().getFile(editor.document) - ?: return CompletableFuture.completedFuture(null) - val params = - if (lookupString.isNullOrEmpty()) - AutocompleteParams( - uriFor(virtualFile), - Position(position.line, position.character), - if (triggerKind == InlineCompletionTriggerKind.INVOKE) - AutocompleteTriggerKind.INVOKE.value - else AutocompleteTriggerKind.AUTOMATIC.value) - else - AutocompleteParams( - uriFor(virtualFile), - Position(position.line, position.character), - AutocompleteTriggerKind.AUTOMATIC.value, - SelectedCompletionInfo( - lookupString, - if (startPosition < 0) Range(position, position) - else Range(Position(lineNumber, startPosition), position))) - notifyApplication(project, CodyStatus.AutocompleteInProgress) - - val resultOuter = CompletableFuture() - CodyAgentService.withAgent(project) { agent -> - if (triggerKind == InlineCompletionTriggerKind.INVOKE && - IgnoreOracle.getInstance(project).policyForUri(virtualFile.url, agent).get() != - IgnorePolicy.USE) { - ActionInIgnoredFileNotification.maybeNotify(project) - resetApplication(project) - resultOuter.cancel(true) - } else { - val completions = agent.server.autocompleteExecute(params) - - // Important: we have to `.cancel()` the original `CompletableFuture` from lsp4j. As soon - // as we use `thenAccept()` we get a new instance of `CompletableFuture` which does - // not - // correctly propagate the cancellation to the agent. - cancellationToken.onCancellationRequested { completions.cancel(true) } - - ApplicationManager.getApplication().executeOnPooledThread { - completions - .handle { result, error -> - if (error != null) { - if (triggerKind == InlineCompletionTriggerKind.INVOKE || - !UpgradeToCodyProNotification.isFirstRLEOnAutomaticAutocompletionsShown) { - handleError(project, error) - } - } else if (result != null && result.items.isNotEmpty()) { - UpgradeToCodyProNotification.isFirstRLEOnAutomaticAutocompletionsShown = false - UpgradeToCodyProNotification.autocompleteRateLimitError.set(null) - processAutocompleteResult(editor, offset, triggerKind, result, cancellationToken) - } - null - } - .exceptionally { error: Throwable? -> - if (!(error is CancellationException || error is CompletionException)) { - logger.warn("failed autocomplete request $params", error) - } - null - } - .completeOnTimeout(null, 3, TimeUnit.SECONDS) - .thenRun { // This is a terminal operation, so we needn't call get(). - resetApplication(project) - resultOuter.complete(null) - } + originalText, + logger) { autocompleteResult -> + processAutocompleteResult( + editor, offset, triggerKind, autocompleteResult, cancellationToken) } - } - } - cancellationToken.onCancellationRequested { resultOuter.cancel(true) } - return resultOuter - } - - private fun handleError(project: Project, error: Throwable?) { - if (error is ResponseErrorException) { - if (error.toErrorCode() == ErrorCode.RateLimitError) { - val rateLimitError = error.toRateLimitError() - UpgradeToCodyProNotification.autocompleteRateLimitError.set(rateLimitError) - UpgradeToCodyProNotification.isFirstRLEOnAutomaticAutocompletionsShown = true - ApplicationManager.getApplication().executeOnPooledThread { - UpgradeToCodyProNotification.notify(error.toRateLimitError(), project) - } - } - } } private fun processAutocompleteResult( @@ -511,21 +396,6 @@ class CodyAutocompleteManager { return fontSize + lineSpacing + extraMargin } - private fun findLastCommonSuffixElementPosition( - stringToFindSuffixIn: String, - suffix: String - ): Int { - var i = 0 - while (i <= suffix.length) { - val partY = suffix.substring(0, suffix.length - i) - if (stringToFindSuffixIn.endsWith(partY)) { - return stringToFindSuffixIn.length - (suffix.length - i) - } - i++ - } - return 0 - } - private fun cancelCurrentJob(project: Project?) { currentJob.get().abort() project?.let { resetApplication(it) } diff --git a/src/main/kotlin/com/sourcegraph/cody/autocomplete/CodyInlineCompletionProvider.kt b/src/main/kotlin/com/sourcegraph/cody/autocomplete/CodyInlineCompletionProvider.kt new file mode 100644 index 0000000000..4209a01410 --- /dev/null +++ b/src/main/kotlin/com/sourcegraph/cody/autocomplete/CodyInlineCompletionProvider.kt @@ -0,0 +1,164 @@ +package com.sourcegraph.cody.autocomplete + +import com.intellij.codeInsight.inline.completion.InlineCompletionElement +import com.intellij.codeInsight.inline.completion.InlineCompletionEvent +import com.intellij.codeInsight.inline.completion.InlineCompletionProvider +import com.intellij.codeInsight.inline.completion.InlineCompletionRequest +import com.intellij.codeInsight.inline.completion.InlineCompletionSuggestion +import com.intellij.codeInsight.inline.completion.elements.InlineCompletionGrayTextElement +import com.intellij.openapi.application.ApplicationInfo +import com.intellij.openapi.application.ApplicationManager +import com.intellij.openapi.application.ReadAction +import com.intellij.openapi.client.ClientSessionsManager +import com.intellij.openapi.command.WriteCommandAction +import com.intellij.openapi.diagnostic.Logger +import com.intellij.openapi.editor.Editor +import com.intellij.openapi.editor.event.DocumentEvent +import com.intellij.openapi.project.Project +import com.intellij.openapi.util.TextRange +import com.intellij.util.concurrency.annotations.RequiresReadLock +import com.sourcegraph.cody.agent.CodyAgentService +import com.sourcegraph.cody.agent.protocol.AutocompleteResult +import com.sourcegraph.cody.agent.protocol.CompletionItemParams +import com.sourcegraph.cody.statusbar.CodyStatusService.Companion.resetApplication +import com.sourcegraph.cody.vscode.CancellationToken +import com.sourcegraph.cody.vscode.InlineCompletionTriggerKind +import com.sourcegraph.cody.vscode.IntelliJTextDocument +import com.sourcegraph.config.ConfigUtil +import com.sourcegraph.utils.CodyEditorUtil.getTextRange +import com.sourcegraph.utils.CodyEditorUtil.isImplicitAutocompleteEnabledForEditor +import com.sourcegraph.utils.CodyFormatter +import java.util.concurrent.CompletableFuture +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.AtomicReference + +@JvmInline value class InlineCompletionProviderID(val id: String) + +class CodyInlineCompletionProvider : InlineCompletionProvider { + private val logger = Logger.getInstance(CodyInlineCompletionProvider::class.java) + private val currentJob = AtomicReference(CancellationToken()) + val id = InlineCompletionProviderID("Cody") + + suspend fun getSuggestion(request: InlineCompletionRequest): InlineCompletionSuggestion { + ApplicationManager.getApplication().assertIsNonDispatchThread() + val editor = request.editor + val project = editor.project ?: return InlineCompletionSuggestion.empty() + if (!isImplicitAutocompleteEnabledForEditor(editor)) { + return InlineCompletionSuggestion.empty() + } + val lookupString: String? = null // todo: can we use this provider for lookups? + + cancelCurrentJob(project) + val cancellationToken = CancellationToken() + currentJob.set(cancellationToken) + + val completions = + fetchCompletions( + project, + editor, + InlineCompletionTriggerKind.AUTOMATIC, + cancellationToken, + lookupString) + .completeOnTimeout(null, 1, TimeUnit.SECONDS) + .get() ?: return InlineCompletionSuggestion.empty() + + return InlineCompletionSuggestion.withFlow { + completions.items + .firstNotNullOfOrNull { + WriteCommandAction.runWriteCommandAction( + editor.project) { + val range = getTextRange(editor.document, it.range) + val originalText = editor.document.getText(range) + val cursorOffsetInOriginalText = request.endOffset - range.startOffset + + val formattedCompletionText: String = + if (System.getProperty("cody.autocomplete.enableFormatting") == "false") { + it.insertText + } else { + CodyFormatter.formatStringBasedOnDocument( + it.insertText, project, editor.document, range, request.endOffset) + } + + // ... + + val originalTextBeforeCursor = + originalText.substring(0, cursorOffsetInOriginalText) + val originalTextAfterCursor = originalText.substring(cursorOffsetInOriginalText) + val completionText = + formattedCompletionText + .removePrefix(originalTextBeforeCursor) + .removeSuffix(originalTextAfterCursor) + if (completionText.trim().isBlank()) { + null + } else { + + CodyAgentService.withAgent(project) { agent -> + agent.server.completionSuggested(CompletionItemParams(it.id)) + } + + InlineCompletionGrayTextElement(completionText) + } + } + } + ?.let { emit(it) } + } + } + + @RequiresReadLock + private fun fetchCompletions( + project: Project, + editor: Editor, + triggerKind: InlineCompletionTriggerKind, + cancellationToken: CancellationToken, + lookupString: String?, + ): CompletableFuture { + val textDocument = IntelliJTextDocument(editor, project) + val offset = ReadAction.compute { editor.caretModel.offset } + val lineNumber = editor.document.getLineNumber(offset) + val caretPositionInLine = offset - editor.document.getLineStartOffset(lineNumber) + val originalText = editor.document.getText(TextRange(offset - caretPositionInLine, offset)) + + val result = CompletableFuture() + Utils.triggerAutocompleteAsync( + project, + editor, + offset, + textDocument, + triggerKind, + cancellationToken, + lookupString, + originalText, + logger) { autocompleteResult -> + result.complete(autocompleteResult) + } + return result + } + + private fun cancelCurrentJob(project: Project?) { + currentJob.get().abort() + project?.let { resetApplication(it) } + } + + fun isEnabled(event: InlineCompletionEvent): Boolean { + return isEnabled() + } + + override suspend fun getProposals( + request: InlineCompletionRequest + ): List { + return emptyList() + } + + override fun isEnabled(event: DocumentEvent): Boolean { + return isEnabled() + } + + private fun isEnabled(): Boolean { + val ideVersion = ApplicationInfo.getInstance().build.baselineVersion + val isRemoteDev = ClientSessionsManager.getAppSession()?.isRemote ?: false + return ideVersion >= 233 && + isRemoteDev && + ConfigUtil.isCodyEnabled() && + ConfigUtil.isCodyAutocompleteEnabled() + } +} diff --git a/src/main/kotlin/com/sourcegraph/cody/autocomplete/Utils.kt b/src/main/kotlin/com/sourcegraph/cody/autocomplete/Utils.kt new file mode 100644 index 0000000000..06a9b4da6e --- /dev/null +++ b/src/main/kotlin/com/sourcegraph/cody/autocomplete/Utils.kt @@ -0,0 +1,156 @@ +package com.sourcegraph.cody.autocomplete + +import com.intellij.openapi.application.ApplicationManager +import com.intellij.openapi.diagnostic.Logger +import com.intellij.openapi.editor.Editor +import com.intellij.openapi.fileEditor.FileDocumentManager +import com.intellij.openapi.project.Project +import com.sourcegraph.cody.agent.CodyAgentService +import com.sourcegraph.cody.agent.protocol.AutocompleteParams +import com.sourcegraph.cody.agent.protocol.AutocompleteResult +import com.sourcegraph.cody.agent.protocol.AutocompleteTriggerKind +import com.sourcegraph.cody.agent.protocol.ErrorCode +import com.sourcegraph.cody.agent.protocol.ErrorCodeUtils.toErrorCode +import com.sourcegraph.cody.agent.protocol.ProtocolTextDocument.Companion.uriFor +import com.sourcegraph.cody.agent.protocol.RateLimitError.Companion.toRateLimitError +import com.sourcegraph.cody.agent.protocol.SelectedCompletionInfo +import com.sourcegraph.cody.agent.protocol_extensions.Position +import com.sourcegraph.cody.agent.protocol_generated.Position +import com.sourcegraph.cody.agent.protocol_generated.Range +import com.sourcegraph.cody.ignore.ActionInIgnoredFileNotification +import com.sourcegraph.cody.ignore.IgnoreOracle +import com.sourcegraph.cody.ignore.IgnorePolicy +import com.sourcegraph.cody.statusbar.CodyStatus +import com.sourcegraph.cody.statusbar.CodyStatusService.Companion.notifyApplication +import com.sourcegraph.cody.statusbar.CodyStatusService.Companion.resetApplication +import com.sourcegraph.cody.vscode.CancellationToken +import com.sourcegraph.cody.vscode.InlineCompletionTriggerKind +import com.sourcegraph.cody.vscode.TextDocument +import com.sourcegraph.common.UpgradeToCodyProNotification +import java.util.concurrent.CancellationException +import java.util.concurrent.CompletableFuture +import java.util.concurrent.CompletionException +import java.util.concurrent.TimeUnit +import org.eclipse.lsp4j.jsonrpc.ResponseErrorException + +object Utils { + fun triggerAutocompleteAsync( + project: Project, + editor: Editor, + offset: Int, + textDocument: TextDocument, + triggerKind: InlineCompletionTriggerKind, + cancellationToken: CancellationToken, + lookupString: String?, + originalText: String, + logger: Logger, + successCallback: (AutocompleteResult) -> Unit, + ): CompletableFuture { + val position = textDocument.positionAt(offset) + val lineNumber = editor.document.getLineNumber(offset) + var startPosition = 0 + if (!lookupString.isNullOrEmpty()) { + startPosition = findLastCommonSuffixElementPosition(originalText, lookupString) + } + + val virtualFile = + FileDocumentManager.getInstance().getFile(editor.document) + ?: return CompletableFuture.completedFuture(null) + val params = + if (lookupString.isNullOrEmpty()) + AutocompleteParams( + uriFor(virtualFile), + Position(position.line, position.character), + if (triggerKind == InlineCompletionTriggerKind.INVOKE) + AutocompleteTriggerKind.INVOKE.value + else AutocompleteTriggerKind.AUTOMATIC.value) + else + AutocompleteParams( + uriFor(virtualFile), + Position(position.line, position.character), + AutocompleteTriggerKind.AUTOMATIC.value, + SelectedCompletionInfo( + lookupString, + if (startPosition < 0) Range(position, position) + else Range(Position(lineNumber, startPosition), position))) + notifyApplication(project, CodyStatus.AutocompleteInProgress) + + val resultOuter = CompletableFuture() + CodyAgentService.withAgent(project) { agent -> + if (triggerKind == InlineCompletionTriggerKind.INVOKE && + IgnoreOracle.getInstance(project).policyForUri(virtualFile.url, agent).get() != + IgnorePolicy.USE) { + ActionInIgnoredFileNotification.maybeNotify(project) + resetApplication(project) + resultOuter.cancel(true) + } else { + val completions = agent.server.autocompleteExecute(params) + + // Important: we have to `.cancel()` the original `CompletableFuture` from lsp4j. As soon + // as we use `thenAccept()` we get a new instance of `CompletableFuture` which does + // not + // correctly propagate the cancellation to the agent. + cancellationToken.onCancellationRequested { completions.cancel(true) } + + ApplicationManager.getApplication().executeOnPooledThread { + completions + .handle { result, error -> + if (error != null) { + if (triggerKind == InlineCompletionTriggerKind.INVOKE || + !UpgradeToCodyProNotification.isFirstRLEOnAutomaticAutocompletionsShown) { + handleError(project, error) + } + } else if (result != null && result.items.isNotEmpty()) { + UpgradeToCodyProNotification.isFirstRLEOnAutomaticAutocompletionsShown = false + UpgradeToCodyProNotification.autocompleteRateLimitError.set(null) + successCallback(result) + resultOuter.complete(result) + } + null + } + .exceptionally { error: Throwable? -> + if (!(error is CancellationException || error is CompletionException)) { + logger.warn("failed autocomplete request $params", error) + } + null + } + .completeOnTimeout(null, 3, TimeUnit.SECONDS) + .thenRun { // This is a terminal operation, so we needn't call get(). + resetApplication(project) + resultOuter.complete(null) + } + } + } + } + cancellationToken.onCancellationRequested { resultOuter.cancel(true) } + return resultOuter + } + + private fun handleError(project: Project, error: Throwable?) { + if (error is ResponseErrorException) { + if (error.toErrorCode() == ErrorCode.RateLimitError) { + val rateLimitError = error.toRateLimitError() + UpgradeToCodyProNotification.autocompleteRateLimitError.set(rateLimitError) + UpgradeToCodyProNotification.isFirstRLEOnAutomaticAutocompletionsShown = true + ApplicationManager.getApplication().executeOnPooledThread { + UpgradeToCodyProNotification.notify(error.toRateLimitError(), project) + } + } + } + } + + private fun findLastCommonSuffixElementPosition( + stringToFindSuffixIn: String, + suffix: String + ): Int { + var i = 0 + while (i <= suffix.length) { + val partY = suffix.substring(0, suffix.length - i) + if (stringToFindSuffixIn.endsWith(partY)) { + return stringToFindSuffixIn.length - (suffix.length - i) + } + i++ + } + return 0 + } +} diff --git a/src/main/resources/META-INF/plugin.xml b/src/main/resources/META-INF/plugin.xml index f1cce0abd3..73f1c36ea1 100644 --- a/src/main/resources/META-INF/plugin.xml +++ b/src/main/resources/META-INF/plugin.xml @@ -96,6 +96,7 @@ +