Merge branch 'main' into local-branch-python-add-reset-to-workflow

python/packages/core/AGENTS.md

@@ -80,6 +80,8 @@ agent_framework/
 
 - **`MCPTool`** - Base wrapper that owns the MCP `ClientSession` and exposes the remote server's tools as `FunctionTool`s.
 - **`MCPStdioTool`** / **`MCPStreamableHTTPTool`** / **`MCPWebsocketTool`** - Transport-specific subclasses.
+- **Argument allowlist (`_prepare_call_kwargs`)** - Before each `tools/call`, kwargs are filtered to an **allowlist** built from the tool's declared parameters (`inputSchema.properties`) plus any user-configured extras. Framework runtime kwargs injected through the function-invocation pipeline (e.g. `thread`, `conversation_id`, `chat_options`, `options`, `response_format`) are stripped by default rather than forwarded. A tool that declares no usable `properties` (including schemas with `additionalProperties: true`) forwards only the configured extras. The `_MCP_FRAMEWORK_DENYLIST` is a safety net for framework-named params a server *declares* in its schema (those are dropped); names explicitly opted in via `additional_tool_argument_names` always win. The reserved `_meta` key is extracted as MCP request metadata, never forwarded as an argument.
+- **`additional_tool_argument_names`** (constructor arg on all `MCPTool` subclasses) - Opt extra argument names back into the allowlist. Accepts a `Sequence[str]` (applied to every tool) or a `Mapping[str, Sequence[str]]` keyed by **remote tool name**, where the reserved key `"*"` denotes global extras. It is configured only in user code at construction; there is **no per-call/runtime override**, so a model-issued tool call cannot change which names pass through. To use a server that accepts `additionalProperties: true`, list the extra names here and then either (1) manually extend that tool's `inputSchema` (via the `.functions` list after connecting) so the model is prompted to supply them, or (2) supply the values yourself via `function_invocation_kwargs`. If a name is supplied by both the model and `function_invocation_kwargs`, the model-supplied value wins.
 - **`MCPTaskOptions`** (experimental, `MCP_LONG_RUNNING_TASKS` feature, **frozen**) - Per-tool-instance options controlling the SEP-2663 long-running task lifecycle. When the server advertises a tool with `execution.taskSupport == "required"`, `MCPTool.call_tool` transparently routes through `call_tool_as_task`, which sends an augmented `tools/call`, polls `tasks/get` until terminal, and reinterprets `tasks/result` as a normal `CallToolResult`. Instances are immutable; replace via `MCPTool.task_options = MCPTaskOptions(...)`. Fields:
   - `default_ttl: timedelta | None` — forwarded to the server as `params.task.ttl` (milliseconds). When `None`, the server's default applies.
   - `cancel_remote_task_on_local_cancellation: bool = True` — only gates the `CancelledError` path. Abandonment paths (see below) always cancel.

python/packages/core/agent_framework/_agents.py

@@ -92,12 +92,16 @@ OptionsCoT = TypeVar(
 def _merge_options(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]:
     """Merge two options dicts, with override values taking precedence.
 
+    ``None`` is treated as "unset": ``None`` overrides are skipped so they don't clobber a base
+    value, and the merged result is stripped of any remaining ``None`` values in a final pass so
+    unset options are never forwarded (e.g. an unset ``store`` is left for the service to default).
+
     Args:
         base: The base options dict.
         override: The override options dict (values take precedence).
 
     Returns:
-        A new merged options dict.
+        A new merged options dict containing no ``None`` values.
     """
     result = dict(base)
 
@@ -123,7 +127,7 @@ def _merge_options(base: dict[str, Any], override: dict[str, Any]) -> dict[str,
             result["instructions"] = f"{result['instructions']}\n{value}"
         else:
             result[key] = value
-    return result
+    return {key: value for key, value in result.items() if value is not None}
 
 
 def _sanitize_agent_name(agent_name: str | None) -> str | None:
@@ -460,6 +464,9 @@ class BaseAgent(SerializationMixin):
         if provider_session is None and self.context_providers:
             provider_session = AgentSession()
 
+        # When per-service-call persistence is enabled, the per-service-call middleware owns
+        # HistoryProvider persistence (in both the local and service-managed cases), so skip
+        # them on the once-per-run path to avoid double persistence.
         per_service_call_history_required = self.require_per_service_call_history_persistence and any(
             isinstance(provider, HistoryProvider) for provider in self.context_providers
         )
@@ -686,11 +693,16 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]):  # type: ignore[misc]
             description: A brief description of the agent's purpose.
             context_providers: Context providers to include during agent invocation.
             middleware: List of middleware to intercept agent and function invocations.
-            require_per_service_call_history_persistence: When True, history providers are invoked
-                around each model call instead of once per ``run()`` when the service
-                is not already storing history. If service-side storage is active for
-                the run, the agent skips local history providers and relies on the
-                service-managed conversation instead.
+            require_per_service_call_history_persistence: When True (and a HistoryProvider is
+                present), the provider always persists history via per-service-call middleware,
+                regardless of whether the client stores history server-side. If the client does
+                not store history, the middleware also loads providers around each model call and
+                drives the function loop with a local conversation; if it does, loading is skipped
+                (the service-managed conversation is the source of truth) and the middleware only
+                persists. A warning is logged for providers with ``load_messages=True`` when
+                loading is skipped because service-side storage is active. When no HistoryProvider
+                is present, this flag has no effect (no middleware is installed and nothing is
+                persisted).
             default_options: A TypedDict containing chat options. When using a typed agent like
                 ``Agent[OpenAIChatOptions]``, this enables IDE autocomplete for
                 provider-specific options including temperature, max_tokens, model,
@@ -791,22 +803,20 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]):  # type: ignore[misc]
         self,
         *,
         session: AgentSession | None,
-        options: Mapping[str, Any] | None,
+        conversation_id: str | None,
         service_stores_history: bool,
     ) -> list[HistoryProvider]:
         history_providers = self._get_history_providers()
         if not self.require_per_service_call_history_persistence or not history_providers:
             return []
 
-        conversation_id = (
-            session.service_session_id
-            if session and session.service_session_id
-            else cast(str | None, (options or {}).get("conversation_id") or self.default_options.get("conversation_id"))
-        )
-        if service_stores_history:
-            return []
-
-        if conversation_id is not None:
+        # A live service-managed session id takes precedence over the resolved conversation id.
+        if session and session.service_session_id:
+            conversation_id = session.service_session_id
+        # Without service-side storage the middleware persists locally and drives the function
+        # loop with a local sentinel, which cannot be reconciled with an existing service-managed
+        # conversation. When the service stores history, an existing conversation id is expected.
+        if conversation_id is not None and not service_stores_history:
             raise AgentInvalidRequestException(
                 "require_per_service_call_history_persistence cannot be used "
                 "with an existing service-managed conversation."
@@ -1167,18 +1177,34 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]):  # type: ignore[misc]
 
         input_messages = normalize_messages(messages)
 
-        # `store` in runtime or agent options takes precedence over client-level storage
-        # indicators. An explicit `store=False` forces local (in-memory) history injection,
-        # even if the client is configured to use service-side storage by default.
-        store_ = opts.get("store", self.default_options.get("store", getattr(self.client, "STORES_BY_DEFAULT", False)))
+        # Combine agent-level defaults with runtime options up front so the decisions below read
+        # `store` from a single place rather than introspecting both dicts. _merge_options applies
+        # the same precedence used for the actual client call (runtime wins; unset/None falls back
+        # to the agent default).
+        effective_options = _merge_options(self.default_options, opts)
+
+        # `store` in runtime or agent options takes precedence over the client's default
+        # storage behavior. An explicit `store=False` forces local (in-memory) history
+        # injection even when the client stores server-side by default; an explicit
+        # `store=True` forces service-side storage. A `store=None`/unset value means the
+        # service falls back to its own default.
+        explicit_store = effective_options.get("store")
+        # Internal behavior hint: will the service own history for this run? Only when the
+        # user left `store` unset do we fall back to the client's STORES_BY_DEFAULT.
+        service_stores_history = (
+            explicit_store if explicit_store is not None else getattr(self.client, "STORES_BY_DEFAULT", False)
+        )
+        # Resolve conversation_id from the same combined view so an agent-level default is honored
+        # when the runtime omits it (a live session id still takes precedence below).
+        effective_conversation_id = effective_options.get("conversation_id")
         # Auto-inject InMemoryHistoryProvider when session is provided, no context providers
         # registered, and no service-side storage indicators
         if (
             session is not None
             and not self.context_providers
             and not session.service_session_id
-            and not opts.get("conversation_id")
-            and not store_
+            and not effective_conversation_id
+            and not service_stores_history
         ):
             self.context_providers.append(InMemoryHistoryProvider())
 
