Python: Fix store=False not overriding client default (#4569)

* Fix store=False not overriding client default

* Address comments

* Fix unit tests

* Fix integration tests

* Fix tests
This commit is contained in:
Tao Chen
2026-03-10 11:44:59 -07:00
committed by GitHub
Unverified
parent c15f075412
commit 55fc882ca8
3 changed files with 375 additions and 203 deletions
@@ -33,7 +33,13 @@ from ._clients import BaseChatClient, SupportsChatGetResponse
from ._mcp import LOG_LEVEL_MAPPING, MCPTool
from ._middleware import AgentMiddlewareLayer, MiddlewareTypes
from ._serialization import SerializationMixin
from ._sessions import AgentSession, BaseContextProvider, BaseHistoryProvider, InMemoryHistoryProvider, SessionContext
from ._sessions import (
AgentSession,
BaseContextProvider,
BaseHistoryProvider,
InMemoryHistoryProvider,
SessionContext,
)
from ._tools import (
FunctionInvocationLayer,
FunctionTool,
@@ -532,7 +538,14 @@ class BaseAgent(SerializationMixin):
if stream_callback is None:
# Use non-streaming mode
return (await self.run(input_text, stream=False, session=parent_session, **forwarded_kwargs)).text
return (
await self.run(
input_text,
stream=False,
session=parent_session,
**forwarded_kwargs,
)
).text
# Use streaming mode - accumulate updates and create final response
response_updates: list[AgentResponseUpdate] = []
@@ -951,7 +964,9 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
**ctx["filtered_kwargs"],
)
def _propagate_conversation_id(update: AgentResponseUpdate) -> AgentResponseUpdate:
def _propagate_conversation_id(
update: AgentResponseUpdate,
) -> AgentResponseUpdate:
"""Eagerly propagate conversation_id to session as updates arrive.
This ensures session.service_session_id is set even when the user
@@ -975,8 +990,8 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
return self._finalize_response_updates(updates, response_format=rf)
return (
ResponseStream # type: ignore[reportUnknownMemberType]
.from_awaitable(_get_stream())
ResponseStream
.from_awaitable(_get_stream()) # type: ignore[reportUnknownMemberType]
.map(
transform=partial(
map_chat_to_agent_update,
@@ -1002,7 +1017,9 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
)
@staticmethod
def _extract_conversation_id_from_streaming_response(response: AgentResponse[Any]) -> str | None:
def _extract_conversation_id_from_streaming_response(
response: AgentResponse[Any],
) -> str | None:
"""Extract conversation_id from streaming raw updates, if present."""
raw = response.raw_representation
if raw is None:
@@ -1039,6 +1056,10 @@ 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)))
# Auto-inject InMemoryHistoryProvider when session is provided, no context providers
# registered, and no service-side storage indicators
if (
@@ -1046,8 +1067,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
and not self.context_providers
and not session.service_session_id
and not opts.get("conversation_id")
and not opts.get("store")
and not (getattr(self.client, "STORES_BY_DEFAULT", False) and opts.get("store") is not False)
and not store_
):
self.context_providers.append(InMemoryHistoryProvider())
+177 -41
View File
@@ -50,7 +50,9 @@ async def test_agent_run_with_content(agent: SupportsAgentRun) -> None:
async def test_agent_run_streaming(agent: SupportsAgentRun) -> None:
async def collect_updates(updates: AsyncIterable[AgentResponseUpdate]) -> list[AgentResponseUpdate]:
async def collect_updates(
updates: AsyncIterable[AgentResponseUpdate],
) -> list[AgentResponseUpdate]:
return [u async for u in updates]
updates = await collect_updates(agent.run("test", stream=True))
@@ -72,7 +74,9 @@ async def test_chat_client_agent_init(client: SupportsChatGetResponse) -> None:
assert agent.description == "Test"
async def test_chat_client_agent_init_with_name(client: SupportsChatGetResponse) -> None:
async def test_chat_client_agent_init_with_name(
client: SupportsChatGetResponse,
) -> None:
agent_id = str(uuid4())
agent = Agent(client=client, id=agent_id, name="Test Agent", description="Test")
@@ -108,7 +112,13 @@ async def test_chat_client_agent_streaming_response_format_from_default_options(
json_text = '{"greeting": "Hello"}'
client.streaming_responses.append( # type: ignore[attr-defined]
[ChatResponseUpdate(contents=[Content.from_text(json_text)], role="assistant", finish_reason="stop")]
[
ChatResponseUpdate(
contents=[Content.from_text(json_text)],
role="assistant",
finish_reason="stop",
)
]
)
agent = Agent(client=client, default_options={"response_format": Greeting})
@@ -134,7 +144,13 @@ async def test_chat_client_agent_streaming_response_format_from_run_options(
json_text = '{"greeting": "Hi"}'
client.streaming_responses.append( # type: ignore[attr-defined]
[ChatResponseUpdate(contents=[Content.from_text(json_text)], role="assistant", finish_reason="stop")]
[
ChatResponseUpdate(
contents=[Content.from_text(json_text)],
role="assistant",
finish_reason="stop",
)
]
)
agent = Agent(client=client)
@@ -149,14 +165,18 @@ async def test_chat_client_agent_streaming_response_format_from_run_options(
assert result.value.greeting == "Hi"
async def test_chat_client_agent_create_session(client: SupportsChatGetResponse) -> None:
async def test_chat_client_agent_create_session(
client: SupportsChatGetResponse,
) -> None:
agent = Agent(client=client)
session = agent.create_session()
assert isinstance(session, AgentSession)
async def test_chat_client_agent_prepare_session_and_messages(client: SupportsChatGetResponse) -> None:
async def test_chat_client_agent_prepare_session_and_messages(
client: SupportsChatGetResponse,
) -> None:
from agent_framework._sessions import InMemoryHistoryProvider
agent = Agent(client=client, context_providers=[InMemoryHistoryProvider()])
@@ -175,7 +195,9 @@ async def test_chat_client_agent_prepare_session_and_messages(client: SupportsCh
assert result_messages[1].text == "Test"
async def test_prepare_session_does_not_mutate_agent_chat_options(client: SupportsChatGetResponse) -> None:
async def test_prepare_session_does_not_mutate_agent_chat_options(
client: SupportsChatGetResponse,
) -> None:
tool = {"type": "code_interpreter"}
agent = Agent(client=client, tools=[tool])
@@ -195,7 +217,9 @@ async def test_prepare_session_does_not_mutate_agent_chat_options(client: Suppor
assert len(agent.default_options["tools"]) == 1
async def test_chat_client_agent_run_with_session(chat_client_base: SupportsChatGetResponse) -> None:
async def test_chat_client_agent_run_with_session(
chat_client_base: SupportsChatGetResponse,
) -> None:
mock_response = ChatResponse(
messages=[Message(role="assistant", contents=[Content.from_text("test response")])],
conversation_id="123",
@@ -391,7 +415,9 @@ async def test_chat_client_agent_streaming_session_history_saved_without_get_fin
assert chat_messages[1].text == "Hello Alice!"
async def test_chat_client_agent_update_session_messages(client: SupportsChatGetResponse) -> None:
async def test_chat_client_agent_update_session_messages(
client: SupportsChatGetResponse,
) -> None:
from agent_framework._sessions import InMemoryHistoryProvider
agent = Agent(client=client)
@@ -410,7 +436,9 @@ async def test_chat_client_agent_update_session_messages(client: SupportsChatGet
assert chat_messages[1].text == "test response"
async def test_chat_client_agent_update_session_conversation_id_missing(client: SupportsChatGetResponse) -> None:
async def test_chat_client_agent_update_session_conversation_id_missing(
client: SupportsChatGetResponse,
) -> None:
agent = Agent(client=client)
session = agent.get_session(service_session_id="123")
@@ -418,7 +446,9 @@ async def test_chat_client_agent_update_session_conversation_id_missing(client:
assert session.service_session_id == "123"
async def test_chat_client_agent_default_author_name(client: SupportsChatGetResponse) -> None:
async def test_chat_client_agent_default_author_name(
client: SupportsChatGetResponse,
) -> None:
# Name is not specified here, so default name should be used
agent = Agent(client=client)
@@ -427,7 +457,9 @@ async def test_chat_client_agent_default_author_name(client: SupportsChatGetResp
assert result.messages[0].author_name == "UnnamedAgent"
async def test_chat_client_agent_author_name_as_agent_name(client: SupportsChatGetResponse) -> None:
async def test_chat_client_agent_author_name_as_agent_name(
client: SupportsChatGetResponse,
) -> None:
# Name is specified here, so it should be used as author name
agent = Agent(client=client, name="TestAgent")
@@ -436,11 +468,17 @@ async def test_chat_client_agent_author_name_as_agent_name(client: SupportsChatG
assert result.messages[0].author_name == "TestAgent"
async def test_chat_client_agent_author_name_is_used_from_response(chat_client_base: SupportsChatGetResponse) -> None:
async def test_chat_client_agent_author_name_is_used_from_response(
chat_client_base: SupportsChatGetResponse,
) -> None:
chat_client_base.run_responses = [
ChatResponse(
messages=[
Message(role="assistant", contents=[Content.from_text("test response")], author_name="TestAuthor")
Message(
role="assistant",
contents=[Content.from_text("test response")],
author_name="TestAuthor",
)
]
)
]
@@ -476,7 +514,9 @@ class MockContextProvider(BaseContextProvider):
self.new_messages.extend(context.response.messages)
async def test_chat_agent_context_providers_model_before_run(client: SupportsChatGetResponse) -> None:
async def test_chat_agent_context_providers_model_before_run(
client: SupportsChatGetResponse,
) -> None:
"""Test that context providers' before_run is called during agent run."""
mock_provider = MockContextProvider(messages=[Message(role="system", text="Test context instructions")])
agent = Agent(client=client, context_providers=[mock_provider])
@@ -486,7 +526,9 @@ async def test_chat_agent_context_providers_model_before_run(client: SupportsCha
assert mock_provider.before_run_called
async def test_chat_agent_context_providers_after_run(chat_client_base: SupportsChatGetResponse) -> None:
async def test_chat_agent_context_providers_after_run(
chat_client_base: SupportsChatGetResponse,
) -> None:
"""Test that context providers' after_run is called during agent run."""
mock_provider = MockContextProvider()
chat_client_base.run_responses = [
@@ -505,7 +547,9 @@ async def test_chat_agent_context_providers_after_run(chat_client_base: Supports
assert mock_provider.last_service_session_id == "test-thread-id"
async def test_chat_agent_context_providers_messages_adding(client: SupportsChatGetResponse) -> None:
async def test_chat_agent_context_providers_messages_adding(
client: SupportsChatGetResponse,
) -> None:
"""Test that context providers' after_run is called during agent run."""
mock_provider = MockContextProvider()
agent = Agent(client=client, context_providers=[mock_provider])
@@ -517,10 +561,16 @@ async def test_chat_agent_context_providers_messages_adding(client: SupportsChat
assert len(mock_provider.new_messages) >= 2
async def test_chat_agent_context_instructions_in_messages(client: SupportsChatGetResponse) -> None:
async def test_chat_agent_context_instructions_in_messages(
client: SupportsChatGetResponse,
) -> None:
"""Test that AI context instructions are included in messages."""
mock_provider = MockContextProvider(messages=[Message(role="system", text="Context-specific instructions")])
agent = Agent(client=client, instructions="Agent instructions", context_providers=[mock_provider])
agent = Agent(
client=client,
instructions="Agent instructions",
context_providers=[mock_provider],
)
# We need to test the _prepare_session_and_messages method directly
session_context, _ = await agent._prepare_session_and_messages( # type: ignore[reportPrivateUsage]
@@ -537,10 +587,16 @@ async def test_chat_agent_context_instructions_in_messages(client: SupportsChatG
# instructions system message is added by a client
async def test_chat_agent_no_context_instructions(client: SupportsChatGetResponse) -> None:
async def test_chat_agent_no_context_instructions(
client: SupportsChatGetResponse,
) -> None:
"""Test behavior when AI context has no instructions."""
mock_provider = MockContextProvider()
agent = Agent(client=client, instructions="Agent instructions", context_providers=[mock_provider])
agent = Agent(
client=client,
instructions="Agent instructions",
context_providers=[mock_provider],
)
session_context, _ = await agent._prepare_session_and_messages( # type: ignore[reportPrivateUsage]
session=None, input_messages=[Message(role="user", text="Hello")]
@@ -553,7 +609,9 @@ async def test_chat_agent_no_context_instructions(client: SupportsChatGetRespons
assert messages[0].text == "Hello"
async def test_chat_agent_run_stream_context_providers(client: SupportsChatGetResponse) -> None:
async def test_chat_agent_run_stream_context_providers(
client: SupportsChatGetResponse,
) -> None:
"""Test that context providers work with run method."""
mock_provider = MockContextProvider(messages=[Message(role="system", text="Stream context instructions")])
agent = Agent(client=client, context_providers=[mock_provider])
@@ -571,7 +629,9 @@ async def test_chat_agent_run_stream_context_providers(client: SupportsChatGetRe
assert mock_provider.after_run_called
async def test_chat_agent_context_providers_with_service_session_id(chat_client_base: SupportsChatGetResponse) -> None:
async def test_chat_agent_context_providers_with_service_session_id(
chat_client_base: SupportsChatGetResponse,
) -> None:
"""Test context providers with service-managed session."""
mock_provider = MockContextProvider()
chat_client_base.run_responses = [
@@ -604,7 +664,9 @@ async def test_chat_agent_as_tool_basic(client: SupportsChatGetResponse) -> None
assert hasattr(tool, "input_model")
async def test_chat_agent_as_tool_custom_parameters(client: SupportsChatGetResponse) -> None:
async def test_chat_agent_as_tool_custom_parameters(
client: SupportsChatGetResponse,
) -> None:
"""Test as_tool with custom parameters."""
agent = Agent(client=client, name="TestAgent", description="Original description")
@@ -652,7 +714,9 @@ async def test_chat_agent_as_tool_no_name(client: SupportsChatGetResponse) -> No
agent.as_tool()
async def test_chat_agent_as_tool_function_execution(client: SupportsChatGetResponse) -> None:
async def test_chat_agent_as_tool_function_execution(
client: SupportsChatGetResponse,
) -> None:
"""Test that the generated FunctionTool can be executed."""
agent = Agent(client=client, name="TestAgent", description="Test agent")
@@ -666,7 +730,9 @@ async def test_chat_agent_as_tool_function_execution(client: SupportsChatGetResp
assert result == "test response" # From mock chat client
async def test_chat_agent_as_tool_with_stream_callback(client: SupportsChatGetResponse) -> None:
async def test_chat_agent_as_tool_with_stream_callback(
client: SupportsChatGetResponse,
) -> None:
"""Test as_tool with stream callback functionality."""
agent = Agent(client=client, name="StreamingAgent")
@@ -689,7 +755,9 @@ async def test_chat_agent_as_tool_with_stream_callback(client: SupportsChatGetRe
assert result == expected_text
async def test_chat_agent_as_tool_with_custom_arg_name(client: SupportsChatGetResponse) -> None:
async def test_chat_agent_as_tool_with_custom_arg_name(
client: SupportsChatGetResponse,
) -> None:
"""Test as_tool with custom argument name."""
agent = Agent(client=client, name="CustomArgAgent")
@@ -700,7 +768,9 @@ async def test_chat_agent_as_tool_with_custom_arg_name(client: SupportsChatGetRe
assert result == "test response"
async def test_chat_agent_as_tool_with_async_stream_callback(client: SupportsChatGetResponse) -> None:
async def test_chat_agent_as_tool_with_async_stream_callback(
client: SupportsChatGetResponse,
) -> None:
"""Test as_tool with async stream callback functionality."""
agent = Agent(client=client, name="AsyncStreamingAgent")
@@ -723,7 +793,9 @@ async def test_chat_agent_as_tool_with_async_stream_callback(client: SupportsCha
assert result == expected_text
async def test_chat_agent_as_tool_name_sanitization(client: SupportsChatGetResponse) -> None:
async def test_chat_agent_as_tool_name_sanitization(
client: SupportsChatGetResponse,
) -> None:
"""Test as_tool name sanitization."""
test_cases = [
("Invoice & Billing Agent", "Invoice_Billing_Agent"),
@@ -741,7 +813,9 @@ async def test_chat_agent_as_tool_name_sanitization(client: SupportsChatGetRespo
assert tool.name == expected_tool_name, f"Expected {expected_tool_name}, got {tool.name} for input {agent_name}"
async def test_chat_agent_as_tool_propagate_session_true(client: SupportsChatGetResponse) -> None:
async def test_chat_agent_as_tool_propagate_session_true(
client: SupportsChatGetResponse,
) -> None:
"""Test that propagate_session=True forwards the parent's session to the sub-agent."""
agent = Agent(client=client, name="SubAgent", description="Sub agent")
tool = agent.as_tool(propagate_session=True)
@@ -767,7 +841,9 @@ async def test_chat_agent_as_tool_propagate_session_true(client: SupportsChatGet
assert captured_session.state["shared_key"] == "shared_value"
async def test_chat_agent_as_tool_propagate_session_false_by_default(client: SupportsChatGetResponse) -> None:
async def test_chat_agent_as_tool_propagate_session_false_by_default(
client: SupportsChatGetResponse,
) -> None:
"""Test that propagate_session defaults to False and does not forward the session."""
agent = Agent(client=client, name="SubAgent", description="Sub agent")
tool = agent.as_tool() # default: propagate_session=False
@@ -789,7 +865,9 @@ async def test_chat_agent_as_tool_propagate_session_false_by_default(client: Sup
assert captured_session is None
async def test_chat_agent_as_tool_propagate_session_shares_state(client: SupportsChatGetResponse) -> None:
async def test_chat_agent_as_tool_propagate_session_shares_state(
client: SupportsChatGetResponse,
) -> None:
"""Test that shared session allows the sub-agent to read and write parent's state."""
agent = Agent(client=client, name="SubAgent", description="Sub agent")
tool = agent.as_tool(propagate_session=True)
@@ -858,13 +936,20 @@ async def test_chat_agent_with_local_mcp_tools(client: SupportsChatGetResponse)
# Test agent with MCP tools in constructor
with contextlib.suppress(Exception):
agent = Agent(client=client, name="TestAgent", description="Test agent", tools=[mock_mcp_tool])
agent = Agent(
client=client,
name="TestAgent",
description="Test agent",
tools=[mock_mcp_tool],
)
# Test async context manager with MCP tools
async with agent:
pass
async def test_mcp_tools_not_duplicated_when_passed_as_runtime_tools(chat_client_base: Any) -> None:
async def test_mcp_tools_not_duplicated_when_passed_as_runtime_tools(
chat_client_base: Any,
) -> None:
"""Test that MCP tool functions from self.mcp_tools are not duplicated when already present in runtime tools."""
captured_options: list[dict[str, Any]] = []
@@ -925,7 +1010,11 @@ async def test_agent_tool_receives_session_in_kwargs(chat_client_base: Any) -> N
messages=Message(
role="assistant",
contents=[
Content.from_function_call(call_id="1", name="echo_session_info", arguments='{"text": "hello"}')
Content.from_function_call(
call_id="1",
name="echo_session_info",
arguments='{"text": "hello"}',
)
],
)
),
@@ -935,7 +1024,11 @@ async def test_agent_tool_receives_session_in_kwargs(chat_client_base: Any) -> N
agent = Agent(client=chat_client_base, tools=[echo_session_info])
session = agent.create_session()
result = await agent.run("hello", session=session, options={"additional_function_arguments": {"session": session}})
result = await agent.run(
"hello",
session=session,
options={"additional_function_arguments": {"session": session}},
)
assert result.text == "done"
assert captured.get("has_session") is True
@@ -1377,7 +1470,9 @@ def test_chat_agent_calls_update_agent_name_on_client():
@pytest.mark.asyncio
async def test_chat_agent_context_provider_adds_tools_when_agent_has_none(chat_client_base: SupportsChatGetResponse):
async def test_chat_agent_context_provider_adds_tools_when_agent_has_none(
chat_client_base: SupportsChatGetResponse,
):
"""Test that context provider tools are used when agent has no default tools."""
@tool
@@ -1439,7 +1534,9 @@ async def test_chat_agent_context_provider_adds_instructions_when_agent_has_none
# region STORES_BY_DEFAULT tests
async def test_stores_by_default_skips_inmemory_injection(client: SupportsChatGetResponse) -> None:
async def test_stores_by_default_skips_inmemory_injection(
client: SupportsChatGetResponse,
) -> None:
"""Client with STORES_BY_DEFAULT=True should not auto-inject InMemoryHistoryProvider."""
from agent_framework._sessions import InMemoryHistoryProvider
@@ -1455,7 +1552,9 @@ async def test_stores_by_default_skips_inmemory_injection(client: SupportsChatGe
assert not any(isinstance(p, InMemoryHistoryProvider) for p in agent.context_providers)
async def test_stores_by_default_false_injects_inmemory(client: SupportsChatGetResponse) -> None:
async def test_stores_by_default_false_injects_inmemory(
client: SupportsChatGetResponse,
) -> None:
"""Client with STORES_BY_DEFAULT=False (default) should auto-inject InMemoryHistoryProvider."""
from agent_framework._sessions import InMemoryHistoryProvider
@@ -1468,7 +1567,9 @@ async def test_stores_by_default_false_injects_inmemory(client: SupportsChatGetR
assert any(isinstance(p, InMemoryHistoryProvider) for p in agent.context_providers)
async def test_stores_by_default_with_store_false_injects_inmemory(client: SupportsChatGetResponse) -> None:
async def test_stores_by_default_with_store_false_injects_inmemory(
client: SupportsChatGetResponse,
) -> None:
"""Client with STORES_BY_DEFAULT=True but store=False should still inject InMemoryHistoryProvider."""
from agent_framework._sessions import InMemoryHistoryProvider
@@ -1483,7 +1584,42 @@ async def test_stores_by_default_with_store_false_injects_inmemory(client: Suppo
assert any(isinstance(p, InMemoryHistoryProvider) for p in agent.context_providers)
# endregion
async def test_store_true_skips_inmemory_injection(
client: SupportsChatGetResponse,
) -> None:
"""Explicitly setting store=True should not auto-inject InMemoryHistoryProvider."""
from agent_framework._sessions import InMemoryHistoryProvider
agent = Agent(client=client)
session = agent.create_session()
await agent.run("Hello", session=session, options={"store": True})
# User explicitly enabled server storage, so InMemoryHistoryProvider should not be injected
assert not any(isinstance(p, InMemoryHistoryProvider) for p in agent.context_providers)
async def test_stores_by_default_with_store_false_in_default_options_injects_inmemory(
client: SupportsChatGetResponse,
) -> None:
"""Client with STORES_BY_DEFAULT=True but store=False in default_options should inject InMemoryHistoryProvider.
This covers the regression where store=False is set via Agent(..., default_options={"store": False})
with no per-run override while the client has STORES_BY_DEFAULT=True.
"""
from agent_framework._sessions import InMemoryHistoryProvider
client.STORES_BY_DEFAULT = True # type: ignore[attr-defined]
# Set store=False at agent initialization via default_options, not at run-time
agent = Agent(client=client, default_options={"store": False})
session = agent.create_session()
# Run without any per-run options override
await agent.run("Hello", session=session)
# User explicitly disabled server storage in default_options, so InMemoryHistoryProvider should be injected
assert any(isinstance(p, InMemoryHistoryProvider) for p in agent.context_providers)
# endregion