pi-observational-memory 1.0.4 → 2.1.1

package/src/hooks/observer-trigger.ts

@@ -0,0 +1,89 @@
+import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import {
+	firstRawIdAfter,
+	getMemoryState,
+	lastObservationCoverEndIdx,
+	rawTailEntriesBetween,
+	rawTokensSinceLastBound,
+} from "../branch.js";
+import { observationsToPromptLines, runObserver } from "../observer.js";
+import type { Runtime } from "../runtime.js";
+import { serializeBranchEntries } from "../serialize.js";
+import { estimateStringTokens } from "../tokens.js";
+import { OBSERVATION_CUSTOM_TYPE, type ObservationEntryData } from "../types.js";
+
+export function registerObserverTrigger(pi: ExtensionAPI, runtime: Runtime): void {
+	pi.on("turn_end", (_event, ctx) => {
+		runtime.ensureConfig(ctx.cwd);
+		if (runtime.observerInFlight) return;
+
+		const entries = ctx.sessionManager.getBranch() as Parameters<typeof rawTokensSinceLastBound>[0];
+		const tokens = rawTokensSinceLastBound(entries);
+		if (tokens < runtime.config.observationThresholdTokens) return;
+
+		const lastBoundIdx = lastObservationCoverEndIdx(entries);
+		const coversFromId = firstRawIdAfter(entries, lastBoundIdx);
+		if (!coversFromId) return;
+
+		const leafId = ctx.sessionManager.getLeafId();
+		if (!leafId) return;
+		const coversUpToId = leafId;
+
+		const { reflections, committedObs, pendingObs } = getMemoryState(entries);
+		const priorObservationLines = observationsToPromptLines([...committedObs, ...pendingObs]);
+
+		const chunkEntries = rawTailEntriesBetween(entries, coversFromId, coversUpToId);
+		if (chunkEntries.length === 0) return;
+		const chunk = serializeBranchEntries(chunkEntries);
+		if (!chunk.trim()) return;
+
+		if (ctx.hasUI) ctx.ui.notify(
+			`Observational memory: observer running on ~${tokens.toLocaleString()}-token chunk`,
+			"info",
+		);
+
+		void runtime.launchObserverTask(ctx, "observer", async () => {
+			const resolved = await runtime.resolveModel(ctx as any);
+			if (!resolved.ok) {
+				if (!runtime.resolveFailureNotified && ctx.hasUI && ctx.ui) {
+					ctx.ui.notify(
+						`Observational memory: observer skipped — ${resolved.reason}`,
+						"warning",
+					);
+					runtime.resolveFailureNotified = true;
+				}
+				return;
+			}
+			runtime.resolveFailureNotified = false;
+
+			const records = await runObserver({
+				model: resolved.model as any,
+				apiKey: resolved.apiKey,
+				headers: resolved.headers,
+				priorReflections: reflections,
+				priorObservations: priorObservationLines,
+				chunk,
+			});
+			if (!records || records.length === 0) {
+				if (ctx.hasUI && ctx.ui) ctx.ui.notify(
+					"Observational memory: observer returned no observations",
+					"warning",
+				);
+				return;
+			}
+
+			const observationTokens = records.reduce((sum, r) => sum + estimateStringTokens(r.content), 0);
+			const data: ObservationEntryData = {
+				records,
+				coversFromId,
+				coversUpToId,
+				tokenCount: observationTokens,
+			};
+			pi.appendEntry(OBSERVATION_CUSTOM_TYPE, data);
+			if (ctx.hasUI && ctx.ui) ctx.ui.notify(
+				`Observational memory: ${records.length} observation${records.length === 1 ? "" : "s"} recorded (~${observationTokens.toLocaleString()} tokens)`,
+				"info",
+			);
+		});
+	});
+}

package/src/ids.ts

