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:

backend/src/app.py

@@ -20,7 +20,6 @@ import asyncio
 from contextlib import asynccontextmanager
 import os
 from pathlib import Path
-import sys
 import uuid
 
 # project_root is used for static files mounting
@@ -669,28 +668,52 @@ def _authorize_websocket(websocket: WebSocket, resource: str, action: str) -> bo
     try:
         from .core.auth.jwt import decode_token
         from .core.database import SessionLocal
-        from .models.auth import User, Role
+        from .models.auth import User
 
         payload = decode_token(ws_token)
         username = payload.get("sub")
         if not isinstance(username, str) or not username:
+            logger.explore(
+                "WebSocket authorization — token missing 'sub' claim",
+                payload={"resource": resource, "action": action},
+                error="JWT payload has no valid subject",
+            )
             return False
 
         db = SessionLocal()
         try:
             user = db.query(User).filter(User.username == username).first()
             if not user:
+                logger.explore(
+                    "WebSocket authorization — user not found in database",
+                    payload={"resource": resource, "action": action, "username": username},
+                    error="Authenticated user missing from local DB",
+                )
                 return False
             if not getattr(user, "is_active", True):
+                logger.explore(
+                    "WebSocket authorization — user is inactive",
+                    payload={"resource": resource, "action": action, "username": username},
+                    error="User account is disabled",
+                )
                 return False
 
-            # Admin bypass via is_admin flag
+            # is_admin is the sole administrative authority. The database migration
+            # backfills the legacy Admin role before this authorization path runs.
             for role in user.roles:
                 if getattr(role, "is_admin", False):
                     return True
                 for perm in role.permissions:
                     if perm.resource == resource and perm.action == action:
                         return True
+
+            logger.explore(
+                "WebSocket authorization denied — no matching permission or admin role",
+                payload={"resource": resource, "action": action, "user": username,
+                         "roles": [r.name for r in user.roles],
+                         "is_active": getattr(user, "is_active", True)},
+                error="No matching RBAC permission",
+            )
             return False
         finally:
             db.close()

backend/src/dependencies.py

@@ -831,12 +831,8 @@ def has_permission(resource: str, action: str):
                 if perm.resource == resource and perm.action == action:
                     return current_user
 
-        # Special case for Admin role (full access) — uses is_admin flag, not name string.
-        # Fallback to role.name == "Admin" for roles created before is_admin migration.
-        if any(
-            getattr(role, "is_admin", False) or role.name == "Admin"
-            for role in current_user.roles
-        ):
+        # is_admin is the single source of truth for the administrative bypass.
+        if any(getattr(role, "is_admin", False) for role in current_user.roles):
             return current_user
 
         from .core.auth.logger import log_security_event

backend/src/scripts/create_admin.py

@@ -44,10 +44,18 @@ def create_admin(username, password, email=None):
             admin_role = db.query(Role).filter(Role.name == "Admin").first()
             if not admin_role:
                 logger.reason("Creating Admin role")
-                admin_role = Role(name="Admin", description="System Administrator")
+                admin_role = Role(
+                    name="Admin",
+                    description="System Administrator",
+                    is_admin=True,
+                )
                 db.add(admin_role)
                 db.commit()
                 db.refresh(admin_role)
+            elif not admin_role.is_admin:
+                logger.reason("Marking existing Admin role as administrative")
+                admin_role.is_admin = True
+                db.commit()
 
             # 2. Check if user already exists
             existing_user = db.query(User).filter(User.username == username).first()

backend/tests/api/test_storage.py

@@ -37,7 +37,7 @@ def _make_client(overrides: dict | None = None) -> tuple[TestClient, MagicMock]:
         id="admin-1", username="admin", email="admin@x.com",
         auth_source="LOCAL",
         created_at=__import__("datetime").datetime.now(),
-        roles=[RoleSchema(id="r1", name="Admin", description="", permissions=[])],
+        roles=[RoleSchema(id="r1", name="Admin", description="", is_admin=True, permissions=[])],
     )
 
     mock_plugin_loader = MagicMock()

