promptkit/coordinate_lifecycle_contract_test.go

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)
	}
}