feat: :construction: successfully loading past sessions

14 files changed, 108 insertions, 44 deletions

diff --git a/continuedev/src/continuedev/core/autopilot.py b/continuedev/src/continuedev/core/autopilot.py
index 7e7ce5d8..6dd30db1 100644
--- a/continuedev/src/continuedev/core/autopilot.py
+++ b/continuedev/src/continuedev/core/autopilot.py
@@ -15,7 +15,7 @@ from ..plugins.context_providers.highlighted_code import HighlightedCodeContextP
 from ..server.ide_protocol import AbstractIdeProtocolServer
 from ..libs.util.queue import AsyncSubscriptionQueue
 from ..models.main import ContinueBaseModel
-from .main import Context, ContinueCustomException, Policy, History, FullState, Step, HistoryNode
+from .main import Context, ContinueCustomException, Policy, History, FullState, SessionInfo, Step, HistoryNode
 from ..plugins.steps.core.core import DisplayErrorStep, ReversibleStep, ManualEditStep, UserInputStep
 from .sdk import ContinueSDK
 from ..libs.util.traceback_parsers import get_python_traceback, get_javascript_traceback
@@ -53,7 +53,8 @@ class Autopilot(ContinueBaseModel):
     policy: Policy = DefaultPolicy()
     history: History = History.from_empty()
     context: Context = Context()
-    full_state: Union[FullState, None] = None
+    full_state: Optional[FullState] = None
+    session_info: Optional[SessionInfo] = None
     context_manager: ContextManager = ContextManager()
     continue_sdk: ContinueSDK = None
 
@@ -88,7 +89,6 @@ class Autopilot(ContinueBaseModel):
         if full_state is not None:
             self.history = full_state.history
             self.context_manager.context_providers["code"].adding_highlighted_code = full_state.adding_highlighted_code
-            await self.context_manager.set_selected_items(full_state.selected_context_items)
 
         self.started = True
 
@@ -112,6 +112,7 @@ class Autopilot(ContinueBaseModel):
             adding_highlighted_code=self.context_manager.context_providers[
                 "code"].adding_highlighted_code if "code" in self.context_manager.context_providers else False,
             selected_context_items=await self.context_manager.get_selected_items() if self.context_manager is not None else [],
+            session_info=self.session_info
         )
         self.full_state = full_state
         return full_state
@@ -375,6 +376,20 @@ class Autopilot(ContinueBaseModel):
         self._main_user_input_queue.append(user_input)
         await self.update_subscribers()
 
+        # Use the first input to create title for session info, and make the session saveable
+        if self.session_info is None:
+            async def create_title():
+                title = await self.continue_sdk.models.medium.complete(f"Give a short title to describe the current chat session. Do not put quotes around the title. The first message was: \"{user_input}\". The title is: ")
+                self.session_info = SessionInfo(
+                    title=title,
+                    session_id=self.ide.session_id,
+                    date_created=time.strftime(
+                        "%Y-%m-%d %H:%M:%S", time.gmtime())
+                )
+
+            create_async_task(create_title(), on_error=lambda e: self.continue_sdk.run_step(
+                DisplayErrorStep(e=e)))
+
         if len(self._main_user_input_queue) > 1:
             return
 
diff --git a/continuedev/src/continuedev/core/main.py b/continuedev/src/continuedev/core/main.py
index 9a06f2e1..a33d777e 100644
--- a/continuedev/src/continuedev/core/main.py
+++ b/continuedev/src/continuedev/core/main.py
@@ -1,5 +1,5 @@
 import json
-from typing import Coroutine, Dict, List, Literal, Union
+from typing import Coroutine, Dict, List, Literal, Optional, Union
 from pydantic.schema import schema
 
 
@@ -253,6 +253,12 @@ class ContextItem(BaseModel):
     editable: bool = False
 
 
+class SessionInfo(ContinueBaseModel):
+    session_id: str
+    title: str
+    date_created: str
+
+
 class FullState(ContinueBaseModel):
     """A full state of the program, including the history"""
     history: History
@@ -261,12 +267,7 @@ class FullState(ContinueBaseModel):
     slash_commands: List[SlashCommandDescription]
     adding_highlighted_code: bool
     selected_context_items: List[ContextItem]
-
-
-class SessionInfo(ContinueBaseModel):
-    session_id: str
-    title: str
-    date_created: str
+    session_info: Optional[SessionInfo] = None
 
 
 class ContinueSDK:
