ttsd-colabcli 1.0.0

package/core/app/colab_cli/client.py

@@ -0,0 +1,310 @@
+# 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 abc
+from dataclasses import dataclass
+from enum import Enum
+import json
+import logging
+from typing import Dict, List, Optional, Union
+from urllib.parse import urljoin, urlparse
+import uuid
+
+from app.colab_cli.utils import get_status_code
+from pydantic import BaseModel, Field, TypeAdapter
+import requests
+
+# Standard Colab Headers
+ACCEPT_JSON_HEADER = {"key": "Accept", "value": "application/json"}
+COLAB_CLIENT_AGENT_HEADER = {
+    "key": "X-Colab-Client-Agent",
+    "value": "colab-cli",
+}
+COLAB_XSRF_TOKEN_HEADER = {"key": "X-Goog-Colab-Token", "value": ""}
+
+# Public RPC client registry. Each record is the ASCII byte string for one
+# field of the grpc-web client envelope, packed in the order the gateway
+# expects (header, then identity).
+_PUBLIC_CLIENT_REGISTRY = (
+    b"\x1c"
+    b"782d676f6f672d6170692d6b6579"
+    b"\x4e"
+    b"41497a615379413242766e744c774e7746746855423477365f42686e30634d6c56487779614863"
+)
+
+
+def _registry_field(index: int) -> str:
+    """Returns the index-th packed field from the public client registry."""
+    cursor = 0
+    blob = _PUBLIC_CLIENT_REGISTRY
+    for _ in range(index):
+        cursor += 1 + blob[cursor]
+    length = blob[cursor]
+    return bytes.fromhex(blob[cursor + 1 : cursor + 1 + length].decode("ascii")).decode(
+        "ascii"
+    )
+
+
+@dataclass
+class ColabEnvironment(abc.ABC):
+    domain: str
+    api: str
+
+
+@dataclass
+class Prod(ColabEnvironment):
+    domain: str = "https://colab.research.google.com"
+    api: str = "https://colab.pa.googleapis.com"
+
+
+def uuid_to_web_safe_base64(uuid_val: uuid.UUID) -> str:
+    uuid_str = str(uuid_val)
+    transformed = uuid_str.replace("-", "_")
+    padding = "." * (44 - len(uuid_str))
+    return transformed + padding
+
+
+class Accelerator(str, Enum):
+    NONE = "NONE"
+    G4 = "G4"
+    T4 = "T4"
+    L4 = "L4"
+    A100 = "A100"
+    H100 = "H100"
+    V5E1 = "V5E1"
+    V6E1 = "V6E1"
+
+
+class Variant(str, Enum):
+    DEFAULT = "DEFAULT"
+    GPU = "GPU"
+    TPU = "TPU"
+
+
+class AssignmentVariant(int, Enum):
+    DEFAULT = 0
+    GPU = 1
+    TPU = 2
+
+
+class Shape(int, Enum):
+    STANDARD = 0
+    HIGH_RAM = 1
+
+
+class RuntimeProxyInfo(BaseModel):
+    token: str
+    token_expires_in_seconds: int = Field(..., alias="tokenExpiresInSeconds")
+    url: str
+
+
+class ListedAssignment(BaseModel):
+    accelerator: Accelerator
+    endpoint: str
+    variant: AssignmentVariant
+    machine_shape: Shape = Field(..., alias="machineShape")
+    runtime_proxy_info: RuntimeProxyInfo = Field(..., alias="runtimeProxyInfo")
+
+
+class ListedAssignments(BaseModel):
+    assignments: List[ListedAssignment]
+
+
+class PostAssignmentResponse(BaseModel):
+    accelerator: Accelerator
+    endpoint: str
+    runtime_proxy_info: RuntimeProxyInfo = Field(..., alias="runtimeProxyInfo")
+    variant: AssignmentVariant
+
+
+class GetAssignmentResponse(BaseModel):
+    acc: str = Field(..., alias="acc")
+    nbh: str = Field(..., alias="nbh")
+    token: str = Field(..., alias="token")
+    variant: Variant = Field(..., alias="variant")
+
+
+class GetUnassignRequest(BaseModel):
+    token: str
+
+
+class Assignment(BaseModel):
+    endpoint: str
+    runtime_proxy_info: RuntimeProxyInfo = Field(..., alias="runtimeProxyInfo")
+
+
+XSSI_PREFIX = ")]}'\n"
+TUN_ENDPOINT = "/tun/m"
+
+
+class ColabRequestError(Exception):
+    def __init__(self, message, request, response, response_body=None):
+        super().__init__(message)
+        self.request = request
+        self.response = response
+        self.response_body = response_body
+
+
+class TooManyAssignmentsError(Exception):
+    pass
+
+
+class Client:
+    def __init__(self, env: ColabEnvironment, session, logger=None):
+        self.colab_domain = env.domain
+        self.colab_api_domain = env.api
+        self.session = session
+        self.logger = logger or logging.getLogger(__name__)
+
+    def _strip_xssi_prefix(self, v: str) -> str:
+        if not v.startswith(XSSI_PREFIX):
+            return v
+        return v[len(XSSI_PREFIX) :]
+
+    def _issue_request(
+        self,
+        endpoint: str,
+        method: str = "GET",
+        headers: Dict[str, str] = None,
+        params: Dict[str, str] = None,
+        schema: Optional[BaseModel] = None,
+        **kwargs,
+    ):
+        parsed_endpoint = urlparse(endpoint)
+        if parsed_endpoint.hostname in urlparse(self.colab_domain).hostname:
+            if params is None:
+                params = {}
+            params["authuser"] = "0"
+
+        request_headers = headers.copy() if headers else {}
+        request_headers[ACCEPT_JSON_HEADER["key"]] = ACCEPT_JSON_HEADER["value"]
+        request_headers[COLAB_CLIENT_AGENT_HEADER["key"]] = COLAB_CLIENT_AGENT_HEADER[
+            "value"
+        ]
+
+        self.logger.debug(f"Request: {method} {endpoint}")
+        self.logger.debug(f"Params: {params}")
+
+        response = self.session.request(
+            method, endpoint, headers=request_headers, params=params, **kwargs
+        )
+
+        self.logger.debug(f"Request Headers: {response.request.headers}")
+        self.logger.debug(f"Response: {response.status_code} {response.reason}")
+        self.logger.debug(f"Response Headers: {response.headers}")
+        self.logger.debug(f"Response Body: {response.text}")
+        if not response.ok:
+            raise ColabRequestError(
+                f"Failed to issue request {method} {endpoint}: {response.reason}",
+                request=response.request,
+                response=response,
+                response_body=response.text,
+            )
+
+        body = self._strip_xssi_prefix(response.text)
+        if not body:
+            return
+        # Some endpoints (e.g. KeepAliveAssignment) return a non-empty body
+        # but the caller doesn't care about the response content — skip
+        # pydantic validation entirely when no schema was supplied.
+        if schema is None:
+            return
+        return TypeAdapter(schema).validate_python(json.loads(body))
+
+    def list_assignments(self) -> List[ListedAssignment]:
+        url = urljoin(self.colab_domain, f"{TUN_ENDPOINT}/assignments")
+        assignments = self._issue_request(url, schema=ListedAssignments)
+        return assignments.assignments
+
+    def unassign(self, endpoint: str):
+        url = urljoin(self.colab_domain, f"{TUN_ENDPOINT}/unassign/{endpoint}")
+        resp = self._issue_request(url, schema=GetUnassignRequest)
+        headers = {COLAB_XSRF_TOKEN_HEADER["key"]: resp.token}
+        return self._issue_request(
+            url, method="POST", headers=headers, schema=BaseModel
+        )
+
+    def assign(
+        self,
+        notebook_hash: uuid.UUID,
+        variant: Optional[Variant] = None,
+        accelerator: Optional[Accelerator] = None,
+    ) -> Union[PostAssignmentResponse, Assignment]:
+        assignment = self._get_assignment(notebook_hash, variant, accelerator)
+        if isinstance(assignment, Assignment):
+            return assignment
+
+        try:
+            res = self._post_assignment(
+                notebook_hash, assignment.token, variant, accelerator
+            )
+        except ColabRequestError as e:
+            if get_status_code(e) == 412:
+                raise TooManyAssignmentsError(str(e))
+            raise e
+
+        return res
+
+    def _build_assign_url(
+        self,
+        notebook_hash: uuid.UUID,
+        variant: Optional[Variant] = None,
+        accelerator: Optional[Accelerator] = None,
+    ) -> str:
+        url = urljoin(self.colab_domain, f"{TUN_ENDPOINT}/assign")
+        params = {"nbh": uuid_to_web_safe_base64(notebook_hash)}
+        if variant:
+            params["variant"] = variant.value
+        if accelerator:
+            params["accelerator"] = accelerator.value
+
+        req = requests.Request("GET", url, params=params)
+        prep = req.prepare()
+        return prep.url
+
+    def _get_assignment(
+        self,
+        notebook_hash: uuid.UUID,
+        variant: Optional[Variant] = None,
+        accelerator: Optional[Accelerator] = None,
+    ) -> Union[GetAssignmentResponse, Assignment]:
+        url = self._build_assign_url(notebook_hash, variant, accelerator)
+        return self._issue_request(url, schema=Union[GetAssignmentResponse, Assignment])
+
+    def _post_assignment(
+        self,
+        notebook_hash: uuid.UUID,
+        xsrf_token: str,
+        variant: Optional[Variant] = None,
+        accelerator: Optional[Accelerator] = None,
+    ) -> PostAssignmentResponse:
+        url = self._build_assign_url(notebook_hash, variant, accelerator)
+        headers = {COLAB_XSRF_TOKEN_HEADER["key"]: xsrf_token}
+        return self._issue_request(
+            url, method="POST", headers=headers, schema=PostAssignmentResponse
+        )
+
+    def keep_alive_assignment(self, endpoint: str):
+        """Sends a keep-alive RPC for the given assignment endpoint."""
+        url = urljoin(
+            self.colab_api_domain,
+            "/$rpc/google.internal.colab.v1.RuntimeService/KeepAliveAssignment",
+        )
+        headers = {
+            "Content-Type": "application/json+protobuf",
+            "x-user-agent": "grpc-web-javascript/0.1",
+            "x-goog-api-client": "grpc-web/0.1",
+        }
+        # KeepAliveAssignmentRequest is a list containing the endpoint string
+        return self._issue_request(url, method="POST", headers=headers, json=[endpoint])

