@opengeni/db 0.31.1 → 0.36.0

package/src/task-notes.ts

@@ -0,0 +1,178 @@
+import {
+  TASK_NOTE_LIST_DEFAULT_LIMIT,
+  TASK_NOTE_LIST_MAX_LIMIT,
+  TaskNote,
+  TaskNoteKind,
+  TaskNoteListResponse,
+  TaskNoteMutationResult,
+  TaskNoteReason,
+  TaskNoteText,
+  type TaskNoteKind as TaskNoteKindType,
+} from "@opengeni/contracts";
+import { sql } from "drizzle-orm";
+import type { Database } from "./database";
+import { rawRows, withRlsContext } from "./database";
+
+export type TaskNoteAttemptClaims = {
+  accountId: string;
+  workspaceId: string;
+  sessionId: string;
+  turnId: string;
+  attemptId: string;
+  executionGeneration: number;
+};
+
+export type CreateTaskNoteInput = TaskNoteAttemptClaims & {
+  operationId: string;
+  kind: TaskNoteKindType;
+  text: string;
+  expiresInDays: number;
+};
+
+export type ArchiveTaskNoteInput = TaskNoteAttemptClaims & {
+  operationId: string;
+  noteId: string;
+  expectedVersion: number;
+  reason: string;
+};
+
+export type ListTaskNotesInput = TaskNoteAttemptClaims & {
+  includeArchived?: boolean;
+  limit?: number;
+};
+
+type TaskNoteRow = {
+  note_id: string;
+  root_session_id: string;
+  kind: string;
+  note_text: string;
+  status: string;
+  version: number;
+  expires_at: Date | string;
+  created_at: Date | string;
+  updated_at: Date | string;
+  archived_at: Date | string | null;
+  actor_kind: string;
+  source_session_id: string;
+  source_turn_id: string;
+  replayed?: boolean;
+};
+
+function iso(value: Date | string): string {
+  return (value instanceof Date ? value : new Date(value)).toISOString();
+}
+
+function noteFromRow(row: TaskNoteRow) {
+  return TaskNote.parse({
+    id: row.note_id,
+    rootSessionId: row.root_session_id,
+    kind: row.kind,
+    text: row.note_text,
+    status: row.status,
+    version: row.version,
+    expiresAt: iso(row.expires_at),
+    createdAt: iso(row.created_at),
+    updatedAt: iso(row.updated_at),
+    archivedAt: row.archived_at === null ? null : iso(row.archived_at),
+    provenance: {
+      actorKind: row.actor_kind,
+      sourceSessionId: row.source_session_id,
+      sourceTurnId: row.source_turn_id,
+    },
+  });
+}
+
+function assertAttemptClaims(claims: TaskNoteAttemptClaims): void {
+  if (!Number.isSafeInteger(claims.executionGeneration) || claims.executionGeneration < 1) {
+    throw new Error("Task notes require a positive exact execution generation");
+  }
+}
+
+export async function createTaskNote(db: Database, input: CreateTaskNoteInput) {
+  assertAttemptClaims(input);
+  TaskNoteKind.parse(input.kind);
+  TaskNoteText.parse(input.text);
+  if (
+    !Number.isSafeInteger(input.expiresInDays) ||
+    input.expiresInDays < 1 ||
+    input.expiresInDays > 30
+  ) {
+    throw new Error("Task note expiry must be 1-30 whole days");
+  }
+  const rows = await withRlsContext(
+    db,
+    { accountId: input.accountId, workspaceId: input.workspaceId },
+    async (tx) =>
+      await rawRows<TaskNoteRow>(
+        tx,
+        sql`SELECT * FROM create_task_note_for_attempt(
+          ${input.accountId}::uuid,
+          ${input.workspaceId}::uuid,
+          ${input.sessionId}::uuid,
+          ${input.turnId}::uuid,
+          ${input.attemptId}::uuid,
+          ${input.executionGeneration}::integer,
+          ${input.operationId}::uuid,
+          ${input.kind}::text,
+          ${input.text}::text,
+          ${input.expiresInDays}::integer
+        )`,
+      ),
+  );
+  if (rows.length !== 1) throw new Error("Task note create returned no durable receipt");
+  return TaskNoteMutationResult.parse({ note: noteFromRow(rows[0]!), replayed: rows[0]!.replayed });
+}
+
+export async function archiveTaskNote(db: Database, input: ArchiveTaskNoteInput) {
+  assertAttemptClaims(input);
+  TaskNoteReason.parse(input.reason);
+  const rows = await withRlsContext(
+    db,
+    { accountId: input.accountId, workspaceId: input.workspaceId },
+    async (tx) =>
+      await rawRows<TaskNoteRow>(
+        tx,
+        sql`SELECT * FROM archive_task_note_for_attempt(
+          ${input.accountId}::uuid,
+          ${input.workspaceId}::uuid,
+          ${input.sessionId}::uuid,
+          ${input.turnId}::uuid,
+          ${input.attemptId}::uuid,
+          ${input.executionGeneration}::integer,
+          ${input.operationId}::uuid,
+          ${input.noteId}::uuid,
+          ${input.expectedVersion}::integer,
+          ${input.reason}::text
+        )`,
+      ),
+  );
+  if (rows.length !== 1) throw new Error("Task note archive returned no durable receipt");
+  return TaskNoteMutationResult.parse({ note: noteFromRow(rows[0]!), replayed: rows[0]!.replayed });
+}
+
+export async function listTaskNotes(db: Database, input: ListTaskNotesInput) {
+  assertAttemptClaims(input);
+  const limit = input.limit ?? TASK_NOTE_LIST_DEFAULT_LIMIT;
+  if (!Number.isSafeInteger(limit) || limit < 1 || limit > TASK_NOTE_LIST_MAX_LIMIT) {
+    throw new Error(`Task note list limit must be 1-${TASK_NOTE_LIST_MAX_LIMIT}`);
+  }
+  const rows = await withRlsContext(
+    db,
+    { accountId: input.accountId, workspaceId: input.workspaceId },
+    async (tx) =>
+      await rawRows<TaskNoteRow>(
+        tx,
+        sql`SELECT * FROM list_task_notes_for_attempt(
+          ${input.accountId}::uuid,
+          ${input.workspaceId}::uuid,
+          ${input.sessionId}::uuid,
+          ${input.turnId}::uuid,
+          ${input.attemptId}::uuid,
+          ${input.executionGeneration}::integer,
+          ${input.includeArchived ?? false}::boolean,
+          ${limit}::integer
+        )`,
+      ),
+  );
+  return TaskNoteListResponse.parse({ notes: rows.map(noteFromRow) });
+}

