refactor(integrations): shared agent plugin runtimes and native adapters (#7203)
This commit is contained in:
Kartik
@@ -1,24 +1,17 @@
|
||||
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
|
||||
import MemoryClient from "mem0ai";
|
||||
import { loadConfig, CONFIG_DIR } from "./config/index.ts";
|
||||
import { loadConfig } from "./config/index.ts";
|
||||
import { detectAppId, detectRunId, resolveSearchFilters } from "./memory/scoping.ts";
|
||||
import { formatMemoryList } from "./memory/formatting.ts";
|
||||
import { registerMemoryTool } from "./memory/tools.ts";
|
||||
import { registerCommands } from "./commands.ts";
|
||||
import { setupAutoCapture } from "./capture/index.ts";
|
||||
import { MEMORY_POLICY } from "./prompt.ts";
|
||||
import { DREAM_PROTOCOL } from "./dream/prompt.ts";
|
||||
import {
|
||||
incrementSessionCount,
|
||||
checkCheapGates,
|
||||
checkMemoryGate,
|
||||
acquireDreamLock,
|
||||
releaseDreamLock,
|
||||
recordDreamCompletion,
|
||||
} from "./dream/index.ts";
|
||||
import { captureEvent } from "./telemetry.ts";
|
||||
import * as os from "node:os";
|
||||
import type { ScopeContext } from "./types.ts";
|
||||
import { createMemoryLifecycle } from "../../agent-plugin-core/typescript/src/lifecycle.ts";
|
||||
|
||||
export { buildRecallContext } from "../../agent-plugin-core/typescript/src/lifecycle.ts";
|
||||
|
||||
export function resolveUserId(configUserId: string): string {
|
||||
if (configUserId) return configUserId;
|
||||
@@ -27,31 +20,6 @@ export function resolveUserId(configUserId: string): string {
|
||||
try { return os.userInfo().username; } catch { return "default"; }
|
||||
}
|
||||
|
||||
/**
|
||||
* Build the auto-recall context block for a turn: search memory with the user's
|
||||
* prompt and format the top matches so they are guaranteed in context instead of
|
||||
* relying on the agent to call the tool. Best-effort — returns "" when disabled,
|
||||
* the prompt is blank, nothing matches, or the search fails; it must never block
|
||||
* the turn.
|
||||
*/
|
||||
export async function buildRecallContext(
|
||||
prompt: string,
|
||||
enabled: boolean,
|
||||
search: (query: string) => Promise<{ results?: unknown[] }>,
|
||||
): Promise<string> {
|
||||
if (!enabled) return "";
|
||||
const q = prompt.trim();
|
||||
if (!q) return "";
|
||||
try {
|
||||
const res = await search(q);
|
||||
const memories = (res.results ?? []) as Parameters<typeof formatMemoryList>[0];
|
||||
if (memories.length === 0) return "";
|
||||
return `<mem0-relevant-memories>\nRetrieved automatically for the current request. This is a shallow first pass — search mem0_memory for more if you need it.\n${formatMemoryList(memories)}\n</mem0-relevant-memories>`;
|
||||
} catch {
|
||||
return "";
|
||||
}
|
||||
}
|
||||
|
||||
export default function mem0Extension(pi: ExtensionAPI): void {
|
||||
const config = loadConfig();
|
||||
|
||||
@@ -73,20 +41,21 @@ export default function mem0Extension(pi: ExtensionAPI): void {
|
||||
}
|
||||
|
||||
const telemetryCtx = { apiKey: config.apiKey };
|
||||
const lifecycle = createMemoryLifecycle();
|
||||
|
||||
// ── Register tool + commands + auto-capture ─────────────────────────
|
||||
registerMemoryTool(pi, mem0, config, getScopeCtx, telemetryCtx);
|
||||
registerCommands(pi, mem0, config, getScopeCtx, telemetryCtx);
|
||||
setupAutoCapture(pi, mem0, config, getScopeCtx, telemetryCtx);
|
||||
setupAutoCapture(pi, mem0, config, getScopeCtx, telemetryCtx, lifecycle);
|
||||
|
||||
captureEvent("pi.plugin.registered", {
|
||||
auto_capture: config.autoCapture,
|
||||
dream_enabled: config.dream.enabled,
|
||||
default_scope: config.defaultScope,
|
||||
}, telemetryCtx);
|
||||
|
||||
// ── session_start: detect project + session, reconstruct scope ──────
|
||||
pi.on("session_start", async (_event, ctx) => {
|
||||
lifecycle.beginSession();
|
||||
scopeCtx.appId = detectAppId(ctx.cwd);
|
||||
|
||||
const sessionFile = ctx.sessionManager?.getSessionFile?.();
|
||||
@@ -96,86 +65,29 @@ export default function mem0Extension(pi: ExtensionAPI): void {
|
||||
scopeCtx.userId = config.userId;
|
||||
}
|
||||
|
||||
if (config.dream.enabled) {
|
||||
incrementSessionCount(CONFIG_DIR, scopeCtx.runId);
|
||||
}
|
||||
|
||||
captureEvent("pi.session.start", {}, telemetryCtx);
|
||||
});
|
||||
|
||||
// ── before_agent_start: append memory policy + auto-dream trigger ───
|
||||
let dreamTriggered = false;
|
||||
let dreamChecked = false;
|
||||
|
||||
// ── before_agent_start: append memory policy and recall ─────────────
|
||||
pi.on("before_agent_start", async (event, _ctx) => {
|
||||
let extra = MEMORY_POLICY;
|
||||
|
||||
// Guaranteed retrieval: prefetch memories relevant to this prompt so the
|
||||
// agent always has them, rather than depending on it to call the tool.
|
||||
const recall = await buildRecallContext(
|
||||
const recall = await lifecycle.recall(
|
||||
event.prompt ?? "",
|
||||
config.contextInjection,
|
||||
(q) => mem0.search(q, { filters: resolveSearchFilters("project", scopeCtx) }),
|
||||
);
|
||||
if (recall) extra += "\n\n" + recall;
|
||||
|
||||
if (config.dream.enabled && config.dream.auto && !dreamTriggered && !dreamChecked) {
|
||||
const gates = checkCheapGates(CONFIG_DIR, config.dream);
|
||||
if (gates.proceed) {
|
||||
try {
|
||||
const filters = resolveSearchFilters("project", scopeCtx);
|
||||
const result = await mem0.getAll({ filters });
|
||||
const count = result.count ?? (result.results ?? []).length;
|
||||
dreamChecked = true;
|
||||
const memGate = checkMemoryGate(count, config.dream);
|
||||
|
||||
if (memGate.pass && acquireDreamLock(CONFIG_DIR)) {
|
||||
dreamTriggered = true;
|
||||
extra += "\n\n" + DREAM_PROTOCOL;
|
||||
captureEvent("pi.dream.triggered", { memory_count: count }, telemetryCtx);
|
||||
}
|
||||
} catch {
|
||||
// Transient error — retry next turn
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
systemPrompt: (event.systemPrompt ?? "") + "\n\n" + extra,
|
||||
};
|
||||
});
|
||||
|
||||
// ── agent_end: dream completion check ───────────────────────────────
|
||||
pi.on("agent_end", async (event) => {
|
||||
if (!dreamTriggered) return;
|
||||
|
||||
const messages = event.messages ?? [];
|
||||
const hadWriteAction = messages.some((m) => {
|
||||
if (m.role !== "assistant") return false;
|
||||
const content = Array.isArray(m.content) ? m.content : [];
|
||||
return content.some(
|
||||
(block: any) =>
|
||||
block.type === "tool_use" &&
|
||||
block.name === "mem0_memory" &&
|
||||
["add", "delete", "delete_all"].includes(block.input?.action),
|
||||
);
|
||||
});
|
||||
|
||||
if (hadWriteAction) {
|
||||
recordDreamCompletion(CONFIG_DIR);
|
||||
captureEvent("pi.dream.completed", {}, telemetryCtx);
|
||||
}
|
||||
|
||||
releaseDreamLock(CONFIG_DIR);
|
||||
dreamTriggered = false;
|
||||
});
|
||||
|
||||
// ── session_shutdown: release dream lock if still held ──────────────
|
||||
// ── session_shutdown ────────────────────────────────────────────────
|
||||
pi.on("session_shutdown", async () => {
|
||||
captureEvent("pi.session.stop", {}, telemetryCtx);
|
||||
if (dreamTriggered) {
|
||||
releaseDreamLock(CONFIG_DIR);
|
||||
dreamTriggered = false;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user