From 7cd2e482d16bdaa7c93c3f6c26b74a2b48990c4c Mon Sep 17 00:00:00 2001 From: Yuval Dinodia <102706514+yetval@users.noreply.github.com> Date: Sat, 11 Jul 2026 19:26:22 -0400 Subject: [PATCH] fix(agents): lease session worktrees during runs (#103263) * fix(agents): lease managed worktrees across a run with an exclusive removal claim A resumed session's managed worktree could be removed mid-run by idle GC, a separate remover, or a manual worktrees.gc, destroying in-flight work; a session whose worktree was already removed was admitted anyway and ran without its checkout. Runs now take a shared run-lease and removal takes an exclusive claim, both held as rows in the shared state_leases table so admission and removal serialize through one BEGIN IMMEDIATE transaction. Lease rows carry an opaque token plus pid and process start time; dead or reused-pid owners are pruned inside the same transaction. Removal claims before any cleanliness or snapshot work and rejects while a live run lease exists; a run releases only its own token. The worktree id is resolved from the session binding or by containment over both the run cwd and workspace dir, so a workspace-only child still leases. The git worktree lock is retained as a secondary raw-Git guard and refcounted per process so a parent and a same-process child do not unlock each other. * fix(agents): serialize worktree git lock cleanup * test(agents): seed worktree races in sqlite --------- Co-authored-by: Peter Steinberger --- src/agents/agent-command.ts | 224 ++++++----- src/agents/worktrees/git-lock.ts | 52 +++ src/agents/worktrees/registry.ts | 274 ++++++++++++++ src/agents/worktrees/run-lease.test.ts | 331 ++++++++++++++++ src/agents/worktrees/run-lease.ts | 352 ++++++++++++++++++ .../worktrees/service.remove-lease.test.ts | 104 ++++++ src/agents/worktrees/service.ts | 191 +++++----- src/commands/agent.worktree-race.test.ts | 204 ++++++++++ 8 files changed, 1535 insertions(+), 197 deletions(-) create mode 100644 src/agents/worktrees/git-lock.ts create mode 100644 src/agents/worktrees/run-lease.test.ts create mode 100644 src/agents/worktrees/run-lease.ts create mode 100644 src/agents/worktrees/service.remove-lease.test.ts create mode 100644 src/commands/agent.worktree-race.test.ts diff --git a/src/agents/agent-command.ts b/src/agents/agent-command.ts index 66d1b41f01a5..66bf30538efe 100644 --- a/src/agents/agent-command.ts +++ b/src/agents/agent-command.ts @@ -175,6 +175,7 @@ import { resolveCandidateThinkingLevel, resolveEffectiveAgentRuntime } from "./t import { resolveAgentTimeoutMs } from "./timeout.js"; import { hasNonzeroUsage } from "./usage.js"; import { ensureAgentWorkspace } from "./workspace.js"; +import { acquireWorktreeRunLease, resolveWorktreeIdForPath } from "./worktrees/run-lease.js"; const log = createSubsystemLogger("agents/agent-command"); @@ -902,61 +903,77 @@ async function prepareAgentCommandExecution(opts: AgentCommandOpts, runtime: Run if (opts.thinkingOnce && !thinkOnce) { throw new Error(`Invalid one-shot thinking level. Use one of: ${thinkingLevelsHint}.`); } - await ensureAgentWorkspace({ - dir: workspaceDirRaw, - ensureBootstrapFiles: !agentCfg?.skipBootstrap, - skipOptionalBootstrapFiles: agentCfg?.skipOptionalBootstrapFiles, - }); - const runId = opts.runId?.trim() || sessionId; - const { getAcpSessionManager } = await loadAcpManagerRuntime(); - const acpManager = getAcpSessionManager(); - const acpResolution = sessionKey - ? acpManager.resolveSession({ - cfg, - sessionKey, - }) - : null; - const body = - !isRawModelRun && acpResolution?.kind === "ready" - ? resolveAcpPromptBody(message, opts.internalEvents) - : prependInternalEventContext(message, opts.internalEvents); - const transcriptBody = - opts.transcriptMessage ?? resolveInternalEventTranscriptBody(message, opts.internalEvents); - - return { - opts: commandOpts, - body, - transcriptBody, - cfg, - configuredThinkingCatalog, - normalizedSpawned, - agentCfg, - thinkOverride, - thinkOnce, - verboseOverride, - timeoutMs, - runTimeoutOverrideMs, - sessionId, - sessionKey, + // Lease the managed worktree before any workspace-dependent preparation so idle + // GC or a separate remover cannot delete the checkout while it is being set up; + // a session bound to a removed worktree fails closed here. Released on any prepare + // failure below and in the run's outer cleanup. + const resolvedCwd = cwd ? resolveUserPath(cwd) : undefined; + const worktreeId = await resolveWorktreeIdForPath({ sessionEntry: sessionEntryRaw, - sessionStore, - storePath, - isNewSession, - persistedThinking, - persistedVerbose, - sessionAgentId, - outboundSession, - workspaceDir, - cwd: cwd ? resolveUserPath(cwd) : undefined, - agentDir, - pluginsEnabled, - manifestMetadataSnapshot, - modelManifestContext, - runId, - isSubagentLane, - acpManager, - acpResolution, - }; + candidatePaths: [resolvedCwd ?? workspaceDir, workspaceDir], + }); + const runLease = worktreeId ? await acquireWorktreeRunLease(worktreeId) : undefined; + try { + await ensureAgentWorkspace({ + dir: workspaceDirRaw, + ensureBootstrapFiles: !agentCfg?.skipBootstrap, + skipOptionalBootstrapFiles: agentCfg?.skipOptionalBootstrapFiles, + }); + const runId = opts.runId?.trim() || sessionId; + const { getAcpSessionManager } = await loadAcpManagerRuntime(); + const acpManager = getAcpSessionManager(); + const acpResolution = sessionKey + ? acpManager.resolveSession({ + cfg, + sessionKey, + }) + : null; + const body = + !isRawModelRun && acpResolution?.kind === "ready" + ? resolveAcpPromptBody(message, opts.internalEvents) + : prependInternalEventContext(message, opts.internalEvents); + const transcriptBody = + opts.transcriptMessage ?? resolveInternalEventTranscriptBody(message, opts.internalEvents); + + return { + opts: commandOpts, + body, + transcriptBody, + cfg, + configuredThinkingCatalog, + normalizedSpawned, + agentCfg, + thinkOverride, + thinkOnce, + verboseOverride, + timeoutMs, + runTimeoutOverrideMs, + sessionId, + sessionKey, + sessionEntry: sessionEntryRaw, + sessionStore, + storePath, + isNewSession, + persistedThinking, + persistedVerbose, + sessionAgentId, + outboundSession, + workspaceDir, + cwd: resolvedCwd, + agentDir, + pluginsEnabled, + manifestMetadataSnapshot, + modelManifestContext, + runId, + isSubagentLane, + acpManager, + acpResolution, + runLease, + }; + } catch (error) { + await runLease?.release(); + throw error; + } } async function agentCommandInternal( @@ -1008,9 +1025,9 @@ async function agentCommandInternal( pluginsEnabled, manifestMetadataSnapshot, modelManifestContext, + runLease, } = prepared; let lifecycleGeneration = opts.lifecycleGeneration ?? captureAgentRunLifecycleGeneration(runId); - assertAgentRunLifecycleGenerationCurrent(lifecycleGeneration); const effectiveCwd = cwd ? resolveUserPath(cwd) : workspaceDir; let sessionEntry = prepared.sessionEntry; const sessionStateActor = classifySessionStateActor({ @@ -1026,7 +1043,6 @@ async function agentCommandInternal( let trackedRestartRecoveryDeliveryContext = false; let currentRunDeliveryContext: DeliveryContext | undefined; const preparedSessionId = sessionEntry?.sessionId; - const sessionStoreRuntime = storePath && sessionKey ? await loadSessionStoreRuntime() : undefined; const internalModelRunTargets = initialOpts.modelRun === true && suppressVisibleSessionEffects ? new Map() @@ -1043,49 +1059,56 @@ async function agentCommandInternal( ); } - // Reset marks its mutation before interrupting work. An aborted run must not - // queue behind that mutation or reset would wait on the run holding the queue. - const sessionWorkAdmission = await beginSessionWorkAdmission({ - scope: storePath ?? `agent:${sessionAgentId}`, - identities: [sessionKey, sessionId], - signal: opts.abortSignal, - onInterrupt: () => lifecycleAbortController.abort(createAgentRunRestartAbortError()), - assertAllowed: () => { - const currentEntry = - sessionStoreRuntime && storePath && sessionKey - ? sessionStoreRuntime.loadSessionEntry({ - storePath, - sessionKey, - readConsistency: "latest", - }) - : sessionEntry; - if (!currentEntry && preparedSessionId) { - throw new Error(`Session "${sessionKey ?? sessionId}" changed while starting work. Retry.`); - } - const matchesIntentionalRollover = - isNewSession && currentEntry?.sessionId === preparedSessionId; - if (currentEntry && currentEntry.sessionId !== sessionId && !matchesIntentionalRollover) { - throw new Error(`Session "${sessionKey ?? sessionId}" changed while starting work. Retry.`); - } - const archivedSessionError = resolveSessionWorkStartError( - sessionKey ?? sessionId, - currentEntry, - ); - if (archivedSessionError) { - throw new Error(archivedSessionError); - } - sessionEntry = currentEntry; - if (sessionStore && sessionKey) { - if (currentEntry) { - sessionStore[sessionKey] = currentEntry; - } else { - delete sessionStore[sessionKey]; - } - } - }, - }); - + let sessionWorkAdmission: Awaited> | undefined; try { + assertAgentRunLifecycleGenerationCurrent(lifecycleGeneration); + const sessionStoreRuntime = + storePath && sessionKey ? await loadSessionStoreRuntime() : undefined; + // Reset marks its mutation before interrupting work. An aborted run must not + // queue behind that mutation or reset would wait on the run holding the queue. + sessionWorkAdmission = await beginSessionWorkAdmission({ + scope: storePath ?? `agent:${sessionAgentId}`, + identities: [sessionKey, sessionId], + signal: opts.abortSignal, + onInterrupt: () => lifecycleAbortController.abort(createAgentRunRestartAbortError()), + assertAllowed: () => { + const currentEntry = + sessionStoreRuntime && storePath && sessionKey + ? sessionStoreRuntime.loadSessionEntry({ + storePath, + sessionKey, + readConsistency: "latest", + }) + : sessionEntry; + if (!currentEntry && preparedSessionId) { + throw new Error( + `Session "${sessionKey ?? sessionId}" changed while starting work. Retry.`, + ); + } + const matchesIntentionalRollover = + isNewSession && currentEntry?.sessionId === preparedSessionId; + if (currentEntry && currentEntry.sessionId !== sessionId && !matchesIntentionalRollover) { + throw new Error( + `Session "${sessionKey ?? sessionId}" changed while starting work. Retry.`, + ); + } + const archivedSessionError = resolveSessionWorkStartError( + sessionKey ?? sessionId, + currentEntry, + ); + if (archivedSessionError) { + throw new Error(archivedSessionError); + } + sessionEntry = currentEntry; + if (sessionStore && sessionKey) { + if (currentEntry) { + sessionStore[sessionKey] = currentEntry; + } else { + delete sessionStore[sessionKey]; + } + } + }, + }); return await sessionWorkAdmission.run(async () => { if (opts.deliver === true) { const sendPolicy = resolveSendPolicy({ @@ -2821,7 +2844,10 @@ async function agentCommandInternal( } }); } finally { - sessionWorkAdmission.release(); + if (runLease) { + await runLease.release(); + } + sessionWorkAdmission?.release(); if (internalModelRunTargets) { // Compaction may rotate a private session identity. Remove every owned // SQLite row only after delivery; transcript and trajectory rows cascade. diff --git a/src/agents/worktrees/git-lock.ts b/src/agents/worktrees/git-lock.ts new file mode 100644 index 000000000000..96ea233026fd --- /dev/null +++ b/src/agents/worktrees/git-lock.ts @@ -0,0 +1,52 @@ +import path from "node:path"; +import { isPidDefinitelyDead } from "../../shared/pid-alive.js"; +import { commandError, listGitWorktrees, runGit } from "./git.js"; +import type { ManagedWorktreeRecord } from "./types.js"; + +const OPENCLAW_LOCK_PATTERN = /^openclaw pid=(\d+)$/; + +export type LockState = + | { kind: "none" } + | { kind: "live"; pid: number } + | { kind: "dead"; pid: number } + | { kind: "foreign"; reason: string }; + +export async function lockState(record: ManagedWorktreeRecord): Promise { + const entry = (await listGitWorktrees(record.repoRoot)).find( + (candidate) => path.resolve(candidate.path) === path.resolve(record.path), + ); + if (!entry || entry.lockedReason === undefined) { + return { kind: "none" }; + } + const match = OPENCLAW_LOCK_PATTERN.exec(entry.lockedReason); + if (!match) { + return { kind: "foreign", reason: entry.lockedReason }; + } + const pid = Number(match[1]); + // A cross-user (EPERM) OpenClaw lock is treated as live so a run's checkout is + // never removed under it; only an ESRCH/zombie owner counts as dead. + return isPidDefinitelyDead(pid) ? { kind: "dead", pid } : { kind: "live", pid }; +} + +export async function lockWorktreeForProcess(record: ManagedWorktreeRecord): Promise { + const result = await runGit(record.repoRoot, [ + "worktree", + "lock", + "--reason", + `openclaw pid=${process.pid}`, + record.path, + ]); + if (result.code !== 0) { + const state = await lockState(record); + if (state.kind !== "live" || state.pid !== process.pid) { + throw commandError("git worktree lock", result); + } + } +} + +export async function unlockWorktree(record: ManagedWorktreeRecord): Promise { + const result = await runGit(record.repoRoot, ["worktree", "unlock", record.path]); + if (result.code !== 0) { + throw commandError("git worktree unlock", result); + } +} diff --git a/src/agents/worktrees/registry.ts b/src/agents/worktrees/registry.ts index 36be5d3c1c0e..59bd4e0257ef 100644 --- a/src/agents/worktrees/registry.ts +++ b/src/agents/worktrees/registry.ts @@ -1,6 +1,7 @@ import type { DatabaseSync } from "node:sqlite"; import type { Insertable, Selectable } from "kysely"; import { executeSqliteQuerySync, getNodeSqliteKysely } from "../../infra/kysely-sync.js"; +import { isLockOwnerDefinitelyStale } from "../../infra/stale-lock-file.js"; import type { DB as OpenClawStateKyselyDatabase } from "../../state/openclaw-state-db.generated.js"; import { openOpenClawStateDatabase, @@ -11,6 +12,7 @@ import type { ManagedWorktreeOwnerKind, ManagedWorktreeRecord } from "./types.js type WorktreesTable = OpenClawStateKyselyDatabase["worktrees"]; type WorktreeRow = Selectable; type WorktreeRegistryDatabase = Pick; +type WorktreeLeaseDatabase = Pick; function dbFor(env: NodeJS.ProcessEnv): DatabaseSync { return openOpenClawStateDatabase({ env }).db; @@ -20,6 +22,10 @@ function kyselyFor(db: DatabaseSync) { return getNodeSqliteKysely(db); } +function kyselyLeaseFor(db: DatabaseSync) { + return getNodeSqliteKysely(db); +} + function rowToRecord(row: WorktreeRow): ManagedWorktreeRecord { return { id: row.id, @@ -164,3 +170,271 @@ export function deleteRegistryWorktree(env: NodeJS.ProcessEnv, id: string): void executeSqliteQuerySync(db, kyselyFor(db).deleteFrom("worktrees").where("id", "=", id)); }); } + +const WORKTREE_RUN_LEASE_SCOPE_PREFIX = "worktree-run:"; +const WORKTREE_REMOVING_LEASE_KEY = "__removing__"; + +export type RunLeaseOwnerChecks = { + isPidDefinitelyDead?: (pid: number) => boolean; + getProcessStartTime?: (pid: number) => number | null; +}; + +function worktreeRunLeaseScope(worktreeId: string): string { + return `${WORKTREE_RUN_LEASE_SCOPE_PREFIX}${worktreeId}`; +} + +function parseLeaseOwnerPayload(payloadJson: string | null): { pid?: number; starttime?: number } { + if (!payloadJson) { + return {}; + } + try { + const parsed = JSON.parse(payloadJson) as Record; + return { + pid: typeof parsed.pid === "number" ? parsed.pid : undefined, + starttime: typeof parsed.starttime === "number" ? parsed.starttime : undefined, + }; + } catch { + return {}; + } +} + +type ScopeLeaseState = { livePids: number[]; removingToken?: string }; + +function collectLiveRunLeases( + db: DatabaseSync, + k: ReturnType, + scope: string, + checks: RunLeaseOwnerChecks, +): ScopeLeaseState { + const rows = executeSqliteQuerySync( + db, + k + .selectFrom("state_leases") + .select(["lease_key", "owner", "payload_json"]) + .where("scope", "=", scope), + ).rows; + const livePids: number[] = []; + const staleKeys: string[] = []; + let removingToken: string | undefined; + for (const row of rows) { + const payload = parseLeaseOwnerPayload(row.payload_json); + const stale = isLockOwnerDefinitelyStale({ + payload, + isPidDefinitelyDead: checks.isPidDefinitelyDead, + getProcessStartTime: checks.getProcessStartTime, + }); + if (row.lease_key === WORKTREE_REMOVING_LEASE_KEY) { + // A removal marker whose remover process died before finalize must self-heal, + // otherwise a still-live worktree stays permanently unadmittable. A live marker + // carries the owning claim token so a competing remover is rejected. + if (stale) { + staleKeys.push(row.lease_key); + } else { + removingToken = row.owner; + } + continue; + } + if (stale) { + staleKeys.push(row.lease_key); + continue; + } + if (payload.pid !== undefined) { + livePids.push(payload.pid); + } + } + if (staleKeys.length > 0) { + executeSqliteQuerySync( + db, + k.deleteFrom("state_leases").where("scope", "=", scope).where("lease_key", "in", staleKeys), + ); + } + return { livePids, ...(removingToken !== undefined ? { removingToken } : {}) }; +} + +export function admitWorktreeRunLeaseRow( + env: NodeJS.ProcessEnv, + params: { + worktreeId: string; + token: string; + pid: number; + startTime: number | null; + now: number; + checks?: RunLeaseOwnerChecks; + }, +): void { + runOpenClawStateWriteTransaction( + (database) => { + const db = database.db; + const k = kyselyLeaseFor(db); + const scope = worktreeRunLeaseScope(params.worktreeId); + const record = executeSqliteQuerySync( + db, + k + .selectFrom("worktrees") + .select(["path", "removed_at"]) + .where("id", "=", params.worktreeId), + ).rows[0]; + const worktreePath = record?.path ?? params.worktreeId; + if (!record || record.removed_at != null) { + throw new Error(`managed worktree was removed: ${worktreePath}`); + } + const { removingToken } = collectLiveRunLeases(db, k, scope, params.checks ?? {}); + if (removingToken !== undefined) { + throw new Error(`managed worktree was removed: ${worktreePath}`); + } + executeSqliteQuerySync( + db, + k.insertInto("state_leases").values({ + scope, + lease_key: params.token, + owner: `${params.pid}:${params.startTime ?? ""}`, + expires_at: null, + heartbeat_at: null, + payload_json: JSON.stringify({ + pid: params.pid, + starttime: params.startTime ?? undefined, + }), + created_at: params.now, + updated_at: params.now, + }), + ); + }, + { env }, + ); +} + +export function claimWorktreeRemovalRow( + env: NodeJS.ProcessEnv, + params: { + worktreeId: string; + token: string; + force: boolean; + pid: number; + startTime: number | null; + now: number; + checks?: RunLeaseOwnerChecks; + }, +): void { + runOpenClawStateWriteTransaction( + (database) => { + const db = database.db; + const k = kyselyLeaseFor(db); + const scope = worktreeRunLeaseScope(params.worktreeId); + const { livePids, removingToken } = collectLiveRunLeases(db, k, scope, params.checks ?? {}); + if (!params.force && livePids.length > 0) { + throw new Error(`worktree is busy: locked by live pid ${livePids[0]}`); + } + // The removal claim is exclusive: a live marker owned by a different token means + // another remover is mid-operation, so this remover must not enter it too. + if (removingToken !== undefined && removingToken !== params.token) { + throw new Error("worktree removal is already in progress"); + } + const payloadJson = JSON.stringify({ + pid: params.pid, + starttime: params.startTime ?? undefined, + }); + executeSqliteQuerySync( + db, + k + .insertInto("state_leases") + .values({ + scope, + lease_key: WORKTREE_REMOVING_LEASE_KEY, + owner: params.token, + expires_at: null, + heartbeat_at: null, + payload_json: payloadJson, + created_at: params.now, + updated_at: params.now, + }) + .onConflict((conflict) => + conflict.columns(["scope", "lease_key"]).doUpdateSet({ + owner: params.token, + payload_json: payloadJson, + updated_at: params.now, + }), + ), + ); + }, + { env }, + ); +} + +export function releaseWorktreeRunLeaseRow( + env: NodeJS.ProcessEnv, + worktreeId: string, + token: string, +): void { + const db = dbFor(env); + runOpenClawStateWriteTransaction( + () => { + executeSqliteQuerySync( + db, + kyselyLeaseFor(db) + .deleteFrom("state_leases") + .where("scope", "=", worktreeRunLeaseScope(worktreeId)) + .where("lease_key", "=", token), + ); + }, + { env }, + ); +} + +export function finalizeWorktreeRemovalRows(env: NodeJS.ProcessEnv, worktreeId: string): void { + const db = dbFor(env); + runOpenClawStateWriteTransaction( + () => { + executeSqliteQuerySync( + db, + kyselyLeaseFor(db) + .deleteFrom("state_leases") + .where("scope", "=", worktreeRunLeaseScope(worktreeId)), + ); + }, + { env }, + ); +} + +export function abortWorktreeRemovalRow( + env: NodeJS.ProcessEnv, + worktreeId: string, + token: string, +): void { + const db = dbFor(env); + runOpenClawStateWriteTransaction( + () => { + // Owner-scoped: only the claim that still owns the marker may clear it, so a slow + // remover cannot delete a marker a newer remover established after replacing it. + executeSqliteQuerySync( + db, + kyselyLeaseFor(db) + .deleteFrom("state_leases") + .where("scope", "=", worktreeRunLeaseScope(worktreeId)) + .where("lease_key", "=", WORKTREE_REMOVING_LEASE_KEY) + .where("owner", "=", token), + ); + }, + { env }, + ); +} + +export function hasLiveWorktreeRunLeaseRow( + env: NodeJS.ProcessEnv, + worktreeId: string, + checks?: RunLeaseOwnerChecks, +): boolean { + return runOpenClawStateWriteTransaction( + (database) => { + const db = database.db; + const k = kyselyLeaseFor(db); + const { livePids } = collectLiveRunLeases( + db, + k, + worktreeRunLeaseScope(worktreeId), + checks ?? {}, + ); + return livePids.length > 0; + }, + { env }, + ); +} diff --git a/src/agents/worktrees/run-lease.test.ts b/src/agents/worktrees/run-lease.test.ts new file mode 100644 index 000000000000..0942f6be70d1 --- /dev/null +++ b/src/agents/worktrees/run-lease.test.ts @@ -0,0 +1,331 @@ +import { execFile } from "node:child_process"; +import fs from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; +import { promisify } from "node:util"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { closeOpenClawStateDatabaseForTest } from "../../state/openclaw-state-db.js"; +import { lockState, unlockWorktree } from "./git-lock.js"; +import { + admitWorktreeRunLeaseRow, + getRegistryWorktree, + releaseWorktreeRunLeaseRow, +} from "./registry.js"; +import { + __testing, + abortWorktreeRemoval, + acquireWorktreeRunLease, + claimWorktreeRemoval, + hasLiveWorktreeRunLease, + resolveWorktreeIdForPath, +} from "./run-lease.js"; +import { ManagedWorktreeService } from "./service.js"; + +const execFileAsync = promisify(execFile); + +async function git(cwd: string, ...args: string[]): Promise { + const { stdout } = await execFileAsync("git", ["-C", cwd, ...args], { encoding: "utf8" }); + return stdout.trim(); +} + +async function initializeRepository(root: string): Promise { + const repo = path.join(root, "repo"); + await fs.mkdir(repo, { recursive: true }); + await git(repo, "init", "-b", "main"); + await git(repo, "config", "user.name", "OpenClaw Test"); + await git(repo, "config", "user.email", "openclaw-test@example.invalid"); + await fs.writeFile(path.join(repo, "README.md"), "base\n"); + await git(repo, "add", "README.md"); + await git(repo, "commit", "-m", "initial"); + return await fs.realpath(repo); +} + +describe("worktree run lease", () => { + let root: string; + let repo: string; + let env: NodeJS.ProcessEnv; + let service: ManagedWorktreeService; + + beforeEach(async () => { + const tempRoot = await fs.realpath(os.tmpdir()); + root = await fs.mkdtemp(path.join(tempRoot, "openclaw-run-lease-")); + repo = await initializeRepository(root); + env = { ...process.env, OPENCLAW_STATE_DIR: path.join(root, "openclaw-state") }; + service = new ManagedWorktreeService({ env }); + }); + + afterEach(async () => { + __testing.resetForTest(); + closeOpenClawStateDatabaseForTest(); + await fs.rm(root, { recursive: true, force: true }); + }); + + async function createSessionWorktree(): Promise<{ id: string; path: string }> { + const created = await service.create({ + repoRoot: repo, + name: "run-lease-session", + ownerKind: "session", + ownerId: "agent:main:run-lease", + }); + return { id: created.id, path: created.path }; + } + + it("shares one worktree across concurrent runs and refcounts the git lock", async () => { + const created = await createSessionWorktree(); + const parent = await acquireWorktreeRunLease(created.id, { env }); + const child = await acquireWorktreeRunLease(created.id, { env }); + + const record = getRegistryWorktree(env, created.id); + expect(record).toBeDefined(); + expect(await lockState(record!)).toEqual({ kind: "live", pid: process.pid }); + expect(hasLiveWorktreeRunLease(env, created.id)).toBe(true); + + await parent.release(); + expect(await lockState(record!)).toEqual({ kind: "live", pid: process.pid }); + expect(hasLiveWorktreeRunLease(env, created.id)).toBe(true); + + await child.release(); + expect(await lockState(record!)).toEqual({ kind: "none" }); + expect(hasLiveWorktreeRunLease(env, created.id)).toBe(false); + }); + + it("resolves the worktree id for a nested workspace path with no session binding", async () => { + const created = await createSessionWorktree(); + const nested = path.join(created.path, "workspace"); + await fs.mkdir(nested); + + const resolved = await resolveWorktreeIdForPath({ candidatePaths: [nested], env }); + expect(resolved).toBe(created.id); + + const lease = await acquireWorktreeRunLease(created.id, { env }); + expect(() => + claimWorktreeRemoval(env, { worktreeId: created.id, token: "remover", force: false }), + ).toThrow("worktree is busy"); + await lease.release(); + }); + + it("prunes a dead owner lease so removal can proceed", async () => { + const created = await createSessionWorktree(); + admitWorktreeRunLeaseRow(env, { + worktreeId: created.id, + token: "dead-owner", + pid: 987_654, + startTime: 4242, + now: 1, + }); + __testing.setDeadPidResolverForTest((pid) => pid === 987_654); + + expect(hasLiveWorktreeRunLease(env, created.id)).toBe(false); + expect(() => + claimWorktreeRemoval(env, { worktreeId: created.id, token: "remover", force: false }), + ).not.toThrow(); + }); + + it("prunes a reused pid whose start time no longer matches", async () => { + const created = await createSessionWorktree(); + admitWorktreeRunLeaseRow(env, { + worktreeId: created.id, + token: "reused-pid", + pid: process.pid, + startTime: 111, + now: 1, + }); + __testing.setProcessStartTimeResolverForTest(() => 222); + + expect(hasLiveWorktreeRunLease(env, created.id)).toBe(false); + + const lease = await acquireWorktreeRunLease(created.id, { env }); + expect(lease.token).not.toBe("reused-pid"); + expect(hasLiveWorktreeRunLease(env, created.id)).toBe(true); + await lease.release(); + expect(hasLiveWorktreeRunLease(env, created.id)).toBe(false); + }); + + it("rejects removal of a live lease unless forced", async () => { + const created = await createSessionWorktree(); + const lease = await acquireWorktreeRunLease(created.id, { env }); + + expect(() => + claimWorktreeRemoval(env, { worktreeId: created.id, token: "remover", force: false }), + ).toThrow("worktree is busy"); + expect(() => + claimWorktreeRemoval(env, { worktreeId: created.id, token: "remover", force: true }), + ).not.toThrow(); + await lease.release(); + }); + + it("fails admission once a removal claim is held", async () => { + const created = await createSessionWorktree(); + claimWorktreeRemoval(env, { worktreeId: created.id, token: "remover", force: true }); + + await expect(acquireWorktreeRunLease(created.id, { env })).rejects.toThrow( + `managed worktree was removed: ${created.path}`, + ); + }); + + it("recovers admission when the remover died before finalizing the removal", async () => { + const created = await createSessionWorktree(); + claimWorktreeRemoval(env, { worktreeId: created.id, token: "remover", force: true }); + __testing.setDeadPidResolverForTest((pid) => pid === process.pid); + + const lease = await acquireWorktreeRunLease(created.id, { env }); + expect(lease.token).toBeTruthy(); + await lease.release(); + }); + + it("rejects a second live remover until the first releases, even with force", async () => { + const created = await createSessionWorktree(); + claimWorktreeRemoval(env, { worktreeId: created.id, token: "remover-a", force: false }); + + expect(() => + claimWorktreeRemoval(env, { worktreeId: created.id, token: "remover-b", force: false }), + ).toThrow("worktree removal is already in progress"); + expect(() => + claimWorktreeRemoval(env, { worktreeId: created.id, token: "remover-b", force: true }), + ).toThrow("worktree removal is already in progress"); + + abortWorktreeRemoval(env, created.id, "remover-a"); + expect(() => + claimWorktreeRemoval(env, { worktreeId: created.id, token: "remover-b", force: false }), + ).not.toThrow(); + }); + + it("recovers a transient release failure within a single release call", async () => { + const created = await createSessionWorktree(); + const lease = await acquireWorktreeRunLease(created.id, { env }); + const record = getRegistryWorktree(env, created.id)!; + + let attempts = 0; + __testing.setReleaseRowImplForTest((rowEnv, id, token) => { + attempts += 1; + if (attempts === 1) { + throw new Error("simulated state database failure"); + } + releaseWorktreeRunLeaseRow(rowEnv, id, token); + }); + + await lease.release(); + expect(attempts).toBeGreaterThanOrEqual(2); + expect(hasLiveWorktreeRunLease(env, created.id)).toBe(false); + expect(await lockState(record)).toEqual({ kind: "none" }); + }); + + it("retains the lease and git guard across sustained delete failures, freeing on a lifecycle retry", async () => { + const created = await createSessionWorktree(); + const lease = await acquireWorktreeRunLease(created.id, { env }); + const record = getRegistryWorktree(env, created.id)!; + + let fail = true; + __testing.setReleaseRowImplForTest((rowEnv, id, token) => { + if (fail) { + throw new Error("simulated state database failure"); + } + releaseWorktreeRunLeaseRow(rowEnv, id, token); + }); + + await lease.release(); + expect(hasLiveWorktreeRunLease(env, created.id)).toBe(true); + expect(await lockState(record)).toEqual({ kind: "live", pid: process.pid }); + expect(() => + claimWorktreeRemoval(env, { worktreeId: created.id, token: "remover", force: false }), + ).toThrow("worktree is busy"); + + fail = false; + await __testing.drainPendingCleanupsForTest(); + expect(hasLiveWorktreeRunLease(env, created.id)).toBe(false); + expect(await lockState(record)).toEqual({ kind: "none" }); + expect(() => + claimWorktreeRemoval(env, { worktreeId: created.id, token: "remover", force: false }), + ).not.toThrow(); + }); + + it("serializes overlapping same-process acquisitions so the guard holds until the last release", async () => { + const created = await createSessionWorktree(); + const record = getRegistryWorktree(env, created.id)!; + + const [first, second] = await Promise.all([ + acquireWorktreeRunLease(created.id, { env }), + acquireWorktreeRunLease(created.id, { env }), + ]); + + expect(await lockState(record)).toEqual({ kind: "live", pid: process.pid }); + await first.release(); + expect(await lockState(record)).toEqual({ kind: "live", pid: process.pid }); + await second.release(); + expect(await lockState(record)).toEqual({ kind: "none" }); + }); + + it("retains the git guard when unlock fails, releasing it on a lifecycle retry", async () => { + const created = await createSessionWorktree(); + const lease = await acquireWorktreeRunLease(created.id, { env }); + const record = getRegistryWorktree(env, created.id)!; + + let failUnlock = true; + __testing.setUnlockImplForTest(async (rec) => { + if (failUnlock) { + throw new Error("simulated git unlock failure"); + } + await unlockWorktree(rec); + }); + + await lease.release(); + expect(hasLiveWorktreeRunLease(env, created.id)).toBe(false); + expect(await lockState(record)).toEqual({ kind: "live", pid: process.pid }); + + failUnlock = false; + await __testing.drainPendingCleanupsForTest(); + expect(await lockState(record)).toEqual({ kind: "none" }); + }); + + it("does not let a failed cleanup unlock a newer holder generation", async () => { + const created = await createSessionWorktree(); + const first = await acquireWorktreeRunLease(created.id, { env }); + const record = getRegistryWorktree(env, created.id)!; + + let failUnlock = true; + __testing.setUnlockImplForTest(async (rec) => { + if (failUnlock) { + throw new Error("simulated git unlock failure"); + } + await unlockWorktree(rec); + }); + + await first.release(); + expect(await lockState(record)).toEqual({ kind: "live", pid: process.pid }); + + const second = await acquireWorktreeRunLease(created.id, { env }); + failUnlock = false; + await __testing.drainPendingCleanupsForTest(); + expect(await lockState(record)).toEqual({ kind: "live", pid: process.pid }); + + await second.release(); + expect(await lockState(record)).toEqual({ kind: "none" }); + }); + + it("fails closed when a session's authoritative worktree binding is removed", async () => { + const created = await createSessionWorktree(); + await service.remove({ id: created.id, reason: "manual-delete", force: true }); + + await expect( + resolveWorktreeIdForPath({ + sessionEntry: { worktree: { id: created.id } }, + candidatePaths: [], + env, + }), + ).rejects.toThrow("managed worktree was removed"); + }); + + it("does not let a superseded remover clear a newer removal claim", async () => { + const created = await createSessionWorktree(); + claimWorktreeRemoval(env, { worktreeId: created.id, token: "remover-a", force: false }); + + __testing.setDeadPidResolverForTest((pid) => pid === process.pid); + claimWorktreeRemoval(env, { worktreeId: created.id, token: "remover-b", force: false }); + __testing.setDeadPidResolverForTest(null); + + abortWorktreeRemoval(env, created.id, "remover-a"); + await expect(acquireWorktreeRunLease(created.id, { env })).rejects.toThrow( + "managed worktree was removed", + ); + }); +}); diff --git a/src/agents/worktrees/run-lease.ts b/src/agents/worktrees/run-lease.ts new file mode 100644 index 000000000000..4a707a9e75c9 --- /dev/null +++ b/src/agents/worktrees/run-lease.ts @@ -0,0 +1,352 @@ +import { randomUUID } from "node:crypto"; +import fs from "node:fs/promises"; +import path from "node:path"; +import { createSubsystemLogger } from "../../logging/subsystem.js"; +import { getFileLockProcessStartTime } from "../../shared/pid-alive.js"; +import { lockWorktreeForProcess, unlockWorktree } from "./git-lock.js"; +import { + abortWorktreeRemovalRow, + admitWorktreeRunLeaseRow, + claimWorktreeRemovalRow, + finalizeWorktreeRemovalRows, + getRegistryWorktree, + hasLiveWorktreeRunLeaseRow, + listRegistryWorktrees, + releaseWorktreeRunLeaseRow, + type RunLeaseOwnerChecks, +} from "./registry.js"; + +const log = createSubsystemLogger("agents/worktrees"); + +const RELEASE_MAX_ATTEMPTS = 3; + +export type WorktreeRunLease = { + id: string; + token: string; + release: () => Promise; +}; + +type HeldWorktreeLock = { refcount: number; gitLocked: boolean }; + +// The git lock is a per-process single-holder resource; a parent and a same-process +// child that share one worktree refcount it here so a child release does not unlock +// the parent's still-live checkout. +const heldGitLocks = new Map(); +const gitLockTransitionTails = new Map>(); +let ownerChecks: RunLeaseOwnerChecks = {}; +let resolveSelfStartTime = getFileLockProcessStartTime; +let releaseRunLeaseRow = releaseWorktreeRunLeaseRow; +let unlockWorktreeImpl = unlockWorktree; + +// A cleanup that could not finish (persistent state-database delete or git unlock +// failure) is retained here so the process keeps ownership of it and retries on the +// next lease acquisition and at exit, instead of stranding the row and git guard. +type LeaseCleanup = { + env: NodeJS.ProcessEnv; + id: string; + token: string; + rowDeleted: boolean; + refcountReleased: boolean; + gitUnlockPending: boolean; +}; +const pendingLeaseCleanups = new Set(); +let exitCleanupRegistered = false; + +function errorMessage(error: unknown): string { + return error instanceof Error ? error.message : String(error); +} + +async function withGitLockTransition(id: string, operation: () => Promise): Promise { + const previous = gitLockTransitionTails.get(id) ?? Promise.resolve(); + let finish!: () => void; + const current = new Promise((resolve) => { + finish = resolve; + }); + const tail = previous.then(() => current); + gitLockTransitionTails.set(id, tail); + await previous; + try { + return await operation(); + } finally { + finish(); + if (gitLockTransitionTails.get(id) === tail) { + gitLockTransitionTails.delete(id); + } + } +} + +async function retainGitLock(env: NodeJS.ProcessEnv, id: string): Promise { + await withGitLockTransition(id, async () => { + const held = heldGitLocks.get(id) ?? { refcount: 0, gitLocked: false }; + const needsLock = held.refcount === 0 && !held.gitLocked; + held.refcount += 1; + heldGitLocks.set(id, held); + if (!needsLock) { + return; + } + const record = getRegistryWorktree(env, id); + if (!record) { + return; + } + try { + await lockWorktreeForProcess(record); + held.gitLocked = true; + } catch (error) { + log.warn(`worktree git lock unavailable for ${id}: ${errorMessage(error)}`); + } + }); +} + +async function releaseGitLock(cleanup: LeaseCleanup): Promise { + return await withGitLockTransition(cleanup.id, async () => { + let held = heldGitLocks.get(cleanup.id); + if (!cleanup.refcountReleased) { + cleanup.refcountReleased = true; + if (held) { + held.refcount -= 1; + } + } + held = heldGitLocks.get(cleanup.id); + if (!held) { + cleanup.gitUnlockPending = false; + return true; + } + if (held.refcount > 0) { + // A newer holder adopted a guard whose prior unlock failed. Its own final + // release now owns the unlock; stale cleanup must not drop that generation. + cleanup.gitUnlockPending = false; + return true; + } + if (!held.gitLocked) { + heldGitLocks.delete(cleanup.id); + cleanup.gitUnlockPending = false; + return true; + } + const record = getRegistryWorktree(cleanup.env, cleanup.id); + if (!record) { + heldGitLocks.delete(cleanup.id); + cleanup.gitUnlockPending = false; + return true; + } + try { + await unlockWorktreeImpl(record); + } catch (error) { + cleanup.gitUnlockPending = true; + log.warn(`failed to unlock worktree ${cleanup.id}: ${errorMessage(error)}`); + return false; + } + heldGitLocks.delete(cleanup.id); + cleanup.gitUnlockPending = false; + return true; + }); +} + +async function realpathOrSelf(candidate: string): Promise { + try { + return await fs.realpath(candidate); + } catch { + return path.resolve(candidate); + } +} + +export async function resolveWorktreeIdForPath(params: { + sessionEntry?: { worktree?: { id: string } }; + candidatePaths: Array; + env?: NodeJS.ProcessEnv; +}): Promise { + const env = params.env ?? process.env; + const boundId = params.sessionEntry?.worktree?.id; + if (boundId !== undefined) { + // The session's stored binding is authoritative: if that worktree is gone the + // run must fail closed rather than silently continue as an unmanaged directory. + const record = getRegistryWorktree(env, boundId); + if (!record || record.removedAt !== undefined) { + throw new Error(`managed worktree was removed: ${record?.path ?? boundId}`); + } + return boundId; + } + const records = listRegistryWorktrees(env).filter((record) => record.removedAt === undefined); + if (records.length === 0) { + return undefined; + } + const bases = new Map(); + for (const record of records) { + bases.set(record.id, await realpathOrSelf(record.path)); + } + const seen = new Set(); + for (const candidate of params.candidatePaths) { + if (!candidate || seen.has(candidate)) { + continue; + } + seen.add(candidate); + const real = await realpathOrSelf(candidate); + for (const record of records) { + const base = bases.get(record.id); + if (base && (real === base || real.startsWith(`${base}${path.sep}`))) { + return record.id; + } + } + } + return undefined; +} + +function deleteRunLeaseRowWithRetries(cleanup: LeaseCleanup): boolean { + for (let attempt = 1; attempt <= RELEASE_MAX_ATTEMPTS; attempt += 1) { + try { + releaseRunLeaseRow(cleanup.env, cleanup.id, cleanup.token); + return true; + } catch (error) { + log.warn( + `failed to release worktree run lease for ${cleanup.id} (attempt ${attempt}): ${errorMessage(error)}`, + ); + } + } + return false; +} + +// Drives a lease cleanup as far as it can and returns true only once both the token +// row and the git guard are released. Keeps everything until each step succeeds so a +// removal stays correctly blocked while cleanup is still owed. +async function runLeaseCleanup(cleanup: LeaseCleanup): Promise { + if (!cleanup.rowDeleted) { + if (!deleteRunLeaseRowWithRetries(cleanup)) { + return false; + } + cleanup.rowDeleted = true; + } + return await releaseGitLock(cleanup); +} + +async function drainPendingLeaseCleanups(): Promise { + for (const cleanup of pendingLeaseCleanups) { + if (await runLeaseCleanup(cleanup)) { + pendingLeaseCleanups.delete(cleanup); + } + } +} + +function ensureExitCleanupRegistered(): void { + if (exitCleanupRegistered) { + return; + } + exitCleanupRegistered = true; + // A row that never deleted keeps its worktree unremovable until this process ends; + // delete it synchronously on exit so a live-pid lease row does not linger. + process.on("exit", () => { + for (const cleanup of pendingLeaseCleanups) { + if (!cleanup.rowDeleted) { + try { + releaseRunLeaseRow(cleanup.env, cleanup.id, cleanup.token); + } catch { + // Best effort at exit; the dead pid also lets a later process prune it. + } + } + } + }); +} + +export async function acquireWorktreeRunLease( + id: string, + opts: { env?: NodeJS.ProcessEnv } = {}, +): Promise { + const env = opts.env ?? process.env; + ensureExitCleanupRegistered(); + // Retry any cleanup a prior run could not finish before starting a new one. + await drainPendingLeaseCleanups(); + const token = randomUUID(); + const pid = process.pid; + const startTime = resolveSelfStartTime(pid); + admitWorktreeRunLeaseRow(env, { + worktreeId: id, + token, + pid, + startTime, + now: Date.now(), + checks: ownerChecks, + }); + // Serialize refcount and Git transitions so a cleanup retry cannot unlock a + // newer same-process holder after a prior generation's unlock failed. + await retainGitLock(env, id); + const cleanup: LeaseCleanup = { + env, + id, + token, + rowDeleted: false, + refcountReleased: false, + gitUnlockPending: false, + }; + let released = false; + return { + id, + token, + release: async () => { + if (released) { + return; + } + released = true; + if (!(await runLeaseCleanup(cleanup))) { + pendingLeaseCleanups.add(cleanup); + } + }, + }; +} + +export function claimWorktreeRemoval( + env: NodeJS.ProcessEnv, + params: { worktreeId: string; token: string; force: boolean }, +): void { + const pid = process.pid; + claimWorktreeRemovalRow(env, { + ...params, + pid, + startTime: resolveSelfStartTime(pid), + now: Date.now(), + checks: ownerChecks, + }); +} + +export function finalizeWorktreeRemoval(env: NodeJS.ProcessEnv, worktreeId: string): void { + finalizeWorktreeRemovalRows(env, worktreeId); +} + +export function abortWorktreeRemoval( + env: NodeJS.ProcessEnv, + worktreeId: string, + token: string, +): void { + abortWorktreeRemovalRow(env, worktreeId, token); +} + +export function hasLiveWorktreeRunLease(env: NodeJS.ProcessEnv, worktreeId: string): boolean { + return hasLiveWorktreeRunLeaseRow(env, worktreeId, ownerChecks); +} + +const testing = { + setProcessStartTimeResolverForTest(resolver: ((pid: number) => number | null) | null): void { + resolveSelfStartTime = resolver ?? getFileLockProcessStartTime; + ownerChecks = { ...ownerChecks, getProcessStartTime: resolver ?? undefined }; + }, + setDeadPidResolverForTest(resolver: ((pid: number) => boolean) | null): void { + ownerChecks = { ...ownerChecks, isPidDefinitelyDead: resolver ?? undefined }; + }, + setReleaseRowImplForTest(impl: typeof releaseWorktreeRunLeaseRow | null): void { + releaseRunLeaseRow = impl ?? releaseWorktreeRunLeaseRow; + }, + setUnlockImplForTest(impl: typeof unlockWorktree | null): void { + unlockWorktreeImpl = impl ?? unlockWorktree; + }, + async drainPendingCleanupsForTest(): Promise { + await drainPendingLeaseCleanups(); + }, + resetForTest(): void { + heldGitLocks.clear(); + gitLockTransitionTails.clear(); + pendingLeaseCleanups.clear(); + ownerChecks = {}; + resolveSelfStartTime = getFileLockProcessStartTime; + releaseRunLeaseRow = releaseWorktreeRunLeaseRow; + unlockWorktreeImpl = unlockWorktree; + }, +}; + +export { testing as __testing }; diff --git a/src/agents/worktrees/service.remove-lease.test.ts b/src/agents/worktrees/service.remove-lease.test.ts new file mode 100644 index 000000000000..74dfeb5c6b08 --- /dev/null +++ b/src/agents/worktrees/service.remove-lease.test.ts @@ -0,0 +1,104 @@ +import { execFile } from "node:child_process"; +import fs from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; +import { promisify } from "node:util"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { closeOpenClawStateDatabaseForTest } from "../../state/openclaw-state-db.js"; +import { getRegistryWorktree } from "./registry.js"; +import { __testing, acquireWorktreeRunLease } from "./run-lease.js"; +import { IDLE_GC_MS, ManagedWorktreeService } from "./service.js"; + +const execFileAsync = promisify(execFile); + +async function git(cwd: string, ...args: string[]): Promise { + const { stdout } = await execFileAsync("git", ["-C", cwd, ...args], { encoding: "utf8" }); + return stdout.trim(); +} + +async function initializeRepository(root: string): Promise { + const repo = path.join(root, "repo"); + await fs.mkdir(repo, { recursive: true }); + await git(repo, "init", "-b", "main"); + await git(repo, "config", "user.name", "OpenClaw Test"); + await git(repo, "config", "user.email", "openclaw-test@example.invalid"); + await fs.writeFile(path.join(repo, "README.md"), "base\n"); + await git(repo, "add", "README.md"); + await git(repo, "commit", "-m", "initial"); + return await fs.realpath(repo); +} + +describe("ManagedWorktreeService removal against a live run lease", () => { + let root: string; + let repo: string; + let env: NodeJS.ProcessEnv; + let now: number; + let service: ManagedWorktreeService; + + beforeEach(async () => { + const tempRoot = await fs.realpath(os.tmpdir()); + root = await fs.mkdtemp(path.join(tempRoot, "openclaw-remove-lease-")); + repo = await initializeRepository(root); + env = { ...process.env, OPENCLAW_STATE_DIR: path.join(root, "openclaw-state") }; + now = 1_700_000_000_000; + service = new ManagedWorktreeService({ env, now: () => now }); + }); + + afterEach(async () => { + __testing.resetForTest(); + closeOpenClawStateDatabaseForTest(); + await fs.rm(root, { recursive: true, force: true }); + }); + + async function createSessionWorktree(): Promise<{ id: string; path: string }> { + const created = await service.create({ + repoRoot: repo, + name: "removal-session", + ownerKind: "session", + ownerId: "agent:main:removal", + }); + return { id: created.id, path: created.path }; + } + + it("rejects removal before snapshotting while a run lease is live", async () => { + const created = await createSessionWorktree(); + const lease = await acquireWorktreeRunLease(created.id, { env }); + + await expect(service.remove({ id: created.id, reason: "manual-delete" })).rejects.toThrow( + "worktree is busy", + ); + expect(getRegistryWorktree(env, created.id)?.snapshotRef).toBeUndefined(); + expect(getRegistryWorktree(env, created.id)?.removedAt).toBeUndefined(); + expect(await fs.stat(created.path)).toBeTruthy(); + + await lease.release(); + expect((await service.remove({ id: created.id, reason: "manual-delete" })).removed).toBe(true); + await expect(fs.stat(created.path)).rejects.toMatchObject({ code: "ENOENT" }); + }); + + it("rejects a concurrent second remover while the first holds the claim", async () => { + const created = await createSessionWorktree(); + const first = service.remove({ id: created.id, reason: "manual-delete" }); + + await expect(service.remove({ id: created.id, reason: "manual-delete" })).rejects.toThrow( + /already in progress|unknown active worktree/, + ); + + expect((await first).removed).toBe(true); + await expect(fs.stat(created.path)).rejects.toMatchObject({ code: "ENOENT" }); + }); + + it("skips idle garbage collection for a worktree with a live run lease", async () => { + const created = await createSessionWorktree(); + now += IDLE_GC_MS + 1; + const lease = await acquireWorktreeRunLease(created.id, { env }); + + const skipped = await service.gc(); + expect(skipped.removed).toEqual([]); + expect(getRegistryWorktree(env, created.id)?.removedAt).toBeUndefined(); + + await lease.release(); + const collected = await service.gc(); + expect(collected.removed).toContain(created.id); + }); +}); diff --git a/src/agents/worktrees/service.ts b/src/agents/worktrees/service.ts index 9fd45f46e9b7..88c272a567ec 100644 --- a/src/agents/worktrees/service.ts +++ b/src/agents/worktrees/service.ts @@ -6,6 +6,7 @@ import path from "node:path"; import { resolveStateDir } from "../../config/paths.js"; import { createSubsystemLogger } from "../../logging/subsystem.js"; import { runCommandWithTimeout } from "../../process/exec.js"; +import { lockState, lockWorktreeForProcess, unlockWorktree } from "./git-lock.js"; import { commandError, listGitWorktrees, @@ -26,6 +27,12 @@ import { listRegistryWorktrees, updateRegistryWorktree, } from "./registry.js"; +import { + abortWorktreeRemoval, + claimWorktreeRemoval, + finalizeWorktreeRemoval, + hasLiveWorktreeRunLease, +} from "./run-lease.js"; import type { CreateManagedWorktreeParams, ManagedWorktreeBranch, @@ -51,7 +58,6 @@ export class WorktreeSnapshotError extends Error { } } const SNAPSHOT_REF_PREFIX = "refs/openclaw/snapshots"; -const OPENCLAW_LOCK_PATTERN = /^openclaw pid=(\d+)$/; const log = createSubsystemLogger("agents/worktrees"); type ServiceOptions = { @@ -63,12 +69,6 @@ type ManagedWorktreeGcParams = { isOwnerActive?: (ownerKind: ManagedWorktreeOwnerKind, ownerId: string) => boolean; }; -type LockState = - | { kind: "none" } - | { kind: "live"; pid: number } - | { kind: "dead"; pid: number } - | { kind: "foreign"; reason: string }; - function resultMessage(result: GitResult): string { return (result.stderr || result.stdout).trim().split("\n").slice(-12).join("\n"); } @@ -306,30 +306,6 @@ async function runSetupScript(repoRoot: string, worktreePath: string): Promise { - const entry = (await listGitWorktrees(record.repoRoot)).find( - (candidate) => path.resolve(candidate.path) === path.resolve(record.path), - ); - if (!entry || entry.lockedReason === undefined) { - return { kind: "none" }; - } - const match = OPENCLAW_LOCK_PATTERN.exec(entry.lockedReason); - if (!match) { - return { kind: "foreign", reason: entry.lockedReason }; - } - const pid = Number(match[1]); - return processIsAlive(pid) ? { kind: "live", pid } : { kind: "dead", pid }; -} - async function snapshotWorktree(record: ManagedWorktreeRecord, reason: string): Promise { const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-worktree-index-")); const indexPath = path.join(tempDir, "index"); @@ -581,19 +557,7 @@ export class ManagedWorktreeService { async acquire(id: string): Promise { const record = this.requireLiveRecord(id); - const result = await runGit(record.repoRoot, [ - "worktree", - "lock", - "--reason", - `openclaw pid=${process.pid}`, - record.path, - ]); - if (result.code !== 0) { - const state = await lockState(record); - if (state.kind !== "live" || state.pid !== process.pid) { - throw commandError("git worktree lock", result); - } - } + await lockWorktreeForProcess(record); const lastActiveAt = this.now(); updateRegistryWorktree(this.env, id, { lastActiveAt }); return { ...record, lastActiveAt }; @@ -612,10 +576,7 @@ export class ManagedWorktreeService { return; } if (state.kind !== "none") { - const result = await runGit(record.repoRoot, ["worktree", "unlock", record.path]); - if (result.code !== 0) { - throw commandError("git worktree unlock", result); - } + await unlockWorktree(record); } } @@ -623,50 +584,64 @@ export class ManagedWorktreeService { id: string; reason: string; force?: boolean; + claimToken?: string; }): Promise { const record = this.requireLiveRecord(params.id); - const state = await lockState(record); - if ((state.kind === "live" || state.kind === "foreign") && !params.force) { - throw new Error( - state.kind === "live" - ? `worktree is locked by live OpenClaw pid ${state.pid}` - : `worktree has a foreign lock${state.reason ? `: ${state.reason}` : ""}`, - ); - } - if (state.kind !== "none") { - await requireGit(record.repoRoot, ["worktree", "unlock", record.path]); - } - let snapshotRef = record.snapshotRef; - let snapshotError: string | undefined; + const force = params.force ?? false; + // Claim removal before any cleanliness or snapshot work so a live run lease + // rejects it and an admitted run cannot start once the claim is held. The + // opaque token makes the claim exclusive against competing removers; a caller + // that already claimed (removeIfLossless) passes its token to keep one claim. + const claimToken = params.claimToken ?? randomUUID(); + claimWorktreeRemoval(this.env, { worktreeId: record.id, token: claimToken, force }); try { - snapshotRef = await snapshotWorktree(record, params.reason); - updateRegistryWorktree(this.env, record.id, { snapshotRef }); - } catch (error) { - snapshotError = error instanceof Error ? error.message : String(error); - if (!params.force) { - throw new WorktreeSnapshotError(snapshotError, { cause: error }); + const state = await lockState(record); + if ((state.kind === "live" || state.kind === "foreign") && !force) { + throw new Error( + state.kind === "live" + ? `worktree is locked by live OpenClaw pid ${state.pid}` + : `worktree has a foreign lock${state.reason ? `: ${state.reason}` : ""}`, + ); } + if (state.kind !== "none") { + await requireGit(record.repoRoot, ["worktree", "unlock", record.path]); + } + let snapshotRef = record.snapshotRef; + let snapshotError: string | undefined; + try { + snapshotRef = await snapshotWorktree(record, params.reason); + updateRegistryWorktree(this.env, record.id, { snapshotRef }); + } catch (error) { + snapshotError = error instanceof Error ? error.message : String(error); + if (!force) { + throw new WorktreeSnapshotError(snapshotError, { cause: error }); + } + } + const removed = await runGit(record.repoRoot, ["worktree", "remove", "--force", record.path]); + if (removed.code !== 0) { + throw commandError("git worktree remove", removed); + } + const branchDelete = await runGit(record.repoRoot, ["branch", "-D", record.branch]); + if (branchDelete.code !== 0) { + throw commandError("git branch -D", branchDelete); + } + await requireGit(record.repoRoot, ["worktree", "prune"]); + await removeEmptyParents( + path.dirname(record.path), + path.join(resolveStateDir(this.env), "worktrees"), + ); + const removedAt = this.now(); + updateRegistryWorktree(this.env, record.id, { removedAt, snapshotRef }); + finalizeWorktreeRemoval(this.env, record.id); + return { + removed: true, + ...(snapshotRef ? { snapshotRef } : {}), + ...(snapshotError ? { snapshotError } : {}), + }; + } catch (error) { + abortWorktreeRemoval(this.env, record.id, claimToken); + throw error; } - const removed = await runGit(record.repoRoot, ["worktree", "remove", "--force", record.path]); - if (removed.code !== 0) { - throw commandError("git worktree remove", removed); - } - const branchDelete = await runGit(record.repoRoot, ["branch", "-D", record.branch]); - if (branchDelete.code !== 0) { - throw commandError("git branch -D", branchDelete); - } - await requireGit(record.repoRoot, ["worktree", "prune"]); - await removeEmptyParents( - path.dirname(record.path), - path.join(resolveStateDir(this.env), "worktrees"), - ); - const removedAt = this.now(); - updateRegistryWorktree(this.env, record.id, { removedAt, snapshotRef }); - return { - removed: true, - ...(snapshotRef ? { snapshotRef } : {}), - ...(snapshotError ? { snapshotError } : {}), - }; } async restore(params: { id: string }): Promise { @@ -709,6 +684,9 @@ export class ManagedWorktreeService { } const lastActiveAt = this.now(); updateRegistryWorktree(this.env, params.id, { removedAt: undefined, lastActiveAt }); + // Clear any lease rows or removal marker stranded by a crash between git removal + // and finalize so the restored worktree admits runs again. + finalizeWorktreeRemoval(this.env, params.id); const restored = { ...record, lastActiveAt }; delete restored.removedAt; return restored; @@ -716,19 +694,33 @@ export class ManagedWorktreeService { async removeIfLossless(id: string): Promise { const record = this.requireLiveRecord(id); - const status = await requireGit(record.path, ["status", "--porcelain"]); - const unpushed = await requireGit(record.path, [ - "log", - "HEAD", - "--not", - "--remotes", - "--oneline", - ]); - await this.release(id); - if (status || unpushed) { + const claimToken = randomUUID(); + try { + claimWorktreeRemoval(this.env, { worktreeId: id, token: claimToken, force: false }); + } catch { + // A live run lease or a competing remover holds the worktree; a lossless + // auto-cleanup must not race it. return false; } - await this.remove({ id, reason: "run-end" }); + try { + const status = await requireGit(record.path, ["status", "--porcelain"]); + const unpushed = await requireGit(record.path, [ + "log", + "HEAD", + "--not", + "--remotes", + "--oneline", + ]); + if (status || unpushed) { + abortWorktreeRemoval(this.env, id, claimToken); + return false; + } + } catch (error) { + abortWorktreeRemoval(this.env, id, claimToken); + throw error; + } + await this.release(id); + await this.remove({ id, reason: "run-end", claimToken }); return true; } @@ -770,6 +762,9 @@ export class ManagedWorktreeService { ) { continue; } + if (hasLiveWorktreeRunLease(this.env, record.id)) { + continue; + } const state = await lockState(record); if (state.kind === "live" || state.kind === "foreign") { continue; diff --git a/src/commands/agent.worktree-race.test.ts b/src/commands/agent.worktree-race.test.ts new file mode 100644 index 000000000000..34bff7964d5e --- /dev/null +++ b/src/commands/agent.worktree-race.test.ts @@ -0,0 +1,204 @@ +import { execFile } from "node:child_process"; +import fsSync from "node:fs"; +import fs from "node:fs/promises"; +import path from "node:path"; +import { promisify } from "node:util"; +import { withTempHome } from "openclaw/plugin-sdk/test-env"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import "./agent-command.test-mocks.js"; +import { ensureAgentWorkspace } from "../agents/workspace.js"; +import { getRegistryWorktree } from "../agents/worktrees/registry.js"; +import { managedWorktrees } from "../agents/worktrees/service.js"; +import { upsertSqliteSessionEntry } from "../config/sessions/session-accessor.sqlite.js"; +import { clearSessionStoreCacheForTest } from "../config/sessions/store.js"; +import type { OpenClawConfig } from "../config/types.openclaw.js"; +import { closeOpenClawStateDatabaseForTest } from "../state/openclaw-state-db.js"; +import { testing as agentCommandTesting } from "./agent.js"; +import { createThrowingTestRuntime } from "./test-runtime-config-helpers.js"; + +const configIoMocks = vi.hoisted(() => ({ + loadConfig: vi.fn(), + readConfigFileSnapshotForWrite: vi.fn(), +})); +const pluginRegistryMocks = vi.hoisted(() => ({ + ensurePluginRegistryLoaded: vi.fn(), +})); + +vi.mock("../config/io.js", () => ({ + getRuntimeConfig: configIoMocks.loadConfig, + loadConfig: configIoMocks.loadConfig, + readConfigFileSnapshotForWrite: configIoMocks.readConfigFileSnapshotForWrite, +})); + +vi.mock("../plugins/runtime/runtime-registry-loader.js", () => ({ + ensurePluginRegistryLoaded: pluginRegistryMocks.ensurePluginRegistryLoaded, +})); + +const execFileAsync = promisify(execFile); +const runtime = createThrowingTestRuntime(); +const sessionKey = "agent:main:worktree-race"; + +function recordProof(line: string): void { + const out = process.env.OPENCLAW_PROOF_OUT; + if (out) { + fsSync.appendFileSync(out, `${line}\n`); + } +} + +async function git(cwd: string, ...args: string[]): Promise { + await execFileAsync("git", ["-C", cwd, ...args], { encoding: "utf8" }); +} + +async function initializeRepository(root: string): Promise { + const repo = path.join(root, "repo"); + await fs.mkdir(repo, { recursive: true }); + await git(repo, "init", "-b", "main"); + await git(repo, "config", "user.name", "OpenClaw Test"); + await git(repo, "config", "user.email", "openclaw-test@example.invalid"); + await fs.writeFile(path.join(repo, "README.md"), "base\n"); + await git(repo, "add", "README.md"); + await git(repo, "commit", "-m", "initial"); + return await fs.realpath(repo); +} + +function mockConfig(home: string, storePath: string): OpenClawConfig { + const cfg = { + agents: { + defaults: { + model: { primary: "anthropic/claude-opus-4-6" }, + models: { "anthropic/claude-opus-4-6": {} }, + workspace: path.join(home, "openclaw"), + }, + }, + session: { store: storePath, mainKey: "main" }, + } as OpenClawConfig; + configIoMocks.loadConfig.mockReturnValue(cfg); + return cfg; +} + +async function seedSession( + storePath: string, + spawnedCwd: string, + worktree?: { id: string; branch: string; repoRoot: string }, +): Promise { + await upsertSqliteSessionEntry( + { agentId: "main", sessionKey, storePath }, + { + sessionId: "session-worktree-race", + updatedAt: Date.now(), + spawnedCwd, + ...(worktree ? { worktree } : {}), + }, + ); +} + +async function createSessionWorktree( + home: string, +): Promise<{ id: string; path: string; branch: string; repoRoot: string }> { + const repo = await initializeRepository(home); + const created = await managedWorktrees.create({ + repoRoot: repo, + name: "race-session", + ownerKind: "session", + ownerId: sessionKey, + }); + return { id: created.id, path: created.path, branch: created.branch, repoRoot: created.repoRoot }; +} + +describe("agent command worktree admission", () => { + beforeEach(() => { + vi.mocked(ensureAgentWorkspace).mockClear(); + vi.mocked(ensureAgentWorkspace).mockResolvedValue({ dir: "" }); + clearSessionStoreCacheForTest(); + }); + + afterEach(() => { + closeOpenClawStateDatabaseForTest(); + }); + + it("holds the lease through workspace preparation so a racing removal is rejected", async () => { + await withTempHome(async (home) => { + const storePath = path.join(home, "sessions.json"); + mockConfig(home, storePath); + const created = await createSessionWorktree(home); + const nested = path.join(created.path, "workspace"); + await fs.mkdir(nested); + await seedSession(storePath, nested); + + let releasePause = () => {}; + const paused = new Promise((resolve) => { + releasePause = resolve; + }); + let reachPause = () => {}; + const pauseReached = new Promise((resolve) => { + reachPause = resolve; + }); + vi.mocked(ensureAgentWorkspace).mockImplementationOnce(async (params) => { + reachPause(); + await paused; + return { dir: params?.dir ?? "" }; + }); + + const preparing = agentCommandTesting.prepareAgentCommandExecution( + { message: "resume in worktree", sessionKey }, + runtime, + ); + await pauseReached; + + let removalDuringPreparation: string; + try { + const removed = await managedWorktrees.remove({ id: created.id, reason: "idle-gc" }); + removalDuringPreparation = `removed=${removed.removed}`; + } catch (error) { + removalDuringPreparation = `rejected: ${(error as Error).message}`; + } + const checkoutSurvivedPreparation = fsSync.existsSync(nested); + recordProof(`removal during production workspace preparation: ${removalDuringPreparation}`); + recordProof(`managed checkout survived preparation: ${checkoutSurvivedPreparation}`); + + releasePause(); + const prepared = await preparing; + await prepared.runLease?.release(); + + expect(removalDuringPreparation).toContain("worktree is busy"); + expect(checkoutSurvivedPreparation).toBe(true); + expect((await managedWorktrees.remove({ id: created.id, reason: "idle-gc" })).removed).toBe( + true, + ); + }); + }); + + it("fails closed before workspace setup when the session's bound worktree was removed", async () => { + await withTempHome(async (home) => { + const storePath = path.join(home, "sessions.json"); + mockConfig(home, storePath); + const created = await createSessionWorktree(home); + const nested = path.join(created.path, "workspace"); + await fs.mkdir(nested); + await seedSession(storePath, nested, { + id: created.id, + branch: created.branch, + repoRoot: created.repoRoot, + }); + await managedWorktrees.remove({ id: created.id, reason: "manual-delete", force: true }); + expect(getRegistryWorktree(process.env, created.id)?.removedAt).toBeDefined(); + + let preparationResult: string; + try { + await agentCommandTesting.prepareAgentCommandExecution( + { message: "resume in worktree", sessionKey }, + runtime, + ); + preparationResult = "preparation proceeded without its checkout"; + } catch (error) { + preparationResult = `preparation fails: ${(error as Error).message}`; + } + const workspaceSetupRan = vi.mocked(ensureAgentWorkspace).mock.calls.length > 0; + recordProof(`admission for a removed authoritative binding: ${preparationResult}`); + recordProof(`workspace setup ran on the removed worktree: ${workspaceSetupRan}`); + + expect(preparationResult).toContain("managed worktree was removed"); + expect(workspaceSetupRan).toBe(false); + }); + }); +});