fix(#6936): persist token telemetry when agents are cancelled

cmd/fullsend/main.go

@@ -18,9 +18,34 @@ type exitCoder interface {
 	ExitCode() int
 }
 
+// signalContext returns a context that is cancelled on the first
+// SIGINT/SIGTERM, plus a cleanup function that stops signal forwarding.
+// Subsequent signals are absorbed to prevent the default terminate handler
+// from killing the process before cleanup (metrics, telemetry) completes.
+//
+// GitHub Actions sends SIGINT, waits ~7.5 s, then SIGTERM; without absorbing
+// the second signal the default handler terminates the process before the
+// metrics/telemetry flush path finishes (#6936).
+//
+// signal.NotifyContext stops listening after the first signal, which
+// re-enables the default "terminate" handler for subsequent deliveries.
+// Keeping our own channel registered prevents that.
+func signalContext() (context.Context, func()) {
+	ctx, cancel := context.WithCancel(context.Background())
+	sigCh := make(chan os.Signal, 1)
+	signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
+	go func() {
+		<-sigCh // First signal: cancel the context.
+		cancel()
+		for range sigCh { // Subsequent signals: absorbed.
+		}
+	}()
+	return ctx, func() { signal.Stop(sigCh) }
+}
+
 func main() {
-	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
-	defer stop()
+	ctx, cleanup := signalContext()
+	defer cleanup()
 
 	if err := cli.Execute(ctx); err != nil {
 		fmt.Fprintf(os.Stderr, "Error: %v\n", err)

cmd/fullsend/main_test.go

@@ -0,0 +1,109 @@
+package main
+
+import (
+	"os"
+	"os/signal"
+	"syscall"
+	"testing"
+	"time"
+)
+
+// TestSignalContext_CancelsOnFirstSignal verifies that signalContext returns
+// a context that is cancelled when the process receives the first SIGINT.
+// This is the primary invariant for #6936: the first signal must cancel the
+// context so the cleanup path (metrics, telemetry) can run.
+func TestSignalContext_CancelsOnFirstSignal(t *testing.T) {
+	ctx, cleanup := signalContext()
+	defer cleanup()
+
+	// Send SIGINT to ourselves.
+	proc, err := os.FindProcess(os.Getpid())
+	if err != nil {
+		t.Fatalf("FindProcess: %v", err)
+	}
+	if err := proc.Signal(syscall.SIGINT); err != nil {
+		t.Fatalf("Signal: %v", err)
+	}
+
+	select {
+	case <-ctx.Done():
+		// Expected: context cancelled on first signal.
+	case <-time.After(2 * time.Second):
+		t.Fatal("context was not cancelled within 2s after SIGINT")
+	}
+}
+
+// TestSignalContext_AbsorbsSubsequentSignals verifies that after the first
+// signal cancels the context, further SIGINT/SIGTERM deliveries are absorbed
+// (they do not kill the process). This is the second invariant for #6936:
+// GitHub Actions sends SIGINT then SIGTERM ~7.5 s later; the second signal
+// must not trigger the default terminate handler.
+func TestSignalContext_AbsorbsSubsequentSignals(t *testing.T) {
+	ctx, cleanup := signalContext()
+	defer cleanup()
+
+	proc, err := os.FindProcess(os.Getpid())
+	if err != nil {
+		t.Fatalf("FindProcess: %v", err)
+	}
+
+	// First signal: cancels the context.
+	if err := proc.Signal(syscall.SIGINT); err != nil {
+		t.Fatalf("first Signal: %v", err)
+	}
+
+	select {
+	case <-ctx.Done():
+		// Good — context cancelled.
+	case <-time.After(2 * time.Second):
+		t.Fatal("context was not cancelled after first SIGINT")
+	}
+
+	// Second signal: must be absorbed (not kill the process).
+	if err := proc.Signal(syscall.SIGTERM); err != nil {
+		t.Fatalf("second Signal: %v", err)
+	}
+
+	// If we reach this point without dying, the signal was absorbed.
+	// Give a short window for the signal to be delivered and handled.
+	time.Sleep(50 * time.Millisecond)
+}
+
+// TestSignalContext_CleanupStopsForwarding verifies that after cleanup() is
+// called, signal forwarding to signalContext's channel is disabled: a signal
+// sent after cleanup() must not cancel ctx. A probe registered independently
+// via signal.Notify confirms the signal was actually delivered by the OS
+// (so the assertion isn't vacuously true because the signal never arrived),
+// without relying on ctx or signalContext's own (now-detached) channel and
+// without killing the test process.
+func TestSignalContext_CleanupStopsForwarding(t *testing.T) {
+	ctx, cleanup := signalContext()
+	cleanup()
+
+	probe := make(chan os.Signal, 1)
+	signal.Notify(probe, syscall.SIGINT)
+	defer signal.Stop(probe)
+
+	proc, err := os.FindProcess(os.Getpid())
+	if err != nil {
+		t.Fatalf("FindProcess: %v", err)
+	}
+	if err := proc.Signal(syscall.SIGINT); err != nil {
+		t.Fatalf("Signal: %v", err)
+	}
+
+	select {
+	case <-probe:
+		// Expected: the OS still delivered SIGINT to an independent
+		// receiver, so the process is alive and the signal was sent.
+	case <-time.After(2 * time.Second):
+		t.Fatal("probe did not receive SIGINT within 2s; signal was never delivered")
+	}
+
+	select {
+	case <-ctx.Done():
+		t.Fatal("context was cancelled after cleanup(); signal.Stop did not disable forwarding")
+	default:
+		// Expected: forwarding disabled, context still active.
+	}
+}

internal/cli/run.go

@@ -2253,6 +2253,15 @@ func runAgent(ctx context.Context, agentName, fullsendDir, outputBase, targetRep
 		// Accumulate behavioral metrics across iterations.
 		aggregateRunMetrics(&aggMetrics, &metrics, iteration)
 
+		if cancelled, cancelExitCode, cancelledErr := handleRunCancellation(
+			ctx, runErr, iteration, exitCode, rt.System(), rt.Name(),
+			&metrics, aggMetrics, runDir, agentSpan, attachIterationContent,
+			printer, lastIterElapsed,
+		); cancelled {
+			lastExitCode = cancelExitCode
+			return cancelledErr
+		}
+
 		if runErr != nil {
 			attachIterationContent("error")
 			finalizeAgentSpan(agentSpan, runErr, iteration, exitCode, rt.System(), rt.Name(), &metrics, "")
@@ -3676,6 +3685,52 @@ func transcriptErrorMessage(te agentruntime.TranscriptError) string {
 	return truncateStatusMsgTo(te.DisplayMessage(), maxSpanEventMsgLen)
 }
 
+// handleRunCancellation short-circuits the per-iteration loop in runAgent on
+// context cancellation: it persists partial metrics and finalizes the agent
+// span immediately, before extraction and validation that would be
+// pointless on a dead sandbox. GitHub Actions cancellation (SIGTERM)
+// terminates the process shortly after — writing metrics here ensures the
+// artifact upload step (if: always()) captures the partial usage data
+// (#6936).
+//
+// NOTE: TotalCostUSD will be zero in the persisted metrics because dollar
+// cost is only available from the terminal ResultEvent, which a cancelled
+// run never emits. Token counts (input, output, cache_read, cache_creation)
+// are captured via the deferred TokensEvent and will be non-zero. See #6936
+// for background.
+//
+// cancelled is false when ctx is still live, in which case the caller's
+// normal control flow continues unchanged; the other return values are
+// meaningless in that case.
+func handleRunCancellation(
+	ctx context.Context,
+	runErr error,
+	iteration, exitCode int,
+	system, runtimeName string,
+	metrics *agentruntime.RunMetrics,
+	aggMetrics aggregateMetrics,
+	runDir string,
+	agentSpan trace.Span,
+	attachIterationContent func(finishReason string),
+	printer *ui.Printer,
+	lastIterElapsed time.Duration,
+) (cancelled bool, lastExitCode int, err error) {
+	cancelErr := ctx.Err()
+	if cancelErr == nil {
+		return false, 0, nil
+	}
+	if runErr == nil {
+		runErr = cancelErr
+	}
+	attachIterationContent("error")
+	finalizeAgentSpan(agentSpan, runErr, iteration, exitCode, system, runtimeName, metrics, "")
+	printer.StepWarn(fmt.Sprintf("Run cancelled (iteration %d, %.1fs elapsed)", iteration, lastIterElapsed.Seconds()))
+	if writeErr := writeMetricsJSON(runDir, aggMetrics); writeErr != nil {
+		printer.StepWarn("Failed to write metrics.json: " + writeErr.Error())
+	}
+	return true, exitCode, fmt.Errorf("run cancelled (iteration %d): %w", iteration, runErr)
+}
+
 // finalizeAgentSpan records the end-of-iteration attributes and status on an
 // agent span and ends it. transcriptErr is non-empty when the transcript
 // reported a failure the process exit code did not (#2786): exit_code keeps