zopu-code/packages/backend/convex/conversationMessages.ts

import { makeFunctionReference } from "convex/server";
import { ConvexError, v } from "convex/values";

import type { Doc, Id } from "./_generated/dataModel";
import {
  env,
  internalAction,
  internalMutation,
  internalQuery,
  mutation,
  query,
} from "./_generated/server";
import { requireOrganizationMember } from "./authz";

const MAX_ATTEMPTS = 3;

const runTurnRef = makeFunctionReference<
  "action",
  { turnId: Id<"conversationTurns">; attempt: number }
>("conversationMessages:runTurn");
const getTurnRef = makeFunctionReference<
  "query",
  { turnId: Id<"conversationTurns"> },
  {
    turn: Doc<"conversationTurns">;
    organizationId: Id<"organizations">;
    user: Doc<"conversationMessages">;
    assistant: Doc<"conversationMessages">;
    attachments: Doc<"conversationAttachments">[];
  } | null
>("conversationMessages:getTurn");
const markProcessingRef = makeFunctionReference<
  "mutation",
  {
    turnId: Id<"conversationTurns">;
    attempt: number;
    leaseOwner: string;
  },
  boolean
>("conversationMessages:markProcessing");
const completeTurnRef = makeFunctionReference<
  "mutation",
  {
    turnId: Id<"conversationTurns">;
    attempt: number;
    leaseOwner: string;
    submissionId: string;
    text: string;
  },
  boolean
>("conversationMessages:completeTurn");
const failTurnRef = makeFunctionReference<
  "mutation",
  {
    turnId: Id<"conversationTurns">;
    attempt: number;
    leaseOwner: string;
    error: string;
    retry: boolean;
  },
  boolean
>("conversationMessages:failTurn");

export const generateUploadUrl = mutation({
  args: {},
  handler: async (ctx): Promise<string> => {
    if (!(await ctx.auth.getUserIdentity())) {
      throw new ConvexError("Authentication required");
    }
    return await ctx.storage.generateUploadUrl();
  },
});

export const send = mutation({
  args: {
    clientRequestId: v.string(),
    images: v.array(
      v.object({
        filename: v.optional(v.string()),
        mimeType: v.string(),
        storageId: v.id("_storage"),
      })
    ),
    organizationId: v.id("organizations"),
    rawText: v.string(),
  },
  handler: async (ctx, args) => {
    await requireOrganizationMember(ctx, args.organizationId);
    let conversation = await ctx.db
      .query("conversations")
      .withIndex("by_organizationId", (q) =>
        q.eq("organizationId", args.organizationId)
      )
      .unique();
    if (!conversation) {
      const conversationId = await ctx.db.insert("conversations", {
        createdAt: Date.now(),
        organizationId: args.organizationId,
      });
      conversation = await ctx.db.get(conversationId);
    }
    if (!conversation) {
      throw new Error("Conversation could not be created");
    }

    const existing = await ctx.db
      .query("conversationTurns")
      .withIndex("by_conversationId_and_clientRequestId", (q) =>
        q
          .eq("conversationId", conversation._id)
          .eq("clientRequestId", args.clientRequestId)
      )
      .unique();
    if (existing) {
      const user = await ctx.db
        .query("conversationMessages")
        .withIndex("by_turnId_and_role", (q) =>
          q.eq("turnId", existing._id).eq("role", "user")
        )
        .unique();
      return { messageId: user?._id ?? null, turnId: existing._id };
    }

    const createdAt = Date.now();
    const lastMessage = await ctx.db
      .query("conversationMessages")
      .withIndex("by_conversationId_and_ordinal", (q) =>
        q.eq("conversationId", conversation._id)
      )
      .order("desc")
      .first();
    const ordinal = (lastMessage?.ordinal ?? -1) + 1;
    const turnId = await ctx.db.insert("conversationTurns", {
      attemptNumber: 1,
      clientRequestId: args.clientRequestId,
      conversationId: conversation._id,
      createdAt,
      status: "queued",
    });
    const messageId = await ctx.db.insert("conversationMessages", {
      content: args.rawText,
      conversationId: conversation._id,
      createdAt,
      ordinal,
      role: "user",
      turnId,
    });
    for (const image of args.images) {
      await ctx.db.insert("conversationAttachments", {
        ...image,
        createdAt,
        messageId,
      });
    }
    await ctx.db.insert("conversationMessages", {
      content: "",
      conversationId: conversation._id,
      createdAt: createdAt + 1,
      ordinal: ordinal + 1,
      role: "assistant",
      turnId,
    });
    await ctx.scheduler.runAfter(0, runTurnRef, { attempt: 1, turnId });
    return { messageId, turnId };
  },
});

