@shanepadgett/tau-agent 0.33.1 → 0.35.0

package/extensions/working-memory/index.ts

@@ -1,341 +0,0 @@
-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 WorkingMemoryCheckpointDetailsV2 } 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. Retain one or more valuable user or visible assistant messages, auto-read full source or carry file structure as needed, defer conditionally relevant files, and distill exploration findings into one compact continuation note.";
-
-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, WorkingMemoryCheckpointDetailsV2>({
-					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 only for user messages and visible assistant text. Tool calls, tool results, hidden reasoning, and framework messages cannot be retained.",
-						"Every checkpoint must retain at least one relevant referenced message. Keep task framing, constraints, decisions, and immediate work chain when needed; continuation and file tiers supplement retained conversation and cannot replace it.",
-						"Choose one file tier: readFiles auto-reads source into the next turn when its body is needed; outlineFiles carries symbols and locations for later scoped inspection; deferFiles records inactive conditional paths. Do not read a file merely to decide whether to outline it.",
-						"Choose readFiles instead of outlineFiles when next work will require the complete file. Do not list a path in more than one file tier.",
-						"Use continuation as a working note for resuming mid-task. Carry durable decisions, concrete findings, live reasoning, unresolved questions, remaining work, and next action in as much detail as needed to continue without rereading discarded results. Do not make it a user-facing status update or narrate the checkpoint.",
-					],
-					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,
-							executionStarted: context.executionStarted,
-						});
-					},
-					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

@@ -1,235 +0,0 @@
-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;
-	messages: ContextMessage[];
-}
-
-const REFERENCE_CATALOG_TYPE = "tau.working-memory.references";
-
-export function buildMemoryCatalog(branch: readonly SessionEntry[]): Map<string, MemoryUnit> {
-	const catalog = new Map<string, MemoryUnit>();
-	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 text = message.content.filter((block) => block.type === "text");
-			if (text.length > 0) {
-				const ref = `m:${entry.id}`;
-				catalog.set(ref, {
-					ref,
-					label: "assistant",
-					preview: previewAssistant(text),
-					order,
-					messages: [{ ...message, content: text }],
-				});
-			}
-			continue;
-		}
-		if (message.role !== "user") continue;
-		const ref = `m:${entry.id}`;
-		catalog.set(ref, {
-			ref,
-			label: "user",
-			preview: previewUser(message),
-			order,
-			messages: [message],
-		});
-	}
-	return catalog;
-}
-
-export function collectPrunedRowIds(branch: readonly SessionEntry[], before: number): string[] {
-	const rowIds = new Set<string>();
-	for (let index = 0; index < before; index += 1) {
-		const entry = branch[index];
-		if (!entry) continue;
-		for (const message of sessionEntryToContextMessages(entry)) {
-			if (message.role === "assistant") {
-				for (const block of message.content) {
-					if (block.type === "toolCall" && block.name !== "working_memory") rowIds.add(block.id);
-				}
-			}
-			if (
-				message.role === "custom" &&
-				(message.customType === "tau.explore.outline" || message.customType === "tau.autoread") &&
-				isRecord(message.details) &&
-				typeof message.details.rowId === "string"
-			) {
-				rowIds.add(message.details.rowId);
-			}
-		}
-	}
-	return [...rowIds];
-}
-
-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);
-	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(() => []);
-	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 !== "user" && message.role !== "assistant") return [];
-		const unit = catalog.get(`m:${entry.id}`);
-		return unit && currentMessage.role === message.role ? [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,
-		},
-	];
-	let deferredAssistantUnits: MemoryUnit[] = [];
-	let pendingToolCallIds = new Set<string>();
-	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 (message.role === "assistant") {
-			pendingToolCallIds = new Set(
-				message.content.filter((block) => block.type === "toolCall").map((block) => block.id),
-			);
-			if (pendingToolCallIds.size > 0) {
-				deferredAssistantUnits = [...unique.values()];
-				continue;
-			}
-		}
-		if (message.role === "toolResult" && pendingToolCallIds.delete(message.toolCallId)) {
-			if (pendingToolCallIds.size > 0) continue;
-			for (const unit of deferredAssistantUnits) unique.set(unit.ref, unit);
-			deferredAssistantUnits = [];
-		}
-		if (unique.size === 0) continue;
-		projected.push({
-			role: "custom",
-			customType: REFERENCE_CATALOG_TYPE,
-			content: [...unique.values()]
-				.sort((left, right) => left.order - right.order)
-				.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);
-	}
-	return "";
-}
-
-function previewUser(message: Extract<ContextMessage, { role: "user" }>): string {
-	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) : "";
-}
-
-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);
-}