@opengeni/db 0.2.2 → 0.3.0

package/src/index.ts

@@ -34,6 +34,7 @@ import type {
   SessionGoal,
   SessionGoalCreatedBy,
   SessionGoalStatus,
+  SessionMcpServerMetadata,
   SessionStatus,
   SessionTurn,
   SessionTurnSource,
@@ -348,6 +349,7 @@ export const allWorkspacePermissions: Permission[] = [
   "api_keys:manage",
   "environments:manage",
   "environments:use",
+  "mcp_servers:attach",
   "goals:manage",
   "enrollments:read",
   "enrollments:manage",
@@ -1200,6 +1202,33 @@ export type EnabledMcpCapabilityServer = {
   headersEncrypted?: Record<string, string>;
 };
 
+export type CreateSessionMcpServerInput = {
+  id: string;
+  name?: string | null;
+  url: string;
+  allowedTools?: string[] | null;
+  timeoutMs?: number | null;
+  cacheToolsList?: boolean | null;
+  headersEncrypted?: Record<string, string>;
+};
+
+export type UpdateSessionMcpServerCredentialsInput = {
+  id: string;
+  headersEncrypted: Record<string, string>;
+};
+
+export type UpdateSessionMcpServerCredentialsResult = {
+  servers: SessionMcpServerMetadata[];
+  missingIds: string[];
+};
+
+export type SessionMcpServerForRun = SessionMcpServerMetadata & {
+  allowedTools?: string[];
+  timeoutMs?: number;
+  cacheToolsList?: boolean;
+  headers: Record<string, string>;
+};
+
 export type EnqueueSessionTurnInput = {
   accountId: string;
   workspaceId: string;
@@ -3085,6 +3114,155 @@ function mapWorkspaceEnvironmentVariableMetadata(row: {
   };
 }
 
+function mapSessionMcpServerMetadata(row: typeof schema.sessionMcpServers.$inferSelect): SessionMcpServerMetadata {
+  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: Database,
+  workspaceId: string,
+  sessionIds: string[],
+): Promise<Map<string, SessionMcpServerMetadata[]>> {
+  const grouped = new Map<string, SessionMcpServerMetadata[]>();
+  if (sessionIds.length === 0) {
+    return grouped;
+  }
+  const rows = await db.select().from(schema.sessionMcpServers)
+    .where(and(
+      eq(schema.sessionMcpServers.workspaceId, workspaceId),
+      inArray(schema.sessionMcpServers.sessionId, sessionIds),
+    ))
+    .orderBy(asc(schema.sessionMcpServers.createdAt), asc(schema.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: Database, input: {
+  accountId: string;
+  workspaceId: string;
+  sessionId: string;
+  servers: CreateSessionMcpServerInput[];
+}): Promise<SessionMcpServerMetadata[]> {
+  if (input.servers.length === 0) {
+    return [];
+  }
+  const rows = await db.insert(schema.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);
+}
+
+export async function createSessionMcpServers(db: Database, input: {
+  accountId: string;
+  workspaceId: string;
+  sessionId: string;
+  servers: CreateSessionMcpServerInput[];
+}): Promise<SessionMcpServerMetadata[]> {
+  return await withRlsContext(db, { accountId: input.accountId, workspaceId: input.workspaceId }, async (scopedDb) =>
+    await insertSessionMcpServers(scopedDb, input)
+  );
+}
+
+export async function listSessionMcpServerMetadata(db: Database, workspaceId: string, sessionId: string): Promise<SessionMcpServerMetadata[]> {
+  return await withWorkspaceRls(db, workspaceId, async (scopedDb) => {
+    const grouped = await sessionMcpServerMetadataForSessions(scopedDb, workspaceId, [sessionId]);
+    return grouped.get(sessionId) ?? [];
+  });
+}
+
+export async function updateSessionMcpServerCredentials(db: Database, input: {
+  workspaceId: string;
+  sessionId: string;
+  updates: UpdateSessionMcpServerCredentialsInput[];
+}): Promise<UpdateSessionMcpServerCredentialsResult> {
+  return await withWorkspaceRls(db, input.workspaceId, async (scopedDb) => await scopedDb.transaction(async (tx) =>
+    await updateSessionMcpServerCredentialsInTransaction(tx, input)
+  ));
+}
+
+async function updateSessionMcpServerCredentialsInTransaction(
+  tx: Pick<Database, "update">,
+  input: {
+    workspaceId: string;
+    sessionId: string;
+    updates: UpdateSessionMcpServerCredentialsInput[];
+  },
+): Promise<UpdateSessionMcpServerCredentialsResult> {
+  const servers: SessionMcpServerMetadata[] = [];
+  const missingIds: string[] = [];
+  for (const update of input.updates) {
+    const [row] = await tx.update(schema.sessionMcpServers)
+      .set({
+        headersEncrypted: update.headersEncrypted,
+        credentialVersion: sql`${schema.sessionMcpServers.credentialVersion} + 1`,
+        updatedAt: new Date(),
+      })
+      .where(and(
+        eq(schema.sessionMcpServers.workspaceId, input.workspaceId),
+        eq(schema.sessionMcpServers.sessionId, input.sessionId),
+        eq(schema.sessionMcpServers.serverId, update.id),
+      ))
+      .returning();
+    if (!row) {
+      missingIds.push(update.id);
+    } else {
+      servers.push(mapSessionMcpServerMetadata(row));
+    }
+  }
+  return { servers, missingIds };
+}
+
+export async function listSessionMcpServersForRun(
+  db: Database,
+  workspaceId: string,
+  sessionId: string,
+  encryptionKey: Uint8Array,
+): Promise<SessionMcpServerForRun[]> {
+  return await withWorkspaceRls(db, workspaceId, async (scopedDb) => {
+    const rows = await scopedDb.select().from(schema.sessionMcpServers)
+      .where(and(
+        eq(schema.sessionMcpServers.workspaceId, workspaceId),
+        eq(schema.sessionMcpServers.sessionId, sessionId),
+      ))
+      .orderBy(asc(schema.sessionMcpServers.createdAt), asc(schema.sessionMcpServers.serverId));
+    return rows.map((row) => {
+      let headers: Record<string, string>;
+      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,
+      };
+    });
+  });
+}
+
 export async function createSession(db: Database, input: {
   accountId: string;
   workspaceId: string;
@@ -3103,6 +3281,7 @@ export async function createSession(db: Database, input: {
   // shared spawn passes the parent's sandboxGroupId so both run in ONE box.
   sandboxGroupId?: string | null;
   sandboxOs?: SandboxOs;
+  mcpServers?: CreateSessionMcpServerInput[];
 }): Promise<Session> {
   // Generate the id up front so the same uuid can seed sandbox_group_id for a
   // singleton group (sandbox_group_id cannot SQL-default to id).
@@ -3129,7 +3308,13 @@ export async function createSession(db: Database, 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);
   });
 }
 
@@ -3160,6 +3345,7 @@ export async function createSessionWithIdempotencyKey(db: Database, input: {
   // (group === the new row's own id); a shared spawn passes the parent's group.
   sandboxGroupId?: string | null;
   sandboxOs?: SandboxOs;
+  mcpServers?: CreateSessionMcpServerInput[];
 }): Promise<{ session: Session; created: boolean }> {
   // Generate the id up front so the same uuid can seed sandbox_group_id for a
   // singleton group (sandbox_group_id cannot SQL-default to id).
@@ -3187,7 +3373,13 @@ export async function createSessionWithIdempotencyKey(db: Database, input: {
       where: sql`${schema.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(schema.sessions).where(and(
       eq(schema.sessions.workspaceId, input.workspaceId),
@@ -3199,7 +3391,8 @@ export async function createSessionWithIdempotencyKey(db: Database, input: {
       // than silently returning a phantom.
       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 };
   });
 }
 
@@ -3209,14 +3402,18 @@ export async function getSessionByCreateIdempotencyKey(db: Database, workspaceId
       eq(schema.sessions.workspaceId, workspaceId),
       eq(schema.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) ?? []);
   });
 }
 
 export async function getSession(db: Database, workspaceId: string, sessionId: string): Promise<Session | null> {
   return await withWorkspaceRls(db, workspaceId, async (scopedDb) => {
     const [row] = await scopedDb.select().from(schema.sessions).where(and(eq(schema.sessions.workspaceId, workspaceId), eq(schema.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) ?? []);
   });
 }
 
@@ -3262,7 +3459,8 @@ export async function listSessions(db: Database, workspaceId: string, limit = 50
       .where(eq(schema.sessions.workspaceId, workspaceId))
       .orderBy(desc(schema.sessions.createdAt), desc(schema.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) ?? []));
   });
 }
 
@@ -3294,16 +3492,63 @@ export async function requireSession(db: Database, workspaceId: string, sessionI
   return session;
 }
 
-export async function listSessionEvents(db: Database, workspaceId: string, sessionId: string, after = 0, limit = 500): Promise<SessionEvent[]> {
+export type ListSessionEventsOptions = {
+  after?: number;
+  before?: number;
+  limit?: number;
+};
+
+const POSTGRES_INT_MAX = 2_147_483_647;
+
+export async function listSessionEvents(db: Database, workspaceId: string, sessionId: string): Promise<SessionEvent[]>;
+export async function listSessionEvents(db: Database, workspaceId: string, sessionId: string, after: number, limit?: number): Promise<SessionEvent[]>;
+export async function listSessionEvents(db: Database, workspaceId: string, sessionId: string, options: ListSessionEventsOptions): Promise<SessionEvent[]>;
+export async function listSessionEvents(
+  db: Database,
+  workspaceId: string,
+  sessionId: string,
+  afterOrOptions: number | ListSessionEventsOptions = 0,
+  legacyLimit = 500,
+): Promise<SessionEvent[]> {
+  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 !== undefined && Number.isFinite(options.before);
+  const before = hasBefore ? Math.floor(options.before as number) : undefined;
+
   return await withWorkspaceRls(db, workspaceId, async (scopedDb) => {
+    const filters: SQL[] = [
+      eq(schema.sessionEvents.workspaceId, workspaceId),
+      eq(schema.sessionEvents.sessionId, sessionId),
+      gt(schema.sessionEvents.sequence, after),
+    ];
+    if (before !== undefined && before <= POSTGRES_INT_MAX) {
+      filters.push(lt(schema.sessionEvents.sequence, before));
+    }
     const rows = await scopedDb.select().from(schema.sessionEvents)
-      .where(and(eq(schema.sessionEvents.workspaceId, workspaceId), eq(schema.sessionEvents.sessionId, sessionId), gt(schema.sessionEvents.sequence, after)))
-      .orderBy(asc(schema.sessionEvents.sequence))
+      .where(and(...filters))
+      .orderBy(hasBefore ? desc(schema.sessionEvents.sequence) : asc(schema.sessionEvents.sequence))
       .limit(limit);
-    return rows.map(mapEvent);
+    return (hasBefore ? rows.reverse() : rows).map(mapEvent);
   });
 }
 
+function normalizeEventSequence(value: number | undefined, fallback: number): number {
+  if (value === undefined || !Number.isFinite(value)) {
+    return fallback;
+  }
+  return Math.floor(value);
+}
+
+function normalizeEventLimit(value: number | undefined, fallback: number): number {
+  if (value === undefined || !Number.isFinite(value)) {
+    return fallback;
+  }
+  return Math.max(0, Math.floor(value));
+}
+
 export async function getSessionEvent(db: Database, workspaceId: string, eventId: string): Promise<SessionEvent | null> {
   return await withWorkspaceRls(db, workspaceId, async (scopedDb) => {
     const [row] = await scopedDb.select().from(schema.sessionEvents).where(and(eq(schema.sessionEvents.workspaceId, workspaceId), eq(schema.sessionEvents.id, eventId))).limit(1);
@@ -3384,8 +3629,8 @@ export async function getSessionHistoryItems(db: Database, workspaceId: string,
 
 /**
  * The LIVE conversation-truth read path: only active rows, position-ordered.
- * After a client-side context compaction this returns [active summary,
- * ...active recent tail]; with no compaction yet it equals
+ * After a client-side context compaction this returns [retained user messages,
+ * active summary]; with no compaction yet it equals
  * getSessionHistoryItems. The model-facing read path uses this so superseded
  * (summarized-away) prefix rows are excluded while the full transcript stays in
  * the table as an audit trail.
@@ -3555,8 +3800,15 @@ export async function applyContextCompaction(db: Database, input: {
   turnId?: string | null;
   /** Active prefix rows with position < boundaryPosition get superseded. */
   boundaryPosition: number;
-  /** Fractional position for the new summary row (must be < boundaryPosition). */
+  /** Position for the new summary row. Old boundary mode uses a fractional half-step before the kept tail. */
   summaryPosition: number;
+  /**
+   * Optional replacement rows inserted after superseding the old active set.
+   * Used by Codex-parity client compaction to rebuild active history as retained
+   * user messages plus one summary. These rows are synthetic replay rows, so
+   * they intentionally do not inherit the current compaction turn id.
+   */
+  replacementItems?: Array<{ position: number; item: Record<string, unknown> }>;
   summaryItem: Record<string, unknown>;
 }): Promise<void> {
   await withRlsContext(db, { accountId: input.accountId, workspaceId: input.workspaceId }, async (scopedDb) => {
@@ -3569,6 +3821,20 @@ export async function applyContextCompaction(db: Database, input: {
           eq(schema.sessionHistoryItems.active, true),
           lt(schema.sessionHistoryItems.position, input.boundaryPosition),
         ));
+      if (input.replacementItems && input.replacementItems.length > 0) {
+        await tx.insert(schema.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: [schema.sessionHistoryItems.workspaceId, schema.sessionHistoryItems.sessionId, schema.sessionHistoryItems.position],
+          set: { active: true },
+        });
+      }
       // Insert the summary at its FRACTIONAL position. The supersede step above
       // also sets active=false for any rows with position < boundaryPosition —
       // which on a RETRY includes the summary itself (it sits below the
@@ -7275,7 +7541,11 @@ export async function appendSessionEventsAndUpdateSession(db: Database, workspac
   }));
 }
 
-export async function appendSessionEventsWithLockedSessionUpdate(db: Database, workspaceId: string, sessionId: string, build: (session: Session) => {
+type LockedSessionUpdateContext = {
+  updateSessionMcpServerCredentials: (updates: UpdateSessionMcpServerCredentialsInput[]) => Promise<UpdateSessionMcpServerCredentialsResult>;
+};
+
+type LockedSessionUpdateResult = {
   events: AppendEventInput[];
   update?: {
     resources?: ResourceRef[];
@@ -7285,13 +7555,22 @@ export async function appendSessionEventsWithLockedSessionUpdate(db: Database, w
     status?: SessionStatus;
     activeTurnId?: string | null;
   };
-}): Promise<SessionEvent[]> {
+};
+
+export async function appendSessionEventsWithLockedSessionUpdate(
+  db: Database,
+  workspaceId: string,
+  sessionId: string,
+  build: (session: Session, context: LockedSessionUpdateContext) => LockedSessionUpdateResult | Promise<LockedSessionUpdateResult>,
+): Promise<SessionEvent[]> {
   return await withWorkspaceRls(db, workspaceId, async (scopedDb) => await scopedDb.transaction(async (tx) => {
     const [sessionRow] = await tx.select().from(schema.sessions).where(and(eq(schema.sessions.workspaceId, workspaceId), eq(schema.sessions.id, sessionId))).for("update").limit(1);
     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 [];
     }
@@ -7330,7 +7609,7 @@ export function sessionSubject(workspaceId: string, sessionId: string): string {
   return `workspaces.${workspaceId}.sessions.${sessionId}.events`;
 }
 
-function mapSession(row: typeof schema.sessions.$inferSelect): Session {
+function mapSession(row: typeof schema.sessions.$inferSelect, mcpServers: SessionMcpServerMetadata[] = []): Session {
   return {
     id: row.id,
     accountId: row.accountId,
@@ -7353,6 +7632,7 @@ function mapSession(row: typeof schema.sessions.$inferSelect): Session {
     activeEpoch: Number(row.activeEpoch),
     environmentId: row.environmentId,
     firstPartyMcpPermissions: (row.firstPartyMcpPermissions as Permission[] | null) ?? null,
+    mcpServers,
     parentSessionId: row.parentSessionId ?? null,
     createIdempotencyKey: row.createIdempotencyKey ?? null,
     temporalWorkflowId: row.temporalWorkflowId,

package/src/schema.ts

@@ -283,6 +283,28 @@ export const sessions = pgTable("sessions", {
   createIdempotency: uniqueIndex("sessions_workspace_create_idempotency_idx").on(table.workspaceId, table.createIdempotencyKey).where(sql`${table.createIdempotencyKey} is not null`),
 }));
 
+export const sessionMcpServers = pgTable("session_mcp_servers", {
+  id: uuid("id").primaryKey().defaultRandom(),
+  accountId: uuid("account_id").notNull().references(() => managedAccounts.id, { onDelete: "cascade" }),
+  workspaceId: uuid("workspace_id").notNull().references(() => workspaces.id, { onDelete: "cascade" }),
+  sessionId: uuid("session_id").notNull().references(() => sessions.id, { onDelete: "cascade" }),
+  serverId: text("server_id").notNull(),
+  name: text("name"),
+  url: text("url").notNull(),
+  allowedTools: jsonb("allowed_tools").$type<string[]>(),
+  timeoutMs: integer("timeout_ms"),
+  cacheToolsList: boolean("cache_tools_list").notNull().default(false),
+  // Map of header name -> AES-GCM ciphertext. Values are decrypted only by the
+  // worker's run-preparation path and never returned by API helpers.
+  headersEncrypted: jsonb("headers_encrypted").$type<Record<string, string>>().notNull().default({}),
+  credentialVersion: integer("credential_version").notNull().default(1),
+  createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(),
+  updatedAt: timestamp("updated_at", { withTimezone: true }).notNull().defaultNow(),
+}, (table) => ({
+  sessionServer: uniqueIndex("session_mcp_servers_session_server_idx").on(table.workspaceId, table.sessionId, table.serverId),
+  session: index("session_mcp_servers_session_idx").on(table.workspaceId, table.sessionId),
+}));
+
 export const files = pgTable("files", {
   id: uuid("id").primaryKey().defaultRandom(),
   accountId: uuid("account_id").notNull().references(() => managedAccounts.id, { onDelete: "cascade" }),