checkpoint! protocol reform and it works now

37 files changed, 864 insertions, 612 deletions

diff --git a/continuedev/src/continuedev/core/agent.py b/continuedev/src/continuedev/core/agent.py
index 329e3d4c..7f7466a2 100644
--- a/continuedev/src/continuedev/core/agent.py
+++ b/continuedev/src/continuedev/core/agent.py
@@ -13,7 +13,6 @@ from .sdk import ContinueSDK
 
 
 class Agent(ContinueBaseModel):
-    llm: LLM
     policy: Policy
     ide: AbstractIdeProtocolServer
     history: History = History.from_empty()
@@ -31,27 +30,24 @@ class Agent(ContinueBaseModel):
     def get_full_state(self) -> FullState:
         return FullState(history=self.history, active=self._active, user_input_queue=self._main_user_input_queue)
 
-    def on_update(self, callback: Callable[["FullState"], None]):
+    def on_update(self, callback: Coroutine["FullState", None, None]):
         """Subscribe to changes to state"""
         self._on_update_callbacks.append(callback)
 
-    def update_subscribers(self):
+    async def update_subscribers(self):
         full_state = self.get_full_state()
         for callback in self._on_update_callbacks:
-            callback(full_state)
-
-    def __get_step_params(self, step: "Step"):
-        return ContinueSDK(agent=self, llm=self.llm.with_system_message(step.system_message))
+            await callback(full_state)
 
     def give_user_input(self, input: str, index: int):
-        self._user_input_queue.post(index, input)
+        self._user_input_queue.post(str(index), input)
 
     async def wait_for_user_input(self) -> str:
         self._active = False
-        self.update_subscribers()
-        user_input = await self._user_input_queue.get(self.history.current_index)
+        await self.update_subscribers()
+        user_input = await self._user_input_queue.get(str(self.history.current_index))
         self._active = True
-        self.update_subscribers()
+        await self.update_subscribers()
         return user_input
 
     _manual_edits_buffer: List[FileEditWithFullContents] = []
@@ -62,9 +58,9 @@ class Agent(ContinueBaseModel):
                 current_step = self.history.get_current().step
                 self.history.step_back()
                 if issubclass(current_step.__class__, ReversibleStep):
-                    await current_step.reverse(self.__get_step_params(current_step))
+                    await current_step.reverse(ContinueSDK(self))
 
-                self.update_subscribers()
+                await self.update_subscribers()
         except Exception as e:
             print(e)
 
@@ -94,17 +90,17 @@ class Agent(ContinueBaseModel):
 
         # Run step
         self._step_depth += 1
-        observation = await step(self.__get_step_params(step))
+        observation = await step(ContinueSDK(self))
         self._step_depth -= 1
 
         # Add observation to history
         self.history.get_current().observation = observation
 
         # Update its description
-        step._set_description(await step.describe(self.llm))
+        step._set_description(await step.describe(ContinueSDK(self)))
 
         # Call all subscribed callbacks
-        self.update_subscribers()
+        await self.update_subscribers()
 
         return observation
 
@@ -138,7 +134,7 @@ class Agent(ContinueBaseModel):
 
         # Doing this so active can make it to the frontend after steps are done. But want better state syncing tools
         for callback in self._on_update_callbacks:
-            callback(None)
+            await callback(None)
 
     async def run_from_observation(self, observation: Observation):
         next_step = self.policy.next(self.history)
@@ -158,7 +154,7 @@ class Agent(ContinueBaseModel):
 
     async def accept_user_input(self, user_input: str):
         self._main_user_input_queue.append(user_input)
-        self.update_subscribers()
+        await self.update_subscribers()
 
         if len(self._main_user_input_queue) > 1:
             return
@@ -167,7 +163,7 @@ class Agent(ContinueBaseModel):
         # Just run the step that takes user input, and
         # then up to the policy to decide how to deal with it.
         self._main_user_input_queue.pop(0)
-        self.update_subscribers()
+        await self.update_subscribers()
         await self.run_from_step(UserInputStep(user_input=user_input))
 
         while len(self._main_user_input_queue) > 0:
diff --git a/continuedev/src/continuedev/core/env.py b/continuedev/src/continuedev/core/env.py
index 6267ed60..edd3297c 100644
--- a/continuedev/src/continuedev/core/env.py
+++ b/continuedev/src/continuedev/core/env.py
@@ -8,6 +8,10 @@ def get_env_var(var_name: str):
 
 
 def save_env_var(var_name: str, var_value: str):
+    if not os.path.exists('.env'):
+        with open('.env', 'w') as f:
+            f.write(f'{var_name}="{var_value}"\n')
+        return
     with open('.env', 'r') as f:
         lines = f.readlines()
     with open('.env', 'w') as f:
diff --git a/continuedev/src/continuedev/core/main.py b/continuedev/src/continuedev/core/main.py
index 51fcd299..6be5139b 100644
--- a/continuedev/src/continuedev/core/main.py
+++ b/continuedev/src/continuedev/core/main.py
@@ -72,7 +72,7 @@ class ContinueSDK:
     pass
 
 
-class SequentialStep:
+class Models:
     pass
 
 
@@ -94,7 +94,7 @@ class Step(ContinueBaseModel):
     class Config:
         copy_on_model_validation = False
 
-    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
+    async def describe(self, models: Models) -> Coroutine[str, None, None]:
         if self._description is not None:
             return self._description
         return "Running step: " + self.name
@@ -135,6 +135,16 @@ class Step(ContinueBaseModel):
         return SequentialStep(steps=steps)
 
 
+class SequentialStep(Step):
+    steps: list[Step]
+    hide: bool = True
+
+    async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
+        for step in self.steps:
+            observation = await sdk.run_step(step)
+        return observation
+
+
 class ValidatorObservation(Observation):
     passed: bool
     observation: Observation
diff --git a/continuedev/src/continuedev/core/policy.py b/continuedev/src/continuedev/core/policy.py
index 07101576..c0ba0f4f 100644
--- a/continuedev/src/continuedev/core/policy.py
+++ b/continuedev/src/continuedev/core/policy.py
@@ -1,20 +1,16 @@
 from typing import List, Tuple, Type
 
-
-from ..models.main import ContinueBaseModel
-
 from ..libs.steps.ty import CreatePipelineStep
 from .main import Step, Validator, History, Policy
 from .observation import Observation, TracebackObservation, UserInputObservation
 from ..libs.steps.main import EditHighlightedCodeStep, SolveTracebackStep, RunCodeStep, FasterEditHighlightedCodeStep, StarCoderEditHighlightedCodeStep
 from ..libs.steps.nate import WritePytestsStep, CreateTableStep
-from ..libs.steps.chroma import AnswerQuestionChroma, EditFileChroma
+# from ..libs.steps.chroma import AnswerQuestionChroma, EditFileChroma
 from ..libs.steps.continue_step import ContinueStepStep
 
 
 class DemoPolicy(Policy):
     ran_code_last: bool = False
-    cmd: str
 
     def next(self, history: History) -> Step:
         observation = history.last_observation()
@@ -26,18 +22,15 @@ class DemoPolicy(Policy):
                 return CreatePipelineStep()
             elif "/table" in observation.user_input:
                 return CreateTableStep(sql_str=" ".join(observation.user_input.split(" ")[1:]))
-            elif "/ask" in observation.user_input:
-                return AnswerQuestionChroma(question=" ".join(observation.user_input.split(" ")[1:]))
-            elif "/edit" in observation.user_input:
-                return EditFileChroma(request=" ".join(observation.user_input.split(" ")[1:]))
+            # elif "/ask" in observation.user_input:
+            #     return AnswerQuestionChroma(question=" ".join(observation.user_input.split(" ")[1:]))
+            # elif "/edit" in observation.user_input:
+            #     return EditFileChroma(request=" ".join(observation.user_input.split(" ")[1:]))
             elif "/step" in observation.user_input:
                 return ContinueStepStep(prompt=" ".join(observation.user_input.split(" ")[1:]))
             return StarCoderEditHighlightedCodeStep(user_input=observation.user_input)
 
         state = history.get_current()
-        if state is None or not self.ran_code_last:
-            self.ran_code_last = True
-            return RunCodeStep(cmd=self.cmd)
 
         if observation is not None and isinstance(observation, TracebackObservation):
             self.ran_code_last = False
diff --git a/continuedev/src/continuedev/core/sdk.py b/continuedev/src/continuedev/core/sdk.py
index 750b335d..6ae0be04 100644
--- a/continuedev/src/continuedev/core/sdk.py
+++ b/continuedev/src/continuedev/core/sdk.py
@@ -4,6 +4,7 @@ from ..models.filesystem_edit import FileSystemEdit, AddFile, DeleteFile, AddDir
 from ..models.filesystem import RangeInFile
 from ..libs.llm import LLM
 from ..libs.llm.hf_inference_api import HuggingFaceInferenceAPI
+from ..libs.llm.openai import OpenAI
 from .observation import Observation
 from ..server.ide_protocol import AbstractIdeProtocolServer
 from .main import History, Step
@@ -29,20 +30,20 @@ class Models:
             'HUGGING_FACE_TOKEN', 'Please enter your Hugging Face token')
         return HuggingFaceInferenceAPI(api_key=api_key)
 
+    async def gpt35(self):
+        api_key = await self.sdk.get_user_secret(
+            'OPENAI_API_KEY', 'Please enter your OpenAI API key')
+        return OpenAI(api_key=api_key, default_model="gpt-3.5-turbo")
+
 
 class ContinueSDK:
     """The SDK provided as parameters to a step"""
-    llm: LLM
     ide: AbstractIdeProtocolServer
     steps: ContinueSDKSteps
     models: Models
     __agent: Agent
 
-    def __init__(self, agent: Agent, llm: Union[LLM, None] = None):
-        if llm is None:
-            self.llm = agent.llm
-        else:
-            self.llm = llm
+    def __init__(self, agent: Agent):
         self.ide = agent.ide
         self.__agent = agent
         self.steps = ContinueSDKSteps(self)
