@opengeni/db 1.0.1 → 3.2.0-canary.0

package/src/attached-browser-devices.ts

@@ -8,7 +8,7 @@ import {
   type AttachedBrowserDeviceListResponse as AttachedBrowserDeviceListResponseValue,
   type AttachedBrowserInventorySnapshot as AttachedBrowserInventorySnapshotValue,
 } from "@opengeni/contracts";
-import { and, asc, desc, eq, inArray, isNull } from "drizzle-orm";
+import { and, asc, desc, eq, inArray, isNull, sql } from "drizzle-orm";
 import { type Database, withRlsContext } from "./database";
 import {
   advanceWorkspaceInteractionRevision,
@@ -16,6 +16,12 @@ import {
 } from "./interaction-revisions";
 import * as schema from "./schema";
 
+const ATTACHED_DEVICE_LIVE_LIFECYCLES = ["starting", "active", "suspending", "restoring"] as const;
+// `ending` stays off this list so an in-flight `/end` can finish physical
+// teardown (`stopCapture` + helper exit) after Chrome reconnects.
+
+const CONTROLLER_TRANSITION_EXPIRED = "controller_transition_expired";
+
 type DeviceRow = typeof schema.attachedBrowserDevices.$inferSelect;
 
 const CONSISTENT_READ = {
@@ -156,13 +162,29 @@ export async function reconcileAttachedBrowserInventory(
           .where(eq(schema.attachedBrowserDevices.workspaceId, input.workspaceId))
           .for("update");
         const byId = new Map(rows.map((row) => [row.id, row]));
+        const generationChangedDeviceIds = new Set<string>();
         let changed = false;
         for (const device of snapshot.devices) {
           const existing = byId.get(device.id);
+          let reclaimingRevokedEnrollment = false;
           if (existing && existing.enrollmentId !== input.enrollmentId) {
-            throw new AttachedBrowserInventoryConflictError(
-              "attached browser device id is already owned by another enrollment",
-            );
+            const [previousEnrollment] = await tx
+              .select({ status: schema.enrollments.status })
+              .from(schema.enrollments)
+              .where(
+                and(
+                  eq(schema.enrollments.id, existing.enrollmentId),
+                  eq(schema.enrollments.accountId, input.accountId),
+                  eq(schema.enrollments.workspaceId, input.workspaceId),
+                ),
+              )
+              .limit(1);
+            if (previousEnrollment?.status !== "revoked") {
+              throw new AttachedBrowserInventoryConflictError(
+                "attached browser device id is already owned by another active enrollment",
+              );
+            }
+            reclaimingRevokedEnrollment = true;
           }
           if (!existing) {
             await tx.insert(schema.attachedBrowserDevices).values({
@@ -189,10 +211,15 @@ export async function reconcileAttachedBrowserInventory(
             changed = true;
             continue;
           }
-          const materialChanged = !sameAnnouncement(existing, device);
+          if (existing.connectionGeneration !== device.connectionGeneration) {
+            generationChangedDeviceIds.add(device.id);
+          }
+          const materialChanged =
+            reclaimingRevokedEnrollment || !sameAnnouncement(existing, device);
           await tx
             .update(schema.attachedBrowserDevices)
             .set({
+              enrollmentId: input.enrollmentId,
               name: device.name,
               profileLabel: device.profileLabel,
               browserName: device.browserName,
@@ -217,6 +244,15 @@ export async function reconcileAttachedBrowserInventory(
           changed ||= materialChanged;
         }
 
+        if (
+          await terminalizeStaleAttachedDeviceSessions(tx, {
+            workspaceId: input.workspaceId,
+            deviceIds: [...generationChangedDeviceIds],
+          })
+        ) {
+          changed = true;
+        }
+
         const announced = new Set(snapshot.devices.map((device) => device.id));
         const disconnectedIds = rows
           .filter(
@@ -275,6 +311,153 @@ export async function reconcileAttachedBrowserInventory(
   );
 }
 
+/** A new physical Chrome generation cannot inherit a prior controller token.
+ *  Mark every still-live BrowserSession/ComputerSession on that exact device
+ *  lost in place. Never rewrite placementInstanceId onto the successor. */
+async function terminalizeStaleAttachedDeviceSessions(
+  tx: Database,
+  input: { workspaceId: string; deviceIds: readonly string[] },
+): Promise<boolean> {
+  if (input.deviceIds.length === 0) return false;
+
+  const browserRows = await tx
+    .select({ id: schema.browserSessions.id })
+    .from(schema.browserSessions)
+    .where(
+      and(
+        eq(schema.browserSessions.workspaceId, input.workspaceId),
+        eq(schema.browserSessions.placementKind, "attached_device"),
+        inArray(schema.browserSessions.deviceId, [...input.deviceIds]),
+        inArray(schema.browserSessions.lifecycle, [...ATTACHED_DEVICE_LIVE_LIFECYCLES]),
+      ),
+    );
+  const computerRows = await tx
+    .select({ id: schema.computerSessions.id })
+    .from(schema.computerSessions)
+    .where(
+      and(
+        eq(schema.computerSessions.workspaceId, input.workspaceId),
+        eq(schema.computerSessions.placementKind, "attached_device"),
+        inArray(schema.computerSessions.deviceId, [...input.deviceIds]),
+        inArray(schema.computerSessions.lifecycle, [...ATTACHED_DEVICE_LIVE_LIFECYCLES]),
+      ),
+    );
+  if (browserRows.length === 0 && computerRows.length === 0) return false;
+
+  const browserIds = browserRows.map((row) => row.id).sort();
+  const computerIds = computerRows.map((row) => row.id).sort();
+  const resourceIds = [...browserIds, ...computerIds];
+  const operations = await tx
+    .select({ operationId: schema.interactionOperations.operationId })
+    .from(schema.interactionOperations)
+    .where(
+      and(
+        eq(schema.interactionOperations.workspaceId, input.workspaceId),
+        inArray(schema.interactionOperations.resourceId, resourceIds),
+        inArray(schema.interactionOperations.state, ["prepared", "dispatched"]),
+      ),
+    );
+  const operationIds = operations.map((row) => row.operationId).sort();
+  for (const operationId of operationIds) {
+    await tx.execute(sql`
+      select operation_id from interaction_operations
+      where workspace_id = ${input.workspaceId} and operation_id = ${operationId}
+      for update
+    `);
+  }
+  for (const browserSessionId of browserIds) {
+    await tx.execute(sql`
+      select id from browser_sessions
+      where workspace_id = ${input.workspaceId} and id = ${browserSessionId}
+      for update
+    `);
+  }
+  for (const computerSessionId of computerIds) {
+    await tx.execute(sql`
+      select id from computer_sessions
+      where workspace_id = ${input.workspaceId} and id = ${computerSessionId}
+      for update
+    `);
+  }
+
+  const now = new Date();
+  if (operationIds.length > 0) {
+    await tx
+      .update(schema.interactionOperations)
+      .set({
+        state: "outcome_unknown",
+        errorCode: "outcome_unknown",
+        errorMessage: "Attached Chrome connection generation changed",
+        errorRetryable: false,
+        errorDetails: { reason: "connection_generation_changed" },
+        settledAt: now,
+        updatedAt: now,
+      })
+      .where(
+        and(
+          eq(schema.interactionOperations.workspaceId, input.workspaceId),
+          inArray(schema.interactionOperations.operationId, operationIds),
+          inArray(schema.interactionOperations.state, ["prepared", "dispatched"]),
+        ),
+      );
+  }
+  if (browserIds.length > 0) {
+    await tx
+      .update(schema.browserSessions)
+      .set({
+        lifecycle: "lost",
+        failureCode: CONTROLLER_TRANSITION_EXPIRED,
+        updatedAt: now,
+      })
+      .where(
+        and(
+          eq(schema.browserSessions.workspaceId, input.workspaceId),
+          inArray(schema.browserSessions.id, browserIds),
+          inArray(schema.browserSessions.lifecycle, [...ATTACHED_DEVICE_LIVE_LIFECYCLES]),
+        ),
+      );
+  }
+  if (computerIds.length > 0) {
+    await tx
+      .update(schema.computerSessions)
+      .set({
+        lifecycle: "lost",
+        failureCode: CONTROLLER_TRANSITION_EXPIRED,
+        updatedAt: now,
+      })
+      .where(
+        and(
+          eq(schema.computerSessions.workspaceId, input.workspaceId),
+          inArray(schema.computerSessions.id, computerIds),
+          inArray(schema.computerSessions.lifecycle, [...ATTACHED_DEVICE_LIVE_LIFECYCLES]),
+        ),
+      );
+  }
+  return true;
+}
+
+/** Lock the device row first (same order as inventory) and prove the live
+ *  physical generation still matches this controller placement. */
+export async function attachedDeviceGenerationMatches(
+  tx: Database,
+  input: { workspaceId: string; deviceId: string; placementInstanceId: string },
+): Promise<boolean> {
+  const [device] = await tx
+    .select({
+      connectionGeneration: schema.attachedBrowserDevices.connectionGeneration,
+    })
+    .from(schema.attachedBrowserDevices)
+    .where(
+      and(
+        eq(schema.attachedBrowserDevices.workspaceId, input.workspaceId),
+        eq(schema.attachedBrowserDevices.id, input.deviceId),
+      ),
+    )
+    .for("update")
+    .limit(1);
+  return device?.connectionGeneration === input.placementInstanceId;
+}
+
 /** Immediately mark every live Chrome profile on a cleanly disconnected agent
  *  offline. Historical endpoint rows remain discoverable and reconnect in place. */
 export async function disconnectAttachedBrowserDevices(
@@ -336,8 +519,53 @@ export async function listAttachedBrowserDevices(
           asc(schema.attachedBrowserDevices.name),
           asc(schema.attachedBrowserDevices.id),
         );
+      const inventories = await scopedDb
+        .select({
+          enrollmentId: schema.attachedBrowserInventories.enrollmentId,
+          bridgeGeneration: schema.attachedBrowserInventories.bridgeGeneration,
+          revision: schema.attachedBrowserInventories.revision,
+          observedAt: schema.attachedBrowserInventories.observedAt,
+          enrollmentStatus: schema.enrollments.status,
+          connectionLeaseExpiresAt: schema.enrollments.connectionLeaseExpiresAt,
+          wentOfflineAt: schema.enrollments.wentOfflineAt,
+        })
+        .from(schema.attachedBrowserInventories)
+        .innerJoin(
+          schema.enrollments,
+          and(
+            eq(schema.enrollments.id, schema.attachedBrowserInventories.enrollmentId),
+            eq(schema.enrollments.workspaceId, schema.attachedBrowserInventories.workspaceId),
+          ),
+        )
+        .where(eq(schema.attachedBrowserInventories.workspaceId, input.workspaceId))
+        .orderBy(
+          desc(schema.attachedBrowserInventories.observedAt),
+          asc(schema.attachedBrowserInventories.enrollmentId),
+        );
+      const connectedProfilesByEnrollment = new Map<string, number>();
+      for (const row of rows) {
+        if (row.disconnectedAt !== null) continue;
+        connectedProfilesByEnrollment.set(
+          row.enrollmentId,
+          (connectedProfilesByEnrollment.get(row.enrollmentId) ?? 0) + 1,
+        );
+      }
       return AttachedBrowserDeviceListResponse.parse({
         revision: await readWorkspaceInteractionRevision(scopedDb, input.workspaceId),
+        bridges: inventories.map((inventory) => ({
+          enrollmentId: inventory.enrollmentId,
+          state:
+            inventory.enrollmentStatus === "active" &&
+            inventory.wentOfflineAt === null &&
+            inventory.connectionLeaseExpiresAt !== null &&
+            inventory.connectionLeaseExpiresAt.getTime() > Date.now()
+              ? "online"
+              : "offline",
+          bridgeGeneration: inventory.bridgeGeneration,
+          inventoryRevision: inventory.revision,
+          connectedProfileCount: connectedProfilesByEnrollment.get(inventory.enrollmentId) ?? 0,
+          lastSeenAt: inventory.observedAt.toISOString(),
+        })),
         devices: rows.map(deviceFromRow),
       });
     },