Keep message progress on abort

package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@ztimson/ai-utils",
-	"version": "1.6.5",
+	"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);