@friggframework/devtools 2.0.0--canary.397.155fecd.0 → 2.0.0--canary.398.c9e5d61.0

package/infrastructure/serverless-template.js

@@ -1,7 +1,14 @@
 const path = require('path');
 const fs = require('fs');
 
-// Function to find the actual path to node_modules
+/**
+ * Find the actual path to node_modules directory
+ * Tries multiple methods to locate node_modules:
+ * 1. Traversing up from current directory
+ * 2. Using npm root command
+ * 3. Looking for package.json and adjacent node_modules
+ * @returns {string} Path to node_modules directory
+ */
 const findNodeModulesPath = () => {
     try {
         // Method 1: Try to find node_modules by traversing up from current directory
@@ -75,7 +82,12 @@ const findNodeModulesPath = () => {
     }
 };
 
-// Function to modify handler paths to point to the correct node_modules
+/**
+ * Modify handler paths to point to the correct node_modules location
+ * Only modifies paths when running in offline mode
+ * @param {Object} functions - Serverless functions configuration object
+ * @returns {Object} Modified functions object with updated handler paths
+ */
 const modifyHandlerPaths = (functions) => {
     // Check if we're running in offline mode
     const isOffline = process.argv.includes('offline');
@@ -94,7 +106,8 @@ const modifyHandlerPaths = (functions) => {
         const functionDef = modifiedFunctions[functionName];
         if (functionDef?.handler?.includes('node_modules/')) {
             // Replace node_modules/ with the actual path to node_modules/
-            functionDef.handler = functionDef.handler.replace('node_modules/', '../node_modules/');
+            const relativePath = path.relative(process.cwd(), nodeModulesPath);
+            functionDef.handler = functionDef.handler.replace('node_modules/', `${relativePath}/`);
             console.log(`Updated handler for ${functionName}: ${functionDef.handler}`);
         }
     }
@@ -102,6 +115,329 @@ const modifyHandlerPaths = (functions) => {
     return modifiedFunctions;
 };
 
+/**
+ * Create VPC infrastructure resources for CloudFormation
+ * Creates VPC, subnets, NAT gateway, route tables, and security groups
+ * @param {Object} AppDefinition - Application definition object
+ * @param {Object} AppDefinition.vpc - VPC configuration
+ * @param {string} [AppDefinition.vpc.cidrBlock='10.0.0.0/16'] - CIDR block for VPC
+ * @returns {Object} CloudFormation resources for VPC infrastructure
+ */
+const createVPCInfrastructure = (AppDefinition) => {
+    const vpcResources = {
+        // VPC
+        FriggVPC: {
+            Type: 'AWS::EC2::VPC',
+            Properties: {
+                CidrBlock: AppDefinition.vpc.cidrBlock || '10.0.0.0/16',
+                EnableDnsHostnames: true,
+                EnableDnsSupport: true,
+                Tags: [
+                    { Key: 'Name', Value: '${self:service}-${self:provider.stage}-vpc' }
+                ]
+            }
+        },
+
+        // Internet Gateway
+        FriggInternetGateway: {
+            Type: 'AWS::EC2::InternetGateway',
+            Properties: {
+                Tags: [
+                    { Key: 'Name', Value: '${self:service}-${self:provider.stage}-igw' }
+                ]
+            }
+        },
+
+        // Attach Internet Gateway to VPC
+        FriggVPCGatewayAttachment: {
+            Type: 'AWS::EC2::VPCGatewayAttachment',
+            Properties: {
+                VpcId: { Ref: 'FriggVPC' },
+                InternetGatewayId: { Ref: 'FriggInternetGateway' }
+            }
+        },
+
+        // Public Subnet for NAT Gateway
+        FriggPublicSubnet: {
+            Type: 'AWS::EC2::Subnet',
+            Properties: {
+                VpcId: { Ref: 'FriggVPC' },
+                CidrBlock: '10.0.1.0/24',
+                AvailabilityZone: { 'Fn::Select': [0, { 'Fn::GetAZs': '' }] },
+                MapPublicIpOnLaunch: true,
+                Tags: [
+                    { Key: 'Name', Value: '${self:service}-${self:provider.stage}-public-subnet' }
+                ]
+            }
+        },
+
+        // Private Subnet 1 for Lambda
+        FriggPrivateSubnet1: {
+            Type: 'AWS::EC2::Subnet',
+            Properties: {
+                VpcId: { Ref: 'FriggVPC' },
+                CidrBlock: '10.0.2.0/24',
+                AvailabilityZone: { 'Fn::Select': [0, { 'Fn::GetAZs': '' }] },
+                Tags: [
+                    { Key: 'Name', Value: '${self:service}-${self:provider.stage}-private-subnet-1' }
+                ]
+            }
+        },
+
+        // Private Subnet 2 for Lambda (different AZ for redundancy)
+        FriggPrivateSubnet2: {
+            Type: 'AWS::EC2::Subnet',
+            Properties: {
+                VpcId: { Ref: 'FriggVPC' },
+                CidrBlock: '10.0.3.0/24',
+                AvailabilityZone: { 'Fn::Select': [1, { 'Fn::GetAZs': '' }] },
+                Tags: [
+                    { Key: 'Name', Value: '${self:service}-${self:provider.stage}-private-subnet-2' }
+                ]
+            }
+        },
+
+        // Elastic IP for NAT Gateway
+        FriggNATGatewayEIP: {
+            Type: 'AWS::EC2::EIP',
+            Properties: {
+                Domain: 'vpc',
+                Tags: [
+                    { Key: 'Name', Value: '${self:service}-${self:provider.stage}-nat-eip' }
+                ]
+            },
+            DependsOn: 'FriggVPCGatewayAttachment'
+        },
+
+        // NAT Gateway for private subnet internet access
+        FriggNATGateway: {
+            Type: 'AWS::EC2::NatGateway',
+            Properties: {
+                AllocationId: { 'Fn::GetAtt': ['FriggNATGatewayEIP', 'AllocationId'] },
+                SubnetId: { Ref: 'FriggPublicSubnet' },
+                Tags: [
+                    { Key: 'Name', Value: '${self:service}-${self:provider.stage}-nat-gateway' }
+                ]
+            }
+        },
+
+        // Public Route Table
+        FriggPublicRouteTable: {
+            Type: 'AWS::EC2::RouteTable',
+            Properties: {
+                VpcId: { Ref: 'FriggVPC' },
+                Tags: [
+                    { Key: 'Name', Value: '${self:service}-${self:provider.stage}-public-rt' }
+                ]
+            }
+        },
+
+        // Public Route to Internet Gateway
+        FriggPublicRoute: {
+            Type: 'AWS::EC2::Route',
+            Properties: {
+                RouteTableId: { Ref: 'FriggPublicRouteTable' },
+                DestinationCidrBlock: '0.0.0.0/0',
+                GatewayId: { Ref: 'FriggInternetGateway' }
+            },
+            DependsOn: 'FriggVPCGatewayAttachment'
+        },
+
+        // Associate Public Subnet with Public Route Table
+        FriggPublicSubnetRouteTableAssociation: {
+            Type: 'AWS::EC2::SubnetRouteTableAssociation',
+            Properties: {
+                SubnetId: { Ref: 'FriggPublicSubnet' },
+                RouteTableId: { Ref: 'FriggPublicRouteTable' }
+            }
+        },
+
+        // Private Route Table
+        FriggPrivateRouteTable: {
+            Type: 'AWS::EC2::RouteTable',
+            Properties: {
+                VpcId: { Ref: 'FriggVPC' },
+                Tags: [
+                    { Key: 'Name', Value: '${self:service}-${self:provider.stage}-private-rt' }
+                ]
+            }
+        },
+
+        // Private Route to NAT Gateway
+        FriggPrivateRoute: {
+            Type: 'AWS::EC2::Route',
+            Properties: {
+                RouteTableId: { Ref: 'FriggPrivateRouteTable' },
+                DestinationCidrBlock: '0.0.0.0/0',
+                NatGatewayId: { Ref: 'FriggNATGateway' }
+            }
+        },
+
+        // Associate Private Subnet 1 with Private Route Table
+        FriggPrivateSubnet1RouteTableAssociation: {
+            Type: 'AWS::EC2::SubnetRouteTableAssociation',
+            Properties: {
+                SubnetId: { Ref: 'FriggPrivateSubnet1' },
+                RouteTableId: { Ref: 'FriggPrivateRouteTable' }
+            }
+        },
+
+        // Associate Private Subnet 2 with Private Route Table
+        FriggPrivateSubnet2RouteTableAssociation: {
+            Type: 'AWS::EC2::SubnetRouteTableAssociation',
+            Properties: {
+                SubnetId: { Ref: 'FriggPrivateSubnet2' },
+                RouteTableId: { Ref: 'FriggPrivateRouteTable' }
+            }
+        },
+
+        // Security Group for Lambda functions
+        FriggLambdaSecurityGroup: {
+            Type: 'AWS::EC2::SecurityGroup',
+            Properties: {
+                GroupDescription: 'Security group for Frigg Lambda functions',
+                VpcId: { Ref: 'FriggVPC' },
+                SecurityGroupEgress: [
+                    {
+                        IpProtocol: 'tcp',
+                        FromPort: 443,
+                        ToPort: 443,
+                        CidrIp: '0.0.0.0/0',
+                        Description: 'HTTPS outbound'
+                    },
+                    {
+                        IpProtocol: 'tcp',
+                        FromPort: 80,
+                        ToPort: 80,
+                        CidrIp: '0.0.0.0/0',
+                        Description: 'HTTP outbound'
+                    },
+                    {
+                        IpProtocol: 'tcp',
+                        FromPort: 53,
+                        ToPort: 53,
+                        CidrIp: '0.0.0.0/0',
+                        Description: 'DNS TCP'
+                    },
+                    {
+                        IpProtocol: 'udp',
+                        FromPort: 53,
+                        ToPort: 53,
+                        CidrIp: '0.0.0.0/0',
+                        Description: 'DNS UDP'
+                    }
+                ],
+                Tags: [
+                    { Key: 'Name', Value: '${self:service}-${self:provider.stage}-lambda-sg' }
+                ]
+            }
+        }
+    };
+
+    // Add VPC Endpoints for cost optimization
+    if (AppDefinition.vpc.enableVPCEndpoints !== false) {
+        // S3 Gateway Endpoint (free)
+        vpcResources.FriggS3VPCEndpoint = {
+            Type: 'AWS::EC2::VPCEndpoint',
+            Properties: {
+                VpcId: { Ref: 'FriggVPC' },
+                ServiceName: 'com.amazonaws.${self:provider.region}.s3',
+                VpcEndpointType: 'Gateway',
+                RouteTableIds: [
+                    { Ref: 'FriggPrivateRouteTable' }
+                ]
+            }
+        };
+
+        // DynamoDB Gateway Endpoint (free)
+        vpcResources.FriggDynamoDBVPCEndpoint = {
+            Type: 'AWS::EC2::VPCEndpoint',
+            Properties: {
+                VpcId: { Ref: 'FriggVPC' },
+                ServiceName: 'com.amazonaws.${self:provider.region}.dynamodb',
+                VpcEndpointType: 'Gateway',
+                RouteTableIds: [
+                    { Ref: 'FriggPrivateRouteTable' }
+                ]
+            }
+        };
+
+        // KMS Interface Endpoint (paid, but useful if using KMS)
+        if (AppDefinition.encryption?.useDefaultKMSForFieldLevelEncryption === true) {
+            vpcResources.FriggKMSVPCEndpoint = {
+                Type: 'AWS::EC2::VPCEndpoint',
+                Properties: {
+                    VpcId: { Ref: 'FriggVPC' },
+                    ServiceName: 'com.amazonaws.${self:provider.region}.kms',
+                    VpcEndpointType: 'Interface',
+                    SubnetIds: [
+                        { Ref: 'FriggPrivateSubnet1' },
+                        { Ref: 'FriggPrivateSubnet2' }
+                    ],
+                    SecurityGroupIds: [
+                        { Ref: 'FriggVPCEndpointSecurityGroup' }
+                    ],
+                    PrivateDnsEnabled: true
+                }
+            };
+        }
+
+        // Secrets Manager Interface Endpoint (paid, but useful for secrets)
+        vpcResources.FriggSecretsManagerVPCEndpoint = {
+            Type: 'AWS::EC2::VPCEndpoint',
+            Properties: {
+                VpcId: { Ref: 'FriggVPC' },
+                ServiceName: 'com.amazonaws.${self:provider.region}.secretsmanager',
+                VpcEndpointType: 'Interface',
+                SubnetIds: [
+                    { Ref: 'FriggPrivateSubnet1' },
+                    { Ref: 'FriggPrivateSubnet2' }
+                ],
+                SecurityGroupIds: [
+                    { Ref: 'FriggVPCEndpointSecurityGroup' }
+                ],
+                PrivateDnsEnabled: true
+            }
+        };
+
+        // Security Group for VPC Endpoints
+        vpcResources.FriggVPCEndpointSecurityGroup = {
+            Type: 'AWS::EC2::SecurityGroup',
+            Properties: {
+                GroupDescription: 'Security group for Frigg VPC Endpoints',
+                VpcId: { Ref: 'FriggVPC' },
+                SecurityGroupIngress: [
+                    {
+                        IpProtocol: 'tcp',
+                        FromPort: 443,
+                        ToPort: 443,
+                        SourceSecurityGroupId: { Ref: 'FriggLambdaSecurityGroup' },
+                        Description: 'HTTPS from Lambda'
+                    }
+                ],
+                Tags: [
+                    { Key: 'Name', Value: '${self:service}-${self:provider.stage}-vpc-endpoint-sg' }
+                ]
+            }
+        };
+    }
+
+    return vpcResources;
+};
+
+/**
+ * Compose a complete serverless framework configuration from app definition
+ * @param {Object} AppDefinition - Application definition object
+ * @param {string} [AppDefinition.name] - Application name
+ * @param {string} [AppDefinition.provider='aws'] - Cloud provider
+ * @param {Array} AppDefinition.integrations - Array of integration definitions
+ * @param {Object} [AppDefinition.vpc] - VPC configuration
+ * @param {Object} [AppDefinition.encryption] - KMS encryption configuration
+ * @param {Object} [AppDefinition.ssm] - SSM parameter store configuration
+ * @param {Object} [AppDefinition.websockets] - WebSocket configuration
+ * @param {boolean} [AppDefinition.websockets.enable=false] - Enable WebSocket support for live update streaming
+ * @returns {Object} Complete serverless framework configuration
+ */
 const composeServerlessDefinition = (AppDefinition) => {
     const definition = {
         frameworkVersion: '>=3.17.0',
@@ -180,26 +516,6 @@ const composeServerlessDefinition = (AppDefinition) => {
             },
         },
         functions: {
-            defaultWebsocket: {
-                handler: 'node_modules/@friggframework/core/handlers/routers/websocket.handler',
-                events: [
-                    {
-                        websocket: {
-                            route: '$connect',
-                        },
-                    },
-                    {
-                        websocket: {
-                            route: '$default',
-                        },
-                    },
-                    {
-                        websocket: {
-                            route: '$disconnect',
-                        },
-                    },
-                ],
-            },
             auth: {
                 handler: 'node_modules/@friggframework/core/handlers/routers/auth.handler',
                 events: [
@@ -245,7 +561,7 @@ const composeServerlessDefinition = (AppDefinition) => {
                     Type: 'AWS::SQS::Queue',
                     Properties: {
                         QueueName:
-                            'internal-error-queue-${self:provider.stage}',
+                            '${self:service}-internal-error-queue-${self:provider.stage}',
                         MessageRetentionPeriod: 300,
                     },
                 },
@@ -329,6 +645,7 @@ const composeServerlessDefinition = (AppDefinition) => {
         },
     };
 
+
     // KMS Configuration based on App Definition
     if (AppDefinition.encryption?.useDefaultKMSForFieldLevelEncryption === true) {
         // Add KMS IAM permissions
@@ -347,12 +664,181 @@ const composeServerlessDefinition = (AppDefinition) => {
         // Add serverless-kms-grants plugin
         definition.plugins.push('serverless-kms-grants');
 
-        // Configure KMS grants with default key
+        // Configure KMS grants with discovered default key
         definition.custom.kmsGrants = {
-            kmsKeyId: '*'
+            kmsKeyId: '${env:AWS_DISCOVERY_KMS_KEY_ID}'
         };
     }
 
+    // VPC Configuration based on App Definition
+    if (AppDefinition.vpc?.enable === true) {
+        // Add VPC-related IAM permissions
+        definition.provider.iamRoleStatements.push({
+            Effect: 'Allow',
+            Action: [
+                'ec2:CreateNetworkInterface',
+                'ec2:DescribeNetworkInterfaces',
+                'ec2:DeleteNetworkInterface',
+                'ec2:AttachNetworkInterface',
+                'ec2:DetachNetworkInterface'
+            ],
+            Resource: '*'
+        });
+
+        // Default approach: Use AWS Discovery to find existing VPC resources
+        if (AppDefinition.vpc.createNew === true) {
+            // Option 1: Create new VPC infrastructure (explicit opt-in)
+            const vpcConfig = {};
+
+            if (AppDefinition.vpc.securityGroupIds) {
+                // User provided custom security groups
+                vpcConfig.securityGroupIds = AppDefinition.vpc.securityGroupIds;
+            } else {
+                // Use auto-created security group
+                vpcConfig.securityGroupIds = [{ Ref: 'FriggLambdaSecurityGroup' }];
+            }
+
+            if (AppDefinition.vpc.subnetIds) {
+                // User provided custom subnets
+                vpcConfig.subnetIds = AppDefinition.vpc.subnetIds;
+            } else {
+                // Use auto-created private subnets
+                vpcConfig.subnetIds = [
+                    { Ref: 'FriggPrivateSubnet1' },
+                    { Ref: 'FriggPrivateSubnet2' }
+                ];
+            }
+
+            // Set VPC config for Lambda functions
+            definition.provider.vpc = vpcConfig;
+
+            // Add VPC infrastructure resources to CloudFormation
+            const vpcResources = createVPCInfrastructure(AppDefinition);
+            Object.assign(definition.resources.Resources, vpcResources);
+        } else {
+            // Option 2: Use AWS Discovery (default behavior)
+            // VPC configuration using discovered or explicitly provided resources
+            const vpcConfig = {
+                securityGroupIds: AppDefinition.vpc.securityGroupIds || 
+                                  ['${env:AWS_DISCOVERY_SECURITY_GROUP_ID}'],
+                subnetIds: AppDefinition.vpc.subnetIds || [
+                    '${env:AWS_DISCOVERY_SUBNET_ID_1}',
+                    '${env:AWS_DISCOVERY_SUBNET_ID_2}'
+                ]
+            };
+
+            // Set VPC config for Lambda functions
+            definition.provider.vpc = vpcConfig;
+
+            // Add NAT Gateway for Lambda internet access (required for external API calls)
+            // Even with existing VPC, Lambdas need guaranteed internet access for integrations
+            definition.resources.Resources.FriggNATGatewayEIP = {
+                Type: 'AWS::EC2::EIP',
+                Properties: {
+                    Domain: 'vpc',
+                    Tags: [
+                        { Key: 'Name', Value: '${self:service}-${self:provider.stage}-nat-eip' }
+                    ]
+                }
+            };
+
+            definition.resources.Resources.FriggNATGateway = {
+                Type: 'AWS::EC2::NatGateway',
+                Properties: {
+                    AllocationId: { 'Fn::GetAtt': ['FriggNATGatewayEIP', 'AllocationId'] },
+                    SubnetId: '${env:AWS_DISCOVERY_PUBLIC_SUBNET_ID}', // Discovery finds public subnet
+                    Tags: [
+                        { Key: 'Name', Value: '${self:service}-${self:provider.stage}-nat-gateway' }
+                    ]
+                }
+            };
+
+            // Create route table for Lambda subnets to use NAT Gateway
+            definition.resources.Resources.FriggLambdaRouteTable = {
+                Type: 'AWS::EC2::RouteTable',
+                Properties: {
+                    VpcId: '${env:AWS_DISCOVERY_VPC_ID}',
+                    Tags: [
+                        { Key: 'Name', Value: '${self:service}-${self:provider.stage}-lambda-rt' }
+                    ]
+                }
+            };
+
+            definition.resources.Resources.FriggNATRoute = {
+                Type: 'AWS::EC2::Route',
+                Properties: {
+                    RouteTableId: { Ref: 'FriggLambdaRouteTable' },
+                    DestinationCidrBlock: '0.0.0.0/0',
+                    NatGatewayId: { Ref: 'FriggNATGateway' }
+                }
+            };
+
+            // Associate Lambda subnets with NAT Gateway route table
+            definition.resources.Resources.FriggSubnet1RouteAssociation = {
+                Type: 'AWS::EC2::SubnetRouteTableAssociation',
+                Properties: {
+                    SubnetId: '${env:AWS_DISCOVERY_SUBNET_ID_1}',
+                    RouteTableId: { Ref: 'FriggLambdaRouteTable' }
+                }
+            };
+
+            definition.resources.Resources.FriggSubnet2RouteAssociation = {
+                Type: 'AWS::EC2::SubnetRouteTableAssociation',
+                Properties: {
+                    SubnetId: '${env:AWS_DISCOVERY_SUBNET_ID_2}',
+                    RouteTableId: { Ref: 'FriggLambdaRouteTable' }
+                }
+            };
+
+            // Add VPC endpoints for AWS service optimization (optional but recommended)
+            if (AppDefinition.vpc.enableVPCEndpoints !== false) {
+                definition.resources.Resources.VPCEndpointS3 = {
+                    Type: 'AWS::EC2::VPCEndpoint',
+                    Properties: {
+                        VpcId: '${env:AWS_DISCOVERY_VPC_ID}',
+                        ServiceName: 'com.amazonaws.${self:provider.region}.s3',
+                        VpcEndpointType: 'Gateway',
+                        RouteTableIds: [{ Ref: 'FriggLambdaRouteTable' }]
+                    }
+                };
+
+                definition.resources.Resources.VPCEndpointDynamoDB = {
+                    Type: 'AWS::EC2::VPCEndpoint',
+                    Properties: {
+                        VpcId: '${env:AWS_DISCOVERY_VPC_ID}',
+                        ServiceName: 'com.amazonaws.${self:provider.region}.dynamodb',
+                        VpcEndpointType: 'Gateway',
+                        RouteTableIds: [{ Ref: 'FriggLambdaRouteTable' }]
+                    }
+                };
+            }
+        }
+    }
+
+    // SSM Parameter Store Configuration based on App Definition  
+    if (AppDefinition.ssm?.enable === true) {
+        // Add AWS Parameters and Secrets Lambda Extension layer
+        definition.provider.layers = [
+            'arn:aws:lambda:${self:provider.region}:177933569100:layer:AWS-Parameters-and-Secrets-Lambda-Extension:11'
+        ];
+
+        // Add SSM IAM permissions
+        definition.provider.iamRoleStatements.push({
+            Effect: 'Allow',
+            Action: [
+                'ssm:GetParameter',
+                'ssm:GetParameters',
+                'ssm:GetParametersByPath'
+            ],
+            Resource: [
+                'arn:aws:ssm:${self:provider.region}:*:parameter/${self:service}/${self:provider.stage}/*'
+            ]
+        });
+
+        // Add environment variable for SSM parameter prefix
+        definition.provider.environment.SSM_PARAMETER_PREFIX = '/${self:service}/${self:provider.stage}';
+    }
+
     // Add integration-specific functions and resources
     for (const integration of AppDefinition.integrations) {
         const integrationName = integration.Definition.name;
@@ -419,6 +905,51 @@ const composeServerlessDefinition = (AppDefinition) => {
         definition.custom[queueReference] = queueName;
     }
 
+    // Check if AWS discovery environment variables are missing and error
+    const needsDiscovery = AppDefinition.vpc?.enable || 
+                          AppDefinition.encryption?.useDefaultKMSForFieldLevelEncryption ||
+                          AppDefinition.ssm?.enable;
+    
+    if (needsDiscovery && !process.env.AWS_DISCOVERY_VPC_ID) {
+        console.error('❌ AWS discovery environment variables not found');
+        console.error('');
+        console.error('🚨 Your AppDefinition requires AWS discovery but variables are missing:');
+        if (AppDefinition.vpc?.enable) console.error('   ❌ VPC support (vpc.enable: true)');
+        if (AppDefinition.encryption?.useDefaultKMSForFieldLevelEncryption) console.error('   ❌ KMS encryption (encryption.useDefaultKMSForFieldLevelEncryption: true)');
+        if (AppDefinition.ssm?.enable) console.error('   ❌ SSM parameters (ssm.enable: true)');
+        console.error('');
+        console.error('💡 Run AWS discovery before building:');
+        console.error('   node node_modules/@friggframework/devtools/infrastructure/run-discovery.js');
+        console.error('');
+        console.error('🔧 Or use the build command which runs discovery automatically:');
+        console.error('   npx frigg build');
+        throw new Error('AWS discovery required but environment variables not found');
+    }
+
+    // Add websocket function if enabled
+    if (AppDefinition.websockets?.enable === true) {
+        definition.functions.defaultWebsocket = {
+            handler: 'node_modules/@friggframework/core/handlers/routers/websocket.handler',
+            events: [
+                {
+                    websocket: {
+                        route: '$connect',
+                    },
+                },
+                {
+                    websocket: {
+                        route: '$default',
+                    },
+                },
+                {
+                    websocket: {
+                        route: '$disconnect',
+                    },
+                },
+            ],
+        };
+    }
+
     // Modify handler paths to point to the correct node_modules location
     definition.functions = modifyHandlerPaths(definition.functions);