Skip to content

Commit 1be9e3c

Browse files
RoyLinRoyLin
authored andcommitted
feat(sae): wire SaeJudge into the pipeline + ACL config
- Pipeline::with_sae + routing: LlmActivations events are judged by the SAE tier (not the rule chain); an SAE escalation defers to the deep L3 agent. - config.rs: a sae { dict, escalate_at, block_at } ACL block loads the labeled feature dictionary and wires the tier; the napi/pyo3 SDKs get it via evaluate(). - SaeJudge::from_path. 49 lib tests pass (pipeline routing + config integration).
1 parent 018191e commit 1be9e3c

3 files changed

Lines changed: 129 additions & 1 deletion

File tree

src/config.rs

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
//!
1111
//! llm { url = "http://llm:18051/v1" model = "glm" key = "..." timeout_s = 30 } # L2 (optional)
1212
//! agent { bin = "a3s-code" skills = "./skills" timeout_s = 120 } # L3 (optional)
13+
//! sae { dict = "features.json" escalate_at = 0.3 block_at = 0.6 } # mech-interp (optional)
1314
//! deny { egress = "egress.txt" file = "file.txt" exec = "exec.txt" } # sinks (optional)
1415
//!
1516
//! rules = [
@@ -21,6 +22,7 @@
2122
use crate::enforce::Enforcer;
2223
use crate::pipeline::{Judge, Pipeline};
2324
use crate::rules::{default_rules, LiveRules, RuleEngine, RuleSpec};
25+
use crate::sae::SaeJudge;
2426
use crate::verdict::Severity;
2527
use crate::{AgentJudge, LlmJudge};
2628
use serde::Deserialize;
@@ -41,6 +43,9 @@ pub struct SdkConfig {
4143
pub llm: Option<LlmCfg>,
4244
/// L3 — a deep a3s-code agent investigation. Omit if you don't run L3.
4345
pub agent: Option<AgentCfg>,
46+
/// SAE mechanistic-interpretability tier — judges model-output `LlmActivations` events (from
47+
/// a3s-power's in-enclave tap) against a labeled feature dictionary. Omit if you don't run it.
48+
pub sae: Option<SaeCfg>,
4449
/// Deny-file sinks the kernel guards read. Omit to judge without enforcing.
4550
pub deny: Option<DenyCfg>,
4651
/// L1 site rules, evaluated before the built-in defaults (first match wins).
@@ -63,6 +68,16 @@ pub struct AgentCfg {
6368
pub timeout_s: Option<u64>,
6469
}
6570

71+
#[derive(Debug, Deserialize)]
72+
pub struct SaeCfg {
73+
/// Path to the labeled feature dictionary JSON: `{ "8801": {concept, category, weight, severity?}, ... }`.
74+
pub dict: String,
75+
/// harmful ≥ this → escalate (default 0.30).
76+
pub escalate_at: Option<f32>,
77+
/// harmful ≥ this → block (default 0.60).
78+
pub block_at: Option<f32>,
79+
}
80+
6681
#[derive(Debug, Default, Deserialize)]
6782
pub struct DenyCfg {
6883
pub egress: Option<String>,
@@ -110,6 +125,13 @@ impl SdkConfig {
110125
self.fail_closed,
111126
)));
112127
}
128+
if let Some(s) = self.sae {
129+
let mut judge = SaeJudge::from_path(Path::new(&s.dict))?;
130+
if let (Some(e), Some(b)) = (s.escalate_at, s.block_at) {
131+
judge = judge.thresholds(e, b);
132+
}
133+
pipeline = pipeline.with_sae(Arc::new(judge));
134+
}
113135

