サーバーアクション（Vercel AI）をuseMutationに噛ませたい

業務で必要になったので作ってみました。

以下ざくっと実装です、本当に最低限だけ。

useServerAction

import { StreamableValue, readStreamableValue } from "ai/rsc";
import { useCallback, useState } from "react";

type UseMutationOptions<TData, TVariables> = {
  mutationFn: (variables: TVariables) => Promise<StreamableValue<TData>>;
};

type UseMutationResult<TData, TVariables> = {
  data?: TData;
  isPending: boolean;
  mutate: (variables: TVariables) => void;
};

export default function useServerAction<TData = unknown, TVariables = void>({
  mutationFn,
}: UseMutationOptions<TData, TVariables>): UseMutationResult<
  TData,
  TVariables
> {
  const [data, setData] = useState<TData>();
  const [isPending, setIsPending] = useState(false);
  const mutate = useCallback<UseMutationResult<TData, TVariables>["mutate"]>(
    async (variables) => {
      setIsPending(true);

      try {
        const data = await mutationFn(variables);

        for await (const partialData of readStreamableValue(data)) {
          if (partialData) {
            setData(partialData);
          }
        }
      } finally {
        setIsPending(false);
      }
    },
    [mutationFn],
  );

  return {
    data,
    isPending,
    mutate,
  };
}

サーバーアクション（Vercel AI）

"use server";
import { openai } from "@ai-sdk/openai";
import { streamObject } from "ai";
import { StreamableValue, createStreamableValue } from "ai/rsc";
import { z } from "zod";

type GenerateData = Partial<{
  fuga: string;
}>;

export default async function generate(): Promise<
  StreamableValue<GenerateData>
> {
  const stream = createStreamableValue<GenerateData>();

  (async (): Promise<void> => {
    const { partialObjectStream } = await streamObject({
      model: openai("gpt-4o"),
      prompt: "hoge",
      schema: z.object({
        fuga: z.string().describe("piyo"),
      }),
    });

    for await (const partialObject of partialObjectStream) {
      stream.update(partialObject);
    }

    stream.done();
  })();

  return stream.value;
}

クライアントコンポーネント

"use client";
import serverAction from "./actions";

export default function ClientComponent(): JSX.Element {
  const {
    data,
    isPending,
    mutate,
  } = useServerAction({ mutationFn: serverAction });

  ...
}

そんな感じです。