feat(ssh): accept public key for --identity

docs/06_ssh_access.md

@@ -1,5 +1,6 @@
 ---
 log:
+2026-11-09: `--identity/-i` now accepts either a `.pub` public key or a private key path: `_resolve_pubkey` reads the `.pub` file directly when present before falling back to `_pubkey_from_identity` (`ssh-keygen -y -f`), and `_ssh_base_args` strips `.pub` for `ssh -i` when the sibling private key exists on disk.
 2026-10-29: Added `colab ssh [-s NAME] -- CMD`, which runs one command on an existing session and exits with its status (a PTY only when stdin and stdout are terminals; neither it nor its ProxyCommand child creates a runtime). Its handshake carries `ssh-mode-cmd`, and the update banner is skipped for `ssh` when stdout is not a terminal. Added `tests/test_ssh_cmd.py`.
 2026-10-29: HTTP 429 now shows the server's SSH session limit ("N of N SSH sessions in use") when the body carries one.
 2026-10-29: The ProxyCommand that `colab ssh` starts now keeps a non-default global `--auth` and `--config` (as an absolute path), so `colab --config ALT ssh -s S` bridges S from ALT instead of the default sessions file.
@@ -34,7 +35,7 @@ colab ssh [OPTIONS] [-- CMD [ARGS]...]
 |---|---|---|---|
 | `-s`, `--session` | str | auto | Session to connect to. If omitted, uses your only active session, auto-creates one when you have none (not with `-- CMD`), or errors when you have several. |
 | `--proxy-mode` | bool | False | Act as an OpenSSH `ProxyCommand`-compatible WebSocket↔stdio bridge (reads stdin, writes stdout) for `~/.ssh/config`. Every other flag still applies. |
-| `-i`, `--identity` | str | auto | Private key for the public key sent to Colab. Default: first of `~/.ssh/id_ed25519`, `id_ecdsa`. |
+| `-i`, `--identity` | str | auto | SSH key (`.pub` public key or private key path) for the public key sent to Colab. Default: first of `~/.ssh/id_ed25519.pub`, `id_ecdsa.pub`. |
 | `--gpu` | str | None | GPU accelerator for a runtime this command creates (T4, L4, G4, H100, A100). |
 | `--tpu` | str | None | TPU accelerator for a runtime this command creates (v5e1, v6e1). |
 | `--high-mem` | bool | False | Request high-RAM when this command auto-creates a runtime (ignored when connecting to an existing session). |
@@ -138,8 +139,10 @@ session (never creating one); every word from CMD on goes to CMD, flags too.
 
 ### Unit tests (`tests/test_ssh.py`)
 1. WebSocket URL construction (`wss` for https, `ws` for http; token query param).
-2. Pubkey resolution — `--identity` (via `ssh-keygen -y -f`) and the `~/.ssh`
-   default scan; missing-key and missing-identity exit paths.
+2. Pubkey resolution — `--identity` (reads the matching `.pub` file first,
+   whether passed with or without `.pub`, before falling back to `ssh-keygen -y
+   -f`) and the `~/.ssh` default scan; missing-key and missing-identity exit
+   paths.
 3. The full status→message map (400/401/403/404/429/502/other/none).
 4. Shell quoting for the `ProxyCommand` string.
 5. Session resolution (existing vs missing).

integration/repro_ssh/test.sh

@@ -165,6 +165,20 @@ if ! echo "$E2E_OUT" | grep -q "$MARKER"; then
 fi
 echo "[SUCCESS] Handshake + pubkey-auth + bridge + remote exec all work."
 
