@shanepadgett/tau-agent 0.30.0 → 0.31.1

package/extensions/tau-help/help.md

@@ -36,7 +36,7 @@ Adds `/context` to set branch-local reusable repository work scopes from `.pi/co
 
 ## working-memory
 
-Gives agent `working_memory` for selective hard checkpoints. Model-only references identify useful conversation evidence and complete tool exchanges. Requested source files return as structural outlines, while deferred files remain cheap conditional reminders. Everything else before checkpoint leaves future model input without changing saved session. Advisory reminders begin at 40k active-context tokens. Run `/prune` to request reassessment manually.
+Gives agent `working_memory` for selective hard checkpoints. Model-only references identify useful user messages and visible assistant text. Requested source files return as structural outlines, deferred files remain cheap conditional reminders, and a continuation note carries conclusions extracted from exploration. Tool history and full file reads leave future model input without changing saved session. Advisory reminders begin at 40k active-context tokens. Run `/prune` to request reassessment manually.
 
 ## explore
 

package/extensions/working-memory/README.md

@@ -2,7 +2,7 @@
 
 Working Memory gives agent selective checkpoints without changing saved conversation.
 
-`working_memory` keeps referenced conversation evidence and complete tool exchanges, carries requested source files as structural outlines, and records deferred files as conditional reminders. Everything else before checkpoint leaves future model context.
+`working_memory` keeps referenced user messages and visible assistant text, carries requested source files as structural outlines, and records deferred files as conditional reminders. Tool history and full file reads leave model context; useful findings from them are distilled into a compact continuation note.
 
 Automatic reminders begin at 40,000 active-context tokens and remain advisory. Run `/prune` to request reassessment manually.
 

package/extensions/working-memory/checkpoint.ts

@@ -3,7 +3,7 @@ import { type ExtensionAPI, type ExtensionContext, type SessionEntry } from "@ea
 import { type Static, Type } from "typebox";
 import { truncateBoundedHead } from "../../shared/bounded-text-result.ts";
 import { requestOutlineInjections, type PreparedOutlineInjection } from "../../shared/outline-injection.ts";
-import { buildMemoryCatalog } from "./memory.ts";
+import { buildMemoryCatalog, collectPrunedRowIds } from "./memory.ts";
 import { WORKING_MEMORY_TOOL, type DeferredFile, type WorkingMemoryCheckpointDetailsV1 } from "./state.ts";
 
 const PATH = Type.String({ minLength: 1, maxLength: 500, pattern: "\\S" });
@@ -58,7 +58,7 @@ export async function executeWorkingMemory(options: ExecuteWorkingMemoryOptions)
 			const unit = catalog.get(ref);
 			return unit ? [unit] : [];
 		})
-		.sort((left, right) => left.order - right.order || left.suborder - right.suborder);
+		.sort((left, right) => left.order - right.order);
 	const retainedRefs = retained.map((unit) => unit.ref);
 	const warnings = requestedRefs
 		.filter((ref) => !catalog.has(ref))
@@ -89,9 +89,8 @@ export async function executeWorkingMemory(options: ExecuteWorkingMemoryOptions)
 	}
 
 	const anchorIndex = findAnchorEntry(branch, options.toolCallId);
-	const retainedSet = new Set(retainedRefs);
 	const preAnchorUnits = [...catalog.values()].filter((unit) => unit.order < anchorIndex);
-	const prunedRowIds = [...new Set(preAnchorUnits.flatMap((unit) => (retainedSet.has(unit.ref) ? [] : unit.rowIds)))];
+	const prunedRowIds = collectPrunedRowIds(branch, anchorIndex);
 	const details: WorkingMemoryCheckpointDetailsV1 = {
 		v: 1,
 		anchorToolCallId: options.toolCallId,

package/extensions/working-memory/index.ts

@@ -22,7 +22,7 @@ import { replayWorkingMemoryState, WORKING_MEMORY_TOOL, type WorkingMemoryCheckp
 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.";
+	"Create a selective hard checkpoint for future model context. Keep valuable user and visible assistant messages, carry file structure as outlines, defer conditionally relevant files without reading them, and distill exploration findings into one compact continuation note.";
 
 interface NudgeState {
 	anchorToolCallId: string | undefined;
@@ -98,9 +98,9 @@ export default function workingMemoryExtension(pi: ExtensionAPI): void {
 					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.",
+						"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.",
+						"Outline files that remain in the repository working set, defer paths whose relevance is conditional, and discard exploration history after extracting its useful conclusions.",
+						"Everything before working_memory leaves future model context unless selected in keep. Put findings from tools, durable decisions, constraints, unresolved matters, and next action in continuation without duplicating retained messages.",
 					],
 					parameters: workingMemoryParameters,
 					executionMode: "sequential",

package/extensions/working-memory/memory.ts

@@ -9,23 +9,13 @@ export interface MemoryUnit {
 	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;
@@ -35,61 +25,56 @@ export function buildMemoryCatalog(branch: readonly SessionEntry[]): Map<string,
 			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 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(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),
+					preview: previewAssistant(text),
 					order,
-					suborder: index + 1,
-					messages: [{ ...message, content: [call] }, result],
-					rowIds: [call.id],
+					messages: [{ ...message, content: text }],
 				});
 			}
 			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;
+		if (message.role !== "user") continue;
 		const ref = `m:${entry.id}`;
 		catalog.set(ref, {
 			ref,
-			label: message.role === "custom" ? message.customType : message.role,
-			preview: previewMessage(message),
+			label: "user",
+			preview: previewUser(message),
 			order,
-			suborder: 0,
 			messages: [message],
-			rowIds: outlineRowIds(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" &&
+				isRecord(message.details) &&
+				typeof message.details.rowId === "string"
+			) {
+				rowIds.add(message.details.rowId);
+			}
+		}
+	}
+	return [...rowIds];
+}
+
 export function projectWorkingMemory(
 	messages: readonly ContextMessage[],
 	state: ActiveWorkingMemoryState,
@@ -109,7 +94,7 @@ export function projectWorkingMemory(
 			const unit = catalog.get(ref);
 			return unit ? [unit] : [];
 		})
-		.sort((left, right) => left.order - right.order || left.suborder - right.suborder);
+		.sort((left, right) => left.order - right.order);
 	const projected: ContextMessage[] = [];
 	const projectedRefs: MemoryUnit[][] = [];
 	for (const unit of retained) {
@@ -137,12 +122,6 @@ function referenceCurrentMessages(
 		(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 [];
@@ -150,13 +129,9 @@ function referenceCurrentMessages(
 		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) ?? [];
+		if (message.role !== "user" && message.role !== "assistant") return [];
 		const unit = catalog.get(`m:${entry.id}`);
-		return unit ? [unit] : [];
+		return unit && currentMessage.role === message.role ? [unit] : [];
 	});
 }
 
@@ -174,17 +149,33 @@ function addReferenceCatalog(
 			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 || left.suborder - right.suborder)
+				.sort((left, right) => left.order - right.order)
 				.map((unit) => `${unit.ref} (${unit.label}): ${unit.preview}`)
 				.join("\n"),
 			display: false,
@@ -224,36 +215,14 @@ function findAnchor(messages: readonly ContextMessage[], id: string): number {
 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 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 {

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@shanepadgett/tau-agent",
-	"version": "0.30.0",
+	"version": "0.31.1",
 	"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.30.0",
+		"@shanepadgett/tau-tui": "0.31.1",
 		"@vscode/tree-sitter-wasm": "0.3.1",
 		"image-size": "2.0.2",
 		"smol-toml": "1.7.0",