growqr-backend/src/actors/conversation/conversation-actor.ts

import { asc, eq } from "drizzle-orm";
import { actor, event, queue } from "rivetkit";
import { db as drizzleDb } from "rivetkit/db/drizzle";
import { streamConversationResponse } from "./agent.js";
import { migrateConversationDb } from "./migrations.js";
import {
  conversationMessages,
  conversationSchema,
} from "./schema.js";
import type {
  ConversationMessage,
  ConversationQueueMessage,
  ConversationResponseEvent,
  ConversationStatus,
} from "./types.js";

const buildId = (prefix: string) =>
  `${prefix}-${Date.now()}-${Math.random().toString(16).slice(2)}`;

function now() {
  return Date.now();
}

function conversationIdFromKey(key: unknown[]) {
  return String(key[1] ?? key[0] ?? "default");
}

function userIdFromKey(key: unknown[]) {
  return String(key[0] ?? "");
}

function toPublicMessage(row: typeof conversationMessages.$inferSelect): ConversationMessage {
  return {
    id: row.id,
    conversationId: row.conversationId,
    role: row.role,
    sender: row.sender,
    content: row.content,
    createdAt:
      row.createdAt instanceof Date ? row.createdAt.getTime() : Number(row.createdAt),
  };
}

export const conversationActor = actor({
  // Keep only small runtime state here. Message history lives in actor-local
  // SQLite via Drizzle so this actor can grow without bloating c.state.
  state: {
    status: {
      state: "idle",
      updatedAt: Date.now(),
    } as ConversationStatus,
  },

  db: drizzleDb({
    schema: conversationSchema,
  }),

  queues: {
    message: queue<ConversationQueueMessage>(),
  },

  events: {
    messageAdded: event<ConversationMessage>(),
    status: event<ConversationStatus>(),
    response: event<ConversationResponseEvent>(),
  },

  onCreate: async (c) => {
    await migrateConversationDb(c.db);
  },

  onWake: async (c) => {
    await migrateConversationDb(c.db);
  },

  run: async (c) => {
    for await (const queued of c.queue.iter()) {
      const { body } = queued;
      if (!body?.text || typeof body.text !== "string") continue;

      const conversationId = conversationIdFromKey(c.key);
      const sender = body.sender?.trim() || "User";

      const userMessage: ConversationMessage = {
        id: buildId("user"),
        conversationId,
        role: "user",
        sender,
        content: body.text.trim(),
        createdAt: now(),
      };

      await c.db.insert(conversationMessages).values({
        ...userMessage,
        createdAt: new Date(userMessage.createdAt),
      });
      c.broadcast("messageAdded", userMessage);

      const historyRows = await c.db
        .select()
        .from(conversationMessages)
        .where(eq(conversationMessages.conversationId, conversationId))
        .orderBy(asc(conversationMessages.createdAt));
      const history = historyRows.map(toPublicMessage);

      const assistantMessage: ConversationMessage = {
        id: buildId("assistant"),
        conversationId,
        role: "assistant",
        sender: "GrowQR",
        content: "",
        createdAt: now(),
      };

      await c.db.insert(conversationMessages).values({
        ...assistantMessage,
        createdAt: new Date(assistantMessage.createdAt),
      });
      c.broadcast("messageAdded", assistantMessage);

      c.state.status = { state: "thinking", updatedAt: now() };
      c.broadcast("status", c.state.status);

      try {
        const result = streamConversationResponse(history, {
          userId: body.context?.userId ?? userIdFromKey(c.key),
          conversationId,
          missionInstanceId: body.context?.missionInstanceId,
          missionId: body.context?.missionId,
          stageId: body.context?.stageId,
          source: body.context?.source,
        });

        let content = "";
        for await (const delta of result.textStream) {
          if (c.aborted) break;

          content += delta;
          c.broadcast("response", {
            messageId: assistantMessage.id,
            delta,
            content,
            done: false,
          });
        }

        assistantMessage.content = content || assistantMessage.content;
        await c.db
          .update(conversationMessages)
          .set({ content: assistantMessage.content })
          .where(eq(conversationMessages.id, assistantMessage.id));

        c.broadcast("response", {
          messageId: assistantMessage.id,
          delta: "",
          content: assistantMessage.content,
          done: true,
        });

        c.state.status = { state: "idle", updatedAt: now() };
        c.broadcast("status", c.state.status);
      } catch (error) {
        const errorMessage =
          error instanceof Error ? error.message : "Unknown conversation error";

        assistantMessage.content =
          assistantMessage.content ||
          "I hit a snag while responding. Please try again.";

        await c.db
          .update(conversationMessages)
          .set({ content: assistantMessage.content })
          .where(eq(conversationMessages.id, assistantMessage.id));

        c.state.status = {
          state: "error",
          updatedAt: now(),
          error: errorMessage,
        };

        c.broadcast("response", {
          messageId: assistantMessage.id,
          delta: "",
          content: assistantMessage.content,
          done: true,
          error: errorMessage,
        });
        c.broadcast("status", c.state.status);
      }
    }
  },

  actions: {
    sendMessage: async (c, input: ConversationQueueMessage) => {
      await c.queue.send("message", input);
      return { queued: true };
    },

    getHistory: async (c): Promise<ConversationMessage[]> => {
      const conversationId = conversationIdFromKey(c.key);
      const rows = await c.db
        .select()
        .from(conversationMessages)
        .where(eq(conversationMessages.conversationId, conversationId))
        .orderBy(asc(conversationMessages.createdAt));

      return rows.map(toPublicMessage);
    },

    addMessage: async (
      c,
      input: { role: "user" | "assistant"; content: string; sender?: string; id?: string },
    ): Promise<ConversationMessage> => {
      const conversationId = conversationIdFromKey(c.key);
      const message: ConversationMessage = {
        id: input.id ?? buildId(input.role),
        conversationId,
        role: input.role,
        sender: input.sender?.trim() || (input.role === "assistant" ? "GrowQR" : "User"),
        content: input.content,
        createdAt: now(),
      };

      await c.db.insert(conversationMessages).values({
        ...message,
        createdAt: new Date(message.createdAt),
      });
      c.broadcast("messageAdded", message);
      return message;
    },

    clearHistory: async (c) => {
      const conversationId = conversationIdFromKey(c.key);
      await c.db
        .delete(conversationMessages)
        .where(eq(conversationMessages.conversationId, conversationId));
      c.state.status = { state: "idle", updatedAt: now() };
      c.broadcast("status", c.state.status);
      return { ok: true };
    },

    getStatus: (c): ConversationStatus => c.state.status,
  },
});