[Feature] MCP and HTTP tools, agentic tutorial, see-also pointers

docs/source/reference/llms_envs.rst

@@ -1,5 +1,7 @@
 :orphan:
 
+.. _llm_envs:
+
 .. currentmodule:: torchrl.envs.llm
 
 LLM Environments
@@ -159,3 +161,70 @@ trades rich display for portability.
     ReplError
     JupyterRepl
     SubprocessRepl
+
+Built-in tools and adapters
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. currentmodule:: torchrl.envs.llm.agentic
+
+.. autosummary::
+    :toctree: generated/
+    :template: rl_template.rst
+
+    ToolCompose
+    DispatchResult
+    PythonTool
+    ShellTool
+    FileReadTool
+    StopTool
+    HttpTool
+    MCPServerConfig
+    MCPToolset
+    RateLimiter
+    as_tool
+
+Migration from legacy tool transforms
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Existing code built on :mod:`torchrl.envs.llm.transforms` keeps working:
+no ``DeprecationWarning`` is emitted in this release. Each legacy class
+has a ``.. seealso::`` block in its docstring pointing at the
+recommended replacement, summarised here.
+
+.. list-table:: Legacy transform → agentic counterpart
+   :header-rows: 1
+   :widths: 30 30 40
+
+   * - Legacy
+     - Agentic
+     - Adapter recipe
+   * - ``ExecuteToolsInOrder``
+     - :class:`ToolCompose`
+     - Replace at the env stack level. ``ToolCompose`` runs calls
+       concurrently; pin sequential execution per-tool with
+       :class:`RateLimiter` ``max_concurrent=1`` if you depend on
+       ordering.
+   * - ``PythonInterpreter``
+     - :class:`PythonTool` + :class:`Sandbox` + :class:`Repl`
+     - For a soft migration, lift the existing transform: ``as_tool(PythonInterpreter(persistent=True), name="python", input_schema=...)``.
+   * - ``SimpleToolTransform``
+     - Native :class:`Tool` subclass
+     - Or ``as_tool(transform, name=..., input_schema=...)``.
+   * - ``BrowserTransform``
+     - :func:`tools.as_tool` of the existing transform
+     - A native :class:`Tool` for browser automation may land later;
+       until then the adapter is the recommended path.
+   * - ``MCPToolTransform``
+     - :class:`MCPToolset`
+     - One :class:`Tool` per remote tool, schemas auto-discovered.
+       Register it with ``ToolCompose.add_toolset`` so discovery,
+       dispatch, and shutdown share one event loop.
+   * - ``XMLBlockParser`` / ``JSONCallParser``
+     - :class:`parsers.XMLToolCallParser` / :class:`parsers.JSONToolCallParser`
+     - Same syntax; the agentic versions enforce a stable ``call_id``.
+   * - ``ToolService`` / ``ToolRegistry``
+     - The ``tools=[...]`` argument to :class:`ToolCompose`
+     - The registry pattern collapses into the compose container.
+
+For a guided walkthrough, see the
+:ref:`agentic ChatEnv tutorial <llm_agentic>`.

test/llm/test_agentic.py

@@ -15,8 +15,10 @@
 import asyncio
 import json
 import sys
+import threading
 import time as _time
 import warnings
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from typing import ClassVar
 
 import pytest
@@ -58,7 +60,8 @@
     _has_sandbox_exec,
     _profile,
 )
