lock: add support for locking stdenv + flakerefs (#2465)

Make the `stdenv` (and other flakerefs) lockable and updateable. This
makes it possible to update the stdenv with a regular `devbox update`
and simplifies the logic for how the stdenv commit is chosen:

1. If there's no stdenv flakeref in devbox.lock, resolve
   github:NixOS/nixpkgs/nixpkgs-unstable to a locked ref and store it
   in the lockfile.
2. Otherwise, use the ref in the lockfile.

internal/devbox/devbox.go

@@ -46,6 +46,7 @@ import (
 	"go.jetpack.io/devbox/internal/shellgen"
 	"go.jetpack.io/devbox/internal/telemetry"
 	"go.jetpack.io/devbox/internal/ux"
+	"go.jetpack.io/devbox/nix/flake"
 )
 
 const (
@@ -202,8 +203,14 @@ func (d *Devbox) ConfigHash() (string, error) {
 	return cachehash.Bytes(buf.Bytes()), nil
 }
 
-func (d *Devbox) NixPkgsCommitHash() string {
-	return d.cfg.NixPkgsCommitHash()
+func (d *Devbox) Stdenv() flake.Ref {
+	return flake.Ref{
+		Type:  flake.TypeGitHub,
+		Owner: "NixOS",
+		Repo:  "nixpkgs",
+		Ref:   "nixpkgs-unstable",
+		Rev:   d.cfg.NixPkgsCommitHash(),
+	}
 }
 
 func (d *Devbox) Generate(ctx context.Context) error {

internal/devbox/packages.go

@@ -28,6 +28,7 @@ import (
 	"go.jetpack.io/devbox/internal/setup"
 	"go.jetpack.io/devbox/internal/shellgen"
 	"go.jetpack.io/devbox/internal/telemetry"
+	"go.jetpack.io/devbox/nix/flake"
 	"go.jetpack.io/pkg/auth"
 
 	"go.jetpack.io/devbox/internal/boxcli/usererr"
@@ -101,10 +102,17 @@ func (d *Devbox) Add(ctx context.Context, pkgsNames []string, opts devopt.AddOpt
 			// This means it didn't validate and we don't want to fallback to legacy
 			// Just propagate the error.
 			return err
-		} else if _, err := nix.Search(d.lockfile.LegacyNixpkgsPath(pkg.Raw)); err != nil {
-			// This means it looked like a devbox package or attribute path, but we
-			// could not find it in search or in the legacy nixpkgs path.
-			return usererr.New("Package %s not found", pkg.Raw)
+		} else {
+			installable := flake.Installable{
+				Ref:      d.lockfile.Stdenv(),
+				AttrPath: pkg.Raw,
+			}
+			_, err := nix.Search(installable.String())
+			if err != nil {
+				// This means it looked like a devbox package or attribute path, but we
+				// could not find it in search or in the legacy nixpkgs path.
+				return usererr.New("Package %s not found", pkg.Raw)
+			}
 		}
 
 		ux.Finfof(d.stderr, "Adding package %q to devbox.json\n", packageNameForConfig)

internal/devbox/update.go

@@ -65,6 +65,9 @@ func (d *Devbox) Update(ctx context.Context, opts devopt.UpdateOpts) error {
 		}
 	}
 
+	if err := d.updateStdenv(); err != nil {
+		return err
+	}
 	if err := d.ensureStateIsUpToDate(ctx, update); err != nil {
 		return err
 	}
@@ -103,6 +106,15 @@ func (d *Devbox) inputsToUpdate(
 	return pkgsToUpdate, nil
 }
 
+func (d *Devbox) updateStdenv() error {
+	err := d.lockfile.Remove(d.Stdenv().String())
+	if err != nil {
+		return err
+	}
+	d.lockfile.Stdenv() // will re-resolve the stdenv flake
+	return nil
+}
+
 func (d *Devbox) updateDevboxPackage(pkg *devpkg.Package) error {
 	resolved, err := d.lockfile.FetchResolvedPackage(pkg.Raw)
 	if err != nil {

internal/devconfig/configfile/file.go

@@ -94,11 +94,8 @@ func (c *ConfigFile) Equals(other *ConfigFile) bool {
 }
 
 func (c *ConfigFile) NixPkgsCommitHash() string {
-	// The commit hash for nixpkgs-unstable on 2023-10-25 from status.nixos.org
-	const DefaultNixpkgsCommit = "75a52265bda7fd25e06e3a67dee3f0354e73243c"
-
-	if c == nil || c.Nixpkgs == nil || c.Nixpkgs.Commit == "" {
-		return DefaultNixpkgsCommit
+	if c == nil || c.Nixpkgs == nil {
+		return ""
 	}
 	return c.Nixpkgs.Commit
 }

internal/devpkg/package.go

@@ -151,9 +151,13 @@ func newPackage(raw string, isInstallable func() bool, locker lock.Locker) *Pack
 		return pkg
 	}
 
-	// We currently don't lock flake references in devbox.lock, so there's
-	// nothing to resolve.
-	pkg.resolve = sync.OnceValue(func() error { return nil })
+	pkg.resolve = sync.OnceValue(func() error {
+		// Don't lock flakes that are local paths.
+		if parsed.Ref.Type == flake.TypePath {
+			return nil
+		}
+		return resolve(pkg)
+	})
 	pkg.setInstallable(parsed, locker.ProjectDir())
 	pkg.outputs = outputs{selectedNames: strings.Split(parsed.Outputs, ",")}
 	pkg.Patch = pkgNeedsPatch(pkg.CanonicalName(), configfile.PatchAuto)

internal/devpkg/package_test.go

@@ -12,6 +12,7 @@ import (
 	"github.com/samber/lo"
 	"go.jetpack.io/devbox/internal/lock"
 	"go.jetpack.io/devbox/internal/nix"
+	"go.jetpack.io/devbox/nix/flake"
 )
 
 const nixCommitHash = "hsdafkhsdafhas"
@@ -108,12 +109,13 @@ func (l *lockfile) ProjectDir() string {
 	return l.projectDir
 }
 
-func (l *lockfile) LegacyNixpkgsPath(pkg string) string {
-	return fmt.Sprintf(
-		"github:NixOS/nixpkgs/%s#%s",
-		nixCommitHash,
-		pkg,
-	)
+func (l *lockfile) Stdenv() flake.Ref {
+	return flake.Ref{
+		Type:  flake.TypeGitHub,
+		Owner: "NixOS",
+		Repo:  "nixpkgs",
+		Rev:   nixCommitHash,
+	}
 }
 
 func (l *lockfile) Get(pkg string) *lock.Package {
@@ -128,7 +130,10 @@ func (l *lockfile) Resolve(pkg string) (*lock.Package, error) {
 		return &lock.Package{Resolved: pkg}, nil
 	default:
 		return &lock.Package{
-			Resolved: l.LegacyNixpkgsPath(pkg),
+			Resolved: flake.Installable{
+				Ref:      l.Stdenv(),
+				AttrPath: pkg,
+			}.String(),
 		}, nil
 	}
 }

internal/lock/interfaces.go

@@ -3,16 +3,18 @@
 
 package lock
 
+import "go.jetpack.io/devbox/nix/flake"
+
 type devboxProject interface {
 	ConfigHash() (string, error)
-	NixPkgsCommitHash() string
+	Stdenv() flake.Ref
 	AllPackageNamesIncludingRemovedTriggerPackages() []string
 	ProjectDir() string
 }
 
 type Locker interface {
 	Get(string) *Package
-	LegacyNixpkgsPath(string) string
+	Stdenv() flake.Ref
 	ProjectDir() string
 	Resolve(string) (*Package, error)
 }

internal/lock/lockfile.go

@@ -5,17 +5,18 @@ package lock
 
 import (
 	"context"
-	"fmt"
 	"io/fs"
+	"maps"
 	"path/filepath"
+	"slices"
 	"strings"
 
 	"github.com/pkg/errors"
-	"github.com/samber/lo"
 	"go.jetpack.io/devbox/internal/cachehash"
 	"go.jetpack.io/devbox/internal/devpkg/pkgtype"
 	"go.jetpack.io/devbox/internal/nix"
 	"go.jetpack.io/devbox/internal/searcher"
+	"go.jetpack.io/devbox/nix/flake"
 	"go.jetpack.io/pkg/runx/impl/types"
 
 	"go.jetpack.io/devbox/internal/cuecfg"
@@ -74,25 +75,32 @@ func (f *File) Remove(pkgs ...string) error {
 // This avoids writing values that may need to be removed in case of error.
 func (f *File) Resolve(pkg string) (*Package, error) {
 	entry, hasEntry := f.Packages[pkg]
+	if hasEntry && entry.Resolved != "" {
+		return f.Packages[pkg], nil
+	}
 
-	if !hasEntry || entry.Resolved == "" {
-		locked := &Package{}
-		var err error
-		if _, _, versioned := searcher.ParseVersionedPackage(pkg); pkgtype.IsRunX(pkg) || versioned {
-			locked, err = f.FetchResolvedPackage(pkg)
-			if err != nil {
-				return nil, err
-			}
-		} else if IsLegacyPackage(pkg) {
-			// These are legacy packages without a version. Resolve to nixpkgs with
-			// whatever hash is in the devbox.json
-			locked = &Package{
-				Resolved: f.LegacyNixpkgsPath(pkg),
-				Source:   nixpkgSource,
-			}
+	locked := &Package{}
+	_, _, versioned := searcher.ParseVersionedPackage(pkg)
+	if pkgtype.IsRunX(pkg) || versioned || pkgtype.IsFlake(pkg) {
+		resolved, err := f.FetchResolvedPackage(pkg)
+		if err != nil {
+			return nil, err
+		}
+		if resolved != nil {
+			locked = resolved
+		}
+	} else if IsLegacyPackage(pkg) {
+		// These are legacy packages without a version. Resolve to nixpkgs with
+		// whatever hash is in the devbox.json
+		locked = &Package{
+			Resolved: flake.Installable{
+				Ref:      f.Stdenv(),
+				AttrPath: pkg,
+			}.String(),
+			Source: nixpkgSource,
 		}
-		f.Packages[pkg] = locked
 	}
+	f.Packages[pkg] = locked
 
 	return f.Packages[pkg], nil
 }
@@ -133,12 +141,17 @@ func (f *File) Save() error {
 	return cuecfg.WriteFile(lockFilePath(f.devboxProject.ProjectDir()), f)
 }
 
-func (f *File) LegacyNixpkgsPath(pkg string) string {
-	return fmt.Sprintf(
-		"github:NixOS/nixpkgs/%s#%s",
-		f.NixPkgsCommitHash(),
-		pkg,
-	)
+func (f *File) Stdenv() flake.Ref {
+	unlocked := f.devboxProject.Stdenv()
+	pkg, err := f.Resolve(unlocked.String())
+	if err != nil {
+		return unlocked
+	}
+	ref, err := flake.ParseRef(pkg.Resolved)
+	if err != nil {
+		return unlocked
+	}
+	return ref
 }
 
 func (f *File) Get(pkg string) *Package {
@@ -174,10 +187,11 @@ func IsLegacyPackage(pkg string) bool {
 // Tidy ensures that the lockfile has the set of packages corresponding to the devbox.json config.
 // It gets rid of older packages that are no longer needed.
 func (f *File) Tidy() {
-	f.Packages = lo.PickByKeys(
-		f.Packages,
-		f.devboxProject.AllPackageNamesIncludingRemovedTriggerPackages(),
-	)
+	keep := f.devboxProject.AllPackageNamesIncludingRemovedTriggerPackages()
+	keep = append(keep, f.devboxProject.Stdenv().String())
+	maps.DeleteFunc(f.Packages, func(key string, pkg *Package) bool {
+		return !slices.Contains(keep, key)
+	})
 }
 
 // IsUpToDateAndInstalled returns true if the lockfile is up to date and the

internal/lock/resolve.go

@@ -18,6 +18,7 @@ import (
 	"go.jetpack.io/devbox/internal/nix"
 	"go.jetpack.io/devbox/internal/redact"
 	"go.jetpack.io/devbox/internal/searcher"
+	"go.jetpack.io/devbox/nix/flake"
 	"golang.org/x/sync/errgroup"
 )
 
@@ -29,7 +30,17 @@ import (
 // to update because it would be slow and wasteful.
 func (f *File) FetchResolvedPackage(pkg string) (*Package, error) {
 	if pkgtype.IsFlake(pkg) {
-		return nil, nil
+		installable, err := flake.ParseInstallable(pkg)
+		if err != nil {
+			return nil, fmt.Errorf("package %q: %v", pkg, err)
+		}
+		installable.Ref, err = lockFlake(context.TODO(), installable.Ref)
+		if err != nil {
+			return nil, err
+		}
+		return &Package{
+			Resolved: installable.String(),
+		}, nil
 	}
 
 	name, version, _ := searcher.ParseVersionedPackage(pkg)
@@ -194,3 +205,24 @@ func buildLockSystemInfos(pkg *searcher.PackageVersion) (map[string]*SystemInfo,
 	}
 	return sysInfos, nil
 }
+
+func lockFlake(ctx context.Context, ref flake.Ref) (flake.Ref, error) {
+	if ref.Locked() {
+		return ref, nil
+	}
+
+	// Nix requires a NAR hash for GitHub flakes to be locked. A Devbox lock
+	// file is a bit more lenient and only requires a revision so that we
+	// don't need to download the nixpkgs source for cached packages. If the
+	// search index is ever able to return the NAR hash then we can remove
+	// this check.
+	if ref.Type == flake.TypeGitHub && (ref.Rev != "") {
+		return ref, nil
+	}
+
+	meta, err := nix.ResolveFlake(ctx, ref)
+	if err != nil {
+		return ref, err
+	}
+	return meta.Locked, nil
+}

internal/nix/build.go

@@ -20,6 +20,9 @@ type BuildArgs struct {
 
 func Build(ctx context.Context, args *BuildArgs, installables ...string) error {
 	defer debug.FunctionTimer().End()
+
+	FixInstallableArgs(installables)
+
 	// --impure is required for allowUnfreeEnv/allowInsecureEnv to work.
 	cmd := command("build", "--impure")
 	cmd.Args = appendArgs(cmd.Args, args.Flags)

internal/nix/flake.go

@@ -0,0 +1,30 @@
+package nix
+
+import (
+	"context"
+	"encoding/json"
+
+	"go.jetpack.io/devbox/nix/flake"
+)
+
+type FlakeMetadata struct {
+	Description string    `json:"description"`
+	Original    flake.Ref `json:"original"`
+	Resolved    flake.Ref `json:"resolved"`
+	Locked      flake.Ref `json:"locked"`
+	Path        string    `json:"path"`
+}
+
+func ResolveFlake(ctx context.Context, ref flake.Ref) (FlakeMetadata, error) {
+	cmd := command("flake", "metadata", "--json", ref)
+	out, err := cmd.Output(ctx)
+	if err != nil {
+		return FlakeMetadata{}, err
+	}
+	meta := FlakeMetadata{}
+	err = json.Unmarshal(out, &meta)
+	if err != nil {
+		return FlakeMetadata{}, err
+	}
+	return meta, nil
+}