fix(cleanup): remove the temp directory of a scan that stops early

src/skillspector/cleanup.py

@@ -6,8 +6,12 @@
 import os
 import shutil
 import stat
-from collections.abc import Callable
+from collections.abc import Callable, Iterator, Mapping
+from contextlib import contextmanager
 from pathlib import Path
+from typing import Any
+
+from langchain_core.callbacks import BaseCallbackHandler
 
 from skillspector.python_ast import clear_python_ast_cache
 
@@ -43,6 +47,66 @@ def remove_temp_tree(path: str | Path) -> None:
     shutil.rmtree(path, onexc=_retry_writable)
 
 
+SCAN_TEMP_DIR_PREFIX = "skillspector_"
+
+
+def _is_scan_temp_dir(value: object) -> bool:
+    """Return whether ``value`` names an existing, non-symlink scan temp directory."""
+    if not isinstance(value, str) or not value:
+        return False
+    path = Path(value)
+    if not path.name.startswith(SCAN_TEMP_DIR_PREFIX):
+        return False
+    try:
+        return path.is_dir() and not path.is_symlink() and not os.path.isjunction(path)
+    except OSError:
+        return False
+
+
+class TempDirTracker(BaseCallbackHandler):
+    """Remember the temp directory a graph run materializes, as soon as it exists.
+
+    ``resolve_input`` reports the directory in its output as ``temp_dir_for_cleanup``,
+    but a run that raises or is interrupted returns no result for
+    :func:`cleanup_result`. Pass the tracker in the run's ``callbacks`` and use
+    :meth:`removing_on_error` around the run.
+    """
+
+    run_inline = True
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.temp_dir: str | None = None
+
+    def on_chain_end(self, outputs: Any, **kwargs: Any) -> None:
+        """Record ``temp_dir_for_cleanup`` from a node or graph output.
+
+        The recorded path is later removed recursively, so only a value that
+        looks like a scan temp dir is accepted: an existing directory, not a
+        symlink, whose name carries the ``skillspector_`` prefix that
+        ``InputHandler`` gives ``mkdtemp``. Anything else is ignored and never
+        replaces a path already recorded.
+        """
+        if isinstance(outputs, Mapping):
+            temp_dir = outputs.get("temp_dir_for_cleanup")
+            if _is_scan_temp_dir(temp_dir):
+                self.temp_dir = temp_dir
+
+    def remove(self) -> None:
+        """Remove the recorded temp directory, if any."""
+        if self.temp_dir:
+            remove_temp_tree(self.temp_dir)
+
+    @contextmanager
+    def removing_on_error(self) -> Iterator[None]:
+        """Remove the recorded temp directory if the wrapped run raises or is interrupted."""
+        try:
+            yield
+        except BaseException:
+            self.remove()
+            raise
+
+
 def cleanup_result(result: dict[str, object]) -> None:
     """Release scan-local resources and remove a temp dir if set."""
     python_ast_cache_key = result.get("python_ast_cache_key")

src/skillspector/cli.py

@@ -44,7 +44,7 @@
 from rich.tree import Tree
 
 from skillspector import __version__, transitive
-from skillspector.cleanup import cleanup_result
+from skillspector.cleanup import TempDirTracker, cleanup_result
 from skillspector.constants import RISK_THRESHOLD
 from skillspector.graph_proxy import graph
 from skillspector.input_handler import validate_local_input_path
@@ -1512,8 +1512,13 @@ def _run_graph_scan(
     if initial_inspection_ledger:
         state["inspection_ledger"] = initial_inspection_ledger
     trace_config = _build_trace_config(input_path, format, no_llm)
+    # A scan that raises or is interrupted returns no result for the caller to
+    # clean up, so remove the temp directory resolve_input made here instead.
+    temp_dir_tracker = TempDirTracker()
+    trace_config["callbacks"] = [temp_dir_tracker]
     if not stream_progress:
-        return cast(dict[str, object], graph.invoke(state, config=trace_config))
+        with temp_dir_tracker.removing_on_error():
+            return cast(dict[str, object], graph.invoke(state, config=trace_config))
 
     analyzer_node_ids = _wired_analyzer_node_ids()
     total_analyzers = len(analyzer_node_ids)
@@ -1531,6 +1536,7 @@ def _run_graph_scan(
             console=err_console,
             transient=True,
         ) as progress,
+        temp_dir_tracker.removing_on_error(),
     ):
         warnings.filterwarnings("ignore", category=UserWarning, module="pydantic")
         task_id = progress.add_task("Resolving input...", total=total_steps)
