Tighten prepared execution credential handling

2026-07-30 19:06:17 +00:00
parent 6112c2af0c
commit 2ba0146e5d
11 changed files with 180 additions and 56 deletions
--- a/backends.go
+++ b/backends.go
@@ -40,12 +40,12 @@ type Backend struct {
 	// calls allowed for this backend within one Engine. Zero leaves the backend
 	// unlimited. A negative value makes NewEngine fail with ErrInvalidConfig.
 	ConcurrencyLimit int
-	// QueueCapacity controls how many additional Run calls may be admitted
-	// beyond ConcurrencyLimit. Nil uses 1024 when ConcurrencyLimit is positive;
-	// a pointer uses its exact value, including zero. The pointed-to value must
-	// be non-negative, and QueueCapacity must be nil when ConcurrencyLimit is
-	// zero. Their sum must fit in an int. WithBackend copies the value and does
-	// not retain the pointer.
+	// QueueCapacity controls how many additional Run or RunPrepared calls may
+	// be admitted beyond ConcurrencyLimit. Nil uses 1024 when ConcurrencyLimit
+	// is positive; a pointer uses its exact value, including zero. The pointed-to
+	// value must be non-negative, and QueueCapacity must be nil when
+	// ConcurrencyLimit is zero. Their sum must fit in an int. WithBackend copies
+	// the value and does not retain the pointer.
 	QueueCapacity *int
 }

--- a/doc.go
+++ b/doc.go
@@ -3,23 +3,25 @@
 //
 // Applications construct an [Engine] with [NewEngine], select filesystem or
 // in-memory sources and optional engine-scoped [Backend] registrations, and
-// call [Engine.Prepare] or [Engine.Run]. Concrete registries, repositories,
-// validators, and the built-in OpenAI-compatible client remain internal
-// implementation details.
+// call [Engine.Prepare], [Engine.PrepareExecution], [Engine.Run], or
+// [Engine.RunPrepared]. Concrete registries, repositories, validators, and the
+// built-in OpenAI-compatible client remain internal implementation details.
 //
 // # Concurrency and ownership
 //
-// An Engine supports concurrent Prepare and Run calls. Engine-local backend
-// policies bound admitted Run calls and model generations where configured,
-// while different backend pools and unlimited backends continue independently.
-// An injected [LLMClient] or [ArtifactReader] can therefore still receive
-// concurrent calls and must be safe for that use.
+// An Engine supports concurrent Prepare, PrepareExecution, Run, and RunPrepared
+// calls. Engine-local backend policies bound admitted Run and RunPrepared calls
+// and model generations where configured, while different backend pools and
+// unlimited backends continue independently. An injected [LLMClient] or
+// [ArtifactReader] can therefore still receive concurrent calls and must be
+// safe for that use.
 //
-// NewEngine copies in-memory profiles and backend definitions. Prepare and Run
-// copy request maps, slices, pointer values, and JSON-compatible extra
-// parameters before using them. Returned values and values passed to extension
-// interfaces are likewise isolated from engine state. Callers own those copies
-// and may mutate them after the call that supplied or returned them.
+// NewEngine copies in-memory profiles and backend definitions. Prepare,
+// PrepareExecution, and Run copy request maps, slices, pointer values, and
+// JSON-compatible extra parameters before using them. Returned values and
+// values passed to extension interfaces are likewise isolated from engine
+// state. Callers own those copies and may mutate them after the call that
+// supplied or returned them.
 //
 // # Security and sensitive data
 //
@@ -45,10 +47,11 @@
 // [GenerateResponse], [ExecutionTargetPresence], and the string value types
 // used by those values.
 //
-// Construction values, including [Config], [Backend], [RunRequest],
-// [ArtifactRef], [ExecutionTargetOverride], [Profile], and
-// [OpenAICompatibleProfileConfig], do not have stable JSON representations.
-// Direct API keys are nevertheless excluded from JSON for every public value.
+// Construction and handle values, including [Config], [Backend], [RunRequest],
+// [ArtifactRef], [ExecutionTargetOverride], [Profile],
+// [OpenAICompatibleProfileConfig], and [PreparedExecution], do not have stable
+// JSON representations. Direct API keys are nevertheless excluded from JSON
+// for every public value.
 //
 // JSON timestamps use time.Time's RFC 3339 encoding and are omitted when zero.
 // PreparedRun and RunResult durations are encoded as integer milliseconds in
