Recommended pattern for initializing useQuery with server loaded value

export default async function ChatServer({
  chatId,
}: {
  chatId: ChatId | null;
}) {
  const initialMessages = chatId
    ? await fetchQuery(api.chats.listMessagesByChatId, {
        chatId,
      })
    : undefined;
  return (
    <ChatClient
      initialChatId={chatId}
      initialMessages={initialMessages}
    />
  );
}

export default function ChatClient({
  initialChatId,
  initialMessages,
}: {
  initialChatId: ChatId | null;
  initialMessages?: ChatMessage[];
}) {
  const messages = useQueryWithInit(
    api.chats.listMessagesByChatId,
    initialMessages,
    chatId ? { chatId } : "skip",
  );

  return (
    <>
      <div className="w-full h-full flex flex-col">
        <Content
          messages={messages}
        />
      </div>
    </>
  );
}

export function useQueryWithInit<T extends FunctionReference<"query">>(
  queryFunction: T,
  initialValue: FunctionReturnType<T> | undefined,
  ...args: OptionalRestArgsOrSkip<T>
): FunctionReturnType<T> | undefined {
  const queryResult = useQuery(queryFunction, ...args);

  const hasReceivedValue = useRef(false);

  if (args[0] === "skip") {
    return null;
  }

  if (queryResult !== undefined) {
    // Mark that we've received a value
    hasReceivedValue.current = true;
    // Return the query result immediately
    return queryResult;
  }

  if (!hasReceivedValue.current) {
    return initialValue;
  }

  return undefined;
}