@@ -1188,10 +1214,30 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]):  # type: ignore[misc]
 
         per_service_call_history_providers = self._resolve_per_service_call_history_providers(
             session=active_session,
-            options=opts,
-            service_stores_history=bool(store_),
+            conversation_id=effective_conversation_id,
+            service_stores_history=service_stores_history,
         )
 
+        # When require_per_service_call_history_persistence is set together with a
+        # HistoryProvider, the per-service-call middleware (installed below) always persists
+        # the provider. ``service_stores_history`` only selects how the middleware behaves:
+        # - service does not store: the middleware also loads providers and drives the function
+        #   loop with a local sentinel conversation id, or
+        # - service stores: the middleware skips loading (the service owns history) and simply
+        #   persists each service call while the real conversation id flows through.
+        # In the service-managed case loading is skipped, so warn for providers that expect to load.
+        history_providers = self._get_history_providers()
+        if self.require_per_service_call_history_persistence and history_providers and service_stores_history:
+            for provider in history_providers:
+                if provider.load_messages:
+                    logger.warning(
+                        "HistoryProvider '%s' has load_messages=True but the chat client stores history "
+                        "server-side; skipping local history load and relying on the service-managed "
+                        "conversation. Set store=False to load from the provider, or load_messages=False "
+                        "to silence this warning.",
+                        provider.source_id,
+                    )
+
         session_context, chat_options = await self._prepare_session_and_messages(
             session=active_session,
             input_messages=input_messages,
@@ -1265,8 +1311,8 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]):  # type: ignore[misc]
         }
         if model is not None:
             run_opts["model"] = model
-        # Remove None values and merge with chat_options
-        run_opts = {k: v for k, v in run_opts.items() if v is not None}
+        # _merge_options strips unset (None) options, so e.g. an unset `store` is not forwarded
+        # and the service decides its own default.
         co = _merge_options(chat_options, run_opts)
 
         # Build session_messages from session context: context messages + input messages
@@ -1280,6 +1326,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]):  # type: ignore[misc]
                 agent=self,
                 session=active_session,
                 providers=per_service_call_history_providers,
+                service_stores_history=service_stores_history,
             )
             existing_middleware = effective_client_kwargs.get("middleware")
             if isinstance(existing_middleware, Sequence) and not isinstance(existing_middleware, (str, bytes)):
@@ -1319,7 +1366,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]):  # type: ignore[misc]
             "input_messages": input_messages,
             "session_messages": session_messages,
             "agent_name": agent_name,
-            "suppress_response_id": bool(per_service_call_history_providers),
+            "suppress_response_id": bool(per_service_call_history_providers) and not service_stores_history,
             "chat_options": co,
             "compaction_strategy": compaction_strategy or self.compaction_strategy,
             "tokenizer": tokenizer or self.tokenizer,
@@ -1413,11 +1460,15 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]):  # type: ignore[misc]
             options=options or {},
         )
 
+        # When per-service-call persistence is enabled, the per-service-call middleware owns
+        # HistoryProvider loading (it loads locally when the service does not store history, or
+        # relies on the service when it does), so skip them on the once-per-run before_run path.
         per_service_call_history_required = self.require_per_service_call_history_persistence and bool(
             self._get_history_providers()
         )
 
-        # Run before_run providers (forward order, skip HistoryProvider when per-service-call persistence owns history)
+        # Run before_run providers (forward order, skip HistoryProvider when per-service-call
+        # persistence owns loading)
         for provider in self.context_providers:
             if per_service_call_history_required and isinstance(provider, HistoryProvider):
                 continue

python/packages/core/agent_framework/_clients.py

@@ -604,10 +604,13 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]):
                 and dict literals are accepted without specialized option typing.
             context_providers: Context providers to include during agent invocation.
             middleware: List of middleware to intercept agent and function invocations.
-            require_per_service_call_history_persistence: Whether to require per-service-call
-                chat history persistence. When enabled, history providers are invoked around
-                each model call instead of once per ``run()`` when the service is not already
-                storing history.
+            require_per_service_call_history_persistence: When enabled (and a HistoryProvider is
+                present), the provider always persists history after each model call. If the
+                client does not store history server-side, history providers are also loaded and
+                injected around each model call; if it does, provider loading is skipped and the
+                service-managed conversation is the source of truth (persistence still happens
+                after each model call). When no HistoryProvider is present, this flag has no
+                effect (no middleware is installed and nothing is persisted).
             function_invocation_configuration: Optional function invocation configuration override.
             compaction_strategy: Optional agent-level compaction override. When omitted,
                 client-level compaction defaults remain in effect for each call.

python/packages/core/agent_framework/_mcp.py

@@ -70,6 +70,31 @@ class MCPSpecificApproval(TypedDict, total=False):
 
 _MCP_REMOTE_NAME_KEY = "_mcp_remote_name"
 _MCP_NORMALIZED_NAME_KEY = "_mcp_normalized_name"
+# Reserved key in an ``additional_tool_argument_names`` mapping that applies its
+# values to every tool on the server rather than a single named tool.
+_MCP_GLOBAL_EXTRA_ARGS_KEY = "*"
+# Framework kwargs that flow through the function-invocation pipeline (via
+# ``FunctionInvocationContext.kwargs``) but must never be forwarded to an MCP
+# server: they are internal objects that the MCP SDK cannot serialize. They are
+# dropped as a safety net when a tool declares one of them in its schema, unless
+# the user explicitly opts the name back in via ``additional_tool_argument_names``
+# (explicit extras always win over the denylist).
+# - chat_options/tools/tool_choice/session/thread: framework runtime objects.
+# - conversation_id: internal tracking ID used by services like Azure AI.
+# - options: metadata/store used by AG-UI for Azure AI client requirements.
+# - response_format: a Pydantic model class for structured output (not serializable).
+# - _meta: reserved key extracted separately as MCP request metadata.
+_MCP_FRAMEWORK_DENYLIST: frozenset[str] = frozenset({
+    "chat_options",
+    "tools",
+    "tool_choice",
+    "session",
+    "thread",
+    "conversation_id",
+    "options",
+    "response_format",
+    "_meta",
+})
 _mcp_call_headers: contextvars.ContextVar[dict[str, str]] = contextvars.ContextVar("_mcp_call_headers")
 MCP_DEFAULT_TIMEOUT = 30
 MCP_DEFAULT_SSE_READ_TIMEOUT = 60 * 5
