promptkit/execution_budget_contract_test.go

package promptkit_test

import (
	"context"
	"errors"
	"net/http"
	"testing"
	"testing/synctest"
	"time"

	"go.maximumdirect.net/promptkit"
)

func TestOperationBudgetCoversPreparationAndReleasesAdmission(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		release := make(chan struct{})
		reader := &capacityArtifactReader{entered: make(chan struct{}, 2), release: release}
		client := newCapacityGateClient(closedCapacityChannel(), 2)
		engine := newBackendCapacityEngine(t, client, 1, capacityInt(0), reader)
		parent := context.Background()
		req := capacityInputRequest("http://example.invalid/v1")
		req.OperationTimeout = time.Second
		_, err := engine.Run(parent, req)
		if !errors.Is(err, promptkit.ErrOperationTimeout) || !errors.Is(err, context.DeadlineExceeded) || !errors.Is(err, promptkit.ErrArtifactLoad) || parent.Err() != nil {
			t.Fatalf("preparation error=%v parent=%v", err, parent.Err())
		}
		close(release)
		if result, err := engine.Run(parent, req); err != nil || result == nil {
			t.Fatalf("capacity not released: %v", err)
		}
	})
}

func TestOperationBudgetCoversGenerationCapacityWait(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		release := make(chan struct{})
		client := newCapacityGateClient(release, 2)
		engine := newBackendCapacityEngine(t, client, 1, capacityInt(1), nil)
		results := make(chan capacityRunResult, 1)
		go runCapacityRequest(engine, context.Background(), promptkit.RunRequest{PromptID: "prompt"}, results)
		<-client.started
		_, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Second})
		if !errors.Is(err, promptkit.ErrOperationTimeout) || !errors.Is(err, promptkit.ErrLLMGenerate) || !errors.Is(err, context.DeadlineExceeded) {
			t.Fatalf("capacity wait error=%v", err)
		}
		close(release)
		if result := <-results; result.err != nil {
			t.Fatal(result.err)
		}
		if _, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt"}); err != nil {
			t.Fatal(err)
		}
	})
}

func TestOperationBudgetInterruptsRetryAndRepair(t *testing.T) {
	for _, repair := range []bool{false, true} {
		t.Run(map[bool]string{false: "retry", true: "repair"}[repair], func(t *testing.T) {
			synctest.Test(t, func(t *testing.T) {
				calls := 0
				config := simpleConfig("message")
				one := 1
				config.TransientRetries = &one
				config.HTTPClient = &http.Client{Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
					calls++
					if repair && calls == 1 {
						return retryHTTPResponse(200, "not json"), nil
					}
					if repair {
						<-r.Context().Done()
						return nil, r.Context().Err()
					}
					response := retryHTTPResponse(503, "")
					response.Header = make(http.Header)
					response.Header.Set("Retry-After", "5")
					return response, nil
				})}
				engine, err := promptkit.NewEngine(config)
				if err != nil {
					t.Fatal(err)
				}
				req := simpleGenerationRequest()
				req.OperationTimeout = time.Second
				if repair {
					req.Validation = &promptkit.OutputContract{Format: promptkit.FormatJSON, ValidationMode: promptkit.ValidationJSON, RepairAttempts: 1}
				}
				_, err = engine.Run(context.Background(), req)
				if !errors.Is(err, promptkit.ErrOperationTimeout) || !errors.Is(err, promptkit.ErrLLMGenerate) || !errors.Is(err, context.DeadlineExceeded) {
					t.Fatalf("error=%v calls=%d", err, calls)
				}
			})
		})
	}
}

func TestOperationOutcomeFrozenBeforeTerminalObserver(t *testing.T) {
	for _, split := range []bool{false, true} {
		t.Run(map[bool]string{false: "run", true: "prepared"}[split], func(t *testing.T) {
			synctest.Test(t, func(t *testing.T) {
				engine := observationEngine(t, observationClientFunc(func(context.Context, promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
					return &promptkit.GenerateResponse{Content: "ok"}, nil
				}))
				var terminal promptkit.ExecutionEvent
				req := promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Second, Observer: func(event promptkit.ExecutionEvent) {
					if event.Kind == promptkit.ExecutionFinished {
						time.Sleep(2 * time.Second)
						terminal = event
					}
				}}
				var result *promptkit.RunResult
				var err error
				if split {
					var handle *promptkit.PreparedExecution
					handle, err = engine.PrepareExecution(context.Background(), req)
					if err == nil {
						result, err = engine.RunPrepared(context.Background(), handle)
					}
				} else {
					result, err = engine.Run(context.Background(), req)
				}
				if err != nil || result == nil || terminal.Outcome != promptkit.ExecutionOutcomeSuccess {
					t.Fatalf("result=%v error=%v terminal=%+v", result, err, terminal)
				}
			})
		})
	}
}

func TestNotificationConsumesOperationBudget(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		engine := observationEngine(t, observationClientFunc(func(context.Context, promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
			t.Fatal("generation ran")
			return nil, nil
		}))
		_, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Second, OnPrepared: func(ctx context.Context, _ promptkit.PreparedRun) { <-ctx.Done() }})
		if !errors.Is(err, promptkit.ErrOperationTimeout) || !errors.Is(err, promptkit.ErrLLMGenerate) || !errors.Is(err, context.DeadlineExceeded) {
			t.Fatal(err)
		}
	})
}

func TestOperationTimeoutDoesNotRelabelOtherFailures(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		failure := errors.New("independent failure")
		engine := observationEngine(t, observationClientFunc(func(ctx context.Context, _ promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
			<-ctx.Done()
			return nil, failure
		}))
		var terminal promptkit.ExecutionEvent
		_, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Second, Observer: func(event promptkit.ExecutionEvent) {
			if event.Kind == promptkit.ExecutionFinished {
				terminal = event
			}
		}})
		if !errors.Is(err, failure) || errors.Is(err, promptkit.ErrOperationTimeout) || terminal.Failure != promptkit.ExecutionFailureGeneration {
			t.Fatalf("error=%v terminal=%+v", err, terminal)
		}
	})
}

func TestOperationTimeoutTerminalObserverSeesFinalFailure(t *testing.T) {
	for _, split := range []bool{false, true} {
		t.Run(map[bool]string{false: "run", true: "prepared"}[split], func(t *testing.T) {
			synctest.Test(t, func(t *testing.T) {
				engine := observationEngine(t, observationClientFunc(func(ctx context.Context, _ promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
					<-ctx.Done()
					return nil, ctx.Err()
				}))
				var terminal promptkit.ExecutionEvent
				req := promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Second, Observer: func(event promptkit.ExecutionEvent) {
					if event.Kind == promptkit.ExecutionFinished {
						time.Sleep(time.Second)
						terminal = event
					}
				}}
				var err error
				if split {
					var handle *promptkit.PreparedExecution
					handle, err = engine.PrepareExecution(context.Background(), req)
					if err == nil {
						_, err = engine.RunPrepared(context.Background(), handle)
					}
				} else {
					_, err = engine.Run(context.Background(), req)
				}
				if !errors.Is(err, promptkit.ErrOperationTimeout) || !errors.Is(err, context.DeadlineExceeded) || terminal.Failure != promptkit.ExecutionFailureDeadline {
					t.Fatalf("error=%v terminal=%+v", err, terminal)
				}
			})
		})
	}
}