storytellurx-v1/app/game-engine/narrative-provider.tsx

"use client";

import type {
  ErrorCode,
  GenerateRequestBody,
  StreamErrorPayload,
} from "../api/chat/route.ts";

import { type UIMessage, useChat } from "@ai-sdk/react";
import {
  type DataPartSchemas,
  dataPartSchemas,
  type GameState,
  MessageMetadata,
  MessageMetadataSchema,
} from "../game-engine/game-schemas.ts";
import {
  createContext,
  type ReactNode,
  useCallback,
  useContext,
  useEffect,
  useMemo,
  useRef,
  useState,
} from "react";
import { useChatHistory } from "../hooks/use-chat-history.ts";
import {
  addFullGameStateAsync,
  addFullGameStateSync,
  getGameState,
} from "../game-engine/game-state.ts";
import {
  DefaultChatTransport,
  type InferUIDataParts,
  parsePartialJson,
} from "ai";
import { useAsyncMemo } from "../hooks/use-async-memo.ts";
import {
  AVAILABLE_CLIENT_MODELS,
  mockFetchToLocalAPI,
} from "./models-client.ts";

type NarrativeContextValue =
  & ReturnType<
    typeof useChat<unknown, DataPartSchemas>
  >
  & {
    isNewChat: boolean;
    locale: GenerateRequestBody["locale"];
    gameState: GameState;
    uiVariant: "classical" | "terminal";
    localModelProgress?: { text: string; progress: number };
    errorCode: ErrorCode | "PARSE_FAILED" | null;
    errorDetails: StreamErrorPayload | null;
    clearError: () => void;
    autoRetryAttempt: number;
    autoRetryMax: number;
  };

export const NarrativeContext = createContext<
  NarrativeContextValue | undefined
>(
  undefined,
);

export type NarrativeMessage = UIMessage<
  MessageMetadata,
  InferUIDataParts<DataPartSchemas>
>;

export interface NarrativeProviderProps {
  children: ReactNode;
  apiEndpoint: string;
  initialMessages: NarrativeMessage[];
  initialGameState: GameState;
  locale: GenerateRequestBody["locale"];
  sessionId: string;
  modelId?: string;
}

