[loop-cycle-47] refactor: replace bare sqlite3.connect() with context managers (#148) (#155)

src/infrastructure/events/bus.py

@@ -10,6 +10,7 @@ import json
 import logging
 import sqlite3
 from collections.abc import Callable, Coroutine
+from contextlib import closing
 from dataclasses import dataclass, field
 from datetime import UTC, datetime
 from pathlib import Path
@@ -99,14 +100,11 @@ class EventBus:
         if self._persistence_db_path is None:
             return
         self._persistence_db_path.parent.mkdir(parents=True, exist_ok=True)
-        conn = sqlite3.connect(str(self._persistence_db_path))
-        try:
+        with closing(sqlite3.connect(str(self._persistence_db_path))) as conn:
             conn.execute("PRAGMA journal_mode=WAL")
             conn.execute("PRAGMA busy_timeout=5000")
             conn.executescript(_EVENTS_SCHEMA)
             conn.commit()
-        finally:
-            conn.close()
 
     def _get_persistence_conn(self) -> sqlite3.Connection | None:
         """Get a connection to the persistence database."""
@@ -123,27 +121,26 @@ class EventBus:
         if conn is None:
             return
         try:
-            task_id = event.data.get("task_id", "")
-            agent_id = event.data.get("agent_id", "")
-            conn.execute(
-                "INSERT OR IGNORE INTO events "
-                "(id, event_type, source, task_id, agent_id, data, timestamp) "
-                "VALUES (?, ?, ?, ?, ?, ?, ?)",
-                (
-                    event.id,
-                    event.type,
-                    event.source,
-                    task_id,
-                    agent_id,
-                    json.dumps(event.data),
-                    event.timestamp,
-                ),
-            )
-            conn.commit()
+            with closing(conn):
+                task_id = event.data.get("task_id", "")
+                agent_id = event.data.get("agent_id", "")
+                conn.execute(
+                    "INSERT OR IGNORE INTO events "
+                    "(id, event_type, source, task_id, agent_id, data, timestamp) "
+                    "VALUES (?, ?, ?, ?, ?, ?, ?)",
+                    (
+                        event.id,
+                        event.type,
+                        event.source,
+                        task_id,
+                        agent_id,
+                        json.dumps(event.data),
+                        event.timestamp,
+                    ),
+                )
+                conn.commit()
         except Exception as exc:
             logger.debug("Failed to persist event: %s", exc)
-        finally:
-            conn.close()
 
     # ── Replay ───────────────────────────────────────────────────────────
 
@@ -170,40 +167,39 @@ class EventBus:
             return []
 
         try:
-            conditions = []
-            params: list = []
+            with closing(conn):
+                conditions = []
+                params: list = []
 
-            if event_type:
-                conditions.append("event_type = ?")
-                params.append(event_type)
-            if source:
-                conditions.append("source = ?")
-                params.append(source)
-            if task_id:
-                conditions.append("task_id = ?")
-                params.append(task_id)
+                if event_type:
+                    conditions.append("event_type = ?")
+                    params.append(event_type)
+                if source:
+                    conditions.append("source = ?")
+                    params.append(source)
+                if task_id:
+                    conditions.append("task_id = ?")
+                    params.append(task_id)
 
-            where = " AND ".join(conditions) if conditions else "1=1"
-            sql = f"SELECT * FROM events WHERE {where} ORDER BY timestamp DESC LIMIT ?"
-            params.append(limit)
+                where = " AND ".join(conditions) if conditions else "1=1"
+                sql = f"SELECT * FROM events WHERE {where} ORDER BY timestamp DESC LIMIT ?"
+                params.append(limit)
 
-            rows = conn.execute(sql, params).fetchall()
+                rows = conn.execute(sql, params).fetchall()
 
-            return [
-                Event(
-                    id=row["id"],
-                    type=row["event_type"],
-                    source=row["source"],
-                    data=json.loads(row["data"]) if row["data"] else {},
-                    timestamp=row["timestamp"],
-                )
-                for row in rows
-            ]
+                return [
+                    Event(
+                        id=row["id"],
+                        type=row["event_type"],
+                        source=row["source"],
+                        data=json.loads(row["data"]) if row["data"] else {},
+                        timestamp=row["timestamp"],
+                    )
+                    for row in rows
+                ]
         except Exception as exc:
             logger.debug("Failed to replay events: %s", exc)
             return []
-        finally:
-            conn.close()
 
     # ── Subscribe / Publish ──────────────────────────────────────────────
 

src/infrastructure/models/registry.py

@@ -11,6 +11,7 @@ model roles (student, teacher, judge/PRM) run on dedicated resources.
 import logging
 import sqlite3
 import threading
+from contextlib import closing
 from dataclasses import dataclass
 from datetime import UTC, datetime
 from enum import StrEnum
@@ -105,23 +106,22 @@ class ModelRegistry:
     def _load_from_db(self) -> None:
         """Bootstrap cache from SQLite."""
         try:
-            conn = _get_conn()
-            for row in conn.execute("SELECT * FROM custom_models WHERE active = 1").fetchall():
-                self._models[row["name"]] = CustomModel(
-                    name=row["name"],
-                    format=ModelFormat(row["format"]),
-                    path=row["path"],
-                    role=ModelRole(row["role"]),
-                    context_window=row["context_window"],
-                    description=row["description"],
-                    registered_at=row["registered_at"],
-                    active=bool(row["active"]),
-                    default_temperature=row["default_temperature"],
-                    max_tokens=row["max_tokens"],
-                )
-            for row in conn.execute("SELECT * FROM agent_model_assignments").fetchall():
-                self._agent_assignments[row["agent_id"]] = row["model_name"]
-            conn.close()
+            with closing(_get_conn()) as conn:
+                for row in conn.execute("SELECT * FROM custom_models WHERE active = 1").fetchall():
+                    self._models[row["name"]] = CustomModel(
+                        name=row["name"],
+                        format=ModelFormat(row["format"]),
+                        path=row["path"],
+                        role=ModelRole(row["role"]),
+                        context_window=row["context_window"],
+                        description=row["description"],
+                        registered_at=row["registered_at"],
+                        active=bool(row["active"]),
+                        default_temperature=row["default_temperature"],
+                        max_tokens=row["max_tokens"],
+                    )
+                for row in conn.execute("SELECT * FROM agent_model_assignments").fetchall():
+                    self._agent_assignments[row["agent_id"]] = row["model_name"]
         except Exception as exc:
             logger.warning("Failed to load model registry from DB: %s", exc)
 
@@ -130,29 +130,28 @@ class ModelRegistry:
     def register(self, model: CustomModel) -> CustomModel:
         """Register a new custom model."""
         with self._lock:
-            conn = _get_conn()
-            conn.execute(
-                """
-                INSERT OR REPLACE INTO custom_models
-                    (name, format, path, role, context_window, description,
-                     registered_at, active, default_temperature, max_tokens)
-                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
-                """,
-                (
-                    model.name,
-                    model.format.value,
-                    model.path,
-                    model.role.value,
-                    model.context_window,
-                    model.description,
-                    model.registered_at,
-                    int(model.active),
-                    model.default_temperature,
-                    model.max_tokens,
-                ),
-            )
-            conn.commit()
-            conn.close()
+            with closing(_get_conn()) as conn:
+                conn.execute(
+                    """
+                    INSERT OR REPLACE INTO custom_models
+                        (name, format, path, role, context_window, description,
+                         registered_at, active, default_temperature, max_tokens)
+                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                    """,
+                    (
+                        model.name,
+                        model.format.value,
+                        model.path,
+                        model.role.value,
+                        model.context_window,
+                        model.description,
+                        model.registered_at,
+                        int(model.active),
+                        model.default_temperature,
+                        model.max_tokens,
+                    ),
+                )
+                conn.commit()
             self._models[model.name] = model
             logger.info("Registered model: %s (%s)", model.name, model.format.value)
             return model
@@ -162,11 +161,10 @@ class ModelRegistry:
         with self._lock:
             if name not in self._models:
                 return False
-            conn = _get_conn()
-            conn.execute("DELETE FROM custom_models WHERE name = ?", (name,))
-            conn.execute("DELETE FROM agent_model_assignments WHERE model_name = ?", (name,))
-            conn.commit()
-            conn.close()
+            with closing(_get_conn()) as conn:
+                conn.execute("DELETE FROM custom_models WHERE name = ?", (name,))
+                conn.execute("DELETE FROM agent_model_assignments WHERE model_name = ?", (name,))
+                conn.commit()
             del self._models[name]
             # Remove any agent assignments using this model
             self._agent_assignments = {
@@ -193,13 +191,12 @@ class ModelRegistry:
             return False
         with self._lock:
             model.active = active
-            conn = _get_conn()
-            conn.execute(
-                "UPDATE custom_models SET active = ? WHERE name = ?",
-                (int(active), name),
-            )
-            conn.commit()
-            conn.close()
+            with closing(_get_conn()) as conn:
+                conn.execute(
+                    "UPDATE custom_models SET active = ? WHERE name = ?",
+                    (int(active), name),
+                )
+                conn.commit()
         return True
 
     # ── Agent-model assignments ────────────────────────────────────────────
@@ -210,17 +207,16 @@ class ModelRegistry:
             return False
         with self._lock:
             now = datetime.now(UTC).isoformat()
-            conn = _get_conn()
-            conn.execute(
-                """
-                INSERT OR REPLACE INTO agent_model_assignments
-                    (agent_id, model_name, assigned_at)
-                VALUES (?, ?, ?)
-                """,
-                (agent_id, model_name, now),
-            )
-            conn.commit()
-            conn.close()
+            with closing(_get_conn()) as conn:
+                conn.execute(
+                    """
+                    INSERT OR REPLACE INTO agent_model_assignments
+                        (agent_id, model_name, assigned_at)
+                    VALUES (?, ?, ?)
+                    """,
+                    (agent_id, model_name, now),
+                )
+                conn.commit()
             self._agent_assignments[agent_id] = model_name
             logger.info("Assigned model %s to agent %s", model_name, agent_id)
         return True
@@ -230,13 +226,12 @@ class ModelRegistry:
         with self._lock:
             if agent_id not in self._agent_assignments:
                 return False
-            conn = _get_conn()
-            conn.execute(
-                "DELETE FROM agent_model_assignments WHERE agent_id = ?",
-                (agent_id,),
-            )
-            conn.commit()
-            conn.close()
+            with closing(_get_conn()) as conn:
+                conn.execute(
+                    "DELETE FROM agent_model_assignments WHERE agent_id = ?",
+                    (agent_id,),
+                )
+                conn.commit()
             del self._agent_assignments[agent_id]
         return True