ttsd-colabcli 1.0.0

package/core/app/colab_cli/console.py

@@ -0,0 +1,172 @@
+# 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 json
+import logging
+import os
+import signal
+import sys
+import termios
+import threading
+import time
+import tty
+from urllib.parse import urlparse
+
+import websocket
+
+from app.colab_cli.state import SessionState
+
+logger = logging.getLogger(__name__)
+
+# Global flag to stop the read thread when the websocket closes
+_is_running = False
+_last_error = None
+
+# When stdin is piped and reaches EOF, we send "exit\n" to the remote shell and
+# then wait this many seconds for any remaining output (the shell's goodbye,
+# tmux teardown messages, etc.) to flush before closing the websocket from the
+# client side. Empirically 0.5s is enough for the typical /colab/tty backend
+# wrapped in tmux + bash; bumping it just delays exit, lowering it risks
+# truncating tail output.
+PIPED_EOF_GRACE_SECONDS = 0.5
+
+
+def on_message(ws, message):
+    """Callback for when a message is received from the server."""
+    try:
+        data = json.loads(message)
+        if "data" in data:
+            # The backend sends raw ANSI escape sequences and string content.
+            # We write it directly to stdout buffer to avoid python print() formatting.
+            sys.stdout.buffer.write(data["data"].encode("utf-8"))
+            sys.stdout.buffer.flush()
+    except Exception as e:
+        logger.debug(f"Error parsing message: {e}")
+
+
+def on_error(ws, error):
+    """Callback for when a websocket error occurs."""
+    global _last_error
+    _last_error = error
+    logger.error(f"WebSocket Error: {error}")
+
+
+def on_close(ws, close_status_code, close_msg):
+    """Callback for when the websocket is closed."""
+    global _is_running
+    _is_running = False
+
+
+def send_terminal_size(ws):
+    """Sends the current terminal size to the remote backend."""
+    try:
+        size = os.get_terminal_size()
+        payload = json.dumps({"cols": size.columns, "rows": size.lines})
+        ws.send(payload)
+    except Exception as e:
+        logger.debug(f"Failed to send terminal size: {e}")
+
+
+def on_open(ws):
+    """Callback for when the websocket connection is opened."""
+    global _is_running
+    _is_running = True
+
+    # Send initial terminal size
+    send_terminal_size(ws)
+
+    # Setup the background thread to read from stdin
+    def read_stdin():
+        is_tty = sys.stdin.isatty()
+        while _is_running:
+            try:
+                # Read a single character (or escape sequence byte)
+                char = sys.stdin.read(1)
+                if not char:
+                    if not is_tty:
+                        # Piped input has reached EOF. The remote /colab/tty
+                        # endpoint wraps bash in tmux which intercepts \x04
+                        # (Ctrl-D) as a literal character, so it never exits.
+                        # Instead send "exit\n" so bash voluntarily terminates,
+                        # wait a short grace period for the shell's goodbye
+                        # output to drain back to us, then close the websocket
+                        # ourselves to guarantee the client unblocks.
+                        try:
+                            ws.send(json.dumps({"data": "exit\n"}))
+                        except Exception:
+                            pass
+                        time.sleep(PIPED_EOF_GRACE_SECONDS)
+                        try:
+                            ws.close()
+                        except Exception:
+                            pass
+                    break
+                ws.send(json.dumps({"data": char}))
+            except Exception:
+                break
+
+    thread = threading.Thread(target=read_stdin, daemon=True)
+    thread.start()
+
+
+def connect_console(session: SessionState):
+    """
+    Connects to the Colab TTY endpoint and sets up a raw terminal session.
+    """
+    global _is_running, _last_error
+    _last_error = None
+
+    # Construct the WebSocket URL from the base URL
+    parsed = urlparse(session.url)
+    ws_scheme = "wss" if parsed.scheme == "https" else "ws"
+    ws_url = f"{ws_scheme}://{parsed.netloc}/colab/tty?colab-runtime-proxy-token={session.token}"
+
+    is_tty = sys.stdin.isatty()
+    fd = sys.stdin.fileno() if is_tty else None
+    old_settings = termios.tcgetattr(fd) if is_tty else None
+
+    ws = websocket.WebSocketApp(
+        url=ws_url,
+        on_open=on_open,
+        on_message=on_message,
+        on_error=on_error,
+        on_close=on_close,
+    )
+
+    def handle_sigwinch(signum, frame):
+        """Handle window resize events."""
+        if _is_running:
+            send_terminal_size(ws)
+
+    try:
+        if is_tty:
+            tty.setraw(fd, termios.TCSANOW)
+            signal.signal(signal.SIGWINCH, handle_sigwinch)
+
+        # This is a blocking call until the connection is closed
+        ws.run_forever()
+
+        if _last_error:
+            # Re-raise or wrap terminal errors
+            err_msg = str(_last_error)
+            if "404" in err_msg or "401" in err_msg:
+                # We raise a standard exception that the caller can recognize
+                raise RuntimeError(f"Connection failed: {err_msg}")
+    finally:
+        if is_tty:
+            # Always ensure the terminal is restored to its original state
+            termios.tcsetattr(fd, termios.TCSANOW, old_settings)
+            # Restore the default signal handler for resize
+            signal.signal(signal.SIGWINCH, signal.SIG_DFL)
+        print("\r\nConnection closed.")

