pi-observational-memory 2.1.3 → 2.3.0

package/README.md

@@ -112,6 +112,16 @@ Settings live in Pi's `settings.json` — globally at `~/.pi/agent/settings.json
 }
 ```
 
+To run the background memory work (observer, reflector, pruner) on a cheaper / faster model than your main coding agent — often the single biggest cost lever the extension exposes — add `compactionModel`:
+
+```json
+{
+  "observational-memory": {
+    "compactionModel": { "provider": "openrouter", "id": "google/gemma-4-31b-it" }
+  }
+}
+```
+
 The five settings most worth knowing:
 
 | Setting | Default | What it controls |

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-observational-memory",
-  "version": "2.1.3",
+  "version": "2.3.0",
   "description": "Observational memory extension for pi — cache-friendly tiered compaction with observations and reflections.",
   "type": "module",
   "license": "MIT",
@@ -30,18 +30,22 @@
     "README.md"
   ],
   "scripts": {
-    "typecheck": "tsc --noEmit"
+    "typecheck": "tsc --noEmit",
+    "test": "vitest run"
   },
   "peerDependencies": {
-    "@mariozechner/pi-coding-agent": "*",
+    "@mariozechner/pi-agent-core": "*",
     "@mariozechner/pi-ai": "*",
-    "@mariozechner/pi-agent-core": "*"
+    "@mariozechner/pi-coding-agent": "*",
+    "@mariozechner/pi-tui": "*"
   },
   "devDependencies": {
-    "@mariozechner/pi-ai": "^0.66.1",
     "@mariozechner/pi-agent-core": "^0.66.1",
+    "@mariozechner/pi-ai": "^0.66.1",
     "@mariozechner/pi-coding-agent": "^0.66.1",
+    "@mariozechner/pi-tui": "^0.66.1",
     "@types/node": "^22.0.0",
-    "typescript": "^5.6.0"
+    "typescript": "^5.6.0",
+    "vitest": "^4.1.5"
   }
 }

package/src/branch.ts

@@ -9,7 +9,7 @@ import {
 } from "./types.js";
 import { estimateEntryTokens } from "./tokens.js";
 
-type Entry = {
+export type Entry = {
 	type: string;
 	id: string;
 	timestamp?: string;
@@ -25,10 +25,50 @@ type Entry = {
 
 const RAW_TYPES = new Set(["message", "custom_message", "branch_summary"]);
 
+export function isSourceEntry(entry: Entry): boolean {
+	return RAW_TYPES.has(entry.type);
+}
+
 function isObservationEntry(entry: Entry): boolean {
 	return entry.type === "custom" && entry.customType === OBSERVATION_CUSTOM_TYPE;
 }
 
+export type RecallObservationMatch =
+	| {
+			status: "ok";
+			observation: ObservationRecord;
+			observationEntryId: string;
+			sourceEntryIds: string[];
+			sourceEntries: Entry[];
+		}
+	| {
+			status: "no_source";
+			observation: ObservationRecord;
+			observationEntryId: string;
+		}
+	| {
+			status: "source_unavailable";
+			observation: ObservationRecord;
+			observationEntryId: string;
+			sourceEntryIds: string[];
+			missingSourceEntryIds: string[];
+			nonSourceEntryIds: string[];
+		};
+
+export type RecallObservationSourcesResult =
+	| {
+			status: "not_found";
+			observationId: string;
+			matches: [];
+			collision: false;
+		}
+	| {
+			status: "found";
+			observationId: string;
+			matches: RecallObservationMatch[];
+			collision: boolean;
+		};
+
 export function findLastCompactionIndex(entries: Entry[]): number {
 	for (let i = entries.length - 1; i >= 0; i--) {
 		if (entries[i].type === "compaction") return i;
@@ -103,11 +143,89 @@ export function rawTailEntriesBetween(entries: Entry[], fromId: string, untilId:
 
 	const result: Entry[] = [];
 	for (let i = fromIdx; i <= untilIdx; i++) {
-		if (RAW_TYPES.has(entries[i].type)) result.push(entries[i]);
+		if (isSourceEntry(entries[i])) result.push(entries[i]);
 	}
 	return result;
 }
 
+function uniqueIds(ids: string[]): string[] {
+	return Array.from(new Set(ids));
+}
+
+function resolveSourceEntries(entries: Entry[], sourceEntryIds: string[]): {
+	status: "ok" | "source_unavailable";
+	sourceEntryIds: string[];
+	sourceEntries: Entry[];
+	missingSourceEntryIds: string[];
+	nonSourceEntryIds: string[];
+} {
+	const requested = uniqueIds(sourceEntryIds);
+	const requestedSet = new Set(requested);
+	const entriesById = new Map(entries.map((entry) => [entry.id, entry]));
+	const missingSourceEntryIds = requested.filter((id) => !entriesById.has(id));
+	const nonSourceEntryIds = requested.filter((id) => {
+		const entry = entriesById.get(id);
+		return entry !== undefined && !isSourceEntry(entry);
+	});
+	if (missingSourceEntryIds.length > 0 || nonSourceEntryIds.length > 0) {
+		return {
+			status: "source_unavailable",
+			sourceEntryIds: requested,
+			sourceEntries: [],
+			missingSourceEntryIds,
+			nonSourceEntryIds,
+		};
+	}
+
+	const sourceEntries = entries.filter((entry) => requestedSet.has(entry.id));
+	return {
+		status: "ok",
+		sourceEntryIds: sourceEntries.map((entry) => entry.id),
+		sourceEntries,
+		missingSourceEntryIds: [],
+		nonSourceEntryIds: [],
+	};
+}
+
+export function recallObservationSources(entries: Entry[], observationId: string): RecallObservationSourcesResult {
+	const matches: RecallObservationMatch[] = [];
+	for (const entry of entries) {
+		if (!isObservationEntry(entry)) continue;
+		if (!isObservationEntryData(entry.data)) continue;
+		for (const observation of entry.data.records) {
+			if (observation.id !== observationId) continue;
+			if (!observation.sourceEntryIds || observation.sourceEntryIds.length === 0) {
+				matches.push({ status: "no_source", observation, observationEntryId: entry.id });
+				continue;
+			}
+
+			const resolved = resolveSourceEntries(entries, observation.sourceEntryIds);
+			if (resolved.status === "source_unavailable") {
+				matches.push({
+					status: "source_unavailable",
+					observation,
+					observationEntryId: entry.id,
+					sourceEntryIds: resolved.sourceEntryIds,
+					missingSourceEntryIds: resolved.missingSourceEntryIds,
+					nonSourceEntryIds: resolved.nonSourceEntryIds,
+				});
+				continue;
+			}
+
+			matches.push({
+				status: "ok",
+				observation,
+				observationEntryId: entry.id,
+				sourceEntryIds: resolved.sourceEntryIds,
+				sourceEntries: resolved.sourceEntries,
+			});
+		}
+	}
+
+	if (matches.length === 0) return { status: "not_found", observationId, matches: [], collision: false };
+	return { status: "found", observationId, matches, collision: matches.length > 1 };
+}
+
 function getPriorMemoryDetails(entries: Entry[]): MemoryDetails | undefined {
 	const idx = findLastCompactionIndex(entries);
 	if (idx === -1) return undefined;

package/src/compaction.ts

@@ -323,7 +323,7 @@ export function renderSummary(reflections: Reflection[], observations: Observati
 		parts.push(`## Reflections\n${reflections.join("\n")}`);
 	}
 	if (observations.length > 0) {
-		const body = observations.map((o) => `${o.timestamp} [${o.relevance}] ${o.content}`).join("\n");
+		const body = observationsToPromptLines(observations).join("\n");
 		parts.push(`## Observations\n${body}`);
 	}
 

