@opengeni/db 3.7.4 → 4.0.0-canary.0

package/src/session-queue-commands.ts

@@ -44,10 +44,13 @@ import {
   canonicalSessionCommandHash,
   evaluateSessionControl,
   lockSessionEventWriteRows,
+  lockSessionPromptCommandOperation,
   lockWorkspaceInferenceControl,
   lockWorkspaceInferenceControlForAdmission,
+  getCompletedSessionPromptCommandReceipt,
   registerInternalUpdateWakeInTransaction,
   reserveSessionCommandReceipt,
+  reserveSessionPromptCommandReceipt,
   registerSessionWorkflowWakeInTransaction,
   type SessionCommandActor,
   type SessionCommandReceiptRow,
@@ -210,6 +213,122 @@ function mapSubmittedPromptTurn(row: typeof schema.sessionTurns.$inferSelect): S
   };
 }
 
+async function loadSubmittedHumanPromptReplay(
+  db: Database,
+  input: {
+    workspaceId: string;
+    sessionId: string;
+    subjectId: string;
+    delivery: "send" | "steer";
+    expectedDraftRevision?: number | null;
+    receipt: SessionCommandReceiptRow;
+  },
+): Promise<SubmitHumanPromptResult> {
+  const turnId = String(input.receipt.result.turnId ?? "");
+  const acceptedEventId = String(input.receipt.result.acceptedEventId ?? "");
+  const eventIds = Array.isArray(input.receipt.result.eventIds)
+    ? input.receipt.result.eventIds.filter((id): id is string => typeof id === "string")
+    : [];
+  const wakeRevision = Number(input.receipt.result.wakeRevision ?? 0);
+  if (!turnId || !acceptedEventId || wakeRevision < 1) {
+    throw new SessionControlInvariantError("Replayed prompt receipt is incomplete");
+  }
+  const replayEvents = await db
+    .select()
+    .from(schema.sessionEvents)
+    .where(
+      and(
+        eq(schema.sessionEvents.workspaceId, input.workspaceId),
+        eq(schema.sessionEvents.sessionId, input.sessionId),
+        inArray(schema.sessionEvents.id, eventIds),
+      ),
+    )
+    .orderBy(asc(schema.sessionEvents.sequence));
+  const [replayTurn] = await db
+    .select()
+    .from(schema.sessionTurns)
+    .where(
+      and(
+        eq(schema.sessionTurns.workspaceId, input.workspaceId),
+        eq(schema.sessionTurns.sessionId, input.sessionId),
+        eq(schema.sessionTurns.id, turnId),
+      ),
+    )
+    .limit(1);
+  const events = replayEvents.map(mapSubmittedPromptEvent);
+  const accepted = events.find((event) => event.id === acceptedEventId);
+  if (!accepted || !replayTurn || events.length !== eventIds.length) {
+    throw new SessionControlInvariantError("Replayed prompt rows are incomplete");
+  }
+  const replayDraft =
+    input.expectedDraftRevision === null || input.expectedDraftRevision === undefined
+      ? null
+      : await getComposerDraftInTransaction(db, {
+          workspaceId: input.workspaceId,
+          sessionId: input.sessionId,
+          subjectId: input.subjectId,
+        });
+  return {
+    receipt: input.receipt,
+    routing:
+      input.receipt.result.routing === "accepted_for_execution" ||
+      input.receipt.result.routing === "queued_for_execution" ||
+      input.receipt.result.routing === "accepted_for_steering"
+        ? input.receipt.result.routing
+        : input.delivery === "steer"
+          ? "accepted_for_steering"
+          : // Pre-routing receipts cannot reconstruct admission-time activity
+            // or Pause truth. Queue is the only conservative placement; a
+            // later turn.started event moves it into chat without duplication.
+            "queued_for_execution",
+    queueVersion: Number(input.receipt.appliedQueueVersion),
+    accepted,
+    events,
+    turn: mapSubmittedPromptTurn(replayTurn),
+    draft: replayDraft,
+    acceptedEventId,
+    eventIds,
+    turnId,
+    wakeRevision,
+    interruptionCount: Number(input.receipt.result.interruptionCount ?? 0),
+    workspaceControlEventId:
+      typeof input.receipt.result.workspaceControlEventId === "string"
+        ? input.receipt.result.workspaceControlEventId
+        : null,
+    replay: true,
+  };
+}
+
+export async function replaySubmittedHumanPromptFromBoundaryReceipt(
+  db: Database,
+  input: {
+    workspaceId: string;
+    sessionId: string;
+    subjectId: string;
+    actor: SessionCommandActor;
+    operationKey: string;
+    delivery: "send" | "steer";
+    boundaryRequestHash: string;
+    expectedDraftRevision?: number | null;
+    /** Wait for an overlapping same-key transaction before checking committed replay truth. */
+    serializeOperation?: boolean;
+  },
+): Promise<SubmitHumanPromptResult | null> {
+  if (input.serializeOperation) {
+    await lockSessionPromptCommandOperation(db, input);
+  }
+  const receipt = await getCompletedSessionPromptCommandReceipt(db, {
+    workspaceId: input.workspaceId,
+    actor: input.actor,
+    action: input.delivery === "steer" ? "prompt.steer" : "prompt.send",
+    targetSessionId: input.sessionId,
+    operationKey: input.operationKey,
+    boundaryRequestHash: input.boundaryRequestHash,
+  });
+  if (!receipt) return null;
+  return await loadSubmittedHumanPromptReplay(db, { ...input, receipt });
+}
+
 export type AgentInternalUpdateCommandResult = {
   receipt: SessionCommandReceiptRow;
   updateId: string;
@@ -1608,6 +1727,8 @@ export async function submitHumanPromptInTransaction(
     subjectLabel?: string;
     actor: SessionCommandActor;
     operationKey: string;
+    /** Stable parsed boundary-request identity used before mutable enrichment. */
+    boundaryRequestHash?: string;
     delivery: "send" | "steer";
     controlEtag?: string | null;
     expectedDraftRevision?: number | null;
@@ -1642,6 +1763,9 @@ export async function submitHumanPromptInTransaction(
       id: string;
       headersEncrypted: Record<string, string>;
     }>;
+    /** Trusted database-only admission seam. It runs only after a completed
+     * operation replay has been ruled out; throwing rolls the prompt back. */
+    beforeFreshPromptCommit?: (tx: Database) => Promise<void>;
     /** Request-scoped callers bound the control prefix wait; lifecycle callers omit it. */
     controlLockTimeoutMs?: number;
   },
@@ -1708,91 +1832,41 @@ export async function submitHumanPromptInTransaction(
         }
       : {}),
   });
