CodeWhisperer CodeScan - Handle file outside project (#3779)

Author: abhikuhikar

jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/sessionconfig/CodeScanSessionConfig.kt

@@ -33,7 +33,8 @@ sealed class CodeScanSessionConfig(
     private val selectedFile: VirtualFile,
     private val project: Project
 ) {
-    val projectRoot = project.guessProjectDir() ?: error("Cannot guess base directory for project ${project.name}")
+    var projectRoot = project.guessProjectDir() ?: error("Cannot guess base directory for project ${project.name}")
+        private set
 
     abstract val sourceExt: String
 
@@ -64,23 +65,38 @@ sealed class CodeScanSessionConfig(
 
         LOG.debug { "Creating payload. File selected as root for the context truncation: ${selectedFile.path}" }
 
-        val (includedSourceFiles, payloadSize, totalLines, _) = includeDependencies()
+        val payloadMetadata = when (selectedFile.path.startsWith(projectRoot.path)) {
+            true -> includeDependencies()
+            false -> {
+                // Set project root as the parent of the selected file.
+                projectRoot = selectedFile.parent
+                includeFileOutsideProjectRoot()
+            }
+        }
 
         // Copy all the included source files to the source zip
-        val srcZip = zipFiles(includedSourceFiles.map { Path.of(it) })
+        val srcZip = zipFiles(payloadMetadata.sourceFiles.map { Path.of(it) })
         val payloadContext = PayloadContext(
             selectedFile.programmingLanguage().toTelemetryType(),
-            totalLines,
-            includedSourceFiles.size,
+            payloadMetadata.linesScanned,
+            payloadMetadata.sourceFiles.size,
             Instant.now().toEpochMilli() - start,
-            includedSourceFiles.mapNotNull { Path.of(it).toFile().toVirtualFile() },
-            payloadSize,
+            payloadMetadata.sourceFiles.mapNotNull { Path.of(it).toFile().toVirtualFile() },
+            payloadMetadata.payloadSize,
             srcZip.length()
         )
 
         return Payload(payloadContext, srcZip)
     }
 
+    open fun includeFileOutsideProjectRoot(): PayloadMetadata =
+        // Handle the case where the selected file is outside the project root.
+        PayloadMetadata(
+            setOf(selectedFile.path),
+            selectedFile.length,
+            Files.lines(selectedFile.toNioPath()).count().toLong()
+        )
+
     open fun includeDependencies(): PayloadMetadata {
         val includedSourceFiles = mutableSetOf<String>()
         var currentTotalFileSize = 0L

jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererPythonCodeScanTest.kt

@@ -149,6 +149,42 @@ class CodeWhispererPythonCodeScanTest : CodeWhispererCodeScanTestBase(PythonCode
         assertThat(filesInZip).isEqualTo(2)
     }
 
+    @Test
+    fun `test createPayload for file outside project`() {
+        val fileOutsideProjectPy = projectRule.fixture.addFileToProject(
+            "../fileOutsideProject.py",
+            """
+                import numpy as np
+                import util
+                a = 1
+                """
+        ).virtualFile
+        val totalSize = fileOutsideProjectPy.length
+        val totalLines = fileOutsideProjectPy.toNioPath().toFile().readLines().size.toLong()
+        sessionConfigSpy = spy(CodeScanSessionConfig.create(fileOutsideProjectPy, project) as PythonCodeScanSessionConfig)
+
+        val payload = sessionConfigSpy.createPayload()
+        assertNotNull(payload)
+        assertThat(payload.context.totalFiles).isEqualTo(1)
+
+        assertThat(payload.context.scannedFiles.size).isEqualTo(1)
+        assertThat(payload.context.scannedFiles).containsExactly(fileOutsideProjectPy)
+
+        assertThat(payload.context.srcPayloadSize).isEqualTo(totalSize)
+        assertThat(payload.context.language).isEqualTo(CodewhispererLanguage.Python)
+        assertThat(payload.context.totalLines).isEqualTo(totalLines)
+        assertNotNull(payload.srcZip)
+
+        val bufferedInputStream = BufferedInputStream(payload.srcZip.inputStream())
+        val zis = ZipInputStream(bufferedInputStream)
+        var filesInZip = 0
+        while (zis.nextEntry != null) {
+            filesInZip += 1
+        }
+
+        assertThat(filesInZip).isEqualTo(1)
+    }
+
     @Test
     fun `e2e happy path integration test`() {
         assertE2ERunsSuccessfully(sessionConfigSpy, project, totalLines, 3, totalSize, 2)