Langgraph Help: Writing awaitable actions in nodes

Question

I'm trying to follow the langgraph documentation and create a basic graph as below. The main caveat of my graph is that I don't have an LLM, but just make API calls that aggregate as a response. The problem I'm facing is with writing awaitable Actions to make the API calls in the graph nodes.

NOTE: This same code works if I change it to its non-awaited counterparts.

Minimal reproduction Code:

class State(TypedDict):
    messages: Annotated[list[str], add_messages]


class ReturnNodeValue:
    def __init__(self, assistantId: str):
        self.assistantId = assistantId

    async def __call__(self, state: State, **kwargs) -> Any:
        response = await self.custom_api_call(state["messages"][-1], self.assistantId)
        return {"messages": [response]}

    async def custom_api_call(self, message: str, assistantId: str) -> str:
        return await callApi(message, assistantId)


class GraphBuilder:
    def __init__(self):
        self.graph_builder = StateGraph(State)
        self.buildGraph()

    def buildGraph(self, assistantId1, assistantId2):
        self.graph_builder.add_edge(START, "node1")
        self.graph_builder.add_edge("node1", "node2")
        self.graph_builder.add_edge("node2", END)

        self.graph_builder.add_node("node1", ReturnNodeValue(assistantId1))
        self.graph_builder.add_node("node2", ReturnNodeValue(assistantId2))

    async def executeGraph(self):
        try:
            initial_state = {"messages": ["Hellow"]}
            app = self.graph_builder.compile()
            result = await app.ainvoke(initial_state)
            # Error on above line: "Dispatcher.dispatch_forever.<locals>.<lambda>() got an unexpected keyword argument 'context'"
            return result
        except Exception as e:
            print(f"An error occurred: {e}")
            raise


graph = GraphBuilder()
graph.buildGraph("abc123", "abc234")
result = graph.executeGraph()

Error: Dispatcher.dispatch_forever..() got an unexpected keyword argument 'context'

The error happens when I hit the app.ainvoke(), and as mentioned in the note, the code works perfectly fine if I don't await anything.

Any help would greatly be appreciated, Thank you :)