@agi-cli/server 0.1.120 → 0.1.121

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';

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

@@ -0,0 +1,93 @@
+import type { getDb } from '@agi-cli/database';
+import { 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 { UsageData, ProviderMetadata } from '../session/db-operations.ts';
+import type { StepFinishEvent } from './types.ts';
+
+export function createStepFinishHandler(
+	opts: RunOpts,
+	db: Awaited<ReturnType<typeof getDb>>,
+	getStepIndex: () => number,
+	incrementStepIndex: () => number,
+	getCurrentPartId: () => string | null,
+	updateCurrentPartId: (id: string | null) => void,
+	updateAccumulated: (text: string) => void,
+	sharedCtx: ToolAdapterContext,
+	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();
+		const currentPartId = getCurrentPartId();
+		const stepIndex = getStepIndex();
+
+		try {
+			if (currentPartId) {
+				await db
+					.update(messageParts)
+					.set({ completedAt: finishedAt })
+					.where(eq(messageParts.id, currentPartId));
+			}
+		} catch {}
+
+		if (step.usage) {
+			try {
+				await updateSessionTokensIncrementalFn(
+					step.usage,
+					step.experimental_providerMetadata,
+					opts,
+					db,
+				);
+			} catch {}
+
+			try {
+				await updateMessageTokensIncrementalFn(
+					step.usage,
+					step.experimental_providerMetadata,
+					opts,
+					db,
+				);
+			} 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();
+			sharedCtx.stepIndex = newStepIndex;
+			updateCurrentPartId(null);
+			updateAccumulated('');
+		} catch {}
+	};
+}

package/src/runtime/stream/types.ts

@@ -0,0 +1,17 @@
+import type { UsageData, ProviderMetadata } from '../session/db-operations.ts';
+
+export type StepFinishEvent = {
+	usage?: UsageData;
+	finishReason?: string;
+	response?: unknown;
+	experimental_providerMetadata?: ProviderMetadata;
+};
+
+export type FinishEvent = {
+	usage?: Pick<UsageData, 'inputTokens' | 'outputTokens' | 'totalTokens'>;
+	finishReason?: string;
+};
+
+export type AbortEvent = {
+	steps: unknown[];
+};

package/src/runtime/{tool-context.ts → tools/context.ts}

@@ -1,7 +1,7 @@
 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';
+import { publish } from '../../events/bus.ts';
 
 export type StepExecutionState = {
 	chain: Promise<void>;

package/src/runtime/{tool-context-setup.ts → tools/setup.ts}

@@ -1,7 +1,7 @@
 import type { getDb } from '@agi-cli/database';
-import { time } from './debug.ts';
-import type { ToolAdapterContext } from '../tools/adapter.ts';
-import type { RunOpts } from './session-queue.ts';
+import { time } from '../debug/index.ts';
+import type { ToolAdapterContext } from '../../tools/adapter.ts';
+import type { RunOpts } from '../session/queue.ts';
 
 export type RunnerToolContext = ToolAdapterContext & { stepIndex: number };
 

package/src/runtime/{token-utils.ts → utils/token.ts}

@@ -1,6 +1,6 @@
 import { catalog } from '@agi-cli/sdk';
-import { debugLog } from './debug.ts';
-import type { ProviderName } from './provider.ts';
+import { debugLog } from '../debug/index.ts';
+import type { ProviderName } from '../provider/index.ts';
 
 /**
  * Gets the maximum output tokens allowed for a given provider/model combination.

package/src/tools/adapter.ts

@@ -3,18 +3,18 @@ import { messageParts, sessions } from '@agi-cli/database/schema';
 import { eq } from 'drizzle-orm';
 import { publish } from '../events/bus.ts';
 import type { DiscoveredTool } from '@agi-cli/sdk';
-import { getCwd, setCwd, joinRelative } from '../runtime/cwd.ts';
+import { getCwd, setCwd, joinRelative } from '../runtime/utils/cwd.ts';
 import type {
 	ToolAdapterContext,
 	StepExecutionState,
-} from '../runtime/tool-context.ts';
+} from '../runtime/tools/context.ts';
 import { isToolError } from '@agi-cli/sdk/tools/error';
 import {
 	toClaudeCodeName,
 	requiresClaudeCodeNaming,
-} from '../runtime/tool-mapping.ts';
+} from '../runtime/tools/mapping.ts';
 
-export type { ToolAdapterContext } from '../runtime/tool-context.ts';
+export type { ToolAdapterContext } from '../runtime/tools/context.ts';
 
 type ToolExecuteSignature = Tool['execute'] extends (
 	input: infer Input,