@agi-cli/server 0.1.58 → 0.1.61

package/src/runtime/stream-handlers.ts

@@ -8,10 +8,26 @@ 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?: { inputTokens?: number; outputTokens?: number };
+	usage?: UsageData;
 	finishReason?: string;
 	response?: unknown;
+	experimental_providerMetadata?: ProviderMetadata;
 };
 
 type FinishEvent = {
@@ -35,10 +51,22 @@ 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,
+		opts: RunOpts,
+		db: Awaited<ReturnType<typeof getDb>>,
+	) => Promise<void>,
+	updateMessageTokensIncrementalFn: (
+		usage: UsageData,
+		providerMetadata: ProviderMetadata | undefined,
+		opts: RunOpts,
+		db: Awaited<ReturnType<typeof getDb>>,
+	) => Promise<void>,
 ) {
 	return async (step: StepFinishEvent) => {
 		const finishedAt = Date.now();
@@ -50,116 +78,94 @@ export function createStepFinishHandler(
 				.update(messageParts)
 				.set({ completedAt: finishedAt })
 				.where(eq(messageParts.id, currentPartId));
-		} catch {}
+		} catch (err) {
+			console.error('[createStepFinishHandler] Failed to update part', err);
+		}
 
-		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 },
-				});
+		// Update tokens incrementally
+		if (step.usage) {
+			try {
+				await updateSessionTokensIncrementalFn(
+					step.usage,
+					step.experimental_providerMetadata,
+					opts,
+					db,
+				);
+				await updateMessageTokensIncrementalFn(
+					step.usage,
+					step.experimental_providerMetadata,
+					opts,
+					db,
+				);
+			} catch (err) {
+				console.error('[createStepFinishHandler] Token update failed', err);
 			}
-		} catch {}
+		}
 
-		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 {}
+		// Publish step-finished event
+		publish('stream:step-finished', {
+			sessionId: opts.sessionId,
+			messageId: opts.assistantMessageId,
+			assistantMessageId: opts.assistantMessageId,
+			stepIndex,
+			finishReason: step.finishReason,
+			usage: step.usage,
+		});
+
+		incrementStepIndex();
 	};
 }
 
 /**
- * Creates the onError handler for the stream
+ * Creates the onFinish handler for the stream
  */
-export function createErrorHandler(
+export function createFinishHandler(
 	opts: RunOpts,
 	db: Awaited<ReturnType<typeof getDb>>,
-	getStepIndex: () => number,
-	sharedCtx: ToolAdapterContext,
+	completeAssistantMessageFn: (
+		fin: FinishEvent,
+		opts: RunOpts,
+		db: Awaited<ReturnType<typeof getDb>>,
+	) => Promise<void>,
+	_getAccumulated: () => string,
+	_abortController: AbortController,
 ) {
-	return async (err: unknown) => {
-		const errorPayload = toErrorPayload(err);
-		const isApiError = APICallError.isInstance(err);
-		const stepIndex = getStepIndex();
+	return async (fin: FinishEvent) => {
+		try {
+			await completeAssistantMessageFn(fin, opts, db);
 
-		// 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(),
-		});
+			const msgRows = await db
+				.select()
+				.from(messages)
+				.where(eq(messages.id, opts.assistantMessageId));
 
-		// 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));
+			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),
+				);
+			}
 
-		// Publish enhanced error event
-		publish({
-			type: 'error',
-			sessionId: opts.sessionId,
-			payload: {
+			publish('stream:finished', {
+				sessionId: opts.sessionId,
 				messageId: opts.assistantMessageId,
-				partId: errorPartId,
-				error: errorPayload.message,
-				errorType: errorPayload.type,
-				details: errorPayload.details,
-				isAborted: false,
-			},
-		});
+				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,
+				messageId: opts.assistantMessageId,
+				error: toErrorPayload(err),
+			});
+		}
 	};
 }
 
