PaperClipAI/server/src/services/workspace-source-lock.ts

import { AsyncLocalStorage } from "node:async_hooks";
import { withDedicatedDbConnection, type Db } from "@paperclipai/db";
import { sql } from "drizzle-orm";
import path from "node:path";
import { resolvePaperclipInstanceRoot } from "../home-paths.js";
import { canonicalRepositoryRoot } from "./execution-workspace-source-roots.js";
import { readWorkspaceGitSourceRoots } from "./workspace-git-source-roots.js";
import { withWorkspaceLockKeepalive } from "./native-runtime/native-workspace-lock-keepalive.js";

type SourceTransaction = Parameters<Parameters<Db["transaction"]>[0]>[0];
type SourceLockScope = { db: Db; companyId: string; keys: Set<string>; transaction: SourceTransaction; active: boolean };
const sourceLocks = new AsyncLocalStorage<SourceLockScope>();
export class IsolatedRepositorySourceBusy extends Error {}
export type WorkspaceSourceLockWait = { waitMs: number; signal?: AbortSignal };

async function acquireQueuedSourceLocks(db: Db, transaction: SourceTransaction, keys: string[], wait: WorkspaceSourceLockWait) {
  wait.signal?.throwIfAborted();
  const deadline = performance.now() + wait.waitMs;
  const [settings] = await transaction.execute(sql`select current_setting('lock_timeout') as lock_timeout,
    pg_backend_pid() as pid, backend_start::text from pg_stat_activity where pid=pg_backend_pid()`);
  let cancelling: Promise<unknown> | undefined;
  const cancel = () => {
    cancelling ??= withDedicatedDbConnection(db, async control => {
      const timer = setTimeout(() => { void control.$client.end({ timeout: 0 }); }, 1000);
      try {
        await control.execute(sql`select pg_cancel_backend(pid) from pg_stat_activity
          where pid=${Number(settings.pid)} and backend_start=${String(settings.backend_start)}::timestamptz
            and datname=current_database() and usename=current_user`);
      } finally { clearTimeout(timer); }
    }).catch(() => undefined);
  };
  wait.signal?.addEventListener("abort", cancel, { once: true });
  try {
    for (const key of keys) {
      wait.signal?.throwIfAborted();
      const remaining = Math.ceil(deadline - performance.now());
      if (remaining <= 0) throw new IsolatedRepositorySourceBusy("workspace_source_busy");
      await transaction.execute(sql`select set_config('lock_timeout', ${`${remaining}ms`}, true)`);
      await transaction.execute(sql`select pg_advisory_xact_lock(hashtextextended(${key}, 0))`);
    }
    wait.signal?.throwIfAborted();
    if (performance.now() >= deadline) throw new IsolatedRepositorySourceBusy("workspace_source_busy");
    await transaction.execute(sql`select set_config('lock_timeout', ${String(settings.lock_timeout)}, true)`);
    wait.signal?.throwIfAborted();
  } catch (error) {
    wait.signal?.throwIfAborted();
    if ((error as { cause?: { code?: string }; code?: string }).cause?.code === "55P03"
      || (error as { code?: string }).code === "55P03") throw new IsolatedRepositorySourceBusy("workspace_source_busy");
    throw error;
  } finally {
    wait.signal?.removeEventListener("abort", cancel);
    await cancelling;
  }
}

// Timers and detached promises keep the async context that created them, so
// work can see a scope whose transaction has committed and whose dedicated
// connection is closed. Only a live scope holds locks.
function heldScope(): SourceLockScope | undefined {
  const scope = sourceLocks.getStore();
  return scope?.active ? scope : undefined;
}

/** Acquires source locks in a new owner scope and ignores an inherited one.
 * Use it at the root of independent work, such as a run's workspace
 * preparation, that can start inside another operation's source scope.
 * Joining that scope would tie these locks to the other transaction. */
export function withIndependentWorkspaceSourceLocks<T>(db: Db,
  input: { companyId: string; issueIds?: readonly string[]; projectIds?: readonly string[] },
  action: (transaction: SourceTransaction) => Promise<T>, tryOnly = false, wait?: WorkspaceSourceLockWait): Promise<T> {
  return sourceLocks.exit(() => withWorkspaceSourceLocks(db, input, action, tryOnly, wait));
}

export function withIsolatedRepositorySourceLock<T>(db: Db, companyId: string, issueId: string,
  action: () => Promise<T>, tryOnly = false): Promise<T> {
  return withWorkspaceSourceLocks(db, { companyId, issueIds: [issueId] }, action, tryOnly);
}

