Verify coordinated execution controls and observers

coordinate_controls_contract_test.go

@@ -0,0 +1,86 @@
+package promptkit_test
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"testing/synctest"
+	"time"
+
+	"go.maximumdirect.net/promptkit"
+)
+
+func TestCoordinateUsesIndependentOperationBudgets(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		engine := newBackendCapacityEngine(t, observationClientFunc(func(_ context.Context, req promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
+			if req.Prompt.SessionID == "continuation" {
+				time.Sleep(2 * time.Second)
+			}
+			return &promptkit.GenerateResponse{Content: "answer", Usage: promptkit.TokenUsage{PromptTokens: 2}}, nil
+		}), 1, capacityInt(0), nil)
+		parent := context.Background()
+		initial, err := engine.Run(parent, promptkit.RunRequest{PromptID: "prompt", CaptureExchange: true})
+		if err != nil {
+			t.Fatal(err)
+		}
+		var group promptkit.ExecutionGroup
+		var child context.Context
+		var details []promptkit.PreparedRun
+		var requestEvents, coordinatorEvents []promptkit.ExecutionEvent
+		start := time.Now()
+		result, err := engine.Coordinate(parent, promptkit.CoordinateRequest{Start: initial.Exchange, AdditionalTurns: 2, Observer: promptkit.ComposeObservers(group.Observe, func(e promptkit.ExecutionEvent) { coordinatorEvents = append(coordinatorEvents, e) }), Next: func(ctx context.Context, turn promptkit.ConversationTurn) (*promptkit.RunRequest, error) {
+			if turn.Result != nil {
+				if ctx.Err() != nil || child == nil || !errors.Is(child.Err(), context.Canceled) {
+					t.Fatal("child was not cleaned up before callback")
+				}
+				if _, err := engine.Run(ctx, promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Second}); err != nil {
+					t.Fatalf("assessment: %v", err)
+				}
+			}
+			if turn.RemainingTurns == 0 {
+				return nil, nil
+			}
+			return &promptkit.RunRequest{PromptID: "prompt", SessionID: "continuation", OperationTimeout: 3 * time.Second, Observer: func(e promptkit.ExecutionEvent) { requestEvents = append(requestEvents, e) }, OnPrepared: func(ctx context.Context, prepared promptkit.PreparedRun) {
+				child = ctx
+				details = append(details, prepared)
+			}}, nil
+		}})
+		if err != nil || result.Outcome != promptkit.CoordinationStopped || time.Since(start) != 4*time.Second || len(details) != 2 || len(details[0].Messages) != 3 || len(details[1].Messages) != 5 {
+			t.Fatalf("result=%+v error=%v details=%v elapsed=%v", result, err, len(details), time.Since(start))
+		}
+		summary := group.Snapshot()
+		if len(requestEvents) != 6 || len(coordinatorEvents) != 6 || len(summary.Members) != 2 || summary.Attempts != 2 || summary.Usage.PromptTokens.Value != 4 || !summary.TerminalCoverageComplete {
+			t.Fatalf("summary=%+v request=%d coordinator=%d", summary, len(requestEvents), len(coordinatorEvents))
+		}
+	})
+}
+
+func TestCoordinateTimeoutLeavesParentAndSiblingUsable(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		engine := newBackendCapacityEngine(t, observationClientFunc(func(ctx context.Context, req promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
+			if req.Prompt.SessionID == "slow" {
+				<-ctx.Done()
+				return nil, ctx.Err()
+			}
+			return &promptkit.GenerateResponse{Content: "answer"}, nil
+		}), 1, capacityInt(0), nil)
+		parent := context.Background()
+		initial, err := engine.Run(parent, promptkit.RunRequest{PromptID: "prompt", CaptureExchange: true})
+		if err != nil {
+			t.Fatal(err)
+		}
+		var details promptkit.PreparedRun
+		calls := 0
+		result, err := engine.Coordinate(parent, promptkit.CoordinateRequest{Start: initial.Exchange, AdditionalTurns: 1, Next: func(context.Context, promptkit.ConversationTurn) (*promptkit.RunRequest, error) {
+			calls++
+			return &promptkit.RunRequest{PromptID: "prompt", SessionID: "slow", OperationTimeout: time.Second, OnPrepared: func(_ context.Context, p promptkit.PreparedRun) { details = p }}, nil
+		}})
+		if !errors.Is(err, promptkit.ErrOperationTimeout) || !errors.Is(err, context.DeadlineExceeded) || calls != 1 || result.Outcome != "" || result.Undelivered != nil || details.ParentRunID != initial.RunID || parent.Err() != nil {
+			t.Fatalf("result=%+v error=%v callbacks=%d", result, err, calls)
+		}
+		sibling, err := engine.Run(parent, promptkit.RunRequest{PromptID: "prompt", ContinueFrom: initial.Exchange})
+		if err != nil || sibling.Exchange == nil {
+			t.Fatalf("sibling=%v error=%v", sibling, err)
+		}
+	})
+}