+echo "[*] End-to-end: colab ssh -i <key>.pub -- CMD (public key path)."
+PUB_MARKER="ssh-pub-ok-$$-$RANDOM"
+PUB_OUT=$(
+  timeout 120 uv run colab $AUTH_FLAGS --config "$SESSION_FILE" ssh \
+    -s "$SESSION_NAME" -i "$KEY.pub" -- "whoami; echo $PUB_MARKER" </dev/null 2>&1
+)
+PUB_RC=$?
+echo "$PUB_OUT"
+if [ "$PUB_RC" -ne 0 ] || ! echo "$PUB_OUT" | grep -q "$PUB_MARKER"; then
+  echo "[FAILURE] colab ssh -i <key>.pub -- CMD failed (rc=$PUB_RC)."
+  exit 1
+fi
+echo "[SUCCESS] --identity with .pub suffix works end-to-end."
+
 echo "[*] Negative: an RSA key is rejected by the server (non-interactive)."
 ssh-keygen -t rsa -b 2048 -N "" -f "$RSA_KEY" >/dev/null
 RSA_OUT=$(

src/colab_cli/commands/ssh.py

@@ -31,9 +31,10 @@
 and ``--rm`` stops the runtime when you disconnect.
 
 ``--identity/-i`` overrides the default key order (``~/.ssh/id_ed25519`` ->
-``id_ecdsa``); the public key is derived via ``ssh-keygen -y -f`` and sent in
-the ``X-Colab-Ssh-Pubkey`` header. RSA keys are rejected by the server, so
-``id_rsa`` is not auto-selected.
+``id_ecdsa``); it accepts either a public key (``.pub``) or a private key path,
+reading the ``.pub`` file directly when present and falling back to
+``ssh-keygen -y -f`` to populate the ``X-Colab-Ssh-Pubkey`` header. RSA keys are
+rejected by the server, so ``id_rsa`` is not auto-selected.
 """
 
 import contextlib
@@ -102,21 +103,22 @@ def _pubkey_from_identity(identity: str) -> str:
         raise typer.Exit(code=2)
     pubkey = res.stdout.strip()
     if not pubkey:
-        typer.echo(
-            f"[colab] ssh-keygen produced no key for {identity}.", err=True
-        )
+        typer.echo(f"[colab] ssh-keygen produced no key for {identity}.", err=True)
         raise typer.Exit(code=2)
     return pubkey
 
 
 def _resolve_pubkey(identity: Optional[str]) -> str:
     """Returns the public key for the ``X-Colab-Ssh-Pubkey`` header.
 
-    With ``identity``, derives it from that private key; otherwise scans
-    ``~/.ssh`` for the first existing ``id_<type>.pub`` in preference order.
+    With ``identity``, reads the matching ``.pub`` file directly if present
+    (whether ``identity`` was passed with or without ``.pub``), falling back to
+    deriving it from the private key via ``_pubkey_from_identity``; otherwise
+    scans ``~/.ssh`` for the first existing ``id_<type>.pub`` in preference
+    order.
 
     Args:
-      identity: Path to a private key, or None to scan ``~/.ssh``.
+      identity: Path to a public or private key, or None to scan ``~/.ssh``.
 
     Returns:
       The public key text to send verbatim in the header.
@@ -125,7 +127,18 @@ def _resolve_pubkey(identity: Optional[str]) -> str:
       typer.Exit: If no usable key is found (exit code 2).
     """
     if identity:
-        return _pubkey_from_identity(identity)
+        expanded = Path(os.path.expanduser(identity))
+        if expanded.suffix == ".pub":
+            pub_path = expanded
+            priv_path = expanded.with_suffix("")
+        else:
+            priv_path = expanded
+            pub_path = expanded.with_suffix(".pub")
+        if pub_path.exists():
+            if pubkey := pub_path.read_text().strip():
+                return pubkey
+        target = str(priv_path) if priv_path.exists() else identity
+        return _pubkey_from_identity(target)
 
     ssh_dir = Path(os.path.expanduser("~/.ssh"))
     for name in _KEY_TYPES:
@@ -476,7 +489,12 @@ def _ssh_base_args(proxy_command: str, identity: Optional[str]) -> list[str]:
         "LogLevel=ERROR",
     ]
     if identity:
-        args.extend(["-i", os.path.expanduser(identity)])
+        ssh_identity = os.path.expanduser(identity)
+        if ssh_identity.endswith(".pub"):
+            priv_path = ssh_identity.removesuffix(".pub")
+            if os.path.exists(priv_path):
+                ssh_identity = priv_path
+        args.extend(["-i", ssh_identity])
     return args
 
 
@@ -560,9 +578,7 @@ def _select_proxy_session(
     """
     if session and not _session_exists(session):
         with contextlib.redirect_stdout(sys.stderr):
