fix(plugin-sdk): harden diagnostic alias cleanup

Kaspre · Kaspre · commit 8da195379c28 · 2026-05-26T20:37:18.000-04:00
diff --git a/src/plugin-sdk/root-alias.cjs b/src/plugin-sdk/root-alias.cjs
@@ -147,14 +147,55 @@ function onDiagnosticEventFromSharedState(listener) {
   };
 }
 
+function snapshotDiagnosticListeners(state) {
+  return state && state.listeners instanceof Set ? new Set(state.listeners) : null;
+}
+
+function removeAddedDiagnosticListeners(beforeListeners) {
+  const state = getDiagnosticEventsState(false);
+  if (!state || !(state.listeners instanceof Set)) {
+    return;
+  }
+  if (!beforeListeners) {
+    state.listeners.clear();
+    return;
+  }
+  for (const listener of state.listeners) {
+    if (!beforeListeners.has(listener)) {
+      state.listeners.delete(listener);
+    }
+  }
+}
+
+function trySubscribeDiagnosticEvents(diagnosticEvents, listener, beforeListeners) {
+  try {
+    const unsubscribe = diagnosticEvents.onDiagnosticEvent(listener);
+    if (typeof unsubscribe === "function") {
+      return unsubscribe;
+    }
+  } catch {
+    // Fall back to shared state if a stale dist chunk exposes a broken wrapper.
+  }
+  removeAddedDiagnosticListeners(beforeListeners);
+  return null;
+}
+
 function onDiagnosticEvent(listener) {
   const beforeState = getDiagnosticEventsState(false);
+  const beforeListeners = snapshotDiagnosticListeners(beforeState);
   const beforeSize = beforeState?.listeners?.size;
   const diagnosticEvents = loadDiagnosticEventsModule();
   if (!diagnosticEvents || typeof diagnosticEvents.onDiagnosticEvent !== "function") {
     return onDiagnosticEventFromSharedState(listener);
   }
-  const unsubscribeDiagnosticEvents = diagnosticEvents.onDiagnosticEvent(listener);
+  const unsubscribeDiagnosticEvents = trySubscribeDiagnosticEvents(
+    diagnosticEvents,
+    listener,
+    beforeListeners,
+  );
+  if (!unsubscribeDiagnosticEvents) {
+    return onDiagnosticEventFromSharedState(listener);
+  }
   const afterState = getDiagnosticEventsState(false);
   if (afterState && afterState.listeners.size > (beforeSize ?? 0)) {
     return unsubscribeDiagnosticEvents;
@@ -163,8 +204,11 @@ function onDiagnosticEvent(listener) {
   // diagnostic module in a separate graph from the active core emitter.
   const unsubscribeSharedState = onDiagnosticEventFromSharedState(listener);
   return () => {
-    unsubscribeDiagnosticEvents();
-    unsubscribeSharedState();
+    try {
+      unsubscribeDiagnosticEvents();
+    } finally {
+      unsubscribeSharedState();
+    }
   };
 }
 
diff --git a/src/plugins/contracts/plugin-sdk-root-alias.test.ts b/src/plugins/contracts/plugin-sdk-root-alias.test.ts
@@ -11,6 +11,7 @@ const rootSdk = require(rootAliasPath) as Record<string, unknown>;
 const rootAliasSource = fs.readFileSync(rootAliasPath, "utf-8");
 const compatPath = fileURLToPath(new URL("../../plugin-sdk/compat.ts", import.meta.url));
 const packageJsonPath = fileURLToPath(new URL("../../../package.json", import.meta.url));
+const diagnosticEventsStateKey = Symbol.for("openclaw.diagnosticEvents.state.v1");
 const legacyRootExportNames = [
   "registerContextEngine",
   "buildMemorySystemPromptAddition",
@@ -36,6 +37,10 @@ type EmptySchema = {
       };
 };
 
+type DiagnosticEventsStateFixture = {
+  listeners: Set<(event: { type: string }, metadata: { trusted: boolean }) => void>;
+};
+
 function requirePropertyDescriptor(
   target: Record<string, unknown>,
   propertyName: string,
@@ -183,6 +188,56 @@ function loadDiagnosticEventsAlias(distEntries: string[]) {
   });
 }
 
+function ensureDiagnosticEventsStateFixture(
+  context: Record<PropertyKey, unknown>,
+): DiagnosticEventsStateFixture {
+  const existing = context[diagnosticEventsStateKey] as DiagnosticEventsStateFixture | undefined;
+  if (existing) {
+    return existing;
+  }
+  const state = vm.runInNewContext(
+    `({
+      marker: Symbol.for("openclaw.diagnosticEvents.state.v1"),
+      enabled: true,
+      seq: 0,
+      listeners: new Set(),
+      dispatchDepth: 0,
+      asyncQueue: [],
+      asyncDrainScheduled: false,
+      asyncDroppedEvents: 0,
+      asyncDroppedTrustedEvents: 0,
+      asyncDroppedUntrustedEvents: 0,
+      asyncDroppedPriorityEvents: 0,
+    })`,
+    context,
+  ) as DiagnosticEventsStateFixture;
+  Object.defineProperty(context, diagnosticEventsStateKey, {
+    configurable: true,
+    enumerable: false,
+    value: state,
+    writable: false,
+  });
+  return state;
+}
+
+function requireDiagnosticEventsStateFixture(
+  lazyModule: ReturnType<typeof loadRootAliasWithStubs>,
+): DiagnosticEventsStateFixture {
+  const state = lazyModule.globalContext[diagnosticEventsStateKey] as
+    | DiagnosticEventsStateFixture
+    | undefined;
+  if (!state) {
+    throw new Error("expected diagnostic events state fixture");
+  }
+  return state;
+}
+
+function emitFixtureDiagnosticEvent(state: DiagnosticEventsStateFixture): void {
+  for (const registered of state.listeners) {
+    registered({ type: "model.usage" }, { trusted: false });
+  }
+}
+
 function expectDiagnosticEventAccessor(lazyModule: ReturnType<typeof loadRootAliasWithStubs>) {
   expect(
     typeof (lazyModule.moduleExports.onDiagnosticEvent as (listener: () => void) => () => void)(
@@ -555,6 +610,106 @@ describe("plugin-sdk root alias", () => {
     expect(unsubscribeCount).toBe(1);
   });
 
+  it("falls back and removes stale diagnostic listeners when the dist subscription is invalid", () => {
+    const seen: string[] = [];
+    let lazyModule!: ReturnType<typeof loadRootAliasWithStubs>;
+    const preexistingListener = (): void => undefined;
+    lazyModule = loadRootAliasWithStubs({
+      aliasPath: createDistAliasPath(),
+      distEntries: ["diagnostic-events-W3Hz61fI.js"],
+      monolithicExports: {
+        onDiagnosticEvent(listener: (event: { type: string }) => void): undefined {
+          const state = ensureDiagnosticEventsStateFixture(lazyModule.globalContext);
+          state.listeners.add((event, metadata) => {
+            if (!metadata.trusted) {
+              listener(event);
+            }
+          });
+          return undefined;
+        },
+      },
+    });
+    const state = ensureDiagnosticEventsStateFixture(lazyModule.globalContext);
+    state.listeners.add(preexistingListener);
+
+    const unsubscribe = (
+      lazyModule.moduleExports.onDiagnosticEvent as (
+        listener: (event: { type: string }) => void,
+      ) => () => void
+    )((event) => {
+      seen.push(event.type);
+    });
+
+    expect(state.listeners.size).toBe(2);
+    expect(state.listeners.has(preexistingListener)).toBe(true);
+    emitFixtureDiagnosticEvent(state);
+    unsubscribe();
+
+    expect(seen).toEqual(["model.usage"]);
+    expect(state.listeners.size).toBe(1);
+    expect(state.listeners.has(preexistingListener)).toBe(true);
+  });
+
+  it("falls back to shared diagnostic state when the dist subscription throws", () => {
+    const seen: string[] = [];
+    let subscribeCount = 0;
+    const lazyModule = loadRootAliasWithStubs({
+      aliasPath: createDistAliasPath(),
+      distEntries: ["diagnostic-events-W3Hz61fI.js"],
+      monolithicExports: {
+        onDiagnosticEvent(): never {
+          subscribeCount += 1;
+          throw new Error("stale diagnostic subscription");
+        },
+      },
+    });
+
+    const unsubscribe = (
+      lazyModule.moduleExports.onDiagnosticEvent as (
+        listener: (event: { type: string }) => void,
+      ) => () => void
+    )((event) => {
+      seen.push(event.type);
+    });
+    const state = requireDiagnosticEventsStateFixture(lazyModule);
+
+    expect(subscribeCount).toBe(1);
+    expect(state.listeners.size).toBe(1);
+    emitFixtureDiagnosticEvent(state);
+    unsubscribe();
+
+    expect(seen).toEqual(["model.usage"]);
+    expect(state.listeners.size).toBe(0);
+  });
+
+  it("removes the shared-state fallback listener when diagnostic cleanup throws", () => {
+    let diagnosticUnsubscribeCount = 0;
+    const lazyModule = loadRootAliasWithStubs({
+      aliasPath: createDistAliasPath(),
+      distEntries: ["diagnostic-events-W3Hz61fI.js"],
+      monolithicExports: {
+        onDiagnosticEvent(): () => void {
+          return () => {
+            diagnosticUnsubscribeCount += 1;
+            throw new Error("diagnostic cleanup failed");
+          };
+        },
+      },
+    });
+
+    const unsubscribe = (
+      lazyModule.moduleExports.onDiagnosticEvent as (
+        listener: (event: { type: string }) => void,
+      ) => () => void
+    )(() => undefined);
+    const state = requireDiagnosticEventsStateFixture(lazyModule);
+
+    expect(state.listeners.size).toBe(1);
+    expect(() => unsubscribe()).toThrow("diagnostic cleanup failed");
+    expect(diagnosticUnsubscribeCount).toBe(1);
+    expect(state.listeners.size).toBe(0);
+  });
+
   it("bridges diagnostic listeners through shared process state when the lazy module is isolated", () => {
     const seen: string[] = [];
     const lazyModule = loadDiagnosticEventsAlias(["diagnostic-events-W3Hz61fI.js"]);
