ttsd-colabcli 1.0.0

package/core/app/colab_cli/repl.py

@@ -0,0 +1,173 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import datetime
+from typing import Any, List, Optional
+
+from prompt_toolkit import PromptSession
+from prompt_toolkit.history import InMemoryHistory
+from prompt_toolkit.key_binding import KeyBindings
+from prompt_toolkit.lexers import PygmentsLexer
+from prompt_toolkit.styles import Style
+from pygments.lexers.python import PythonLexer
+from rich.console import Console
+from rich.text import Text
+
+from app.colab_cli.runtime import ColabRuntime
+from app.colab_cli.utils import handle_image
+
+console = Console()
+
+
+class ColabREPL:
+    def __init__(
+        self,
+        runtime: ColabRuntime,
+        session_name: Optional[str] = None,
+        history_logger: Optional[Any] = None,
+        output_image: Optional[str] = None,
+    ):
+        self.runtime = runtime
+        self.session_name = session_name
+        self.history_logger = history_logger
+        self.output_image = output_image
+        self.kb = KeyBindings()
+        self.console = console
+        self.repl_history: List[dict] = []
+
+        @self.kb.add("enter")
+        def _(event):
+            event.current_buffer.validate_and_handle()
+
+        @self.kb.add("escape", "enter")
+        @self.kb.add("c-j")
+        def _(event):
+            event.current_buffer.insert_text("\n")
+
+        self.session = PromptSession(
+            history=InMemoryHistory(),
+            lexer=PygmentsLexer(PythonLexer),
+            include_default_pygments_style=False,
+            key_bindings=self.kb,
+            multiline=True,
+        )
+        self.style = Style.from_dict(
+            {
+                "prompt": "bold blue",
+                "continuation": "#888888",
+            }
+        )
+
+    def print_info(self, message: str):
+        self.console.print(f"[bold blue][*][/bold blue] {message}")
+
+    def print_error(self, message: str):
+        self.console.print(f"[bold red][!][/bold red] {message}")
+
+    def display_output(self, output: dict):
+        if "text" in output:
+            self.console.print(Text.from_ansi(output["text"]), end="")
+        elif "data" in output:
+            data = output["data"]
+
+            # Check for images first
+            image_displayed = False
+            for mime_type in ["image/png", "image/jpeg"]:
+                if mime_type in data:
+                    handle_image(
+                        data[mime_type], mime_type, target_path=self.output_image
+                    )
+                    image_displayed = True
+                    break
+
+            if "text/plain" in data:
+                text = data["text/plain"]
+                # Skip generic IPython object reprs if we already showed an image
+                if image_displayed and any(
+                    x in text for x in ["<IPython.core.display.Image", "<Figure size"]
+                ):
+                    return
+                self.console.print(Text.from_ansi(text))
+        elif output.get("output_type") == "error":
+            ename = output.get("ename", "Error")
+            evalue = output.get("evalue", "")
+            traceback = output.get("traceback", [])
+            if traceback:
+                self.console.print(Text.from_ansi("".join(traceback)))
+            else:
+                self.print_error(f"{ename}: {evalue}")
+
+    def execute(self, code: str):
+        if self.session_name:
+            from app.colab_cli.common import state
+
+            s = state.store.get(self.session_name)
+            if s:
+                s.last_execution = (
+                    "REPL",
+                    None,
+                    datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+                )
+                state.store.add(s)
+
+        try:
+            outputs = self.runtime.execute_code(
+                code, output_hook=lambda o: self.display_output(o)
+            )
+            # Ensure next prompt starts on a newline after streaming
+            print()
+
+            self.repl_history.append({"input": code, "outputs": outputs or []})
+            if self.history_logger and self.session_name:
+                self.history_logger.log_event(
+                    self.session_name,
+                    "execution",
+                    {"code": code, "outputs": outputs or []},
+                )
+        except Exception as e:
+            self.print_error(f"Execution failed: {e}")
+
+    def run(self):
+        self.console.print("Python 3 (Google Colab Runtime)\nType /quit to exit.")
+
+        while True:
+            try:
+                result = self.session.prompt(
+                    ">>> ",
+                    style=self.style,
+                )
+
+                if result is None:
+                    continue
+
+                code = result.strip()
+
+                if not code:
+                    continue
+
+                if code.lower() in ("/quit", "quit()", "exit()"):
+                    break
+
+                self.execute(code)
+
+            except EOFError:
+                break
+            except KeyboardInterrupt:
+                print()
+                continue
+            except Exception as e:
+                self.print_error(f"REPL Error: {e}")
+
+        self.print_info("Goodbye!")
+        self.runtime.stop()

