From f673262edb7769f0f8edc98cd403f8f421fc98e7 Mon Sep 17 00:00:00 2001 From: mkzung <103102868+mkzung@users.noreply.github.com> Date: Sat, 7 Nov 2026 21:20:52 +0000 Subject: [PATCH 1/2] fix(client): preserve caller headers and content type overrides --- python/restate/client.py | 5 ++- tests/client_headers.py | 67 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 69 insertions(+), 3 deletions(-) create mode 100644 tests/client_headers.py diff --git a/python/restate/client.py b/python/restate/client.py index 8c24ab3e..2934ae1a 100644 --- a/python/restate/client.py +++ b/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 diff --git a/tests/client_headers.py b/tests/client_headers.py new file mode 100644 index 00000000..6a3732d7 --- /dev/null +++ b/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) From 50cf9bf51943b4f955e24ca7b7fc70c619f1bc86 Mon Sep 17 00:00:00 2001 From: Alex Castillo Date: Tue, 3 Nov 2026 09:16:35 -0500 Subject: [PATCH 2/2] fix(pydantic): import restate.ext.pydantic without the pydantic-ai mcp extra _toolset.py imported pydantic_ai.mcp at module level, so importing restate.ext.pydantic failed unless pydantic-ai-slim[mcp] was installed, although _agent.py already treats MCP support as optional. Move RestateMCPToolset into its own module that is only imported lazily. --- python/restate/ext/pydantic/_agent.py | 2 +- python/restate/ext/pydantic/_mcp_toolset.py | 100 ++++++++++++++++++++ python/restate/ext/pydantic/_toolset.py | 86 +---------------- 3 files changed, 102 insertions(+), 86 deletions(-) create mode 100644 python/restate/ext/pydantic/_mcp_toolset.py diff --git a/python/restate/ext/pydantic/_agent.py b/python/restate/ext/pydantic/_agent.py index 22efe247..a9630249 100644 --- a/python/restate/ext/pydantic/_agent.py +++ b/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: diff --git a/python/restate/ext/pydantic/_mcp_toolset.py b/python/restate/ext/pydantic/_mcp_toolset.py new file mode 100644 index 00000000..2a0b28b4 --- /dev/null +++ b/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 diff --git a/python/restate/ext/pydantic/_toolset.py b/python/restate/ext/pydantic/_toolset.py index 582ed80a..7852f4d9 100644 --- a/python/restate/ext/pydantic/_toolset.py +++ b/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