turul_mcp_aws_lambda/
server.rs

1//! Lambda MCP Server Implementation
2//!
3//! This module provides the main Lambda MCP server implementation that mirrors
4//! the architecture of turul-mcp-server but adapted for AWS Lambda deployment.
5
6use std::collections::HashMap;
7use std::sync::Arc;
8
9use tracing::{debug, info};
10
11use turul_http_mcp_server::{ServerConfig, StreamConfig, StreamManager};
12use turul_mcp_protocol::{Implementation, ServerCapabilities};
13use turul_mcp_server::{
14    McpCompletion, McpElicitation, McpLogger, McpNotification, McpPrompt, McpResource, McpRoot,
15    McpSampling, McpTool, handlers::McpHandler, session::SessionManager,
16};
17use turul_mcp_session_storage::BoxedSessionStorage;
18
19use crate::error::Result;
20use crate::handler::LambdaMcpHandler;
21
22#[cfg(feature = "cors")]
23use crate::cors::CorsConfig;
24
25/// Main Lambda MCP server
26///
27/// This server stores all configuration and can create Lambda handlers when needed.
28/// It mirrors the architecture of McpServer but is designed for Lambda deployment.
29#[allow(dead_code)]
30pub struct LambdaMcpServer {
31    /// Server implementation information
32    pub implementation: Implementation,
33    /// Server capabilities
34    pub capabilities: ServerCapabilities,
35    /// Registered tools
36    tools: HashMap<String, Arc<dyn McpTool>>,
37    /// Registered resources
38    resources: HashMap<String, Arc<dyn McpResource>>,
39    /// Registered prompts
40    prompts: HashMap<String, Arc<dyn McpPrompt>>,
41    /// Registered elicitations
42    elicitations: HashMap<String, Arc<dyn McpElicitation>>,
43    /// Registered sampling providers
44    sampling: HashMap<String, Arc<dyn McpSampling>>,
45    /// Registered completion providers
46    completions: HashMap<String, Arc<dyn McpCompletion>>,
47    /// Registered loggers
48    loggers: HashMap<String, Arc<dyn McpLogger>>,
49    /// Registered root providers
50    root_providers: HashMap<String, Arc<dyn McpRoot>>,
51    /// Registered notification providers
52    notifications: HashMap<String, Arc<dyn McpNotification>>,
53    /// Registered handlers
54    handlers: HashMap<String, Arc<dyn McpHandler>>,
55    /// Configured roots
56    roots: Vec<turul_mcp_protocol::roots::Root>,
57    /// Optional client instructions
58    instructions: Option<String>,
59    /// Session manager for state persistence
60    session_manager: Arc<SessionManager>,
61    /// Session storage backend (shared between SessionManager and handler)
62    session_storage: Arc<BoxedSessionStorage>,
63    /// Strict MCP lifecycle enforcement
64    strict_lifecycle: bool,
65    /// Server configuration
66    server_config: ServerConfig,
67    /// Enable SSE streaming
68    enable_sse: bool,
69    /// Stream configuration
70    stream_config: StreamConfig,
71    /// CORS configuration (if enabled)
72    #[cfg(feature = "cors")]
73    cors_config: Option<CorsConfig>,
74    /// Middleware stack for request/response interception
75    middleware_stack: turul_http_mcp_server::middleware::MiddlewareStack,
76    /// Custom route registry (e.g., .well-known endpoints)
77    route_registry: Arc<turul_http_mcp_server::RouteRegistry>,
78    /// Optional task runtime for MCP task support
79    task_runtime: Option<Arc<turul_mcp_server::TaskRuntime>>,
80    /// Stable fingerprint of the registered tool set for session versioning
81    tool_fingerprint: String,
82    /// Dynamic tool registry (only in Dynamic mode)
83    #[cfg(feature = "dynamic-tools")]
84    tool_registry: Option<Arc<turul_mcp_server::ToolRegistry>>,
85    /// Whether cross-instance coordination is enabled (explicit storage was provided)
86    #[cfg(feature = "dynamic-tools")]
87    coordination_enabled: bool,
88}
89
90impl LambdaMcpServer {
91    /// Create a new Lambda MCP server (use builder instead)
92    #[allow(clippy::too_many_arguments)]
93    pub(crate) fn new(
94        implementation: Implementation,
95        capabilities: ServerCapabilities,
96        tools: HashMap<String, Arc<dyn McpTool>>,
97        resources: HashMap<String, Arc<dyn McpResource>>,
98        prompts: HashMap<String, Arc<dyn McpPrompt>>,
99        elicitations: HashMap<String, Arc<dyn McpElicitation>>,
100        sampling: HashMap<String, Arc<dyn McpSampling>>,
101        completions: HashMap<String, Arc<dyn McpCompletion>>,
102        loggers: HashMap<String, Arc<dyn McpLogger>>,
103        root_providers: HashMap<String, Arc<dyn McpRoot>>,
104        notifications: HashMap<String, Arc<dyn McpNotification>>,
105        handlers: HashMap<String, Arc<dyn McpHandler>>,
106        roots: Vec<turul_mcp_protocol::roots::Root>,
107        instructions: Option<String>,
108        session_storage: Arc<BoxedSessionStorage>,
109        strict_lifecycle: bool,
110        server_config: ServerConfig,
111        enable_sse: bool,
112        stream_config: StreamConfig,
113        #[cfg(feature = "cors")] cors_config: Option<CorsConfig>,
114        middleware_stack: turul_http_mcp_server::middleware::MiddlewareStack,
115        route_registry: Arc<turul_http_mcp_server::RouteRegistry>,
116        task_runtime: Option<Arc<turul_mcp_server::TaskRuntime>>,
117        tool_fingerprint: String,
118        #[cfg(feature = "dynamic-tools")] dynamic_tools: bool,
119        #[cfg(feature = "dynamic-tools")]
120        server_state_storage: Option<Arc<dyn turul_mcp_server_state_storage::ServerStateStorage>>,
121    ) -> Self {
122        // Create session manager with server capabilities
123        let session_manager = Arc::new(SessionManager::with_storage_and_timeouts(
124            Arc::clone(&session_storage),
125            capabilities.clone(),
126            std::time::Duration::from_secs(30 * 60), // 30 minutes default
127            std::time::Duration::from_secs(60),      // 1 minute cleanup interval
128        ));
129
130        // Coordination enabled only for shared backends that can be accessed by multiple instances.
131        // InMemory and SQLite are local-only. Only PostgreSQL and DynamoDB are shared.
132        #[cfg(feature = "dynamic-tools")]
133        let coordination_enabled = server_state_storage
134            .as_ref()
135            .map(|s| matches!(s.backend_name(), "PostgreSQL" | "DynamoDB"))
136            .unwrap_or(false);
137
138        // Create ToolRegistry when dynamic mode is enabled
139        #[cfg(feature = "dynamic-tools")]
140        let tool_registry = if dynamic_tools {
141            let storage = server_state_storage.unwrap_or_else(|| {
142                Arc::new(turul_mcp_server_state_storage::InMemoryServerStateStorage::new())
143            });
144            Some(Arc::new(turul_mcp_server::ToolRegistry::new(
145                tools.clone(),
146                session_manager.clone(),
147                storage,
148            )))
149        } else {
150            None
151        };
152
153        Self {
154            implementation,
155            capabilities,
156            tools,
157            resources,
158            prompts,
159            elicitations,
160            sampling,
161            completions,
162            loggers,
163            root_providers,
164            notifications,
165            handlers,
166            roots,
167            instructions,
168            session_manager,
169            session_storage,
170            strict_lifecycle,
171            server_config,
172            enable_sse,
173            stream_config,
174            #[cfg(feature = "cors")]
175            cors_config,
176            middleware_stack,
177            route_registry,
178            task_runtime,
179            tool_fingerprint,
180            #[cfg(feature = "dynamic-tools")]
181            tool_registry,
182            #[cfg(feature = "dynamic-tools")]
183            coordination_enabled,
184        }
185    }
186
187    /// Get a reference to the server capabilities.
188    pub fn capabilities(&self) -> &ServerCapabilities {
189        &self.capabilities
190    }
191
192    /// Create a Lambda handler ready for use with Lambda runtime
193    ///
194    /// This is equivalent to McpServer::run_http() but creates a handler instead of running a server.
195    pub async fn handler(&self) -> Result<LambdaMcpHandler> {
196        info!(
197            "Creating Lambda MCP handler: {} v{}",
198            self.implementation.name, self.implementation.version
199        );
200        info!("Session management: enabled with automatic cleanup");
201
202        if self.enable_sse {
203            info!("SSE notifications: enabled for Lambda responses");
204
205            // ⚠️ GUARDRAIL: SSE enabled without streaming feature
206            #[cfg(not(feature = "streaming"))]
207            {
208                use tracing::warn;
209                warn!("⚠️  SSE is enabled but 'streaming' feature is not available!");
210                warn!(
211                    "   For real SSE streaming, use handle_streaming() with run_with_streaming_response"
212                );
213                warn!(
214                    "   Current handle() method will return SSE snapshots, not real-time streams"
215                );
216                warn!("   To enable streaming: add 'streaming' feature to turul-mcp-aws-lambda");
217            }
218        }
219
220        // Start session cleanup task (same as MCP server)
221        let _cleanup_task = self.session_manager.clone().start_cleanup_task();
222
223        // Sync tool registry with shared storage on startup (coordination mode only)
224        #[cfg(feature = "dynamic-tools")]
225        if self.coordination_enabled {
226            if let Some(ref registry) = self.tool_registry {
227                use tracing::warn;
228                match registry.sync_from_storage().await {
229                    Ok(_) => {
230                        info!("Dynamic: synced tool registry with shared storage");
231                    }
232                    Err(e) => {
233                        warn!(error = %e, "Dynamic: failed to sync with shared storage on cold start");
234                    }
235                }
236            }
237        }
238
239        // Cold-start recovery: handler() is called once per Lambda cold start from main().
240        // The returned LambdaMcpHandler is Clone'd for each request — recovery runs exactly once.
241        if let Some(ref runtime) = self.task_runtime {
242            match runtime.recover_stuck_tasks().await {
243                Ok(recovered) if !recovered.is_empty() => {
244                    info!(
245                        count = recovered.len(),
246                        "Recovered stuck tasks from previous invocations"
247                    );
248                }
249                Err(e) => {
250                    use tracing::warn;
251                    warn!(error = %e, "Failed to recover stuck tasks on startup");
252                }
253                _ => {}
254            }
255        }
256
257        // Create stream manager for SSE
258        let stream_manager = Arc::new(StreamManager::with_config(
259            self.session_storage.clone(),
260            self.stream_config.clone(),
261        ));
262
263        // Set up SSE event bridge between SessionManager and StreamManager.
264        // This ensures that events from ToolRegistry.broadcast_notification()
265        // (which go through SessionManager) are forwarded to StreamManager where
266        // registered POST SSE connections can receive them.
267        {
268            let bridge_stream_manager = Arc::clone(&stream_manager);
269            let mut global_events = self.session_manager.subscribe_all_session_events();
270
271            tokio::spawn(async move {
272                debug!("Lambda SSE Event Bridge: started");
273
274                while let Ok((session_id, event)) = global_events.recv().await {
275                    if let turul_mcp_server::session::SessionEvent::Custom {
276                        event_type,
277                        data,
278                    } = event
279                    {
280                        if let Err(e) = bridge_stream_manager
281                            .broadcast_to_session(&session_id, event_type, data)
282                            .await
283                        {
284                            debug!(
285                                "Lambda SSE Bridge: broadcast to session {} failed: {} (normal if no active connections)",
286                                session_id, e
287                            );
288                        }
289                    }
290                }
291
292                debug!("Lambda SSE Event Bridge: stopped");
293            });
294        }
295
296        // Create JSON-RPC dispatcher
297        use turul_mcp_json_rpc_server::JsonRpcDispatcher;
298        let mut dispatcher = JsonRpcDispatcher::new();
299
300        // Create session-aware initialize handler (reuse MCP server handler)
301        use turul_mcp_server::SessionAwareInitializeHandler;
302        let init_handler = SessionAwareInitializeHandler::new(
303            self.implementation.clone(),
304            self.capabilities.clone(),
305            self.instructions.clone(),
306            self.session_manager.clone(),
307            self.strict_lifecycle,
308            self.tool_fingerprint.clone(),
309        );
310        dispatcher.register_method("initialize".to_string(), init_handler);
311
312        // Create session-aware tools/list handler (reuse MCP server handler)
313        use turul_mcp_server::ListToolsHandler;
314        let mut list_handler = ListToolsHandler::new_with_session_manager(
315            self.tools.clone(),
316            self.session_manager.clone(),
317            self.strict_lifecycle,
318            self.task_runtime.is_some(),
319        );
320        #[cfg(feature = "dynamic-tools")]
321        if let Some(ref registry) = self.tool_registry {
322            list_handler = list_handler.with_tool_registry(Arc::clone(registry));
323        }
324        dispatcher.register_method("tools/list".to_string(), list_handler);
325
326        // Create session-aware tool handler for tools/call (reuse MCP server handler)
327        use turul_mcp_server::SessionAwareToolHandler;
328        let mut tool_handler = SessionAwareToolHandler::new(
329            self.tools.clone(),
330            self.session_manager.clone(),
331            self.strict_lifecycle,
332        );
333        if let Some(ref runtime) = self.task_runtime {
334            tool_handler = tool_handler.with_task_runtime(Arc::clone(runtime));
335        }
336        #[cfg(feature = "dynamic-tools")]
337        if let Some(ref registry) = self.tool_registry {
338            tool_handler = tool_handler.with_tool_registry(Arc::clone(registry));
339        }
340        dispatcher.register_method("tools/call".to_string(), tool_handler);
341
342        // Register all MCP handlers with session awareness (reuse MCP server bridge)
343        use turul_mcp_server::SessionAwareMcpHandlerBridge;
344        for (method, handler) in &self.handlers {
345            let bridge_handler = SessionAwareMcpHandlerBridge::new(
346                handler.clone(),
347                self.session_manager.clone(),
348                self.strict_lifecycle,
349            );
350            dispatcher.register_method(method.clone(), bridge_handler);
351        }
352
353        // Register notifications/initialized handler — required for strict lifecycle.
354        // Without this, clients can never complete the MCP handshake.
355        use turul_mcp_server::handlers::InitializedNotificationHandler;
356        let initialized_handler = InitializedNotificationHandler::new(self.session_manager.clone());
357        let initialized_bridge = SessionAwareMcpHandlerBridge::new(
358            Arc::new(initialized_handler),
359            self.session_manager.clone(),
360            self.strict_lifecycle,
361        );
362        dispatcher.register_method("notifications/initialized".to_string(), initialized_bridge);
363
364        // Create the Lambda handler with all components and middleware
365        let middleware_stack = Arc::new(self.middleware_stack.clone());
366
367        let handler = LambdaMcpHandler::with_middleware_and_fingerprint(
368            self.server_config.clone(),
369            Arc::new(dispatcher),
370            self.session_storage.clone(),
371            stream_manager,
372            self.stream_config.clone(),
373            self.capabilities.clone(),
374            middleware_stack,
375            self.enable_sse,
376            Arc::clone(&self.route_registry),
377            Some(self.tool_fingerprint.clone()),
378        );
379
380        #[cfg(feature = "dynamic-tools")]
381        let handler = if let Some(ref registry) = self.tool_registry {
382            handler.with_tool_registry(Arc::clone(registry))
383        } else {
384            handler
385        };
386
387        Ok(handler)
388    }
389
390    /// Get information about the session storage backend
391    pub fn session_storage_info(&self) -> &str {
392        "Session storage configured"
393    }
394}
turul_mcp_aws_lambda/server.rs

turul_mcp_aws_lambda/
server.rs
