gongxue-base/apps/admin/src/components/AiChat/provider.ts

import {
  AbstractChatProvider,
  XRequest,
  type TransformMessage,
  type XRequestOptions,
} from '@ant-design/x-sdk';
import { usePermissionStore } from '../../store/permission/permissionStore';
import { useUserStore } from '../../store/user/userStore';
import type { AiArtifactSchema, AiChatInput, AiChatMessage, AiReviewSchema, AiSseChunk } from './types';
import { emptyAssistant, parseSsePayload, reduceAiSseMessage } from './sseReducer';

export { parseSsePayload, reduceAiSseMessage };

export async function authenticatedFetch(
  input: RequestInfo | URL,
  init?: RequestInit,
): Promise<Response> {
  const headers = new Headers(init?.headers);
  const token = useUserStore.getState().token;
  if (token) headers.set('Authorization', `Bearer ${token}`);
  headers.set('Accept', 'text/event-stream');
  let requestInput = input;
  let requestInit = init;
  if (typeof init?.body === 'string') {
    try {
      const body = JSON.parse(init.body) as AiChatInput;
      if (body.regenerateMessageId) {
        requestInput = `${String(input).replace(/\/stream$/, '')}/messages/${body.regenerateMessageId}/regenerate/stream`;
        requestInit = {
          ...init,
          body: JSON.stringify({
            clientRequestId: body.clientRequestId,
            reasoningEffort: body.reasoningEffort,
          }),
        };
      } else if (body.editMessageId) {
        requestInput = `${String(input).replace(/\/stream$/, '')}/messages/${body.editMessageId}/edit/stream`;
        requestInit = {
          ...init,
          body: JSON.stringify({
            content: body.message,
            clientRequestId: body.clientRequestId,
            reasoningEffort: body.reasoningEffort,
          }),
        };
      } else if (body.formSubmission) {
        requestInput = `${String(input).replace(/\/conversations\/\d+\/stream$/, '')}/forms/${body.formSubmission.formId}/submit/stream`;
        requestInit = {
          ...init,
          body: JSON.stringify({
            clientRequestId: body.clientRequestId,
            values: body.formSubmission.values,
            reasoningEffort: body.reasoningEffort,
          }),
        };
      } else if (body.reviewSubmission) {
        requestInput = `${String(input).replace(/\/conversations\/\d+\/stream$/, '')}/reviews/${body.reviewSubmission.reviewId}/submit/stream`;
        requestInit = {
          ...init,
          body: JSON.stringify({
            clientRequestId: body.clientRequestId,
            reasoningEffort: body.reasoningEffort,
          }),
        };
      } else {
        const {
          localAttachments: _localAttachments,
          reloadMessage: _reloadMessage,
          regenerateMessageId: _regenerateMessageId,
          editMessageId: _editMessageId,
          formSubmission: _formSubmission,
          reviewSubmission: _reviewSubmission,
          ...payload
        } = body;
        requestInit = { ...init, body: JSON.stringify(payload) };
      }
    } catch {
      requestInit = init;
    }
  }
  const response = await fetch(requestInput, { ...requestInit, headers });
  if (response.status === 401) {
    // 提示由登录页读取展示：直接弹 toast 会被跳转销毁
    sessionStorage.setItem('login_expired_hint', '1');
    useUserStore.getState().logout();
    usePermissionStore.getState().clearPermissions();
    window.location.href = '/login';
  }
  return response;
}

export class GongxueAiChatProvider extends AbstractChatProvider<
  AiChatMessage,
  AiChatInput,
  AiSseChunk
> {
  /** Routes events that target another (already streamed) message. */
  onExternalReview?: (messageId: number, review: AiReviewSchema) => void;
  onExternalArtifact?: (messageId: number, artifact: AiArtifactSchema) => void;

  constructor(url: string, onSettled?: (result?: { ok: boolean; aborted?: boolean }) => void) {
    super({
      request: XRequest<AiChatInput, AiSseChunk, AiChatMessage>(url, {
        manual: true,
        fetch: authenticatedFetch,
        timeout: 15_000,
        streamTimeout: 1_800_000,
        callbacks: {
          onUpdate: () => undefined,
          onSuccess: () => onSettled?.({ ok: true }),
          onError: (error) =>
            onSettled?.({
              ok: false,
              aborted: error?.name === 'AbortError',
            }),
        },
      }),
    });
  }

  transformParams(
    requestParams: Partial<AiChatInput>,
    options: XRequestOptions<AiChatInput, AiSseChunk, AiChatMessage>,
  ): AiChatInput {
    return {
      ...options.params,
      message: requestParams.message?.trim() || '',
      attachmentIds: requestParams.attachmentIds ?? [],
      skillKey: requestParams.skillKey ?? null,
      clientRequestId: requestParams.clientRequestId || crypto.randomUUID(),
      reasoningEffort: requestParams.reasoningEffort,
      localAttachments: requestParams.localAttachments,
      formSubmission: requestParams.formSubmission,
      reviewSubmission: requestParams.reviewSubmission,
      regenerateMessageId: requestParams.regenerateMessageId,
      editMessageId: requestParams.editMessageId,
      reloadMessage: requestParams.reloadMessage,
    };
  }

  transformLocalMessage(requestParams: Partial<AiChatInput>): AiChatMessage | AiChatMessage[] {
    if (requestParams.editMessageId) {
      // 编辑消息不需要新增用户气泡，store 里已原位更新原消息。
      return [];
    }
    if (requestParams.formSubmission) {
      return {
        role: 'user',
        content: '',
        reasoningContent: '',
        toolRuns: [],
        attachments: requestParams.localAttachments ?? [],
        metadata: {
          a2uiSubmit: {
            formTitle: requestParams.formSubmission.formTitle,
          },
        },
      };
    }
    if (requestParams.reviewSubmission) {
      return {
        role: 'user',
        content: '',
        reasoningContent: '',
        toolRuns: [],
        attachments: requestParams.localAttachments ?? [],
        metadata: {
          a2uiReviewSubmit: {
            reviewTitle: requestParams.reviewSubmission.reviewTitle,
          },
        },
      };
    }
    return {
      role: 'user',
      content: requestParams.message?.trim() || '',
      reasoningContent: '',
      toolRuns: [],
      attachments: requestParams.localAttachments ?? [],
    };
  }

  transformMessage(info: TransformMessage<AiChatMessage, AiSseChunk>): AiChatMessage {
    const { event, payload } = parseSsePayload(info.chunk);
    if (
      event === 'ui.artifact' &&
      payload.artifact &&
      typeof payload.messageId === 'number' &&
      info.originMessage?.id !== payload.messageId
    ) {
      this.onExternalArtifact?.(payload.messageId, payload.artifact);
      return info.originMessage ?? emptyAssistant();
    }
    if (
      event === 'ui.form' &&
      payload.form &&
      typeof payload.messageId === 'number' &&
      info.originMessage?.id !== payload.messageId
    ) {
      this.onExternalArtifact?.(payload.messageId, {
        id: payload.form.id,
        type: 'form',
        status: payload.form.status ?? 'pending',
        messageId: payload.messageId,
        payload: payload.form,
      });
      return info.originMessage ?? emptyAssistant();
    }
    if (
      event === 'ui.review' &&
      payload.review &&
      typeof payload.messageId === 'number' &&
      info.originMessage?.id !== payload.messageId
    ) {
      // The submitted review belongs to the original assistant message;
      // do not merge it into the message currently being streamed.
      this.onExternalReview?.(payload.messageId, payload.review);
      return info.originMessage ?? emptyAssistant();
    }
    return reduceAiSseMessage(info.originMessage, info.chunk);
  }
}