kongbrain 0.1.4 → 0.2.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "kongbrain",
-  "version": "0.1.4",
+  "version": "0.2.0",
   "description": "Graph-backed persistent memory engine for OpenClaw. Replaces the default context window with SurrealDB + vector embeddings that learn across sessions.",
   "type": "module",
   "license": "MIT",

package/src/causal.ts

@@ -86,7 +86,7 @@ export async function linkCausalEdges(
       // Store chain metadata
       await store.queryExec(`CREATE causal_chain CONTENT $data`, {
         data: {
-          session_id: sessionId,
+          session_id: String(sessionId),
           trigger_memory: triggerId,
           outcome_memory: outcomeId,
           description_memory: descriptionId,

package/src/context-engine.ts

@@ -45,6 +45,7 @@ import { evaluateRetrieval, getStagedItems } from "./retrieval-quality.js";
 import { shouldRunCheck, runCognitiveCheck } from "./cognitive-check.js";
 import { checkACANReadiness } from "./acan.js";
 import { predictQueries, prefetchContext } from "./prefetch.js";
+import { runDeferredCleanup } from "./deferred-cleanup.js";
 import { swallow } from "./errors.js";
 
 const __dirname = dirname(fileURLToPath(import.meta.url));
@@ -122,6 +123,7 @@ export class KongBrainContextEngine implements ContextEngine {
       store.consolidateMemories((text) => embeddings.embed(text)),
       store.garbageCollectMemories(),
       checkACANReadiness(store),
+      // Deferred cleanup is triggered on first afterTurn() when complete() is available
     ]).catch(e => swallow.warn("bootstrap:maintenance", e));
 
     return { bootstrapped: true };
@@ -340,7 +342,13 @@ export class KongBrainContextEngine implements ContextEngine {
     const session = this.state.getSession(sessionKey);
     if (!session) return;
 
-    const { store } = this.state;
+    const { store, embeddings } = this.state;
+
+    // Deferred cleanup: run once on first turn when complete() is available
+    if (session.userTurnCount <= 1 && typeof this.state.complete === "function") {
+      runDeferredCleanup(store, embeddings, this.state.complete)
+        .catch(e => swallow.warn("afterTurn:deferredCleanup", e));
+    }
 
     // Ingest new messages from this turn (OpenClaw skips ingest() when afterTurn exists)
     const newMessages = params.messages.slice(params.prePromptMessageCount);

package/src/daemon-manager.ts

@@ -9,6 +9,7 @@
  */
 import type { SurrealConfig, EmbeddingConfig } from "./config.js";
 import type { TurnData, PriorExtractions } from "./daemon-types.js";
+import type { CompleteFn } from "./state.js";
 import { SurrealStore } from "./surreal.js";
 import { EmbeddingService } from "./embeddings.js";
 import { swallow } from "./errors.js";
@@ -35,7 +36,7 @@ export function startMemoryDaemon(
   surrealConfig: SurrealConfig,
   embeddingConfig: EmbeddingConfig,
   sessionId: string,
-  llmConfig?: { provider?: string; model?: string },
+  complete: CompleteFn,
 ): MemoryDaemon {
   // Daemon-local DB and embedding instances (separate connections)
   let store: SurrealStore | null = null;
@@ -79,13 +80,6 @@ export function startMemoryDaemon(
     if (!store || !embeddings) return;
     if (turns.length < 2) return;
 
-    const provider = llmConfig?.provider;
-    const modelId = llmConfig?.model;
-    if (!provider || !modelId) {
-      swallow.warn("daemon:extraction", new Error("Missing llmProvider/llmModel"));
-      return;
-    }
-
     // Merge incoming prior state
     if (incomingPrior) {
       for (const name of incomingPrior.conceptNames) {
@@ -116,22 +110,12 @@ export function startMemoryDaemon(
 
     const systemPrompt = buildSystemPrompt(thinking.length > 0, retrievedMemories.length > 0, priorState);
 
-    const { completeSimple, getModel } = await import("@mariozechner/pi-ai");
-    const model = (getModel as any)(provider, modelId);
-
-    const response = await completeSimple(model, {
-      systemPrompt,
-      messages: [{
-        role: "user",
-        timestamp: Date.now(),
-        content: sections.join("\n\n"),
-      }],
+    const response = await complete({
+      system: systemPrompt,
+      messages: [{ role: "user", content: sections.join("\n\n") }],
     });
 
-    const responseText = response.content
-      .filter((c: any) => c.type === "text")
-      .map((c: any) => c.text)
-      .join("");
+    const responseText = response.text;
 
     const jsonMatch = responseText.match(/\{[\s\S]*\}/);
     if (!jsonMatch) return;

package/src/deferred-cleanup.ts

@@ -15,20 +15,49 @@ import { buildSystemPrompt, buildTranscript, writeExtractionResults } from "./me
 import type { PriorExtractions } from "./daemon-types.js";
 import { swallow } from "./errors.js";
 
+// Process-global flag — deferred cleanup runs AT MOST ONCE per process.
+// Using Symbol.for so it survives Jiti re-importing this module.
+const RAN_KEY = Symbol.for("kongbrain.deferredCleanup.ran");
+
 /**
  * Find and process orphaned sessions. Runs with a 30s total timeout.
  * Fire-and-forget from session_start — does not block the new session.
+ * Only runs once per process lifetime.
  */
 export async function runDeferredCleanup(
   store: SurrealStore,
   embeddings: EmbeddingService,
   complete: CompleteFn,
+): Promise<number> {
+  // Once per process — never re-run even if first run times out
+  if ((globalThis as any)[RAN_KEY]) return 0;
+  (globalThis as any)[RAN_KEY] = true;
+
+  try {
+    return await runDeferredCleanupInner(store, embeddings, complete);
+  } catch (e) {
+    swallow.warn("deferredCleanup:outer", e);
+    return 0;
+  }
+}
+
+async function runDeferredCleanupInner(
+  store: SurrealStore,
+  embeddings: EmbeddingService,
+  complete: CompleteFn,
 ): Promise<number> {
   if (!store.isAvailable()) return 0;
 
-  const orphaned = await store.getOrphanedSessions(3).catch(() => []);
+  const orphaned = await store.getOrphanedSessions(10).catch(() => []);
   if (orphaned.length === 0) return 0;
 
+  // Immediately claim all orphaned sessions so no concurrent run can pick them up
+  await Promise.all(
+    orphaned.map(s =>
+      store.markSessionEnded(s.id).catch(e => swallow("deferred:claim", e))
+    )
+  );
+
   let processed = 0;
 
   const cleanup = async () => {
@@ -42,10 +71,10 @@ export async function runDeferredCleanup(
     }
   };
 
-  // 30s timeout — don't hold up the new session forever
+  // 90s timeout — each session needs ~6s (2 LLM calls), 10 sessions ≈ 60s
   await Promise.race([
     cleanup(),
-    new Promise<void>(resolve => setTimeout(resolve, 30_000)),
+    new Promise<void>(resolve => setTimeout(resolve, 90_000)),
   ]);
 
   return processed;
@@ -57,25 +86,15 @@ async function processOrphanedSession(
   embeddings: EmbeddingService,
   complete: CompleteFn,
 ): Promise<void> {
-  // Find the OpenClaw session ID from turns stored in this session
-  // (turns use the OpenClaw session_id, not the surreal record ID)
-  const sessionTurns = await store.queryFirst<{ session_id: string }>(
-    `SELECT session_id FROM turn WHERE session_id != NONE ORDER BY created_at DESC LIMIT 1`,
-  ).catch(() => []);
-
-  // Load turns for extraction
-  // We need to find turns associated with this DB session via the part_of edge
+  // Load turns for extraction via part_of edges (turn->part_of->session)
   const turns = await store.queryFirst<{ role: string; text: string; tool_name?: string }>(
-    `SELECT role, text, tool_name FROM turn
-     WHERE session_id IN (SELECT VALUE out FROM part_of WHERE in = $sid)
-        OR session_id = $sid
+    `SELECT role, text, tool_name, created_at FROM turn
+     WHERE id IN (SELECT VALUE in FROM part_of WHERE out = $sid)
      ORDER BY created_at ASC LIMIT 50`,
     { sid: surrealSessionId },
   ).catch(() => []);
 
   if (turns.length < 2) {
-    // Nothing to extract, just mark complete
-    await store.markSessionEnded(surrealSessionId).catch(e => swallow("deferred:markEmpty", e));
     return;
   }
 
@@ -86,12 +105,14 @@ async function processOrphanedSession(
   const systemPrompt = buildSystemPrompt(false, false, priorState);
 
   try {
+    console.warn(`[deferred] extracting session ${surrealSessionId} (${turns.length} turns, transcript ${transcript.length} chars)`);
     const response = await complete({
       system: systemPrompt,
       messages: [{ role: "user", content: `[TRANSCRIPT]\n${transcript.slice(0, 60000)}` }],
     });
 
     const responseText = response.text;
+    console.warn(`[deferred] extraction response: ${responseText.length} chars`);
     const jsonMatch = responseText.match(/\{[\s\S]*\}/);
     if (jsonMatch) {
       let result: Record<string, any>;
@@ -103,10 +124,14 @@ async function processOrphanedSession(
         } catch { result = {}; }
       }
 
-      if (Object.keys(result).length > 0) {
-        const sessionId = surrealSessionId; // Use DB ID as session reference
-        await writeExtractionResults(result, sessionId, store, embeddings, priorState);
+      const keys = Object.keys(result);
+      console.warn(`[deferred] parsed ${keys.length} keys: ${keys.join(", ")}`);
+      if (keys.length > 0) {
+        await writeExtractionResults(result, surrealSessionId, store, embeddings, priorState);
+        console.warn(`[deferred] wrote extraction results for ${surrealSessionId}`);
       }
+    } else {
+      console.warn(`[deferred] no JSON found in response`);
     }
   } catch (e) {
     swallow.warn("deferredCleanup:extraction", e);
@@ -125,6 +150,7 @@ async function processOrphanedSession(
     });
 
     const handoffText = handoffResponse.text.trim();
+    console.warn(`[deferred] handoff response: ${handoffText.length} chars`);
     if (handoffText.length > 20) {
       let emb: number[] | null = null;
       if (embeddings.isAvailable()) {
@@ -135,7 +161,4 @@ async function processOrphanedSession(
   } catch (e) {
     swallow.warn("deferredCleanup:handoff", e);
   }
-
-  // Mark session as cleaned up
-  await store.markSessionEnded(surrealSessionId).catch(e => swallow("deferred:markDone", e));
 }

package/src/embeddings.ts

@@ -15,7 +15,9 @@ export class EmbeddingService {
 
   constructor(private readonly config: EmbeddingConfig) {}
 
-  async initialize(): Promise<void> {
+  /** Initialize the embedding model. Returns true if freshly loaded, false if already ready. */
+  async initialize(): Promise<boolean> {
+    if (this.ready) return false;
     if (!existsSync(this.config.modelPath)) {
       throw new Error(
         `Embedding model not found at: ${this.config.modelPath}\n  Download BGE-M3 GGUF or set EMBED_MODEL_PATH`,
@@ -33,6 +35,7 @@ export class EmbeddingService {
     this.model = await llama.loadModel({ modelPath: this.config.modelPath });
     this.ctx = await this.model.createEmbeddingContext();
     this.ready = true;
+    return true;
   }
 
   async embed(text: string): Promise<number[]> {

package/src/index.ts

@@ -6,12 +6,13 @@
  */
 
 import { readFile } from "node:fs/promises";
-import { join } from "node:path";
+import { existsSync } from "node:fs";
+import { join, dirname } from "node:path";
 import { definePluginEntry } from "openclaw/plugin-sdk/plugin-entry";
 import { parsePluginConfig } from "./config.js";
 import { SurrealStore } from "./surreal.js";
 import { EmbeddingService } from "./embeddings.js";
-import { GlobalPluginState } from "./state.js";
+import { GlobalPluginState, type CompleteFn } from "./state.js";
 import { KongBrainContextEngine } from "./context-engine.js";
 import { createRecallToolDef } from "./tools/recall.js";
 import { createCoreMemoryToolDef } from "./tools/core-memory.js";
@@ -32,11 +33,28 @@ import { writeHandoffFileSync } from "./handoff-file.js";
 import { runDeferredCleanup } from "./deferred-cleanup.js";
 import { swallow } from "./errors.js";
 
-let globalState: GlobalPluginState | null = null;
+// Use process-global symbols so state survives Jiti re-importing the module.
+// Jiti may load this file multiple times (fresh module scope each time),
+// but process.env and Symbol.for() are process-wide singletons.
+const GLOBAL_KEY = Symbol.for("kongbrain.globalState");
+const REGISTERED_KEY = Symbol.for("kongbrain.registered");
+
+function getGlobalState(): GlobalPluginState | null {
+  return (globalThis as any)[GLOBAL_KEY] ?? null;
+}
+function setGlobalState(state: GlobalPluginState): void {
+  (globalThis as any)[GLOBAL_KEY] = state;
+}
+function isRegistered(): boolean {
+  return (globalThis as any)[REGISTERED_KEY] === true;
+}
+function markRegistered(): void {
+  (globalThis as any)[REGISTERED_KEY] = true;
+}
+
 let shutdownPromise: Promise<void> | null = null;
 let registeredExitHandler: (() => void) | null = null;
 let registeredSyncExitHandler: (() => void) | null = null;
-let registered = false;
 
 /**
  * Run the critical session-end extraction for all active sessions.
@@ -294,13 +312,79 @@ export default definePluginEntry({
     const config = parsePluginConfig(api.pluginConfig as Record<string, unknown> | undefined);
     const logger = api.logger;
 
-    // Initialize shared resources — reuse existing globalState if register() is called
-    // multiple times (OpenClaw may invoke the factory more than once). Hooks from the
-    // first register() hold a closure over globalState, so replacing it would orphan them.
+    // Initialize shared resources — reuse existing state if register() is called
+    // multiple times (OpenClaw may invoke the factory more than once, and Jiti may
+    // re-import the module creating fresh module scope). Process-global symbols
+    // ensure a single instance survives across module reloads.
+    let globalState = getGlobalState();
     if (!globalState) {
       const store = new SurrealStore(config.surreal);
       const embeddings = new EmbeddingService(config.embedding);
-      globalState = new GlobalPluginState(config, store, embeddings, api.runtime.complete);
+      // Build a CompleteFn using pi-ai directly since api.runtime.complete
+      // is not available in OpenClaw 2026.3.24 (unreleased feature).
+      const apiRef = api;
+      // Resolve pi-ai from openclaw's node_modules. pi-ai is ESM-only so
+      // require() can't load it. Walk up from process.argv[1] to find it,
+      // then lazy-load via import() on first use.
+      let piAi: { getModel: any; completeSimple: any } | null = null;
+      let piAiPath: string | null = null;
+      {
+        let dir = dirname(process.argv[1] || __filename);
+        for (let i = 0; i < 10; i++) {
+          const candidate = join(dir, "node_modules", "@mariozechner", "pi-ai", "dist", "index.js");
+          if (existsSync(candidate)) { piAiPath = candidate; break; }
+          const parent = dirname(dir);
+          if (parent === dir) break;
+          dir = parent;
+        }
+      }
+
+      const complete: CompleteFn = async (params) => {
+        // Try runtime.complete first (future-proof for when it ships)
+        if (typeof apiRef.runtime?.complete === "function") {
+          return apiRef.runtime.complete(params);
+        }
+        if (!piAi) {
+          if (!piAiPath) {
+            throw new Error("LLM completion not available: @mariozechner/pi-ai not found and runtime.complete missing");
+          }
+          piAi = await import(piAiPath);
+        }
+        // Fall back to calling pi-ai directly (runtime.complete not in OpenClaw 2026.3.24)
+        const provider = params.provider ?? apiRef.runtime.agent.defaults.provider;
+        const modelId = params.model ?? apiRef.runtime.agent.defaults.model;
+        const model = piAi!.getModel(provider, modelId);
+        if (!model) {
+          throw new Error(`Model "${modelId}" not found for provider "${provider}"`);
+        }
+        // Resolve auth via OpenClaw's runtime (handles profiles, env vars, etc.)
+        const cfg = apiRef.runtime.config.loadConfig();
+        const auth = await apiRef.runtime.modelAuth.getApiKeyForModel({ model, cfg });
+        // Build context
+        const now = Date.now();
+        const messages: any[] = params.messages.map(m =>
+          m.role === "user"
+            ? { role: "user", content: m.content, timestamp: now }
+            : { role: "assistant", content: [{ type: "text", text: m.content }],
+                api: model.api, provider: model.provider, model: model.id,
+                usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 } },
+                stopReason: "stop", timestamp: now }
+        );
+        const context = { systemPrompt: params.system, messages };
+        // Pass apiKey directly in options so the provider can use it
+        const response = await piAi!.completeSimple(model, context, {
+          apiKey: auth.apiKey,
+        });
+        let text = "";
+        let thinking: string | undefined;
+        for (const block of response.content) {
+          if (block.type === "text") text += block.text;
+          else if ((block as any).type === "thinking") thinking = (thinking ?? "") + (block as any).thinking;
+        }
+        return { text, thinking, usage: { input: response.usage.input, output: response.usage.output } };
+      };
+      globalState = new GlobalPluginState(config, store, embeddings, complete);
+      setGlobalState(globalState);
     }
     globalState.workspaceDir = api.resolvePath(".");
     globalState.enqueueSystemEvent = (text, opts) =>
@@ -312,19 +396,19 @@ export default definePluginEntry({
     api.registerContextEngine("kongbrain", async () => {
       const { store, embeddings } = state;
 
-      // Connect to SurrealDB
+      // Connect to SurrealDB (no-op if already connected)
       try {
-        await store.initialize();
-        logger.info(`SurrealDB connected: ${config.surreal.url}`);
+        const freshConnect = await store.initialize();
+        if (freshConnect) logger.info(`SurrealDB connected: ${config.surreal.url}`);
       } catch (e) {
         logger.error(`SurrealDB connection failed: ${e}`);
         throw e;
       }
 
-      // Initialize BGE-M3 embeddings
+      // Initialize BGE-M3 embeddings (no-op if already loaded)
       try {
-        await embeddings.initialize();
-        logger.info(`BGE-M3 embeddings initialized: ${config.embedding.modelPath}`);
+        const freshEmbed = await embeddings.initialize();
+        if (freshEmbed) logger.info(`BGE-M3 embeddings initialized: ${config.embedding.modelPath}`);
       } catch (e) {
         logger.warn(`Embeddings init failed — running in degraded mode: ${e}`);
       }
@@ -339,7 +423,7 @@ export default definePluginEntry({
 
     // ── Hook handlers (register once — register() may be called multiple times) ──
 
-    if (!registered) {
+    if (!isRegistered()) {
       api.on("before_prompt_build", createBeforePromptBuildHandler(globalState));
       api.on("before_tool_call", createBeforeToolCallHandler(globalState));
       api.on("after_tool_call", createAfterToolCallHandler(globalState));
@@ -348,7 +432,8 @@ export default definePluginEntry({
 
     // ── Session lifecycle (also register once) ─────────────────────────
 
-    if (!registered) api.on("session_start", async (event) => {
+    if (!isRegistered()) api.on("session_start", async (event) => {
+      const globalState = getGlobalState();
       if (!globalState) return;
       const sessionKey = event.sessionKey ?? event.sessionId;
       const session = globalState.getOrCreateSession(sessionKey, event.sessionId);
@@ -377,7 +462,7 @@ export default definePluginEntry({
           config.surreal,
           config.embedding,
           session.sessionId,
-          { provider: api.runtime.agent.defaults.provider, model: api.runtime.agent.defaults.model },
+          globalState!.complete,
         );
       } catch (e) {
         swallow.warn("index:startDaemon", e);
@@ -398,31 +483,32 @@ export default definePluginEntry({
       setReflectionContextWindow(200000);
 
       // Check for recent graduation event (from a previous session)
-      detectGraduationEvent(store, session, globalState!)
+      detectGraduationEvent(globalState!.store, session, globalState!)
         .catch(e => swallow("index:graduationDetect", e));
 
       // Synthesize wakeup briefing (background, non-blocking)
       // The briefing is stored and later injected via assemble()'s systemPromptAddition
-      synthesizeWakeup(store, globalState!.complete, session.sessionId, globalState!.workspaceDir)
+      synthesizeWakeup(globalState!.store, globalState!.complete, session.sessionId, globalState!.workspaceDir)
         .then(briefing => {
           if (briefing) (session as any)._wakeupBriefing = briefing;
         })
         .catch(e => swallow.warn("index:wakeup", e));
 
       // Startup cognition (background)
-      synthesizeStartupCognition(store, globalState!.complete)
+      synthesizeStartupCognition(globalState!.store, globalState!.complete)
         .then(cognition => {
           if (cognition) (session as any)._startupCognition = cognition;
         })
         .catch(e => swallow.warn("index:startupCognition", e));
 
       // Deferred cleanup: extract knowledge from orphaned sessions (background)
-      runDeferredCleanup(store, embeddings, globalState!.complete)
+      runDeferredCleanup(globalState!.store, globalState!.embeddings, globalState!.complete)
         .then(n => { if (n > 0) logger.info(`Deferred cleanup: processed ${n} orphaned session(s)`); })
         .catch(e => swallow.warn("index:deferredCleanup", e));
     });
 
-    if (!registered) api.on("session_end", async (event) => {
+    if (!isRegistered()) api.on("session_end", async (event) => {
+      const globalState = getGlobalState();
       if (!globalState) return;
       const sessionKey = event.sessionKey ?? event.sessionId;
       const session = globalState.getSession(sessionKey);
@@ -434,7 +520,7 @@ export default definePluginEntry({
 
       session.cleanedUp = true;
       if (session.surrealSessionId) {
-        await store.markSessionEnded(session.surrealSessionId)
+        await globalState.store.markSessionEnded(session.surrealSessionId)
           .catch(e => swallow.warn("session_end:markEnded", e));
       }
 
@@ -458,8 +544,9 @@ export default definePluginEntry({
 
     // Sync exit handler: writes handoff file for all uncleaned sessions
     const syncExitHandler = () => {
-      if (!globalState?.workspaceDir) return;
-      const sessions = [...(globalState as any).sessions.values()] as import("./state.js").SessionState[];
+      const gs = getGlobalState();
+      if (!gs?.workspaceDir) return;
+      const sessions = [...(gs as any).sessions.values()] as import("./state.js").SessionState[];
       for (const session of sessions) {
         if (session.cleanedUp) continue;
         writeHandoffFileSync({
@@ -469,21 +556,22 @@ export default definePluginEntry({
           lastUserText: session.lastUserText.slice(0, 500),
           lastAssistantText: session.lastAssistantText.slice(0, 500),
           unextractedTokens: session.newContentTokens,
-        }, globalState!.workspaceDir!);
+        }, gs.workspaceDir!);
       }
     };
 
     // Async exit handler: full cleanup for SIGTERM (gateway/daemon mode)
     const asyncExitHandler = () => {
-      if (!globalState) return;
-      const sessions = [...(globalState as any).sessions.values()] as import("./state.js").SessionState[];
+      const gs = getGlobalState();
+      if (!gs) return;
+      const sessions = [...(gs as any).sessions.values()] as import("./state.js").SessionState[];
       if (sessions.length === 0 && !shutdownPromise) return;
 
-      const cleanups = sessions.map(s => runSessionCleanup(s, globalState!));
+      const cleanups = sessions.map(s => runSessionCleanup(s, gs));
       if (shutdownPromise) cleanups.push(shutdownPromise);
 
       const done = Promise.allSettled(cleanups).then(() => {
-        globalState?.shutdown().catch(() => {});
+        gs.shutdown().catch(() => {});
       });
 
       done.then(() => process.exit(0)).catch(() => process.exit(1));
@@ -494,9 +582,6 @@ export default definePluginEntry({
     process.on("exit", syncExitHandler);
     process.once("SIGTERM", asyncExitHandler);
 
-    if (!registered) {
-      logger.info("KongBrain plugin registered");
-      registered = true;
-    }
+    markRegistered();
   },
 });

package/src/state.ts

@@ -1,12 +1,32 @@
-import type { PluginCompleteParams, PluginCompleteResult } from "openclaw/plugin-sdk";
 import type { KongBrainConfig } from "./config.js";
 import type { SurrealStore } from "./surreal.js";
 import type { EmbeddingService } from "./embeddings.js";
 import type { AdaptiveConfig } from "./orchestrator.js";
 import type { MemoryDaemon } from "./daemon-manager.js";
 
+/** Parameters for an LLM completion call. */
+export type CompleteParams = {
+  system?: string;
+  messages: { role: "user" | "assistant"; content: string }[];
+  provider?: string;
+  model?: string;
+  temperature?: number;
+  maxTokens?: number;
+  reasoning?: "none" | "low" | "medium" | "high";
+};
+
+/** Result of an LLM completion call. */
+export type CompleteResult = {
+  text: string;
+  thinking?: string;
+  usage?: { input: number; output: number };
+  provider?: string;
+  model?: string;
+  stopReason?: string;
+};
+
 /** Provider-agnostic LLM completion function. */
-export type CompleteFn = (params: PluginCompleteParams) => Promise<PluginCompleteResult>;
+export type CompleteFn = (params: CompleteParams) => Promise<CompleteResult>;
 
 // --- Per-session mutable state ---
 
@@ -78,7 +98,7 @@ export class GlobalPluginState {
   readonly config: KongBrainConfig;
   readonly store: SurrealStore;
   readonly embeddings: EmbeddingService;
-  readonly complete: CompleteFn;
+  complete: CompleteFn;
   workspaceDir?: string;
   enqueueSystemEvent?: EnqueueSystemEventFn;
   private sessions = new Map<string, SessionState>();

package/src/surreal.ts

@@ -105,19 +105,28 @@ export class SurrealStore {
   private config: SurrealConfig;
   private reconnecting: Promise<void> | null = null;
   private shutdownFlag = false;
+  private initialized = false;
 
   constructor(config: SurrealConfig) {
     this.config = config;
     this.db = new Surreal();
   }
 
-  async initialize(): Promise<void> {
+  /** Connect and run schema. Returns true if a new connection was made, false if already initialized. */
+  async initialize(): Promise<boolean> {
+    // Only connect once — subsequent calls are no-ops.
+    // This prevents register()/factory re-invocations from disrupting
+    // in-flight operations (deferred cleanup, daemon extraction).
+    // Don't check isConnected — ensureConnected() handles reconnection.
+    if (this.initialized) return false;
     await this.db.connect(this.config.url, {
       namespace: this.config.ns,
       database: this.config.db,
       authentication: { username: this.config.user, password: this.config.pass },
     });
     await this.runSchema();
+    this.initialized = true;
+    return true;
   }
 
   markShutdown(): void {
@@ -218,16 +227,43 @@ export class SurrealStore {
     }
   }
 
+  /** Returns true if an error is a connection-level failure worth retrying. */
+  private isConnectionError(e: unknown): boolean {
+    const msg = String((e as any)?.message ?? e);
+    return msg.includes("must be connected") || msg.includes("ConnectionUnavailable");
+  }
+
+  /** Run a query function with one retry on connection errors. */
+  private async withRetry<T>(fn: () => Promise<T>): Promise<T> {
+    try {
+      return await fn();
+    } catch (e) {
+      if (!this.isConnectionError(e)) throw e;
+      // Connection died — force a fresh connection (close stale socket first)
+      this.initialized = false;
+      try { await this.db?.close(); } catch { /* ignore */ }
+      this.db = new Surreal();
+      await this.db.connect(this.config.url, {
+        namespace: this.config.ns,
+        database: this.config.db,
+        authentication: { username: this.config.user, password: this.config.pass },
+      });
+      return await fn();
+    }
+  }
+
   // ── Query helpers ──────────────────────────────────────────────────────
 
   async queryFirst<T>(sql: string, bindings?: Record<string, unknown>): Promise<T[]> {
     await this.ensureConnected();
-    const ns = this.config.ns;
-    const dbName = this.config.db;
-    const fullSql = `USE NS ${ns} DB ${dbName}; ${patchOrderByFields(sql)}`;
-    const result = await this.db.query<[T[]]>(fullSql, bindings);
-    const rows = Array.isArray(result) ? result[result.length - 1] : result;
-    return (Array.isArray(rows) ? rows : []).filter(Boolean);
+    return this.withRetry(async () => {
+      const ns = this.config.ns;
+      const dbName = this.config.db;
+      const fullSql = `USE NS ${ns} DB ${dbName}; ${patchOrderByFields(sql)}`;
+      const result = await this.db.query<[T[]]>(fullSql, bindings);
+      const rows = Array.isArray(result) ? result[result.length - 1] : result;
+      return (Array.isArray(rows) ? rows : []).filter(Boolean);
+    });
   }
 
   async queryMulti<T = unknown>(
@@ -235,20 +271,24 @@ export class SurrealStore {
     bindings?: Record<string, unknown>,
   ): Promise<T | undefined> {
     await this.ensureConnected();
-    const ns = this.config.ns;
-    const dbName = this.config.db;
-    const fullSql = `USE NS ${ns} DB ${dbName}; ${patchOrderByFields(sql)}`;
-    const raw = await this.db.query(fullSql, bindings);
-    const flat = (raw as unknown[]).flat();
-    return flat[flat.length - 1] as T | undefined;
+    return this.withRetry(async () => {
+      const ns = this.config.ns;
+      const dbName = this.config.db;
+      const fullSql = `USE NS ${ns} DB ${dbName}; ${patchOrderByFields(sql)}`;
+      const raw = await this.db.query(fullSql, bindings);
+      const flat = (raw as unknown[]).flat();
+      return flat[flat.length - 1] as T | undefined;
+    });
   }
 
   async queryExec(sql: string, bindings?: Record<string, unknown>): Promise<void> {
     await this.ensureConnected();
-    const ns = this.config.ns;
-    const dbName = this.config.db;
-    const fullSql = `USE NS ${ns} DB ${dbName}; ${patchOrderByFields(sql)}`;
-    await this.db.query(fullSql, bindings);
+    return this.withRetry(async () => {
+      const ns = this.config.ns;
+      const dbName = this.config.db;
+      const fullSql = `USE NS ${ns} DB ${dbName}; ${patchOrderByFields(sql)}`;
+      await this.db.query(fullSql, bindings);
+    });
   }
 
   private async safeQuery(