From 8fd4556e2729be877953794bd8db57c121967225 Mon Sep 17 00:00:00 2001
From: Carl-Robert Linnupuu <carlrobertoh@gmail.com>
Date: Wed, 9 Oct 2024 15:16:37 +0300
Subject: [PATCH] feat: keep previous files in context while chatting

---
 .../actions/IncludeFilesInContextAction.java  |  5 ++--
 .../chat/ChatToolWindowTabPanel.java          | 25 +++++++++++++++++--
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/src/main/java/ee/carlrobert/codegpt/actions/IncludeFilesInContextAction.java b/src/main/java/ee/carlrobert/codegpt/actions/IncludeFilesInContextAction.java
index d8bb3bc3..d936bdd4 100644
--- a/src/main/java/ee/carlrobert/codegpt/actions/IncludeFilesInContextAction.java
+++ b/src/main/java/ee/carlrobert/codegpt/actions/IncludeFilesInContextAction.java
@@ -6,7 +6,6 @@ import static ee.carlrobert.codegpt.settings.IncludedFilesSettingsState.DEFAULT_
 import static ee.carlrobert.codegpt.settings.IncludedFilesSettingsState.DEFAULT_REPEATABLE_CONTEXT;
 import static java.lang.String.format;
 
-import com.intellij.icons.AllIcons;
 import com.intellij.openapi.actionSystem.AnAction;
 import com.intellij.openapi.actionSystem.AnActionEvent;
 import com.intellij.openapi.actionSystem.DataContext;
@@ -149,7 +148,9 @@ public class IncludeFilesInContextAction extends AnAction {
 
     private String getVirtualFileContent(VirtualFile virtualFile) {
       try {
-        return new String(Files.readAllBytes(Paths.get(virtualFile.getPath())));
+        if (!virtualFile.isDirectory()) {
+          return new String(Files.readAllBytes(Paths.get(virtualFile.getPath())));
+        }
       } catch (IOException ex) {
         LOG.error(ex);
       }
diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java
index d1817ab9..b4dee22d 100644
--- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java
+++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java
@@ -2,6 +2,7 @@ package ee.carlrobert.codegpt.toolwindow.chat;
 
 import static ee.carlrobert.codegpt.ui.UIUtil.createScrollPaneWithSmartScroller;
 import static java.lang.String.format;
+import static java.util.Collections.emptyList;
 
 import com.intellij.openapi.Disposable;
 import com.intellij.openapi.application.ApplicationManager;
@@ -36,11 +37,15 @@ import ee.carlrobert.codegpt.ui.textarea.UserInputPanel;
 import ee.carlrobert.codegpt.util.EditorUtil;
 import ee.carlrobert.codegpt.util.file.FileUtil;
 import java.awt.BorderLayout;
+import java.io.File;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 import java.util.UUID;
+import java.util.stream.Stream;
 import javax.swing.JComponent;
 import javax.swing.JPanel;
 import kotlin.Unit;
@@ -137,11 +142,27 @@ public class ChatToolWindowTabPanel implements Disposable {
         message.setReferencedFilePaths(referencedFilePaths);
         message.setUserMessage(message.getPrompt());
 
-        totalTokensPanel.updateReferencedFilesTokens(referencedFiles);
-
         chatToolWindowPanel.ifPresent(panel -> panel.clearNotifications(project));
+      } else {
+        referencedFiles = conversation.getMessages().stream()
+            .flatMap(prevMessage -> {
+              if (prevMessage.getReferencedFilePaths() != null) {
+                return prevMessage.getReferencedFilePaths().stream();
+              }
+              return Stream.empty();
+            })
+            .map(filePath -> {
+              try {
+                return new ReferencedFile(new File(filePath));
+              } catch (Exception e) {
+                return null;
+              }
+            })
+            .filter(Objects::nonNull)
+            .toList();
       }
       totalTokensPanel.updateConversationTokens(conversation);
+      totalTokensPanel.updateReferencedFilesTokens(referencedFiles);
 
       var userMessagePanel = new UserMessagePanel(project, message, this);
       var attachedFilePath = CodeGPTKeys.IMAGE_ATTACHMENT_FILE_PATH.get(project);