File size: 4,131 Bytes
eb3f11e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | import { AsyncLocalStorage } from "node:async_hooks";
import { hasSameContextEngineInstance } from "../../context-engine/registry.js";
import type { ContextEngine } from "../../context-engine/types.js";
import { formatErrorMessage } from "../../infra/errors.js";
import { AsyncWorkScope, trackAsyncWork } from "../../shared/async-work-scope.js";
import type { createSessionMaintenanceOwner } from "../session-maintenance/coordinator.js";
import { log } from "./logger.js";
export type ContextEngineMaintenanceResources = {
closeFactoryWork: () => Promise<void>;
release: () => Promise<void>;
};
/** Maintenance owns cooperating descendants through the operation's actual settlement. */
export async function runContextEngineMaintenanceWork(
run: () => Promise<void>,
signal: AbortSignal,
): Promise<void> {
const work = new AsyncWorkScope();
const context = work.run(() => AsyncLocalStorage.snapshot());
// Accepted background work follows maintenance shutdown, not foreground completion.
const cancel = () => context(() => work.beginClose(signal.reason));
signal.addEventListener("abort", cancel, { once: true });
if (signal.aborted) {
cancel();
}
try {
await work.track(run);
} finally {
try {
// Normal completion must not abort work that returned an early result.
await AsyncWorkScope.runWhenAllIdle(
() => [work],
() => context(() => work.drain()),
);
} finally {
signal.removeEventListener("abort", cancel);
}
}
}
export async function disposeDeferredMaintenanceContextEngine(
params: {
contextEngine: ContextEngine;
runInContext: ReturnType<typeof AsyncLocalStorage.snapshot>;
factoryResourceOwners?: ReadonlySet<ContextEngineMaintenanceResources>;
},
maintenance: Pick<ReturnType<typeof createSessionMaintenanceOwner>, "run" | "signal">,
): Promise<void> {
const failures: unknown[] = [];
const resources = [...(params.factoryResourceOwners ?? [])];
try {
await params.runInContext(() =>
maintenance.run(() =>
runContextEngineMaintenanceWork(async () => {
const disposal = (async () => {
await params.contextEngine.dispose?.();
})();
const factoryWork = resources.map(({ closeFactoryWork }) =>
trackAsyncWork(closeFactoryWork),
);
const outcomes = await Promise.allSettled([disposal, ...factoryWork]);
for (const outcome of outcomes) {
if (outcome.status === "rejected") {
failures.push(outcome.reason);
}
}
}, maintenance.signal),
),
);
} catch (error) {
failures.push(error);
}
// The maintenance owner retains leases until engine and factory descendants have joined.
const releases = await Promise.allSettled(resources.map(async (owner) => await owner.release()));
for (const outcome of releases) {
if (outcome.status === "rejected") {
failures.push(outcome.reason);
}
}
for (const error of failures) {
log.warn("context engine dispose failed after deferred maintenance", {
errorMessage: formatErrorMessage(error),
});
}
}
type ContextEngineFactoryWork = {
contextEngine: ContextEngine;
factoryResourceOwners: Set<ContextEngineMaintenanceResources>;
};
/** Shared engine instances retain every factory lifetime until their final disposer starts. */
export function mergeContextEngineFactoryWork(
params: ContextEngineFactoryWork,
activeEngine: ContextEngine,
activeResources: Set<ContextEngineMaintenanceResources>,
superseded?: ContextEngineFactoryWork,
): Set<ContextEngineMaintenanceResources> {
if (superseded && hasSameContextEngineInstance(superseded.contextEngine, params.contextEngine)) {
for (const resources of superseded.factoryResourceOwners) {
params.factoryResourceOwners.add(resources);
}
}
if (hasSameContextEngineInstance(params.contextEngine, activeEngine)) {
for (const resources of params.factoryResourceOwners) {
activeResources.add(resources);
}
return activeResources;
}
return params.factoryResourceOwners;
}
|