From fc6e2a0f05034f21af3a7d91e882900b778fb3df Mon Sep 17 00:00:00 2001
From: Max Peintner <max@caos.ch>
Date: Mon, 1 Jun 2026 16:49:07 +0200
Subject: [PATCH] fix(login): retry logic for session creation after
 registration on NotFound (#12189)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

After user registration, the backend projections may not be up to date
yet when the Login UI immediately tries to create a session. This
results in a `QUERY-Dfbg2` ("User could not be found") error even though
the user was created successfully.

This adds retry logic with backoff (500ms/1s/2s, up to 3 attempts)
around `createSessionAndUpdateCookie` in the registration flow. Only
`NotFound` errors are retried — other errors are thrown immediately.

Closes #12173

---------

Co-authored-by: Livio Spring <9405495+livio-a@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
---
 apps/login/src/lib/server/register.ts | 38 ++++++++++++++++++++++++---
 1 file changed, 35 insertions(+), 3 deletions(-)

diff --git a/apps/login/src/lib/server/register.ts b/apps/login/src/lib/server/register.ts
index 374a16138b..998719d09e 100644
--- a/apps/login/src/lib/server/register.ts
+++ b/apps/login/src/lib/server/register.ts
@@ -2,17 +2,49 @@
 
 import { createSessionAndUpdateCookie, createSessionForIdpAndUpdateCookie } from "@/lib/server/cookie";
 import { addHumanUser, addIDPLink, getLoginSettings, getUserByID, listAuthenticationMethodTypes } from "@/lib/zitadel";
-import { create } from "@zitadel/client";
+import { Code, ConnectError, Duration, create } from "@zitadel/client";
 import { Factors } from "@zitadel/proto/zitadel/session/v2/session_pb";
-import { ChecksJson, ChecksSchema } from "@zitadel/proto/zitadel/session/v2/session_service_pb";
+import { Checks, ChecksJson, ChecksSchema } from "@zitadel/proto/zitadel/session/v2/session_service_pb";
 import crypto from "crypto";
 import { getTranslations } from "next-intl/server";
 import { cookies, headers } from "next/headers";
 import { completeFlowOrGetUrl } from "../client";
 import { getOrSetFingerprintId } from "../fingerprint";
+import { createLogger } from "../logger";
 import { getServiceConfig } from "../service-url";
 import { checkEmailVerification, checkMFAFactors } from "../verify-helper";
 
+const logger = createLogger("register");
+
+const MAX_SESSION_RETRIES = 3;
+const RETRY_DELAYS_MS = [500, 1000, 2000];
+
+/**
+ * After user creation, backend projections (users, login_names) may not be updated yet.
+ * This helper retries createSessionAndUpdateCookie on NotFound errors with increasing delays.
+ */
+async function createSessionWithRetry(command: { checks: Checks; requestId: string | undefined; lifetime?: Duration }) {
+  let lastError: unknown;
+  for (let attempt = 0; attempt < MAX_SESSION_RETRIES; attempt++) {
+    try {
+      return await createSessionAndUpdateCookie(command);
+    } catch (error) {
+      lastError = error;
+      const isNotFound = error instanceof ConnectError && error.code === Code.NotFound;
+      const isLastAttempt = attempt + 1 >= MAX_SESSION_RETRIES;
+      if (!isNotFound || isLastAttempt) {
+        throw error;
+      }
+      const delay = RETRY_DELAYS_MS[attempt] ?? 2000;
+      logger.warn(
+        `Session creation failed with NotFound (attempt ${attempt + 1}/${MAX_SESSION_RETRIES}), retrying in ${delay}ms...`,
+      );
+      await new Promise((resolve) => setTimeout(resolve, delay));
+    }
+  }
+  throw lastError;
+}
+
 type RegisterUserCommand = {
   email: string;
   firstName: string;
@@ -74,7 +106,7 @@ export async function registerUser(
 
   const checks = create(ChecksSchema, checkPayload);
 
-  const result = await createSessionAndUpdateCookie({
+  const result = await createSessionWithRetry({
     checks,
     requestId: command.requestId,
     lifetime: command.password ? loginSettings?.passwordCheckLifetime : undefined,