@@ -2527,21 +2533,29 @@ def _scan_skill(
         active_visited.add(transitive.canonicalize_source_identity(input_path))
     except ValueError:
         pass
-    return _scan_transitive(
-        initial_result=result,
-        format=format,
-        no_llm=no_llm,
-        max_depth=transitive_depth,
-        transitive_allow_prefix=transitive_allow_prefix,
-        transitive_deny_prefix=transitive_deny_prefix,
-        baseline=baseline,
-        show_suppressed=show_suppressed,
-        visited=active_visited,
-        scan_cache=transitive_cache,
-        yara_dir=yara_dir,
-        traversal=transitive_traversal,
-        source_local_only=source_local_only,
-    )
+    # The root graph has returned, so its tracker no longer guards the root's
+    # temp dir. If the transitive phase is interrupted or raises, nothing is
+    # returned for the caller's cleanup_result, so remove it here. On success
+    # the merged result carries the same temp_dir_for_cleanup for the caller.
+    try:
+        return _scan_transitive(
+            initial_result=result,
+            format=format,
+            no_llm=no_llm,
+            max_depth=transitive_depth,
+            transitive_allow_prefix=transitive_allow_prefix,
+            transitive_deny_prefix=transitive_deny_prefix,
+            baseline=baseline,
+            show_suppressed=show_suppressed,
+            visited=active_visited,
+            scan_cache=transitive_cache,
+            yara_dir=yara_dir,
+            traversal=transitive_traversal,
+            source_local_only=source_local_only,
+        )
+    except BaseException:
+        cleanup_result(result)
+        raise
 
 
 def _multi_skill_public_record_count(result: dict[str, object]) -> int:

src/skillspector/mcp_server.py

@@ -32,7 +32,7 @@
 from typing import TYPE_CHECKING, Any
 
 from skillspector import __version__
-from skillspector.cleanup import cleanup_result
+from skillspector.cleanup import TempDirTracker, cleanup_result
 from skillspector.constants import RISK_THRESHOLD
 from skillspector.graph import graph
 from skillspector.graph_proxy import restore_package_graph_export
@@ -153,10 +153,14 @@ async def run_scan(
     )
 
     result: dict[str, Any] | None = None
+    # A cancelled or failed scan returns no result to clean up; the tracker still
+    # knows the temp directory resolve_input made.
+    temp_dir_tracker = TempDirTracker()
     try:
         result = await graph.ainvoke(
             state,
             config={
+                "callbacks": [temp_dir_tracker],
                 "run_name": "skillspector-mcp-scan",
                 "tags": ["skillspector", "mcp"],
                 "metadata": {
@@ -228,6 +232,8 @@ async def run_scan(
     finally:
         if result is not None:
             cleanup_result(result)
+        else:
+            temp_dir_tracker.remove()
 
 
 def build_server(name: str = "skillspector", *, allow_local_targets: bool = False) -> FastMCP:

tests/unit/test_cleanup.py

@@ -11,7 +11,7 @@
 
 import pytest
 
-from skillspector.cleanup import _retry_writable, cleanup_result
+from skillspector.cleanup import TempDirTracker, _retry_writable, cleanup_result
 from skillspector.input_handler import InputHandler
 
 
@@ -163,3 +163,111 @@ def test_input_handler_cleanup_removes_read_only_git_objects(
 
     assert not temp_dir.exists()
     assert handler.temp_dir_for_cleanup() is None
+
+
+def _scan_temp_dir(root: Path) -> Path:
+    """Create a directory named like the ones ``InputHandler`` materializes."""
+    path = root / "skillspector_abc123"
+    path.mkdir()
+    (path / "SKILL.md").write_text("# Skill\n", encoding="utf-8")
+    return path
+
+
+def test_tracker_records_a_scan_temp_dir(tmp_path: Path) -> None:
+    """A node output naming an existing skillspector_ directory is recorded."""
+    temp_dir = _scan_temp_dir(tmp_path)
+    tracker = TempDirTracker()
+
+    tracker.on_chain_end({"temp_dir_for_cleanup": str(temp_dir)})
+
+    assert tracker.temp_dir == str(temp_dir)
+
+
+@pytest.mark.parametrize("later", [None, "", 42, ["x"]])
+def test_tracker_keeps_recorded_path_when_a_later_output_has_none(
+    tmp_path: Path, later: object
+) -> None:
+    """A later output without a usable path does not erase the recorded one."""
+    temp_dir = _scan_temp_dir(tmp_path)
+    tracker = TempDirTracker()
+    tracker.on_chain_end({"temp_dir_for_cleanup": str(temp_dir)})
+
+    tracker.on_chain_end({"temp_dir_for_cleanup": later})
+    tracker.on_chain_end({"other": "value"})
+    tracker.on_chain_end("not a mapping")
+
+    assert tracker.temp_dir == str(temp_dir)
+
+
+def test_tracker_ignores_paths_that_are_not_scan_temp_dirs(tmp_path: Path) -> None:
+    """Only an existing, non-symlink skillspector_ directory can become a deletion target."""
+    user_dir = tmp_path / "my-skill"
+    user_dir.mkdir()
+    missing = tmp_path / "skillspector_missing"
+    a_file = tmp_path / "skillspector_file"
+    a_file.write_text("x", encoding="utf-8")
+    tracker = TempDirTracker()
+
+    for candidate in (user_dir, missing, a_file):
+        tracker.on_chain_end({"temp_dir_for_cleanup": str(candidate)})
+
+    assert tracker.temp_dir is None
+    assert user_dir.exists()
+
+
+@pytest.mark.skipif(not hasattr(os, "symlink"), reason="symlinks unavailable")
+def test_tracker_ignores_a_symlink_named_like_a_scan_temp_dir(tmp_path: Path) -> None:
+    """A symlink is never recorded, even with the scan prefix and a directory target."""
+    target = tmp_path / "keep"
+    target.mkdir()
+    link = tmp_path / "skillspector_link"
+    try:
+        link.symlink_to(target, target_is_directory=True)
+    except OSError:
+        pytest.skip("cannot create symlinks here")
+    tracker = TempDirTracker()
+
+    tracker.on_chain_end({"temp_dir_for_cleanup": str(link)})
+
+    assert tracker.temp_dir is None
+
+
+def test_tracker_remove_twice_is_a_no_op(tmp_path: Path) -> None:
+    """Removing an already-removed directory does not raise."""
+    temp_dir = _scan_temp_dir(tmp_path)
+    tracker = TempDirTracker()
+    tracker.on_chain_end({"temp_dir_for_cleanup": str(temp_dir)})
+
+    tracker.remove()
+    tracker.remove()
+
+    assert not temp_dir.exists()
+
+
+@pytest.mark.parametrize("error", [KeyboardInterrupt, RuntimeError])
+def test_removing_on_error_reraises_the_original_exception(
+    tmp_path: Path, error: type[BaseException]
+) -> None:
+    """The wrapped run's exception propagates unchanged after the directory is removed."""
+    temp_dir = _scan_temp_dir(tmp_path)
+    tracker = TempDirTracker()
+    original = error("stopped")
+
+    with pytest.raises(error) as raised:
+        with tracker.removing_on_error():
+            tracker.on_chain_end({"temp_dir_for_cleanup": str(temp_dir)})
+            raise original
+
+    assert raised.value is original
+    assert not temp_dir.exists()
+
+
+def test_removing_on_error_leaves_the_directory_on_success(tmp_path: Path) -> None:
+    """A run that completes leaves removal to the caller's cleanup_result."""
+    temp_dir = _scan_temp_dir(tmp_path)
+    tracker = TempDirTracker()
+
+    with tracker.removing_on_error():
+        tracker.on_chain_end({"temp_dir_for_cleanup": str(temp_dir)})
+
+    assert temp_dir.exists()