refactor: centralize cleanup and risk threshold

src/skillspector/cleanup.py

@@ -0,0 +1,13 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Shared cleanup helpers for SkillSpector."""
+
+import shutil
+
+
+def cleanup_result(result: dict[str, object]) -> None:
+    """Remove temp dir from graph result if set."""
+    temp_dir = result.get("temp_dir_for_cleanup")
+    if temp_dir and isinstance(temp_dir, str):
+        shutil.rmtree(temp_dir, ignore_errors=True)

src/skillspector/cli.py

@@ -23,7 +23,6 @@
 
 import json
 import os
-import shutil
 import sys
 from enum import StrEnum
 from pathlib import Path
@@ -34,6 +33,8 @@
 from rich.console import Console
 
 from skillspector import __version__
+from skillspector.cleanup import cleanup_result
+from skillspector.constants import RISK_THRESHOLD
 from skillspector.graph import graph
 from skillspector.logging_config import get_logger, set_level
 from skillspector.multi_skill import MultiSkillDetectionResult, detect_skills
@@ -161,13 +162,6 @@ def _write_result(
             print(report_body)
 
 
-def _cleanup_result(result: dict[str, object]) -> None:
-    """Remove temp dir from graph result if set."""
-    temp_dir = result.get("temp_dir_for_cleanup")
-    if temp_dir and isinstance(temp_dir, str):
-        shutil.rmtree(temp_dir, ignore_errors=True)
-
-
 @app.command()
 def scan(
     input_path: Annotated[
@@ -317,7 +311,7 @@ def scan(
 
         _write_result(result, output, format)
 
-        if (result.get("risk_score") or 0) > 50:
+        if (result.get("risk_score") or 0) > RISK_THRESHOLD:
             raise typer.Exit(code=1)
     except typer.Exit:
         raise
@@ -332,7 +326,7 @@ def scan(
         raise typer.Exit(code=2) from e
     finally:
         if result is not None:
-            _cleanup_result(result)
+            cleanup_result(result)
 
 
 def _build_trace_config(input_path: str, format: FormatChoice, no_llm: bool) -> RunnableConfig:
@@ -433,7 +427,7 @@ def _scan_multi_skill(
             if "error" not in result:
                 _write_result(result, None, format)
 
-    if max_score > 50:
+    if max_score > RISK_THRESHOLD:
         raise typer.Exit(code=1)
 
 
@@ -555,7 +549,7 @@ def baseline(
         raise typer.Exit(code=2) from e
     finally:
         if result is not None:
-            _cleanup_result(result)
+            cleanup_result(result)
 
 
 if __name__ == "__main__":

src/skillspector/constants.py

@@ -26,6 +26,8 @@
 MAX_INPUT_TOKENS_PCT = 0.75
 # Fallback context length when no metadata API or registry entry is available.
 DEFAULT_CONTEXT_LENGTH = 128_000
+# Risk score threshold above which a scan is treated as unsafe.
+RISK_THRESHOLD = 50
 
 # Default-model selection lives on each provider (see providers/<name>/provider.py
 # for ``DEFAULT_MODEL`` and ``SLOT_DEFAULTS``).  The active provider's

src/skillspector/mcp_server.py

@@ -27,10 +27,11 @@
 
 from __future__ import annotations
 
-import shutil
 from typing import TYPE_CHECKING, Any
 
 from skillspector import __version__
+from skillspector.cleanup import cleanup_result
+from skillspector.constants import RISK_THRESHOLD
 from skillspector.graph import graph
 from skillspector.logging_config import get_logger
 from skillspector.providers import resolve_provider_credentials
@@ -42,9 +43,6 @@
 
 VALID_FORMATS = ("json", "markdown", "sarif", "terminal")
 
-# Mirrors the CLI: a scan scoring above this is treated as unsafe (CLI exits 1).
-RISK_THRESHOLD = 50
-
 
 async def run_scan(
     target: str,
@@ -126,9 +124,7 @@ async def run_scan(
         }
     finally:
         if result is not None:
-            temp_dir = result.get("temp_dir_for_cleanup")
-            if temp_dir and isinstance(temp_dir, str):
-                shutil.rmtree(temp_dir, ignore_errors=True)
+            cleanup_result(result)
 
 
 def build_server(name: str = "skillspector") -> FastMCP: