From d4f1a1c6a4fd68e8f61e19e86a4a23fdb5653125 Mon Sep 17 00:00:00 2001
From: Rune Olsen <rune@rune.pm>
Date: Wed, 7 Jan 2026 08:01:33 +0100
Subject: [PATCH] Fixed bug in ctrl+c handling during response streaming

---
 oai.py | 61 ++++++++++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 47 insertions(+), 14 deletions(-)

diff --git a/oai.py b/oai.py
index e679ce5..be8c1b7 100644
--- a/oai.py
+++ b/oai.py
@@ -4985,26 +4985,59 @@ All queries are read-only. INSERT/UPDATE/DELETE are not allowed."""
             # PROCESS FINAL RESPONSE
             # ========================================================================
             full_response = ""
+            stream_interrupted = False
+
             if is_streaming:
                 try:
                     with Live("", console=console, refresh_per_second=10, auto_refresh=True) as live:
-                        for chunk in response:
-                            if hasattr(chunk, 'error') and chunk.error:
-                                console.print(f"\n[bold red]Stream error: {chunk.error.message}[/]")
-                                app_logger.error(f"Stream error: {chunk.error.message}")
-                                break
-                            if hasattr(chunk.choices[0].delta, 'content') and chunk.choices[0].delta.content:
-                                content_chunk = chunk.choices[0].delta.content
-                                full_response += content_chunk
-                                md = Markdown(full_response)
-                                live.update(md)
+                        try:
+                            for chunk in response:
+                                if hasattr(chunk, 'error') and chunk.error:
+                                    console.print(f"\n[bold red]Stream error: {chunk.error.message}[/]")
+                                    app_logger.error(f"Stream error: {chunk.error.message}")
+                                    break
+                                if hasattr(chunk.choices[0].delta, 'content') and chunk.choices[0].delta.content:
+                                    content_chunk = chunk.choices[0].delta.content
+                                    full_response += content_chunk
+                                    md = Markdown(full_response)
+                                    live.update(md)
+                        except KeyboardInterrupt:
+                            stream_interrupted = True
+                            console.print("\n[bold yellow]⚠️  Streaming interrupted![/]")
+                            app_logger.info("Streaming interrupted by user (Ctrl+C)")
+                        except Exception as stream_error:
+                            stream_interrupted = True
+                            console.print(f"\n[bold red]Stream error: {stream_error}[/]")
+                            app_logger.error(f"Stream processing error: {stream_error}")
                     
-                    console.print("")
+                    if not stream_interrupted:
+                        console.print("")
                     
                 except KeyboardInterrupt:
-                    console.print("\n[bold yellow]Streaming cancelled![/]")
-                    app_logger.info("Streaming cancelled by user")
-                    continue
+                    # Outer interrupt handler (in case inner misses it)
+                    stream_interrupted = True
+                    console.print("\n[bold yellow]⚠️  Streaming interrupted![/]")
+                    app_logger.info("Streaming interrupted by user (outer)")
+                except Exception as e:
+                    stream_interrupted = True
+                    console.print(f"\n[bold red]Error during streaming: {e}[/]")
+                    app_logger.error(f"Streaming error: {e}")
+                
+                # If stream was interrupted, skip processing and continue to next prompt
+                if stream_interrupted:
+                    if full_response:
+                        console.print(f"\n[dim yellow]Partial response received ({len(full_response)} chars). Discarding...[/]")
+                    console.print("[dim blue]💡 Ready for next prompt[/]\n")
+                    app_logger.info("Stream cleanup completed, returning to prompt")
+                    
+                    # Force close the response if possible
+                    try:
+                        if hasattr(response, 'close'):
+                            response.close()
+                    except:
+                        pass
+                    
+                    continue  # Now it's safe to continue
             else:
                 full_response = response.choices[0].message.content if response.choices else ""
                 console.print(f"\r{' ' * 50}\r", end="")