Preserve operation ownership in terminal observations

nested_error_contract_test.go

@@ -49,3 +49,51 @@ func TestNestedEngineErrorKeepsOwningOperation(t *testing.T) {
 		t.Fatalf("artifact error = %v", err)
 	}
 }
+
+func TestObservedNestedErrorsKeepTheParentOperation(t *testing.T) {
+	child := newSimpleEngine(t, "child")
+	_, childErr := child.Run(context.Background(), promptkit.RunRequest{PromptID: "missing"})
+	for _, mode := range []string{"run", "prepared", "repair", "artifact"} {
+		t.Run(mode, func(t *testing.T) {
+			var events []promptkit.ExecutionEvent
+			req := promptkit.RunRequest{PromptID: "prompt", Observer: func(e promptkit.ExecutionEvent) { events = append(events, e) }}
+			engine := newSimpleEngine(t, "parent", promptkit.WithLLMClient(&fakeLLMClient{err: childErr}))
+			want, owner := promptkit.ExecutionFailureGeneration, promptkit.ErrLLMGenerate
+			if mode == "artifact" {
+				engine = newSimpleEngine(t, "parent", promptkit.WithArtifactReader(nestedErrorReader{cause: childErr}))
+				req.Inputs = map[string]promptkit.ArtifactRef{"extra": promptkit.Inline("data")}
+				want, owner = promptkit.ExecutionFailureArtifactLoad, promptkit.ErrArtifactLoad
+			}
+			if mode == "repair" {
+				calls := 0
+				engine = newSimpleEngine(t, "parent", promptkit.WithLLMClient(observationClientFunc(func(context.Context, promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
+					calls++
+					if calls == 1 {
+						return &promptkit.GenerateResponse{Content: "invalid JSON"}, nil
+					}
+					return nil, childErr
+				})))
+				req.Validation = &promptkit.OutputContract{Format: promptkit.FormatJSON, ValidationMode: promptkit.ValidationJSON, RepairAttempts: 1}
+			}
+			var err error
+			if mode == "prepared" {
+				handle, prepErr := engine.PrepareExecution(context.Background(), req)
+				if prepErr != nil {
+					t.Fatal(prepErr)
+				}
+				_, err = engine.RunPrepared(context.Background(), handle)
+			} else {
+				_, err = engine.Run(context.Background(), req)
+			}
+			if !errors.Is(err, owner) || !errors.Is(err, childErr) {
+				t.Fatalf("lost owner/cause: %v", err)
+			}
+			if len(events) == 0 || events[len(events)-1].Failure != want {
+				t.Fatalf("terminal classification: %+v", events)
+			}
+			if mode != "artifact" && events[len(events)-2].Failure != want {
+				t.Fatalf("attempt classification: %+v", events)
+			}
+		})
+	}
+}