allframe_forge/templates/
websocket.rs

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

allframe_forge/templates/
websocket.rs
