allframe_forge/templates/
websocket.rs

1//! WebSocket Gateway archetype templates
2//!
3//! Templates for generating real-time WebSocket services with channel-based messaging.
4
5use crate::config::ProjectConfig;
6
7/// Convert a string to PascalCase
8fn to_pascal_case(s: &str) -> String {
9    s.split(|c| c == '-' || c == '_')
10        .map(|word| {
11            let mut chars = word.chars();
12            match chars.next() {
13                None => String::new(),
14                Some(first) => first.to_uppercase().chain(chars).collect(),
15            }
16        })
17        .collect()
18}
19
20/// Generate Cargo.toml for websocket gateway project
21pub fn cargo_toml(config: &ProjectConfig) -> String {
22    let ws = config.websocket_gateway.as_ref().unwrap();
23    let name = &config.name;
24
25    format!(
26        r#"[package]
27name = "{name}"
28version = "0.1.0"
29edition = "2021"
30rust-version = "1.75"
31description = "{display_name}"
32
33[dependencies]
34# AllFrame
35allframe-core = {{ version = "0.1", features = ["resilience", "otel"] }}
36
37# Web Framework & WebSocket
38axum = {{ version = "0.7", features = ["ws"] }}
39tokio-tungstenite = "0.26"
40tower = "0.5"
41tower-http = {{ version = "0.6", features = ["trace", "cors"] }}
42
43# Async
44tokio = {{ version = "1", features = ["full"] }}
45async-trait = "0.1"
46futures = "0.3"
47
48# Serialization
49serde = {{ version = "1.0", features = ["derive"] }}
50serde_json = "1.0"
51
52# Error handling
53thiserror = "2.0"
54anyhow = "1.0"
55
56# Tracing & Metrics
57tracing = "0.1"
58tracing-subscriber = {{ version = "0.3", features = ["env-filter"] }}
59opentelemetry = {{ version = "0.27", features = ["metrics"] }}
60
61# Utilities
62chrono = {{ version = "0.4", features = ["serde"] }}
63uuid = {{ version = "1.0", features = ["v4", "serde"] }}
64dashmap = "6.0"
65dotenvy = "0.15"
66
67[dev-dependencies]
68tokio-test = "0.4"
69
70[[bin]]
71name = "{name}"
72path = "src/main.rs"
73"#,
74        name = name,
75        display_name = ws.display_name,
76    )
77}
78
79/// Generate main.rs
80pub fn main_rs(config: &ProjectConfig) -> String {
81    let ws = config.websocket_gateway.as_ref().unwrap();
82    let pascal_name = to_pascal_case(&ws.service_name);
83
84    format!(
85        r#"//! {display_name}
86//!
87//! A WebSocket gateway service for real-time bidirectional communication.
88
89use std::sync::Arc;
90use tracing::info;
91
92mod config;
93mod error;
94mod domain;
95mod application;
96mod infrastructure;
97mod presentation;
98
99use config::Config;
100use application::{pascal_name}Hub;
101use infrastructure::HealthServer;
102
103#[tokio::main]
104async fn main() -> anyhow::Result<()> {{
105    // Load environment variables
106    dotenvy::dotenv().ok();
107
108    // Initialize tracing
109    tracing_subscriber::fmt()
110        .with_env_filter(
111            tracing_subscriber::EnvFilter::from_default_env()
112                .add_directive(tracing::Level::INFO.into()),
113        )
114        .init();
115
116    // Load configuration
117    let config = Config::from_env();
118    info!("Starting {display_name}");
119    info!("Configured channels: {{:?}}", config.channels.iter().map(|c| &c.name).collect::<Vec<_>>());
120
121    // Create WebSocket hub
122    let hub = Arc::new({pascal_name}Hub::new(config.clone()));
123
124    // Start hub background tasks
125    let hub_handle = {{
126        let hub = hub.clone();
127        tokio::spawn(async move {{
128            hub.run().await
129        }})
130    }};
131
132    // Start health server in background
133    let health_port = config.server.health_port;
134    let health_handle = tokio::spawn(async move {{
135        let health_server = HealthServer::new(health_port);
136        health_server.run().await
137    }});
138
139    // Create router and start WebSocket server
140    let app = presentation::create_router(hub.clone());
141
142    info!("Starting WebSocket server on port {{}}", config.server.http_port);
143    let listener = tokio::net::TcpListener::bind(
144        format!("0.0.0.0:{{}}", config.server.http_port)
145    ).await?;
146    axum::serve(listener, app).await?;
147
148    hub_handle.abort();
149    health_handle.abort();
150    info!("WebSocket gateway shutdown complete");
151    Ok(())
152}}
153"#,
154        pascal_name = pascal_name,
155        display_name = ws.display_name,
156    )
157}
158
159/// Generate config.rs
160pub fn config_rs(config: &ProjectConfig) -> String {
161    let ws = config.websocket_gateway.as_ref().unwrap();
162
163    let channel_configs: Vec<String> = ws
164        .channels
165        .iter()
166        .map(|ch| {
167            format!(
168                r#"            ChannelConfig {{
169                name: "{}".to_string(),
170                description: "{}".to_string(),
171                authenticated: {},
172            }}"#,
173                ch.name, ch.description, ch.authenticated
174            )
175        })
176        .collect();
177
178    format!(
179        r#"//! Service configuration
180
181use std::env;
182
183/// Main configuration
184#[derive(Debug, Clone)]
185pub struct Config {{
186    pub server: ServerConfig,
187    pub channels: Vec<ChannelConfig>,
188    pub max_connections_per_client: u32,
189    pub heartbeat_interval_secs: u64,
190    pub connection_timeout_secs: u64,
191}}
192
193/// Server configuration
194#[derive(Debug, Clone)]
195pub struct ServerConfig {{
196    pub http_port: u16,
197    pub health_port: u16,
198}}
199
200/// Channel configuration
201#[derive(Debug, Clone)]
202pub struct ChannelConfig {{
203    pub name: String,
204    pub description: String,
205    pub authenticated: bool,
206}}
207
208impl Config {{
209    pub fn from_env() -> Self {{
210        Self {{
211            server: ServerConfig {{
212                http_port: env::var("PORT")
213                    .unwrap_or_else(|_| "{http_port}".to_string())
214                    .parse()
215                    .expect("PORT must be a number"),
216                health_port: env::var("HEALTH_PORT")
217                    .unwrap_or_else(|_| "{health_port}".to_string())
218                    .parse()
219                    .expect("HEALTH_PORT must be a number"),
220            }},
221            channels: vec![
222{channel_configs}
223            ],
224            max_connections_per_client: env::var("MAX_CONNECTIONS_PER_CLIENT")
225                .unwrap_or_else(|_| "{max_connections}".to_string())
226                .parse()
227                .expect("MAX_CONNECTIONS_PER_CLIENT must be a number"),
228            heartbeat_interval_secs: env::var("HEARTBEAT_INTERVAL_SECS")
229                .unwrap_or_else(|_| "{heartbeat_interval}".to_string())
230                .parse()
231                .expect("HEARTBEAT_INTERVAL_SECS must be a number"),
232            connection_timeout_secs: env::var("CONNECTION_TIMEOUT_SECS")
233                .unwrap_or_else(|_| "{connection_timeout}".to_string())
234                .parse()
235                .expect("CONNECTION_TIMEOUT_SECS must be a number"),
236        }}
237    }}
238}}
239"#,
240        http_port = ws.server.http_port,
241        health_port = ws.server.health_port,
242        channel_configs = channel_configs.join(",\n"),
243        max_connections = ws.max_connections_per_client,
244        heartbeat_interval = ws.heartbeat_interval_secs,
245        connection_timeout = ws.connection_timeout_secs,
246    )
247}
248
249/// Generate error.rs
250pub fn error_rs(config: &ProjectConfig) -> String {
251    let ws = config.websocket_gateway.as_ref().unwrap();
252    let pascal_name = to_pascal_case(&ws.service_name);
253
254    format!(
255        r#"//! Error types
256
257use thiserror::Error;
258
259/// WebSocket gateway errors
260#[derive(Error, Debug)]
261pub enum {pascal_name}Error {{
262    #[error("Connection error: {{0}}")]
263    Connection(String),
264
265    #[error("Channel not found: {{0}}")]
266    ChannelNotFound(String),
267
268    #[error("Authentication required")]
269    AuthRequired,
270
271    #[error("Connection limit exceeded")]
272    ConnectionLimitExceeded,
273
274    #[error("Message send error: {{0}}")]
275    SendError(String),
276
277    #[error("Internal error: {{0}}")]
278    Internal(String),
279}}
280"#,
281        pascal_name = pascal_name,
282    )
283}
284
285/// Generate domain/mod.rs
286pub fn domain_mod(_config: &ProjectConfig) -> String {
287    r#"//! Domain layer
288
289pub mod messages;
290pub mod connection;
291
292pub use messages::*;
293pub use connection::*;
294"#
295    .to_string()
296}
297
298/// Generate domain/messages.rs
299pub fn domain_messages(config: &ProjectConfig) -> String {
300    let ws = config.websocket_gateway.as_ref().unwrap();
301    let pascal_name = to_pascal_case(&ws.service_name);
302
303    format!(
304        r#"//! WebSocket message types
305
306use chrono::{{DateTime, Utc}};
307use serde::{{Deserialize, Serialize}};
308use uuid::Uuid;
309
310/// Client to server message
311#[derive(Debug, Clone, Serialize, Deserialize)]
312#[serde(tag = "type", content = "payload")]
313pub enum ClientMessage {{
314    /// Subscribe to a channel
315    Subscribe {{ channel: String }},
316    /// Unsubscribe from a channel
317    Unsubscribe {{ channel: String }},
318    /// Send a message to a channel
319    Publish {{ channel: String, data: serde_json::Value }},
320    /// Ping for keepalive
321    Ping,
322}}
323
324/// Server to client message
325#[derive(Debug, Clone, Serialize, Deserialize)]
326#[serde(tag = "type", content = "payload")]
327pub enum ServerMessage {{
328    /// Subscription confirmed
329    Subscribed {{ channel: String }},
330    /// Unsubscription confirmed
331    Unsubscribed {{ channel: String }},
332    /// Message from a channel
333    Message(ChannelMessage),
334    /// Pong response
335    Pong,
336    /// Error message
337    Error {{ code: String, message: String }},
338}}
339
340/// Channel message with metadata
341#[derive(Debug, Clone, Serialize, Deserialize)]
342pub struct ChannelMessage {{
343    pub id: Uuid,
344    pub channel: String,
345    pub data: serde_json::Value,
346    pub sender_id: Option<Uuid>,
347    pub timestamp: DateTime<Utc>,
348}}
349
350impl ChannelMessage {{
351    pub fn new(channel: String, data: serde_json::Value, sender_id: Option<Uuid>) -> Self {{
352        Self {{
353            id: Uuid::new_v4(),
354            channel,
355            data,
356            sender_id,
357            timestamp: Utc::now(),
358        }}
359    }}
360}}
361
362/// {pascal_name} event for internal use
363#[derive(Debug, Clone)]
364pub enum HubEvent {{
365    ClientConnected {{ client_id: Uuid }},
366    ClientDisconnected {{ client_id: Uuid }},
367    Subscribe {{ client_id: Uuid, channel: String }},
368    Unsubscribe {{ client_id: Uuid, channel: String }},
369    Broadcast {{ channel: String, message: ChannelMessage }},
370}}
371"#,
372        pascal_name = pascal_name,
373    )
374}
375
376/// Generate domain/connection.rs
377pub fn domain_connection(_config: &ProjectConfig) -> String {
378    r#"//! Connection management
379
380use chrono::{DateTime, Utc};
381use uuid::Uuid;
382use std::collections::HashSet;
383
384/// Represents a connected client
385#[derive(Debug, Clone)]
386pub struct ClientConnection {
387    pub id: Uuid,
388    pub connected_at: DateTime<Utc>,
389    pub last_heartbeat: DateTime<Utc>,
390    pub subscriptions: HashSet<String>,
391    pub user_id: Option<Uuid>,
392}
393
394impl ClientConnection {
395    pub fn new() -> Self {
396        let now = Utc::now();
397        Self {
398            id: Uuid::new_v4(),
399            connected_at: now,
400            last_heartbeat: now,
401            subscriptions: HashSet::new(),
402            user_id: None,
403        }
404    }
405
406    pub fn with_user(mut self, user_id: Uuid) -> Self {
407        self.user_id = Some(user_id);
408        self
409    }
410
411    pub fn update_heartbeat(&mut self) {
412        self.last_heartbeat = Utc::now();
413    }
414
415    pub fn subscribe(&mut self, channel: &str) {
416        self.subscriptions.insert(channel.to_string());
417    }
418
419    pub fn unsubscribe(&mut self, channel: &str) {
420        self.subscriptions.remove(channel);
421    }
422
423    pub fn is_subscribed(&self, channel: &str) -> bool {
424        self.subscriptions.contains(channel)
425    }
426}
427
428impl Default for ClientConnection {
429    fn default() -> Self {
430        Self::new()
431    }
432}
433"#
434    .to_string()
435}
436
437/// Generate application/mod.rs
438pub fn application_mod(_config: &ProjectConfig) -> String {
439    r#"//! Application layer
440
441pub mod hub;
442
443pub use hub::*;
444"#
445    .to_string()
446}
447
448/// Generate application/hub.rs
449pub fn application_hub(config: &ProjectConfig) -> String {
450    let ws = config.websocket_gateway.as_ref().unwrap();
451    let pascal_name = to_pascal_case(&ws.service_name);
452
453    format!(
454        r#"//! WebSocket hub for managing connections and channels
455
456use std::sync::Arc;
457use dashmap::DashMap;
458use tokio::sync::{{broadcast, mpsc}};
459use tracing::{{info, warn, debug}};
460use uuid::Uuid;
461
462use crate::config::Config;
463use crate::domain::{{ClientConnection, ChannelMessage, HubEvent, ServerMessage}};
464use crate::error::{pascal_name}Error;
465
466/// Channel for sending messages to clients
467pub type ClientSender = mpsc::UnboundedSender<ServerMessage>;
468
469/// WebSocket hub managing all connections and channels
470pub struct {pascal_name}Hub {{
471    config: Config,
472    /// Connected clients
473    connections: DashMap<Uuid, ClientSender>,
474    /// Channel subscriptions: channel_name -> set of client_ids
475    subscriptions: DashMap<String, DashMap<Uuid, ()>>,
476    /// Broadcast channel for hub events
477    event_tx: broadcast::Sender<HubEvent>,
478}}
479
480impl {pascal_name}Hub {{
481    pub fn new(config: Config) -> Self {{
482        let (event_tx, _) = broadcast::channel(1024);
483
484        // Pre-create channels from config
485        let subscriptions = DashMap::new();
486        for channel in &config.channels {{
487            subscriptions.insert(channel.name.clone(), DashMap::new());
488        }}
489
490        Self {{
491            config,
492            connections: DashMap::new(),
493            subscriptions,
494            event_tx,
495        }}
496    }}
497
498    /// Run the hub event loop
499    pub async fn run(&self) {{
500        let mut rx = self.event_tx.subscribe();
501
502        loop {{
503            match rx.recv().await {{
504                Ok(event) => {{
505                    self.handle_event(event).await;
506                }}
507                Err(broadcast::error::RecvError::Lagged(n)) => {{
508                    warn!("Hub event receiver lagged by {{}} messages", n);
509                }}
510                Err(broadcast::error::RecvError::Closed) => {{
511                    info!("Hub event channel closed");
512                    break;
513                }}
514            }}
515        }}
516    }}
517
518    async fn handle_event(&self, event: HubEvent) {{
519        match event {{
520            HubEvent::ClientConnected {{ client_id }} => {{
521                debug!(client_id = %client_id, "Client connected");
522            }}
523            HubEvent::ClientDisconnected {{ client_id }} => {{
524                debug!(client_id = %client_id, "Client disconnected");
525                self.cleanup_client(client_id);
526            }}
527            HubEvent::Subscribe {{ client_id, channel }} => {{
528                if let Some(subs) = self.subscriptions.get(&channel) {{
529                    subs.insert(client_id, ());
530                    debug!(client_id = %client_id, channel = %channel, "Client subscribed");
531                }}
532            }}
533            HubEvent::Unsubscribe {{ client_id, channel }} => {{
534                if let Some(subs) = self.subscriptions.get(&channel) {{
535                    subs.remove(&client_id);
536                    debug!(client_id = %client_id, channel = %channel, "Client unsubscribed");
537                }}
538            }}
539            HubEvent::Broadcast {{ channel, message }} => {{
540                self.broadcast_to_channel(&channel, message).await;
541            }}
542        }}
543    }}
544
545    fn cleanup_client(&self, client_id: Uuid) {{
546        self.connections.remove(&client_id);
547        for entry in self.subscriptions.iter() {{
548            entry.value().remove(&client_id);
549        }}
550    }}
551
552    async fn broadcast_to_channel(&self, channel: &str, message: ChannelMessage) {{
553        if let Some(subscribers) = self.subscriptions.get(channel) {{
554            let msg = ServerMessage::Message(message);
555            for entry in subscribers.iter() {{
556                let client_id = *entry.key();
557                if let Some(sender) = self.connections.get(&client_id) {{
558                    let _ = sender.send(msg.clone());
559                }}
560            }}
561        }}
562    }}
563
564    /// Register a new client connection
565    pub fn register_client(&self, client_id: Uuid, sender: ClientSender) {{
566        self.connections.insert(client_id, sender);
567        let _ = self.event_tx.send(HubEvent::ClientConnected {{ client_id }});
568    }}
569
570    /// Remove a client connection
571    pub fn unregister_client(&self, client_id: Uuid) {{
572        let _ = self.event_tx.send(HubEvent::ClientDisconnected {{ client_id }});
573    }}
574
575    /// Subscribe client to a channel
576    pub fn subscribe(&self, client_id: Uuid, channel: &str) -> Result<(), {pascal_name}Error> {{
577        if !self.subscriptions.contains_key(channel) {{
578            return Err({pascal_name}Error::ChannelNotFound(channel.to_string()));
579        }}
580
581        let _ = self.event_tx.send(HubEvent::Subscribe {{
582            client_id,
583            channel: channel.to_string(),
584        }});
585
586        // Send confirmation
587        if let Some(sender) = self.connections.get(&client_id) {{
588            let _ = sender.send(ServerMessage::Subscribed {{
589                channel: channel.to_string(),
590            }});
591        }}
592
593        Ok(())
594    }}
595
596    /// Unsubscribe client from a channel
597    pub fn unsubscribe(&self, client_id: Uuid, channel: &str) {{
598        let _ = self.event_tx.send(HubEvent::Unsubscribe {{
599            client_id,
600            channel: channel.to_string(),
601        }});
602
603        if let Some(sender) = self.connections.get(&client_id) {{
604            let _ = sender.send(ServerMessage::Unsubscribed {{
605                channel: channel.to_string(),
606            }});
607        }}
608    }}
609
610    /// Publish a message to a channel
611    pub fn publish(&self, channel: &str, data: serde_json::Value, sender_id: Option<Uuid>) -> Result<(), {pascal_name}Error> {{
612        if !self.subscriptions.contains_key(channel) {{
613            return Err({pascal_name}Error::ChannelNotFound(channel.to_string()));
614        }}
615
616        let message = ChannelMessage::new(channel.to_string(), data, sender_id);
617        let _ = self.event_tx.send(HubEvent::Broadcast {{
618            channel: channel.to_string(),
619            message,
620        }});
621
622        Ok(())
623    }}
624
625    /// Send pong response to a client
626    pub fn send_pong(&self, client_id: Uuid) {{
627        if let Some(sender) = self.connections.get(&client_id) {{
628            let _ = sender.send(ServerMessage::Pong);
629        }}
630    }}
631
632    /// Get the number of connected clients
633    pub fn connection_count(&self) -> usize {{
634        self.connections.len()
635    }}
636
637    /// Get available channels
638    pub fn channels(&self) -> Vec<String> {{
639        self.config.channels.iter().map(|c| c.name.clone()).collect()
640    }}
641}}
642"#,
643        pascal_name = pascal_name,
644    )
645}
646
647/// Generate infrastructure/mod.rs
648pub fn infrastructure_mod(_config: &ProjectConfig) -> String {
649    r#"//! Infrastructure layer
650
651pub mod health;
652
653pub use health::*;
654"#
655    .to_string()
656}
657
658/// Generate infrastructure/health.rs
659pub fn infrastructure_health(_config: &ProjectConfig) -> String {
660    r#"//! Health check server
661
662use axum::{routing::get, Router};
663use tracing::info;
664
665/// Health check server for Kubernetes probes
666pub struct HealthServer {
667    port: u16,
668}
669
670impl HealthServer {
671    pub fn new(port: u16) -> Self {
672        Self { port }
673    }
674
675    pub async fn run(&self) -> Result<(), std::io::Error> {
676        let app = Router::new()
677            .route("/health", get(health))
678            .route("/ready", get(ready));
679
680        let addr = format!("0.0.0.0:{}", self.port);
681        info!("Health server listening on {}", addr);
682
683        let listener = tokio::net::TcpListener::bind(&addr).await?;
684        axum::serve(listener, app).await
685    }
686}
687
688async fn health() -> &'static str {
689    "OK"
690}
691
692async fn ready() -> &'static str {
693    "OK"
694}
695"#
696    .to_string()
697}
698
699/// Generate presentation/mod.rs
700pub fn presentation_mod(_config: &ProjectConfig) -> String {
701    r#"//! Presentation layer
702
703pub mod handlers;
704
705pub use handlers::*;
706"#
707    .to_string()
708}
709
710/// Generate presentation/handlers.rs
711pub fn presentation_handlers(config: &ProjectConfig) -> String {
712    let ws = config.websocket_gateway.as_ref().unwrap();
713    let pascal_name = to_pascal_case(&ws.service_name);
714
715    format!(
716        r#"//! WebSocket handlers
717
718use std::sync::Arc;
719use axum::{{
720    extract::{{ws::{{Message, WebSocket, WebSocketUpgrade}}, State}},
721    response::IntoResponse,
722    routing::get,
723    Router,
724}};
725use futures::{{SinkExt, StreamExt}};
726use tokio::sync::mpsc;
727use tracing::{{error, info, debug}};
728use uuid::Uuid;
729
730use crate::application::{pascal_name}Hub;
731use crate::domain::{{ClientConnection, ClientMessage, ServerMessage}};
732
733type AppState = Arc<{pascal_name}Hub>;
734
735/// Create the WebSocket router
736pub fn create_router(hub: Arc<{pascal_name}Hub>) -> Router {{
737    Router::new()
738        .route("/ws", get(ws_handler))
739        .route("/channels", get(list_channels))
740        .route("/stats", get(stats))
741        .with_state(hub)
742}}
743
744async fn ws_handler(
745    ws: WebSocketUpgrade,
746    State(hub): State<AppState>,
747) -> impl IntoResponse {{
748    ws.on_upgrade(move |socket| handle_socket(socket, hub))
749}}
750
751async fn handle_socket(socket: WebSocket, hub: Arc<{pascal_name}Hub>) {{
752    let connection = ClientConnection::new();
753    let client_id = connection.id;
754
755    info!(client_id = %client_id, "WebSocket connection established");
756
757    let (mut ws_sender, mut ws_receiver) = socket.split();
758
759    // Create channel for sending messages to this client
760    let (tx, mut rx) = mpsc::unbounded_channel::<ServerMessage>();
761
762    // Register client with hub
763    hub.register_client(client_id, tx);
764
765    // Spawn task to forward messages from hub to WebSocket
766    let send_task = tokio::spawn(async move {{
767        while let Some(msg) = rx.recv().await {{
768            match serde_json::to_string(&msg) {{
769                Ok(text) => {{
770                    if ws_sender.send(Message::Text(text.into())).await.is_err() {{
771                        break;
772                    }}
773                }}
774                Err(e) => {{
775                    error!("Failed to serialize message: {{}}", e);
776                }}
777            }}
778        }}
779    }});
780
781    // Process incoming messages
782    while let Some(result) = ws_receiver.next().await {{
783        match result {{
784            Ok(Message::Text(text)) => {{
785                match serde_json::from_str::<ClientMessage>(&text) {{
786                    Ok(msg) => handle_client_message(msg, client_id, &hub),
787                    Err(e) => {{
788                        debug!("Invalid message format: {{}}", e);
789                    }}
790                }}
791            }}
792            Ok(Message::Close(_)) => {{
793                break;
794            }}
795            Ok(Message::Ping(_)) => {{
796                // Axum handles ping/pong automatically
797            }}
798            Ok(_) => {{
799                // Ignore other message types
800            }}
801            Err(e) => {{
802                error!(client_id = %client_id, error = %e, "WebSocket error");
803                break;
804            }}
805        }}
806    }}
807
808    // Cleanup
809    hub.unregister_client(client_id);
810    send_task.abort();
811
812    info!(client_id = %client_id, "WebSocket connection closed");
813}}
814
815fn handle_client_message(msg: ClientMessage, client_id: Uuid, hub: &{pascal_name}Hub) {{
816    match msg {{
817        ClientMessage::Subscribe {{ channel }} => {{
818            if let Err(e) = hub.subscribe(client_id, &channel) {{
819                error!(client_id = %client_id, error = %e, "Subscribe failed");
820            }}
821        }}
822        ClientMessage::Unsubscribe {{ channel }} => {{
823            hub.unsubscribe(client_id, &channel);
824        }}
825        ClientMessage::Publish {{ channel, data }} => {{
826            if let Err(e) = hub.publish(&channel, data, Some(client_id)) {{
827                error!(client_id = %client_id, error = %e, "Publish failed");
828            }}
829        }}
830        ClientMessage::Ping => {{
831            hub.send_pong(client_id);
832        }}
833    }}
834}}
835
836async fn list_channels(State(hub): State<AppState>) -> axum::Json<Vec<String>> {{
837    axum::Json(hub.channels())
838}}
839
840async fn stats(State(hub): State<AppState>) -> axum::Json<serde_json::Value> {{
841    axum::Json(serde_json::json!({{
842        "connections": hub.connection_count(),
843        "channels": hub.channels(),
844    }}))
845}}
846"#,
847        pascal_name = pascal_name,
848    )
849}
850
851/// Generate README.md
852pub fn readme(config: &ProjectConfig) -> String {
853    let ws = config.websocket_gateway.as_ref().unwrap();
854    let name = &config.name;
855
856    let channel_table: Vec<String> = ws
857        .channels
858        .iter()
859        .map(|ch| {
860            format!(
861                "| {} | {} | {} |",
862                ch.name,
863                ch.description,
864                if ch.authenticated { "Yes" } else { "No" }
865            )
866        })
867        .collect();
868
869    format!(
870        r#"# {display_name}
871
872A WebSocket gateway service built with AllFrame for real-time bidirectional communication.
873
874## Features
875
876- **Channel-based Messaging**: Subscribe/publish to named channels
877- **Connection Management**: Automatic heartbeat and timeout handling
878- **Scalable**: DashMap-based concurrent connection handling
879- **OpenTelemetry**: Distributed tracing and metrics
880- **Health Checks**: Kubernetes-ready liveness and readiness probes
881
882## Prerequisites
883
884- Rust 1.75+
885
886## Configuration
887
888Set the following environment variables:
889
890```bash
891# Server
892PORT=8080
893HEALTH_PORT=8081
894
895# WebSocket
896MAX_CONNECTIONS_PER_CLIENT=5
897HEARTBEAT_INTERVAL_SECS=30
898CONNECTION_TIMEOUT_SECS=60
899```
900
901## Channels
902
903| Name | Description | Auth Required |
904|------|-------------|---------------|
905{channel_table}
906
907## Running
908
909```bash
910# Development
911cargo run
912
913# Production
914cargo build --release
915./target/release/{name}
916```
917
918## WebSocket API
919
920### Connect
921
922```javascript
923const ws = new WebSocket('ws://localhost:8080/ws');
924```
925
926### Message Types
927
928#### Subscribe to Channel
929```json
930{{"type": "Subscribe", "payload": {{"channel": "events"}}}}
931```
932
933#### Unsubscribe from Channel
934```json
935{{"type": "Unsubscribe", "payload": {{"channel": "events"}}}}
936```
937
938#### Publish to Channel
939```json
940{{"type": "Publish", "payload": {{"channel": "events", "data": {{"key": "value"}}}}}}
941```
942
943#### Ping
944```json
945{{"type": "Ping"}}
946```
947
948### Server Responses
949
950```json
951// Subscription confirmed
952{{"type": "Subscribed", "payload": {{"channel": "events"}}}}
953
954// Message received
955{{"type": "Message", "payload": {{"id": "...", "channel": "events", "data": {{...}}, "timestamp": "..."}}}}
956
957// Pong
958{{"type": "Pong"}}
959
960// Error
961{{"type": "Error", "payload": {{"code": "CHANNEL_NOT_FOUND", "message": "..."}}}}
962```
963
964## REST Endpoints
965
966| Method | Endpoint | Description |
967|--------|----------|-------------|
968| GET | /ws | WebSocket upgrade endpoint |
969| GET | /channels | List available channels |
970| GET | /stats | Connection statistics |
971
972## Architecture
973
974```
975┌───────────────────────────────────────────────────────────┐
976│                   WebSocket Gateway                        │
977│  ┌───────────┐    ┌──────────┐    ┌────────────────────┐  │
978│  │  Handler  │───▶│   Hub    │───▶│   Subscriptions    │  │
979│  └───────────┘    └──────────┘    └────────────────────┘  │
980│        │               │                   │              │
981│        ▼               ▼                   ▼              │
982│  ┌───────────┐    ┌──────────┐    ┌────────────────────┐  │
983│  │ WebSocket │    │ Broadcast│    │    Connections     │  │
984│  │  Upgrade  │    │  Channel │    │      (DashMap)     │  │
985│  └───────────┘    └──────────┘    └────────────────────┘  │
986└───────────────────────────────────────────────────────────┘
987```
988
989## License
990
991MIT
992"#,
993        display_name = ws.display_name,
994        name = name,
995        channel_table = channel_table.join("\n"),
996    )
997}
998
999/// Generate Dockerfile
1000pub fn dockerfile(config: &ProjectConfig) -> String {
1001    let name = &config.name;
1002
1003    format!(
1004        r#"FROM rust:1.75-slim as builder
1005
1006WORKDIR /app
1007
1008# Install dependencies
1009RUN apt-get update && apt-get install -y \
1010    pkg-config \
1011    libssl-dev \
1012    && rm -rf /var/lib/apt/lists/*
1013
1014# Copy manifests
1015COPY Cargo.toml Cargo.lock ./
1016
1017# Create dummy main to cache dependencies
1018RUN mkdir src && echo "fn main() {{}}" > src/main.rs
1019RUN cargo build --release
1020RUN rm -rf src
1021
1022# Copy source
1023COPY src ./src
1024
1025# Build
1026RUN touch src/main.rs && cargo build --release
1027
1028# Runtime image
1029FROM debian:bookworm-slim
1030
1031RUN apt-get update && apt-get install -y \
1032    ca-certificates \
1033    libssl3 \
1034    && rm -rf /var/lib/apt/lists/*
1035
1036WORKDIR /app
1037
1038COPY --from=builder /app/target/release/{name} /app/
1039
1040ENV PORT=8080
1041ENV HEALTH_PORT=8081
1042EXPOSE 8080 8081
1043
1044CMD ["/app/{name}"]
1045"#,
1046        name = name,
1047    )
1048}
1049
1050#[cfg(test)]
1051mod tests {
1052    use super::*;
1053
1054    #[test]
1055    fn test_to_pascal_case() {
1056        assert_eq!(to_pascal_case("ws_gateway"), "WsGateway");
1057        assert_eq!(to_pascal_case("real-time"), "RealTime");
1058        assert_eq!(to_pascal_case("simple"), "Simple");
1059    }
1060}
allframe_forge/templates/websocket.rs

allframe_forge/templates/
websocket.rs