package/core/app/colab_cli/contents.py

@@ -0,0 +1,93 @@
+# 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
+from urllib.parse import quote
+
+import requests
+
+from app.colab_cli.state import SessionState
+from app.colab_cli.utils import get_status_code
+
+
+class ContentsClient:
+    def __init__(self, session_state: SessionState):
+        self.base_url = session_state.url.rstrip("/")
+        self.token = session_state.token
+
+    def _request(
+        self, method: str, path: str, params: dict = None, json_data: dict = None
+    ):
+        # Quote the path, but don't encode slashes so directory paths stay intact
+        quoted_path = quote(path.strip("/"), safe="/")
+        url = f"{self.base_url}/api/contents/{quoted_path}"
+
+        req_params = {"authuser": "0", "colab-runtime-proxy-token": self.token}
+        if params:
+            req_params.update(params)
+
+        response = requests.request(method, url, params=req_params, json=json_data)
+
+        if get_status_code(response) == 404:
+            raise FileNotFoundError(f"File or directory not found: {path}")
+
+        response.raise_for_status()
+
+        # DELETE doesn't return JSON
+        if method == "DELETE":
+            return None
+
+        return response.json()
+
+    def list_dir(self, path: str):
+        return self._request("GET", path)
+
+    def upload(self, local_path: str, remote_path: str):
+        with open(local_path, "rb") as f:
+            content = f.read()
+
+        b64_content = base64.b64encode(content).decode("ascii")
+        filename = remote_path.split("/")[-1]
+
+        payload = {
+            "name": filename,
+            "path": remote_path,
+            "type": "file",
+            "format": "base64",
+            "content": b64_content,
+            "chunk": 1,
+        }
+
+        return self._request("PUT", remote_path, json_data=payload)
+
+    def download(self, remote_path: str, local_path: str):
+        data = self._request("GET", remote_path, params={"content": "1"})
+
+        if data.get("type") == "directory":
+            raise IsADirectoryError(f"Cannot download a directory: {remote_path}")
+
+        content = data.get("content", "")
+        fmt = data.get("format")
+
+        if fmt == "base64":
+            content_bytes = base64.b64decode(content)
+        else:
+            # Assume text if it's not base64 explicitly encoded
+            content_bytes = str(content).encode("utf-8")
+
+        with open(local_path, "wb") as f:
+            f.write(content_bytes)
+
+    def rm(self, remote_path: str):
+        self._request("DELETE", remote_path)

package/core/app/colab_cli/converter.py

@@ -0,0 +1,184 @@
+# 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 nbformat
+from typing import Any, Dict, List
+import json
+import os
+import uuid
+
+
+def export_history(events: List[Dict[str, Any]], session_name: str, output_path: str):
+    """
+    Exports history based on file extension.
+    """
+    ext = os.path.splitext(output_path)[1].lower()
+
+    if ext == ".ipynb":
+        nb = convert_history_to_ipynb(events, session_name)
+        with open(output_path, "w", encoding="utf-8") as f:
+            nbformat.write(nb, f)
+
+    elif ext == ".jsonl":
+        with open(output_path, "w", encoding="utf-8") as f:
+            for event in events:
+                f.write(json.dumps(event) + "\n")
+
+    elif ext == ".md":
+        with open(output_path, "w", encoding="utf-8") as f:
+            f.write(f"# Colab Session: {session_name}\n\n")
+            for event in events:
+                ts = event.get("timestamp", "").split(".")[0].replace("T", " ")
+                etype = event.get("event_type")
+                if etype == "execution":
+                    code = event.get("code", "")
+                    f.write(f"### Execution ({ts})\n```python\n{code}\n```\n\n")
+                    for o in event.get("outputs", []):
+                        if "text" in o:
+                            f.write(f"**Output**:\n```\n{o['text']}```\n\n")
+                elif etype == "session_created":
+                    f.write(
+                        f"## Session Created: {ts}\n- Endpoint: `{event.get('endpoint')}`\n\n"
+                    )
+                elif etype == "file_operation":
+                    f.write(
+                        f"*File Operation*: `{event.get('op')}` on `{event.get('path', event.get('remote', ''))}`\n\n"
+                    )
+
+    elif ext == ".txt":
+        with open(output_path, "w", encoding="utf-8") as f:
+            f.write(f"Colab Session: {session_name}\n" + "=" * 20 + "\n\n")
+            for event in events:
+                ts = event.get("timestamp", "").split(".")[0].replace("T", " ")
+                etype = event.get("event_type", "unknown")
+                f.write(f"[{ts}] {etype.upper()}: ")
+                if etype == "execution":
+                    f.write(event.get("code", "").strip() + "\n")
+                else:
+                    f.write(str(event) + "\n")
+
+    else:
+        print(f"[colab] Unsupported export format: {ext}")
+        return
+
+    print(f"[colab] Exported history to '{output_path}'.")
+
+
+def convert_history_to_ipynb(
+    events: List[Dict[str, Any]], session_name: str
+) -> nbformat.NotebookNode:
+    """
+    Converts a list of session events to a Jupyter Notebook (v4).
+    """
+    nb = nbformat.v4.new_notebook()
+    nb.metadata.kernelspec = {
+        "display_name": "Python 3 (Google Colab)",
+        "language": "python",
+        "name": "python3",
+    }
+
+    title = f"# Colab Session: {session_name}\nGenerated from colab-cli history log."
+    cell = nbformat.v4.new_markdown_cell(title)
+    cell.id = str(uuid.uuid4())
+    nb.cells.append(cell)
+
+    for event in events:
+        etype = event.get("event_type")
+        ts = event.get("timestamp", "").split(".")[0].replace("T", " ")
+
+        if etype == "session_created":
+            meta = f"**Session Created**: {ts}\n- Endpoint: `{event.get('endpoint')}`\n- Hardware: `{event.get('accelerator')}`"
+            cell = nbformat.v4.new_markdown_cell(meta)
+            cell.id = str(uuid.uuid4())
+            nb.cells.append(cell)
+
+        elif etype == "execution":
+            code = event.get("code", "")
+            # Check for shell commands (starting with ! or from piped console)
+            # If it's a raw shell command from a console pipe, wrap it in %%bash if it doesn't have !
+            if event.get("source") == "piped" and not code.startswith("!"):
+                code = "%%bash\n" + code
+
+            outputs = _map_outputs(event.get("outputs", []))
+            cell = nbformat.v4.new_code_cell(code, outputs=outputs)
+            cell.id = str(uuid.uuid4())
+            nb.cells.append(cell)
+
+        elif etype == "automation":
+            op = event.get("op")
+            code = event.get("code", "")
+            cell = nbformat.v4.new_markdown_cell(f"### Automation: {op} ({ts})")
+            cell.id = str(uuid.uuid4())
+            nb.cells.append(cell)
+            if code:
+                # Get the result from the next event if it's automation_result
+                cell = nbformat.v4.new_code_cell(code)
+                cell.id = str(uuid.uuid4())
+                nb.cells.append(cell)
+
+        elif etype == "automation_result":
+            # We can attach these outputs to the previous automation cell if we were more clever,
+            # but for now let's just ensure we capture the output.
+            if event.get("outputs"):
+                cell = nbformat.v4.new_code_cell(
+                    "# Result of previous automation",
+                    outputs=_map_outputs(event.get("outputs")),
+                )
+                cell.id = str(uuid.uuid4())
+                nb.cells.append(cell)
+
+        elif etype == "file_operation":
+            cell = nbformat.v4.new_markdown_cell(
+                f"*File Operation*: `{event.get('op')}` on `{event.get('path', event.get('remote', ''))}`"
+            )
+            cell.id = str(uuid.uuid4())
+            nb.cells.append(cell)
+
+        elif etype == "stdin_request":
+            cell = nbformat.v4.new_markdown_cell(
+                f"> **Input Requested**: {event.get('prompt')}"
+            )
+            cell.id = str(uuid.uuid4())
+            nb.cells.append(cell)
+
+        elif etype == "input_reply":
+            cell = nbformat.v4.new_markdown_cell(
+                f"> **User Input**: `{event.get('value')}`"
+            )
+            cell.id = str(uuid.uuid4())
+            nb.cells.append(cell)
+
+    return nb
+
+
+def _map_outputs(outputs: List[Dict[str, Any]]) -> List[nbformat.NotebookNode]:
+    nb_outputs = []
+    for o in outputs:
+        otype = o.get("output_type")
+        if "text" in o:
+            nb_outputs.append(
+                nbformat.v4.new_output("stream", name="stdout", text=o["text"])
+            )
+        elif "data" in o:
+            nb_outputs.append(nbformat.v4.new_output("display_data", data=o["data"]))
+        elif otype == "error":
+            nb_outputs.append(
+                nbformat.v4.new_output(
+                    "error",
+                    ename=o.get("ename", "Error"),
+                    evalue=o.get("evalue", ""),
+                    traceback=o.get("traceback", []),
+                )
+            )
+    return nb_outputs

