adding sqlite3 queue

Author: CamDavidsonPilon

src/mcp_utils/queue.py

@@ -3,6 +3,8 @@
 """
 
 import logging
+import sqlite3
+import time
 from typing import Protocol
 
 import msgspec
@@ -108,3 +110,105 @@ def clear_session(self, session_id: str) -> None:
         queue_key = self._get_queue_key(session_id)
         self.redis.delete(queue_key)
         logger.debug(f"Redis: Clearing session: {session_id}")
+
+
+class SQLiteResponseQueue(ResponseQueueProtocol):
+    """
+    A SQLite-backed queue implementation for MCP responses.
+    Each session is a logical queue keyed by `session_id`.
+
+    Uses a simple table and transactional pop to ensure single-delivery.
+    """
+
+    def __init__(self, db_path: str = ":memory:"):
+        """Initialize and create tables if needed.
+
+        Args:
+            db_path: Path to the SQLite database file. Defaults to in-memory.
+        """
+        # isolation_level=None -> autocommit mode so we can manage BEGIN/COMMIT explicitly
+        self.conn = sqlite3.connect(
+            db_path, timeout=30, check_same_thread=False, isolation_level=None
+        )
+        self.conn.execute(
+            """
+            CREATE TABLE IF NOT EXISTS response_queue (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                session_id TEXT NOT NULL,
+                payload TEXT NOT NULL,
+                created_at REAL DEFAULT (strftime('%s','now'))
+            )
+            """
+        )
+        self.conn.execute(
+            """
+            CREATE INDEX IF NOT EXISTS idx_response_queue_session_id
+            ON response_queue(session_id, id)
+            """
+        )
+
+    def push_response(self, session_id: str, response: MCPResponse) -> None:
+        """Insert a response payload for a session."""
+        payload = msgspec.json.encode(response).decode()
+        logger.debug(f"SQLite: Saving response for session: {session_id}: {payload}")
+        self.conn.execute(
+            "INSERT INTO response_queue (session_id, payload) VALUES (?, ?)",
+            (session_id, payload),
+        )
+
+    def _pop_one(self, session_id: str) -> str | None:
+        """Atomically pop the oldest payload for the session, if any."""
+        try:
+            cur = self.conn.cursor()
+            cur.execute("BEGIN IMMEDIATE")
+            row = cur.execute(
+                "SELECT id, payload FROM response_queue WHERE session_id = ? ORDER BY id LIMIT 1",
+                (session_id,),
+            ).fetchone()
+            if not row:
+                cur.execute("ROLLBACK")
+                return None
+            row_id, payload = row
+            cur.execute("DELETE FROM response_queue WHERE id = ?", (row_id,))
+            cur.execute("COMMIT")
+            return payload
+        except sqlite3.Error as e:
+            logger.error(f"SQLite pop error: {e}")
+            try:
+                self.conn.execute("ROLLBACK")
+            except Exception:
+                pass
+            return None
+
+    def wait_for_response(self, session_id: str, timeout: float | None = None) -> str | None:
+        """
+        Wait for and pop the next response for a session.
+
+        If timeout is None, waits indefinitely using polling.
+        If timeout is 0, returns immediately if none available.
+        """
+        # Immediate, non-blocking attempt
+        payload = self._pop_one(session_id)
+        if payload is not None:
+            return payload
+
+        if timeout == 0:
+            return None
+
+        # Poll until available or timeout
+        start = time.time()
+        while True:
+            payload = self._pop_one(session_id)
+            if payload is not None:
+                return payload
+            if timeout is not None and (time.time() - start) >= timeout:
+                return None
+            time.sleep(0.1)
+
+    def clear_session(self, session_id: str) -> None:
+        """Remove all queued items for a session."""
+        logger.debug(f"SQLite: Clearing session: {session_id}")
+        self.conn.execute(
+            "DELETE FROM response_queue WHERE session_id = ?",
+            (session_id,),
+        )

tests/test_core.py

@@ -26,7 +26,7 @@
 )
 
 
-class TestResponseQueue(ResponseQueueProtocol):
+class DemoResponseQueue(ResponseQueueProtocol):
     """Mock response queue for testing."""
 
     def __init__(self) -> None:
@@ -58,13 +58,13 @@ def clear_session(self, session_id: str) -> None:
 
 
 @pytest.fixture
-def response_queue() -> TestResponseQueue:
+def response_queue() -> DemoResponseQueue:
     """Create a test response queue."""
-    return TestResponseQueue()
+    return DemoResponseQueue()
 
 
 @pytest.fixture
-def server(response_queue: TestResponseQueue) -> MCPServer:
+def server(response_queue: DemoResponseQueue) -> MCPServer:
     """Create a test MCP server."""
     return MCPServer(
         name="Test Server",

tests/test_queue_sqlite.py

@@ -0,0 +1,71 @@
+"""Tests for SQLiteResponseQueue implementation."""
+
+import json
+from typing import Any
+
+import pytest
+
+from mcp_utils.queue import SQLiteResponseQueue
+from mcp_utils.schema import MCPResponse
+
+
+@pytest.fixture()
+def q(tmp_path) -> SQLiteResponseQueue:
+    return SQLiteResponseQueue(db_path=str(tmp_path / "queue.db"))
+
+
+def build_response(resp_id: str | int, result: dict[str, Any] | None = None) -> MCPResponse:
+    return MCPResponse(jsonrpc="2.0", id=resp_id, result=result or {})
+
+
+def test_push_and_pop_single(q: SQLiteResponseQueue) -> None:
+    session = "s1"
+    resp = build_response("1", {"x": 1})
+    q.push_response(session, resp)
+
+    data = q.wait_for_response(session, timeout=0)
+    assert data is not None
+    payload = json.loads(data)
+    assert payload["jsonrpc"] == "2.0"
+    assert payload["id"] == "1"
+    assert payload["result"] == {"x": 1}
+
+
+def test_timeout_behavior(q: SQLiteResponseQueue) -> None:
+    session = "s-timeout"
+    assert q.wait_for_response(session, timeout=0) is None
+    assert q.wait_for_response(session, timeout=0.1) is None
+
+
+def test_clear_session(q: SQLiteResponseQueue) -> None:
+    session = "s-clear"
+    q.push_response(session, build_response("1"))
+    q.push_response(session, build_response("2"))
+    q.clear_session(session)
+    assert q.wait_for_response(session, timeout=0) is None
+
+
+def test_fifo_ordering(q: SQLiteResponseQueue) -> None:
+    session = "s-fifo"
+    q.push_response(session, build_response("1", {"n": 1}))
+    q.push_response(session, build_response("2", {"n": 2}))
+
+    d1 = json.loads(q.wait_for_response(session, timeout=0))
+    d2 = json.loads(q.wait_for_response(session, timeout=0))
+    assert d1["id"] == "1"
+    assert d2["id"] == "2"
+
+
+def test_isolation_between_sessions(q: SQLiteResponseQueue) -> None:
+    s1, s2 = "s-a", "s-b"
+    q.push_response(s1, build_response("a1"))
+    q.push_response(s2, build_response("b1"))
+
+    d1 = json.loads(q.wait_for_response(s1, timeout=0))
+    assert d1["id"] == "a1"
+    # s1 should be empty now
+    assert q.wait_for_response(s1, timeout=0) is None
+
+    d2 = json.loads(q.wait_for_response(s2, timeout=0))
+    assert d2["id"] == "b1"
+