Add prepared execution lifecycle to the runner

This commit is contained in:
2026-07-30 18:10:28 +00:00
parent c301eb8d55
commit 49fe402dd2
5 changed files with 877 additions and 29 deletions

View File

@@ -18,13 +18,19 @@ type visit struct {
ptr uintptr
}
// Copy validates and deeply copies a JSON-compatible value while preserving
// compatible concrete map, slice, array, scalar, and number types.
func Copy(src any) (any, error) {
return copyValue(reflect.ValueOf(src), "value", make(map[visit]struct{}), true)
}
// CopyMap validates and deeply copies an extra-parameter map while preserving
// compatible concrete map, slice, array, scalar, and number types.
func CopyMap(src map[string]any) (map[string]any, error) {
if src == nil {
return nil, nil
}
copied, err := copyValue(reflect.ValueOf(src), "extra_params", make(map[visit]struct{}))
copied, err := copyValue(reflect.ValueOf(src), "extra_params", make(map[visit]struct{}), false)
if err != nil {
return nil, err
}
@@ -35,7 +41,12 @@ func CopyMap(src map[string]any) (map[string]any, error) {
return out, nil
}
func copyValue(value reflect.Value, path string, seen map[visit]struct{}) (any, error) {
func copyValue(
value reflect.Value,
path string,
seen map[visit]struct{},
allowEmptyMapKeys bool,
) (any, error) {
if !value.IsValid() {
return nil, nil
}
@@ -43,7 +54,7 @@ func copyValue(value reflect.Value, path string, seen map[visit]struct{}) (any,
if value.IsNil() {
return nil, nil
}
return copyValue(value.Elem(), path, seen)
return copyValue(value.Elem(), path, seen, allowEmptyMapKeys)
}
if !value.CanInterface() {
return nil, fmt.Errorf("%s: value cannot be copied", path)
@@ -88,22 +99,27 @@ func copyValue(value reflect.Value, path string, seen map[visit]struct{}) (any,
}
seen[current] = struct{}{}
defer delete(seen, current)
return copyValue(value.Elem(), path, seen)
return copyValue(value.Elem(), path, seen, allowEmptyMapKeys)
case reflect.Map:
return copyMapValue(value, path, seen)
return copyMapValue(value, path, seen, allowEmptyMapKeys)
case reflect.Slice:
if value.IsNil() {
return nil, nil
}
return copySequenceValue(value, path, seen)
return copySequenceValue(value, path, seen, allowEmptyMapKeys)
case reflect.Array:
return copySequenceValue(value, path, seen)
return copySequenceValue(value, path, seen, allowEmptyMapKeys)
default:
return nil, fmt.Errorf("%s: unsupported JSON value type %s", path, value.Type())
}
}
func copyMapValue(value reflect.Value, path string, seen map[visit]struct{}) (any, error) {
func copyMapValue(
value reflect.Value,
path string,
seen map[visit]struct{},
allowEmptyMapKeys bool,
) (any, error) {
if value.IsNil() {
return nil, nil
}
@@ -133,10 +149,10 @@ func copyMapValue(value reflect.Value, path string, seen map[visit]struct{}) (an
elementType := value.Type().Elem()
for _, key := range keys {
name := key.String()
if name == "" {
if name == "" && !allowEmptyMapKeys {
return nil, fmt.Errorf("%s: map key must not be empty", path)
}
copied, err := copyValue(value.MapIndex(key), path+"."+name, seen)
copied, err := copyValue(value.MapIndex(key), path+"."+name, seen, allowEmptyMapKeys)
if err != nil {
return nil, err
}
@@ -171,7 +187,12 @@ func copyMapValue(value reflect.Value, path string, seen map[visit]struct{}) (an
return out, nil
}
func copySequenceValue(value reflect.Value, path string, seen map[visit]struct{}) (any, error) {
func copySequenceValue(
value reflect.Value,
path string,
seen map[visit]struct{},
allowEmptyMapKeys bool,
) (any, error) {
var current visit
if value.Kind() == reflect.Slice {
current = visit{typ: value.Type(), ptr: value.Pointer()}
@@ -186,7 +207,12 @@ func copySequenceValue(value reflect.Value, path string, seen map[visit]struct{}
preserveType := true
elementType := value.Type().Elem()
for i := 0; i < value.Len(); i++ {
copied, err := copyValue(value.Index(i), fmt.Sprintf("%s[%d]", path, i), seen)
copied, err := copyValue(
value.Index(i),
fmt.Sprintf("%s[%d]", path, i),
seen,
allowEmptyMapKeys,
)
if err != nil {
return nil, err
}

View File

@@ -44,6 +44,21 @@ func TestCopyMapPreservesTypesAndIsolatesMutations(t *testing.T) {
}
}
func TestCopyAllowsEmptyObjectKeysAndIsolatesMutations(t *testing.T) {
nested := map[string]any{"": []any{"original"}}
copiedValue, err := jsonvalue.Copy(nested)
if err != nil {
t.Fatalf("copy value: %v", err)
}
nested[""].([]any)[0] = "changed"
copied := copiedValue.(map[string]any)
if got := copied[""].([]any)[0]; got != "original" {
t.Fatalf("copied value was not isolated: %v", got)
}
}
func TestCopyMapRejectsInvalidValues(t *testing.T) {
cyclicMap := map[string]any{}
cyclicMap["self"] = cyclicMap

View File

@@ -0,0 +1,300 @@
package usecase
import (
"context"
"fmt"
"sync"
"time"
"gitea.maximumdirect.net/eric/promptkit/internal/domain"
"gitea.maximumdirect.net/eric/promptkit/internal/jsonvalue"
"gitea.maximumdirect.net/eric/promptkit/internal/validate"
)
type preparedExecutionState uint8
const (
preparedExecutionReady preparedExecutionState = iota
preparedExecutionClaimed
preparedExecutionDiscarded
)
// PreparedExecution owns one frozen, single-use runner execution.
type PreparedExecution struct {
owner *Runner
mu sync.Mutex
state preparedExecutionState
details *domain.PreparedRun
payload *preparedExecutionPayload
}
type preparedExecutionPayload struct {
prepared *domain.PreparedRun
validation validate.PreparedValidation
}
// PrepareExecution completes preparation without generation or admission and
// returns a runner-bound, single-use execution.
func (r *Runner) PrepareExecution(ctx context.Context, req domain.RunRequest) (*PreparedExecution, error) {
state, err := r.resolvePreparation(ctx, req, time.Now().UTC())
if err != nil {
return nil, err
}
validationPlan, err := r.prepareValidation(ctx, state.effectiveContract)
if err != nil {
return nil, err
}
structuredOutput, err := r.structuredOutputFromValidationPlan(
state.definition,
state.effectiveContract,
validationPlan,
)
if err != nil {
return nil, err
}
prepared, err := r.completePreparationWithStructuredOutput(ctx, req, state, structuredOutput)
if err != nil {
return nil, err
}
executionSnapshot, err := clonePreparedRun(prepared)
if err != nil {
prepared.EffectiveModelParams.APIKey = ""
return nil, fmt.Errorf("%w: failed to copy prepared execution: %v", ErrInvalidRequest, err)
}
executionSnapshot.StructuredOutput = prepared.StructuredOutput
prepared.EffectiveModelParams.APIKey = ""
details, err := clonePreparedRun(executionSnapshot)
if err != nil {
executionSnapshot.EffectiveModelParams.APIKey = ""
return nil, fmt.Errorf("%w: failed to copy prepared execution details: %v", ErrInvalidRequest, err)
}
details.EffectiveModelParams.APIKey = ""
return &PreparedExecution{
owner: r,
state: preparedExecutionReady,
details: details,
payload: &preparedExecutionPayload{
prepared: executionSnapshot,
validation: validationPlan,
},
}, nil
}
func (r *Runner) prepareValidation(
ctx context.Context,
contract domain.OutputContract,
) (validate.PreparedValidation, error) {
if r.validator == nil {
return noOpPreparedValidation{contract: contract}, nil
}
preparer, ok := r.validator.(validate.ValidationPreparer)
if !ok {
return nil, fmt.Errorf("%w: validator does not support prepared validation", ErrValidation)
}
plan, err := preparer.PrepareValidation(ctx, contract)
if err != nil {
return nil, fmt.Errorf("%w: %w", ErrValidation, err)
}
if plan == nil {
return nil, fmt.Errorf("%w: validator returned nil prepared validation", ErrValidation)
}
return plan, nil
}
func (r *Runner) structuredOutputFromValidationPlan(
def *domain.PromptDefinition,
contract domain.OutputContract,
plan validate.PreparedValidation,
) (*domain.StructuredOutputSpec, error) {
if contract.ValidationMode != domain.ValidationJSONSchema {
return nil, nil
}
schemaDocument := plan.SchemaDocument()
if schemaDocument == nil {
if r.validator == nil {
return nil, nil
}
return nil, fmt.Errorf("%w: prepared json_schema validation has no schema document", ErrValidation)
}
return structuredOutputSpec(def, schemaDocument), nil
}
// Details returns a fresh credential-redacted copy of the prepared run.
func (p *PreparedExecution) Details() *domain.PreparedRun {
if p == nil {
return nil
}
p.mu.Lock()
detailsSnapshot := p.details
p.mu.Unlock()
details, err := clonePreparedRun(detailsSnapshot)
if err != nil {
return nil
}
return details
}
// Discard invalidates an unclaimed execution and drops its private payload.
func (p *PreparedExecution) Discard() {
if p == nil {
return
}
p.mu.Lock()
if p.state != preparedExecutionReady {
p.mu.Unlock()
return
}
p.state = preparedExecutionDiscarded
payload := p.payload
p.payload = nil
p.mu.Unlock()
payload.clear()
}
// RunPrepared claims and executes one prepared execution owned by this runner.
func (r *Runner) RunPrepared(ctx context.Context, prepared *PreparedExecution) (*domain.RunResult, error) {
payload, err := prepared.claim(r)
if err != nil {
return nil, err
}
defer payload.clear()
runID, err := newRunID()
if err != nil {
return nil, fmt.Errorf("failed to create run id: %w", err)
}
start := time.Now().UTC()
target := payload.prepared.EffectiveModelParams
if err := validateAPIKey(target.APIKeyEnv, target.APIKey, target.APIKeyRequired); err != nil {
return nil, fmt.Errorf("%w: %w", ErrInvalidRequest, err)
}
release, err := r.admitRun(ctx, target.BackendID)
if err != nil {
return nil, err
}
defer release()
return r.executePreparedRun(ctx, payload.prepared, runID, start, func(
ctx context.Context,
artifact *domain.Artifact,
attemptsUsed int,
) (domain.ValidationResult, error) {
result, validationErr := payload.validation.Validate(ctx, artifact)
if validationErr != nil {
return domain.ValidationResult{}, validationErr
}
result.RepairAttempts = attemptsUsed
return result, nil
})
}
func (p *PreparedExecution) claim(owner *Runner) (*preparedExecutionPayload, error) {
if p == nil || owner == nil || p.owner != owner {
return nil, fmt.Errorf("%w: prepared execution does not belong to this runner", ErrInvalidRequest)
}
p.mu.Lock()
defer p.mu.Unlock()
if p.state != preparedExecutionReady || p.payload == nil {
return nil, fmt.Errorf("%w: prepared execution is not ready", ErrInvalidRequest)
}
p.state = preparedExecutionClaimed
payload := p.payload
p.payload = nil
return payload, nil
}
func (p *preparedExecutionPayload) clear() {
if p == nil {
return
}
if p.prepared != nil {
p.prepared.EffectiveModelParams.APIKey = ""
}
p.prepared = nil
p.validation = nil
}
type noOpPreparedValidation struct {
contract domain.OutputContract
}
func (p noOpPreparedValidation) Validate(
ctx context.Context,
_ *domain.Artifact,
) (domain.ValidationResult, error) {
if err := ctx.Err(); err != nil {
return domain.ValidationResult{}, err
}
return domain.ValidationResult{
Status: domain.ValidationSkipped,
Mode: p.contract.ValidationMode,
SchemaPath: p.contract.SchemaPath,
RepairAttempts: p.contract.RepairAttempts,
IsValid: true,
}, nil
}
func (noOpPreparedValidation) SchemaDocument() any {
return nil
}
func clonePreparedRun(source *domain.PreparedRun) (*domain.PreparedRun, error) {
if source == nil {
return nil, nil
}
copied := *source
extraParams, err := jsonvalue.CopyMap(source.EffectiveModelParams.ExtraParams)
if err != nil {
return nil, err
}
copied.EffectiveModelParams.ExtraParams = extraParams
if source.InputHashes != nil {
copied.InputHashes = make(map[string]string, len(source.InputHashes))
for name, hash := range source.InputHashes {
copied.InputHashes[name] = hash
}
}
if source.Messages != nil {
copied.Messages = make([]domain.RenderedMessage, len(source.Messages))
for i, message := range source.Messages {
copied.Messages[i] = message
if message.CacheControl != nil {
cacheControl := *message.CacheControl
copied.Messages[i].CacheControl = &cacheControl
}
}
}
if source.StructuredOutput != nil {
structuredOutput := *source.StructuredOutput
copied.StructuredOutput = &structuredOutput
if source.StructuredOutput.JSONSchema != nil {
jsonSchema := *source.StructuredOutput.JSONSchema
copied.StructuredOutput.JSONSchema = &jsonSchema
schema, copyErr := cloneJSONValue(source.StructuredOutput.JSONSchema.Schema)
if copyErr != nil {
return nil, copyErr
}
copied.StructuredOutput.JSONSchema.Schema = schema
}
}
return &copied, nil
}
func cloneJSONValue(source any) (any, error) {
return jsonvalue.Copy(source)
}

View File

@@ -0,0 +1,468 @@
package usecase
import (
"context"
"errors"
"os"
"reflect"
"testing"
"gitea.maximumdirect.net/eric/promptkit/internal/domain"
"gitea.maximumdirect.net/eric/promptkit/internal/validate"
)
type recordingPreparedValidation struct {
contract domain.OutputContract
schemaDocument any
results []domain.ValidationResult
errs []error
artifacts []string
}
func (p *recordingPreparedValidation) Validate(
_ context.Context,
artifact *domain.Artifact,
) (domain.ValidationResult, error) {
p.artifacts = append(p.artifacts, string(artifact.Body))
index := len(p.artifacts) - 1
if index < len(p.errs) && p.errs[index] != nil {
return domain.ValidationResult{}, p.errs[index]
}
if len(p.results) == 0 {
return domain.ValidationResult{
Status: domain.ValidationPassed,
Mode: p.contract.ValidationMode,
IsValid: true,
}, nil
}
if index >= len(p.results) {
index = len(p.results) - 1
}
return p.results[index], nil
}
func (p *recordingPreparedValidation) SchemaDocument() any {
return p.schemaDocument
}
type recordingValidationPreparer struct {
plan *recordingPreparedValidation
prepareErr error
prepareCalls int
directValidateCalls int
}
func (v *recordingValidationPreparer) Validate(
context.Context,
*domain.Artifact,
domain.OutputContract,
) (domain.ValidationResult, error) {
v.directValidateCalls++
return domain.ValidationResult{}, errors.New("live validation must not be used")
}
func (v *recordingValidationPreparer) PrepareValidation(
_ context.Context,
contract domain.OutputContract,
) (validate.PreparedValidation, error) {
v.prepareCalls++
if v.prepareErr != nil {
return nil, v.prepareErr
}
v.plan.contract = contract
return v.plan, nil
}
func TestRunnerPrepareExecutionCompletesWithoutAdmissionOrGeneration(t *testing.T) {
schemaDocument := map[string]any{
"type": "object",
"properties": map[string]any{
"value": map[string]any{"type": "string"},
"": map[string]any{"type": "boolean"},
},
}
def := promptDef(domain.FormatJSON, domain.ValidationJSONSchema, 0)
def.Validation.SchemaPath = "schema.json"
reader := defaultArtifactReader()
renderer := &fakeRenderer{rendered: &domain.RenderedPrompt{
SessionID: "prepared-session",
Messages: []domain.RenderedMessage{{
Role: "user",
Content: "original message",
}},
}}
llmClient := &fakeLLM{forbid: true}
validator := &recordingValidationPreparer{
plan: &recordingPreparedValidation{schemaDocument: schemaDocument},
}
admitter := &fakeRunAdmitter{}
profile := defaultExecutionProfile()
profile.ExtraParams = map[string]any{
"metadata": map[string]any{"source": "original"},
}
runner := NewRunner(
&fakePromptRepo{def: def},
&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": profile}},
nil,
reader,
renderer,
llmClient,
validator,
admitter,
)
prepared, err := runner.PrepareExecution(context.Background(), domain.RunRequest{
PromptID: "p",
ProfileID: "exec",
APIKey: "direct-test-key",
Inputs: singleInputRef(),
})
if err != nil {
t.Fatalf("prepare execution: %v", err)
}
defer prepared.Discard()
if validator.prepareCalls != 1 || validator.directValidateCalls != 0 {
t.Fatalf(
"validation calls=(prepare=%d direct=%d), want (1, 0)",
validator.prepareCalls,
validator.directValidateCalls,
)
}
if reader.calls != 1 || renderer.calls != 1 {
t.Fatalf("completion calls=(artifact=%d render=%d), want (1, 1)", reader.calls, renderer.calls)
}
if len(admitter.backendIDs) != 0 || llmClient.calls != 0 {
t.Fatalf("prepare invoked execution collaborators: admission=%v generation=%d", admitter.backendIDs, llmClient.calls)
}
first := prepared.Details()
if first == nil {
t.Fatal("prepared details are nil")
}
if first.EffectiveModelParams.APIKey != "" {
t.Fatal("prepared details retained the direct API key")
}
if first.StructuredOutput == nil ||
first.StructuredOutput.JSONSchema == nil ||
!reflect.DeepEqual(first.StructuredOutput.JSONSchema.Schema, schemaDocument) {
t.Fatalf("prepared details have unexpected structured output: %#v", first.StructuredOutput)
}
first.Messages[0].Content = "caller mutation"
first.InputHashes["input"] = "caller mutation"
first.EffectiveModelParams.ExtraParams["metadata"].(map[string]any)["source"] = "caller mutation"
first.StructuredOutput.JSONSchema.Schema.(map[string]any)["type"] = "string"
renderer.rendered.Messages[0].Content = "source mutation"
second := prepared.Details()
if second.Messages[0].Content != "original message" ||
second.InputHashes["input"] == "caller mutation" ||
second.EffectiveModelParams.ExtraParams["metadata"].(map[string]any)["source"] != "original" ||
second.StructuredOutput.JSONSchema.Schema.(map[string]any)["type"] != "object" {
t.Fatalf("details did not preserve an independent snapshot: %#v", second)
}
}
func TestRunnerRunPreparedRechecksEnvironmentCredentialBeforeAdmission(t *testing.T) {
const environmentName = "PROMPTKIT_PREPARED_EXECUTION_TEST_KEY"
t.Setenv(environmentName, "available-during-preparation")
profile := defaultExecutionProfile()
profile.APIKeyEnv = environmentName
validator := &recordingValidationPreparer{plan: &recordingPreparedValidation{}}
admitter := &fakeRunAdmitter{}
llmClient := &fakeLLM{resp: &domain.GenerateResponse{Content: "unexpected"}}
runner := NewRunner(
&fakePromptRepo{def: promptDef(domain.FormatText, domain.ValidationNone, 0)},
&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": profile}},
nil,
defaultArtifactReader(),
defaultRenderer(),
llmClient,
validator,
admitter,
)
prepared, err := runner.PrepareExecution(context.Background(), domain.RunRequest{
PromptID: "p",
ProfileID: "exec",
Inputs: singleInputRef(),
})
if err != nil {
t.Fatalf("prepare execution: %v", err)
}
if err := os.Unsetenv(environmentName); err != nil {
t.Fatalf("unset credential environment: %v", err)
}
result, err := runner.RunPrepared(context.Background(), prepared)
if result != nil {
t.Fatalf("credential failure returned partial result: %+v", result)
}
if !errors.Is(err, ErrInvalidRequest) || !errors.Is(err, ErrAPIKeyEnvMissing) {
t.Fatalf("credential error identities are missing: %v", err)
}
if len(admitter.backendIDs) != 0 || llmClient.calls != 0 {
t.Fatalf("credential failure reached admission or generation: admission=%v generation=%d", admitter.backendIDs, llmClient.calls)
}
if _, err := runner.RunPrepared(context.Background(), prepared); !errors.Is(err, ErrInvalidRequest) {
t.Fatalf("credential failure did not consume execution: %v", err)
}
}
func TestRunnerRunPreparedUsesFrozenValidationForInitialAndRepairOutputs(t *testing.T) {
validator := &recordingValidationPreparer{
plan: &recordingPreparedValidation{
results: []domain.ValidationResult{
{
Status: domain.ValidationFailed,
Mode: domain.ValidationJSON,
Errors: []string{"invalid"},
IsValid: false,
},
{
Status: domain.ValidationPassed,
Mode: domain.ValidationJSON,
IsValid: true,
},
},
},
}
repairer := &fakeRepairer{
responses: []*domain.GenerateResponse{{Content: `{"repaired":true}`}},
}
admitter := &fakeRunAdmitter{}
reader := defaultArtifactReader()
renderer := defaultRenderer()
runner := NewRunnerWithRepairer(
&fakePromptRepo{def: promptDef(domain.FormatJSON, domain.ValidationJSON, 1)},
&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": defaultExecutionProfile()}},
nil,
reader,
renderer,
&fakeLLM{resp: &domain.GenerateResponse{Content: `{"broken":true}`}},
validator,
repairer,
admitter,
)
prepared, err := runner.PrepareExecution(context.Background(), domain.RunRequest{
PromptID: "p",
ProfileID: "exec",
Inputs: singleInputRef(),
})
if err != nil {
t.Fatalf("prepare execution: %v", err)
}
result, err := runner.RunPrepared(context.Background(), prepared)
if err != nil {
t.Fatalf("run prepared: %v", err)
}
if !reflect.DeepEqual(validator.plan.artifacts, []string{`{"broken":true}`, `{"repaired":true}`}) {
t.Fatalf("prepared validation artifacts=%#v", validator.plan.artifacts)
}
if validator.directValidateCalls != 0 || repairer.calls != 1 {
t.Fatalf("validation/repair calls=(direct=%d repair=%d), want (0, 1)", validator.directValidateCalls, repairer.calls)
}
if result.Validation.Status != domain.ValidationPassed || result.Validation.RepairAttempts != 1 {
t.Fatalf("unexpected repaired validation result: %+v", result.Validation)
}
if admitter.releaseCalls != 1 {
t.Fatalf("admission releases=%d, want 1", admitter.releaseCalls)
}
if reader.calls != 1 || renderer.calls != 1 {
t.Fatalf("execution reopened preparation sources: artifact=%d render=%d", reader.calls, renderer.calls)
}
}
func TestRunnerRunPreparedReleasesAdmissionAcrossExecutionErrors(t *testing.T) {
generationFailure := errors.New("generation failed")
validationFailure := errors.New("validation failed")
repairFailure := errors.New("repair failed")
tests := []struct {
name string
generationErr error
validation *recordingPreparedValidation
repairer *fakeRepairer
wantError error
}{
{
name: "generation failure",
generationErr: generationFailure,
validation: &recordingPreparedValidation{},
wantError: ErrLLMGenerate,
},
{
name: "validation failure",
validation: &recordingPreparedValidation{
errs: []error{validationFailure},
},
wantError: ErrValidation,
},
{
name: "repair failure",
validation: &recordingPreparedValidation{
results: []domain.ValidationResult{{
Status: domain.ValidationFailed,
Mode: domain.ValidationJSON,
Errors: []string{"invalid"},
IsValid: false,
}},
},
repairer: &fakeRepairer{err: repairFailure},
wantError: ErrValidation,
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
def := promptDef(domain.FormatJSON, domain.ValidationJSON, 1)
if test.repairer == nil {
def.Validation.RepairAttempts = 0
}
validator := &recordingValidationPreparer{plan: test.validation}
admitter := &fakeRunAdmitter{}
runner := NewRunnerWithRepairer(
&fakePromptRepo{def: def},
&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": defaultExecutionProfile()}},
nil,
defaultArtifactReader(),
defaultRenderer(),
&fakeLLM{
resp: &domain.GenerateResponse{Content: `{"value":true}`},
err: test.generationErr,
},
validator,
test.repairer,
admitter,
)
prepared, err := runner.PrepareExecution(context.Background(), domain.RunRequest{
PromptID: "p",
ProfileID: "exec",
Inputs: singleInputRef(),
})
if err != nil {
t.Fatalf("prepare execution: %v", err)
}
result, err := runner.RunPrepared(context.Background(), prepared)
if result != nil || !errors.Is(err, test.wantError) {
t.Fatalf("run prepared=(%+v, %v), want %v", result, err, test.wantError)
}
if len(admitter.backendIDs) != 1 || admitter.releaseCalls != 1 {
t.Fatalf(
"admission calls=%#v releases=%d, want one each",
admitter.backendIDs,
admitter.releaseCalls,
)
}
})
}
}
func TestRunnerPreparedExecutionOwnershipUseAndDiscard(t *testing.T) {
newRunner := func() *Runner {
return NewRunner(
&fakePromptRepo{def: promptDef(domain.FormatText, domain.ValidationNone, 0)},
&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": defaultExecutionProfile()}},
nil,
defaultArtifactReader(),
defaultRenderer(),
&fakeLLM{resp: &domain.GenerateResponse{Content: "ok"}},
&recordingValidationPreparer{plan: &recordingPreparedValidation{}},
nil,
)
}
request := domain.RunRequest{
PromptID: "p",
ProfileID: "exec",
Inputs: singleInputRef(),
}
owner := newRunner()
prepared, err := owner.PrepareExecution(context.Background(), request)
if err != nil {
t.Fatalf("prepare execution: %v", err)
}
if _, err := newRunner().RunPrepared(context.Background(), prepared); !errors.Is(err, ErrInvalidRequest) {
t.Fatalf("foreign runner error=%v, want ErrInvalidRequest", err)
}
if _, err := owner.RunPrepared(context.Background(), prepared); err != nil {
t.Fatalf("owner run prepared: %v", err)
}
if _, err := owner.RunPrepared(context.Background(), prepared); !errors.Is(err, ErrInvalidRequest) {
t.Fatalf("second owner run error=%v, want ErrInvalidRequest", err)
}
if prepared.Details() == nil {
t.Fatal("details unavailable after execution")
}
discarded, err := owner.PrepareExecution(context.Background(), request)
if err != nil {
t.Fatalf("prepare discarded execution: %v", err)
}
discarded.Discard()
discarded.Discard()
if _, err := owner.RunPrepared(context.Background(), discarded); !errors.Is(err, ErrInvalidRequest) {
t.Fatalf("discarded execution error=%v, want ErrInvalidRequest", err)
}
if discarded.Details() == nil {
t.Fatal("details unavailable after discard")
}
}
func TestRunnerPreparedExecutionWithoutValidatorSkipsValidation(t *testing.T) {
runner := NewRunner(
&fakePromptRepo{def: promptDef(domain.FormatJSON, domain.ValidationJSON, 0)},
&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": defaultExecutionProfile()}},
nil,
defaultArtifactReader(),
defaultRenderer(),
&fakeLLM{resp: &domain.GenerateResponse{Content: `{}`}},
nil,
nil,
)
prepared, err := runner.PrepareExecution(context.Background(), domain.RunRequest{
PromptID: "p",
ProfileID: "exec",
Inputs: singleInputRef(),
})
if err != nil {
t.Fatalf("prepare execution: %v", err)
}
result, err := runner.RunPrepared(context.Background(), prepared)
if err != nil {
t.Fatalf("run prepared: %v", err)
}
if result.Validation.Status != domain.ValidationSkipped || !result.Validation.IsValid {
t.Fatalf("unexpected no-validator result: %+v", result.Validation)
}
}
func TestRunnerPrepareExecutionRequiresValidationPreparer(t *testing.T) {
reader := defaultArtifactReader()
runner := NewRunner(
&fakePromptRepo{def: promptDef(domain.FormatText, domain.ValidationBasic, 0)},
&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": defaultExecutionProfile()}},
nil,
reader,
defaultRenderer(),
&fakeLLM{forbid: true},
&fakeValidator{},
nil,
)
prepared, err := runner.PrepareExecution(context.Background(), domain.RunRequest{
PromptID: "p",
ProfileID: "exec",
Inputs: singleInputRef(),
})
if prepared != nil || !errors.Is(err, ErrValidation) {
t.Fatalf("prepare execution=(%+v, %v), want ErrValidation", prepared, err)
}
if reader.calls != 0 {
t.Fatalf("unsupported validator allowed completion, artifact calls=%d", reader.calls)
}
}

View File

@@ -131,26 +131,39 @@ func (r *Runner) Run(ctx context.Context, req domain.RunRequest) (*domain.RunRes
return nil, err
}
if r.admitter != nil {
release, admitErr := r.admitter.Admit(ctx, state.effectiveModel.BackendID)
if admitErr != nil {
if errors.Is(admitErr, capacity.ErrCapacityExceeded) {
return nil, fmt.Errorf(
"backend %q admission: %w",
state.effectiveModel.BackendID,
admitErr,
)
}
return nil, admitErr
}
defer release()
release, err := r.admitRun(ctx, state.effectiveModel.BackendID)
if err != nil {
return nil, err
}
defer release()
prepared, err := r.completePreparation(ctx, req, state)
if err != nil {
return nil, err
}
return r.executePreparedRun(ctx, prepared, runID, start, func(
ctx context.Context,
artifact *domain.Artifact,
attemptsUsed int,
) (domain.ValidationResult, error) {
return r.validateOutput(ctx, artifact, prepared.OutputContract, attemptsUsed)
})
}
type preparedValidationFunc func(
context.Context,
*domain.Artifact,
int,
) (domain.ValidationResult, error)
func (r *Runner) executePreparedRun(
ctx context.Context,
prepared *domain.PreparedRun,
runID string,
start time.Time,
validateArtifact preparedValidationFunc,
) (*domain.RunResult, error) {
genResp, err := r.llm.Generate(ctx, domain.GenerateRequest{
Prompt: domain.RenderedPrompt{SessionID: prepared.SessionID, Messages: prepared.Messages},
Target: prepared.EffectiveModelParams,
@@ -165,7 +178,7 @@ func (r *Runner) Run(ctx context.Context, req domain.RunRequest) (*domain.RunRes
}
outputArtifact := buildOutputArtifact(genResp.Content, prepared.OutputContract.Format)
validationResult, err := r.validateOutput(ctx, &outputArtifact, prepared.OutputContract, 0)
validationResult, err := validateArtifact(ctx, &outputArtifact, 0)
if err != nil {
return nil, fmt.Errorf("%w: %w", ErrValidation, err)
}
@@ -195,7 +208,7 @@ func (r *Runner) Run(ctx context.Context, req domain.RunRequest) (*domain.RunRes
genResp = repairResp
outputArtifact = buildOutputArtifact(genResp.Content, prepared.OutputContract.Format)
validationResult, err = r.validateOutput(ctx, &outputArtifact, prepared.OutputContract, attemptsUsed)
validationResult, err = validateArtifact(ctx, &outputArtifact, attemptsUsed)
if err != nil {
return nil, fmt.Errorf("%w: %w", ErrValidation, err)
}
@@ -324,7 +337,15 @@ func (r *Runner) completePreparation(
if err != nil {
return nil, err
}
return r.completePreparationWithStructuredOutput(ctx, req, state, structuredOutput)
}
func (r *Runner) completePreparationWithStructuredOutput(
ctx context.Context,
req domain.RunRequest,
state *preparationState,
structuredOutput *domain.StructuredOutputSpec,
) (*domain.PreparedRun, error) {
resolvedInputs := make(map[string]*domain.Artifact, len(req.Inputs))
inputHashes := make(map[string]string, len(req.Inputs))
for name, ref := range req.Inputs {
@@ -374,6 +395,20 @@ func (r *Runner) completePreparation(
}, nil
}
func (r *Runner) admitRun(ctx context.Context, backendID string) (func(), error) {
if r.admitter == nil {
return func() {}, nil
}
release, err := r.admitter.Admit(ctx, backendID)
if err != nil {
if errors.Is(err, capacity.ErrCapacityExceeded) {
return nil, fmt.Errorf("backend %q admission: %w", backendID, err)
}
return nil, err
}
return release, nil
}
func (r *Runner) resolveStructuredOutput(ctx context.Context, def *domain.PromptDefinition, contract domain.OutputContract) (*domain.StructuredOutputSpec, error) {
if contract.ValidationMode != domain.ValidationJSONSchema {
return nil, nil
@@ -389,14 +424,18 @@ func (r *Runner) resolveStructuredOutput(ctx context.Context, def *domain.Prompt
return nil, fmt.Errorf("%w: failed to load json schema for structured output: %v", ErrValidation, err)
}
return structuredOutputSpec(def, schemaDoc), nil
}
func structuredOutputSpec(def *domain.PromptDefinition, schemaDocument any) *domain.StructuredOutputSpec {
return &domain.StructuredOutputSpec{
Type: domain.StructuredOutputJSONSchema,
JSONSchema: &domain.StructuredOutputJSONSpec{
Name: deriveStructuredSchemaName(def.ID, def.Version),
Strict: true,
Schema: schemaDoc,
Schema: schemaDocument,
},
}, nil
}
}
func deriveStructuredSchemaName(promptID string, promptVersion string) string {