@agi-cli/server 0.1.58 → 0.1.61

package/src/runtime/cache-optimizer.ts

@@ -0,0 +1,93 @@
+import type { ModelMessage } from 'ai';
+
+type SystemMessage =
+	| string
+	| Array<{
+			type: 'text';
+			text: string;
+			cache_control?: { type: 'ephemeral' };
+	  }>;
+
+interface ContentPart {
+	type: string;
+	[key: string]: unknown;
+	providerOptions?: {
+		anthropic?: {
+			cacheControl?: { type: 'ephemeral' };
+		};
+	};
+}
+
+/**
+ * Adds cache control to messages for prompt caching optimization.
+ * Anthropic supports caching for system messages, tools, and long context.
+ */
+export function addCacheControl(
+	provider: string,
+	system: string | undefined,
+	messages: ModelMessage[],
+): {
+	system?: SystemMessage;
+	messages: ModelMessage[];
+} {
+	// Only Anthropic supports prompt caching currently
+	if (provider !== 'anthropic') {
+		return { system, messages };
+	}
+
+	// Convert system to cacheable format if it's long enough
+	let cachedSystem: SystemMessage | undefined = system;
+	if (system && system.length > 1024) {
+		// Anthropic requires 1024+ tokens for Claude Sonnet/Opus
+		cachedSystem = [
+			{
+				type: 'text',
+				text: system,
+				cache_control: { type: 'ephemeral' as const },
+			},
+		];
+	}
+
+	// Anthropic cache_control limits:
+	// - Max 4 cache blocks total
+	// - System message: 1 block
+	// - Tools: 2 blocks (read, write)
+	// - Last user message: 1 block
+	// Total: 4 blocks
+
+	// Add cache control to the last user message if conversation is long
+	// This caches the conversation history up to that point
+	if (messages.length >= 3) {
+		const cachedMessages = [...messages];
+
+		// Find second-to-last user message (not the current one)
+		const userIndices = cachedMessages
+			.map((m, i) => (m.role === 'user' ? i : -1))
+			.filter((i) => i >= 0);
+
+		if (userIndices.length >= 2) {
+			const targetIndex = userIndices[userIndices.length - 2];
+			const targetMsg = cachedMessages[targetIndex];
+
+			if (Array.isArray(targetMsg.content)) {
+				// Add cache control to the last content part of that message
+				const lastPart = targetMsg.content[targetMsg.content.length - 1];
+				if (lastPart && typeof lastPart === 'object' && 'type' in lastPart) {
+					(lastPart as ContentPart).providerOptions = {
+						anthropic: { cacheControl: { type: 'ephemeral' } },
+					};
+				}
+			}
+		}
+
+		return {
+			system: cachedSystem,
+			messages: cachedMessages,
+		};
+	}
+
+	return {
+		system: cachedSystem,
+		messages,
+	};
+}

package/src/runtime/context-optimizer.ts

