writer-work-flow/apps/web/lib/stream/useDraftStream.ts

"use client";

import { useCallback, useReducer, useRef } from "react";

import { API_BASE_PUBLIC } from "@/lib/api/config";
import {
  SseFrameBuffer,
  initialStreamState,
  reduceStream,
  type StreamState,
} from "./sse";

type Action =
  | { type: "start" }
  | { type: "events"; events: ReturnType<SseFrameBuffer["push"]> }
  | { type: "abort" }
  | { type: "fail"; code: string; message: string }
  | { type: "reset"; text: string };

function reducer(state: StreamState, action: Action): StreamState {
  switch (action.type) {
    case "start":
      return { phase: "streaming", text: "", error: null };
    case "events":
      return action.events.reduce(reduceStream, state);
    case "abort":
      return { ...state, phase: "aborted" };
    case "fail":
      return {
        ...state,
        phase: "error",
        error: { code: action.code, message: action.message },
      };
    case "reset":
      return { ...initialStreamState, text: action.text };
    default:
      return state;
  }
}

export interface UseDraftStream {
  state: StreamState;
  isStreaming: boolean;
  start: (
    projectId: string,
    chapterNo: number,
    directive?: string,
  ) => Promise<void>;
  stop: () => void;
  reset: (text: string) => void;
}

// 消费 POST .../draft 的 SSE 流。用 fetch+ReadableStream（EventSource 不支持 POST）。
// "停" = abort 连接（已生成部分保留在 state.text，由调用方落入草稿）。
export function useDraftStream(): UseDraftStream {
  const [state, dispatch] = useReducer(reducer, initialStreamState);
  const controllerRef = useRef<AbortController | null>(null);

  const stop = useCallback(() => {
    controllerRef.current?.abort();
    controllerRef.current = null;
    dispatch({ type: "abort" });
  }, []);

  const reset = useCallback((text: string) => {
    dispatch({ type: "reset", text });
  }, []);

  const start = useCallback(
    async (projectId: string, chapterNo: number, directive?: string) => {
      const controller = new AbortController();
      controllerRef.current = controller;
      dispatch({ type: "start" });
      try {
        // 仅当本章指令非空才发 JSON body（无指令退回裸 POST，与旧后端调用兼容）。
        const trimmed = directive?.trim();
        const init: RequestInit =
          trimmed && trimmed.length > 0
            ? {
                method: "POST",
                headers: {
                  Accept: "text/event-stream",
                  "Content-Type": "application/json",
                },
                body: JSON.stringify({ directive: trimmed }),
                signal: controller.signal,
              }
            : {
                method: "POST",
                headers: { Accept: "text/event-stream" },
                signal: controller.signal,
              };
        const res = await fetch(
          `${API_BASE_PUBLIC}/projects/${projectId}/chapters/${chapterNo}/draft`,
          init,
        );
        if (!res.ok || !res.body) {
          // 流前错误（如无凭据 → 503 LLM_UNAVAILABLE，JSON 信封而非帧）。
          let code = "STREAM_FAILED";
          let message = `写章请求失败（${res.status}）`;
          try {
            const body = (await res.json()) as {
              error?: { code?: string; message?: string };
            };
            if (body.error?.code) code = body.error.code;
            if (body.error?.message) message = body.error.message;
          } catch {
            // 非 JSON 信封，沿用默认文案。
          }
          dispatch({ type: "fail", code, message });
          return;
        }
        const reader = res.body.getReader();
        const decoder = new TextDecoder();
        const buffer = new SseFrameBuffer();
        for (;;) {
          const { value, done } = await reader.read();
          if (done) break;
          const events = buffer.push(decoder.decode(value, { stream: true }));
          if (events.length > 0) dispatch({ type: "events", events });
        }
      } catch (err: unknown) {
        if (err instanceof DOMException && err.name === "AbortError") {
          // 用户主动停止：state 已置 aborted。
          return;
        }
        const message = err instanceof Error ? err.message : "未知网络错误";
        dispatch({ type: "fail", code: "NETWORK", message });
      } finally {
        controllerRef.current = null;
      }
    },
    [],
  );

  return {
    state,
    isStreaming: state.phase === "streaming",
    start,
    stop,
    reset,
  };
}