metrale_bench/benchmarks/ssm_poison/
driver.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: The SSM state poisoning gate's driver: replay the probe script
4//! against one server, compare every replay with round 0, then run the
5//! tool-call path probe.
6//!
7//! Owner: bench, SSM poisoning gate.
8//! Invariants:
9//! - A replay round that errors is scored
10//!   [`super::compare::RoundVerdict::Unmeasured`]; the run goes on.
11//! - The run fails on a reference round that misses the script anchors, on a
12//!   short run, on any collapsed, unmeasured or zero-cache replay, and on a
13//!   tool call that differs between predecessor paths.
14
15use crate::hardware::Sensitivity;
16use std::future::Future;
17use std::time::{Duration, Instant};
18
19use anyhow::{Context, Result, bail};
20use serde_json::{Value, json};
21
22use crate::benchmark::{Benchmark, BenchmarkDescriptor};
23use crate::benchmarks::{one_line, transcript::Transcript};
24use crate::http;
25use crate::metadata::PluginMetadata;
26use crate::params::{ParamKind, ParamSpec, ParamValue, ParamValues};
27use crate::plugin::{Plugin, PluginHandle};
28use crate::result::{BenchmarkResult, LogLine, RunStatus, Verdict};
29
30use super::compare::{self, RoundVerdict};
31use super::probe;
32use super::score::RoundRecord;
33use super::toolcall::{self, Path, PathResult};
34
35const SUMMARY: &str = "Replayed conversations must come back byte-identical";
36pub const METADATA: PluginMetadata = PluginMetadata::metrale(SUMMARY);
37
38/// 2026-09-26: The default replay count. The gate's BENCH.toml entry bounds the
39/// recorded `rounds` metric at exactly 12.
40pub const DEFAULT_ROUNDS: usize = 12;
41
42pub const DESCRIPTOR: BenchmarkDescriptor = BenchmarkDescriptor {
43    id: "ssm-state-poisoning-gate",
44    name: "SSM State Poisoning Gate",
45    summary: SUMMARY,
46    detail: "Replays a fixed 4-turn conversation script 12 times against one prefix-cached \
47             server at temperature 0, comparing every turn against the first round. It splits \
48             two failure classes: restore JITTER (same finish reason, length within bounds) is \
49             a healthy engine property — Marconi restores the same token from alternating \
50             anchors between rounds, so accumulation differs and turns 2-4 come back reworded — \
51             and is recorded but passed; restore POISONING collapses the output (early-EOS \
52             stubs or runaway, the exact signature the batch4 stack shipped 2026-08-11) and \
53             FAILS the gate. Any collapsed or unmeasured round fails; jitter only records. \
54             Two self-honesty checks keep the gate from passing vacuously: the reference \
55             round must satisfy the script's semantic anchors (poisoning deterministic from \
56             round 0 would otherwise certify as Invariant), and every replay's first turn \
57             must attest a nonzero prefix-cache restore (with caching off, the gate would \
58             green-light a path it never exercised). Serves with --serve-override \
59             ssm_cache_slots=256 so the snapshot pool cannot evict mid-run (churn is noise).",
60    duration_hint: "~8–10 min",
61    expected_secs: 150,
62    updated: "2026-08-13",
63    needs_confirmation: false,
64    // 2026-09-26: Not pinned to a checkpoint: replay stability is asked of any
65    // served model.
66    intended_for: None,
67    threshold_params: &[],
68    // 2026-09-26: The verdict reads transcripts and cache counts, never timing.
69    sensitivity: Sensitivity::Correctness,
70    ctor: || Box::new(SsmPoison::default()),
71};
72
73#[derive(Default, Clone, Copy, PartialEq, Eq, Debug)]
74enum Phase {
75    #[default]
76    Baseline,
77    Replay,
78    /// 2026-09-26: Path-independence of tool calls: one target turn, reached by
79    /// each `toolcall::Path`.
80    ToolPath,
81    Score,
82    Done,
83}
84
85#[derive(Default)]
86pub struct SsmPoison {
87    handle: Option<PluginHandle>,
88    phase: Phase,
89    rounds: usize,
90    /// 2026-09-26: Divergences found by the tool-call path probe.
91    tool_divergences: Vec<toolcall::Divergence>,
92    /// 2026-09-26: Whether the direct path called a tool on the target. Logged;
93    /// not part of the verdict.
94    tool_reference_called: bool,
95    max_tokens: usize,
96    timeout: Duration,
97    started: Option<Instant>,
98    probed: bool,
99    /// 2026-09-26: Round 0's transcripts, one per turn: the reference.
100    reference: Vec<Transcript>,
101    /// 2026-09-26: One record per replay round: verdict plus the turn-1 cache
102    /// count the vacuity check reads.
103    replays: Vec<RoundRecord>,
104}
105
106impl SsmPoison {
107    fn handle(&self) -> Result<&PluginHandle> {
108        self.handle.as_ref().context("benchmark was not loaded")
109    }
110
111    fn elapsed(&self) -> Duration {
112        self.started.map(|s| s.elapsed()).unwrap_or_default()
113    }
114
115    /// 2026-09-26: probe + reference + N replays + tool paths + score.
116    fn total_steps(&self) -> u64 {
117        self.rounds as u64 + 4
118    }
119
120    fn steps_done(&self) -> u64 {
121        match self.phase {
122            Phase::Baseline => 1,
123            Phase::Replay => 2 + self.replays.len() as u64,
124            Phase::ToolPath => 2 + self.rounds as u64,
125            Phase::Score => 3 + self.rounds as u64,
126            Phase::Done => self.total_steps(),
127        }
128    }
129
130    fn frame(&self, phase: &str, line: Option<LogLine>) -> BenchmarkResult {
131        let mut f = BenchmarkResult::running(phase, self.elapsed())
132            .with_progress(self.steps_done(), self.total_steps());
133        if let Some(line) = line {
134            f = f.log_line(line);
135        }
136        f
137    }
138
139    /// 2026-09-26: One chat turn against the endpoint. The error is returned,
140    /// so the caller decides between a failed round and a failed run.
141    async fn turn(&self, messages: &[Value]) -> Result<Transcript> {
142        let handle = self.handle()?;
143        let target = handle.target();
144        let body = probe::request_body(&target.model, messages, self.max_tokens);
145        let outcome = http::chat_stream(target, &body, self.timeout)
146            .await
147            .context("chat request failed")?;
148        Ok(Transcript::from(&outcome))
149    }
150
151    /// 2026-09-26: Replay the whole script from scratch. Returns the per-turn
152    /// transcripts, or the first error that stopped the replay.
153    async fn replay_script(&self, label: &str) -> Result<Vec<Transcript>> {
154        let handle = self.handle()?.clone();
155        let mut messages: Vec<Value> = Vec::with_capacity(probe::TURNS.len() * 2);
156        let mut transcripts = Vec::with_capacity(probe::TURNS.len());
157        for (i, turn) in probe::TURNS.iter().enumerate() {
158            handle.check_cancelled()?;
159            let content = if i == 0 {
160                probe::first_turn()
161            } else {
162                turn.to_string()
163            };
164            messages.push(json!({"role": "user", "content": content}));
165            handle.status(format!(
166                "{label} · turn {}/{turns}",
167                i + 1,
168                turns = probe::TURNS.len()
169            ));
170            let t = self.turn(&messages).await?;
171            messages.push(json!({"role": "assistant", "content": t.text.clone()}));
172            transcripts.push(t);
173        }
174        Ok(transcripts)
175    }
176
177    /// 2026-09-26: Issue the target turn along each `Path`. The user turns of
178    /// the paths differ only in the interposed one.
179    async fn run_tool_paths(&self) -> Result<Vec<PathResult>> {
180        let handle = self.handle()?.clone();
181        let mut out = Vec::with_capacity(Path::ALL.len());
182        for path in Path::ALL {
183            handle.check_cancelled()?;
184            handle.status(format!("tool path · {}", path.label()));
185            let mut messages: Vec<Value> =
186                vec![json!({"role": "user", "content": probe::first_turn()})];
187            let ack = self.tool_turn(&messages).await?;
188            messages.push(json!({"role": "assistant", "content": ack.text}));
189            if let Some(interposed) = path.interposed() {
190                messages.push(json!({"role": "user", "content": interposed}));
191                let t = self.tool_turn(&messages).await?;
192                // 2026-09-26: The interposed turn's calls are not compared, and
193                // only its text is appended to the history.
194                messages.push(json!({"role": "assistant", "content": t.text}));
195            }
196            messages.push(json!({"role": "user", "content": toolcall::TARGET}));
197            let target_reply = self.tool_turn(&messages).await?;
198            out.push(PathResult {
199                path,
200                target: target_reply,
201            });
202        }
203        Ok(out)
204    }
205
206    /// 2026-09-26: One turn of the tool-path probe: same transport as the replay
207    /// probe, but the body offers tools and leaves the choice to the model.
208    async fn tool_turn(&self, messages: &[Value]) -> Result<Transcript> {
209        let handle = self.handle()?;
210        let target = handle.target();
211        let body = toolcall::request_body(&target.model, messages, self.max_tokens);
212        let outcome = http::chat_stream(target, &body, self.timeout)
213            .await
214            .context("tool-path chat request failed")?;
215        Ok(Transcript::from(&outcome))
216    }
217
218    /// 2026-09-26: The replay score and verdict (`score::score`,
219    /// `score::verdict`); `next` adds the tool-path finding.
220    pub(super) fn scored(&self) -> (super::score::Score, Verdict) {
221        let s = super::score::score(&self.replays);
222        let v = super::score::verdict(&s, self.rounds);
223        (s, v)
224    }
225}
226
227impl Plugin for SsmPoison {
228    fn metadata(&self) -> &'static PluginMetadata {
229        &METADATA
230    }
231
232    fn load(&mut self, handle: PluginHandle) -> impl Future<Output = Result<()>> + Send {
233        self.handle = Some(handle);
234        self.started = Some(Instant::now());
235        async { Ok(()) }
236    }
237}
238
239impl Benchmark for SsmPoison {
240    fn descriptor(&self) -> &'static BenchmarkDescriptor {
241        &DESCRIPTOR
242    }
243
244    fn parameters(&self) -> Vec<ParamSpec> {
245        vec![
246            ParamSpec::new(
247                "rounds",
248                "Replay rounds",
249                "How many times the script is replayed after the reference round. The incident \
250                 poisoned rounds 8-9 of 10; BENCH.toml pins the gate at 12.",
251                ParamKind::Int { min: 3, max: 30 },
252                ParamValue::Int(DEFAULT_ROUNDS as i64),
253            ),
254            ParamSpec::new(
255                "max_tokens",
256                "Max tokens per turn",
257                "Output budget per turn. Sized so the COLLAPSE_RATIO_CEIL (2.0) stays \
258                 reachable: the script's longest answer is a short paragraph (~a few hundred \
259                 tokens), and a runaway replay must be able to reach twice the reference \
260                 length before the budget clamps it. The old 256 clamped runaways below the \
261                 ceiling on any turn past 128 tokens. A reference turn that hits this budget \
262                 fails the reference anchors, so a budget still too small is loud, not silent.",
263                ParamKind::Int { min: 32, max: 4096 },
264                ParamValue::Int(1024),
265            ),
266            ParamSpec::new(
267                "request_timeout_s",
268                "Request timeout",
269                "Seconds before a single turn is abandoned; the round scores Unmeasured.",
270                ParamKind::Int { min: 10, max: 3600 },
271                ParamValue::Int(300),
272            ),
273        ]
274    }
275
276    fn configure(&mut self, values: &ParamValues) -> Result<()> {
277        let specs = self.parameters();
278        values.validate_against(&specs)?;
279        self.rounds = values.usize("rounds")?;
280        self.max_tokens = values.usize("max_tokens")?;
281        self.timeout = Duration::from_secs(values.usize("request_timeout_s")? as u64);
282        self.probed = false;
283        self.phase = Phase::Baseline;
284        self.reference.clear();
285        self.replays.clear();
286        Ok(())
287    }
288
289    async fn next(&mut self) -> Result<BenchmarkResult> {
290        let handle = self.handle()?.clone();
291        handle.check_cancelled()?;
292
293        if !self.probed {
294            self.probed = true;
295            http::probe(handle.target(), Duration::from_secs(10))
296                .await
297                .context("endpoint probe failed — check the target URL and port")?;
298            if self.rounds < 3 {
299                bail!("need at least 3 replay rounds");
300            }
301            return Ok(self.frame(
302                "probe",
303                Some(LogLine::info(format!(
304                    "{} · model {} · {} replay rounds, prefix caching under test",
305                    handle.target().base_url,
306                    handle.target().model,
307                    self.rounds
308                ))),
309            ));
310        }
311
312        match self.phase {
313            Phase::Baseline => {
314                self.reference = self.replay_script("reference").await?;
315                // 2026-09-26: The replays are held only to round 0, so round 0
316                // is held to the script: output that is wrong the same way in
317                // every round would otherwise score Invariant.
318                let violations = probe::validate_reference(&self.reference);
319                if !violations.is_empty() {
320                    self.phase = Phase::Done;
321                    let reason = format!(
322                        "reference round failed the script anchors ({}) — replays compared \
323                         against it would certify nothing",
324                        violations.join("; ")
325                    );
326                    let (s, _) = self.scored();
327                    let mut metrics = super::report::metrics(&s);
328                    metrics.insert("sum_wall_s".to_string(), self.elapsed().as_secs_f64());
329                    return Ok(BenchmarkResult {
330                        status: RunStatus::Completed,
331                        ..BenchmarkResult::running("reference", self.elapsed())
332                    }
333                    .with_metrics(metrics)
334                    .with_verdict(Verdict::fail(reason.clone()))
335                    .log_line(LogLine::error(one_line(reason))));
336                }
337                self.phase = Phase::Replay;
338                Ok(self.frame(
339                    "reference",
340                    Some(LogLine::info(format!(
341                        "reference round captured: {} turns, script anchors hold",
342                        self.reference.len()
343                    ))),
344                ))
345            }
346            Phase::Replay => {
347                let n = self.replays.len() + 1;
348                // 2026-09-26: Turn 1's cached-token count is recorded;
349                // `score::verdict` fails a replay whose turn 1 reports 0.
350                let (v, turn1_cached) = match self.replay_script(&format!("replay {n}")).await {
351                    Ok(replay) => {
352                        let cached = replay.first().map(|t| t.cached_prompt_tokens);
353                        (compare::compare_round(&self.reference, &replay), cached)
354                    }
355                    Err(e) => (
356                        RoundVerdict::Unmeasured {
357                            reason: one_line(format!("{e:#}")),
358                        },
359                        None,
360                    ),
361                };
362                let line = match &v {
363                    RoundVerdict::Collapsed { turns } => Some(LogLine::error(format!(
364                        "replay {n} COLLAPSED — restored state produced degenerate output on \
365                         turns {:?} (the SSM state poisoning signature)",
366                        turns.iter().map(|t| t.turn).collect::<Vec<_>>()
367                    ))),
368                    RoundVerdict::Jittered { turns } => Some(LogLine::info(format!(
369                        "replay {n} jittered (healthy) on turns {:?} — restore anchor \
370                         selection varies between rounds",
371                        turns.iter().map(|t| t.turn).collect::<Vec<_>>()
372                    ))),
373                    _ => None,
374                };
375                // 2026-09-26: A collapse or jitter line keeps priority; the
376                // vacuity line is logged only when there is no other line.
377                let line = if line.is_none() && turn1_cached == Some(0) {
378                    Some(LogLine::error(format!(
379                        "replay {n} restored 0 cached prompt tokens on turn 1 — the round \
380                         never exercised the prefix restore path"
381                    )))
382                } else {
383                    line
384                };
385                self.replays.push(RoundRecord {
386                    round: n,
387                    verdict: v,
388                    turn1_cached,
389                });
390                if self.replays.len() >= self.rounds {
391                    self.phase = Phase::ToolPath;
392                }
393                Ok(self.frame(&format!("replay {n}"), line))
394            }
395            Phase::ToolPath => {
396                self.phase = Phase::Score;
397                let results = self.run_tool_paths().await?;
398                self.tool_divergences = toolcall::divergences(&results);
399                self.tool_reference_called = toolcall::reference_called(&results);
400                let line = if let Some(d) = self.tool_divergences.first() {
401                    LogLine::error(one_line(format!(
402                        "TOOL CALL DEPENDS ON HISTORY — {} ({} of {} perturbed paths)",
403                        d.describe(),
404                        self.tool_divergences.len(),
405                        Path::ALL.len() - 1
406                    )))
407                } else if self.tool_reference_called {
408                    LogLine::error(
409                        "the direct path itself called a tool on the irrelevance target — the                          reference every other path is compared against is already wrong"
410                            .to_string(),
411                    )
412                } else {
413                    LogLine::info(format!(
414                        "tool calls identical across all {} predecessor paths",
415                        Path::ALL.len()
416                    ))
417                };
418                Ok(self.frame("tool paths", Some(line)))
419            }
420            Phase::Score => {
421                let (s, v) = self.scored();
422                // 2026-09-26: Replay and tool-path findings are independent;
423                // either alone fails the run.
424                let v = if let Some(d) = self.tool_divergences.first() {
425                    Verdict::fail(format!(
426                        concat!(
427                            "TOOL CALLS DEPEND ON HISTORY: {} — the same request ",
428                            "answered differently depending only on what preceded ",
429                            "it, so a known-answer score is not reproducible under ",
430                            "reordering"
431                        ),
432                        d.describe()
433                    ))
434                } else {
435                    v
436                };
437                self.phase = Phase::Done;
438                let line = LogLine::info(one_line(format!(
439                    "{} replays: {} invariant · {} jittered · {} collapsed · {} unmeasured",
440                    s.rounds, s.invariant, s.jittered, s.collapsed, s.unmeasured
441                )));
442                // 2026-09-26: `sum_wall_s` is wall-clock state, so it is added
443                // here rather than in the pure `report::metrics`.
444                let mut metrics = super::report::metrics(&s);
445                metrics.insert("sum_wall_s".to_string(), self.elapsed().as_secs_f64());
446                metrics.insert(
447                    "tool_path_divergences".to_string(),
448                    self.tool_divergences.len() as f64,
449                );
450                metrics.insert("tool_paths".to_string(), toolcall::Path::ALL.len() as f64);
451                Ok(BenchmarkResult {
452                    status: RunStatus::Completed,
453                    ..BenchmarkResult::running("done", self.elapsed())
454                }
455                .with_progress(self.total_steps(), self.total_steps())
456                .with_summary(super::report::summary(&s))
457                .with_table(super::report::table(&s))
458                .with_metrics(metrics)
459                .with_verdict(v)
460                .log_line(line))
461            }
462            Phase::Done => bail!("next() was called after the run finished"),
463        }
464    }
465}
466
467#[cfg(test)]
468#[path = "driver_tests.rs"]
469mod driver_tests;