package/core/app/colab_cli/runtime.py

@@ -0,0 +1,262 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import time
+from typing import Any, Callable, Dict, List, Optional
+
+import jupyter_kernel_client
+import requests
+
+
+class ColabRuntime:
+    def __init__(
+        self,
+        url: str,
+        token: str,
+        session_name: Optional[str] = None,
+        history: Optional[Any] = None,
+        kernel_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        on_kernel_started: Optional[Callable[[str], None]] = None,
+        on_session_started: Optional[Callable[[str], None]] = None,
+    ):
+        self.url = url
+        self.token = token
+        self.session_name = session_name
+        self.history = history
+        self.kernel_id = kernel_id
+        self.session_id = session_id
+        self.on_kernel_started = on_kernel_started
+        self.on_session_started = on_session_started
+        self._kernel_client = None
+        self.colab_request_hook: Optional[Callable[[Dict[str, Any], Any], None]] = None
+
+    def _apply_ws_hook(self):
+        wsclient = self._kernel_client._manager.client
+        original_on_message = wsclient.kernel_socket.on_message
+
+        def hooked_on_message(s_ws, message):
+            if not self.colab_request_hook:
+                return original_on_message(s_ws, message)
+
+            try:
+                from jupyter_kernel_client.wsclient import JupyterSubprotocol
+
+                if wsclient._subprotocol == JupyterSubprotocol.DEFAULT:
+                    from jupyter_kernel_client.wsclient import (
+                        deserialize_msg_from_ws_default,
+                    )
+
+                    deserialize_msg = deserialize_msg_from_ws_default(message)
+                elif wsclient._subprotocol == JupyterSubprotocol.V1:
+                    from jupyter_kernel_client.wsclient import (
+                        deserialize_msg_from_ws_v1,
+                    )
+
+                    channel, msg_list = deserialize_msg_from_ws_v1(message)
+                    deserialize_msg = wsclient.session.deserialize(msg_list)
+                else:
+                    deserialize_msg = None
+
+                if deserialize_msg:
+                    msg_type = deserialize_msg.get("msg_type")
+                    if msg_type == "colab_request":
+                        # We pass the deserialized msg and the wsclient to the hook
+                        if self.colab_request_hook(deserialize_msg, wsclient):
+                            # If the hook returns True, we intercept and do NOT pass to original
+                            return
+
+            except Exception as e:
+                logging.debug(f"Error in colab_request hook: {e}")
+
+            # Call original for all other messages
+            original_on_message(s_ws, message)
+
+        wsclient.kernel_socket.on_message = hooked_on_message
+
+    @property
+    def kernel_client(self):
+        if not self._kernel_client:
+            retries = 3
+            backoff = 2
+            last_err = None
+
+            for i in range(retries):
+                try:
+                    client_kwargs = {
+                        "subprotocol": jupyter_kernel_client.JupyterSubprotocol.DEFAULT,
+                        "extra_params": {"colab-runtime-proxy-token": self.token},
+                    }
+                    if self.session_id:
+                        # WSSession (Session) expects 'session' for the ID
+                        client_kwargs["session"] = self.session_id
+
+                    self._kernel_client = jupyter_kernel_client.KernelClient(
+                        server_url=self.url,
+                        token=self.token,
+                        kernel_id=self.kernel_id,
+                        client_kwargs=client_kwargs,
+                        headers={
+                            "X-Colab-Client-Agent": "colab-cli",
+                            "X-Colab-Runtime-Proxy-Token": self.token,
+                        },
+                    )
+                    # Force _own_kernel to False. This prevents jupyter-kernel-client
+                    # from automatically deleting the kernel when the client is closed or deleted.
+                    self._kernel_client._own_kernel = False
+
+                    self._kernel_client.start()
+                    self._apply_ws_hook()
+
+                    # Capture IDs if we started fresh
+                    if not self.kernel_id and self._kernel_client.id:
+                        self.kernel_id = self._kernel_client.id
+                        if self.on_kernel_started:
+                            self.on_kernel_started(self.kernel_id)
+
+                    if (
+                        not self.session_id
+                        and self._kernel_client._manager.client.session.session
+                    ):
+                        self.session_id = (
+                            self._kernel_client._manager.client.session.session
+                        )
+                        if self.on_session_started:
+                            self.on_session_started(self.session_id)
+                    break
+                except (
+                    requests.exceptions.ReadTimeout,
+                    requests.exceptions.ConnectTimeout,
+                ) as e:
+                    last_err = e
+                    if i < retries - 1:
+                        sleep_time = backoff ** (i + 1)
+                        logging.debug(
+                            f"Kernel startup timeout, retrying in {sleep_time}s..."
+                            f" ({i + 1}/{retries})"
+                        )
+                        time.sleep(sleep_time)
+                    else:
+                        raise last_err
+                except Exception as e:
+                    raise e
+
+        return self._kernel_client
+
+    def restart(
+        self,
+        timeout: Optional[float] = None,
+    ):
+        self.kernel_client.restart(timeout=timeout)
+
+    def execute_code(
+        self,
+        code: str,
+        allow_stdin: bool = False,
+        stdin_hook: Any = None,
+        output_hook: Optional[Callable[[Dict[str, Any]], None]] = None,
+        timeout: Optional[float] = None,
+    ) -> List[Dict[str, Any]]:
+        # ``jupyter_kernel_client`` defaults ``timeout`` to ``REQUEST_TIMEOUT``
+        # (10 seconds) on both ``execute`` and ``execute_interactive``. That
+        # value is a wall-clock budget that shrinks every time the poll loop
+        # iterates -- as long as iopub/stdin events arrive back-to-back the
+        # call survives, but a single >10s quiet stretch (e.g. a kernel
+        # blocked on ``input_request`` while the user OAuths in the browser)
+        # will raise ``TimeoutError`` even though the underlying execution is
+        # still healthy. Callers that know they need a longer ceiling can
+        # pass ``timeout=`` here; otherwise we forward whatever the upstream
+        # default is (currently 10s).
+        kwargs = {"allow_stdin": allow_stdin}
+        if timeout is not None:
+            kwargs["timeout"] = timeout
+
+        # Wrap stdin_hook to log inputs
+        original_stdin_hook = stdin_hook
+
+        def wrapped_stdin_hook(prompt):
+            if self.history and self.session_name:
+                self.history.log_event(
+                    self.session_name, "stdin_request", {"prompt": prompt}
+                )
+
+            res = original_stdin_hook(prompt) if original_stdin_hook else input(prompt)
+
+            if self.history and self.session_name:
+                self.history.log_event(self.session_name, "input_reply", {"value": res})
+            return res
+
+        if allow_stdin:
+            kwargs["stdin_hook"] = wrapped_stdin_hook
+
+        if output_hook:
+            # If we have an output hook, we use execute_interactive and manage buffering ourselves
+            outputs = []
+
+            def wrapped_output_hook(msg):
+                from jupyter_kernel_client.client import (
+                    output_hook as default_output_hook,
+                )
+
+                # Update local outputs list using the default logic
+                new_indexes = default_output_hook(outputs, msg)
+                # If new outputs were added, call our streaming hook with the new data
+                if new_indexes:
+                    for idx in sorted(new_indexes):
+                        if idx < len(outputs):
+                            output_hook(outputs[idx])
+
+            reply = self.kernel_client.execute_interactive(
+                code, output_hook=wrapped_output_hook, **kwargs
+            )
+            # execute_interactive returns the raw reply message
+            reply_content = reply["content"] if reply else {"status": "error"}
+        else:
+            reply = self.kernel_client.execute(code, **kwargs)
+            if not reply:
+                return []
+            outputs = reply.get("outputs", [])
+            reply_content = reply
+
+        # If there's an error status but no error in outputs, synthesize one
+        if reply_content.get("status") == "error":
+            has_error_output = any(o.get("output_type") == "error" for o in outputs)
+            if not has_error_output:
+                outputs.append(
+                    {
+                        "output_type": "error",
+                        "ename": reply_content.get("ename", "Error"),
+                        "evalue": reply_content.get("evalue", "Unknown error"),
+                        "traceback": reply_content.get("traceback", []),
+                    }
+                )
+
+        return outputs
+
+    def stop(self, shutdown_kernel: bool = False):
+        if self._kernel_client:
+            try:
+                # We manage kernel lifecycle explicitly.
+                # To prevent automatic shutdown, we bypass the manager's stop() and
+                # directly close the channels and socket.
+                client = self._kernel_client._manager.client
+                client.stop_channels()
+                if client.kernel_socket:
+                    client.kernel_socket.close()
+
+                if shutdown_kernel:
+                    self._kernel_client._manager.shutdown_kernel(now=True)
+            except Exception:
+                logging.exception("Error stopping kernel client")

