@agi-cli/server 0.1.61 → 0.1.63

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,
+			},
+		});
 	};
 }