kongbrain 0.1.2 → 0.1.3

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "kongbrain",
-  "version": "0.1.2",
+  "version": "0.1.3",
   "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/context-engine.ts

@@ -103,10 +103,13 @@ export class KongBrainContextEngine implements ContextEngine {
         .catch(e => swallow.warn("bootstrap:linkTaskToProject", e));
 
       const surrealSessionId = await store.createSession(session.agentId);
+      await store.markSessionActive(surrealSessionId)
+        .catch(e => swallow.warn("bootstrap:markActive", e));
       await store.linkSessionToTask(surrealSessionId, session.taskId)
         .catch(e => swallow.warn("bootstrap:linkSessionToTask", e));
 
-      // Update session with the DB-assigned session ID
+      // Store the DB session ID for cleanup tracking
+      session.surrealSessionId = surrealSessionId;
       session.lastUserTurnId = "";
     } catch (e) {
       swallow.error("bootstrap:5pillar", e);
@@ -248,8 +251,10 @@ export class KongBrainContextEngine implements ContextEngine {
         });
 
         if (turnId) {
-          await store.relate(turnId, "part_of", session.sessionId)
-            .catch(e => swallow.warn("ingest:relate", e));
+          if (session.surrealSessionId) {
+            await store.relate(turnId, "part_of", session.surrealSessionId)
+              .catch(e => swallow.warn("ingest:relate", e));
+          }
 
           // Link to previous user turn for responds_to edge
           if (role === "assistant" && session.lastUserTurnId) {
@@ -257,8 +262,8 @@ export class KongBrainContextEngine implements ContextEngine {
               .catch(e => swallow.warn("ingest:responds_to", e));
           }
 
-          // Extract and link concepts for user turns
-          if (role === "user" && worthEmbedding) {
+          // Extract and link concepts for both user and assistant turns
+          if (worthEmbedding) {
             extractAndLinkConcepts(turnId, text, this.state)
               .catch(e => swallow.warn("ingest:concepts", e));
           }
@@ -381,8 +386,10 @@ export class KongBrainContextEngine implements ContextEngine {
       session.newContentTokens += Math.ceil(session.lastAssistantText.length / 4);
     }
 
-    // Flush to daemon when token threshold is reached
-    if (session.daemon && session.newContentTokens >= session.DAEMON_TOKEN_THRESHOLD) {
+    // Flush to daemon when token threshold OR turn count threshold is reached
+    const tokenReady = session.newContentTokens >= session.DAEMON_TOKEN_THRESHOLD;
+    const turnReady = session.userTurnCount >= session.lastDaemonFlushTurnCount + 3;
+    if (session.daemon && (tokenReady || turnReady)) {
       try {
         const recentTurns = await store.getSessionTurns(session.sessionId, 20);
         const turnData = recentTurns.map(t => ({
@@ -404,6 +411,7 @@ export class KongBrainContextEngine implements ContextEngine {
         );
 
         session.newContentTokens = 0;
+        session.lastDaemonFlushTurnCount = session.userTurnCount;
         session.pendingThinking.length = 0;
       } catch (e) {
         swallow.warn("afterTurn:daemonBatch", e);

package/src/daemon-manager.ts

@@ -1,20 +1,18 @@
 /**
- * Daemon Manager — spawns and manages the memory daemon worker thread.
+ * Daemon Manager — runs memory extraction in-process.
  *
- * Provides a clean interface for sending turn batches, querying status,
- * and graceful shutdown. Used by the session lifecycle hooks.
- *
- * Ported from kongbrain — takes config as params instead of env globals.
+ * Originally used a Worker thread, but OpenClaw loads plugins via jiti
+ * (TypeScript only, no compiled JS), and Node's Worker constructor requires
+ * .js files. Refactored to run extraction async in the main thread.
+ * The extraction is I/O-bound (LLM calls + DB writes), not CPU-bound,
+ * so in-process execution is fine.
  */
-import { Worker } from "node:worker_threads";
-import { join, dirname } from "node:path";
-import { fileURLToPath } from "node:url";
 import type { SurrealConfig, EmbeddingConfig } from "./config.js";
-import type { DaemonMessage, DaemonResponse, DaemonWorkerData, TurnData, PriorExtractions } from "./daemon-types.js";
+import type { TurnData, PriorExtractions } from "./daemon-types.js";
+import { SurrealStore } from "./surreal.js";
+import { EmbeddingService } from "./embeddings.js";
 import { swallow } from "./errors.js";
 
-const __dirname = dirname(fileURLToPath(import.meta.url));
-
 export type { TurnData } from "./daemon-types.js";
 
 export interface MemoryDaemon {
@@ -25,11 +23,11 @@ export interface MemoryDaemon {
     retrievedMemories: { id: string; text: string }[],
     priorExtractions?: PriorExtractions,
   ): void;
-  /** Request current daemon status (async, waits for response). */
-  getStatus(): Promise<DaemonResponse & { type: "status" }>;
-  /** Graceful shutdown: waits for current extraction, then terminates. */
+  /** Request current daemon status. */
+  getStatus(): Promise<{ type: "status"; extractedTurns: number; pendingBatches: number; errors: number }>;
+  /** Graceful shutdown: waits for current extraction, then cleans up. */
   shutdown(timeoutMs?: number): Promise<void>;
-  /** Synchronous: how many turns has the daemon already extracted? */
+  /** How many turns has the daemon already extracted? */
   getExtractedTurnCount(): number;
 }
 
@@ -39,106 +37,192 @@ export function startMemoryDaemon(
   sessionId: string,
   llmConfig?: { provider?: string; model?: string },
 ): MemoryDaemon {
-  const workerData: DaemonWorkerData = {
-    surrealConfig,
-    embeddingConfig,
-    sessionId,
-    llmProvider: llmConfig?.provider,
-    llmModel: llmConfig?.model,
+  // Daemon-local DB and embedding instances (separate connections)
+  let store: SurrealStore | null = null;
+  let embeddings: EmbeddingService | null = null;
+  let initialized = false;
+  let initFailed = false;
+  let processing = false;
+  let shuttingDown = false;
+  let extractedTurnCount = 0;
+  let errorCount = 0;
+
+  const priorState: PriorExtractions = {
+    conceptNames: [], artifactPaths: [], skillNames: [],
   };
 
-  const worker = new Worker(join(__dirname, "memory-daemon.js"), { workerData });
+  // Lazy init — connect on first batch, not at startup
+  async function ensureInit(): Promise<boolean> {
+    if (initialized) return true;
+    if (initFailed) return false;
+    try {
+      store = new SurrealStore(surrealConfig);
+      await store.initialize();
+      embeddings = new EmbeddingService(embeddingConfig);
+      await embeddings.initialize();
+      initialized = true;
+      return true;
+    } catch (e) {
+      swallow.warn("daemon:init", e);
+      initFailed = true;
+      return false;
+    }
+  }
 
-  let extractedTurnCount = 0;
-  let terminated = false;
-
-  let pendingStatusResolve: ((resp: DaemonResponse & { type: "status" }) => void) | null = null;
-
-  worker.on("message", (msg: DaemonResponse) => {
-    switch (msg.type) {
-      case "extraction_complete":
-        extractedTurnCount = msg.extractedTurnCount;
-        break;
-      case "status":
-        if (pendingStatusResolve) {
-          pendingStatusResolve(msg as DaemonResponse & { type: "status" });
-          pendingStatusResolve = null;
-        }
-        break;
-      case "error":
-        swallow.warn("daemon-manager:worker-error", new Error(msg.message));
-        break;
+  // Import extraction logic lazily to avoid circular deps
+  async function runExtraction(
+    turns: TurnData[],
+    thinking: string[],
+    retrievedMemories: { id: string; text: string }[],
+    incomingPrior?: PriorExtractions,
+  ): Promise<void> {
+    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) {
+        if (!priorState.conceptNames.includes(name)) priorState.conceptNames.push(name);
+      }
+      for (const path of incomingPrior.artifactPaths) {
+        if (!priorState.artifactPaths.includes(path)) priorState.artifactPaths.push(path);
+      }
+      for (const name of incomingPrior.skillNames) {
+        if (!priorState.skillNames.includes(name)) priorState.skillNames.push(name);
+      }
     }
-  });
 
-  worker.on("error", (err) => {
-    swallow.warn("daemon-manager:worker-thread-error", err);
-  });
+    // Dynamically import the extraction helpers from memory-daemon
+    const { buildSystemPrompt, buildTranscript, writeExtractionResults } = await import("./memory-daemon.js");
 
-  worker.on("exit", (code) => {
-    terminated = true;
-    if (code !== 0) {
-      swallow.warn("daemon-manager:worker-exit", new Error(`Daemon exited with code ${code}`));
+    const transcript = buildTranscript(turns);
+    const sections: string[] = [`[TRANSCRIPT]\n${transcript.slice(0, 60000)}`];
+
+    if (thinking.length > 0) {
+      sections.push(`[THINKING]\n${thinking.slice(-8).join("\n---\n").slice(0, 4000)}`);
+    }
+
+    if (retrievedMemories.length > 0) {
+      const memList = retrievedMemories.map(m => `${m.id}: ${String(m.text).slice(0, 200)}`).join("\n");
+      sections.push(`[RETRIEVED MEMORIES]\nMark any that have been fully addressed/fixed/completed.\n${memList}`);
     }
-  });
+
+    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 responseText = response.content
+      .filter((c: any) => c.type === "text")
+      .map((c: any) => c.text)
+      .join("");
+
+    const jsonMatch = responseText.match(/\{[\s\S]*\}/);
+    if (!jsonMatch) return;
+
+    let result: Record<string, any>;
+    try {
+      result = JSON.parse(jsonMatch[0]);
+    } catch {
+      try {
+        result = JSON.parse(jsonMatch[0].replace(/,\s*([}\]])/g, "$1"));
+      } catch {
+        result = {};
+        const fields = ["causal", "monologue", "resolved", "concepts", "corrections", "preferences", "artifacts", "decisions", "skills"];
+        for (const field of fields) {
+          const fieldMatch = jsonMatch[0].match(new RegExp(`"${field}"\\s*:\\s*(\\[[\\s\\S]*?\\])(?=\\s*[,}]\\s*"[a-z]|\\s*\\}$)`, "m"));
+          if (fieldMatch) {
+            try { result[field] = JSON.parse(fieldMatch[1]); } catch { /* skip */ }
+          }
+        }
+        if (Object.keys(result).length === 0) return;
+      }
+    }
+
+    const counts = await writeExtractionResults(result, sessionId, store, embeddings, priorState);
+    extractedTurnCount = turns.length;
+  }
+
+  // Pending batch (only keep latest — newer batch supersedes older)
+  let pendingBatch: {
+    turns: TurnData[];
+    thinking: string[];
+    retrievedMemories: { id: string; text: string }[];
+    priorExtractions?: PriorExtractions;
+  } | null = null;
+
+  async function processPending(): Promise<void> {
+    if (processing || shuttingDown) return;
+    while (pendingBatch) {
+      processing = true;
+      const batch = pendingBatch;
+      pendingBatch = null;
+      try {
+        await runExtraction(batch.turns, batch.thinking, batch.retrievedMemories, batch.priorExtractions);
+      } catch (e) {
+        errorCount++;
+        swallow.warn("daemon:extraction", e);
+      } finally {
+        processing = false;
+      }
+    }
+  }
 
   return {
     sendTurnBatch(turns, thinking, retrievedMemories, priorExtractions) {
-      if (terminated) return;
-      try {
-        worker.postMessage({
-          type: "turn_batch",
-          turns,
-          thinking,
-          retrievedMemories,
-          sessionId,
-          priorExtractions,
-        } satisfies DaemonMessage);
-      } catch (e) { swallow.warn("daemon-manager:sendBatch", e); }
+      if (shuttingDown) return;
+      pendingBatch = { turns, thinking, retrievedMemories, priorExtractions };
+      // Fire-and-forget: init if needed, then process
+      ensureInit()
+        .then(ok => { if (ok) return processPending(); })
+        .catch(e => swallow.warn("daemon:sendBatch", e));
     },
 
     async getStatus() {
-      if (terminated) return { type: "status" as const, extractedTurns: extractedTurnCount, pendingBatches: 0, errors: 0 };
-      return new Promise<DaemonResponse & { type: "status" }>((resolve) => {
-        const timer = setTimeout(() => {
-          pendingStatusResolve = null;
-          resolve({ type: "status", extractedTurns: extractedTurnCount, pendingBatches: -1, errors: -1 });
-        }, 5000);
-        pendingStatusResolve = (resp) => {
-          clearTimeout(timer);
-          resolve(resp);
-        };
-        worker.postMessage({ type: "status_request" } satisfies DaemonMessage);
-      });
+      return {
+        type: "status" as const,
+        extractedTurns: extractedTurnCount,
+        pendingBatches: pendingBatch ? 1 : 0,
+        errors: errorCount,
+      };
     },
 
     async shutdown(timeoutMs = 45_000) {
-      if (terminated) return;
-      return new Promise<void>((resolve) => {
-        const timer = setTimeout(() => {
-          worker.terminate().catch(() => {});
-          terminated = true;
-          resolve();
-        }, timeoutMs);
-
-        const onMessage = (msg: DaemonResponse) => {
-          if (msg.type === "shutdown_complete") {
-            clearTimeout(timer);
-            worker.removeListener("message", onMessage);
-            terminated = true;
-            resolve();
-          }
-        };
-        worker.on("message", onMessage);
-
-        try {
-          worker.postMessage({ type: "shutdown" } satisfies DaemonMessage);
-        } catch {
-          clearTimeout(timer);
-          terminated = true;
-          resolve();
-        }
-      });
+      shuttingDown = true;
+      // Wait for current extraction to finish
+      if (processing) {
+        await Promise.race([
+          new Promise<void>(resolve => {
+            const check = setInterval(() => {
+              if (!processing) { clearInterval(check); resolve(); }
+            }, 100);
+          }),
+          new Promise<void>(resolve => setTimeout(resolve, timeoutMs)),
+        ]);
+      }
+      // Clean up daemon-local connections
+      await Promise.allSettled([
+        store?.dispose(),
+        embeddings?.dispose(),
+      ]).catch(() => {});
+      store = null;
+      embeddings = null;
     },
 
     getExtractedTurnCount() {

package/src/daemon-types.ts

@@ -1,9 +1,6 @@
 /**
  * Shared types for the memory daemon system.
- * Imported by both the worker thread (memory-daemon.ts) and the
- * main thread manager (daemon-manager.ts).
  */
-import type { SurrealConfig, EmbeddingConfig } from "./config.js";
 
 export interface TurnData {
   role: string;
@@ -13,53 +10,9 @@ export interface TurnData {
   file_paths?: string[];
 }
 
-/** Data passed to the worker thread via workerData. */
-export interface DaemonWorkerData {
-  surrealConfig: SurrealConfig;
-  embeddingConfig: EmbeddingConfig;
-  sessionId: string;
-  /** LLM provider name (resolved from OpenClaw config at daemon start). */
-  llmProvider?: string;
-  /** LLM model ID (resolved from OpenClaw config at daemon start). */
-  llmModel?: string;
-}
-
-/** Previously extracted item names — for dedup across daemon runs. */
+/** Previously extracted item names for dedup across daemon runs. */
 export interface PriorExtractions {
   conceptNames: string[];
   artifactPaths: string[];
   skillNames: string[];
 }
-
-/** Messages from main thread -> daemon worker. */
-export type DaemonMessage =
-  | {
-      type: "turn_batch";
-      turns: TurnData[];
-      thinking: string[];
-      retrievedMemories: { id: string; text: string }[];
-      sessionId: string;
-      priorExtractions?: PriorExtractions;
-    }
-  | { type: "shutdown" }
-  | { type: "status_request" };
-
-/** Messages from daemon worker -> main thread. */
-export type DaemonResponse =
-  | {
-      type: "extraction_complete";
-      extractedTurnCount: number;
-      causalCount: number;
-      monologueCount: number;
-      resolvedCount: number;
-      conceptCount: number;
-      correctionCount: number;
-      preferenceCount: number;
-      artifactCount: number;
-      decisionCount: number;
-      skillCount: number;
-      extractedNames?: PriorExtractions;
-    }
-  | { type: "status"; extractedTurns: number; pendingBatches: number; errors: number }
-  | { type: "shutdown_complete" }
-  | { type: "error"; message: string };

package/src/deferred-cleanup.ts

@@ -0,0 +1,141 @@
+/**
+ * Deferred Cleanup — extract knowledge from orphaned sessions.
+ *
+ * When the process dies abruptly (Ctrl+C×2), session cleanup never runs.
+ * On next session start, this module finds orphaned sessions (started but
+ * never marked cleanup_completed), loads their turns, runs daemon extraction,
+ * generates a handoff note, and marks them complete.
+ *
+ * Turns are already persisted via afterTurn → ingest. This just processes them.
+ */
+import type { SurrealStore } from "./surreal.js";
+import type { EmbeddingService } from "./embeddings.js";
+import type { CompleteFn } from "./state.js";
+import { buildSystemPrompt, buildTranscript, writeExtractionResults } from "./memory-daemon.js";
+import type { PriorExtractions } from "./daemon-types.js";
+import { swallow } from "./errors.js";
+
+/**
+ * Find and process orphaned sessions. Runs with a 30s total timeout.
+ * Fire-and-forget from session_start — does not block the new session.
+ */
+export async function runDeferredCleanup(
+  store: SurrealStore,
+  embeddings: EmbeddingService,
+  complete: CompleteFn,
+): Promise<number> {
+  if (!store.isAvailable()) return 0;
+
+  const orphaned = await store.getOrphanedSessions(3).catch(() => []);
+  if (orphaned.length === 0) return 0;
+
+  let processed = 0;
+
+  const cleanup = async () => {
+    for (const session of orphaned) {
+      try {
+        await processOrphanedSession(session.id, store, embeddings, complete);
+        processed++;
+      } catch (e) {
+        swallow.warn("deferredCleanup:session", e);
+      }
+    }
+  };
+
+  // 30s timeout — don't hold up the new session forever
+  await Promise.race([
+    cleanup(),
+    new Promise<void>(resolve => setTimeout(resolve, 30_000)),
+  ]);
+
+  return processed;
+}
+
+async function processOrphanedSession(
+  surrealSessionId: string,
+  store: SurrealStore,
+  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
+  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
+     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;
+  }
+
+  // Run daemon extraction
+  const priorState: PriorExtractions = { conceptNames: [], artifactPaths: [], skillNames: [] };
+  const turnData = turns.map(t => ({ role: t.role, text: t.text, tool_name: t.tool_name }));
+  const transcript = buildTranscript(turnData);
+  const systemPrompt = buildSystemPrompt(false, false, priorState);
+
+  try {
+    const response = await complete({
+      system: systemPrompt,
+      messages: [{ role: "user", content: `[TRANSCRIPT]\n${transcript.slice(0, 60000)}` }],
+    });
+
+    const responseText = response.text;
+    const jsonMatch = responseText.match(/\{[\s\S]*\}/);
+    if (jsonMatch) {
+      let result: Record<string, any>;
+      try {
+        result = JSON.parse(jsonMatch[0]);
+      } catch {
+        try {
+          result = JSON.parse(jsonMatch[0].replace(/,\s*([}\]])/g, "$1"));
+        } catch { result = {}; }
+      }
+
+      if (Object.keys(result).length > 0) {
+        const sessionId = surrealSessionId; // Use DB ID as session reference
+        await writeExtractionResults(result, sessionId, store, embeddings, priorState);
+      }
+    }
+  } catch (e) {
+    swallow.warn("deferredCleanup:extraction", e);
+  }
+
+  // Generate handoff note
+  try {
+    const lastTurns = turns.slice(-15);
+    const turnSummary = lastTurns
+      .map(t => `[${t.role}] ${t.text.slice(0, 200)}`)
+      .join("\n");
+
+    const handoffResponse = await complete({
+      system: "Summarize this session for handoff to your next self. What was worked on, what's unfinished, what to remember. 2-3 sentences. Write in first person.",
+      messages: [{ role: "user", content: turnSummary }],
+    });
+
+    const handoffText = handoffResponse.text.trim();
+    if (handoffText.length > 20) {
+      let emb: number[] | null = null;
+      if (embeddings.isAvailable()) {
+        try { emb = await embeddings.embed(handoffText); } catch { /* ok */ }
+      }
+      await store.createMemory(handoffText, emb, 8, "handoff", surrealSessionId);
+    }
+  } catch (e) {
+    swallow.warn("deferredCleanup:handoff", e);
+  }
+
+  // Mark session as cleaned up
+  await store.markSessionEnded(surrealSessionId).catch(e => swallow("deferred:markDone", e));
+}

package/src/handoff-file.ts

@@ -0,0 +1,55 @@
+/**
+ * Sync handoff file — last-resort session continuity bridge.
+ *
+ * When the process dies (Ctrl+C×2), there's no async cleanup window.
+ * This module writes a minimal JSON snapshot synchronously on exit
+ * so the next session's wakeup has context even before deferred
+ * extraction runs.
+ */
+import { readFileSync, writeFileSync, unlinkSync, existsSync } from "node:fs";
+import { join } from "node:path";
+
+const HANDOFF_FILENAME = ".kongbrain-handoff.json";
+
+export interface HandoffFileData {
+  sessionId: string;
+  timestamp: string;
+  userTurnCount: number;
+  lastUserText: string;
+  lastAssistantText: string;
+  unextractedTokens: number;
+}
+
+/**
+ * Synchronously write a handoff file. Safe to call from process.on("exit").
+ */
+export function writeHandoffFileSync(
+  data: HandoffFileData,
+  workspaceDir: string,
+): void {
+  try {
+    const path = join(workspaceDir, HANDOFF_FILENAME);
+    writeFileSync(path, JSON.stringify(data, null, 2), "utf-8");
+  } catch {
+    // Best-effort — sync exit handler, can't log async
+  }
+}
+
+/**
+ * Read and delete the handoff file. Returns null if not found.
+ */
+export function readAndDeleteHandoffFile(
+  workspaceDir: string,
+): HandoffFileData | null {
+  const path = join(workspaceDir, HANDOFF_FILENAME);
+  if (!existsSync(path)) return null;
+  try {
+    const raw = readFileSync(path, "utf-8");
+    unlinkSync(path);
+    return JSON.parse(raw) as HandoffFileData;
+  } catch {
+    // Corrupted or deleted between check and read
+    try { unlinkSync(path); } catch { /* ignore */ }
+    return null;
+  }
+}