Skip to main content

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, error_code,
6};
7use serde_json::Value;
8
9use crate::handler::HandlerIdentity;
10use crate::package::LoadedModelPack;
11
12/// Typed invoke error.
13///
14/// Replaces the previous `Result<Value, String>` contract. The `kind`
15/// selects a stable IPC error code and a fixed public message; `detail`
16/// carries host-only diagnostic text (e.g. for stderr logs) and is **never**
17/// forwarded to IPC clients, so guests cannot exfiltrate arbitrary content
18/// through the error path.
19#[derive(Debug, Clone)]
20pub struct InvokeError {
21    kind: InvokeErrorKind,
22    detail: Option<String>,
23}
24
25/// Maximum byte length of the host-only `detail` string. Guests can fully
26/// control this payload via the WIT `handler-error` variant, so the host
27/// truncates it to bound memory and stderr noise. The limit is enforced
28/// on a UTF-8 char boundary so the stored string stays valid.
29pub const MAX_DETAIL_BYTES: usize = 4 * 1024;
30
31/// Stable categorisation of invoke failures.
32///
33/// The four guest-reported variants (`InvalidModel`, `InvalidInput`,
34/// `UnsupportedCapability`, `ExecutionFailed`) correspond 1:1 to the
35/// WIT `handler-error` variants. They share the same stable IPC code
36/// (`handlerInternalError`) for backwards compatibility with v1/v2
37/// clients, but carry distinct fixed public messages and are
38/// distinguishable host-side for logging and diagnostics.
39///
40/// Marked `#[non_exhaustive]` so future variants (e.g. for new WIT
41/// `handler-error` entries or host-side failure modes) can be added
42/// without breaking downstream exhaustive `match` arms. This preserves
43/// the patch-level version guarantee even though the enum is part of
44/// the crate's public API surface.
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46#[non_exhaustive]
47pub enum InvokeErrorKind {
48    /// Host-side input serialisation or size check failed.
49    Internal,
50    /// Fuel budget or epoch deadline was hit. Retryable.
51    Timeout,
52    /// Wasmtime trap (unreachable, OOB, stack overflow, …).
53    Trap,
54    /// Handler output exceeded [`MAX_IO_BYTES`](crate::MAX_IO_BYTES).
55    OutputTooLarge,
56    /// Handler output failed JSON deserialisation on the host side.
57    InvalidOutput,
58    /// Guest reported `invalid-model` via the WIT `handler-error`
59    /// variant. The variant detail is stored in [`InvokeError::detail`]
60    /// for host logs only.
61    InvalidModel,
62    /// Guest reported `invalid-input` via the WIT `handler-error`
63    /// variant. The variant detail is stored in [`InvokeError::detail`]
64    /// for host logs only.
65    InvalidInput,
66    /// Guest reported `unsupported-capability` via the WIT
67    /// `handler-error` variant. The variant detail is stored in
68    /// [`InvokeError::detail`] for host logs only.
69    UnsupportedCapability,
70    /// Guest reported `execution-failed` via the WIT `handler-error`
71    /// variant. The variant detail is stored in [`InvokeError::detail`]
72    /// for host logs only.
73    ExecutionFailed,
74}
75
76impl InvokeError {
77    /// Create a new typed error with no host detail.
78    pub const fn new(kind: InvokeErrorKind) -> Self {
79        Self { kind, detail: None }
80    }
81
82    /// Create a new typed error carrying host-only diagnostic text.
83    ///
84    /// `detail` is intended for `eprintln!` logs and **must not** be sent
85    /// to IPC clients. Guests can fully control this string via the WIT
86    /// `handler-error` payload, so it cannot be trusted for security
87    /// decisions. It is truncated to [`MAX_DETAIL_BYTES`] on a UTF-8 char
88    /// boundary so a malicious guest cannot grow host memory unboundedly
89    /// through the error path.
90    pub fn with_detail(kind: InvokeErrorKind, detail: impl Into<String>) -> Self {
91        Self {
92            kind,
93            detail: Some(truncate_to_bytes(detail.into(), MAX_DETAIL_BYTES)),
94        }
95    }
96
97    /// Error category.
98    pub const fn kind(&self) -> InvokeErrorKind {
99        self.kind
100    }
101
102    /// Host-only diagnostic text. Never sent to IPC clients.
103    pub fn detail(&self) -> Option<&str> {
104        self.detail.as_deref()
105    }
106
107    /// Stable IPC error code. Backwards-compatible with the v1/v2 wire
108    /// format produced by the previous `map_invoke_error` string matching.
109    ///
110    /// All four guest-reported WIT `handler-error` variants
111    /// (`invalid-model`, `invalid-input`, `unsupported-capability`,
112    /// `execution-failed`) collapse to `handlerInternalError` on the wire
113    /// to preserve compatibility with v1/v2 clients. The host still
114    /// distinguishes them internally via [`InvokeError::kind`] for
115    /// logging and diagnostics.
116    pub const fn stable_code(&self) -> &'static str {
117        match self.kind {
118            InvokeErrorKind::Internal => error_code::HANDLER_INTERNAL_ERROR,
119            InvokeErrorKind::Timeout => error_code::HANDLER_TIMEOUT,
120            InvokeErrorKind::Trap => error_code::HANDLER_TRAP,
121            InvokeErrorKind::OutputTooLarge => error_code::HANDLER_OUTPUT_TOO_LARGE,
122            InvokeErrorKind::InvalidOutput => error_code::HANDLER_INVALID_OUTPUT,
123            // Guest-reported WIT `handler-error` variants all collapse to
124            // `handlerInternalError` on the wire, matching the previous
125            // `map_invoke_error` behaviour that mapped
126            // `handlerExecutionFailed: ...` to `handlerInternalError`.
127            InvokeErrorKind::InvalidModel
128            | InvokeErrorKind::InvalidInput
129            | InvokeErrorKind::UnsupportedCapability
130            | InvokeErrorKind::ExecutionFailed => error_code::HANDLER_INTERNAL_ERROR,
131        }
132    }
133
134    /// Fixed public message. Never contains guest-supplied content.
135    pub const fn public_message(&self) -> &'static str {
136        match self.kind {
137            InvokeErrorKind::Internal => "internal runtime error",
138            InvokeErrorKind::Timeout => "handler exceeded the wall-clock deadline",
139            InvokeErrorKind::Trap => "handler trapped",
140            InvokeErrorKind::OutputTooLarge => "handler output exceeded the size limit",
141            InvokeErrorKind::InvalidOutput => "handler output was not valid JSON",
142            InvokeErrorKind::InvalidModel => "handler rejected the model configuration",
143            InvokeErrorKind::InvalidInput => "handler rejected the input",
144            InvokeErrorKind::UnsupportedCapability => "handler does not support the capability",
145            InvokeErrorKind::ExecutionFailed => "handler execution failed",
146        }
147    }
148
149    /// Whether the caller may retry the same request on a fresh handler.
150    pub const fn retryable(&self) -> bool {
151        matches!(self.kind, InvokeErrorKind::Timeout)
152    }
153}
154
155impl std::fmt::Display for InvokeError {
156    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
157        match &self.detail {
158            Some(detail) => write!(f, "{}: {}", self.stable_code(), detail),
159            None => f.write_str(self.stable_code()),
160        }
161    }
162}
163
164impl std::error::Error for InvokeError {}
165
166/// Truncate `s` to at most `max_bytes` on a UTF-8 char boundary.
167///
168/// `String::truncate` panics on a non-char boundary, so we walk backwards
169/// from `max_bytes` until `is_char_boundary` succeeds. The result is always
170/// valid UTF-8 and never longer than `max_bytes`.
171fn truncate_to_bytes(s: String, max_bytes: usize) -> String {
172    if s.len() <= max_bytes {
173        return s;
174    }
175    let mut end = max_bytes;
176    while end > 0 && !s.is_char_boundary(end) {
177        end -= 1;
178    }
179    let mut truncated = s;
180    truncated.truncate(end);
181    truncated
182}
183
184/// Minimal host-side log sink for invoke diagnostics.
185///
186/// Production code uses [`StderrLogSink`]. Downstream test harnesses can
187/// implement this trait to capture and verify log content without touching
188/// stderr. Keeping this trait tiny avoids pulling in a full logging
189/// framework while still making the runtime's only log call testable.
190///
191/// The sink receives a single pre-formatted message per invoke error. The
192/// message is constructed from the already-truncated
193/// [`InvokeError::detail`], so a malicious 16 KiB guest error payload can
194/// never produce a 16 KiB log line.
195pub trait HostLogSink: Send + Sync + std::fmt::Debug {
196    /// Emit a single log line. The implementation decides where it goes.
197    fn emit(&self, message: &str);
198}
199
200/// Default [`HostLogSink`] writing to stderr via `eprintln!`.
201#[derive(Debug, Default, Clone)]
202pub struct StderrLogSink;
203
204impl HostLogSink for StderrLogSink {
205    fn emit(&self, message: &str) {
206        eprintln!("{message}");
207    }
208}
209
210/// Consumers can implement this trait to add business-specific invocation logic.
211pub trait InvokeHandler: Send + Sync + std::fmt::Debug {
212    fn invoke(&self, capability: &str, input: &Value) -> Result<Value, InvokeError>;
213}
214
215/// Engine-side response produced by [`RuntimeEngine`]. The IPC layer converts
216/// this to a v1 [`RuntimeResponse`] or v2 [`RuntimeResponseV2`] based on the
217/// request's `api_version`.
218///
219/// This type is part of the 1.x stable API because it is the return type of
220/// [`RuntimeEngine::handle`]. Downstream consumers that embed the engine
221/// (rather than using the `rill-runtime` CLI) call `handle` and then convert
222/// the result to the appropriate wire version.
223#[derive(Debug, Clone)]
224#[non_exhaustive]
225pub enum EngineResponse {
226    Handshake {
227        request_id: String,
228        runtime_version: String,
229        model_pack_id: String,
230        model_pack_version: String,
231        capabilities: Vec<String>,
232        handler: Option<HandlerIdentity>,
233    },
234    Health {
235        request_id: String,
236        healthy: bool,
237        model_pack_id: String,
238        model_pack_version: String,
239    },
240    Result {
241        request_id: String,
242        output: Value,
243    },
244    Error {
245        request_id: String,
246        code: String,
247        message: String,
248        retryable: bool,
249    },
250}
251
252impl EngineResponse {
253    /// Convert to a v1 wire response. Handler identity fields are dropped.
254    pub fn to_v1(&self, api_version: u32) -> RuntimeResponse {
255        match self {
256            Self::Handshake {
257                request_id,
258                runtime_version,
259                model_pack_id,
260                model_pack_version,
261                capabilities,
262                ..
263            } => RuntimeResponse::Handshake {
264                request_id: request_id.clone(),
265                api_version,
266                runtime_version: runtime_version.clone(),
267                model_pack_id: model_pack_id.clone(),
268                model_pack_version: model_pack_version.clone(),
269                capabilities: capabilities.clone(),
270            },
271            Self::Health {
272                request_id,
273                healthy,
274                model_pack_id,
275                model_pack_version,
276            } => RuntimeResponse::Health {
277                request_id: request_id.clone(),
278                api_version,
279                healthy: *healthy,
280                model_pack_id: model_pack_id.clone(),
281                model_pack_version: model_pack_version.clone(),
282            },
283            Self::Result { request_id, output } => RuntimeResponse::Result {
284                request_id: request_id.clone(),
285                api_version,
286                output: output.clone(),
287            },
288            Self::Error {
289                request_id,
290                code,
291                message,
292                retryable,
293            } => RuntimeResponse::Error {
294                request_id: request_id.clone(),
295                api_version,
296                code: code.clone(),
297                message: message.clone(),
298                retryable: *retryable,
299            },
300        }
301    }
302
303    /// Convert to a v2 wire response. If no handler is loaded, handler fields
304    /// are filled with empty/zero values and effective_capabilities equals the
305    /// model capabilities.
306    pub fn to_v2(&self, api_version: u32) -> RuntimeResponseV2 {
307        match self {
308            Self::Handshake {
309                request_id,
310                runtime_version,
311                model_pack_id,
312                model_pack_version,
313                capabilities,
314                handler,
315            } => {
316                let (handler_id, handler_version, handler_api_version, effective) = match handler {
317                    Some(h) => (
318                        h.handler_id.clone(),
319                        h.handler_version.clone(),
320                        h.handler_api_version,
321                        h.effective_capabilities.clone(),
322                    ),
323                    None => (String::new(), String::new(), 0, capabilities.clone()),
324                };
325                RuntimeResponseV2::Handshake {
326                    request_id: request_id.clone(),
327                    api_version,
328                    runtime_version: runtime_version.clone(),
329                    model_pack_id: model_pack_id.clone(),
330                    model_pack_version: model_pack_version.clone(),
331                    capabilities: capabilities.clone(),
332                    handler_id,
333                    handler_version,
334                    handler_api_version,
335                    effective_capabilities: effective,
336                }
337            }
338            Self::Health {
339                request_id,
340                healthy,
341                model_pack_id,
342                model_pack_version,
343            } => RuntimeResponseV2::Health {
344                request_id: request_id.clone(),
345                api_version,
346                healthy: *healthy,
347                model_pack_id: model_pack_id.clone(),
348                model_pack_version: model_pack_version.clone(),
349            },
350            Self::Result { request_id, output } => RuntimeResponseV2::Result {
351                request_id: request_id.clone(),
352                api_version,
353                output: output.clone(),
354            },
355            Self::Error {
356                request_id,
357                code,
358                message,
359                retryable,
360            } => RuntimeResponseV2::Error {
361                request_id: request_id.clone(),
362                api_version,
363                code: code.clone(),
364                message: message.clone(),
365                retryable: *retryable,
366            },
367        }
368    }
369}
370
371#[derive(Debug, Clone)]
372pub struct RuntimeEngine {
373    pack: LoadedModelPack,
374    invoke_handler: Option<Arc<dyn InvokeHandler>>,
375    handler_identity: Option<HandlerIdentity>,
376    effective_capabilities: Vec<String>,
377    log_sink: Arc<dyn HostLogSink>,
378}
379
380impl RuntimeEngine {
381    pub fn new(pack: LoadedModelPack) -> Self {
382        Self {
383            pack,
384            invoke_handler: None,
385            handler_identity: None,
386            effective_capabilities: Vec::new(),
387            log_sink: Arc::new(StderrLogSink),
388        }
389    }
390
391    pub fn with_invoke_handler(mut self, handler: Arc<dyn InvokeHandler>) -> Self {
392        self.invoke_handler = Some(handler);
393        self
394    }
395
396    /// Replace the default [`StderrLogSink`] with a custom sink. Tests
397    /// inject a capturing sink to verify log bounds and content
398    /// without capturing stderr.
399    pub fn with_log_sink(mut self, sink: Arc<dyn HostLogSink>) -> Self {
400        self.log_sink = sink;
401        self
402    }
403
404    /// Attach handler identity and effective capabilities for IPC v2 handshake.
405    pub fn with_handler_identity(mut self, identity: HandlerIdentity) -> Self {
406        self.effective_capabilities = identity.effective_capabilities.clone();
407        self.handler_identity = Some(identity);
408        self
409    }
410
411    /// Effective capability set (intersection of model and handler). Empty when
412    /// no handler is loaded.
413    pub fn effective_capabilities(&self) -> &[String] {
414        &self.effective_capabilities
415    }
416
417    /// Handler identity if a handler was loaded.
418    pub fn handler_identity(&self) -> Option<&HandlerIdentity> {
419        self.handler_identity.as_ref()
420    }
421
422    pub fn handle(&self, request: RuntimeRequest) -> EngineResponse {
423        let request_id = request.request_id().to_string();
424        if request_id.is_empty() || request_id.len() > 128 {
425            return self.error(
426                request_id,
427                error_code::INVALID_REQUEST_ID,
428                "invalid request id",
429                false,
430            );
431        }
432        let api_version = request.api_version();
433        if !(MIN_RUNTIME_API_VERSION..=RUNTIME_API_VERSION).contains(&api_version) {
434            return self.error(
435                request_id,
436                error_code::INCOMPATIBLE_API_VERSION,
437                "runtime API version is not supported",
438                false,
439            );
440        }
441
442        match request {
443            RuntimeRequest::Handshake {
444                request_id,
445                client_name,
446                client_version,
447                ..
448            } => {
449                if client_name.is_empty()
450                    || client_name.len() > 96
451                    || client_version.is_empty()
452                    || client_version.len() > 48
453                {
454                    return self.error(
455                        request_id,
456                        error_code::INVALID_CLIENT_IDENTITY,
457                        "invalid client identity",
458                        false,
459                    );
460                }
461                EngineResponse::Handshake {
462                    request_id,
463                    runtime_version: env!("CARGO_PKG_VERSION").into(),
464                    model_pack_id: self.pack.manifest.id.clone(),
465                    model_pack_version: self.pack.manifest.version.clone(),
466                    capabilities: self.pack.manifest.capabilities.clone(),
467                    handler: self.handler_identity.clone(),
468                }
469            }
470            RuntimeRequest::Health { request_id, .. } => EngineResponse::Health {
471                request_id,
472                healthy: true,
473                model_pack_id: self.pack.manifest.id.clone(),
474                model_pack_version: self.pack.manifest.version.clone(),
475            },
476            RuntimeRequest::Invoke {
477                request_id,
478                capability,
479                input,
480                ..
481            } => {
482                if !self.is_capability_allowed(&capability) {
483                    return self.error(
484                        request_id,
485                        error_code::UNSUPPORTED_CAPABILITY,
486                        "capability is not in the effective set",
487                        false,
488                    );
489                }
490                let Some(handler) = &self.invoke_handler else {
491                    return self.error(
492                        request_id,
493                        error_code::NO_INVOKE_HANDLER,
494                        "no invoke handler registered",
495                        false,
496                    );
497                };
498                match handler.invoke(&capability, &input) {
499                    Ok(output) => EngineResponse::Result { request_id, output },
500                    Err(invoke_err) => {
501                        // Log host-side detail (if any) for debugging; the
502                        // IPC message is always the fixed public string so
503                        // guests cannot exfiltrate content via the error
504                        // payload. The detail is already truncated to
505                        // [`MAX_DETAIL_BYTES`] by [`InvokeError::with_detail`],
506                        // so a 16 KiB guest payload can never produce a
507                        // 16 KiB log line. This is the single log call for
508                        // invoke errors; the WASM adapter must not also
509                        // log the same error (see audit 5.2).
510                        if let Some(detail) = invoke_err.detail() {
511                            self.log_sink.emit(&format!(
512                                "rill-runtime: invoke {} -> {} (detail: {})",
513                                capability,
514                                invoke_err.stable_code(),
515                                detail
516                            ));
517                        }
518                        self.error(
519                            request_id,
520                            invoke_err.stable_code(),
521                            invoke_err.public_message(),
522                            invoke_err.retryable(),
523                        )
524                    }
525                }
526            }
527        }
528    }
529
530    /// Checks the capability against the effective set when a handler is loaded,
531    /// or against the model pack's declared capabilities when no handler is
532    /// loaded (for backwards compatibility with built-in handlers selected by
533    /// the binary).
534    fn is_capability_allowed(&self, capability: &str) -> bool {
535        if !self.effective_capabilities.is_empty() {
536            self.effective_capabilities.iter().any(|c| c == capability)
537        } else {
538            self.pack
539                .manifest
540                .capabilities
541                .iter()
542                .any(|c| c == capability)
543        }
544    }
545
546    fn error(
547        &self,
548        request_id: String,
549        code: &str,
550        message: &str,
551        retryable: bool,
552    ) -> EngineResponse {
553        EngineResponse::Error {
554            request_id,
555            code: code.into(),
556            message: message.into(),
557            retryable,
558        }
559    }
560}
561
562#[cfg(test)]
563mod tests {
564    use rill_runtime_protocol::{MODEL_PACK_FORMAT_VERSION, ModelPackManifest};
565    use std::sync::Mutex;
566
567    use super::*;
568    use crate::handler::builtin::LINEAR_REGRESSION_CAPABILITY;
569
570    /// Test-only [`HostLogSink`] that captures every emitted message in a
571    /// `Mutex<Vec<String>>`. Tests inspect the captured messages to verify
572    /// log bounds, content, and deduplication without touching stderr.
573    ///
574    /// This type is `pub(crate)` and lives inside `#[cfg(test)]` so it
575    /// never appears in the public API surface. Downstream test harnesses
576    /// that need similar functionality should implement [`HostLogSink`]
577    /// directly.
578    #[derive(Debug, Default)]
579    pub(crate) struct CapturingLogSink {
580        messages: Mutex<Vec<String>>,
581    }
582
583    impl CapturingLogSink {
584        /// Create an empty capturing sink.
585        pub(crate) fn new() -> Self {
586            Self::default()
587        }
588
589        /// Return a snapshot of all captured messages in emission order.
590        pub(crate) fn messages(&self) -> Vec<String> {
591            self.messages
592                .lock()
593                .expect("CapturingLogSink poisoned")
594                .clone()
595        }
596
597        /// Total byte length of all captured messages. Useful for asserting
598        /// that a 16 KiB guest error did not produce a 16 KiB log.
599        #[allow(dead_code)]
600        pub(crate) fn total_bytes(&self) -> usize {
601            self.messages
602                .lock()
603                .expect("CapturingLogSink poisoned")
604                .iter()
605                .map(String::len)
606                .sum()
607        }
608
609        /// Drop all captured messages.
610        #[allow(dead_code)]
611        pub(crate) fn clear(&self) {
612            self.messages
613                .lock()
614                .expect("CapturingLogSink poisoned")
615                .clear();
616        }
617    }
618
619    impl HostLogSink for CapturingLogSink {
620        fn emit(&self, message: &str) {
621            self.messages
622                .lock()
623                .expect("CapturingLogSink poisoned")
624                .push(message.to_string());
625        }
626    }
627
628    fn engine() -> RuntimeEngine {
629        RuntimeEngine::new(LoadedModelPack {
630            manifest: ModelPackManifest {
631                format_version: MODEL_PACK_FORMAT_VERSION,
632                id: "rillml.example.default".into(),
633                version: "0.7.0".into(),
634                runtime_api_version: RUNTIME_API_VERSION,
635                min_runtime_version: "0.7.0".into(),
636                publisher_key_id: "test".into(),
637                capabilities: vec!["rillml.example".into()],
638            },
639            model: serde_json::json!({}),
640        })
641    }
642
643    #[test]
644    fn handshake_reports_loaded_pack() {
645        let response = engine().handle(RuntimeRequest::Handshake {
646            request_id: "hello".into(),
647            api_version: RUNTIME_API_VERSION,
648            client_name: "example-host".into(),
649            client_version: "0.9.0".into(),
650        });
651        assert!(matches!(
652            response,
653            EngineResponse::Handshake { model_pack_id, .. }
654                if model_pack_id == "rillml.example.default"
655        ));
656    }
657
658    #[test]
659    fn incompatible_api_is_a_typed_error() {
660        let response = engine().handle(RuntimeRequest::Health {
661            request_id: "health".into(),
662            api_version: RUNTIME_API_VERSION + 1,
663        });
664        assert!(matches!(
665            response,
666            EngineResponse::Error { code, .. } if code == "incompatibleApiVersion"
667        ));
668    }
669
670    #[test]
671    fn invoke_without_handler_returns_no_invoke_handler_error() {
672        let response = engine().handle(RuntimeRequest::Invoke {
673            request_id: "invoke-1".into(),
674            api_version: RUNTIME_API_VERSION,
675            capability: "rillml.example".into(),
676            input: serde_json::json!({}),
677        });
678        assert!(matches!(
679            response,
680            EngineResponse::Error { code, .. } if code == "noInvokeHandler"
681        ));
682    }
683
684    #[test]
685    fn invoke_rejects_capability_not_declared_by_signed_manifest() {
686        let response = engine().handle(RuntimeRequest::Invoke {
687            request_id: "invoke-undeclared".into(),
688            api_version: RUNTIME_API_VERSION,
689            capability: "undeclared.capability".into(),
690            input: serde_json::json!({}),
691        });
692        assert!(matches!(
693            response,
694            EngineResponse::Error { code, .. } if code == "unsupportedCapability"
695        ));
696    }
697
698    #[test]
699    fn v1_handshake_omits_handler_fields() {
700        let identity = HandlerIdentity {
701            handler_id: "org.example.handler".into(),
702            handler_version: "1.0.0".into(),
703            handler_api_version: 1,
704            effective_capabilities: vec!["rillml.example".into()],
705        };
706        let engine = engine().with_handler_identity(identity);
707        let response = engine.handle(RuntimeRequest::Handshake {
708            request_id: "v1-test".into(),
709            api_version: 1,
710            client_name: "v1-host".into(),
711            client_version: "0.6.0".into(),
712        });
713        let v1 = response.to_v1(1);
714        let json = serde_json::to_string(&v1).unwrap();
715        assert!(!json.contains("handlerId"));
716        assert!(!json.contains("effectiveCapabilities"));
717    }
718
719    #[test]
720    fn v2_handshake_includes_handler_fields() {
721        let identity = HandlerIdentity {
722            handler_id: "org.example.handler".into(),
723            handler_version: "1.0.0".into(),
724            handler_api_version: 1,
725            effective_capabilities: vec!["rillml.example".into()],
726        };
727        let engine = engine().with_handler_identity(identity);
728        let response = engine.handle(RuntimeRequest::Handshake {
729            request_id: "v2-test".into(),
730            api_version: 2,
731            client_name: "v2-host".into(),
732            client_version: "0.7.0".into(),
733        });
734        let v2 = response.to_v2(2);
735        let json = serde_json::to_string(&v2).unwrap();
736        assert!(json.contains("\"handlerId\":\"org.example.handler\""));
737        assert!(json.contains("\"handlerApiVersion\":1"));
738        assert!(json.contains("\"effectiveCapabilities\":[\"rillml.example\"]"));
739    }
740
741    #[test]
742    fn v2_handshake_without_handler_has_empty_fields() {
743        let response = engine().handle(RuntimeRequest::Handshake {
744            request_id: "v2-no-handler".into(),
745            api_version: 2,
746            client_name: "v2-host".into(),
747            client_version: "0.7.0".into(),
748        });
749        let v2 = response.to_v2(2);
750        match v2 {
751            RuntimeResponseV2::Handshake {
752                handler_id,
753                handler_version,
754                handler_api_version,
755                effective_capabilities,
756                ..
757            } => {
758                assert!(handler_id.is_empty());
759                assert!(handler_version.is_empty());
760                assert_eq!(handler_api_version, 0);
761                assert_eq!(effective_capabilities, vec!["rillml.example"]);
762            }
763            _ => panic!("expected handshake"),
764        }
765    }
766
767    #[test]
768    fn linear_regression_handler_validates_and_predicts() {
769        use crate::handler::builtin::LinearRegressionInvokeHandler;
770
771        let pack = LoadedModelPack {
772            manifest: ModelPackManifest {
773                format_version: MODEL_PACK_FORMAT_VERSION,
774                id: "rillml.example.default".into(),
775                version: "0.7.0".into(),
776                runtime_api_version: RUNTIME_API_VERSION,
777                min_runtime_version: "0.7.0".into(),
778                publisher_key_id: "test".into(),
779                capabilities: vec![LINEAR_REGRESSION_CAPABILITY.into()],
780            },
781            model: serde_json::json!({
782                "kind": "linearRegression",
783                "weights": [0.5, -0.25],
784                "intercept": 1.0
785            }),
786        };
787        let handler = LinearRegressionInvokeHandler::from_pack(&pack).unwrap();
788        let engine = RuntimeEngine::new(pack).with_invoke_handler(Arc::new(handler));
789        let response = engine.handle(RuntimeRequest::Invoke {
790            request_id: "invoke-linear".into(),
791            api_version: RUNTIME_API_VERSION,
792            capability: LINEAR_REGRESSION_CAPABILITY.into(),
793            input: serde_json::json!({"features": [4.0, 2.0]}),
794        });
795        assert!(matches!(
796            response,
797            EngineResponse::Result { output, .. } if output["prediction"] == 2.5
798        ));
799    }
800
801    #[test]
802    fn invoke_error_stable_codes_match_wire_format() {
803        // Every kind must map to the exact IPC code expected by v1/v2
804        // clients, preserving backwards compatibility with the previous
805        // `map_invoke_error` string matching.
806        assert_eq!(
807            InvokeError::new(InvokeErrorKind::Trap).stable_code(),
808            "handlerTrap"
809        );
810        assert_eq!(
811            InvokeError::new(InvokeErrorKind::Timeout).stable_code(),
812            "handlerTimeout"
813        );
814        assert_eq!(
815            InvokeError::new(InvokeErrorKind::OutputTooLarge).stable_code(),
816            "handlerOutputTooLarge"
817        );
818        assert_eq!(
819            InvokeError::new(InvokeErrorKind::InvalidOutput).stable_code(),
820            "handlerInvalidOutput"
821        );
822        assert_eq!(
823            InvokeError::new(InvokeErrorKind::Internal).stable_code(),
824            "handlerInternalError"
825        );
826        // All four guest-reported WIT handler-error variants collapse to
827        // handlerInternalError on the wire, matching the previous
828        // `map_invoke_error` behaviour. The host distinguishes them
829        // internally via `kind()` for logging, but v1/v2 clients see
830        // the same code.
831        for kind in [
832            InvokeErrorKind::InvalidModel,
833            InvokeErrorKind::InvalidInput,
834            InvokeErrorKind::UnsupportedCapability,
835            InvokeErrorKind::ExecutionFailed,
836        ] {
837            assert_eq!(
838                InvokeError::new(kind).stable_code(),
839                "handlerInternalError",
840                "{kind:?} must map to handlerInternalError for v1/v2 compat"
841            );
842        }
843    }
844
845    #[test]
846    fn invoke_error_retryable_only_for_timeout() {
847        assert!(InvokeError::new(InvokeErrorKind::Timeout).retryable());
848        for kind in [
849            InvokeErrorKind::Trap,
850            InvokeErrorKind::OutputTooLarge,
851            InvokeErrorKind::InvalidOutput,
852            InvokeErrorKind::Internal,
853            InvokeErrorKind::InvalidModel,
854            InvokeErrorKind::InvalidInput,
855            InvokeErrorKind::UnsupportedCapability,
856            InvokeErrorKind::ExecutionFailed,
857        ] {
858            assert!(
859                !InvokeError::new(kind).retryable(),
860                "{kind:?} must not be retryable"
861            );
862        }
863    }
864
865    #[test]
866    fn invoke_error_guest_variants_have_distinct_public_messages() {
867        // Each guest variant carries a fixed public message that never
868        // contains guest-supplied content. The messages are distinct so
869        // operators can distinguish variants in host logs.
870        let messages = [
871            InvokeError::new(InvokeErrorKind::InvalidModel).public_message(),
872            InvokeError::new(InvokeErrorKind::InvalidInput).public_message(),
873            InvokeError::new(InvokeErrorKind::UnsupportedCapability).public_message(),
874            InvokeError::new(InvokeErrorKind::ExecutionFailed).public_message(),
875        ];
876        // All distinct.
877        for i in 0..messages.len() {
878            for j in (i + 1)..messages.len() {
879                assert_ne!(messages[i], messages[j], "public messages must be distinct");
880            }
881        }
882        // None contain guest content markers.
883        for msg in messages {
884            assert!(!msg.contains("detail"));
885            assert!(!msg.contains("guest"));
886        }
887    }
888
889    #[test]
890    fn invoke_error_public_message_never_contains_detail() {
891        // Guest can fully control the detail string; the public message
892        // must always be the fixed constant.
893        let err = InvokeError::with_detail(
894            InvokeErrorKind::ExecutionFailed,
895            "SECRET-TOKEN-LEAK-ATTEMPT guest-controlled-payload",
896        );
897        assert_eq!(err.public_message(), "handler execution failed");
898        assert_eq!(err.stable_code(), "handlerInternalError");
899        assert_eq!(
900            err.detail(),
901            Some("SECRET-TOKEN-LEAK-ATTEMPT guest-controlled-payload")
902        );
903        // The Display impl is for host logs only; the IPC layer must
904        // never send `err.to_string()` to clients.
905        assert!(err.to_string().contains("SECRET-TOKEN-LEAK-ATTEMPT"));
906        // The public_message is what the IPC layer actually sends.
907        assert!(!err.public_message().contains("SECRET"));
908    }
909
910    #[test]
911    fn invoke_error_without_detail_has_no_detail() {
912        let err = InvokeError::new(InvokeErrorKind::Trap);
913        assert_eq!(err.kind(), InvokeErrorKind::Trap);
914        assert_eq!(err.detail(), None);
915        assert_eq!(err.stable_code(), "handlerTrap");
916        assert_eq!(err.to_string(), "handlerTrap");
917    }
918
919    #[test]
920    fn invoke_error_detail_is_truncated_to_4kib_on_char_boundary() {
921        // A malicious guest tries to grow host memory via an oversized
922        // error payload. The host must truncate to MAX_DETAIL_BYTES on a
923        // UTF-8 char boundary.
924        let huge = "A".repeat(MAX_DETAIL_BYTES * 4);
925        let err = InvokeError::with_detail(InvokeErrorKind::ExecutionFailed, huge);
926        let detail = err.detail().expect("detail must be stored");
927        assert!(
928            detail.len() <= MAX_DETAIL_BYTES,
929            "detail length {} must not exceed {}",
930            detail.len(),
931            MAX_DETAIL_BYTES
932        );
933        // Truncation must land on a char boundary (the string is valid UTF-8
934        // by construction, but the test guards against a future unsafe path).
935        assert!(detail.chars().all(|c| c == 'A'));
936    }
937
938    #[test]
939    fn invoke_error_detail_truncation_respects_multibyte_chars() {
940        // Multi-byte UTF-8 must not be split mid-codepoint. Use 3-byte
941        // CJK characters so the MAX_DETAIL_BYTES boundary lands inside a
942        // character; the result must back up to the previous char boundary.
943        let emoji = "🌟".repeat(MAX_DETAIL_BYTES); // each '🌟' is 4 bytes
944        let err = InvokeError::with_detail(InvokeErrorKind::ExecutionFailed, emoji);
945        let detail = err.detail().expect("detail must be stored");
946        assert!(detail.len() <= MAX_DETAIL_BYTES);
947        // Every stored character must be a complete '🌟'.
948        for c in detail.chars() {
949            assert_eq!(c, '🌟');
950        }
951    }
952
953    /// Minimal handler that always returns the supplied error, for
954    /// exercising the engine's invoke error path without a WASM sandbox.
955    #[derive(Debug)]
956    struct FailingHandler {
957        err: InvokeError,
958    }
959
960    impl InvokeHandler for FailingHandler {
961        fn invoke(&self, _capability: &str, _input: &Value) -> Result<Value, InvokeError> {
962            Err(self.err.clone())
963        }
964    }
965
966    #[test]
967    fn engine_invoke_error_does_not_leak_guest_detail_in_message() {
968        // A malicious guest tries to exfiltrate a token via the WIT
969        // handler-error payload. The IPC Error.message field must be
970        // the fixed public string, not the guest-supplied detail.
971        let err = InvokeError::with_detail(
972            InvokeErrorKind::ExecutionFailed,
973            "leak-attempt:SECRET-TOKEN",
974        );
975        let pack = LoadedModelPack {
976            manifest: ModelPackManifest {
977                format_version: MODEL_PACK_FORMAT_VERSION,
978                id: "rillml.example.default".into(),
979                version: "0.7.0".into(),
980                runtime_api_version: RUNTIME_API_VERSION,
981                min_runtime_version: "0.7.0".into(),
982                publisher_key_id: "test".into(),
983                capabilities: vec!["rillml.example".into()],
984            },
985            model: serde_json::json!({}),
986        };
987        let sink = Arc::new(CapturingLogSink::new());
988        let engine = RuntimeEngine::new(pack)
989            .with_invoke_handler(Arc::new(FailingHandler { err }))
990            .with_log_sink(sink.clone());
991        let response = engine.handle(RuntimeRequest::Invoke {
992            request_id: "leak-test".into(),
993            api_version: RUNTIME_API_VERSION,
994            capability: "rillml.example".into(),
995            input: serde_json::json!({}),
996        });
997        match response {
998            EngineResponse::Error {
999                code,
1000                message,
1001                retryable,
1002                ..
1003            } => {
1004                assert_eq!(code, "handlerInternalError");
1005                assert_eq!(message, "handler execution failed");
1006                assert!(!retryable);
1007                // The guest-supplied detail must NOT appear anywhere in
1008                // the IPC response fields.
1009                assert!(!message.contains("SECRET"));
1010                assert!(!message.contains("leak-attempt"));
1011            }
1012            _ => panic!("expected EngineResponse::Error"),
1013        }
1014        // The host log line does contain the (truncated) detail for
1015        // operator diagnostics, but the detail is host-only — it never
1016        // reaches the IPC `message` field. This assertion documents that
1017        // the log sink received exactly one message referencing the
1018        // secret, proving the detail was captured host-side.
1019        let messages = sink.messages();
1020        assert_eq!(
1021            messages.len(),
1022            1,
1023            "the engine must log the invoke error exactly once"
1024        );
1025        assert!(messages[0].contains("SECRET-TOKEN"));
1026    }
1027
1028    /// Verifies audit 5.2: a 16 KiB guest error payload must not produce
1029    /// a 16 KiB log line. The host constructs `InvokeError::with_detail`
1030    /// (which truncates to `MAX_DETAIL_BYTES`) before logging, so the
1031    /// captured log message must be well under 16 KiB.
1032    #[test]
1033    fn engine_log_does_not_emit_oversized_guest_detail() {
1034        let huge_detail = "X".repeat(MAX_DETAIL_BYTES * 4); // 16 KiB
1035        let err = InvokeError::with_detail(InvokeErrorKind::ExecutionFailed, huge_detail);
1036        let pack = LoadedModelPack {
1037            manifest: ModelPackManifest {
1038                format_version: MODEL_PACK_FORMAT_VERSION,
1039                id: "rillml.example.default".into(),
1040                version: "0.7.0".into(),
1041                runtime_api_version: RUNTIME_API_VERSION,
1042                min_runtime_version: "0.7.0".into(),
1043                publisher_key_id: "test".into(),
1044                capabilities: vec!["rillml.example".into()],
1045            },
1046            model: serde_json::json!({}),
1047        };
1048        let sink = Arc::new(CapturingLogSink::new());
1049        let engine = RuntimeEngine::new(pack)
1050            .with_invoke_handler(Arc::new(FailingHandler { err }))
1051            .with_log_sink(sink.clone());
1052        let _ = engine.handle(RuntimeRequest::Invoke {
1053            request_id: "oversized".into(),
1054            api_version: RUNTIME_API_VERSION,
1055            capability: "rillml.example".into(),
1056            input: serde_json::json!({}),
1057        });
1058        let messages = sink.messages();
1059        assert_eq!(messages.len(), 1, "exactly one log line expected");
1060        let log_line = &messages[0];
1061        // The log line consists of a fixed prefix + the truncated detail.
1062        // The detail is at most MAX_DETAIL_BYTES; the prefix is small.
1063        // 16 KiB must never appear in the log.
1064        assert!(
1065            log_line.len() < MAX_DETAIL_BYTES * 2,
1066            "log line length {} must be well under 2x MAX_DETAIL_BYTES ({}); \
1067             a 16 KiB guest payload must not produce a 16 KiB log",
1068            log_line.len(),
1069            MAX_DETAIL_BYTES * 2
1070        );
1071        // The detail portion (after the prefix) must not exceed the cap.
1072        assert!(
1073            log_line.len() < MAX_DETAIL_BYTES + 256,
1074            "log line length {} must be < MAX_DETAIL_BYTES + prefix overhead",
1075            log_line.len()
1076        );
1077    }
1078
1079    /// Verifies audit 5.2: the same invoke error must not be logged
1080    /// twice. The WASM adapter must not log the error if the engine
1081    /// already logs it; this test uses a `FailingHandler` (no WASM
1082    /// adapter) and confirms exactly one log line per invoke.
1083    #[test]
1084    fn engine_logs_invoke_error_exactly_once() {
1085        let err = InvokeError::with_detail(
1086            InvokeErrorKind::UnsupportedCapability,
1087            "capability foo not supported",
1088        );
1089        let pack = LoadedModelPack {
1090            manifest: ModelPackManifest {
1091                format_version: MODEL_PACK_FORMAT_VERSION,
1092                id: "rillml.example.default".into(),
1093                version: "0.7.0".into(),
1094                runtime_api_version: RUNTIME_API_VERSION,
1095                min_runtime_version: "0.7.0".into(),
1096                publisher_key_id: "test".into(),
1097                capabilities: vec!["rillml.example".into()],
1098            },
1099            model: serde_json::json!({}),
1100        };
1101        let sink = Arc::new(CapturingLogSink::new());
1102        let engine = RuntimeEngine::new(pack)
1103            .with_invoke_handler(Arc::new(FailingHandler { err }))
1104            .with_log_sink(sink.clone());
1105        let _ = engine.handle(RuntimeRequest::Invoke {
1106            request_id: "once".into(),
1107            api_version: RUNTIME_API_VERSION,
1108            capability: "rillml.example".into(),
1109            input: serde_json::json!({}),
1110        });
1111        assert_eq!(
1112            sink.messages().len(),
1113            1,
1114            "the engine must log the invoke error exactly once, not twice"
1115        );
1116    }
1117
1118    /// Verifies audit 5.2: a trap backtrace (which can be very long)
1119    /// must be truncated before logging. The `FailingHandler` simulates
1120    /// a trap with a long backtrace-like detail string.
1121    #[test]
1122    fn engine_log_traps_backtrace_is_truncated() {
1123        let fake_backtrace = "trap: unreachable\n".repeat(1024); // ~17 KiB
1124        let err = InvokeError::with_detail(InvokeErrorKind::Trap, fake_backtrace);
1125        let pack = LoadedModelPack {
1126            manifest: ModelPackManifest {
1127                format_version: MODEL_PACK_FORMAT_VERSION,
1128                id: "rillml.example.default".into(),
1129                version: "0.7.0".into(),
1130                runtime_api_version: RUNTIME_API_VERSION,
1131                min_runtime_version: "0.7.0".into(),
1132                publisher_key_id: "test".into(),
1133                capabilities: vec!["rillml.example".into()],
1134            },
1135            model: serde_json::json!({}),
1136        };
1137        let sink = Arc::new(CapturingLogSink::new());
1138        let engine = RuntimeEngine::new(pack)
1139            .with_invoke_handler(Arc::new(FailingHandler { err }))
1140            .with_log_sink(sink.clone());
1141        let _ = engine.handle(RuntimeRequest::Invoke {
1142            request_id: "trap-trunc".into(),
1143            api_version: RUNTIME_API_VERSION,
1144            capability: "rillml.example".into(),
1145            input: serde_json::json!({}),
1146        });
1147        let messages = sink.messages();
1148        assert_eq!(messages.len(), 1);
1149        let log_line = &messages[0];
1150        assert!(
1151            log_line.len() < MAX_DETAIL_BYTES + 256,
1152            "trap backtrace log must be truncated; got {} bytes",
1153            log_line.len()
1154        );
1155    }
1156
1157    /// Verifies that all four guest WIT variants flow through the engine
1158    /// with the correct `kind()` and fixed public message, while the
1159    /// stable IPC code stays `handlerInternalError` for v1/v2 compat.
1160    #[test]
1161    fn engine_preserves_guest_variant_kind_for_all_wit_variants() {
1162        for (kind, expected_message) in [
1163            (
1164                InvokeErrorKind::InvalidModel,
1165                "handler rejected the model configuration",
1166            ),
1167            (InvokeErrorKind::InvalidInput, "handler rejected the input"),
1168            (
1169                InvokeErrorKind::UnsupportedCapability,
1170                "handler does not support the capability",
1171            ),
1172            (InvokeErrorKind::ExecutionFailed, "handler execution failed"),
1173        ] {
1174            let err = InvokeError::with_detail(kind, "guest detail");
1175            let pack = LoadedModelPack {
1176                manifest: ModelPackManifest {
1177                    format_version: MODEL_PACK_FORMAT_VERSION,
1178                    id: "rillml.example.default".into(),
1179                    version: "0.7.0".into(),
1180                    runtime_api_version: RUNTIME_API_VERSION,
1181                    min_runtime_version: "0.7.0".into(),
1182                    publisher_key_id: "test".into(),
1183                    capabilities: vec!["rillml.example".into()],
1184                },
1185                model: serde_json::json!({}),
1186            };
1187            let engine =
1188                RuntimeEngine::new(pack).with_invoke_handler(Arc::new(FailingHandler { err }));
1189            let response = engine.handle(RuntimeRequest::Invoke {
1190                request_id: "variant".into(),
1191                api_version: RUNTIME_API_VERSION,
1192                capability: "rillml.example".into(),
1193                input: serde_json::json!({}),
1194            });
1195            match response {
1196                EngineResponse::Error { code, message, .. } => {
1197                    assert_eq!(
1198                        code, "handlerInternalError",
1199                        "{kind:?}: stable code must stay handlerInternalError"
1200                    );
1201                    assert_eq!(
1202                        message, expected_message,
1203                        "{kind:?}: public message mismatch"
1204                    );
1205                }
1206                _ => panic!("{kind:?}: expected EngineResponse::Error"),
1207            }
1208        }
1209    }
1210}