vibefast-cli 0.7.12 → 0.7.14

package/recipes/chatbot/packages/backend/convex/chatbotAgent.ts

@@ -1,216 +1,25 @@
 import { getAuthUserId } from '@convex-dev/auth/server';
 import { PersistentTextStreaming } from '@convex-dev/persistent-text-streaming';
-import type {
-  CoreMessage,
-  FilePart,
-  ImagePart,
-  TextPart,
-  ToolCallUnion,
-  ToolResultUnion,
-} from 'ai';
 import { v } from 'convex/values';
 
+import { handleRunStreaming } from './chatbot/streaming';
+import { ensureAgentSession, getSessionForConversation } from './chatbot/sessions';
+import { createTelemetryTracker } from './chatbot/telemetry';
 import { components, internal } from './_generated/api';
-import type { Doc, Id } from './_generated/dataModel';
+import type { Id } from './_generated/dataModel';
 import {
   internalAction,
   internalMutation,
   internalQuery,
-  type MutationCtx,
   mutation,
-  type QueryCtx,
   query,
 } from './_generated/server';
-import { agent, DEFAULT_MODEL, resolveChatModel } from './agents';
+import { DEFAULT_MODEL } from './agents';
 import { enforceAgentRateLimit } from './lib/rateLimit';
-import { getChatTelemetryMode, type TelemetryMode } from './lib/telemetry';
-
 const persistentTextStreaming = new PersistentTextStreaming(
   components.persistentTextStreaming,
 );
 