@@ -135,6 +160,34 @@ def _build_prefixed_mcp_name(
     return f"{normalized_prefix}_{trimmed_name}" if trimmed_name else normalized_prefix
 
 
+def _normalize_additional_tool_argument_names(
+    additional_tool_argument_names: Sequence[str] | Mapping[str, Sequence[str]] | None,
+) -> tuple[set[str], dict[str, set[str]]]:
+    """Split user-supplied extra argument names into global and per-tool sets.
+
+    Accepts either a sequence (applied to every tool) or a mapping keyed by remote
+    tool name, where the reserved key ``"*"`` is treated as global. Mapping values
+    may be a sequence or a single string. Returns a
+    ``(global_extras, per_tool_extras)`` tuple.
+    """
+    if additional_tool_argument_names is None:
+        return set(), {}
+    if isinstance(additional_tool_argument_names, str):
+        return {additional_tool_argument_names}, {}
+    if isinstance(additional_tool_argument_names, Mapping):
+        global_extras: set[str] = set()
+        per_tool_extras: dict[str, set[str]] = {}
+        for tool_name, names in additional_tool_argument_names.items():
+            # Treat a bare string value as a single name rather than iterating its characters.
+            names_set = {names} if isinstance(names, str) else set(names)
+            if tool_name == _MCP_GLOBAL_EXTRA_ARGS_KEY:
+                global_extras.update(names_set)
+            else:
+                per_tool_extras[tool_name] = names_set
+        return global_extras, per_tool_extras
+    return set(additional_tool_argument_names), {}
+
+
 def _inject_otel_into_mcp_meta(meta: dict[str, Any] | None = None) -> dict[str, Any] | None:
     """Inject OpenTelemetry trace context into MCP request _meta via the global propagator(s)."""
     carrier: dict[str, str] = {}
@@ -294,6 +347,7 @@ class MCPTool:
         client: SupportsChatGetResponse | None = None,
         additional_properties: dict[str, Any] | None = None,
         task_options: MCPTaskOptions | None = None,
+        additional_tool_argument_names: Sequence[str] | Mapping[str, Sequence[str]] | None = None,
     ) -> None:
         """Initialize the MCP Tool base.
 
@@ -328,6 +382,10 @@ class MCPTool:
             task_options: Options controlling how long-running MCP tasks are driven for
                 tools that advertise ``execution.taskSupport == "required"``. When ``None``,
                 the defaults from :class:`MCPTaskOptions` are used.
+            additional_tool_argument_names: Extra argument names to forward to the MCP server
+                in addition to each tool's declared parameters. A ``Sequence[str]`` applies to
+                every tool; a ``Mapping[str, Sequence[str]]`` is keyed by remote tool name with
+                ``"*"`` as a global key. See the transport subclasses for full details.
         """
         self.name = name
         self.description = description or ""
@@ -355,6 +413,10 @@ class MCPTool:
         self._functions: list[FunctionTool] = []
         self._tool_call_meta_by_name: dict[str, dict[str, Any]] = {}
         self._tool_task_support_by_name: dict[str, str] = {}
+        self._tool_param_names_by_name: dict[str, set[str]] = {}
+        self._global_extra_arg_names, self._tool_extra_arg_names = _normalize_additional_tool_argument_names(
+            additional_tool_argument_names
+        )
         self.is_connected: bool = False
         self._tools_loaded: bool = False
         self._prompts_loaded: bool = False
@@ -1229,6 +1291,7 @@ class MCPTool:
         existing_names = {func.name for func in self._functions}
         tool_call_meta_by_name: dict[str, dict[str, Any]] = {}
         tool_task_support_by_name: dict[str, str] = {}
+        tool_param_names_by_name: dict[str, set[str]] = {}
 
         params: types.PaginatedRequestParams | None = None
         while True:
@@ -1271,14 +1334,6 @@ class MCPTool:
                 if task_support is not None:
                     tool_task_support_by_name[tool.name] = task_support
 
-                normalized_name = _normalize_mcp_name(tool.name)
-                local_name = _build_prefixed_mcp_name(normalized_name, self.tool_name_prefix)
-
-                # Skip if already loaded
-                if local_name in existing_names:
-                    continue
-
-                approval_mode = self._determine_approval_mode(local_name, normalized_name, tool.name)
                 # Normalize inputSchema: ensure "properties" exists for object schemas.
                 # Some MCP servers (e.g. zero-argument tools) omit "properties",
                 # which causes OpenAI API to reject the schema with a 400 error.
@@ -1288,6 +1343,24 @@ class MCPTool:
                 if input_schema.get("type") == "object" and "properties" not in input_schema:
                     input_schema["properties"] = {}
 
+                # Register declared param names before the existing-tool skip below so that
+                # reloads (e.g. notifications/tools/list_changed) preserve the allowlist for
+                # tools that are already loaded, consistent with tool_call_meta_by_name and
+                # tool_task_support_by_name above.
+                schema_properties = input_schema.get("properties")
+                tool_param_names_by_name[tool.name] = (
+                    set(cast(dict[str, Any], schema_properties)) if isinstance(schema_properties, dict) else set()
+                )
+
+                normalized_name = _normalize_mcp_name(tool.name)
+                local_name = _build_prefixed_mcp_name(normalized_name, self.tool_name_prefix)
+
+                # Skip if already loaded
+                if local_name in existing_names:
+                    continue
+
+                approval_mode = self._determine_approval_mode(local_name, normalized_name, tool.name)
+
                 async def _call_tool_with_runtime_kwargs(
                     ctx: FunctionInvocationContext,
                     *,
@@ -1320,6 +1393,7 @@ class MCPTool:
 
         self._tool_call_meta_by_name = tool_call_meta_by_name
         self._tool_task_support_by_name = tool_task_support_by_name
+        self._tool_param_names_by_name = tool_param_names_by_name
 
     async def _close_on_owner(self) -> None:
         # Cancel any pending reload tasks before tearing down the session.
@@ -1530,10 +1604,14 @@ class MCPTool:
                 raise ToolExecutionException(f"Failed to call tool '{tool_name}'.", inner_exception=ex) from ex
         raise ToolExecutionException(f"Failed to call tool '{tool_name}' after retries.")
 
+    def _resolved_extra_args(self, tool_name: str) -> set[str]:
+        """Return the user-configured extra argument names allowed for a tool."""
+        return self._global_extra_arg_names | self._tool_extra_arg_names.get(tool_name, set())
+
     def _prepare_call_kwargs(
         self, tool_name: str, kwargs: dict[str, Any]
     ) -> tuple[dict[str, Any], dict[str, Any] | None]:
-        """Filter framework-only kwargs and build the merged MCP request metadata."""
+        """Filter kwargs down to the tool's arguments and build the merged MCP request metadata."""
         raw_user_meta: object | None = kwargs.get("_meta")
         user_meta: dict[str, Any] | None = None
         if raw_user_meta is not None and not isinstance(raw_user_meta, dict):
@@ -1546,27 +1624,28 @@ class MCPTool:
                     raise ToolExecutionException("MCP tool metadata provided via _meta must use string keys.")
                 user_meta[key] = value
 
-        # Filter out framework kwargs that cannot be serialized by the MCP SDK.
-        # These are internal objects passed through the function invocation pipeline
-        # that should not be forwarded to external MCP servers.
-        # conversation_id is an internal tracking ID used by services like Azure AI.
-        # options contains metadata/store used by AG-UI for Azure AI client requirements.
-        # response_format is a Pydantic model class used for structured output (not serializable).
+        # Allowlist: forward only the tool's declared parameters (from inputSchema.properties)
+        # plus any user-configured extra argument names. Everything else - notably the
+        # framework runtime kwargs injected through the function-invocation pipeline - is
+        # stripped so it is never forwarded to the MCP server. Tools that declare no usable
+        # properties forward only the user-configured extras.
+        #
+        # The extra names come exclusively from additional_tool_argument_names, which is set in
+        # user code at construction time; there is no per-call override, so a model-issued tool
+        # call cannot change which names are allowed through.
+        #
+        # The framework denylist acts as a safety net for keys a server *declares* in its
+        # schema that collide with internal, non-serializable framework objects (e.g. a tool
+        # that declares a parameter literally named "thread"): such declared-but-denylisted
+        # keys are dropped. Names the user explicitly opts in via additional_tool_argument_names
+        # always win. The reserved _meta key is handled separately above and never forwarded as
+        # an argument.
+        declared = self._tool_param_names_by_name.get(tool_name, set())
+        extras = self._resolved_extra_args(tool_name)
         filtered_kwargs = {
             k: v
             for k, v in kwargs.items()
-            if k
-            not in {
-                "chat_options",
-                "tools",
-                "tool_choice",
-                "session",
-                "thread",
-                "conversation_id",
-                "options",
-                "response_format",
-                "_meta",
-            }
+            if k != "_meta" and (k in extras or (k in declared and k not in _MCP_FRAMEWORK_DENYLIST))
         }
 
         # Some MCP proxies require their tools/list metadata to be echoed on tools/call.
@@ -2142,6 +2221,7 @@ class MCPStdioTool(MCPTool):
         client: SupportsChatGetResponse | None = None,
         additional_properties: dict[str, Any] | None = None,
         task_options: MCPTaskOptions | None = None,
+        additional_tool_argument_names: Sequence[str] | Mapping[str, Sequence[str]] | None = None,
         **kwargs: Any,
     ) -> None:
         """Initialize the MCP stdio tool.
@@ -2188,6 +2268,20 @@ class MCPStdioTool(MCPTool):
             client: The chat client to use for sampling.
             task_options: Options for tools that advertise
                 ``execution.taskSupport == "required"``. See :class:`MCPTaskOptions`.
+            additional_tool_argument_names: Extra argument names to forward to the MCP server in
+                addition to each tool's declared parameters (from its ``inputSchema.properties``).
+                By default only declared parameters are sent; framework runtime kwargs injected
+                through the function-invocation pipeline are stripped. Use this to opt specific
+                keys back in. Accepts either a ``Sequence[str]`` applied to every tool, or a
+                ``Mapping[str, Sequence[str]]`` keyed by remote tool name where the reserved key
+                ``"*"`` applies to every tool. This is configured only here in user code; there is
+                no per-call override, so a model-issued tool call cannot change which names pass
+                through. To use a server that accepts ``additionalProperties: true``, list the
+                extra names here and then either (1) manually extend that tool's ``inputSchema``
+                (via the ``.functions`` list after connecting) so the model is prompted to supply
+                them, or (2) supply the values yourself through ``function_invocation_kwargs``. If
+                a name is supplied via both the model and ``function_invocation_kwargs``, the
+                model-supplied value wins.
             kwargs: Any extra arguments to pass to the stdio client.
         """
         super().__init__(
@@ -2205,6 +2299,7 @@ class MCPStdioTool(MCPTool):
             parse_prompt_results=parse_prompt_results,
             request_timeout=request_timeout,
             task_options=task_options,
+            additional_tool_argument_names=additional_tool_argument_names,
         )
         self.command = command
         self.args = args or []
