test_sync_async_wrappers.py — langchain Source File

Source Code

"""Tests for sync/async middleware composition with wrap_tool_call and awrap_tool_call.

These tests verify the desired behavior:
1. If middleware defines both sync and async -> use both on respective paths
2. If middleware defines only sync -> use on sync path, raise NotImplementedError on async path
3. If middleware defines only async -> use on async path, raise NotImplementedError on sync path
"""

from collections.abc import Awaitable, Callable
from typing import Any

import pytest
from langchain_core.messages import HumanMessage, ToolCall, ToolMessage
from langchain_core.tools import tool
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.types import Command

from langchain.agents.factory import create_agent
from langchain.agents.middleware.types import AgentMiddleware, ToolCallRequest, wrap_tool_call
from tests.unit_tests.agents.model import FakeToolCallingModel


@tool
def search(query: str) -> str:
    """Search for information."""
    return f"Results for: {query}"


@tool
def calculator(expression: str) -> str:
    """Calculate an expression."""
    return f"Calculated: {expression}"


class TestSyncAsyncMiddlewareComposition:
    """Test sync/async middleware composition behavior."""

    def test_sync_only_middleware_works_on_sync_path(self) -> None:
        """Middleware with only sync wrap_tool_call works on sync path."""
        call_log = []

        class SyncOnlyMiddleware(AgentMiddleware):
            def wrap_tool_call(
                self,
                request: ToolCallRequest,
                handler: Callable[[ToolCallRequest], ToolMessage | Command[Any]],
            ) -> ToolMessage | Command[Any]:
                call_log.append("sync_called")
                return handler(request)

        model = FakeToolCallingModel(
            tool_calls=[
                [ToolCall(name="search", args={"query": "test"}, id="1")],
                [],
            ]
        )

        agent = create_agent(
            model=model,
            tools=[search],
// ... (431 more lines)