export const listForCurrentOrganization = query({
  args: { organizationId: v.id("organizations") },
  handler: async (ctx, args) => {
    await requireOrganizationMember(ctx, args.organizationId);
    const conversation = await ctx.db
      .query("conversations")
      .withIndex("by_organizationId", (q) =>
        q.eq("organizationId", args.organizationId)
      )
      .unique();
    if (!conversation) {
      return [];
    }
    const messages = await ctx.db
      .query("conversationMessages")
      .withIndex("by_conversationId_and_ordinal", (q) =>
        q.eq("conversationId", conversation._id)
      )
      .order("asc")
      .take(200);
    return await Promise.all(
      messages.map(async (message) => {
        const turn = await ctx.db.get(message.turnId);
        const attachments = await ctx.db
          .query("conversationAttachments")
          .withIndex("by_messageId", (q) => q.eq("messageId", message._id))
          .collect();
        return {
          attachments: await Promise.all(
            attachments.map(async (attachment) => ({
              filename: attachment.filename ?? null,
              id: String(attachment._id),
              mediaType: attachment.mimeType,
              url: await ctx.storage.getUrl(attachment.storageId),
            }))
          ),
          error: turn?.error ?? null,
          messageId: String(message._id),
          rawText: message.content,
          role: message.role,
          status: turn?.status ?? "failed",
        };
      })
    );
  },
});

export const getTurn = internalQuery({
  args: { turnId: v.id("conversationTurns") },
  handler: async (ctx, args) => {
    const turn = await ctx.db.get(args.turnId);
    if (!turn) {
      return null;
    }
    const conversation = await ctx.db.get(turn.conversationId);
    const user = await ctx.db
      .query("conversationMessages")
      .withIndex("by_turnId_and_role", (q) =>
        q.eq("turnId", turn._id).eq("role", "user")
      )
      .unique();
    const assistant = await ctx.db
      .query("conversationMessages")
      .withIndex("by_turnId_and_role", (q) =>
        q.eq("turnId", turn._id).eq("role", "assistant")
      )
      .unique();
    if (!conversation || !user || !assistant) {
      return null;
    }
    const attachments = await ctx.db
      .query("conversationAttachments")
      .withIndex("by_messageId", (q) => q.eq("messageId", user._id))
      .collect();
    return {
      assistant,
      attachments,
      organizationId: conversation.organizationId,
      turn,
      user,
    };
  },
});

export const markProcessing = internalMutation({
  args: {
    attempt: v.number(),
    leaseOwner: v.string(),
    turnId: v.id("conversationTurns"),
  },
  handler: async (ctx, args): Promise<boolean> => {
    const turn = await ctx.db.get(args.turnId);
    if (
      !turn ||
      turn.status !== "queued" ||
      (turn.attemptNumber ?? 1) !== args.attempt
    ) {
      return false;
    }
    await ctx.db.patch(turn._id, {
      error: undefined,
      leaseExpiresAt: Date.now() + 60_000,
      leaseOwner: args.leaseOwner,
      status: "processing",
    });
    return true;
  },
});

export const completeTurn = internalMutation({
  args: {
    attempt: v.number(),
    leaseOwner: v.string(),
    submissionId: v.string(),
    text: v.string(),
    turnId: v.id("conversationTurns"),
  },
  handler: async (ctx, args): Promise<boolean> => {
    const turn = await ctx.db.get(args.turnId);
    if (
      !turn ||
      turn.status !== "processing" ||
      (turn.attemptNumber ?? 1) !== args.attempt ||
      turn.leaseOwner !== args.leaseOwner ||
      (turn.leaseExpiresAt !== undefined && turn.leaseExpiresAt < Date.now())
    ) {
      return false;
    }
    const assistant = await ctx.db
      .query("conversationMessages")
      .withIndex("by_turnId_and_role", (q) =>
        q.eq("turnId", args.turnId).eq("role", "assistant")
      )
      .unique();
    if (assistant) {
      await ctx.db.patch(assistant._id, { content: args.text });
    }
    await ctx.db.patch(args.turnId, {
      completedAt: Date.now(),
      error: undefined,
      leaseExpiresAt: undefined,
      leaseOwner: undefined,
      status: "completed",
      submissionId: args.submissionId,
    });
    return true;
  },
});

export const failTurn = internalMutation({
  args: {
    attempt: v.number(),
    error: v.string(),
    leaseOwner: v.string(),
    retry: v.boolean(),
    turnId: v.id("conversationTurns"),
  },
  handler: async (ctx, args): Promise<boolean> => {
    const turn = await ctx.db.get(args.turnId);
    if (
      !turn ||
      turn.status !== "processing" ||
      (turn.attemptNumber ?? 1) !== args.attempt ||
      turn.leaseOwner !== args.leaseOwner ||
      (turn.leaseExpiresAt !== undefined && turn.leaseExpiresAt < Date.now())
    ) {
      return false;
    }
    await ctx.db.patch(turn._id, {
      attemptNumber: args.retry ? args.attempt + 1 : args.attempt,
      completedAt: args.retry ? undefined : Date.now(),
      error: args.error,
      leaseExpiresAt: undefined,
      leaseOwner: undefined,
      status: args.retry ? "queued" : "failed",
    });
    return true;
  },
});

