Record proxied exit codes and subcommands in command telemetry

cmd/extension.go

@@ -5,7 +5,6 @@ import (
 	"errors"
 	"fmt"
 	"os"
-	"os/exec"
 	"strings"
 	"time"
 
@@ -67,15 +66,11 @@ func dispatchExtension(ctx context.Context, cfg *env.Env, tel *telemetry.Client,
 	start := time.Now()
 	runErr := extension.Invoke(ctx, ext, extArgs, runCtx)
 
-	exitCode, errorMsg := 0, ""
+	exitCode, errorMsg := ExitCode(runErr), ""
 	if runErr != nil {
-		exitCode, errorMsg = 1, runErr.Error()
-		var exitErr *exec.ExitError
-		if errors.As(runErr, &exitErr) {
-			exitCode = exitErr.ExitCode()
-		}
+		errorMsg = runErr.Error()
 	}
-	tel.EmitCommand(ctx, "ext:"+name, nil, time.Since(start).Milliseconds(), exitCode, errorMsg)
+	tel.EmitCommand(ctx, "ext:"+name, "", nil, time.Since(start).Milliseconds(), exitCode, errorMsg)
 
 	return runErr
 }

cmd/instrument_test.go

@@ -0,0 +1,103 @@
+package cmd
+
+import (
+	"errors"
+	"os/exec"
+	"runtime"
+	"strconv"
+	"strings"
+	"testing"
+
+	"github.com/localstack/lstk/internal/output"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// realExitError produces a genuine *exec.ExitError with the given code, the
+// same error shape awscli.Exec returns when a proxied tool exits non-zero.
+func realExitError(t *testing.T, code int) error {
+	t.Helper()
+	var c *exec.Cmd
+	if runtime.GOOS == "windows" {
+		c = exec.Command("cmd", "/c", "exit", strconv.Itoa(code))
+	} else {
+		c = exec.Command("sh", "-c", "exit "+strconv.Itoa(code))
+	}
+	err := c.Run()
+	require.Error(t, err)
+	return err
+}
+
+func TestExitCode(t *testing.T) {
+	t.Run("nil error is 0", func(t *testing.T) {
+		assert.Equal(t, 0, ExitCode(nil))
+	})
+
+	t.Run("plain error is 1", func(t *testing.T) {
+		assert.Equal(t, 1, ExitCode(errors.New("boom")))
+	})
+
+	t.Run("proxied exit code unwraps through SilentError", func(t *testing.T) {
+		err := output.NewSilentError(realExitError(t, 252))
+		assert.Equal(t, 252, ExitCode(err))
+	})
+
+	t.Run("json envelope ExitCodeError code is used", func(t *testing.T) {
+		err := output.NewSilentError(&output.ExitCodeError{Err: errors.New("confirmation required"), Code: 3})
+		assert.Equal(t, 3, ExitCode(err))
+	})
+}
+
+func TestProxySubcommand(t *testing.T) {
+	tests := []struct {
+		name string
+		args []string
+		want string
+	}{
+		{
+			name: "service and operation",
+			args: []string{"s3", "ls"},
+			want: "s3 ls",
+		},
+		{
+			name: "caps at two tokens so values are never recorded",
+			args: []string{"s3", "cp", "file.txt", "s3://bucket"},
+			want: "s3 cp",
+		},
+		{
+			name: "single token",
+			args: []string{"plan"},
+			want: "plan",
+		},
+		{
+			name: "empty args",
+			args: nil,
+			want: "",
+		},
+		{
+			name: "leading double-dash flag stops collection so flag values are never recorded",
+			args: []string{"--region", "us-east-1", "s3", "ls"},
+			want: "",
+		},
+		{
+			name: "single-dash flag stops collection",
+			args: []string{"plan", "-json"},
+			want: "plan",
+		},
+		{
+			name: "lstk global flags are stripped first",
+			args: []string{"--non-interactive", "s3", "ls"},
+			want: "s3 ls",
+		},
+		{
+			name: "overlong token is truncated",
+			args: []string{strings.Repeat("a", 100), "ls"},
+			want: strings.Repeat("a", 64) + " ls",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equal(t, tt.want, proxySubcommand(tt.args))
+		})
+	}
+}

cmd/proxy.go

@@ -51,6 +51,30 @@ func stripGlobalFlags(args []string) ([]string, globalFlags) {
 	return out, gf
 }
 
+// proxySubcommand returns the leading subcommand tokens of a proxy command's
+// raw args for telemetry, e.g. "s3 ls" for `lstk aws s3 ls s3://bucket`. Only
+// leading non-flag tokens are collected: collection stops at the first
+// flag-like arg so a flag's value can never be mistaken for a subcommand, and
+// is capped at two tokens of at most 64 runes each so free-form values are
+// never recorded.
+func proxySubcommand(args []string) string {
+	args, _ = stripGlobalFlags(args)
+	tokens := make([]string, 0, 2)
+	for _, arg := range args {
+		if strings.HasPrefix(arg, "-") {
+			break
+		}
+		if r := []rune(arg); len(r) > 64 {
+			arg = string(r[:64])
+		}
+		tokens = append(tokens, arg)
+		if len(tokens) == 2 {
+			break
+		}
+	}
+	return strings.Join(tokens, " ")
+}
+
 // jsonPrecedesCommandName reports whether --json (or --json=<value>) appears in
 // the raw command line before the literal token calledAs — the resolved proxy
 // command's own name/alias (e.g. "aws", "terraform"/"tf", "az"). Proxy commands

cmd/root.go

@@ -2,8 +2,10 @@ package cmd
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"os"
+	"os/exec"
 	"path/filepath"
 	"strings"
 	"time"
@@ -407,6 +409,28 @@ func commandDisplayName(c *cobra.Command) string {
 	return strings.TrimPrefix(c.CommandPath(), c.Root().Name()+" ")
 }
 
+// ExitCode maps a command error to the exit code the lstk process terminates
+// with: a proxied tool's *exec.ExitError carries that tool's exact code, an
+// output.ExitCodeError carries the --json exit-code convention (3
+// CONFIRMATION_REQUIRED, 4 AUTH_REQUIRED), anything else collapses to 1.
+// errors.As unwraps through the SilentError wrapper to reach either type.
+// main.go and instrumentCommands both use this, so the telemetry exit_code
+// always matches the real process exit code.
+func ExitCode(err error) int {
+	if err == nil {
+		return 0
+	}
+	var exitErr *exec.ExitError
+	if errors.As(err, &exitErr) {
+		return exitErr.ExitCode()
+	}
+	var codeErr *output.ExitCodeError
+	if errors.As(err, &codeErr) {
+		return codeErr.Code
+	}
+	return 1
+}
+
 // instrumentCommands walks the Cobra command tree and wraps every RunE with telemetry emission.
 func instrumentCommands(cmd *cobra.Command, tel *telemetry.Client) {
 	walkCommandsWithRunE(cmd, func(c *cobra.Command) {
@@ -427,14 +451,21 @@ func instrumentCommands(cmd *cobra.Command, tel *telemetry.Client) {
 				flags = append(flags, "--"+f.Name)
 			})
 
-			exitCode := 0
+			// Proxy commands disable flag parsing, so their wrapped tool's
+			// subcommand is invisible in the command path; record its leading
+			// tokens so failures are attributable to a service/operation.
+			subcommand := ""
+			if c.DisableFlagParsing {
+				subcommand = proxySubcommand(args)
+			}
+
+			exitCode := ExitCode(runErr)
 			errorMsg := ""
 			if runErr != nil {
-				exitCode = 1
 				errorMsg = runErr.Error()
 			}
 
-			tel.EmitCommand(c.Context(), commandDisplayName(c), flags, time.Since(startTime).Milliseconds(), exitCode, errorMsg)
+			tel.EmitCommand(c.Context(), commandDisplayName(c), subcommand, flags, time.Since(startTime).Milliseconds(), exitCode, errorMsg)
 
 			return runErr
 		}

internal/telemetry/events.go

@@ -38,10 +38,13 @@ type CommandEvent struct {
 	Result      CommandResult     `json:"result"`
 }
 
-// CommandParameters holds the command name and set flags.
+// CommandParameters holds the command name and set flags. Subcommand carries
+// the leading service/operation tokens of a proxied tool invocation (e.g.
+// "s3 ls" for `lstk aws s3 ls`); empty for lstk's own commands.
 type CommandParameters struct {
-	Command string   `json:"command"`
-	Flags   []string `json:"flags"`
+	Command    string   `json:"command"`
+	Subcommand string   `json:"subcommand,omitempty"`
+	Flags      []string `json:"flags"`
 }
 
 // CommandResult holds the outcome of a command invocation.
@@ -106,10 +109,10 @@ func (c *Client) GetEnvironment(ctx context.Context) Environment {
 
 // EmitCommand emits an lstk_command telemetry event. The Environment block is
 // populated automatically from the client state.
-func (c *Client) EmitCommand(ctx context.Context, command string, flags []string, durationMS int64, exitCode int, errorMsg string) {
+func (c *Client) EmitCommand(ctx context.Context, command, subcommand string, flags []string, durationMS int64, exitCode int, errorMsg string) {
 	c.Emit(ctx, "lstk_command", ToMap(CommandEvent{
 		Environment: c.GetEnvironment(ctx),
-		Parameters:  CommandParameters{Command: command, Flags: flags},
+		Parameters:  CommandParameters{Command: command, Subcommand: subcommand, Flags: flags},
 		Result: CommandResult{
 			DurationMS: durationMS,
 			ExitCode:   exitCode,

internal/telemetry/events_test.go

@@ -68,7 +68,7 @@ func TestEmitCommand_SendsCorrectEventNameAndStructure(t *testing.T) {
 	tel, ch := captureEvents(t)
 
 	tel.SetAuthToken("ls-token")
-	tel.EmitCommand(context.Background(), "start", []string{"--non-interactive"}, 1200, 0, "")
+	tel.EmitCommand(context.Background(), "start", "", []string{"--non-interactive"}, 1200, 0, "")
 
 	got := drainEvent(t, tel, ch)
 
@@ -102,7 +102,7 @@ func TestEmitCommand_SendsCorrectEventNameAndStructure(t *testing.T) {
 func TestEmitCommand_IncludesErrorMsgOnFailure(t *testing.T) {
 	tel, ch := captureEvents(t)
 
-	tel.EmitCommand(context.Background(), "start", nil, 50, 1, "port 4566 already in use")
+	tel.EmitCommand(context.Background(), "start", "", nil, 50, 1, "port 4566 already in use")
 
 	got := drainEvent(t, tel, ch)
 	payload := got["payload"].(map[string]any)
@@ -111,6 +111,32 @@ func TestEmitCommand_IncludesErrorMsgOnFailure(t *testing.T) {
 	assert.InDelta(t, 1, result["exit_code"], 0)
 }
 
+func TestEmitCommand_RecordsSubcommandAndRealExitCode(t *testing.T) {
+	tel, ch := captureEvents(t)
+
+	tel.EmitCommand(context.Background(), "aws", "s3 ls", nil, 80, 252, "exit status 252")
+
+	got := drainEvent(t, tel, ch)
+	payload := got["payload"].(map[string]any)
+	params := payload["parameters"].(map[string]any)
+	assert.Equal(t, "aws", params["command"])
+	assert.Equal(t, "s3 ls", params["subcommand"])
+	result := payload["result"].(map[string]any)
+	assert.InDelta(t, 252, result["exit_code"], 0)
+}
+
+func TestEmitCommand_OmitsSubcommandWhenEmpty(t *testing.T) {
+	tel, ch := captureEvents(t)
+
+	tel.EmitCommand(context.Background(), "start", "", nil, 80, 0, "")
+
+	got := drainEvent(t, tel, ch)
+	payload := got["payload"].(map[string]any)
+	params := payload["parameters"].(map[string]any)
+	_, present := params["subcommand"]
+	assert.False(t, present, "empty subcommand should be omitted from the payload")
+}
+
 func TestEmitCommand_IsNoOpWhenDisabled(t *testing.T) {
 	received := make(chan struct{}, 1)
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -119,7 +145,7 @@ func TestEmitCommand_IsNoOpWhenDisabled(t *testing.T) {
 	defer srv.Close()
 
 	tel := New(srv.URL, true) // disabled
-	tel.EmitCommand(context.Background(), "start", nil, 0, 0, "")
+	tel.EmitCommand(context.Background(), "start", "", nil, 0, 0, "")
 	tel.Close()
 
 	select {

main.go

@@ -2,35 +2,22 @@ package main
 
 import (
 	"context"
-	"errors"
 	"os"
-	"os/exec"
 	"os/signal"
 	"syscall"
 
 	"github.com/localstack/lstk/cmd"
-	"github.com/localstack/lstk/internal/output"
 )
 
 func main() {
 	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
 	defer cancel()
 
 	if err := cmd.Execute(ctx); err != nil {
-		// A proxied tool (aws, terraform, cdk, sam, az, extensions) exited
-		// non-zero: propagate its exact code rather than collapsing to 1.
-		var exitErr *exec.ExitError
-		if errors.As(err, &exitErr) {
-			os.Exit(exitErr.ExitCode())
-		}
-		// A JSON-capable command failed after rendering its error envelope to
-		// stdout: use the --json exit-code convention (3 CONFIRMATION_REQUIRED,
-		// 4 AUTH_REQUIRED, 1 otherwise) attached by wrapCommandsWithJSONEnvelope.
-		// errors.As unwraps through the SilentError wrapper to reach it.
-		var codeErr *output.ExitCodeError
-		if errors.As(err, &codeErr) {
-			os.Exit(codeErr.Code)
-		}
-		os.Exit(1)
+		// A proxied tool's exit code (aws, terraform, cdk, sam, az, extensions)
+		// and the --json exit-code convention are propagated exactly; anything
+		// else collapses to 1. See cmd.ExitCode, which telemetry shares so the
+		// recorded exit_code matches the real one.
+		os.Exit(cmd.ExitCode(err))
 	}
 }