@opengeni/db 3.7.4 → 4.0.0-canary.0

package/src/model-context-snapshots.ts

@@ -0,0 +1,151 @@
+import {
+  MODEL_CONTEXT_SNAPSHOT_MAX_UTF8_BYTES,
+  ModelContextSnapshot,
+  SessionModelContextResponse,
+} from "@opengeni/contracts";
+import { and, desc, eq, sql } from "drizzle-orm";
+import type { Database } from "./database";
+import { withRlsContext } from "./database";
+import * as schema from "./schema";
+
+export class ModelContextSnapshotAuthorityError extends Error {
+  readonly code = "model_context_snapshot_authority_mismatch";
+
+  constructor() {
+    super("Model context snapshot does not match the exact durable execution attempt");
+    this.name = "ModelContextSnapshotAuthorityError";
+  }
+}
+
+export async function persistModelContextSnapshot(
+  db: Database,
+  input: {
+    accountId: string;
+    workspaceId: string;
+    sessionId: string;
+    turnId: string;
+    attemptId: string;
+    executionGeneration: number;
+    snapshot: ModelContextSnapshot;
+  },
+): Promise<void> {
+  const snapshot = ModelContextSnapshot.parse(input.snapshot);
+  if (Buffer.byteLength(JSON.stringify(snapshot), "utf8") > MODEL_CONTEXT_SNAPSHOT_MAX_UTF8_BYTES) {
+    throw new RangeError("Model context snapshot exceeds the 16 MiB persist limit");
+  }
+  await withRlsContext(
+    db,
+    { accountId: input.accountId, workspaceId: input.workspaceId },
+    async (scopedDb) =>
+      await scopedDb.transaction(async (tx) => {
+        const [attempt] = await tx
+          .select({
+            accountId: schema.sessionTurnAttempts.accountId,
+            workspaceId: schema.sessionTurnAttempts.workspaceId,
+            sessionId: schema.sessionTurnAttempts.sessionId,
+            turnId: schema.sessionTurnAttempts.turnId,
+            executionGeneration: schema.sessionTurnAttempts.executionGeneration,
+          })
+          .from(schema.sessionTurnAttempts)
+          .where(
+            and(
+              eq(schema.sessionTurnAttempts.workspaceId, input.workspaceId),
+              eq(schema.sessionTurnAttempts.id, input.attemptId),
+            ),
+          )
+          .limit(1);
+        if (
+          !attempt ||
+          attempt.accountId !== input.accountId ||
+          attempt.sessionId !== input.sessionId ||
+          attempt.turnId !== input.turnId ||
+          attempt.executionGeneration !== input.executionGeneration
+        ) {
+          throw new ModelContextSnapshotAuthorityError();
+        }
+
+        const now = new Date(snapshot.capturedAt);
+        await tx
+          .insert(schema.sessionAttemptModelContextSnapshots)
+          .values({
+            attemptId: input.attemptId,
+            accountId: input.accountId,
+            workspaceId: input.workspaceId,
+            sessionId: input.sessionId,
+            turnId: input.turnId,
+            executionGeneration: input.executionGeneration,
+            requestIndex: snapshot.requestIndex,
+            capturedAt: now,
+            snapshot,
+            createdAt: now,
+            updatedAt: now,
+          })
+          .onConflictDoUpdate({
+            target: schema.sessionAttemptModelContextSnapshots.attemptId,
+            set: {
+              requestIndex: snapshot.requestIndex,
+              capturedAt: now,
+              snapshot,
+              updatedAt: now,
+            },
+            setWhere: sql`${schema.sessionAttemptModelContextSnapshots.requestIndex} <= ${snapshot.requestIndex}`,
+          });
+      }),
+  );
+}
+
+export async function getLatestSessionModelContext(
+  db: Database,
+  input: { accountId: string; workspaceId: string; sessionId: string },
+): Promise<SessionModelContextResponse> {
+  return await withRlsContext(
+    db,
+    { accountId: input.accountId, workspaceId: input.workspaceId },
+    async (scopedDb) => {
+      const [row] = await scopedDb
+        .select({
+          attemptId: schema.sessionAttemptModelContextSnapshots.attemptId,
+          turnId: schema.sessionAttemptModelContextSnapshots.turnId,
+          snapshot: schema.sessionAttemptModelContextSnapshots.snapshot,
+        })
+        .from(schema.sessionAttemptModelContextSnapshots)
+        .innerJoin(
+          schema.sessionTurns,
+          and(
+            eq(schema.sessionTurns.id, schema.sessionAttemptModelContextSnapshots.turnId),
+            eq(
+              schema.sessionTurns.workspaceId,
+              schema.sessionAttemptModelContextSnapshots.workspaceId,
+            ),
+            eq(schema.sessionTurns.sessionId, schema.sessionAttemptModelContextSnapshots.sessionId),
+          ),
+        )
+        .where(
+          and(
+            eq(schema.sessionAttemptModelContextSnapshots.workspaceId, input.workspaceId),
+            eq(schema.sessionAttemptModelContextSnapshots.sessionId, input.sessionId),
+          ),
+        )
+        .orderBy(
+          desc(schema.sessionTurns.position),
+          desc(schema.sessionAttemptModelContextSnapshots.executionGeneration),
+          desc(schema.sessionAttemptModelContextSnapshots.requestIndex),
+        )
+        .limit(1);
+      if (!row) {
+        return {
+          sessionId: input.sessionId,
+          attemptId: null,
+          turnId: null,
+          snapshot: null,
+        };
+      }
+      return {
+        sessionId: input.sessionId,
+        attemptId: row.attemptId,
+        turnId: row.turnId,
+        snapshot: ModelContextSnapshot.parse(row.snapshot),
+      };
+    },
+  );
+}

