ensure_agent_sdk.py 9.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225
  1. #!/usr/bin/env python3
  2. """SessionStart bootstrap: ensure claude_agent_sdk is importable for the
  3. agentic commit reviewer.
  4. If claude_agent_sdk already imports in the current python3, this is a no-op.
  5. Otherwise it creates a venv at ~/.claude/security/agent-sdk-venv and installs
  6. the SDK there. security_reminder_hook.py prepends that venv's site-packages to
  7. sys.path before attempting the SDK import, so the venv is used as a
  8. fallback only when the system install is missing.
  9. The venv lives under ~/.claude/security/ (same dir the plugin already uses
  10. for per-session state) so it persists across plugin updates — rebuilding
  11. on every update is 30-60s of wasted work for a package that changes far
  12. less often than the plugin does.
  13. """
  14. from __future__ import annotations
  15. import importlib.util
  16. import json
  17. import os
  18. import subprocess
  19. import sys
  20. import time
  21. from pathlib import Path
  22. # Outcome codes for the sdk_bootstrap metric. Values are stable for telemetry.
  23. NOOP_SYSTEM = 0 # claude_agent_sdk already importable in system python
  24. NOOP_VENV = 1 # venv already built and SDK imports from it
  25. BUILT = 2 # venv created + SDK pip-installed this run
  26. BUILD_FAILED = 3 # venv create or pip install raised/timed out
  27. SKIP_WIN32 = 4 # Windows; consumer glob doesn't handle Lib/ layout
  28. SKIP_SENTINEL = 5 # another SessionStart is currently building
  29. def _sdk_on_syspath() -> bool:
  30. # find_spec is ~10ms; actually importing the SDK pulls in
  31. # transitive deps and costs ~800ms — too heavy for a
  32. # per-SessionStart no-op check that most sessions hit.
  33. try:
  34. return importlib.util.find_spec("claude_agent_sdk") is not None
  35. except Exception:
  36. return False
  37. def _plugin_version_int() -> int:
  38. # Same encoding as security_reminder_hook._read_plugin_version_int so
  39. # metrics rows from both hooks join on pv.
  40. try:
  41. p = Path(__file__).parent.parent / ".claude-plugin" / "plugin.json"
  42. v = json.loads(p.read_text())["version"]
  43. major, minor, patch = (int(x) for x in v.split(".")[:3])
  44. return major * 10000 + minor * 100 + patch
  45. except Exception:
  46. return 0
  47. def main() -> tuple[int, str, str]:
  48. """Run the bootstrap. Returns (outcome, err_phase, err_kind).
  49. err_phase / err_kind are non-empty only on BUILD_FAILED — they let
  50. telemetry split bootstrap failures by root cause.
  51. """
  52. # Windows venv layout (Lib/site-packages, no python* subdir) isn't
  53. # handled by the consumer's glob in security_reminder_hook.py; skip the
  54. # bootstrap entirely rather than build a venv that's never read.
  55. if sys.platform == "win32":
  56. return SKIP_WIN32, "", ""
  57. if _sdk_on_syspath():
  58. return NOOP_SYSTEM, "", ""
  59. state_dir = Path(
  60. os.environ.get("SECURITY_WARNINGS_STATE_DIR")
  61. or os.path.expanduser("~/.claude/security")
  62. )
  63. venv = state_dir / "agent-sdk-venv"
  64. venv_py = venv / "bin" / "python"
  65. # Another SessionStart (concurrent CC instance, same plugin) may already
  66. # be building. The sentinel lives NEXT TO the venv, not inside it —
  67. # `python -m venv --clear` wipes the target dir's contents, so an
  68. # in-venv sentinel would be deleted the instant we create the venv.
  69. # Stale sentinels (>5min) from a SIGKILL'd build are ignored.
  70. sentinel = state_dir / "agent-sdk-venv.building"
  71. if sentinel.exists():
  72. try:
  73. if time.time() - sentinel.stat().st_mtime < 300:
  74. return SKIP_SENTINEL, "", ""
  75. sentinel.unlink(missing_ok=True)
  76. except OSError:
  77. return SKIP_SENTINEL, "", ""
  78. # If a venv already exists and its python can import the SDK, done.
  79. if venv_py.exists():
  80. try:
  81. r = subprocess.run(
  82. [str(venv_py), "-c", "import claude_agent_sdk"],
  83. capture_output=True, timeout=10,
  84. )
  85. if r.returncode == 0:
  86. return NOOP_VENV, "", ""
  87. except Exception:
  88. pass # broken venv; rebuild below
  89. err_phase = ""
  90. err_kind = ""
  91. we_own_sentinel = False
  92. try:
  93. state_dir.mkdir(parents=True, exist_ok=True)
  94. # O_EXCL makes the sentinel an atomic lock — if two SessionStarts
  95. # race past the exists() check above, only one creates it.
  96. try:
  97. os.close(os.open(sentinel, os.O_CREAT | os.O_EXCL | os.O_WRONLY))
  98. except FileExistsError:
  99. return SKIP_SENTINEL, "", ""
  100. we_own_sentinel = True
  101. err_phase = "venv"
  102. subprocess.run(
  103. [sys.executable, "-m", "venv", "--clear", str(venv)],
  104. capture_output=True, timeout=60, check=True,
  105. )
  106. # Some machines route pip through a private registry; we
  107. # don't pass --index-url here so we inherit that default. Outside
  108. # the user's machine, pip's own default registry applies — that's the same
  109. # exposure the user would have running `pip install` themselves, so
  110. # we're not widening the supply-chain surface.
  111. err_phase = "pip"
  112. subprocess.run(
  113. [str(venv_py), "-m", "pip", "install", "--quiet",
  114. "--disable-pip-version-check", "claude-agent-sdk"],
  115. capture_output=True, timeout=120, check=True,
  116. )
  117. return BUILT, "", ""
  118. except subprocess.CalledProcessError as e:
  119. # Capture a stderr fingerprint so telemetry can split BUILD_FAILED by
  120. # root cause (no-network, package-not-found, dns-fail, etc.).
  121. # Categorize first, then keep a short raw tail for the long tail of
  122. # unexpected modes.
  123. stderr_b = e.stderr or b""
  124. if isinstance(stderr_b, bytes):
  125. stderr_str = stderr_b.decode("utf-8", errors="replace")
  126. else:
  127. stderr_str = str(stderr_b)
  128. s = stderr_str.lower()
  129. if "no matching distribution" in s or "could not find a version" in s:
  130. err_kind = "pip_no_match"
  131. elif "name or service not known" in s or "name resolution" in s \
  132. or "nodename nor servname" in s or "temporary failure in name" in s:
  133. err_kind = "dns_fail"
  134. elif "connection refused" in s or "connection reset" in s:
  135. err_kind = "conn_refused"
  136. elif "ssl" in s and ("verify" in s or "certificate" in s):
  137. err_kind = "ssl_verify"
  138. elif "permission denied" in s or "read-only file system" in s:
  139. err_kind = "perm_denied"
  140. elif "no module named pip" in s or "no module named ensurepip" in s:
  141. err_kind = "no_pip"
  142. elif "no space left" in s or "disk quota" in s:
  143. err_kind = "disk_full"
  144. elif "proxy" in s and ("authent" in s or "tunnel" in s or "407" in s):
  145. err_kind = "proxy_auth"
  146. elif "timeout" in s or "timed out" in s:
  147. err_kind = "stderr_timeout"
  148. else:
  149. # First 60 chars of the last non-empty stderr line — bounded to
  150. # stay inside CC's metric value-length budget. Real failure modes
  151. # we haven't categorized show up here as a low-cardinality bucket.
  152. tail = next(
  153. (ln.strip() for ln in reversed(stderr_str.splitlines()) if ln.strip()),
  154. "",
  155. )[:60]
  156. err_kind = f"other:{tail}" if tail else "other"
  157. return BUILD_FAILED, err_phase, err_kind
  158. except subprocess.TimeoutExpired:
  159. return BUILD_FAILED, err_phase, "subprocess_timeout"
  160. except Exception as e:
  161. return BUILD_FAILED, err_phase, f"exc:{type(e).__name__}"
  162. finally:
  163. # Only remove the sentinel if THIS process created it. The
  164. # FileExistsError path above means another process owns the lock;
  165. # unconditionally unlinking here would delete its sentinel and let
  166. # a third concurrent SessionStart `venv --clear` over the in-flight
  167. # build.
  168. if we_own_sentinel:
  169. sentinel.unlink(missing_ok=True)
  170. if __name__ == "__main__":
  171. # Tell the harness this is async — venv create + pip install can take
  172. # 30-60s on a cold cache, well past the default sync hook timeout.
  173. # SessionStart runs before the user's first prompt; doing this in the
  174. # background means the first commit-review of the session usually finds
  175. # the venv ready.
  176. print(json.dumps({"async": True, "asyncTimeout": 180000}), flush=True)
  177. t0 = time.perf_counter()
  178. try:
  179. outcome, err_phase, err_kind = main()
  180. except Exception as exc:
  181. outcome, err_phase, err_kind = (
  182. BUILD_FAILED, "main", f"exc:{type(exc).__name__}"
  183. )
  184. # CC's async-hook registry scans stdout line-by-line after process exit
  185. # and takes the FIRST non-{"async":...} JSON line as the hook response;
  186. # its `metrics` key is forwarded to the hook metrics event on the
  187. # next attachments pass. Must be a single line — the registry splits on
  188. # \n and json-parses each independently. Values must be bool|number OR
  189. # short strings (CC accepts string metric values if they're not
  190. # null). Stay inside the 10-key emit cap.
  191. metrics: dict[str, object] = {
  192. "sdk_bootstrap": outcome,
  193. "sdk_bootstrap_ms": round((time.perf_counter() - t0) * 1000),
  194. }
  195. if err_kind:
  196. # Truncate defensively; categorized values are <40 chars but the
  197. # `other:<tail>` mode could be longer. err_phase may be empty for
  198. # pre-venv failures (state_dir.mkdir perm-denied, sentinel O_EXCL
  199. # raising a non-FileExistsError OSError) — emit as "pre" so the
  200. # err_kind isn't silently dropped.
  201. metrics["sdk_bootstrap_phase"] = (err_phase or "pre")[:16]
  202. metrics["sdk_bootstrap_err"] = err_kind[:96]
  203. pv = _plugin_version_int()
  204. if pv:
  205. metrics["pv"] = pv
  206. print(json.dumps({"metrics": metrics}), flush=True)