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 `invokeFailed`.
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("handlerCapabilityMismatch") {
334        ("handlerCapabilityMismatch", false)
335    } else {
336        ("invokeFailed", false)
337    }
338}
339
340#[cfg(test)]
341mod tests {
342    use rill_runtime_protocol::{MODEL_PACK_FORMAT_VERSION, ModelPackManifest};
343
344    use super::*;
345    use crate::handler::builtin::LINEAR_REGRESSION_CAPABILITY;
346
347    fn engine() -> RuntimeEngine {
348        RuntimeEngine::new(LoadedModelPack {
349            manifest: ModelPackManifest {
350                format_version: MODEL_PACK_FORMAT_VERSION,
351                id: "rillml.example.default".into(),
352                version: "0.7.0".into(),
353                runtime_api_version: RUNTIME_API_VERSION,
354                min_runtime_version: "0.7.0".into(),
355                publisher_key_id: "test".into(),
356                capabilities: vec!["rillml.example".into()],
357            },
358            model: serde_json::json!({}),
359        })
360    }
361
362    #[test]
363    fn handshake_reports_loaded_pack() {
364        let response = engine().handle(RuntimeRequest::Handshake {
365            request_id: "hello".into(),
366            api_version: RUNTIME_API_VERSION,
367            client_name: "example-host".into(),
368            client_version: "0.9.0".into(),
369        });
370        assert!(matches!(
371            response,
372            EngineResponse::Handshake { model_pack_id, .. }
373                if model_pack_id == "rillml.example.default"
374        ));
375    }
376
377    #[test]
378    fn incompatible_api_is_a_typed_error() {
379        let response = engine().handle(RuntimeRequest::Health {
380            request_id: "health".into(),
381            api_version: RUNTIME_API_VERSION + 1,
382        });
383        assert!(matches!(
384            response,
385            EngineResponse::Error { code, .. } if code == "incompatibleApiVersion"
386        ));
387    }
388
389    #[test]
390    fn invoke_without_handler_returns_no_invoke_handler_error() {
391        let response = engine().handle(RuntimeRequest::Invoke {
392            request_id: "invoke-1".into(),
393            api_version: RUNTIME_API_VERSION,
394            capability: "rillml.example".into(),
395            input: serde_json::json!({}),
396        });
397        assert!(matches!(
398            response,
399            EngineResponse::Error { code, .. } if code == "noInvokeHandler"
400        ));
401    }
402
403    #[test]
404    fn invoke_rejects_capability_not_declared_by_signed_manifest() {
405        let response = engine().handle(RuntimeRequest::Invoke {
406            request_id: "invoke-undeclared".into(),
407            api_version: RUNTIME_API_VERSION,
408            capability: "undeclared.capability".into(),
409            input: serde_json::json!({}),
410        });
411        assert!(matches!(
412            response,
413            EngineResponse::Error { code, .. } if code == "unsupportedCapability"
414        ));
415    }
416
417    #[test]
418    fn v1_handshake_omits_handler_fields() {
419        let identity = HandlerIdentity {
420            handler_id: "org.example.handler".into(),
421            handler_version: "1.0.0".into(),
422            handler_api_version: 1,
423            effective_capabilities: vec!["rillml.example".into()],
424        };
425        let engine = engine().with_handler_identity(identity);
426        let response = engine.handle(RuntimeRequest::Handshake {
427            request_id: "v1-test".into(),
428            api_version: 1,
429            client_name: "v1-host".into(),
430            client_version: "0.6.0".into(),
431        });
432        let v1 = response.to_v1(1);
433        let json = serde_json::to_string(&v1).unwrap();
434        assert!(!json.contains("handlerId"));
435        assert!(!json.contains("effectiveCapabilities"));
436    }
437
438    #[test]
439    fn v2_handshake_includes_handler_fields() {
440        let identity = HandlerIdentity {
441            handler_id: "org.example.handler".into(),
442            handler_version: "1.0.0".into(),
443            handler_api_version: 1,
444            effective_capabilities: vec!["rillml.example".into()],
445        };
446        let engine = engine().with_handler_identity(identity);
447        let response = engine.handle(RuntimeRequest::Handshake {
448            request_id: "v2-test".into(),
449            api_version: 2,
450            client_name: "v2-host".into(),
451            client_version: "0.7.0".into(),
452        });
453        let v2 = response.to_v2(2);
454        let json = serde_json::to_string(&v2).unwrap();
455        assert!(json.contains("\"handlerId\":\"org.example.handler\""));
456        assert!(json.contains("\"handlerApiVersion\":1"));
457        assert!(json.contains("\"effectiveCapabilities\":[\"rillml.example\"]"));
458    }
459
460    #[test]
461    fn v2_handshake_without_handler_has_empty_fields() {
462        let response = engine().handle(RuntimeRequest::Handshake {
463            request_id: "v2-no-handler".into(),
464            api_version: 2,
465            client_name: "v2-host".into(),
466            client_version: "0.7.0".into(),
467        });
468        let v2 = response.to_v2(2);
469        match v2 {
470            RuntimeResponseV2::Handshake {
471                handler_id,
472                handler_version,
473                handler_api_version,
474                effective_capabilities,
475                ..
476            } => {
477                assert!(handler_id.is_empty());
478                assert!(handler_version.is_empty());
479                assert_eq!(handler_api_version, 0);
480                assert_eq!(effective_capabilities, vec!["rillml.example"]);
481            }
482            _ => panic!("expected handshake"),
483        }
484    }
485
486    #[test]
487    fn linear_regression_handler_validates_and_predicts() {
488        use crate::handler::builtin::LinearRegressionInvokeHandler;
489
490        let pack = LoadedModelPack {
491            manifest: ModelPackManifest {
492                format_version: MODEL_PACK_FORMAT_VERSION,
493                id: "rillml.example.default".into(),
494                version: "0.7.0".into(),
495                runtime_api_version: RUNTIME_API_VERSION,
496                min_runtime_version: "0.7.0".into(),
497                publisher_key_id: "test".into(),
498                capabilities: vec![LINEAR_REGRESSION_CAPABILITY.into()],
499            },
500            model: serde_json::json!({
501                "kind": "linearRegression",
502                "weights": [0.5, -0.25],
503                "intercept": 1.0
504            }),
505        };
506        let handler = LinearRegressionInvokeHandler::from_pack(&pack).unwrap();
507        let engine = RuntimeEngine::new(pack).with_invoke_handler(Arc::new(handler));
508        let response = engine.handle(RuntimeRequest::Invoke {
509            request_id: "invoke-linear".into(),
510            api_version: RUNTIME_API_VERSION,
511            capability: LINEAR_REGRESSION_CAPABILITY.into(),
512            input: serde_json::json!({"features": [4.0, 2.0]}),
513        });
514        assert!(matches!(
515            response,
516            EngineResponse::Result { output, .. } if output["prediction"] == 2.5
517        ));
518    }
519}
rill_runtime/server.rs

rill_runtime/
server.rs
