Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 37 additions & 18 deletions experimental/ssh/internal/client/agentshim.go
Original file line number Diff line number Diff line change
Expand Up @@ -50,15 +50,18 @@ const (
ugExecutable = "ug"

// pins the version of uv downloaded
uvBaseURL = "https://github.com/astral-sh/uv/releases/download/0.12.18/"
uvVersion = "0.12.18"
uvBaseURL = "https://github.com/astral-sh/uv/releases/download/" + uvVersion + "/"

// SHA-256 of each pinned uv release tarball
uvX64LinuxChecksum = "89eadd7c76fc063887959510d5ba0ab1264dfd5f1143b925ddb73021a40acf16"
uvArm64LinuxChecksum = "afb6291f3f0a6b4521fc67b947822506c41dde5b60d2189dd8f3695b2ac8c9e7"

// overrides for uv tool installs so they are isolated to a directory we control
uvToolDir = agentDepsDir + "/uv_tools"
uvToolBinDir = agentDepsDir + "/uv_tools_bin"
// overrides for uv tool installs so they are isolated to a directory we control.
// Keyed by ugCommit so bumping the pin installs it afresh on a long-lived cluster
// rather than reusing the previous commit's install.
uvToolDir = agentDepsDir + "/uv_tools-" + ugCommit
uvToolBinDir = agentDepsDir + "/uv_tools_bin-" + ugCommit

// pins the version of node/npm downloaded
nodeVersion = "v24.21.0"
Expand Down Expand Up @@ -464,10 +467,17 @@ func bootstrapToolchain(ctx context.Context, home string) (string, error) {
return ugPath, nil
}

func toolchainReady() bool {
_, ugErr := exec.LookPath(ugExecutable)
// toolchainReady reports whether the pinned Unity Gateway CLI and some npm are
// installed. npm (and uv) shipped with the image are used as-is, but ug must be the
// pinned commit, so it is checked at its versioned path rather than looked up on PATH.
func toolchainReady(home string) bool {
_, npmErr := exec.LookPath(npmExecutable)
return ugErr == nil && npmErr == nil
return ugInstalled(home) && npmErr == nil
}

func ugInstalled(home string) bool {
_, err := os.Stat(filepath.Join(home, uvToolBinDir, ugExecutable))
return err == nil
}

func ensureToolchain(ctx context.Context, home, arch string) error {
Expand All @@ -478,11 +488,10 @@ func ensureToolchain(ctx context.Context, home, arch string) error {
uvToolBinDirAbsolute := filepath.Join(home, uvToolBinDir)
prependPath(ctx, uvToolBinDirAbsolute)
for _, spec := range []archiveSpec{uvSpec, nodeSpec} {
binDir := filepath.Join(home, agentDepsDir, spec.dirName, spec.binDirName)
prependPath(ctx, binDir)
prependPath(ctx, spec.binDir(home))
}

if toolchainReady() {
if toolchainReady(home) {
return nil
}
unlock, err := acquireSetupLock(ctx, home)
Expand All @@ -491,7 +500,7 @@ func ensureToolchain(ctx context.Context, home, arch string) error {
}
defer unlock()
// Another client may have finished the install while we waited for the lock.
if toolchainReady() {
if toolchainReady(home) {
return nil
}

Expand All @@ -506,7 +515,7 @@ func ensureToolchain(ctx context.Context, home, arch string) error {
}
}

if _, err := exec.LookPath(ugExecutable); err != nil {
if !ugInstalled(home) {
if err := runStep(ctx, "Installing Unity Gateway CLI", func(out io.Writer) error {
env := os.Environ()
env = append(env, "UV_TOOL_DIR="+uvToolDirAbsolute, "UV_TOOL_BIN_DIR="+uvToolBinDirAbsolute)
Expand Down Expand Up @@ -535,11 +544,22 @@ type archiveSpec struct {
archiveURL string
binaryName string
dirName string
// sub-directory inside dirName where the binary is located, "" if it's the same as dirName
// installed under dirName-version, so a version bump installs afresh on a
// long-lived cluster instead of reusing the previous version
version string
// sub-directory inside the install dir where the binary is located, "" if it's the install dir itself
binDirName string
checksum string
}

func (s archiveSpec) installDirName() string {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Non-blocking: nothing removes the previous install once a bump lands, so every pin change leaves the old node-* (~190 MB extracted), uv-* and uv_tools-<commit> / uv_tools_bin-<commit> dirs behind in ~/.agent-shim/deps, plus the pre-versioning node / uv / uv_tools* dirs. On the long-lived dedicated clusters this is for, that keeps growing until the cluster restarts.

Could be a follow-up: after a successful install, under the setup lock, prune the sibling <dirName>-* / uv_tools*-* dirs that aren't the current version. One thing to watch: a session from an older CLI version still running on the same cluster may be using its own dir, so pruning would have to skip anything in use, or be limited to the unversioned legacy dirs.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah I wanted to add cleanup but I skipped it for now since I didn't have a clear idea how to not cleanup directories from older CLI versions that are still in use

return s.dirName + "-" + s.version
}

func (s archiveSpec) binDir(home string) string {
return filepath.Join(home, agentDepsDir, s.installDirName(), s.binDirName)
}

// ensure the specified binary is installed, installing it if it's not already
// present. Extraction output is written to out. Returns the path to the directory
// containing the installed binary
Expand All @@ -549,9 +569,8 @@ func ensureBinary(ctx context.Context, home string, out io.Writer, spec archiveS
}

depsRoot := filepath.Join(home, agentDepsDir)
depsDir := filepath.Join(depsRoot, spec.dirName)
binDir := filepath.Join(depsDir, spec.binDirName)
if _, err := os.Stat(filepath.Join(binDir, spec.binaryName)); err == nil {
depsDir := filepath.Join(depsRoot, spec.installDirName())
if _, err := os.Stat(filepath.Join(spec.binDir(home), spec.binaryName)); err == nil {
return nil
}

Expand Down Expand Up @@ -588,7 +607,7 @@ func ensureBinary(ctx context.Context, home string, out io.Writer, spec archiveS
}

func uvArchiveSpec(goarch string) archiveSpec {
spec := archiveSpec{dirName: "uv", binaryName: "uv"}
spec := archiveSpec{dirName: "uv", version: uvVersion, binaryName: "uv"}
switch goarch {
case "amd64":
spec.archiveURL = uvBaseURL + "uv-x86_64-unknown-linux-gnu.tar.gz"
Expand Down Expand Up @@ -749,7 +768,7 @@ func downloadVerified(ctx context.Context, spec archiveSpec) (string, error) {
}

func nodeArchiveSpec(goarch string) archiveSpec {
spec := archiveSpec{binDirName: "bin", dirName: "node", binaryName: "npm"}
spec := archiveSpec{binDirName: "bin", dirName: "node", version: nodeVersion, binaryName: "npm"}
switch goarch {
case "amd64":
spec.archiveURL = nodeBaseURL + "node-" + nodeVersion + "-linux-x64.tar.xz"
Expand Down
55 changes: 44 additions & 11 deletions experimental/ssh/internal/client/agentshim_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -190,6 +190,15 @@ func TestEnsureBinary(t *testing.T) {
defer srv.Close()

home := t.TempDir()
// Older installs: an unversioned one from a CLI before versioned installs, and a
// previous version. Neither counts as the pinned version being installed.
depsRoot := filepath.Join(home, agentDepsDir)
for _, dir := range []string{"node", "node-v1"} {
writeFakeExec(t, filepath.Join(depsRoot, dir, "bin"), "npm", ":")
}
for _, dir := range []string{"uv", "uv-1"} {
writeFakeExec(t, filepath.Join(depsRoot, dir), "uv", ":")
}
for _, c := range []struct {
spec archiveSpec
wantBinDir string
Expand All @@ -200,16 +209,18 @@ func TestEnsureBinary(t *testing.T) {
binaryName: "npm",
checksum: hex.EncodeToString(nodeSum[:]),
dirName: "node",
version: "v2",
binDirName: "bin",
}, wantBinDir: filepath.Join(home, agentDepsDir, "node", "bin"),
}, wantBinDir: filepath.Join(home, agentDepsDir, "node-v2", "bin"),
},
{
spec: archiveSpec{
archiveURL: srv.URL + "/uv.tar.gz",
binaryName: "uv",
checksum: hex.EncodeToString(uvSum[:]),
dirName: "uv",
}, wantBinDir: filepath.Join(home, agentDepsDir, "uv"),
version: "2",
}, wantBinDir: filepath.Join(home, agentDepsDir, "uv-2"),
},
} {
t.Run("downloads, extracts, and returns the bin dir ("+c.spec.binaryName+")", func(t *testing.T) {
Expand Down Expand Up @@ -255,8 +266,8 @@ func TestEnsureToolchain(t *testing.T) {
t.Run(fmt.Sprintf("no-op when the toolchain is already present (staleLock=%v)", staleLock), func(t *testing.T) {
home := t.TempDir()
writeFakeExec(t, filepath.Join(home, uvToolBinDir), ugExecutable, ":")
writeFakeExec(t, filepath.Join(home, agentDepsDir, "node", "bin"), "npm", ":")
writeFakeExec(t, filepath.Join(home, agentDepsDir, "uv"), "uv", fakeUvRecorder)
writeFakeExec(t, nodeArchiveSpec("amd64").binDir(home), "npm", ":")
writeFakeExec(t, uvArchiveSpec("amd64").binDir(home), "uv", fakeUvRecorder)
if staleLock {
// left by a session killed mid-install
writeFakeExec(t, filepath.Join(home, agentRootDir), setupLockName, "")
Expand All @@ -274,8 +285,8 @@ func TestEnsureToolchain(t *testing.T) {
home := t.TempDir()
// uv and npm are present; ug is not, so only ug gets installed and no
// download happens
writeFakeExec(t, filepath.Join(home, agentDepsDir, "uv"), "uv", fakeUvRecorder)
writeFakeExec(t, filepath.Join(home, agentDepsDir, "node", "bin"), "npm", ":")
writeFakeExec(t, uvArchiveSpec("amd64").binDir(home), "uv", fakeUvRecorder)
writeFakeExec(t, nodeArchiveSpec("amd64").binDir(home), "npm", ":")
marker := filepath.Join(t.TempDir(), "uv.log")
t.Setenv("FAKE_UV_LOG", marker)
t.Setenv("PATH", t.TempDir())
Expand All @@ -290,11 +301,33 @@ func TestEnsureToolchain(t *testing.T) {
assert.NoFileExists(t, filepath.Join(home, agentRootDir, setupLockName), "the setup lock is released once setup finishes")
})

t.Run("replaces a Unity Gateway CLI from another commit with the pinned one", func(t *testing.T) {
home := t.TempDir()
writeFakeExec(t, uvArchiveSpec("amd64").binDir(home), "uv", fakeUvRecorder)
writeFakeExec(t, nodeArchiveSpec("amd64").binDir(home), "npm", ":")
// Installs from an older pin and from a CLI before installs were versioned. Both
// put a ug on PATH, but neither is the pinned commit.
depsRoot := filepath.Join(home, agentDepsDir)
oldBinDirs := []string{filepath.Join(depsRoot, "uv_tools_bin"), filepath.Join(depsRoot, "uv_tools_bin-oldcommit")}
for _, dir := range oldBinDirs {
writeFakeExec(t, dir, ugExecutable, ":")
}
marker := filepath.Join(t.TempDir(), "uv.log")
t.Setenv("FAKE_UV_LOG", marker)
t.Setenv("PATH", strings.Join(oldBinDirs, string(os.PathListSeparator)))

require.NoError(t, ensureToolchain(cmdio.MockDiscard(t.Context()), home, "amd64"))

out, err := os.ReadFile(marker)
require.NoError(t, err)
assert.Contains(t, string(out), "tool install git+https://github.com/"+ugRepo+"@"+ugCommit)
})

t.Run("recheck after acquiring the lock skips a redundant install", func(t *testing.T) {
home := t.TempDir()
// uv is present and records if it ever installs; it must not be called once a
// concurrent client finishes the toolchain while we wait for the lock.
writeFakeExec(t, filepath.Join(home, agentDepsDir, "uv"), "uv", fakeUvRecorder)
writeFakeExec(t, uvArchiveSpec("amd64").binDir(home), "uv", fakeUvRecorder)
marker := filepath.Join(t.TempDir(), "uv.log")
t.Setenv("FAKE_UV_LOG", marker)
t.Setenv("PATH", t.TempDir())
Expand All @@ -310,7 +343,7 @@ func TestEnsureToolchain(t *testing.T) {
// The other client finishes: publish ug+npm into the PATH dir, then release
// the lock so the waiting ensureToolchain can proceed.
writeFakeExec(t, filepath.Join(home, uvToolBinDir), ugExecutable, ":")
writeFakeExec(t, filepath.Join(home, agentDepsDir, "node", "bin"), "npm", ":")
writeFakeExec(t, nodeArchiveSpec("amd64").binDir(home), "npm", ":")
unlock()

require.NoError(t, <-done)
Expand Down Expand Up @@ -635,7 +668,7 @@ func TestBootstrapToolchain(t *testing.T) {
t.Run("returns the path to the Unity Gateway CLI", func(t *testing.T) {
home := t.TempDir()
writeFakeExec(t, filepath.Join(home, uvToolBinDir), ugExecutable, ":")
writeFakeExec(t, filepath.Join(home, agentDepsDir, "node", "bin"), npmExecutable, ":")
writeFakeExec(t, nodeArchiveSpec("amd64").binDir(home), npmExecutable, ":")
t.Setenv("PATH", t.TempDir())

ugPath, err := bootstrapToolchain(cmdio.MockDiscard(t.Context()), home)
Expand All @@ -646,8 +679,8 @@ func TestBootstrapToolchain(t *testing.T) {
t.Run("reports UNITY_GATEWAY_CLI_NOT_FOUND when the install leaves no CLI on PATH", func(t *testing.T) {
home := t.TempDir()
// uv "succeeds" without installing anything, so the post-setup lookup fails.
writeFakeExec(t, filepath.Join(home, agentDepsDir, "uv"), "uv", ":")
writeFakeExec(t, filepath.Join(home, agentDepsDir, "node", "bin"), npmExecutable, ":")
writeFakeExec(t, uvArchiveSpec("amd64").binDir(home), "uv", ":")
writeFakeExec(t, nodeArchiveSpec("amd64").binDir(home), npmExecutable, ":")
t.Setenv("PATH", t.TempDir())

_, err := bootstrapToolchain(cmdio.MockDiscard(t.Context()), home)
Expand Down
Loading