feat(integrate): add --run-init-hook flag to run init_hook when reopening an editor

internal/boxcli/integrate.go

@@ -21,9 +21,10 @@ import (
 )
 
 type integrateCmdFlags struct {
-	config    configFlags
-	debugmode bool
-	ideName   string
+	config      configFlags
+	debugmode   bool
+	ideName     string
+	runInitHook bool
 }
 
 func integrateCmd() *cobra.Command {
@@ -53,6 +54,10 @@ func integrateVSCodeCmd() *cobra.Command {
 	}
 	command.Flags().BoolVar(&flags.debugmode, "debugmode", false, "enable debug outputs to a file.")
 	command.Flags().StringVar(&flags.ideName, "ide", "code", "name of the currently open editor to reopen after it's closed.")
+	command.Flags().BoolVar(
+		&flags.runInitHook, "run-init-hook", false,
+		"run the project's init_hook and include the variables it sets in the editor's environment.",
+	)
 	flags.config.register(command)
 
 	return command
@@ -95,9 +100,20 @@ func runIntegrateVSCodeCmd(cmd *cobra.Command, flags integrateCmdFlags) error {
 		dbug.logToFile(err.Error())
 		return err
 	}
-	// Get env variables of a devbox shell
+	// Get env variables of a devbox shell. The editor is launched directly (not
+	// through a devbox shell), so the init hook never runs and the variables it
+	// sets are missing from the reopened environment unless --run-init-hook is
+	// set, in which case we source the init hook and capture its env too. This
+	// is opt-in because init hooks can be slow or have side effects. See issue
+	// #2703.
 	dbug.logToFile("Computing devbox environment")
-	envVars, err := box.EnvVars(cmd.Context())
+	var envVars []string
+	if flags.runInitHook {
+		dbug.logToFile("Running init hook to capture its environment")
+		envVars, err = box.EnvVarsWithInitHook(cmd.Context())
+	} else {
+		envVars, err = box.EnvVars(cmd.Context())
+	}
 	if err != nil {
 		dbug.logToFile(err.Error())
 		return err

internal/devbox/inithookenv.go

@@ -0,0 +1,110 @@
+// Copyright 2024 Jetify Inc. and contributors. All rights reserved.
+// Use of this source code is governed by the license in the LICENSE file.
+
+package devbox
+
+import (
+	"context"
+	"io"
+	"log/slog"
+	"os/exec"
+	"runtime/trace"
+	"strings"
+
+	"github.com/pkg/errors"
+	"go.jetify.com/devbox/internal/devbox/devopt"
+	"go.jetify.com/devbox/internal/envir"
+	"go.jetify.com/devbox/internal/fileutil"
+	"go.jetify.com/devbox/internal/shellgen"
+)
+
+// EnvVarsWithInitHook returns the environment variables for the Devbox
+// environment, including any variables set (or modified) by the project's
+// init hook.
+//
+// Unlike EnvVars, which deliberately excludes the init hook, this sources the
+// init hook in a subshell and captures the resulting environment. It is meant
+// for integrations that launch a program directly, without going through a
+// devbox shell (or `devbox run`) that would otherwise source the init hook —
+// for example the VSCode "Reopen in Devbox" action. See issue #2703.
+//
+// If the init hook fails, the environment without the hook's modifications is
+// returned rather than erroring, so the integration keeps working.
+func (d *Devbox) EnvVarsWithInitHook(ctx context.Context) ([]string, error) {
+	ctx, task := trace.NewTask(ctx, "devboxEnvVarsWithInitHook")
+	defer task.End()
+
+	env, err := d.ensureStateIsUpToDateAndComputeEnv(ctx, devopt.EnvOptions{})
+	if err != nil {
+		return nil, err
+	}
+
+	// Persist the init hook (and scripts) to disk so we can source the hooks
+	// file below.
+	if err := shellgen.WriteScriptsToFiles(d); err != nil {
+		return nil, err
+	}
+
+	hooksPath := shellgen.ScriptPath(d.ProjectDir(), shellgen.HooksFilename)
+	withHooks, err := captureEnvWithInitHook(ctx, hooksPath, env, d.stderr)
+	if err != nil {
+		// Don't fail the whole integration if the init hook errors. Fall back
+		// to the environment without the hook's modifications.
+		slog.Debug("failed to run init hook while computing env", "err", err)
+		return envir.MapToPairs(env), nil
+	}
+	return envir.MapToPairs(withHooks), nil
+}
+
+// captureEnvWithInitHook sources the init hook script at hooksPath using the
+// given base environment and returns the resulting environment. The init
+// hook's own stdout is redirected to hookStderr so it cannot corrupt the
+// captured environment dump.
+func captureEnvWithInitHook(
+	ctx context.Context,
+	hooksPath string,
+	baseEnv map[string]string,
+	hookStderr io.Writer,
+) (map[string]string, error) {
+	if !fileutil.Exists(hooksPath) {
+		// No hooks file; nothing to source.
+		return baseEnv, nil
+	}
+
+	// Source the init hook, then print the resulting environment NUL-separated
+	// so values containing newlines stay intact.
+	//
+	//   - The hooks path is passed as a positional parameter ($1) rather than
+	//     interpolated into the script, so special characters in the path (e.g.
+	//     spaces, quotes, $(), backticks) can't change shell parsing.
+	//   - The hook's own stdout is redirected to stderr (1>&2) so only the env
+	//     dump reaches stdout and can't corrupt it.
+	//   - awk's POSIX ENVIRON is used instead of `env -0`: the latter isn't
+	//     supported by macOS' default /usr/bin/env, whereas awk is portable.
+	const script = `. "$1" 1>&2
+exec awk 'BEGIN { for (k in ENVIRON) printf "%s=%s%c", k, ENVIRON[k], 0 }'`
+	cmd := exec.CommandContext(ctx, "sh", "-c", script, "sh", hooksPath)
+	cmd.Env = envir.MapToPairs(baseEnv)
+	cmd.Stderr = hookStderr
+	out, err := cmd.Output()
+	if err != nil {
+		return nil, errors.WithStack(err)
+	}
+	return parseNulEnv(out), nil
+}
+
+// parseNulEnv parses NUL-separated KEY=VALUE pairs into a map.
+func parseNulEnv(b []byte) map[string]string {
+	result := map[string]string{}
+	for _, pair := range strings.Split(string(b), "\x00") {
+		if pair == "" {
+			continue
+		}
+		key, value, ok := strings.Cut(pair, "=")
+		if !ok {
+			continue
+		}
+		result[key] = value
+	}
+	return result
+}

internal/devbox/inithookenv_test.go

@@ -0,0 +1,81 @@
+// Copyright 2024 Jetify Inc. and contributors. All rights reserved.
+// Use of this source code is governed by the license in the LICENSE file.
+
+package devbox
+
+import (
+	"bytes"
+	"os"
+	"path/filepath"
+	"testing"
+)
+
+func TestParseNulEnv(t *testing.T) {
+	in := []byte("FOO=bar\x00MULTI=line1\nline2\x00EMPTY=\x00")
+	got := parseNulEnv(in)
+	want := map[string]string{
+		"FOO":   "bar",
+		"MULTI": "line1\nline2",
+		"EMPTY": "",
+	}
+	if len(got) != len(want) {
+		t.Fatalf("got %d entries, want %d: %#v", len(got), len(want), got)
+	}
+	for k, v := range want {
+		if got[k] != v {
+			t.Errorf("key %q: got %q, want %q", k, got[k], v)
+		}
+	}
+}
+
+func TestCaptureEnvWithInitHook(t *testing.T) {
+	dir := t.TempDir()
+	hooksPath := filepath.Join(dir, "hooks.sh")
+	// The init hook sets a new var, modifies an existing one, and prints to
+	// stdout (which must not leak into the captured env).
+	hookBody := "echo 'hello from init hook'\n" +
+		"export TEST=true\n" +
+		"export BASE=modified\n"
+	if err := os.WriteFile(hooksPath, []byte(hookBody), 0o755); err != nil {
+		t.Fatal(err)
+	}
+
+	baseEnv := map[string]string{
+		"BASE": "original",
+		"PATH": os.Getenv("PATH"),
+	}
+
+	var hookStdout bytes.Buffer
+	got, err := captureEnvWithInitHook(t.Context(), hooksPath, baseEnv, &hookStdout)
+	if err != nil {
+		t.Fatalf("captureEnvWithInitHook returned error: %v", err)
+	}
+
+	if got["TEST"] != "true" {
+		t.Errorf("TEST: got %q, want %q", got["TEST"], "true")
+	}
+	if got["BASE"] != "modified" {
+		t.Errorf("BASE: got %q, want %q (init hook should override base env)", got["BASE"], "modified")
+	}
+
+	// The init hook's stdout must not corrupt the captured environment.
+	if _, ok := got["hello from init hook"]; ok {
+		t.Errorf("init hook stdout leaked into captured env: %#v", got)
+	}
+}
+
+func TestCaptureEnvWithInitHook_NoHooksFile(t *testing.T) {
+	baseEnv := map[string]string{"FOO": "bar"}
+	got, err := captureEnvWithInitHook(
+		t.Context(),
+		filepath.Join(t.TempDir(), "does-not-exist.sh"),
+		baseEnv,
+		&bytes.Buffer{},
+	)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if got["FOO"] != "bar" {
+		t.Errorf("expected base env to be returned unchanged, got %#v", got)
+	}
+}