refactor: address PR review on streaming logs

- collapse upload_logs/decline_logs config fields into a single
  Optional[bool] (tri-state); projection to share_logs/decline_logs
  happens at the setup_streaming call site.
- streaming.py: replace module-level globals + atexit teardown with a
  StreamingLogs context manager. set_run_status disappears entirely —
  __exit__ infers the run status from the exception that closed the
  with block. set_report_run_id is now an instance method. Logger
  handler wiring iterates over (cli_logger, sdk_logger) instead of
  repeating itself.
- log_uploader.py: drop _LEVEL_MAP, use logging.getLevelName directly.
  Wire format changes WARN/ERROR-for-CRITICAL to WARNING/CRITICAL.
- log_uploader.py: tidy BatchedLogUploader.stop so the final _flush
  always runs and the thread shutdown only runs when there is a thread.
- socketcli.py: wrap main_code body in 'with setup_streaming(...) as
  streaming:'; cli()'s exception handlers no longer need to set status
  before re-raising.
- tests updated for the CM API and the new level strings.

Author: barslev

socketsecurity/config.py

@@ -139,8 +139,9 @@ class CliConfig:
     ignore_commit_files: bool = False
     disable_blocking: bool = False
     disable_ignore: bool = False
-    upload_logs: bool = False
-    decline_logs: bool = False
+    # Tri-state log-upload preference: True = --upload-logs, False = --no-upload-logs,
+    # None = neither (server-side override decides).
+    upload_logs: Optional[bool] = None
     strict_blocking: bool = False
     integration_type: IntegrationType = "api"
     integration_org_slug: Optional[str] = None
@@ -216,6 +217,9 @@ def from_args(cls, args_list: Optional[List[str]] = None) -> 'CliConfig':
 
         if args.upload_logs and args.decline_logs:
             parser.error("--upload-logs and --no-upload-logs are mutually exclusive")