backend/tests/test_app_ws_auth.py

@@ -9,11 +9,12 @@
 
 from pathlib import Path
 import sys
+from types import SimpleNamespace
 
 sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
 
 import pytest
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import MagicMock, patch
 
 
 class TestAuthenticateWebsocket:
@@ -117,4 +118,47 @@ class TestAuthenticateWebsocketApiKeyException:
             msl.return_value = db
             assert await _authenticate_websocket(ws, "ws/logs") is False
     # #endregion Test.AppModule.TestApikeyDbException
+
+
+# #region Test.AppModule.AuthorizeWebsocket [C:3] [TYPE Class] [SEMANTICS test,app,ws,authorization,rbac]
+# @BRIEF Verify WebSocket authorization grants access only through is_admin or explicit permissions.
+# @RELATION BINDS_TO -> [App.AppModule.AuthorizeWebsocket]
+class TestAuthorizeWebsocket:
+    def _authorize(self, user, resource="tasks", action="READ"):
+        from src.app import _authorize_websocket
+
+        ws = MagicMock()
+        ws.query_params = {"token": "valid.jwt"}
+        db = MagicMock()
+        db.query.return_value.filter.return_value.first.return_value = user
+        with (
+            patch("src.core.auth.jwt.decode_token", return_value={"sub": "testuser"}),
+            patch("src.core.database.SessionLocal", return_value=db),
+        ):
+            result = _authorize_websocket(ws, resource, action)
+        db.close.assert_called_once()
+        return result
+
+    def test_is_admin_role_bypasses_permission_check(self):
+        user = SimpleNamespace(
+            is_active=True,
+            roles=[SimpleNamespace(name="Admin", is_admin=True, permissions=[])],
+        )
+        assert self._authorize(user) is True
+
+    def test_legacy_admin_name_without_flag_does_not_bypass_permissions(self):
+        user = SimpleNamespace(
+            is_active=True,
+            roles=[SimpleNamespace(name="Admin", is_admin=False, permissions=[])],
+        )
+        assert self._authorize(user) is False
+
+    def test_explicit_permission_allows_websocket_access(self):
+        permission = SimpleNamespace(resource="tasks", action="READ")
+        user = SimpleNamespace(
+            is_active=True,
+            roles=[SimpleNamespace(name="Operator", is_admin=False, permissions=[permission])],
+        )
+        assert self._authorize(user) is True
+# #endregion Test.AppModule.AuthorizeWebsocket
 # #endregion Test.AppModule.WsAuth

frontend/src/lib/components/tasks/TaskResultPanel.svelte

@@ -1,13 +1,19 @@
-<!-- #region Tasks.TaskResultPanel [C:2] [TYPE Component] [SEMANTICS task, result, summary, status, plugin] -->
+<!-- #region Tasks.TaskResultPanel [C:3] [TYPE Component] [SEMANTICS task, result, summary, status, plugin] -->
 <!-- @ingroup Tasks -->
 <!-- @BRIEF Displays decision-oriented task outcome summaries with optional technical details. -->
 <!-- @LAYER UI -->
 <!-- @RELATION DEPENDS_ON -> [EXT:frontend:i18n] -->
+<!-- @RELATION CALLS -> [Services.StorageService.DownloadFileFunction] -->
 <!-- @UX_STATE Empty -> No task selected, placeholder shown. -->
 <!-- @UX_STATE Loaded -> Task result displayed with status color coding. -->
+<!-- @UX_STATE Downloading -> Selected artifact button disabled with busy state. -->
+<!-- @UX_RECOVERY Download error -> Toast notification with failure reason; user can retry via the same button. -->
 <script lang="ts">
   import { t } from '$lib/i18n/index.svelte.js';
+  import { log } from '$lib/cot-logger';
+  import { notifications } from '$lib/toasts.svelte.js';
   import { EmptyState } from '$lib/ui';
+  import { downloadFile } from '../../../services/storageService';
   import type { FailedDashboardEntry, MigrationTaskResult } from '$types/dashboard';
 
   let {
@@ -20,6 +26,7 @@
   const migrationResult = $derived(
     pluginId === 'superset-migration' ? result as MigrationTaskResult : null,
   );
+  let downloadingArchives = $state(new Set<string>());
 
   function phaseLabel(phase?: string): string {
     const labels: Record<string, string> = {
@@ -31,8 +38,25 @@
     return phase ? labels[phase] || phase : '';
   }
 
-  function archiveUrl(path: string): string {
-    return `/api/storage/download/migrations_failed/${path}`;
+  /**
+   * Authenticated download of a migration failure artifact.
+   * Uses bearer-token fetch (downloadFile from storageService) instead of raw <a href>
+   * which cannot send Authorization headers and returns 401.
+   */
+  async function downloadArchive(label: string, path: string): Promise<void> {
+    if (downloadingArchives.has(path)) return;
+    downloadingArchives = new Set([...downloadingArchives, path]);
+    log('Tasks.TaskResultPanel', 'REASON', 'Downloading migration artifact', { label, path });
+    try {
+      await downloadFile('migrations_failed', path);
+      log('Tasks.TaskResultPanel', 'REFLECT', 'Artifact downloaded', { label, path });
+    } catch (e: unknown) {
+      const msg = e instanceof Error ? e.message : 'Download failed';
+      log('Tasks.TaskResultPanel', 'EXPLORE', 'Artifact download failed', { label, path }, msg);
+      notifications.error(`${label}: ${msg}`);
+    } finally {
+      downloadingArchives = new Set([...downloadingArchives].filter((entry) => entry !== path));
+    }
   }
 
   const migrationStatusLabel = $derived(
@@ -182,8 +206,8 @@
               {/if}
               <div class="mt-3 flex flex-wrap items-center gap-x-4 gap-y-2">
                 <button type="button" class="text-xs font-medium text-primary hover:text-primary-hover" onclick={() => onshowrelatedlogs(failed)}>{$t.tasks?.show_related_logs || 'Show related logs'}</button>
-                {#if failed.source_archive_path}<a href={archiveUrl(failed.source_archive_path)} download class="text-xs font-medium text-primary hover:text-primary-hover">{$t.tasks?.download_source_archive || 'Download source export ZIP'}</a>{/if}
-                {#if failed.archive_path}<a href={archiveUrl(failed.archive_path)} download class="text-xs font-medium text-primary hover:text-primary-hover">{$t.tasks?.download_transformed_archive || 'Download transformed ZIP'}</a>{/if}
+                {#if failed.source_archive_path}<button type="button" class="text-xs font-medium text-primary hover:text-primary-hover disabled:cursor-not-allowed disabled:opacity-60" disabled={downloadingArchives.has(failed.source_archive_path)} aria-busy={downloadingArchives.has(failed.source_archive_path)} onclick={() => downloadArchive($t.tasks?.download_source_archive || 'Source ZIP', failed.source_archive_path)}>{$t.tasks?.download_source_archive || 'Download source export ZIP'}</button>{/if}
+                {#if failed.archive_path}<button type="button" class="text-xs font-medium text-primary hover:text-primary-hover disabled:cursor-not-allowed disabled:opacity-60" disabled={downloadingArchives.has(failed.archive_path)} aria-busy={downloadingArchives.has(failed.archive_path)} onclick={() => downloadArchive($t.tasks?.download_transformed_archive || 'Transformed ZIP', failed.archive_path)}>{$t.tasks?.download_transformed_archive || 'Download transformed ZIP'}</button>{/if}
               </div>
             </li>
           {/each}

frontend/src/lib/components/tasks/__tests__/TaskResultPanel.migration.test.ts

@@ -1,10 +1,20 @@
 // #region Test.Tasks.TaskResultPanel.Migration [C:2] [TYPE Module] [SEMANTICS test,tasks,migration,result]
 // @BRIEF Migration result panel shows successful and failed dashboard names plus Superset error text.
 // @RELATION BINDS_TO -> [Tasks.TaskResultPanel]
-import { describe, expect, it, vi } from 'vitest';
-import { render, screen } from '@testing-library/svelte';
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+import { fireEvent, render, screen, waitFor } from '@testing-library/svelte';
 import TaskResultPanel from '../TaskResultPanel.svelte';
 
+const { downloadFile, log, error } = vi.hoisted(() => ({
+  downloadFile: vi.fn(),
+  log: vi.fn(),
+  error: vi.fn(),
+}));
+
+vi.mock('../../../../services/storageService', () => ({ downloadFile }));
+vi.mock('$lib/cot-logger', () => ({ log, getTraceId: () => 'test-trace-id' }));
+vi.mock('$lib/toasts.svelte.js', () => ({ notifications: { error } }));
+
 vi.mock('$lib/i18n/index.svelte.js', () => ({
   t: {
     subscribe: (fn: (_v: Record<string, unknown>) => void) => {
@@ -49,6 +59,10 @@ vi.mock('$lib/i18n/index.svelte.js', () => ({
 }));
 
 describe('TaskResultPanel migration summary', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
   it('shows diagnostic fields, raw Superset response, artifacts, and related-log action', async () => {
     const onshowrelatedlogs = vi.fn();
     const task = {
@@ -95,10 +109,45 @@ describe('TaskResultPanel migration summary', () => {
     expect(screen.getByText('HTTP 422')).toBeTruthy();
     expect(screen.getByText(/Exception: SupersetAPIError/)).toBeTruthy();
     expect(screen.getByText('Raw Superset response')).toBeTruthy();
-    expect(screen.getByRole('link', { name: 'Download source export ZIP' }).getAttribute('href')).toContain('3.source.zip');
-    expect(screen.getByRole('link', { name: 'Download transformed ZIP' }).getAttribute('href')).toContain('3.zip');
+    expect(screen.getByRole('button', { name: 'Download source export ZIP' })).toBeTruthy();
+    expect(screen.getByRole('button', { name: 'Download transformed ZIP' })).toBeTruthy();
     screen.getByRole('button', { name: 'Show related logs' }).click();
     expect(onshowrelatedlogs).toHaveBeenCalledWith(expect.objectContaining({ id: 3, phase: 'import' }));
   });
+
+  it('downloads artifacts through the authenticated storage service and blocks duplicate clicks', async () => {
+    let resolveDownload: (() => void) | undefined;
+    downloadFile.mockImplementation(() => new Promise<void>((resolve) => { resolveDownload = resolve; }));
+    const task = {
+      plugin_id: 'superset-migration',
+      result: { status: 'FAILED', failed_dashboards: [{ id: 3, title: 'Revenue', error: 'failed', archive_path: 'task-1/3.zip' }] },
+    };
+
+    render(TaskResultPanel, { props: { task } });
+    const button = screen.getByRole('button', { name: 'Download transformed ZIP' });
+    await fireEvent.click(button);
+    await fireEvent.click(button);
+
+    expect(downloadFile).toHaveBeenCalledTimes(1);
+    expect(downloadFile).toHaveBeenCalledWith('migrations_failed', 'task-1/3.zip');
+    expect((button as HTMLButtonElement).disabled).toBe(true);
+    resolveDownload?.();
+    await waitFor(() => expect((button as HTMLButtonElement).disabled).toBe(false));
+  });
+
+  it('reports authenticated download failures and allows retry', async () => {
+    downloadFile.mockRejectedValueOnce(new Error('Forbidden'));
+    const task = {
+      plugin_id: 'superset-migration',
+      result: { status: 'FAILED', failed_dashboards: [{ id: 3, title: 'Revenue', error: 'failed', archive_path: 'task-1/3.zip' }] },
+    };
+
+    render(TaskResultPanel, { props: { task } });
+    const button = screen.getByRole('button', { name: 'Download transformed ZIP' });
+    await fireEvent.click(button);
+
+    await waitFor(() => expect(error).toHaveBeenCalledWith('Download transformed ZIP: Forbidden'));
+    expect((button as HTMLButtonElement).disabled).toBe(false);
+  });
 });
 // #endregion Test.Tasks.TaskResultPanel.Migration