| 1 | //! Prompt-bytes family: the model-visible prefix of a real first request. |
| 2 | //! |
| 3 | //! `docs/CACHE.md`: the system prompt and tool catalog are a session-pinned |
| 4 | //! KV-cache prefix, so a migration that moves any contributor (a tool adapter, |
| 5 | //! an MCP catalog, a slash command) must not change those bytes by accident. |
| 6 | //! Each case runs one real turn through the engine with the scripted provider |
| 7 | //! and captures the first `MessageRequest` — the exact prefix production would |
| 8 | //! send — then pins it three ways: |
| 9 | //! |
| 10 | //! - `<case>.golden.json`: sha256 of the system prompt, the tool catalog and |
| 11 | //! the combined prefix (bytes as serialized, key order included), sizes, and |
| 12 | //! the tool names in catalog order; |
| 13 | //! - `<case>.system.golden.txt`: the flat system prompt, readable in review; |
| 14 | //! - `<case>.tools.golden.json`: the tool catalog, readable in review. |
| 15 | //! |
| 16 | //! Masks are only the genuinely host-specific facts: temp paths, and the |
| 17 | //! environment block's `platform` / `shell` lines. Tools the catalog admits |
| 18 | //! only after probing the host for a binary are removed from the pinned |
| 19 | //! catalog and pinned one by one in `host_probed_tools.golden.json`, checked |
| 20 | //! whenever the probe succeeds on the running host. |
| 21 | |
| 22 | use std::collections::BTreeMap; |
| 23 | |
| 24 | use serde_json::{Value, json}; |
| 25 | |
| 26 | use super::events::run_scripted_turn; |
| 27 | use super::golden::{self, Failures, Sandbox}; |
| 28 | |
| 29 | const FAMILY: &str = "prompt"; |
| 30 | |
| 31 | /// Tools registered only when a host binary or OS service is found |
| 32 | /// (`tool_catalog.rs` python probe, `registry.rs` pandoc / OCR probes, |
| 33 | /// `js_execution`'s node probe). |
| 34 | const HOST_PROBED_TOOLS: &[&str] = &[ |
| 35 | "code_execution", |
| 36 | "image_ocr", |
| 37 | "js_execution", |
| 38 | "pandoc_convert", |
| 39 | ]; |
| 40 | |
| 41 | fn default_script() -> Value { |
| 42 | json!([{ "events": [ |
| 43 | { "type": "message_start", "message": { |
| 44 | "id": "msg_prompt", "type": "message", "role": "assistant", "content": [], |
| 45 | "model": "conformance", "stop_reason": null, "stop_sequence": null, |
| 46 | "usage": { "input_tokens": 0, "output_tokens": 0 } } }, |
| 47 | { "type": "content_block_start", "index": 0, "content_block": { "type": "text", "text": "" } }, |
| 48 | { "type": "content_block_delta", "index": 0, "delta": { "type": "text_delta", "text": "ok" } }, |
| 49 | { "type": "content_block_stop", "index": 0 }, |
| 50 | { "type": "message_delta", "delta": { "stop_reason": "end_turn" } }, |
| 51 | { "type": "message_stop" } |
| 52 | ] }]) |
| 53 | } |
| 54 | |
| 55 | struct Prefix { |
| 56 | system_json: String, |
| 57 | system_flat: String, |
| 58 | system_blocks: usize, |
| 59 | tools: Vec<Value>, |
| 60 | probed: BTreeMap<String, Value>, |
| 61 | } |
| 62 | |
| 63 | fn mask_host_lines(text: &str) -> String { |
| 64 | let platform = regex::Regex::new(r"(?m)^- platform: .*$").expect("platform regex"); |
| 65 | let shell = regex::Regex::new(r"(?m)^- shell: .*$").expect("shell regex"); |
| 66 | let text = platform.replace_all(text, "- platform: <PLATFORM>"); |
| 67 | shell.replace_all(&text, "- shell: <SHELL>").into_owned() |
| 68 | } |
| 69 | |
| 70 | /// Apply [`mask_host_lines`] inside every string of a JSON tree (the lines |
| 71 | /// live inside JSON strings, where serialization escapes their newlines). |
| 72 | fn mask_host_value(value: &mut Value) { |
| 73 | match value { |
| 74 | Value::String(text) => *text = mask_host_lines(text), |
| 75 | Value::Array(items) => items.iter_mut().for_each(mask_host_value), |
| 76 | Value::Object(map) => map.values_mut().for_each(mask_host_value), |
| 77 | _ => {} |
| 78 | } |
| 79 | } |
| 80 | |
| 81 | fn capture_prefix(name: &str, case: &Value) -> Result<Prefix, String> { |
| 82 | let sandbox = Sandbox::new(case); |
| 83 | let script = case |
| 84 | .get("provider_script") |
| 85 | .cloned() |
| 86 | .unwrap_or_else(default_script); |
| 87 | let (record, provider) = run_scripted_turn(&sandbox, case, &script); |
| 88 | if let Some(error) = record.failure { |
| 89 | return Err(error); |
| 90 | } |
| 91 | let requests = provider.captured(); |
| 92 | let first = requests |
| 93 | .first() |
| 94 | .ok_or_else(|| format!("{name}: the engine sent no model request"))?; |
| 95 | let mut masker = sandbox.masker(&[]); |
| 96 | |
| 97 | let mut system = serde_json::to_value(&first.system).expect("system serializes"); |
| 98 | masker.value(&mut system); |
| 99 | mask_host_value(&mut system); |
| 100 | let system_json = serde_json::to_string(&system).expect("system json"); |
| 101 | let system_flat = first |
| 102 | .system |
| 103 | .as_ref() |
| 104 | .map(crate::prompts::system_prompt_flat_text) |
| 105 | .unwrap_or_default(); |
| 106 | let system_flat = mask_host_lines(&masker.text(&system_flat)); |
| 107 | let system_blocks = match &first.system { |
| 108 | Some(codewhale_models::SystemPrompt::Blocks(blocks)) => blocks.len(), |
| 109 | Some(codewhale_models::SystemPrompt::Text(_)) => 1, |
| 110 | None => 0, |
| 111 | }; |
| 112 | if system_blocks == 0 { |
| 113 | return Err(format!("{name}: the captured system prefix is empty")); |
| 114 | } |
| 115 | |
| 116 | let mut tools = Vec::new(); |
| 117 | let mut probed = BTreeMap::new(); |
| 118 | for tool in first.tools.iter().flatten() { |
| 119 | let mut value = serde_json::to_value(tool).expect("tool serializes"); |
| 120 | masker.value(&mut value); |
| 121 | mask_host_value(&mut value); |
| 122 | if HOST_PROBED_TOOLS.contains(&tool.name.as_str()) { |
| 123 | probed.insert(tool.name.clone(), value); |
| 124 | } else { |
| 125 | tools.push(value); |
| 126 | } |
| 127 | } |
| 128 | if tools.is_empty() && probed.is_empty() { |
| 129 | return Err(format!("{name}: the captured tool catalog is empty")); |
| 130 | } |
| 131 | Ok(Prefix { |
| 132 | system_json, |
| 133 | system_flat, |
| 134 | system_blocks, |
| 135 | tools, |
| 136 | probed, |
| 137 | }) |
| 138 | } |
| 139 | |
| 140 | fn summary(prefix: &Prefix) -> Value { |
| 141 | let tools_json = serde_json::to_string(&prefix.tools).expect("tools json"); |
| 142 | let combined = format!("{}\n{}", prefix.system_json, tools_json); |
| 143 | json!({ |
| 144 | "system_sha256": crate::hashing::sha256_hex(prefix.system_json.as_bytes()), |
| 145 | "system_bytes": prefix.system_json.len(), |
| 146 | "system_blocks": prefix.system_blocks, |
| 147 | "tools_sha256": crate::hashing::sha256_hex(tools_json.as_bytes()), |
| 148 | "tools_bytes": tools_json.len(), |
| 149 | "tool_count": prefix.tools.len(), |
| 150 | "tool_names": prefix.tools.iter().map(|tool| tool["name"].clone()).collect::<Vec<_>>(), |
| 151 | "prefix_sha256": crate::hashing::sha256_hex(combined.as_bytes()), |
| 152 | "masks": ["<WORKSPACE>", "<HOME>", "<TMP>", "- platform: <PLATFORM>", "- shell: <SHELL>"], |
| 153 | "host_probed_tools_excluded": HOST_PROBED_TOOLS, |
| 154 | }) |
| 155 | } |
| 156 | |
| 157 | #[test] |
| 158 | fn model_visible_prefix_bytes_match_goldens() { |
| 159 | let dir = golden::family_dir(FAMILY); |
| 160 | let names = golden::case_names(FAMILY); |
| 161 | let mut failures = Failures::default(); |
| 162 | let mut probed_seen: BTreeMap<String, Value> = BTreeMap::new(); |
| 163 | for name in &names { |
| 164 | let case = golden::read_case(FAMILY, name); |
| 165 | let prefix = match capture_prefix(name, &case) { |
| 166 | Ok(prefix) => prefix, |
| 167 | Err(message) => { |
| 168 | failures.push(name, message); |
| 169 | continue; |
| 170 | } |
| 171 | }; |
| 172 | // A prefix that differs between two identical sessions is a cache |
| 173 | // bug in its own right (map iteration order, a clock, a counter). |
| 174 | match capture_prefix(name, &case) { |
| 175 | Ok(again) if again.system_json == prefix.system_json && again.tools == prefix.tools => { |
| 176 | } |
| 177 | Ok(_) => { |
| 178 | failures.push( |
| 179 | name, |
| 180 | "the prefix differs between two identical sessions in one process", |
| 181 | ); |
| 182 | continue; |
| 183 | } |
| 184 | Err(message) => { |
| 185 | failures.push(name, message); |
| 186 | continue; |
| 187 | } |
| 188 | } |
| 189 | for (tool, definition) in &prefix.probed { |
| 190 | if let Some(previous) = probed_seen.get(tool) |
| 191 | && previous != definition |
| 192 | { |
| 193 | failures.push( |
| 194 | name, |
| 195 | format!("host-probed tool `{tool}` differs across cases"), |
| 196 | ); |
| 197 | } |
| 198 | probed_seen.insert(tool.clone(), definition.clone()); |
| 199 | } |
| 200 | failures.record( |
| 201 | name, |
| 202 | golden::check_golden( |
| 203 | &dir.join(format!("{name}.golden.json")), |
| 204 | &golden::pretty(&summary(&prefix)), |
| 205 | ), |
| 206 | ); |
| 207 | failures.record( |
| 208 | name, |
| 209 | golden::check_golden( |
| 210 | &dir.join(format!("{name}.system.golden.txt")), |
| 211 | &format!("{}\n", prefix.system_flat), |
| 212 | ), |
| 213 | ); |
| 214 | failures.record( |
| 215 | name, |
| 216 | golden::check_golden( |
| 217 | &dir.join(format!("{name}.tools.golden.json")), |
| 218 | &golden::pretty(&Value::Array(prefix.tools.clone())), |
| 219 | ), |
| 220 | ); |
| 221 | } |
| 222 | |
| 223 | // Host-probed tools: every definition this host produced must match the |
| 224 | // pinned one. A tool this host cannot offer is skipped, not failed. |
| 225 | let probed_path = dir.join("host_probed_tools.golden.json"); |
| 226 | let mut pinned: BTreeMap<String, Value> = std::fs::read_to_string(&probed_path) |
| 227 | .ok() |
| 228 | .and_then(|text| serde_json::from_str(&text).ok()) |
| 229 | .unwrap_or_default(); |
| 230 | let mut drift = Vec::new(); |
| 231 | for (tool, definition) in &probed_seen { |
| 232 | if pinned.get(tool) != Some(definition) { |
| 233 | drift.push(tool.clone()); |
| 234 | pinned.insert(tool.clone(), definition.clone()); |
| 235 | } |
| 236 | } |
| 237 | if !drift.is_empty() { |
| 238 | let merged = Value::Object(pinned.into_iter().collect()); |
| 239 | failures.record( |
| 240 | "host_probed_tools", |
| 241 | golden::check_golden(&probed_path, &golden::pretty(&merged)).map_err(|message| { |
| 242 | format!("host-probed tool definition(s) {drift:?} drifted: {message}") |
| 243 | }), |
| 244 | ); |
| 245 | } |
| 246 | let skipped: Vec<&str> = HOST_PROBED_TOOLS |
| 247 | .iter() |
| 248 | .copied() |
| 249 | .filter(|tool| !probed_seen.contains_key(*tool)) |
| 250 | .collect(); |
| 251 | if !skipped.is_empty() { |
| 252 | eprintln!( |
| 253 | "conformance: host cannot offer {skipped:?}; their pinned definitions were not checked" |
| 254 | ); |
| 255 | } |
| 256 | failures.finish(FAMILY, names.len()); |
| 257 | } |
| 258 |