diff --git a/tests/test_run/test_websocket_socket.py b/tests/test_run/test_websocket_socket.py
index daf531c..324b60f 100644
--- a/tests/test_run/test_websocket_socket.py
+++ b/tests/test_run/test_websocket_socket.py
@@ -28,7 +28,14 @@
     TestSuiteExecution,
     TestSuiteMetadata,
 )
-from th_cli.test_run.socket_schemas import TestCaseUpdate, TestRunUpdate, TestStepUpdate, TestSuiteUpdate, TestUpdate
+from th_cli.test_run.socket_schemas import (
+    TestCaseUpdate,
+    TestLogRecord,
+    TestRunUpdate,
+    TestStepUpdate,
+    TestSuiteUpdate,
+    TestUpdate,
+)
 from th_cli.test_run.websocket import TestRunSocket
 
 # ---------------------------------------------------------------------------
@@ -304,7 +311,7 @@ async def test_step_update_routed_correctly(self):
             ),
         )
         with patch.object(s, "_TestRunSocket__log_test_step_update") as mock_fn:
-            await s._TestRunSocket__handle_test_update(socket=AsyncMock(), update=update)
+            await s._TestRunSocket__handle_test_update(update=update)
 
         mock_fn.assert_called_once()
 
@@ -319,7 +326,7 @@ async def test_case_update_routed_correctly(self):
             body=TestCaseUpdate(state="passed", test_case_execution_index=0, test_suite_execution_index=0),
         )
         with patch.object(s, "_TestRunSocket__log_test_case_update") as mock_fn:
-            await s._TestRunSocket__handle_test_update(socket=AsyncMock(), update=update)
+            await s._TestRunSocket__handle_test_update(update=update)
 
         mock_fn.assert_called_once()
 
@@ -333,28 +340,55 @@ async def test_suite_update_routed_correctly(self):
             body=TestSuiteUpdate(state="passed", test_suite_execution_index=0),
         )
         with patch.object(s, "_TestRunSocket__log_test_suite_update") as mock_fn:
-            await s._TestRunSocket__handle_test_update(socket=AsyncMock(), update=update)
+            await s._TestRunSocket__handle_test_update(update=update)
 
         mock_fn.assert_called_once()
 
     @pytest.mark.asyncio
-    async def test_run_update_executing_does_not_close_socket(self):
+    async def test_run_update_executing_leaves_run_not_finished(self):
         s = _make_socket()
-        mock_socket = AsyncMock()
 
         update = TestUpdate(test_type="test_run", body=TestRunUpdate(state="executing", test_run_execution_id=1))
         with patch.object(s, "_TestRunSocket__log_test_run_update", new_callable=AsyncMock):
-            await s._TestRunSocket__handle_test_update(socket=mock_socket, update=update)
+            await s._TestRunSocket__handle_test_update(update=update)
 
-        mock_socket.close.assert_not_called()
+        assert s._run_finished is False
 
     @pytest.mark.asyncio
-    async def test_run_update_non_executing_closes_socket(self):
+    async def test_run_update_non_executing_marks_run_finished(self):
         s = _make_socket()
-        mock_socket = AsyncMock()
 
         update = TestUpdate(test_type="test_run", body=TestRunUpdate(state="passed", test_run_execution_id=1))
         with patch.object(s, "_TestRunSocket__log_test_run_update", new_callable=AsyncMock):
-            await s._TestRunSocket__handle_test_update(socket=mock_socket, update=update)
+            await s._TestRunSocket__handle_test_update(update=update)
+
+        assert s._run_finished is True
+
+    @pytest.mark.asyncio
+    async def test_run_update_pending_leaves_run_not_finished(self):
+        # Regression test: "pending" is non-terminal (backend's TestRun.completed()
+        # excludes both PENDING and EXECUTING), so it must not close the socket.
+        # A prior implementation used a negation check (`state != "executing"`)
+        # that misclassified any non-"executing" state, including "pending", as
+        # terminal.
+        s = _make_socket()
+
+        update = TestUpdate(test_type="test_run", body=TestRunUpdate(state="pending", test_run_execution_id=1))
+        with patch.object(s, "_TestRunSocket__log_test_run_update", new_callable=AsyncMock):
+            await s._TestRunSocket__handle_test_update(update=update)
+
+        assert s._run_finished is False
+
+    @pytest.mark.asyncio
+    async def test_handle_log_record_logs_every_record(self):
+        s = _make_socket()
+        records = [
+            TestLogRecord(level="INFO", timestamp=0.0, message=f"msg{i}") for i in range(3)
+        ]
+
+        with patch("th_cli.test_run.websocket.logger") as mock_logger:
+            await s._TestRunSocket__handle_log_record(records)
 
