@agi-cli/server 0.1.61 → 0.1.62

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@agi-cli/server",
-	"version": "0.1.61",
+	"version": "0.1.62",
 	"description": "HTTP API server for AGI CLI",
 	"type": "module",
 	"main": "./src/index.ts",
@@ -29,8 +29,8 @@
 		"typecheck": "tsc --noEmit"
 	},
 	"dependencies": {
-		"@agi-cli/sdk": "0.1.61",
-		"@agi-cli/database": "0.1.61",
+		"@agi-cli/sdk": "0.1.62",
+		"@agi-cli/database": "0.1.62",
 		"drizzle-orm": "^0.44.5",
 		"hono": "^4.9.9"
 	},

package/src/runtime/cache-optimizer.ts

@@ -1,23 +1,5 @@
 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.
@@ -27,7 +9,13 @@ export function addCacheControl(
 	system: string | undefined,
 	messages: ModelMessage[],
 ): {
-	system?: SystemMessage;
+	system?:
+		| string
+		| Array<{
+				type: 'text';
+				text: string;
+				cache_control?: { type: 'ephemeral' };
+		  }>;
 	messages: ModelMessage[];
 } {
 	// Only Anthropic supports prompt caching currently
@@ -36,7 +24,7 @@ export function addCacheControl(
 	}
 
 	// Convert system to cacheable format if it's long enough
-	let cachedSystem: SystemMessage | undefined = system;
+	let cachedSystem: any = system;
 	if (system && system.length > 1024) {
 		// Anthropic requires 1024+ tokens for Claude Sonnet/Opus
 		cachedSystem = [
@@ -73,21 +61,55 @@ export function addCacheControl(
 				// 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 = {
+					(lastPart as any).providerOptions = {
 						anthropic: { cacheControl: { type: 'ephemeral' } },
 					};
 				}
 			}
 		}
 
-		return {
-			system: cachedSystem,
-			messages: cachedMessages,
-		};
+		return { system: cachedSystem, messages: cachedMessages };
+	}
+
+	return { system: cachedSystem, messages };
+}
+
+/**
+ * Truncates old messages to reduce context size while keeping recent context.
+ * Strategy: Keep system message + last N messages
+ */
+export function truncateHistory(
+	messages: ModelMessage[],
+	maxMessages = 20,
+): ModelMessage[] {
+	if (messages.length <= maxMessages) {
+		return messages;
+	}
+
+	// Keep the most recent messages
+	return messages.slice(-maxMessages);
+}
+
+/**
+ * Estimates token count (rough approximation: ~4 chars per token)
+ */
+export function estimateTokens(text: string): number {
+	return Math.ceil(text.length / 4);
+}
+
+/**
+ * Summarizes tool results if they're too long
+ */
+export function summarizeToolResult(result: unknown, maxLength = 5000): string {
+	const str = typeof result === 'string' ? result : JSON.stringify(result);
+
+	if (str.length <= maxLength) {
+		return str;
 	}
 
-	return {
-		system: cachedSystem,
-		messages,
-	};
+	// Truncate and add indicator
+	return (
+		str.slice(0, maxLength) +
+		`\n\n[... truncated ${str.length - maxLength} characters]`
+	);
 }

package/src/runtime/db-operations.ts