@@ -2284,6 +2379,7 @@ class MCPStreamableHTTPTool(MCPTool):
         http_client: AsyncClient | None = None,
         header_provider: Callable[[dict[str, Any]], dict[str, str]] | None = None,
         task_options: MCPTaskOptions | None = None,
+        additional_tool_argument_names: Sequence[str] | Mapping[str, Sequence[str]] | None = None,
         **kwargs: Any,
     ) -> None:
         """Initialize the MCP streamable HTTP tool.
@@ -2338,6 +2434,20 @@ class MCPStreamableHTTPTool(MCPTool):
                 agent middleware) without creating a separate ``httpx.AsyncClient``.
             task_options: Options for tools that advertise
                 ``execution.taskSupport == "required"``. See :class:`MCPTaskOptions`.
+            additional_tool_argument_names: Extra argument names to forward to the MCP server in
+                addition to each tool's declared parameters (from its ``inputSchema.properties``).
+                By default only declared parameters are sent; framework runtime kwargs injected
+                through the function-invocation pipeline are stripped. Use this to opt specific
+                keys back in. Accepts either a ``Sequence[str]`` applied to every tool, or a
+                ``Mapping[str, Sequence[str]]`` keyed by remote tool name where the reserved key
+                ``"*"`` applies to every tool. This is configured only here in user code; there is
+                no per-call override, so a model-issued tool call cannot change which names pass
+                through. To use a server that accepts ``additionalProperties: true``, list the
+                extra names here and then either (1) manually extend that tool's ``inputSchema``
+                (via the ``.functions`` list after connecting) so the model is prompted to supply
+                them, or (2) supply the values yourself through ``function_invocation_kwargs``. If
+                a name is supplied via both the model and ``function_invocation_kwargs``, the
+                model-supplied value wins.
             kwargs: Additional keyword arguments (accepted for backward compatibility but not used).
         """
         super().__init__(
@@ -2355,6 +2465,7 @@ class MCPStreamableHTTPTool(MCPTool):
             parse_prompt_results=parse_prompt_results,
             request_timeout=request_timeout,
             task_options=task_options,
+            additional_tool_argument_names=additional_tool_argument_names,
         )
         self.url = url
         self.terminate_on_close = terminate_on_close
@@ -2481,6 +2592,7 @@ class MCPWebsocketTool(MCPTool):
         client: SupportsChatGetResponse | None = None,
         additional_properties: dict[str, Any] | None = None,
         task_options: MCPTaskOptions | None = None,
+        additional_tool_argument_names: Sequence[str] | Mapping[str, Sequence[str]] | None = None,
         **kwargs: Any,
     ) -> None:
         """Initialize the MCP WebSocket tool.
@@ -2525,6 +2637,20 @@ class MCPWebsocketTool(MCPTool):
             client: The chat client to use for sampling.
             task_options: Options for tools that advertise
                 ``execution.taskSupport == "required"``. See :class:`MCPTaskOptions`.
+            additional_tool_argument_names: Extra argument names to forward to the MCP server in
+                addition to each tool's declared parameters (from its ``inputSchema.properties``).
+                By default only declared parameters are sent; framework runtime kwargs injected
+                through the function-invocation pipeline are stripped. Use this to opt specific
+                keys back in. Accepts either a ``Sequence[str]`` applied to every tool, or a
+                ``Mapping[str, Sequence[str]]`` keyed by remote tool name where the reserved key
+                ``"*"`` applies to every tool. This is configured only here in user code; there is
+                no per-call override, so a model-issued tool call cannot change which names pass
+                through. To use a server that accepts ``additionalProperties: true``, list the
+                extra names here and then either (1) manually extend that tool's ``inputSchema``
+                (via the ``.functions`` list after connecting) so the model is prompted to supply
+                them, or (2) supply the values yourself through ``function_invocation_kwargs``. If
+                a name is supplied via both the model and ``function_invocation_kwargs``, the
+                model-supplied value wins.
             kwargs: Any extra arguments to pass to the WebSocket client.
         """
         super().__init__(
@@ -2542,6 +2668,7 @@ class MCPWebsocketTool(MCPTool):
             parse_prompt_results=parse_prompt_results,
             request_timeout=request_timeout,
             task_options=task_options,
+            additional_tool_argument_names=additional_tool_argument_names,
         )
         self.url = url
         self._client_kwargs = kwargs

python/packages/core/agent_framework/_sessions.py

@@ -16,6 +16,7 @@ from __future__ import annotations
 import asyncio
 import copy
 import json
+import logging
 import threading
 import uuid
 import weakref
@@ -36,6 +37,8 @@ if TYPE_CHECKING:
     from ._middleware import MiddlewareTypes
 
 
+logger = logging.getLogger("agent_framework")
+
 # Registry of known types for state deserialization
 _STATE_TYPE_REGISTRY: dict[str, type] = {}
 
@@ -580,6 +583,7 @@ class PerServiceCallHistoryPersistingMiddleware(ChatMiddleware):
         agent: SupportsAgentRun,
         session: AgentSession,
         providers: Sequence[HistoryProvider],
+        service_stores_history: bool = False,
     ) -> None:
         """Initialize the middleware.
 
@@ -587,10 +591,16 @@ class PerServiceCallHistoryPersistingMiddleware(ChatMiddleware):
             agent: The agent that owns the history providers.
             session: The active session for the current run.
             providers: The history providers participating in per-service-call persistence.
+            service_stores_history: When True, the chat client stores history server-side. The
+                middleware then skips loading providers and leaves the real conversation id
+                untouched, persisting each service call without driving the function loop with a
+                local sentinel. When False, the middleware loads providers and uses a local
+                sentinel conversation id so the function loop runs without service-side storage.
         """
         self._agent = agent
         self._session = session
         self._providers = list(providers)
+        self._service_stores_history = service_stores_history
 
     async def _prepare_service_call_context(self, messages: Sequence[Message]) -> SessionContext:
         """Create a per-call SessionContext and load history providers into it."""
@@ -602,6 +612,9 @@ class PerServiceCallHistoryPersistingMiddleware(ChatMiddleware):
         )
         for source_id, source_messages in context_messages.items():
             service_call_context.extend_messages(source_id, source_messages)
+        # When the service stores history, it owns loading; the providers are write-only sinks.
+        if self._service_stores_history:
+            return service_call_context
         for provider in self._providers:
             if not provider.load_messages:
                 continue
@@ -652,17 +665,35 @@ class PerServiceCallHistoryPersistingMiddleware(ChatMiddleware):
         response: ChatResponse,
     ) -> ChatResponse:
         """Persist a model response and apply the local follow-up sentinel when needed."""
