diff --git a/scripts/e2e/lib/kitchen-sink-plugin/assertions.mjs b/scripts/e2e/lib/kitchen-sink-plugin/assertions.mjs index be6d6b66968..3499ee115b3 100644 --- a/scripts/e2e/lib/kitchen-sink-plugin/assertions.mjs +++ b/scripts/e2e/lib/kitchen-sink-plugin/assertions.mjs @@ -317,7 +317,7 @@ const expectMissing = (listValue, expected, field) => { } }; -const INVALID_PROBE_DIAGNOSTIC_SURFACE_MODES = new Set(["full", "conformance", "adversarial"]); +const INVALID_PROBE_DIAGNOSTIC_SURFACE_MODES = new Set(["full", "adversarial"]); const requiredFullDiagnosticCanaries = new Set([ "agent tool result middleware must be a function", "trusted tool policy registration requires id, description, and evaluate()", diff --git a/src/agents/subagent-registry-memory.test.ts b/src/agents/subagent-registry-memory.test.ts new file mode 100644 index 00000000000..822eac1354c --- /dev/null +++ b/src/agents/subagent-registry-memory.test.ts @@ -0,0 +1,86 @@ +import { afterEach, describe, expect, it } from "vitest"; +import { + getSubagentRunsForChildSession, + getSubagentRunsForCollectorGroup, + subagentRuns, +} from "./subagent-registry-memory.js"; +import type { SubagentRunRecord } from "./subagent-registry.types.js"; + +function createRun(runId: string, childSessionKey: string): SubagentRunRecord { + return { + runId, + childSessionKey, + requesterSessionKey: "agent:main:main", + requesterDisplayKey: "main", + task: runId, + cleanup: "keep", + createdAt: 1, + }; +} + +afterEach(() => { + subagentRuns.clear(); +}); + +describe("subagent run memory indexes", () => { + it("tracks child-session generations across replacement, deletion, and clear", () => { + const first = createRun("run-first", "agent:main:subagent:shared"); + const second = createRun("run-second", "agent:main:subagent:shared"); + subagentRuns.set(first.runId, first); + subagentRuns.set(second.runId, second); + + expect([...getSubagentRunsForChildSession(first.childSessionKey)]).toEqual([first, second]); + + const replacement = createRun("run-first", "agent:main:subagent:replacement"); + subagentRuns.set(replacement.runId, replacement); + expect([...getSubagentRunsForChildSession(first.childSessionKey)]).toEqual([second]); + expect([...getSubagentRunsForChildSession(replacement.childSessionKey)]).toEqual([replacement]); + + subagentRuns.delete(second.runId); + expect([...getSubagentRunsForChildSession(first.childSessionKey)]).toEqual([]); + + subagentRuns.clear(); + expect([...getSubagentRunsForChildSession(replacement.childSessionKey)]).toEqual([]); + }); + + it("tracks every collector group member across replacement, deletion, and clear", () => { + const requesterSessionKey = "agent:main:requester"; + const first = { + ...createRun("run-first", "agent:main:subagent:first"), + requesterSessionKey, + collect: true, + groupId: "swarm:shared", + collectorCompletion: { status: "done" as const }, + }; + const incomplete = { + ...createRun("run-incomplete", "agent:main:subagent:incomplete"), + requesterSessionKey, + collect: true, + groupId: "swarm:shared", + }; + subagentRuns.set(first.runId, first); + subagentRuns.set(incomplete.runId, incomplete); + + expect([...getSubagentRunsForCollectorGroup(requesterSessionKey, "swarm:shared")]).toEqual([ + [first.runId, first], + [incomplete.runId, incomplete], + ]); + + const replacement = { ...first, groupId: "swarm:replacement" }; + subagentRuns.set(replacement.runId, replacement); + expect([...getSubagentRunsForCollectorGroup(requesterSessionKey, "swarm:shared")]).toEqual([ + [incomplete.runId, incomplete], + ]); + expect([...getSubagentRunsForCollectorGroup(requesterSessionKey, "swarm:replacement")]).toEqual( + [[replacement.runId, replacement]], + ); + + subagentRuns.delete(incomplete.runId); + expect([...getSubagentRunsForCollectorGroup(requesterSessionKey, "swarm:shared")]).toEqual([]); + + subagentRuns.clear(); + expect([...getSubagentRunsForCollectorGroup(requesterSessionKey, "swarm:replacement")]).toEqual( + [], + ); + }); +}); diff --git a/src/agents/subagent-registry-memory.ts b/src/agents/subagent-registry-memory.ts index 4e0cbf88408..c61eb5893de 100644 --- a/src/agents/subagent-registry-memory.ts +++ b/src/agents/subagent-registry-memory.ts @@ -9,17 +9,74 @@ import type { SubagentRunRecord } from "./subagent-registry.types.js"; // Preflight consults the collector lookup on every Gateway agent request, so it // must stay O(1) regardless of retained collector records. The map subclass // maintains the index through every existing mutation path (registry, run -// manager, tests); collect/childSessionKey are fixed at registration, so -// in-place field edits never require re-indexing. +// manager, tests); collector identity and childSessionKey are fixed at +// registration, so in-place lifecycle field edits never require re-indexing. const collectorRunIdByChildSessionKey = new Map(); +const runsByChildSessionKey = new Map>(); +const runsByCollectorGroupKey = new Map>(); + +function collectorGroupKey(entry: SubagentRunRecord): string | undefined { + if (entry.collect !== true || !entry.groupId) { + return undefined; + } + return JSON.stringify([ + entry.swarmRequesterSessionKey ?? entry.requesterSessionKey, + entry.groupId, + ]); +} + +function removeRunFromChildSessionIndex(runId: string, entry: SubagentRunRecord) { + const sessionRuns = runsByChildSessionKey.get(entry.childSessionKey); + if (sessionRuns?.get(runId) !== entry) { + return; + } + sessionRuns.delete(runId); + if (sessionRuns.size === 0) { + runsByChildSessionKey.delete(entry.childSessionKey); + } +} + +function removeRunFromCollectorGroupIndex(runId: string, entry: SubagentRunRecord) { + const key = collectorGroupKey(entry); + if (!key) { + return; + } + const groupRuns = runsByCollectorGroupKey.get(key); + if (groupRuns?.get(runId) !== entry) { + return; + } + groupRuns.delete(runId); + if (groupRuns.size === 0) { + runsByCollectorGroupKey.delete(key); + } +} class SubagentRunMap extends Map { override set(runId: string, entry: SubagentRunRecord): this { const prev = this.get(runId); - if (prev?.collect === true && prev.childSessionKey) { - collectorRunIdByChildSessionKey.delete(prev.childSessionKey); + if (prev) { + removeRunFromChildSessionIndex(runId, prev); + removeRunFromCollectorGroupIndex(runId, prev); + if (prev.collect === true && prev.childSessionKey) { + collectorRunIdByChildSessionKey.delete(prev.childSessionKey); + } } super.set(runId, entry); + let sessionRuns = runsByChildSessionKey.get(entry.childSessionKey); + if (!sessionRuns) { + sessionRuns = new Map(); + runsByChildSessionKey.set(entry.childSessionKey, sessionRuns); + } + sessionRuns.set(runId, entry); + const groupKey = collectorGroupKey(entry); + if (groupKey) { + let groupRuns = runsByCollectorGroupKey.get(groupKey); + if (!groupRuns) { + groupRuns = new Map(); + runsByCollectorGroupKey.set(groupKey, groupRuns); + } + groupRuns.set(runId, entry); + } if (entry.collect === true && entry.childSessionKey) { collectorRunIdByChildSessionKey.set(entry.childSessionKey, runId); } @@ -28,6 +85,10 @@ class SubagentRunMap extends Map { override delete(runId: string): boolean { const prev = this.get(runId); + if (prev) { + removeRunFromChildSessionIndex(runId, prev); + removeRunFromCollectorGroupIndex(runId, prev); + } if ( prev?.collect === true && prev.childSessionKey && @@ -41,11 +102,29 @@ class SubagentRunMap extends Map { override clear(): void { super.clear(); collectorRunIdByChildSessionKey.clear(); + runsByChildSessionKey.clear(); + runsByCollectorGroupKey.clear(); } } export const subagentRuns: Map = new SubagentRunMap(); +/** Iterate live generations for one child session without scanning the registry. */ +export function getSubagentRunsForChildSession( + childSessionKey: string, +): Iterable { + return runsByChildSessionKey.get(childSessionKey)?.values() ?? []; +} + +/** Iterate live collector members for one requester/group archive decision. */ +export function getSubagentRunsForCollectorGroup( + requesterSessionKey: string, + groupId: string, +): Iterable<[string, SubagentRunRecord]> { + const key = JSON.stringify([requesterSessionKey, groupId]); + return runsByCollectorGroupKey.get(key)?.entries() ?? []; +} + /** Resolve a collector tombstone that reserves its child session from ordinary turns. */ export function findSwarmCollectorSession(childSessionKey?: string): SubagentRunRecord | undefined { const key = childSessionKey?.trim(); diff --git a/src/agents/subagent-registry-sweep-kill.ts b/src/agents/subagent-registry-sweep-kill.ts index c17e9aa8079..35474cab38d 100644 --- a/src/agents/subagent-registry-sweep-kill.ts +++ b/src/agents/subagent-registry-sweep-kill.ts @@ -21,11 +21,11 @@ import { } from "./subagent-session-reconciliation.js"; function findNextSubagentRunCreatedAt( - runs: Map, + candidates: Iterable, entry: SubagentRunRecord, ): number | undefined { let nextCreatedAt = entry.killReconciliation?.supersededAt; - for (const candidate of runs.values()) { + for (const candidate of candidates) { if ( candidate.runId === entry.runId || candidate.childSessionKey !== entry.childSessionKey || @@ -38,22 +38,10 @@ function findNextSubagentRunCreatedAt( return nextCreatedAt; } -function isStableCancellation(task: TaskRecord | undefined) { - return task?.status === "cancelled" && !isProvisionalSubagentKillTask(task); -} - -function isUnstableTask(task: TaskRecord | undefined) { - return ( - task !== undefined && - (task.status === "queued" || task.status === "running" || isProvisionalSubagentKillTask(task)) - ); -} - -export function resolveSubagentTaskForRun( - runs: Map, +function resolveSubagentTaskForRunGeneration( entry: SubagentRunRecord, + nextRunCreatedAt: number | undefined, ) { - const nextRunCreatedAt = findNextSubagentRunCreatedAt(runs, entry); const generationStartedAt = entry.sessionStartedAt ?? entry.createdAt; return findDetachedTaskRun({ runId: entry.taskRunId ?? entry.runId, @@ -70,6 +58,27 @@ export function resolveSubagentTaskForRun( }); } +function isStableCancellation(task: TaskRecord | undefined) { + return task?.status === "cancelled" && !isProvisionalSubagentKillTask(task); +} + +function isUnstableTask(task: TaskRecord | undefined) { + return ( + task !== undefined && + (task.status === "queued" || task.status === "running" || isProvisionalSubagentKillTask(task)) + ); +} + +export function resolveSubagentTaskForRun( + candidates: Iterable, + entry: SubagentRunRecord, +) { + return resolveSubagentTaskForRunGeneration( + entry, + findNextSubagentRunCreatedAt(candidates, entry), + ); +} + function resolveCompletionFromTerminalTask(task: TaskRecord | undefined, entry: SubagentRunRecord) { if ( !task || @@ -108,6 +117,7 @@ export async function reconcileProvisionalSubagentKill(params: { ) => Promise; retireSupersededRun: (runId: string, entry: SubagentRunRecord) => Promise; startSubagentAnnounceCleanupFlow: (runId: string, entry: SubagentRunRecord) => boolean; + getRunsForChildSession: (childSessionKey: string) => Iterable; warn: (message: string, meta?: Record) => void; }): Promise { const { entry, now, runId, runs } = params; @@ -115,9 +125,22 @@ export async function reconcileProvisionalSubagentKill(params: { if (!killReconciliation) { return false; } - const taskResolution = resolveSubagentTaskForRun(runs, entry); + // The child-session index stays current across awaits. Re-read it at each + // decision boundary so a newly registered generation can supersede this run. + const resolveGeneration = () => { + const nextRunCreatedAt = findNextSubagentRunCreatedAt( + params.getRunsForChildSession(entry.childSessionKey), + entry, + ); + return { + nextRunCreatedAt, + taskResolution: resolveSubagentTaskForRunGeneration(entry, nextRunCreatedAt), + }; + }; + const initialGeneration = resolveGeneration(); + const taskResolution = initialGeneration.taskResolution; const task = taskResolution.task; - const nextRunCreatedAt = findNextSubagentRunCreatedAt(runs, entry); + const nextRunCreatedAt = initialGeneration.nextRunCreatedAt; const hasStableTaskCancellation = isStableCancellation(task); const killedAt = killReconciliation.killedAt; const isCurrentKill = () => @@ -205,7 +228,7 @@ export async function reconcileProvisionalSubagentKill(params: { if (!isCurrentKill()) { return false; } - const taskAfterResolution = resolveSubagentTaskForRun(runs, entry); + const taskAfterResolution = resolveGeneration().taskResolution; const taskAfter = taskAfterResolution.task; const stableCancellationWonDuringCompletion = isStableCancellation(taskAfter) && completionEndedAt >= killedAt; @@ -216,7 +239,7 @@ export async function reconcileProvisionalSubagentKill(params: { if (!isCurrentKill()) { return false; } - const taskBeforeResolution = resolveSubagentTaskForRun(runs, entry); + const taskBeforeResolution = resolveGeneration().taskResolution; const taskBefore = taskBeforeResolution.task; const stableTaskCancellationAfterReconciliation = isStableCancellation(taskBefore); const taskNeedsStabilization = @@ -239,7 +262,7 @@ export async function reconcileProvisionalSubagentKill(params: { suppressDelivery: true, }); if (finalizedTasks.length === 0) { - const taskAfterResolution = resolveSubagentTaskForRun(runs, entry); + const taskAfterResolution = resolveGeneration().taskResolution; const taskAfter = taskAfterResolution.task; if (taskAfterResolution.lookup === "available" && isUnstableTask(taskAfter)) { params.warn("killed task was not stabilized during sweep", { @@ -264,7 +287,7 @@ export async function reconcileProvisionalSubagentKill(params: { return false; } } - if (findNextSubagentRunCreatedAt(runs, entry) !== undefined) { + if (resolveGeneration().nextRunCreatedAt !== undefined) { await params.retireSupersededRun(runId, entry); return true; } diff --git a/src/agents/subagent-registry-sweeper.ts b/src/agents/subagent-registry-sweeper.ts index 103fbcb70de..02484e8f730 100644 --- a/src/agents/subagent-registry-sweeper.ts +++ b/src/agents/subagent-registry-sweeper.ts @@ -88,6 +88,11 @@ export function createSubagentRegistrySweeper(params: { runContextEngineSubagentEnded: (params: ContextEngineSubagentEndedParams) => Promise; notifyContextEngineSubagentEnded: (params: ContextEngineSubagentEndedParams) => Promise; retireSupersededRun: (runId: string, entry: SubagentRunRecord) => Promise; + getRunsForChildSession: (childSessionKey: string) => Iterable; + getRunsForCollectorGroup: ( + requesterSessionKey: string, + groupId: string, + ) => Iterable<[string, SubagentRunRecord]>; warn: (message: string, meta?: Record) => void; }) { const { runs, resumedRuns } = params; @@ -234,10 +239,17 @@ export function createSubagentRegistrySweeper(params: { const storeCache: SubagentSessionStoreCache = new Map(); let mutated = false; const mutatedRunIds = new Set(); - const archivedCollectorGroups = new Set(); - const suspendedEntries = [...runs.entries()].filter(([, entry]) => - isSuspendedPendingFinalDelivery(entry), - ); + const collectorArchiveCandidates = new Map< + string, + { requesterSessionKey: string; groupId: string } + >(); + const suspendedEntries: Array<[string, SubagentRunRecord]> = []; + for (const pair of runs.entries()) { + const [, entry] = pair; + if (isSuspendedPendingFinalDelivery(entry)) { + suspendedEntries.push(pair); + } + } const pressureDiscardRunIds = new Set(); if (suspendedEntries.length > SUSPENDED_DELIVERY_HARD_CAP) { const pressureCount = Math.max( @@ -357,6 +369,7 @@ export function createSubagentRegistrySweeper(params: { completeSubagentRunWithRecovery: params.completeSubagentRunWithRecovery, retireSupersededRun: params.retireSupersededRun, startSubagentAnnounceCleanupFlow: params.startSubagentAnnounceCleanupFlow, + getRunsForChildSession: params.getRunsForChildSession, warn: params.warn, }); if (reconciled) { @@ -396,97 +409,12 @@ export function createSubagentRegistrySweeper(params: { const groupKey = groupId ? JSON.stringify([swarmRequesterSessionKey, groupId]) : undefined; - if (!groupKey || archivedCollectorGroups.has(groupKey)) { - continue; - } - const groupEntries = [...runs.entries()].filter( - ([, candidate]) => - candidate.collect === true && - (candidate.swarmRequesterSessionKey ?? candidate.requesterSessionKey) === - swarmRequesterSessionKey && - candidate.groupId === groupId, - ); - if ( - groupEntries.some( - ([, candidate]) => - !candidate.collectorCompletion || - candidate.collectorLaunchCleanupPending === true || - candidate.archiveAtMs === undefined || - candidate.archiveAtMs > now, - ) - ) { - continue; - } - let deleteFailed = false; - for (const [candidateRunId, candidate] of groupEntries) { - try { - await deleteSession(candidate.childSessionKey); - } catch (error) { - params.warn("sessions.delete failed during collector group sweep; keeping group", { - runId: candidateRunId, - childSessionKey: candidate.childSessionKey, - groupId, - error, - }); - deleteFailed = true; - break; - } - } - if (deleteFailed) { - continue; - } - let attachmentCleanupFailed = false; - for (const [candidateRunId, candidate] of groupEntries) { - if (await safeRemoveAttachmentsDir(candidate)) { - continue; - } - params.warn("attachment cleanup failed during collector group sweep; keeping group", { - runId: candidateRunId, - childSessionKey: candidate.childSessionKey, + if (groupKey && groupId) { + collectorArchiveCandidates.set(groupKey, { + requesterSessionKey: swarmRequesterSessionKey, groupId, }); - attachmentCleanupFailed = true; - break; } - if (attachmentCleanupFailed) { - continue; - } - let contextCleanupFailed = false; - for (const [candidateRunId, candidate] of groupEntries) { - if ( - candidate.cleanup === "delete" || - typeof candidate.contextEngineCleanupCompletedAt === "number" - ) { - continue; - } - try { - await params.runContextEngineSubagentEnded(sweptContext(candidate)); - candidate.contextEngineCleanupCompletedAt = Date.now(); - params.persist(candidateRunId); - } catch (error) { - params.warn( - "context-engine cleanup failed during collector group sweep; keeping group", - { - runId: candidateRunId, - childSessionKey: candidate.childSessionKey, - groupId, - error, - }, - ); - contextCleanupFailed = true; - break; - } - } - if (contextCleanupFailed) { - continue; - } - for (const [candidateRunId] of groupEntries) { - params.clearPendingLifecycleError(candidateRunId); - runs.delete(candidateRunId); - mutatedRunIds.add(candidateRunId); - } - archivedCollectorGroups.add(groupKey); - mutated = true; continue; } if (!entry.archiveAtMs && entry.cleanup === "keep" && entry.spawnMode !== "session") { @@ -532,6 +460,108 @@ export function createSubagentRegistrySweeper(params: { await params.notifyContextEngineSubagentEnded(sweptContext(entry)); }); } + for (const { requesterSessionKey, groupId } of collectorArchiveCandidates.values()) { + // Earlier sweep work may await while group membership changes. Read the + // mutation-owned index once, after per-run collector cleanup has settled. + const groupEntries = [...params.getRunsForCollectorGroup(requesterSessionKey, groupId)]; + if ( + groupEntries.some( + ([, candidate]) => + !candidate.collectorCompletion || + candidate.collectorLaunchCleanupPending === true || + candidate.archiveAtMs === undefined || + candidate.archiveAtMs > now, + ) + ) { + continue; + } + let deleteFailed = false; + for (const [candidateRunId, candidate] of groupEntries) { + try { + await deleteSession(candidate.childSessionKey); + } catch (error) { + params.warn("sessions.delete failed during collector group sweep; keeping group", { + runId: candidateRunId, + childSessionKey: candidate.childSessionKey, + groupId, + error, + }); + deleteFailed = true; + break; + } + } + if (deleteFailed) { + continue; + } + let attachmentCleanupFailed = false; + for (const [candidateRunId, candidate] of groupEntries) { + if (await safeRemoveAttachmentsDir(candidate)) { + continue; + } + params.warn("attachment cleanup failed during collector group sweep; keeping group", { + runId: candidateRunId, + childSessionKey: candidate.childSessionKey, + groupId, + }); + attachmentCleanupFailed = true; + break; + } + if (attachmentCleanupFailed) { + continue; + } + let contextCleanupFailed = false; + for (const [candidateRunId, candidate] of groupEntries) { + if ( + candidate.cleanup === "delete" || + typeof candidate.contextEngineCleanupCompletedAt === "number" + ) { + continue; + } + try { + await params.runContextEngineSubagentEnded(sweptContext(candidate)); + candidate.contextEngineCleanupCompletedAt = Date.now(); + params.persist(candidateRunId); + } catch (error) { + params.warn( + "context-engine cleanup failed during collector group sweep; keeping group", + { + runId: candidateRunId, + childSessionKey: candidate.childSessionKey, + groupId, + error, + }, + ); + contextCleanupFailed = true; + break; + } + } + if (contextCleanupFailed) { + continue; + } + // Cleanup awaits can admit a new collector or replace an existing run. + // Delete only the exact group snapshot whose resources were cleaned. + const expectedGroupEntries = new Map(groupEntries); + const liveGroupEntries = [...params.getRunsForCollectorGroup(requesterSessionKey, groupId)]; + if ( + liveGroupEntries.length !== groupEntries.length || + liveGroupEntries.some( + ([candidateRunId, candidate]) => + expectedGroupEntries.get(candidateRunId) !== candidate || + !candidate.collectorCompletion || + candidate.collectorLaunchCleanupPending === true || + candidate.archiveAtMs === undefined || + candidate.archiveAtMs > now, + ) + ) { + continue; + } + for (const [candidateRunId] of liveGroupEntries) { + params.clearPendingLifecycleError(candidateRunId); + runs.delete(candidateRunId); + mutatedRunIds.add(candidateRunId); + } + mutated = true; + } params.sweepPendingLifecycle(now); if (mutated) { diff --git a/src/agents/subagent-registry.test.ts b/src/agents/subagent-registry.test.ts index 8cc1f13a431..51101a13a8b 100644 --- a/src/agents/subagent-registry.test.ts +++ b/src/agents/subagent-registry.test.ts @@ -479,6 +479,200 @@ describe("subagent registry seam flow", () => { ).toHaveLength(2); }); + it("keeps collector archive groups scoped to their requester", async () => { + const now = Date.now(); + for (const [requesterSessionKey, archiveAtMs] of [ + ["agent:main:requester-one", now - 1], + ["agent:main:requester-two", now + 1_000], + ] as const) { + mod.addSubagentRunForTests({ + runId: `run-${requesterSessionKey}`, + childSessionKey: `${requesterSessionKey}:subagent:collector`, + requesterSessionKey, + task: "retain requester-scoped collector groups", + cleanup: "delete", + createdAt: now - 10_000, + endedAt: now - 5_000, + cleanupCompletedAt: now - 4_000, + archiveAtMs, + collect: true, + groupId: "swarm:shared-group-id", + collectorCompletion: { status: "done" }, + }); + } + + await mod.testing.sweepOnceForTests(); + + expect(mod.getSubagentRunByRunId("run-agent:main:requester-one")).toBeUndefined(); + expect(mod.getSubagentRunByRunId("run-agent:main:requester-two")).toBeDefined(); + }); + + it("keeps completed collectors while any group member is incomplete", async () => { + const now = Date.now(); + mod.addSubagentRunForTests({ + runId: "run-collector-complete", + childSessionKey: "agent:main:subagent:collector-complete", + task: "completed collector", + createdAt: now - 10_000, + endedAt: now - 5_000, + archiveAtMs: now - 1, + collect: true, + groupId: "swarm:incomplete-member", + collectorCompletion: { status: "done" }, + }); + mod.addSubagentRunForTests({ + runId: "run-collector-incomplete", + childSessionKey: "agent:main:subagent:collector-incomplete", + task: "incomplete collector", + createdAt: now - 9_000, + endedAt: now - 4_000, + archiveAtMs: now + 1_000, + collect: true, + groupId: "swarm:incomplete-member", + }); + + await mod.testing.sweepOnceForTests(); + + expect(mod.getSubagentRunByRunId("run-collector-complete")).toBeDefined(); + expect(mod.getSubagentRunByRunId("run-collector-incomplete")).toBeDefined(); + }); + + it("refreshes collector membership after awaited sweep work", async () => { + const now = Date.now(); + let releaseFirstDelete: (() => void) | undefined; + let shouldBlockDelete = true; + mocks.callGateway.mockImplementation((request: { method?: string }) => { + if (request.method !== "sessions.delete" || !shouldBlockDelete) { + return Promise.resolve({}); + } + shouldBlockDelete = false; + return new Promise>((resolve) => { + releaseFirstDelete = () => resolve({}); + }); + }); + mod.addSubagentRunForTests({ + runId: "run-archive-blocker", + childSessionKey: "agent:main:subagent:archive-blocker", + task: "hold the sweep before collector archival", + cleanup: "delete", + createdAt: now - 10_000, + endedAt: now - 5_000, + cleanupCompletedAt: now - 4_000, + archiveAtMs: now - 1, + }); + mod.addSubagentRunForTests({ + runId: "run-collector-before-await", + childSessionKey: "agent:main:subagent:collector-before-await", + task: "completed collector present at sweep start", + createdAt: now - 10_000, + endedAt: now - 5_000, + archiveAtMs: now - 1, + collect: true, + groupId: "swarm:late-member", + collectorCompletion: { status: "done" }, + }); + + const sweep = mod.testing.runSweeperTickForTests(); + await waitForFast(() => expect(releaseFirstDelete).toBeTypeOf("function")); + mod.addSubagentRunForTests({ + runId: "run-collector-after-await", + childSessionKey: "agent:main:subagent:collector-after-await", + task: "incomplete collector registered during sweep", + createdAt: now, + collect: true, + groupId: "swarm:late-member", + }); + releaseFirstDelete?.(); + await sweep; + + expect(mod.getSubagentRunByRunId("run-collector-before-await")).toBeDefined(); + expect(mod.getSubagentRunByRunId("run-collector-after-await")).toBeDefined(); + }); + + it("revalidates collector membership after collector cleanup awaits", async () => { + const now = Date.now(); + let releaseCollectorDelete: (() => void) | undefined; + mocks.callGateway.mockImplementation((request: { method?: string }) => { + if (request.method !== "sessions.delete" || releaseCollectorDelete) { + return Promise.resolve({}); + } + return new Promise>((resolve) => { + releaseCollectorDelete = () => resolve({}); + }); + }); + mod.addSubagentRunForTests({ + runId: "run-collector-cleanup-snapshot", + childSessionKey: "agent:main:subagent:collector-cleanup-snapshot", + task: "completed collector present before cleanup", + createdAt: now - 10_000, + endedAt: now - 5_000, + archiveAtMs: now - 1, + collect: true, + groupId: "swarm:cleanup-race", + collectorCompletion: { status: "done" }, + }); + + const sweep = mod.testing.runSweeperTickForTests(); + await waitForFast(() => expect(releaseCollectorDelete).toBeTypeOf("function")); + mod.addSubagentRunForTests({ + runId: "run-collector-added-during-cleanup", + childSessionKey: "agent:main:subagent:collector-added-during-cleanup", + task: "incomplete collector registered during cleanup", + createdAt: now, + collect: true, + groupId: "swarm:cleanup-race", + }); + releaseCollectorDelete?.(); + await sweep; + + expect(mod.getSubagentRunByRunId("run-collector-cleanup-snapshot")).toBeDefined(); + expect(mod.getSubagentRunByRunId("run-collector-added-during-cleanup")).toBeDefined(); + }); + + it("keeps a collector replaced during collector cleanup awaits", async () => { + const now = Date.now(); + let releaseCollectorDelete: (() => void) | undefined; + mocks.callGateway.mockImplementation((request: { method?: string }) => { + if (request.method !== "sessions.delete" || releaseCollectorDelete) { + return Promise.resolve({}); + } + return new Promise>((resolve) => { + releaseCollectorDelete = () => resolve({}); + }); + }); + mod.addSubagentRunForTests({ + runId: "run-collector-replaced-during-cleanup", + childSessionKey: "agent:main:subagent:collector-before-replacement", + task: "collector before replacement", + createdAt: now - 10_000, + endedAt: now - 5_000, + archiveAtMs: now - 1, + collect: true, + groupId: "swarm:replacement-race", + collectorCompletion: { status: "done" }, + }); + + const sweep = mod.testing.runSweeperTickForTests(); + await waitForFast(() => expect(releaseCollectorDelete).toBeTypeOf("function")); + mod.addSubagentRunForTests({ + runId: "run-collector-replaced-during-cleanup", + childSessionKey: "agent:main:subagent:collector-after-replacement", + task: "collector after replacement", + createdAt: now, + endedAt: now, + archiveAtMs: now - 1, + collect: true, + groupId: "swarm:replacement-race", + collectorCompletion: { status: "done" }, + }); + releaseCollectorDelete?.(); + await sweep; + + expect( + mod.getSubagentRunByRunId("run-collector-replaced-during-cleanup")?.childSessionKey, + ).toBe("agent:main:subagent:collector-after-replacement"); + }); + it("keeps collector groups while any member owes failed-launch cleanup", async () => { const now = Date.now(); mod.addSubagentRunForTests({ diff --git a/src/agents/subagent-registry.ts b/src/agents/subagent-registry.ts index a965545e298..49bb4fb8208 100644 --- a/src/agents/subagent-registry.ts +++ b/src/agents/subagent-registry.ts @@ -31,7 +31,11 @@ import { } from "./subagent-registry-helpers.js"; import { createSubagentRegistryLifecycleController } from "./subagent-registry-lifecycle.js"; import { createSubagentRegistryListener } from "./subagent-registry-listener.js"; -import { subagentRuns } from "./subagent-registry-memory.js"; +import { + getSubagentRunsForChildSession, + getSubagentRunsForCollectorGroup, + subagentRuns, +} from "./subagent-registry-memory.js"; import { createSubagentRegistryPublicApi } from "./subagent-registry-public-api.js"; import { createSubagentRegistryRestorer } from "./subagent-registry-restore.js"; import { @@ -93,7 +97,7 @@ function persistSubagentRunsOrThrow(...runIds: string[]) { } function findSubagentTaskForRun(entry: SubagentRunRecord) { - return resolveSubagentTaskForRun(subagentRuns, entry); + return resolveSubagentTaskForRun(getSubagentRunsForChildSession(entry.childSessionKey), entry); } export function scheduleSubagentOrphanRecovery(params?: { delayMs?: number; maxRetries?: number }) { @@ -404,6 +408,8 @@ const subagentSweeper = createSubagentRegistrySweeper({ runContextEngineSubagentEnded: contextCleanup.runContextEngineSubagentEnded, notifyContextEngineSubagentEnded: contextCleanup.notifyContextEngineSubagentEnded, retireSupersededRun: retireSupersededSubagentRun, + getRunsForChildSession: getSubagentRunsForChildSession, + getRunsForCollectorGroup: getSubagentRunsForCollectorGroup, warn: (message, meta) => log.warn(message, meta), }); diff --git a/test/scripts/kitchen-sink-plugin-assertions.test.ts b/test/scripts/kitchen-sink-plugin-assertions.test.ts index b1629c6fd79..7583ebf7fc5 100644 --- a/test/scripts/kitchen-sink-plugin-assertions.test.ts +++ b/test/scripts/kitchen-sink-plugin-assertions.test.ts @@ -71,11 +71,13 @@ function runAssertInstalled({ diagnostics = [], env = {}, inspectPayload, + surfaceMode = "full", }: { allInspectPayload?: unknown; diagnostics?: Array<{ level: string; message: string }>; env?: NodeJS.ProcessEnv; inspectPayload?: ReturnType; + surfaceMode?: string; } = {}) { const label = `diagnostics-${process.pid}-${Date.now()}-${Math.random().toString(16).slice(2)}`; const pluginId = "openclaw-kitchen-sink-fixture"; @@ -121,7 +123,7 @@ function runAssertInstalled({ KITCHEN_SINK_LABEL: label, KITCHEN_SINK_SOURCE: "npm", KITCHEN_SINK_SPEC: "npm:@openclaw/kitchen-sink@latest", - KITCHEN_SINK_SURFACE_MODE: "full", + KITCHEN_SINK_SURFACE_MODE: surfaceMode, KITCHEN_SINK_TMP_DIR: scratchRoot, }, }); @@ -313,6 +315,18 @@ describe("kitchen-sink plugin assertions", () => { expect(result.status).toBe(0); }); + it("rejects diagnostics in conformance mode", () => { + const result = runAssertInstalled({ + diagnostics: diagnosticErrors(["plugin must declare contracts.tools for: kitchen-sink-tool"]), + surfaceMode: "conformance", + }); + + expect(result.status).not.toBe(0); + expect(result.stderr).toContain( + "unexpected kitchen-sink diagnostic errors: plugin must declare contracts.tools for: kitchen-sink-tool", + ); + }); + it("requires kitchen-sink plugins to appear in inspect-all output", () => { const result = runAssertInstalled({ allInspectPayload: [fullSurfaceInspectPayload("other-plugin")], diff --git a/test/scripts/plugin-prerelease-test-plan.test.ts b/test/scripts/plugin-prerelease-test-plan.test.ts index 352194074a3..3ae9266039a 100644 --- a/test/scripts/plugin-prerelease-test-plan.test.ts +++ b/test/scripts/plugin-prerelease-test-plan.test.ts @@ -180,7 +180,7 @@ describe("scripts/lib/plugin-prerelease-test-plan.mjs", () => { expect(assertionsScript).toContain("assertClawHubExternalInstallContract"); expect(assertionsScript).toContain("expectedErrorMessages"); expect(assertionsScript).toContain( - 'const INVALID_PROBE_DIAGNOSTIC_SURFACE_MODES = new Set(["full", "conformance", "adversarial"]);', + 'const INVALID_PROBE_DIAGNOSTIC_SURFACE_MODES = new Set(["full", "adversarial"]);', ); expect(assertionsScript).toContain("!INVALID_PROBE_DIAGNOSTIC_SURFACE_MODES.has(surfaceMode)"); expect(readFileSync("scripts/e2e/lib/clawhub-fixture-server.cjs", "utf8")).toContain(