Skip to content

fix(cli): create SSH keys locally and register only the public key - #231

Merged
k223kim merged 4 commits into
mainfrom
k223kim/ssh-local-key
Nov 4, 2026
Merged

k223kim merged 4 commits into
mainfrom
k223kim/ssh-local-key

Conversation

@k223kim

@k223kim k223kim commented Nov 4, 2026 •

Copy link
Copy Markdown
Contributor

The CLI now sets up SSH the same way the web UI's setup script does: the key is created on the user's machine with ssh-keygen, and its public key is registered with CreateSSHPublicKey (POST /v1/ssh-keys) through the SDK's generated client, which sends the lightning login token as Authorization: Bearer.

What changed

All three commands now share ensure_ssh_key() in cli/utils/ssh_connection.py:

  1. If ~/.ssh/lightning_ed25519 is missing, it runs ssh-keygen -t ed25519 -N "" -C lightning@<hostname>. If only the .pub file is missing, it recreates it with ssh-keygen -y. Once ~/.ssh/lightning_ed25519 is there, it's never replaced.
  2. It calls CreateSSHPublicKey on every run, naming the key after its comment (lightning@<hostname> by default), as the web UI does. The CLI doesn't check the account's keys first: the server returns a key the account already has, unchanged and without a new row, even at the key cap, and compares by fingerprint. The CLI prints nothing about registration, because a new key and a repeat look the same. The create isn't retried, because once the client's retries run out it raises a bare Exception that drops the server's message. If the server refuses, the CLI shows its message as is, for example the per-user key limit.

Per command:

  • lightning ssh configure and the legacy lightning configure ssh run the steps above, then add a Host <studio name> entry to ~/.ssh/config. They skip this when a Host or Match block for that name already has .ssh/lightning_ed25519 as an IdentityFile. That check is a port of the web UI's awk check. Before, the check was a substring match on Host <name>, which also matched Host <name>-2 and left an entry pointing at lightning_rsa alone. --overwrite no longer replaces anything. It prints a note saying to delete the key to make a new one, and its help text says it has no effect.
  • lightning studio ssh runs the steps above and then ssh -i ~/.ssh/lightning_ed25519 s_<id>@ssh.lightning.ai. Before, it used a different key, ~/.lightning/lightning_rsa. The "retry with fresh keys" branch is removed because there's nothing left to refresh. If ssh itself is missing, the command now raises a ClickException instead of a RuntimeError. -o values are passed as single arguments, so -o "StrictHostKeyChecking no" now works.
  • The API key check (auth.api_key or auth.load()) is removed from all three commands. It refused users signed in with only the CLI login token. require_auth_header() still stops them before any key work when there are no credentials.
  • lightning ssh generate and lightning generate ssh print the same entry, which now points at lightning_ed25519.
  • lightning job ssh, lightning mmt ssh and lightning studio connect use the same helper, so they get the new key too. In job ssh and mmt ssh the only change is that the force_download retry is gone. Jobs are out of scope here, and their connect path wasn't checked.

Matching the web UI's setup script

The CLI uses the same key file (~/.ssh/lightning_ed25519), comment (lightning@<hostname>), default name, and ~/.ssh/config lines, indented by two spaces. This means the browser flow and the CLI share one key on each machine. The entry checks the gateway's host key, as the script's does: StrictHostKeyChecking accept-new with UserKnownHostsFile ~/.ssh/lightning_known_hosts, so ssh records the key on the first connection and refuses a changed one after that. There are two differences:

  • No passphrase. The script lets ssh-keygen ask for one. The CLI passes -N "" so that lightning studio ssh works without a terminal, as the old lightning_rsa key did. A key that the script already made with a passphrase is used as is. If its .pub file is missing, ssh-keygen -y asks for the passphrase on the terminal.
  • The host alias. The script writes Host ssh.lightning.ai. The CLI writes Host <studio name>, with User s_<id> and Hostname ssh.lightning.ai, as it did before, so that ssh <studio name> keeps working. The port is always 22, so there is no Port line.

Old lightning_rsa entries

A machine set up before this change may still have a Host <studio name> block whose only IdentityFile is lightning_rsa. That's why lightning ssh configure adds the new entry after it instead of treating the host as set up. ssh uses every IdentityFile from every block that matches, even with IdentitiesOnly yes, so it offers both keys. I checked this with ssh -G (OpenSSH 9.9), and a test checks it too. The old block comes first, so its User still wins. That's the same Studio unless a Studio name was reused for a different Studio.

How the key is generated

The CLI runs ssh-keygen rather than using a Python library:

  • The SDK has no crypto dependency. Writing an OpenSSH private key in Python would mean adding cryptography for every SDK user.
  • Anyone who can use the key needs ssh, and ssh-keygen ships with it: OpenSSH on Linux and macOS, and the OpenSSH client that Windows 10 (1809 and later) and 11 install by default. subprocess passes -N "" to the Windows ssh-keygen as an empty argument. If ssh-keygen isn't found, the CLI says "ssh-keygen was not found. Install OpenSSH and run this again." and registers nothing.
  • It's the same tool the web UI's script uses, so both flows produce the same kind of key.
  • stdin is closed when generating the key, so ssh-keygen can't prompt. If another process creates the key between the existence check and the ssh-keygen call, ssh-keygen refuses to overwrite it.

