cli-sessions: share and fork storage rules

DEVELOPMENT.md

@@ -245,6 +245,21 @@ pnpm test
 This runs the web, web environment, and local development-tool tests. They
 should pass against the local PostgreSQL database.
 
+### CLI session storage verification
+
+After migrating the local database, run the session-copy and account-cleanup
+regressions with:
+
+```bash
+pnpm --filter web exec jest --runInBand src/routers/cli-sessions-router.test.ts src/lib/r2/cli-sessions.test.ts src/lib/user/deletion-queue/handlers/cli-v1.test.ts
+```
+
+Share publication refreshes a stable link with the current session snapshot.
+The blob URL columns identify the current copy folder. `pending_copy_ids`
+records folders awaiting deletion; later publication and deletion lifecycle
+calls retry that work before admitting another copy.
+
+
 ## Common Development Commands
 
 | Command | Description |

apps/web/src/app/api/upload-cli-session-blob-v2/route.ts

@@ -6,9 +6,7 @@ import { db } from '@kilocode/web-shared/lib/drizzle';
 import { eq, and } from 'drizzle-orm';
 import { cliSessions } from '@kilocode/db/schema';
 import { generateSignedUploadUrl } from '@/lib/r2/cli-sessions';
-import { BLOB_TYPES } from '@/routers/cli-sessions-router';
-
-const MAX_CONTENT_LENGTH = 5 * 1024 * 1024;
+import { BLOB_TYPES, CLI_SESSION_BLOB_MAX_BYTES } from '@/lib/cli-sessions/session-copies';
 
 const BodySchema = z.object({
   session_id: z.uuid(),
@@ -17,7 +15,10 @@ const BodySchema = z.object({
     .number()
     .int()
     .positive()
-    .max(MAX_CONTENT_LENGTH, `Content length must not exceed ${MAX_CONTENT_LENGTH} bytes (5MB)`),
+    .max(
+      CLI_SESSION_BLOB_MAX_BYTES,
+      `Content length must not exceed ${CLI_SESSION_BLOB_MAX_BYTES} bytes (5MB)`
+    ),
 });
 
 export async function POST(request: NextRequest) {

apps/web/src/lib/cli-sessions/session-copies.ts

@@ -0,0 +1,323 @@
+import 'server-only';
+
+import { and, count, eq, isNotNull, sql } from 'drizzle-orm';
+import { TRPCError } from '@trpc/server';
+import { cliSessions, sharedCliSessions } from '@kilocode/db/schema';
+import { db, type DrizzleTransaction } from '@kilocode/web-shared/lib/drizzle';
+import {
+  SessionBlobCopyTooLargeError,
+  copyBlobs,
+  deleteBlobs,
+  type BlobUrlKeys,
+  type FileName,
+  type FolderName,
+} from '@/lib/r2/cli-sessions';
+
+export const BLOB_TYPES = [
+  'api_conversation_history',
+  'task_metadata',
+  'ui_messages',
+  'git_state',
+] as const satisfies readonly FileName[];
+
+/** Largest blob the CLI session upload route accepts, per blob type. */
+export const CLI_SESSION_BLOB_MAX_BYTES = 5 * 1024 * 1024;
+
+/** A copy carries every blob type, so it is bounded by the upload cap times the blob count. */
+export const CLI_SESSION_COPY_MAX_BYTES = BLOB_TYPES.length * CLI_SESSION_BLOB_MAX_BYTES;
+
+/** Shared sessions one account may hold at a time. */
+export const CLI_SESSION_SHARE_MAX_PER_OWNER = 200;
+
+/** Forked sessions one account may hold at a time. */
+export const CLI_SESSION_FORK_MAX_PER_OWNER = 200;
+
+/**
+ * Every object key of one session folder. Deleting the whole set is idempotent,
+ * so callers do not have to know which blob types a session actually has.
+ */
+export function sessionBlobRefs(
+  folderName: FolderName
+): { folderName: FolderName; filename: FileName }[] {
+  return BLOB_TYPES.map(filename => ({ folderName, filename }));
+}
+
+/**
+ * Serialize the copy admission of one account. `pg_advisory_xact_lock` releases
+ * on commit or rollback, so the lock never outlives the transaction that guards
+ * the count-then-insert. A hash collision only serializes two accounts; it
+ * never admits a copy that should have been refused.
+ */
+async function lockOwnerCopies(tx: DrizzleTransaction, kiloUserId: string): Promise<void> {
+  await tx.execute(
+    sql`SELECT pg_advisory_xact_lock(hashtext('cli_session_copy'), hashtext(${kiloUserId}))`
+  );
+}
+
+export type PublishSharedSessionCopyParams = {
+  kiloUserId: string;
+  sessionId: string;
+  sharedState: string;
+};
+
+/** One row per blob type; a missing copy is stored as NULL rather than left out. */
+export function blobUrlColumns(copiedBlobs: Partial<Record<BlobUrlKeys, string>>) {
+  return {
+    api_conversation_history_blob_url: copiedBlobs.api_conversation_history_blob_url ?? null,
+    task_metadata_blob_url: copiedBlobs.task_metadata_blob_url ?? null,
+    ui_messages_blob_url: copiedBlobs.ui_messages_blob_url ?? null,
+    git_state_blob_url: copiedBlobs.git_state_blob_url ?? null,
+  };
+}
+
+type SharedSessionCopyRow = Pick<
+  typeof sharedCliSessions.$inferSelect,
+  | 'share_id'
+  | 'api_conversation_history_blob_url'
+  | 'task_metadata_blob_url'
+  | 'ui_messages_blob_url'
+  | 'git_state_blob_url'
+  | 'pending_copy_ids'
+>;
+
+export function sharedSessionCurrentCopyIds(share: SharedSessionCopyRow): string[] {
+  const copyIds = new Set<string>();
+
+  for (const blobUrl of [
+    share.api_conversation_history_blob_url,
+    share.task_metadata_blob_url,
+    share.ui_messages_blob_url,
+    share.git_state_blob_url,
+  ]) {
+    const [folderName, copyId] = blobUrl?.split('/') ?? [];
+    if (folderName === 'shared-sessions' && copyId) {
+      copyIds.add(copyId);
+    }
+  }
+
+  if (copyIds.size === 0) {
+    copyIds.add(share.share_id);
+  }
+
+  return [...copyIds];
+}
+
+/** Current storage folder for a link, used when copying it into a fork. */
+export function sharedSessionCurrentCopyId(share: SharedSessionCopyRow): string {
+  return sharedSessionCurrentCopyIds(share)[0];
+}
+
+/** Current and pending shared-session folders, used by cleanup paths. */
+export function sharedSessionCopyIds(share: SharedSessionCopyRow): string[] {
+  return [...new Set([...sharedSessionCurrentCopyIds(share), ...share.pending_copy_ids])];
+}
+
+/**
+ * Reserve the session row for a fork, or refuse once the account holds
+ * {@link CLI_SESSION_FORK_MAX_PER_OWNER} forks. The row is inserted with no
+ * blob urls; the caller adds them once the copy succeeds, so an admitted fork
+ * can never publish a row that points at nothing.
+ */
+export async function reserveForkedSession(
+  kiloUserId: string,
+  values: Omit<typeof cliSessions.$inferInsert, 'session_id' | 'is_fork'>
+): Promise<string> {
+  return db.transaction(async tx => {
+    await lockOwnerCopies(tx, kiloUserId);
+
+    const [totals] = await tx
+      .select({ forkCount: count() })
+      .from(cliSessions)
+      .where(and(eq(cliSessions.kilo_user_id, kiloUserId), eq(cliSessions.is_fork, true)));
+
+    if (totals.forkCount >= CLI_SESSION_FORK_MAX_PER_OWNER) {
+      throw new TRPCError({
+        code: 'BAD_REQUEST',
+        message: `An account can hold ${CLI_SESSION_FORK_MAX_PER_OWNER} forked sessions. Delete a forked session before forking another.`,
+      });
+    }
+
+    const sessionId = crypto.randomUUID();
+    await tx.insert(cliSessions).values({ session_id: sessionId, is_fork: true, ...values });
+    return sessionId;
+  });
+}
+
+/** Copy every blob of one session into another folder under the copy size limit. */
+export async function copySessionBlobs(
+  sourceId: string,
+  sourceFolder: FolderName,
+  destinationId: string,
+  destinationFolder: FolderName
+) {
+  try {
+    return await copyBlobs(
+      sourceId,
+      sourceFolder,
+      destinationId,
+      destinationFolder,
+      [...BLOB_TYPES],
+      {
+        maxBytes: CLI_SESSION_COPY_MAX_BYTES,
+      }
+    );
+  } catch (error) {
+    if (error instanceof SessionBlobCopyTooLargeError) {
+      throw new TRPCError({
+        code: 'BAD_REQUEST',
+        message: `Session blobs are larger than the ${CLI_SESSION_COPY_MAX_BYTES} byte copy limit.`,
+      });
+    }
+    throw error;
+  }
+}
+
+/**
+ * Publish a fresh copy for an existing link or create a link after its first
+ * copy completes. The owner lock keeps allowance accounting atomic, while the
+ * row lock serializes refreshes and deletion of one published link.
+ */
+export async function publishSharedSessionCopy({
+  kiloUserId,
+  sessionId,
+  sharedState,
+}: PublishSharedSessionCopyParams): Promise<{ share_id: string; session_id: string }> {
+  let copyId: string | undefined;
+  let previousCopyIds: string[] = [];
+
+  let shareId: string;
+  try {
+    shareId = await db.transaction(async tx => {
+      await lockOwnerCopies(tx, kiloUserId);
+
+      const [existing] = await tx
+        .select({
+          share_id: sharedCliSessions.share_id,
+          api_conversation_history_blob_url: sharedCliSessions.api_conversation_history_blob_url,
+          task_metadata_blob_url: sharedCliSessions.task_metadata_blob_url,
+          ui_messages_blob_url: sharedCliSessions.ui_messages_blob_url,
+          git_state_blob_url: sharedCliSessions.git_state_blob_url,
+          pending_copy_ids: sharedCliSessions.pending_copy_ids,
+        })
+        .from(sharedCliSessions)
+        .where(
+          and(
+            eq(sharedCliSessions.session_id, sessionId),
+            eq(sharedCliSessions.kilo_user_id, kiloUserId),
+            eq(sharedCliSessions.shared_state, sharedState)
+          )
+        )
+        .limit(1)
+        .for('update');
+
+      if (!existing) {
+        const [totals] = await tx
+          .select({ shareCount: count() })
+          .from(sharedCliSessions)
+          .where(
+            and(
+              eq(sharedCliSessions.kilo_user_id, kiloUserId),
+              isNotNull(sharedCliSessions.session_id)
+            )
+          );
+
+        if (totals.shareCount >= CLI_SESSION_SHARE_MAX_PER_OWNER) {
+          throw new TRPCError({
+            code: 'BAD_REQUEST',
+            message: `An account can hold ${CLI_SESSION_SHARE_MAX_PER_OWNER} shared sessions. Delete a shared session before sharing another.`,
+          });
+        }
+      }
+
+      if (existing && existing.pending_copy_ids.length > 0) {
+        await deleteSharedSessionCopies(existing.pending_copy_ids);
+      }
+
+      shareId = existing?.share_id ?? crypto.randomUUID();
+      copyId = crypto.randomUUID();
+      const copiedBlobs = await copySessionBlobs(sessionId, 'sessions', copyId, 'shared-sessions');
+
+      if (existing) {
+        previousCopyIds = sharedSessionCurrentCopyIds(existing);
+        await tx
+          .update(sharedCliSessions)
+          .set({
+            ...blobUrlColumns(copiedBlobs),
+            pending_copy_ids: previousCopyIds,
+          })
+          .where(eq(sharedCliSessions.share_id, shareId));
+      } else {
+        await tx.insert(sharedCliSessions).values({
+          share_id: shareId,
+          session_id: sessionId,
+          kilo_user_id: kiloUserId,
+          shared_state: sharedState,
+          ...blobUrlColumns(copiedBlobs),
+        });
+      }
+
+      return shareId;
+    });
+  } catch (error) {
+    if (copyId) {
+      try {
+        await deleteSharedSessionCopies([copyId]);
+      } catch {
+        // A cleanup failure must not replace the error that caused the rollback.
+      }
+    }
+    throw error;
+  }
+
+  for (const previousCopyId of previousCopyIds) {
+    try {
+      await deleteSharedSessionCopies([previousCopyId]);
+      await clearPendingSharedSessionCopy(shareId, previousCopyId);
+    } catch {
+      // The row retains this folder id so a later refresh or account cleanup can retry it.
+    }
+  }
+
+  return { share_id: shareId, session_id: sessionId };
+}
+
+/**
+ * Undo a fork reservation whose copy did not complete. The row belongs to the
+ * request that inserted it, so removing it together with whatever the copy
+ * wrote leaves nothing behind.
+ */
+export async function discardReservedFork(sessionId: string): Promise<void> {
+  try {
+    await deleteBlobs(sessionId, sessionBlobRefs('sessions'));
+  } catch {
+    // The row delete below still runs; a cleanup failure must not replace the
+    // error that caused the discard.
+  }
+  await db.delete(cliSessions).where(eq(cliSessions.session_id, sessionId));
+}
+
+/** Delete the shared-sessions object copies in the given copy folders. */
+export async function deleteSharedSessionCopies(copyIds: string[]): Promise<void> {
+  for (const copyId of copyIds) {
+    await deleteBlobs(copyId, sessionBlobRefs('shared-sessions'));
+  }
+}
+
+async function clearPendingSharedSessionCopy(shareId: string, copyId: string): Promise<void> {
+  await db.transaction(async tx => {
+    const [share] = await tx
+      .select({ pending_copy_ids: sharedCliSessions.pending_copy_ids })
+      .from(sharedCliSessions)
+      .where(eq(sharedCliSessions.share_id, shareId))
+      .for('update');
+
+    if (!share) return;
+
+    await tx
+      .update(sharedCliSessions)
+      .set({
+        pending_copy_ids: share.pending_copy_ids.filter(pendingId => pendingId !== copyId),
+      })
+      .where(eq(sharedCliSessions.share_id, shareId));
+  });
+}