cozsweet-frontend-nextjs/src/stores/chat/machine/send-flow.ts

import { type DoneActorEvent, type ErrorActorEvent } from "xstate";

import { todayString } from "@/utils/date";
import { Logger } from "@/utils/logger";

import {
  applyHttpSendOutput,
  beginPendingReply,
  finishPendingReply,
  type HttpSendOutput,
} from "../helper/send-state";
import { createClientUiMessageIdentity } from "../ui-message";
import { historyMachineSetup } from "./history-flow";
import { createChatActorActionSetup } from "./setup";

const sendLog = new Logger("StoresChatChatSendFlow");
const mediaLog = new Logger("StoresChatChatMediaFlow");

const enqueueMessageAction = historyMachineSetup.sendTo(
  "messageQueue",
  ({ event }) => event,
);

const appendGuestUserMessageAction = historyMachineSetup.assign(
  ({ context, event }) => {
    if (event.type !== "ChatSendMessage") return {};

    const today = todayString();
    sendLog.debug("[chat-machine] appendGuestUserMessage", {
      contentLength: event.content.length,
      contentPreview: event.content.slice(0, 50),
      quotaMode: "server controlled",
      isReplyingAI: context.isReplyingAI,
    });

    return {
      messages: [
        ...context.messages,
        {
          ...createClientUiMessageIdentity("message"),
          content: event.content,
          isFromAI: false,
          date: today,
        },
      ],
      outgoingMessageRevision: context.outgoingMessageRevision + 1,
      upgradePromptVisible: false,
    upgradeReason: null,
    paymentGuidance: null,
    };
  },
);

const appendUserMessageAction = historyMachineSetup.assign(
  ({ context, event }) => {
    if (event.type !== "ChatSendMessage") return {};

    const today = todayString();
    sendLog.debug("[chat-machine] appendUserMessage", {
      contentLength: event.content.length,
      contentPreview: event.content.slice(0, 50),
      isReplyingAI: context.isReplyingAI,
    });

    return {
      messages: [
        ...context.messages,
        {
          ...createClientUiMessageIdentity("message"),
          content: event.content,
          isFromAI: false,
          date: today,
        },
      ],
      outgoingMessageRevision: context.outgoingMessageRevision + 1,
      upgradePromptVisible: false,
    upgradeReason: null,
    paymentGuidance: null,
    };
  },
);

const appendQueuedSendErrorMessageAction = historyMachineSetup.assign(
  ({ context, event }) => {
    if (event.type !== "ChatQueuedSendError") return {};
    const characterErrorCode = event.characterErrorCode ?? null;
    const unavailable = characterErrorCode === "CHARACTER_STATE_UNAVAILABLE";
    const messages = [
      ...context.messages,
      {
        ...createClientUiMessageIdentity("error"),
        content: unavailable
          ? "This character is temporarily unavailable. Please try again shortly."
          : "Something went wrong. Try sending again?",
        isFromAI: true,
        date: todayString(),
      },
    ];
    return {
      messages,
      ...finishPendingReply(context),
      characterErrorCode,
      canSendMessage:
        characterErrorCode === "CHARACTER_DISABLED" ||
        characterErrorCode === "CHARACTER_NOT_FOUND"
          ? false
          : context.canSendMessage,
    };
  },
);

const markQueuedSendStartedAction = historyMachineSetup.assign(
  ({ context }) => beginPendingReply(context),
);

const applyQueuedHttpOutputAction = historyMachineSetup.assign(
  ({ context, event }) => {
    if (event.type !== "ChatQueuedHttpDone") return {};
    return applyHttpSendOutput(context, event.output);
  },
);

const clearUpgradePromptAction = historyMachineSetup.assign(() => ({
  upgradePromptVisible: false,
  upgradeReason: null,
  canSendMessage: true,
  requiredCredits: 0,
  shortfallCredits: 0,
  paymentGuidance: null,
}));