@@ -0,0 +1,206 @@
+import type { ModelMessage } from 'ai';
+
+/**
+ * Optimizes message context by deduplicating file reads and pruning old tool results.
+ */
+
+interface FileRead {
+	messageIndex: number;
+	partIndex: number;
+	path: string;
+}
+
+interface ToolPart {
+	type: string;
+	input?: {
+		path?: string;
+		filePattern?: string;
+		pattern?: string;
+	};
+	output?: unknown;
+	[key: string]: unknown;
+}
+
+/**
+ * Deduplicates file read results, keeping only the latest version of each file.
+ *
+ * Strategy:
+ * - Track all file reads (read, grep, glob tools)
+ * - For files read multiple times, remove older results
+ * - Keep only the most recent read of each file
+ */
+export function deduplicateFileReads(messages: ModelMessage[]): ModelMessage[] {
+	const fileReads = new Map<string, FileRead[]>();
+
+	// First pass: identify all file reads and their locations
+	messages.forEach((msg, msgIdx) => {
+		if (msg.role !== 'assistant' || !Array.isArray(msg.content)) return;
+
+		msg.content.forEach((part, partIdx) => {
+			if (!part || typeof part !== 'object') return;
+			if (!('type' in part)) return;
+
+			const toolType = part.type as string;
+
+			// Check if this is a file read tool (read, grep, glob)
+			if (!toolType.startsWith('tool-')) return;
+
+			const toolName = toolType.replace('tool-', '');
+			if (!['read', 'grep', 'glob'].includes(toolName)) return;
+
+			// Extract file path from input
+			const toolPart = part as ToolPart;
+			const input = toolPart.input;
+			if (!input) return;
+
+			const path = input.path || input.filePattern || input.pattern;
+			if (!path) return;
+
+			// Track this file read
+			if (!fileReads.has(path)) {
+				fileReads.set(path, []);
+			}
+			fileReads
+				.get(path)
+				?.push({ messageIndex: msgIdx, partIndex: partIdx, path });
+		});
+	});
+
+	// Second pass: identify reads to remove (all but the latest for each file)
+	const readsToRemove = new Set<string>();
+
+	for (const [_path, reads] of fileReads) {
+		if (reads.length <= 1) continue;
+
+		// Sort by message index descending (latest first)
+		reads.sort((a, b) => b.messageIndex - a.messageIndex);
+
+		// Remove all but the first (latest)
+		for (let i = 1; i < reads.length; i++) {
+			const read = reads[i];
+			readsToRemove.add(`${read.messageIndex}-${read.partIndex}`);
+		}
+	}
+
+	// Third pass: rebuild messages without removed reads
+	return messages.map((msg, msgIdx) => {
+		if (msg.role !== 'assistant' || !Array.isArray(msg.content)) return msg;
+
+		const filteredContent = msg.content.filter((_part, partIdx) => {
+			const key = `${msgIdx}-${partIdx}`;
+			return !readsToRemove.has(key);
+		});
+
+		return {
+			...msg,
+			content: filteredContent,
+		};
+	});
+}
+
+/**
+ * Prunes old tool results to reduce context size.
+ *
+ * Strategy:
+ * - Keep only the last N tool results
+ * - Preserve tool calls but remove their output
+ * - Keep text parts intact
+ */
+export function pruneToolResults(
+	messages: ModelMessage[],
+	maxToolResults = 30,
+): ModelMessage[] {
+	// Collect all tool result locations
+	const toolResults: Array<{ messageIndex: number; partIndex: number }> = [];
+
+	messages.forEach((msg, msgIdx) => {
+		if (msg.role !== 'assistant' || !Array.isArray(msg.content)) return;
+
+		msg.content.forEach((part, partIdx) => {
+			if (!part || typeof part !== 'object') return;
+			if (!('type' in part)) return;
+
+			const toolType = part.type as string;
+			if (!toolType.startsWith('tool-')) return;
+
+			// Check if this has output
+			const toolPart = part as ToolPart;
+			const hasOutput = toolPart.output !== undefined;
+			if (!hasOutput) return;
+
+			toolResults.push({ messageIndex: msgIdx, partIndex: partIdx });
+		});
+	});
+
+	// If under limit, no pruning needed
+	if (toolResults.length <= maxToolResults) {
+		return messages;
+	}
+
+	// Keep only the last N tool results
+	const toKeep = new Set<string>();
+	const keepCount = Math.min(maxToolResults, toolResults.length);
+	const keepStart = toolResults.length - keepCount;
+
+	for (let i = keepStart; i < toolResults.length; i++) {
+		const result = toolResults[i];
+		toKeep.add(`${result.messageIndex}-${result.partIndex}`);
+	}
+
+	// Rebuild messages, removing old tool outputs
+	return messages.map((msg, msgIdx) => {
+		if (msg.role !== 'assistant' || !Array.isArray(msg.content)) return msg;
+
+		const processedContent = msg.content.map((part, partIdx) => {
+			if (!part || typeof part !== 'object') return part;
+			if (!('type' in part)) return part;
+
+			const toolPart = part as ToolPart;
+			const toolType = toolPart.type;
+			if (!toolType.startsWith('tool-')) return part;
+
+			const key = `${msgIdx}-${partIdx}`;
+			const hasOutput = toolPart.output !== undefined;
+
+			// If this tool result should be pruned, remove its output
+			if (hasOutput && !toKeep.has(key)) {
+				return {
+					...part,
+					output: '[pruned to save context]',
+				};
+			}
+
+			return part;
+		});
+
+		return {
+			...msg,
+			content: processedContent,
+		};
+	});
+}
+
+/**
+ * Applies all context optimizations:
+ * 1. Deduplicate file reads
+ * 2. Prune old tool results
+ */
+export function optimizeContext(
+	messages: ModelMessage[],
+	options: {
+		deduplicateFiles?: boolean;
+		maxToolResults?: number;
+	} = {},
+): ModelMessage[] {
+	let optimized = messages;
+
+	if (options.deduplicateFiles !== false) {
+		optimized = deduplicateFileReads(optimized);
+	}
+
+	if (options.maxToolResults !== undefined) {
+		optimized = pruneToolResults(optimized, options.maxToolResults);
+	}
+
+	return optimized;
+}

package/src/runtime/db-operations.ts

@@ -3,8 +3,103 @@ import { messages, messageParts, sessions } from '@agi-cli/database/schema';
 import { eq } from 'drizzle-orm';
 import type { RunOpts } from './session-queue.ts';
 
