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) => 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) => {
    const controller = new AbortController();
    controllerRef.current = controller;
    dispatch({ type: "start" });
    try {
      const res = await fetch(
        `${API_BASE_PUBLIC}/projects/${projectId}/chapters/${chapterNo}/draft`,
        {
          method: "POST",
          headers: { Accept: "text/event-stream" },
          signal: controller.signal,
        },
      );
      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,
  };
}