@shanepadgett/tau-agent 0.28.1 → 0.30.0

package/extensions/working-memory/index.ts

@@ -0,0 +1,338 @@
+import {
+	defineTool,
+	type ExtensionAPI,
+	type ExtensionContext,
+	type SessionEntry,
+} from "@earendil-works/pi-coding-agent";
+import { emitTauEvent, onTauEvent } from "../../shared/events.ts";
+import { loadTauExtensionSettings } from "../../shared/settings/load.ts";
+import { createToolRowStateStore } from "../../shared/tool-row-state.ts";
+import { executeWorkingMemory, workingMemoryParameters } from "./checkpoint.ts";
+import { projectWorkingMemory } from "./memory.ts";
+import {
+	parseWorkingMemoryNudge,
+	renderWorkingMemoryCall,
+	renderWorkingMemoryNudge,
+	renderWorkingMemoryResult,
+	type WorkingMemoryNudgeDetails,
+} from "./render.ts";
+import workingMemorySettings from "./settings.ts";
+import { replayWorkingMemoryState, WORKING_MEMORY_TOOL, type WorkingMemoryCheckpointDetailsV1 } from "./state.ts";
+
+const NUDGE_TYPE = "tau.working-memory.nudge";
+const BASELINE_TYPE = "tau.working-memory.nudge-baseline";
+const TOOL_DESCRIPTION =
+	"Create a selective hard checkpoint for future model context. Keep valuable referenced evidence and complete tool exchanges, carry file structure as outlines, defer conditionally relevant files without reading them, and provide one compact continuation note. Never checkpoint useful evidence merely to reread it.";
+
+interface NudgeState {
+	anchorToolCallId: string | undefined;
+	suppressedThroughTokens: number | undefined;
+	highestBoundaryTokens: number;
+}
+
+export default function workingMemoryExtension(pi: ExtensionAPI): void {
+	let enabled = false;
+	let generation = 0;
+	let interval = workingMemorySettings.defaults.nudgeEveryTokens;
+	let instructions = workingMemorySettings.defaults.nudgeInstructions;
+	let toolRegistered = false;
+	let commandRegistered = false;
+	let visualRows = new Set<string>();
+	let nudgeState: NudgeState = { anchorToolCallId: undefined, suppressedThroughTokens: 0, highestBoundaryTokens: 0 };
+	const rowState = createToolRowStateStore(pi, "working-memory.tool-row-state");
+
+	pi.registerMessageRenderer<WorkingMemoryNudgeDetails>(NUDGE_TYPE, (message, _options, theme) =>
+		renderWorkingMemoryNudge(message.details, theme),
+	);
+
+	const pushVisualSnapshot = () => {
+		emitTauEvent(pi, "tau:tool-row-state.snapshot", {
+			states: [...visualRows].map((rowId) => ({ rowId, state: "pruned" as const })),
+		});
+	};
+	onTauEvent(
+		pi,
+		"working-memory.tool-row-state-producer",
+		"tau:tool-row-state.snapshot.requested",
+		pushVisualSnapshot,
+	);
+
+	const invalidate = () => {
+		generation += 1;
+	};
+	const setToolActive = (active: boolean) => {
+		if (!toolRegistered) return;
+		const activeTools = pi.getActiveTools();
+		const currentlyActive = activeTools.includes(WORKING_MEMORY_TOOL);
+		if (active === currentlyActive) return;
+		pi.setActiveTools(
+			active ? [...activeTools, WORKING_MEMORY_TOOL] : activeTools.filter((name) => name !== WORKING_MEMORY_TOOL),
+		);
+	};
+	const syncBranch = (ctx: ExtensionContext) => {
+		const branch = ctx.sessionManager.getBranch();
+		const state = replayWorkingMemoryState(branch, enabled);
+		visualRows = new Set(state.prunedRowIds);
+		nudgeState = reconstructNudgeState(branch, state.latestAnchorToolCallId);
+		pushVisualSnapshot();
+	};
+
+	pi.on("session_start", async (_event, ctx) => {
+		invalidate();
+		enabled = false;
+		setToolActive(false);
+		visualRows.clear();
+		pushVisualSnapshot();
+		const current = generation;
+		const settings = await loadTauExtensionSettings(ctx, workingMemorySettings);
+		if (current !== generation) return;
+		enabled = settings.enabled;
+		interval = settings.nudgeEveryTokens;
+		instructions = settings.nudgeInstructions;
+		if (enabled && !toolRegistered) {
+			pi.registerTool(
+				defineTool<typeof workingMemoryParameters, WorkingMemoryCheckpointDetailsV1>({
+					name: WORKING_MEMORY_TOOL,
+					label: WORKING_MEMORY_TOOL,
+					description: TOOL_DESCRIPTION,
+					promptSnippet: "Reassess and selectively checkpoint active working memory",
+					promptGuidelines: [
+						"Use working_memory when stale evidence has accumulated or a memory reminder asks for reassessment; continue coherent exploration when current evidence remains useful.",
+						"A hidden working-memory reference catalog provides keep refs. Message refs preserve conversational evidence; tool refs preserve one complete call/result exchange.",
+						"Keep active or expensive evidence, outline files when structure suffices, defer paths whose relevance is conditional, and discard evidence with no expected value.",
+						"Everything before working_memory leaves future model context unless selected in keep. Put durable decisions, constraints, unresolved matters, and next action in continuation without duplicating retained evidence.",
+					],
+					parameters: workingMemoryParameters,
+					executionMode: "sequential",
+					async execute(toolCallId, params, signal, _onUpdate, ctx) {
+						const execution = await executeWorkingMemory({
+							pi,
+							toolCallId,
+							params,
+							signal,
+							ctx,
+							generation,
+							currentGeneration: () => generation,
+						});
+						for (const outline of execution.outlines) pi.sendMessage(outline, { deliverAs: "steer" });
+						return execution.result;
+					},
+					renderCall(args, theme, context) {
+						return renderWorkingMemoryCall(args, theme, {
+							rowState,
+							rowId: context.toolCallId,
+							invalidate: context.invalidate,
+							lastComponent: context.lastComponent,
+						});
+					},
+					renderResult(result, options, theme, context) {
+						return renderWorkingMemoryResult(result, options.expanded, theme, context.lastComponent);
+					},
+				}),
+			);
+			toolRegistered = true;
+		}
+		setToolActive(enabled);
+		if (enabled && !commandRegistered) {
+			pi.registerCommand("prune", {
+				description: "Ask agent to reassess working memory and continue",
+				async handler(args, commandContext) {
+					if (!enabled) {
+						commandContext.ui.notify("Working memory is disabled.", "info");
+						return;
+					}
+					if (args.trim().length > 0) {
+						commandContext.ui.notify("Usage: /prune", "info");
+						return;
+					}
+					const anchor = replayWorkingMemoryState(
+						commandContext.sessionManager.getBranch(),
+						true,
+					).latestAnchorToolCallId;
+					pi.sendMessage<WorkingMemoryNudgeDetails>(
+						{
+							customType: NUDGE_TYPE,
+							content: manualInstruction(),
+							display: true,
+							details: {
+								v: 1,
+								kind: "manual",
+								tokens: null,
+								boundaryTokens: null,
+								reminder: null,
+								tier: null,
+								tierCount: null,
+								anchorToolCallId: anchor ?? null,
+							},
+						},
+						{ deliverAs: "steer", triggerTurn: true },
+					);
+				},
+			});
+			commandRegistered = true;
+		}
+		syncBranch(ctx);
+	});
+
+	pi.on("before_agent_start", (_event, ctx) => {
+		if (!enabled) return undefined;
+		const usage = ctx.getContextUsage();
+		if (!usage || usage.tokens === null || !Number.isFinite(usage.tokens)) return undefined;
+		const tokens = Math.max(0, Math.floor(usage.tokens));
+		const reminder = Math.floor(tokens / interval);
+		if (reminder < 1) return undefined;
+		const instruction = instructions[Math.min(reminder, instructions.length) - 1] ?? instructions[0];
+		return {
+			message: {
+				customType: NUDGE_TYPE,
+				content: automaticInstruction(instruction),
+				display: false,
+				details: {
+					v: 1,
+					kind: "automatic",
+					tokens,
+					boundaryTokens: reminder * interval,
+					reminder,
+					tier: Math.min(reminder, instructions.length),
+					tierCount: instructions.length,
+					anchorToolCallId:
+						replayWorkingMemoryState(ctx.sessionManager.getBranch(), true).latestAnchorToolCallId ?? null,
+				},
+			},
+		};
+	});
+
+	pi.on("turn_end", (event, ctx) => {
+		if (!enabled || event.toolResults.length === 0) return undefined;
+		const usage = ctx.getContextUsage();
+		if (!usage || usage.tokens === null || !Number.isFinite(usage.tokens)) return undefined;
+		const tokens = Math.max(0, Math.floor(usage.tokens));
+		const branch = ctx.sessionManager.getBranch();
+		const anchor = replayWorkingMemoryState(branch, true).latestAnchorToolCallId;
+		if (anchor !== nudgeState.anchorToolCallId) nudgeState = reconstructNudgeState(branch, anchor);
+		if (anchor !== undefined && nudgeState.suppressedThroughTokens === undefined) {
+			const floor = Math.floor(tokens / interval) * interval;
+			pi.appendEntry(BASELINE_TYPE, { v: 1, anchorToolCallId: anchor, suppressedThroughTokens: floor });
+			nudgeState.suppressedThroughTokens = floor;
+			nudgeState.highestBoundaryTokens = floor;
+			return undefined;
+		}
+		const reminder = Math.floor(tokens / interval);
+		if (reminder < 1) return undefined;
+		const boundaryTokens = reminder * interval;
+		if (boundaryTokens <= nudgeState.highestBoundaryTokens) return undefined;
+		const tier = Math.min(reminder, instructions.length);
+		const instruction = instructions[tier - 1] ?? instructions[0];
+		pi.sendMessage<WorkingMemoryNudgeDetails>(
+			{
+				customType: NUDGE_TYPE,
+				content: automaticInstruction(instruction),
+				display: true,
+				details: {
+					v: 1,
+					kind: "automatic",
+					tokens,
+					boundaryTokens,
+					reminder,
+					tier,
+					tierCount: instructions.length,
+					anchorToolCallId: anchor ?? null,
+				},
+			},
+			{ deliverAs: "steer" },
+		);
+		nudgeState.highestBoundaryTokens = boundaryTokens;
+		return undefined;
+	});
+
+	pi.on("context", (event, ctx) => {
+		if (!enabled) return undefined;
+		const branch = ctx.sessionManager.getBranch();
+		const state = replayWorkingMemoryState(branch, true);
+		const messages = projectWorkingMemory(event.messages, state, branch, ctx.sessionManager.buildContextEntries());
+		if (!setsEqual(visualRows, state.prunedRowIds)) {
+			visualRows = new Set(state.prunedRowIds);
+			pushVisualSnapshot();
+		}
+		return { messages };
+	});
+
+	pi.on("session_tree", (_event, ctx) => {
+		invalidate();
+		syncBranch(ctx);
+	});
+	pi.on("session_compact", (_event, ctx) => {
+		invalidate();
+		syncBranch(ctx);
+	});
+	pi.on("session_shutdown", () => {
+		invalidate();
+		enabled = false;
+		setToolActive(false);
+		visualRows.clear();
+		pushVisualSnapshot();
+	});
+}
+
+function reconstructNudgeState(branch: readonly SessionEntry[], anchor: string | undefined): NudgeState {
+	let suppressedThroughTokens = anchor === undefined ? 0 : undefined;
+	let highestBoundaryTokens = 0;
+	let anchorResultIndex = -1;
+	if (anchor !== undefined) {
+		anchorResultIndex = branch.findIndex(
+			(entry) =>
+				entry.type === "message" &&
+				entry.message.role === "toolResult" &&
+				entry.message.toolName === WORKING_MEMORY_TOOL &&
+				entry.message.toolCallId === anchor,
+		);
+	}
+	for (let index = anchorResultIndex + 1; index < branch.length; index += 1) {
+		const entry = branch[index];
+		if (!entry) continue;
+		if (entry.type === "custom" && entry.customType === BASELINE_TYPE) {
+			const baseline = parseBaseline(entry.data);
+			if (baseline !== undefined && baseline.anchorToolCallId === anchor) {
+				suppressedThroughTokens = baseline.suppressedThroughTokens;
+				highestBoundaryTokens = Math.max(highestBoundaryTokens, baseline.suppressedThroughTokens);
+			}
+			continue;
+		}
+		if (entry.type !== "custom_message" || entry.customType !== NUDGE_TYPE) continue;
+		const details = parseWorkingMemoryNudge(entry.details);
+		if (details?.kind !== "automatic" || details.anchorToolCallId !== (anchor ?? null)) continue;
+		highestBoundaryTokens = Math.max(highestBoundaryTokens, details.boundaryTokens);
+	}
+	return { anchorToolCallId: anchor, suppressedThroughTokens, highestBoundaryTokens };
+}
+
+function parseBaseline(value: unknown): { anchorToolCallId: string; suppressedThroughTokens: number } | undefined {
+	if (!value || typeof value !== "object" || Array.isArray(value)) return undefined;
+	const record = value as Record<string, unknown>;
+	if (
+		record.v !== 1 ||
+		typeof record.anchorToolCallId !== "string" ||
+		!Number.isSafeInteger(record.suppressedThroughTokens) ||
+		(record.suppressedThroughTokens as number) < 0
+	) {
+		return undefined;
+	}
+	return {
+		anchorToolCallId: record.anchorToolCallId,
+		suppressedThroughTokens: record.suppressedThroughTokens as number,
+	};
+}
+
+function automaticInstruction(instruction: string): string {
+	return `Internal working-memory instruction. Follow silently without mentioning token counts or context management. ${instruction} Use working_memory only when a selective checkpoint improves future work; never prune useful evidence merely to reread it.`;
+}
+
+function manualInstruction(): string {
+	return "Internal working-memory instruction. Follow silently. Reassess active evidence, use working_memory if a selective checkpoint improves future work, then continue unfinished work. Do not mention this request.";
+}
+
+function setsEqual(left: ReadonlySet<string>, right: ReadonlySet<string>): boolean {
+	if (left.size !== right.size) return false;
+	for (const value of left) if (!right.has(value)) return false;
+	return true;
+}

