fix(telemetry): capture log tail through upload

Author: shubham-stepsecurity

internal/telemetry/execution_deadline_test.go

@@ -24,12 +24,14 @@ func TestExecutionDeadlineFromEnv(t *testing.T) {
 	}
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
+			// Always set the var (to "" for the unset cases) so the test is
+			// hermetic: a value inherited from the host or a CI job must not
+			// leak into the "unset" cases. os.Getenv treats truly-unset and
+			// empty identically, so "" exercises the same fall-through path.
 			if tc.set {
 				t.Setenv("STEPSEC_MAX_EXECUTION_DURATION", tc.env)
 			} else {
-				// Belt-and-braces: t.Setenv resets at test exit, but
-				// nothing in this package mutates the env at init.
-				_ = tc.env
+				t.Setenv("STEPSEC_MAX_EXECUTION_DURATION", "")
 			}
 			got := ExecutionDeadlineFromEnv()
 			if got != tc.want {
@@ -63,8 +65,12 @@ func TestExecutionDeadline_EnvThenConfigThenDefault(t *testing.T) {
 	}
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
+			// Always set the var (to "" when setEnv is false) so an inherited
+			// host/CI value can't bypass the config-fallback cases under test.
 			if tc.setEnv {
 				t.Setenv("STEPSEC_MAX_EXECUTION_DURATION", tc.env)
+			} else {
+				t.Setenv("STEPSEC_MAX_EXECUTION_DURATION", "")
 			}
 			got := ExecutionDeadline(tc.configVal)
 			if got != tc.want {

internal/telemetry/log_tail_emitter.go

@@ -79,6 +79,30 @@ func (e *logTailEmitter) MaybeAttach(info *RunStatusInfo) {
 	info.LogTailGzipBase64 = encoded
 }
 
+// ForceAttach attaches the current tail regardless of the throttle window.
+// Used for the single final post-upload snapshot so the run-status row carries
+// the most complete tail — including the upload's own output and the
+// completion line — which a throttled MaybeAttach would otherwise drop on a
+// sub-2-minute run. Updates lastSent so any later MaybeAttach still respects
+// the interval. No-op on nil capture / empty buffer.
+func (e *logTailEmitter) ForceAttach(info *RunStatusInfo) {
+	if e == nil || e.capture == nil || info == nil {
+		return
+	}
+	tail := e.capture.Tail(captureTailBytes)
+	if len(tail) == 0 {
+		return
+	}
+	encoded, err := gzipBase64(tail)
+	if err != nil {
+		return
+	}
+	e.mu.Lock()
+	e.lastSent = e.now()
+	e.mu.Unlock()
+	info.LogTailGzipBase64 = encoded
+}
+
 // gzipBase64 gzips b at default compression and base64-encodes the
 // result. Wrapped here rather than inlined so the emitter and tests
 // share a single encoding pipeline.

internal/telemetry/log_tail_emitter_test.go

@@ -63,16 +63,65 @@ func TestLogTailEmitter_ThrottlesAttachments(t *testing.T) {
 	}
 }
 
+// ForceAttach must attach the current tail even inside the throttle window —
+// the final post-upload snapshot depends on it so the last lines (the upload's
+// own output) aren't dropped by the throttle on a sub-2-minute run.
+func TestLogTailEmitter_ForceAttachBypassesThrottle(t *testing.T) {
+	lc := &LogCapture{ring: newRingBuffer(64 * 1024)}
+	lc.ring.Write([]byte("startup output\n"))
+
+	clk := &stepClock{t: time.Unix(1_700_000_000, 0)}
+	em := newEmitterWithFakeClock(lc, 2*time.Minute, clk)
+
+	// First MaybeAttach consumes the throttle slot.
+	var snap RunStatusInfo
+	em.MaybeAttach(&snap)
+	if snap.LogTailGzipBase64 == "" {
+		t.Fatal("first MaybeAttach should attach")
+	}
+
+	// Still within the throttle window: MaybeAttach skips, ForceAttach does not.
+	clk.advance(10 * time.Second)
+	lc.ring.Write([]byte("upload completed\n"))
+
+	var throttled RunStatusInfo
+	em.MaybeAttach(&throttled)
+	if throttled.LogTailGzipBase64 != "" {
+		t.Fatal("MaybeAttach within the window should skip")
+	}
+
+	var forced RunStatusInfo
+	em.ForceAttach(&forced)
+	if forced.LogTailGzipBase64 == "" {
+		t.Fatal("ForceAttach must attach even within the throttle window")
+	}
+	if decoded := decodeTail(t, forced.LogTailGzipBase64); !strings.Contains(decoded, "upload completed") {
+		t.Fatalf("ForceAttach must reflect the latest buffer; got %q", decoded)
+	}
+
+	// ForceAttach updated lastSent, so a follow-on MaybeAttach still throttles.
+	clk.advance(10 * time.Second)
+	var after RunStatusInfo
+	em.MaybeAttach(&after)
+	if after.LogTailGzipBase64 != "" {
+		t.Fatal("MaybeAttach should still be throttled after ForceAttach")
+	}
+}
+
 func TestLogTailEmitter_NilSafe(t *testing.T) {
-	// Nil receiver, nil capture, and nil info should all be no-ops, not panics.
+	// Nil receiver, nil capture, and nil info should all be no-ops, not panics —
+	// for both MaybeAttach and ForceAttach.
 	var em *logTailEmitter
 	em.MaybeAttach(&RunStatusInfo{}) // nil receiver
+	em.ForceAttach(&RunStatusInfo{}) // nil receiver
 
 	em2 := newLogTailEmitter(nil, 2*time.Minute)
 	em2.MaybeAttach(&RunStatusInfo{}) // nil capture
+	em2.ForceAttach(&RunStatusInfo{}) // nil capture
 
 	em3 := newLogTailEmitter(&LogCapture{ring: newRingBuffer(1024)}, 2*time.Minute)
 	em3.MaybeAttach(nil) // nil info
+	em3.ForceAttach(nil) // nil info
 }
 
 func TestRingBuffer_TailRespectsWraparound(t *testing.T) {

internal/telemetry/logcapture.go

@@ -175,6 +175,17 @@ func (lc *LogCapture) Finalize() string {
 	return base64.StdEncoding.EncodeToString(lc.ringBytesLocked())
 }
 
+// SnapshotBase64 returns the base64-encoded buffer contents WITHOUT stopping
+// capture, so a caller can embed the session-so-far in the telemetry payload
+// while the capture keeps recording (e.g. through the upload that follows).
+// The real teardown — closing the pipe and restoring os.Stderr — stays in
+// Finalize, which the caller still defers. Safe to call during active capture.
+func (lc *LogCapture) SnapshotBase64() string {
+	lc.mu.Lock()
+	defer lc.mu.Unlock()
+	return base64.StdEncoding.EncodeToString(lc.ringBytesLocked())
+}
+
 // Tail returns the last n captured bytes as a fresh slice. Safe to call
 // concurrently with active capture; returns nil if the buffer is empty
 // or n ≤ 0. Used by heartbeat posts to ship the most recent diagnostic

internal/telemetry/phase_deadline.go

@@ -2,6 +2,8 @@ package telemetry
 
 import (
 	"context"
+	"os"
+	"strings"
 	"time"
 
 	"github.com/step-security/dev-machine-guard/internal/progress"
@@ -13,13 +15,15 @@ import (
 // production heartbeat data — they exist to bound the pathological tail,
 // not to clip normal scans.
 //
-// Order of overrides at a single phase site:
-//  1. STEPSEC_PHASE_BUDGET_<NAME> env var (Go duration, e.g. "10m") — if set
-//  2. this map's entry
-//  3. defaultPhaseBudget (5m) when neither is set
+// Order of overrides at a single phase site (resolved by resolvePhaseBudget):
+//  1. STEPSEC_PHASE_BUDGET_<NAME> env var — <NAME> is the phase name
+//     upper-cased (e.g. node_scan → STEPSEC_PHASE_BUDGET_NODE_SCAN=20m).
+//     "0" or "off" disables the deadline; an unparseable value is ignored.
+//  2. this map's entry (an explicit 0 here also disables the deadline)
+//  3. defaultPhaseBudget (5m) when the phase has no entry above
 //
-// A budget of 0 disables the per-phase deadline for that phase (the parent
-// scan deadline still applies).
+// "Disabled" means no per-phase deadline; the parent scan deadline
+// (STEPSEC_MAX_SCAN_DURATION) and the whole-process watchdog still apply.
 var phaseBudgets = map[string]time.Duration{
 	"device_info":      30 * time.Second,
 	"ide_scan":         2 * time.Minute,
@@ -35,18 +39,48 @@ var phaseBudgets = map[string]time.Duration{
 
 const defaultPhaseBudget = 5 * time.Minute
 
+// resolvePhaseBudget computes a phase's effective budget, honoring (in order)
+// the STEPSEC_PHASE_BUDGET_<NAME> env override, the phaseBudgets map entry,
+// then defaultPhaseBudget. The bool is false when the phase has NO deadline:
+// either source can disable it ("0"/"off" in the env var, or an explicit 0 in
+// the map). A missing map entry — distinct from an entry of 0 — falls back to
+// defaultPhaseBudget. An unparseable/empty env value is ignored so a typo in
+// an unattended scheduler context can't strip every phase's deadline.
+func resolvePhaseBudget(name string) (time.Duration, bool) {
+	if v := os.Getenv("STEPSEC_PHASE_BUDGET_" + strings.ToUpper(name)); v != "" {
+		switch {
+		case v == "0" || v == "off":
+			return 0, false
+		default:
+			if d, err := time.ParseDuration(v); err == nil && d > 0 {
+				return d, true
+			}
+			// Unparseable or non-positive: ignore and fall through.
+		}
+	}
+	if budget, ok := phaseBudgets[name]; ok {
+		if budget <= 0 {
+			return 0, false // explicit 0 in the map disables the deadline
+		}
+		return budget, true
+	}
+	return defaultPhaseBudget, true
+}
+
 // startPhase opens a new phase and returns a derived context that carries
 // the phase's budget as its deadline. Callers must invoke endPhase (or
 // otherwise call the returned cancel func) before opening the next phase.
+// When the phase budget is disabled (see resolvePhaseBudget) the returned
+// context carries no per-phase deadline — only the parent scan deadline.
 //
 // The caller continues to own postPhase() — endPhase only handles the
 // tracker.Finish + cancel + deadline-overrun log line so the per-phase
 // edit at each site stays small.
 func startPhase(parent context.Context, tracker *PhaseTracker, name string) (context.Context, context.CancelFunc) {
 	tracker.Start(name)
-	budget := phaseBudgets[name]
-	if budget == 0 {
-		budget = defaultPhaseBudget
+	budget, hasDeadline := resolvePhaseBudget(name)
+	if !hasDeadline {
+		return context.WithCancel(parent)
 	}
 	return context.WithTimeout(parent, budget)
 }
@@ -62,10 +96,9 @@ func startPhase(parent context.Context, tracker *PhaseTracker, name string) (con
 func endPhase(phaseCtx context.Context, cancel context.CancelFunc,
 	tracker *PhaseTracker, log *progress.Logger, name string) {
 	if phaseCtx.Err() == context.DeadlineExceeded {
-		budget := phaseBudgets[name]
-		if budget == 0 {
-			budget = defaultPhaseBudget
-		}
+		// Only reachable when the phase had a deadline, so hasDeadline is true
+		// and budget is the real value used to set the timeout.
+		budget, _ := resolvePhaseBudget(name)
 		log.Warn("phase %s exceeded budget %s — continuing with partial results", name, budget)
 	}
 	cancel()

internal/telemetry/phase_deadline_test.go

@@ -0,0 +1,99 @@
+package telemetry
+
+import (
+	"context"
+	"testing"
+	"time"
+)
+
+// resolvePhaseBudget honors STEPSEC_PHASE_BUDGET_<NAME> > map entry > default,
+// and treats an env "0"/"off" or an explicit map 0 as "no deadline". Every
+// subtest sets the env var explicitly (to "" when it shouldn't apply) so an
+// inherited host/CI value can't make the result nondeterministic.
+func TestResolvePhaseBudget(t *testing.T) {
+	cases := []struct {
+		name        string
+		phase       string
+		env         string // value for STEPSEC_PHASE_BUDGET_<PHASE>
+		wantBudget  time.Duration
+		wantEnabled bool
+	}{
+		{"env override wins over map", "node_scan", "20m", 20 * time.Minute, true},
+		{"env off disables", "node_scan", "off", 0, false},
+		{"env 0 disables", "node_scan", "0", 0, false},
+		{"env junk ignored, falls to map", "node_scan", "junk", 15 * time.Minute, true},
+		{"env negative ignored, falls to map", "node_scan", "-5m", 15 * time.Minute, true},
+		{"map entry used when env empty", "ide_scan", "", 2 * time.Minute, true},
+		{"unlisted phase falls to default", "totally_made_up_phase", "", defaultPhaseBudget, true},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			// Hermetic: set the exact env key this phase reads (to "" when the
+			// case shouldn't have an override).
+			t.Setenv("STEPSEC_PHASE_BUDGET_"+upper(tc.phase), tc.env)
+			gotBudget, gotEnabled := resolvePhaseBudget(tc.phase)
+			if gotEnabled != tc.wantEnabled {
+				t.Fatalf("enabled = %v, want %v", gotEnabled, tc.wantEnabled)
+			}
+			if tc.wantEnabled && gotBudget != tc.wantBudget {
+				t.Errorf("budget = %v, want %v", gotBudget, tc.wantBudget)
+			}
+		})
+	}
+}
+
+// An explicit 0 in the phaseBudgets map disables the deadline — distinct from a
+// missing entry, which falls back to the default. (Copilot finding: the old
+// `if budget == 0 { budget = default }` conflated the two.)
+func TestResolvePhaseBudget_ExplicitZeroInMapDisables(t *testing.T) {
+	const phase = "zzz_explicit_zero_phase"
+	phaseBudgets[phase] = 0
+	t.Cleanup(func() { delete(phaseBudgets, phase) })
+	t.Setenv("STEPSEC_PHASE_BUDGET_"+upper(phase), "") // no env override
+
+	d, enabled := resolvePhaseBudget(phase)
+	if enabled || d != 0 {
+		t.Errorf("explicit map 0 must disable the deadline: got (%v, enabled=%v)", d, enabled)
+	}
+}
+
+// startPhase must produce a context with NO deadline when the budget is
+// disabled (only the parent scan deadline applies), and a deadline bounded by
+// the budget otherwise.
+func TestStartPhase_DeadlineWiring(t *testing.T) {
+	tracker := NewPhaseTracker()
+
+	t.Run("disabled budget → no phase deadline", func(t *testing.T) {
+		t.Setenv("STEPSEC_PHASE_BUDGET_NODE_SCAN", "off")
+		ctx, cancel := startPhase(context.Background(), tracker, "node_scan")
+		defer cancel()
+		if _, ok := ctx.Deadline(); ok {
+			t.Error("disabled phase budget must not set a context deadline")
+		}
+	})
+
+	t.Run("enabled budget → bounded deadline", func(t *testing.T) {
+		t.Setenv("STEPSEC_PHASE_BUDGET_NODE_SCAN", "20m")
+		ctx, cancel := startPhase(context.Background(), tracker, "node_scan")
+		defer cancel()
+		dl, ok := ctx.Deadline()
+		if !ok {
+			t.Fatal("enabled phase budget must set a context deadline")
+		}
+		if remaining := time.Until(dl); remaining <= 0 || remaining > 20*time.Minute {
+			t.Errorf("deadline %v out of expected ~20m window", remaining)
+		}
+	})
+}
+
+// upper uppercases an ASCII phase name for the env-var key, mirroring
+// strings.ToUpper without importing it into the test for one call.
+func upper(s string) string {
+	b := []byte(s)
+	for i := range b {
+		if b[i] >= 'a' && b[i] <= 'z' {
+			b[i] -= 'a' - 'A'
+		}
+	}
+	return string(b)
+}