supipowers 2.0.2 → 2.2.0

package/src/mempalace/hooks.ts

@@ -3,17 +3,73 @@ import { normalizeSystemPromptBlocks } from "../platform/system-prompt.js";
 import type { SupipowersConfig } from "../types.js";
 import { createMempalaceBridge, type MempalaceBridgeFacade } from "./bridge.js";
 import { resolveDefaultWing, resolveMempalaceConfig, type ResolvedMempalaceConfig } from "./config.js";
+import { resolveInstalledBridgeScriptPath } from "./runtime.js";
 import { getEventStore as getContextEventStore, getSessionId as getContextSessionId } from "../context-mode/hooks.js";
 import { buildCompactionCheckpoint, buildShutdownDiary } from "./session-summary.js";
+import { snapshotMempalaceInstall } from "./installer-helper.js";
 
 export interface MempalaceHooksDeps {
   createBridge?: (config: ResolvedMempalaceConfig, cwd: string) => MempalaceBridgeFacade;
   getEventStore?: () => Parameters<typeof buildCompactionCheckpoint>[0]["eventStore"];
   getSessionId?: () => string;
   now?: () => string;
+  snapshotInstall?: (paths: Platform["paths"], cwd: string, config: SupipowersConfig) => { ready: boolean };
 }
 
-const wakeUpCache = new Map<string, string>();
+/** Maximum number of (sessionId × wing × palace) entries to keep in memory. */
+const HOOK_CACHE_LRU_CAP = 64;
+
+/** Insertion-ordered bounded LRU. Drops the least-recently-used entry on overflow. */
+class BoundedLRU<K, V> {
+  private readonly inner = new Map<K, V>();
+
+  constructor(private readonly cap: number) {}
+
+  get(key: K): V | undefined {
+    if (!this.inner.has(key)) return undefined;
+    const value = this.inner.get(key) as V;
+    this.inner.delete(key);
+    this.inner.set(key, value);
+    return value;
+  }
+
+  set(key: K, value: V): void {
+    if (this.inner.has(key)) {
+      this.inner.delete(key);
+    } else if (this.inner.size >= this.cap) {
+      const oldest = this.inner.keys().next().value;
+      if (oldest !== undefined) this.inner.delete(oldest);
+    }
+    this.inner.set(key, value);
+  }
+
+  delete(key: K): boolean {
+    return this.inner.delete(key);
+  }
+
+  clear(): void {
+    this.inner.clear();
+  }
+
+  keys(): IterableIterator<K> {
+    return this.inner.keys();
+  }
+}
+
+function warnHookStateFallback(platform: Platform, message: string): void {
+  const logger = (platform as { logger?: { warn?: (message: string) => void } }).logger;
+  if (typeof logger?.warn === "function") {
+    logger.warn(message);
+    return;
+  }
+  console.warn(message);
+}
+
+function hookTimeoutSeconds(timeoutMs: number): number {
+  return Math.max(1, Math.floor(timeoutMs / 1000));
+}
+
+const wakeUpCache = new BoundedLRU<string, string>(HOOK_CACHE_LRU_CAP);
 
 /**
  * Per-session turn counter for wake-up cadence gating. The full wake-up block
@@ -21,7 +77,7 @@ const wakeUpCache = new Map<string, string>();
  * `mempalace.budgets.wakeUpInjectionEvery`); other turns get a one-line
  * refresher. Cleared on session_start / session_switch.
  */
-const turnCounters = new Map<string, number>();
+const turnCounters = new BoundedLRU<string, number>(HOOK_CACHE_LRU_CAP);
 
 /** Test-only: reset cadence state between cases. */
 export function _resetMempalaceHookState(): void {
@@ -91,6 +147,15 @@ function setupGuidanceBlock(resolved: ResolvedMempalaceConfig, wing: string): st
   ].join("\n");
 }
 
