Tighten prepared execution credential handling

prepared_execution_contract_test.go

@@ -310,6 +310,75 @@ func TestPreparedExecutionConcurrentClaimAllowsOneGeneration(t *testing.T) {
 	}
 }
 
+func TestPreparedExecutionRunAndDiscardRaceHasOneWinner(t *testing.T) {
+	const attempts = 32
+
+	for i := 0; i < attempts; i++ {
+		client := &preparedRecordingClient{
+			response: &promptkit.GenerateResponse{Content: "ok"},
+		}
+		engine := newPreparedContractEngine(t, client, "race content")
+		prepared, err := engine.PrepareExecution(
+			context.Background(),
+			promptkit.RunRequest{PromptID: "prepared"},
+		)
+		if err != nil {
+			t.Fatalf("attempt %d prepare execution: %v", i, err)
+		}
+
+		start := make(chan struct{})
+		type outcome struct {
+			result *promptkit.RunResult
+			err    error
+		}
+		runOutcome := make(chan outcome, 1)
+		discardDone := make(chan struct{})
+
+		go func() {
+			<-start
+			result, runErr := engine.RunPrepared(context.Background(), prepared)
+			runOutcome <- outcome{result: result, err: runErr}
+		}()
+		go func() {
+			<-start
+			prepared.Discard()
+			close(discardDone)
+		}()
+
+		close(start)
+		runResult := <-runOutcome
+		<-discardDone
+		calls := len(client.snapshot())
+		switch {
+		case runResult.err == nil:
+			if runResult.result == nil || calls != 1 {
+				t.Fatalf(
+					"attempt %d run won with outcome=(%+v, %v), generation calls=%d",
+					i,
+					runResult.result,
+					runResult.err,
+					calls,
+				)
+			}
+		case errors.Is(runResult.err, promptkit.ErrInvalidRequest):
+			if runResult.result != nil || calls != 0 {
+				t.Fatalf(
+					"attempt %d discard won with outcome=(%+v, %v), generation calls=%d",
+					i,
+					runResult.result,
+					runResult.err,
+					calls,
+				)
+			}
+		default:
+			t.Fatalf("attempt %d unexpected run outcome=(%+v, %v)", i, runResult.result, runResult.err)
+		}
+		if prepared.Details().PromptID != "prepared" {
+			t.Fatalf("attempt %d details unavailable after race", i)
+		}
+	}
+}
+
 func TestPreparedExecutionDiscardAndFormattingDoNotExposePrivateState(t *testing.T) {
 	const (
 		directCredential = "pk-test-direct-credential-41f7"
@@ -355,9 +424,6 @@ func TestPreparedExecutionDiscardAndFormattingDoNotExposePrivateState(t *testing
 	if err != nil {
 		t.Fatalf("marshal opaque handle: %v", err)
 	}
-	if string(payload) != "{}" {
-		t.Fatalf("opaque handle JSON=%s, want {}", payload)
-	}
 	assertPreparedPrivateValuesAbsent(t, string(payload), directCredential, renderedContent)
 
 	detailsBefore := prepared.Details()