114136
let deny = self.deny.unwrap_or_default();
115137
let enforcer = Enforcer::new(
@@ -132,3 +154,38 @@ fn parse_severity(s: &str) -> Severity {
132154
_ => Severity::High,
133155
}
134156
}
157+
158+
#[cfg(test)]
159+
mod tests {
160+
use super::*;
161+
use crate::verdict::{Tier, Verdict};
162+
163+
#[test]
164+
fn sae_block_wires_the_tier_and_judges_activations() {
165+
// A labeled feature dictionary on disk.
166+
let dir = std::env::temp_dir();
167+
let dict_path = dir.join("sentry_sae_cfg_test_dict.json");
168+
std::fs::write(
169+
&dict_path,
170+
r#"{"8801":{"concept":"exploit-code-synthesis","category":"cyber_offense","weight":0.9,"severity":"high"}}"#,
171+
)
172+
.unwrap();
173+
174+
let acl = format!(
175+
"fail_closed = false\nsae {{ dict = \"{}\" }}\n",
176+
dict_path.display()
177+
);
178+
let (pipeline, _enforcer) = SdkConfig::from_acl(&acl).unwrap().build().unwrap();
179+
180+
// A model-output activations line with the harmful feature firing → blocked + explained.
181+
let line = r#"{"identity":{"agent":"deep-finance"},"event":{"LlmActivations":{"pid":1,"layer":18,"features":[[8801,0.95]]}}}"#;
182+
let ev = crate::event::ObservedEvent::parse(line).expect("parses");
183+
let d = pipeline.evaluate(&ev);
184+
assert_eq!(d.verdict, Verdict::Block);
185+
assert_eq!(d.tier, Tier::Sae);
186+
let ex = d.explain.expect("explainability present");
187+
assert_eq!(ex.drivers[0].concept, "exploit-code-synthesis");
188+
189+
let _ = std::fs::remove_file(&dict_path);
190+
}
191+
}

src/pipeline.rs