package/src/hooks/compaction-hook.ts

@@ -8,7 +8,7 @@ import {
 import { renderSummary, runPruner, runReflector } from "../compaction.js";
 import { observationsToPromptLines, runObserver } from "../observer.js";
 import type { Runtime } from "../runtime.js";
-import { serializeBranchEntries } from "../serialize.js";
+import { serializeSourceAddressedBranchEntries } from "../serialize.js";
 import { estimateStringTokens } from "../tokens.js";
 import {
 	OBSERVATION_CUSTOM_TYPE,
@@ -58,8 +58,8 @@ export function registerCompactionHook(pi: ExtensionAPI, runtime: Runtime): void
 			let gapObservationData: ObservationEntryData | null = null;
 			const gap = gapRawEntries(entries, firstKeptEntryId);
 			if (gap.length > 0) {
-				const gapChunk = serializeBranchEntries(gap);
-				if (gapChunk.trim()) {
+				const { text: gapChunk, sourceEntryIds } = serializeSourceAddressedBranchEntries(gap);
+				if (gapChunk.trim() && sourceEntryIds.length > 0) {
 					const gapFromId = gap[0].id;
 					const gapUpToId = gap[gap.length - 1].id;
 					const priorObservationLines = observationsToPromptLines([
@@ -79,6 +79,7 @@ export function registerCompactionHook(pi: ExtensionAPI, runtime: Runtime): void
 						priorReflections: memoryState.reflections,
 						priorObservations: priorObservationLines,
 						chunk: gapChunk,
+						allowedSourceEntryIds: sourceEntryIds,
 						signal,
 					});
 					const gapPromise: Promise<void> = gapCall.then(() => undefined, () => undefined);

package/src/hooks/observer-trigger.ts

@@ -8,7 +8,7 @@ import {
 } from "../branch.js";
 import { observationsToPromptLines, runObserver } from "../observer.js";
 import type { Runtime } from "../runtime.js";
-import { serializeBranchEntries } from "../serialize.js";
+import { serializeSourceAddressedBranchEntries } from "../serialize.js";
 import { estimateStringTokens } from "../tokens.js";
 import { OBSERVATION_CUSTOM_TYPE, type ObservationEntryData } from "../types.js";
 
@@ -34,8 +34,8 @@ export function registerObserverTrigger(pi: ExtensionAPI, runtime: Runtime): voi
 
 		const chunkEntries = rawTailEntriesBetween(entries, coversFromId, coversUpToId);
 		if (chunkEntries.length === 0) return;
-		const chunk = serializeBranchEntries(chunkEntries);
-		if (!chunk.trim()) return;
+		const { text: chunk, sourceEntryIds } = serializeSourceAddressedBranchEntries(chunkEntries);
+		if (!chunk.trim() || sourceEntryIds.length === 0) return;
 
 		if (ctx.hasUI) ctx.ui.notify(
 			`Observational memory: observer running on ~${tokens.toLocaleString()}-token chunk`,
@@ -63,6 +63,7 @@ export function registerObserverTrigger(pi: ExtensionAPI, runtime: Runtime): voi
 				priorReflections: reflections,
 				priorObservations: priorObservationLines,
 				chunk,
+				allowedSourceEntryIds: sourceEntryIds,
 			});
 			if (!records || records.length === 0) {
 				if (ctx.hasUI && ctx.ui) ctx.ui.notify(

package/src/index.ts

@@ -5,6 +5,7 @@ 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";
+import { registerRecallTool } from "./tools/recall-observation.js";
 
 export default function observationalMemory(pi: ExtensionAPI) {
 	const runtime = new Runtime();
@@ -15,4 +16,5 @@ export default function observationalMemory(pi: ExtensionAPI) {
 
 	registerStatusCommand(pi, runtime);
 	registerViewCommand(pi, runtime);
+	registerRecallTool(pi);
 }

package/src/observer.ts

@@ -14,6 +14,7 @@ interface RunObserverArgs {
 	priorReflections: string[];
 	priorObservations: string[];
 	chunk: string;
+	allowedSourceEntryIds: string[];
 	signal?: AbortSignal;
 }
 
@@ -36,6 +37,15 @@ const RecordObservationsSchema = Type.Object({
 				description: "Single-line plain prose. No markdown, no tags, no embedded timestamp.",
 			}),
 			relevance: RelevanceSchema,
+			sourceEntryIds: Type.Array(
+				Type.String({ minLength: 1 }),
+				{
+					minItems: 1,
+					description:
+						"Exact source entry ids from the chunk that directly support this observation. " +
+						"Use only ids shown in '[Source entry id: ...]' labels; never invent ids.",
+				},
+			),
 		}),
 		{ description: "Batch of new observations. May be empty only if the tool is not called at all." },
 	),
@@ -47,8 +57,25 @@ function joinOrEmpty(items: string[]): string {
 	return items.length ? items.join("\n") : "(none yet)";
 }
 
+export function normalizeSourceEntryIds(
+	sourceEntryIds: readonly string[] | undefined,
+	allowedSourceEntryIds: readonly string[],
+): string[] | undefined {
+	if (!sourceEntryIds || sourceEntryIds.length === 0) return undefined;
+	const allowedOrder = new Map<string, number>();
+	for (let i = 0; i < allowedSourceEntryIds.length; i++) allowedOrder.set(allowedSourceEntryIds[i], i);
+
+	const seen = new Set<string>();
+	for (const id of sourceEntryIds) {
+		if (!allowedOrder.has(id)) return undefined;
+		seen.add(id);
+	}
+	if (seen.size === 0) return undefined;
+	return Array.from(seen).sort((a, b) => (allowedOrder.get(a) ?? 0) - (allowedOrder.get(b) ?? 0));
+}
+
 export async function runObserver(args: RunObserverArgs): Promise<ObservationRecord[] | undefined> {
-	const { model, apiKey, headers, priorReflections, priorObservations, chunk, signal } = args;
+	const { model, apiKey, headers, priorReflections, priorObservations, chunk, allowedSourceEntryIds, signal } = args;
 	const conversation = chunk.trim();
 	if (!conversation) return undefined;
 
@@ -65,7 +92,13 @@ export async function runObserver(args: RunObserverArgs): Promise<ObservationRec
 		execute: async (_id, params: RecordObservationsArgs) => {
 			let added = 0;
 			let duplicates = 0;
+			let rejected = 0;
 			for (const obs of params.observations) {
+				const sourceEntryIds = normalizeSourceEntryIds(obs.sourceEntryIds, allowedSourceEntryIds);
+				if (!sourceEntryIds) {
+					rejected++;
+					continue;
+				}
 				const content = truncateRecordContent(obs.content);
 				const id = hashId(content);
 				if (accumulated.has(id)) {
@@ -77,15 +110,20 @@ export async function runObserver(args: RunObserverArgs): Promise<ObservationRec
 					content,
 					timestamp: obs.timestamp,
 					relevance: obs.relevance as Relevance,
+					sourceEntryIds,
 				});
 				added++;
 			}
+			const rejectedPart = rejected > 0
+				? ` ${rejected} observation${rejected === 1 ? "" : "s"} rejected for missing or invalid sourceEntryIds.`
+				: "";
 			const ack =
 				`Recorded ${added} new observation${added === 1 ? "" : "s"} ` +
-				(duplicates > 0 ? `(${duplicates} duplicate${duplicates === 1 ? "" : "s"} skipped). ` : ". ") +
-				`Total so far this run: ${accumulated.size}. ` +
+				(duplicates > 0 ? `(${duplicates} duplicate${duplicates === 1 ? "" : "s"} skipped).` : ".") +
+				rejectedPart +
+				` 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 } };
+			return { content: [{ type: "text", text: ack }], details: { added, duplicates, rejected, total: accumulated.size } };
 		},
 	};
 

package/src/prompts.ts

@@ -97,7 +97,7 @@ Your job is to compress a chunk of recent conversation into timestamped, rated o
 You receive:
 - Current reflections (long-lived facts already crystallized).
 - Current observations (already-recorded observations, each shown as "[id] YYYY-MM-DD HH:MM [relevance] content").
-- A new chunk of conversation with inline message timestamps formatted as "[User @ YYYY-MM-DD HH:MM]:", "[Assistant @ ...]:", "[Tool result for <name> @ ...]:".
+- A new chunk of conversation with source entry labels and inline message timestamps. Each source block starts with "[Source entry id: <id>]" followed by content formatted as "[User @ YYYY-MM-DD HH:MM]:", "[Assistant @ ...]:", "[Tool result for <name> @ ...]:", custom messages, or branch summaries.
 - A current local time fallback for observations that have no obvious message timestamp.
 
 How you work:
@@ -110,6 +110,9 @@ How you work:
 What to emit:
 - Produce NEW observations for the new chunk only. Do not restate facts already present in reflections or current observations unless something has materially changed.
 - Use the timestamp from the relevant conversation message. Fall back to current local time ONLY when no message timestamp applies.
+- For every observation, include sourceEntryIds: the smallest exact set of "[Source entry id: ...]" ids that directly support the observation.
+- Never invent source entry ids. Use only ids printed in the chunk. If an observation spans multiple turns or tool results, include every supporting source entry id.
+- Observations with missing, empty, or invalid sourceEntryIds will be rejected and not recorded, so do not call record_observations until you can cite valid source ids.
 - Group repeated similar tool calls into a single observation rather than one per call.
 - Skip routine, low-information events. It is fine to emit zero observations if the chunk carries no new information — in that case, simply do not call the tool and end with a plain-text confirmation.
 
@@ -267,6 +270,8 @@ export function buildPrunerPassGuidance(pass: number, maxPasses: number): string
 export const CONTEXT_USAGE_INSTRUCTIONS = `These are condensed memories from earlier in this session.
 
 - Reflections: stable, long-lived facts about the user, project, decisions, and constraints.
-- Observations: timestamped events from the conversation history, in chronological order.
+- Observations: timestamped events from the conversation history, in chronological order. Observation lines include ids in brackets.
 
-Treat these as past records. When entries conflict, the most recent observation reflects the latest known state. Work that prior observations describe as completed should not be redone unless the user explicitly asks to revisit it.`;
+Treat these as past records. When entries conflict, the most recent observation reflects the latest known state. Work that prior observations describe as completed should not be redone unless the user explicitly asks to revisit it.
+
+When exact source context is needed for precision or traceability, use the recall tool with the relevant observation id. Do not use recall as broad search or inject raw source unless it is needed.`;

package/src/serialize.ts

@@ -14,35 +14,78 @@ function formatTimestamp(v: number | string | undefined): string {
 	return Number.isNaN(d.getTime()) ? "????-??-?? ??:??" : fmtLocal(d);
 }
 
+function formatRecallTimestamp(...values: Array<number | string | undefined>): string {
+	for (const v of values) {
+		if (v === undefined) continue;
+		const d = new Date(v);
+		if (!Number.isNaN(d.getTime())) return fmtLocal(d);
+	}
+	return "Unknown time";
+}
+
+function textAndPlaceholders(
+	content: unknown,
+	options: { omitRedactedThinking?: boolean; includeThinking?: boolean } = {},
+): string {
+	if (typeof content === "string") return content;
+	if (!Array.isArray(content)) return "[non-text content omitted]";
+
+	const parts: string[] = [];
+	for (const block of content as Array<Record<string, unknown>>) {
+		if (!block || typeof block !== "object") {
+			parts.push("[non-text content omitted]");
+			continue;
+		}
+		if (block.type === "text" && typeof block.text === "string") {
+			parts.push(block.text);
+			continue;
+		}
+		if (block.type === "thinking") {
+			if (options.omitRedactedThinking && block.redacted === true) continue;
+			if (options.includeThinking && typeof block.thinking === "string") {
+				parts.push(`[thinking: ${block.thinking}]`);
+				continue;
+			}
+			parts.push("[non-text content omitted]");
+			continue;
+		}
+		if (block.type === "toolCall" && typeof block.name === "string") {
+			parts.push(`[${block.name}(${JSON.stringify(block.arguments ?? {})})]`);
+			continue;
+		}
+		parts.push("[non-text content omitted]");
+	}
+	return parts.join("\n");
+}
+
+function textOnly(content: string | Array<{ type?: string; text?: string }>): string {
+	if (typeof content === "string") return content;
+	return content
+		.filter((b): b is TextContent => b?.type === "text" && typeof b.text === "string")
+		.map((b) => b.text)
+		.join("\n");
+}
+
 export function serializeConversation(messages: Message[]): string {
 	return messages
 		.map((msg): string | null => {
 			const time = formatTimestamp(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");
+				const text = textOnly(msg.content);
 				return `[User @ ${time}]: ${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");
+				const body = textAndPlaceholders(msg.content, {
+					includeThinking: true,
+					omitRedactedThinking: true,
+				})
+					.split("\n")
+					.filter(Boolean)
+					.join("\n");
 				if (!body) return null;
 				return `[Assistant @ ${time}]: ${body}`;
 			}
-			const text = msg.content
-				.filter((b): b is TextContent => b.type === "text")
-				.map((b) => b.text)
-				.join("\n");
+			const text = textOnly(msg.content);
 			return `[Tool result for ${(msg as ToolResultMessage).toolName} @ ${time}]: ${text}`;
 		})
 		.filter((line): line is string => line !== null)
@@ -62,8 +105,9 @@ export function truncateRecordContent(content: string): string {
 	return `${head} … [truncated ${dropped} chars]`;
 }
 
-type RenderableEntry = {
+export type RenderableEntry = {
 	type: string;
+	id?: string;
 	timestamp?: string;
 	message?: unknown;
 	customType?: string;
@@ -71,6 +115,26 @@ type RenderableEntry = {
 	summary?: unknown;
 };
 
+function renderCustomMessage(entry: RenderableEntry, options: { recallFormat: boolean }): string {
+	const time = options.recallFormat ? formatRecallTimestamp(entry.timestamp) : formatTimestamp(entry.timestamp);
+	const text = options.recallFormat
+		? textAndPlaceholders(entry.content)
+		: typeof entry.content === "string"
+			? entry.content
+			: Array.isArray(entry.content)
+				? (entry.content as Array<{ type?: string; text?: string }>)
+						.filter((b) => b?.type === "text" && typeof b.text === "string")
+						.map((b) => b.text as string)
+						.join("\n")
+				: "";
+	if (options.recallFormat) {
+		const origin = entry.customType ? `Custom message (${entry.customType})` : "Custom message";
+		return `[${origin} @ ${time}]: ${text}`;
+	}
+	const tag = entry.customType ? `Custom (${entry.customType})` : "Custom";
+	return `[${tag} @ ${time}]: ${text}`;
+}
+
 export function serializeBranchEntries(entries: RenderableEntry[]): string {
 	const blocks: string[] = [];
 	for (const entry of entries) {
@@ -80,18 +144,7 @@ export function serializeBranchEntries(entries: RenderableEntry[]): string {
 			continue;
 		}
 		if (entry.type === "custom_message") {
-			const time = formatTimestamp(entry.timestamp);
-			let text = "";
-			if (typeof entry.content === "string") {
-				text = entry.content;
-			} else if (Array.isArray(entry.content)) {
-				text = (entry.content as Array<{ type?: string; text?: string }>)
-					.filter((b) => b?.type === "text" && typeof b.text === "string")
-					.map((b) => b.text as string)
-					.join("\n");
-			}
-			const tag = entry.customType ? `Custom (${entry.customType})` : "Custom";
-			blocks.push(`[${tag} @ ${time}]: ${text}`);
+			blocks.push(renderCustomMessage(entry, { recallFormat: false }));
 			continue;
 		}
 		if (entry.type === "branch_summary" && typeof entry.summary === "string") {
@@ -101,3 +154,63 @@ export function serializeBranchEntries(entries: RenderableEntry[]): string {
 	}
 	return blocks.join("\n\n");
 }
+
+export type SourceAddressedSerialization = {
+	text: string;
+	sourceEntryIds: string[];
+};
+
+function isSourceRenderableEntry(entry: RenderableEntry): boolean {
+	return entry.type === "message" || entry.type === "custom_message" || entry.type === "branch_summary";
+}
+
+export function serializeSourceAddressedBranchEntries(entries: RenderableEntry[]): SourceAddressedSerialization {
+	const blocks: string[] = [];
+	const sourceEntryIds: string[] = [];
+	for (const entry of entries) {
+		if (!entry.id || !isSourceRenderableEntry(entry)) continue;
+		const rendered = serializeBranchEntries([entry]);
+		if (!rendered.trim()) continue;
+		sourceEntryIds.push(entry.id);
+		blocks.push(`[Source entry id: ${entry.id}]\n${rendered}`);
+	}
+	return { text: blocks.join("\n\n"), sourceEntryIds };
+}
+
+function renderRecallMessage(entry: RenderableEntry): string | null {
+	if (!entry.message || typeof entry.message !== "object") return null;
+	const msg = entry.message as Message;
+	const time = formatRecallTimestamp(msg.timestamp, entry.timestamp);
+	if (msg.role === "user") {
+		return `[User @ ${time}]: ${textAndPlaceholders(msg.content)}`;
+	}
+	if (msg.role === "assistant") {
+		const body = textAndPlaceholders(msg.content, {
+			includeThinking: true,
+			omitRedactedThinking: true,
+		})
+			.split("\n")
+			.filter(Boolean)
+			.join("\n");
+		if (!body) return null;
+		return `[Assistant @ ${time}]: ${body}`;
+	}
+	return `[Tool result: ${(msg as ToolResultMessage).toolName} @ ${time}]: ${textAndPlaceholders(msg.content)}`;
+}
+
+export function renderRecallSourceEntry(entry: RenderableEntry): string | null {
+	if (entry.type === "message") return renderRecallMessage(entry);
+	if (entry.type === "custom_message") return renderCustomMessage(entry, { recallFormat: true });
+	if (entry.type === "branch_summary" && typeof entry.summary === "string") {
+		const time = formatRecallTimestamp(entry.timestamp);
+		return `[Branch summary @ ${time}]: ${entry.summary}`;
+	}
+	return null;
+}
+
+export function renderRecallSourceEntries(entries: RenderableEntry[]): string {
+	return entries
+		.map(renderRecallSourceEntry)
+		.filter((block): block is string => block !== null && block.trim().length > 0)
+		.join("\n\n");
+}

package/src/tools/recall-observation.ts

@@ -0,0 +1,409 @@
+import { Type } from "@mariozechner/pi-ai";
+import type { Message, ToolResultMessage } from "@mariozechner/pi-ai";
+import { defineTool, type ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import type { AgentToolResult } from "@mariozechner/pi-agent-core";
+import { Text } from "@mariozechner/pi-tui";
+import {
+	recallObservationSources,
+	type Entry,
+	type RecallObservationMatch,
+	type RecallObservationSourcesResult,
+} from "../branch.js";
+import { renderRecallSourceEntries, renderRecallSourceEntry } from "../serialize.js";
+import { estimateEntryTokens } from "../tokens.js";
+import type { ObservationRecord } from "../types.js";
+
+export const RECALL_OBSERVATION_TOOL_NAME = "recall";
+
+const OBSERVATION_ID_PATTERN = /^[a-f0-9]{12}$/;
+
+type RecallObservationToolStatus =
+	| "ok"
+	| "invalid_id"
+	| "not_found"
+	| "no_source"
+	| "source_unavailable";
+
+type ObservationDetails = Pick<ObservationRecord, "id" | "content" | "timestamp" | "relevance">;
+
+export type RecallSourceEntryDetails = {
+	id: string;
+	origin: string;
+	timestamp: string;
+	tokens: number;
+	qualifiers: string[];
+	content?: string;
+};
+
+type RecallObservationMatchDetails = {
+	status: RecallObservationMatch["status"];
+	observationEntryId: string;
+	observation: ObservationDetails;
+	sourceEntryIds?: string[];
+	sourceEntries?: RecallSourceEntryDetails[];
+	missingSourceEntryIds?: string[];
+	nonSourceEntryIds?: string[];
+	sourceCharacterCount?: number;
+};
+
+export type RecallObservationToolDetails = {
+	status: RecallObservationToolStatus;
+	observationId: string;
+	collision: boolean;
+	matches: RecallObservationMatchDetails[];
+	sourceCharacterCount?: number;
+	message?: string;
+};
+
+function pad(n: number): string {
+	return n.toString().padStart(2, "0");
+}
+
+function fmtLocal(d: Date): string {
+	return `${d.getFullYear()}-${pad(d.getMonth() + 1)}-${pad(d.getDate())} ${pad(d.getHours())}:${pad(d.getMinutes())}`;
+}
+
+function formatDisplayTimestamp(...values: Array<number | string | undefined>): string {
+	for (const v of values) {
+		if (v === undefined) continue;
+		const d = new Date(v);
+		if (!Number.isNaN(d.getTime())) return fmtLocal(d);
+	}
+	return "Unknown time";
+}
+
+function textContentBlocks(content: unknown): Array<Record<string, unknown>> {
+	return Array.isArray(content) ? content.filter((block): block is Record<string, unknown> => !!block && typeof block === "object") : [];
+}
+
+function uniqueStrings(items: string[]): string[] {
+	return Array.from(new Set(items));
+}
+
+function sourceOriginAndQualifiers(entry: Entry): { origin: string; timestamp: string; qualifiers: string[] } {
+	if (entry.type === "message" && entry.message && typeof entry.message === "object") {
+		const msg = entry.message as Message;
+		const timestamp = formatDisplayTimestamp(msg.timestamp, entry.timestamp);
+		if (msg.role === "user") return { origin: "User", timestamp, qualifiers: [] };
+		if (msg.role === "assistant") {
+			const toolCalls = uniqueStrings(
+				textContentBlocks(msg.content)
+					.filter((block) => block.type === "toolCall" && typeof block.name === "string")
+					.map((block) => block.name as string),
+			);
+			return {
+				origin: "Assistant",
+				timestamp,
+				qualifiers: toolCalls.length > 0 ? [`tool calls: ${toolCalls.join(", ")}`] : [],
+			};
+		}
+		const toolName = (msg as ToolResultMessage).toolName;
+		return { origin: `Tool result: ${typeof toolName === "string" && toolName ? toolName : "unknown"}`, timestamp, qualifiers: [] };
+	}
+
+	if (entry.type === "custom_message") {
+		return {
+			origin: "Custom message",
+			timestamp: formatDisplayTimestamp(entry.timestamp),
+			qualifiers: typeof entry.customType === "string" && entry.customType ? [`custom: ${entry.customType}`] : [],
+		};
+	}
+
+	if (entry.type === "branch_summary") {
+		return { origin: "Branch summary", timestamp: formatDisplayTimestamp(entry.timestamp), qualifiers: [] };
+	}
+
+	return { origin: entry.type || "Entry", timestamp: formatDisplayTimestamp(entry.timestamp), qualifiers: [] };
+}
+
+function renderSourceEntryContentOnly(entry: Entry): string | undefined {
+	const rendered = renderRecallSourceEntry(entry);
+	return rendered?.replace(/^\[[^\]]+\]:\s?/, "") || undefined;
+}
+
+function sourceEntryDetails(entry: Entry, includeContent: boolean): RecallSourceEntryDetails {
+	const { origin, timestamp, qualifiers } = sourceOriginAndQualifiers(entry);
+	const content = renderSourceEntryContentOnly(entry);
+	return {
+		id: entry.id,
+		origin,
+		timestamp,
+		tokens: estimateEntryTokens(entry),
+		qualifiers,
+		...(includeContent && content ? { content } : {}),
+	};
+}
+
+function observationDetails(observation: ObservationRecord): ObservationDetails {
+	return {
+		id: observation.id,
+		content: observation.content,
+		timestamp: observation.timestamp,
+		relevance: observation.relevance,
+	};
+}
+
+function matchDetails(match: RecallObservationMatch, sourceText?: string, includeSourceContent = true): RecallObservationMatchDetails {
+	if (match.status === "ok") {
+		return {
+			status: "ok",
+			observationEntryId: match.observationEntryId,
+			observation: observationDetails(match.observation),
+			sourceEntryIds: match.sourceEntryIds,
+			sourceEntries: match.sourceEntries.map((entry) => sourceEntryDetails(entry, includeSourceContent)),
+			sourceCharacterCount: sourceText?.length ?? 0,
+		};
+	}
+	if (match.status === "source_unavailable") {
+		return {
+			status: "source_unavailable",
+			observationEntryId: match.observationEntryId,
+			observation: observationDetails(match.observation),
+			sourceEntryIds: match.sourceEntryIds,
+			missingSourceEntryIds: match.missingSourceEntryIds,
+			nonSourceEntryIds: match.nonSourceEntryIds,
+		};
+	}
+	return {
+		status: "no_source",
+		observationEntryId: match.observationEntryId,
+		observation: observationDetails(match.observation),
+	};
+}
+
+function textResult(text: string, details: RecallObservationToolDetails) {
+	return {
+		content: [{ type: "text" as const, text }],
+		details,
+	};
+}
+
+function aggregateStatus(matches: RecallObservationMatch[]): RecallObservationToolStatus {
+	if (matches.some((match) => match.status === "ok")) return "ok";
+	if (matches.some((match) => match.status === "source_unavailable")) return "source_unavailable";
+	return "no_source";
+}
+
+function friendlyNoSourceMessage(observationId: string): string {
+	return `Observation ${observationId} has no source entries associated with it. This can happen for legacy observations created before source recall was available.`;
+}
+
+function friendlySourceUnavailableMessage(match: Extract<RecallObservationMatch, { status: "source_unavailable" }>): string {
+	const missing = match.missingSourceEntryIds.length > 0 ? ` missing: ${match.missingSourceEntryIds.join(", ")}` : "";
+	const nonSource = match.nonSourceEntryIds.length > 0 ? ` non-source: ${match.nonSourceEntryIds.join(", ")}` : "";
+	return `Observation ${match.observation.id} has source entries associated, but some are unavailable on the current branch or are not source-renderable.${missing}${nonSource}`;
+}
+
+function renderFoundResult(result: Extract<RecallObservationSourcesResult, { status: "found" }>): ReturnType<typeof textResult> {
+	const sections: string[] = [];
+	const detailsMatches: RecallObservationMatchDetails[] = [];
+	let sourceCharacterCount = 0;
+
+	if (result.collision) {
+		sections.push(`Multiple observations share id ${result.observationId}; returning all matching source results from the current branch.`);
+	}
+
+	for (const match of result.matches) {
+		if (match.status === "ok") {
+			const sourceText = renderRecallSourceEntries(match.sourceEntries);
+			sourceCharacterCount += sourceText.length;
+			detailsMatches.push(matchDetails(match, sourceText));
+			if (sourceText.trim()) sections.push(sourceText);
+			else sections.push(`Observation ${match.observation.id} has source entries associated, but they rendered no text content.`);
+			continue;
+		}
+
+		if (match.status === "source_unavailable") {
+			detailsMatches.push(matchDetails(match));
+			sections.push(friendlySourceUnavailableMessage(match));
+			continue;
+		}
+
+		detailsMatches.push(matchDetails(match));
+		sections.push(friendlyNoSourceMessage(match.observation.id));
+	}
+
+	const text = sections.join("\n\n");
+	return textResult(text, {
+		status: aggregateStatus(result.matches),
+		observationId: result.observationId,
+		collision: result.collision,
+		matches: detailsMatches,
+		sourceCharacterCount,
+	});
+}
+
+function plural(n: number, singular: string, pluralForm = `${singular}s`): string {
+	return `${n.toLocaleString()} ${n === 1 ? singular : pluralForm}`;
+}
+
+function sourceEntriesFromDetails(details: RecallObservationToolDetails): RecallSourceEntryDetails[] {
+	return details.matches.flatMap((match) => match.sourceEntries ?? []);
+}
+
+function tokenSummary(tokens: number): string {
+	return `~${tokens.toLocaleString()} ${tokens === 1 ? "token" : "tokens"}`;
+}
+
+function statusIcon(details: RecallObservationToolDetails): string {
+	if (details.status === "ok") return details.collision ? "⚠" : "✓";
+	return "×";
+}
+
+function statusSummary(details: RecallObservationToolDetails): string {
+	if (details.status === "invalid_id") return "invalid id";
+	if (details.status === "not_found") return "not found";
+	if (details.status === "source_unavailable") return "source unavailable";
+	if (details.status === "no_source") return "no source";
+	return details.collision ? "recalled · id collision" : "recalled";
+}
+
+export function formatRecallHeaderForTui(details: RecallObservationToolDetails): string {
+	const parts = [`${statusIcon(details)} ${statusSummary(details)}`];
+	if (details.matches.length > 0) parts.push(plural(details.matches.length, "match", "matches"));
+	const sources = sourceEntriesFromDetails(details);
+	if (sources.length > 0) parts.push(plural(sources.length, "source entry", "source entries"));
+	const tokens = sources.reduce((sum, source) => sum + source.tokens, 0);
+	if (tokens > 0) parts.push(tokenSummary(tokens));
+	return parts.join(" · ");
+}
+
+function sourceLabel(source: RecallSourceEntryDetails): string {
+	return source.origin ? `${source.origin[0].toLowerCase()}${source.origin.slice(1)}` : "entry";
+}
+
+function sourceMetadataLine(source: RecallSourceEntryDetails): string {
+	const qualifiers = source.qualifiers.length > 0 ? ` · ${source.qualifiers.join(" · ")}` : "";
+	return `✓ ${sourceLabel(source)} · ${source.timestamp} · entry ${source.id} · ${tokenSummary(source.tokens)}${qualifiers}`;
+}
+
+function observationLine(observation: ObservationDetails): string {
+	return `✓ observation · ${observation.timestamp} · [${observation.relevance}] · ${observation.content}`;
+}
+
+function indentContent(content: string): string {
+	return content
+		.split("\n")
+		.map((line) => `    ${line}`)
+		.join("\n");
+}
+
+function unavailableSourceLine(match: RecallObservationMatchDetails): string {
+	const parts: string[] = [];
+	if (match.missingSourceEntryIds && match.missingSourceEntryIds.length > 0) {
+		parts.push(`missing: ${match.missingSourceEntryIds.join(", ")}`);
+	}
+	if (match.nonSourceEntryIds && match.nonSourceEntryIds.length > 0) {
+		parts.push(`non-source: ${match.nonSourceEntryIds.join(", ")}`);
+	}
+	return `× source unavailable${parts.length > 0 ? ` · ${parts.join(" · ")}` : ""}`;
+}
+
+function matchLines(match: RecallObservationMatchDetails, expanded: boolean): string[] {
+	const lines = [observationLine(match.observation), ""];
+	if (match.status === "ok") {
+		for (const source of match.sourceEntries ?? []) {
+			lines.push(sourceMetadataLine(source));
+			if (expanded && source.content) {
+				lines.push(indentContent(source.content));
+				lines.push("");
+			}
+		}
+		return lines;
+	}
+	if (match.status === "source_unavailable") return [...lines, unavailableSourceLine(match)];
+	return [...lines, "× no source · legacy/unattributed observation"];
+}
+
+export function formatRecallResultForTui(result: AgentToolResult<RecallObservationToolDetails>, expanded: boolean): string {
+	const details = result.details;
+	if (!details) {
+		const text = result.content
+			.filter((part): part is { type: "text"; text: string } => part.type === "text" && typeof part.text === "string")
+			.map((part) => part.text)
+			.join("\n");
+		return text || "recall";
+	}
+
+	const lines: string[] = [];
+	if (details.matches.length > 0) {
+		for (const match of details.matches) {
+			if (lines.length > 0) lines.push("");
+			lines.push(...matchLines(match, expanded));
+		}
+	} else if (details.message) {
+		lines.push(details.message);
+	}
+	if (!expanded && details.matches.some((match) => match.status === "ok" && (match.sourceEntries?.length ?? 0) > 0)) {
+		lines.push("", "(Ctrl+O to expand)");
+	}
+	return lines.join("\n").trimEnd();
+}
+
+export function formatRecallCallForTui(id: string | undefined): string {
+	return `recall ${id ?? "..."}`;
+}
+
+export function formatRecallRenderedResultForTui(result: AgentToolResult<RecallObservationToolDetails>, expanded: boolean): string {
+	const body = formatRecallResultForTui(result, expanded);
+	const header = result.details ? formatRecallHeaderForTui(result.details) : undefined;
+	if (header && body) return `\n${header}\n\n${body}`;
+	if (header) return `\n${header}`;
+	return body ? `\n${body}` : "";
+}
+
+export const recallObservationTool = defineTool({
+	name: RECALL_OBSERVATION_TOOL_NAME,
+	label: "Recall observation source",
+	description: "Recall exact source entries for an observational-memory observation id on the current branch.",
+	promptSnippet: "Recall exact source entries for a compacted observational-memory observation id.",
+	promptGuidelines: [
+		"Use recall when a compacted observation id needs exact source context or the user asks what supports a remembered claim.",
+		"This is not general search: pass a specific observation id from the compacted Observations list.",
+		"Do not call recall for broad transcript browsing or off-branch history.",
+	],
+	parameters: Type.Object({
+		id: Type.String({
+			pattern: "^[a-f0-9]{12}$",
+			description: "12-character lowercase hex observational-memory observation id.",
+		}),
+	}),
+	renderCall(args) {
+		return new Text(formatRecallCallForTui(args.id), 0, 0);
+	},
+	renderResult(result, options) {
+		return new Text(formatRecallRenderedResultForTui(result as AgentToolResult<RecallObservationToolDetails>, options.expanded), 0, 0);
+	},
+	async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
+		const observationId = params.id;
+		if (!OBSERVATION_ID_PATTERN.test(observationId)) {
+			const message = `Observation id must be 12 lowercase hex characters. Received: ${observationId}`;
+			return textResult(message, {
+				status: "invalid_id",
+				observationId,
+				collision: false,
+				matches: [],
+				message,
+			});
+		}
+
+		const branchEntries = ctx.sessionManager.getBranch() as Entry[];
+		const result = recallObservationSources(branchEntries, observationId);
+		if (result.status === "not_found") {
+			const message = `No observation with id ${observationId} was found on the current branch.`;
+			return textResult(message, {
+				status: "not_found",
+				observationId,
+				collision: false,
+				matches: [],
+				message,
+			});
+		}
+
+		return renderFoundResult(result);
+	},
+});
+
+export function registerRecallTool(pi: ExtensionAPI): void {
+	pi.registerTool(recallObservationTool);
+}

package/src/types.ts

@@ -9,6 +9,7 @@ export interface ObservationRecord {
 	content: string;
 	timestamp: string;
 	relevance: Relevance;
+	sourceEntryIds?: string[];
 }
 
 export type Reflection = string;
@@ -34,11 +35,19 @@ function isRelevance(v: unknown): v is Relevance {
 function isObservationRecord(v: unknown): v is ObservationRecord {
 	if (!v || typeof v !== "object") return false;
 	const o = v as Record<string, unknown>;
+	if (
+		typeof o.id !== "string" ||
+		typeof o.content !== "string" ||
+		typeof o.timestamp !== "string" ||
+		!isRelevance(o.relevance)
+	) {
+		return false;
+	}
+	if (o.sourceEntryIds === undefined) return true;
 	return (
-		typeof o.id === "string" &&
-		typeof o.content === "string" &&
-		typeof o.timestamp === "string" &&
-		isRelevance(o.relevance)
+		Array.isArray(o.sourceEntryIds) &&
+		o.sourceEntryIds.length > 0 &&
+		o.sourceEntryIds.every((id) => typeof id === "string" && id.length > 0)
 	);
 }