-  const reserved = await reserveSessionCommandReceipt(db, {
-    accountId: input.accountId,
-    workspaceId: input.workspaceId,
-    actor: input.actor,
-    action: input.delivery === "steer" ? "prompt.steer" : "prompt.send",
-    targetSessionId: input.sessionId,
-    targetTurnId: null,
-    operationKey: input.operationKey,
-    canonicalRequestHash: requestHash,
-  });
+  const action = input.delivery === "steer" ? "prompt.steer" : "prompt.send";
+  const reserved = input.boundaryRequestHash
+    ? await reserveSessionPromptCommandReceipt(db, {
+        accountId: input.accountId,
+        workspaceId: input.workspaceId,
+        actor: input.actor,
+        action,
+        targetSessionId: input.sessionId,
+        operationKey: input.operationKey,
+        canonicalRequestHash: requestHash,
+        boundaryRequestHash: input.boundaryRequestHash,
+      })
+    : await reserveSessionCommandReceipt(db, {
+        accountId: input.accountId,
+        workspaceId: input.workspaceId,
+        actor: input.actor,
+        action,
+        targetSessionId: input.sessionId,
+        targetTurnId: null,
+        operationKey: input.operationKey,
+        canonicalRequestHash: requestHash,
+      });
   if (reserved.replay && reserved.receipt.appliedQueueVersion !== null) {
-    const turnId = String(reserved.receipt.result.turnId ?? "");
-    const acceptedEventId = String(reserved.receipt.result.acceptedEventId ?? "");
-    const eventIds = Array.isArray(reserved.receipt.result.eventIds)
-      ? reserved.receipt.result.eventIds.filter((id): id is string => typeof id === "string")
-      : [];
-    const wakeRevision = Number(reserved.receipt.result.wakeRevision ?? 0);
-    if (!turnId || !acceptedEventId || wakeRevision < 1) {
-      throw new SessionControlInvariantError("Replayed prompt receipt is incomplete");
-    }
-    const replayEvents = await db
-      .select()
-      .from(schema.sessionEvents)
-      .where(
-        and(
-          eq(schema.sessionEvents.workspaceId, input.workspaceId),
-          eq(schema.sessionEvents.sessionId, input.sessionId),
-          inArray(schema.sessionEvents.id, eventIds),
-        ),
-      )
-      .orderBy(asc(schema.sessionEvents.sequence));
-    const [replayTurn] = await db
-      .select()
-      .from(schema.sessionTurns)
-      .where(
-        and(
-          eq(schema.sessionTurns.workspaceId, input.workspaceId),
-          eq(schema.sessionTurns.sessionId, input.sessionId),
-          eq(schema.sessionTurns.id, turnId),
-        ),
-      )
-      .limit(1);
-    const events = replayEvents.map(mapSubmittedPromptEvent);
-    const accepted = events.find((event) => event.id === acceptedEventId);
-    if (!accepted || !replayTurn || events.length !== eventIds.length) {
-      throw new SessionControlInvariantError("Replayed prompt rows are incomplete");
-    }
-    const replayDraft =
-      input.expectedDraftRevision === null || input.expectedDraftRevision === undefined
-        ? null
-        : await getComposerDraftInTransaction(db, {
-            workspaceId: input.workspaceId,
-            sessionId: input.sessionId,
-            subjectId: input.subjectId,
-          });
-    return {
+    return await loadSubmittedHumanPromptReplay(db, {
+      workspaceId: input.workspaceId,
+      sessionId: input.sessionId,
+      subjectId: input.subjectId,
+      delivery: input.delivery,
+      ...(input.expectedDraftRevision !== undefined
+        ? { expectedDraftRevision: input.expectedDraftRevision }
+        : {}),
       receipt: reserved.receipt,
-      routing:
-        reserved.receipt.result.routing === "accepted_for_execution" ||
-        reserved.receipt.result.routing === "queued_for_execution" ||
-        reserved.receipt.result.routing === "accepted_for_steering"
-          ? reserved.receipt.result.routing
-          : input.delivery === "steer"
-            ? "accepted_for_steering"
-            : // Pre-routing receipts cannot reconstruct admission-time activity
-              // or Pause truth. Queue is the only conservative placement; a
-              // later turn.started event moves it into chat without duplication.
-              "queued_for_execution",
-      queueVersion: Number(reserved.receipt.appliedQueueVersion),
-      accepted,
-      events,
-      turn: mapSubmittedPromptTurn(replayTurn),
-      draft: replayDraft,
-      acceptedEventId,
-      eventIds,
-      turnId,
-      wakeRevision,
-      interruptionCount: Number(reserved.receipt.result.interruptionCount ?? 0),
-      workspaceControlEventId:
-        typeof reserved.receipt.result.workspaceControlEventId === "string"
-          ? reserved.receipt.result.workspaceControlEventId
-          : null,
-      replay: true,
-    };
+    });
   }
