Client keeps the caller's headers; restate.ext.pydantic imports without the MCP extra

python/restate/client.py

@@ -85,9 +85,8 @@ async def do_call(
         else:
             content_type = target_handler.handler_io.content_type
 
-        if headers is None:
-            headers = {}
-        if headers.get("Content-Type") is None and content_type is not None:
+        headers = dict(headers) if headers is not None else {}
+        if content_type is not None and not any(name.lower() == "content-type" for name in headers):
             headers["Content-Type"] = content_type
 
         service = target_handler.service_tag.name

python/restate/ext/pydantic/_agent.py

@@ -122,7 +122,7 @@ def set_context(toolset: AbstractToolset[AgentDepsT]) -> AbstractToolset[AgentDe
             try:
                 from pydantic_ai.mcp import MCPToolset
 
-                from ._toolset import RestateMCPToolset
+                from ._mcp_toolset import RestateMCPToolset
             except ImportError:
                 pass
             else:

python/restate/ext/pydantic/_mcp_toolset.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, replace
+from typing import Any
+
+from restate import RunOptions, SdkInternalBaseException
+from restate.extensions import current_context
+
+from pydantic_ai import ToolDefinition
+from pydantic_ai._run_context import AgentDepsT
+from pydantic_ai.exceptions import UserError
+from pydantic_ai.mcp import MCPToolset
+from pydantic_ai.tools import RunContext
+from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool
+from pydantic_ai.toolsets.wrapper import WrapperToolset
+
+from ._serde import PydanticTypeAdapter
+
+
+@dataclass
+class RestateMCPGetToolsContextRunResult:
+    """A simple wrapper for tool results to be used with Restate's run_typed."""
+
+    output: dict[str, ToolDefinition]
+
+
+MCP_GET_TOOLS_SERDE = PydanticTypeAdapter(RestateMCPGetToolsContextRunResult)
+
+
+@dataclass
+class RestateMCPToolRunResult:
+    """A simple wrapper for tool results to be used with Restate's run_typed."""
+
+    output: Any
+
+
+MCP_RUN_SERDE = PydanticTypeAdapter(RestateMCPToolRunResult)
+
+
+class RestateMCPToolset(WrapperToolset[AgentDepsT]):
+    """A wrapper for MCPToolset that integrates with Restate."""
+
+    def __init__(self, wrapped: MCPToolset[AgentDepsT], run_options: RunOptions):
+        super().__init__(wrapped)
+        self._wrapped = wrapped
+        self.get_tools_options = replace(run_options, serde=MCP_GET_TOOLS_SERDE)
+        self.run_tools_options = replace(run_options, serde=MCP_RUN_SERDE)
+
+    def visit_and_replace(
+        self, visitor: Callable[[AbstractToolset[AgentDepsT]], AbstractToolset[AgentDepsT]]
+    ) -> AbstractToolset[AgentDepsT]:
+        return visitor(self)
+
+    async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
+        async def get_tools_in_context() -> RestateMCPGetToolsContextRunResult:
+            res = await self._wrapped.get_tools(ctx)
+            # ToolsetTool is not serializable as it holds a SchemaValidator
+            # (which is also the same for every MCP tool so unnecessary to pass along the wire every time),
+            # so we just return the ToolDefinitions and wrap them in ToolsetTool outside of the activity.
+            return RestateMCPGetToolsContextRunResult(output={name: tool.tool_def for name, tool in res.items()})
+
+        context = current_context()
+        if context is None:
+            raise UserError(
+                "A toolset cannot be used without a Restate context. Make sure to run it within an agent or a run context."
+            )
+
+        try:
+            tool_defs = await context.run_typed("get mcp tools", get_tools_in_context, self.get_tools_options)
+            return {name: self.tool_for_tool_def(tool_def, ctx=ctx) for name, tool_def in tool_defs.output.items()}
+        except SdkInternalBaseException as e:
+            raise Exception("Internal error during get_tools call") from e
+
+    def tool_for_tool_def(self, tool_def: ToolDefinition, *, ctx: RunContext[AgentDepsT]) -> ToolsetTool[AgentDepsT]:
+        assert isinstance(self.wrapped, MCPToolset)
+        return self.wrapped.tool_for_tool_def(tool_def, ctx=ctx)
+
+    async def call_tool(
+        self,
+        name: str,
+        tool_args: dict[str, Any],
+        ctx: RunContext[AgentDepsT],
+        tool: ToolsetTool[AgentDepsT],
+    ) -> Any:
+        async def call_tool_in_context() -> RestateMCPToolRunResult:
+            res = await self._wrapped.call_tool(name, tool_args, ctx, tool)
+            return RestateMCPToolRunResult(output=res)
+
+        context = current_context()
+        if context is None:
+            raise UserError(
+                "A toolset cannot be used without a Restate context. Make sure to run it within an agent or a run context."
+            )
+        try:
+            res = await context.run_typed(f"Calling mcp tool {name}", call_tool_in_context, self.run_tools_options)
+        except SdkInternalBaseException as e:
+            raise Exception("Internal error during tool call") from e
+
+        return res.output

python/restate/ext/pydantic/_toolset.py

@@ -4,14 +4,12 @@
 from dataclasses import dataclass, replace
 from typing import Any, Literal
 
-from restate import RunOptions, SdkInternalBaseException, TerminalError
+from restate import RunOptions, TerminalError
 from restate.exceptions import SdkInternalException
 from restate.extensions import current_context
 
-from pydantic_ai import ToolDefinition
 from pydantic_ai._run_context import AgentDepsT
 from pydantic_ai.exceptions import ApprovalRequired, CallDeferred, ModelRetry, UserError
-from pydantic_ai.mcp import MCPToolset
 from pydantic_ai.tools import RunContext
 from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool
 from pydantic_ai.toolsets.wrapper import WrapperToolset
@@ -32,26 +30,6 @@ class RestateContextRunResult:
 CONTEXT_RUN_SERDE = PydanticTypeAdapter(RestateContextRunResult)
 
 
-@dataclass
-class RestateMCPGetToolsContextRunResult:
-    """A simple wrapper for tool results to be used with Restate's run_typed."""
-
-    output: dict[str, ToolDefinition]
-
-
-MCP_GET_TOOLS_SERDE = PydanticTypeAdapter(RestateMCPGetToolsContextRunResult)
-
-
-@dataclass
-class RestateMCPToolRunResult:
-    """A simple wrapper for tool results to be used with Restate's run_typed."""
-
-    output: Any
-
-
-MCP_RUN_SERDE = PydanticTypeAdapter(RestateMCPToolRunResult)
-
-
 class RestateContextRunToolSet(WrapperToolset[AgentDepsT]):
     """A toolset that automatically wraps tool calls with restate's `ctx.run_typed()`."""
 
@@ -114,65 +92,3 @@ def visit_and_replace(
         self, visitor: Callable[[AbstractToolset[AgentDepsT]], AbstractToolset[AgentDepsT]]
     ) -> AbstractToolset[AgentDepsT]:
         return visitor(self)
-
-
-class RestateMCPToolset(WrapperToolset[AgentDepsT]):
-    """A wrapper for MCPToolset that integrates with Restate."""
-
-    def __init__(self, wrapped: MCPToolset[AgentDepsT], run_options: RunOptions):
-        super().__init__(wrapped)
-        self._wrapped = wrapped
-        self.get_tools_options = replace(run_options, serde=MCP_GET_TOOLS_SERDE)
-        self.run_tools_options = replace(run_options, serde=MCP_RUN_SERDE)
-
-    def visit_and_replace(
-        self, visitor: Callable[[AbstractToolset[AgentDepsT]], AbstractToolset[AgentDepsT]]
-    ) -> AbstractToolset[AgentDepsT]:
-        return visitor(self)
-
-    async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
-        async def get_tools_in_context() -> RestateMCPGetToolsContextRunResult:
-            res = await self._wrapped.get_tools(ctx)
-            # ToolsetTool is not serializable as it holds a SchemaValidator
-            # (which is also the same for every MCP tool so unnecessary to pass along the wire every time),
-            # so we just return the ToolDefinitions and wrap them in ToolsetTool outside of the activity.
-            return RestateMCPGetToolsContextRunResult(output={name: tool.tool_def for name, tool in res.items()})
-
-        context = current_context()
-        if context is None:
-            raise UserError(
-                "A toolset cannot be used without a Restate context. Make sure to run it within an agent or a run context."
-            )
-
-        try:
-            tool_defs = await context.run_typed("get mcp tools", get_tools_in_context, self.get_tools_options)
-            return {name: self.tool_for_tool_def(tool_def, ctx=ctx) for name, tool_def in tool_defs.output.items()}
-        except SdkInternalBaseException as e:
-            raise Exception("Internal error during get_tools call") from e
-
-    def tool_for_tool_def(self, tool_def: ToolDefinition, *, ctx: RunContext[AgentDepsT]) -> ToolsetTool[AgentDepsT]:
-        assert isinstance(self.wrapped, MCPToolset)
-        return self.wrapped.tool_for_tool_def(tool_def, ctx=ctx)
-
-    async def call_tool(
-        self,
-        name: str,
-        tool_args: dict[str, Any],
-        ctx: RunContext[AgentDepsT],
-        tool: ToolsetTool[AgentDepsT],
-    ) -> Any:
-        async def call_tool_in_context() -> RestateMCPToolRunResult:
-            res = await self._wrapped.call_tool(name, tool_args, ctx, tool)
-            return RestateMCPToolRunResult(output=res)
-
-        context = current_context()
-        if context is None:
-            raise UserError(
-                "A toolset cannot be used without a Restate context. Make sure to run it within an agent or a run context."
-            )
-        try:
-            res = await context.run_typed(f"Calling mcp tool {name}", call_tool_in_context, self.run_tools_options)
-        except SdkInternalBaseException as e:
-            raise Exception("Internal error during tool call") from e
-
-        return res.output

tests/client_headers.py

@@ -0,0 +1,67 @@
+#
+#  Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH
+#
+#  This file is part of the Restate SDK for Python,
+#  which is released under the MIT license.
+#
+#  You can find a copy of the license in file LICENSE in the root
+#  directory of this repository or package, or at
+#  https://github.com/restatedev/sdk-typescript/blob/main/LICENSE
+#
+"""Per-call headers are reusable across handlers with different content types."""
+
+import httpx2
+import pytest
+
+import restate
+from restate.client import Client
+from restate.serde import BytesSerde
+
+
+@pytest.fixture(scope="session")
+def anyio_backend():
+    return "asyncio"
+
+
+echo = restate.Service("echo")
+
+
+@echo.handler()
+async def text(ctx: restate.Context, value: str) -> str:
+    return value
+
+
+@echo.handler(
+    accept="application/octet-stream",
+    content_type="application/octet-stream",
+    input_serde=BytesSerde(),
+    output_serde=BytesSerde(),
+)
+async def binary(ctx: restate.Context, value: bytes) -> bytes:
+    return value
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize("header_name", [None, "Content-Type", "content-type", "CONTENT-TYPE"])
+async def test_call_headers_are_preserved(header_name: str | None):
+    headers = {"X-Correlation-Id": "shared"}
+    if header_name is not None:
+        headers[header_name] = "application/custom"
+    original = headers.copy()
+    content_types: list[list[str]] = []
+
+    def respond(request: httpx2.Request) -> httpx2.Response:
+        content_types.append(request.headers.get_list("content-type"))
+        return httpx2.Response(200, content=request.content)
+
+    async with httpx2.AsyncClient(base_url="http://restate.test", transport=httpx2.MockTransport(respond)) as http:
+        client = Client(http)
+        assert await client.service_call(text, "hello", headers=headers) == "hello"
+        assert await client.service_call(binary, b"hello", headers=headers) == b"hello"
+
+    expected = (
+        [["application/json"], ["application/octet-stream"]]
+        if header_name is None
+        else [["application/custom"], ["application/custom"]]
+    )
+    assert (headers, content_types) == (original, expected)