vibefast-cli 0.2.2 → 0.2.4

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

@@ -0,0 +1,1085 @@
+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);
+  },
+});