@agi-cli/server 0.1.55

package/src/runtime/stream-handlers.ts

@@ -0,0 +1,275 @@
+import type { getDb } from '@agi-cli/database';
+import { messages, messageParts } from '@agi-cli/database/schema';
+import { eq } from 'drizzle-orm';
+import { APICallError } from 'ai';
+import { publish } from '../events/bus.ts';
+import { estimateModelCostUsd } from '@agi-cli/sdk';
+import { toErrorPayload } from './error-handling.ts';
+import type { RunOpts } from './session-queue.ts';
+import type { ToolAdapterContext } from '../tools/adapter.ts';
+
+type StepFinishEvent = {
+	usage?: { inputTokens?: number; outputTokens?: number };
+	finishReason?: string;
+	response?: unknown;
+};
+
+type FinishEvent = {
+	usage?: {
+		inputTokens?: number;
+		outputTokens?: number;
+		totalTokens?: number;
+	};
+	finishReason?: string;
+};
+
+type AbortEvent = {
+	steps: unknown[];
+};
+
+/**
+ * Creates the onStepFinish handler for the stream
+ */
+export function createStepFinishHandler(
+	opts: RunOpts,
+	db: Awaited<ReturnType<typeof getDb>>,
+	getCurrentPartId: () => string,
+	getStepIndex: () => number,
+	sharedCtx: ToolAdapterContext,
+	updateCurrentPartId: (id: string) => void,
+	updateAccumulated: (text: string) => void,
+	incrementStepIndex: () => number,
+) {
+	return async (step: StepFinishEvent) => {
+		const finishedAt = Date.now();
+		const currentPartId = getCurrentPartId();
+		const stepIndex = getStepIndex();
+
+		try {
+			await db
+				.update(messageParts)
+				.set({ completedAt: finishedAt })
+				.where(eq(messageParts.id, currentPartId));
+		} catch {}
+
+		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 {}
+
+		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 onError handler for the stream
+ */
+export function createErrorHandler(
+	opts: RunOpts,
+	db: Awaited<ReturnType<typeof getDb>>,
+	getStepIndex: () => number,
+	sharedCtx: ToolAdapterContext,
+) {
+	return async (err: unknown) => {
+		const errorPayload = toErrorPayload(err);
+		const isApiError = APICallError.isInstance(err);
+		const stepIndex = getStepIndex();
+
+		// 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(),
+		});
+
+		// 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 enhanced error event
+		publish({
+			type: 'error',
+			sessionId: opts.sessionId,
+			payload: {
+				messageId: opts.assistantMessageId,
+				partId: errorPartId,
+				error: errorPayload.message,
+				errorType: errorPayload.type,
+				details: errorPayload.details,
+				isAborted: false,
+			},
+		});
+	};
+}
+
+/**
+ * Creates the onAbort handler for the stream
+ */
+export function createAbortHandler(
+	opts: RunOpts,
+	db: Awaited<ReturnType<typeof getDb>>,
+	getStepIndex: () => number,
+	sharedCtx: ToolAdapterContext,
+) {
+	return async ({ steps }: AbortEvent) => {
+		const stepIndex = getStepIndex();
+
+		// 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,
+				partId: abortPartId,
+				error: 'Generation stopped by user',
+				errorType: 'abort',
+				isAborted: true,
+				stepsCompleted: steps.length,
+			},
+		});
+	};
+}
+
+/**
+ * Creates the onFinish handler for the stream
+ */
+export function createFinishHandler(
+	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>,
+) {
+	return async (fin: FinishEvent) => {
+		try {
+			await ensureFinishToolCalled();
+		} catch {}
+
+		try {
+			await updateSessionTokensFn(fin, opts, db);
+		} catch {}
+
+		try {
+			await completeAssistantMessageFn(fin, opts, db);
+		} catch {}
+
+		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,
+			},
+		});
+	};
+}

package/src/runtime/token-utils.ts

@@ -0,0 +1,35 @@
+import { catalog } from '@agi-cli/sdk';
+import { debugLog } from './debug.ts';
+import type { ProviderName } from './provider.ts';
+
+/**
+ * Gets the maximum output tokens allowed for a given provider/model combination.
+ * Returns undefined if the information is not available in the catalog.
+ */
+export function getMaxOutputTokens(
+	provider: ProviderName,
+	modelId: string,
+): number | undefined {
+	try {
+		const providerCatalog = catalog[provider];
+		if (!providerCatalog) {
+			debugLog(`[maxOutputTokens] No catalog found for provider: ${provider}`);
+			return undefined;
+		}
+		const modelInfo = providerCatalog.models.find((m) => m.id === modelId);
+		if (!modelInfo) {
+			debugLog(
+				`[maxOutputTokens] No model info found for: ${modelId} in provider: ${provider}`,
+			);
+			return undefined;
+		}
+		const outputLimit = modelInfo.limit?.output;
+		debugLog(
+			`[maxOutputTokens] Provider: ${provider}, Model: ${modelId}, Limit: ${outputLimit}`,
+		);
+		return outputLimit;
+	} catch (err) {
+		debugLog(`[maxOutputTokens] Error looking up limit: ${err}`);
+		return undefined;
+	}
+}