const toBase64 = (buffer: ArrayBuffer): string => {
  let binary = "";
  for (const byte of new Uint8Array(buffer)) {
    binary += String.fromCodePoint(byte);
  }
  return btoa(binary);
};

export const runTurn = internalAction({
  args: { attempt: v.number(), turnId: v.id("conversationTurns") },
  handler: async (ctx, args): Promise<null> => {
    const leaseOwner = `conversation:${args.turnId}:${args.attempt}`;
    const turn = await ctx.runQuery(getTurnRef, { turnId: args.turnId });
    if (
      !turn ||
      !(await ctx.runMutation(markProcessingRef, {
        attempt: args.attempt,
        leaseOwner,
        turnId: args.turnId,
      }))
    ) {
      return null;
    }
    try {
      const flueUrl = (env as typeof env & { readonly FLUE_URL?: string })
        .FLUE_URL;
      if (!flueUrl) {
        throw new Error("FLUE_URL is not configured in Convex");
      }
      const images = await Promise.all(
        turn.attachments.map(async (attachment) => {
          const url = await ctx.storage.getUrl(attachment.storageId);
          const response = url ? await fetch(url) : null;
          if (!response?.ok) {
            throw new Error("Conversation image could not be loaded");
          }
          return {
            data: toBase64(await response.arrayBuffer()),
            mimeType: attachment.mimeType,
            type: "image" as const,
          };
        })
      );
      const endpoint = new URL(
        `agents/zopu/${encodeURIComponent(String(turn.organizationId))}`,
        `${flueUrl.replace(/\/+$/u, "")}/`
      );
      endpoint.searchParams.set("wait", "result");
      const response = await fetch(endpoint, {
        body: JSON.stringify({ images, message: turn.user.content }),
        headers: {
          authorization: `Bearer ${env.FLUE_DB_TOKEN}`,
          "content-type": "application/json",
          "x-zopu-organization-id": String(turn.organizationId),
          "x-zopu-request-id": turn.turn.clientRequestId,
        },
        method: "POST",
      });
      const payload: unknown = await response.json().catch(() => null);
      if (
        !response.ok ||
        typeof payload !== "object" ||
        payload === null ||
        !("submissionId" in payload) ||
        typeof payload.submissionId !== "string" ||
        !("result" in payload) ||
        typeof payload.result !== "object" ||
        payload.result === null ||
        !("text" in payload.result) ||
        typeof payload.result.text !== "string"
      ) {
        throw new Error(`Flue turn failed (${response.status})`);
      }
      await ctx.runMutation(completeTurnRef, {
        attempt: args.attempt,
        leaseOwner,
        submissionId: payload.submissionId,
        text: payload.result.text,
        turnId: args.turnId,
      });
    } catch (error) {
      const retry = args.attempt < MAX_ATTEMPTS;
      await ctx.runMutation(failTurnRef, {
        attempt: args.attempt,
        error: error instanceof Error ? error.message : String(error),
        leaseOwner,
        retry,
        turnId: args.turnId,
      });
      if (retry) {
        await ctx.scheduler.runAfter(args.attempt * 1000, runTurnRef, {
          attempt: args.attempt + 1,
          turnId: args.turnId,
        });
      }
    }
    return null;
  },
});

export const reconcileExpiredTurns = internalMutation({
  args: {},
  handler: async (ctx): Promise<{ reconciled: number }> => {
    const expired = await ctx.db
      .query("conversationTurns")
      .withIndex("by_status_and_leaseExpiresAt", (q) =>
        q.eq("status", "processing").lt("leaseExpiresAt", Date.now())
      )
      .collect();
    for (const turn of expired) {
      const attempt = turn.attemptNumber ?? 1;
      const retry = attempt < MAX_ATTEMPTS;
      await ctx.db.patch(turn._id, {
        attemptNumber: retry ? attempt + 1 : attempt,
        completedAt: retry ? undefined : Date.now(),
        error: "Conversation worker lease expired",
        leaseExpiresAt: undefined,
        leaseOwner: undefined,
        status: retry ? "queued" : "failed",
      });
      if (retry) {
        await ctx.scheduler.runAfter(0, runTurnRef, {
          attempt: attempt + 1,
          turnId: turn._id,
        });
      }
    }
    return { reconciled: expired.length };
  },
});