fix: harden agent startup and websocket auth

agent/requirements.txt

@@ -34,7 +34,8 @@ openpyxl
 
 # DB for LangGraph checkpoint
 psycopg2-binary
-psycopg>=3.1
+# Bundle libpq with the agent environment; plain psycopg requires an OS-level libpq.
+psycopg[binary]>=3.1
 
 # Retry/utility
 tenacity>=8.0.0

agent/src/ss_tools/agent/langgraph_setup.py

@@ -9,7 +9,7 @@
 
 import inspect as _inspect
 import os
-from urllib.parse import urlsplit
+from urllib.parse import urlsplit, urlunsplit
 
 from langchain_openai import ChatOpenAI
 from langgraph.checkpoint.memory import InMemorySaver
@@ -56,20 +56,91 @@ _CHECKPOINTER_INIT = False
 _CHECKPOINTER_CONN = None
 
 
-# #region AgentChat.LangGraph.Setup.InitCheckpointer [C:3] [TYPE Function] [SEMANTICS agent-chat,langgraph,checkpointer,postgres]
+# #region AgentChat.LangGraph.Setup.RedactDbUrl [C:2] [TYPE Function] [SEMANTICS agent-chat,diagnostics,security]
 # @ingroup AgentChat
-# @BRIEF Initialize AsyncPostgresSaver from DATABASE_URL env var.
+# @BRIEF Redact password from a database URL for safe diagnostic logging.
+# @INVARIANT Never returns raw password or full credentials in the output string.
+def _redact_db_url(url: str) -> str:
+    """Return a copy of `url` with the password replaced by '***'."""
+    if not url:
+        return "<empty>"
+    try:
+        parsed = urlsplit(url)
+        if not parsed.scheme or not parsed.hostname:
+            return "<invalid-url>"
+        host = f"[{parsed.hostname}]" if ":" in parsed.hostname else parsed.hostname
+        port = f":{parsed.port}" if parsed.port else ""
+        user = f"{parsed.username}:***@" if parsed.username else ""
+        return urlunsplit((parsed.scheme, f"{user}{host}{port}", parsed.path, "", ""))
+    except Exception:
+        return "<unparseable-url>"
+# #endregion AgentChat.LangGraph.Setup.RedactDbUrl
+
+
+# #region AgentChat.LangGraph.Setup.InitCheckpointer [C:4] [TYPE Function] [SEMANTICS agent-chat,langgraph,checkpointer,postgres]
+# @ingroup AgentChat
+# @BRIEF Initialize AsyncPostgresSaver from DATABASE_URL env var with validation and diagnostics.
 # @SIDE_EFFECT Connects to PostgreSQL; creates checkpointer table via setup().
+# @INVARIANT DATABASE_URL must be a valid PostgreSQL URI before connection is attempted.
+# @RATIONALE Production agent crash-loop was caused by an empty or misconfigured DATABASE_URL
+#            that resolved to a nonexistent Unix socket. Pre-connection validation + redacted
+#            diagnostics (host, port, db name — never password) allows operators to debug
+#            configuration failures without exposing secrets in logs.
 async def init_checkpointer() -> None:
     global _CHECKPOINTER, _CHECKPOINTER_INIT, _CHECKPOINTER_CONN
     if _CHECKPOINTER_INIT:
         return
     db_url = os.getenv("DATABASE_URL")