package/core/app/colab_cli/state.py

@@ -0,0 +1,156 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import contextlib
+import json
+import os
+from datetime import datetime
+from typing import Dict, Optional, Tuple, Iterator, IO
+
+import filelock
+from pydantic import BaseModel
+
+
+class SessionState(BaseModel):
+    name: str
+    token: str
+    url: str
+    endpoint: str
+    variant: str = "DEFAULT"
+    accelerator: str = "NONE"
+    kernel_id: Optional[str] = None
+    session_id: Optional[str] = None
+    last_execution: Optional[Tuple[str, Optional[str], str]] = None
+    running: Optional[str] = None
+    keep_alive_pid: Optional[int] = None
+
+
+class Settings(BaseModel):
+    update_url: str = "https://pypi.org/pypi/google-colab-cli/json"
+    last_check: Optional[datetime] = None
+    enable_update_check: bool = True
+    # Highest version seen on the update source; cached for the banner.
+    latest_version: Optional[str] = None
+
+
+class _LockedFileStore:
+    def __init__(self, path: str):
+        self.path = path
+        self.lock_path = "%s.lock" % self.path
+        # ReadWriteLock gives us shared (concurrent) readers and exclusive
+        # writers -- the cross-platform equivalent of fcntl LOCK_SH/LOCK_EX.
+        # is_singleton=False keeps each store's lock independent: with the
+        # default (True), two StateStore instances for the same path in one
+        # process are merged into a single reentrant lock, whose reentrancy
+        # guard then raises RuntimeError when two threads contend for the write
+        # lock. We want them to actually serialize via the underlying file lock.
+        self._rwlock = filelock.ReadWriteLock(self.lock_path, is_singleton=False)
+        self._ensure_dir()
+
+    def _ensure_dir(self):
+        os.makedirs(os.path.dirname(self.path), exist_ok=True)
+
+    def _write_data(self, f: IO, data: str):
+        f.seek(0)
+        f.truncate()
+        f.write(data)
+        f.flush()
+        os.fsync(f.fileno())
+
+    @contextlib.contextmanager
+    def _lock_shared(self) -> Iterator[Optional[IO]]:
+        if not os.path.exists(self.path):
+            yield None
+            return
+        with self._rwlock.read_lock():
+            with open(self.path, "r") as f:
+                yield f
+
+    @contextlib.contextmanager
+    def _lock_exclusive(self) -> Iterator[IO]:
+        with self._rwlock.write_lock():
+            with open(self.path, "a+") as f:
+                yield f
+
+
+class SettingsStore(_LockedFileStore):
+    def __init__(self, path: Optional[str] = None):
+        if not path:
+            path = os.path.expanduser("~/.config/colab-cli/settings.json")
+        super().__init__(path)
+
+    def load(self) -> Settings:
+        with self._lock_shared() as f:
+            if f is None:
+                return Settings()
+            try:
+                content = f.read()
+                if not content or content.isspace():
+                    return Settings()
+                data = json.loads(content)
+                return Settings.model_validate(data)
+            except Exception:
+                return Settings()
+
+    def save(self, settings: Settings):
+        with self._lock_exclusive() as f:
+            self._write_data(f, settings.model_dump_json(indent=2))
+
+
+class StateStore(_LockedFileStore):
+    def __init__(self, path: Optional[str] = None):
+        if not path:
+            path = os.path.expanduser("~/.config/colab-cli/sessions.json")
+        super().__init__(path)
+
+    def _load_raw(self, f) -> Dict[str, SessionState]:
+        try:
+            f.seek(0)
+            content = f.read()
+            if not content or content.isspace():
+                return {}
+            data = json.loads(content)
+            return {k: SessionState(**v) for k, v in data.items()}
+        except Exception:
+            return {}
+
+    def _save_raw(self, f, sessions: Dict[str, SessionState]):
+        content = json.dumps({k: v.model_dump() for k, v in sessions.items()}, indent=2)
+        self._write_data(f, content)
+
+    def add(self, state: SessionState):
+        with self._lock_exclusive() as f:
+            sessions = self._load_raw(f)
+            sessions[state.name] = state
+            self._save_raw(f, sessions)
+
+    def get(self, name: str) -> Optional[SessionState]:
+        with self._lock_shared() as f:
+            if f is None:
+                return None
+            sessions = self._load_raw(f)
+            return sessions.get(name)
+
+    def remove(self, name: str):
+        with self._lock_exclusive() as f:
+            sessions = self._load_raw(f)
+            if name in sessions:
+                del sessions[name]
+                self._save_raw(f, sessions)
+
+    def list(self) -> Dict[str, SessionState]:
+        with self._lock_shared() as f:
+            if f is None:
+                return {}
+            return self._load_raw(f)