package/core/app/colab_cli/history.py

@@ -0,0 +1,65 @@
+# 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
+import json
+import os
+from typing import Any, Dict, List
+
+
+class HistoryLogger:
+    def __init__(self, log_dir: str = "~/.config/colab-cli/history"):
+        self.log_dir = os.path.expanduser(log_dir)
+        os.makedirs(self.log_dir, exist_ok=True)
+
+    def _get_log_path(self, session_name: str) -> str:
+        return os.path.join(self.log_dir, f"{session_name}.jsonl")
+
+    def log_event(self, session_name: str, event_type: str, data: Dict[str, Any]):
+        """
+        Appends a structured event to the session's history file.
+
+        event_types:
+          - session_created
+          - session_terminated
+          - execution (code + outputs)
+          - input_requested (stdin prompts/replies)
+          - file_operation (ls, rm, upload, download)
+          - automation (auth, install, drivemount)
+        """
+        log_path = self._get_log_path(session_name)
+        event = {
+            "timestamp": datetime.datetime.now(datetime.timezone.utc).isoformat(),
+            "event_type": event_type,
+            **data,
+        }
+        with open(log_path, "a", encoding="utf-8") as f:
+            f.write(json.dumps(event) + "\n")
+
+    def list_sessions(self) -> List[str]:
+        if not os.path.exists(self.log_dir):
+            return []
+        return [f[:-6] for f in os.listdir(self.log_dir) if f.endswith(".jsonl")]
+
+    def get_history(self, session_name: str) -> List[Dict[str, Any]]:
+        log_path = self._get_log_path(session_name)
+        if not os.path.exists(log_path):
+            return []
+
+        history = []
+        with open(log_path, "r", encoding="utf-8") as f:
+            for line in f:
+                if line.strip():
+                    history.append(json.loads(line))
+        return history

package/core/app/colab_cli/oauth_config.json

@@ -0,0 +1,11 @@
+{
+  "installed": {
+    "client_id": "764086051850-6qr4p6gpi6hn506pt8ejuq83di341hur.apps.googleusercontent.com",
+    "project_id": "colab-cli",
+    "auth_uri": "https://accounts.google.com/o/oauth2/auth",
+    "token_uri": "https://oauth2.googleapis.com/token",
+    "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
+    "client_secret": "d-FL95Q19q7MQmFpd7hHD0Ty",
+    "redirect_uris": ["http://localhost"]
+  }
+}