rill_runtime/
server.rs

1use std::sync::Arc;
2
3use rill_runtime_protocol::{RUNTIME_API_VERSION, RuntimeRequest, RuntimeResponse};
4use serde_json::Value;
5
6use crate::package::LoadedModelPack;
7
8/// 消费方实现此 trait 处理 Invoke 请求。
9/// RillML runtime 不提供默认实现。
10pub trait InvokeHandler: Send + Sync + std::fmt::Debug {
11    fn invoke(&self, capability: &str, input: &Value) -> Result<Value, String>;
12}
13
14#[derive(Debug, Clone)]
15pub struct RuntimeEngine {
16    pack: LoadedModelPack,
17    invoke_handler: Option<Arc<dyn InvokeHandler>>,
18}
19
20impl RuntimeEngine {
21    pub fn new(pack: LoadedModelPack) -> Self {
22        Self {
23            pack,
24            invoke_handler: None,
25        }
26    }
27
28    pub fn with_invoke_handler(mut self, handler: Arc<dyn InvokeHandler>) -> Self {
29        self.invoke_handler = Some(handler);
30        self
31    }
32
33    pub fn handle(&self, request: RuntimeRequest) -> RuntimeResponse {
34        let request_id = request.request_id().to_string();
35        if request_id.is_empty() || request_id.len() > 128 {
36            return self.error(request_id, "invalidRequestId", "invalid request id", false);
37        }
38        if request.api_version() != RUNTIME_API_VERSION {
39            return self.error(
40                request_id,
41                "incompatibleApiVersion",
42                "runtime API version is not supported",
43                false,
44            );
45        }
46
47        match request {
48            RuntimeRequest::Handshake {
49                request_id,
50                client_name,
51                client_version,
52                ..
53            } => {
54                if client_name.is_empty()
55                    || client_name.len() > 96
56                    || client_version.is_empty()
57                    || client_version.len() > 48
58                {
59                    return self.error(
60                        request_id,
61                        "invalidClientIdentity",
62                        "invalid client identity",
63                        false,
64                    );
65                }
66                RuntimeResponse::Handshake {
67                    request_id,
68                    api_version: RUNTIME_API_VERSION,
69                    runtime_version: env!("CARGO_PKG_VERSION").into(),
70                    model_pack_id: self.pack.manifest.id.clone(),
71                    model_pack_version: self.pack.manifest.version.clone(),
72                    capabilities: self.pack.manifest.capabilities.clone(),
73                }
74            }
75            RuntimeRequest::Health { request_id, .. } => RuntimeResponse::Health {
76                request_id,
77                api_version: RUNTIME_API_VERSION,
78                healthy: true,
79                model_pack_id: self.pack.manifest.id.clone(),
80                model_pack_version: self.pack.manifest.version.clone(),
81            },
82            RuntimeRequest::Invoke {
83                request_id,
84                capability,
85                input,
86                ..
87            } => match &self.invoke_handler {
88                Some(handler) => match handler.invoke(&capability, &input) {
89                    Ok(output) => RuntimeResponse::Result {
90                        request_id,
91                        api_version: RUNTIME_API_VERSION,
92                        output,
93                    },
94                    Err(message) => self.error(request_id, "invokeFailed", &message, false),
95                },
96                None => self.error(
97                    request_id,
98                    "noInvokeHandler",
99                    "no invoke handler registered",
100                    false,
101                ),
102            },
103        }
104    }
105
106    fn error(
107        &self,
108        request_id: String,
109        code: &str,
110        message: &str,
111        retryable: bool,
112    ) -> RuntimeResponse {
113        RuntimeResponse::Error {
114            request_id,
115            api_version: RUNTIME_API_VERSION,
116            code: code.into(),
117            message: message.into(),
118            retryable,
119        }
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    use rill_runtime_protocol::{MODEL_PACK_FORMAT_VERSION, ModelPackManifest};
126
127    use super::*;
128
129    fn engine() -> RuntimeEngine {
130        RuntimeEngine::new(LoadedModelPack {
131            manifest: ModelPackManifest {
132                format_version: MODEL_PACK_FORMAT_VERSION,
133                id: "rillml.example.default".into(),
134                version: "0.5.0".into(),
135                runtime_api_version: RUNTIME_API_VERSION,
136                min_runtime_version: "0.5.0".into(),
137                publisher_key_id: "test".into(),
138                capabilities: vec!["rillml.example".into()],
139            },
140            model: serde_json::json!({}),
141        })
142    }
143
144    #[test]
145    fn handshake_reports_loaded_pack() {
146        let response = engine().handle(RuntimeRequest::Handshake {
147            request_id: "hello".into(),
148            api_version: RUNTIME_API_VERSION,
149            client_name: "example-host".into(),
150            client_version: "0.9.0".into(),
151        });
152        assert!(matches!(
153            response,
154            RuntimeResponse::Handshake { model_pack_id, .. }
155                if model_pack_id == "rillml.example.default"
156        ));
157    }
158
159    #[test]
160    fn incompatible_api_is_a_typed_error() {
161        let response = engine().handle(RuntimeRequest::Health {
162            request_id: "health".into(),
163            api_version: RUNTIME_API_VERSION + 1,
164        });
165        assert!(matches!(
166            response,
167            RuntimeResponse::Error { code, .. } if code == "incompatibleApiVersion"
168        ));
169    }
170
171    #[test]
172    fn invoke_without_handler_returns_no_invoke_handler_error() {
173        let response = engine().handle(RuntimeRequest::Invoke {
174            request_id: "invoke-1".into(),
175            api_version: RUNTIME_API_VERSION,
176            capability: "rillml.example".into(),
177            input: serde_json::json!({}),
178        });
179        assert!(matches!(
180            response,
181            RuntimeResponse::Error { code, .. } if code == "noInvokeHandler"
182        ));
183    }
184}
rill_runtime/server.rs

rill_runtime/
server.rs
