Class JupyterKernelWebSocket

python/e2b_code_interpreter/messaging.py:31–214

class JupyterKernelWebSocket(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)

    url: str

    _cells: Dict[str, Cell] = {}
    _waiting_for_replies: Dict[str, DeferredFuture] = PrivateAttr(default_factory=dict)
    _queue_in: Queue = PrivateAttr(default_factory=Queue)
    _queue_out: Queue = PrivateAttr(default_factory=Queue)
    _process_cleanup: List[Callable[[], Any]] = PrivateAttr(default_factory=list)
    _closed: bool = PrivateAttr(default=False)

    def process_messages(self):
        while True:
            data = self._queue_out.get()

            # logger.debug(f"WebSocket received message: {data}".strip())
            self._receive_message(json.loads(data))
            self._queue_out.task_done()

    def connect(self, timeout: float = TIMEOUT):
        started = threading.Event()
        stopped = threading.Event()
        self._process_cleanup.append(stopped.set)

        threading.Thread(
            target=self.process_messages, daemon=True, name="e2b-process-messages"
        ).start()

        threading.Thread(
            target=WebSocket(
                url=self.url,
                queue_in=self._queue_in,
                queue_out=self._queue_out,
                started=started,
                stopped=stopped,
            ).run,
            daemon=True,
            name="e2b-code-interpreter-websocket",
        ).start()

        # logger.info("WebSocket waiting to start")

        try:
            start_time = time.time()
            while (
                not started.is_set()
                and time.time() - start_time < timeout
                and not self._closed
            ):
                time.sleep(0.1)

            if not started.is_set():
                # logger.error("WebSocket failed to start")
                raise TimeoutException("WebSocket failed to start")
        except BaseException as e:
            self.close()
            raise Exception(f"WebSocket failed to start: {e}") from e