-        if response.conversation_id is not None and not is_local_history_conversation_id(response.conversation_id):
+        if (
+            not self._service_stores_history
+            and response.conversation_id is not None
+            and not is_local_history_conversation_id(response.conversation_id)
+        ):
             raise ChatClientInvalidResponseException(
                 "require_per_service_call_history_persistence cannot be used "
                 "when the chat client returns a real conversation_id."
             )
 
+        # In storing mode the service is expected to echo a conversation id that the next run
+        # resumes from. If it comes back empty, the provider still captures this turn but there is
+        # no service id to load from next time, so cross-turn history can be lost silently. Warn
+        # every time so this uncommon, easy-to-miss failure mode cannot fail quietly.
+        if self._service_stores_history and response.conversation_id is None:
+            logger.warning(
+                "require_per_service_call_history_persistence is enabled with a chat client that "
+                "stores history server-side, but the client returned no conversation_id; cross-turn "
+                "history may not resume. Set store=False to load and resume from the HistoryProvider "
+                "instead."
+            )
+
         await self._persist_service_call_response(
             service_call_context=service_call_context,
             response=response,
         )
-        if _response_contains_follow_up_request(response):
+        # The local sentinel only applies when the service does not store history; when it does,
+        # the real conversation id already drives function-loop continuation.
+        if not self._service_stores_history and _response_contains_follow_up_request(response):
             response.mark_internal_conversation_id()
             response.conversation_id = LOCAL_HISTORY_CONVERSATION_ID
         return response
