Middleware

AgentScope provides a flexible middleware system that allows developers to intercept and modify the execution of various operations. Currently, middleware support is available for tool execution in the Toolkit class.

The middleware system follows an onion model, where each middleware wraps around the previous one, forming layers. This allows developers to:

• Perform pre-processing before the operation

• Intercept and modify responses during execution

• Perform post-processing after the operation completes

• Skip the operation execution entirely based on conditions

Tip

Future versions of AgentScope will expand middleware support to other components such as agents and models.

import asyncio
from typing import AsyncGenerator, Callable

from agentscope.message import TextBlock, ToolUseBlock
from agentscope.tool import ToolResponse, Toolkit

Tool Execution Middleware

The Toolkit class supports middleware for tool execution via the register_middleware method. Each middleware can intercept the tool call and modify the input or output.

Middleware Signature

A middleware function should have the following signature:

async def middleware(
    kwargs: dict,
    next_handler: Callable,
) -> AsyncGenerator[ToolResponse, None]:
    # Access parameters from kwargs
    tool_call = kwargs["tool_call"]

    # Pre-processing
    # ...

    # Call the next middleware or tool function
    async for response in await next_handler(**kwargs):
        # Post-processing
        yield response

Middleware Parameters

Parameter	Type	Description
kwargs	dict	Context parameters. Currently, includes tool_call (ToolUseBlock). May include additional parameters in future versions.
next_handler	Callable	A callable that accepts kwargs dict and returns a coroutine yielding AsyncGenerator of ToolResponse objects
Returns	AsyncGenerator[ToolResponse, None]	An async generator that yields ToolResponse objects

Basic Example

Here is a simple middleware that logs tool calls:

async def logging_middleware(
    kwargs: dict,
    next_handler: Callable,
) -> AsyncGenerator[ToolResponse, None]:
    """A middleware that logs tool execution."""
    # Access the tool call from kwargs
    tool_call = kwargs["tool_call"]

    # Pre-processing: log before tool execution
    print(f"[Middleware] Calling tool: {tool_call['name']}")
    print(f"[Middleware] Input: {tool_call['input']}")

    # Call the next handler (either another middleware or the actual tool)
    async for response in await next_handler(**kwargs):
        # Post-processing: log the response
        print(f"[Middleware] Response: {response.content[0]['text']}")
        yield response

    # This will execute after all responses are yielded
    print(f"[Middleware] Tool {tool_call['name']} completed")

Let’s register this middleware with a toolkit and test it:

async def search_tool(query: str) -> ToolResponse:
    """A simple search tool.

    Args:
        query (`str`):
            The search query.

    Returns:
        `ToolResponse`:
            The search result.
    """
    return ToolResponse(
        content=[
            TextBlock(
                type="text",
                text=f"Search results for '{query}'",
            ),
        ],
    )


async def example_logging_middleware() -> None:
    """Example of using logging middleware."""
    # Create a toolkit and register the tool
    toolkit = Toolkit()
    toolkit.register_tool_function(search_tool)

    # Register the middleware
    toolkit.register_middleware(logging_middleware)

    # Call the tool
    result = await toolkit.call_tool_function(
        ToolUseBlock(
            type="tool_use",
            id="1",
            name="search_tool",
            input={"query": "AgentScope"},
        ),
    )

    async for response in result:
        print(f"\n[Final] {response.content[0]['text']}\n")


print("=" * 60)
print("Example 1: Logging Middleware")
print("=" * 60)
asyncio.run(example_logging_middleware())

============================================================
Example 1: Logging Middleware
============================================================
[Middleware] Calling tool: search_tool
[Middleware] Input: {'query': 'AgentScope'}
[Middleware] Response: Search results for 'AgentScope'

[Final] Search results for 'AgentScope'

[Middleware] Tool search_tool completed

Modifying Input and Output

Middleware can also modify the tool call input and the response content:

async def transform_middleware(
    kwargs: dict,
    next_handler: Callable,
) -> AsyncGenerator[ToolResponse, None]:
    """A middleware that transforms input and output."""
    # Access the tool call from kwargs
    tool_call = kwargs["tool_call"]

    # Pre-processing: modify the input
    original_query = tool_call["input"]["query"]
    tool_call["input"]["query"] = f"[TRANSFORMED] {original_query}"

    async for response in await next_handler(**kwargs):
        # Post-processing: modify the response
        original_text = response.content[0]["text"]
        response.content[0]["text"] = f"{original_text} [MODIFIED]"
        yield response


async def example_transform_middleware() -> None:
    """Example of transforming middleware."""
    toolkit = Toolkit()
    toolkit.register_tool_function(search_tool)
    toolkit.register_middleware(transform_middleware)

    result = await toolkit.call_tool_function(
        ToolUseBlock(
            type="tool_use",
            id="2",
            name="search_tool",
            input={"query": "middleware"},
        ),
    )

    async for response in result:
        print(f"Result: {response.content[0]['text']}")


print("\n" + "=" * 60)
print("Example 2: Transform Middleware")
print("=" * 60)
asyncio.run(example_transform_middleware())

============================================================
Example 2: Transform Middleware
============================================================
Result: Search results for '[TRANSFORMED] middleware' [MODIFIED]

Authorization Middleware

You can use middleware to implement authorization checks and skip tool execution if not authorized:

async def authorization_middleware(
    kwargs: dict,
    next_handler: Callable,
) -> AsyncGenerator[ToolResponse, None]:
    """A middleware that checks authorization."""
    # Access the tool call from kwargs
    tool_call = kwargs["tool_call"]

    # Check if the tool is authorized (simple example)
    authorized_tools = {"search_tool"}

    if tool_call["name"] not in authorized_tools:
        # Skip execution and return error directly
        print(f"[Auth] Tool {tool_call['name']} is not authorized")
        yield ToolResponse(
            content=[
                TextBlock(
                    type="text",
                    text=f"Error: Tool '{tool_call['name']}' is not authorized",  # noqa: E501
                ),
            ],
        )
        return

    # Tool is authorized, proceed
    print(f"[Auth] Tool {tool_call['name']} is authorized")
    async for response in await next_handler(**kwargs):
        yield response


async def unauthorized_tool(data: str) -> ToolResponse:
    """An unauthorized tool.

    Args:
        data (`str`):
            Some data.

    Returns:
        `ToolResponse`:
            The result.
    """
    return ToolResponse(
        content=[TextBlock(type="text", text=f"Processing {data}")],
    )


async def example_authorization_middleware() -> None:
    """Example of authorization middleware."""
    toolkit = Toolkit()
    toolkit.register_tool_function(search_tool)
    toolkit.register_tool_function(unauthorized_tool)
    toolkit.register_middleware(authorization_middleware)

    # Try authorized tool
    print("\nCalling authorized tool:")
    result = await toolkit.call_tool_function(
        ToolUseBlock(
            type="tool_use",
            id="3",
            name="search_tool",
            input={"query": "test"},
        ),
    )
    async for response in result:
        print(f"Result: {response.content[0]['text']}")

    # Try unauthorized tool
    print("\nCalling unauthorized tool:")
    result = await toolkit.call_tool_function(
        ToolUseBlock(
            type="tool_use",
            id="4",
            name="unauthorized_tool",
            input={"data": "test"},
        ),
    )
    async for response in result:
        print(f"Result: {response.content[0]['text']}")


print("\n" + "=" * 60)
print("Example 3: Authorization Middleware")
print("=" * 60)
asyncio.run(example_authorization_middleware())

============================================================
Example 3: Authorization Middleware
============================================================

Calling authorized tool:
[Auth] Tool search_tool is authorized
Result: Search results for 'test'

Calling unauthorized tool:
[Auth] Tool unauthorized_tool is not authorized
Result: Error: Tool 'unauthorized_tool' is not authorized

Multiple Middleware (Onion Model)

When multiple middleware are registered, they form an onion-like structure. The execution order follows the onion model:

• Pre-processing: Executes in the order middleware are registered

• Post-processing: Executes in reverse order (inner to outer)

This is because the actual tool response object is passed through the middleware chain, and each middleware modifies it in place.

async def middleware_1(
    kwargs: dict,
    next_handler: Callable,
) -> AsyncGenerator[ToolResponse, None]:
    """First middleware."""
    # Access the tool call from kwargs
    tool_call = kwargs["tool_call"]

    # Pre-processing
    print("[M1] Pre-processing")
    tool_call["input"]["query"] += " [M1]"

    async for response in await next_handler(**kwargs):
        # Post-processing
        response.content[0]["text"] += " [M1]"
        print("[M1] Post-processing")
        yield response


async def middleware_2(
    kwargs: dict,
    next_handler: Callable,
) -> AsyncGenerator[ToolResponse, None]:
    """Second middleware."""
    # Access the tool call from kwargs
    tool_call = kwargs["tool_call"]

    # Pre-processing
    print("[M2] Pre-processing")
    tool_call["input"]["query"] += " [M2]"

    async for response in await next_handler(**kwargs):
        # Post-processing
        response.content[0]["text"] += " [M2]"
        print("[M2] Post-processing")
        yield response


async def example_multiple_middleware() -> None:
    """Example of multiple middleware."""
    toolkit = Toolkit()
    toolkit.register_tool_function(search_tool)

    # Register middleware in order
    toolkit.register_middleware(middleware_1)
    toolkit.register_middleware(middleware_2)

    result = await toolkit.call_tool_function(
        ToolUseBlock(
            type="tool_use",
            id="5",
            name="search_tool",
            input={"query": "test"},
        ),
    )

    async for response in result:
        print(f"\nFinal result: {response.content[0]['text']}")


print("\n" + "=" * 60)
print("Example 4: Multiple Middleware (Onion Model)")
print("=" * 60)
print("\nExecution flow:")
print("M1 Pre → M2 Pre → Tool → M2 Post → M1 Post")
print()
asyncio.run(example_multiple_middleware())

============================================================
Example 4: Multiple Middleware (Onion Model)
============================================================

Execution flow:
M1 Pre → M2 Pre → Tool → M2 Post → M1 Post

[M1] Pre-processing
[M2] Pre-processing
[M2] Post-processing
[M1] Post-processing

Final result: Search results for 'test [M1] [M2]' [M2] [M1]

Use Cases

The middleware system is useful for various scenarios:

• Logging and Monitoring: Track tool usage and performance

• Authorization: Control access to specific tools

• Rate Limiting: Limit the frequency of tool calls

• Caching: Cache tool responses for repeated calls

• Error Handling: Add retry logic or graceful degradation

• Input Validation: Validate and sanitize tool inputs

• Output Transformation: Format or filter tool outputs

• Metrics Collection: Collect statistics about tool usage

Note

• Middleware are applied in the order they are registered

• The same ToolResponse object is passed through the middleware chain and modified in place

• Middleware can completely skip tool execution by not calling next_handler

• All middleware must be async generator functions that yield ToolResponse objects