232 lines
8.5 KiB
Python
232 lines
8.5 KiB
Python
"""disk-cleanup plugin — auto-cleanup of ephemeral Hermes session files.
|
|
|
|
Wires three behaviours:
|
|
|
|
1. ``post_tool_call`` hook — inspects ``write_file`` / ``patch`` / ``terminal``
|
|
tool calls for newly-created paths matching test/temp patterns under
|
|
``HERMES_HOME`` and tracks them silently. Zero agent compliance required.
|
|
2. ``on_session_end`` hook — when any test files were auto-tracked during the
|
|
just-finished turn, runs :func:`disk_cleanup.quick` and logs one line to
|
|
``$HERMES_HOME/disk-cleanup/cleanup.log``.
|
|
3. ``/disk-cleanup`` slash command — manual ``status``, ``dry-run``, ``quick``,
|
|
``deep``, ``track``, ``forget``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import re
|
|
import shlex
|
|
import threading
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Dict, List, Optional, Set
|
|
|
|
from . import disk_cleanup as dg
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# Per-task set of test files newly tracked this turn, keyed by task_id (or
|
|
# session_id as fallback) so on_session_end can decide whether to run cleanup.
|
|
# Locked: post_tool_call can fire concurrently on parallel tool calls.
|
|
_recent_test_tracks: Dict[str, Set[str]] = {}
|
|
_lock = threading.Lock()
|
|
|
|
_TERMINAL_PATH_REGEX = re.compile(r"(?:^|\s)(/[^\s'\"`]+|\~/[^\s'\"`]+)")
|
|
|
|
|
|
# --- Path extraction from tool calls ----------------------------------------
|
|
|
|
def _extract_path_arg(args: Dict[str, Any], result: str) -> Set[str]:
|
|
"""write_file and patch: only the single-file ``path`` arg. patch mostly edits
|
|
existing files; re-tracking is a no-op (track() dedups), so this is safe."""
|
|
path = args.get("path")
|
|
return {path} if isinstance(path, str) and path else set()
|
|
|
|
|
|
def _extract_paths_from_terminal(args: Dict[str, Any], result: str) -> Set[str]:
|
|
"""Best-effort: pull candidate filesystem paths from a terminal command and
|
|
its output; ``guess_category`` / ``is_safe_path`` filter them afterwards."""
|
|
paths: Set[str] = set()
|
|
cmd = args.get("command") or ""
|
|
if isinstance(cmd, str) and cmd:
|
|
try: # tokenise — catches `touch /tmp/hermes-x/test_foo.py`
|
|
paths.update(tok for tok in shlex.split(cmd, posix=True) if tok.startswith(("/", "~")))
|
|
except ValueError:
|
|
pass
|
|
# Only scan the result text if it's a reasonable size (avoid 50KB dumps).
|
|
if isinstance(result, str) and len(result) < 4096:
|
|
paths.update(_TERMINAL_PATH_REGEX.findall(result))
|
|
return paths
|
|
|
|
|
|
_PATH_EXTRACTORS: Dict[str, Callable[[Dict[str, Any], str], Set[str]]] = {
|
|
"write_file": _extract_path_arg,
|
|
"patch": _extract_path_arg,
|
|
"terminal": _extract_paths_from_terminal,
|
|
}
|
|
|
|
|
|
# --- Hooks ------------------------------------------------------------------
|
|
|
|
def _on_post_tool_call(
|
|
tool_name: str = "",
|
|
args: Optional[Dict[str, Any]] = None,
|
|
result: Any = None,
|
|
task_id: str = "",
|
|
session_id: str = "",
|
|
tool_call_id: str = "",
|
|
**_: Any,
|
|
) -> None:
|
|
"""Auto-track ephemeral files created by recent tool calls. Best-effort, never raises."""
|
|
extractor = _PATH_EXTRACTORS.get(tool_name)
|
|
if not isinstance(args, dict) or extractor is None:
|
|
return
|
|
for path_str in extractor(args, result if isinstance(result, str) else ""):
|
|
try:
|
|
p = Path(path_str).expanduser()
|
|
except Exception:
|
|
continue
|
|
if not p.exists():
|
|
continue
|
|
category = dg.guess_category(p)
|
|
if category is not None and dg.track(str(p), category, silent=True) and category == "test":
|
|
with _lock:
|
|
_recent_test_tracks.setdefault(task_id or session_id or "default", set()).add(str(p))
|
|
|
|
|
|
def _on_session_end(
|
|
session_id: str = "",
|
|
completed: bool = True,
|
|
interrupted: bool = False,
|
|
**_: Any,
|
|
) -> None:
|
|
"""Run quick cleanup if any test files were tracked during this turn."""
|
|
# Drain the session bucket plus every task-scoped bucket: subagents record
|
|
# into their own task_id buckets and should be cleaned up at session end too.
|
|
with _lock:
|
|
had_tracks = bool(_recent_test_tracks.pop(session_id or "default", None) or _recent_test_tracks)
|
|
_recent_test_tracks.clear()
|
|
if not had_tracks:
|
|
return
|
|
|
|
try:
|
|
summary = dg.quick()
|
|
except Exception as exc:
|
|
logger.debug("disk-cleanup quick cleanup failed: %s", exc)
|
|
return
|
|
|
|
if summary["deleted"] or summary["empty_dirs"]:
|
|
dg._log(
|
|
f"AUTO_QUICK (session_end): deleted={summary['deleted']} "
|
|
f"dirs={summary['empty_dirs']} freed={dg.fmt_size(summary['freed'])}"
|
|
)
|
|
|
|
|
|
# --- Slash command ----------------------------------------------------------
|
|
|
|
_HELP_TEXT = """\
|
|
/disk-cleanup — ephemeral-file cleanup
|
|
|
|
Subcommands:
|
|
status Per-category breakdown + top-10 largest
|
|
dry-run Preview what quick/deep would delete
|
|
quick Run safe cleanup now (no prompts)
|
|
deep Run quick, then list items that need prompts
|
|
track <path> <category> Manually add a path to tracking
|
|
forget <path> Stop tracking a path (does not delete)
|
|
|
|
Categories: temp | test | research | download | chrome-profile | cron-output | other
|
|
|
|
All operations are scoped to HERMES_HOME and /tmp/hermes-*.
|
|
Test files are auto-tracked on write_file / terminal and auto-cleaned at session end.
|
|
"""
|
|
|
|
|
|
def _fmt_summary(summary: Dict[str, Any]) -> str:
|
|
base = (
|
|
f"[disk-cleanup] Cleaned {summary['deleted']} files + "
|
|
f"{summary['empty_dirs']} empty dirs, freed {dg.fmt_size(summary['freed'])}."
|
|
)
|
|
if summary.get("errors"):
|
|
base += f"\n {len(summary['errors'])} error(s); see cleanup.log."
|
|
return base
|
|
|
|
|
|
def _item_block(header: str, items: List[Dict], indent: str) -> List[str]:
|
|
"""``header`` formatted with ``{n}`` / ``{size}``, followed by one line per item."""
|
|
size = dg.fmt_size(sum(i["size"] for i in items))
|
|
return [header.format(n=len(items), size=size)] + [f"{indent}[{item['category']}] {item['path']}" for item in items]
|
|
|
|
|
|
def _cmd_dry_run(argv: List[str]) -> str:
|
|
auto, prompt = dg.dry_run()
|
|
lines = ["Dry-run preview (nothing deleted):"]
|
|
lines += _item_block(" Auto-delete : {n} files ({size})", auto, " ")
|
|
lines += _item_block(" Needs prompt: {n} files ({size})", prompt, " ")
|
|
lines.append(f"\n Total potential: {dg.fmt_size(sum(i['size'] for i in auto + prompt))}")
|
|
return "\n".join(lines)
|
|
|
|
|
|
def _cmd_deep(argv: List[str]) -> str:
|
|
# In-session deep can't prompt interactively — show what quick cleaned plus
|
|
# the items that WOULD need confirmation.
|
|
quick_summary = dg.quick()
|
|
_auto, prompt_items = dg.dry_run()
|
|
lines = [_fmt_summary(quick_summary)]
|
|
if prompt_items:
|
|
lines += _item_block("\n{n} item(s) need confirmation ({size}):", prompt_items, " ")
|
|
lines.append("\nRun `/disk-cleanup forget <path>` to skip, or delete manually via terminal.")
|
|
return "\n".join(lines)
|
|
|
|
|
|
def _cmd_track(argv: List[str]) -> str:
|
|
if len(argv) < 3:
|
|
return "Usage: /disk-cleanup track <path> <category>"
|
|
path_arg, category = argv[1], argv[2]
|
|
if category not in dg.ALLOWED_CATEGORIES:
|
|
return f"Unknown category '{category}'. Allowed: {sorted(dg.ALLOWED_CATEGORIES)}"
|
|
if dg.track(path_arg, category, silent=True):
|
|
return f"Tracked {path_arg} as '{category}'."
|
|
return f"Not tracked (already present, missing, or outside HERMES_HOME): {path_arg}"
|
|
|
|
|
|
def _cmd_forget(argv: List[str]) -> str:
|
|
if len(argv) < 2:
|
|
return "Usage: /disk-cleanup forget <path>"
|
|
n = dg.forget(argv[1])
|
|
return (
|
|
f"Removed {n} tracking entr{'y' if n == 1 else 'ies'} for {argv[1]}."
|
|
if n else f"Not found in tracking: {argv[1]}"
|
|
)
|
|
|
|
|
|
_SUBCOMMANDS: Dict[str, Callable[[List[str]], str]] = {
|
|
"status": lambda argv: dg.format_status(dg.status()),
|
|
"dry-run": _cmd_dry_run,
|
|
"quick": lambda argv: _fmt_summary(dg.quick()),
|
|
"deep": _cmd_deep,
|
|
"track": _cmd_track,
|
|
"forget": _cmd_forget,
|
|
}
|
|
|
|
|
|
def _handle_slash(raw_args: str) -> Optional[str]:
|
|
argv = raw_args.strip().split()
|
|
if not argv or argv[0] in {"help", "-h", "--help"}:
|
|
return _HELP_TEXT
|
|
handler = _SUBCOMMANDS.get(argv[0])
|
|
if handler is None:
|
|
return f"Unknown subcommand: {argv[0]}\n\n{_HELP_TEXT}"
|
|
return handler(argv)
|
|
|
|
|
|
def register(ctx) -> None:
|
|
ctx.register_hook("post_tool_call", _on_post_tool_call)
|
|
ctx.register_hook("on_session_end", _on_session_end)
|
|
ctx.register_command(
|
|
"disk-cleanup",
|
|
handler=_handle_slash,
|
|
description="Track and clean up ephemeral Hermes session files.",
|
|
)
|