Update state from within node

[tonxxd]

Hi, i have a long running function that needs to be execute by the agent. I would like to send multiple ai messages from within the node with updates about the state of the
function execution.
I tried to do something like

let runnable: CompiledStateGraph<any>;

const config = {
  configurable: {
    thread_id: "1",
  },
};

const graph = new MessageGraph();

// ...
graph.addNode("data", async (state: BaseMessage[], { config }) => {
  const results = await compute({
    onUpdate: (u) => {
      runnable.updateState(config, [
        new AIMessage(message.content, {
          type: "text",
        }),
      ]);
    },
  });

  return [new AIMessage("done")];
});

const checkpointer = new MemorySaver();

runnable = graph.compile({
  checkpointer,
});

for await (const event of await runnable.stream(messages, config)) {
  console.log(
    `event received`);
}

But the event is not received by the stream callback..
Any ideas? Is there a way to update the state from within the node, before returning?