diff --git a/src/codegate/pipeline/output.py b/src/codegate/pipeline/output.py
index 4f0706c2..b7ce419b 100644
--- a/src/codegate/pipeline/output.py
+++ b/src/codegate/pipeline/output.py
@@ -169,8 +169,9 @@ async def process_stream(
         finally:
             # NOTE: Don't use await in finally block, it will break the stream
             # Don't flush the buffer if we assume we'll call the pipeline again
-            if cleanup_sensitive is False and finish_stream:
-                self._record_to_db()
+            if cleanup_sensitive is False:
+                if finish_stream:
+                    self._record_to_db()
                 return
 
             # TODO figure out what's the logic here.
diff --git a/src/codegate/providers/copilot/provider.py b/src/codegate/providers/copilot/provider.py
index b1b944ed..e5787efc 100644
--- a/src/codegate/providers/copilot/provider.py
+++ b/src/codegate/providers/copilot/provider.py
@@ -871,16 +871,15 @@ async def stream_iterator():
                 while not self.stream_queue.empty():
                     incoming_record = await self.stream_queue.get()
                     for choice in incoming_record.choices:
-                        if choice.finish_reason and \
-                                choice.finish_reason in ["stop", "length", "content_filter", "tool_calls"]:
+                        if choice.finish_reason and choice.finish_reason is not None:
                             self.finish_stream = True
                     yield incoming_record
 
             # needs to be set as the flag gets reset on finish_data
             finish_stream_flag = any(
-                choice.get("finish_reason") == "stop"
+                choice.finish_reason is not None
                 for record in list(self.stream_queue._queue)
-                for choice in record.get("content", {}).get("choices", [])
+                for choice in record.choices
             )
             async for record in self.output_pipeline_instance.process_stream(
                 stream_iterator(),