@opengeni/db 0.9.3 → 0.12.0

package/src/migrate.ts

@@ -4,6 +4,10 @@ import { fileURLToPath } from "node:url";
 import postgres from "postgres";
 
 const DEFAULT_DATABASE_URL = "postgres://opengeni:opengeni@127.0.0.1:5432/opengeni";
+const DEFAULT_MAX_NESTED_AGENT_DEPTH = 3;
+const MAX_NESTED_AGENT_DEPTH = 2_147_483_647;
+const batchedBackfillDirective =
+  /^-- opengeni:batched-backfill batch-size=(\d+) lock-timeout=(\d+(?:ms|s|min)) statement-timeout=(\d+(?:ms|s|min))$/;
 const concurrentIndexDirective = /^-- opengeni:concurrent-index lock-timeout=(\d+(?:ms|s|min))$/;
 const concurrentIndexStatement =
   /^CREATE\s+(?:UNIQUE\s+)?INDEX\s+CONCURRENTLY\s+(?:(IF\s+NOT\s+EXISTS)\s+)?(?:"((?:[^"]|"")+)"|([A-Za-z_][A-Za-z0-9_]*))\s+ON\b/is;
@@ -24,6 +28,15 @@ export interface ConcurrentIndexMigration {
   statement: string;
 }
 
