trpc `useSubscription` doesn't get React state updates

I'm building a Next.js app with trpc that requires real time comunication. It has been going well, especially since my useSubscription calls have only been updating state rather than reading it so far.

I have a ContextProvider with some state which uses a web socket subscription through the useSubscription method. The isue is, the onData method depends on the Provider's state, but it doesn't get updated, only ever getting the initial state passed into the useState function

function Component() {
    const [state, setState] = useState('initial')
    const triggerSocket = trpc.trigger.useMutation()

    trpc.socket.useSubscription({
        onData() {
            console.log(state) // initial
        }
    })

    return <button onClick={async () => {
        setState('updated')
        await triggerSocket.mutateAsync()
    }}>Click me</button>
}

Am I missing something? Is this a trpc bug or am I doing something wrong? Is there some way to get around this issue?

I managed to get it to work by swapping my useState's with useRef's. That worked well enough for now, still, if anyone can get useState to work, feel free to leave an answer here