eric/promptkit
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Tighten prepared execution credential handling

Browse Source
This commit is contained in:
Eric Rakestraw
2026-07-30 19:06:17 +00:00
parent 6112c2af0c
commit 2ba0146e5d
11 changed files with 180 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
internal/jsonvalue/jsonvalue.go
View File

@@ -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 (

12
internal/usecase/prepared_execution.go
View File

@@ -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 {

42
internal/usecase/prepared_execution_test.go
View File

@@ -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{

17
internal/usecase/runner.go
View File

@@ -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,

5
internal/usecase/runner_test.go
View File

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