-    pg_url = db_url.replace("postgresql+psycopg2://", "postgres://").replace("postgresql://", "postgres://")
-    _CHECKPOINTER_CONN = await psycopg.AsyncConnection.connect(pg_url, autocommit=True, row_factory=dict_row)
-    _CHECKPOINTER = AsyncPostgresSaver(_CHECKPOINTER_CONN)
-    await _CHECKPOINTER.setup()
+    if not db_url or not db_url.strip():
+        logger.explore(
+            "DATABASE_URL env var is missing or empty — checkpointer cannot be initialized",
+            payload={"env_var": "DATABASE_URL"},
+            error="DATABASE_URL is not set",
+        )
+        raise RuntimeError("DATABASE_URL is not set. PostgreSQL checkpointer requires a valid database URI.")
+
+    # Redact password from URL for safe diagnostic logging.
+    _sanitized = _redact_db_url(db_url)
+    logger.reason(
+        "Initializing PostgreSQL checkpointer",
+        payload={"sanitized_url": _sanitized},
+    )
+
+    pg_url: str = db_url.replace("postgresql+psycopg2://", "postgres://").replace("postgresql://", "postgres://")
+    if not pg_url.startswith("postgres://") and not pg_url.startswith("postgresql://"):
+        logger.explore(
+            "DATABASE_URL does not look like a PostgreSQL connection string",
+            payload={"sanitized_url": _sanitized, "parsed_scheme": pg_url.split("://")[0] if "://" in pg_url else "none"},
+            error="Not a PostgreSQL URL",
+        )
+        raise RuntimeError(f"DATABASE_URL must be a PostgreSQL URI. Got invalid scheme.")
+
+    try:
+        _CHECKPOINTER_CONN = await psycopg.AsyncConnection.connect(pg_url, autocommit=True, row_factory=dict_row)
+    except Exception as e:
+        logger.explore(
+            "Failed to connect to PostgreSQL checkpointer",
+            payload={"sanitized_url": _sanitized, "exception_type": type(e).__name__},
+            error="PostgreSQL connection failed; inspect database service availability and credentials",
+        )
+        raise
+
+    try:
+        _CHECKPOINTER = AsyncPostgresSaver(_CHECKPOINTER_CONN)
+        await _CHECKPOINTER.setup()
+    except Exception as e:
+        logger.explore(
+            "Checkpointer setup() failed",
+            payload={"sanitized_url": _sanitized, "exception_type": type(e).__name__},
+            error="PostgreSQL checkpointer setup failed; inspect database schema permissions",
+        )
+        raise
+
     _CHECKPOINTER_INIT = True
+    logger.reason(
+        "PostgreSQL checkpointer initialized successfully",
+        payload={"sanitized_url": _sanitized},
+    )
 # #endregion AgentChat.LangGraph.Setup.InitCheckpointer
 
 _llm_config: dict | None = None

agent/tests/test_agent/test_langgraph_setup.py

@@ -64,6 +64,44 @@ def test_llm_diagnostics_redacts_api_key_and_path():
 # #endregion Test.AgentChat.TestLlmDiagnostics
 
 
+# #region Test.AgentChat.TestCheckpointerDiagnostics [C:3] [TYPE Class] [SEMANTICS test,agent,checkpointer,security]
+# @BRIEF Verify checkpointer diagnostics redact credentials and reject invalid configuration safely.
+# @RELATION BINDS_TO -> [AgentChat.LangGraph.Setup.InitCheckpointer]
+class TestCheckpointerDiagnostics:
+    def test_redact_db_url_removes_password_query_and_fragment(self):
+        from ss_tools.agent.langgraph_setup import _redact_db_url
+
+        redacted = _redact_db_url("postgresql://user:secret@[::1]:5432/ss_tools?sslpassword=hidden#fragment")
+        assert redacted == "postgresql://user:***@[::1]:5432/ss_tools"
+        assert "secret" not in redacted
+        assert "hidden" not in redacted
+
+    def test_redact_db_url_handles_unparseable_value(self):
+        from ss_tools.agent.langgraph_setup import _redact_db_url
+
+        assert _redact_db_url("not a database url") == "<invalid-url>"
+
+    @pytest.mark.anyio
+    async def test_init_checkpointer_rejects_missing_database_url(self, monkeypatch):
+        import ss_tools.agent.langgraph_setup as ls
+
+        monkeypatch.delenv("DATABASE_URL", raising=False)
+        ls._CHECKPOINTER_INIT = False
+        with pytest.raises(RuntimeError, match="DATABASE_URL is not set"):
+            await ls.init_checkpointer()
+
+    @pytest.mark.anyio
+    async def test_init_checkpointer_redacts_connection_failure(self, monkeypatch):
+        import ss_tools.agent.langgraph_setup as ls
+
+        monkeypatch.setenv("DATABASE_URL", "postgresql://agent:secret@db:5432/ss_tools?sslpassword=hidden")
+        ls._CHECKPOINTER_INIT = False
+        with patch("ss_tools.agent.langgraph_setup.psycopg.AsyncConnection.connect", new=AsyncMock(side_effect=Exception("secret"))):
+            with pytest.raises(Exception, match="secret"):
+                await ls.init_checkpointer()
+# #endregion Test.AgentChat.TestCheckpointerDiagnostics
+
+
 # #region Test.AgentChat.TestCreateAgent [C:2] [TYPE Function]
 # @BRIEF Test create_agent with various LLM config states.
 class TestCreateAgent: