src/context-engine/registry.resources.ts

import { AsyncLocalStorage } from "node:async_hooks";
import { sanitizeForLog } from "../../packages/terminal-core/src/ansi.js";
import { getPluginValueInstance } from "../plugins/plugin-instance-scope.js";
import {
  collectRegistryInvocationInstances,
  PluginInvocationScope,
} from "../plugins/plugin-invocation-scope.js";
import type { ContextEngineRegistration } from "../plugins/registry-contribution-types.js";
import {
  getPluginRegistryInspectionResources,
  type PluginRegistryInspectionResources,
} from "../plugins/registry-inspection-resources.js";
import type { PluginRegistry } from "../plugins/registry-types.js";
import {
  AsyncWorkScope,
  captureAsyncWorkTracker,
  getAsyncWorkSignal,
} from "../shared/async-work-scope.js";
import { createDeferredCore } from "../shared/deferred.js";
import { resolveGlobalSingleton } from "../shared/global-singleton.js";
import type { ContextEngine } from "./types.js";

// Adoption preserves entry identity and the original registration source.
const registrationSources = resolveGlobalSingleton(
  Symbol.for("openclaw.contextEngineRegistrationSources"),
  () => new WeakMap<ContextEngineRegistration, PluginRegistryInspectionResources>(),
);

export function recordContextEngineRegistrationSource(
  registration: ContextEngineRegistration,
  registry: PluginRegistry,
): void {
  const resources = getPluginRegistryInspectionResources(registry);
  if (resources) {
    registrationSources.set(registration, resources);
  }
}

export class ContextEngineFactoryResources {
  private cleanupInvocations?: ReturnType<PluginInvocationScope["beginCleanup"]>;
  readonly work = new AsyncWorkScope();
  readonly context = this.work.run(() => AsyncLocalStorage.snapshot());
  readonly cleanupWork = new AsyncWorkScope();
  readonly cleanupContext = this.cleanupWork.run(() => AsyncLocalStorage.snapshot());
  private readonly parentSignal = getAsyncWorkSignal();
  private readonly abort = () => {
    this.context(() => this.work.beginClose(this.parentSignal?.reason));
    this.cleanupContext(() => this.cleanupWork.beginClose(this.parentSignal?.reason));
  };

  constructor(
    private readonly claims: readonly { release: () => Promise<void> }[],
    private readonly invocations?: PluginInvocationScope,
  ) {
    this.parentSignal?.addEventListener("abort", this.abort, { once: true });
    if (this.parentSignal?.aborted) {
      this.abort();
    }
  }

  run<T>(operation: () => T | Promise<T>): Promise<T> {
    return this.work.track(() => this.context(() => this.invoke(operation)));
  }

  runCleanup<T>(operation: () => T): T {
    this.beginCleanup();
    try {
      return this.cleanupWork.run(() =>
        this.cleanupContext(() =>
          this.cleanupInvocations ? this.cleanupInvocations.scope.run(operation) : operation(),
        ),
      );
    } finally {
      this.context(() => this.work.beginClose());
    }
  }

  beginCleanup(): void {
    this.cleanupInvocations ??= this.invocations?.beginCleanup();
  }

  private invoke<T>(operation: () => T): T {
    return this.invocations ? this.invocations.run(operation) : operation();
  }

  wrap<T>(value: T): T {
    return this.invocations ? this.invocations.wrap(value) : value;
  }

  async release(): Promise<void> {
    this.parentSignal?.removeEventListener("abort", this.abort);
    this.invocations?.release();
    let failure: { error: unknown } | undefined;
    try {
      await this.cleanupInvocations?.release();
    } catch (error) {
      failure = { error };
    }
    // An uncovered donor stays held while primary cleanup finishes, even if that cleanup fails.
    for (const claim of this.claims) {
      try {
        await claim.release();
      } catch (error) {
        failure ??= { error };
      }
    }
    if (failure) {
      throw failure.error;
    }
  }
}

function retainContextEngineFactorySource(
  registry: PluginRegistry,
  registration: ContextEngineRegistration | undefined,
  primary: PluginRegistryInspectionResources | undefined,
  abandon: ContextEngineFactoryFailureCleanup,
): ContextEngineFactoryResources | undefined {
  const claims: Array<{ release: () => Promise<void> }> = [];
  try {
    if (primary) {
      claims.push(primary.retain());
    }
    const source = registration && registrationSources.get(registration);
    if (source && !primary?.coversSource(source)) {
      claims.push(source.retain());
    }
    const instances = collectRegistryInvocationInstances(registry);
    if (claims.length === 0 && instances.size === 0) {
      return undefined;
    }
    const invocations = primary
      ? primary.createInvocationScope(registry)
      : new PluginInvocationScope(registry, instances, {
          // A managed factory owns a logical consumer even on a caller-owned root view.
          retained:
            registration !== undefined &&
            getPluginValueInstance(registration.factory) !== undefined,
        });
    return new ContextEngineFactoryResources(claims, invocations);
  } catch (error) {
    if (claims.length > 0) {
      abandon(new ContextEngineFactoryResources(claims));
    }
    throw error;
  }
}

