diff --git a/priv/python/albedo_plugins/colab.py b/priv/python/albedo_plugins/colab.py
index 6d46e278..b0e693eb 100644
--- a/priv/python/albedo_plugins/colab.py
+++ b/priv/python/albedo_plugins/colab.py
@@ -9,7 +9,11 @@ Colab exposes, run as ssh's ProxyCommand by `albedo_plugins.ws` (see `connect`).
from __future__ import annotations
import asyncio
+import base64
+import http.client
import json
+import re
+import secrets
import shlex
import sys
import urllib.error
@@ -21,7 +25,7 @@ from typing import Any
import albedo_ssh
from albedo_api import PythonApi, Record
-from albedo_plugins.remote import remote
+from albedo_plugins.remote import RemoteError, remote
API_BASE = "https://colab.research.google.com"
CONNECT_TIMEOUT = 60.0
@@ -43,6 +47,10 @@ class ColabError(RuntimeError):
self.status = status
+class ColabSshSessionError(ColabError):
+ """The runtime refused ssh because its single session slot is occupied."""
+
+
class ColabInstance(Record):
"""One Google Colab runtime instance."""
@@ -390,13 +398,178 @@ async def connect(
"ALBEDO_WS_HEADERS": json.dumps({"X-Colab-Ssh-Pubkey": public}),
},
)
- return await remote.connect(
- host,
- remote_cwd=remote_cwd or "/content",
- python=python or "/usr/bin/python3",
- modules=modules,
- timeout=timeout,
+ try:
+ return await remote.connect(
+ host,
+ remote_cwd=remote_cwd or "/content",
+ python=python or "/usr/bin/python3",
+ modules=modules,
+ timeout=timeout,
+ )
+ except RemoteError as error:
+ try:
+ occupied = await asyncio.to_thread(
+ _ssh_session_occupied, socket_url, public, min(timeout, 10.0)
+ )
+ except (OSError, http.client.HTTPException):
+ occupied = False
+ if occupied:
+ raise ColabSshSessionError(
+ "Colab allows one ssh session per runtime; a previously held session "
+ "was not released. This is server-side and does not clear on retry. "
+ "Use the HTTP file routes to keep working: colab.contents(), "
+ "colab.read(), colab.pull() (and colab.runtime_status() for liveness). "
+ "stop()/unassign is the only forceful reset and destroys the VM.",
+ 429,
+ ) from error
+ raise
+
+
+def _ssh_session_occupied(url: str, public: str, timeout: float) -> bool:
+ target = urllib.parse.urlsplit(url)
+ transport = (
+ http.client.HTTPSConnection
+ if target.scheme == "wss"
+ else http.client.HTTPConnection
+ )
+ connection = transport(target.netloc, timeout=timeout)
+ try:
+ connection.request(
+ "GET",
+ f"{target.path}?{target.query}",
+ headers={
+ "Upgrade": "websocket",
+ "Connection": "Upgrade",
+ "Sec-WebSocket-Key": base64.b64encode(secrets.token_bytes(16)).decode(),
+ "Sec-WebSocket-Version": "13",
+ "X-Colab-Ssh-Pubkey": public,
+ },
+ )
+ response = connection.getresponse()
+ return response.status == 429 and b"already-active-session" in response.read(
+ 4096
+ )
+ finally:
+ connection.close()
+
+
+async def _runtime_url(instance: str | ColabInstance, path: str, **query: str) -> str:
+ live = _find(await instances(), instance)
+ if not (live.url and live.token):
+ raise ColabError(f"Colab instance {live.id} has no runtime proxy yet")
+ target = urllib.parse.urlsplit(live.url)
+ return urllib.parse.urlunsplit(
+ (
+ target.scheme,
+ target.netloc,
+ path,
+ urllib.parse.urlencode({"colab-runtime-proxy-token": live.token, **query}),
+ "",
+ )
+ )
+
+
+def _runtime_open(url: str, headers: dict[str, str] | None = None) -> Any:
+ request = urllib.request.Request(
+ url, headers={"X-Colab-Client-Agent": "colab-cli", **(headers or {})}
+ )
+ try:
+ return urllib.request.urlopen(request, timeout=API_TIMEOUT)
+ except urllib.error.HTTPError as error:
+ status = error.code
+ error.close()
+ raise ColabError(
+ f"Colab runtime HTTP request failed ({status})", status
+ ) from None
+ except urllib.error.URLError:
+ raise ColabError(
+ "network error connecting to the Colab runtime proxy"
+ ) from None
+
+
+def _runtime_json(url: str) -> Any:
+ with _runtime_open(url) as response:
+ return json.load(response)
+
+
+async def runtime_status(instance: str | ColabInstance) -> Record:
+ """Jupyter liveness over HTTPS, independent of the ssh session slot."""
+ url = await _runtime_url(instance, "/api/status")
+ return Record(await asyncio.to_thread(_runtime_json, url))
+
+
+async def contents(
+ instance: str | ColabInstance, path: str = "", *, content: bool = True
+) -> Record:
+ """Jupyter metadata, directory entries, or text-file content over HTTPS.
+
+ Paths are relative to Jupyter's root, not the ssh working directory.
+ Use content=False for file metadata without loading the file.
+ """
+ route = "/api/contents/" + urllib.parse.quote(path.lstrip("/"), safe="/")
+ url = await _runtime_url(
+ instance, route, content="1" if content else "0", format="text"
)
+ return Record(await asyncio.to_thread(_runtime_json, url))
+
+
+async def read(instance: str | ColabInstance, path: str) -> str:
+ """Read one text file over HTTPS; use pull() for large or binary files."""
+ result = await contents(instance, path)
+ if result.get("type") != "file" or not isinstance(result.get("content"), str):
+ raise ColabError(f"not a text file: {path!r}")
+ return result["content"]
+
+
+def _pull(url: str, local: Path, resume: bool) -> Path:
+ have = local.stat().st_size if resume and local.exists() else 0
+ headers = {"Range": f"bytes={have}-"} if have else {}
+ with _runtime_open(url, headers) as response:
+ total = None
+ if response.status == 206:
+ match = re.fullmatch(
+ r"bytes (\d+)-(\d+)/(\d+)", response.headers.get("Content-Range", "")
+ )
+ if not match or int(match[1]) != have or int(match[2]) + 1 != int(match[3]):
+ raise ColabError(
+ "invalid Content-Range from Colab; local file left unchanged"
+ )
+ total = int(match[3])
+ elif response.status == 200:
+ # An ignored Range must replace the partial file, not append to it.
+ have = 0
+ length = response.headers.get("Content-Length")
+ total = int(length) if length is not None else None
+ else:
+ raise ColabError(f"unexpected Colab download response ({response.status})")
+ local.parent.mkdir(parents=True, exist_ok=True)
+ with local.open("ab" if have else "wb") as output:
+ while chunk := response.read(1 << 20):
+ output.write(chunk)
+ have += len(chunk)
+ if total is not None and have != total:
+ raise ColabError(
+ f"incomplete Colab download: got {have} of {total} bytes; retry to resume"
+ )
+ return local
+
+
+async def pull(
+ instance: str | ColabInstance,
+ remote: str,
+ local: str | Path,
+ *,
+ resume: bool = True,
+) -> Path:
+ """Stream a Jupyter-root-relative file to disk over HTTPS, without ssh.
+
+ Partial downloads remain on disk for retry. Resume assumes the remote file
+ has not changed; use resume=False for mutable files. A 416 response is an
+ error, not proof that the local file is complete.
+ """
+ route = "/files/" + urllib.parse.quote(remote.lstrip("/"), safe="/")
+ url = await _runtime_url(instance, route)
+ return await asyncio.to_thread(_pull, url, Path(local).expanduser(), resume)
class Colab:
@@ -407,6 +580,10 @@ class Colab:
launch = staticmethod(launch)
stop = staticmethod(stop)
connect = staticmethod(connect)
+ runtime_status = staticmethod(runtime_status)
+ contents = staticmethod(contents)
+ read = staticmethod(read)
+ pull = staticmethod(pull)
colab = Colab()
@@ -419,6 +596,7 @@ def setup(api: PythonApi) -> dict[str, object]:
return {
"colab": colab,
"ColabError": ColabError,
+ "ColabSshSessionError": ColabSshSessionError,
"ColabInstance": ColabInstance,
"ColabStatus": ColabStatus,
}
diff --git a/robot-docs/colab.md b/robot-docs/colab.md
index 39b4ec5b..269b4dee 100644
--- a/robot-docs/colab.md
+++ b/robot-docs/colab.md
@@ -24,6 +24,24 @@ runtimes from the kernel and reaches one through `remote`.
written to `~/.ssh/config`. a missing ssh public key (`~/.ssh/id_ed25519.pub`,
`id_ecdsa.pub`, or `ssh_key=`) is an error before anything connects.
`stop()` drops the alias; it does not close a connection already open on it.
+- **occupied ssh slot**: only after `remote.connect()` fails, a bounded websocket
+ handshake probe checks for HTTP 429 with `already-active-session`. this raises
+ `ColabSshSessionError` (a `ColabError`, `status=429`). unrelated refusals and
+ probe failures preserve the original `RemoteError`; successful connects cost
+ no extra request. a previously held server-side session may survive a bridge
+ restart. retrying does not release it. `stop()`/unassign is the only forceful
+ reset and **destroys the VM**; never do that to recover files from a live job.
+- **HTTP recovery**: `await colab.runtime_status(vm)` reads Jupyter `/api/status`;
+ `await colab.contents(vm, "content")` lists a directory;
+ `await colab.contents(vm, path, content=False)` reads metadata (including size);
+ `await colab.read(vm, path)` returns a text file in memory;
+ `await colab.pull(vm, path, local, resume=True)` streams `/files/<path>` to disk
+ in 1 MiB chunks and returns its local `Path`. these read-only routes refresh
+ the runtime proxy token and work independently of ssh. paths are relative to
+ **Jupyter's root**, not `/content`. partial downloads remain for retry; resume
+ assumes unchanged remote content, so use `resume=False` for mutable files.
+ an ignored Range restarts rather than appending; invalid ranges and HTTP 416
+ fail instead of claiming completeness. runtime HTTP errors omit credential URLs.
- **tests**: `test/harness/colab_plugin_test.py` (the plugin against a faked
HTTP layer), `ws_test.py` (frames, and the bridge in a real
subprocess) and `google_oauth_test.gleam` (the refresh against a local token
diff --git a/src/albedo/harness/extensions/colab/extension.gleam b/src/albedo/harness/extensions/colab/extension.gleam
index 01806b65..db56efa4 100644
--- a/src/albedo/harness/extensions/colab/extension.gleam
+++ b/src/albedo/harness/extensions/colab/extension.gleam
@@ -30,7 +30,13 @@ const instructions =
<> " and await colab.status() checks credentials."
<> " rem = await colab.connect(vm) (or await vm.connect()) returns what remote.connect() does,"
<> " for the VM in /content over Colab's ssh proxy; it needs an ssh public key in ~/.ssh or"
- <> " ssh_key=."
+ <> " ssh_key=. ColabSshSessionError (a ColabError, status=429) means the runtime's single ssh"
+ <> " session is still held server-side; retrying does not release it. HTTP works independently:"
+ <> " await colab.runtime_status(vm) checks Jupyter liveness; await colab.contents(vm, path)"
+ <> " lists a directory (content=False gives metadata); await colab.read(vm, path) reads text;"
+ <> " await colab.pull(vm, remote, local, resume=True) streams a file to disk in 1 MiB chunks."
+ <> " Paths are relative to Jupyter's root, not /content; resume assumes unchanged content."
+ <> " stop()/unassign forcefully resets ssh but destroys the VM; use HTTP to recover files first."
<> " Always call await vm.stop() (or await colab.stop(vm)) when done"
<> " to release the runtime and preserve compute units. Failures raise ColabError."
diff --git a/test/harness/colab_plugin_test.py b/test/harness/colab_plugin_test.py
index f6d4038d..02e459a6 100644
--- a/test/harness/colab_plugin_test.py
+++ b/test/harness/colab_plugin_test.py
@@ -1,13 +1,21 @@
-"""Unit tests for the Google Colab tool plugin (albedo_plugins.colab)."""
+"""Colab protocol tests: local refusals and interrupted downloads need no live GPU.
+
+The daemon cannot deterministically reproduce an occupied Colab ssh slot or a
+proxy that ignores Range; these tests exercise those responses locally.
+"""
from __future__ import annotations
import asyncio
+from contextlib import contextmanager
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import json
import os
from pathlib import Path
import sys
import tempfile
+import threading
+import urllib.parse
import unittest
from unittest.mock import AsyncMock, patch
@@ -46,6 +54,35 @@ def instance(**fields) -> colab.ColabInstance:
)
+@contextmanager
+def runtime_proxy(status, body, headers=None):
+ requests = []
+
+ class Handler(BaseHTTPRequestHandler):
+ def do_GET(self):
+ requests.append((self.path, self.headers))
+ self.send_response(status)
+ for key, value in (headers or {}).items():
+ self.send_header(key, value)
+ self.end_headers()
+ self.wfile.write(body)
+
+ def log_message(self, format, *args):
+ pass
+
+ server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
+ thread = threading.Thread(
+ target=server.serve_forever, kwargs={"poll_interval": 0.01}, daemon=True
+ )
+ thread.start()
+ try:
+ yield f"http://127.0.0.1:{server.server_port}", requests
+ finally:
+ server.shutdown()
+ server.server_close()
+ thread.join()
+
+
class ColabPluginTest(unittest.IsolatedAsyncioTestCase):
def setUp(self):
self.home = tempfile.mkdtemp(prefix="albedo-colab-home-")
@@ -289,6 +326,168 @@ class ColabPluginTest(unittest.IsolatedAsyncioTestCase):
self.assertIs(colab.HostErrorType, RuntimeError)
loop.close()
self.assertEqual(
- sorted(exports), ["ColabError", "ColabInstance", "ColabStatus", "colab"]
+ sorted(exports),
+ [
+ "ColabError",
+ "ColabInstance",
+ "ColabSshSessionError",
+ "ColabStatus",
+ "colab",
+ ],
)
self.assertIs(exports["colab"], colab.colab)
+
+ async def test_failed_connect_classifies_only_the_specific_refusal(self):
+ (Path(self.home) / ".ssh" / "id_ed25519.pub").write_text(
+ "ssh-ed25519 AAAA test"
+ )
+ for status, body, occupied in (
+ (429, b'{"error":"already-active-session"}', True),
+ (429, b'{"error":"rate-limit"}', False),
+ (502, b"already-active-session", False),
+ (101, b"", False),
+ ):
+ with (
+ self.subTest(status=status, body=body),
+ runtime_proxy(status, body) as (url, requests),
+ ):
+ live = instance(url=url, token="secret/never-log")
+ original = colab.RemoteError("original failure")
+ with (
+ patch.object(colab, "instances", AsyncMock(return_value=[live])),
+ patch.object(
+ colab.remote, "connect", AsyncMock(side_effect=original)
+ ),
+ ):
+ with self.assertRaises(
+ colab.ColabSshSessionError if occupied else colab.RemoteError
+ ) as caught:
+ await colab.connect(live)
+ if occupied:
+ assert isinstance(caught.exception, colab.ColabSshSessionError)
+ self.assertEqual(caught.exception.status, 429)
+ self.assertNotIn(live.token, str(caught.exception))
+ self.assertIn("destroys the VM", str(caught.exception))
+ else:
+ self.assertIs(caught.exception, original)
+ path, headers = requests[0]
+ self.assertIn("/colab/ssh?", path)
+ self.assertEqual(headers["Upgrade"], "websocket")
+ self.assertEqual(headers["X-Colab-Ssh-Pubkey"], "ssh-ed25519 AAAA test")
+ self.assertEqual(headers["Sec-WebSocket-Version"], "13")
+
+ async def test_probe_network_failure_preserves_original_error(self):
+ (Path(self.home) / ".ssh" / "id_ed25519.pub").write_text(
+ "ssh-ed25519 AAAA test"
+ )
+ live = instance(url="https://vm", token="secret")
+ original = colab.RemoteError("original")
+ with (
+ patch.object(colab, "instances", AsyncMock(return_value=[live])),
+ patch.object(colab.remote, "connect", AsyncMock(side_effect=original)),
+ patch.object(
+ colab, "_ssh_session_occupied", side_effect=OSError("probe failure")
+ ),
+ ):
+ with self.assertRaises(colab.RemoteError) as caught:
+ await colab.connect(live)
+ self.assertIs(caught.exception, original)
+
+ async def test_success_does_not_probe(self):
+ (Path(self.home) / ".ssh" / "id_ed25519.pub").write_text(
+ "ssh-ed25519 AAAA test"
+ )
+ live = instance(url="https://vm", token="secret")
+ with (
+ patch.object(colab, "instances", AsyncMock(return_value=[live])),
+ patch.object(colab.remote, "connect", AsyncMock(return_value="connection")),
+ patch.object(colab, "_ssh_session_occupied") as probe,
+ ):
+ self.assertEqual(await colab.connect(live), "connection")
+ probe.assert_not_called()
+
+ async def test_runtime_helpers_refresh_and_quote_paths_without_ssh(self):
+ response = {"type": "file", "content": "hello"}
+ for name, path, route in (
+ ("status", "", "/api/status"),
+ ("contents", "a ?#%/b", "/api/contents/a%20%3F%23%25/b"),
+ ("read", "text", "/api/contents/text"),
+ ):
+ with runtime_proxy(200, json.dumps(response).encode()) as (url, requests):
+ live = instance(url=url, token="secret")
+ with (
+ patch.object(
+ colab, "instances", AsyncMock(return_value=[live])
+ ) as listing,
+ patch.object(colab.remote, "connect", AsyncMock()) as ssh,
+ ):
+ if name == "status":
+ result = await colab.runtime_status(live)
+ elif name == "contents":
+ result = await colab.contents(live, path)
+ else:
+ result = await colab.read(live, path)
+ self.assertEqual(result, "hello" if name == "read" else response)
+ listing.assert_awaited_once()
+ ssh.assert_not_called()
+ self.assertEqual(urllib.parse.urlsplit(requests[0][0]).path, route)
+
+ async def test_pull_resume_and_ignored_range(self):
+ for status, body, headers in (
+ (206, b"def", {"Content-Range": "bytes 3-5/6"}),
+ (200, b"abcdef", {"Content-Length": "6"}),
+ ):
+ with (
+ self.subTest(status=status),
+ runtime_proxy(status, body, headers) as (url, requests),
+ ):
+ live = instance(url=url, token="secret")
+ local = Path(self.home) / "download"
+ local.write_bytes(b"abc")
+ with patch.object(colab, "instances", AsyncMock(return_value=[live])):
+ result = await colab.pull(live, "remote ?#", local)
+ self.assertEqual(result, local)
+ self.assertEqual(local.read_bytes(), b"abcdef")
+ self.assertEqual(requests[0][1]["Range"], "bytes=3-")
+ self.assertIn("/files/remote%20%3F%23?", requests[0][0])
+
+ async def test_pull_rejects_bad_ranges_and_http_errors_without_overwriting(self):
+ for status, headers in (
+ (206, {"Content-Range": "bytes 0-2/3"}),
+ (206, {"Content-Range": "garbage"}),
+ (416, {}),
+ (403, {}),
+ ):
+ with (
+ self.subTest(status=status),
+ runtime_proxy(status, b"xyz", headers) as (url, _),
+ ):
+ local = Path(self.home) / "download"
+ local.write_bytes(b"abc")
+ with self.assertRaises(colab.ColabError) as caught:
+ await asyncio.to_thread(
+ colab._pull, url + "?token=secret", local, True
+ )
+ self.assertEqual(local.read_bytes(), b"abc")
+ self.assertNotIn("secret", str(caught.exception))
+
+ async def test_pull_detects_truncation_and_keeps_partial_file(self):
+ with runtime_proxy(206, b"d", {"Content-Range": "bytes 3-5/6"}) as (url, _):
+ local = Path(self.home) / "download"
+ local.write_bytes(b"abc")
+ with self.assertRaisesRegex(colab.ColabError, "incomplete"):
+ await asyncio.to_thread(colab._pull, url, local, True)
+ self.assertEqual(local.read_bytes(), b"abcd")
+
+ async def test_pull_fresh_multichunk_file_replaces_local_when_resume_disabled(self):
+ body = b"x" * ((2 << 20) + 17)
+ with runtime_proxy(200, body, {"Content-Length": str(len(body))}) as (
+ url,
+ requests,
+ ):
+ local = Path(self.home) / "nested" / "download"
+ local.parent.mkdir()
+ local.write_bytes(b"old")
+ await asyncio.to_thread(colab._pull, url, local, False)
+ self.assertEqual(local.read_bytes(), body)
+ self.assertNotIn("Range", requests[0][1])