Implement ToolLimitationNode in stateGraph

hi @Dev

if this is your full implementation (source):

limit = ToolCallLimitMiddleware(
    thread_limit=None,
    run_limit=DEFAULT_TOOL_LIMIT,
    exit_behavior=DEFAULT_EXIT_BEHAVIOR,
)
def limit_check_node(state: GraphState, runtime):
        """
        MUST be placed after tool execution.
        This is the only correct place where tool call count is updated.
        """
        updates = limit.after_model(state, runtime) or {}

        if updates.get("jump_to") == "end":
            return {"__end__": True}

        return updates

?

after_model doesn’t return jump_to. It only retuns this:

    def after_model(self, state: ModelCallLimitState, runtime: Runtime) -> dict[str, Any] | None:  # noqa: ARG002
        """Increment model call counts after a model call.

        Args:
            state: The current agent state.
            runtime: The langgraph runtime.

        Returns:
            State updates with incremented call counts.
        """
        return {
            "thread_model_call_count": state.get("thread_model_call_count", 0) + 1,
            "run_model_call_count": state.get("run_model_call_count", 0) + 1,
        }

you have to create another node that is before your llm call and call limit.before_model there/

See this (source):

from langchain.agents.middleware.model_call_limit import ModelCallLimitMiddleware, ModelCallLimitExceededError

limit = ModelCallLimitMiddleware(thread_limit=10, run_limit=5, exit_behavior="end")

def before_model_node(state, runtime):
    updates = limit.before_model(state, runtime) or {}
    # If updates includes {"jump_to": "end"}, add your conditional edge to END.
    return updates

def model_node(state):
    # ... call your model ...
    return {"messages": [...]}  # your AIMessage(s)

def after_model_node(state, runtime):
    return limit.after_model(state, runtime) or {}