@agi-cli/server 0.1.119 → 0.1.121

package/src/runtime/session/branch.ts

@@ -0,0 +1,277 @@
+import { eq, asc } from 'drizzle-orm';
+import type { DB } from '@agi-cli/database';
+import { sessions, messages, messageParts } from '@agi-cli/database/schema';
+import { publish } from '../../events/bus.ts';
+import type { ProviderId } from '@agi-cli/sdk';
+
+type SessionRow = typeof sessions.$inferSelect;
+
+export type CreateBranchInput = {
+	db: DB;
+	parentSessionId: string;
+	fromMessageId: string;
+	provider?: ProviderId;
+	model?: string;
+	agent?: string;
+	title?: string;
+	projectPath: string;
+};
+
+export type BranchResult = {
+	session: SessionRow;
+	parentSessionId: string;
+	branchPointMessageId: string;
+	copiedMessages: number;
+	copiedParts: number;
+};
+
+export async function createBranch({
+	db,
+	parentSessionId,
+	fromMessageId,
+	provider,
+	model,
+	agent,
+	title,
+	projectPath,
+}: CreateBranchInput): Promise<BranchResult> {
+	const parentRows = await db
+		.select()
+		.from(sessions)
+		.where(eq(sessions.id, parentSessionId));
+
+	if (!parentRows.length) {
+		throw new Error('Parent session not found');
+	}
+
+	const parent = parentRows[0];
+
+	if (parent.projectPath !== projectPath) {
+		throw new Error('Parent session not found in this project');
+	}
+
+	const branchPointRows = await db
+		.select()
+		.from(messages)
+		.where(eq(messages.id, fromMessageId));
+
+	if (!branchPointRows.length) {
+		throw new Error('Branch point message not found');
+	}
+
+	const branchPoint = branchPointRows[0];
+
+	if (branchPoint.sessionId !== parentSessionId) {
+		throw new Error('Branch point message does not belong to parent session');
+	}
+
+	const allMessages = await db
+		.select()
+		.from(messages)
+		.where(eq(messages.sessionId, parentSessionId))
+		.orderBy(asc(messages.createdAt));
+
+	const branchPointIndex = allMessages.findIndex((m) => m.id === fromMessageId);
+	if (branchPointIndex === -1) {
+		throw new Error('Branch point message not found in session');
+	}
+
+	const messagesToCopy = allMessages.slice(0, branchPointIndex + 1);
+
+	const newSessionId = crypto.randomUUID();
+	const now = Date.now();
+
+	const newSession: typeof sessions.$inferInsert = {
+		id: newSessionId,
+		title: title || `Branch of ${parent.title || 'Untitled'}`,
+		agent: agent || parent.agent,
+		provider: provider || parent.provider,
+		model: model || parent.model,
+		projectPath: parent.projectPath,
+		createdAt: now,
+		lastActiveAt: now,
+		parentSessionId,
+		branchPointMessageId: fromMessageId,
+		sessionType: 'branch',
+	};
+
+	await db.insert(sessions).values(newSession);
+
+	const messageIdMap = new Map<string, string>();
+	let copiedParts = 0;
+
+	for (const msg of messagesToCopy) {
+		const newMessageId = crypto.randomUUID();
+		messageIdMap.set(msg.id, newMessageId);
+
+		const newMessage: typeof messages.$inferInsert = {
+			id: newMessageId,
+			sessionId: newSessionId,
+			role: msg.role,
+			status: msg.status,
+			agent: msg.agent,
+			provider: msg.provider,
+			model: msg.model,
+			createdAt: msg.createdAt,
+			completedAt: msg.completedAt,
+			latencyMs: msg.latencyMs,
+			promptTokens: msg.promptTokens,
+			completionTokens: msg.completionTokens,
+			totalTokens: msg.totalTokens,
+			cachedInputTokens: msg.cachedInputTokens,
+			reasoningTokens: msg.reasoningTokens,
+			error: msg.error,
+			errorType: msg.errorType,
+			errorDetails: msg.errorDetails,
+			isAborted: msg.isAborted,
+		};
+
+		await db.insert(messages).values(newMessage);
+
+		const parts = await db
+			.select()
+			.from(messageParts)
+			.where(eq(messageParts.messageId, msg.id))
+			.orderBy(asc(messageParts.index));
+
+		for (const part of parts) {
+			const newPart: typeof messageParts.$inferInsert = {
+				id: crypto.randomUUID(),
+				messageId: newMessageId,
+				index: part.index,
+				stepIndex: part.stepIndex,
+				type: part.type,
+				content: part.content,
+				agent: part.agent,
+				provider: part.provider,
+				model: part.model,
+				startedAt: part.startedAt,
+				completedAt: part.completedAt,
+				compactedAt: part.compactedAt,
+				toolName: part.toolName,
+				toolCallId: part.toolCallId,
+				toolDurationMs: part.toolDurationMs,
+			};
+
+			await db.insert(messageParts).values(newPart);
+			copiedParts++;
+		}
+	}
+
+	const result: SessionRow = {
+		...newSession,
+		totalInputTokens: null,
+		totalOutputTokens: null,
+		totalCachedTokens: null,
+		totalReasoningTokens: null,
+		totalToolTimeMs: null,
+		toolCountsJson: null,
+		contextSummary: null,
+		lastCompactedAt: null,
+	};
+
+	publish({
+		type: 'session.created',
+		sessionId: newSessionId,
+		payload: result,
+	});
+
+	return {
+		session: result,
+		parentSessionId,
+		branchPointMessageId: fromMessageId,
+		copiedMessages: messagesToCopy.length,
+		copiedParts,
+	};
+}
+
+export type ListBranchesResult = Array<{
+	session: SessionRow;
+	branchPointMessageId: string | null;
+	branchPointPreview: string | null;
+	createdAt: number;
+}>;
+
+export async function listBranches(
+	db: DB,
+	sessionId: string,
+	projectPath: string,
+): Promise<ListBranchesResult> {
+	const branches = await db
+		.select()
+		.from(sessions)
+		.where(eq(sessions.parentSessionId, sessionId))
+		.orderBy(asc(sessions.createdAt));
+
+	const results: ListBranchesResult = [];
+
+	for (const branch of branches) {
+		if (branch.projectPath !== projectPath) continue;
+
+		let preview: string | null = null;
+
+		if (branch.branchPointMessageId) {
+			const msgRows = await db
+				.select()
+				.from(messages)
+				.where(eq(messages.id, branch.branchPointMessageId));
+
+			if (msgRows.length > 0) {
+				const parts = await db
+					.select()
+					.from(messageParts)
+					.where(eq(messageParts.messageId, branch.branchPointMessageId))
+					.orderBy(asc(messageParts.index));
+
+				for (const part of parts) {
+					if (part.type === 'text') {
+						try {
+							const content = JSON.parse(part.content || '{}');
+							if (content.text) {
+								preview = content.text.slice(0, 100);
+								break;
+							}
+						} catch {}
+					}
+				}
+			}
+		}
+
+		results.push({
+			session: branch,
+			branchPointMessageId: branch.branchPointMessageId,
+			branchPointPreview: preview,
+			createdAt: branch.createdAt,
+		});
+	}
+
+	return results;
+}
+
+export async function getParentSession(
+	db: DB,
+	sessionId: string,
+	projectPath: string,
+): Promise<SessionRow | null> {
+	const sessionRows = await db
+		.select()
+		.from(sessions)
+		.where(eq(sessions.id, sessionId));
+
+	if (!sessionRows.length) return null;
+
+	const session = sessionRows[0];
+	if (!session.parentSessionId) return null;
+
+	const parentRows = await db
+		.select()
+		.from(sessions)
+		.where(eq(sessions.id, session.parentSessionId));
+
+	if (!parentRows.length) return null;
+
+	const parent = parentRows[0];
+	if (parent.projectPath !== projectPath) return null;
+
+	return parent;
+}