package/src/runtime/tool-context-setup.ts

@@ -0,0 +1,58 @@
+import type { getDb } from '@agi-cli/database';
+import { messageParts } from '@agi-cli/database/schema';
+import { eq } from 'drizzle-orm';
+import { time } from './debug.ts';
+import type { ToolAdapterContext } from '../tools/adapter.ts';
+import type { RunOpts } from './session-queue.ts';
+
+export type RunnerToolContext = ToolAdapterContext & { stepIndex: number };
+
+/**
+ * Sets up the shared tool context for a run, including the index counter
+ * and first tool call tracking.
+ */
+export async function setupToolContext(
+	opts: RunOpts,
+	db: Awaited<ReturnType<typeof getDb>>,
+) {
+	const firstToolTimer = time('runner:first-tool-call');
+	let firstToolSeen = false;
+
+	const sharedCtx: RunnerToolContext = {
+		nextIndex: async () => 0,
+		stepIndex: 0,
+		sessionId: opts.sessionId,
+		messageId: opts.assistantMessageId,
+		assistantPartId: opts.assistantPartId,
+		db,
+		agent: opts.agent,
+		provider: opts.provider,
+		model: opts.model,
+		projectRoot: opts.projectRoot,
+		onFirstToolCall: () => {
+			if (firstToolSeen) return;
+			firstToolSeen = true;
+			firstToolTimer.end();
+		},
+	};
+
+	let counter = 0;
+	try {
+		const existing = await db
+			.select()
+			.from(messageParts)
+			.where(eq(messageParts.messageId, opts.assistantMessageId));
+		if (existing.length) {
+			const indexes = existing.map((p) => Number(p.index ?? 0));
+			const maxIndex = Math.max(...indexes);
+			if (Number.isFinite(maxIndex)) counter = maxIndex;
+		}
+	} catch {}
+
+	sharedCtx.nextIndex = () => {
+		counter += 1;
+		return counter;
+	};
+
+	return { sharedCtx, firstToolTimer, firstToolSeen: () => firstToolSeen };
+}

package/src/runtime/tool-context.ts

@@ -0,0 +1,72 @@
+import { eq } from 'drizzle-orm';
+import type { DB } from '@agi-cli/database';
+import { messageParts } from '@agi-cli/database/schema';
+import { publish } from '../events/bus.ts';
+
+export type ToolAdapterContext = {
+	sessionId: string;
+	messageId: string;
+	assistantPartId: string;
+	db: DB;
+	agent: string;
+	provider: string;
+	model: string;
+	projectRoot: string;
+	nextIndex: () => number | Promise<number>;
+	stepIndex?: number;
+	onFirstToolCall?: () => void;
+};
+
+export function extractFinishText(input: unknown): string | undefined {
+	if (typeof input === 'string') return input;
+	if (!input || typeof input !== 'object') return undefined;
+	const obj = input as Record<string, unknown>;
+	if (typeof obj.text === 'string') return obj.text;
+	if (
+		obj.input &&
+		typeof (obj.input as Record<string, unknown>).text === 'string'
+	)
+		return String((obj.input as Record<string, unknown>).text);
+	return undefined;
+}
+
+export async function appendAssistantText(
+	ctx: ToolAdapterContext,
+	text: string,
+): Promise<void> {
+	try {
+		const rows = await ctx.db
+			.select()
+			.from(messageParts)
+			.where(eq(messageParts.id, ctx.assistantPartId));
+		let previous = '';
+		if (rows.length) {
+			try {
+				const parsed = JSON.parse(rows[0]?.content ?? '{}');
+				if (parsed && typeof parsed.text === 'string') previous = parsed.text;
+			} catch {}
+		}
+		const addition = text.startsWith(previous)
+			? text.slice(previous.length)
+			: text;
+		if (addition.length) {
+			const payload: Record<string, unknown> = {
+				messageId: ctx.messageId,
+				partId: ctx.assistantPartId,
+				delta: addition,
+			};
+			if (ctx.stepIndex !== undefined) payload.stepIndex = ctx.stepIndex;
+			publish({
+				type: 'message.part.delta',
+				sessionId: ctx.sessionId,
+				payload,
+			});
+		}
+		await ctx.db
+			.update(messageParts)
+			.set({ content: JSON.stringify({ text }) })
+			.where(eq(messageParts.id, ctx.assistantPartId));
+	} catch {
+		// ignore to keep run alive if we can't persist the text
+	}
+}