@@ -681,8 +712,12 @@ class PerServiceCallHistoryPersistingMiddleware(ChatMiddleware):
                 result type for streaming or non-streaming execution.
         """
         service_call_context = await self._prepare_service_call_context(context.messages)
-        context.messages = service_call_context.get_messages(include_input=True)
-        self._strip_local_conversation_id(context)
+        # When the service stores history, leave the outgoing messages and the real conversation
+        # id untouched (pass-through); the middleware only persists. Otherwise reconstruct the
+        # outgoing messages from the loaded local history and strip the local sentinel.
+        if not self._service_stores_history:
+            context.messages = service_call_context.get_messages(include_input=True)
+            self._strip_local_conversation_id(context)
 
         await call_next()
 

python/packages/core/tests/core/test_agents.py

@@ -3,6 +3,7 @@
 import contextlib
 import inspect
 import json
+import logging
 from collections.abc import AsyncIterable, Awaitable, Callable, MutableSequence, Sequence
 from typing import Any, cast
 from unittest.mock import AsyncMock, MagicMock, patch
@@ -42,6 +43,8 @@ from agent_framework._mcp import MCPTool, _build_prefixed_mcp_name, _normalize_m
 from agent_framework._middleware import FunctionInvocationContext
 from agent_framework.exceptions import AgentInvalidRequestException, ChatClientInvalidResponseException
 
+from .conftest import MockBaseChatClient
+
 
 class _FixedTokenizer:
     def __init__(self, token_count: int) -> None:
@@ -609,6 +612,7 @@ async def test_streaming_per_service_call_persistence_hides_response_id_from_aft
 
 async def test_per_service_call_persistence_uses_real_service_storage_when_client_stores_by_default(
     chat_client_base: SupportsChatGetResponse,
+    caplog: pytest.LogCaptureFixture,
 ) -> None:
     provider = _RecordingHistoryProvider()
 
@@ -649,15 +653,22 @@ async def test_per_service_call_persistence_uses_real_service_storage_when_clien
         require_per_service_call_history_persistence=True,
     )
 
-    result = await agent.run("What's the weather in Seattle?", session=session)
+    with caplog.at_level(logging.WARNING, logger="agent_framework"):
+        result = await agent.run("What's the weather in Seattle?", session=session)
 
     provider_state = session.state[provider.source_id]
 
     assert result.text == "It is sunny in Seattle."
     assert result.response_id == "resp_call_2"
     assert chat_client_base.call_count == 2
+    # The service owns the conversation, so the provider never loads (issue #5798).
     assert "get_call_count" not in provider_state
-    assert "save_call_count" not in provider_state
+    # Persistence is owned by the per-service-call middleware: it persists once per service call
+    # (issue #5798: the provider must never be silently bypassed when the service stores history).
+    # This run makes two service calls (function call + final answer), so it persists twice.
+    assert provider_state["save_call_count"] == 2
+    # load_messages=True while the service stores history surfaces a warning.
+    assert any("load_messages" in record.message for record in caplog.records)
     assert session.service_session_id == "resp_service_managed"
 
 
@@ -1996,6 +2007,19 @@ def test_merge_options_none_values_ignored():
     assert result["key2"] == "value2"
 
 
+def test_merge_options_drops_none_base_values():
+    """Test _merge_options strips None values so unset options are never forwarded."""
+    base = {"store": None, "temperature": 0.5}
+    override = {"top_p": 0.9}
+
+    result = _merge_options(base, override)
+
+    # An unset base value (e.g. store=None from default_options) must not survive the merge.
+    assert "store" not in result
+    assert result["temperature"] == 0.5
+    assert result["top_p"] == 0.9
+
+
 def test_merge_options_runtime_model_overrides_default_model() -> None:
     """Test _merge_options lets a runtime model override a default model."""
     result = _merge_options({"model": "default-model"}, {"model": "runtime-model"})
@@ -2658,3 +2682,449 @@ async def test_as_tool_raises_on_user_input_request(client: SupportsChatGetRespo
     assert len(exc_info.value.contents) == 1
     assert exc_info.value.contents[0].type == "oauth_consent_request"
     assert exc_info.value.contents[0].consent_link == "https://login.microsoftonline.com/consent"
+
+
+# region Per-service-call history persistence scenario matrix
+#
+# The driving field is ``require_per_service_call_history_persistence``. Every scenario runs a
+# single agent run that makes **two service calls** -- a function call followed by a final
+# completion -- so the *timing* of persistence is observable:
+#
+# * When the flag is ``True``, the per-service-call middleware persists the provider **after each
+#   service call**. So the function-call turn is already saved by the time the second (final)
+#   service call starts. This holds regardless of whether the chat client stores history
+#   server-side (the bug in issue #5798 was that a storing client silently bypassed persistence).
+# * When the flag is ``False``, the provider persists **once, at the end of the run** -- nothing is
+#   saved between the two service calls.
+#
+# ``SpyChatClient.saves_before_call`` records ``provider.save_calls`` at the start of every service
+# call, so ``[0, 1]`` means "the function-call turn was persisted before the final call" and
+# ``[0, 0]`` means "no persistence happened mid-run". The client's ``store`` / ``STORES_BY_DEFAULT``
+# only selects *how* the middleware behaves -- never *whether* the provider persists.
+
+_PSC_SERVICE_CONVERSATION_ID = "svc-conversation"
+
+_psc_stream_params = pytest.mark.parametrize("stream", [False, True], ids=["sync", "stream"])
+
+
+@tool(name="lookup_weather", approval_mode="never_require")
+def _psc_lookup_weather(location: str) -> str:
+    return f"Weather in {location}: sunny"
+
+
+def _psc_function_call_script() -> list[tuple[str, ...]]:
+    """A fresh function-call-then-final-completion script (the client mutates it)."""
+    return [
+        ("call", "call_1", "lookup_weather", '{"location": "Seattle"}'),
+        ("text", "It is sunny in Seattle."),
+    ]
+
+
+class _PscSpyHistoryProvider(HistoryProvider):
+    """In-memory history provider that records load/save calls for assertions."""
+
+    def __init__(self, source_id: str = "spy_history", **kwargs: Any) -> None:
+        super().__init__(source_id, **kwargs)
+        self._messages: list[Message] = []
+        self.get_calls: int = 0
+        self.save_calls: int = 0
+        self.saved_batches: list[list[Message]] = []
+
+    async def get_messages(
+        self, session_id: str | None, *, state: dict[str, Any] | None = None, **kwargs: Any
+    ) -> list[Message]:
+        self.get_calls += 1
+        return list(self._messages)
+
+    async def save_messages(
+        self,
+        session_id: str | None,
+        messages: Sequence[Message],
+        *,
+        state: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        self.save_calls += 1
+        self.saved_batches.append(list(messages))
+        self._messages.extend(messages)
+
+    @property
+    def stored_messages(self) -> list[Message]:
+        return list(self._messages)
+
+
+class _PscSpyChatClient(MockBaseChatClient):
+    """Chat client that scripts a function-call/final-completion sequence.
+
+    It records, at the start of each service call, how many provider saves have already happened
+    (``saves_before_call``), what messages it received, and what options it saw. When the effective
+    ``store`` is truthy it returns a stable ``conversation_id`` to mimic a server-managed
+    conversation, so the framework propagates ``session.service_session_id``.
+    """
+
+    def __init__(
+        self,
+        *,
+        provider: _PscSpyHistoryProvider,
+        stores_by_default: bool = False,
+        script: list[tuple[str, ...]] | None = None,
+        echo_conversation_id: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.STORES_BY_DEFAULT = stores_by_default  # type: ignore[attr-defined]
+        self._provider = provider
+        self._script = list(script) if script is not None else [("text", "ok")]
+        self._echo_conversation_id = echo_conversation_id
+        self.received_messages: list[list[Message]] = []
+        self.received_options: list[dict[str, Any]] = []
+        self.saves_before_call: list[int] = []
+
+    def _effective_store(self, options: dict[str, Any]) -> bool:
+        store = options.get("store")
+        if store is None:
+            return bool(self.STORES_BY_DEFAULT)
+        return bool(store)
+
+    def _next_contents(self) -> list[Content]:
+        turn = self._script.pop(0) if self._script else ("text", "ok")
+        if turn[0] == "call":
+            _, call_id, name, args = turn
+            return [Content.from_function_call(call_id=call_id, name=name, arguments=args)]
+        return [Content.from_text(turn[1])]
+
+    def _inner_get_response(  # type: ignore[override]
+        self,
+        *,
+        messages: MutableSequence[Message],
+        stream: bool,
+        options: dict[str, Any],
+        **kwargs: Any,
+    ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
+        self.received_messages.append(list(messages))
+        self.received_options.append(dict(options))
+        self.saves_before_call.append(self._provider.save_calls)
+        store_and_echo = self._effective_store(options) and self._echo_conversation_id
+        conv_id = _PSC_SERVICE_CONVERSATION_ID if store_and_echo else None
+        contents = self._next_contents()
+
+        if stream:
+
+            async def _stream() -> AsyncIterable[ChatResponseUpdate]:
+                self.call_count += 1
+                yield ChatResponseUpdate(
+                    contents=contents,
+                    role="assistant",
+                    finish_reason="stop",
+                    conversation_id=conv_id,
+                )
+
+            def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
+                response = ChatResponse.from_updates(updates, output_format_type=options.get("response_format"))
+                if conv_id:
+                    response.conversation_id = conv_id
+                return response
+
+            return ResponseStream(_stream(), finalizer=_finalize)
+
+        async def _get() -> ChatResponse:
+            self.call_count += 1
+            return ChatResponse(
+                messages=Message(role="assistant", contents=contents),
+                conversation_id=conv_id,
+            )
+
+        return _get()
+
+
+def _psc_build_agent(
+    client: _PscSpyChatClient,
+    provider: _PscSpyHistoryProvider,
+    *,
+    require_per_service_call_history_persistence: bool,
+    default_options: dict[str, Any] | None = None,
+) -> Agent:
+    kwargs: dict[str, Any] = {}
+    if default_options is not None:
+        kwargs["default_options"] = default_options
+    return Agent(
+        client=client,
+        tools=[_psc_lookup_weather],
+        context_providers=[provider],
+        require_per_service_call_history_persistence=require_per_service_call_history_persistence,
+        **kwargs,
+    )
+
+
+async def _psc_run(agent: Agent, text: str, session: AgentSession, *, stream: bool) -> str:
+    if stream:
+        chunks: list[str] = []
+        async for update in agent.run(text, session=session, stream=True):
+            chunks.append(update.text or "")
+        return "".join(chunks)
+    result = await agent.run(text, session=session)
+    return result.text
+
+
+# driver=True (the contract under test): persistence happens per service call
+
+
+@_psc_stream_params
+async def test_psc_flag_on_store_false_persists_after_each_service_call(stream: bool) -> None:
+    """Mode A (flag on, service does not store): function-call turn is persisted before the final call."""
+    provider = _PscSpyHistoryProvider()
+    client = _PscSpyChatClient(provider=provider, stores_by_default=False, script=_psc_function_call_script())
+    agent = _psc_build_agent(client, provider, require_per_service_call_history_persistence=True)
+    session = agent.create_session()
+
+    text = await _psc_run(agent, "What's the weather in Seattle?", session, stream=stream)
+
+    assert text == "It is sunny in Seattle."
+    # Two service calls: function call, then final completion.
+    assert client.call_count == 2
+    # The contract: the function-call turn was persisted *before* the second service call started.
+    assert client.saves_before_call == [0, 1]
+    assert provider.save_calls == 2
+    # Mode A loads local history (the middleware injects it before each service call).
+    assert provider.get_calls >= 1
+    # No service-side storage, so no conversation id is propagated.
+    assert session.service_session_id is None
+
+
+@_psc_stream_params
+async def test_psc_flag_on_stores_by_default_persists_after_each_service_call(
+    stream: bool, caplog: pytest.LogCaptureFixture
+) -> None:
+    """Mode B (flag on, service stores by default): still persists per service call, but skips load (issue #5798)."""
+    provider = _PscSpyHistoryProvider()  # load_messages=True by default
+    client = _PscSpyChatClient(provider=provider, stores_by_default=True, script=_psc_function_call_script())
+    agent = _psc_build_agent(client, provider, require_per_service_call_history_persistence=True)
+    session = agent.create_session()
+
+    with caplog.at_level(logging.WARNING, logger="agent_framework"):
+        text = await _psc_run(agent, "What's the weather in Seattle?", session, stream=stream)
+
+    assert text == "It is sunny in Seattle."
+    assert client.call_count == 2
+    # The invariant the bug violated: persistence still happens per service call when the service stores.
+    assert client.saves_before_call == [0, 1]
+    assert provider.save_calls == 2
+    # The service owns loading, so the provider is never asked to load.
+    assert provider.get_calls == 0
+    # A warning surfaces the bypassed load (load_messages=True).
+    assert any("load_messages" in record.message for record in caplog.records)
+    # The real service conversation id propagates to the session.
+    assert session.service_session_id == _PSC_SERVICE_CONVERSATION_ID
+
+
+@_psc_stream_params
+async def test_psc_flag_on_store_only_provider_no_load_no_warning(
+    stream: bool, caplog: pytest.LogCaptureFixture
+) -> None:
+    """Mode B with a store-only provider (load_messages=False): persists per call, no load, no warning."""
+    provider = _PscSpyHistoryProvider(load_messages=False)
+    client = _PscSpyChatClient(provider=provider, stores_by_default=True, script=_psc_function_call_script())
+    agent = _psc_build_agent(client, provider, require_per_service_call_history_persistence=True)
+    session = agent.create_session()
+
+    with caplog.at_level(logging.WARNING, logger="agent_framework"):
+        await _psc_run(agent, "What's the weather in Seattle?", session, stream=stream)
+
+    assert client.saves_before_call == [0, 1]
+    assert provider.save_calls == 2
+    assert provider.get_calls == 0
+    assert not any("load_messages" in record.message for record in caplog.records)
+
+
+@_psc_stream_params
+async def test_psc_flag_on_store_false_override_behaves_as_mode_a(stream: bool) -> None:
+    """Flag on + storing client but store=False override: falls back to Mode A (local, per call)."""
+    provider = _PscSpyHistoryProvider()
+    client = _PscSpyChatClient(provider=provider, stores_by_default=True, script=_psc_function_call_script())
+    agent = _psc_build_agent(
+        client, provider, require_per_service_call_history_persistence=True, default_options={"store": False}
+    )
+    session = agent.create_session()
+
+    await _psc_run(agent, "What's the weather in Seattle?", session, stream=stream)
+
+    assert client.saves_before_call == [0, 1]
+    assert provider.save_calls == 2
+    assert provider.get_calls >= 1
+    # store=False forces local handling, so no real service conversation id.
+    assert session.service_session_id is None
+
+
+@_psc_stream_params
+async def test_psc_flag_on_store_none_treated_as_absent(stream: bool, caplog: pytest.LogCaptureFixture) -> None:
+    """Flag on + storing client + explicit store=None: None is "unset", so the storing default applies (Mode B)."""
+    provider = _PscSpyHistoryProvider()
+    client = _PscSpyChatClient(provider=provider, stores_by_default=True, script=_psc_function_call_script())
+    agent = _psc_build_agent(
+        client, provider, require_per_service_call_history_persistence=True, default_options={"store": None}
+    )
+    session = agent.create_session()
+
+    with caplog.at_level(logging.WARNING, logger="agent_framework"):
+        await _psc_run(agent, "What's the weather in Seattle?", session, stream=stream)
+
+    assert client.saves_before_call == [0, 1]
+    assert provider.save_calls == 2
+    assert provider.get_calls == 0
+    assert session.service_session_id == _PSC_SERVICE_CONVERSATION_ID
+    assert any("load_messages" in record.message for record in caplog.records)
+    # store=None must not be forwarded to the client; the service decides its own default.
+    assert all("store" not in options for options in client.received_options)
+
+
+@_psc_stream_params
+async def test_psc_flag_on_respects_store_outputs_flag(stream: bool) -> None:
+    """Flag on: the provider's store_inputs/store_outputs flags still apply per service call."""
+    provider = _PscSpyHistoryProvider(store_inputs=True, store_outputs=False)
+    client = _PscSpyChatClient(provider=provider, stores_by_default=True, script=_psc_function_call_script())
+    agent = _psc_build_agent(client, provider, require_per_service_call_history_persistence=True)
+    session = agent.create_session()
+
+    await _psc_run(agent, "What's the weather in Seattle?", session, stream=stream)
+
+    assert provider.save_calls == 2
+    # Outputs disabled, so no assistant/tool-call messages were stored, only user/tool inputs.
+    assert provider.stored_messages
+    assert all(message.role != "assistant" for message in provider.stored_messages)
+
+
+# driver=False (control): persistence happens once, at the end of the run
+
+
+@_psc_stream_params
+async def test_psc_flag_off_store_false_persists_once_at_end(stream: bool) -> None:
+    """Flag off + non-storing client: nothing is persisted mid-run; one save at the end."""
+    provider = _PscSpyHistoryProvider()
+    client = _PscSpyChatClient(provider=provider, stores_by_default=False, script=_psc_function_call_script())
+    agent = _psc_build_agent(client, provider, require_per_service_call_history_persistence=False)
+    session = agent.create_session()
+
+    text = await _psc_run(agent, "What's the weather in Seattle?", session, stream=stream)
+
+    assert text == "It is sunny in Seattle."
+    assert client.call_count == 2
+    # The control contract: no save happened between the function call and the final completion.
+    assert client.saves_before_call == [0, 0]
+    assert provider.save_calls == 1
+
+
+@_psc_stream_params
+async def test_psc_flag_off_stores_by_default_persists_once_at_end(stream: bool) -> None:
+    """Flag off + storing client: once-per-run persistence, and the service conversation id propagates."""
+    provider = _PscSpyHistoryProvider()
+    client = _PscSpyChatClient(provider=provider, stores_by_default=True, script=_psc_function_call_script())
+    agent = _psc_build_agent(client, provider, require_per_service_call_history_persistence=False)
+    session = agent.create_session()
+
+    await _psc_run(agent, "What's the weather in Seattle?", session, stream=stream)
+
+    assert client.saves_before_call == [0, 0]
+    assert provider.save_calls == 1
+    assert session.service_session_id == _PSC_SERVICE_CONVERSATION_ID
+
+
+@_psc_stream_params
+async def test_psc_flag_on_storing_with_existing_conversation_id_does_not_raise(stream: bool) -> None:
+    """Allow side of the guard: flag on + storing client + an existing conversation_id resumes (no raise).
+
+    The non-storing path raises on an existing service-managed conversation id, but with a storing
+    client the run must proceed and the service conversation id must propagate to the session.
+    """
+    provider = _PscSpyHistoryProvider()
+    client = _PscSpyChatClient(provider=provider, stores_by_default=True, script=_psc_function_call_script())
+    agent = _psc_build_agent(client, provider, require_per_service_call_history_persistence=True)
+    session = agent.create_session()
+
+    if stream:
+        chunks: list[str] = []
+        async for update in agent.run(
+            "What's the weather in Seattle?",
+            session=session,
+            stream=True,
+            options={"conversation_id": "existing_conversation"},
+        ):
+            chunks.append(update.text or "")
+        text = "".join(chunks)
+    else:
+        result = await agent.run(
+            "What's the weather in Seattle?",
+            session=session,
+            options={"conversation_id": "existing_conversation"},
+        )
+        text = result.text
+
+    assert text == "It is sunny in Seattle."
+    # Persistence still happens per service call, and the real service id propagates to the session.
+    assert provider.save_calls == 2
+    assert provider.get_calls == 0
+    assert session.service_session_id == _PSC_SERVICE_CONVERSATION_ID
+
+
+@_psc_stream_params
+async def test_psc_flag_on_storing_two_runs_same_session(stream: bool) -> None:
+    """Storing mode across two runs on one session: persistence keeps happening, id is stable, no load.
+
+    The second run exercises the precedence branch where the session already carries a
+    service_session_id, which must continue to skip provider loading and keep persisting.
+    """
+    provider = _PscSpyHistoryProvider()
+    client = _PscSpyChatClient(provider=provider, stores_by_default=True, script=_psc_function_call_script())
+    agent = _psc_build_agent(client, provider, require_per_service_call_history_persistence=True)
+    session = agent.create_session()
+
+    await _psc_run(agent, "What's the weather in Seattle?", session, stream=stream)
+
+    assert provider.save_calls == 2
+    assert provider.get_calls == 0
+    first_run_service_id = session.service_session_id
+    assert first_run_service_id == _PSC_SERVICE_CONVERSATION_ID
+
+    # Reset the scripted client for a second run on the same session.
+    client._script = _psc_function_call_script()
+    client.call_count = 0
+    client.saves_before_call = []
+
+    await _psc_run(agent, "And in Portland?", session, stream=stream)
+
+    # Persistence keeps happening on the second run (two more saves), still per service call.
+    assert client.saves_before_call == [2, 3]
+    assert provider.save_calls == 4
+    # Loading stays skipped and the service conversation id stays stable across runs.
+    assert provider.get_calls == 0
+    assert session.service_session_id == first_run_service_id
+
+
+@_psc_stream_params
+async def test_psc_flag_on_storing_without_conversation_id_warns_every_call(
+    stream: bool, caplog: pytest.LogCaptureFixture
+) -> None:
+    """Storing mode but the client returns no conversation_id: warn on every service call.
+
+    Without an echoed conversation id the next run has nothing to resume from, so cross-turn
+    history can be lost silently. The warning fires per service call (no dedup) so the uncommon
+    failure mode cannot pass unnoticed.
+    """
+    provider = _PscSpyHistoryProvider()
+    client = _PscSpyChatClient(
+        provider=provider,
+        stores_by_default=True,
+        script=_psc_function_call_script(),
+        echo_conversation_id=False,
+    )
+    agent = _psc_build_agent(client, provider, require_per_service_call_history_persistence=True)
+    session = agent.create_session()
+
+    with caplog.at_level(logging.WARNING, logger="agent_framework"):
+        await _psc_run(agent, "What's the weather in Seattle?", session, stream=stream)
+
+    # Persistence still happens, but no service id is captured to resume from.
+    assert provider.save_calls == 2
+    assert session.service_session_id is None
+    # Two service calls -> the warning is emitted twice (one per call, not deduped).
+    missing_id_warnings = [r for r in caplog.records if "returned no conversation_id" in r.message]
+    assert len(missing_id_warnings) == 2

python/packages/core/tests/core/test_mcp.py

@@ -30,6 +30,7 @@ from agent_framework._mcp import (
     MCPTool,
     _build_prefixed_mcp_name,
     _get_input_model_from_mcp_prompt,
+    _normalize_additional_tool_argument_names,
     _normalize_mcp_name,
     _should_propagate_cancelled_error,
     logger,
@@ -6057,3 +6058,205 @@ async def test_max_wait_interrupts_long_poll_sleep(monkeypatch: pytest.MonkeyPat
 
 
 # endregion
+
+
+# region additional_tool_argument_names / allowlist filtering
+
+
+def test_normalize_additional_tool_argument_names_none() -> None:
+    global_extras, per_tool = _normalize_additional_tool_argument_names(None)
+    assert global_extras == set()
+    assert per_tool == {}
+
+
+def test_normalize_additional_tool_argument_names_sequence() -> None:
+    global_extras, per_tool = _normalize_additional_tool_argument_names(["a", "b", "a"])
+    assert global_extras == {"a", "b"}
+    assert per_tool == {}
+
+
+def test_normalize_additional_tool_argument_names_single_string() -> None:
+    # A bare string must be treated as a single name, not split into characters.
+    global_extras, per_tool = _normalize_additional_tool_argument_names("conversation_id")
+    assert global_extras == {"conversation_id"}
+    assert per_tool == {}
+
+
+def test_normalize_additional_tool_argument_names_mapping_with_global_key() -> None:
+    global_extras, per_tool = _normalize_additional_tool_argument_names({
+        "*": ["g1"],
+        "tool_a": ["a1", "a2"],
+        "tool_b": ["b1"],
+    })
+    assert global_extras == {"g1"}
+    assert per_tool == {"tool_a": {"a1", "a2"}, "tool_b": {"b1"}}
+
+
+def test_normalize_additional_tool_argument_names_mapping_with_string_values() -> None:
+    # A bare string mapping value is a single name, not an iterable of characters.
+    global_extras, per_tool = _normalize_additional_tool_argument_names({
+        "*": "conversation_id",
+        "tool_a": "custom",
+    })
+    assert global_extras == {"conversation_id"}
+    assert per_tool == {"tool_a": {"custom"}}
+
+
+def test_prepare_call_kwargs_strips_undeclared_arguments() -> None:
+    server = MCPTool(name="test_server")
+    server._tool_param_names_by_name = {"test_tool": {"param"}}
+
+    filtered, meta = server._prepare_call_kwargs(
+        "test_tool",
+        {"param": "value", "conversation_id": "c", "thread": object(), "unexpected": 1},
+    )
+
+    assert filtered == {"param": "value"}
+    assert meta is None
+
+
+def test_prepare_call_kwargs_global_extras_allowed() -> None:
+    server = MCPTool(name="test_server", additional_tool_argument_names=["conversation_id"])
+    server._tool_param_names_by_name = {"test_tool": {"param"}}
+
+    filtered, _ = server._prepare_call_kwargs(
+        "test_tool",
+        {"param": "value", "conversation_id": "c", "options": {}},
+    )
+
+    assert filtered == {"param": "value", "conversation_id": "c"}
+
+
+def test_prepare_call_kwargs_per_tool_and_global_extras() -> None:
+    server = MCPTool(
+        name="test_server",
+        additional_tool_argument_names={"*": ["conversation_id"], "test_tool": ["custom"]},
+    )
+    server._tool_param_names_by_name = {"test_tool": {"param"}, "other_tool": {"x"}}
+
+    filtered, _ = server._prepare_call_kwargs(
+        "test_tool",
+        {"param": "v", "conversation_id": "c", "custom": "y", "thread": object()},
+    )
+    assert filtered == {"param": "v", "conversation_id": "c", "custom": "y"}
+
+    # The per-tool extra does not leak to other tools; the global one still applies.
+    filtered_other, _ = server._prepare_call_kwargs(
+        "other_tool",
+        {"x": 1, "conversation_id": "c", "custom": "y"},
+    )
+    assert filtered_other == {"x": 1, "conversation_id": "c"}
+
+
+def test_prepare_call_kwargs_denylist_guards_server_declared_names() -> None:
+    # The denylist is a safety net for framework-named params a server *declares* in its
+    # schema: they are dropped so internal objects never leak. Names explicitly opted in
+    # via extras always win.
+    server = MCPTool(name="test_server", additional_tool_argument_names=["conversation_id"])
+    server._tool_param_names_by_name = {"test_tool": {"param", "thread"}}
+
+    filtered, _ = server._prepare_call_kwargs(
+        "test_tool",
+        {"param": "v", "thread": object(), "conversation_id": "c"},
+    )
+    # "thread" is declared by the schema but denylisted -> dropped; conversation_id opted in -> kept.
+    assert filtered == {"param": "v", "conversation_id": "c"}
+
+
+def test_prepare_call_kwargs_extras_override_denylist() -> None:
+    # Opting a denylisted framework name back in via extras takes precedence over the
+    # denylist safety net. "thread" is on the framework denylist, but an explicit extra wins.
+    server = MCPTool(name="test_server", additional_tool_argument_names=["thread"])
+    server._tool_param_names_by_name = {"test_tool": {"param"}}
+
+    sentinel = object()
+    filtered, _ = server._prepare_call_kwargs(
+        "test_tool",
+        {"param": "v", "thread": sentinel, "conversation_id": "c"},
+    )
+    # "thread" opted in via extras -> kept despite the denylist; conversation_id is denylisted,
+    # not declared, and not opted in -> dropped.
+    assert filtered == {"param": "v", "thread": sentinel}
+
+
+def test_prepare_call_kwargs_zero_arg_tool_passes_no_arguments() -> None:
+    server = MCPTool(name="test_server")
+    server._tool_param_names_by_name = {"test_tool": set()}
+
+    filtered, _ = server._prepare_call_kwargs(
+        "test_tool",
+        {"conversation_id": "c", "thread": object(), "stray": 1},
+    )
+    assert filtered == {}
+
+
+def test_prepare_call_kwargs_unknown_tool_passes_only_global_extras() -> None:
+    server = MCPTool(name="test_server", additional_tool_argument_names=["conversation_id"])
+    # No entry in _tool_param_names_by_name for this tool name.
+
+    filtered, _ = server._prepare_call_kwargs(
+        "unknown_tool",
+        {"conversation_id": "c", "other": 1},
+    )
+    assert filtered == {"conversation_id": "c"}
+
+
+def test_prepare_call_kwargs_extracts_meta() -> None:
+    server = MCPTool(name="test_server")
+    server._tool_param_names_by_name = {"test_tool": {"param"}}
+
+    filtered, meta = server._prepare_call_kwargs(
+        "test_tool",
+        {"param": "v", "_meta": {"trace": "abc"}},
+    )
+    assert filtered == {"param": "v"}
+    assert meta is not None
+    assert meta.get("trace") == "abc"
+
+
+async def test_call_tool_forwards_only_declared_arguments() -> None:
+    """End-to-end: framework runtime kwargs are stripped before reaching the server."""
+
+    class TestServer(MCPTool):
+        async def connect(self):
+            self.session = Mock(spec=ClientSession)
+            self.session.list_tools = AsyncMock(
+                return_value=types.ListToolsResult(
+                    tools=[
+                        types.Tool(
+                            name="test_tool",
+                            description="Test tool",
+                            inputSchema={
+                                "type": "object",
+                                "properties": {"param": {"type": "string"}},
+                                "required": ["param"],
+                            },
+                        )
+                    ]
+                )
+            )
+            self.session.call_tool = AsyncMock(
+                return_value=types.CallToolResult(content=[types.TextContent(type="text", text="ok")])
+            )
+
+        def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
+            return None
+
+    server = TestServer(name="test_server", additional_tool_argument_names=["conversation_id"])
+    async with server:
+        await server.load_tools()
+        session_mock = server.session
+        await server.call_tool(
+            "test_tool",
+            param="value",
+            conversation_id="c",
+            thread=object(),
+            response_format=object(),
+        )
+
+        session_mock.call_tool.assert_called_once()
+        _, call_kwargs = session_mock.call_tool.call_args
+        assert call_kwargs["arguments"] == {"param": "value", "conversation_id": "c"}
+
+
+# endregion