package/src/runtime/{db-operations.ts → session/db-operations.ts}

@@ -1,7 +1,7 @@
 import type { getDb } from '@agi-cli/database';
 import { messages, messageParts, sessions } from '@agi-cli/database/schema';
 import { eq } from 'drizzle-orm';
-import type { RunOpts } from './session-queue.ts';
+import type { RunOpts } from './queue.ts';
 
 export type UsageData = {
 	inputTokens?: number;

package/src/runtime/{session-manager.ts → session/manager.ts}

@@ -8,7 +8,7 @@ import {
 	ensureProviderEnv,
 	type ProviderId,
 } from '@agi-cli/sdk';
-import { publish } from '../events/bus.ts';
+import { publish } from '../../events/bus.ts';
 
 type SessionRow = typeof sessions.$inferSelect;
 

package/src/runtime/{session-queue.ts → session/queue.ts}

@@ -1,5 +1,5 @@
-import type { ProviderName } from './provider.ts';
-import { publish } from '../events/bus.ts';
+import type { ProviderName } from '../provider/index.ts';
+import { publish } from '../../events/bus.ts';
 
 export type RunOpts = {
 	sessionId: string;

package/src/runtime/stream/abort-handler.ts

@@ -0,0 +1,65 @@
+import type { getDb } from '@agi-cli/database';
+import { messages, messageParts } from '@agi-cli/database/schema';
+import { eq } from 'drizzle-orm';
+import { publish } from '../../events/bus.ts';
+import type { RunOpts } from '../session/queue.ts';
+import type { ToolAdapterContext } from '../../tools/adapter.ts';
+import type { AbortEvent } from './types.ts';
+
+export function createAbortHandler(
+	opts: RunOpts,
+	db: Awaited<ReturnType<typeof getDb>>,
+	getStepIndex: () => number,
+	sharedCtx: ToolAdapterContext,
+) {
+	return async ({ steps }: AbortEvent) => {
+		const stepIndex = getStepIndex();
+
+		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(),
+		});
+
+		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({
+			type: 'error',
+			sessionId: opts.sessionId,
+			payload: {
+				messageId: opts.assistantMessageId,
+				partId: abortPartId,
+				error: 'Generation stopped by user',
+				errorType: 'abort',
+				isAborted: true,
+				stepsCompleted: steps.length,
+			},
+		});
+	};
+}