package/core/app/colab_cli/utils.py

@@ -0,0 +1,85 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import base64
+import logging
+import sys
+import tempfile
+
+
+from typing import Optional
+
+
+def get_status_code(e: Exception) -> Optional[int]:
+    """Safely extracts status code from various exception types."""
+    if hasattr(e, "response") and e.response is not None:
+        if hasattr(e.response, "status_code"):
+            return e.response.status_code
+    if hasattr(e, "status_code"):
+        return e.status_code
+    return None
+
+
+def is_terminal_error(e: Exception) -> bool:
+    """Checks if an exception indicates a lost session (404/401)."""
+    code = get_status_code(e)
+    if code in (404, 401):
+        return True
+    # Some exceptions from jupyter-kernel-client might wrap the real one or be different
+    err_msg = str(e)
+    if "404" in err_msg or "401" in err_msg:
+        return True
+    return False
+
+
+def print_kitty(image_bytes: bytes):
+    """
+    Outputs an image using the Kitty Graphics Protocol.
+    Expects PNG bytes.
+
+    No-op when stdout is not a TTY: the escape sequence is meaningless to a
+    file/pipe and visually corrupts captured output (e.g. when piping
+    `colab exec` into a shell tool, redirecting to a log file, or running
+    under non-Kitty terminals). Callers still get the image via
+    `handle_image`'s file-write path.
+    """
+    if not sys.stdout.isatty():
+        return
+    try:
+        b64_data = base64.b64encode(image_bytes).decode("ascii")
+        sys.stdout.write("\n\033_Ga=T,f=100;")
+        sys.stdout.write(b64_data)
+        sys.stdout.write("\033\\\n")
+        sys.stdout.flush()
+    except Exception:
+        logging.exception("Kitty rendering failed")
+
+
+def handle_image(image_b64: str, mime_type: str = "image/png", target_path: str = None):
+    image_bytes = base64.b64decode(image_b64)
+    # Print inline using Kitty protocol
+    print_kitty(image_bytes)
+
+    if target_path:
+        # If a target path is specified, save it there
+        with open(target_path, "wb") as f:
+            f.write(image_bytes)
+        print(f"\n[Image saved to: {target_path}]")
+    else:
+        # Save to temp file as fallback
+        ext = mime_type.split("/")[-1]
+        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=f".{ext}")
+        tmp.write(image_bytes)
+        tmp.close()
+        print(f"\n[Image saved to: {tmp.name}]")