774 lines
25 KiB
Go
774 lines
25 KiB
Go
package promptkit_test
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"os"
|
|
"reflect"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"testing/fstest"
|
|
"time"
|
|
|
|
"gitea.maximumdirect.net/eric/promptkit"
|
|
)
|
|
|
|
func TestPreparedExecutionFreezesSourcesAndReturnsIndependentDetails(t *testing.T) {
|
|
promptSource := preparedPromptSource("original")
|
|
profileSource := preparedProfileSource("original-model")
|
|
schemaSource := preparedSchemaSource()
|
|
reader := &mutablePreparedArtifactReader{
|
|
body: "original artifact",
|
|
hash: "original-input-hash",
|
|
}
|
|
client := &preparedRecordingClient{
|
|
response: &promptkit.GenerateResponse{Content: `{"value":3}`},
|
|
}
|
|
engine, err := promptkit.NewEngine(
|
|
promptkit.Config{},
|
|
promptkit.WithPromptFS(promptSource, "."),
|
|
promptkit.WithProfileFS(profileSource, "."),
|
|
promptkit.WithSchemaFS(schemaSource, "."),
|
|
promptkit.WithArtifactReader(reader),
|
|
promptkit.WithLLMClient(client),
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("construct engine: %v", err)
|
|
}
|
|
|
|
temperature := 0.25
|
|
extraParams := map[string]any{
|
|
"nested": map[string]any{"source": "original"},
|
|
}
|
|
request := promptkit.RunRequest{
|
|
PromptID: "prepared",
|
|
Inputs: map[string]promptkit.ArtifactRef{
|
|
"input": promptkit.Inline("original request input"),
|
|
},
|
|
Vars: map[string]string{"label": "original variable"},
|
|
Execution: &promptkit.ExecutionTargetOverride{
|
|
Temperature: &temperature,
|
|
ExtraParams: extraParams,
|
|
},
|
|
}
|
|
preparationContext, cancelPreparation := context.WithCancel(context.Background())
|
|
prepared, err := engine.PrepareExecution(preparationContext, request)
|
|
if err != nil {
|
|
t.Fatalf("prepare execution: %v", err)
|
|
}
|
|
cancelPreparation()
|
|
|
|
request.PromptID = "changed"
|
|
request.Inputs["input"] = promptkit.Inline("changed request input")
|
|
request.Vars["label"] = "changed variable"
|
|
temperature = 1.5
|
|
extraParams["nested"].(map[string]any)["source"] = "changed"
|
|
promptSource["prompt.yaml"] = &fstest.MapFile{Data: []byte(`id: changed`)}
|
|
profileSource["profile.yaml"] = &fstest.MapFile{Data: []byte(`id: changed`)}
|
|
schemaSource["schema.json"] = &fstest.MapFile{Data: []byte(`{
|
|
"$schema": "https://json-schema.org/draft/2020-12/schema",
|
|
"title": "changed root",
|
|
"type": "string"
|
|
}`)}
|
|
schemaSource["value.json"] = &fstest.MapFile{Data: []byte(`{
|
|
"$schema": "https://json-schema.org/draft/2020-12/schema",
|
|
"type": "string"
|
|
}`)}
|
|
reader.set("changed artifact", "changed-input-hash")
|
|
|
|
first := prepared.Details()
|
|
first.Messages[0].Content = "changed details"
|
|
first.InputHashes["input"] = "changed-details-hash"
|
|
first.EffectiveModelParams.ExtraParams["nested"].(map[string]any)["source"] = "changed details"
|
|
first.StructuredOutput.JSONSchema.Schema.(map[string]any)["title"] = "changed details"
|
|
|
|
second := prepared.Details()
|
|
if second.Messages[0].Content != "Input=original artifact Label=original variable" {
|
|
t.Fatalf("details message changed: %q", second.Messages[0].Content)
|
|
}
|
|
if second.InputHashes["input"] != "original-input-hash" {
|
|
t.Fatalf("details input hash changed: %q", second.InputHashes["input"])
|
|
}
|
|
if second.EffectiveModelParams.Model != "original-model" ||
|
|
second.EffectiveModelParams.Temperature != 0.25 ||
|
|
second.EffectiveModelParams.ExtraParams["nested"].(map[string]any)["source"] != "original" {
|
|
t.Fatalf("details target changed: %+v", second.EffectiveModelParams)
|
|
}
|
|
schema := second.StructuredOutput.JSONSchema.Schema.(map[string]any)
|
|
if schema["title"] != "original root" {
|
|
t.Fatalf("details schema changed: %#v", schema)
|
|
}
|
|
|
|
result, err := engine.RunPrepared(context.Background(), prepared)
|
|
if err != nil {
|
|
t.Fatalf("run prepared after preparation-context cancellation: %v", err)
|
|
}
|
|
if result.Validation.Status != promptkit.ValidationPassed || !result.Validation.IsValid {
|
|
t.Fatalf("frozen schema did not validate original output: %+v", result.Validation)
|
|
}
|
|
if reader.callCount() != 1 {
|
|
t.Fatalf("execution reopened artifact source: calls=%d", reader.callCount())
|
|
}
|
|
|
|
requests := client.snapshot()
|
|
if len(requests) != 1 {
|
|
t.Fatalf("generation calls=%d, want 1", len(requests))
|
|
}
|
|
generated := requests[0]
|
|
if generated.Prompt.Messages[0].Content != second.Messages[0].Content ||
|
|
generated.Target.Model != second.EffectiveModelParams.Model ||
|
|
!reflect.DeepEqual(generated.Target.ExtraParams, second.EffectiveModelParams.ExtraParams) ||
|
|
!reflect.DeepEqual(generated.StructuredOutput, second.StructuredOutput) {
|
|
t.Fatalf("generation did not use frozen details:\nrequest=%+v\ndetails=%+v", generated, second)
|
|
}
|
|
if result.PromptID != second.PromptID ||
|
|
result.PromptVersion != second.PromptVersion ||
|
|
result.PromptHash != second.PromptHash ||
|
|
result.SessionID != second.SessionID ||
|
|
result.RenderedPromptHash != second.RenderedPromptHash ||
|
|
result.SelectedProfileID != second.SelectedProfileID ||
|
|
result.SelectedBackendID != second.SelectedBackendID ||
|
|
!reflect.DeepEqual(result.EffectiveModelParams, second.EffectiveModelParams) ||
|
|
!reflect.DeepEqual(result.InputHashes, second.InputHashes) {
|
|
t.Fatalf("result provenance does not match details:\nresult=%+v\ndetails=%+v", result, second)
|
|
}
|
|
}
|
|
|
|
func TestPreparedExecutionLifecycleAndEngineBinding(t *testing.T) {
|
|
ownerClient := &preparedRecordingClient{
|
|
response: &promptkit.GenerateResponse{Content: "ok"},
|
|
}
|
|
owner := newPreparedContractEngine(t, ownerClient, "owner content")
|
|
foreign := newPreparedContractEngine(t, &preparedRecordingClient{
|
|
response: &promptkit.GenerateResponse{Content: "unexpected"},
|
|
}, "foreign content")
|
|
|
|
prepared, err := owner.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prepared"})
|
|
if err != nil {
|
|
t.Fatalf("prepare execution: %v", err)
|
|
}
|
|
copied := *prepared
|
|
|
|
var nilEngine *promptkit.Engine
|
|
if result, err := nilEngine.RunPrepared(context.Background(), prepared); result != nil ||
|
|
!errors.Is(err, promptkit.ErrInvalidConfig) {
|
|
t.Fatalf("nil engine result=(%+v, %v), want ErrInvalidConfig", result, err)
|
|
}
|
|
if result, err := foreign.RunPrepared(context.Background(), prepared); result != nil ||
|
|
!errors.Is(err, promptkit.ErrInvalidRequest) {
|
|
t.Fatalf("foreign engine result=(%+v, %v), want ErrInvalidRequest", result, err)
|
|
}
|
|
if result, err := owner.RunPrepared(context.Background(), nil); result != nil ||
|
|
!errors.Is(err, promptkit.ErrInvalidRequest) {
|
|
t.Fatalf("nil handle result=(%+v, %v), want ErrInvalidRequest", result, err)
|
|
}
|
|
if result, err := owner.RunPrepared(context.Background(), &promptkit.PreparedExecution{}); result != nil ||
|
|
!errors.Is(err, promptkit.ErrInvalidRequest) {
|
|
t.Fatalf("zero handle result=(%+v, %v), want ErrInvalidRequest", result, err)
|
|
}
|
|
|
|
result, err := owner.RunPrepared(context.Background(), &copied)
|
|
if err != nil || result == nil {
|
|
t.Fatalf("owner run prepared=(%+v, %v), want success", result, err)
|
|
}
|
|
for name, handle := range map[string]*promptkit.PreparedExecution{
|
|
"original": prepared,
|
|
"copy": &copied,
|
|
} {
|
|
if result, err := owner.RunPrepared(context.Background(), handle); result != nil ||
|
|
!errors.Is(err, promptkit.ErrInvalidRequest) {
|
|
t.Fatalf("%s reused handle result=(%+v, %v), want ErrInvalidRequest", name, result, err)
|
|
}
|
|
if handle.Details().PromptID != "prepared" {
|
|
t.Fatalf("%s details unavailable after execution", name)
|
|
}
|
|
}
|
|
if len(ownerClient.snapshot()) != 1 {
|
|
t.Fatalf("owner generation calls=%d, want 1", len(ownerClient.snapshot()))
|
|
}
|
|
|
|
collaboratorFailure := errors.New("prepared collaborator failure")
|
|
failingClient := &preparedRecordingClient{err: collaboratorFailure}
|
|
failingEngine := newPreparedContractEngine(t, failingClient, "failure content")
|
|
failing, err := failingEngine.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prepared"})
|
|
if err != nil {
|
|
t.Fatalf("prepare failing execution: %v", err)
|
|
}
|
|
if result, err := failingEngine.RunPrepared(context.Background(), failing); result != nil ||
|
|
!errors.Is(err, promptkit.ErrLLMGenerate) ||
|
|
!errors.Is(err, collaboratorFailure) {
|
|
t.Fatalf("generation failure result=(%+v, %v), want public and collaborator identities", result, err)
|
|
}
|
|
if result, err := failingEngine.RunPrepared(context.Background(), failing); result != nil ||
|
|
!errors.Is(err, promptkit.ErrInvalidRequest) {
|
|
t.Fatalf("failed execution was reusable: result=(%+v, %v)", result, err)
|
|
}
|
|
|
|
cancellationRelease := make(chan struct{})
|
|
cancellationStarted := make(chan struct{}, 1)
|
|
cancelingEngine := newPreparedContractEngine(t, &preparedRecordingClient{
|
|
response: &promptkit.GenerateResponse{Content: "unexpected"},
|
|
started: cancellationStarted,
|
|
release: cancellationRelease,
|
|
}, "cancellation content")
|
|
canceling, err := cancelingEngine.PrepareExecution(
|
|
context.Background(),
|
|
promptkit.RunRequest{PromptID: "prepared"},
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("prepare canceled execution: %v", err)
|
|
}
|
|
executionContext, cancelExecution := context.WithCancel(context.Background())
|
|
type canceledOutcome struct {
|
|
result *promptkit.RunResult
|
|
err error
|
|
}
|
|
canceledResult := make(chan canceledOutcome, 1)
|
|
go func() {
|
|
result, runErr := cancelingEngine.RunPrepared(executionContext, canceling)
|
|
canceledResult <- canceledOutcome{result: result, err: runErr}
|
|
}()
|
|
select {
|
|
case <-cancellationStarted:
|
|
case <-time.After(5 * time.Second):
|
|
t.Fatal("timed out waiting for cancelable generation")
|
|
}
|
|
cancelExecution()
|
|
select {
|
|
case outcome := <-canceledResult:
|
|
if outcome.result != nil ||
|
|
!errors.Is(outcome.err, promptkit.ErrLLMGenerate) ||
|
|
!errors.Is(outcome.err, context.Canceled) {
|
|
t.Fatalf(
|
|
"canceled execution=(%+v, %v), want generation and context identities",
|
|
outcome.result,
|
|
outcome.err,
|
|
)
|
|
}
|
|
case <-time.After(5 * time.Second):
|
|
t.Fatal("timed out waiting for canceled execution")
|
|
}
|
|
if result, err := cancelingEngine.RunPrepared(context.Background(), canceling); result != nil ||
|
|
!errors.Is(err, promptkit.ErrInvalidRequest) {
|
|
t.Fatalf("canceled execution was reusable: result=(%+v, %v)", result, err)
|
|
}
|
|
}
|
|
|
|
func TestPreparedExecutionConcurrentClaimAllowsOneGeneration(t *testing.T) {
|
|
release := make(chan struct{})
|
|
client := &preparedRecordingClient{
|
|
response: &promptkit.GenerateResponse{Content: "ok"},
|
|
started: make(chan struct{}, 1),
|
|
release: release,
|
|
}
|
|
engine := newPreparedContractEngine(t, client, "concurrent content")
|
|
prepared, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prepared"})
|
|
if err != nil {
|
|
t.Fatalf("prepare execution: %v", err)
|
|
}
|
|
|
|
type outcome struct {
|
|
result *promptkit.RunResult
|
|
err error
|
|
}
|
|
outcomes := make(chan outcome, 2)
|
|
for i := 0; i < 2; i++ {
|
|
go func() {
|
|
result, runErr := engine.RunPrepared(context.Background(), prepared)
|
|
outcomes <- outcome{result: result, err: runErr}
|
|
}()
|
|
}
|
|
|
|
select {
|
|
case <-client.started:
|
|
case <-time.After(5 * time.Second):
|
|
t.Fatal("timed out waiting for generation")
|
|
}
|
|
select {
|
|
case loser := <-outcomes:
|
|
if loser.result != nil || !errors.Is(loser.err, promptkit.ErrInvalidRequest) {
|
|
t.Fatalf("concurrent loser=(%+v, %v), want ErrInvalidRequest", loser.result, loser.err)
|
|
}
|
|
case <-time.After(5 * time.Second):
|
|
t.Fatal("timed out waiting for rejected concurrent claim")
|
|
}
|
|
|
|
close(release)
|
|
select {
|
|
case winner := <-outcomes:
|
|
if winner.err != nil || winner.result == nil {
|
|
t.Fatalf("concurrent winner=(%+v, %v), want success", winner.result, winner.err)
|
|
}
|
|
case <-time.After(5 * time.Second):
|
|
t.Fatal("timed out waiting for successful concurrent claim")
|
|
}
|
|
if len(client.snapshot()) != 1 {
|
|
t.Fatalf("generation calls=%d, want 1", len(client.snapshot()))
|
|
}
|
|
}
|
|
|
|
func TestPreparedExecutionRunAndDiscardRaceHasOneWinner(t *testing.T) {
|
|
const attempts = 32
|
|
|
|
for i := 0; i < attempts; i++ {
|
|
client := &preparedRecordingClient{
|
|
response: &promptkit.GenerateResponse{Content: "ok"},
|
|
}
|
|
engine := newPreparedContractEngine(t, client, "race content")
|
|
prepared, err := engine.PrepareExecution(
|
|
context.Background(),
|
|
promptkit.RunRequest{PromptID: "prepared"},
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("attempt %d prepare execution: %v", i, err)
|
|
}
|
|
|
|
start := make(chan struct{})
|
|
type outcome struct {
|
|
result *promptkit.RunResult
|
|
err error
|
|
}
|
|
runOutcome := make(chan outcome, 1)
|
|
discardDone := make(chan struct{})
|
|
|
|
go func() {
|
|
<-start
|
|
result, runErr := engine.RunPrepared(context.Background(), prepared)
|
|
runOutcome <- outcome{result: result, err: runErr}
|
|
}()
|
|
go func() {
|
|
<-start
|
|
prepared.Discard()
|
|
close(discardDone)
|
|
}()
|
|
|
|
close(start)
|
|
runResult := <-runOutcome
|
|
<-discardDone
|
|
calls := len(client.snapshot())
|
|
switch {
|
|
case runResult.err == nil:
|
|
if runResult.result == nil || calls != 1 {
|
|
t.Fatalf(
|
|
"attempt %d run won with outcome=(%+v, %v), generation calls=%d",
|
|
i,
|
|
runResult.result,
|
|
runResult.err,
|
|
calls,
|
|
)
|
|
}
|
|
case errors.Is(runResult.err, promptkit.ErrInvalidRequest):
|
|
if runResult.result != nil || calls != 0 {
|
|
t.Fatalf(
|
|
"attempt %d discard won with outcome=(%+v, %v), generation calls=%d",
|
|
i,
|
|
runResult.result,
|
|
runResult.err,
|
|
calls,
|
|
)
|
|
}
|
|
default:
|
|
t.Fatalf("attempt %d unexpected run outcome=(%+v, %v)", i, runResult.result, runResult.err)
|
|
}
|
|
if prepared.Details().PromptID != "prepared" {
|
|
t.Fatalf("attempt %d details unavailable after race", i)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestPreparedExecutionDiscardAndFormattingDoNotExposePrivateState(t *testing.T) {
|
|
const (
|
|
directCredential = "pk-test-direct-credential-41f7"
|
|
renderedContent = "rendered-content-sentinel-98d2"
|
|
)
|
|
client := &preparedRecordingClient{
|
|
response: &promptkit.GenerateResponse{Content: "generated output"},
|
|
}
|
|
engine, err := promptkit.NewEngine(
|
|
promptkit.Config{},
|
|
promptkit.WithPromptFS(contractPromptFS("prepared", "profile", renderedContent), "."),
|
|
promptkit.WithProfiles(promptkit.Profile{
|
|
ID: "profile",
|
|
Endpoint: "http://example.test/v1",
|
|
Model: "model",
|
|
APIKeyRequired: true,
|
|
}),
|
|
promptkit.WithLLMClient(client),
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("construct engine: %v", err)
|
|
}
|
|
prepared, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{
|
|
PromptID: "prepared",
|
|
APIKey: directCredential,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("prepare execution: %v", err)
|
|
}
|
|
|
|
formattedValues := []string{
|
|
fmt.Sprint(prepared),
|
|
fmt.Sprintf("%+v", prepared),
|
|
fmt.Sprintf("%#v", prepared),
|
|
}
|
|
for _, formatted := range formattedValues {
|
|
if formatted != "promptkit.PreparedExecution{opaque}" {
|
|
t.Fatalf("unexpected opaque formatting: %q", formatted)
|
|
}
|
|
assertPreparedPrivateValuesAbsent(t, formatted, directCredential, renderedContent)
|
|
}
|
|
payload, err := json.Marshal(prepared)
|
|
if err != nil {
|
|
t.Fatalf("marshal opaque handle: %v", err)
|
|
}
|
|
assertPreparedPrivateValuesAbsent(t, string(payload), directCredential, renderedContent)
|
|
|
|
detailsBefore := prepared.Details()
|
|
detailsJSON, err := json.Marshal(detailsBefore)
|
|
if err != nil {
|
|
t.Fatalf("marshal prepared details: %v", err)
|
|
}
|
|
assertPreparedPrivateValuesAbsent(t, string(detailsJSON), directCredential)
|
|
|
|
prepared.Discard()
|
|
prepared.Discard()
|
|
result, lifecycleErr := engine.RunPrepared(context.Background(), prepared)
|
|
if result != nil || !errors.Is(lifecycleErr, promptkit.ErrInvalidRequest) {
|
|
t.Fatalf("discarded execution result=(%+v, %v), want ErrInvalidRequest", result, lifecycleErr)
|
|
}
|
|
assertPreparedPrivateValuesAbsent(t, lifecycleErr.Error(), directCredential, renderedContent)
|
|
if !reflect.DeepEqual(prepared.Details(), detailsBefore) {
|
|
t.Fatal("details changed after discard")
|
|
}
|
|
|
|
executed, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{
|
|
PromptID: "prepared",
|
|
APIKey: directCredential,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("prepare execution for request inspection: %v", err)
|
|
}
|
|
executionResult, err := engine.RunPrepared(context.Background(), executed)
|
|
if err != nil {
|
|
t.Fatalf("run execution for request inspection: %v", err)
|
|
}
|
|
requests := client.snapshot()
|
|
if len(requests) != 1 || requests[0].APIKey != directCredential {
|
|
t.Fatalf("direct credential did not reach only the client credential field: %#v", requests)
|
|
}
|
|
requestJSON, err := json.Marshal(requests[0])
|
|
if err != nil {
|
|
t.Fatalf("marshal captured generate request: %v", err)
|
|
}
|
|
for _, value := range []string{
|
|
fmt.Sprint(requests[0]),
|
|
fmt.Sprintf("%+v", requests[0]),
|
|
fmt.Sprintf("%#v", requests[0]),
|
|
string(requestJSON),
|
|
fmt.Sprint(executionResult),
|
|
} {
|
|
assertPreparedPrivateValuesAbsent(t, value, directCredential)
|
|
}
|
|
resultJSON, err := json.Marshal(executionResult)
|
|
if err != nil {
|
|
t.Fatalf("marshal execution result: %v", err)
|
|
}
|
|
assertPreparedPrivateValuesAbsent(t, string(resultJSON), directCredential)
|
|
|
|
var nilHandle *promptkit.PreparedExecution
|
|
nilHandle.Discard()
|
|
if !reflect.DeepEqual(nilHandle.Details(), promptkit.PreparedRun{}) {
|
|
t.Fatalf("nil handle details=%+v, want zero value", nilHandle.Details())
|
|
}
|
|
zeroHandle := &promptkit.PreparedExecution{}
|
|
zeroHandle.Discard()
|
|
if !reflect.DeepEqual(zeroHandle.Details(), promptkit.PreparedRun{}) {
|
|
t.Fatalf("zero handle details=%+v, want zero value", zeroHandle.Details())
|
|
}
|
|
}
|
|
|
|
func TestPreparedExecutionCredentialCapacityAndTimingBoundaries(t *testing.T) {
|
|
t.Run("credential is rechecked before generation", func(t *testing.T) {
|
|
const (
|
|
environmentName = "PROMPTKIT_PREPARED_CONTRACT_KEY"
|
|
environmentKey = "environment-credential-sentinel"
|
|
)
|
|
t.Setenv(environmentName, environmentKey)
|
|
client := &preparedRecordingClient{
|
|
response: &promptkit.GenerateResponse{Content: "unexpected"},
|
|
}
|
|
engine, err := promptkit.NewEngine(
|
|
promptkit.Config{},
|
|
promptkit.WithPromptFS(contractPromptFS("prepared", "profile", "content"), "."),
|
|
promptkit.WithProfileFS(preparedCredentialProfileSource(environmentName), "."),
|
|
promptkit.WithLLMClient(client),
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("construct credential engine: %v", err)
|
|
}
|
|
prepared, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prepared"})
|
|
if err != nil {
|
|
t.Fatalf("prepare credential execution: %v", err)
|
|
}
|
|
if err := os.Unsetenv(environmentName); err != nil {
|
|
t.Fatalf("unset credential environment: %v", err)
|
|
}
|
|
|
|
result, err := engine.RunPrepared(context.Background(), prepared)
|
|
if result != nil ||
|
|
!errors.Is(err, promptkit.ErrInvalidRequest) ||
|
|
!errors.Is(err, promptkit.ErrAPIKeyEnvMissing) {
|
|
t.Fatalf("credential execution=(%+v, %v), want credential identities", result, err)
|
|
}
|
|
if len(client.snapshot()) != 0 {
|
|
t.Fatalf("credential failure reached generation: %d calls", len(client.snapshot()))
|
|
}
|
|
assertPreparedPrivateValuesAbsent(t, err.Error(), environmentKey)
|
|
if result, err := engine.RunPrepared(context.Background(), prepared); result != nil ||
|
|
!errors.Is(err, promptkit.ErrInvalidRequest) {
|
|
t.Fatalf("credential failure did not consume handle: result=(%+v, %v)", result, err)
|
|
}
|
|
})
|
|
|
|
t.Run("preparation does not admit and execution timing starts after retention", func(t *testing.T) {
|
|
release := make(chan struct{})
|
|
client := newCapacityGateClient(release, 4)
|
|
engine := newBackendCapacityEngine(t, client, 1, capacityInt(0), nil)
|
|
|
|
activeRun := make(chan capacityRunResult, 1)
|
|
go runCapacityRequest(
|
|
engine,
|
|
context.Background(),
|
|
promptkit.RunRequest{PromptID: "prompt"},
|
|
activeRun,
|
|
)
|
|
awaitCapacityRequest(t, client.started)
|
|
|
|
prepared, err := engine.PrepareExecution(
|
|
context.Background(),
|
|
promptkit.RunRequest{PromptID: "prompt"},
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("prepare while capacity is full: %v", err)
|
|
}
|
|
if _, _, calls := client.snapshot(); calls != 1 {
|
|
t.Fatalf("preparation invoked generation: calls=%d", calls)
|
|
}
|
|
if result, err := engine.RunPrepared(context.Background(), prepared); result != nil ||
|
|
!errors.Is(err, promptkit.ErrCapacityExceeded) {
|
|
t.Fatalf("capacity execution=(%+v, %v), want ErrCapacityExceeded", result, err)
|
|
} else {
|
|
var capacityErr *promptkit.CapacityError
|
|
if !errors.As(err, &capacityErr) || capacityErr == nil || capacityErr.BackendID != "limited" {
|
|
t.Fatalf("capacity execution=%v, want limited CapacityError", err)
|
|
}
|
|
}
|
|
if result, err := engine.RunPrepared(context.Background(), prepared); result != nil ||
|
|
!errors.Is(err, promptkit.ErrInvalidRequest) {
|
|
t.Fatalf("capacity rejection did not consume handle: result=(%+v, %v)", result, err)
|
|
}
|
|
if prepared.Details().PromptID != "prompt" {
|
|
t.Fatal("details unavailable after capacity rejection")
|
|
}
|
|
|
|
close(release)
|
|
activeOutcome := awaitCapacityRun(t, activeRun)
|
|
if activeOutcome.err != nil || activeOutcome.result == nil {
|
|
t.Fatalf("active run outcome=(%+v, %v), want success", activeOutcome.result, activeOutcome.err)
|
|
}
|
|
|
|
timed, err := engine.PrepareExecution(
|
|
context.Background(),
|
|
promptkit.RunRequest{PromptID: "prompt"},
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("prepare timed execution: %v", err)
|
|
}
|
|
details := timed.Details()
|
|
time.Sleep(25 * time.Millisecond)
|
|
executionFloor := time.Now().UTC()
|
|
result, err := engine.RunPrepared(context.Background(), timed)
|
|
if err != nil {
|
|
t.Fatalf("run timed execution: %v", err)
|
|
}
|
|
if result.StartTime.Before(executionFloor) ||
|
|
!result.StartTime.After(details.EndTime) ||
|
|
result.EndTime.Before(result.StartTime) ||
|
|
result.Duration != result.EndTime.Sub(result.StartTime) {
|
|
t.Fatalf(
|
|
"execution timing includes preparation or retention: details_end=%s floor=%s result=%+v",
|
|
details.EndTime,
|
|
executionFloor,
|
|
result,
|
|
)
|
|
}
|
|
if _, _, calls := client.snapshot(); calls != 2 {
|
|
t.Fatalf("generation calls=%d, want active and timed executions only", calls)
|
|
}
|
|
})
|
|
}
|
|
|
|
type mutablePreparedArtifactReader struct {
|
|
mu sync.Mutex
|
|
body string
|
|
hash string
|
|
calls int
|
|
}
|
|
|
|
func (r *mutablePreparedArtifactReader) Read(
|
|
_ context.Context,
|
|
_ promptkit.ArtifactRef,
|
|
) (*promptkit.Artifact, error) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
r.calls++
|
|
return &promptkit.Artifact{
|
|
Body: []byte(r.body),
|
|
Hash: r.hash,
|
|
}, nil
|
|
}
|
|
|
|
func (r *mutablePreparedArtifactReader) set(body, hash string) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
r.body = body
|
|
r.hash = hash
|
|
}
|
|
|
|
func (r *mutablePreparedArtifactReader) callCount() int {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
return r.calls
|
|
}
|
|
|
|
type preparedRecordingClient struct {
|
|
mu sync.Mutex
|
|
response *promptkit.GenerateResponse
|
|
err error
|
|
requests []promptkit.GenerateRequest
|
|
started chan struct{}
|
|
release <-chan struct{}
|
|
}
|
|
|
|
func (c *preparedRecordingClient) Generate(
|
|
ctx context.Context,
|
|
request promptkit.GenerateRequest,
|
|
) (*promptkit.GenerateResponse, error) {
|
|
c.mu.Lock()
|
|
c.requests = append(c.requests, request)
|
|
c.mu.Unlock()
|
|
|
|
if c.started != nil {
|
|
c.started <- struct{}{}
|
|
}
|
|
if c.release != nil {
|
|
select {
|
|
case <-c.release:
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
}
|
|
}
|
|
if c.err != nil {
|
|
return nil, c.err
|
|
}
|
|
return c.response, nil
|
|
}
|
|
|
|
func (c *preparedRecordingClient) snapshot() []promptkit.GenerateRequest {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
return append([]promptkit.GenerateRequest(nil), c.requests...)
|
|
}
|
|
|
|
func newPreparedContractEngine(
|
|
t *testing.T,
|
|
client promptkit.LLMClient,
|
|
message string,
|
|
) *promptkit.Engine {
|
|
t.Helper()
|
|
engine, err := promptkit.NewEngine(
|
|
promptkit.Config{},
|
|
promptkit.WithPromptFS(contractPromptFS("prepared", "profile", message), "."),
|
|
promptkit.WithProfiles(promptkit.Profile{
|
|
ID: "profile",
|
|
Endpoint: "http://example.test/v1",
|
|
Model: "model",
|
|
}),
|
|
promptkit.WithLLMClient(client),
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("construct prepared execution engine: %v", err)
|
|
}
|
|
return engine
|
|
}
|
|
|
|
func preparedPromptSource(label string) fstest.MapFS {
|
|
return fstest.MapFS{
|
|
"prompt.yaml": &fstest.MapFile{Data: []byte(`id: prepared
|
|
version: "1"
|
|
default_profile: profile
|
|
inputs:
|
|
- name: input
|
|
required: true
|
|
messages:
|
|
- role: user
|
|
content: 'Input={{input "input"}} Label={{.label}}'
|
|
description: ` + label + `
|
|
output:
|
|
format: json
|
|
validation_mode: json_schema
|
|
schema_path: schema.json
|
|
`)},
|
|
}
|
|
}
|
|
|
|
func preparedProfileSource(model string) fstest.MapFS {
|
|
return fstest.MapFS{
|
|
"profile.yaml": &fstest.MapFile{Data: []byte(`id: profile
|
|
endpoint: http://example.test/v1
|
|
model: ` + model + `
|
|
`)},
|
|
}
|
|
}
|
|
|
|
func preparedCredentialProfileSource(environmentName string) fstest.MapFS {
|
|
return fstest.MapFS{
|
|
"profile.yaml": &fstest.MapFile{Data: []byte(`id: profile
|
|
endpoint: http://example.test/v1
|
|
model: model
|
|
api_key_env: ` + environmentName + `
|
|
`)},
|
|
}
|
|
}
|
|
|
|
func preparedSchemaSource() fstest.MapFS {
|
|
return fstest.MapFS{
|
|
"schema.json": &fstest.MapFile{Data: []byte(`{
|
|
"$schema": "https://json-schema.org/draft/2020-12/schema",
|
|
"title": "original root",
|
|
"type": "object",
|
|
"required": ["value"],
|
|
"properties": {
|
|
"value": {"$ref": "value.json"}
|
|
}
|
|
}`)},
|
|
"value.json": &fstest.MapFile{Data: []byte(`{
|
|
"$schema": "https://json-schema.org/draft/2020-12/schema",
|
|
"type": "integer",
|
|
"minimum": 2
|
|
}`)},
|
|
}
|
|
}
|
|
|
|
func assertPreparedPrivateValuesAbsent(t *testing.T, value string, privateValues ...string) {
|
|
t.Helper()
|
|
for _, privateValue := range privateValues {
|
|
if strings.Contains(value, privateValue) {
|
|
t.Fatalf("value exposed private data %q: %s", privateValue, value)
|
|
}
|
|
}
|
|
}
|