fix session not started issue

Author: liqul

taskweaver/ces/environment.py

@@ -304,8 +304,6 @@ def execute_code(
     ) -> ExecutionResult:
         exec_id = get_id(prefix="exec") if exec_id is None else exec_id
         session = self._get_session(session_id)
-        if session.kernel_status == "pending":
-            self.start_session(session_id)
 
         session.execution_count += 1
         execution_index = session.execution_count
@@ -438,6 +436,8 @@ def _get_session(
             )
             os.makedirs(new_session.session_dir, exist_ok=True)
             self.session_dict[session_id] = new_session
+        elif session_id not in self.session_dict:
+            raise ValueError(f"Session {session_id} not found.")
 
         return self.session_dict.get(session_id, None)
 

taskweaver/code_interpreter/code_executor.py

@@ -63,6 +63,7 @@ def __init__(
         self.plugin_loaded: bool = False
         self.config = config
         self.tracing = tracing
+        self.session_variables = {}
 
     @tracing_decorator
     def execute_code(self, exec_id: str, code: str) -> ExecutionResult:
@@ -75,6 +76,9 @@ def execute_code(self, exec_id: str, code: str) -> ExecutionResult:
             with get_tracer().start_as_current_span("load_plugin"):
                 self.load_plugin()
                 self.plugin_loaded = True
+        
+        # update session variables
+        self.exec_client.update_session_var(self.session_variables)
 
         with get_tracer().start_as_current_span("run_code"):
             self.tracing.set_span_attribute("code", code)
@@ -104,6 +108,9 @@ def execute_code(self, exec_id: str, code: str) -> ExecutionResult:
         self.tracing.set_span_attribute("result", self.format_code_output(result, with_code=False))
 
         return result
+    
+    def update_session_var(self, session_var_dict: dict) -> None:
+        self.session_variables.update(session_var_dict)
 
     def _save_file(
         self,

taskweaver/code_interpreter/code_interpreter/code_interpreter.py

@@ -122,7 +122,7 @@ def get_intro(self) -> str:
 
     def update_session_variables(self, session_variables: Dict[str, str]):
         self.logger.info(f"Updating session variables: {session_variables}")
-        self.executor.exec_client.update_session_var(session_variables)
+        self.executor.update_session_var(session_variables)
 
     @tracing_decorator
     def reply(

taskweaver/code_interpreter/code_interpreter_plugin_only/code_interpreter_plugin_only.py

@@ -57,7 +57,7 @@ def get_intro(self) -> str:
         return self.intro.format(plugin_description=self.plugin_description)
 
     def update_session_variables(self, session_variables: dict) -> None:
-        self.executor.exec_client.update_session_variables(session_variables)
+        self.executor.update_session_var(session_variables)
 
     @tracing_decorator
     def reply(