Lines changed: 64 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
//! so ready sooner — is authoritative. High-risk events thus get the deep look without paying the
1212
//! serial L2+L3 latency.
1313
14-
use crate::event::ObservedEvent;
14+
use crate::event::{Event, ObservedEvent};
1515
use crate::verdict::{Decision, Severity, Tier, Verdict};
1616
use std::sync::atomic::{AtomicU64, Ordering};
1717
use std::sync::Arc;
@@ -36,6 +36,9 @@ pub struct Pipeline {
3636
l1: Arc<dyn Judge>,
3737
l2: Option<Arc<dyn Judge>>,
3838
l3: Option<Arc<dyn Judge>>,
39+
/// The SAE mechanistic-interpretability tier — judges model-output `LlmActivations` events from
40+
/// a3s-power's in-enclave feature emission (a separate domain from the observer-event rule chain).
41+
sae: Option<Arc<dyn Judge>>,
3942
fail_closed: bool,
4043
speculate_above: Option<Severity>,
4144
/// Live count of speculative L3 threads (incl. ones detached after an L2 short-circuit), so we
@@ -50,6 +53,7 @@ impl Pipeline {
5053
l1,
5154
l2: None,
5255
l3: None,
56+
sae: None,
5357
fail_closed: false,
5458
speculate_above: None,
5559
l3_inflight: Arc::new(AtomicU64::new(0)),
@@ -73,6 +77,13 @@ impl Pipeline {
7377
self
7478
}
7579

80+
/// The SAE mechanistic-interpretability tier. Model-output `LlmActivations` events route here
81+
/// instead of the L1 rule chain; an SAE escalation can still defer to the deep L3 agent.
82+
pub fn with_sae(mut self, sae: Arc<dyn Judge>) -> Self {
83+
self.sae = Some(sae);
84+
self
85+
}
86+
7687
/// Treat an unresolved escalation (no deeper tier available) as `Block` instead of `Allow`.
7788
pub fn fail_closed(mut self, yes: bool) -> Self {
7889
self.fail_closed = yes;
@@ -86,8 +97,26 @@ impl Pipeline {
8697
self
8798
}
8899

100+
/// Route a model-output `LlmActivations` event to the SAE tier (with L3-on-escalate). Returns
101+
/// `None` for non-activation events, or when no SAE tier is configured (they take the rule chain).
102+
fn judge_model_output(&self, ev: &ObservedEvent) -> Option<Decision> {
103+
if !matches!(ev.event, Event::LlmActivations { .. }) {
104+
return None;
105+
}
106+
let sae = self.sae.as_ref()?;
107+
let d = sae.judge(ev);
108+
Some(match (d.verdict, &self.l3) {
109+
(Verdict::Escalate, Some(l3)) => l3.judge(ev),
110+
(Verdict::Escalate, None) => self.resolve_unescalated(d),
111+
_ => d,
112+
})
113+
}
114+
89115
/// Run the event through the tiers and return the deciding [`Decision`].
90116
pub fn evaluate(&self, ev: &ObservedEvent) -> Decision {
117+
if let Some(d) = self.judge_model_output(ev) {
118+
return d;
119+
}
91120
let d1 = self.l1.judge(ev);
92121
if d1.verdict != Verdict::Escalate {
93122
return d1;
@@ -106,6 +135,12 @@ impl Pipeline {
106135
/// blocks the event stream. The worker then calls [`evaluate`](Pipeline::evaluate) (L1 re-runs in
107136
/// µs, then L2/L3) on the escalated event.
108137
pub fn classify_l1(&self, ev: &ObservedEvent) -> Decision {
138+
// Model-output activations are judged by the SAE tier, not the rule engine.
139+
if let Some(sae) = &self.sae {
140+
if matches!(ev.event, Event::LlmActivations { .. }) {
141+
return sae.judge(ev);
142+
}
143+
}
109144
self.l1.judge(ev)
110145
}
111146

@@ -237,6 +272,34 @@ mod tests {
237272
assert_eq!(p.evaluate(&ev()).verdict, Verdict::Block);
238273
}
239274

275+
#[test]
276+
fn model_activations_route_to_sae_tier() {
277+
use crate::event::{Event, Identity};
278+
let act = ObservedEvent {
279+
identity: Identity::default(),
280+
provider: None,
281+
event: Event::LlmActivations {
282+
pid: 1,
283+
layer: 18,
284+
features: vec![],
285+
},
286+
raw: String::new(),
287+
};
288+
// Activations are judged by the SAE tier, not the rule chain.
289+
let p = Pipeline::new(Arc::new(Fixed(Tier::Rules, Verdict::Allow)))
290+
.with_sae(Arc::new(Fixed(Tier::Sae, Verdict::Block)));
291+
let d = p.evaluate(&act);
292+
assert_eq!(d.verdict, Verdict::Block);
293+
assert_eq!(d.tier, Tier::Sae);
294+
// A non-activation event still takes the L1 rule chain.
295+
assert_eq!(p.evaluate(&ev()).verdict, Verdict::Allow);
296+
// An SAE escalation defers to the deep L3 agent when present.
297+
let p2 = Pipeline::new(Arc::new(Fixed(Tier::Rules, Verdict::Allow)))
298+
.with_sae(Arc::new(Fixed(Tier::Sae, Verdict::Escalate)))
299+
.with_l3(Arc::new(Fixed(Tier::Agent, Verdict::Block)));
300+
assert_eq!(p2.evaluate(&act).tier, Tier::Agent);
301+
}
302+
240303
#[test]
241304
fn escalates_l1_to_l2_to_l3() {
242305
let p = Pipeline::new(Arc::new(Fixed(Tier::Rules, Verdict::Escalate)))

src/sae.rs

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ use crate::pipeline::Judge;
1818
use crate::verdict::{Decision, Driver, SaeScore, Severity, Tier, Verdict};
1919
use serde::Deserialize;
2020
use std::collections::{BTreeMap, HashMap};
21+
use std::path::Path;
2122

2223
/// One SAE feature's safety meaning — turns an anonymous feature id into a named, weighted concept.
2324
/// Produced offline by probing the SAE on a labeled safety set + causal validation (ablate the
@@ -71,6 +72,13 @@ impl SaeJudge {
7172
Ok(Self::new(dict))
7273
}
7374

75+
/// Load the labeled feature dictionary from a JSON file on disk.
76+
pub fn from_path(path: &Path) -> anyhow::Result<Self> {
77+
let json = std::fs::read_to_string(path)
78+
.map_err(|e| anyhow::anyhow!("reading SAE feature dict {}: {e}", path.display()))?;
79+
Self::from_json(&json)
80+
}
81+
7482
pub fn thresholds(mut self, escalate_at: f32, block_at: f32) -> Self {
7583
self.escalate_at = escalate_at;
7684
self.block_at = block_at;

0 commit comments

Comments
 (0)