Files
pi-file-claiming/index.ts
2026-06-19 12:46:02 -04:00

679 lines
19 KiB
TypeScript

/**
* index.ts — File Claiming Extension entry point.
*
* This extension provides a file-claiming mechanism that tracks active
* claims (read / write / exclusive) on file paths so tools and extensions
* can detect conflicts before operating on files.
*
* ## Quick start
*
* ```ts
* import type { FileClaim, ClaimOwner } from "./src/lock-types";
* ```
*
* ## Events
*
* The extension emits the following events on the shared extension bus
* (`pi.events`):
*
* | Event | Payload | Description |
* |--------------------|---------------|---------------------------------|
* | `claim:acquired` | `ClaimEvent` | A claim was successfully acquired. |
* | `claim:released` | `ClaimEvent` | A claim was released. |
* | `claim:conflicted` | `ClaimEvent` | A claim could not be acquired. |
* | `claim:expired` | `ClaimEvent` | A claim was auto-expired. |
*
* ## Configuration
*
* Options are managed through the `config.ts` module and can be
* persisted to `{lockDir}/config.json`. Use `/file-claiming-config`
* to inspect or update values at runtime.
*
* @module file-claiming
*/
import type {
ExtensionAPI,
ExtensionCommandContext,
ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import type {
ClaimConflict,
ClaimEvent,
ClaimEventType,
ClaimOwner,
ClaimRegistry,
ClaimResult,
FileClaim,
LockEntry,
PathLockType,
} from "./src/lock-types";
import {
type FileClaimingConfig,
getConfig,
setConfig,
loadConfigFromFile,
saveConfigToFile,
getConfigFilePath,
createDefaultConfig,
} from "./src/config";
import {
injectLockClaimingIntoPrompt,
buildLockClaimingGuidelines,
buildLockClaimingToolSnippets,
} from "./src/system-prompt";
import {
buildDiagnosticCollection,
formatDiagnostics,
} from "./src/diagnostics";
import { registerLockTools } from "./src/tools";
import {
createLockNotificationHandler,
claimEventToNotification,
} from "./src/notifications";
import { updateLockStatus, persistLockState } from "./src/user-interaction";
import {
createLockCommandHandler,
createNotifyCommandHandler,
} from "./src/user-interaction";
import {
acquireLock,
autoClaim,
isFileLocked,
getLockInfo,
isMutationTool,
shouldAutoClaim,
handleToolLock,
buildBlockingError,
checkToolBlocking,
cleanupExpiredLocks,
type AcquireLockOptions,
type LockAcquisitionResult,
} from "./src/lock-acquisition";
import { registerEventHandlers } from "./src/event-handlers";
// ---------------------------------------------------------------------------
// In-memory registry implementation
// ---------------------------------------------------------------------------
// We use plain objects as maps because ES Map does not serialize well and
// we may persist registry snapshots via pi.appendEntry later.
/** Claims indexed by ID. */
let claimsById: Record<string, FileClaim> = {};
/** Lock entries indexed by path. */
let locksByPath: Record<string, LockEntry[]> = {};
/** Reference to the interval handle for the expiry sweeper. */
let sweepTimer: ReturnType<typeof setInterval> | undefined;
/** Shared event bus (set during init). */
let emitEvent:
| ((type: ClaimEventType, claim: FileClaim, conflict?: ClaimConflict) => void)
| undefined;
// ---------------------------------------------------------------------------
// Conflict detection logic
// ---------------------------------------------------------------------------
/**
* Returns whether a proposed lock type is compatible with the set of locks
* already held on a path, excluding locks owned by `owner`.
*/
function isCompatible(
existing: LockEntry[],
requested: PathLockType,
owner: ClaimOwner,
): { ok: boolean; blockers: LockEntry[] } {
const blockers: LockEntry[] = [];
for (const entry of existing) {
// Skip locks held by the same owner — they already hold it.
if (entry.owner.type === owner.type && entry.owner.id === owner.id) {
continue;
}
// Exclusive blocks everything.
if (entry.lockType === "exclusive" || requested === "exclusive") {
blockers.push(entry);
continue;
}
// Write blocks write and is blocked by write.
if (entry.lockType === "write" || requested === "write") {
blockers.push(entry);
}
// Read is always compatible with other reads.
}
return { ok: blockers.length === 0, blockers };
}
// ---------------------------------------------------------------------------
// Registry API
// ---------------------------------------------------------------------------
const registry: ClaimRegistry = {
get claims(): Record<string, FileClaim> {
return claimsById;
},
get locks(): Record<string, LockEntry[]> {
return locksByPath;
},
getActiveClaims(path: string): FileClaim[] {
return Object.values(claimsById).filter(
(c) => c.path === path && c.status === "active",
);
},
getLocks(path: string): LockEntry[] {
return locksByPath[path] ?? [];
},
checkConflict(
path: string,
lockType: PathLockType,
owner: ClaimOwner,
): ClaimConflict | undefined {
const existing = locksByPath[path] ?? [];
const { ok, blockers } = isCompatible(existing, lockType, owner);
if (ok) return undefined;
const blockingClaims = blockers
.map((b) => claimsById[b.claimId])
.filter(Boolean);
return {
path,
severity: lockType === "read" ? "info" : "warning",
blockedClaim: {
id: "",
path,
lockType,
status: "conflicted",
owner,
createdAt: new Date().toISOString(),
updatedAt: new Date().toISOString(),
},
blockingClaims,
message: `Cannot acquire "${lockType}" lock on "${path}": blocked by ${blockers.length} existing lock(s)`,
};
},
acquire(claim: FileClaim): ClaimResult {
const existing = locksByPath[claim.path] ?? [];
const { ok, blockers } = isCompatible(
existing,
claim.lockType,
claim.owner,
);
if (!ok) {
const blockingClaims = blockers
.map((b) => claimsById[b.claimId])
.filter(Boolean);
const conflict: ClaimConflict = {
path: claim.path,
severity: claim.lockType === "read" ? "info" : "warning",
blockedClaim: claim,
blockingClaims,
message: `Cannot acquire "${claim.lockType}" lock on "${claim.path}": blocked by ${blockers.length} existing lock(s)`,
};
emitEvent?.("claim:conflicted", claim, conflict);
return { success: false, claim, conflict };
}
// Store the claim
const now = new Date().toISOString();
const stored: FileClaim = {
...claim,
status: "active",
createdAt: claim.createdAt || now,
updatedAt: now,
};
claimsById[stored.id] = stored;
// Add lock entry
const entry: LockEntry = {
path: stored.path,
lockType: stored.lockType,
claimId: stored.id,
owner: stored.owner,
acquiredAt: now,
};
locksByPath[stored.path] = [...(locksByPath[stored.path] ?? []), entry];
emitEvent?.("claim:acquired", stored);
return { success: true, claim: stored };
},
release(claimId: string): boolean {
const claim = claimsById[claimId];
if (!claim) return false;
claim.status = "released";
claim.updatedAt = new Date().toISOString();
// Remove lock entries for this claim
const pathEntries = locksByPath[claim.path] ?? [];
locksByPath[claim.path] = pathEntries.filter((e) => e.claimId !== claimId);
if ((locksByPath[claim.path] ?? []).length === 0) {
delete locksByPath[claim.path];
}
emitEvent?.("claim:released", claim);
return true;
},
releaseAllByOwner(owner: ClaimOwner): void {
const toRelease = Object.values(claimsById).filter(
(c) =>
c.owner.type === owner.type &&
c.owner.id === owner.id &&
c.status === "active",
);
for (const claim of toRelease) {
registry.release(claim.id);
}
},
};
// ---------------------------------------------------------------------------
// Expiry sweeper
// ---------------------------------------------------------------------------
function sweepExpiredClaims(): void {
const now = Date.now();
const toExpire = Object.values(claimsById).filter((c) => {
if (c.status !== "active") return false;
if (!c.expiresAt) return false;
return new Date(c.expiresAt).getTime() <= now;
});
for (const claim of toExpire) {
claim.status = "expired";
claim.updatedAt = new Date().toISOString();
// Remove lock entries
const pathEntries = locksByPath[claim.path] ?? [];
locksByPath[claim.path] = pathEntries.filter((e) => e.claimId !== claim.id);
if ((locksByPath[claim.path] ?? []).length === 0) {
delete locksByPath[claim.path];
}
emitEvent?.("claim:expired", claim);
}
}
// ---------------------------------------------------------------------------
// Config-driven helpers
// ---------------------------------------------------------------------------
/**
* Build the current agent-owner context from a turn event or tool event context.
* Used by `releaseOnTurnEnd` and `blockedTools` logic.
*/
function ownerFromContext(ctx: ExtensionContext): ClaimOwner | undefined {
// The agent is the primary owner during tool execution.
return {
type: "agent",
id: "main",
sessionId: ctx.sessionManager.getSessionFile() ?? undefined,
};
}
/**
* Release all agent-held claims when `releaseOnTurnEnd` is enabled.
*/
function releaseAgentClaimsOnTurnEnd(ctx: ExtensionContext): void {
const config = getConfig();
if (!config.releaseOnTurnEnd) return;
const owner = ownerFromContext(ctx);
if (owner) {
registry.releaseAllByOwner(owner);
}
}
/**
* Check whether a given tool name should be blocked based on `blockedTools`
* and the current lock state.
*/
function isToolBlocked(toolName: string): boolean {
const config = getConfig();
if (config.blockedTools.length === 0) return false;
if (!config.blockedTools.includes(toolName)) return false;
// Check if there are any active claims
return Object.values(claimsById).some((c) => c.status === "active");
}
// ---------------------------------------------------------------------------
// Lock acquisition exports
// ---------------------------------------------------------------------------
/**
* Re-export lock acquisition functions from the lock-acquisition module
* for use by other extensions and tools.
*/
export {
acquireLock,
autoClaim,
isFileLocked,
getLockInfo,
isMutationTool,
shouldAutoClaim,
handleToolLock,
buildBlockingError,
checkToolBlocking,
cleanupExpiredLocks,
type AcquireLockOptions,
type LockAcquisitionResult,
} from "./src/lock-acquisition";
// ---------------------------------------------------------------------------
// Reset for testability
// ---------------------------------------------------------------------------
/**
* Reset all in-memory state. Primarily intended for testing.
*/
export function resetRegistry(): void {
claimsById = {};
locksByPath = {};
}
/**
* Returns a reference to the in-memory {@link ClaimRegistry} for direct
* use by other extensions or tools.
*/
export function getClaimRegistry(): ClaimRegistry {
return registry;
}
// ---------------------------------------------------------------------------
// Extension factory
// ---------------------------------------------------------------------------
/**
* Pi extension factory for the File Claiming extension.
*
* Handles:
* - Configuration loading and command
* - Sweeper interval
* - Turn-end claim release
* - Tool blocking when claims are active
* - Diagnostics widget
*
* @param pi - The {@link ExtensionAPI} instance provided by Pi.
*/
export default function (pi: ExtensionAPI): void {
// -----------------------------------------------------------------------
// Event emitter helper
// -----------------------------------------------------------------------
emitEvent = (
type: ClaimEventType,
claim: FileClaim,
conflict?: ClaimConflict,
) => {
const event: ClaimEvent = {
type,
claim,
conflict,
timestamp: new Date().toISOString(),
};
pi.events.emit(type, event);
};
// -----------------------------------------------------------------------
// Session lifecycle
// -----------------------------------------------------------------------
pi.on("session_start", async (_event: unknown, ctx: ExtensionContext) => {
// 1. Load persisted configuration
await loadConfigFromFile();
const config = getConfig();
// 2. Start the expiry sweeper (idempotent)
if (!sweepTimer) {
sweepTimer = setInterval(sweepExpiredClaims, 5_000);
}
// 3. Set up notification handler
const notificationHandler = createLockNotificationHandler(
ctx.ui,
pi.events,
);
// 4. Show footer status if enabled
if (config.showDiagnostics && ctx.hasUI) {
ctx.ui.setWidget("file-claiming-diagnostics", undefined);
ctx.ui.setStatus("file-claiming", "Claims: 0 active");
}
// 5. Persist initial lock state
persistLockState(pi);
// 6. Register dedicated event handlers from event-handlers module
registerEventHandlers(pi, ctx);
});
pi.on("session_shutdown", async () => {
if (sweepTimer) {
clearInterval(sweepTimer);
sweepTimer = undefined;
}
// Release all claims
registry.releaseAllByOwner({ type: "agent", id: "main" });
cleanupExpiredLocks();
});
// -----------------------------------------------------------------------
// Turn-end: release agent claims
// -----------------------------------------------------------------------
pi.on("turn_end", async (_event: unknown, ctx: ExtensionContext) => {
releaseAgentClaimsOnTurnEnd(ctx);
updateDiagnosticsWidget(ctx);
updateLockStatus(ctx.ui, registry);
});
function updateDiagnosticsWidget(ctx: ExtensionContext): void {
const config = getConfig();
if (!config.showDiagnostics || !ctx.hasUI) return;
const activeCount = Object.values(registry.claims).filter(
(c) => c.status === "active",
).length;
ctx.ui.setStatus("file-claiming", `Claims: ${activeCount} active`);
}
// -----------------------------------------------------------------------
// /file-claiming-config command
// -----------------------------------------------------------------------
pi.registerCommand("file-claiming-config", {
description:
"View or update File Claiming extension configuration. " +
"Usage: /file-claiming-config [key=value ...] — omitting arguments shows current config. " +
"Keys: autoReleaseTTL, releaseOnTurnEnd, lockDir, blockedTools, showDiagnostics.",
handler: async (args: string, ctx: ExtensionCommandContext) => {
const trimmed = args.trim();
// No arguments → display current config
if (!trimmed) {
const cfg = getConfig();
const lines = [
"── File Claiming Configuration ──",
` autoReleaseTTL: ${cfg.autoReleaseTTL} ms ${cfg.autoReleaseTTL === createDefaultConfig().autoReleaseTTL ? "(default)" : ""}`,
` releaseOnTurnEnd: ${cfg.releaseOnTurnEnd} ${cfg.releaseOnTurnEnd === createDefaultConfig().releaseOnTurnEnd ? "(default)" : ""}`,
` lockDir: ${cfg.lockDir}`,
` blockedTools: [${cfg.blockedTools.join(", ")}] ${jsonEqual(cfg.blockedTools, createDefaultConfig().blockedTools) ? "(default)" : ""}`,
` showDiagnostics: ${cfg.showDiagnostics} ${cfg.showDiagnostics === createDefaultConfig().showDiagnostics ? "(default)" : ""}`,
` config file: ${getConfigFilePath(cfg.lockDir)}`,
"",
` Active claims: ${Object.values(claimsById).filter((c) => c.status === "active").length}`,
"──────────────────────────────",
"",
" Set values: /file-claiming-config key=value [key=value ...]",
" Examples:",
" /file-claiming-config autoReleaseTTL=600000",
" /file-claiming-config releaseOnTurnEnd=false",
' /file-claiming-config blockedTools=["edit","write","bash"]',
" /file-claiming-config showDiagnostics=false",
];
ctx.ui.notify(lines.join("\n"), "info");
return;
}
// Parse key=value pairs
const pairs = trimmed.split(/\s+/);
const updates: Record<string, unknown> = {};
for (const pair of pairs) {
const eqIdx = pair.indexOf("=");
if (eqIdx === -1) {
ctx.ui.notify(`Invalid syntax: "${pair}" — use key=value`, "error");
return;
}
const key = pair.slice(0, eqIdx);
const rawValue = pair.slice(eqIdx + 1);
// Coerce value to the expected type
updates[key] = coerceConfigValue(key, rawValue);
}
const result = setConfig(updates as Partial<FileClaimingConfig>);
if (!result.valid) {
ctx.ui.notify(
`Validation errors:\n${result.errors.join("\n")}`,
"error",
);
return;
}
// Persist to disk
try {
await saveConfigToFile();
ctx.ui.notify("Configuration updated and saved.", "info");
} catch (err) {
ctx.ui.notify(
`Configuration updated in memory but failed to persist: ${err}`,
"warning",
);
}
// Restart sweeper with new config if autoReleaseTTL changed
updateDiagnosticsWidget(ctx);
},
});
// -----------------------------------------------------------------------
// Register /file-claiming-locks command for interactive lock management
// -----------------------------------------------------------------------
const lockCommandHandler = createLockCommandHandler(() => registry);
pi.registerCommand("file-claiming-locks", {
description:
"Interactive lock management. Use /file-claiming-locks for interactive mode, " +
"or /file-claiming-locks [list|check|release|reset] [path] [id] for direct commands.",
handler: lockCommandHandler,
});
// -----------------------------------------------------------------------
// Context event: inject diagnostic messages
// -----------------------------------------------------------------------
pi.on(
"context",
async (
event: {
messages: import("@earendil-works/pi-agent-core").AgentMessage[];
},
ctx: ExtensionContext,
) => {
const config = getConfig();
if (!config.showDiagnostics) return {};
const activeClaims = Object.values(claimsById).filter(
(c) => c.status === "active",
);
if (activeClaims.length === 0) return {};
const collection = buildDiagnosticCollection(registry);
const diagnosticLines: string[] = [
"## File Claiming Lock Status",
"",
`Active claims: ${collection.count}`,
"",
];
for (const [uri, items] of collection.diagnostics) {
for (const item of items) {
diagnosticLines.push(
`- **${item.lockType} lock** on \`${item.uri}\` by ${item.tool ?? item.source}${item.message}`,
);
}
}
const diagnosticMessage: import("@earendil-works/pi-agent-core").AgentMessage =
{
role: "user",
content: [{ type: "text", text: diagnosticLines.join("\n") }],
id: `file-claiming-diagnostics-${Date.now()}`,
timestamp: Date.now(),
};
return { messages: [...event.messages, diagnosticMessage] };
},
);
}
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
/**
* Deep-compare two values for equality (simple JSON comparison).
*/
function jsonEqual(a: unknown, b: unknown): boolean {
return JSON.stringify(a) === JSON.stringify(b);
}
/**
* Coerce a command-line string value to the expected type for a given config key.
*/
function coerceConfigValue(key: string, raw: string): unknown {
// Booleans
if (raw === "true") return true;
if (raw === "false") return false;
// Numbers
if (/^-?\d+(\.\d+)?$/.test(raw)) {
const num = Number(raw);
if (Number.isFinite(num)) return num;
}
// Arrays — parse as JSON
if (raw.startsWith("[") && raw.endsWith("]")) {
try {
return JSON.parse(raw);
} catch {
// Fall through to return raw string
}
}
// Everything else is a string
return raw;
}