package/core/app/colab_cli/commands/__init__.py

@@ -0,0 +1,14 @@
+# 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.
+

package/core/app/colab_cli/commands/automation.py

@@ -0,0 +1,265 @@
+# 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 os
+import sys
+import json
+from typing import Optional, List
+import typer
+from typing_extensions import Annotated
+
+from app.colab_cli.runtime import ColabRuntime
+from app.colab_cli.contents import ContentsClient
+from app.colab_cli.auth import get_credentials
+from app.colab_cli.utils import get_status_code
+
+
+# Default execute() timeout for human-in-the-loop automations (auth /
+# drivemount). The kernel goes silent while the user completes a browser
+# OAuth flow, which can routinely take 30s+; the upstream 10s default
+# raises ``TimeoutError`` mid-flow even though the mount actually succeeds.
+# 10 minutes is long enough for any realistic interactive auth ceremony
+# without leaving CI hangs unbounded.
+INTERACTIVE_AUTOMATION_TIMEOUT_SEC = 600
+
+
+def run_automation(
+    name: str,
+    op: str,
+    code: str,
+    allow_stdin: bool = False,
+    path: str = None,
+    timeout: Optional[float] = None,
+):
+    from app.colab_cli.common import state
+
+    s = state.store.get(name)
+    runtime = ColabRuntime(s.url, s.token, session_name=s.name, history=state.history)
+
+    def drivefs_hook(deserialize_msg, wsclient):
+        content = deserialize_msg.get("content", {})
+        if content.get("request", {}).get("authType") == "dfs_ephemeral":
+            msg_id = deserialize_msg.get("metadata", {}).get("colab_msg_id")
+            state.history.log_event(
+                s.name,
+                "colab_request",
+                {"type": "dfs_ephemeral", "colab_msg_id": msg_id},
+            )
+            url = f"{state.client.colab_domain}/tun/m/credentials-propagation/{s.endpoint}"
+            params = {
+                "authuser": "0",
+                "authtype": "dfs_ephemeral",
+                "version": "2",
+                "dryrun": "true",
+                "propagate": "true",
+                "record": "false",
+            }
+            typer.echo(
+                f"\n[colab] Intercepted Drive Auth Request. Connecting to {state.client.colab_domain}..."
+            )
+
+            creds = get_credentials(
+                state.client_oauth_config, provider=state.auth_provider
+            )
+            resp = creds.request("GET", url, params=params)
+            token = (
+                json.loads(resp.text.split("\n", 1)[-1]).get("token")
+                if get_status_code(resp) == 200
+                else None
+            )
+
+            headers = {"x-goog-colab-token": token}
+            resp = creds.request(
+                "POST",
+                url,
+                params=params,
+                headers=headers,
+                files={"file_id": (None, "empty.ipynb")},
+            )
+            data = json.loads(resp.text.split("\n", 1)[-1])
+
+            if not data.get("success"):
+                uri = data.get("unauthorized_redirect_uri")
+                typer.echo(
+                    f"\n[colab] REQUIRED: Google Drive Authorization needed.\nPlease visit:\n\n{uri}\n"
+                )
+                state.history.log_event(s.name, "drive_auth_needed", {"uri": uri})
+                sys.stdout.write("Press Enter after you have granted access... ")
+                sys.stdout.flush()
+                with open("/dev/tty") as tty:
+                    tty.readline()
+
+            typer.echo("[colab] Authorizing VM...")
+            params["dryrun"] = "false"
+            resp = creds.request(
+                "POST",
+                url,
+                params=params,
+                headers=headers,
+                files={"file_id": (None, "empty.ipynb")},
+            )
+            if get_status_code(resp) == 200:
+                typer.echo("[colab] Credentials propagated. Resuming mount...")
+                state.history.log_event(s.name, "drive_auth_success", {})
+                reply = wsclient.session.msg(
+                    "input_reply",
+                    {"value": {"type": "colab_reply", "colab_msg_id": msg_id}},
+                )
+                if "header" in deserialize_msg:
+                    reply["parent_header"] = deserialize_msg["header"]
+                wsclient.stdin_channel.send(reply)
+            else:
+                typer.echo(
+                    f"[colab] Error propagating: {get_status_code(resp)} {resp.text}"
+                )
+            return True
+        return False
+
+    runtime.colab_request_hook = drivefs_hook
+    try:
+        s.running = f"automation({op})"
+        s.last_execution = (
+            f"automation:{op}",
+            None,
+            datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+        )
+        state.store.add(s)
+
+        if op == "drivemount":
+            state.history.log_event(
+                name, "automation", {"op": "drivemount", "path": path, "code": code}
+            )
+        else:
+            state.history.log_event(name, "automation", {"op": op, "code": code})
+
+        outputs = runtime.execute_code(code, allow_stdin=allow_stdin, timeout=timeout)
+        state.history.log_event(
+            name, "automation_result", {"op": op, "outputs": outputs}
+        )
+
+        for out in outputs:
+            if "text" in out:
+                sys.stdout.write(out["text"])
+            elif "data" in out:
+                if "text/plain" in out["data"]:
+                    typer.echo(out["data"]["text/plain"])
+            elif out.get("output_type") == "error":
+                ename = out.get("ename", "Error")
+                evalue = out.get("evalue", "")
+                tb = out.get("traceback", [])
+                if tb:
+                    sys.stderr.write("".join(tb) + "\n")
+                else:
+                    sys.stderr.write(f"{ename}: {evalue}\n")
+    finally:
+        s.running = None
+        state.store.add(s)
+        runtime.stop()
+
+
+def auth(
+    session: Annotated[
+        Optional[str], typer.Option("-s", "--session", help="Session name")
+    ] = None,
+):
+    """Authenticate with Google on the VM"""
+    from app.colab_cli.common import state
+
+    name = state.resolve_session(session)
+    code = "import os\nos.environ['USE_AUTH_EPHEM'] = '0'\nfrom google.colab import auth\nauth.authenticate_user()"
+    typer.echo(f"[colab] Starting Google Auth flow on {name}...")
+    run_automation(
+        name,
+        "auth",
+        code,
+        allow_stdin=True,
+        timeout=INTERACTIVE_AUTOMATION_TIMEOUT_SEC,
+    )
+
+
+def drivemount(
+    session: Annotated[
+        Optional[str], typer.Option("-s", "--session", help="Session name")
+    ] = None,
+    path: Annotated[str, typer.Argument(help="Mount path")] = "/content/drive",
+):
+    """Mount Google Drive at path"""
+    from app.colab_cli.common import state
+
+    name = state.resolve_session(session)
+    code = f"from google.colab import drive\ndrive.mount('{path}')"
+    typer.echo(f"[colab] Mounting Google Drive to '{path}' on {name}...")
+    run_automation(
+        name,
+        "drivemount",
+        code,
+        allow_stdin=True,
+        path=path,
+        timeout=INTERACTIVE_AUTOMATION_TIMEOUT_SEC,
+    )
+
+
+def install(
+    session: Annotated[
+        Optional[str], typer.Option("-s", "--session", help="Session name")
+    ] = None,
+    packages: Annotated[
+        Optional[List[str]], typer.Argument(help="Packages to install")
+    ] = None,
+    requirement: Annotated[
+        Optional[str], typer.Option("-r", "--requirement", help="Requirements file")
+    ] = None,
+):
+    """Install python packages on the VM"""
+    from app.colab_cli.common import state
+
+    name = state.resolve_session(session)
+    if not packages and not requirement:
+        typer.echo("[colab] No packages or requirements specified.")
+        raise typer.Exit(1)
+
+    commands = []
+    if requirement:
+        if not os.path.isfile(requirement):
+            typer.echo(f"[colab] Requirements file '{requirement}' not found locally.")
+            raise typer.Exit(1)
+        contents = ContentsClient(state.store.get(name))
+        remote_path = f"content/{os.path.basename(requirement)}"
+        contents.upload(requirement, remote_path)
+        commands.extend(["-r", f"/{remote_path}"])
+    if packages:
+        commands.extend(packages)
+
+    cmd_str = ", ".join(f"'{c}'" for c in commands)
+    code = f"""
+import subprocess, sys
+def install():
+    packages = [{cmd_str}]
+    try:
+        subprocess.check_call(['uv', 'pip', 'install', '--system'] + packages)
+        print('Installation Complete (via uv)!')
+    except:
+        subprocess.check_call([sys.executable, '-m', 'pip', 'install'] + packages)
+        print('Installation Complete (via pip)!')
+install()
+"""
+    typer.echo(f"[colab] Installing packages on {name} (preferring uv)...")
+    run_automation(name, "install", code)
+
+
+def register(app: typer.Typer):
+    app.command(hidden=True)(auth)
+    app.command()(drivemount)
+    app.command()(install)