diff --git a/api/server/experimental.js b/api/server/experimental.js index ee6cf8c8b77..194b01ebb71 100644 --- a/api/server/experimental.js +++ b/api/server/experimental.js @@ -54,6 +54,7 @@ const { GenerationJobManager, createAgentEventTerminalHandler, startCodeEnvironmentLifecycleReconciler, + createClusteredPRAutomationReaper, waitForKeyvRedisClient, createCodeApiUploadRegistry, cacheConfig, @@ -80,6 +81,7 @@ const { seedDatabase, recordAuditEntry, sweepOrphanedPreviews, + reapPRAutomations, } = require('~/models'); const { checkMigrations } = require('./services/start/migration'); const initializeMCPs = require('./services/initializeMCPs'); @@ -407,6 +409,9 @@ if (cluster.isMaster) { // Tear down stream resources before shared caches and telemetry exporters shut down. registerShutdownTask('generation job manager', destroyGenerationJobManager, { priority: 100 }); const expiredFileSweep = createClusteredFileSweep(cacheConfig.USE_REDIS, startExpiredFileSweep); + const prAutomationReaper = createClusteredPRAutomationReaper(cacheConfig.USE_REDIS, { + reapPRAutomations, + }); const SCHEDULE_ENGINE_OPTIONAL_METHODS = new Set(['GET', 'HEAD', 'OPTIONS', 'DELETE']); const rejectScheduleWritesUntilReady = (req, res, next) => { @@ -438,6 +443,7 @@ if (cluster.isMaster) { if (msg.type === 'file-retention-sweep-worker') { logger.info(wrapLogMessage(`Worker ${process.pid} is assigned file-retention sweep`)); expiredFileSweep.assign(); + prAutomationReaper.assign(); } }); const startServer = async () => { @@ -505,6 +511,7 @@ if (cluster.isMaster) { // principal) still merges DB `__base__` overrides, which must not drive which hook // modules load in every worker (matches api/server/index.js's baseOnly usage). const baseAppConfig = await getAppConfig({ baseOnly: true }); + prAutomationReaper.start(baseAppConfig); await configureSubagentTaskRouting(baseAppConfig?.endpoints?.agents?.subagentActivity); registerBackgroundTaskShutdown({ interruptGraceMs: baseAppConfig?.endpoints?.agents?.backgroundTasks?.shutdownInterruptGraceMs, diff --git a/api/server/index.js b/api/server/index.js index 535e7a46c50..6307c561c8f 100644 --- a/api/server/index.js +++ b/api/server/index.js @@ -62,6 +62,7 @@ const { createAgentEventTerminalHandler, createScheduleWriteGate, startCodeEnvironmentLifecycleReconciler, + startPRAutomationReaper, waitForKeyvRedisClient, warnOnUnreachableDeliveryPaths, createCodeApiUploadRegistry, @@ -73,6 +74,7 @@ const { getRoleByName, seedDatabase, recordAuditEntry, + reapPRAutomations, } = require('~/models'); const initializeOAuthReconnectManager = require('./services/initializeOAuthReconnectManager'); const { capabilityContextMiddleware } = require('./middleware/roles/capabilities'); @@ -236,6 +238,7 @@ const startServer = async () => { logger.error('[sweepOrphanedPreviews] Background sweep failed:', err); }); const appConfig = await getAppConfig({ baseOnly: true }); + startPRAutomationReaper({ reapPRAutomations }, appConfig); await configureSubagentTaskRouting(appConfig?.endpoints?.agents?.subagentActivity); registerBackgroundTaskShutdown({ interruptGraceMs: appConfig?.endpoints?.agents?.backgroundTasks?.shutdownInterruptGraceMs, diff --git a/packages/api/src/agents/automation.spec.ts b/packages/api/src/agents/automation.spec.ts new file mode 100644 index 00000000000..6c5eb9b8239 --- /dev/null +++ b/packages/api/src/agents/automation.spec.ts @@ -0,0 +1,91 @@ +import { AppService } from '@librechat/data-schemas'; +import { + startPRAutomationReaper, + reapOrphanedPRAutomations, + createClusteredPRAutomationReaper, +} from './automation'; +import { isLeader } from '~/cluster'; + +jest.mock('~/cluster', () => ({ isLeader: jest.fn() })); + +const emptyPage = () => jest.fn().mockResolvedValue({ scanned: 0, removed: 0 }); + +describe('PR automation reaper', () => { + let stop: (() => void) | undefined; + + afterEach(() => { + stop?.(); + stop = undefined; + jest.mocked(isLeader).mockReset(); + jest.clearAllTimers(); + jest.useRealTimers(); + }); + + test('follows the cursor until every page was read', async () => { + const reapPRAutomations = jest + .fn() + .mockResolvedValueOnce({ scanned: 100, removed: 2, next: 'a' }) + .mockResolvedValueOnce({ scanned: 3, removed: 1 }); + + expect(await reapOrphanedPRAutomations({ reapPRAutomations })).toBe(3); + expect(reapPRAutomations.mock.calls).toEqual([[{ after: undefined }], [{ after: 'a' }]]); + }); + + test('reaps only on the elected replica, then again every hour', async () => { + jest.useFakeTimers(); + const leader = jest.mocked(isLeader); + leader.mockResolvedValueOnce(false).mockResolvedValueOnce(true); + const reapPRAutomations = emptyPage(); + + stop = startPRAutomationReaper({ reapPRAutomations }); + await jest.advanceTimersByTimeAsync(0); + expect(reapPRAutomations).not.toHaveBeenCalled(); + + await jest.advanceTimersByTimeAsync(60 * 60_000); + expect(leader).toHaveBeenCalledTimes(2); + expect(reapPRAutomations).toHaveBeenCalledTimes(1); + }); + + test('runs at the interval the yaml sets', async () => { + jest.useFakeTimers(); + jest.mocked(isLeader).mockResolvedValue(true); + const reapPRAutomations = emptyPage(); + const appConfig = await AppService({ + config: { endpoints: { agents: { prAutomation: { reaperIntervalMinutes: 15 } } } }, + }); + + stop = startPRAutomationReaper({ reapPRAutomations }, appConfig); + await jest.advanceTimersByTimeAsync(0); + await jest.advanceTimersByTimeAsync(15 * 60_000); + expect(reapPRAutomations).toHaveBeenCalledTimes(2); + }); + + test('without Redis, waits in a worker until the primary assigns it', async () => { + jest.useFakeTimers(); + jest.mocked(isLeader).mockResolvedValue(true); + const reapPRAutomations = emptyPage(); + const reaper = createClusteredPRAutomationReaper(false, { reapPRAutomations }); + + stop = reaper.stop; + reaper.start(); + await jest.advanceTimersByTimeAsync(0); + expect(reapPRAutomations).not.toHaveBeenCalled(); + + reaper.assign(); + stop = reaper.stop; + await jest.advanceTimersByTimeAsync(0); + expect(reapPRAutomations).toHaveBeenCalledTimes(1); + }); + + test('with Redis, every worker starts and leader election picks one', async () => { + jest.useFakeTimers(); + jest.mocked(isLeader).mockResolvedValue(true); + const reapPRAutomations = emptyPage(); + const reaper = createClusteredPRAutomationReaper(true, { reapPRAutomations }); + + reaper.start(); + stop = reaper.stop; + await jest.advanceTimersByTimeAsync(0); + expect(reapPRAutomations).toHaveBeenCalledTimes(1); + }); +}); diff --git a/packages/api/src/agents/automation.ts b/packages/api/src/agents/automation.ts new file mode 100644 index 00000000000..e42b0c56e70 --- /dev/null +++ b/packages/api/src/agents/automation.ts @@ -0,0 +1,90 @@ +import { logger, runAsSystem } from '@librechat/data-schemas'; +import { DEFAULT_PR_AUTOMATION_REAPER_MINUTES } from 'librechat-data-provider'; +import type { AppConfig, PRAutomationMethods } from '@librechat/data-schemas'; +import { isLeader } from '~/cluster'; + +let reapTimer: NodeJS.Timeout | undefined; +let reapInFlight: Promise | undefined; + +type ReaperMethods = Pick; + +/** Reads every PR automation record once and removes those whose conversation or owner is gone. */ +export async function reapOrphanedPRAutomations({ + reapPRAutomations: reapPage, +}: ReaperMethods): Promise { + let removed = 0; + let after: string | undefined; + do { + const page = await reapPage({ after }); + removed += page.removed; + after = page.next; + } while (after != null); + return removed; +} + +/** + * Conversations also leave through the retention index, which runs no application code, so the + * elected replica periodically removes the PR automation records they leave behind. Returns a + * function that stops the schedule. + */ +export function startPRAutomationReaper(methods: ReaperMethods, appConfig?: AppConfig): () => void { + if (reapTimer != null) return stopPRAutomationReaper; + const run = (): void => { + if (reapInFlight != null) return; + const current = isLeader() + .then(async (leader) => { + if (!leader) return; + await runAsSystem(() => reapOrphanedPRAutomations(methods)); + }) + .catch((error) => { + logger.error('[prAutomation] reaping orphaned records failed:', error); + }) + .finally(() => { + if (reapInFlight === current) reapInFlight = undefined; + }); + reapInFlight = current; + }; + run(); + const minutes = + appConfig?.endpoints?.agents?.prAutomation?.reaperIntervalMinutes ?? + DEFAULT_PR_AUTOMATION_REAPER_MINUTES; + reapTimer = setInterval(run, minutes * 60_000); + reapTimer.unref(); + return stopPRAutomationReaper; +} + +function stopPRAutomationReaper(): void { + clearInterval(reapTimer); + reapTimer = undefined; +} + +/** + * Starts the reaper in one worker of a clustered server. Leader election needs Redis; without + * it every worker counts as the leader, so the worker the primary assigns the file-retention + * sweep to runs the reaper too. `start` passes the loaded config, `assign` marks the worker. + */ +export function createClusteredPRAutomationReaper( + distributed: boolean, + methods: ReaperMethods, +): { start: (appConfig?: AppConfig) => void; assign: () => void; stop: () => void } { + let eligible = distributed; + let ready = false; + let config: AppConfig | undefined; + let stop: (() => void) | undefined; + const maybeStart = () => { + if (!eligible || !ready || stop != null) return; + stop = startPRAutomationReaper(methods, config); + }; + return { + start: (appConfig) => { + config = appConfig; + ready = true; + maybeStart(); + }, + assign: () => { + eligible = true; + maybeStart(); + }, + stop: () => stop?.(), + }; +} diff --git a/packages/api/src/agents/index.ts b/packages/api/src/agents/index.ts index b8c34f34abd..0d4bd369414 100644 --- a/packages/api/src/agents/index.ts +++ b/packages/api/src/agents/index.ts @@ -1,5 +1,6 @@ export * from './avatars'; export * from './attachments'; +export * from './automation'; export * from './chain'; export * from './callerCapabilities'; export * from './client'; diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index 9b8f8fcb8e3..aed97fc810b 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -69,6 +69,14 @@ import { } from './types/pullRequest'; import type { PullRequestProvider } from './types/pullRequest'; import { isActionTool } from './types/tools'; +import { + MAX_PR_AUTOMATION_BOTS, + PR_AUTOMATION_TRUST_LEVELS, + DEFAULT_PR_AUTOMATION_BOTS, + MIN_PR_AUTOMATION_REAPER_MINUTES, + MAX_PR_AUTOMATION_REAPER_MINUTES, + DEFAULT_PR_AUTOMATION_REAPER_MINUTES, +} from './types/prAutomation'; import { apiBaseUrl } from './api-endpoints'; import { FileSources } from './types/files'; import { MCPServersSchema } from './mcp'; @@ -2014,6 +2022,38 @@ export const agentsEndpointSchema = baseEndpointSchema sharedCacheSlots: z.number().int().min(64).max(1_000_000).optional().default(65_536), }) .optional(), + /** Auto-fix CI and address review comments on a conversation's pull request. Off unless an + * administrator opts in; a user can only narrow `maxTrust`, never widen it. */ + prAutomation: z + .object({ + enabled: z.boolean().optional().default(false), + /** Fix rounds one pull request may run before the automation stops. */ + maxRounds: z.number().int().min(1).max(20).optional().default(5), + /** Wall-clock minutes from the first automated round before the automation stops. */ + maxMinutes: z.number().int().min(5).max(1440).optional().default(120), + /** Bots one user may approve for one repository. */ + maxBots: z + .number() + .int() + .min(1) + .max(MAX_PR_AUTOMATION_BOTS) + .optional() + .default(DEFAULT_PR_AUTOMATION_BOTS), + /** Widest author trust a user may select: `approvedBots` acts on the linked account and + * allowlisted bots only, `collaborators` adds write-access users, `anyone` acts on all. */ + maxTrust: z.enum(PR_AUTOMATION_TRUST_LEVELS).optional().default('approvedBots'), + /** Minutes between passes that remove records whose conversation or owner is gone, + * including conversations the retention index deleted. Runs whether or not `enabled` + * is set, so records from an earlier opt-in are still cleaned up. */ + reaperIntervalMinutes: z + .number() + .int() + .min(MIN_PR_AUTOMATION_REAPER_MINUTES) + .max(MAX_PR_AUTOMATION_REAPER_MINUTES) + .optional() + .default(DEFAULT_PR_AUTOMATION_REAPER_MINUTES), + }) + .optional(), /** Conversational background-task delivery policy. Automatic completion wakeups are * enabled unless an administrator explicitly restores poll-only behavior. */ backgroundTasks: z diff --git a/packages/data-provider/src/index.ts b/packages/data-provider/src/index.ts index 1c9dc09f08e..af2d9de1bd9 100644 --- a/packages/data-provider/src/index.ts +++ b/packages/data-provider/src/index.ts @@ -66,6 +66,7 @@ export type { export * from './types/subagents'; export * from './types/background'; export * from './types/pullRequest'; +export * from './types/prAutomation'; export * from './types/queuedTurns'; /* access permissions */ export * from './accessPermissions'; diff --git a/packages/data-provider/src/prAutomation.spec.ts b/packages/data-provider/src/prAutomation.spec.ts new file mode 100644 index 00000000000..5999f6c0255 --- /dev/null +++ b/packages/data-provider/src/prAutomation.spec.ts @@ -0,0 +1,89 @@ +import type { PRAutomationTrustLevel } from './types/prAutomation'; +import { clampPRAutomationTrust } from './types/prAutomation'; +import { configSchema } from './config'; + +describe('clampPRAutomationTrust', () => { + it('defaults to the narrowest level when the user chose nothing', () => { + expect(clampPRAutomationTrust(undefined, 'anyone')).toBe('approvedBots'); + }); + + it.each<[PRAutomationTrustLevel, PRAutomationTrustLevel, PRAutomationTrustLevel]>([ + ['anyone', 'approvedBots', 'approvedBots'], + ['collaborators', 'approvedBots', 'approvedBots'], + ['anyone', 'collaborators', 'collaborators'], + ['approvedBots', 'anyone', 'approvedBots'], + ['collaborators', 'anyone', 'collaborators'], + ['anyone', 'anyone', 'anyone'], + ])('requested %s under a %s ceiling gives %s', (requested, ceiling, expected) => { + expect(clampPRAutomationTrust(requested, ceiling)).toBe(expected); + }); + + it('falls back to the narrowest level for a value it does not know', () => { + expect(clampPRAutomationTrust('everyone' as unknown as PRAutomationTrustLevel, 'anyone')).toBe( + 'approvedBots', + ); + }); +}); + +describe('agent PR automation config', () => { + const parse = (prAutomation?: unknown) => + configSchema.safeParse({ version: '1.0', endpoints: { agents: { prAutomation } } }); + + it('is absent when not configured, which leaves it off', () => { + const absent = configSchema.parse({ version: '1.0', endpoints: { agents: {} } }); + expect(absent.endpoints?.agents?.prAutomation).toBeUndefined(); + }); + + it('defaults to disabled with the narrowest ceiling', () => { + const empty = configSchema.parse({ + version: '1.0', + endpoints: { agents: { prAutomation: {} } }, + }); + expect(empty.endpoints?.agents?.prAutomation).toEqual({ + enabled: false, + maxRounds: 5, + maxMinutes: 120, + maxBots: 20, + maxTrust: 'approvedBots', + reaperIntervalMinutes: 60, + }); + }); + + it.each([4, 1441, 1.5])('rejects a %s minute reaper interval', (reaperIntervalMinutes) => { + expect(parse({ reaperIntervalMinutes }).success).toBe(false); + }); + + it('accepts a shorter reaper interval', () => { + expect(parse({ reaperIntervalMinutes: 15 }).success).toBe(true); + }); + + it('accepts an administrator raising the limits and the ceiling', () => { + const result = parse({ + enabled: true, + maxRounds: 10, + maxMinutes: 240, + maxTrust: 'collaborators', + }); + expect(result.success).toBe(true); + }); + + it.each([0, 21, 1.5])('rejects %s fix rounds', (maxRounds) => { + expect(parse({ maxRounds }).success).toBe(false); + }); + + it.each([4, 1441, 1.5])('rejects a %s minute window', (maxMinutes) => { + expect(parse({ maxMinutes }).success).toBe(false); + }); + + it.each([0, 101, 1.5])('rejects %s approved bots per repository', (maxBots) => { + expect(parse({ maxBots }).success).toBe(false); + }); + + it('accepts the largest approved bot limit', () => { + expect(parse({ maxBots: 100 }).success).toBe(true); + }); + + it.each(['everyone', 'bots', ''])('rejects %s as a trust ceiling', (maxTrust) => { + expect(parse({ maxTrust }).success).toBe(false); + }); +}); diff --git a/packages/data-provider/src/types/prAutomation.ts b/packages/data-provider/src/types/prAutomation.ts new file mode 100644 index 00000000000..967a8c78808 --- /dev/null +++ b/packages/data-provider/src/types/prAutomation.ts @@ -0,0 +1,50 @@ +export const PR_AUTOMATION_STATES = ['idle', 'waiting', 'fixing', 'needs_user', 'stopped'] as const; + +/** + * Stable, machine-readable reasons an automation stopped. The client maps each + * code to localized copy; none of them carries provider text or a comment body. + */ +export const PR_AUTOMATION_STOP_CODES = [ + 'user_stopped', + 'round_cap', + 'time_cap', + 'fork_pull_request', + 'untrusted_head_commit', + 'approval_expired', + 'pull_request_closed', + 'disabled_by_admin', +] as const; + +/** Largest per-repository bot allowlist a deployment can configure. */ +export const MAX_PR_AUTOMATION_BOTS = 100; +export const DEFAULT_PR_AUTOMATION_BOTS = 20; +/** Minutes between passes that remove records whose conversation or owner is gone. */ +export const DEFAULT_PR_AUTOMATION_REAPER_MINUTES = 60; +export const MIN_PR_AUTOMATION_REAPER_MINUTES = 5; +export const MAX_PR_AUTOMATION_REAPER_MINUTES = 1440; + +/** Ordered from narrowest to widest: who may be acted on. */ +export const PR_AUTOMATION_TRUST_LEVELS = ['approvedBots', 'collaborators', 'anyone'] as const; + +export type PRAutomationState = (typeof PR_AUTOMATION_STATES)[number]; +export type PRAutomationStopCode = (typeof PR_AUTOMATION_STOP_CODES)[number]; +export type PRAutomationTrustLevel = (typeof PR_AUTOMATION_TRUST_LEVELS)[number]; + +/** + * An administrator's ceiling bounds what a user may choose. A user can only + * narrow it, so the result is whichever level is narrower. + */ +export function clampPRAutomationTrust( + requested: PRAutomationTrustLevel | undefined, + ceiling: PRAutomationTrustLevel, +): PRAutomationTrustLevel { + if (requested == null) { + return PR_AUTOMATION_TRUST_LEVELS[0]; + } + const requestedRank = PR_AUTOMATION_TRUST_LEVELS.indexOf(requested); + const ceilingRank = PR_AUTOMATION_TRUST_LEVELS.indexOf(ceiling); + if (requestedRank < 0) { + return PR_AUTOMATION_TRUST_LEVELS[0]; + } + return PR_AUTOMATION_TRUST_LEVELS[Math.min(requestedRank, ceilingRank)]; +} diff --git a/packages/data-schemas/src/app/service.spec.ts b/packages/data-schemas/src/app/service.spec.ts index 92d59a5aa76..3863b274529 100644 --- a/packages/data-schemas/src/app/service.spec.ts +++ b/packages/data-schemas/src/app/service.spec.ts @@ -220,6 +220,18 @@ describe('loadFiltersConfig', () => { }); }); +describe('AppService PR automation config', () => { + it('carries the yaml reaper interval onto the app config', async () => { + const config = { + endpoints: { [EModelEndpoint.agents]: { prAutomation: { reaperIntervalMinutes: 15 } } }, + } as DeepPartial; + + const result = await AppService({ config }); + + expect(result.endpoints?.[EModelEndpoint.agents]?.prAutomation?.reaperIntervalMinutes).toBe(15); + }); +}); + describe('AppService assistants config', () => { it('preserves configured Assistants API versions', async () => { const config = { diff --git a/packages/data-schemas/src/methods/conversation.spec.ts b/packages/data-schemas/src/methods/conversation.spec.ts index ecea2ae3596..04175ea32ac 100644 --- a/packages/data-schemas/src/methods/conversation.spec.ts +++ b/packages/data-schemas/src/methods/conversation.spec.ts @@ -19,6 +19,7 @@ import type { } from '../types'; import { ConversationMethods, createConversationMethods } from './conversation'; import { tenantStorage, runAsSystem } from '~/config/tenantContext'; +import { createPRAutomationMethods } from './prAutomation'; import { createChatProjectMethods } from './chatProject'; import { createModels } from '../models'; @@ -3727,6 +3728,154 @@ describe('Conversation Operations', () => { }); describe('deleteConvos', () => { + it('removes the PR automation record of a deleted conversation and keeps the others', async () => { + const conversationId = uuidv4(); + const keptConversationId = uuidv4(); + const PRAutomation = mongoose.models.PRAutomation; + for (const id of [conversationId, keptConversationId]) { + await Conversation.create({ conversationId: id, user: 'user123', endpoint: 'agents' }); + await PRAutomation.create({ user: 'user123', conversationId: id }); + } + + await deleteConvos('user123', { conversationId }); + + expect(await PRAutomation.countDocuments({ conversationId })).toBe(0); + expect(await PRAutomation.countDocuments({ conversationId: keptConversationId })).toBe(1); + }); + + it('keeps a record enabled for the conversation saved again before the cleanup', async () => { + const conversationId = uuidv4(); + const PRAutomation = mongoose.models.PRAutomation; + const deleted = await Conversation.create({ + conversationId, + user: 'user123', + endpoint: 'agents', + }); + await PRAutomation.create({ + user: 'user123', + conversationId, + conversationRef: String(deleted._id), + }); + const original = Conversation.deleteMany.bind(Conversation); + const spy = jest.spyOn(Conversation, 'deleteMany').mockImplementationOnce((( + ...args: Parameters + ) => + original(...args).then(async (result) => { + const renewed = await Conversation.create({ + conversationId, + user: 'user123', + endpoint: 'agents', + }); + await PRAutomation.updateOne( + { user: 'user123', conversationId }, + { $set: { conversationRef: String(renewed._id) } }, + ); + return result; + })) as unknown as typeof Conversation.deleteMany); + try { + await deleteConvos('user123', { conversationId }); + } finally { + spy.mockRestore(); + } + expect(await PRAutomation.countDocuments({ user: 'user123', conversationId })).toBe(1); + }); + + describe('when the conversation delete fails', () => { + const user = 'user123'; + const setup = async () => { + const conversationId = uuidv4(); + await Conversation.create({ conversationId, user, endpoint: EModelEndpoint.agents }); + const automation = createPRAutomationMethods(mongoose); + await mongoose.models.PRAutomation.create({ + user, + conversationId, + repository: 'acme/one', + pullNumber: 1, + }); + jest.spyOn(Conversation, 'deleteMany').mockRejectedValueOnce(new Error('delete failed')); + await expect(deleteConvos(user, { conversationId })).rejects.toThrow('delete failed'); + return { conversationId, automation }; + }; + + it('keeps the PR automation record', async () => { + const { conversationId, automation } = await setup(); + expect(await automation.getPRAutomation({ userId: user, conversationId })).not.toBeNull(); + expect(await Conversation.exists({ conversationId })).not.toBeNull(); + }); + + it('leaves the record exactly as it was, so the kept conversation keeps working', async () => { + const { conversationId, automation } = await setup(); + expect(await automation.getPRAutomation({ userId: user, conversationId })).toMatchObject({ + state: 'idle', + repository: 'acme/one', + pullNumber: 1, + }); + }); + }); + + it('writes nothing for a conversation that never had an automation', async () => { + const conversationId = uuidv4(); + await Conversation.create({ conversationId, user: 'user123', endpoint: 'agents' }); + const PRAutomation = mongoose.models.PRAutomation; + const writes = ['updateOne', 'updateMany', 'bulkWrite', 'insertMany', 'create'].map( + (method) => jest.spyOn(PRAutomation, method as 'updateOne'), + ); + try { + await deleteConvos('user123', { conversationId }); + for (const write of writes) { + expect(write).not.toHaveBeenCalled(); + } + expect(await PRAutomation.countDocuments({ user: 'user123', conversationId })).toBe(0); + } finally { + writes.forEach((write) => write.mockRestore()); + } + }); + + it('still reports a deletion that committed when the record cleanup keeps failing', async () => { + const conversationId = uuidv4(); + await Conversation.create({ conversationId, user: 'user123', endpoint: 'agents' }); + const PRAutomation = mongoose.models.PRAutomation; + await PRAutomation.create({ user: 'user123', conversationId }); + const cleanup = jest + .spyOn(PRAutomation, 'deleteMany') + .mockRejectedValue(new Error('cleanup unavailable')); + try { + await expect(deleteConvos('user123', { conversationId })).resolves.toMatchObject({ + deletedCount: 1, + }); + expect(cleanup.mock.calls.length).toBeGreaterThan(1); + } finally { + cleanup.mockRestore(); + } + }); + + it('removes the PR automation record of a root that a previous attempt already deleted', async () => { + /** The reaper judges the owner too, so the record names a real account id. */ + const owner = new mongoose.Types.ObjectId().toString(); + const rootId = uuidv4(); + const childId = uuidv4(); + const PRAutomation = mongoose.models.PRAutomation; + await Conversation.create({ + conversationId: childId, + user: owner, + endpoint: EModelEndpoint.agents, + subagentThread: { + rootConversationId: rootId, + parentConversationId: rootId, + parentMessageId: 'message-1', + parentToolCallId: 'call-1', + subagentType: 'agent-child', + subagentKind: 'agent', + depth: 1, + }, + }); + await PRAutomation.create({ user: owner, conversationId: rootId }); + + await deleteConvos(owner, { conversationId: rootId }, { allowEmpty: true }); + + expect(await PRAutomation.countDocuments({ conversationId: rootId })).toBe(0); + }); + it('retires queued-turn work before each conversation deletion wave', async () => { const conversationId = uuidv4(); await Conversation.create({ diff --git a/packages/data-schemas/src/methods/conversation.ts b/packages/data-schemas/src/methods/conversation.ts index 350178cafe2..809a483af04 100644 --- a/packages/data-schemas/src/methods/conversation.ts +++ b/packages/data-schemas/src/methods/conversation.ts @@ -53,6 +53,7 @@ import { isAgentFadingTier, isAgentFadingTierEntries } from '~/utils/fading'; import { isCompactionSemanticIndexProjection } from '~/types/compaction'; import { withoutMeiliIndexing } from '~/models/plugins/mongoMeili'; import { tenantSafeBulkWrite } from '~/utils/tenantBulkWrite'; +import { createPRAutomationMethods } from './prAutomation'; import { isValidObjectIdString } from '~/utils/objectId'; import { decrementTagCounts } from './conversationTag'; import logger from '~/config/winston'; @@ -668,6 +669,7 @@ export function createConversationMethods( deps?: ConversationMethodDeps, ): ConversationMethods { let legacyReceiptExpiryCursor: Types.ObjectId | undefined; + const prAutomation = createPRAutomationMethods(mongoose); /** * Stamps a real assistant reply with a strictly increasing server value. @@ -4025,7 +4027,7 @@ export function createConversationMethods( const userFilter = { ...filter, user }; type DeletionConversation = Pick< IConversation, - 'conversationId' | 'tenantId' | 'chatProjectId' | 'tags' + '_id' | 'conversationId' | 'tenantId' | 'chatProjectId' | 'tags' >; const retryCascadeOperation = async (operation: () => PromiseLike | T): Promise => { let lastError: unknown; @@ -4141,6 +4143,15 @@ export function createConversationMethods( await options?.beforeDelete?.(waveIds); await deps?.prepareAgentTriggerConversationResultErasure?.(user, waveIds); const result = await Conversation.deleteMany({ user, conversationId: { $in: waveIds } }); + /** Only after the delete committed, so a failed delete keeps the user's automation. + * A claim refuses any record whose conversation is gone and removes it, so a cleanup + * that still fails after its retry converges on the next claim. */ + const waveRefs = wave.map((conversation) => String(conversation._id)); + await retryCascadeOperation(() => + prAutomation.deletePRAutomations(user, waveIds, waveRefs), + ).catch((error) => + logger.warn('[deleteConvos] PR automation cleanup deferred to its next claim', error), + ); if (result.deletedCount > 0) { /** Result erasure is irreversible. Keep receipts intact when a * pre-delete hook or the conversation delete itself fails, so a @@ -4174,6 +4185,12 @@ export function createConversationMethods( ]; if (recoveryConversationIds.length > 0) { + /** The deleted documents are unknown here, so only records whose conversation is gone go. */ + await retryCascadeOperation(() => + prAutomation.reapPRAutomationsOf(user, recoveryConversationIds), + ).catch((error) => + logger.warn('[deleteConvos] PR automation cleanup deferred to its next claim', error), + ); await deps?.deleteAgentQueuedTurns?.( user, recoveryConversationIds.map((conversationId) => ({ diff --git a/packages/data-schemas/src/methods/index.ts b/packages/data-schemas/src/methods/index.ts index 8be120e86dc..b30a5fd9d0d 100644 --- a/packages/data-schemas/src/methods/index.ts +++ b/packages/data-schemas/src/methods/index.ts @@ -43,6 +43,8 @@ import { MAX_TOOL_FAVORITES, type ToolFavoriteMethods, } from './favorite'; +/* PR Automation */ +import { createPRAutomationMethods, type PRAutomationMethods } from './prAutomation'; /* Agent Categories */ import { createAgentCategoryMethods, type AgentCategoryMethods } from './agentCategory'; /* Agent API Keys */ @@ -264,6 +266,7 @@ export type AllMethods = ToolApprovalGrantStorage & FileMethods & MemoryMethods & ToolFavoriteMethods & + PRAutomationMethods & AgentCategoryMethods & AgentApiKeyMethods & MCPServerMethods & @@ -504,6 +507,7 @@ export function createMethods( ...createFileMethods(mongoose), ...createMemoryMethods(mongoose), ...createToolFavoriteMethods(mongoose), + ...createPRAutomationMethods(mongoose), ...createAgentCategoryMethods(mongoose), ...createAgentApiKeyMethods(mongoose), ...createMCPServerMethods(mongoose, { @@ -578,6 +582,7 @@ export type { ProjectFileRecord, ProjectFilesOptions, ToolFavoriteMethods, + PRAutomationMethods, AgentCategoryMethods, AgentApiKeyMethods, MCPServerMethods, diff --git a/packages/data-schemas/src/methods/prAutomation.spec.ts b/packages/data-schemas/src/methods/prAutomation.spec.ts new file mode 100644 index 00000000000..a8bc9445918 --- /dev/null +++ b/packages/data-schemas/src/methods/prAutomation.spec.ts @@ -0,0 +1,1944 @@ +import mongoose from 'mongoose'; +import { logger, createModels } from '..'; +import { MongoMemoryServer } from 'mongodb-memory-server'; +import type { PRAutomationStopCode } from 'librechat-data-provider'; +import type { PRAutomationEventStopCode } from '~/types/prAutomation'; +import { runAsSystem, tenantStorage } from '~/config/tenantContext'; +import { createPRAutomationMethods } from './prAutomation'; + +logger.silent = true; + +let PRAutomation: mongoose.Model; +let methods: ReturnType; +let mongoServer: MongoMemoryServer; + +const userId = new mongoose.Types.ObjectId().toString(); +const otherUserId = new mongoose.Types.ObjectId().toString(); +const key = { userId, conversationId: 'convo-1' }; +const limits = { maxRounds: 3, maxMinutes: 60 }; +const maxBots = 3; +const head = (n: number) => (n === 1 ? 'a' : String(n)).repeat(40).slice(0, 40); +const pullOne = { repository: 'acme/one', pullNumber: 1 }; +const pullTwo = { repository: 'acme/one', pullNumber: 2 }; +const otherRepository = { repository: 'acme/two', pullNumber: 2 }; + +const mismatch = { ok: false, error: { code: 'binding_mismatch' } }; +/** Enables and returns the record, for setups that expect the enable to succeed. */ +type EnableParams = Parameters[0]; +/** The widest ceiling, so setups that pass a trust level store it as given. */ +const enableOpen = ( + params: Omit & { maxTrust?: EnableParams['maxTrust'] }, +) => methods.enablePRAutomation({ maxTrust: 'anyone', ...params }); +const enableWith = async (params: Parameters[0]) => { + const result = await enableOpen(params); + if (!result.ok) { + throw new Error(`enable failed: ${result.error.code}`); + } + return result.value; +}; +const enable = (binding = pullOne) => enableWith({ ...key, binding }); +/** The current run of a record, as a webhook handler would read it before claiming. */ +const currentEpoch = async (target = key) => + (await methods.getPRAutomation(target))?.epoch ?? 'none'; +const claim = async ( + n: number, + { binding = pullOne, ...extra }: { now?: Date; binding?: typeof pullOne; epoch?: string } = {}, +) => + methods.claimPRAutomationRound({ + ...key, + ...limits, + headSha: head(n), + ...extra, + binding: { ...binding, epoch: extra.epoch ?? (await currentEpoch()) }, + }); +/** Claims a round that must succeed and returns the record it started. */ +const startRound = async (n: number) => { + const result = await claim(n); + if (!result.ok) { + throw new Error(`claim ${n} failed: ${result.error.code}`); + } + return result.value; +}; +const settle = (round: number, runId: string, state: 'waiting' | 'needs_user' = 'waiting') => + methods.settlePRAutomationRound({ ...key, round, runId, state }); +/** Test setup that does not depend on the settle method under test. */ +const toWaiting = () => + PRAutomation.updateOne( + { user: userId, conversationId: key.conversationId }, + { $set: { state: 'waiting' } }, + ); +const addBot = async (id: number, login?: string, limit = maxBots, repository = 'acme/one') => + methods.addPRAutomationBot(key, login == null ? { id } : { id, login }, limit, { + repository, + epoch: (await methods.getPRAutomation(key))?.epoch ?? 'none', + }); +const keyFilterOf = ({ userId: user, conversationId }: typeof key) => ({ user, conversationId }); +const seedConversation = (fields: Record = {}) => + mongoose.models.Conversation.create({ + conversationId: key.conversationId, + user: userId, + endpoint: 'agents', + ...fields, + }); + +/** + * Runs `during` right after the next `findOneAndUpdate` of `model` resolves, which for the + * default is the claim's atomic write, before the claim checks again. + */ +function interleave( + during: () => Promise, + model: mongoose.Model = PRAutomation, + method: 'findOneAndUpdate' | 'findById' | 'findOne' = 'findOneAndUpdate', +) { + const original = (model[method] as typeof PRAutomation.findOneAndUpdate).bind(model); + return jest.spyOn(model, method as 'findOneAndUpdate').mockImplementationOnce((( + ...args: Parameters + ) => { + const query = original(...args); + const lean = query.lean.bind(query); + query.lean = ((...leanArgs: Parameters) => { + const result = lean(...leanArgs); + return { + ...result, + then: (resolve: (value: unknown) => unknown, reject: (error: unknown) => unknown) => + result + .then(async (value: unknown) => { + await during(); + return value; + }) + .then(resolve, reject), + }; + }) as typeof query.lean; + return query; + }) as typeof PRAutomation.findOneAndUpdate); +} +/** + * Runs `during` right before the `nth` owner read. An enable reads the owner once before its + * write and once after it, so the second read of an enable is where its rollback is decided. + */ +function beforeOwnerRead(during: () => Promise, nth = 2) { + const User = mongoose.models.User; + const original = User.findById.bind(User); + let reads = 0; + return jest.spyOn(User, 'findById').mockImplementation((( + ...args: Parameters + ) => { + const query = original(...args); + if (++reads !== nth) { + return query; + } + const lean = query.lean.bind(query); + query.lean = ((...leanArgs: Parameters) => ({ + then: (resolve: (value: unknown) => unknown, reject: (error: unknown) => unknown) => + during() + .then(() => lean(...leanArgs)) + .then(resolve, reject), + })) as unknown as typeof query.lean; + return query; + }) as typeof User.findById); +} +/** Runs `during` right before the next `replaceOne` of the record, which is an enable's rollback. */ +function beforeRollback(during: () => Promise) { + const original = PRAutomation.replaceOne.bind(PRAutomation); + return jest.spyOn(PRAutomation, 'replaceOne').mockImplementationOnce((( + ...args: Parameters + ) => ({ + then: (resolve: (value: unknown) => unknown, reject: (error: unknown) => unknown) => + during() + .then(() => original(...args)) + .then(resolve, reject), + })) as unknown as typeof PRAutomation.replaceOne); +} +/** What an event handler captured about the run it acts on. */ +const fenceFor = async (binding = pullOne) => ({ + ...binding, + epoch: (await methods.getPRAutomation(key))?.epoch as string, +}); +const startAccountDeletionNow = () => + mongoose.models.User.updateOne( + { _id: userId }, + { $set: { agentTriggerDeletionStartedAt: new Date() } }, + ); +const cancelAccountDeletion = () => + mongoose.models.User.updateOne({ _id: userId }, { $unset: { agentTriggerDeletionStartedAt: 1 } }); + +beforeAll(async () => { + mongoServer = await MongoMemoryServer.create(); + await mongoose.connect(mongoServer.getUri()); + createModels(mongoose); + PRAutomation = mongoose.models.PRAutomation; + await PRAutomation.syncIndexes(); + methods = createPRAutomationMethods(mongoose); +}); + +afterAll(async () => { + await mongoose.disconnect(); + await mongoServer.stop(); +}); + +beforeEach(async () => { + await PRAutomation.deleteMany({}); + await mongoose.models.Conversation.deleteMany({}); + await mongoose.models.User.deleteMany({}); + await mongoose.models.User.create({ _id: userId, email: 'owner@example.com', provider: 'local' }); + await seedConversation(); +}); + +describe('getPRAutomation', () => { + test('returns null for a conversation that never enabled it', async () => { + expect(await methods.getPRAutomation(key)).toBeNull(); + }); +}); + +describe('enablePRAutomation', () => { + test('creates an idle record on the narrowest trust level', async () => { + const record = await enable(); + expect(record).toMatchObject({ + conversationId: 'convo-1', + state: 'idle', + round: 0, + trust: 'approvedBots', + trustedBots: [], + }); + }); + + test('is idempotent and keeps one record per user and conversation', async () => { + await enable(); + await enable(); + expect(await PRAutomation.countDocuments({ user: userId })).toBe(1); + }); + + test('does not change an active record that is enabled again', async () => { + await enable(); + await claim(1); + const again = await enable(); + expect(again).toMatchObject({ state: 'fixing', round: 1 }); + }); + + test('restarts a stopped record with a fresh round counter and window', async () => { + await enable(); + await claim(1); + await methods.stopPRAutomation(key, 'user_stopped'); + + const restarted = await enable(); + expect(restarted).toMatchObject({ state: 'idle', round: 0 }); + expect(restarted.stopCode).toBeUndefined(); + expect(restarted.startedAt).toBeUndefined(); + expect(restarted.lastHeadSha).toBeUndefined(); + }); + + test('lets a restarted record claim a head it claimed before the stop', async () => { + await enable(); + await claim(1); + await methods.stopPRAutomation(key, 'user_stopped'); + await enable(); + expect((await claim(1)).ok).toBe(true); + }); + + test('creates a record no wider than the administrator ceiling', async () => { + const record = await enableWith({ + ...key, + binding: pullOne, + trust: 'anyone', + maxTrust: 'collaborators', + }); + expect(record.trust).toBe('collaborators'); + }); + + test('rebinds a record no wider than the administrator ceiling', async () => { + await enable(pullOne); + const rebound = await enableWith({ + ...key, + binding: otherRepository, + trust: 'anyone', + maxTrust: 'approvedBots', + }); + expect(rebound.trust).toBe('approvedBots'); + }); + + test('applies a narrower trust level passed with a rebind of a running record', async () => { + await enableWith({ ...key, binding: pullOne, trust: 'anyone' }); + const rebound = await enableWith({ ...key, binding: otherRepository, trust: 'approvedBots' }); + expect(rebound.trust).toBe('approvedBots'); + }); + + test('applies a trust level passed with a rebind to another pull request', async () => { + await enableWith({ ...key, binding: pullOne, trust: 'anyone' }); + const rebound = await enableWith({ ...key, binding: pullTwo, trust: 'collaborators' }); + expect(rebound.trust).toBe('collaborators'); + }); + + test('enables for a stored conversation whose id is longer than 256 characters', async () => { + const conversationId = 'c'.repeat(300); + await seedConversation({ conversationId }); + const result = await methods.enablePRAutomation({ userId, conversationId, maxTrust: 'anyone' }); + expect(result.ok).toBe(true); + }); + + test('keeps records of different users apart', async () => { + await enable(); + expect(await methods.getPRAutomation({ ...key, userId: otherUserId })).toBeNull(); + }); + + test('clears approved bots when the conversation is bound to a different repository', async () => { + await enableOpen({ ...key, binding: pullOne }); + await addBot(101); + const rebound = await enableWith({ ...key, binding: otherRepository }); + expect(rebound).toMatchObject({ ...otherRepository, trustedBots: [] }); + }); + + test('refuses a bot approved before any repository was bound', async () => { + await enableOpen(key); + expect(await addBot(101)).toEqual({ ok: false, error: { code: 'binding_mismatch' } }); + expect((await methods.getPRAutomation(key))?.trustedBots).toEqual([]); + }); + + test('keeps approved bots when the same repository is bound to another pull request', async () => { + await enableOpen({ ...key, binding: pullOne }); + await addBot(101); + const rebound = await enableWith({ ...key, binding: pullTwo }); + expect(rebound).toMatchObject({ pullNumber: 2, trustedBots: [{ id: 101 }] }); + }); +}); + +describe('claiming while the owner is being deleted', () => { + const startAccountDeletion = () => + mongoose.models.User.updateOne( + { _id: userId }, + { $set: { agentTriggerDeletionStartedAt: new Date() } }, + ); + + test('refuses a claim once account deletion has started', async () => { + await enable(pullOne); + await startAccountDeletion(); + expect(await claim(1)).toEqual({ ok: false, error: { code: 'owner_inactive' } }); + }); + + test('keeps the record and spends no round, so a cancelled deletion leaves it as it was', async () => { + await enable(pullOne); + await startAccountDeletion(); + await claim(1); + expect(await methods.getPRAutomation(key)).toMatchObject({ state: 'idle', round: 0 }); + }); + + test('undoes a rejected claim even when a trust change landed after it', async () => { + await enable(pullOne); + const spy = interleave(async () => { + await methods.setPRAutomationTrust(key, 'collaborators', 'anyone'); + await startAccountDeletion(); + }); + try { + expect(await claim(1)).toEqual({ ok: false, error: { code: 'owner_inactive' } }); + } finally { + spy.mockRestore(); + } + await cancelAccountDeletion(); + expect(await methods.getPRAutomation(key)).toMatchObject({ + state: 'idle', + round: 0, + claimedHeads: [], + trust: 'collaborators', + }); + expect((await claim(1)).ok).toBe(true); + }); + + test('hands the round back when the check after the claim fails', async () => { + await enable(pullOne); + const Conversation = mongoose.models.Conversation; + const original = Conversation.aggregate.bind(Conversation); + let calls = 0; + const spy = jest.spyOn(Conversation, 'aggregate').mockImplementation((( + ...args: Parameters + ) => { + if (++calls === 2) { + throw new Error('synthetic read failure'); + } + return original(...args); + }) as typeof Conversation.aggregate); + try { + await expect(claim(1)).rejects.toThrow('synthetic read failure'); + } finally { + spy.mockRestore(); + } + expect(await methods.getPRAutomation(key)).toMatchObject({ state: 'idle', round: 0 }); + expect((await claim(1)).ok).toBe(true); + }); + + test('refuses a claim for an owner that no longer exists', async () => { + await enable(pullOne); + await mongoose.models.User.deleteMany({}); + expect(await claim(1)).toEqual({ ok: false, error: { code: 'owner_inactive' } }); + }); + + test('removes the record of an owner that no longer exists', async () => { + await enable(pullOne); + await mongoose.models.User.deleteMany({}); + await claim(1); + expect(await methods.getPRAutomation(key)).toBeNull(); + }); + + test('claims again once the deletion was cancelled', async () => { + await enable(pullOne); + await startAccountDeletion(); + await claim(1); + await mongoose.models.User.updateOne( + { _id: userId }, + { $unset: { agentTriggerDeletionStartedAt: 1 } }, + ); + expect((await claim(1)).ok).toBe(true); + }); +}); + +describe('a deletion that lands while a claim is in flight', () => { + test('does not hand out a round for a conversation deleted after the first check', async () => { + await enable(pullOne); + const spy = interleave(() => mongoose.models.Conversation.deleteMany({})); + try { + expect(await claim(1)).toEqual({ ok: false, error: { code: 'conversation_gone' } }); + } finally { + spy.mockRestore(); + } + }); + + test('does not hand out a round once account deletion started after the first check', async () => { + await enable(pullOne); + const spy = interleave(() => + mongoose.models.User.updateOne( + { _id: userId }, + { $set: { agentTriggerDeletionStartedAt: new Date() } }, + ), + ); + try { + expect(await claim(1)).toEqual({ ok: false, error: { code: 'owner_inactive' } }); + } finally { + spy.mockRestore(); + } + }); + + test('leaves no round running when the owner deletion started mid-claim', async () => { + await enable(pullOne); + const spy = interleave(() => + mongoose.models.User.updateOne( + { _id: userId }, + { $set: { agentTriggerDeletionStartedAt: new Date() } }, + ), + ); + try { + await claim(1); + } finally { + spy.mockRestore(); + } + expect((await methods.getPRAutomation(key))?.state).not.toBe('fixing'); + }); + + /** Runs one claim whose owner deletion starts between its write and its second check. */ + const claimDuringDeletion = async (n: number) => { + const spy = interleave(startAccountDeletionNow); + try { + return await claim(n); + } finally { + spy.mockRestore(); + } + }; + + test('lets the same head claim again once a deletion that started mid-claim is cancelled', async () => { + await enable(pullOne); + await claimDuringDeletion(1); + await cancelAccountDeletion(); + expect((await claim(1)).ok).toBe(true); + }); + + test('does not spend a round on a claim it handed back', async () => { + await enable(pullOne); + await claimDuringDeletion(1); + expect((await methods.getPRAutomation(key))?.round).toBe(0); + }); + + test('does not open the time window for a claim it handed back', async () => { + await enable(pullOne); + await claimDuringDeletion(1); + expect((await methods.getPRAutomation(key))?.startedAt).toBeUndefined(); + }); + + test('hands a later claim back to exactly where the previous round left the record', async () => { + await enable(pullOne); + await startRound(1); + await toWaiting(); + await claimDuringDeletion(2); + expect(await methods.getPRAutomation(key)).toMatchObject({ + state: 'waiting', + round: 1, + lastHeadSha: head(1), + claimedHeads: [head(1)], + }); + }); +}); + +describe('a conversation deleted and saved again under the same id', () => { + const recreate = async () => { + await mongoose.models.Conversation.deleteMany({}); + await seedConversation(); + }; + + test('refuses a claim for the record of the earlier conversation', async () => { + await enable(pullOne); + await recreate(); + expect(await claim(1)).toEqual({ ok: false, error: { code: 'conversation_gone' } }); + expect(await methods.getPRAutomation(key)).toBeNull(); + }); + + test('keeps the new record when a claim for the old one refuses after a re-enable', async () => { + await enable(pullOne); + const before = await methods.getPRAutomation(key); + const spy = interleave( + async () => { + await recreate(); + await PRAutomation.deleteMany({}); + await enable(pullOne); + }, + PRAutomation, + 'findOne', + ); + try { + expect(await claim(1, { epoch: before?.epoch })).toEqual({ + ok: false, + error: { code: 'not_found' }, + }); + } finally { + spy.mockRestore(); + } + const renewed = await methods.getPRAutomation(key); + expect(renewed?.version).toBe(before?.version); + expect(renewed?.epoch).not.toBe(before?.epoch); + }); + + test('reaps only the listed records whose conversation is gone', async () => { + await enable(pullOne); + const orphan = { userId, conversationId: 'orphan-1' }; + await PRAutomation.create({ user: userId, conversationId: orphan.conversationId }); + expect( + await methods.reapPRAutomationsOf(userId, [key.conversationId, orphan.conversationId]), + ).toBe(1); + expect(await methods.getPRAutomation(key)).not.toBeNull(); + expect(await methods.getPRAutomation(orphan)).toBeNull(); + }); + + test('starts a new run for the new conversation when enabled', async () => { + await enableWith({ ...key, binding: pullOne, trust: 'anyone' }); + await addBot(7); + await recreate(); + const record = await enableWith({ ...key, binding: otherRepository }); + const conversation = await mongoose.models.Conversation.findOne({ + conversationId: key.conversationId, + }).lean<{ _id: unknown }>(); + expect(record).toMatchObject({ + ...otherRepository, + trust: 'approvedBots', + trustedBots: [], + round: 0, + conversationRef: String(conversation?._id), + }); + expect((await claim(1, { binding: otherRepository })).ok).toBe(true); + }); +}); + +describe('enabling under a lowered trust ceiling', () => { + test('restarts a stopped record at the current ceiling when no trust is requested', async () => { + await enableWith({ ...key, binding: pullOne, trust: 'anyone' }); + await methods.stopPRAutomation(key, 'user_stopped'); + await enableWith({ ...key, binding: pullOne, maxTrust: 'approvedBots' }); + expect(await methods.getPRAutomation(key)).toMatchObject({ trust: 'approvedBots' }); + }); + + test('narrows a running record when it is enabled again', async () => { + await enableWith({ ...key, binding: pullOne, trust: 'anyone' }); + await enableWith({ ...key, binding: pullOne, maxTrust: 'collaborators' }); + expect(await methods.getPRAutomation(key)).toMatchObject({ trust: 'collaborators' }); + }); +}); + +describe('enabling for a conversation or owner that no longer exists', () => { + test('does not leave a record behind for a deleted conversation', async () => { + await mongoose.models.Conversation.deleteMany({}); + await enableOpen({ ...key, binding: pullOne }); + expect(await methods.getPRAutomation(key)).toBeNull(); + }); + + test('reports that the conversation is gone', async () => { + await mongoose.models.Conversation.deleteMany({}); + expect(await enableOpen({ ...key, binding: pullOne })).toEqual({ + ok: false, + error: { code: 'conversation_gone' }, + }); + }); + + /** Runs `during` right after the enable's write, before it checks again. */ + const interleaveEnable = (during: () => Promise) => beforeOwnerRead(during); + + test('rejects an enable once account deletion started mid-write', async () => { + const spy = interleaveEnable(startAccountDeletionNow); + try { + expect(await enableOpen({ ...key, binding: pullOne })).toEqual({ + ok: false, + error: { code: 'owner_inactive' }, + }); + } finally { + spy.mockRestore(); + } + }); + + test('leaves nothing claimable after a cancelled deletion that an enable crossed', async () => { + const spy = interleaveEnable(startAccountDeletionNow); + try { + await enableOpen({ ...key, binding: pullOne }); + } finally { + spy.mockRestore(); + } + await cancelAccountDeletion(); + expect(await claim(1)).toEqual({ ok: false, error: { code: 'not_found' } }); + }); + + test('does not roll back an enable that landed after the crossing enable was rejected', async () => { + await enable(pullOne); + const spy = interleaveEnable(startAccountDeletionNow); + /** After the rejected enable reads the owner, the deletion is cancelled and a newer + * enable rebinds the record before the rejected one undoes its write. */ + const rollback = beforeRollback(async () => { + await cancelAccountDeletion(); + await enable(pullTwo); + }); + try { + await enableOpen({ ...key, binding: otherRepository }); + } finally { + spy.mockRestore(); + rollback.mockRestore(); + } + expect((await methods.getPRAutomation(key))?.pullNumber).toBe(2); + }); + + test('never brings back an older binding when an enable crosses a deletion', async () => { + await enable(pullOne); + const spy = interleave(async () => { + await enable(otherRepository); + await startAccountDeletionNow(); + }); + try { + await enableOpen({ ...key, binding: pullTwo }); + } finally { + spy.mockRestore(); + } + await cancelAccountDeletion(); + expect(await methods.getPRAutomation(key)).not.toMatchObject(pullOne); + }); + + test('leaves nothing claimable when two enables cross the same deletion', async () => { + await enable(pullOne); + /** + * Owner reads, in order: A before, A after, B before, B after. A writes, B writes over it, + * the deletion starts, A sees it and rolls back first, then B sees it and rolls back. + */ + const User = mongoose.models.User; + const original = User.findById.bind(User); + let reads = 0; + let reachedB = () => {}; + const atB = new Promise((resolve) => (reachedB = resolve)); + let releaseB = () => {}; + const gateB = new Promise((resolve) => (releaseB = resolve)); + let enableB: Promise = Promise.resolve(); + const before: Record Promise> = { + 2: async () => { + enableB = enableOpen({ ...key, binding: otherRepository }); + await atB; + await startAccountDeletionNow(); + }, + 4: async () => { + reachedB(); + await gateB; + }, + }; + const spy = jest.spyOn(User, 'findById').mockImplementation((( + ...args: Parameters + ) => { + const hook = before[++reads]; + const query = original(...args); + if (hook == null) { + return query; + } + const lean = query.lean.bind(query); + query.lean = ((...leanArgs: Parameters) => ({ + then: (resolve: (value: unknown) => unknown, reject: (error: unknown) => unknown) => + hook() + .then(() => lean(...leanArgs)) + .then(resolve, reject), + })) as unknown as typeof query.lean; + return query; + }) as typeof User.findById); + try { + await enableOpen({ ...key, binding: pullTwo }); + releaseB(); + await enableB; + } finally { + spy.mockRestore(); + } + await cancelAccountDeletion(); + expect(await claim(1, { binding: pullTwo })).not.toMatchObject({ ok: true }); + }); + + /** + * Enables `binding` and runs `during` at the enable's second owner read: its writes are done + * and its rollback is next. Account deletion starts right after `during`. + */ + const enableCrossedBy = async (during: () => Promise, binding = pullTwo) => { + const spy = beforeOwnerRead(async () => { + await during(); + await startAccountDeletionNow(); + }); + try { + return await enableOpen({ ...key, binding }); + } finally { + spy.mockRestore(); + await cancelAccountDeletion(); + } + }; + + test('keeps a user stop that lands between an enable reading the record and writing it', async () => { + await enable(pullOne); + const spy = interleave( + () => methods.stopPRAutomation(key, 'user_stopped'), + PRAutomation, + 'findOne', + ); + try { + expect(await enableOpen({ ...key, binding: pullTwo })).toEqual({ + ok: false, + error: { code: 'conflict' }, + }); + } finally { + spy.mockRestore(); + } + expect(await methods.getPRAutomation(key)).toMatchObject({ + ...pullOne, + state: 'stopped', + stopCode: 'user_stopped', + }); + }); + + test('moves the version on every write', async () => { + const versions: number[] = []; + const read = async () => versions.push((await methods.getPRAutomation(key))?.version ?? -1); + await enable(pullOne); + await read(); + await enable(pullOne); + await read(); + await methods.setPRAutomationTrust(key, 'collaborators', 'anyone'); + await read(); + await addBot(7); + await read(); + const run = await startRound(1); + await read(); + await settle(run.round, run.runId); + await read(); + await methods.removePRAutomationBot(key, 7, await fenceFor()); + await read(); + await methods.stopPRAutomation(key, 'user_stopped'); + await read(); + expect(versions.every((version, index) => index === 0 || version > versions[index - 1])).toBe( + true, + ); + }); + + test('does not write over a record disabled and enabled again after its read', async () => { + await enable(pullOne); + const before = await methods.getPRAutomation(key); + const spy = interleave( + async () => { + await methods.disablePRAutomation(key); + await enable(pullOne); + }, + PRAutomation, + 'findOne', + ); + try { + expect(await enableOpen({ ...key, binding: pullTwo })).toEqual({ + ok: false, + error: { code: 'conflict' }, + }); + } finally { + spy.mockRestore(); + } + const renewed = await methods.getPRAutomation(key); + expect(renewed?.version).toBe(before?.version); + expect(renewed).toMatchObject(pullOne); + expect(renewed?.pending).toBeUndefined(); + expect((await claim(1)).ok).toBe(true); + }); + + test('keeps a user stop that lands while a rejected enable rolls back', async () => { + await enable(pullOne); + const result = await enableCrossedBy(() => methods.stopPRAutomation(key, 'user_stopped')); + expect(result).toEqual({ ok: false, error: { code: 'owner_inactive' } }); + expect(await methods.getPRAutomation(key)).toMatchObject({ + state: 'stopped', + stopCode: 'user_stopped', + }); + expect(await claim(1)).not.toMatchObject({ ok: true }); + }); + + test('keeps a trust change that lands while a rejected enable rolls back', async () => { + await enable(pullOne); + const result = await enableCrossedBy(() => + methods.setPRAutomationTrust(key, 'collaborators', 'anyone'), + ); + expect(result).toEqual({ ok: false, error: { code: 'owner_inactive' } }); + expect(await methods.getPRAutomation(key)).toMatchObject({ trust: 'collaborators' }); + expect(await claim(1, { binding: pullTwo })).not.toMatchObject({ ok: true }); + await enable(pullOne); + expect(await methods.getPRAutomation(key)).toMatchObject({ + trust: 'collaborators', + ...pullOne, + }); + expect((await claim(1)).ok).toBe(true); + }); + + test('keeps a trust change on a record a rejected first enable inserted', async () => { + await PRAutomation.deleteMany({}); + await enableCrossedBy( + () => methods.setPRAutomationTrust(key, 'collaborators', 'anyone'), + pullOne, + ); + expect(await methods.getPRAutomation(key)).toMatchObject({ trust: 'collaborators' }); + expect(await claim(1)).not.toMatchObject({ ok: true }); + }); + + test('does not report success when another enable replaced its write before it committed', async () => { + await enable(pullOne); + /** Another enable's write lands between this enable's second check and its commit. */ + const spy = beforeOwnerRead(() => + PRAutomation.collection.updateOne( + { user: userId, conversationId: key.conversationId }, + { $set: { epoch: 'other-attempt', pending: 'other-attempt' }, $inc: { version: 1 } }, + ), + ); + try { + const result = await methods.enablePRAutomation({ + ...key, + binding: pullTwo, + maxTrust: 'anyone', + }); + expect(result.ok).toBe(false); + } finally { + spy.mockRestore(); + } + }); + + test('lets the next enable recover a record an interrupted enable left pending', async () => { + await enable(pullOne); + await PRAutomation.updateOne( + { user: userId, conversationId: key.conversationId }, + { $set: { pending: 'interrupted-attempt' } }, + ); + await enable(pullOne); + expect((await claim(1)).ok).toBe(true); + }); + + test('keeps a record the crossing enable did not change', async () => { + await enable(pullOne); + await startRound(1); + const spy = interleaveEnable(startAccountDeletionNow); + try { + await enableOpen({ ...key, binding: pullOne }); + } finally { + spy.mockRestore(); + } + expect(await methods.getPRAutomation(key)).toMatchObject({ state: 'fixing', round: 1 }); + }); + + test('does not leave a record behind for an owner that no longer exists', async () => { + await mongoose.models.User.deleteMany({}); + await enableOpen({ ...key, binding: pullOne }); + expect(await methods.getPRAutomation(key)).toBeNull(); + }); +}); + +describe('removing records in bulk', () => { + test('removes every listed conversation, across more than one batch', async () => { + const ids = Array.from({ length: 2500 }, (_, index) => `convo-${index}`); + await PRAutomation.insertMany(ids.map((conversationId) => ({ user: userId, conversationId }))); + await methods.deletePRAutomations(userId, ids); + expect(await PRAutomation.countDocuments({ user: userId })).toBe(0); + }); + + test('bounds each write instead of sending one unbounded list', async () => { + const ids = Array.from({ length: 2500 }, (_, index) => `convo-${index}`); + const spy = jest.spyOn(PRAutomation, 'deleteMany'); + try { + await methods.deletePRAutomations(userId, ids); + const sizes = spy.mock.calls.map( + ([filter]) => (filter as { conversationId: { $in: string[] } }).conversationId.$in.length, + ); + expect(Math.max(...sizes)).toBeLessThanOrEqual(1000); + } finally { + spy.mockRestore(); + } + }); + + test('leaves other users alone', async () => { + await PRAutomation.create({ user: 'someone-else', conversationId: 'convo-1' }); + await methods.deletePRAutomations(userId, ['convo-1']); + expect(await PRAutomation.countDocuments({ user: 'someone-else' })).toBe(1); + }); +}); + +describe('claiming for a subagent thread whose root conversation is gone', () => { + const childKey = { userId, conversationId: 'child-1' }; + const seedChild = () => + mongoose.models.Conversation.create({ + conversationId: childKey.conversationId, + user: userId, + endpoint: 'agents', + subagentThread: { + rootConversationId: key.conversationId, + parentConversationId: key.conversationId, + parentMessageId: 'message-1', + parentToolCallId: 'call-1', + subagentType: 'agent-child', + subagentKind: 'agent', + depth: 1, + }, + }); + const claimChild = async () => + methods.claimPRAutomationRound({ + ...childKey, + ...limits, + binding: { ...pullOne, epoch: await currentEpoch(childKey) }, + headSha: head(1), + }); + + const recreateRoot = async () => { + await mongoose.models.Conversation.deleteOne({ conversationId: key.conversationId }); + await seedConversation(); + }; + + test('refuses and reaps a child whose root was deleted and saved again', async () => { + await seedChild(); + await enableOpen({ ...childKey, binding: pullOne }); + await recreateRoot(); + expect(await methods.reapPRAutomations()).toEqual({ scanned: 1, removed: 1 }); + await enableOpen({ ...childKey, binding: pullOne }); + await recreateRoot(); + expect(await claimChild()).toEqual({ ok: false, error: { code: 'conversation_gone' } }); + }); + + test('refuses a grandchild whose middle ancestor was deleted and saved again', async () => { + const leafKey = { userId, conversationId: 'leaf-1' }; + const thread = (parent: string) => ({ + rootConversationId: key.conversationId, + parentConversationId: parent, + parentMessageId: 'message-1', + parentToolCallId: 'call-1', + subagentType: 'agent-child', + subagentKind: 'agent', + depth: 1, + }); + await seedChild(); + await mongoose.models.Conversation.create({ + conversationId: leafKey.conversationId, + user: userId, + endpoint: 'agents', + subagentThread: { ...thread(childKey.conversationId), depth: 2 }, + }); + await enableOpen({ ...leafKey, binding: pullOne }); + await mongoose.models.Conversation.deleteOne({ conversationId: childKey.conversationId }); + await seedChild(); + expect( + await methods.claimPRAutomationRound({ + ...leafKey, + ...limits, + binding: { ...pullOne, epoch: await currentEpoch(leafKey) }, + headSha: head(1), + }), + ).toEqual({ ok: false, error: { code: 'conversation_gone' } }); + }); + + test('starts a child over for the root saved again when it is enabled', async () => { + await seedChild(); + await enableOpen({ ...childKey, binding: pullOne, trust: 'anyone' }); + const before = await methods.getPRAutomation(childKey); + await recreateRoot(); + await enableOpen({ ...childKey, binding: pullOne }); + const after = await methods.getPRAutomation(childKey); + expect(after?.epoch).not.toBe(before?.epoch); + expect(after?.lineageRefs).not.toEqual(before?.lineageRefs); + expect(after).toMatchObject({ trust: 'approvedBots', round: 0 }); + expect((await claimChild()).ok).toBe(true); + }); + + test('rejects a claim for a child once its root was deleted', async () => { + await seedChild(); + await enableOpen({ ...childKey, binding: pullOne }); + await mongoose.models.Conversation.deleteOne({ conversationId: key.conversationId }); + expect(await claimChild()).toEqual({ ok: false, error: { code: 'conversation_gone' } }); + }); + + test('removes the child record so nothing is left to claim later', async () => { + await seedChild(); + await enableOpen({ ...childKey, binding: pullOne }); + await mongoose.models.Conversation.deleteOne({ conversationId: key.conversationId }); + await claimChild(); + expect(await methods.getPRAutomation(childKey)).toBeNull(); + }); + + test('rejects a claim for a child whose root passed its retention date', async () => { + await seedChild(); + await enableOpen({ ...childKey, binding: pullOne }); + await mongoose.models.Conversation.updateOne( + { conversationId: key.conversationId }, + { $set: { expiredAt: new Date(Date.now() - 60_000) } }, + ); + expect(await claimChild()).toEqual({ ok: false, error: { code: 'conversation_gone' } }); + }); + + test('rejects a claim for a grandchild once its parent was deleted but its root remains', async () => { + await seedChild(); + await mongoose.models.Conversation.create({ + conversationId: 'grandchild-1', + user: userId, + endpoint: 'agents', + subagentThread: { + rootConversationId: key.conversationId, + parentConversationId: childKey.conversationId, + parentMessageId: 'message-2', + parentToolCallId: 'call-2', + subagentType: 'agent-child', + subagentKind: 'agent', + depth: 2, + }, + }); + const grandKey = { userId, conversationId: 'grandchild-1' }; + await enableOpen({ ...grandKey, binding: pullOne }); + await mongoose.models.Conversation.deleteOne({ conversationId: childKey.conversationId }); + expect( + await methods.claimPRAutomationRound({ + ...grandKey, + ...limits, + binding: { ...pullOne, epoch: await currentEpoch(grandKey) }, + headSha: head(1), + }), + ).toEqual({ ok: false, error: { code: 'conversation_gone' } }); + }); + + test('rejects a claim for a descendant once an ancestor between its parent and root was deleted', async () => { + await seedChild(); + const thread = (conversationId: string, parentConversationId: string, depth: number) => + mongoose.models.Conversation.create({ + conversationId, + user: userId, + endpoint: 'agents', + subagentThread: { + rootConversationId: key.conversationId, + parentConversationId, + parentMessageId: `message-${depth}`, + parentToolCallId: `call-${depth}`, + subagentType: 'agent-child', + subagentKind: 'agent', + depth, + }, + }); + await thread('grandchild-1', childKey.conversationId, 2); + await thread('great-1', 'grandchild-1', 3); + const greatKey = { userId, conversationId: 'great-1' }; + await enableOpen({ ...greatKey, binding: pullOne }); + await mongoose.models.Conversation.deleteOne({ conversationId: childKey.conversationId }); + expect( + await methods.claimPRAutomationRound({ + ...greatKey, + ...limits, + binding: { ...pullOne, epoch: await currentEpoch(greatKey) }, + headSha: head(1), + }), + ).toEqual({ ok: false, error: { code: 'conversation_gone' } }); + }); + + test("does not let another tenant's conversation stand in for a deleted parent", async () => { + await mongoose.models.Conversation.deleteMany({}); + const inTenant = (tenantId: string, run: () => Promise) => + tenantStorage.run({ tenantId }, run); + await inTenant('tenant-b', () => seedConversation()); + await mongoose.models.User.deleteMany({}); + await inTenant('tenant-a', () => + mongoose.models.User.create({ _id: userId, email: 'owner@example.com', provider: 'local' }), + ); + await inTenant('tenant-a', async () => { + await seedConversation(); + await seedChild(); + await enableWith({ ...childKey, binding: pullOne }); + await mongoose.models.Conversation.deleteOne({ conversationId: key.conversationId }); + expect(await claimChild()).toEqual({ ok: false, error: { code: 'conversation_gone' } }); + }); + }); + + test('accepts a claim for a child while its root is active', async () => { + await seedChild(); + await enableOpen({ ...childKey, binding: pullOne }); + expect((await claimChild()).ok).toBe(true); + }); +}); + +describe('pull request numbers', () => { + test.each([1.5, 0.5, 2.0000001])('rejects %s as a pull request number', async (pullNumber) => { + await expect(enable({ repository: 'acme/one', pullNumber })).rejects.toThrow(); + expect(await methods.getPRAutomation(key)).toBeNull(); + }); + + test('still accepts a whole pull request number', async () => { + expect(await enable({ repository: 'acme/one', pullNumber: 42 })).toMatchObject({ + pullNumber: 42, + }); + }); +}); + +describe('index guarantees', () => { + test('keeps one record per tenant for the same user and conversation id', async () => { + const inTenant = (tenantId: string, run: () => Promise) => + tenantStorage.run({ tenantId }, run); + await inTenant('tenant-a', () => PRAutomation.create(keyFilterOf(key))); + await inTenant('tenant-b', () => PRAutomation.create(keyFilterOf(key))); + await expect(inTenant('tenant-b', () => PRAutomation.create(keyFilterOf(key)))).rejects.toThrow( + /E11000/, + ); + expect(await runAsSystem(() => PRAutomation.countDocuments({}))).toBe(2); + }); + + test('builds the unique index before the first write when automatic indexing is off', async () => { + await PRAutomation.collection.dropIndexes(); + const fresh = createPRAutomationMethods(mongoose); + await fresh.enablePRAutomation({ ...key, binding: pullOne, maxTrust: 'anyone' }); + const indexes = await PRAutomation.collection.indexes(); + expect( + indexes.some( + (index) => index.unique === true && index.key.user === 1 && index.key.conversationId === 1, + ), + ).toBe(true); + }); +}); + +describe('binding a pull request', () => { + test('keeps a stopped record fenced when the new binding is rejected', async () => { + await enable(pullOne); + await methods.stopPRAutomation(key, 'user_stopped'); + + await expect(enable({ repository: 'acme/one', pullNumber: 0 })).rejects.toThrow(); + + expect(await methods.getPRAutomation(key)).toMatchObject({ + state: 'stopped', + stopCode: 'user_stopped', + pullNumber: 1, + }); + }); + + test('revives a stopped record onto the pull request it is bound to now', async () => { + await enable(pullOne); + await methods.stopPRAutomation(key, 'user_stopped'); + expect(await enable(pullTwo)).toMatchObject({ state: 'idle', pullNumber: 2, round: 0 }); + }); + + test('starts a fresh run when bound to another pull request in the same repository', async () => { + await enableOpen({ ...key, binding: pullOne }); + await startRound(1); + await toWaiting(); + await startRound(2); + + const rebound = await enableWith({ ...key, binding: pullTwo }); + expect(rebound).toMatchObject({ pullNumber: 2, state: 'idle', round: 0, claimedHeads: [] }); + expect(rebound.startedAt).toBeUndefined(); + expect(rebound.lastHeadSha).toBeUndefined(); + }); + + test('lets the new pull request claim a head the old one claimed', async () => { + await enableOpen({ ...key, binding: pullOne }); + await startRound(1); + await enableOpen({ ...key, binding: pullTwo }); + expect((await claim(1, { binding: pullTwo })).ok).toBe(true); + }); + + test('does not carry the old pull request round count into the new one', async () => { + await enableOpen({ ...key, binding: pullOne }); + for (let round = 1; round <= limits.maxRounds; round++) { + await startRound(round); + await toWaiting(); + } + await enableOpen({ ...key, binding: pullTwo }); + expect((await claim(9, { binding: pullTwo })).ok).toBe(true); + }); + + test('ignores a completion from the round of the pull request it replaced', async () => { + await enableOpen({ ...key, binding: pullOne }); + const running = await startRound(1); + await enableOpen({ ...key, binding: pullTwo }); + + expect(await settle(1, running.runId)).toBeNull(); + expect((await methods.getPRAutomation(key))?.state).toBe('idle'); + }); + + test('leaves a run alone when it is bound to the pair it already has', async () => { + await enableOpen({ ...key, binding: pullOne }); + await startRound(1); + const again = await enableWith({ ...key, binding: pullOne }); + expect(again).toMatchObject({ state: 'fixing', round: 1 }); + }); + + /** Guard, not a proven regression: the interleaving it protects against is timing dependent. */ + test('never leaves a repository and pull request pair nobody asked for', async () => { + await enable(); + const pairs = [pullOne, otherRepository]; + await Promise.all( + Array.from({ length: 24 }, (_, index) => enableOpen({ ...key, binding: pairs[index % 2] })), + ); + const record = await methods.getPRAutomation(key); + expect(pairs).toContainEqual({ + repository: record?.repository, + pullNumber: record?.pullNumber, + }); + }); +}); + +describe('claimPRAutomationRound', () => { + test('starts a round, counts it and records the head', async () => { + await enable(); + expect(await claim(1)).toMatchObject({ + ok: true, + value: { state: 'fixing', round: 1, lastHeadSha: head(1) }, + }); + }); + + test('rejects a second delivery for a head that was already claimed', async () => { + await enable(); + await claim(1); + await toWaiting(); + expect(await claim(1)).toEqual({ ok: false, error: { code: 'stale_head' } }); + expect((await methods.getPRAutomation(key))?.round).toBe(1); + }); + + test('rejects a delayed delivery of an earlier head after a later head was claimed', async () => { + await enable(); + await claim(1); + await toWaiting(); + await claim(2); + await toWaiting(); + expect(await claim(1)).toEqual({ ok: false, error: { code: 'stale_head' } }); + expect((await methods.getPRAutomation(key))?.round).toBe(2); + }); + + test('reports a retryable conflict when the running round settles before the fallback read', async () => { + await enable(); + await claim(1); + const spy = interleave(toWaiting); + try { + expect(await claim(2)).toEqual({ ok: false, error: { code: 'conflict' } }); + } finally { + spy.mockRestore(); + } + }); + + test('refuses a head that is not a commit SHA before writing anything', async () => { + await enable(); + await expect( + methods.claimPRAutomationRound({ + ...key, + ...limits, + binding: { ...pullOne, epoch: await currentEpoch() }, + headSha: 'a'.repeat(65), + }), + ).rejects.toThrow(RangeError); + expect((await methods.getPRAutomation(key))?.claimedHeads).toEqual([]); + }); + + test('treats a head in upper case as the head it already claimed', async () => { + await enable(); + await claim(1); + await toWaiting(); + const upper = await methods.claimPRAutomationRound({ + ...key, + ...limits, + binding: { ...pullOne, epoch: await currentEpoch() }, + headSha: 'A'.repeat(40), + }); + expect(upper).toEqual({ ok: false, error: { code: 'stale_head' } }); + }); + + test('accepts a SHA-256 head', async () => { + await enable(); + const result = await methods.claimPRAutomationRound({ + ...key, + ...limits, + binding: { ...pullOne, epoch: await currentEpoch() }, + headSha: 'b'.repeat(64), + }); + expect(result.ok).toBe(true); + }); + + test('rejects a delivery from a run that was stopped and restarted on the same pull request', async () => { + await enable(pullOne); + const stale = await currentEpoch(); + await methods.stopPRAutomation(key, 'user_stopped'); + await enable(pullOne); + expect(await claim(1, { epoch: stale })).toEqual(mismatch); + }); + + test('rejects a claim while a round is already running', async () => { + await enable(); + await claim(1); + expect(await claim(2)).toEqual({ ok: false, error: { code: 'not_active' } }); + }); + + test('stops at the round cap and never goes past it', async () => { + await enable(); + for (let round = 1; round <= limits.maxRounds; round++) { + expect((await claim(round)).ok).toBe(true); + await toWaiting(); + } + expect(await claim(9)).toEqual({ ok: false, error: { code: 'round_cap' } }); + expect((await methods.getPRAutomation(key))?.round).toBe(limits.maxRounds); + }); + + test('lets exactly one of several concurrent deliveries claim the same head', async () => { + await enable(); + const results = await Promise.all(Array.from({ length: 8 }, () => claim(1))); + expect(results.filter((result) => result.ok)).toHaveLength(1); + expect((await methods.getPRAutomation(key))?.round).toBe(1); + }); + + test('never exceeds the cap under concurrent deliveries for different heads', async () => { + await enable(); + const results = await Promise.all(Array.from({ length: 8 }, (_, index) => claim(index + 1))); + expect(results.filter((result) => result.ok)).toHaveLength(1); + expect((await methods.getPRAutomation(key))?.round).toBeLessThanOrEqual(limits.maxRounds); + }); + + test('stops once the wall-clock window has passed', async () => { + await enable(); + const start = new Date('2026-01-01T00:00:00Z'); + await claim(1, { now: start }); + await toWaiting(); + + const later = new Date(start.getTime() + (limits.maxMinutes + 1) * 60_000); + expect(await claim(2, { now: later })).toEqual({ ok: false, error: { code: 'time_cap' } }); + }); + + test('reports not_found for a conversation with no record', async () => { + expect(await claim(1)).toEqual({ ok: false, error: { code: 'not_found' } }); + }); + + test('rejects a claim on a stopped record', async () => { + await enable(); + await methods.stopPRAutomation(key, 'user_stopped'); + expect(await claim(1)).toEqual({ ok: false, error: { code: 'not_active' } }); + }); + + test('rejects a delivery for the pull request the record was rebound away from', async () => { + await enable(pullOne); + await enable(pullTwo); + expect(await claim(1, { binding: pullOne })).toEqual(mismatch); + expect((await methods.getPRAutomation(key))?.round).toBe(0); + }); + + test('rejects a delivery for a repository the record is not bound to', async () => { + await enable(pullOne); + expect(await claim(1, { binding: otherRepository })).toEqual(mismatch); + }); + + test('rejects a claim on a record that is bound to no pull request', async () => { + await enableOpen(key); + expect(await claim(1)).toEqual(mismatch); + }); + + test('does not open the time window when the claim is refused', async () => { + await enable(pullOne); + const refused = await methods.claimPRAutomationRound({ + ...key, + binding: { ...pullOne, epoch: await currentEpoch() }, + maxRounds: 0, + maxMinutes: limits.maxMinutes, + headSha: head(1), + }); + expect(refused).toEqual({ ok: false, error: { code: 'round_cap' } }); + expect((await methods.getPRAutomation(key))?.startedAt).toBeUndefined(); + }); + + test('opens the time window with the first round that is claimed', async () => { + await enable(pullOne); + const start = new Date('2026-01-01T00:00:00Z'); + await claim(1, { now: start }); + expect((await methods.getPRAutomation(key))?.startedAt).toEqual(start); + }); + + test('keeps the first start when a later round is claimed', async () => { + await enable(pullOne); + const start = new Date('2026-01-01T00:00:00Z'); + await claim(1, { now: start }); + await toWaiting(); + await claim(2, { now: new Date(start.getTime() + 10 * 60_000) }); + expect((await methods.getPRAutomation(key))?.startedAt).toEqual(start); + }); + + test('does not open the time window for a delivery it rejected', async () => { + await enable(pullOne); + await enable(pullTwo); + await claim(1, { binding: pullOne }); + expect((await methods.getPRAutomation(key))?.startedAt).toBeUndefined(); + }); +}); + +describe('claiming for a conversation that no longer exists', () => { + const gone = { ok: false, error: { code: 'conversation_gone' } }; + + test('rejects a claim once the conversation was removed by the retention index', async () => { + await enable(); + await mongoose.models.Conversation.deleteMany({}); + expect(await claim(1)).toEqual(gone); + }); + + test('removes the record so nothing is left to claim later', async () => { + await enable(); + await mongoose.models.Conversation.deleteMany({}); + await claim(1); + expect(await methods.getPRAutomation(key)).toBeNull(); + }); + + test('rejects a claim for a conversation that expired but is not yet purged', async () => { + await enable(); + await mongoose.models.Conversation.deleteMany({}); + await seedConversation({ expiredAt: new Date(Date.now() - 60_000) }); + expect(await claim(1)).toEqual(gone); + }); + + test('accepts a claim while the conversation retention date is still ahead', async () => { + await enable(); + await mongoose.models.Conversation.updateOne( + { conversationId: key.conversationId }, + { $set: { expiredAt: new Date(Date.now() + 3_600_000) } }, + ); + expect((await claim(1)).ok).toBe(true); + }); + + test('does not spend a round for a conversation that is gone', async () => { + await enable(); + await mongoose.models.Conversation.deleteMany({}); + await claim(1); + await seedConversation(); + await enable(); + expect(await startRound(1)).toMatchObject({ round: 1 }); + }); + + test('keeps reporting not_found when there is no record at all', async () => { + await mongoose.models.Conversation.deleteMany({}); + expect(await claim(1)).toEqual({ ok: false, error: { code: 'not_found' } }); + }); + + test('only looks at the conversation of the same owner', async () => { + await enable(); + await mongoose.models.Conversation.deleteMany({}); + await seedConversation({ user: otherUserId }); + expect(await claim(1)).toEqual(gone); + }); +}); + +describe('reaping records the retention index left behind', () => { + test('removes an idle record once its conversation was purged, with no delivery', async () => { + await enable(); + await mongoose.models.Conversation.deleteMany({}); + expect(await methods.reapPRAutomations()).toEqual({ scanned: 1, removed: 1 }); + expect(await methods.getPRAutomation(key)).toBeNull(); + }); + + test('removes a record whose conversation is past its retention date', async () => { + await enable(); + await mongoose.models.Conversation.updateOne( + { conversationId: key.conversationId }, + { $set: { expiredAt: new Date(Date.now() - 60_000) } }, + ); + await methods.reapPRAutomations(); + expect(await methods.getPRAutomation(key)).toBeNull(); + }); + + test('keeps live records and records whose owner deletion can still be cancelled', async () => { + await enable(); + const liveKey = { userId, conversationId: 'convo-live' }; + await seedConversation({ conversationId: liveKey.conversationId }); + await enableOpen({ ...liveKey, binding: pullOne }); + await mongoose.models.Conversation.deleteMany({ conversationId: key.conversationId }); + await startAccountDeletionNow(); + expect(await methods.reapPRAutomations()).toEqual({ scanned: 2, removed: 0 }); + await cancelAccountDeletion(); + expect(await methods.reapPRAutomations()).toEqual({ scanned: 2, removed: 1 }); + expect(await methods.getPRAutomation(liveKey)).not.toBeNull(); + }); + + test('judges each record by its own tenant when the same chat exists in two', async () => { + const inTenant = (tenantId: string, run: () => Promise) => + tenantStorage.run({ tenantId }, run); + await mongoose.models.Conversation.deleteMany({}); + await inTenant('tenant-b', () => seedConversation()); + await inTenant('tenant-a', () => PRAutomation.create({ ...keyFilterOf(key), ...pullOne })); + await inTenant('tenant-b', () => PRAutomation.create({ ...keyFilterOf(key), ...pullOne })); + expect(await runAsSystem(() => methods.reapPRAutomations())).toEqual({ + scanned: 2, + removed: 1, + }); + const left = await runAsSystem(() => PRAutomation.find({}).lean<{ tenantId?: string }[]>()); + expect(left.map((record) => record.tenantId)).toEqual(['tenant-b']); + }); + + test('keeps a record renewed between its check and its delete', async () => { + await enable(); + await mongoose.models.Conversation.deleteMany({}); + const original = PRAutomation.deleteOne.bind(PRAutomation); + const spy = jest.spyOn(PRAutomation, 'deleteOne').mockImplementationOnce((( + ...args: Parameters + ) => ({ + then: (resolve: (value: unknown) => unknown, reject: (error: unknown) => unknown) => + seedConversation() + .then(() => enable()) + .then(() => original(...args)) + .then(resolve, reject), + })) as unknown as typeof PRAutomation.deleteOne); + try { + expect(await methods.reapPRAutomations()).toEqual({ scanned: 1, removed: 0 }); + } finally { + spy.mockRestore(); + } + expect(await methods.getPRAutomation(key)).toMatchObject(pullOne); + expect((await claim(1)).ok).toBe(true); + }); + + test('checks a few records at a time within a page', async () => { + const ids = Array.from({ length: 20 }, (_, index) => `orphan-${index}`); + await PRAutomation.insertMany(ids.map((conversationId) => ({ user: userId, conversationId }))); + const Conversation = mongoose.models.Conversation; + const original = Conversation.aggregate.bind(Conversation); + let inFlight = 0; + let peak = 0; + const spy = jest.spyOn(Conversation, 'aggregate').mockImplementation((( + ...args: Parameters + ) => { + inFlight++; + peak = Math.max(peak, inFlight); + const query = original(...args); + const exec = query.exec.bind(query); + query.exec = (() => + exec().finally(() => { + inFlight--; + })) as typeof query.exec; + return query; + }) as typeof Conversation.aggregate); + try { + expect(await methods.reapPRAutomations()).toEqual({ scanned: 20, removed: 20 }); + } finally { + spy.mockRestore(); + } + expect(peak).toBeGreaterThan(1); + expect(peak).toBeLessThanOrEqual(8); + }); + + test('removes a record of a conversation deleted and saved again under the same id', async () => { + await enable(); + await mongoose.models.Conversation.deleteMany({}); + await seedConversation(); + expect(await methods.reapPRAutomations()).toEqual({ scanned: 1, removed: 1 }); + expect(await methods.getPRAutomation(key)).toBeNull(); + }); + + test('pages through every record', async () => { + const ids = Array.from({ length: 5 }, (_, index) => `orphan-${index}`); + await PRAutomation.insertMany(ids.map((conversationId) => ({ user: userId, conversationId }))); + let next: string | undefined; + let removed = 0; + let pages = 0; + do { + const result = await methods.reapPRAutomations({ after: next, limit: 2 }); + removed += result.removed; + next = result.next; + pages++; + } while (next != null); + expect(pages).toBe(3); + expect(removed).toBe(5); + expect(await PRAutomation.countDocuments({})).toBe(0); + }); +}); + +describe('settlePRAutomationRound', () => { + test('settles the round that owns it, including from needs_user', async () => { + await enable(); + const { runId } = await startRound(1); + expect(await settle(1, runId, 'needs_user')).toMatchObject({ state: 'needs_user', round: 1 }); + expect(await settle(1, runId, 'waiting')).toMatchObject({ state: 'waiting', round: 1 }); + }); + + test('ignores a completion that belongs to an earlier round', async () => { + await enable(); + const { runId } = await startRound(1); + await settle(1, runId); + await claim(2); + + expect(await settle(1, runId, 'needs_user')).toBeNull(); + expect(await methods.getPRAutomation(key)).toMatchObject({ state: 'fixing', round: 2 }); + }); + + test('does not let a stale completion open the way for an overlapping claim', async () => { + await enable(); + const { runId } = await startRound(1); + await settle(1, runId); + await claim(2); + await settle(1, runId); + expect(await claim(3)).toEqual({ ok: false, error: { code: 'not_active' } }); + }); + + test('ignores a completion from a run that was stopped and restarted', async () => { + await enable(); + const previous = await startRound(1); + await methods.stopPRAutomation(key, 'user_stopped'); + await enable(); + await startRound(1); + + expect(await settle(1, previous.runId, 'needs_user')).toBeNull(); + expect(await methods.getPRAutomation(key)).toMatchObject({ state: 'fixing', round: 1 }); + }); + + test('ignores a completion from a record that was disabled and enabled again', async () => { + await enable(); + const previous = await startRound(1); + await methods.disablePRAutomation(key); + await enable(); + await startRound(1); + + expect(await settle(1, previous.runId, 'needs_user')).toBeNull(); + expect(await methods.getPRAutomation(key)).toMatchObject({ state: 'fixing', round: 1 }); + }); + + test('does not settle a record that never started a round', async () => { + await enable(); + expect(await settle(0, 'no-run')).toBeNull(); + expect((await methods.getPRAutomation(key))?.state).toBe('idle'); + }); + + test('does not revive a stopped record', async () => { + await enable(); + const { runId } = await startRound(1); + await methods.stopPRAutomation(key, 'user_stopped'); + expect(await settle(1, runId)).toBeNull(); + expect((await methods.getPRAutomation(key))?.state).toBe('stopped'); + }); + + test('returns null for a conversation with no record', async () => { + expect(await settle(1, 'missing-run')).toBeNull(); + }); +}); + +describe('stopPRAutomation', () => { + test('records the stop code', async () => { + await enable(); + expect(await methods.stopPRAutomation(key, 'round_cap', await fenceFor())).toMatchObject({ + state: 'stopped', + stopCode: 'round_cap', + }); + }); + + test('stops a running round without waiting for it to settle', async () => { + await enable(); + await claim(1); + expect(await methods.stopPRAutomation(key, 'user_stopped')).toMatchObject({ + state: 'stopped', + stopCode: 'user_stopped', + }); + }); + + test('keeps the first stop code when it is stopped again', async () => { + await enable(); + await methods.stopPRAutomation(key, 'round_cap', await fenceFor()); + expect(await methods.stopPRAutomation(key, 'user_stopped')).toBeNull(); + expect((await methods.getPRAutomation(key))?.stopCode).toBe('round_cap'); + }); + + test('rejects a code outside the stop code list', async () => { + await enable(); + await expect( + methods.stopPRAutomation( + key, + 'because' as unknown as PRAutomationEventStopCode, + await fenceFor(), + ), + ).rejects.toThrow(); + }); + + test('returns null for a conversation with no record', async () => { + expect(await methods.stopPRAutomation(key, 'user_stopped')).toBeNull(); + }); + + test('ignores an event-driven stop for a pull request the record was rebound away from', async () => { + await enable(pullOne); + await enable(pullTwo); + expect(await methods.stopPRAutomation(key, 'pull_request_closed', await fenceFor())).toBeNull(); + expect(await methods.getPRAutomation(key)).toMatchObject({ state: 'idle', pullNumber: 2 }); + }); + + test('ignores an event-driven stop for another repository', async () => { + await enable(pullOne); + expect( + await methods.stopPRAutomation(key, 'pull_request_closed', await fenceFor(otherRepository)), + ).toBeNull(); + expect((await methods.getPRAutomation(key))?.state).toBe('idle'); + }); + + test('applies an event-driven stop for the pull request the record is bound to', async () => { + await enable(pullOne); + expect( + await methods.stopPRAutomation(key, 'pull_request_closed', await fenceFor()), + ).toMatchObject({ + state: 'stopped', + stopCode: 'pull_request_closed', + }); + }); + + test('refuses an event-driven stop that names no pull request', async () => { + await enable(pullOne); + await expect( + ( + methods.stopPRAutomation as ( + target: typeof key, + code: PRAutomationStopCode, + ) => Promise + )(key, 'pull_request_closed'), + ).rejects.toThrow(RangeError); + expect((await methods.getPRAutomation(key))?.state).toBe('idle'); + }); + + test('applies a user stop that names a pull request the record was rebound away from', async () => { + await enable(pullOne); + await enable(pullTwo); + const stop = methods.stopPRAutomation as ( + target: typeof key, + code: PRAutomationStopCode, + binding: typeof pullOne, + ) => Promise; + expect(await stop(key, 'user_stopped', pullOne)).toMatchObject({ state: 'stopped' }); + }); + + test('ignores an event stop from a run that was stopped and restarted on the same pull request', async () => { + await enable(pullOne); + const previous = await fenceFor(); + await methods.stopPRAutomation(key, 'user_stopped'); + await enable(pullOne); + expect(await methods.stopPRAutomation(key, 'pull_request_closed', previous)).toBeNull(); + }); + + test('ignores an event stop from a record that was disabled and enabled again', async () => { + await enable(pullOne); + const previous = await fenceFor(); + await methods.disablePRAutomation(key); + await enable(pullOne); + expect(await methods.stopPRAutomation(key, 'pull_request_closed', previous)).toBeNull(); + }); + + test('lets a user stop win even after the record was rebound', async () => { + await enable(pullOne); + await enable(pullTwo); + expect(await methods.stopPRAutomation(key, 'user_stopped')).toMatchObject({ + state: 'stopped', + stopCode: 'user_stopped', + }); + }); +}); + +describe('setPRAutomationTrust', () => { + test('changes the trust level', async () => { + await enable(); + const updated = await methods.setPRAutomationTrust(key, 'collaborators', 'anyone'); + expect(updated?.trust).toBe('collaborators'); + }); + + test('stores no level wider than the administrator ceiling', async () => { + await enable(); + const updated = await methods.setPRAutomationTrust(key, 'anyone', 'approvedBots'); + expect(updated?.trust).toBe('approvedBots'); + }); + + test('rejects a level outside the enum', async () => { + await enable(); + await expect( + methods.setPRAutomationTrust(key, 'everyone' as unknown as 'anyone', 'anyone'), + ).rejects.toThrow(); + }); +}); + +describe('approved bots', () => { + test('keeps a bot reapproved after the record was rebound away and back', async () => { + await enable(pullOne); + const stale = { repository: 'acme/one', epoch: await currentEpoch() }; + await enable(otherRepository); + await enable(pullOne); + await addBot(7); + await methods.removePRAutomationBot(key, 7, stale); + expect((await methods.getPRAutomation(key))?.trustedBots).toEqual([{ id: 7 }]); + }); + + test('keeps approved bots when the repository is named with different casing', async () => { + await enable(pullOne); + await addBot(7); + await enable({ repository: 'Acme/One', pullNumber: 1 }); + expect((await methods.getPRAutomation(key))?.trustedBots).toEqual([{ id: 7 }]); + }); + + test('claims for the bound repository named with different casing', async () => { + await enable(pullOne); + expect((await claim(1, { binding: { repository: 'ACME/one', pullNumber: 1 } })).ok).toBe(true); + }); + + test('refuses an approval from before the record was rebound away and back', async () => { + await enable(pullOne); + const stale = { repository: 'acme/one', epoch: (await methods.getPRAutomation(key))!.epoch }; + await enable(otherRepository); + await enable(pullOne); + expect(await methods.addPRAutomationBot(key, { id: 7 }, maxBots, stale)).toEqual(mismatch); + }); + + test('refuses a bot id that is not a whole number', async () => { + await enable(); + await expect(addBot(7.5)).rejects.toThrow(); + expect((await methods.getPRAutomation(key))?.trustedBots).toEqual([]); + }); + + test('adds a bot by numeric id', async () => { + await enable(); + expect(await addBot(101, 'review-bot[bot]')).toMatchObject({ + ok: true, + value: { trustedBots: [{ id: 101, login: 'review-bot[bot]' }] }, + }); + }); + + test('is idempotent by id even when the login was renamed', async () => { + await enable(); + await addBot(101, 'old-name[bot]'); + expect((await addBot(101, 'new-name[bot]')).ok).toBe(true); + const record = await methods.getPRAutomation(key); + expect(record?.trustedBots).toEqual([{ id: 101, login: 'old-name[bot]' }]); + }); + + test('refuses a bot beyond the limit it is given, with a stable code', async () => { + await enable(); + for (let id = 1; id <= maxBots; id++) { + expect({ id, ok: (await addBot(id)).ok }).toEqual({ id, ok: true }); + } + expect(await addBot(9999)).toEqual({ ok: false, error: { code: 'bot_limit' } }); + }); + + test('accepts a bot that is already approved when the list is full', async () => { + await enable(); + for (let id = 1; id <= maxBots; id++) { + await addBot(id); + } + expect((await addBot(1)).ok).toBe(true); + }); + + test('admits more bots once a higher limit is passed', async () => { + await enable(); + for (let id = 1; id <= maxBots; id++) { + await addBot(id); + } + expect((await addBot(maxBots + 1, undefined, maxBots + 2)).ok).toBe(true); + }); + + test('rejects a limit below one instead of storing without a cap', async () => { + await enable(); + await expect(addBot(1, undefined, 0)).rejects.toThrow(RangeError); + }); + + test('reports not_found when the conversation has no record', async () => { + expect(await addBot(101)).toEqual({ ok: false, error: { code: 'not_found' } }); + }); + + test('refuses a bot approved for a repository the record is not bound to', async () => { + await enable(pullOne); + expect(await addBot(101, undefined, maxBots, 'acme/two')).toEqual({ + ok: false, + error: { code: 'binding_mismatch' }, + }); + expect((await methods.getPRAutomation(key))?.trustedBots).toEqual([]); + }); + + test('refuses an approval that was authorized before the record was rebound', async () => { + await enable(pullOne); + await enable(otherRepository); + expect(await addBot(101, undefined, maxBots, 'acme/one')).toEqual({ + ok: false, + error: { code: 'binding_mismatch' }, + }); + expect((await methods.getPRAutomation(key))?.trustedBots).toEqual([]); + }); + + test('still treats an already approved bot as success for the bound repository', async () => { + await enable(pullOne); + await addBot(101); + expect((await addBot(101)).ok).toBe(true); + }); + + test('removes a bot by id and leaves the others', async () => { + await enable(); + await addBot(1); + await addBot(2); + const record = await methods.removePRAutomationBot(key, 1, { + repository: 'acme/one', + epoch: await currentEpoch(), + }); + expect(record?.trustedBots).toEqual([{ id: 2 }]); + }); + + test('keeps an approval made for another repository when a removal is delayed', async () => { + await enable(pullOne); + await addBot(7); + await enable(otherRepository); + await addBot(7, undefined, maxBots, 'acme/two'); + const record = await methods.removePRAutomationBot(key, 7, { + repository: 'acme/one', + epoch: await currentEpoch(), + }); + expect(record).toBeNull(); + expect((await methods.getPRAutomation(key))?.trustedBots).toEqual([{ id: 7 }]); + }); + + test('removes an approval for the repository the record is bound to', async () => { + await enable(pullOne); + await addBot(7); + const record = await methods.removePRAutomationBot(key, 7, { + repository: 'acme/one', + epoch: await currentEpoch(), + }); + expect(record?.trustedBots).toEqual([]); + }); + + test('keeps each conversation allowlist separate', async () => { + await enable(); + await seedConversation({ conversationId: 'convo-2' }); + await enableOpen({ userId, conversationId: 'convo-2' }); + await addBot(101); + const other = await methods.getPRAutomation({ userId, conversationId: 'convo-2' }); + expect(other?.trustedBots).toEqual([]); + }); +}); + +describe('disablePRAutomation', () => { + test('removes the record and reports it', async () => { + await enable(); + expect(await methods.disablePRAutomation(key)).toEqual({ removed: true }); + expect(await methods.getPRAutomation(key)).toBeNull(); + }); + + test('reports nothing removed when there was no record', async () => { + expect(await methods.disablePRAutomation(key)).toEqual({ removed: false }); + }); + + test('still removes a record a user stopped', async () => { + await enable(pullOne); + await methods.stopPRAutomation(key, 'user_stopped'); + expect(await methods.disablePRAutomation(key)).toEqual({ removed: true }); + }); +}); diff --git a/packages/data-schemas/src/methods/prAutomation.ts b/packages/data-schemas/src/methods/prAutomation.ts new file mode 100644 index 00000000000..1ab20edd520 --- /dev/null +++ b/packages/data-schemas/src/methods/prAutomation.ts @@ -0,0 +1,996 @@ +import { randomUUID } from 'node:crypto'; +import { + MAX_SUBAGENT_DEPTH, + MAX_PR_AUTOMATION_BOTS, + PR_AUTOMATION_STOP_CODES, + PR_AUTOMATION_TRUST_LEVELS, + clampPRAutomationTrust, +} from 'librechat-data-provider'; +import type { PRAutomationTrustLevel } from 'librechat-data-provider'; +import type { PRAutomationStopCode } from 'librechat-data-provider'; +import type { PRAutomationState } from 'librechat-data-provider'; +import type { PipelineStage } from 'mongoose'; +import type * as t from '~/types/prAutomation'; +import { isValidObjectIdString } from '~/utils/objectId'; +import { createIndexesWithRetry } from '~/utils/retry'; + +const PROJECTION = '-_id -__v'; +/** Records removed per write, so a cleanup never holds one unbounded `$in`. */ +const DELETE_BATCH = 1000; +/** A full SHA-1 or SHA-256 commit id; anything else is refused before it is stored. */ +const COMMIT_SHA = /^(?:[0-9a-f]{40}|[0-9a-f]{64})$/i; +const DUPLICATE_KEY = 11000; +/** Liveness checks the reaper runs at once within a page. */ +const REAP_CONCURRENCY = 8; + +/** + * `owner_inactive` is an account deletion in progress, which can still be cancelled, so the + * record is kept. `owner_gone` and `conversation_gone` can never run again. + */ +type Liveness = 'live' | 'conversation_gone' | 'owner_gone' | 'owner_inactive'; +type Gone = Extract; +/** + * What a liveness check is judged against. `tenantId`, when the key is present, is the record's + * tenant (`null` for one stored without one). `conversationRef` is the `_id` of the conversation + * document the record was enabled for: a conversation deleted and saved again under the same id + * is a different document, so a record of the old one does not belong to it. `lineageRefs` + * does the same for a subagent thread's ancestors, parent first and root last. + */ +type LivenessScope = { tenantId?: string | null; conversationRef?: string; lineageRefs?: string[] }; +/** The conversation documents a live check found: the record's own and its ancestors'. */ +type Lineage = { conversationRef: string; lineageRefs: string[] }; +const CLAIMABLE_STATES: PRAutomationState[] = ['idle', 'waiting']; +const SETTLEABLE_STATES: PRAutomationState[] = ['fixing', 'needs_user']; +const SETTLED_STATES: PRAutomationState[] = ['waiting', 'needs_user']; +/** Everything that belongs to one run on one pull request, under the given new epoch. */ +const runReset = (epoch: string, fields: Record = {}) => ({ + $set: { state: 'idle', round: 0, claimedHeads: [], epoch, ...fields }, + $unset: { stopCode: '', startedAt: '', lastHeadSha: '', runId: '' }, +}); +/** + * Every write to a record goes through `$inc: { version: 1 }`. A write that acts on a record it + * read first names the version it read, so any change in between, a stop or a trust change + * included, makes it miss instead of being overwritten. + */ +const BUMP = { $inc: { version: 1 } }; +/** + * One record as a write saw it. The version alone is not an identity: a record deleted and + * enabled again starts its count over, so a fence names the document `_id` as well. + */ +type RecordFence = { _id: unknown; version: number }; +const REAP_FIELDS = { + _id: 1, + user: 1, + tenantId: 1, + conversationId: 1, + conversationRef: 1, + lineageRefs: 1, + version: 1, +}; +type ReapCandidate = RecordFence & { + user: string; + tenantId?: string; + conversationId: string; + conversationRef?: string; + lineageRefs?: string[]; +}; +/** GitHub compares `owner/name` without case, so records store and match the lower case form. */ +const canonicalRepository = (repository: string) => repository.toLowerCase(); +const LINK_FIELDS = [ + '_id', + 'conversationId', + 'user', + 'tenantId', + 'expiredAt', + 'subagentThread.rootConversationId', + 'subagentThread.parentConversationId', +]; + +/** + * One aggregation that reads a conversation and up to `MAX_SUBAGENT_DEPTH` ancestors, one + * `$lookup` per generation, so a liveness check is one round trip at any depth. `$graphLookup` + * would be shorter but Amazon DocumentDB rejects it. Each lookup is served by the unique + * `{ conversationId, user, tenantId }` index. Nested lookups bypass the tenant plugin, so + * ownership, tenant and retention are all judged on the result. The first match is scoped by + * the tenant plugin inside a tenant context; a caller without one, such as the reaper, names the + * record's tenant, since the same user and conversation id may exist in more than one tenant. + */ +function lineagePipeline( + { userId, conversationId }: t.PRAutomationKey, + scope?: LivenessScope, +): PipelineStage[] { + const tenant = scope != null && 'tenantId' in scope ? { tenantId: scope.tenantId ?? null } : {}; + const stages: PipelineStage[] = [ + { $match: { user: userId, conversationId, ...tenant } }, + { $limit: 1 }, + { + $project: { + _id: 0, + l0: Object.fromEntries(LINK_FIELDS.map((field) => [field, `$${field}`])), + }, + }, + ]; + for (let level = 1; level <= MAX_SUBAGENT_DEPTH; level++) { + const as = `l${level}`; + stages.push( + { + $lookup: { + from: 'conversations', + localField: `l${level - 1}.subagentThread.parentConversationId`, + foreignField: 'conversationId', + as, + }, + }, + { $unwind: { path: `$${as}`, preserveNullAndEmptyArrays: true } }, + ); + } + const keep: Record = {}; + for (let level = 0; level <= MAX_SUBAGENT_DEPTH; level++) { + for (const field of LINK_FIELDS) { + keep[`l${level}.${field}`] = 1; + } + } + stages.push({ $project: { _id: 0, ...keep } }); + return stages; +} + +const sameRefs = (a: string[], b: string[]) => + a.length === b.length && a.every((ref, index) => ref === b[index]); + +const isGone = (liveness: Liveness): liveness is Gone => + liveness === 'conversation_gone' || liveness === 'owner_gone'; + +export function createPRAutomationMethods(mongoose: typeof import('mongoose')): { + getPRAutomation: (key: t.PRAutomationKey) => Promise; + enablePRAutomation: (params: t.EnablePRAutomationParams) => Promise; + disablePRAutomation: (key: t.PRAutomationKey) => Promise<{ removed: boolean }>; + claimPRAutomationRound: ( + params: t.ClaimPRAutomationRoundParams, + ) => Promise; + settlePRAutomationRound: ( + params: t.SettlePRAutomationRoundParams, + ) => Promise; + stopPRAutomation: { + (key: t.PRAutomationKey, stopCode: 'user_stopped'): Promise; + ( + key: t.PRAutomationKey, + stopCode: t.PRAutomationEventStopCode, + fence: t.PRAutomationRunFence, + ): Promise; + }; + deletePRAutomations: ( + userId: string, + conversationIds?: string[], + conversationRefs?: string[], + ) => Promise; + reapPRAutomations: (params?: t.ReapPRAutomationsParams) => Promise; + reapPRAutomationsOf: (userId: string, conversationIds: string[]) => Promise; + setPRAutomationTrust: ( + key: t.PRAutomationKey, + trust: PRAutomationTrustLevel, + maxTrust: PRAutomationTrustLevel, + ) => Promise; + addPRAutomationBot: ( + key: t.PRAutomationKey, + bot: t.IPRAutomationBot, + maxBots: number, + fence: t.PRAutomationBotFence, + ) => Promise; + removePRAutomationBot: ( + key: t.PRAutomationKey, + botId: number, + fence: t.PRAutomationBotFence, + ) => Promise; +} { + let indexesPromise: Promise | null = null; + + /** + * The unique `{ user, tenantId, conversationId }` index is what keeps concurrent first-time enables + * from inserting two records that are then claimed independently and exceed the round cap. + * With `MONGO_AUTO_INDEX=false` the schema declaration alone never builds it, so it is + * ensured before the first write. + */ + function ensureIndexes(): Promise { + if (!indexesPromise) { + indexesPromise = createIndexesWithRetry(mongoose.models.PRAutomation).catch((error) => { + indexesPromise = null; + throw error; + }); + } + return indexesPromise; + } + + const keyFilter = ({ userId, conversationId }: t.PRAutomationKey) => ({ + user: userId, + conversationId, + }); + + /** + * Whether a round may run for this record right now: the owner exists and no account + * deletion has started (the same durable marker agent triggers use), and the conversation + * and every conversation that controls it still exist and are inside their retention. A + * subagent thread is controlled by every ancestor up to its root, and a deletion removes + * one generation before it discovers the next, so the whole parent chain is checked, read in + * one aggregation. The chain is live only when every link belongs to the owner, is inside + * its retention and the walk ends at the recorded root. The walk is bounded by + * `MAX_SUBAGENT_DEPTH`; a longer chain fails closed. + * Deletions do not fence the record; enable and claim ask this on both sides of their write + * instead. + */ + async function checkLiveness(key: t.PRAutomationKey, scope?: LivenessScope): Promise { + return (await inspectLiveness(key, scope)).liveness; + } + + /** `checkLiveness`, also naming the conversation document a live check found. */ + async function inspectLiveness( + key: t.PRAutomationKey, + scope?: LivenessScope, + ): Promise<{ liveness: Liveness } & Partial> { + const judged = (liveness: Liveness) => ({ liveness }); + if (!isValidObjectIdString(key.userId)) { + return judged('owner_inactive'); + } + const Conversation = mongoose.models.Conversation; + type Link = { + _id?: unknown; + conversationId: string; + user: string; + tenantId?: string | null; + expiredAt?: Date | null; + subagentThread?: { rootConversationId?: string; parentConversationId?: string }; + }; + const [owner, rows] = await Promise.all([ + mongoose.models.User.findById(key.userId) + .select('agentTriggerDeletionStartedAt') + .lean<{ agentTriggerDeletionStartedAt?: Date }>(), + Conversation.aggregate>(lineagePipeline(key, scope)), + ]); + if (owner == null) { + return judged('owner_gone'); + } + if (owner.agentTriggerDeletionStartedAt != null) { + return judged('owner_inactive'); + } + /** + * Links of other users or other tenants never count: a lookup by `conversationId` is scoped + * by neither. The first link came through the tenant-scoped `$match`, so it names the tenant. + */ + const now = Date.now(); + const tenant = rows[0]?.l0?.tenantId ?? null; + const links = new Map(); + for (const row of rows) { + for (const link of Object.values(row)) { + const active = link?.expiredAt == null || new Date(link.expiredAt).getTime() > now; + const owned = + link != null && String(link.user) === key.userId && (link.tenantId ?? null) === tenant; + if (link != null && owned && active) { + links.set(link.conversationId, link); + } + } + } + const conversation = links.get(key.conversationId); + if (conversation == null) { + return judged('conversation_gone'); + } + let current: Link | undefined = conversation; + const lineageRefs: string[] = []; + for (let hops = 0; current?.subagentThread?.parentConversationId != null; hops++) { + if (hops >= MAX_SUBAGENT_DEPTH) { + return judged('conversation_gone'); + } + current = links.get(current.subagentThread.parentConversationId); + if (current != null) { + lineageRefs.push(String(current._id)); + } + } + const root = conversation.subagentThread?.rootConversationId; + if (current == null || (root != null && current.conversationId !== root)) { + return judged('conversation_gone'); + } + const conversationRef = String(conversation._id); + if (scope?.conversationRef != null && scope.conversationRef !== conversationRef) { + return judged('conversation_gone'); + } + /** An ancestor deleted and saved again under the same id no longer controls the record. */ + if (scope?.lineageRefs != null && !sameRefs(scope.lineageRefs, lineageRefs)) { + return judged('conversation_gone'); + } + return { liveness: 'live', conversationRef, lineageRefs }; + } + + /** A missing owner reads as inactive to callers; the distinction only decides removal. */ + const goneCode = (liveness: Gone) => + liveness === 'owner_gone' ? 'owner_inactive' : 'conversation_gone'; + + /** + * A record whose conversation or owner is gone can never run again, so it is removed. This + * also converges a cleanup that failed after the delete committed. The delete names the + * record and version read before the liveness check, so a record renewed after the check is kept; without + * one nothing is removed and the reaper decides later. + */ + async function refuseClaim( + key: t.PRAutomationKey, + liveness: Exclude, + fence: RecordFence | null | undefined, + ): Promise { + if (!isGone(liveness)) { + return { ok: false, error: { code: liveness } }; + } + const removed = + fence == null + ? { deletedCount: 0 } + : await mongoose.models.PRAutomation.deleteOne({ + ...keyFilter(key), + _id: fence._id, + version: fence.version, + }); + return { + ok: false, + error: { code: removed.deletedCount > 0 ? goneCode(liveness) : 'not_found' }, + }; + } + + /** + * Clamps a valid level to the administrator ceiling. An unknown level is refused rather than + * narrowed, so a malformed request fails loudly instead of being stored as something else. + */ + function boundTrust( + trust: PRAutomationTrustLevel, + maxTrust: PRAutomationTrustLevel, + ): PRAutomationTrustLevel { + if (!PR_AUTOMATION_TRUST_LEVELS.includes(trust)) { + throw new RangeError('Unknown PR automation trust level'); + } + return clampPRAutomationTrust(trust, maxTrust); + } + + function assertStopCode(stopCode: PRAutomationStopCode): void { + if (!PR_AUTOMATION_STOP_CODES.includes(stopCode)) { + throw new RangeError('A stopped PR automation needs a known stop code'); + } + } + + /** Absence is a normal answer here; a query failure throws. */ + async function getPRAutomation(key: t.PRAutomationKey): Promise { + const PRAutomation = mongoose.models.PRAutomation; + return PRAutomation.findOne(keyFilter(key)).select(PROJECTION).lean(); + } + + /** + * Idempotent. An existing active record is returned unchanged. A stopped + * record is the user explicitly turning the automation back on, so it starts + * a fresh run: the round counter, the time window, the claimed heads and the + * run identity reset, which they could not do from a webhook-triggered turn. + * + * `binding` names the pull request. Each case below is one conditional write that + * carries the whole pair and the run reset together, so concurrent requests cannot + * leave a pair nobody asked for, and a stopped record is revived only by the write + * that binds it: a binding that is rejected leaves it stopped. A different repository + * clears the approved bots, which belong to one repository, and starts a fresh run. + * A different pull request in the same repository keeps the bots and starts a fresh + * run, because the round count, the time window and the claimed heads describe the + * previous pull request. + * + * Liveness is checked on both sides of the writes, like a claim. An enable for a conversation + * or owner that is already gone writes nothing. One that passed the first check but lost a + * race with the deletion and its cleanup removes the record it wrote, so a late enable cannot + * recreate a record nobody will clean up. + * + * `maxTrust` is the administrator ceiling, resolved by the caller. A requested trust level + * is clamped to it before any write, and a stored level above it is narrowed, so an enabled + * record never keeps a level the deployment now forbids. + */ + async function enablePRAutomation({ + trust: requestedTrust, + maxTrust, + binding: requested, + ...key + }: t.EnablePRAutomationParams): Promise { + const binding = + requested == null + ? undefined + : { ...requested, repository: canonicalRepository(requested.repository) }; + const trust = requestedTrust == null ? undefined : boundTrust(requestedTrust, maxTrust); + const { liveness: before, conversationRef, lineageRefs } = await inspectLiveness(key); + if (before !== 'live') { + return { ok: false, error: { code: isGone(before) ? goneCode(before) : before } }; + } + await ensureIndexes(); + const PRAutomation = mongoose.models.PRAutomation; + const filter = keyFilter(key); + /** + * The enable reads the record once and makes at most one write, fenced on the version it + * read, so a stop, a trust change or another enable that lands in between makes it report a + * retryable `conflict` instead of reviving or overwriting what that change left. The write + * marks the record `pending` with the attempt, and nothing claims a pending record. Only an + * attempt that passes its second liveness check clears the mark; one that sees an account + * deletion start puts the record back as it read it, fenced on the version it wrote. A + * preimage another rejected attempt left is still pending when it is restored, so two + * enables crossing the same deletion leave nothing claimable. + */ + const attempt = randomUUID(); + const observed = await PRAutomation.findOne(filter).lean(); + let written: number; + let recordId: unknown = observed?._id; + let restartedRun = true; + const fresh = (version: number) => ({ + ...filter, + state: 'idle', + round: 0, + trustedBots: [], + claimedHeads: [], + epoch: attempt, + pending: attempt, + version, + conversationRef, + lineageRefs, + trust: clampPRAutomationTrust(trust, maxTrust), + ...binding, + }); + /** A record of an earlier conversation saved under the same id starts over for this one. */ + const replaced = + (observed?.conversationRef != null && observed.conversationRef !== conversationRef) || + (observed?.lineageRefs != null && !sameRefs(observed.lineageRefs, lineageRefs ?? [])); + if (observed == null) { + written = 1; + try { + recordId = (await PRAutomation.create(fresh(written)))._id; + } catch (error) { + if ((error as { code?: number }).code === DUPLICATE_KEY) { + return { ok: false, error: { code: 'conflict' } }; + } + throw error; + } + } else if (replaced) { + written = observed.version + 1; + const applied = await PRAutomation.replaceOne( + { ...filter, _id: observed._id, version: observed.version }, + { + ...fresh(written), + ...(observed.tenantId != null && { tenantId: observed.tenantId }), + }, + { runValidators: true }, + ); + if (applied.matchedCount === 0) { + return { ok: false, error: { code: 'conflict' } }; + } + } else { + written = observed.version + 1; + const restart = enableRestart(observed, attempt, binding, trust ?? observed.trust, maxTrust); + restartedRun = restart != null; + /** An enable also narrows a stored level the current ceiling no longer allows. */ + const narrowed = clampPRAutomationTrust(observed.trust, maxTrust); + const update: { $set?: Record; $unset?: Record } = + restart ?? (narrowed === observed.trust ? {} : { $set: { trust: narrowed } }); + const applied = await PRAutomation.updateOne( + { ...filter, _id: observed._id, version: observed.version }, + { ...update, $set: { ...update.$set, conversationRef, lineageRefs }, ...BUMP }, + { runValidators: true }, + ); + if (applied.matchedCount === 0) { + return { ok: false, error: { code: 'conflict' } }; + } + } + const after = await checkLiveness(key, { conversationRef, lineageRefs }); + if (isGone(after)) { + await PRAutomation.deleteOne({ ...filter, _id: recordId, version: written }); + return { ok: false, error: { code: goneCode(after) } }; + } + if (after !== 'live') { + if (restartedRun) { + await undoEnable(filter, { _id: recordId, version: written }, observed); + } + return { ok: false, error: { code: after } }; + } + if (restartedRun) { + const committed = await PRAutomation.updateOne( + { ...filter, _id: recordId, version: written, pending: attempt }, + { $unset: { pending: '' }, ...BUMP }, + ); + if (committed.matchedCount === 0) { + return { ok: false, error: { code: 'conflict' } }; + } + } + const record = await getPRAutomation(key); + if (record == null) { + throw new Error('PR automation record missing after enable'); + } + return { ok: true, value: record }; + } + + /** + * The run-starting write an enable makes on an existing record, or `null` when the record is + * already running on the requested binding. A stopped record, a different binding, or a record + * another attempt left pending starts a fresh run under the attempt's epoch. A different + * repository also clears the approved bots, which belong to one repository. The run starts + * with the requested trust, or the stored one, clamped to the current ceiling. + */ + function enableRestart( + current: t.IPRAutomation, + attempt: string, + binding: t.PRAutomationBinding | undefined, + trust: PRAutomationTrustLevel, + maxTrust: PRAutomationTrustLevel, + ): ReturnType | null { + const otherRepository = binding != null && current.repository !== binding.repository; + const restart = + current.state === 'stopped' || + current.pending != null || + otherRepository || + (binding != null && current.pullNumber !== binding.pullNumber); + if (!restart) { + return null; + } + return runReset(attempt, { + trust: clampPRAutomationTrust(trust, maxTrust), + ...binding, + ...(otherRepository && { trustedBots: [] }), + pending: attempt, + }); + } + + /** + * Puts the record back as the enable read it when an account deletion started during the + * write, which can still be cancelled. It applies only while the record still carries the + * version the enable wrote, so anything that changed it since, a stop, a trust change or a + * later enable, is kept, and the record then stays pending, which nothing claims. + */ + async function undoEnable( + key: ReturnType, + written: RecordFence, + previous: (t.IPRAutomation & { _id: unknown }) | null, + ): Promise { + const PRAutomation = mongoose.models.PRAutomation; + const guard = { ...key, _id: written._id, version: written.version }; + if (previous == null) { + await PRAutomation.deleteOne(guard); + return; + } + await PRAutomation.replaceOne(guard, { ...previous, version: written.version + 1 }); + } + + async function disablePRAutomation(key: t.PRAutomationKey): Promise<{ removed: boolean }> { + const PRAutomation = mongoose.models.PRAutomation; + const result = await PRAutomation.deleteOne(keyFilter(key)); + return { removed: result.deletedCount > 0 }; + } + + /** + * Starts one fix round, atomically. The round count, the time window, the + * state and the head are all conditions of a single `findOneAndUpdate`, so two + * deliveries racing for the same record cannot both pass the cap. A head that + * any earlier round already claimed is rejected, so a delayed delivery of a + * superseded commit cannot spend the budget. The persisted counter is what + * stops a new turn from resetting the cap. Each claim gets a `runId`, which + * the round must present to settle. The delivery names the pull request it is + * for, and the claim applies only while the record is still bound to it, so an + * event that arrives after a rebind cannot spend the new pull request's budget. + * + * A force-push back to an earlier SHA is rejected the same way; the user + * restarts the automation to work on it. + */ + async function claimPRAutomationRound({ + binding, + maxRounds, + maxMinutes, + headSha: rawHeadSha, + now = new Date(), + ...key + }: t.ClaimPRAutomationRoundParams): Promise { + if (!COMMIT_SHA.test(rawHeadSha)) { + throw new RangeError('headSha must be a full commit SHA'); + } + /** One Git object, one identity: hex case must not let a commit spend two rounds. */ + const headSha = rawHeadSha.toLowerCase(); + const PRAutomation = mongoose.models.PRAutomation; + const repository = canonicalRepository(binding.repository); + const filter = { + ...keyFilter(key), + repository, + pullNumber: binding.pullNumber, + epoch: binding.epoch, + }; + const cutoff = new Date(now.getTime() - maxMinutes * 60_000); + const runId = randomUUID(); + + /** + * Conversations also leave through the retention index, which never runs + * `deleteConvos`, so liveness is checked here rather than trusted to a cascade. + */ + const seen = await PRAutomation.findOne(keyFilter(key)) + .select({ _id: 1, version: 1, conversationRef: 1, lineageRefs: 1 }) + .lean(); + const before = await checkLiveness(key, { + conversationRef: seen?.conversationRef, + lineageRefs: seen?.lineageRefs, + }); + if (before !== 'live') { + return refuseClaim(key, before, seen); + } + + /** + * The window opens with the first round that is actually claimed. `$min` sets `startedAt` + * only when it is missing and keeps the earlier time otherwise, so the window is part of the + * same atomic write: a claim that fails or is abandoned leaves nothing behind to count + * against `maxMinutes`. + */ + const claimed = await PRAutomation.findOneAndUpdate( + { + ...filter, + state: { $in: CLAIMABLE_STATES }, + pending: { $exists: false }, + round: { $lt: maxRounds }, + $or: [{ startedAt: { $exists: false } }, { startedAt: { $gte: cutoff } }], + claimedHeads: { $ne: headSha }, + }, + { + $min: { startedAt: now }, + $inc: { round: 1, version: 1 }, + $set: { state: 'fixing', lastHeadSha: headSha, runId }, + $push: { claimedHeads: headSha }, + $unset: { stopCode: '' }, + }, + { new: true, select: PROJECTION }, + ).lean(); + if (claimed != null) { + /** + * A deletion that started between the first check and the write is seen here. The + * whole claim is undone under its own run id, so it never starts work and a cancelled + * deletion leaves the record as it was: the round, the window and the head are not + * spent. Rounds and heads are claimed one for one and the window opens with the first + * round, so the state before the claim follows from the claimed record. A record whose + * conversation or owner is gone is removed. + */ + let after: Liveness; + try { + after = await checkLiveness(key, { + conversationRef: claimed.conversationRef, + lineageRefs: claimed.lineageRefs, + }); + } catch (error) { + /** A check that fails leaves no answer to act on, so the round is handed back. */ + await releaseClaim(key, claimed, runId).catch(() => undefined); + throw error; + } + if (after === 'live') { + return { ok: true, value: { ...claimed, runId } }; + } + const released = await releaseClaim(key, claimed, runId); + /** Only a claim that undid itself knows the version it left; otherwise the reaper decides. */ + return refuseClaim(key, after, released); + } + + const current = await getPRAutomation(key); + if (current == null) { + return { ok: false, error: { code: 'not_found' } }; + } + if ( + current.repository !== repository || + current.pullNumber !== binding.pullNumber || + current.epoch !== binding.epoch + ) { + return { ok: false, error: { code: 'binding_mismatch' } }; + } + if (!CLAIMABLE_STATES.includes(current.state)) { + return { ok: false, error: { code: 'not_active' } }; + } + if (current.round >= maxRounds) { + return { ok: false, error: { code: 'round_cap' } }; + } + if (current.startedAt != null && current.startedAt < cutoff) { + return { ok: false, error: { code: 'time_cap' } }; + } + if (current.pending != null) { + /** An enable is still deciding; it either commits this run or replaces it. */ + return { ok: false, error: { code: 'conflict' } }; + } + if (current.claimedHeads.includes(headSha)) { + return { ok: false, error: { code: 'stale_head' } }; + } + /** Every condition passes now, so the record changed between the write and this read. */ + return { ok: false, error: { code: 'conflict' } }; + } + + /** + * Undoes a claim that will not run. It is fenced on the claim's own run identity rather than + * the version: a trust or bot change made meanwhile touches nothing the claim owns, and must + * not strand the round. Rounds and heads are claimed one for one and the window opens with the + * first round, so the state before the claim follows from the claimed record. + */ + async function releaseClaim( + key: t.PRAutomationKey, + claimed: t.IPRAutomation, + runId: string, + ): Promise { + const previousHeads = claimed.claimedHeads.slice(0, -1); + const lastHeadSha = previousHeads[previousHeads.length - 1]; + return mongoose.models.PRAutomation.findOneAndUpdate( + { ...keyFilter(key), runId, state: 'fixing', round: claimed.round }, + { + ...BUMP, + $set: { + state: lastHeadSha == null ? 'idle' : 'waiting', + round: claimed.round - 1, + claimedHeads: previousHeads, + ...(lastHeadSha != null && { lastHeadSha }), + }, + $unset: { + runId: '', + ...(lastHeadSha == null && { startedAt: '', lastHeadSha: '' }), + }, + }, + { new: true, projection: { _id: 1, version: 1 } }, + ).lean(); + } + + /** + * Reports where the round that owns the record ended up. The transition only + * applies while that round is still the current one, identified by its + * `runId`, so a completion that is retried after a later round was claimed, or + * after the run was stopped, restarted or rebound, cannot move the newer run + * to `waiting` and open the way for an overlapping claim. Returns `null` when + * nothing changed: no record, a stopped record, or a round that is not the + * current one. + */ + async function settlePRAutomationRound({ + round, + runId, + state, + ...key + }: t.SettlePRAutomationRoundParams): Promise { + if (!SETTLED_STATES.includes(state)) { + throw new RangeError( + 'A round settles as waiting or needs_user; use stopPRAutomation to stop', + ); + } + const PRAutomation = mongoose.models.PRAutomation; + return PRAutomation.findOneAndUpdate( + { ...keyFilter(key), state: { $in: SETTLEABLE_STATES }, round, runId }, + { $set: { state }, ...BUMP }, + { new: true, select: PROJECTION }, + ).lean(); + } + + /** + * Stops the automation whatever round is running. A stop code is required: the client + * maps it to the reason it shows. The first stop wins, and a stopped record leaves that + * state only through `enablePRAutomation`. + * + * A user stop is unconditional, so it is never lost to a race; any fence passed with it is + * ignored. Every other code comes from an event about one run, and a delayed one must not + * stop the run that replaced it, so it names the pull request and the epoch it saw and + * applies only while both are current. A restart on the same pull request starts a new + * epoch, so a stop meant for the previous run cannot end the new one. + */ + async function stopPRAutomation( + key: t.PRAutomationKey, + stopCode: PRAutomationStopCode, + fence?: t.PRAutomationRunFence, + ): Promise { + assertStopCode(stopCode); + const user = stopCode === 'user_stopped'; + if (!user && (fence?.repository == null || fence.pullNumber == null || !fence.epoch)) { + throw new RangeError('An event-driven stop must name the pull request and run it is for'); + } + const condition = + user || fence == null + ? {} + : { + repository: canonicalRepository(fence.repository), + pullNumber: fence.pullNumber, + epoch: fence.epoch, + }; + const PRAutomation = mongoose.models.PRAutomation; + return PRAutomation.findOneAndUpdate( + { ...keyFilter(key), ...condition, state: { $ne: 'stopped' } }, + { $set: { state: 'stopped', stopCode }, ...BUMP }, + { new: true, select: PROJECTION }, + ).lean(); + } + + /** + * Removes every record of a user, or of the listed conversations, in bounded batches. + * Cleanup only: it creates nothing, so a deployment that never used the feature pays one + * indexed delete that matches no document. `conversationRefs`, aligned with + * `conversationIds`, names the conversation documents that were deleted, so a record enabled + * for a conversation saved again under the same id meanwhile is kept. + */ + async function deletePRAutomations( + userId: string, + conversationIds?: string[], + conversationRefs?: string[], + ): Promise { + const PRAutomation = mongoose.models.PRAutomation; + if (PRAutomation == null || conversationIds?.length === 0) { + return; + } + if (conversationIds == null) { + await PRAutomation.deleteMany({ user: userId }); + return; + } + for (let start = 0; start < conversationIds.length; start += DELETE_BATCH) { + const refs = conversationRefs?.slice(start, start + DELETE_BATCH); + await PRAutomation.deleteMany({ + user: userId, + conversationId: { $in: conversationIds.slice(start, start + DELETE_BATCH) }, + ...(refs != null && { + $or: [{ conversationRef: { $in: refs } }, { conversationRef: { $exists: false } }], + }), + }); + } + } + + /** + * Removes records whose conversation, an ancestor of it, or owner is gone, one page at a time + * in `_id` order. The retention index deletes conversations without running `deleteConvos`, + * so an idle record would otherwise outlive its conversation; this is the periodic pass that + * catches it. A deletion in progress keeps its records, as it does for a claim. + */ + async function reapPRAutomations({ + after, + limit = 100, + }: t.ReapPRAutomationsParams = {}): Promise { + const page = await mongoose.models.PRAutomation.find( + after == null ? {} : { _id: { $gt: after } }, + ) + .sort({ _id: 1 }) + .limit(limit) + .select(REAP_FIELDS) + .lean(); + const removed = await reapCandidates(page); + const last = page.length === limit ? page[page.length - 1] : undefined; + return { + scanned: page.length, + removed, + ...(last == null ? {} : { next: String(last._id) }), + }; + } + + /** + * Reaps the records of the listed conversations only, for a cleanup that knows the ids but + * not the conversation documents they belonged to: a record whose conversation was saved again + * under the same id belongs to the new one, is live, and is kept. + */ + async function reapPRAutomationsOf(userId: string, conversationIds: string[]): Promise { + let removed = 0; + for (let start = 0; start < conversationIds.length; start += DELETE_BATCH) { + const records = await mongoose.models.PRAutomation.find({ + user: userId, + conversationId: { $in: conversationIds.slice(start, start + DELETE_BATCH) }, + }) + .select(REAP_FIELDS) + .lean(); + removed += await reapCandidates(records); + } + return removed; + } + + /** Checks a few records at a time, so a large batch neither runs serially nor floods the database. */ + async function reapCandidates(records: ReapCandidate[]): Promise { + let removed = 0; + for (let start = 0; start < records.length; start += REAP_CONCURRENCY) { + const counts = await Promise.all( + records.slice(start, start + REAP_CONCURRENCY).map(reapCandidate), + ); + for (const count of counts) { + removed += count; + } + } + return removed; + } + + async function reapCandidate(record: ReapCandidate): Promise { + const key = { userId: String(record.user), conversationId: record.conversationId }; + const scope = { + tenantId: record.tenantId, + conversationRef: record.conversationRef, + lineageRefs: record.lineageRefs, + }; + if (!isGone(await checkLiveness(key, scope))) { + return 0; + } + /** A record written after the check, a renewed enable included, has a new version. */ + const result = await mongoose.models.PRAutomation.deleteOne({ + _id: record._id, + ...keyFilter(key), + version: record.version, + }); + return result.deletedCount; + } + + /** `maxTrust` is the administrator ceiling; a wider request is stored as the ceiling. */ + async function setPRAutomationTrust( + key: t.PRAutomationKey, + trust: PRAutomationTrustLevel, + maxTrust: PRAutomationTrustLevel, + ): Promise { + const PRAutomation = mongoose.models.PRAutomation; + return PRAutomation.findOneAndUpdate( + keyFilter(key), + { $set: { trust: boundTrust(trust, maxTrust) }, ...BUMP }, + { new: true, select: PROJECTION, runValidators: true }, + ).lean(); + } + + /** + * Idempotent by numeric id. The limit is the configured `maxBots`, resolved by + * the caller, and cannot exceed the schema ceiling. A login is stored for + * display only. The approval names the repository and the epoch it was authorized + * against and applies only while both are current, so an approval that overlaps a + * rebind, including one away and back to the same repository that cleared the list, + * cannot add a bot nobody approved for the current binding. + */ + async function addPRAutomationBot( + key: t.PRAutomationKey, + bot: t.IPRAutomationBot, + maxBots: number, + fence: t.PRAutomationBotFence, + ): Promise { + const repository = canonicalRepository(fence.repository); + const { epoch } = fence; + if (!Number.isInteger(maxBots) || maxBots < 1 || maxBots > MAX_PR_AUTOMATION_BOTS) { + throw new RangeError(`maxBots must be an integer from 1 to ${MAX_PR_AUTOMATION_BOTS}`); + } + const PRAutomation = mongoose.models.PRAutomation; + const updated = await PRAutomation.findOneAndUpdate( + { + ...keyFilter(key), + repository, + epoch, + 'trustedBots.id': { $ne: bot.id }, + [`trustedBots.${maxBots - 1}`]: { $exists: false }, + }, + { $push: { trustedBots: bot }, ...BUMP }, + { new: true, select: PROJECTION, runValidators: true }, + ).lean(); + if (updated != null) { + return { ok: true, value: updated }; + } + + const current = await getPRAutomation(key); + if (current == null) { + return { ok: false, error: { code: 'not_found' } }; + } + if (current.repository !== repository || current.epoch !== epoch) { + return { ok: false, error: { code: 'binding_mismatch' } }; + } + if (current.trustedBots.some((existing) => existing.id === bot.id)) { + return { ok: true, value: current }; + } + return { ok: false, error: { code: 'bot_limit' } }; + } + + /** + * Names the repository and epoch like an approval does, so a delayed removal cannot delete + * an approval made for a later binding, including one away and back to the same repository. + */ + async function removePRAutomationBot( + key: t.PRAutomationKey, + botId: number, + { repository, epoch }: t.PRAutomationBotFence, + ): Promise { + const PRAutomation = mongoose.models.PRAutomation; + return PRAutomation.findOneAndUpdate( + { ...keyFilter(key), repository: canonicalRepository(repository), epoch }, + { $pull: { trustedBots: { id: botId } }, ...BUMP }, + { new: true, select: PROJECTION }, + ).lean(); + } + + return { + getPRAutomation, + enablePRAutomation, + disablePRAutomation, + claimPRAutomationRound, + settlePRAutomationRound, + stopPRAutomation, + deletePRAutomations, + reapPRAutomations, + reapPRAutomationsOf, + setPRAutomationTrust, + addPRAutomationBot, + removePRAutomationBot, + }; +} + +export type PRAutomationMethods = ReturnType; diff --git a/packages/data-schemas/src/methods/user.methods.spec.ts b/packages/data-schemas/src/methods/user.methods.spec.ts index d9318173d1b..f77b27b94ce 100644 --- a/packages/data-schemas/src/methods/user.methods.spec.ts +++ b/packages/data-schemas/src/methods/user.methods.spec.ts @@ -8,6 +8,9 @@ import { import type * as t from '~/types'; import { createToolApprovalGrantModel } from '~/models/toolApprovalGrant'; import { createUserMethods, USER_DELETION_FENCE_STALE_MS } from './user'; +import { runAsSystem, tenantStorage } from '~/config/tenantContext'; +import { createPRAutomationModel } from '~/models/prAutomation'; +import { createConversationModel } from '~/models/convo'; import balanceSchema from '~/schema/balance'; import userSchema from '~/schema/user'; @@ -51,6 +54,8 @@ beforeAll(async () => { /** Initialize methods */ methods = createUserMethods(mongoose); createToolApprovalGrantModel(mongoose); + createPRAutomationModel(mongoose); + createConversationModel(mongoose); }); afterAll(async () => { @@ -2199,6 +2204,135 @@ describe('User Methods - Database Tests', () => { }); }); +describe('PR automation cleanup on account deletion', () => { + it('removes the records in every tenant when the deletion runs in one', async () => { + const user = await User.create({ email: 'delete-pr-tenants@example.com', provider: 'local' }); + const id = user._id.toString(); + const PRAutomation = mongoose.models.PRAutomation; + const inTenant = (tenantId: string, run: () => Promise) => + tenantStorage.run({ tenantId }, run); + await inTenant('tenant-a', () => PRAutomation.create({ user: id, conversationId: 'chat-a' })); + await inTenant('tenant-b', () => PRAutomation.create({ user: id, conversationId: 'chat-a' })); + await inTenant('tenant-a', () => methods.deleteUserById(id)); + expect(await runAsSystem(() => PRAutomation.countDocuments({ user: id }))).toBe(0); + }); + + it('keeps the records exactly as they were when the account delete fails', async () => { + const user = await User.create({ email: 'delete-pr-fail@example.com', provider: 'local' }); + const id = user._id.toString(); + const PRAutomation = mongoose.models.PRAutomation; + await PRAutomation.create({ user: id, conversationId: 'chat-a', repository: 'acme/one' }); + const deletion = jest.spyOn(User, 'deleteOne').mockRejectedValueOnce(new Error('db down')); + try { + await expect(methods.deleteUserById(id)).rejects.toThrow('db down'); + expect(await User.exists({ _id: user._id })).not.toBeNull(); + expect(await PRAutomation.findOne({ user: id }).lean()).toMatchObject({ + state: 'idle', + repository: 'acme/one', + }); + } finally { + deletion.mockRestore(); + } + }); + + it('invalidates the auth cache before a cleanup that fails', async () => { + enableAuthUserDocCache(); + const user = await User.create({ email: 'delete-pr-cache@example.com', provider: 'local' }); + const id = user._id.toString(); + const cache = { + get: jest.fn().mockResolvedValue(['auth-cache-key-a']), + set: jest.fn().mockResolvedValue(true), + delete: jest.fn().mockResolvedValue(true), + }; + const methodsWithCache = createUserMethods(mongoose, { getCache: () => cache }); + const cleanup = jest + .spyOn(mongoose.models.PRAutomation, 'deleteMany') + .mockRejectedValueOnce(new Error('synthetic cleanup failure')); + try { + await methodsWithCache.deleteUserById(id); + expect(cache.delete).toHaveBeenCalledWith(`${AUTH_USER_DOC_BY_ID_PREFIX}:${id}`); + } finally { + cleanup.mockRestore(); + } + }); + + it('reports the account deleted when only the cleanup after the delete fails', async () => { + const user = await User.create({ email: 'delete-pr-committed@example.com', provider: 'local' }); + const id = user._id.toString(); + const PRAutomation = mongoose.models.PRAutomation; + await PRAutomation.create({ user: id, conversationId: 'chat-a' }); + const cleanup = jest + .spyOn(PRAutomation, 'deleteMany') + .mockRejectedValueOnce(new Error('synthetic cleanup failure')); + try { + await expect(methods.deleteUserById(id)).resolves.toMatchObject({ deletedCount: 1 }); + } finally { + cleanup.mockRestore(); + } + }); + + it('converges on a retry when the cleanup fails after the account was deleted', async () => { + const user = await User.create({ email: 'delete-pr-retry@example.com', provider: 'local' }); + const id = user._id.toString(); + const PRAutomation = mongoose.models.PRAutomation; + await PRAutomation.create({ user: id, conversationId: 'chat-a' }); + const cleanup = jest + .spyOn(PRAutomation, 'deleteMany') + .mockRejectedValueOnce(new Error('synthetic cleanup failure')); + try { + await methods.deleteUserById(id); + expect(await User.exists({ _id: user._id })).toBeNull(); + expect(await PRAutomation.countDocuments({ user: id })).toBe(1); + await expect(methods.deleteUserById(id)).resolves.toMatchObject({ deletedCount: 0 }); + expect(await PRAutomation.countDocuments({ user: id })).toBe(0); + } finally { + cleanup.mockRestore(); + } + }); + + it('removes the records after the account delete committed, not before', async () => { + const user = await User.create({ email: 'delete-pr-order@example.com', provider: 'local' }); + const id = user._id.toString(); + const PRAutomation = mongoose.models.PRAutomation; + await PRAutomation.create({ user: id, conversationId: 'chat-a' }); + const cleanup = jest.spyOn(PRAutomation, 'deleteMany'); + const deletion = jest.spyOn(User, 'deleteOne'); + try { + await methods.deleteUserById(id); + expect(deletion.mock.invocationCallOrder[0]).toBeLessThan( + cleanup.mock.invocationCallOrder[0], + ); + } finally { + cleanup.mockRestore(); + deletion.mockRestore(); + } + }); + + it('removes the records of an account that no longer exists', async () => { + const id = new mongoose.Types.ObjectId().toString(); + const PRAutomation = mongoose.models.PRAutomation; + await PRAutomation.create({ user: id, conversationId: 'chat-a' }); + + await expect(methods.deleteUserById(id)).resolves.toMatchObject({ deletedCount: 0 }); + + expect(await PRAutomation.countDocuments({ user: id })).toBe(0); + }); + + it('removes the deleted user records and leaves other users alone', async () => { + const user = await User.create({ email: 'delete-pr@example.com', provider: 'local' }); + const id = user._id.toString(); + const PRAutomation = mongoose.models.PRAutomation; + await PRAutomation.create({ user: id, conversationId: 'chat-a' }); + await PRAutomation.create({ user: id, conversationId: 'chat-b' }); + await PRAutomation.create({ user: 'another-user', conversationId: 'chat-a' }); + + await expect(methods.deleteUserById(id)).resolves.toMatchObject({ deletedCount: 1 }); + + expect(await PRAutomation.countDocuments({ user: id })).toBe(0); + expect(await PRAutomation.countDocuments({ user: 'another-user' })).toBe(1); + }); +}); + describe('personal grant cleanup before account deletion', () => { async function owner() { const user = await User.create({ email: 'delete-grants@example.com', provider: 'local' }); diff --git a/packages/data-schemas/src/methods/user.ts b/packages/data-schemas/src/methods/user.ts index cc950283441..64bdc2a4528 100644 --- a/packages/data-schemas/src/methods/user.ts +++ b/packages/data-schemas/src/methods/user.ts @@ -16,6 +16,8 @@ import type { } from '~/types'; import type { TwoFactorEnrollmentGuard, TwoFactorEnrollmentUpdate } from '~/types'; import type { CacheStore } from '~/types'; +import { createPRAutomationMethods } from './prAutomation'; +import { runAsSystem } from '~/config/tenantContext'; import { evictAuthUserDocs } from '~/utils/eviction'; import { escapeRegExp } from '~/utils/string'; import { signPayload } from '~/crypto'; @@ -221,6 +223,8 @@ export function createUserMethods( return normalized; } + const prAutomation = createPRAutomationMethods(mongoose); + /** * Search for a single user based on partial data and return matching user document as plain object. * Email fields in searchCriteria are automatically normalized to lowercase for case-insensitive matching. @@ -644,10 +648,26 @@ export function createUserMethods( const User = mongoose.models.User; await mongoose.models.ToolApprovalGrant?.deleteMany({ user: userId }); const result = await User.deleteOne({ _id: userId }); + if (result.deletedCount > 0) { + /** Before any fallible cleanup, so a stale cached `req.user` cannot outlive the account. */ + await invalidateAuthUserDocCache(userId); + } + /** Removed only after the account delete committed, so a failed delete keeps the user's + * automation. Once the account is gone a cleanup failure must not report the deletion as + * failed: the next claim removes any record whose owner no longer exists. A retry on an + * account that was already gone has committed nothing, so its failure still throws. */ + try { + /** The account is gone in every tenant, so its automation is removed from all of them. */ + await runAsSystem(() => prAutomation.deletePRAutomations(userId)); + } catch (error) { + if (result.deletedCount === 0) { + throw error; + } + logger.warn('[deleteUserById] PR automation cleanup deferred to its next claim', error); + } if (result.deletedCount === 0) { return { deletedCount: 0, message: 'No user found with that ID.' }; } - await invalidateAuthUserDocCache(userId); return { deletedCount: result.deletedCount, message: 'User was deleted successfully.' }; } catch (error: unknown) { const errorMessage = error instanceof Error ? error.message : 'Unknown error'; diff --git a/packages/data-schemas/src/models/index.ts b/packages/data-schemas/src/models/index.ts index 1d999c870ff..e65ee8c97c9 100644 --- a/packages/data-schemas/src/models/index.ts +++ b/packages/data-schemas/src/models/index.ts @@ -12,6 +12,7 @@ import { createSkillSyncStatusModel } from './skillSyncStatus'; import { createConversationTagModel } from './conversationTag'; import { createCodeEnvironmentModel } from './codeEnvironment'; import { createAgentCategoryModel } from './agentCategory'; +import { createPRAutomationModel } from './prAutomation'; import { createResourceLockModel } from './resourceLock'; import { createChatProjectModel } from './chatProject'; import { createAgentApiKeyModel } from './agentApiKey'; @@ -87,6 +88,7 @@ export function createModels(mongoose: typeof import('mongoose')): { ToolCall: ReturnType; MemoryEntry: ReturnType; ToolFavorite: ReturnType; + PRAutomation: ReturnType; AccessRole: ReturnType; AclEntry: ReturnType; SystemGrant: ReturnType; @@ -139,6 +141,7 @@ export function createModels(mongoose: typeof import('mongoose')): { ToolCall: createToolCallModel(mongoose), MemoryEntry: createMemoryModel(mongoose), ToolFavorite: createToolFavoriteModel(mongoose), + PRAutomation: createPRAutomationModel(mongoose), AccessRole: createAccessRoleModel(mongoose), AclEntry: createAclEntryModel(mongoose), SystemGrant: createSystemGrantModel(mongoose), diff --git a/packages/data-schemas/src/models/prAutomation.ts b/packages/data-schemas/src/models/prAutomation.ts new file mode 100644 index 00000000000..3c13fcdc8c8 --- /dev/null +++ b/packages/data-schemas/src/models/prAutomation.ts @@ -0,0 +1,14 @@ +import { Model } from 'mongoose'; +import type { IPRAutomationDocument } from '~/types/prAutomation'; +import { applyTenantIsolation } from '~/models/plugins/tenantIsolation'; +import prAutomationSchema from '~/schema/prAutomation'; + +export function createPRAutomationModel( + mongoose: typeof import('mongoose'), +): Model { + applyTenantIsolation(prAutomationSchema); + return ( + mongoose.models.PRAutomation || + mongoose.model('PRAutomation', prAutomationSchema) + ); +} diff --git a/packages/data-schemas/src/schema/prAutomation.ts b/packages/data-schemas/src/schema/prAutomation.ts new file mode 100644 index 00000000000..cbf2c704bc6 --- /dev/null +++ b/packages/data-schemas/src/schema/prAutomation.ts @@ -0,0 +1,94 @@ +import { Schema } from 'mongoose'; +import { randomUUID } from 'node:crypto'; +import { PR_AUTOMATION_STATES } from 'librechat-data-provider'; +import { MAX_PR_AUTOMATION_BOTS } from 'librechat-data-provider'; +import { PR_AUTOMATION_STOP_CODES } from 'librechat-data-provider'; +import { PR_AUTOMATION_TRUST_LEVELS } from 'librechat-data-provider'; +import type { IPRAutomationDocument } from '~/types/prAutomation'; + +const trustedBotSchema = new Schema( + { + id: { + type: Number, + required: true, + min: 1, + validate: { + validator: Number.isSafeInteger, + message: 'A bot account id is a positive whole number', + }, + }, + login: { type: String, maxlength: 128 }, + }, + { _id: false }, +); + +const prAutomationSchema: Schema = new Schema( + { + /** Served by the unique `{ user, tenantId, conversationId }` index below; a second index only adds + * write cost. */ + user: { + type: String, + required: true, + }, + tenantId: { + type: String, + index: true, + }, + /** Unbounded, like the conversation schema it mirrors, so every stored conversation can enable. */ + conversationId: { type: String, required: true }, + /** The `_id` of the conversation document the record was enabled for. */ + conversationRef: { type: String, maxlength: 64 }, + /** The `_id` of each ancestor of a subagent thread, parent first and root last. */ + lineageRefs: { type: [{ type: String, maxlength: 64 }], default: undefined }, + repository: { type: String, maxlength: 256 }, + pullNumber: { + type: Number, + min: 1, + validate: { + validator: Number.isInteger, + message: 'A pull request number is a whole number', + }, + }, + state: { + type: String, + enum: PR_AUTOMATION_STATES, + default: 'idle', + required: true, + }, + stopCode: { type: String, enum: PR_AUTOMATION_STOP_CODES }, + trust: { + type: String, + enum: PR_AUTOMATION_TRUST_LEVELS, + default: 'approvedBots', + required: true, + }, + trustedBots: { + type: [trustedBotSchema], + default: [], + validate: { + validator: (bots: unknown[]) => bots.length <= MAX_PR_AUTOMATION_BOTS, + message: `At most ${MAX_PR_AUTOMATION_BOTS} approved bots`, + }, + }, + round: { type: Number, default: 0, min: 0, required: true }, + startedAt: { type: Date }, + lastHeadSha: { type: String, maxlength: 64 }, + claimedHeads: { type: [{ type: String, maxlength: 64 }], default: [] }, + runId: { type: String, maxlength: 64 }, + epoch: { type: String, maxlength: 64, required: true, default: () => randomUUID() }, + pending: { type: String, maxlength: 64 }, + /** Moves on every write, so a write based on an earlier read can be fenced on it. */ + version: { type: Number, default: 0, min: 0, required: true }, + }, + { timestamps: true }, +); + +/** + * One automation per conversation, keyed like the conversation it belongs to: + * `{ conversationId, user, tenantId }` is unique there, so the same user and + * conversation id can exist in two tenants. The unique index makes enabling + * idempotent and race-free within a tenant. + */ +prAutomationSchema.index({ user: 1, tenantId: 1, conversationId: 1 }, { unique: true }); + +export default prAutomationSchema; diff --git a/packages/data-schemas/src/types/index.ts b/packages/data-schemas/src/types/index.ts index eba6570762e..a918aa6f5a7 100644 --- a/packages/data-schemas/src/types/index.ts +++ b/packages/data-schemas/src/types/index.ts @@ -31,6 +31,7 @@ export * from './pluginAuth'; /* Memories */ export * from './memory'; export * from './favorite'; +export * from './prAutomation'; /* Prompts */ export * from './prompts'; /* Skills */ diff --git a/packages/data-schemas/src/types/prAutomation.ts b/packages/data-schemas/src/types/prAutomation.ts new file mode 100644 index 00000000000..2f24f86393e --- /dev/null +++ b/packages/data-schemas/src/types/prAutomation.ts @@ -0,0 +1,159 @@ +import type { PRAutomationTrustLevel } from 'librechat-data-provider'; +import type { PRAutomationStopCode } from 'librechat-data-provider'; +import type { PRAutomationState } from 'librechat-data-provider'; +import type { Document } from 'mongoose'; + +/** A bot the user approved for one repository, matched by its numeric account id. */ +export interface IPRAutomationBot { + id: number; + /** Display only. Logins can be renamed, so identity checks never read it. */ + login?: string; +} + +export interface IPRAutomation { + user: string; + tenantId?: string; + conversationId: string; + /** + * The `_id` of the conversation document the record was enabled for. A conversation deleted + * and saved again under the same id is a different document, and the record does not carry over. + */ + conversationRef?: string; + /** The same for each ancestor of a subagent thread, parent first and root last. */ + lineageRefs?: string[]; + /** `owner/name` of the repository the pull request belongs to. */ + repository?: string; + pullNumber?: number; + state: PRAutomationState; + stopCode?: PRAutomationStopCode; + trust: PRAutomationTrustLevel; + /** Approved bots for this repository only. */ + trustedBots: IPRAutomationBot[]; + /** Fix rounds started so far. Persisted so a new turn cannot reset it. */ + round: number; + /** First automated round, the start of the wall-clock window. */ + startedAt?: Date; + /** The head the latest round works on. */ + lastHeadSha?: string; + /** + * Identifies the round that is running. Only that round may settle the record, so a + * completion from a run that was stopped, restarted or rebound to another pull request + * cannot move the run that replaced it. + */ + runId?: string; + /** + * Identifies the run: a new value on every enable that starts one, every restart and every + * rebind. An event-driven stop names it, so a stop meant for an earlier run on the same pull + * request cannot end the current one. + */ + epoch: string; + /** + * Set to the enable attempt that wrote the record until that attempt confirms the owner and + * conversation are still live. Nothing claims a pending record. + */ + pending?: string; + /** + * Every head a round has been claimed for since the last restart. A delayed + * delivery of an earlier head is rejected against this list. It holds at most + * one entry per round, so the round cap bounds it. + */ + claimedHeads: string[]; + /** Moves on every write; a write based on an earlier read is fenced on it. */ + version: number; + createdAt?: Date; + updatedAt?: Date; +} + +export interface IPRAutomationDocument extends IPRAutomation, Document {} + +export interface PRAutomationKey { + userId: string; + conversationId: string; +} + +/** The pull request a conversation works on. The pair is always written together. */ +export interface PRAutomationBinding { + repository: string; + pullNumber: number; +} + +/** What an event handler saw about the run it acts on. */ +export interface PRAutomationRunFence extends PRAutomationBinding { + epoch: string; +} + +/** What an approval was authorized against: the bound repository and its run epoch. */ +export interface PRAutomationBotFence { + repository: string; + epoch: string; +} + +export type PRAutomationEventStopCode = Exclude; + +export type EnablePRAutomationResult = + | { ok: true; value: IPRAutomation } + /** `conflict`: a concurrent enable replaced this one before it committed; retry. */ + | { ok: false; error: { code: 'conversation_gone' | 'owner_inactive' | 'conflict' } }; + +export interface EnablePRAutomationParams extends PRAutomationKey { + trust?: PRAutomationTrustLevel; + /** The administrator ceiling (`prAutomation.maxTrust`); `trust` is clamped to it. */ + maxTrust: PRAutomationTrustLevel; + binding?: PRAutomationBinding; +} + +export interface SettlePRAutomationRoundParams extends PRAutomationKey { + /** The round that is reporting. A completion for any other round is ignored. */ + round: number; + /** The `runId` the claim returned. A completion for any other run is ignored. */ + runId: string; + state: 'waiting' | 'needs_user'; +} + +export interface ClaimPRAutomationRoundParams extends PRAutomationKey { + /** + * The pull request the delivery is for. A delivery for a pull request the record is no + * longer bound to cannot claim a round, so an event that arrives after a rebind is rejected + * instead of spending the new pull request's budget. The epoch the delivery was admitted + * under is part of it, so a delivery for a run that was restarted is rejected too. + */ + binding: PRAutomationRunFence; + maxRounds: number; + maxMinutes: number; + /** The head the round will work on; recorded so a stale event can be recognized. */ + headSha: string; + now?: Date; +} + +export type PRAutomationClaimErrorCode = + | 'not_found' + | 'not_active' + | 'binding_mismatch' + | 'conversation_gone' + | 'owner_inactive' + | 'round_cap' + | 'time_cap' + | 'stale_head' + /** The record changed while the claim was being decided; the delivery may be retried. */ + | 'conflict'; + +export type ClaimPRAutomationRoundResult = + | { ok: true; value: IPRAutomation & { runId: string } } + | { ok: false; error: { code: PRAutomationClaimErrorCode } }; + +export type PRAutomationBotResult = + | { ok: true; value: IPRAutomation } + | { ok: false; error: { code: 'not_found' | 'bot_limit' | 'binding_mismatch' } }; + +export interface ReapPRAutomationsParams { + /** Continue after this record, as returned in `next` by the previous call. */ + after?: string; + limit?: number; +} + +export interface ReapPRAutomationsResult { + scanned: number; + removed: number; + /** Where the next call continues; absent once every record has been read. */ + next?: string; +}