| 1 | package agent |
| 2 | |
| 3 | import ( |
| 4 | "crypto/sha256" |
| 5 | "encoding/json" |
| 6 | "fmt" |
| 7 | "sort" |
| 8 | |
| 9 | "reasonix/internal/event" |
| 10 | "reasonix/internal/provider" |
| 11 | ) |
| 12 | |
| 13 | // PrefixShape hashes the portions of the request prefix that influence |
| 14 | // provider-side prompt-cache reuse. Comparing snapshots across turns |
| 15 | // lets us explain *why* a cache miss happened. |
| 16 | type PrefixShape struct { |
| 17 | SystemHash string |
| 18 | ToolsHash string |
| 19 | PrefixHash string |
| 20 | LogRewriteVersion int |
| 21 | ToolSchemaTokens int |
| 22 | SessionContextDigest string |
| 23 | } |
| 24 | |
| 25 | // CacheDiagnostics is a type alias for event.CacheDiagnostics so the agent |
| 26 | // can construct and compare diagnostics without importing event itself in |
| 27 | // every call site, while still assigning to event.Event.CacheDiagnostics. |
| 28 | type CacheDiagnostics = event.CacheDiagnostics |
| 29 | |
| 30 | func shortHash(v any) string { |
| 31 | b, _ := json.Marshal(v) |
| 32 | h := sha256.Sum256(b) |
| 33 | return fmt.Sprintf("%x", h[:8]) |
| 34 | } |
| 35 | |
| 36 | // CaptureShape takes a snapshot of the current prefix state. |
| 37 | func CaptureShape(systemPrompt string, schemas []provider.ToolSchema, rewriteVersion int) PrefixShape { |
| 38 | normalizedSchemas := normalizeToolSchemas(schemas) |
| 39 | toolsJSON, _ := json.Marshal(normalizedSchemas) |
| 40 | return PrefixShape{ |
| 41 | SystemHash: shortHash(systemPrompt), |
| 42 | ToolsHash: shortHash(string(toolsJSON)), |
| 43 | PrefixHash: shortHash(map[string]any{ |
| 44 | "system": systemPrompt, |
| 45 | "tools": string(toolsJSON), |
| 46 | }), |
| 47 | LogRewriteVersion: rewriteVersion, |
| 48 | ToolSchemaTokens: estimateTokens(string(toolsJSON)), |
| 49 | } |
| 50 | } |
| 51 | |
| 52 | func normalizeToolSchemas(schemas []provider.ToolSchema) []provider.ToolSchema { |
| 53 | out := make([]provider.ToolSchema, len(schemas)) |
| 54 | copy(out, schemas) |
| 55 | sort.Slice(out, func(i, j int) bool { |
| 56 | if out[i].Name != out[j].Name { |
| 57 | return out[i].Name < out[j].Name |
| 58 | } |
| 59 | if out[i].Description != out[j].Description { |
| 60 | return out[i].Description < out[j].Description |
| 61 | } |
| 62 | return string(out[i].Parameters) < string(out[j].Parameters) |
| 63 | }) |
| 64 | return out |
| 65 | } |
| 66 | |
| 67 | // CompareShape returns diagnostics describing what changed between two shapes. |
| 68 | // contentReasons is the set of provider-visible rewrite reasons (e.g. |
| 69 | // "compact_auto", "snip", "rewind_truncate") drained from the Session since |
| 70 | // prev was captured — see Session.DrainContentRewriteReasons. It is the sole |
| 71 | // source of rewrite-caused reasons: a bare LogRewriteVersion change with no |
| 72 | // drained reason means only local-only metadata was touched (a decision |
| 73 | // receipt, tool-call preview/resolution, or an Edited-message replace), which |
| 74 | // never reaches the provider and so must not be reported as a cache change. |
| 75 | func CompareShape(prev, cur PrefixShape, usage *provider.Usage, contentReasons []string) CacheDiagnostics { |
| 76 | reasons := []string{} |
| 77 | if prev.SystemHash != "" && prev.SystemHash != cur.SystemHash { |
| 78 | reasons = append(reasons, "system") |
| 79 | } |
| 80 | if prev.ToolsHash != "" && prev.ToolsHash != cur.ToolsHash { |
| 81 | reasons = append(reasons, "tools") |
| 82 | } |
| 83 | if prev.SessionContextDigest != cur.SessionContextDigest { |
| 84 | reasons = append(reasons, "session_context") |
| 85 | } |
| 86 | reasons = append(reasons, contentReasons...) |
| 87 | var miss, hit int |
| 88 | if usage != nil { |
| 89 | miss = usage.CacheMissTokens |
| 90 | hit = usage.CacheHitTokens |
| 91 | } |
| 92 | return CacheDiagnostics{ |
| 93 | PrefixHash: cur.PrefixHash, |
| 94 | PrefixChanged: len(reasons) > 0, |
| 95 | PrefixChangeReasons: reasons, |
| 96 | SystemHash: cur.SystemHash, |
| 97 | ToolsHash: cur.ToolsHash, |
| 98 | LogRewriteVersion: cur.LogRewriteVersion, |
| 99 | ToolSchemaTokens: cur.ToolSchemaTokens, |
| 100 | CacheMissTokens: miss, |
| 101 | CacheHitTokens: hit, |
| 102 | } |
| 103 | } |
| 104 | |
| 105 | // estimateTokens gives a rough token count from byte length. |
| 106 | // A proper tokenizer would be more accurate, but for diagnostic |
| 107 | // purposes a byte-based estimate is sufficient and zero-alloc. |
| 108 | func estimateTokens(s string) int { |
| 109 | // ~4 chars per token is a workable heuristic for code-heavy JSON. |
| 110 | if len(s) == 0 { |
| 111 | return 0 |
| 112 | } |
| 113 | return len(s) / 4 |
| 114 | } |
| 115 | |
| 116 | // SchemaTokenCosts returns per-tool token cost estimates for display. |
| 117 | func SchemaTokenCosts(schemas []provider.ToolSchema) []ToolSchemaCost { |
| 118 | out := make([]ToolSchemaCost, 0, len(schemas)) |
| 119 | for _, s := range schemas { |
| 120 | b, _ := json.Marshal(s) |
| 121 | out = append(out, ToolSchemaCost{Name: s.Name, Tokens: estimateTokens(string(b))}) |
| 122 | } |
| 123 | return out |
| 124 | } |
| 125 | |
| 126 | // ToolSchemaCost is a per-tool token cost estimate for diagnostic display. |
| 127 | type ToolSchemaCost struct { |
| 128 | Name string |
| 129 | Tokens int |
| 130 | } |
| 131 |