use anyhow::{bail, Context, Result};
use regex::Regex;
use serde::{Deserialize, Serialize};
use std::collections::BTreeSet;
use std::path::PathBuf;
use crate::analysis_pipeline::{compile_nonempty_regex, require_name};
use crate::hashing::sha256;
pub const CORRECTION_POLICY_SCHEMA_VERSION: u32 = 1;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum CorrectionPolicySource {
Embedded,
File { path: PathBuf },
Inline,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct CorrectionPolicySpec {
pub schema_version: u32,
pub name: String,
pub version: String,
pub categories: Vec<CorrectionCategorySpec>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct CorrectionCategorySpec {
pub name: String,
pub patterns: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CorrectionPolicyIdentity {
pub name: String,
pub version: String,
pub sha256: String,
pub source: CorrectionPolicySource,
}
#[derive(Debug, Clone)]
pub struct CorrectionPolicy {
identity: CorrectionPolicyIdentity,
rules: Vec<(String, Regex)>,
category_count: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CorrectionHit<'policy> {
pub category: &'policy str,
pub matched_text: String,
pub match_start_char: usize,
pub match_end_char_exclusive: usize,
}
impl CorrectionPolicySpec {
pub(crate) fn compile_in_memory(
self,
source: CorrectionPolicySource,
) -> Result<CorrectionPolicy> {
let digest = sha256(&canonical_digest_input(&self));
self.compile(source, digest)
}
pub fn compile(
self,
source: CorrectionPolicySource,
digest: String,
) -> Result<CorrectionPolicy> {
if self.schema_version != CORRECTION_POLICY_SCHEMA_VERSION {
bail!(
"unsupported correction policy schema_version {}; this build understands {}. \
Upgrade aise, or set schema_version = {}",
self.schema_version,
CORRECTION_POLICY_SCHEMA_VERSION,
CORRECTION_POLICY_SCHEMA_VERSION
);
}
require_name("correction policy name", &self.name)?;
require_name("correction policy version", &self.version)?;
if self.categories.is_empty() {
bail!(
"correction policy '{}' defines no categories; add at least one [[categories]] \
table with a name and one or more patterns",
self.name
);
}
let mut seen = BTreeSet::new();
let mut rules = Vec::new();
for category in &self.categories {
require_name("correction category name", &category.name)?;
if !seen.insert(category.name.clone()) {
bail!(
"duplicate correction category '{}'; merge its patterns into one \
[[categories]] table, because the first match wins and the second would \
never be reached",
category.name
);
}
if category.patterns.is_empty() {
bail!(
"correction category '{}' has no patterns; give it at least one, or remove \
the category",
category.name
);
}
for pattern in &category.patterns {
compile_nonempty_regex("correction", &category.name, pattern)?;
}
let joined = format!("(?i){}", category.patterns.join("|"));
let regex = Regex::new(&joined).with_context(|| {
format!(
"correction category '{}' does not compile once its {} patterns are combined; \
each is valid alone, so check for an unbalanced group spanning them",
category.name,
category.patterns.len()
)
})?;
if regex.is_match("") {
bail!(
"correction category '{}' matches empty text after its patterns are combined; \
every category must require at least one input character",
category.name
);
}
rules.push((category.name.clone(), regex));
}
Ok(CorrectionPolicy {
identity: CorrectionPolicyIdentity {
name: self.name,
version: self.version,
sha256: digest,
source,
},
rules,
category_count: self.categories.len(),
})
}
}
impl CorrectionPolicy {
#[cfg(test)]
pub fn parse_toml(source_text: &str, source: CorrectionPolicySource) -> Result<Self> {
let spec: CorrectionPolicySpec = toml::from_str(source_text).with_context(|| {
format!(
"failed to parse correction policy from {}",
describe_source(&source)
)
})?;
let digest = sha256(source_text.as_bytes());
spec.compile(source, digest)
}
pub fn classify(&self, text: &str) -> Option<CorrectionHit<'_>> {
self.classify_fragments(std::iter::once((0, text)))
}
fn classify_fragments<'policy, 'text, I>(
&'policy self,
fragments: I,
) -> Option<CorrectionHit<'policy>>
where
I: Iterator<Item = (usize, &'text str)> + Clone,
{
self.rules.iter().find_map(|(category, regex)| {
fragments.clone().find_map(|(fragment_start_char, text)| {
regex.find(text).map(|matched| {
let match_start_char =
fragment_start_char.saturating_add(text[..matched.start()].chars().count());
let matched_text = matched.as_str().to_string();
let match_end_char_exclusive =
match_start_char.saturating_add(matched_text.chars().count());
CorrectionHit {
category: category.as_str(),
matched_text,
match_start_char,
match_end_char_exclusive,
}
})
})
})
}
pub fn identity(&self) -> &CorrectionPolicyIdentity {
&self.identity
}
pub fn category_count(&self) -> usize {
self.category_count
}
pub fn rules(&self) -> &[(String, Regex)] {
&self.rules
}
}
pub const EMBEDDED_POLICY_NAME: &str = "corrections";
#[derive(Debug, Clone)]
pub struct ResolvedCorrectionPolicySet {
policies: Vec<CorrectionPolicy>,
}
impl ResolvedCorrectionPolicySet {
pub fn from_policies(policies: Vec<CorrectionPolicy>) -> Self {
Self { policies }
}
pub fn classify(&self, text: &str) -> Option<(&CorrectionPolicyIdentity, CorrectionHit<'_>)> {
self.policies
.iter()
.find_map(|policy| policy.classify(text).map(|hit| (policy.identity(), hit)))
}
pub(crate) fn classify_fragments<'policy, 'text, I>(
&'policy self,
fragments: I,
) -> Option<(&'policy CorrectionPolicyIdentity, CorrectionHit<'policy>)>
where
I: Iterator<Item = (usize, &'text str)> + Clone,
{
self.policies.iter().find_map(|policy| {
policy
.classify_fragments(fragments.clone())
.map(|hit| (policy.identity(), hit))
})
}
pub fn receipts(&self) -> Vec<CapabilityReceipt> {
self.policies
.iter()
.map(|policy| {
let identity = policy.identity();
CapabilityReceipt {
name: identity.name.clone(),
version: identity.version.clone(),
sha256: identity.sha256.clone(),
}
})
.collect()
}
pub fn policies(&self) -> &[CorrectionPolicy] {
&self.policies
}
pub fn is_empty(&self) -> bool {
self.policies.is_empty()
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CapabilityReceipt {
pub name: String,
pub version: String,
pub sha256: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessageClassificationReport {
pub policies: Vec<CapabilityReceipt>,
pub matches: Vec<crate::models::MessageClassificationMatch>,
}
pub(crate) fn embedded_policy() -> Result<CorrectionPolicy> {
crate::message_classification::MessageClassificationPolicySpec::parse_toml(
EMBEDDED_POLICY_TOML,
)?
.compile(
EMBEDDED_POLICY_NAME.to_string(),
env!("CARGO_PKG_VERSION").to_string(),
CorrectionPolicySource::Embedded,
EMBEDDED_POLICY_TOML.as_bytes(),
)
}
pub(crate) const EMBEDDED_POLICY_TOML: &str =
include_str!("../skills/ai-session-search/aise-capability.toml");
pub(crate) fn canonical_digest_input_len(spec: &CorrectionPolicySpec) -> usize {
let mut bytes = 0_usize;
let mut add = |value: &str| {
bytes = bytes.saturating_add(std::mem::size_of::<u64>() + value.len());
};
add(&spec.schema_version.to_string());
add(&spec.name);
add(&spec.version);
for category in &spec.categories {
add(&category.name);
for pattern in &category.patterns {
add(pattern);
}
}
bytes
}
fn canonical_digest_input(spec: &CorrectionPolicySpec) -> Vec<u8> {
let mut out = Vec::new();
let mut push = |value: &str| {
out.extend_from_slice(&(value.len() as u64).to_le_bytes());
out.extend_from_slice(value.as_bytes());
};
push(&spec.schema_version.to_string());
push(&spec.name);
push(&spec.version);
for category in &spec.categories {
push(&category.name);
for pattern in &category.patterns {
push(pattern);
}
}
out
}
#[cfg(test)]
fn describe_source(source: &CorrectionPolicySource) -> String {
match source {
CorrectionPolicySource::Embedded => "the policy embedded in this executable".to_string(),
CorrectionPolicySource::File { path } => path.display().to_string(),
CorrectionPolicySource::Inline => "typed inline capability parameters".to_string(),
}
}
#[cfg(test)]
mod tests {
#[test]
fn a_policy_may_have_one_category_or_many_but_never_zero() {
let spec = |count: usize| CorrectionPolicySpec {
schema_version: CORRECTION_POLICY_SCHEMA_VERSION,
name: "boundary".to_string(),
version: "1.0.0".to_string(),
categories: (0..count)
.map(|index| CorrectionCategorySpec {
name: format!("c{index}"),
patterns: vec![format!(r"\bmarker{index}\b")],
})
.collect(),
};
let error = spec(0)
.compile_in_memory(CorrectionPolicySource::Embedded)
.expect_err("a policy with no categories matches nothing, silently");
assert!(format!("{error:#}").contains("[[categories]]"), "{error:#}");
for count in [1_usize, 64] {
let policy = spec(count)
.compile_in_memory(CorrectionPolicySource::Embedded)
.unwrap_or_else(|error| panic!("{count} categories must compile: {error:#}"));
assert_eq!(policy.category_count(), count);
let hit = policy
.classify(&format!("marker{}", count - 1))
.expect("the final category still matches");
assert_eq!(hit.category, format!("c{}", count - 1));
}
}
#[test]
fn a_pattern_matching_empty_text_is_refused() {
let spec = CorrectionPolicySpec {
schema_version: CORRECTION_POLICY_SCHEMA_VERSION,
name: "greedy".to_string(),
version: "1.0.0".to_string(),
categories: vec![CorrectionCategorySpec {
name: "everything".to_string(),
patterns: vec![".*".to_string()],
}],
};
let error = spec
.compile_in_memory(CorrectionPolicySource::Embedded)
.expect_err("a rule that matches the empty string matches every message");
let message = format!("{error:#}");
assert!(
message.contains(".*"),
"name the offending pattern: {message}"
);
}
#[test]
fn blank_names_and_versions_are_refused_by_field() {
for (name, version, category, expected) in [
("", "1.0.0", "c", "correction policy name"),
(" ", "1.0.0", "c", "correction policy name"),
("ok", "", "c", "correction policy version"),
("ok", " ", "c", "correction policy version"),
("ok", "1.0.0", " ", "correction category name"),
] {
let spec = CorrectionPolicySpec {
schema_version: CORRECTION_POLICY_SCHEMA_VERSION,
name: name.to_string(),
version: version.to_string(),
categories: vec![CorrectionCategorySpec {
name: category.to_string(),
patterns: vec![r"\bx\b".to_string()],
}],
};
let error = spec
.compile_in_memory(CorrectionPolicySource::Embedded)
.expect_err(&format!(
"{name:?}/{version:?}/{category:?} must be refused"
));
let message = format!("{error:#}");
assert!(
message.contains(expected),
"the message must name WHICH field is blank: {message}"
);
}
}
use super::*;
fn policy(body: &str) -> String {
format!(
"schema_version = 1\nname = \"fixture\"\nversion = \"1.0.0\"\n\n{body}",
body = body
)
}
fn compile(body: &str) -> Result<CorrectionPolicy> {
CorrectionPolicy::parse_toml(&policy(body), CorrectionPolicySource::Embedded)
}
fn message(result: Result<CorrectionPolicy>) -> String {
format!(
"{:#}",
result.expect_err("policy should have been rejected")
)
}
#[test]
fn declaration_order_is_evaluation_order() {
let compiled = compile(
r#"
[[categories]]
name = "regression"
patterns = ['''\byou broke\b''']
[[categories]]
name = "incomplete"
patterns = ['''\bstill need\b''']
[[categories]]
name = "other"
patterns = ['''\bstop\b''']
"#,
)
.unwrap();
let order: Vec<&str> = compiled
.rules()
.iter()
.map(|(category, _)| category.as_str())
.collect();
assert_eq!(
order,
["regression", "incomplete", "other"],
"categories must evaluate in declaration order, not sorted order"
);
}
#[test]
fn first_matching_category_wins_and_later_ones_are_not_reported() {
let compiled = compile(
r#"
[[categories]]
name = "specific"
patterns = ['''\byou broke the build\b''']
[[categories]]
name = "catch_all"
patterns = ['''\bbroke\b''']
"#,
)
.unwrap();
let hit = compiled.classify("you broke the build again").unwrap();
assert_eq!(hit.category, "specific");
assert_eq!(hit.matched_text, "you broke the build");
let hit = compiled.classify("the deploy broke overnight").unwrap();
assert_eq!(hit.category, "catch_all");
}
#[test]
fn patterns_within_a_category_are_ored() {
let compiled = compile(
r#"
[[categories]]
name = "skip_step"
patterns = ['''\byou forgot\b''', '''\byou missed\b''', '''\byou skipped\b''']
"#,
)
.unwrap();
for (text, expected) in [
("you forgot the tests", "you forgot"),
("you missed a step", "you missed"),
("you skipped the lint", "you skipped"),
] {
let hit = compiled.classify(text).unwrap();
assert_eq!(hit.category, "skip_step");
assert_eq!(hit.matched_text, expected);
}
assert!(compiled.classify("everything looks fine").is_none());
}
#[test]
fn categories_are_case_insensitive_and_report_the_leftmost_match() {
let compiled = compile(
r#"
[[categories]]
name = "skip_step"
patterns = ['''\byou forgot\b''', '''\byou missed\b''']
"#,
)
.unwrap();
assert_eq!(
compiled
.classify("You Forgot the tests")
.unwrap()
.matched_text,
"You Forgot",
"matching is case-insensitive, and matched_text preserves the ORIGINAL casing"
);
assert_eq!(
compiled
.classify("you missed X, and you forgot Y")
.unwrap()
.matched_text,
"you missed",
"within a category the leftmost match in the message wins, not the first pattern"
);
}
#[test]
fn matched_text_is_the_substring_not_the_rule() {
let compiled = compile(
r#"
[[categories]]
name = "regression"
patterns = ['''\byou (deleted|removed|reverted)\b''']
"#,
)
.unwrap();
let hit = compiled.classify("é wait, you reverted my fix").unwrap();
assert_eq!(hit.matched_text, "you reverted");
assert_eq!(hit.match_start_char, 8);
assert_eq!(hit.match_end_char_exclusive, 20);
assert_ne!(
hit.matched_text, r"\byou (deleted|removed|reverted)\b",
"the field must not carry the regex source"
);
}
#[test]
fn identity_digests_the_exact_source_bytes() {
let body = r#"
[[categories]]
name = "regression"
patterns = ['''\byou broke\b''']
"#;
let first = compile(body).unwrap();
let second = compile(body).unwrap();
assert_eq!(
first.identity().sha256,
second.identity().sha256,
"identical bytes must digest identically"
);
let commented = compile(&format!("# a note\n{body}")).unwrap();
assert_ne!(
first.identity().sha256,
commented.identity().sha256,
"a byte-level edit must change the digest even when the parsed values match"
);
assert_eq!(first.identity().name, "fixture");
assert_eq!(first.identity().version, "1.0.0");
assert_eq!(first.identity().source, CorrectionPolicySource::Embedded);
}
#[test]
fn unknown_schema_version_is_rejected_by_name() {
let source = "schema_version = 2\nname = \"f\"\nversion = \"1\"\n\n[[categories]]\nname = \"c\"\npatterns = ['''x''']\n";
let err = CorrectionPolicy::parse_toml(source, CorrectionPolicySource::Embedded);
let err = message(err);
assert!(err.contains("schema_version 2"), "{err}");
assert!(
err.contains("Upgrade aise"),
"must say what to do, not only what is wrong: {err}"
);
}
#[test]
fn unknown_field_is_a_hard_error_not_a_silent_drop() {
let err = compile(
r#"
[[categories]]
name = "regression"
pattern = ['''\byou broke\b''']
"#,
);
let err = message(err);
assert!(
err.contains("pattern"),
"a misspelled key must be named, not ignored: {err}"
);
}
#[test]
fn duplicate_category_names_are_rejected_because_the_second_is_unreachable() {
let err = compile(
r#"
[[categories]]
name = "regression"
patterns = ['''\byou broke\b''']
[[categories]]
name = "regression"
patterns = ['''\byou deleted\b''']
"#,
);
let err = message(err);
assert!(
err.contains("duplicate correction category 'regression'"),
"{err}"
);
assert!(err.contains("never be reached"), "must explain why: {err}");
}
#[test]
fn structural_emptiness_is_rejected_at_every_level() {
let no_categories = CorrectionPolicy::parse_toml(
"schema_version = 1\nname = \"f\"\nversion = \"1\"\ncategories = []\n",
CorrectionPolicySource::Embedded,
);
let no_categories = message(no_categories);
assert!(
no_categories.contains("defines no categories"),
"{no_categories}"
);
let no_patterns = message(compile(
"[[categories]]\nname = \"regression\"\npatterns = []\n",
));
assert!(no_patterns.contains("has no patterns"), "{no_patterns}");
let blank_name = message(compile(
"[[categories]]\nname = \" \"\npatterns = ['''x''']\n",
));
assert!(blank_name.contains("must not be empty"), "{blank_name}");
let blank_policy_name = message(CorrectionPolicy::parse_toml(
"schema_version = 1\nname = \"\"\nversion = \"1\"\n\n[[categories]]\nname = \"c\"\npatterns = ['''x''']\n",
CorrectionPolicySource::Embedded,
));
assert!(
blank_policy_name.contains("must not be empty"),
"{blank_policy_name}"
);
}
#[test]
fn invalid_and_empty_matching_regexes_are_rejected() {
let invalid = message(compile(
"[[categories]]\nname = \"c\"\npatterns = ['''(unclosed''']\n",
));
assert!(invalid.contains("invalid correction regex"), "{invalid}");
let empty_matching = message(compile(
"[[categories]]\nname = \"c\"\npatterns = ['''a*''']\n",
));
assert!(
empty_matching.contains("matches empty text"),
"{empty_matching}"
);
assert!(
empty_matching.contains("a*"),
"and it names WHICH pattern, since a category may hold several: {empty_matching}"
);
let empty_pattern = message(compile(
"[[categories]]\nname = \"c\"\npatterns = ['''''']\n",
));
assert!(
empty_pattern.contains("must not be empty"),
"{empty_pattern}"
);
}
#[test]
fn the_default_resolution_compiles_to_todays_built_in_rules() {
let config = crate::config::Config::default();
let resolved = ResolvedCorrectionPolicySet::from_policies(vec![embedded_policy().unwrap()]);
let new_rules: Vec<(String, String)> = resolved.policies()[0]
.rules()
.iter()
.map(|(category, regex)| (category.clone(), regex.as_str().to_string()))
.collect();
let old_rules: Vec<(String, String)> = crate::analytics::compile_patterns(&config)
.unwrap()
.iter()
.map(|(category, regex)| (category.clone(), regex.as_str().to_string()))
.collect();
assert_eq!(
new_rules, old_rules,
"the embedded policy must be indistinguishable from the built-ins it replaces"
);
assert_eq!(resolved.receipts()[0].name, EMBEDDED_POLICY_NAME);
}
#[test]
fn malformed_toml_names_where_it_came_from() {
let err = CorrectionPolicy::parse_toml(
"schema_version = ",
CorrectionPolicySource::File {
path: PathBuf::from("/skills/mine/corrections/policy.toml"),
},
);
let err = message(err);
assert!(
err.contains("/skills/mine/corrections/policy.toml"),
"a parse failure must name the file: {err}"
);
}
}