merge queue: checking #1332 on main (0d8f2e0), stacked on 4 pull requests ahead

engine/hooks/_runner/dispatch.py

@@ -1,11 +1,12 @@
 from __future__ import annotations
 
 import argparse
-import concurrent.futures
 import glob
+import importlib.util
 import json
 import os
 import re
+import signal
 import subprocess
 import sys
 import time
@@ -16,7 +17,11 @@
 import run
 
 DEFAULT_BUDGET = 59.5
+DEFAULT_SUBAGENT_STOP_EVENT_BUDGET = 9.0
+DEFAULT_SUBAGENT_STOP_DETECTOR_BUDGET = 4.0
 OPT_OUT_KEY = "subagent_stop"
+DISPATCH_KEY = "dispatch"
+DEFAULT_DISPATCH_ENTRY = "detect:detect"
 MIRRORED_EVENTS = {"SubagentStop": "Stop"}
 DIRECT_RE = re.compile(r"^python3 \$HOME/\.(claude|cursor|codex)/hooks/([^/\s]+)/([^/\s]+\.py)((?:\s+.*)?)$")
 MESSAGE_KEYS = ("reason", "message", "additionalContext", "additional_context", "user_message")
@@ -157,6 +162,11 @@ def _load_manifest_hooks(
                 opt_out = manifest.get(OPT_OUT_KEY)
                 if isinstance(opt_out, dict) and opt_out.get("inherit") is False:
                     continue
+                dispatch_entry = _dispatch_entry(manifest)
+                if dispatch_entry is None:
+                    continue
+            else:
+                dispatch_entry = None
             for entry in entries:
                 matcher = None
                 if not mirrored and isinstance(entry, dict):
@@ -178,11 +188,26 @@ def _load_manifest_hooks(
                             "args": list(args),
                             "timeout": hook.get("timeout"),
                             "matcher": matcher,
+                            "dispatch_entry": dispatch_entry,
                         }
                     )
     return records, warnings
 
 
+def _dispatch_entry(manifest: dict) -> str | None:
+    config = manifest.get(OPT_OUT_KEY)
+    if not isinstance(config, dict):
+        return None
+    raw = config.get(DISPATCH_KEY)
+    if raw is True:
+        return DEFAULT_DISPATCH_ENTRY
+    if isinstance(raw, dict):
+        entry = raw.get("entry")
+        if isinstance(entry, str) and entry.strip():
+            return entry.strip()
+    return None
+
+
 def _parse_command(command: str, harness: str) -> tuple[str, str, str] | None:
     match = DIRECT_RE.fullmatch(command)
     if not match or match.group(1) != harness:
@@ -211,7 +236,22 @@ def _payload(stdin: bytes) -> dict[str, object]:
     return parsed if isinstance(parsed, dict) else {}
 
 
-def _run_one(hooks_root: str, python: str, record: dict, stdin: bytes, budget: float) -> dict:
+def _run_one(
+    hooks_root: str,
+    python: str,
+    record: dict,
+    stdin: bytes,
+    budget: float,
+    parsed_event: dict[str, object] | None = None,
+) -> dict:
+    if record.get("dispatch_entry") and _entrypoint_exists(hooks_root, record["hook"], str(record["dispatch_entry"])):
+        return _run_sdk_one(hooks_root, record, stdin, budget, parsed_event)
+    return _run_subprocess_one(hooks_root, python, record, stdin, budget, fallback=True)
+
+
+def _run_subprocess_one(
+    hooks_root: str, python: str, record: dict, stdin: bytes, budget: float, fallback: bool = False
+) -> dict:
     hook, script = record["hook"], record["script"]
     started = time.monotonic()
     script_path = os.path.join(hooks_root, hook, script)
