feat(amazonq): implement inlineCompletionWithReferences

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt

@@ -29,6 +29,8 @@
 import kotlinx.coroutines.isActive
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
+import org.eclipse.lsp4j.Position
+import org.eclipse.lsp4j.TextDocumentIdentifier
 import software.amazon.awssdk.core.exception.SdkServiceException
 import software.amazon.awssdk.core.util.DefaultSdkAutoConstructList
 import software.amazon.awssdk.services.codewhispererruntime.model.CodeWhispererRuntimeException
@@ -53,6 +55,11 @@
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererConnection
 import software.aws.toolkits.jetbrains.services.amazonq.SUPPLEMENTAL_CONTEXT_TIMEOUT
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument.InlineCompletionContext
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument.InlineCompletionTriggerKind
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument.InlineCompletionWithReferencesParams
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.util.FileUriUtil.toUriString
 import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
 import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorManager
@@ -205,6 +212,7 @@
             return
         }
 
+        handleInlineCompletion(editor)
         invokeCodeWhispererInBackground(requestContext)
     }
 
@@ -527,6 +535,15 @@
         return nextStates
     }
 
+    internal fun handleInlineCompletion(editor: Editor) {
+        editor.project?.let { project ->
+            AmazonQLspService.executeIfRunning(project) { server ->
+                val params = buildInlineCompletionParams(editor)
+                server.inlineCompletionWithReferences(params)
+            }
+        }
+    }
+
     private fun initStates(
         requestContext: RequestContext,
         responseContext: ResponseContext,
@@ -706,6 +723,19 @@
         return states
     }
 
