fix(autohooks): close ariel's connections on the failure path too

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:

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.

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"]

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())