Testing

Unit tests: what each new test checks, and what was run
  • New tests in tests/api/test_ssh_key_api.py call the real generated client, mocked only at urllib3. They check that the create is POST /v1/ssh-keys with {name, publicKey}. They check that the create sends Authorization: Bearer <login token>, that a refusal carries the server's message, and that a 500 on create isn't retried.
  • New tests in tests/cli/utils/test_ssh_connection.py use the real ssh-keygen in a temporary home directory. They check that the key is ed25519 with no passphrase, mode 0600 and comment lightning@<hostname>. They also cover these cases:
    • A second run registers the same key again and leaves the private key alone.
    • A key file already in ~/.ssh, such as one the web UI's script made, is registered without being replaced.
    • A missing .pub file is recreated.
    • The refusal is shown.
    • A missing ssh-keygen is explained.
    • Nothing runs without login.
  • tests/cli/ssh/test_configure.py covers these cases:
    • Both configure commands create, register and write the exact entry.
    • Running a command again, even with --overwrite, keeps the key, registers it again and leaves the config alone.
    • The refusal is shown, and no config is written.
    • ssh -G reads a new entry as StrictHostKeyChecking accept-new with the Lightning known-hosts file.
    • The entry is added after an old lightning_rsa-only entry, and ssh -G then offers both keys.
    • The Host/Match check works.
  • tests/cli/studio/test_ssh.py checks the exact ssh command, including -o values, and that both runs register the same key. The job and MMT retry tests are updated for the removed retry.
  • Run: tests/api/test_ssh_key_api.py, tests/cli/utils/test_ssh_connection.py, tests/cli/ssh/, tests/cli/studio/test_ssh.py, tests/cli/studio/test_connect.py, tests/cli/job/test_ssh.py, tests/cli/mmt/test_ssh.py (78 passed), and tests/cli/test_entrypoint.py with tests/cli/test_temporarily_unavailable.py (216 passed). pre-commit run passes on the changed files, including ruff, ruff-format and mypy.

End to end, against a local control plane. The CLI from this branch was run against a local control plane built from the latest platform master, with the SSH gateway on localhost:2222 and a Studio running. The CLI signed in with a local session in LIGHTNING_AUTH_TOKEN, sent as Authorization: Bearer. Every case below passed.

Demo: https://www.loom.com/share/6c229fdb9e5c4a1d91bb03d41110a0ee

Cases and what each one checked

lightning studio ssh and the ~/.ssh/config entry always name ssh.lightning.ai, so each connection added -o HostName=localhost -o Port=2222.

  1. A fresh key, through to a shell. With no ~/.ssh/lightning_ed25519, lightning ssh configure created an ed25519 key with mode 0600 and the comment lightning@<hostname>. The account gained one users.ssh_keys row, holding exactly that public key. lightning studio ssh then opened a shell in the Studio, and the gateway logged the session.
  2. The key is reused. Running configure again printed no "Created SSH key", kept the same fingerprint and added no row.
  3. A key made by the web UI's setup script. A key that ssh-setup.sh made was registered by configure without being replaced: the fingerprint and the file's modification time didn't change. It then connected. A second script-made key was pasted into the web UI's SSH dialog first. Running configure afterwards left it as one row with the same ID.
  4. The config entry checks the host key. ssh ssh-ui-review saved the gateway's host key to ~/.ssh/lightning_known_hosts on the first connection and connected without a prompt on the second. With that line replaced by a different key, ssh refused to connect.
  5. An old lightning_rsa entry. With a Host ssh-ui-review block naming only lightning_rsa in ~/.ssh/config, configure added the new entry after it, and the alias connected.
  6. The server's refusal is shown. With a key the server refused, configure exited with code 1 and printed "Could not add the SSH key … to your Lightning account:" followed by the server's message. ~/.ssh/config was unchanged.
  7. A deleted key. After DELETE /v1/ssh-keys/{id}, plain ssh was refused with Permission denied (publickey). Running lightning studio ssh afterwards registered the key again and connected, because the CLI registers on every run. A key deleted in the web UI therefore comes back on a machine that runs lightning studio ssh again.

Not tested:

  • Windows wasn't run. CI's CLI tests run on Ubuntu and macOS only.
  • A key with a passphrase whose .pub file is missing (the interactive ssh-keygen -y prompt) wasn't run.

— Claude Opus 5.5, on direct instruction; reviewed before posting

@k223kim
k223kim marked this pull request as ready for review November 4, 2026 14:53
Earlier tests in the cli/ suite can leave rich's console in color mode, so the error box's borders arrive wrapped in ANSI codes and split the message.

Co-authored-by: Cursor <cursoragent@cursor.com>
@k223kim
k223kim merged commit 2fcefd4 into main Nov 4, 2026
24 checks passed
@k223kim
k223kim deleted the k223kim/ssh-local-key branch November 4, 2026 16:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants