Testing Guards & Policies

use chio_kernel_core::{Guard, GuardContext, KernelCoreError, Verdict};

/// Admission-time version of the document application's word-limit rule.
pub struct WordLimitGuard {
    pub max_words: Option<usize>,
}

impl Guard for WordLimitGuard {
    fn name(&self) -> &str {
        "word-limit"
    }

    fn evaluate(&self, ctx: &GuardContext<'_>) -> Result<Verdict, KernelCoreError> {
        let request = ctx.request;
        if request.server_id != "documents" || request.tool_name != "count_words" {
            return Ok(Verdict::Allow);
        }
        let Some(limit) = self.max_words.filter(|n| (1..=10_000).contains(n)) else {
            return Ok(Verdict::Deny);
        };
        let Some(text) = request.arguments.get("text").and_then(|v| v.as_str()) else {
            return Ok(Verdict::Deny);
        };
        Ok(if text.split_whitespace().take(limit + 1).count() > limit {
            Verdict::Deny
        } else {
            Verdict::Allow
        })
    }
}

use chio_core_types::capability::scope::ChioScope;
use chio_kernel_core::{Guard, GuardContext, PortableToolCallRequest, Verdict};
use guard_tests::WordLimitGuard;
use serde_json::{json, Value};

fn check(server: &str, tool: &str, arguments: Value, limit: Option<usize>) -> Verdict {
    let request = PortableToolCallRequest {
        request_id: "unit-document-1".into(),
        tool_name: tool.into(),
        server_id: server.into(),
        agent_id: "unit-agent".into(),
        arguments,
    };
    let scope = ChioScope::default();
    let context = GuardContext {
        request: &request,
        scope: &scope,
        agent_id: &request.agent_id,
        server_id: &request.server_id,
        session_filesystem_roots: None,
        matched_grant_index: None,
    };
    WordLimitGuard { max_words: limit }
        .evaluate(&context)
        .unwrap()
}

#[test]
fn six_words_are_allowed() {
    assert_eq!(
        check(
            "documents",
            "count_words",
            json!({
                "text": "One governed tool, two protocol surfaces."
            }),
            Some(8)
        ),
        Verdict::Allow
    );
}

#[test]
fn eleven_words_are_denied() {
    assert_eq!(
        check(
            "documents",
            "count_words",
            json!({
                "text": "This document has eleven words and should exceed the configured limit."
            }),
            Some(8)
        ),
        Verdict::Deny
    );
}

#[test]
fn changing_the_limit_allows_the_same_document() {
    assert_eq!(
        check(
            "documents",
            "count_words",
            json!({
                "text": "This document has eleven words and should exceed the configured limit."
            }),
            Some(16)
        ),
        Verdict::Allow
    );
}

#[test]
fn exactly_eight_words_are_allowed() {
    assert_eq!(
        check(
            "documents",
            "count_words",
            json!({
                "text": "one two three four five six seven eight"
            }),
            Some(8)
        ),
        Verdict::Allow
    );
}

#[test]
fn whitespace_is_not_a_model_token_count() {
    assert_eq!(
        check(
            "documents",
            "count_words",
            json!({
                "text": "one\ttwo\nthree   four"
            }),
            Some(4)
        ),
        Verdict::Allow
    );
}

#[test]
fn missing_or_invalid_configuration_denies() {
    for limit in [None, Some(0), Some(10_001), Some(usize::MAX)] {
        assert_eq!(
            check("documents", "count_words", json!({"text": "one"}), limit),
            Verdict::Deny
        );
    }
}

#[test]
fn missing_or_non_string_text_denies() {
    for arguments in [json!({}), json!({"text": null}), json!({"text": 12})] {
        assert_eq!(
            check("documents", "count_words", arguments, Some(8)),
            Verdict::Deny
        );
    }
}

#[test]
fn unrelated_calls_pass_this_guard_only() {
    assert_eq!(
        check("other", "count_words", json!({}), None),
        Verdict::Allow
    );
    assert_eq!(
        check("documents", "other_tool", json!({}), None),
        Verdict::Allow
    );
}

#[test]
fn whitespace_and_unicode_preserve_the_word_limit_invariant() {
    for limit in [1usize, 2, 7, 32, 127] {
        for count in [0, limit.saturating_sub(1), limit, limit + 1, limit + 9] {
            for separator in [" ", "\t", "\n", "\u{2003}"] {
                let text = vec!["café"; count].join(separator);
                let verdict = check(
                    "documents",
                    "count_words",
                    json!({"text":text}),
                    Some(limit),
                );
                assert_eq!(verdict == Verdict::Allow, count <= limit);
            }
        }
    }
}

use chio_core_types::capability::runtime_attestation::RuntimeAssuranceTier;
use chio_policy::evaluate::RuntimeAttestationContext;
use chio_policy::{compile_policy, evaluate, Decision, EvaluationAction, HushSpec};

const PREFERRED: &str = include_str!("../policies/preferred.yaml");
const REQUIRED: &str = include_str!("../policies/required.yaml");

fn document_action() -> EvaluationAction {
    EvaluationAction {
        action_type: "tool_call".into(),
        target: Some("count_words".into()),
        ..Default::default()
    }
}

#[test]
fn missing_preferred_assurance_warns() {
    let spec = HushSpec::parse(PREFERRED).unwrap();
    let result = evaluate(&spec, &document_action());
    assert_eq!(result.decision, Decision::Warn);
    assert_eq!(
        result.matched_rule.as_deref(),
        Some("rules.tool_access.prefer_runtime_assurance_tier")
    );
}

#[test]
fn missing_required_assurance_denies() {
    let spec = HushSpec::parse(REQUIRED).unwrap();
    let result = evaluate(&spec, &document_action());
    assert_eq!(result.decision, Decision::Deny);
}

#[test]
fn valid_attested_context_satisfies_both_policies() {
    let mut action = document_action();
    action.runtime_attestation = Some(RuntimeAttestationContext {
        tier: RuntimeAssuranceTier::Attested,
        valid: true,
        verifier: None,
        workload_identity: None,
    });
    for yaml in [PREFERRED, REQUIRED] {
        assert_eq!(
            evaluate(&HushSpec::parse(yaml).unwrap(), &action).decision,
            Decision::Allow
        );
    }
}

#[test]
fn an_invalid_attestation_does_not_satisfy_a_requirement() {
    let mut action = document_action();
    action.runtime_attestation = Some(RuntimeAttestationContext {
        tier: RuntimeAssuranceTier::Attested,
        valid: false,
        verifier: None,
        workload_identity: None,
    });
    assert_eq!(
        evaluate(&HushSpec::parse(REQUIRED).unwrap(), &action).decision,
        Decision::Deny
    );
}

#[test]
fn compiling_a_preference_does_not_create_a_required_constraint() {
    let preferred = compile_policy(&HushSpec::parse(PREFERRED).unwrap()).unwrap();
    let required = compile_policy(&HushSpec::parse(REQUIRED).unwrap()).unwrap();
    assert_eq!(preferred.default_scope.grants.len(), 1);
    assert!(preferred.default_scope.grants[0].constraints.is_empty());
    assert_eq!(required.default_scope.grants.len(), 1);
    assert_eq!(required.default_scope.grants[0].constraints.len(), 1);
}