[loop-cycle-50] refactor: replace bare sqlite3.connect() with context managers batch 2 (#157) (#180)

src/infrastructure/models/registry.py

@@ -11,7 +11,8 @@ model roles (student, teacher, judge/PRM) run on dedicated resources.
 import logging
 import sqlite3
 import threading
-from contextlib import closing
+from collections.abc import Generator
+from contextlib import closing, contextmanager
 from dataclasses import dataclass
 from datetime import UTC, datetime
 from enum import StrEnum
@@ -61,36 +62,37 @@ class CustomModel:
             self.registered_at = datetime.now(UTC).isoformat()
 
 
-def _get_conn() -> sqlite3.Connection:
+@contextmanager
+def _get_conn() -> Generator[sqlite3.Connection, None, None]:
     DB_PATH.parent.mkdir(parents=True, exist_ok=True)
-    conn = sqlite3.connect(str(DB_PATH))
-    conn.row_factory = sqlite3.Row
-    conn.execute("PRAGMA journal_mode=WAL")
-    conn.execute("PRAGMA busy_timeout=5000")
-    conn.execute("""
-        CREATE TABLE IF NOT EXISTS custom_models (
-            name            TEXT PRIMARY KEY,
-            format          TEXT NOT NULL,
-            path            TEXT NOT NULL,
-            role            TEXT NOT NULL DEFAULT 'general',
-            context_window  INTEGER NOT NULL DEFAULT 4096,
-            description     TEXT NOT NULL DEFAULT '',
-            registered_at   TEXT NOT NULL,
-            active          INTEGER NOT NULL DEFAULT 1,
-            default_temperature REAL NOT NULL DEFAULT 0.7,
-            max_tokens      INTEGER NOT NULL DEFAULT 2048
-        )
-        """)
-    conn.execute("""
-        CREATE TABLE IF NOT EXISTS agent_model_assignments (
-            agent_id    TEXT PRIMARY KEY,
-            model_name  TEXT NOT NULL,
-            assigned_at TEXT NOT NULL,
-            FOREIGN KEY (model_name) REFERENCES custom_models(name)
-        )
-        """)
-    conn.commit()
-    return conn
+    with closing(sqlite3.connect(str(DB_PATH))) as conn:
+        conn.row_factory = sqlite3.Row
+        conn.execute("PRAGMA journal_mode=WAL")
+        conn.execute("PRAGMA busy_timeout=5000")
+        conn.execute("""
+            CREATE TABLE IF NOT EXISTS custom_models (
+                name            TEXT PRIMARY KEY,
+                format          TEXT NOT NULL,
+                path            TEXT NOT NULL,
+                role            TEXT NOT NULL DEFAULT 'general',
+                context_window  INTEGER NOT NULL DEFAULT 4096,
+                description     TEXT NOT NULL DEFAULT '',
+                registered_at   TEXT NOT NULL,
+                active          INTEGER NOT NULL DEFAULT 1,
+                default_temperature REAL NOT NULL DEFAULT 0.7,
+                max_tokens      INTEGER NOT NULL DEFAULT 2048
+            )
+            """)
+        conn.execute("""
+            CREATE TABLE IF NOT EXISTS agent_model_assignments (
+                agent_id    TEXT PRIMARY KEY,
+                model_name  TEXT NOT NULL,
+                assigned_at TEXT NOT NULL,
+                FOREIGN KEY (model_name) REFERENCES custom_models(name)
+            )
+            """)
+        conn.commit()
+        yield conn
 
 
 class ModelRegistry:
@@ -106,7 +108,7 @@ class ModelRegistry:
     def _load_from_db(self) -> None:
         """Bootstrap cache from SQLite."""
         try:
-            with closing(_get_conn()) as conn:
+            with _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"],
@@ -130,7 +132,7 @@ class ModelRegistry:
     def register(self, model: CustomModel) -> CustomModel:
         """Register a new custom model."""
         with self._lock:
-            with closing(_get_conn()) as conn:
+            with _get_conn() as conn:
                 conn.execute(
                     """
                     INSERT OR REPLACE INTO custom_models
@@ -161,7 +163,7 @@ class ModelRegistry:
         with self._lock:
             if name not in self._models:
                 return False
-            with closing(_get_conn()) as conn:
+            with _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()
@@ -191,7 +193,7 @@ class ModelRegistry:
             return False
         with self._lock:
             model.active = active
-            with closing(_get_conn()) as conn:
+            with _get_conn() as conn:
                 conn.execute(
                     "UPDATE custom_models SET active = ? WHERE name = ?",
                     (int(active), name),
@@ -207,7 +209,7 @@ class ModelRegistry:
             return False
         with self._lock:
             now = datetime.now(UTC).isoformat()
-            with closing(_get_conn()) as conn:
+            with _get_conn() as conn:
                 conn.execute(
                     """
                     INSERT OR REPLACE INTO agent_model_assignments
@@ -226,7 +228,7 @@ class ModelRegistry:
         with self._lock:
             if agent_id not in self._agent_assignments:
                 return False
-            with closing(_get_conn()) as conn:
+            with _get_conn() as conn:
                 conn.execute(
                     "DELETE FROM agent_model_assignments WHERE agent_id = ?",
                     (agent_id,),