From c95d6fcfece78ca68e838e6abc5334651b615f2f Mon Sep 17 00:00:00 2001 From: Eric Rakestraw Date: Thu, 2 Jul 2026 23:03:50 +0000 Subject: [PATCH] Preserve cache control in rendered prompts --- internal/domain/domain.go | 5 +- internal/domain/prepared_run_test.go | 49 +++++++++++++++ internal/format/prepared_run.go | 7 +++ internal/format/prepared_run_test.go | 93 ++++++++++++++++++++++++++++ internal/prompt/go_renderer.go | 13 +++- internal/prompt/renderer_test.go | 60 ++++++++++++++++++ internal/usecase/runner.go | 8 +++ internal/usecase/runner_test.go | 55 ++++++++++++++++ 8 files changed, 286 insertions(+), 4 deletions(-) diff --git a/internal/domain/domain.go b/internal/domain/domain.go index 8a0a4d0..94b98dd 100644 --- a/internal/domain/domain.go +++ b/internal/domain/domain.go @@ -194,8 +194,9 @@ type RenderedPrompt struct { // RenderedMessage is a single message in a rendered prompt. type RenderedMessage struct { - Role string `json:"role"` - Content string `json:"content"` + Role string `json:"role"` + Content string `json:"content"` + CacheControl *CacheControl `json:"cache_control,omitempty"` } // GenerateRequest is the internal request passed to the LLM client. diff --git a/internal/domain/prepared_run_test.go b/internal/domain/prepared_run_test.go index 7d90883..9fe1923 100644 --- a/internal/domain/prepared_run_test.go +++ b/internal/domain/prepared_run_test.go @@ -53,3 +53,52 @@ func TestPreparedRunJSONDoesNotIncludeSecretValues(t *testing.T) { } } } + +func TestPreparedRunJSONIncludesMessageCacheControlOnlyWhenPresent(t *testing.T) { + prepared := PreparedRun{ + PromptID: "prompt.id", + SelectedProfileID: "local-fast", + EffectiveModelParams: ExecutionTarget{ + Endpoint: "http://llm/v1", + Model: "gpt-test", + }, + RenderedPromptHash: "rendered-hash", + Messages: []RenderedMessage{ + { + Role: "system", + Content: "You are helpful.", + CacheControl: &CacheControl{ + Type: CacheControlEphemeral, + TTL: "1h", + }, + }, + {Role: "user", Content: "Summarize this."}, + }, + } + + b, err := json.Marshal(prepared) + if err != nil { + t.Fatalf("marshal failed: %v", err) + } + + var decoded struct { + Messages []map[string]any `json:"messages"` + } + if err := json.Unmarshal(b, &decoded); err != nil { + t.Fatalf("unmarshal failed: %v", err) + } + if len(decoded.Messages) != 2 { + t.Fatalf("expected 2 messages, got %d", len(decoded.Messages)) + } + + cacheControl, ok := decoded.Messages[0]["cache_control"].(map[string]any) + if !ok { + t.Fatalf("expected cache_control on first message, got %#v", decoded.Messages[0]) + } + if cacheControl["type"] != string(CacheControlEphemeral) || cacheControl["ttl"] != "1h" { + t.Fatalf("unexpected cache_control payload: %#v", cacheControl) + } + if _, ok := decoded.Messages[1]["cache_control"]; ok { + t.Fatalf("expected second message to omit cache_control, got %#v", decoded.Messages[1]) + } +} diff --git a/internal/format/prepared_run.go b/internal/format/prepared_run.go index 9800b02..690af67 100644 --- a/internal/format/prepared_run.go +++ b/internal/format/prepared_run.go @@ -151,6 +151,13 @@ func (textPreparedRunFormatter) Format(prepared *domain.PreparedRun) ([]byte, er messages := byRole[role] for i, msg := range messages { fmt.Fprintf(&b, " - message: %d\n", i+1) + if msg.CacheControl != nil { + fmt.Fprintf(&b, " cache_control: %s", msg.CacheControl.Type) + if msg.CacheControl.TTL != "" { + fmt.Fprintf(&b, " ttl=%s", msg.CacheControl.TTL) + } + fmt.Fprintln(&b) + } fmt.Fprintln(&b, " content: |") content := msg.Content if content == "" { diff --git a/internal/format/prepared_run_test.go b/internal/format/prepared_run_test.go index 48191bd..4527062 100644 --- a/internal/format/prepared_run_test.go +++ b/internal/format/prepared_run_test.go @@ -62,6 +62,58 @@ func TestTextFormatterDoesNotIncludeResolvedAPIKeyValue(t *testing.T) { } } +func TestTextFormatterIncludesMessageCacheControlBeforeContent(t *testing.T) { + prepared := samplePreparedRun() + prepared.Messages = []domain.RenderedMessage{ + { + Role: "system", + Content: "System guidance.", + CacheControl: &domain.CacheControl{ + Type: domain.CacheControlEphemeral, + TTL: "1h", + }, + }, + {Role: "user", Content: "Summarize the transcript."}, + } + + out, err := FormatPreparedRun(prepared, PreparedRunFormatText) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + s := string(out) + if !strings.Contains(s, " system:\n - message: 1\n cache_control: ephemeral ttl=1h\n content: |") { + t.Fatalf("expected system message cache control before content, got:\n%s", s) + } + if strings.Count(s, "cache_control:") != 1 { + t.Fatalf("expected exactly one cache_control line, got:\n%s", s) + } +} + +func TestTextFormatterOmitsEmptyCacheControlTTL(t *testing.T) { + prepared := samplePreparedRun() + prepared.Messages = []domain.RenderedMessage{ + { + Role: "system", + Content: "System guidance.", + CacheControl: &domain.CacheControl{ + Type: domain.CacheControlEphemeral, + }, + }, + } + + out, err := FormatPreparedRun(prepared, PreparedRunFormatText) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + s := string(out) + if !strings.Contains(s, " cache_control: ephemeral\n") { + t.Fatalf("expected cache_control line without ttl, got:\n%s", s) + } + if strings.Contains(s, "ttl=") { + t.Fatalf("expected empty ttl to be omitted, got:\n%s", s) + } +} + func TestJSONFormatterEmitsValidJSONAndIncludesPreparedRunFields(t *testing.T) { prepared := samplePreparedRun() @@ -98,6 +150,47 @@ func TestJSONFormatterEmitsValidJSONAndIncludesPreparedRunFields(t *testing.T) { } } +func TestJSONFormatterIncludesMessageCacheControlOnlyWhenPresent(t *testing.T) { + prepared := samplePreparedRun() + prepared.Messages = []domain.RenderedMessage{ + { + Role: "system", + Content: "System guidance.", + CacheControl: &domain.CacheControl{ + Type: domain.CacheControlEphemeral, + TTL: "1h", + }, + }, + {Role: "user", Content: "Summarize the transcript."}, + } + + out, err := FormatPreparedRun(prepared, PreparedRunFormatJSON) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + + var decoded struct { + Messages []map[string]any `json:"messages"` + } + if err := json.Unmarshal(out, &decoded); err != nil { + t.Fatalf("expected valid json output, got %v", err) + } + if len(decoded.Messages) != 2 { + t.Fatalf("expected 2 messages, got %d", len(decoded.Messages)) + } + + cacheControl, ok := decoded.Messages[0]["cache_control"].(map[string]any) + if !ok { + t.Fatalf("expected first message cache_control, got %#v", decoded.Messages[0]) + } + if cacheControl["type"] != string(domain.CacheControlEphemeral) || cacheControl["ttl"] != "1h" { + t.Fatalf("unexpected cache_control payload: %#v", cacheControl) + } + if _, ok := decoded.Messages[1]["cache_control"]; ok { + t.Fatalf("expected second message to omit cache_control, got %#v", decoded.Messages[1]) + } +} + func TestJSONFormatterDoesNotIncludeResolvedAPIKeyValue(t *testing.T) { const secret = "super-secret-api-key" t.Setenv("SCRIPTORIUM_API_KEY", secret) diff --git a/internal/prompt/go_renderer.go b/internal/prompt/go_renderer.go index 92c273d..dc317fc 100644 --- a/internal/prompt/go_renderer.go +++ b/internal/prompt/go_renderer.go @@ -75,8 +75,9 @@ func (r *goRenderer) Render(ctx context.Context, definition *domain.PromptDefini } renderedMessages = append(renderedMessages, domain.RenderedMessage{ - Role: tmplMsg.Role, - Content: buf.String(), + Role: tmplMsg.Role, + Content: buf.String(), + CacheControl: cloneCacheControl(tmplMsg.CacheControl), }) } @@ -84,3 +85,11 @@ func (r *goRenderer) Render(ctx context.Context, definition *domain.PromptDefini Messages: renderedMessages, }, nil } + +func cloneCacheControl(in *domain.CacheControl) *domain.CacheControl { + if in == nil { + return nil + } + out := *in + return &out +} diff --git a/internal/prompt/renderer_test.go b/internal/prompt/renderer_test.go index ca3c2a3..5b360a5 100644 --- a/internal/prompt/renderer_test.go +++ b/internal/prompt/renderer_test.go @@ -78,6 +78,66 @@ func TestGoRenderer_Render(t *testing.T) { } }) + t.Run("copying cache control to rendered messages", func(t *testing.T) { + def := &domain.PromptDefinition{ + Inputs: []domain.PromptInput{{Name: "transcript", Required: true}}, + Templates: []domain.PromptMessageTemplate{ + { + Role: "system", + Content: "You are concise.", + CacheControl: &domain.CacheControl{ + Type: domain.CacheControlEphemeral, + TTL: "1h", + }, + }, + {Role: "user", Content: "Analyze this: {{input \"transcript\"}}"}, + }, + } + + res, err := renderer.Render(ctx, def, inputs, vars) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(res.Messages) != 2 { + t.Fatalf("expected 2 messages, got %d", len(res.Messages)) + } + if res.Messages[0].CacheControl == nil { + t.Fatal("expected rendered cache control") + } + if res.Messages[0].CacheControl.Type != domain.CacheControlEphemeral { + t.Fatalf("unexpected cache control type: %q", res.Messages[0].CacheControl.Type) + } + if res.Messages[0].CacheControl.TTL != "1h" { + t.Fatalf("unexpected cache control ttl: %q", res.Messages[0].CacheControl.TTL) + } + if res.Messages[1].CacheControl != nil { + t.Fatalf("expected no cache control on second message, got %#v", res.Messages[1].CacheControl) + } + }) + + t.Run("rendered cache control does not alias source template", func(t *testing.T) { + source := &domain.CacheControl{Type: domain.CacheControlEphemeral, TTL: "1h"} + def := &domain.PromptDefinition{ + Inputs: []domain.PromptInput{{Name: "transcript", Required: true}}, + Templates: []domain.PromptMessageTemplate{ + {Role: "system", Content: "You are concise.", CacheControl: source}, + }, + } + + res, err := renderer.Render(ctx, def, inputs, vars) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if res.Messages[0].CacheControl == source { + t.Fatal("expected rendered cache control to be cloned") + } + + res.Messages[0].CacheControl.TTL = "" + if source.TTL != "1h" { + t.Fatalf("source cache control was mutated, ttl=%q", source.TTL) + } + }) + t.Run("accessing vars", func(t *testing.T) { def := &domain.PromptDefinition{ Inputs: []domain.PromptInput{{Name: "transcript", Required: true}}, diff --git a/internal/usecase/runner.go b/internal/usecase/runner.go index 4091403..508acb2 100644 --- a/internal/usecase/runner.go +++ b/internal/usecase/runner.go @@ -428,6 +428,14 @@ func hashRenderedPrompt(p domain.RenderedPrompt) string { b.WriteString(msg.Role) b.WriteByte('\n') b.WriteString(msg.Content) + if msg.CacheControl != nil { + b.WriteString("\ncache_control.type=") + b.WriteString(string(msg.CacheControl.Type)) + if msg.CacheControl.TTL != "" { + b.WriteString("\ncache_control.ttl=") + b.WriteString(msg.CacheControl.TTL) + } + } b.WriteString("\n---\n") } h := sha256.Sum256([]byte(b.String())) diff --git a/internal/usecase/runner_test.go b/internal/usecase/runner_test.go index 580db21..c9d6447 100644 --- a/internal/usecase/runner_test.go +++ b/internal/usecase/runner_test.go @@ -577,6 +577,61 @@ func TestDeriveStructuredSchemaName(t *testing.T) { } } +func TestHashRenderedPromptIncludesCacheControlWhenPresent(t *testing.T) { + uncached := domain.RenderedPrompt{Messages: []domain.RenderedMessage{ + {Role: "system", Content: "sys"}, + {Role: "user", Content: "usr"}, + }} + wantLegacyHash := hashString("system\nsys\n---\nuser\nusr\n---\n") + if got := hashRenderedPrompt(uncached); got != wantLegacyHash { + t.Fatalf("expected no-cache hash to preserve legacy input, got %q want %q", got, wantLegacyHash) + } + + withCache := domain.RenderedPrompt{Messages: []domain.RenderedMessage{ + { + Role: "system", + Content: "sys", + CacheControl: &domain.CacheControl{ + Type: domain.CacheControlEphemeral, + TTL: "1h", + }, + }, + {Role: "user", Content: "usr"}, + }} + alsoWithCache := domain.RenderedPrompt{Messages: []domain.RenderedMessage{ + { + Role: "system", + Content: "sys", + CacheControl: &domain.CacheControl{ + Type: domain.CacheControlEphemeral, + TTL: "1h", + }, + }, + {Role: "user", Content: "usr"}, + }} + withoutTTL := domain.RenderedPrompt{Messages: []domain.RenderedMessage{ + { + Role: "system", + Content: "sys", + CacheControl: &domain.CacheControl{ + Type: domain.CacheControlEphemeral, + }, + }, + {Role: "user", Content: "usr"}, + }} + + cachedHash := hashRenderedPrompt(withCache) + if cachedHash == hashRenderedPrompt(uncached) { + t.Fatal("expected cache control to change rendered prompt hash") + } + if cachedHash != hashRenderedPrompt(alsoWithCache) { + t.Fatal("expected identical cache control metadata to produce stable hash") + } + if cachedHash == hashRenderedPrompt(withoutTTL) { + t.Fatal("expected ttl changes to affect rendered prompt hash") + } +} + func TestRunnerRunSuccessful(t *testing.T) { promptRepo := &fakePromptRepo{def: promptDef(domain.FormatMarkdown, domain.ValidationBasic, 0)} execRepo := &fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": defaultExecutionProfile()}}