File size: 4,894 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 114 115 116 117 118 119 120 121 122 123 124 125 126 | import { AsyncLocalStorage } from "node:async_hooks";
import type { ContextEngine } from "../../context-engine/types.js";
import { formatErrorMessage } from "../../infra/errors.js";
import {
AsyncWorkScope,
captureAsyncWorkTracker,
getAsyncWorkSignal,
} from "../../shared/async-work-scope.js";
import { createDeferredCore } from "../../shared/deferred.js";
import type { PreparedModelRuntimeLease } from "../prepared-model-runtime.js";
import type { ContextEngineMaintenanceResources } from "./context-engine-maintenance-work.js";
import { log } from "./logger.js";
import type { EmbeddedAgentCompactResult } from "./types.js";
export type ForegroundCompactionOwner = {
adoptLease: (lease: PreparedModelRuntimeLease) => ContextEngineMaintenanceResources;
resolveEngine: (factory: () => Promise<ContextEngine>) => Promise<ContextEngine>;
captureContext: () => void;
transferEngine: () => void;
};
/** Keeps physical cleanup separate from the compaction result and transcript fences. */
export async function runForegroundCompactionWork(
run: (owner: ForegroundCompactionOwner) => Promise<EmbeddedAgentCompactResult>,
): Promise<EmbeddedAgentCompactResult> {
const result = createDeferredCore<EmbeddedAgentCompactResult>();
const trackOwner = captureAsyncWorkTracker();
const parentSignal = getAsyncWorkSignal();
void trackOwner(async () => {
const work = new AsyncWorkScope();
const factoryWork = new AsyncWorkScope();
const cleanupWork = new AsyncWorkScope();
let runInContext = work.run(() => AsyncLocalStorage.snapshot());
let factoryContext = factoryWork.run(() => AsyncLocalStorage.snapshot());
let cleanupContext = cleanupWork.run(() => AsyncLocalStorage.snapshot());
let foregroundClosed: Promise<void> | undefined;
const closeForegroundWork = () =>
(foregroundClosed ??= AsyncWorkScope.runWhenAllIdle(
() => [work],
() => runInContext(() => work.drain()),
));
let factoryClosed: Promise<void> | undefined;
const closeFactoryWork = () => (factoryClosed ??= factoryContext(() => factoryWork.drain()));
let lease: PreparedModelRuntimeLease | undefined;
let engine: ContextEngine | undefined;
let engineTransferred = false;
const closeFromParent = () => {
runInContext(() => work.beginClose(parentSignal?.reason));
if (!engineTransferred) {
factoryContext(() => factoryWork.beginClose(parentSignal?.reason));
}
cleanupContext(() => cleanupWork.beginClose(parentSignal?.reason));
};
parentSignal?.addEventListener("abort", closeFromParent, { once: true });
if (parentSignal?.aborted) {
closeFromParent();
}
try {
result.resolve(
await work.track(() =>
run({
adoptLease: (acquired) => {
lease = acquired;
return {
closeFactoryWork,
release: async () => {
// Fast maintenance still shares any admitted foreground preparation tails.
await closeForegroundWork();
await acquired[Symbol.asyncDispose]();
},
};
},
resolveEngine: (factory) =>
factoryWork.track(async () => {
factoryContext = AsyncLocalStorage.snapshot();
engine = await factory();
return engine;
}),
captureContext: () => {
runInContext = AsyncLocalStorage.snapshot();
},
transferEngine: () => {
engine = undefined;
engineTransferred = true;
},
}),
),
);
} catch (error) {
result.reject(error);
} finally {
try {
await closeForegroundWork();
if (!engineTransferred) {
cleanupContext = factoryContext(() =>
cleanupWork.run(() => AsyncLocalStorage.snapshot()),
);
// Factory services can need their disposer to start before their signal closes.
const disposal = cleanupContext(() =>
cleanupWork.track(async () => {
try {
await engine?.dispose?.();
} catch (error) {
log.warn("context engine dispose failed", {
errorMessage: formatErrorMessage(error),
});
}
}),
);
await Promise.all([disposal, closeFactoryWork()]);
await AsyncWorkScope.runWhenAllIdle(
() => [cleanupWork],
() => cleanupContext(() => cleanupWork.drain()),
);
}
} finally {
parentSignal?.removeEventListener("abort", closeFromParent);
if (!engineTransferred) {
await lease?.[Symbol.asyncDispose]();
}
}
}
}).catch(result.reject);
return await result.promise;
}
|