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::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}
77
78impl LambdaMcpServer {
79    /// Create a new Lambda MCP server (use builder instead)
80    #[allow(clippy::too_many_arguments)]
81    pub(crate) fn new(
82        implementation: Implementation,
83        capabilities: ServerCapabilities,
84        tools: HashMap<String, Arc<dyn McpTool>>,
85        resources: HashMap<String, Arc<dyn McpResource>>,
86        prompts: HashMap<String, Arc<dyn McpPrompt>>,
87        elicitations: HashMap<String, Arc<dyn McpElicitation>>,
88        sampling: HashMap<String, Arc<dyn McpSampling>>,
89        completions: HashMap<String, Arc<dyn McpCompletion>>,
90        loggers: HashMap<String, Arc<dyn McpLogger>>,
91        root_providers: HashMap<String, Arc<dyn McpRoot>>,
92        notifications: HashMap<String, Arc<dyn McpNotification>>,
93        handlers: HashMap<String, Arc<dyn McpHandler>>,
94        roots: Vec<turul_mcp_protocol::roots::Root>,
95        instructions: Option<String>,
96        session_storage: Arc<BoxedSessionStorage>,
97        strict_lifecycle: bool,
98        server_config: ServerConfig,
99        enable_sse: bool,
100        stream_config: StreamConfig,
101        #[cfg(feature = "cors")] cors_config: Option<CorsConfig>,
102        middleware_stack: turul_http_mcp_server::middleware::MiddlewareStack,
103    ) -> Self {
104        // Create session manager with server capabilities
105        let session_manager = Arc::new(SessionManager::with_storage_and_timeouts(
106            Arc::clone(&session_storage),
107            capabilities.clone(),
108            std::time::Duration::from_secs(30 * 60), // 30 minutes default
109            std::time::Duration::from_secs(60),      // 1 minute cleanup interval
110        ));
111
112        Self {
113            implementation,
114            capabilities,
115            tools,
116            resources,
117            prompts,
118            elicitations,
119            sampling,
120            completions,
121            loggers,
122            root_providers,
123            notifications,
124            handlers,
125            roots,
126            instructions,
127            session_manager,
128            session_storage,
129            strict_lifecycle,
130            server_config,
131            enable_sse,
132            stream_config,
133            #[cfg(feature = "cors")]
134            cors_config,
135            middleware_stack,
136        }
137    }
138
139    /// Create a Lambda handler ready for use with Lambda runtime
140    ///
141    /// This is equivalent to McpServer::run_http() but creates a handler instead of running a server.
142    pub async fn handler(&self) -> Result<LambdaMcpHandler> {
143        info!(
144            "Creating Lambda MCP handler: {} v{}",
145            self.implementation.name, self.implementation.version
146        );
147        info!("Session management: enabled with automatic cleanup");
148
149        if self.enable_sse {
150            info!("SSE notifications: enabled for Lambda responses");
151
152            // ⚠️ GUARDRAIL: SSE enabled without streaming feature
153            #[cfg(not(feature = "streaming"))]
154            {
155                use tracing::warn;
156                warn!("⚠️  SSE is enabled but 'streaming' feature is not available!");
157                warn!(
158                    "   For real SSE streaming, use handle_streaming() with run_with_streaming_response"
159                );
160                warn!(
161                    "   Current handle() method will return SSE snapshots, not real-time streams"
162                );
163                warn!("   To enable streaming: add 'streaming' feature to turul-mcp-aws-lambda");
164            }
165        }
166
167        // Start session cleanup task (same as MCP server)
168        let _cleanup_task = self.session_manager.clone().start_cleanup_task();
169
170        // Create stream manager for SSE
171        let stream_manager = Arc::new(StreamManager::with_config(
172            self.session_storage.clone(),
173            self.stream_config.clone(),
174        ));
175
176        // Create JSON-RPC dispatcher
177        use turul_mcp_json_rpc_server::JsonRpcDispatcher;
178        let mut dispatcher = JsonRpcDispatcher::new();
179
180        // Create session-aware initialize handler (reuse MCP server handler)
181        use turul_mcp_server::SessionAwareInitializeHandler;
182        let init_handler = SessionAwareInitializeHandler::new(
183            self.implementation.clone(),
184            self.capabilities.clone(),
185            self.instructions.clone(),
186            self.session_manager.clone(),
187            self.strict_lifecycle,
188        );
189        dispatcher.register_method("initialize".to_string(), init_handler);
190
191        // Create tools/list handler (reuse MCP server handler)
192        use turul_mcp_server::ListToolsHandler;
193        let list_handler = ListToolsHandler::new(self.tools.clone());
194        dispatcher.register_method("tools/list".to_string(), list_handler);
195
196        // Create session-aware tool handler for tools/call (reuse MCP server handler)
197        use turul_mcp_server::SessionAwareToolHandler;
198        let tool_handler = SessionAwareToolHandler::new(
199            self.tools.clone(),
200            self.session_manager.clone(),
201            self.strict_lifecycle,
202        );
203        dispatcher.register_method("tools/call".to_string(), tool_handler);
204
205        // Register all MCP handlers with session awareness (reuse MCP server bridge)
206        use turul_mcp_server::SessionAwareMcpHandlerBridge;
207        for (method, handler) in &self.handlers {
208            let bridge_handler = SessionAwareMcpHandlerBridge::new(
209                handler.clone(),
210                self.session_manager.clone(),
211                self.strict_lifecycle,
212            );
213            dispatcher.register_method(method.clone(), bridge_handler);
214        }
215
216        // Create the Lambda handler with all components and middleware
217        let middleware_stack = Arc::new(self.middleware_stack.clone());
218
219        Ok(LambdaMcpHandler::with_middleware(
220            self.server_config.clone(),
221            Arc::new(dispatcher),
222            self.session_storage.clone(),
223            stream_manager,
224            self.stream_config.clone(),
225            self.capabilities.clone(),
226            middleware_stack,
227            self.enable_sse,
228        ))
229    }
230
231    /// Get information about the session storage backend
232    pub fn session_storage_info(&self) -> &str {
233        "Session storage configured"
234    }
235}