fix: honor type caps and MCP restarts

Author: Coding-Dev-Tools

engraphis/core/recall.py

@@ -352,7 +352,10 @@ def recall(self, query: str, flt: Optional[SearchFilter] = None, *, k: int = 8,
             )
             if (
                 not prompt_only
-                or len(recs) >= prompt_target
+                or (
+                    len(recs) >= prompt_target
+                    and _mtype_limits_can_fill(recs, effective_limits, prompt_target)
+                )
                 or arm_candidate_k >= candidate_ceiling
                 or not can_expand
             ):
@@ -1408,6 +1411,25 @@ def _apply_mtype_limits(
     return selected, drops
 
 
+def _mtype_limits_can_fill(
+    records: dict[str, MemoryRecord], limits: dict[MemoryType, int], target: int,
+) -> bool:
+    """Whether the fetched prompt-safe records can fill ``target`` after type caps."""
+    if not limits:
+        return True
+    selected = 0
+    counts: dict[MemoryType, int] = {}
+    for record in records.values():
+        limit = limits.get(record.mtype)
+        if limit is not None and counts.get(record.mtype, 0) >= limit:
+            continue
+        selected += 1
+        counts[record.mtype] = counts.get(record.mtype, 0) + 1
+        if selected >= target:
+            return True
+    return False
+
+
 def _type_aware_rerank_pool(
     candidates: list[Candidate],
     limits: dict[MemoryType, int],

integrations/pi/index.ts

@@ -51,6 +51,7 @@ export default function engraphisPiExtension(pi: ExtensionAPI) {
 	const discoveredActions = new Map<string, DiscoveredAction>();
 
 	const call = async (name: string, args: Record<string, unknown>, signal?: AbortSignal) => {
+		const generation = client.generation();
 		try {
 			const result = await client.callTool(name, args, signal);
 			if (name === "engraphis_discover_actions") {
@@ -65,6 +66,9 @@ export default function engraphisPiExtension(pi: ExtensionAPI) {
 			}
 			return formatMcpResult(result);
 		} catch (error) {
+			// The MCP client closes an unhealthy transport before this catch runs. Capabilities
+			// are signed by that subprocess, so a restart makes every cached action invalid.
+			if (client.generation() !== generation) discoveredActions.clear();
 			if (name === "engraphis_execute_action" && !(error instanceof EngraphisMcpToolError)) {
 				throw new Error(
 					"Engraphis action outcome is unknown because the local connection failed. " +

integrations/pi/src/mcp-client.ts

@@ -53,6 +53,11 @@ export class EngraphisMcpClient {
 
 	constructor(private readonly config: EngraphisRuntimeConfig) {}
 
+	/** Changes whenever this extension closes a transport and invalidates server-issued state. */
+	generation(): number {
+		return this.lifecycle;
+	}
+
 	async connect(): Promise<Client> {
 		if (this.client) return this.client;
 		if (this.connecting) return this.connecting;

integrations/pi/test/extension.test.ts

@@ -126,6 +126,41 @@ test("requires a fresh discovery and explicit Pi approval for every advanced act
 	}
 });
 
+test("clears discovered actions after an MCP transport reset", async () => {
+	const originalCall = EngraphisMcpClient.prototype.callTool;
+	const originalGeneration = EngraphisMcpClient.prototype.generation;
+	let generation = 0;
+	EngraphisMcpClient.prototype.generation = function () { return generation; };
+	EngraphisMcpClient.prototype.callTool = async function (name: string) {
+		if (name === "engraphis_discover_actions") {
+			return { content: [{ type: "text", text: JSON.stringify({ actions: [{
+				capability_id: "cap_restart", canonical_action: "retire",
+				schema_digest: "1234567890abcdef", side_effect: "state_change", title: "Retire memory",
+			}] }) }] };
+		}
+		generation += 1;
+		throw new Error("stdio transport closed");
+	};
+	try {
+		const { tools } = extensionHarness();
+		const discover = tools.find((tool) => tool.name === "engraphis_discover_actions")!;
+		const recall = tools.find((tool) => tool.name === "engraphis_recall_context")!;
+		const execute = tools.find((tool) => tool.name === "engraphis_execute_action")!;
+		await discover.execute("discover", { task: "retire stale memory" }, undefined);
+		await assert.rejects(recall.execute("recall", { query: "trigger reset" }, undefined));
+		await assert.rejects(
+			execute.execute("action", {
+				arguments: {}, capability_id: "cap_restart", schema_digest: "1234567890abcdef",
+			}, undefined, undefined, { hasUI: true, ui: { confirm: async () => true } }),
+			/not issued by the current Engraphis discovery session/,
+		);
+	} finally {
+		EngraphisMcpClient.prototype.callTool = originalCall;
+		EngraphisMcpClient.prototype.generation = originalGeneration;
+	}
+});
+
+
 test("fails closed when Pi cannot present an action approval dialog", async () => {
 	const original = EngraphisMcpClient.prototype.callTool;
 	EngraphisMcpClient.prototype.callTool = async function (name: string) {

tests/test_recall.py

@@ -1,7 +1,7 @@
 from engraphis.backends import DeterministicEmbedder, NumpyVectorIndex
 from engraphis.backends.reranker import IdentityReranker
-from engraphis.core.interfaces import MemoryRecord, Scope, SearchFilter
-from engraphis.core.recall import RecallEngine, _absolute_retrieval_support
+from engraphis.core.interfaces import MemoryRecord, MemoryType, Scope, SearchFilter
+from engraphis.core.recall import RecallEngine, _absolute_retrieval_support, _mtype_limits_can_fill
 from engraphis.core.retrieval_policy import ProfileConfig
 from engraphis.core.store import Store
 
@@ -13,6 +13,17 @@ class _SemanticTestEmbedder(DeterministicEmbedder):
     embedding_mode = "semantic"
 
 
+def test_prompt_candidate_expansion_accounts_for_memory_type_caps():
+    semantic = MemoryRecord(id="mem_semantic", content="", mtype=MemoryType.SEMANTIC)
+    procedural = MemoryRecord(id="mem_procedural", content="", mtype=MemoryType.PROCEDURAL)
+    limits = {MemoryType.SEMANTIC: 0}
+
+    assert not _mtype_limits_can_fill({semantic.id: semantic}, limits, 1)
+    assert _mtype_limits_can_fill(
+        {semantic.id: semantic, procedural.id: procedural}, limits, 1,
+    )
+
+
 def _engine():
     store = Store(":memory:")
     emb = DeterministicEmbedder(256)