-from torchrl.envs.llm.agentic.tools import as_tool
+from torchrl.envs.llm.agentic.tools import as_tool, HttpTool
+from torchrl.envs.llm.agentic.tools.mcp import _has_mcp, _MCPTool, MCPServerConfig
 from torchrl.envs.llm.transforms import (
     IncrementalTokenizer,
     PythonInterpreter,
@@ -752,6 +755,38 @@ def test_lookup_by_name(self):
         assert "stop" in compose
         assert compose["stop"].name == "stop"
 
+    def test_toolset_uses_compose_event_loop_for_full_lifecycle(self):
+        class _ToolsetTool(_Echo):
+            def __init__(self, toolset):
+                self.toolset = toolset
+
+            async def run(self, args, ctx):
+                assert asyncio.get_running_loop() is self.toolset.loop
+                return await super().run(args, ctx)
+
+        class _Toolset:
+            def __init__(self):
+                self.loop = None
+                self.closed_on_loop = False
+                self.tools = (_ToolsetTool(self),)
+
+            async def open(self):
+                self.loop = asyncio.get_running_loop()
+
+            async def close(self):
+                self.closed_on_loop = asyncio.get_running_loop() is self.loop
+
+        toolset = _Toolset()
+        compose = ToolCompose(tools=[], parser=XMLToolCallParser())
+        compose.add_toolset(toolset)
+        env = TransformedEnv(ChatEnv(batch_size=(1,), input_mode="history"), compose)
+        obs = env.reset(TensorDict({"query": "?"}, batch_size=(1,)))
+        _push_assistant(obs, '<tool name="echo">{"value": 1}</tool>')
+        nxt = env.step(obs)
+        assert not bool(nxt.get(("next", "agentic", "any_error")).item())
+        env.close()
+        assert toolset.closed_on_loop
+
     def test_parallel_dispatch_wall_time(self):
         env = _agentic_env([_Sleeper("a"), _Sleeper("b"), _Sleeper("c")])
         obs = env.reset(TensorDict({"query": "go"}, batch_size=(1,)))
@@ -1173,3 +1208,155 @@ def _process_batch_item(self, content, index):
         tool = as_tool(_Legacy(), name="legacy")
         _run(tool.run({}, ToolContext(call_id="c", batch_index=3)))
         assert indices == [3]
+
+
+# ----- MCP and HTTP tools -----
+
+
+class TestMCPToolset:
+    def test_construction_requires_mcp_package(self):
+        if not _has_mcp:
+            with pytest.raises(ImportError):
+                from torchrl.envs.llm.agentic.tools import MCPToolset
+
+                MCPToolset(MCPServerConfig(command="true"))
+        else:
+            # When the package is installed we can at least construct
+            # without opening a session.
+            from torchrl.envs.llm.agentic.tools import MCPToolset
+
+            pool = MCPToolset(MCPServerConfig(command="true"))
+            assert pool.tools == ()
+
+    def test_server_config(self):
+        cfg = MCPServerConfig(command="npx", args=("@browsermcp/mcp@latest",))
+        assert cfg.command == "npx"
+        assert cfg.args == ("@browsermcp/mcp@latest",)
+
+    def test_tool_call_runs_on_opening_loop(self):
+        class _Response:
+            content = []
+            isError = False
+
+        class _Session:
+            async def call_tool(self, name, args):
+                assert name == "remote"
+                assert args == {"value": 1}
+                return _Response()
+
+        class _Toolset:
+            request_timeout = 1.0
+            _session = _Session()
+            _loop = None
+
+        async def go():
+            toolset = _Toolset()
+            toolset._loop = asyncio.get_running_loop()
+            tool = _MCPTool(
+                toolset=toolset,
+                remote_name="remote",
+                description="",
+                input_schema={"type": "object"},
+                exposed_name="remote",
+            )
+            return await tool.run({"value": 1}, ToolContext(call_id="mcp-call"))
+
+        result = _run(go())
+        assert not result.is_error
+
+
+class TestHttpTool:
+    def test_blocks_disallowed_host(self):
+        async def go():
+            tool = HttpTool(allowed_hosts=("api.example.com",))
+            await tool.setup()
+            res = await tool.run(
+                {"url": "https://other-host.example/foo"},
+                ToolContext(call_id="c"),
+            )
+            assert res.is_error
+            assert "allowed_hosts" in res.text
+            await tool.teardown()
+
+        _run(go())
+
+    def test_protocol_conformance(self):
+        tool = HttpTool()
+        # Sanity: it walks like a Tool.
+        assert tool.name == "http"
+        assert callable(tool.run)
+        assert "url" in tool.input_schema["properties"]
+
+    @pytest.mark.parametrize("scheme", ["file", "ftp"])
+    def test_blocks_non_http_schemes(self, scheme):
+        result = _run(
+            HttpTool().run(
+                {"url": f"{scheme}:///etc/hosts"},
+                ToolContext(call_id="scheme"),
+            )
+        )
+        assert result.is_error
+        assert result.meta["blocked_scheme"] == scheme
+
+    def test_blocks_redirect_to_disallowed_host(self):
+        class _TargetHandler(BaseHTTPRequestHandler):
+            def do_GET(self):
+                self.send_response(200)
+                self.end_headers()
+                self.wfile.write(b"private target")
+
+            def log_message(self, format, *args):
+                pass
+
+        target = ThreadingHTTPServer(("127.0.0.1", 0), _TargetHandler)
+
+        class _RedirectHandler(BaseHTTPRequestHandler):
+            def do_GET(self):
+                self.send_response(302)
+                self.send_header(
+                    "Location", f"http://localhost:{target.server_port}/secret"
+                )
+                self.end_headers()
+
+            def log_message(self, format, *args):
+                pass
+
+        redirect = ThreadingHTTPServer(("127.0.0.1", 0), _RedirectHandler)
+        threads = [
+            threading.Thread(target=server.serve_forever)
+            for server in (target, redirect)
+        ]
+        for thread in threads:
+            thread.start()
+        try:
+            exact = _run(
+                HttpTool(
+                    allowed_hosts=("127.0.0.1",),
+                    max_response_bytes=len(b"private target"),
+                ).run(
+                    {"url": f"http://127.0.0.1:{target.server_port}/"},
+                    ToolContext(call_id="exact"),
+                )
+            )
+            assert not exact.meta["truncated"]
+            assert exact.text == "private target"
+
+            result = _run(
+                HttpTool(allowed_hosts=("127.0.0.1",)).run(
+                    {"url": f"http://127.0.0.1:{redirect.server_port}/"},
+                    ToolContext(call_id="redirect"),
+                )
+            )
+            assert result.is_error
+            assert "redirect host" in result.text
+            assert "private target" not in result.text
+        finally:
+            for server in (redirect, target):
+                server.shutdown()
+                server.server_close()
+            for thread in threads:
+                thread.join()
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

torchrl/envs/llm/agentic/__init__.py

@@ -33,14 +33,27 @@
 )
 from .rate_limit import RateLimiter
 from .schema import json_schema_from_pydantic, validate_args
