1use std::collections::BTreeSet;
9use std::sync::Arc;
10
11use harn_vm::{Vm, VmError, VmValue};
12
13use crate::error::HostlibError;
14
15pub type SyncHandler = Arc<dyn Fn(&[VmValue]) -> Result<VmValue, HostlibError> + Send + Sync>;
19
20#[derive(Clone)]
24pub struct RegisteredBuiltin {
25 pub name: &'static str,
27 pub module: &'static str,
29 pub method: &'static str,
31 pub handler: SyncHandler,
33}
34
35impl std::fmt::Debug for RegisteredBuiltin {
36 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
37 f.debug_struct("RegisteredBuiltin")
38 .field("name", &self.name)
39 .field("module", &self.module)
40 .field("method", &self.method)
41 .finish()
42 }
43}
44
45#[derive(Default)]
47pub struct BuiltinRegistry {
48 builtins: Vec<RegisteredBuiltin>,
49 command_policy_builtins: BTreeSet<&'static str>,
50}
51
52impl BuiltinRegistry {
53 pub fn new() -> Self {
55 Self::default()
56 }
57
58 pub fn register(&mut self, builtin: RegisteredBuiltin) {
60 self.builtins.push(builtin);
61 }
62
63 pub fn register_unimplemented(
66 &mut self,
67 name: &'static str,
68 module: &'static str,
69 method: &'static str,
70 ) {
71 let handler: SyncHandler =
72 Arc::new(move |_args| Err(HostlibError::Unimplemented { builtin: name }));
73 self.register(RegisteredBuiltin {
74 name,
75 module,
76 method,
77 handler,
78 });
79 }
80
81 pub(crate) fn register_fn(
85 &mut self,
86 module: &'static str,
87 name: &'static str,
88 method: &'static str,
89 runner: fn(&[VmValue]) -> Result<VmValue, HostlibError>,
90 ) {
91 let handler: SyncHandler = Arc::new(runner);
92 self.register(RegisteredBuiltin {
93 name,
94 module,
95 method,
96 handler,
97 });
98 }
99
100 pub(crate) fn register_gated_fn(
104 &mut self,
105 module: &'static str,
106 name: &'static str,
107 method: &'static str,
108 runner: fn(&[VmValue]) -> Result<VmValue, HostlibError>,
109 ) {
110 self.register(RegisteredBuiltin {
111 name,
112 module,
113 method,
114 handler: crate::tools::permissions::gated_handler(name, runner),
115 });
116 }
117
118 pub(crate) fn register_gated_command_fn(
121 &mut self,
122 module: &'static str,
123 name: &'static str,
124 method: &'static str,
125 runner: fn(&[VmValue]) -> Result<VmValue, HostlibError>,
126 ) {
127 self.register_gated_fn(module, name, method, runner);
128 self.command_policy_builtins.insert(name);
129 }
130
131 fn uses_command_policy(&self, name: &str) -> bool {
132 self.command_policy_builtins.contains(name)
133 }
134
135 pub fn iter(&self) -> impl Iterator<Item = &RegisteredBuiltin> {
137 self.builtins.iter()
138 }
139
140 pub fn len(&self) -> usize {
142 self.builtins.len()
143 }
144
145 pub fn is_empty(&self) -> bool {
147 self.builtins.is_empty()
148 }
149
150 pub fn find(&self, name: &str) -> Option<&RegisteredBuiltin> {
152 self.builtins.iter().find(|b| b.name == name)
153 }
154}
155
156pub trait HostlibCapability: 'static {
160 fn module_name(&self) -> &'static str;
162
163 fn register_builtins(&self, registry: &mut BuiltinRegistry);
165}
166
167pub struct HostlibRegistry {
173 builtins: BuiltinRegistry,
174 modules: Vec<&'static str>,
175}
176
177impl Default for HostlibRegistry {
178 fn default() -> Self {
179 Self::new()
180 }
181}
182
183impl HostlibRegistry {
184 pub fn new() -> Self {
187 Self {
188 builtins: BuiltinRegistry::new(),
189 modules: Vec::new(),
190 }
191 }
192
193 #[must_use]
195 pub fn with<C: HostlibCapability>(mut self, capability: C) -> Self {
196 let module = capability.module_name();
197 capability.register_builtins(&mut self.builtins);
198 self.modules.push(module);
199 self
200 }
201
202 pub fn register_into_vm(&mut self, vm: &mut Vm) {
204 for builtin in self.builtins.iter().cloned() {
205 let module = builtin.module;
206 let method = builtin.method;
207 harn_vm::stdlib::host::register_callable_host_operation(
208 module,
209 method,
210 "Hostlib schema-backed operation registered at runtime.",
211 );
212 let handler = builtin.handler.clone();
213 if self.builtins.uses_command_policy(builtin.name) {
214 vm.register_async_builtin(builtin.name, move |ctx, args| {
215 let handler = handler.clone();
216 async move {
217 let request = crate::schemas::validate_request_args(
218 builtin.name,
219 module,
220 method,
221 &args,
222 )
223 .map_err(VmError::from)?;
224 let params = request.as_dict().ok_or_else(|| {
225 VmError::Runtime(format!(
226 "{}: validated request must be a dict",
227 builtin.name
228 ))
229 })?;
230 if let Some(mocked) = harn_vm::stdlib::host::dispatch_mock_hostlib_call(
231 module, method, params,
232 ) {
233 return mocked;
234 }
235 let caller = serde_json::json!({
236 "surface": "hostlib",
237 "builtin": builtin.name,
238 "module": module,
239 "method": method,
240 "session_id": harn_vm::current_agent_session_id(),
241 });
242 match harn_vm::orchestration::run_command_policy_preflight_with_ctx(
243 Some(&ctx),
244 params,
245 caller,
246 )
247 .await?
248 {
249 harn_vm::orchestration::CommandPolicyPreflight::Blocked {
250 status,
251 message,
252 context,
253 decisions,
254 } => {
255 let response = harn_vm::orchestration::blocked_command_response(
256 params, status, &message, context, decisions,
257 );
258 crate::schemas::validate_response(
259 builtin.name,
260 module,
261 method,
262 crate::tools::policy_blocked_run_command_response(response),
263 )
264 .map_err(VmError::from)
265 }
266 harn_vm::orchestration::CommandPolicyPreflight::Proceed {
267 params,
268 context,
269 decisions,
270 } => {
271 let rewritten = VmValue::dict(params.clone());
275 let validated = crate::schemas::validate_request_args(
276 builtin.name,
277 module,
278 method,
279 &[rewritten],
280 )
281 .map_err(VmError::from)?;
282 let result = handler(&[validated]).map_err(VmError::from)?;
283 if crate::tools::run_command_request_is_background(¶ms) {
284 return crate::schemas::validate_response(
285 builtin.name,
286 module,
287 method,
288 result,
289 )
290 .map_err(VmError::from);
291 }
292 let result =
293 harn_vm::orchestration::run_command_policy_postflight_with_ctx(
294 Some(&ctx),
295 ¶ms,
296 result,
297 context,
298 decisions,
299 )
300 .await?;
301 crate::schemas::validate_response(
302 builtin.name,
303 module,
304 method,
305 result,
306 )
307 .map_err(VmError::from)
308 }
309 }
310 }
311 });
312 } else {
313 vm.register_builtin(
314 builtin.name,
315 move |args, _out| -> Result<VmValue, VmError> {
316 let request = crate::schemas::validate_request_args(
317 builtin.name,
318 module,
319 method,
320 args,
321 )
322 .map_err(VmError::from)?;
323 let validated_args = [request.clone()];
324 if let Some(params) = request.as_dict() {
325 if let Some(mocked) = harn_vm::stdlib::host::dispatch_mock_hostlib_call(
326 module, method, params,
327 ) {
328 return mocked;
329 }
330 }
331 handler(&validated_args).map_err(VmError::from)
332 },
333 );
334 }
335 }
336 }
337
338 pub fn builtins(&self) -> &BuiltinRegistry {
341 &self.builtins
342 }
343
344 pub fn modules(&self) -> &[&'static str] {
346 &self.modules
347 }
348}
349
350#[cfg(test)]
351mod tests {
352 use super::*;
353
354 #[test]
355 fn unimplemented_builtins_route_through_error() {
356 let mut registry = BuiltinRegistry::new();
357 registry.register_unimplemented("hostlib_demo", "demo", "ping");
358 let entry = registry.find("hostlib_demo").expect("registered");
359 let err = (entry.handler)(&[]).expect_err("should be unimplemented");
360 assert!(
361 matches!(err, HostlibError::Unimplemented { builtin } if builtin == "hostlib_demo")
362 );
363 }
364
365 #[test]
366 fn registry_records_modules_in_order() {
367 struct First;
368 impl HostlibCapability for First {
369 fn module_name(&self) -> &'static str {
370 "first"
371 }
372 fn register_builtins(&self, _registry: &mut BuiltinRegistry) {}
373 }
374 struct Second;
375 impl HostlibCapability for Second {
376 fn module_name(&self) -> &'static str {
377 "second"
378 }
379 fn register_builtins(&self, _registry: &mut BuiltinRegistry) {}
380 }
381
382 let registry = HostlibRegistry::new().with(First).with(Second);
383 assert_eq!(registry.modules(), &["first", "second"]);
384 }
385}