package/src/video-generation.ts

@@ -1,6 +1,7 @@
 import {
   GENERATED_VIDEO_MAX_BYTES,
   GeneratedVideoReceipt,
+  GROK_IMAGINE_VIDEO_1_5_MODEL_ID,
   SEEDANCE_2_5_MODEL_ID,
   VideoGenerationPolicy,
   VideoGenerationOperationSummary,
@@ -236,8 +237,13 @@ export async function admitVideoGenerationOperation(
         }
         await tx
           .insert(schema.workspaceVideoGenerationQuotas)
-          .values({ workspaceId: input.workspaceId, accountId: input.accountId })
-          .onConflictDoNothing({ target: schema.workspaceVideoGenerationQuotas.workspaceId });
+          .values({
+            workspaceId: input.workspaceId,
+            accountId: input.accountId,
+          })
+          .onConflictDoNothing({
+            target: schema.workspaceVideoGenerationQuotas.workspaceId,
+          });
         const [quota] = await tx
           .select()
           .from(schema.workspaceVideoGenerationQuotas)
@@ -311,6 +317,44 @@ export async function admitVideoGenerationOperation(
   );
 }
 
+/** Persist a rotated provider token before it is used for another remote call. */
+export async function rotateVideoGenerationCredential(
+  db: Database,
+  input: {
+    accountId: string;
+    workspaceId: string;
+    operationId: string;
+    requestDigest: string;
+    expectedCredentialEncrypted: string;
+    credentialEncrypted: string;
+  },
+): Promise<VideoGenerationOperation> {
+  const [row] = await db
+    .update(schema.videoGenerationOperations)
+    .set({
+      credentialEncrypted: input.credentialEncrypted,
+      updatedAt: new Date(),
+    })
+    .where(
+      and(
+        eq(schema.videoGenerationOperations.accountId, input.accountId),
+        eq(schema.videoGenerationOperations.workspaceId, input.workspaceId),
+        eq(schema.videoGenerationOperations.id, input.operationId),
+        eq(schema.videoGenerationOperations.requestDigest, input.requestDigest),
+        eq(schema.videoGenerationOperations.credentialEncrypted, input.expectedCredentialEncrypted),
+        inArray(schema.videoGenerationOperations.status, [
+          "accepted",
+          "submission_uncertain",
+          "provider_started",
+          "retaining",
+        ]),
+      ),
+    )
+    .returning();
+  if (!row) throw new VideoGenerationConflictError("Video provider credential changed");
+  return row;
+}
+
 /**
  * Seal a reference-bearing operation only after every immutable staging object
  * has been uploaded and independently verified by the worker.  The source
@@ -422,11 +466,17 @@ export async function markVideoGenerationAcceptedInTransaction(
     operation.connectionId === null &&
     (operation.creditState === "debited" ||
       (operation.creditState === "not_applicable" && operation.pricedCostMicros === 0));
+  const supergrokAuthorized =
+    operation.fundingSource === "supergrok_subscription" &&
+    operation.connectionId === null &&
+    operation.credentialEncrypted !== null &&
+    operation.pricedCostMicros === 0 &&
+    operation.creditState === "not_applicable";
   if (
     !policy ||
     policy.fundingSource !== operation.fundingSource ||
     !policy.enabledModelIds.includes(operation.modelId) ||
-    (!workspaceGatewayAuthorized && !managedAuthorized)
+    (!workspaceGatewayAuthorized && !managedAuthorized && !supergrokAuthorized)
   ) {
     throw new VideoGenerationConflictError("Video generation authorization changed");
   }
@@ -684,7 +734,10 @@ export async function markVideoGenerationSubmissionIntent(
 
 export async function markVideoGenerationProviderStarted(
   db: Database,
-  input: OperationMutationInput & { providerJobId: string; nextReconcileAt: Date },
+  input: OperationMutationInput & {
+    providerJobId: string;
+    nextReconcileAt: Date;
+  },
 ): Promise<VideoGenerationOperation> {
   return await mutateOperation(db, input, async (tx, operation) => {
     if (
@@ -720,7 +773,11 @@ export async function markVideoGenerationProviderStarted(
 
 export async function rescheduleVideoGenerationOperation(
   db: Database,
-  input: OperationMutationInput & { owner?: string; nextReconcileAt: Date; error?: string },
+  input: OperationMutationInput & {
+    owner?: string;
+    nextReconcileAt: Date;
+    error?: string;
+  },
 ): Promise<boolean> {
   return await withRlsContext(
     db,
@@ -832,7 +889,10 @@ export type SettleVideoGenerationReadyInput = OperationMutationInput & {
 export async function settleVideoGenerationReady(
   db: Database,
   input: SettleVideoGenerationReadyInput,
-): Promise<{ operation: VideoGenerationOperation; artifact: GeneratedVideoArtifact }> {
+): Promise<{
+  operation: VideoGenerationOperation;
+  artifact: GeneratedVideoArtifact;
+}> {
   return await withRlsContext(
     db,
     { accountId: input.accountId, workspaceId: input.workspaceId },
@@ -931,6 +991,8 @@ export async function settleVideoGenerationReady(
             providerRequestExpiresAt: null,
             credentialEncrypted: null,
             requestEncrypted: null,
+            boundedPublicReason: null,
+            lastError: null,
             privateDataEraseAfter: now,
             reconcileLeaseOwner: null,
             reconcileLeaseExpiresAt: null,
@@ -950,7 +1012,10 @@ export async function getGeneratedVideoArtifact(
   db: Database,
   workspaceId: string,
   artifactId: string,
-): Promise<{ artifact: GeneratedVideoArtifact; file: typeof schema.files.$inferSelect } | null> {
+): Promise<{
+  artifact: GeneratedVideoArtifact;
+  file: typeof schema.files.$inferSelect;
+} | null> {
   return await withWorkspaceRls(db, workspaceId, async (scopedDb) => {
     const [row] = await scopedDb
       .select({ artifact: schema.generatedVideoArtifacts, file: schema.files })
@@ -971,13 +1036,21 @@ export async function listSessionGeneratedVideoArtifacts(
   db: Database,
   workspaceId: string,
   sessionId: string,
-): Promise<Array<{ artifact: GeneratedVideoArtifact; file: typeof schema.files.$inferSelect }>> {
+): Promise<
+  Array<{
+    artifact: GeneratedVideoArtifact;
+    file: typeof schema.files.$inferSelect;
+  }>
+> {
   return await withWorkspaceRls(
     db,
     workspaceId,
     async (scopedDb) =>
       await scopedDb
-        .select({ artifact: schema.generatedVideoArtifacts, file: schema.files })
+        .select({
+          artifact: schema.generatedVideoArtifacts,
+          file: schema.files,
+        })
         .from(schema.generatedVideoArtifacts)
         .innerJoin(schema.files, eq(schema.files.id, schema.generatedVideoArtifacts.primaryFileId))
         .where(
@@ -1026,7 +1099,9 @@ export async function markVideoGenerationReferenceCleaned(
             eq(schema.videoGenerationReferences.stagingObjectKey, input.stagingObjectKey),
           ),
         )
-        .returning({ operationId: schema.videoGenerationReferences.operationId });
+        .returning({
+          operationId: schema.videoGenerationReferences.operationId,
+        });
       if (updated) return true;
       const [existing] = await scopedDb
         .select({ cleanedAt: schema.videoGenerationReferences.cleanedAt })
@@ -1177,7 +1252,10 @@ export async function mediaGenerationResultForStoredOperation(
     },
     sandboxPath: `/workspace/generated-videos/${retained.artifact.sandboxFilename}`,
   });
-  return mediaGenerationResultForOperation({ operation, readyReceipt: receipt });
+  return mediaGenerationResultForOperation({
+    operation,
+    readyReceipt: receipt,
+  });
 }
 
 export async function getVideoGenerationOperationSummary(
@@ -1399,7 +1477,7 @@ function policyFromRow(
 
 function assertSupportedModels(modelIds: readonly string[]): void {
   for (const modelId of modelIds) {
-    if (modelId !== SEEDANCE_2_5_MODEL_ID) {
+    if (modelId !== SEEDANCE_2_5_MODEL_ID && modelId !== GROK_IMAGINE_VIDEO_1_5_MODEL_ID) {
       throw new Error(`Unknown video generation model: ${modelId}`);
     }
   }

package/src/workspace-artifacts.ts

@@ -320,7 +320,7 @@ async function assertAttemptAuthority(
           session.first_party_mcp_permissions IS NULL
           OR session.first_party_mcp_permissions @> '["artifacts:publish"]'::jsonb
         )
-      FOR UPDATE OF session
+      FOR NO KEY UPDATE OF session
     ), locked_turn AS MATERIALIZED (
       SELECT turn.id, turn.account_id, turn.workspace_id, turn.session_id,
         turn.active_attempt_id, turn.execution_generation, turn.initiator_subject_id