Python: [BREAKING] PR2 — Wire context provider pipeline, remove old types, update all consumers (#3850)

* PR2: Wire context provider pipeline and update all internal consumers

- Replace AgentThread with AgentSession across all packages
- Replace ContextProvider with BaseContextProvider across all packages
- Replace context_provider param with context_providers (Sequence)
- Replace thread= with session= in run() signatures
- Replace get_new_thread() with create_session()
- Add get_session(service_session_id) to agent interface
- DurableAgentThread -> DurableAgentSession
- Remove _notify_thread_of_new_messages from WorkflowAgent
- Wire before_run/after_run context provider pipeline in RawAgent
- Auto-inject InMemoryHistoryProvider when no providers configured

* fix: update all tests for context provider pipeline, fix lazy-loaders, remove old test files

* refactor: update all sample files for context provider pipeline (AgentThread→AgentSession, ContextProvider→BaseContextProvider)

* fix: update remaining ag-ui references (client docstring, getting_started sample)

* fix: make get_session service_session_id keyword-only to avoid confusion with session_id

* refactor: rename _RunContext.thread_messages to session_messages

* refactor: remove _threads.py, _memory.py, and old provider files; migrate devui to use plain message lists

* rename: remove _new_ prefix from test files

* refactor: rewrite SlidingWindowChatMessageStore as SlidingWindowHistoryProvider(InMemoryHistoryProvider)

* fix: read full history from session state directly instead of reaching into provider internals

* fix: update stale .pyi stubs, sample imports, and README references for new provider types

* fix: remove stale message_store, _notify_thread_of_new_messages, and session_id.key references in samples

* refactor: merge context_providers and sessions sample folders into sessions, remove aggregate_context_provider

* refactor: UserInfoMemory stores state in session.state instead of instance attributes

* feat: add Pydantic BaseModel support to session state serialization

Pydantic models stored in session.state are now automatically serialized
via model_dump() and restored via model_validate() during to_dict()/from_dict()
round-trips. Models are auto-registered on first serialization; use
register_state_type() for cold-start deserialization.

Also export register_state_type as a public API.

* fix mem0

* Update sample README links and descriptions for session terminology

- Replace 'thread' with 'session' in sample descriptions across all READMEs
- Update file links for renamed samples (mem0_sessions, redis_sessions, etc.)
- Fix Threads section → Sessions section in main samples/README.md
- Update tools, middleware, workflows, durabletask, azure_functions READMEs
- Update architecture diagrams in concepts/tools/README.md
- Update migration guides (autogen, semantic-kernel)

* Fix broken Redis README link to renamed sample

* Fix Mem0 OSS client search: pass scoping params as direct kwargs

AsyncMemory (OSS) expects user_id/agent_id/run_id as direct kwargs,
while AsyncMemoryClient (Platform) expects them in a filters dict.
Adds tests for both client types.

Port of fix from #3844 to new Mem0ContextProvider.

* Fix rebase issues: restore missing _conversation_state.py and checkpoint decode logic

- Add back _conversation_state.py (encode/decode_chat_messages) lost in rebase
- Fix on_checkpoint_restore to decode cache/conversation with decode_chat_messages
- Fix on_checkpoint_restore to use decode_checkpoint_value for pending requests
- Add tests/workflow/__init__.py for relative import support
- Fix test_agent_executor checkpoint selection (checkpoints[1] not superstep)

* Add STORES_BY_DEFAULT ClassVar to skip redundant InMemoryHistoryProvider injection

Chat clients that store history server-side by default (OpenAI Responses API,
Azure AI Agent) now declare STORES_BY_DEFAULT = True. The agent checks this
during auto-injection and skips InMemoryHistoryProvider unless the user
explicitly sets store=False.

* Fix broken markdown links in azure_ai and redis READMEs

* Fix getting-started samples to use session API instead of removed thread/ContextProvider API

* updates to workflow as agent

* fix group chat import

* Rename Thread→Session throughout, fix service_session_id propagation, remove stale AGUIThread

- Fix: Propagate conversation_id from ChatResponse back to session.service_session_id
  in both streaming and non-streaming paths in _agents.py
- Rename AgentThreadException → AgentSessionException
- Remove stale AGUIThread from ag_ui lazy-loader
- Rename use_service_thread → use_service_session in ag-ui package
- Rename test functions from *_thread_* to *_session_*
- Rename sample files from *_thread* to *_session*
- Update docstrings and comments: thread → session
- Update _mcp.py kwargs filter: add 'session' alongside 'thread'
- Fix ContinuationToken docstring example: thread=thread → session=session
- Fix _clients.py docstring: 'Agent threads' → 'Agent sessions'

* Fix broken markdown links after thread→session file renames

* fix azure ai test
This commit is contained in:
Eduard van Valkenburg
2026-02-12 22:00:32 +01:00
committed by GitHub
Unverified
parent 0c67dbbce5
commit 1e350ea22f
312 changed files with 6669 additions and 11423 deletions
+1 -1
View File
@@ -18,7 +18,7 @@ Durable execution support for long-running agent workflows using Azure Durable F
### State Management
- **`DurableAgentState`** - State container for durable agents
- **`DurableAgentThread`** - Thread management for durable agents
- **`DurableAgentSession`** - Session management for durable agents
- **`DurableAIAgentOrchestrationContext`** - Orchestration context
### Callbacks
@@ -45,7 +45,7 @@ from ._durable_agent_state import (
)
from ._entities import AgentEntity, AgentEntityStateProviderMixin
from ._executors import DurableAgentExecutor
from ._models import AgentSessionId, DurableAgentThread, RunRequest
from ._models import AgentSessionId, DurableAgentSession, RunRequest
from ._orchestration_context import DurableAIAgentOrchestrationContext
from ._response_utils import ensure_response_format, load_agent_response
from ._shim import DurableAIAgent
@@ -79,6 +79,7 @@ __all__ = [
"DurableAIAgentOrchestrationContext",
"DurableAIAgentWorker",
"DurableAgentExecutor",
"DurableAgentSession",
"DurableAgentState",
"DurableAgentStateContent",
"DurableAgentStateData",
@@ -99,7 +100,6 @@ __all__ = [
"DurableAgentStateUriContent",
"DurableAgentStateUsage",
"DurableAgentStateUsageContent",
"DurableAgentThread",
"DurableStateFields",
"RunRequest",
"__version__",
@@ -16,7 +16,7 @@ from abc import ABC, abstractmethod
from datetime import datetime, timezone
from typing import Any, Generic, TypeVar
from agent_framework import AgentResponse, AgentThread, Content, Message, get_logger
from agent_framework import AgentResponse, AgentSession, Content, Message, get_logger
from durabletask.client import TaskHubGrpcClient
from durabletask.entities import EntityInstanceId
from durabletask.task import CompletableTask, CompositeTask, OrchestrationContext, Task
@@ -24,7 +24,7 @@ from pydantic import BaseModel
from ._constants import DEFAULT_MAX_POLL_RETRIES, DEFAULT_POLL_INTERVAL_SECONDS
from ._durable_agent_state import DurableAgentState
from ._models import AgentSessionId, DurableAgentThread, RunRequest
from ._models import AgentSessionId, DurableAgentSession, RunRequest
from ._response_utils import ensure_response_format, load_agent_response
logger = get_logger("agent_framework.durabletask.executors")
@@ -114,7 +114,7 @@ class DurableAgentExecutor(ABC, Generic[TaskT]):
self,
agent_name: str,
run_request: RunRequest,
thread: AgentThread | None = None,
session: AgentSession | None = None,
) -> TaskT:
"""Execute the durable agent.
@@ -123,20 +123,20 @@ class DurableAgentExecutor(ABC, Generic[TaskT]):
"""
raise NotImplementedError
def get_new_thread(self, agent_name: str, **kwargs: Any) -> DurableAgentThread:
"""Create a new DurableAgentThread with random session ID."""
def get_new_session(self, agent_name: str, **kwargs: Any) -> DurableAgentSession:
"""Create a new DurableAgentSession with random session ID."""
session_id = self._create_session_id(agent_name)
return DurableAgentThread.from_session_id(session_id, **kwargs)
return DurableAgentSession.from_session_id(session_id, **kwargs)
def _create_session_id(
self,
agent_name: str,
thread: AgentThread | None = None,
session: AgentSession | None = None,
) -> AgentSessionId:
"""Create the AgentSessionId for the execution."""
if isinstance(thread, DurableAgentThread) and thread.session_id is not None:
return thread.session_id
# Create new session ID - either no thread provided or it's a regular AgentThread
if isinstance(session, DurableAgentSession) and session.durable_session_id is not None:
return session.durable_session_id
# Create new session ID - either no session provided or it's a regular AgentSession
key = self.generate_unique_id()
return AgentSessionId(name=agent_name, key=key)
@@ -217,7 +217,7 @@ class ClientAgentExecutor(DurableAgentExecutor[AgentResponse]):
self,
agent_name: str,
run_request: RunRequest,
thread: AgentThread | None = None,
session: AgentSession | None = None,
) -> AgentResponse:
"""Execute the agent via the durabletask client.
@@ -231,14 +231,14 @@ class ClientAgentExecutor(DurableAgentExecutor[AgentResponse]):
Args:
agent_name: Name of the agent to execute
run_request: The run request containing message and optional response format
thread: Optional conversation thread (creates new if not provided)
session: Optional conversation session (creates new if not provided)
Returns:
AgentResponse: The agent's response after execution completes, or an immediate
acknowledgement if wait_for_response is False
"""
# Signal the entity with the request
entity_id = self._signal_agent_entity(agent_name, run_request, thread)
entity_id = self._signal_agent_entity(agent_name, run_request, session)
# If fire-and-forget mode, return immediately without polling
if not run_request.wait_for_response:
@@ -258,20 +258,20 @@ class ClientAgentExecutor(DurableAgentExecutor[AgentResponse]):
self,
agent_name: str,
run_request: RunRequest,
thread: AgentThread | None,
session: AgentSession | None,
) -> EntityInstanceId:
"""Signal the agent entity with a run request.
Args:
agent_name: Name of the agent to execute
run_request: The run request containing message and optional response format
thread: Optional conversation thread
session: Optional conversation session
Returns:
entity_id
"""
# Get or create session ID
session_id = self._create_session_id(agent_name, thread)
session_id = self._create_session_id(agent_name, session)
# Create the entity ID
entity_id = EntityInstanceId(
@@ -460,7 +460,7 @@ class OrchestrationAgentExecutor(DurableAgentExecutor[DurableAgentTask]):
self,
agent_name: str,
run_request: RunRequest,
thread: AgentThread | None = None,
session: AgentSession | None = None,
) -> DurableAgentTask:
"""Execute the agent via orchestration context.
@@ -470,13 +470,13 @@ class OrchestrationAgentExecutor(DurableAgentExecutor[DurableAgentTask]):
Args:
agent_name: Name of the agent to execute
run_request: The run request containing message and optional response format
thread: Optional conversation thread (creates new if not provided)
session: Optional conversation session (creates new if not provided)
Returns:
DurableAgentTask: A task wrapping the entity call that yields AgentResponse
"""
# Resolve session
session_id = self._create_session_id(agent_name, thread)
session_id = self._create_session_id(agent_name, session)
# Create the entity ID
entity_id = EntityInstanceId(
@@ -10,13 +10,12 @@ from __future__ import annotations
import inspect
import json
import uuid
from collections.abc import MutableMapping
from dataclasses import dataclass, field
from datetime import datetime, timezone
from importlib import import_module
from typing import TYPE_CHECKING, Any, cast
from agent_framework import AgentThread
from agent_framework import AgentSession
from ._constants import REQUEST_RESPONSE_FORMAT_TEXT
@@ -274,65 +273,57 @@ class AgentSessionId:
raise ValueError(f"Invalid agent session ID format: {session_id_string}")
class DurableAgentThread(AgentThread):
"""Durable agent thread that tracks the owning :class:`AgentSessionId`."""
class DurableAgentSession(AgentSession):
"""Durable agent session that tracks the owning :class:`AgentSessionId`."""
_SERIALIZED_SESSION_ID_KEY = "durable_session_id"
def __init__(
self,
*,
session_id: AgentSessionId | None = None,
durable_session_id: AgentSessionId | None = None,
session_id: str | None = None,
service_session_id: str | None = None,
**kwargs: Any,
) -> None:
super().__init__(**kwargs)
self._session_id: AgentSessionId | None = session_id
super().__init__(session_id=session_id, service_session_id=service_session_id, **kwargs)
self._session_id_value: AgentSessionId | None = durable_session_id
@property
def session_id(self) -> AgentSessionId | None:
return self._session_id
def durable_session_id(self) -> AgentSessionId | None:
return self._session_id_value
@session_id.setter
def session_id(self, value: AgentSessionId | None) -> None:
self._session_id = value
@durable_session_id.setter
def durable_session_id(self, value: AgentSessionId | None) -> None:
self._session_id_value = value
@classmethod
def from_session_id(
cls,
session_id: AgentSessionId,
**kwargs: Any,
) -> DurableAgentThread:
return cls(session_id=session_id, **kwargs)
) -> DurableAgentSession:
return cls(durable_session_id=session_id, **kwargs)
async def serialize(self, **kwargs: Any) -> dict[str, Any]:
state = await super().serialize(**kwargs)
if self._session_id is not None:
state[self._SERIALIZED_SESSION_ID_KEY] = str(self._session_id)
def to_dict(self) -> dict[str, Any]:
state = super().to_dict()
if self._session_id_value is not None:
state[self._SERIALIZED_SESSION_ID_KEY] = str(self._session_id_value)
return state
@classmethod
async def deserialize(
cls,
serialized_thread_state: MutableMapping[str, Any],
*,
message_store: Any = None,
**kwargs: Any,
) -> DurableAgentThread:
state_payload = dict(serialized_thread_state)
def from_dict(cls, data: dict[str, Any]) -> DurableAgentSession:
state_payload = dict(data)
session_id_value = state_payload.pop(cls._SERIALIZED_SESSION_ID_KEY, None)
thread = await super().deserialize(
state_payload,
message_store=message_store,
**kwargs,
session = super().from_dict(state_payload)
# We need to create a DurableAgentSession from the base AgentSession
durable_session = cls(
session_id=session.session_id,
service_session_id=session.service_session_id,
)
if not isinstance(thread, DurableAgentThread):
raise TypeError("Deserialized thread is not a DurableAgentThread instance")
if session_id_value is None:
return thread
if not isinstance(session_id_value, str):
raise ValueError("durable_session_id must be a string when present in serialized state")
thread.session_id = AgentSessionId.parse(session_id_value)
return thread
durable_session.state.update(session.state)
if session_id_value is not None:
if not isinstance(session_id_value, str):
raise ValueError("durable_session_id must be a string when present in serialized state")
durable_session._session_id_value = AgentSessionId.parse(session_id_value)
return durable_session
@@ -12,10 +12,10 @@ from __future__ import annotations
from abc import ABC, abstractmethod
from typing import Any, Generic, Literal, TypeVar
from agent_framework import AgentThread, Message, SupportsAgentRun
from agent_framework import AgentSession, Message, SupportsAgentRun
from ._executors import DurableAgentExecutor
from ._models import DurableAgentThread
from ._models import DurableAgentSession
# TypeVar for the task type returned by executors
# Covariant because TaskT only appears in return positions (output)
@@ -89,7 +89,7 @@ class DurableAIAgent(SupportsAgentRun, Generic[TaskT]):
messages: str | Message | list[str] | list[Message] | None = None,
*,
stream: Literal[False] = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
options: dict[str, Any] | None = None,
) -> TaskT:
"""Execute the agent via the injected provider.
@@ -98,7 +98,7 @@ class DurableAIAgent(SupportsAgentRun, Generic[TaskT]):
messages: The message(s) to send to the agent
stream: Whether to use streaming for the response (must be False)
DurableAgents do not support streaming mode.
thread: Optional agent thread for conversation context
session: Optional agent session for conversation context
options: Optional options dictionary. Supported keys include
``response_format``, ``enable_tool_calls``, and ``wait_for_response``.
Additional keys are forwarded to the agent execution.
@@ -129,12 +129,19 @@ class DurableAIAgent(SupportsAgentRun, Generic[TaskT]):
return self._executor.run_durable_agent(
agent_name=self.name,
run_request=run_request,
thread=thread,
session=session,
)
def get_new_thread(self, **kwargs: Any) -> DurableAgentThread:
"""Create a new agent thread via the provider."""
return self._executor.get_new_thread(self.name, **kwargs)
def create_session(self, **kwargs: Any) -> DurableAgentSession:
"""Create a new agent session via the provider."""
return self._executor.get_new_session(self.name, **kwargs)
def get_session(self, **kwargs: Any) -> AgentSession:
"""Retrieve an existing session via the provider.
For durable agents, sessions do not use `service_session_id` so this is not used.
"""
return self._executor.get_new_session(self.name, **kwargs)
def _normalize_messages(self, messages: str | Message | list[str] | list[Message] | None) -> str:
"""Convert supported message inputs to a single string.
@@ -39,9 +39,9 @@ class TestSingleAgent:
def test_single_interaction(self):
"""Test a single interaction with the agent."""
agent = self.agent_client.get_agent("Joker")
thread = agent.get_new_thread()
session = agent.create_session()
response = agent.run("Tell me a short joke about programming.", thread=thread)
response = agent.run("Tell me a short joke about programming.", session=session)
assert response is not None
assert response.text is not None
@@ -50,33 +50,33 @@ class TestSingleAgent:
def test_conversation_continuity(self):
"""Test that conversation context is maintained across turns."""
agent = self.agent_client.get_agent("Joker")
thread = agent.get_new_thread()
session = agent.create_session()
# First turn: Ask for a joke about a specific topic
response1 = agent.run("Tell me a joke about cats.", thread=thread)
response1 = agent.run("Tell me a joke about cats.", session=session)
assert response1 is not None
assert len(response1.text) > 0
# Second turn: Ask a follow-up that requires context
response2 = agent.run("Can you make it funnier?", thread=thread)
response2 = agent.run("Can you make it funnier?", session=session)
assert response2 is not None
assert len(response2.text) > 0
# The agent should understand "it" refers to the previous joke
def test_multiple_threads(self):
"""Test that different threads maintain separate contexts."""
def test_multiple_sessions(self):
"""Test that different sessions maintain separate contexts."""
agent = self.agent_client.get_agent("Joker")
# Create two separate threads
thread1 = agent.get_new_thread()
thread2 = agent.get_new_thread()
# Create two separate sessions
session1 = agent.create_session()
session2 = agent.create_session()
assert thread1.session_id != thread2.session_id
assert session1.durable_session_id != session2.durable_session_id
# Send different messages to each thread
response1 = agent.run("Tell me a joke about dogs.", thread=thread1)
response2 = agent.run("Tell me a joke about birds.", thread=thread2)
# Send different messages to each session
response1 = agent.run("Tell me a joke about dogs.", session=session1)
response2 = agent.run("Tell me a joke about birds.", session=session2)
assert response1 is not None
assert response2 is not None
@@ -47,9 +47,9 @@ class TestMultiAgent:
def test_weather_agent_with_tool(self):
"""Test weather agent with weather tool execution."""
agent = self.agent_client.get_agent(WEATHER_AGENT_NAME)
thread = agent.get_new_thread()
session = agent.create_session()
response = agent.run("What's the weather in Seattle?", thread=thread)
response = agent.run("What's the weather in Seattle?", session=session)
assert response is not None
assert response.text is not None
@@ -66,9 +66,9 @@ class TestMultiAgent:
def test_math_agent_with_tool(self):
"""Test math agent with calculation tool execution."""
agent = self.agent_client.get_agent(MATH_AGENT_NAME)
thread = agent.get_new_thread()
session = agent.create_session()
response = agent.run("Calculate a 20% tip on a $50 bill.", thread=thread)
response = agent.run("Calculate a 20% tip on a $50 bill.", session=session)
assert response is not None
assert response.text is not None
@@ -85,11 +85,11 @@ class TestMultiAgent:
def test_multiple_calls_to_same_agent(self):
"""Test multiple sequential calls to the same agent."""
agent = self.agent_client.get_agent(WEATHER_AGENT_NAME)
thread = agent.get_new_thread()
session = agent.create_session()
# Multiple weather queries
response1 = agent.run("What's the weather in Chicago?", thread=thread)
response2 = agent.run("And what about Los Angeles?", thread=thread)
response1 = agent.run("What's the weather in Chicago?", session=session)
response2 = agent.run("And what about Los Angeles?", session=session)
assert response1 is not None
assert response2 is not None
@@ -70,7 +70,7 @@ class TestSampleReliableStreaming:
async def _stream_from_redis(
self,
thread_id: str,
session_key: str,
cursor: str | None = None,
timeout: float = 30.0,
) -> tuple[str, bool, str]:
@@ -78,7 +78,7 @@ class TestSampleReliableStreaming:
Stream responses from Redis using the sample's RedisStreamResponseHandler.
Args:
thread_id: The conversation/thread ID to stream from
session_key: The conversation/thread ID to stream from
cursor: Optional cursor to resume from
timeout: Maximum time to wait for stream completion
@@ -92,7 +92,7 @@ class TestSampleReliableStreaming:
async with await self._get_stream_handler() as stream_handler: # type: ignore[reportUnknownMemberType]
try:
async for chunk in stream_handler.read_stream(thread_id, cursor): # type: ignore[reportUnknownMemberType]
async for chunk in stream_handler.read_stream(session_key, cursor): # type: ignore[reportUnknownMemberType]
if time.time() - start_time > timeout:
break
@@ -124,15 +124,15 @@ class TestSampleReliableStreaming:
assert travel_planner is not None
assert travel_planner.name == "TravelPlanner"
# Create a new thread
thread = travel_planner.get_new_thread()
assert thread.session_id is not None
assert thread.session_id.key is not None
thread_id = str(thread.session_id.key)
# Create a new session
session = travel_planner.create_session()
assert session.durable_session_id is not None
assert session.durable_session_id.key is not None
session_key = str(session.durable_session_id.key)
# Start agent run with wait_for_response=False for non-blocking execution
travel_planner.run(
"Plan a 1-day trip to Seattle in 1 sentence", thread=thread, options={"wait_for_response": False}
"Plan a 1-day trip to Seattle in 1 sentence", session=session, options={"wait_for_response": False}
)
# Poll Redis stream with retries to handle race conditions
@@ -146,7 +146,7 @@ class TestSampleReliableStreaming:
while retry_count < max_retries and not is_complete:
text, is_complete, last_cursor = asyncio.run(
self._stream_from_redis(thread_id, cursor=cursor, timeout=10.0)
self._stream_from_redis(session_key, cursor=cursor, timeout=10.0)
)
accumulated_text += text
cursor = last_cursor # Resume from last position on next read
@@ -166,7 +166,7 @@ class TestSampleReliableStreaming:
# Verify we got content
assert len(accumulated_text) > 0, (
f"Expected text content but got empty string for thread_id: {thread_id} after {retry_count} retries"
f"Expected text content but got empty string for session_key: {session_key} after {retry_count} retries"
)
assert "seattle" in accumulated_text.lower(), f"Expected 'seattle' in response but got: {accumulated_text}"
assert is_complete, "Expected stream to be complete"
@@ -175,13 +175,13 @@ class TestSampleReliableStreaming:
"""Test streaming with cursor-based resumption."""
# Get the TravelPlanner agent
travel_planner = self.agent_client.get_agent("TravelPlanner")
thread = travel_planner.get_new_thread()
assert thread.session_id is not None
assert thread.session_id.key is not None
thread_id = str(thread.session_id.key)
session = travel_planner.create_session()
assert session.durable_session_id is not None
assert session.durable_session_id.key is not None
session_key = str(session.durable_session_id.key)
# Start agent run
travel_planner.run("What's the weather like?", thread=thread, options={"wait_for_response": False})
travel_planner.run("What's the weather like?", session=session, options={"wait_for_response": False})
# Wait for agent to start writing
time.sleep(3)
@@ -194,7 +194,7 @@ class TestSampleReliableStreaming:
chunk_count = 0
# Read just first 2 chunks
async for chunk in stream_handler.read_stream(thread_id): # type: ignore[reportUnknownMemberType]
async for chunk in stream_handler.read_stream(session_key): # type: ignore[reportUnknownMemberType]
last_entry_id = chunk.entry_id # type: ignore[reportUnknownMemberType]
if chunk.text: # type: ignore[reportUnknownMemberType]
accumulated_text += chunk.text # type: ignore[reportUnknownMemberType]
@@ -207,7 +207,7 @@ class TestSampleReliableStreaming:
partial_text, cursor = asyncio.run(get_partial_stream())
# Resume from cursor
remaining_text, _, _ = asyncio.run(self._stream_from_redis(thread_id, cursor=cursor))
remaining_text, _, _ = asyncio.run(self._stream_from_redis(session_key, cursor=cursor))
# Verify we got some initial content
assert len(partial_text) > 0
@@ -1,11 +1,11 @@
# Copyright (c) Microsoft. All rights reserved.
"""Unit tests for AgentSessionId and DurableAgentThread."""
"""Unit tests for AgentSessionId and DurableAgentSession."""
import pytest
from agent_framework import AgentThread
from agent_framework import AgentSession
from agent_framework_durabletask._models import AgentSessionId, DurableAgentThread
from agent_framework_durabletask._models import AgentSessionId, DurableAgentSession
class TestAgentSessionId:
@@ -121,154 +121,162 @@ class TestAgentSessionId:
assert "Invalid agent session ID format" in str(exc_info.value)
class TestDurableAgentThread:
"""Test suite for DurableAgentThread."""
class TestDurableAgentSession:
"""Test suite for DurableAgentSession."""
def test_init_with_session_id(self) -> None:
"""Test DurableAgentThread initialization with session ID."""
def test_init_with_durable_session_id(self) -> None:
"""Test DurableAgentSession initialization with durable session ID."""
session_id = AgentSessionId(name="TestAgent", key="test-key")
thread = DurableAgentThread(session_id=session_id)
session = DurableAgentSession(durable_session_id=session_id)
assert thread.session_id is not None
assert thread.session_id == session_id
assert session.durable_session_id is not None
assert session.durable_session_id == session_id
def test_init_without_session_id(self) -> None:
"""Test DurableAgentThread initialization without session ID."""
thread = DurableAgentThread()
def test_init_without_durable_session_id(self) -> None:
"""Test DurableAgentSession initialization without durable session ID."""
session = DurableAgentSession()
assert thread.session_id is None
assert session.durable_session_id is None
def test_session_id_setter(self) -> None:
"""Test setting a session ID to an existing thread."""
thread = DurableAgentThread()
assert thread.session_id is None
def test_durable_session_id_setter(self) -> None:
"""Test setting a durable session ID to an existing session."""
session = DurableAgentSession()
assert session.durable_session_id is None
session_id = AgentSessionId(name="TestAgent", key="test-key")
thread.session_id = session_id
session.durable_session_id = session_id
assert thread.session_id is not None
assert thread.session_id == session_id
assert thread.session_id.name == "TestAgent"
assert session.durable_session_id is not None
assert session.durable_session_id == session_id
assert session.durable_session_id.name == "TestAgent"
def test_from_session_id(self) -> None:
"""Test creating DurableAgentThread from session ID."""
"""Test creating DurableAgentSession from session ID."""
session_id = AgentSessionId(name="TestAgent", key="test-key")
thread = DurableAgentThread.from_session_id(session_id)
session = DurableAgentSession.from_session_id(session_id)
assert isinstance(thread, DurableAgentThread)
assert thread.session_id is not None
assert thread.session_id == session_id
assert thread.session_id.name == "TestAgent"
assert thread.session_id.key == "test-key"
assert isinstance(session, DurableAgentSession)
assert session.durable_session_id is not None
assert session.durable_session_id == session_id
assert session.durable_session_id.name == "TestAgent"
assert session.durable_session_id.key == "test-key"
def test_from_session_id_with_service_thread_id(self) -> None:
"""Test creating DurableAgentThread with service thread ID."""
def test_from_session_id_with_service_session_id(self) -> None:
"""Test creating DurableAgentSession with service session ID."""
session_id = AgentSessionId(name="TestAgent", key="test-key")
thread = DurableAgentThread.from_session_id(session_id, service_thread_id="service-123")
session = DurableAgentSession.from_session_id(session_id, service_session_id="service-123")
assert thread.session_id is not None
assert thread.session_id == session_id
assert thread.service_thread_id == "service-123"
assert session.durable_session_id is not None
assert session.durable_session_id == session_id
assert session.service_session_id == "service-123"
async def test_serialize_with_session_id(self) -> None:
"""Test serialization includes session ID."""
def test_to_dict_with_durable_session_id(self) -> None:
"""Test serialization includes durable session ID."""
session_id = AgentSessionId(name="TestAgent", key="test-key")
thread = DurableAgentThread(session_id=session_id)
session = DurableAgentSession(durable_session_id=session_id)
serialized = await thread.serialize()
serialized = session.to_dict()
assert isinstance(serialized, dict)
assert "durable_session_id" in serialized
assert serialized["durable_session_id"] == "@TestAgent@test-key"
async def test_serialize_without_session_id(self) -> None:
"""Test serialization without session ID."""
thread = DurableAgentThread()
def test_to_dict_without_durable_session_id(self) -> None:
"""Test serialization without durable session ID."""
session = DurableAgentSession()
serialized = await thread.serialize()
serialized = session.to_dict()
assert isinstance(serialized, dict)
assert "durable_session_id" not in serialized
async def test_deserialize_with_session_id(self) -> None:
"""Test deserialization restores session ID."""
def test_from_dict_with_durable_session_id(self) -> None:
"""Test deserialization restores durable session ID."""
serialized = {
"service_thread_id": "thread-123",
"type": "session",
"session_id": "session-123",
"service_session_id": "service-123",
"state": {},
"durable_session_id": "@TestAgent@test-key",
}
thread = await DurableAgentThread.deserialize(serialized)
session = DurableAgentSession.from_dict(serialized)
assert isinstance(thread, DurableAgentThread)
assert thread.session_id is not None
assert thread.session_id.name == "TestAgent"
assert thread.session_id.key == "test-key"
assert thread.service_thread_id == "thread-123"
assert isinstance(session, DurableAgentSession)
assert session.durable_session_id is not None
assert session.durable_session_id.name == "TestAgent"
assert session.durable_session_id.key == "test-key"
assert session.service_session_id == "service-123"
async def test_deserialize_without_session_id(self) -> None:
"""Test deserialization without session ID."""
def test_from_dict_without_durable_session_id(self) -> None:
"""Test deserialization without durable session ID."""
serialized = {
"service_thread_id": "thread-456",
"type": "session",
"session_id": "session-456",
"service_session_id": "service-456",
"state": {},
}
thread = await DurableAgentThread.deserialize(serialized)
session = DurableAgentSession.from_dict(serialized)
assert isinstance(thread, DurableAgentThread)
assert thread.session_id is None
assert thread.service_thread_id == "thread-456"
assert isinstance(session, DurableAgentSession)
assert session.durable_session_id is None
assert session.session_id == "session-456"
async def test_round_trip_serialization(self) -> None:
"""Test round-trip serialization preserves session ID."""
def test_round_trip_serialization(self) -> None:
"""Test round-trip serialization preserves durable session ID."""
session_id = AgentSessionId(name="TestAgent", key="test-key-789")
original = DurableAgentThread(session_id=session_id)
original = DurableAgentSession(durable_session_id=session_id)
serialized = await original.serialize()
restored = await DurableAgentThread.deserialize(serialized)
serialized = original.to_dict()
restored = DurableAgentSession.from_dict(serialized)
assert isinstance(restored, DurableAgentThread)
assert restored.session_id is not None
assert restored.session_id.name == session_id.name
assert restored.session_id.key == session_id.key
assert isinstance(restored, DurableAgentSession)
assert restored.durable_session_id is not None
assert restored.durable_session_id.name == session_id.name
assert restored.durable_session_id.key == session_id.key
async def test_deserialize_invalid_session_id_type(self) -> None:
"""Test deserialization with invalid session ID type raises error."""
def test_from_dict_invalid_durable_session_id_type(self) -> None:
"""Test deserialization with invalid durable session ID type raises error."""
serialized = {
"service_thread_id": "thread-123",
"type": "session",
"session_id": "session-123",
"state": {},
"durable_session_id": 12345, # Invalid type
}
with pytest.raises(ValueError, match="durable_session_id must be a string"):
await DurableAgentThread.deserialize(serialized)
DurableAgentSession.from_dict(serialized)
class TestAgentThreadCompatibility:
"""Test suite for compatibility between AgentThread and DurableAgentThread."""
class TestAgentSessionCompatibility:
"""Test suite for compatibility between AgentSession and DurableAgentSession."""
async def test_agent_thread_serialize(self) -> None:
"""Test that base AgentThread can be serialized."""
thread = AgentThread()
def test_agent_session_to_dict(self) -> None:
"""Test that base AgentSession can be serialized."""
session = AgentSession()
serialized = await thread.serialize()
serialized = session.to_dict()
assert isinstance(serialized, dict)
assert "service_thread_id" in serialized
assert "session_id" in serialized
async def test_agent_thread_deserialize(self) -> None:
"""Test that base AgentThread can be deserialized."""
thread = AgentThread()
serialized = await thread.serialize()
def test_agent_session_from_dict(self) -> None:
"""Test that base AgentSession can be deserialized."""
session = AgentSession()
serialized = session.to_dict()
restored = await AgentThread.deserialize(serialized)
restored = AgentSession.from_dict(serialized)
assert isinstance(restored, AgentThread)
assert restored.service_thread_id == thread.service_thread_id
assert isinstance(restored, AgentSession)
assert restored.session_id == session.session_id
async def test_durable_thread_is_agent_thread(self) -> None:
"""Test that DurableAgentThread is an AgentThread."""
thread = DurableAgentThread()
def test_durable_session_is_agent_session(self) -> None:
"""Test that DurableAgentSession is an AgentSession."""
session = DurableAgentSession()
assert isinstance(thread, AgentThread)
assert isinstance(thread, DurableAgentThread)
assert isinstance(session, AgentSession)
assert isinstance(session, DurableAgentSession)
class TestModelIntegration:
@@ -281,19 +289,19 @@ class TestModelIntegration:
assert session_id_str.startswith("@AgentEntity@")
async def test_thread_with_session_preserves_on_serialization(self) -> None:
"""Test that thread with session ID preserves it through serialization."""
def test_session_with_durable_id_preserves_on_serialization(self) -> None:
"""Test that session with durable session ID preserves it through serialization."""
session_id = AgentSessionId(name="TestAgent", key="preserved-key")
thread = DurableAgentThread.from_session_id(session_id)
session = DurableAgentSession.from_session_id(session_id)
# Serialize and deserialize
serialized = await thread.serialize()
restored = await DurableAgentThread.deserialize(serialized)
serialized = session.to_dict()
restored = DurableAgentSession.from_dict(serialized)
# Session ID should be preserved
assert restored.session_id is not None
assert restored.session_id.name == "TestAgent"
assert restored.session_id.key == "preserved-key"
# Durable session ID should be preserved
assert restored.durable_session_id is not None
assert restored.durable_session_id.name == "TestAgent"
assert restored.durable_session_id.key == "preserved-key"
if __name__ == "__main__":
@@ -11,7 +11,7 @@ from unittest.mock import Mock
import pytest
from agent_framework import SupportsAgentRun
from agent_framework_durabletask import DurableAgentThread, DurableAIAgentClient
from agent_framework_durabletask import DurableAgentSession, DurableAIAgentClient
from agent_framework_durabletask._constants import DEFAULT_MAX_POLL_RETRIES, DEFAULT_POLL_INTERVAL_SECONDS
from agent_framework_durabletask._shim import DurableAIAgent
@@ -80,22 +80,22 @@ class TestDurableAIAgentClientIntegration:
assert hasattr(agent, "run")
assert callable(agent.run)
def test_client_agent_can_create_threads(self, agent_client: DurableAIAgentClient) -> None:
"""Verify agent from client can create DurableAgentThread instances."""
def test_client_agent_can_create_sessions(self, agent_client: DurableAIAgentClient) -> None:
"""Verify agent from client can create DurableAgentSession instances."""
agent = agent_client.get_agent("assistant")
thread = agent.get_new_thread()
session = agent.create_session()
assert isinstance(thread, DurableAgentThread)
assert isinstance(session, DurableAgentSession)
def test_client_agent_thread_with_parameters(self, agent_client: DurableAIAgentClient) -> None:
"""Verify agent can create threads with custom parameters."""
def test_client_agent_session_with_parameters(self, agent_client: DurableAIAgentClient) -> None:
"""Verify agent can create sessions with custom parameters."""
agent = agent_client.get_agent("assistant")
thread = agent.get_new_thread(service_thread_id="client-session-123")
session = agent.create_session(service_session_id="client-session-123")
assert isinstance(thread, DurableAgentThread)
assert thread.service_thread_id == "client-session-123"
assert isinstance(session, DurableAgentSession)
assert session.service_session_id == "client-session-123"
class TestDurableAIAgentClientPollingConfiguration:
@@ -16,7 +16,7 @@ from durabletask.entities import EntityInstanceId
from durabletask.task import Task
from pydantic import BaseModel
from agent_framework_durabletask import DurableAgentThread
from agent_framework_durabletask import DurableAgentSession
from agent_framework_durabletask._constants import DEFAULT_MAX_POLL_RETRIES, DEFAULT_POLL_INTERVAL_SECONDS
from agent_framework_durabletask._executors import (
ClientAgentExecutor,
@@ -106,42 +106,42 @@ def configure_failed_entity_task(mock_entity_task: Mock) -> Any:
return _configure
class TestExecutorThreadCreation:
"""Test that executors properly create DurableAgentThread with parameters."""
class TestExecutorSessionCreation:
"""Test that executors properly create DurableAgentSession with parameters."""
def test_client_executor_creates_durable_thread(self, mock_client: Mock) -> None:
"""Verify ClientAgentExecutor creates DurableAgentThread instances."""
def test_client_executor_creates_durable_session(self, mock_client: Mock) -> None:
"""Verify ClientAgentExecutor creates DurableAgentSession instances."""
executor = ClientAgentExecutor(mock_client)
thread = executor.get_new_thread("test_agent")
session = executor.get_new_session("test_agent")
assert isinstance(thread, DurableAgentThread)
assert isinstance(session, DurableAgentSession)
def test_client_executor_forwards_kwargs_to_thread(self, mock_client: Mock) -> None:
"""Verify ClientAgentExecutor forwards kwargs to DurableAgentThread creation."""
def test_client_executor_forwards_kwargs_to_session(self, mock_client: Mock) -> None:
"""Verify ClientAgentExecutor forwards kwargs to DurableAgentSession creation."""
executor = ClientAgentExecutor(mock_client)
thread = executor.get_new_thread("test_agent", service_thread_id="client-123")
session = executor.get_new_session("test_agent", service_session_id="client-123")
assert isinstance(thread, DurableAgentThread)
assert thread.service_thread_id == "client-123"
assert isinstance(session, DurableAgentSession)
assert session.service_session_id == "client-123"
def test_orchestration_executor_creates_durable_thread(
def test_orchestration_executor_creates_durable_session(
self, orchestration_executor: OrchestrationAgentExecutor
) -> None:
"""Verify OrchestrationAgentExecutor creates DurableAgentThread instances."""
thread = orchestration_executor.get_new_thread("test_agent")
"""Verify OrchestrationAgentExecutor creates DurableAgentSession instances."""
session = orchestration_executor.get_new_session("test_agent")
assert isinstance(thread, DurableAgentThread)
assert isinstance(session, DurableAgentSession)
def test_orchestration_executor_forwards_kwargs_to_thread(
def test_orchestration_executor_forwards_kwargs_to_session(
self, orchestration_executor: OrchestrationAgentExecutor
) -> None:
"""Verify OrchestrationAgentExecutor forwards kwargs to DurableAgentThread creation."""
thread = orchestration_executor.get_new_thread("test_agent", service_thread_id="orch-456")
"""Verify OrchestrationAgentExecutor forwards kwargs to DurableAgentSession creation."""
session = orchestration_executor.get_new_session("test_agent", service_session_id="orch-456")
assert isinstance(thread, DurableAgentThread)
assert thread.service_thread_id == "orch-456"
assert isinstance(session, DurableAgentSession)
assert session.service_session_id == "orch-456"
class TestClientAgentExecutorRun:
@@ -353,18 +353,18 @@ class TestOrchestrationAgentExecutorRun:
# Verify request dict
assert request_dict_arg == sample_run_request.to_dict()
def test_orchestration_executor_uses_thread_session_id(
def test_orchestration_executor_uses_session_durable_id(
self,
mock_orchestration_context: Mock,
orchestration_executor: OrchestrationAgentExecutor,
sample_run_request: RunRequest,
) -> None:
"""Verify executor uses thread's session ID when provided."""
# Create thread with specific session ID
"""Verify executor uses session's durable session ID when provided."""
# Create session with specific durable session ID
session_id = AgentSessionId(name="test_agent", key="specific-key-123")
thread = DurableAgentThread.from_session_id(session_id)
session = DurableAgentSession.from_session_id(session_id)
result = orchestration_executor.run_durable_agent("test_agent", sample_run_request, thread=thread)
result = orchestration_executor.run_durable_agent("test_agent", sample_run_request, session=session)
# Verify call_entity was called with the specific key
call_args = mock_orchestration_context.call_entity.call_args
@@ -11,7 +11,7 @@ from unittest.mock import Mock
import pytest
from agent_framework import SupportsAgentRun
from agent_framework_durabletask import DurableAgentThread
from agent_framework_durabletask import DurableAgentSession
from agent_framework_durabletask._orchestration_context import DurableAIAgentOrchestrationContext
from agent_framework_durabletask._shim import DurableAIAgent
@@ -74,24 +74,24 @@ class TestDurableAIAgentOrchestrationContextIntegration:
assert hasattr(agent, "run")
assert callable(agent.run)
def test_orchestration_agent_can_create_threads(self, agent_context: DurableAIAgentOrchestrationContext) -> None:
"""Verify agent from context can create DurableAgentThread instances."""
def test_orchestration_agent_can_create_sessions(self, agent_context: DurableAIAgentOrchestrationContext) -> None:
"""Verify agent from context can create DurableAgentSession instances."""
agent = agent_context.get_agent("assistant")
thread = agent.get_new_thread()
session = agent.create_session()
assert isinstance(thread, DurableAgentThread)
assert isinstance(session, DurableAgentSession)
def test_orchestration_agent_thread_with_parameters(
def test_orchestration_agent_session_with_parameters(
self, agent_context: DurableAIAgentOrchestrationContext
) -> None:
"""Verify agent can create threads with custom parameters."""
"""Verify agent can create sessions with custom parameters."""
agent = agent_context.get_agent("assistant")
thread = agent.get_new_thread(service_thread_id="orch-session-456")
session = agent.create_session(service_session_id="orch-session-456")
assert isinstance(thread, DurableAgentThread)
assert thread.service_thread_id == "orch-session-456"
assert isinstance(session, DurableAgentSession)
assert session.service_session_id == "orch-session-456"
if __name__ == "__main__":
+24 -24
View File
@@ -13,7 +13,7 @@ import pytest
from agent_framework import Message, SupportsAgentRun
from pydantic import BaseModel
from agent_framework_durabletask import DurableAgentThread
from agent_framework_durabletask import DurableAgentSession
from agent_framework_durabletask._executors import DurableAgentExecutor
from agent_framework_durabletask._models import RunRequest
from agent_framework_durabletask._shim import DurableAgentProvider, DurableAIAgent
@@ -30,7 +30,7 @@ def mock_executor() -> Mock:
"""Create a mock executor for testing."""
mock = Mock(spec=DurableAgentExecutor)
mock.run_durable_agent = Mock(return_value=None)
mock.get_new_thread = Mock(return_value=DurableAgentThread())
mock.get_new_session = Mock(return_value=DurableAgentSession())
# Mock get_run_request to create actual RunRequest objects
def create_run_request(
@@ -124,14 +124,14 @@ class TestDurableAIAgentMessageNormalization:
class TestDurableAIAgentParameterFlow:
"""Test that parameters flow correctly through the shim to executor."""
def test_run_forwards_thread_parameter(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
"""Verify run forwards thread parameter to executor."""
thread = DurableAgentThread(service_thread_id="test-thread")
test_agent.run("message", thread=thread)
def test_run_forwards_session_parameter(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
"""Verify run forwards session parameter to executor."""
session = DurableAgentSession(service_session_id="test-session")
test_agent.run("message", session=session)
mock_executor.run_durable_agent.assert_called_once()
_, kwargs = mock_executor.run_durable_agent.call_args
assert kwargs["thread"] == thread
assert kwargs["session"] == session
def test_run_forwards_response_format(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
"""Verify run forwards response_format parameter to executor."""
@@ -171,29 +171,29 @@ class TestDurableAISupportsAgentRunCompliance:
assert agent.name == "my_agent"
class TestDurableAIAgentThreadManagement:
"""Test thread creation and management."""
class TestDurableAIAgentSessionManagement:
"""Test session creation and management."""
def test_get_new_thread_delegates_to_executor(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
"""Verify get_new_thread delegates to executor."""
mock_thread = DurableAgentThread()
mock_executor.get_new_thread.return_value = mock_thread
def test_create_session_delegates_to_executor(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
"""Verify create_session delegates to executor."""
mock_session = DurableAgentSession()
mock_executor.get_new_session.return_value = mock_session
thread = test_agent.get_new_thread()
session = test_agent.create_session()
mock_executor.get_new_thread.assert_called_once_with("test_agent")
assert thread == mock_thread
mock_executor.get_new_session.assert_called_once_with("test_agent")
assert session == mock_session
def test_get_new_thread_forwards_kwargs(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
"""Verify get_new_thread forwards kwargs to executor."""
mock_thread = DurableAgentThread(service_thread_id="thread-123")
mock_executor.get_new_thread.return_value = mock_thread
def test_create_session_forwards_kwargs(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
"""Verify create_session forwards kwargs to executor."""
mock_session = DurableAgentSession(service_session_id="session-123")
mock_executor.get_new_session.return_value = mock_session
test_agent.get_new_thread(service_thread_id="thread-123")
test_agent.create_session(service_session_id="session-123")
mock_executor.get_new_thread.assert_called_once()
_, kwargs = mock_executor.get_new_thread.call_args
assert kwargs["service_thread_id"] == "thread-123"
mock_executor.get_new_session.assert_called_once()
_, kwargs = mock_executor.get_new_session.call_args
assert kwargs["service_session_id"] == "session-123"
class TestDurableAgentProviderInterface: