From 92a06e0ea34eb5387e257a9b6463174a5d4f0bee Mon Sep 17 00:00:00 2001
From: Owen <owen@txv.io>
Date: Wed, 13 May 2026 14:00:20 -0700
Subject: [PATCH] Handle jit mode with syncs

---
 server/routers/olm/buildConfiguration.ts | 18 +++++------
 server/routers/olm/sync.ts               | 39 ++++++++++++++++++++++--
 2 files changed, 44 insertions(+), 13 deletions(-)

diff --git a/server/routers/olm/buildConfiguration.ts b/server/routers/olm/buildConfiguration.ts
index 4182725d3..640031bca 100644
--- a/server/routers/olm/buildConfiguration.ts
+++ b/server/routers/olm/buildConfiguration.ts
@@ -27,11 +27,11 @@ export async function buildSiteConfigurationForOlmClient(
 ) {
     const siteConfigurations: {
         siteId: number;
-        name?: string
-        endpoint?: string
-        publicKey?: string
-        serverIP?: string | null
-        serverPort?: number | null
+        name?: string;
+        endpoint?: string;
+        publicKey?: string;
+        serverIP?: string | null;
+        serverPort?: number | null;
         remoteSubnets?: string[];
         aliases: Alias[];
     }[] = [];
@@ -79,7 +79,6 @@ export async function buildSiteConfigurationForOlmClient(
                 )
             );
 
-
         if (jitMode) {
             // Add site configuration to the array
             siteConfigurations.push({
@@ -109,10 +108,9 @@ export async function buildSiteConfigurationForOlmClient(
             continue;
         }
 
-        if (!site.publicKey || site.publicKey == "") { // the site is not ready to accept new peers
-            logger.warn(
-                `Site ${site.siteId} has no public key, skipping`
-            );
+        if (!site.publicKey || site.publicKey == "") {
+            // the site is not ready to accept new peers
+            logger.warn(`Site ${site.siteId} has no public key, skipping`);
             continue;
         }
 
diff --git a/server/routers/olm/sync.ts b/server/routers/olm/sync.ts
index c994b2c73..46e1fbd88 100644
--- a/server/routers/olm/sync.ts
+++ b/server/routers/olm/sync.ts
@@ -9,16 +9,50 @@ import {
 import { buildSiteConfigurationForOlmClient } from "./buildConfiguration";
 import { sendToClient } from "#dynamic/routers/ws";
 import logger from "@server/logger";
-import { eq, inArray } from "drizzle-orm";
+import { count, eq, inArray } from "drizzle-orm";
 import config from "@server/lib/config";
 import { canCompress } from "@server/lib/clientVersionChecks";
+import { build } from "@server/build";
 
 export async function sendOlmSyncMessage(olm: Olm, client: Client) {
+    // Get all sites data
+    const sitesCountResult = await db
+        .select({ count: count() })
+        .from(sites)
+        .innerJoin(
+            clientSitesAssociationsCache,
+            eq(sites.siteId, clientSitesAssociationsCache.siteId)
+        )
+        .where(eq(clientSitesAssociationsCache.clientId, client.clientId));
+
+    // Extract the count value from the result array
+    const sitesCount =
+        sitesCountResult.length > 0 ? sitesCountResult[0].count : 0;
+
+    // Prepare an array to store site configurations
+    logger.debug(
+        `[handleOlmRegisterMessage] Found ${sitesCount} sites for client ${client.clientId}`,
+        { orgId: client.orgId }
+    );
+
+    let jitMode = false;
+    if (sitesCount > 250 && build == "saas") {
+        // THIS IS THE MAX ON THE BUSINESS TIER
+        // we have too many sites
+        // If we have too many sites we need to drop into fully JIT mode by not sending any of the sites
+        logger.info(
+            `[handleOlmRegisterMessage] Too many sites (${sitesCount}), dropping into JIT mode`,
+            { orgId: client.orgId }
+        );
+        jitMode = true;
+    }
+
     // NOTE: WE ARE HARDCODING THE RELAY PARAMETER TO FALSE HERE BUT IN THE REGISTER MESSAGE ITS DEFINED BY THE CLIENT
     const siteConfigurations = await buildSiteConfigurationForOlmClient(
         client,
         client.pubKey,
-        false
+        false,
+        jitMode
     );
 
     // Get all exit nodes from sites where the client has peers
@@ -82,7 +116,6 @@ export async function sendOlmSyncMessage(olm: Olm, client: Client) {
                 exitNodes: exitNodesData
             }
         },
-
         {
             compress: canCompress(olm.version, "olm")
         }