-            return _auto_create_session(
-                gpu, tpu, name=session, high_mem=high_mem
-            ), True
+            return _auto_create_session(gpu, tpu, name=session, high_mem=high_mem), True
     return _resolve_session(session), False
 
 
@@ -720,9 +736,9 @@ def ssh(
             "--identity",
             "-i",
             help=(
-                "SSH private key whose public key is sent in the "
-                "X-Colab-Ssh-Pubkey header (default: first of "
-                "~/.ssh/id_ed25519, id_ecdsa)."
+                "SSH key (public `.pub` or private key path) whose public key "
+                "is sent in the X-Colab-Ssh-Pubkey header (default: first of "
+                "~/.ssh/id_ed25519.pub, id_ecdsa.pub)."
             ),
         ),
     ] = None,

tests/test_ssh.py

@@ -76,23 +76,68 @@ def test_build_ws_url_scheme(url, scheme):
 # --- Pubkey resolution -------------------------------------------------------
 
 
-def test_resolve_pubkey_with_identity_calls_ssh_keygen(mocker, tmp_path):
+@pytest.mark.parametrize(
+    "pass_pub_suffix", [True, False], ids=["with-pub", "without-pub"]
+)
+def test_resolve_pubkey_with_identity_reads_pub_file_first(
+    mocker, tmp_path, pass_pub_suffix
+):
+    """When `<key>.pub` exists, `_resolve_pubkey` reads it directly (whether
+    `--identity` was passed with or without `.pub`) before calling
+    `_pubkey_from_identity`."""
+    key = tmp_path / "id_test"
+    pub = tmp_path / "id_test.pub"
+    key.write_text("(fake private key)")
+    fake_pub = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIDfake user@host"
+    pub.write_text(fake_pub + "\n")
+    spy_from_identity = mocker.spy(ssh_module, "_pubkey_from_identity")
+    mock_run = mocker.patch("subprocess.run")
+
+    arg = str(pub if pass_pub_suffix else key)
+    out = ssh_module._resolve_pubkey(arg)
+    assert out == fake_pub
+    spy_from_identity.assert_not_called()
+    mock_run.assert_not_called()
+
+
+def test_resolve_pubkey_with_pub_only_identity_needs_no_private_key(mocker, tmp_path):
+    """A standalone `.pub` file (no private key on disk) works without
+    invoking `_pubkey_from_identity`."""
+    pub = tmp_path / "agent_only.pub"
+    fake_pub = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIAGENT user@host"
+    pub.write_text(fake_pub + "\n")
+    spy_from_identity = mocker.spy(ssh_module, "_pubkey_from_identity")
+
+    assert ssh_module._resolve_pubkey(str(pub)) == fake_pub
+    spy_from_identity.assert_not_called()
+
+
+@pytest.mark.parametrize("pass_pub_suffix", [False, True], ids=["priv-arg", "pub-arg"])
+def test_resolve_pubkey_with_identity_calls_ssh_keygen_when_no_pub_file(
+    mocker, tmp_path, pass_pub_suffix
+):
+    """When no `.pub` file exists on disk, `_resolve_pubkey` falls back to
+    `_pubkey_from_identity` (`ssh-keygen -y -f`) using the private key path."""
     key = tmp_path / "id_test"
     key.write_text("(fake private key)")
     fake_pub = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIDfake user@host"
     mock_run = mocker.patch(
         "subprocess.run",
         return_value=MagicMock(stdout=fake_pub + "\n", returncode=0),
     )
-    out = ssh_module._resolve_pubkey(str(key))
+    arg = f"{key}.pub" if pass_pub_suffix else str(key)
+    out = ssh_module._resolve_pubkey(arg)
     assert out == fake_pub
     args, _ = mock_run.call_args
     assert args[0][:3] == ["ssh-keygen", "-y", "-f"]
     assert args[0][3] == str(key)
 
 
-def test_resolve_pubkey_missing_identity_exits(tmp_path):
-    missing = tmp_path / "no-such-key"
+@pytest.mark.parametrize(
+    "arg_name", ["no-such-key", "no-such-key.pub"], ids=["priv", "pub"]
+)
+def test_resolve_pubkey_missing_identity_exits(tmp_path, arg_name):
+    missing = tmp_path / arg_name
     with pytest.raises(typer.Exit) as exc_info:
         ssh_module._resolve_pubkey(str(missing))
     assert exc_info.value.exit_code == 2