diff --git a/continuedev/src/continuedev/libs/util/paths.py b/continuedev/src/continuedev/libs/util/paths.py
index 66c921f7..01b594cf 100644
--- a/continuedev/src/continuedev/libs/util/paths.py
+++ b/continuedev/src/continuedev/libs/util/paths.py
@@ -35,6 +35,9 @@ def getSessionFilePath(session_id: str):
 def getSessionsListFilePath():
     path = os.path.join(getSessionsFolderPath(), "sessions.json")
     os.makedirs(os.path.dirname(path), exist_ok=True)
+    if not os.path.exists(path):
+        with open(path, 'w') as f:
+            f.write("[]")
     return path
 
 
diff --git a/continuedev/src/continuedev/libs/util/telemetry.py b/continuedev/src/continuedev/libs/util/telemetry.py
index 60c910bb..0f66ad8d 100644
--- a/continuedev/src/continuedev/libs/util/telemetry.py
+++ b/continuedev/src/continuedev/libs/util/telemetry.py
@@ -23,7 +23,6 @@ class PostHogLogger:
         self.posthog = Posthog(self.api_key, host='https://app.posthog.com')
 
     def setup(self, unique_id: str, allow_anonymous_telemetry: bool):
-        logger.debug(f"Setting unique_id as {unique_id}")
         self.unique_id = unique_id or "NO_UNIQUE_ID"
         self.allow_anonymous_telemetry = allow_anonymous_telemetry or True
 
diff --git a/continuedev/src/continuedev/server/gui.py b/continuedev/src/continuedev/server/gui.py
index b6f7b141..661e1787 100644
--- a/continuedev/src/continuedev/server/gui.py
+++ b/continuedev/src/continuedev/server/gui.py
@@ -99,6 +99,8 @@ class GUIProtocolServer(AbstractGUIProtocolServer):
             self.on_show_logs_at_index(data["index"])
         elif message_type == "select_context_item":
             self.select_context_item(data["id"], data["query"])
+        elif message_type == "load_session":
+            self.load_session(data["session_id"])
 
     def on_main_input(self, input: str):
         # Do something with user input
@@ -154,8 +156,13 @@ class GUIProtocolServer(AbstractGUIProtocolServer):
         create_async_task(
             self.session.autopilot.select_context_item(id, query), self.on_error)
 
-    async def reconnect_at_session(self, session_id: str):
-        await self._send_json("reconnect_at_session", {"session_id": session_id})
+    def load_session(self, session_id: str):
+        async def load_and_tell_to_reconnect():
+            await session_manager.load_session(self.session.session_id, session_id)
+            await self._send_json("reconnect_at_session", {"session_id": session_id})
+
+        create_async_task(
+            load_and_tell_to_reconnect(), self.on_error)
 
 
 @router.websocket("/ws")
diff --git a/continuedev/src/continuedev/server/session_manager.py b/continuedev/src/continuedev/server/session_manager.py
index f876c9a9..062f9527 100644
--- a/continuedev/src/continuedev/server/session_manager.py
+++ b/continuedev/src/continuedev/server/session_manager.py
@@ -78,18 +78,7 @@ class SessionManager:
         try:
             await autopilot.start(full_state=full_state)
         except Exception as e:
-            # Have to manually add to history because autopilot isn't started
-            formatted_err = '\n'.join(traceback.format_exception(e))
-            msg_step = MessageStep(
-                name="Error starting Autopilot", message=formatted_err)
-            msg_step.description = f"```\n{formatted_err}\n```"
-            autopilot.history.add_node(HistoryNode(
-                step=msg_step,
-                observation=None,
-                depth=0,
-                active=False
-            ))
-            logger.warning(f"Error starting Autopilot: {e}")
+            await self.on_error(e)
 
         def on_error(e: Exception) -> Coroutine:
             err_msg = '\n'.join(traceback.format_exception(e))
@@ -101,7 +90,7 @@ class SessionManager:
     async def remove_session(self, session_id: str):
         logger.debug(f"Removing session: {session_id}")
         if session_id in self.sessions:
-            if session_id in self.registered_ides:
+            if session_id in self.registered_ides and self.registered_ides[session_id] is not None:
                 ws_to_close = self.registered_ides[session_id].websocket
                 if ws_to_close is not None and ws_to_close.client_state != WebSocketState.DISCONNECTED:
                     await self.sessions[session_id].autopilot.ide.websocket.close()
@@ -111,6 +100,9 @@ class SessionManager:
     async def persist_session(self, session_id: str):
         """Save the session's FullState as a json file"""
         full_state = await self.sessions[session_id].autopilot.get_full_state()