-type AgentSessionDoc = Doc<'agentSessions'>;
-type AgentTools = NonNullable<(typeof agent)['options']['tools']>;
-
-const SENTENCE_DELIMITER = /[.!?]/;
-const MAX_PENDING_CHARS = 180;
-
-const TOOL_STATUS_LABELS: Record<string, string> = {
-  knowledgeRetrieval: 'Looking through knowledge…',
-  tavilySearch: 'Searching the web…',
-  userProfile: 'Checking your profile…',
-};
-
-function toolLabel(toolName: string) {
-  return TOOL_STATUS_LABELS[toolName] ?? `Using ${toolName}`;
-}
-
-function mergeToolResults(
-  toolCalls: ToolCallUnion<AgentTools>[],
-  toolResults: ToolResultUnion<AgentTools>[],
-) {
-  if (!toolCalls?.length) {
-    return [];
-  }
-
-  const resultByCallId = new Map<string, unknown>();
-
-  for (const result of toolResults ?? []) {
-    const key =
-      (result as { toolCallId?: string }).toolCallId ??
-      (result as { id?: string }).id;
-    if (key) {
-      resultByCallId.set(key, (result as { result?: unknown }).result);
-    }
-  }
-
-  return toolCalls.map((call) => {
-    const callId =
-      (call as { id?: string }).id ??
-      (call as { toolCallId?: string }).toolCallId ??
-      call.toolName;
-
-    return {
-      toolName: call.toolName ?? callId ?? 'tool',
-      args: (call as { args?: unknown }).args ?? null,
-      result: callId ? resultByCallId.get(callId) : undefined,
-    };
-  });
-}
-
-async function buildUserContent(
-  ctx: any,
-  message: Doc<'messages'>,
-): Promise<string | (TextPart | ImagePart | FilePart)[]> {
-  const attachments = message.attachments ?? [];
-  if (attachments.length === 0) {
-    return message.text ?? '';
-  }
-
-  const parts: (TextPart | ImagePart)[] = [];
-  if (message.text) {
-    parts.push({ type: 'text', text: message.text });
-  }
-
-  for (const attachment of attachments) {
-    try {
-      const url = await ctx.storage.getUrl(attachment.storageId);
-      if (url) {
-        parts.push({ type: 'image', image: new URL(url) });
-      }
-    } catch (error) {
-      console.error('[chatbotAgent] Failed to resolve attachment URL', {
-        storageId: attachment.storageId,
-        error,
-      });
-    }
-  }
-
-  return parts.length > 0 ? parts : (message.text ?? '');
-}
-
-const DEFAULT_TELEMETRY_SAMPLE = 0.1;
-
-type TelemetryCheckpoint = {
-  label: string;
-  at: number;
-  data?: Record<string, unknown>;
-};
-
-type TelemetryTracker = {
-  mark: (label: string, data?: Record<string, unknown>) => void;
-  finalize: (status: 'ok' | 'error', extra?: Record<string, unknown>) => void;
-};
-
-function createTelemetryTracker(
-  event: string,
-  baseContext: Record<string, unknown>,
-  sampleProbability = DEFAULT_TELEMETRY_SAMPLE,
-): TelemetryTracker | null {
-  const mode: TelemetryMode = getChatTelemetryMode();
-  const shouldTrack =
-    mode === 'debug' ||
-    (mode === 'sample' && Math.random() < sampleProbability);
-
-  if (!shouldTrack) {
-    return null;
-  }
-
-  const start = Date.now();
-  const checkpoints: TelemetryCheckpoint[] = [{ label: 'start', at: start }];
-
-  return {
-    mark(label, data) {
-      checkpoints.push({ label, at: Date.now(), data });
-    },
-    finalize(status, extra) {
-      const finishedAt = Date.now();
-      const timeline = checkpoints.map((checkpoint) => ({
-        label: checkpoint.label,
-        msFromStart: checkpoint.at - start,
-        ...(checkpoint.data ?? {}),
-      }));
-
-      console.log('[chatbotAgent][telemetry]', {
-        event,
-        mode,
-        status,
-        totalMs: finishedAt - start,
-        ...baseContext,
-        ...(extra ?? {}),
-        timeline,
-      });
-    },
-  };
-}
-
-async function getSessionForConversation(
-  ctx: Pick<MutationCtx, 'db'> | Pick<QueryCtx, 'db'>,
-  conversationId: Id<'conversations'>,
-) {
-  return await ctx.db
-    .query('agentSessions')
-    .withIndex('by_conversationId', (q) =>
-      q.eq('conversationId', conversationId),
-    )
-    .first();
-}
-
-async function ensureAgentSession(
-  ctx: MutationCtx,
-  conversationId: Id<'conversations'>,
-  userId: Id<'users'>,
-  preferredModel?: string,
-): Promise<{ session: AgentSessionDoc; wasCreated: boolean }> {
-  const existing = await getSessionForConversation(ctx, conversationId);
-  if (existing) {
-    return { session: existing, wasCreated: false };
-  }
-
-  const { threadId } = await agent.createThread(ctx, {
-    userId,
-  });
-
-  const sessionId = await ctx.db.insert('agentSessions', {
-    conversationId,
-    userId,
-    agentThreadId: threadId,
-    activeStreamId: undefined,
-    lastUserMessageId: undefined,
-    lastAssistantMessageId: undefined,
-    preferredModel: preferredModel ?? DEFAULT_MODEL,
-    updatedAt: Date.now(),
-    lastErrorMessage: undefined,
-    lastErrorAt: undefined,
-  });
-
-  const session = await ctx.db.get(sessionId);
-  if (!session) {
-    throw new Error('Failed to create agent session');
-  }
-  return { session, wasCreated: true };
-}
-
 export const startAgentStream = mutation({
   args: {
     conversationId: v.id('conversations'),
@@ -411,461 +220,7 @@ export const runStreaming = internalAction({
     preferredModel: v.optional(v.string()),
     agentSessionId: v.id('agentSessions'),
   },
-  handler: async (ctx, args) => {
-    const telemetry = createTelemetryTracker('runStreaming', {
-      streamId: args.streamId,
-      conversationId: args.conversationId,
-      agentSessionId: args.agentSessionId,
-    });
-
-    let telemetryStatus: 'ok' | 'error' = 'ok';
-    const telemetryExtra: Record<string, unknown> = {};
-    let failureHandled = false;
-
-    const inferErrorMessage = (raw: unknown): string | undefined => {
-      if (!raw) return undefined;
-      if (typeof raw === 'string') {
-        try {
-          const parsed = JSON.parse(raw);
-          if (parsed && typeof parsed === 'object') {
-            if (typeof parsed.error === 'string') {
-              return parsed.error;
-            }
-            if (typeof parsed.message === 'string') {
-              return parsed.message;
-            }
-          }
-        } catch {
-          // ignore JSON parse failures, fall back to raw string if concise
-        }
-        const trimmed = raw.trim();
-        if (trimmed.length > 0 && trimmed.length <= 300) {
-          return trimmed;
-        }
-        return undefined;
-      }
-      if (raw instanceof Error) {
-        return raw.message;
-      }
-      return undefined;
-    };
-
-    const markFailure = async (
-      maybeMessage: unknown,
-      extra?: Record<string, unknown>,
-    ) => {
-      telemetryStatus = 'error';
-      if (extra) {
-        Object.assign(telemetryExtra, extra);
-      }
-      failureHandled = true;
-      const message =
-        inferErrorMessage(maybeMessage) ??
-        'Sorry, I ran into a technical issue while replying. Please try again.';
-
-      try {
-        await ctx.runMutation(
-          components.persistentTextStreaming.lib.setStreamStatus,
-          {
-            streamId: args.streamId,
-            status: 'error',
-          },
-        );
-      } catch (statusError) {
-        console.error(
-          '[chatbotAgent] Failed to update stream status to error',
-          {
-            streamId: args.streamId,
-            statusError,
-          },
-        );
-      }
-
-      await ctx.runMutation(internal.chatbotAgent.markStreamFailed, {
-        agentSessionId: args.agentSessionId,
-        conversationId: args.conversationId,
-        streamId: args.streamId,
-        errorMessage: message,
-      });
-    };
-
-    try {
-      telemetry?.mark('conversation.lookup');
-      const conversation = await ctx.runQuery(
-        internal.chatbotAgent.getConversationForStreaming,
-        { conversationId: args.conversationId },
-      );
-      if (!conversation) {
-        await markFailure('Conversation not found.', {
-          reason: 'conversationMissing',
-        });
-        return;
-      }
-
-      telemetry?.mark('session.lookup');
-      const session = await ctx.runQuery(
-        internal.chatbotAgent.getAgentSessionById,
-        { agentSessionId: args.agentSessionId },
-      );
-      if (!session) {
-        await markFailure('Agent session missing.', {
-          reason: 'sessionMissing',
-        });
-        return;
-      }
-
-      telemetry?.mark('message.lookup');
-      const message = await ctx.runQuery(
-        internal.chatbotAgent.getMessageForStreaming,
-        { messageId: args.userMessageId },
-      );
-      if (!message || message.conversationId !== args.conversationId) {
-        await markFailure('User message missing.', {
-          reason: 'userMessageMissing',
-        });
-        return;
-      }
-
-      const { model, provider, modelId } = resolveChatModel(
-        args.preferredModel ?? session.preferredModel ?? undefined,
-      );
-      telemetryExtra.provider = provider;
-      telemetryExtra.modelId = modelId;
-      telemetry?.mark('model.resolved', { provider, modelId });
-
-      const userContent = await buildUserContent(ctx, message);
-      telemetry?.mark('content.resolved', {
-        hasAttachments: Boolean(message.attachments?.length),
-      });
-
-      const { thread } = await agent.continueThread(ctx, {
-        threadId: args.agentThreadId,
-        userId: conversation.userId,
-      });
-      telemetry?.mark('agent.threadReady');
-
-      const messages: CoreMessage[] = [
-        {
-          role: 'user',
-          content: userContent as string | (TextPart | ImagePart | FilePart)[],
-        },
-      ];
-
-      let streamResult;
-      try {
-        streamResult = await thread.streamText(
-          {
-            model,
-            messages,
-          },
-          {
-            storageOptions: {
-              saveMessages: 'promptAndOutput',
-            },
-          },
-        );
-      } catch (streamInitError) {
-        console.error('[chatbotAgent] Failed to initialize stream', {
-          streamInitError,
-          provider,
-          modelId,
-        });
-        await markFailure(streamInitError, {
-          reason: 'streamInitFailed',
-          provider,
-          modelId,
-        });
-        return;
-      }
-
-      telemetry?.mark('agent.stream.started');
-
-      const activeToolCalls = new Map<
-        string,
-        { toolName: string; label: string }
-      >();
-
-      type StatusRecord =
-        | null
-        | { type: 'thinking'; label: string | null }
-        | {
-            type: 'tool';
-            label: string | null;
-            toolName: string | null;
-            toolCallId: string | null;
-          };
-
-      let lastStatus: StatusRecord = { type: 'thinking', label: 'Thinking…' };
-
-      const applyStatus = async (
-        status:
-          | { type: 'thinking'; label?: string }
-          | {
-              type: 'tool';
-              toolName: string;
-              label?: string;
-              toolCallId?: string;
-            }
-          | null,
-      ) => {
-        const normalized: StatusRecord = status
-          ? status.type === 'tool'
-            ? {
-                type: 'tool',
-                label: status.label ?? null,
-                toolName: status.toolName,
-                toolCallId: status.toolCallId ?? null,
-              }
-            : {
-                type: 'thinking',
-                label: status.label ?? null,
-              }
-          : null;
-
-        if (normalized === null && lastStatus === null) {
-          return;
-        }
-
-        if (
-          normalized &&
-          lastStatus &&
-          normalized.type === lastStatus.type &&
-          normalized.label === lastStatus.label &&
-          (normalized.type !== 'tool' ||
-            (lastStatus.type === 'tool' &&
-              normalized.toolName === lastStatus.toolName &&
-              normalized.toolCallId === lastStatus.toolCallId))
-        ) {
-          return;
-        }
-
-        lastStatus = normalized;
-
-        const payload = normalized
-          ? {
-              type: normalized.type,
-              label: normalized.label ?? undefined,
-              toolName:
-                normalized.type === 'tool'
-                  ? (normalized.toolName ?? undefined)
-                  : undefined,
-            }
-          : null;
-
-        await ctx.runMutation(internal.chatbotAgent.updateSessionStatus, {
-          agentSessionId: args.agentSessionId,
-          status: payload,
-        });
-      };
-
-      let finalResponse = '';
-      let pendingChunk = '';
-      let wroteChunk = false;
-
-      const flushChunk = async (final: boolean) => {
-        const textToWrite = pendingChunk;
-        pendingChunk = '';
-        if (!final && textToWrite.length === 0) {
-          return;
-        }
-        await ctx.runMutation(components.persistentTextStreaming.lib.addChunk, {
-          streamId: args.streamId,
-          text: textToWrite,
-          final,
-        });
-        wroteChunk = true;
-      };
-
-      const STREAM_TIMEOUT = 120000;
-      const streamPromise = (async () => {
-        try {
-          for await (const part of streamResult.fullStream) {
-            switch (part.type) {
-              case 'text-delta': {
-                const delta = part.textDelta;
-                if (!delta) {
-                  break;
-                }
-                finalResponse += delta;
-                pendingChunk += delta;
-
-                if (activeToolCalls.size === 0 && lastStatus !== null) {
-                  await applyStatus(null);
-                }
-
-                if (
-                  SENTENCE_DELIMITER.test(delta) ||
-                  pendingChunk.length >= MAX_PENDING_CHARS
-                ) {
-                  await flushChunk(false);
-                }
-                break;
-              }
-              case 'tool-call-streaming-start': {
-                activeToolCalls.set(part.toolCallId, {
-                  toolName: part.toolName,
-                  label: toolLabel(part.toolName),
-                });
-                await applyStatus({
-                  type: 'tool',
-                  toolName: part.toolName,
-                  label: toolLabel(part.toolName),
-                  toolCallId: part.toolCallId,
-                });
-                break;
-              }
-              case 'tool-call': {
-                const callId =
-                  (part as { toolCallId?: string }).toolCallId ??
-                  (part as { id?: string }).id ??
-                  part.toolName;
-                activeToolCalls.set(callId, {
-                  toolName: part.toolName,
-                  label: toolLabel(part.toolName),
-                });
-                await applyStatus({
-                  type: 'tool',
-                  toolName: part.toolName,
-                  label: toolLabel(part.toolName),
-                  toolCallId: callId,
-                });
-                break;
-              }
-              case 'tool-result': {
-                const callId =
-                  (part as { toolCallId?: string }).toolCallId ??
-                  (part as { id?: string }).id ??
-                  null;
-                if (callId) {
-                  activeToolCalls.delete(callId);
-                } else {
-                  activeToolCalls.clear();
-                }
-
-                if (activeToolCalls.size > 0) {
-                  const [nextId, next] = activeToolCalls.entries().next()
-                    .value as [string, { toolName: string; label: string }];
-                  await applyStatus({
-                    type: 'tool',
-                    toolName: next.toolName,
-                    label: next.label,
-                    toolCallId: nextId,
-                  });
-                } else {
-                  await applyStatus({
-                    type: 'thinking',
-                    label: 'Thinking…',
-                  });
-                }
-                break;
-              }
-              default:
-                break;
-            }
-          }
-        } catch (streamError) {
-          throw streamError;
-        }
-      })();
-
-      const timeoutPromise = new Promise((_, reject) => {
-        setTimeout(() => {
-          reject(new Error('Stream timed out after 2 minutes'));
-        }, STREAM_TIMEOUT);
-      });
-
-      try {
-        await Promise.race([streamPromise, timeoutPromise]);
-      } catch (streamError) {
-        console.error('[chatbotAgent] Stream failed or timed out', {
-          streamError,
-          streamId: args.streamId,
-        });
-        telemetryExtra.streamError =
-          streamError instanceof Error
-            ? streamError.message
-            : String(streamError);
-        await markFailure(streamError, {
-          reason: 'streamReadFailed',
-        });
-        return;
-      }
-
-      if (pendingChunk.length > 0) {
-        await flushChunk(true);
-      } else {
-        await ctx.runMutation(
-          components.persistentTextStreaming.lib.setStreamStatus,
-          {
-            streamId: args.streamId,
-            status: wroteChunk ? 'done' : 'done',
-          },
-        );
-      }
-
-      await applyStatus(null);
-
-      telemetry?.mark('agent.stream.completed', {
-        responseLength: finalResponse.length,
-      });
-
-      const [usage, toolCalls, toolResults] = await Promise.all([
-        streamResult.usage.catch(() => undefined),
-        streamResult.toolCalls.catch(() => [] as ToolCallUnion<AgentTools>[]),
-        streamResult.toolResults.catch(
-          () => [] as ToolResultUnion<AgentTools>[],
-        ),
-      ]);
-
-      const serializedTools = mergeToolResults(
-        toolCalls ?? [],
-        toolResults ?? [],
-      );
-
-      await ctx.runMutation(internal.chatbotAgent.saveAssistantMessage, {
-        agentSessionId: args.agentSessionId,
-        conversationId: args.conversationId,
-        streamId: args.streamId,
-        text: finalResponse,
-        provider,
-        modelId,
-        usage: usage
-          ? {
-              promptTokens: usage.promptTokens ?? undefined,
-              completionTokens: usage.completionTokens ?? undefined,
-              totalTokens: usage.totalTokens ?? undefined,
-            }
-          : undefined,
-        toolCalls: serializedTools,
-      });
-
-      telemetryExtra.responseLength = finalResponse.length;
-      telemetryStatus = 'ok';
-    } catch (error) {
-      console.error('[chatbotAgent] runStreaming caught error', {
-        error,
-        failureHandled,
-        streamId: args.streamId,
-        conversationId: args.conversationId,
-      });
-
-      if (!failureHandled) {
-        try {
-          await markFailure(error);
-          console.log('[chatbotAgent] markFailure completed successfully');
-        } catch (markFailureError) {
-          console.error('[chatbotAgent] markFailure itself failed!', {
-            markFailureError,
-            originalError: error,
-          });
-        }
-      }
-      telemetryExtra.error =
-        error instanceof Error ? error.message : String(error);
-    } finally {
-      telemetry?.finalize(telemetryStatus, telemetryExtra);
-    }
-  },
+  handler: async (ctx, args) => handleRunStreaming(ctx, args),
 });
 
 export const updateSessionStatus = internalMutation({