fix: small pg related fixes for stability

application/app.py

@@ -65,7 +65,19 @@ ensure_database_ready(
 # Own the vector DB's schema here too, so the retrieval hot path is pure reads
 # instead of re-running DDL for every source of every request.
 if settings.AUTO_VECTOR_SCHEMA:
-    ensure_vector_schema(logger=logging.getLogger("application.app"))
+    _vector_schema_log = logging.getLogger("application.app")
+    try:
+        ensure_vector_schema(logger=_vector_schema_log)
+    except Exception:
+        # The vector DB is often a separate cluster. This runs at import time,
+        # so re-raising would stop gunicorn and every Celery worker from
+        # booting -- taking auth, chat history and webhooks down over a fault
+        # that only affects retrieval. PGVectorStore re-checks the schema on
+        # its write path, so degrading here loses nothing.
+        _vector_schema_log.exception(
+            "ensure_vector_schema failed; retrieval is degraded until the "
+            "vector database is reachable and its width matches EMBEDDINGS_NAME."
+        )
 
 from application.agents.default_tools import (  # noqa: E402
     validate_default_chat_tools,

application/graphrag/store.py

@@ -126,7 +126,11 @@ class GraphStore:
         ``PGVectorStore``, so a retrieval that touches both pays one checkout
         each instead of two fresh connect handshakes.
         """
-        if self._connection is None or self._connection.closed:
+        if self._connection is not None and self._connection.closed:
+            # Hand the dead connection back before replacing it; an unreturned
+            # checkout is a pool slot lost for the life of the process.
+            self.close()
+        if self._connection is None:
             if self._pool_max_size > 0:
                 self._connection = pgconn.pool_for(
                     self._connection_string, self._pool_max_size

application/sandbox/daytona.py

@@ -561,7 +561,7 @@ class DaytonaSandbox(CodeSandbox):
                 handle.sandbox.fs.create_folder(parent, "755", request_timeout=self._default_timeout)
             except Exception as folder_exc:  # noqa: BLE001 - folder may already exist
                 logger.debug("put_file parent folder create returned: %s", folder_exc)
-        handle.sandbox.fs.upload_file(data, remote)
+        handle.sandbox.fs.upload_file(data, remote, timeout=int(self._default_timeout))
 
     def get_file(self, session_id: str, path: str) -> bytes:
         """Download ``path`` from the session workspace as bytes, capped at ``max_file_bytes``."""
@@ -616,7 +616,9 @@ class DaytonaSandbox(CodeSandbox):
         size = getattr(info, "size", None)
         if size is not None and size > self._max_file_bytes:
             raise IOError(f"file too large: {size} > {self._max_file_bytes} bytes")
-        return handle.sandbox.fs.download_file(remote)
+        # download_file takes its timeout positionally and dispatches on
+        # isinstance(arg, int): a float would be read as a destination path.
+        return handle.sandbox.fs.download_file(remote, int(self._default_timeout))
 
     def list_files(self, session_id: str) -> List[str]:
         """List workspace-relative file paths for ``session_id`` (recursive, never escapes the workspace)."""

application/storage/db/bootstrap.py

@@ -153,7 +153,9 @@ def ensure_vector_schema(*, logger: Optional[logging.Logger] = None) -> None:
     started = time.monotonic()
     # A plain connection, never the store's pool: this can run pre-fork under
     # ``gunicorn --preload``, and an inherited pooled socket is a broken one.
-    conn = psycopg.connect(dsn)
+    # Bounded: a suspended/unreachable vector cluster must fail this hook
+    # rather than hang boot until a liveness probe kills the process.
+    conn = psycopg.connect(dsn, connect_timeout=10)
     try:
         cursor = conn.cursor()
         try:

application/vectorstore/pgvector.py

@@ -97,7 +97,12 @@ class PGVectorStore(BaseVectorStore):
 
     def _get_connection(self):
         """Get or create this store's connection, pooled unless pooling is off."""
-        if self._connection is None or self._connection.closed:
+        if self._connection is not None and self._connection.closed:
+            # Hand the dead connection back before replacing it: psycopg_pool
+            # never reclaims a checkout that is not returned, so dropping it
+            # costs the pool a slot for the life of the process.
+            self.close()
+        if self._connection is None:
             if self._pool_max_size > 0:
                 self._connection = _pool_for(
                     self._connection_string, self._pool_max_size

tests/graphrag/test_store.py

@@ -743,6 +743,29 @@ class TestGraphStorePooling:
         pooled_conn.rollback.assert_not_called()
         pool.putconn.assert_called_once_with(pooled_conn)
 
+    def test_a_dead_pooled_connection_is_returned_before_being_replaced(
+        self, monkeypatch
+    ):
+        # Same contract as ``PGVectorStore``: a connection that dies while this
+        # store holds it must go back to the pool, or the slot is lost for the
+        # life of the process. Extraction holds one store across the whole
+        # per-chunk LLM loop, which is exactly when a backend gets reaped.
+        pool, pooled_conn = self._fake_pool()
+        monkeypatch.setattr(store_module.pgconn, "pool_for", lambda dsn, n: pool)
+        monkeypatch.setitem(pgconn._POOLS, POOL_DSN, pool)
+        store = self._store()
+        store._get_connection()
+        replacement = MagicMock()
+        replacement.closed = False
+        pool.getconn.return_value = replacement
+
+        pooled_conn.closed = True
+        conn = store._get_connection()
+
+        assert conn is replacement
+        pool.putconn.assert_called_once_with(pooled_conn)
+        assert pool.getconn.call_count == 2
+
     def test_legacy_path_connects_directly_and_closes(self, monkeypatch):
         def _never(dsn, n):
             raise AssertionError("pooling is off; no pool must be built")

tests/sandbox/test_daytona.py

@@ -437,6 +437,28 @@ def test_put_file_uploads_under_workspace(sandbox):
     assert remote.endswith("/docsgpt-sandbox/out/data.csv")
 
 
+def test_put_file_bounds_the_upload(sandbox):
+    # Every other fs.* call is bounded at default_timeout; without this the
+    # upload falls back to the SDK's 30-minute default and a toolbox proxy that
+    # accepts but never answers pins the chat stream for that long.
+    sandbox.open("conv-1")
+    _, created = sandbox._client.created[0]
+    sandbox.put_file("conv-1", "out/data.csv", b"a,b\n")
+    assert created.fs.upload_file.call_args.kwargs["timeout"] == 60
+
+
+def test_get_file_bounds_the_download(sandbox):
+    # download_file takes its timeout positionally and dispatches on
+    # isinstance(arg, int) -- a float would be read as a local destination path.
+    sandbox.open("conv-1")
+    _, created = sandbox._client.created[0]
+    created.fs.download_file.return_value = b"hello"
+    sandbox.get_file("conv-1", "out/data.csv")
+    args = created.fs.download_file.call_args.args
+    assert args[1] == 60
+    assert isinstance(args[1], int)
+
+
 def test_put_file_creates_parent_dirs(sandbox):
     sandbox.open("conv-1")
     _, created = sandbox._client.created[0]

tests/storage/db/test_bootstrap_vector_schema.py

@@ -8,6 +8,7 @@ a ``documents`` table built for a different model silently retrieves garbage.
 
 from __future__ import annotations
 
+import ast
 from pathlib import Path
 from unittest.mock import MagicMock, patch
 
@@ -91,7 +92,11 @@ class TestEnsureVectorSchemaCreates:
     def test_creates_the_vector_schema_under_the_advisory_lock(self, vector_settings):
         connect, conn, cursor, vector_schema, graph_schema, _ = self._run()
 
-        connect.assert_called_once_with("postgresql://user:pass@localhost/db")
+        # A bounded connect: an unreachable or suspended vector DB must fail the
+        # hook, not hang boot until a liveness probe kills the process.
+        connect.assert_called_once_with(
+            "postgresql://user:pass@localhost/db", connect_timeout=10
+        )
         statements = " ".join(str(c) for c in cursor.execute.call_args_list)
         assert "pg_advisory_xact_lock" in statements
         vector_schema.assert_called_once()
@@ -187,6 +192,29 @@ class TestBootGating:
         assert "if settings.AUTO_VECTOR_SCHEMA:" in source
         assert "ensure_vector_schema(" in source
 
+    def test_the_boot_hook_cannot_take_the_process_down(self):
+        """A vector-DB fault must degrade retrieval, not crash-loop the app.
+
+        The hook runs at import time, so an exception escaping it stops
+        gunicorn and every Celery worker from booting -- taking auth, chat
+        history and webhooks down with retrieval. ``PGVectorStore`` re-checks
+        the schema on its write path, so failing soft here loses nothing.
+        """
+        tree = ast.parse(_APP_PY.read_text())
+
+        def _calls_hook(node) -> bool:
+            return any(
+                isinstance(n, ast.Call)
+                and isinstance(n.func, ast.Name)
+                and n.func.id == "ensure_vector_schema"
+                for n in ast.walk(node)
+            )
+
+        assert any(
+            isinstance(node, ast.Try) and node.handlers and _calls_hook(node)
+            for node in ast.walk(tree)
+        ), "ensure_vector_schema() at import time must be wrapped in try/except"
+
     def test_setting_defaults_on(self):
         from application.core.settings import Settings
 

tests/vectorstore/test_pgvector_schema_and_pool.py

@@ -290,6 +290,29 @@ class TestConnectionPooling:
         pooled_conn.rollback.assert_not_called()
         fake_pool.putconn.assert_called_once_with(pooled_conn)
 
+    def test_a_dead_pooled_connection_is_returned_before_being_replaced(
+        self, monkeypatch
+    ):
+        # psycopg_pool never reclaims a checkout that is not handed back, so
+        # replacing a connection that died in our hands would burn a pool slot
+        # permanently: after PGVECTOR_POOL_MAX_SIZE such events every checkout
+        # blocks for the pool timeout and then fails.
+        store, fake_pool, pooled_conn = self._pooled_store(monkeypatch)
+        monkeypatch.setitem(
+            pgvector_module._POOLS, store._connection_string, fake_pool
+        )
+        store._get_connection()
+        replacement = MagicMock()
+        replacement.closed = False
+        fake_pool.getconn.return_value = replacement
+
+        pooled_conn.closed = True  # server terminated the backend mid-idle
+        conn = store._get_connection()
+
+        assert conn is replacement
+        fake_pool.putconn.assert_called_once_with(pooled_conn)
+        assert fake_pool.getconn.call_count == 2
+
     def test_legacy_path_connects_directly_and_closes(self):
         store, mock_conn, _, _ = _make_store()
         store._pool_max_size = 0