From 623e1b2b86819694898fcee546be703fb3f00674 Mon Sep 17 00:00:00 2001 From: olive Date: Wed, 27 May 2026 01:04:29 -0700 Subject: [PATCH] Add post-tool and session-end hooks for adapters --- Gradata/src/gradata/hooks/adapters/_base.py | 16 +++- Gradata/src/gradata/hooks/adapters/codex.py | 66 ++++++++++++---- Gradata/src/gradata/hooks/adapters/hermes.py | 45 +++++++---- .../src/gradata/hooks/adapters/opencode.py | 76 +++++++++++++++---- Gradata/tests/test_hook_adapters.py | 70 +++++++++++++++++ 5 files changed, 227 insertions(+), 46 deletions(-) diff --git a/Gradata/src/gradata/hooks/adapters/_base.py b/Gradata/src/gradata/hooks/adapters/_base.py index 7ea67bc5..c159f2ee 100644 --- a/Gradata/src/gradata/hooks/adapters/_base.py +++ b/Gradata/src/gradata/hooks/adapters/_base.py @@ -129,13 +129,25 @@ def hook_signature(agent: str, brain_dir: Path) -> str: return f"gradata:{agent}:{brain_dir.resolve().as_posix()}" -def hook_command(brain_dir: Path) -> str: +def hook_command_for_module(brain_dir: Path, module: str) -> str: return ( f"BRAIN_DIR={shlex.quote(str(brain_dir))} " - f"{shlex.quote(sys.executable)} -m gradata.hooks.inject_brain_rules" + f"{shlex.quote(sys.executable)} -m {module}" ) +def hook_command(brain_dir: Path) -> str: + return hook_command_for_module(brain_dir, "gradata.hooks.inject_brain_rules") + + +def post_tool_hook_command(brain_dir: Path) -> str: + return hook_command_for_module(brain_dir, "gradata.hooks.auto_correct") + + +def session_end_hook_command(brain_dir: Path) -> str: + return hook_command_for_module(brain_dir, "gradata.hooks.session_close") + + def mcp_command(brain_dir: Path) -> list[str]: return [sys.executable, "-m", "gradata.mcp_server", "--brain-dir", str(brain_dir)] diff --git a/Gradata/src/gradata/hooks/adapters/codex.py b/Gradata/src/gradata/hooks/adapters/codex.py index dd73525a..c123f558 100644 --- a/Gradata/src/gradata/hooks/adapters/codex.py +++ b/Gradata/src/gradata/hooks/adapters/codex.py @@ -9,12 +9,13 @@ WRITE_TOOL_ALIASES, InstallResult, _normalize_tool_name, - contains_signature, extract_from_edit_args, extract_from_write_args, failure, hook_command, hook_signature, + post_tool_hook_command, + session_end_hook_command, ) AGENT = "codex" @@ -70,32 +71,65 @@ def _toml_string(value: str) -> str: return json.dumps(value) +def _hook_table_has_signature(text: str, table_name: str, signature: str) -> bool: + """Return True if a specific TOML hook table already contains signature.""" + current: list[str] = [] + in_table = False + + def flush() -> bool: + return in_table and any(signature in line for line in current) + + for line in text.splitlines(keepends=True): + stripped = line.lstrip() + if stripped.startswith("[[") or stripped.startswith("["): + if flush(): + return True + current = [line] + in_table = stripped.startswith(f"[[hooks.{table_name}]]") + elif in_table: + current.append(line) + return flush() + + def install(brain_dir: Path, agent_config_path: Path) -> InstallResult: try: sig = hook_signature(AGENT, brain_dir) - if contains_signature(agent_config_path, sig): - return InstallResult( - AGENT, agent_config_path, "already_present", "hook already present" - ) existing = ( agent_config_path.read_text(encoding="utf-8") if agent_config_path.exists() else "" ) - block = ( - "\n[[hooks.pre_tool]]\n" - f"id = {_toml_string(sig)}\n" - f"command = {_toml_string(hook_command(brain_dir))}\n" + blocks: list[str] = [] + for table_name, command in ( + ("pre_tool", hook_command(brain_dir)), + ("post_tool", post_tool_hook_command(brain_dir)), + ("session_end", session_end_hook_command(brain_dir)), + ): + if _hook_table_has_signature(existing, table_name, sig): + continue + blocks.append( + f"\n[[hooks.{table_name}]]\n" + f"id = {_toml_string(sig)}\n" + f"command = {_toml_string(command)}\n" + ) + if not blocks: + return InstallResult( + AGENT, agent_config_path, "already_present", "hooks already present" + ) + atomic_write_text(agent_config_path, existing.rstrip() + "".join(blocks)) + return InstallResult( + AGENT, + agent_config_path, + "added", + "installed pre_tool, post_tool, and session_end hooks", ) - atomic_write_text(agent_config_path, existing.rstrip() + block) - return InstallResult(AGENT, agent_config_path, "added", "installed pre_tool hook") except Exception as exc: return failure(AGENT, agent_config_path, exc) def uninstall(brain_dir: Path, agent_config_path: Path) -> InstallResult: - """Reverse install: drop the [[hooks.pre_tool]] block carrying our signature. + """Reverse install: drop hook blocks carrying our signature. Operates on the raw TOML text — walks line-by-line, identifies the - [[hooks.pre_tool]] table that contains our signature, and removes that + hook tables that contain our signature, and removes those table + its keys. Preserves all other tables verbatim. """ try: @@ -131,7 +165,9 @@ def flush(buf: list[str], is_hook: bool) -> None: # Flush the previous table flush(current_table, current_is_hook) current_table = [line] - current_is_hook = stripped.startswith("[[hooks.pre_tool]]") + current_is_hook = stripped.startswith( + ("[[hooks.pre_tool]]", "[[hooks.post_tool]]", "[[hooks.session_end]]") + ) else: if current_table: current_table.append(line) @@ -148,7 +184,7 @@ def flush(buf: list[str], is_hook: bool) -> None: new_text = "".join(out_lines).rstrip() + "\n" atomic_write_text(agent_config_path, new_text) return InstallResult( - AGENT, agent_config_path, "removed", f"removed {removed} [[hooks.pre_tool]] block" + AGENT, agent_config_path, "removed", f"removed {removed} hook block" ) except Exception as exc: return failure(AGENT, agent_config_path, exc) diff --git a/Gradata/src/gradata/hooks/adapters/hermes.py b/Gradata/src/gradata/hooks/adapters/hermes.py index ada9e275..28437ed2 100644 --- a/Gradata/src/gradata/hooks/adapters/hermes.py +++ b/Gradata/src/gradata/hooks/adapters/hermes.py @@ -8,12 +8,13 @@ WRITE_TOOL_ALIASES, InstallResult, _normalize_tool_name, - contains_signature, extract_from_edit_args, extract_from_write_args, failure, hook_command, hook_signature, + post_tool_hook_command, + session_end_hook_command, ) AGENT = "hermes" @@ -149,11 +150,6 @@ def _format_scalar(value) -> str: def install(brain_dir: Path, agent_config_path: Path) -> InstallResult: try: sig = hook_signature(AGENT, brain_dir) - if contains_signature(agent_config_path, sig): - return InstallResult( - AGENT, agent_config_path, "already_present", "hook already present" - ) - existing = ( agent_config_path.read_text(encoding="utf-8") if agent_config_path.exists() else "" ) @@ -167,14 +163,28 @@ def install(brain_dir: Path, agent_config_path: Path) -> InstallResult: # Claude-Code names (pre_tool_use / post_tool_use / session_end) here # results in Hermes silently ignoring the entries (warning only). See # Gradata/gradata#190 for the install-UX epic. - pre_tool_call = _migrate_legacy_event(hooks, "pre_tool_use", "pre_tool_call") - if any(isinstance(entry, dict) and entry.get("id") == sig for entry in pre_tool_call): + added: list[str] = [] + for current_key, legacy_key, command in ( + ("pre_tool_call", "pre_tool_use", hook_command(brain_dir)), + ("post_tool_call", "post_tool_use", post_tool_hook_command(brain_dir)), + ("on_session_end", "session_end", session_end_hook_command(brain_dir)), + ): + entries = _migrate_legacy_event(hooks, legacy_key, current_key) + if any(isinstance(entry, dict) and entry.get("id") == sig for entry in entries): + continue + entries.append({"id": sig, "command": command}) + added.append(current_key) + if not added: return InstallResult( - AGENT, agent_config_path, "already_present", "hook already present" + AGENT, agent_config_path, "already_present", "hooks already present" ) - pre_tool_call.append({"id": sig, "command": hook_command(brain_dir)}) atomic_write_text(agent_config_path, _dump_simple_yaml(data)) - return InstallResult(AGENT, agent_config_path, "added", "installed pre_tool_call hook") + return InstallResult( + AGENT, + agent_config_path, + "added", + "installed pre_tool_call, post_tool_call, and on_session_end hooks", + ) except Exception as exc: return failure(AGENT, agent_config_path, exc) @@ -232,7 +242,7 @@ def extract_correction( def uninstall(brain_dir: Path, agent_config_path: Path) -> InstallResult: - """Reverse install: drop signature-matching entries from hooks.pre_tool_call. + """Reverse install: drop signature-matching hook entries. Hermes uses YAML, so we can't reuse the generic JSON helper. Idempotent. Preserves user-owned entries and other hook events. @@ -253,8 +263,15 @@ def uninstall(brain_dir: Path, agent_config_path: Path) -> InstallResult: return InstallResult(AGENT, agent_config_path, "already_present", "no hooks block") removed = 0 - # Check both current and legacy event names. - for key in ("pre_tool_call", "pre_tool_use"): + # Check current and legacy event names. + for key in ( + "pre_tool_call", + "post_tool_call", + "on_session_end", + "pre_tool_use", + "post_tool_use", + "session_end", + ): entries = hooks.get(key) if not isinstance(entries, list): continue diff --git a/Gradata/src/gradata/hooks/adapters/opencode.py b/Gradata/src/gradata/hooks/adapters/opencode.py index 0a51789d..1deea5fb 100644 --- a/Gradata/src/gradata/hooks/adapters/opencode.py +++ b/Gradata/src/gradata/hooks/adapters/opencode.py @@ -12,7 +12,9 @@ failure, hook_command, hook_signature, + post_tool_hook_command, read_json, + session_end_hook_command, write_json, ) @@ -49,26 +51,70 @@ def install(brain_dir: Path, agent_config_path: Path) -> InstallResult: sig = hook_signature(AGENT, brain_dir) data = read_json(agent_config_path) hooks = data.setdefault("hooks", {}) - pre_tool = hooks.setdefault("preTool", []) - if any(sig in str(item) for item in pre_tool): + added: list[str] = [] + for key, command in ( + ("preTool", hook_command(brain_dir)), + ("postTool", post_tool_hook_command(brain_dir)), + ("sessionEnd", session_end_hook_command(brain_dir)), + ): + entries = hooks.setdefault(key, []) + if not isinstance(entries, list): + entries = [] + hooks[key] = entries + if any(sig in str(item) for item in entries): + continue + entries.append({"id": sig, "command": command}) + added.append(key) + if not added: return InstallResult( - AGENT, agent_config_path, "already_present", "hook already present" + AGENT, agent_config_path, "already_present", "hooks already present" ) - pre_tool.append({"id": sig, "command": hook_command(brain_dir)}) write_json(agent_config_path, data) - return InstallResult(AGENT, agent_config_path, "added", "installed preTool hook") + return InstallResult( + AGENT, + agent_config_path, + "added", + "installed preTool, postTool, and sessionEnd hooks", + ) except Exception as exc: return failure(AGENT, agent_config_path, exc) def uninstall(brain_dir: Path, agent_config_path: Path) -> InstallResult: - """Reverse install: drop signature-matching entries from hooks.preTool.""" - from gradata.hooks.adapters._base import uninstall_from_list_in_dict - - return uninstall_from_list_in_dict( - agent=AGENT, - brain_dir=brain_dir, - agent_config_path=agent_config_path, - outer_key="hooks", - inner_key="preTool", - ) + """Reverse install: drop signature-matching entries from OpenCode hooks.""" + try: + if not agent_config_path.is_file(): + return InstallResult( + AGENT, agent_config_path, "already_present", "config file does not exist" + ) + sig = hook_signature(AGENT, brain_dir) + data = read_json(agent_config_path) + hooks = data.get("hooks") + if not isinstance(hooks, dict): + return InstallResult(AGENT, agent_config_path, "already_present", "no hooks block") + + removed = 0 + for key in ("preTool", "postTool", "sessionEnd"): + entries = hooks.get(key) + if not isinstance(entries, list): + continue + kept = [] + for entry in entries: + if sig in str(entry): + removed += 1 + continue + kept.append(entry) + if kept: + hooks[key] = kept + else: + hooks.pop(key, None) + + if removed == 0: + return InstallResult(AGENT, agent_config_path, "already_present", "hook not present") + + if not hooks: + data.pop("hooks", None) + write_json(agent_config_path, data) + return InstallResult(AGENT, agent_config_path, "removed", f"removed {removed} hook entry") + except Exception as exc: + return failure(AGENT, agent_config_path, exc) diff --git a/Gradata/tests/test_hook_adapters.py b/Gradata/tests/test_hook_adapters.py index bccdb3bd..64b29fc5 100644 --- a/Gradata/tests/test_hook_adapters.py +++ b/Gradata/tests/test_hook_adapters.py @@ -52,6 +52,76 @@ def test_codex_adapter_writes_valid_toml_with_quoted_brain_path(tmp_path: Path) assert brain_dir.as_posix() in hook["id"] +def test_codex_adapter_installs_post_tool_and_session_end_hooks(tmp_path: Path) -> None: + brain_dir = tmp_path / "brain" + brain_dir.mkdir() + config_path = adapter_config_path("codex") + + result = get_adapter("codex").install(brain_dir, config_path) + + assert result.action == "added" + parsed = tomllib.loads(config_path.read_text(encoding="utf-8")) + hooks = parsed["hooks"] + assert "gradata.hooks.inject_brain_rules" in hooks["pre_tool"][0]["command"] + assert "gradata.hooks.auto_correct" in hooks["post_tool"][0]["command"] + assert "gradata.hooks.session_close" in hooks["session_end"][0]["command"] + + +@pytest.mark.parametrize( + ("agent", "expected_hooks"), + [ + ( + "hermes", + { + "pre_tool_call": "gradata.hooks.inject_brain_rules", + "post_tool_call": "gradata.hooks.auto_correct", + "on_session_end": "gradata.hooks.session_close", + }, + ), + ( + "opencode", + { + "preTool": "gradata.hooks.inject_brain_rules", + "postTool": "gradata.hooks.auto_correct", + "sessionEnd": "gradata.hooks.session_close", + }, + ), + ], +) +def test_json_and_yaml_adapters_install_post_tool_and_session_end_hooks( + tmp_path: Path, agent: str, expected_hooks: dict[str, str] +) -> None: + brain_dir = tmp_path / "brain" + brain_dir.mkdir() + config_path = adapter_config_path(agent) + + result = get_adapter(agent).install(brain_dir, config_path) + + assert result.action == "added" + text = config_path.read_text(encoding="utf-8") + for hook_name, module in expected_hooks.items(): + assert hook_name in text + assert module in text + + +@pytest.mark.parametrize("agent", ["codex", "hermes", "opencode"]) +def test_adapter_uninstall_removes_post_tool_and_session_end_hooks( + tmp_path: Path, agent: str +) -> None: + brain_dir = tmp_path / "brain" + brain_dir.mkdir() + adapter = get_adapter(agent) + config_path = adapter_config_path(agent) + + assert adapter.install(brain_dir, config_path).action == "added" + assert adapter.uninstall(brain_dir, config_path).action == "removed" + + text = config_path.read_text(encoding="utf-8") if config_path.exists() else "" + assert "gradata.hooks.auto_correct" not in text + assert "gradata.hooks.session_close" not in text + assert f"gradata:{agent}:" not in text + + def test_adapter_install_does_not_touch_real_user_config(tmp_path: Path) -> None: real_config = _REAL_HOME / ".codex" / "config.toml" before = real_config.read_text(encoding="utf-8") if real_config.exists() else None