execute.go

package ax

import (
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"time"

	"github.com/spf13/cobra"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/trace"

	"github.com/rshade/ax-go/internal/cli"
	internaltelemetry "github.com/rshade/ax-go/internal/telemetry"
)

// ExecuteOption configures Execute.
type ExecuteOption func(*executeConfig)

type executeConfig struct {
	stdin           io.Reader
	stdout          io.Writer
	stderr          io.Writer
	env             func(string) string
	stdoutIsTTY     *bool
	version         string
	shutdownTimeout time.Duration
}

// WithStdin sets the input stream for Cobra.
func WithStdin(r io.Reader) ExecuteOption {
	return func(cfg *executeConfig) {
		cfg.stdin = r
	}
}

// WithStdout sets the machine payload output stream.
func WithStdout(w io.Writer) ExecuteOption {
	return func(cfg *executeConfig) {
		cfg.stdout = w
	}
}

// WithStderr sets the operational output stream.
func WithStderr(w io.Writer) ExecuteOption {
	return func(cfg *executeConfig) {
		cfg.stderr = w
	}
}

// WithEnv sets the environment lookup used by Execute.
func WithEnv(env func(string) string) ExecuteOption {
	return func(cfg *executeConfig) {
		cfg.env = env
	}
}

// WithStdoutIsTTY overrides TTY detection, primarily for tests.
func WithStdoutIsTTY(isTTY bool) ExecuteOption {
	return func(cfg *executeConfig) {
		cfg.stdoutIsTTY = &isTTY
	}
}

// WithVersion sets the tool version reported in schema and error envelopes.
// When omitted or empty, Execute falls back to ResolveVersion(""), which
// resolves build metadata and is never empty.
func WithVersion(version string) ExecuteOption {
	return func(cfg *executeConfig) {
		cfg.version = version
	}
}

// WithTelemetryShutdownTimeout sets the OTel shutdown timeout.
func WithTelemetryShutdownTimeout(timeout time.Duration) ExecuteOption {
	return func(cfg *executeConfig) {
		cfg.shutdownTimeout = timeout
	}
}

// Execute wraps Cobra execution with AX mode resolution, idempotency, schema,
// error-envelope, and telemetry lifecycle behavior. It returns a deterministic
// exit code and leaves process termination to the caller.
//
// The version reported in __schema output and error envelopes comes from
// WithVersion. When WithVersion is not supplied, Execute falls back to
// ResolveVersion("") — link-time injection, then Go build metadata, then the
// "0.0.0-unknown" sentinel — so the version surfaced to agents is never empty.
//
// When the command returns an *Error, Execute normalizes a copy of it (filling
// in trace ID, tool, and version) before writing the envelope to stderr; the
// caller's *Error value is never mutated.
func Execute(ctx context.Context, root *cobra.Command, opts ...ExecuteOption) int {
	cfg := executeConfig{
		stdin:           os.Stdin,
		stdout:          os.Stdout,
		stderr:          os.Stderr,
		env:             os.Getenv,
		shutdownTimeout: defaultTelemetryShutdownTimeout,
	}
	for _, opt := range opts {
		opt(&cfg)
	}
	if cfg.env == nil {
		cfg.env = os.Getenv
	}
	if cfg.stderr == nil {
		cfg.stderr = os.Stderr
	}
	if cfg.version == "" {
		cfg.version = ResolveVersion("")
	}
	// Serialize all writes to stderr. OTel exporters, zerolog hooks, and the
	// shutdown diagnostic may write concurrently; a mutex writer prevents
	// interleaved or torn output lines.
	cfg.stderr = internaltelemetry.NewLockedWriter(cfg.stderr)

	ctx, telemetry, _ := StartTelemetry(
		ctx,
		WithTelemetryEnv(cfg.env),
		WithTelemetryStderr(cfg.stderr),
		WithTelemetryServiceName(root.Name()),
		WithTelemetryServiceVersion(cfg.version),
		WithTelemetryShutdownBudget(cfg.shutdownTimeout),
	)
	defer func() {
		shutdownCtx, cancel := context.WithTimeout(context.Background(), cfg.shutdownTimeout)
		defer cancel()
		if shutdownErr := telemetry.Shutdown(shutdownCtx); shutdownErr != nil {
			fmt.Fprintf(cfg.stderr, "ax: otel shutdown failed: %s\n",
				internaltelemetry.SanitizeDiagnostic(shutdownErr.Error()))
		}
	}()

	ctx, span := otel.Tracer("github.com/rshade/ax-go").Start(ctx, root.Name())
	defer span.End()

	prepareCommand(root, cfg)
	root.SetIn(cfg.stdin)
	root.SetOut(cfg.stdout)
	root.SetErr(cfg.stderr)

	if executeErr := root.ExecuteContext(ctx); executeErr != nil {
		span.SetStatus(codes.Error, executeErr.Error())
		axErr := normalizeExecuteError(root.Context(), root.Name(), cfg.version, executeErr)
		_ = WriteError(cfg.stderr, axErr)
		return axErr.ExitCode()
	}

	return ExitSuccess
}

