diff --git a/th_cli/test_run/logging.py b/th_cli/test_run/logging.py
index bf3363b..becc7f5 100644
--- a/th_cli/test_run/logging.py
+++ b/th_cli/test_run/logging.py
@@ -62,8 +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."""
@@ -101,6 +100,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 23ab7f2..594108d 100644
--- a/th_cli/test_run/logs_http_server.py
+++ b/th_cli/test_run/logs_http_server.py
@@ -26,10 +26,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"
ENDPOINT_STATUS = "/api/status"
@@ -42,8 +43,6 @@ 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()
elif self.path == ENDPOINT_STATUS:
self.serve_status()
else:
@@ -65,41 +64,6 @@ def serve_status(self):
self.end_headers()
self.wfile.write(body)
- 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
-
- 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()
-
- CHUNK_SIZE = 65536 # 64 KB
- 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}")
-
def stream_logs(self):
"""Stream logs (and tree events) using Server-Sent Events (SSE)."""
logger.info("Client connected for log stream")
@@ -183,6 +147,13 @@ def stream_logs(self):
if not self._send_sse_event("tree_update", entry):
client_disconnected = True
+ elif entry_type == "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": entry.get("run_id")}):
+ client_disconnected = True
+
else:
# Regular log entry (no "type" key, or type=="log").
if not self._send_sse_event("log", entry):
@@ -222,14 +193,34 @@ 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"
with open(template_path, "r", encoding="utf-8") as f:
html_template = f.read()
- html_content = html_template.format(test_run_title=html.escape(test_run_title))
+ # Replace placeholders
+ html_content = html_template.format(
+ 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}")
html_content = f"""
@@ -271,7 +262,6 @@ def start(
tree_state: dict,
test_run_title: str = "Test Execution",
local_ip: Optional[str] = None,
- log_file_path: Optional[str] = None,
tree_lock: Optional[threading.Lock] = None,
):
"""Start the HTTP server for log streaming.
@@ -283,7 +273,6 @@ def start(
so clients connecting after init_tree() receive a current snapshot.
test_run_title: Title shown in the browser UI.
local_ip: LAN IP used for display purposes.
- log_file_path: Path to the on-disk log file for download.
tree_lock: Lock protecting tree_state reads/writes.
"""
try:
@@ -296,8 +285,8 @@ def start(
self.server.tree_lock = tree_lock
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.start_time = datetime.datetime.now().isoformat()
+ self.server.run_id = None
logger.info(f"Logs HTTP server configured for test run: {test_run_title}")
@@ -316,6 +305,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 the HTTP server."""
if self.server:
diff --git a/th_cli/test_run/websocket.py b/th_cli/test_run/websocket.py
index a368ce6..887c9be 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 .logging import get_log_stream_handler
from .prompt_manager import handle_file_upload_request, handle_prompt
@@ -57,6 +59,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__(
@@ -69,6 +90,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]] = {}
@@ -86,9 +108,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):
@@ -104,7 +136,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
@@ -113,7 +151,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}")
@@ -130,7 +168,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
@@ -140,7 +178,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):
@@ -149,14 +187,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)
@@ -316,9 +353,16 @@ def __log_test_step_update(self, update: TestStepUpdate) -> None:
step_idx=update.test_step_execution_index,
)
- 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]