-        mock_socket.close.assert_called_once()
+        assert mock_logger.log.call_count == 3
+        for record in records:
+            mock_logger.log.assert_any_call(record.level, record.message)
diff --git a/th_cli/commands/run_tests.py b/th_cli/commands/run_tests.py
index 6b7532a..beeac27 100644
--- a/th_cli/commands/run_tests.py
+++ b/th_cli/commands/run_tests.py
@@ -224,6 +224,7 @@ async def run_tests(
             execution_pics=execution_pics,
             project_id=project_id,
         )
+        test_logging.set_download_run_id(new_test_run.id)
         if _contains_webrtc_two_way_talk(selected_tests_dict):
             _webrtc_handler = TwoWayTalkHandler(port=8999)
             _webrtc_handler.start_waiting()
diff --git a/th_cli/test_run/log_stream_handler.py b/th_cli/test_run/log_stream_handler.py
index 0c56aca..dc1cd67 100644
--- a/th_cli/test_run/log_stream_handler.py
+++ b/th_cli/test_run/log_stream_handler.py
@@ -28,7 +28,7 @@ class LogStreamHandler:
 
     def __init__(self, port: int = 8998):
         """Initialize the log stream handler.
-        
+
         Args:
             port: Port number for the HTTP server (default: 8998)
         """
@@ -36,48 +36,64 @@ def __init__(self, port: int = 8998):
         self.http_server = LogsHTTPServer(port=port)
         self.log_queue: queue.Queue = queue.Queue(maxsize=1000)
         self.is_running = False
-        self.log_file_path: Optional[str] = None
-        
-    def start(self, test_run_title: str = "Test Execution", log_file_path: Optional[str] = None) -> str:
+
+    def start(self, test_run_title: str = "Test Execution") -> str:
         """Start the log streaming HTTP server.
-        
+
         Args:
             test_run_title: Title of the test run for display
-            log_file_path: Path to the log file for download functionality
-            
+
         Returns:
             URL where logs can be viewed
         """
         if self.is_running:
             logger.warning("Log stream handler already running")
             return self._get_log_viewer_url()
-        
+
         try:
-            # Store log file path for download functionality
-            self.log_file_path = log_file_path
-            
             # Get local IP address
             local_ip = self._get_local_ip()
-            
+
             # Start HTTP server
             self.http_server.start(
                 log_queue=self.log_queue,
                 test_run_title=test_run_title,
                 local_ip=local_ip,
-                log_file_path=log_file_path,
             )
-            
+
             self.is_running = True
-            
+
             viewer_url = f"http://{local_ip}:{self.port}"
             logger.info(f"Log stream viewer started: {viewer_url}")
-            
+
             return viewer_url
-            
+
         except Exception as e:
             logger.error(f"Failed to start log stream handler: {e}")
             raise
-    
+
+    def set_run_id(self, run_id: int) -> None:
+        """Tell the HTTP server which run's log to link "Download Logs" to,
+        once the run has been created (its id isn't known when the server
+        starts).
+        """
+        if not self.is_running:
+            return
+
+        self.http_server.set_run_id(run_id)
+
+        # A viewer may already be connected (the run_id is typically set
+        # only *after* the viewer URL was printed and likely opened), so
+        # also push it through the existing SSE stream as a control message
+        # - a future/refreshed page load will pick it up from the HTTP
+        # server attribute above, but an already-open one only sees this.
+        try:
+            self.log_queue.put_nowait({"__event__": "run_id", "run_id": run_id})
+        except queue.Full:
+            # Best-effort: a future page load/refresh will still pick up
+            # the run id via the HTTP server attribute set above.
+            pass
+
     def stop(self):
         """Stop the log streaming HTTP server."""
         if not self.is_running:
diff --git a/th_cli/test_run/log_viewer.html b/th_cli/test_run/log_viewer.html
index 4e38ca6..718b45e 100644
--- a/th_cli/test_run/log_viewer.html
+++ b/th_cli/test_run/log_viewer.html
@@ -184,6 +184,8 @@
             cursor: pointer;
             transition: all 0.2s;
             font-weight: 500;
+            display: inline-block;
+            text-decoration: none;
         }}
 
         .btn:hover {{
@@ -496,7 +498,7 @@
                     <div class="control-group">
                         <button id="clearBtn" class="btn" onclick="clearLogs()">Clear</button>
                         <button id="autoScrollBtn" class="btn active" onclick="toggleAutoScroll()">Auto-scroll</button>
-                        <button id="downloadBtn" class="btn download" onclick="downloadLogs()">Download Logs</button>
+                        <a id="downloadBtn" class="btn download" href="#" onclick="return false;" title="Available once the run starts" style="opacity:0.5;cursor:default;">Download Logs</a>
                         <input type="text" id="searchBox" class="search-box" placeholder="Search logs..." oninput="filterLogs()">
                         <span class="log-count">Logs: <span id="logCount">0</span></span>
                     </div>
@@ -522,6 +524,18 @@
     </div>
 
     <script>
+        // Injected server-side: the run's id (or null if the run hasn't been
+        // created yet - reassigned later via the 'run_id' SSE event below if
+        // it arrives after this page has already loaded, which is the
+        // common case: the viewer URL is printed and typically opened
+        // before the run is created) and the backend's configured hostname
+        // (or null if it was just "localhost"/"127.0.0.1" - meaningless
+        // once embedded in a page that may be opened from a different
+        // device than the one running the CLI). See updateDownloadLink()
+        // below.
+        let RUN_ID = {run_id};
+        const BACKEND_HOST = {backend_host};
+
         let eventSource = null;
         let autoScroll = true;
         let logCount = 0;
@@ -532,29 +546,59 @@
         let lastBatchTime = 0;
         const BATCH_INTERVAL_MS = 50;
         const MAX_BATCH_SIZE = 50;
+        // Cap how many entries are kept in memory/rendered in the DOM. A run
+        // can produce hundreds of thousands of log lines; retaining all of
+        // them here (unbounded array + DOM nodes) is what freezes the tab.
+        // The "Download Logs" button reads the full file from disk directly,
+        // so it isn't affected by this cap.
+        const MAX_RETAINED_LOGS = 5000;
+        // Cap how many entries can sit unrendered in the queue. Without this,
+        // a fast burst (hundreds of thousands of lines) makes the queue grow
+        // into a multi-minute backlog: the view still eventually renders
+        // everything, but always shows stale content instead of what's
+        // happening right now - as bad as no live view at all for reading
+        // context to answer a prompt while a test is running. Dropping the
+        // oldest not-yet-rendered entries keeps the view within a few
+        // seconds of real-time, at the cost of some historical entries never
+        // appearing live during a burst (the full history is unaffected -
+        // it's still in the saved log file and via "Download Logs").
+        const MAX_PENDING_QUEUE = 2000;
 
         function connectToLogStream() {{
             const statusIndicator = document.getElementById('statusIndicator');
             const statusText = document.getElementById('statusText');
-            
+
             eventSource = new EventSource('/api/logs/stream');
-            
+
             eventSource.addEventListener('connected', function(e) {{
                 console.log('Connected to log stream');
                 statusIndicator.className = 'status-indicator connected';
                 statusText.textContent = 'Connected';
             }});
-            
+
+            eventSource.addEventListener('run_id', function(e) {{
+                // Arrives once the run is created, which is typically after
+                // this page has already loaded (RUN_ID was baked in as null
+                // at that point) - re-resolve the download link now that
+                // it's known, without requiring a page reload.
+                const data = JSON.parse(e.data);
+                RUN_ID = data.run_id;
+                updateDownloadLink();
+            }});
+
             eventSource.addEventListener('log', function(e) {{
                 const logEntry = JSON.parse(e.data);
                 logBatchQueue.push(logEntry);
+                if (logBatchQueue.length > MAX_PENDING_QUEUE) {{
+                    logBatchQueue.splice(0, logBatchQueue.length - MAX_PENDING_QUEUE);
+                }}
                 scheduleBatchProcessing();
             }});
-            
+
             eventSource.addEventListener('keepalive', function(e) {{
                 // Silent keepalive
             }});
-            
+
             eventSource.addEventListener('end', function(e) {{
                 console.log('Log stream ended');
                 statusIndicator.className = 'status-indicator disconnected';
@@ -616,17 +660,30 @@
             const fragment = document.createDocumentFragment();
             const batchSize = Math.min(logBatchQueue.length, MAX_BATCH_SIZE);
             const batch = logBatchQueue.splice(0, batchSize);
-            
+
             batch.forEach(entry => {{
                 allLogs.push(entry);
                 logCount++;
-                
+
                 const logDiv = createLogElement(entry);
                 fragment.appendChild(logDiv);
             }});
-            
+
             container.appendChild(fragment);
-            
+
+            // Trim oldest entries once past the retention cap, so memory and
+            // DOM size stay bounded regardless of total run volume.
+            if (allLogs.length > MAX_RETAINED_LOGS) {{
+                const excess = allLogs.length - MAX_RETAINED_LOGS;
+                allLogs.splice(0, excess);
+                for (let i = 0; i < excess; i++) {{
+                    const oldest = container.firstElementChild;
+                    if (oldest) {{
+                        container.removeChild(oldest);
+                    }}
+                }}
+            }}
+
             document.getElementById('logCount').textContent = logCount;
             
             if (autoScroll) {{
@@ -727,69 +784,6 @@
             }});
         }}
 
-        function downloadLogs() {{
-            // First try server download (if server is still running)
-            fetch('/download_logs')
-                .then(response => {{
-                    if (response.ok) {{
-                        return response.blob();
-                    }}
-                    throw new Error('Server not available');
-                }})
-                .then(blob => {{
-                    // Server download successful
-                    const url = window.URL.createObjectURL(blob);
-                    const a = document.createElement('a');
-                    a.href = url;
-                    a.download = 'test_run_logs.log';
-                    document.body.appendChild(a);
-                    a.click();
-                    window.URL.revokeObjectURL(url);
-                    document.body.removeChild(a);
-                }})
-                .catch(error => {{
-                    // Server not available, use client-side download
-                    console.log('Server unavailable, using client-side download');
-                    downloadLogsClientSide();
-                }});
-        }}
-
-        function downloadLogsClientSide() {{
-            // Build log content from stored logs
-            let logContent = '';
-            allLogs.forEach(entry => {{
-                const timestamp = entry.timestamp || new Date().toISOString();
-                const level = entry.level || 'INFO';
-                const message = entry.message || '';
-                // Replace literal \n strings with actual newlines, then add entry
-                const cleanMessage = message.replace(/\\n/g, '\n');
-                logContent += `${{timestamp}} [${{level}}] ${{cleanMessage}}\n`;
-            }});
-
-            if (logContent === '') {{
-                alert('No logs to download');
-                return;
-            }}
-
-            // Create blob and download
-            const blob = new Blob([logContent], {{ type: 'text/plain;charset=utf-8' }});
-            const url = window.URL.createObjectURL(blob);
-            const a = document.createElement('a');
-            a.href = url;
-            
-            // Generate filename with current timestamp
-            const now = new Date();
-            const dateStr = now.toISOString().slice(0,19).replace(/[T:]/g, '-');
-            a.download = `test_run_logs_${{dateStr}}.log`;
-            
-            document.body.appendChild(a);
-            a.click();
-            window.URL.revokeObjectURL(url);
-            document.body.removeChild(a);
-            
-            console.log(`Downloaded ${{allLogs.length}} log entries client-side`);
-        }}
-
         function escapeHtml(text) {{
             const div = document.createElement('div');
             div.textContent = text;
@@ -813,7 +807,29 @@
             }}
         }});
 
+        function updateDownloadLink() {{
+            if (RUN_ID === null) {{
+                return;
+            }}
+            // Prefer the backend's own configured hostname when it's a real
+            // address; otherwise fall back to whatever host this page was
+            // actually loaded from (correct whenever the CLI and backend
+            // run on the same machine, which is the common case - see the
+            // comment on BACKEND_HOST above).
+            const host = BACKEND_HOST || window.location.hostname;
+            const btn = document.getElementById('downloadBtn');
+            // download=true tells the backend to send Content-Disposition:
+            // attachment, so the browser downloads the file directly
+            // instead of just rendering the raw log text in the tab.
+            btn.href = `${{window.location.protocol}}//${{host}}/api/v1/test_run_executions/${{RUN_ID}}/log?download=true`;
+            btn.removeAttribute('onclick');
+            btn.title = '';
+            btn.style.opacity = '';
+            btn.style.cursor = '';
+        }}
+
         window.addEventListener('load', function() {{
+            updateDownloadLink();
             connectToLogStream();
         }});
 
diff --git a/th_cli/test_run/logging.py b/th_cli/test_run/logging.py
index 05fdcd6..4d7bdbc 100644
--- a/th_cli/test_run/logging.py
+++ b/th_cli/test_run/logging.py
@@ -62,7 +62,7 @@ def configure_logger_for_run(title: str, enable_log_streaming: bool = False) ->
             from th_cli.test_run.log_stream_handler import LogStreamHandler
 
             _log_stream_handler = LogStreamHandler(port=8998)
-            viewer_url = _log_stream_handler.start(test_run_title=title, log_file_path=log_path)
+            viewer_url = _log_stream_handler.start(test_run_title=title)
             # Add custom sink that forwards logs to the stream handler
             def stream_sink(message):
                 """Custom sink that forwards logs to the HTTP stream."""
@@ -91,7 +91,7 @@ def stream_sink(message):
 def stop_log_streaming():
     """Stop the log streaming server if it's running."""
     global _log_stream_handler
-    
+
     if _log_stream_handler:
         try:
             _log_stream_handler.stop()
@@ -102,6 +102,15 @@ def stop_log_streaming():
             _log_stream_handler = None
 
 
+def set_download_run_id(run_id: int) -> None:
+    """Tell the log viewer which run's log to link the "Download Logs" button
+    to, once the run has been created and its id is known (the log-streaming
+    server starts before the run exists, so this can't be known up front).
+    """
+    if _log_stream_handler:
+        _log_stream_handler.set_run_id(run_id)
+
+
 def get_log_stream_url() -> Optional[str]:
     """Get the URL for the log viewer if streaming is enabled.
     
diff --git a/th_cli/test_run/logs_http_server.py b/th_cli/test_run/logs_http_server.py
index 3a14027..da52907 100644
--- a/th_cli/test_run/logs_http_server.py
+++ b/th_cli/test_run/logs_http_server.py
@@ -24,10 +24,11 @@
 
 from loguru import logger
 
+from th_cli.config import config
+
 # HTTP Endpoints
 ENDPOINT_ROOT = "/"
 ENDPOINT_LOGS_STREAM = "/api/logs/stream"
-ENDPOINT_DOWNLOAD_LOGS = "/download_logs"
 
 
 class LogStreamingHandler(BaseHTTPRequestHandler):
@@ -39,51 +40,10 @@ def do_GET(self):
             self.serve_log_viewer()
         elif self.path == ENDPOINT_LOGS_STREAM:
             self.stream_logs()
-        elif self.path == ENDPOINT_DOWNLOAD_LOGS:
-            self.download_logs()
         else:
             logger.warning(f"404 for GET {self.path}")
             self.send_error(404)
     
-    def download_logs(self):
-        """Serve the log file for download using chunked streaming."""
-        log_file_path = getattr(self.server, "log_file_path", None)
-        
-        if not log_file_path or not Path(log_file_path).exists():
-            self.send_error(404, "Log file not found")
-            return
-        
-        try:
-            file_path = Path(log_file_path)
-            filename = file_path.name
-            file_size = file_path.stat().st_size
-            
-            # Send headers
-            self.send_response(200)
-            self.send_header("Content-Type", "text/plain; charset=utf-8")
-            self.send_header("Content-Disposition", f'attachment; filename="{filename}"')
-            self.send_header("Content-Length", str(file_size))
-            self.send_header("Access-Control-Allow-Origin", "*")
-            self.end_headers()
-            
-            # Stream file in chunks to avoid loading entire file into memory
-            CHUNK_SIZE = 65536  # 64KB chunks
-            bytes_sent = 0
-
-            with open(log_file_path, 'rb') as f:
-                while chunk := f.read(CHUNK_SIZE):
-                    self.wfile.write(chunk)
-                    bytes_sent += len(chunk)
-
-            logger.info(f"Log file downloaded: {filename} ({bytes_sent} bytes)")
-            
-        except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError):
-            # Client disconnected during download - normal, not an error
-            logger.debug("Client disconnected during log file download")
-        except Exception as e:
-            logger.error(f"Error serving log file: {e}")
-            # Note: Can't send error response because headers were already sent
-
     def stream_logs(self):
         """Stream logs using Server-Sent Events (SSE)."""
         logger.info("Client connected for log stream")
@@ -116,20 +76,30 @@ def stream_logs(self):
                 try:
                     # Get log entry from queue with timeout
                     log_entry = log_queue.get(timeout=1.0)
-                    
+
                     if log_entry is None:  # Signal to stop
                         logger.info("Received end-of-stream signal for logs")
                         self._send_sse_event("end", {"message": "Log stream ended"})
                         break
 
+                    if isinstance(log_entry, dict) and log_entry.get("__event__") == "run_id":
+                        # Control message, not a log line - lets an already-
+                        # connected viewer pick up the run id once it's
+                        # known, instead of only a fresh page load.
+                        if not self._send_sse_event("run_id", {"run_id": log_entry["run_id"]}):
+                            logger.debug("Client disconnected while streaming")
+                            client_disconnected = True
+                            break
+                        continue
+
                     # Send log entry as SSE event
                     if not self._send_sse_event("log", log_entry):
                         logger.debug("Client disconnected while streaming")
                         client_disconnected = True
                         break
-                        
+
                     sent_count += 1
-                    
+
                     if sent_count % 100 == 0:
                         logger.debug(f"Sent {sent_count} log entries to client")
 
@@ -173,7 +143,22 @@ def serve_log_viewer(self):
         """Serve the log viewer HTML page."""
         # Get configuration from server
         test_run_title = getattr(self.server, "test_run_title", "Test Execution")
-        
+        run_id = getattr(self.server, "run_id", None)
+
+        # config.hostname is only meaningful to the download link when it's a
+        # real, routable address. This page can be opened from a different
+        # device than the one running the CLI (that's why this server binds
+        # to the LAN IP in the first place) - if config.hostname is just
+        # "localhost" (the common case when the CLI and backend run on the
+        # same machine as each other), embedding it here would tell a remote
+        # browser to download from *itself*. Leave it unset in that case so
+        # the page falls back to whatever host the browser actually used to
+        # reach it (correct whenever the CLI and backend share a machine,
+        # which is the common case); keep it when it's a real configured
+        # address (correct when the CLI talks to a genuinely separate
+        # backend host).
+        backend_host = None if config.hostname in ("localhost", "127.0.0.1") else config.hostname
+
         # Read HTML template from file
         try:
             template_path = Path(__file__).parent / "log_viewer.html"
@@ -182,7 +167,9 @@ def serve_log_viewer(self):
 
             # Replace placeholders
             html_content = html_template.format(
-                test_run_title=html.escape(test_run_title)
+                test_run_title=html.escape(test_run_title),
+                run_id=json.dumps(run_id),
+                backend_host=json.dumps(backend_host),
             )
         except Exception as e:
             logger.error(f"Failed to load HTML template: {e}")
@@ -231,15 +218,13 @@ def start(
         log_queue: queue.Queue,
         test_run_title: str = "Test Execution",
         local_ip: Optional[str] = None,
-        log_file_path: Optional[str] = None,
     ):
         """Start HTTP server for log streaming.
-        
+
         Args:
             log_queue: Queue containing log entries to stream
             test_run_title: Title of the test run for display
             local_ip: Local IP address for display (defaults to localhost)
-            log_file_path: Path to log file for download functionality
         """
         try:
             # Use ThreadingHTTPServer for better concurrency
@@ -250,7 +235,7 @@ def start(
             self.server.log_queue = log_queue
             self.server.test_run_title = test_run_title
             self.server.local_ip = local_ip or "localhost"
-            self.server.log_file_path = log_file_path
+            self.server.run_id = None
 
             logger.info(f"Logs HTTP server configured for test run: {test_run_title}")
 
@@ -269,6 +254,13 @@ def run_server():
             logger.error(f"Failed to start logs HTTP server: {e}")
             raise
 
+    def set_run_id(self, run_id: int) -> None:
+        """Set the run id the "Download Logs" link should point to, once the
+        run has been created (it doesn't exist yet when the server starts).
+        """
+        if self.server is not None:
+            self.server.run_id = run_id
+
     def stop(self):
         """Stop HTTP server."""
         if self.server:
diff --git a/th_cli/test_run/websocket.py b/th_cli/test_run/websocket.py
index e60da91..e946814 100644
--- a/th_cli/test_run/websocket.py
+++ b/th_cli/test_run/websocket.py
@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+import asyncio
+
 import click
 import websockets
 from loguru import logger
@@ -37,7 +39,7 @@
     colorize_state,
 )
 from th_cli.config import config
-from th_cli.shared_constants import MessageTypeEnum
+from th_cli.shared_constants import MessageTypeEnum, TestStateEnum
 
 from .prompt_manager import handle_file_upload_request, handle_prompt
 from .socket_schemas import (
@@ -56,6 +58,25 @@
 
 WEBSOCKET_MAX_MESSAGE_SIZE = 32 * 1024 * 1024  # 32MB
 
+# After the test run reaches a terminal state, the backend may still have a
+# trailing batch of log records queued/in-flight (it flushes and broadcasts
+# any pending log entries *after* sending the terminal state update - see
+# TestLogHandler.finish()/TestUIObserver.complete_tasks() on the backend).
+# Keep draining for a short grace period instead of closing immediately, so
+# that trailing batch isn't dropped by a socket we already hung up on.
+DRAIN_TIMEOUT_S = 5.0
+
+# Yield to the event loop every N log records while processing one batch, so
+# a very large batch doesn't block the websocket read loop for its entire
+# duration.
+LOG_RECORD_YIELD_INTERVAL = 200
+
+# TestRun states that are not yet finished, mirroring the backend's own
+# TestRun.completed() contract (state not in [PENDING, EXECUTING]). Anything
+# else is terminal - checked explicitly rather than negating "executing" so
+# this can't misclassify a non-terminal state (e.g. PENDING) as terminal.
+NON_TERMINAL_RUN_STATES = (TestStateEnum.PENDING, TestStateEnum.EXECUTING)
+
 
 class TestRunSocket:
     def __init__(
@@ -68,6 +89,7 @@ def __init__(
         self.project_config_dict = project_config_dict or {}
         self.two_way_talk_handler = two_way_talk_handler
         self._chip_server_info_displayed = False
+        self._run_finished = False
         # Track test step errors for logging
         # Key: (suite_index, case_index), Value: list of error strings from all steps
         self.test_case_step_errors: dict[tuple[int, int], list[str]] = {}
@@ -85,9 +107,19 @@ async def connect_websocket(self) -> None:
                 try:
                     while True:
                         try:
-                            message = await socket.recv()
+                            if self._run_finished:
+                                # Drain any trailing messages for a short grace
+                                # period instead of closing the instant the
+                                # terminal state update arrives.
+                                message = await asyncio.wait_for(socket.recv(), timeout=DRAIN_TIMEOUT_S)
+                            else:
+                                message = await socket.recv()
                         except websockets.exceptions.ConnectionClosedOK:
                             break
+                        except asyncio.TimeoutError:
+                            # No more trailing messages arrived during the
+                            # drain grace period - safe to close now.
+                            break
 
                         # skip messages that are bytes, as we're expecting a string.\
                         if not isinstance(message, str):
@@ -103,7 +135,13 @@ async def connect_websocket(self) -> None:
                             click.echo(colorize_error(f"Received invalid socket message: {message}"), err=True)
                             click.echo(colorize_error(e.json()), err=True)
                 finally:
-                    pass  # Cleanup if needed
+                    if self._run_finished:
+                        try:
+                            await socket.close()
+                        except websockets.exceptions.ConnectionClosedError:
+                            # Backend closed connection without completing handshake
+                            # This is acceptable as test run completed successfully
+                            pass
         except websockets.exceptions.ConnectionClosed:
             # Handle case where backend doesn't complete close handshake properly
             # This can happen with long-running test executions
@@ -112,7 +150,7 @@ async def connect_websocket(self) -> None:
 
     async def __handle_incoming_socket_message(self, socket: WebSocketClientProtocol, message: SocketMessage) -> None:
         if isinstance(message.payload, TestUpdate):
-            await self.__handle_test_update(socket=socket, update=message.payload)
+            await self.__handle_test_update(update=message.payload)
         elif isinstance(message.payload, PromptRequest):
             # Debug: log the message type
             logger.debug(f"Received prompt with type: {message.type}")
@@ -129,7 +167,7 @@ async def __handle_incoming_socket_message(self, socket: WebSocketClientProtocol
                     two_way_talk_handler=self.two_way_talk_handler,
                 )
         elif message.type == MessageTypeEnum.TEST_LOG_RECORDS and isinstance(message.payload, list):
-            self.__handle_log_record(message.payload)
+            await self.__handle_log_record(message.payload)
         elif isinstance(message.payload, TimeOutNotification):
             # ignore time_out_notification as we handle timeout our selves
             pass
@@ -139,7 +177,7 @@ async def __handle_incoming_socket_message(self, socket: WebSocketClientProtocol
                 err=True,
             )
 
-    async def __handle_test_update(self, socket: WebSocketClientProtocol, update: TestUpdate) -> None:
+    async def __handle_test_update(self, update: TestUpdate) -> None:
         if isinstance(update.body, TestStepUpdate):
             self.__log_test_step_update(update.body)
         elif isinstance(update.body, TestCaseUpdate):
@@ -148,14 +186,13 @@ async def __handle_test_update(self, socket: WebSocketClientProtocol, update: Te
             self.__log_test_suite_update(update.body)
         elif isinstance(update.body, TestRunUpdate):
             await self.__log_test_run_update(update.body)
-            if update.body.state != "executing":
-                # Test run ended disconnect.
-                try:
-                    await socket.close()
-                except websockets.exceptions.ConnectionClosedError:
-                    # Backend closed connection without completing handshake
-                    # This is acceptable as test run completed successfully
-                    pass
+            if update.body.state not in NON_TERMINAL_RUN_STATES:
+                # Test run ended. Don't close immediately - the backend may
+                # still be flushing/broadcasting a trailing batch of log
+                # entries after this message; let the read loop keep
+                # draining for a short grace period (see DRAIN_TIMEOUT_S)
+                # before actually closing.
+                self._run_finished = True
 
     async def __log_test_run_update(self, update: TestRunUpdate) -> None:
         # Display CHIP server info when test run starts executing (SDK container already running)
@@ -284,9 +321,16 @@ def __log_test_step_update(self, update: TestStepUpdate) -> None:
             self.test_case_step_errors.setdefault(case_key, []).extend(update.errors)
             logger.debug(f"Tracked {len(update.errors)} error(s) for test case {case_key}: {update.errors}")
 
-    def __handle_log_record(self, records: list[TestLogRecord]) -> None:
-        for record in records:
+    async def __handle_log_record(self, records: list[TestLogRecord]) -> None:
+        # Batches can contain tens of thousands of entries after a large test
+        # case run. Yield periodically instead of logging the whole batch in
+        # one uninterrupted stretch, so the websocket read loop (and any
+        # other pending work, e.g. prompt handling) doesn't stall for the
+        # entire duration of processing one message.
+        for i, record in enumerate(records):
             logger.log(record.level, record.message)
+            if (i + 1) % LOG_RECORD_YIELD_INTERVAL == 0:
+                await asyncio.sleep(0)
 
     def __suite(self, index: int) -> TestSuiteExecution:
         return self.run.test_suite_executions[index]