package/src/runtime/stream/error-handler.ts

@@ -0,0 +1,200 @@
+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 { toErrorPayload } from '../errors/handling.ts';
+import type { RunOpts } from '../session/queue.ts';
+import type { ToolAdapterContext } from '../../tools/adapter.ts';
+import { pruneSession, performAutoCompaction } from '../message/compaction.ts';
+import { debugLog } from '../debug/index.ts';
+import { enqueueAssistantRun } from '../session/queue.ts';
+
+export function createErrorHandler(
+	opts: RunOpts,
+	db: Awaited<ReturnType<typeof getDb>>,
+	getStepIndex: () => number,
+	sharedCtx: ToolAdapterContext,
+	retryCallback?: (sessionId: string) => Promise<void>,
+) {
+	return async (err: unknown) => {
+		const errorPayload = toErrorPayload(err);
+		const isApiError = APICallError.isInstance(err);
+		const stepIndex = getStepIndex();
+
+		const errObj = err as Record<string, unknown>;
+		const nestedError = (errObj?.error as Record<string, unknown>)?.error as
+			| Record<string, unknown>
+			| undefined;
+		const errorCode =
+			(errObj?.code as string) ?? (nestedError?.code as string) ?? '';
+		const errorType =
+			(errObj?.apiErrorType as string) ?? (nestedError?.type as string) ?? '';
+		const fullErrorStr = JSON.stringify(err).toLowerCase();
+
+		const isPromptTooLong =
+			fullErrorStr.includes('prompt is too long') ||
+			fullErrorStr.includes('maximum context length') ||
+			fullErrorStr.includes('too many tokens') ||
+			fullErrorStr.includes('context_length_exceeded') ||
+			fullErrorStr.includes('request too large') ||
+			fullErrorStr.includes('exceeds the model') ||
+			fullErrorStr.includes('context window') ||
+			fullErrorStr.includes('input is too long') ||
+			errorCode === 'context_length_exceeded' ||
+			errorType === 'invalid_request_error';
+
+		debugLog(
+			`[stream-handlers] isPromptTooLong: ${isPromptTooLong}, errorCode: ${errorCode}, errorType: ${errorType}`,
+		);
+
+		if (isPromptTooLong && !opts.isCompactCommand) {
+			debugLog(
+				'[stream-handlers] Prompt too long detected, auto-compacting...',
+			);
+			let compactionSucceeded = false;
+			try {
+				const publishWrapper = (event: {
+					type: string;
+					sessionId: string;
+					payload: Record<string, unknown>;
+				}) => {
+					publish(event as Parameters<typeof publish>[0]);
+				};
+				const compactResult = await performAutoCompaction(
+					db,
+					opts.sessionId,
+					opts.assistantMessageId,
+					publishWrapper,
+					opts.provider,
+					opts.model,
+				);
+				if (compactResult.success) {
+					debugLog(
+						`[stream-handlers] Auto-compaction succeeded: ${compactResult.summary?.slice(0, 100)}...`,
+					);
+					compactionSucceeded = true;
+				} else {
+					debugLog(
+						`[stream-handlers] Auto-compaction failed: ${compactResult.error}, falling back to prune`,
+					);
+					const pruneResult = await pruneSession(db, opts.sessionId);
+					debugLog(
+						`[stream-handlers] Fallback pruned ${pruneResult.pruned} parts, saved ~${pruneResult.saved} tokens`,
+					);
+					compactionSucceeded = pruneResult.pruned > 0;
+				}
+			} catch (compactErr) {
+				debugLog(
+					`[stream-handlers] Auto-compact error: ${compactErr instanceof Error ? compactErr.message : String(compactErr)}`,
+				);
+			}
+
+			if (compactionSucceeded) {
+				await db
+					.update(messages)
+					.set({
+						status: 'completed',
+					})
+					.where(eq(messages.id, opts.assistantMessageId));
+
+				publish({
+					type: 'message.completed',
+					sessionId: opts.sessionId,
+					payload: {
+						id: opts.assistantMessageId,
+						autoCompacted: true,
+					},
+				});
+
+				if (retryCallback) {
+					debugLog('[stream-handlers] Triggering retry after compaction...');
+					const newAssistantMessageId = crypto.randomUUID();
+					await db.insert(messages).values({
+						id: newAssistantMessageId,
+						sessionId: opts.sessionId,
+						role: 'assistant',
+						status: 'pending',
+						agent: opts.agent,
+						provider: opts.provider,
+						model: opts.model,
+						createdAt: Date.now(),
+					});
+
+					publish({
+						type: 'message.created',
+						sessionId: opts.sessionId,
+						payload: { id: newAssistantMessageId, role: 'assistant' },
+					});
+
+					enqueueAssistantRun(
+						{
+							...opts,
+							assistantMessageId: newAssistantMessageId,
+						},
+						retryCallback,
+					);
+				} else {
+					debugLog(
+						'[stream-handlers] No retryCallback provided, cannot auto-retry',
+					);
+				}
+
+				return;
+			}
+		}
+
+		const errorPartId = crypto.randomUUID();
+		const displayMessage =
+			isPromptTooLong && !opts.isCompactCommand
+				? `${errorPayload.message}. Context auto-compacted - please retry your message.`
+				: errorPayload.message;
+		await db.insert(messageParts).values({
+			id: errorPartId,
+			messageId: opts.assistantMessageId,
+			index: await sharedCtx.nextIndex(),
+			stepIndex,
+			type: 'error',
+			content: JSON.stringify({
+				message: displayMessage,
+				type: errorPayload.type,
+				details: errorPayload.details,
+				isAborted: false,
+			}),
+			agent: opts.agent,
+			provider: opts.provider,
+			model: opts.model,
+			startedAt: Date.now(),
+			completedAt: Date.now(),
+		});
+
+		await db
+			.update(messages)
+			.set({
+				status: 'error',
+				error: displayMessage,
+				errorType: errorPayload.type,
+				errorDetails: JSON.stringify({
+					...errorPayload.details,
+					isApiError,
+					autoCompacted: isPromptTooLong && !opts.isCompactCommand,
+				}),
+				isAborted: false,
+			})
+			.where(eq(messages.id, opts.assistantMessageId));
+
+		publish({
+			type: 'error',
+			sessionId: opts.sessionId,
+			payload: {
+				messageId: opts.assistantMessageId,
+				partId: errorPartId,
+				error: displayMessage,
+				errorType: errorPayload.type,
+				details: errorPayload.details,
+				isAborted: false,
+				autoCompacted: isPromptTooLong && !opts.isCompactCommand,
+			},
+		});
+	};
+}

