diff --git a/src/skillspector/cleanup.py b/src/skillspector/cleanup.py new file mode 100644 index 00000000..ded8f994 --- /dev/null +++ b/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) diff --git a/src/skillspector/cli.py b/src/skillspector/cli.py index deae40a6..4beb9c23 100644 --- a/src/skillspector/cli.py +++ b/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__": diff --git a/src/skillspector/constants.py b/src/skillspector/constants.py index 375992c7..c7e0d5e6 100644 --- a/src/skillspector/constants.py +++ b/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//provider.py # for ``DEFAULT_MODEL`` and ``SLOT_DEFAULTS``). The active provider's diff --git a/src/skillspector/mcp_server.py b/src/skillspector/mcp_server.py index 444b75fc..ef77ddd1 100644 --- a/src/skillspector/mcp_server.py +++ b/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: