add heartbeat for streaming responses to prevent idle timeout of 10 mins

aws-samples · Jan 16, 2025 · 44dd09f · 44dd09f
1 parent de518e5
commit 44dd09f
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 2 deletions.
diff --git a/lex-web-ui/src/store/actions.js b/lex-web-ui/src/store/actions.js
@@ -710,11 +710,11 @@ export default {
 
           wsClient.onmessage = (event) => {
             if(event.data!=='/stop/' && context.getters.isStartingTypingWsMessages()){
-              console.info("streaming ", context.getters.isStartingTypingWsMessages());
+              console.info('Streaming: ', context.getters.isStartingTypingWsMessages());
               context.commit('pushWebSocketMessage',event.data);
               context.dispatch('typingWsMessages')
             }else{
-              console.info('stopping streaming');
+              console.info('Currently not streaming');
             }
           }
         }
@@ -1304,6 +1304,41 @@ export default {
 
     const signedUrl = Signer.signUrl(context.state.config.lex.streamingWebSocketEndpoint+'?sessionId='+sessionId, accessInfo, serviceInfo);
     wsClient = new WebSocket(signedUrl);
+
+    // Add heartbeat logic
+    const HEARTBEAT_INTERVAL = 540000; // 9 minutes
+    const MAX_DURATION = 7200000; // 2 hours
+    const startTime = Date.now();
+    let heartbeatTimer = null;
+
+    function startHeartbeat() {
+        if (wsClient.readyState === WebSocket.OPEN) {
+            const elapsedTime = Date.now() - startTime;
+            if (elapsedTime < MAX_DURATION) {
+                const pingMessage = JSON.stringify({ action: 'ping' });
+                wsClient.send(pingMessage);
+                console.log('Sending Ping:', new Date().toISOString());
+                heartbeatTimer = setTimeout(startHeartbeat, HEARTBEAT_INTERVAL);
+            } else {
+                console.log('Stopped sending pings after reaching 2-hour limit.');
+                clearTimeout(heartbeatTimer);
+            }
+        }
+    }
+    wsClient.onopen = () => {
+      console.log('WebSocket Connected');
+      startHeartbeat();
+    };
+
+    wsClient.onclose = () => {
+        console.log('WebSocket Closed');
+        clearTimeout(heartbeatTimer);
+    };
+
+    wsClient.onerror = (error) => {
+        console.log('WebSocket Error', error.message);
+        clearTimeout(heartbeatTimer);
+    };
   },
   typingWsMessages(context){
     if (context.getters.wsMessagesCurrentIndex()<context.getters.wsMessagesLength()-1){

diff --git a/templates/streaming-support.yaml b/templates/streaming-support.yaml
@@ -169,11 +169,34 @@ Resources:
       IntegrationUri: 
         Fn::Sub:
             arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${StreamingLambda.Arn}/invocations    
+
+  PingRoute:
+    Type: AWS::ApiGatewayV2::Route
+    Properties:
+      ApiId: !Ref StreamingWebSocket
+      RouteKey: ping
+      OperationName: PingRoute
+      Target: !Join
+        - '/'
+        - - 'integrations'
+          - !Ref PingIntegration
+
+  PingIntegration:
+    Type: AWS::ApiGatewayV2::Integration
+    Properties:
+      ApiId: !Ref StreamingWebSocket
+      Description: Ping Integration
+      IntegrationType: MOCK
+      PassthroughBehavior: WHEN_NO_MATCH
+      RequestTemplates:
+        application/json: '{"statusCode":200}'
+
 
   Deployment:
     Type: AWS::ApiGatewayV2::Deployment
     DependsOn:
     - ConnectRoute
+    - PingRoute
     Properties:
       ApiId: !Ref StreamingWebSocket