wrap response to support trigger close

examples/flask/hello.py

@@ -45,6 +45,7 @@ def mask_event(eventmodel):
     'GET_METADATA': get_metadata,
     'CAPTURE_OUTGOING_REQUESTS': False,
 }
+
 app.wsgi_app = MoesifMiddleware(app.wsgi_app, moesif_settings)
 
 @app.route('/')
@@ -53,7 +54,12 @@ def index():
 
 @app.route('/hello')
 def hello():
-    return 'Hello, world'
+      # Create a response object
+    response = Response('Hello, world')
+
+    # Attach a function to call_on_close
+    response.call_on_close(lambda: print("Response has been sent to the client."))
+    return response;
 
 @app.route('/user/<username>')
 def show_user_profile(username):
@@ -80,10 +86,16 @@ def show_post(post_id):
     }
 ]
 
+@app.route('/error', methods=['GET'])
+def error():
+    raise ValueError("This is a test error!")
+
 @app.route('/todo/api/v1.0/tasks', methods=['GET'])
 def get_tasks():
     json_body = json.dumps({'tasks': tasks})
-    return Response(response=json_body, status=201, mimetype='application/json')
+    response = Response(response=json_body, status=201, mimetype='application/json')
+    response.call_on_close(lambda: print("Response has been sent to the client."));
+    return response
 
 
 @app.route('/todo/api/v1.0/tasks/<int:task_id>', methods=['GET'])

moesifwsgi/middleware.py

@@ -172,7 +172,8 @@ def _start_response(status, response_headers, *args):
           blocked_by = governed_response['blocked_by']
         else:
           # trigger next step in the process
-          response_chunks = event_info.finish_response(self.app(environ, _start_response))
+          original_response = self.app(environ, _start_response);
+          response_chunks = event_info.finish_response(original_response)
 
         # Add response chunks and response headers to the environ
         environ["moesif.response_body_chunks"] = response_chunks
@@ -186,10 +187,31 @@ def _start_response(status, response_headers, *args):
         self.add_user_and_metadata(event_info, environ, response_headers_mapping)
 
         try:
-            return response_chunks
+            return self.wrap_response(response_chunks, original_response)
         finally:
             self.process_and_add_event_if_required(event_info, environ, response_headers_mapping, blocked_by)
 
+    def wrap_response(self, response_chunks, original_response):
+        class WrappedResponse:
+            def __init__(self, chunks, original):
+                print("creating wrapped response")
+                self.chunks = chunks
+                self.original = original
+
+            def __iter__(self):
+                return iter(self.chunks)
+
+            def close(self):
+                print("response close is triggered")
+                if hasattr(self.original, 'close'):
+                    self.original.close()
+
+            def __getattr__(self, name):
+                # Delegate attribute access to the original response
+                return getattr(self.original, name)
+
+        return WrappedResponse(response_chunks, original_response)
+
     def prepare_event_info(self, environ, start_response, request_time):
         event_info = DataHolder(
             self.settings.get("DISABLED_TRANSACTION_ID", False),