Langgraph Studio token streaming doesn't work after resume, langgraphjs Gemini

[fridaystreet]

Hi,

I have a streaming setup with token streaming using gemini. It worls absolutely fine until I interrupt and resume. After a resume it no longer streams the tokens. I have debugged the llm calls and can see that they are indeed using the runManager callbacks in the streamResponseChuncks and I have dug all the way into the run manager calls and can see it getting all the way through to posting the events. It doesn't look like that is having any issues.

As I say, prior to the resume, the streaming works fine. I am storing the compiled graph instance in a map inside a class and then retrieving it upon resume using the thread id and calling the stream with the resume payload and the config passed in by studio. I was adding some extra stuff to the config for our own executions, thought maybe this was breaking it, but I've completely removed that and am passing the exact config object directly from the studio stream request and no joy

I put together the standalone test below just to confirm that it definitely works and this test works. It still streams tokens after the interrupt. But I can't for the life of me work out what's breaking it in my actual code, apart from this retrieval of the compiled graph from a cache map. It feels like the Gemini adapter and the call backs are all working, but for somereason they are possibly posting the tokes to the wrong (old) run? Is that possible?.

I can't post up my actual code as it's part of a complex system, I'll try and isolate it a bit more and see if I can replicate the full process of the stored graph in cache map.

Just wondered if anyone had seen anything similar and knew of any gotchas in and around this?

Cheers

working test code

// Tool-based streaming interrupt test - replicates React agent ask question behavior
import { 
  StateGraph, 
  START, 
  END, 
  Annotation, 
  MessagesAnnotation,
  messagesStateReducer,
  interrupt,
  Command
} from "@langchain/langgraph";
import { ChatGoogleGenerativeAI } from "@langchain/google-genai";
import { HumanMessage, AIMessage, BaseMessage, ToolMessage, AIMessageChunk } from "@langchain/core/messages";
import { MemorySaver } from "@langchain/langgraph";
import { DynamicStructuredTool } from "@langchain/core/tools";
import { z } from "zod";

// State structure matching React agent pattern
const StreamingTestState = Annotation.Root({
  ...MessagesAnnotation.spec,
  internalMessages: Annotation<BaseMessage[]>({
    reducer: messagesStateReducer,
    default: () => []
  })
});

// Create ask question tool that triggers interrupt (like your React agent)
const askQuestionTool = new DynamicStructuredTool({
  name: "ask_question",
  description: "Ask the user a question and wait for their response",
  schema: z.object({
    question: z.string().describe("The question to ask the user")
  }),
  func: async ({ question }) => {
    // Return the question - the interrupt will be handled in the tools node
    return `Asked question: ${question}`;
  }
});

// LLM with tools bound (like React agent)
const llm = new ChatGoogleGenerativeAI({
  model: "gemini-2.5-flash",
  apiKey: "<your apikey>",
  temperature: 0,
  streaming: true
}).bindTools([askQuestionTool]);

// Agent node - immediately asks question to test streaming after interrupt
async function agentNode(state: typeof StreamingTestState.State, config?: any) {
  const { internalMessages, messages } = state;
  
  const updateInternalMessages = [...internalMessages];
  if (messages.length && !internalMessages.find((m: any) => m.id === messages[messages.length-1].id)) {
    updateInternalMessages.push(messages[messages.length-1]);
  }

  // Simple system prompt: just ask a question immediately
  const systemPrompt = new HumanMessage({
    content: `You are a helpful assistant. Immediately ask the user a follow-up question using the ask_question tool. Don't provide a long answer first - just use the ask_question tool right away.`
  });
  
  const messagesWithSystem = [systemPrompt, ...updateInternalMessages];

  // This should stream tokens
  const response = await llm.invoke(messagesWithSystem, config) as AIMessage;
  
  updateInternalMessages.push(response);
  return { 
    internalMessages: updateInternalMessages,
    messages: [response]
  };
}

// Tools node - handles tool calls and creates interrupts
async function toolsNode(state: typeof StreamingTestState.State) {
  const { internalMessages } = state;
  const lastMessage = internalMessages[internalMessages.length - 1] as AIMessage;
  
  if (!lastMessage.tool_calls || lastMessage.tool_calls.length === 0) {
    throw new Error("No tool calls found");
  }

  const toolMessages: ToolMessage[] = [];
  
  for (const toolCall of lastMessage.tool_calls) {
    if (toolCall.name === "ask_question") {
      // Execute the tool to get the question
      const result = await askQuestionTool.invoke(toolCall);
      const question = toolCall.args.question;
      
      // Create interrupt metadata like your React agent
      const interruptData = {
        interrupt: {
          interrupts: [{ value: question }]
        }
      };
      
      const toolMessage = new ToolMessage({
        content: result,
        tool_call_id: toolCall.id!,
        metadata: interruptData
      });
      toolMessages.push(toolMessage);
    }
  }
  
  return {
    internalMessages: toolMessages
  };
}

// Human node - handles the interrupt (like your React agent)
function humanNode(state: typeof StreamingTestState.State): Command {
  const lastToolMessage: any = state.internalMessages[state.internalMessages.length-1];
  
  let msg = 'Ready for user input';
  if (lastToolMessage?.metadata?.interrupt?.interrupts?.length) {
    const interruptData: any = lastToolMessage.metadata.interrupt.interrupts[0];
    msg = interruptData.value;
  }
  
  const response = interrupt(msg);
  
  return new Command({
    update: {
      "messages": [
        new HumanMessage({ content: JSON.stringify(response) })
      ]
    }
  });
}

// Router - determines next step based on last message
function shouldContinue(state: typeof StreamingTestState.State) {
  const { internalMessages } = state;
  const lastMessage = internalMessages[internalMessages.length - 1];
  
  // If it's a tool message with interrupt, go to human
  if (lastMessage instanceof ToolMessage && lastMessage.metadata?.interrupt) {
    console.log("Routing to HUMAN for interrupt");
    return 'human';
  }
  
  // If agent has tool calls, go to tools
  if ((lastMessage instanceof AIMessageChunk) && lastMessage.tool_calls && lastMessage.tool_calls.length > 0) {
    console.log("Routing to TOOLS");
    return 'tools';
  }
  
  // Otherwise end
  console.log("Routing to END");
  return END;
}

// Build the graph (mimicking React agent structure)
function createStreamingTestGraph() {
  const graph = new StateGraph(StreamingTestState)
    .addNode("agent", agentNode)
    .addNode("tools", toolsNode)
    .addNode("human", humanNode)
    .addEdge(START, "agent")
    .addEdge("human", "agent")
    .addConditionalEdges("agent", shouldContinue, {
      tools: "tools",
      [END]: END
    })
    .addConditionalEdges("tools", shouldContinue, {
      human: "human",
      agent: "agent",
      [END]: END
    });

  const checkpointer = new MemorySaver();
  return graph.compile({ checkpointer });
}

export const streamingInterruptTest = createStreamingTestGraph();

// Test: 
// 1. Send message like "Tell me about cats"
// 2. Agent streams response, then calls ask_question tool
// 3. Tool triggers interrupt with question
// 4. You respond to question
// 5. Agent should stream final response (THIS IS WHERE BUG OCCURS)

[fridaystreet]

fixed it, in case anyone else gets here, it seems you can't reassign the original config object at all, so any { ...config, ... } seems to break the callbacks for some reason.

Well breaks the callbacks after a resume. I was able to spread the config prior to the first call, but doing it on subsequent calls when resuming, it breaks.

Anyway I've just refactored my code to not reassign the original config object coming from studio.