Skip to main content

nexara_runtime/
lib.rs

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