/** One instance may come from both factories; every source stays held through shared cleanup. */
export async function disposeContextEngineSources(
  engine: ContextEngine | undefined,
  sources: readonly ContextEngineFactoryResources[],
  dispose: () => void | Promise<void> = () => engine?.dispose?.(),
): Promise<void> {
  if (sources.length === 0) {
    await dispose();
    return;
  }
  for (const source of sources) {
    source.beginCleanup();
  }
  // Start instance cleanup before retiring the factory lifetime that it may need to stop.
  const cleanup = sources[0]!.cleanupWork.track(() => sources[0]!.runCleanup(dispose));
  for (const source of sources) {
    source.context(() => source.work.beginClose());
  }
  const [engineResult] = await Promise.allSettled([cleanup]);
  const scopes = sources.flatMap((source) => [source.work, source.cleanupWork]);
  await AsyncWorkScope.runWhenAllIdle(
    () => scopes,
    () => {
      for (const source of sources) {
        source.cleanupContext(() => source.cleanupWork.beginClose());
      }
    },
  );
  await AsyncWorkScope.runWhenAllIdle(
    () => scopes,
    () =>
      Promise.all(
        sources.flatMap((source) => [
          source.context(() => source.work.drain()),
          source.cleanupContext(() => source.cleanupWork.drain()),
        ]),
      ),
  );
  const released = await Promise.allSettled(sources.map((source) => source.release()));
  if (engineResult.status === "rejected") {
    throw engineResult.reason;
  }
  for (const result of released) {
    if (result.status === "rejected") {
      throw result.reason;
    }
  }
}

type ContextEngineFactoryFailureCleanup = (
  source: ContextEngineFactoryResources | undefined,
) => void;

/** Report selection separately while its admitted owner joins failed-factory cleanup. */
export async function runContextEngineFactoryResolution<T>(
  resolve: (abandon: ContextEngineFactoryFailureCleanup) => Promise<T>,
  onCleanupFailure?: () => void,
): Promise<T> {
  const reportCleanupFailure = onCleanupFailure
    ? AsyncLocalStorage.bind(onCleanupFailure)
    : undefined;
  const result = createDeferredCore<T>();
  const failedSources: Promise<void>[] = [];
  const abandon: ContextEngineFactoryFailureCleanup = (source) => {
    if (source) {
      failedSources.push(
        disposeContextEngineSources(undefined, [source]).catch((error: unknown) => {
          reportCleanupFailure?.();
          console.warn(
            `[context-engine] Failed factory resource cleanup: ${sanitizeForLog(String(error))}`,
          );
        }),
      );
    }
  };
  // Admit construction before either factory can escape into asynchronous work.
  void captureAsyncWorkTracker()(async () => {
    try {
      result.resolve(await resolve(abandon));
    } catch (error) {
      result.reject(error);
    } finally {
      await Promise.all(failedSources);
    }
  }).catch(result.reject);
  return await result.promise;
}

export function retainLogicalTurnContextEngineSources(
  registry: PluginRegistry,
  fallback: ContextEngineRegistration | undefined,
  configured: ContextEngineRegistration | undefined,
  abandon: ContextEngineFactoryFailureCleanup,
): {
  fallback: ContextEngineFactoryResources | undefined;
  configured?: ContextEngineFactoryResources;
  configuredFailure?: { error: unknown };
} {
  const primary = getPluginRegistryInspectionResources(registry);
  const fallbackSource = retainContextEngineFactorySource(registry, fallback, primary, abandon);
  try {
    const configuredSource =
      configured?.lifecycle === "runtime"
        ? retainContextEngineFactorySource(registry, configured, primary, abandon)
        : undefined;
    return { fallback: fallbackSource, configured: configuredSource };
  } catch (error) {
    return { fallback: fallbackSource, configuredFailure: { error } };
  }
}

export async function resolveContextEngineFactory<T extends { engine: ContextEngine }>(
  source: ContextEngineFactoryResources | undefined,
  owners: Map<ContextEngine, ContextEngineFactoryResources[]>,
  create: () => Promise<T>,
): Promise<T> {
  const ref = await (source ? source.run(create) : create());
  if (source) {
    const retained = owners.get(ref.engine) ?? [];
    retained.push(source);
    owners.set(ref.engine, retained);
  }
  return ref;
}

/** Foreground engine resolution shares the same factory execution and physical resource owner. */
export async function createContextEngineWithResources<T>(
  registry: PluginRegistry,
  registration: ContextEngineRegistration,
  create: (source: ContextEngineFactoryResources | undefined) => Promise<T>,
): Promise<T> {
  return await runContextEngineFactoryResolution(async (abandon) => {
    const source = retainContextEngineFactorySource(
      registry,
      registration,
      getPluginRegistryInspectionResources(registry),
      abandon,
    );
    try {
      return await (source ? source.run(() => create(source)) : create(undefined));
    } catch (error) {
      abandon(source);
      throw error;
    }
  });
}