+type UsageData = {
+	inputTokens?: number;
+	outputTokens?: number;
+	totalTokens?: number;
+	cachedInputTokens?: number;
+	reasoningTokens?: number;
+};
+
+interface ProviderMetadata {
+	openai?: {
+		cachedPromptTokens?: number;
+	};
+	[key: string]: unknown;
+}
+
+/**
+ * Updates session token counts incrementally after each step.
+ * Note: onStepFinish.usage is CUMULATIVE per message, so we compute DELTA and add to session.
+ */
+export async function updateSessionTokensIncremental(
+	usage: UsageData,
+	providerMetadata: ProviderMetadata | undefined,
+	opts: RunOpts,
+	db: Awaited<ReturnType<typeof getDb>>,
+) {
+	if (!usage) return;
+
+	// Read session totals
+	const sessRows = await db
+		.select()
+		.from(sessions)
+		.where(eq(sessions.id, opts.sessionId));
+
+	if (sessRows.length === 0 || !sessRows[0]) return;
+
+	const sess = sessRows[0];
+	const priorInputSess = Number(sess.totalInputTokens ?? 0);
+	const priorOutputSess = Number(sess.totalOutputTokens ?? 0);
+	const priorCachedSess = Number(sess.totalCachedTokens ?? 0);
+	const priorReasoningSess = Number(sess.totalReasoningTokens ?? 0);
+
+	// Read current message totals to compute delta
+	const msgRows = await db
+		.select()
+		.from(messages)
+		.where(eq(messages.id, opts.assistantMessageId));
+
+	const msg = msgRows[0];
+	const priorPromptMsg = Number(msg?.promptTokens ?? 0);
+	const priorCompletionMsg = Number(msg?.completionTokens ?? 0);
+	const priorCachedMsg = Number(msg?.cachedInputTokens ?? 0);
+	const priorReasoningMsg = Number(msg?.reasoningTokens ?? 0);
+
+	// Treat usage as cumulative per-message for this step
+	const cumPrompt =
+		usage.inputTokens != null ? Number(usage.inputTokens) : priorPromptMsg;
+	const cumCompletion =
+		usage.outputTokens != null
+			? Number(usage.outputTokens)
+			: priorCompletionMsg;
+	const cumReasoning =
+		usage.reasoningTokens != null
+			? Number(usage.reasoningTokens)
+			: priorReasoningMsg;
+
+	const cumCached =
+		usage.cachedInputTokens != null
+			? Number(usage.cachedInputTokens)
+			: providerMetadata?.openai?.cachedPromptTokens != null
+				? Number(providerMetadata.openai.cachedPromptTokens)
+				: priorCachedMsg;
+
+	// Compute deltas for this step; clamp to 0 in case provider reports smaller values
+	const deltaInput = Math.max(0, cumPrompt - priorPromptMsg);
+	const deltaOutput = Math.max(0, cumCompletion - priorCompletionMsg);
+	const deltaCached = Math.max(0, cumCached - priorCachedMsg);
+	const deltaReasoning = Math.max(0, cumReasoning - priorReasoningMsg);
+
+	const nextInputSess = priorInputSess + deltaInput;
+	const nextOutputSess = priorOutputSess + deltaOutput;
+	const nextCachedSess = priorCachedSess + deltaCached;
+	const nextReasoningSess = priorReasoningSess + deltaReasoning;
+
+	await db
+		.update(sessions)
+		.set({
+			totalInputTokens: nextInputSess,
+			totalOutputTokens: nextOutputSess,
+			totalCachedTokens: nextCachedSess,
+			totalReasoningTokens: nextReasoningSess,
+		})
+		.where(eq(sessions.id, opts.sessionId));
+}
+
 /**
  * Updates session token counts after a run completes.
+ * @deprecated Use updateSessionTokensIncremental for per-step tracking
  */
 export async function updateSessionTokens(
 	fin: { usage?: { inputTokens?: number; outputTokens?: number } },
@@ -36,59 +131,98 @@ export async function updateSessionTokens(
 }
 
 /**
- * Marks an assistant message as complete with token usage information.
+ * Updates message token counts incrementally after each step.
+ * Note: onStepFinish.usage is CUMULATIVE per message, so we REPLACE values, not add.
+ */
+export async function updateMessageTokensIncremental(
+	usage: UsageData,
+	providerMetadata: ProviderMetadata | undefined,
+	opts: RunOpts,
+	db: Awaited<ReturnType<typeof getDb>>,
+) {
+	if (!usage) return;
+
+	const msgRows = await db
+		.select()
+		.from(messages)
+		.where(eq(messages.id, opts.assistantMessageId));
+
+	if (msgRows.length > 0 && msgRows[0]) {
+		const msg = msgRows[0];
+		const priorPrompt = Number(msg.promptTokens ?? 0);
+		const priorCompletion = Number(msg.completionTokens ?? 0);
+		const priorCached = Number(msg.cachedInputTokens ?? 0);
+		const priorReasoning = Number(msg.reasoningTokens ?? 0);
+
+		// Treat usage as cumulative per-message - REPLACE not ADD
+		const nextPrompt =
+			usage.inputTokens != null ? Number(usage.inputTokens) : priorPrompt;
+		const nextCompletion =
+			usage.outputTokens != null ? Number(usage.outputTokens) : priorCompletion;
+		const nextReasoning =
+			usage.reasoningTokens != null
+				? Number(usage.reasoningTokens)
+				: priorReasoning;
+
+		const nextCached =
+			usage.cachedInputTokens != null
+				? Number(usage.cachedInputTokens)
+				: providerMetadata?.openai?.cachedPromptTokens != null
+					? Number(providerMetadata.openai.cachedPromptTokens)
+					: priorCached;
+
+		await db
+			.update(messages)
+			.set({
+				promptTokens: nextPrompt,
+				completionTokens: nextCompletion,
+				cachedInputTokens: nextCached,
+				reasoningTokens: nextReasoning,
+			})
+			.where(eq(messages.id, opts.assistantMessageId));
+	}
+}
+
+/**
+ * Completes the assistant message after the run finishes.
+ * Used to finalize timing but NOT tokens, which are already incremental.
  */
 export async function completeAssistantMessage(
-	fin: {
+	_fin: {
 		usage?: {
 			inputTokens?: number;
 			outputTokens?: number;
-			totalTokens?: number;
 		};
 	},
 	opts: RunOpts,
 	db: Awaited<ReturnType<typeof getDb>>,
 ) {
-	const vals: Record<string, unknown> = {
-		status: 'complete',
-		completedAt: Date.now(),
-	};
+	const msgRow = await db
+		.select()
+		.from(messages)
+		.where(eq(messages.id, opts.assistantMessageId));
 
-	if (fin.usage) {
-		vals.promptTokens = fin.usage.inputTokens;
-		vals.completionTokens = fin.usage.outputTokens;
-		vals.totalTokens =
-			fin.usage.totalTokens ??
-			(vals.promptTokens as number) + (vals.completionTokens as number);
+	if (msgRow.length > 0) {
+		await db
+			.update(messages)
+			.set({
+				finishedAt: new Date(),
+			})
+			.where(eq(messages.id, opts.assistantMessageId));
 	}
-
-	await db
-		.update(messages)
-		.set(vals)
-		.where(eq(messages.id, opts.assistantMessageId));
 }
 
-/**
- * Removes empty text parts from an assistant message.
- */
-export async function cleanupEmptyTextParts(
-	opts: RunOpts,
+export async function createMessagePart(
+	partData: {
+		messageId: number;
+		contentType: 'text' | 'tool' | 'other';
+		toolName?: string | null;
+		toolArgs?: unknown;
+		toolResult?: unknown;
+		textContent?: string | null;
+		stepIndex?: number | null;
+	},
 	db: Awaited<ReturnType<typeof getDb>>,
 ) {
-	const parts = await db
-		.select()
-		.from(messageParts)
-		.where(eq(messageParts.messageId, opts.assistantMessageId));
-
-	for (const p of parts) {
-		if (p.type === 'text') {
-			let t = '';
-			try {
-				t = JSON.parse(p.content || '{}')?.text || '';
-			} catch {}
-			if (!t || t.length === 0) {
-				await db.delete(messageParts).where(eq(messageParts.id, p.id));
-			}
-		}
-	}
+	await db.insert(messageParts).values(partData);
 }

package/src/runtime/history-truncator.ts

@@ -0,0 +1,26 @@
+import type { ModelMessage } from 'ai';
+
+/**
+ * Truncates conversation history to keep only the most recent messages.
+ * This helps manage context window size and improves performance.
+ *
+ * Strategy:
+ * - Keep only the last N messages
+ * - Preserve message pairs (assistant + user responses) when possible
+ * - Always keep at least the system message if present
+ */
+export function truncateHistory(
+	messages: ModelMessage[],
+	maxMessages: number,
+): ModelMessage[] {
+	if (messages.length <= maxMessages) {
+		return messages;
+	}
+
+	// Calculate how many messages to keep
+	const keepCount = Math.min(maxMessages, messages.length);
+	const startIndex = messages.length - keepCount;
+
+	// Return the most recent messages
+	return messages.slice(startIndex);
+}