+  await input.beforeFreshPromptCommit?.(db);
 
   const before = await evaluateSessionControl(db, input.workspaceId, input.sessionId, {
     workspaceControl,
@@ -2434,6 +2508,7 @@ export async function submitHumanPromptInTransaction(
     turnVersion: turn.version,
     ...(nextDraft ? { draftRevision: nextDraft.revision } : {}),
     result: {
+      ...(input.boundaryRequestHash ? { boundaryRequestHash: input.boundaryRequestHash } : {}),
       turnId,
       acceptedEventId,
       eventIds,

package/src/tool-gateway-approvals.ts

@@ -0,0 +1,168 @@
+import { and, eq, gt, isNull, sql } from "drizzle-orm";
+import type { ToolGatewayIdentity } from "@opengeni/contracts";
+import type { Database } from "./database";
+import { withWorkspaceSubjectRls } from "./database";
+import * as schema from "./schema";
+
+const MAX_LIVE_APPROVALS_PER_SUBJECT = 100;
+
+export class ToolGatewayApprovalRateLimitError extends Error {
+  readonly name = "ToolGatewayApprovalRateLimitError";
+}
+
+export class ToolGatewayApprovalOperationStartedError extends Error {
+  readonly name = "ToolGatewayApprovalOperationStartedError";
+}
+
+export async function issueToolGatewayApproval(
+  db: Database,
+  input: {
+    tokenHash: string;
+    accountId: string;
+    workspaceId: string;
+    subjectId: string;
+    operationId: string;
+    catalogDigest: string;
+    identity: ToolGatewayIdentity;
+    argumentsDigest: string;
+    approvalAuthorityDigest: string;
+    expiresAt: Date;
+  },
+): Promise<void> {
+  await withWorkspaceSubjectRls(db, input.workspaceId, input.subjectId, async (scopedDb) => {
+    await lockApprovalSubject(scopedDb, input.workspaceId, input.subjectId);
+    await lockApprovalOperation(scopedDb, input.workspaceId, input.subjectId, input.operationId);
+    const [existing] = await scopedDb
+      .select({
+        tokenHash: schema.toolGatewayApprovalCapabilities.tokenHash,
+        consumedAt: schema.toolGatewayApprovalCapabilities.consumedAt,
+      })
+      .from(schema.toolGatewayApprovalCapabilities)
+      .where(
+        and(
+          eq(schema.toolGatewayApprovalCapabilities.workspaceId, input.workspaceId),
+          eq(schema.toolGatewayApprovalCapabilities.subjectId, input.subjectId),
+          eq(schema.toolGatewayApprovalCapabilities.operationId, input.operationId),
+        ),
+      )
+      .limit(1);
+    if (existing?.consumedAt) {
+      throw new ToolGatewayApprovalOperationStartedError(
+        "The tool operation has already consumed its approval",
+      );
+    }
+    if (existing) {
+      // Reapproval before execution is allowed when catalog or provider
+      // authority changed. Once consumed, the row is retained permanently as
+      // the operation tombstone above and is never replaced.
+      await scopedDb
+        .delete(schema.toolGatewayApprovalCapabilities)
+        .where(eq(schema.toolGatewayApprovalCapabilities.tokenHash, existing.tokenHash));
+    }
+    await scopedDb.execute(sql`
+      delete from tool_gateway_approval_capabilities
+      where token_hash in (
+        select token_hash
+        from tool_gateway_approval_capabilities
+        where workspace_id = ${input.workspaceId}
+          and subject_id = ${input.subjectId}
+          and expires_at <= clock_timestamp()
+          and consumed_at is null
+        order by expires_at, token_hash
+        limit 128
+      )
+    `);
+    const [count] = await scopedDb
+      .select({ value: sql<number>`count(*)::integer` })
+      .from(schema.toolGatewayApprovalCapabilities)
+      .where(
+        and(
+          eq(schema.toolGatewayApprovalCapabilities.workspaceId, input.workspaceId),
+          eq(schema.toolGatewayApprovalCapabilities.subjectId, input.subjectId),
+          isNull(schema.toolGatewayApprovalCapabilities.consumedAt),
+          gt(schema.toolGatewayApprovalCapabilities.expiresAt, new Date()),
+        ),
+      );
+    if ((count?.value ?? 0) >= MAX_LIVE_APPROVALS_PER_SUBJECT) {
+      throw new ToolGatewayApprovalRateLimitError("Too many live tool approvals");
+    }
+    await scopedDb.insert(schema.toolGatewayApprovalCapabilities).values({
+      tokenHash: input.tokenHash,
+      accountId: input.accountId,
+      workspaceId: input.workspaceId,
+      subjectId: input.subjectId,
+      operationId: input.operationId,
+      catalogDigest: input.catalogDigest,
+      serverId: input.identity.serverId,
+      toolName: input.identity.toolName,
+      argumentsDigest: input.argumentsDigest,
+      authorityDigest: input.approvalAuthorityDigest,
+      expiresAt: input.expiresAt,
+    });
+  });
+}
+
+export async function consumeToolGatewayApproval(
+  db: Database,
+  input: {
+    tokenHash: string;
+    accountId: string;
+    workspaceId: string;
+    subjectId: string;
+    operationId: string;
+    catalogDigest: string;
+    identity: ToolGatewayIdentity;
+    argumentsDigest: string;
+    approvalAuthorityDigest: string;
+  },
+): Promise<boolean> {
+  return await withWorkspaceSubjectRls(db, input.workspaceId, input.subjectId, async (scopedDb) => {
+    await lockApprovalOperation(scopedDb, input.workspaceId, input.subjectId, input.operationId);
+    const [consumed] = await scopedDb
+      .update(schema.toolGatewayApprovalCapabilities)
+      .set({ consumedAt: new Date() })
+      .where(
+        and(
+          eq(schema.toolGatewayApprovalCapabilities.tokenHash, input.tokenHash),
+          eq(schema.toolGatewayApprovalCapabilities.accountId, input.accountId),
+          eq(schema.toolGatewayApprovalCapabilities.workspaceId, input.workspaceId),
+          eq(schema.toolGatewayApprovalCapabilities.subjectId, input.subjectId),
+          eq(schema.toolGatewayApprovalCapabilities.operationId, input.operationId),
+          eq(schema.toolGatewayApprovalCapabilities.catalogDigest, input.catalogDigest),
+          eq(schema.toolGatewayApprovalCapabilities.serverId, input.identity.serverId),
+          eq(schema.toolGatewayApprovalCapabilities.toolName, input.identity.toolName),
+          eq(schema.toolGatewayApprovalCapabilities.argumentsDigest, input.argumentsDigest),
+          eq(schema.toolGatewayApprovalCapabilities.authorityDigest, input.approvalAuthorityDigest),
+          isNull(schema.toolGatewayApprovalCapabilities.consumedAt),
+          gt(schema.toolGatewayApprovalCapabilities.expiresAt, new Date()),
+        ),
+      )
+      .returning({ tokenHash: schema.toolGatewayApprovalCapabilities.tokenHash });
+    return consumed !== undefined;
+  });
+}
+
+async function lockApprovalSubject(
+  db: Database,
+  workspaceId: string,
+  subjectId: string,
+): Promise<void> {
+  await db.execute(
+    sql`select pg_advisory_xact_lock(
+      hashtextextended(${`tool-gateway-approval-subject:${workspaceId}:${subjectId}`}, 0)
+    )`,
+  );
+}
+
+async function lockApprovalOperation(
+  db: Database,
+  workspaceId: string,
+  subjectId: string,
+  operationId: string,
+): Promise<void> {
+  await db.execute(
+    sql`select pg_advisory_xact_lock(
+      hashtextextended(${`tool-gateway-approval:${workspaceId}:${subjectId}:${operationId}`}, 0)
+    )`,
+  );
+}