스트림이 중단될 때 onEnd가 호출되지 않아요

스트림이 중단될 때 onEnd가 호출되지 않아요

toUIMessageStream에 onEnd 콜백을 썼는데 스트림이 중단(abort)될 때 콜백이 실행되지 않는 문제와 해결 방법을 알려드려요.

출처: 문서

본문

문제 (Issue)

toUIMessageStream에 onEnd 콜백을 사용할 때, 스트림이 중단되면 콜백이 실행되지 않을 수 있어요. 이는 abort 핸들러가 즉시 응답을 종료해서 onEnd 콜백이 실행될 기회가 없기 때문이에요.

// abort 시 onEnd가 호출되지 않는 서버 측 코드
export async function POST(req: Request) {
  const { messages } = await req.json();

  const result = streamText({
    model: __MODEL__,
    messages: await convertToModelMessages(messages),
    abortSignal: req.signal,
  });

  return createUIMessageStreamResponse({
    stream: toUIMessageStream({
      stream: result.stream,
      onEnd: async ({ isAborted }) => {
        // 스트림이 중단되면 이것은 호출되지 않아요!
        if (isAborted) {
          console.log('Stream was aborted');
          // 중단 특정 정리 처리
        } else {
          console.log('Stream completed normally');
          // 정상 완료 처리
        }
      },
    }),
  });
}

배경 (Background)

스트림이 중단되면 응답이 즉시 종료돼요. 적절한 처리가 없으면 onEnd 콜백이 실행될 기회가 없어서, 부분 결과 저장이나 abort 이벤트 로깅 같은 중요한 정리 작업이 실행되지 못해요.

해결 방법 (Solution)

createUIMessageStreamResponse 구성에 consumeSseStream: consumeStream을 추가하세요. 그러면 abort 이벤트가 제대로 포착되어 onEnd 콜백으로 전달돼, 스트림이 중단돼도 콜백이 실행될 수 있어요.

// 다른 imports...
import {
  consumeStream,
  createUIMessageStreamResponse,
  toUIMessageStream,
} from 'ai';

export async function POST(req: Request) {
  const { messages } = await req.json();

  const result = streamText({
    model: __MODEL__,
    messages: await convertToModelMessages(messages),
    abortSignal: req.signal,
  });

  return createUIMessageStreamResponse({
    stream: toUIMessageStream({
      stream: result.stream,
      onEnd: async ({ isAborted }) => {
        // 이제 중단돼도 이게 호출될 거예요!
        if (isAborted) {
          console.log('Stream was aborted');
          // 중단 특정 정리 처리
        } else {
          console.log('Stream completed normally');
          // 정상 완료 처리
        }
      },
    }),
    consumeSseStream: consumeStream, // 이게 abort 시 onEnd가 호출되게 해줘요
  });
}

더 알아보기 (Learn more)