Runnable with_listeners only hits first LLM invoke/ainvoke call in StateGraph Node

[colintkyle]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

from langgraph.graph import StateGraph, END
from langchain_openai import AzureChatOpenAI
from my_repo import LLMBuilder
from typing import TypedDict
from langchain_core.prompts import PromptTemplate
import logging

import sys

logging.basicConfig(
    level=logging.DEBUG, 
    format="%(levelname)s: %(message)s",
    stream=sys.stdout, 
    force=True   
)

logger = logging.getLogger(__name__)

class State(TypedDict):
    data: str

def listener(run_record):
    logger.info(run_record.inputs)

class MyGraph:
    def __init__(self):
        # Create a StateGraph instance from LangChain.
        self.graph = StateGraph(State)
        self.graph.add_node("main", self.main)

        self.graph.set_entry_point("main")
        self.graph.add_edge("main", END)

        # Iterate over all nodes in the graph and attach listeners
        for node_name in self.graph.nodes.keys():
            runnable = self.graph.nodes.get(node_name).runnable
            # Ensure the runnable has the 'with_listeners' method
            if hasattr(runnable, 'with_listeners'):
                try:
                    new_runnable = runnable.with_listeners(on_end=listener)
                    new_node = self.graph.nodes.get(node_name)._replace(runnable=new_runnable)
                    self.graph.nodes[node_name] = new_node
                except Exception as e:
                    print(f"Failed to attach listener to runnable '{node_name}': {e}")
            else:
                print(f"Runnable '{node_name}' does not support attaching listeners.")

    async def run(self,state):
        app = self.graph.compile()
        return await app.ainvoke(state)
    
    def main(self,state):

        llm: AzureChatOpenAI = LLMBuilder.build_llm()

        dialog_prompt = PromptTemplate.from_template(
            """
            {user}
            """
        )

        chain = dialog_prompt | llm

        response1 = chain.invoke({'user':"user1"})
        state['data'] += response1.content
        response2 = chain.invoke({'user':"user2"})
        state['data'] += response2.content
        return state


mg = MyGraph()
await mg.run(State({'data':""}))
# This outputs only:
# INFO: {'user': 'user1'}

Description

You can see in the minimal example here I'm creating a StateGraph and attaching the listener which will print the input arguments. I attach the listener to a node that makes two invoke calls.

It's a bit confusing that it's not the inputs to the node (state) that are getting logged - it's only the invoke/ainvoke calls. However, I do see the advantage of being able to get the invoke calls, however...

My listener only gets called the first time invoke is called. So I'm not able to record details of the LLM call.

I get that I'm not really following the design pattern here and could add another node and edge, but the documentation doesn't explain that this way of building a StateGraph would be a problem. Unfortunately, I have a large amount of code written in this style, does anyone have any suggestions for getting listeners to work for multiple llm calls? I'm also wondering if I should raise an issue, or am just misunderstanding how StateGraph is supposed to work.

Thanks!

System Info

langchain==0.3.17
langchain-community==0.3.16
langchain-core==0.3.33
langchain-openai==0.2.14
langchain-text-splitters==0.3.5