package/extensions/working-memory/memory.ts

@@ -0,0 +1,266 @@
+import { sessionEntryToContextMessages, type ContextEvent, type SessionEntry } from "@earendil-works/pi-coding-agent";
+import { isContextProjectionMessage, isLegacyContextMessage } from "../../shared/context-messages.ts";
+import type { ActiveWorkingMemoryState } from "./state.ts";
+
+type ContextMessage = ContextEvent["messages"][number];
+
+export interface MemoryUnit {
+	ref: string;
+	label: string;
+	preview: string;
+	order: number;
+	suborder: number;
+	messages: ContextMessage[];
+	rowIds: string[];
+}
+
+const EXCLUDED_CUSTOM_TYPES = new Set(["tau.autoread", "tau.runtime-context", "tau.working-memory.nudge"]);
+const REFERENCE_CATALOG_TYPE = "tau.working-memory.references";
+
+export function buildMemoryCatalog(branch: readonly SessionEntry[]): Map<string, MemoryUnit> {
+	const catalog = new Map<string, MemoryUnit>();
+	const results = new Map<string, ContextMessage>();
+	for (const entry of branch) {
+		for (const message of sessionEntryToContextMessages(entry)) {
+			if (message.role === "toolResult") results.set(message.toolCallId, message);
+		}
+	}
+
+	for (let order = 0; order < branch.length; order += 1) {
+		const entry = branch[order];
+		if (!entry) continue;
+		const message = sessionEntryToContextMessages(entry)[0];
+		if (!message) continue;
+		if (message.role === "assistant") {
+			const calls = message.content.filter((block) => block.type === "toolCall");
+			const frameworkCall = calls.some((call) => call.name === "working_memory");
+			if (frameworkCall) continue;
+			const prose = message.content.filter((block) => block.type !== "toolCall");
+			if (prose.length > 0) {
+				const ref = `m:${entry.id}`;
+				catalog.set(ref, {
+					ref,
+					label: "assistant",
+					preview: previewAssistant(prose),
+					order,
+					suborder: 0,
+					messages: [{ ...message, content: prose }],
+					rowIds: [],
+				});
+			}
+			for (let index = 0; index < calls.length; index += 1) {
+				const call = calls[index];
+				if (!call) continue;
+				const result = results.get(call.id);
+				if (result?.role !== "toolResult" || result.toolName !== call.name) continue;
+				const ref = `t:${entry.id}:${index + 1}`;
+				catalog.set(ref, {
+					ref,
+					label: `tool ${call.name}`,
+					preview: previewTool(call.arguments),
+					order,
+					suborder: index + 1,
+					messages: [{ ...message, content: [call] }, result],
+					rowIds: [call.id],
+				});
+			}
+			continue;
+		}
+		if (message.role === "toolResult") continue;
+		if (message.role === "custom" && EXCLUDED_CUSTOM_TYPES.has(message.customType)) continue;
+		if (
+			message.role === "custom" &&
+			message.customType === "tau.injected-context" &&
+			message.details &&
+			typeof message.details === "object" &&
+			(message.details as Record<string, unknown>).source === "context"
+		)
+			continue;
+		const ref = `m:${entry.id}`;
+		catalog.set(ref, {
+			ref,
+			label: message.role === "custom" ? message.customType : message.role,
+			preview: previewMessage(message),
+			order,
+			suborder: 0,
+			messages: [message],
+			rowIds: outlineRowIds(message),
+		});
+	}
+	return catalog;
+}
+
+export function projectWorkingMemory(
+	messages: readonly ContextMessage[],
+	state: ActiveWorkingMemoryState,
+	branch: readonly SessionEntry[],
+	contextEntries: readonly SessionEntry[],
+): ContextMessage[] {
+	const catalog = buildMemoryCatalog(branch);
+	const refsByMessage = referenceCurrentMessages(messages, contextEntries, catalog);
+	if (state.latestAnchorToolCallId === undefined) {
+		return addReferenceCatalog(messages, refsByMessage);
+	}
+	const anchorIndex = findAnchor(messages, state.latestAnchorToolCallId);
+	if (anchorIndex < 0) return addReferenceCatalog(messages, refsByMessage);
+
+	const retained = state.retainedRefs
+		.flatMap((ref) => {
+			const unit = catalog.get(ref);
+			return unit ? [unit] : [];
+		})
+		.sort((left, right) => left.order - right.order || left.suborder - right.suborder);
+	const projected: ContextMessage[] = [];
+	const projectedRefs: MemoryUnit[][] = [];
+	for (const unit of retained) {
+		projected.push(...unit.messages);
+		projectedRefs.push(...unit.messages.map((_, index) => (index === unit.messages.length - 1 ? [unit] : [])));
+	}
+	for (let index = anchorIndex; index < messages.length; index += 1) {
+		const message = messages[index];
+		if (!message) continue;
+		projected.push(index === anchorIndex ? sanitizeAnchor(message, state.latestAnchorToolCallId) : message);
+		projectedRefs.push(refsByMessage[index] ?? []);
+	}
+	return addReferenceCatalog(projected, projectedRefs);
+}
+
+function referenceCurrentMessages(
+	messages: readonly ContextMessage[],
+	entries: readonly SessionEntry[],
+	catalog: ReadonlyMap<string, MemoryUnit>,
+): MemoryUnit[][] {
+	const generated = entries
+		.flatMap((entry) => sessionEntryToContextMessages(entry).map((message) => ({ entry, message })))
+		.filter(({ message }) => !isLegacyContextMessage(message));
+	const sessionMessages = messages.filter(
+		(message) => !isContextProjectionMessage(message) && !isLegacyContextMessage(message),
+	);
+	if (generated.length !== sessionMessages.length) return messages.map(() => []);
+	const toolRefs = new Map<string, MemoryUnit[]>();
+	for (const unit of catalog.values()) {
+		const result = unit.messages.find((message) => message.role === "toolResult");
+		if (result?.role !== "toolResult") continue;
+		toolRefs.set(result.toolCallId, [...(toolRefs.get(result.toolCallId) ?? []), unit]);
+	}
+	let generatedIndex = 0;
+	return messages.map((currentMessage) => {
+		if (isContextProjectionMessage(currentMessage) || isLegacyContextMessage(currentMessage)) return [];
+		const current = generated[generatedIndex];
+		generatedIndex += 1;
+		if (!current) return [];
+		const { entry, message } = current;
+		if (message.role === "assistant") {
+			const unit = catalog.get(`m:${entry.id}`);
+			return unit && currentMessage.role === "assistant" ? [unit] : [];
+		}
+		if (message.role === "toolResult") return toolRefs.get(message.toolCallId) ?? [];
+		const unit = catalog.get(`m:${entry.id}`);
+		return unit ? [unit] : [];
+	});
+}
+
+function addReferenceCatalog(
+	messages: readonly ContextMessage[],
+	unitsByMessage: readonly MemoryUnit[][],
+): ContextMessage[] {
+	const projected: ContextMessage[] = [
+		{
+			role: "custom",
+			customType: REFERENCE_CATALOG_TYPE,
+			content:
+				"Working-memory references follow their source messages. Use these exact refs in working_memory.keep; do not repeat them in responses.",
+			display: false,
+			timestamp: 0,
+		},
+	];
+	for (let index = 0; index < messages.length; index += 1) {
+		const message = messages[index];
+		if (!message) continue;
+		projected.push(message);
+		const unique = new Map((unitsByMessage[index] ?? []).map((unit) => [unit.ref, unit]));
+		if (unique.size === 0) continue;
+		projected.push({
+			role: "custom",
+			customType: REFERENCE_CATALOG_TYPE,
+			content: [...unique.values()]
+				.sort((left, right) => left.order - right.order || left.suborder - right.suborder)
+				.map((unit) => `${unit.ref} (${unit.label}): ${unit.preview}`)
+				.join("\n"),
+			display: false,
+			timestamp: 0,
+		});
+	}
+	return projected;
+}
+
+function sanitizeAnchor(message: ContextMessage, id: string): ContextMessage {
+	if (message.role !== "assistant") return message;
+	return {
+		...message,
+		content: message.content.map((block) =>
+			block.type === "toolCall" && block.id === id && block.name === "working_memory"
+				? { ...block, arguments: { checkpoint: true } }
+				: block,
+		),
+	};
+}
+
+function findAnchor(messages: readonly ContextMessage[], id: string): number {
+	for (let index = messages.length - 1; index >= 0; index -= 1) {
+		const message = messages[index];
+		if (
+			message?.role === "assistant" &&
+			message.content.some(
+				(block) => block.type === "toolCall" && block.id === id && block.name === "working_memory",
+			)
+		) {
+			return index;
+		}
+	}
+	return -1;
+}
+
+function previewAssistant(blocks: ReadonlyArray<{ type: string }>): string {
+	for (const block of blocks) {
+		if (block.type === "text" && "text" in block && typeof block.text === "string") return compact(block.text);
+		if (block.type === "thinking" && "thinking" in block && typeof block.thinking === "string") {
+			return compact(block.thinking);
+		}
+	}
+	return "";
+}
+
+function previewTool(args: Record<string, unknown>): string {
+	return compact(JSON.stringify(args));
+}
+
+function previewMessage(message: ContextMessage): string {
+	if (message.role === "user" || message.role === "custom") {
+		if (typeof message.content === "string") return compact(message.content);
+		const text = message.content.find((part) => part.type === "text");
+		return text?.type === "text" ? compact(text.text) : "";
+	}
+	if (message.role === "bashExecution") return compact(`${message.command}: ${message.output}`);
+	if (message.role === "assistant") return previewAssistant(message.content);
+	if (message.role === "toolResult") {
+		const text = message.content.find((part) => part.type === "text");
+		return text?.type === "text" ? compact(text.text) : "";
+	}
+	return compact(message.summary);
+}
+
+function outlineRowIds(message: ContextMessage): string[] {
+	if (message.role !== "custom" || message.customType !== "tau.explore.outline") return [];
+	if (!isRecord(message.details) || typeof message.details.rowId !== "string") return [];
+	return [message.details.rowId];
+}
+
+function compact(value: string): string {
+	const text = value.replaceAll(/\s+/g, " ").trim();
+	return text.length <= 120 ? text : `${text.slice(0, 119)}…`;
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+	return typeof value === "object" && value !== null && !Array.isArray(value);
+}