package/src/new-session-drafts.ts

@@ -357,23 +357,10 @@ export async function rememberNewSessionSelectionInTransaction(
   return Boolean(updated);
 }
 
-/**
- * Replace one exact accepted draft with the next-create safe seed. The row is
- * identity-serialized before the revision check so even the absent-row
- * revision-zero case cannot race a first save. Exact create callers receive a
- * typed conflict; legacy revision-only callers retain the old false/no-op
- * result. A missing row is an idempotent no-op.
- */
-export async function seedNewSessionDraftInTransaction(
+async function lockCurrentNewSessionDraft(
   db: Database,
-  input: {
-    workspaceId: string;
-    subjectId: string;
-    expectedRevision: number;
-    expectedSnapshot?: NewSessionDraftSnapshot;
-    acceptedSelection?: AcceptedNewSessionSelection;
-  },
-): Promise<boolean> {
+  input: { workspaceId: string; subjectId: string },
+): Promise<NewSessionDraftRow | undefined> {
   await lockNewSessionDraftIdentity(db, input.workspaceId, input.subjectId);
   const [current] = await db
     .select()
@@ -386,14 +373,23 @@ export async function seedNewSessionDraftInTransaction(
     )
     .for("update")
     .limit(1);
+  return current;
+}
+
+async function lockExactNewSessionDraft(
+  db: Database,
+  input: {
+    workspaceId: string;
+    subjectId: string;
+    expectedRevision: number;
+    expectedSnapshot: NewSessionDraftSnapshot;
+  },
+): Promise<NewSessionDraftRow> {
+  const current = await lockCurrentNewSessionDraft(db, input);
   if (!current || current.revision !== input.expectedRevision) {
-    if (input.expectedSnapshot) {
-      throw new NewSessionDraftConflictError(current?.revision ?? 0);
-    }
-    return false;
+    throw new NewSessionDraftConflictError(current?.revision ?? 0);
   }
   if (
-    input.expectedSnapshot &&
     stableJson({
       text: current.text,
       resources: current.resources,
@@ -407,6 +403,55 @@ export async function seedNewSessionDraftInTransaction(
   ) {
     throw new NewSessionDraftConflictError(current.revision);
   }
+  return current;
+}
+
+/**
+ * Fence a fresh session shell on the exact actor-private draft it represents.
+ * The caller keeps this row lock through its create transaction, so a draft
+ * that was already stale cannot leave a committed, uninitialized session.
+ * The later atomic initializer still performs the authoritative consume.
+ */
+export async function assertExactNewSessionDraftInTransaction(
+  db: Database,
+  input: {
+    workspaceId: string;
+    subjectId: string;
+    expectedRevision: number;
+    expectedSnapshot: NewSessionDraftSnapshot;
+  },
+): Promise<void> {
+  await lockExactNewSessionDraft(db, input);
+}
+
+/**
+ * Replace one exact accepted draft with the next-create safe seed. The row is
+ * identity-serialized before the revision check so even the absent-row
+ * revision-zero case cannot race a first save. Exact create callers receive a
+ * typed conflict; legacy revision-only callers retain the old false/no-op
+ * result. A missing row is an idempotent no-op.
+ */
+export async function seedNewSessionDraftInTransaction(
+  db: Database,
+  input: {
+    workspaceId: string;
+    subjectId: string;
+    expectedRevision: number;
+    expectedSnapshot?: NewSessionDraftSnapshot;
+    acceptedSelection?: AcceptedNewSessionSelection;
+  },
+): Promise<boolean> {
+  const current = input.expectedSnapshot
+    ? await lockExactNewSessionDraft(db, {
+        workspaceId: input.workspaceId,
+        subjectId: input.subjectId,
+        expectedRevision: input.expectedRevision,
+        expectedSnapshot: input.expectedSnapshot,
+      })
+    : await lockCurrentNewSessionDraft(db, input);
+  if (!current || current.revision !== input.expectedRevision) {
+    return false;
+  }
 
   const parsedOptions = NewSessionDraftOptions.safeParse(publicNewSessionDraftOptions(current));
   const options = parsedOptions.success ? parsedOptions.data : {};

package/src/organization-membership-lifecycle.ts

@@ -1,4 +1,5 @@
 import {
+  CreateAdditionalOrganizationResponse,
   CreateOrganizationResponse,
   ListOrganizationAdministrationMembersResponse,
   ListOrganizationInvitationsResponse,
@@ -18,6 +19,7 @@ import {
   OrganizationRetentionPolicy,
   OrganizationSummary,
   RevokeOrganizationWorkspaceMemberResponse,
+  type CreateAdditionalOrganizationResponse as CreateAdditionalOrganizationResponseType,
   type CreateOrganizationResponse as CreateOrganizationResponseType,
   type OrganizationAdministrationMember as OrganizationAdministrationMemberType,
   type OrganizationAdministrationOverview as OrganizationAdministrationOverviewType,
@@ -74,6 +76,32 @@ export async function createManagedOrganization(
   });
 }
 
+export async function createAdditionalManagedOrganization(
+  db: Database,
+  input: {
+    subjectId: string;
+    subjectLabel: string;
+    name: string;
+    workspaceName: string;
+    operationId: string;
+  },
+): Promise<CreateAdditionalOrganizationResponseType> {
+  return await db.transaction(async (tx) => {
+    await setSubjectRlsContext(tx as unknown as Database, input.subjectId);
+    const [row] = await rawRows<{ result: unknown }>(
+      tx,
+      sql`select create_additional_managed_organization(
+        ${input.subjectId},
+        ${input.subjectLabel},
+        ${input.name},
+        ${input.workspaceName},
+        ${input.operationId}::uuid
+      ) as result`,
+    );
+    return CreateAdditionalOrganizationResponse.parse(row?.result);
+  });
+}
+
 /**
  * Bounded replay of one exact organization-lifecycle transaction after a
  * PostgreSQL deadlock abort.
@@ -208,7 +236,10 @@ async function settleOrganizationMembershipProtocols(
   db: Database,
   settlements: OrganizationMembershipProtocolSettlement[],
 ): Promise<void> {
-  const [prior] = await rawRows<{ subject_id: string; initiating_human_subject_id: string }>(
+  const [prior] = await rawRows<{
+    subject_id: string;
+    initiating_human_subject_id: string;
+  }>(
     db,
     sql`select
       coalesce(current_setting('opengeni.subject_id', true), '') as subject_id,
@@ -436,7 +467,9 @@ async function runOrganizationWorkspaceCommand(
     { accountId: command.organizationId, workspaceId: null },
     async (scopedDb) => {
       await setSubjectRlsContext(scopedDb, command.actorSubjectId);
-      const [row] = await rawRows<{ result: OrganizationWorkspaceCommandResult }>(
+      const [row] = await rawRows<{
+        result: OrganizationWorkspaceCommandResult;
+      }>(
         scopedDb,
         sql`select organization_workspace_command(${JSON.stringify(command)}::jsonb) as result`,
       );
@@ -455,7 +488,10 @@ export async function createOrganizationWorkspace(
     operationId: string;
   },
 ): Promise<OrganizationWorkspaceAccessType> {
-  const command = await runOrganizationWorkspaceCommand(db, { action: "create", ...input });
+  const command = await runOrganizationWorkspaceCommand(db, {
+    action: "create",
+    ...input,
+  });
   const overview = await getOrganizationAdministrationOverview(db, input);
   const workspace = overview.workspaces.find((candidate) => candidate.id === command.workspaceId);
   if (!workspace) throw new Error("Created organization workspace is missing from its overview");
@@ -615,7 +651,9 @@ export async function assertActiveManagedHumanOrganizationMembership(
   db: Database,
   input: { accountId: string; subjectId: string },
 ): Promise<number | null> {
-  const [row] = await rawRows<{ authorization_revision: number | string | null }>(
+  const [row] = await rawRows<{
+    authorization_revision: number | string | null;
+  }>(
     db,
     sql`select assert_active_managed_human_organization_membership(
       ${input.accountId}::uuid,
@@ -1098,7 +1136,11 @@ export async function previewOrganizationRetentionDeletions(
 
 export async function claimOrganizationRetentionDeletion(
   db: Database,
-  input: { organizationId: string; operationId: string; excludedMembershipIds?: string[] },
+  input: {
+    organizationId: string;
+    operationId: string;
+    excludedMembershipIds?: string[];
+  },
 ): Promise<OrganizationRetentionDeletionClaimType | null> {
   return await runRetentionCapability(db, input.organizationId, async (scopedDb) => {
     const excludedMembershipIds = input.excludedMembershipIds ?? [];