From 5f4b41c501e1fce676d6569f4340a285f2e3a1a4 Mon Sep 17 00:00:00 2001
From: Jarrod Johnson <jjohnson2@lenovo.com>
Date: Tue, 13 Oct 2026 09:29:08 -0400
Subject: [PATCH] Cleanup and handle reconnection of nodeconsole -tv

Also fix screenshot to vnconly target
---
 confluent_client/bin/nodeconsole  |  8 +++++++-
 confluent_client/confluent/vnc.py | 20 ++++++++++++++++----
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/confluent_client/bin/nodeconsole b/confluent_client/bin/nodeconsole
index 2662c8e4..5bd0def1 100755
--- a/confluent_client/bin/nodeconsole
+++ b/confluent_client/bin/nodeconsole
@@ -1030,7 +1030,7 @@ async def relay_keypresses(keys, modifiers=None):
     for node in focused_nodes:
         vnclient = vncclientsbynode.get(node)
         if not vnclient:
-            sys.stderr.write(f"No VNC client for node {node}, unable to relay keypress\n")
+            #sys.stderr.write(f"No VNC client for node {node}, unable to relay keypress\n")
             continue
         try:
             await vnclient.send_keypresses(keys, modifiers)
@@ -1108,6 +1108,9 @@ async def do_vnc(node, url, outputfile=None):
                         # without the timeout, we lose track of proxmox reset
                         # with the timeout, we falsely assume dead on stale
                         break
+                    except vnc.VNCReceiverTerminated:
+                        # The VNC receiver task has terminated unexpectedly
+                        break
                     retries = 5
                     # Save as PNG using PIL/pillow
                     outfile = io.BytesIO()
@@ -1118,6 +1121,9 @@ async def do_vnc(node, url, outputfile=None):
                     else:
                         keeprunning = False
                         break
+                    if not streaming:
+                        keeprunning = False
+                        break
         except ValueError as e:
             draw_node(node, None, str(e), '', cwidth, cheight)
             retries -= 1
diff --git a/confluent_client/confluent/vnc.py b/confluent_client/confluent/vnc.py
index a93c9506..0265830e 100644
--- a/confluent_client/confluent/vnc.py
+++ b/confluent_client/confluent/vnc.py
@@ -7,6 +7,8 @@ import threading
 import time
 import zlib
 
+class VNCReceiverTerminated(Exception):
+    pass
 # This results in an RGBA organization of pixels
 MYPIXFORMAT = bytearray([
     32,  # bits per pixel
@@ -141,9 +143,17 @@ class VNCClient:
         self.writer.write(data)
         return data
 
-    async def get_screenshot(self):
-        while self._updating:
-            await asyncio.sleep(0.1)
+    async def get_screenshot(self, oneshot=False):
+        # Oneshot is wanting a full screenshot, right now
+        # this means we wait up to half a second getting screenshots,
+        # then return whatever the result is after the accumulated updates over the half second
+        if oneshot:
+            await asyncio.sleep(0.5)  # Wait up to half a second for updates
+        else:
+            while self._updating and not self.receiver.done():
+                await asyncio.sleep(0.1)
+        if self.receiver.done():
+            raise VNCReceiverTerminated('VNC receiver task has terminated unexpectedly')
         await asyncio.sleep(0)
         if self.framebuffer is None:
             raise Exception('No framebuffer data available')
@@ -152,6 +162,8 @@ class VNCClient:
 
     async def _do_vnc_handshake(self):
         rfbver = await self.reader.readline()
+        if not rfbver:
+            raise Exception('Connection failed to VNC')
         if not rfbver.startswith(b'RFB 003.008'):
             self.writer.close()
             await self.writer.wait_closed()
@@ -224,7 +236,7 @@ class VNCClient:
                 else:
                     raise Exception(f'Unknown message type: {message_type}')
             except Exception as e:
-                print(f"Error in receive loop: {e}")
+                #print(f"Error in receive loop: {e}")
                 break        
 
     async def _handle_framebuffer_update(self):