fix(self-host): harden smoke/e2e gates for long brain init and CSRF

Re-login when JWT expires during brain-init wait, send Hermes CSRF on
browser-origin agent-api POSTs, and assert current_database() against the
selected connection instead of hardcoding dba_agent.

Co-authored-by: Venkat SF <venkatesh.sakamuri@stayflexi.com>

Author: cursoragent

scripts/self-host/e2e-agent-check.py

@@ -48,8 +48,29 @@ def main() -> int:
     conn = sys.argv[1] if len(sys.argv) > 1 else None
 
     opener = build_opener(HTTPCookieProcessor(CookieJar()))
+    csrf_token: str | None = None
+    csrf_header = "X-Hermes-CSRF-Token"
 
-    def req(url: str, data=None, *, origin: str | None = None, timeout: int = 180):
+    def fetch_agent_csrf() -> str | None:
+        nonlocal csrf_token
+        r = urllib.request.Request(
+            f"{frontend}/agent-api/api/auth/status",
+            headers={"Accept": "application/json"},
+        )
+        with opener.open(r, timeout=30) as resp:
+            data = json.loads(resp.read().decode() or "{}")
+        csrf_token = data.get("csrf_token") or None
+        return csrf_token
+
+    def req(
+        url: str,
+        data=None,
+        *,
+        origin: str | None = None,
+        timeout: int = 180,
+        _retried: bool = False,
+    ):
+        nonlocal csrf_token
         body = None
         headers: dict[str, str] = {}
         if data is not None:
@@ -58,15 +79,37 @@ def req(url: str, data=None, *, origin: str | None = None, timeout: int = 180):
         if origin:
             headers["Origin"] = origin
             headers["Referer"] = origin.rstrip("/") + "/"
+        # Browser Origin POSTs to /agent-api need the Hermes CSRF token once
+        # trusted-auth is on — same contract as src/lib/api/agentClient.js.
+        if data is not None and "/agent-api/" in url:
+            if not csrf_token:
+                fetch_agent_csrf()
+            if csrf_token:
+                headers[csrf_header] = csrf_token
         r = urllib.request.Request(
             url, data=body, headers=headers, method="POST" if data is not None else "GET"
         )
-        with opener.open(r, timeout=timeout) as resp:
-            raw = resp.read().decode() or "null"
-            return json.loads(raw)
+        try:
+            with opener.open(r, timeout=timeout) as resp:
+                raw = resp.read().decode() or "null"
+                return json.loads(raw)
+        except urllib.error.HTTPError as e:
+            if e.code == 403 and not _retried and "/agent-api/" in url and data is not None:
+                csrf_token = None
+                fetch_agent_csrf()
+                return req(url, data, origin=origin, timeout=timeout, _retried=True)
+            raise
 
     print("→ login")
-    req(f"{backend}/auth/login", {"email": email, "password": password})
+    # Prefer the frontend proxy so cookies match the Host /agent-api auth_request uses.
+    try:
+        req(
+            f"{frontend}/api/auth/login",
+            {"email": email, "password": password},
+            origin=frontend,
+        )
+    except Exception:
+        req(f"{backend}/auth/login", {"email": email, "password": password})
 
     if not conn:
         conns = req(f"{backend}/connections")
@@ -81,6 +124,27 @@ def req(url: str, data=None, *, origin: str | None = None, timeout: int = 180):
         return 1
     print(f"→ connection {conn}")
 
+    # Resolve the expected current_database() value from the connection record.
+    # Hardcoding dba_agent falsely fails when the demo seed connection (demo_shop)
+    # is selected — the agent is correct; the gate was wrong.
+    expected_db = "dba_agent"
+    try:
+        conns = req(f"{backend}/connections")
+        items = conns if isinstance(conns, list) else (conns.get("connections") or conns.get("items") or [])
+        for c in items:
+            cid = c.get("connectionId") or c.get("id")
+            if str(cid) == str(conn):
+                expected_db = (
+                    c.get("databaseName")
+                    or c.get("database")
+                    or c.get("dbName")
+                    or expected_db
+                )
+                break
+    except Exception as e:
+        print(f"WARN: could not resolve expected DB name ({e}); defaulting to {expected_db}")
+    print(f"→ expected current_database() = {expected_db}")
+
     # ── Agent tab ──────────────────────────────────────────────────────────
     print("\n=== Agent tab (browser → /agent-api → DeepSQL Agent → MCP) ===")
     bridge = req(f"{backend}/agent/session", {"connectionId": conn})
@@ -184,14 +248,14 @@ def req(url: str, data=None, *, origin: str | None = None, timeout: int = 180):
     )
     seen_failures = [m for m in failure_markers if m in answer_l]
     called_sql = any("execute_sql" in t for t in tools)
-    answered = "dba_agent" in answer_l
+    answered = expected_db.lower() in answer_l
 
     agent_ok = answered and called_sql and not seen_failures
     if not agent_ok:
         if not called_sql:
             print("AGENT_FAIL: execute_sql was never called")
         if not answered:
-            print("AGENT_FAIL: reply lacks the expected database name 'dba_agent'")
+            print(f"AGENT_FAIL: reply lacks the expected database name '{expected_db}'")
         if seen_failures:
             print(f"AGENT_FAIL: reply reports tool failure {seen_failures}")
     print("AGENT_OK", agent_ok)

scripts/self-host/smoke-test.sh

@@ -80,33 +80,37 @@ fi
 base="http://localhost:${DEEPSQL_BACKEND_PORT}/api"
 cookie_jar="$(mktemp)"
 trap 'rm -f "$cookie_jar"' EXIT
-# Retried rather than attempted once. The backend answers /actuator/health UP before it
-# serves logins, so this script -- the command install.sh recommends running next -- used
-# to abort on a perfectly good install with a bare `curl: (22) 401`. Because curl runs
-# under `set -e` with -f, that exit happened before the error message below could print,
-# so the failure named neither the endpoint nor the reason.
-login_json=""
-login_deadline=$((SECONDS + 120))
-while (( SECONDS < login_deadline )); do
-  # Login through the frontend proxy so the cookie jar matches the Host the
-  # browser (and /agent-api auth_request) will use. A jar filled against
-  # :8080 alone has made nginx's auth_request return 401 even when /auth/me
-  # on the backend would succeed with the same cookie.
-  if login_json="$(curl -fsS -c "$cookie_jar" -H 'Content-Type: application/json' \
-       -X POST "http://localhost:${DEEPSQL_FRONTEND_PORT}/api/auth/login" \
-       -d "{\"email\":\"${DEEPSQL_SMOKE_EMAIL}\",\"password\":\"${DEEPSQL_SMOKE_PASSWORD}\"}" 2>/dev/null)"; then
-    break
-  fi
-  # Fall back to hitting the backend directly (older installs / no frontend).
-  if login_json="$(curl -fsS -c "$cookie_jar" -H 'Content-Type: application/json' \
-       -X POST "$base/auth/login" \
-       -d "{\"email\":\"${DEEPSQL_SMOKE_EMAIL}\",\"password\":\"${DEEPSQL_SMOKE_PASSWORD}\"}" 2>/dev/null)"; then
-    break
-  fi
-  echo "Waiting for the backend to accept logins..."
-  sleep 5
-done
 
+# Login through the frontend proxy so the cookie jar matches the Host the
+# browser (and /agent-api auth_request) will use. A jar filled against
+# :8080 alone has made nginx's auth_request return 401 even when /auth/me
+# on the backend would succeed with the same cookie.
+# Retried rather than attempted once. The backend answers /actuator/health UP
+# before it serves logins, so this script used to abort on a good install with
+# a bare `curl: (22) 401` under `set -e` + curl -f.
+smoke_login() {
+  local deadline=$((SECONDS + "${1:-120}"))
+  local body=""
+  while (( SECONDS < deadline )); do
+    if body="$(curl -fsS -c "$cookie_jar" -b "$cookie_jar" -H 'Content-Type: application/json' \
+         -X POST "http://localhost:${DEEPSQL_FRONTEND_PORT}/api/auth/login" \
+         -d "{\"email\":\"${DEEPSQL_SMOKE_EMAIL}\",\"password\":\"${DEEPSQL_SMOKE_PASSWORD}\"}" 2>/dev/null)"; then
+      printf '%s' "$body"
+      return 0
+    fi
+    if body="$(curl -fsS -c "$cookie_jar" -b "$cookie_jar" -H 'Content-Type: application/json' \
+         -X POST "$base/auth/login" \
+         -d "{\"email\":\"${DEEPSQL_SMOKE_EMAIL}\",\"password\":\"${DEEPSQL_SMOKE_PASSWORD}\"}" 2>/dev/null)"; then
+      printf '%s' "$body"
+      return 0
+    fi
+    echo "Waiting for the backend to accept logins..."
+    sleep 5
+  done
+  return 1
+}
+
+login_json="$(smoke_login 120 || true)"
 if [[ "$login_json" != *"\"email\""* ]]; then
   echo "Error: login failed during smoke test." >&2
   echo "$login_json" >&2
@@ -163,8 +167,27 @@ if [[ "$DEEPSQL_SMOKE_WAIT_FOR_INIT" == "true" ]]; then
   echo "This calls the LLM once per schema batch, so several minutes is normal."
   deadline=$((SECONDS + DEEPSQL_SMOKE_INIT_TIMEOUT_SECONDS))
   last_report=""
+  init_json=""
   while (( SECONDS < deadline )); do
-    init_json="$(curl -fsS -b "$cookie_jar" "$base/connections/${connection_id}/init-status")"
+    # Do not use curl -f here: brain init often outlives the JWT (~15m), and a
+    # 401 under set -e aborted the smoke mid-progress with no recovery path.
+    init_code="$(curl -sS -o /tmp/deepsql-smoke-init.json -w '%{http_code}' \
+      -b "$cookie_jar" -c "$cookie_jar" \
+      "$base/connections/${connection_id}/init-status" || echo "000")"
+    if [[ "$init_code" == "401" ]]; then
+      echo "  [${SECONDS}s] session expired during brain init — re-logging in..."
+      if ! smoke_login 60 >/dev/null; then
+        echo "Error: re-login failed while waiting for brain init." >&2
+        exit 1
+      fi
+      continue
+    fi
+    if [[ "$init_code" != "200" ]]; then
+      echo "Error: init-status returned HTTP ${init_code}." >&2
+      cat /tmp/deepsql-smoke-init.json 2>/dev/null >&2 || true
+      exit 1
+    fi
+    init_json="$(cat /tmp/deepsql-smoke-init.json 2>/dev/null || true)"
     init_stage="$(printf '%s' "$init_json" | sed -n 's/.*"currentStage":"\([^"]*\)".*/\1/p')"
     init_progress="$(printf '%s' "$init_json" | sed -n 's/.*"progressPercent":\([0-9][0-9]*\).*/\1/p')"
     init_message="$(printf '%s' "$init_json" | sed -n 's/.*"stageMessage":"\([^"]*\)".*/\1/p')"