diff --git a/continuedev/src/continuedev/libs/steps/chroma.py b/continuedev/src/continuedev/libs/steps/chroma.py
index f13a2bab..39424c5c 100644
--- a/continuedev/src/continuedev/libs/steps/chroma.py
+++ b/continuedev/src/continuedev/libs/steps/chroma.py
@@ -40,7 +40,7 @@ class AnswerQuestionChroma(Step):
 
             Here is the answer:""")
 
-        answer = sdk.llm.complete(prompt)
+        answer = (await sdk.models.gpt35()).complete(prompt)
         print(answer)
         self._answer = answer
 
diff --git a/continuedev/src/continuedev/libs/steps/core/core.py b/continuedev/src/continuedev/libs/steps/core/core.py
index 0338d635..14b3cb80 100644
--- a/continuedev/src/continuedev/libs/steps/core/core.py
+++ b/continuedev/src/continuedev/libs/steps/core/core.py
@@ -4,27 +4,18 @@ from textwrap import dedent
 from typing import Coroutine, List, Union
 from ...llm.prompt_utils import MarkdownStyleEncoderDecoder
 
-from ...util.traceback_parsers import parse_python_traceback
-
 from ....models.filesystem_edit import EditDiff, FileEditWithFullContents, FileSystemEdit
 from ....models.filesystem import FileSystem, RangeInFile, RangeInFileWithContents
-from ...llm import LLM
 from ....core.observation import Observation, TextObservation, TracebackObservation, UserInputObservation
-from ....core.main import Step
+from ....core.main import Step, SequentialStep
 
 
 class ContinueSDK:
     pass
 
 
-class SequentialStep(Step):
-    steps: list[Step]
-    hide: bool = True
-
-    async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
-        for step in self.steps:
-            observation = await sdk.run_step(step)
-        return observation
+class Models:
+    pass
 
 
 class ReversibleStep(Step):
@@ -52,7 +43,7 @@ def ShellCommandsStep(Step):
     cwd: str | None = None
     name: str = "Run Shell Commands"
 
-    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
+    async def describe(self, models: Models) -> Coroutine[str, None, None]:
         return "\n".join(self.cmds)
 
     async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
@@ -81,13 +72,13 @@ class EditCodeStep(Step):
     _prompt: Union[str, None] = None
     _completion: Union[str, None] = None
 
-    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
+    async def describe(self, models: Models) -> Coroutine[str, None, None]:
         if self._edit_diffs is None:
             return "Editing files: " + ", ".join(map(lambda rif: rif.filepath, self.range_in_files))
         elif len(self._edit_diffs) == 0:
             return "No edits made"
         else:
-            return llm.complete(dedent(f"""{self._prompt}{self._completion}
+            return (await models.gpt35()).complete(dedent(f"""{self._prompt}{self._completion}
 
                 Maximally concise summary of changes in bullet points (can use markdown):
             """))
@@ -102,7 +93,7 @@ class EditCodeStep(Step):
         code_string = enc_dec.encode()
         prompt = self.prompt.format(code=code_string)
 
-        completion = sdk.llm.complete(prompt)
+        completion = (await sdk.models.gpt35()).complete(prompt)
 
         # Temporarily doing this to generate description.
         self._prompt = prompt
@@ -127,7 +118,7 @@ class EditFileStep(Step):
     prompt: str
     hide: bool = True
 
-    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
+    async def describe(self, models: Models) -> Coroutine[str, None, None]:
         return "Editing file: " + self.filepath
 
     async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
@@ -145,7 +136,7 @@ class ManualEditStep(ReversibleStep):
 
     hide: bool = True
 
-    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
+    async def describe(self, models: Models) -> Coroutine[str, None, None]:
         return "Manual edit step"
         # TODO - only handling FileEdit here, but need all other types of FileSystemEdits
         # Also requires the merge_file_edit function
@@ -181,7 +172,7 @@ class UserInputStep(Step):
     name: str = "User Input"
     hide: bool = True
 
-    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
+    async def describe(self, models: Models) -> Coroutine[str, None, None]:
         return self.user_input
 
     async def run(self, sdk: ContinueSDK) -> Coroutine[UserInputObservation, None, None]:
@@ -194,7 +185,7 @@ class WaitForUserInputStep(Step):
 
     _description: Union[str, None] = None
 
-    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
+    async def describe(self, models: Models) -> Coroutine[str, None, None]:
         return self.prompt
 
     async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
@@ -207,7 +198,7 @@ class WaitForUserConfirmationStep(Step):
     prompt: str
     name: str = "Waiting for user confirmation"
 
-    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
+    async def describe(self, models: Models) -> Coroutine[str, None, None]:
         return self.prompt
 
     async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
diff --git a/continuedev/src/continuedev/libs/steps/draft/dlt.py b/continuedev/src/continuedev/libs/steps/draft/dlt.py
index 5ba5692a..460aa0cc 100644
--- a/continuedev/src/continuedev/libs/steps/draft/dlt.py
+++ b/continuedev/src/continuedev/libs/steps/draft/dlt.py
@@ -10,7 +10,7 @@ class SetupPipelineStep(Step):
     api_description: str  # e.g. "I want to load data from the weatherapi.com API"
 
     async def run(self, sdk: ContinueSDK):
-        source_name = sdk.llm.complete(
+        source_name = (await sdk.models.gpt35()).complete(
             f"Write a snake_case name for the data source described by {self.api_description}: ").strip()
         filename = f'{source_name}.py'
 
diff --git a/continuedev/src/continuedev/libs/steps/main.py b/continuedev/src/continuedev/libs/steps/main.py
index c8a85800..70c0d4b8 100644
--- a/continuedev/src/continuedev/libs/steps/main.py
+++ b/continuedev/src/continuedev/libs/steps/main.py
@@ -11,7 +11,7 @@ from ...core.observation import Observation, TextObservation, TracebackObservati
 from ..llm.prompt_utils import MarkdownStyleEncoderDecoder
 from textwrap import dedent
 from ...core.main import Step
-from ...core.sdk import ContinueSDK
+from ...core.sdk import ContinueSDK, Models
 from ...core.observation import Observation
 import subprocess
 from .core.core import EditCodeStep
@@ -20,7 +20,7 @@ from .core.core import EditCodeStep
 class RunCodeStep(Step):
     cmd: str
 
-    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
+    async def describe(self, models: Models) -> Coroutine[str, None, None]:
         return f"Ran command: `{self.cmd}`"
 
     async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
@@ -59,7 +59,7 @@ class RunCommandStep(Step):
     name: str = "Run command"
     _description: str = None
 
-    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
+    async def describe(self, models: Models) -> Coroutine[str, None, None]:
         if self._description is not None:
             return self._description
         return self.cmd
@@ -125,7 +125,7 @@ class FasterEditHighlightedCodeStep(Step):
         Here is the description of changes to make:
 """)
 
-    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
+    async def describe(self, models: Models) -> Coroutine[str, None, None]:
         return "Editing highlighted code"
 
     async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
@@ -154,7 +154,7 @@ class FasterEditHighlightedCodeStep(Step):
         for rif in rif_with_contents:
             rif_dict[rif.filepath] = rif.contents
 
-        completion = sdk.llm.complete(prompt)
+        completion = (await sdk.models.gpt35()).complete(prompt)
 
         # Temporarily doing this to generate description.
         self._prompt = prompt
@@ -215,10 +215,10 @@ class FasterEditHighlightedCodeStep(Step):
 
 class StarCoderEditHighlightedCodeStep(Step):
     user_input: str
-    hide = True
+    hide = False
     _prompt: str = "<commit_before>{code}<commit_msg>{user_request}<commit_after>"
 
-    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
+    async def describe(self, models: Models) -> Coroutine[str, None, None]:
         return "Editing highlighted code"
 
     async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
@@ -271,7 +271,7 @@ This is the user request:
 This is the code after being changed to perfectly satisfy the user request:
     """)
 
-    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
+    async def describe(self, models: Models) -> Coroutine[str, None, None]:
         return "Editing highlighted code"
 
     async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
@@ -293,7 +293,7 @@ This is the code after being changed to perfectly satisfy the user request:
 class FindCodeStep(Step):
     prompt: str
 
-    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
+    async def describe(self, models: Models) -> Coroutine[str, None, None]:
         return "Finding code"
 
     async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
@@ -307,7 +307,7 @@ class UserInputStep(Step):
 class SolveTracebackStep(Step):
     traceback: Traceback
 
-    async def describe(self, llm: LLM) -> Coroutine[str, None, None]:
+    async def describe(self, models: Models) -> Coroutine[str, None, None]:
         return f"```\n{self.traceback.full_traceback}\n```"
 
     async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
diff --git a/continuedev/src/continuedev/libs/steps/migration.py b/continuedev/src/continuedev/libs/steps/migration.py
index f044a60f..7b70422d 100644
--- a/continuedev/src/continuedev/libs/steps/migration.py
+++ b/continuedev/src/continuedev/libs/steps/migration.py
@@ -15,7 +15,7 @@ class MigrationStep(Step):
         recent_edits = await sdk.ide.get_recent_edits(self.edited_file)
         recent_edits_string = "\n\n".join(
             map(lambda x: x.to_string(), recent_edits))
-        description = await sdk.llm.complete(f"{recent_edits_string}\n\nGenerate a short description of the migration made in the above changes:\n")
+        description = await (await sdk.models.gpt35()).complete(f"{recent_edits_string}\n\nGenerate a short description of the migration made in the above changes:\n")
         await sdk.run_step(RunCommandStep(cmd=f"cd libs && poetry run alembic revision --autogenerate -m {description}"))
         migration_file = f"libs/alembic/versions/{?}.py"
         contents = await sdk.ide.readFile(migration_file)
diff --git a/continuedev/src/continuedev/libs/steps/nate.py b/continuedev/src/continuedev/libs/steps/nate.py
index a0e728e5..2f84e9d7 100644
--- a/continuedev/src/continuedev/libs/steps/nate.py
+++ b/continuedev/src/continuedev/libs/steps/nate.py
@@ -45,7 +45,7 @@ Here are additional instructions:
 Here is a complete set of pytest unit tests:
 
         """)
-        # tests = sdk.llm.complete(prompt)
+        # tests = (await sdk.models.gpt35()).complete(prompt)
         tests = '''
 import pytest
 
@@ -169,9 +169,9 @@ export class Order {
   tracking_number: string;
 }'''
         time.sleep(2)
-        # orm_entity = sdk.llm.complete(
+        # orm_entity = (await sdk.models.gpt35()).complete(
         #     f"{self.sql_str}\n\nWrite a TypeORM entity called {entity_name} for this table, importing as necessary:")
-        # sdk.llm.complete("What is the name of the entity?")
+        # (await sdk.models.gpt35()).complete("What is the name of the entity?")
         await sdk.apply_filesystem_edit(AddFile(filepath=f"/Users/natesesti/Desktop/continue/extension/examples/python/MyProject/src/entity/{entity_name}.ts", content=orm_entity))
         await sdk.ide.setFileOpen(f"/Users/natesesti/Desktop/continue/extension/examples/python/MyProject/src/entity/{entity_name}.ts", True)
 
diff --git a/continuedev/src/continuedev/libs/steps/pytest.py b/continuedev/src/continuedev/libs/steps/pytest.py
index b4e6dfd2..2e83ae2d 100644
--- a/continuedev/src/continuedev/libs/steps/pytest.py
+++ b/continuedev/src/continuedev/libs/steps/pytest.py
@@ -33,5 +33,5 @@ class WritePytestsStep(Step):
 
             Here is a complete set of pytest unit tests:
         """)
-        tests = sdk.llm.complete(prompt)
+        tests = (await sdk.models.gpt35()).complete(prompt)
         await sdk.apply_filesystem_edit(AddFile(filepath=path, content=tests))
diff --git a/continuedev/src/continuedev/libs/steps/ty.py b/continuedev/src/continuedev/libs/steps/ty.py
index 5ff03f04..9dde7c86 100644
--- a/continuedev/src/continuedev/libs/steps/ty.py
+++ b/continuedev/src/continuedev/libs/steps/ty.py
@@ -18,7 +18,7 @@ class SetupPipelineStep(Step):
     api_description: str  # e.g. "I want to load data from the weatherapi.com API"
 
     async def run(self, sdk: ContinueSDK):
-        # source_name = sdk.llm.complete(
+        # source_name = (await sdk.models.gpt35()).complete(
         #     f"Write a snake_case name for the data source described by {self.api_description}: ").strip()
         filename = f'/Users/natesesti/Desktop/continue/extension/examples/python/{source_name}.py'
 
diff --git a/continuedev/src/continuedev/server/ide.py b/continuedev/src/continuedev/server/ide.py
index dd1dc463..50296841 100644
--- a/continuedev/src/continuedev/server/ide.py
+++ b/continuedev/src/continuedev/server/ide.py
@@ -1,5 +1,6 @@
 # This is a separate server from server/main.py
 import asyncio
+import json
 import os
 from typing import Any, Dict, List, Type, TypeVar, Union
 import uuid
@@ -90,31 +91,33 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
     def __init__(self, session_manager: SessionManager):
         self.session_manager = session_manager
 
-    async def _send_json(self, data: Any):
-        await self.websocket.send_json(data)
+    async def _send_json(self, message_type: str, data: Any):
+        await self.websocket.send_json({
+            "messageType": message_type,
+            "data": data
+        })
 
     async def _receive_json(self, message_type: str) -> Any:
         return await self.sub_queue.get(message_type)
 
     async def _send_and_receive_json(self, data: Any, resp_model: Type[T], message_type: str) -> T:
-        await self._send_json(data)
+        await self._send_json(message_type, data)
         resp = await self._receive_json(message_type)
         return resp_model.parse_obj(resp)
 
-    async def handle_json(self, data: Any):
-        t = data["messageType"]
-        if t == "openNotebook":
+    async def handle_json(self, message_type: str, data: Any):
+        if message_type == "openNotebook":
             await self.openNotebook()
-        elif t == "setFileOpen":
+        elif message_type == "setFileOpen":
             await self.setFileOpen(data["filepath"], data["open"])
-        elif t == "fileEdits":
+        elif message_type == "fileEdits":
             fileEdits = list(
                 map(lambda d: FileEditWithFullContents.parse_obj(d), data["fileEdits"]))
             self.onFileEdits(fileEdits)
-        elif t in ["highlightedCode", "openFiles", "readFile", "editFile", "workspaceDirectory"]:
-            self.sub_queue.post(t, data)
+        elif message_type in ["highlightedCode", "openFiles", "readFile", "editFile", "workspaceDirectory"]:
+            self.sub_queue.post(message_type, data)
         else:
-            raise ValueError("Unknown message type", t)
+            raise ValueError("Unknown message type", message_type)
 
     # ------------------------------- #
     # Request actions in IDE, doesn't matter which Session
@@ -123,24 +126,21 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
 
     async def setFileOpen(self, filepath: str, open: bool = True):
         # Agent needs access to this.