@@ -249,6 +289,8 @@ def _run_one(hooks_root: str, python: str, record: dict, stdin: bytes, budget: f
         run._delete_findings_file(findings_path)
     result_outcome = outcome.classify(exit_code, stdout, stderr, timed_out)
     row = run._row(hooks_root, hook, script, stdin, result_outcome, exit_code, started, stdout, stderr, rule_ids)
+    if fallback:
+        row["dispatch_path"] = "subprocess_fallback"
     return {
         "hook": hook,
         "script": script,
@@ -260,12 +302,158 @@ def _run_one(hooks_root: str, python: str, record: dict, stdin: bytes, budget: f
     }
 
 
+def _sdk_dirs(hooks_root: str) -> tuple[str, str]:
+    engine_hooks = os.path.abspath(hooks_root)
+    sdk_dir = os.path.join(engine_hooks, "_sdk")
+    return engine_hooks, sdk_dir
+
+
+def _prefer_import_path(path: str, index: int) -> None:
+    try:
+        sys.path.remove(path)
+    except ValueError:
+        pass
+    sys.path.insert(index, path)
+
+
+def _load_entrypoint(hooks_root: str, hook: str, entry: str):
+    module_name, sep, function_name = entry.partition(":")
+    if not sep or not module_name or not function_name:
+        raise ValueError(f"dispatch entry must be MODULE:FUNCTION, got {entry!r}")
+    module_path = _entrypoint_path(hooks_root, hook, module_name)
+    if not os.path.isfile(module_path):
+        raise FileNotFoundError(module_path)
+    _engine_hooks, sdk_dir = _sdk_dirs(hooks_root)
+    _prefer_import_path(sdk_dir, 0)
+    _prefer_import_path(os.path.dirname(module_path), 1)
+    unique_name = f"_catstack_dispatch_{hook.replace('-', '_')}_{abs(hash(module_path))}"
+    spec = importlib.util.spec_from_file_location(unique_name, module_path)
+    if spec is None or spec.loader is None:
+        raise ImportError(f"could not load {module_path}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    target = getattr(module, function_name)
+    if not callable(target):
+        raise TypeError(f"{entry!r} did not resolve to a callable")
+    return target
+
+
+def _entrypoint_path(hooks_root: str, hook: str, module_name: str) -> str:
+    return os.path.join(hooks_root, hook, *module_name.split(".")) + ".py"
+
+
+def _entrypoint_exists(hooks_root: str, hook: str, entry: str) -> bool:
+    module_name, sep, function_name = entry.partition(":")
+    return bool(sep and module_name and function_name and os.path.isfile(_entrypoint_path(hooks_root, hook, module_name)))
+
+
+def _event(stdin: bytes, harness: str, event_name: str) -> dict[str, object]:
+    raw = stdin.decode("utf-8", errors="replace")
+    try:
+        parsed = json.loads(raw)
+    except json.JSONDecodeError as exc:
+        parsed = {
+            "_raw_payload": raw,
+            "_payload_error": f"the hook payload is not JSON ({exc})",
+            "_payload_error_type": type(exc).__name__,
+            "_payload_error_detail": str(exc),
+        }
+    if not isinstance(parsed, dict):
+        parsed = {"_raw_payload": raw, "_payload_error": "the hook payload is not a JSON object"}
+    parsed.setdefault("_raw_payload", raw)
+    parsed.setdefault("_catstack_harness", harness)
+    parsed.setdefault("hook_event_name", event_name)
+    return parsed
+
+
+def _render_findings(hooks_root: str, hook: str, harness: str, event_name: str, event: dict[str, object], findings: list):
+    _engine_hooks, sdk_dir = _sdk_dirs(hooks_root)
+    _prefer_import_path(sdk_dir, 0)
+    from modes import effective_finding_modes
+    from render import render
+    from runtime import _renderable_findings
+
+    mode, _mode_source, finding_modes = effective_finding_modes(hook, event, findings)
+    rendered_mode, rendered_findings = _renderable_findings(finding_modes, findings, mode)
+    stdout_text, stderr_text, exit_code = render(harness, event_name, rendered_mode, rendered_findings)
+    return stdout_text.encode(), stderr_text.encode(), exit_code, [finding.rule_id for finding in findings]
+
+
+def _run_sdk_one(
+    hooks_root: str,
+    record: dict,
+    stdin: bytes,
+    budget: float,
+    parsed_event: dict[str, object] | None = None,
+) -> dict:
+    hook, script = record["hook"], record["script"]
+    started = time.monotonic()
+    stdout = b""
+    stderr = b""
+    exit_code = 0
+    timed_out = False
+    rule_ids: list[str] = []
+    harness = run._harness(hooks_root)
+    payload = dict(parsed_event) if parsed_event is not None else _event(stdin, harness, "SubagentStop")
+    previous_handler = None
+    try:
+        detect = _load_entrypoint(hooks_root, hook, str(record["dispatch_entry"]))
+        if hasattr(signal, "SIGALRM"):
+            previous_handler = signal.getsignal(signal.SIGALRM)
+
+            def _timeout(_signum, _frame):
+                raise TimeoutError(f"{hook}/{record['dispatch_entry']} timed out after {budget}s")
+
+            signal.signal(signal.SIGALRM, _timeout)
+            signal.setitimer(signal.ITIMER_REAL, budget)
+        findings = detect(payload)
+        if not isinstance(findings, list):
+            raise TypeError("detect(event) must return a list")
+        stdout, stderr, exit_code, rule_ids = _render_findings(hooks_root, hook, harness, "SubagentStop", payload, findings)
+        lines = payload.get("_catstack_stderr_lines")
+        if isinstance(lines, list):
+            stderr += "".join(f"{line}\n" for line in lines if isinstance(line, str) and line).encode()
+    except TimeoutError:
+        timed_out = True
+        exit_code = 1
+        stdout = b""
+        stderr = f"catstack-hook-runner: {hook}/{script} timed out after {budget}s\n".encode()
+    except BaseException as exc:
+        if isinstance(exc, KeyboardInterrupt):
+            raise
+        if isinstance(exc, SystemExit):
+            exit_code = exc.code if isinstance(exc.code, int) else 1
+        else:
+            exit_code = 0
+        stdout = b""
+        stderr = f"catstack-hook-error {hook}: {type(exc).__name__}: {exc}\n".encode()
+    finally:
+        if hasattr(signal, "SIGALRM"):
+            signal.setitimer(signal.ITIMER_REAL, 0)
+            if previous_handler is not None:
+                signal.signal(signal.SIGALRM, previous_handler)
+    result_outcome = outcome.classify(exit_code, stdout, stderr, timed_out)
+    row = run._row(hooks_root, hook, script, stdin, result_outcome, exit_code, started, stdout, stderr, rule_ids)
+    row["dispatch_path"] = "sdk" if not timed_out and not stderr.startswith(b"catstack-hook-error ") else "sdk_error"
+    return {
+        "hook": hook,
+        "script": script,
+        "outcome": result_outcome,
+        "exit_code": exit_code,
+        "stdout": stdout,
+        "stderr": stderr,
+        "row": row,
+    }
+
+
 def _merge_stdout(results: list[dict]) -> tuple[bytes, bytes, int]:
     blocked = [r for r in results if r["outcome"] == "blocked"]
     spoke = [r for r in results if r["stdout"].strip()]
     if blocked:
         return _merge_blocks(blocked)
     if spoke:
+        if len(spoke) == 1:
+            return spoke[0]["stdout"], b"", spoke[0]["exit_code"]
         contexts = [note for r in spoke for note in (_note(r),) if note]
         payload: dict[str, object] = {"continue": True}
         if contexts:
@@ -340,11 +528,9 @@ def run_dispatch(
         ).encode()
         return 1, b"", stderr, []
 
-    results: list[dict] = []
-    with concurrent.futures.ThreadPoolExecutor(max_workers=len(records)) as pool:
-        futures = [pool.submit(_run_one, hooks_root, python, record, stdin, budget) for record in records]
-        for future in futures:
-            results.append(future.result())
+    detector_budget = min(budget, DEFAULT_SUBAGENT_STOP_DETECTOR_BUDGET) if event == "SubagentStop" else budget
+    parsed_event = _event(stdin, harness, event) if event == "SubagentStop" else None
+    results = [_run_one(hooks_root, python, record, stdin, detector_budget, parsed_event) for record in records]
 
     stdout, block_notes, exit_code = _merge_stdout(results)
     stderr_lines = [r["stderr"] for r in results if r["stderr"]]
@@ -355,7 +541,7 @@ def run_dispatch(
 def _parse_args(argv: list[str]) -> argparse.Namespace:
     parser = argparse.ArgumentParser()
     parser.add_argument("--event", required=True)
-    parser.add_argument("--timeout", type=float, default=DEFAULT_BUDGET)
+    parser.add_argument("--timeout", type=float)
     return parser.parse_args(argv)
 
 
@@ -364,7 +550,10 @@ def main(argv: list[str] | None = None) -> int:
     stdin = sys.stdin.buffer.read()
     hooks_root = run._hooks_root()
     harness = run._harness(hooks_root)
-    exit_code, stdout, stderr, rows = run_dispatch(hooks_root, harness, args.event, stdin, args.timeout)
+    budget = args.timeout
+    if budget is None:
+        budget = DEFAULT_SUBAGENT_STOP_EVENT_BUDGET if args.event == "SubagentStop" else DEFAULT_BUDGET
+    exit_code, stdout, stderr, rows = run_dispatch(hooks_root, harness, args.event, stdin, budget)
     metrics_path = run._metrics_path()
     metrics_errors = b"".join(run._write_metrics(row, metrics_path) for row in rows)
     sys.stdout.buffer.write(stdout)