@@ -169,107 +175,116 @@ export function createErrorHandler(
 export function createAbortHandler(
 	opts: RunOpts,
 	db: Awaited<ReturnType<typeof getDb>>,
-	getStepIndex: () => number,
-	sharedCtx: ToolAdapterContext,
+	_abortController: AbortController,
 ) {
-	return async ({ steps }: AbortEvent) => {
-		const stepIndex = getStepIndex();
+	return async (_event: AbortEvent) => {
+		try {
+			await db
+				.update(messages)
+				.set({ status: 'aborted', finishedAt: new Date() })
+				.where(eq(messages.id, opts.assistantMessageId));
 
-		// 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(),
-		});
+			publish('stream:aborted', {
+				sessionId: opts.sessionId,
+				messageId: opts.assistantMessageId,
+				assistantMessageId: opts.assistantMessageId,
+			});
+		} catch (err) {
+			console.error('[createAbortHandler] Error in onAbort', err);
+		}
+	};
+}
 
-		// 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));
+/**
+ * Creates the error handler for the stream
+ */
+export function createErrorHandler(
+	opts: RunOpts,
+	db: Awaited<ReturnType<typeof getDb>>,
+) {
+	return async (err: unknown) => {
+		console.error('[createErrorHandler] Stream error:', err);
 
-		// Publish abort event
-		publish({
-			type: 'error',
-			sessionId: opts.sessionId,
-			payload: {
+		try {
+			let errorMessage = 'Unknown error';
+			let errorType = 'UNKNOWN_ERROR';
+			let errorStack: string | undefined;
+
+			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;
+			}
+
+			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,
-				partId: abortPartId,
-				error: 'Generation stopped by user',
-				errorType: 'abort',
-				isAborted: true,
-				stepsCompleted: steps.length,
-			},
-		});
+				assistantMessageId: opts.assistantMessageId,
+				error: {
+					message: errorMessage,
+					type: errorType,
+					stack: errorStack,
+				},
+			});
+		} catch (dbErr) {
+			console.error('[createErrorHandler] Failed to save error to DB', dbErr);
+		}
 	};
 }
 
 /**
- * Creates the onFinish handler for the stream
+ * Creates the text delta handler for the stream
  */
-export function createFinishHandler(
+export function createTextHandler(
 	opts: RunOpts,
 	db: Awaited<ReturnType<typeof getDb>>,
-	ensureFinishToolCalled: () => Promise<void>,
-	updateSessionTokensFn: (
-		fin: FinishEvent,
-		opts: RunOpts,
-		db: Awaited<ReturnType<typeof getDb>>,
-	) => Promise<void>,
-	completeAssistantMessageFn: (
-		fin: FinishEvent,
-		opts: RunOpts,
-		db: Awaited<ReturnType<typeof getDb>>,
-	) => Promise<void>,
+	getCurrentPartId: () => string,
+	getStepIndex: () => number,
+	_updateCurrentPartId: (id: string) => void,
+	updateAccumulated: (text: string) => void,
+	getAccumulated: () => string,
 ) {
-	return async (fin: FinishEvent) => {
-		try {
-			await ensureFinishToolCalled();
-		} catch {}
+	return async (textDelta: string) => {
+		const currentPartId = getCurrentPartId();
+		const stepIndex = getStepIndex();
 
-		try {
-			await updateSessionTokensFn(fin, opts, db);
-		} catch {}
+		// Accumulate the text
+		const accumulated = getAccumulated() + textDelta;
+		updateAccumulated(accumulated);
 
 		try {
-			await completeAssistantMessageFn(fin, opts, db);
-		} catch {}
+			if (currentPartId) {
+				await db
+					.update(messageParts)
+					.set({ content: accumulated })
+					.where(eq(messageParts.id, currentPartId));
+			}
 
-		const costUsd = fin.usage
-			? estimateModelCostUsd(opts.provider, opts.model, fin.usage)
-			: undefined;
-		publish({
-			type: 'message.completed',
-			sessionId: opts.sessionId,
-			payload: {
-				id: opts.assistantMessageId,
-				usage: fin.usage,
-				costUsd,
-				finishReason: fin.finishReason,
-			},
-		});
+			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);
+		}
 	};
 }