@@ -11,20 +11,13 @@ type UsageData = {
 	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,
+	providerMetadata: Record<string, any> | undefined,
 	opts: RunOpts,
 	db: Awaited<ReturnType<typeof getDb>>,
 ) {
@@ -136,7 +129,7 @@ export async function updateSessionTokens(
  */
 export async function updateMessageTokensIncremental(
 	usage: UsageData,
-	providerMetadata: ProviderMetadata | undefined,
+	providerMetadata: Record<string, any> | undefined,
 	opts: RunOpts,
 	db: Awaited<ReturnType<typeof getDb>>,
 ) {
@@ -155,74 +148,86 @@ export async function updateMessageTokensIncremental(
 		const priorReasoning = Number(msg.reasoningTokens ?? 0);
 
 		// Treat usage as cumulative per-message - REPLACE not ADD
-		const nextPrompt =
+		const cumPrompt =
 			usage.inputTokens != null ? Number(usage.inputTokens) : priorPrompt;
-		const nextCompletion =
+		const cumCompletion =
 			usage.outputTokens != null ? Number(usage.outputTokens) : priorCompletion;
-		const nextReasoning =
+		const cumReasoning =
 			usage.reasoningTokens != null
 				? Number(usage.reasoningTokens)
 				: priorReasoning;
 
-		const nextCached =
+		const cumCached =
 			usage.cachedInputTokens != null
 				? Number(usage.cachedInputTokens)
 				: providerMetadata?.openai?.cachedPromptTokens != null
 					? Number(providerMetadata.openai.cachedPromptTokens)
 					: priorCached;
 
+		const cumTotal =
+			usage.totalTokens != null
+				? Number(usage.totalTokens)
+				: cumPrompt + cumCompletion + cumReasoning;
+
 		await db
 			.update(messages)
 			.set({
-				promptTokens: nextPrompt,
-				completionTokens: nextCompletion,
-				cachedInputTokens: nextCached,
-				reasoningTokens: nextReasoning,
+				promptTokens: cumPrompt,
+				completionTokens: cumCompletion,
+				totalTokens: cumTotal,
+				cachedInputTokens: cumCached,
+				reasoningTokens: cumReasoning,
 			})
 			.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.
+ * Marks an assistant message as complete.
+ * Token usage is tracked incrementally via updateMessageTokensIncremental().
  */
 export async function completeAssistantMessage(
-	_fin: {
+	fin: {
 		usage?: {
 			inputTokens?: number;
 			outputTokens?: number;
+			totalTokens?: number;
 		};
 	},
 	opts: RunOpts,
 	db: Awaited<ReturnType<typeof getDb>>,
 ) {
-	const msgRow = await db
-		.select()
-		.from(messages)
+	// Only mark as complete - tokens are already tracked incrementally
+	await db
+		.update(messages)
+		.set({
+			status: 'complete',
+			completedAt: Date.now(),
+		})
 		.where(eq(messages.id, opts.assistantMessageId));
-
-	if (msgRow.length > 0) {
-		await db
-			.update(messages)
-			.set({
-				finishedAt: new Date(),
-			})
-			.where(eq(messages.id, opts.assistantMessageId));
-	}
 }
 
-export async function createMessagePart(
-	partData: {
-		messageId: number;
-		contentType: 'text' | 'tool' | 'other';
-		toolName?: string | null;
-		toolArgs?: unknown;
-		toolResult?: unknown;
-		textContent?: string | null;
-		stepIndex?: number | null;
-	},
+/**
+ * Removes empty text parts from an assistant message.
+ */
+export async function cleanupEmptyTextParts(
+	opts: RunOpts,
 	db: Awaited<ReturnType<typeof getDb>>,
 ) {
-	await db.insert(messageParts).values(partData);
+	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));
+			}
+		}
+	}
 }

package/src/runtime/runner.ts

@@ -1,4 +1,4 @@
-import { streamText } from 'ai';
+import { hasToolCall, streamText } from 'ai';
 import { loadConfig } from '@agi-cli/sdk';
 import { getDb } from '@agi-cli/database';
 import { messageParts } from '@agi-cli/database/schema';
@@ -7,8 +7,11 @@ import { resolveModel } from './provider.ts';
 import { resolveAgentConfig } from './agent-registry.ts';
 import { composeSystemPrompt } from './prompt.ts';
 import { discoverProjectTools } from '@agi-cli/sdk';
-import { publish } from '../events/bus.ts';
+import { adaptTools } from '../tools/adapter.ts';
+import { publish, subscribe } from '../events/bus.ts';
+import { debugLog, time } from './debug.ts';
 import { buildHistoryMessages } from './history-builder.ts';
+import { toErrorPayload } from './error-handling.ts';
 import { getMaxOutputTokens } from './token-utils.ts';
 import {
 	type RunOpts,
@@ -19,11 +22,15 @@ import {
 	dequeueJob,
 	cleanupSession,
 } from './session-queue.ts';
-import { setupToolContext } from './tool-context-setup.ts';
+import {
+	setupToolContext,
+	type RunnerToolContext,
+} from './tool-context-setup.ts';
 import {
 	updateSessionTokensIncremental,
 	updateMessageTokensIncremental,
 	completeAssistantMessage,
+	cleanupEmptyTextParts,
 } from './db-operations.ts';
 import {
 	createStepFinishHandler,
@@ -31,38 +38,175 @@ import {
 	createAbortHandler,
 	createFinishHandler,
 } from './stream-handlers.ts';
-import { addCacheControl } from './cache-optimizer.ts';
-import { optimizeContext } from './context-optimizer.ts';
-import { truncateHistory } from './history-truncator.ts';
 
 /**
- * Main runner that executes the LLM streaming loop with tools
+ * Enqueues an assistant run for processing.
+ */
+export function enqueueAssistantRun(opts: Omit<RunOpts, 'abortSignal'>) {
+	enqueueRun(opts, processQueue);
+}
+
+/**
+ * Aborts an active session.
+ */
+export function abortSession(sessionId: string) {
+	abortSessionQueue(sessionId);
+}
+
+/**
+ * Processes the queue of assistant runs for a session.
+ */
+async function processQueue(sessionId: string) {
+	const state = getRunnerState(sessionId);
+	if (!state) return;
+	if (state.running) return;
+	setRunning(sessionId, true);
+
+	while (state.queue.length > 0) {
+		const job = dequeueJob(sessionId);
+		if (!job) break;
+		try {
+			await runAssistant(job);
+		} catch (_err) {
+			// Swallow to keep the loop alive; event published by runner
+		}
+	}
+
+	setRunning(sessionId, false);
+	cleanupSession(sessionId);
+}
+
+/**
+ * Ensures the finish tool is called if not already observed.
+ */
+async function ensureFinishToolCalled(
+	finishObserved: boolean,
+	toolset: ReturnType<typeof adaptTools>,
+	sharedCtx: RunnerToolContext,
+	stepIndex: number,
+) {
+	if (finishObserved || !toolset?.finish?.execute) return;
+
+	const finishInput = {} as const;
+	const callOptions = { input: finishInput } as const;
+
+	sharedCtx.stepIndex = stepIndex;
+
+	try {
+		await toolset.finish.onInputStart?.(callOptions as never);
+	} catch {}
+
+	try {
+		await toolset.finish.onInputAvailable?.(callOptions as never);
+	} catch {}
+
+	await toolset.finish.execute(finishInput, {} as never);
+}
+
+/**
+ * Main function to run the assistant for a given request.
  */
-export async function runAssistant(opts: RunOpts) {
-	const db = await getDb();
-	const config = await loadConfig();
-	const [provider, modelName] = opts.model.split('/', 2);
-	const model = resolveModel(provider, modelName);
-
-	// Build agent + system prompt
-	const agentConfig = resolveAgentConfig(opts.agent);
-	const availableTools = await discoverProjectTools(config.project.root);
-	const system = composeSystemPrompt(agentConfig, availableTools);
-
-	// Build message history
-	const history = await buildHistoryMessages(opts, db);
-
-	// Setup tool context
-	const toolContext = await setupToolContext(opts, db);
-	const { tools, sharedCtx } = toolContext;
-
-	// State
-	let currentPartId = sharedCtx.assistantPartId;
-	let stepIndex = sharedCtx.stepIndex;
+async function runAssistant(opts: RunOpts) {
+	const cfgTimer = time('runner:loadConfig+db');
+	const cfg = await loadConfig(opts.projectRoot);
+	const db = await getDb(cfg.projectRoot);
+	cfgTimer.end();
+
+	const agentTimer = time('runner:resolveAgentConfig');
+	const agentCfg = await resolveAgentConfig(cfg.projectRoot, opts.agent);
+	agentTimer.end({ agent: opts.agent });
+
+	const agentPrompt = agentCfg.prompt || '';
+
+	const historyTimer = time('runner:buildHistory');
+	const history = await buildHistoryMessages(db, opts.sessionId);
+	historyTimer.end({ messages: history.length });
+
+	const isFirstMessage = history.length === 0;
+
+	const systemTimer = time('runner:composeSystemPrompt');
+	const { getAuth } = await import('@agi-cli/sdk');
+	const { getProviderSpoofPrompt } = await import('./prompt.ts');
+	const auth = await getAuth(opts.provider, cfg.projectRoot);
+	const needsSpoof = auth?.type === 'oauth';
+	const spoofPrompt = needsSpoof
+		? getProviderSpoofPrompt(opts.provider)
+		: undefined;
+
+	let system: string;
+	let additionalSystemMessages: Array<{ role: 'system'; content: string }> = [];
+
+	if (spoofPrompt) {
+		system = spoofPrompt;
+		const fullPrompt = await composeSystemPrompt({
+			provider: opts.provider,
+			model: opts.model,
+			projectRoot: cfg.projectRoot,
+			agentPrompt,
+			oneShot: opts.oneShot,
+			spoofPrompt: undefined,
+			includeProjectTree: isFirstMessage,
+		});
+		additionalSystemMessages = [{ role: 'system', content: fullPrompt }];
+	} else {
+		system = await composeSystemPrompt({
+			provider: opts.provider,
+			model: opts.model,
+			projectRoot: cfg.projectRoot,
+			agentPrompt,
+			oneShot: opts.oneShot,
+			spoofPrompt: undefined,
+			includeProjectTree: isFirstMessage,
+		});
+	}
+	systemTimer.end();
+	debugLog('[system] composed prompt (provider+base+agent):');
+	debugLog(system);
+
+	const toolsTimer = time('runner:discoverTools');
+	const allTools = await discoverProjectTools(cfg.projectRoot);
+	toolsTimer.end({ count: allTools.length });
+	const allowedNames = new Set([
+		...(agentCfg.tools || []),
+		'finish',
+		'progress_update',
+	]);
+	const gated = allTools.filter((t) => allowedNames.has(t.name));
+	const messagesWithSystemInstructions = [
+		...(isFirstMessage ? additionalSystemMessages : []),
+		...history,
+	];
+
+	const { sharedCtx, firstToolTimer, firstToolSeen } = await setupToolContext(
+		opts,
+		db,
+	);
+	const toolset = adaptTools(gated, sharedCtx, opts.provider);
+
+	const modelTimer = time('runner:resolveModel');
+	const model = await resolveModel(opts.provider, opts.model, cfg);
+	modelTimer.end();
+
+	const maxOutputTokens = getMaxOutputTokens(opts.provider, opts.model);
+
+	let currentPartId = opts.assistantPartId;
 	let accumulated = '';
-	const abortController = new AbortController();
+	let stepIndex = 0;
+
+	let finishObserved = false;
+	const unsubscribeFinish = subscribe(opts.sessionId, (evt) => {
+		if (evt.type !== 'tool.result') return;
+		try {
+			const name = (evt.payload as { name?: string } | undefined)?.name;
+			if (name === 'finish') finishObserved = true;
+		} catch {}
+	});
 
-	// State getters/setters
+	const streamStartTimer = time('runner:first-delta');
+	let firstDeltaSeen = false;
+	debugLog(`[streamText] Calling with maxOutputTokens: ${maxOutputTokens}`);
+
+	// State management helpers
 	const getCurrentPartId = () => currentPartId;
 	const getStepIndex = () => stepIndex;
 	const updateCurrentPartId = (id: string) => {
@@ -71,10 +215,12 @@ export async function runAssistant(opts: RunOpts) {
 	const updateAccumulated = (text: string) => {
 		accumulated = text;
 	};
-	const getAccumulated = () => accumulated;
-	const incrementStepIndex = () => ++stepIndex;
+	const incrementStepIndex = () => {
+		stepIndex += 1;
+		return stepIndex;
+	};
 
-	// Handlers
+	// Create stream handlers
 	const onStepFinish = createStepFinishHandler(
 		opts,
 		db,
@@ -88,102 +234,105 @@ export async function runAssistant(opts: RunOpts) {
 		updateMessageTokensIncremental,
 	);
 
+	const onError = createErrorHandler(opts, db, getStepIndex, sharedCtx);
+
+	const onAbort = createAbortHandler(opts, db, getStepIndex, sharedCtx);
+
 	const onFinish = createFinishHandler(
 		opts,
 		db,
+		() => ensureFinishToolCalled(finishObserved, toolset, sharedCtx, stepIndex),
 		completeAssistantMessage,
-		getAccumulated,
-		abortController,
 	);
 
-	const _onAbort = createAbortHandler(opts, db, abortController);
-	const onError = createErrorHandler(opts, db);
+	// Apply optimizations: deduplication, pruning, cache control, and truncation
+	const { addCacheControl, truncateHistory } = await import(
+		'./cache-optimizer.ts'
+	);
+	const { optimizeContext } = await import('./context-optimizer.ts');
 
-	// Context optimization
-	const contextOptimized = optimizeContext(history, {
+	// 1. Optimize context (deduplicate file reads, prune old tool results)
+	const contextOptimized = optimizeContext(messagesWithSystemInstructions, {
 		deduplicateFiles: true,
 		maxToolResults: 30,
 	});
 
-	// Truncate history
+	// 2. Truncate history
 	const truncatedMessages = truncateHistory(contextOptimized, 20);
 
-	// Add cache control
+	// 3. Add cache control
 	const { system: cachedSystem, messages: optimizedMessages } = addCacheControl(
-		opts.provider,
+		opts.provider as any,
 		system,
 		truncatedMessages,
 	);
 
 	try {
-		const maxTokens = getMaxOutputTokens(provider, modelName);
-		const result = await streamText({
+		// @ts-expect-error this is fine 🔥
+		const result = streamText({
 			model,
-			system: cachedSystem,
+			tools: toolset,
+			...(cachedSystem ? { system: cachedSystem } : {}),
 			messages: optimizedMessages,
-			tools,
-			maxSteps: 50,
-			maxTokens,
-			temperature: agentConfig.temperature ?? 0.7,
-			abortSignal: abortController.signal,
+			...(maxOutputTokens ? { maxOutputTokens } : {}),
+			abortSignal: opts.abortSignal,
+			stopWhen: hasToolCall('finish'),
 			onStepFinish,
+			onError,
+			onAbort,
 			onFinish,
-			experimental_continueSteps: true,
 		});
 
-		// Process the stream
 		for await (const delta of result.textStream) {
-			if (abortController.signal.aborted) break;
-
-			accumulated += delta;
-			if (currentPartId) {
-				await db
-					.update(messageParts)
-					.set({ content: accumulated })
-					.where(eq(messageParts.id, currentPartId));
+			if (!delta) continue;
+			if (!firstDeltaSeen) {
+				firstDeltaSeen = true;
+				streamStartTimer.end();
 			}
-
-			publish('stream:text-delta', {
+			accumulated += delta;
+			publish({
+				type: 'message.part.delta',
 				sessionId: opts.sessionId,
-				messageId: opts.assistantMessageId,
-				assistantMessageId: opts.assistantMessageId,
-				stepIndex,
-				textDelta: delta,
-				fullText: accumulated,
+				payload: {
+					messageId: opts.assistantMessageId,
+					partId: currentPartId,
+					stepIndex,
+					delta,
+				},
 			});
+			await db
+				.update(messageParts)
+				.set({ content: JSON.stringify({ text: accumulated }) })
+				.where(eq(messageParts.id, currentPartId));
 		}
-	} catch (err) {
-		await onError(err);
+	} catch (error) {
+		const errorPayload = toErrorPayload(error);
+		await db
+			.update(messageParts)
+			.set({
+				content: JSON.stringify({
+					text: accumulated,
+					error: errorPayload.message,
+				}),
+			})
+			.where(eq(messageParts.messageId, opts.assistantMessageId));
+		publish({
+			type: 'error',
+			sessionId: opts.sessionId,
+			payload: {
+				messageId: opts.assistantMessageId,
+				error: errorPayload.message,
+				details: errorPayload.details,
+			},
+		});
+		throw error;
 	} finally {
-		setRunning(opts.sessionId, false);
-		dequeueJob(opts.sessionId);
+		if (!firstToolSeen()) firstToolTimer.end({ skipped: true });
+		try {
+			unsubscribeFinish();
+		} catch {}
+		try {
+			await cleanupEmptyTextParts(opts, db);
+		} catch {}
 	}
 }
-
-/**
- * Enqueues an assistant run
- */
-export async function enqueueAssistantRun(opts: RunOpts) {
-	return enqueueRun(opts);
-}
-
-/**
- * Aborts a running session
- */
-export async function abortSession(sessionId: number) {
-	return abortSessionQueue(sessionId);
-}
-
-/**
- * Gets the current runner state for a session
- */
-export function getSessionState(sessionId: number) {
-	return getRunnerState(sessionId);
-}
-
-/**
- * Cleanup session resources
- */
-export function cleanupSessionResources(sessionId: number) {
-	return cleanupSession(sessionId);
-}

package/src/runtime/stream-handlers.ts

@@ -8,26 +8,17 @@ import { toErrorPayload } from './error-handling.ts';
 import type { RunOpts } from './session-queue.ts';
 import type { ToolAdapterContext } from '../tools/adapter.ts';
 
-interface ProviderMetadata {
-	openai?: {
-		cachedPromptTokens?: number;
-	};
-	[key: string]: unknown;
-}
-
-interface UsageData {
-	inputTokens?: number;
-	outputTokens?: number;
-	totalTokens?: number;
-	cachedInputTokens?: number;
-	reasoningTokens?: number;
-}
-
 type StepFinishEvent = {
-	usage?: UsageData;
+	usage?: {
+		inputTokens?: number;
+		outputTokens?: number;
+		totalTokens?: number;
+		cachedInputTokens?: number;
+		reasoningTokens?: number;
+	};
 	finishReason?: string;
 	response?: unknown;
-	experimental_providerMetadata?: ProviderMetadata;
+	experimental_providerMetadata?: Record<string, any>;
 };
 
 type FinishEvent = {
@@ -51,19 +42,19 @@ export function createStepFinishHandler(
 	db: Awaited<ReturnType<typeof getDb>>,
 	getCurrentPartId: () => string,
 	getStepIndex: () => number,
-	_sharedCtx: ToolAdapterContext,
-	_updateCurrentPartId: (id: string) => void,
-	_updateAccumulated: (text: string) => void,
+	sharedCtx: ToolAdapterContext,
+	updateCurrentPartId: (id: string) => void,
+	updateAccumulated: (text: string) => void,
 	incrementStepIndex: () => number,
 	updateSessionTokensIncrementalFn: (
-		usage: UsageData,
-		providerMetadata: ProviderMetadata | undefined,
+		usage: any,
+		providerMetadata: Record<string, any> | undefined,
 		opts: RunOpts,
 		db: Awaited<ReturnType<typeof getDb>>,
 	) => Promise<void>,
 	updateMessageTokensIncrementalFn: (
-		usage: UsageData,
-		providerMetadata: ProviderMetadata | undefined,
+		usage: any,
+		providerMetadata: Record<string, any> | undefined,
 		opts: RunOpts,
 		db: Awaited<ReturnType<typeof getDb>>,
 	) => Promise<void>,
@@ -78,11 +69,9 @@ export function createStepFinishHandler(
 				.update(messageParts)
 				.set({ completedAt: finishedAt })
 				.where(eq(messageParts.id, currentPartId));
-		} catch (err) {
-			console.error('[createStepFinishHandler] Failed to update part', err);
-		}
+		} catch {}
 
-		// Update tokens incrementally
+		// Update token counts incrementally after each step
 		if (step.usage) {
 			try {
 				await updateSessionTokensIncrementalFn(
@@ -91,81 +80,126 @@ export function createStepFinishHandler(
 					opts,
 					db,
 				);
+			} catch {}
+
+			try {
 				await updateMessageTokensIncrementalFn(
 					step.usage,
 					step.experimental_providerMetadata,
 					opts,
 					db,
 				);
-			} catch (err) {
-				console.error('[createStepFinishHandler] Token update failed', err);
-			}
+			} catch {}
 		}
 
-		// Publish step-finished event
-		publish('stream:step-finished', {
-			sessionId: opts.sessionId,
-			messageId: opts.assistantMessageId,
-			assistantMessageId: opts.assistantMessageId,
-			stepIndex,
-			finishReason: step.finishReason,
-			usage: step.usage,
-		});
+		try {
+			publish({
+				type: 'finish-step',
+				sessionId: opts.sessionId,
+				payload: {
+					stepIndex,
+					usage: step.usage,
+					finishReason: step.finishReason,
+					response: step.response,
+				},
+			});
+			if (step.usage) {
+				publish({
+					type: 'usage',
+					sessionId: opts.sessionId,
+					payload: { stepIndex, ...step.usage },
+				});
+			}
+		} catch {}
 
-		incrementStepIndex();
+		try {
+			const newStepIndex = incrementStepIndex();
+			const newPartId = crypto.randomUUID();
+			const index = await sharedCtx.nextIndex();
+			const nowTs = Date.now();
+			await db.insert(messageParts).values({
+				id: newPartId,
+				messageId: opts.assistantMessageId,
+				index,
+				stepIndex: newStepIndex,
+				type: 'text',
+				content: JSON.stringify({ text: '' }),
+				agent: opts.agent,
+				provider: opts.provider,
+				model: opts.model,
+				startedAt: nowTs,
+			});
+			updateCurrentPartId(newPartId);
+			sharedCtx.assistantPartId = newPartId;
+			sharedCtx.stepIndex = newStepIndex;
+			updateAccumulated('');
+		} catch {}
 	};
 }
 
 /**
- * Creates the onFinish handler for the stream
+ * Creates the onError handler for the stream
  */
-export function createFinishHandler(
+export function createErrorHandler(
 	opts: RunOpts,
 	db: Awaited<ReturnType<typeof getDb>>,
-	completeAssistantMessageFn: (
-		fin: FinishEvent,
-		opts: RunOpts,
-		db: Awaited<ReturnType<typeof getDb>>,
-	) => Promise<void>,
-	_getAccumulated: () => string,
-	_abortController: AbortController,
+	getStepIndex: () => number,
+	sharedCtx: ToolAdapterContext,
 ) {
-	return async (fin: FinishEvent) => {
-		try {
-			await completeAssistantMessageFn(fin, opts, db);
+	return async (err: unknown) => {
+		const errorPayload = toErrorPayload(err);
+		const isApiError = APICallError.isInstance(err);
+		const stepIndex = getStepIndex();
 
-			const msgRows = await db
-				.select()
-				.from(messages)
-				.where(eq(messages.id, opts.assistantMessageId));
+		// Create error part for UI display
+		const errorPartId = crypto.randomUUID();
+		await db.insert(messageParts).values({
+			id: errorPartId,
+			messageId: opts.assistantMessageId,
+			index: await sharedCtx.nextIndex(),
+			stepIndex,
+			type: 'error',
+			content: JSON.stringify({
+				message: errorPayload.message,
+				type: errorPayload.type,
+				details: errorPayload.details,
+				isAborted: false,
+			}),
+			agent: opts.agent,
+			provider: opts.provider,
+			model: opts.model,
+			startedAt: Date.now(),
+			completedAt: Date.now(),
+		});
 
-			let estimatedCost = 0;
-			if (msgRows.length > 0 && msgRows[0]) {
-				const msg = msgRows[0];
-				estimatedCost = estimateModelCostUsd(
-					opts.provider,
-					opts.model,
-					Number(msg.promptTokens ?? 0),
-					Number(msg.completionTokens ?? 0),
-				);
-			}
+		// Update message status
+		await db
+			.update(messages)
+			.set({
+				status: 'error',
+				error: errorPayload.message,
+				errorType: errorPayload.type,
+				errorDetails: JSON.stringify({
+					...errorPayload.details,
+					isApiError,
+				}),
+				isAborted: false,
+			})
+			.where(eq(messages.id, opts.assistantMessageId));
 
-			publish('stream:finished', {
-				sessionId: opts.sessionId,
-				messageId: opts.assistantMessageId,
-				assistantMessageId: opts.assistantMessageId,
-				usage: fin.usage,
-				finishReason: fin.finishReason,
-				estimatedCost,
-			});
-		} catch (err) {
-			console.error('[createFinishHandler] Error in onFinish', err);
-			publish('stream:error', {
-				sessionId: opts.sessionId,
+		// Publish enhanced error event
+		publish({
+			type: 'error',
+			sessionId: opts.sessionId,
+			payload: {
 				messageId: opts.assistantMessageId,
-				error: toErrorPayload(err),
-			});
-		}
+				partId: errorPartId,
+				error: errorPayload.message,
+				errorType: errorPayload.type,
+				details: errorPayload.details,
+				isAborted: false,
+			},
+		});
 	};
 }
 
@@ -175,116 +209,116 @@ export function createFinishHandler(
 export function createAbortHandler(
 	opts: RunOpts,
 	db: Awaited<ReturnType<typeof getDb>>,
-	_abortController: AbortController,
+	getStepIndex: () => number,
+	sharedCtx: ToolAdapterContext,
 ) {
-	return async (_event: AbortEvent) => {
-		try {
-			await db
-				.update(messages)
-				.set({ status: 'aborted', finishedAt: new Date() })
-				.where(eq(messages.id, opts.assistantMessageId));
+	return async ({ steps }: AbortEvent) => {
+		const stepIndex = getStepIndex();
 
-			publish('stream:aborted', {
-				sessionId: opts.sessionId,
+		// Create abort part for UI
+		const abortPartId = crypto.randomUUID();
+		await db.insert(messageParts).values({
+			id: abortPartId,
+			messageId: opts.assistantMessageId,
+			index: await sharedCtx.nextIndex(),
+			stepIndex,
+			type: 'error',
+			content: JSON.stringify({
+				message: 'Generation stopped by user',
+				type: 'abort',
+				isAborted: true,
+				stepsCompleted: steps.length,
+			}),
+			agent: opts.agent,
+			provider: opts.provider,
+			model: opts.model,
+			startedAt: Date.now(),
+			completedAt: Date.now(),
+		});
+
+		// Store abort info
+		await db
+			.update(messages)
+			.set({
+				status: 'error',
+				error: 'Generation stopped by user',
+				errorType: 'abort',
+				errorDetails: JSON.stringify({
+					stepsCompleted: steps.length,
+					abortedAt: Date.now(),
+				}),
+				isAborted: true,
+			})
+			.where(eq(messages.id, opts.assistantMessageId));
+
+		// Publish abort event
+		publish({
+			type: 'error',
+			sessionId: opts.sessionId,
+			payload: {
 				messageId: opts.assistantMessageId,
-				assistantMessageId: opts.assistantMessageId,
-			});
-		} catch (err) {
-			console.error('[createAbortHandler] Error in onAbort', err);
-		}
+				partId: abortPartId,
+				error: 'Generation stopped by user',
+				errorType: 'abort',
+				isAborted: true,
+				stepsCompleted: steps.length,
+			},
+		});
 	};
 }
 
 /**
- * Creates the error handler for the stream
+ * Creates the onFinish handler for the stream
  */
-export function createErrorHandler(
+export function createFinishHandler(
 	opts: RunOpts,
 	db: Awaited<ReturnType<typeof getDb>>,
+	ensureFinishToolCalled: () => Promise<void>,
+	completeAssistantMessageFn: (
+		fin: FinishEvent,
+		opts: RunOpts,
+		db: Awaited<ReturnType<typeof getDb>>,
+	) => Promise<void>,
 ) {
-	return async (err: unknown) => {
-		console.error('[createErrorHandler] Stream error:', err);
-
+	return async (fin: FinishEvent) => {
 		try {
-			let errorMessage = 'Unknown error';
-			let errorType = 'UNKNOWN_ERROR';
-			let errorStack: string | undefined;
+			await ensureFinishToolCalled();
+		} catch {}
 
-			if (err instanceof APICallError) {
-				errorMessage = err.message;
-				errorType = 'API_CALL_ERROR';
-				errorStack = err.stack;
-			} else if (err instanceof Error) {
-				errorMessage = err.message;
-				errorType = err.name || 'ERROR';
-				errorStack = err.stack;
-			} else if (typeof err === 'string') {
-				errorMessage = err;
-			}
+		// Note: Token updates are handled incrementally in onStepFinish
+		// Do NOT add fin.usage here as it would cause double-counting
 
-			await db
-				.update(messages)
-				.set({
-					status: 'error',
-					finishedAt: new Date(),
-					error: errorMessage,
-				})
-				.where(eq(messages.id, opts.assistantMessageId));
-
-			publish('stream:error', {
-				sessionId: opts.sessionId,
-				messageId: opts.assistantMessageId,
-				assistantMessageId: opts.assistantMessageId,
-				error: {
-					message: errorMessage,
-					type: errorType,
-					stack: errorStack,
-				},
-			});
-		} catch (dbErr) {
-			console.error('[createErrorHandler] Failed to save error to DB', dbErr);
-		}
-	};
-}
+		try {
+			await completeAssistantMessageFn(fin, opts, db);
+		} catch {}
 
-/**
- * Creates the text delta handler for the stream
- */
-export function createTextHandler(
-	opts: RunOpts,
-	db: Awaited<ReturnType<typeof getDb>>,
-	getCurrentPartId: () => string,
-	getStepIndex: () => number,
-	_updateCurrentPartId: (id: string) => void,
-	updateAccumulated: (text: string) => void,
-	getAccumulated: () => string,
-) {
-	return async (textDelta: string) => {
-		const currentPartId = getCurrentPartId();
-		const stepIndex = getStepIndex();
+		// Use session totals from DB for accurate cost calculation
+		const sessRows = await db
+			.select()
+			.from(messages)
+			.where(eq(messages.id, opts.assistantMessageId));
 
-		// Accumulate the text
-		const accumulated = getAccumulated() + textDelta;
-		updateAccumulated(accumulated);
+		const usage = sessRows[0]
+			? {
+					inputTokens: Number(sessRows[0].promptTokens ?? 0),
+					outputTokens: Number(sessRows[0].completionTokens ?? 0),
+					totalTokens: Number(sessRows[0].totalTokens ?? 0),
+				}
+			: fin.usage;
 
-		try {
-			if (currentPartId) {
-				await db
-					.update(messageParts)
-					.set({ content: accumulated })
-					.where(eq(messageParts.id, currentPartId));
-			}
+		const costUsd = usage
+			? estimateModelCostUsd(opts.provider, opts.model, usage)
+			: undefined;
 
-			publish('stream:text-delta', {
-				sessionId: opts.sessionId,
-				messageId: opts.assistantMessageId,
-				assistantMessageId: opts.assistantMessageId,
-				stepIndex,
-				textDelta,
-				fullText: accumulated,
-			});
-		} catch (err) {
-			console.error('[createTextHandler] Error updating text part', err);
-		}
+		publish({
+			type: 'message.completed',
+			sessionId: opts.sessionId,
+			payload: {
+				id: opts.assistantMessageId,
+				usage,
+				costUsd,
+				finishReason: fin.finishReason,
+			},
+		});
 	};
 }