--- a/docs/internal/overview.md
+++ b/docs/internal/overview.md
@@ -19,7 +19,7 @@ contributor workflow and validation.
 | `internal/domain` | Defines internal framework values for requests, artifacts, prompt definitions, profiles, execution targets, rendering, generation, and validation. | [Domain declarations](../../internal/domain/domain.go) |
 | `internal/defaults` | Defines application-neutral framework constants and constructs the default execution target. It contains no CLI, server, or inbound HTTP limits. | [Framework defaults](../../internal/defaults/defaults.go) |
 | `internal/filecatalog` | Provides deterministic YAML discovery and path helpers for operating-system filesystems and `fs.FS` sources. | [File catalog](../../internal/filecatalog/catalog.go) |
-| `internal/jsonvalue` | Validates and deeply copies JSON-compatible extra-parameter trees while preserving supported concrete value types. | [JSON values](../../internal/jsonvalue/jsonvalue.go) |
+| `internal/jsonvalue` | Validates and deeply copies JSON-compatible extra-parameter and prepared-schema trees while preserving supported concrete value types. | [JSON values](../../internal/jsonvalue/jsonvalue.go) |
 | `internal/promptdef` | Loads strictly decoded, validated prompt definitions from filesystem and `fs.FS` sources, including version selection and contained file-backed message content. | [Framework formats](../formats.md), [prompt-definition repository](../../internal/promptdef/filesystem_repository.go) |
 | `internal/profile` | Loads strictly decoded, validated execution profiles, including backend selection, from filesystem and `fs.FS` sources and composes repositories with error-preserving fallback. | [Framework formats](../formats.md), [profile repositories](../../internal/profile/filesystem_repository.go) |
 | `internal/profile/builtin` | Embeds the built-in profile catalog, whose entries select OpenRouter, and combines it with an optional primary repository. | [Built-in catalog](../formats.md#built-in-profile-catalog), [repository](../../internal/profile/builtin/repository.go) |
--- a/engine.go
+++ b/engine.go
@@ -149,12 +149,14 @@ type engineOptions struct {
 	artifactSource  bool
 }

-// WithLLMClient replaces the built-in model client used by [Engine.Run].
+// WithLLMClient replaces the built-in model client used by [Engine.Run] and
+// [Engine.RunPrepared].
 //
 // A nil client makes NewEngine fail with ErrInvalidConfig. The Engine schedules
 // Generate calls according to the selected backend's capacity policy, but the
 // client may still be called concurrently across different backend pools or for
-// unlimited backends. The client is not used by [Engine.Prepare].
+// unlimited backends. The client is not used by [Engine.Prepare] or
+// [Engine.PrepareExecution].
 func WithLLMClient(client LLMClient) Option {
 	return optionFunc(func(options *engineOptions) error {
 		if client == nil {
--- a/internal/jsonvalue/jsonvalue.go
+++ b/internal/jsonvalue/jsonvalue.go
@@ -1,5 +1,5 @@
 // Package jsonvalue validates and defensively copies JSON-compatible value
-// trees used by public configuration and request boundaries.
+// trees used by configuration, request, and prepared-state boundaries.
 package jsonvalue

 import (
--- a/internal/usecase/prepared_execution.go
+++ b/internal/usecase/prepared_execution.go
@@ -31,6 +31,7 @@ type PreparedExecution struct {
 type preparedExecutionPayload struct {
 	prepared   *domain.PreparedRun
 	validation validate.PreparedValidation
+	directKey  string
 }

 // PrepareExecution completes preparation without generation or admission and
@@ -61,18 +62,13 @@ func (r *Runner) PrepareExecution(ctx context.Context, req domain.RunRequest) (*

 	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,
@@ -81,6 +77,7 @@ func (r *Runner) PrepareExecution(ctx context.Context, req domain.RunRequest) (*
 		payload: &preparedExecutionPayload{
 			prepared:   executionSnapshot,
 			validation: validationPlan,
+			directKey:  state.effectiveModel.APIKey,
 		},
 	}, nil
 }
@@ -175,7 +172,7 @@ func (r *Runner) RunPrepared(ctx context.Context, prepared *PreparedExecution) (
 	start := time.Now().UTC()

 	target := payload.prepared.EffectiveModelParams
-	if err := validateAPIKey(target.APIKeyEnv, target.APIKey, target.APIKeyRequired); err != nil {
+	if err := validateAPIKey(target.APIKeyEnv, payload.directKey, target.APIKeyRequired); err != nil {
 		return nil, fmt.Errorf("%w: %w", ErrInvalidRequest, err)
 	}

@@ -185,7 +182,7 @@ func (r *Runner) RunPrepared(ctx context.Context, prepared *PreparedExecution) (
 	}
 	defer release()

-	return r.executePreparedRun(ctx, payload.prepared, runID, start, func(
+	return r.executePreparedRun(ctx, payload.prepared, payload.directKey, runID, start, func(
 		ctx context.Context,
 		artifact *domain.Artifact,
 		attemptsUsed int,
@@ -224,6 +221,7 @@ func (p *preparedExecutionPayload) clear() {
 	}
 	p.prepared = nil
 	p.validation = nil
+	p.directKey = ""
 }

 type noOpPreparedValidation struct {
--- a/internal/usecase/prepared_execution_test.go
+++ b/internal/usecase/prepared_execution_test.go
@@ -210,6 +210,48 @@ func TestRunnerRunPreparedRechecksEnvironmentCredentialBeforeAdmission(t *testin
 	}
 }

+func TestRunnerRunPreparedKeepsDirectCredentialOutOfMetadata(t *testing.T) {
+	const directKey = "direct-prepared-test-key"
+
+	profile := defaultExecutionProfile()
+	profile.APIKeyRequired = true
+	client := &fakeLLM{resp: &domain.GenerateResponse{Content: "ok"}}
+	runner := NewRunner(
+		&fakePromptRepo{def: promptDef(domain.FormatText, domain.ValidationNone, 0)},
+		&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": profile}},
+		nil,
+		defaultArtifactReader(),
+		defaultRenderer(),
+		client,
+		&recordingValidationPreparer{plan: &recordingPreparedValidation{}},
+		nil,
+	)
+
+	prepared, err := runner.PrepareExecution(context.Background(), domain.RunRequest{
+		PromptID:  "p",
+		ProfileID: "exec",
+		APIKey:    directKey,
+		Inputs:    singleInputRef(),
+	})
+	if err != nil {
+		t.Fatalf("prepare execution: %v", err)
+	}
+	if details := prepared.Details(); details.EffectiveModelParams.APIKey != "" {
+		t.Fatal("prepared details retained direct credential")
+	}
+
+	result, err := runner.RunPrepared(context.Background(), prepared)
+	if err != nil {
+		t.Fatalf("run prepared: %v", err)
+	}
+	if client.lastReq.Target.APIKey != directKey {
+		t.Fatal("generation did not receive direct credential")
+	}
+	if result.EffectiveModelParams.APIKey != "" {
+		t.Fatal("run result retained direct credential")
+	}
+}
+
 func TestRunnerRunPreparedUsesFrozenValidationForInitialAndRepairOutputs(t *testing.T) {
 	validator := &recordingValidationPreparer{
 		plan: &recordingPreparedValidation{
--- a/internal/usecase/runner.go
+++ b/internal/usecase/runner.go
@@ -141,8 +141,9 @@ func (r *Runner) Run(ctx context.Context, req domain.RunRequest) (*domain.RunRes
 	if err != nil {
 		return nil, err
 	}
+	directAPIKey := state.effectiveModel.APIKey

-	return r.executePreparedRun(ctx, prepared, runID, start, func(
+	return r.executePreparedRun(ctx, prepared, directAPIKey, runID, start, func(
 		ctx context.Context,
 		artifact *domain.Artifact,
 		attemptsUsed int,
@@ -160,13 +161,16 @@ type preparedValidationFunc func(
 func (r *Runner) executePreparedRun(
 	ctx context.Context,
 	prepared *domain.PreparedRun,
+	directAPIKey string,
 	runID string,
 	start time.Time,
 	validateArtifact preparedValidationFunc,
 ) (*domain.RunResult, error) {
+	executionTarget := prepared.EffectiveModelParams
+	executionTarget.APIKey = directAPIKey
 	genResp, err := r.llm.Generate(ctx, domain.GenerateRequest{
 		Prompt:           domain.RenderedPrompt{SessionID: prepared.SessionID, Messages: prepared.Messages},
-		Target:           prepared.EffectiveModelParams,
+		Target:           executionTarget,
 		TargetPresence:   prepared.TargetPresence,
 		StructuredOutput: prepared.StructuredOutput,
 	})
@@ -192,7 +196,7 @@ func (r *Runner) executePreparedRun(
 				PreviousOutput:   genResp.Content,
 				ValidationErrors: validationResult.Errors,
 				SessionID:        prepared.SessionID,
-				Target:           prepared.EffectiveModelParams,
+				Target:           executionTarget,
 				StructuredOutput: prepared.StructuredOutput,
 				Attempt:          attemptsUsed,
 				MaxAttempts:      prepared.OutputContract.RepairAttempts,
@@ -216,6 +220,7 @@ func (r *Runner) executePreparedRun(
 	}

 	end := time.Now().UTC()
+	executionTarget.APIKey = ""

 	return &domain.RunResult{
 		RunID:                runID,
@@ -231,7 +236,7 @@ func (r *Runner) executePreparedRun(
 		SelectedBackendID:    prepared.SelectedBackendID,
 		ModelName:            prepared.EffectiveModelParams.Model,
 		Endpoint:             prepared.EffectiveModelParams.Endpoint,
-		EffectiveModelParams: prepared.EffectiveModelParams,
+		EffectiveModelParams: executionTarget,
 		InputHashes:          prepared.InputHashes,
 		Usage:                genResp.Usage,
 		StartTime:            start,
@@ -375,13 +380,15 @@ func (r *Runner) completePreparationWithStructuredOutput(
 	}

 	end := time.Now().UTC()
+	effectiveModel := state.effectiveModel
+	effectiveModel.APIKey = ""
 	return &domain.PreparedRun{
 		PromptID:             state.definition.ID,
 		PromptVersion:        state.definition.Version,
 		PromptHash:           state.promptDefinitionHash,
 		SelectedProfileID:    state.selectedProfileID,
 		SelectedBackendID:    state.effectiveModel.BackendID,
-		EffectiveModelParams: state.effectiveModel,
+		EffectiveModelParams: effectiveModel,
 		TargetPresence:       state.targetPresence,
 		OutputContract:       state.effectiveContract,
 		StructuredOutput:     structuredOutput,
--- a/internal/usecase/runner_test.go
+++ b/internal/usecase/runner_test.go
@@ -1757,7 +1757,7 @@ func TestRunnerRunDirectAPIKeyBypassesMissingEnvAndReachesLLM(t *testing.T) {
 	llmClient := &fakeLLM{resp: &domain.GenerateResponse{Content: "ok"}}
 	runner := NewRunner(promptRepo, execRepo, nil, defaultArtifactReader(), defaultRenderer(), llmClient, nil, nil)

-	_, err := runner.Run(context.Background(), domain.RunRequest{
+	result, err := runner.Run(context.Background(), domain.RunRequest{
 		PromptID:  "p",
 		ProfileID: "exec",
 		APIKey:    directKey,
@@ -1769,6 +1769,9 @@ func TestRunnerRunDirectAPIKeyBypassesMissingEnvAndReachesLLM(t *testing.T) {
 	if llmClient.lastReq.Target.APIKey != directKey {
 		t.Fatalf("expected direct API key to reach LLM request")
 	}
+	if result.EffectiveModelParams.APIKey != "" {
+		t.Fatal("run result retained direct API key")
+	}
 	if llmClient.lastReq.Target.APIKeyEnv != "PROMPTKIT_MISSING_KEY" {
 		t.Fatalf("expected api_key_env name to remain on target, got %q", llmClient.lastReq.Target.APIKeyEnv)
 	}
--- a/prepared_execution_contract_test.go
+++ b/prepared_execution_contract_test.go
@@ -310,6 +310,75 @@ func TestPreparedExecutionConcurrentClaimAllowsOneGeneration(t *testing.T) {
 	}
 }

+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"
@@ -355,9 +424,6 @@ func TestPreparedExecutionDiscardAndFormattingDoNotExposePrivateState(t *testing
 	if err != nil {
 		t.Fatalf("marshal opaque handle: %v", err)
 	}
-	if string(payload) != "{}" {
-		t.Fatalf("opaque handle JSON=%s, want {}", payload)
-	}
 	assertPreparedPrivateValuesAbsent(t, string(payload), directCredential, renderedContent)

 	detailsBefore := prepared.Details()
--- a/types.go
+++ b/types.go
@@ -51,7 +51,8 @@ const (
 	// ValidationPassed means the generated output satisfied its contract.
 	ValidationPassed ValidationStatus = "passed"
 	// ValidationFailed means validation completed and rejected the generated
-	// output. Engine.Run returns this status in a result, not as an error.
+	// output. Engine.Run and Engine.RunPrepared return this status in a result,
+	// not as an error.
 	ValidationFailed ValidationStatus = "failed"
 	// ValidationSkipped means ValidationNone selected no content check.
 	ValidationSkipped ValidationStatus = "skipped"
@@ -267,10 +268,11 @@ type Artifact struct {
 // ArtifactReader resolves a prompt input reference into its content.
 //
 // Read may be called concurrently. It must honor ctx cancellation to make
-// Prepare and Run responsive to cancellation. The engine passes a copied ref
-// and immediately copies the returned Artifact.Body; it does not retain either
-// value. Readers supply artifact metadata, and the engine assigns an input-map
-// name only when the returned artifact name is empty.
+// Prepare, PrepareExecution, and Run responsive to cancellation. The engine
+// passes a copied ref and immediately copies the returned Artifact.Body; it
+// does not retain either value. Readers supply artifact metadata, and the
+// engine assigns an input-map name only when the returned artifact name is
+// empty.
 //
 // An injected reader owns any application-specific path containment,
 // authorization, content-size, and content-type policy. It must protect
@@ -567,25 +569,26 @@ type StructuredOutputJSONSpec struct {
 	Schema any `json:"schema"`
 }

-// LLMClient executes rendered prompts for [Engine.Run].
+// LLMClient executes rendered prompts for [Engine.Run] and
+// [Engine.RunPrepared].
 //
 // Generate is scheduled according to the resolved backend's capacity policy.
 // It may still be called concurrently for different backend pools or unlimited
 // backends. Cancellation while waiting for capacity can prevent Generate from
 // being called. Once invoked, it must honor context cancellation to make Run
-// responsive to cancellation. The request and all nested maps, slices, and
-// pointers are client-owned copies and may be mutated or retained without
-// affecting engine state.
+// and RunPrepared responsive to cancellation. The request and all nested maps,
+// slices, and pointers are client-owned copies and may be mutated or retained
+// without affecting engine state.
 //
 // Generate receives rendered messages and may receive a direct API key. A
 // client must protect those values and any raw output in its logging, storage,
 // and retained copies. It is responsible for the cancellation behavior of any
 // work it starts and for synchronizing access to retained or shared data.
 //
-// A returned error makes Run return ErrLLMGenerate while preserving the client
-// error through errors.Is. A nil response with a nil error also produces
-// ErrLLMGenerate. Promptkit copies the non-nil response before returning from
-// Run.
+// A returned error makes Run or RunPrepared return ErrLLMGenerate while
+// preserving the client error through errors.Is. A nil response with a nil
+// error also produces ErrLLMGenerate. Promptkit copies the non-nil response
+// before returning from either method.
 type LLMClient interface {
 	Generate(context.Context, GenerateRequest) (*GenerateResponse, error)
 }