storytellurx-v1/app/api/chat/route.ts

import { AI_PERSONAS, AIPersonaId } from "../../game-engine/personas.ts";
import {
  getModelDefinition,
  type ModelId,
} from "../../game-engine/models-server.ts";
import {
  EroticNarrativeResponseSchema,
  MessageMetadata,
  SystemNarrativeResponseSchema,
  VisualPromptInstructionSchema,
} from "../../game-engine/game-schemas.ts";
import { I18NLocale } from "../../actions/locale.ts";
import { constructSystemPrompt } from "../../game-engine/prompt-constructor.ts";
import {
  convertToModelMessages,
  createUIMessageStream,
  createUIMessageStreamResponse,
  Output,
  streamText,
} from "ai";
import {
  addFullGameStateAsync,
  getGameState,
} from "../../game-engine/game-state.ts";
import type { NarrativeMessage } from "../../game-engine/narrative-provider.tsx";
import {
  applyMindControlRewrites,
  includeUserHintForSystem,
  inludeSystemAsAssistant,
} from "../../game-engine/generation-utils.ts";
import {
  generateImageAction,
  type ImageGenerationPayload,
} from "../../actions/generate-image.ts";

export const maxDuration = 60;

export type ErrorCode =
  | "QUOTA_EXCEEDED"
  | "PROVIDER_UNAVAILABLE"
  | "MODEL_NOT_FOUND"
  | "UPSTREAM_ERROR"
  | "UNKNOWN";

export interface StreamErrorPayload {
  code: ErrorCode;
  retryable: boolean;
  retryAfterSeconds?: number;
  quotaModel?: string;
  quotaLimit?: number;
}

function extractQuotaDetails(error: unknown): {
  retryAfterSeconds?: number;
  quotaModel?: string;
  quotaLimit?: number;
} {
  const msg = String(error);

  // Extract "Please retry in 47.747102362s"
  const retryMatch = msg.match(/retry in (\d+(?:\.\d+)?)s/i);
  const retryAfterSeconds = retryMatch
    ? Math.ceil(parseFloat(retryMatch[1]))
    : undefined;

  // Extract "limit: 20, model: gemini-2.5-flash"
  const limitMatch = msg.match(/limit:\s*(\d+)/);
  const modelMatch = msg.match(/model:\s*([\w.-]+)/);

  return {
    retryAfterSeconds,
    quotaModel: modelMatch?.[1],
    quotaLimit: limitMatch ? parseInt(limitMatch[1]) : undefined,
  };
}

function classifyError(error: unknown): StreamErrorPayload {
  const msg = String(error);
  if (
    msg.includes("quota") || msg.includes("rate limit") ||
    msg.includes("429") || msg.includes("RESOURCE_EXHAUSTED")
  ) {
    return {
      code: "QUOTA_EXCEEDED",
      retryable: false,
      ...extractQuotaDetails(error),
    };
  }
  if (msg.includes("Forbidden") || msg.includes("not available") || msg.includes("403")) {
    return { code: "PROVIDER_UNAVAILABLE", retryable: false };
  }
  if (msg.includes("Not Found") || msg.includes("404")) {
    return { code: "MODEL_NOT_FOUND", retryable: false };
  }
  if (msg.includes("internal error") || msg.includes("RetryError") || msg.includes("500")) {
    return { code: "UPSTREAM_ERROR", retryable: true };
  }
  return { code: "UNKNOWN", retryable: true };
}

function isAbortError(error: unknown): boolean {
  const msg = String(error);
  return msg.includes("AbortError") || msg.includes("cancelled") || msg.includes("aborted");
}

export interface GenerateRequestBody {
  id: string;
  messages: NarrativeMessage[];
  locale: I18NLocale;
  modelId: ModelId;
}

