context.go

// Copyright 2026 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package jiri

import (
	"context"
	"io"
	"os"

	"go.fuchsia.dev/jiri/cmdline"
	"go.fuchsia.dev/jiri/color"
	"go.fuchsia.dev/jiri/envvar"
	"go.fuchsia.dev/jiri/log"
)

type checkoutContextKey struct{}

// WithCheckout returns a copy of ctx that carries c, retrievable with
// CheckoutFromContext.
func WithCheckout(ctx context.Context, c *Checkout) context.Context {
	return context.WithValue(ctx, checkoutContextKey{}, c)
}

// CheckoutFromContext returns the Checkout attached to ctx by WithCheckout.
//
// It panics if ctx has no Checkout, since every jiri code path that takes a
// ctx is expected to run within a jiri command invocation, and silently
// falling back to a zero-value Checkout would hide bugs.
func CheckoutFromContext(ctx context.Context) *Checkout {
	c, ok := ctx.Value(checkoutContextKey{}).(*Checkout)
	if !ok || c == nil {
		panic("jiri.CheckoutFromContext: context has no *jiri.Checkout attached; use jiri.WithCheckout")
	}
	return c
}

// NewContext returns a copy of ctx that carries the checkout enclosing the
// current directory (or the one named by flags.Root), along with the logger
// and color configured by flags. ctx must already carry a cmdline.Env.
func NewContext(ctx context.Context, flags TopLevelFlags) (context.Context, error) {
	return newContext(ctx, flags, true /* requiresRoot */)
}

// NewContextWithoutRoot is like NewContext, but doesn't search for or require
// a .jiri_root checkout root.
//
// CAUTION: This should ONLY be used in limited contexts (such as self-update)
// where a .jiri_root checkout root is intentionally not required.
func NewContextWithoutRoot(ctx context.Context, flags TopLevelFlags) (context.Context, error) {
	return newContext(ctx, flags, false /* requiresRoot */)
}

func newContext(ctx context.Context, flags TopLevelFlags, requiresRoot bool) (context.Context, error) {
	env := cmdlineEnv(ctx)
	cf := color.EnableColor(flags.Color)
	if cf != color.ColorAuto && cf != color.ColorAlways && cf != color.ColorNever {
		return nil, env.UsageErrorf("invalid value of -color flag")
	}
	c := color.NewColor(cf)

	loggerLevel := log.InfoLevel
	if flags.QuietVerbose {
		loggerLevel = log.WarningLevel
	} else if flags.TraceVerbose {
		loggerLevel = log.TraceLevel
	} else if flags.DebugVerbose {
		loggerLevel = log.DebugLevel
	}
	if flags.ProgressWindowSize < 1 {
		flags.ProgressWindowSize = 1
	} else if flags.ProgressWindowSize > 10 {
		flags.ProgressWindowSize = 10
	}
	logger := log.NewLogger(loggerLevel, c, flags.ShowProgress, flags.ProgressWindowSize, flags.TimeLogThreshold, Stdout(ctx), Stderr(ctx))

	ctx = color.WithColor(ctx, c)
	ctx = log.WithLogger(ctx, logger)
	ctx = WithRunState(ctx)
	checkout, err := newCheckout(ctx, flags, requiresRoot)
	if err != nil {
		return nil, err
	}
	if checkout.Root != "" && env.Vars[PreservePathEnv] == "" {
		// Prepend .jiri_root/bin to the PATH, so execing a binary will invoke
		// the one in that directory, if it exists. This is crucial for jiri
		// subcommands, where we want to invoke the binary that jiri installed,
		// not whatever is in the user's PATH.
		//
		// Both the os env and the ctx's env need updating, so that execing a
		// binary through os/exec behaves the same with or without EnvVars.
		newPath := envvar.PrependUniqueToken(env.Vars["PATH"], string(os.PathListSeparator), checkout.BinDir())
		if err := os.Setenv("PATH", newPath); err != nil {
			return nil, err
		}
		newEnv := *env
		newEnv.Vars = envvar.CopyMap(env.Vars)
		newEnv.Vars["PATH"] = newPath
		ctx = cmdline.AddEnvToContext(ctx, &newEnv)
	}
	return WithCheckout(ctx, checkout), nil
}

// The functions below read from the cmdline.Env that cmdline attaches to
// every command's context. They exist so that call sites don't each have to
// fetch the Env and deal with it being missing.

func cmdlineEnv(ctx context.Context) *cmdline.Env {
	env := cmdline.EnvFromContext(ctx)
	if env == nil {
		panic("context has no *cmdline.Env attached; use cmdline.AddEnvToContext")
	}
	return env
}

// Stdin returns the standard input stream for the current command.
func Stdin(ctx context.Context) io.Reader {
	if r := cmdlineEnv(ctx).Stdin; r != nil {
		return r
	}
	return os.Stdin
}

// Stdout returns the standard output stream for the current command.
func Stdout(ctx context.Context) io.Writer {
	if w := cmdlineEnv(ctx).Stdout; w != nil {
		return w
	}
	return os.Stdout
}

// Stderr returns the standard error stream for the current command.
func Stderr(ctx context.Context) io.Writer {
	if w := cmdlineEnv(ctx).Stderr; w != nil {
		return w
	}
	return os.Stderr
}

// EnvVars returns the environment variables that subprocesses should run
// with. Callers must not modify the returned map.
func EnvVars(ctx context.Context) map[string]string {
	return cmdlineEnv(ctx).Vars
}

// WithStdout returns a copy of ctx whose standard output stream is w.
func WithStdout(ctx context.Context, w io.Writer) context.Context {
	env := *cmdlineEnv(ctx)
	env.Stdout = w
	return cmdline.AddEnvToContext(ctx, &env)
}

// WithStderr returns a copy of ctx whose standard error stream is w.
func WithStderr(ctx context.Context, w io.Writer) context.Context {
	env := *cmdlineEnv(ctx)
	env.Stderr = w
	return cmdline.AddEnvToContext(ctx, &env)
}

// TimerPush starts timing an operation called name, if ctx has a timer.
func TimerPush(ctx context.Context, name string) {
	cmdlineEnv(ctx).TimerPush(name)
}

// TimerPop finishes timing the operation most recently started by
// TimerPush.
func TimerPop(ctx context.Context) {
	cmdlineEnv(ctx).TimerPop()
}

// UsageErrorf prints the message built from format and args, followed by the
// current command's usage, to the command's stderr. It always returns
// cmdline.ErrUsage so that the command exits with the usage-error code.
func UsageErrorf(ctx context.Context, format string, args ...any) error {
	return cmdlineEnv(ctx).UsageErrorf(format, args...)
}