Skip to main content

nexara_runtime/
lib.rs

1use async_trait::async_trait;
2use nexara_core::{
3    ActionClass, AuditOutcome, AuditRecord, AuditSink, EffectiveTrustPolicy, NexaraError,
4    PolicyContract, PolicyDecision, PolicyEvaluation, PolicyEvaluationDecision, ToolBroker,
5    ToolCallRequest, ToolCallResult, ToolDescriptor, ToolRegistry, ToolSelectionRequest,
6    ToolUsageSignalProvider, TracingAuditSink, TrustProfile, hash_result, matched_rules_metadata,
7};
8use serde::{Deserialize, Serialize};
9use std::sync::Arc;
10use std::time::Instant;
11use tokio::sync::Semaphore;
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct NexaraRuntimeConfig {
15    pub max_tools_per_request: usize,
16    pub max_payload_bytes: usize,
17    pub max_concurrent_calls: usize,
18    pub default_scopes: Vec<String>,
19    pub tool_selection_enabled: bool,
20}
21
22impl Default for NexaraRuntimeConfig {
23    fn default() -> Self {
24        Self {
25            max_tools_per_request: 8,
26            max_payload_bytes: 64 * 1024,
27            max_concurrent_calls: 8,
28            default_scopes: vec!["read".to_string()],
29            tool_selection_enabled: true,
30        }
31    }
32}
33
34pub trait ToolCatalogProvider: Send + Sync {
35    fn tools(&self) -> Vec<ToolDescriptor>;
36}
37
38#[derive(Default)]
39pub struct CompositeToolCatalog {
40    providers: Vec<Arc<dyn ToolCatalogProvider>>,
41}
42
43impl CompositeToolCatalog {
44    pub fn new(providers: Vec<Arc<dyn ToolCatalogProvider>>) -> Self {
45        Self { providers }
46    }
47}
48
49impl ToolCatalogProvider for CompositeToolCatalog {
50    fn tools(&self) -> Vec<ToolDescriptor> {
51        self.providers
52            .iter()
53            .flat_map(|provider| provider.tools())
54            .collect()
55    }
56}
57
58#[derive(Debug, Clone)]
59pub struct StaticToolCatalog {
60    tools: Vec<ToolDescriptor>,
61}
62
63impl StaticToolCatalog {
64    pub fn new(tools: Vec<ToolDescriptor>) -> Self {
65        Self { tools }
66    }
67}
68
69impl ToolCatalogProvider for StaticToolCatalog {
70    fn tools(&self) -> Vec<ToolDescriptor> {
71        self.tools.clone()
72    }
73}
74
75#[async_trait]
76pub trait HostToolExecutor<C>: Send + Sync {
77    async fn call_host_tool(
78        &self,
79        tool: &ToolDescriptor,
80        request: ToolCallRequest,
81        context: C,
82    ) -> Result<ToolCallResult, NexaraError>;
83}
84
85pub trait TrustPolicyResolver<C>: Send + Sync {
86    fn resolve(&self, context: &C) -> EffectiveTrustPolicy;
87}
88
89#[async_trait]
90pub trait RuntimeLifecycleHooks<C>: Send + Sync {
91    async fn before_call(&self, _tool: &ToolDescriptor, _request: &ToolCallRequest, _context: &C) {}
92
93    async fn after_call(
94        &self,
95        _tool: &ToolDescriptor,
96        _request: &ToolCallRequest,
97        _outcome: Result<&ToolCallResult, &NexaraError>,
98        _context: &C,
99    ) {
100    }
101}
102
103#[derive(Debug, Default)]
104pub struct NoopRuntimeLifecycleHooks;
105
106#[async_trait]
107impl<C> RuntimeLifecycleHooks<C> for NoopRuntimeLifecycleHooks where C: Send + Sync {}
108
109#[derive(Debug, Clone)]
110pub struct StaticTrustPolicyResolver {
111    profile: TrustProfile,
112}
113
114impl StaticTrustPolicyResolver {
115    pub fn new(profile: TrustProfile) -> Self {
116        Self { profile }
117    }
118}
119
120impl Default for StaticTrustPolicyResolver {
121    fn default() -> Self {
122        Self {
123            profile: TrustProfile::Observe,
124        }
125    }
126}
127
128impl<C> TrustPolicyResolver<C> for StaticTrustPolicyResolver {
129    fn resolve(&self, _context: &C) -> EffectiveTrustPolicy {
130        EffectiveTrustPolicy::for_profile(self.profile)
131    }
132}
133
134pub struct NexaraRuntime<C> {
135    config: NexaraRuntimeConfig,
136    catalog: Arc<dyn ToolCatalogProvider>,
137    executor: Arc<dyn HostToolExecutor<C>>,
138    trust_resolver: Arc<dyn TrustPolicyResolver<C>>,
139    audit_sink: Arc<dyn AuditSink>,
140    lifecycle_hooks: Arc<dyn RuntimeLifecycleHooks<C>>,
141    broker: ToolBroker,
142    policy_contract: Option<PolicyContract>,
143    call_semaphore: Arc<Semaphore>,
144}
145
146impl<C> NexaraRuntime<C>
147where
148    C: Clone + Send + Sync + 'static,
149{
150    pub fn new(
151        config: NexaraRuntimeConfig,
152        catalog: Arc<dyn ToolCatalogProvider>,
153        executor: Arc<dyn HostToolExecutor<C>>,
154        trust_resolver: Arc<dyn TrustPolicyResolver<C>>,
155    ) -> Self {
156        let max_concurrent_calls = config.max_concurrent_calls;
157        Self {
158            config,
159            catalog,
160            executor,
161            trust_resolver,
162            audit_sink: Arc::new(TracingAuditSink),
163            lifecycle_hooks: Arc::new(NoopRuntimeLifecycleHooks),
164            broker: ToolBroker::new(),
165            policy_contract: None,
166            call_semaphore: Arc::new(Semaphore::new(max_concurrent_calls)),
167        }
168    }
169
170    pub fn with_audit_sink(mut self, audit_sink: Arc<dyn AuditSink>) -> Self {
171        self.audit_sink = audit_sink;
172        self
173    }
174
175    pub fn with_broker(mut self, broker: ToolBroker) -> Self {
176        self.broker = broker;
177        self
178    }
179
180    pub fn with_usage_signals(
181        mut self,
182        usage_signals: Box<dyn ToolUsageSignalProvider>,
183        max_adjustment: f32,
184    ) -> Self {
185        self.broker = self
186            .broker
187            .with_usage_signals(usage_signals, max_adjustment);
188        self
189    }
190
191    pub fn with_lifecycle_hooks(mut self, hooks: Arc<dyn RuntimeLifecycleHooks<C>>) -> Self {
192        self.lifecycle_hooks = hooks;
193        self
194    }
195
196    pub fn with_policy_contract(mut self, policy_contract: PolicyContract) -> Self {
197        self.policy_contract = Some(policy_contract);
198        self
199    }
200
201    pub fn with_verified_policy_contract(
202        mut self,
203        policy_contract: PolicyContract,
204        verifier: impl FnOnce(&PolicyContract) -> Result<(), NexaraError>,
205    ) -> Result<Self, NexaraError> {
206        verifier(&policy_contract)?;
207        self.policy_contract = Some(policy_contract);
208        Ok(self)
209    }
210
211    pub fn registry(&self) -> ToolRegistry {
212        let mut registry = ToolRegistry::new();
213        for tool in self.catalog.tools() {
214            registry.register(tool);
215        }
216        registry
217    }
218
219    pub fn list_tools(&self, prompt: Option<&str>, scopes: &[String]) -> Vec<ToolDescriptor> {
220        let registry = self.registry();
221        if self.config.tool_selection_enabled && prompt.is_some() {
222            self.broker.select(
223                &registry,
224                &ToolSelectionRequest {
225                    prompt: prompt.unwrap_or_default().to_string(),
226                    scopes: scopes.to_vec(),
227                    max_tools: self.config.max_tools_per_request,
228                },
229            )
230        } else {
231            registry
232                .all()
233                .into_iter()
234                .filter(|tool| tool.enabled)
235                .take(self.config.max_tools_per_request)
236                .collect()
237        }
238    }
239
240    pub async fn call_tool(
241        &self,
242        mut request: ToolCallRequest,
243        context: C,
244    ) -> Result<ToolCallResult, NexaraError> {
245        let started = Instant::now();
246        let payload_size = serde_json::to_vec(&request.params)
247            .map(|payload| payload.len())
248            .unwrap_or(usize::MAX);
249        if payload_size > self.config.max_payload_bytes {
250            return Err(NexaraError::PayloadTooLarge(payload_size));
251        }
252        let registry = self.registry();
253        let tool = registry
254            .get(&request.tool)
255            .cloned()
256            .ok_or(NexaraError::ToolNotFound)?;
257        let policy = self.trust_resolver.resolve(&context);
258        let policy_evaluation = self.evaluate_policy_contract(&tool, &request);
259        if let Some(evaluation) = &policy_evaluation
260            && let Some(error) = policy_error_for_evaluation(evaluation)
261        {
262            let mut record = AuditRecord::now(
263                "unknown",
264                tool.name.clone(),
265                policy_decision_for_evaluation(evaluation),
266                AuditOutcome::Rejected,
267            );
268            record.duration_ms = started.elapsed().as_millis() as u64;
269            record
270                .metadata
271                .extend(matched_rules_metadata(&evaluation.matched_rules));
272            record.metadata.insert(
273                "policy.explanation".to_string(),
274                evaluation.explanation.clone(),
275            );
276            let _ = self.audit_sink.record(record);
277            return Err(error);
278        }
279        self.enforce_policy(&tool, &mut request, &policy)?;
280        let permit = self
281            .call_semaphore
282            .try_acquire()
283            .map_err(|_| NexaraError::ConcurrencyLimitExceeded)?;
284        self.lifecycle_hooks
285            .before_call(&tool, &request, &context)
286            .await;
287        let result = self
288            .executor
289            .call_host_tool(&tool, request.clone(), context.clone())
290            .await;
291        drop(permit);
292
293        let (decision, outcome, hash) = match &result {
294            Ok(result) => (
295                PolicyDecision::Allowed,
296                AuditOutcome::Success,
297                Some(hash_result(&result.result)),
298            ),
299            Err(_) => (PolicyDecision::Allowed, AuditOutcome::Error, None),
300        };
301        let mut record = AuditRecord::now("unknown", tool.name.clone(), decision, outcome);
302        record.duration_ms = started.elapsed().as_millis() as u64;
303        record.result_hash = hash;
304        if let Some(policy_evaluation) = policy_evaluation {
305            record
306                .metadata
307                .extend(matched_rules_metadata(&policy_evaluation.matched_rules));
308            record.metadata.insert(
309                "policy.explanation".to_string(),
310                policy_evaluation.explanation,
311            );
312        }
313        let _ = self.audit_sink.record(record);
314        self.lifecycle_hooks
315            .after_call(&tool, &request, result.as_ref(), &context)
316            .await;
317        result
318    }
319
320    fn enforce_policy(
321        &self,
322        tool: &ToolDescriptor,
323        request: &mut ToolCallRequest,
324        policy: &EffectiveTrustPolicy,
325    ) -> Result<(), NexaraError> {
326        let confirmed = request
327            .params
328            .get("_confirmed")
329            .and_then(|value| value.as_bool())
330            .unwrap_or(false);
331        if let Some(object) = request.params.as_object_mut() {
332            object.remove("_confirmed");
333        }
334        match tool.action_class {
335            ActionClass::Read if policy.allow_read => Ok(()),
336            ActionClass::Read => Err(NexaraError::TrustPolicyDenied(
337                "read tools are disabled by trust policy".to_string(),
338            )),
339            ActionClass::Write if !policy.allow_write => Err(NexaraError::TrustPolicyDenied(
340                "write tools are disabled by trust policy".to_string(),
341            )),
342            ActionClass::Write if policy.require_confirmation_for_write && !confirmed => Err(
343                NexaraError::ConfirmationRequired("write tool requires confirmation".to_string()),
344            ),
345            ActionClass::Write => Ok(()),
346            ActionClass::Execute if !policy.allow_execute => Err(NexaraError::TrustPolicyDenied(
347                "execute tools are disabled by trust policy".to_string(),
348            )),
349            ActionClass::Execute if policy.require_confirmation_for_execute && !confirmed => Err(
350                NexaraError::ConfirmationRequired("execute tool requires confirmation".to_string()),
351            ),
352            ActionClass::Execute => Ok(()),
353        }
354    }
355
356    fn evaluate_policy_contract(
357        &self,
358        tool: &ToolDescriptor,
359        request: &ToolCallRequest,
360    ) -> Option<PolicyEvaluation> {
361        let Some(contract) = &self.policy_contract else {
362            return None;
363        };
364        let confirmed = request
365            .params
366            .get("_confirmed")
367            .and_then(|value| value.as_bool())
368            .unwrap_or(false);
369        Some(contract.evaluate_tool(tool, confirmed))
370    }
371}
372
373fn policy_error_for_evaluation(evaluation: &PolicyEvaluation) -> Option<NexaraError> {
374    match evaluation.decision {
375        PolicyEvaluationDecision::Allowed => None,
376        PolicyEvaluationDecision::ConfirmationRequired => Some(NexaraError::ConfirmationRequired(
377            evaluation.explanation.clone(),
378        )),
379        PolicyEvaluationDecision::Denied
380        | PolicyEvaluationDecision::NoMatchingAllow
381        | PolicyEvaluationDecision::DescriptorMissingCapabilities => Some(
382            NexaraError::TrustPolicyDenied(evaluation.explanation.clone()),
383        ),
384    }
385}
386
387fn policy_decision_for_evaluation(evaluation: &PolicyEvaluation) -> PolicyDecision {
388    match evaluation.decision {
389        PolicyEvaluationDecision::Allowed => PolicyDecision::Allowed,
390        PolicyEvaluationDecision::Denied => PolicyDecision::DeniedPolicyContract,
391        PolicyEvaluationDecision::NoMatchingAllow => PolicyDecision::DeniedNoMatchingAllow,
392        PolicyEvaluationDecision::DescriptorMissingCapabilities => {
393            PolicyDecision::DeniedMissingCapabilities
394        }
395        PolicyEvaluationDecision::ConfirmationRequired => {
396            PolicyDecision::DeniedPolicyConfirmationRequired
397        }
398    }
399}