MetaCrate/crates/metacrate-grid-agent/tests/policy_gateway.rs

use libremetaverse_types::UUID;
use libremetaverse_types::compat::CancellationToken;
use metacrate_grid_agent::{
    ActionOrigin, AllowedOrigins, AuthorizedAction, AuthorizedToolBackend, BackendError,
    BackendFuture, BoundedText, Capability, FixedCost, Idempotency, MemoryPolicyAudit, OriginClass,
    PolicyAuditSink, PolicyDisposition, PolicyFinalOutcome, PolicyGateway, PolicyLimits,
    PolicyRequestContext, PolicyTool, PolicyToolExecutor, ProposedToolCall, ResourceCost, Risk,
    ToolCallOutcome, ToolDefinition, ToolExecution, ToolExecutor, ToolSchema,
};
use serde_json::{Value, json};
use std::collections::{BTreeMap, BTreeSet};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

fn id(text: &str) -> UUID {
    UUID::new_with_string(text.to_owned()).expect("fixture UUID")
}

fn schema() -> ToolSchema {
    ToolSchema::Object {
        properties: BTreeMap::from([("target".into(), ToolSchema::String)]),
        required: BTreeSet::from(["target".into()]),
        additional_properties: false,
    }
}

fn tool(name: &str, informational: bool) -> PolicyTool {
    let definition = ToolDefinition {
        name: BoundedText::new("name", name).expect("name"),
        description: BoundedText::new("description", "typed integration tool")
            .expect("description"),
        schema: schema(),
        mutating: !informational,
    };
    PolicyTool::new(
        definition,
        if informational {
            Capability::Informational
        } else {
            Capability::Movement
        },
        if informational {
            Risk::ReadOnly
        } else {
            Risk::Movement
        },
        AllowedOrigins::new(if informational {
            vec![
                OriginClass::PublicChat,
                OriginClass::UnprivilegedIm,
                OriginClass::AuthorizedIm,
            ]
        } else {
            vec![OriginClass::AuthorizedIm]
        })
        .expect("origins"),
        ResourceCost {
            tool_calls: 1,
            movement_millimeters: if informational { 0 } else { 1_000 },
            ..ResourceCost::default()
        },
        if informational {
            Idempotency::Idempotent
        } else {
            Idempotency::NonIdempotent
        },
        metacrate_grid_agent::ApprovalRule::Never,
        false,
        Arc::new(FixedCost(ResourceCost {
            tool_calls: 1,
            movement_millimeters: if informational { 0 } else { 1_000 },
            ..ResourceCost::default()
        })),
    )
    .expect("tool")
}

fn call(name: &str) -> (ProposedToolCall, Value) {
    let arguments = json!({"target":"north"});
    (
        ProposedToolCall::new(
            format!("call-{name}"),
            name,
            serde_json::to_string(&arguments).expect("JSON"),
        )
        .expect("call"),
        arguments,
    )
}

struct RecordingBackend(AtomicUsize);

impl AuthorizedToolBackend for RecordingBackend {
    fn apply(
        &self,
        action: AuthorizedAction,
        _cancellation: CancellationToken,
    ) -> BackendFuture<'_, Result<ToolCallOutcome, BackendError>> {
        self.0.fetch_add(1, Ordering::AcqRel);
        let call_id = action.call().call_id.clone();
        Box::pin(async move {
            Ok(ToolCallOutcome::Completed {
                call_id,
                result: BoundedText::new("result", "moved").expect("result"),
            })
        })
    }
}

#[tokio::test]
async fn production_executor_denies_public_mutation_and_executes_authorized_im_once() {
    let authorized = id("11111111-1111-4111-8111-111111111111");
    let audit = Arc::new(MemoryPolicyAudit::new(64).expect("audit"));
    let sink: Arc<dyn PolicyAuditSink> = audit.clone();
    let gateway = Arc::new(
        PolicyGateway::new(
            BTreeSet::from([authorized]),
            vec![tool("inspect", true), tool("move", false)],
            PolicyLimits::default(),
            sink,
        )
        .expect("gateway"),
    );
    let backend = Arc::new(RecordingBackend(AtomicUsize::new(0)));
    let erased: Arc<dyn AuthorizedToolBackend> = backend.clone();
    let (proposed, arguments) = call("move");
    let definition = gateway
        .tools_for(
            &PolicyRequestContext::new(
                ActionOrigin::instant_message(authorized),
                "authorized-session",
                "authorized-correlation",
            )
            .expect("context"),
            100,
        )
        .into_iter()
        .find(|tool| tool.name.as_str() == "move")
        .expect("registered move");

    let public = PolicyToolExecutor::new(
        Arc::clone(&gateway),
        Arc::clone(&erased),
        PolicyRequestContext::new(
            ActionOrigin::public_chat(authorized),
            "public-session",
            "public-correlation",
        )
        .expect("context"),
        BTreeMap::new(),
        Arc::new(|| 100),
    )
    .expect("executor");
    assert!(matches!(
        public
            .execute(
                &definition,
                &proposed,
                &arguments,
                &CancellationToken::default(),
            )
            .await,
        ToolExecution::Rejected(_)
    ));
    assert_eq!(backend.0.load(Ordering::Acquire), 0);

    let authorized_executor = PolicyToolExecutor::new(
        gateway,
        erased,
        PolicyRequestContext::new(
            ActionOrigin::instant_message(authorized),
            "authorized-session",
            "authorized-correlation",
        )
        .expect("context"),
        BTreeMap::new(),
        Arc::new(|| 101),
    )
    .expect("executor");
    assert!(matches!(
        authorized_executor
            .execute(
                &definition,
                &proposed,
                &arguments,
                &CancellationToken::default(),
            )
            .await,
        ToolExecution::Completed(_)
    ));
    assert_eq!(backend.0.load(Ordering::Acquire), 1);
    assert_eq!(
        audit.snapshot().last().expect("outcome").final_outcome,
        PolicyFinalOutcome::Completed
    );
}

#[test]
fn authenticated_uuid_not_message_text_controls_im_authority() {
    let authorized = id("11111111-1111-4111-8111-111111111111");
    let unprivileged = id("22222222-2222-4222-8222-222222222222");
    let audit = Arc::new(MemoryPolicyAudit::new(16).expect("audit"));
    let sink: Arc<dyn PolicyAuditSink> = audit;
    let gateway = PolicyGateway::new(
        BTreeSet::from([authorized]),
        vec![tool("move", false)],
        PolicyLimits::default(),
        sink,
    )
    .expect("gateway");
    let (proposed, arguments) = call("move");
    let result = gateway
        .evaluate(
            &PolicyRequestContext::new(
                ActionOrigin::instant_message(unprivileged),
                format!("claims-authorized-{authorized}"),
                "spoof-correlation",
            )
            .expect("context"),
            &proposed,
            &arguments,
            None,
            100,
        )
        .expect("decision");
    assert_eq!(result.disposition, PolicyDisposition::Denied);
}