Preserve cancellation in aggregate smoke tests

integrations/ckprompt/smoke_aggregate_test.go

@@ -3,9 +3,11 @@ package ckprompt
 import (
 	"bytes"
 	"context"
+	"encoding/json"
 	"errors"
 	"io"
 	"reflect"
+	"strings"
 	"testing"
 
 	"go.maximumdirect.net/ckconfig"
@@ -109,6 +111,79 @@ func TestSmokeAggregateCoordinatesConcreteRequestsAndRetainsErrors(t *testing.T)
 	}
 }
 
+type cancelProgressWriter struct {
+	bytes.Buffer
+	cancel context.CancelFunc
+}
+
+func (writer *cancelProgressWriter) Write(p []byte) (int, error) {
+	writer.cancel()
+	return writer.Buffer.Write(p)
+}
+
+func TestSmokeAggregateCancellationDuringProgress(t *testing.T) {
+	providerErr := errors.New("provider failed")
+	for _, test := range []struct {
+		name     string
+		err      error
+		cancelAt int
+	}{
+		{name: "partial success", cancelAt: 1},
+		{name: "partial provider failure", err: providerErr, cancelAt: 1},
+		{name: "last completed test", cancelAt: 4},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			var kinds []promptkit.SmokeTestKind
+			app := newPromptTestApp(t, Command(Suite[struct{}, struct{}]{
+				SmokeTest: func(_ context.Context, request SmokeRequest, _, _ struct{}) (commandkit.Acquired[*promptkit.SmokeTestReport], error) {
+					kinds = append(kinds, request.Request.Kind)
+					acquired := commandkit.Acquired[*promptkit.SmokeTestReport]{Value: &promptkit.SmokeTestReport{Kind: request.Request.Kind, Completed: true}}
+					if len(kinds) == test.cancelAt {
+						acquired.Diagnostics = []commandkit.Diagnostic{{Severity: commandkit.SeverityProgress, Write: func(w io.Writer) error { _, err := io.WriteString(w, "progress evidence\n"); return err }}}
+						acquired.Explanation = commandkit.Explanation{Message: "provider explanation"}
+						return acquired, test.err
+					}
+					return acquired, nil
+				},
+			}))
+			var output bytes.Buffer
+			diagnostics := &cancelProgressWriter{cancel: cancel}
+			outcome := app.Execute(ctx, []string{"promptkit", "test", "all", "--progress", "--format=json"}, commandkit.Streams{Out: &output, Err: diagnostics})
+			if len(kinds) != test.cancelAt || !strings.Contains(diagnostics.String(), "progress evidence") {
+				t.Fatalf("kinds=%v diagnostics=%q", kinds, diagnostics.String())
+			}
+			if test.cancelAt < 4 {
+				if outcome.Kind != commandkit.OutcomeOperational || !errors.Is(outcome.Cause, context.Canceled) || (test.err != nil && (!errors.Is(outcome.Cause, test.err) || !strings.Contains(diagnostics.String(), "provider explanation"))) {
+					t.Fatalf("partial outcome=%#v diagnostics=%q", outcome, diagnostics.String())
+				}
+			} else if outcome.Kind != commandkit.OutcomeSuccess || outcome.Cause != nil {
+				t.Fatalf("completed outcome=%#v", outcome)
+			}
+			var report struct {
+				Tests []struct {
+					Kind   string `json:"kind"`
+					Report *struct {
+						Completed bool `json:"completed"`
+					} `json:"report"`
+				} `json:"tests"`
+			}
+			if err := json.Unmarshal(output.Bytes(), &report); err != nil {
+				t.Fatal(err)
+			}
+			if len(report.Tests) != len(kinds) {
+				t.Fatalf("report=%s", output.String())
+			}
+			for i, entry := range report.Tests {
+				if entry.Kind != string(kinds[i]) || entry.Report == nil || !entry.Report.Completed {
+					t.Fatalf("report entry=%#v", entry)
+				}
+			}
+		})
+	}
+}
+
 func TestSmokeAggregateCancellationRetainsPartialResults(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	t.Cleanup(cancel)

integrations/ckprompt/smoke_command.go

@@ -173,6 +173,9 @@ func runSmokeAggregate[I, S any](invocation commandkit.Invocation, suite Suite[I
 			}
 		}
 		if err := promptKitContextError(invocation.Context); err != nil {
+			if len(results) < len(kinds) && !errors.Is(errors.Join(operationErrors...), err) {
+				operationErrors = append(operationErrors, err)
+			}
 			break
 		}
 	}