Add request API key support

2026-07-04 16:55:01 +00:00
parent 32e2433628
commit 3ad247039b
13 changed files with 284 additions and 8 deletions
--- a/engine_test.go
+++ b/engine_test.go
@@ -4,6 +4,8 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"net/http"
+	"net/http/httptest"
 	"os"
 	"path/filepath"
 	"reflect"
@@ -205,6 +207,7 @@ func TestRunSucceedsWithInjectedLLMClient(t *testing.T) {
 }

 func TestRunPassesPreparedRequestToInjectedLLMClient(t *testing.T) {
+	const directKey = "direct-injected-key"
 	fake := &fakeLLMClient{
 		response: &scriptorium.GenerateResponse{Content: "ok"},
 	}
@@ -214,6 +217,7 @@ func TestRunPassesPreparedRequestToInjectedLLMClient(t *testing.T) {

 	_, err := engine.Run(context.Background(), scriptorium.RunRequest{
 		PromptID: "generic.markdown_summary",
+		APIKey:   directKey,
 		Inputs: map[string]scriptorium.ArtifactRef{
 			"transcript": scriptorium.Inline("Rin opens the gate."),
 			"glossary":   scriptorium.Inline("gate: A guarded passage."),
@@ -244,6 +248,159 @@ func TestRunPassesPreparedRequestToInjectedLLMClient(t *testing.T) {
 	if req.StructuredOutput != nil {
 		t.Fatalf("did not expect structured output for markdown prompt: %+v", req.StructuredOutput)
 	}
+	if req.APIKey != directKey {
+		t.Fatalf("expected direct key on injected generate request")
+	}
+	payload, err := json.Marshal(req)
+	if err != nil {
+		t.Fatalf("expected generate request to marshal, got %v", err)
+	}
+	if strings.Contains(string(payload), directKey) {
+		t.Fatalf("generate request JSON leaked direct API key: %s", payload)
+	}
+}
+
+func TestRunUsesDirectAPIKeyWithDefaultLLMClient(t *testing.T) {
+	const directKey = "direct-public-key"
+	const missingEnv = "SCRIPTORIUM_PUBLIC_DIRECT_MISSING"
+	t.Setenv(missingEnv, "")
+
+	var gotAuth string
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		gotAuth = r.Header.Get("Authorization")
+		if r.URL.Path != "/v1/chat/completions" {
+			t.Errorf("unexpected path: %s", r.URL.Path)
+		}
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(`{
+  "choices": [{"message": {"role": "assistant", "content": "# Summary\n\nDone."}}],
+  "usage": {"prompt_tokens": 3, "completion_tokens": 4, "total_tokens": 7}
+}`))
+	}))
+	defer server.Close()
+
+	profileDir := t.TempDir()
+	writePublicProfileFileWithAPIKeyEnv(t, profileDir, "direct-auth", server.URL+"/v1", "test-model", missingEnv)
+	engine, err := scriptorium.NewEngine(scriptorium.Config{
+		PromptDir:  "./examples/prompts",
+		ProfileDir: profileDir,
+		SchemaDir:  "./examples/schemas",
+	})
+	if err != nil {
+		t.Fatalf("expected engine construction to succeed, got %v", err)
+	}
+
+	result, err := engine.Run(context.Background(), scriptorium.RunRequest{
+		PromptID:  "generic.markdown_summary",
+		ProfileID: "direct-auth",
+		APIKey:    directKey,
+		Inputs: map[string]scriptorium.ArtifactRef{
+			"transcript": scriptorium.Inline("Rin opens the gate."),
+			"glossary":   scriptorium.Inline("gate: A guarded passage."),
+		},
+	})
+	if err != nil {
+		t.Fatalf("expected run with direct API key to succeed, got %v", err)
+	}
+	if gotAuth != "Bearer "+directKey {
+		t.Fatalf("unexpected Authorization header: %q", gotAuth)
+	}
+	if result.Usage.TotalTokens != 7 {
+		t.Fatalf("unexpected usage: %+v", result.Usage)
+	}
+
+	payload, err := json.Marshal(result)
+	if err != nil {
+		t.Fatalf("expected run result to marshal, got %v", err)
+	}
+	if strings.Contains(string(payload), directKey) {
+		t.Fatalf("run result JSON leaked direct API key: %s", payload)
+	}
+}
+
+func TestPrepareDirectAPIKeyBypassesMissingEnvWithoutLeakingOrHashing(t *testing.T) {
+	const missingEnv = "SCRIPTORIUM_PUBLIC_PREPARE_MISSING"
+	const firstKey = "first-direct-key"
+	const secondKey = "second-direct-key"
+	t.Setenv(missingEnv, "")
+
+	profileDir := t.TempDir()
+	writePublicProfileFileWithAPIKeyEnv(t, profileDir, "direct-prepare", "http://localhost:8000/v1", "test-model", missingEnv)
+	engine, err := scriptorium.NewEngine(scriptorium.Config{
+		PromptDir:  "./examples/prompts",
+		ProfileDir: profileDir,
+		SchemaDir:  "./examples/schemas",
+	})
+	if err != nil {
+		t.Fatalf("expected engine construction to succeed, got %v", err)
+	}
+
+	baseReq := scriptorium.RunRequest{
+		PromptID:  "generic.markdown_summary",
+		ProfileID: "direct-prepare",
+		Inputs: map[string]scriptorium.ArtifactRef{
+			"transcript": scriptorium.Inline("Rin opens the gate."),
+			"glossary":   scriptorium.Inline("gate: A guarded passage."),
+		},
+	}
+	firstReq := baseReq
+	firstReq.APIKey = firstKey
+	firstPrepared, err := engine.Prepare(context.Background(), firstReq)
+	if err != nil {
+		t.Fatalf("expected prepare with direct API key to succeed, got %v", err)
+	}
+	secondReq := baseReq
+	secondReq.APIKey = secondKey
+	secondPrepared, err := engine.Prepare(context.Background(), secondReq)
+	if err != nil {
+		t.Fatalf("expected prepare with alternate direct API key to succeed, got %v", err)
+	}
+
+	if firstPrepared.PromptHash != secondPrepared.PromptHash {
+		t.Fatalf("direct API keys changed prompt hash: %q vs %q", firstPrepared.PromptHash, secondPrepared.PromptHash)
+	}
+	if firstPrepared.RenderedPromptHash != secondPrepared.RenderedPromptHash {
+		t.Fatalf("direct API keys changed rendered prompt hash: %q vs %q", firstPrepared.RenderedPromptHash, secondPrepared.RenderedPromptHash)
+	}
+
+	payload, err := json.Marshal(firstPrepared)
+	if err != nil {
+		t.Fatalf("expected prepared run to marshal, got %v", err)
+	}
+	if strings.Contains(string(payload), firstKey) {
+		t.Fatalf("prepared run JSON leaked direct API key: %s", payload)
+	}
+}
+
+func TestMissingCredentialsFailClearlyWhenProfileRequiresAuth(t *testing.T) {
+	const missingEnv = "SCRIPTORIUM_PUBLIC_AUTH_MISSING"
+	t.Setenv(missingEnv, "")
+
+	profileDir := t.TempDir()
+	writePublicProfileFileWithAPIKeyEnv(t, profileDir, "requires-auth", "http://localhost:8000/v1", "test-model", missingEnv)
+	engine, err := scriptorium.NewEngine(scriptorium.Config{
+		PromptDir:  "./examples/prompts",
+		ProfileDir: profileDir,
+		SchemaDir:  "./examples/schemas",
+	})
+	if err != nil {
+		t.Fatalf("expected engine construction to succeed, got %v", err)
+	}
+
+	_, err = engine.Prepare(context.Background(), scriptorium.RunRequest{
+		PromptID:  "generic.markdown_summary",
+		ProfileID: "requires-auth",
+		Inputs: map[string]scriptorium.ArtifactRef{
+			"transcript": scriptorium.Inline("Rin opens the gate."),
+			"glossary":   scriptorium.Inline("gate: A guarded passage."),
+		},
+	})
+	if !errors.Is(err, scriptorium.ErrInvalidRequest) {
+		t.Fatalf("expected invalid request for missing credentials, got %v", err)
+	}
+	if err == nil || !strings.Contains(err.Error(), missingEnv) {
+		t.Fatalf("expected missing env name in error, got %v", err)
+	}
 }

 func TestRunValidationFailureReturnsResult(t *testing.T) {
@@ -693,6 +850,18 @@ model: ` + model + `
 	}
 }

+func writePublicProfileFileWithAPIKeyEnv(t *testing.T, dir, id, endpoint, model, apiKeyEnv string) {
+	t.Helper()
+	data := `id: ` + id + `
+endpoint: ` + endpoint + `
+model: ` + model + `
+api_key_env: ` + apiKeyEnv + `
+`
+	if err := os.WriteFile(filepath.Join(dir, id+".yaml"), []byte(data), 0o644); err != nil {
+		t.Fatalf("failed to write profile fixture: %v", err)
+	}
+}
+
 type fakeLLMClient struct {
 	response *scriptorium.GenerateResponse
 	err      error