├── .gitignore ├── Dockerfile ├── LICENSE ├── README.md ├── architecture.png ├── build_and_push.sh ├── deploy.py ├── docs ├── audience-benefits.md ├── conclusion.md ├── linkedin-post.md └── technical-analysis.md ├── package-lock.json ├── package.json ├── payload.json ├── pyproject.toml ├── src ├── auth │ ├── auth.py │ ├── cognito.ts │ ├── index.ts │ ├── lambdaAuth.ts │ └── oauthMetadata.ts ├── client.ts ├── logging.js ├── mcpClientUtils.ts ├── mcpTransport.ts └── server.ts ├── tsconfig.json └── uv.lock /.gitignore: -------------------------------------------------------------------------------- 1 | # Logs 2 | logs 3 | *.log 4 | npm-debug.log* 5 | yarn-debug.log* 6 | yarn-error.log* 7 | lerna-debug.log* 8 | .pnpm-debug.log* 9 | 10 | # Diagnostic reports (https://nodejs.org/api/report.html) 11 | report.[0-9]*.[0-9]*.[0-9]*.[0-9]*.json 12 | 13 | # Runtime data 14 | pids 15 | *.pid 16 | *.seed 17 | *.pid.lock 18 | 19 | # Directory for instrumented libs generated by jscoverage/JSCover 20 | lib-cov 21 | 22 | # Coverage directory used by tools like istanbul 23 | coverage 24 | *.lcov 25 | 26 | # nyc test coverage 27 | .nyc_output 28 | 29 | # Grunt intermediate storage (https://gruntjs.com/creating-plugins#storing-task-files) 30 | .grunt 31 | 32 | # Bower dependency directory (https://bower.io/) 33 | bower_components 34 | 35 | # node-waf configuration 36 | .lock-wscript 37 | 38 | # Compiled binary addons (https://nodejs.org/api/addons.html) 39 | build/Release 40 | 41 | # Dependency directories 42 | node_modules/ 43 | jspm_packages/ 44 | 45 | # Snowpack dependency directory (https://snowpack.dev/) 46 | web_modules/ 47 | 48 | # TypeScript cache 49 | *.tsbuildinfo 50 | 51 | # Optional npm cache directory 52 | .npm 53 | 54 | # Optional eslint cache 55 | .eslintcache 56 | 57 | # Optional stylelint cache 58 | .stylelintcache 59 | 60 | # Microbundle cache 61 | .rpt2_cache/ 62 | .rts2_cache_cjs/ 63 | .rts2_cache_es/ 64 | .rts2_cache_umd/ 65 | 66 | # Optional REPL history 67 | .node_repl_history 68 | 69 | # Output of 'npm pack' 70 | *.tgz 71 | 72 | # Yarn Integrity file 73 | .yarn-integrity 74 | 75 | # dotenv environment variable files 76 | .env 77 | .env.development.local 78 | .env.test.local 79 | .env.production.local 80 | .env.local 81 | 82 | # parcel-bundler cache (https://parceljs.org/) 83 | .cache 84 | .parcel-cache 85 | 86 | # Next.js build output 87 | .next 88 | out 89 | 90 | # Nuxt.js build / generate output 91 | .nuxt 92 | dist 93 | 94 | # Gatsby files 95 | .cache/ 96 | # Comment in the public line in if your project uses Gatsby and not Next.js 97 | # https://nextjs.org/blog/next-9-1#public-directory-support 98 | # public 99 | 100 | # vuepress build output 101 | .vuepress/dist 102 | 103 | # vuepress v2.x temp and cache directory 104 | .temp 105 | .cache 106 | 107 | # vitepress build output 108 | **/.vitepress/dist 109 | 110 | # vitepress cache directory 111 | **/.vitepress/cache 112 | 113 | # Docusaurus cache and generated files 114 | .docusaurus 115 | 116 | # Serverless directories 117 | .serverless/ 118 | 119 | # FuseBox cache 120 | .fusebox/ 121 | 122 | # DynamoDB Local files 123 | .dynamodb/ 124 | 125 | # TernJS port file 126 | .tern-port 127 | 128 | # Stores VSCode versions used for testing VSCode extensions 129 | .vscode-test 130 | 131 | # yarn v2 132 | .yarn/cache 133 | .yarn/unplugged 134 | .yarn/build-state.yml 135 | .yarn/install-state.gz 136 | .pnp.* 137 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | # Use Node.js 18 as the base image for Lambda 2 | FROM public.ecr.aws/lambda/nodejs:18 3 | 4 | # Set working directory 5 | WORKDIR ${LAMBDA_TASK_ROOT} 6 | 7 | # Copy package files and TypeScript config 8 | COPY package*.json ./ 9 | COPY tsconfig.json ./ 10 | 11 | # Install dependencies - including dev dependencies for building 12 | RUN npm install 13 | 14 | # Copy source code and tests 15 | COPY src/ ./src/ 16 | 17 | # Build TypeScript files 18 | RUN npm run build && ls -la dist/ 19 | 20 | # The Lambda handler 21 | CMD [ "dist/server.handler" ] -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Streamable MCP Server on AWS Lambda with Multiple Authorization Options 2 | 3 | This project implements a Model Context Protocol (MCP) server as a containerized application on AWS Lambda, accessible via Amazon API Gateway. It showcases the [`Streamable-HTTP`](https://modelcontextprotocol.io/specification/2025-03-26/basic/transports#streamable-http) transport along with multiple authorization options: 4 | 5 | 1. **OAuth 2.1 Authorization** through AWS Cognito 6 | 2. **Lambda Authorizer** for simpler token-based authorization 7 | 8 | The MCP server in this repo: 9 | - Uses session management via the `Mcp-Session-id` header 10 | - Supports both OAuth 2.1 and Lambda authorizer methods 11 | - Provides tools to analyze Amazon Bedrock usage 12 | 13 | Both server and client are written in TypeScript, with the server deployed as a container on Lambda. 14 | 15 | >Note: As of 4/22/2025 Lambda supports HTTP Streaming for Node.js managed runtime ([documentation](https://docs.aws.amazon.com/lambda/latest/dg/configuration-response-streaming.html)). This implementation leverages that capability to create a fully serverless MCP deployment. 16 | 17 | ## Key Features 18 | 19 | - **Standards Compliance**: Implements both Streamable-HTTP transport and OAuth 2.1 authorization specs 20 | - **Serverless Deployment**: Runs on AWS Lambda and API Gateway for scalability 21 | - **Multiple Authorization Options**: Supports both OAuth 2.1 and Lambda authorizer methods 22 | - **Secure Authentication**: Uses AWS Cognito for OAuth 2.1 authentication 23 | - **Discovery Support**: Implements OAuth discovery flow per RFC9728 24 | - **Analytics Tool**: Provides Bedrock usage analysis tool 25 | 26 | ## Architecture 27 | 28 | ![Architecture Diagram](architecture.png) 29 | 30 | *Image credit: [https://github.com/aal80/simple-mcp-server-on-lambda](https://github.com/aal80/simple-mcp-server-on-lambda)* 31 | 32 | The architecture consists of: 33 | 1. **API Gateway**: HTTP API endpoint with streaming support 34 | 2. **Lambda Function**: Containerized MCP server with OAuth authorization 35 | 3. **ECR**: Container storage 36 | 4. **Cognito**: OAuth 2.1 authentication provider 37 | 5. **CloudWatch Logs**: Server and Bedrock usage logging 38 | 6. **Bedrock**: The foundation model service 39 | 40 | ## Authorization Options 41 | 42 | This project supports two authorization methods: 43 | 44 | ### 1. OAuth 2.1 Authorization 45 | 46 | Uses AWS Cognito as the OAuth provider with full OAuth 2.1 compliance: 47 | 48 | - Standards-compliant authorization with Cognito 49 | - PKCE support for enhanced security 50 | - JWT verification via JWKS 51 | - Token refresh with rotation 52 | - RFC9728-compliant OAuth discovery 53 | 54 | To deploy with OAuth authorization: 55 | 56 | ```bash 57 | python deploy.py \ 58 | --function-name mcp-server \ 59 | --role-arn \ 60 | --region us-east-1 \ 61 | --memory 2048 \ 62 | --timeout 300 \ 63 | --api-gateway \ 64 | --api-name mcp-server-api \ 65 | --stage-name prod \ 66 | --auth-method oauth \ 67 | --cognito-user-pool-id \ 68 | --cognito-domain \ 69 | --cognito-client-ids 70 | ``` 71 | 72 | ### 2. Lambda Authorizer 73 | 74 | Uses a separate Lambda function to authorize requests: 75 | 76 | - Simple token-based authorization 77 | - API Gateway integration 78 | - Accepts any properly formatted Bearer token 79 | 80 | To deploy with Lambda authorizer: 81 | 82 | ```bash 83 | python deploy.py \ 84 | --function-name mcp-server \ 85 | --role-arn \ 86 | --region us-east-1 \ 87 | --memory 2048 \ 88 | --timeout 300 \ 89 | --api-gateway \ 90 | --api-name mcp-server-api \ 91 | --stage-name prod \ 92 | --auth-method lambda \ 93 | --lambda-authorizer-name mcp-authorizer 94 | ``` 95 | 96 | ## Prerequisites 97 | 98 | 1. **Node.js & npm:** v18.x or later 99 | 2. **Python 3.11:** For the deployment script 100 | 3. **Docker:** For container building 101 | 4. **AWS Account & CLI:** With appropriate credentials 102 | 5. **CloudWatch Setup:** [Model invocation logs](https://docs.aws.amazon.com/bedrock/latest/userguide/model-invocation-logging.html#setup-cloudwatch-logs-destination) configured 103 | 6. **IAM Permissions:** 104 | - CloudWatch Logs read access ([CloudWatch Logs Read-Only Access](https://docs.aws.amazon.com/aws-managed-policy/latest/reference/CloudWatchLogsReadOnlyAccess.html)) 105 | - Cost Explorer read access ([Billing policy examples](https://docs.aws.amazon.com/awsaccountbilling/latest/aboutv2/billing-example-policies.html)) 106 | 107 | ## Quick Start Guide 108 | 109 | ### Local Development 110 | 111 | 1. **Install dependencies:** 112 | ```bash 113 | npm install 114 | ``` 115 | 116 | 2. **Start the server locally:** 117 | ```bash 118 | npx tsx src/server.ts 119 | ``` 120 | The server will start on `http://localhost:3000` by default. 121 | 122 | 3. **Run the client in another terminal:** 123 | ```bash 124 | npx tsx src/client.ts 125 | ``` 126 | 127 | 4. **Optional authentication:** 128 | ```bash 129 | # In the client, authenticate with: 130 | > auth login 131 | ``` 132 | 133 | ### AWS Deployment 134 | 135 | 1. **Install Python dependencies:** 136 | ```bash 137 | # Install uv 138 | curl -LsSf https://astral.sh/uv/install.sh | sh 139 | export PATH="$HOME/.local/bin:$PATH" 140 | 141 | # Create a virtual environment and install dependencies 142 | uv venv --python 3.11 && source .venv/bin/activate && uv pip install --requirement pyproject.toml 143 | ``` 144 | 145 | 2. **Configure AWS Cognito:** 146 | - Create a User Pool in AWS Cognito Console 147 | - Create an App Client with authorization code grant flow 148 | - Configure callback URL (typically http://localhost:8000/callback) 149 | - Add allowed scopes: openid, profile 150 | 151 | 3. **Set environment variables:** 152 | ```bash 153 | # Server environment variables 154 | export COGNITO_REGION="us-east-1" 155 | export COGNITO_USER_POOL_ID="" 156 | export COGNITO_ALLOWED_CLIENT_IDS="" 157 | 158 | # COGNITO_DOMAIN can be either full domain or just the prefix 159 | # Full domain example: 160 | export COGNITO_DOMAIN=".auth.us-east-1.amazoncognito.com" 161 | # Or just the prefix (domain part will be built automatically): 162 | export COGNITO_DOMAIN="" 163 | 164 | # Client environment variables 165 | export OAUTH_CLIENT_ID="" 166 | export OAUTH_REDIRECT_URI="http://localhost:8000/callback" 167 | ``` 168 | 169 | 4. **Run deployment script:** 170 | ```bash 171 | python deploy.py \ 172 | --function-name mcp-server \ 173 | --role-arn \ 174 | --region us-east-1 \ 175 | --memory 2048 \ 176 | --timeout 300 \ 177 | --api-gateway \ 178 | --api-name mcp-server-api \ 179 | --stage-name prod \ 180 | --cognito-user-pool-id \ 181 | --cognito-domain \ 182 | --cognito-client-ids 183 | ``` 184 | 185 | 5. **Connect client to deployed server:** 186 | ```bash 187 | export MCP_SERVER_URL="https://.execute-api..amazonaws.com/prod/mcp" 188 | 189 | # For OAuth 2.1 authentication 190 | export AUTH_METHOD=oauth 191 | export OAUTH_CLIENT_ID="" 192 | export OAUTH_REDIRECT_URI="http://localhost:8000/callback" 193 | 194 | # Or for Lambda authorizer 195 | export AUTH_METHOD=lambda 196 | 197 | npx tsx src/client.ts 198 | 199 | # At the client prompt (for OAuth) 200 | > auth login 201 | > connect 202 | 203 | # Or for Lambda authorizer 204 | > connect 205 | ``` 206 | 207 | ## MCP Protocol Implementation 208 | 209 | This project implements two key MCP specifications: 210 | 211 | 1. **Streamable-HTTP Transport** - Enables stateful connections over HTTP with session management 212 | 2. **OAuth 2.1 Authorization** - Secures API access with standards-compliant authorization 213 | 214 | ### Streamable-HTTP Sequence Diagram 215 | 216 | ```mermaid 217 | sequenceDiagram 218 | participant Client 219 | participant Server 220 | 221 | Client->>Server: POST /mcp - initialize request 222 | Server->>Client: 200 OK with Mcp-Session-Id 223 | Note over Client,Server: Client stores session ID 224 | 225 | Client->>Server: POST /mcp - tools/list (with session ID) 226 | Server->>Client: 200 OK with tools list 227 | 228 | Client->>Server: POST /mcp - tools/call (with session ID) 229 | Server->>Client: 200 OK with streaming content 230 | 231 | Client->>Server: DELETE /mcp (with session ID) 232 | Server->>Client: 200 OK - Session terminated 233 | ``` 234 | 235 | ### OAuth 2.1 Authorization Flow 236 | 237 | ```mermaid 238 | sequenceDiagram 239 | participant Client 240 | participant MCP Server 241 | participant Cognito 242 | 243 | Client->>MCP Server: Initial Request (no token) 244 | MCP Server->>Client: 401 Unauthorized + WWW-Authenticate header 245 | Note over Client: Extracts resource_metadata_uri 246 | 247 | Client->>MCP Server: GET /.well-known/oauth-protected-resource 248 | MCP Server->>Client: Resource Metadata (auth server info) 249 | 250 | Client->>Cognito: GET /.well-known/openid-configuration 251 | Cognito->>Client: OAuth/OIDC Endpoints 252 | 253 | Note over Client: Generate PKCE parameters 254 | 255 | Client->>Cognito: Authenticate (with PKCE) 256 | Cognito->>Client: Access + Refresh Tokens 257 | 258 | Client->>MCP Server: Request with Bearer Token 259 | MCP Server->>Cognito: Verify token 260 | Cognito->>MCP Server: Token validation 261 | 262 | alt Token Valid 263 | MCP Server->>Client: 200 Success + Session ID 264 | else Invalid Token 265 | MCP Server->>Client: 401 Unauthorized 266 | end 267 | 268 | Note over Client: When token expires 269 | Client->>Cognito: Refresh token request 270 | Cognito->>Client: New access token 271 | ``` 272 | 273 | ## Authentication Setup 274 | 275 | ### Configuring AWS Cognito 276 | 277 | 1. **Create a Cognito User Pool** in the AWS Console 278 | 2. **Create an App Client:** 279 | - Set Authorization code grant flow 280 | - Add allowed scopes: openid, profile 281 | - Configure callback URL as http://localhost:8000/callback 282 | 283 | ### Client Authentication Commands 284 | 285 | ```bash 286 | # Login with OAuth 2.1 (with PKCE) 287 | > auth login 288 | 289 | # Check authentication status 290 | > auth status 291 | 292 | # Manually refresh token 293 | > auth refresh 294 | 295 | # Log out 296 | > auth logout 297 | ``` 298 | 299 | ### Authentication Features 300 | 301 | - **OAuth 2.1 Compliance**: Standards-compliant authorization 302 | - **PKCE Support**: Enhanced security for public clients 303 | - **JWT Verification**: Token verification via JWKS 304 | - **Token Refresh**: Automatic refresh with rotation 305 | - **Discovery**: RFC9728-compliant OAuth discovery with Cognito-specific adaptations 306 | - **Graceful Fallbacks**: Robust error handling with multiple discovery mechanisms 307 | 308 | ## Using the Bedrock Report Tool 309 | 310 | ```bash 311 | # Default parameters (us-east-1, standard log group, 1 day) 312 | > bedrock-report 313 | 314 | # Custom parameters 315 | > bedrock-report us-east-1 /aws/bedrock/modelinvocations 7 123456789012 316 | ``` 317 | 318 | This will produce a detailed report showing: 319 | - Total requests, tokens, and usage 320 | - Daily breakdowns 321 | - Regional summaries 322 | - Model-specific usage 323 | - User/role summaries 324 | 325 | Example output: 326 | ``` 327 | Tool result: 328 | Bedrock Daily Usage Report (Past 7 days - us-east-1) 329 | Total Requests: 13060 330 | Total Input Tokens: 2992387 331 | Total Completion Tokens: 254124 332 | Total Tokens: 3246511 333 | 334 | --- Daily Totals --- 335 | 2025-04-06: Requests=8330, Input=1818253, Completion=171794, Total=1990047 336 | 2025-04-07: Requests=4669, Input=936299, Completion=71744, Total=1008043 337 | ... 338 | 339 | --- Model Summary --- 340 | nova-lite-v1:0: Requests=93, Input=177416, Completion=30331, Total=207747 341 | titan-embed-text-v1: Requests=62, Input=845, Completion=0, Total=845 342 | ... 343 | ``` 344 | 345 | ## Available MCP Tools 346 | 347 | The server implements these tools: 348 | 349 | 1. **greet**: A simple greeting tool that returns a personalized message 350 | 2. **multi-greet**: A tool that sends multiple greetings with delays 351 | 3. **get_bedrock_usage_report**: A comprehensive Bedrock usage analysis tool 352 | 353 | ## Client Command Reference 354 | 355 | | Command | Description | 356 | |---------|-------------| 357 | | `connect [url]` | Connect to MCP server (default or specified URL) | 358 | | `disconnect` | Disconnect from server | 359 | | `terminate-session` | Terminate the current session | 360 | | `reconnect` | Reconnect to the server | 361 | | `list-tools` | List available tools on the server | 362 | | `call-tool [args]` | Call a tool with optional JSON arguments | 363 | | `greet [name]` | Call the greet tool with optional name | 364 | | `multi-greet [name]` | Call the multi-greet tool with notifications | 365 | | `list-resources` | List available resources | 366 | | `bedrock-report [region] [log_group] [days] [account_id]` | Get Bedrock usage report | 367 | | `auth login` | Authenticate with the MCP server using OAuth 2.1 | 368 | | `auth logout` | Clear the stored token | 369 | | `auth status` | Show current authentication status | 370 | | `auth refresh` | Force refresh the access token | 371 | | `set-auth-method ` | Set authorization method (oauth, lambda, or auto) | 372 | | `debug on off` | Enable or disable debug logging | 373 | | `help` | Show help information | 374 | | `quit` | Exit the program | 375 | 376 | ## Advanced Usage 377 | 378 | ### Direct API Calls with curl 379 | 380 | ```bash 381 | # Replace with your access token if using authentication 382 | curl -XPOST "https://.execute-api..amazonaws.com/prod/mcp" \ 383 | -H "Content-Type: application/json" \ 384 | -H "Accept: application/json, text/event-stream" \ 385 | -H "Authorization: Bearer " \ 386 | -d '{ 387 | "jsonrpc": "2.0", 388 | "method": "initialize", 389 | "params": { 390 | "clientInfo": { "name": "curl-client", "version": "1.0" }, 391 | "protocolVersion": "2025-03-26", 392 | "capabilities": {} 393 | }, 394 | "id": "init-1" 395 | }' | cat 396 | ``` 397 | 398 | ### Environment Variables 399 | 400 | #### Server Variables 401 | - `COGNITO_REGION`: AWS region where Cognito is deployed 402 | - `COGNITO_USER_POOL_ID`: ID of the Cognito user pool 403 | - `COGNITO_ALLOWED_CLIENT_IDS`: Comma-separated list of allowed client IDs 404 | - `COGNITO_DOMAIN`: Cognito domain for the user pool (can be full domain like "my-domain.auth.region.amazoncognito.com" or just the prefix like "my-domain") 405 | 406 | #### Client Variables 407 | - `MCP_SERVER_URL`: URL of the MCP server 408 | - `AUTH_METHOD`: Authorization method to use (`oauth`, `lambda`, or `auto`) 409 | - `OAUTH_CLIENT_ID`: Cognito app client ID 410 | - `OAUTH_REDIRECT_URI`: Redirect URI for OAuth flow 411 | - `MCP_CLIENT_DEBUG`: Set to 'true' to enable debug logging 412 | - `ACCESS_TOKEN`: Directly provide an access token (optional) 413 | 414 | ## Troubleshooting 415 | 416 | ### Common Issues 417 | 418 | 1. **Authentication Failures**: 419 | - Ensure Cognito user pool and client IDs are correctly configured 420 | - Verify the user credentials being used 421 | - Try `auth refresh` or `auth login` to obtain fresh tokens 422 | - Check Cognito domain format - either provide the full domain (e.g., "my-domain.auth.us-east-1.amazoncognito.com") or just the prefix (e.g., "my-domain") 423 | - If you see DNS errors like "getaddrinfo ENOTFOUND", check that your domain format is correct; for user pool IDs with format "region_ID", the correct domain prefix is "region-id" (lowercase with hyphen) 424 | - Note that Cognito uses a different OpenID discovery URL format than standard OAuth providers; this implementation handles this automatically by trying both formats 425 | 426 | 2. **Connection Problems**: 427 | - Verify the server URL is correct with proper `/mcp` suffix 428 | - Check network connectivity to the API Gateway endpoint 429 | - Ensure AWS region is correctly specified in environment variables 430 | 431 | 3. **Permission Errors**: 432 | - Verify Lambda has appropriate IAM permissions for CloudWatch Logs access 433 | - Check that the user/role has Cost Explorer permissions 434 | 435 | 4. **Deployment Issues**: 436 | - Ensure Docker is running and properly configured for `linux/amd64` builds 437 | - Check that AWS CLI has necessary permissions for ECR, Lambda, and API Gateway 438 | 439 | ### Debug Mode 440 | 441 | Enable debug logging to troubleshoot connection issues: 442 | 443 | ```bash 444 | # In client: 445 | > debug on 446 | 447 | # Or as environment variable: 448 | export MCP_CLIENT_DEBUG=true 449 | ``` 450 | 451 | ## Project Structure 452 | 453 | - `/src/server.ts` - MCP server implementation 454 | - `/src/client.ts` - MCP client with OAuth support 455 | - `/src/auth/` - Authorization components 456 | - `cognito.ts` - JWT verification logic 457 | - `oauthMetadata.ts` - OAuth resource metadata 458 | - `oauth2Client.ts` - OAuth client implementation 459 | 460 | ## Known Issues 461 | 462 | - Client occasionally fails to call tools, which can be resolved by reconnecting or restarting 463 | - Token refresh may sometimes require manual intervention with `auth refresh` 464 | - Large responses from the Bedrock usage tool may be truncated in certain environments 465 | 466 | ## Acknowledgments 467 | 468 | This project was developed with the support of: 469 | 470 | - [Model Context Protocol](https://modelcontextprotocol.io/) for the communication protocol 471 | - [AWS Lambda](https://aws.amazon.com/lambda/) for serverless computing 472 | - [Amazon API Gateway](https://aws.amazon.com/api-gateway/) for API management 473 | - [Amazon Bedrock](https://aws.amazon.com/bedrock/) for foundation models 474 | - [Amazon Cognito](https://aws.amazon.com/cognito/) for authentication 475 | - [Node.js](https://nodejs.org/) and [TypeScript](https://www.typescriptlang.org/) for implementation 476 | - [Express.js](https://expressjs.com/) for the web server framework 477 | 478 | ## License 479 | 480 | This project is licensed under the terms in the LICENSE file included in this repository. -------------------------------------------------------------------------------- /architecture.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aarora79/streamable-mcp-serverless/927793e6b14e196eb0d466160b3f0d1ec39e84b0/architecture.png -------------------------------------------------------------------------------- /build_and_push.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Exit on error 4 | set -e 5 | 6 | # Configuration 7 | AWS_REGION="us-east-1" 8 | ECR_REPO_NAME="understand-bedrock-spend-mcp-server" 9 | AWS_ACCOUNT_ID=$(aws sts get-caller-identity --query Account --output text) 10 | ECR_REPO_URI="$AWS_ACCOUNT_ID.dkr.ecr.$AWS_REGION.amazonaws.com/$ECR_REPO_NAME" 11 | 12 | echo "🔐 Logging in to Amazon ECR..." 13 | aws ecr get-login-password --region $AWS_REGION | docker login --username AWS --password-stdin "$AWS_ACCOUNT_ID.dkr.ecr.$AWS_REGION.amazonaws.com" 14 | 15 | # Create repository if it doesn't exist 16 | echo "📦 Creating ECR repository if it doesn't exist..." 17 | aws ecr describe-repositories --repository-names "$ECR_REPO_NAME" --region "$AWS_REGION" || \ 18 | aws ecr create-repository --repository-name "$ECR_REPO_NAME" --region "$AWS_REGION" 19 | 20 | # Build the Docker image using buildx for x86_64 architecture 21 | echo "🏗️ Building Docker image with buildx for x86_64 architecture..." 22 | docker buildx build --platform linux/amd64 --load -t "$ECR_REPO_NAME" . 23 | 24 | # Tag the image 25 | echo "🏷️ Tagging image..." 26 | docker tag "$ECR_REPO_NAME":latest "$ECR_REPO_URI":latest 27 | 28 | # Push the image to ECR 29 | echo "⬆️ Pushing image to ECR..." 30 | docker push "$ECR_REPO_URI":latest 31 | 32 | echo "✅ Successfully built and pushed image to:" 33 | echo "$ECR_REPO_URI:latest" -------------------------------------------------------------------------------- /deploy.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | import argparse 3 | import boto3 4 | import time 5 | import sys 6 | import subprocess 7 | import os 8 | import random 9 | import json 10 | import uuid 11 | import zipfile 12 | import io 13 | 14 | def wait_for_function_update_completion(lambda_client, function_name): 15 | """ 16 | Wait for a Lambda function to complete any in-progress updates. 17 | 18 | Args: 19 | lambda_client: The boto3 Lambda client 20 | function_name: The name of the Lambda function 21 | """ 22 | max_attempts = 30 23 | for attempt in range(max_attempts): 24 | try: 25 | # Get the current state of the function 26 | response = lambda_client.get_function(FunctionName=function_name) 27 | last_update_status = response['Configuration'].get('LastUpdateStatus') 28 | 29 | # If LastUpdateStatus is not present or is "Successful", the update is complete 30 | if not last_update_status or last_update_status == 'Successful': 31 | print(f"Function update completed successfully") 32 | return True 33 | 34 | # If the update failed, report the error 35 | if last_update_status == 'Failed': 36 | failure_reason = response['Configuration'].get('LastUpdateStatusReason', 'Unknown error') 37 | print(f"Function update failed: {failure_reason}") 38 | return False 39 | 40 | # Still in progress, wait and retry 41 | print(f"Function update status: {last_update_status}. Waiting...") 42 | time.sleep(2) 43 | 44 | except Exception as e: 45 | # Specific check for ThrottlingException to avoid unnecessary retries 46 | if hasattr(e, 'response') and e.response.get('Error', {}).get('Code') == 'ThrottlingException': 47 | print(f"Throttling exception encountered: {str(e)}. Waiting before retry...") 48 | time.sleep(5 + random.random()) # Wait a bit longer for throttling 49 | elif isinstance(e, lambda_client.exceptions.ResourceNotFoundException): 50 | print(f"Function {function_name} not found during status check.") 51 | return False # Function doesn't exist, can't complete update 52 | else: 53 | print(f"Error checking function status: {str(e)}") 54 | time.sleep(2) # General retry delay 55 | 56 | print(f"Timed out waiting for function update to complete") 57 | return False 58 | 59 | 60 | def build_and_push_container(): 61 | """ 62 | Calls the build_and_push.sh script to build and push a Docker container to ECR. 63 | 64 | Returns: 65 | str: The ECR image URI if successful, None otherwise 66 | """ 67 | print("=" * 80) 68 | print("No image URI provided. Building and pushing container to ECR...") 69 | print("=" * 80) 70 | 71 | try: 72 | # Ensure the script is executable 73 | script_path = "./build_and_push.sh" 74 | if not os.path.isfile(script_path): 75 | raise FileNotFoundError(f"Script {script_path} not found. Please make sure it exists in the current directory.") 76 | 77 | # Execute the build_and_push.sh script 78 | result = subprocess.run( 79 | [script_path], 80 | stdout=subprocess.PIPE, 81 | stderr=subprocess.PIPE, 82 | text=True, 83 | check=True 84 | ) 85 | 86 | # Extract the ECR image URI from the script output 87 | # Assuming the script outputs the ECR URI as the last line or in a specific format 88 | output_lines = result.stdout.strip().split('\n') 89 | ecr_uri = output_lines[-1].strip() 90 | 91 | # Validate the URI (basic check) 92 | if not (ecr_uri.startswith("https://") or 93 | ".dkr.ecr." in ecr_uri or 94 | ".amazonaws.com/" in ecr_uri): 95 | print(f"Warning: The returned URI '{ecr_uri}' doesn't look like a valid ECR URI.") 96 | print("Full script output:") 97 | print(result.stdout) 98 | return None 99 | 100 | print(f"Successfully built and pushed container to: {ecr_uri}") 101 | return ecr_uri 102 | 103 | except subprocess.CalledProcessError as e: 104 | print(f"Error executing build_and_push.sh: {e}") 105 | print(f"Script output: {e.stdout}") 106 | print(f"Script error: {e.stderr}") 107 | return None 108 | except Exception as e: 109 | print(f"Error during build and push process: {str(e)}") 110 | return None 111 | 112 | def format_cognito_domain(domain, user_pool_id, region): 113 | """ 114 | Formats the Cognito domain to ensure it is properly constructed. 115 | 116 | Args: 117 | domain (str): The provided domain (can be prefix or full domain) 118 | user_pool_id (str): The Cognito user pool ID 119 | region (str): The AWS region 120 | 121 | Returns: 122 | str: The properly formatted domain 123 | """ 124 | if not domain: 125 | # If no domain provided, construct from user pool ID 126 | if user_pool_id: 127 | parts = user_pool_id.split('_') 128 | if len(parts) == 2: 129 | region_prefix = parts[0] 130 | id_part = parts[1].lower() 131 | domain = f"{region_prefix}-{id_part}" 132 | else: 133 | # Fallback if format is unexpected 134 | domain = user_pool_id.lower().replace('_', '-') 135 | 136 | # Check if this is a full domain or just a prefix 137 | if domain and not (domain.endswith('.amazoncognito.com') or '.auth.' in domain): 138 | # Just a prefix, add the full domain 139 | domain = f"{domain}.auth.{region}.amazoncognito.com" 140 | 141 | return domain 142 | def create_lambda_authorizer_function(lambda_client, function_name, role_arn, region): 143 | """ 144 | Create a Lambda authorizer function using the original auth.py code. 145 | 146 | Args: 147 | lambda_client: The boto3 Lambda client 148 | function_name: Name for the Lambda authorizer function 149 | role_arn: ARN of the Lambda execution role 150 | region: AWS region 151 | 152 | Returns: 153 | str: The ARN of the created Lambda function 154 | """ 155 | print("=" * 80) 156 | print(f"Creating Lambda authorizer function: {function_name}") 157 | print("=" * 80) 158 | 159 | try: 160 | # Create a ZIP deployment package with auth.py 161 | zip_buffer = io.BytesIO() 162 | with zipfile.ZipFile(zip_buffer, 'w') as zip_file: 163 | with open('src/auth/auth.py', 'rb') as auth_file: 164 | zip_file.writestr('auth.py', auth_file.read()) 165 | 166 | zip_buffer.seek(0) 167 | 168 | # Check if function already exists 169 | try: 170 | lambda_client.get_function(FunctionName=function_name) 171 | function_exists = True 172 | print(f"Lambda authorizer function {function_name} already exists. Updating...") 173 | except lambda_client.exceptions.ResourceNotFoundException: 174 | function_exists = False 175 | print(f"Lambda authorizer function {function_name} does not exist. Creating new function...") 176 | 177 | if function_exists: 178 | # Update existing function 179 | response = lambda_client.update_function_code( 180 | FunctionName=function_name, 181 | ZipFile=zip_buffer.read() 182 | ) 183 | else: 184 | # Create new function 185 | response = lambda_client.create_function( 186 | FunctionName=function_name, 187 | Runtime='python3.9', 188 | Role=role_arn, 189 | Handler='auth.lambda_handler', 190 | Code={'ZipFile': zip_buffer.read()}, 191 | Timeout=10 192 | ) 193 | 194 | # Get the function ARN 195 | function_arn = response['FunctionArn'] 196 | print(f"Lambda authorizer function deployed: {function_arn}") 197 | 198 | # Wait for function to be active 199 | print("Waiting for Lambda authorizer function to be ready...") 200 | function_state = "" 201 | max_attempts = 10 202 | attempts = 0 203 | 204 | while function_state != "Active" and attempts < max_attempts: 205 | time.sleep(5) 206 | attempts += 1 207 | function_info = lambda_client.get_function(FunctionName=function_name) 208 | function_state = function_info['Configuration']['State'] 209 | print(f"Current state: {function_state}") 210 | 211 | if function_state == "Active": 212 | print(f"Successfully deployed Lambda authorizer function: {function_name}") 213 | return function_arn 214 | else: 215 | print(f"Function deployment did not reach Active state in time. Last state: {function_state}") 216 | return None 217 | 218 | except Exception as e: 219 | print(f"Error deploying Lambda authorizer function: {str(e)}") 220 | return None 221 | 222 | def attach_lambda_authorizer(apigateway_client, api_id, authorizer_name, function_arn, region): 223 | """ 224 | Attach a Lambda authorizer to an API Gateway. 225 | 226 | Args: 227 | apigateway_client: The boto3 API Gateway client 228 | api_id: The ID of the API Gateway 229 | authorizer_name: The name for the authorizer 230 | function_arn: The ARN of the Lambda authorizer function 231 | region: AWS region 232 | 233 | Returns: 234 | str: The ID of the created authorizer 235 | """ 236 | print("=" * 80) 237 | print(f"Attaching Lambda authorizer to API Gateway: {api_id}") 238 | print("=" * 80) 239 | 240 | try: 241 | # Check if authorizer exists 242 | authorizer_id = None 243 | try: 244 | authorizers = apigateway_client.get_authorizers(ApiId=api_id) 245 | for authorizer in authorizers.get('Items', []): 246 | if authorizer['Name'] == authorizer_name: 247 | authorizer_id = authorizer['AuthorizerId'] 248 | print(f"Found existing authorizer: {authorizer_id}") 249 | break 250 | except Exception as e: 251 | print(f"Error checking existing authorizers: {str(e)}") 252 | 253 | # Format the Lambda function ARN for API Gateway invocation 254 | # Convert from: arn:aws:lambda:region:account-id:function:function-name 255 | # To: arn:aws:apigateway:region:lambda:path/2015-03-31/functions/arn:aws:lambda:region:account-id:function:function-name/invocations 256 | sts_client = boto3.client('sts', region_name=region) 257 | account_id = sts_client.get_caller_identity()['Account'] 258 | 259 | # Extract function name from ARN 260 | function_name = function_arn.split(':')[-1] 261 | if function_name.startswith('function:'): 262 | function_name = function_name[9:] # Remove 'function:' prefix 263 | 264 | # Construct the proper invocation URI for API Gateway 265 | invocation_uri = f"arn:aws:apigateway:{region}:lambda:path/2015-03-31/functions/arn:aws:lambda:{region}:{account_id}:function:{function_name}/invocations" 266 | 267 | print(f"Using Lambda invocation URI: {invocation_uri}") 268 | 269 | # Create or update authorizer 270 | if authorizer_id: 271 | # Update existing authorizer 272 | response = apigateway_client.update_authorizer( 273 | ApiId=api_id, 274 | AuthorizerId=authorizer_id, 275 | AuthorizerPayloadFormatVersion='2.0', 276 | AuthorizerType='REQUEST', 277 | AuthorizerUri=invocation_uri, 278 | EnableSimpleResponses=True, 279 | IdentitySource=['$request.header.Authorization'] 280 | ) 281 | authorizer_id = response['AuthorizerId'] 282 | print(f"Updated Lambda authorizer: {authorizer_id}") 283 | else: 284 | # Create new authorizer 285 | response = apigateway_client.create_authorizer( 286 | ApiId=api_id, 287 | AuthorizerPayloadFormatVersion='2.0', 288 | AuthorizerType='REQUEST', 289 | AuthorizerUri=invocation_uri, 290 | EnableSimpleResponses=True, 291 | IdentitySource=['$request.header.Authorization'], 292 | Name=authorizer_name 293 | ) 294 | authorizer_id = response['AuthorizerId'] 295 | print(f"Created Lambda authorizer: {authorizer_id}") 296 | 297 | # Add Lambda permission for API Gateway to invoke the authorizer 298 | lambda_client = boto3.client('lambda', region_name=region) 299 | 300 | # Source ARN for the authorizer 301 | source_arn = f"arn:aws:execute-api:{region}:{account_id}:{api_id}/authorizers/{authorizer_id}" 302 | 303 | try: 304 | # Remove potentially existing permission first to avoid conflicts 305 | lambda_client.remove_permission( 306 | FunctionName=function_name, 307 | StatementId=f'apigateway-auth-{api_id}' 308 | ) 309 | except Exception: 310 | pass # Ignore if permission doesn't exist 311 | 312 | # Add permission for API Gateway to invoke the authorizer 313 | lambda_client.add_permission( 314 | FunctionName=function_name, 315 | StatementId=f'apigateway-auth-{api_id}', 316 | Action='lambda:InvokeFunction', 317 | Principal='apigateway.amazonaws.com', 318 | SourceArn=source_arn 319 | ) 320 | 321 | print(f"Added permission for API Gateway to invoke Lambda authorizer") 322 | 323 | # Update routes to use the authorizer 324 | routes = apigateway_client.get_routes(ApiId=api_id) 325 | for route in routes.get('Items', []): 326 | # Skip public paths 327 | if any(route['RouteKey'].endswith(path) for path in [ 328 | '/.well-known/oauth-protected-resource', 329 | '/.well-known/oauth-authorization-server', 330 | '/.well-known/openid-configuration' 331 | ]): 332 | continue 333 | 334 | try: 335 | apigateway_client.update_route( 336 | ApiId=api_id, 337 | RouteId=route['RouteId'], 338 | AuthorizerId=authorizer_id, 339 | AuthorizationType='CUSTOM' 340 | ) 341 | print(f"Updated route {route['RouteKey']} to use Lambda authorizer") 342 | except Exception as e: 343 | print(f"Error updating route {route['RouteKey']}: {str(e)}") 344 | 345 | return authorizer_id 346 | 347 | except Exception as e: 348 | print(f"Error attaching Lambda authorizer: {str(e)}") 349 | return None 350 | def deploy_lambda_container(ecr_image_uri, function_name, role_arn, bedrock_role_arn, region="us-east-1", memory_size=1024, timeout=90, api_gateway=False, api_name=None, stage_name="prod", cognito_user_pool_id=None, cognito_domain=None, cognito_client_ids=None, auth_method="oauth"): 351 | """ 352 | Deploy a container from ECR as a Lambda function with optional API Gateway. 353 | 354 | Args: 355 | ecr_image_uri (str): URI of the ECR image to deploy 356 | function_name (str): Name of the main Lambda function 357 | role_arn (str): ARN of the Lambda execution role 358 | bedrock_role_arn (str): ARN of the role to use when invoking Bedrock models 359 | region (str): AWS region to deploy the Lambda function 360 | memory_size (int): Memory size in MB for the main Lambda function 361 | timeout (int): Timeout in seconds for the main Lambda function 362 | api_gateway (bool): Whether to create an API Gateway for the Lambda 363 | api_name (str): Name for the API Gateway (defaults to function-name-api) 364 | stage_name (str): API Gateway stage name 365 | cognito_user_pool_id (str): ID of the Cognito user pool (for OAuth method) 366 | cognito_domain (str): Cognito domain for the user pool (for OAuth method) 367 | cognito_client_ids (str): Comma-separated list of allowed client IDs (for OAuth method) 368 | auth_method (str): Authorization method to use ('oauth' or 'lambda') 369 | """ 370 | print("=" * 80) 371 | print(f"Deploying Lambda function {function_name} in region {region}...") 372 | print("=" * 80) 373 | 374 | # Initialize the Lambda client with specified region 375 | lambda_client = boto3.client('lambda', region_name=region) 376 | 377 | try: 378 | # Check if function already exists 379 | try: 380 | lambda_client.get_function(FunctionName=function_name) 381 | function_exists = True 382 | print(f"Function {function_name} already exists. Updating...") 383 | except lambda_client.exceptions.ResourceNotFoundException: 384 | function_exists = False 385 | print(f"Function {function_name} does not exist. Creating new function...") 386 | 387 | # Create or update Lambda function 388 | if function_exists: 389 | # Update function code with retry logic 390 | max_code_retries = 5 391 | for attempt in range(max_code_retries): 392 | try: 393 | response = lambda_client.update_function_code( 394 | FunctionName=function_name, 395 | ImageUri=ecr_image_uri, 396 | Publish=True 397 | ) 398 | print(f"Function code update initiated successfully") 399 | break 400 | except lambda_client.exceptions.ResourceConflictException as e: 401 | if attempt < max_code_retries - 1: 402 | wait_time = (2 ** attempt) + (random.random() * 0.5) # Exponential backoff with jitter 403 | print(f"Update already in progress. Waiting {wait_time:.2f} seconds before retrying...") 404 | time.sleep(wait_time) 405 | else: 406 | raise e 407 | 408 | # Wait for function code update to complete before updating configuration 409 | print("Waiting for function code update to complete...") 410 | wait_for_function_update_completion(lambda_client, function_name) 411 | 412 | # Update configuration with retry logic 413 | max_config_retries = 5 414 | for attempt in range(max_config_retries): 415 | try: 416 | lambda_client.update_function_configuration( 417 | FunctionName=function_name, 418 | Timeout=timeout, 419 | MemorySize=memory_size 420 | ) 421 | print(f"Function configuration updated successfully") 422 | break 423 | except lambda_client.exceptions.ResourceConflictException as e: 424 | if attempt < max_config_retries - 1: 425 | wait_time = (2 ** attempt) + (random.random() * 0.5) # Exponential backoff with jitter 426 | print(f"Update already in progress. Waiting {wait_time:.2f} seconds before retrying...") 427 | time.sleep(wait_time) 428 | else: 429 | raise e 430 | else: 431 | # Create new function 432 | env_vars = {} 433 | if bedrock_role_arn is not None: 434 | env_vars['BEDROCK_ROLE_ARN'] = bedrock_role_arn 435 | 436 | # Add authorization environment variables 437 | env_vars['AUTH_METHOD'] = auth_method 438 | 439 | # Add Cognito environment variables if using OAuth 440 | if auth_method == 'oauth' and cognito_user_pool_id: 441 | env_vars['COGNITO_USER_POOL_ID'] = cognito_user_pool_id 442 | env_vars['COGNITO_REGION'] = region 443 | 444 | if cognito_domain or cognito_user_pool_id: 445 | formatted_domain = format_cognito_domain(cognito_domain, cognito_user_pool_id, region) 446 | env_vars['COGNITO_DOMAIN'] = formatted_domain 447 | print(f"Setting COGNITO_DOMAIN: {formatted_domain}") 448 | 449 | if cognito_client_ids: 450 | env_vars['COGNITO_ALLOWED_CLIENT_IDS'] = cognito_client_ids 451 | 452 | env = { 453 | 'Variables': env_vars 454 | } 455 | 456 | response = lambda_client.create_function( 457 | FunctionName=function_name, 458 | PackageType='Image', 459 | Code={ 460 | 'ImageUri': ecr_image_uri 461 | }, 462 | Role=role_arn, 463 | Timeout=timeout, 464 | MemorySize=memory_size, 465 | Environment=env 466 | ) 467 | 468 | # Wait for function to be active 469 | print("Waiting for Lambda function to be ready...") 470 | function_state = "" 471 | max_attempts = 10 472 | attempts = 0 473 | 474 | while function_state != "Active" and attempts < max_attempts: 475 | time.sleep(5) 476 | attempts += 1 477 | function_info = lambda_client.get_function(FunctionName=function_name) 478 | function_state = function_info['Configuration']['State'] 479 | print(f"Current state: {function_state}") 480 | 481 | if function_state == "Active": 482 | print(f"Successfully deployed Lambda function: {function_name}") 483 | function_arn = function_info['Configuration']['FunctionArn'] 484 | print(f"Function ARN: {function_arn}") 485 | 486 | if api_gateway: 487 | # Setup API Gateway 488 | return deploy_api_gateway(function_name, function_arn, region, api_name=api_name, stage_name=stage_name) 489 | else: 490 | print("No API Gateway requested. Lambda deployment complete.") 491 | return True 492 | else: 493 | print(f"Function deployment did not reach Active state in time. Last state: {function_state}") 494 | return False 495 | 496 | except Exception as e: 497 | print(f"Error deploying Lambda function: {str(e)}") 498 | return False 499 | finally: 500 | print("Lambda deployment process completed.") 501 | def deploy_api_gateway(function_name, function_arn, region, api_name=None, stage_name="prod"): 502 | """ 503 | Deploy an API Gateway v2 HTTP API with Lambda integration. 504 | 505 | Args: 506 | function_name (str): The backend Lambda function name 507 | function_arn (str): The backend Lambda function ARN 508 | region (str): AWS region 509 | api_name (str): Name for the API Gateway 510 | stage_name (str): API Gateway stage name 511 | 512 | Returns: 513 | bool: True if successful, False otherwise 514 | """ 515 | if api_name is None: 516 | api_name = f"{function_name}-api" 517 | 518 | print("=" * 80) 519 | print(f"Deploying API Gateway ({api_name}) for: {function_name}") 520 | print("=" * 80) 521 | 522 | # Initialize clients 523 | apigateway_client = boto3.client('apigatewayv2', region_name=region) 524 | lambda_client = boto3.client('lambda', region_name=region) 525 | sts_client = boto3.client('sts', region_name=region) 526 | account_id = sts_client.get_caller_identity()['Account'] 527 | 528 | try: 529 | # Step 1: Create or get existing API 530 | api_id = None 531 | 532 | # Check if API already exists with this name 533 | try: 534 | response = apigateway_client.get_apis() 535 | for api in response.get('Items', []): 536 | if api['Name'] == api_name: 537 | api_id = api['ApiId'] 538 | print(f"Found existing API Gateway: {api_id}") 539 | break 540 | except Exception as e: 541 | print(f"Error checking existing APIs: {str(e)}") 542 | 543 | # Create new API if needed 544 | if not api_id: 545 | try: 546 | # Create HTTP API 547 | response = apigateway_client.create_api( 548 | Name=api_name, 549 | ProtocolType='HTTP', 550 | CorsConfiguration={ 551 | 'AllowOrigins': ['*'], 552 | 'AllowMethods': ['*'], 553 | 'AllowHeaders': ['*'], 554 | 'MaxAge': 86400 555 | } 556 | ) 557 | api_id = response['ApiId'] 558 | print(f"Created new API Gateway: {api_id}") 559 | except Exception as e: 560 | print(f"Error creating API Gateway: {str(e)}") 561 | return False 562 | 563 | # Step 2: Create or update integration with Lambda function 564 | integration_id = None 565 | 566 | try: 567 | # Check for existing integrations 568 | response = apigateway_client.get_integrations(ApiId=api_id) 569 | for integration in response.get('Items', []): 570 | if integration.get('IntegrationUri') == function_arn: 571 | integration_id = integration['IntegrationId'] 572 | print(f"Found existing Lambda integration: {integration_id}") 573 | break 574 | 575 | if not integration_id: 576 | # Create new integration 577 | response = apigateway_client.create_integration( 578 | ApiId=api_id, 579 | IntegrationType='AWS_PROXY', 580 | IntegrationMethod='POST', 581 | PayloadFormatVersion='2.0', 582 | IntegrationUri=function_arn, 583 | TimeoutInMillis=30000 584 | ) 585 | integration_id = response['IntegrationId'] 586 | print(f"Created new Lambda integration: {integration_id}") 587 | except Exception as e: 588 | print(f"Error setting up Lambda integration: {str(e)}") 589 | return False 590 | 591 | # Step 3: Create or update routes for the API 592 | # Define the specific routes to configure 593 | route_keys_to_ensure = { 594 | 'GET /': {}, 595 | 'GET /docs': {}, 596 | 'GET /{proxy+}': {}, 597 | 'POST /{proxy+}': {}, 598 | 'DELETE /{proxy+}': {} 599 | } 600 | 601 | # Get existing routes 602 | existing_routes = {} 603 | try: 604 | paginator = apigateway_client.get_paginator('get_routes') 605 | for page in paginator.paginate(ApiId=api_id): 606 | for route in page.get('Items', []): 607 | existing_routes[route['RouteKey']] = route['RouteId'] 608 | except Exception as e: 609 | print(f"Warning: Could not retrieve existing routes: {str(e)}") 610 | 611 | for route_key, config in route_keys_to_ensure.items(): 612 | try: 613 | target = f'integrations/{integration_id}' 614 | if route_key in existing_routes: 615 | route_id = existing_routes[route_key] 616 | print(f"Updating existing route: {route_key} (ID: {route_id})") 617 | apigateway_client.update_route( 618 | ApiId=api_id, 619 | RouteId=route_id, 620 | RouteKey=route_key, 621 | Target=target 622 | ) 623 | else: 624 | print(f"Creating route: {route_key}") 625 | response = apigateway_client.create_route( 626 | ApiId=api_id, 627 | RouteKey=route_key, 628 | Target=target 629 | ) 630 | print(f"Created route: {route_key} (ID: {response['RouteId']})") 631 | except Exception as e: 632 | print(f"Error creating/updating route {route_key}: {str(e)}") 633 | 634 | # Step 4: Add Lambda permission for API Gateway to invoke the Lambda 635 | try: 636 | # Source ARN needs to cover all defined routes/methods 637 | source_arn = f"arn:aws:execute-api:{region}:{account_id}:{api_id}/*/*" 638 | statement_id_backend = f'apigateway-invoke-lambda-{api_id}' 639 | 640 | try: 641 | # Remove potentially existing permission first to avoid conflicts 642 | lambda_client.remove_permission( 643 | FunctionName=function_name, 644 | StatementId=statement_id_backend, 645 | ) 646 | print(f"Removed existing invoke permission for Lambda {function_name} (StatementId: {statement_id_backend})") 647 | except lambda_client.exceptions.ResourceNotFoundException: 648 | print(f"No existing invoke permission found for Lambda {function_name} (StatementId: {statement_id_backend}), proceeding to add.") 649 | except Exception as e: 650 | print(f"Warning: Could not remove potentially existing permission for Lambda {function_name} (StatementId: {statement_id_backend}): {e}") 651 | 652 | # Add the permission 653 | lambda_client.add_permission( 654 | FunctionName=function_name, 655 | StatementId=statement_id_backend, 656 | Action='lambda:InvokeFunction', 657 | Principal='apigateway.amazonaws.com', 658 | SourceArn=source_arn 659 | ) 660 | print(f"Added/Updated permission for API Gateway to invoke Lambda: {function_name}") 661 | except Exception as e: 662 | print(f"Error setting Lambda permission: {str(e)}") 663 | return False 664 | 665 | # Step 5: Deploy the API to a stage 666 | try: 667 | # Check if stage exists 668 | stage_exists = False 669 | try: 670 | apigateway_client.get_stage(ApiId=api_id, StageName=stage_name) 671 | stage_exists = True 672 | print(f"Stage {stage_name} already exists") 673 | except apigateway_client.exceptions.NotFoundException: 674 | pass 675 | 676 | if not stage_exists: 677 | response = apigateway_client.create_stage( 678 | ApiId=api_id, 679 | StageName=stage_name, 680 | AutoDeploy=True 681 | ) 682 | print(f"Created stage: {stage_name}") 683 | except Exception as e: 684 | print(f"Error creating API stage: {str(e)}") 685 | return False 686 | 687 | # Print the API URL 688 | api_url = f"https://{api_id}.execute-api.{region}.amazonaws.com/{stage_name}" 689 | print("\n" + "=" * 80) 690 | print(f"API Gateway successfully deployed!") 691 | print(f"API URL: {api_url}") 692 | print(f"Lambda Function ARN: {function_arn}") 693 | print("\nMCP Authorization is now implemented directly in the Lambda function") 694 | print("Use the /.well-known/oauth-protected-resource endpoint to discover authorization metadata") 695 | print("=" * 80) 696 | 697 | return True 698 | 699 | except Exception as e: 700 | print(f"Error deploying API Gateway: {str(e)}") 701 | return False 702 | 703 | if __name__ == '__main__': 704 | parser = argparse.ArgumentParser(description='Deploy a container from ECR as a Lambda function') 705 | parser.add_argument('--image-uri', required=False, help='ECR image URI to deploy (if not provided, will build and push container)') 706 | parser.add_argument('--function-name', required=True, help='Name for the Lambda function') 707 | parser.add_argument('--role-arn', required=True, help='ARN of the Lambda execution role') 708 | parser.add_argument('--bedrock-role-arn', required=False, help='ARN of the role to use when invoking Bedrock models') 709 | parser.add_argument('--region', default='us-east-1', help='AWS region to deploy the Lambda function (default: us-east-1)') 710 | parser.add_argument('--memory', type=int, default=2048, help='Memory size in MB (default: 2048)') 711 | parser.add_argument('--timeout', type=int, default=300, help='Timeout in seconds (default: 300)') 712 | parser.add_argument('--api-gateway', action='store_true', help='Create an API Gateway') 713 | parser.add_argument('--api-name', help='Name for the API Gateway (defaults to function-name-api)') 714 | parser.add_argument('--stage-name', default='prod', help='API Gateway stage name (default: prod)') 715 | parser.add_argument('--auth-method', choices=['oauth', 'lambda'], default='oauth', help='Authorization method to use (oauth or lambda)') 716 | parser.add_argument('--lambda-authorizer-name', help='Name for the Lambda authorizer function (required if auth-method is lambda)') 717 | parser.add_argument('--cognito-user-pool-id', help='AWS Cognito User Pool ID for authorization') 718 | parser.add_argument('--cognito-domain', help='AWS Cognito domain name for authorization') 719 | parser.add_argument('--cognito-client-ids', help='Comma-separated list of allowed Cognito client IDs') 720 | 721 | args = parser.parse_args() 722 | 723 | # Validate auth method arguments 724 | if args.auth_method == 'lambda' and not args.lambda_authorizer_name: 725 | print("Error: --lambda-authorizer-name is required when using --auth-method=lambda") 726 | sys.exit(1) 727 | 728 | if args.auth_method == 'oauth' and not args.cognito_user_pool_id: 729 | print("Warning: --cognito-user-pool-id is recommended when using --auth-method=oauth") 730 | 731 | # Determine if we need to build and push a container or use the provided URI 732 | ecr_image_uri = args.image_uri 733 | if not ecr_image_uri: 734 | ecr_image_uri = build_and_push_container() 735 | if not ecr_image_uri: 736 | print("Failed to build and push container. Exiting.") 737 | sys.exit(1) 738 | else: 739 | print("=" * 80) 740 | print(f"Using provided image URI: {ecr_image_uri}") 741 | print("=" * 80) 742 | 743 | # Deploy the Lambda function with the image URI 744 | success = deploy_lambda_container( 745 | ecr_image_uri, 746 | args.function_name, 747 | args.role_arn, 748 | args.bedrock_role_arn, 749 | args.region, 750 | args.memory, 751 | args.timeout, 752 | args.api_gateway, 753 | args.api_name, 754 | args.stage_name, 755 | args.cognito_user_pool_id if args.auth_method == 'oauth' else None, 756 | args.cognito_domain if args.auth_method == 'oauth' else None, 757 | args.cognito_client_ids if args.auth_method == 'oauth' else None, 758 | args.auth_method 759 | ) 760 | 761 | if not success: 762 | sys.exit(1) 763 | 764 | # If using Lambda authorizer, deploy it 765 | if args.auth_method == 'lambda' and args.api_gateway: 766 | lambda_client = boto3.client('lambda', region_name=args.region) 767 | apigateway_client = boto3.client('apigatewayv2', region_name=args.region) 768 | 769 | # Get the API ID 770 | api_id = None 771 | api_name = args.api_name or f"{args.function_name}-api" 772 | 773 | try: 774 | response = apigateway_client.get_apis() 775 | for api in response.get('Items', []): 776 | if api['Name'] == api_name: 777 | api_id = api['ApiId'] 778 | break 779 | except Exception as e: 780 | print(f"Error getting API ID: {str(e)}") 781 | sys.exit(1) 782 | 783 | if not api_id: 784 | print(f"Could not find API with name: {api_name}") 785 | sys.exit(1) 786 | 787 | # Deploy the Lambda authorizer 788 | authorizer_arn = create_lambda_authorizer_function( 789 | lambda_client, 790 | args.lambda_authorizer_name, 791 | args.role_arn, 792 | args.region 793 | ) 794 | 795 | if not authorizer_arn: 796 | print("Failed to deploy Lambda authorizer. Exiting.") 797 | sys.exit(1) 798 | 799 | # Attach the authorizer to the API Gateway 800 | authorizer_id = attach_lambda_authorizer( 801 | apigateway_client, 802 | api_id, 803 | args.lambda_authorizer_name, 804 | authorizer_arn, 805 | args.region 806 | ) 807 | 808 | if not authorizer_id: 809 | print("Failed to attach Lambda authorizer to API Gateway. Exiting.") 810 | sys.exit(1) 811 | 812 | print("\n" + "=" * 80) 813 | print(f"Lambda authorizer successfully deployed and attached to API Gateway!") 814 | print(f"API ID: {api_id}") 815 | print(f"Authorizer ID: {authorizer_id}") 816 | print("=" * 80) 817 | -------------------------------------------------------------------------------- /docs/audience-benefits.md: -------------------------------------------------------------------------------- 1 | # Streamable MCP Server: Benefits for Different Audiences 2 | 3 | ## For Developers 4 | 5 | ### Technical Benefits 6 | - **Simplified Integration**: The MCP protocol provides a standardized way to build and connect AI tools, reducing custom integration work. 7 | - **TypeScript Implementation**: Full TypeScript support with type definitions for better developer experience. 8 | - **Session Management**: Built-in session handling through HTTP headers eliminates the need to implement custom session logic. 9 | - **Resumability**: Support for reconnecting and continuing from where you left off, improving resilience. 10 | - **Notification System**: Built-in support for server-to-client notifications for real-time updates. 11 | 12 | ### Development Workflow Benefits 13 | - **Local Development**: Easy local testing with the same code that runs in production. 14 | - **SDK Support**: Leverages the MCP SDK for both client and server implementations. 15 | - **Containerization**: Docker-based deployment simplifies environment consistency. 16 | - **Example Tools**: Includes working examples of tools that can be used as templates. 17 | 18 | ### Code Examples 19 | ```typescript 20 | // Server-side tool implementation 21 | server.tool( 22 | 'greet', 23 | 'A simple greeting tool', 24 | { 25 | name: z.string().describe('Name to greet'), 26 | }, 27 | async ({ name }): Promise => { 28 | return { 29 | content: [{ type: 'text', text: `Hello, ${name}!` }], 30 | }; 31 | } 32 | ); 33 | 34 | // Client-side tool invocation 35 | await callTool('greet', { name: 'User' }); 36 | ``` 37 | 38 | ## For DevOps Engineers 39 | 40 | ### Infrastructure Benefits 41 | - **Serverless Architecture**: No servers to manage or scale. 42 | - **Auto-scaling**: Automatically handles varying loads without configuration. 43 | - **Cost Optimization**: Pay only for actual usage, with no idle resources. 44 | - **Containerized Deployment**: Consistent environments across development and production. 45 | - **Infrastructure as Code**: Deployment script creates all necessary AWS resources. 46 | 47 | ### Operational Benefits 48 | - **Simplified Monitoring**: Integration with CloudWatch for logs and metrics. 49 | - **API Management**: API Gateway provides rate limiting, authentication, and usage plans. 50 | - **Deployment Automation**: Single command deployment process. 51 | - **Resource Cleanup**: Automated resource management and cleanup. 52 | 53 | ### Deployment Example 54 | ```bash 55 | python deploy.py \ 56 | --function-name bedrock-spend-mcp-server \ 57 | --role-arn \ 58 | --region us-east-1 \ 59 | --memory 2048 \ 60 | --timeout 300 \ 61 | --api-gateway \ 62 | --api-name mcp-server-api \ 63 | --stage-name prod 64 | ``` 65 | 66 | ## For Business Stakeholders 67 | 68 | ### Strategic Benefits 69 | - **Cost Visibility**: Real-time insights into AI model usage and costs. 70 | - **Resource Optimization**: Identify opportunities to optimize model usage and reduce costs. 71 | - **Scalability**: Handles growth without infrastructure investments. 72 | - **Reduced Time-to-Market**: Faster deployment of AI tools and capabilities. 73 | - **Standards Compliance**: Alignment with open standards for better interoperability. 74 | 75 | ### Financial Benefits 76 | - **Lower Infrastructure Costs**: Serverless model eliminates the need for dedicated servers. 77 | - **Operational Efficiency**: Reduced DevOps overhead for maintaining infrastructure. 78 | - **Usage-Based Pricing**: Pay only for what you use, with no upfront commitments. 79 | - **Cost Monitoring**: Built-in tools for tracking and analyzing AI model usage costs. 80 | 81 | ### Example Use Case: Bedrock Spend Analysis 82 | The included Bedrock usage analysis tool provides: 83 | - Total requests and token usage across all models 84 | - Daily usage patterns to identify trends 85 | - Usage breakdown by region, model, and user 86 | - Cost allocation insights for different teams or projects 87 | 88 | ## For AI Engineers 89 | 90 | ### AI Integration Benefits 91 | - **Standardized Tool Interface**: Consistent way to expose AI capabilities as tools. 92 | - **Model Usage Analytics**: Built-in analytics for understanding model usage patterns. 93 | - **Streaming Responses**: Support for streaming responses from AI models. 94 | - **Notification System**: Real-time updates during long-running AI operations. 95 | - **Session Context**: Maintain context across multiple interactions. 96 | 97 | ### AI Development Benefits 98 | - **Tool Abstraction**: Abstract AI capabilities behind a consistent tool interface. 99 | - **Prompt Management**: Built-in support for managing and versioning prompts. 100 | - **Resource Access**: Standardized way to access external resources for AI models. 101 | - **Testing Framework**: Consistent way to test AI tool implementations. 102 | 103 | ## Conclusion 104 | 105 | The Streamable MCP Server implementation on AWS Lambda and API Gateway provides significant benefits across different roles in an organization: 106 | 107 | - **Developers** get a standardized, type-safe way to build and connect AI tools 108 | - **DevOps Engineers** benefit from simplified infrastructure management and deployment 109 | - **Business Stakeholders** gain cost visibility and reduced infrastructure expenses 110 | - **AI Engineers** have a consistent framework for exposing AI capabilities as tools 111 | 112 | By leveraging the new Streamable HTTP enhancement to the MCP protocol, this implementation enables serverless deployment patterns that were previously challenging, opening up new possibilities for building and deploying AI tools in a cost-effective, scalable manner. -------------------------------------------------------------------------------- /docs/conclusion.md: -------------------------------------------------------------------------------- 1 | # Streamable MCP Server on AWS Lambda: Conclusion 2 | 3 | ## Project Summary 4 | 5 | This project demonstrates a significant advancement in the Model Context Protocol (MCP) ecosystem by implementing the new Streamable HTTP enhancement. It showcases a serverless deployment pattern for MCP servers using AWS Lambda and Amazon API Gateway, with a practical application for Amazon Bedrock spend analysis. 6 | 7 | ## Key Innovations 8 | 9 | 1. **Streamable HTTP Implementation**: The project implements the Streamable HTTP enhancement to the MCP protocol, enabling stateful sessions over standard HTTP without requiring Server-Sent Events (SSE) support from API Gateway. 10 | 11 | 2. **Serverless Deployment Pattern**: The implementation demonstrates how to deploy an MCP server in a serverless environment, leveraging AWS Lambda's HTTP streaming capabilities for Node.js runtimes. 12 | 13 | 3. **Session Management**: The server maintains session state through HTTP headers (`Mcp-Session-id`), allowing for stateful interactions in a stateless serverless environment. 14 | 15 | 4. **Practical AI Tool**: The implementation includes a practical tool for analyzing Amazon Bedrock usage and costs, demonstrating the real-world utility of MCP servers. 16 | 17 | ## Technical Highlights 18 | 19 | - **TypeScript Implementation**: Both server and client are implemented in TypeScript, providing type safety and modern development practices. 20 | - **Containerization**: The server is packaged as a Docker container for deployment on AWS Lambda. 21 | - **Express.js Integration**: Uses Express.js for HTTP routing and middleware. 22 | - **AWS SDK Integration**: Integrates with AWS SDK for CloudWatch Logs access. 23 | - **Automated Deployment**: Includes a Python script for automated deployment to AWS. 24 | 25 | ## Architectural Strengths 26 | 27 | - **Scalability**: Leverages AWS Lambda's auto-scaling capabilities. 28 | - **Cost-Effectiveness**: Pay-per-use pricing model eliminates idle resource costs. 29 | - **Operational Simplicity**: No server management or scaling configuration required. 30 | - **Security**: API Gateway provides authentication, authorization, and rate limiting. 31 | 32 | ## Value Proposition 33 | 34 | This implementation provides value to multiple stakeholders: 35 | 36 | - **Developers**: Simplified integration with AI tools through a standardized protocol. 37 | - **DevOps Engineers**: Reduced operational overhead through serverless architecture. 38 | - **Business Stakeholders**: Cost visibility and optimization for AI model usage. 39 | - **AI Engineers**: Standardized framework for exposing AI capabilities as tools. 40 | 41 | ## Future Directions 42 | 43 | The project lays the groundwork for several potential enhancements: 44 | 45 | 1. **Persistent Storage**: Replace the in-memory transport store with a persistent solution like DynamoDB or Redis. 46 | 2. **Additional Tools**: Expand the available tools to cover more AI use cases. 47 | 3. **Authentication Integration**: Integrate with AWS Cognito or other identity providers. 48 | 4. **Multi-Region Deployment**: Extend the deployment to multiple AWS regions for global availability. 49 | 5. **Enhanced Monitoring**: Add custom CloudWatch metrics for better operational visibility. 50 | 51 | ## Conclusion 52 | 53 | The Streamable MCP Server implementation represents a significant step forward in making AI tools more accessible and easier to deploy. By leveraging the new Streamable HTTP enhancement to the MCP protocol and AWS serverless technologies, it provides a scalable, cost-effective solution for deploying MCP servers without the operational overhead of traditional infrastructure. 54 | 55 | This project serves as both a reference implementation and a practical tool, demonstrating the power and flexibility of the MCP ecosystem while providing immediate utility through its Bedrock usage analysis capabilities. 56 | 57 | The combination of modern development practices, serverless architecture, and standardized protocols creates a compelling solution that addresses real-world challenges in deploying and managing AI tools in production environments. -------------------------------------------------------------------------------- /docs/linkedin-post.md: -------------------------------------------------------------------------------- 1 | # Streamable HTTP in Action: Deploying MCP Servers on AWS Lambda and API Gateway 2 | 3 | I'm excited to share a project that demonstrates the new **Streamable HTTP** enhancement to the Model Context Protocol (MCP), showcasing a powerful serverless deployment pattern for AI tools. 4 | 5 | ## What is this project? 6 | 7 | This implementation deploys an MCP server as a containerized application on AWS Lambda, accessible via Amazon API Gateway. The server provides tools for analyzing Amazon Bedrock usage and spend, making it easy to monitor and understand your AI model consumption costs. 8 | 9 | ## Why is this significant? 10 | 11 | The Model Context Protocol (MCP) is an open standard for communication between AI applications and tools. The new **Streamable HTTP** enhancement is a game-changer because: 12 | 13 | 1. It enables MCP servers to operate as independent processes handling multiple client connections 14 | 2. It implements session management via HTTP headers (`Mcp-Session-id`) 15 | 3. It allows for serverless deployments on AWS Lambda and API Gateway without requiring SSE support 16 | 17 | ## Technical Implementation 18 | 19 | - **Server & Client**: Both written in TypeScript using the MCP SDK 20 | - **Containerization**: Packaged as a Docker container for Lambda deployment 21 | - **Session Management**: Uses `Mcp-Session-id` header for maintaining client state 22 | - **Deployment**: Automated via Python script to set up ECR, Lambda, and API Gateway 23 | - **Tools**: Provides Bedrock usage analysis tools that query CloudWatch logs 24 | 25 | ## Why this architecture matters 26 | 27 | This pattern solves a critical challenge: As of April 2025, Lambda supports response streaming for Node.js 18, this combined with the fact that MCP protocol now supports Streamable HTTP enables robust MCP server deployments in serverless environments (without requiring Service Side Events (SSE) support in the API Gateway). 28 | 29 | The architecture provides: 30 | - Cost-effective serverless scaling 31 | - No infrastructure management overhead 32 | - Simplified deployment through containerization 33 | - Secure API access with API key authentication 34 | 35 | This example demonstrates how the MCP ecosystem continues to evolve, making AI tools more accessible and easier to deploy in production environments. 36 | 37 | #MCP #AWS #Serverless #APIGateway #Lambda #AmazonBedrock #AI #CloudComputing #TypeScript -------------------------------------------------------------------------------- /docs/technical-analysis.md: -------------------------------------------------------------------------------- 1 | # Technical Analysis: Streamable MCP Server on AWS Lambda 2 | 3 | ## Overview of the Model Context Protocol (MCP) and Streamable HTTP 4 | 5 | The Model Context Protocol (MCP) is an open standard for communication between AI applications and tools. The new `Streamable-HTTP` enhancement (as of the 2025-03-26 specification) represents a significant advancement in how MCP servers can be deployed, particularly in serverless environments. 6 | 7 | ## Implementation Details 8 | 9 | ### Server Implementation (`server.ts`) 10 | 11 | The server implementation showcases several key aspects of the Streamable HTTP enhancement: 12 | 13 | 1. **Session Management**: 14 | ```typescript 15 | // Map to store transports by session ID 16 | const transports: { [sessionId: string]: StreamableHTTPServerTransport } = {}; 17 | 18 | // Check for existing session ID in request headers 19 | const sessionId = req.headers['mcp-session-id'] as string | undefined; 20 | ``` 21 | 22 | 2. **Transport Creation and Management**: 23 | ```typescript 24 | transport = new StreamableHTTPServerTransport({ 25 | sessionIdGenerator: () => randomUUID(), 26 | eventStore, // Enable resumability 27 | onsessioninitialized: (sessionId) => { 28 | // Store the transport by session ID when session is initialized 29 | console.log(`Session initialized with ID: ${sessionId}`); 30 | transports[sessionId] = transport; 31 | } 32 | }); 33 | ``` 34 | 35 | 3. **HTTP Endpoint Handling**: 36 | ```typescript 37 | app.post('/prod/mcp', async (req: Request, res: Response) => { 38 | // Handle POST requests for JSON-RPC messages 39 | }); 40 | 41 | app.get('/prod/mcp', async (req: Request, res: Response) => { 42 | // Handle GET requests for SSE streams 43 | }); 44 | 45 | app.delete('/prod/mcp', async (req: Request, res: Response) => { 46 | // Handle DELETE requests for session termination 47 | }); 48 | ``` 49 | 50 | 4. **Lambda Integration**: 51 | ```typescript 52 | const isLambda = !!process.env.LAMBDA_TASK_ROOT; 53 | 54 | if (isLambda) { 55 | // Start the server 56 | handler = serverlessExpress({ app }); 57 | } 58 | ``` 59 | 60 | ### Client Implementation (`client.ts`) 61 | 62 | The client implementation demonstrates how to interact with a Streamable HTTP MCP server: 63 | 64 | 1. **Transport Creation**: 65 | ```typescript 66 | transport = new StreamableHTTPClientTransport( 67 | new URL(serverUrl), 68 | { 69 | sessionId: sessionId, 70 | debug: true, 71 | onnotification: (rawNotification) => { 72 | console.log('RAW NOTIFICATION RECEIVED:', JSON.stringify(rawNotification, null, 2)); 73 | } 74 | } 75 | ); 76 | ``` 77 | 78 | 2. **Session Management**: 79 | ```typescript 80 | // Connect the client 81 | await client.connect(transport); 82 | sessionId = transport.sessionId 83 | console.log('Transport created with session ID:', sessionId); 84 | ``` 85 | 86 | 3. **Session Termination**: 87 | ```typescript 88 | async function terminateSession(): Promise { 89 | // ... 90 | await transport.terminateSession(); 91 | // ... 92 | } 93 | ``` 94 | 95 | ## Technical Innovations 96 | 97 | ### 1. Stateful Sessions in a Stateless Environment 98 | 99 | The implementation solves a fundamental challenge: maintaining stateful sessions in a stateless serverless environment. This is achieved through: 100 | 101 | - Generating a unique session ID for each client connection 102 | - Storing the transport instance in a server-side map keyed by session ID 103 | - Including the session ID in HTTP headers for subsequent requests 104 | 105 | ### 2. Resumability 106 | 107 | The implementation supports resumability, allowing clients to reconnect and continue from where they left off: 108 | 109 | ```typescript 110 | // Check for Last-Event-ID header for resumability 111 | const lastEventId = req.headers['last-event-id'] as string | undefined; 112 | if (lastEventId) { 113 | console.log(`Client reconnecting with Last-Event-ID: ${lastEventId}`); 114 | } 115 | ``` 116 | 117 | ### 3. Lambda and API Gateway Integration 118 | 119 | The project demonstrates how to deploy an MCP server on AWS Lambda and API Gateway, which traditionally has been challenging due to: 120 | 121 | - Lambda's ephemeral nature 122 | - API Gateway's lack of native support for Server-Sent Events (SSE) 123 | 124 | The Streamable HTTP enhancement addresses these challenges by: 125 | 126 | - Using standard HTTP methods (POST, GET, DELETE) for different aspects of the protocol 127 | - Implementing session management via HTTP headers 128 | - Leveraging Lambda's HTTP streaming capabilities for Node.js runtimes 129 | 130 | ## Deployment Architecture 131 | 132 | The deployment architecture consists of: 133 | 134 | 1. **API Gateway**: Provides the HTTP API endpoint 135 | 2. **Lambda Function**: Runs the containerized MCP server 136 | 3. **ECR**: Stores the Docker container image 137 | 4. **CloudWatch Logs**: Collects and stores logs from both the server and Bedrock usage 138 | 139 | This architecture provides several benefits: 140 | 141 | - **Scalability**: Lambda automatically scales based on request volume 142 | - **Cost-effectiveness**: Pay only for actual usage 143 | - **Simplified Operations**: No server management required 144 | - **Security**: API Gateway provides authentication and authorization 145 | 146 | ## Practical Application: Bedrock Usage Analysis 147 | 148 | The implementation includes tools for analyzing Amazon Bedrock usage, demonstrating a practical application of the MCP server: 149 | 150 | ```typescript 151 | server.tool( 152 | "get_bedrock_usage_report", 153 | "Get Bedrock daily usage report (non-streaming)", 154 | // Schema definition... 155 | async (args, extra): Promise => { 156 | // Implementation... 157 | } 158 | ); 159 | ``` 160 | 161 | This tool aggregates and analyzes Bedrock usage data, providing insights into: 162 | 163 | - Total requests and token usage 164 | - Daily usage patterns 165 | - Usage by region, model, and user 166 | 167 | ## Conclusion 168 | 169 | This implementation demonstrates how the Streamable HTTP enhancement to the MCP protocol enables new deployment patterns, particularly in serverless environments. By leveraging AWS Lambda and API Gateway, it provides a scalable, cost-effective solution for deploying MCP servers without the operational overhead of managing traditional server infrastructure. 170 | 171 | The project serves as a reference implementation for developers looking to deploy their own MCP servers in serverless environments, showcasing best practices for session management, error handling, and integration with AWS services. -------------------------------------------------------------------------------- /package.json: -------------------------------------------------------------------------------- 1 | { 2 | "devDependencies": { 3 | "@types/aws-lambda": "^8.10.149", 4 | "@types/express": "^4.17.21", 5 | "@types/jest": "^29.5.12", 6 | "@types/node": "^20.11.24", 7 | "jest": "^29.7.0", 8 | "ts-jest": "^29.1.2", 9 | "tsx": "^4.19.3", 10 | "typescript": "^5.8.3" 11 | }, 12 | "dependencies": { 13 | "@aws-sdk/client-bedrock-runtime": "^3.799.0", 14 | "@aws-sdk/client-cloudwatch-logs": "^3.787.0", 15 | "@codegenie/serverless-express": "^4.12.0", 16 | "@modelcontextprotocol/sdk": "^1.11.0", 17 | "@types/jsonwebtoken": "^9.0.9", 18 | "express": "^5.1.0", 19 | "https-proxy-agent": "^7.0.6", 20 | "jose": "^6.0.10", 21 | "jsonwebtoken": "^9.0.2", 22 | "jwks-rsa": "^3.2.0", 23 | "log4js": "^6.9.1", 24 | "node-fetch": "^3.3.2", 25 | "zod": "^3.24.3" 26 | }, 27 | "scripts": { 28 | "start": "tsx src/server.ts", 29 | "client": "tsx src/client.ts", 30 | "build": "tsc", 31 | "test": "jest", 32 | "test:watch": "jest --watch", 33 | "test:coverage": "jest --coverage" 34 | } 35 | } 36 | -------------------------------------------------------------------------------- /payload.json: -------------------------------------------------------------------------------- 1 | { 2 | "jsonrpc": "2.0", 3 | "method": "initialize", 4 | "params": { 5 | "clientInfo": { "name": "curl-client", "version": "1.0" }, 6 | "protocolVersion": "2025-03-26", 7 | "capabilities": {} 8 | }, 9 | "id": "some-random-id" 10 | } -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [project] 2 | name = "mcp-server" 3 | version = "0.1.0" 4 | description = "MCP server on AWS Lambda" 5 | readme = "README.md" 6 | requires-python = ">=3.11,<3.13" 7 | license = {file = "LICENSE"} 8 | authors = [ 9 | {name = "Amit Arora", email = "aa1603@georgetown.edu"} 10 | ] 11 | dependencies = [ 12 | "boto3>=1.37.37", 13 | ] 14 | 15 | 16 | 17 | [build-system] 18 | requires = ["hatchling"] 19 | build-backend = "hatchling.build" 20 | 21 | [tool.hatch.build.targets.wheel] 22 | packages = ["app", "lambda"] 23 | -------------------------------------------------------------------------------- /src/auth/auth.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | 4 | def lambda_handler(event, context): 5 | print(f"Received event: {json.dumps(event)}") # Log the event to see its structure 6 | 7 | # --- Get ARN --- 8 | # Try methodArn first (REST API), then routeArn (HTTP API) 9 | method_arn = event.get('methodArn') 10 | if not method_arn: 11 | print("Could not find 'methodArn' in the event. Trying 'routeArn'.") 12 | method_arn = event.get('routeArn') 13 | if not method_arn: 14 | print("Could not find 'methodArn' or 'routeArn'. Cannot generate policy.") 15 | # Return an unauthorized response for HTTP API Lambda authorizers instead of raising Exception 16 | # Raising an exception results in a 500 error for the client 17 | return {"isAuthorized": False} 18 | # For REST API: raise Exception('Unauthorized') # Or return deny policy 19 | print(f"Using ARN: {method_arn}") 20 | 21 | # --- Get Token --- 22 | # For HTTP API REQUEST authorizer (Payload 2.0), the identity source is an array 23 | raw_auth_header = None 24 | identity_source = event.get('identitySource') 25 | if identity_source and isinstance(identity_source, list) and len(identity_source) > 0: 26 | raw_auth_header = identity_source[0] 27 | print(f"Found identitySource value: {raw_auth_header}") 28 | # Check headers directly for Payload 1.0 or other cases 29 | elif event.get('headers'): 30 | # Header names might be lowercased by API Gateway 31 | auth_header_key = next((k for k in event['headers'] if k.lower() == 'authorization'), None) 32 | if auth_header_key: 33 | raw_auth_header = event['headers'][auth_header_key] 34 | print(f"Found Authorization header value: {raw_auth_header}") 35 | 36 | # Fallback check for authorizationToken (REST API Lambda Authorizer - TOKEN type) 37 | elif event.get('authorizationToken'): 38 | raw_auth_header = event.get('authorizationToken') 39 | print(f"Found authorizationToken value: {raw_auth_header}") 40 | else: 41 | print("Could not find token in identitySource, headers, or authorizationToken field.") 42 | 43 | # Extract token assuming 'Bearer ' format 44 | token = None 45 | if raw_auth_header: 46 | parts = raw_auth_header.split() 47 | if len(parts) == 2 and parts[0].lower() == 'bearer': 48 | token = parts[1] 49 | print("Successfully extracted Bearer token.") 50 | else: 51 | print("Authorization header format is not 'Bearer '.") 52 | # Treat non-bearer token as invalid for this example 53 | pass # token remains None 54 | else: 55 | print("No raw authorization header value found.") 56 | 57 | # Check if token was extracted 58 | if not token: 59 | print("Token not found or invalid format.") 60 | # Return unauthorized for HTTP API simple response 61 | return {"isAuthorized": False} 62 | # For REST API: raise Exception('Unauthorized') # Or return deny policy 63 | 64 | # --- Placeholder for actual token validation --- 65 | # Replace this with your actual token validation logic (e.g., JWT verification, DB lookup) 66 | print(f"Validating token (placeholder): {token}") 67 | is_valid_token = True # Replace with actual validation result 68 | 69 | if is_valid_token: 70 | print("Token is valid (placeholder). Authorizing request.") 71 | # For HTTP API simple response, just return True 72 | return {"isAuthorized": True} 73 | # --- For REST API IAM Policy response (Example) --- 74 | # principal_id = "user|a1b2c3d4" # Unique identifier for the user 75 | # policy = AuthPolicy(principal_id, context.aws_request_id.replace('-', '')) 76 | # policy.restApiId = api_gateway_arn_tmp[0] 77 | # policy.region = api_gateway_arn_tmp[1] 78 | # policy.stage = api_gateway_arn_tmp[2] 79 | # policy.allowAllMethods() # Or specify allowed methods/resources 80 | # auth_response = policy.build() 81 | # Add context if needed: 82 | # auth_response['context'] = { 83 | # 'stringKey': 'stringval', 84 | # 'numberKey': 123, 85 | # 'booleanKey': True 86 | # } 87 | # return auth_response 88 | else: 89 | print("Token is invalid (placeholder). Denying request.") 90 | # Return unauthorized for HTTP API simple response 91 | return {"isAuthorized": False} 92 | # For REST API: raise Exception('Unauthorized') # Or return deny policy 93 | 94 | 95 | # --- Example AuthPolicy class for REST API IAM Policy response --- 96 | # (Include this class if you need to return IAM policies for REST APIs) 97 | # class AuthPolicy(object): 98 | # # ... (Implementation of AuthPolicy class) ... 99 | # pass 100 | -------------------------------------------------------------------------------- /src/auth/cognito.ts: -------------------------------------------------------------------------------- 1 | import type { JWTVerifyResult } from 'jose'; 2 | import { Request } from 'express'; 3 | 4 | import jwksClient from 'jwks-rsa'; 5 | import * as https from 'https'; 6 | import { promisify } from 'util'; 7 | import { logger, authLogger } from '../logging'; 8 | 9 | /** 10 | * Configuration for AWS Cognito integration 11 | */ 12 | export interface CognitoConfig { 13 | /** AWS Cognito Region */ 14 | region: string; 15 | /** AWS Cognito User Pool ID */ 16 | userPoolId: string; 17 | /** Client IDs allowed to access this server */ 18 | allowedClientIds: string[]; 19 | /** Resource server identifier (e.g., API identifier in Cognito) */ 20 | resourceId?: string; 21 | /** Required scopes (if any) */ 22 | requiredScopes?: string[]; 23 | } 24 | 25 | /** 26 | * Creates the JWKS URL for a given Cognito user pool 27 | */ 28 | function getJwksUrl(region: string, userPoolId: string): string { 29 | return `https://cognito-idp.${region}.amazonaws.com/${userPoolId}/.well-known/jwks.json`; 30 | } 31 | 32 | interface JWK { 33 | alg: string; 34 | e: string; 35 | kid: string; 36 | kty: string; 37 | n: string; 38 | use: string; 39 | } 40 | 41 | interface JWKS { 42 | keys: JWK[]; 43 | } 44 | 45 | /** 46 | * Class for handling Cognito JWT verification 47 | */ 48 | export class CognitoAuthorizer { 49 | private jwksUrl: string; 50 | private jwksClient: any; 51 | private config: CognitoConfig; 52 | private issuer: string; 53 | private isLambda: boolean; 54 | private jwksCache: JWKS | null = null; 55 | private jwksCacheTime = 0; 56 | private readonly jwksCacheMaxAge = 3600000; // 1 hour in milliseconds 57 | 58 | /** 59 | * Creates a new instance of CognitoAuthorizer 60 | */ 61 | constructor(config: CognitoConfig) { 62 | this.config = config; 63 | this.jwksUrl = getJwksUrl(config.region, config.userPoolId); 64 | this.issuer = `https://cognito-idp.${config.region}.amazonaws.com/${config.userPoolId}`; 65 | this.isLambda = !!process.env.LAMBDA_TASK_ROOT; 66 | 67 | // Use jwks-rsa client which is CommonJS compatible 68 | this.jwksClient = jwksClient({ 69 | jwksUri: this.jwksUrl, 70 | cache: true, 71 | cacheMaxEntries: 5, 72 | cacheMaxAge: 600000 // 10 minutes 73 | }); 74 | 75 | console.log(`Initialized Cognito Authorizer with JWKS URL: ${this.jwksUrl}`); 76 | } 77 | 78 | /** 79 | * Extracts the authorization token from the request headers 80 | */ 81 | extractToken(req: Request): string | null { 82 | const authHeader = req.headers.authorization; 83 | if (!authHeader || !authHeader.startsWith('Bearer ')) { 84 | return null; 85 | } 86 | 87 | return authHeader.substring(7); 88 | } 89 | 90 | /** 91 | * Fetch the JSON Web Key Set (JWKS) from the Cognito endpoint 92 | * Uses a memory cache to avoid frequent network requests 93 | * 94 | * @returns Promise resolving to the JWKS 95 | * @throws Error if fetching the JWKS fails 96 | */ 97 | private async fetchJWKS(): Promise { 98 | try { 99 | // Use cached JWKS if available and not expired 100 | const now = Date.now(); 101 | if (this.jwksCache && now - this.jwksCacheTime < this.jwksCacheMaxAge) { 102 | return this.jwksCache; 103 | } 104 | 105 | authLogger.debug('Fetching JWKS from Cognito', { url: this.jwksUrl }); 106 | 107 | // Fetch the JWKS with a Promise-based approach 108 | const jwksData = await new Promise((resolve, reject) => { 109 | https.get(this.jwksUrl, (res) => { 110 | if (res.statusCode !== 200) { 111 | reject(new Error(`Failed to fetch JWKS: HTTP ${res.statusCode}`)); 112 | return; 113 | } 114 | 115 | let data = ''; 116 | res.on('data', (chunk) => { 117 | data += chunk; 118 | }); 119 | 120 | res.on('end', () => { 121 | try { 122 | const jwks = JSON.parse(data) as JWKS; 123 | resolve(jwks); 124 | } catch (e) { 125 | reject(new Error(`Failed to parse JWKS: ${e instanceof Error ? e.message : String(e)}`)); 126 | } 127 | }); 128 | }).on('error', (e) => { 129 | reject(new Error(`Failed to fetch JWKS: ${e.message}`)); 130 | }); 131 | }); 132 | 133 | // Update the cache 134 | this.jwksCache = jwksData; 135 | this.jwksCacheTime = now; 136 | 137 | authLogger.debug('Successfully fetched JWKS', { 138 | keyCount: jwksData.keys.length, 139 | firstKeyId: jwksData.keys[0]?.kid 140 | }); 141 | 142 | return jwksData; 143 | } catch (error) { 144 | authLogger.error('Error fetching JWKS', { 145 | error: error instanceof Error ? error.message : String(error), 146 | url: this.jwksUrl 147 | }); 148 | throw error; 149 | } 150 | } 151 | 152 | /** 153 | * Verify the signature of a JWT token using the RSA public key from JWKS 154 | * 155 | * @param token - The JWT token to verify 156 | * @param header - The parsed JWT header containing the key ID (kid) 157 | * @param signature - The signature part of the JWT 158 | * @returns Promise resolving to boolean indicating if signature is valid 159 | * @throws Error if verification fails 160 | */ 161 | private async verifySignature(token: string, header: any, signature: Buffer): Promise { 162 | try { 163 | // Get the signing key from JWKS using the key ID in the token header 164 | const kid = header.kid; 165 | if (!kid) { 166 | throw new Error('Token header missing key ID (kid)'); 167 | } 168 | 169 | // Fetch JWKS and find the matching key 170 | const jwks = await this.fetchJWKS(); 171 | const key = jwks.keys.find(k => k.kid === kid); 172 | 173 | if (!key) { 174 | throw new Error(`Unable to find key with ID: ${kid}`); 175 | } 176 | 177 | // If the key is not an RSA key or not for signature verification 178 | if (key.kty !== 'RSA' || key.use !== 'sig') { 179 | throw new Error(`Invalid key type: ${key.kty} or use: ${key.use}`); 180 | } 181 | 182 | // Use the jwks-rsa client to get the signing key 183 | const getSigningKey = promisify(this.jwksClient.getSigningKey).bind(this.jwksClient); 184 | const signingKey = await getSigningKey(kid); 185 | const publicKey = signingKey.getPublicKey(); 186 | 187 | // Verify the token with the public key 188 | // Using Node's built-in crypto for verification 189 | const crypto = require('crypto'); 190 | const verifier = crypto.createVerify('RSA-SHA256'); 191 | 192 | // The content to verify is the header and payload parts of the JWT 193 | const parts = token.split('.'); 194 | const signedContent = parts[0] + '.' + parts[1]; 195 | verifier.update(signedContent); 196 | 197 | // Verify the signature 198 | return verifier.verify( 199 | publicKey, 200 | signature, 201 | 'base64url' 202 | ); 203 | } catch (error) { 204 | authLogger.error('Signature verification failed', { 205 | error: error instanceof Error ? error.message : String(error) 206 | }); 207 | throw error; 208 | } 209 | } 210 | 211 | /** 212 | * Verifies a JWT token against the Cognito JWKS 213 | * 214 | * This method validates the token format, expiration, signature, 215 | * and checks additional claims like token_use, client_id, and required scopes. 216 | * 217 | * @param token - The JWT token to verify 218 | * @returns A promise resolving to the verification result with payload and header 219 | * @throws Error if the token is invalid, expired, or missing required claims 220 | */ 221 | async verifyToken(token: string): Promise { 222 | try { 223 | authLogger.debug('Starting JWT token verification'); 224 | 225 | // Basic token format validation 226 | const tokenParts = token.split('.'); 227 | if (tokenParts.length !== 3) { 228 | authLogger.warn('Invalid token format (not three parts)'); 229 | throw new Error('Invalid token format'); 230 | } 231 | 232 | // Decode header 233 | let header: any; 234 | try { 235 | const headerJson = Buffer.from(tokenParts[0], 'base64url').toString('utf8'); 236 | header = JSON.parse(headerJson); 237 | } catch (e) { 238 | authLogger.warn('Failed to parse token header'); 239 | throw new Error('Invalid token header'); 240 | } 241 | 242 | // Verify token algorithm 243 | if (header.alg !== 'RS256') { 244 | authLogger.warn(`Unsupported token algorithm: ${header.alg}`); 245 | throw new Error(`Unsupported token algorithm: ${header.alg}`); 246 | } 247 | 248 | // Decode payload 249 | let payload: any; 250 | try { 251 | const payloadJson = Buffer.from(tokenParts[1], 'base64url').toString('utf8'); 252 | payload = JSON.parse(payloadJson); 253 | } catch (e) { 254 | authLogger.warn('Failed to parse token payload'); 255 | throw new Error('Invalid token payload'); 256 | } 257 | 258 | // Decode signature 259 | let signature: Buffer; 260 | try { 261 | signature = Buffer.from(tokenParts[2], 'base64url'); 262 | } catch (e) { 263 | authLogger.warn('Failed to decode token signature'); 264 | throw new Error('Invalid token signature'); 265 | } 266 | 267 | // Log token info (redacted for security) 268 | authLogger.debug('Token info', { 269 | iss: payload.iss, 270 | sub: typeof payload.sub === 'string' ? payload.sub.substring(0, 8) + '...' : undefined, 271 | client_id: payload.client_id, 272 | token_use: payload.token_use, 273 | exp: payload.exp, 274 | iat: payload.iat, 275 | username: typeof payload.username === 'string' ? payload.username.substring(0, 3) + '...' : undefined 276 | }); 277 | 278 | // Verify expiration 279 | const nowInSeconds = Math.floor(Date.now() / 1000); 280 | if (!payload.exp) { 281 | authLogger.warn('Token missing expiration claim'); 282 | throw new Error('Token missing expiration time'); 283 | } 284 | 285 | if (payload.exp < nowInSeconds) { 286 | const expiredAt = new Date(payload.exp * 1000).toISOString(); 287 | const currentTime = new Date().toISOString(); 288 | authLogger.warn(`Token expired at ${expiredAt}, current time is ${currentTime}`); 289 | throw new Error(`Token expired at ${expiredAt}`); 290 | } 291 | 292 | // Verify not before (nbf) if present 293 | if (payload.nbf && payload.nbf > nowInSeconds) { 294 | const notBefore = new Date(payload.nbf * 1000).toISOString(); 295 | const currentTime = new Date().toISOString(); 296 | authLogger.warn(`Token not valid until ${notBefore}, current time is ${currentTime}`); 297 | throw new Error(`Token not valid until ${notBefore}`); 298 | } 299 | 300 | // Verify issuer if configured 301 | if (this.config.userPoolId) { 302 | const expectedIssuer = `https://cognito-idp.${this.config.region}.amazonaws.com/${this.config.userPoolId}`; 303 | if (payload.iss !== expectedIssuer) { 304 | authLogger.warn(`Invalid token issuer: ${payload.iss}, expected: ${expectedIssuer}`); 305 | throw new Error(`Invalid token issuer: ${payload.iss}`); 306 | } 307 | } 308 | 309 | // FULL SIGNATURE VERIFICATION WITH JWKS 310 | authLogger.debug('Verifying token signature with JWKS'); 311 | const isSignatureValid = await this.verifySignature(token, header, signature); 312 | 313 | if (!isSignatureValid) { 314 | authLogger.warn('Invalid token signature'); 315 | throw new Error('Invalid token signature'); 316 | } 317 | 318 | // Create a result similar to jose's jwtVerify 319 | const result = { 320 | payload: payload, 321 | protectedHeader: header 322 | }; 323 | 324 | // Validate additional claims 325 | authLogger.debug('Validating additional token claims'); 326 | 327 | // Check for token_use (must be 'access') 328 | if (payload.token_use !== 'access') { 329 | authLogger.warn(`Invalid token use: ${payload.token_use} (expected 'access')`); 330 | throw new Error(`Invalid token use: ${payload.token_use}`); 331 | } 332 | 333 | // Check client ID is in allowed list 334 | if (payload.client_id && !this.config.allowedClientIds.includes(payload.client_id as string)) { 335 | authLogger.warn(`Unauthorized client ID: ${payload.client_id}`); 336 | throw new Error(`Unauthorized client ID: ${payload.client_id}`); 337 | } 338 | 339 | // Check required scopes 340 | if (this.config.requiredScopes && this.config.requiredScopes.length > 0) { 341 | const tokenScopes = typeof payload.scope === 'string' 342 | ? payload.scope.split(' ') 343 | : []; 344 | 345 | authLogger.debug('Checking token scopes', { 346 | tokenScopes, 347 | requiredScopes: this.config.requiredScopes 348 | }); 349 | 350 | const hasAllRequiredScopes = this.config.requiredScopes.every( 351 | scope => tokenScopes.includes(scope) 352 | ); 353 | 354 | if (!hasAllRequiredScopes) { 355 | authLogger.warn('Token missing required scopes'); 356 | throw new Error('Token missing required scopes'); 357 | } 358 | } 359 | 360 | authLogger.debug('Token verification completed successfully'); 361 | return result; 362 | } catch (error) { 363 | authLogger.error('Token verification failed', { error }); 364 | throw error; 365 | } 366 | } 367 | 368 | /** 369 | * Authorized middleware factory for Express 370 | * 371 | * Creates an Express middleware function that validates JWT tokens 372 | * from the Authorization header and attaches the verified claims to the request. 373 | * 374 | * @returns Express middleware function for authentication 375 | */ 376 | createAuthMiddleware() { 377 | return async (req: Request, res: any, next: any) => { 378 | try { 379 | authLogger.debug({ 380 | message: 'Processing authentication for request', 381 | path: req.path, 382 | method: req.method, 383 | hasAuthHeader: !!req.headers.authorization 384 | }); 385 | 386 | const token = this.extractToken(req); 387 | if (!token) { 388 | authLogger.debug('No token found in authorization header'); 389 | return this.handleUnauthorized(req, res, 'Bearer token missing'); 390 | } 391 | 392 | authLogger.debug('Token extracted, attempting verification'); 393 | const verifyResult = await this.verifyToken(token); 394 | 395 | authLogger.info({ 396 | message: 'Token verified successfully', 397 | subject: verifyResult.payload.sub, 398 | username: verifyResult.payload.username, 399 | clientId: verifyResult.payload.client_id, 400 | tokenUse: verifyResult.payload.token_use 401 | }); 402 | 403 | // Attach the verified claims to the request for later use 404 | (req as any).user = verifyResult.payload; 405 | 406 | // Attach helper method to check scopes 407 | (req as any).hasScope = (scope: string): boolean => { 408 | const tokenScopes = typeof verifyResult.payload.scope === 'string' 409 | ? verifyResult.payload.scope.split(' ') 410 | : []; 411 | 412 | return tokenScopes.includes(scope); 413 | }; 414 | 415 | next(); 416 | } catch (error) { 417 | authLogger.error({ 418 | message: 'Authentication failed', 419 | error 420 | }); 421 | 422 | return this.handleUnauthorized(req, res, error); 423 | } 424 | }; 425 | } 426 | 427 | /** 428 | * Handle unauthorized access according to the MCP Auth specification 429 | * 430 | * This method generates a proper WWW-Authenticate header as required by the MCP Auth spec, 431 | * including the resource metadata URI and optional error information. 432 | * 433 | * @param req - The Express request object 434 | * @param res - The Express response object 435 | * @param error - Optional error message or object to include in the response 436 | * @returns The response object with appropriate status and headers 437 | */ 438 | handleUnauthorized(req: Request, res: any, error?: string | Error) { 439 | // Get the protocol and host from request headers or config 440 | const defaultProtocol = this.isLambda ? 'https' : 'http'; 441 | const protocol = req.headers['x-forwarded-proto'] || defaultProtocol; 442 | const host = req.headers.host || req.hostname; 443 | 444 | // Get the stage from the request path or environment 445 | let stage = ''; 446 | 447 | // If running in Lambda with API Gateway, we need the stage name 448 | if (this.isLambda) { 449 | // Extract from original URL if available 450 | const url = req.originalUrl || req.url || ''; 451 | 452 | // Check if URL starts with a stage path like /prod/ 453 | if (url.startsWith('/prod/')) { 454 | stage = '/prod'; 455 | } else { 456 | // Default to /prod for Lambda environments 457 | stage = '/prod'; 458 | } 459 | } 460 | 461 | // Construct the resource metadata URL with stage if needed 462 | const resourceMetadataUrl = `${protocol}://${host}${stage}/.well-known/oauth-protected-resource`; 463 | 464 | // Build the WWW-Authenticate header according to RFC7235 and OAuth 2.0 specs 465 | let wwwAuthHeader = `Bearer realm="${this.issuer}", resource_metadata_uri="${resourceMetadataUrl}"`; 466 | 467 | // Add error details if provided 468 | if (error) { 469 | const errorMsg = error instanceof Error ? error.message : error; 470 | wwwAuthHeader += `, error="invalid_token", error_description="${errorMsg}"`; 471 | } 472 | 473 | // Log the response for debugging 474 | authLogger.debug({ 475 | message: 'Sending unauthorized response', 476 | resourceMetadataUrl, 477 | error: error instanceof Error ? error.message : error, 478 | headers: { 479 | 'WWW-Authenticate': wwwAuthHeader 480 | } 481 | }); 482 | 483 | // Set WWW-Authenticate header as required by MCP Auth 484 | res.setHeader('WWW-Authenticate', wwwAuthHeader); 485 | 486 | // Add CORS headers to ensure the client can read the response 487 | res.setHeader('Access-Control-Allow-Origin', '*'); 488 | res.setHeader('Access-Control-Allow-Headers', 'Authorization, Content-Type'); 489 | res.setHeader('Access-Control-Expose-Headers', 'WWW-Authenticate'); 490 | 491 | // Return a JSON-RPC 2.0 error response 492 | return res.status(401).json({ 493 | jsonrpc: '2.0', 494 | error: { 495 | code: -32001, 496 | message: error instanceof Error ? 497 | `Unauthorized: ${error.message}` : 498 | 'Unauthorized: Bearer token required', 499 | }, 500 | id: null, 501 | }); 502 | } 503 | } 504 | 505 | /** 506 | * Creates a configured Cognito authorizer based on environment variables 507 | */ 508 | export function createCognitoAuthorizer(): CognitoAuthorizer { 509 | const region = process.env.COGNITO_REGION || 'us-east-1'; 510 | const userPoolId = process.env.COGNITO_USER_POOL_ID || 'us-east-1_IgrnjnCts'; 511 | 512 | // Parse allowed client IDs 513 | const clientIds = (process.env.COGNITO_ALLOWED_CLIENT_IDS || 514 | '7dmhq1mos3d41k85u10inppvvi,3o7qjhi1d6nap68pdnrobgdurh').split(','); 515 | 516 | // Optional scopes 517 | const scopes = process.env.COGNITO_REQUIRED_SCOPES?.split(',') || []; 518 | 519 | return new CognitoAuthorizer({ 520 | region, 521 | userPoolId, 522 | allowedClientIds: clientIds, 523 | requiredScopes: scopes, 524 | resourceId: process.env.COGNITO_RESOURCE_ID, 525 | }); 526 | } -------------------------------------------------------------------------------- /src/auth/index.ts: -------------------------------------------------------------------------------- 1 | /** 2 | * Authorization Module Index 3 | * 4 | * This file exports the authorization methods supported by the MCP server. 5 | * It provides a factory function to create the appropriate authorizer based on configuration. 6 | */ 7 | 8 | import { CognitoAuthorizer, createCognitoAuthorizer } from './cognito'; 9 | import { LambdaAuthorizer, createLambdaAuthorizer } from './lambdaAuth'; 10 | 11 | /** 12 | * Creates an authorizer based on the specified type 13 | * 14 | * @param type The type of authorizer to create ('oauth' or 'lambda') 15 | * @returns The appropriate authorizer instance 16 | */ 17 | export function createAuthorizer(type: 'oauth' | 'lambda') { 18 | if (type === 'oauth') { 19 | return createCognitoAuthorizer(); 20 | } else { 21 | return createLambdaAuthorizer(); 22 | } 23 | } 24 | 25 | export { CognitoAuthorizer, LambdaAuthorizer }; 26 | -------------------------------------------------------------------------------- /src/auth/lambdaAuth.ts: -------------------------------------------------------------------------------- 1 | /** 2 | * Lambda Authorizer Integration 3 | * 4 | * This file provides integration with the Lambda authorizer approach. 5 | * It creates a middleware that bypasses authorization checks in the Express app 6 | * since API Gateway with Lambda authorizer will handle authorization before 7 | * the request reaches the Lambda function. 8 | */ 9 | import { Request, Response, NextFunction } from 'express'; 10 | import { logger, authLogger } from '../logging'; 11 | 12 | /** 13 | * Configuration for Lambda authorizer integration 14 | */ 15 | export interface LambdaAuthConfig { 16 | /** Whether Lambda authorizer is enabled */ 17 | enabled: boolean; 18 | } 19 | 20 | /** 21 | * Class for handling Lambda authorizer integration 22 | */ 23 | export class LambdaAuthorizer { 24 | private config: LambdaAuthConfig; 25 | private isLambda: boolean; 26 | 27 | /** 28 | * Creates a new instance of LambdaAuthorizer 29 | */ 30 | constructor(config: LambdaAuthConfig) { 31 | this.config = config; 32 | this.isLambda = !!process.env.LAMBDA_TASK_ROOT; 33 | 34 | console.log(`Initialized Lambda Authorizer with enabled=${config.enabled}`); 35 | } 36 | 37 | /** 38 | * Creates a middleware that bypasses authorization checks 39 | * since API Gateway with Lambda authorizer will handle authorization 40 | * before the request reaches the Lambda function. 41 | */ 42 | createAuthMiddleware() { 43 | return async (req: Request, res: Response, next: NextFunction) => { 44 | try { 45 | authLogger.debug({ 46 | message: 'Lambda authorizer middleware processing request', 47 | path: req.path, 48 | method: req.method, 49 | hasAuthHeader: !!req.headers.authorization 50 | }); 51 | 52 | // In Lambda environment, API Gateway with Lambda authorizer 53 | // will handle authorization before the request reaches this function 54 | if (this.isLambda) { 55 | authLogger.debug('Running in Lambda environment - authorization handled by API Gateway'); 56 | next(); 57 | return; 58 | } 59 | 60 | // In development environment, we can do basic validation 61 | // or just pass through for testing purposes 62 | const token = this.extractToken(req); 63 | if (!token) { 64 | authLogger.debug('No token found in authorization header'); 65 | return this.handleUnauthorized(req, res, 'Bearer token missing'); 66 | } 67 | 68 | // For development, we accept any token 69 | // In production, the Lambda authorizer will handle validation 70 | authLogger.info({ 71 | message: 'Development mode: Accepting token without validation', 72 | token: token.substring(0, 8) + '...' 73 | }); 74 | 75 | next(); 76 | } catch (error) { 77 | authLogger.error({ 78 | message: 'Error in Lambda authorizer middleware', 79 | error 80 | }); 81 | 82 | return this.handleUnauthorized(req, res, error); 83 | } 84 | }; 85 | } 86 | 87 | /** 88 | * Extracts the authorization token from the request headers 89 | */ 90 | extractToken(req: Request): string | null { 91 | const authHeader = req.headers.authorization; 92 | if (!authHeader || !authHeader.startsWith('Bearer ')) { 93 | return null; 94 | } 95 | 96 | return authHeader.substring(7); 97 | } 98 | 99 | /** 100 | * Handle unauthorized access 101 | */ 102 | handleUnauthorized(req: Request, res: Response, error?: string | Error) { 103 | const errorMsg = error instanceof Error ? error.message : error; 104 | 105 | authLogger.debug({ 106 | message: 'Sending unauthorized response', 107 | error: errorMsg 108 | }); 109 | 110 | // Return a JSON-RPC 2.0 error response 111 | return res.status(401).json({ 112 | jsonrpc: '2.0', 113 | error: { 114 | code: -32001, 115 | message: error instanceof Error ? 116 | `Unauthorized: ${error.message}` : 117 | 'Unauthorized: Bearer token required', 118 | }, 119 | id: null, 120 | }); 121 | } 122 | } 123 | 124 | /** 125 | * Creates a configured Lambda authorizer 126 | */ 127 | export function createLambdaAuthorizer(): LambdaAuthorizer { 128 | return new LambdaAuthorizer({ 129 | enabled: true 130 | }); 131 | } 132 | -------------------------------------------------------------------------------- /src/auth/oauthMetadata.ts: -------------------------------------------------------------------------------- 1 | /** 2 | * Implementation of OAuth 2.0 Protected Resource Metadata (RFC9728) 3 | * https://datatracker.ietf.org/doc/html/rfc9728 4 | * 5 | * This implementation provides the discovery mechanism for OAuth 2.0 clients 6 | * to locate the appropriate authorization server for a protected resource. 7 | */ 8 | 9 | import { Request, Response } from 'express'; 10 | import { logger, authLogger } from '../logging'; 11 | 12 | /** 13 | * OAuth 2.0 Protected Resource Metadata configuration 14 | */ 15 | export interface ResourceMetadataConfig { 16 | /** Resource identifier */ 17 | resourceId: string; 18 | /** List of authorization servers */ 19 | authorizationServers: AuthorizationServer[]; 20 | /** Additional metadata properties */ 21 | additionalMetadata?: Record; 22 | } 23 | 24 | /** 25 | * Authorization server configuration 26 | */ 27 | export interface AuthorizationServer { 28 | /** The authorization server URL (issuer) */ 29 | url: string; 30 | /** Authorization server metadata URL (optional, will be derived if not provided) */ 31 | metadataUrl?: string; 32 | /** Authorization scopes required for this resource (optional) */ 33 | scopes?: string[]; 34 | } 35 | 36 | /** 37 | * Creates a response handler for the /.well-known/oauth-protected-resource endpoint 38 | * according to RFC9728 (OAuth 2.0 Protected Resource Metadata) 39 | */ 40 | export function createResourceMetadataHandler(config: ResourceMetadataConfig) { 41 | return (req: Request, res: Response) => { 42 | try { 43 | authLogger.debug({ 44 | message: 'Resource metadata request received', 45 | path: req.path, 46 | method: req.method, 47 | ip: req.ip 48 | }); 49 | 50 | // Build the OAuth 2.0 Protected Resource Metadata response 51 | const metadata = { 52 | resource: config.resourceId, 53 | authorization_servers: config.authorizationServers.map(server => { 54 | // If metadataUrl is not provided, construct it from the server URL 55 | const metadataUrl = server.metadataUrl || 56 | `${server.url}/.well-known/oauth-authorization-server`; 57 | 58 | const serverEntry: Record = { 59 | issuer: server.url, 60 | authorization_server_metadata_url: metadataUrl 61 | }; 62 | 63 | // Include scopes if provided 64 | if (server.scopes && server.scopes.length > 0) { 65 | serverEntry.scopes_supported = server.scopes; 66 | } 67 | 68 | return serverEntry; 69 | }), 70 | ...config.additionalMetadata 71 | }; 72 | 73 | // Set CORS headers to ensure the metadata is accessible from any origin 74 | res.header('Access-Control-Allow-Origin', '*'); 75 | res.header('Access-Control-Allow-Methods', 'GET, OPTIONS'); 76 | res.header('Access-Control-Allow-Headers', 'Content-Type, Authorization'); 77 | res.header('Access-Control-Max-Age', '86400'); // 24 hours 78 | 79 | // Set cache-control to enable caching (recommended by RFC9728) 80 | res.header('Cache-Control', 'public, max-age=3600'); 81 | 82 | // Set content type to application/json 83 | res.header('Content-Type', 'application/json'); 84 | 85 | authLogger.debug({ 86 | message: 'Returning resource metadata', 87 | resourceId: config.resourceId, 88 | authServers: config.authorizationServers.length 89 | }); 90 | 91 | // Return the metadata as JSON 92 | res.json(metadata); 93 | } catch (error) { 94 | authLogger.error({ 95 | message: 'Error generating OAuth resource metadata', 96 | error 97 | }); 98 | 99 | // Return a properly formatted OAuth error response 100 | res.status(500).json({ 101 | error: 'server_error', 102 | error_description: 'An error occurred while generating the OAuth resource metadata' 103 | }); 104 | } 105 | }; 106 | } 107 | 108 | /** 109 | * Creates default resource metadata configuration from environment variables 110 | */ 111 | export function getDefaultResourceMetadataConfig(): ResourceMetadataConfig { 112 | const region = process.env.COGNITO_REGION || 'us-east-1'; 113 | const userPoolId = process.env.COGNITO_USER_POOL_ID || 'us-east-1_IgrnjnCts'; 114 | // Construct the Cognito domain following the format: 115 | // https://{domain-prefix}.auth.{region}.amazoncognito.com 116 | // For user pool IDs like 'us-east-1_IgrnjnCts', the domain prefix should be 'us-east-1-igrnjncts' 117 | let userPoolDomain = process.env.COGNITO_DOMAIN; 118 | 119 | if (!userPoolDomain) { 120 | // If domain isn't set, construct it from user pool ID 121 | const parts = userPoolId.split('_'); 122 | if (parts.length === 2) { 123 | const regionPrefix = parts[0]; 124 | const idPart = parts[1].toLowerCase(); 125 | userPoolDomain = `${regionPrefix}-${idPart}.auth.${region}.amazoncognito.com`; 126 | } else { 127 | // Fallback if format is unexpected 128 | userPoolDomain = `${userPoolId.toLowerCase().replace('_', '-')}.auth.${region}.amazoncognito.com`; 129 | } 130 | } else if (!userPoolDomain.includes('.auth.') && !userPoolDomain.includes('.amazoncognito.com')) { 131 | // If domain is provided but doesn't include the full domain suffix, add it 132 | // This handles cases where just the prefix is provided (e.g., "us-east-1-igrnjncts") 133 | userPoolDomain = `${userPoolDomain}.auth.${region}.amazoncognito.com`; 134 | } 135 | 136 | // Log the constructed domain for debugging 137 | logger.debug(`Using Cognito domain: ${userPoolDomain}`); 138 | 139 | // Get the resource ID (API identifier) 140 | const resourceId = process.env.COGNITO_RESOURCE_ID || 'mcp-server-api'; 141 | 142 | // Build the authorization server URL 143 | const authServerUrl = `https://${userPoolDomain}`; 144 | 145 | // Get supported scopes from environment variables or use defaults 146 | const scopes = process.env.COGNITO_SUPPORTED_SCOPES ? 147 | process.env.COGNITO_SUPPORTED_SCOPES.split(',') : 148 | ['openid', 'profile', 'email']; 149 | 150 | // Log the configuration 151 | logger.info('Creating default OAuth 2.0 Protected Resource Metadata configuration', { 152 | resourceId, 153 | authServerUrl, 154 | region, 155 | userPoolId 156 | }); 157 | 158 | // For Cognito, we need to provide both the domain URL and the service URL 159 | // The OIDC discovery endpoint is available at the service URL, not the domain URL 160 | const serviceUrl = `https://cognito-idp.${region}.amazonaws.com/${userPoolId}`; 161 | 162 | return { 163 | resourceId, 164 | authorizationServers: [ 165 | { 166 | url: authServerUrl, 167 | // Use Cognito-specific metadata URL format for OpenID Connect discovery 168 | metadataUrl: `${serviceUrl}/.well-known/openid-configuration`, 169 | scopes 170 | } 171 | ], 172 | additionalMetadata: { 173 | protocol_version: "2025-03-26", 174 | description: "MCP Server Protected Resource", 175 | api_documentation: "https://modelcontextprotocol.io/specification/draft/basic/authorization" 176 | } 177 | }; 178 | } -------------------------------------------------------------------------------- /src/logging.js: -------------------------------------------------------------------------------- 1 | import log4js from 'log4js'; 2 | import { unlinkSync } from 'fs'; 3 | 4 | // Check if we're running in Lambda 5 | const isLambda = !!process.env.LAMBDA_TASK_ROOT; 6 | 7 | // Set log file path based on environment 8 | // In Lambda, use /tmp which is writable 9 | const LOG_FILE = isLambda ? '/tmp/server.log' : './server.log'; 10 | 11 | // Try to clean up old log files if they exist 12 | // try { unlinkSync(LOG_FILE); } catch (e) { }; 13 | 14 | // Define custom JSON layout for structured logging 15 | const jsonLayout = { 16 | type: 'pattern', 17 | pattern: '%d{ISO8601_WITH_TZ_OFFSET} %p %c %m', 18 | // Transform log events to structured JSON format 19 | transform: (logEvent) => { 20 | const { startTime, categoryName, level, data } = logEvent; 21 | 22 | // Basic log structure 23 | const result = { 24 | timestamp: new Date(startTime).toISOString(), 25 | level: level.levelStr, 26 | category: categoryName, 27 | }; 28 | 29 | // Handle different message formats 30 | if (data.length === 1) { 31 | // Simple string message 32 | if (typeof data[0] === 'string') { 33 | result.message = data[0]; 34 | } 35 | // Object that's already structured 36 | else if (typeof data[0] === 'object') { 37 | Object.assign(result, data[0]); 38 | } 39 | } else if (data.length > 1) { 40 | // Format with message and data 41 | result.message = data[0]; 42 | result.data = data.slice(1); 43 | } 44 | 45 | // Include error stack if present 46 | if (logEvent.error) { 47 | result.error = { 48 | message: logEvent.error.message, 49 | name: logEvent.error.name, 50 | stack: logEvent.error.stack 51 | }; 52 | } 53 | 54 | return JSON.stringify(result); 55 | } 56 | }; 57 | 58 | // Define appenders based on environment 59 | const appenders = { 60 | stdout: { 61 | type: 'stdout', 62 | enableCallStack: true, 63 | layout: { 64 | type: 'pattern', 65 | pattern: '%[%p [%f{1}:%l:%M] %m%]' 66 | } 67 | }, 68 | stdoutJson: { 69 | type: 'stdout', 70 | layout: jsonLayout 71 | } 72 | }; 73 | 74 | // Only add file appenders if not in Lambda 75 | if (!isLambda) { 76 | appenders.file = { 77 | type: 'file', 78 | filename: LOG_FILE, 79 | enableCallStack: true, 80 | layout: { 81 | type: 'pattern', 82 | pattern: '%[%p [%f{1}:%l:%M] %m%]' 83 | } 84 | }; 85 | appenders.fileJson = { 86 | type: 'file', 87 | filename: LOG_FILE + '.json', 88 | layout: jsonLayout 89 | }; 90 | } 91 | 92 | // Configure log4js with appropriate appenders 93 | log4js.configure({ 94 | appenders: appenders, 95 | categories: { 96 | default: { 97 | appenders: isLambda ? ['stdout'] : ['stdout', 'fileJson'], 98 | level: 'debug', 99 | enableCallStack: true 100 | }, 101 | auth: { 102 | appenders: isLambda ? ['stdout'] : ['stdout', 'fileJson'], 103 | level: 'debug', 104 | enableCallStack: true 105 | }, 106 | api: { 107 | appenders: isLambda ? ['stdout'] : ['stdout', 'fileJson'], 108 | level: 'debug', 109 | enableCallStack: true 110 | } 111 | } 112 | }); 113 | 114 | // Create logger instances for different components 115 | export const logger = log4js.getLogger(); 116 | export const authLogger = log4js.getLogger('auth'); 117 | export const apiLogger = log4js.getLogger('api'); 118 | 119 | // Helper function to log HTTP requests in a structured way 120 | export const logHttpRequest = (req, componentName = 'api') => { 121 | const logData = { 122 | type: 'request', 123 | method: req.method, 124 | path: req.path, 125 | headers: { 126 | // Include only relevant headers and redact sensitive information 127 | 'content-type': req.headers['content-type'], 128 | 'mcp-session-id': req.headers['mcp-session-id'], 129 | 'mcp-protocol-version': req.headers['mcp-protocol-version'], 130 | // Indicate auth header presence without revealing token 131 | 'authorization': req.headers.authorization ? 'Bearer [redacted]' : undefined, 132 | }, 133 | ip: req.ip, 134 | userAgent: req.headers['user-agent'], 135 | component: componentName 136 | }; 137 | 138 | log4js.getLogger(componentName).info(logData); 139 | 140 | return logData; 141 | }; 142 | 143 | // Helper function to log HTTP responses in a structured way 144 | export const logHttpResponse = (res, duration, componentName = 'api') => { 145 | const logData = { 146 | type: 'response', 147 | statusCode: res.statusCode, 148 | duration: duration, 149 | component: componentName 150 | }; 151 | 152 | log4js.getLogger(componentName).info(logData); 153 | 154 | return logData; 155 | }; 156 | 157 | // Export default logger for backwards compatibility 158 | export default logger; -------------------------------------------------------------------------------- /src/mcpClientUtils.ts: -------------------------------------------------------------------------------- 1 | /** 2 | * MCP Client Utilities 3 | * 4 | * This file provides utility functions for the MCP client, including: 5 | * - OAuth authentication helpers (PKCE, token management) 6 | * - Secure token storage and retrieval 7 | * - Input/output utilities 8 | * 9 | * These utilities abstract away the complexity of authentication, 10 | * security, and user interaction for the main client code. 11 | */ 12 | 13 | import * as fs from 'fs'; 14 | import * as path from 'path'; 15 | import * as os from 'os'; 16 | import { createHash, randomBytes } from 'crypto'; 17 | import { createInterface } from 'node:readline'; 18 | import { exec } from 'child_process'; 19 | 20 | // Path to store tokens and PKCE parameters 21 | export const configDir = path.join(os.homedir(), '.mcp-client'); 22 | export const tokenStore = path.join(configDir, 'tokens.json'); 23 | export const pkceStore = path.join(configDir, 'pkce.json'); 24 | 25 | // Make sure the config directory exists 26 | if (!fs.existsSync(configDir)) { 27 | fs.mkdirSync(configDir, { recursive: true }); 28 | } 29 | 30 | /** 31 | * Token-related types 32 | */ 33 | 34 | export interface TokenData { 35 | access_token: string; 36 | token_type: string; 37 | expires_in: number; 38 | refresh_token?: string; 39 | id_token?: string; 40 | iat?: number; 41 | [key: string]: any; 42 | } 43 | 44 | export interface PKCEParams { 45 | codeVerifier: string; 46 | codeChallenge: string; 47 | state: string; 48 | redirectUri: string; 49 | timestamp?: number; 50 | } 51 | 52 | /** 53 | * Helper function to prompt for user input 54 | */ 55 | export function promptForInput(question: string, readline: ReturnType): Promise { 56 | return new Promise((resolve) => { 57 | readline.question(question, (answer) => { 58 | resolve(answer); 59 | }); 60 | }); 61 | } 62 | 63 | /** 64 | * Generate a cryptographically secure random string of specified length 65 | * Used for PKCE code verifier and state parameter 66 | */ 67 | export function generateRandomString(length: number = 64): string { 68 | return randomBytes(length).toString('base64url').substring(0, length); 69 | } 70 | 71 | /** 72 | * Generates a code challenge using SHA-256 hash as required by PKCE 73 | */ 74 | export function generateCodeChallenge(codeVerifier: string): string { 75 | return createHash('sha256') 76 | .update(codeVerifier) 77 | .digest('base64url') 78 | .replace(/\+/g, '-') 79 | .replace(/\//g, '_') 80 | .replace(/=/g, ''); 81 | } 82 | 83 | /** 84 | * Store PKCE parameters for use in authorization code exchange 85 | */ 86 | export function storePKCEParams(params: PKCEParams): void { 87 | try { 88 | // Add timestamp for expiration 89 | const pkceData = { 90 | ...params, 91 | timestamp: Date.now() 92 | }; 93 | 94 | fs.writeFileSync(pkceStore, JSON.stringify(pkceData, null, 2)); 95 | } catch (error) { 96 | console.error('Error storing PKCE parameters:', error); 97 | } 98 | } 99 | 100 | /** 101 | * Get stored PKCE parameters 102 | */ 103 | export function getPKCEParams(): PKCEParams | null { 104 | try { 105 | if (fs.existsSync(pkceStore)) { 106 | const data = fs.readFileSync(pkceStore, 'utf8'); 107 | const pkceData = JSON.parse(data) as PKCEParams; 108 | 109 | // Check if PKCE data is expired (10 minutes) 110 | const now = Date.now(); 111 | if (pkceData.timestamp && now - pkceData.timestamp > 10 * 60 * 1000) { 112 | // PKCE data expired, delete it 113 | fs.unlinkSync(pkceStore); 114 | return null; 115 | } 116 | 117 | return pkceData; 118 | } 119 | } catch (error) { 120 | console.error('Error reading PKCE parameters:', error); 121 | } 122 | return null; 123 | } 124 | 125 | /** 126 | * Clear stored PKCE parameters 127 | */ 128 | export function clearPKCEParams(): void { 129 | try { 130 | if (fs.existsSync(pkceStore)) { 131 | fs.unlinkSync(pkceStore); 132 | } 133 | } catch (error) { 134 | console.error('Error clearing PKCE parameters:', error); 135 | } 136 | } 137 | 138 | /** 139 | * Store a token in the token store 140 | */ 141 | export function storeToken(token: TokenData): void { 142 | try { 143 | // Add current timestamp if not present 144 | if (!token.iat) { 145 | token.iat = Math.floor(Date.now() / 1000); 146 | } 147 | 148 | fs.writeFileSync(tokenStore, JSON.stringify(token, null, 2)); 149 | } catch (error) { 150 | console.error('Error storing token:', error); 151 | } 152 | } 153 | 154 | /** 155 | * Get a stored token from the token store 156 | */ 157 | export function getStoredToken(): TokenData | null { 158 | try { 159 | if (fs.existsSync(tokenStore)) { 160 | const data = fs.readFileSync(tokenStore, 'utf8'); 161 | return JSON.parse(data); 162 | } 163 | } catch (error) { 164 | console.error('Error reading token:', error); 165 | } 166 | return null; 167 | } 168 | 169 | /** 170 | * Check if a token is expired 171 | */ 172 | export function isTokenExpired(token: TokenData): boolean { 173 | // If there's no expiration or timestamp, consider it expired 174 | if (!token.expires_in || !token.iat) { 175 | return true; 176 | } 177 | 178 | // Check if the token has expired (with 60-second buffer) 179 | const expirationTime = (token.iat + token.expires_in - 60); 180 | const currentTime = Math.floor(Date.now() / 1000); 181 | 182 | return currentTime > expirationTime; 183 | } 184 | 185 | /** 186 | * Execute a command as a child process 187 | * 188 | * This is a promisified version of child_process.exec for use in async/await code. 189 | * Currently not used in the client code (which uses execSync directly), 190 | * but kept for potential future use in asynchronous command execution. 191 | * 192 | * @param command - The shell command to execute 193 | * @returns A promise that resolves with the command output or rejects with error 194 | */ 195 | export function executeCommand(command: string): Promise { 196 | return new Promise((resolve, reject) => { 197 | exec(command, (error, stdout, stderr) => { 198 | if (error) { 199 | reject({ error, stderr }); 200 | return; 201 | } 202 | resolve(stdout.trim()); 203 | }); 204 | }); 205 | } 206 | 207 | /** 208 | * Opens the default browser with the provided URL 209 | */ 210 | export function openBrowser(url: string): void { 211 | const command = process.platform === 'win32' ? 'start' : 212 | process.platform === 'darwin' ? 'open' : 'xdg-open'; 213 | 214 | console.log(`Opening browser to: ${url}`); 215 | exec(`${command} "${url}"`); 216 | } 217 | 218 | /** 219 | * Parse a URL query string into key-value pairs 220 | */ 221 | export function parseQueryString(queryString: string): Record { 222 | const params: Record = {}; 223 | const searchParams = new URLSearchParams(queryString); 224 | 225 | for (const [key, value] of searchParams.entries()) { 226 | params[key] = value; 227 | } 228 | 229 | return params; 230 | } 231 | 232 | /** 233 | * Escape special characters to prevent command injection 234 | */ 235 | export function escapeShellArg(arg: string): string { 236 | return arg.replace(/['"\\]/g, '\\$&'); 237 | } -------------------------------------------------------------------------------- /src/mcpTransport.ts: -------------------------------------------------------------------------------- 1 | /** 2 | * MCP Transport Layer 3 | * 4 | * This file provides transport layer utilities for the MCP client, including: 5 | * - HTTP/HTTPS request handling 6 | * - Streamable-HTTP format processing 7 | * - MCP protocol messaging 8 | * 9 | * It abstracts away the complexity of the MCP transport protocol 10 | * and provides a reliable communication layer with error handling. 11 | */ 12 | import * as http from 'http'; 13 | import * as https from 'https'; 14 | import { URL } from 'url'; 15 | 16 | /** 17 | * Response object returned by makeRequest 18 | */ 19 | export interface RequestResponse { 20 | statusCode: number; 21 | headers: Record; 22 | body: string; 23 | } 24 | 25 | /** 26 | * MCP Request type for more strongly typed requests 27 | */ 28 | export interface McpRequest { 29 | jsonrpc: string; 30 | method: string; 31 | params: Record; 32 | id: string; 33 | } 34 | 35 | /** 36 | * Make an HTTP/HTTPS request with comprehensive error handling and debug logging 37 | * 38 | * @param method - HTTP method (GET, POST, etc.) 39 | * @param url - The URL to request 40 | * @param headers - HTTP headers to include 41 | * @param body - Optional request body 42 | * @param debug - Whether to enable debug logging 43 | * @returns Promise resolving to response with status code, headers, and body 44 | */ 45 | export async function makeRequest( 46 | method: string, 47 | url: string, 48 | headers: Record, 49 | body?: any, 50 | debug: boolean = false 51 | ): Promise { 52 | return new Promise((resolve, reject) => { 53 | // Parse the URL to determine whether to use HTTP or HTTPS 54 | const parsedUrl = new URL(url); 55 | const httpModule = parsedUrl.protocol === 'https:' ? https : http; 56 | 57 | // Always ensure we have the proper content type and accept headers 58 | headers = { 59 | 'Content-Type': 'application/json', 60 | 'Accept': 'application/json, text/event-stream', 61 | ...headers // This allows overriding the defaults if provided 62 | }; 63 | 64 | // Add protocol version header for MCP 65 | headers['MCP-Protocol-Version'] = '2025-03-26'; 66 | 67 | // Prepare request options 68 | const options: http.RequestOptions = { 69 | method, 70 | headers, 71 | timeout: 30000, // 30 second timeout 72 | }; 73 | 74 | if (debug) { 75 | console.log('\n-------- REQUEST --------'); 76 | console.log(`${method} ${url}`); 77 | console.log('Headers:', JSON.stringify(headers, null, 2)); 78 | if (body) { 79 | console.log('Body:', typeof body === 'string' ? body : JSON.stringify(body, null, 2)); 80 | } 81 | console.log('--------------------------\n'); 82 | } 83 | 84 | // Create the request 85 | const req = httpModule.request(url, options, (res) => { 86 | let responseBody = ''; 87 | 88 | // Log the response status and headers if in debug mode 89 | if (debug) { 90 | console.log('\n-------- RESPONSE --------'); 91 | console.log(`Status: ${res.statusCode} ${res.statusMessage}`); 92 | console.log('Headers:', JSON.stringify(res.headers, null, 2)); 93 | } 94 | 95 | // Collect the response body 96 | res.on('data', (chunk) => { 97 | responseBody += chunk; 98 | if (debug) { 99 | console.log(`Received chunk: ${chunk.length} bytes`); 100 | } 101 | }); 102 | 103 | // Process the complete response 104 | res.on('end', () => { 105 | if (debug) { 106 | try { 107 | // Try to parse as JSON for prettier logging 108 | const parsed = JSON.parse(responseBody); 109 | console.log('Body (JSON):', JSON.stringify(parsed, null, 2)); 110 | } catch { 111 | // If not JSON, log as text with length 112 | console.log(`Body (Text - ${responseBody.length} bytes):`); 113 | if (responseBody.length < 1000) { 114 | console.log(responseBody); 115 | } else { 116 | console.log(responseBody.substring(0, 1000) + '... [truncated]'); 117 | } 118 | } 119 | console.log('---------------------------\n'); 120 | } 121 | 122 | // Check for error status codes 123 | if (res.statusCode && (res.statusCode < 200 || res.statusCode >= 300)) { 124 | const error: any = new Error(`HTTP ${res.statusCode}: ${res.statusMessage}`); 125 | error.statusCode = res.statusCode; 126 | error.headers = res.headers; 127 | error.body = responseBody; 128 | 129 | // Special handling for common error codes 130 | switch (res.statusCode) { 131 | case 401: 132 | // For 401 errors, extract resource metadata from WWW-Authenticate header 133 | if (res.headers['www-authenticate']) { 134 | const wwwAuth = res.headers['www-authenticate'] as string; 135 | error.authHeader = wwwAuth; 136 | 137 | // Extract the resource metadata URI 138 | const metadataMatch = wwwAuth.match(/resource_metadata_uri="([^"]+)"/); 139 | if (metadataMatch && metadataMatch[1]) { 140 | error.resourceMetadataUri = metadataMatch[1]; 141 | console.log('\nAuthentication required:'); 142 | console.log(`WWW-Authenticate: ${wwwAuth}`); 143 | console.log(`Resource metadata available at: ${metadataMatch[1]}`); 144 | console.log('Please run "auth login" to authenticate'); 145 | } 146 | } 147 | break; 148 | 149 | case 403: 150 | console.log('\nAuthorization failed: Insufficient permissions'); 151 | try { 152 | const errorBody = JSON.parse(responseBody); 153 | if (errorBody?.error?.message) { 154 | console.log(`Error details: ${errorBody.error.message}`); 155 | } 156 | } catch {} 157 | break; 158 | 159 | case 400: 160 | console.log('\nBad request: The server could not understand the request'); 161 | try { 162 | const errorBody = JSON.parse(responseBody); 163 | if (errorBody?.error?.message) { 164 | console.log(`Error details: ${errorBody.error.message}`); 165 | } 166 | } catch {} 167 | break; 168 | 169 | case 429: 170 | console.log('\nToo many requests: Rate limit exceeded'); 171 | // Check for Retry-After header 172 | if (res.headers['retry-after']) { 173 | console.log(`Please retry after ${res.headers['retry-after']} seconds`); 174 | } 175 | break; 176 | 177 | case 500: 178 | case 502: 179 | case 503: 180 | case 504: 181 | console.log('\nServer error: The server encountered an error processing the request'); 182 | break; 183 | } 184 | 185 | reject(error); 186 | return; 187 | } 188 | 189 | // Extract all headers 190 | const responseHeaders: Record = {}; 191 | Object.keys(res.headers).forEach(key => { 192 | const value = res.headers[key]; 193 | if (value !== undefined) { 194 | responseHeaders[key] = Array.isArray(value) ? value[0] : value; 195 | } 196 | }); 197 | 198 | // Resolve with the complete response 199 | resolve({ 200 | statusCode: res.statusCode || 200, 201 | headers: responseHeaders, 202 | body: responseBody 203 | }); 204 | }); 205 | }); 206 | 207 | // Handle request errors 208 | req.on('error', (error: NodeJS.ErrnoException) => { 209 | console.error(`Network error: ${error.message}`); 210 | 211 | // Special handling for common network errors 212 | if (error.code === 'ECONNREFUSED') { 213 | console.log('Could not connect to the server. Please check:'); 214 | console.log('1. Is the server running?'); 215 | console.log('2. Is the URL correct?'); 216 | console.log('3. Is there a firewall blocking the connection?'); 217 | } else if (error.code === 'ENOTFOUND') { 218 | console.log('Could not resolve the hostname. Please check:'); 219 | console.log('1. Is the URL correct?'); 220 | console.log('2. Is your internet connection working?'); 221 | } else if (error.code === 'ETIMEDOUT') { 222 | console.log('The connection timed out. Please check:'); 223 | console.log('1. Is the server running?'); 224 | console.log('2. Is the network slow or overloaded?'); 225 | } 226 | 227 | reject(error); 228 | }); 229 | 230 | // Send the request body if provided 231 | if (body) { 232 | const bodyData = typeof body === 'string' ? body : JSON.stringify(body); 233 | req.write(bodyData); 234 | } 235 | 236 | // End the request 237 | req.end(); 238 | }); 239 | } 240 | 241 | /** 242 | * Determine if content is in MCP Streamable-HTTP format 243 | * 244 | * @param content - Content to check for streaming format 245 | * @returns boolean indicating if content appears to be in streaming format 246 | */ 247 | export function isStreamingFormat(content: string): boolean { 248 | // Check for MCP Streamable-HTTP format indicators 249 | return content.trim().startsWith('event:') || 250 | content.includes('\nevent:') || 251 | content.includes('\ndata:'); 252 | } 253 | 254 | /** 255 | * Parse Streamable-HTTP formatted messages 256 | * Extracts data from MCP's streaming format, handling both single and multi-line events 257 | * 258 | * @param streamContent - Raw streaming content from server 259 | * @returns Array of parsed JSON objects extracted from data fields 260 | */ 261 | export function parseStreamingMessages(streamContent: string): any[] { 262 | const results: any[] = []; 263 | 264 | // Split the content into individual event blocks 265 | const eventBlocks = streamContent.split(/\n\n+/g).filter(block => block.trim()); 266 | 267 | for (const block of eventBlocks) { 268 | const lines = block.split('\n'); 269 | let dataContent = ''; 270 | let eventType = ''; 271 | let eventId = ''; 272 | 273 | for (const line of lines) { 274 | if (line.startsWith('data:')) { 275 | dataContent = line.substring(5).trim(); 276 | } else if (line.startsWith('event:')) { 277 | eventType = line.substring(6).trim(); 278 | } else if (line.startsWith('id:')) { 279 | eventId = line.substring(3).trim(); 280 | } 281 | } 282 | 283 | if (dataContent) { 284 | try { 285 | const parsedData = JSON.parse(dataContent); 286 | results.push(parsedData); 287 | } catch (e) { 288 | console.warn(`Could not parse SSE data as JSON: ${dataContent}`); 289 | // Still add the raw data if we can't parse it 290 | results.push({ rawData: dataContent }); 291 | } 292 | } 293 | } 294 | 295 | return results; 296 | } 297 | 298 | /** 299 | * Creates an MCP initialize request object 300 | */ 301 | export function createInitializeRequest(clientName: string, clientVersion: string): McpRequest { 302 | return { 303 | jsonrpc: '2.0', 304 | method: 'initialize', 305 | params: { 306 | clientInfo: { 307 | name: clientName, 308 | version: clientVersion 309 | }, 310 | protocolVersion: '2025-03-26', 311 | capabilities: {} 312 | }, 313 | id: 'init-1' 314 | }; 315 | } 316 | 317 | /** 318 | * Creates an MCP tools/list request object 319 | */ 320 | export function createListToolsRequest(): McpRequest { 321 | return { 322 | jsonrpc: '2.0', 323 | method: 'tools/list', 324 | params: {}, 325 | id: 'tools-1' 326 | }; 327 | } 328 | 329 | /** 330 | * Creates an MCP tools/call request object 331 | */ 332 | export function createToolCallRequest(toolName: string, args: Record = {}): McpRequest { 333 | return { 334 | jsonrpc: '2.0', 335 | method: 'tools/call', 336 | params: { 337 | name: toolName, 338 | arguments: args 339 | }, 340 | id: `tool-${Date.now()}` 341 | }; 342 | } 343 | 344 | /** 345 | * Creates an MCP resources/list request object 346 | */ 347 | export function createListResourcesRequest(): McpRequest { 348 | return { 349 | jsonrpc: '2.0', 350 | method: 'resources/list', 351 | params: {}, 352 | id: 'resources-1' 353 | }; 354 | } -------------------------------------------------------------------------------- /src/server.ts: -------------------------------------------------------------------------------- 1 | import express, { Request, Response, NextFunction } from 'express'; 2 | import { randomUUID } from 'node:crypto'; 3 | import { z } from 'zod'; 4 | 5 | import { CloudWatchLogsClient, FilterLogEventsCommand } from "@aws-sdk/client-cloudwatch-logs"; 6 | import serverlessExpress from '@codegenie/serverless-express'; 7 | 8 | import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; 9 | import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js"; 10 | import { InMemoryEventStore } from "@modelcontextprotocol/sdk/examples/shared/inMemoryEventStore.js"; 11 | 12 | import { CallToolResult, GetPromptResult, isInitializeRequest, ReadResourceResult } from "@modelcontextprotocol/sdk/types.js"; 13 | 14 | // Import auth modules 15 | import { createAuthorizer } from './auth'; 16 | import { getDefaultResourceMetadataConfig, createResourceMetadataHandler } from './auth/oauthMetadata'; 17 | import { logger, authLogger, apiLogger, logHttpRequest, logHttpResponse } from './logging'; 18 | 19 | // Create an MCP server with implementation details 20 | const server = new McpServer({ 21 | name: 'bedrock-usage-stats-http-server', 22 | version: '1.0.1', 23 | }, { capabilities: { logging: {} } }); 24 | 25 | // Get the authorization method from environment variable 26 | const authMethod = process.env.AUTH_METHOD || 'oauth'; 27 | console.log(`Using authorization method: ${authMethod}`); 28 | 29 | // Create the appropriate authorizer 30 | const authorizer = createAuthorizer(authMethod as 'oauth' | 'lambda'); 31 | server.tool( 32 | 'greet', 33 | 'A simple greeting tool', 34 | { 35 | name: z.string().describe('Name to greet'), 36 | }, 37 | async ({ name }): Promise => { 38 | return { 39 | content: [ 40 | { 41 | type: 'text', 42 | text: `Hello, ${name}!`, 43 | }, 44 | ], 45 | }; 46 | } 47 | ); 48 | 49 | // Register a tool that sends multiple greetings with notifications 50 | server.tool( 51 | 'multi-greet', 52 | 'A tool that sends different greetings with delays between them', 53 | { 54 | name: z.string().describe('Name to greet'), 55 | }, 56 | async ({ name }, { sendNotification }): Promise => { 57 | const sleep = (ms: number) => new Promise(resolve => setTimeout(resolve, ms)); 58 | 59 | await sendNotification({ 60 | method: "notifications/message", 61 | params: { level: "debug", data: `Starting multi-greet for ${name}` } 62 | }); 63 | 64 | await sleep(1000); // Wait 1 second before first greeting 65 | 66 | await sendNotification({ 67 | method: "notifications/message", 68 | params: { level: "info", data: `Sending first greeting to ${name}` } 69 | }); 70 | 71 | await sleep(1000); // Wait another second before second greeting 72 | 73 | await sendNotification({ 74 | method: "notifications/message", 75 | params: { level: "info", data: `Sending second greeting to ${name}` } 76 | }); 77 | 78 | return { 79 | content: [ 80 | { 81 | type: 'text', 82 | text: `Good morning, ${name}!`, 83 | } 84 | ], 85 | }; 86 | } 87 | ); 88 | 89 | // Register a simple prompt 90 | server.prompt( 91 | 'greeting-template', 92 | 'A simple greeting prompt template', 93 | { 94 | name: z.string().describe('Name to include in greeting'), 95 | }, 96 | async ({ name }): Promise => { 97 | return { 98 | messages: [ 99 | { 100 | role: 'user', 101 | content: { 102 | type: 'text', 103 | text: `Please greet ${name} in a friendly manner.`, 104 | }, 105 | }, 106 | ], 107 | }; 108 | } 109 | ); 110 | 111 | // --- Bedrock Logs Tool --- 112 | 113 | // Assume getAwsServiceClient is defined elsewhere or add its definition if needed 114 | // Example placeholder: 115 | async function getAwsServiceClient(clientConstructor: any, region: string, accountId?: string) { 116 | console.log(`Placeholder: Creating AWS client ${clientConstructor.name} for region ${region} (Account: ${accountId || 'default'})`); 117 | // Replace with actual AWS SDK client creation logic (e.g., using credentials) 118 | return new clientConstructor({ region }); 119 | } 120 | 121 | 122 | const BedrockLogsParamsSchema = z.object({ 123 | region: z.string().describe("AWS region for CloudWatch Logs"), 124 | log_group_name: z.string().describe("Name of the CloudWatch log group containing Bedrock logs"), 125 | days: z.number().int().positive().describe("Number of past days to fetch logs for"), 126 | aws_account_id: z.string().optional().describe("AWS Account ID if using cross-account access"), 127 | }); 128 | 129 | // Helper: Get Bedrock Logs (Simplified version of Python's get_bedrock_logs) 130 | // Note: Does not replicate pandas DataFrame creation/manipulation. Returns array of log objects. 131 | const getBedrockLogs = async (params: z.infer) => { 132 | console.log("getBedrockLogs, params=", params); 133 | const client = await getAwsServiceClient(CloudWatchLogsClient, params.region, params.aws_account_id); 134 | 135 | const endTime = new Date(); 136 | const startTime = new Date(); 137 | startTime.setDate(endTime.getDate() - params.days); 138 | 139 | const startTimeMs = startTime.getTime(); 140 | const endTimeMs = endTime.getTime(); 141 | 142 | let filteredLogs: any[] = []; // Using any for simplicity here, define a specific type if preferred 143 | let nextToken: string | undefined; 144 | 145 | try { 146 | do { 147 | const command = new FilterLogEventsCommand({ 148 | logGroupName: params.log_group_name, 149 | startTime: startTimeMs, 150 | endTime: endTimeMs, 151 | nextToken: nextToken, 152 | }); 153 | 154 | const response = await client.send(command); 155 | const events = response.events || []; 156 | 157 | for (const event of events) { 158 | try { 159 | const message = JSON.parse(event.message); 160 | const inputTokens = message.input?.inputTokenCount ?? 0; 161 | const outputTokens = message.output?.outputTokenCount ?? 0; 162 | filteredLogs.push({ 163 | timestamp: message.timestamp, 164 | region: message.region, 165 | modelId: message.modelId, 166 | userId: message.identity?.arn, 167 | inputTokens: inputTokens, 168 | completionTokens: outputTokens, 169 | totalTokens: inputTokens + outputTokens, 170 | }); 171 | } catch (jsonError) { 172 | console.warn("Skipping non-JSON log message:", event.message); 173 | } 174 | } 175 | nextToken = response.nextToken; 176 | } while (nextToken); 177 | 178 | console.log(`Found ${filteredLogs.length} Bedrock log events.`); 179 | return filteredLogs; 180 | 181 | } catch (error: any) { 182 | if (error.name === 'ResourceNotFoundException') { 183 | console.error(`Log group '${params.log_group_name}' not found in region ${params.region}.`); 184 | return []; 185 | } 186 | console.error("Error retrieving Bedrock logs:", error); 187 | throw new Error(`Failed to retrieve Bedrock logs: ${error.message}`); 188 | } 189 | }; 190 | 191 | // Tool: Bedrock Daily Usage Report (Non-Streaming Version) 192 | server.tool( 193 | "get_bedrock_usage_report", 194 | "Get Bedrock daily usage report (non-streaming)", 195 | // Define schema directly as an object literal 196 | { 197 | region: z.string().describe("AWS region for CloudWatch Logs"), 198 | log_group_name: z.string().describe("Name of the CloudWatch log group containing Bedrock logs"), 199 | days: z.number().int().positive().describe("Number of past days to fetch logs for"), 200 | aws_account_id: z.string().optional().describe("AWS Account ID if using cross-account access"), 201 | }, 202 | // Standard async function returning the full report 203 | async (args, extra): Promise => { 204 | console.log("Executing get_bedrock_report (non-streaming) with args:", args); 205 | try { 206 | const logs = await getBedrockLogs(args); 207 | 208 | if (!logs || logs.length === 0) { 209 | return { content: [{ type: "text", text: "No Bedrock usage data found for the specified period." }] }; 210 | } 211 | 212 | // --- Data Aggregation (Same as before) --- 213 | const dailyStats: { [date: string]: { regions: any, models: any, users: any, requests: number, inputTokens: number, completionTokens: number, totalTokens: number } } = {}; 214 | let totalRequests = 0; 215 | let totalInputTokens = 0; 216 | let totalCompletionTokens = 0; 217 | let totalTokens = 0; 218 | const regionSummary: { [region: string]: { requests: number, input: number, completion: number, total: number } } = {}; 219 | const modelSummary: { [modelId: string]: { requests: number, input: number, completion: number, total: number } } = {}; 220 | const userSummary: { [userId: string]: { requests: number, input: number, completion: number, total: number } } = {}; 221 | 222 | logs.forEach(log => { 223 | totalRequests++; 224 | totalInputTokens += log.inputTokens || 0; 225 | totalCompletionTokens += log.completionTokens || 0; 226 | totalTokens += log.totalTokens || 0; 227 | const date = new Date(log.timestamp).toISOString().split('T')[0]; 228 | if (!dailyStats[date]) { 229 | dailyStats[date] = { regions: {}, models: {}, users: {}, requests: 0, inputTokens: 0, completionTokens: 0, totalTokens: 0 }; 230 | } 231 | dailyStats[date].requests++; 232 | dailyStats[date].inputTokens += log.inputTokens || 0; 233 | dailyStats[date].completionTokens += log.completionTokens || 0; 234 | dailyStats[date].totalTokens += log.totalTokens || 0; 235 | if (!regionSummary[log.region]) regionSummary[log.region] = { requests: 0, input: 0, completion: 0, total: 0 }; 236 | regionSummary[log.region].requests++; 237 | regionSummary[log.region].input += log.inputTokens || 0; 238 | regionSummary[log.region].completion += log.completionTokens || 0; 239 | regionSummary[log.region].total += log.totalTokens || 0; 240 | const simpleModelId = log.modelId?.includes('.') ? log.modelId.split('.').pop()! : log.modelId?.split('/').pop() || 'unknown'; 241 | if (!modelSummary[simpleModelId]) modelSummary[simpleModelId] = { requests: 0, input: 0, completion: 0, total: 0 }; 242 | modelSummary[simpleModelId].requests++; 243 | modelSummary[simpleModelId].input += log.inputTokens || 0; 244 | modelSummary[simpleModelId].completion += log.completionTokens || 0; 245 | modelSummary[simpleModelId].total += log.totalTokens || 0; 246 | if (log.userId) { 247 | if (!userSummary[log.userId]) userSummary[log.userId] = { requests: 0, input: 0, completion: 0, total: 0 }; 248 | userSummary[log.userId].requests++; 249 | userSummary[log.userId].input += log.inputTokens || 0; 250 | userSummary[log.userId].completion += log.completionTokens || 0; 251 | userSummary[log.userId].total += log.totalTokens || 0; 252 | } 253 | }); 254 | 255 | // --- Formatting Output (Single String) --- 256 | let output = `Bedrock Daily Usage Report (Past ${args.days} days - ${args.region})\n`; 257 | output += `Total Requests: ${totalRequests}\n`; 258 | output += `Total Input Tokens: ${totalInputTokens}\n`; 259 | output += `Total Completion Tokens: ${totalCompletionTokens}\n`; 260 | output += `Total Tokens: ${totalTokens}\n`; 261 | 262 | output += "\n--- Daily Totals ---\n"; 263 | Object.entries(dailyStats).sort(([dateA], [dateB]) => dateA.localeCompare(dateB)).forEach(([date, stats]) => { 264 | output += `${date}: Requests=${stats.requests}, Input=${stats.inputTokens}, Completion=${stats.completionTokens}, Total=${stats.totalTokens}\n`; 265 | }); 266 | 267 | output += "\n--- Region Summary ---\n"; 268 | Object.entries(regionSummary).forEach(([region, stats]) => { 269 | output += `${region}: Requests=${stats.requests}, Input=${stats.input}, Completion=${stats.completion}, Total=${stats.total}\n`; 270 | }); 271 | 272 | output += "\n--- Model Summary ---\n"; 273 | Object.entries(modelSummary).forEach(([model, stats]) => { 274 | output += `${model}: Requests=${stats.requests}, Input=${stats.input}, Completion=${stats.completion}, Total=${stats.total}\n`; 275 | }); 276 | 277 | if (Object.keys(userSummary).length > 0) { 278 | output += "\n--- User Summary ---\n"; 279 | Object.entries(userSummary).forEach(([user, stats]) => { 280 | output += `${user}: Requests=${stats.requests}, Input=${stats.input}, Completion=${stats.completion}, Total=${stats.total}\n`; 281 | }); 282 | } 283 | 284 | // Return the full report string 285 | return { content: [{ type: "text", text: output.trimEnd() }] }; 286 | 287 | } catch (error: any) { 288 | console.error("Error in get_bedrock_report tool:", error); 289 | // Return an error result 290 | return { content: [{ type: "text", text: `Error getting Bedrock report: ${error.message}` }], isError: true }; 291 | } 292 | } 293 | ); 294 | 295 | 296 | // Create a simple resource at a fixed URI 297 | server.resource( 298 | 'greeting-resource', 299 | 'https://example.com/greetings/default', 300 | { mimeType: 'text/plain' }, 301 | async (): Promise => { 302 | return { 303 | contents: [ 304 | { 305 | uri: 'https://example.com/greetings/default', 306 | text: 'Hello, world!', 307 | }, 308 | ], 309 | }; 310 | } 311 | ); 312 | 313 | const app = express(); 314 | app.use(express.json()); 315 | 316 | // Create Cognito authorizer based on environment variables 317 | // const authorizer = createCognitoAuthorizer(); 318 | 319 | // Setup OAuth 2.0 Protected Resource Metadata endpoint 320 | const resourceMetadataConfig = getDefaultResourceMetadataConfig(); 321 | 322 | // Register the OAuth resource metadata endpoint at the root and at /prod to handle both locally and in Lambda 323 | app.get('/.well-known/oauth-protected-resource', createResourceMetadataHandler(resourceMetadataConfig)); 324 | app.get('/prod/.well-known/oauth-protected-resource', createResourceMetadataHandler(resourceMetadataConfig)); 325 | 326 | // Add authorization server metadata discovery endpoint 327 | // This is a convenience endpoint to redirect to the actual Cognito OIDC configuration 328 | app.get('/.well-known/oauth-authorization-server', (req: Request, res: Response) => { 329 | try { 330 | authLogger.debug({ 331 | message: 'Authorization server metadata request received', 332 | path: req.path 333 | }); 334 | 335 | const authServer = resourceMetadataConfig.authorizationServers[0]; 336 | if (!authServer || !authServer.metadataUrl) { 337 | authLogger.error('No authorization server metadata URL available'); 338 | return res.status(404).json({ 339 | error: 'not_found', 340 | error_description: 'Authorization server metadata not available' 341 | }); 342 | } 343 | 344 | // Set CORS headers 345 | res.header('Access-Control-Allow-Origin', '*'); 346 | res.header('Access-Control-Allow-Methods', 'GET, OPTIONS'); 347 | res.header('Access-Control-Allow-Headers', 'Content-Type, Authorization'); 348 | 349 | // Log the redirect 350 | authLogger.info(`Redirecting to authorization server metadata: ${authServer.metadataUrl}`); 351 | 352 | // Redirect to the authorization server's metadata URL (typically Cognito's OIDC configuration) 353 | res.redirect(authServer.metadataUrl); 354 | } catch (error) { 355 | authLogger.error({ 356 | message: 'Error processing authorization server metadata request', 357 | error 358 | }); 359 | res.status(500).json({ 360 | error: 'server_error', 361 | error_description: 'An error occurred while processing the authorization server metadata request' 362 | }); 363 | } 364 | }); 365 | 366 | // Also handle the discovery endpoint at the /prod path 367 | app.get('/prod/.well-known/oauth-authorization-server', (req: Request, res: Response) => { 368 | try { 369 | const authServer = resourceMetadataConfig.authorizationServers[0]; 370 | if (!authServer || !authServer.metadataUrl) { 371 | return res.status(404).json({ 372 | error: 'not_found', 373 | error_description: 'Authorization server metadata not available' 374 | }); 375 | } 376 | 377 | // Set CORS headers 378 | res.header('Access-Control-Allow-Origin', '*'); 379 | res.header('Access-Control-Allow-Methods', 'GET, OPTIONS'); 380 | res.header('Access-Control-Allow-Headers', 'Content-Type, Authorization'); 381 | 382 | // Redirect to the authorization server's metadata URL 383 | res.redirect(authServer.metadataUrl); 384 | } catch (error) { 385 | res.status(500).json({ 386 | error: 'server_error', 387 | error_description: 'An error occurred while processing the authorization server metadata request' 388 | }); 389 | } 390 | }); 391 | 392 | // Track MCP Protocol Version 393 | app.use((req, res, next) => { 394 | const protocolVersion = req.headers['mcp-protocol-version']; 395 | if (protocolVersion) { 396 | logger.debug(`Client requested MCP Protocol Version: ${protocolVersion}`); 397 | // You can use this to handle version-specific behavior if needed 398 | } 399 | next(); 400 | }); 401 | 402 | // Create middleware for logging requests 403 | app.use((req: Request, res: Response, next: NextFunction) => { 404 | // Use enhanced structured logging 405 | logHttpRequest(req); 406 | 407 | // Capture response timing and status code 408 | const startTime = Date.now(); 409 | 410 | // Intercept res.end to log response 411 | const originalEnd = res.end; 412 | res.end = function(...args) { 413 | const duration = Date.now() - startTime; 414 | logHttpResponse(res, duration); 415 | return originalEnd.apply(res, args); 416 | }; 417 | 418 | next(); 419 | }); 420 | 421 | // Map to store transports by session ID 422 | // TODO: Use a more efficient transport, like a Redis-based transport or maybe DynamoDB-based transport 423 | const transports: { [sessionId: string]: StreamableHTTPServerTransport } = {}; 424 | 425 | // Define routes that skip authentication 426 | const publicPaths = [ 427 | '/.well-known/oauth-protected-resource', 428 | '/prod/.well-known/oauth-protected-resource', 429 | '/.well-known/oauth-authorization-server', 430 | '/prod/.well-known/oauth-authorization-server', 431 | '/.well-known/openid-configuration', 432 | '/prod/.well-known/openid-configuration', 433 | ]; 434 | 435 | // Authentication middleware - only applied to protected routes 436 | const authMiddleware = (req: Request, res: Response, next: NextFunction) => { 437 | // Skip authentication for public paths 438 | if (publicPaths.some(path => req.path.startsWith(path))) { 439 | return next(); 440 | } 441 | 442 | // Skip auth check for initialize requests (we'll validate them later) 443 | if (req.method === 'POST' && !req.headers['mcp-session-id'] && 444 | req.body && req.body.method === 'initialize') { 445 | console.log('Initialization request detected - authentication check skipped at this stage'); 446 | console.log('isInitializeRequest result:', isInitializeRequest(req.body)); 447 | console.log('Request method check:', req.body.method === 'initialize'); 448 | console.log('Request body:', JSON.stringify(req.body)); 449 | return next(); 450 | } 451 | 452 | // For all other requests, check authentication 453 | console.log(`Checking authentication for request to ${req.path} with method ${req.method}`); 454 | console.log(`Headers: ${JSON.stringify(req.headers)}`); 455 | return authorizer.createAuthMiddleware()(req, res, next); 456 | }; 457 | 458 | // Apply auth middleware 459 | app.use(authMiddleware); 460 | 461 | // Use '/mcp' for local development, '/prod/mcp' for Lambda 462 | const isLambda = !!process.env.LAMBDA_TASK_ROOT; 463 | const mcpEndpoint = isLambda ? '/prod/mcp' : '/mcp'; 464 | 465 | app.post(mcpEndpoint, async (req: Request, res: Response) => { 466 | authLogger.debug({ message: 'Received MCP request', body: req.body }); 467 | console.log(`MCP request received at ${mcpEndpoint}, method: ${req.method}, sessionId: ${req.headers['mcp-session-id']}`); 468 | console.log(`Request body: ${JSON.stringify(req.body)}`); 469 | console.log(`Is initialize request: ${req.body && isInitializeRequest(req.body)}`); 470 | 471 | try { 472 | // Check for existing session ID 473 | const sessionId = req.headers['mcp-session-id'] as string | undefined; 474 | let transport: StreamableHTTPServerTransport; 475 | 476 | if (sessionId && transports[sessionId]) { 477 | // Reuse existing transport 478 | console.log(`Reusing existing transport for session ID: ${sessionId}`); 479 | transport = transports[sessionId]; 480 | } else if (!sessionId && req.body && req.body.method === 'initialize') { 481 | // New initialization request 482 | console.log('Processing new initialization request'); 483 | const eventStore = new InMemoryEventStore(); 484 | transport = new StreamableHTTPServerTransport({ 485 | sessionIdGenerator: () => randomUUID(), 486 | eventStore, // Enable resumability 487 | onsessioninitialized: (sessionId) => { 488 | // Store the transport by session ID when session is initialized 489 | // This avoids race conditions where requests might come in before the session is stored 490 | console.log(`Session initialized with ID: ${sessionId}`); 491 | apiLogger.info({ message: 'Session initialized', sessionId }); 492 | transports[sessionId] = transport; 493 | } 494 | }); 495 | 496 | // Set up onclose handler to clean up transport when closed 497 | transport.onclose = () => { 498 | const sid = transport.sessionId; 499 | if (sid && transports[sid]) { 500 | console.log(`Transport closed for session ID: ${sid}`); 501 | apiLogger.info({ message: 'Transport closed', sessionId: sid }); 502 | delete transports[sid]; 503 | } 504 | }; 505 | 506 | // Connect the transport to the MCP server BEFORE handling the request 507 | // so responses can flow back through the same transport 508 | console.log('Connecting transport to MCP server...'); 509 | await server.connect(transport); 510 | 511 | console.log('Handling initialization request...'); 512 | await transport.handleRequest(req, res, req.body); 513 | console.log('Initialization request handled'); 514 | return; // Already handled 515 | } else { 516 | // Invalid request - no session ID or not initialization request 517 | console.log('Invalid request - no session ID or not initialization request'); 518 | res.status(400).json({ 519 | jsonrpc: '2.0', 520 | error: { 521 | code: -32000, 522 | message: 'Bad Request: No valid session ID provided', 523 | }, 524 | id: null, 525 | }); 526 | return; 527 | } 528 | 529 | // Handle the request with existing transport - no need to reconnect 530 | // The existing transport is already connected to the server 531 | await transport.handleRequest(req, res, req.body); 532 | } catch (error) { 533 | logger.error({ message: 'Error handling MCP request', error }); 534 | if (!res.headersSent) { 535 | res.status(500).json({ 536 | jsonrpc: '2.0', 537 | error: { 538 | code: -32603, 539 | message: 'Internal server error', 540 | }, 541 | id: null, 542 | }); 543 | } 544 | } 545 | }); 546 | 547 | // Handle GET requests for SSE streams (using built-in support from StreamableHTTP) 548 | app.get(mcpEndpoint, async (req: Request, res: Response) => { 549 | const sessionId = req.headers['mcp-session-id'] as string | undefined; 550 | if (!sessionId || !transports[sessionId]) { 551 | apiLogger.warn({ message: 'Invalid session ID', sessionId }); 552 | res.status(400).send('Invalid or missing session ID'); 553 | return; 554 | } 555 | 556 | // Check for Last-Event-ID header for resumability 557 | const lastEventId = req.headers['last-event-id'] as string | undefined; 558 | if (lastEventId) { 559 | apiLogger.info({ message: 'Client reconnecting', sessionId, lastEventId }); 560 | } else { 561 | apiLogger.info({ message: 'Establishing new SSE stream', sessionId }); 562 | } 563 | 564 | const transport = transports[sessionId]; 565 | await transport.handleRequest(req, res); 566 | }); 567 | 568 | // Handle DELETE requests for session termination (according to MCP spec) 569 | app.delete(mcpEndpoint, async (req: Request, res: Response) => { 570 | const sessionId = req.headers['mcp-session-id'] as string | undefined; 571 | if (!sessionId || !transports[sessionId]) { 572 | apiLogger.warn({ message: 'Invalid session ID for termination', sessionId }); 573 | res.status(400).send('Invalid or missing session ID'); 574 | return; 575 | } 576 | 577 | apiLogger.info({ message: 'Received session termination request', sessionId }); 578 | 579 | try { 580 | const transport = transports[sessionId]; 581 | await transport.handleRequest(req, res); 582 | } catch (error) { 583 | apiLogger.error({ message: 'Error handling session termination', sessionId, error }); 584 | if (!res.headersSent) { 585 | res.status(500).send('Error processing session termination'); 586 | } 587 | } 588 | }); 589 | 590 | let handler: any; // Declare handler outside the if block 591 | 592 | if (isLambda) { 593 | // Start the server 594 | logger.info('Starting MCP server in Lambda environment'); 595 | handler = serverlessExpress({ app }); 596 | 597 | } 598 | else { 599 | const PORT = 3000; 600 | app.listen(PORT, () => { 601 | logger.info(`MCP Streamable HTTP Server listening on port ${PORT}`); 602 | logger.info(`OAuth Protected Resource metadata available at: http://localhost:${PORT}/.well-known/oauth-protected-resource`); 603 | }); 604 | 605 | // Handle server shutdown 606 | process.on('SIGINT', async () => { 607 | logger.info('Shutting down server...'); 608 | 609 | // Close all active transports to properly clean up resources 610 | for (const sessionId in transports) { 611 | try { 612 | logger.info({ message: 'Closing transport', sessionId }); 613 | await transports[sessionId].close(); 614 | delete transports[sessionId]; 615 | } catch (error) { 616 | logger.error({ message: 'Error closing transport', sessionId, error }); 617 | } 618 | } 619 | await server.close(); 620 | logger.info('Server shutdown complete'); 621 | process.exit(0); 622 | }); 623 | } 624 | 625 | // Export handler at the top level if it was assigned 626 | if (handler) { 627 | exports.handler = handler; 628 | } -------------------------------------------------------------------------------- /tsconfig.json: -------------------------------------------------------------------------------- 1 | { 2 | "compilerOptions": { 3 | "target": "es2020", 4 | "strict": false, 5 | "preserveConstEnums": true, 6 | "sourceMap": false, 7 | "module": "commonjs", 8 | "moduleResolution": "node", 9 | "esModuleInterop": true, 10 | "skipLibCheck": true, 11 | "forceConsistentCasingInFileNames": true, 12 | "isolatedModules": true, 13 | "outDir": "./dist", 14 | "rootDir": "./src", 15 | "allowJs": true, 16 | "resolveJsonModule": true, 17 | "baseUrl": ".", 18 | "paths": { 19 | "@modelcontextprotocol/sdk/*": ["node_modules/@modelcontextprotocol/sdk/dist/*"] 20 | } 21 | }, 22 | "include": ["src/**/*.ts"], 23 | "exclude": ["node_modules", "**/*.test.ts", "dist", "src/client.ts"] 24 | } --------------------------------------------------------------------------------