refactor: move memory plugin state into plugin host

src/agents/system-prompt.ts

@@ -2,7 +2,7 @@ import { createHmac, createHash } from "node:crypto";
 import type { ReasoningLevel, ThinkLevel } from "../auto-reply/thinking.js";
 import { SILENT_REPLY_TOKEN } from "../auto-reply/tokens.js";
 import type { MemoryCitationsMode } from "../config/types.memory.js";
-import { buildMemoryPromptSection } from "../memory/prompt-section.js";
+import { buildMemoryPromptSection } from "../plugins/memory-state.js";
 import { listDeliverableMessageChannels } from "../utils/message-channel.js";
 import type { ResolvedTimeFormat } from "./date-time.js";
 import type { EmbeddedContextFile } from "./pi-embedded-helpers.js";

src/auto-reply/reply/agent-runner-memory.ts

@@ -25,7 +25,7 @@ import {
 import { readSessionMessages } from "../../gateway/session-utils.fs.js";
 import { logVerbose } from "../../globals.js";
 import { registerAgentRunContext } from "../../infra/agent-events.js";
-import { resolveMemoryFlushPlan } from "../../memory/flush-plan.js";
+import { resolveMemoryFlushPlan } from "../../plugins/memory-state.js";
 import type { TemplateContext } from "../templating.js";
 import type { VerboseLevel } from "../thinking.js";
 import type { GetReplyOptions } from "../types.js";

src/memory/flush-plan.test.ts

@@ -1,49 +0,0 @@
-import { afterEach, describe, expect, it } from "vitest";
-import {
-  clearMemoryFlushPlanResolver,
-  getMemoryFlushPlanResolver,
-  registerMemoryFlushPlanResolver,
-  resolveMemoryFlushPlan,
-  restoreMemoryFlushPlanResolver,
-} from "./flush-plan.js";
-
-describe("memory flush plan registry", () => {
-  afterEach(() => {
-    clearMemoryFlushPlanResolver();
-  });
-
-  it("returns null when no resolver is registered", () => {
-    expect(resolveMemoryFlushPlan({})).toBeNull();
-  });
-
-  it("uses the registered resolver", () => {
-    registerMemoryFlushPlanResolver(() => ({
-      softThresholdTokens: 1,
-      forceFlushTranscriptBytes: 2,
-      reserveTokensFloor: 3,
-      prompt: "prompt",
-      systemPrompt: "system",
-      relativePath: "memory/test.md",
-    }));
-
-    expect(resolveMemoryFlushPlan({})?.relativePath).toBe("memory/test.md");
-  });
-
-  it("restoreMemoryFlushPlanResolver swaps resolver state", () => {
-    registerMemoryFlushPlanResolver(() => ({
-      softThresholdTokens: 1,
-      forceFlushTranscriptBytes: 2,
-      reserveTokensFloor: 3,
-      prompt: "first",
-      systemPrompt: "first",
-      relativePath: "memory/first.md",
-    }));
-    const current = getMemoryFlushPlanResolver();
-
-    clearMemoryFlushPlanResolver();
-    expect(resolveMemoryFlushPlan({})).toBeNull();
-
-    restoreMemoryFlushPlanResolver(current);
-    expect(resolveMemoryFlushPlan({})?.relativePath).toBe("memory/first.md");
-  });
-});

src/memory/flush-plan.ts

@@ -1,42 +0,0 @@
-import type { OpenClawConfig } from "../config/config.js";
-
-export type MemoryFlushPlan = {
-  softThresholdTokens: number;
-  forceFlushTranscriptBytes: number;
-  reserveTokensFloor: number;
-  prompt: string;
-  systemPrompt: string;
-  relativePath: string;
-};
-
-export type MemoryFlushPlanResolver = (params: {
-  cfg?: OpenClawConfig;
-  nowMs?: number;
-}) => MemoryFlushPlan | null;
-
-let _resolver: MemoryFlushPlanResolver | undefined;
-
-export function registerMemoryFlushPlanResolver(resolver: MemoryFlushPlanResolver): void {
-  _resolver = resolver;
-}
-
-export function resolveMemoryFlushPlan(params: {
-  cfg?: OpenClawConfig;
-  nowMs?: number;
-}): MemoryFlushPlan | null {
-  return _resolver?.(params) ?? null;
-}
-
-export function getMemoryFlushPlanResolver(): MemoryFlushPlanResolver | undefined {
-  return _resolver;
-}
-
-export function restoreMemoryFlushPlanResolver(
-  resolver: MemoryFlushPlanResolver | undefined,
-): void {
-  _resolver = resolver;
-}
-
-export function clearMemoryFlushPlanResolver(): void {
-  _resolver = undefined;
-}

src/memory/prompt-section.test.ts

@@ -1,64 +0,0 @@
-import { describe, it, expect, beforeEach } from "vitest";
-import {
-  registerMemoryPromptSection,
-  buildMemoryPromptSection,
-  clearMemoryPromptSection,
-  _resetMemoryPromptSection,
-} from "./prompt-section.js";
-
-describe("memory prompt section registry", () => {
-  beforeEach(() => {
-    _resetMemoryPromptSection();
-  });
-
-  it("returns empty array when no builder is registered", () => {
-    const result = buildMemoryPromptSection({
-      availableTools: new Set(["memory_search", "memory_get"]),
-    });
-    expect(result).toEqual([]);
-  });
-
-  it("delegates to the registered builder", () => {
-    registerMemoryPromptSection(({ availableTools }) => {
-      if (!availableTools.has("memory_search")) {
-        return [];
-      }
-      return ["## Custom Memory", "Use custom memory tools.", ""];
-    });
-
-    const result = buildMemoryPromptSection({
-      availableTools: new Set(["memory_search"]),
-    });
-    expect(result).toEqual(["## Custom Memory", "Use custom memory tools.", ""]);
-  });
-
-  it("passes citationsMode to the builder", () => {
-    registerMemoryPromptSection(({ citationsMode }) => {
-      return [`citations: ${citationsMode ?? "default"}`];
-    });
-
-    expect(
-      buildMemoryPromptSection({
-        availableTools: new Set(),
-        citationsMode: "off",
-      }),
-    ).toEqual(["citations: off"]);
-  });
-
-  it("last registration wins", () => {
-    registerMemoryPromptSection(() => ["first"]);
-    registerMemoryPromptSection(() => ["second"]);
-
-    const result = buildMemoryPromptSection({ availableTools: new Set() });
-    expect(result).toEqual(["second"]);
-  });
-
-  it("clearMemoryPromptSection resets the builder", () => {
-    registerMemoryPromptSection(() => ["stale section"]);
-    expect(buildMemoryPromptSection({ availableTools: new Set() })).toEqual(["stale section"]);
-
-    clearMemoryPromptSection();
-
-    expect(buildMemoryPromptSection({ availableTools: new Set() })).toEqual([]);
-  });
-});

src/memory/prompt-section.ts

@@ -1,42 +0,0 @@
-import type { MemoryCitationsMode } from "../config/types.memory.js";
-
-/**
- * Callback that the active memory plugin provides to build
- * its section of the agent system prompt.
- */
-export type MemoryPromptSectionBuilder = (params: {
-  availableTools: Set<string>;
-  citationsMode?: MemoryCitationsMode;
-}) => string[];
-
-// Module-level singleton — only one memory plugin can be active (exclusive slot).
-let _builder: MemoryPromptSectionBuilder | undefined;
-
-export function registerMemoryPromptSection(builder: MemoryPromptSectionBuilder): void {
-  _builder = builder;
-}
-
-export function buildMemoryPromptSection(params: {
-  availableTools: Set<string>;
-  citationsMode?: MemoryCitationsMode;
-}): string[] {
-  return _builder?.(params) ?? [];
-}
-
-/** Return the current builder (used by the plugin cache to snapshot state). */
-export function getMemoryPromptSectionBuilder(): MemoryPromptSectionBuilder | undefined {
-  return _builder;
-}
-
-/** Restore a previously-snapshotted builder (used on plugin cache hits). */
-export function restoreMemoryPromptSection(builder: MemoryPromptSectionBuilder | undefined): void {
-  _builder = builder;
-}
-
-/** Clear the registered builder (called on plugin reload and in tests). */
-export function clearMemoryPromptSection(): void {
-  _builder = undefined;
-}
-
-/** @deprecated Use {@link clearMemoryPromptSection}. */
-export const _resetMemoryPromptSection = clearMemoryPromptSection;

src/plugin-sdk/memory-core.ts

@@ -30,6 +30,9 @@ export { shortenHomeInString, shortenHomePath } from "../utils.js";
 export type { OpenClawConfig } from "../config/config.js";
 export type { MemoryCitationsMode } from "../config/types.memory.js";
 export type { MemorySearchResult } from "../memory/types.js";
-export type { MemoryFlushPlan, MemoryFlushPlanResolver } from "../memory/flush-plan.js";
-export type { MemoryPromptSectionBuilder } from "../memory/prompt-section.js";
+export type {
+  MemoryFlushPlan,
+  MemoryFlushPlanResolver,
+  MemoryPromptSectionBuilder,
+} from "../plugins/memory-state.js";
 export type { OpenClawPluginApi } from "../plugins/types.js";

src/plugins/loader.test.ts

@@ -3,8 +3,6 @@ import os from "node:os";
 import path from "node:path";
 import { afterAll, afterEach, describe, expect, it } from "vitest";
 import { emitDiagnosticEvent, resetDiagnosticEventsForTest } from "../infra/diagnostic-events.js";
-import { registerMemoryFlushPlanResolver, resolveMemoryFlushPlan } from "../memory/flush-plan.js";
-import { buildMemoryPromptSection, registerMemoryPromptSection } from "../memory/prompt-section.js";
 import { withEnv } from "../test-utils/env.js";
 import { clearPluginCommands, getPluginCommandSpecs } from "./command-registry-state.js";
 import { clearPluginDiscoveryCache } from "./discovery.js";
@@ -12,6 +10,12 @@ import { getGlobalHookRunner, resetGlobalHookRunner } from "./hook-runner-global
 import { createHookRunner } from "./hooks.js";
 import { __testing, clearPluginLoaderCache, loadOpenClawPlugins } from "./loader.js";
 import { clearPluginManifestRegistryCache } from "./manifest-registry.js";
+import {
+  buildMemoryPromptSection,
+  registerMemoryFlushPlanResolver,
+  registerMemoryPromptSection,
+  resolveMemoryFlushPlan,
+} from "./memory-state.js";
 import { createEmptyPluginRegistry } from "./registry.js";
 import {
   getActivePluginRegistry,

src/plugins/loader.ts

@@ -8,16 +8,6 @@ import type { PluginInstallRecord } from "../config/types.plugins.js";
 import type { GatewayRequestHandler } from "../gateway/server-methods/types.js";
 import { openBoundaryFileSync } from "../infra/boundary-file-read.js";
 import { createSubsystemLogger } from "../logging/subsystem.js";
-import {
-  clearMemoryFlushPlanResolver,
-  getMemoryFlushPlanResolver,
-  restoreMemoryFlushPlanResolver,
-} from "../memory/flush-plan.js";
-import {
-  clearMemoryPromptSection,
-  getMemoryPromptSectionBuilder,
-  restoreMemoryPromptSection,
-} from "../memory/prompt-section.js";
 import { resolveUserPath } from "../utils.js";
 import { inspectBundleMcpRuntimeSupport } from "./bundle-mcp.js";
 import { clearPluginCommands } from "./command-registry-state.js";
@@ -32,6 +22,12 @@ import { discoverOpenClawPlugins } from "./discovery.js";
 import { initializeGlobalHookRunner } from "./hook-runner-global.js";
 import { clearPluginInteractiveHandlers } from "./interactive.js";
 import { loadPluginManifestRegistry } from "./manifest-registry.js";
+import {
+  clearMemoryPluginState,
+  getMemoryFlushPlanResolver,
+  getMemoryPromptSectionBuilder,
+  restoreMemoryPluginState,
+} from "./memory-state.js";
 import { isPathInside, safeStatSync } from "./path-safety.js";
 import { createPluginRegistry, type PluginRecord, type PluginRegistry } from "./registry.js";
 import { resolvePluginCacheInputs } from "./roots.js";
@@ -129,8 +125,7 @@ const LAZY_RUNTIME_REFLECTION_KEYS = [
 export function clearPluginLoaderCache(): void {
   registryCache.clear();
   openAllowlistWarningCache.clear();
-  clearMemoryFlushPlanResolver();
-  clearMemoryPromptSection();
+  clearMemoryPluginState();
 }
 
 const defaultLogger = () => createSubsystemLogger("plugins");
@@ -714,8 +709,10 @@ export function loadOpenClawPlugins(options: PluginLoadOptions = {}): PluginRegi
   if (cacheEnabled) {
     const cached = getCachedPluginRegistry(cacheKey);
     if (cached) {
-      restoreMemoryPromptSection(cached.memoryPromptBuilder);
-      restoreMemoryFlushPlanResolver(cached.memoryFlushPlanResolver);
+      restoreMemoryPluginState({
+        promptBuilder: cached.memoryPromptBuilder,
+        flushPlanResolver: cached.memoryFlushPlanResolver,
+      });
       if (shouldActivate) {
         activatePluginRegistry(cached.registry, cacheKey);
       }
@@ -728,8 +725,7 @@ export function loadOpenClawPlugins(options: PluginLoadOptions = {}): PluginRegi
   if (shouldActivate) {
     clearPluginCommands();
     clearPluginInteractiveHandlers();
-    clearMemoryFlushPlanResolver();
-    clearMemoryPromptSection();
+    clearMemoryPluginState();
   }
 
   // Lazy: avoid creating the Jiti loader when all plugins are disabled (common in unit tests).
@@ -1242,14 +1238,18 @@ export function loadOpenClawPlugins(options: PluginLoadOptions = {}): PluginRegi
       }
       // Snapshot loads should not replace process-global runtime prompt state.
       if (!shouldActivate) {
-        restoreMemoryPromptSection(previousMemoryPromptBuilder);
-        restoreMemoryFlushPlanResolver(previousMemoryFlushPlanResolver);
+        restoreMemoryPluginState({
+          promptBuilder: previousMemoryPromptBuilder,
+          flushPlanResolver: previousMemoryFlushPlanResolver,
+        });
       }
       registry.plugins.push(record);
       seenIds.set(pluginId, candidate.origin);
     } catch (err) {
-      restoreMemoryPromptSection(previousMemoryPromptBuilder);
-      restoreMemoryFlushPlanResolver(previousMemoryFlushPlanResolver);
+      restoreMemoryPluginState({
+        promptBuilder: previousMemoryPromptBuilder,
+        flushPlanResolver: previousMemoryFlushPlanResolver,
+      });
       recordPluginError({
         logger,
         registry,

src/plugins/memory-state.test.ts

@@ -0,0 +1,105 @@
+import { afterEach, describe, expect, it } from "vitest";
+import {
+  _resetMemoryPluginState,
+  buildMemoryPromptSection,
+  clearMemoryPluginState,
+  getMemoryFlushPlanResolver,
+  getMemoryPromptSectionBuilder,
+  registerMemoryFlushPlanResolver,
+  registerMemoryPromptSection,
+  resolveMemoryFlushPlan,
+  restoreMemoryPluginState,
+} from "./memory-state.js";
+
+describe("memory plugin state", () => {
+  afterEach(() => {
+    clearMemoryPluginState();
+  });
+
+  it("returns empty defaults when no memory plugin state is registered", () => {
+    expect(resolveMemoryFlushPlan({})).toBeNull();
+    expect(buildMemoryPromptSection({ availableTools: new Set(["memory_search"]) })).toEqual([]);
+  });
+
+  it("delegates prompt building to the registered memory plugin", () => {
+    registerMemoryPromptSection(({ availableTools }) => {
+      if (!availableTools.has("memory_search")) {
+        return [];
+      }
+      return ["## Custom Memory", "Use custom memory tools.", ""];
+    });
+
+    expect(buildMemoryPromptSection({ availableTools: new Set(["memory_search"]) })).toEqual([
+      "## Custom Memory",
+      "Use custom memory tools.",
+      "",
+    ]);
+  });
+
+  it("passes citations mode through to the prompt builder", () => {
+    registerMemoryPromptSection(({ citationsMode }) => [
+      `citations: ${citationsMode ?? "default"}`,
+    ]);
+
+    expect(
+      buildMemoryPromptSection({
+        availableTools: new Set(),
+        citationsMode: "off",
+      }),
+    ).toEqual(["citations: off"]);
+  });
+
+  it("uses the registered flush plan resolver", () => {
+    registerMemoryFlushPlanResolver(() => ({
+      softThresholdTokens: 1,
+      forceFlushTranscriptBytes: 2,
+      reserveTokensFloor: 3,
+      prompt: "prompt",
+      systemPrompt: "system",
+      relativePath: "memory/test.md",
+    }));
+
+    expect(resolveMemoryFlushPlan({})?.relativePath).toBe("memory/test.md");
+  });
+
+  it("restoreMemoryPluginState swaps both prompt and flush state", () => {
+    registerMemoryPromptSection(() => ["first"]);
+    registerMemoryFlushPlanResolver(() => ({
+      softThresholdTokens: 1,
+      forceFlushTranscriptBytes: 2,
+      reserveTokensFloor: 3,
+      prompt: "first",
+      systemPrompt: "first",
+      relativePath: "memory/first.md",
+    }));
+    const snapshot = {
+      promptBuilder: getMemoryPromptSectionBuilder(),
+      flushPlanResolver: getMemoryFlushPlanResolver(),
+    };
+
+    _resetMemoryPluginState();
+    expect(buildMemoryPromptSection({ availableTools: new Set() })).toEqual([]);
+    expect(resolveMemoryFlushPlan({})).toBeNull();
+
+    restoreMemoryPluginState(snapshot);
+    expect(buildMemoryPromptSection({ availableTools: new Set() })).toEqual(["first"]);
+    expect(resolveMemoryFlushPlan({})?.relativePath).toBe("memory/first.md");
+  });
+
+  it("clearMemoryPluginState resets both registries", () => {
+    registerMemoryPromptSection(() => ["stale section"]);
+    registerMemoryFlushPlanResolver(() => ({
+      softThresholdTokens: 1,
+      forceFlushTranscriptBytes: 2,
+      reserveTokensFloor: 3,
+      prompt: "prompt",
+      systemPrompt: "system",
+      relativePath: "memory/stale.md",
+    }));
+
+    clearMemoryPluginState();
+
+    expect(buildMemoryPromptSection({ availableTools: new Set() })).toEqual([]);
+    expect(resolveMemoryFlushPlan({})).toBeNull();
+  });
+});

src/plugins/memory-state.ts

@@ -0,0 +1,70 @@
+import type { OpenClawConfig } from "../config/config.js";
+import type { MemoryCitationsMode } from "../config/types.memory.js";
+
+export type MemoryPromptSectionBuilder = (params: {
+  availableTools: Set<string>;
+  citationsMode?: MemoryCitationsMode;
+}) => string[];
+
+export type MemoryFlushPlan = {
+  softThresholdTokens: number;
+  forceFlushTranscriptBytes: number;
+  reserveTokensFloor: number;
+  prompt: string;
+  systemPrompt: string;
+  relativePath: string;
+};
+
+export type MemoryFlushPlanResolver = (params: {
+  cfg?: OpenClawConfig;
+  nowMs?: number;
+}) => MemoryFlushPlan | null;
+
+type MemoryPluginState = {
+  promptBuilder?: MemoryPromptSectionBuilder;
+  flushPlanResolver?: MemoryFlushPlanResolver;
+};
+
+const memoryPluginState: MemoryPluginState = {};
+
+export function registerMemoryPromptSection(builder: MemoryPromptSectionBuilder): void {
+  memoryPluginState.promptBuilder = builder;
+}
+
+export function buildMemoryPromptSection(params: {
+  availableTools: Set<string>;
+  citationsMode?: MemoryCitationsMode;
+}): string[] {
+  return memoryPluginState.promptBuilder?.(params) ?? [];
+}
+
+export function getMemoryPromptSectionBuilder(): MemoryPromptSectionBuilder | undefined {
+  return memoryPluginState.promptBuilder;
+}
+
+export function registerMemoryFlushPlanResolver(resolver: MemoryFlushPlanResolver): void {
+  memoryPluginState.flushPlanResolver = resolver;
+}
+
+export function resolveMemoryFlushPlan(params: {
+  cfg?: OpenClawConfig;
+  nowMs?: number;
+}): MemoryFlushPlan | null {
+  return memoryPluginState.flushPlanResolver?.(params) ?? null;
+}
+
+export function getMemoryFlushPlanResolver(): MemoryFlushPlanResolver | undefined {
+  return memoryPluginState.flushPlanResolver;
+}
+
+export function restoreMemoryPluginState(state: MemoryPluginState): void {
+  memoryPluginState.promptBuilder = state.promptBuilder;
+  memoryPluginState.flushPlanResolver = state.flushPlanResolver;
+}
+
+export function clearMemoryPluginState(): void {
+  memoryPluginState.promptBuilder = undefined;
+  memoryPluginState.flushPlanResolver = undefined;
+}
+
+export const _resetMemoryPluginState = clearMemoryPluginState;

src/plugins/registry.ts

@@ -8,13 +8,12 @@ import type {
 } from "../gateway/server-methods/types.js";
 import { registerInternalHook } from "../hooks/internal-hooks.js";
 import type { HookEntry } from "../hooks/types.js";
-import { registerMemoryFlushPlanResolver } from "../memory/flush-plan.js";
-import { registerMemoryPromptSection } from "../memory/prompt-section.js";
 import { resolveUserPath } from "../utils.js";
 import { registerPluginCommand, validatePluginCommandDefinition } from "./command-registration.js";
 import { normalizePluginHttpPath } from "./http-path.js";
 import { findOverlappingPluginHttpRoute } from "./http-route-overlap.js";
 import { registerPluginInteractiveHandler } from "./interactive.js";
+import { registerMemoryFlushPlanResolver, registerMemoryPromptSection } from "./memory-state.js";
 import { normalizeRegisteredProvider } from "./provider-validation.js";
 import { createEmptyPluginRegistry } from "./registry-empty.js";
 import { withPluginRuntimePluginIdScope } from "./runtime/gateway-request-scope.js";

src/plugins/types.ts

@@ -1407,12 +1407,10 @@ export type OpenClawPluginApi = {
   ) => void;
   /** Register the system prompt section builder for this memory plugin (exclusive slot). */
   registerMemoryPromptSection: (
-    builder: import("../memory/prompt-section.js").MemoryPromptSectionBuilder,
+    builder: import("./memory-state.js").MemoryPromptSectionBuilder,
   ) => void;
   /** Register the pre-compaction flush plan resolver for this memory plugin (exclusive slot). */
-  registerMemoryFlushPlan: (
-    resolver: import("../memory/flush-plan.js").MemoryFlushPlanResolver,
-  ) => void;
+  registerMemoryFlushPlan: (resolver: import("./memory-state.js").MemoryFlushPlanResolver) => void;
   resolvePath: (input: string) => string;
   /** Register a lifecycle hook handler */
   on: <K extends PluginHookName>(