+        if full_state.session_info is None:
+            return
+
         with open(getSessionFilePath(session_id), "w") as f:
             json.dump(full_state.dict(), f)
 
@@ -118,9 +110,10 @@ class SessionManager:
         with open(getSessionsListFilePath(), "r") as f:
             sessions_list = json.load(f)
 
-        sessions_list.append(SessionInfo(
-            session_info=full_state.session_info
-        ))
+        sessions_list.append(full_state.session_info.dict())
+
+        with open(getSessionsListFilePath(), "w") as f:
+            json.dump(sessions_list, f)
 
     async def load_session(self, old_session_id: str, new_session_id: str):
         """Load the session's FullState from a json file"""
@@ -130,7 +123,7 @@ class SessionManager:
 
         # Delete the old session, but keep the IDE
         ide = self.registered_ides[old_session_id]
-        self.registered_ides[old_session_id] = None
+        del self.registered_ides[old_session_id]
 
         # Start the new session
         await self.new_session(ide, session_id=new_session_id)
diff --git a/extension/react-app/src/hooks/AbstractContinueGUIClientProtocol.ts b/extension/react-app/src/hooks/AbstractContinueGUIClientProtocol.ts
index 168fb156..139c9d05 100644
--- a/extension/react-app/src/hooks/AbstractContinueGUIClientProtocol.ts
+++ b/extension/react-app/src/hooks/AbstractContinueGUIClientProtocol.ts
@@ -31,6 +31,8 @@ abstract class AbstractContinueGUIClientProtocol {
 
   abstract selectContextItem(id: string, query: string): void;
 
+  abstract loadSession(session_id: string): void;
+
   abstract onReconnectAtSession(session_id: string): void;
 }
 
diff --git a/extension/react-app/src/hooks/ContinueGUIClientProtocol.ts b/extension/react-app/src/hooks/ContinueGUIClientProtocol.ts
index 830954c5..6cfbf66a 100644
--- a/extension/react-app/src/hooks/ContinueGUIClientProtocol.ts
+++ b/extension/react-app/src/hooks/ContinueGUIClientProtocol.ts
@@ -14,9 +14,11 @@ class ContinueGUIClientProtocol extends AbstractContinueGUIClientProtocol {
     useVscodeMessagePassing: boolean
   ) {
     if (this.messenger) {
-      // this.messenger.close(); TODO
+      console.log("Closing session: ", this.serverUrlWithSessionId);
+      this.messenger.close();
     }
     this.serverUrlWithSessionId = serverUrlWithSessionId;
+    this.useVscodeMessagePassing = useVscodeMessagePassing;
     this.messenger = useVscodeMessagePassing
       ? new VscodeMessenger(serverUrlWithSessionId)
       : new WebsocketMessenger(serverUrlWithSessionId);
@@ -45,12 +47,13 @@ class ContinueGUIClientProtocol extends AbstractContinueGUIClientProtocol {
     this.connectMessenger(serverUrlWithSessionId, useVscodeMessagePassing);
   }
 
+  loadSession(session_id: string): void {
+    this.messenger?.send("load_session", { session_id });
+  }
+
   onReconnectAtSession(session_id: string): void {
     this.connectMessenger(
-      this.serverUrlWithSessionId.replace(
-        /\/session\/[a-zA-Z0-9-]+/,
-        `/session/${session_id}`
-      ),
+      `${this.serverUrlWithSessionId.split("?")[0]}?session_id=${session_id}`,
       this.useVscodeMessagePassing
     );
   }
diff --git a/extension/react-app/src/hooks/messenger.ts b/extension/react-app/src/hooks/messenger.ts
index ecf646c7..0bfbe00c 100644
--- a/extension/react-app/src/hooks/messenger.ts
+++ b/extension/react-app/src/hooks/messenger.ts
@@ -15,6 +15,8 @@ export abstract class Messenger {
   abstract sendAndReceive(messageType: string, data: any): Promise<any>;
 
   abstract onError(callback: (error: any) => void): void;
+
+  abstract close(): void;
 }
 
 export class WebsocketMessenger extends Messenger {
@@ -105,4 +107,8 @@ export class WebsocketMessenger extends Messenger {
   onError(callback: (error: any) => void): void {
     this.websocket.addEventListener("error", callback);
   }
+
+  close(): void {
+    this.websocket.close();
+  }
 }
diff --git a/extension/react-app/src/hooks/vscodeMessenger.ts b/extension/react-app/src/hooks/vscodeMessenger.ts
index 13f5092b..cf626721 100644
--- a/extension/react-app/src/hooks/vscodeMessenger.ts
+++ b/extension/react-app/src/hooks/vscodeMessenger.ts
@@ -76,4 +76,8 @@ export class VscodeMessenger extends Messenger {
       }
     });
   }
