diff --git a/autohooks/__main__.py b/autohooks/__main__.py index 776281d4..203134c9 100644 --- a/autohooks/__main__.py +++ b/autohooks/__main__.py @@ -31,6 +31,7 @@ import argparse import asyncio +import contextlib import json import logging import os @@ -194,100 +195,107 @@ def _close_ariel() -> None: asyncio.run(connection_manager.close_all()) - if ns.command == "daemon": - from autohooks.daemon import run_daemon - from autohooks.source import SqliteSource - - assert cfg.source is not None # checked above, before ariel was loaded - source = SqliteSource.from_config(cfg) - asyncio.run( - run_daemon( - cfg, - source, - mem, - graph, - rag, - max_iterations=1 if ns.once else None, - # S10: give the daemon a per-layer resolver for assistant-role - # persona dispatches. - resolve=lambda layer: resolve_layer(app, layer, cfg.user_id), + try: + if ns.command == "daemon": + from autohooks.daemon import run_daemon + from autohooks.source import SqliteSource + + assert cfg.source is not None # checked above, before ariel was loaded + source = SqliteSource.from_config(cfg) + asyncio.run( + run_daemon( + cfg, + source, + mem, + graph, + rag, + max_iterations=1 if ns.once else None, + # S10: give the daemon a per-layer resolver for assistant-role + # persona dispatches. + resolve=lambda layer: resolve_layer(app, layer, cfg.user_id), + ) ) - ) - _close_ariel() - return 0 - - if ns.command == "dispatch": - from hooks.external import dispatch_event - - try: - extra = json.loads(ns.payload) if ns.payload else {} - except json.JSONDecodeError: - logger.exception("invalid --payload JSON") - return 2 - if not isinstance(extra, dict): - logger.error("--payload must be a JSON object") - return 2 - since, until = _dispatch_window(ns) - d_layer = _dispatch_layer(cfg, extra) - if d_layer != cfg.layer: - mem, graph, rag = resolve_layer(app, d_layer, cfg.user_id) - result = asyncio.run( - dispatch_event( - ns.event, - d_layer, - cfg.user_id, - {"since": since, "until": until, "persona_owner": cfg.persona_owner, **extra}, - mem, - graph, - rag, + return 0 + + if ns.command == "dispatch": + from hooks.external import dispatch_event + + try: + extra = json.loads(ns.payload) if ns.payload else {} + except json.JSONDecodeError: + logger.exception("invalid --payload JSON") + return 2 + if not isinstance(extra, dict): + logger.error("--payload must be a JSON object") + return 2 + since, until = _dispatch_window(ns) + d_layer = _dispatch_layer(cfg, extra) + if d_layer != cfg.layer: + mem, graph, rag = resolve_layer(app, d_layer, cfg.user_id) + result = asyncio.run( + dispatch_event( + ns.event, + d_layer, + cfg.user_id, + {"since": since, "until": until, "persona_owner": cfg.persona_owner, **extra}, + mem, + graph, + rag, + ) ) - ) - _close_ariel() - print(json.dumps(result, ensure_ascii=False)) + print(json.dumps(result, ensure_ascii=False)) + return 0 + + if ns.command == "recall": + from features.recall import recall_protocol + + blocks = asyncio.run(recall_protocol(mem, rag, cfg.user_id, query=ns.query, budget=int(ns.budget))) + if ns.format == "json": + print(json.dumps({"blocks": blocks}, ensure_ascii=False)) + else: + print(_render_recall_md(blocks)) + return 0 + + if ns.command == "context": + # S18 tail context_assembly: pre-response assembly (top-5 relevant + + # recent L1) — the adapter calls it BEFORE the agent answers and + # inserts the output. + from autohooks.context import assemble_context, render_context_json, render_context_md + + assembly = asyncio.run(assemble_context(mem, rag, cfg.user_id, text=ns.text, budget=int(ns.budget))) + if ns.format == "json": + print(render_context_json(assembly)) + else: + print(render_context_md(assembly)) + return 0 + + if ns.command == "recap": + from features.continuity import session_recap + + blocks = asyncio.run(session_recap(mem, cfg.user_id, budget=int(ns.budget))) + if ns.format == "json": + print(json.dumps({"blocks": blocks}, ensure_ascii=False)) + else: + print(_render_recall_md(blocks)) + return 0 + + from autohooks.inject import run_inject + + out = asyncio.run(run_inject(cfg, mem, graph, rag, text=ns.text, fmt=ns.format, blocks=ns.blocks, with_recap=ns.with_recap)) + print(out) return 0 - - if ns.command == "recall": - from features.recall import recall_protocol - - blocks = asyncio.run(recall_protocol(mem, rag, cfg.user_id, query=ns.query, budget=int(ns.budget))) + except BaseException: + # Close on the failure path too. `_close_ariel()` used to sit after each + # `asyncio.run(...)`, so an exception skipped it; aiosqlite's worker + # threads are non-daemon, `threading._shutdown` joins them at exit, and + # the process then prints a traceback and never returns — systemd keeps + # the unit `active`, so Restart=on-failure never fires. Suppressed here + # only so a failing close cannot replace the command's own exception. + with contextlib.suppress(Exception): + _close_ariel() + raise + else: _close_ariel() - if ns.format == "json": - print(json.dumps({"blocks": blocks}, ensure_ascii=False)) - else: - print(_render_recall_md(blocks)) - return 0 - - if ns.command == "context": - # S18 tail context_assembly: pre-response assembly (top-5 relevant + - # recent L1) — the adapter calls it BEFORE the agent answers and - # inserts the output. - from autohooks.context import assemble_context, render_context_json, render_context_md - - assembly = asyncio.run(assemble_context(mem, rag, cfg.user_id, text=ns.text, budget=int(ns.budget))) - _close_ariel() - if ns.format == "json": - print(render_context_json(assembly)) - else: - print(render_context_md(assembly)) - return 0 - - if ns.command == "recap": - from features.continuity import session_recap - - blocks = asyncio.run(session_recap(mem, cfg.user_id, budget=int(ns.budget))) - _close_ariel() - if ns.format == "json": - print(json.dumps({"blocks": blocks}, ensure_ascii=False)) - else: - print(_render_recall_md(blocks)) - return 0 - - from autohooks.inject import run_inject - - out = asyncio.run(run_inject(cfg, mem, graph, rag, text=ns.text, fmt=ns.format, blocks=ns.blocks, with_recap=ns.with_recap)) - _close_ariel() - print(out) - return 0 def _render_recall_md(blocks: list[dict[str, Any]]) -> str: diff --git a/shared/connection.py b/shared/connection.py index 51fcaf7b..8f8e87dc 100644 --- a/shared/connection.py +++ b/shared/connection.py @@ -390,12 +390,21 @@ def _connect() -> sqlite3.Connection: async def close_all(self) -> None: """Close all open connections (on shutdown).""" for name, conn in self._conns.items(): + # Two SEPARATE suppressions on purpose (F97). Sharing one meant a + # failing `PRAGMA optimize` — a locked or half-dead handle is enough + # — skipped the `close()` below it, and an unclosed aiosqlite + # connection keeps a NON-daemon worker thread, which + # `threading._shutdown` joins forever at interpreter exit. That is + # the same hang `main()` now prevents on the command path; closing + # here must not reintroduce it. The optimize is an optimisation; the + # close is the contract. with contextlib.suppress(Exception): # PRAGMA optimize lets SQLite run ANALYZE if the query planner # needs it — cheap, recommended before closing a connection. await conn.execute("PRAGMA optimize") + with contextlib.suppress(Exception): await conn.close() - logger.debug("closed connection %s", name) + logger.debug("closed connection %s", name) self._conns.clear() self._paths.clear() # The write dispatcher owns a connector of its own, built on this manager. diff --git a/tests/test_autohooks_cli.py b/tests/test_autohooks_cli.py index 46cc74a6..b17acd0a 100644 --- a/tests/test_autohooks_cli.py +++ b/tests/test_autohooks_cli.py @@ -95,3 +95,41 @@ def test_main_dispatch_invalid_payload_exits_2(tmp_path, monkeypatch): ], ) assert cli.main() == 2 + + +def test_main_closes_connections_when_command_raises(tmp_path, monkeypatch): + """A failing command must still close ariel's connections. + + The close used to sit on the line after `asyncio.run(...)`, so an exception + from the command skipped it; the non-daemon aiosqlite workers then blocked + interpreter shutdown (systemd kept the unit `active`, so Restart=on-failure + never fired). Close must run on the exception path too, and the original + exception must still reach the caller. + """ + from autohooks import appctx + from shared.connection import connection_manager + + cfg = tmp_path / "agent.yaml" + cfg.write_text(f"data_dir: {tmp_path}\nuser_id: default\nlayer: user\n", encoding="utf-8") + # The exit watchdog is a process-wide faulthandler timer that would fire + # mid-suite 120 s later; these tests do not hang, so do not arm it. + monkeypatch.setenv("ARIEL_CLI_GUARD", "0") + + closed: list[str] = [] + + async def fake_close_all() -> None: + closed.append("closed") + + monkeypatch.setattr(connection_manager, "close_all", fake_close_all) + monkeypatch.setattr(appctx, "build_app_context", lambda: object()) + monkeypatch.setattr(appctx, "resolve_layer", lambda app, layer, uid: (None, None, None)) + + async def boom(*args, **kwargs): + raise RuntimeError("boom: command failed mid-run") + + monkeypatch.setattr("features.recall.recall_protocol", boom) + + with pytest.raises(RuntimeError, match="boom: command failed mid-run"): + cli.main(["recall", "--config", str(cfg)]) + + assert closed == ["closed"] diff --git a/tests/test_shared/test_connection.py b/tests/test_shared/test_connection.py index 5454098c..ec2211bf 100644 --- a/tests/test_shared/test_connection.py +++ b/tests/test_shared/test_connection.py @@ -155,3 +155,38 @@ async def t(): assert row[0] == 1 asyncio.run(t()) + + +def test_close_all_closes_even_if_optimize_fails(tmp_path): + """F97 tail: `PRAGMA optimize` and `close()` must not share one suppress. + + `close_all` used both inside a single `contextlib.suppress(Exception)`. If the + optimize raised — a locked or half-dead handle is enough — the `close()` on + the next line never ran, the aiosqlite worker thread (which is NON-daemon) + stayed alive, and the interpreter hung at exit with the unit still "active" + under systemd. The close is the load-bearing part; the optimize is not. + """ + from shared.connection import AsyncConnectionManager + + class _Stubborn: + """A connection whose PRAGMA raises; close() must still be reached.""" + + def __init__(self) -> None: + self.closed = False + + async def execute(self, sql: str, *a: object) -> None: + if "optimize" in sql: + raise RuntimeError("database is locked") + + async def close(self) -> None: + self.closed = True + + async def t() -> None: + cm = AsyncConnectionManager(base_dir=str(tmp_path)) + conn = _Stubborn() + cm._conns["memory.db"] = conn # type: ignore[assignment] + await cm.close_all() + assert cm._conns == {}, "the registry is cleared regardless" + assert conn.closed is True, "close() must run even though PRAGMA optimize raised" + + asyncio.run(t())