rill_runtime/
server.rs

1use std::sync::Arc;
2
3use rill_runtime_protocol::{
4    MIN_RUNTIME_API_VERSION, RUNTIME_API_VERSION, RuntimeRequest, RuntimeResponse,
5    RuntimeResponseV2,
6};
7use serde_json::Value;
8
9use crate::handler::HandlerIdentity;
10use crate::package::LoadedModelPack;
11
12/// Consumers can implement this trait to add business-specific invocation logic.
13pub trait InvokeHandler: Send + Sync + std::fmt::Debug {
14    fn invoke(&self, capability: &str, input: &Value) -> Result<Value, String>;
15}
16
17/// Internal response type produced by [`RuntimeEngine`]. The IPC layer converts
18/// this to a v1 [`RuntimeResponse`] or v2 [`RuntimeResponseV2`] based on the
19/// request's `api_version`.
20#[derive(Debug, Clone)]
21pub enum EngineResponse {
22    Handshake {
23        request_id: String,
24        runtime_version: String,
25        model_pack_id: String,
26        model_pack_version: String,
27        capabilities: Vec<String>,
28        handler: Option<HandlerIdentity>,
29    },
30    Health {
31        request_id: String,
32        healthy: bool,
33        model_pack_id: String,
34        model_pack_version: String,
35    },
36    Result {
37        request_id: String,
38        output: Value,
39    },
40    Error {
41        request_id: String,
42        code: String,
43        message: String,
44        retryable: bool,
45    },
46}
47
48impl EngineResponse {
49    /// Convert to a v1 wire response. Handler identity fields are dropped.
50    pub fn to_v1(&self, api_version: u32) -> RuntimeResponse {
51        match self {
52            Self::Handshake {
53                request_id,
54                runtime_version,
55                model_pack_id,
56                model_pack_version,
57                capabilities,
58                ..
59            } => RuntimeResponse::Handshake {
60                request_id: request_id.clone(),
61                api_version,
62                runtime_version: runtime_version.clone(),
63                model_pack_id: model_pack_id.clone(),
64                model_pack_version: model_pack_version.clone(),
65                capabilities: capabilities.clone(),
66            },
67            Self::Health {
68                request_id,
69                healthy,
70                model_pack_id,
71                model_pack_version,
72            } => RuntimeResponse::Health {
73                request_id: request_id.clone(),
74                api_version,
75                healthy: *healthy,
76                model_pack_id: model_pack_id.clone(),
77                model_pack_version: model_pack_version.clone(),
78            },
79            Self::Result { request_id, output } => RuntimeResponse::Result {
80                request_id: request_id.clone(),
81                api_version,
82                output: output.clone(),
83            },
84            Self::Error {
85                request_id,
86                code,
87                message,
88                retryable,
89            } => RuntimeResponse::Error {
90                request_id: request_id.clone(),
91                api_version,
92                code: code.clone(),
93                message: message.clone(),
94                retryable: *retryable,
95            },
96        }
97    }
98
99    /// Convert to a v2 wire response. If no handler is loaded, handler fields
100    /// are filled with empty/zero values and effective_capabilities equals the
101    /// model capabilities.
102    pub fn to_v2(&self, api_version: u32) -> RuntimeResponseV2 {
103        match self {
104            Self::Handshake {
105                request_id,
106                runtime_version,
107                model_pack_id,
108                model_pack_version,
109                capabilities,
110                handler,
111            } => {
112                let (handler_id, handler_version, handler_api_version, effective) = match handler {
113                    Some(h) => (
114                        h.handler_id.clone(),
115                        h.handler_version.clone(),
116                        h.handler_api_version,
117                        h.effective_capabilities.clone(),
118                    ),
119                    None => (String::new(), String::new(), 0, capabilities.clone()),
120                };
121                RuntimeResponseV2::Handshake {
122                    request_id: request_id.clone(),
123                    api_version,
124                    runtime_version: runtime_version.clone(),
125                    model_pack_id: model_pack_id.clone(),
126                    model_pack_version: model_pack_version.clone(),
127                    capabilities: capabilities.clone(),
128                    handler_id,
129                    handler_version,
130                    handler_api_version,
131                    effective_capabilities: effective,
132                }
133            }
134            Self::Health {
135                request_id,
136                healthy,
137                model_pack_id,
138                model_pack_version,
139            } => RuntimeResponseV2::Health {
140                request_id: request_id.clone(),
141                api_version,
142                healthy: *healthy,
143                model_pack_id: model_pack_id.clone(),
144                model_pack_version: model_pack_version.clone(),
145            },
146            Self::Result { request_id, output } => RuntimeResponseV2::Result {
147                request_id: request_id.clone(),
148                api_version,
149                output: output.clone(),
150            },
151            Self::Error {
152                request_id,
153                code,
154                message,
155                retryable,
156            } => RuntimeResponseV2::Error {
157                request_id: request_id.clone(),
158                api_version,
159                code: code.clone(),
160                message: message.clone(),
161                retryable: *retryable,
162            },
163        }
164    }
165}
166
167#[derive(Debug, Clone)]
168pub struct RuntimeEngine {
169    pack: LoadedModelPack,
170    invoke_handler: Option<Arc<dyn InvokeHandler>>,
171    handler_identity: Option<HandlerIdentity>,
172    effective_capabilities: Vec<String>,
173}
174
175impl RuntimeEngine {
176    pub fn new(pack: LoadedModelPack) -> Self {
177        Self {
178            pack,
179            invoke_handler: None,
180            handler_identity: None,
181            effective_capabilities: Vec::new(),
182        }
183    }
184
185    pub fn with_invoke_handler(mut self, handler: Arc<dyn InvokeHandler>) -> Self {
186        self.invoke_handler = Some(handler);
187        self
188    }
189
190    /// Attach handler identity and effective capabilities for IPC v2 handshake.
191    pub fn with_handler_identity(mut self, identity: HandlerIdentity) -> Self {
192        self.effective_capabilities = identity.effective_capabilities.clone();
193        self.handler_identity = Some(identity);
194        self
195    }
196
197    /// Effective capability set (intersection of model and handler). Empty when
198    /// no handler is loaded.
199    pub fn effective_capabilities(&self) -> &[String] {
200        &self.effective_capabilities
201    }
202
203    /// Handler identity if a handler was loaded.
204    pub fn handler_identity(&self) -> Option<&HandlerIdentity> {
205        self.handler_identity.as_ref()
206    }
207
208    pub fn handle(&self, request: RuntimeRequest) -> EngineResponse {
209        let request_id = request.request_id().to_string();
210        if request_id.is_empty() || request_id.len() > 128 {
211            return self.error(request_id, "invalidRequestId", "invalid request id", false);
212        }
213        let api_version = request.api_version();
214        if !(MIN_RUNTIME_API_VERSION..=RUNTIME_API_VERSION).contains(&api_version) {
215            return self.error(
216                request_id,
217                "incompatibleApiVersion",
218                "runtime API version is not supported",
219                false,
220            );
221        }
222
223        match request {
224            RuntimeRequest::Handshake {
225                request_id,
226                client_name,
227                client_version,
228                ..
229            } => {
230                if client_name.is_empty()
231                    || client_name.len() > 96
232                    || client_version.is_empty()
233                    || client_version.len() > 48
234                {
235                    return self.error(
236                        request_id,
237                        "invalidClientIdentity",
238                        "invalid client identity",
239                        false,
240                    );
241                }
242                EngineResponse::Handshake {
243                    request_id,
244                    runtime_version: env!("CARGO_PKG_VERSION").into(),
245                    model_pack_id: self.pack.manifest.id.clone(),
246                    model_pack_version: self.pack.manifest.version.clone(),
247                    capabilities: self.pack.manifest.capabilities.clone(),
248                    handler: self.handler_identity.clone(),
249                }
250            }
251            RuntimeRequest::Health { request_id, .. } => EngineResponse::Health {
252                request_id,
253                healthy: true,
254                model_pack_id: self.pack.manifest.id.clone(),
255                model_pack_version: self.pack.manifest.version.clone(),
256            },
257            RuntimeRequest::Invoke {
258                request_id,
259                capability,
260                input,
261                ..
262            } => {
263                if !self.is_capability_allowed(&capability) {
264                    return self.error(
265                        request_id,
266                        "unsupportedCapability",
267                        "capability is not in the effective set",
268                        false,
269                    );
270                }
271                let Some(handler) = &self.invoke_handler else {
272                    return self.error(
273                        request_id,
274                        "noInvokeHandler",
275                        "no invoke handler registered",
276                        false,
277                    );
278                };
279                match handler.invoke(&capability, &input) {
280                    Ok(output) => EngineResponse::Result { request_id, output },
281                    Err(message) => {
282                        let (code, retryable) = map_invoke_error(&message);
283                        self.error(request_id, code, &message, retryable)
284                    }
285                }
286            }
287        }
288    }
289
290    /// Checks the capability against the effective set when a handler is loaded,
291    /// or against the model pack's declared capabilities when no handler is
292    /// loaded (for backwards compatibility with built-in handlers selected by
293    /// the binary).
294    fn is_capability_allowed(&self, capability: &str) -> bool {
295        if !self.effective_capabilities.is_empty() {
296            self.effective_capabilities.iter().any(|c| c == capability)
297        } else {
298            self.pack
299                .manifest
300                .capabilities
301                .iter()
302                .any(|c| c == capability)
303        }
304    }
305
306    fn error(
307        &self,
308        request_id: String,
309        code: &str,
310        message: &str,
311        retryable: bool,
312    ) -> EngineResponse {
313        EngineResponse::Error {
314            request_id,
315            code: code.into(),
316            message: message.into(),
317            retryable,
318        }
319    }
320}
321
322/// Maps a handler error message to a stable error code. Recognised codes are
323/// extracted from the message prefix; unknown errors map to `handlerInternalError`.
324fn map_invoke_error(message: &str) -> (&'static str, bool) {
325    if message.starts_with("handlerTrap") {
326        ("handlerTrap", false)
327    } else if message.starts_with("handlerTimeout") {
328        ("handlerTimeout", true)
329    } else if message.starts_with("handlerOutputTooLarge") {
330        ("handlerOutputTooLarge", false)
331    } else if message.starts_with("handlerInvalidOutput") {
332        ("handlerInvalidOutput", false)
333    } else if message.starts_with("handlerInternalError") {
334        ("handlerInternalError", false)
335    } else if message.starts_with("handlerExecutionFailed") {
336        // Handler returned an error via the WIT result type.
337        ("handlerInternalError", false)
338    } else {
339        ("handlerInternalError", false)
340    }
341}
342
343#[cfg(test)]
344mod tests {
345    use rill_runtime_protocol::{MODEL_PACK_FORMAT_VERSION, ModelPackManifest};
346
347    use super::*;
348    use crate::handler::builtin::LINEAR_REGRESSION_CAPABILITY;
349
350    fn engine() -> RuntimeEngine {
351        RuntimeEngine::new(LoadedModelPack {
352            manifest: ModelPackManifest {
353                format_version: MODEL_PACK_FORMAT_VERSION,
354                id: "rillml.example.default".into(),
355                version: "0.7.0".into(),
356                runtime_api_version: RUNTIME_API_VERSION,
357                min_runtime_version: "0.7.0".into(),
358                publisher_key_id: "test".into(),
359                capabilities: vec!["rillml.example".into()],
360            },
361            model: serde_json::json!({}),
362        })
363    }
364
365    #[test]
366    fn handshake_reports_loaded_pack() {
367        let response = engine().handle(RuntimeRequest::Handshake {
368            request_id: "hello".into(),
369            api_version: RUNTIME_API_VERSION,
370            client_name: "example-host".into(),
371            client_version: "0.9.0".into(),
372        });
373        assert!(matches!(
374            response,
375            EngineResponse::Handshake { model_pack_id, .. }
376                if model_pack_id == "rillml.example.default"
377        ));
378    }
379
380    #[test]
381    fn incompatible_api_is_a_typed_error() {
382        let response = engine().handle(RuntimeRequest::Health {
383            request_id: "health".into(),
384            api_version: RUNTIME_API_VERSION + 1,
385        });
386        assert!(matches!(
387            response,
388            EngineResponse::Error { code, .. } if code == "incompatibleApiVersion"
389        ));
390    }
391
392    #[test]
393    fn invoke_without_handler_returns_no_invoke_handler_error() {
394        let response = engine().handle(RuntimeRequest::Invoke {
395            request_id: "invoke-1".into(),
396            api_version: RUNTIME_API_VERSION,
397            capability: "rillml.example".into(),
398            input: serde_json::json!({}),
399        });
400        assert!(matches!(
401            response,
402            EngineResponse::Error { code, .. } if code == "noInvokeHandler"
403        ));
404    }
405
406    #[test]
407    fn invoke_rejects_capability_not_declared_by_signed_manifest() {
408        let response = engine().handle(RuntimeRequest::Invoke {
409            request_id: "invoke-undeclared".into(),
410            api_version: RUNTIME_API_VERSION,
411            capability: "undeclared.capability".into(),
412            input: serde_json::json!({}),
413        });
414        assert!(matches!(
415            response,
416            EngineResponse::Error { code, .. } if code == "unsupportedCapability"
417        ));
418    }
419
420    #[test]
421    fn v1_handshake_omits_handler_fields() {
422        let identity = HandlerIdentity {
423            handler_id: "org.example.handler".into(),
424            handler_version: "1.0.0".into(),
425            handler_api_version: 1,
426            effective_capabilities: vec!["rillml.example".into()],
427        };
428        let engine = engine().with_handler_identity(identity);
429        let response = engine.handle(RuntimeRequest::Handshake {
430            request_id: "v1-test".into(),
431            api_version: 1,
432            client_name: "v1-host".into(),
433            client_version: "0.6.0".into(),
434        });
435        let v1 = response.to_v1(1);
436        let json = serde_json::to_string(&v1).unwrap();
437        assert!(!json.contains("handlerId"));
438        assert!(!json.contains("effectiveCapabilities"));
439    }
440
441    #[test]
442    fn v2_handshake_includes_handler_fields() {
443        let identity = HandlerIdentity {
444            handler_id: "org.example.handler".into(),
445            handler_version: "1.0.0".into(),
446            handler_api_version: 1,
447            effective_capabilities: vec!["rillml.example".into()],
448        };
449        let engine = engine().with_handler_identity(identity);
450        let response = engine.handle(RuntimeRequest::Handshake {
451            request_id: "v2-test".into(),
452            api_version: 2,
453            client_name: "v2-host".into(),
454            client_version: "0.7.0".into(),
455        });
456        let v2 = response.to_v2(2);
457        let json = serde_json::to_string(&v2).unwrap();
458        assert!(json.contains("\"handlerId\":\"org.example.handler\""));
459        assert!(json.contains("\"handlerApiVersion\":1"));
460        assert!(json.contains("\"effectiveCapabilities\":[\"rillml.example\"]"));
461    }
462
463    #[test]
464    fn v2_handshake_without_handler_has_empty_fields() {
465        let response = engine().handle(RuntimeRequest::Handshake {
466            request_id: "v2-no-handler".into(),
467            api_version: 2,
468            client_name: "v2-host".into(),
469            client_version: "0.7.0".into(),
470        });
471        let v2 = response.to_v2(2);
472        match v2 {
473            RuntimeResponseV2::Handshake {
474                handler_id,
475                handler_version,
476                handler_api_version,
477                effective_capabilities,
478                ..
479            } => {
480                assert!(handler_id.is_empty());
481                assert!(handler_version.is_empty());
482                assert_eq!(handler_api_version, 0);
483                assert_eq!(effective_capabilities, vec!["rillml.example"]);
484            }
485            _ => panic!("expected handshake"),
486        }
487    }
488
489    #[test]
490    fn linear_regression_handler_validates_and_predicts() {
491        use crate::handler::builtin::LinearRegressionInvokeHandler;
492
493        let pack = LoadedModelPack {
494            manifest: ModelPackManifest {
495                format_version: MODEL_PACK_FORMAT_VERSION,
496                id: "rillml.example.default".into(),
497                version: "0.7.0".into(),
498                runtime_api_version: RUNTIME_API_VERSION,
499                min_runtime_version: "0.7.0".into(),
500                publisher_key_id: "test".into(),
501                capabilities: vec![LINEAR_REGRESSION_CAPABILITY.into()],
502            },
503            model: serde_json::json!({
504                "kind": "linearRegression",
505                "weights": [0.5, -0.25],
506                "intercept": 1.0
507            }),
508        };
509        let handler = LinearRegressionInvokeHandler::from_pack(&pack).unwrap();
510        let engine = RuntimeEngine::new(pack).with_invoke_handler(Arc::new(handler));
511        let response = engine.handle(RuntimeRequest::Invoke {
512            request_id: "invoke-linear".into(),
513            api_version: RUNTIME_API_VERSION,
514            capability: LINEAR_REGRESSION_CAPABILITY.into(),
515            input: serde_json::json!({"features": [4.0, 2.0]}),
516        });
517        assert!(matches!(
518            response,
519            EngineResponse::Result { output, .. } if output["prediction"] == 2.5
520        ));
521    }
522
523    #[test]
524    fn map_invoke_error_recognizes_handler_trap() {
525        let (code, retryable) = map_invoke_error("handlerTrap: unreachable");
526        assert_eq!(code, "handlerTrap");
527        assert!(!retryable);
528    }
529
530    #[test]
531    fn map_invoke_error_recognizes_handler_timeout() {
532        let (code, retryable) = map_invoke_error("handlerTimeout: epoch deadline exceeded");
533        assert_eq!(code, "handlerTimeout");
534        assert!(retryable);
535    }
536
537    #[test]
538    fn map_invoke_error_recognizes_handler_output_too_large() {
539        let (code, retryable) = map_invoke_error("handlerOutputTooLarge: output exceeds 1 MiB");
540        assert_eq!(code, "handlerOutputTooLarge");
541        assert!(!retryable);
542    }
543
544    #[test]
545    fn map_invoke_error_recognizes_handler_invalid_output() {
546        let (code, retryable) =
547            map_invoke_error("handlerInvalidOutput: expected value at line 1 column 1");
548        assert_eq!(code, "handlerInvalidOutput");
549        assert!(!retryable);
550    }
551
552    #[test]
553    fn map_invoke_error_recognizes_handler_internal_error() {
554        let (code, retryable) = map_invoke_error("handlerInternalError: lock poisoned");
555        assert_eq!(code, "handlerInternalError");
556        assert!(!retryable);
557    }
558
559    #[test]
560    fn map_invoke_error_recognizes_handler_execution_failed() {
561        let (code, retryable) = map_invoke_error("handlerExecutionFailed: guest returned error");
562        assert_eq!(code, "handlerInternalError");
563        assert!(!retryable);
564    }
565
566    #[test]
567    fn map_invoke_error_maps_capability_mismatch_to_internal() {
568        // handlerCapabilityMismatch is a load-phase error (RFC §6.2) and is
569        // never produced by invoke(); it falls through to handlerInternalError.
570        let (code, retryable) = map_invoke_error("handlerCapabilityMismatch: cap not declared");
571        assert_eq!(code, "handlerInternalError");
572        assert!(!retryable);
573    }
574
575    #[test]
576    fn map_invoke_error_maps_unknown_to_internal_error() {
577        let (code, retryable) = map_invoke_error("unknown error");
578        assert_eq!(code, "handlerInternalError");
579        assert!(!retryable);
580    }
581}
rill_runtime/server.rs

rill_runtime/
server.rs