const handleGeneration = async (
  request: Request,
  body: GenerateRequestBody,
) => {
  const { messages: originalMessages, locale, modelId } = body;
  const enrichedMessages = await addFullGameStateAsync(originalMessages);
  const currentGameState = getGameState(enrichedMessages);
  const messagesForLLM = applyMindControlRewrites(
    originalMessages,
    enrichedMessages,
  );

  const actingPersonaId: AIPersonaId =
    currentGameState.personaId === "storytellurx_pause"
      ? "storytellurx_pause"
      : currentGameState.personaId ?? "storytellurx_gateway";

  const persona = AI_PERSONAS.find((p) => p.id === actingPersonaId) ??
    AI_PERSONAS.find((p) => p.id === "storytellurx_gateway")!;
  const schemaForThisTurn = persona.type === "system"
    ? SystemNarrativeResponseSchema
    : EroticNarrativeResponseSchema;

  const { getModel, features, parameters } = getModelDefinition(modelId);

  const systemInstruction = await constructSystemPrompt(
    currentGameState,
    actingPersonaId,
    locale,
    features,
  );

  const model = getModel();

  const result = streamText({
    model,
    abortSignal: request.signal,
    messages: convertToModelMessages(
      includeUserHintForSystem(
        features?.developerInstructionDisabled
          ? inludeSystemAsAssistant(
            messagesForLLM,
            systemInstruction,
            schemaForThisTurn,
            features,
          )
          : messagesForLLM,
        features,
      ),
    ),
    system: features?.developerInstructionDisabled
      ? undefined
      : systemInstruction,
    experimental_output: features?.structuredOutputDisabled
      ? undefined
      : Output.object({ schema: schemaForThisTurn }),
    maxRetries: 1,
    ...parameters ?? {},
  });

  result.consumeStream().catch((error) => {
    if (!isAbortError(error)) {
      console.error("consumeStream error:", classifyError(error).code);
    }
  });

  return createUIMessageStreamResponse({
    stream: createUIMessageStream({
      execute: ({ writer }) => {
        const imageStream = createUIMessageStream({
          execute: async ({ writer }) => {
            let instructionTriggered = false;
            try {
              for await (
                const partial of result.experimental_partialOutputStream
              ) {
                if (instructionTriggered) continue;
                if (
                  "visualPromptInstruction" in partial &&
                  partial.visualPromptInstruction &&
                  ([partial.gameStateDiff?.modality, currentGameState.modality]
                    .includes("text-image"))
                ) {
                  const validation = VisualPromptInstructionSchema.safeParse(
                    partial.visualPromptInstruction?.slice(0, 99),
                  );
                  console.log(validation);
                  if (validation.success) {
                    instructionTriggered = true;
                    writer.write({
                      type: "tool-call",
                      toolCallId: "image",
                      toolName: "generateImageAction",
                      args: { isLoadingImage: true },
                    });
                    const simpleHistory = messagesForLLM.map((m) => ({
                      role: m.role,
                      content: m.parts.filter((p) => p.type === "text")
                        .map((p) => p.text)
                        .join(""),
                    }));
                    const imageResult = await generateImageAction(
                      validation.data,
                      simpleHistory,
                    );
                    const dataPayload: ImageGenerationPayload = {
                      isLoadingImage: false,
                    };
                    if ("data" in imageResult && imageResult.data.output) {
                      dataPayload.imageUrl = imageResult.data.output
                          .startsWith("data:")
                        ? imageResult.data.output
                        : `data:image/jpeg;base64,${imageResult.data.output}`;
                    }
                    writer.write({
                      type: "tool-result",
                      toolCallId: "image",
                      result: dataPayload,
                    });
                  }
                }
              }
            } catch (error) {
              console.error(
                "Error processing partial output stream for image generation:",
                error,
              );
              writer.write({
                type: "tool-result",
                toolCallId: "image",
                result: { error },
              });
            }
          },
        });

        const stream = result.toUIMessageStream({
          originalMessages,
          onError: (error) => {
            if (isAbortError(error)) return "";
            const classified = classifyError(error);
            console.error(`[${classified.code}]`, error);
            return JSON.stringify(classified);
          },
          messageMetadata: ({ part }) => {
            if (part.type === "finish") {
              return {
                model: {
                  id: typeof model === "string" ? model : model.modelId,
                  parameters: features,
                },
                createdAt: Date.now(),
                system: { tokenCount: Date.now(), content: systemInstruction },
                totalTokenCount: part.totalUsage.totalTokens ?? 1,
                tokenCount: 1,
              } as Partial<MessageMetadata>;
            }
          },
        });

        try {
          writer.merge(imageStream);
          writer.merge(stream);
        } catch {
          // client disconnected — stream destination closed
        }
      },
    }),
  });
};

export async function POST(req: Request) {
  try {
    const body: GenerateRequestBody = await req.json();
    return await handleGeneration(req, body);
  } catch (error) {
    if (isAbortError(error)) return new Response(null, { status: 499 });
    const classified = classifyError(error);
    console.error(`[${classified.code}] Chat API Error:`, error);
    return Response.json({ error: classified }, {
      status: 500,
    });
  }
}