+export type MigrationRuntimeOptions = {
+  maxNestedAgentDepth?: number;
+};
+
+type DeploymentDepthPolicy = {
+  maxNestedAgentDepth: number;
+  source: "deployment" | "default";
+};
+
 /** A bare Postgres identifier (schema/role name) safe to interpolate into DDL. */
 function assertIdentifier(name: string, value: string): string {
   if (!/^[A-Za-z_][A-Za-z0-9_]*$/.test(value)) {
@@ -95,11 +108,69 @@ export function parseConcurrentIndexMigration(
   };
 }
 
+function parseBatchedBackfillMigration(
+  file: string,
+  sqlText: string,
+): { batchSize: number; lockTimeout: string; statementTimeout: string; statement: string } | null {
+  const lines = sqlText.replaceAll("\r\n", "\n").split("\n");
+  const directiveIndex = /^-- deployment-mode: (?:rolling|maintenance)$/.test(
+    lines[0]?.trim() ?? "",
+  )
+    ? 1
+    : 0;
+  const directive = batchedBackfillDirective.exec(lines[directiveIndex]?.trim() ?? "");
+  if (!directive) return null;
+  const statement = lines
+    .slice(directiveIndex + 1)
+    .filter((line) => !line.trim().startsWith("--"))
+    .join("\n")
+    .trim();
+  const withoutTrailingSemicolon = statement.endsWith(";")
+    ? statement.slice(0, -1).trimEnd()
+    : statement;
+  const batchSize = Number(directive[1]!);
+  if (
+    !Number.isSafeInteger(batchSize) ||
+    batchSize < 1 ||
+    batchSize > 10_000 ||
+    !/^WITH\b/is.test(withoutTrailingSemicolon) ||
+    !/\bUPDATE\b/is.test(withoutTrailingSemicolon) ||
+    !/\bRETURNING\b/is.test(withoutTrailingSemicolon) ||
+    !new RegExp(`\\bLIMIT\\s+${batchSize}\\b`, "i").test(withoutTrailingSemicolon) ||
+    withoutTrailingSemicolon.includes(";")
+  ) {
+    throw new Error(
+      `${file}: opengeni:batched-backfill requires one bounded WITH ... UPDATE ... RETURNING statement whose LIMIT matches batch-size`,
+    );
+  }
+  return {
+    batchSize,
+    lockTimeout: directive[2]!,
+    statementTimeout: directive[3]!,
+    statement,
+  };
+}
+
 async function executeMigrationFile(
   sql: postgres.Sql,
   file: string,
   sqlText: string,
 ): Promise<void> {
+  const batchedBackfill = parseBatchedBackfillMigration(file, sqlText);
+  if (batchedBackfill) {
+    await sql`select set_config('lock_timeout', ${batchedBackfill.lockTimeout}, false)`;
+    await sql`select set_config('statement_timeout', ${batchedBackfill.statementTimeout}, false)`;
+    try {
+      for (;;) {
+        const result = await sql.unsafe(batchedBackfill.statement);
+        if (result.length === 0) break;
+      }
+    } finally {
+      await sql`select set_config('statement_timeout', '0', false)`;
+      await sql`select set_config('lock_timeout', '0', false)`;
+    }
+    return;
+  }
   const concurrentIndex = parseConcurrentIndexMigration(file, sqlText);
   if (!concurrentIndex) {
     await sql.unsafe(sqlText);
@@ -127,6 +198,53 @@ async function executeMigrationFile(
   }
 }
 
+function deploymentDepthPolicy(
+  options: MigrationRuntimeOptions | undefined,
+): DeploymentDepthPolicy {
+  const raw =
+    options === undefined
+      ? process.env.OPENGENI_MAX_NESTED_AGENT_DEPTH?.trim() || undefined
+      : options.maxNestedAgentDepth;
+  if (raw === undefined) {
+    return { maxNestedAgentDepth: DEFAULT_MAX_NESTED_AGENT_DEPTH, source: "default" };
+  }
+  const value = typeof raw === "number" ? raw : Number(raw);
+  if (
+    !Number.isSafeInteger(value) ||
+    value < 0 ||
+    value > MAX_NESTED_AGENT_DEPTH ||
+    (typeof raw === "string" && !/^(0|[1-9][0-9]*)$/.test(raw))
+  ) {
+    throw new Error(
+      `OPENGENI_MAX_NESTED_AGENT_DEPTH must be a non-negative 32-bit integer: ${raw}`,
+    );
+  }
+  return { maxNestedAgentDepth: value, source: "deployment" };
+}
+
+async function persistDeploymentDepthPolicy(
+  sql: postgres.Sql,
+  policy: DeploymentDepthPolicy,
+): Promise<void> {
+  const [relation] = await sql<{ exists: boolean }[]>`
+    select to_regclass('nested_agent_depth_configuration') is not null as exists
+  `;
+  if (!relation?.exists) return;
+  await sql`
+    insert into "nested_agent_depth_configuration" (
+      "singleton", "max_nested_agent_depth", "policy_source", "updated_at"
+    ) values (true, ${policy.maxNestedAgentDepth}, ${policy.source}, now())
+    on conflict ("singleton") do update
+    set "max_nested_agent_depth" = excluded."max_nested_agent_depth",
+        "policy_source" = excluded."policy_source",
+        "updated_at" = now()
+    where "nested_agent_depth_configuration"."max_nested_agent_depth"
+            is distinct from excluded."max_nested_agent_depth"
+       or "nested_agent_depth_configuration"."policy_source"
+            is distinct from excluded."policy_source"
+  `;
+}
+
 /**
  * Apply the OpenGeni SQL migration chain.
  *
@@ -160,9 +278,11 @@ export async function migrate(
     process.env.OPENGENI_DATABASE_URL ??
     DEFAULT_DATABASE_URL,
   schema: string | undefined = process.env.OPENGENI_DB_SCHEMA?.trim() || undefined,
+  runtimeOptions?: MigrationRuntimeOptions,
 ): Promise<void> {
   const migrationsDir = join(dirname(fileURLToPath(import.meta.url)), "../drizzle");
   const files = (await readdir(migrationsDir)).filter((file) => file.endsWith(".sql")).sort();
+  const depthPolicy = deploymentDepthPolicy(runtimeOptions);
   const sql = postgres(databaseUrl, { max: 1 });
   try {
     // Serialize concurrent migrate() runs; the session-level lock is released
@@ -177,6 +297,8 @@ export async function migrate(
       await sql.unsafe(`CREATE SCHEMA IF NOT EXISTS "opengeni_private"`);
       await sql.unsafe(`SET search_path = "${schema}", "opengeni_private", "public"`);
     }
+    await sql`select set_config('opengeni.max_nested_agent_depth', ${String(depthPolicy.maxNestedAgentDepth)}, false)`;
+    await sql`select set_config('opengeni.nested_agent_depth_policy_source', ${depthPolicy.source}, false)`;
     await sql.unsafe(
       `CREATE TABLE IF NOT EXISTS "schema_migrations" ("name" text PRIMARY KEY, "applied_at" timestamptz NOT NULL DEFAULT now())`,
     );
@@ -190,6 +312,9 @@ export async function migrate(
       await executeMigrationFile(sql, file, sqlText);
       await sql`INSERT INTO "schema_migrations" ("name") VALUES (${file}) ON CONFLICT DO NOTHING`;
     }
+    // Reconcile even when all migration names were already recorded. This is
+    // the only supported way to change deployment policy in a live database.
+    await persistDeploymentDepthPolicy(sql, depthPolicy);
   } finally {
     await sql.end();
   }
@@ -203,8 +328,12 @@ export async function migrate(
  * `targetSchema` undefined → `public` → standalone behavior. Thin wrapper over
  * `migrate` so there is one migration engine.
  */
-export async function runMigrations(adminConnection: string, targetSchema?: string): Promise<void> {
-  await migrate(adminConnection, targetSchema);
+export async function runMigrations(
+  adminConnection: string,
+  targetSchema?: string,
+  runtimeOptions?: MigrationRuntimeOptions,
+): Promise<void> {
+  await migrate(adminConnection, targetSchema, runtimeOptions);
 }
 
 if (import.meta.main) {

package/src/new-session-drafts.ts

@@ -0,0 +1,144 @@
+import type {
+  NewSessionDraftOptions,
+  ReasoningEffort,
+  ResourceRef,
+  ToolRef,
+} from "@opengeni/contracts";
+import { and, eq } from "drizzle-orm";
+import type { Database } from "./index";
+import * as schema from "./schema";
+
+export type NewSessionDraftRow = typeof schema.newSessionDrafts.$inferSelect;
+
+export class NewSessionDraftConflictError extends Error {
+  readonly name = "NewSessionDraftConflictError";
+
+  constructor(readonly currentRevision: number) {
+    super("New-session draft changed in another client");
+  }
+}
+
+export class NewSessionDraftAccessError extends Error {
+  readonly name = "NewSessionDraftAccessError";
+
+  constructor() {
+    super("New-session draft access changed");
+  }
+}
+
+export async function getNewSessionDraftInTransaction(
+  db: Database,
+  input: { workspaceId: string; subjectId: string; lock?: boolean },
+): Promise<NewSessionDraftRow | null> {
+  const query = db
+    .select()
+    .from(schema.newSessionDrafts)
+    .where(
+      and(
+        eq(schema.newSessionDrafts.workspaceId, input.workspaceId),
+        eq(schema.newSessionDrafts.subjectId, input.subjectId),
+      ),
+    )
+    .limit(1);
+  const rows = input.lock ? await query.for("update") : await query;
+  return rows[0] ?? null;
+}
+
+export async function saveNewSessionDraftInTransaction(
+  db: Database,
+  input: {
+    accountId: string;
+    workspaceId: string;
+    subjectId: string;
+    expectedRevision: number;
+    text: string;
+    resources: ResourceRef[];
+    tools: ToolRef[];
+    model: string;
+    reasoningEffort: ReasoningEffort;
+    options: NewSessionDraftOptions;
+    /** API-key and delegated service subjects have no workspace-membership row. */
+    requireWorkspaceMembership?: boolean;
+  },
+): Promise<NewSessionDraftRow> {
+  if (input.requireWorkspaceMembership !== false) {
+    // Serialize with removeWorkspaceMember(), which takes FOR UPDATE before it
+    // deletes private rows and the membership. A save that wins first commits
+    // before removal's cleanup; a removal that wins first leaves no membership
+    // for a stale, already-authorized request to recreate after re-invitation.
+    const [membership] = await db
+      .select({ id: schema.workspaceMemberships.id })
+      .from(schema.workspaceMemberships)
+      .where(
+        and(
+          eq(schema.workspaceMemberships.workspaceId, input.workspaceId),
+          eq(schema.workspaceMemberships.subjectId, input.subjectId),
+        ),
+      )
+      .for("key share")
+      .limit(1);
+    if (!membership) throw new NewSessionDraftAccessError();
+  }
+  const current = await getNewSessionDraftInTransaction(db, { ...input, lock: true });
+  const currentRevision = current?.revision ?? 0;
+  if (currentRevision !== input.expectedRevision) {
+    throw new NewSessionDraftConflictError(currentRevision);
+  }
+
+  const revision = currentRevision + 1;
+  const values = {
+    accountId: input.accountId,
+    workspaceId: input.workspaceId,
+    subjectId: input.subjectId,
+    revision,
+    text: input.text,
+    resources: input.resources,
+    tools: input.tools,
+    model: input.model,
+    reasoningEffort: input.reasoningEffort,
+    sessionOptions: input.options,
+    updatedAt: new Date(),
+  };
+  if (current) {
+    const [saved] = await db
+      .update(schema.newSessionDrafts)
+      .set(values)
+      .where(eq(schema.newSessionDrafts.id, current.id))
+      .returning();
+    if (!saved) throw new Error("New-session draft did not save");
+    return saved;
+  }
+
+  // SELECT FOR UPDATE cannot lock an absent key. Two first saves may therefore
+  // race; ON CONFLICT keeps the loser transaction usable so it can report the
+  // winner's authoritative revision instead of leaking a unique violation.
+  const [inserted] = await db
+    .insert(schema.newSessionDrafts)
+    .values(values)
+    .onConflictDoNothing({
+      target: [schema.newSessionDrafts.workspaceId, schema.newSessionDrafts.subjectId],
+    })
+    .returning();
+  if (inserted) return inserted;
+  const raced = await getNewSessionDraftInTransaction(db, { ...input, lock: true });
+  throw new NewSessionDraftConflictError(raced?.revision ?? 0);
+}
+
+/** Delete only the submitted revision; a newer sibling-tab revision survives. */
+export async function consumeNewSessionDraftInTransaction(
+  db: Database,
+  input: { workspaceId: string; subjectId: string; expectedRevision: number },
+): Promise<boolean> {
+  if (input.expectedRevision === 0) return false;
+  const deleted = await db
+    .delete(schema.newSessionDrafts)
+    .where(
+      and(
+        eq(schema.newSessionDrafts.workspaceId, input.workspaceId),
+        eq(schema.newSessionDrafts.subjectId, input.subjectId),
+        eq(schema.newSessionDrafts.revision, input.expectedRevision),
+      ),
+    )
+    .returning({ id: schema.newSessionDrafts.id });
+  return deleted.length > 0;
+}