1use crate::audit::{AuditLogger, SecurityEvent};
4use crate::score::RiskScorer;
5use std::sync::Arc;
6
7#[derive(Clone)]
9pub struct TRonQuery {
10 pub(crate) audit: Arc<AuditLogger>,
11}
12
13impl TRonQuery {
14 pub async fn recent_events(&self, limit: usize) -> Vec<SecurityEvent> {
16 self.audit.recent(limit).await
17 }
18
19 pub async fn agent_risk_score(&self, agent_id: &str) -> f64 {
21 RiskScorer::score(&self.audit, agent_id).await
22 }
23
24 pub async fn total_events(&self) -> usize {
26 self.audit.total_count().await
27 }
28
29 pub async fn total_denials(&self) -> usize {
31 self.audit.deny_count().await
32 }
33
34 pub async fn agent_audit(&self, agent_id: &str, limit: usize) -> Vec<SecurityEvent> {
36 self.audit.agent_events(agent_id, limit).await
37 }
38
39 pub fn verify_chain(&self) -> libro::Result<()> {
41 self.audit.verify_chain()
42 }
43
44 pub fn chain_review(&self) -> libro::ChainReview {
46 self.audit.chain_review()
47 }
48
49 pub fn chain_len(&self) -> usize {
51 self.audit.chain_len()
52 }
53}
54
55#[cfg(test)]
56mod tests {
57 use crate::{DefaultAction, TRon, TRonConfig};
58
59 fn permissive_config() -> TRonConfig {
60 TRonConfig {
61 default_unknown_agent: DefaultAction::Allow,
62 default_unknown_tool: DefaultAction::Allow,
63 scan_payloads: false,
64 analyze_patterns: false,
65 ..Default::default()
66 }
67 }
68
69 #[tokio::test]
70 async fn query_api_initial_state() {
71 let tron = TRon::new(TRonConfig::default());
72 let query = tron.query();
73 assert_eq!(query.total_events().await, 0);
74 assert_eq!(query.total_denials().await, 0);
75 assert!(query.recent_events(10).await.is_empty());
76 }
77
78 #[tokio::test]
79 async fn query_after_checks() {
80 let tron = TRon::new(permissive_config());
81 let query = tron.query();
82
83 let call = crate::gate::ToolCall {
85 agent_id: "agent-1".to_string(),
86 tool_name: "tarang_probe".to_string(),
87 params: serde_json::json!({}),
88 timestamp: chrono::Utc::now(),
89 };
90 for _ in 0..5 {
91 tron.check(&call).await;
92 }
93
94 assert_eq!(query.total_events().await, 5);
95 assert_eq!(query.total_denials().await, 0);
96
97 let events = query.recent_events(3).await;
98 assert_eq!(events.len(), 3);
99 }
100
101 #[tokio::test]
102 async fn query_risk_score_after_denials() {
103 let tron = TRon::new(TRonConfig::default());
104 let query = tron.query();
105
106 let call = crate::gate::ToolCall {
108 agent_id: "bad-agent".to_string(),
109 tool_name: "anything".to_string(),
110 params: serde_json::json!({}),
111 timestamp: chrono::Utc::now(),
112 };
113 for _ in 0..5 {
114 let v = tron.check(&call).await;
115 assert!(v.is_denied());
116 }
117
118 assert_eq!(query.total_denials().await, 5);
119 assert_eq!(query.agent_risk_score("bad-agent").await, 1.0);
120 assert_eq!(query.agent_risk_score("nobody").await, 0.0);
121 }
122
123 #[tokio::test]
124 async fn query_agent_audit_trail() {
125 let tron = TRon::new(TRonConfig::default());
126 let query = tron.query();
127
128 let call_a = crate::gate::ToolCall {
130 agent_id: "agent-a".to_string(),
131 tool_name: "tool".to_string(),
132 params: serde_json::json!({}),
133 timestamp: chrono::Utc::now(),
134 };
135 let call_b = crate::gate::ToolCall {
136 agent_id: "agent-b".to_string(),
137 tool_name: "tool".to_string(),
138 params: serde_json::json!({}),
139 timestamp: chrono::Utc::now(),
140 };
141 for _ in 0..3 {
142 tron.check(&call_a).await;
143 }
144 for _ in 0..7 {
145 tron.check(&call_b).await;
146 }
147
148 let trail_a = query.agent_audit("agent-a", 100).await;
149 let trail_b = query.agent_audit("agent-b", 100).await;
150 assert_eq!(trail_a.len(), 3);
151 assert_eq!(trail_b.len(), 7);
152
153 assert_eq!(query.agent_audit("agent-b", 2).await.len(), 2);
155 }
156
157 #[tokio::test]
158 async fn query_chain_verification() {
159 let tron = TRon::new(permissive_config());
160 let query = tron.query();
161
162 let call = crate::gate::ToolCall {
163 agent_id: "agent-1".to_string(),
164 tool_name: "tool".to_string(),
165 params: serde_json::json!({}),
166 timestamp: chrono::Utc::now(),
167 };
168 for _ in 0..10 {
169 tron.check(&call).await;
170 }
171
172 assert!(query.verify_chain().is_ok());
173 assert_eq!(query.chain_len(), 10);
174
175 let review = query.chain_review();
176 assert_eq!(review.entry_count, 10);
177 }
178}