+function wakeFailureBlock(resolved: ResolvedMempalaceConfig, wing: string, error: string): string {
+  return [
+    "# MemPalace memory",
+    `- palace: ${resolved.palacePath}`,
+    `- default wing: ${wing}`,
+    `- Wake-up failed: ${error}`,
+  ].join("\n");
+}
+
 function wakeUpBlock(resolved: ResolvedMempalaceConfig, wing: string, text: string): string {
   const excerpt = truncateByTokenBudget(text, resolved.budgets.wakeUpTokens);
   const lines = [
@@ -139,19 +204,33 @@ function extractUserPrompt(event: unknown): string {
 /** Minimum prompt length below which we skip auto-search (saves a bridge call). */
 const AUTO_SEARCH_MIN_PROMPT_CHARS = 15;
 
-/** Cap on the search query length. Long prompts are truncated to the first N chars. */
-const AUTO_SEARCH_QUERY_MAX_CHARS = 500;
-
 /**
- * Returns `true` when `prompt` is a trivial acknowledgement that does not warrant
- * a memory search (saves the bridge round-trip for "yes", "ok", "thanks", etc.).
+ * Returns `true` when `prompt` warrants a MemPalace auto-search.
+ * Rules applied in order:
+ *   1. Too-short or trivial filler word → skip (saves the bridge round-trip).
+ *   2. Contains "?" or starts with a question word → search.
+ *   3. Contains a memory/recall signal → search.
+ *   4. Starts with a clearly imperative verb and has no search signal → skip.
+ *   5. Ambiguous → search (preserves recall on uncertain prompts).
  */
-function isTrivialPrompt(prompt: string): boolean {
+function shouldAutoSearchPrompt(prompt: string): boolean {
   const normalized = prompt.toLowerCase().replace(/[^\p{L}\p{N}\s]+/gu, " ").replace(/\s+/g, " ").trim();
-  if (normalized.length < AUTO_SEARCH_MIN_PROMPT_CHARS) return true;
-  // Conservative wordlist: only strip obvious filler.
+  // Rule 1: trivial length or obvious filler.
+  if (normalized.length < AUTO_SEARCH_MIN_PROMPT_CHARS) return false;
   const TRIVIAL = new Set(["yes", "no", "ok", "okay", "thanks", "thank you", "great", "cool", "go", "continue", "proceed"]);
-  return TRIVIAL.has(normalized);
+  if (TRIVIAL.has(normalized)) return false;
+  // Rule 2: explicit question signals → search.
+  if (prompt.includes("?")) return true;
+  const QUESTION_PREFIXES = ["what", "why", "when", "who", "where", "how", "which", "do", "does", "is", "are", "can", "should"];
+  if (QUESTION_PREFIXES.some(p => normalized.startsWith(p + " ") || normalized === p)) return true;
+  // Rule 3: memory/recall signal words → search.
+  const RECALL_SIGNALS = ["remember", "recall", "decided", "decision", "chose", "last time", "previously", "earlier", "before"];
+  if (RECALL_SIGNALS.some(s => normalized.includes(s))) return true;
+  // Rule 4: clearly imperative verb at start, no search signal above → skip.
+  const IMPERATIVE_PREFIXES = ["fix", "add", "remove", "delete", "run", "update", "refactor", "rename", "move", "write", "create", "make", "implement", "build"];
+  if (IMPERATIVE_PREFIXES.some(p => normalized.startsWith(p + " ") || normalized === p)) return false;
+  // Rule 5: ambiguous → search.
+  return true;
 }
 
 interface SearchHit {
@@ -169,13 +248,13 @@ function pickHits(result: unknown): SearchHit[] {
 }
 
 /** Score-gated relevance check so we don't inject low-quality matches as noise. */
-function isRelevantHit(hit: SearchHit): boolean {
+function isRelevantHit(hit: SearchHit, similarityFloor: number, bm25Floor: number): boolean {
   const sim = typeof hit.similarity === "number" ? hit.similarity : null;
   const bm25 = typeof hit.bm25_score === "number" ? hit.bm25_score : null;
-  // Either signal must clear a low bar. Mempalace's `similarity` is ~1.0 for
+  // Either signal must clear the configured floor. similarity is ~1.0 for
   // perfect, ~0.5 for "kinda related"; bm25 is unbounded but >0.3 is meaningful.
-  if (sim !== null && sim >= 0.55) return true;
-  if (bm25 !== null && bm25 >= 0.3) return true;
+  if (sim !== null && sim >= similarityFloor) return true;
+  if (bm25 !== null && bm25 >= bm25Floor) return true;
   return false;
 }
 
@@ -210,12 +289,38 @@ export function registerMempalaceHooks(
 ): void {
   if (!config.mempalace.enabled) return;
 
-  const clearAll = () => {
-    wakeUpCache.clear();
-    turnCounters.clear();
+  const snapshotInstall = deps.snapshotInstall ?? snapshotMempalaceInstall;
+  const isInstallReady = (cwd: string): boolean => snapshotInstall(platform.paths, cwd, config).ready;
+
+  const bridgeRuntime = {
+    resolveBridgeScriptPath: () => resolveInstalledBridgeScriptPath(platform.paths),
+  };
+
+  const clearSessionState = (sessionId: string): void => {
+    for (const key of [...wakeUpCache.keys()]) {
+      if (key.startsWith(`${sessionId}|`)) wakeUpCache.delete(key);
+    }
+    for (const key of [...turnCounters.keys()]) {
+      if (key.startsWith(`${sessionId}|`)) turnCounters.delete(key);
+    }
+  };
+
+  const clearForSession = (event: unknown): void => {
+    const source = typeof event === "object" && event !== null ? event as { sessionId?: unknown; previousSessionId?: unknown } : null;
+    const sessionId = typeof source?.sessionId === "string" && source.sessionId.length > 0 ? source.sessionId : null;
+    const previousSessionId =
+      typeof source?.previousSessionId === "string" && source.previousSessionId.length > 0 ? source.previousSessionId : null;
+    if (sessionId === null && previousSessionId === null) {
+      warnHookStateFallback(platform, "[mempalace hooks] session event missing sessionId — clearing all hook state");
+      wakeUpCache.clear();
+      turnCounters.clear();
+      return;
+    }
+    if (sessionId !== null) clearSessionState(sessionId);
+    if (previousSessionId !== null) clearSessionState(previousSessionId);
   };
-  platform.on("session_start", clearAll);
-  platform.on("session_switch", clearAll);
+  platform.on("session_start", clearForSession);
+  platform.on("session_switch", clearForSession);
 
   platform.on("before_agent_start", async (event: unknown, ctx: unknown) => {
     const wakeUpEnabled = config.mempalace.hooks.wakeUp;
@@ -234,11 +339,14 @@ export function registerMempalaceHooks(
     const sessionId = sessionIdFrom(event, ctx);
     const cacheKey = `${sessionId}|${wing}|${resolved.palacePath}`;
     const basePrompt = currentSystemPromptBlocks(event, ctx);
+    if (!isInstallReady(cwd)) {
+      return appendPrompt(basePrompt, setupGuidanceBlock(resolved, wing));
+    }
     const userPrompt = extractUserPrompt(event);
 
     const bridge = deps.createBridge
       ? deps.createBridge(resolved, cwd)
-      : createMempalaceBridge({ cwd, config: resolved });
+      : createMempalaceBridge({ cwd, config: resolved, runtime: bridgeRuntime });
 
     // Cadence gating: wake-up dump on turn 1 and every Nth turn; refresher
     // otherwise. Saves ~750 tokens/turn average for a default cadence of 10.
@@ -248,42 +356,78 @@ export function registerMempalaceHooks(
     turnCounters.set(turnKey, turnCount);
     const isFullInjectionTurn = turnCount === 1 || turnCount % cadence === 0;
 
-    // Run the cached generic wake_up and the per-prompt search in parallel.
-    // wake_up is cached for the session; auto-search is fresh every turn.
-    const wakePromise = (async (): Promise<string> => {
-      if (!isFullInjectionTurn) return wakeUpRefresher(resolved, wing);
+    // On cadence-gated turns: one bridge.execute call handles both wake and search
+    // (wake_up_and_search). On non-injection turns: refresher string + optional
+    // separate search call. Either path issues at most one bridge call per turn.
+    const timeoutSeconds = hookTimeoutSeconds(resolved.timeouts.hookMs); /* seconds; bridge multiplies by 1000 */
+    const wantsSearch = autoSearchEnabled && shouldAutoSearchPrompt(userPrompt);
+    // Let upstream MemPalace extract the salient question/tail from long prompts.
+    const query = wantsSearch ? userPrompt : undefined;
+
+    let wakeBlock: string;
+    let searchBlock: string | null = null;
+
+    if (!isFullInjectionTurn) {
+      // Non-injection turns: lightweight refresher, no wake bridge call needed.
+      wakeBlock = wakeUpRefresher(resolved, wing);
+      if (wantsSearch) {
+        try {
+          const result = await bridge.execute({ action: "search", query: query!, wing, limit: 3, timeout: timeoutSeconds });
+          if (result.ok) {
+            const { autoSearchSimilarityFloor, autoSearchBm25Floor } = resolved.budgets;
+            const hits = pickHits(result.result).filter(hit => isRelevantHit(hit, autoSearchSimilarityFloor, autoSearchBm25Floor));
+            searchBlock = autoSearchBlock(hits, resolved.budgets.autoSearchTokens);
+          }
+        } catch {
+          // Auto-search is best-effort. A failure here must never block the turn.
+        }
+      }
+    } else {
       const cached = wakeUpCache.get(cacheKey);
-      if (cached) return cached;
-      const wake = await bridge.execute({ action: "wake_up", wing, timeout: resolved.timeouts.hookMs });
-      const block = wake.ok
-        ? wakeUpBlock(resolved, wing, wakeText(wake.result))
-        : setupGuidanceBlock(resolved, wing);
-      wakeUpCache.set(cacheKey, block);
-      return block;
-    })();
-
-    const searchPromise = (async (): Promise<string | null> => {
-      if (!autoSearchEnabled) return null;
-      if (!userPrompt || isTrivialPrompt(userPrompt)) return null;
-      const query = userPrompt.slice(0, AUTO_SEARCH_QUERY_MAX_CHARS);
-      try {
-        const result = await bridge.execute({
-          action: "search",
-          query,
+      if (cached) {
+        // Wake block is already cached — only issue a search call if warranted.
+        wakeBlock = cached;
+        if (wantsSearch) {
+          try {
+            const result = await bridge.execute({ action: "search", query: query!, wing, limit: 3, timeout: timeoutSeconds });
+            if (result.ok) {
+              const { autoSearchSimilarityFloor, autoSearchBm25Floor } = resolved.budgets;
+              const hits = pickHits(result.result).filter(hit => isRelevantHit(hit, autoSearchSimilarityFloor, autoSearchBm25Floor));
+              searchBlock = autoSearchBlock(hits, resolved.budgets.autoSearchTokens);
+            }
+          } catch {
+            // Auto-search is best-effort. A failure here must never block the turn.
+          }
+        }
+      } else {
+        // Cache miss: batch wake + search into one bridge call.
+        const batchResult = await bridge.execute({
+          action: "wake_up_and_search",
           wing,
-          limit: 3,
-          timeout: resolved.timeouts.hookMs,
+          timeout: timeoutSeconds,
+          ...(query !== undefined ? { query, limit: 3 } : {}),
         });
-        if (!result.ok) return null;
-        const hits = pickHits(result.result).filter(isRelevantHit);
-        return autoSearchBlock(hits, resolved.budgets.autoSearchTokens);
-      } catch {
-        // Auto-search is best-effort. A failure here must never block the turn.
-        return null;
+        const composite = batchResult.ok ? (batchResult.result as Record<string, unknown>) : null;
+        const compositeWake = composite !== null ? (composite.wake as Record<string, unknown> | null | undefined) : undefined;
+        const wakeError = composite !== null && typeof composite.wake_error === "string" ? composite.wake_error : "wake_up failed";
+        const block = compositeWake != null
+          ? wakeUpBlock(resolved, wing, wakeText(compositeWake))
+          : batchResult.ok
+            ? wakeFailureBlock(resolved, wing, wakeError)
+            : setupGuidanceBlock(resolved, wing);
+        wakeUpCache.set(cacheKey, block);
+        wakeBlock = block;
+
+        // Extract search hits from the composite result (only if search was requested).
+        const compositeSearch = composite !== null ? composite.search : undefined;
+        if (compositeSearch != null && autoSearchEnabled) {
+          const { autoSearchSimilarityFloor, autoSearchBm25Floor } = resolved.budgets;
+          const hits = pickHits(compositeSearch).filter(hit => isRelevantHit(hit, autoSearchSimilarityFloor, autoSearchBm25Floor));
+          searchBlock = autoSearchBlock(hits, resolved.budgets.autoSearchTokens);
+        }
       }
-    })();
+    }
 
-    const [wakeBlock, searchBlock] = await Promise.all([wakePromise, searchPromise]);
     const combined = searchBlock ? `${wakeBlock}\n${searchBlock}` : wakeBlock;
     return appendPrompt(basePrompt, combined);
   });
@@ -292,6 +436,7 @@ export function registerMempalaceHooks(
     platform.on("session_before_compact", async (_event: unknown, ctx: unknown) => {
       try {
         const cwd = contextCwd(ctx);
+        if (!isInstallReady(cwd)) return undefined;
         const resolved = resolveMempalaceConfig(config, cwd, platform.paths);
         let wing: string;
         try {
@@ -311,7 +456,7 @@ export function registerMempalaceHooks(
         });
         const bridge = deps.createBridge
           ? deps.createBridge(resolved, cwd)
-          : createMempalaceBridge({ cwd, config: resolved });
+          : createMempalaceBridge({ cwd, config: resolved, runtime: bridgeRuntime });
         await bridge.execute({
           action: "add_drawer",
           wing: checkpoint.metadata.wing,
@@ -319,7 +464,7 @@ export function registerMempalaceHooks(
           content: checkpoint.content,
           added_by: checkpoint.metadata.added_by,
           source_file: checkpoint.metadata.source_file,
-          timeout: resolved.timeouts.hookMs,
+          timeout: hookTimeoutSeconds(resolved.timeouts.hookMs), /* seconds; bridge multiplies by 1000 */
         });
       } catch {
         // Compaction must never be cancelled by MemPalace checkpoint failures.
@@ -332,6 +477,7 @@ export function registerMempalaceHooks(
     platform.on("session_shutdown", async (_event: unknown, ctx: unknown) => {
       try {
         const cwd = contextCwd(ctx);
+        if (!isInstallReady(cwd)) return undefined;
         const resolved = resolveMempalaceConfig(config, cwd, platform.paths);
         let wing: string;
         try {
@@ -351,7 +497,7 @@ export function registerMempalaceHooks(
         });
         const bridge = deps.createBridge
           ? deps.createBridge(resolved, cwd)
-          : createMempalaceBridge({ cwd, config: resolved });
+          : createMempalaceBridge({ cwd, config: resolved, runtime: bridgeRuntime });
         await bridge.execute({
           action: "diary_write",
           agent_name: diary.metadata.agent_name,
@@ -359,7 +505,7 @@ export function registerMempalaceHooks(
           topic: diary.metadata.topic,
           entry: diary.entry,
           source_file: diary.metadata.source_file,
-          timeout: resolved.timeouts.hookMs,
+          timeout: hookTimeoutSeconds(resolved.timeouts.hookMs), /* seconds; bridge multiplies by 1000 */
         });
       } catch {
         // Shutdown must never be delayed or failed by MemPalace diary writes.
@@ -367,4 +513,6 @@ export function registerMempalaceHooks(
       return undefined;
     });
   }
+
+  platform.on("session_shutdown", clearForSession);
 }

package/src/mempalace/installer-helper.ts

@@ -6,7 +6,7 @@ import type { SupipowersConfig } from "../types.js";
 import { createMempalaceBridge, type MempalaceBridgeFacade } from "./bridge.js";
 import { resolveDefaultWing, resolveMempalaceConfig } from "./config.js";
 import {
-  resolveBridgeScriptPath,
+  resolveInstalledBridgeScriptPath,
   resolveManagedVenvPaths,
   setupMempalaceRuntime,
   type ProcessRunner,
@@ -44,7 +44,7 @@ export function snapshotMempalaceInstall(
   const managedBinDir = paths.global("bin");
   const uvBinary = process.platform === "win32" ? "uv.exe" : "uv";
   const uvPath = path.join(managedBinDir, uvBinary);
-  const bridge = resolveBridgeScriptPath();
+  const bridge = resolveInstalledBridgeScriptPath(paths);
 
   const uvInstalled = existsSync(uvPath);
   const venvInstalled = existsSync(venv.python);
@@ -85,7 +85,7 @@ export async function runMempalaceSetup(
 ): Promise<SetupMempalaceRuntimeResult> {
   const config = options.config ?? DEFAULT_CONFIG;
   const resolved = resolveMempalaceConfig(config, options.cwd, options.paths);
-  const bridge = resolveBridgeScriptPath();
+  const bridge = resolveInstalledBridgeScriptPath(options.paths);
   if (!bridge.ok) {
     return { ok: false, error: bridge.error };
   }
@@ -109,6 +109,16 @@ export interface MempalaceInitState {
 function isWingPresent(result: unknown, wing: string): boolean {
   if (!result || typeof result !== "object") return false;
   const record = result as Record<string, unknown>;
+
+  // tool_list_wings returns `{ wings: { <name>: <count>, ... } }`. The
+  // dict shape is the canonical one from mempalace.mcp_server. Older
+  // / partial responses may carry array shapes (items/results), so we
+  // accept both rather than coupling tightly.
+  const wings = record.wings;
+  if (wings && typeof wings === "object" && !Array.isArray(wings)) {
+    if (Object.prototype.hasOwnProperty.call(wings, wing)) return true;
+  }
+
   const candidates = [record.wings, record.items, record.results];
   for (const list of candidates) {
     if (!Array.isArray(list)) continue;
@@ -145,7 +155,11 @@ export async function checkMempalaceProjectInitialized(options: {
   } catch {
     wing = "project";
   }
-  const bridge = options.bridge ?? createMempalaceBridge({ cwd: options.cwd, config: resolved });
+  const bridge = options.bridge ?? createMempalaceBridge({
+    cwd: options.cwd,
+    config: resolved,
+    runtime: { resolveBridgeScriptPath: () => resolveInstalledBridgeScriptPath(options.paths) },
+  });
   const result = await bridge.execute({ action: "list_wings" });
   if (!result.ok) {
     return {

package/src/mempalace/python/mempalace_bridge.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 """One-shot JSON bridge for the native supipowers MemPalace tool.
 
-Action mapping (auditable API surface) — wired against MemPalace 3.3.4:
+Action mapping (auditable API surface):
 - version: stdlib importlib.metadata only; does not import MemPalace runtime modules.
 - 29 MCP-equivalent actions dispatch to ``mempalace.mcp_server.tool_<action>``
   (with the ``traverse → tool_traverse_graph`` rename). The mcp_server
@@ -129,6 +129,8 @@ def _handle_version(params: Dict[str, Any], options: Dict[str, Any]) -> Dict[str
 
 def _apply_palace_path(options: Dict[str, Any]) -> None:
     """Set MEMPALACE_PALACE_PATH so MemPalace's MempalaceConfig picks it up."""
+    # NOTE: mutates os.environ — safe only because we run in a fresh process per
+    # call. Any future daemon/persistent-process refactor MUST re-evaluate this.
     palace = options.get("palacePath") or options.get("palace")
     if palace and isinstance(palace, str) and palace:
         os.environ["MEMPALACE_PALACE_PATH"] = os.path.expanduser(palace)
@@ -203,6 +205,42 @@ def _rename(mapping: Dict[str, str]) -> Callable[[Dict[str, Any]], Dict[str, Any
     return extract
 
 
+
+# ── diary_write: source_file embedding ───────────────────────────────────
+#
+# tool_diary_write(agent_name, entry, topic, wing) does not accept source_file.
+# When a source_file is provided, embed it as a deterministic first-line prefix
+# "[source: <source_file>]" in the entry text so that diary_read content
+# carries the linkage. This is the Option-2 convention documented in
+# session-summary.ts: the caller supplies source_file in params; the bridge
+# embeds it; the round-trip proves it via content inspection.
+DIARY_SOURCE_PREFIX = "[source: "
+
+# Upstream sanitize_content default in mempalace.config (3.3.5). Mirrored in
+# src/mempalace/upstream-limits.ts (MEMPALACE_MAX_CONTENT_LENGTH). Bump both
+# when upgrading.
+_MAX_CONTENT_LENGTH = 100_000
+
+
+def _diary_write_extractor(params: Dict[str, Any]) -> Dict[str, Any]:
+    kwargs = _select("agent_name", "topic", "wing")(params)
+    entry: str = params.get("entry") or ""
+    source_file = params.get("source_file")
+    if source_file and isinstance(source_file, str):
+        # MemPalace's tool_diary_write calls sanitize_content(entry), which
+        # raises ValueError if len(entry) > MAX_CONTENT_LENGTH. Adding the
+        # deterministic `[source: ...]\n` prefix can push a TS-valid entry
+        # past that limit and turn a previously-valid request into a domain
+        # failure. Reserve the prefix budget here by clipping the entry tail
+        # so the prefixed payload always fits.
+        prefix = f"{DIARY_SOURCE_PREFIX}{source_file}]\n"
+        budget = max(0, _MAX_CONTENT_LENGTH - len(prefix))
+        if len(entry) > budget:
+            entry = entry[:budget]
+        entry = f"{prefix}{entry}"
+    kwargs["entry"] = entry
+    return kwargs
+
 # ── MCP-equivalent action dispatch ───────────────────────────────────────
 #
 # Each entry maps our action -> (function name in mempalace.mcp_server,
@@ -224,7 +262,7 @@ MCP_TOOL_DISPATCH: Dict[str, "tuple[str, Callable[[Dict[str, Any]], Dict[str, An
     "delete_drawer": ("tool_delete_drawer", _select("drawer_id")),
     # Knowledge graph: MemPalace uses `entity` for the subject in queries.
     "kg_query": ("tool_kg_query", _rename({"subject": "entity", "as_of": "as_of", "direction": "direction"})),
-    "kg_add": ("tool_kg_add", _select("subject", "predicate", "object", "valid_from", "source_closet")),
+    "kg_add": ("tool_kg_add", _select("subject", "predicate", "object", "valid_from", "valid_to", "source_file", "source_drawer_id")),
     "kg_invalidate": ("tool_kg_invalidate", _select("subject", "predicate", "object", "ended")),
     "kg_timeline": ("tool_kg_timeline", _rename({"subject": "entity"})),
     "kg_stats": ("tool_kg_stats", lambda p: {}),
@@ -239,7 +277,9 @@ MCP_TOOL_DISPATCH: Dict[str, "tuple[str, Callable[[Dict[str, Any]], Dict[str, An
     "delete_tunnel": ("tool_delete_tunnel", _select("tunnel_id")),
     # follow_tunnels: MemPalace uses wing/room (not source_wing/source_room).
     "follow_tunnels": ("tool_follow_tunnels", _rename({"source_wing": "wing", "source_room": "room"})),
-    "diary_write": ("tool_diary_write", _select("agent_name", "entry", "topic", "wing")),
+    # diary_write: source_file is embedded as a prefix in the entry text by
+    # _diary_write_extractor — tool_diary_write does not accept source_file natively.
+    "diary_write": ("tool_diary_write", _diary_write_extractor),
     "diary_read": ("tool_diary_read", _select("agent_name", "wing")),
     "hook_settings": ("tool_hook_settings", lambda p: {}),
     "memories_filed_away": ("tool_memories_filed_away", lambda p: {}),
@@ -287,6 +327,84 @@ def _handle_wake_up(params: Dict[str, Any], options: Dict[str, Any]) -> Dict[str
     return _ok({"text": text}, options)
 
 
+
+# ── wake_up_and_search: composite action (saves one python process spawn) ─
+
+
+def _handle_wake_up_and_search(params: Dict[str, Any], options: Dict[str, Any]) -> Dict[str, Any]:
+    """Batch wake_up + optional search into a single python process.
+
+    Returns {"wake": {"text": ...}, "search": <tool_search payload or null>}.
+    Either half may fail independently; the other half is still returned.
+    """
+    _apply_palace_path(options)
+    palace = options.get("palacePath")
+
+    # ── Wake half ─────────────────────────────────────────────────────────────
+    wake_result: Any = None
+    wake_error: "str | None" = None
+    try:
+        layers = _import_or_raise("mempalace.layers")
+        stack = layers.MemoryStack(palace_path=palace) if palace else layers.MemoryStack()
+        wake_kwargs: Dict[str, Any] = {}
+        if params.get("wing"):
+            wake_kwargs["wing"] = params["wing"]
+        text = _wrap_runtime_errors("MemoryStack.wake_up", lambda: stack.wake_up(**wake_kwargs))
+        wake_result = {"text": text}
+    except BridgeDomainError as exc:
+        wake_error = f"{exc.code}: {exc.message}"
+    except Exception as exc:  # pragma: no cover — defensive isolation
+        wake_error = str(exc)
+
+    # ── Search half (skipped when query is absent or empty) ───────────────────
+    search_result: Any = None
+    search_error: "str | None" = None
+    query = params.get("query")
+    if query and isinstance(query, str) and query.strip():
+        try:
+            module = _import_or_raise("mempalace.mcp_server")
+            tool_search = getattr(module, "tool_search", None)
+            if not callable(tool_search):
+                raise BridgeDomainError(
+                    "mempalace_missing",
+                    "mempalace.mcp_server.tool_search is not callable.",
+                    "Upgrade the managed MemPalace runtime via `/supi:memory setup`.",
+                )
+            search_kwargs: Dict[str, Any] = {"query": query}
+            if params.get("wing"):
+                search_kwargs["wing"] = params["wing"]
+            if params.get("room"):
+                search_kwargs["room"] = params["room"]
+            if isinstance(params.get("limit"), int):
+                search_kwargs["limit"] = params["limit"]
+            raw = _wrap_runtime_errors(
+                "mempalace.mcp_server.tool_search",
+                lambda: tool_search(**search_kwargs),
+            )
+            normalized = _to_jsonable(raw)
+            if isinstance(normalized, dict) and normalized.get("ok") is False and isinstance(normalized.get("error"), dict):
+                err = normalized["error"]
+                code = err.get("code") if isinstance(err.get("code"), str) else "mempalace_runtime_error"
+                message = err.get("message") if isinstance(err.get("message"), str) else json.dumps(err, sort_keys=True)
+                search_error = f"{code}: {message}"
+            else:
+                search_result = normalized
+        except BridgeDomainError as exc:
+            # Surface as partial error so the caller can distinguish a real
+            # search failure from "no query / no hits". Mirrors the wake half.
+            search_error = f"{exc.code}: {exc.message}"
+        except Exception as exc:  # pragma: no cover — defensive isolation
+            # Same contract: never let one half kill the other, but the caller
+            # gets a string they can render or log.
+            search_error = str(exc)
+
+    payload: Dict[str, Any] = {"wake": wake_result, "search": search_result}
+    if wake_error is not None:
+        payload["wake_error"] = wake_error
+    if search_error is not None:
+        payload["search_error"] = search_error
+    return _ok(payload, options)
+
 # ── Native CLI args builders ──────────────────────────────────────────────
 
 
@@ -307,6 +425,8 @@ def _make_cli_args_mine(params: Dict[str, Any]) -> "list[str]":
         args.append("--include-ignored")
     if params.get("no_gitignore"):
         args.append("--no-gitignore")
+    if params.get("dry_run"):
+        args.append("--dry-run")
     if params.get("extract"):
         args.append("--extract")
     return args
@@ -331,6 +451,10 @@ def _make_cli_args_repair(params: Dict[str, Any]) -> "list[str]":
         args.append("--yes")
     if params.get("mode"):
         args.extend(["--mode", str(params["mode"])])
+    if params.get("source"):
+        args.extend(["--source", str(params["source"])])
+    if params.get("archive_existing"):
+        args.append("--archive-existing")
     if params.get("dry_run"):
         args.append("--dry-run")
     return args
@@ -396,6 +520,7 @@ def _make_cli_handler(action: str) -> Callable[[Dict[str, Any], Dict[str, Any]],
 DISPATCH: Dict[str, Callable[[Dict[str, Any], Dict[str, Any]], Dict[str, Any]]] = {
     "version": _handle_version,
     "wake_up": _handle_wake_up,
+    "wake_up_and_search": _handle_wake_up_and_search,
 }
 for _action in MCP_TOOL_DISPATCH:
     DISPATCH[_action] = _make_mcp_handler(_action)