-from .tools import as_tool, FileReadTool, PythonTool, ShellTool, StopSignal, StopTool
+from .tools import (
+    as_tool,
+    FileReadTool,
+    HttpTool,
+    MCPServerConfig,
+    MCPToolset,
+    PythonTool,
+    ShellTool,
+    StopSignal,
+    StopTool,
+)
 
 __all__ = [
     "DispatchResult",
     "FileReadTool",
     "FileRefPart",
+    "HttpTool",
     "ImagePart",
     "JsonPart",
+    "MCPServerConfig",
+    "MCPToolset",
     "ParseResult",
     "ParsedCall",
     "PythonTool",

torchrl/envs/llm/agentic/compose.py

@@ -176,6 +176,7 @@ def __init__(
         self._validate_inputs = validate_inputs
         self._setup_done = False
         self._async_runner: _AsyncRunner | None = None
+        self._toolsets: list[Any] = []
 
     # ----- introspection helpers -----
 
@@ -273,11 +274,47 @@ def clone(self) -> ToolCompose:
             validate_inputs=self._validate_inputs,
         )
 
+    def add_toolset(self, toolset: Any) -> None:
+        """Open an async toolset and add its discovered tools.
+
+        The toolset is opened on the same persistent event loop used for tool
+        dispatch and is closed with this compose. It must expose async
+        ``open()`` / ``close()`` methods and a ``tools`` iterable after opening.
+
+        Args:
+            toolset: An async tool provider such as
+                :class:`~torchrl.envs.llm.agentic.MCPToolset`.
+        """
+        if self._setup_done:
+            raise RuntimeError("cannot add a toolset after tool setup")
+        self._run_async(toolset.open())
+        try:
+            tools = list(toolset.tools)
+            for tool in tools:
+                if not _is_tool(tool):
+                    raise TypeError(
+                        "toolset returned a non-Tool object: "
+                        f"{type(tool).__name__!r}"
+                    )
+            names = [tool.name for tool in tools]
+            duplicates = set(names) & set(self._tools_by_name)
+            duplicates.update(name for name in names if names.count(name) > 1)
+            if duplicates:
+                raise ValueError(f"duplicate tool names: {sorted(duplicates)!r}")
+        except Exception:
+            self._run_async(toolset.close())
+            raise
+        for tool in tools:
+            self.append_tool(tool)
+        self._toolsets.append(toolset)
+
     # ----- lifecycle -----
 
     async def _setup_tools(self) -> None:
         if self._setup_done:
             return
+        for toolset in self._toolsets:
+            await toolset.open()
         for tool in self._tool_list:
             try:
                 await tool.setup()
@@ -293,6 +330,11 @@ async def _teardown_tools(self) -> None:
                 torchrl_logger.exception(
                     "tool %r teardown raised; continuing", tool.name
                 )
+        for toolset in reversed(self._toolsets):
+            try:
+                await toolset.close()
+            except Exception:  # pragma: no cover
+                torchrl_logger.exception("toolset close raised; continuing")
         self._setup_done = False
 
     def close(self) -> None:  # type: ignore[override]