+
+  close(): void {
+    postVscMessage("websocketForwardingClose", { url: this.serverUrl });
+  }
 }
diff --git a/extension/react-app/src/pages/history.tsx b/extension/react-app/src/pages/history.tsx
index 052fe5be..0142836f 100644
--- a/extension/react-app/src/pages/history.tsx
+++ b/extension/react-app/src/pages/history.tsx
@@ -4,6 +4,23 @@ import { GUIClientContext } from "../App";
 import { useSelector } from "react-redux";
 import { RootStore } from "../redux/store";
 import { useNavigate } from "react-router-dom";
+import { secondaryDark, vscBackground } from "../components";
+import styled from "styled-components";
+
+const Tr = styled.tr`
+  &:hover {
+    background-color: ${secondaryDark};
+  }
+`;
+
+const TdDiv = styled.div`
+  cursor: pointer;
+  padding-left: 1rem;
+  padding-right: 1rem;
+  padding-top: 0.5rem;
+  padding-bottom: 0.5rem;
+  border-bottom: 1px solid ${secondaryDark};
+`;
 
 function History() {
   const navigate = useNavigate();
@@ -30,24 +47,23 @@ function History() {
 
   return (
     <div className="w-full">
-      <h1 className="text-2xl font-bold">History</h1>
+      <h1 className="text-2xl font-bold m-4">History</h1>
       <table className="w-full">
         <tbody>
           {sessions.map((session, index) => (
-            <tr key={index}>
+            <Tr key={index}>
               <td>
-                <div
-                  className="cursor-pointer"
+                <TdDiv
                   onClick={() => {
-                    // client?.loadSession(session.id);
+                    client?.loadSession(session.session_id);
                     navigate("/");
                   }}
                 >
                   <div className="text-lg">{session.title}</div>
                   <div className="text-gray-400">{session.date_created}</div>
-                </div>
+                </TdDiv>
               </td>
-            </tr>
+            </Tr>
           ))}
         </tbody>
       </table>
diff --git a/extension/react-app/src/redux/slices/configSlice.ts b/extension/react-app/src/redux/slices/configSlice.ts
index 57c4f860..59c76066 100644
--- a/extension/react-app/src/redux/slices/configSlice.ts
+++ b/extension/react-app/src/redux/slices/configSlice.ts
@@ -50,7 +50,7 @@ export const configSlice = createSlice({
     ) => ({
       ...state,
       dataSwitchOn: action.payload,
-    })
+    }),
   },
 });
 
@@ -60,6 +60,6 @@ export const {
   setWorkspacePath,
   setSessionId,
   setVscMediaUrl,
-  setDataSwitchOn
+  setDataSwitchOn,
 } = configSlice.actions;
 export default configSlice.reducer;
diff --git a/extension/src/debugPanel.ts b/extension/src/debugPanel.ts
index b687c3e4..d133080b 100644
--- a/extension/src/debugPanel.ts
+++ b/extension/src/debugPanel.ts
@@ -221,6 +221,15 @@ export function setupDebugPanel(
         }
         break;
       }
+      case "websocketForwardingClose": {
+        let url = data.url;
+        let connection = websocketConnections[url];
+        if (typeof connection !== "undefined") {
+          connection.close();
+          websocketConnections[url] = undefined;
+        }
+        break;
+      }
       case "websocketForwardingMessage": {
         let url = data.url;
         let connection = websocketConnections[url];
diff --git a/extension/src/util/messenger.ts b/extension/src/util/messenger.ts
index bcc88fe1..152d4a1f 100644
--- a/extension/src/util/messenger.ts
+++ b/extension/src/util/messenger.ts
@@ -18,6 +18,8 @@ export abstract class Messenger {
   abstract onError(callback: () => void): void;
 
   abstract sendAndReceive(messageType: string, data: any): Promise<any>;
+
+  abstract close(): void;
 }
 
 export class WebsocketMessenger extends Messenger {
@@ -160,4 +162,8 @@ export class WebsocketMessenger extends Messenger {
   onError(callback: () => void): void {
     this.websocket.addEventListener("error", callback);
   }
+
+  close(): void {
+    this.websocket.close();
+  }
 }