vibefast-cli 0.7.0 → 0.7.1

package/2/packages/backend/convex/chatbotAgent.ts

@@ -1,1085 +0,0 @@
-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 { components, internal } from './_generated/api';
-import type { Doc, 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 { 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'),
-    userMessageId: v.id('messages'),
-    preferredModel: v.optional(v.string()),
-  },
-  handler: async (ctx, args) => {
-    const telemetry = createTelemetryTracker('startAgentStream', {
-      conversationId: args.conversationId,
-      userMessageId: args.userMessageId,
-      requestedModel: args.preferredModel ?? null,
-    });
-
-    const userId = await getAuthUserId(ctx);
-    telemetry?.mark('auth.resolved', { hasUser: Boolean(userId) });
-
-    try {
-      if (!userId) {
-        throw new Error('Not authenticated');
-      }
-
-      const conversation = await ctx.db.get(args.conversationId);
-      if (!conversation || conversation.userId !== userId) {
-        throw new Error('Conversation not found');
-      }
-      telemetry?.mark('conversation.loaded');
-
-      const message = await ctx.db.get(args.userMessageId);
-      if (!message || message.conversationId !== args.conversationId) {
-        throw new Error('User message not found for conversation');
-      }
-      telemetry?.mark('message.loaded', {
-        hasAttachments: Boolean(message.attachments?.length),
-      });
-
-      if (message.authorType !== 'user') {
-        throw new Error('Only user messages can start a streaming response');
-      }
-
-      await enforceAgentRateLimit(ctx.db, userId);
-      telemetry?.mark('rateLimit.checked');
-
-      const { session, wasCreated } = await ensureAgentSession(
-        ctx,
-        args.conversationId,
-        userId,
-        args.preferredModel ?? undefined,
-      );
-      telemetry?.mark('session.ready', { wasCreated });
-
-      const streamId = await persistentTextStreaming.createStream(ctx);
-      telemetry?.mark('stream.created');
-
-      await ctx.db.patch(session._id, {
-        activeStreamId: streamId,
-        lastUserMessageId: args.userMessageId,
-        preferredModel:
-          args.preferredModel ?? session.preferredModel ?? DEFAULT_MODEL,
-        updatedAt: Date.now(),
-        lastErrorMessage: undefined,
-        lastErrorAt: undefined,
-        currentStatus: {
-          type: 'thinking',
-          label: 'Thinking…',
-          startedAt: Date.now(),
-        },
-      });
-      telemetry?.mark('session.patched');
-
-      console.log('[chatbotAgent] startAgentStream', {
-        conversationId: args.conversationId,
-        userMessageId: args.userMessageId,
-        requestedModel: args.preferredModel ?? null,
-        persistedPreferredModel:
-          args.preferredModel ?? session.preferredModel ?? DEFAULT_MODEL,
-        streamId,
-        agentSessionId: session._id,
-      });
-
-      await ctx.scheduler.runAfter(0, internal.chatbotAgent.runStreaming, {
-        streamId,
-        conversationId: args.conversationId,
-        agentThreadId: session.agentThreadId,
-        userMessageId: args.userMessageId,
-        preferredModel:
-          args.preferredModel ?? session.preferredModel ?? undefined,
-        agentSessionId: session._id,
-      });
-      telemetry?.mark('stream.scheduled');
-      telemetry?.finalize('ok', { streamId, agentSessionId: session._id });
-
-      return { streamId };
-    } catch (error) {
-      telemetry?.finalize('error', {
-        error:
-          error instanceof Error
-            ? error.message
-            : typeof error === 'string'
-              ? error
-              : 'Unknown error',
-      });
-      throw error;
-    }
-  },
-});
-
-export const getAgentStreamBody = query({
-  args: {
-    conversationId: v.id('conversations'),
-    streamId: v.string(),
-  },
-  handler: async (ctx, args) => {
-    const userId = await getAuthUserId(ctx);
-    if (!userId) {
-      throw new Error('Not authenticated');
-    }
-
-    const conversation = await ctx.db.get(args.conversationId);
-    if (!conversation || conversation.userId !== userId) {
-      throw new Error('Conversation not found');
-    }
-
-    const session = await getSessionForConversation(ctx, args.conversationId);
-    if (!session || session.activeStreamId !== args.streamId) {
-      // Allow reading historical streams that were saved on the message metadata.
-      const messageWithStream = await ctx.db
-        .query('messages')
-        .withIndex('by_conversationId', (q) =>
-          q.eq('conversationId', args.conversationId),
-        )
-        .filter((q) => q.eq(q.field('metadata.streamId'), args.streamId))
-        .first();
-
-      if (!messageWithStream) {
-        console.warn(
-          '[chatbotAgent] getAgentStreamBody fallback returning empty body',
-          {
-            conversationId: args.conversationId,
-            streamId: args.streamId,
-          },
-        );
-        return '';
-      }
-    }
-
-    return await persistentTextStreaming.getStreamBody(
-      ctx,
-      args.streamId as any,
-    );
-  },
-});
-
-export const getAgentSession = query({
-  args: {
-    conversationId: v.id('conversations'),
-  },
-  handler: async (ctx, args) => {
-    const userId = await getAuthUserId(ctx);
-    if (!userId) {
-      throw new Error('Not authenticated');
-    }
-
-    const conversation = await ctx.db.get(args.conversationId);
-    if (!conversation || conversation.userId !== userId) {
-      throw new Error('Conversation not found');
-    }
-
-    const session = await getSessionForConversation(ctx, args.conversationId);
-    if (!session) {
-      return null;
-    }
-
-    return {
-      activeStreamId: session.activeStreamId ?? null,
-      preferredModel: session.preferredModel ?? DEFAULT_MODEL,
-      agentThreadId: session.agentThreadId,
-      lastErrorMessage: session.lastErrorMessage ?? null,
-      lastErrorAt: session.lastErrorAt ?? null,
-      lastUserMessageId: session.lastUserMessageId ?? null,
-      currentStatus: session.currentStatus
-        ? {
-            type: session.currentStatus.type,
-            label: session.currentStatus.label ?? null,
-            toolName: session.currentStatus.toolName ?? null,
-            startedAt: session.currentStatus.startedAt,
-          }
-        : null,
-    };
-  },
-});
-
-export const runStreaming = internalAction({
-  args: {
-    streamId: v.string(),
-    conversationId: v.id('conversations'),
-    agentThreadId: v.string(),
-    userMessageId: v.id('messages'),
-    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);
-    }
-  },
-});
-
-export const updateSessionStatus = internalMutation({
-  args: {
-    agentSessionId: v.id('agentSessions'),
-    status: v.union(
-      v.null(),
-      v.object({
-        type: v.union(v.literal('thinking'), v.literal('tool')),
-        label: v.optional(v.string()),
-        toolName: v.optional(v.string()),
-      }),
-    ),
-  },
-  handler: async (ctx, { agentSessionId, status }) => {
-    await ctx.db.patch(agentSessionId, {
-      currentStatus: status
-        ? {
-            type: status.type,
-            label: status.label ?? undefined,
-            toolName: status.toolName ?? undefined,
-            startedAt: Date.now(),
-          }
-        : undefined,
-    });
-  },
-});
-
-export const markStreamFailed = internalMutation({
-  args: {
-    agentSessionId: v.id('agentSessions'),
-    conversationId: v.id('conversations'),
-    streamId: v.string(),
-    errorMessage: v.optional(v.string()),
-  },
-  handler: async (ctx, args) => {
-    console.log('[chatbotAgent] markStreamFailed called', {
-      agentSessionId: args.agentSessionId,
-      conversationId: args.conversationId,
-      streamId: args.streamId,
-      errorMessage: args.errorMessage,
-    });
-
-    const session = await ctx.db.get(args.agentSessionId);
-
-    if (!session) {
-      console.error('[chatbotAgent] markStreamFailed: session not found!', {
-        agentSessionId: args.agentSessionId,
-      });
-      return;
-    }
-
-    const lastUserMessageId = session?.lastUserMessageId;
-
-    if (lastUserMessageId) {
-      try {
-        await ctx.db.delete(lastUserMessageId);
-        console.log('[chatbotAgent] Deleted failed user message', {
-          messageId: lastUserMessageId,
-        });
-      } catch (error) {
-        console.warn('[chatbotAgent] Failed to delete failed user message', {
-          messageId: lastUserMessageId,
-          error,
-        });
-      }
-    }
-
-    await ctx.db.patch(args.agentSessionId, {
-      activeStreamId: undefined,
-      lastUserMessageId: undefined,
-      lastErrorMessage: args.errorMessage ?? 'An error occurred',
-      lastErrorAt: Date.now(),
-      updatedAt: Date.now(),
-      currentStatus: undefined,
-    });
-
-    console.log('[chatbotAgent] Session patched with error state', {
-      agentSessionId: args.agentSessionId,
-      lastErrorAt: Date.now(),
-      lastErrorMessage: args.errorMessage,
-    });
-  },
-});
-
-export const saveAssistantMessage = internalMutation({
-  args: {
-    agentSessionId: v.id('agentSessions'),
-    conversationId: v.id('conversations'),
-    streamId: v.string(),
-    text: v.optional(v.string()),
-    provider: v.optional(v.string()),
-    modelId: v.optional(v.string()),
-    usage: v.optional(
-      v.object({
-        promptTokens: v.optional(v.number()),
-        completionTokens: v.optional(v.number()),
-        totalTokens: v.optional(v.number()),
-      }),
-    ),
-    toolCalls: v.optional(
-      v.array(
-        v.object({
-          toolName: v.string(),
-          args: v.any(),
-          result: v.optional(v.any()),
-        }),
-      ),
-    ),
-  },
-  handler: async (ctx, args) => {
-    const telemetry = createTelemetryTracker('saveAssistantMessage', {
-      streamId: args.streamId,
-      agentSessionId: args.agentSessionId,
-      conversationId: args.conversationId,
-    });
-
-    try {
-      telemetry?.mark('lookup.session');
-      const session = await ctx.db.get(args.agentSessionId);
-      telemetry?.mark('lookup.conversation');
-      const conversation = await ctx.db.get(args.conversationId);
-
-      if (!session || !conversation) {
-        telemetry?.finalize('error', {
-          reason: 'missingData',
-          sessionExists: Boolean(session),
-          conversationExists: Boolean(conversation),
-        });
-        console.warn(
-          '[chatbotAgent] Unable to save assistant message: missing data',
-          {
-            sessionExists: !!session,
-            conversationExists: !!conversation,
-          },
-        );
-        return;
-      }
-
-      const trimmedText = args.text?.trim();
-
-      telemetry?.mark('message.insert.started', {
-        hasText: Boolean(trimmedText),
-        hasTools: Boolean(args.toolCalls?.length),
-      });
-      const messageId = await ctx.db.insert('messages', {
-        conversationId: args.conversationId,
-        userId: conversation.userId,
-        authorType: 'bot',
-        text: trimmedText && trimmedText.length > 0 ? trimmedText : undefined,
-        aiProvider: args.provider,
-        aiModel: args.modelId,
-        toolCalls: args.toolCalls ?? undefined,
-        metadata: {
-          streamId: args.streamId,
-          usage: args.usage
-            ? {
-                promptTokens: args.usage.promptTokens ?? 0,
-                completionTokens: args.usage.completionTokens ?? 0,
-                totalTokens: args.usage.totalTokens ?? 0,
-              }
-            : undefined,
-          status: 'complete',
-        },
-      });
-      telemetry?.mark('message.insert.completed', { messageId });
-
-      await ctx.db.patch(args.agentSessionId, {
-        activeStreamId: undefined,
-        lastAssistantMessageId: messageId,
-        lastErrorMessage: undefined,
-        lastErrorAt: undefined,
-        updatedAt: Date.now(),
-        currentStatus: undefined,
-      });
-      telemetry?.mark('session.patched');
-      telemetry?.finalize('ok', { messageId });
-    } catch (error) {
-      telemetry?.finalize('error', {
-        error:
-          error instanceof Error
-            ? error.message
-            : typeof error === 'string'
-              ? error
-              : 'Unknown error',
-      });
-      throw error;
-    }
-  },
-});
-
-export const getConversationForStreaming = internalQuery({
-  args: {
-    conversationId: v.id('conversations'),
-  },
-  handler: async (ctx, args) => {
-    return await ctx.db.get(args.conversationId);
-  },
-});
-
-export const getMessageForStreaming = internalQuery({
-  args: {
-    messageId: v.id('messages'),
-  },
-  handler: async (ctx, args) => {
-    return await ctx.db.get(args.messageId);
-  },
-});
-
-export const getAgentSessionById = internalQuery({
-  args: {
-    agentSessionId: v.id('agentSessions'),
-  },
-  handler: async (ctx, args) => {
-    return await ctx.db.get(args.agentSessionId);
-  },
-});