messaging.py

import datetime
import json
import logging
import uuid
import asyncio

from asyncio import Queue
from typing import (
    Dict,
    Optional,
    Union,
)
from pydantic import StrictStr
from websockets.client import WebSocketClientProtocol, connect

from api.models.error import Error
from api.models.logs import Stdout, Stderr
from api.models.result import Result
from api.models.output import (
    EndOfExecution,
    NumberOfExecutions,
    OutputType,
    UnexpectedEndOfExecution,
)
from errors import ExecutionError
from envs import get_envs

logger = logging.getLogger(__name__)


class Execution:
    def __init__(self, in_background: bool = False):
        self.queue = Queue[
            Union[
                Result,
                Error,
                Stdout,
                Stderr,
                EndOfExecution,
                NumberOfExecutions,
                UnexpectedEndOfExecution,
            ]
        ]()
        self.input_accepted = False
        self.errored = False
        self.in_background = in_background


class ContextWebSocket:
    _ws: Optional[WebSocketClientProtocol] = None
    _receive_task: Optional[asyncio.Task] = None
    _global_env_vars: Optional[Dict[StrictStr, str]] = None
    _cleanup_task: Optional[asyncio.Task] = None

    def __init__(self, context_id: str, session_id: str, language: str, cwd: str):
        self.language = language
        self.cwd = cwd
        self.context_id = context_id
        self.url = f"ws://localhost:8888/api/kernels/{context_id}/channels"
        self.session_id = session_id
        self._executions: Dict[str, Execution] = {}
        self._lock = asyncio.Lock()

    async def connect(self):
        logger.debug(f"WebSocket connecting to {self.url}")

        ws_logger = logger.getChild("websockets.client")
        ws_logger.setLevel(logging.ERROR)

        self._ws = await connect(
            self.url,
            max_size=None,
            max_queue=None,
            logger=ws_logger,
        )

        logger.info(f"WebSocket connected to {self.url}")
        self._receive_task = asyncio.create_task(
            self._receive_message(),
            name="receive_message",
        )

    async def reconnect(self, max_retries: int = 5, retry_delay: float = 0.1):
        """Reconnect the WebSocket if it's disconnected with retry logic."""
        logger.info(f"Attempting to reconnect WebSocket {self.context_id}")

        # Close existing connection if any
        if self._ws is not None:
            try:
                await self._ws.close()
            except Exception as e:
                logger.warning(f"Error closing existing WebSocket: {e}")

        # Cancel existing receive task if any
        if self._receive_task is not None and not self._receive_task.done():
            self._receive_task.cancel()
            try:
                await self._receive_task
            except asyncio.CancelledError:
                pass

        # Reset WebSocket and task references
        self._ws = None
        self._receive_task = None

        # Attempt to reconnect with fixed delay
        for attempt in range(max_retries):
            try:
                await self.connect()
                logger.info(
                    f"Successfully reconnected WebSocket {self.context_id} on attempt {attempt + 1}"
                )
                return True
            except Exception as e:
                if attempt < max_retries - 1:
                    logger.warning(
                        f"Reconnection attempt {attempt + 1} failed: {e}. Retrying in {retry_delay}s..."
                    )
                    await asyncio.sleep(retry_delay)
                else:
                    logger.error(
                        f"Failed to reconnect WebSocket {self.context_id} after {max_retries} attempts: {e}"
                    )
                    return False

        return False

    def is_connected(self) -> bool:
        """Check if the WebSocket is connected and healthy."""
        return (
            self._ws is not None
            and not self._ws.closed
            and self._receive_task is not None
            and not self._receive_task.done()
        )

    async def ensure_connected(self):
        """Ensure WebSocket is connected, reconnect if necessary."""
        if not self.is_connected():
            logger.warning(
                f"WebSocket {self.context_id} is not connected, attempting to reconnect"
            )
            success = await self.reconnect()
            if not success:
                raise Exception(f"Failed to reconnect WebSocket {self.context_id}")

    def _get_execute_request(
        self, msg_id: str, code: Union[str, StrictStr], background: bool
    ) -> str:
        return json.dumps(
            {
                "header": {
                    "msg_id": msg_id,
                    "username": "e2b",
                    "session": self.session_id,
                    "msg_type": "execute_request",
                    "version": "5.3",
                    "date": datetime.datetime.now(datetime.timezone.utc).isoformat(),
                },
                "parent_header": {},
                "metadata": {
                    "trusted": True,
                    "deletedCells": [],
                    "recordTiming": False,
                    "cellId": str(uuid.uuid4()),
                },
                "content": {
                    "code": code,
                    "silent": background,
                    "store_history": True,
                    "user_expressions": {},
                    "stop_on_error": True,
                    "allow_stdin": False,
                },
            }
        )

    def _set_env_var_snippet(self, key: str, value: str) -> str:
        """Get environment variable set command for the current language."""
        if self.language == "python":
            return f"import os; os.environ['{key}'] = '{value}'"
        elif self.language in ["javascript", "typescript"]:
            return f"process.env['{key}'] = '{value}'"
        elif self.language == "deno":
            return f"Deno.env.set('{key}', '{value}')"
        elif self.language == "r":
            return f'Sys.setenv({key} = "{value}")'
        elif self.language == "java":
            return f'System.setProperty("{key}", "{value}");'
        elif self.language == "bash":
            return f"export {key}='{value}'"
        return ""

    def _delete_env_var_snippet(self, key: str) -> str:
        """Get environment variable delete command for the current language."""
        if self.language == "python":
            return f"import os; del os.environ['{key}']"
        elif self.language in ["javascript", "typescript"]:
            return f"delete process.env['{key}']"
        elif self.language == "deno":
            return f"Deno.env.delete('{key}')"
        elif self.language == "r":
            return f"Sys.unsetenv('{key}')"
        elif self.language == "java":
            return f'System.clearProperty("{key}");'
        elif self.language == "bash":
            return f"unset {key}"
        return ""

    def _set_env_vars_code(self, env_vars: Dict[StrictStr, str]) -> str:
        """Build environment variable code for the current language."""
        env_commands = []
        for k, v in env_vars.items():
            command = self._set_env_var_snippet(k, v)
            if command:
                env_commands.append(command)

        return "\n".join(env_commands)

    def _reset_env_vars_code(self, env_vars: Dict[StrictStr, str]) -> str:
        """Build environment variable cleanup code for the current language."""
        cleanup_commands = []

        for key in env_vars:
            # Check if this var exists in global env vars
            if self._global_env_vars and key in self._global_env_vars:
                # Reset to global value
                value = self._global_env_vars[key]
                command = self._set_env_var_snippet(key, value)
            else:
                # Remove the variable
                command = self._delete_env_var_snippet(key)

            if command:
                cleanup_commands.append(command)

        return "\n".join(cleanup_commands)

    def _get_code_indentation(self, code: str) -> str:
        """Get the indentation from the first non-empty line of code."""
        if not code or not code.strip():
            return ""

        lines = code.split("\n")
        for line in lines:
            if line.strip():  # First non-empty line
                return line[: len(line) - len(line.lstrip())]

        return ""

    def _indent_code_with_level(self, code: str, indent_level: str) -> str:
        """Apply the given indentation level to each line of code."""
        if not code or not indent_level:
            return code

        lines = code.split("\n")
        indented_lines = []

        for line in lines:
            if line.strip():  # Non-empty lines
                indented_lines.append(indent_level + line)
            else:
                indented_lines.append(line)

        return "\n".join(indented_lines)

    async def _cleanup_env_vars(self, env_vars: Dict[StrictStr, str]):
        """Clean up environment variables in a separate execution request."""
        message_id = str(uuid.uuid4())
        self._executions[message_id] = Execution(in_background=True)

        try:
            cleanup_code = self._reset_env_vars_code(env_vars)
            if cleanup_code:
                logger.info(f"Cleaning up env vars: {cleanup_code}")
                # Ensure WebSocket is connected before sending cleanup request
                await self.ensure_connected()
                request = self._get_execute_request(message_id, cleanup_code, True)
                if self._ws is None:
                    raise Exception("WebSocket not connected")
                await self._ws.send(request)

                async for item in self._wait_for_result(message_id):
                    if isinstance(item, dict) and item.get("type") == "error":
                        logger.error(f"Error during env var cleanup: {item}")
        finally:
            del self._executions[message_id]

    async def _wait_for_result(self, message_id: str):
        queue = self._executions[message_id].queue

        while True:
            output = await queue.get()
            if output.type == OutputType.END_OF_EXECUTION:
                break

            if output.type == OutputType.UNEXPECTED_END_OF_EXECUTION:
                logger.error(f"Unexpected end of execution for code ({message_id})")
                yield Error(
                    name="UnexpectedEndOfExecution",
                    value="Connection to the execution was closed before the execution was finished",
                    traceback="",
                )
                break

            yield output.model_dump(exclude_none=True)

    async def change_current_directory(
        self, path: Union[str, StrictStr], language: str
    ):
        message_id = str(uuid.uuid4())
        self._executions[message_id] = Execution(in_background=True)

        # Ensure WebSocket is connected before changing directory
        await self.ensure_connected()

        if language == "python":
            request = self._get_execute_request(message_id, f"%cd {path}", True)
        elif language == "deno":
            request = self._get_execute_request(
                message_id, f"Deno.chdir('{path}')", True
            )
        elif language == "js":
            request = self._get_execute_request(
                message_id, f"process.chdir('{path}')", True
            )
        elif language == "r":
            request = self._get_execute_request(message_id, f"setwd('{path}')", True)
        # This does not actually change the working directory, but sets the user.dir property
        elif language == "java":
            request = self._get_execute_request(
                message_id, f'System.setProperty("user.dir", "{path}");', True
            )
        else:
            return

        if self._ws is None:
            raise Exception("WebSocket not connected")

        await self._ws.send(request)

        async for item in self._wait_for_result(message_id):
            if isinstance(item, dict) and item.get("type") == "error":
                raise ExecutionError(f"Error during execution: {item}")

    async def execute(
        self,
        code: Union[str, StrictStr],
        env_vars: Dict[StrictStr, str],
        access_token: str,
    ):
        message_id = str(uuid.uuid4())
        self._executions[message_id] = Execution()

        # Ensure WebSocket is connected before executing
        await self.ensure_connected()

        async with self._lock:
            # Wait for any pending cleanup task to complete
            if self._cleanup_task and not self._cleanup_task.done():
                logger.debug("Waiting for pending cleanup task to complete")
                try:
                    await self._cleanup_task
                except Exception as e:
                    logger.warning(f"Cleanup task failed: {e}")
                finally:
                    self._cleanup_task = None

            # Get the indentation level from the code
            code_indent = self._get_code_indentation(code)

            # Build the complete code snippet with env vars
            complete_code = code

            global_env_vars_snippet = ""
            env_vars_snippet = ""

            if self._global_env_vars is None:
                self._global_env_vars = await get_envs(access_token=access_token)
                global_env_vars_snippet = self._set_env_vars_code(self._global_env_vars)

            if env_vars:
                env_vars_snippet = self._set_env_vars_code(env_vars)

            if global_env_vars_snippet or env_vars_snippet:
                indented_env_code = self._indent_code_with_level(
                    f"{global_env_vars_snippet}\n{env_vars_snippet}", code_indent
                )
                complete_code = f"{indented_env_code}\n{complete_code}"

            logger.info(
                f"Sending code for the execution ({message_id}): {complete_code}"
            )
            request = self._get_execute_request(message_id, complete_code, False)

            # Send the code for execution
            if self._ws is None:
                raise Exception("WebSocket not connected")
            await self._ws.send(request)

            # Stream the results
            async for item in self._wait_for_result(message_id):
                yield item

            del self._executions[message_id]

            # Clean up env vars in a separate request after the main code has run
            if env_vars:
                self._cleanup_task = asyncio.create_task(
                    self._cleanup_env_vars(env_vars)
                )

    async def _receive_message(self):
        if not self._ws:
            logger.error("No WebSocket connection")
            return

        try:
            async for message in self._ws:
                await self._process_message(json.loads(message))
        except Exception as e:
            logger.error(f"WebSocket received error while receiving messages: {str(e)}")

            # Attempt to reconnect when connection drops
            logger.info("Attempting to reconnect due to connection loss...")
            reconnect_success = await self.reconnect()

            if reconnect_success:
                logger.info("Successfully reconnected after connection loss")
                # Continue receiving messages with the new connection
                try:
                    async for message in self._ws:
                        await self._process_message(json.loads(message))
                except Exception as reconnect_e:
                    logger.error(f"Error in reconnected WebSocket: {str(reconnect_e)}")

            # Mark all pending executions as failed due to connection loss
            for execution in self._executions.values():
                await execution.queue.put(
                    Error(
                        name="ConnectionLost",
                        value="WebSocket connection was lost during execution",
                        traceback="",
                    )
                )
                await execution.queue.put(UnexpectedEndOfExecution())

    async def _process_message(self, data: dict):
        """
        Process messages from the WebSocket

        Message types:
        https://jupyter-client.readthedocs.io/en/stable/messaging.html

        :param data: The message data
        """
        if (
            data["msg_type"] == "status"
            and data["content"]["execution_state"] == "restarting"
        ):
            logger.error("Context is restarting")
            for execution in self._executions.values():
                await execution.queue.put(
                    Error(
                        name="ContextRestarting",
                        value="Context was restarted",
                        traceback="",
                    )
                )
                await execution.queue.put(EndOfExecution())
            return

        parent_msg_ig = data["parent_header"].get("msg_id", None)
        if parent_msg_ig is None:
            logger.warning("Parent message ID not found. %s", data)
            return

        execution = self._executions.get(parent_msg_ig)
        if not execution:
            return

        queue = execution.queue
        if data["msg_type"] == "error":
            logger.debug(
                f"Execution {parent_msg_ig} finished execution with error: {data['content']['ename']}: {data['content']['evalue']}"
            )

            if execution.errored:
                return

            execution.errored = True
            await queue.put(
                Error(
                    name=data["content"]["ename"],
                    value=data["content"]["evalue"],
                    traceback="".join(data["content"]["traceback"]),
                )
            )

        elif data["msg_type"] == "stream":
            if data["content"]["name"] == "stdout":
                logger.debug(
                    f"Execution {parent_msg_ig} received stdout: {data['content']['text']}"
                )
                await queue.put(
                    Stdout(
                        text=data["content"]["text"], timestamp=data["header"]["date"]
                    )
                )

            elif data["content"]["name"] == "stderr":
                logger.debug(
                    f"Execution {parent_msg_ig} received stderr: {data['content']['text']}"
                )
                await queue.put(
                    Stderr(
                        text=data["content"]["text"], timestamp=data["header"]["date"]
                    )
                )

        elif data["msg_type"] in "display_data":
            result = Result(is_main_result=False, data=data["content"]["data"])
            logger.debug(
                f"Execution {parent_msg_ig} received display data with following formats: {result.formats()}"
            )
            await queue.put(result)

        elif data["msg_type"] == "execute_result":
            result = Result(is_main_result=True, data=data["content"]["data"])
            logger.debug(
                f"Execution {parent_msg_ig} received execution result with following formats: {result.formats()}"
            )
            await queue.put(result)

        elif data["msg_type"] == "status":
            if data["content"]["execution_state"] == "busy" and execution.in_background:
                logger.debug(f"Execution {parent_msg_ig} started execution")
                execution.input_accepted = True

            if data["content"]["execution_state"] == "idle":
                if execution.input_accepted:
                    logger.debug(f"Execution {parent_msg_ig} finished execution")
                    await queue.put(EndOfExecution())

            elif data["content"]["execution_state"] == "error":
                logger.debug(f"Execution {parent_msg_ig} finished execution with error")
                await queue.put(
                    Error(
                        name=data["content"]["ename"],
                        value=data["content"]["evalue"],
                        traceback="".join(data["content"]["traceback"]),
                    )
                )
                await queue.put(EndOfExecution())

        elif data["msg_type"] == "execute_reply":
            if data["content"]["status"] == "error":
                logger.debug(f"Execution {parent_msg_ig} finished execution with error")

                if execution.errored:
                    return

                execution.errored = True
                await queue.put(
                    Error(
                        name=data["content"].get("ename", ""),
                        value=data["content"].get("evalue", ""),
                        traceback="".join(data["content"].get("traceback", [])),
                    )
                )
            elif data["content"]["status"] == "abort":
                logger.debug(f"Execution {parent_msg_ig} was aborted")
                await queue.put(
                    Error(
                        name="ExecutionAborted",
                        value="Execution was aborted",
                        traceback="",
                    )
                )
                await queue.put(EndOfExecution())
            elif data["content"]["status"] == "ok":
                pass

        elif data["msg_type"] == "execute_input":
            logger.debug(f"Input accepted for {parent_msg_ig}")
            await queue.put(
                NumberOfExecutions(execution_count=data["content"]["execution_count"])
            )
            execution.input_accepted = True
        else:
            logger.warning(f"[UNHANDLED MESSAGE TYPE]: {data['msg_type']}")

    async def close(self):
        logger.debug(f"Closing WebSocket {self.context_id}")

        if self._ws is not None:
            await self._ws.close()

        if self._receive_task is not None:
            self._receive_task.cancel()

        # Cancel any pending cleanup task
        if self._cleanup_task and not self._cleanup_task.done():
            self._cleanup_task.cancel()
            try:
                await self._cleanup_task
            except asyncio.CancelledError:
                pass

        for execution in self._executions.values():
            execution.queue.put_nowait(UnexpectedEndOfExecution())