const appendGuestUserImageAction = historyMachineSetup.assign(
  ({ context, event }) => {
    if (event.type !== "ChatSendImage") return {};

    const today = todayString();
    mediaLog.debug("[chat-machine] appendGuestUserImage", {
      oldMessagesCount: context.messages.length,
      quotaMode: "server controlled",
    });
    return {
      messages: [
        ...context.messages,
        {
          ...createClientUiMessageIdentity("image"),
          content: "[Image]",
          isFromAI: false,
          date: today,
          imageUrl: event.imageBase64,
        },
      ],
      outgoingMessageRevision: context.outgoingMessageRevision + 1,
      ...beginPendingReply(context),
      upgradePromptVisible: false,
      upgradeReason: null,
    };
  },
);

const appendUserImageAction = historyMachineSetup.assign(
  ({ context, event }) => {
    if (event.type !== "ChatSendImage") return {};

    const today = todayString();
    mediaLog.debug("[chat-machine] appendUserImage", {
      oldMessagesCount: context.messages.length,
    });

    return {
      messages: [
        ...context.messages,
        {
          ...createClientUiMessageIdentity("image"),
          content: "[Image]",
          isFromAI: false,
          date: today,
          imageUrl: event.imageBase64,
        },
      ],
      outgoingMessageRevision: context.outgoingMessageRevision + 1,
      ...beginPendingReply(context),
      upgradePromptVisible: false,
      upgradeReason: null,
    };
  },
);

export const sendMachineSetup = historyMachineSetup.extend({
  actions: {
    enqueueMessage: enqueueMessageAction,
    appendGuestUserMessage: appendGuestUserMessageAction,
    appendUserMessage: appendUserMessageAction,
    appendQueuedSendErrorMessage: appendQueuedSendErrorMessageAction,
    markQueuedSendStarted: markQueuedSendStartedAction,
    applyQueuedHttpOutput: applyQueuedHttpOutputAction,
    clearUpgradePrompt: clearUpgradePromptAction,
    appendGuestUserImage: appendGuestUserImageAction,
    appendUserImage: appendUserImageAction,
  },
});

const sendMessageDoneActionSetup =
  createChatActorActionSetup<DoneActorEvent<HttpSendOutput>>();
const actorErrorActionSetup = createChatActorActionSetup<ErrorActorEvent>();

const applyHttpSendOutputAction = sendMessageDoneActionSetup.assign(
  ({ context, event }) => applyHttpSendOutput(context, event.output),
);

const markHttpSendFailedAction = actorErrorActionSetup.assign({
  isReplyingAI: false,
});

export const guestReadyState = sendMachineSetup.createStateConfig({
  on: {
    ChatSendMessage: {
      actions: ["appendGuestUserMessage", "enqueueMessage"],
      guard: ({ context, event }) =>
        context.canSendMessage && event.content.trim().length > 0,
    },
    ChatSendImage: {
      actions: "appendGuestUserImage",
      target: "sending",
    },
  },
});

export const guestSendingState = sendMachineSetup.createStateConfig({
  invoke: {
    src: "sendMessageHttp",
    input: ({ context, event }) => ({
      characterId: context.characterId,
      content: event.type === "ChatSendMessage" ? event.content : "",
    }),
    onDone: {
      target: "ready",
      actions: applyHttpSendOutputAction,
    },
    onError: {
      target: "ready",
      actions: markHttpSendFailedAction,
    },
  },
});

export const userSendingViaHttpState = sendMachineSetup.createStateConfig({
  invoke: {
    src: "sendMessageHttp",
    input: ({ context, event }) => ({
      characterId: context.characterId,
      content: event.type === "ChatSendMessage" ? event.content : "",
    }),
    onDone: {
      target: "ready",
      actions: applyHttpSendOutputAction,
    },
    onError: {
      target: "ready",
      actions: markHttpSendFailedAction,
    },
  },
});