Skip to content

Commit

Permalink
Add CodyInlineCompletionProvider as an alternative for inlay completi…
Browse files Browse the repository at this point in the history
…ons (#2304)

Fixes https://linear.app/sourcegraph/issue/CODY-3742.

<!-- start git-machete generated -->

# Based on PR #2303

## Full chain of PRs as of 2024-09-16

* PR #2304:
  `mkondratek/feat/completion-provider` ➔ `mkondratek/chore/ui-fixes`
* PR #2303:
  `mkondratek/chore/ui-fixes` ➔ `main`

<!-- end git-machete generated -->



## Test plan

AUTOMATIC
1. Run IDE in Remote Development mode 
2. Open a project, open a file
3. Hit Enter to trigger autocomplete 
Expected: autocompletion suggested

INVOKE
1. `shift + option + \` on **macOS**
Expected: autocompletion suggested
  • Loading branch information
mkondratek authored Oct 4, 2024
1 parent 2446c5c commit e84c6ce
Show file tree
Hide file tree
Showing 9 changed files with 469 additions and 146 deletions.
8 changes: 2 additions & 6 deletions build.gradle.kts
Original file line number Diff line number Diff line change
Expand Up @@ -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()
}
}
Expand Down Expand Up @@ -499,6 +499,7 @@ tasks {

buildPlugin {
dependsOn(project.tasks.getByPath("buildCody"))
composedJar.get().exclude("com/intellij/codeInsight/inline/completion/**")
from(
fileTree(buildCodyDir) {
include("*")
Expand Down Expand Up @@ -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")
Expand Down
Original file line number Diff line number Diff line change
@@ -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
}
Original file line number Diff line number Diff line change
@@ -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<InlineCompletionElement>

class Default(override val suggestionFlow: Flow<InlineCompletionElement>) :
InlineCompletionSuggestion()

companion object {
fun empty(): InlineCompletionSuggestion = Default(emptyFlow())

fun withFlow(
buildSuggestion: suspend FlowCollector<InlineCompletionElement>.() -> Unit
): InlineCompletionSuggestion {
return Default(flow(buildSuggestion))
}
}
}
Original file line number Diff line number Diff line change
@@ -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
Original file line number Diff line number Diff line change
@@ -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
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -27,37 +27,22 @@ 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
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
Expand All @@ -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
Expand Down Expand Up @@ -160,6 +140,10 @@ class CodyAutocompleteManager {
}
return
}
val isRemoteDev = ClientSessionsManager.getAppSession()?.isRemote ?: false
if (isRemoteDev) {
return
}
if (isTriggeredImplicitly && !isImplicitAutocompleteEnabledForEditor(editor)) {
return
}
Expand Down Expand Up @@ -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<Void?> {
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<Void?>()
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<T>` from lsp4j. As soon
// as we use `thenAccept()` we get a new instance of `CompletableFuture<Void>` 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(
Expand Down Expand Up @@ -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) }
Expand Down
Loading

0 comments on commit e84c6ce

Please sign in to comment.