+    private fun buildInlineCompletionParams(editor: Editor): InlineCompletionWithReferencesParams =
+        InlineCompletionWithReferencesParams(
+            context = InlineCompletionContext(
+                triggerKind = InlineCompletionTriggerKind.Invoke
+            )
+        ).apply {
+            textDocument = TextDocumentIdentifier(toUriString(editor.virtualFile))
+            position = Position(
+                editor.caretModel.primaryCaret.visualPosition.line,
+                editor.caretModel.primaryCaret.offset
+            )
+        }
+
     private fun addPopupChildDisposables(popup: JBPopup) {
         codeInsightSettingsFacade.disableCodeInsightUntil(popup)
 

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageServer.kt

@@ -9,13 +9,18 @@ import org.eclipse.lsp4j.jsonrpc.services.JsonRequest
 import org.eclipse.lsp4j.services.LanguageServer
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.UpdateCredentialsPayload
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.dependencies.DidChangeDependencyPathsParams
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument.InlineCompletionListWithReferences
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument.InlineCompletionWithReferencesParams
 import java.util.concurrent.CompletableFuture
 
 /**
  * Remote interface exposed by the Amazon Q language server
  */
 @Suppress("unused")
 interface AmazonQLanguageServer : LanguageServer {
+    @JsonRequest("aws/textDocument/inlineCompletionWithReferences")
+    fun inlineCompletionWithReferences(params: InlineCompletionWithReferencesParams): CompletableFuture<InlineCompletionListWithReferences>
+
     @JsonNotification("aws/didChangeDependencyPaths")
     fun didChangeDependencyPaths(params: DidChangeDependencyPathsParams): CompletableFuture<Unit>
 

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/textDocument/InlineCompletionContext.kt

@@ -0,0 +1,9 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument
+
+data class InlineCompletionContext(
+    var triggerKind: InlineCompletionTriggerKind,
+    var selectedCompletionInfo: SelectedCompletionInfo? = null,
+)

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/textDocument/InlineCompletionItem.kt

@@ -0,0 +1,10 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument
+
+data class InlineCompletionItem(
+    var itemId: String,
+    var insertText: String,
+    var references: List<InlineCompletionReference>,
+)

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/textDocument/InlineCompletionListWithReferences.kt

@@ -0,0 +1,9 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument
+
+data class InlineCompletionListWithReferences(
+    var items: List<InlineCompletionItem>,
+    var sessionId: String,
+)

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/textDocument/InlineCompletionReference.kt

@@ -0,0 +1,11 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument
+
+data class InlineCompletionReference(
+    var referenceName: String,
+    var referenceUrl: String,
+    var licenseName: String,
+    var position: InlineCompletionReferencePosition,
+)

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/textDocument/InlineCompletionReferencePosition.kt

@@ -0,0 +1,9 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument
+
+data class InlineCompletionReferencePosition(
+    var startCharacter: Int = 0,
+    var endCharacter: Int = 0,
+)

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/textDocument/InlineCompletionTriggerKind.kt

@@ -0,0 +1,9 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument
+
+enum class InlineCompletionTriggerKind(val value: Int) {
+    Invoke(0),
+    Automatic(1),
+}

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/textDocument/InlineCompletionWithReferencesParams.kt

@@ -0,0 +1,10 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument
+
+import org.eclipse.lsp4j.TextDocumentPositionAndWorkDoneProgressParams
+
+data class InlineCompletionWithReferencesParams(
+    var context: InlineCompletionContext,
+) : TextDocumentPositionAndWorkDoneProgressParams()

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/textDocument/SelectedCompletionInfo.kt

@@ -0,0 +1,11 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument
+
+import org.eclipse.lsp4j.Range
+
+data class SelectedCompletionInfo(
+    var text: String,
+    var range: Range,
+)

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/textdocument/TextDocumentServiceHandler.kt

@@ -52,6 +52,27 @@
             FileDocumentManagerListener.TOPIC,
             this
         )
+
+        // open files on startup
+        val fileEditorManager = FileEditorManager.getInstance(project)
+        fileEditorManager.openFiles.forEach { file ->
+            handleFileOpened(file)
+        }
+    }
+
+    private fun handleFileOpened(file: VirtualFile) {
+        AmazonQLspService.executeIfRunning(project) { languageServer ->
+            toUriString(file)?.let { uri ->
+                languageServer.textDocumentService.didOpen(
+                    DidOpenTextDocumentParams().apply {
+                        textDocument = TextDocumentItem().apply {
+                            this.uri = uri
+                            text = file.inputStream.readAllBytes().decodeToString()
+                        }
+                    }
+                )
+            }
+        }
     }
 
     override fun beforeDocumentSaving(document: Document) {
@@ -99,18 +120,7 @@
         source: FileEditorManager,
         file: VirtualFile,
     ) {
-        AmazonQLspService.executeIfRunning(project) { languageServer ->
-            toUriString(file)?.let { uri ->
-                languageServer.textDocumentService.didOpen(
-                    DidOpenTextDocumentParams().apply {
-                        textDocument = TextDocumentItem().apply {
-                            this.uri = uri
-                            text = file.inputStream.readAllBytes().decodeToString()
-                        }
-                    }
-                )
-            }
-        }
+        handleFileOpened(file)
     }
 
     override fun fileClosed(

plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/textdocument/TextDocumentServiceHandlerTest.kt

@@ -9,6 +9,7 @@ import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.components.serviceIfCreated
 import com.intellij.openapi.editor.Document
 import com.intellij.openapi.fileEditor.FileDocumentManager
+import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.newvfs.events.VFileContentChangeEvent
@@ -43,6 +44,7 @@ import java.util.concurrent.CompletableFuture
 
 class TextDocumentServiceHandlerTest {
     private lateinit var project: Project
+    private lateinit var mockFileEditorManager: FileEditorManager
     private lateinit var mockLanguageServer: AmazonQLanguageServer
     private lateinit var mockTextDocumentService: TextDocumentService
     private lateinit var sut: TextDocumentServiceHandler
@@ -90,6 +92,11 @@ class TextDocumentServiceHandlerTest {
         every { messageBus.connect(any<Disposable>()) } returns mockConnection
         every { mockConnection.subscribe(any(), any()) } just runs
 
+        // Mock FileEditorManager
+        mockFileEditorManager = mockk<FileEditorManager>()
+        every { mockFileEditorManager.openFiles } returns emptyArray()
+        every { project.getService(FileEditorManager::class.java) } returns mockFileEditorManager
+
         sut = TextDocumentServiceHandler(project, mockk())
     }
 
@@ -126,6 +133,25 @@ class TextDocumentServiceHandlerTest {
         }
     }
 
+    @Test
+    fun `didOpen runs on service init`() = runTest {
+        val uri = URI.create("file:///test/path/file.txt")
+        val content = "test content"
+        val file = createMockVirtualFile(uri, content)
+
+        every { mockFileEditorManager.openFiles } returns arrayOf(file)
+
+        sut = TextDocumentServiceHandler(project, mockk())
+
+        val paramsSlot = slot<DidOpenTextDocumentParams>()
+        verify { mockTextDocumentService.didOpen(capture(paramsSlot)) }
+
+        with(paramsSlot.captured.textDocument) {
+            assertEquals(normalizeFileUri(uri.toString()), this.uri)
+            assertEquals(content, text)
+        }
+    }
+
     @Test
     fun `didOpen runs on fileOpened`() = runTest {
         // Create test file