Keep message progress on abort

package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@ztimson/ai-utils",
-	"version": "1.6.4",
+	"version": "1.6.6",
 	"description": "AI Utility library",
 	"author": "Zak Timson",
 	"license": "MIT",

src/ai.ts

@@ -4,7 +4,7 @@ import { Audio } from './audio.ts';
 import {Vision} from './vision.ts';
 
 export type AbortablePromise<T> = Promise<T> & {
-	abort: () => any
+	abort: (keep?: boolean) => any
 };
 
 export type AiOptions = {

src/llm.ts

@@ -265,7 +265,7 @@ class LLM {
 		};
 	}
 
-	private setupAgent(agents: Agent[] = [], allAgents: Agent[], history: LLMMessage[], aborts: (() => void)[], depth = 0, delegateState: {resp: string | null}): AiTool[] {
+	private setupAgent(agents: Agent[] = [], allAgents: Agent[], history: LLMMessage[], aborts: ((keep?: boolean) => void)[], depth = 0, delegateState: {resp: string | null}): AiTool[] {
 		return agents.map(a => {
 			const toolName = `${a.delegate ? '' : 'sub'}agent_${snakeCase(a.name)}`;
 			return {
@@ -397,11 +397,13 @@ ${a.system}`,
 		if(!this.models[m]) throw new Error(`Model does not exist: ${m}`);
 		let request: AbortablePromise<string> | null = null;
 		let aborted = false;
-		const nestedAborts: (() => void)[] = [];
-		const abort = () => {
+		let keepOnAbort = true;
+		const nestedAborts: ((keep?: boolean) => void)[] = [];
+		const abort = (keep = true) => {
 			aborted = true;
-			request?.abort?.();
-			nestedAborts.forEach(a => a());
+			keepOnAbort = keep;
+			request?.abort?.(keep);
+			nestedAborts.forEach(a => a(keep));
 		};
 
 		let promise: any;
@@ -411,9 +413,25 @@ ${a.system}`,
 			let tools: AiTool[] = options.tools || this.ai.options.llm?.tools || [];
 			const prompts: string[] = [];
 			let history = options.history || [];
+			const historyStart = history.length;
 			const files = options.files || [];
 			if(message || files.length) history.push({role: 'user', content: message || '', timestamp: Date.now()});
 
+			// Accumulate streamed text so it can be committed to history if aborted mid-generation
+			let partialText = '';
+			const onStream = options.stream;
+			const stream = (chunk: {text?: string, tool?: string, done?: true}) => {
+				if(chunk.text) partialText += chunk.text;
+				return onStream?.(chunk);
+			};
+
+			/** Commit (keep) or discard this turn's progress on abort, then throw */
+			const abortNow = (): never => {
+				if(keepOnAbort) { if(partialText) history.push({role: 'assistant', content: partialText, timestamp: Date.now()}); }
+				else history.splice(historyStart, history.length - historyStart);
+				throw Object.assign(new Error('Aborted'), {name: 'AbortError'});
+			};
+
 			// MCP
 			const mcp = options.mcp || this.ai.options?.llm?.mcp;
 			if(mcp?.length) {
@@ -441,8 +459,8 @@ ${a.system}`,
 				const mems = mem.memory instanceof MemoryCache ? mem.memory.memories : mem.memory;
 				if(mems.length) {
 					if(mem.inject) {
-						const pool = 15; // candidates considered, cheap since only refs are listed
-						const budget = mem.maxTokens ?? 2000; // actual content injected
+						const pool = 15;
+						const budget = mem.maxTokens ?? 2000;
 						const relevant = await this.memoryManager.recollect(message, mem.memory, pool);
 
 						let used = 0;
@@ -481,7 +499,7 @@ Linked: ${makeUnique([...r.links, ...r.backlinks]).join(', ')}
 				}
 			}
 
-			if(aborted) throw Object.assign(new Error('Aborted'), {name: 'AbortError'});
+			if(aborted) abortNow();
 
 			const lastMsg = history[history.length - 1];
 			if(files.length && lastMsg?.role === 'user') lastMsg.files = files;
@@ -500,11 +518,17 @@ Linked: ${makeUnique([...r.links, ...r.backlinks]).join(', ')}
 			const toolTimings = new Map<string, {duration: number, tps: number}>();
 			tools = this.wrapToolTiming(tools, toolTimings);
 
-			if(aborted) throw Object.assign(new Error('Aborted'), {name: 'AbortError'});
+			if(aborted) abortNow();
 
 			prompts.unshift(options.system || this.ai.options.llm?.system || '');
-			request = this.models[m].ask('', {...options, tools, system: prompts.filter(Boolean).join('\n\n')});
-			let resp = await request;
+			request = this.models[m].ask('', {...options, tools, stream, system: prompts.filter(Boolean).join('\n\n')});
+			let resp: string;
+			try {
+				resp = await request;
+			} catch(err: any) {
+				if(aborted) return abortNow();
+				throw err;
+			}
 
 			// Strip the file injection shim
 			restores.forEach(({msg, content}) => msg.content = content);

src/memory.ts

@@ -4,7 +4,7 @@ import {AiTool} from './tools.ts';
 import {KDPoint, KDTree} from './kd-tree.ts';
 import {escapeRegex} from '@ztimson/utils';
 
-const MERGE_THRESHOLD = 0.88;
+const MERGE_THRESHOLD = 0.12;
 const PENDING_HEADING = '## Pending';
 const GENERIC_TEMPLATE = `# {{Title}}
 
@@ -191,13 +191,13 @@ export type MemoryOptions = {
 }
 
 export class MemoryManager {
-	private recentlyTouched = new Map<string, number>();
-
+	private mergeLock: Promise<any> = Promise.resolve();
 	private queues = new Map<string, {
 		dirty: boolean,
 		request: {abort?: () => void} | null,
 		task: Promise<void>,
 	}>();
+	private recentlyTouched = new Map<string, number>();
 
 	tools = {
 		forget: (memories: Memory[] | MemoryCache): AiTool => ({
@@ -351,10 +351,6 @@ ${ghosts.length ? `${ghosts.map(g => `- ${g}: (Ghost)`).join('\n')}` : ''}`,
 		return memories.map(m => ({name: m.name, description: m.description}));
 	}
 
-	/** Find the nearest node above the similarity threshold and fold the smaller/less-connected one into
-	 * the other. Journals are exempt — they're partitioned by date, not topic, and merging across weeks
-	 * would wreck the timeline. Returns 'merged' if `node` absorbed another (caller should re-run the doc
-	 * agent), 'absorbed' if `node` itself got folded away (caller should stop touching it), or null. */
 	private async checkMerge(node: Memory, memories: Memory[] | MemoryCache, options: LLMRequest, threshold = MERGE_THRESHOLD): Promise<Memory | null> {
 		if (!node.embedding?.length || node.name.startsWith('Journal/')) return null;
 		const store = this.access(memories);
@@ -404,8 +400,9 @@ ${ghosts.length ? `${ghosts.map(g => `- ${g}: (Ghost)`).join('\n')}` : ''}`,
 			do {
 				entry.dirty = false;
 				await this.docAgent(current, store.list, options, entry);
-				const merged = await this.checkMerge(current, memories, options);
-				if (merged) { current = merged; entry.dirty = true; }
+				this.mergeLock = this.mergeLock.then(() => this.checkMerge(current, memories, options));
+				const merged = await this.mergeLock;
+				if(merged) current = merged;
 			} while (entry.dirty);
 		})().finally(() => {
 			this.queues.delete(key);
@@ -432,11 +429,9 @@ ${ghosts.length ? `${ghosts.map(g => `- ${g}: (Ghost)`).join('\n')}` : ''}`,
 If it has a "## Pending" section, fold all new material into the appropriate part, resolve overlap, then remove the section entirely. If no section, just tidy per the rules below.
 
 Use this loose structure, adapting headings to what the content needs:
-
-# Title
-## Summary
-## Details
-## Related
+\`\`\`markdown
+${GENERIC_TEMPLATE}
+\`\`\`
 
 Rules:
 - Contradictions: newer facts always win — delete outdated statements entirely
@@ -483,11 +478,9 @@ ${currentBody}
 			system: `You are a knowledge base editor merging two overlapping Obsidian documents into one. Newer facts win on contradiction.
 
 Structure loosely:
-
-# Title
-## Summary
-## Details
-## Related
+\`\`\`markdown
+${GENERIC_TEMPLATE}
+\`\`\`
 
 Combine both documents, resolve duplication and contradictions.
 
@@ -617,16 +610,16 @@ ${stripHeader(b.content)}
 			touched.push(node);
 		}
 
-		for (const node of touched) {
+		await Promise.all(touched.map(async node => {
 			const [e] = await this.llm.embedding(`${node.description}\n\n${stripHeader(node.content)}`.trim());
 			if (e) node.embedding = e.embedding;
 			this.touch(node.name);
-		}
+		}));
 
 		if (touched.length) {
 			store.commit();
 			(pending as any).content = `Saved to ${touched.map(n => `[[${n.name}]]`).join(', ')}`;
-			await Promise.all(touched.map(node => this.reconcile(node, memories, options).catch(() => {})));
+			Promise.all(touched.map(node => this.reconcile(node, memories, options).catch(() => {})));
 		} else {
 			(pending as any).content = 'Nothing worth remembering.';
 		}