Cover coordinator lifecycle boundaries

coordinate_lifecycle_contract_test.go

@@ -0,0 +1,170 @@
+package promptkit_test
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"testing"
+	"time"
+
+	"go.maximumdirect.net/promptkit"
+)
+
+func TestCoordinateReleasesCapacityBeforeCallback(t *testing.T) {
+	client := &preparedRecordingClient{responses: []*promptkit.GenerateResponse{{Content: "initial"}, {Content: ""}, {Content: " "}, {Content: "validator"}}}
+	engine := newBackendCapacityEngine(t, client, 1, capacityInt(0), nil)
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	initial, err := engine.Run(ctx, promptkit.RunRequest{PromptID: "prompt", CaptureExchange: true})
+	if err != nil {
+		t.Fatal(err)
+	}
+	calls := 0
+	result, err := engine.Coordinate(ctx, promptkit.CoordinateRequest{Start: initial.Exchange, AdditionalTurns: 1, Next: func(ctx context.Context, turn promptkit.ConversationTurn) (*promptkit.RunRequest, error) {
+		calls++
+		if turn.Result == nil {
+			return &promptkit.RunRequest{PromptID: "prompt", Validation: &promptkit.OutputContract{Format: promptkit.FormatText, ValidationMode: promptkit.ValidationBasic, RepairAttempts: 1}}, nil
+		}
+		if turn.Result.Validation.Status != promptkit.ValidationFailed || turn.Result.Validation.RepairAttempts != 1 || turn.RemainingTurns != 0 {
+			t.Fatalf("delivered turn = %#v", turn)
+		}
+		_, err := engine.Run(ctx, promptkit.RunRequest{PromptID: "prompt"})
+		return nil, err
+	}})
+	if err != nil || result.Outcome != promptkit.CoordinationStopped || calls != 2 || len(client.snapshot()) != 4 {
+		t.Fatalf("outcome=%#v error=%v callbacks=%d calls=%d", result, err, calls, len(client.snapshot()))
+	}
+}
+
+func TestCoordinatePreservesCapacityAndProviderErrors(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	sourceEngine := newSimpleEngine(t, "source", promptkit.WithLLMClient(&fakeLLMClient{response: &promptkit.GenerateResponse{Content: "answer"}}))
+	initial, err := sourceEngine.Run(ctx, promptkit.RunRequest{PromptID: "prompt", CaptureExchange: true})
+	if err != nil {
+		t.Fatal(err)
+	}
+	release := make(chan struct{})
+	client := newCapacityGateClient(release, 1)
+	engine := newBackendCapacityEngine(t, client, 1, capacityInt(0), nil)
+	done := make(chan capacityRunResult, 1)
+	go runCapacityRequest(engine, ctx, promptkit.RunRequest{PromptID: "prompt"}, done)
+	select {
+	case <-client.started:
+	case <-ctx.Done():
+		t.Fatal(ctx.Err())
+	}
+	request := promptkit.CoordinateRequest{Start: initial.Exchange, AdditionalTurns: 1, Next: func(context.Context, promptkit.ConversationTurn) (*promptkit.RunRequest, error) {
+		return &promptkit.RunRequest{PromptID: "prompt"}, nil
+	}}
+	_, err = engine.Coordinate(ctx, request)
+	var capacity *promptkit.CapacityError
+	if !errors.As(err, &capacity) || capacity.BackendID != "limited" {
+		t.Fatalf("capacity error = %v", err)
+	}
+	close(release)
+	if finished := awaitCapacityRun(t, done); finished.err != nil {
+		t.Fatal(finished.err)
+	}
+	providerEngine := newBuiltInGenerationErrorEngine(t, http.StatusUnprocessableEntity, `{"error":{"message":"synthetic"}}`)
+	_, err = providerEngine.Coordinate(ctx, request)
+	var provider *promptkit.GenerationError
+	if !errors.As(err, &provider) || provider.StatusCode() != http.StatusUnprocessableEntity {
+		t.Fatalf("provider error = %v", err)
+	}
+}
+
+func TestCoordinateRetriesAndRepairsDeliverOnce(t *testing.T) {
+	initialEngine := newSimpleEngine(t, "source", promptkit.WithLLMClient(&fakeLLMClient{response: &promptkit.GenerateResponse{Content: "answer"}}))
+	initial, err := initialEngine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", CaptureExchange: true})
+	if err != nil {
+		t.Fatal(err)
+	}
+	responses := []*http.Response{observedHTTPResponse(503, `{}`), observedHTTPResponse(200, `{"choices":[{"message":{"content":"invalid"}}]}`), observedHTTPResponse(200, `{"choices":[{"message":{"content":"{\"ok\":true}"}}]}`)}
+	calls := 0
+	one := 1
+	config := simpleConfig("follow-up")
+	config.TransientRetries = &one
+	config.HTTPClient = &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
+		response := responses[calls]
+		calls++
+		return response, nil
+	})}
+	engine, err := promptkit.NewEngine(config)
+	if err != nil {
+		t.Fatal(err)
+	}
+	callbacks := 0
+	result, err := engine.Coordinate(context.Background(), promptkit.CoordinateRequest{Start: initial.Exchange, AdditionalTurns: 1, Next: func(_ context.Context, turn promptkit.ConversationTurn) (*promptkit.RunRequest, error) {
+		callbacks++
+		if turn.Result == nil {
+			return &promptkit.RunRequest{PromptID: "prompt", Validation: &promptkit.OutputContract{Format: promptkit.FormatJSON, ValidationMode: promptkit.ValidationJSON, RepairAttempts: 1}}, nil
+		}
+		if !turn.Result.Validation.IsValid || turn.RemainingTurns != 0 {
+			t.Fatalf("turn = %#v", turn)
+		}
+		return nil, nil
+	}})
+	if err != nil || result.Outcome != promptkit.CoordinationStopped || callbacks != 2 || calls != 3 {
+		t.Fatalf("result=%#v error=%v callbacks=%d attempts=%d", result, err, callbacks, calls)
+	}
+}
+
+func TestCoordinateCancellationDelivery(t *testing.T) {
+	for _, cancelInside := range []bool{false, true} {
+		t.Run(map[bool]string{false: "before delivery", true: "inside callback"}[cancelInside], func(t *testing.T) {
+			client := &fakeLLMClient{response: &promptkit.GenerateResponse{Content: "answer"}}
+			engine := newSimpleEngine(t, "question", promptkit.WithLLMClient(client))
+			initial, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", CaptureExchange: true})
+			if err != nil {
+				t.Fatal(err)
+			}
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			calls := 0
+			result, err := engine.Coordinate(ctx, promptkit.CoordinateRequest{Start: initial.Exchange, AdditionalTurns: 1, Next: func(_ context.Context, turn promptkit.ConversationTurn) (*promptkit.RunRequest, error) {
+				calls++
+				if turn.Result != nil {
+					cancel()
+					return nil, nil
+				}
+				return &promptkit.RunRequest{PromptID: "prompt", Observer: func(event promptkit.ExecutionEvent) {
+					if !cancelInside && event.Kind == promptkit.ExecutionFinished {
+						cancel()
+					}
+				}}, nil
+			}})
+			if !errors.Is(err, context.Canceled) || result.Outcome != "" || (result.Undelivered != nil) == cancelInside {
+				t.Fatalf("result=%#v error=%v callbacks=%d", result, err, calls)
+			}
+			if !cancelInside && (calls != 1 || result.Undelivered.Exchange == nil) {
+				t.Fatalf("undelivered result=%#v calls=%d", result, calls)
+			}
+			if cancelInside && calls != 2 {
+				t.Fatalf("callbacks=%d", calls)
+			}
+		})
+	}
+}
+
+func TestCoordinateFailureLeavesSiblingUsable(t *testing.T) {
+	client := &fakeLLMClient{response: &promptkit.GenerateResponse{Content: "answer"}}
+	engine := newSimpleEngine(t, "question", promptkit.WithLLMClient(client))
+	ctx := context.Background()
+	initial, err := engine.Run(ctx, promptkit.RunRequest{PromptID: "prompt", CaptureExchange: true})
+	if err != nil {
+		t.Fatal(err)
+	}
+	sentinel := errors.New("provider")
+	client.err = sentinel
+	result, err := engine.Coordinate(ctx, promptkit.CoordinateRequest{Start: initial.Exchange, AdditionalTurns: 1, Next: func(context.Context, promptkit.ConversationTurn) (*promptkit.RunRequest, error) {
+		return &promptkit.RunRequest{PromptID: "prompt"}, nil
+	}})
+	if !errors.Is(err, sentinel) || !errors.Is(err, promptkit.ErrLLMGenerate) || result.Outcome != "" || result.Undelivered != nil {
+		t.Fatalf("failure=%#v, %v", result, err)
+	}
+	client.err = nil
+	if sibling, err := engine.Run(ctx, promptkit.RunRequest{PromptID: "prompt", ContinueFrom: initial.Exchange}); err != nil || sibling.Exchange == nil {
+		t.Fatalf("sibling=%#v, %v", sibling, err)
+	}
+}