+        upload_logs: Optional[bool] = (
+            True if args.upload_logs else False if args.decline_logs else None
+        )
 
         if args.reach_exclude_paths:
             logging.warning(
@@ -287,8 +291,7 @@ def from_args(cls, args_list: Optional[List[str]] = None) -> 'CliConfig':
             'ignore_commit_files': args.ignore_commit_files,
             'disable_blocking': args.disable_blocking,
             'disable_ignore': args.disable_ignore,
-            'upload_logs': args.upload_logs,
-            'decline_logs': args.decline_logs,
+            'upload_logs': upload_logs,
             'strict_blocking': args.strict_blocking,
             'integration_type': args.integration,
             'pending_head': args.pending_head,

socketsecurity/core/log_uploader.py

@@ -26,14 +26,6 @@
 
 _FLUSH_GUARD = threading.local()
 
-_LEVEL_MAP = {
-    logging.DEBUG: "DEBUG",
-    logging.INFO: "INFO",
-    logging.WARNING: "WARN",
-    logging.ERROR: "ERROR",
-    logging.CRITICAL: "ERROR",
-}
-
 
 def _now_str() -> str:
     return datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
@@ -69,12 +61,10 @@ def start(self) -> None:
         self._thread.start()
 
     def stop(self, timeout: float = 2.0) -> None:
-        if self._thread is None:
-            self._flush()
-            return
-        self._stop.set()
-        self._thread.join(timeout=timeout)
-        self._thread = None
+        if self._thread is not None:
+            self._stop.set()
+            self._thread.join(timeout=timeout)
+            self._thread = None
         self._flush()
 
     def _run(self) -> None:
@@ -114,7 +104,7 @@ def emit(self, record: logging.LogRecord) -> None:
         try:
             self._uploader.add({
                 "timestamp": _now_str(),
-                "level": _LEVEL_MAP.get(record.levelno, "INFO"),
+                "level": logging.getLevelName(record.levelno),
                 "message": self.format(record),
                 "context": self._context,
             })

socketsecurity/core/streaming.py

@@ -1,33 +1,118 @@
-"""Wire the server log streaming pipeline for one CLI run.
-
-`setup_streaming` registers the run with the backend, attaches handlers that
-route the CLI's own log output through both the local terminal and a batched
-uploader, and forces the loggers into DEBUG so the upload captures everything
-regardless of local terminal verbosity.
-
-Returns a teardown callable to invoke on exit (typically via `atexit.register`).
-Returns None if registration failed; in that case nothing was wired up.
+"""Server log streaming pipeline for one CLI run.
+
+`setup_streaming` returns a `StreamingLogs` context manager. On enter it
+registers the run with the backend, attaches handlers that route the CLI's
+own log output through both the local terminal and a batched uploader, and
+forces the loggers into DEBUG so the upload captures everything regardless
+of local terminal verbosity. On exit it tears the handlers back down and
+finalizes the run; the status sent to finalize is inferred from the
+exception that closed the `with` block (success / failure / cancelled).
+
+If registration fails the manager becomes a no-op — nothing is wired up
+and __exit__ does nothing.
 """
 
 import logging
-from typing import Callable, Optional
+from typing import Optional
 
 from .cli_client import CliClient
 from .cli_run import finalize_cli_run, register_cli_run
 from .log_uploader import BatchedLogUploader, UploadingLogHandler
 
-_run_status: str = "success"
-_report_run_id: Optional[str] = None
-
-
-def set_run_status(status: str) -> None:
-    global _run_status
-    _run_status = status
-
 
-def set_report_run_id(report_run_id: Optional[str]) -> None:
-    global _report_run_id
-    _report_run_id = report_run_id
+class StreamingLogs:
+    def __init__(
+        self,
+        *,
+        client: CliClient,
+        cli_logger: logging.Logger,
+        sdk_logger: logging.Logger,
+        client_version: str,
+        share_logs: bool,
+        decline_logs: bool,
+        enable_debug: bool,
+    ):
+        self._client = client
+        self._loggers = (cli_logger, sdk_logger)
+        self._client_version = client_version
+        self._share_logs = share_logs
+        self._decline_logs = decline_logs
+        self._enable_debug = enable_debug
+
+        self._run_id: Optional[str] = None
+        self._report_run_id: Optional[str] = None
+        self._uploader: Optional[BatchedLogUploader] = None
+        self._upload_handler: Optional[UploadingLogHandler] = None
+        self._terminal_handler: Optional[logging.StreamHandler] = None
+        self._saved_levels: tuple = ()
+        self._saved_propagate: tuple = ()
+
+    def set_report_run_id(self, report_run_id: Optional[str]) -> None:
+        self._report_run_id = report_run_id
+
+    def __enter__(self) -> "StreamingLogs":
+        self._run_id = register_cli_run(
+            self._client,
+            client_version=self._client_version,
+            share_logs=self._share_logs,
+            decline_logs=self._decline_logs,
+        )
+        cli_logger = self._loggers[0]
+        if not self._run_id:
+            cli_logger.debug("server log streaming not active for this run")
+            return self
+
+        self._uploader = BatchedLogUploader(self._client, self._run_id)
+        self._uploader.start()
+        self._upload_handler = UploadingLogHandler(self._uploader, context="socket-python-cli")
+        self._upload_handler.setFormatter(logging.Formatter("%(message)s"))
+
+        self._terminal_handler = logging.StreamHandler()
+        self._terminal_handler.setLevel(logging.DEBUG if self._enable_debug else logging.INFO)
+        self._terminal_handler.setFormatter(logging.Formatter("%(asctime)s: %(message)s"))
+
+        self._saved_levels = tuple(lg.level for lg in self._loggers)
+        self._saved_propagate = tuple(lg.propagate for lg in self._loggers)
+        for lg in self._loggers:
+            lg.setLevel(logging.DEBUG)
+            lg.propagate = False
+            lg.addHandler(self._terminal_handler)
+            lg.addHandler(self._upload_handler)
+
+        cli_logger.debug(f"server log streaming enabled (run_id={self._run_id})")
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
+        if self._run_id is None:
+            return False
+
+        status = self._status_for_exit(exc_type, exc_val)
+        for lg in self._loggers:
+            lg.removeHandler(self._upload_handler)
+        self._uploader.stop()
+        finalize_cli_run(
+            self._client,
+            self._run_id,
+            status=status,
+            report_run_id=self._report_run_id,
+        )
+        for lg in self._loggers:
+            lg.removeHandler(self._terminal_handler)
+        for lg, level, propagate in zip(self._loggers, self._saved_levels, self._saved_propagate):
+            lg.setLevel(level)
+            lg.propagate = propagate
+        return False
+
+    @staticmethod
+    def _status_for_exit(exc_type, exc_val) -> str:
+        if exc_type is None:
+            return "success"
+        if issubclass(exc_type, KeyboardInterrupt):
+            return "cancelled"
+        # SystemExit with code 0 / None is a clean exit; non-zero codes signal failure.
+        if issubclass(exc_type, SystemExit) and not getattr(exc_val, "code", None):
+            return "success"
+        return "failure"
 
 
 def setup_streaming(
@@ -39,49 +124,13 @@ def setup_streaming(
     share_logs: bool,
     decline_logs: bool,
     enable_debug: bool,
-) -> Optional[Callable[[], None]]:
-    run_id = register_cli_run(
-        client,
+) -> StreamingLogs:
+    return StreamingLogs(
+        client=client,
+        cli_logger=cli_logger,
+        sdk_logger=sdk_logger,
         client_version=client_version,
         share_logs=share_logs,
         decline_logs=decline_logs,
+        enable_debug=enable_debug,
     )
-    if not run_id:
-        cli_logger.debug("server log streaming not active for this run")
-        return None
-
-    log_uploader = BatchedLogUploader(client, run_id)
-    log_uploader.start()
-    upload_handler = UploadingLogHandler(log_uploader, context="socket-python-cli")
-    upload_handler.setFormatter(logging.Formatter("%(message)s"))
-
-    terminal_handler = logging.StreamHandler()
-    terminal_handler.setLevel(logging.DEBUG if enable_debug else logging.INFO)
-    terminal_handler.setFormatter(logging.Formatter("%(asctime)s: %(message)s"))
-
-    saved_levels = (cli_logger.level, sdk_logger.level)
-    saved_propagate = (cli_logger.propagate, sdk_logger.propagate)
-    cli_logger.setLevel(logging.DEBUG)
-    sdk_logger.setLevel(logging.DEBUG)
-    cli_logger.propagate = False
-    sdk_logger.propagate = False
-    cli_logger.addHandler(terminal_handler)
-    sdk_logger.addHandler(terminal_handler)
-    cli_logger.addHandler(upload_handler)
-    sdk_logger.addHandler(upload_handler)
-
-    cli_logger.debug(f"server log streaming enabled (run_id={run_id})")
-
-    def teardown() -> None:
-        cli_logger.removeHandler(upload_handler)
-        sdk_logger.removeHandler(upload_handler)
-        log_uploader.stop()
-        finalize_cli_run(client, run_id, status=_run_status, report_run_id=_report_run_id)
-        cli_logger.removeHandler(terminal_handler)
-        sdk_logger.removeHandler(terminal_handler)
-        cli_logger.setLevel(saved_levels[0])
-        sdk_logger.setLevel(saved_levels[1])
-        cli_logger.propagate = saved_propagate[0]
-        sdk_logger.propagate = saved_propagate[1]
-
-    return teardown