@@ -0,0 +1,5 @@
+import { createHash } from "node:crypto";
+
+export function hashId(content: string): string {
+	return createHash("sha256").update(content).digest("hex").slice(0, 12);
+}

package/src/index.ts

@@ -1,336 +1,18 @@
-import { completeSimple, type Message, type TextContent, type ToolCall, type ToolResultMessage } from "@mariozechner/pi-ai";
 import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
-import { convertToLlm, SettingsManager } from "@mariozechner/pi-coding-agent";
-import { DEFAULTS, loadConfig } from "./config.js";
-import type { Config } from "./config.js";
-import { CONTEXT_USAGE_INSTRUCTIONS, OBSERVER_SYSTEM, REFLECTOR_SYSTEM } from "./prompts.js";
-import { estimateRawTailTokens, estimateTokens, extractText } from "./tokens.js";
-import type { MemoryDetails, MemoryState } from "./types.js";
-import { isMemoryDetails } from "./types.js";
-
-function utcDate(epochMs: number): string {
-	if (!Number.isFinite(epochMs)) return "????-??-??";
-	return new Date(epochMs).toISOString().slice(0, 10);
-}
-
-function utcTime(epochMs: number): string {
-	if (!Number.isFinite(epochMs)) return "??:??";
-	return new Date(epochMs).toISOString().slice(11, 16);
-}
-
-function serializeWithTimestamps(messages: Message[]): string {
-	return messages
-		.map((msg): string | null => {
-			const time = utcTime(msg.timestamp);
-			if (msg.role === "user") {
-				const text =
-					typeof msg.content === "string"
-						? msg.content
-						: msg.content
-							.filter((b): b is TextContent => b.type === "text")
-							.map((b) => b.text)
-							.join("\n");
-				return `[User @ ${time} UTC]: ${text}`;
-			}
-			if (msg.role === "assistant") {
-				const parts = msg.content.map((b) => {
-					if (b.type === "text") return b.text;
-					if (b.type === "thinking") return b.redacted ? "" : `[thinking: ${b.thinking}]`;
-					if (b.type === "toolCall") return `[${b.name}(${JSON.stringify(b.arguments)})]`;
-					return "";
-				});
-				const body = parts.filter(Boolean).join("\n");
-				if (!body) return null;
-				return `[Assistant @ ${time} UTC]: ${body}`;
-			}
-			// toolResult
-			const text = msg.content
-				.filter((b): b is TextContent => b.type === "text")
-				.map((b) => b.text)
-				.join("\n");
-			return `[Tool result for ${(msg as ToolResultMessage).toolName} @ ${time} UTC]: ${text}`;
-		})
-		.filter((line): line is string => line !== null)
-		.join("\n\n");
-}
+import { registerStatusCommand } from "./commands/status.js";
+import { registerViewCommand } from "./commands/view.js";
+import { registerCompactionHook } from "./hooks/compaction-hook.js";
+import { registerCompactionTrigger } from "./hooks/compaction-trigger.js";
+import { registerObserverTrigger } from "./hooks/observer-trigger.js";
+import { Runtime } from "./runtime.js";
 
 export default function observationalMemory(pi: ExtensionAPI) {
-	let config: Config = { ...DEFAULTS };
-	let state: MemoryState = { observations: "", reflections: "" };
-	let compactInFlight = false;
-
-	pi.on("session_start", (_event, ctx) => {
-		config = loadConfig(ctx.cwd);
-		state = { observations: "", reflections: "" };
-		compactInFlight = false;
-
-		const entries = ctx.sessionManager.getBranch();
-		for (let i = entries.length - 1; i >= 0; i--) {
-			const entry = entries[i];
-			if (entry.type === "compaction" && isMemoryDetails(entry.details)) {
-				state.observations = entry.details.observations;
-				state.reflections = entry.details.reflections;
-				break;
-			}
-		}
-	});
-
-	pi.on("agent_end", (_event, ctx) => {
-		if (compactInFlight) return;
-
-		const entries = ctx.sessionManager.getBranch();
-		const tokens = estimateRawTailTokens(entries);
-		if (tokens < config.observationThreshold) return;
-
-		compactInFlight = true;
-		setTimeout(() => {
-			if (!ctx.isIdle()) {
-				compactInFlight = false;
-				return;
-			}
-			try {
-				ctx.compact({
-					onComplete: () => {
-						compactInFlight = false;
-						if (ctx.hasUI) ctx.ui.notify("Observational memory: compaction complete", "info");
-					},
-					onError: (error) => {
-						compactInFlight = false;
-						if (ctx.hasUI) ctx.ui.notify(`Observational memory: ${error.message}`, "error");
-					},
-				});
-			} catch (error) {
-				compactInFlight = false;
-				const msg = error instanceof Error ? error.message : String(error);
-				if (ctx.hasUI) ctx.ui.notify(`Observational memory: compact threw: ${msg}`, "error");
-			}
-		}, 0);
-	});
-
-	pi.on("session_before_compact", async (event, ctx) => {
-		const { preparation, signal } = event;
-		const { messagesToSummarize, turnPrefixMessages, firstKeptEntryId, tokensBefore } = preparation;
-
-		let model = ctx.model;
-		if (config.compactionModel) {
-			const configured = ctx.modelRegistry.find(config.compactionModel.provider, config.compactionModel.id);
-			if (configured) {
-				model = configured;
-			} else if (ctx.hasUI) {
-				ctx.ui.notify(
-					`Observational memory: configured model ${config.compactionModel.provider}/${config.compactionModel.id} not found, using session model`,
-					"warning",
-				);
-			}
-		}
-		if (!model) return;
-
-		const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
-		if (!auth.ok || !auth.apiKey) return;
-
-		const allMessages = [...messagesToSummarize, ...turnPrefixMessages];
-		if (allMessages.length === 0) return;
-
-		const now = new Date();
-		const dateStr = utcDate(now.getTime());
-		const timeStr = utcTime(now.getTime());
-
-		const llmMessages = convertToLlm(allMessages);
-		const conversationText = serializeWithTimestamps(llmMessages);
-
-		let dateRangeNote = "";
-		if (llmMessages.length > 0) {
-			const timestamps = llmMessages.map((m) => m.timestamp);
-			const firstTs = timestamps.reduce((a, b) => Math.min(a, b));
-			const lastTs = timestamps.reduce((a, b) => Math.max(a, b));
-			const firstMsgDate = utcDate(firstTs);
-			const lastMsgDate = utcDate(lastTs);
-			dateRangeNote =
-				firstMsgDate === lastMsgDate
-					? ` Messages in this batch are from ${firstMsgDate} (UTC).`
-					: ` Messages in this batch span ${firstMsgDate} to ${lastMsgDate} (UTC).`;
-		}
-
-		ctx.ui.notify("Observational memory: running observer...", "info");
-
-		try {
-			const observerOptions = model.reasoning
-				? { apiKey: auth.apiKey, headers: auth.headers, maxTokens: 4096, signal, reasoning: "high" as const }
-				: { apiKey: auth.apiKey, headers: auth.headers, maxTokens: 4096, signal };
-
-			const observerResponse = await completeSimple(
-				model,
-				{
-					systemPrompt: OBSERVER_SYSTEM,
-					messages: [
-						{
-							role: "user" as const,
-							content: [
-								{
-									type: "text" as const,
-									text: `Today is ${dateStr}, current time is ${timeStr} UTC.${dateRangeNote}\n\n<current-reflections>\n${state.reflections || "(none yet)"}\n</current-reflections>\n\n<current-observations>\n${state.observations || "(none yet)"}\n</current-observations>\n\nCompress the following conversation into new observations:\n\n<conversation>\n${conversationText}\n</conversation>`,
-								},
-							],
-							timestamp: Date.now(),
-						},
-					],
-				},
-				observerOptions,
-			);
-
-			const newObservations = extractText(observerResponse);
-			if (!newObservations.trim()) return;
-
-			state.observations = state.observations
-				? `${state.observations}\n\n${newObservations}`
-				: newObservations;
-		} catch (error) {
-			const msg = error instanceof Error ? error.message : String(error);
-			if (ctx.hasUI) ctx.ui.notify(`Observer failed: ${msg}`, "error");
-			return;
-		}
-
-		if (estimateTokens(state.observations) > config.reflectionThreshold) {
-			ctx.ui.notify("Observational memory: running reflector...", "info");
-
-			try {
-				const reflectorOptions = model.reasoning
-					? { apiKey: auth.apiKey, headers: auth.headers, maxTokens: 8192, signal, reasoning: "high" as const }
-					: { apiKey: auth.apiKey, headers: auth.headers, maxTokens: 8192, signal };
-
-				const reflectorResponse = await completeSimple(
-					model,
-					{
-						systemPrompt: REFLECTOR_SYSTEM,
-						messages: [
-							{
-								role: "user" as const,
-								content: [
-									{
-										type: "text" as const,
-										text: `Today is ${dateStr} (UTC).\n\n<current-reflections>\n${state.reflections || "(none yet)"}\n</current-reflections>\n\n<current-observations>\n${state.observations}\n</current-observations>\n\nGarbage-collect these observations. Promote long-lived facts to reflections, prune what's no longer needed, keep what's still active.`,
-									},
-								],
-								timestamp: Date.now(),
-							},
-						],
-					},
-					reflectorOptions,
-				);
-
-				const output = extractText(reflectorResponse);
-				const reflectionsMatch = output.match(/<reflections>\n?([\s\S]*?)\n?<\/reflections>/);
-				const observationsMatch = output.match(/<observations>\n?([\s\S]*?)\n?<\/observations>/);
-
-				if (reflectionsMatch) state.reflections = reflectionsMatch[1].trim();
-				if (observationsMatch) state.observations = observationsMatch[1].trim();
-			} catch (error) {
-				const msg = error instanceof Error ? error.message : String(error);
-				if (ctx.hasUI) ctx.ui.notify(`Reflector failed: ${msg}`, "warning");
-			}
-		}
-
-		let summary = "";
-		if (state.reflections) {
-			summary += `<reflections>\n${state.reflections}\n</reflections>\n\n`;
-		}
-		if (state.observations) {
-			summary += `<observations>\n${state.observations}\n</observations>`;
-		}
-
-		if (!summary.trim()) return;
-
-		summary += `\n\n${CONTEXT_USAGE_INSTRUCTIONS}`;
-
-		const details: MemoryDetails = {
-			type: "observational-memory",
-			version: 1,
-			observations: state.observations,
-			reflections: state.reflections,
-		};
-
-		return {
-			compaction: {
-				summary,
-				firstKeptEntryId,
-				tokensBefore,
-				details,
-			},
-		};
-	});
-
-	pi.registerCommand("om-status", {
-		description: "Show observational memory status",
-		handler: async (_args, ctx) => {
-			const entries = ctx.sessionManager.getBranch();
-			const rawTokens = estimateRawTailTokens(entries);
-			const obsTokens = estimateTokens(state.observations);
-			const refTokens = estimateTokens(state.reflections);
-			const keepRecentTokens = SettingsManager.create(ctx.cwd).getCompactionKeepRecentTokens();
-
-			const lines = [
-				"── Observational Memory ──",
-				`Raw messages:  ~${rawTokens.toLocaleString()} tokens`,
-				`Observations:  ~${obsTokens.toLocaleString()} tokens`,
-				`Reflections:   ~${refTokens.toLocaleString()} tokens`,
-				"",
-				"── Parameters ──",
-				`Observation threshold: ${config.observationThreshold.toLocaleString()}`,
-				`Reflection threshold:  ${config.reflectionThreshold.toLocaleString()} (interpreted as observations token budget)`,
-				`Keep recent tokens:    ${keepRecentTokens.toLocaleString()} (pi compaction, how many tokens in raw messages)`,
-			];
-
-			ctx.ui.notify(lines.join("\n"), "info");
-		},
-	});
-
-	pi.registerCommand("om-view", {
-		description: "Print full observational memory contents (--full to include raw messages)",
-		handler: async (args, ctx) => {
-			const full = args.includes("--full");
-			const sections: string[] = [];
-
-			sections.push("── Reflections ──");
-			sections.push(state.reflections || "(none)");
-			sections.push("");
-			sections.push("── Observations ──");
-			sections.push(state.observations || "(none)");
-
-			if (full) {
-				const entries = ctx.sessionManager.getBranch();
-				let startIndex = 0;
-				for (let i = entries.length - 1; i >= 0; i--) {
-					const entry = entries[i];
-					if (entry.type === "compaction") {
-						const keptId = entry.firstKeptEntryId;
-						let found = false;
-						for (let j = 0; j < entries.length; j++) {
-							if (entries[j].id === keptId) {
-								startIndex = j;
-								found = true;
-								break;
-							}
-						}
-						if (!found) startIndex = i + 1;
-						break;
-					}
-				}
-
-				const rawMessages = entries
-					.slice(startIndex)
-					.filter((e): e is typeof e & { type: "message"; message: unknown } => e.type === "message")
-					.map((e) => e.message);
+	const runtime = new Runtime();
 
-				sections.push("");
-				sections.push("── Raw Messages ──");
-				if (rawMessages.length > 0) {
-					sections.push(serializeWithTimestamps(convertToLlm(rawMessages)));
-				} else {
-					sections.push("(none)");
-				}
-			}
+	registerObserverTrigger(pi, runtime);
+	registerCompactionTrigger(pi, runtime);
+	registerCompactionHook(pi, runtime);
 
-			ctx.ui.notify(sections.join("\n"), "info");
-		},
-	});
+	registerStatusCommand(pi, runtime);
+	registerViewCommand(pi, runtime);
 }

