test_tool_runtime_with_custom_state() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/test_injected_runtime_create_agent.py lines 285–327

def test_tool_runtime_with_custom_state() -> None:
    """Test ToolRuntime works with custom state schemas."""

    class CustomState(AgentState[Any]):
        custom_field: str

    runtime_state = {}

    @tool
    def custom_state_tool(x: int, runtime: ToolRuntime) -> str:
        """Tool that accesses custom state."""
        runtime_state["custom_field"] = runtime.state.get("custom_field", "not found")
        return f"Custom: {x}"

    class CustomMiddleware(AgentMiddleware):
        state_schema = CustomState

    agent = create_agent(
        model=FakeToolCallingModel(
            tool_calls=[
                [{"args": {"x": 10}, "id": "custom_call", "name": "custom_state_tool"}],
                [],
            ]
        ),
        tools=[custom_state_tool],
        system_prompt="You are a helpful assistant.",
        middleware=[CustomMiddleware()],
    )

    result = agent.invoke(
        {
            "messages": [HumanMessage("Test custom state")],
            "custom_field": "custom_value",
        }
    )

    # Verify custom field was accessible
    assert runtime_state["custom_field"] == "custom_value"

    # Verify tool executed
    tool_message = result["messages"][2]
    assert isinstance(tool_message, ToolMessage)
    assert tool_message.content == "Custom: 10"