@shanepadgett/tau-agent 0.28.0 → 0.29.0

package/extensions/working-memory/memory.ts

@@ -0,0 +1,248 @@
+import { sessionEntryToContextMessages, type ContextEvent, type SessionEntry } from "@earendil-works/pi-coding-agent";
+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;
+		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 })),
+	);
+	if (generated.length !== messages.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]);
+	}
+	return generated.map(({ entry, message }, index) => {
+		if (message.role === "assistant") {
+			const unit = catalog.get(`m:${entry.id}`);
+			return unit && messages[index]?.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);
+}

package/extensions/working-memory/render.ts

@@ -0,0 +1,178 @@
+import type { AgentToolResult, Theme } from "@earendil-works/pi-coding-agent";
+import { Text, truncateToWidth, visibleWidth } from "@earendil-works/pi-tui";
+import { Marker } from "@shanepadgett/tau-tui";
+import { formatToolRowTitle, type ToolRowStateStore } from "../../shared/tool-row-state.ts";
+import type { WorkingMemoryInput } from "./checkpoint.ts";
+import { parseWorkingMemoryDetails } from "./state.ts";
+
+const MAX_ITEMS = 20;
+const MAX_LINE = 240;
+const MAX_TEXT = 4_000;
+
+export type WorkingMemoryNudgeDetails =
+	| {
+			v: 1;
+			kind: "automatic";
+			tokens: number;
+			boundaryTokens: number;
+			reminder: number;
+			tier: number;
+			tierCount: number;
+			anchorToolCallId: string | null;
+	  }
+	| {
+			v: 1;
+			kind: "manual";
+			tokens: null;
+			boundaryTokens: null;
+			reminder: null;
+			tier: null;
+			tierCount: null;
+			anchorToolCallId: string | null;
+	  };
+
+export function parseWorkingMemoryNudge(value: unknown): WorkingMemoryNudgeDetails | undefined {
+	if (!isRecord(value) || value.v !== 1 || (value.kind !== "automatic" && value.kind !== "manual")) return undefined;
+	if (value.anchorToolCallId !== null && !nonempty(value.anchorToolCallId)) return undefined;
+	if (value.kind === "manual") {
+		if (
+			value.tokens !== null ||
+			value.boundaryTokens !== null ||
+			value.reminder !== null ||
+			value.tier !== null ||
+			value.tierCount !== null
+		) {
+			return undefined;
+		}
+		return {
+			v: 1,
+			kind: "manual",
+			tokens: null,
+			boundaryTokens: null,
+			reminder: null,
+			tier: null,
+			tierCount: null,
+			anchorToolCallId: value.anchorToolCallId,
+		};
+	}
+	if (
+		!count(value.tokens, 0) ||
+		!count(value.boundaryTokens, 1) ||
+		!count(value.reminder, 1) ||
+		!count(value.tier, 1, 5) ||
+		!count(value.tierCount, 1, 5) ||
+		value.boundaryTokens > value.tokens ||
+		value.tier > value.tierCount ||
+		value.tier !== Math.min(value.reminder, value.tierCount)
+	) {
+		return undefined;
+	}
+	return {
+		v: 1,
+		kind: "automatic",
+		tokens: value.tokens,
+		boundaryTokens: value.boundaryTokens,
+		reminder: value.reminder,
+		tier: value.tier,
+		tierCount: value.tierCount,
+		anchorToolCallId: value.anchorToolCallId,
+	};
+}
+
+export function renderWorkingMemoryNudge(details: unknown, theme: Theme): Marker | undefined {
+	const parsed = parseWorkingMemoryNudge(details);
+	if (!parsed) return undefined;
+	return new Marker({
+		theme,
+		state: "muted",
+		label: "Memory:",
+		parts:
+			parsed.kind === "manual" ? ["Reassessment requested."] : [formatTokens(parsed.boundaryTokens), "Reassess."],
+	});
+}
+
+export function renderWorkingMemoryCall(
+	args: WorkingMemoryInput,
+	theme: Theme,
+	context: {
+		rowState: ToolRowStateStore;
+		rowId: string;
+		invalidate: () => void;
+		lastComponent: unknown;
+	},
+): Text {
+	context.rowState.watch(context.rowId, context.invalidate);
+	const component = context.lastComponent instanceof Text ? context.lastComponent : new Text("", 0, 0);
+	const kept = Array.isArray(args.keep) ? args.keep.length : 0;
+	const outlined = Array.isArray(args.outlineFiles) ? args.outlineFiles.length : 0;
+	const deferred = Array.isArray(args.deferFiles) ? args.deferFiles.length : 0;
+	component.setText(
+		`${formatToolRowTitle(context.rowState, context.rowId, "working_memory", theme)} ${theme.fg("muted", `${kept} kept · ${outlined} outline · ${deferred} deferred`)}`,
+	);
+	return component;
+}
+
+export function renderWorkingMemoryResult(
+	result: AgentToolResult<unknown>,
+	expanded: boolean,
+	theme: Theme,
+	lastComponent: unknown,
+): Text {
+	const component = lastComponent instanceof Text ? lastComponent : new Text("", 0, 0);
+	const details = parseWorkingMemoryDetails(result.details);
+	if (!details) {
+		component.setText(
+			theme.fg("warning", bounded(firstText(result) || "working_memory returned invalid details", 1_000)),
+		);
+		return component;
+	}
+	let text = theme.fg(
+		details.warnings.length === 0 ? "success" : "warning",
+		`Checkpoint · kept ${details.retainedRefs.length} · outlined ${details.outlinedFiles.length} · deferred ${details.deferredFiles.length} · removed ${details.removedUnits}${details.warnings.length === 0 ? "" : ` · warnings ${details.warnings.length}`}`,
+	);
+	if (expanded) {
+		const lines = [
+			bounded(firstText(result), 1_500),
+			...details.retainedLabels.map((item) => `${item.ref} ${item.label}: ${item.preview}`),
+			...details.outlinedFiles.map((file) => `outlined: ${file.path}`),
+			...details.deferredFiles.map((file) => `deferred: ${file.path} — ${file.reason}; when ${file.relevantWhen}`),
+			...details.warnings.map((warning) => `warning: ${warning}`),
+		].filter(Boolean);
+		let remaining = MAX_TEXT - visibleWidth(text);
+		for (const line of lines.slice(0, MAX_ITEMS)) {
+			if (remaining <= 1) break;
+			const styled = theme.fg("dim", bounded(line, MAX_LINE));
+			text += `\n${truncateToWidth(styled, remaining - 1, "…")}`;
+			remaining -= 1 + Math.min(visibleWidth(styled), remaining - 1);
+		}
+		if (lines.length > MAX_ITEMS && remaining > 10)
+			text += `\n${theme.fg("dim", `… ${lines.length - MAX_ITEMS} more`)}`;
+	}
+	component.setText(text);
+	return component;
+}
+
+function firstText(result: AgentToolResult<unknown>): string {
+	for (const part of result.content) if (part.type === "text") return part.text;
+	return "";
+}
+
+function bounded(text: string, maximum: number): string {
+	return text.length <= maximum ? text : `${text.slice(0, maximum - 1)}…`;
+}
+
+function formatTokens(tokens: number): string {
+	return tokens < 1_000 ? `${tokens}` : `${Math.round(tokens / 1_000)}k`;
+}
+
+function count(value: unknown, minimum: number, maximum = Number.MAX_SAFE_INTEGER): value is number {
+	return Number.isSafeInteger(value) && (value as number) >= minimum && (value as number) <= maximum;
+}
+
+function nonempty(value: unknown): value is string {
+	return typeof value === "string" && value.length > 0;
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+	return typeof value === "object" && value !== null && !Array.isArray(value);
+}

package/extensions/working-memory/settings.ts

@@ -0,0 +1,38 @@
+import { Type } from "typebox";
+import { defineTauExtensionSettings } from "../../shared/settings/define.ts";
+
+const DEFAULT_NUDGE_INSTRUCTIONS: [string, ...string[]] = [
+	"Reassess working memory. Continue coherent exploration when its evidence remains useful; otherwise prune known dead ends, obsolete outputs, and other context with no expected value.",
+	"Context is materially larger. Prune stale or bulky irrelevant evidence when safe, but keep active working evidence that would otherwise need to be reread.",
+	"Strongly reassess before more broad work. Remove accumulated waste and carry useful information at cheapest sufficient fidelity without scrubbing the active working set.",
+];
+
+export default defineTauExtensionSettings({
+	key: "workingMemory",
+	defaults: {
+		enabled: true as boolean,
+		nudgeEveryTokens: 40_000 as number,
+		nudgeInstructions: DEFAULT_NUDGE_INSTRUCTIONS,
+	},
+	schema: Type.Object(
+		{
+			enabled: Type.Optional(Type.Boolean({ default: true, description: "Enable working-memory checkpoints." })),
+			nudgeEveryTokens: Type.Optional(
+				Type.Integer({
+					default: 40_000,
+					minimum: 1,
+					description: "Active-context token interval between advisory working-memory reminders.",
+				}),
+			),
+			nudgeInstructions: Type.Optional(
+				Type.Array(Type.String({ minLength: 1, maxLength: 2_000 }), {
+					default: DEFAULT_NUDGE_INSTRUCTIONS,
+					minItems: 1,
+					maxItems: 5,
+					description: "Ordered advisory working-memory instructions. Later reminders repeat final instruction.",
+				}),
+			),
+		},
+		{ additionalProperties: false },
+	),
+});

package/extensions/working-memory/state.ts

@@ -0,0 +1,152 @@
+import type { SessionEntry } from "@earendil-works/pi-coding-agent";
+
+export const WORKING_MEMORY_TOOL = "working_memory";
+
+export interface DeferredFile {
+	path: string;
+	reason: string;
+	relevantWhen: string;
+}
+
+export interface RetainedLabel {
+	ref: string;
+	label: string;
+	preview: string;
+}
+
+export interface WorkingMemoryCheckpointDetailsV1 {
+	v: 1;
+	anchorToolCallId: string;
+	retainedRefs: string[];
+	retainedLabels: RetainedLabel[];
+	prunedRowIds: string[];
+	outlinedFiles: Array<{ path: string; rowId: string }>;
+	deferredFiles: DeferredFile[];
+	removedUnits: number;
+	warnings: string[];
+}
+
+export interface ActiveWorkingMemoryState {
+	latestAnchorToolCallId: string | undefined;
+	retainedRefs: readonly string[];
+	prunedRowIds: ReadonlySet<string>;
+}
+
+export function parseWorkingMemoryDetails(value: unknown): WorkingMemoryCheckpointDetailsV1 | undefined {
+	if (!isRecord(value) || value.v !== 1 || !nonempty(value.anchorToolCallId)) return undefined;
+	const keys = [
+		"v",
+		"anchorToolCallId",
+		"retainedRefs",
+		"retainedLabels",
+		"prunedRowIds",
+		"outlinedFiles",
+		"deferredFiles",
+		"removedUnits",
+		"warnings",
+	];
+	if (Object.keys(value).length !== keys.length || !keys.every((key) => Object.hasOwn(value, key))) return undefined;
+	const retainedRefs = uniqueStrings(value.retainedRefs);
+	const prunedRowIds = uniqueStrings(value.prunedRowIds);
+	const warnings = strings(value.warnings);
+	const retainedLabels = parseRetainedLabels(value.retainedLabels);
+	const outlinedFiles = parseOutlinedFiles(value.outlinedFiles);
+	const deferredFiles = parseDeferredFiles(value.deferredFiles);
+	if (
+		!retainedRefs ||
+		!prunedRowIds ||
+		!warnings ||
+		!retainedLabels ||
+		retainedLabels.length !== retainedRefs.length ||
+		retainedLabels.some((label, index) => label.ref !== retainedRefs[index]) ||
+		!outlinedFiles ||
+		!deferredFiles ||
+		!Number.isSafeInteger(value.removedUnits) ||
+		(value.removedUnits as number) < 0
+	) {
+		return undefined;
+	}
+	return {
+		v: 1,
+		anchorToolCallId: value.anchorToolCallId,
+		retainedRefs,
+		retainedLabels,
+		prunedRowIds,
+		outlinedFiles,
+		deferredFiles,
+		removedUnits: value.removedUnits as number,
+		warnings,
+	};
+}
+
+export function replayWorkingMemoryState(branch: readonly SessionEntry[], enabled: boolean): ActiveWorkingMemoryState {
+	let prunedRowIds = new Set<string>();
+	let latestAnchorToolCallId: string | undefined;
+	let retainedRefs: readonly string[] = [];
+	if (enabled) {
+		for (const entry of branch) {
+			if (entry.type !== "message" || entry.message.role !== "toolResult") continue;
+			if (entry.message.toolName !== WORKING_MEMORY_TOOL) continue;
+			const details = parseWorkingMemoryDetails(entry.message.details);
+			if (!details || details.anchorToolCallId !== entry.message.toolCallId) continue;
+			prunedRowIds = new Set(details.prunedRowIds);
+			latestAnchorToolCallId = details.anchorToolCallId;
+			retainedRefs = details.retainedRefs;
+		}
+	}
+	return { latestAnchorToolCallId, retainedRefs, prunedRowIds };
+}
+
+function parseRetainedLabels(value: unknown): RetainedLabel[] | undefined {
+	if (!Array.isArray(value)) return undefined;
+	const labels: RetainedLabel[] = [];
+	for (const item of value) {
+		if (!isRecord(item) || !nonempty(item.ref) || !nonempty(item.label) || typeof item.preview !== "string") {
+			return undefined;
+		}
+		labels.push({ ref: item.ref, label: item.label, preview: item.preview });
+	}
+	return labels;
+}
+
+function parseOutlinedFiles(value: unknown): Array<{ path: string; rowId: string }> | undefined {
+	if (!Array.isArray(value)) return undefined;
+	const files: Array<{ path: string; rowId: string }> = [];
+	for (const item of value) {
+		if (!isRecord(item) || !nonempty(item.path) || !nonempty(item.rowId)) return undefined;
+		files.push({ path: item.path, rowId: item.rowId });
+	}
+	return files;
+}
+
+function parseDeferredFiles(value: unknown): DeferredFile[] | undefined {
+	if (!Array.isArray(value)) return undefined;
+	const files: DeferredFile[] = [];
+	for (const item of value) {
+		if (!isRecord(item) || !nonempty(item.path) || !nonempty(item.reason) || !nonempty(item.relevantWhen)) {
+			return undefined;
+		}
+		files.push({ path: item.path, reason: item.reason, relevantWhen: item.relevantWhen });
+	}
+	return files;
+}
+
+function uniqueStrings(value: unknown): string[] | undefined {
+	const parsed = strings(value, true);
+	if (!parsed || new Set(parsed).size !== parsed.length) return undefined;
+	return parsed;
+}
+
+function strings(value: unknown, requireNonempty = false): string[] | undefined {
+	if (!Array.isArray(value)) return undefined;
+	if (!value.every((item) => typeof item === "string" && (!requireNonempty || item.length > 0))) return undefined;
+	return [...value];
+}
+
+function nonempty(value: unknown): value is string {
+	return typeof value === "string" && value.length > 0;
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+	return typeof value === "object" && value !== null && !Array.isArray(value);
+}

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@shanepadgett/tau-agent",
-	"version": "0.28.0",
+	"version": "0.29.0",
 	"description": "Tau is a custom agentic harness built with pi extensions",
 	"type": "module",
 	"main": "./src/index.ts",
@@ -35,7 +35,7 @@
 	],
 	"dependencies": {
 		"@ast-grep/wasm": "0.45.0",
-		"@shanepadgett/tau-tui": "0.28.0",
+		"@shanepadgett/tau-tui": "0.29.0",
 		"@vscode/tree-sitter-wasm": "0.3.1",
 		"image-size": "2.0.2",
 		"smol-toml": "1.7.0",