@opengeni/db 0.2.2 → 0.4.0

package/dist/index.js

@@ -26,6 +26,7 @@ import {
   sessionEvents,
   sessionGoals,
   sessionHistoryItems,
+  sessionMcpServers,
   sessionRecordings,
   sessionTurns,
   sessions,
@@ -38,7 +39,7 @@ import {
   workspaceMemberships,
   workspacePacks,
   workspaces
-} from "./chunk-PSX56ZTL.js";
+} from "./chunk-T2U4H4Z2.js";
 import {
   migrate,
   runMigrations
@@ -144,12 +145,64 @@ function sanitizeEventPayload(payload) {
   }
   if (payload && typeof payload === "object") {
     const entries = Object.entries(payload).map(
-      ([key, value]) => [sanitizeEventString(key), sanitizeEventPayload(value)]
+      ([key, value]) => [sanitizeEventString(key), sanitizeSensitiveEventField(key, value)]
     );
     return Object.fromEntries(entries);
   }
   return payload;
 }
+function sanitizeSensitiveEventField(key, value) {
+  if (key === "mcpServers") {
+    return sanitizeSessionMcpServerList(value);
+  }
+  if (key === "mcpCredentialUpdates") {
+    return sanitizeMcpCredentialUpdateList(value);
+  }
+  return sanitizeEventPayload(value);
+}
+function sanitizeSessionMcpServerList(value) {
+  if (!Array.isArray(value)) {
+    return sanitizeEventPayload(value);
+  }
+  return value.map((item) => {
+    if (!isPlainObject(item)) {
+      return sanitizeEventPayload(item);
+    }
+    const { headers, headersEncrypted, ...rest } = item;
+    const cleaned = sanitizeEventPayload(rest);
+    const headerNames = safeHeaderNames(headers) ?? safeHeaderNames(headersEncrypted);
+    if (headerNames) {
+      cleaned.headerNames = headerNames;
+    }
+    return cleaned;
+  });
+}
+function sanitizeMcpCredentialUpdateList(value) {
+  if (!Array.isArray(value)) {
+    return sanitizeEventPayload(value);
+  }
+  return value.map((item) => {
+    if (!isPlainObject(item)) {
+      return sanitizeEventPayload(item);
+    }
+    const { headers, headersEncrypted, ...rest } = item;
+    const cleaned = sanitizeEventPayload(rest);
+    const headerNames = safeHeaderNames(headers) ?? safeHeaderNames(headersEncrypted);
+    if (headerNames) {
+      cleaned.headerNames = headerNames;
+    }
+    return cleaned;
+  });
+}
+function safeHeaderNames(value) {
+  if (!isPlainObject(value)) {
+    return null;
+  }
+  return Object.keys(value).map(sanitizeEventString).sort();
+}
+function isPlainObject(value) {
+  return Boolean(value && typeof value === "object" && !Array.isArray(value));
+}
 
 // src/index.ts
 import { sql as sql2 } from "drizzle-orm";
@@ -388,6 +441,7 @@ var allWorkspacePermissions = [
   "api_keys:manage",
   "environments:manage",
   "environments:use",
+  "mcp_servers:attach",
   "goals:manage",
   "enrollments:read",
   "enrollments:manage"
@@ -2093,6 +2147,111 @@ function mapWorkspaceEnvironmentVariableMetadata(row) {
     updatedAt: row.updatedAt.toISOString()
   };
 }