package/src/observer.ts

@@ -0,0 +1,143 @@
+import { agentLoop, type AgentContext, type AgentLoopConfig, type AgentTool } from "@mariozechner/pi-agent-core";
+import type { Message, Model } from "@mariozechner/pi-ai";
+import { Type } from "@mariozechner/pi-ai";
+import type { Static } from "@sinclair/typebox";
+import { hashId } from "./ids.js";
+import { OBSERVER_SYSTEM } from "./prompts.js";
+import { nowTimestamp, truncateRecordContent } from "./serialize.js";
+import type { ObservationRecord, Relevance } from "./types.js";
+
+interface RunObserverArgs {
+	model: Model<any>;
+	apiKey: string;
+	headers?: Record<string, string>;
+	priorReflections: string[];
+	priorObservations: string[];
+	chunk: string;
+	signal?: AbortSignal;
+}
+
+const RelevanceSchema = Type.Union([
+	Type.Literal("low"),
+	Type.Literal("medium"),
+	Type.Literal("high"),
+	Type.Literal("critical"),
+]);
+
+const RecordObservationsSchema = Type.Object({
+	observations: Type.Array(
+		Type.Object({
+			timestamp: Type.String({
+				pattern: "^\\d{4}-\\d{2}-\\d{2} \\d{2}:\\d{2}$",
+				description: "Observation time in local 'YYYY-MM-DD HH:MM' format.",
+			}),
+			content: Type.String({
+				minLength: 1,
+				description: "Single-line plain prose. No markdown, no tags, no embedded timestamp.",
+			}),
+			relevance: RelevanceSchema,
+		}),
+		{ description: "Batch of new observations. May be empty only if the tool is not called at all." },
+	),
+});
+
+type RecordObservationsArgs = Static<typeof RecordObservationsSchema>;
+
+function joinOrEmpty(items: string[]): string {
+	return items.length ? items.join("\n") : "(none yet)";
+}
+
+export async function runObserver(args: RunObserverArgs): Promise<ObservationRecord[] | undefined> {
+	const { model, apiKey, headers, priorReflections, priorObservations, chunk, signal } = args;
+	const conversation = chunk.trim();
+	if (!conversation) return undefined;
+
+	const accumulated = new Map<string, ObservationRecord>();
+
+	const recordObservations: AgentTool<typeof RecordObservationsSchema> = {
+		name: "record_observations",
+		label: "Record observations",
+		description:
+			"Record a batch of new observations distilled from the conversation chunk. " +
+			"Call this multiple times as you work through the chunk. Stop calling when coverage is complete, " +
+			"then emit a short plain-text confirmation to end the run.",
+		parameters: RecordObservationsSchema,
+		execute: async (_id, params: RecordObservationsArgs) => {
+			let added = 0;
+			let duplicates = 0;
+			for (const obs of params.observations) {
+				const content = truncateRecordContent(obs.content);
+				const id = hashId(content);
+				if (accumulated.has(id)) {
+					duplicates++;
+					continue;
+				}
+				accumulated.set(id, {
+					id,
+					content,
+					timestamp: obs.timestamp,
+					relevance: obs.relevance as Relevance,
+				});
+				added++;
+			}
+			const ack =
+				`Recorded ${added} new observation${added === 1 ? "" : "s"} ` +
+				(duplicates > 0 ? `(${duplicates} duplicate${duplicates === 1 ? "" : "s"} skipped). ` : ". ") +
+				`Total so far this run: ${accumulated.size}. ` +
+				`Continue if the chunk still has uncovered content; otherwise stop calling the tool and emit a short plain-text confirmation.`;
+			return { content: [{ type: "text", text: ack }], details: { added, duplicates, total: accumulated.size } };
+		},
+	};
+
+	const now = nowTimestamp();
+	const userText = `Current local time: ${now}
+
+CURRENT REFLECTIONS:
+${joinOrEmpty(priorReflections)}
+
+CURRENT OBSERVATIONS:
+${joinOrEmpty(priorObservations)}
+
+Compress the following new conversation chunk into observations by calling record_observations one or more times. Do not restate facts already present in current reflections or current observations. Prefer inline conversation timestamps when assigning times; fall back to the current local time above only if no message timestamp applies. Stop calling the tool and reply with a short plain-text confirmation once the chunk is fully covered.
+
+NEW CONVERSATION CHUNK:
+${conversation}`;
+
+	const prompts: Message[] = [
+		{
+			role: "user",
+			content: [{ type: "text", text: userText }],
+			timestamp: Date.now(),
+		},
+	];
+
+	const context: AgentContext = {
+		systemPrompt: OBSERVER_SYSTEM,
+		messages: [],
+		tools: [recordObservations as AgentTool<any>],
+	};
+
+	const reasoning = (model as { reasoning?: unknown }).reasoning;
+	const config: AgentLoopConfig = {
+		model,
+		apiKey,
+		headers,
+		maxTokens: 4096,
+		convertToLlm: (msgs) => msgs as Message[],
+		toolExecution: "sequential",
+		...(reasoning ? { reasoning: "high" as const } : {}),
+	};
+
+	const stream = agentLoop(prompts, context, config, signal);
+	for await (const _event of stream) {
+		// Drain events; the tool's execute already collects records.
+	}
+	await stream.result();
+
+	if (accumulated.size === 0) return undefined;
+	return Array.from(accumulated.values());
+}
+
+export function observationsToPromptLines(records: ObservationRecord[]): string[] {
+	return records.map((r) => `[${r.id}] ${r.timestamp} [${r.relevance}] ${r.content}`);
+}