@@ -146,9 +191,7 @@ def test_resolve_pubkey_default_scan_key_order(
     if present:
         content = f"ssh-key-content-for-{present}\n"
         (ssh_dir / present).write_text(content)
-    monkeypatch.setattr(
-        "os.path.expanduser", lambda p: p.replace("~", str(fake_home))
-    )
+    monkeypatch.setattr("os.path.expanduser", lambda p: p.replace("~", str(fake_home)))
     if expect_found:
         assert ssh_module._resolve_pubkey(None) == content.strip()
     else:
@@ -488,18 +531,14 @@ def test_ssh_proxy_mode_calls_websocket(mock_common_state, mocker):
     connect = mocker.patch.object(
         ssh_module, "_connect_websocket", return_value=fake_ws
     )
-    bridge = mocker.patch.object(
-        ssh_module, "_bridge_proxy_mode", return_value=0
-    )
+    bridge = mocker.patch.object(ssh_module, "_bridge_proxy_mode", return_value=0)
     ssh_subprocess = mocker.patch.object(ssh_module, "_run_interactive_ssh")
 
     result = runner.invoke(app, ["ssh", "--proxy-mode", "-s", "s1"])
     assert result.exit_code == 0
     connect.assert_called_once()
     args, _ = connect.call_args
-    assert args[0].startswith(
-        "wss://abc-foo.colab.googleusercontent.com/colab/ssh"
-    )
+    assert args[0].startswith("wss://abc-foo.colab.googleusercontent.com/colab/ssh")
     assert args[1] == fake_pub
     bridge.assert_called_once_with(fake_ws)
     ssh_subprocess.assert_not_called()
@@ -546,9 +585,7 @@ def fake_connect(url, pubkey):
         captured["url"] = url
         return MagicMock()
 
-    mocker.patch.object(
-        ssh_module, "_connect_websocket", side_effect=fake_connect
-    )
+    mocker.patch.object(ssh_module, "_connect_websocket", side_effect=fake_connect)
     mocker.patch.object(ssh_module, "_bridge_proxy_mode", return_value=0)
 
     runner.invoke(app, ["ssh", "--proxy-mode", "-s", "s1"])
@@ -575,9 +612,7 @@ def test_ssh_handshake_400_emits_actionable_message(
         ssh_module, "_resolve_pubkey", return_value="ssh-rsa AAAAfake u@h"
     )
 
-    err = websocket.WebSocketBadStatusException(
-        "Handshake status 400 Bad Request", 400
-    )
+    err = websocket.WebSocketBadStatusException("Handshake status 400 Bad Request", 400)
     err.status_code = 400
     err.resp_body = resp_body
     mocker.patch.object(websocket.WebSocket, "connect", side_effect=err)

tests/test_ssh_workdir.py

@@ -73,3 +73,22 @@ def test_interactive_ssh_builds_args(mocker, identity):
     else:
         assert "-i" not in args
         assert "--identity" not in joined
+
+
+@pytest.mark.parametrize("priv_exists", [True, False], ids=["priv-on-disk", "pub-only"])
+def test_ssh_base_args_strips_pub_when_sibling_private_key_exists(
+    tmp_path, priv_exists
+):
+    """When `--identity` ends with `.pub` and the sibling private key exists on
+    disk, pass the private key path to `ssh -i` (OpenSSH rejects a `0644` `.pub`
+    file as an unprotected private key unless the key is loaded in ssh-agent);
+    when only the `.pub` file exists, pass the `.pub` path as-is."""
+    priv = tmp_path / "id_ed25519"
+    pub = tmp_path / "id_ed25519.pub"
+    pub.write_text("ssh-ed25519 AAAAfake user@host\n")
+    if priv_exists:
+        priv.write_text("(private key)")
+
+    args = ssh_module._ssh_base_args("proxy", str(pub))
+    idx = args.index("-i")
+    assert args[idx + 1] == str(priv if priv_exists else pub)