Python: PR1 — New session and context provider types (side-by-side) #3763
Conversation
markwallace-microsoft
added
the
documentation
eavanvalkenburg
changed the title
There was a problem hiding this comment.
Pull request overview
Adds the new “hooks-pattern” context/session abstractions (per ADR 0016) side-by-side with the current thread/context APIs, along with initial provider ports (Redis, Mem0, Azure AI Search) and corresponding test coverage.
Changes:
- Introduces new core session/context types (
SessionContext,BaseContextProvider,BaseHistoryProvider,AgentSession,InMemoryHistoryProvider) with deep (de)serialization support. - Adds side-by-side “new pattern” external providers: Redis context/history, Mem0 context, Azure AI Search context (all
_-prefixed). - Adds extensive unit tests for the new core types and each provider, plus an accepted ADR documenting the design.
Reviewed changes
Copilot reviewed 12 out of 13 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| python/packages/core/agent_framework/_sessions.py | New core session/context provider types and state serialization helpers. |
| python/packages/core/tests/core/test_sessions.py | Unit tests for SessionContext, provider bases, AgentSession, and InMemoryHistoryProvider. |
| python/packages/redis/agent_framework_redis/_context_provider.py | New-pattern Redis context provider built on BaseContextProvider. |
| python/packages/redis/agent_framework_redis/_history_provider.py | New-pattern Redis history provider built on BaseHistoryProvider. |
| python/packages/redis/agent_framework_redis/init.py | Exports new _RedisContextProvider / _RedisHistoryProvider symbols. |
| python/packages/redis/tests/test_new_providers.py | Unit tests for new Redis providers and BaseHistoryProvider hook defaults. |
| python/packages/mem0/agent_framework_mem0/_context_provider.py | New-pattern Mem0 context provider built on BaseContextProvider. |
| python/packages/mem0/agent_framework_mem0/init.py | Exports new _Mem0ContextProvider symbol. |
| python/packages/mem0/tests/test_mem0_new_context_provider.py | Unit tests for new Mem0 context provider. |
| python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py | New-pattern Azure AI Search context provider built on BaseContextProvider. |
| python/packages/azure-ai-search/agent_framework_azure_ai_search/init.py | Exports new _AzureAISearchContextProvider symbol. |
| python/packages/azure-ai-search/tests/test_aisearch_new_context_provider.py | Unit tests for new Azure AI Search provider behaviors (semantic mode + cleanup). |
| docs/decisions/0016-python-context-middleware.md | ADR documenting the new unified context/session/provider design and rollout plan. |
Python Test Coverage Report •
Python Unit Test Overview
|
|||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
New types in _sessions.py (no changes to existing code): - SessionContext: per-invocation state with extend_messages/get_messages/ extend_instructions/extend_tools and read-only response property - _ContextProviderBase: base class with before_run/after_run hooks - _HistoryProviderBase: storage base with load/store flags, abstract get_messages/save_messages, default before_run/after_run - AgentSession: lightweight session with state dict, to_dict/from_dict - InMemoryHistoryProvider: built-in provider storing in session.state 35 unit tests covering all classes and configuration flags.
…rialization - Make before_run/after_run parameters keyword-only - InMemoryHistoryProvider stores ChatMessage objects directly (no per-cycle serialization) - Deep serialization via to_dict/from_dict only at session boundary - State type registry for automatic deserialization of registered types - Updated tests for new serialization approach
- _RedisContextProvider(BaseContextProvider) - Redis search/vector context - _RedisHistoryProvider(BaseHistoryProvider) - Redis-backed message storage - _Mem0ContextProvider(BaseContextProvider) - Mem0 semantic memory - _AzureAISearchContextProvider(BaseContextProvider) - Azure AI Search (semantic + agentic) All use temporary _ prefix names for side-by-side coexistence with existing providers. Will be renamed in PR2 when old ContextProvider/ChatMessageStore are removed.
- 32 tests for _RedisContextProvider and _RedisHistoryProvider - 29 tests for _Mem0ContextProvider - 17 tests for _AzureAISearchContextProvider
- Move module docstring before imports in _sessions.py (review comment) - Import TYPE_CHECKING unconditionally in Redis _context_provider.py (NameError on Python <3.12) - Fix Mem0 test_init_auto_creates_client_when_none to patch at class level
eavanvalkenburg
force-pushed
the
adr-0016-python-context-middleware
branch
from
6a1ddf3 to
7cf583c
Compare
Set attribution marker in additional_properties for each message added via extend_messages(), matching the tool attribution pattern. Uses setdefault to preserve any existing attribution.
…ssages - SessionContext.extend_messages now accepts source as str or object with source_id attribute; when an object is passed, its class name is recorded as source_type in the attribution dict - Messages are shallow-copied before attribution is added so callers' original objects are never mutated - Filter framework-internal keys (attribution) from A2A wire metadata to prevent leaking internal state over the wire
| return value | ||
|
|
||
|
|
||
| def _deserialize_value(value: Any) -> Any: |
There was a problem hiding this comment.
When a dict has a "type" key but the type isn't in _STATE_TYPE_REGISTRY, the function silently falls through and returns a plain dict. Users who store custom typed objects in session state will get dicts back after deserialization with no warning, leading to confusing AttributeError failures downstream, if I am understanding this correctly.
Should we log a warning when a "type" key exists but no deserializer is registered, and wrap from_dict() calls in try/except with context?
| ) | ||
| text_results = await self.redis_index.query(query) | ||
| return cast(list[dict[str, Any]], text_results) | ||
| except Exception as exc: # pragma: no cover |
There was a problem hiding this comment.
All exceptions (connection errors, auth failures, timeouts, programming bugs) are flattened into ServiceInvalidRequestError("Redis text search failed"). Doesn't this make it impossible for callers to distinguish between transient / permanent failures? How should retry logic be handled?
|
|
||
| async with self._redis_client.pipeline(transaction=True) as pipe: | ||
| for serialized in serialized_messages: | ||
| await pipe.rpush(key, serialized) # type: ignore[misc] |
There was a problem hiding this comment.
The rpush is inside a pipeline transaction, but llen + ltrim happen outside it as separate operations. Under concurrent writes, the ltrim can calculate against a stale count and trim messages that another caller just pushed. Could the ltrim move inside the pipeline? ltrim(key, -max_messages, -1) is idempotent and safe to call unconditionally:
async with self._redis_client.pipeline(transaction=True) as pipe:
for serialized in serialized_messages:
await pipe.rpush(key, serialized)
if self.max_messages is not None:
await pipe.ltrim(key, -self.max_messages, -1)
await pipe.execute()| Returns: | ||
| List of stored ChatMessage objects in chronological order. | ||
| """ | ||
| key = self._redis_key(session_id) |
There was a problem hiding this comment.
One corrupted JSON entry in Redis will raise JSONDecodeError and kill the entire get_messages call, losing access to the whole session history. This is especially likely after schema migrations where old entries have a different format. Should this do per-message error handling and skip corrupted entries with a warning?
|
|
||
| def _redis_key(self, session_id: str | None) -> str: | ||
| """Get the Redis key for a given session's messages.""" | ||
| return f"{self.key_prefix}:{session_id or 'default'}" |
There was a problem hiding this comment.
When session_id is None, this falls back to "default", meaning all sessions without explicit IDs share a single Redis key. Could this silently blend data across unrelated sessions? Should it at least log a warning?
| """Async context manager entry.""" | ||
| return self | ||
|
|
||
| async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None: |
There was a problem hiding this comment.
aexit has an empty body, but AsyncSearchIndex holds Redis connections internally. The history provider properly implements aclose() -- should this provider close its resources too? Users of async with get a false sense of cleanup otherwise.
| if not input_text.strip(): | ||
| return | ||
|
|
||
| memories = await self._redis_search(text=input_text, session_id=context.session_id) |
There was a problem hiding this comment.
_redis_search is called with session_id=context.session_id, which adds thread_id and conversation_id filters (line 357-362). This means search results are scoped to the current session. Is the expectation that the same session_id is reused across turns of a conversation? If so this works fine. But if callers create a fresh AgentSession() per invocation (which auto-generates a new UUID), the search would never match previously stored data. Is that the expected usage pattern, or should the session-scoping be optional here?
5 participants
Summary
Implements the new context provider types from ADR 0016 PR #3609 (
docs/decisions/0016-python-context-middleware.md), creating them side-by-side with the existingContextProviderandChatMessageStoreclasses. No existing code is modified — this is purely additive.What is included
Core types (
packages/core/agent_framework/_sessions.py)SessionContext— Per-invocation state passed through the provider pipelineBaseContextProvider— Base class withbefore_run()/after_run()hooks (renamed toContextProviderin PR2)BaseHistoryProvider— Specialized base for conversation history with load/store flags (renamed toHistoryProviderin PR2)AgentSession— Lightweight session container withto_dict()/from_dict()serializationInMemoryHistoryProvider— Built-in stateless in-memory history providerExternal package providers
_RedisContextProvider(BaseContextProvider)— Redis search/vector context (portsRedisProvider)_RedisHistoryProvider(BaseHistoryProvider)— Redis-backed message storage (portsRedisChatMessageStore)_Mem0ContextProvider(BaseContextProvider)— Mem0 semantic memory (portsMem0Provider)_AzureAISearchContextProvider(BaseContextProvider)— Azure AI Search semantic + agentic modes (portsAzureAISearchContextProvider)All external package providers use temporary
_prefix names to avoid collision with existing classes.Tests
Related issues
This PR lays the groundwork for the following issues from #3575, but does not yet resolve them — full resolution requires PR2 (agent integration + old type removal):
AgentSessionis created here; replacingAgentThreadhappens in PR2create_session()/get_session()API designed in ADR; agent integration in PR2AgentSession.to_dict()/from_dict()implemented here; agent wiring in PR2BaseHistoryProviderwithload_messages/store_*flags implemented here; agent integration in PR2BaseHistoryProviderreplaces it; old removal in PR2Key design decisions
before_run/after_run) instead of middleware chain — simpler, nonext()to forgetInMemoryHistoryProvider— all data lives insession.state, not on the provider instanceAgentSession.to_dict()usesSerializationProtocolto serializeChatMessageobjects; deserialization via type registrySee ADR 0016 PR #3609 for the full design rationale.