use parking_lot::RwLock;
use std::collections::HashMap;
use std::future::Future;
use std::hash::{Hash, Hasher};
use std::path::Path;
use std::pin::Pin;
use std::process::Stdio;
use std::sync::Arc;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum InterruptMode {
Never,
#[default]
ProseOnly,
ToolOnly,
Always,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ScopeToken {
Text,
Thinking,
Tool {
name: String,
globs: Vec<String>,
},
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RuleSource {
BuiltinDefaults,
Project,
User,
}
#[derive(Debug, Clone)]
pub struct Rule {
pub name: String,
pub content: String,
pub description: Option<String>,
pub condition: Vec<regex::Regex>,
pub scope: Vec<ScopeToken>,
pub interrupt_mode: InterruptMode,
pub globs: Vec<String>,
pub always_apply: bool,
pub source: RuleSource,
pub ast_condition: Option<String>,
}
pub trait RuleRegistry: Send + Sync + 'static {
fn rules<'a>(&'a self) -> Pin<Box<dyn Future<Output = Vec<Rule>> + Send + 'a>>;
fn mark_injected(&self, _name: &str, _turn: u64) {}
fn injected_records(&self) -> Vec<(String, u64)> {
vec![]
}
fn restore(&self, _records: Vec<(String, u64)>) {}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum MatchSource {
Text,
Thinking,
Tool,
}
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
struct BufferKey {
source: MatchSource,
tool_name: Option<String>,
}
#[derive(Debug, Clone)]
pub struct TtsrMatchContext {
pub source: MatchSource,
pub file_paths: Vec<String>,
pub tool_name: Option<String>,
pub file_contents: Vec<(String, String)>,
}
#[derive(Debug, Clone)]
pub struct AstRule {
pub name: String,
pub pattern: String,
pub file_scope: Vec<String>,
pub interrupt_mode: InterruptMode,
}
pub type AstMatcherFn = dyn Fn(&str, &str) -> bool + Send + Sync;
fn default_sg_matcher() -> Box<AstMatcherFn> {
Box::new(|pattern: &str, content: &str| {
let mut tmp = std::env::temp_dir();
let unique = format!(
"ttsr-ast-{}-{}.snap",
std::process::id(),
content_digest(content)
);
tmp.push(unique);
if std::fs::write(&tmp, content).is_err() {
return false;
}
let matched = run_sg_match(pattern, &tmp).unwrap_or(false);
let _ = std::fs::remove_file(&tmp);
matched
})
}
fn run_sg_match(pattern: &str, target: &Path) -> std::io::Result<bool> {
let output = std::process::Command::new("sg")
.arg("run")
.arg("-p")
.arg(pattern)
.arg("--json")
.arg(target)
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.output();
let output = match output {
Ok(o) => o,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(false),
Err(_) => return Ok(false),
};
Ok(!output.stdout.is_empty())
}
pub struct TtsrAstMatcher {
rules: Vec<AstRule>,
seen_digests: HashMap<String, u64>,
matcher: Box<AstMatcherFn>,
}
impl std::fmt::Debug for TtsrAstMatcher {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("TtsrAstMatcher")
.field("rules_count", &self.rules.len())
.field("seen_digests_count", &self.seen_digests.len())
.finish_non_exhaustive()
}
}
impl TtsrAstMatcher {
pub fn new(rules: Vec<AstRule>) -> Self {
Self {
rules,
seen_digests: HashMap::new(),
matcher: default_sg_matcher(),
}
}
pub fn with_matcher(rules: Vec<AstRule>, matcher: Box<AstMatcherFn>) -> Self {
Self {
rules,
seen_digests: HashMap::new(),
matcher,
}
}
pub fn rule_count(&self) -> usize {
self.rules.len()
}
pub fn clear_dedup(&mut self) {
self.seen_digests.clear();
}
pub fn check_tool_snapshot(&mut self, file_path: &str, content: &str) -> Option<String> {
if self.rules.is_empty() {
return None;
}
let candidates: Vec<&AstRule> = self
.rules
.iter()
.filter(|r| file_scope_matches(&r.file_scope, file_path))
.collect();
if candidates.is_empty() {
return None;
}
let digest = content_digest(content);
if self.seen_digests.get(file_path) == Some(&digest) {
return None;
}
for rule in candidates {
if (self.matcher)(&rule.pattern, content) {
self.seen_digests.insert(file_path.to_string(), digest);
return Some(rule.name.clone());
}
}
self.seen_digests.insert(file_path.to_string(), digest);
None
}
}
fn content_digest(content: &str) -> u64 {
let mut hasher = std::collections::hash_map::DefaultHasher::new();
content.hash(&mut hasher);
hasher.finish()
}
fn file_scope_matches(scope: &[String], file_path: &str) -> bool {
if scope.is_empty() {
return true;
}
scope.iter().any(|g| {
glob::Pattern::new(g)
.map(|p| p.matches(file_path))
.unwrap_or(false)
})
}
pub struct TtsrEngine {
rules: Arc<dyn RuleRegistry>,
buffers: RwLock<HashMap<BufferKey, Vec<String>>>,
settings: TtsrSettings,
ast_matcher: RwLock<Option<TtsrAstMatcher>>,
}
impl std::fmt::Debug for TtsrEngine {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("TtsrEngine")
.field("settings", &self.settings)
.finish_non_exhaustive()
}
}
#[derive(Debug, Clone)]
pub struct TtsrSettings {
pub enabled: bool,
pub interrupt_mode: InterruptMode,
pub builtin_rules: bool,
pub max_retries_per_turn: u32,
}
impl Default for TtsrSettings {
fn default() -> Self {
Self {
enabled: false,
interrupt_mode: InterruptMode::ProseOnly,
builtin_rules: true,
max_retries_per_turn: 3,
}
}
}
impl TtsrEngine {
pub fn new(rules: Arc<dyn RuleRegistry>, settings: TtsrSettings) -> Self {
Self {
rules,
buffers: RwLock::new(HashMap::new()),
settings,
ast_matcher: RwLock::new(None),
}
}
pub fn with_ast_matcher(
rules: Arc<dyn RuleRegistry>,
settings: TtsrSettings,
ast_matcher: TtsrAstMatcher,
) -> Self {
Self {
rules,
buffers: RwLock::new(HashMap::new()),
settings,
ast_matcher: RwLock::new(Some(ast_matcher)),
}
}
pub fn set_ast_matcher(&self, matcher: TtsrAstMatcher) {
*self.ast_matcher.write() = Some(matcher);
}
pub fn clear_ast_matcher(&self) {
*self.ast_matcher.write() = None;
}
pub fn reset_buffers(&self) {
self.buffers.write().clear();
}
pub fn check_delta(&self, delta: &str, ctx: &TtsrMatchContext) -> Vec<Rule> {
if !self.settings.enabled {
return vec![];
}
let key = self.buffer_key(ctx);
let mut buffers = self.buffers.write();
let buf = buffers.entry(key).or_default();
buf.push(delta.to_string());
let full: String = buf.concat();
let mut matched = self.match_buffer(&full, ctx);
if matches!(ctx.source, MatchSource::Tool) && !ctx.file_contents.is_empty() {
let ast_matches = self.check_ast_against_contents(ctx);
for ast_match in ast_matches {
if !matched.iter().any(|r| r.name == ast_match.name) {
matched.push(ast_match);
}
}
}
matched
}
pub fn check_snapshot(&self, snapshot: &str, ctx: &TtsrMatchContext) -> Vec<Rule> {
if !self.settings.enabled {
return vec![];
}
let key = self.buffer_key(ctx);
let mut buffers = self.buffers.write();
buffers.insert(key, vec![snapshot.to_string()]);
let mut matched = self.match_buffer(snapshot, ctx);
if !ctx.file_contents.is_empty() {
let ast_matches = self.check_ast_against_contents(ctx);
for ast_match in ast_matches {
if !matched.iter().any(|r| r.name == ast_match.name) {
matched.push(ast_match);
}
}
}
matched
}
pub fn injected_records(&self) -> Vec<(String, u64)> {
self.rules.injected_records()
}
fn buffer_key(&self, ctx: &TtsrMatchContext) -> BufferKey {
BufferKey {
source: ctx.source,
tool_name: if matches!(ctx.source, MatchSource::Tool) {
ctx.tool_name.clone()
} else {
None
},
}
}
fn check_ast_against_contents(&self, ctx: &TtsrMatchContext) -> Vec<Rule> {
let mut guard = self.ast_matcher.write();
let matcher = match guard.as_mut() {
Some(m) => m,
None => return Vec::new(),
};
let mut matched = Vec::new();
for (path, content) in &ctx.file_contents {
if let Some(rule_name) = matcher.check_tool_snapshot(path, content)
&& let Some(rule) = self.lookup_rule(&rule_name)
{
matched.push(rule);
}
}
matched
}
fn lookup_rule(&self, name: &str) -> Option<Rule> {
let rules: Vec<Rule> = futures::executor::block_on(self.rules.rules());
rules.into_iter().find(|r| r.name == name)
}
fn match_buffer(&self, buf: &str, ctx: &TtsrMatchContext) -> Vec<Rule> {
let mut matched = Vec::new();
let rules: Vec<Rule> = futures::executor::block_on(self.rules.rules());
for rule in rules {
if !self.scope_matches(&rule, ctx) {
continue;
}
let mode = if matches!(rule.interrupt_mode, InterruptMode::Never) {
self.settings.interrupt_mode
} else {
rule.interrupt_mode
};
if !self.mode_allows(mode, ctx.source) {
continue;
}
if !rule.condition.iter().any(|re| re.is_match(buf)) {
continue;
}
matched.push(rule);
}
matched
}
fn scope_matches(&self, rule: &Rule, ctx: &TtsrMatchContext) -> bool {
if rule.scope.is_empty() {
return true;
}
for token in &rule.scope {
match token {
ScopeToken::Text => {
if matches!(ctx.source, MatchSource::Text) {
return true;
}
}
ScopeToken::Thinking => {
if matches!(ctx.source, MatchSource::Thinking) {
return true;
}
}
ScopeToken::Tool { name, globs } => {
if !matches!(ctx.source, MatchSource::Tool) {
continue;
}
if matches!(ctx.tool_name.as_ref(), Some(tool_name) if tool_name != name) {
continue;
}
if !globs.is_empty() {
let any_match = ctx.file_paths.iter().any(|fp| {
globs.iter().any(|g| {
g.strip_suffix("/*")
.map(|prefix| fp.starts_with(prefix))
.unwrap_or_else(|| g == fp)
})
});
if !any_match {
continue;
}
}
return true;
}
}
}
false
}
fn mode_allows(&self, mode: InterruptMode, source: MatchSource) -> bool {
match mode {
InterruptMode::Never => false,
InterruptMode::ProseOnly => matches!(source, MatchSource::Text),
InterruptMode::ToolOnly => matches!(source, MatchSource::Tool),
InterruptMode::Always => true,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use regex::Regex;
use std::pin::Pin;
struct StaticRegistry {
rules: Vec<Rule>,
injections: RwLock<Vec<(String, u64)>>,
}
impl RuleRegistry for StaticRegistry {
fn rules<'a>(&'a self) -> Pin<Box<dyn Future<Output = Vec<Rule>> + Send + 'a>> {
Box::pin(std::future::ready(self.rules.clone()))
}
fn mark_injected(&self, name: &str, turn: u64) {
self.injections.write().push((name.to_string(), turn));
}
fn injected_records(&self) -> Vec<(String, u64)> {
self.injections.read().clone()
}
fn restore(&self, records: Vec<(String, u64)>) {
*self.injections.write() = records;
}
}
fn make_rule(name: &str, pattern: &str) -> Rule {
Rule {
name: name.to_string(),
content: format!("Do not use {pattern}."),
description: Some(format!("Forbids {pattern}")),
condition: vec![Regex::new(pattern).unwrap()],
scope: vec![],
interrupt_mode: InterruptMode::ProseOnly,
globs: vec![],
always_apply: false,
source: RuleSource::BuiltinDefaults,
ast_condition: None,
}
}
fn substring_matcher(pattern: &str, content: &str) -> bool {
content.contains(pattern)
}
fn make_ast_rule(name: &str, pattern: &str, scope: Vec<String>) -> AstRule {
AstRule {
name: name.to_string(),
pattern: pattern.to_string(),
file_scope: scope,
interrupt_mode: InterruptMode::Always,
}
}
#[test]
fn test_check_delta_matches_simple_pattern() {
let rules = Arc::new(StaticRegistry {
rules: vec![make_rule("no-todo", r"TODO:")],
injections: RwLock::new(Vec::new()),
});
let engine = TtsrEngine::new(
rules,
TtsrSettings {
enabled: true,
..Default::default()
},
);
let ctx = TtsrMatchContext {
source: MatchSource::Text,
file_paths: vec![],
tool_name: None,
file_contents: vec![],
};
let results = engine.check_delta("This code is almost ", &ctx);
assert!(results.is_empty());
let results = engine.check_delta("TODO: fix later", &ctx);
assert_eq!(results.len(), 1);
assert_eq!(results[0].name, "no-todo");
}
#[test]
fn test_check_delta_respects_disabled() {
let rules = Arc::new(StaticRegistry {
rules: vec![make_rule("no-todo", r"TODO:")],
injections: RwLock::new(Vec::new()),
});
let engine = TtsrEngine::new(
rules,
TtsrSettings {
enabled: false, ..Default::default()
},
);
let ctx = TtsrMatchContext {
source: MatchSource::Text,
file_paths: vec![],
tool_name: None,
file_contents: vec![],
};
let results = engine.check_delta("TODO: fix later", &ctx);
assert!(results.is_empty(), "disabled engine must return no matches");
}
#[test]
fn test_scope_filter_respects_tool_scope() {
let rules = Arc::new(StaticRegistry {
rules: vec![Rule {
name: "edit-only-rule".to_string(),
content: "Only for edit tool".to_string(),
description: None,
condition: vec![Regex::new("bad").unwrap()],
scope: vec![ScopeToken::Tool {
name: "edit".to_string(),
globs: vec![],
}],
interrupt_mode: InterruptMode::Always,
globs: vec![],
always_apply: false,
source: RuleSource::BuiltinDefaults,
ast_condition: None,
}],
injections: RwLock::new(Vec::new()),
});
let engine = TtsrEngine::new(
rules,
TtsrSettings {
enabled: true,
..Default::default()
},
);
let text_ctx = TtsrMatchContext {
source: MatchSource::Text,
file_paths: vec![],
tool_name: None,
file_contents: vec![],
};
assert!(engine.check_delta("bad code", &text_ctx).is_empty());
let tool_ctx = TtsrMatchContext {
source: MatchSource::Tool,
file_paths: vec![],
tool_name: Some("edit".to_string()),
file_contents: vec![],
};
assert!(!engine.check_delta("bad code", &tool_ctx).is_empty());
let write_ctx = TtsrMatchContext {
source: MatchSource::Tool,
file_paths: vec![],
tool_name: Some("write".to_string()),
file_contents: vec![],
};
assert!(engine.check_delta("bad code", &write_ctx).is_empty());
}
#[test]
fn test_reset_buffers_clears_accumulation() {
let rules = Arc::new(StaticRegistry {
rules: vec![make_rule("no-todo", r"TODO:")],
injections: RwLock::new(Vec::new()),
});
let engine = TtsrEngine::new(
rules,
TtsrSettings {
enabled: true,
..Default::default()
},
);
let ctx = TtsrMatchContext {
source: MatchSource::Text,
file_paths: vec![],
tool_name: None,
file_contents: vec![],
};
engine.check_delta("TODO", &ctx);
engine.reset_buffers();
let results = engine.check_delta(":", &ctx);
assert!(results.is_empty(), "buffer was reset — TODO should be gone");
}
#[test]
fn test_prose_only_mode_ignores_tool_source() {
let rules = Arc::new(StaticRegistry {
rules: vec![make_rule("no-bad", r"bad")],
injections: RwLock::new(Vec::new()),
});
let engine = TtsrEngine::new(
rules,
TtsrSettings {
enabled: true,
interrupt_mode: InterruptMode::ProseOnly,
..Default::default()
},
);
let text_ctx = TtsrMatchContext {
source: MatchSource::Text,
file_paths: vec![],
tool_name: None,
file_contents: vec![],
};
assert!(!engine.check_delta("bad code", &text_ctx).is_empty());
let tool_ctx = TtsrMatchContext {
source: MatchSource::Tool,
file_paths: vec![],
tool_name: Some("edit".to_string()),
file_contents: vec![],
};
assert!(engine.check_delta("bad code", &tool_ctx).is_empty());
}
#[test]
fn test_ast_match_detects_pattern() {
let ast_rules = vec![make_ast_rule(
"no-box-leak",
"Box::leak",
vec!["*.rs".to_string()],
)];
let mut matcher = TtsrAstMatcher::with_matcher(ast_rules, Box::new(substring_matcher));
let content = "fn main() {\n let _ = Box::leak(Box::new(0));\n}\n";
let result = matcher.check_tool_snapshot("src/main.rs", content);
assert_eq!(result.as_deref(), Some("no-box-leak"));
}
#[test]
fn test_ast_match_no_false_positive() {
let ast_rules = vec![make_ast_rule(
"no-box-leak",
"Box::leak",
vec!["*.rs".to_string()],
)];
let mut matcher = TtsrAstMatcher::with_matcher(ast_rules, Box::new(substring_matcher));
let content = "fn main() {\n println!(\"clean code\");\n}\n";
let result = matcher.check_tool_snapshot("src/main.rs", content);
assert!(result.is_none(), "pattern absent — must not match");
let result = matcher.check_tool_snapshot("src/main.rs", content);
assert!(result.is_none());
let edited = "fn main() {\n println!(\"clean code v2\");\n}\n";
let result = matcher.check_tool_snapshot("src/main.rs", edited);
assert!(result.is_none());
}
#[test]
fn test_ast_match_respects_file_scope() {
let ast_rules = vec![
make_ast_rule("no-rs-leak", "Box::leak", vec!["*.rs".to_string()]),
make_ast_rule("no-ts-leak", "Box::leak", vec!["*.ts".to_string()]),
];
let mut matcher = TtsrAstMatcher::with_matcher(ast_rules, Box::new(substring_matcher));
let ts_content = "export const x = Box::leak(new Object());\n";
let result = matcher.check_tool_snapshot("app/index.ts", ts_content);
assert_eq!(result.as_deref(), Some("no-ts-leak"));
let md_content = "Documentation note: Box::leak is forbidden.\n";
let result = matcher.check_tool_snapshot("docs/notes.md", md_content);
assert!(
result.is_none(),
"scope filter must exclude out-of-scope files"
);
let mut permissive = TtsrAstMatcher::with_matcher(
vec![make_ast_rule("global", "forbidden-token", vec![])],
Box::new(substring_matcher),
);
let result = permissive.check_tool_snapshot("any/path.xyz", "has forbidden-token here");
assert_eq!(result.as_deref(), Some("global"));
}
#[test]
fn test_engine_ast_integration_via_tool_delta() {
let registry_rules = vec![Rule {
name: "no-box-leak".to_string(),
content: "Do not call Box::leak.".to_string(),
description: None,
condition: vec![],
scope: vec![ScopeToken::Tool {
name: "write".to_string(),
globs: vec![],
}],
interrupt_mode: InterruptMode::Always,
globs: vec![],
always_apply: false,
source: RuleSource::BuiltinDefaults,
ast_condition: Some("Box::leak".to_string()),
}];
let registry: Arc<dyn RuleRegistry> = Arc::new(StaticRegistry {
rules: registry_rules,
injections: RwLock::new(Vec::new()),
});
let ast_matcher = TtsrAstMatcher::with_matcher(
vec![make_ast_rule(
"no-box-leak",
"Box::leak",
vec!["*.rs".to_string()],
)],
Box::new(substring_matcher),
);
let engine = TtsrEngine::with_ast_matcher(
registry,
TtsrSettings {
enabled: true,
..Default::default()
},
ast_matcher,
);
let ctx = TtsrMatchContext {
source: MatchSource::Tool,
file_paths: vec!["src/main.rs".to_string()],
tool_name: Some("write".to_string()),
file_contents: vec![(
"src/main.rs".to_string(),
"fn main() { let _ = Box::leak(Box::new(1)); }\n".to_string(),
)],
};
let matched = engine.check_delta("editing src/main.rs", &ctx);
assert_eq!(matched.len(), 1);
assert_eq!(matched[0].name, "no-box-leak");
}
}