func prepareCommand(root *cobra.Command, cfg executeConfig) {
	root.SilenceUsage = true
	root.SilenceErrors = true

	cli.EnsurePersistentStringFlag(root, cli.FlagFormat, "", "output format: json or human")
	cli.EnsurePersistentBoolFlag(root, cli.FlagDryRun, false, "emit the envelope without side effects")
	cli.EnsurePersistentStringFlag(
		root,
		cli.FlagIdempotencyKey,
		"",
		"opaque key used to prevent duplicate-create retries",
	)
	ensureSchemaCommand(root, cfg.version)
	wrapPersistentPreRun(root, cfg)
}

func ensureSchemaCommand(root *cobra.Command, version string) {
	for _, command := range root.Commands() {
		if command.Name() == schemaCommandName {
			return
		}
	}
	root.AddCommand(NewSchemaCommand(root, WithSchemaVersion(version)))
}

func wrapPersistentPreRun(root *cobra.Command, cfg executeConfig) {
	previousE := root.PersistentPreRunE
	previous := root.PersistentPreRun

	root.PersistentPreRunE = func(cmd *cobra.Command, args []string) error {
		format := cli.LookupFlagString(cmd, cli.FlagFormat)
		dryRun := cli.LookupFlagBool(cmd, cli.FlagDryRun)
		idempotencyKey := cli.LookupFlagString(cmd, cli.FlagIdempotencyKey)
		if idempotencyKey == "" {
			idempotencyKey = NewIdempotencyKey()
		}

		stdoutIsTTY := stdoutIsTerminal()
		if cfg.stdoutIsTTY != nil {
			stdoutIsTTY = *cfg.stdoutIsTTY
		}

		mode, err := ResolveMode(format, cfg.env("AGENT_MODE"), stdoutIsTTY)
		if err != nil {
			return NewError(cmd.Context(), "validation_error", err.Error(), WithErrorExitCode(ExitValidation))
		}

		ctx := cmd.Context()
		ctx = WithMode(ctx, mode)
		ctx = WithDryRun(ctx, dryRun)
		ctx = WithIdempotencyKey(ctx, idempotencyKey)
		cmd.SetContext(ctx)
		trace.SpanFromContext(ctx).SetName(cmd.CommandPath())

		if previousE != nil {
			if preRunErr := previousE(cmd, args); preRunErr != nil {
				return preRunErr
			}
		}
		if previous != nil {
			previous(cmd, args)
		}
		return nil
	}
}

// normalizeExecuteError fills empty envelope fields (trace ID, tool, version,
// schema version) on the error the command returned. A caller-supplied *Error
// is copied first: the caller owns that value and Execute must not mutate it,
// so normalization lands on the copy while the caller's fields stay untouched.
func normalizeExecuteError(ctx context.Context, tool, version string, err error) *Error {
	var axErr *Error
	if errors.As(err, &axErr) {
		normalized := *axErr
		if normalized.TraceID == "" {
			normalized.TraceID = TraceIDFromContext(ctx)
		}
		if normalized.Tool == "" {
			normalized.Tool = tool
		}
		if normalized.Version == "" {
			normalized.Version = version
		}
		if normalized.SchemaVersion == "" {
			normalized.SchemaVersion = ErrorSchemaVersion
		}
		return &normalized
	}

	return NewError(
		ctx,
		"internal_error",
		err.Error(),
		WithErrorTool(tool),
		WithErrorVersion(version),
		WithErrorExitCode(ErrorExitCode(err)),
	)
}