| 1 | package extension |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "crypto/sha256" |
| 6 | "encoding/hex" |
| 7 | "encoding/json" |
| 8 | "sort" |
| 9 | "testing" |
| 10 | |
| 11 | "reasonix/internal/provider" |
| 12 | ) |
| 13 | |
| 14 | // hashTestSnapshot builds a snapshot with two tools, an interceptor, and a |
| 15 | // claimed slot so every accessor has content to defend. |
| 16 | func hashTestSnapshot(t *testing.T, prompt string, readDesc string) *RuntimeSnapshot { |
| 17 | t.Helper() |
| 18 | b := NewBuilder().WithSystemPrompt(prompt).WithGeneration(3) |
| 19 | b.AddContributor(staticContributor("mix", |
| 20 | Contribution{Kind: KindTool, ID: "read_file", Source: src(ScopeBuiltin, "", "builtin"), Payload: schemaPayload("read_file", readDesc)}, |
| 21 | Contribution{Kind: KindTool, ID: "bash", Source: src(ScopeBuiltin, "", "builtin"), Payload: schemaPayload("bash", "run commands")}, |
| 22 | Contribution{Kind: KindInterceptor, ID: string(PointToolBefore), Priority: 0, Source: src(ScopePlugin, "pa", "plugin"), Payload: "i"}, |
| 23 | Contribution{Kind: KindStrategy, ID: "strat", Source: src(ScopePlugin, "pa", "plugin"), Payload: claimPayload{slots: []Slot{SlotContext}}}, |
| 24 | )) |
| 25 | snap, _, err := b.Build(context.Background()) |
| 26 | if err != nil { |
| 27 | t.Fatalf("Build failed: %v", err) |
| 28 | } |
| 29 | return snap |
| 30 | } |
| 31 | |
| 32 | // TestSnapshotImmutable: a snapshot is shared across turns and frontends; |
| 33 | // every accessor must return copies so no consumer can edit another |
| 34 | // consumer's view. |
| 35 | func TestSnapshotImmutable(t *testing.T) { |
| 36 | snap := hashTestSnapshot(t, "prompt", "read files") |
| 37 | |
| 38 | schemas := snap.ToolSchemas() |
| 39 | originalDesc := schemas[0].Description |
| 40 | schemas[0].Description = "mutated" |
| 41 | if snap.ToolSchemas()[0].Description != originalDesc { |
| 42 | t.Fatal("mutating ToolSchemas() result changed the snapshot") |
| 43 | } |
| 44 | |
| 45 | chains := snap.InterceptorChain() |
| 46 | chains[PointToolBefore][0].Priority = -999 |
| 47 | delete(chains, PointToolBefore) |
| 48 | if len(snap.InterceptorChain()) != 1 { |
| 49 | t.Fatal("deleting from InterceptorChain() result changed the snapshot") |
| 50 | } |
| 51 | if snap.InterceptorChain()[PointToolBefore][0].Priority != 0 { |
| 52 | t.Fatal("mutating a chained contribution changed the snapshot") |
| 53 | } |
| 54 | |
| 55 | repl := snap.Replacements() |
| 56 | repl[SlotSystemPrompt] = src(ScopePlugin, "evil", "plugin") |
| 57 | delete(repl, SlotContext) |
| 58 | if len(snap.Replacements()) != 1 { |
| 59 | t.Fatal("mutating Replacements() result changed the snapshot") |
| 60 | } |
| 61 | if _, ok := snap.Replacements()[SlotSystemPrompt]; ok { |
| 62 | t.Fatal("injected a slot into the snapshot via the returned map") |
| 63 | } |
| 64 | |
| 65 | // The frozen catalog refuses growth. |
| 66 | defer func() { |
| 67 | if recover() == nil { |
| 68 | t.Fatal("Add on snapshot catalog did not panic") |
| 69 | } |
| 70 | }() |
| 71 | snap.Catalog().Add(Contribution{Kind: KindTool, ID: "evil", Source: src(ScopeBuiltin, "", "builtin")}) |
| 72 | } |
| 73 | |
| 74 | // canonicalCacheHash recomputes the documented hash form independently: |
| 75 | // sha256-hex over the JSON of {"systemPrompt":..., "toolSchemas":[...]} with |
| 76 | // schemas sorted by (name, description, parameters). The test must match the |
| 77 | // implementation without calling it, so a drift in either is caught. |
| 78 | func canonicalCacheHash(t *testing.T, prompt string, schemas []provider.ToolSchema) (systemHash, toolsHash, cacheHash string) { |
| 79 | t.Helper() |
| 80 | sorted := make([]provider.ToolSchema, len(schemas)) |
| 81 | copy(sorted, schemas) |
| 82 | sort.Slice(sorted, func(i, j int) bool { |
| 83 | if sorted[i].Name != sorted[j].Name { |
| 84 | return sorted[i].Name < sorted[j].Name |
| 85 | } |
| 86 | if sorted[i].Description != sorted[j].Description { |
| 87 | return sorted[i].Description < sorted[j].Description |
| 88 | } |
| 89 | return string(sorted[i].Parameters) < string(sorted[j].Parameters) |
| 90 | }) |
| 91 | sum := sha256.Sum256([]byte(prompt)) |
| 92 | systemHash = hex.EncodeToString(sum[:]) |
| 93 | toolsJSON, err := json.Marshal(sorted) |
| 94 | if err != nil { |
| 95 | t.Fatalf("marshal schemas: %v", err) |
| 96 | } |
| 97 | sum = sha256.Sum256(toolsJSON) |
| 98 | toolsHash = hex.EncodeToString(sum[:]) |
| 99 | combined, err := json.Marshal(struct { |
| 100 | SystemPrompt string `json:"systemPrompt"` |
| 101 | ToolSchemas []provider.ToolSchema `json:"toolSchemas"` |
| 102 | }{SystemPrompt: prompt, ToolSchemas: sorted}) |
| 103 | if err != nil { |
| 104 | t.Fatalf("marshal combined: %v", err) |
| 105 | } |
| 106 | sum = sha256.Sum256(combined) |
| 107 | cacheHash = hex.EncodeToString(sum[:]) |
| 108 | return systemHash, toolsHash, cacheHash |
| 109 | } |
| 110 | |
| 111 | // TestCacheHashCanonicalForm: CacheHash is a protocol — other tooling |
| 112 | // recomputes it to detect prefix drift — so it must equal the independently |
| 113 | // specified canonical form, byte for byte. |
| 114 | func TestCacheHashCanonicalForm(t *testing.T) { |
| 115 | snap := hashTestSnapshot(t, "prompt v1", "read files") |
| 116 | wantSystem, wantTools, wantCache := canonicalCacheHash(t, snap.SystemPrompt(), snap.ToolSchemas()) |
| 117 | if snap.CacheHash() != wantCache { |
| 118 | t.Fatalf("CacheHash = %s, want %s", snap.CacheHash(), wantCache) |
| 119 | } |
| 120 | gotSystem, gotTools := snap.CacheShape() |
| 121 | if gotSystem != wantSystem || gotTools != wantTools { |
| 122 | t.Fatalf("CacheShape = (%s, %s), want (%s, %s)", gotSystem, gotTools, wantSystem, wantTools) |
| 123 | } |
| 124 | } |
| 125 | |
| 126 | // TestCacheHashStability: identical inputs hash identically (the provider |
| 127 | // cache survives a snapshot rebuild), and any provider-visible change moves |
| 128 | // the hash. |
| 129 | func TestCacheHashStability(t *testing.T) { |
| 130 | a := hashTestSnapshot(t, "prompt v1", "read files") |
| 131 | b := hashTestSnapshot(t, "prompt v1", "read files") |
| 132 | if a.CacheHash() != b.CacheHash() { |
| 133 | t.Fatal("identical builds produced different CacheHash") |
| 134 | } |
| 135 | |
| 136 | changed := hashTestSnapshot(t, "prompt v1", "read files thoroughly") |
| 137 | if changed.CacheHash() == a.CacheHash() { |
| 138 | t.Fatal("changing one schema description did not change CacheHash") |
| 139 | } |
| 140 | _, aTools := a.CacheShape() |
| 141 | changedSystem, changedTools := changed.CacheShape() |
| 142 | if changedTools == aTools { |
| 143 | t.Fatal("changing one schema description did not change the tools hash") |
| 144 | } |
| 145 | if aSys, _ := a.CacheShape(); changedSystem != aSys { |
| 146 | t.Fatal("a tools-only change moved the system hash") |
| 147 | } |
| 148 | |
| 149 | changedPrompt := hashTestSnapshot(t, "prompt v2", "read files") |
| 150 | if changedPrompt.CacheHash() == a.CacheHash() { |
| 151 | t.Fatal("changing the system prompt did not change CacheHash") |
| 152 | } |
| 153 | } |
| 154 |