+function mapSessionMcpServerMetadata(row) {
+  return {
+    id: row.serverId,
+    name: row.name ?? null,
+    url: row.url,
+    headerNames: Object.keys(row.headersEncrypted ?? {}).sort(),
+    credentialVersion: Number(row.credentialVersion)
+  };
+}
+async function sessionMcpServerMetadataForSessions(db, workspaceId, sessionIds) {
+  const grouped = /* @__PURE__ */ new Map();
+  if (sessionIds.length === 0) {
+    return grouped;
+  }
+  const rows = await db.select().from(sessionMcpServers).where(and(
+    eq(sessionMcpServers.workspaceId, workspaceId),
+    inArray(sessionMcpServers.sessionId, sessionIds)
+  )).orderBy(asc(sessionMcpServers.createdAt), asc(sessionMcpServers.serverId));
+  for (const row of rows) {
+    const list = grouped.get(row.sessionId) ?? [];
+    list.push(mapSessionMcpServerMetadata(row));
+    grouped.set(row.sessionId, list);
+  }
+  return grouped;
+}
+async function insertSessionMcpServers(db, input) {
+  if (input.servers.length === 0) {
+    return [];
+  }
+  const rows = await db.insert(sessionMcpServers).values(input.servers.map((server) => ({
+    accountId: input.accountId,
+    workspaceId: input.workspaceId,
+    sessionId: input.sessionId,
+    serverId: server.id,
+    name: server.name ?? null,
+    url: server.url,
+    allowedTools: server.allowedTools ?? null,
+    timeoutMs: server.timeoutMs ?? null,
+    cacheToolsList: server.cacheToolsList ?? false,
+    headersEncrypted: server.headersEncrypted ?? {}
+  }))).returning();
+  return rows.map(mapSessionMcpServerMetadata);
+}
+async function createSessionMcpServers(db, input) {
+  return await withRlsContext(
+    db,
+    { accountId: input.accountId, workspaceId: input.workspaceId },
+    async (scopedDb) => await insertSessionMcpServers(scopedDb, input)
+  );
+}
+async function listSessionMcpServerMetadata(db, workspaceId, sessionId) {
+  return await withWorkspaceRls(db, workspaceId, async (scopedDb) => {
+    const grouped = await sessionMcpServerMetadataForSessions(scopedDb, workspaceId, [sessionId]);
+    return grouped.get(sessionId) ?? [];
+  });
+}
+async function updateSessionMcpServerCredentials(db, input) {
+  return await withWorkspaceRls(db, input.workspaceId, async (scopedDb) => await scopedDb.transaction(
+    async (tx) => await updateSessionMcpServerCredentialsInTransaction(tx, input)
+  ));
+}
+async function updateSessionMcpServerCredentialsInTransaction(tx, input) {
+  const servers = [];
+  const missingIds = [];
+  for (const update of input.updates) {
+    const [row] = await tx.update(sessionMcpServers).set({
+      headersEncrypted: update.headersEncrypted,
+      credentialVersion: sql`${sessionMcpServers.credentialVersion} + 1`,
+      updatedAt: /* @__PURE__ */ new Date()
+    }).where(and(
+      eq(sessionMcpServers.workspaceId, input.workspaceId),
+      eq(sessionMcpServers.sessionId, input.sessionId),
+      eq(sessionMcpServers.serverId, update.id)
+    )).returning();
+    if (!row) {
+      missingIds.push(update.id);
+    } else {
+      servers.push(mapSessionMcpServerMetadata(row));
+    }
+  }
+  return { servers, missingIds };
+}
+async function listSessionMcpServersForRun(db, workspaceId, sessionId, encryptionKey) {
+  return await withWorkspaceRls(db, workspaceId, async (scopedDb) => {
+    const rows = await scopedDb.select().from(sessionMcpServers).where(and(
+      eq(sessionMcpServers.workspaceId, workspaceId),
+      eq(sessionMcpServers.sessionId, sessionId)
+    )).orderBy(asc(sessionMcpServers.createdAt), asc(sessionMcpServers.serverId));
+    return rows.map((row) => {
+      let headers;
+      try {
+        headers = Object.fromEntries(Object.entries(row.headersEncrypted ?? {}).map(([name, stored]) => [name, decryptEnvironmentValue(encryptionKey, stored)]));
+      } catch {
+        throw new Error("session MCP server credential decryption failed");
+      }
+      return {
+        ...mapSessionMcpServerMetadata(row),
+        ...row.allowedTools ? { allowedTools: row.allowedTools } : {},
+        ...row.timeoutMs ? { timeoutMs: row.timeoutMs } : {},
+        ...row.cacheToolsList ? { cacheToolsList: row.cacheToolsList } : {},
+        headers
+      };
+    });
+  });
+}
 async function createSession(db, input) {
   const id = crypto.randomUUID();
   return await withRlsContext(db, { accountId: input.accountId, workspaceId: input.workspaceId }, async (scopedDb) => {
@@ -2110,6 +2269,7 @@ async function createSession(db, input) {
       sandboxGroupId: input.sandboxGroupId ?? id,
       environmentId: input.environmentId ?? null,
       firstPartyMcpPermissions: input.firstPartyMcpPermissions ?? null,
+      instructions: input.instructions ?? null,
       parentSessionId: input.parentSessionId ?? null,
       createIdempotencyKey: input.createIdempotencyKey ?? null,
       status: "queued"
@@ -2117,7 +2277,13 @@ async function createSession(db, input) {
     if (!row) {
       throw new Error("Failed to create session");
     }
-    return mapSession(row);
+    const mcpServers = await insertSessionMcpServers(scopedDb, {
+      accountId: input.accountId,
+      workspaceId: input.workspaceId,
+      sessionId: row.id,
+      servers: input.mcpServers ?? []
+    });
+    return mapSession(row, mcpServers);
   });
 }
 async function createSessionWithIdempotencyKey(db, input) {
@@ -2137,6 +2303,7 @@ async function createSessionWithIdempotencyKey(db, input) {
       sandboxGroupId: input.sandboxGroupId ?? id,
       environmentId: input.environmentId ?? null,
       firstPartyMcpPermissions: input.firstPartyMcpPermissions ?? null,
+      instructions: input.instructions ?? null,
       parentSessionId: input.parentSessionId ?? null,
       createIdempotencyKey: input.createIdempotencyKey,
       status: "queued"
@@ -2145,7 +2312,13 @@ async function createSessionWithIdempotencyKey(db, input) {
       where: sql`${sessions.createIdempotencyKey} is not null`
     }).returning();
     if (inserted) {
-      return { session: mapSession(inserted), created: true };
+      const mcpServers = await insertSessionMcpServers(scopedDb, {
+        accountId: input.accountId,
+        workspaceId: input.workspaceId,
+        sessionId: inserted.id,
+        servers: input.mcpServers ?? []
+      });
+      return { session: mapSession(inserted, mcpServers), created: true };
     }
     const [existing] = await scopedDb.select().from(sessions).where(and(
       eq(sessions.workspaceId, input.workspaceId),
@@ -2154,7 +2327,8 @@ async function createSessionWithIdempotencyKey(db, input) {
     if (!existing) {
       throw new Error("Failed to create session under idempotency key");
     }
-    return { session: mapSession(existing), created: false };
+    const grouped = await sessionMcpServerMetadataForSessions(scopedDb, input.workspaceId, [existing.id]);
+    return { session: mapSession(existing, grouped.get(existing.id) ?? []), created: false };
   });
 }
 async function getSessionByCreateIdempotencyKey(db, workspaceId, createIdempotencyKey) {
@@ -2163,13 +2337,17 @@ async function getSessionByCreateIdempotencyKey(db, workspaceId, createIdempoten
       eq(sessions.workspaceId, workspaceId),
       eq(sessions.createIdempotencyKey, createIdempotencyKey)
     )).limit(1);
-    return row ? mapSession(row) : null;
+    if (!row) return null;
+    const grouped = await sessionMcpServerMetadataForSessions(scopedDb, workspaceId, [row.id]);
+    return mapSession(row, grouped.get(row.id) ?? []);
   });
 }
 async function getSession(db, workspaceId, sessionId) {
   return await withWorkspaceRls(db, workspaceId, async (scopedDb) => {
     const [row] = await scopedDb.select().from(sessions).where(and(eq(sessions.workspaceId, workspaceId), eq(sessions.id, sessionId))).limit(1);
-    return row ? mapSession(row) : null;
+    if (!row) return null;
+    const grouped = await sessionMcpServerMetadataForSessions(scopedDb, workspaceId, [row.id]);
+    return mapSession(row, grouped.get(row.id) ?? []);
   });
 }
 async function getAnySessionInGroup(db, workspaceId, sandboxGroupId) {
@@ -2187,7 +2365,8 @@ async function listDistinctEnvironmentIdsInGroup(db, workspaceId, sandboxGroupId
 async function listSessions(db, workspaceId, limit = 50) {
   return await withWorkspaceRls(db, workspaceId, async (scopedDb) => {
     const rows = await scopedDb.select().from(sessions).where(eq(sessions.workspaceId, workspaceId)).orderBy(desc(sessions.createdAt), desc(sessions.id)).limit(limit);
-    return rows.map(mapSession);
+    const grouped = await sessionMcpServerMetadataForSessions(scopedDb, workspaceId, rows.map((row) => row.id));
+    return rows.map((row) => mapSession(row, grouped.get(row.id) ?? []));
   });
 }
 async function countActiveSessionsForWorkspace(db, workspaceId) {
@@ -2208,12 +2387,38 @@ async function requireSession(db, workspaceId, sessionId) {
   }
   return session;
 }
-async function listSessionEvents(db, workspaceId, sessionId, after = 0, limit = 500) {
+var POSTGRES_INT_MAX = 2147483647;
+async function listSessionEvents(db, workspaceId, sessionId, afterOrOptions = 0, legacyLimit = 500) {
+  const options = typeof afterOrOptions === "number" ? { after: afterOrOptions, limit: legacyLimit } : afterOrOptions;
+  const after = normalizeEventSequence(options.after, 0);
+  const limit = normalizeEventLimit(options.limit, 500);
+  const hasBefore = options.before !== void 0 && Number.isFinite(options.before);
+  const before = hasBefore ? Math.floor(options.before) : void 0;
   return await withWorkspaceRls(db, workspaceId, async (scopedDb) => {
-    const rows = await scopedDb.select().from(sessionEvents).where(and(eq(sessionEvents.workspaceId, workspaceId), eq(sessionEvents.sessionId, sessionId), gt(sessionEvents.sequence, after))).orderBy(asc(sessionEvents.sequence)).limit(limit);
-    return rows.map(mapEvent);
+    const filters = [
+      eq(sessionEvents.workspaceId, workspaceId),
+      eq(sessionEvents.sessionId, sessionId),
+      gt(sessionEvents.sequence, after)
+    ];
+    if (before !== void 0 && before <= POSTGRES_INT_MAX) {
+      filters.push(lt(sessionEvents.sequence, before));
+    }
+    const rows = await scopedDb.select().from(sessionEvents).where(and(...filters)).orderBy(hasBefore ? desc(sessionEvents.sequence) : asc(sessionEvents.sequence)).limit(limit);
+    return (hasBefore ? rows.reverse() : rows).map(mapEvent);
   });
 }
+function normalizeEventSequence(value, fallback) {
+  if (value === void 0 || !Number.isFinite(value)) {
+    return fallback;
+  }
+  return Math.floor(value);
+}
+function normalizeEventLimit(value, fallback) {
+  if (value === void 0 || !Number.isFinite(value)) {
+    return fallback;
+  }
+  return Math.max(0, Math.floor(value));
+}
 async function getSessionEvent(db, workspaceId, eventId) {
   return await withWorkspaceRls(db, workspaceId, async (scopedDb) => {
     const [row] = await scopedDb.select().from(sessionEvents).where(and(eq(sessionEvents.workspaceId, workspaceId), eq(sessionEvents.id, eventId))).limit(1);
@@ -2348,6 +2553,20 @@ async function applyContextCompaction(db, input) {
         eq(sessionHistoryItems.active, true),
         lt(sessionHistoryItems.position, input.boundaryPosition)
       ));
+      if (input.replacementItems && input.replacementItems.length > 0) {
+        await tx.insert(sessionHistoryItems).values(input.replacementItems.map((entry) => ({
+          accountId: input.accountId,
+          workspaceId: input.workspaceId,
+          sessionId: input.sessionId,
+          turnId: null,
+          position: entry.position,
+          item: sanitizeEventPayload(entry.item),
+          active: true
+        }))).onConflictDoUpdate({
+          target: [sessionHistoryItems.workspaceId, sessionHistoryItems.sessionId, sessionHistoryItems.position],
+          set: { active: true }
+        });
+      }
       await tx.insert(sessionHistoryItems).values({
         accountId: input.accountId,
         workspaceId: input.workspaceId,
@@ -2816,6 +3035,28 @@ async function commitWarmingToWarm(db, input) {
     }
   );
 }
+async function recordWarmingSandboxCreated(db, input) {
+  return await withRlsContext(
+    db,
+    { accountId: input.accountId, workspaceId: input.workspaceId },
+    async (scopedDb) => {
+      const resumeStateJson = input.resumeState == null ? null : JSON.stringify(input.resumeState);
+      const rows = await scopedDb.execute(sql`
+        update sandbox_leases set
+          instance_id       = ${input.instanceId},
+          resume_backend_id = ${input.resumeBackendId ?? null},
+          resume_state      = ${resumeStateJson}::jsonb,
+          expires_at        = now() + (${String(input.leaseTtlMs)} || ' milliseconds')::interval,
+          updated_at        = now()
+        where workspace_id = ${input.workspaceId} and sandbox_group_id = ${input.sandboxGroupId}
+          and liveness = 'warming' and lease_epoch = ${input.expectedEpoch}
+        returning *
+      `);
+      if (rows.length === 0) return { recorded: false, lease: null };
+      return { recorded: true, lease: mapLeaseRow(rows[0]) };
+    }
+  );
+}
 async function failWarmingToCold(db, input) {
   await withRlsContext(
     db,
@@ -2943,7 +3184,21 @@ async function reapStaleLeaseHolders(db, input) {
           resume_backend_id = null, resume_state = null,
           data_plane_url = null, terminal_data_plane_url = null, updated_at = now()
         where workspace_id = ${input.workspaceId}
-          and liveness = 'warming' and expires_at < now()
+          and liveness = 'warming' and expires_at < now() and instance_id is null
+        returning id
+      `);
+    const warmingDrain = await tx.execute(sql`
+        update sandbox_leases set
+          liveness = 'draining',
+          refcount = 0,
+          turn_holders = 0,
+          viewer_holders = 0,
+          data_plane_url = null,
+          terminal_data_plane_url = null,
+          expires_at = now() - interval '1 millisecond',
+          updated_at = now()
+        where workspace_id = ${input.workspaceId}
+          and liveness = 'warming' and expires_at < now() and instance_id is not null
         returning id
       `);
     const drainable = await rawRows(tx, sql`
@@ -2953,7 +3208,7 @@ async function reapStaleLeaseHolders(db, input) {
       `);
     return {
       reapedViewers: reaped.length,
-      warmingReset: warmingReset.length,
+      warmingReset: warmingReset.length + warmingDrain.length,
       drained: drainable.map((r) => ({
         workspaceId: input.workspaceId,
         sandboxGroupId: r.sandbox_group_id,
@@ -2988,6 +3243,19 @@ async function listMeterableWarmLeases(db) {
     backend: r.backend
   }));
 }
+async function listLiveModalSandboxLeaseAttributions(db) {
+  const rows = await rawRows(db, sql`
+    select lease_id, workspace_id, sandbox_group_id, instance_id, liveness
+    from opengeni_private.list_live_modal_sandbox_leases()
+  `);
+  return rows.map((r) => ({
+    leaseId: r.lease_id,
+    workspaceId: r.workspace_id,
+    sandboxGroupId: r.sandbox_group_id,
+    instanceId: r.instance_id,
+    liveness: r.liveness
+  }));
+}
 async function reArmDrainingLease(db, input) {
   return await withRlsContext(
     db,
@@ -4482,7 +4750,9 @@ async function appendSessionEventsWithLockedSessionUpdate(db, workspaceId, sessi
     if (!sessionRow) {
       throw new Error(`Session not found: ${sessionId}`);
     }
-    const built = build(mapSession(sessionRow));
+    const built = await build(mapSession(sessionRow), {
+      updateSessionMcpServerCredentials: async (updates) => await updateSessionMcpServerCredentialsInTransaction(tx, { workspaceId, sessionId, updates })
+    });
     if (built.events.length === 0) {
       return [];
     }
@@ -4519,7 +4789,7 @@ async function appendSessionEventsWithLockedSessionUpdate(db, workspaceId, sessi
 function sessionSubject(workspaceId, sessionId) {
   return `workspaces.${workspaceId}.sessions.${sessionId}.events`;
 }
-function mapSession(row) {
+function mapSession(row, mcpServers = []) {
   return {
     id: row.id,
     accountId: row.accountId,
@@ -4528,6 +4798,7 @@ function mapSession(row) {
     initialMessage: row.initialMessage,
     title: row.title ?? null,
     titleSource: row.titleSource ?? null,
+    instructions: row.instructions ?? null,
     resources: row.resources,
     tools: row.tools,
     metadata: row.metadata,
@@ -4542,6 +4813,7 @@ function mapSession(row) {
     activeEpoch: Number(row.activeEpoch),
     environmentId: row.environmentId,
     firstPartyMcpPermissions: row.firstPartyMcpPermissions ?? null,
+    mcpServers,
     parentSessionId: row.parentSessionId ?? null,
     createIdempotencyKey: row.createIdempotencyKey ?? null,
     temporalWorkflowId: row.temporalWorkflowId,
@@ -4962,6 +5234,7 @@ export {
   createScheduledTaskRun,
   createSession,
   createSessionGoal,
+  createSessionMcpServers,
   createSessionWithIdempotencyKey,
   createSocialConnection,
   createSocialPost,
@@ -5053,6 +5326,7 @@ export {
   listEnrollments,
   listGitHubInstallationIdsForWorkspace,
   listGitHubInstallationsForWorkspace,
+  listLiveModalSandboxLeaseAttributions,
   listMeterableWarmLeases,
   listOpenPtySessions,
   listPackInstallations,
@@ -5062,6 +5336,8 @@ export {
   listScheduledTasks,
   listSessionEvents,
   listSessionIdsInGroup,
+  listSessionMcpServerMetadata,
+  listSessionMcpServersForRun,
   listSessionTurns,
   listSessions,
   listSocialConnections,
@@ -5099,6 +5375,7 @@ export {
   recordStreamAcknowledgment,
   recordStripeWebhookEvent,
   recordUsageEvent,
+  recordWarmingSandboxCreated,
   registerDbBinding,
   registerWorkspacePack,
   releaseLeaseHolder,
@@ -5145,6 +5422,7 @@ export {
   updateScheduledTask,
   updateScheduledTaskRun,
   updateSessionGoal,
+  updateSessionMcpServerCredentials,
   updateSessionTitle,
   updateWorkspace,
   updateWorkspaceEnvironment,