From 07e9cd36a98da78cace04c99f9146ec4cda6d82b Mon Sep 17 00:00:00 2001
From: Reginald Alfret <reginaldalfret@gmail.com>
Date: Fri, 25 Sep 2026 10:31:21 +0530
Subject: [PATCH] fix(flask): handle missing segment gracefully in
 after_request

When an earlier Flask before_request extension short-circuits or raises, X-Ray before_request might not create a segment. Guard _after_request to return the response gracefully when no segment is open, matching the pattern in _teardown_request.

Fixes #405
---
 aws_xray_sdk/ext/flask/middleware.py | 15 +++++++++++----
 tests/ext/flask/test_flask.py        | 17 +++++++++++++++++
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/aws_xray_sdk/ext/flask/middleware.py b/aws_xray_sdk/ext/flask/middleware.py
index fdc3b32f..aca5a6ea 100644
--- a/aws_xray_sdk/ext/flask/middleware.py
+++ b/aws_xray_sdk/ext/flask/middleware.py
@@ -67,10 +67,17 @@ def _before_request(self):
             segment.put_http_meta(http.CLIENT_IP, req.remote_addr)
 
     def _after_request(self, response):
-        if self.in_lambda_ctx:
-            segment = self._recorder.current_subsegment()
-        else:
-            segment = self._recorder.current_segment()
+        segment = None
+        try:
+            if self.in_lambda_ctx:
+                segment = self._recorder.current_subsegment()
+            else:
+                segment = self._recorder.current_segment()
+        except Exception:
+            pass
+        if not segment:
+            return response
+
         segment.put_http_meta(http.STATUS, response.status_code)
 
         origin_header = segment.get_origin_trace_header()
diff --git a/tests/ext/flask/test_flask.py b/tests/ext/flask/test_flask.py
index 283c0e89..e39f95c3 100644
--- a/tests/ext/flask/test_flask.py
+++ b/tests/ext/flask/test_flask.py
@@ -238,3 +238,20 @@ def subsegment():
     app_client.get(path)
     segment = recorder.emitter.pop()
     assert not segment  # Segment should be none because it's created and ended by the middleware
+
+
+def test_missing_segment_in_after_request():
+    new_recorder = get_new_stubbed_recorder()
+    new_recorder.configure(service='test', sampling=False)
+    new_app = Flask(__name__)
+
+    @new_app.before_request
+    def early_before_request():
+        return 'Early response', 400
+
+    XRayMiddleware(new_app, new_recorder)
+    app_client = new_app.test_client()
+
+    response = app_client.get('/early')
+    assert response.status_code == 400
+    assert response.data == b'Early response'