export async function withWorkspaceSourceLocks<T>(db: Db,
  input: { companyId: string; issueIds?: readonly string[]; projectIds?: readonly string[] },
  action: (transaction: SourceTransaction) => Promise<T>, tryOnly = false, wait?: WorkspaceSourceLockWait): Promise<T> {
  if (wait && (!Number.isSafeInteger(wait.waitMs) || wait.waitMs < 1 || wait.waitMs > 2_147_483_647))
    throw new Error("workspace_source_wait_invalid");
  const keys = [...new Set([
    ...(input.issueIds ?? []).map(id => `isolated_repository_source:${input.companyId}:${id}`),
    ...(input.projectIds ?? []).map(id => `managed_project_source:${input.companyId}:${id}`),
  ])].sort();
  if (keys.length > 128) throw new Error("workspace_source_scope_limit");
  const held = heldScope();
  if (wait && held?.db === db) throw new Error("workspace_source_wait_requires_independent_scope");
  const acquire = async (scope: SourceLockScope, nonblocking: boolean) => {
    for (const key of keys) {
      if (scope.keys.has(key)) continue;
      const lock = nonblocking ? sql`pg_try_advisory_xact_lock(hashtextextended(${key}, 0))`
        : sql`pg_advisory_xact_lock(hashtextextended(${key}, 0))`;
      const rows = await scope.transaction.execute(sql`select ${lock} as acquired`);
      if (nonblocking && !rows[0]?.acquired) throw new IsolatedRepositorySourceBusy("workspace_source_busy");
      scope.keys.add(key);
    }
  };
  if (held?.db === db) {
    if (new Set([...held.keys, ...keys]).size > 128) throw new Error("workspace_source_scope_limit");
    if (held.companyId !== input.companyId) throw new Error("workspace_source_company_mismatch");
    const onlyProjectKeys = keys.filter(key => !held.keys.has(key)).every(key => key.startsWith("managed_project_source:"));
    const alreadyHasProject = [...held.keys].some(key => key.startsWith("managed_project_source:"));
    await acquire(held, !onlyProjectKeys || alreadyHasProject);
    return action(held.transaction);
  }
  return withDedicatedDbConnection(db, connection => connection.transaction(async transaction => {
    const scope: SourceLockScope = { db, companyId: input.companyId, transaction, keys: new Set<string>(), active: true };
    try {
      if (wait) {
        await acquireQueuedSourceLocks(db, transaction, keys, wait);
        for (const key of keys) scope.keys.add(key);
      } else await acquire(scope, tryOnly);
      const rows = await transaction.execute(sql`select setting::bigint as idle_timeout_ms from pg_settings where name = 'idle_in_transaction_session_timeout'`);
      return await withWorkspaceLockKeepalive(transaction, Number(rows[0]?.idle_timeout_ms),
        () => sourceLocks.run(scope, () => action(transaction)));
    } finally {
      scope.active = false;
    }
  }));
}

export async function managedProjectSourceIdentity(companyId: string, cwd: string): Promise<string | null> {
  const root = await canonicalRepositoryRoot(path.join(resolvePaperclipInstanceRoot(), "projects"));
  const target = await canonicalRepositoryRoot(cwd);
  const relative = path.relative(root, target);
  if (relative === ".." || relative.startsWith(`..${path.sep}`) || path.isAbsolute(relative)) return null;
  const [owner, projectId] = relative.split(path.sep);
  if (owner !== companyId) throw new Error("workspace_source_company_mismatch");
  return projectId && /^[a-f0-9-]{36}$/.test(projectId) ? projectId : null;
}

export async function isolatedTaskSourceIdentity(companyId: string, cwd: string): Promise<string | null> {
  const root = await canonicalRepositoryRoot(resolvePaperclipInstanceRoot());
  const relative = path.relative(root, await canonicalRepositoryRoot(cwd)).split(path.sep);
  if (!["isolated-workspaces", "isolated-workspace-attempts", "isolated-repository-sources"].includes(relative[0])) return null;
  if (relative[0] === "isolated-repository-sources" && relative[3] !== "checkout") return null;
  if (relative[1] !== companyId) throw new Error("workspace_source_company_mismatch");
  return relative[2] && /^[a-f0-9-]{36}$/.test(relative[2]) ? relative[2] : null;
}

export async function retainManagedProjectSource(db: Db, companyId: string, cwd: string): Promise<void> {
  const projectId = await managedProjectSourceIdentity(companyId, cwd);
  if (!projectId) return;
  if (heldScope()?.db !== db) throw new Error("workspace_source_admission_scope_required");
  await withWorkspaceSourceLocks(db, { companyId, projectIds: [projectId] }, async () => undefined, true);
}

export async function workspaceSourceIdentities(companyId: string, cwd: string) {
  cwd = await canonicalRepositoryRoot(cwd);
  const git = await readWorkspaceGitSourceRoots(cwd, { ancestors: true });
  const roots = [cwd, ...(git ? [git.root, git.gitDirectory.path, git.commonDirectory.path,
    ...git.objectDirectories.map(directory => directory.path)] : [])];
  const projectIds = new Set<string>(), issueIds = new Set<string>(), unownedPaths: string[] = [];
  for (const target of new Set(roots)) {
    const projectId = await managedProjectSourceIdentity(companyId, target);
    const issueId = await isolatedTaskSourceIdentity(companyId, target);
    if (projectId) projectIds.add(projectId);
    if (issueId) issueIds.add(issueId);
    if (!projectId && !issueId) unownedPaths.push(target);
  }
  return { projectIds: [...projectIds].sort(), issueIds: [...issueIds].sort(), unownedPaths, git };
}

export async function retainCurrentWorkspaceSource(cwd: string): Promise<void> {
  const scope = heldScope();
  if (!scope) return;
  const identities = await workspaceSourceIdentities(scope.companyId, cwd);
  await withWorkspaceSourceLocks(scope.db, { companyId: scope.companyId, ...identities }, async () => undefined, true);
}

export async function assertCurrentWorkspaceSourceHeld(cwd: string): Promise<void> {
  const scope = heldScope();
  if (!scope) throw new Error("workspace_source_admission_scope_required");
  const identities = await workspaceSourceIdentities(scope.companyId, cwd);
  if (identities.projectIds.some(id => !scope.keys.has(`managed_project_source:${scope.companyId}:${id}`))
    || identities.issueIds.some(id => !scope.keys.has(`isolated_repository_source:${scope.companyId}:${id}`))) {
    throw new IsolatedRepositorySourceBusy("workspace_source_busy");
  }
}