Skip to main content

procwire_client/control/
init.rs

1//! `$init` message builder.
2//!
3//! The `$init` message is sent via stdout to tell the parent:
4//! - The pipe path for data plane connection
5//! - The schema (methods and events with their IDs)
6//! - Protocol version
7//!
8//! # Example
9//!
10//! ```
11//! use procwire_client::control::{InitSchema, ResponseType, build_init_message};
12//!
13//! let mut schema = InitSchema::new();
14//! schema.add_method("echo", 1, ResponseType::Result);
15//! schema.add_method("generate", 2, ResponseType::Stream);
16//! schema.add_event("progress", 3);
17//!
18//! let json = build_init_message("/tmp/procwire.sock", &schema);
19//! assert!(json.contains("$init"));
20//! ```
21
22use serde::Serialize;
23use serde_json::json;
24use std::collections::HashMap;
25
26/// Response type for a method.
27#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
28#[serde(rename_all = "lowercase")]
29pub enum ResponseType {
30    /// Returns a single result.
31    Result,
32    /// Returns a stream of chunks.
33    Stream,
34    /// Returns only an acknowledgment.
35    Ack,
36    /// Fire-and-forget, no response.
37    None,
38}
39
40/// Method definition with ID and response type.
41#[derive(Debug, Clone)]
42pub struct MethodSchema {
43    /// Assigned method ID (1-65534).
44    pub id: u16,
45    /// Expected response type.
46    pub response: ResponseType,
47}
48
49/// Event definition with ID.
50#[derive(Debug, Clone)]
51pub struct EventSchema {
52    /// Assigned event ID (1-65534).
53    pub id: u16,
54}
55
56/// Schema describing available methods and events.
57///
58/// Used to build the `$init` message that tells the parent
59/// what methods and events this worker supports.
60#[derive(Debug, Clone, Default)]
61pub struct InitSchema {
62    /// Map of method names to their definitions.
63    pub methods: HashMap<String, MethodSchema>,
64    /// Map of event names to their definitions.
65    pub events: HashMap<String, EventSchema>,
66}
67
68impl InitSchema {
69    /// Create a new empty schema.
70    pub fn new() -> Self {
71        Self::default()
72    }
73
74    /// Add a method to the schema.
75    ///
76    /// # Arguments
77    ///
78    /// * `name` - Method name (e.g., "echo")
79    /// * `id` - Assigned method ID (must be 1-65534)
80    /// * `response` - Expected response type
81    pub fn add_method(&mut self, name: &str, id: u16, response: ResponseType) {
82        self.methods
83            .insert(name.to_string(), MethodSchema { id, response });
84    }
85
86    /// Add an event to the schema.
87    ///
88    /// # Arguments
89    ///
90    /// * `name` - Event name (e.g., "progress")
91    /// * `id` - Assigned event ID (must be 1-65534)
92    pub fn add_event(&mut self, name: &str, id: u16) {
93        self.events.insert(name.to_string(), EventSchema { id });
94    }
95
96    /// Get a method by name.
97    pub fn get_method(&self, name: &str) -> Option<&MethodSchema> {
98        self.methods.get(name)
99    }
100
101    /// Get an event by name.
102    pub fn get_event(&self, name: &str) -> Option<&EventSchema> {
103        self.events.get(name)
104    }
105
106    /// Check if schema is empty.
107    pub fn is_empty(&self) -> bool {
108        self.methods.is_empty() && self.events.is_empty()
109    }
110}
111
112/// Protocol version string.
113pub const PROTOCOL_VERSION: &str = "2.0.0";
114
115/// Build the `$init` JSON-RPC message.
116///
117/// This creates the JSON message that should be sent to the parent
118/// via stdout to initialize the connection.
119///
120/// # Arguments
121///
122/// * `pipe_path` - Path to the pipe/socket for data plane
123/// * `schema` - Schema describing methods and events
124///
125/// # Returns
126///
127/// JSON string ready to be written to stdout.
128pub fn build_init_message(pipe_path: &str, schema: &InitSchema) -> String {
129    let methods: serde_json::Map<String, serde_json::Value> = schema
130        .methods
131        .iter()
132        .map(|(name, m)| {
133            (
134                name.clone(),
135                json!({
136                    "id": m.id,
137                    "response": m.response
138                }),
139            )
140        })
141        .collect();
142
143    let events: serde_json::Map<String, serde_json::Value> = schema
144        .events
145        .iter()
146        .map(|(name, e)| (name.clone(), json!({ "id": e.id })))
147        .collect();
148
149    let msg = json!({
150        "jsonrpc": "2.0",
151        "method": "$init",
152        "params": {
153            "pipe": pipe_path,
154            "schema": {
155                "methods": methods,
156                "events": events
157            },
158            "version": PROTOCOL_VERSION
159        }
160    });
161
162    serde_json::to_string(&msg).expect("JSON serialization should not fail")
163}
164
165// Legacy types for backward compatibility
166// TODO: Remove after migration
167
168/// The `$init` message (legacy struct, prefer `build_init_message`).
169#[derive(Debug, Serialize)]
170pub struct InitMessage {
171    jsonrpc: &'static str,
172    method: &'static str,
173    params: InitParams,
174}
175
176/// Parameters for the `$init` message (legacy).
177#[derive(Debug, Serialize)]
178#[serde(rename_all = "camelCase")]
179pub struct InitParams {
180    /// Path to the named pipe or Unix socket.
181    pub pipe_path: String,
182    /// Schema describing available methods and events.
183    pub schema: Schema,
184}
185
186/// Schema (legacy, prefer `InitSchema`).
187#[derive(Debug, Default, Serialize)]
188pub struct Schema {
189    /// Map of method names to their definitions.
190    pub methods: HashMap<String, MethodDef>,
191    /// Map of event names to their definitions.
192    pub events: HashMap<String, EventDef>,
193}
194
195/// Method definition (legacy).
196#[derive(Debug, Serialize)]
197pub struct MethodDef {
198    /// Assigned method ID.
199    pub id: u16,
200}
201
202/// Event definition (legacy).
203#[derive(Debug, Serialize)]
204pub struct EventDef {
205    /// Assigned event ID.
206    pub id: u16,
207}
208
209impl InitMessage {
210    /// Create a new `$init` message (legacy).
211    pub fn new(pipe_path: String, schema: Schema) -> Self {
212        Self {
213            jsonrpc: "2.0",
214            method: "$init",
215            params: InitParams { pipe_path, schema },
216        }
217    }
218}
219
220#[cfg(test)]
221mod tests {
222    use super::*;
223
224    #[test]
225    fn test_build_init_message_format() {
226        let mut schema = InitSchema::new();
227        schema.add_method("echo", 1, ResponseType::Result);
228        schema.add_method("generate", 2, ResponseType::Stream);
229        schema.add_event("progress", 3);
230
231        let msg = build_init_message("/tmp/test.sock", &schema);
232        let parsed: serde_json::Value = serde_json::from_str(&msg).unwrap();
233
234        assert_eq!(parsed["jsonrpc"], "2.0");
235        assert_eq!(parsed["method"], "$init");
236        assert_eq!(parsed["params"]["pipe"], "/tmp/test.sock");
237        assert_eq!(parsed["params"]["version"], "2.0.0");
238    }
239
240    #[test]
241    fn test_method_schema() {
242        let mut schema = InitSchema::new();
243        schema.add_method("echo", 1, ResponseType::Result);
244
245        let msg = build_init_message("/tmp/test.sock", &schema);
246        let parsed: serde_json::Value = serde_json::from_str(&msg).unwrap();
247
248        assert_eq!(parsed["params"]["schema"]["methods"]["echo"]["id"], 1);
249        assert_eq!(
250            parsed["params"]["schema"]["methods"]["echo"]["response"],
251            "result"
252        );
253    }
254
255    #[test]
256    fn test_event_schema() {
257        let mut schema = InitSchema::new();
258        schema.add_event("progress", 5);
259
260        let msg = build_init_message("/tmp/test.sock", &schema);
261        let parsed: serde_json::Value = serde_json::from_str(&msg).unwrap();
262
263        assert_eq!(parsed["params"]["schema"]["events"]["progress"]["id"], 5);
264    }
265
266    #[test]
267    fn test_response_types() {
268        let mut schema = InitSchema::new();
269        schema.add_method("result_method", 1, ResponseType::Result);
270        schema.add_method("stream_method", 2, ResponseType::Stream);
271        schema.add_method("ack_method", 3, ResponseType::Ack);
272        schema.add_method("none_method", 4, ResponseType::None);
273
274        let msg = build_init_message("/tmp/test.sock", &schema);
275        let parsed: serde_json::Value = serde_json::from_str(&msg).unwrap();
276
277        let methods = &parsed["params"]["schema"]["methods"];
278        assert_eq!(methods["result_method"]["response"], "result");
279        assert_eq!(methods["stream_method"]["response"], "stream");
280        assert_eq!(methods["ack_method"]["response"], "ack");
281        assert_eq!(methods["none_method"]["response"], "none");
282    }
283
284    #[test]
285    fn test_empty_schema() {
286        let schema = InitSchema::new();
287        assert!(schema.is_empty());
288
289        let msg = build_init_message("/tmp/test.sock", &schema);
290        let parsed: serde_json::Value = serde_json::from_str(&msg).unwrap();
291
292        assert!(parsed["params"]["schema"]["methods"]
293            .as_object()
294            .unwrap()
295            .is_empty());
296        assert!(parsed["params"]["schema"]["events"]
297            .as_object()
298            .unwrap()
299            .is_empty());
300    }
301
302    #[test]
303    fn test_schema_get_method() {
304        let mut schema = InitSchema::new();
305        schema.add_method("echo", 1, ResponseType::Result);
306
307        let method = schema.get_method("echo").unwrap();
308        assert_eq!(method.id, 1);
309        assert_eq!(method.response, ResponseType::Result);
310
311        assert!(schema.get_method("nonexistent").is_none());
312    }
313
314    #[test]
315    fn test_schema_get_event() {
316        let mut schema = InitSchema::new();
317        schema.add_event("progress", 5);
318
319        let event = schema.get_event("progress").unwrap();
320        assert_eq!(event.id, 5);
321
322        assert!(schema.get_event("nonexistent").is_none());
323    }
324}