package/src/runtime/stream/finish-handler.ts

@@ -0,0 +1,123 @@
+import type { getDb } from '@agi-cli/database';
+import { messages, messageParts } from '@agi-cli/database/schema';
+import { eq } from 'drizzle-orm';
+import { publish } from '../../events/bus.ts';
+import { estimateModelCostUsd } from '@agi-cli/sdk';
+import type { RunOpts } from '../session/queue.ts';
+import {
+	pruneSession,
+	isOverflow,
+	getModelLimits,
+	type TokenUsage,
+	markSessionCompacted,
+} from '../message/compaction.ts';
+import { debugLog } from '../debug/index.ts';
+import type { FinishEvent } from './types.ts';
+
+export function createFinishHandler(
+	opts: RunOpts,
+	db: Awaited<ReturnType<typeof getDb>>,
+	completeAssistantMessageFn: (
+		fin: FinishEvent,
+		opts: RunOpts,
+		db: Awaited<ReturnType<typeof getDb>>,
+	) => Promise<void>,
+) {
+	return async (fin: FinishEvent) => {
+		try {
+			await completeAssistantMessageFn(fin, opts, db);
+		} catch {}
+
+		if (opts.isCompactCommand && fin.finishReason !== 'error') {
+			const assistantParts = await db
+				.select()
+				.from(messageParts)
+				.where(eq(messageParts.messageId, opts.assistantMessageId));
+			const hasTextContent = assistantParts.some(
+				(p) => p.type === 'text' && p.content && p.content !== '{"text":""}',
+			);
+
+			if (!hasTextContent) {
+				debugLog(
+					'[stream-handlers] /compact finished but no summary generated, skipping compaction marking',
+				);
+			} else {
+				try {
+					debugLog(
+						`[stream-handlers] /compact complete, marking session compacted`,
+					);
+					const result = await markSessionCompacted(
+						db,
+						opts.sessionId,
+						opts.assistantMessageId,
+					);
+					debugLog(
+						`[stream-handlers] Compacted ${result.compacted} parts, saved ~${result.saved} tokens`,
+					);
+				} catch (err) {
+					debugLog(
+						`[stream-handlers] Compaction failed: ${err instanceof Error ? err.message : String(err)}`,
+					);
+				}
+			}
+		}
+
+		const sessRows = await db
+			.select()
+			.from(messages)
+			.where(eq(messages.id, opts.assistantMessageId));
+
+		const usage = sessRows[0]
+			? {
+					inputTokens: Number(sessRows[0].promptTokens ?? 0),
+					outputTokens: Number(sessRows[0].completionTokens ?? 0),
+					totalTokens: Number(sessRows[0].totalTokens ?? 0),
+					cachedInputTokens: Number(sessRows[0].cachedInputTokens ?? 0),
+				}
+			: fin.usage;
+
+		const costUsd = usage
+			? estimateModelCostUsd(opts.provider, opts.model, usage)
+			: undefined;
+
+		if (usage) {
+			try {
+				const limits = getModelLimits(opts.provider, opts.model);
+				if (limits) {
+					const tokenUsage: TokenUsage = {
+						input: usage.inputTokens ?? 0,
+						output: usage.outputTokens ?? 0,
+						cacheRead:
+							(usage as { cachedInputTokens?: number }).cachedInputTokens ?? 0,
+					};
+
+					if (isOverflow(tokenUsage, limits)) {
+						debugLog(
+							`[stream-handlers] Context overflow detected, triggering prune for session ${opts.sessionId}`,
+						);
+						pruneSession(db, opts.sessionId).catch((err) => {
+							debugLog(
+								`[stream-handlers] Prune failed: ${err instanceof Error ? err.message : String(err)}`,
+							);
+						});
+					}
+				}
+			} catch (err) {
+				debugLog(
+					`[stream-handlers] Overflow check failed: ${err instanceof Error ? err.message : String(err)}`,
+				);
+			}
+		}
+
+		publish({
+			type: 'message.completed',
+			sessionId: opts.sessionId,
+			payload: {
+				id: opts.assistantMessageId,
+				usage,
+				costUsd,
+				finishReason: fin.finishReason,
+			},
+		});
+	};
+}

package/src/runtime/stream/handlers.ts

@@ -0,0 +1,5 @@
+export { createStepFinishHandler } from './step-finish.ts';
+export { createErrorHandler } from './error-handler.ts';
+export { createAbortHandler } from './abort-handler.ts';
+export { createFinishHandler } from './finish-handler.ts';
+export type { StepFinishEvent, FinishEvent, AbortEvent } from './types.ts';