export const NarrativeProvider = ({
  children,
  apiEndpoint,
  initialMessages,
  locale,
  sessionId,
  modelId,
}: NarrativeProviderProps) => {
  const [localModelProgress, setLocalModelProgress] = useState<
    { text: string; progress: number }
  >();
  const [errorCode, setErrorCode] = useState<
    ErrorCode | "PARSE_FAILED" | null
  >(null);
  const clearError = useCallback(() => setErrorCode(null), []);
  const errorCodeRef = useRef(errorCode);
  errorCodeRef.current = errorCode;

  const chat = useChat({
    id: sessionId,
    messages: initialMessages,
    dataPartSchemas,
    messageMetadataSchema: MessageMetadataSchema,
    transport: new DefaultChatTransport({
      api: apiEndpoint,
      fetch: (input, init) =>
        AVAILABLE_CLIENT_MODELS.find((v) => v.id === modelId)
          ? mockFetchToLocalAPI(input, init, {
            onProgress: setLocalModelProgress,
          })
          : globalThis.fetch(input, init),
      prepareRequest: (
        { id, messages, credentials, headers, body, requestMetadata: _ },
      ) => ({
        credentials,
        headers,
        body: { ...body, id, messages, locale, modelId } as GenerateRequestBody,
      }),
    }),
    onFinish: async ({ message }) => {
      // Clear any previous error on successful finish
      if (errorCodeRef.current === "PARSE_FAILED") setErrorCode(null);

      const text = message.parts.filter((v) => v.type === "text").at(0)?.text;
      const value = await parsePartialJson(
        message.parts.filter((v) => v.type === "text").at(0)?.text,
      );
      if (text && value.state !== "failed-parse") {
        await saveMessage(sessionId, message);
      } else {
        setErrorCode("PARSE_FAILED");
        chat.setMessages(chat.messages.slice(0, -1));
      }
    },
  });

  // Extract structured error code and details from chat.error
  const { derivedErrorCode, errorDetails } = useMemo(() => {
    if (errorCode) return { derivedErrorCode: errorCode, errorDetails: null };
    if (chat.error) {
      try {
        const parsed: StreamErrorPayload = JSON.parse(chat.error.message);
        if (parsed.code) {
          return { derivedErrorCode: parsed.code, errorDetails: parsed };
        }
      } catch {
        const msg = chat.error.message;
        if (msg.includes("QUOTA_EXCEEDED")) {
          return {
            derivedErrorCode: "QUOTA_EXCEEDED" as const,
            errorDetails: null,
          };
        }
        if (msg.includes("PROVIDER_UNAVAILABLE")) {
          return {
            derivedErrorCode: "PROVIDER_UNAVAILABLE" as const,
            errorDetails: null,
          };
        }
        if (msg.includes("MODEL_NOT_FOUND")) {
          return {
            derivedErrorCode: "MODEL_NOT_FOUND" as const,
            errorDetails: null,
          };
        }
        if (msg.includes("UPSTREAM_ERROR")) {
          return {
            derivedErrorCode: "UPSTREAM_ERROR" as const,
            errorDetails: null,
          };
        }
      }
      return { derivedErrorCode: "UNKNOWN" as const, errorDetails: null };
    }
    return { derivedErrorCode: null, errorDetails: null };
  }, [chat.error, errorCode]);

  // Auto-retry with exponential backoff for transient upstream errors
  const AUTO_RETRY_MAX = 3;
  const [autoRetryAttempt, setAutoRetryAttempt] = useState(0);

  useEffect(() => {
    if (derivedErrorCode !== "UPSTREAM_ERROR") {
      // Reset counter when error clears or changes to a non-retryable type
      if (autoRetryAttempt > 0) setAutoRetryAttempt(0);
      return;
    }
    if (autoRetryAttempt >= AUTO_RETRY_MAX) return; // exhausted — show error

    const delay = Math.pow(2, autoRetryAttempt) * 1500; // 1.5s, 3s, 6s
    const timer = setTimeout(() => {
      setAutoRetryAttempt((prev) => prev + 1);
      chat.reload();
    }, delay);
    return () => clearTimeout(timer);
  }, [derivedErrorCode, autoRetryAttempt]);

  // Reset retry counter on successful completion
  useEffect(() => {
    if (chat.status === "ready" && !chat.error && autoRetryAttempt > 0) {
      setAutoRetryAttempt(0);
    }
  }, [chat.status, chat.error]);

  // For UI: suppress error display while auto-retrying
  const visibleErrorCode = derivedErrorCode === "UPSTREAM_ERROR" &&
      autoRetryAttempt < AUTO_RETRY_MAX
    ? null
    : derivedErrorCode;

  const { saveMessage } = useChatHistory<NarrativeMessage>({
    onHistoryInit: (history) => {
      if (history?.[sessionId]) chat.setMessages(history?.[sessionId]);
    },
  });

  const initialMessages2 = useMemo(
    () => addFullGameStateSync(initialMessages),
    [initialMessages],
  );

  const messages = useAsyncMemo(
    () => addFullGameStateAsync(chat.messages),
    [chat.messages],
    initialMessages2,
  );

  const value = useMemo(() => {
    const gameState = getGameState(messages ?? []);
    return ({
      ...chat,
      messages,
      localModelProgress: chat.status !== "ready"
        ? localModelProgress
        : undefined,
      locale,
      uiVariant: gameState.personaId !== "storytellurx_gateway"
        ? "classical"
        : "terminal",
      gameState,
      isNewChat: chat.messages.length === 2,
      errorCode: visibleErrorCode,
      errorDetails,
      clearError,
      autoRetryAttempt,
      autoRetryMax: AUTO_RETRY_MAX,
    });
  }, [chat, messages, localModelProgress, locale, visibleErrorCode, clearError, autoRetryAttempt, errorDetails]);

  return (
    <NarrativeContext.Provider value={value}>
      {children}
    </NarrativeContext.Provider>
  );
};

export const useNarrative = <T extends boolean = false>(
  props?: { allowNullable?: T },
): T extends true ? NarrativeContextValue | null : NarrativeContextValue => {
  const context = useContext(NarrativeContext);

  if (context === undefined && !props?.allowNullable) {
    throw new Error("useNarrative must be used within a NarrativeProvider");
  }

  // Explicitly handle the undefined case by converting it to null
  return (context === undefined ? null : context) as T extends true
    ? NarrativeContextValue | null
    : NarrativeContextValue;
};