openclaw / src /agents /embedded-agent-runner /compact.foreground-work.ts

Add files using upload-large-folder tool

eb3f11e verified about 1 month ago

4.89 kB

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;
}

	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;
	}