-        await self.websocket.send_json({
-            "messageType": "setFileOpen",
+        await self._send_json("setFileOpen", {
             "filepath": filepath,
             "open": open
         })
 
     async def openNotebook(self):
         session_id = self.session_manager.new_session(self)
-        await self._send_json({
-            "messageType": "openNotebook",
+        await self._send_json("openNotebook", {
             "sessionId": session_id
         })
 
     async def showSuggestionsAndWait(self, suggestions: List[FileEdit]) -> bool:
         ids = [str(uuid.uuid4()) for _ in suggestions]
         for i in range(len(suggestions)):
-            self._send_json({
-                "messageType": "showSuggestion",
+            self._send_json("showSuggestion", {
                 "suggestion": suggestions[i],
                 "suggestionId": ids[i]
             })
@@ -210,8 +210,7 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
 
     async def saveFile(self, filepath: str):
         """Save a file"""
-        await self._send_json({
-            "messageType": "saveFile",
+        await self._send_json("saveFile", {
             "filepath": filepath
         })
 
@@ -293,10 +292,17 @@ ideProtocolServer = IdeProtocolServer(session_manager)
 async def websocket_endpoint(websocket: WebSocket):
     await websocket.accept()
     print("Accepted websocket connection from, ", websocket.client)
-    await websocket.send_json({"messageType": "connected"})
+    await websocket.send_json({"messageType": "connected", "data": {}})
     ideProtocolServer.websocket = websocket
     while True:
-        data = await websocket.receive_json()
-        await ideProtocolServer.handle_json(data)
+        message = await websocket.receive_text()
+        message = json.loads(message)
+
+        if "messageType" not in message or "data" not in message:
+            continue
+        message_type = message["messageType"]
+        data = message["data"]
+
+        await ideProtocolServer.handle_json(message_type, data)
 
     await websocket.close()
diff --git a/continuedev/src/continuedev/server/main.py b/continuedev/src/continuedev/server/main.py
index 11ad1d8f..e87d5fa9 100644
--- a/continuedev/src/continuedev/server/main.py
+++ b/continuedev/src/continuedev/server/main.py
@@ -32,7 +32,7 @@ args = parser.parse_args()
 
 
 def run_server():
-    uvicorn.run(app, host="0.0.0.0", port=args.port, log_config="logging.ini")
+    uvicorn.run(app, host="0.0.0.0", port=args.port)
 
 
 if __name__ == "__main__":
diff --git a/continuedev/src/continuedev/server/notebook.py b/continuedev/src/continuedev/server/notebook.py
index c5dcea31..edb61a45 100644
--- a/continuedev/src/continuedev/server/notebook.py
+++ b/continuedev/src/continuedev/server/notebook.py
@@ -1,18 +1,12 @@
-from fastapi import FastAPI, Depends, Header, WebSocket, APIRouter
-from typing import Any, Dict, List, Union
-from uuid import uuid4
+import json
+from fastapi import Depends, Header, WebSocket, APIRouter
+from typing import Any, Type, TypeVar, Union
 from pydantic import BaseModel
 from uvicorn.main import Server
 
-from ..models.filesystem_edit import FileEditWithFullContents
-from ..core.policy import DemoPolicy
-from ..core.main import FullState, History, Step
-from ..core.agent import Agent
-from ..libs.steps.nate import ImplementAbstractMethodStep
-from ..core.observation import Observation
-from ..libs.llm.openai import OpenAI
-from .ide_protocol import AbstractIdeProtocolServer
-from ..core.env import get_env_var
+from .session_manager import SessionManager, session_manager, Session
+from .notebook_protocol import AbstractNotebookProtocolServer
+from ..libs.util.queue import AsyncSubscriptionQueue
 import asyncio
 import nest_asyncio
 nest_asyncio.apply()
@@ -36,160 +30,99 @@ class AppStatus:
 Server.handle_exit = AppStatus.handle_exit
 
 
-class Session:
-    session_id: str
-    agent: Agent
-    ws: Union[WebSocket, None]
-
-    def __init__(self, session_id: str, agent: Agent):
-        self.session_id = session_id
-        self.agent = agent
-        self.ws = None
-
-
-class DemoAgent(Agent):
-    first_seen: bool = False
-    cumulative_edit_string = ""
-
-    def handle_manual_edits(self, edits: List[FileEditWithFullContents]):
-        for edit in edits:
-            self.cumulative_edit_string += edit.fileEdit.replacement
-            self._manual_edits_buffer.append(edit)
-            # Note that you're storing a lot of unecessary data here. Can compress into EditDiffs on the spot, and merge.
-            # self._manual_edits_buffer = merge_file_edit(self._manual_edits_buffer, edit)
-            # FOR DEMO PURPOSES
-            if edit.fileEdit.filepath.endswith("filesystem.py") and "List" in self.cumulative_edit_string and ":" in edit.fileEdit.replacement:
-                self.cumulative_edit_string = ""
-                asyncio.create_task(self.run_from_step(
-                    ImplementAbstractMethodStep()))
-
-
-class SessionManager:
-    sessions: Dict[str, Session] = {}
-    _event_loop: Union[asyncio.BaseEventLoop, None] = None
-
-    def get_session(self, session_id: str) -> Session:
-        if session_id not in self.sessions:
-            raise KeyError("Session ID not recognized")
-        return self.sessions[session_id]
-
-    def new_session(self, ide: AbstractIdeProtocolServer) -> str:
-        cmd = "python3 /Users/natesesti/Desktop/continue/extension/examples/python/main.py"
-        agent = DemoAgent(llm=OpenAI(api_key=get_env_var("OPENAI_API_KEY")),
-                          policy=DemoPolicy(cmd=cmd), ide=ide)
-        session_id = str(uuid4())
-        session = Session(session_id=session_id, agent=agent)
-        self.sessions[session_id] = session
+def session(x_continue_session_id: str = Header("anonymous")) -> Session:
+    return session_manager.get_session(x_continue_session_id)
 
-        def on_update(state: FullState):
-            session_manager.send_ws_data(session_id, {
-                "messageType": "state",
-                "state": agent.get_full_state().dict()
-            })
 
-        agent.on_update(on_update)
-        asyncio.create_task(agent.run_policy())
-        return session_id
+def websocket_session(session_id: str) -> Session:
+    return session_manager.get_session(session_id)
 
-    def remove_session(self, session_id: str):
-        del self.sessions[session_id]
 
-    def register_websocket(self, session_id: str, ws: WebSocket):
-        self.sessions[session_id].ws = ws
-        print("Registered websocket for session", session_id)
+T = TypeVar("T", bound=BaseModel)
 
-    def send_ws_data(self, session_id: str, data: Any):
-        if self.sessions[session_id].ws is None:
-            print(f"Session {session_id} has no websocket")
-            return
+# You should probably abstract away the websocket stuff into a separate class
 
-        async def a():
-            await self.sessions[session_id].ws.send_json(data)
 
-        # Run coroutine in background
-        if self._event_loop is None or self._event_loop.is_closed():
-            self._event_loop = asyncio.new_event_loop()
-            self._event_loop.run_until_complete(a())
-            self._event_loop.close()
-        else:
-            self._event_loop.run_until_complete(a())
-            self._event_loop.close()
+class NotebookProtocolServer(AbstractNotebookProtocolServer):
+    websocket: WebSocket
+    session: Session
+    sub_queue: AsyncSubscriptionQueue = AsyncSubscriptionQueue()
 
+    def __init__(self, session: Session):
+        self.session = session
 
-session_manager = SessionManager()
+    async def _send_json(self, data: Any):
+        await self.websocket.send_json(data)
 
+    async def _receive_json(self, message_type: str) -> Any:
+        return await self.sub_queue.get(message_type)
 
-def session(x_continue_session_id: str = Header("anonymous")) -> Session:
-    return session_manager.get_session(x_continue_session_id)
+    async def _send_and_receive_json(self, data: Any, resp_model: Type[T], message_type: str) -> T:
+        await self._send_json(data)
+        resp = await self._receive_json(message_type)
+        return resp_model.parse_obj(resp)
 
+    def handle_json(self, message_type: str, data: Any):
+        try:
+            if message_type == "main_input":
+                self.on_main_input(data["input"])
+            elif message_type == "step_user_input":
+                self.on_step_user_input(data["input"], data["index"])
+            elif message_type == "refinement_input":
+                self.on_refinement_input(data["input"], data["index"])
+            elif message_type == "reverse_to_index":
+                self.on_reverse_to_index(data["index"])
+        except Exception as e:
+            print(e)
 
-def websocket_session(session_id: str) -> Session:
-    return session_manager.get_session(session_id)
+    async def send_state_update(self):
+        state = self.session.agent.get_full_state().dict()
+        await self._send_json({
+            "messageType": "state_update",
+            "state": state
+        })
 
+    def on_main_input(self, input: str):
+        # Do something with user input
+        asyncio.create_task(self.session.agent.accept_user_input(input))
 
-class StartSessionBody(BaseModel):
-    config_file_path: Union[str, None]
+    def on_reverse_to_index(self, index: int):
+        # Reverse the history to the given index
+        asyncio.create_task(self.session.agent.reverse_to_index(index))
 
+    def on_step_user_input(self, input: str, index: int):
+        asyncio.create_task(
+            self.session.agent.give_user_input(input, index))
 
-class StartSessionResp(BaseModel):
-    session_id: str
+    def on_refinement_input(self, input: str, index: int):
+        asyncio.create_task(
+            self.session.agent.accept_refinement_input(input, index))
 
 
 @router.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket, session: Session = Depends(websocket_session)):
     await websocket.accept()
 
+    print("Session started")
     session_manager.register_websocket(session.session_id, websocket)
-    data = await websocket.receive_text()
+    protocol = NotebookProtocolServer(session)
+    protocol.websocket = websocket
+
     # Update any history that may have happened before connection
-    await websocket.send_json({
-        "messageType": "state",
-        "state": session_manager.get_session(session.session_id).agent.get_full_state().dict()
-    })
-    print("Session started", data)
+    await protocol.send_state_update()
+
     while AppStatus.should_exit is False:
-        data = await websocket.receive_json()
-        print("Received data", data)
+        message = await websocket.receive_json()
+        print("Received message", message)
+        if type(message) is str:
+            message = json.loads(message)
 
-        if "messageType" not in data:
+        if "messageType" not in message or "data" not in message:
             continue
-        messageType = data["messageType"]
+        message_type = message["messageType"]
+        data = message["data"]
 
-        try:
-            if messageType == "main_input":
-                # Do something with user input
-                asyncio.create_task(
-                    session.agent.accept_user_input(data["value"]))
-            elif messageType == "step_user_input":
-                asyncio.create_task(
-                    session.agent.give_user_input(data["value"], data["index"]))
-            elif messageType == "refinement_input":
-                asyncio.create_task(
-                    session.agent.accept_refinement_input(data["value"], data["index"]))
-            elif messageType == "reverse":
-                # Reverse the history to the given index
-                asyncio.create_task(
-                    session.agent.reverse_to_index(data["index"]))
-        except Exception as e:
-            print(e)
+        protocol.handle_json(message_type, data)
 
     print("Closing websocket")
     await websocket.close()
-
-
-@router.post("/run")
-def request_run(step: Step, session=Depends(session)):
-    """Tell an agent to take a specific action."""
-    asyncio.create_task(session.agent.run_from_step(step))
-    return "Success"
-
-
-@router.get("/history")
-def get_history(session=Depends(session)) -> History:
-    return session.agent.history
-
-
-@router.post("/observation")
-def post_observation(observation: Observation, session=Depends(session)):
-    asyncio.create_task(session.agent.run_from_observation(observation))
-    return "Success"
diff --git a/continuedev/src/continuedev/server/notebook_protocol.py b/continuedev/src/continuedev/server/notebook_protocol.py
new file mode 100644
index 00000000..c2be82e0
--- /dev/null
+++ b/continuedev/src/continuedev/server/notebook_protocol.py
@@ -0,0 +1,28 @@
+from typing import Any
+from abc import ABC, abstractmethod
+
+
+class AbstractNotebookProtocolServer(ABC):
+    @abstractmethod
+    async def handle_json(self, data: Any):
+        """Handle a json message"""
+
+    @abstractmethod
+    def on_main_input(self, input: str):
+        """Called when the user inputs something"""
+
+    @abstractmethod
+    def on_reverse_to_index(self, index: int):
+        """Called when the user requests reverse to a previous index"""
+
+    @abstractmethod
+    def on_refinement_input(self, input: str, index: int):
+        """Called when the user inputs a refinement"""
+
+    @abstractmethod
+    def on_step_user_input(self, input: str, index: int):
+        """Called when the user inputs a step"""
+
+    @abstractmethod
+    async def send_state_update(self, state: dict):
+        """Send a state update to the client"""
diff --git a/continuedev/src/continuedev/server/session_manager.py b/continuedev/src/continuedev/server/session_manager.py
new file mode 100644
index 00000000..b48c21b7
--- /dev/null
+++ b/continuedev/src/continuedev/server/session_manager.py
@@ -0,0 +1,101 @@
+from fastapi import WebSocket
+from typing import Any, Dict, List, Union
+from uuid import uuid4
+
+from ..models.filesystem_edit import FileEditWithFullContents
+from ..core.policy import DemoPolicy
+from ..core.main import FullState
+from ..core.agent import Agent
+from ..libs.steps.nate import ImplementAbstractMethodStep
+from .ide_protocol import AbstractIdeProtocolServer
+import asyncio
+import nest_asyncio
+nest_asyncio.apply()
+
+
+class Session:
+    session_id: str
+    agent: Agent
+    ws: Union[WebSocket, None]
+
+    def __init__(self, session_id: str, agent: Agent):
+        self.session_id = session_id
+        self.agent = agent
+        self.ws = None
+
+
+class DemoAgent(Agent):
+    first_seen: bool = False
+    cumulative_edit_string = ""
+
+    def handle_manual_edits(self, edits: List[FileEditWithFullContents]):
+        for edit in edits:
+            self.cumulative_edit_string += edit.fileEdit.replacement
+            self._manual_edits_buffer.append(edit)
+            # Note that you're storing a lot of unecessary data here. Can compress into EditDiffs on the spot, and merge.
+            # self._manual_edits_buffer = merge_file_edit(self._manual_edits_buffer, edit)
+            # FOR DEMO PURPOSES
+            if edit.fileEdit.filepath.endswith("filesystem.py") and "List" in self.cumulative_edit_string and ":" in edit.fileEdit.replacement:
+                self.cumulative_edit_string = ""
+                asyncio.create_task(self.run_from_step(
+                    ImplementAbstractMethodStep()))
+
+
+class SessionManager:
+    sessions: Dict[str, Session] = {}
+    _event_loop: Union[asyncio.BaseEventLoop, None] = None
+
+    def get_session(self, session_id: str) -> Session:
+        if session_id not in self.sessions:
+            raise KeyError("Session ID not recognized")
+        return self.sessions[session_id]
+
+    def new_session(self, ide: AbstractIdeProtocolServer) -> str:
+        agent = DemoAgent(policy=DemoPolicy(), ide=ide)
+        session_id = str(uuid4())
+        session = Session(session_id=session_id, agent=agent)
+        self.sessions[session_id] = session
+
+        async def on_update(state: FullState):
+            await session_manager.send_ws_data(session_id, "state_update", {
+                "state": agent.get_full_state().dict()
+            })
+
+        agent.on_update(on_update)
+        asyncio.create_task(agent.run_policy())
+        return session_id
+
+    def remove_session(self, session_id: str):
+        del self.sessions[session_id]
+
+    def register_websocket(self, session_id: str, ws: WebSocket):
+        self.sessions[session_id].ws = ws
+        print("Registered websocket for session", session_id)
+
+    async def send_ws_data(self, session_id: str, message_type: str, data: Any):
+        if self.sessions[session_id].ws is None:
+            print(f"Session {session_id} has no websocket")
+            return
+
+        async def a():
+            await self.sessions[session_id].ws.send_json({
+                "messageType": message_type,
+                "data": data
+            })
+
+        # Run coroutine in background
+        await self.sessions[session_id].ws.send_json({
+            "messageType": message_type,
+            "data": data
+        })
+        return
+        if self._event_loop is None or self._event_loop.is_closed():
+            self._event_loop = asyncio.new_event_loop()
+            self._event_loop.run_until_complete(a())
+            self._event_loop.close()
+        else:
+            self._event_loop.run_until_complete(a())
+            self._event_loop.close()
+
+
+session_manager = SessionManager()
diff --git a/extension/package-lock.json b/extension/package-lock.json
index 20ac24be..04af09d3 100644
--- a/extension/package-lock.json
+++ b/extension/package-lock.json
@@ -28,6 +28,7 @@
         "@types/node": "16.x",
         "@types/node-fetch": "^2.6.2",
         "@types/vscode": "^1.74.0",
+        "@types/ws": "^8.5.4",
         "@typescript-eslint/eslint-plugin": "^5.45.0",
         "@typescript-eslint/parser": "^5.45.0",
         "@vscode/test-electron": "^2.2.0",
@@ -2027,6 +2028,15 @@
       "integrity": "sha512-LyeCIU3jb9d38w0MXFwta9r0Jx23ugujkAxdwLTNCyspdZTKUc43t7ppPbCiPoQ/Ivd/pnDFZrb4hWd45wrsgA==",
       "dev": true
     },
+    "node_modules/@types/ws": {
+      "version": "8.5.4",
+      "resolved": "https://registry.npmjs.org/@types/ws/-/ws-8.5.4.tgz",
+      "integrity": "sha512-zdQDHKUgcX/zBc4GrwsE/7dVdAD8JR4EuiAXiiUhhfyIJXXb2+PrGshFyeXWQPMmmZ2XxgaqclgpIC7eTXc1mg==",
+      "dev": true,
+      "dependencies": {
+        "@types/node": "*"
+      }
+    },
     "node_modules/@typescript-eslint/eslint-plugin": {
       "version": "5.48.2",
       "resolved": "https://registry.npmjs.org/@typescript-eslint/eslint-plugin/-/eslint-plugin-5.48.2.tgz",
@@ -9246,6 +9256,15 @@
       "integrity": "sha512-LyeCIU3jb9d38w0MXFwta9r0Jx23ugujkAxdwLTNCyspdZTKUc43t7ppPbCiPoQ/Ivd/pnDFZrb4hWd45wrsgA==",
       "dev": true
     },
+    "@types/ws": {
+      "version": "8.5.4",
+      "resolved": "https://registry.npmjs.org/@types/ws/-/ws-8.5.4.tgz",
+      "integrity": "sha512-zdQDHKUgcX/zBc4GrwsE/7dVdAD8JR4EuiAXiiUhhfyIJXXb2+PrGshFyeXWQPMmmZ2XxgaqclgpIC7eTXc1mg==",
+      "dev": true,
+      "requires": {
+        "@types/node": "*"
+      }
+    },
     "@typescript-eslint/eslint-plugin": {
       "version": "5.48.2",
       "resolved": "https://registry.npmjs.org/@typescript-eslint/eslint-plugin/-/eslint-plugin-5.48.2.tgz",
diff --git a/extension/package.json b/extension/package.json
index dc0192c3..c96655a9 100644
--- a/extension/package.json
+++ b/extension/package.json
@@ -148,7 +148,7 @@
   },
   "scripts": {
     "vscode:prepublish": "npm run esbuild-base -- --minify",
-    "esbuild-base": "esbuild ./src/extension.ts --bundle --outfile=out/extension.js --external:vscode --format=cjs --platform=node",
+    "esbuild-base": "rm -rf ./out && esbuild ./src/extension.ts --bundle --outfile=out/extension.js --external:vscode --format=cjs --platform=node",
     "esbuild": "rm -rf ./out && npm run esbuild-base -- --sourcemap",
     "esbuild-watch": "npm run esbuild-base -- --sourcemap --watch",
     "test-compile": "tsc -p ./",
@@ -160,9 +160,9 @@
     "pretest": "npm run compile && npm run lint",
     "lint": "eslint src --ext ts",
     "test": "node ./out/test/runTest.js",
-    "package": "cp ./config/prod_config.json ./config/config.json && mkdir -p ./build && vsce package --out ./build && chmod 777 ./build/continue-0.0.2.vsix && cp ./config/dev_config.json ./config/config.json",
+    "package": "cp ./config/prod_config.json ./config/config.json && mkdir -p ./build && vsce package --out ./build && chmod 777 ./build/continue-0.0.5.vsix && cp ./config/dev_config.json ./config/config.json",
     "full-package": "cd ../continuedev && poetry build && cp ./dist/continuedev-0.1.0-py3-none-any.whl ../extension/scripts/continuedev-0.1.0-py3-none-any.whl && cd ../extension && npm run typegen && npm run clientgen && cd react-app && npm run build && cd .. && npm run package",
-    "install-extension": "code --install-extension ./build/continue-0.0.1.vsix",
+    "install-extension": "code --install-extension ./build/continue-0.0.5.vsix",
     "uninstall": "code --uninstall-extension .continue",
     "reinstall": "rm -rf ./build && npm run package && npm run uninstall && npm run install-extension"
   },
@@ -173,6 +173,7 @@
     "@types/node": "16.x",
     "@types/node-fetch": "^2.6.2",
     "@types/vscode": "^1.74.0",
+    "@types/ws": "^8.5.4",
     "@typescript-eslint/eslint-plugin": "^5.45.0",
     "@typescript-eslint/parser": "^5.45.0",
     "@vscode/test-electron": "^2.2.0",
diff --git a/extension/react-app/src/hooks/ContinueNotebookClientProtocol.ts b/extension/react-app/src/hooks/ContinueNotebookClientProtocol.ts
new file mode 100644
index 00000000..75fd7373
--- /dev/null
+++ b/extension/react-app/src/hooks/ContinueNotebookClientProtocol.ts
@@ -0,0 +1,13 @@
+abstract class AbstractContinueNotebookClientProtocol {
+  abstract sendMainInput(input: string): void;
+
+  abstract reverseToIndex(index: number): void;
+
+  abstract sendRefinementInput(input: string, index: number): void;
+
+  abstract sendStepUserInput(input: string, index: number): void;
+
+  abstract onStateUpdate(state: any): void;
+}
+
+export default AbstractContinueNotebookClientProtocol;
diff --git a/extension/react-app/src/hooks/messenger.ts b/extension/react-app/src/hooks/messenger.ts
new file mode 100644
index 00000000..e2a0bab8
--- /dev/null
+++ b/extension/react-app/src/hooks/messenger.ts
@@ -0,0 +1,108 @@
+// console.log("Websocket import");
+// const WebSocket = require("ws");
+
+export abstract class Messenger {
+  abstract send(messageType: string, data: object): void;
+
+  abstract onMessageType(
+    messageType: string,
+    callback: (data: object) => void
+  ): void;
+
+  abstract onMessage(callback: (messageType: string, data: any) => void): void;
+
+  abstract onOpen(callback: () => void): void;
+
+  abstract onClose(callback: () => void): void;
+
+  abstract sendAndReceive(messageType: string, data: any): Promise<any>;
+}
+
+export class WebsocketMessenger extends Messenger {
+  websocket: WebSocket;
+  private onMessageListeners: {
+    [messageType: string]: ((data: object) => void)[];
+  } = {};
+  private onOpenListeners: (() => void)[] = [];
+  private onCloseListeners: (() => void)[] = [];
+  private serverUrl: string;
+
+  _newWebsocket(): WebSocket {
+    // // Dynamic import, because WebSocket is builtin with browser, but not with node. And can't use require in browser.
+    // if (typeof process === "object") {
+    //   console.log("Using node");
+    //   // process is only available in Node
+    //   var WebSocket = require("ws");
+    // }
+
+    const newWebsocket = new WebSocket(this.serverUrl);
+    for (const listener of this.onOpenListeners) {
+      this.onOpen(listener);
+    }
+    for (const listener of this.onCloseListeners) {
+      this.onClose(listener);
+    }
+    for (const messageType in this.onMessageListeners) {
+      for (const listener of this.onMessageListeners[messageType]) {
+        this.onMessageType(messageType, listener);
+      }
+    }
+    return newWebsocket;
+  }
+
+  constructor(serverUrl: string) {
+    super();
+    this.serverUrl = serverUrl;
+    this.websocket = this._newWebsocket();
+  }
+
+  send(messageType: string, data: object) {
+    const payload = JSON.stringify({ messageType, data });
+    if (this.websocket.readyState === this.websocket.OPEN) {
+      this.websocket.send(payload);
+    } else {
+      if (this.websocket.readyState !== this.websocket.CONNECTING) {
+        this.websocket = this._newWebsocket();
+      }
+      this.websocket.addEventListener("open", () => {
+        this.websocket.send(payload);
+      });
+    }
+  }
+
+  sendAndReceive(messageType: string, data: any): Promise<any> {
+    return new Promise((resolve, reject) => {
+      const eventListener = (data: any) => {
+        // THIS ISN"T GETTING CALLED
+        resolve(data);
+        this.websocket.removeEventListener("message", eventListener);
+      };
+      this.onMessageType(messageType, eventListener);
+      this.send(messageType, data);
+    });
+  }
+
+  onMessageType(messageType: string, callback: (data: any) => void): void {
+    this.websocket.addEventListener("message", (event: any) => {
+      const msg = JSON.parse(event.data);
+      if (msg.messageType === messageType) {
+        callback(msg.data);
+      }
+    });
+  }
+
+  onMessage(callback: (messageType: string, data: any) => void): void {
+    this.websocket.addEventListener("message", (event) => {
+      const msg = JSON.parse(event.data);
+      callback(msg.messageType, msg.data);
+    });
+  }
+
+  onOpen(callback: () => void): void {
+    this.websocket.addEventListener("open", callback);
+  }
+
+  onClose(callback: () => void): void {
+    this.websocket.addEventListener("close", callback);
+  }
+}
diff --git a/extension/react-app/src/hooks/useContinueNotebookProtocol.ts b/extension/react-app/src/hooks/useContinueNotebookProtocol.ts
new file mode 100644
index 00000000..d5ffbf09
--- /dev/null
+++ b/extension/react-app/src/hooks/useContinueNotebookProtocol.ts
@@ -0,0 +1,49 @@
+import AbstractContinueNotebookClientProtocol from "./ContinueNotebookClientProtocol";
+// import { Messenger, WebsocketMessenger } from "../../../src/util/messenger";
+import { Messenger, WebsocketMessenger } from "./messenger";
+import { VscodeMessenger } from "./vscodeMessenger";
+
+class ContinueNotebookClientProtocol extends AbstractContinueNotebookClientProtocol {
+  messenger: Messenger;
+  // Server URL must contain the session ID param
+  serverUrlWithSessionId: string;
+
+  constructor(
+    serverUrlWithSessionId: string,
+    useVscodeMessagePassing: boolean = false
+  ) {
+    super();
+    this.serverUrlWithSessionId = serverUrlWithSessionId;
+    if (useVscodeMessagePassing) {
+      this.messenger = new VscodeMessenger(serverUrlWithSessionId);
+    } else {
+      this.messenger = new WebsocketMessenger(serverUrlWithSessionId);
+    }
+  }
+
+  sendMainInput(input: string) {
+    this.messenger.send("main_input", { input });
+  }
+
+  reverseToIndex(index: number) {
+    this.messenger.send("reverse_to_index", { index });
+  }
+
+  sendRefinementInput(input: string, index: number) {
+    this.messenger.send("refinement_input", { input, index });
+  }
+
+  sendStepUserInput(input: string, index: number) {
+    this.messenger.send("step_user_input", { input, index });
+  }
+
+  onStateUpdate(callback: (state: any) => void) {
+    this.messenger.onMessageType("state_update", (data: any) => {
+      if (data.state) {
+        callback(data.state);
+      }
+    });
+  }
+}
+
+export default ContinueNotebookClientProtocol;
diff --git a/extension/react-app/src/hooks/useWebsocket.ts b/extension/react-app/src/hooks/useWebsocket.ts
index 6e8e68fa..b98be577 100644
--- a/extension/react-app/src/hooks/useWebsocket.ts
+++ b/extension/react-app/src/hooks/useWebsocket.ts
@@ -1,158 +1,39 @@
 import React, { useEffect, useState } from "react";
 import { RootStore } from "../redux/store";
 import { useSelector } from "react-redux";
+import ContinueNotebookClientProtocol from "./useContinueNotebookProtocol";
 import { postVscMessage } from "../vscode";
 
-abstract class Messenger {
-  abstract send(data: string): void;
-}
-
-class VscodeMessenger extends Messenger {
-  url: string;
-
-  constructor(
-    url: string,
-    onMessage: (message: { data: any }) => void,
-    onOpen: (messenger: Messenger) => void,
-    onClose: (messenger: Messenger) => void
-  ) {
-    super();
-    this.url = url;
-    window.addEventListener("message", (event: any) => {
-      switch (event.data.type) {
-        case "websocketForwardingMessage":
-          onMessage(event.data);
-          break;
-        case "websocketForwardingOpen":
-          onOpen(this);
-          break;
-        case "websocketForwardingClose":
-          onClose(this);
-          break;
-      }
-    });
-
-    postVscMessage("websocketForwardingOpen", { url: this.url });
-  }
-
-  send(data: string) {
-    postVscMessage("websocketForwardingMessage", {
-      message: data,
-      url: this.url,
-    });
-  }
-}
-
-class WebsocketMessenger extends Messenger {
-  websocket: WebSocket;
-  constructor(
-    websocket: WebSocket,
-    onMessage: (message: { data: any }) => void,
-    onOpen: (messenger: Messenger) => void,
-    onClose: (messenger: Messenger) => void
-  ) {
-    super();
-    this.websocket = websocket;
-
-    websocket.addEventListener("close", () => {
-      onClose(this);
-    });
-
-    websocket.addEventListener("open", () => {
-      onOpen(this);
-    });
-
-    websocket.addEventListener("message", (event) => {
-      onMessage(event.data);
-    });
-  }
-
-  static async connect(
-    url: string,
-    sessionId: string,
-    onMessage: (message: { data: any }) => void,
-    onOpen: (messenger: Messenger) => void,
-    onClose: (messenger: Messenger) => void
-  ): Promise<WebsocketMessenger> {
-    const ws = new WebSocket(url);
-
-    return new Promise((resolve, reject) => {
-      ws.addEventListener("open", () => {
-        resolve(new WebsocketMessenger(ws, onMessage, onOpen, onClose));
-      });
-    });
-  }
-
-  send(data: string) {
-    this.websocket.send(JSON.stringify(data));
-  }
-}
-
-function useContinueWebsocket(
-  serverUrl: string,
-  onMessage: (message: { data: any }) => void,
-  useVscodeMessagePassing: boolean = true
-) {
+function useContinueNotebookProtocol(useVscodeMessagePassing: boolean = false) {
   const sessionId = useSelector((state: RootStore) => state.config.sessionId);
-  const [websocket, setWebsocket] = useState<Messenger | undefined>(undefined);
+  const serverHttpUrl = useSelector((state: RootStore) => state.config.apiUrl);
+  const [client, setClient] = useState<
+    ContinueNotebookClientProtocol | undefined
+  >(undefined);
 
-  async function connect() {
-    while (!sessionId) {
-      await new Promise((resolve) => setTimeout(resolve, 300));
+  useEffect(() => {
+    if (!sessionId || !serverHttpUrl) {
+      if (useVscodeMessagePassing) {
+        postVscMessage("onLoad", {});
+      }
+      setClient(undefined);
+      return;
     }
 
-    console.log("Creating websocket", sessionId);
-    console.log("Using vscode message passing", useVscodeMessagePassing);
-
-    const onClose = (messenger: Messenger) => {
-      console.log("Websocket closed");
-      setWebsocket(undefined);
-    };
-
-    const onOpen = (messenger: Messenger) => {
-      console.log("Websocket opened");
-      messenger.send(JSON.stringify({ sessionId }));
-    };
-
-    const url =
-      serverUrl.replace("http", "ws") +
+    const serverUrlWithSessionId =
+      serverHttpUrl.replace("http", "ws") +
       "/notebook/ws?session_id=" +
       encodeURIComponent(sessionId);
 
-    const messenger: Messenger = useVscodeMessagePassing
-      ? new VscodeMessenger(url, onMessage, onOpen, onClose)
-      : await WebsocketMessenger.connect(
-          url,
-          sessionId,
-          onMessage,
-          onOpen,
-          onClose
-        );
-
-    setWebsocket(messenger);
-
-    return messenger;
-  }
-
-  async function getConnection() {
-    if (!websocket) {
-      return await connect();
-    }
-    return websocket;
-  }
-
-  async function send(message: object) {
-    let ws = await getConnection();
-    ws.send(JSON.stringify(message));
-  }
-
-  useEffect(() => {
-    if (!sessionId) {
-      return;
-    }
-    connect();
-  }, [sessionId]);
-
-  return { send };
+    console.log("Creating websocket", serverUrlWithSessionId);
+    console.log("Using vscode message passing", useVscodeMessagePassing);
+    const newClient = new ContinueNotebookClientProtocol(
+      serverUrlWithSessionId,
+      useVscodeMessagePassing
+    );
+    setClient(newClient);
+  }, [sessionId, serverHttpUrl]);
+
+  return client;
 }
-export default useContinueWebsocket;
+export default useContinueNotebookProtocol;
diff --git a/extension/react-app/src/hooks/vscodeMessenger.ts b/extension/react-app/src/hooks/vscodeMessenger.ts
new file mode 100644
index 00000000..746c4302
--- /dev/null
+++ b/extension/react-app/src/hooks/vscodeMessenger.ts
@@ -0,0 +1,68 @@
+import { postVscMessage } from "../vscode";
+// import { Messenger } from "../../../src/util/messenger";
+import { Messenger } from "./messenger";
+
+export class VscodeMessenger extends Messenger {
+  serverUrl: string;
+
+  constructor(serverUrl: string) {
+    super();
+    this.serverUrl = serverUrl;
+    postVscMessage("websocketForwardingOpen", { url: this.serverUrl });
+  }
+
+  send(messageType: string, data: object) {
+    postVscMessage("websocketForwardingMessage", {
+      message: { messageType, data },
+      url: this.serverUrl,
+    });
+  }
+
+  onMessageType(messageType: string, callback: (data: object) => void): void {
+    window.addEventListener("message", (event: any) => {
+      if (event.data.type === "websocketForwardingMessage") {
+        if (event.data.message.messageType === messageType) {
+          callback(event.data.message.data);
+        }
+      }
+    });
+  }
+
+  onMessage(callback: (messageType: string, data: any) => void): void {
+    window.addEventListener("message", (event: any) => {
+      if (event.data.type === "websocketForwardingMessage") {
+        callback(event.data.message.messageType, event.data.message.data);
+      }
+    });
+  }
+
+  sendAndReceive(messageType: string, data: any): Promise<any> {
+    return new Promise((resolve) => {
+      const handler = (event: any) => {
+        if (event.data.type === "websocketForwardingMessage") {
+          if (event.data.message.messageType === messageType) {
+            window.removeEventListener("message", handler);
+            resolve(event.data.message.data);
+          }
+        }
+      };
+      window.addEventListener("message", handler);
+      this.send(messageType, data);
+    });
+  }
+
+  onOpen(callback: () => void): void {
+    window.addEventListener("message", (event: any) => {
+      if (event.data.type === "websocketForwardingOpen") {
+        callback();
+      }
+    });
+  }
+  onClose(callback: () => void): void {
+    window.addEventListener("message", (event: any) => {
+      if (event.data.type === "websocketForwardingClose") {
+        callback();
+      }
+    });
+  }
+}
diff --git a/extension/react-app/src/tabs/notebook.tsx b/extension/react-app/src/tabs/notebook.tsx
index a9c69c5b..02c9ff31 100644
--- a/extension/react-app/src/tabs/notebook.tsx
+++ b/extension/react-app/src/tabs/notebook.tsx
@@ -14,6 +14,7 @@ import StepContainer from "../components/StepContainer";
 import { useSelector } from "react-redux";
 import { RootStore } from "../redux/store";
 import useContinueWebsocket from "../hooks/useWebsocket";
+import useContinueNotebookProtocol from "../hooks/useWebsocket";
 
 let TopNotebookDiv = styled.div`
   display: grid;
@@ -33,8 +34,6 @@ interface NotebookProps {
 }
 
 function Notebook(props: NotebookProps) {
-  const serverUrl = useSelector((state: RootStore) => state.config.apiUrl);
-
   const [waitingForSteps, setWaitingForSteps] = useState(false);
   const [userInputQueue, setUserInputQueue] = useState<string[]>([]);
   const [history, setHistory] = useState<History | undefined>();
@@ -157,30 +156,17 @@ function Notebook(props: NotebookProps) {
   // } as any
   // );
 
-  const { send: websocketSend } = useContinueWebsocket(serverUrl, (msg) => {
-    let data = JSON.parse(msg.data);
-    if (data.messageType === "state") {
-      setWaitingForSteps(data.state.active);
-      setHistory(data.state.history);
-      setUserInputQueue(data.state.user_input_queue);
-    }
-  });
+  const client = useContinueNotebookProtocol();
 
-  // useEffect(() => {
-  //   (async () => {
-  //     if (sessionId && props.firstObservation) {
-  //       let resp = await fetch(serverUrl + "/observation", {
-  //         method: "POST",
-  //         headers: new Headers({
-  //           "x-continue-session-id": sessionId,
-  //         }),
-  //         body: JSON.stringify({
-  //           observation: props.firstObservation,
-  //         }),
-  //       });
-  //     }
-  //   })();
-  // }, [props.firstObservation]);
+  useEffect(() => {
+    console.log("CLIENT ON STATE UPDATE: ", client, client?.onStateUpdate);
+    client?.onStateUpdate((state) => {
+      console.log("Received state update: ", state);
+      setWaitingForSteps(state.active);
+      setHistory(state.history);
+      setUserInputQueue(state.user_input_queue);
+    });
+  }, [client]);
 
   const mainTextInputRef = useRef<HTMLTextAreaElement>(null);
 
@@ -201,14 +187,12 @@ function Notebook(props: NotebookProps) {
 
   const onMainTextInput = () => {
     if (mainTextInputRef.current) {
-      let value = mainTextInputRef.current.value;
+      if (!client) return;
+      let input = mainTextInputRef.current.value;
       setWaitingForSteps(true);
-      websocketSend({
-        messageType: "main_input",
-        value: value,
-      });
+      client.sendMainInput(input);
       setUserInputQueue((queue) => {
-        return [...queue, value];
+        return [...queue, input];
       });
       mainTextInputRef.current.value = "";
       mainTextInputRef.current.style.height = "";
@@ -216,17 +200,20 @@ function Notebook(props: NotebookProps) {
   };
 
   const onStepUserInput = (input: string, index: number) => {
+    if (!client) return;
     console.log("Sending step user input", input, index);
-    websocketSend({
-      messageType: "step_user_input",
-      value: input,
-      index,
-    });
+    client.sendStepUserInput(input, index);
   };
 
   // const iterations = useSelector(selectIterations);
   return (
     <TopNotebookDiv>
+      {typeof client === "undefined" && (
+        <>
+          <Loader></Loader>
+          <p>Server disconnected</p>
+        </>
+      )}
       {history?.timeline.map((node: HistoryNode, index: number) => {
         return (
           <StepContainer
@@ -237,17 +224,10 @@ function Notebook(props: NotebookProps) {
             inFuture={index > history?.current_index}
             historyNode={node}
             onRefinement={(input: string) => {
-              websocketSend({
-                messageType: "refinement_input",
-                value: input,
-                index,
-              });
+              client?.sendRefinementInput(input, index);
             }}
             onReverse={() => {
-              websocketSend({
-                messageType: "reverse",
-                index,
-              });
+              client?.reverseToIndex(index);
             }}
           />
         );
diff --git a/extension/react-app/src/vscode/index.ts b/extension/react-app/src/vscode/index.ts
index 7e373cd9..0785aa4d 100644
--- a/extension/react-app/src/vscode/index.ts
+++ b/extension/react-app/src/vscode/index.ts
@@ -5,6 +5,7 @@ declare const vscode: any;
 
 export function postVscMessage(type: string, data: any) {
   if (typeof vscode === "undefined") {
+    console.log("Unable to send message: vscode is undefined");
     return;
   }
   vscode.postMessage({
diff --git a/extension/react-app/tsconfig.json b/extension/react-app/tsconfig.json
index 3d0a51a8..940a9359 100644
--- a/extension/react-app/tsconfig.json
+++ b/extension/react-app/tsconfig.json
@@ -16,6 +16,6 @@
     "noEmit": true,
     "jsx": "react-jsx"
   },
-  "include": ["src"],
+  "include": ["src", "../src/util/messenger.ts"],
   "references": [{ "path": "./tsconfig.node.json" }]
 }
diff --git a/extension/scripts/continuedev-0.1.0-py3-none-any.whl b/extension/scripts/continuedev-0.1.0-py3-none-any.whl
index d1483db9..2019c904 100644
--- a/extension/scripts/continuedev-0.1.0-py3-none-any.whl
+++ b/extension/scripts/continuedev-0.1.0-py3-none-any.whl
diff --git a/extension/src/activation/activate.ts b/extension/src/activation/activate.ts
index a0aa560b..712ffe13 100644
--- a/extension/src/activation/activate.ts
+++ b/extension/src/activation/activate.ts
@@ -10,7 +10,7 @@ import { getContinueServerUrl } from "../bridge";
 
 export let extensionContext: vscode.ExtensionContext | undefined = undefined;
 
-export let ideProtocolClient: IdeProtocolClient | undefined = undefined;
+export let ideProtocolClient: IdeProtocolClient;
 
 export function activateExtension(
   context: vscode.ExtensionContext,
@@ -24,7 +24,7 @@ export function activateExtension(
   let serverUrl = getContinueServerUrl();
 
   ideProtocolClient = new IdeProtocolClient(
-    serverUrl.replace("http", "ws") + "/ide/ws",
+    `${serverUrl.replace("http", "ws")}/ide/ws`,
     context
   );
 
diff --git a/extension/src/activation/environmentSetup.ts b/extension/src/activation/environmentSetup.ts
index 93a471ff..ad6ac71b 100644
--- a/extension/src/activation/environmentSetup.ts
+++ b/extension/src/activation/environmentSetup.ts
@@ -28,18 +28,7 @@ async function runCommand(cmd: string): Promise<[string, string | undefined]> {
   return [stdout, stderr];
 }
 
-async function getPythonCmdAssumingInstalled() {
-  const [, stderr] = await runCommand("python3 --version");
-  if (stderr) {
-    return "python";
-  }
-  return "python3";
-}
-
-async function setupPythonEnv() {
-  console.log("Setting up python env for Continue extension...");
-  // First check that python3 is installed
-
+async function getPythonPipCommands() {
   var [stdout, stderr] = await runCommand("python3 --version");
   let pythonCmd = "python3";
   if (stderr) {
@@ -58,28 +47,77 @@ async function setupPythonEnv() {
     }
   }
   let pipCmd = pythonCmd.endsWith("3") ? "pip3" : "pip";
+  return [pythonCmd, pipCmd];
+}
 
+function getActivateUpgradeCommands(pythonCmd: string, pipCmd: string) {
   let activateCmd = ". env/bin/activate";
   let pipUpgradeCmd = `${pipCmd} install --upgrade pip`;
   if (process.platform == "win32") {
     activateCmd = ".\\env\\Scripts\\activate";
     pipUpgradeCmd = `${pythonCmd} -m pip install --upgrade pip`;
   }
+  return [activateCmd, pipUpgradeCmd];
+}
 
-  let command = `cd ${path.join(
+function checkEnvExists() {
+  const envBinActivatePath = path.join(
     getExtensionUri().fsPath,
-    "scripts"
-  )} && ${pythonCmd} -m venv env && ${activateCmd} && ${pipUpgradeCmd} && ${pipCmd} install -r requirements.txt`;
-  var [stdout, stderr] = await runCommand(command);
-  if (stderr) {
-    throw new Error(stderr);
+    "scripts",
+    "env",
+    "bin",
+    "activate"
+  );
+  return fs.existsSync(envBinActivatePath);
+}
+
+async function setupPythonEnv() {
+  console.log("Setting up python env for Continue extension...");
+
+  // Assemble the command to create the env
+  const [pythonCmd, pipCmd] = await getPythonPipCommands();
+  const [activateCmd, pipUpgradeCmd] = getActivateUpgradeCommands(
+    pythonCmd,
+    pipCmd
+  );
+  const createEnvCommand = [
+    `cd ${path.join(getExtensionUri().fsPath, "scripts")}`,
+    `${pythonCmd} -m venv env`,
+  ].join(" && ");
+
+  // Repeat until it is successfully created (sometimes it fails to generate the bin, need to try again)
+  while (true) {
+    const [, stderr] = await runCommand(createEnvCommand);
+    if (stderr) {
+      throw new Error(stderr);
+    }
+    if (checkEnvExists()) {
+      break;
+    } else {
+      // Remove the env and try again
+      const removeCommand = `rm -rf ${path.join(
+        getExtensionUri().fsPath,
+        "scripts",
+        "env"
+      )}`;
+      await runCommand(removeCommand);
+    }
   }
   console.log(
     "Successfully set up python env at ",
     getExtensionUri().fsPath + "/scripts/env"
   );
 
-  await startContinuePythonServer();
+  const installRequirementsCommand = [
+    `cd ${path.join(getExtensionUri().fsPath, "scripts")}`,
+    activateCmd,
+    pipUpgradeCmd,
+    `${pipCmd} install -r requirements.txt`,
+  ].join(" && ");
+  const [, stderr] = await runCommand(installRequirementsCommand);
+  if (stderr) {
+    throw new Error(stderr);
+  }
 }
 
 function readEnvFile(path: string) {
@@ -116,29 +154,19 @@ function writeEnvFile(path: string, key: string, value: string) {
 }
 
 export async function startContinuePythonServer() {
+  await setupPythonEnv();
+
   // Check vscode settings
   let serverUrl = getContinueServerUrl();
   if (serverUrl !== "http://localhost:8000") {
     return;
   }
 
-  let envFile = path.join(getExtensionUri().fsPath, "scripts", ".env");
-  let openai_api_key: string | undefined =
-    readEnvFile(envFile)["OPENAI_API_KEY"];
-  while (typeof openai_api_key === "undefined" || openai_api_key === "") {
-    openai_api_key = await vscode.window.showInputBox({
-      prompt: "Enter your OpenAI API key",
-      placeHolder: "Enter your OpenAI API key",
-    });
-    // Write to .env file
-  }
-  writeEnvFile(envFile, "OPENAI_API_KEY", openai_api_key);
-
   console.log("Starting Continue python server...");
 
   // Check if already running by calling /health
   try {
-    let response = await fetch(serverUrl + "/health");
+    const response = await fetch(serverUrl + "/health");
     if (response.status === 200) {
       console.log("Continue python server already running");
       return;
@@ -152,15 +180,18 @@ export async function startContinuePythonServer() {
     pythonCmd = "python";
   }
 
+  // let command = `cd ${path.join(
+  //   getExtensionUri().fsPath,
+  //   "scripts"
+  // )} && ${activateCmd} && cd env/lib/python3.11/site-packages && ${pythonCmd} -m continuedev.server.main`;
   let command = `cd ${path.join(
     getExtensionUri().fsPath,
     "scripts"
   )} && ${activateCmd} && cd .. && ${pythonCmd} -m scripts.run_continue_server`;
   try {
     // exec(command);
-    let child = spawn(command, {
+    const child = spawn(command, {
       shell: true,
-      detached: true,
     });
     child.stdout.on("data", (data: any) => {
       console.log(`stdout: ${data}`);
@@ -194,11 +225,6 @@ export function isPythonEnvSetup(): boolean {
   return fs.existsSync(path.join(pathToEnvCfg));
 }
 
-export async function setupExtensionEnvironment() {
-  console.log("Setting up environment for Continue extension...");
-  await Promise.all([setupPythonEnv()]);
-}
-
 export async function downloadPython3() {
   // Download python3 and return the command to run it (python or python3)
   let os = process.platform;
diff --git a/extension/src/commands.ts b/extension/src/commands.ts
index 18f08e31..aeeb4b4f 100644
--- a/extension/src/commands.ts
+++ b/extension/src/commands.ts
@@ -62,11 +62,11 @@ const commandsMap: { [command: string]: (...args: any) => any } = {
   "continue.acceptSuggestion": acceptSuggestionCommand,
   "continue.rejectSuggestion": rejectSuggestionCommand,
   "continue.openDebugPanel": () => {
-    ideProtocolClient?.openNotebook();
+    ideProtocolClient.openNotebook();
   },
   "continue.focusContinueInput": async () => {
     if (!debugPanelWebview) {
-      await ideProtocolClient?.openNotebook();
+      await ideProtocolClient.openNotebook();
     }
     debugPanelWebview?.postMessage({
       type: "focusContinueInput",
diff --git a/extension/src/continueIdeClient.ts b/extension/src/continueIdeClient.ts
index 35eb668d..477d1420 100644
--- a/extension/src/continueIdeClient.ts
+++ b/extension/src/continueIdeClient.ts
@@ -10,30 +10,28 @@ import {
 } from "./suggestions";
 import { debugPanelWebview, setupDebugPanel } from "./debugPanel";
 import { FileEditWithFullContents } from "../schema/FileEditWithFullContents";
-const util = require("util");
-const exec = util.promisify(require("child_process").exec);
-const WebSocket = require("ws");
 import fs = require("fs");
+import { WebsocketMessenger } from "./util/messenger";
 
 class IdeProtocolClient {
-  private _ws: WebSocket | null = null;
-  private _panels: Map<string, vscode.WebviewPanel> = new Map();
-  private readonly _serverUrl: string;
-  private readonly _context: vscode.ExtensionContext;
+  private messenger: WebsocketMessenger | null = null;
+  private panels: Map<string, vscode.WebviewPanel> = new Map();
+  private readonly context: vscode.ExtensionContext;
 
   private _makingEdit = 0;
 
   constructor(serverUrl: string, context: vscode.ExtensionContext) {
-    this._context = context;
-    this._serverUrl = serverUrl;
-    let ws = new WebSocket(serverUrl);
-    this._ws = ws;
-    ws.onclose = () => {
-      this._ws = null;
-    };
-    ws.on("message", (data: any) => {
-      this.handleMessage(JSON.parse(data));
+    this.context = context;
+
+    let messenger = new WebsocketMessenger(serverUrl);
+    this.messenger = messenger;
+    messenger.onClose(() => {
+      this.messenger = null;
+    });
+    messenger.onMessage((messageType, data) => {
+      this.handleMessage(messageType, data);
     });
+
     // Setup listeners for any file changes in open editors
     vscode.workspace.onDidChangeTextDocument((event) => {
       if (this._makingEdit === 0) {
@@ -58,125 +56,52 @@ class IdeProtocolClient {
             };
           }
         );
-        this.send("fileEdits", { fileEdits });
+        this.messenger?.send("fileEdits", { fileEdits });
       } else {
         this._makingEdit--;
       }
     });
   }
 
-  async isConnected() {
-    if (this._ws === null || this._ws.readyState !== WebSocket.OPEN) {
-      let ws = new WebSocket(this._serverUrl);
-      ws.onclose = () => {
-        this._ws = null;
-      };
-      ws.on("message", (data: any) => {
-        this.handleMessage(JSON.parse(data));
-      });
-      this._ws = ws;
-
-      return new Promise((resolve, reject) => {
-        ws.addEventListener("open", () => {
-          resolve(null);
-        });
-      });
-    }
-  }
-
-  async startCore() {
-    var { stdout, stderr } = await exec(
-      "cd /Users/natesesti/Desktop/continue/continue && poetry shell"
-    );
-    if (stderr) {
-      throw new Error(stderr);
-    }
-    var { stdout, stderr } = await exec(
-      "cd .. && uvicorn continue.src.server.main:app --reload --reload-dir continue"
-    );
-    if (stderr) {
-      throw new Error(stderr);
-    }
-    var { stdout, stderr } = await exec("python3 -m continue.src.libs.ide");
-    if (stderr) {
-      throw new Error(stderr);
-    }
-  }
-
-  async send(messageType: string, data: object) {
-    await this.isConnected();
-    let msg = JSON.stringify({ messageType, ...data });
-    this._ws!.send(msg);
-    console.log("Sent message", msg);
-  }
-
-  async receiveMessage(messageType: string): Promise<any> {
-    await this.isConnected();
-    console.log("Connected to websocket");
-    return await new Promise((resolve, reject) => {
-      if (!this._ws) {
-        reject("Not connected to websocket");
-      }
-      this._ws!.onmessage = (event: any) => {
-        let message = JSON.parse(event.data);
-        console.log("RECEIVED MESSAGE", message);
-        if (message.messageType === messageType) {
-          resolve(message);
-        }
-      };
-    });
-  }
-
-  async sendAndReceive(message: any, messageType: string): Promise<any> {
-    try {
-      await this.send(messageType, message);
-      let msg = await this.receiveMessage(messageType);
-      console.log("Received message", msg);
-      return msg;
-    } catch (e) {
-      console.log("Error sending message", e);
-    }
-  }
-
-  async handleMessage(message: any) {
-    switch (message.messageType) {
+  async handleMessage(messageType: string, data: any) {
+    switch (messageType) {
       case "highlightedCode":
-        this.send("highlightedCode", {
+        this.messenger?.send("highlightedCode", {
           highlightedCode: this.getHighlightedCode(),
         });
         break;
       case "workspaceDirectory":
-        this.send("workspaceDirectory", {
+        this.messenger?.send("workspaceDirectory", {
           workspaceDirectory: this.getWorkspaceDirectory(),
         });
       case "openFiles":
-        this.send("openFiles", {
+        this.messenger?.send("openFiles", {
           openFiles: this.getOpenFiles(),
         });
         break;
       case "readFile":
-        this.send("readFile", {
-          contents: this.readFile(message.filepath),
+        this.messenger?.send("readFile", {
+          contents: this.readFile(data.filepath),
         });
         break;
       case "editFile":
-        let fileEdit = await this.editFile(message.edit);
-        this.send("editFile", {
+        const fileEdit = await this.editFile(data.edit);
+        this.messenger?.send("editFile", {
           fileEdit,
         });
         break;
       case "saveFile":
-        this.saveFile(message.filepath);
+        this.saveFile(data.filepath);
         break;
       case "setFileOpen":
-        this.openFile(message.filepath);
+        this.openFile(data.filepath);
         // TODO: Close file
         break;
       case "openNotebook":
       case "connected":
         break;
       default:
-        throw Error("Unknown message type:" + message.messageType);
+        throw Error("Unknown message type:" + messageType);
     }
   }
   getWorkspaceDirectory() {
@@ -209,17 +134,20 @@ class IdeProtocolClient {
   // Initiate Request
 
   closeNotebook(sessionId: string) {
-    this._panels.get(sessionId)?.dispose();
-    this._panels.delete(sessionId);
+    this.panels.get(sessionId)?.dispose();
+    this.panels.delete(sessionId);
   }
 
   async openNotebook() {
     console.log("OPENING NOTEBOOK");
-    let resp = await this.sendAndReceive({}, "openNotebook");
-    let sessionId = resp.sessionId;
+    if (this.messenger === null) {
+      console.log("MESSENGER IS NULL");
+    }
+    const resp = await this.messenger?.sendAndReceive("openNotebook", {});
+    const sessionId = resp.sessionId;
     console.log("SESSION ID", sessionId);
 
-    let column = getRightViewColumn();
+    const column = getRightViewColumn();
     const panel = vscode.window.createWebviewPanel(
       "continue.debugPanelView",
       "Continue",
@@ -231,9 +159,9 @@ class IdeProtocolClient {
     );
 
     // And set its HTML content
-    panel.webview.html = setupDebugPanel(panel, this._context, sessionId);
+    panel.webview.html = setupDebugPanel(panel, this.context, sessionId);
 
-    this._panels.set(sessionId, panel);
+    this.panels.set(sessionId, panel);
   }
 
   acceptRejectSuggestion(accept: boolean, key: SuggestionRanges) {
diff --git a/extension/src/debugPanel.ts b/extension/src/debugPanel.ts
index 4192595c..a295085f 100644
--- a/extension/src/debugPanel.ts
+++ b/extension/src/debugPanel.ts
@@ -16,7 +16,6 @@ import {
 import { sendTelemetryEvent, TelemetryEvent } from "./telemetry";
 import { RangeInFile, SerializedDebugContext } from "./client";
 import { addFileSystemToDebugContext } from "./util/util";
-const WebSocket = require("ws");
 
 class StreamManager {
   private _fullText: string = "";
@@ -108,15 +107,15 @@ class WebsocketConnection {
     this._onOpen = onOpen;
     this._onClose = onClose;
 
-    this._ws.onmessage = (event) => {
+    this._ws.addEventListener("message", (event) => {
       this._onMessage(event.data);
-    };
-    this._ws.onclose = () => {
+    });
+    this._ws.addEventListener("close", () => {
       this._onClose();
-    };
-    this._ws.onopen = () => {
+    });
+    this._ws.addEventListener("open", () => {
       this._onOpen();
-    };
+    });
   }
 
   public send(message: string) {
@@ -230,6 +229,19 @@ export function setupDebugPanel(
           apiUrl: getContinueServerUrl(),
           sessionId,
         });
+
+        // // Listen for changes to server URL in settings
+        // vscode.workspace.onDidChangeConfiguration((event) => {
+        //   if (event.affectsConfiguration("continue.serverUrl")) {
+        //     debugPanelWebview?.postMessage({
+        //       type: "onLoad",
+        //       vscMachineId: vscode.env.machineId,
+        //       apiUrl: getContinueServerUrl(),
+        //       sessionId,
+        //     });
+        //   }
+        // });
+
         break;
       }
 
diff --git a/extension/src/extension.ts b/extension/src/extension.ts
index e0b94278..88af0d19 100644
--- a/extension/src/extension.ts
+++ b/extension/src/extension.ts
@@ -4,7 +4,6 @@
 
 import * as vscode from "vscode";
 import {
-  setupExtensionEnvironment,
   isPythonEnvSetup,
   startContinuePythonServer,
 } from "./activation/environmentSetup";
@@ -26,11 +25,7 @@ export function activate(context: vscode.ExtensionContext) {
       cancellable: false,
     },
     async () => {
-      if (isPythonEnvSetup()) {
-        await startContinuePythonServer();
-      } else {
-        await setupExtensionEnvironment();
-      }
+      await startContinuePythonServer();
       dynamicImportAndActivate(context, true);
     }
   );
diff --git a/extension/src/test/runTest.ts b/extension/src/test/runTest.ts
index 27b3ceb2..e810ed5b 100644
--- a/extension/src/test/runTest.ts
+++ b/extension/src/test/runTest.ts
@@ -1,23 +1,23 @@
-import * as path from 'path';
+import * as path from "path";
 
-import { runTests } from '@vscode/test-electron';
+import { runTests } from "@vscode/test-electron";
 
 async function main() {
-	try {
-		// The folder containing the Extension Manifest package.json
-		// Passed to `--extensionDevelopmentPath`
-		const extensionDevelopmentPath = path.resolve(__dirname, '../../');
+  try {
+    // The folder containing the Extension Manifest package.json
+    // Passed to `--extensionDevelopmentPath`
+    const extensionDevelopmentPath = path.resolve(__dirname, "../../");
 
-		// The path to test runner
-		// Passed to --extensionTestsPath
-		const extensionTestsPath = path.resolve(__dirname, './suite/index');
+    // The path to test runner
+    // Passed to --extensionTestsPath
+    const extensionTestsPath = path.resolve(__dirname, "./suite/index");
 
-		// Download VS Code, unzip it and run the integration test
-		await runTests({ extensionDevelopmentPath, extensionTestsPath });
-	} catch (err) {
-		console.error('Failed to run tests');
-		process.exit(1);
-	}
+    // Download VS Code, unzip it and run the integration test
+    await runTests({ extensionDevelopmentPath, extensionTestsPath });
+  } catch (err) {
+    console.error("Failed to run tests");
+    process.exit(1);
+  }
 }
 
 main();
diff --git a/extension/src/util/messenger.ts b/extension/src/util/messenger.ts
new file mode 100644
index 00000000..6f8bb29d
--- /dev/null
+++ b/extension/src/util/messenger.ts
@@ -0,0 +1,108 @@
+console.log("Websocket import");
+const WebSocket = require("ws");
+
+export abstract class Messenger {
+  abstract send(messageType: string, data: object): void;
+
+  abstract onMessageType(
+    messageType: string,
+    callback: (data: object) => void
+  ): void;
+
+  abstract onMessage(callback: (messageType: string, data: any) => void): void;
+
+  abstract onOpen(callback: () => void): void;
+
+  abstract onClose(callback: () => void): void;
+
+  abstract sendAndReceive(messageType: string, data: any): Promise<any>;
+}
+
+export class WebsocketMessenger extends Messenger {
+  websocket: WebSocket;
+  private onMessageListeners: {
+    [messageType: string]: ((data: object) => void)[];
+  } = {};
+  private onOpenListeners: (() => void)[] = [];
+  private onCloseListeners: (() => void)[] = [];
+  private serverUrl: string;
+
+  _newWebsocket(): WebSocket {
+    // // Dynamic import, because WebSocket is builtin with browser, but not with node. And can't use require in browser.
+    // if (typeof process === "object") {
+    //   console.log("Using node");
+    //   // process is only available in Node
+    //   var WebSocket = require("ws");
+    // }
+
+    const newWebsocket = new WebSocket(this.serverUrl);
+    for (const listener of this.onOpenListeners) {
+      this.onOpen(listener);
+    }
+    for (const listener of this.onCloseListeners) {
+      this.onClose(listener);
+    }
+    for (const messageType in this.onMessageListeners) {
+      for (const listener of this.onMessageListeners[messageType]) {
+        this.onMessageType(messageType, listener);
+      }
+    }
+    return newWebsocket;
+  }
+
+  constructor(serverUrl: string) {
+    super();
+    this.serverUrl = serverUrl;
+    this.websocket = this._newWebsocket();
+  }
+
+  send(messageType: string, data: object) {
+    const payload = JSON.stringify({ messageType, data });
+    if (this.websocket.readyState === this.websocket.OPEN) {
+      this.websocket.send(payload);
+    } else {
+      if (this.websocket.readyState !== this.websocket.CONNECTING) {
+        this.websocket = this._newWebsocket();
+      }
+      this.websocket.addEventListener("open", () => {
+        this.websocket.send(payload);
+      });
+    }
+  }
+
+  sendAndReceive(messageType: string, data: any): Promise<any> {
+    return new Promise((resolve, reject) => {
+      const eventListener = (data: any) => {
+        // THIS ISN"T GETTING CALLED
+        resolve(data);
+        this.websocket.removeEventListener("message", eventListener);
+      };
+      this.onMessageType(messageType, eventListener);
+      this.send(messageType, data);
+    });
+  }
+
+  onMessageType(messageType: string, callback: (data: any) => void): void {
+    this.websocket.addEventListener("message", (event: any) => {
+      const msg = JSON.parse(event.data);
+      if (msg.messageType === messageType) {
+        callback(msg.data);
+      }
+    });
+  }
+
+  onMessage(callback: (messageType: string, data: any) => void): void {
+    this.websocket.addEventListener("message", (event) => {
+      const msg = JSON.parse(event.data);
+      callback(msg.messageType, msg.data);
+    });
+  }
+
+  onOpen(callback: () => void): void {
+    this.websocket.addEventListener("open", callback);
+  }
+
+  onClose(callback: () => void): void {
+    this.websocket.addEventListener("close", callback);
+  }
+}